Add a db.CompactHeadWithoutTruncation function (#638)

* Add a db.CompactHeadWithoutTruncation function

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>

* Apply suggestions from code review

Co-authored-by: Arve Knudsen <arve.knudsen@gmail.com>

---------

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>
Co-authored-by: Arve Knudsen <arve.knudsen@gmail.com>

Author: codesome

tsdb/db.go

@@ -1249,7 +1249,7 @@ func (db *DB) Compact(ctx context.Context) (returnErr error) {
 		// We do need to wait for any overlapping appenders that started previously to finish.
 		db.head.WaitForAppendersOverlapping(rh.MaxTime())
 
-		if err := db.compactHead(rh); err != nil {
+		if err := db.compactHead(rh, true); err != nil {
 			return fmt.Errorf("compact head: %w", err)
 		}
 		// Consider only successful compactions for WAL truncation.
@@ -1286,7 +1286,7 @@ func (db *DB) CompactHead(head *RangeHead) error {
 	db.cmtx.Lock()
 	defer db.cmtx.Unlock()
 
-	if err := db.compactHead(head); err != nil {
+	if err := db.compactHead(head, true); err != nil {
 		return fmt.Errorf("compact head: %w", err)
 	}
 
@@ -1296,6 +1296,18 @@ func (db *DB) CompactHead(head *RangeHead) error {
 	return nil
 }
 
+// CompactHeadWithoutTruncation compacts the given RangeHead but does not truncate the
+// in-memory data and the WAL related to this compaction.
+func (db *DB) CompactHeadWithoutTruncation(head *RangeHead) error {
+	db.cmtx.Lock()
+	defer db.cmtx.Unlock()
+
+	if err := db.compactHead(head, false); err != nil {
+		return fmt.Errorf("compact head without truncation: %w", err)
+	}
+	return nil
+}
+
 // CompactOOOHead compacts the OOO Head.
 func (db *DB) CompactOOOHead(ctx context.Context) error {
 	db.cmtx.Lock()
@@ -1400,7 +1412,7 @@ func (db *DB) compactOOO(dest string, oooHead *OOOCompactionHead) (_ []ulid.ULID
 
 // compactHead compacts the given RangeHead.
 // The compaction mutex should be held before calling this method.
-func (db *DB) compactHead(head *RangeHead) error {
+func (db *DB) compactHead(head *RangeHead, truncateMemory bool) error {
 	uid, err := db.compactor.Write(db.dir, head, head.MinTime(), head.BlockMaxTime(), nil)
 	if err != nil {
 		return fmt.Errorf("persist head block: %w", err)
@@ -1415,6 +1427,9 @@ func (db *DB) compactHead(head *RangeHead) error {
 		}
 		return fmt.Errorf("reloadBlocks blocks: %w", err)
 	}
+	if !truncateMemory {
+		return nil
+	}
 	if err = db.head.truncateMemory(head.BlockMaxTime()); err != nil {
 		return fmt.Errorf("head memory truncate: %w", err)
 	}

tsdb/db_test.go

@@ -7057,3 +7057,59 @@ func TestAbortBlockCompactions(t *testing.T) {
 	require.True(t, db.head.compactable(), "head should be compactable")
 	require.Equal(t, 4, compactions, "expected 4 compactions to be completed")
 }
+
+func TestCompactHeadWithoutTruncation(t *testing.T) {
+	setupDB := func() *DB {
+		db := openTestDB(t, nil, nil)
+		t.Cleanup(func() {
+			require.NoError(t, db.Close())
+		})
+		db.DisableCompactions()
+
+		// Add samples to the head.
+		lbls := labels.FromStrings("foo", "bar")
+		app := db.Appender(context.Background())
+		_, err := app.Append(0, lbls, 0, 0)
+		require.NoError(t, err)
+		_, err = app.Append(0, lbls, DefaultBlockDuration/2, float64(DefaultBlockDuration/2))
+		require.NoError(t, err)
+		_, err = app.Append(0, lbls, DefaultBlockDuration-1, float64(DefaultBlockDuration-1))
+		require.NoError(t, err)
+		_, err = app.Append(0, lbls, 2*DefaultBlockDuration, float64(2*DefaultBlockDuration))
+		require.NoError(t, err)
+		require.NoError(t, app.Commit())
+
+		return db
+	}
+
+	testQuery := func(db *DB, expSamples []chunks.Sample) {
+		rh := NewRangeHead(db.Head(), math.MinInt64, math.MaxInt64)
+		q, err := NewBlockQuerier(rh, math.MinInt64, math.MaxInt64)
+		require.NoError(t, err)
+		ss := query(t, q, labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"))
+		require.Equal(t, map[string][]chunks.Sample{`{foo="bar"}`: expSamples}, ss)
+	}
+
+	{ // Compact the head with truncation.
+		db := setupDB()
+		rh := NewRangeHead(db.Head(), 0, DefaultBlockDuration-1)
+		require.NoError(t, db.CompactHead(rh))
+		// Samples got truncated from the head.
+		testQuery(db, []chunks.Sample{
+			sample{t: 2 * DefaultBlockDuration, f: float64(2 * DefaultBlockDuration)},
+		})
+	}
+
+	{ // Compact the head without truncation.
+		db := setupDB()
+		rh := NewRangeHead(db.Head(), 0, DefaultBlockDuration-1)
+		require.NoError(t, db.CompactHeadWithoutTruncation(rh))
+		// All samples still exist in the head.
+		testQuery(db, []chunks.Sample{
+			sample{t: 0, f: 0},
+			sample{t: DefaultBlockDuration / 2, f: float64(DefaultBlockDuration / 2)},
+			sample{t: DefaultBlockDuration - 1, f: float64(DefaultBlockDuration - 1)},
+			sample{t: 2 * DefaultBlockDuration, f: float64(2 * DefaultBlockDuration)},
+		})
+	}
+}