kvserver: register MultiRaft service with DRPC server

Enable MultiRaft service on the DRPC server in addition to gRPC. Controlled
by `rpc.experimental_drpc.enabled` (off by default).

Epic: CRDB-48925
Release note: None

Author: shubhamdhama

pkg/kv/kvserver/BUILD.bazel

@@ -262,6 +262,7 @@ go_library(
         "@com_github_prometheus_client_model//go",
         "@com_github_raduberinde_btreemap//:btreemap",
         "@io_opentelemetry_go_otel//attribute",
+        "@io_storj_drpc//:drpc",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_x_time//rate",
     ],

pkg/kv/kvserver/client_merge_test.go

@@ -2482,6 +2482,7 @@ func TestStoreReplicaGCAfterMerge(t *testing.T) {
 		nodedialer.New(tc.Servers[0].RPCContext(),
 			gossip.AddressResolver(tc.Servers[0].GossipI().(*gossip.Gossip))),
 		nil, /* grpcServer */
+		nil, /* drpcServer */
 		(*node_rac2.AdmittedPiggybacker)(nil),
 		nil, /* PiggybackedAdmittedResponseScheduler */
 		nil, /* knobs */

pkg/kv/kvserver/client_raft_test.go

@@ -3436,6 +3436,7 @@ func TestReplicaGCRace(t *testing.T) {
 		tc.Servers[0].Clock(),
 		nodedialer.New(tc.Servers[0].RPCContext(), gossip.AddressResolver(fromStore.Gossip())),
 		nil, /* grpcServer */
+		nil, /* drpcServer */
 		(*node_rac2.AdmittedPiggybacker)(nil),
 		nil, /* PiggybackedAdmittedResponseScheduler */
 		nil, /* knobs */
@@ -3828,6 +3829,7 @@ func TestReplicateRemovedNodeDisruptiveElection(t *testing.T) {
 		nodedialer.New(tc.Servers[0].RPCContext(),
 			gossip.AddressResolver(tc.GetFirstStoreFromServer(t, 0).Gossip())),
 		nil, /* grpcServer */
+		nil, /* drpcServer */
 		(*node_rac2.AdmittedPiggybacker)(nil),
 		nil, /* PiggybackedAdmittedResponseScheduler */
 		nil, /* knobs */

pkg/kv/kvserver/raft_transport.go

@@ -36,6 +36,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
+	drpc "storj.io/drpc"
 )
 
 const (
@@ -78,7 +79,7 @@ type RaftMessageResponseStream interface {
 // Send. Note that the default implementation of grpc.Stream for server
 // responses (grpc.serverStream) is not safe for concurrent calls to Send.
 type lockedRaftMessageResponseStream struct {
-	wrapped MultiRaft_RaftMessageBatchServer
+	wrapped RPCMultiRaft_RaftMessageBatchStream
 	sendMu  syncutil.Mutex
 }
 
@@ -235,9 +236,7 @@ func NewDummyRaftTransport(
 	resolver := func(roachpb.NodeID) (net.Addr, roachpb.Locality, error) {
 		return nil, roachpb.Locality{}, errors.New("dummy resolver")
 	}
-	return NewRaftTransport(ambient, st, nil, clock, nodedialer.New(nil, resolver), nil,
-		nil, nil, nil,
-	)
+	return NewRaftTransport(ambient, st, nil, clock, nodedialer.New(nil, resolver), nil, nil, nil, nil, nil)
 }
 
 // NewRaftTransport creates a new RaftTransport.
@@ -248,6 +247,7 @@ func NewRaftTransport(
 	clock *hlc.Clock,
 	dialer *nodedialer.Dialer,
 	grpcServer *grpc.Server,
+	drpcMux drpc.Mux,
 	piggybackReader node_rac2.PiggybackMsgReader,
 	piggybackedResponseScheduler PiggybackedAdmittedResponseScheduler,
 	knobs *RaftTransportTestingKnobs,
@@ -271,6 +271,9 @@ func NewRaftTransport(
 	if grpcServer != nil {
 		RegisterMultiRaftServer(grpcServer, t)
 	}
+	if drpcMux != nil {
+		_ = DRPCRegisterMultiRaft(drpcMux, t.AsDRPCServer())
+	}
 	return t
 }
 
@@ -375,8 +378,29 @@ func newRaftMessageResponse(
 	return resp
 }
 
+type drpcRaftTransport RaftTransport
+
+// AsDRPCServer returns the DRPC server implementation for the Raft service.
+func (t *RaftTransport) AsDRPCServer() DRPCMultiRaftServer {
+	return (*drpcRaftTransport)(t)
+}
+
+// RaftMessageBatch proxies the incoming requests to the listening server interface.
+func (t *drpcRaftTransport) RaftMessageBatch(
+	stream DRPCMultiRaft_RaftMessageBatchStream,
+) (lastErr error) {
+	return (*RaftTransport)(t).raftMessageBatch(stream)
+}
+
 // RaftMessageBatch proxies the incoming requests to the listening server interface.
 func (t *RaftTransport) RaftMessageBatch(stream MultiRaft_RaftMessageBatchServer) (lastErr error) {
+	return t.raftMessageBatch(stream)
+}
+
+// raftMessageBatch is the shared implementation for RaftMessageBatch for both gRPC and DRPC.
+func (t *RaftTransport) raftMessageBatch(
+	stream RPCMultiRaft_RaftMessageBatchStream,
+) (lastErr error) {
 	errCh := make(chan error, 1)
 
 	// Node stopping error is caught below in the select.
@@ -444,10 +468,25 @@ func (t *RaftTransport) RaftMessageBatch(stream MultiRaft_RaftMessageBatchServer
 	}
 }
 
+// DelegateRaftSnapshot handles incoming delegated snapshot requests and passes
+// the request to pass off to the sender store. Errors during the snapshots
+// process are sent back as a response.
+func (t *drpcRaftTransport) DelegateRaftSnapshot(
+	stream DRPCMultiRaft_DelegateRaftSnapshotStream,
+) error {
+	return (*RaftTransport)(t).delegateRaftSnapshot(stream)
+}
+
 // DelegateRaftSnapshot handles incoming delegated snapshot requests and passes
 // the request to pass off to the sender store. Errors during the snapshots
 // process are sent back as a response.
 func (t *RaftTransport) DelegateRaftSnapshot(stream MultiRaft_DelegateRaftSnapshotServer) error {
+	return t.delegateRaftSnapshot(stream)
+}
+
+// delegateRaftSnapshot is the shared implementation for DelegateRaftSnapshot
+// for both gRPC and DRPC.
+func (t *RaftTransport) delegateRaftSnapshot(stream RPCMultiRaft_DelegateRaftSnapshotStream) error {
 	ctx, cancel := t.stopper.WithCancelOnQuiesce(stream.Context())
 	defer cancel()
 	req, err := stream.Recv()
@@ -494,8 +533,18 @@ func (t *RaftTransport) InternalDelegateRaftSnapshot(
 	return incomingMessageHandler.HandleDelegatedSnapshot(ctx, req)
 }
 
+// RaftSnapshot handles incoming streaming snapshot requests.
+func (t *drpcRaftTransport) RaftSnapshot(stream DRPCMultiRaft_RaftSnapshotStream) error {
+	return (*RaftTransport)(t).raftSnapshot(stream)
+}
+
 // RaftSnapshot handles incoming streaming snapshot requests.
 func (t *RaftTransport) RaftSnapshot(stream MultiRaft_RaftSnapshotServer) error {
+	return t.raftSnapshot(stream)
+}
+
+// raftSnapshot is the shared implementation for RaftSnapshot for both gRPC and DRPC.
+func (t *RaftTransport) raftSnapshot(stream RPCMultiRaft_RaftSnapshotStream) error {
 	ctx, cancel := t.stopper.WithCancelOnQuiesce(stream.Context())
 	defer cancel()
 	req, err := stream.Recv()
@@ -547,7 +596,7 @@ func (t *RaftTransport) StopOutgoingMessage(storeID roachpb.StoreID) {
 // lost and a new instance of processQueue will be started by the next message
 // to be sent.
 func (t *RaftTransport) processQueue(
-	q *raftSendQueue, stream MultiRaft_RaftMessageBatchClient, _ rpcbase.ConnectionClass,
+	q *raftSendQueue, stream RPCMultiRaft_RaftMessageBatchClient, _ rpcbase.ConnectionClass,
 ) error {
 	errCh := make(chan error, 1)
 

pkg/kv/kvserver/raft_transport_test.go

@@ -182,7 +182,10 @@ func (rttc *raftTransportTestContext) AddNodeWithoutGossip(
 	manual := hlc.NewHybridManualClock()
 	clock := hlc.NewClockForTesting(manual)
 	rttc.clocks[nodeID] = clockWithManualSource{manual: manual, clock: clock}
-	grpcServer, err := rpc.NewServer(context.Background(), rttc.nodeRPCContext)
+	ctx := context.Background()
+	grpcServer, err := rpc.NewServer(ctx, rttc.nodeRPCContext)
+	require.NoError(rttc.t, err)
+	drpcServer, err := rpc.NewDRPCServer(ctx, rttc.nodeRPCContext)
 	require.NoError(rttc.t, err)
 	transport := kvserver.NewRaftTransport(
 		log.MakeTestingAmbientCtxWithNewTracer(),
@@ -191,6 +194,7 @@ func (rttc *raftTransportTestContext) AddNodeWithoutGossip(
 		clock,
 		nodedialer.New(rttc.nodeRPCContext, gossip.AddressResolver(rttc.gossip)),
 		grpcServer,
+		drpcServer,
 		piggybacker,
 		piggybackedResponseScheduler,
 		knobs,

pkg/kv/kvserver/raft_transport_unit_test.go

@@ -53,6 +53,8 @@ func TestRaftTransportStartNewQueue(t *testing.T) {
 
 	grpcServer, err := rpc.NewServer(ctx, rpcC)
 	require.NoError(t, err)
+	drpcServer, err := rpc.NewDRPCServer(ctx, rpcC)
+	require.NoError(t, err)
 	// RegisterMultiRaftServer(grpcServer, mrs)
 
 	var addr net.Addr
@@ -71,6 +73,7 @@ func TestRaftTransportStartNewQueue(t *testing.T) {
 		hlc.NewClockForTesting(nil),
 		nodedialer.New(rpcC, resolver),
 		grpcServer,
+		drpcServer,
 		(*node_rac2.AdmittedPiggybacker)(nil),
 		nil, /* PiggybackedAdmittedResponseScheduler */
 		nil, /* knobs */

pkg/kv/kvserver/store_test.go

@@ -186,7 +186,9 @@ func createTestStoreWithoutStart(
 	}
 	rpcContext := rpc.NewContext(ctx, rpcOpts)
 	stopper.SetTracer(cfg.AmbientCtx.Tracer)
-	server, err := rpc.NewServer(ctx, rpcContext) // never started
+	grpcServer, err := rpc.NewServer(ctx, rpcContext) // never started
+	require.NoError(t, err)
+	drpcServer, err := rpc.NewDRPCServer(ctx, rpcContext) // never started
 	require.NoError(t, err)
 
 	// Some tests inject their own Gossip and StorePool, via
@@ -235,7 +237,8 @@ func createTestStoreWithoutStart(
 		stopper,
 		cfg.Clock,
 		cfg.NodeDialer,
-		server,
+		grpcServer,
+		drpcServer,
 		(*node_rac2.AdmittedPiggybacker)(nil),
 		nil, /* PiggybackedAdmittedResponseScheduler */
 		nil, /* knobs */
@@ -246,7 +249,7 @@ func createTestStoreWithoutStart(
 		supportGracePeriod := rpcContext.StoreLivenessWithdrawalGracePeriod()
 		options := storeliveness.NewOptions(livenessInterval, heartbeatInterval, supportGracePeriod)
 		transport := storeliveness.NewTransport(
-			cfg.AmbientCtx, stopper, cfg.Clock, cfg.NodeDialer, server, nil, /* knobs */
+			cfg.AmbientCtx, stopper, cfg.Clock, cfg.NodeDialer, grpcServer, nil, /* knobs */
 		)
 		knobs := cfg.TestingKnobs.StoreLivenessKnobs
 		cfg.StoreLiveness = storeliveness.NewNodeContainer(stopper, options, transport, knobs)

pkg/server/server.go

@@ -645,6 +645,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		clock,
 		kvNodeDialer,
 		grpcServer.Server,
+		drpcServer.DRPCServer,
 		admittedPiggybacker,
 		storesForRACv2,
 		raftTransportKnobs,