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

This commit fixes a long-standing bug around the flow registry that
could make the drain loop be stuck forever. Drain process works by
draining several components in a loop until each component reports that
there was no more remaining work when the drain iteration was initiated.
One of the components to be drained is the flow registry: namely, we
want to make sure that there are no more remote flows present on the
node. We track that by having a map from the `FlowID` to the `flowEntry`
object.

Previously, it was possible for a `flowEntry` to become "stale" and
remain in the map forever. In particular, this was the case when
- `ConnectInboundStream` was called before the flow was scheduled
- the gRPC "handshake" failed in `ConnectInboundStream` (most likely due
to a network fluke)
- the flow never arrived (perhaps it was canceled before
`Flow.StartInternal` is called), or it arrived too late when the
registry was marked as "draining".

With such a scenario we would create a `flowEntry` with ref count of
zero and add it to the map in `ConnectInboundStream`, but no one would
ever remove it. This commit fixes this oversight by adjusting the ref
counting logic a bit so that we always hold a reference throughout (and
only until the end of) `ConnectInboundStream`.

Release note (bug fix): A rare bug with distributed plans shutdown has
been fixed that previously could make the graceful drain of cockroach
nodes be retrying forever. The bug has been present since before 22.1.
The drain process is affected by this bug if you see messages in the
logs like `drain details: distSQL execution flows:` with non-zero number
of flows that isn't going down over long period of time.

pkg/sql/flowinfra/BUILD.bazel

@@ -120,6 +120,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
 }
 
@@ -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
@@ -367,6 +369,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()
@@ -396,15 +399,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
@@ -434,7 +437,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
@@ -464,7 +467,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
@@ -504,7 +507,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)
+		}
 	})
 }