kv: extract KVBatch service from Internal service and remove api_drpc_hacky

This change removes the old manually-created `api_drpc_hacky.go` file and
uses the properly auto-generated `api_drpc.pb.go` now that the DRPC
compiler setup is working.

We also extract the `KVBatch` service from the `Internal` service, which
has largely been an ad-hoc collection of internal RPCs. This enables:
- Easier Maintenance: Smaller services are simpler to manage and update.
- Smoother DRPC Rollout: We can move to DRPC step-by-step with these
  smaller services, rather than doing `Internal` all at once.
- Subsystems get client interfaces focused only on what they need.

We also refactor `batchStreamImpl` to directly call `RecvMsg` on the
provided stream, which is enabled by the change in codegen to include a
gRPC-compatible `RecvMsg`.

---

Some context on why we use `RecvMsg`: Commit 34e3abf optimized the
`BatchStream` implementation by changing how the `BatchRequest` and its
initial `Requests` slice capacity are allocated. `RecvMsg` allows the
caller to provide a pre-allocated `BatchRequest` instance. This pattern,
when used with a pre-allocated instance of a `BatchRequest` header and
first request together, helps reduce per-message allocations compared to a
`Recv()` method that might allocate a new message object on each call

Author: shubhamdhama

pkg/kv/kvpb/BUILD.bazel

@@ -8,21 +8,17 @@ load(":gen.bzl", "batch_gen")
 go_library(
     name = "kvpb",
     srcs = [
-        ":gen-batch-generated",  # keep
-        ":gen-errordetailtype-stringer",  # keep
-        ":gen-method-stringer",  # keep
         "ambiguous_result_error.go",
         "api.go",
-        # DRPC protobuf file (api_drpc.pb.go) is currently generated manually.
-        # TODO (chandrat): Remove this once DRPC protobuf generation is
-        # integrated into the build process.
-        "api_drpc_hacky.go",
         "batch.go",
         "data.go",
         "errors.go",
         "method.go",
         "node_decommissioned_error.go",
         "replica_unavailable_error.go",
+        ":gen-batch-generated",  # keep
+        ":gen-errordetailtype-stringer",  # keep
+        ":gen-method-stringer",  # keep
     ],
     embed = [":kvpb_go_proto"],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvpb",
@@ -52,8 +48,6 @@ go_library(
         "@com_github_gogo_protobuf//types",
         "@com_github_gogo_status//:status",
         "@com_github_golang_mock//gomock",  # keep
-        "@io_storj_drpc//:drpc",
-        "@io_storj_drpc//drpcerr",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",  # keep
     ],

pkg/kv/kvpb/api.proto

@@ -3618,6 +3618,19 @@ message JoinNodeResponse {
   roachpb.Version active_version = 4;
 }
 
+// KVBatch defines Batch RPCs of Internal service (see below) for DRPC use only.
+//
+// During the gRPC to DRPC migration:
+// - gRPC clients will continue to use the full 'Internal' service client.
+// - DRPC clients targeting Batch* RPCs will use clients generated from this.
+//
+// TODO(server): Remove these methods from the 'Internal' service once the
+// migration to DRPC complete, making this service the sole definition.
+service KVBatch {
+  rpc Batch(BatchRequest) returns (BatchResponse) {}
+  rpc BatchStream(stream BatchRequest) returns (stream BatchResponse) {}
+}
+
 // Batch and RangeFeed service implemented by nodes for KV API requests.
 service Internal {
   rpc Batch (BatchRequest) returns (BatchResponse) {}

pkg/kv/kvpb/api_drpc_hacky.go

@@ -167,7 +167,7 @@ func (n *Dialer) DialInternalClient(
 		client = &unaryDRPCBatchServiceToInternalAdapter{
 			useStreamPoolClient:      useStreamPoolClient,
 			RestrictedInternalClient: client, // for RangeFeed only
-			drpcClient:               kvpb.NewDRPCBatchClient(dconn),
+			drpcClient:               kvpb.NewDRPCKVBatchClient(dconn),
 			drpcStreamPool:           drpcBatchStreamPool,
 		}
 		return client, nil

pkg/rpc/nodedialer/nodedialer.go

@@ -17,7 +17,7 @@ import (
 type unaryDRPCBatchServiceToInternalAdapter struct {
 	useStreamPoolClient bool
 	rpc.RestrictedInternalClient
-	drpcClient     kvpb.DRPCBatchClient
+	drpcClient     kvpb.DRPCKVBatchClient
 	drpcStreamPool *rpc.DRPCBatchStreamPool
 }
 

pkg/rpc/nodedialer/nodedialer_drpc.go

@@ -334,5 +334,5 @@ type DRPCBatchStreamClient = streamClient[*kvpb.BatchRequest, *kvpb.BatchRespons
 
 // newDRPCBatchStream constructs a BatchStreamClient from a drpc.Conn.
 func newDRPCBatchStream(ctx context.Context, dc drpc.Conn) (DRPCBatchStreamClient, error) {
-	return kvpb.NewDRPCBatchClient(dc).BatchStream(ctx)
+	return kvpb.NewDRPCKVBatchClient(dc).BatchStream(ctx)
 }

pkg/rpc/stream_pool.go

@@ -67,7 +67,7 @@ func TestDRPCBatchServer(t *testing.T) {
 
 		desc := c.LookupRangeOrFatal(t, c.ScratchRange(t))
 
-		client := kvpb.NewDRPCBatchClient(conn)
+		client := kvpb.NewDRPCKVBatchClient(conn)
 		ba := &kvpb.BatchRequest{}
 		ba.RangeID = desc.RangeID
 		var ok bool
@@ -121,7 +121,7 @@ func TestStreamContextCancel(t *testing.T) {
 	}()
 
 	desc := c.LookupRangeOrFatal(t, c.ScratchRange(t))
-	client := kvpb.NewDRPCBatchClient(conn)
+	client := kvpb.NewDRPCKVBatchClient(conn)
 
 	singleRequest := func() {
 		streamCtx, streamCtxCancel := context.WithCancel(ctx)

pkg/server/drpc_test.go

@@ -1867,18 +1867,10 @@ 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 {
-	return n.batchStreamImpl(stream, func(ba *kvpb.BatchRequest) error {
-		return stream.RecvMsg(ba)
-	})
+	return n.batchStreamImpl(stream)
 }
 
-func (n *Node) batchStreamImpl(
-	stream interface {
-		Context() context.Context
-		Send(response *kvpb.BatchResponse) error
-	},
-	recvMsg func(*kvpb.BatchRequest) error,
-) error {
+func (n *Node) batchStreamImpl(stream kvpb.RPCKVBatch_BatchStreamStream) error {
 	ctx := stream.Context()
 	for {
 		argsAlloc := new(struct {
@@ -1888,7 +1880,7 @@ func (n *Node) batchStreamImpl(
 		args := &argsAlloc.args
 		args.Requests = argsAlloc.reqs[:0]
 
-		err := recvMsg(args)
+		err := stream.RecvMsg(args)
 		if err != nil {
 			if errors.Is(err, io.EOF) {
 				return nil
@@ -1907,7 +1899,7 @@ func (n *Node) batchStreamImpl(
 	}
 }
 
-func (n *Node) AsDRPCBatchServer() kvpb.DRPCBatchServer {
+func (n *Node) AsDRPCKVBatchServer() kvpb.DRPCKVBatchServer {
 	return (*drpcNode)(n)
 }
 
@@ -1919,12 +1911,8 @@ func (n *drpcNode) Batch(
 	return (*Node)(n).Batch(ctx, request)
 }
 
-func (n *drpcNode) BatchStream(stream kvpb.DRPCBatch_BatchStreamStream) error {
-	return (*Node)(n).batchStreamImpl(stream, func(ba *kvpb.BatchRequest) error {
-		return stream.(interface {
-			RecvMsg(request *kvpb.BatchRequest) error
-		}).RecvMsg(ba)
-	})
+func (n *drpcNode) BatchStream(stream kvpb.DRPCKVBatch_BatchStreamStream) error {
+	return (*Node)(n).batchStreamImpl(stream)
 }
 
 // spanForRequest is the retval of setupSpanForIncomingRPC. It groups together a

pkg/server/node.go

@@ -997,7 +997,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		cfg.LicenseEnforcer,
 	)
 	kvpb.RegisterInternalServer(grpcServer.Server, node)
-	if err := kvpb.DRPCRegisterBatch(drpcServer.mux, node.AsDRPCBatchServer()); err != nil {
+	if err := kvpb.DRPCRegisterKVBatch(drpcServer.mux, node.AsDRPCKVBatchServer()); err != nil {
 		return nil, err
 	}
 	kvserver.RegisterPerReplicaServer(grpcServer.Server, node.perReplicaServer)