Merge #153012

153012: rpc: drpc interceptors run task within stopper r=shubhamdhama a=Nukitt

This change adds support for running tasks within stoppers for DRPC interceptors. 
Now, DRPC unary and streaming RPCs run their handlers inside Stopper tasks, ensuring requests are tracked and rejected once draining begins.

Epic: CRDB-49359
Fixes: #144371
Release note: None

Co-authored-by: Nukitt <nukit.tailor@cockroachlabs.com>

Author: craig[bot]

pkg/rpc/BUILD.bazel

@@ -114,6 +114,7 @@ go_test(
         "context_test.go",
         "datadriven_test.go",
         "down_node_test.go",
+        "drpc_test.go",
         "heartbeat_test.go",
         "helpers_test.go",
         "main_test.go",
@@ -180,6 +181,7 @@ go_test(
         "@com_github_prometheus_client_model//go",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
+        "@io_storj_drpc//:drpc",
         "@io_storj_drpc//drpcctx",
         "@io_storj_drpc//drpcmetadata",
         "@io_storj_drpc//drpcmux",

pkg/rpc/drpc.go

@@ -161,6 +161,41 @@ type drpcServer struct {
 	drpc.Mux
 }
 
+// makeStopperInterceptors returns unary and stream interceptors that run
+// incoming RPCs in stopper tasks.
+func makeStopperInterceptors(
+	rpcCtx *Context,
+) (drpcmux.UnaryServerInterceptor, drpcmux.StreamServerInterceptor) {
+	unary := func(
+		ctx context.Context, req interface{}, rpc string, handler drpcmux.UnaryHandler,
+	) (interface{}, error) {
+		var resp interface{}
+		if err := rpcCtx.Stopper.RunTaskWithErr(ctx, rpc, func(ctx context.Context) error {
+			var err error
+			resp, err = handler(ctx, req)
+			return err
+		}); err != nil {
+			return nil, err
+		}
+		return resp, nil
+	}
+
+	stream := func(
+		stream drpc.Stream, rpc string, handler drpcmux.StreamHandler,
+	) (interface{}, error) {
+		var resp interface{}
+		if err := rpcCtx.Stopper.RunTaskWithErr(stream.Context(), rpc, func(ctx context.Context) error {
+			var err error
+			resp, err = handler(stream)
+			return err
+		}); err != nil {
+			return nil, err
+		}
+		return resp, nil
+	}
+	return unary, stream
+}
+
 // NewDRPCServer creates a new DRPCServer with the provided rpc context.
 func NewDRPCServer(_ context.Context, rpcCtx *Context, opts ...ServerOption) (DRPCServer, error) {
 	d := &drpcServer{}
@@ -173,6 +208,16 @@ func NewDRPCServer(_ context.Context, rpcCtx *Context, opts ...ServerOption) (DR
 	var unaryInterceptors []drpcmux.UnaryServerInterceptor
 	var streamInterceptors []drpcmux.StreamServerInterceptor
 
+	// These interceptors run in the order they're appended. The first
+	// interceptor added becomes the outermost wrapper around the handler.
+
+	// We start with an interceptor that ensures every RPC executes inside a
+	// stopper task. Running the handler in a stopper task lets the stopper
+	// keep track of in-flight RPCs and reject new ones once draining begins.
+	stopUnary, stopStream := makeStopperInterceptors(rpcCtx)
+	unaryInterceptors = append(unaryInterceptors, stopUnary)
+	streamInterceptors = append(streamInterceptors, stopStream)
+
 	if !rpcCtx.ContextOptions.Insecure {
 		a := kvAuth{
 			sv: &rpcCtx.Settings.SV,
@@ -199,19 +244,6 @@ func NewDRPCServer(_ context.Context, rpcCtx *Context, opts ...ServerOption) (DR
 	})
 	d.Mux = mux
 
-	// NB: any server middleware (server interceptors in gRPC parlance) would go
-	// here:
-	//     dmux = whateverMiddleware1(dmux)
-	//     dmux = whateverMiddleware2(dmux)
-	//     ...
-	//
-	// Each middleware must implement the Handler interface:
-	//
-	//   HandleRPC(stream Stream, rpc string) error
-	//
-	// where Stream
-	// See here for an example:
-	// https://github.com/bryk-io/pkg/blob/4da5fbfef47770be376e4022eab5c6c324984bf7/net/drpc/server.go#L91-L101
 	return d, nil
 }
 

pkg/rpc/drpc_test.go

@@ -0,0 +1,79 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package rpc
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
+	"github.com/stretchr/testify/require"
+	"storj.io/drpc"
+)
+
+// dummyStream is a minimal implementation of drpc.Stream used for testing the
+// stream interceptor.
+type dummyStream struct {
+	ctx context.Context
+}
+
+func (d dummyStream) Context() context.Context                  { return d.ctx }
+func (d dummyStream) MsgSend(drpc.Message, drpc.Encoding) error { return nil }
+func (d dummyStream) MsgRecv(drpc.Message, drpc.Encoding) error { return nil }
+func (d dummyStream) CloseSend() error                          { return nil }
+func (d dummyStream) Close() error                              { return nil }
+
+// TestMakeStopperInterceptors verifies that the stopper interceptors allow RPCs
+// to run before the stopper quiesces and reject them afterward.
+func TestMakeStopperInterceptors(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
+
+	rpcCtx := &Context{ContextOptions: ContextOptions{Stopper: stopper}}
+
+	unaryInterceptor, streamInterceptor := makeStopperInterceptors(rpcCtx)
+
+	// Before quiesce runs.
+	called := false
+	_, err := unaryInterceptor(ctx, nil, "test", func(ctx context.Context, req interface{}) (interface{}, error) {
+		called = true
+		return nil, nil
+	})
+	require.NoError(t, err)
+	require.True(t, called)
+
+	called = false
+	_, err = streamInterceptor(dummyStream{ctx: ctx}, "test", func(stream drpc.Stream) (interface{}, error) {
+		called = true
+		return nil, nil
+	})
+	require.NoError(t, err)
+	require.True(t, called)
+
+	// After quiesce, RPCs are rejected.
+	stopper.Quiesce(ctx)
+
+	called = false
+	_, err = unaryInterceptor(ctx, nil, "test", func(ctx context.Context, req interface{}) (interface{}, error) {
+		called = true
+		return nil, nil
+	})
+	require.ErrorIs(t, err, stop.ErrUnavailable)
+	require.False(t, called)
+
+	called = false
+	_, err = streamInterceptor(dummyStream{ctx: ctx}, "test", func(stream drpc.Stream) (interface{}, error) {
+		called = true
+		return nil, nil
+	})
+	require.ErrorIs(t, err, stop.ErrUnavailable)
+	require.False(t, called)
+}