From 18c8ff9497baa12dce80c380a61d838dfe5f4b46 Mon Sep 17 00:00:00 2001 From: Shubham Dhama Date: Wed, 13 May 2026 05:33:40 +0000 Subject: [PATCH 1/4] do-not-merge: enable DRPC --- pkg/sql/tests/sysbench_test.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pkg/sql/tests/sysbench_test.go b/pkg/sql/tests/sysbench_test.go index 905cf83665eb..700a09593cbe 100644 --- a/pkg/sql/tests/sysbench_test.go +++ b/pkg/sql/tests/sysbench_test.go @@ -190,9 +190,10 @@ func newTestCluster( } tc := serverutils.StartCluster(b, nodes, base.TestClusterArgs{ ServerArgs: base.TestServerArgs{ - Settings: st, - CacheSize: cacheSize, - Knobs: knobs, + Settings: st, + DefaultDRPCOption: base.TestDRPCEnabled, + CacheSize: cacheSize, + Knobs: knobs, }}, ) if nodes > 1 { From 50a2734d5796cf004e7da5ab6f9da17b6f6e864b Mon Sep 17 00:00:00 2001 From: Shubham Dhama Date: Wed, 13 May 2026 07:55:11 +0000 Subject: [PATCH 2/4] scripts: add bump-drpc.sh helper Adds a small helper that bumps the DRPC replace directive in go.mod to a specific commit. It takes a GitHub commit URL (cockroachdb or shubhamdhama fork), resolves the pseudo-version with `go list`, rewrites the replace line, runs `go mod tidy`, and regenerates the Bazel files via `./dev generate bazel --mirror`. Bumping DRPC was a multi-step ritual; the script removes the chance of forgetting one of the steps. Release note: None Epic: none --- scripts/bump-drpc.sh | 80 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100755 scripts/bump-drpc.sh diff --git a/scripts/bump-drpc.sh b/scripts/bump-drpc.sh new file mode 100755 index 000000000000..24b3f3aee9e6 --- /dev/null +++ b/scripts/bump-drpc.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash + +# Copyright 2026 The Cockroach Authors. +# +# Use of this software is governed by the CockroachDB Software License +# included in the /LICENSE file. + +# This script bumps the DRPC dependency in go.mod to a given commit. +# +# Usage: +# +# ./scripts/bump-drpc.sh +# +# Where is a GitHub commit link, e.g.: +# https://github.com/cockroachdb/drpc/commit/d57f8922c21431976751f40127e59cda71768431 +# https://github.com/shubhamdhama/drpc/commit/d57f8922c21431976751f40127e59cda71768431 + +set -euo pipefail + +if [ -z "${1-}" ]; then + echo "Usage: $0 " >&2 + echo " e.g. $0 https://github.com/cockroachdb/drpc/commit/abc123" >&2 + exit 1 +fi + +COMMIT_URL="$1" + +# Parse the org and SHA from the commit URL. +# Expected format: https://github.com//drpc/commit/ +if [[ "$COMMIT_URL" =~ github\.com/([^/]+)/drpc/commit/([0-9a-f]+) ]]; then + ORG="${BASH_REMATCH[1]}" + SHA="${BASH_REMATCH[2]}" +else + echo "Error: could not parse commit URL: $COMMIT_URL" >&2 + echo "Expected format: https://github.com//drpc/commit/" >&2 + exit 1 +fi + +MODULE="github.com/${ORG}/drpc" +echo "Organization: $ORG" +echo "Commit SHA: $SHA" +echo "Module: $MODULE" +echo + +# Resolve the pseudo-version for this commit. +echo "Resolving pseudo-version..." +VERSION=$(go list -m -json "${MODULE}@${SHA}" | jq -r '.Version') +if [ -z "$VERSION" ] || [ "$VERSION" == "null" ]; then + echo "Error: could not resolve version for ${MODULE}@${SHA}" >&2 + exit 1 +fi +echo "Resolved version: $VERSION" +echo + +# Update the replace directive in go.mod. +OLD_REPLACE=$(grep '^replace storj.io/drpc =>' go.mod) +NEW_REPLACE="replace storj.io/drpc => ${MODULE} ${VERSION}" + +if [ -z "$OLD_REPLACE" ]; then + echo "Error: could not find 'replace storj.io/drpc =>' in go.mod" >&2 + exit 1 +fi + +echo "Old: $OLD_REPLACE" +echo "New: $NEW_REPLACE" +echo + +sed -i "s|^replace storj.io/drpc =>.*|${NEW_REPLACE}|" go.mod + +# Tidy up. +go mod tidy + +# Regenerate Bazel files. +echo "Running dev generate bazel --mirror..." +./dev generate bazel --mirror + +echo +echo "Done. DRPC bumped to ${SHA:0:12} (${VERSION})." +echo "Modified files:" +git diff --name-only From 2b1ebfd24b795b75c0c8966165e374c0f4f03209 Mon Sep 17 00:00:00 2001 From: Shubham Dhama Date: Wed, 15 Apr 2026 07:48:36 +0000 Subject: [PATCH 3/4] rpc: wire up DRPC stream multiplexing behind an env gate Bumps DRPC to pick up stream multiplexing from cockroachdb/drpc#58, where a single transport connection can carry multiple concurrent streams. The multiplexing dial path is guarded by COCKROACH_EXPERIMENTAL_DRPC_MUX_ENABLED, defaulted to false. With the gate off, DialDRPC keeps using drpcpool, which checks out a connection per active stream and dials a new one when none is idle. The upgraded multiplexing-capable library is therefore exercised on every connection, but the multiplexing capability itself stays unused: every stream still ends up on its own underlying connection, matching the prior behavior. This keeps the initial scope small. The dial-mux path is in place so it can be turned on for validation, with a follow-up to flip the default once we are confident. Both paths now share the same dial-option setup, so the mux path also picks up the client metrics and the request-recording gate that the pool path was already attaching. Release note: None Epic: none --- DEPS.bzl | 6 +++--- build/bazelutil/distdir_files.bzl | 2 +- go.mod | 2 +- go.sum | 4 ++-- pkg/rpc/drpc.go | 26 ++++++++++++++++++++++++-- pkg/server/BUILD.bazel | 1 - pkg/server/drpc_test.go | 7 +------ 7 files changed, 32 insertions(+), 16 deletions(-) diff --git a/DEPS.bzl b/DEPS.bzl index a9c126c554c0..855066af3fed 100644 --- a/DEPS.bzl +++ b/DEPS.bzl @@ -11080,10 +11080,10 @@ def go_deps(): name = "io_storj_drpc", build_file_proto_mode = "disable_global", importpath = "storj.io/drpc", - sha256 = "089481375939240f2fd5c2caa87ecde398bb2a9fb7d7998a2cee8ad0624a7ea1", - strip_prefix = "github.com/cockroachdb/drpc@v0.0.0-20260406142218-6c77a9e470d3", + sha256 = "7c658dfbb3603b4351e4d0d91b91a9787d868348915582f547446e9f96b6ac11", + strip_prefix = "github.com/shubhamdhama/drpc@v0.0.0-20260511155932-95375d1e96eb", urls = [ - "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20260406142218-6c77a9e470d3.zip", + "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shubhamdhama/drpc/com_github_shubhamdhama_drpc-v0.0.0-20260511155932-95375d1e96eb.zip", ], ) go_repository( diff --git a/build/bazelutil/distdir_files.bzl b/build/bazelutil/distdir_files.bzl index 86d75e3f269c..81b8b73be483 100644 --- a/build/bazelutil/distdir_files.bzl +++ b/build/bazelutil/distdir_files.bzl @@ -349,7 +349,6 @@ DISTDIR_FILES = { "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/crlfmt/com_github_cockroachdb_crlfmt-v0.0.0-20221214225007-b2fc5c302548.zip": "fedc01bdd6d964da0425d5eaac8efadc951e78e13f102292cc0774197f09ab63", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/crlib/com_github_cockroachdb_crlib-v0.0.0-20251122031428-fe658a2dbda1.zip": "fa361e52b072ce18ac1d103e2556851906477361b5fe688745d29b1efabff3cb", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/datadriven/com_github_cockroachdb_datadriven-v1.0.3-0.20251123150250-ddff6747b112.zip": "6c1ae8a9550d3a92f85d7371e5c0caf922dc0ff1425e299359814b6a9a587f14", - "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20260406142218-6c77a9e470d3.zip": "089481375939240f2fd5c2caa87ecde398bb2a9fb7d7998a2cee8ad0624a7ea1", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/errors/com_github_cockroachdb_errors-v1.12.1-0.20251010171200-64801262cd6f.zip": "4df66cc44791d4290071696abf179dc6df7b94b4cb5d29a20f39c6bf522c60ee", "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", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/gogoproto/com_github_cockroachdb_gogoproto-v1.3.3-0.20241216150617-2358cdb156a1.zip": "bf052c9a7f9e23fb3ec7e9f3b7201cfc264c18ed6da0d662952d276dbc339003", @@ -936,6 +935,7 @@ DISTDIR_FILES = { "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/sergi/go-diff/com_github_sergi_go_diff-v1.0.0.zip": "287218ffcd136dbb28ce99a2f162048d8dfa6f97b524c17797964aacde2f8f52", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shirou/gopsutil/v3/com_github_shirou_gopsutil_v3-v3.21.12.zip": "ea6f8b430cee40870d8d454aaa5d4c22e84d217a2548a3f755b91a96b1c67a88", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shopspring/decimal/com_github_shopspring_decimal-v1.2.0.zip": "65c34c248e7f736cadf03a7caa0c0870d15499eb593f933fe106c96c2b7699a7", + "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shubhamdhama/drpc/com_github_shubhamdhama_drpc-v0.0.0-20260511155932-95375d1e96eb.zip": "7c658dfbb3603b4351e4d0d91b91a9787d868348915582f547446e9f96b6ac11", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shurcooL/httpfs/com_github_shurcool_httpfs-v0.0.0-20190707220628-8d4bc4ba7749.zip": "a2079dbd8c236262ecbb22312467265fbbddd9b5ee789531c5f7f24fbdda174b", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shurcooL/sanitized_anchor_name/com_github_shurcool_sanitized_anchor_name-v1.0.0.zip": "0af034323e0627a9e94367f87aa50ce29e5b165d54c8da2926cbaffd5834f757", "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/shurcooL/vfsgen/com_github_shurcool_vfsgen-v0.0.0-20200824052919-0d455de96546.zip": "98198ecd8f122d1266ff2db193f1aae8a88f2f299bfc34b06ef356694cca537d", diff --git a/go.mod b/go.mod index 12ab5dea0082..ad5f5b1a42f9 100644 --- a/go.mod +++ b/go.mod @@ -561,7 +561,7 @@ replace github.com/docker/docker => github.com/moby/moby v25.0.14+incompatible replace github.com/gogo/protobuf => github.com/cockroachdb/gogoproto v1.3.3-0.20241216150617-2358cdb156a1 -replace storj.io/drpc => github.com/cockroachdb/drpc v0.0.0-20260406142218-6c77a9e470d3 +replace storj.io/drpc => github.com/shubhamdhama/drpc v0.0.0-20260511155932-95375d1e96eb // Note: This forked dependency adds a commit that opens up some // private APIs to enable us to make some perf improvements to diff --git a/go.sum b/go.sum index 39426f1d2668..a5343cfb2815 100644 --- a/go.sum +++ b/go.sum @@ -571,8 +571,6 @@ github.com/cockroachdb/datadriven v0.0.0-20190809214429-80d97fb3cbaa/go.mod h1:z github.com/cockroachdb/datadriven v1.0.2/go.mod h1:a9RdTaap04u637JoCzcUoIcDmvwSUtcUFtT/C3kJlTU= github.com/cockroachdb/datadriven v1.0.3-0.20251123150250-ddff6747b112 h1:T1++5Vt0/4/IWZ1mHmUYl7fhQnz50QhNWIY+ITvLLIM= github.com/cockroachdb/datadriven v1.0.3-0.20251123150250-ddff6747b112/go.mod h1:jsaKMvD3RBCATk1/jbUZM8C9idWBJME9+VRZ5+Liq1g= -github.com/cockroachdb/drpc v0.0.0-20260406142218-6c77a9e470d3 h1:Tn7JOQhDXUACzk732VOFSlh9jOlTYEcpZS23MHwyorI= -github.com/cockroachdb/drpc v0.0.0-20260406142218-6c77a9e470d3/go.mod h1:jeOpcnHh3/TvQ7K2kYXyLUMOMoMpuxp5vtdmtNPVW+8= github.com/cockroachdb/errors v1.12.1-0.20251010171200-64801262cd6f h1:lUmJxzb2/ukuRIvKTaNkvuj5LwlX4u/KxnI3zmx1SSw= github.com/cockroachdb/errors v1.12.1-0.20251010171200-64801262cd6f/go.mod h1:SvzfYNNBshAVbZ8wzNc/UPK3w1vf0dKDUP41ucAIf7g= github.com/cockroachdb/go-test-teamcity v0.0.0-20191211140407-cff980ad0a55 h1:YqzBA7tf8Gv8Oz0BbBsPenqkyjiohS7EUIwi7p1QJCU= @@ -2150,6 +2148,8 @@ github.com/shirou/gopsutil/v3 v3.21.12/go.mod h1:BToYZVTlSVlfazpDDYFnsVZLaoRG+g8 github.com/shopspring/decimal v0.0.0-20180709203117-cd690d0c9e24/go.mod h1:M+9NzErvs504Cn4c5DxATwIqPbtswREoFCre64PpcG4= github.com/shopspring/decimal v1.2.0 h1:abSATXmQEYyShuxI4/vyW3tV1MrKAJzCZ/0zLUXYbsQ= github.com/shopspring/decimal v1.2.0/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o= +github.com/shubhamdhama/drpc v0.0.0-20260511155932-95375d1e96eb h1:40AIbuzpZfhXsWfARr5M3YqcRLCHLs/rhH16EDq0ZCg= +github.com/shubhamdhama/drpc v0.0.0-20260511155932-95375d1e96eb/go.mod h1:jeOpcnHh3/TvQ7K2kYXyLUMOMoMpuxp5vtdmtNPVW+8= github.com/shurcooL/httpfs v0.0.0-20190707220628-8d4bc4ba7749/go.mod h1:ZY1cvUeJuFPAdZ/B6v7RHavJWZn2YPVFQ1OSXhCGOkg= github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc= github.com/shurcooL/vfsgen v0.0.0-20181202132449-6a9ea43bcacd/go.mod h1:TrYk7fJVaAttu97ZZKrO9UbRa8izdowaMIZcxYMbVaw= diff --git a/pkg/rpc/drpc.go b/pkg/rpc/drpc.go index 5d6c8ce74db7..7780003b488d 100644 --- a/pkg/rpc/drpc.go +++ b/pkg/rpc/drpc.go @@ -15,6 +15,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/kv/kvpb" "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/rpc/rpcbase" + "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/cockroach/pkg/util/stop" "github.com/cockroachdb/cockroach/pkg/util/tracing/drpcinterceptor" @@ -30,6 +31,8 @@ import ( "storj.io/drpc/drpcwire" ) +var envExperimentalDRPCMuxEnabled = envutil.EnvOrDefaultBool("COCKROACH_EXPERIMENTAL_DRPC_MUX_ENABLED", false) + // Default idle connection timeout for DRPC connections in the pool. var defaultDRPCConnIdleTimeout = 5 * time.Minute @@ -60,9 +63,16 @@ func DialDRPC( return ShouldRecordRequestMetricsDRPC(rpcCtx.Settings) } - drpcDialOptions = append(drpcDialOptions, drpcclient.WithMetrics(cm)) drpcDialOptions = append(drpcDialOptions, - drpcclient.WithShouldRecordFunc(shouldRecordFunc)) + drpcclient.WithMetrics(cm), + drpcclient.WithShouldRecordFunc(shouldRecordFunc), + ) + + if envExperimentalDRPCMuxEnabled { + log.Dev.Infof(ctx, "dialing DRPC mux connection to %s", target) + return dialDRPCMux(ctx, target, drpcDialOptions) + } + log.Dev.Infof(ctx, "dialing DRPC non-mux connection to %s", target) drpcPoolMetrics := rpcCtx.DRPCPoolMetrics() // TODO(server): could use connection class instead of empty key here. @@ -93,6 +103,18 @@ func DialDRPC( } } +// dialDRPCMux establishes a single DRPC transport connection that carries +// all streams for the peer multiplexed over it. +func dialDRPCMux( + ctx context.Context, target string, drpcDialOptions []drpcclient.DialOption, +) (drpc.Conn, error) { + conn, err := drpcclient.DialContext(ctx, target, drpcDialOptions...) + if err != nil { + return nil, err + } + return drpcclient.NewClientConnWithOptions(ctx, conn, drpcDialOptions...) +} + // drpcDialOptionsInternal is similar to grpcDialOptionsInternal but for // drpc connections. func (rpcCtx *Context) drpcDialOptionsInternal( diff --git a/pkg/server/BUILD.bazel b/pkg/server/BUILD.bazel index 5f4cf039ead9..36a8d0b44927 100644 --- a/pkg/server/BUILD.bazel +++ b/pkg/server/BUILD.bazel @@ -665,7 +665,6 @@ go_test( "@io_storj_drpc//:drpc", "@io_storj_drpc//drpcclient", "@io_storj_drpc//drpcconn", - "@io_storj_drpc//drpcmanager", "@io_storj_drpc//drpcmigrate", "@org_golang_google_grpc//:grpc", "@org_golang_google_grpc//codes", diff --git a/pkg/server/drpc_test.go b/pkg/server/drpc_test.go index b5c8d46eda61..76d471d1249a 100644 --- a/pkg/server/drpc_test.go +++ b/pkg/server/drpc_test.go @@ -24,7 +24,6 @@ import ( "storj.io/drpc" "storj.io/drpc/drpcclient" "storj.io/drpc/drpcconn" - "storj.io/drpc/drpcmanager" "storj.io/drpc/drpcmigrate" ) @@ -115,11 +114,7 @@ func TestStreamContextCancel(t *testing.T) { tlsCfg = tlsCfg.Clone() tlsCfg.ServerName = "*.local" tlsConn := tls.Client(rawconn, tlsCfg) - conn := drpcconn.NewWithOptions(tlsConn, drpcconn.Options{ - Manager: drpcmanager.Options{ - SoftCancel: true, // don't close the transport when stream context is canceled - }, - }) + conn := drpcconn.NewWithOptions(tlsConn, drpcconn.Options{}) defer func() { require.NoError(t, conn.Close()) }() From cc1edf7d48e71e67c3d0344f10b09fec2818d086 Mon Sep 17 00:00:00 2001 From: Shubham Dhama Date: Wed, 13 May 2026 08:43:08 +0000 Subject: [PATCH 4/4] do-not-merge: always on DRPC for unit tests --- pkg/testutils/serverutils/test_server_shim.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pkg/testutils/serverutils/test_server_shim.go b/pkg/testutils/serverutils/test_server_shim.go index cd2e065c1ced..2a7f5f192b2a 100644 --- a/pkg/testutils/serverutils/test_server_shim.go +++ b/pkg/testutils/serverutils/test_server_shim.go @@ -626,8 +626,7 @@ func ShouldEnableDRPC( // Benchmarks need deterministic behavior; skip random DRPC selection. return base.TestDRPCDisabled } - rng, _ := randutil.NewTestRand() - enableDRPC = rng.Intn(2) == 0 + enableDRPC = true case base.TestDRPCUnset: return base.TestDRPCDisabled }