go-otel/internal/collector/collector.go

// Package collector implements an OpenTelemetry consumer that stores telemetry in history
// and broadcasts it to connected clients.
package collector

import (
	"container/ring"
	"context"
	"encoding/json"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.tornberg.me/go-otel/internal/issues"
	kube "git.tornberg.me/go-otel/internal/kubernetes"
	"git.tornberg.me/go-otel/internal/prometheus"
	"git.tornberg.me/go-otel/internal/sourcemaps"
	"git.tornberg.me/go-otel/internal/storage"
	"git.tornberg.me/go-otel/internal/websocket"

	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/pdata/plog"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/pdata/ptrace"
)

// Collector handles the collection, storage and broadcasting of telemetry data.
type Collector struct {
	broadcaster  websocket.Broadcaster
	historySize  int
	history      *ring.Ring
	historyMutex sync.RWMutex
	store        storage.Store
	storeChan    chan websocket.HistoryEntry
	subscribers  map[chan websocket.HistoryEntry]func(websocket.HistoryEntry) bool
	subMutex     sync.Mutex
	waitLimiter  *ipRateLimiter
	issues       *issues.Tracker

	ingestLimiter *ipRateLimiter
	ingestKey     string

	notifyIssuesFn func()

	// statsQuerier is the ClickHouse store when configured; nil otherwise.
	// The stats endpoints return 501 when it's nil so the UI can hide those
	// views instead of erroring.
	statsQuerier storage.StatsQuerier
	sourceMaps   *sourcemaps.Resolver
	kubeWatcher  *kube.Watcher
	prometheus   *prometheus.Client
}

// New creates a new Collector instance.
func New(broadcaster websocket.Broadcaster, historySize int, store storage.Store) *Collector {
	c := &Collector{
		broadcaster: broadcaster,
		historySize: historySize,
		history:     ring.New(historySize),
		store:       store,
		storeChan:   make(chan websocket.HistoryEntry, 10000),
		subscribers: make(map[chan websocket.HistoryEntry]func(websocket.HistoryEntry) bool),
		// Allow bursts of a few immediate waits, then steady-state ~1 req/sec per IP.
		waitLimiter: newIPRateLimiter(1, 5),
		issues:      issues.NewTracker(),
		// Generous enough for normal exporter batching, but enough to stop a
		// runaway client (e.g. a frontend error handler that itself throws,
		// looping) from hammering the collector.
		ingestLimiter: newIPRateLimiter(20, 50),
	}

	if store != nil {
		// Bounded worker pool prevents concurrent ClickHouse INSERTs from
		// multiplying query and merge memory under bursty ingest traffic.
		numWorkers := 4
		if value, err := strconv.Atoi(os.Getenv("OTEL_STORE_WORKERS")); err == nil && value > 0 && value <= 32 {
			numWorkers = value
		}
		for i := 0; i < numWorkers; i++ {
			go func() {
				for entry := range c.storeChan {
					if err := c.store.Append(context.Background(), entry); err != nil {
						log.Printf("failed to append to store: %v", err)
					}
				}
			}()
		}
	}

	return c
}

// SetIngestKey configures an optional shared key that HTTP OTLP ingestion
// (traces/metrics/logs) must present, via a '?key=' query parameter or an
// 'X-Otel-Ingest-Key' header. An empty key (the default) disables the check.
// This is meant to scope/rate-limit a public ingest endpoint (e.g. one fed
// by a browser script), not to provide real authentication.
func (c *Collector) SetIngestKey(key string) {
	c.ingestKey = key
}

// SetStatsQuerier enables the aggregate stats endpoints (error trends,
// release comparisons) when a store supporting them is configured.
func (c *Collector) SetStatsQuerier(q storage.StatsQuerier) {
	c.statsQuerier = q
}

func (c *Collector) SetKubernetesWatcher(w *kube.Watcher)     { c.kubeWatcher = w }
func (c *Collector) SetPrometheusClient(p *prometheus.Client) { c.prometheus = p }
func (c *Collector) ListRollouts() []kube.Rollout {
	if c.kubeWatcher == nil {
		return nil
	}
	return c.kubeWatcher.List()
}
func (c *Collector) GetRollout(id string) (kube.Rollout, bool) {
	if c.kubeWatcher == nil {
		return kube.Rollout{}, false
	}
	return c.kubeWatcher.Get(id)
}
func (c *Collector) KubernetesEnabled() bool { return c.kubeWatcher != nil && c.kubeWatcher.Enabled() }
func (c *Collector) KubernetesError() string {
	if c.kubeWatcher == nil {
		return "kubernetes integration is not configured"
	}
	return c.kubeWatcher.Error()
}

// ListIssues returns deduplicated error issues, most recently seen first.
func (c *Collector) ListIssues(limit int) []issues.Issue {
	list := c.issues.ListFiltered(issues.ListOptions{Limit: limit}).Issues
	if c.sourceMaps != nil {
		for i := range list {
			resolveIssue(&list[i], c.sourceMaps)
		}
	}
	return list
}

// ListIssuesFiltered returns a filtered page of deduplicated issues.
func (c *Collector) ListIssuesFiltered(opts issues.ListOptions) issues.ListPage {
	page := c.issues.ListFiltered(opts)
	if c.sourceMaps != nil {
		for i := range page.Issues {
			resolveIssue(&page.Issues[i], c.sourceMaps)
		}
	}
	return page
}

// GetIssue returns a single issue by id.
func (c *Collector) GetIssue(id string) (issues.Issue, bool) {
	issue, ok := c.issues.Get(id)
	if ok && c.sourceMaps != nil {
		resolveIssue(&issue, c.sourceMaps)
	}
	return issue, ok
}

func (c *Collector) SetSourceMaps(resolver *sourcemaps.Resolver) { c.sourceMaps = resolver }

func resolveIssue(issue *issues.Issue, resolver *sourcemaps.Resolver) {
	for i := range issue.Samples {
		issue.Samples[i].ResolvedStacktrace = resolver.ResolveStack(issue.Samples[i].Release, issue.Samples[i].Stacktrace)
	}
}

// SetIssueStatus updates the triage state of an issue.
func (c *Collector) SetIssueStatus(id, status string) (issues.Issue, bool) {
	return c.issues.SetStatus(id, status)
}

func (c *Collector) SimilarIssues(id string, limit int) ([]issues.SimilarIssue, bool) {
	return c.issues.Similar(id, limit)
}

func (c *Collector) ListPatterns() []issues.Pattern { return c.issues.Patterns() }

func (c *Collector) SetSimilarityFeedback(a, b, decision string) bool {
	return c.issues.SetSimilarityFeedback(a, b, decision)
}

// SetBroadcaster updates the broadcaster used by the collector.
func (c *Collector) SetBroadcaster(broadcaster websocket.Broadcaster) {
	c.broadcaster = broadcaster
	if notifier, ok := broadcaster.(interface{ NotifyIssues() }); ok {
		c.notifyIssuesFn = notifier.NotifyIssues
	}
}

// notifyIssues pings WebSocket clients that new error telemetry arrived, so
// the UI can refresh its Issues view without waiting for its poll interval.
func (c *Collector) notifyIssues() {
	if c.notifyIssuesFn != nil {
		c.notifyIssuesFn()
	}
}

func (c *Collector) addToHistory(dataType websocket.DataType, data json.RawMessage) {
	entry := websocket.HistoryEntry{
		Type:      string(dataType),
		Timestamp: time.Now().UnixNano(),
		Data:      data,
	}

	c.historyMutex.Lock()
	c.history.Value = entry
	c.history = c.history.Next()
	c.historyMutex.Unlock()

	if c.store != nil {
		select {
		case c.storeChan <- entry:
		default:
			log.Printf("collector: store write queue full, dropping entry to preserve ingestion")
		}
	}

	c.notifySubscribers(entry)
}

func (c *Collector) notifySubscribers(entry websocket.HistoryEntry) {
	c.subMutex.Lock()
	defer c.subMutex.Unlock()

	for ch, filter := range c.subscribers {
		if filter(entry) {
			select {
			case ch <- entry:
				// Successfully sent
			default:
				// Channel blocked, skip
			}
		}
	}
}

// Subscribe adds a new history subscriber with the given filter.
// It returns a channel for receiving filtered history entries and a cleanup function.
func (c *Collector) Subscribe(filter func(websocket.HistoryEntry) bool) (chan websocket.HistoryEntry, func()) {
	ch := make(chan websocket.HistoryEntry, 1)
	c.subMutex.Lock()
	c.subscribers[ch] = filter
	c.subMutex.Unlock()

	cleanup := func() {
		c.subMutex.Lock()
		delete(c.subscribers, ch)
		c.subMutex.Unlock()
		close(ch)
	}

	return ch, cleanup
}

// GetHistory returns the last 'limit' history entries.
func (c *Collector) GetHistory(limit int) []websocket.HistoryEntry {
	if limit <= 0 {
		limit = c.historySize
	}

	c.historyMutex.RLock()
	defer c.historyMutex.RUnlock()

	var history []websocket.HistoryEntry
	c.history.Do(func(v interface{}) {
		if entry, ok := v.(websocket.HistoryEntry); ok && entry.Type != "" {
			history = append(history, entry)
		}
	})

	// Return only the last 'limit' items if we have more
	if len(history) > limit {
		history = history[len(history)-limit:]
	}

	return history
}

// SearchHistory searches through stored telemetry data using the provided options.
// When a persistent store is configured, the full stored history is searched;
// otherwise it falls back to the in-memory ring buffer.
func (c *Collector) SearchHistory(ctx context.Context, opts storage.SearchOptions, limit int) ([]websocket.HistoryEntry, error) {
	if limit <= 0 {
		limit = 50
	}

	if c.store != nil {
		return c.store.Search(ctx, opts, limit)
	}

	// Fallback to in-memory if no store. Use the full ring buffer, not just
	// 'limit' entries, so type/time filters see everything available.
	history := c.GetHistory(0)
	var filtered []websocket.HistoryEntry
	// Lowercase query for case-insensitive search
	lowerQuery := strings.ToLower(opts.Query)

	for _, entry := range history {
		if opts.Query != "" && !strings.Contains(strings.ToLower(string(entry.Data)), lowerQuery) {
			continue
		}
		if opts.Type != "" && entry.Type != opts.Type {
			continue
		}
		if opts.Since != 0 && entry.Timestamp < opts.Since {
			continue
		}
		if opts.Until != 0 && entry.Timestamp > opts.Until {
			continue
		}
		filtered = append(filtered, entry)
	}

	// Respect limit and return latest messages first
	if len(filtered) > limit {
		filtered = filtered[len(filtered)-limit:]
	}

	// Reverse to match Store's Search behavior (latest first)
	for i, j := 0, len(filtered)-1; i < j; i, j = i+1, j-1 {
		filtered[i], filtered[j] = filtered[j], filtered[i]
	}

	return filtered, nil
}

// ConsumeTraces implements the consumer.Traces interface.
func (c *Collector) ConsumeTraces(ctx context.Context, td ptrace.Traces) error {
	marshaler := &ptrace.JSONMarshaler{}
	buf, err := marshaler.MarshalTraces(td)
	if err != nil {
		log.Printf("failed to marshal traces: %v", err)
		return err
	}

	// Store and broadcast raw bytes as json.RawMessage to avoid unsafe interface{} usage
	c.addToHistory(websocket.DataTypeTraces, buf)
	if extractSpanErrors(c.issues, td, time.Now().UnixNano()) {
		c.notifyIssues()
	}

	entry := websocket.Entry{
		Type: string(websocket.DataTypeTraces),
		Data: buf,
	}

	c.broadcaster.Broadcast(&entry)
	return nil
}

// ConsumeMetrics implements the consumer.Metrics interface.
func (c *Collector) ConsumeMetrics(ctx context.Context, md pmetric.Metrics) error {
	marshaler := &pmetric.JSONMarshaler{}
	buf, err := marshaler.MarshalMetrics(md)
	if err != nil {
		log.Printf("failed to marshal metrics: %v", err)
		return err
	}

	// Store and broadcast raw bytes as json.RawMessage to avoid unsafe interface{} usage
	c.addToHistory(websocket.DataTypeMetrics, buf)

	entry := websocket.Entry{
		Type: string(websocket.DataTypeMetrics),
		Data: buf,
	}

	c.broadcaster.Broadcast(&entry)
	return nil
}

// ConsumeLogs implements the consumer.Logs interface.
func (c *Collector) ConsumeLogs(ctx context.Context, ld plog.Logs) error {
	marshaler := &plog.JSONMarshaler{}
	buf, err := marshaler.MarshalLogs(ld)
	if err != nil {
		log.Printf("failed to marshal logs: %v", err)
		return err
	}

	// Store and broadcast raw bytes as json.RawMessage to avoid unsafe interface{} usage
	c.addToHistory(websocket.DataTypeLogs, buf)
	if extractLogErrors(c.issues, ld, time.Now().UnixNano()) {
		c.notifyIssues()
	}

	entry := websocket.Entry{
		Type: string(websocket.DataTypeLogs),
		Data: buf,
	}

	c.broadcaster.Broadcast(&entry)

	return nil
}

// Capabilities returns the capabilities of the collector.
func (c *Collector) Capabilities() consumer.Capabilities {
	return consumer.Capabilities{MutatesData: false}
}