Fix backfill TSDB replay with tests for it

Signed-off-by: Ganesh Vernekar <[email protected]>

Author: codesome

pkg/ingester/ingester_v2.go

@@ -1088,8 +1088,15 @@ func (i *Ingester) closeAllTSDB() {
 		}(userDB)
 	}
 
-	// Wait until all Close() completed
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		i.closeAllBackfillTSDBs()
+	}()
+
 	i.userStatesMtx.Unlock()
+
+	// Wait until all Close() completed
 	wg.Wait()
 }
 

pkg/ingester/ingester_v2_backfill.go

@@ -124,7 +124,6 @@ func (i *Ingester) getOrCreateBackfillTSDB(userID string, ts int64) (*tsdbBucket
 	insertIdx := len(userBuckets)
 	for idx, b := range userBuckets {
 		if ts >= b.bucketStart && ts < b.bucketEnd {
-			fmt.Println("bucket found")
 			return b, userBuckets, nil
 		}
 
@@ -201,7 +200,7 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 		}
 
 		for bucketID, bucketName := range bucketNames {
-			if bucketName.IsDir() {
+			if !bucketName.IsDir() {
 				continue
 			}
 
@@ -229,7 +228,7 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 			}
 
 			wg.Add(1)
-			go func(bucketID int, userID, bucketName string) {
+			go func(bucketID int, userID, bucketName, dbDir string) {
 				defer wg.Done()
 				defer openGate.Done()
 				defer func(ts time.Time) {
@@ -241,7 +240,7 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 					level.Error(util.Logger).Log("msg", "unable to get bucket range", "err", err, "user", userID, "bucketName", bucketName)
 					return
 				}
-				db, err := i.createNewTSDB(userID, filepath.Join(userID, bucketName), (end-start)*2, (end-start)*2, prometheus.NewRegistry())
+				db, err := i.createNewTSDB(userID, dbDir, (end-start)*2, (end-start)*2, prometheus.NewRegistry())
 				if err != nil {
 					level.Error(util.Logger).Log("msg", "unable to open user backfill TSDB", "err", err, "user", userID)
 					return
@@ -258,7 +257,7 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 				// Append at the end, we will sort it at the end.
 				i.TSDBState.backfillDBs.tsdbs[userID] = append(i.TSDBState.backfillDBs.tsdbs[userID], bucket)
 				i.TSDBState.backfillDBs.tsdbsMtx.Unlock()
-			}(bucketID, userID, bucketName.Name())
+			}(bucketID, userID, bucketName.Name(), dbPath)
 		}
 
 		if runErr != nil {
@@ -344,6 +343,30 @@ func (i *Ingester) backfillSelect(ctx context.Context, userID string, from, thro
 	return result, nil
 }
 
+func (i *Ingester) closeAllBackfillTSDBs() {
+	i.TSDBState.backfillDBs.tsdbsMtx.Lock()
+	defer i.TSDBState.backfillDBs.tsdbsMtx.Unlock()
+
+	wg := &sync.WaitGroup{}
+
+	for userID, buckets := range i.TSDBState.backfillDBs.tsdbs {
+		for _, bucket := range buckets {
+			wg.Add(1)
+			go func(uid string, db *userTSDB) {
+				defer wg.Done()
+
+				if err := db.Close(); err != nil {
+					level.Warn(util.Logger).Log("msg", "unable to close backfill TSDB", "user", uid, "bucket_dir", db.Dir(), "err", err)
+				}
+			}(userID, bucket.db)
+		}
+	}
+
+	wg.Wait()
+	// Clear all DBs irrespective of them closing.
+	i.TSDBState.backfillDBs.tsdbs = make(map[string][]*tsdbBucket)
+}
+
 // Assumes 1h bucket range for . TODO(codesome): protect stuff with locks.
 type backfillTSDBs struct {
 	tsdbsMtx sync.RWMutex

pkg/ingester/ingester_v2_test.go

@@ -1820,11 +1820,11 @@ func TestIngesterV2BackfillPushAndQuery(t *testing.T) {
 	ts := 100 * time.Hour.Milliseconds()
 	ingestSample(ts, 0, false)
 
-	// 99.5h
+	// 99.5h, to the main TSDB.
 	ts = 99*time.Hour.Milliseconds() + 30*time.Minute.Milliseconds()
 	ingestSample(ts, 0, false)
 
-	// 99h
+	// 99h, to the main TSDB.
 	ts = 99 * time.Hour.Milliseconds()
 	ingestSample(ts, 0, false)
 
@@ -1848,12 +1848,12 @@ func TestIngesterV2BackfillPushAndQuery(t *testing.T) {
 	ts = (100-12)*time.Hour.Milliseconds() + 1*time.Millisecond.Milliseconds()
 	ingestSample(ts, 4, false)
 
-	// 100h-12h, out of bounds even for backfill.
-	ts = (100 - 12) * time.Hour.Milliseconds()
+	// 100h-backfillLimit, out of bounds even for backfill.
+	ts = 100*time.Hour.Milliseconds() - backfillLimit.Milliseconds()
 	ingestSample(ts, 4, true)
 
-	// 100h-13h, out of bounds even for backfill.
-	ts = (100 - 13) * time.Hour.Milliseconds()
+	// 99h-backfillLimit, out of bounds even for backfill.
+	ts = 99*time.Hour.Milliseconds() - backfillLimit.Milliseconds()
 	ingestSample(ts, 4, true)
 
 	// Query back all the samples.
@@ -1865,4 +1865,31 @@ func TestIngesterV2BackfillPushAndQuery(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, res)
 	assert.Equal(t, expectedIngested, res.Timeseries)
+
+	// Restart to check if we can still query backfill TSDBs.
+
+	// Stop ingester.
+	require.NoError(t, services.StopAndAwaitTerminated(context.Background(), i))
+
+	overrides, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
+	require.NoError(t, err)
+
+	i, err = NewV2(i.cfg, defaultClientTestConfig(), overrides, nil)
+	require.NoError(t, err)
+	require.NoError(t, services.StartAndAwaitRunning(context.Background(), i))
+
+	// Wait until it's ACTIVE
+	test.Poll(t, 10*time.Millisecond, ring.ACTIVE, func() interface{} {
+		return i.lifecycler.GetState()
+	})
+
+	// Query back all the samples.
+	res, err = i.v2Query(ctx, &client.QueryRequest{
+		StartTimestampMs: math.MinInt64,
+		EndTimestampMs:   math.MaxInt64,
+		Matchers:         []*client.LabelMatcher{{Type: client.REGEX_MATCH, Name: labels.MetricName, Value: ".*"}},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, res)
+	assert.Equal(t, expectedIngested, res.Timeseries)
 }