Skip to content

Commit 7a0d830

Browse files
ClamChowderTiDBClamChowderTiDB
committed
client: pre-encode request before sending to batch-send-loop (tikv#1841) …
X-Original-Commit: 567951320485b0cf1136a9f6a5b9086c8e675172
1 parent 43c1391 commit 7a0d830

3 files changed

Lines changed: 304 additions & 1 deletion

File tree

‎internal/client/client_batch.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,79 @@ import (
6464
"google.golang.org/grpc/metadata"
6565
)
6666

67+
// globalEncodedMsgDataPool is used to pool pre-encoded message data for batch commands.
68+
var globalEncodedMsgDataPool = grpc.NewSharedBufferPool()
69+
70+
type encodedBatchCmd struct {
71+
// implement isBatchCommandsRequest_Request_Cmd
72+
tikvpb.BatchCommandsRequest_Request_Empty
73+
// pre-encoded message data
74+
data *[]byte
75+
}
76+
77+
func (p *encodedBatchCmd) MarshalTo(data []byte) (int, error) {
78+
if p.data == nil {
79+
return 0, errors.New("message data has been released")
80+
} else if len(data) < len(*p.data) {
81+
return 0, errors.New("no enough space to marshal message")
82+
}
83+
n := copy(data, *p.data)
84+
return n, nil
85+
}
86+
87+
func (p *encodedBatchCmd) Size() int {
88+
if p.data == nil {
89+
return 0
90+
}
91+
return len(*p.data)
92+
}
93+
94+
// encodeRequestCmd encodes the `req.Cmd` into a `encodedBatchCmd` and updates the `req.Cmd` to it in place.
95+
//
96+
// SAFETY: This function is called just after `tikvrpc.Request.ToBatchCommandsRequest()` and before constructing
97+
// `batchCommandsEntry`. Note that the `req` argument is the output of `ToBatchCommandsRequest`, which is a new object
98+
// of `tikvpb.BatchCommandsRequest_Request`. So `req.Cmd` can be only modified inside the batch client by
99+
// `reuseRequestData` after the `req` has been sent. So there should be no data race on `req.Cmd` via public API
100+
// (`SendRequest` and `SendRequestAsync`). When writing unit tests, please make sure it's only called once for one `req`
101+
// (do NOT reuse `req`).
102+
func encodeRequestCmd(req *tikvpb.BatchCommandsRequest_Request) error {
103+
if req.Cmd == nil {
104+
// req.Cmd might be nil in unit tests.
105+
return nil
106+
}
107+
if encoded, ok := req.Cmd.(*encodedBatchCmd); ok {
108+
if encoded.data == nil {
109+
return errors.New("request has already been encoded and sent")
110+
}
111+
return nil
112+
}
113+
data := globalEncodedMsgDataPool.Get(req.Cmd.Size())
114+
n, err := req.Cmd.MarshalTo(data)
115+
if err != nil {
116+
globalEncodedMsgDataPool.Put(&data)
117+
return errors.WithStack(err)
118+
} else if n != len(data) {
119+
globalEncodedMsgDataPool.Put(&data)
120+
return errors.Errorf("unexpected marshaled size: got %d, want %d", n, len(data))
121+
}
122+
req.Cmd = &encodedBatchCmd{data: &data}
123+
return nil
124+
}
125+
126+
// reuseRequestData puts back all pre-encoded message data in the request to the `globalEncodedMsgDataPool`. The
127+
// returned count is used for testing.
128+
func reuseRequestData(req *tikvpb.BatchCommandsRequest) int {
129+
count := 0
130+
for _, r := range req.Requests {
131+
if cmd, ok := r.Cmd.(*encodedBatchCmd); ok && cmd.data != nil {
132+
globalEncodedMsgDataPool.Put(cmd.data)
133+
cmd.data = nil
134+
count++
135+
}
136+
}
137+
return count
138+
}
139+
67140
type batchCommandsEntry struct {
68141
ctx context.Context
69142
req *tikvpb.BatchCommandsRequest_Request
@@ -784,6 +857,8 @@ func (c *batchCommandsClient) available() int64 {
784857
}
785858

786859
func (c *batchCommandsClient) send(forwardedHost string, req *tikvpb.BatchCommandsRequest) {
860+
defer reuseRequestData(req)
861+
787862
err := c.initBatchClient(forwardedHost)
788863
if err != nil {
789864
logutil.BgLogger().Warn(
@@ -1102,6 +1177,9 @@ func sendBatchRequest(
11021177
timeout time.Duration,
11031178
priority uint64,
11041179
) (*tikvrpc.Response, error) {
1180+
if err := encodeRequestCmd(req); err != nil {
1181+
return nil, err
1182+
}
11051183
entry := &batchCommandsEntry{
11061184
ctx: ctx,
11071185
req: req,

‎internal/client/client_batch_test.go‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
// Copyright 2026 TiKV Authors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package client
16+
17+
import (
18+
"sync"
19+
"sync/atomic"
20+
"testing"
21+
22+
"github.com/pingcap/kvproto/pkg/kvrpcpb"
23+
"github.com/pingcap/kvproto/pkg/tikvpb"
24+
"github.com/stretchr/testify/require"
25+
)
26+
27+
func TestEncodedBatchCmd_SizeAndMarshalTo(t *testing.T) {
28+
data := []byte{0x01, 0x02, 0x03, 0x04, 0x05}
29+
cmd := &encodedBatchCmd{data: &data}
30+
31+
// Verify Size returns correct length
32+
require.Equal(t, 5, cmd.Size())
33+
34+
// Verify MarshalTo copies data correctly
35+
buf := make([]byte, 5)
36+
n, err := cmd.MarshalTo(buf)
37+
require.NoError(t, err)
38+
require.Equal(t, 5, n)
39+
require.Equal(t, data, buf)
40+
}
41+
42+
func TestEncodeRequestCmd_Basic(t *testing.T) {
43+
cmd := &tikvpb.BatchCommandsRequest_Request_Get{
44+
Get: &kvrpcpb.GetRequest{
45+
Key: []byte("test-key"),
46+
},
47+
}
48+
batch1 := &tikvpb.BatchCommandsRequest{Requests: []*tikvpb.BatchCommandsRequest_Request{{Cmd: cmd}}}
49+
batch2 := &tikvpb.BatchCommandsRequest{Requests: []*tikvpb.BatchCommandsRequest_Request{{Cmd: cmd}}}
50+
51+
err := encodeRequestCmd(batch1.Requests[0])
52+
require.NoError(t, err)
53+
54+
// Verify req.Cmd is now a encodedBatchCmd
55+
_, ok := batch1.Requests[0].Cmd.(*encodedBatchCmd)
56+
require.True(t, ok, "req.Cmd should be converted to *encodedBatchCmd")
57+
58+
// Verify two batches have the same size after encoding
59+
require.Equal(t, batch1.Size(), batch2.Size())
60+
61+
// Verify marshaled data matches
62+
data1 := make([]byte, batch1.Size())
63+
n1, err := batch1.MarshalTo(data1)
64+
require.NoError(t, err)
65+
require.Equal(t, batch1.Size(), n1)
66+
67+
data2 := make([]byte, batch2.Size())
68+
n2, err := batch2.MarshalTo(data2)
69+
require.NoError(t, err)
70+
require.Equal(t, batch2.Size(), n2)
71+
72+
require.Equal(t, data1, data2)
73+
}
74+
75+
func TestEncodeRequestCmd_PoolReuse(t *testing.T) {
76+
// Encode cmd of an Empty request
77+
req1 := &tikvpb.BatchCommandsRequest_Request{
78+
Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{
79+
Empty: &tikvpb.BatchCommandsEmptyRequest{},
80+
},
81+
}
82+
err := encodeRequestCmd(req1)
83+
require.NoError(t, err)
84+
85+
// Simulate send/reuse cycle
86+
batchReq := &tikvpb.BatchCommandsRequest{Requests: []*tikvpb.BatchCommandsRequest_Request{req1}}
87+
reuseRequestData(batchReq)
88+
89+
// Encode cmd of a Get request
90+
req2 := &tikvpb.BatchCommandsRequest_Request{
91+
Cmd: &tikvpb.BatchCommandsRequest_Request_Get{
92+
Get: &kvrpcpb.GetRequest{
93+
Key: []byte("test-key"),
94+
},
95+
},
96+
}
97+
err = encodeRequestCmd(req2)
98+
require.NoError(t, err)
99+
}
100+
101+
func TestEncodeRequestCmd_AfterPoolReturn(t *testing.T) {
102+
req := &tikvpb.BatchCommandsRequest_Request{
103+
Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{
104+
Empty: &tikvpb.BatchCommandsEmptyRequest{},
105+
},
106+
}
107+
108+
// First encode - should succeed
109+
err := encodeRequestCmd(req)
110+
require.NoError(t, err)
111+
112+
// Simulate send: marshal and return to pool
113+
batch := &tikvpb.BatchCommandsRequest{Requests: []*tikvpb.BatchCommandsRequest_Request{req}}
114+
_, err = batch.Marshal()
115+
require.NoError(t, err)
116+
reuseRequestData(batch)
117+
118+
// Second encode on same request - should fail
119+
err = encodeRequestCmd(req)
120+
require.Error(t, err)
121+
require.Contains(t, err.Error(), "already been encoded and sent")
122+
}
123+
124+
func TestReuseRequestData_Basic(t *testing.T) {
125+
126+
// Normal request
127+
normalReq := &tikvpb.BatchCommandsRequest_Request{
128+
Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{
129+
Empty: &tikvpb.BatchCommandsEmptyRequest{},
130+
},
131+
}
132+
133+
// Encoded request
134+
encodedReq := &tikvpb.BatchCommandsRequest_Request{
135+
Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{
136+
Empty: &tikvpb.BatchCommandsEmptyRequest{},
137+
},
138+
}
139+
err := encodeRequestCmd(encodedReq)
140+
require.NoError(t, err)
141+
142+
// Batch with mixed requests
143+
batch := &tikvpb.BatchCommandsRequest{
144+
Requests: []*tikvpb.BatchCommandsRequest_Request{normalReq, encodedReq},
145+
}
146+
147+
// reuseRequestData should only affect the encoded request
148+
count := reuseRequestData(batch)
149+
require.Equal(t, 1, count)
150+
_, ok := normalReq.Cmd.(*encodedBatchCmd)
151+
require.False(t, ok, "Normal request should not be converted")
152+
_, ok = encodedReq.Cmd.(*encodedBatchCmd)
153+
require.True(t, ok, "Encoded request should remain encoded")
154+
_, err = encodedReq.Marshal()
155+
require.Error(t, err, "Marshaling released data should fail")
156+
}
157+
158+
func TestReuseRequestData_DoubleReturn(t *testing.T) {
159+
req := &tikvpb.BatchCommandsRequest_Request{
160+
Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{
161+
Empty: &tikvpb.BatchCommandsEmptyRequest{},
162+
},
163+
}
164+
165+
// Encode the request
166+
err := encodeRequestCmd(req)
167+
require.NoError(t, err)
168+
169+
// First reuse - should succeed and return 1
170+
batch := &tikvpb.BatchCommandsRequest{Requests: []*tikvpb.BatchCommandsRequest_Request{req}}
171+
count := reuseRequestData(batch)
172+
require.Equal(t, 1, count)
173+
174+
// Verify data is now nil
175+
encodedCmd := req.Cmd.(*encodedBatchCmd)
176+
require.Nil(t, encodedCmd.data)
177+
178+
// Second reuse - should return 0 (nil data is not returned to pool)
179+
count = reuseRequestData(batch)
180+
require.Equal(t, 0, count)
181+
}
182+
183+
func TestEncodedMsgDataPool_ConcurrentSafety(t *testing.T) {
184+
var wg sync.WaitGroup
185+
errors := make([]error, 50)
186+
successCount := int32(0)
187+
188+
for i := 0; i < 50; i++ {
189+
wg.Add(1)
190+
go func(idx int) {
191+
defer wg.Done()
192+
for j := 0; j < 100; j++ {
193+
req := &tikvpb.BatchCommandsRequest_Request{
194+
Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{
195+
Empty: &tikvpb.BatchCommandsEmptyRequest{},
196+
},
197+
}
198+
err := encodeRequestCmd(req)
199+
if err != nil {
200+
errors[idx] = err
201+
return
202+
}
203+
204+
atomic.AddInt32(&successCount, 1)
205+
batch := &tikvpb.BatchCommandsRequest{Requests: []*tikvpb.BatchCommandsRequest_Request{req}}
206+
_, err = batch.Marshal()
207+
if err != nil {
208+
errors[idx] = err
209+
return
210+
}
211+
reuseRequestData(batch)
212+
}
213+
}(i)
214+
}
215+
216+
wg.Wait()
217+
for _, err := range errors {
218+
require.NoError(t, err)
219+
}
220+
require.Equal(t, int32(5000), successCount)
221+
}

‎internal/client/client_test.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -670,11 +670,11 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
670670
server.Stop()
671671
}()
672672

673-
req := &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: &coprocessor.Request{}}}
674673
conn, err := client.getConnArray(addr, true)
675674
assert.Nil(t, err)
676675
// send some request, it should be success.
677676
for i := 0; i < 100; i++ {
677+
req := &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: &coprocessor.Request{}}}
678678
_, err = sendBatchRequest(context.Background(), addr, "", conn.batchConn, req, time.Second*20, 0)
679679
require.NoError(t, err)
680680
}
@@ -685,6 +685,7 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
685685

686686
// send some request, it should be failed since server is down.
687687
for i := 0; i < 10; i++ {
688+
req := &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: &coprocessor.Request{}}}
688689
_, err = sendBatchRequest(context.Background(), addr, "", conn.batchConn, req, time.Millisecond*100, 0)
689690
require.Error(t, err)
690691
time.Sleep(time.Millisecond * time.Duration(rand.Intn(300)))
@@ -728,6 +729,7 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
728729

729730
// send some request, it should be success again.
730731
for i := 0; i < 100; i++ {
732+
req := &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: &coprocessor.Request{}}}
731733
_, err = sendBatchRequest(context.Background(), addr, "", conn.batchConn, req, time.Second*20, 0)
732734
require.NoError(t, err)
733735
}
@@ -1040,6 +1042,8 @@ func TestFastFailWhenNoAvailableConn(t *testing.T) {
10401042
// mock all client a in recreate.
10411043
c.lockForRecreate()
10421044
}
1045+
1046+
req = &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: &coprocessor.Request{}}}
10431047
start := time.Now()
10441048
timeout := time.Second
10451049
_, err = sendBatchRequest(context.Background(), addr, "", conn.batchConn, req, timeout, 0)

0 commit comments

Comments
 (0)