feat: add ingestion services and test helpers

Author: sirrobot01

internal/ingestion/service.go

@@ -0,0 +1,354 @@
+package ingestion
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/sirrobot01/mcpulse/internal/config"
+	"github.com/sirrobot01/mcpulse/internal/database"
+	"github.com/sirrobot01/mcpulse/internal/metrics"
+	"github.com/sirrobot01/mcpulse/internal/models"
+	"go.uber.org/zap"
+	"golang.org/x/time/rate"
+)
+
+// ErrRateLimited indicates that the ingestion request exceeded the configured rate limit.
+var ErrRateLimited = errors.New("rate limit exceeded")
+
+// ErrBufferFull indicates that the ingestion buffer is full and cannot accept more work.
+var ErrBufferFull = errors.New("ingestion buffer is full")
+
+const (
+	eventTypeToolCall = "tool_call"
+)
+
+// Event represents a real-time ingestion event that can be streamed to subscribers.
+type Event struct {
+	Type     string
+	Metric   models.ToolCallMetric
+	ServerID string
+}
+
+type subscriber struct {
+	id       int64
+	serverID string
+	topics   map[string]struct{}
+	ch       chan Event
+}
+
+// Service implements ingestion logic including validation, sanitization, persistence,
+// and optional real-time fan-out.
+type Service struct {
+	cfg     config.IngestionConfig
+	privacy config.PrivacyConfig
+	repo    *database.Repository
+	metrics *metrics.Metrics
+	logger  *zap.Logger
+
+	limiter *rate.Limiter
+
+	subsMu       sync.RWMutex
+	subscribers  map[int64]*subscriber
+	nextSubID    atomic.Int64
+	subscriberCh int
+}
+
+// NewService constructs a new ingestion service.
+func NewService(
+	cfg config.IngestionConfig,
+	privacy config.PrivacyConfig,
+	repo *database.Repository,
+	metricsCollector *metrics.Metrics,
+	logger *zap.Logger,
+) *Service {
+	var limiter *rate.Limiter
+	if cfg.RateLimit.Enabled && cfg.RateLimit.RequestsPerSecond > 0 {
+		limiter = rate.NewLimiter(rate.Limit(cfg.RateLimit.RequestsPerSecond), cfg.RateLimit.Burst)
+	}
+
+	s := &Service{
+		cfg:          cfg,
+		privacy:      privacy,
+		repo:         repo,
+		metrics:      metricsCollector,
+		logger:       logger,
+		limiter:      limiter,
+		subscribers:  make(map[int64]*subscriber),
+		subscriberCh: max(cfg.BatchSize, 64),
+	}
+
+	return s
+}
+
+// Ingest validates, sanitizes, and persists the provided metrics batch.
+func (s *Service) Ingest(ctx context.Context, batch []models.ToolCallMetric) (models.IngestResponse, error) {
+	resp := models.IngestResponse{}
+
+	if len(batch) == 0 {
+		return resp, nil
+	}
+
+	if s.limiter != nil && !s.limiter.AllowN(time.Now(), len(batch)) {
+		s.metrics.RecordIngest(eventTypeToolCall, false)
+		return resp, ErrRateLimited
+	}
+
+	validMetrics := make([]models.ToolCallMetric, 0, len(batch))
+	for idx := range batch {
+		metric := &batch[idx]
+		if err := s.prepareMetric(metric); err != nil {
+			resp.Rejected++
+			resp.Errors = append(resp.Errors, models.IngestError{
+				Index:  idx,
+				Reason: err.Error(),
+			})
+			s.metrics.RecordIngestError("validation")
+			continue
+		}
+
+		validMetrics = append(validMetrics, *metric)
+		s.metrics.RecordIngest(eventTypeToolCall, true)
+	}
+
+	if len(validMetrics) == 0 {
+		return resp, nil
+	}
+
+	start := time.Now()
+	if err := s.repo.InsertToolCallBatch(ctx, validMetrics); err != nil {
+		s.metrics.RecordIngestError("storage")
+		s.logger.Error("failed to insert tool call batch", zap.Error(err))
+		return resp, fmt.Errorf("failed to persist metrics: %w", err)
+	}
+	s.metrics.IngestDuration.WithLabelValues("persist_tool_calls").Observe(time.Since(start).Seconds())
+
+	resp.Accepted = len(validMetrics)
+
+	// Update server metadata for each unique server ID in the batch.
+	s.upsertServers(ctx, validMetrics)
+
+	// Broadcast to subscribers for real-time streaming.
+	for _, metric := range validMetrics {
+		s.broadcast(Event{
+			Type:     eventTypeToolCall,
+			Metric:   metric,
+			ServerID: metric.ServerID,
+		})
+	}
+
+	return resp, nil
+}
+
+// Subscribe registers a subscriber for real-time ingestion events. The caller MUST call the
+// returned cancel function when they are done consuming events.
+func (s *Service) Subscribe(serverID string, topics []string) (<-chan Event, func()) {
+	ch := make(chan Event, s.subscriberCh)
+	sub := &subscriber{
+		id:       s.nextSubID.Add(1),
+		serverID: serverID,
+		topics:   make(map[string]struct{}),
+		ch:       ch,
+	}
+
+	if len(topics) == 0 {
+		sub.topics[eventTypeToolCall] = struct{}{}
+	} else {
+		for _, topic := range topics {
+			sub.topics[strings.ToLower(strings.TrimSpace(topic))] = struct{}{}
+		}
+	}
+
+	s.subsMu.Lock()
+	s.subscribers[sub.id] = sub
+	s.subsMu.Unlock()
+
+	cancel := func() {
+		s.subsMu.Lock()
+		delete(s.subscribers, sub.id)
+		s.subsMu.Unlock()
+		close(ch)
+	}
+
+	return ch, cancel
+}
+
+// prepareMetric validates and sanitizes a single metric entry.
+func (s *Service) prepareMetric(metric *models.ToolCallMetric) error {
+	if metric.ID == "" {
+		return errors.New("missing required field: id")
+	}
+	if metric.ServerID == "" {
+		return errors.New("missing required field: server_id")
+	}
+	if metric.ToolName == "" {
+		return errors.New("missing required field: tool_name")
+	}
+	if metric.Status == "" {
+		return errors.New("missing required field: status")
+	}
+	if metric.DurationMS < 0 {
+		return errors.New("duration_ms must be >= 0")
+	}
+
+	if metric.Timestamp.IsZero() {
+		metric.Timestamp = time.Now().UTC()
+	} else {
+		metric.Timestamp = metric.Timestamp.UTC()
+	}
+
+	status := strings.ToLower(metric.Status)
+	switch status {
+	case "success", "error", "timeout":
+	default:
+		return fmt.Errorf("invalid status: %s", metric.Status)
+	}
+	metric.Status = status
+
+	if metric.Status == "error" && (metric.ErrorMessage == nil || strings.TrimSpace(*metric.ErrorMessage) == "") {
+		return errors.New("error_message required when status=error")
+	}
+
+	if metric.SessionID != nil {
+		sessionID := strings.TrimSpace(*metric.SessionID)
+		if sessionID == "" {
+			metric.SessionID = nil
+		} else if s.privacy.HashSessionIDs {
+			hashed := hashString(sessionID)
+			metric.SessionID = &hashed
+		} else {
+			metric.SessionID = &sessionID
+		}
+	}
+
+	if len(metric.Parameters) > 0 {
+		metric.Parameters = s.sanitizeParameters(metric.Parameters)
+	}
+
+	return nil
+}
+
+func (s *Service) sanitizeParameters(values map[string]interface{}) map[string]interface{} {
+	if !s.privacy.SanitizeParameters || len(s.privacy.SensitiveKeys) == 0 {
+		return values
+	}
+
+	sanitized := make(map[string]interface{}, len(values))
+	for key, value := range values {
+		if s.isSensitiveKey(key) {
+			sanitized[key] = "[REDACTED]"
+			continue
+		}
+
+		switch typed := value.(type) {
+		case map[string]interface{}:
+			sanitized[key] = s.sanitizeParameters(typed)
+		case []interface{}:
+			sanitized[key] = s.sanitizeSlice(typed)
+		default:
+			sanitized[key] = value
+		}
+	}
+	return sanitized
+}
+
+func (s *Service) sanitizeSlice(values []interface{}) []interface{} {
+	sanitized := make([]interface{}, len(values))
+	for idx, value := range values {
+		switch typed := value.(type) {
+		case map[string]interface{}:
+			sanitized[idx] = s.sanitizeParameters(typed)
+		case []interface{}:
+			sanitized[idx] = s.sanitizeSlice(typed)
+		default:
+			sanitized[idx] = value
+		}
+	}
+	return sanitized
+}
+
+func (s *Service) isSensitiveKey(key string) bool {
+	lower := strings.ToLower(key)
+	for _, sensitive := range s.privacy.SensitiveKeys {
+		if lower == strings.ToLower(sensitive) {
+			return true
+		}
+	}
+	return false
+}
+
+func (s *Service) upsertServers(ctx context.Context, metrics []models.ToolCallMetric) {
+	seen := make(map[string]models.ServerInfo)
+
+	for _, metric := range metrics {
+		info, exists := seen[metric.ServerID]
+		if !exists {
+			info = models.ServerInfo{
+				ID:        metric.ServerID,
+				Name:      metric.ServerID,
+				FirstSeen: metric.Timestamp,
+				LastSeen:  metric.Timestamp,
+				Metadata:  map[string]interface{}{},
+			}
+		}
+
+		if metric.Timestamp.Before(info.FirstSeen) {
+			info.FirstSeen = metric.Timestamp
+		}
+		if metric.Timestamp.After(info.LastSeen) {
+			info.LastSeen = metric.Timestamp
+		}
+		seen[metric.ServerID] = info
+	}
+
+	for _, server := range seen {
+		if err := s.repo.UpsertServer(ctx, server); err != nil {
+			s.logger.Warn("failed to upsert server info during ingestion",
+				zap.String("server_id", server.ID),
+				zap.Error(err),
+			)
+		}
+	}
+}
+
+func (s *Service) broadcast(event Event) {
+	s.subsMu.RLock()
+	defer s.subsMu.RUnlock()
+
+	for _, sub := range s.subscribers {
+		if sub.serverID != "" && sub.serverID != event.ServerID {
+			continue
+		}
+		if _, ok := sub.topics[event.Type]; !ok {
+			continue
+		}
+
+		select {
+		case sub.ch <- event:
+		default:
+			// Drop message if subscriber is too slow; avoid blocking ingestion.
+			s.logger.Debug("dropping realtime event for slow subscriber",
+				zap.String("server_id", event.ServerID),
+				zap.String("type", event.Type),
+			)
+		}
+	}
+}
+
+func hashString(input string) string {
+	sum := sha256.Sum256([]byte(input))
+	return hex.EncodeToString(sum[:])
+}
+
+func max(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}