// Copyright The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package agent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math"
	"path/filepath"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/model"
	"github.com/prometheus/common/promslog"
	"github.com/stretchr/testify/require"

	"github.com/prometheus/prometheus/model/exemplar"
	"github.com/prometheus/prometheus/model/histogram"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/tsdb"
	"github.com/prometheus/prometheus/tsdb/chunks"
	"github.com/prometheus/prometheus/tsdb/record"
	"github.com/prometheus/prometheus/tsdb/tsdbutil"
	"github.com/prometheus/prometheus/tsdb/wlog"
	"github.com/prometheus/prometheus/util/testutil"
)

func TestDB_InvalidSeries(t *testing.T) {
	s := createTestAgentDB(t, nil, DefaultOptions())
	defer s.Close()

	app := s.Appender(context.Background())

	t.Run("Samples", func(t *testing.T) {
		_, err := app.Append(0, labels.Labels{}, 0, 0)
		require.ErrorIs(t, err, tsdb.ErrInvalidSample, "should reject empty labels")

		_, err = app.Append(0, labels.FromStrings("a", "1", "a", "2"), 0, 0)
		require.ErrorIs(t, err, tsdb.ErrInvalidSample, "should reject duplicate labels")
	})

	t.Run("Histograms", func(t *testing.T) {
		_, err := app.AppendHistogram(0, labels.Labels{}, 0, tsdbutil.GenerateTestHistograms(1)[0], nil)
		require.ErrorIs(t, err, tsdb.ErrInvalidSample, "should reject empty labels")

		_, err = app.AppendHistogram(0, labels.FromStrings("a", "1", "a", "2"), 0, tsdbutil.GenerateTestHistograms(1)[0], nil)
		require.ErrorIs(t, err, tsdb.ErrInvalidSample, "should reject duplicate labels")
	})

	t.Run("Exemplars", func(t *testing.T) {
		sRef, err := app.Append(0, labels.FromStrings("a", "1"), 0, 0)
		require.NoError(t, err, "should not reject valid series")

		_, err = app.AppendExemplar(0, labels.EmptyLabels(), exemplar.Exemplar{})
		require.EqualError(t, err, "unknown series ref when trying to add exemplar: 0")

		e := exemplar.Exemplar{Labels: labels.FromStrings("a", "1", "a", "2")}
		_, err = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
		require.ErrorIs(t, err, tsdb.ErrInvalidExemplar, "should reject duplicate labels")

		e = exemplar.Exemplar{Labels: labels.FromStrings("a_somewhat_long_trace_id", "nYJSNtFrFTY37VR7mHzEE/LIDt7cdAQcuOzFajgmLDAdBSRHYPDzrxhMA4zz7el8naI/AoXFv9/e/G0vcETcIoNUi3OieeLfaIRQci2oa")}
		_, err = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
		require.ErrorIs(t, err, storage.ErrExemplarLabelLength, "should reject too long label length")

		// Inverse check
		e = exemplar.Exemplar{Labels: labels.FromStrings("a", "1"), Value: 20, Ts: 10, HasTs: true}
		_, err = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
		require.NoError(t, err, "should not reject valid exemplars")
	})
}

func createTestAgentDB(t testing.TB, reg prometheus.Registerer, opts *Options) *DB {
	t.Helper()

	dbDir := t.TempDir()

	rs := remote.NewStorage(promslog.NewNopLogger(), reg, startTime, dbDir, time.Second*30, nil, false)
	t.Cleanup(func() {
		require.NoError(t, rs.Close())
	})

	db, err := Open(promslog.NewNopLogger(), reg, rs, dbDir, opts)
	require.NoError(t, err)
	return db
}

// TestConcurrentAppendSameLabels verifies that concurrent appends for the same
// label set produce exactly one series in memory and one series record in the WAL.
func TestConcurrentAppendSameLabels(t *testing.T) {
	opts := DefaultOptions()
	opts.OutOfOrderTimeWindow = math.MaxInt64
	db := createTestAgentDB(t, nil, opts)
	lset := labels.FromStrings("__name__", "test_metric")

	const n = 100
	var wg sync.WaitGroup
	start := make(chan struct{})

	wg.Add(n)
	for range n {
		go func() {
			defer wg.Done()
			app := db.Appender(context.Background())
			<-start
			_, err := app.Append(0, lset, 1000, 1.0)
			require.NoError(t, err)
			require.NoError(t, app.Commit())
		}()
	}
	close(start)
	wg.Wait()

	var total int
	for i := range db.series.size {
		db.series.locks[i].RLock()
		total += len(db.series.series[i])
		db.series.locks[i].RUnlock()
	}
	require.Equal(t, 1, total)
	require.NoError(t, db.Close())

	sr, err := wlog.NewSegmentsReader(db.wal.Dir())
	require.NoError(t, err)
	defer func() { require.NoError(t, sr.Close()) }()

	r := wlog.NewReader(sr)
	dec := record.NewDecoder(labels.NewSymbolTable(), promslog.NewNopLogger())
	var walSeries int
	for r.Next() {
		rec := r.Record()
		if dec.Type(rec) == record.Series {
			var s []record.RefSeries
			s, err = dec.Series(rec, s)
			require.NoError(t, err)
			walSeries += len(s)
		}
	}
	require.NoError(t, r.Err())
	require.Equal(t, 1, walSeries)
}

func TestUnsupportedFunctions(t *testing.T) {
	s := createTestAgentDB(t, nil, DefaultOptions())
	defer s.Close()

	t.Run("Querier", func(t *testing.T) {
		_, err := s.Querier(0, 0)
		require.Equal(t, err, ErrUnsupported)
	})

	t.Run("ChunkQuerier", func(t *testing.T) {
		_, err := s.ChunkQuerier(0, 0)
		require.Equal(t, err, ErrUnsupported)
	})

	t.Run("ExemplarQuerier", func(t *testing.T) {
		_, err := s.ExemplarQuerier(context.TODO())
		require.Equal(t, err, ErrUnsupported)
	})
}

func TestCommit(t *testing.T) {
	const (
		numDatapoints = 1000
		numHistograms = 100
		numSeries     = 8
	)

	s := createTestAgentDB(t, nil, DefaultOptions())
	app := s.Appender(context.TODO())

	lbls := labelsForTest(t.Name(), numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for i := range numDatapoints {
			sample := chunks.GenerateSamples(0, 1)
			ref, err := app.Append(0, lset, sample[0].T(), sample[0].F())
			require.NoError(t, err)

			e := exemplar.Exemplar{
				Labels: lset,
				Ts:     sample[0].T() + int64(i),
				Value:  sample[0].F(),
				HasTs:  true,
			}
			_, err = app.AppendExemplar(ref, lset, e)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		customBucketHistograms := tsdbutil.GenerateTestCustomBucketsHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), customBucketHistograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		customBucketFloatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, customBucketFloatHistograms[i])
			require.NoError(t, err)
		}
	}

	require.NoError(t, app.Commit())
	require.NoError(t, s.Close())

	sr, err := wlog.NewSegmentsReader(s.wal.Dir())
	require.NoError(t, err)
	defer func() {
		require.NoError(t, sr.Close())
	}()

	// Read records from WAL and check for expected count of series, samples, and exemplars.
	var (
		r   = wlog.NewReader(sr)
		dec = record.NewDecoder(labels.NewSymbolTable(), promslog.NewNopLogger())

		walSeriesCount, walSamplesCount, walExemplarsCount, walHistogramCount, walFloatHistogramCount int
	)
	for r.Next() {
		rec := r.Record()
		switch dec.Type(rec) {
		case record.Series:
			var series []record.RefSeries
			series, err = dec.Series(rec, series)
			require.NoError(t, err)
			walSeriesCount += len(series)

		case record.Samples, record.SamplesV2:
			var samples []record.RefSample
			samples, err = dec.Samples(rec, samples)
			require.NoError(t, err)
			walSamplesCount += len(samples)

		case record.HistogramSamples, record.CustomBucketsHistogramSamples, record.HistogramSamplesV2:
			var histograms []record.RefHistogramSample
			histograms, err = dec.HistogramSamples(rec, histograms)
			require.NoError(t, err)
			walHistogramCount += len(histograms)

		case record.FloatHistogramSamples, record.CustomBucketsFloatHistogramSamples, record.FloatHistogramSamplesV2:
			var floatHistograms []record.RefFloatHistogramSample
			floatHistograms, err = dec.FloatHistogramSamples(rec, floatHistograms)
			require.NoError(t, err)
			walFloatHistogramCount += len(floatHistograms)

		case record.Exemplars:
			var exemplars []record.RefExemplar
			exemplars, err = dec.Exemplars(rec, exemplars)
			require.NoError(t, err)
			walExemplarsCount += len(exemplars)

		default:
		}
	}

	// Check that the WAL contained the same number of committed series/samples/exemplars.
	require.Equal(t, numSeries*5, walSeriesCount, "unexpected number of series")
	require.Equal(t, numSeries*numDatapoints, walSamplesCount, "unexpected number of samples")
	require.Equal(t, numSeries*numDatapoints, walExemplarsCount, "unexpected number of exemplars")
	require.Equal(t, numSeries*numHistograms*2, walHistogramCount, "unexpected number of histograms")
	require.Equal(t, numSeries*numHistograms*2, walFloatHistogramCount, "unexpected number of float histograms")

	// Check that we can get another appender after this - see https://github.com/prometheus/prometheus/issues/17800.
	_ = s.Appender(context.TODO())
}

func TestRollback(t *testing.T) {
	const (
		numDatapoints = 1000
		numHistograms = 100
		numSeries     = 8
	)

	s := createTestAgentDB(t, nil, DefaultOptions())
	app := s.Appender(context.TODO())

	lbls := labelsForTest(t.Name(), numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for range numDatapoints {
			sample := chunks.GenerateSamples(0, 1)
			_, err := app.Append(0, lset, sample[0].T(), sample[0].F())
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	// Do a rollback, which should clear uncommitted data. A followup call to
	// commit should persist nothing to the WAL.
	require.NoError(t, app.Rollback())
	require.NoError(t, app.Commit())
	require.NoError(t, s.Close())

	sr, err := wlog.NewSegmentsReader(s.wal.Dir())
	require.NoError(t, err)
	defer func() {
		require.NoError(t, sr.Close())
	}()

	// Read records from WAL and check for expected count of series and samples.
	var (
		r   = wlog.NewReader(sr)
		dec = record.NewDecoder(labels.NewSymbolTable(), promslog.NewNopLogger())

		walSeriesCount, walSamplesCount, walHistogramCount, walFloatHistogramCount, walExemplarsCount int
	)
	for r.Next() {
		rec := r.Record()
		switch dec.Type(rec) {
		case record.Series:
			var series []record.RefSeries
			series, err = dec.Series(rec, series)
			require.NoError(t, err)
			walSeriesCount += len(series)

		case record.Samples, record.SamplesV2:
			var samples []record.RefSample
			samples, err = dec.Samples(rec, samples)
			require.NoError(t, err)
			walSamplesCount += len(samples)

		case record.Exemplars:
			var exemplars []record.RefExemplar
			exemplars, err = dec.Exemplars(rec, exemplars)
			require.NoError(t, err)
			walExemplarsCount += len(exemplars)

		case record.HistogramSamples, record.CustomBucketsHistogramSamples, record.HistogramSamplesV2:
			var histograms []record.RefHistogramSample
			histograms, err = dec.HistogramSamples(rec, histograms)
			require.NoError(t, err)
			walHistogramCount += len(histograms)

		case record.FloatHistogramSamples, record.CustomBucketsFloatHistogramSamples, record.FloatHistogramSamplesV2:
			var floatHistograms []record.RefFloatHistogramSample
			floatHistograms, err = dec.FloatHistogramSamples(rec, floatHistograms)
			require.NoError(t, err)
			walFloatHistogramCount += len(floatHistograms)

		default:
		}
	}

	// Check that only series get stored after calling Rollback.
	require.Equal(t, numSeries*5, walSeriesCount, "series should have been written to WAL")
	require.Equal(t, 0, walSamplesCount, "samples should not have been written to WAL")
	require.Equal(t, 0, walExemplarsCount, "exemplars should not have been written to WAL")
	require.Equal(t, 0, walHistogramCount, "histograms should not have been written to WAL")
	require.Equal(t, 0, walFloatHistogramCount, "float histograms should not have been written to WAL")
}

func TestFullTruncateWAL(t *testing.T) {
	const (
		numDatapoints = 1000
		numHistograms = 100
		numSeries     = 800
		lastTs        = 500
	)

	reg := prometheus.NewRegistry()
	opts := DefaultOptions()
	opts.TruncateFrequency = time.Minute * 2

	s := createTestAgentDB(t, reg, opts)
	defer func() {
		require.NoError(t, s.Close())
	}()
	app := s.Appender(context.TODO())

	lbls := labelsForTest(t.Name(), numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for range numDatapoints {
			_, err := app.Append(0, lset, int64(lastTs), 0)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(lastTs), histograms[i], nil)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(lastTs), histograms[i], nil)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(lastTs), nil, floatHistograms[i])
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, int64(lastTs), nil, floatHistograms[i])
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	// Truncate WAL with mint to GC all the samples.
	s.truncate(lastTs + 1)

	m := gatherFamily(t, reg, "prometheus_agent_deleted_series")
	require.Equal(t, float64(numSeries*5), m.Metric[0].Gauge.GetValue(), "agent wal truncate mismatch of deleted series count")
}

func TestPartialTruncateWAL(t *testing.T) {
	const (
		numDatapoints = 1000
		numSeries     = 800
	)

	opts := DefaultOptions()

	reg := prometheus.NewRegistry()
	s := createTestAgentDB(t, reg, opts)
	defer func() {
		require.NoError(t, s.Close())
	}()
	app := s.Appender(context.TODO())

	// Create first batch of 800 series with 1000 data-points with a fixed lastTs as 500.
	var lastTs int64 = 500
	lbls := labelsForTest(t.Name()+"batch-1", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for range numDatapoints {
			_, err := app.Append(0, lset, lastTs, 0)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_histogram_batch-1", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram_batch-1", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_float_histogram_batch-1", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram_batch-1", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	// Create second batch of 800 series with 1000 data-points with a fixed lastTs as 600.
	lastTs = 600
	lbls = labelsForTest(t.Name()+"batch-2", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for range numDatapoints {
			_, err := app.Append(0, lset, lastTs, 0)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_histogram_batch-2", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram_batch-2", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_float_histogram_batch-2", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram_batch-2", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numDatapoints)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
			require.NoError(t, err)
		}
		require.NoError(t, app.Commit())
	}

	// Truncate WAL with mint to GC only the first batch of 800 series and retaining 2nd batch of 800 series.
	s.truncate(lastTs - 1)

	m := gatherFamily(t, reg, "prometheus_agent_deleted_series")
	require.Len(t, m.Metric, 1)
	require.Equal(t, float64(numSeries*5), m.Metric[0].Gauge.GetValue(), "agent wal truncate mismatch of deleted series count")
}

func TestWALReplay(t *testing.T) {
	const (
		numDatapoints = 1000
		numHistograms = 100
		numSeries     = 8
		lastTs        = 500
	)

	s := createTestAgentDB(t, nil, DefaultOptions())
	app := s.Appender(context.TODO())

	lbls := labelsForTest(t.Name(), numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for range numDatapoints {
			_, err := app.Append(0, lset, lastTs, 0)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numHistograms)

		for i := range numHistograms {
			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	require.NoError(t, app.Commit())
	require.NoError(t, s.Close())

	// Hack: s.wal.Dir() is the /wal subdirectory of the original storage path.
	// We need the original directory so we can recreate the storage for replay.
	storageDir := filepath.Dir(s.wal.Dir())

	reg := prometheus.NewRegistry()
	replayStorage, err := Open(s.logger, reg, nil, storageDir, s.opts)
	if err != nil {
		t.Fatalf("unable to create storage for the agent: %v", err)
	}
	defer func() {
		require.NoError(t, replayStorage.Close())
	}()

	// Check if all the series are retrieved back from the WAL.
	m := gatherFamily(t, reg, "prometheus_agent_active_series")
	require.Equal(t, float64(numSeries*5), m.Metric[0].Gauge.GetValue(), "agent wal replay mismatch of active series count")

	// Check if lastTs of the samples retrieved from the WAL is retained.
	metrics := replayStorage.series.series
	for i := range metrics {
		mp := metrics[i]
		for _, v := range mp {
			require.Equal(t, v.lastTs, int64(lastTs))
		}
	}
}

func TestLockfile(t *testing.T) {
	tsdbutil.TestDirLockerUsage(t, func(t *testing.T, data string, createLock bool) (*tsdbutil.DirLocker, testutil.Closer) {
		logger := promslog.NewNopLogger()
		reg := prometheus.NewRegistry()
		rs := remote.NewStorage(logger, reg, startTime, data, time.Second*30, nil, false)
		t.Cleanup(func() {
			require.NoError(t, rs.Close())
		})

		opts := DefaultOptions()
		opts.NoLockfile = !createLock

		// Create the DB. This should create lockfile and its metrics.
		db, err := Open(logger, nil, rs, data, opts)
		require.NoError(t, err)

		return db.locker, testutil.NewCallbackCloser(func() {
			require.NoError(t, db.Close())
		})
	})
}

func Test_ExistingWAL_NextRef(t *testing.T) {
	dbDir := t.TempDir()
	rs := remote.NewStorage(promslog.NewNopLogger(), nil, startTime, dbDir, time.Second*30, nil, false)
	defer func() {
		require.NoError(t, rs.Close())
	}()

	db, err := Open(promslog.NewNopLogger(), nil, rs, dbDir, DefaultOptions())
	require.NoError(t, err)

	seriesCount := 10

	// Append <seriesCount> series
	app := db.Appender(context.Background())
	for i := range seriesCount {
		lset := labels.FromStrings(model.MetricNameLabel, fmt.Sprintf("series_%d", i))
		_, err := app.Append(0, lset, 0, 100)
		require.NoError(t, err)
	}

	histogramCount := 10
	histograms := tsdbutil.GenerateTestHistograms(histogramCount)
	// Append <histogramCount> series
	for i := range histogramCount {
		lset := labels.FromStrings(model.MetricNameLabel, fmt.Sprintf("histogram_%d", i))
		_, err := app.AppendHistogram(0, lset, 0, histograms[i], nil)
		require.NoError(t, err)
	}
	require.NoError(t, app.Commit())

	// Truncate the WAL to force creation of a new segment.
	require.NoError(t, db.truncate(0))
	require.NoError(t, db.Close())

	// Create a new storage and see what nextRef is initialized to.
	db, err = Open(promslog.NewNopLogger(), nil, rs, dbDir, DefaultOptions())
	require.NoError(t, err)
	defer func() {
		require.NoError(t, db.Close())
	}()

	require.Equal(t, uint64(seriesCount+histogramCount), db.nextRef.Load(), "nextRef should be equal to the number of series written across the entire WAL")
}

func Test_validateOptions(t *testing.T) {
	t.Run("Apply defaults to zero values", func(t *testing.T) {
		opts := validateOptions(&Options{})
		require.Equal(t, DefaultOptions(), opts)
	})

	t.Run("Defaults are already valid", func(t *testing.T) {
		require.Equal(t, DefaultOptions(), validateOptions(nil))
	})

	t.Run("MaxWALTime should not be lower than TruncateFrequency", func(t *testing.T) {
		opts := validateOptions(&Options{
			MaxWALTime:        int64(time.Hour / time.Millisecond),
			TruncateFrequency: 2 * time.Hour,
		})
		require.Equal(t, int64(2*time.Hour/time.Millisecond), opts.MaxWALTime)
	})
}

func startTime() (int64, error) {
	return time.Now().Unix() * 1000, nil
}

// Create series for tests.
func labelsForTest(lName string, seriesCount int) [][]labels.Label {
	var series [][]labels.Label

	for i := range seriesCount {
		lset := []labels.Label{
			{Name: "a", Value: lName},
			{Name: "instance", Value: "localhost" + strconv.Itoa(i)},
			{Name: "job", Value: "prometheus"},
		}
		series = append(series, lset)
	}

	return series
}

func gatherFamily(t *testing.T, reg prometheus.Gatherer, familyName string) *dto.MetricFamily {
	t.Helper()

	families, err := reg.Gather()
	require.NoError(t, err, "failed to gather metrics")

	for _, f := range families {
		if f.GetName() == familyName {
			return f
		}
	}

	t.Fatalf("could not find family %s", familyName)

	return nil
}

func TestStorage_DuplicateExemplarsIgnored(t *testing.T) {
	s := createTestAgentDB(t, nil, DefaultOptions())
	app := s.Appender(context.Background())
	defer s.Close()

	sRef, err := app.Append(0, labels.FromStrings("a", "1"), 0, 0)
	require.NoError(t, err, "should not reject valid series")

	// Write a few exemplars to our appender and call Commit().
	// If the Labels, Value or Timestamp are different than the last exemplar,
	// then a new one should be appended; Otherwise, it should be skipped.
	e := exemplar.Exemplar{Labels: labels.FromStrings("a", "1"), Value: 20, Ts: 10, HasTs: true}
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)

	e.Labels = labels.FromStrings("b", "2")
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)

	e.Value = 42
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)

	e.Ts = 25
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)
	_, _ = app.AppendExemplar(sRef, labels.EmptyLabels(), e)

	require.NoError(t, app.Commit())

	// Read back what was written to the WAL.
	var walExemplarsCount int
	sr, err := wlog.NewSegmentsReader(s.wal.Dir())
	require.NoError(t, err)
	defer sr.Close()
	r := wlog.NewReader(sr)

	dec := record.NewDecoder(labels.NewSymbolTable(), promslog.NewNopLogger())
	for r.Next() {
		rec := r.Record()
		if dec.Type(rec) == record.Exemplars {
			var exemplars []record.RefExemplar
			exemplars, err = dec.Exemplars(rec, exemplars)
			require.NoError(t, err)
			walExemplarsCount += len(exemplars)
		}
	}

	// We had 9 calls to AppendExemplar but only 4 of those should have gotten through.
	require.Equal(t, 4, walExemplarsCount)
}

func TestDBAllowOOOSamples(t *testing.T) {
	const (
		numDatapoints = 5
		numHistograms = 5
		numSeries     = 4
		offset        = 100
	)

	reg := prometheus.NewRegistry()
	opts := DefaultOptions()
	opts.OutOfOrderTimeWindow = math.MaxInt64
	s := createTestAgentDB(t, reg, opts)
	app := s.Appender(context.TODO())

	// Let's add some samples in the [offset, offset+numDatapoints) range.
	lbls := labelsForTest(t.Name(), numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		for i := offset; i < numDatapoints+offset; i++ {
			ref, err := app.Append(0, lset, int64(i), float64(i))
			require.NoError(t, err)

			e := exemplar.Exemplar{
				Labels: lset,
				Ts:     int64(i) * 2,
				Value:  float64(i),
				HasTs:  true,
			}
			_, err = app.AppendExemplar(ref, lset, e)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numHistograms)

		for i := offset; i < numDatapoints+offset; i++ {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i-offset], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numHistograms)

		for i := offset; i < numDatapoints+offset; i++ {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i-offset], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)

		for i := offset; i < numDatapoints+offset; i++ {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i-offset])
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram", numSeries)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numHistograms)

		for i := offset; i < numDatapoints+offset; i++ {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i-offset])
			require.NoError(t, err)
		}
	}

	require.NoError(t, app.Commit())
	m := gatherFamily(t, reg, "prometheus_agent_samples_appended_total")
	require.Equal(t, float64(20), m.Metric[0].Counter.GetValue(), "agent wal mismatch of total appended samples")
	require.Equal(t, float64(80), m.Metric[1].Counter.GetValue(), "agent wal mismatch of total appended histograms")
	require.NoError(t, s.Close())

	// Hack: s.wal.Dir() is the /wal subdirectory of the original storage path.
	// We need the original directory so we can recreate the storage for replay.
	storageDir := filepath.Dir(s.wal.Dir())

	// Replay the storage so that the lastTs for each series is recorded.
	reg2 := prometheus.NewRegistry()
	db, err := Open(s.logger, reg2, nil, storageDir, s.opts)
	if err != nil {
		t.Fatalf("unable to create storage for the agent: %v", err)
	}

	app = db.Appender(context.Background())

	// Now the lastTs will have been recorded successfully.
	// Let's try appending twice as many OOO samples in the [0, numDatapoints) range.
	lbls = labelsForTest(t.Name()+"_histogram", numSeries*2)
	for _, l := range lbls {
		lset := labels.New(l...)

		for i := range numDatapoints {
			ref, err := app.Append(0, lset, int64(i), float64(i))
			require.NoError(t, err)

			e := exemplar.Exemplar{
				Labels: lset,
				Ts:     int64(i) * 2,
				Value:  float64(i),
				HasTs:  true,
			}
			_, err = app.AppendExemplar(ref, lset, e)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_histogram", numSeries*2)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestHistograms(numHistograms)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_histogram", numSeries*2)
	for _, l := range lbls {
		lset := labels.New(l...)

		histograms := tsdbutil.GenerateTestCustomBucketsHistograms(numHistograms)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_float_histogram", numSeries*2)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	lbls = labelsForTest(t.Name()+"_custom_buckets_float_histogram", numSeries*2)
	for _, l := range lbls {
		lset := labels.New(l...)

		floatHistograms := tsdbutil.GenerateTestCustomBucketsFloatHistograms(numHistograms)

		for i := range numDatapoints {
			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
			require.NoError(t, err)
		}
	}

	require.NoError(t, app.Commit())
	m = gatherFamily(t, reg2, "prometheus_agent_samples_appended_total")
	require.Equal(t, float64(40), m.Metric[0].Counter.GetValue(), "agent wal mismatch of total appended samples")
	require.Equal(t, float64(160), m.Metric[1].Counter.GetValue(), "agent wal mismatch of total appended histograms")
	require.NoError(t, db.Close())
}

func TestDBOutOfOrderTimeWindow(t *testing.T) {
	tc := []struct {
		outOfOrderTimeWindow, firstTs, secondTs int64
		expectedError                           error
	}{
		{0, 100, 101, nil},
		{0, 100, 100, storage.ErrOutOfOrderSample},
		{0, 100, 99, storage.ErrOutOfOrderSample},
		{100, 100, 1, nil},
		{100, 100, 0, storage.ErrOutOfOrderSample},
	}

	for _, c := range tc {
		t.Run(fmt.Sprintf("outOfOrderTimeWindow=%d, firstTs=%d, secondTs=%d, expectedError=%s", c.outOfOrderTimeWindow, c.firstTs, c.secondTs, c.expectedError), func(t *testing.T) {
			reg := prometheus.NewRegistry()
			opts := DefaultOptions()
			opts.OutOfOrderTimeWindow = c.outOfOrderTimeWindow
			s := createTestAgentDB(t, reg, opts)
			app := s.Appender(context.TODO())

			lbls := labelsForTest(t.Name()+"_histogram", 1)
			lset := labels.New(lbls[0]...)
			_, err := app.AppendHistogram(0, lset, c.firstTs, tsdbutil.GenerateTestHistograms(1)[0], nil)
			require.NoError(t, err)
			err = app.Commit()
			require.NoError(t, err)
			_, err = app.AppendHistogram(0, lset, c.secondTs, tsdbutil.GenerateTestHistograms(1)[0], nil)
			require.ErrorIs(t, err, c.expectedError)

			lbls = labelsForTest(t.Name(), 1)
			lset = labels.New(lbls[0]...)
			_, err = app.Append(0, lset, c.firstTs, 0)
			require.NoError(t, err)
			err = app.Commit()
			require.NoError(t, err)
			_, err = app.Append(0, lset, c.secondTs, 0)
			require.ErrorIs(t, err, c.expectedError)

			expectedAppendedSamples := float64(2)
			if c.expectedError != nil {
				expectedAppendedSamples = 1
			}
			m := gatherFamily(t, reg, "prometheus_agent_samples_appended_total")
			require.Equal(t, expectedAppendedSamples, m.Metric[0].Counter.GetValue(), "agent wal mismatch of total appended samples")
			require.Equal(t, expectedAppendedSamples, m.Metric[1].Counter.GetValue(), "agent wal mismatch of total appended histograms")
			require.NoError(t, s.Close())
		})
	}
}

type walSample struct {
	t    int64
	f    float64
	h    *histogram.Histogram
	lbls labels.Labels
	ref  storage.SeriesRef
}

// NOTE(bwplotka): This test is testing behaviour of storage.Appender interface against its invariants (see
// storage.Appender comment) around validation of the order of samples within a single Appender. This results
// in a slight bug in AppendSTZero* methods. We are leaving it as-is given the planned removal of AppenderV1 as
// per https://github.com/prometheus/prometheus/issues/17632.
func TestDBStartTimestampSamplesIngestion(t *testing.T) {
	t.Parallel()

	type appendableSample struct {
		t            int64
		st           int64
		v            float64
		lbls         labels.Labels
		h            *histogram.Histogram
		expectsError bool
	}

	testHistograms := tsdbutil.GenerateTestHistograms(2)
	zeroHistogram := &histogram.Histogram{}

	lbls := labelsForTest(t.Name(), 1)
	defLbls := labels.New(lbls[0]...)

	testCases := []struct {
		name                string
		inputSamples        []appendableSample
		expectedSamples     []walSample
		expectedSeriesCount int
	}{
		{
			name: "in order ct+normal sample/floatSamples",
			inputSamples: []appendableSample{
				{t: 100, st: 1, v: 10, lbls: defLbls},
				{t: 101, st: 1, v: 10, lbls: defLbls},
			},
			expectedSamples: []walSample{
				{t: 1, f: 0, lbls: defLbls, ref: 1},
				{t: 100, f: 10, lbls: defLbls, ref: 1},
				{t: 101, f: 10, lbls: defLbls, ref: 1},
			},
		},
		{
			name: "ST+float && ST+histogram samples",
			inputSamples: []appendableSample{
				{
					t:    100,
					st:   30,
					v:    20,
					lbls: defLbls,
				},
				{
					t:    300,
					st:   230,
					h:    testHistograms[0],
					lbls: defLbls,
				},
			},
			expectedSamples: []walSample{
				{t: 30, f: 0, lbls: defLbls, ref: 1},
				{t: 100, f: 20, lbls: defLbls, ref: 1},
				{t: 230, h: zeroHistogram, lbls: defLbls, ref: 1},
				{t: 300, h: testHistograms[0], lbls: defLbls, ref: 1},
			},
			expectedSeriesCount: 1,
		},
		{
			name: "ST+float && ST+histogram samples with error",
			inputSamples: []appendableSample{
				{
					// invalid ST
					t:            100,
					st:           100,
					v:            10,
					lbls:         defLbls,
					expectsError: true,
				},
				{
					// invalid ST histogram
					t:            300,
					st:           300,
					h:            testHistograms[0],
					lbls:         defLbls,
					expectsError: true,
				},
			},
			expectedSamples: []walSample{
				{t: 100, f: 10, lbls: defLbls, ref: 1},
				{t: 300, h: testHistograms[0], lbls: defLbls, ref: 1},
			},
			expectedSeriesCount: 0,
		},
		{
			name: "In order ct+normal sample/histogram",
			inputSamples: []appendableSample{
				{t: 100, h: testHistograms[0], st: 1, lbls: defLbls},
				{t: 101, h: testHistograms[1], st: 1, lbls: defLbls},
			},
			expectedSamples: []walSample{
				{t: 1, h: &histogram.Histogram{}, lbls: defLbls, ref: 1},
				{t: 100, h: testHistograms[0], lbls: defLbls, ref: 1},
				{t: 101, h: testHistograms[1], lbls: defLbls, ref: 1},
			},
		},
		{
			name: "ct+normal then OOO sample/float",
			inputSamples: []appendableSample{
				{t: 60_000, st: 40_000, v: 10, lbls: defLbls},
				{t: 120_000, st: 40_000, v: 10, lbls: defLbls},
				{t: 180_000, st: 40_000, v: 10, lbls: defLbls},
				{t: 50_000, st: 40_000, v: 10, lbls: defLbls},
			},
			expectedSamples: []walSample{
				{t: 40_000, f: 0, lbls: defLbls, ref: 1},
				{t: 60_000, f: 10, lbls: defLbls, ref: 1},
				{t: 120_000, f: 10, lbls: defLbls, ref: 1},
				{t: 180_000, f: 10, lbls: defLbls, ref: 1},
				{t: 50_000, f: 10, lbls: defLbls, ref: 1}, // OOO sample.
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			reg := prometheus.NewRegistry()
			opts := DefaultOptions()
			opts.OutOfOrderTimeWindow = 360_000
			s := createTestAgentDB(t, reg, opts)
			app := s.Appender(context.TODO())

			for _, sample := range tc.inputSamples {
				// We supposed to write a Histogram to the WAL
				if sample.h != nil {
					_, err := app.AppendHistogramSTZeroSample(0, sample.lbls, sample.t, sample.st, zeroHistogram, nil)
					if !errors.Is(err, storage.ErrOutOfOrderST) {
						require.Equal(t, sample.expectsError, err != nil, "expected error: %v, got: %v", sample.expectsError, err)
					}

					_, err = app.AppendHistogram(0, sample.lbls, sample.t, sample.h, nil)
					require.NoError(t, err)
				} else {
					// We supposed to write a float sample to the WAL
					_, err := app.AppendSTZeroSample(0, sample.lbls, sample.t, sample.st)
					if !errors.Is(err, storage.ErrOutOfOrderST) {
						require.Equal(t, sample.expectsError, err != nil, "expected error: %v, got: %v", sample.expectsError, err)
					}

					_, err = app.Append(0, sample.lbls, sample.t, sample.v)
					require.NoError(t, err)
				}
			}

			require.NoError(t, app.Commit())
			// Close the DB to ensure all data is flushed to the WAL
			require.NoError(t, s.Close())

			// Check that we don't have any OOO samples in the WAL by checking metrics
			families, err := reg.Gather()
			require.NoError(t, err, "failed to gather metrics")
			for _, f := range families {
				if f.GetName() == "prometheus_agent_out_of_order_samples_total" {
					t.Fatalf("unexpected metric %s", f.GetName())
				}
			}

			got := readWALSamples(t, s.wal.Dir())
			testutil.RequireEqualWithOptions(t, tc.expectedSamples, got, cmp.Options{cmp.AllowUnexported(walSample{})})
		})
	}
}

func TestDuplicateSeriesRefsByHash(t *testing.T) {
	dbDir := t.TempDir()
	opts := DefaultOptions()
	rs1 := remote.NewStorage(promslog.NewNopLogger(), nil, startTime, dbDir, time.Second*30, nil, false)
	db, err := Open(promslog.NewNopLogger(), nil, rs1, dbDir, opts)
	require.NoError(t, err)

	app := db.Appender(context.Background())

	metricNames := []string{"foo", "bar", "baz", "blerg"}
	originalSeriesRefs := make([]chunks.HeadSeriesRef, 0, len(metricNames))
	for _, metricName := range metricNames {
		lbls := labels.FromMap(map[string]string{"__name__": metricName})

		ref, err := app.Append(storage.SeriesRef(0), lbls, int64(0), 10.0)
		require.NoError(t, err)
		originalSeriesRefs = append(originalSeriesRefs, chunks.HeadSeriesRef(ref))
		ref2, err := app.Append(ref, lbls, int64(10), 100.0)
		require.NoError(t, err)
		require.Equal(t, ref, ref2)
	}
	require.NoError(t, app.Commit())

	// Forcefully create a bunch of new segments to force a truncation.
	for range 3 {
		_, err := db.wal.NextSegmentSync()
		require.NoError(t, err)
	}
	// No series should be deleted yet
	require.Empty(t, db.deleted)

	// Truncate at 1 ms higher than the highest timestamp.
	err = db.truncate(11)
	require.NoError(t, err)

	// The original SeriesRefs should be considered deleted.
	for _, ref := range originalSeriesRefs {
		require.Nil(t, db.series.GetByID(ref))
		require.Contains(t, db.deleted, ref)
	}

	duplicateSeriesRefs := make([]chunks.HeadSeriesRef, 0, len(metricNames))
	for _, metricName := range metricNames {
		lbls := labels.FromMap(map[string]string{"__name__": metricName})
		ref, err := app.Append(storage.SeriesRef(0), lbls, int64(20), 10.0)
		require.NoError(t, err)
		duplicateSeriesRefs = append(duplicateSeriesRefs, chunks.HeadSeriesRef(ref))
	}
	require.NoError(t, app.Commit())

	// The duplicate SeriesRefs should be in series.
	for _, ref := range duplicateSeriesRefs {
		require.NotNil(t, db.series.GetByID(ref))
	}

	// Close the WAL before we have a chance to remove the original RefIDs.
	// Both db and rs1 must be closed to release all file handles before
	// reopening the same directory — important on Windows.
	require.NoError(t, db.Close())
	require.NoError(t, rs1.Close())

	rs2 := remote.NewStorage(promslog.NewNopLogger(), nil, startTime, dbDir, time.Second*30, nil, false)
	t.Cleanup(func() { require.NoError(t, rs2.Close()) })
	db, err = Open(promslog.NewNopLogger(), nil, rs2, dbDir, opts)
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, db.Close()) })

	// The original SeriesRefs should be in series.
	for _, ref := range originalSeriesRefs {
		require.NotNil(t, db.series.GetByID(ref))
		require.NotContains(t, db.deleted, ref)
	}

	// The duplicated SeriesRefs should be considered deleted.
	for _, ref := range duplicateSeriesRefs {
		require.Nil(t, db.series.GetByID(ref))
		require.Contains(t, db.deleted, ref)
	}
}

func readWALSamples(t *testing.T, walDir string) []walSample {
	t.Helper()
	sr, err := wlog.NewSegmentsReader(walDir)
	require.NoError(t, err)
	defer func(sr io.ReadCloser) {
		err := sr.Close()
		require.NoError(t, err)
	}(sr)

	r := wlog.NewReader(sr)
	dec := record.NewDecoder(labels.NewSymbolTable(), promslog.NewNopLogger())

	var (
		samples    []record.RefSample
		histograms []record.RefHistogramSample

		lastSeries    record.RefSeries
		outputSamples = make([]walSample, 0)
	)

	for r.Next() {
		rec := r.Record()
		switch dec.Type(rec) {
		case record.Series:
			series, err := dec.Series(rec, nil)
			require.NoError(t, err)
			lastSeries = series[0]
		case record.Samples, record.SamplesV2:
			samples, err = dec.Samples(rec, samples[:0])
			require.NoError(t, err)
			for _, s := range samples {
				outputSamples = append(outputSamples, walSample{
					t:    s.T,
					f:    s.V,
					lbls: lastSeries.Labels.Copy(),
					ref:  storage.SeriesRef(lastSeries.Ref),
				})
			}
		case record.HistogramSamples:
			histograms, err = dec.HistogramSamples(rec, histograms[:0])
			require.NoError(t, err)
			for _, h := range histograms {
				outputSamples = append(outputSamples, walSample{
					t:    h.T,
					h:    h.H,
					lbls: lastSeries.Labels.Copy(),
					ref:  storage.SeriesRef(lastSeries.Ref),
				})
			}
		}
	}
	return outputSamples
}

func BenchmarkGetOrCreate(b *testing.B) {
	// NOTE: This benchmarks appenderBase, so it does not matter if it's V1 or V2.
	const n = 1_000

	b.Run("new", func(b *testing.B) {
		s := createTestAgentDB(b, nil, DefaultOptions())
		defer s.Close()
		app := s.Appender(context.Background()).(*appender)

		// Fixed-size label set. Before each pass through the set we GC all series
		// (they are created with lastTs==math.MinInt64, so mint=math.MaxInt64
		// evicts everything) so every timed getOrCreate call takes the creation
		// path. This keeps the stripe-series table at a stable size regardless of
		// b.N, preventing per-op cost from growing with the benchmark iteration
		// count.
		lbls := make([]labels.Labels, n)
		for i, l := range labelsForTest("benchmark_new", n) {
			lbls[i] = labels.New(l...)
		}

		b.ResetTimer()

		for i := range b.N {
			if i%n == 0 && i > 0 {
				b.StopTimer()
				_ = s.series.GC(math.MaxInt64, false)
				b.StartTimer()
			}
			app.getOrCreate(0, lbls[i%n])
		}
	})

	b.Run("existing", func(b *testing.B) {
		s := createTestAgentDB(b, nil, DefaultOptions())
		defer s.Close()
		app := s.Appender(context.Background()).(*appender)

		lbls := make([]labels.Labels, n)
		for i, l := range labelsForTest("benchmark_existing", n) {
			lbls[i] = labels.New(l...)
		}

		// Pre-populate all series so every timed call finds an existing series.
		for _, l := range lbls {
			app.getOrCreate(0, l)
		}

		b.ResetTimer()

		for i := range b.N {
			app.getOrCreate(0, lbls[i%n])
		}
	})
}
