fathom/internal/store/schema.go

package store

import (
	"sort"

	"git.z.gd/zzg/fathom/internal/model"
)

// Column names. These are the identifiers the query grammar accepts, so they
// are part of Fathom's user-facing surface, not just its file format.
const (
	ColTraceID      = "trace_id"
	ColSpanID       = "span_id"
	ColParentSpanID = "parent_span_id"
	ColName         = "name"
	ColService      = "service"
	ColKind         = "kind"
	ColStatus       = "status"
	ColStart        = "start_ts"
	ColEnd          = "end_ts"
	ColDuration     = "duration"
	ColSampleRate   = "sample_rate"
	ColInput        = "input"
	ColOutput       = "output"

	ColEventName = "event_name"
	ColSeverity  = "severity"
	ColTS        = "ts"
	ColBody      = "body"

	// ColTime is the mandatory time predicate's name in the query grammar. It
	// resolves to start_ts on spans and ts on logs, so a query does not have to
	// know which table it is scanning.
	ColTime = "time"
)

// TimeColumn returns the column a table's mandatory time range filters on.
func TimeColumn(table string) string {
	if table == model.TableLogs {
		return ColTS
	}
	return ColStart
}

// chunkBuilder accumulates rows and produces a Chunk. It owns the dictionary,
// so every string interned during a batch shares one table and the codes stay
// dense.
type chunkBuilder struct {
	dict *Dict
	rows int

	cols   []*Column
	byName map[string]*Column

	attrKeys  []uint32
	attrVals  []uint32
	attrTypes []uint8
}

func newChunkBuilder() *chunkBuilder {
	return &chunkBuilder{dict: NewDict(), byName: make(map[string]*Column)}
}

func (b *chunkBuilder) col(name string, t ColType) *Column {
	if c, ok := b.byName[name]; ok {
		return c
	}
	c := &Column{Name: name, Type: t}
	b.cols = append(b.cols, c)
	b.byName[name] = c
	return c
}

func (b *chunkBuilder) putDict(name, v string) {
	c := b.col(name, ColDict)
	c.Codes = append(c.Codes, b.dict.Put(v))
}

func (b *chunkBuilder) putI64(name string, v int64) {
	c := b.col(name, ColInt64)
	c.I64 = append(c.I64, v)
}

func (b *chunkBuilder) putU8(name string, v uint8) {
	c := b.col(name, ColUint8)
	c.U8 = append(c.U8, v)
}

func (b *chunkBuilder) putF64(name string, v float64) {
	c := b.col(name, ColFloat64)
	c.F64 = append(c.F64, v)
}

func (b *chunkBuilder) putBlob(name string, v BlobRef) {
	c := b.col(name, ColBlobRef)
	c.Blobs = append(c.Blobs, v)
}

// putAttrs appends a row's attributes to the side table and records its window.
// Keys are sorted so that an identical attribute map always encodes to
// identical bytes, which makes chunks reproducible and round-trip tests exact.
func (b *chunkBuilder) putAttrs(a model.Attrs) {
	start := len(b.attrKeys)
	keys := make([]string, 0, len(a))
	for k := range a {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := a[k]
		b.attrKeys = append(b.attrKeys, b.dict.Put(k))
		b.attrVals = append(b.attrVals, b.dict.Put(v.Str))
		b.attrTypes = append(b.attrTypes, uint8(v.Type))
	}
	b.putI64(colAttrStart, int64(start))
	b.putI64(colAttrCount, int64(len(keys)))
}

func (b *chunkBuilder) build() *Chunk {
	c := &Chunk{
		Dict:      b.dict,
		Rows:      b.rows,
		AttrKeys:  b.attrKeys,
		AttrVals:  b.attrVals,
		AttrTypes: b.attrTypes,
	}
	c.Cols = make([]Column, len(b.cols))
	for i, col := range b.cols {
		c.Cols[i] = *col
	}
	return c
}

// buildSpanChunk encodes spans, writing their large text to blob first so the
// chunk can reference it. A blob write that fails is reported rather than
// silently dropping the span's payload.
func buildSpanChunk(spans []model.Span, blob *BlobFile) (*Chunk, error) {
	b := newChunkBuilder()
	for i := range spans {
		s := &spans[i]

		inRef, err := blob.Append(s.Input)
		if err != nil {
			return nil, err
		}
		outRef, err := blob.Append(s.Output)
		if err != nil {
			return nil, err
		}

		b.putDict(ColTraceID, s.TraceID)
		b.putDict(ColSpanID, s.SpanID)
		b.putDict(ColParentSpanID, s.ParentSpanID)
		b.putDict(ColName, s.Name)
		b.putDict(ColService, s.Service)
		b.putU8(ColKind, uint8(s.Kind))
		b.putU8(ColStatus, uint8(s.Status))
		b.putI64(ColStart, s.StartTS)
		b.putI64(ColEnd, s.EndTS)
		// Duration is stored rather than derived. It is the single most queried
		// and sorted-on quantity in the UI, and materialising it costs a varint
		// per span while saving a subtraction per row on every scan that orders
		// by it.
		b.putI64(ColDuration, s.DurationMicros())
		rate := s.SampleRate
		if rate <= 0 {
			rate = 1
		}
		b.putF64(ColSampleRate, rate)
		b.putBlob(ColInput, inRef)
		b.putBlob(ColOutput, outRef)
		b.putAttrs(s.Attrs)
		b.rows++
	}
	return b.build(), nil
}

func buildLogChunk(recs []model.LogRecord, blob *BlobFile) (*Chunk, error) {
	b := newChunkBuilder()
	for i := range recs {
		r := &recs[i]

		bodyRef, err := blob.Append(r.Body)
		if err != nil {
			return nil, err
		}

		b.putDict(ColTraceID, r.TraceID)
		b.putDict(ColSpanID, r.SpanID)
		b.putDict(ColEventName, r.EventName)
		b.putU8(ColSeverity, r.Severity)
		b.putI64(ColTS, r.TS)
		b.putBlob(ColBody, bodyRef)
		b.putAttrs(r.Attrs)
		b.rows++
	}
	return b.build(), nil
}

// readAttrs materialises row i's attribute map.
//
// This allocates, so the scanner avoids it: predicates over attributes compare
// dictionary codes against the side table directly. It exists for the read
// paths that genuinely need the whole map — trace detail and row output.
func readAttrs(c *Chunk, i int) model.Attrs {
	startCol, ok1 := c.Col(colAttrStart)
	countCol, ok2 := c.Col(colAttrCount)
	if !ok1 || !ok2 || i >= len(startCol.I64) || i >= len(countCol.I64) {
		return nil
	}
	start, n := int(startCol.I64[i]), int(countCol.I64[i])
	if n == 0 {
		return nil
	}
	if start < 0 || start+n > len(c.AttrKeys) {
		return nil // damaged chunk; a query should lose attributes, not crash
	}
	a := make(model.Attrs, n)
	for j := start; j < start+n; j++ {
		a[c.Dict.Get(c.AttrKeys[j])] = model.Value{
			Type: model.ValueType(c.AttrTypes[j]),
			Str:  c.Dict.Get(c.AttrVals[j]),
		}
	}
	return a
}

// readSpan materialises row i as a Span, without its blobs. Callers that need
// input and output resolve them through the shard's BlobFile, which is the
// laziness the split storage exists to allow.
func readSpan(c *Chunk, i int) model.Span {
	d := func(name string) string {
		col, ok := c.Col(name)
		if !ok || i >= len(col.Codes) {
			return ""
		}
		return c.Dict.Get(col.Codes[i])
	}
	i64 := func(name string) int64 {
		col, ok := c.Col(name)
		if !ok || i >= len(col.I64) {
			return 0
		}
		return col.I64[i]
	}
	u8 := func(name string) uint8 {
		col, ok := c.Col(name)
		if !ok || i >= len(col.U8) {
			return 0
		}
		return col.U8[i]
	}
	rate := float64(1)
	if col, ok := c.Col(ColSampleRate); ok && i < len(col.F64) {
		rate = col.F64[i]
	}
	return model.Span{
		TraceID:      d(ColTraceID),
		SpanID:       d(ColSpanID),
		ParentSpanID: d(ColParentSpanID),
		Name:         d(ColName),
		Service:      d(ColService),
		Kind:         model.Kind(u8(ColKind)),
		Status:       model.Status(u8(ColStatus)),
		StartTS:      i64(ColStart),
		EndTS:        i64(ColEnd),
		SampleRate:   rate,
		Attrs:        readAttrs(c, i),
	}
}

func readLog(c *Chunk, i int) model.LogRecord {
	d := func(name string) string {
		col, ok := c.Col(name)
		if !ok || i >= len(col.Codes) {
			return ""
		}
		return c.Dict.Get(col.Codes[i])
	}
	var ts int64
	if col, ok := c.Col(ColTS); ok && i < len(col.I64) {
		ts = col.I64[i]
	}
	var sev uint8
	if col, ok := c.Col(ColSeverity); ok && i < len(col.U8) {
		sev = col.U8[i]
	}
	return model.LogRecord{
		TraceID:   d(ColTraceID),
		SpanID:    d(ColSpanID),
		EventName: d(ColEventName),
		Severity:  sev,
		TS:        ts,
		Attrs:     readAttrs(c, i),
	}
}

// blobAt returns row i's reference in the named blob column.
func blobAt(c *Chunk, name string, i int) BlobRef {
	col, ok := c.Col(name)
	if !ok || i >= len(col.Blobs) {
		return BlobRef{}
	}
	return col.Blobs[i]
}