[WIP] rpc: recycle request/response memory across gRPC calls

pkg/rpc/context.go

@@ -1397,6 +1397,10 @@ func (rpcCtx *Context) dialOptsLocal() ([]grpc.DialOption, error) {
 		return nil, err
 	}
 
+	if rpcCtx.rpcCompression {
+		dialOpts = append(dialOpts, grpc.WithDefaultCallOptions(grpc.UseCompressor((snappyCompressor{}).Name())))
+	}
+
 	dialOpts = append(dialOpts, grpc.WithContextDialer(
 		func(ctx context.Context, _ string) (net.Conn, error) {
 			return rpcCtx.loopbackDialFn(ctx)

pkg/rpc/stream_pool.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -24,6 +25,7 @@ import (
 type streamClient[Req, Resp any] interface {
 	Send(Req) error
 	Recv() (Resp, error)
+	grpc.ClientStream
 }
 
 // streamConstructor creates a new gRPC stream client over the provided client
@@ -74,6 +76,8 @@ type pooledStream[Req, Resp any, Conn comparable] struct {
 
 	reqC  chan Req
 	respC chan result[Resp]
+
+	reqMsg *grpcutil.PreparedMsg
 }
 
 func newPooledStream[Req, Resp any, Conn comparable](
@@ -89,6 +93,7 @@ func newPooledStream[Req, Resp any, Conn comparable](
 		streamCancel: streamCancel,
 		reqC:         make(chan Req),
 		respC:        make(chan result[Resp], 1),
+		reqMsg:       new(grpcutil.PreparedMsg),
 	}
 }
 
@@ -101,7 +106,12 @@ func (s *pooledStream[Req, Resp, Conn]) run(ctx context.Context) {
 func (s *pooledStream[Req, Resp, Conn]) runOnce(ctx context.Context) (loop bool) {
 	select {
 	case req := <-s.reqC:
-		err := s.stream.Send(req)
+		err := s.reqMsg.Encode(s.stream, req)
+		if err != nil {
+			s.respC <- result[Resp]{err: err}
+			return false
+		}
+		err = s.stream.SendMsg(s.reqMsg.AsGrpc())
 		if err != nil {
 			// From grpc.ClientStream.SendMsg:
 			// > On error, SendMsg aborts the stream.

pkg/server/node.go

@@ -59,6 +59,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
+	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/limit"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -1876,6 +1877,7 @@ func (n *Node) Batch(ctx context.Context, args *kvpb.BatchRequest) (*kvpb.BatchR
 // BatchStream implements the kvpb.InternalServer interface.
 func (n *Node) BatchStream(stream kvpb.Internal_BatchStreamServer) error {
 	ctx := stream.Context()
+	respMsg := new(grpcutil.PreparedMsg)
 	for {
 		argsAlloc := new(struct {
 			args kvpb.BatchRequest
@@ -1898,7 +1900,11 @@ func (n *Node) BatchStream(stream kvpb.Internal_BatchStreamServer) error {
 		if err != nil {
 			return err
 		}
-		err = stream.Send(br)
+		err = respMsg.Encode(stream, br)
+		if err != nil {
+			return err
+		}
+		err = stream.SendMsg(respMsg.AsGrpc())
 		if err != nil {
 			return err
 		}

pkg/util/grpcutil/BUILD.bazel

@@ -3,11 +3,13 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "grpcutil",
     srcs = [
+        "compressor_info.go",
         "fast_metadata.go",
         "grpc_err_redaction.go",
         "grpc_log.go",
         "grpc_log_legacy.go",
         "grpc_util.go",
+        "prepared_message.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/util/grpcutil",
     visibility = ["//visibility:public"],
@@ -17,11 +19,14 @@ go_library(
         "//pkg/util/log",
         "//pkg/util/log/severity",
         "//pkg/util/netutil",
+        "//pkg/util/protoutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_errors//errbase",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_gogo_status//:status",
+        "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//encoding",
         "@org_golang_google_grpc//grpclog",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",

pkg/util/grpcutil/compressor_info.go

@@ -0,0 +1,70 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package grpcutil
+
+import (
+	"context"
+	"unsafe"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/encoding"
+)
+
+// RPCInfo exports grpc.rpcInfo.
+type RPCInfo struct {
+	FailFast      bool
+	PreloaderInfo *CompressorInfo
+}
+
+// CompressorInfo exports grpc.compressorInfo.
+type CompressorInfo struct {
+	Codec encoding.Codec
+	Cp    grpc.Compressor
+	Comp  encoding.Compressor
+}
+
+// From runtime/runtime2.go:eface
+type eface struct {
+	typ, data unsafe.Pointer
+}
+
+// RPCInfoFromContext extracts the RPCInfo from the context.
+func RPCInfoFromContext(ctx context.Context) (*RPCInfo, bool) {
+	v := ctx.Value(grpcInfoContextKeyObj)
+	if v == nil {
+		return nil, false
+	}
+	return (*RPCInfo)((*eface)(unsafe.Pointer(&v)).data), true
+}
+
+// grpcInfoContextKeyObj is a copy of a value with the Go type
+// `grpc.rpcInfoContextKey{}`. We cannot construct an object of that type
+// directly, but we can "steal" it by forcing the grpc to give it to us:
+// `grpc.PreparedMsg.Encode` gives an instance of this object as parameter to
+// the `Value` method of the context you give it as argument. We use a custom
+// implementation of that to "steal" the argument of type `rpcInfoContextKey{}`
+// given to us that way.
+//
+// This is the same trick that we pull with grpcIncomingKeyObj.
+var grpcInfoContextKeyObj = func() interface{} {
+	var s fakeStream
+	_ = (*grpc.PreparedMsg)(nil).Encode(&s, nil)
+	if s.recordedKey == nil {
+		panic("PreparedMsg.Encode did not request a key")
+	}
+	return s.recordedKey
+}()
+
+type fakeStream struct {
+	fakeContext
+}
+
+var _ grpc.Stream = (*fakeStream)(nil)
+
+func (s *fakeStream) Context() context.Context { return &s.fakeContext }
+
+func (*fakeStream) SendMsg(interface{}) error { panic("unused") }
+func (*fakeStream) RecvMsg(interface{}) error { panic("unused") }