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"
	"strings"
	"sync"
	"time"

	"git.tornberg.me/go-otel/internal/issues"
	"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
	subscribers  map[chan websocket.HistoryEntry]func(websocket.HistoryEntry) bool
	subMutex     sync.Mutex
	waitLimiter  *ipRateLimiter
	issues       *issues.Tracker
}

// New creates a new Collector instance.
func New(broadcaster websocket.Broadcaster, historySize int, store storage.Store) *Collector {
	return &Collector{
		broadcaster: broadcaster,
		historySize: historySize,
		history:     ring.New(historySize),
		store:       store,
		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(),
	}
}

// ListIssues returns deduplicated error issues, most recently seen first.
func (c *Collector) ListIssues(limit int) []issues.Issue {
	return c.issues.List(limit)
}

// GetIssue returns a single issue by id.
func (c *Collector) GetIssue(id string) (issues.Issue, bool) {
	return c.issues.Get(id)
}

// SetBroadcaster updates the broadcaster used by the collector.
func (c *Collector) SetBroadcaster(broadcaster websocket.Broadcaster) {
	c.broadcaster = broadcaster
}

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 {
		// Asynchronously append to store to avoid blocking ingestion
		go func() {
			if err := c.store.Append(context.Background(), entry); err != nil {
				log.Printf("failed to append to store: %v", err)
			}
		}()
	}

	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)
	extractSpanErrors(c.issues, td, time.Now().UnixNano())

	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)
	extractLogErrors(c.issues, ld, time.Now().UnixNano())

	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}
}