[4.2] fix: reuse remote backend after early stop

pkg/sql/compile/remoterunClient_test.go

@@ -522,6 +522,141 @@ func TestRemoteRun(t *testing.T) {
 	assert.Error(t, err)
 }
 
+func TestIssue27757RemoteRunCancellationClassification(t *testing.T) {
+	oldRuntime := runtime.ServiceRuntime("")
+	testRuntime := runtime.DefaultRuntime()
+	runtime.SetupServiceBasedRuntime("", testRuntime)
+	t.Cleanup(func() {
+		runtime.SetupServiceBasedRuntime("", oldRuntime)
+	})
+	catalog.SetupDefines("")
+
+	substantiveErr := moerr.NewInternalErrorNoCtx("remote execution failed")
+	tests := []struct {
+		name        string
+		cancelQuery bool
+		cancelCause error
+		wantErr     error
+		wantEvent   process.PipelineEventType
+		poison      bool
+	}{
+		{
+			name:      "internal early stop is successful and reuses backend",
+			wantEvent: process.EventEnd,
+		},
+		{
+			name:        "query cancellation is terminal and poisons backend",
+			cancelQuery: true,
+			wantErr:     context.Canceled,
+			wantEvent:   process.EventError,
+			poison:      true,
+		},
+		{
+			name:        "substantive error is terminal and poisons backend",
+			cancelCause: substantiveErr,
+			wantErr:     substantiveErr,
+			wantEvent:   process.EventError,
+			poison:      true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			proc := testutil.NewProcess(t)
+			queryCtx := proc.Base.GetContextBase().BuildQueryCtx(proc.GetTopContext())
+			_, cancelQuery := process.GetQueryCtxFromProc(proc)
+			t.Cleanup(cancelQuery)
+			proc.BuildPipelineContext(queryCtx)
+			txnCli, txnOp := newTestTxnClientAndOp(ctrl)
+			proc.Base.TxnClient = txnCli
+			proc.Base.TxnOperator = txnOp
+
+			responses := make(chan morpc.Message, 2)
+			stream := mock_morpc.NewMockStream(ctrl)
+			stream.EXPECT().Receive().Return(responses, nil)
+			stream.EXPECT().ID().Return(uint64(27757)).AnyTimes()
+			stream.EXPECT().Send(gomock.Any(), gomock.Any()).DoAndReturn(
+				func(_ context.Context, request morpc.Message) error {
+					message := request.(*pipeline.Message)
+					switch message.GetCmd() {
+					case pipeline.Method_PipelineMessage:
+						if tt.cancelQuery {
+							cancelQuery()
+						} else {
+							proc.Cancel(tt.cancelCause)
+						}
+					case pipeline.Method_StopSending:
+						responses <- &pipeline.Message{
+							Id:                   27757,
+							Cmd:                  pipeline.Method_PipelineMessage,
+							Sid:                  pipeline.Status_MessageEnd,
+							AcceptedTeardownMode: pipeline.StreamTeardownMode_FinishAck,
+						}
+					case pipeline.Method_PipelineStreamFinish:
+						responses <- &pipeline.Message{
+							Id:                   27757,
+							Cmd:                  pipeline.Method_PipelineStreamFinishAck,
+							Sid:                  pipeline.Status_MessageEnd,
+							AcceptedTeardownMode: pipeline.StreamTeardownMode_FinishAck,
+						}
+					}
+					return nil
+				}).AnyTimes()
+			stream.EXPECT().Close(tt.poison).Return(nil)
+			testRuntime.SetGlobalVariables(runtime.PipelineClient, &testPipelineClient{
+				genStream: func(context.Context, string) (morpc.Stream, error) {
+					return stream, nil
+				},
+			})
+
+			c := NewCompile(
+				"local-cn:6002",
+				"test",
+				"select id from ivf_entries order by distance limit 10",
+				"",
+				"",
+				newStubEngine(),
+				proc,
+				nil,
+				false,
+				nil,
+				time.Now(),
+			)
+			c.anal = &AnalyzeModule{qry: &plan.Query{}}
+
+			reg := process.NewPipelineEdge(1, 0)
+			root := connector.NewArgument().WithReg(reg)
+			t.Cleanup(root.Release)
+			s := &Scope{
+				Magic:         Remote,
+				Proc:          proc,
+				RootOp:        root,
+				ScopeAnalyzer: &ScopeAnalyzer{},
+				NodeInfo:      engine.Node{Addr: "remote-cn:6002", Mcpu: 1},
+			}
+
+			runErr := s.RemoteRun(c)
+			if tt.wantErr == nil {
+				require.NoError(t, runErr)
+			} else {
+				require.ErrorIs(t, runErr, tt.wantErr)
+			}
+			select {
+			case signal := <-reg.Ch2:
+				_, terminalErr := signal.Action()
+				require.Equal(t, tt.wantEvent, signal.EventType)
+				if tt.wantErr == nil {
+					require.NoError(t, terminalErr)
+				} else {
+					require.ErrorIs(t, terminalErr, tt.wantErr)
+				}
+			case <-time.After(time.Second):
+				t.Fatal("remote cleanup did not terminate its receiver")
+			}
+		})
+	}
+}
+
 func TestRemoteRunFailureReleasesPendingRetainedDispatchAttach(t *testing.T) {
 	oldRuntime := runtime.ServiceRuntime("")
 	testRuntime := runtime.DefaultRuntime()

pkg/sql/compile/scope.go

@@ -629,21 +629,21 @@ func (s *Scope) RemoteRun(c *Compile) error {
 	sender, err := s.remoteRun(c)
 
 	runErr := err
-	runErr = suppressRemoteRunCancelError(s.Proc.Ctx, runErr)
-	if err != nil && s.Proc.Cancel != nil {
-		cancelErr := runErr
-		if cancelErr == nil {
-			cancelErr = err
-		}
-		s.Proc.Cancel(cancelErr)
+	runErr = suppressRemoteRunCancelError(
+		s.Proc.Ctx,
+		scopeRunQueryContext(s.Proc),
+		runErr,
+	)
+	if runErr != nil && s.Proc.Cancel != nil {
+		s.Proc.Cancel(runErr)
 	}
 	// this clean-up action shouldn't be called before context check.
 	// because the clean-up action will cancel the context, and error will be suppressed.
-	p.CleanRootOperator(s.Proc, err != nil, c.isPrepare, runErr)
+	p.CleanRootOperator(s.Proc, runErr != nil, c.isPrepare, runErr)
 
 	// sender should be closed after cleanup (tell the children-pipeline that query was done).
 	if sender != nil {
-		if err == nil {
+		if runErr == nil {
 			sender.prepareForLocalCleanup()
 		}
 		sender.close()
@@ -1308,14 +1308,121 @@ func logRemoteNotifyCleanupSendFailure(
 		err)
 }
 
-func suppressRemoteRunCancelError(procCtx context.Context, err error) error {
-	if err == nil {
+func scopeRunQueryContext(proc *process.Process) context.Context {
+	if proc == nil || proc.Base == nil {
 		return nil
 	}
-	if procCtx != nil && procCtx.Err() != nil &&
-		(moerr.IsMoErrCode(err, moerr.ErrQueryInterrupted) || errors.Is(err, context.Canceled)) {
-		return nil
+	queryCtx, _ := process.GetQueryCtxFromProc(proc)
+	if queryCtx != nil {
+		return queryCtx
+	}
+	return proc.GetTopContext()
+}
+
+func isScopeCancellationError(err error) bool {
+	if err == nil {
+		return false
+	}
+	// A joined result is cancellation fallout only when every leaf is
+	// cancellation-shaped. One cancellation sibling must not hide a
+	// substantive execution failure.
+	if joined, ok := err.(interface{ Unwrap() []error }); ok {
+		children := joined.Unwrap()
+		if len(children) == 0 {
+			return false
+		}
+		for _, child := range children {
+			if !isScopeCancellationError(child) {
+				return false
+			}
+		}
+		return true
+	}
+	if wrapped, ok := err.(interface{ Unwrap() error }); ok {
+		if child := wrapped.Unwrap(); child != nil {
+			return isScopeCancellationError(child)
+		}
+	}
+	return errors.Is(err, context.Canceled) ||
+		errors.Is(err, context.DeadlineExceeded) ||
+		moerr.IsMoErrCode(err, moerr.ErrQueryInterrupted)
+}
+
+func isScopeCancellationFrom(err error, contextErr error) bool {
+	if err == nil || contextErr == nil {
+		return false
+	}
+	if joined, ok := err.(interface{ Unwrap() []error }); ok {
+		children := joined.Unwrap()
+		if len(children) == 0 {
+			return false
+		}
+		for _, child := range children {
+			if !isScopeCancellationFrom(child, contextErr) {
+				return false
+			}
+		}
+		return true
 	}
+	if wrapped, ok := err.(interface{ Unwrap() error }); ok {
+		if child := wrapped.Unwrap(); child != nil {
+			return isScopeCancellationFrom(child, contextErr)
+		}
+	}
+	return errors.Is(err, contextErr) ||
+		moerr.IsMoErrCode(err, moerr.ErrQueryInterrupted)
+}
+
+// normalizeScopeRunError distinguishes a substantive execution failure from
+// cancellation fallout. An internally canceled pipeline may finish
+// successfully, while query cancellation and substantive cancel causes remain
+// terminal and must poison the remote stream.
+func normalizeScopeRunError(
+	err error,
+	pipelineCtx context.Context,
+	queryCtx context.Context,
+) (error, bool) {
+	if err == nil || !isScopeCancellationError(err) ||
+		pipelineCtx == nil || pipelineCtx.Err() == nil {
+		return err, false
+	}
+	if queryCtx != nil {
+		if queryErr := queryCtx.Err(); queryErr != nil {
+			// WithTimeoutCause keeps DeadlineExceeded in Err and stores only
+			// diagnostics in Cause. Preserve the public timeout classification.
+			if errors.Is(queryErr, context.DeadlineExceeded) {
+				return queryErr, true
+			}
+			if !isScopeCancellationFrom(err, queryErr) {
+				return err, false
+			}
+			if cause := context.Cause(queryCtx); cause != nil {
+				return cause, true
+			}
+			return queryErr, true
+		}
+	}
+
+	// Cancellation is secondary only when every error leaf came from this
+	// pipeline. Preserve an independent error that merely raced cancellation.
+	if !isScopeCancellationFrom(err, pipelineCtx.Err()) {
+		return err, false
+	}
+	if cause := context.Cause(pipelineCtx); cause != nil {
+		err = cause
+	}
+	if isScopeCancellationError(err) && queryCtx != nil && queryCtx.Err() == nil {
+		return nil, true
+	}
+	return err, true
+}
+
+func suppressRemoteRunCancelError(
+	procCtx context.Context,
+	queryCtx context.Context,
+	err error,
+) error {
+	err, _ = normalizeScopeRunError(err, procCtx, queryCtx)
 	return err
 }