kvstorage: Introduce WAGTruncator Testing Knobs

This commit introduces TestingKnobs to the WAGTruncator, and adds
AfterTruncationCallback() that is called after truncation pass has
finished.

Release note: None

Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

Author: iskettaneh

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -89,7 +89,6 @@ go_test(
         "//pkg/util/tracing",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_datadriven//:datadriven",
-        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
         "@org_golang_x_time//rate",
     ],

pkg/kv/kvserver/kvstorage/wag_truncator.go

@@ -24,6 +24,14 @@ import (
 	"golang.org/x/time/rate"
 )
 
+// WAGTruncatorTestingKnobs contains testing knobs for the WAGTruncator.
+type WAGTruncatorTestingKnobs struct {
+	// AfterTruncationCallback, if set, is called after each invocation of
+	// truncateAppliedNodesLive completes. Can be used to synchronize tests
+	// with the background truncation goroutine.
+	AfterTruncationCallback func()
+}
+
 // WAGTruncator truncates applied WAG nodes and clears their associated raft
 // state (log entries and sideloaded files).
 type WAGTruncator struct {
@@ -40,18 +48,24 @@ type WAGTruncator struct {
 	lastTruncatedWAGIndex atomic.Uint64
 	// wakeCh is signaled when there are potential WAG nodes to truncate.
 	wakeCh chan struct{}
+	knobs  WAGTruncatorTestingKnobs
 }
 
 // NewWAGTruncator creates a WAGTruncator.
 func NewWAGTruncator(
-	st *cluster.Settings, eng Engines, seq *wag.Seq, lastIndexAfterStartup uint64,
+	st *cluster.Settings,
+	eng Engines,
+	seq *wag.Seq,
+	lastIndexAfterStartup uint64,
+	knobs WAGTruncatorTestingKnobs,
 ) *WAGTruncator {
 	return &WAGTruncator{
 		st:                        st,
 		eng:                       eng,
 		seq:                       seq,
 		lastWAGIndexBeforeStartup: lastIndexAfterStartup,
 		wakeCh:                    make(chan struct{}, 1),
+		knobs:                     knobs,
 	}
 }
 
@@ -105,6 +119,11 @@ func (t *WAGTruncator) DurabilityAdvancedCallback() {
 // truncateAppliedNodesLive is called by the background goroutine to truncate
 // applied WAG nodes.
 func (t *WAGTruncator) truncateAppliedNodesLive(ctx context.Context) {
+	defer func() {
+		if t.knobs.AfterTruncationCallback != nil {
+			t.knobs.AfterTruncationCallback()
+		}
+	}()
 	startIndex := t.lastTruncatedWAGIndex.Load() + 1
 	if startIndex <= t.lastWAGIndexBeforeStartup {
 		// There are WAG nodes that we can potentially truncate with

pkg/kv/kvserver/kvstorage/wag_truncator_test.go

@@ -8,7 +8,6 @@ package kvstorage
 import (
 	"context"
 	"math"
-	"slices"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -21,11 +20,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
-	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
-	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/time/rate"
 )
@@ -107,28 +104,6 @@ func (e *testEngines) listWAGNodes(t *testing.T) []uint64 {
 	return indices
 }
 
-func eventuallyExpectWAGNodesIndices(t *testing.T, e *testEngines, expected []uint64) {
-	t.Helper()
-	testutils.SucceedsSoon(t, func() error {
-		actual := e.listWAGNodes(t)
-		if slices.Equal(expected, actual) {
-			return nil
-		}
-		return errors.Newf("expected WAG nodes %v, got %v", expected, actual)
-	})
-}
-
-func eventuallyExpectLastTruncatedWAGIndex(t *testing.T, truncator *WAGTruncator, expected uint64) {
-	t.Helper()
-	testutils.SucceedsSoon(t, func() error {
-		actual := truncator.lastTruncatedWAGIndex.Load()
-		if actual == expected {
-			return nil
-		}
-		return errors.Newf("expected lastTruncatedWAGIndex %d, got %d", expected, actual)
-	})
-}
-
 func TestTruncateApplied(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -266,7 +241,7 @@ func TestTruncateApplied(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0, WAGTruncatorTestingKnobs{})
 			tc.setup(t, &e)
 			require.NoError(t, e.stateEngine.Flush())
 			_, err := truncator.truncateAppliedNodes(ctx, 0, /* startIndex */
@@ -296,7 +271,7 @@ func TestTruncateAndClearRaftState(t *testing.T) {
 		t.Run(eventType.String(), func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0, WAGTruncatorTestingKnobs{})
 
 			// Write WAG nodes: init then destroy/subsume at index 20.
 			e.writeWAGNode(t, wagpb.Event{
@@ -428,7 +403,7 @@ func TestTruncateGapHandling(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0, WAGTruncatorTestingKnobs{})
 
 			// Write WAG nodes at indices 2, 4, 6.
 			e.seq.Next()
@@ -551,7 +526,7 @@ func TestTruncateAppliedNodes(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq, 0, WAGTruncatorTestingKnobs{})
 
 			// Write WAG nodes at indices 2, 4, 5.
 			e.seq.Next()
@@ -613,23 +588,33 @@ func TestWAGTruncatorBackground(t *testing.T) {
 		Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventApply,
 	})
 
-	// Start the periodic WAG truncation background task.
-	truncator := NewWAGTruncator(st, e.Engines, &e.seq, e.seq.Load() /* lastIndexAfterStartup */)
+	// Start the periodic WAG truncation background task with a knob that
+	// signals when truncation completes, so the test can synchronize
+	// deterministically instead of polling.
+	truncationDone := make(chan struct{}, 1)
+	truncator := NewWAGTruncator(st, e.Engines, &e.seq,
+		e.seq.Load(), /* lastIndexAfterStartup */
+		WAGTruncatorTestingKnobs{
+			AfterTruncationCallback: func() {
+				truncationDone <- struct{}{}
+			},
+		})
 	require.NoError(t, truncator.Start(ctx, stopper))
 	// Create a WAG node after startup.
 	e.writeWAGNode(t, wagpb.Event{
 		Addr: wagpb.MakeAddr(r1, 30), Type: wagpb.EventApply,
 	})
 	require.Equal(t, []uint64{1, 3, 4}, e.listWAGNodes(t))
 
-	flushStateEngineAndSignal := func() {
+	flushAndWaitForTruncation := func() {
 		require.NoError(t, e.StateEngine().Flush())
 		truncator.DurabilityAdvancedCallback()
+		<-truncationDone
 	}
 	// We expect all WAG nodes to be truncated when the state engine is flushed.
-	flushStateEngineAndSignal()
-	eventuallyExpectWAGNodesIndices(t, &e, nil)
-	eventuallyExpectLastTruncatedWAGIndex(t, truncator, 4)
+	flushAndWaitForTruncation()
+	require.Equal(t, ([]uint64)(nil), e.listWAGNodes(t))
+	require.Equal(t, uint64(4), truncator.lastTruncatedWAGIndex.Load())
 
 	// Write two WAG nodes whose events are applied (index <= 100).
 	e.writeWAGNode(t, wagpb.Event{
@@ -642,24 +627,24 @@ func TestWAGTruncatorBackground(t *testing.T) {
 
 	// Now flush the state engine and signal again. Both nodes should be
 	// truncated since their events are applied.
-	flushStateEngineAndSignal()
-	eventuallyExpectWAGNodesIndices(t, &e, nil)
-	eventuallyExpectLastTruncatedWAGIndex(t, truncator, 6)
+	flushAndWaitForTruncation()
+	require.Equal(t, ([]uint64)(nil), e.listWAGNodes(t))
+	require.Equal(t, uint64(6), truncator.lastTruncatedWAGIndex.Load())
 
 	// Write another WAG node but it is NOT applied yet.
 	e.writeWAGNode(t, wagpb.Event{
 		Addr: wagpb.MakeAddr(r1, 200), Type: wagpb.EventApply,
 	})
-	flushStateEngineAndSignal()
+	flushAndWaitForTruncation()
 	// Node 7 should remain because its event isn't applied yet.
-	eventuallyExpectWAGNodesIndices(t, &e, []uint64{7})
-	require.Equal(t, truncator.lastTruncatedWAGIndex.Load(), uint64(6))
+	require.Equal(t, []uint64{7}, e.listWAGNodes(t))
+	require.Equal(t, uint64(6), truncator.lastTruncatedWAGIndex.Load())
 
 	// Advance the applied index past 200 and flush. Now node 7 should be
 	// truncated.
 	require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(),
 		&kvserverpb.RangeAppliedState{RaftAppliedIndex: 200}))
-	flushStateEngineAndSignal()
-	eventuallyExpectWAGNodesIndices(t, &e, nil)
-	eventuallyExpectLastTruncatedWAGIndex(t, truncator, 7)
+	flushAndWaitForTruncation()
+	require.Equal(t, ([]uint64)(nil), e.listWAGNodes(t))
+	require.Equal(t, uint64(7), truncator.lastTruncatedWAGIndex.Load())
 }