Skip to content

Commit 838bcb1

Browse files
craig[bot]shubhamdhama
andcommitted
Merge #145195
145195: kv: extract KVBatch service from Internal service and remove api_drpc_hacky r=rickystewart,cthumuluru-crdb a=shubhamdhama 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 Fixes: #145179 Release note: none Epic: CRDB-48934 Co-authored-by: Shubham Dhama <[email protected]>
2 parents 8bb4bf7 + 13835e9 commit 838bcb1

File tree

13 files changed

+35
-247
lines changed

13 files changed

+35
-247
lines changed

DEPS.bzl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11509,10 +11509,10 @@ def go_deps():
1150911509
name = "io_storj_drpc",
1151011510
build_file_proto_mode = "disable_global",
1151111511
importpath = "storj.io/drpc",
11512-
sha256 = "60246f2907c40615808e4a5af2f6d58a5248802f57a4b7b55dfb03ee5b779b3b",
11513-
strip_prefix = "github.com/cockroachdb/[email protected]20250507084558-a793c5c40d3d",
11512+
sha256 = "67a517d0a6c90586f265135ab60619a06a3be5b5c871c15093865716a6de5b39",
11513+
strip_prefix = "github.com/cockroachdb/[email protected]20250603054748-5b0c5d2c7b38",
1151411514
urls = [
11515-
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250507084558-a793c5c40d3d.zip",
11515+
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250603054748-5b0c5d2c7b38.zip",
1151611516
],
1151711517
)
1151811518
go_repository(

build/bazelutil/distdir_files.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,7 @@ DISTDIR_FILES = {
358358
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/crlfmt/com_github_cockroachdb_crlfmt-v0.0.0-20221214225007-b2fc5c302548.zip": "fedc01bdd6d964da0425d5eaac8efadc951e78e13f102292cc0774197f09ab63",
359359
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/crlib/com_github_cockroachdb_crlib-v0.0.0-20241205160938-4a90b184f49c.zip": "1afc910b4ff270de79eecb42ab7bd5e6404e6128666c6c55e96db9e27d28e69e",
360360
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/datadriven/com_github_cockroachdb_datadriven-v1.0.3-0.20250407164829-2945557346d5.zip": "251593cd9c040fe84a99a3919de7ce6f85030d522159a37d625dc2dea7a4d17f",
361-
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250507084558-a793c5c40d3d.zip": "60246f2907c40615808e4a5af2f6d58a5248802f57a4b7b55dfb03ee5b779b3b",
361+
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250603054748-5b0c5d2c7b38.zip": "67a517d0a6c90586f265135ab60619a06a3be5b5c871c15093865716a6de5b39",
362362
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/errors/com_github_cockroachdb_errors-v1.12.0.zip": "f73d8a5f4d8fcbc4ed61db2b47f17e2601d8b32e9a49c0665667489d9d9d6e7c",
363363
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/go-test-teamcity/com_github_cockroachdb_go_test_teamcity-v0.0.0-20191211140407-cff980ad0a55.zip": "bac30148e525b79d004da84d16453ddd2d5cd20528e9187f1d7dac708335674b",
364364
"https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/gogoproto/com_github_cockroachdb_gogoproto-v1.3.3-0.20241216150617-2358cdb156a1.zip": "bf052c9a7f9e23fb3ec7e9f3b7201cfc264c18ed6da0d662952d276dbc339003",

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -504,7 +504,7 @@ replace golang.org/x/time => github.com/cockroachdb/x-time v0.3.1-0.202305251236
504504

505505
replace github.com/gogo/protobuf => github.com/cockroachdb/gogoproto v1.3.3-0.20241216150617-2358cdb156a1
506506

507-
replace storj.io/drpc => github.com/cockroachdb/drpc v0.0.0-20250507084558-a793c5c40d3d
507+
replace storj.io/drpc => github.com/cockroachdb/drpc v0.0.0-20250603054748-5b0c5d2c7b38
508508

509509
// Note: This forked dependency adds a commit that opens up some
510510
// private APIs to enable us to make some perf improvements to

go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -579,8 +579,8 @@ github.com/cockroachdb/datadriven v0.0.0-20190809214429-80d97fb3cbaa/go.mod h1:z
579579
github.com/cockroachdb/datadriven v1.0.2/go.mod h1:a9RdTaap04u637JoCzcUoIcDmvwSUtcUFtT/C3kJlTU=
580580
github.com/cockroachdb/datadriven v1.0.3-0.20250407164829-2945557346d5 h1:UycK/E0TkisVrQbSoxvU827FwgBBcZ95nRRmpj/12QI=
581581
github.com/cockroachdb/datadriven v1.0.3-0.20250407164829-2945557346d5/go.mod h1:jsaKMvD3RBCATk1/jbUZM8C9idWBJME9+VRZ5+Liq1g=
582-
github.com/cockroachdb/drpc v0.0.0-20250507084558-a793c5c40d3d h1:O43v15qZJIut4hLBaYfuK9gnm9eWpto2Qj/Q6MZeLTY=
583-
github.com/cockroachdb/drpc v0.0.0-20250507084558-a793c5c40d3d/go.mod h1:Y9LZaa8esL1PW2IDMqJE7CFSNq7d5bQ3RI7mGPtmKMg=
582+
github.com/cockroachdb/drpc v0.0.0-20250603054748-5b0c5d2c7b38 h1:7FZw1FKZIOabgsOYY/D8cSj60E7B7h9Tm7xSos/lWF8=
583+
github.com/cockroachdb/drpc v0.0.0-20250603054748-5b0c5d2c7b38/go.mod h1:UWP+upGv1Z+4nWxcdwhv3/wQXSOgCZteytaRVez5PDc=
584584
github.com/cockroachdb/errors v1.9.1/go.mod h1:2sxOtL2WIc096WSZqZ5h8fa17rdDq9HZOZLBCor4mBk=
585585
github.com/cockroachdb/errors v1.12.0 h1:d7oCs6vuIMUQRVbi6jWWWEJZahLCfJpnJSVobd1/sUo=
586586
github.com/cockroachdb/errors v1.12.0/go.mod h1:SvzfYNNBshAVbZ8wzNc/UPK3w1vf0dKDUP41ucAIf7g=

pkg/kv/kvpb/BUILD.bazel

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,21 +8,17 @@ load(":gen.bzl", "batch_gen")
88
go_library(
99
name = "kvpb",
1010
srcs = [
11-
":gen-batch-generated", # keep
12-
":gen-errordetailtype-stringer", # keep
13-
":gen-method-stringer", # keep
1411
"ambiguous_result_error.go",
1512
"api.go",
16-
# DRPC protobuf file (api_drpc.pb.go) is currently generated manually.
17-
# TODO (chandrat): Remove this once DRPC protobuf generation is
18-
# integrated into the build process.
19-
"api_drpc_hacky.go",
2013
"batch.go",
2114
"data.go",
2215
"errors.go",
2316
"method.go",
2417
"node_decommissioned_error.go",
2518
"replica_unavailable_error.go",
19+
":gen-batch-generated", # keep
20+
":gen-errordetailtype-stringer", # keep
21+
":gen-method-stringer", # keep
2622
],
2723
embed = [":kvpb_go_proto"],
2824
importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvpb",
@@ -52,8 +48,6 @@ go_library(
5248
"@com_github_gogo_protobuf//types",
5349
"@com_github_gogo_status//:status",
5450
"@com_github_golang_mock//gomock", # keep
55-
"@io_storj_drpc//:drpc",
56-
"@io_storj_drpc//drpcerr",
5751
"@org_golang_google_grpc//codes",
5852
"@org_golang_google_grpc//metadata", # keep
5953
],

pkg/kv/kvpb/api.proto

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3618,6 +3618,19 @@ message JoinNodeResponse {
36183618
roachpb.Version active_version = 4;
36193619
}
36203620

3621+
// KVBatch defines Batch RPCs of Internal service (see below) for DRPC use only.
3622+
//
3623+
// During the gRPC to DRPC migration:
3624+
// - gRPC clients will continue to use the full 'Internal' service client.
3625+
// - DRPC clients targeting Batch* RPCs will use clients generated from this.
3626+
//
3627+
// TODO(server): Remove these methods from the 'Internal' service once the
3628+
// migration to DRPC complete, making this service the sole definition.
3629+
service KVBatch {
3630+
rpc Batch(BatchRequest) returns (BatchResponse) {}
3631+
rpc BatchStream(stream BatchRequest) returns (stream BatchResponse) {}
3632+
}
3633+
36213634
// Batch and RangeFeed service implemented by nodes for KV API requests.
36223635
service Internal {
36233636
rpc Batch (BatchRequest) returns (BatchResponse) {}

pkg/kv/kvpb/api_drpc_hacky.go

Lines changed: 0 additions & 207 deletions
This file was deleted.

pkg/rpc/nodedialer/nodedialer.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ func (n *Dialer) DialInternalClient(
167167
client = &unaryDRPCBatchServiceToInternalAdapter{
168168
useStreamPoolClient: useStreamPoolClient,
169169
RestrictedInternalClient: client, // for RangeFeed only
170-
drpcClient: kvpb.NewDRPCBatchClient(dconn),
170+
drpcClient: kvpb.NewDRPCKVBatchClient(dconn),
171171
drpcStreamPool: drpcBatchStreamPool,
172172
}
173173
return client, nil

pkg/rpc/nodedialer/nodedialer_drpc.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import (
1717
type unaryDRPCBatchServiceToInternalAdapter struct {
1818
useStreamPoolClient bool
1919
rpc.RestrictedInternalClient
20-
drpcClient kvpb.DRPCBatchClient
20+
drpcClient kvpb.DRPCKVBatchClient
2121
drpcStreamPool *rpc.DRPCBatchStreamPool
2222
}
2323

pkg/rpc/stream_pool.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -334,5 +334,5 @@ type DRPCBatchStreamClient = streamClient[*kvpb.BatchRequest, *kvpb.BatchRespons
334334

335335
// newDRPCBatchStream constructs a BatchStreamClient from a drpc.Conn.
336336
func newDRPCBatchStream(ctx context.Context, dc drpc.Conn) (DRPCBatchStreamClient, error) {
337-
return kvpb.NewDRPCBatchClient(dc).BatchStream(ctx)
337+
return kvpb.NewDRPCKVBatchClient(dc).BatchStream(ctx)
338338
}

0 commit comments

Comments
 (0)