go-otel/internal/collector/handlers.go

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

import (
	"context"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"git.tornberg.me/go-otel/internal/storage"
	"git.tornberg.me/go-otel/internal/websocket"
	"go.opentelemetry.io/collector/pdata/plog"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/pdata/ptrace"
)

// HandleTraces handles OTLP trace requests.
func (c *Collector) HandleTraces(w http.ResponseWriter, r *http.Request) {
	log.Printf("Received traces request: method=%s, content-type=%s", r.Method, r.Header.Get("Content-Type"))
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		log.Printf("failed to read request body: %v", err)
		http.Error(w, "failed to read request body", http.StatusInternalServerError)
		return
	}

	log.Printf("Received traces body, size=%d", len(body))

	var unmarshaler ptrace.Unmarshaler
	contentType := r.Header.Get("Content-Type")
	switch contentType {
	case "application/x-protobuf":
		unmarshaler = &ptrace.ProtoUnmarshaler{}
	case "application/json":
		unmarshaler = &ptrace.JSONUnmarshaler{}
	default:
		unmarshaler = &ptrace.JSONUnmarshaler{} // default to JSON
	}

	traces, err := unmarshaler.UnmarshalTraces(body)
	if err != nil {
		log.Printf("failed to unmarshal traces: %v", err)
		http.Error(w, "failed to unmarshal traces", http.StatusBadRequest)
		return
	}

	log.Printf("Successfully unmarshaled traces, span_count=%d", traces.SpanCount())

	if err := c.ConsumeTraces(context.Background(), traces); err != nil {
		log.Printf("failed to consume traces: %v", err)
		http.Error(w, "failed to consume traces", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

// HandleMetrics handles OTLP metric requests.
func (c *Collector) HandleMetrics(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		log.Printf("failed to read request body: %v", err)
		http.Error(w, "failed to read request body", http.StatusInternalServerError)
		return
	}

	var unmarshaler pmetric.Unmarshaler
	contentType := r.Header.Get("Content-Type")
	switch contentType {
	case "application/x-protobuf":
		unmarshaler = &pmetric.ProtoUnmarshaler{}
	case "application/json":
		unmarshaler = &pmetric.JSONUnmarshaler{}
	default:
		unmarshaler = &pmetric.JSONUnmarshaler{} // default to JSON
	}

	metrics, err := unmarshaler.UnmarshalMetrics(body)
	if err != nil {
		log.Printf("failed to unmarshal metrics: %v", err)
		http.Error(w, "failed to unmarshal metrics", http.StatusBadRequest)
		return
	}

	if err := c.ConsumeMetrics(context.Background(), metrics); err != nil {
		log.Printf("failed to consume metrics: %v", err)
		http.Error(w, "failed to consume metrics", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

// HandleLogs handles OTLP log requests.
func (c *Collector) HandleLogs(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		log.Printf("failed to read request body: %v", err)
		http.Error(w, "failed to read request body", http.StatusInternalServerError)
		return
	}

	var unmarshaler plog.Unmarshaler
	contentType := r.Header.Get("Content-Type")
	switch contentType {
	case "application/x-protobuf":
		unmarshaler = &plog.ProtoUnmarshaler{}
	case "application/json":
		unmarshaler = &plog.JSONUnmarshaler{}
	default:
		unmarshaler = &plog.JSONUnmarshaler{} // default to JSON
	}

	logs, err := unmarshaler.UnmarshalLogs(body)
	if err != nil {
		log.Printf("failed to unmarshal logs: %v", err)
		http.Error(w, "failed to unmarshal logs", http.StatusBadRequest)
		return
	}

	if err := c.ConsumeLogs(context.Background(), logs); err != nil {
		log.Printf("failed to consume logs: %v", err)
		http.Error(w, "failed to consume logs", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

// parseSearchOptions builds a storage.SearchOptions from common query
// parameters shared by the search/wait/history endpoints: q, type, since, until.
func parseSearchOptions(r *http.Request) storage.SearchOptions {
	q := r.URL.Query()
	opts := storage.SearchOptions{
		Query: q.Get("q"),
		Type:  q.Get("type"),
	}
	if since := q.Get("since"); since != "" {
		if v, err := strconv.ParseInt(since, 10, 64); err == nil {
			opts.Since = v
		}
	}
	if until := q.Get("until"); until != "" {
		if v, err := strconv.ParseInt(until, 10, 64); err == nil {
			opts.Until = v
		}
	}
	return opts
}

func parseLimit(r *http.Request, def int) int {
	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
		if l, err := strconv.Atoi(limitStr); err == nil {
			return l
		}
	}
	return def
}

// HandleHistory returns the telemetry history as JSON. It searches the full
// persistent store when one is configured, so results are not bounded by the
// in-memory ring buffer size.
func (c *Collector) HandleHistory(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	history, err := c.SearchHistory(r.Context(), storage.SearchOptions{}, parseLimit(r, 100))
	if err != nil {
		log.Printf("failed to fetch history: %v", err)
		http.Error(w, "failed to fetch history", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(history); err != nil {
		log.Printf("failed to encode history: %v", err)
		http.Error(w, "failed to encode history", http.StatusInternalServerError)
		return
	}
}

// HandleSearch searches the telemetry history. Supports the free-text 'q'
// parameter plus structured filters: type, since, until (unix nanoseconds).
func (c *Collector) HandleSearch(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	opts := parseSearchOptions(r)
	history, err := c.SearchHistory(r.Context(), opts, parseLimit(r, 100))
	if err != nil {
		log.Printf("failed to search history: %v", err)
		http.Error(w, "failed to search history", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(history); err != nil {
		log.Printf("failed to encode search results: %v", err)
		http.Error(w, "failed to encode history", http.StatusInternalServerError)
		return
	}
}

// HandleWait waits for a specific telemetry signal, optionally restricted by
// the same structured filters as HandleSearch (type, since, until).
func (c *Collector) HandleWait(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if !c.waitLimiter.Allow(clientIP(r)) {
		http.Error(w, "too many wait requests, slow down", http.StatusTooManyRequests)
		return
	}

	opts := parseSearchOptions(r)
	timeoutSec := 30
	if t := r.URL.Query().Get("timeout"); t != "" {
		if ts, err := strconv.Atoi(t); err == nil {
			timeoutSec = ts
		}
	}

	ctx, cancel := context.WithTimeout(r.Context(), time.Duration(timeoutSec)*time.Second)
	defer cancel()

	// First, check if result already exists in recent history
	matches, err := c.SearchHistory(ctx, opts, 1)
	if err == nil && len(matches) > 0 {
		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(matches[0]); err != nil {
			log.Printf("failed to encode match: %v", err)
		}
		return
	}

	// Otherwise, subscribe and wait
	lowerQuery := strings.ToLower(opts.Query)
	ch, cleanup := c.Subscribe(func(entry websocket.HistoryEntry) bool {
		if opts.Type != "" && entry.Type != opts.Type {
			return false
		}
		if opts.Query == "" {
			return true
		}
		return strings.Contains(strings.ToLower(string(entry.Data)), lowerQuery)
	})
	defer cleanup()

	select {
	case entry := <-ch:
		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(entry); err != nil {
			log.Printf("failed to encode signal: %v", err)
		}
	case <-ctx.Done():
		http.Error(w, "timeout waiting for signal", http.StatusGatewayTimeout)
	}
}

// HandleHistoryByType returns telemetry history filtered by type, searching
// the full persistent store when one is configured.
func (c *Collector) HandleHistoryByType(dataType string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}

		opts := storage.SearchOptions{Type: dataType}
		history, err := c.SearchHistory(r.Context(), opts, parseLimit(r, 100))
		if err != nil {
			log.Printf("failed to fetch history for type %s: %v", dataType, err)
			http.Error(w, "failed to fetch history", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(history); err != nil {
			log.Printf("failed to encode filtered history: %v", err)
			http.Error(w, "failed to encode history", http.StatusInternalServerError)
			return
		}
	}
}

// HandleTraceByID returns all stored entries referencing the given trace id,
// so agents don't have to rely on the trace id happening to substring-match
// inside a search result.
func (c *Collector) HandleTraceByID(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	traceID := r.PathValue("traceId")
	if traceID == "" {
		http.Error(w, "traceId is required", http.StatusBadRequest)
		return
	}

	entries, err := c.SearchHistory(r.Context(), storage.SearchOptions{Query: traceID}, parseLimit(r, 1000))
	if err != nil {
		log.Printf("failed to fetch trace %s: %v", traceID, err)
		http.Error(w, "failed to fetch trace", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(entries); err != nil {
		log.Printf("failed to encode trace: %v", err)
		http.Error(w, "failed to encode trace", http.StatusInternalServerError)
		return
	}
}

// HandleHealth returns the health status of the collector.
func (c *Collector) HandleHealth(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(map[string]string{"status": "ok"}); err != nil {
		log.Printf("failed to encode health check: %v", err)
	}
}