admin: let a buffered detection completion win over the stream closing (#12379)

The detection waiter selects on the request's completion channel and on the stream's done channel. A worker that finishes a scan and then drops its stream (a restart right after a scan is the usual shape) buffers the completion into the request before the receive loop sees the EOF, and the stream's cleanup then closes done. A waiter that reaches its select after both are ready can take the disconnect branch, and a completed detection, proposals included, is failed as "detection worker disconnected" and retried or recorded as an error. When the stream closes or the liveness tick finds the worker gone, the waiter now takes a completion that is already queued and returns it as the success it is; the detector lease still goes, since the worker that answered is unreachable. A ctx that ended keeps its meaning: the caller stopped wanting the answer. The wait loop moves into awaitDetectionOutcome so a test can hand it both channels already ready; before this change that test fails on about half its runs. A completion that lands after the waiter gave up stays unmatched, since the request's pending entry is gone and the next request has a new id.
2026-11-04 11:07:35 +01:00 · 2026-11-03 12:53:28 -08:00
1 parent 836c60b2c9
commit ab6aa9b645
@@ -587,11 +587,62 @@ func (r *Plugin) runDetectionWithReport(
 		return nil, err
 	}

+	return r.awaitDetectionOutcome(ctx, jobType, requestID, session, state)
+}
+
+// awaitDetectionOutcome waits for the request's completion, the end of the
+// stream it was sent on, a stale worker heartbeat, or ctx. A completion the
+// worker delivered before its stream ended is buffered in state.complete, and
+// the stream's closure can be observed in the same select: that completion
+// wins, so a scan that finished right before its worker went away is reported
+// as the success it was, not as a disconnect. A ctx that ended stands as is:
+// the caller stopped wanting the answer.
+func (r *Plugin) awaitDetectionOutcome(
+	ctx context.Context,
+	jobType string,
+	requestID string,
+	session *streamSession,
+	state *pendingDetectionState,
+) (*DetectionReport, error) {
+	workerID := session.workerID
 	emptyReport := func() *DetectionReport {
 		return &DetectionReport{
 			RequestID: requestID,
 			JobType:   jobType,
-			WorkerID:  detector.WorkerID,
+			WorkerID:  workerID,
+		}
+	}
+	finish := func(complete *plugin_pb.DetectionComplete, ok bool) (*DetectionReport, error) {
+		if !ok {
+			return emptyReport(), fmt.Errorf("detection request %s interrupted", requestID)
+		}
+		proposals := cloneJobProposals(state.proposals)
+		report := &DetectionReport{
+			RequestID: requestID,
+			JobType:   jobType,
+			WorkerID:  workerID,
+			Proposals: proposals,
+			Complete:  complete,
+		}
+		if complete == nil {
+			return report, fmt.Errorf("detection request %s returned no completion state", requestID)
+		}
+		if !complete.Success {
+			return report, fmt.Errorf("detection failed for %s: %s", jobType, complete.ErrorMessage)
+		}
+		return report, nil
+	}
+	// finishIfCompleted takes a completion that is already queued when the
+	// worker is judged lost. The lease still goes: the worker that answered
+	// is gone or unreachable, so the next pick must start fresh.
+	finishIfCompleted := func() (bool, *DetectionReport, error) {
+		select {
+		case complete, ok := <-state.complete:
+			r.clearDetectorLease(jobType, workerID)
+			report, err := finish(complete, ok)
+			return true, report, err
+		default:
+			return false, nil, nil
 		}
 	}

@@ -605,11 +656,11 @@ func (r *Plugin) runDetectionWithReport(
 	for {
 		select {
 		case <-ctx.Done():
-			r.sendCancel(detector.WorkerID, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION, ctx.Err())
+			r.sendCancel(workerID, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION, ctx.Err())
 			r.appendActivity(JobActivity{
 				JobType:    jobType,
 				RequestID:  requestID,
-				WorkerID:   detector.WorkerID,
+				WorkerID:   workerID,
 				Source:     "detector",
 				Stage:      "canceled",
 				Message:    "detection canceled",
@@ -617,32 +668,21 @@ func (r *Plugin) runDetectionWithReport(
 			})
 			return emptyReport(), ctx.Err()
 		case <-session.done:
+			if completed, report, err := finishIfCompleted(); completed {
+				return report, err
+			}
 			return emptyReport(), r.failDetectionWorkerLost(jobType, requestID, session, "", false)
 		case <-liveness.C:
-			if r.isWorkerAvailable(detector.WorkerID) {
+			if r.isWorkerAvailable(workerID) {
 				continue
 			}
+			if completed, report, err := finishIfCompleted(); completed {
+				return report, err
+			}
 			// The stream may still be open, so the worker is told to stop.
 			return emptyReport(), r.failDetectionWorkerLost(jobType, requestID, session, "heartbeat stale", true)
 		case complete, ok := <-state.complete:
-			if !ok {
-				return emptyReport(), fmt.Errorf("detection request %s interrupted", requestID)
-			}
-			proposals := cloneJobProposals(state.proposals)
-			report := &DetectionReport{
-				RequestID: requestID,
-				JobType:   jobType,
-				WorkerID:  detector.WorkerID,
-				Proposals: proposals,
-				Complete:  complete,
-			}
-			if complete == nil {
-				return report, fmt.Errorf("detection request %s returned no completion state", requestID)
-			}
-			if !complete.Success {
-				return report, fmt.Errorf("detection failed for %s: %s", jobType, complete.ErrorMessage)
-			}
-			return report, nil
+			return finish(complete, ok)
 		}
 	}
 }
@@ -404,3 +404,169 @@ func TestScheduledDetectionRetryAvoidsClosedDetectorStream(t *testing.T) {
 	require.Equal(t, "success", p.snapshotSchedulerRun(jobType).lastRunStatus)
 	require.Len(t, sessionA.outgoing, 0, "nothing may be sent on the closed stream")
 }
+
+// The stream delivered the completion and then closed before the waiter
+// reached its select, so both channels are ready at once.
+func TestAwaitDetectionOutcomePrefersBufferedCompletionOverClosedStream(t *testing.T) {
+	p, err := New(Options{})
+	require.NoError(t, err)
+	defer p.Shutdown()
+
+	const jobType = "vacuum"
+	session := newDetectorSession(p, "worker-a", jobType)
+	p.setDetectorLease(jobType, "worker-a")
+	state := &pendingDetectionState{
+		complete: make(chan *plugin_pb.DetectionComplete, 1),
+		jobType:  jobType,
+		workerID: "worker-a",
+	}
+	state.proposals = []*plugin_pb.JobProposal{{ProposalId: "p-1", JobType: jobType}}
+	state.complete <- &plugin_pb.DetectionComplete{RequestId: "detect-1", JobType: jobType, Success: true, TotalProposals: 1}
+	p.WorkerDisconnectForTest("worker-a")
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
+	defer cancel()
+	report, err := p.awaitDetectionOutcome(ctx, jobType, "detect-1", session, state)
+	require.NoError(t, err)
+	require.NotNil(t, report)
+	require.Equal(t, "detect-1", report.RequestID)
+	require.Equal(t, "worker-a", report.WorkerID)
+	require.True(t, report.Complete.GetSuccess())
+	require.Len(t, report.Proposals, 1)
+	require.Equal(t, "p-1", report.Proposals[0].ProposalId)
+	require.Nil(t, findActivity(p.ListActivities(jobType, 50), "detector", "disconnected"), "a completed detection must not be reported as disconnected")
+
+	// The next pick must not be pinned to the gone worker.
+	newDetectorSession(p, "worker-b", jobType)
+	picked, err := p.pickDetector(jobType, nil)
+	require.NoError(t, err)
+	require.Equal(t, "worker-b", picked.WorkerID)
+	require.Equal(t, "worker-b", p.getDetectorLease(jobType))
+}
+
+func TestRunDetectionKeepsCompletionDeliveredBeforeDisconnect(t *testing.T) {
+	p, err := New(Options{})
+	require.NoError(t, err)
+	defer p.Shutdown()
+
+	const jobType = "vacuum"
+	session := newDetectorSession(p, "worker-a", jobType)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
+	defer cancel()
+	type result struct {
+		report *DetectionReport
+		err    error
+	}
+	resultCh := make(chan result, 1)
+	go func() {
+		report, runErr := p.RunDetectionWithReport(ctx, jobType, &plugin_pb.ClusterContext{}, 10)
+		resultCh <- result{report: report, err: runErr}
+	}()
+	request := awaitDetectionRequest(t, session)
+
+	// The worker's stream carries proposals and the completion, then ends:
+	// the worker finished its scan and restarted.
+	p.handleDetectionProposals("worker-a", &plugin_pb.DetectionProposals{
+		RequestId: request.RequestId,
+		JobType:   jobType,
+		Proposals: []*plugin_pb.JobProposal{{ProposalId: "p-1", JobType: jobType}},
+	})
+	p.handleDetectionComplete("worker-a", &plugin_pb.DetectionComplete{RequestId: request.RequestId, JobType: jobType, Success: true, TotalProposals: 1})
+	p.WorkerDisconnectForTest("worker-a")
+
+	var got result
+	select {
+	case got = <-resultCh:
+	case <-time.After(5 * time.Second):
+		t.Fatal("detection did not return")
+	}
+	require.NoError(t, got.err)
+	require.NotNil(t, got.report)
+	require.Equal(t, request.RequestId, got.report.RequestID)
+	require.True(t, got.report.Complete.GetSuccess())
+	require.Len(t, got.report.Proposals, 1)
+	require.Equal(t, "p-1", got.report.Proposals[0].ProposalId)
+	require.Equal(t, 0, p.pendingDetectionCountForTest())
+	require.Nil(t, findActivity(p.ListActivities(jobType, 50), "detector", "disconnected"))
+
+	// The next detection goes to a live worker, not the gone one.
+	next := newDetectorSession(p, "worker-b", jobType)
+	go func() {
+		_, _ = p.RunDetectionWithReport(ctx, jobType, &plugin_pb.ClusterContext{}, 10)
+	}()
+	awaitDetectionRequest(t, next)
+	require.Equal(t, "worker-b", p.getDetectorLease(jobType))
+}
+
+// A completion that lands after the waiter already gave up on the request
+// must not satisfy the next detection of the same job type.
+func TestLateDetectionCompletionDoesNotReachNextDetection(t *testing.T) {
+	p, err := New(Options{})
+	require.NoError(t, err)
+	defer p.Shutdown()
+
+	const jobType = "vacuum"
+	session := newDetectorSession(p, "worker-a", jobType)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
+	defer cancel()
+	errCh := make(chan error, 1)
+	go func() {
+		_, runErr := p.RunDetectionWithReport(ctx, jobType, &plugin_pb.ClusterContext{}, 10)
+		errCh <- runErr
+	}()
+	lost := awaitDetectionRequest(t, session)
+	p.WorkerDisconnectForTest("worker-a")
+	select {
+	case runErr := <-errCh:
+		var disconnected *detectionWorkerDisconnectedError
+		require.True(t, errors.As(runErr, &disconnected), "unexpected error: %v", runErr)
+	case <-time.After(5 * time.Second):
+		t.Fatal("detection did not fail after its worker disconnected")
+	}
+	require.Equal(t, 0, p.pendingDetectionCountForTest())
+
+	next := newDetectorSession(p, "worker-b", jobType)
+	type result struct {
+		report *DetectionReport
+		err    error
+	}
+	resultCh := make(chan result, 1)
+	go func() {
+		report, runErr := p.RunDetectionWithReport(ctx, jobType, &plugin_pb.ClusterContext{}, 10)
+		resultCh <- result{report: report, err: runErr}
+	}()
+	request := awaitDetectionRequest(t, next)
+	require.NotEqual(t, lost.RequestId, request.RequestId)
+
+	// The old worker's answer arrives late, under the old request ID.
+	p.handleDetectionProposals("worker-a", &plugin_pb.DetectionProposals{
+		RequestId: lost.RequestId,
+		JobType:   jobType,
+		Proposals: []*plugin_pb.JobProposal{{ProposalId: "stale", JobType: jobType}},
+	})
+	p.handleDetectionComplete("worker-a", &plugin_pb.DetectionComplete{RequestId: lost.RequestId, JobType: jobType, Success: true, TotalProposals: 1})
+	select {
+	case got := <-resultCh:
+		t.Fatalf("the late completion ended the next detection: %+v %v", got.report, got.err)
+	case <-time.After(300 * time.Millisecond):
+	}
+	require.Equal(t, 1, p.pendingDetectionCountForTest())
+
+	p.handleDetectionProposals("worker-b", &plugin_pb.DetectionProposals{
+		RequestId: request.RequestId,
+		JobType:   jobType,
+		Proposals: []*plugin_pb.JobProposal{{ProposalId: "fresh", JobType: jobType}},
+	})
+	p.handleDetectionComplete("worker-b", &plugin_pb.DetectionComplete{RequestId: request.RequestId, JobType: jobType, Success: true, TotalProposals: 1})
+	select {
+	case got := <-resultCh:
+		require.NoError(t, got.err)
+		require.Equal(t, request.RequestId, got.report.RequestID)
+		require.Len(t, got.report.Proposals, 1)
+		require.Equal(t, "fresh", got.report.Proposals[0].ProposalId)
+	case <-time.After(5 * time.Second):
+		t.Fatal("detection did not complete")
+	}
+}