kvserver: register StoreLiveness service with DRPC server

Enable the `StoreLiveness` service on the DRPC server in addition to gRPC.
This is controlled by `rpc.experimental_drpc.enabled` (off by default).

Note: This only registers the service; the client is not updated to use the
DRPC client, so this service will not have any functional effect.

Epic: CRDB-48925
Release note: None

Author: shubhamdhama

pkg/kv/kvserver/store_test.go

@@ -248,9 +248,10 @@ func createTestStoreWithoutStart(
 		livenessInterval, heartbeatInterval := cfg.StoreLivenessDurations()
 		supportGracePeriod := rpcContext.StoreLivenessWithdrawalGracePeriod()
 		options := storeliveness.NewOptions(livenessInterval, heartbeatInterval, supportGracePeriod)
-		transport := storeliveness.NewTransport(
-			cfg.AmbientCtx, stopper, cfg.Clock, cfg.NodeDialer, grpcServer, nil, /* knobs */
+		transport, err := storeliveness.NewTransport(
+			cfg.AmbientCtx, stopper, cfg.Clock, cfg.NodeDialer, grpcServer, drpcServer, nil, /* knobs */
 		)
+		require.NoError(t, err)
 		knobs := cfg.TestingKnobs.StoreLivenessKnobs
 		cfg.StoreLiveness = storeliveness.NewNodeContainer(stopper, options, transport, knobs)
 	}

pkg/kv/kvserver/storeliveness/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_redact//:redact",
+        "@io_storj_drpc//:drpc",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_x_exp//maps",
     ],

pkg/kv/kvserver/storeliveness/transport.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
+	"storj.io/drpc"
 )
 
 const (
@@ -89,8 +90,9 @@ func NewTransport(
 	clock *hlc.Clock,
 	dialer *nodedialer.Dialer,
 	grpcServer *grpc.Server,
+	drpcMux drpc.Mux,
 	knobs *TransportKnobs,
-) *Transport {
+) (*Transport, error) {
 	if knobs == nil {
 		knobs = &TransportKnobs{}
 	}
@@ -105,7 +107,19 @@ func NewTransport(
 	if grpcServer != nil {
 		slpb.RegisterStoreLivenessServer(grpcServer, t)
 	}
-	return t
+	if drpcMux != nil {
+		if err := slpb.DRPCRegisterStoreLiveness(drpcMux, t.AsDRPCServer()); err != nil {
+			return nil, err
+		}
+	}
+	return t, nil
+}
+
+type drpcTransport Transport
+
+// AsDRPCServer returns the DRPC server implementation for the StoreLiveness service.
+func (t *Transport) AsDRPCServer() slpb.DRPCStoreLivenessServer {
+	return (*drpcTransport)(t)
 }
 
 // Metrics returns metrics tracking this transport.
@@ -123,6 +137,18 @@ func (t *Transport) ListenMessages(storeID roachpb.StoreID, handler MessageHandl
 // only within the same RPC call; in other words, this guarantee does not hold
 // across stream re-connects.
 func (t *Transport) Stream(stream slpb.StoreLiveness_StreamServer) error {
+	return t.stream(stream)
+}
+
+// Stream proxies the incoming requests to the corresponding store's
+// MessageHandler. Messages between a pair of nodes are delivered in order
+// only within the same RPC call; in other words, this guarantee does not hold
+// across stream re-connects.
+func (t *drpcTransport) Stream(stream slpb.DRPCStoreLiveness_StreamStream) error {
+	return (*Transport)(t).stream(stream)
+}
+
+func (t *Transport) stream(stream slpb.RPCStoreLiveness_StreamStream) error {
 	errCh := make(chan error, 1)
 
 	// Node stopping error is caught below in the select.

pkg/kv/kvserver/storeliveness/transport_test.go

@@ -119,14 +119,18 @@ func (tt *transportTester) AddNodeWithoutGossip(
 	tt.clocks[nodeID] = clockWithManualSource{manual: manual, clock: clock}
 	grpcServer, err := rpc.NewServer(context.Background(), tt.nodeRPCContext)
 	require.NoError(tt.t, err)
-	transport := NewTransport(
+	drpcServer, err := rpc.NewDRPCServer(context.Background(), tt.nodeRPCContext)
+	require.NoError(tt.t, err)
+	transport, err := NewTransport(
 		log.MakeTestingAmbientCtxWithNewTracer(),
 		tt.stopper,
 		clock,
 		nodedialer.New(tt.nodeRPCContext, gossip.AddressResolver(tt.gossip)),
 		grpcServer,
+		drpcServer,
 		nil, /* knobs */
 	)
+	require.NoError(tt.t, err)
 	tt.transports[nodeID] = transport
 
 	listener, err := netutil.ListenAndServeGRPC(stopper, grpcServer, util.TestAddr)

pkg/server/server.go

@@ -657,9 +657,13 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		livenessInterval, heartbeatInterval := cfg.StoreLivenessDurations()
 		supportGracePeriod := rpcContext.StoreLivenessWithdrawalGracePeriod()
 		options := storeliveness.NewOptions(livenessInterval, heartbeatInterval, supportGracePeriod)
-		transport := storeliveness.NewTransport(
-			cfg.AmbientCtx, stopper, clock, kvNodeDialer, grpcServer.Server, nil, /* knobs */
+		transport, err := storeliveness.NewTransport(
+			cfg.AmbientCtx, stopper, clock, kvNodeDialer,
+			grpcServer.Server, drpcServer.DRPCServer, nil, /* knobs */
 		)
+		if err != nil {
+			return nil, err
+		}
 		nodeRegistry.AddMetricStruct(transport.Metrics())
 		var knobs *storeliveness.TestingKnobs
 		if storeKnobs := cfg.TestingKnobs.Store; storeKnobs != nil {

pkg/testutils/localtestcluster/local_test_cluster.go

@@ -208,10 +208,13 @@ func (ltc *LocalTestCluster) Start(t testing.TB, initFactory InitFactoryFn) {
 		livenessInterval, heartbeatInterval := cfg.StoreLivenessDurations()
 		supportGracePeriod := cfg.RPCContext.StoreLivenessWithdrawalGracePeriod()
 		options := storeliveness.NewOptions(livenessInterval, heartbeatInterval, supportGracePeriod)
-		transport := storeliveness.NewTransport(
+		transport, err := storeliveness.NewTransport(
 			cfg.AmbientCtx, ltc.stopper, ltc.Clock,
-			nil /* dialer */, nil /* grpcServer */, nil, /* knobs */
+			nil /* dialer */, nil /* grpcServer */, nil /* drpcServer */, nil, /* knobs */
 		)
+		if err != nil {
+			t.Fatal(err)
+		}
 		knobs := cfg.TestingKnobs.StoreLivenessKnobs
 		cfg.StoreLiveness = storeliveness.NewNodeContainer(ltc.stopper, options, transport, knobs)
 	}