traces/modules/livestore/live_store.go

package livestore

import (
	"context"
	"errors"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/google/uuid"
	"github.com/grafana/dskit/kv"
	"github.com/grafana/dskit/ring"
	"github.com/grafana/dskit/services"
	"github.com/grafana/tempo/modules/overrides"
	"github.com/grafana/tempo/pkg/flushqueues"
	"github.com/grafana/tempo/pkg/ingest"
	"github.com/grafana/tempo/pkg/tempopb"
	"github.com/grafana/tempo/pkg/util/shutdownmarker"
	"github.com/grafana/tempo/pkg/validation"
	"github.com/grafana/tempo/tempodb/encoding"
	"github.com/grafana/tempo/tempodb/wal"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/twmb/franz-go/pkg/kadm"
	"github.com/twmb/franz-go/pkg/kgo"
	"go.opentelemetry.io/otel/attribute"
)

const (
	liveStoreServiceName = "live-store"

	// ringAutoForgetUnhealthyPeriods is how many consecutive timeout periods an unhealthy instance
	// in the ring will be automatically removed.
	ringAutoForgetUnhealthyPeriods = 2

	PartitionRingKey  = "livestore-partitions"
	PartitionRingName = "livestore-partitions"

	droppedRecordReasonTooOld           = "too_old"
	droppedRecordReasonDecodingFailed   = "decoding_failed"
	droppedRecordReasonInstanceNotFound = "instance_not_found"
)

var ErrStarting = errors.New("live-store is starting")

var (
	// Queue management metrics
	metricCompleteQueueLength = promauto.NewGauge(prometheus.GaugeOpts{
		Namespace: "tempo_live_store",
		Name:      "complete_queue_length",
		Help:      "Number of wal blocks waiting for completion",
	})

	// Block completion metrics
	metricBlocksCompleted = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: "tempo_live_store",
		Name:      "blocks_completed_total",
		Help:      "The total number of blocks completed",
	})
	metricFailedCompletions = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: "tempo_live_store",
		Name:      "failed_completions_total",
		Help:      "The total number of failed block completions",
	})
	metricCompletionRetries = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: "tempo_live_store",
		Name:      "completion_retries_total",
		Help:      "The total number of retries after a failed completion",
	})
	metricCompletionFailedRetries = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: "tempo_live_store",
		Name:      "completion_failed_retries_total",
		Help:      "The total number of failed retries after a failed completion",
	})
	metricCompletionDuration = promauto.NewHistogram(prometheus.HistogramOpts{
		Namespace:                       "tempo_live_store",
		Name:                            "completion_duration_seconds",
		Help:                            "Records the amount of time to complete a block.",
		Buckets:                         prometheus.ExponentialBuckets(1, 2, 10),
		NativeHistogramBucketFactor:     1.1,
		NativeHistogramMaxBucketNumber:  100,
		NativeHistogramMinResetDuration: 1 * time.Hour,
	})

	// Kafka/Ingest specific metrics
	metricRecordsProcessed = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "tempo_live_store",
		Name:      "records_processed_total",
		Help:      "The total number of kafka records processed per tenant.",
	}, []string{"tenant"})

	metricRecordsDropped = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "tempo_live_store",
		Name:      "records_dropped_total",
		Help:      "The total number of kafka records dropped per tenant.",
	}, []string{"tenant", "reason"})

	// Readiness metrics
	metricCatchUpDuration = promauto.NewGauge(prometheus.GaugeOpts{
		Namespace: "tempo_live_store",
		Name:      "catch_up_duration_seconds",
		Help:      "Time spent catching up at startup",
	})

	metricReady = promauto.NewGauge(prometheus.GaugeOpts{
		Namespace: "tempo_live_store",
		Name:      "ready",
		Help:      "1 if ready to serve queries, 0 otherwise",
	})
)

type LiveStore struct {
	services.Service

	cfg    Config
	logger log.Logger
	reg    prometheus.Registerer

	subservicesWatcher *services.FailureWatcher

	ingestPartitionID         int32
	ingestPartitionLifecycler *ring.PartitionInstanceLifecycler
	livestoreLifecycler       *ring.BasicLifecycler

	client  *kgo.Client
	decoder *ingest.Decoder

	reader *PartitionReader

	// Multi-tenant instances
	instancesMtx          sync.RWMutex
	instances             map[string]*instance
	wal                   *wal.WAL
	completeBlockEncoding encoding.VersionedEncoding
	overrides             overrides.Interface

	// Background processing
	ctx                 context.Context // context for the service. all background processes should exit if this is cancelled
	cancel              func()
	wg                  sync.WaitGroup
	completeQueues      *flushqueues.ExclusiveQueues[*completeOp]
	startupComplete     chan struct{} // channel to signal that the starting function has finished. allows background processes to block until the service is fully started
	lagCancel           context.CancelFunc
	readyErr            atomic.Pointer[error] // nil when ready to serve queries
	lastRecordTimeNanos atomic.Int64          // stores timestamp of last consumed record as UnixNano, -1 means not set
}

func New(cfg Config, overridesService overrides.Interface, logger log.Logger, reg prometheus.Registerer, singlePartition bool) (*LiveStore, error) {
	completeBlockEncoding, walEncoding, encErr := coalesceBlockVersions(&cfg)
	if encErr != nil {
		return nil, encErr
	}
	cfg.WAL.Version = walEncoding.Version()

	ctx, cancel := context.WithCancel(context.Background())

	s := &LiveStore{
		cfg:                   cfg,
		logger:                logger,
		reg:                   reg,
		decoder:               ingest.NewDecoder(),
		completeBlockEncoding: completeBlockEncoding,
		ctx:                   ctx,
		cancel:                cancel,
		instances:             make(map[string]*instance),
		overrides:             overridesService,
		completeQueues:        flushqueues.New[*completeOp](cfg.CompleteBlockConcurrency, metricCompleteQueueLength),
		startupComplete:       make(chan struct{}),
	}

	// Initialize ready state to starting
	s.readyErr.Store(&ErrStarting)
	metricReady.Set(0)
	s.lastRecordTimeNanos.Store(-1)

	var err error
	if singlePartition {
		// For single-binary don't require hostname to identify a partition.
		// Assume partition 0.
		s.ingestPartitionID = 0
	} else {
		s.ingestPartitionID, err = ingest.IngesterPartitionID(cfg.Ring.InstanceID)
		if err != nil {
			return nil, fmt.Errorf("calculating livestore partition ID: %w", err)
		}
	}

	if singlePartition {
		s.cfg.IngestConfig.Kafka.ConsumerGroup = "live-store-zone-a"
	} else {
		// TODO: It's probably easier to just use the ID directly
		//  https://raintank-corp.slack.com/archives/C05CAA0ULUF/p1752847274420489
		s.cfg.IngestConfig.Kafka.ConsumerGroup, err = ingest.LiveStoreConsumerGroupID(cfg.Ring.InstanceID)
		if err != nil {
			return nil, fmt.Errorf("calculating livestore consumer group ID: %w", err)
		}
	}

	// setup partition ring
	partitionRingKV := cfg.PartitionRing.KVStore.Mock
	if partitionRingKV == nil {
		partitionRingKV, err = kv.NewClient(cfg.PartitionRing.KVStore, ring.GetPartitionRingCodec(), kv.RegistererWithKVName(reg, PartitionRingName+"-lifecycler"), logger)
		if err != nil {
			return nil, fmt.Errorf("creating KV store for livestore partition ring: %w", err)
		}
	}

	s.ingestPartitionLifecycler = ring.NewPartitionInstanceLifecycler(
		s.cfg.PartitionRing.ToLifecyclerConfig(s.ingestPartitionID, cfg.Ring.InstanceID),
		PartitionRingName,
		PartitionRingKey,
		partitionRingKV,
		logger,
		prometheus.WrapRegistererWithPrefix("tempo_", reg))

	// setup live store read ring
	ringStore := cfg.Ring.KVStore.Mock
	if ringStore == nil {
		ringStore, err = kv.NewClient(
			cfg.Ring.KVStore,
			ring.GetCodec(),
			kv.RegistererWithKVName(prometheus.WrapRegistererWithPrefix("tempo_", reg), "livestore"),
			s.logger,
		)
		if err != nil {
			return nil, fmt.Errorf("create KV store client: %w", err)
		}
	}

	lifecyclerCfg, err := cfg.Ring.ToLifecyclerConfig(0)
	if err != nil {
		return nil, fmt.Errorf("invalid ring lifecycler config: %w", err)
	}

	// Define lifecycler delegates in reverse order (last to be called defined first because they're
	// chained via "next delegate").
	delegate := ring.BasicLifecyclerDelegate(s)
	delegate = ring.NewLeaveOnStoppingDelegate(delegate, s.logger)
	delegate = ring.NewAutoForgetDelegate(ringAutoForgetUnhealthyPeriods*cfg.Ring.HeartbeatTimeout, delegate, s.logger)

	s.livestoreLifecycler, err = ring.NewBasicLifecycler(lifecyclerCfg, liveStoreServiceName, liveStoreServiceName, ringStore, delegate, s.logger, prometheus.WrapRegistererWithPrefix("tempo_", reg))
	if err != nil {
		return nil, fmt.Errorf("create ring lifecycler: %w", err)
	}

	s.subservicesWatcher = services.NewFailureWatcher()
	s.subservicesWatcher.WatchService(s.ingestPartitionLifecycler)
	s.subservicesWatcher.WatchService(s.livestoreLifecycler)

	s.Service = services.NewBasicService(s.starting, s.running, s.stopping)

	return s, nil
}

func (s *LiveStore) starting(ctx context.Context) error {
	// First of all we have to check if the shutdown marker is set. This needs to be done
	// as first thing because, if found, it may change the behaviour of the live-store startup.

	if _, err := os.Stat(s.cfg.ShutdownMarkerDir); os.IsNotExist(err) {
		err := os.MkdirAll(s.cfg.ShutdownMarkerDir, 0o700)
		if err != nil {
			return fmt.Errorf("failed to create shutdown marker directory: %w", err)
		}
	}

	shutdownMarkerPath := shutdownmarker.GetPath(s.cfg.ShutdownMarkerDir)
	if exists, err := shutdownmarker.Exists(shutdownMarkerPath); err != nil {
		return fmt.Errorf("failed to check live-store shutdown marker: %w", err)
	} else if exists {
		level.Info(s.logger).Log("msg", "detected existing shutdown marker, setting prepare for shutdown", "path", shutdownMarkerPath)
		s.setPrepareShutdown()
	}

	var err error
	s.wal, err = wal.New(&s.cfg.WAL)
	if err != nil {
		return fmt.Errorf("failed to create WAL: %w", err)
	}

	err = s.reloadBlocks()
	if err != nil {
		return fmt.Errorf("failed to reload blocks from wal: %w", err)
	}

	for _, inst := range s.getInstances() {
		err = inst.deleteOldBlocks()
		if err != nil {
			level.Warn(s.logger).Log("msg", "failed to delete old blocks", "err", err, "tenant", inst.tenantID)
		}
	}

	forceFromLookback := len(s.getInstances()) == 0
	if forceFromLookback {
		level.Info(s.logger).Log("msg", "no local data found after reload, will force reading from lookback period")
	}

	err = services.StartAndAwaitRunning(ctx, s.ingestPartitionLifecycler)
	if err != nil {
		return fmt.Errorf("failed to start partition lifecycler: %w", err)
	}

	err = services.StartAndAwaitRunning(ctx, s.livestoreLifecycler)
	if err != nil {
		return fmt.Errorf("failed to start livestore lifecycler: %w", err)
	}

	s.client, err = ingest.NewReaderClient(
		s.cfg.IngestConfig.Kafka,
		ingest.NewReaderClientMetrics(liveStoreServiceName, s.reg),
		s.logger,
	)
	if err != nil {
		return fmt.Errorf("failed to create kafka reader client: %w", err)
	}

	err = ingest.WaitForKafkaBroker(ctx, s.client, s.logger)
	if err != nil {
		return fmt.Errorf("failed to start livestore: %w", err)
	}

	lookbackPeriod := 2 * s.cfg.CompleteBlockTimeout
	s.reader, err = NewPartitionReaderForPusher(s.client, s.ingestPartitionID, s.cfg.IngestConfig.Kafka, s.cfg.CommitInterval, lookbackPeriod, forceFromLookback, s.consume, s.logger, s.reg)
	if err != nil {
		return fmt.Errorf("failed to create partition reader: %w", err)
	}
	err = services.StartAndAwaitRunning(ctx, s.reader)
	if err != nil {
		return fmt.Errorf("failed to start partition reader: %w", err)
	}

	lagCtx, cncl := context.WithCancel(s.ctx)
	s.lagCancel = cncl
	// Start exporting partition lag metrics
	ingest.ExportPartitionLagMetrics(
		lagCtx,
		s.client,
		s.logger,
		s.cfg.IngestConfig,
		func() []int32 { return []int32{s.ingestPartitionID} },
		s.client.ForceMetadataRefresh,
	)

	for i := range s.cfg.CompleteBlockConcurrency {
		idx := i
		s.runInBackground(func() {
			s.globalCompleteLoop(idx)
		})
	}

	// allow background processes to start
	s.startAllBackgroundProcesses()

	// Wait for catch-up before marking ready (if enabled)
	if err := s.waitForCatchUp(ctx); err != nil {
		return fmt.Errorf("failed to catch up: %w", err)
	}

	// Mark as ready at end of starting()
	s.readyErr.Store(nil)
	metricReady.Set(1)
	level.Info(s.logger).Log("msg", "live-store ready to serve queries")

	return nil
}

func (s *LiveStore) running(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return nil
	case err := <-s.subservicesWatcher.Chan():
		return fmt.Errorf("live-store subservice failed: %w", err)
	}
}

func (s *LiveStore) stopping(error) error {
	// Remove partition owner from ring on shutdown if configured.
	// On startup, createPartitionAndRegisterOwner() re-registers the owner immediately.
	if s.cfg.RemoveOwnerOnShutdown {
		level.Info(s.logger).Log("msg", "Unregistering partition owner")
		s.ingestPartitionLifecycler.SetRemoveOwnerOnShutdown(true)
	}

	// Stop the kafka lag background worker.
	s.lagCancel()
	metricReady.Set(0)

	// Stop consuming
	err := services.StopAndAwaitTerminated(context.Background(), s.reader)
	if err != nil {
		level.Warn(s.logger).Log("msg", "failed to stop reader", "err", err)
		return err
	}

	// Reset lag metrics for our partition when stopping
	ingest.ResetLagMetricsForRevokedPartitions(s.cfg.IngestConfig.Kafka.ConsumerGroup, []int32{s.ingestPartitionID})

	// Flush all data to disk
	s.cutAllInstancesToWal()

	// Remove the shutdown marker if it exists since we are shutting down
	shutdownMarkerPath := shutdownmarker.GetPath(s.cfg.ShutdownMarkerDir)
	if err := shutdownmarker.Remove(shutdownMarkerPath); err != nil {
		level.Warn(s.logger).Log("msg", "failed to remove shutdown marker", "path", shutdownMarkerPath, "err", err)
	}

	if s.cfg.holdAllBackgroundProcesses { // nothing to do
		return nil
	}

	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()

	timeout := time.NewTimer(s.cfg.InstanceCleanupPeriod)
	defer timeout.Stop()

	s.stopAllBackgroundProcesses()

	return nil
}

func (s *LiveStore) waitForCatchUp(ctx context.Context) error {
	// If disabled (ReadinessTargetLag == 0), mark ready immediately
	// This preserves backward compatibility
	if s.cfg.ReadinessTargetLag == 0 {
		level.Info(s.logger).Log("msg", "catch-up waiting disabled (readiness_target_lag=0)")
		return nil
	}

	startTime := time.Now()

	ticker := time.NewTicker(time.Second) // Check every second
	defer ticker.Stop()

	level.Info(s.logger).Log(
		"msg", "waiting for Kafka catch-up",
		"target_lag", s.cfg.ReadinessTargetLag,
		"max_wait", s.cfg.ReadinessMaxWait,
	)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			elapsed := time.Since(startTime)

			// Check max wait timeout
			if elapsed >= s.cfg.ReadinessMaxWait {
				level.Warn(s.logger).Log(
					"msg", "max catch-up wait exceeded, proceeding anyway",
					"elapsed", elapsed,
					"max_wait", s.cfg.ReadinessMaxWait,
				)
				metricCatchUpDuration.Set(elapsed.Seconds())
				return nil
			}

			// Calculate current lag
			lag := s.calculateTimeLag(1000)
			if lag == nil {
				level.Debug(s.logger).Log("msg", "catch-up lag could not be determined, waiting")
				continue
			}

			level.Debug(s.logger).Log(
				"msg", "catch-up progress",
				"current_lag", *lag,
				"target_lag", s.cfg.ReadinessTargetLag,
				"elapsed", elapsed,
			)

			if *lag <= s.cfg.ReadinessTargetLag {
				level.Info(s.logger).Log(
					"msg", "caught up with Kafka",
					"final_lag", *lag,
					"target_lag", s.cfg.ReadinessTargetLag,
					"elapsed", elapsed,
				)
				metricCatchUpDuration.Set(elapsed.Seconds())
				return nil
			}
		}
	}
}

// Calculate lag based on parameters populated by PartitionReader
// Edge cases:
// - empty partition = no lag
// - nothing has been fetched yet = indeterminate
// - we know the watermark but nothing has been consumed yet = indeterminate
//
// It takes lagShortcutThreshold to shortcut calculations if the lag is close to the end of the partition.
// To disable the shortcut, set lagShortcutThreshold to a negative value.
func (s *LiveStore) calculateTimeLag(lagShortcutThreshold int64) *time.Duration {
	// Use cached high watermark from fetch responses (avoids extra API call)
	lag := s.reader.lag.Load()
	zero := time.Duration(0)

	// If we haven't performed any fetches yet, we can't determine lag
	if lag < 0 {
		level.Debug(s.logger).Log("msg", "High watermark not set yet")
		return nil
	}

	// Check if we are near end or partition is empty
	// Arbitrary value picked to shortcut calculations
	if lagShortcutThreshold >= 0 && lag <= lagShortcutThreshold {
		level.Debug(s.logger).Log(
			"msg", "At or close to partition end",
			"lag", lag)
		return &zero
	}

	nanos := s.lastRecordTimeNanos.Load()
	if nanos < 0 {
		level.Debug(s.logger).Log("msg", "No last record yet")
		// Haven't consumed records - check if offset at end
		return nil // Not caught up yet, can't determine lag
	}

	// Potential race condition that can result in negative lag?
	// Assuming strictly monotonic timestamps in Kafka, can't cause an issue
	lastRecordTime := time.Unix(0, nanos)
	recordLag := time.Since(lastRecordTime)
	return &recordLag
}

func (s *LiveStore) consume(ctx context.Context, rs recordIter, now time.Time) (*kadm.Offset, error) {
	defer s.decoder.Reset()
	_, span := tracer.Start(ctx, "LiveStore.consume")
	defer span.End()

	recordCount := 0
	var lastRecord *kgo.Record

	cutoff := now.Add(-s.cfg.CompleteBlockTimeout)
	// Process records by tenant
	for !rs.Done() {
		record := rs.Next()
		tenant := string(record.Key)

		if record.Timestamp.Before(cutoff) {
			metricRecordsDropped.WithLabelValues(tenant, droppedRecordReasonTooOld).Inc()
			lastRecord = record
			continue
		}

		s.decoder.Reset()
		pushReq, err := s.decoder.Decode(record.Value)
		if err != nil {
			metricRecordsDropped.WithLabelValues(tenant, droppedRecordReasonDecodingFailed).Inc()
			level.Error(s.logger).Log("msg", "failed to decoded record", "tenant", tenant, "err", err)
			span.RecordError(err)
			lastRecord = record
			continue
		}

		// Get or create tenant instance
		inst, err := s.getOrCreateInstance(tenant)
		if err != nil {
			metricRecordsDropped.WithLabelValues(tenant, droppedRecordReasonInstanceNotFound).Inc()
			level.Error(s.logger).Log("msg", "failed to get instance for tenant", "tenant", tenant, "err", err)
			span.RecordError(err)
			lastRecord = record
			continue
		}

		// Push data to tenant instance
		inst.pushBytes(ctx, record.Timestamp, pushReq)

		// Track partition lag in seconds
		lag := now.Sub(record.Timestamp)
		ingest.SetPartitionLagSeconds(s.cfg.IngestConfig.Kafka.ConsumerGroup, record.Partition, lag)

		metricRecordsProcessed.WithLabelValues(tenant).Inc()
		recordCount++
		lastRecord = record
	}

	span.SetAttributes(attribute.Int("records_count", recordCount))

	if lastRecord == nil {
		return nil, nil
	}

	// Store the timestamp of the last consumed record for lag calculation
	s.lastRecordTimeNanos.Store(lastRecord.Timestamp.UnixNano())

	offset := kadm.NewOffsetFromRecord(lastRecord)
	return &offset, nil
}

func (s *LiveStore) getInstance(tenantID string) (*instance, bool) {
	s.instancesMtx.RLock()
	defer s.instancesMtx.RUnlock()
	inst, ok := s.instances[tenantID]
	return inst, ok
}

func (s *LiveStore) getOrCreateInstance(tenantID string) (*instance, error) {
	s.instancesMtx.RLock()
	inst, ok := s.instances[tenantID]
	s.instancesMtx.RUnlock()

	if ok {
		return inst, nil
	}

	s.instancesMtx.Lock()
	defer s.instancesMtx.Unlock()

	// Double-check in case another goroutine created it
	if inst, ok := s.instances[tenantID]; ok {
		return inst, nil
	}

	// Create new instance
	inst, err := newInstance(tenantID, s.cfg, s.wal, s.completeBlockEncoding, s.overrides, s.logger)
	if err != nil {
		return nil, fmt.Errorf("failed to create instance for tenant %s: %w", tenantID, err)
	}

	s.instances[tenantID] = inst

	s.runInBackground(func() {
		s.perTenantCutToWalLoop(inst)
	})
	s.runInBackground(func() {
		s.perTenantCleanupLoop(inst)
	})

	return inst, nil
}

func (s *LiveStore) getInstances() []*instance {
	s.instancesMtx.RLock()
	defer s.instancesMtx.RUnlock()
	instances := make([]*instance, 0, len(s.instances))
	for _, inst := range s.instances {
		instances = append(instances, inst)
	}
	return instances
}

func (s *LiveStore) cutAllInstancesToWal() {
	instances := s.getInstances()

	for _, instance := range instances {
		s.cutOneInstanceToWal(instance, true)
	}
}

func (s *LiveStore) cutOneInstanceToWal(inst *instance, immediate bool) {
	// Regular trace cuts (live traces -> head block)
	err := inst.cutIdleTraces(immediate)
	if err != nil {
		level.Error(s.logger).Log("msg", "failed to cut idle traces", "tenant", inst.tenantID, "err", err)
	}

	// Regular block cuts
	blockID, err := inst.cutBlocks(immediate)
	if err != nil {
		level.Error(s.logger).Log("msg", "failed to cut blocks", "tenant", inst.tenantID, "err", err)
	}

	// If head block is cut, enqueue complete operation
	if blockID != uuid.Nil {
		err = s.enqueueCompleteOp(inst.tenantID, blockID, false)
		if err != nil {
			level.Error(s.logger).Log("msg", "failed to enqueue complete operation", "tenant", inst.tenantID, "err", err)
			return
		}
	}
}

// CheckReady returns nil if the live-store is ready to serve queries
func (s *LiveStore) CheckReady(_ context.Context) error {
	if err := s.readyErr.Load(); err != nil {
		return *err
	}
	return nil
}

// OnRingInstanceRegister implements ring.BasicLifecyclerDelegate
func (s *LiveStore) OnRingInstanceRegister(_ *ring.BasicLifecycler, _ ring.Desc, _ bool, _ string, _ ring.InstanceDesc) (ring.InstanceState, ring.Tokens) {
	return ring.ACTIVE, nil // no tokens needed for the livestore ring, we just need to be in the ring for service discovery
}

// OnRingInstanceTokens implements ring.BasicLifecyclerDelegate
func (s *LiveStore) OnRingInstanceTokens(*ring.BasicLifecycler, ring.Tokens) {
}

// OnRingInstanceStopping implements ring.BasicLifecyclerDelegate
func (s *LiveStore) OnRingInstanceStopping(*ring.BasicLifecycler) {
}

// OnRingInstanceHeartbeat implements ring.BasicLifecyclerDelegate
func (s *LiveStore) OnRingInstanceHeartbeat(*ring.BasicLifecycler, *ring.Desc, *ring.InstanceDesc) {
}

// FindTraceByID implements tempopb.Querier
func (s *LiveStore) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDRequest) (*tempopb.TraceByIDResponse, error) {
	return withInstance(ctx, s, func(inst *instance) (*tempopb.TraceByIDResponse, error) {
		return inst.FindByTraceID(ctx, req.TraceID, req.AllowPartialTrace)
	})
}

// SearchRecent implements tempopb.Querier
func (s *LiveStore) SearchRecent(ctx context.Context, req *tempopb.SearchRequest) (*tempopb.SearchResponse, error) {
	if s.isLagged(int64(req.End) * 1e9) { // convert seconds to nanoseconds
		return nil, errLagged
	}
	return withInstance(ctx, s, func(inst *instance) (*tempopb.SearchResponse, error) {
		return inst.Search(ctx, req)
	})
}

// SearchBlock implements tempopb.Querier
func (s *LiveStore) SearchBlock(_ context.Context, _ *tempopb.SearchBlockRequest) (*tempopb.SearchResponse, error) {
	return nil, fmt.Errorf("SearchBlock not implemented in livestore")
}

// SearchTags implements tempopb.Querier
func (s *LiveStore) SearchTags(ctx context.Context, req *tempopb.SearchTagsRequest) (*tempopb.SearchTagsResponse, error) {
	return withInstance(ctx, s, func(inst *instance) (*tempopb.SearchTagsResponse, error) {
		return inst.SearchTags(ctx, req.Scope)
	})
}

// SearchTagsV2 implements tempopb.Querier
func (s *LiveStore) SearchTagsV2(ctx context.Context, req *tempopb.SearchTagsRequest) (*tempopb.SearchTagsV2Response, error) {
	return withInstance(ctx, s, func(inst *instance) (*tempopb.SearchTagsV2Response, error) {
		return inst.SearchTagsV2(ctx, req)
	})
}

// SearchTagValues implements tempopb.Querier
func (s *LiveStore) SearchTagValues(ctx context.Context, req *tempopb.SearchTagValuesRequest) (*tempopb.SearchTagValuesResponse, error) {
	return withInstance(ctx, s, func(inst *instance) (*tempopb.SearchTagValuesResponse, error) {
		return inst.SearchTagValues(ctx, req)
	})
}

// SearchTagValuesV2 implements tempopb.Querier
func (s *LiveStore) SearchTagValuesV2(ctx context.Context, req *tempopb.SearchTagValuesRequest) (*tempopb.SearchTagValuesV2Response, error) {
	return withInstance(ctx, s, func(inst *instance) (*tempopb.SearchTagValuesV2Response, error) {
		return inst.SearchTagValuesV2(ctx, req)
	})
}

// PushSpans implements tempopb.MetricsGeneratorServer
func (s *LiveStore) PushSpans(_ context.Context, _ *tempopb.PushSpansRequest) (*tempopb.PushResponse, error) {
	return nil, fmt.Errorf("PushSpans not implemented in livestore")
}

// GetMetrics implements tempopb.MetricsGeneratorServer
func (s *LiveStore) GetMetrics(_ context.Context, _ *tempopb.SpanMetricsRequest) (*tempopb.SpanMetricsResponse, error) {
	// Keep this stub until r241 is fully rolled out. After that, we can remove
	// GetMetrics here by switching LiveStore from MetricsGenerator to MetricsService.
	return nil, fmt.Errorf("GetMetrics not implemented in livestore")
}

// QueryRange implements tempopb.MetricsGeneratorServer
func (s *LiveStore) QueryRange(ctx context.Context, req *tempopb.QueryRangeRequest) (*tempopb.QueryRangeResponse, error) {
	if s.isLagged(int64(req.End)) { // end param is already nanos, no need to convert
		return nil, errLagged
	}
	return withInstance(ctx, s, func(inst *instance) (*tempopb.QueryRangeResponse, error) {
		return inst.QueryRange(ctx, req)
	})
}

var errLagged = errors.New("cannot guarantee complete results")

func (s *LiveStore) isLagged(endNanos int64) bool {
	if !s.cfg.FailOnHighLag { // if config disabled, never lagged
		return false
	}
	lag := s.calculateTimeLag(0)
	if lag == nil { // lag is unknown
		return true // prefer error over potentially incomplete results
	}
	return time.Since(time.Unix(0, endNanos)) < *lag
}

// withInstance extracts the tenant ID from the context, gets the instance,
// and calls the provided function if the instance exists. If the instance
// doesn't exist, it returns the zero value.
func withInstance[T any](ctx context.Context, s *LiveStore, fn func(*instance) (*T, error)) (*T, error) {
	var defaultValue T

	// Check readiness before processing query
	if err := s.CheckReady(ctx); err != nil {
		return &defaultValue, err
	}

	instanceID, err := validation.ExtractValidTenantID(ctx)
	if err != nil {
		return &defaultValue, err
	}

	inst, found := s.getInstance(instanceID)
	if inst == nil || !found {
		return &defaultValue, nil
	}

	return fn(inst)
}