Merge cockroachdb#137108

137108: kv: allocate Batch{Request|Response} header and {requests|responses} together r=nvanbenschoten a=nvanbenschoten

These two commits avoid heap allocations in the common cases by allocating the `BatchRequest` and `BatchResponse` headers and their respective requests or responses together, in a single malloc.

Together, this avoids a pair of heap allocations per KV Batch operation.

```
name                                            old time/op    new time/op    delta
Sysbench/SQL/1node_remote/oltp_read_only-10       3.55ms ±17%    3.54ms ±16%    ~     (p=0.945 n=19+20)
Sysbench/SQL/1node_remote/oltp_write_only-10      2.13ms ±20%    2.13ms ±23%    ~     (p=0.531 n=19+20)
Sysbench/SQL/1node_remote/oltp_read_write-10      5.64ms ± 9%    5.77ms ±24%    ~     (p=0.988 n=18+19)
Sysbench/SQL/1node_remote/oltp_point_select-10     174µs ±10%     177µs ±19%    ~     (p=0.799 n=19+18)
Sysbench/KV/1node_remote/oltp_read_only-10         843µs ±17%     824µs ±15%    ~     (p=0.461 n=20+19)
Sysbench/KV/1node_remote/oltp_write_only-10        784µs ±14%     751µs ±24%    ~     (p=0.226 n=18+18)
Sysbench/KV/1node_remote/oltp_read_write-10       1.75ms ±15%    1.76ms ±21%    ~     (p=0.639 n=17+19)
Sysbench/KV/1node_remote/oltp_point_select-10     35.7µs ±16%    36.4µs ±25%    ~     (p=0.795 n=19+19)

name                                            old alloc/op   new alloc/op   delta
Sysbench/KV/1node_remote/oltp_point_select-10     6.57kB ± 0%    6.56kB ± 0%  -0.26%  (p=0.005 n=20+20)
Sysbench/SQL/1node_remote/oltp_read_only-10       1.15MB ± 1%    1.15MB ± 0%    ~     (p=0.191 n=19+19)
Sysbench/SQL/1node_remote/oltp_write_only-10       573kB ± 4%     573kB ± 4%    ~     (p=0.758 n=20+20)
Sysbench/SQL/1node_remote/oltp_read_write-10      1.64MB ± 1%    1.64MB ± 2%    ~     (p=0.771 n=19+20)
Sysbench/SQL/1node_remote/oltp_point_select-10    30.0kB ± 1%    30.1kB ± 1%    ~     (p=0.632 n=20+19)
Sysbench/KV/1node_remote/oltp_read_only-10         656kB ± 0%     655kB ± 0%    ~     (p=0.067 n=18+20)
Sysbench/KV/1node_remote/oltp_write_only-10        263kB ± 2%     263kB ± 2%    ~     (p=0.961 n=19+16)
Sysbench/KV/1node_remote/oltp_read_write-10        879kB ± 1%     880kB ± 1%    ~     (p=0.495 n=20+20)

name                                            old allocs/op  new allocs/op  delta
Sysbench/KV/1node_remote/oltp_point_select-10       61.0 ± 0%      59.0 ± 0%  -3.28%  (p=0.000 n=20+20)
Sysbench/KV/1node_remote/oltp_read_only-10         1.87k ± 0%     1.84k ± 0%  -1.50%  (p=0.000 n=17+17)
Sysbench/KV/1node_remote/oltp_read_write-10        3.46k ± 0%     3.42k ± 0%  -1.38%  (p=0.000 n=20+19)
Sysbench/KV/1node_remote/oltp_write_only-10        1.59k ± 0%     1.57k ± 0%  -1.32%  (p=0.000 n=20+18)
Sysbench/SQL/1node_remote/oltp_write_only-10       3.98k ± 1%     3.94k ± 1%  -1.00%  (p=0.000 n=20+20)
Sysbench/SQL/1node_remote/oltp_read_only-10        4.78k ± 1%     4.74k ± 1%  -0.76%  (p=0.000 n=19+19)
Sysbench/SQL/1node_remote/oltp_point_select-10       266 ± 1%       264 ± 1%  -0.64%  (p=0.003 n=20+19)
Sysbench/SQL/1node_remote/oltp_read_write-10       8.74k ± 0%     8.68k ± 1%  -0.60%  (p=0.000 n=19+18)
```

Epic: None
Release note: None

Co-authored-by: Nathan VanBenschoten <[email protected]>

pkg/kv/kvpb/gen/main.go

@@ -366,13 +366,49 @@ func (ba *BatchRequest) WriteSummary(b *strings.Builder) {
 		allocTypes[resV.variantName] = allocName
 	}
 
+	fmt.Fprint(f, `
+func allocBatchResponse(nResps int) *BatchResponse {
+	if nResps <= 1 {
+		alloc := new(struct {
+			br    BatchResponse
+			resps [1]ResponseUnion
+		})
+		alloc.br.Responses = alloc.resps[:nResps]
+		return &alloc.br
+	} else if nResps <= 2 {
+		alloc := new(struct {
+			br    BatchResponse
+			resps [2]ResponseUnion
+		})
+		alloc.br.Responses = alloc.resps[:nResps]
+		return &alloc.br
+	} else if nResps <= 4 {
+		alloc := new(struct {
+			br    BatchResponse
+			resps [4]ResponseUnion
+		})
+		alloc.br.Responses = alloc.resps[:nResps]
+		return &alloc.br
+	} else if nResps <= 8 {
+		alloc := new(struct {
+			br    BatchResponse
+			resps [8]ResponseUnion
+		})
+		alloc.br.Responses = alloc.resps[:nResps]
+		return &alloc.br
+	}
+	br := &BatchResponse{}
+	br.Responses = make([]ResponseUnion, nResps)
+	return br
+}
+`)
+
 	fmt.Fprint(f, `
 // CreateReply creates replies for each of the contained requests, wrapped in a
 // BatchResponse. The response objects are batch allocated to minimize
 // allocation overhead.
 func (ba *BatchRequest) CreateReply() *BatchResponse {
-	br := &BatchResponse{}
-	br.Responses = make([]ResponseUnion, len(ba.Requests))
+	br := allocBatchResponse(len(ba.Requests))
 
 	counts := ba.getReqCounts()
 

pkg/server/node.go

@@ -1877,7 +1877,14 @@ func (n *Node) Batch(ctx context.Context, args *kvpb.BatchRequest) (*kvpb.BatchR
 func (n *Node) BatchStream(stream kvpb.Internal_BatchStreamServer) error {
 	ctx := stream.Context()
 	for {
-		args, err := stream.Recv()
+		argsAlloc := new(struct {
+			args kvpb.BatchRequest
+			reqs [1]kvpb.RequestUnion
+		})
+		args := &argsAlloc.args
+		args.Requests = argsAlloc.reqs[:0]
+
+		err := stream.RecvMsg(args)
 		if err != nil {
 			// From grpc.ServerStream.Recv:
 			// > It returns io.EOF when the client has performed a CloseSend.