135 lines
3.2 KiB
Go
135 lines
3.2 KiB
Go
package discovery
|
|
|
|
import (
|
|
"encoding/json"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/nats-io/nats.go"
|
|
mb "github.com/tech/sendico/pkg/messaging/broker"
|
|
"github.com/tech/sendico/pkg/merrors"
|
|
"github.com/tech/sendico/pkg/mlogger"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type RegistryWatcher struct {
|
|
logger mlogger.Logger
|
|
registry *Registry
|
|
kv *KVStore
|
|
watcher nats.KeyWatcher
|
|
|
|
stopOnce sync.Once
|
|
}
|
|
|
|
func NewRegistryWatcher(logger mlogger.Logger, msgBroker mb.Broker, registry *Registry) (*RegistryWatcher, error) {
|
|
if msgBroker == nil {
|
|
return nil, merrors.InvalidArgument("discovery watcher: broker is nil")
|
|
}
|
|
if registry == nil {
|
|
registry = NewRegistry()
|
|
}
|
|
if logger == nil {
|
|
return nil, merrors.InvalidArgument("discovery logger: logger must be provided")
|
|
}
|
|
logger = logger.Named("discovery_watcher")
|
|
provider, ok := msgBroker.(jetStreamProvider)
|
|
if !ok {
|
|
return nil, merrors.Internal("discovery watcher: jetstream not available")
|
|
}
|
|
js := provider.JetStream()
|
|
if js == nil {
|
|
return nil, merrors.Internal("discovery watcher: jetstream not configured")
|
|
}
|
|
store, err := NewKVStore(logger, js, "")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &RegistryWatcher{
|
|
logger: logger,
|
|
registry: registry,
|
|
kv: store,
|
|
}, nil
|
|
}
|
|
|
|
func (w *RegistryWatcher) Start() error {
|
|
if w == nil || w.kv == nil {
|
|
return merrors.Internal("discovery watcher: not configured")
|
|
}
|
|
watcher, err := w.kv.WatchAll()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
w.watcher = watcher
|
|
w.logger.Info("Discovery registry watcher started", zap.String("bucket", w.kv.Bucket()))
|
|
go w.consume(watcher)
|
|
return nil
|
|
}
|
|
|
|
func (w *RegistryWatcher) Stop() {
|
|
if w == nil {
|
|
return
|
|
}
|
|
w.stopOnce.Do(func() {
|
|
if w.watcher != nil {
|
|
_ = w.watcher.Stop()
|
|
}
|
|
w.logger.Info("Discovery registry watcher stopped")
|
|
})
|
|
}
|
|
|
|
func (w *RegistryWatcher) Registry() *Registry {
|
|
if w == nil {
|
|
return nil
|
|
}
|
|
return w.registry
|
|
}
|
|
|
|
func (w *RegistryWatcher) consume(watcher nats.KeyWatcher) {
|
|
if w == nil || watcher == nil {
|
|
return
|
|
}
|
|
initial := true
|
|
initialCount := 0
|
|
for entry := range watcher.Updates() {
|
|
if entry == nil {
|
|
if initial {
|
|
fields := []zap.Field{zap.Int("entries", initialCount)}
|
|
if w.kv != nil {
|
|
fields = append(fields, zap.String("bucket", w.kv.Bucket()))
|
|
}
|
|
w.logger.Info("Discovery registry watcher initial sync complete", fields...)
|
|
initial = false
|
|
}
|
|
continue
|
|
}
|
|
if initial && entry.Operation() == nats.KeyValuePut {
|
|
initialCount++
|
|
}
|
|
switch entry.Operation() {
|
|
case nats.KeyValueDelete, nats.KeyValuePurge:
|
|
key := registryKeyFromKVKey(entry.Key())
|
|
if key != "" {
|
|
if w.registry.Delete(key) {
|
|
w.logger.Info("Discovery registry entry removed", zap.String("key", key))
|
|
}
|
|
}
|
|
continue
|
|
case nats.KeyValuePut:
|
|
default:
|
|
continue
|
|
}
|
|
|
|
var payload RegistryEntry
|
|
if err := json.Unmarshal(entry.Value(), &payload); err != nil {
|
|
w.logger.Warn("Failed to decode discovery KV entry", zap.String("key", entry.Key()), zap.Error(err))
|
|
continue
|
|
}
|
|
result := w.registry.UpsertEntry(payload, time.Now())
|
|
if result.IsNew || result.BecameHealthy {
|
|
fields := append(entryFields(result.Entry), zap.Bool("is_new", result.IsNew), zap.Bool("became_healthy", result.BecameHealthy))
|
|
w.logger.Info("Discovery registry entry updated from KV", fields...)
|
|
}
|
|
}
|
|
}
|