asim: account for follower replica load

Previously, asim only accounted for load on the leaseholder, ignoring
non-leaseholder replicas. This commit updates it to consider all replicas for
RangeUsageInfo and store capacity aggregation. RangeUsageInfo handles
leaseholder checks and clears request CPU and QPS stats for non-leaseholder
replicas.

Epic: none
Release note: none

Author: wenyihu6

pkg/kv/kvserver/asim/state/impl.go

@@ -279,21 +279,18 @@ func (s *state) capacity(storeID StoreID) roachpb.StoreCapacity {
 
 	for _, repl := range s.Replicas(storeID) {
 		rangeID := repl.Range()
-		replicaID := repl.ReplicaID()
 		rng, _ := s.Range(rangeID)
-		if rng.Leaseholder() == replicaID {
-			// TODO(kvoli): We currently only consider load on the leaseholder
-			// replica for a range. The other replicas have an estimate that is
-			// calculated within the allocation algorithm. Adapt this to
-			// support follower reads, when added to the workload generator.
-			usage := s.RangeUsageInfo(rng.RangeID(), storeID)
-			capacity.QueriesPerSecond += usage.QueriesPerSecond
-			capacity.WritesPerSecond += usage.WritesPerSecond
-			capacity.LogicalBytes += usage.LogicalBytes
-			capacity.CPUPerSecond += usage.RequestCPUNanosPerSecond + usage.RaftCPUNanosPerSecond
+		usage := s.RangeUsageInfo(rng.RangeID(), storeID)
+		// RangeUsageInfo should return valid usage depending on leaseholder. No
+		// special handling needed here.
+		capacity.QueriesPerSecond += usage.QueriesPerSecond
+		capacity.WritesPerSecond += usage.WritesPerSecond
+		capacity.LogicalBytes += usage.LogicalBytes
+		capacity.CPUPerSecond += usage.RequestCPUNanosPerSecond + usage.RaftCPUNanosPerSecond
+		capacity.RangeCount++
+		if leaseholder, _ := s.LeaseholderStore(rangeID); leaseholder.StoreID() == storeID {
 			capacity.LeaseCount++
 		}
-		capacity.RangeCount++
 	}
 
 	// TODO(kvoli): parameterize the logical to actual used storage bytes. At the
@@ -1025,30 +1022,36 @@ func (s *state) applyLoad(rng *rng, le workload.LoadEvent) {
 	s.loadsplits[store.StoreID()].Record(s.clock.Now(), rng.rangeID, le)
 }
 
-// ReplicaLoad returns the usage information for the Range with ID
-// RangeID on the store with ID StoreID.
+// RangeUsageInfo returns the usage information for the Range with ID RangeID on
+// the store with ID StoreID. If the given store has a replica for the range,
+// this returns the write-bytes-per-second, raft cpu, written keys and logical
+// bytes. If the given store is the leaseholder for the range, then the request
+// cpu and qps is also returned. If the given store does not have a replica for
+// the range, this function panics.
 func (s *state) RangeUsageInfo(rangeID RangeID, storeID StoreID) allocator.RangeUsageInfo {
-	// NB: we only return the actual replica load, if the range leaseholder is
-	// currently on the store given. Otherwise, return an empty, zero counter
-	// value.
-	// TODO(wenyihu6): Track non-leaseholder replica load as well.
-	store, ok := s.LeaseholderStore(rangeID)
+	r, ok := s.Range(rangeID)
 	if !ok {
 		panic(fmt.Sprintf("no leaseholder store found for range %d", storeID))
 	}
 
-	r, _ := s.Range(rangeID)
-	// TODO(kvoli): The requested storeID is not the leaseholder. Non
-	// leaseholder load tracking is not currently supported but is checked by
-	// other components such as hot ranges. In this case, ignore it but we
-	// should also track non leaseholder load. See load.go for more. Return an
-	// empty initialized load counter here.
-	if store.StoreID() != storeID {
-		return allocator.RangeUsageInfo{LogicalBytes: r.Size()}
+	if _, ok = r.Replica(storeID); !ok {
+		panic(fmt.Sprintf("no replica found for range %v on store %v [replicas=%v]",
+			rangeID, storeID, r.Replicas()))
 	}
 
 	usage := s.load[rangeID].Load()
 	usage.LogicalBytes = r.Size()
+	leaseholderStore, ok := s.LeaseholderStore(rangeID)
+	if !ok {
+		panic(fmt.Sprintf("no leaseholder store found for range %v", rangeID))
+	}
+
+	if leaseholderStore.StoreID() != storeID {
+		// See the method comment, we don't include the request cpu or qps for
+		// non-leaseholder replicas.
+		usage.RequestCPUNanosPerSecond = 0
+		usage.QueriesPerSecond = 0
+	}
 	return usage
 }
 

pkg/kv/kvserver/asim/state/state_test.go

@@ -287,23 +287,30 @@ func TestWorkloadApply(t *testing.T) {
 	require.Equal(t, expectedLoad, sc3)
 }
 
-// TestReplicaLoadQPS asserts that the rated replica load accounting maintains
-// the average per second corresponding to the tick clock.
-func TestReplicaLoadQPS(t *testing.T) {
+// TestReplicaLoadRangeUsageInfo asserts that the rated replica load accounting
+// maintains the average per second corresponding to the tick clock.
+func TestReplicaLoadRangeUsageInfo(t *testing.T) {
 	settings := config.DefaultSimulationSettings()
 	s := NewState(settings)
 	start := settings.StartTime
 
 	n1 := s.AddNode()
+	n2 := s.AddNode()
+	n3 := s.AddNode()
 	k1 := Key(100)
 	qps := 1000
 	s1, _ := s.AddStore(n1.NodeID())
+	s2, _ := s.AddStore(n2.NodeID())
+	s3, _ := s.AddStore(n3.NodeID())
 	_, r1, _ := s.SplitRange(k1)
 	s.AddReplica(r1.RangeID(), s1.StoreID(), roachpb.VOTER_FULL)
+	s.AddReplica(r1.RangeID(), s2.StoreID(), roachpb.VOTER_FULL)
+	s.AddReplica(r1.RangeID(), s3.StoreID(), roachpb.VOTER_FULL)
 
 	applyLoadToStats := func(key int64, count int) {
 		for i := 0; i < count; i++ {
-			s.ApplyLoad(workload.LoadBatch{workload.LoadEvent{Key: key, Writes: 1}})
+			s.ApplyLoad(workload.LoadBatch{workload.LoadEvent{
+				Key: key, Writes: 1, WriteSize: 1, RequestCPU: 1, RaftCPU: 1}})
 		}
 	}
 
@@ -317,6 +324,24 @@ func TestReplicaLoadQPS(t *testing.T) {
 	// Assert that the rated avg comes out to rate of queries applied per
 	// second.
 	require.Equal(t, float64(qps), s.RangeUsageInfo(r1.RangeID(), s1.StoreID()).QueriesPerSecond)
+	// The non-leaseholder replicas should have no QPS.
+	require.Equal(t, float64(0), s.RangeUsageInfo(r1.RangeID(), s2.StoreID()).QueriesPerSecond)
+	require.Equal(t, float64(0), s.RangeUsageInfo(r1.RangeID(), s2.StoreID()).QueriesPerSecond)
+	// Similarly, only the leaseholder should have request CPU recorded.
+	require.Equal(t, float64(qps),
+		s.RangeUsageInfo(r1.RangeID(), s1.StoreID()).RequestCPUNanosPerSecond)
+	require.Equal(t, float64(0),
+		s.RangeUsageInfo(r1.RangeID(), s2.StoreID()).RequestCPUNanosPerSecond)
+	require.Equal(t, float64(0),
+		s.RangeUsageInfo(r1.RangeID(), s3.StoreID()).RequestCPUNanosPerSecond)
+	// All the replicas should have identical write bytes/s, which is equal to
+	// the QPS. The replicas should also have identical raft CPU usage.
+	require.Equal(t, float64(qps),
+		s.RangeUsageInfo(r1.RangeID(), s1.StoreID()).RaftCPUNanosPerSecond)
+	require.Equal(t, float64(qps),
+		s.RangeUsageInfo(r1.RangeID(), s2.StoreID()).RaftCPUNanosPerSecond)
+	require.Equal(t, float64(qps),
+		s.RangeUsageInfo(r1.RangeID(), s3.StoreID()).RaftCPUNanosPerSecond)
 }
 
 // TestKeyTranslation asserts that key encoding between roachpb keys and

pkg/kv/kvserver/asim/tests/testdata/non_rand/example_fulldisk

@@ -19,46 +19,46 @@ OK
 # store 5. This is shown below as it sheds replicas.
 plot stat=replicas
 ----
- 342 ┤                                                                    ╭╮╭─╮╭╭╭─╮╭
- 331 ┤                         ╭╮ ╭╭─╮╮╭╮╭╭╮╭╭───╮─╭──╮╮  ╭╮──╭╭───╮─╮╭──╮╯╭╯╭╰╮╭╯│╰─
- 319 ┤                  ╭╭╮╭╮─╭─╭──╯╭│╭╭─────╯╯╰─│╭╯╮╭│╰─╭╮╭───╯╭─╯│╭─╯╭─╰─╯╰╯─╰╯╭╰╰╯
- 308 ┤               ╭───╭──────╯╯╰╰╯╰─╯╰╯  ╰╯   ╰╯ ╰╯╰──╯╰╯  ╰╰╯╯ ╰╯ ╰╯        ╰╯
- 296 ┼───────────────╮───╯╯╯╰╯
- 285 ┤               ╰──╮
- 274 ┤                  │╭─╮
- 262 ┤                  ╰╯ │╭╮
- 251 ┤                     ╰╯╰╮ ╭╮
- 239 ┤                        │╭╯│                     ╭─╮
- 228 ┤                        ╰╯ ╰╮╭╮   ╭╮ ╭─╮         │ │ ╭╮  ╭╮
- 217 ┤                            ╰╯╰───╯│╭╯ │   ╭╮ ╭╮╭╯ ╰─╯╰╮╭╯│     ╭╮
- 205 ┤                                   ╰╯  ╰╮╭─╯╰─╯╰╯      ╰╯ ╰╮ ╭╮ ││
- 194 ┤                                        ││                 ╰─╯╰─╯╰╮╭╮ ╭╮  ╭╮ ╭╮
- 182 ┤                                        ╰╯                        ╰╯╰╮│╰╮╭╯│ │╰
- 171 ┤                                                                     ╰╯ ╰╯ ╰─╯
+ 364 ┤       ╭╭╮╭─╮╮╭╭╮─╭───╮╭─╮╮╭─╭╮───╭╮─╭╭─╮─╭────╮╭╭───────╮──╭─╮─╭──────────────
+ 344 ┤    ╭───╯╭─────╯╰─╯───╰──────╯╰───╯╰──╯╯╰─╯╰╰──╰─╯╯╰────╯╰──╯╯╰─╯╯  ╰───╯ ╰────
+ 324 ┤  ╭─╯╭───╯
+ 304 ┼─╮───╯
+ 283 ┤ ╰╮
+ 263 ┤  │
+ 243 ┤  ╰╮
+ 223 ┤   ╰╮
+ 203 ┤    ╰╮
+ 183 ┤     ╰╮
+ 163 ┤      ╰╮
+ 143 ┤       │
+ 122 ┤       ╰╮
+ 102 ┤        ╰────────╮╭──╮╭─╮╭──╮ ╭──╮ ╭──╮ ╭╮
+  82 ┤                 ╰╯  ╰╯ ╰╯  ╰─╯  ╰─╯  ╰─╯╰──────────────╮╭───╮╭───╮╭────╮ ╭──╮
+  62 ┤                                                        ╰╯   ╰╯   ╰╯    ╰─╯  ╰─
                                           replicas
 initial store values: [s1=300, s2=300, s3=300, s4=300, s5=300] (stddev=0.00, mean=300.00, sum=1500)
-last store values: [s1=342, s2=341, s3=326, s4=328, s5=178] (stddev=62.84, mean=303.00, sum=1515)
+last store values: [s1=363, s2=352, s3=360, s4=362, s5=66] (stddev=117.36, mean=300.60, sum=1503)
 
 # Plot the % of disk storage capacity used. We should see s5 hovering right
 # around 92.5-95% (the storage capacity threshold value).
 plot stat=disk_fraction_used
 ----
- 1.10 ┤                                     ╭╮          ╭╮             ╭╮           ╭╮
- 1.04 ┤                      ╭╮  ╭╮         │╰╮   ╭╮    ││  ╭╮  ╭╮     ││ ╭╮ ╭╮ ╭─╮ │╰
- 0.97 ┤           ╭───╮╭─╮ ╭╮│╰╮╭╯╰╮╭╮╭╮ ╭╮ │ ╰╮╭╮││ ╭╮ │╰╮╭╯╰╮╭╯╰╮ ╭╮ │╰╮││ │╰╮│ │╭╯
- 0.91 ┤    ╭──────╯   ╰╯ ╰─╯╰╯ ╰╯  ╰╯╰╯╰─╯│╭╯  ││╰╯╰─╯│╭╯ ╰╯  ╰╯  ╰─╯╰─╯ ╰╯╰─╯ ╰╯ ╰╯
- 0.84 ┼────╯                              ╰╯   ╰╯     ╰╯
- 0.78 ┤
- 0.71 ┤
- 0.65 ┤
- 0.59 ┤
- 0.52 ┤
- 0.46 ┤
- 0.39 ┤
- 0.33 ┤
- 0.26 ┤                                               ╭╮  ╭╮╮╭─╭╭────╮╭───────────────
- 0.20 ┤               ╭╮╭╭────────────────────────────╯╰──╯╰────╯─╯╰─╰╯──╯
- 0.14 ┼──────────────────╯ ╰╯╰╯
+ 2.51 ┼─╮
+ 2.37 ┤ ╰╮
+ 2.23 ┤  │
+ 2.09 ┤  ╰╮
+ 1.95 ┤   ╰╮
+ 1.81 ┤    │
+ 1.67 ┤    ╰╮
+ 1.53 ┤     ╰╮
+ 1.39 ┤      ╰╮
+ 1.25 ┤       │                           ╭╮                         ╭╮    ╭─╮
+ 1.11 ┤       ╰╮             ╭╮  ╭╮  ╭──╮ │╰╮  ╭╮  ╭╮ ╭─╮ ╭──╮   ╭╮  │╰─╮  │ ╰╮  ╭─╮
+ 0.97 ┤        ╰╮╭───╮╭─╮╭──╮│╰╮╭╯╰─╮│  ╰╮│ ╰─╮│╰──╯╰╮│ ╰─╯  ╰─╮╭╯╰─╮│  ╰╮╭╯  ╰─╮│ ╰──
+ 0.83 ┤         ╰╯   ╰╯ ╰╯  ╰╯ ╰╯   ╰╯   ╰╯   ╰╯     ╰╯        ╰╯╭─╭╰╯───╰╯─────╰╯────
+ 0.69 ┤                              ╭╭─╭─╮╭───────────────────────╯╯╰─╯   ╰╯
+ 0.55 ┤     ╭──╭────────────────────────╯ ╰╯
+ 0.41 ┼────────╯
                                       disk_fraction_used
-initial store values: [s1=0.14, s2=0.14, s3=0.14, s4=0.14, s5=0.83] (stddev=0.28, mean=0.28, sum=1)
-last store values: [s1=0.28, s2=0.26, s3=0.25, s4=0.25, s5=1.00] (stddev=0.30, mean=0.41, sum=2)
+initial store values: [s1=0.41, s2=0.41, s3=0.41, s4=0.41, s5=2.50] (stddev=0.84, mean=0.83, sum=4)
+last store values: [s1=0.87, s2=0.85, s3=0.87, s4=0.87, s5=0.97] (stddev=0.04, mean=0.89, sum=4)

pkg/kv/kvserver/asim/tests/testdata/non_rand/example_multi_store

@@ -23,22 +23,22 @@ OK
 
 plot stat=leases
 ----
- 14.00 ┼─╮
- 13.07 ┤ │
- 12.13 ┤ ╰╮
- 11.20 ┤  │
- 10.27 ┤  ╰╮
-  9.33 ┤   │
+ 14.00 ┼╮
+ 13.07 ┤╰╮
+ 12.13 ┤ │
+ 11.20 ┤ │
+ 10.27 ┤ ╰╮
+  9.33 ┤  ╰╮
   8.40 ┤   │
-  7.47 ┤   │
-  6.53 ┤   │
-  5.60 ┤   ╰╮
-  4.67 ┤    │
-  3.73 ┤    ╰╮
-  2.80 ┤   ╭────────────╭───────────────╮
-  1.87 ┤   ╭────────────╮──────────────╮│───────────────╮
-  0.93 ┤ ╭─╯────────────╭──────────────────────────────────────────────────────────────
-  0.00 ┼────────────────╯───────────────────────────────╯
+  7.47 ┤   ╰╮
+  6.53 ┤    ╰╮
+  5.60 ┤     │
+  4.67 ┤     ╰──────────╮
+  3.73 ┤                │
+  2.80 ┤                │
+  1.87 ┤     ╭──────────╮───────────────╮
+  0.93 ┤╭╭──────────────╭──────────────────────────────────────────────────────────────
+  0.00 ┼────────────────╯───────────────╯
                                              leases
 initial store values: [s1=14, s2=0, s3=0, s4=0, s5=0, s6=0, s7=0, s8=0, s9=0, s10=0, s11=0, s12=0, s13=0, s14=0] (stddev=3.61, mean=1.00, sum=14)
 last store values: [s1=1, s2=1, s3=1, s4=1, s5=1, s6=1, s7=1, s8=1, s9=1, s10=1, s11=1, s12=1, s13=1, s14=1] (stddev=0.00, mean=1.00, sum=14)