kvserver: pass struct into SetRangeAppliedState

The signature of this method was error-prone and verbose. All the
callers copy the RangeAppliedState from ReplicaState fields, or already
have a fully populated RangeAppliedState.

Epic: none
Release note: none

Author: pav-kv

pkg/kv/kvserver/kvserverpb/state.go

@@ -16,3 +16,18 @@ func (ms *MVCCPersistentStats) ToStats() enginepb.MVCCStats {
 func (ms *MVCCPersistentStats) ToStatsPtr() *enginepb.MVCCStats {
 	return (*enginepb.MVCCStats)(ms)
 }
+
+// ToRangeAppliedState converts the ReplicaState to a RangeAppliedState.
+func (m *ReplicaState) ToRangeAppliedState() RangeAppliedState {
+	// NB: RangeAppliedState materializes in the state machine on every raft
+	// commands batch application, and must be byte-to-byte consistent across
+	// replicas. If you need to change the returned fields here, most likely this
+	// needs to be accompanied by a below-raft migration.
+	return RangeAppliedState{
+		RaftAppliedIndex:     m.RaftAppliedIndex,
+		LeaseAppliedIndex:    m.LeaseAppliedIndex,
+		RangeStats:           MVCCPersistentStats(*m.Stats),
+		RaftClosedTimestamp:  m.RaftClosedTimestamp,
+		RaftAppliedIndexTerm: m.RaftAppliedIndexTerm,
+	}
+}

pkg/kv/kvserver/kvstorage/stateloader.go

@@ -8,7 +8,6 @@ package kvstorage
 import (
 	"context"
 
-	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -128,16 +127,10 @@ func (s StateLoader) Save(
 			return enginepb.MVCCStats{}, err
 		}
 	}
-	if err := s.SetRangeAppliedState(
-		ctx,
-		stateRW,
-		state.RaftAppliedIndex,
-		state.LeaseAppliedIndex,
-		state.RaftAppliedIndexTerm,
-		ms,
-		state.RaftClosedTimestamp,
-		nil,
-	); err != nil {
+	state.Stats = ms // no-op, just an acknowledgement that the stats were updated
+
+	as := state.ToRangeAppliedState()
+	if err := s.SetRangeAppliedState(ctx, stateRW, &as); err != nil {
 		return enginepb.MVCCStats{}, err
 	}
 	return *ms, nil
@@ -216,26 +209,8 @@ func (s StateLoader) LoadMVCCStats(
 // keys. We now deem those keys to be "legacy" because they have been replaced
 // by the range applied state key.
 func (s StateLoader) SetRangeAppliedState(
-	ctx context.Context,
-	stateRW StateRW,
-	appliedIndex kvpb.RaftIndex,
-	leaseAppliedIndex kvpb.LeaseAppliedIndex,
-	appliedIndexTerm kvpb.RaftTerm,
-	newMS *enginepb.MVCCStats,
-	raftClosedTimestamp hlc.Timestamp,
-	asAlloc *kvserverpb.RangeAppliedState, // optional
+	ctx context.Context, stateRW StateRW, as *kvserverpb.RangeAppliedState,
 ) error {
-	if asAlloc == nil {
-		asAlloc = new(kvserverpb.RangeAppliedState)
-	}
-	as := asAlloc
-	*as = kvserverpb.RangeAppliedState{
-		RaftAppliedIndex:     appliedIndex,
-		LeaseAppliedIndex:    leaseAppliedIndex,
-		RangeStats:           kvserverpb.MVCCPersistentStats(*newMS),
-		RaftClosedTimestamp:  raftClosedTimestamp,
-		RaftAppliedIndexTerm: appliedIndexTerm,
-	}
 	// The RangeAppliedStateKey is not included in stats. This is also reflected
 	// in ComputeStats.
 	ms := (*enginepb.MVCCStats)(nil)
@@ -253,10 +228,8 @@ func (s StateLoader) SetMVCCStats(
 	if err != nil {
 		return err
 	}
-	alloc := as // reuse
-	return s.SetRangeAppliedState(
-		ctx, stateRW, as.RaftAppliedIndex, as.LeaseAppliedIndex, as.RaftAppliedIndexTerm, newMS,
-		as.RaftClosedTimestamp, alloc)
+	as.RangeStats = kvserverpb.MVCCPersistentStats(*newMS)
+	return s.SetRangeAppliedState(ctx, stateRW, as)
 }
 
 // SetClosedTimestamp overwrites the closed timestamp.
@@ -267,10 +240,8 @@ func (s StateLoader) SetClosedTimestamp(
 	if err != nil {
 		return err
 	}
-	alloc := as // reuse
-	return s.SetRangeAppliedState(
-		ctx, stateRW, as.RaftAppliedIndex, as.LeaseAppliedIndex, as.RaftAppliedIndexTerm,
-		as.RangeStats.ToStatsPtr(), closedTS, alloc)
+	as.RaftClosedTimestamp = closedTS
+	return s.SetRangeAppliedState(ctx, stateRW, as)
 }
 
 // LoadGCThreshold loads the GC threshold.

pkg/kv/kvserver/raft_log_truncator_test.go

@@ -18,11 +18,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
-	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
-	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
@@ -181,8 +179,10 @@ func (r *replicaTruncatorTest) writeRaftStateToEngine(
 func (r *replicaTruncatorTest) writeRaftAppliedIndex(
 	t *testing.T, eng storage.Engine, raftAppliedIndex kvpb.RaftIndex, flush bool,
 ) {
-	require.NoError(t, r.stateLoader.SetRangeAppliedState(context.Background(), eng,
-		raftAppliedIndex, 0, 0, &enginepb.MVCCStats{}, hlc.Timestamp{}, nil))
+	require.NoError(t, r.stateLoader.SetRangeAppliedState(
+		context.Background(), eng,
+		&kvserverpb.RangeAppliedState{RaftAppliedIndex: raftAppliedIndex},
+	))
 	// Flush to make it satisfy the contract of OnlyReadGuaranteedDurable in
 	// Pebble.
 	if flush {

pkg/kv/kvserver/replica_app_batch.go

@@ -704,16 +704,13 @@ func (b *replicaAppBatch) ApplyToStateMachine(ctx context.Context) error {
 }
 
 // addAppliedStateKeyToBatch adds the applied state key to the application
-// batch's RocksDB batch. This records the highest raft and lease index that
-// have been applied as of this batch. It also records the Range's mvcc stats.
+// batch's Pebble batch. This records the highest raft and lease index that have
+// been applied as of this batch. It also records the Range's MVCC stats.
 func (b *replicaAppBatch) addAppliedStateKeyToBatch(ctx context.Context) error {
-	// Set the range applied state, which includes the last applied raft and
-	// lease index along with the mvcc stats, all in one key.
-	loader := &b.r.raftMu.stateLoader
-	return loader.SetRangeAppliedState(
-		ctx, b.batch, b.state.RaftAppliedIndex, b.state.LeaseAppliedIndex, b.state.RaftAppliedIndexTerm,
-		b.state.Stats, b.state.RaftClosedTimestamp, &b.asAlloc,
-	)
+	// Set the range applied state, which includes the last applied raft and lease
+	// index along with the MVCC stats, all in one key.
+	b.asAlloc = b.state.ToRangeAppliedState()
+	return b.r.raftMu.stateLoader.SetRangeAppliedState(ctx, b.batch, &b.asAlloc)
 }
 
 func (b *replicaAppBatch) recordStatsOnCommit() {