From 286dfc70b7c7d22a17068b5038bc13f02bf87d70 Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Thu, 14 Apr 2022 04:27:06 -0400
Subject: [PATCH] tsdb/agent: port grafana/agent#676 (#10587)

* tsdb/agent: port grafana/agent#676

grafana/agent#676 fixed an issue where a loading a WAL with multiple
segments may result in ref ID collision. The starting ref ID for new
series should be set to the highest ref ID across all series records
from all WAL segments.

This fixes an issue where the starting ref ID was incorrectly set to the
highest ref ID found in the newest segment, which may not have any ref
IDs at all if no series records have been appended to it yet.

Signed-off-by: Robert Fratto <robertfratto@gmail.com>

* tsdb/agent: update terminology (s/ref ID/nextRef)

Signed-off-by: Robert Fratto <robertfratto@gmail.com>
---
 tsdb/agent/db.go      |  2 +-
 tsdb/agent/db_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/tsdb/agent/db.go b/tsdb/agent/db.go
index 7d42d14feb..a06d8f1f90 100644
--- a/tsdb/agent/db.go
+++ b/tsdb/agent/db.go
@@ -392,7 +392,7 @@ func (db *DB) replayWAL() error {
 func (db *DB) loadWAL(r *wal.Reader, multiRef map[chunks.HeadSeriesRef]chunks.HeadSeriesRef) (err error) {
 	var (
 		dec     record.Decoder
-		lastRef chunks.HeadSeriesRef
+		lastRef = chunks.HeadSeriesRef(db.nextRef.Load())
 
 		decoded    = make(chan interface{}, 10)
 		errCh      = make(chan error, 1)
diff --git a/tsdb/agent/db_test.go b/tsdb/agent/db_test.go
index e99712f71f..a44b6c7286 100644
--- a/tsdb/agent/db_test.go
+++ b/tsdb/agent/db_test.go
@@ -15,6 +15,7 @@ package agent
 
 import (
 	"context"
+	"fmt"
 	"path/filepath"
 	"strconv"
 	"testing"
@@ -24,6 +25,7 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
+	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/require"
 
 	"github.com/prometheus/prometheus/model/exemplar"
@@ -409,6 +411,41 @@ func TestLockfile(t *testing.T) {
 	})
 }
 
+func Test_ExistingWAL_NextRef(t *testing.T) {
+	dbDir := t.TempDir()
+	rs := remote.NewStorage(log.NewNopLogger(), nil, startTime, dbDir, time.Second*30, nil)
+	defer func() {
+		require.NoError(t, rs.Close())
+	}()
+
+	db, err := Open(log.NewNopLogger(), nil, rs, dbDir, DefaultOptions())
+	require.NoError(t, err)
+
+	seriesCount := 10
+
+	// Append <seriesCount> series
+	app := db.Appender(context.Background())
+	for i := 0; i < seriesCount; i++ {
+		lset := labels.Labels{
+			{Name: model.MetricNameLabel, Value: fmt.Sprintf("series_%d", i)},
+		}
+		_, err := app.Append(0, lset, 0, 100)
+		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(log.NewNopLogger(), nil, rs, dbDir, DefaultOptions())
+	require.NoError(t, err)
+	defer require.NoError(t, db.Close())
+
+	require.Equal(t, uint64(seriesCount), db.nextRef.Load(), "nextRef should be equal to the number of series written across the entire WAL")
+}
+
 func startTime() (int64, error) {
 	return time.Now().Unix() * 1000, nil
 }