159 lines
2.8 KiB
Go
159 lines
2.8 KiB
Go
package registry
|
|
|
|
import (
|
|
"maps"
|
|
"slices"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.wzray.com/homelab/hivemind/internal/types"
|
|
"github.com/rs/zerolog/log"
|
|
)
|
|
|
|
type Registry struct {
|
|
LastUpdate time.Time
|
|
nodes map[string]types.Node
|
|
storage Storage
|
|
lock sync.RWMutex
|
|
self types.Node
|
|
observers []chan<- RegistryEvent
|
|
}
|
|
|
|
func New(storage Storage, self types.Node) *Registry {
|
|
r := &Registry{
|
|
storage: storage,
|
|
nodes: make(map[string]types.Node),
|
|
self: self,
|
|
}
|
|
|
|
var storedData storedConfig
|
|
if err := storage.Load(&storedData); err != nil {
|
|
log.Warn().Err(err).Msg("unable to load registry from storage")
|
|
goto ret
|
|
}
|
|
|
|
r.LastUpdate = time.UnixMilli(storedData.LastUpdate)
|
|
r.nodes = storedData.Nodes
|
|
|
|
ret:
|
|
r.nodes[self.Hostname] = self
|
|
return r
|
|
}
|
|
|
|
func (r *Registry) snapshot() *storedConfig {
|
|
return &storedConfig{
|
|
LastUpdate: r.LastUpdate.UnixMilli(),
|
|
Nodes: maps.Clone(r.nodes),
|
|
}
|
|
}
|
|
|
|
func (r *Registry) notify(event RegistryEvent) {
|
|
for _, c := range r.observers {
|
|
c <- event
|
|
}
|
|
}
|
|
|
|
func (r *Registry) AllNodes() map[string]types.Node {
|
|
r.lock.RLock()
|
|
defer r.lock.RUnlock()
|
|
return maps.Clone(r.nodes)
|
|
}
|
|
|
|
func (r *Registry) Nodes() map[string]types.Node {
|
|
nodes := r.AllNodes()
|
|
delete(nodes, r.self.Hostname)
|
|
return nodes
|
|
}
|
|
|
|
func (r *Registry) ByRole(role types.Role) map[string]types.Node {
|
|
r.lock.RLock()
|
|
defer r.lock.RUnlock()
|
|
|
|
o := make(map[string]types.Node)
|
|
for name, node := range r.nodes {
|
|
if slices.Contains(node.Roles, role) && node.Hostname != r.self.Hostname {
|
|
o[name] = node
|
|
}
|
|
}
|
|
|
|
return o
|
|
}
|
|
|
|
func (r *Registry) AddNode(node types.Node) error {
|
|
r.lock.Lock()
|
|
r.nodes[node.Hostname] = node
|
|
r.LastUpdate = time.Now()
|
|
snapshot := r.snapshot()
|
|
r.lock.Unlock()
|
|
|
|
if err := r.storage.Save(snapshot); err != nil {
|
|
return err
|
|
}
|
|
|
|
r.notify(RegistryEvent{
|
|
EventNodeJoin,
|
|
map[string]types.Node{
|
|
node.Hostname: node,
|
|
},
|
|
})
|
|
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) RemoveNode(node types.Node) error {
|
|
r.lock.Lock()
|
|
delete(r.nodes, node.Hostname)
|
|
r.LastUpdate = time.Now()
|
|
snapshot := r.snapshot()
|
|
r.lock.Unlock()
|
|
|
|
if err := r.storage.Save(snapshot); err != nil {
|
|
return err
|
|
}
|
|
|
|
r.notify(RegistryEvent{
|
|
EventNodeLeave,
|
|
map[string]types.Node{
|
|
node.Hostname: node,
|
|
},
|
|
})
|
|
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) Set(nodes map[string]types.Node) error {
|
|
r.lock.Lock()
|
|
r.nodes = maps.Clone(nodes)
|
|
snapshot := r.snapshot()
|
|
r.lock.Unlock()
|
|
|
|
if err := r.storage.Save(snapshot); err != nil {
|
|
return err
|
|
}
|
|
|
|
r.notify(RegistryEvent{
|
|
EventSet,
|
|
nodes,
|
|
})
|
|
|
|
return nil
|
|
}
|
|
|
|
func (r *Registry) Exists(name string) bool {
|
|
_, ok := r.nodes[name]
|
|
return ok
|
|
}
|
|
|
|
func (r *Registry) Subscribe() <-chan RegistryEvent { // TODO: rename this
|
|
c := make(chan RegistryEvent, 1)
|
|
r.observers = append(r.observers, c)
|
|
return c
|
|
}
|
|
|
|
func (r *Registry) Save() {
|
|
r.lock.RLock()
|
|
snapshot := r.snapshot()
|
|
r.lock.RUnlock()
|
|
|
|
r.storage.Save(snapshot)
|
|
}
|