102 lines
2 KiB
Go
102 lines
2 KiB
Go
package node
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
|
|
toki "git.toki-labs.com/toki/common-proto-socket/go"
|
|
)
|
|
|
|
// NodeEntry represents one connected node.
|
|
type NodeEntry struct {
|
|
NodeID string
|
|
Alias string
|
|
Client *toki.TcpClient
|
|
}
|
|
|
|
// Registry manages all nodes connected to edge.
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
byID map[string]*NodeEntry
|
|
byAlias map[string]*NodeEntry
|
|
}
|
|
|
|
func NewRegistry() *Registry {
|
|
return &Registry{
|
|
byID: make(map[string]*NodeEntry),
|
|
byAlias: make(map[string]*NodeEntry),
|
|
}
|
|
}
|
|
|
|
func (r *Registry) Register(entry *NodeEntry) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.byID[entry.NodeID] = entry
|
|
if entry.Alias != "" {
|
|
r.byAlias[entry.Alias] = entry
|
|
}
|
|
}
|
|
|
|
func (r *Registry) Unregister(nodeID string) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if entry, ok := r.byID[nodeID]; ok {
|
|
if entry.Alias != "" {
|
|
delete(r.byAlias, entry.Alias)
|
|
}
|
|
delete(r.byID, nodeID)
|
|
}
|
|
}
|
|
|
|
func (r *Registry) Get(nodeID string) (*NodeEntry, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
e, ok := r.byID[nodeID]
|
|
return e, ok
|
|
}
|
|
|
|
func (r *Registry) Resolve(ref string) (*NodeEntry, error) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
if ref != "" {
|
|
if entry, ok := r.byID[ref]; ok {
|
|
return entry, nil
|
|
}
|
|
if entry, ok := r.byAlias[ref]; ok {
|
|
return entry, nil
|
|
}
|
|
return nil, fmt.Errorf("node %q not found", ref)
|
|
}
|
|
|
|
if len(r.byID) == 1 {
|
|
for _, entry := range r.byID {
|
|
return entry, nil
|
|
}
|
|
}
|
|
if len(r.byID) == 0 {
|
|
return nil, fmt.Errorf("no nodes connected")
|
|
}
|
|
return nil, fmt.Errorf("multiple nodes connected; select one with /node <id|alias>")
|
|
}
|
|
|
|
func (r *Registry) All() []*NodeEntry {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
out := make([]*NodeEntry, 0, len(r.byID))
|
|
for _, e := range r.byID {
|
|
out = append(out, e)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (r *Registry) Count() int {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return len(r.byID)
|
|
}
|
|
|
|
// Pick is deprecated: use Resolve("") instead for single-node fallback, or Resolve(ref) for explicit selection.
|
|
func (r *Registry) Pick() (*NodeEntry, error) {
|
|
return r.Resolve("")
|
|
}
|