fathom/internal/store/store.go

// Package store implements Fathom's on-disk format: columnar day-shards with
// per-chunk dictionaries and blob sidecars.
package store

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

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

// Options configure a Store.
type Options struct {
	Dir string

	// FlushRows and FlushInterval bound how long a span sits in memory before
	// it is durable. Whichever trips first wins. The point of buffering at all
	// is to pay one fsync per flush instead of one per span.
	FlushRows     int
	FlushInterval time.Duration

	// RetentionDays evicts shards older than this. Zero disables age eviction.
	RetentionDays int

	// MaxBytes evicts oldest-first once a table exceeds this total. Zero
	// disables size eviction.
	MaxBytes int64

	// ReadOnly opens shards without creating, recovering, or writing them.
	//
	// Every process that only reads must set this, because recovery truncates:
	// a CLI opening a shard the server is mid-append to would cut off the chunk
	// being written. It is safe to read a file under concurrent append, since a
	// partial tail fails its length or CRC check and iteration stops there.
	ReadOnly bool

	// Now is injectable so tests can control shard dates and retention without
	// sleeping or touching the system clock.
	Now func() time.Time

	// OnFlush is called after rows become durable and therefore queryable.
	//
	// It deliberately fires on flush rather than on ingest. Spans are buffered,
	// so a notification at arrival would send readers to refetch data that is
	// not yet visible to a scan — they would see nothing and conclude nothing
	// had changed. Flush is the moment the store's answer actually changes.
	//
	// Called synchronously on the flushing goroutine, so it must not block.
	OnFlush func(FlushInfo)
}

// FlushInfo describes what a flush made visible.
type FlushInfo struct {
	Spans    int
	Logs     int
	TraceIDs []string // distinct trace ids touched, for targeted refresh
}

func (o *Options) withDefaults() {
	if o.FlushRows <= 0 {
		o.FlushRows = 1000
	}
	if o.FlushInterval <= 0 {
		o.FlushInterval = time.Second
	}
	if o.Now == nil {
		o.Now = time.Now
	}
}

// Store owns the data directory: shard lifecycle, write buffering, retention.
type Store struct {
	opts Options

	mu      sync.Mutex
	shards  map[string]*Shard // keyed by table/date
	spanBuf []model.Span
	logBuf  []model.LogRecord

	flushErr error // last background flush failure, surfaced on the next write

	stop    chan struct{}
	stopped sync.Once
	wg      sync.WaitGroup
}

func Open(opts Options) (*Store, error) {
	opts.withDefaults()
	if opts.Dir == "" {
		return nil, fmt.Errorf("store: no data directory configured")
	}
	if !opts.ReadOnly {
		if err := os.MkdirAll(opts.Dir, 0o755); err != nil {
			return nil, err
		}
	}
	s := &Store{
		opts:   opts,
		shards: make(map[string]*Shard),
		stop:   make(chan struct{}),
	}
	// A read-only store has nothing to flush and must never evict, so it runs
	// no background goroutine at all.
	if !opts.ReadOnly {
		s.wg.Add(1)
		go s.background()
	}
	return s, nil
}

// shardDate is the shard a timestamp belongs to. Sharding by UTC date rather
// than local keeps shard boundaries stable across DST changes and across
// machines that disagree about their timezone.
func shardDate(tsMicros int64) string {
	return time.UnixMicro(tsMicros).UTC().Format("2006-01-02")
}

// AddSpans buffers spans for the next flush.
func (s *Store) AddSpans(spans []model.Span) error {
	s.mu.Lock()
	if err := s.flushErr; err != nil {
		s.flushErr = nil
		s.mu.Unlock()
		return fmt.Errorf("an earlier background flush failed: %w", err)
	}
	s.spanBuf = append(s.spanBuf, spans...)
	full := len(s.spanBuf) >= s.opts.FlushRows
	s.mu.Unlock()

	if full {
		return s.Flush()
	}
	return nil
}

// AddLogs buffers log records for the next flush.
func (s *Store) AddLogs(recs []model.LogRecord) error {
	s.mu.Lock()
	if err := s.flushErr; err != nil {
		s.flushErr = nil
		s.mu.Unlock()
		return fmt.Errorf("an earlier background flush failed: %w", err)
	}
	s.logBuf = append(s.logBuf, recs...)
	full := len(s.logBuf) >= s.opts.FlushRows
	s.mu.Unlock()

	if full {
		return s.Flush()
	}
	return nil
}

// Flush writes everything buffered and fsyncs. Safe to call concurrently and
// safe to call when there is nothing pending.
func (s *Store) Flush() error {
	if s.opts.ReadOnly {
		return nil
	}
	s.mu.Lock()
	spans, logs := s.spanBuf, s.logBuf
	s.spanBuf, s.logBuf = nil, nil
	s.mu.Unlock()

	if len(spans) == 0 && len(logs) == 0 {
		return nil
	}

	// A batch can straddle midnight, so group by the shard each row belongs to
	// rather than assuming one shard per flush.
	spansByDate := make(map[string][]model.Span)
	for _, sp := range spans {
		d := shardDate(sp.StartTS)
		spansByDate[d] = append(spansByDate[d], sp)
	}
	logsByDate := make(map[string][]model.LogRecord)
	for _, r := range logs {
		d := shardDate(r.TS)
		logsByDate[d] = append(logsByDate[d], r)
	}

	var firstErr error
	fail := func(err error) {
		if firstErr == nil {
			firstErr = err
		}
	}

	info := FlushInfo{Spans: len(spans), Logs: len(logs)}
	seen := make(map[string]struct{}, len(spans))
	for _, sp := range spans {
		if _, ok := seen[sp.TraceID]; !ok && sp.TraceID != "" {
			seen[sp.TraceID] = struct{}{}
			info.TraceIDs = append(info.TraceIDs, sp.TraceID)
		}
	}
	for _, r := range logs {
		if _, ok := seen[r.TraceID]; !ok && r.TraceID != "" {
			seen[r.TraceID] = struct{}{}
			info.TraceIDs = append(info.TraceIDs, r.TraceID)
		}
	}

	for date, batch := range spansByDate {
		sh, err := s.shard(model.TableSpans, date)
		if err != nil {
			fail(err)
			continue
		}
		c, err := buildSpanChunk(batch, sh.Blob())
		if err != nil {
			fail(err)
			continue
		}
		if err := sh.AppendChunk(c); err != nil {
			fail(err)
		}
	}
	for date, batch := range logsByDate {
		sh, err := s.shard(model.TableLogs, date)
		if err != nil {
			fail(err)
			continue
		}
		c, err := buildLogChunk(batch, sh.Blob())
		if err != nil {
			fail(err)
			continue
		}
		if err := sh.AppendChunk(c); err != nil {
			fail(err)
		}
	}

	// Only announce a flush that actually landed. Telling readers to refetch
	// after a failed write would have them re-scan for data that is not there.
	if firstErr == nil && s.opts.OnFlush != nil {
		s.opts.OnFlush(info)
	}
	return firstErr
}

func (s *Store) shard(table, date string) (*Shard, error) {
	key := table + "/" + date
	s.mu.Lock()
	defer s.mu.Unlock()
	if sh, ok := s.shards[key]; ok {
		return sh, nil
	}
	open := OpenShard
	if s.opts.ReadOnly {
		open = OpenShardReadOnly
	}
	sh, err := open(s.opts.Dir, table, date)
	if err != nil {
		return nil, err
	}
	s.shards[key] = sh
	return sh, nil
}

// ShardInfo describes a shard on disk without opening it.
type ShardInfo struct {
	Table string
	Date  string
	Start int64 // inclusive, unix micros
	End   int64 // exclusive
	Bytes int64
}

// ListShards enumerates a table's shards, oldest first. It reads the directory
// rather than any in-memory registry so that shards written by a previous run
// are visible without a restart-time index.
func (s *Store) ListShards(table string) ([]ShardInfo, error) {
	tdir := filepath.Join(s.opts.Dir, table)
	ents, err := os.ReadDir(tdir)
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var out []ShardInfo
	for _, e := range ents {
		name := e.Name()
		if !strings.HasSuffix(name, ".fathom") {
			continue
		}
		date := strings.TrimSuffix(name, ".fathom")
		day, err := time.Parse("2006-01-02", date)
		if err != nil {
			continue // not one of ours
		}
		info := ShardInfo{
			Table: table,
			Date:  date,
			Start: day.UTC().UnixMicro(),
			End:   day.UTC().AddDate(0, 0, 1).UnixMicro(),
		}
		if st, err := e.Info(); err == nil {
			info.Bytes = st.Size()
		}
		if st, err := os.Stat(filepath.Join(tdir, date+".blob")); err == nil {
			info.Bytes += st.Size()
		}
		out = append(out, info)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Date < out[j].Date })
	return out, nil
}

// ShardsOverlapping returns the shards a query's time range touches.
//
// This is the entire query planner. There is no index to consult and no
// statistics to consult it with: the date in the filename is the only thing
// that prunes, and for a store scoped to one machine's agent traffic it prunes
// enough.
func (s *Store) ShardsOverlapping(table string, minTS, maxTS int64) ([]ShardInfo, error) {
	all, err := s.ListShards(table)
	if err != nil {
		return nil, err
	}
	var out []ShardInfo
	for _, sh := range all {
		if sh.End > minTS && sh.Start <= maxTS {
			out = append(out, sh)
		}
	}
	return out, nil
}

// OpenShardForRead returns a shard handle for scanning. Handles are cached and
// shared, so readers must not close what they are given.
func (s *Store) OpenShardForRead(table, date string) (*Shard, error) {
	return s.shard(table, date)
}

// background runs the periodic flush and retention sweep.
func (s *Store) background() {
	defer s.wg.Done()
	flush := time.NewTicker(s.opts.FlushInterval)
	defer flush.Stop()
	// Retention is checked far less often than flush; shards are whole days, so
	// there is nothing to collect at flush cadence.
	sweep := time.NewTicker(time.Hour)
	defer sweep.Stop()

	for {
		select {
		case <-s.stop:
			return
		case <-flush.C:
			if err := s.Flush(); err != nil {
				s.mu.Lock()
				s.flushErr = err
				s.mu.Unlock()
			}
		case <-sweep.C:
			s.Evict()
		}
	}
}

// Evict applies the retention policy: age first, then total size, oldest first.
// It returns the shards it removed.
func (s *Store) Evict() []ShardInfo {
	if s.opts.ReadOnly {
		return nil
	}
	var removed []ShardInfo
	for _, table := range []string{model.TableSpans, model.TableLogs} {
		shards, err := s.ListShards(table)
		if err != nil {
			continue
		}

		keep := shards[:0:0]
		if s.opts.RetentionDays > 0 {
			cutoff := s.opts.Now().UTC().AddDate(0, 0, -s.opts.RetentionDays).Format("2006-01-02")
			for _, sh := range shards {
				if sh.Date < cutoff {
					if s.dropShard(sh) {
						removed = append(removed, sh)
					}
					continue
				}
				keep = append(keep, sh)
			}
		} else {
			keep = append(keep, shards...)
		}

		if s.opts.MaxBytes > 0 {
			var total int64
			for _, sh := range keep {
				total += sh.Bytes
			}
			// keep is oldest-first, so dropping from the front evicts the
			// oldest data, matching Scuba's space-based eviction.
			for i := 0; i < len(keep) && total > s.opts.MaxBytes; i++ {
				if s.dropShard(keep[i]) {
					total -= keep[i].Bytes
					removed = append(removed, keep[i])
				}
			}
		}
	}
	return removed
}

// dropShard closes and unlinks a shard and its blob sidecar. Eviction is an
// unlink precisely because the unit of deletion and the unit of partitioning
// are the same.
func (s *Store) dropShard(info ShardInfo) bool {
	key := info.Table + "/" + info.Date
	s.mu.Lock()
	if sh, ok := s.shards[key]; ok {
		sh.Close()
		delete(s.shards, key)
	}
	s.mu.Unlock()

	tdir := filepath.Join(s.opts.Dir, info.Table)
	err1 := os.Remove(filepath.Join(tdir, info.Date+".fathom"))
	os.Remove(filepath.Join(tdir, info.Date+".blob"))
	return err1 == nil
}

// Close flushes and releases every open shard.
func (s *Store) Close() error {
	s.stopped.Do(func() { close(s.stop) })
	s.wg.Wait()

	err := s.Flush()

	s.mu.Lock()
	defer s.mu.Unlock()
	for k, sh := range s.shards {
		if cerr := sh.Close(); cerr != nil && err == nil {
			err = cerr
		}
		delete(s.shards, k)
	}
	return err
}