kvstorage: add live WAG truncation

This commit adds the ability to perform live WAG truncation. It keeps track of
the last successfully truncated WAG node index, and only attempts to perform a
round of WAG truncation if (1) The state engine flushes, and (2) There is a
higher WAG sequence number than the last one we truncated.

Release note: None

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

Author: iskettaneh

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "//pkg/util/iterutil",
         "//pkg/util/log",
         "//pkg/util/protoutil",
+        "//pkg/util/stop",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@org_golang_x_time//rate",
@@ -86,6 +87,7 @@ 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/store.go

@@ -65,6 +65,11 @@ func (s *Seq) Next() uint64 {
 	return s.index.Add(1)
 }
 
+// Load returns the last used WAG sequence number.
+func (s *Seq) Load() uint64 {
+	return s.index.Load()
+}
+
 // Write puts the WAG node under the specific sequence number into the given
 // writer. The index must have been allocated to the caller by the sequencer.
 func Write(w storage.Writer, index uint64, node wagpb.Node) error {

pkg/kv/kvserver/kvstorage/wag_truncator.go

@@ -8,6 +8,7 @@ package kvstorage
 import (
 	"context"
 	"math"
+	"sync/atomic"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
@@ -17,6 +18,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/errors"
 	"golang.org/x/time/rate"
 )
@@ -28,43 +31,146 @@ import (
 //     replayed and made durable in the StateEngine.
 //   - [online] during the normal operation, truncate the WAG nodes that were
 //     durably applied to the state machine.
-//
-// TODO(ibrahim): Add the periodic truncation logic.
 type WAGTruncator struct {
 	st  *cluster.Settings
 	eng Engines
+	seq *wag.Seq
+	// wakeCh is signaled when there are potential WAG nodes to truncate.
+	wakeCh chan struct{}
+	// lastTruncatedWAGIndex is the index of the last WAG node that was
+	// successfully truncated. This is to be used by online WAG truncation to
+	// quickly seek into the potential WAG nodes to truncate.
+	lastTruncatedWAGIndex atomic.Uint64
 }
 
 // NewWAGTruncator creates a WAGTruncator.
-func NewWAGTruncator(st *cluster.Settings, eng Engines) *WAGTruncator {
-	return &WAGTruncator{st: st, eng: eng}
+func NewWAGTruncator(st *cluster.Settings, eng Engines, seq *wag.Seq) *WAGTruncator {
+	return &WAGTruncator{
+		st:     st,
+		eng:    eng,
+		seq:    seq,
+		wakeCh: make(chan struct{}, 1),
+	}
+}
+
+// Start launches the background goroutine that performs live WAG truncation.
+// It must be called before starting the engine. It sets lastTruncatedWAGIndex,
+// and we need to set it before any WAG node could have been created so that we
+// guarantee that the periodic WAG truncation will start from the first WAG node.
+//
+// TODO(ibrahim): Add a TestKnob for keeping a suffix of the WAG for debugging.
+func (t *WAGTruncator) Start(ctx context.Context, stopper *stop.Stopper) error {
+	t.lastTruncatedWAGIndex.Store(t.seq.Load())
+	return stopper.RunAsyncTask(ctx, "wag-truncation", func(ctx context.Context) {
+		ctx, cancel := stopper.WithCancelOnQuiesce(ctx)
+		defer cancel()
+		for {
+			select {
+			case <-t.wakeCh:
+				t.truncateAppliedNodesLive(ctx)
+			case <-ctx.Done():
+				return
+			}
+		}
+	})
+}
+
+// DurabilityAdvancedCallback is invoked whenever the state engine completes a
+// flush. It checks whether there could possibly be WAG nodes to truncate by
+// comparing lastTruncatedWAGIndex against seq.Load(). If there are potential
+// truncation opportunities, it sends a non-blocking signal to wake the
+// background goroutine. It must return quickly and must not call into the
+// engine to avoid deadlock (see storage.Engine.RegisterFlushCompletedCallback).
+func (t *WAGTruncator) DurabilityAdvancedCallback() {
+	// The key point is the point in time where we call "t.seq.Load()". If the
+	// sequence number is greater than the last truncated WAG sequence number, we
+	// will attempt to truncate if there are any WAG nodes that are ready for
+	// truncation. If it found no WAG nodes that are ready to be truncated, it
+	// will attempt to truncate on every state engine flush until
+	// sequence number == lastTruncatedWAGIndex.
+	seq, lastTruncated := t.seq.Load(), t.lastTruncatedWAGIndex.Load()
+	if seq == lastTruncated {
+		return
+	}
+	if seq < lastTruncated {
+		log.KvExec.Fatalf(context.Background(),
+			"WAG seq %d < lastTruncatedWAGIndex %d", seq, lastTruncated)
+	}
+	select {
+	case t.wakeCh <- struct{}{}:
+	default:
+	}
 }
 
 // TruncateAll truncates all applied WAG nodes. It's meant to be used at engine
 // startup right after we replay the WAG nodes and sync the state engine.
 func (t *WAGTruncator) TruncateAll(ctx context.Context) error {
+	_, err := t.truncateAppliedNodes(ctx, 0)
+	return err
+}
+
+// truncateAppliedNodesLive is called by the background goroutine to truncate
+// applied WAG nodes. Unlike TruncateAll, it only truncates nodes with index
+// after lastTruncatedWAGIndex. This makes it avoid jumping over gaps in WAG
+// node indices.
+//
+// Note that jumping over gaps during WAG truncation should be safe because a
+// WAG node is only truncated if it has been applied to the state engine, and it
+// by definition has all the required events applied. However, we avoid jumping
+// over gaps so that we know the exact index to truncate next, which facilitates
+// seeking past previously deleted WAG garbage.
+func (t *WAGTruncator) truncateAppliedNodesLive(ctx context.Context) {
+	startIndex := t.lastTruncatedWAGIndex.Load() + 1
+	lastTruncated, err := t.truncateAppliedNodes(ctx, startIndex)
+	if err != nil {
+		// TODO(ibrahim): We need to decide if we need to retry truncating this
+		//  WAG node now, or rely on the next state engine flush to do it.
+		log.KvExec.Errorf(ctx, "truncating WAG node: %+v", err)
+	}
+	if lastTruncated >= startIndex {
+		t.lastTruncatedWAGIndex.Store(lastTruncated)
+	}
+}
+
+// truncateAppliedNodes is a helper function used by both offline and online WAG
+// truncation functions. It repeatedly tries to delete a WAG node in a batch and
+// commit that batch. When startIndex is 0, it repeatedly truncates the first
+// available node (used at startup to drain the WAG). When startIndex is greater
+// than 0, it truncates nodes sequentially starting from that index (used during
+// online operation to avoid jumping over gaps in WAG indices).
+//
+// Returns the index of the last successfully truncated node, or 0 if no nodes
+// were truncated.
+func (t *WAGTruncator) truncateAppliedNodes(
+	ctx context.Context, startIndex uint64,
+) (uint64, error) {
 	stateReader := t.eng.StateEngine().NewReader(storage.GuaranteedDurability)
 	defer stateReader.Close()
+	nextIndex := startIndex
+	var lastTruncated uint64
 	for {
 		if err := ctx.Err(); err != nil {
-			return err
+			return lastTruncated, err
 		}
 		b := t.eng.LogEngine().NewWriteBatch()
 		truncated, err := t.truncateAppliedWAGNodeAndClearRaftState(
-			ctx, Raft{RO: t.eng.LogEngine(), WO: b}, stateReader, 0, /* index */
+			ctx, Raft{RO: t.eng.LogEngine(), WO: b}, stateReader, nextIndex,
 		)
 		if err == nil && truncated {
 			err = b.Commit(false /* sync */)
 		}
 		b.Close()
 		if err != nil {
-			return err
+			return lastTruncated, err
 		}
 		if !truncated {
-			break
+			return lastTruncated, nil
+		}
+		if nextIndex > 0 {
+			lastTruncated = nextIndex
+			nextIndex++
 		}
 	}
-	return nil
 }
 
 // truncateAppliedWAGNodeAndClearRaftState deletes a WAG node if all of its

pkg/kv/kvserver/kvstorage/wag_truncator_test.go

@@ -8,6 +8,7 @@ package kvstorage
 import (
 	"context"
 	"math"
+	"slices"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -20,8 +21,11 @@ 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"
 )
@@ -103,6 +107,30 @@ 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)
@@ -240,7 +268,7 @@ func TestTruncateApplied(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq)
 			tc.setup(t, &e)
 			require.NoError(t, e.stateEngine.Flush())
 			require.NoError(t, truncator.TruncateAll(ctx))
@@ -268,7 +296,7 @@ func TestTruncateAndClearRaftState(t *testing.T) {
 		t.Run(eventType.String(), func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq)
 
 			// Write WAG nodes: init then destroy/subsume at index 20.
 			e.writeWAGNode(t, wagpb.Event{
@@ -372,7 +400,7 @@ func TestTruncateGapHandling(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines)
+			truncator := NewWAGTruncator(st, e.Engines, &e.seq)
 
 			// Write WAG nodes at indices 2, 4, 6.
 			e.seq.Next()
@@ -412,3 +440,81 @@ func TestTruncateGapHandling(t *testing.T) {
 		})
 	}
 }
+
+// TestWAGTruncatorBackground verifies that the WAGTruncator background
+// goroutine only truncates WAG nodes when both conditions are met: (1) the
+// state engine has flushed, and (2) there are WAG nodes to truncate (i.e.,
+// seq.Load() > lastTruncatedWAGIndex).
+func TestWAGTruncatorBackground(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	st := cluster.MakeTestingClusterSettings()
+	e := makeTestEngines()
+	defer e.Close()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
+	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
+	sl := MakeStateLoader(r1.RangeID)
+
+	// Initialize replica state so events can be considered applied.
+	require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID))
+	require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(),
+		&kvserverpb.RangeAppliedState{RaftAppliedIndex: 100}))
+
+	// Start the periodic WAG truncation background task.
+	truncator := NewWAGTruncator(st, e.Engines, &e.seq)
+	require.NoError(t, truncator.Start(ctx, stopper))
+
+	flushStateEngineAndSignal := func() {
+		require.NoError(t, e.StateEngine().Flush())
+		truncator.DurabilityAdvancedCallback()
+	}
+
+	// No WAG nodes exist. Flushing the state engine should not cause the
+	// truncator to do anything (seq.Load() == lastTruncatedWAGIndex == 0).
+	flushStateEngineAndSignal()
+	eventuallyExpectWAGNodesIndices(t, &e, nil)
+	require.Equal(t, truncator.lastTruncatedWAGIndex.Load(), uint64(0))
+
+	// Write two WAG nodes whose events are applied (index <= 100).
+	e.writeWAGNode(t, wagpb.Event{
+		Addr: wagpb.MakeAddr(r1, 10), Type: wagpb.EventInit,
+	})
+	e.writeWAGNode(t, wagpb.Event{
+		Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventApply,
+	})
+
+	// WAG nodes exist, but the state engine hasn't flushed yet, so the
+	// GuaranteedDurability reader won't see the replica state. Signal the
+	// truncator without flushing first.
+	truncator.DurabilityAdvancedCallback()
+	eventuallyExpectWAGNodesIndices(t, &e, []uint64{1, 2})
+	require.Equal(t, truncator.lastTruncatedWAGIndex.Load(), uint64(0))
+
+	// Now flush the state engine and signal again. Both nodes should be
+	// truncated since their events are applied (index 10 and 20 <= 100).
+	flushStateEngineAndSignal()
+	eventuallyExpectWAGNodesIndices(t, &e, nil)
+	// The lastTruncatedWAGIndex is stored after truncateAppliedNodes returns,
+	// so there is a small window between the WAG nodes being deleted and the
+	// index being updated. Use an eventually-consistent check.
+	eventuallyExpectLastTruncatedWAGIndex(t, truncator, 2)
+
+	// Write a third WAG node that is NOT applied (index 200 > 100).
+	e.writeWAGNode(t, wagpb.Event{
+		Addr: wagpb.MakeAddr(r1, 200), Type: wagpb.EventApply,
+	})
+	flushStateEngineAndSignal()
+	// Node 3 should remain because its event isn't applied yet.
+	eventuallyExpectWAGNodesIndices(t, &e, []uint64{3})
+	require.Equal(t, truncator.lastTruncatedWAGIndex.Load(), uint64(2))
+
+	// Advance the applied index past 200 and flush. Now node 3 should be
+	// truncated.
+	require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(),
+		&kvserverpb.RangeAppliedState{RaftAppliedIndex: 200}))
+	flushStateEngineAndSignal()
+	eventuallyExpectWAGNodesIndices(t, &e, nil)
+	eventuallyExpectLastTruncatedWAGIndex(t, truncator, 3)
+}