kami
71 lines
1.3 KiB
Go
71 lines
1.3 KiB
Go
package federation
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var ErrUnknownWorker = errors.New("unknown worker")
|
|
|
|
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"`
|
|
}
|
|
|
|
type Registry struct {
|
|
mu sync.Mutex
|
|
workers map[string]Worker
|
|
TTL time.Duration
|
|
}
|
|
|
|
func (r *Registry) init() {
|
|
if r.TTL <= 0 {
|
|
r.TTL = 90 * time.Second
|
|
}
|
|
if r.workers == nil {
|
|
r.workers = map[string]Worker{}
|
|
}
|
|
}
|
|
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
|
|
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()
|
|
out := make([]Worker, 0, len(r.workers))
|
|
for id, w := range r.workers {
|
|
w.Online = now.Sub(w.LastSeen) <= r.TTL
|
|
r.workers[id] = w
|
|
out = append(out, w)
|
|
}
|
|
return out
|
|
}
|