admin: fail an in-flight detection when its worker disconnects (#12376)

* admin: keep a reconnected plugin worker registered when its old stream unwinds

cleanupSession removed whatever session was registered under the worker
ID, so a worker that reconnected while its previous stream was still
unwinding lost the new session and its registry entry to the old
stream's deferred cleanup. Tie cleanup to the stream that is ending and
publish the new stream before the registry upsert so a lingering
predecessor sees itself superseded.

* admin: fail an in-flight detection when its worker disconnects

RunDetectionWithReport waited only on the caller's context and the
DetectionComplete channel, so a detection whose worker stream ended
mid-request could only finish by timeout. On a customer cluster a
planned restart of ec-vacuum-1 during an ec_vacuum detection left the
default lane in 'detecting' with every execution slot idle until the
3600s detection timeout, blocking every job type in that lane.

Bind the request to the stream it was sent on and fail the detection
as soon as that stream closes, whether by disconnect or by a reconnect
that replaced the session. A liveness tick also fails it once the
registry no longer considers the worker available (heartbeat stale),
sending a cancel in that case because the stream may still be open.
The detector lease is dropped so the next pick starts fresh, and a
'disconnected' detector activity is recorded. WorkerDisconnectForTest
now goes through the real session cleanup.

* admin: retry a scheduled detection once when its detector drops

When a detection fails because its worker disconnected, pick a detector
again and retry once within the same iteration while the job type's
runtime budget allows, instead of giving the whole cycle up. Any other
failure, or a failed retry, ends the run as before and the lane moves
on; the next cycle re-detects anyway.

* admin: register a plugin worker's hello and stream as one
 step

WorkerStream published the stream under sessionsMu but ran the registry
upsert after releasing it, so two concurrent streams for one worker ID
could publish their hellos in the opposite order from their sessions:
the registry then described the older hello's capabilities while only
the newer stream was sendable, and detector picks routed on them.

registerSession upserts the hello and installs the stream under one
sessionsMu critical section, and a hello or heartbeat arriving on a
stream that is no longer the registered one is dropped instead of
rewriting the live stream's entry.

* admin: steer a detection retry off the stream that dropped

The retry picked while the lost worker could still be in the registry:
cleanupSession closed done before taking the registry entry down, and a
stream closed by a reconnect keeps its entry on purpose. The lost worker
ranks first on free slots, so the single retry went back to it and
failed on send while a healthy worker sat idle.

The disconnect error now carries the stream it came from and the retry
picks around it; pickDetector only returns workers whose current stream
is open, excluding registry entries whose stream is closed; and
cleanupSession removes the registry entry before closing done, so a
waiter woken by the closure never sees the dead worker as a candidate.

weed/admin/plugin/plugin.go

@@ -27,6 +27,10 @@ const (
 	defaultReconnectDelay      = 5
 	defaultPendingSchemaBuffer = 1
 	adminScriptJobType         = "admin_script"
+	// How often an in-flight detection re-checks that its worker is still
+	// registered and heartbeating; a session that went stale fails the
+	// detection the same way a closed stream does.
+	defaultDetectionLivenessInterval = 30 * time.Second
 )
 
 type Options struct {
@@ -59,6 +63,8 @@ type Plugin struct {
 	outgoingBuffer int
 	sendTimeout    time.Duration
 
+	detectionLivenessInterval time.Duration
+
 	schedulerTick          time.Duration
 	clusterContextProvider func(context.Context) (*plugin_pb.ClusterContext, error)
 	configDefaultsProvider func(config *plugin_pb.PersistedJobTypeConfig) *plugin_pb.PersistedJobTypeConfig
@@ -140,6 +146,23 @@ type pendingDetectionState struct {
 	workerID  string
 }
 
+// detectionWorkerDisconnectedError reports a detection whose worker session
+// ended (stream closed, replaced by a reconnect, or heartbeat went stale)
+// before it answered.
+type detectionWorkerDisconnectedError struct {
+	WorkerID string
+	Reason   string
+	session  *streamSession
+}
+
+func (e *detectionWorkerDisconnectedError) Error() string {
+	msg := fmt.Sprintf("detection worker %s disconnected", e.WorkerID)
+	if e.Reason != "" {
+		msg += ": " + e.Reason
+	}
+	return msg
+}
+
 // DetectionReport captures one detection run including request metadata.
 type DetectionReport struct {
 	RequestID string
@@ -178,6 +201,7 @@ func New(options Options) (*Plugin, error) {
 		registry:                  NewRegistry(),
 		outgoingBuffer:            bufferSize,
 		sendTimeout:               sendTimeout,
+		detectionLivenessInterval: defaultDetectionLivenessInterval,
 		schedulerTick:             schedulerTick,
 		clusterContextProvider:    options.ClusterContextProvider,
 		configDefaultsProvider:    options.ConfigDefaultsProvider,
@@ -290,15 +314,13 @@ func (r *Plugin) WorkerStream(stream plugin_pb.PluginControlService_WorkerStream
 	}
 
 	workerID := hello.WorkerId
-	r.registry.UpsertFromHello(hello)
-
 	session := &streamSession{
 		workerID: workerID,
 		outgoing: make(chan *plugin_pb.AdminToWorkerMessage, r.outgoingBuffer),
 		done:     make(chan struct{}),
 	}
-	r.putSession(session)
-	defer r.cleanupSession(workerID)
+	r.registerSession(session, hello)
+	defer r.cleanupSession(session)
 
 	glog.V(0).Infof("Plugin worker connected: %s (%s)", workerID, hello.Address)
 
@@ -309,7 +331,7 @@ func (r *Plugin) WorkerStream(stream plugin_pb.PluginControlService_WorkerStream
 		sendErrCh <- r.sendLoop(stream.Context(), stream, session)
 	}()
 
-	if err := r.sendAdminHello(workerID); err != nil {
+	if err := r.sendAdminHello(session); err != nil {
 		glog.Warningf("failed to send plugin admin hello to %s: %v", workerID, err)
 	}
 	go r.prefetchDescriptorsFromHello(hello)
@@ -324,7 +346,7 @@ func (r *Plugin) WorkerStream(stream plugin_pb.PluginControlService_WorkerStream
 				recvErrCh <- recvErr
 				return
 			}
-			r.handleWorkerMessage(workerID, message)
+			r.handleWorkerMessage(session, message)
 		}
 	}()
 
@@ -462,11 +484,23 @@ func (r *Plugin) RunDetectionWithReport(
 	jobType string,
 	clusterContext *plugin_pb.ClusterContext,
 	maxResults int32,
+) (*DetectionReport, error) {
+	return r.runDetectionWithReport(ctx, jobType, clusterContext, maxResults, nil)
+}
+
+// runDetectionWithReport is RunDetectionWithReport with a stream the pick
+// must avoid: the one whose closure a retry is recovering from.
+func (r *Plugin) runDetectionWithReport(
+	ctx context.Context,
+	jobType string,
+	clusterContext *plugin_pb.ClusterContext,
+	maxResults int32,
+	avoid *streamSession,
 ) (*DetectionReport, error) {
 	releaseGate := r.acquireDetectionExecutionGate(jobType, false)
 	defer releaseGate()
 
-	detector, err := r.pickDetector(jobType)
+	detector, err := r.pickDetector(jobType, avoid)
 	if err != nil {
 		return nil, err
 	}
@@ -530,7 +564,16 @@ func (r *Plugin) RunDetectionWithReport(
 		},
 	}
 
-	if err := r.sendToWorker(detector.WorkerID, message); err != nil {
+	// The request is bound to the stream it was sent on: if that stream ends
+	// (disconnect or a reconnect with a new session) the answer can never
+	// arrive, so the wait below fails instead of running out the timeout.
+	session, connected := r.getSession(detector.WorkerID)
+	if !connected {
+		err = fmt.Errorf("worker %s is not connected", detector.WorkerID)
+	} else {
+		err = r.sendToSession(session, message)
+	}
+	if err != nil {
 		r.clearDetectorLease(jobType, detector.WorkerID)
 		r.appendActivity(JobActivity{
 			JobType:    jobType,
@@ -544,49 +587,89 @@ func (r *Plugin) RunDetectionWithReport(
 		return nil, err
 	}
 
-	select {
-	case <-ctx.Done():
-		r.sendCancel(detector.WorkerID, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION, ctx.Err())
-		r.appendActivity(JobActivity{
-			JobType:    jobType,
-			RequestID:  requestID,
-			WorkerID:   detector.WorkerID,
-			Source:     "detector",
-			Stage:      "canceled",
-			Message:    "detection canceled",
-			OccurredAt: timeToPtr(time.Now().UTC()),
-		})
+	emptyReport := func() *DetectionReport {
 		return &DetectionReport{
 			RequestID: requestID,
 			JobType:   jobType,
 			WorkerID:  detector.WorkerID,
-		}, ctx.Err()
-	case complete, ok := <-state.complete:
-		if !ok {
-			return &DetectionReport{
+		}
+	}
+
+	livenessInterval := r.detectionLivenessInterval
+	if livenessInterval <= 0 {
+		livenessInterval = defaultDetectionLivenessInterval
+	}
+	liveness := time.NewTicker(livenessInterval)
+	defer liveness.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			r.sendCancel(detector.WorkerID, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION, ctx.Err())
+			r.appendActivity(JobActivity{
+				JobType:    jobType,
+				RequestID:  requestID,
+				WorkerID:   detector.WorkerID,
+				Source:     "detector",
+				Stage:      "canceled",
+				Message:    "detection canceled",
+				OccurredAt: timeToPtr(time.Now().UTC()),
+			})
+			return emptyReport(), ctx.Err()
+		case <-session.done:
+			return emptyReport(), r.failDetectionWorkerLost(jobType, requestID, session, "", false)
+		case <-liveness.C:
+			if r.isWorkerAvailable(detector.WorkerID) {
+				continue
+			}
+			// 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,
-			}, fmt.Errorf("detection request %s interrupted", requestID)
+				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
 		}
-		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
 	}
 }
 
+// failDetectionWorkerLost closes out a detection whose worker session ended
+// before it answered. The detector lease is dropped so the next pick starts
+// fresh; the worker is only told to cancel when its stream may still be open.
+func (r *Plugin) failDetectionWorkerLost(jobType, requestID string, session *streamSession, reason string, cancelOnWorker bool) error {
+	workerID := session.workerID
+	err := &detectionWorkerDisconnectedError{WorkerID: workerID, Reason: reason, session: session}
+	r.clearDetectorLease(jobType, workerID)
+	if cancelOnWorker {
+		r.sendCancel(workerID, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION, err)
+	}
+	glog.Warningf("Plugin detection %s for %s failed: %v", requestID, jobType, err)
+	r.appendActivity(JobActivity{
+		JobType:    jobType,
+		RequestID:  requestID,
+		WorkerID:   workerID,
+		Source:     "detector",
+		Stage:      "disconnected",
+		Message:    err.Error(),
+		OccurredAt: timeToPtr(time.Now().UTC()),
+	})
+	return err
+}
+
 // RunDetection requests one detector worker to produce job proposals for a job type.
 func (r *Plugin) RunDetection(
 	ctx context.Context,
@@ -834,17 +917,20 @@ func (r *Plugin) FilterProposalsWithActiveJobs(jobType string, proposals []*plug
 }
 
 func (r *Plugin) PickDetectorWorker(jobType string) (*WorkerSession, error) {
-	return r.pickDetector(jobType)
+	return r.pickDetector(jobType, nil)
 }
 
 func (r *Plugin) PickExecutorWorker(jobType string) (*WorkerSession, error) {
 	return r.registry.PickExecutor(jobType)
 }
 
-func (r *Plugin) pickDetector(jobType string) (*WorkerSession, error) {
+// pickDetector picks a detector that can actually be sent to: registered,
+// with an open stream, and not the stream in avoid. The registry alone is
+// not enough because a stream can close before its entry is removed.
+func (r *Plugin) pickDetector(jobType string, avoid *streamSession) (*WorkerSession, error) {
 	leasedWorkerID := r.getDetectorLease(jobType)
 	if leasedWorkerID != "" {
-		if worker, ok := r.registry.Get(leasedWorkerID); ok {
+		if worker, ok := r.registry.Get(leasedWorkerID); ok && r.sessionIsOpen(leasedWorkerID, avoid) {
 			if capability := worker.Capabilities[jobType]; capability != nil && capability.CanDetect {
 				return worker, nil
 			}
@@ -852,13 +938,31 @@ func (r *Plugin) pickDetector(jobType string) (*WorkerSession, error) {
 		r.clearDetectorLease(jobType, leasedWorkerID)
 	}
 
-	detector, err := r.registry.PickDetector(jobType)
-	if err != nil {
-		return nil, err
+	var excluded []string
+	for {
+		detector, err := r.registry.PickDetector(jobType, excluded...)
+		if err != nil {
+			return nil, err
+		}
+		if r.sessionIsOpen(detector.WorkerID, avoid) {
+			r.setDetectorLease(jobType, detector.WorkerID)
+			return detector, nil
+		}
+		excluded = append(excluded, detector.WorkerID)
 	}
+}
 
-	r.setDetectorLease(jobType, detector.WorkerID)
-	return detector, nil
+func (r *Plugin) sessionIsOpen(workerID string, avoid *streamSession) bool {
+	session, ok := r.getSession(workerID)
+	if !ok || session == avoid {
+		return false
+	}
+	select {
+	case <-session.done:
+		return false
+	default:
+		return true
+	}
 }
 
 func (r *Plugin) getDetectorLease(jobType string) string {
@@ -920,7 +1024,7 @@ func (r *Plugin) sendCancel(workerID, targetID string, kind plugin_pb.WorkKind,
 	}
 }
 
-func (r *Plugin) sendAdminHello(workerID string) error {
+func (r *Plugin) sendAdminHello(session *streamSession) error {
 	msg := &plugin_pb.AdminToWorkerMessage{
 		RequestId: "",
 		SentAt:    timestamppb.Now(),
@@ -933,7 +1037,7 @@ func (r *Plugin) sendAdminHello(workerID string) error {
 			},
 		},
 	}
-	return r.sendToWorker(workerID, msg)
+	return r.sendToSession(session, msg)
 }
 
 func (r *Plugin) sendLoop(
@@ -960,37 +1064,50 @@ func (r *Plugin) sendLoop(
 	}
 }
 
-func (r *Plugin) sendToWorker(workerID string, message *plugin_pb.AdminToWorkerMessage) error {
+func (r *Plugin) getSession(workerID string) (*streamSession, bool) {
 	r.sessionsMu.RLock()
 	session, ok := r.sessions[workerID]
 	r.sessionsMu.RUnlock()
+	return session, ok
+}
+
+func (r *Plugin) sendToWorker(workerID string, message *plugin_pb.AdminToWorkerMessage) error {
+	session, ok := r.getSession(workerID)
 	if !ok {
 		return fmt.Errorf("worker %s is not connected", workerID)
 	}
+	return r.sendToSession(session, message)
+}
 
+func (r *Plugin) sendToSession(session *streamSession, message *plugin_pb.AdminToWorkerMessage) error {
 	select {
 	case <-r.shutdownCh:
 		return fmt.Errorf("plugin is shutting down")
 	case <-session.done:
-		return fmt.Errorf("worker %s session is closed", workerID)
+		return fmt.Errorf("worker %s session is closed", session.workerID)
 	case session.outgoing <- message:
 		return nil
 	case <-time.After(r.sendTimeout):
-		return fmt.Errorf("timed out sending message to worker %s", workerID)
+		return fmt.Errorf("timed out sending message to worker %s", session.workerID)
 	}
 }
 
-func (r *Plugin) handleWorkerMessage(workerID string, message *plugin_pb.WorkerToAdminMessage) {
+func (r *Plugin) handleWorkerMessage(session *streamSession, message *plugin_pb.WorkerToAdminMessage) {
 	if message == nil {
 		return
 	}
+	workerID := session.workerID
 
 	switch body := message.Body.(type) {
 	case *plugin_pb.WorkerToAdminMessage_Hello:
-		r.registry.UpsertFromHello(body.Hello)
+		if !r.updateRegistryFromSession(session, func() { r.registry.UpsertFromHello(body.Hello) }) {
+			return
+		}
 		r.wakeScheduler()
 	case *plugin_pb.WorkerToAdminMessage_Heartbeat:
-		r.registry.UpdateHeartbeat(workerID, body.Heartbeat)
+		if !r.updateRegistryFromSession(session, func() { r.registry.UpdateHeartbeat(workerID, body.Heartbeat) }) {
+			return
+		}
 	case *plugin_pb.WorkerToAdminMessage_ConfigSchemaResponse:
 		r.handleConfigSchemaResponse(body.ConfigSchemaResponse)
 	case *plugin_pb.WorkerToAdminMessage_DetectionProposals:
@@ -1335,27 +1452,73 @@ func (r *Plugin) safeSendJobCompleted(requestID string, completed *plugin_pb.Job
 	safeSendCh(ch, completed, r.shutdownCh)
 }
 
+// The registry entry for a worker always describes the stream that is
+// currently sendable for it: every registry write for a worker happens under
+// sessionsMu, together with or checked against r.sessions[workerID].
+
 func (r *Plugin) putSession(session *streamSession) {
 	r.sessionsMu.Lock()
 	defer r.sessionsMu.Unlock()
-
-	if old, exists := r.sessions[session.workerID]; exists {
-		old.close()
-	}
-	r.sessions[session.workerID] = session
+	r.publishSessionLocked(session)
 }
 
-func (r *Plugin) cleanupSession(workerID string) {
-	r.registry.Remove(workerID)
+// registerSession publishes a stream and its hello as one step so a second
+// stream for the same worker cannot leave the registry describing the older
+// hello while only the newer stream is sendable.
+func (r *Plugin) registerSession(session *streamSession, hello *plugin_pb.WorkerHello) {
+	r.sessionsMu.Lock()
+	defer r.sessionsMu.Unlock()
+	r.registry.UpsertFromHello(hello)
+	r.publishSessionLocked(session)
+}
+
+// publishSessionLocked installs the stream and closes the one it replaces.
+// Anything waiting on the old stream's done wakes to a registry and session
+// map that already describe the replacement.
+func (r *Plugin) publishSessionLocked(session *streamSession) {
+	old, exists := r.sessions[session.workerID]
+	r.sessions[session.workerID] = session
+	if exists && old != session {
+		old.close()
+	}
+}
+
+// updateRegistryFromSession applies a registry write for a stream's worker
+// only while that stream is the registered one; a superseded stream's late
+// hello or heartbeat must not describe the live stream.
+func (r *Plugin) updateRegistryFromSession(session *streamSession, update func()) bool {
+	r.sessionsMu.RLock()
+	defer r.sessionsMu.RUnlock()
+	if current, ok := r.sessions[session.workerID]; !ok || current != session {
+		return false
+	}
+	update()
+	return true
+}
+
+// cleanupSession ends one stream. Only the stream currently registered for
+// the worker takes the registry entry down with it; a stream that a reconnect
+// already replaced must not evict its successor. The registry entry goes
+// before done closes so nothing woken by the closure can still pick it.
+func (r *Plugin) cleanupSession(session *streamSession) {
+	workerID := session.workerID
 
 	r.sessionsMu.Lock()
-	session, exists := r.sessions[workerID]
-	if exists {
+	current, exists := r.sessions[workerID]
+	superseded := exists && current != session
+	if exists && !superseded {
 		delete(r.sessions, workerID)
-		session.close()
 	}
+	if !superseded {
+		r.registry.Remove(workerID)
+	}
+	session.close()
 	r.sessionsMu.Unlock()
 
+	if superseded {
+		glog.V(1).Infof("Plugin worker stream closed: %s (replaced by a newer session)", workerID)
+		return
+	}
 	glog.V(0).Infof("Plugin worker disconnected: %s", workerID)
 }
 
@@ -1494,5 +1657,9 @@ func (r *Plugin) WorkerConnectForTest(hello *plugin_pb.WorkerHello) {
 
 // WorkerDisconnectForTest simulates a worker disconnecting (test helper).
 func (r *Plugin) WorkerDisconnectForTest(workerID string) {
+	if session, ok := r.getSession(workerID); ok {
+		r.cleanupSession(session)
+		return
+	}
 	r.registry.Remove(workerID)
 }