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.
This commit is contained in:
Chris Lu authored and GitHub committed 2026-11-03 12:28:39 -08:00
1 parent eeb8aefe70
commit 30f0364256
6 files changed
+692 -84

No files matched your search

+233 -66
View File
@@ -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)
}
@@ -0,0 +1,406 @@
package plugin
import (
"context"
"errors"
"testing"
"time"
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
"github.com/stretchr/testify/require"
)
func newDetectorSession(p *Plugin, workerID, jobType string) *streamSession {
p.registry.UpsertFromHello(&plugin_pb.WorkerHello{
WorkerId: workerID,
WorkerInstanceId: workerID + "-instance-" + time.Now().Format("150405.000000"),
Capabilities: []*plugin_pb.JobTypeCapability{{JobType: jobType, CanDetect: true, MaxDetectionConcurrency: 1}},
})
session := &streamSession{workerID: workerID, outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 8), done: make(chan struct{})}
p.putSession(session)
return session
}
func awaitDetectionRequest(t *testing.T, session *streamSession) *plugin_pb.AdminToWorkerMessage {
t.Helper()
select {
case request := <-session.outgoing:
require.NotNil(t, request.GetRunDetectionRequest(), "expected run_detection_request, got %+v", request)
return request
case <-time.After(2 * time.Second):
t.Fatal("no detection request")
return nil
}
}
func findActivity(activities []JobActivity, source, stage string) *JobActivity {
for i := range activities {
if activities[i].Source == source && activities[i].Stage == stage {
return &activities[i]
}
}
return nil
}
func (r *Plugin) detectionInFlightForTest(jobType string) bool {
r.schedulerMu.Lock()
defer r.schedulerMu.Unlock()
return r.detectionInFlight[jobType]
}
func (r *Plugin) pendingDetectionCountForTest() int {
r.pendingDetectionMu.Lock()
defer r.pendingDetectionMu.Unlock()
return len(r.pendingDetection)
}
func TestRunDetectionFailsWhenDetectorDisconnects(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)
p.WorkerDisconnectForTest("worker-a")
var got result
select {
case got = <-resultCh:
case <-time.After(5 * time.Second):
t.Fatal("detection did not fail after its worker disconnected")
}
require.Error(t, got.err)
var lost *detectionWorkerDisconnectedError
require.True(t, errors.As(got.err, &lost), "unexpected error type: %v", got.err)
require.Equal(t, "detection worker worker-a disconnected", got.err.Error())
require.NotNil(t, got.report)
require.Equal(t, request.RequestId, got.report.RequestID)
require.Equal(t, 0, p.pendingDetectionCountForTest())
require.Equal(t, "", p.getDetectorLease(jobType), "detector lease must not pin the lost worker")
activity := findActivity(p.ListActivities(jobType, 50), "detector", "disconnected")
require.NotNil(t, activity, "expected a disconnected detector activity")
require.Equal(t, request.RequestId, activity.RequestID)
require.Equal(t, "worker-a", activity.WorkerID)
// A replacement worker can serve the next detection right away.
next := newDetectorSession(p, "worker-b", jobType)
go func() {
_, _ = p.RunDetectionWithReport(ctx, jobType, &plugin_pb.ClusterContext{}, 10)
}()
awaitDetectionRequest(t, next)
}
func TestRunDetectionFailsWhenDetectorReconnectsWithNewSession(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
}()
awaitDetectionRequest(t, session)
// The worker process restarted: same worker ID, new instance, new stream.
replacement := newDetectorSession(p, "worker-a", jobType)
select {
case runErr := <-errCh:
var lost *detectionWorkerDisconnectedError
require.True(t, errors.As(runErr, &lost), "unexpected error: %v", runErr)
case <-time.After(5 * time.Second):
t.Fatal("detection did not fail after its worker reconnected with a new session")
}
require.Len(t, replacement.outgoing, 0, "the replacement session must not receive a cancel for the old request")
// The old stream unwinding afterwards must not evict the new session.
p.cleanupSession(session)
current, ok := p.getSession("worker-a")
require.True(t, ok)
require.Same(t, replacement, current)
_, registered := p.registry.Get("worker-a")
require.True(t, registered)
}
func TestRunDetectionSurvivesOtherWorkerDisconnect(t *testing.T) {
p, err := New(Options{})
require.NoError(t, err)
defer p.Shutdown()
const jobType = "vacuum"
sessionA := newDetectorSession(p, "worker-a", jobType)
sessionB := newDetectorSession(p, "worker-b", 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
}()
var request *plugin_pb.AdminToWorkerMessage
detector, other := "", ""
select {
case request = <-sessionA.outgoing:
detector, other = "worker-a", "worker-b"
case request = <-sessionB.outgoing:
detector, other = "worker-b", "worker-a"
case <-time.After(2 * time.Second):
t.Fatal("no detection request")
}
require.NotNil(t, request.GetRunDetectionRequest())
p.WorkerDisconnectForTest(other)
select {
case runErr := <-errCh:
t.Fatalf("detection on %s ended after %s disconnected: %v", detector, other, runErr)
case <-time.After(300 * time.Millisecond):
}
require.Equal(t, 1, p.pendingDetectionCountForTest())
p.handleDetectionComplete(detector, &plugin_pb.DetectionComplete{RequestId: request.RequestId, JobType: jobType, Success: true})
select {
case runErr := <-errCh:
require.NoError(t, runErr)
case <-time.After(2 * time.Second):
t.Fatal("detection did not complete")
}
}
func TestRunDetectionFailsWhenDetectorHeartbeatGoesStale(t *testing.T) {
p, err := New(Options{})
require.NoError(t, err)
defer p.Shutdown()
p.registry.staleAfter = 50 * time.Millisecond
p.detectionLivenessInterval = 20 * time.Millisecond
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
}()
request := awaitDetectionRequest(t, session)
select {
case runErr := <-errCh:
require.Error(t, runErr)
require.Equal(t, "detection worker worker-a disconnected: heartbeat stale", runErr.Error())
case <-time.After(5 * time.Second):
t.Fatal("detection did not fail after its worker went stale")
}
// The stream may still be open, so the worker is asked to stop the request.
select {
case message := <-session.outgoing:
cancelReq := message.GetCancelRequest()
require.NotNil(t, cancelReq, "expected cancel_request, got %+v", message)
require.Equal(t, request.RequestId, cancelReq.TargetId)
require.Equal(t, plugin_pb.WorkKind_WORK_KIND_DETECTION, cancelReq.TargetKind)
case <-time.After(2 * time.Second):
t.Fatal("no cancel request after stale worker")
}
}
func TestScheduledDetectionFailsFastOnWorkerDisconnect(t *testing.T) {
p, err := New(Options{})
require.NoError(t, err)
defer p.Shutdown()
p.clusterContextProvider = func(context.Context) (*plugin_pb.ClusterContext, error) {
return &plugin_pb.ClusterContext{}, nil
}
const jobType = "vacuum"
session := newDetectorSession(p, "worker-a", jobType)
done := make(chan struct{})
go func() {
p.runJobTypeIteration(jobType, schedulerPolicy{DetectionTimeout: time.Hour, JobTypeMaxRuntime: time.Hour})
close(done)
}()
awaitDetectionRequest(t, session)
require.True(t, p.detectionInFlightForTest(jobType))
p.WorkerDisconnectForTest("worker-a")
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("scheduled iteration kept waiting after its detector disconnected")
}
require.False(t, p.detectionInFlightForTest(jobType))
require.Equal(t, 0, p.pendingDetectionCountForTest())
require.Equal(t, "error", p.snapshotSchedulerRun(jobType).lastRunStatus)
require.False(t, p.snapshotSchedulerDetection(jobType).lastErrorAt.IsZero())
activities := p.ListActivities(jobType, 50)
require.NotNil(t, findActivity(activities, "detector", "disconnected"))
retrying := findActivity(activities, "admin_scheduler", "retrying")
require.NotNil(t, retrying, "the scheduler should have tried a replacement pick")
require.Contains(t, retrying.Message, "detection worker worker-a disconnected")
failed := findActivity(activities, "admin_scheduler", "failed")
require.NotNil(t, failed, "with no other detector the run must end as failed")
// The next cycle can detect again.
next := newDetectorSession(p, "worker-b", jobType)
go func() {
p.runJobTypeIteration(jobType, schedulerPolicy{DetectionTimeout: time.Hour, JobTypeMaxRuntime: time.Hour})
}()
awaitDetectionRequest(t, next)
}
func TestScheduledDetectionRetriesOnAnotherWorkerAfterDisconnect(t *testing.T) {
p, err := New(Options{})
require.NoError(t, err)
defer p.Shutdown()
p.clusterContextProvider = func(context.Context) (*plugin_pb.ClusterContext, error) {
return &plugin_pb.ClusterContext{}, nil
}
const jobType = "vacuum"
sessions := map[string]*streamSession{
"worker-a": newDetectorSession(p, "worker-a", jobType),
"worker-b": newDetectorSession(p, "worker-b", jobType),
}
done := make(chan struct{})
go func() {
p.runJobTypeIteration(jobType, schedulerPolicy{DetectionTimeout: time.Hour, JobTypeMaxRuntime: time.Hour})
close(done)
}()
first, other := "", ""
select {
case <-sessions["worker-a"].outgoing:
first, other = "worker-a", "worker-b"
case <-sessions["worker-b"].outgoing:
first, other = "worker-b", "worker-a"
case <-time.After(2 * time.Second):
t.Fatal("no detection request")
}
p.WorkerDisconnectForTest(first)
retry := awaitDetectionRequest(t, sessions[other])
p.handleDetectionComplete(other, &plugin_pb.DetectionComplete{RequestId: retry.RequestId, JobType: jobType, Success: true})
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("scheduled iteration did not finish after the retry completed")
}
require.False(t, p.detectionInFlightForTest(jobType))
require.Equal(t, "success", p.snapshotSchedulerRun(jobType).lastRunStatus)
require.False(t, p.snapshotSchedulerDetection(jobType).lastDetectedAt.IsZero())
require.NotNil(t, findActivity(p.ListActivities(jobType, 50), "admin_scheduler", "retrying"))
}
func TestRegisterSessionKeepsRegistryOnLiveStream(t *testing.T) {
p, err := New(Options{})
require.NoError(t, err)
defer p.Shutdown()
helloA := &plugin_pb.WorkerHello{WorkerId: "worker-a", WorkerInstanceId: "a-1", Capabilities: []*plugin_pb.JobTypeCapability{{JobType: "ec_vacuum", CanDetect: true}}}
helloB := &plugin_pb.WorkerHello{WorkerId: "worker-a", WorkerInstanceId: "a-2", Capabilities: []*plugin_pb.JobTypeCapability{{JobType: "s3_lifecycle", CanDetect: true}}}
streamA := &streamSession{workerID: "worker-a", outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 1), done: make(chan struct{})}
streamB := &streamSession{workerID: "worker-a", outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 1), done: make(chan struct{})}
p.registerSession(streamA, helloA)
p.registerSession(streamB, helloB)
select {
case <-streamA.done:
default:
t.Fatal("the replaced stream must be closed")
}
// Stream A's hello and heartbeat arriving after it was replaced must not
// describe the live stream B, and neither must A's own cleanup.
p.handleWorkerMessage(streamA, &plugin_pb.WorkerToAdminMessage{Body: &plugin_pb.WorkerToAdminMessage_Hello{Hello: helloA}})
p.handleWorkerMessage(streamA, &plugin_pb.WorkerToAdminMessage{Body: &plugin_pb.WorkerToAdminMessage_Heartbeat{Heartbeat: &plugin_pb.WorkerHeartbeat{DetectionSlotsTotal: 9}}})
p.cleanupSession(streamA)
worker, ok := p.registry.Get("worker-a")
require.True(t, ok)
require.Equal(t, "a-2", worker.WorkerInstance)
require.Nil(t, worker.Capabilities["ec_vacuum"])
require.NotNil(t, worker.Capabilities["s3_lifecycle"])
require.Nil(t, worker.Heartbeat, "a superseded stream's heartbeat must be dropped")
current, ok := p.getSession("worker-a")
require.True(t, ok)
require.Same(t, streamB, current)
_, err = p.pickDetector("ec_vacuum", nil)
require.Error(t, err, "the capability only stream A advertised must be gone")
p.handleWorkerMessage(streamB, &plugin_pb.WorkerToAdminMessage{Body: &plugin_pb.WorkerToAdminMessage_Heartbeat{Heartbeat: &plugin_pb.WorkerHeartbeat{DetectionSlotsTotal: 2}}})
worker, ok = p.registry.Get("worker-a")
require.True(t, ok)
require.NotNil(t, worker.Heartbeat)
require.EqualValues(t, 2, worker.Heartbeat.DetectionSlotsTotal)
}
func TestScheduledDetectionRetryAvoidsClosedDetectorStream(t *testing.T) {
p, err := New(Options{})
require.NoError(t, err)
defer p.Shutdown()
p.clusterContextProvider = func(context.Context) (*plugin_pb.ClusterContext, error) {
return &plugin_pb.ClusterContext{}, nil
}
const jobType = "vacuum"
sessionA := newDetectorSession(p, "worker-a", jobType)
sessionB := newDetectorSession(p, "worker-b", jobType)
// worker-a ranks first on free detection slots, for the first pick and
// for any retry that still sees it in the registry.
p.registry.UpdateHeartbeat("worker-a", &plugin_pb.WorkerHeartbeat{DetectionSlotsTotal: 4})
p.registry.UpdateHeartbeat("worker-b", &plugin_pb.WorkerHeartbeat{DetectionSlotsTotal: 1})
done := make(chan struct{})
go func() {
p.runJobTypeIteration(jobType, schedulerPolicy{DetectionTimeout: time.Hour, JobTypeMaxRuntime: time.Hour})
close(done)
}()
awaitDetectionRequest(t, sessionA)
// Close the stream while its registry entry is still present: the
// ordering where the retry picks before the disconnect is fully cleaned up.
sessionA.close()
retry := awaitDetectionRequest(t, sessionB)
require.Equal(t, "worker-b", p.getDetectorLease(jobType))
p.handleDetectionComplete("worker-b", &plugin_pb.DetectionComplete{RequestId: retry.RequestId, JobType: jobType, Success: true})
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("scheduled iteration did not finish after the retry completed")
}
require.Equal(t, "success", p.snapshotSchedulerRun(jobType).lastRunStatus)
require.Len(t, sessionA.outgoing, 0, "nothing may be sent on the closed stream")
}
+27 -9
View File
@@ -322,7 +322,6 @@ func (r *Plugin) runJobTypeIteration(jobType string, policy schedulerPolicy) boo
}
defer releaseDetectionLock()
detectionTimeout := policy.DetectionTimeout
remaining := time.Until(start.Add(maxRuntime))
if remaining <= 0 {
r.appendActivity(JobActivity{
@@ -335,16 +334,35 @@ func (r *Plugin) runJobTypeIteration(jobType string, policy schedulerPolicy) boo
r.recordSchedulerRunComplete(jobType, "timeout")
return false
}
if detectionTimeout <= 0 {
detectionTimeout = defaultScheduledDetectionTimeout
}
if detectionTimeout > remaining {
detectionTimeout = remaining
detect := func(avoid *streamSession) (*DetectionReport, error) {
detectionTimeout := policy.DetectionTimeout
if detectionTimeout <= 0 {
detectionTimeout = defaultScheduledDetectionTimeout
}
if remaining := time.Until(start.Add(maxRuntime)); detectionTimeout > remaining {
detectionTimeout = remaining
}
detectCtx, cancelDetect := context.WithTimeout(jobCtx, detectionTimeout)
defer cancelDetect()
return r.runDetectionWithReport(detectCtx, jobType, clusterContext, policy.MaxResults, avoid)
}
detectCtx, cancelDetect := context.WithTimeout(jobCtx, detectionTimeout)
report, err := r.RunDetectionWithReport(detectCtx, jobType, clusterContext, policy.MaxResults)
cancelDetect()
report, err := detect(nil)
// A detector that dropped mid-request gets one replacement pick, steered
// off the stream that dropped, in this iteration while the runtime budget
// allows; anything else waits for the next cycle, which re-detects anyway.
var lost *detectionWorkerDisconnectedError
if errors.As(err, &lost) && time.Until(start.Add(maxRuntime)) > 0 {
r.recordSchedulerDetectionError(jobType, err)
r.appendActivity(JobActivity{
JobType: jobType,
Source: "admin_scheduler",
Message: fmt.Sprintf("scheduled detection retrying on another worker: %v", err),
Stage: "retrying",
OccurredAt: timeToPtr(time.Now().UTC()),
})
report, err = detect(lost.session)
}
releaseDetectionLock()
if err != nil {
r.recordSchedulerDetectionError(jobType, err)
+5 -2
View File
@@ -562,9 +562,11 @@ func TestPickDetectorPrefersLeasedWorker(t *testing.T) {
},
})
pluginSvc.putSession(&streamSession{workerID: "worker-a", outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 1), done: make(chan struct{})})
pluginSvc.putSession(&streamSession{workerID: "worker-b", outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 1), done: make(chan struct{})})
pluginSvc.setDetectorLease("vacuum", "worker-b")
detector, err := pluginSvc.pickDetector("vacuum")
detector, err := pluginSvc.pickDetector("vacuum", nil)
if err != nil {
t.Fatalf("pickDetector: %v", err)
}
@@ -588,9 +590,10 @@ func TestPickDetectorReassignsWhenLeaseIsStale(t *testing.T) {
{JobType: "vacuum", CanDetect: true},
},
})
pluginSvc.putSession(&streamSession{workerID: "worker-a", outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 1), done: make(chan struct{})})
pluginSvc.setDetectorLease("vacuum", "worker-stale")
detector, err := pluginSvc.pickDetector("vacuum")
detector, err := pluginSvc.pickDetector("vacuum", nil)
if err != nil {
t.Fatalf("pickDetector: %v", err)
}
+3 -3
View File
@@ -69,7 +69,7 @@ func TestRequestObjectPreviewReturnsTheWorkersRows(t *testing.T) {
t.Fatalf("request lost its parameters: %+v", asked)
}
pluginSvc.handleWorkerMessage(workerID, &plugin_pb.WorkerToAdminMessage{
pluginSvc.handleWorkerMessage(session, &plugin_pb.WorkerToAdminMessage{
WorkerId: workerID,
Body: &plugin_pb.WorkerToAdminMessage_ObjectPreviewResponse{
ObjectPreviewResponse: &plugin_pb.ObjectPreviewResponse{
@@ -114,7 +114,7 @@ func TestRequestObjectPreviewSurfacesTheWorkersError(t *testing.T) {
}()
request := <-session.outgoing
pluginSvc.handleWorkerMessage(workerID, &plugin_pb.WorkerToAdminMessage{
pluginSvc.handleWorkerMessage(session, &plugin_pb.WorkerToAdminMessage{
WorkerId: workerID,
Body: &plugin_pb.WorkerToAdminMessage_ObjectPreviewResponse{
ObjectPreviewResponse: &plugin_pb.ObjectPreviewResponse{
@@ -187,7 +187,7 @@ func TestObjectPreviewResponseRacesShutdown(t *testing.T) {
}()
go func() {
defer wg.Done()
pluginSvc.handleWorkerMessage(workerID, &plugin_pb.WorkerToAdminMessage{
pluginSvc.handleWorkerMessage(session, &plugin_pb.WorkerToAdminMessage{
WorkerId: workerID,
Body: &plugin_pb.WorkerToAdminMessage_ObjectPreviewResponse{
ObjectPreviewResponse: &plugin_pb.ObjectPreviewResponse{
+18 -4
View File
@@ -272,9 +272,10 @@ func (r *Registry) PickSchemaProvider(jobType string) (*WorkerSession, error) {
return cloneWorkerSession(candidates[0]), nil
}
// PickDetector picks one detector worker for a job type.
func (r *Registry) PickDetector(jobType string) (*WorkerSession, error) {
return r.pickByKind(jobType, true)
// PickDetector picks one detector worker for a job type, skipping any
// worker IDs in excluded.
func (r *Registry) PickDetector(jobType string, excluded ...string) (*WorkerSession, error) {
return r.pickByKind(jobType, true, excluded)
}
// PickExecutor picks one executor worker for a job type.
@@ -322,11 +323,24 @@ func (r *Registry) ListExecutors(jobType string, preferredHosts ...string) ([]*W
return out, nil
}
func (r *Registry) pickByKind(jobType string, detect bool) (*WorkerSession, error) {
func (r *Registry) pickByKind(jobType string, detect bool, excluded []string) (*WorkerSession, error) {
r.mu.Lock()
defer r.mu.Unlock()
candidates := r.collectByKindLocked(jobType, detect, time.Now())
if len(excluded) > 0 {
skip := make(map[string]struct{}, len(excluded))
for _, workerID := range excluded {
skip[workerID] = struct{}{}
}
kept := candidates[:0]
for _, candidate := range candidates {
if _, found := skip[candidate.WorkerID]; !found {
kept = append(kept, candidate)
}
}
candidates = kept
}
if len(candidates) == 0 {
kind := "executor"