kami
120 lines
2.4 KiB
Go
120 lines
2.4 KiB
Go
package federation
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var ErrUnknownWorker = errors.New("unknown worker")
|
|
var ErrUnauthorized = errors.New("worker authentication failed")
|
|
|
|
type Worker struct {
|
|
ID string `json:"id"`
|
|
Address string `json:"address"`
|
|
Capacity int `json:"capacity"`
|
|
LastSeen time.Time `json:"last_seen"`
|
|
Online bool `json:"online"`
|
|
Token string `json:"-"`
|
|
}
|
|
|
|
type Registry struct {
|
|
mu sync.Mutex
|
|
workers map[string]Worker
|
|
TTL time.Duration
|
|
OnOffline func(Worker)
|
|
cursors map[string]uint64
|
|
}
|
|
|
|
func (r *Registry) init() {
|
|
if r.TTL <= 0 {
|
|
r.TTL = 90 * time.Second
|
|
}
|
|
if r.workers == nil {
|
|
r.workers = map[string]Worker{}
|
|
}
|
|
if r.cursors == nil {
|
|
r.cursors = map[string]uint64{}
|
|
}
|
|
}
|
|
func (r *Registry) Register(w Worker) error {
|
|
if w.ID == "" {
|
|
return errors.New("worker id required")
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.init()
|
|
w.LastSeen = time.Now().UTC()
|
|
w.Online = true
|
|
r.workers[w.ID] = w
|
|
if _, ok := r.cursors[w.ID]; !ok {
|
|
r.cursors[w.ID] = 0
|
|
}
|
|
return nil
|
|
}
|
|
func (r *Registry) Authenticate(id, token string) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.init()
|
|
w, ok := r.workers[id]
|
|
if !ok {
|
|
return ErrUnknownWorker
|
|
}
|
|
if w.Token == "" || token == "" || w.Token != token {
|
|
return ErrUnauthorized
|
|
}
|
|
return nil
|
|
}
|
|
func (r *Registry) Cursor(id string) (uint64, error) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.init()
|
|
if _, ok := r.workers[id]; !ok {
|
|
return 0, ErrUnknownWorker
|
|
}
|
|
return r.cursors[id], nil
|
|
}
|
|
func (r *Registry) Ack(id string, cursor uint64) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.init()
|
|
if _, ok := r.workers[id]; !ok {
|
|
return ErrUnknownWorker
|
|
}
|
|
if cursor < r.cursors[id] {
|
|
return errors.New("cursor moved backwards")
|
|
}
|
|
r.cursors[id] = cursor
|
|
return nil
|
|
}
|
|
func (r *Registry) Heartbeat(id string) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.init()
|
|
w, ok := r.workers[id]
|
|
if !ok {
|
|
return ErrUnknownWorker
|
|
}
|
|
w.LastSeen = time.Now().UTC()
|
|
w.Online = true
|
|
r.workers[id] = w
|
|
return nil
|
|
}
|
|
func (r *Registry) Snapshot() []Worker {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.init()
|
|
now := time.Now().UTC()
|
|
out := make([]Worker, 0, len(r.workers))
|
|
for id, w := range r.workers {
|
|
wasOnline := w.Online
|
|
w.Online = now.Sub(w.LastSeen) <= r.TTL
|
|
r.workers[id] = w
|
|
if wasOnline && !w.Online && r.OnOffline != nil {
|
|
go r.OnOffline(w)
|
|
}
|
|
out = append(out, w)
|
|
}
|
|
return out
|
|
}
|