controller: add replica number (#823)

Signed-off-by: nolouch <nolouch@gmail.com>

internal/mockstore/mocktikv/pd.go

@@ -110,6 +110,23 @@ func (c *pdClient) GetTS(context.Context) (int64, int64, error) {
 	return tsMu.physicalTS, tsMu.logicalTS, nil
 }
 
+// GetMinTS returns the minimal ts.
+func (c *pdClient) GetMinTS(ctx context.Context) (int64, int64, error) {
+	return 0, 0, nil
+}
+
+func (c *pdClient) UpdateGCSafePointV2(ctx context.Context, keyspaceID uint32, safePoint uint64) (uint64, error) {
+	panic("unimplemented")
+}
+
+func (c *pdClient) UpdateServiceSafePointV2(ctx context.Context, keyspaceID uint32, serviceID string, ttl int64, safePoint uint64) (uint64, error) {
+	panic("unimplemented")
+}
+
+func (c *pdClient) WatchGCSafePointV2(ctx context.Context, revision int64) (chan []*pdpb.SafePointEvent, error) {
+	panic("unimplemented")
+}
+
 func (c *pdClient) GetLocalTS(ctx context.Context, dcLocation string) (int64, int64, error) {
 	return c.GetTS(ctx)
 }

internal/resourcecontrol/resource_control.go

@@ -31,8 +31,9 @@ import (
 type RequestInfo struct {
 	// writeBytes is the actual write size if the request is a write request,
 	// or -1 if it's a read request.
-	writeBytes int64
-	storeID    uint64
+	writeBytes    int64
+	storeID       uint64
+	replicaNumber int64
 }
 
 // MakeRequestInfo extracts the relevant information from a BatchRequest.
@@ -56,8 +57,7 @@ func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo {
 			writeBytes += int64(len(k))
 		}
 	}
-
-	return &RequestInfo{writeBytes: writeBytes * req.ReplicaNumber, storeID: req.Context.Peer.StoreId}
+	return &RequestInfo{writeBytes: writeBytes, storeID: req.Context.Peer.StoreId, replicaNumber: req.ReplicaNumber}
 }
 
 // IsWrite returns whether the request is a write request.
@@ -71,6 +71,10 @@ func (req *RequestInfo) WriteBytes() uint64 {
 	return uint64(req.writeBytes)
 }
 
+func (req *RequestInfo) ReplicaNumber() int64 {
+	return req.replicaNumber
+}
+
 func (req *RequestInfo) StoreID() uint64 {
 	return req.storeID
 }