go-otel/internal/collector/issue_traces.go

package collector

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"git.tornberg.me/go-otel/internal/issues"
	"git.tornberg.me/go-otel/internal/storage"
)

// occurrenceBuffer bounds queued occurrence writes; beyond it, writes are
// dropped (and counted in the log) rather than slowing ingest.
const occurrenceBuffer = 4096

// initOptionalStores wires the store's optional capabilities: trace lookups
// and persistent issue occurrences.
func (c *Collector) initOptionalStores(store storage.Store) {
	if tq, ok := store.(storage.TraceQuerier); ok {
		c.traceQuerier = tq
	}
	occ, ok := store.(storage.OccurrenceStore)
	if !ok {
		return
	}
	c.occStore = occ
	c.occChan = make(chan storage.Occurrence, occurrenceBuffer)
	go func() {
		for o := range c.occChan {
			if err := occ.AppendOccurrence(context.Background(), o); err != nil {
				log.Printf("failed to persist issue occurrence: %v", err)
			}
		}
	}()
	var dropped atomic.Int64
	c.issues.SetOccurrenceSink(func(issueID, kind, service string, sample issues.Sample) {
		raw, err := json.Marshal(sample)
		if err != nil {
			return
		}
		select {
		case c.occChan <- storage.Occurrence{
			IssueID: issueID, Kind: kind, Service: service, Timestamp: sample.Timestamp,
			TraceID: sample.TraceID, SpanID: sample.SpanID, Sample: raw,
		}:
		default:
			if n := dropped.Add(1); n%1000 == 1 {
				log.Printf("issue occurrence buffer full; dropped %d so far", n)
			}
		}
	})
}

// issueTraceJSON is one trace in the issue overview.
type issueTraceJSON struct {
	TraceID   string                `json:"traceId"`
	Hits      int                   `json:"hits"`
	FirstSeen int64                 `json:"firstSeen"`
	LastSeen  int64                 `json:"lastSeen"`
	Summary   *storage.TraceSummary `json:"summary,omitempty"`
	Samples   []issues.Sample       `json:"samples"`
}

type issueTracesResponse struct {
	IssueID          string           `json:"issueId"`
	TotalOccurrences int              `json:"totalOccurrences"`
	Traces           []issueTraceJSON `json:"traces"`
	Untraced         []issues.Sample  `json:"untraced"`
	HasMore          bool             `json:"hasMore"`
	// NextBefore is the LastSeen boundary to pass as ?before= for the next page.
	NextBefore int64 `json:"nextBefore,omitempty"`
	// Complete is false when only the in-memory samples were available.
	Complete bool `json:"complete"`
}

func decodeSamples(raw []json.RawMessage) []issues.Sample {
	out := make([]issues.Sample, 0, len(raw))
	for _, r := range raw {
		var s issues.Sample
		if json.Unmarshal(r, &s) == nil {
			out = append(out, s)
		}
	}
	return out
}

// HandleIssueTraces returns every trace an issue occurred in with per-trace
// summaries (duration, spans, services, errors) and the occurrences in each,
// in a single request. Spans themselves stay on GET /api/traces/{id}.
func (c *Collector) HandleIssueTraces(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	issue, known := c.GetIssue(id)

	var resp issueTracesResponse
	resp.IssueID = id
	if c.occStore != nil {
		opts := storage.IssueTracesOptions{Limit: parseLimit(r, 50)}
		if b, err := strconv.ParseInt(r.URL.Query().Get("before"), 10, 64); err == nil {
			opts.Before = b
		}
		page, err := c.occStore.IssueTraces(r.Context(), id, opts)
		if err != nil {
			log.Printf("failed to load issue traces %s: %v", id, err)
			http.Error(w, "failed to load issue traces", http.StatusInternalServerError)
			return
		}
		if page.TotalOccurrences == 0 && !known {
			http.Error(w, "issue not found", http.StatusNotFound)
			return
		}
		resp.TotalOccurrences, resp.HasMore, resp.Complete = page.TotalOccurrences, page.HasMore, true
		resp.Untraced = decodeSamples(page.Untraced)
		for _, t := range page.Traces {
			resp.Traces = append(resp.Traces, issueTraceJSON{
				TraceID: t.TraceID, Hits: t.Hits, FirstSeen: t.FirstSeen, LastSeen: t.LastSeen,
				Summary: t.Summary, Samples: decodeSamples(t.Samples),
			})
		}
		if resp.HasMore && len(resp.Traces) > 0 {
			resp.NextBefore = resp.Traces[len(resp.Traces)-1].LastSeen
		}
	} else {
		if !known {
			http.Error(w, "issue not found", http.StatusNotFound)
			return
		}
		resp.TotalOccurrences = issue.Count
		resp.Traces, resp.Untraced = groupInMemory(issue.Samples)
	}

	// Source-map resolution applies to stored samples exactly as it does for
	// the in-memory ones.
	if c.sourceMaps != nil {
		for i := range resp.Traces {
			resolveSamples(resp.Traces[i].Samples, c.sourceMaps)
		}
		resolveSamples(resp.Untraced, c.sourceMaps)
	}
	if resp.Traces == nil {
		resp.Traces = []issueTraceJSON{}
	}
	if resp.Untraced == nil {
		resp.Untraced = []issues.Sample{}
	}

	// Summaries for traces the store didn't summarise (fallback path, or
	// traces older than the summary table): fill from the store when possible.
	if c.traceQuerier != nil {
		var missing []string
		for _, t := range resp.Traces {
			if t.Summary == nil {
				missing = append(missing, t.TraceID)
			}
		}
		if len(missing) > 0 {
			if sums, err := c.traceQuerier.TraceSummaries(r.Context(), missing); err == nil {
				for i := range resp.Traces {
					if s, ok := sums[resp.Traces[i].TraceID]; ok {
						resp.Traces[i].Summary = &s
					}
				}
			}
		}
	}

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

// HandleIssueSources returns all distinct source IPs recorded for an issue.
// This is separate from /traces because trace pages intentionally bound the
// number of raw samples returned per trace.
func (c *Collector) HandleIssueSources(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	if c.occStore == nil {
		http.Error(w, "issue source history requires ClickHouse storage", http.StatusNotImplemented)
		return
	}
	sources, err := c.occStore.IssueSources(r.Context(), id)
	if err != nil {
		log.Printf("failed to load issue sources %s: %v", id, err)
		http.Error(w, "failed to load issue sources", http.StatusInternalServerError)
		return
	}
	if sources == nil {
		sources = []storage.IssueSource{}
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(sources)
}

// groupInMemory builds the same shape from the tracker's bounded samples,
// newest trace first.
func groupInMemory(samples []issues.Sample) ([]issueTraceJSON, []issues.Sample) {
	byTrace := map[string]*issueTraceJSON{}
	var traces []*issueTraceJSON
	var untraced []issues.Sample
	for _, s := range samples {
		if s.TraceID == "" {
			untraced = append(untraced, s)
			continue
		}
		t, ok := byTrace[s.TraceID]
		if !ok {
			t = &issueTraceJSON{TraceID: s.TraceID, FirstSeen: s.Timestamp, LastSeen: s.Timestamp}
			byTrace[s.TraceID] = t
			traces = append(traces, t)
		}
		t.Hits++
		t.Samples = append(t.Samples, s)
		if s.Timestamp < t.FirstSeen {
			t.FirstSeen = s.Timestamp
		}
		if s.Timestamp > t.LastSeen {
			t.LastSeen = s.Timestamp
		}
	}
	sort.Slice(traces, func(i, j int) bool { return traces[i].LastSeen > traces[j].LastSeen })
	out := make([]issueTraceJSON, len(traces))
	for i, t := range traces {
		// Newest first, matching the store-backed path.
		sort.SliceStable(t.Samples, func(a, b int) bool { return t.Samples[a].Timestamp > t.Samples[b].Timestamp })
		out[i] = *t
	}
	sort.Slice(untraced, func(i, j int) bool { return untraced[i].Timestamp > untraced[j].Timestamp })
	return out, untraced
}

// HandleTraceSummaries returns per-trace summaries for ?ids=a,b,c (max 100).
// Unknown ids are omitted; stores without summaries answer an empty object.
func (c *Collector) HandleTraceSummaries(w http.ResponseWriter, r *http.Request) {
	var ids []string
	for _, id := range strings.Split(r.URL.Query().Get("ids"), ",") {
		if id = strings.TrimSpace(id); id != "" {
			ids = append(ids, id)
		}
	}
	if len(ids) > 100 {
		http.Error(w, "at most 100 ids", http.StatusBadRequest)
		return
	}
	out := map[string]storage.TraceSummary{}
	if c.traceQuerier != nil && len(ids) > 0 {
		ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
		defer cancel()
		sums, err := c.traceQuerier.TraceSummaries(ctx, ids)
		if err != nil {
			log.Printf("failed to load trace summaries: %v", err)
			http.Error(w, "failed to load trace summaries", http.StatusInternalServerError)
			return
		}
		out = sums
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(out)
}