kvstorage: batch WAG node deletions

Previously, truncateAppliedWAGNodeAndClearRaftState deleted one WAG node
per batch and committed immediately.

This commit does the following:

1) Rename truncateAppliedWAGNodeAndClearRaftState() to truncateBatch().

2) Introduce a cluster-setting that controls the batch size.

3) Try to fit up-to batchSize deletion in each call to truncateBatch().

Benchmark results:

```
BenchmarkWAGTruncation/batchSize=1                247114             15261 ns/op            1002 B/op         20 allocs/op
BenchmarkWAGTruncation/batchSize=4                396939              5533 ns/op             608 B/op         11 allocs/op
BenchmarkWAGTruncation/batchSize=8                452824              5220 ns/op             549 B/op         10 allocs/op
BenchmarkWAGTruncation/batchSize=16               447286              4400 ns/op             505 B/op          9 allocs/op
BenchmarkWAGTruncation/batchSize=32               512752              2402 ns/op             503 B/op          9 allocs/op
BenchmarkWAGTruncation/batchSize=64               484096              3737 ns/op             481 B/op          9 allocs/op
```

Release note: None
Epic: none

Author: iskettaneh

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -29,6 +29,7 @@ go_library(
         "//pkg/kv/kvserver/spanset",
         "//pkg/raft/raftpb",
         "//pkg/roachpb",
+        "//pkg/settings",
         "//pkg/settings/cluster",
         "//pkg/storage",
         "//pkg/storage/enginepb",

pkg/kv/kvserver/kvstorage/wag_truncator.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -24,6 +25,14 @@ import (
 	"golang.org/x/time/rate"
 )
 
+var wagTruncatorBatchSize = settings.RegisterIntSetting(
+	settings.SystemOnly,
+	"kv.wag.truncator_batch_size",
+	"number of WAG nodes to delete per write batch during truncation",
+	8,
+	settings.IntInRange(1, 1024),
+)
+
 // WAGTruncatorTestingKnobs contains testing knobs for the WAGTruncator.
 type WAGTruncatorTestingKnobs struct {
 	// AfterTruncationCallback is called after each truncation attempt.
@@ -118,39 +127,39 @@ func (t *WAGTruncator) truncateAppliedNodes(ctx context.Context) error {
 	stateReader := t.eng.StateEngine().NewReader(storage.GuaranteedDurability)
 	defer stateReader.Close()
 	for {
-		truncated, err := t.truncateAppliedWAGNodeAndClearRaftState(ctx, stateReader)
+		truncated, err := t.truncateBatch(ctx, stateReader)
 		if err != nil || !truncated {
 			return err
 		}
 	}
 }
 
-// truncateAppliedWAGNodeAndClearRaftState deletes the first WAG node if all of
-// its events have been applied to the state engine. For nodes containing
+// truncateBatch deletes up to a batch-sized prefix of WAG nodes if all of
+// their events have been applied to the state engine. For nodes containing
 // EventDestroy or EventSubsume events, it also clears the corresponding raft
-// log prefix from the engine and the sideloaded entries storage.
+// log prefix from the engine and the sideloaded entries.
 //
-// Returns a bool indicating whether a WAG node was deleted or not.
+// Returns a bool indicating whether some WAG nodes were deleted or not.
 //
 // The caller must provide a stateRO reader with GuaranteedDurability so that
 // only state confirmed flushed to persistent storage is visible. This ensures
 // we never delete a WAG node whose mutations aren't flushed yet. The caller is
 // also responsible for creating and committing/closing the write batch in
 // raft.WO.
-// TODO(ibrahim): Support deleting multiple WAG nodes within the same batch.
-func (t *WAGTruncator) truncateAppliedWAGNodeAndClearRaftState(
-	ctx context.Context, stateRO StateRO,
-) (bool, error) {
+func (t *WAGTruncator) truncateBatch(ctx context.Context, stateRO StateRO) (bool, error) {
+	batchSize := wagTruncatorBatchSize.Get(&t.st.SV)
+	var count int64
 	var iter wag.Iterator
-	truncateIndex := t.truncIndex.Load() + 1
-	iterStartKey := keys.StoreWAGNodeKey(truncateIndex)
+	targetIndex := t.truncIndex.Load() + 1
+
 	b := t.eng.LogEngine().NewWriteBatch()
 	defer b.Close()
-
-	for index, node := range iter.IterFrom(ctx, t.eng.LogEngine(), iterStartKey) {
-		if index != truncateIndex && index > t.initIndex {
+	for index, node := range iter.IterFrom(
+		ctx, t.eng.LogEngine(), keys.StoreWAGNodeKey(targetIndex),
+	) {
+		if index != targetIndex && index > t.initIndex {
 			// We cannot ignore gaps for WAG indices > initIndex.
-			return false, nil
+			break
 		}
 		// TODO(ibrahim): Right now, the canApplyWAGNode function returns a list of
 		// raftCatchUpTargets that are not needed for the purposes of truncation,
@@ -161,7 +170,7 @@ func (t *WAGTruncator) truncateAppliedWAGNodeAndClearRaftState(
 		}
 		if action.apply {
 			// If an event needs to be applied, the WAG node cannot be deleted yet.
-			return false, nil
+			break
 		}
 		if err := wag.Delete(b, index); err != nil {
 			return false, err
@@ -178,13 +187,23 @@ func (t *WAGTruncator) truncateAppliedWAGNodeAndClearRaftState(
 			}
 		}
 
-		if err = b.Commit(false); err != nil {
-			return false, err
+		targetIndex = index + 1
+		count++
+		if count >= batchSize {
+			break
 		}
-		t.truncIndex.Store(index)
-		return true, nil
 	}
-	return false, iter.Error() // either no more WAG nodes or iter hit an error
+	if count == 0 {
+		return false, nil
+	}
+	if err := iter.Error(); err != nil {
+		return false, err
+	}
+	if err := b.Commit(false); err != nil {
+		return false, err
+	}
+	t.truncIndex.Store(targetIndex - 1) // targetIndex is pointing at the last index truncated + 1.
+	return true, nil
 }
 
 // clearReplicaRaftLogAndSideloaded clears raft log entries at or below the given index for

pkg/kv/kvserver/kvstorage/wag_truncator_test.go

@@ -7,6 +7,7 @@ package kvstorage
 
 import (
 	"context"
+	"fmt"
 	"math"
 	"testing"
 
@@ -55,12 +56,12 @@ func (e *testEngines) writeWAGNode(t *testing.T, event wagpb.Event) {
 
 // writeWAGNodesAt writes WAG nodes at the specified WAG sequence indices for
 // the given replica. Each node contains an EventApply event with its raft index
-// set to the WAG entry index + 10.
+// set to the WAG entry index + 1.
 func (e *testEngines) writeWAGNodesAt(t *testing.T, wagIndices []uint64, r roachpb.FullReplicaID) {
 	t.Helper()
 	for idx, wagIdx := range wagIndices {
 		event := wagpb.Event{
-			Addr: wagpb.MakeAddr(r, kvpb.RaftIndex(idx+10)),
+			Addr: wagpb.MakeAddr(r, kvpb.RaftIndex(idx+1)),
 			Type: wagpb.EventApply,
 		}
 		require.NoError(t, wag.Write(
@@ -462,3 +463,115 @@ func TestWAGTruncatorBackground(t *testing.T) {
 	require.Empty(t, e.listWAGNodes(t))
 	require.Equal(t, uint64(7), truncator.truncIndex.Load())
 }
+
+// TestTruncateBatching verifies that truncateBatch() respects the batch size
+// setting.
+//
+// WAG layout: indices [3, 7, 10, 11, 12], and node at index 12 is not durably
+// applied.
+func TestTruncateBatching(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
+	sl := MakeStateLoader(r1.RangeID)
+	wagNodeIndices := []uint64{3, 7, 10, 11, 12}
+	for _, tc := range []struct {
+		initIndex      uint64
+		batchSize      int64
+		wantTruncated  bool
+		wantRemaining  []uint64
+		wantTruncIndex uint64
+	}{
+		{initIndex: 0, batchSize: 1, wantTruncated: false, wantRemaining: []uint64{3, 7, 10, 11, 12}, wantTruncIndex: 0},
+		{initIndex: 0, batchSize: 8, wantTruncated: false, wantRemaining: []uint64{3, 7, 10, 11, 12}, wantTruncIndex: 0},
+		{initIndex: 3, batchSize: 1, wantTruncated: true, wantRemaining: []uint64{7, 10, 11, 12}, wantTruncIndex: 3},
+		{initIndex: 3, batchSize: 8, wantTruncated: true, wantRemaining: []uint64{7, 10, 11, 12}, wantTruncIndex: 3},
+		{initIndex: 7, batchSize: 1, wantTruncated: true, wantRemaining: []uint64{7, 10, 11, 12}, wantTruncIndex: 3},
+		{initIndex: 7, batchSize: 2, wantTruncated: true, wantRemaining: []uint64{10, 11, 12}, wantTruncIndex: 7},
+		{initIndex: 7, batchSize: 8, wantTruncated: true, wantRemaining: []uint64{10, 11, 12}, wantTruncIndex: 7},
+		{initIndex: 10, batchSize: 1, wantTruncated: true, wantRemaining: []uint64{7, 10, 11, 12}, wantTruncIndex: 3},
+		{initIndex: 10, batchSize: 4, wantTruncated: true, wantRemaining: []uint64{12}, wantTruncIndex: 11},
+		// Node 11 isn't applied yet, so it's not truncated.
+		{initIndex: 10, batchSize: 8, wantTruncated: true, wantRemaining: []uint64{12}, wantTruncIndex: 11},
+	} {
+		t.Run("", func(t *testing.T) {
+			st := cluster.MakeTestingClusterSettings()
+			wagTruncatorBatchSize.Override(ctx, &st.SV, tc.batchSize)
+			e := makeTestEngines()
+			defer e.Close()
+			e.writeWAGNodesAt(t, wagNodeIndices, r1)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq)
+			truncator.initIndex = tc.initIndex
+
+			require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID))
+			// Ensure that the last WAG node isn't applied.
+			require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(),
+				&kvserverpb.RangeAppliedState{RaftAppliedIndex: kvpb.RaftIndex(len(wagNodeIndices) - 1)}))
+			require.NoError(t, e.stateEngine.Flush())
+
+			stateReader := e.StateEngine().NewReader(storage.GuaranteedDurability)
+			defer stateReader.Close()
+			truncated, err := truncator.truncateBatch(ctx, stateReader)
+			require.NoError(t, err)
+			require.Equal(t, tc.wantTruncated, truncated)
+			require.Equal(t, tc.wantRemaining, e.listWAGNodes(t))
+			require.Equal(t, tc.wantTruncIndex, truncator.truncIndex.Load())
+		})
+	}
+}
+
+// BenchmarkWAGTruncation measures the cost of truncating WAG nodes at different
+// batch sizes. It uses an in-memory engine so it doesn't really test the real
+// thing, but it should give an idea of the improvement of different batch
+// sizes.
+func BenchmarkWAGTruncation(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ctx := context.Background()
+	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
+	sl := MakeStateLoader(r1.RangeID)
+	for _, batchSize := range []int64{1, 4, 8, 16, 32, 64} {
+		b.Run(fmt.Sprintf("batchSize=%d", batchSize), func(b *testing.B) {
+			b.StopTimer()
+			st := cluster.MakeTestingClusterSettings()
+			wagTruncatorBatchSize.Override(ctx, &st.SV, batchSize)
+			eng := storage.NewDefaultInMemForTesting()
+			defer eng.Close()
+			engines := MakeEngines(eng)
+			var seq wag.Seq
+			truncator := NewWAGTruncator(st, engines, &seq)
+
+			// Write numNodes WAG nodes that are all eligible for
+			// truncation.
+			for j := 0; j < b.N; j++ {
+				index := seq.Next()
+				if err := wag.Write(eng, index, wagpb.Node{
+					Events: []wagpb.Event{{
+						Addr: wagpb.MakeAddr(r1, kvpb.RaftIndex(j+1)),
+						Type: wagpb.EventApply,
+					}},
+				}); err != nil {
+					b.Fatal(err)
+				}
+			}
+			if err := sl.SetRaftReplicaID(ctx, eng, r1.ReplicaID); err != nil {
+				b.Fatal(err)
+			}
+			if err := sl.SetRangeAppliedState(ctx, eng,
+				&kvserverpb.RangeAppliedState{
+					RaftAppliedIndex: kvpb.RaftIndex(b.N + 1),
+				}); err != nil {
+				b.Fatal(err)
+			}
+			if err := eng.Flush(); err != nil {
+				b.Fatal(err)
+			}
+
+			b.StartTimer()
+			err := truncator.truncateAppliedNodes(ctx)
+			b.StopTimer()
+			require.NoError(b, err)
+			require.Equal(b, truncator.truncIndex.Load(), uint64(b.N))
+		})
+	}
+}