kvserver: record follower write bytes in replica load

Previously, write bytes for range load were only tracked on leaseholder replicas
during proposal, leaving write bytes written to follower replicas unaccounted
for. This commit updates the tracking to include write bytes on follower
replicas as well.

Epic: none
Release note: none

Author: wenyihu6

pkg/kv/kvserver/load/replica_load.go

@@ -65,7 +65,11 @@ type ReplicaLoadStats struct {
 	// follower and leaseholder reads.
 	ReadKeysPerSecond float64
 	// WriteBytesPerSecond is the replica's average bytes written per second. A
-	// "Write" is as described in WritesPerSecond.
+	// "Write" is as described in WritesPerSecond. If the replica is a leaseholder,
+	// this is recorded as the bytes that will be written by the replica during the
+	// application of the Raft command including write bytes and ingested bytes for
+	// AddSSTable requests. If the replica is a follower, this is recorded right
+	// before a command is applied to the state machine.
 	WriteBytesPerSecond float64
 	// ReadBytesPerSecond is the replica's average bytes read per second. A "Read" is as
 	// described in ReadsPerSecond.

pkg/kv/kvserver/replica_app_batch.go

@@ -253,16 +253,20 @@ func (b *replicaAppBatch) runPostAddTriggersReplicaOnly(
 	// We don't track these stats in standalone log application since they depend
 	// on whether the proposer is still waiting locally, and this concept does not
 	// apply in a standalone context.
-	//
-	// TODO(irfansharif): This code block can be removed once below-raft
-	// admission control is the only form of IO admission control. It pre-dates
-	// it -- these stats were previously used to deduct IO tokens for follower
-	// writes/ingests without waiting.
-	if !cmd.IsLocal() && !cmd.ApplyAdmissionControl() {
+	if !cmd.IsLocal() {
 		writeBytes, ingestedBytes := cmd.getStoreWriteByteSizes()
-		b.followerStoreWriteBytes.NumEntries++
-		b.followerStoreWriteBytes.WriteBytes += writeBytes
-		b.followerStoreWriteBytes.IngestedBytes += ingestedBytes
+		if writeBytes > 0 || ingestedBytes > 0 {
+			b.r.recordRequestWriteBytes(writeBytes, ingestedBytes)
+		}
+		// TODO(irfansharif): This code block can be removed once below-raft
+		// admission control is the only form of IO admission control. It pre-dates
+		// it -- these stats were previously used to deduct IO tokens for follower
+		// writes/ingests without waiting.
+		if !cmd.ApplyAdmissionControl() {
+			b.followerStoreWriteBytes.NumEntries++
+			b.followerStoreWriteBytes.WriteBytes += writeBytes
+			b.followerStoreWriteBytes.IngestedBytes += ingestedBytes
+		}
 	}
 
 	// MVCC history mutations violate the closed timestamp, modifying data that

pkg/kv/kvserver/replica_raft.go

@@ -96,6 +96,8 @@ var ReplicaLeaderlessUnavailableThreshold = settings.RegisterDurationSettingWith
 //     terminate execution, although it is given no guarantee that the proposal
 //     won't still go on to commit and apply at some later time.
 //   - the proposal's ID.
+//   - the bytes that will be written by the replica during the application of
+//     the Raft command.
 //   - any error obtained during the creation or proposal of the command, in
 //     which case the other returned values are zero.
 func (r *Replica) evalAndPropose(

pkg/kv/kvserver/replica_rankings_test.go

@@ -218,11 +218,11 @@ func TestWriteLoadStatsAccounting(t *testing.T) {
 		ReplicationMode: base.ReplicationManual,
 	}
 	args.ServerArgs.Knobs.Store = &StoreTestingKnobs{DisableCanAckBeforeApplication: true}
-	tc := serverutils.StartCluster(t, 1, args)
+	tc := serverutils.StartCluster(t, 3, args)
+	defer tc.Stopper().Stop(ctx)
 
 	const epsilonAllowed = 5
 
-	defer tc.Stopper().Stop(ctx)
 	ts := tc.Server(0)
 	db := ts.DB()
 	conn := tc.ServerConn(0)
@@ -246,52 +246,82 @@ func TestWriteLoadStatsAccounting(t *testing.T) {
 		{1234, 1234, 1234, 0, 1234 * writeSize, 0},
 	}
 
-	store, err := ts.GetStores().(*Stores).GetStore(ts.GetFirstStoreID())
+	tc.AddVotersOrFatal(t, scratchKey, tc.Target(1), tc.Target(2))
+	s1, err := tc.Server(0).GetStores().(*Stores).GetStore(tc.Server(0).GetFirstStoreID())
 	require.NoError(t, err)
-
-	repl := store.LookupReplica(roachpb.RKey(scratchKey))
-	require.NotNil(t, repl)
+	lhRepl := s1.LookupReplica(roachpb.RKey(scratchKey))
+	require.NotNil(t, lhRepl)
+	s2, err := tc.Server(1).GetStores().(*Stores).GetStore(tc.Server(1).GetFirstStoreID())
+	require.NoError(t, err)
+	followerRepl1 := s2.LookupReplica(roachpb.RKey(scratchKey))
+	require.NotNil(t, followerRepl1)
+	s3, err := tc.Server(2).GetStores().(*Stores).GetStore(tc.Server(2).GetFirstStoreID())
+	require.NoError(t, err)
+	followerRepl2 := s3.LookupReplica(roachpb.RKey(scratchKey))
+	require.NotNil(t, followerRepl2)
 
 	// Disable the consistency checker, to avoid interleaving requests
-	// artificially inflating measurement due to consistency checking.
+	// artificially inflating measurement due to consistency checking. Also,
+	// disable the mvcc gc queue and raft log queue to avoid them issuing
+	// interleaving requests as well.
 	sqlDB.Exec(t, `SET CLUSTER SETTING server.consistency_check.interval = '0'`)
 	sqlDB.Exec(t, `SET CLUSTER SETTING kv.range_split.by_load.enabled = false`)
+	sqlDB.Exec(t, `SET CLUSTER SETTING kv.mvcc_gc_queue.enabled = false`)
+	sqlDB.Exec(t, `SET CLUSTER SETTING kv.raft_log_queue.enabled = false`)
 
 	for _, testCase := range testCases {
 		// Reset the request counts to 0 before sending to clear previous requests.
-		repl.loadStats.Reset()
-
-		requestsBefore := repl.loadStats.TestingGetSum(load.Requests)
-		writesBefore := repl.loadStats.TestingGetSum(load.WriteKeys)
-		readsBefore := repl.loadStats.TestingGetSum(load.ReadKeys)
-		readBytesBefore := repl.loadStats.TestingGetSum(load.ReadBytes)
-		writeBytesBefore := repl.loadStats.TestingGetSum(load.WriteBytes)
+		lhRepl.loadStats.Reset()
+		followerRepl1.loadStats.Reset()
+		followerRepl2.loadStats.Reset()
+
+		requestsBefore := lhRepl.loadStats.TestingGetSum(load.Requests)
+		readsBefore := lhRepl.loadStats.TestingGetSum(load.ReadKeys)
+		lhWritesBefore := lhRepl.loadStats.TestingGetSum(load.WriteKeys)
+		readBytesBefore := lhRepl.loadStats.TestingGetSum(load.ReadBytes)
+		lhWriteBytesBefore := lhRepl.loadStats.TestingGetSum(load.WriteBytes)
+		follower1WriteBytesBefore := followerRepl1.loadStats.TestingGetSum(load.WriteBytes)
+		follower2WriteBytesBefore := followerRepl2.loadStats.TestingGetSum(load.WriteBytes)
 
 		for i := 0; i < testCase.writes; i++ {
 			_, pErr := db.Inc(ctx, scratchKey, 1)
 			require.Nil(t, pErr)
 		}
 		require.Equal(t, 0.0, requestsBefore)
-		require.Equal(t, 0.0, writesBefore)
+		require.Equal(t, 0.0, lhWritesBefore)
 		require.Equal(t, 0.0, readsBefore)
-		require.Equal(t, 0.0, writeBytesBefore)
+		require.Equal(t, 0.0, lhWriteBytesBefore)
 		require.Equal(t, 0.0, readBytesBefore)
-
-		requestsAfter := repl.loadStats.TestingGetSum(load.Requests)
-		writesAfter := repl.loadStats.TestingGetSum(load.WriteKeys)
-		readsAfter := repl.loadStats.TestingGetSum(load.ReadKeys)
-		readBytesAfter := repl.loadStats.TestingGetSum(load.ReadBytes)
-		writeBytesAfter := repl.loadStats.TestingGetSum(load.WriteBytes)
+		require.Equal(t, 0.0, follower1WriteBytesBefore)
+		require.Equal(t, 0.0, follower2WriteBytesBefore)
+
+		require.NoError(t, waitForApplication(
+			ctx,
+			s1.cfg.NodeDialer,
+			lhRepl.GetRangeID(),
+			lhRepl.Desc().Replicas().Descriptors(),
+			lhRepl.GetLeaseAppliedIndex(),
+		))
+
+		requestsAfter := lhRepl.loadStats.TestingGetSum(load.Requests)
+		lhWritesAfter := lhRepl.loadStats.TestingGetSum(load.WriteKeys)
+		readsAfter := lhRepl.loadStats.TestingGetSum(load.ReadKeys)
+		readBytesAfter := lhRepl.loadStats.TestingGetSum(load.ReadBytes)
+		lhWriteBytesAfter := lhRepl.loadStats.TestingGetSum(load.WriteBytes)
+		follower1WriteBytesAfter := followerRepl1.loadStats.TestingGetSum(load.WriteBytes)
+		follower2WriteBytesAfter := followerRepl2.loadStats.TestingGetSum(load.WriteBytes)
 
 		assertGreaterThanInDelta(t, testCase.expectedRQPS, requestsAfter, epsilonAllowed)
-		assertGreaterThanInDelta(t, testCase.expectedWPS, writesAfter, epsilonAllowed)
+		assertGreaterThanInDelta(t, testCase.expectedWPS, lhWritesAfter, epsilonAllowed)
 		assertGreaterThanInDelta(t, testCase.expectedRPS, readsAfter, epsilonAllowed)
 		assertGreaterThanInDelta(t, testCase.expectedRBPS, readBytesAfter, epsilonAllowed)
 		// NB: We assert that the written bytes is greater than the write
 		// batch request size. However the size multiplication factor,
 		// varies between 3 and 5 so we instead assert that it is greater
 		// than the logical bytes.
-		require.GreaterOrEqual(t, writeBytesAfter, testCase.expectedWBPS)
+		require.GreaterOrEqual(t, lhWriteBytesAfter, testCase.expectedWBPS)
+		require.GreaterOrEqual(t, follower1WriteBytesAfter, testCase.expectedWBPS)
+		require.GreaterOrEqual(t, follower2WriteBytesAfter, testCase.expectedWBPS)
 	}
 }