resourcecontrol: fix nil pointer (tikv#900)

Signed-off-by: nolouch <[email protected]>
nolouch · Aug 11, 2023 · 59c6eae · 59c6eae
1 parent e68109d
commit 59c6eae
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/internal/resourcecontrol/resource_control.go b/internal/resourcecontrol/resource_control.go
@@ -50,8 +50,9 @@ func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo {
 			bypass = true
 		}
 	}
+	storeID := req.Context.GetPeer().GetStoreId()
 	if !req.IsTxnWriteRequest() && !req.IsRawWriteRequest() {
-		return &RequestInfo{writeBytes: -1, storeID: req.Context.Peer.StoreId, bypass: bypass}
+		return &RequestInfo{writeBytes: -1, storeID: storeID, bypass: bypass}
 	}
 
 	var writeBytes int64
@@ -69,7 +70,7 @@ func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo {
 			writeBytes += int64(len(k))
 		}
 	}
-	return &RequestInfo{writeBytes: writeBytes, storeID: req.Context.Peer.StoreId, replicaNumber: req.ReplicaNumber, bypass: bypass}
+	return &RequestInfo{writeBytes: writeBytes, storeID: storeID, replicaNumber: req.ReplicaNumber, bypass: bypass}
 }
 
 // IsWrite returns whether the request is a write request.

diff --git a/internal/resourcecontrol/resource_control_test.go b/internal/resourcecontrol/resource_control_test.go
@@ -37,4 +37,12 @@ func TestMakeRequestInfo(t *testing.T) {
 	assert.Equal(t, uint64(3), info.WriteBytes())
 	assert.False(t, info.Bypass())
 	assert.Equal(t, uint64(3), info.StoreID())
+
+	// Test Nil Peer in Context
+	req = &tikvrpc.Request{Type: tikvrpc.CmdCommit, Req: commitReq, ReplicaNumber: 2, Context: kvrpcpb.Context{}}
+	info = MakeRequestInfo(req)
+	assert.True(t, info.IsWrite())
+	assert.Equal(t, uint64(3), info.WriteBytes())
+	assert.False(t, info.Bypass())
+	assert.Equal(t, uint64(0), info.StoreID())
 }