broad-announce/internal/grpcserver/handler.go

package grpcserver

import (
	"context"
	"encoding/json"
	"io"
	"sync"
	"sync/atomic"
	"time"

	pbv1 "git3.techno-world.net/lrosales/broad-announce/gen/go/broadannounce/v1"
	"git3.techno-world.net/lrosales/broad-announce/internal/pipeline"
)

// Server implements broadannounce.v1.IngestServer.
type Server struct {
	pbv1.UnimplementedIngestServer

	// Dep is the shared pipeline dependencies. All goroutines write to it
	// concurrently — it is safe for concurrent use.
	Dep         pipeline.Deps
	MaxInflight int // max concurrent in-flight messages per stream (default 256)
	Logger      interface{ Info(msg string, args ...any) }
}

// StreamAlerts is a bidirectional stream: client sends zero or more Alerts,
// server sends exactly one Ack per Alert received.
//
// Bounded concurrency: each message runs in a goroutine controlled by a
// semaphore channel of size MaxInflight. This is how we achieve 10k+/s on
// a single stream — the bottleneck is the Redis dedupe + NATS publish (~2ms),
// so 256-way concurrency gives us ~128k msg/s theoretical max.
//
// Backpressure: if the sem is full, we immediately send Error.RATE_LIMITED
// without entering the pipeline. The client is expected to honour retry_after_ms.
func (s *Server) StreamAlerts(stream pbv1.Ingest_StreamAlertsServer) error {
	ctx := stream.Context()

	// Authenticate the stream.
	src, err := authenticate(ctx, s.Dep.Sources)
	if err != nil {
		return err
	}
	sourceKey := src.CompanyID + ":" + src.SourceID

	// Per-stream semaphore: bounds concurrent processing to MaxInflight.
	if s.MaxInflight <= 0 {
		s.MaxInflight = 256
	}
	sem := make(chan struct{}, s.MaxInflight)

	// Per-stream ackCh: each goroutine sends its Ack here; the recv loop
	// forwards to the client. Buffer = MaxInflight so goroutines never block
	// on sending (only on the semaphore).
	ackCh := make(chan *pbv1.Ack, s.MaxInflight)

	// Track active goroutines for graceful shutdown.
	var wg sync.WaitGroup

	// streamErr holds the first non-nil error from goroutines or the recv loop.
	var streamErr atomic.Value // holds error
	setErr := func(err error) {
		if err != nil {
			streamErr.CompareAndSwap(nil, err)
		}
	}

	// Increment active-streams gauge.
	s.Dep.Metrics.StreamsActive.Add(1)
	defer s.Dep.Metrics.StreamsActive.Add(-1)

	// goroutine: receive loop — reads from client, dispatches to pipeline.
	go func() {
		for {
			alert, err := stream.Recv()
			if err == io.EOF {
				// Client called CloseSend. Drain in-flight, then exit.
				close(ackCh)
				wg.Wait()
				return
			}
			if err != nil {
				setErr(err)
				close(ackCh)
				wg.Wait()
				return
			}

			// Validate we have a non-empty Alert before spending goroutines.
			if alert.CompanyId == "" || alert.SourceId == "" {
				ackCh <- &pbv1.Ack{
					DedupeKey: alert.DedupeKey,
					Result: &pbv1.Ack_Error{
						Error: &pbv1.Error{
							Code:    pbv1.Error_INVALID,
							Message: "company_id and source_id are required",
						},
					},
				}
				continue
			}

			// Per-stream rate limit (before entering the pipeline).
			// This is a fast rejection path that doesn't consume goroutines.
			rateLimitKey := "grpc_source:" + sourceKey
			allowed, retryAfter, _ := s.Dep.Limiter.Allow(ctx, rateLimitKey, src.RateLimitPerSec)
			if !allowed {
				s.Dep.Metrics.GRPCRateLimited.WithLabelValues(src.SourceID).Inc()
				s.Dep.Metrics.AlertsReceived.WithLabelValues("grpc", "rate_limited").Inc()
				ackCh <- &pbv1.Ack{
					AlertId:   alert.DedupeKey,
					DedupeKey: alert.DedupeKey,
					Result: &pbv1.Ack_Error{
						Error: &pbv1.Error{
							Code:          pbv1.Error_RATE_LIMITED,
							Message:       "per-stream rate limit exceeded",
							RetryAfterMs: int32(retryAfter.Milliseconds()),
						},
					},
				}
				continue
			}

			// Acquire semaphore slot. If full, backpressure immediately.
			select {
			case sem <- struct{}{}:
				// Proceed.
			default:
				s.Dep.Metrics.GRPCRateLimited.WithLabelValues(src.SourceID).Inc()
				s.Dep.Metrics.AlertsReceived.WithLabelValues("grpc", "rate_limited").Inc()
				ackCh <- &pbv1.Ack{
					AlertId:   alert.DedupeKey,
					DedupeKey: alert.DedupeKey,
					Result: &pbv1.Ack_Error{
						Error: &pbv1.Error{
							Code:          pbv1.Error_RATE_LIMITED,
							Message:       "in-flight capacity reached",
							RetryAfterMs: 5, // 5ms backoff, client should retry
						},
					},
				}
				continue
			}

			wg.Add(1)
			go func(alert *pbv1.Alert) {
				defer func() {
					<-sem
					wg.Done()
				}()

				start := time.Now()
				ack := s.processAlert(ctx, alert, src)

				// Record ack latency metric.
				s.Dep.Metrics.GRPCAckLatency.WithLabelValues(src.SourceID).Observe(
					time.Since(start).Seconds())

				// Send Ack to client. Non-blocking — if the channel is full the
				// recv loop has a backlog and we should not block the goroutine.
				select {
				case ackCh <- ack:
				default:
					// Channel full; log and drop.
					s.Dep.Logger.Info("grpc ack channel full, dropping", "alert_id", ack.AlertId)
				}
			}(alert)
		}
	}()

	// recv loop: forward Acks from goroutines to the client.
	// Also watches for goroutine errors.
	for ack := range ackCh {
		if err := stream.Send(ack); err != nil {
			return err
		}
	}

	// Check if the stream exited due to a goroutine error.
	if se := streamErr.Load().(error); se != nil {
		return se
	}
	return nil
}

// processAlert runs the shared pipeline on a single gRPC Alert and returns
// the corresponding Ack. sig="" because gRPC auth is handled by authenticate()
// before entering this function.
func (s *Server) processAlert(ctx context.Context, alert *pbv1.Alert, src *pipeline.SourceConfig) *pbv1.Ack {
	body, err := alertToJSON(alert)
	if err != nil {
		return errorAck(alert.DedupeKey, pbv1.Error_INVALID, "marshal failed: "+err.Error(), 0)
	}

	// Override the Sources map for this call to use the authenticated source.
	// This ensures company_id/source_id from the API key match the alert body.
	sourceKey := src.CompanyID + ":" + src.SourceID
	origSources := s.Dep.Sources
	s.Dep.Sources = map[string]pipeline.SourceConfig{sourceKey: *src}

	res := s.Dep.Process(ctx, body, "" /* no HMAC for gRPC */)

	// Restore the original Sources map.
	s.Dep.Sources = origSources

	if !res.Accepted {
		code := errorCode(res.RejectReason)
		var retryAfter int32
		if res.RejectReason == "rate_limited_source" || res.RejectReason == "rate_limited_company" {
			if n, _ := parseRetryAfter(res.Detail); n > 0 {
				retryAfter = int32(n)
			}
		}
		return errorAck(alert.DedupeKey, code, res.Detail, retryAfter)
	}

	return &pbv1.Ack{
		AlertId:      res.AlertID,
		DedupeKey:    alert.DedupeKey,
		DedupeCount:  res.DedupeCount,
		AcceptedAtMs: time.Now().UnixMilli(),
		Result:       &pbv1.Ack_Ok{Ok: &pbv1.Ok{}},
	}
}

// alertToJSON converts a protobuf Alert to JSON bytes for the pipeline.
func alertToJSON(a *pbv1.Alert) ([]byte, error) {
	m := map[string]any{
		"company_id": a.CompanyId,
		"source_id":  a.SourceId,
	}
	if a.DedupeKey != "" {
		m["dedupe_key"] = a.DedupeKey
	}
	if a.Severity != "" {
		m["severity"] = a.Severity
	}
	if a.Category != "" {
		m["category"] = a.Category
	}
	if a.Title != "" {
		m["title"] = a.Title
	}
	if a.Body != "" {
		m["body"] = a.Body
	}
	if len(a.Data) > 0 {
		m["data"] = a.Data
	}
	if a.ClientTsMs > 0 {
		m["client_ts_ms"] = a.ClientTsMs
	}
	return json.Marshal(m)
}

func errorAck(dedupeKey string, code pbv1.Error_Code, message string, retryAfterMs int32) *pbv1.Ack {
	return &pbv1.Ack{
		DedupeKey: dedupeKey,
		Result: &pbv1.Ack_Error{
			Error: &pbv1.Error{
				Code:          code,
				Message:       message,
				RetryAfterMs: retryAfterMs,
			},
		},
	}
}

// errorCode maps pipeline reject reasons to gRPC Error codes.
func errorCode(reason string) pbv1.Error_Code {
	switch reason {
	case "unknown_source", "bad_signature":
		return pbv1.Error_UNAUTHENTICATED
	case "rate_limited_source", "rate_limited_company":
		return pbv1.Error_RATE_LIMITED
	case "invalid", "invalid_json", "quarantined":
		return pbv1.Error_INVALID
	case "circuit_open", "broker_unavailable", "marshal_failed":
		return pbv1.Error_INTERNAL
	default:
		return pbv1.Error_UNKNOWN
	}
}

func parseRetryAfter(s string) (int, bool) {
	var n int
	var parsed bool
	for _, c := range s {
		if c >= '0' && c <= '9' {
			n = n*10 + int(c-'0')
			parsed = true
		} else {
			break
		}
	}
	return n, parsed
}