Merge #148489 #148501

craig[bot] · shaikzakiriitm · aa-joshi · craig[bot] · commit 8917e45cf3b4 · 2025-06-23T13:13:22.000Z
148489: rpc: wire drpc client interceptors in dialer r=cthumuluru-crdb a=shaikzakiriitm As part of this change, we use clientConn object to wire in client interceptors. clientConn has a supplier which provides `drpc.Conn` connection from a drpc connection pool. The pooled connection object previously returned, encapsulated within closeEntirePoolConn is now replaced with drpcclient.ClientConn so that client interceptors are applied to it. Epic: CRDB-51616 Fixes: #148430 Part of: CRDB-49359 Release note: None 148501: cli(debug.zip): update Datadog link for table dumps r=aa-joshi a=aa-joshi This patch updates the Datadog dashboard and log explorer links which are generated as part of debug.zip Datadog upload. The table dump dashboard link contains filtering based on upload_id along with default timeframe of 1 month. Epic: None Part of: CRDB-51113 Release note: None Co-authored-by: Shaik Zakir Hussain <zakir.shaik@cockroachlabs.com> Co-authored-by: Akshay Joshi <akshay@cockroachlabs.com>
diff --git a/DEPS.bzl b/DEPS.bzl
@@ -11129,10 +11129,10 @@ def go_deps():
         name = "io_storj_drpc",
         build_file_proto_mode = "disable_global",
         importpath = "storj.io/drpc",
-        sha256 = "67a517d0a6c90586f265135ab60619a06a3be5b5c871c15093865716a6de5b39",
-        strip_prefix = "github.com/cockroachdb/drpc@v0.0.0-20250603054748-5b0c5d2c7b38",
+        sha256 = "612016b7a145f386a2163d788e3376cb0f63410b8423f2ef3f723f7aa24c1971",
+        strip_prefix = "github.com/cockroachdb/drpc@v0.0.0-20250618091105-e79a954a2193",
         urls = [
-            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250603054748-5b0c5d2c7b38.zip",
+            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250618091105-e79a954a2193.zip",
         ],
     )
     go_repository(
diff --git a/build/bazelutil/distdir_files.bzl b/build/bazelutil/distdir_files.bzl
@@ -349,7 +349,7 @@ 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-20241205160938-4a90b184f49c.zip": "1afc910b4ff270de79eecb42ab7bd5e6404e6128666c6c55e96db9e27d28e69e",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/datadriven/com_github_cockroachdb_datadriven-v1.0.3-0.20250407164829-2945557346d5.zip": "251593cd9c040fe84a99a3919de7ce6f85030d522159a37d625dc2dea7a4d17f",
-    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250603054748-5b0c5d2c7b38.zip": "67a517d0a6c90586f265135ab60619a06a3be5b5c871c15093865716a6de5b39",
+    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/drpc/com_github_cockroachdb_drpc-v0.0.0-20250618091105-e79a954a2193.zip": "612016b7a145f386a2163d788e3376cb0f63410b8423f2ef3f723f7aa24c1971",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/errors/com_github_cockroachdb_errors-v1.12.0.zip": "f73d8a5f4d8fcbc4ed61db2b47f17e2601d8b32e9a49c0665667489d9d9d6e7c",
     "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",
diff --git a/go.mod b/go.mod
@@ -501,7 +501,7 @@ replace golang.org/x/time => github.com/cockroachdb/x-time v0.3.1-0.202305251236
 
 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-20250603054748-5b0c5d2c7b38
+replace storj.io/drpc => github.com/cockroachdb/drpc v0.0.0-20250618091105-e79a954a2193
 
 // 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
@@ -562,8 +562,8 @@ 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.20250407164829-2945557346d5 h1:UycK/E0TkisVrQbSoxvU827FwgBBcZ95nRRmpj/12QI=
 github.com/cockroachdb/datadriven v1.0.3-0.20250407164829-2945557346d5/go.mod h1:jsaKMvD3RBCATk1/jbUZM8C9idWBJME9+VRZ5+Liq1g=
-github.com/cockroachdb/drpc v0.0.0-20250603054748-5b0c5d2c7b38 h1:7FZw1FKZIOabgsOYY/D8cSj60E7B7h9Tm7xSos/lWF8=
-github.com/cockroachdb/drpc v0.0.0-20250603054748-5b0c5d2c7b38/go.mod h1:UWP+upGv1Z+4nWxcdwhv3/wQXSOgCZteytaRVez5PDc=
+github.com/cockroachdb/drpc v0.0.0-20250618091105-e79a954a2193 h1:c/lx52g2cnFUVPkHqYX+6Nyh+3L/l5ae0SSrkvrUM4M=
+github.com/cockroachdb/drpc v0.0.0-20250618091105-e79a954a2193/go.mod h1:UWP+upGv1Z+4nWxcdwhv3/wQXSOgCZteytaRVez5PDc=
 github.com/cockroachdb/errors v1.9.1/go.mod h1:2sxOtL2WIc096WSZqZ5h8fa17rdDq9HZOZLBCor4mBk=
 github.com/cockroachdb/errors v1.12.0 h1:d7oCs6vuIMUQRVbi6jWWWEJZahLCfJpnJSVobd1/sUo=
 github.com/cockroachdb/errors v1.12.0/go.mod h1:SvzfYNNBshAVbZ8wzNc/UPK3w1vf0dKDUP41ucAIf7g=
diff --git a/pkg/cli/testdata/upload/tables b/pkg/cli/testdata/upload/tables
@@ -34,6 +34,6 @@ Logs API Hook: https://http-intake.logs.us5.datadoghq.com/api/v2/logs
 Logs API Hook: https://http-intake.logs.us5.datadoghq.com/api/v2/logs
 Logs API Hook: https://http-intake.logs.us5.datadoghq.com/api/v2/logs
 Upload ID: abc-20241114000000
-View as logs here: https://us5.datadoghq.com/logs?query=source:debug-zip&upload_id:abc-20241114000000
-View as tables here: https://us5.datadoghq.com/dashboard/ipq-44t-ez8/table-dumps-from-debug-zip?tpl_var_upload_id%5B0%5D=abc-20241114000000
+View as logs here: https://us5.datadoghq.com/logs?query=source:debug-zip upload_id:abc-20241114000000&from_ts=1728950400000&to_ts=1731542400000
+View as tables here:https://us5.datadoghq.com/dashboard/jrz-h9w-5em/table-dumps-from-debug-zip?tpl_var_upload_id=abc-20241114000000&from_ts=1728950400000&to_ts=1731542400000
 debug zip upload debugDir --dd-api-key=dd-api-key --dd-app-key=dd-app-key --cluster=ABC --include=tables
diff --git a/pkg/cli/zip_upload.go b/pkg/cli/zip_upload.go
@@ -752,8 +752,15 @@ func uploadZipTables(ctx context.Context, uploadID string, debugDirPath string)
 	uploadWG.Wait()
 	close(uploadChan)
 
-	fmt.Printf("\nView as tables here: https://us5.datadoghq.com/dashboard/ipq-44t-ez8/table-dumps-from-debug-zip?tpl_var_upload_id%%5B0%%5D=%s\n", uploadID)
-	fmt.Printf("View as logs here: https://us5.datadoghq.com/logs?query=source:debug-zip&upload_id:%s\n", uploadID)
+	toUnixTimestamp := getCurrentTime().UnixMilli()
+	//create timestamp for T-30 days.
+	fromUnixTimestamp := toUnixTimestamp - (30 * 24 * 60 * 60 * 1000)
+
+	fmt.Printf("\nView as tables here:"+
+		"https://us5.datadoghq.com/dashboard/jrz-h9w-5em/table-dumps-from-debug-zip?tpl_var_upload_id=%s&from_ts=%d&to_ts=%d\n",
+		uploadID, fromUnixTimestamp, toUnixTimestamp)
+	fmt.Printf("View as logs here: https://us5.datadoghq.com/logs?query=source:debug-zip upload_id:%s&from_ts=%d&to_ts=%d\n",
+		uploadID, fromUnixTimestamp, toUnixTimestamp)
 	return nil
 }
 
diff --git a/pkg/cli/zip_upload_test.go b/pkg/cli/zip_upload_test.go
@@ -212,6 +212,10 @@ func TestUploadZipEndToEnd(t *testing.T) {
 				includeFlag = "--include=misc"
 			}
 
+			defer testutils.TestingHook(&getCurrentTime, func() time.Time {
+				return time.Date(2024, 11, 14, 0, 0, 0, 0, time.UTC)
+			})()
+
 			debugDir, cleanup := setupZipDir(t, testInput)
 			defer cleanup()
 
diff --git a/pkg/rpc/BUILD.bazel b/pkg/rpc/BUILD.bazel
@@ -77,6 +77,7 @@ go_library(
         "@com_github_vividcortex_ewma//:ewma",
         "@io_opentelemetry_go_otel//attribute",
         "@io_storj_drpc//:drpc",
+        "@io_storj_drpc//drpcclient",
         "@io_storj_drpc//drpcconn",
         "@io_storj_drpc//drpcmanager",
         "@io_storj_drpc//drpcmigrate",
diff --git a/pkg/rpc/context.go b/pkg/rpc/context.go
@@ -52,6 +52,7 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 	"storj.io/drpc"
+	"storj.io/drpc/drpcclient"
 )
 
 // NewServer sets up an RPC server. Depending on the ServerOptions, the Server
@@ -251,6 +252,9 @@ type Context struct {
 	clientUnaryInterceptors  []grpc.UnaryClientInterceptor
 	clientStreamInterceptors []grpc.StreamClientInterceptor
 
+	clientUnaryInterceptorsDRPC  []drpcclient.UnaryClientInterceptor
+	clientStreamInterceptorsDRPC []drpcclient.StreamClientInterceptor
+
 	// loopbackDialFn, when non-nil, is used when the target of the dial
 	// is ourselves (== AdvertiseAddr).
 	//
diff --git a/pkg/rpc/context_testutils.go b/pkg/rpc/context_testutils.go
@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
+	"storj.io/drpc/drpcclient"
 )
 
 // ContextTestingKnobs provides hooks to aid in testing the system. The testing
@@ -38,6 +39,10 @@ type ContextTestingKnobs struct {
 	// unary RPC.
 	UnaryClientInterceptor func(target string, class rpcbase.ConnectionClass) grpc.UnaryClientInterceptor
 
+	UnaryClientInterceptorDRPC func(target string, class rpcbase.ConnectionClass) drpcclient.UnaryClientInterceptor
+
+	StreamClientInterceptorDRPC func(target string, class rpcbase.ConnectionClass) drpcclient.StreamClientInterceptor
+
 	// InjectedLatencyOracle if non-nil contains a map from target address
 	// (server.RPCServingAddr() of a remote node) to artificial latency in
 	// milliseconds to inject. Setting this will cause the server to pause for
diff --git a/pkg/rpc/drpc.go b/pkg/rpc/drpc.go
@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/errors"
 	"storj.io/drpc"
+	"storj.io/drpc/drpcclient"
 	"storj.io/drpc/drpcconn"
 	"storj.io/drpc/drpcmanager"
 	"storj.io/drpc/drpcmigrate"
@@ -40,7 +41,8 @@ func (d *drpcCloseNotifier) CloseNotify(ctx context.Context) <-chan struct{} {
 	return d.conn.Closed()
 }
 
-func dialDRPC(
+// TODO DB Server: unexport this once dial methods are added in rpccontext.
+func DialDRPC(
 	rpcCtx *Context,
 ) func(ctx context.Context, target string, _ rpcbase.ConnectionClass) (drpc.Conn, error) {
 	return func(ctx context.Context, target string, _ rpcbase.ConnectionClass) (drpc.Conn, error) {
@@ -86,13 +88,27 @@ func dialDRPC(
 
 			return conn, nil
 		})
-		// `pooledConn.Close` doesn't tear down any of the underlying TCP
-		// connections but simply marks the pooledConn handle as returning
-		// errors. When we "close" this conn, we want to tear down all of
-		// the connections in the pool (in effect mirroring the behavior of
-		// gRPC where a single conn is shared).
+
+		if rpcCtx.Knobs.UnaryClientInterceptorDRPC != nil {
+			if interceptor := rpcCtx.Knobs.UnaryClientInterceptorDRPC(target, rpcbase.DefaultClass); interceptor != nil {
+				rpcCtx.clientUnaryInterceptorsDRPC = append(rpcCtx.clientUnaryInterceptorsDRPC, interceptor)
+			}
+		}
+		if rpcCtx.Knobs.StreamClientInterceptorDRPC != nil {
+			if interceptor := rpcCtx.Knobs.StreamClientInterceptorDRPC(target, rpcbase.DefaultClass); interceptor != nil {
+				rpcCtx.clientStreamInterceptorsDRPC = append(rpcCtx.clientStreamInterceptorsDRPC, interceptor)
+			}
+		}
+		clientConn, _ := drpcclient.NewClientConnWithOptions(
+			ctx,
+			pooledConn,
+			drpcclient.WithChainUnaryInterceptor(rpcCtx.clientUnaryInterceptorsDRPC...),
+			drpcclient.WithChainStreamInterceptor(rpcCtx.clientStreamInterceptorsDRPC...),
+		)
+
+		// Wrap the clientConn to ensure the entire pool is closed when this connection handle is closed.
 		return &closeEntirePoolConn{
-			Conn: pooledConn,
+			Conn: clientConn,
 			pool: pool,
 		}, nil
 	}
@@ -129,7 +145,7 @@ type drpcServer struct {
 }
 
 // NewDRPCServer creates a new DRPCServer with the provided rpc context.
-func NewDRPCServer(_ context.Context, rpcCtx *Context) (DRPCServer, error) {
+func NewDRPCServer(_ context.Context, _ *Context) (DRPCServer, error) {
 	d := &drpcServer{}
 	mux := drpcmux.New()
 	d.Server = drpcserver.NewWithOptions(mux, drpcserver.Options{
diff --git a/pkg/rpc/peer.go b/pkg/rpc/peer.go
@@ -248,7 +248,7 @@ func newPeer[Conn rpcConn](rpcCtx *Context, k peerKey, peerOpts *peerOptions[Con
 		opts:               &rpcCtx.ContextOptions,
 		peers:              peerOpts.peers,
 		connOptions:        peerOpts.connOptions,
-		drpcDial:           dialDRPC(rpcCtx),
+		drpcDial:           DialDRPC(rpcCtx),
 		newHeartbeatClient: peerOpts.newHeartbeatClient,
 		heartbeatInterval:  rpcCtx.RPCHeartbeatInterval,
 		heartbeatTimeout:   rpcCtx.RPCHeartbeatTimeout,
diff --git a/pkg/server/BUILD.bazel b/pkg/server/BUILD.bazel
@@ -620,6 +620,8 @@ go_test(
         "@com_github_stretchr_testify//require",
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@io_opentelemetry_go_otel//attribute",
+        "@io_storj_drpc//:drpc",
+        "@io_storj_drpc//drpcclient",
         "@io_storj_drpc//drpcconn",
         "@io_storj_drpc//drpcmanager",
         "@io_storj_drpc//drpcmigrate",
diff --git a/pkg/server/drpc_test.go b/pkg/server/drpc_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
@@ -21,6 +22,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
+	"storj.io/drpc"
+	"storj.io/drpc/drpcclient"
 	"storj.io/drpc/drpcconn"
 	"storj.io/drpc/drpcmanager"
 	"storj.io/drpc/drpcmigrate"
@@ -179,3 +182,100 @@ func TestDefaultDRPCOption(t *testing.T) {
 		require.Equal(t, option, enabled)
 	})
 }
+
+func TestDialDRPC_InterceptorsAreSet(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	var unaryInterceptorCalled bool
+	var streamInterceptorCalled bool
+
+	// dummy unary interceptor for testing.
+	mockUnaryInterceptor := func(
+		ctx context.Context,
+		rpc string,
+		enc drpc.Encoding,
+		in, out drpc.Message,
+		cc *drpcclient.ClientConn,
+		invoker drpcclient.UnaryInvoker,
+	) error {
+		unaryInterceptorCalled = true
+		return invoker(ctx, rpc, enc, in, out, cc)
+	}
+
+	// dummy stream interceptor for testing.
+	mockStreamInterceptor := func(
+		ctx context.Context,
+		rpc string,
+		enc drpc.Encoding,
+		cc *drpcclient.ClientConn,
+		streamer drpcclient.Streamer,
+	) (drpc.Stream, error) {
+		streamInterceptorCalled = true
+		return streamer(ctx, rpc, enc, cc)
+	}
+
+	ctx := context.Background()
+	const numNodes = 1
+	args := base.TestClusterArgs{
+		ReplicationMode: base.ReplicationManual,
+		ServerArgs: base.TestServerArgs{
+			Settings: cluster.MakeClusterSettings(),
+			Insecure: true,
+		},
+	}
+
+	rpc.ExperimentalDRPCEnabled.Override(ctx, &args.ServerArgs.Settings.SV, true)
+	c := testcluster.StartTestCluster(t, numNodes, args)
+	defer c.Stopper().Stop(ctx)
+	require.True(t, c.Server(0).RPCContext().Insecure)
+
+	rpcAddr := c.Server(0).RPCAddr()
+
+	// Client setup
+	// Setup a minimal rpcCtx with interceptors
+	rpcContextOptions := rpc.DefaultContextOptions()
+	rpcContextOptions.Stopper = c.Stopper()
+	rpcContextOptions.Settings = c.Server(0).ClusterSettings()
+	rpcCtx := rpc.NewContext(ctx, rpcContextOptions)
+	rpcCtx.ContextOptions = rpc.ContextOptions{Insecure: true}
+	// Adding test interceptors
+	rpcCtx.Knobs = rpc.ContextTestingKnobs{
+		UnaryClientInterceptorDRPC: func(target string, class rpcbase.ConnectionClass) drpcclient.UnaryClientInterceptor {
+			return mockUnaryInterceptor
+		},
+		StreamClientInterceptorDRPC: func(target string, class rpcbase.ConnectionClass) drpcclient.StreamClientInterceptor {
+			return mockStreamInterceptor
+		},
+	}
+	getConn := rpc.DialDRPC(rpcCtx)
+	conn, err := getConn(ctx, rpcAddr, rpcbase.DefaultClass)
+	require.NoError(t, err)
+	defer func() { require.NoError(t, conn.Close()) }()
+	desc := c.LookupRangeOrFatal(t, c.ScratchRange(t))
+
+	// reset flag values
+	unaryInterceptorCalled = false
+	streamInterceptorCalled = false
+
+	client := kvpb.NewDRPCKVBatchClient(conn)
+	ba := &kvpb.BatchRequest{}
+	ba.RangeID = desc.RangeID
+	var ok bool
+	ba.Replica, ok = desc.GetReplicaDescriptor(1)
+	require.True(t, ok)
+	req := &kvpb.LeaseInfoRequest{}
+	req.Key = desc.StartKey.AsRawKey()
+	ba.Add(req)
+	_, err = client.Batch(ctx, ba)
+	require.NoError(t, err)
+	s, err := client.BatchStream(ctx)
+	require.NoError(t, err)
+	err = s.Send(ba)
+	require.NoError(t, err)
+	_, err = s.Recv()
+	require.NoError(t, err)
+	// assert that the interceptors were called
+	require.True(t, unaryInterceptorCalled)
+	require.True(t, streamInterceptorCalled)
+}

Original file line number	Diff line number	Diff line change
`@@ -212,6 +212,10 @@ func TestUploadZipEndToEnd(t *testing.T) {`
`212`	`212`	`includeFlag = "--include=misc"`
`213`	`213`	`}`
`214`	`214`
	`215`	`+ defer testutils.TestingHook(&getCurrentTime, func() time.Time {`
	`216`	`+ return time.Date(2024, 11, 14, 0, 0, 0, 0, time.UTC)`
	`217`	`+ })()`
	`218`	`+`
`215`	`219`	`debugDir, cleanup := setupZipDir(t, testInput)`
`216`	`220`	`defer cleanup()`
`217`	`221`