toki
c46874055a
- Add edge node unit tests and transport package - Add node test client and related artifacts - Update bootstrap, node, and config modules - Add proto generated files - Update Makefile and configuration files
71 lines
1.3 KiB
Go
71 lines
1.3 KiB
Go
package node
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
|
|
toki "git.toki-labs.com/toki/common-proto-socket/go"
|
|
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
// NodeEntry represents one connected node.
|
|
type NodeEntry struct {
|
|
NodeID string
|
|
Client *toki.TcpClient
|
|
Adapters []*iop.AdapterInfo
|
|
}
|
|
|
|
// Registry manages all nodes connected to edge.
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
nodes map[string]*NodeEntry
|
|
}
|
|
|
|
func NewRegistry() *Registry {
|
|
return &Registry{nodes: make(map[string]*NodeEntry)}
|
|
}
|
|
|
|
func (r *Registry) Register(entry *NodeEntry) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.nodes[entry.NodeID] = entry
|
|
}
|
|
|
|
func (r *Registry) Unregister(nodeID string) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
delete(r.nodes, nodeID)
|
|
}
|
|
|
|
func (r *Registry) Get(nodeID string) (*NodeEntry, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
e, ok := r.nodes[nodeID]
|
|
return e, ok
|
|
}
|
|
|
|
func (r *Registry) All() []*NodeEntry {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
out := make([]*NodeEntry, 0, len(r.nodes))
|
|
for _, e := range r.nodes {
|
|
out = append(out, e)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (r *Registry) Count() int {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return len(r.nodes)
|
|
}
|
|
|
|
func (r *Registry) Pick() (*NodeEntry, error) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
for _, e := range r.nodes {
|
|
return e, nil
|
|
}
|
|
return nil, fmt.Errorf("no nodes connected")
|
|
}
|