traces/modules/generator/storage/instance_test.go

package storage

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/go-kit/log"
	"github.com/grafana/dskit/test"
	"github.com/grafana/dskit/user"
	"github.com/prometheus/client_golang/prometheus"
	prometheus_common_config "github.com/prometheus/common/config"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/model/exemplar"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/prompb"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"

	"github.com/grafana/tempo/modules/overrides/histograms"
)

// Verify basic functionality like sending metrics and exemplars, buffering and retrying failed
// requests.
func TestInstance(t *testing.T) {
	t.Parallel()
	var err error
	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))

	mockServer := newMockPrometheusRemoteWriterServer(logger)
	defer mockServer.close()

	var cfg Config
	cfg.RegisterFlagsAndApplyDefaults("", nil)
	cfg.Path = t.TempDir()
	cfg.RemoteWrite = mockServer.remoteWriteConfig()

	instance, err := New(&cfg, &mockOverrides{}, "test-tenant", &noopRegisterer{}, logger)
	require.NoError(t, err)

	// Refuse requests - the WAL should buffer data until requests succeed
	mockServer.refuseRequests.Store(true)

	sendCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Append some data every second
	go poll(sendCtx, time.Second, func() {
		appender := instance.Appender(context.Background())

		lbls := labels.FromMap(map[string]string{"__name__": "my-metrics"})
		ref, err := appender.Append(0, lbls, time.Now().UnixMilli(), 1.0)
		assert.NoError(t, err)

		_, err = appender.AppendExemplar(ref, lbls, exemplar.Exemplar{
			Labels: labels.FromMap(map[string]string{"traceID": "123"}),
			Value:  1.2,
		})
		assert.NoError(t, err)

		if sendCtx.Err() != nil {
			return
		}

		err = appender.Commit()
		assert.NoError(t, err)
	})

	// Wait until remote.Storage has tried at least once to send data
	test.Poll(t, 30*time.Second, true, func() interface{} {
		mockServer.mtx.Lock()
		defer mockServer.mtx.Unlock()

		return mockServer.refusedRequests > 0
	})

	// Allow requests
	mockServer.refuseRequests.Store(false)

	// Shutdown the instance - even though previous requests failed, remote.Storage should flush pending data
	cancel()
	err = instance.Close()
	assert.NoError(t, err)

	// WAL should be empty again
	entries, err := os.ReadDir(cfg.Path)
	assert.NoError(t, err)
	assert.Len(t, entries, 0)

	// Verify we received metrics
	assert.Len(t, mockServer.timeSeries, 1)
	assert.Contains(t, mockServer.timeSeries, "test-tenant")
	// We should have received at least 2 time series: one for the sample and one for the examplar
	assert.GreaterOrEqual(t, len(mockServer.timeSeries["test-tenant"]), 2)
}

// Verify multiple instances function next to each other, don't trample over each other and are isolated.
func TestInstance_multiTenancy(t *testing.T) {
	t.Parallel()
	var err error
	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))

	mockServer := newMockPrometheusRemoteWriterServer(logger)
	defer mockServer.close()

	var cfg Config
	cfg.RegisterFlagsAndApplyDefaults("", nil)
	cfg.Path = t.TempDir()
	cfg.RemoteWrite = mockServer.remoteWriteConfig()

	var instances []Storage

	for i := 0; i < 3; i++ {
		instance, err := New(&cfg, &mockOverrides{}, strconv.Itoa(i), &noopRegisterer{}, logger)
		assert.NoError(t, err)
		instances = append(instances, instance)
	}

	sendCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Append some data every second
	go poll(sendCtx, time.Second, func() {
		for i, instance := range instances {
			appender := instance.Appender(context.Background())

			lbls := labels.FromMap(map[string]string{"__name__": "my-metric"})
			_, err := appender.Append(0, lbls, time.Now().UnixMilli(), float64(i))
			assert.NoError(t, err)

			if sendCtx.Err() != nil {
				return
			}

			err = appender.Commit()
			assert.NoError(t, err)
		}
	})

	// Wait until every tenant received at least one request
	test.Poll(t, 45*time.Second, true, func() interface{} {
		mockServer.mtx.Lock()
		defer mockServer.mtx.Unlock()

		for i := range instances {
			if mockServer.acceptedRequests[strconv.Itoa(i)] == 0 {
				return false
			}
		}
		return true
	})

	cancel()
	for _, instance := range instances {
		// Shutdown the instance - remote write should flush pending data
		err = instance.Close()
		assert.NoError(t, err)
	}

	// WAL should be empty again
	entries, err := os.ReadDir(cfg.Path)
	assert.NoError(t, err)
	require.Len(t, entries, 0)

	for i := range instances {
		lenOk := assert.GreaterOrEqual(t, len(mockServer.timeSeries[strconv.Itoa(i)]), 1, "instance %d did not receive the expected amount of time series", i)
		if lenOk {
			sample := mockServer.timeSeries[strconv.Itoa(i)][0]
			assert.Equal(t, float64(i), sample.GetSamples()[0].GetValue())
		}
	}
}

func TestInstance_cantWriteToWAL(t *testing.T) {
	var cfg Config
	cfg.RegisterFlagsAndApplyDefaults("", nil)

	// We are obviously not allowed to write here
	cfg.Path = "/root"

	// We should be able to attempt to create the instance multiple times
	_, err := New(&cfg, &mockOverrides{}, "test-tenant", &noopRegisterer{}, log.NewNopLogger())
	require.Error(t, err)
	_, err = New(&cfg, &mockOverrides{}, "test-tenant", &noopRegisterer{}, log.NewNopLogger())
	require.Error(t, err)
}

func TestInstance_remoteWriteHeaders(t *testing.T) {
	var err error
	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))

	mockServer := newMockPrometheusRemoteWriterServer(logger)
	defer mockServer.close()

	var cfg Config
	cfg.RegisterFlagsAndApplyDefaults("", nil)
	cfg.Path = t.TempDir()
	cfg.RemoteWrite = mockServer.remoteWriteConfig()

	headers := map[string]string{user.OrgIDHeaderName: "my-other-tenant"}

	instance, err := New(&cfg, &mockOverrides{headers, histograms.HistogramMethodClassic}, "test-tenant", &noopRegisterer{}, logger)
	require.NoError(t, err)

	// Refuse requests - the WAL should buffer data until requests succeed
	mockServer.refuseRequests.Store(true)

	sendCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Append some data every second
	go poll(sendCtx, time.Second, func() {
		appender := instance.Appender(context.Background())

		lbls := labels.FromMap(map[string]string{"__name__": "my-metrics"})
		ref, err := appender.Append(0, lbls, time.Now().UnixMilli(), 1.0)
		assert.NoError(t, err)

		_, err = appender.AppendExemplar(ref, lbls, exemplar.Exemplar{
			Labels: labels.FromMap(map[string]string{"traceID": "123"}),
			Value:  1.2,
		})
		assert.NoError(t, err)

		if sendCtx.Err() != nil {
			return
		}

		err = appender.Commit()
		assert.NoError(t, err)
	})

	// Wait until remote.Storage has tried at least once to send data
	test.Poll(t, 30*time.Second, true, func() interface{} {
		mockServer.mtx.Lock()
		defer mockServer.mtx.Unlock()

		return mockServer.refusedRequests > 0
	})

	// Allow requests
	mockServer.refuseRequests.Store(false)

	// Shutdown the instance - even though previous requests failed, remote.Storage should flush pending data
	cancel()
	err = instance.Close()
	assert.NoError(t, err)

	// WAL should be empty again
	entries, err := os.ReadDir(cfg.Path)
	assert.NoError(t, err)
	assert.Len(t, entries, 0)

	// Verify we received metrics
	assert.Len(t, mockServer.timeSeries, 1)
	assert.Contains(t, mockServer.timeSeries, "my-other-tenant")
	// We should have received at least 2 time series: one for the sample and one for the examplar
	assert.GreaterOrEqual(t, len(mockServer.timeSeries["my-other-tenant"]), 2)
}

type mockPrometheusRemoteWriteServer struct {
	mtx sync.Mutex

	server         *httptest.Server
	refuseRequests *atomic.Bool
	timeSeries     map[string][]prompb.TimeSeries

	// metrics
	refusedRequests  int
	acceptedRequests map[string]int

	logger log.Logger
}

func newMockPrometheusRemoteWriterServer(logger log.Logger) *mockPrometheusRemoteWriteServer {
	logger = log.With(logger, "component", "mockserver")

	m := &mockPrometheusRemoteWriteServer{
		refuseRequests:   atomic.NewBool(false),
		timeSeries:       make(map[string][]prompb.TimeSeries),
		acceptedRequests: map[string]int{},
		logger:           logger,
	}
	m.server = httptest.NewServer(m)
	return m
}

func (m *mockPrometheusRemoteWriteServer) remoteWriteConfig() []config.RemoteWriteConfig {
	rwCfg := config.DefaultRemoteWriteConfig
	rwCfg.URL = &prometheus_common_config.URL{URL: urlMustParse(fmt.Sprintf("%s/receive", m.server.URL))}
	rwCfg.SendExemplars = true
	// Aggressive queue settings to speed up tests
	rwCfg.QueueConfig.BatchSendDeadline = model.Duration(10 * time.Millisecond)
	return []config.RemoteWriteConfig{rwCfg}
}

func (m *mockPrometheusRemoteWriteServer) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	m.mtx.Lock()
	defer m.mtx.Unlock()

	if m.refuseRequests.Load() {
		m.logger.Log("msg", "refusing request")
		m.refusedRequests++
		http.Error(res, "request refused", http.StatusServiceUnavailable)
		return
	}

	tenant := req.Header.Get(user.OrgIDHeaderName)
	m.logger.Log("msg", "received request", "tenant", tenant)
	m.acceptedRequests[tenant]++

	writeRequest, err := remote.DecodeWriteRequest(req.Body)
	if err != nil {
		http.Error(res, err.Error(), http.StatusInternalServerError)
		return
	}

	m.timeSeries[tenant] = append(m.timeSeries[tenant], writeRequest.GetTimeseries()...)
}

func (m *mockPrometheusRemoteWriteServer) close() {
	m.mtx.Lock()
	defer m.mtx.Unlock()

	m.server.Close()
}

// poll executes f every interval until ctx is done or cancelled.
func poll(ctx context.Context, interval time.Duration, f func()) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			f()
		}
	}
}

var _ Overrides = (*mockOverrides)(nil)

type mockOverrides struct {
	headers          map[string]string
	nativeHistograms histograms.HistogramMethod
}

func (m *mockOverrides) MetricsGeneratorRemoteWriteHeaders(string) map[string]string {
	return m.headers
}

func (m *mockOverrides) MetricsGeneratorGenerateNativeHistograms(string) histograms.HistogramMethod {
	return m.nativeHistograms
}

var _ prometheus.Registerer = (*noopRegisterer)(nil)

type noopRegisterer struct{}

func (n *noopRegisterer) Register(prometheus.Collector) error { return nil }

func (n *noopRegisterer) MustRegister(...prometheus.Collector) {}

func (n *noopRegisterer) Unregister(prometheus.Collector) bool { return true }

func TestInstance_emptyTenantValidation(t *testing.T) {
	t.Parallel()

	o := &mockOverrides{}
	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))

	t.Run("empty tenant should return error", func(t *testing.T) {
		var cfg Config
		cfg.RegisterFlagsAndApplyDefaults("", nil)
		cfg.Path = t.TempDir()
		reg := prometheus.NewRegistry()

		// Empty tenant should return error when org ID header is required
		instance, err := New(&cfg, o, "", reg, logger)

		require.Error(t, err)
		require.Nil(t, instance)
		assert.Contains(t, err.Error(), "tenant cannot be empty")
	})

	t.Run("valid tenant should work normally", func(t *testing.T) {
		var cfg Config
		cfg.RegisterFlagsAndApplyDefaults("", nil)
		cfg.Path = t.TempDir()
		reg := prometheus.NewRegistry()

		instance, err := New(&cfg, o, "valid-tenant", reg, logger)

		require.NoError(t, err)
		require.NotNil(t, instance)

		// Verify correct WAL directory structure
		expectedPath := filepath.Join(cfg.Path, "valid-tenant")
		_, err = os.Stat(expectedPath)
		assert.NoError(t, err, "WAL directory should be created for valid tenant")

		instance.Close()
	})
}