sendico/api/pkg/discovery/service.go

package discovery

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"sync"
	"time"

	msg "github.com/tech/sendico/pkg/messaging"
	"github.com/tech/sendico/pkg/messaging/broker"
	cons "github.com/tech/sendico/pkg/messaging/consumer"
	me "github.com/tech/sendico/pkg/messaging/envelope"
	"github.com/tech/sendico/pkg/mlogger"
	"go.uber.org/zap"
)

type RegistryService struct {
	logger    mlogger.Logger
	registry  *Registry
	producer  msg.Producer
	sender    string
	consumers []consumerHandler

	startOnce sync.Once
	stopOnce  sync.Once
}

type consumerHandler struct {
	consumer msg.Consumer
	handler  msg.MessageHandlerT
}

func NewRegistryService(logger mlogger.Logger, broker broker.Broker, producer msg.Producer, registry *Registry, sender string) (*RegistryService, error) {
	if broker == nil {
		return nil, errors.New("discovery registry: broker is nil")
	}
	if registry == nil {
		registry = NewRegistry()
	}
	if logger != nil {
		logger = logger.Named("discovery_registry")
	}
	sender = strings.TrimSpace(sender)
	if sender == "" {
		sender = "discovery"
	}

	serviceConsumer, err := cons.NewConsumer(logger, broker, ServiceAnnounceEvent())
	if err != nil {
		return nil, err
	}
	gatewayConsumer, err := cons.NewConsumer(logger, broker, GatewayAnnounceEvent())
	if err != nil {
		return nil, err
	}
	heartbeatConsumer, err := cons.NewConsumer(logger, broker, HeartbeatEvent())
	if err != nil {
		return nil, err
	}
	lookupConsumer, err := cons.NewConsumer(logger, broker, LookupRequestEvent())
	if err != nil {
		return nil, err
	}

	svc := &RegistryService{
		logger:   logger,
		registry: registry,
		producer: producer,
		sender:   sender,
		consumers: []consumerHandler{
			{consumer: serviceConsumer, handler: func(ctx context.Context, env me.Envelope) error {
				return svc.handleAnnounce(ctx, env)
			}},
			{consumer: gatewayConsumer, handler: func(ctx context.Context, env me.Envelope) error {
				return svc.handleAnnounce(ctx, env)
			}},
			{consumer: heartbeatConsumer, handler: svc.handleHeartbeat},
			{consumer: lookupConsumer, handler: svc.handleLookup},
		},
	}
	return svc, nil
}

func (s *RegistryService) Start() {
	if s == nil {
		return
	}
	s.startOnce.Do(func() {
		for _, ch := range s.consumers {
			ch := ch
			go func() {
				if err := ch.consumer.ConsumeMessages(ch.handler); err != nil && s.logger != nil {
					s.logger.Warn("Discovery consumer stopped with error", zap.Error(err))
				}
			}()
		}
	})
}

func (s *RegistryService) Stop() {
	if s == nil {
		return
	}
	s.stopOnce.Do(func() {
		for _, ch := range s.consumers {
			if ch.consumer != nil {
				ch.consumer.Close()
			}
		}
	})
}

func (s *RegistryService) handleAnnounce(_ context.Context, env me.Envelope) error {
	var payload Announcement
	if err := json.Unmarshal(env.GetData(), &payload); err != nil {
		s.logWarn("Failed to decode discovery announce payload", zap.Error(err))
		return err
	}
	now := time.Now()
	result := s.registry.UpsertFromAnnouncement(payload, now)
	if result.IsNew || result.BecameHealthy {
		s.publishRefresh(result.Entry)
	}
	return nil
}

func (s *RegistryService) handleHeartbeat(_ context.Context, env me.Envelope) error {
	var payload Heartbeat
	if err := json.Unmarshal(env.GetData(), &payload); err != nil {
		s.logWarn("Failed to decode discovery heartbeat payload", zap.Error(err))
		return err
	}
	if payload.ID == "" {
		return nil
	}
	ts := time.Unix(payload.TS, 0)
	if ts.Unix() <= 0 {
		ts = time.Now()
	}
	result, ok := s.registry.UpdateHeartbeat(payload.ID, strings.TrimSpace(payload.Status), ts, time.Now())
	if ok && result.BecameHealthy {
		s.publishRefresh(result.Entry)
	}
	return nil
}

func (s *RegistryService) handleLookup(_ context.Context, env me.Envelope) error {
	if s.producer == nil {
		s.logWarn("Discovery lookup request ignored: producer not configured")
		return nil
	}
	var payload LookupRequest
	if err := json.Unmarshal(env.GetData(), &payload); err != nil {
		s.logWarn("Failed to decode discovery lookup payload", zap.Error(err))
		return err
	}
	resp := s.registry.Lookup(time.Now())
	resp.RequestID = strings.TrimSpace(payload.RequestID)
	if err := s.producer.SendMessage(NewLookupResponseEnvelope(s.sender, resp)); err != nil {
		s.logWarn("Failed to publish discovery lookup response", zap.Error(err))
		return err
	}
	return nil
}

func (s *RegistryService) publishRefresh(entry RegistryEntry) {
	if s == nil || s.producer == nil {
		return
	}
	payload := RefreshEvent{
		Service: entry.Service,
		Rail:    entry.Rail,
		Network: entry.Network,
		Message: "new module available",
	}
	if err := s.producer.SendMessage(NewRefreshUIEnvelope(s.sender, payload)); err != nil {
		s.logWarn("Failed to publish discovery refresh event", zap.Error(err))
	}
}

func (s *RegistryService) logWarn(message string, fields ...zap.Field) {
	if s.logger == nil {
		return
	}
	s.logger.Warn(message, fields...)
}