notify/server_peer_registry.go

package notify

import (
	"errors"
	"net"
	"sort"
	"sync"
)

type serverPeerRegistry struct {
	mu    sync.RWMutex
	peers map[string]*LogicalConn
}

func newServerPeerRegistry() *serverPeerRegistry {
	return &serverPeerRegistry{
		peers: make(map[string]*LogicalConn),
	}
}

func (s *ServerCommon) getPeerRegistry() *serverPeerRegistry {
	if s == nil {
		return nil
	}
	s.mu.Lock()
	if s.peerRegistry == nil {
		s.peerRegistry = newServerPeerRegistry()
	}
	registry := s.peerRegistry
	s.mu.Unlock()
	return registry
}

func (r *serverPeerRegistry) registerClient(client *ClientConn) *LogicalConn {
	return r.registerLogical(logicalConnFromClient(client))
}

func (r *serverPeerRegistry) registerLogical(logical *LogicalConn) *LogicalConn {
	if r == nil || logical == nil {
		return nil
	}
	r.mu.Lock()
	if logicalID := logical.ID(); logicalID != "" {
		r.peers[logicalID] = logical
	}
	r.mu.Unlock()
	return logical
}

func (r *serverPeerRegistry) getLogical(id string) *LogicalConn {
	if r == nil {
		return nil
	}
	r.mu.RLock()
	logical := r.peers[id]
	r.mu.RUnlock()
	return logical
}

func (r *serverPeerRegistry) hasID(id string) bool {
	if r == nil {
		return false
	}
	r.mu.RLock()
	_, ok := r.peers[id]
	r.mu.RUnlock()
	return ok
}

func (r *serverPeerRegistry) renameLogical(logical *LogicalConn, id string) error {
	if r == nil {
		return errors.New("peer registry is nil")
	}
	if logical == nil {
		return errors.New("logical conn is nil")
	}
	if id == "" {
		return errors.New("client id is empty")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if existing, ok := r.peers[id]; ok && existing != logical {
		return errors.New("client id already exists")
	}
	if currentID := logical.ID(); currentID != "" {
		if existing, ok := r.peers[currentID]; ok && existing == logical {
			delete(r.peers, currentID)
		}
	}
	logical.setID(id)
	r.peers[id] = logical
	return nil
}

func (r *serverPeerRegistry) removeLogical(logical *LogicalConn) {
	if r == nil || logical == nil {
		return
	}
	r.mu.Lock()
	if currentID := logical.ID(); currentID != "" {
		if existing, ok := r.peers[currentID]; ok && existing == logical {
			delete(r.peers, currentID)
		}
	}
	for id, existing := range r.peers {
		if existing == logical {
			delete(r.peers, id)
		}
	}
	r.mu.Unlock()
}

func (r *serverPeerRegistry) resolveLogicalBySource(source string) *LogicalConn {
	if r == nil {
		return nil
	}
	r.mu.RLock()
	if logical, ok := r.peers[source]; ok {
		r.mu.RUnlock()
		return logical
	}
	var match *LogicalConn
	for _, logical := range r.peers {
		addr := logical.RemoteAddr()
		if addr == nil {
			continue
		}
		if addr.String() == source {
			if match != nil && match != logical {
				r.mu.RUnlock()
				return nil
			}
			match = logical
		}
	}
	r.mu.RUnlock()
	return match
}

func (r *serverPeerRegistry) logicalList() []*LogicalConn {
	if r == nil {
		return nil
	}
	r.mu.RLock()
	list := make([]*LogicalConn, 0, len(r.peers))
	for _, logical := range r.peers {
		if logical != nil {
			list = append(list, logical)
		}
	}
	r.mu.RUnlock()
	sort.Slice(list, func(i int, j int) bool {
		left := list[i]
		right := list[j]
		if left == nil || right == nil {
			return left != nil
		}
		if left.ID() == right.ID() {
			return addrString(left.RemoteAddr()) < addrString(right.RemoteAddr())
		}
		return left.ID() < right.ID()
	})
	return list
}

func (r *serverPeerRegistry) detachedLogicals() []*LogicalConn {
	list := r.logicalList()
	filtered := make([]*LogicalConn, 0, len(list))
	for _, logical := range list {
		if logical == nil || !logical.logicalTransportDetachedSnapshot() {
			continue
		}
		filtered = append(filtered, logical)
	}
	return filtered
}

func addrString(addr net.Addr) string {
	if addr == nil {
		return ""
	}
	return addr.String()
}