distsql: register DistSQL service with DRPC server

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

This change is part of a series and is similar to: #146926

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/sql/colflow/vectorized_flow.go

@@ -1274,7 +1274,7 @@ var _ flowinfra.InboundStreamHandler = vectorizedInboundStreamHandler{}
 // Run is part of the flowinfra.InboundStreamHandler interface.
 func (s vectorizedInboundStreamHandler) Run(
 	ctx context.Context,
-	stream execinfrapb.DistSQL_FlowStreamServer,
+	stream execinfrapb.RPCDistSQL_FlowStreamStream,
 	_ *execinfrapb.ProducerMessage,
 	_ *flowinfra.FlowBase,
 ) error {

pkg/sql/distsql/server.go

@@ -97,6 +97,12 @@ func NewServer(
 	return ds
 }
 
+type drpcServerImpl ServerImpl
+
+func (ds *ServerImpl) AsDRPCServer() execinfrapb.DRPCDistSQLServer {
+	return (*drpcServerImpl)(ds)
+}
+
 // Start launches workers for the server.
 //
 // Note that the initialization of the server required for performing the
@@ -614,6 +620,13 @@ func (ds *ServerImpl) setupSpanForIncomingRPC(
 		tracing.WithServerSpanKind)
 }
 
+// SetupFlow is part of the execinfrapb.DRPCDistSQLServer interface.
+func (ds *drpcServerImpl) SetupFlow(
+	ctx context.Context, req *execinfrapb.SetupFlowRequest,
+) (*execinfrapb.SimpleResponse, error) {
+	return (*ServerImpl)(ds).SetupFlow(ctx, req)
+}
+
 // SetupFlow is part of the execinfrapb.DistSQLServer interface.
 func (ds *ServerImpl) SetupFlow(
 	ctx context.Context, req *execinfrapb.SetupFlowRequest,
@@ -676,6 +689,13 @@ func (ds *ServerImpl) SetupFlow(
 	return &execinfrapb.SimpleResponse{}, nil
 }
 
+// CancelDeadFlows is part of the execinfrapb.DRPCDistSQLServer interface.
+func (ds *drpcServerImpl) CancelDeadFlows(
+	ctx context.Context, req *execinfrapb.CancelDeadFlowsRequest,
+) (*execinfrapb.SimpleResponse, error) {
+	return (*ServerImpl)(ds).CancelDeadFlows(ctx, req)
+}
+
 // CancelDeadFlows is part of the execinfrapb.DistSQLServer interface.
 func (ds *ServerImpl) CancelDeadFlows(
 	ctx context.Context, req *execinfrapb.CancelDeadFlowsRequest,
@@ -686,7 +706,7 @@ func (ds *ServerImpl) CancelDeadFlows(
 }
 
 func (ds *ServerImpl) flowStreamInt(
-	ctx context.Context, stream execinfrapb.DistSQL_FlowStreamServer,
+	ctx context.Context, stream execinfrapb.RPCDistSQL_FlowStreamStream,
 ) error {
 	// Receive the first message.
 	msg, err := stream.Recv()
@@ -720,8 +740,17 @@ func (ds *ServerImpl) flowStreamInt(
 	return streamStrategy.Run(ctx, stream, msg, f)
 }
 
+// FlowStream is part of the execinfrapb.DRPCDistSQLServer interface.
+func (ds *drpcServerImpl) FlowStream(stream execinfrapb.DRPCDistSQL_FlowStreamStream) error {
+	return (*ServerImpl)(ds).flowStream(stream)
+}
+
 // FlowStream is part of the execinfrapb.DistSQLServer interface.
 func (ds *ServerImpl) FlowStream(stream execinfrapb.DistSQL_FlowStreamServer) error {
+	return ds.flowStream(stream)
+}
+
+func (ds *ServerImpl) flowStream(stream execinfrapb.RPCDistSQL_FlowStreamStream) error {
 	ctx := ds.AnnotateCtx(stream.Context())
 	err := ds.flowStreamInt(ctx, stream)
 	if err != nil && log.V(2) {

pkg/sql/flowinfra/flow_registry.go

@@ -541,7 +541,7 @@ func (fr *FlowRegistry) ConnectInboundStream(
 	ctx context.Context,
 	flowID execinfrapb.FlowID,
 	streamID execinfrapb.StreamID,
-	stream execinfrapb.DistSQL_FlowStreamServer,
+	stream execinfrapb.RPCDistSQL_FlowStreamStream,
 	timeout time.Duration,
 ) (_ *FlowBase, _ InboundStreamHandler, cleanup func(), retErr error) {
 	fr.Lock()

pkg/sql/flowinfra/flow_registry_test.go

@@ -520,7 +520,7 @@ func TestFlowRegistryDrain(t *testing.T) {
 		close(delayCh)
 		injectedErr := errors.New("injected error")
 		serverStream = &delayedErrorServerStream{
-			DistSQL_FlowStreamServer: serverStream,
+			RPCDistSQL_FlowStreamStream: serverStream,
 			// Make rpcCalledCh a buffered channel so that the RPC is not
 			// blocked.
 			rpcCalledCh: make(chan struct{}, 1),
@@ -705,7 +705,7 @@ func TestFlowCancelPartiallyBlocked(t *testing.T) {
 // allows to block (on delayCh) Send() calls which always result in the
 // provided error.
 type delayedErrorServerStream struct {
-	execinfrapb.DistSQL_FlowStreamServer
+	execinfrapb.RPCDistSQL_FlowStreamStream
 	// rpcCalledCh is sent on in the very beginning of every Send() call.
 	rpcCalledCh chan<- struct{}
 	delayCh     <-chan struct{}
@@ -761,10 +761,10 @@ func TestErrorOnSlowHandshake(t *testing.T) {
 
 	rpcCalledCh, delayCh := make(chan struct{}), make(chan struct{})
 	serverStream = &delayedErrorServerStream{
-		DistSQL_FlowStreamServer: serverStream,
-		rpcCalledCh:              rpcCalledCh,
-		delayCh:                  delayCh,
-		err:                      errors.New("dummy error"),
+		RPCDistSQL_FlowStreamStream: serverStream,
+		rpcCalledCh:                 rpcCalledCh,
+		delayCh:                     delayCh,
+		err:                         errors.New("dummy error"),
 	}
 
 	receiver := &distsqlutils.RowBuffer{}

pkg/sql/flowinfra/inbound.go

@@ -25,7 +25,7 @@ type InboundStreamHandler interface {
 	// Run is called once a FlowStream RPC is handled and a stream is obtained to
 	// make this stream accessible to the rest of the flow.
 	Run(
-		ctx context.Context, stream execinfrapb.DistSQL_FlowStreamServer, firstMsg *execinfrapb.ProducerMessage, f *FlowBase,
+		ctx context.Context, stream execinfrapb.RPCDistSQL_FlowStreamStream, firstMsg *execinfrapb.ProducerMessage, f *FlowBase,
 	) error
 	// Timeout is called with an error, which results in the teardown of the
 	// stream strategy with the given error.
@@ -45,7 +45,7 @@ var _ InboundStreamHandler = RowInboundStreamHandler{}
 // Run is part of the InboundStreamHandler interface.
 func (s RowInboundStreamHandler) Run(
 	ctx context.Context,
-	stream execinfrapb.DistSQL_FlowStreamServer,
+	stream execinfrapb.RPCDistSQL_FlowStreamStream,
 	firstMsg *execinfrapb.ProducerMessage,
 	f *FlowBase,
 ) error {
@@ -61,13 +61,13 @@ func (s RowInboundStreamHandler) Timeout(err error) {
 	s.ProducerDone()
 }
 
-// processInboundStream receives rows from a DistSQL_FlowStreamServer and sends
+// processInboundStream receives rows from a RPCDistSQL_FlowStreamStream and sends
 // them to a RowReceiver. Optionally processes an initial StreamMessage that was
 // already received (because the first message contains the flow and stream IDs,
 // it needs to be received before we can get here).
 func processInboundStream(
 	ctx context.Context,
-	stream execinfrapb.DistSQL_FlowStreamServer,
+	stream execinfrapb.RPCDistSQL_FlowStreamStream,
 	firstMsg *execinfrapb.ProducerMessage,
 	dst execinfra.RowReceiver,
 	f *FlowBase,
@@ -90,7 +90,7 @@ func processInboundStream(
 
 func processInboundStreamHelper(
 	ctx context.Context,
-	stream execinfrapb.DistSQL_FlowStreamServer,
+	stream execinfrapb.RPCDistSQL_FlowStreamStream,
 	firstMsg *execinfrapb.ProducerMessage,
 	dst execinfra.RowReceiver,
 	f *FlowBase,
@@ -173,7 +173,7 @@ func processInboundStreamHelper(
 // producer that it doesn't need any more rows and the producer should drain. A
 // signal is sent on stream to the producer to ask it to send metadata.
 func sendDrainSignalToStreamProducer(
-	ctx context.Context, stream execinfrapb.DistSQL_FlowStreamServer,
+	ctx context.Context, stream execinfrapb.RPCDistSQL_FlowStreamStream,
 ) error {
 	log.VEvent(ctx, 1, "sending drain signal to producer")
 	sig := execinfrapb.ConsumerSignal{DrainRequest: &execinfrapb.DrainRequest{}}
@@ -187,7 +187,7 @@ func sendDrainSignalToStreamProducer(
 func processProducerMessage(
 	ctx context.Context,
 	flowBase *FlowBase,
-	stream execinfrapb.DistSQL_FlowStreamServer,
+	stream execinfrapb.RPCDistSQL_FlowStreamStream,
 	dst execinfra.RowReceiver,
 	sd *StreamDecoder,
 	draining *bool,

pkg/sql/flowinfra/utils_test.go

@@ -32,7 +32,7 @@ import (
 func createDummyStream(
 	t *testing.T,
 ) (
-	serverStream execinfrapb.DistSQL_FlowStreamServer,
+	serverStream execinfrapb.RPCDistSQL_FlowStreamStream,
 	clientStream execinfrapb.RPCDistSQL_FlowStreamClient,
 	cleanup func(),
 ) {