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.
This commit is contained in:
Chris Lu authored and GitHub committed 2026-11-03 12:53:28 -08:00
1 parent 836c60b2c9
commit ab6aa9b645
2 files changed
+228 -22

No files matched your search

+62 -22
View File
@@ -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)
}
}
}