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",

pkg/kv/kvserver/kvstorage/wag/store.go

@@ -65,6 +65,11 @@ func (s *Seq) Next() uint64 {
 	return s.index.Add(1)
 }
 
+// Load returns the current WAG sequence number without allocating a new one.
+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,16 +31,77 @@ 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
+	stopper *stop.Stopper
+	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 do nothing if there are no 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(
+	stopper *stop.Stopper, st *cluster.Settings, eng Engines, seq *wag.Seq,
+) *WAGTruncator {
+	return &WAGTruncator{
+		stopper: stopper,
+		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.
+func (t *WAGTruncator) Start(ctx context.Context) error {
+	t.lastTruncatedWAGIndex.Store(t.seq.Load())
+	return t.stopper.RunAsyncTask(ctx, "wag-truncation", func(ctx context.Context) {
+		ctx, cancel := t.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 are 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 that specific WAG node (and all the WAG nodes
+	// sequentially after it). If the sequence number == lastTruncatedWAGIndex,
+	// then there were no WAG nodes to truncate at that time. We will attempt
+	// again when the state engine completes another flush.
+	seq, lastTruncated := t.seq.Load(), t.lastTruncatedWAGIndex.Load()
+	if seq < lastTruncated {
+		log.KvExec.Fatalf(context.Background(),
+			"WAG seq %d < lastTruncatedWAGIndex %d", seq, lastTruncated)
+	}
+	if seq == lastTruncated {
+		return
+	}
+	select {
+	case t.wakeCh <- struct{}{}:
+	default:
+	}
 }
 
 // TruncateAll truncates all applied WAG nodes. It's meant to be used at engine
@@ -133,6 +197,48 @@ func (t *WAGTruncator) truncateAppliedWAGNodeAndClearRaftState(
 	return false, iter.Error()
 }
 
+// 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) {
+	stateReader := t.eng.StateEngine().NewReader(storage.GuaranteedDurability)
+	defer stateReader.Close()
+	for {
+		if err := ctx.Err(); err != nil {
+			return
+		}
+		nextIndex := t.lastTruncatedWAGIndex.Load() + 1
+		b := t.eng.LogEngine().NewWriteBatch()
+		truncated, err := t.truncateAppliedWAGNodeAndClearRaftState(
+			ctx,
+			Raft{RO: t.eng.LogEngine(), WO: b},
+			stateReader,
+			nextIndex,
+		)
+		if err == nil && truncated {
+			err = b.Commit(false /* sync */)
+		}
+		b.Close()
+		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 %d: %+v", nextIndex, err)
+			return
+		}
+		if !truncated {
+			return
+		}
+		t.lastTruncatedWAGIndex.Store(nextIndex)
+	}
+}
+
 // clearReplicaRaftLogAndSideloaded clears raft log entries at or below the given index for
 // a destroyed or subsumed replica, and it also deletes the sideloaded files associated with the
 // deleted entries.

pkg/kv/kvserver/kvstorage/wag_truncator_test.go

@@ -7,7 +7,9 @@ package kvstorage
 
 import (
 	"context"
+	"fmt"
 	"math"
+	"slices"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -20,8 +22,10 @@ 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/stretchr/testify/require"
 	"golang.org/x/time/rate"
 )
@@ -103,10 +107,23 @@ 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 fmt.Errorf("expected WAG nodes %v, got %v", expected, actual)
+	})
+}
+
 func TestTruncateApplied(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
 	st := cluster.MakeTestingClusterSettings()
 
 	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
@@ -240,7 +257,7 @@ func TestTruncateApplied(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines)
+			truncator := NewWAGTruncator(stopper, st, e.Engines, &e.seq)
 			tc.setup(t, &e)
 			require.NoError(t, e.stateEngine.Flush())
 			require.NoError(t, truncator.TruncateAll(ctx))
@@ -258,6 +275,8 @@ func TestTruncateAndClearRaftState(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
 	st := cluster.MakeTestingClusterSettings()
 
 	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
@@ -268,7 +287,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(stopper, st, e.Engines, &e.seq)
 
 			// Write WAG nodes: init then destroy/subsume at index 20.
 			e.writeWAGNode(t, wagpb.Event{
@@ -339,6 +358,8 @@ func TestTruncateGapHandling(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
 	st := cluster.MakeTestingClusterSettings()
 
 	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
@@ -372,7 +393,7 @@ func TestTruncateGapHandling(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
-			truncator := NewWAGTruncator(st, e.Engines)
+			truncator := NewWAGTruncator(stopper, st, e.Engines, &e.seq)
 
 			// Write WAG nodes at indices 2, 4, 6.
 			e.seq.Next()
@@ -412,3 +433,78 @@ 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) // must run before e.Close() (LIFO)
+	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(stopper, st, e.Engines, &e.seq)
+	require.NoError(t, truncator.Start(ctx))
+
+	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)
+	require.Equal(t, truncator.lastTruncatedWAGIndex.Load(), uint64(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)
+	require.Equal(t, truncator.lastTruncatedWAGIndex.Load(), uint64(3))
+}