release-23.1.0: flowinfra: fix a rare bug that could make drain be stuck forever

pkg/sql/flowinfra/BUILD.bazel

@@ -118,6 +118,7 @@ go_test(
         "//pkg/util/timeutil",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_redact//:redact",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/sql/flowinfra/flow_registry.go

@@ -401,7 +401,6 @@ func (fr *FlowRegistry) waitForFlow(
 		waitCh = make(chan struct{})
 		entry.waitCh = waitCh
 	}
-	entry.refCount++
 	fr.Unlock()
 
 	select {
@@ -411,7 +410,6 @@ func (fr *FlowRegistry) waitForFlow(
 	}
 
 	fr.Lock()
-	fr.releaseEntryLocked(id)
 	return entry.flow
 }
 
@@ -421,7 +419,7 @@ func (fr *FlowRegistry) waitForFlow(
 // expectedConnectionTime so that any flows that were registered at the end of
 // the time window have a reasonable amount of time to connect to their
 // consumers, thus unblocking them. All flows that are still running at this
-// point are canceled if cancelStillRunning is true.
+// point are canceled.
 //
 // The FlowRegistry rejects any new flows once it has finished draining.
 //
@@ -524,8 +522,8 @@ func (fr *FlowRegistry) Drain(
 	allFlowsDone <- struct{}{}
 }
 
-// Undrain causes the FlowRegistry to start accepting flows again.
-func (fr *FlowRegistry) Undrain() {
+// undrain causes the FlowRegistry to start accepting flows again.
+func (fr *FlowRegistry) undrain() {
 	fr.Lock()
 	fr.draining = false
 	fr.Unlock()
@@ -556,7 +554,18 @@ func (fr *FlowRegistry) ConnectInboundStream(
 	fr.Lock()
 	entry := fr.getEntryLocked(flowID)
 	flow := entry.flow
+	// Take a reference that is always released at the end of this method. In
+	// the happy case (when we end up returning a flow that we connected to),
+	// that flow also took reference in RegisterFlow, so the ref count won't go
+	// below one; in all other cases we want to make sure to delete the entry
+	// from the registry if we're holding the only reference.
+	entry.refCount++
 	fr.Unlock()
+	defer func() {
+		fr.Lock()
+		defer fr.Unlock()
+		fr.releaseEntryLocked(flowID)
+	}()
 	if flow == nil {
 		// Send the handshake message informing the producer that the consumer has
 		// not been scheduled yet. Another handshake will be sent below once the
@@ -570,13 +579,6 @@ func (fr *FlowRegistry) ConnectInboundStream(
 				MinAcceptedVersion:       execinfra.MinAcceptedVersion,
 			},
 		}); err != nil {
-			// TODO(andrei): We failed to send a message to the producer; we'll return
-			// an error and leave this stream with connected == false so it times out
-			// later. We could call finishInboundStreamLocked() now so that the flow
-			// doesn't wait for the timeout and we could remember the error for the
-			// consumer if the consumer comes later, but I'm not sure what the best
-			// way to do that is. Similarly for the 2nd handshake message below,
-			// except there we already have the consumer and we can push the error.
 			return nil, nil, nil, err
 		}
 		flow = fr.waitForFlow(ctx, flowID, timeout)

pkg/sql/flowinfra/flow_registry_test.go

@@ -26,9 +26,11 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/distsqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/cancelchecker"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/redact"
 )
 
 // lookupFlow returns the registered flow with the given ID. If no such flow is
@@ -366,6 +368,7 @@ func TestHandshake(t *testing.T) {
 // subtests for more details.
 func TestFlowRegistryDrain(t *testing.T) {
 	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
 
 	ctx := context.Background()
 	reg := NewFlowRegistry()
@@ -394,15 +397,15 @@ func TestFlowRegistryDrain(t *testing.T) {
 		time.Sleep(time.Microsecond)
 		reg.UnregisterFlow(id)
 		<-drainDone
-		reg.Undrain()
+		reg.undrain()
 	})
 
 	// DrainTimeout verifies that Drain returns once the timeout expires.
 	t.Run("DrainTimeout", func(t *testing.T) {
 		registerFlow(t, id)
 		reg.Drain(0 /* flowDrainWait */, 0 /* minFlowDrainWait */, nil /* reporter */)
 		reg.UnregisterFlow(id)
-		reg.Undrain()
+		reg.undrain()
 	})
 
 	// AcceptNewFlow verifies that a FlowRegistry continues accepting flows
@@ -432,7 +435,7 @@ func TestFlowRegistryDrain(t *testing.T) {
 		); !testutils.IsError(err, "draining") {
 			t.Fatalf("unexpected error: %v", err)
 		}
-		reg.Undrain()
+		reg.undrain()
 	})
 
 	// MinFlowWait verifies that the FlowRegistry waits a minimum amount of time
@@ -462,7 +465,7 @@ func TestFlowRegistryDrain(t *testing.T) {
 		}
 		reg.UnregisterFlow(id)
 		<-drainDone
-		reg.Undrain()
+		reg.undrain()
 
 		// Case in which a running flow finishes before the minimum wait time. We
 		// attempt to register another flow after the completion of the first flow
@@ -502,7 +505,52 @@ func TestFlowRegistryDrain(t *testing.T) {
 		reg.UnregisterFlow(id)
 		<-drainDone
 
-		reg.Undrain()
+		reg.undrain()
+	})
+
+	// StaleFlowEntry verifies that the Drain doesn't get stuck forever due to a
+	// stale flowEntry in the registry that was created by ConnectInboundStream
+	// which failed its initial handshake. If such entry remains in the
+	// registry, then the Drain process will be stuck forever (#100710).
+	t.Run("StaleFlowEntry", func(t *testing.T) {
+		serverStream, _ /* clientStream */, cleanup := createDummyStream(t)
+		defer cleanup()
+
+		// Create a gRPC stream where we inject an error on each Send call.
+		//
+		// We don't need to delay the RPC call.
+		delayCh := make(chan struct{})
+		close(delayCh)
+		injectedErr := errors.New("injected error")
+		serverStream = &delayedErrorServerStream{
+			DistSQL_FlowStreamServer: serverStream,
+			// Make rpcCalledCh a buffered channel so that the RPC is not
+			// blocked.
+			rpcCalledCh: make(chan struct{}, 1),
+			delayCh:     delayCh,
+			err:         injectedErr,
+		}
+
+		// Attempt to connect an inbound stream for which there is no flow in
+		// the registry. In such case we attempt to send a handshake message
+		// which should fail with the injected error.
+		_, _, _, err := reg.ConnectInboundStream(ctx, id, execinfrapb.StreamID(0), serverStream, 0 /* timeout */)
+		if !errors.Is(err, injectedErr) {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		// Now, the crux of the test: attempt to drain the registry and verify
+		// that there were no flows to drain. If we do see some, then it means
+		// that we have a stale flowEntry that will never get cleaned up.
+		var remaining int
+		reporter := func(r int, _ redact.SafeString) {
+			remaining += r
+		}
+		reg.Drain(0 /* flowDrainWait */, 0 /* minFlowDrainWait */, reporter)
+		reg.undrain()
+		if remaining != 0 {
+			t.Fatalf("expected zero flows reported, found %d", remaining)
+		}
 	})
 }
 

pkg/sql/pgwire/server.go

@@ -517,7 +517,7 @@ func (s *Server) waitConnsDone() (cancelChanMap, chan struct{}, chan struct{}) {
 	connCancelMap := func() cancelChanMap {
 		s.mu.Lock()
 		defer s.mu.Unlock()
-		connCancelMap := make(cancelChanMap)
+		connCancelMap := make(cancelChanMap, len(s.mu.connCancelMap))
 		for done, cancel := range s.mu.connCancelMap {
 			connCancelMap[done] = cancel
 		}