From 8d939dff9959c5a63e3c85ca1d5e8e3691761fca Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Wed, 5 Apr 2023 14:27:50 -0700 Subject: [PATCH] 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 | 1 + pkg/sql/flowinfra/flow_registry.go | 24 +++++----- pkg/sql/flowinfra/flow_registry_test.go | 58 ++++++++++++++++++++++--- 3 files changed, 67 insertions(+), 16 deletions(-) diff --git a/pkg/sql/flowinfra/BUILD.bazel b/pkg/sql/flowinfra/BUILD.bazel index 2ec70c8b51b3..b589393aa863 100644 --- a/pkg/sql/flowinfra/BUILD.bazel +++ b/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", ], ) diff --git a/pkg/sql/flowinfra/flow_registry.go b/pkg/sql/flowinfra/flow_registry.go index c5c91a445cde..8dd96e4a63b1 100644 --- a/pkg/sql/flowinfra/flow_registry.go +++ b/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) diff --git a/pkg/sql/flowinfra/flow_registry_test.go b/pkg/sql/flowinfra/flow_registry_test.go index 00d9e58bae31..5159e0beb8c5 100644 --- a/pkg/sql/flowinfra/flow_registry_test.go +++ b/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,7 +399,7 @@ 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. @@ -404,7 +407,7 @@ func TestFlowRegistryDrain(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) + } }) }