go-otel/internal/collector/collector.go

package collector

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

	"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"
)

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
}

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

func (c *Collector) SetBroadcaster(broadcaster websocket.Broadcaster) {
	c.broadcaster = broadcaster
}

func (c *Collector) addToHistory(dataType websocket.DataType, data interface{}) {
	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
			}
		}
	}
}

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
}

func (c *Collector) GetHistory() []websocket.HistoryEntry {
	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 history
}

func (c *Collector) SearchHistory(ctx context.Context, query string, limit int) ([]websocket.HistoryEntry, error) {
	if limit <= 0 {
		limit = 50
	}

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

	// Fallback to in-memory if no store
	history := c.GetHistory()
	var filtered []websocket.HistoryEntry
	// Lowercase query for case-insensitive search
	lowerQuery := strings.ToLower(query)

	for _, entry := range history {
		if query == "" {
			filtered = append(filtered, entry)
		} else {
			// Basic string filter for in-memory
			data, _ := json.Marshal(entry.Data)
			if strings.Contains(strings.ToLower(string(data)), lowerQuery) {
				filtered = append(filtered, entry)
			}
		}
		// Keep it simple and just filter all then limit, since historySize is usually small
	}

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

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
	}

	// Parse to interface{} for history storage
	var data interface{}
	if err := json.Unmarshal(buf, &data); err != nil {
		log.Printf("failed to unmarshal traces for history: %v", err)
	} else {
		c.addToHistory(websocket.DataTypeTraces, data)
	}

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

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

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
	}

	// Parse to interface{} for history storage
	var data interface{}
	if err := json.Unmarshal(buf, &data); err != nil {
		log.Printf("failed to unmarshal metrics for history: %v", err)
	} else {
		c.addToHistory(websocket.DataTypeMetrics, data)
	}

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

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

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
	}

	// Parse to interface{} for history storage
	var data interface{}
	if err := json.Unmarshal(buf, &data); err != nil {
		log.Printf("failed to unmarshal logs for history: %v", err)
	} else {
		c.addToHistory(websocket.DataTypeLogs, data)
	}

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

	c.broadcaster.Broadcast(&entry)

	return nil
}

// Implement consumer interfaces
func (c *Collector) Capabilities() consumer.Capabilities {
	return consumer.Capabilities{MutatesData: false}
}