mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-11-04 11:07:35 +01:00
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.
1706 lines
49 KiB
Go
1706 lines
49 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/seaweedfs/seaweedfs/weed/glog"
|
|
"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
|
|
"github.com/seaweedfs/seaweedfs/weed/security"
|
|
"github.com/seaweedfs/seaweedfs/weed/util"
|
|
"google.golang.org/protobuf/proto"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
const (
|
|
defaultOutgoingBuffer = 128
|
|
defaultSendTimeout = 5 * time.Second
|
|
defaultHeartbeatInterval = 30
|
|
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 {
|
|
DataDir string
|
|
OutgoingBufferSize int
|
|
SendTimeout time.Duration
|
|
SchedulerTick time.Duration
|
|
ClusterContextProvider func(context.Context) (*plugin_pb.ClusterContext, error)
|
|
LockManager LockManager
|
|
// ConfigDefaultsProvider is an optional callback invoked when a job type's
|
|
// config is being bootstrapped from its descriptor defaults. It can enrich
|
|
// or replace the default config before it is persisted. If nil, descriptor
|
|
// defaults are used as-is.
|
|
ConfigDefaultsProvider func(config *plugin_pb.PersistedJobTypeConfig) *plugin_pb.PersistedJobTypeConfig
|
|
}
|
|
|
|
// JobTypeInfo contains metadata about a plugin job type.
|
|
type JobTypeInfo struct {
|
|
JobType string `json:"job_type"`
|
|
DisplayName string `json:"display_name"`
|
|
Weight int32 `json:"weight"`
|
|
}
|
|
|
|
type Plugin struct {
|
|
plugin_pb.UnimplementedPluginControlServiceServer
|
|
|
|
store *ConfigStore
|
|
registry *Registry
|
|
|
|
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
|
|
lockManager LockManager
|
|
|
|
schedulerMu sync.Mutex
|
|
nextDetectionAt map[string]time.Time
|
|
detectionInFlight map[string]bool
|
|
|
|
detectorLeaseMu sync.Mutex
|
|
detectorLeases map[string]string
|
|
|
|
lanes map[SchedulerLane]*schedulerLaneState
|
|
|
|
schedulerExecMu sync.Mutex
|
|
schedulerExecReservations map[string]int
|
|
adminScriptRunMu sync.RWMutex
|
|
schedulerDetectionMu sync.Mutex
|
|
schedulerDetection map[string]*schedulerDetectionInfo
|
|
schedulerRunMu sync.Mutex
|
|
schedulerRun map[string]*schedulerRunInfo
|
|
schedulerLoopMu sync.Mutex
|
|
schedulerLoopState schedulerLoopState
|
|
schedulerWakeCh chan struct{}
|
|
|
|
dedupeMu sync.Mutex
|
|
recentDedupeByType map[string]map[string]time.Time
|
|
|
|
sessionsMu sync.RWMutex
|
|
sessions map[string]*streamSession
|
|
|
|
pendingSchemaMu sync.Mutex
|
|
pendingSchema map[string]chan *plugin_pb.ConfigSchemaResponse
|
|
|
|
pendingDetectionMu sync.Mutex
|
|
pendingDetection map[string]*pendingDetectionState
|
|
|
|
pendingExecutionMu sync.Mutex
|
|
pendingExecution map[string]chan *plugin_pb.JobCompleted
|
|
|
|
pendingPreviewMu sync.Mutex
|
|
pendingPreview map[string]chan *plugin_pb.ObjectPreviewResponse
|
|
|
|
observations *ObservationStore
|
|
|
|
jobsMu sync.RWMutex
|
|
jobs map[string]*TrackedJob
|
|
// serialize stale job cleanup to avoid duplicate expirations
|
|
staleJobsMu sync.Mutex
|
|
|
|
jobDetailsMu sync.Mutex
|
|
|
|
activitiesMu sync.RWMutex
|
|
activities []JobActivity
|
|
|
|
dirtyJobs bool
|
|
dirtyActivities bool
|
|
persistTicker *time.Ticker
|
|
|
|
ctx context.Context
|
|
ctxCancel context.CancelFunc
|
|
|
|
shutdownCh chan struct{}
|
|
shutdownOnce sync.Once
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
type streamSession struct {
|
|
workerID string
|
|
outgoing chan *plugin_pb.AdminToWorkerMessage
|
|
done chan struct{}
|
|
closeOnce sync.Once
|
|
}
|
|
|
|
type pendingDetectionState struct {
|
|
proposals []*plugin_pb.JobProposal
|
|
complete chan *plugin_pb.DetectionComplete
|
|
jobType string
|
|
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
|
|
JobType string
|
|
WorkerID string
|
|
Proposals []*plugin_pb.JobProposal
|
|
Complete *plugin_pb.DetectionComplete
|
|
}
|
|
|
|
func New(options Options) (*Plugin, error) {
|
|
store, err := NewConfigStore(options.DataDir)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
bufferSize := options.OutgoingBufferSize
|
|
if bufferSize <= 0 {
|
|
bufferSize = defaultOutgoingBuffer
|
|
}
|
|
sendTimeout := options.SendTimeout
|
|
if sendTimeout <= 0 {
|
|
sendTimeout = defaultSendTimeout
|
|
}
|
|
schedulerTick := options.SchedulerTick
|
|
if schedulerTick <= 0 {
|
|
schedulerTick = defaultSchedulerTick
|
|
}
|
|
|
|
lanes := make(map[SchedulerLane]*schedulerLaneState, len(AllLanes()))
|
|
for _, lane := range AllLanes() {
|
|
lanes[lane] = newLaneState(lane)
|
|
}
|
|
|
|
plugin := &Plugin{
|
|
store: store,
|
|
registry: NewRegistry(),
|
|
outgoingBuffer: bufferSize,
|
|
sendTimeout: sendTimeout,
|
|
detectionLivenessInterval: defaultDetectionLivenessInterval,
|
|
schedulerTick: schedulerTick,
|
|
clusterContextProvider: options.ClusterContextProvider,
|
|
configDefaultsProvider: options.ConfigDefaultsProvider,
|
|
lockManager: options.LockManager,
|
|
lanes: lanes,
|
|
sessions: make(map[string]*streamSession),
|
|
pendingSchema: make(map[string]chan *plugin_pb.ConfigSchemaResponse),
|
|
pendingDetection: make(map[string]*pendingDetectionState),
|
|
pendingExecution: make(map[string]chan *plugin_pb.JobCompleted),
|
|
pendingPreview: make(map[string]chan *plugin_pb.ObjectPreviewResponse),
|
|
observations: NewObservationStore(),
|
|
nextDetectionAt: make(map[string]time.Time),
|
|
detectionInFlight: make(map[string]bool),
|
|
detectorLeases: make(map[string]string),
|
|
schedulerExecReservations: make(map[string]int),
|
|
schedulerDetection: make(map[string]*schedulerDetectionInfo),
|
|
schedulerRun: make(map[string]*schedulerRunInfo),
|
|
recentDedupeByType: make(map[string]map[string]time.Time),
|
|
jobs: make(map[string]*TrackedJob),
|
|
activities: make([]JobActivity, 0, 256),
|
|
persistTicker: time.NewTicker(2 * time.Second),
|
|
schedulerWakeCh: make(chan struct{}, 1),
|
|
shutdownCh: make(chan struct{}),
|
|
}
|
|
plugin.ctx, plugin.ctxCancel = context.WithCancel(context.Background())
|
|
|
|
if err := plugin.loadPersistedMonitorState(); err != nil {
|
|
glog.Warningf("Plugin failed to load persisted monitoring state: %v", err)
|
|
}
|
|
|
|
if plugin.clusterContextProvider != nil {
|
|
for _, ls := range plugin.lanes {
|
|
plugin.wg.Add(1)
|
|
go plugin.laneSchedulerLoop(ls)
|
|
}
|
|
}
|
|
plugin.wg.Add(1)
|
|
go plugin.persistenceLoop()
|
|
|
|
return plugin, nil
|
|
}
|
|
|
|
func (r *Plugin) Shutdown() {
|
|
if r.ctxCancel != nil {
|
|
r.ctxCancel()
|
|
}
|
|
if r.persistTicker != nil {
|
|
r.persistTicker.Stop()
|
|
}
|
|
|
|
r.shutdownOnce.Do(func() { close(r.shutdownCh) })
|
|
|
|
r.sessionsMu.Lock()
|
|
for workerID, session := range r.sessions {
|
|
session.close()
|
|
delete(r.sessions, workerID)
|
|
}
|
|
r.sessionsMu.Unlock()
|
|
|
|
r.pendingSchemaMu.Lock()
|
|
for requestID, ch := range r.pendingSchema {
|
|
close(ch)
|
|
delete(r.pendingSchema, requestID)
|
|
}
|
|
r.pendingSchemaMu.Unlock()
|
|
|
|
r.pendingPreviewMu.Lock()
|
|
for requestID, ch := range r.pendingPreview {
|
|
close(ch)
|
|
delete(r.pendingPreview, requestID)
|
|
}
|
|
r.pendingPreviewMu.Unlock()
|
|
|
|
r.pendingDetectionMu.Lock()
|
|
for requestID, state := range r.pendingDetection {
|
|
close(state.complete)
|
|
delete(r.pendingDetection, requestID)
|
|
}
|
|
r.pendingDetectionMu.Unlock()
|
|
|
|
r.pendingExecutionMu.Lock()
|
|
for requestID, ch := range r.pendingExecution {
|
|
close(ch)
|
|
delete(r.pendingExecution, requestID)
|
|
}
|
|
r.pendingExecutionMu.Unlock()
|
|
|
|
r.wg.Wait()
|
|
}
|
|
|
|
func (r *Plugin) WorkerStream(stream plugin_pb.PluginControlService_WorkerStreamServer) error {
|
|
if err := security.CheckGrpcAdminPeer(stream.Context(), security.SigningKey(util.GetViper().GetString("jwt.filer_signing.key"))); err != nil {
|
|
return err
|
|
}
|
|
|
|
first, err := stream.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
return nil
|
|
}
|
|
return fmt.Errorf("receive worker hello: %w", err)
|
|
}
|
|
|
|
hello := first.GetHello()
|
|
if hello == nil {
|
|
return fmt.Errorf("first message must be hello")
|
|
}
|
|
if strings.TrimSpace(hello.WorkerId) == "" {
|
|
return fmt.Errorf("worker_id is required")
|
|
}
|
|
|
|
workerID := hello.WorkerId
|
|
session := &streamSession{
|
|
workerID: workerID,
|
|
outgoing: make(chan *plugin_pb.AdminToWorkerMessage, r.outgoingBuffer),
|
|
done: make(chan struct{}),
|
|
}
|
|
r.registerSession(session, hello)
|
|
defer r.cleanupSession(session)
|
|
|
|
glog.V(0).Infof("Plugin worker connected: %s (%s)", workerID, hello.Address)
|
|
|
|
sendErrCh := make(chan error, 1)
|
|
r.wg.Add(1)
|
|
go func() {
|
|
defer r.wg.Done()
|
|
sendErrCh <- r.sendLoop(stream.Context(), stream, session)
|
|
}()
|
|
|
|
if err := r.sendAdminHello(session); err != nil {
|
|
glog.Warningf("failed to send plugin admin hello to %s: %v", workerID, err)
|
|
}
|
|
go r.prefetchDescriptorsFromHello(hello)
|
|
|
|
recvErrCh := make(chan error, 1)
|
|
r.wg.Add(1)
|
|
go func() {
|
|
defer r.wg.Done()
|
|
for {
|
|
message, recvErr := stream.Recv()
|
|
if recvErr != nil {
|
|
recvErrCh <- recvErr
|
|
return
|
|
}
|
|
r.handleWorkerMessage(session, message)
|
|
}
|
|
}()
|
|
|
|
for {
|
|
select {
|
|
case <-r.shutdownCh:
|
|
return nil
|
|
case err := <-sendErrCh:
|
|
if err != nil && !errors.Is(err, context.Canceled) {
|
|
return err
|
|
}
|
|
return nil
|
|
case recvErr := <-recvErrCh:
|
|
if errors.Is(recvErr, io.EOF) || errors.Is(recvErr, context.Canceled) {
|
|
return nil
|
|
}
|
|
return fmt.Errorf("receive plugin message from %s: %w", workerID, recvErr)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) RequestConfigSchema(ctx context.Context, jobType string, forceRefresh bool) (*plugin_pb.JobTypeDescriptor, error) {
|
|
if !forceRefresh {
|
|
descriptor, err := r.store.LoadDescriptor(jobType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if descriptor != nil {
|
|
return descriptor, nil
|
|
}
|
|
}
|
|
|
|
provider, err := r.registry.PickSchemaProvider(jobType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return r.requestConfigSchema(ctx, jobType, forceRefresh, provider.WorkerID)
|
|
}
|
|
|
|
func (r *Plugin) requestConfigSchema(ctx context.Context, jobType string, forceRefresh bool, workerID string) (*plugin_pb.JobTypeDescriptor, error) {
|
|
requestID, err := newRequestID("schema")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
responseCh := make(chan *plugin_pb.ConfigSchemaResponse, defaultPendingSchemaBuffer)
|
|
r.pendingSchemaMu.Lock()
|
|
r.pendingSchema[requestID] = responseCh
|
|
r.pendingSchemaMu.Unlock()
|
|
defer func() {
|
|
r.pendingSchemaMu.Lock()
|
|
delete(r.pendingSchema, requestID)
|
|
r.pendingSchemaMu.Unlock()
|
|
}()
|
|
|
|
requestMessage := &plugin_pb.AdminToWorkerMessage{
|
|
RequestId: requestID,
|
|
SentAt: timestamppb.Now(),
|
|
Body: &plugin_pb.AdminToWorkerMessage_RequestConfigSchema{
|
|
RequestConfigSchema: &plugin_pb.RequestConfigSchema{
|
|
JobType: jobType,
|
|
ForceRefresh: forceRefresh,
|
|
},
|
|
},
|
|
}
|
|
|
|
if err := r.sendToWorker(workerID, requestMessage); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case response, ok := <-responseCh:
|
|
if !ok {
|
|
return nil, fmt.Errorf("schema request %s interrupted", requestID)
|
|
}
|
|
if response == nil {
|
|
return nil, fmt.Errorf("schema request %s returned empty response", requestID)
|
|
}
|
|
if !response.Success {
|
|
return nil, fmt.Errorf("schema request failed for %s: %s", jobType, response.ErrorMessage)
|
|
}
|
|
if response.GetJobTypeDescriptor() == nil {
|
|
return nil, fmt.Errorf("schema request for %s returned no descriptor", jobType)
|
|
}
|
|
return response.GetJobTypeDescriptor(), nil
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) LoadJobTypeConfig(jobType string) (*plugin_pb.PersistedJobTypeConfig, error) {
|
|
return r.store.LoadJobTypeConfig(jobType)
|
|
}
|
|
|
|
func (r *Plugin) SaveJobTypeConfig(config *plugin_pb.PersistedJobTypeConfig) error {
|
|
if err := r.store.SaveJobTypeConfig(config); err != nil {
|
|
return err
|
|
}
|
|
r.wakeScheduler()
|
|
return nil
|
|
}
|
|
|
|
func (r *Plugin) LoadDescriptor(jobType string) (*plugin_pb.JobTypeDescriptor, error) {
|
|
return r.store.LoadDescriptor(jobType)
|
|
}
|
|
|
|
func (r *Plugin) LoadRunHistory(jobType string) (*JobTypeRunHistory, error) {
|
|
return r.store.LoadRunHistory(jobType)
|
|
}
|
|
|
|
func (r *Plugin) IsConfigured() bool {
|
|
return r.store.IsConfigured()
|
|
}
|
|
|
|
func (r *Plugin) BaseDir() string {
|
|
return r.store.BaseDir()
|
|
}
|
|
|
|
func (r *Plugin) acquireAdminLock(reason string) (func(), error) {
|
|
if r == nil || r.lockManager == nil {
|
|
return func() {}, nil
|
|
}
|
|
release, err := r.lockManager.Acquire(reason)
|
|
if release == nil {
|
|
// Lock manager implementations may return a nil release when
|
|
// disabled; callers always invoke the release function.
|
|
release = func() {}
|
|
}
|
|
return release, err
|
|
}
|
|
|
|
// RunDetectionWithReport requests one detector worker and returns proposals with request metadata.
|
|
func (r *Plugin) RunDetectionWithReport(
|
|
ctx context.Context,
|
|
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, avoid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
requestID, err := newRequestID("detect")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
adminRuntime, adminConfigValues, workerConfigValues, err := r.loadJobTypeConfigPayload(jobType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
lastCompletedRun := r.loadLastSuccessfulRun(jobType)
|
|
if strings.EqualFold(strings.TrimSpace(jobType), adminScriptJobType) {
|
|
lastCompletedRun = r.loadLastCompletedRun(jobType)
|
|
}
|
|
|
|
state := &pendingDetectionState{
|
|
complete: make(chan *plugin_pb.DetectionComplete, 1),
|
|
jobType: jobType,
|
|
workerID: detector.WorkerID,
|
|
}
|
|
r.pendingDetectionMu.Lock()
|
|
r.pendingDetection[requestID] = state
|
|
r.pendingDetectionMu.Unlock()
|
|
defer func() {
|
|
r.pendingDetectionMu.Lock()
|
|
delete(r.pendingDetection, requestID)
|
|
r.pendingDetectionMu.Unlock()
|
|
}()
|
|
|
|
r.appendActivity(JobActivity{
|
|
JobType: jobType,
|
|
RequestID: requestID,
|
|
WorkerID: detector.WorkerID,
|
|
Source: "detector",
|
|
Stage: "requested",
|
|
Message: "detection requested",
|
|
OccurredAt: timeToPtr(time.Now().UTC()),
|
|
Details: map[string]interface{}{
|
|
"max_results": maxResults,
|
|
},
|
|
})
|
|
|
|
message := &plugin_pb.AdminToWorkerMessage{
|
|
RequestId: requestID,
|
|
SentAt: timestamppb.Now(),
|
|
Body: &plugin_pb.AdminToWorkerMessage_RunDetectionRequest{
|
|
RunDetectionRequest: &plugin_pb.RunDetectionRequest{
|
|
RequestId: requestID,
|
|
JobType: jobType,
|
|
DetectionSequence: time.Now().UnixNano(),
|
|
AdminRuntime: adminRuntime,
|
|
AdminConfigValues: adminConfigValues,
|
|
WorkerConfigValues: workerConfigValues,
|
|
ClusterContext: clusterContext,
|
|
LastSuccessfulRun: lastCompletedRun,
|
|
MaxResults: maxResults,
|
|
},
|
|
},
|
|
}
|
|
|
|
// 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,
|
|
RequestID: requestID,
|
|
WorkerID: detector.WorkerID,
|
|
Source: "detector",
|
|
Stage: "failed_to_send",
|
|
Message: err.Error(),
|
|
OccurredAt: timeToPtr(time.Now().UTC()),
|
|
})
|
|
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: 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
|
|
}
|
|
}
|
|
|
|
livenessInterval := r.detectionLivenessInterval
|
|
if livenessInterval <= 0 {
|
|
livenessInterval = defaultDetectionLivenessInterval
|
|
}
|
|
liveness := time.NewTicker(livenessInterval)
|
|
defer liveness.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
r.sendCancel(workerID, requestID, plugin_pb.WorkKind_WORK_KIND_DETECTION, ctx.Err())
|
|
r.appendActivity(JobActivity{
|
|
JobType: jobType,
|
|
RequestID: requestID,
|
|
WorkerID: workerID,
|
|
Source: "detector",
|
|
Stage: "canceled",
|
|
Message: "detection canceled",
|
|
OccurredAt: timeToPtr(time.Now().UTC()),
|
|
})
|
|
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(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:
|
|
return finish(complete, ok)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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,
|
|
jobType string,
|
|
clusterContext *plugin_pb.ClusterContext,
|
|
maxResults int32,
|
|
) ([]*plugin_pb.JobProposal, error) {
|
|
report, err := r.RunDetectionWithReport(ctx, jobType, clusterContext, maxResults)
|
|
if report == nil {
|
|
return nil, err
|
|
}
|
|
return report.Proposals, err
|
|
}
|
|
|
|
// ExecuteJob sends one job to a capable executor worker and waits for completion.
|
|
func (r *Plugin) ExecuteJob(
|
|
ctx context.Context,
|
|
job *plugin_pb.JobSpec,
|
|
clusterContext *plugin_pb.ClusterContext,
|
|
attempt int32,
|
|
) (*plugin_pb.JobCompleted, error) {
|
|
if job == nil {
|
|
return nil, fmt.Errorf("job is nil")
|
|
}
|
|
jobType := strings.TrimSpace(job.JobType)
|
|
if jobType == "" {
|
|
return nil, fmt.Errorf("job_type is required")
|
|
}
|
|
releaseGate := r.acquireDetectionExecutionGate(jobType, true)
|
|
defer releaseGate()
|
|
|
|
executor, err := r.registry.PickExecutor(jobType, JobReplicaHosts(job)...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return r.executeJobWithExecutor(ctx, executor, job, clusterContext, attempt)
|
|
}
|
|
|
|
func (r *Plugin) acquireDetectionExecutionGate(jobType string, execution bool) func() {
|
|
normalizedJobType := strings.ToLower(strings.TrimSpace(jobType))
|
|
if execution && normalizedJobType == adminScriptJobType {
|
|
r.adminScriptRunMu.Lock()
|
|
return func() {
|
|
r.adminScriptRunMu.Unlock()
|
|
}
|
|
}
|
|
if normalizedJobType != adminScriptJobType {
|
|
r.adminScriptRunMu.RLock()
|
|
return func() {
|
|
r.adminScriptRunMu.RUnlock()
|
|
}
|
|
}
|
|
return func() {}
|
|
}
|
|
|
|
func (r *Plugin) executeJobWithExecutor(
|
|
ctx context.Context,
|
|
executor *WorkerSession,
|
|
job *plugin_pb.JobSpec,
|
|
clusterContext *plugin_pb.ClusterContext,
|
|
attempt int32,
|
|
) (*plugin_pb.JobCompleted, error) {
|
|
if executor == nil {
|
|
return nil, fmt.Errorf("executor is nil")
|
|
}
|
|
if job == nil {
|
|
return nil, fmt.Errorf("job is nil")
|
|
}
|
|
if strings.TrimSpace(job.JobType) == "" {
|
|
return nil, fmt.Errorf("job_type is required")
|
|
}
|
|
|
|
if strings.TrimSpace(job.JobId) == "" {
|
|
var err error
|
|
job.JobId, err = newRequestID("job")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
requestID, err := newRequestID("exec")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
adminRuntime, adminConfigValues, workerConfigValues, err := r.loadJobTypeConfigPayload(job.JobType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Apply per-job-type cluster-allocation decoration (e.g. s3_lifecycle
|
|
// divides cluster_deletes_per_second by min(workers, maxJobsPerDetection)
|
|
// and ships the share via ClusterContext.Metadata). No-op for job
|
|
// types without an allocator registered. MaxJobsPerDetection caps
|
|
// the divisor so a singleton job (maxJobs=1) gets the full budget on
|
|
// the single active worker, not 1/N of it.
|
|
clusterContext = r.decorateClusterContextForJob(clusterContext, job.JobType, adminConfigValues, int(adminRuntime.GetMaxJobsPerDetection()))
|
|
|
|
completedCh := make(chan *plugin_pb.JobCompleted, 1)
|
|
r.pendingExecutionMu.Lock()
|
|
r.pendingExecution[requestID] = completedCh
|
|
r.pendingExecutionMu.Unlock()
|
|
defer func() {
|
|
r.pendingExecutionMu.Lock()
|
|
delete(r.pendingExecution, requestID)
|
|
r.pendingExecutionMu.Unlock()
|
|
}()
|
|
|
|
r.trackExecutionStart(requestID, executor.WorkerID, job, attempt)
|
|
|
|
message := &plugin_pb.AdminToWorkerMessage{
|
|
RequestId: requestID,
|
|
SentAt: timestamppb.Now(),
|
|
Body: &plugin_pb.AdminToWorkerMessage_ExecuteJobRequest{
|
|
ExecuteJobRequest: &plugin_pb.ExecuteJobRequest{
|
|
RequestId: requestID,
|
|
Job: job,
|
|
AdminRuntime: adminRuntime,
|
|
AdminConfigValues: adminConfigValues,
|
|
WorkerConfigValues: workerConfigValues,
|
|
ClusterContext: clusterContext,
|
|
Attempt: attempt,
|
|
},
|
|
},
|
|
}
|
|
|
|
if err := r.sendToWorker(executor.WorkerID, message); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
r.sendCancel(executor.WorkerID, requestID, plugin_pb.WorkKind_WORK_KIND_EXECUTION, ctx.Err())
|
|
return nil, ctx.Err()
|
|
case completed, ok := <-completedCh:
|
|
if !ok {
|
|
return nil, fmt.Errorf("execution request %s interrupted", requestID)
|
|
}
|
|
if completed == nil {
|
|
return nil, fmt.Errorf("execution request %s returned empty completion", requestID)
|
|
}
|
|
if !completed.Success {
|
|
return completed, fmt.Errorf("job %s failed: %s", job.JobId, completed.ErrorMessage)
|
|
}
|
|
return completed, nil
|
|
}
|
|
}
|
|
|
|
// HasCapableWorker checks if any non-stale worker has a capability for the given job type.
|
|
func (r *Plugin) HasCapableWorker(jobType string) bool {
|
|
return r.registry.HasCapableWorker(jobType)
|
|
}
|
|
|
|
func (r *Plugin) ListWorkers() []*WorkerSession {
|
|
return r.registry.List()
|
|
}
|
|
|
|
func (r *Plugin) ListKnownJobTypes() ([]JobTypeInfo, error) {
|
|
registryJobTypes := r.registry.JobTypes()
|
|
storedJobTypes, err := r.store.ListJobTypes()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
jobTypeSet := make(map[string]struct{}, len(registryJobTypes)+len(storedJobTypes))
|
|
for _, jobType := range registryJobTypes {
|
|
jobTypeSet[jobType] = struct{}{}
|
|
}
|
|
for _, jobType := range storedJobTypes {
|
|
jobTypeSet[jobType] = struct{}{}
|
|
}
|
|
|
|
jobTypeList := make([]string, 0, len(jobTypeSet))
|
|
for jobType := range jobTypeSet {
|
|
jobTypeList = append(jobTypeList, jobType)
|
|
}
|
|
sort.Strings(jobTypeList)
|
|
|
|
result := make([]JobTypeInfo, 0, len(jobTypeList))
|
|
workers := r.registry.List()
|
|
|
|
// Pre-calculate the best capability for each job type from available workers.
|
|
// Prefer capabilities with non-empty DisplayName, then higher Weight.
|
|
jobTypeToCap := make(map[string]*plugin_pb.JobTypeCapability)
|
|
for _, worker := range workers {
|
|
for jobType, cap := range worker.Capabilities {
|
|
if cap == nil {
|
|
continue
|
|
}
|
|
existing, exists := jobTypeToCap[jobType]
|
|
if !exists || existing == nil {
|
|
jobTypeToCap[jobType] = cap
|
|
continue
|
|
}
|
|
// Preserve existing if it has DisplayName but cap doesn't.
|
|
if existing.DisplayName != "" && cap.DisplayName == "" {
|
|
continue
|
|
}
|
|
// Prefer capabilities with a non-empty DisplayName.
|
|
if existing.DisplayName == "" && cap.DisplayName != "" {
|
|
jobTypeToCap[jobType] = cap
|
|
continue
|
|
}
|
|
// If DisplayName statuses are equal, prefer higher Weight.
|
|
if cap.Weight > existing.Weight {
|
|
jobTypeToCap[jobType] = cap
|
|
}
|
|
}
|
|
}
|
|
|
|
for _, jobType := range jobTypeList {
|
|
info := JobTypeInfo{JobType: jobType}
|
|
|
|
// Get display name and weight from pre-calculated capabilities
|
|
if cap, ok := jobTypeToCap[jobType]; ok && cap != nil {
|
|
if cap.DisplayName != "" {
|
|
info.DisplayName = cap.DisplayName
|
|
}
|
|
info.Weight = cap.Weight
|
|
}
|
|
|
|
// Default display name to job type if not set
|
|
if info.DisplayName == "" {
|
|
info.DisplayName = jobType
|
|
}
|
|
|
|
result = append(result, info)
|
|
}
|
|
|
|
// Sort by weight (descending) then by job type (ascending)
|
|
sort.Slice(result, func(i, j int) bool {
|
|
if result[i].Weight != result[j].Weight {
|
|
return result[i].Weight > result[j].Weight // higher weight first
|
|
}
|
|
return result[i].JobType < result[j].JobType // alphabetical as tiebreaker
|
|
})
|
|
|
|
return result, nil
|
|
}
|
|
|
|
// FilterProposalsWithActiveJobs drops proposals that are already assigned/running.
|
|
func (r *Plugin) FilterProposalsWithActiveJobs(jobType string, proposals []*plugin_pb.JobProposal) ([]*plugin_pb.JobProposal, int) {
|
|
return r.filterProposalsWithActiveJobs(jobType, proposals)
|
|
}
|
|
|
|
func (r *Plugin) PickDetectorWorker(jobType string) (*WorkerSession, error) {
|
|
return r.pickDetector(jobType, nil)
|
|
}
|
|
|
|
func (r *Plugin) PickExecutorWorker(jobType string) (*WorkerSession, error) {
|
|
return r.registry.PickExecutor(jobType)
|
|
}
|
|
|
|
// 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 && r.sessionIsOpen(leasedWorkerID, avoid) {
|
|
if capability := worker.Capabilities[jobType]; capability != nil && capability.CanDetect {
|
|
return worker, nil
|
|
}
|
|
}
|
|
r.clearDetectorLease(jobType, leasedWorkerID)
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
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 {
|
|
r.detectorLeaseMu.Lock()
|
|
defer r.detectorLeaseMu.Unlock()
|
|
return r.detectorLeases[jobType]
|
|
}
|
|
|
|
func (r *Plugin) setDetectorLease(jobType string, workerID string) {
|
|
r.detectorLeaseMu.Lock()
|
|
defer r.detectorLeaseMu.Unlock()
|
|
if jobType == "" || workerID == "" {
|
|
return
|
|
}
|
|
r.detectorLeases[jobType] = workerID
|
|
}
|
|
|
|
func (r *Plugin) clearDetectorLease(jobType string, workerID string) {
|
|
r.detectorLeaseMu.Lock()
|
|
defer r.detectorLeaseMu.Unlock()
|
|
|
|
current := r.detectorLeases[jobType]
|
|
if current == "" {
|
|
return
|
|
}
|
|
if workerID != "" && current != workerID {
|
|
return
|
|
}
|
|
delete(r.detectorLeases, jobType)
|
|
}
|
|
|
|
func (r *Plugin) sendCancel(workerID, targetID string, kind plugin_pb.WorkKind, cause error) {
|
|
if strings.TrimSpace(workerID) == "" || strings.TrimSpace(targetID) == "" {
|
|
return
|
|
}
|
|
|
|
requestID, err := newRequestID("cancel")
|
|
if err != nil {
|
|
requestID = ""
|
|
}
|
|
reason := "request canceled"
|
|
if cause != nil {
|
|
reason = cause.Error()
|
|
}
|
|
|
|
message := &plugin_pb.AdminToWorkerMessage{
|
|
RequestId: requestID,
|
|
SentAt: timestamppb.Now(),
|
|
Body: &plugin_pb.AdminToWorkerMessage_CancelRequest{
|
|
CancelRequest: &plugin_pb.CancelRequest{
|
|
TargetId: targetID,
|
|
TargetKind: kind,
|
|
Reason: reason,
|
|
},
|
|
},
|
|
}
|
|
if err := r.sendToWorker(workerID, message); err != nil {
|
|
glog.V(1).Infof("Plugin failed to send cancel request to worker=%s target=%s: %v", workerID, targetID, err)
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) sendAdminHello(session *streamSession) error {
|
|
msg := &plugin_pb.AdminToWorkerMessage{
|
|
RequestId: "",
|
|
SentAt: timestamppb.Now(),
|
|
Body: &plugin_pb.AdminToWorkerMessage_Hello{
|
|
Hello: &plugin_pb.AdminHello{
|
|
Accepted: true,
|
|
Message: "plugin connected",
|
|
HeartbeatIntervalSeconds: defaultHeartbeatInterval,
|
|
ReconnectDelaySeconds: defaultReconnectDelay,
|
|
},
|
|
},
|
|
}
|
|
return r.sendToSession(session, msg)
|
|
}
|
|
|
|
func (r *Plugin) sendLoop(
|
|
ctx context.Context,
|
|
stream plugin_pb.PluginControlService_WorkerStreamServer,
|
|
session *streamSession,
|
|
) error {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil
|
|
case <-r.shutdownCh:
|
|
return nil
|
|
case <-session.done:
|
|
return nil
|
|
case msg := <-session.outgoing:
|
|
if msg == nil {
|
|
return nil
|
|
}
|
|
if err := stream.Send(msg); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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", session.workerID)
|
|
case session.outgoing <- message:
|
|
return nil
|
|
case <-time.After(r.sendTimeout):
|
|
return fmt.Errorf("timed out sending message to worker %s", session.workerID)
|
|
}
|
|
}
|
|
|
|
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:
|
|
if !r.updateRegistryFromSession(session, func() { r.registry.UpsertFromHello(body.Hello) }) {
|
|
return
|
|
}
|
|
r.wakeScheduler()
|
|
case *plugin_pb.WorkerToAdminMessage_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:
|
|
r.handleDetectionProposals(workerID, body.DetectionProposals)
|
|
case *plugin_pb.WorkerToAdminMessage_DetectionComplete:
|
|
r.handleDetectionComplete(workerID, body.DetectionComplete)
|
|
case *plugin_pb.WorkerToAdminMessage_JobProgressUpdate:
|
|
r.handleJobProgressUpdate(workerID, body.JobProgressUpdate)
|
|
case *plugin_pb.WorkerToAdminMessage_JobCompleted:
|
|
r.handleJobCompleted(body.JobCompleted)
|
|
case *plugin_pb.WorkerToAdminMessage_Observations:
|
|
r.observations.Record(workerID, body.Observations)
|
|
case *plugin_pb.WorkerToAdminMessage_ObjectPreviewResponse:
|
|
r.handleObjectPreviewResponse(body.ObjectPreviewResponse)
|
|
case *plugin_pb.WorkerToAdminMessage_Acknowledge:
|
|
if !body.Acknowledge.Accepted {
|
|
glog.Warningf("Plugin worker %s rejected request %s: %s", workerID, body.Acknowledge.RequestId, body.Acknowledge.Message)
|
|
}
|
|
default:
|
|
// Keep the transport open even if admin does not yet consume all message variants.
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) handleConfigSchemaResponse(response *plugin_pb.ConfigSchemaResponse) {
|
|
if response == nil {
|
|
return
|
|
}
|
|
|
|
if response.Success && response.GetJobTypeDescriptor() != nil {
|
|
jobType := response.JobType
|
|
if jobType == "" {
|
|
jobType = response.GetJobTypeDescriptor().JobType
|
|
}
|
|
if jobType != "" {
|
|
if err := r.store.SaveDescriptor(jobType, response.GetJobTypeDescriptor()); err != nil {
|
|
glog.Warningf("Plugin failed to persist descriptor for %s: %v", jobType, err)
|
|
}
|
|
if err := r.ensureJobTypeConfigFromDescriptor(jobType, response.GetJobTypeDescriptor()); err != nil {
|
|
glog.Warningf("Plugin failed to bootstrap config for %s: %v", jobType, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
r.safeSendSchemaResponse(response.RequestId, response)
|
|
}
|
|
|
|
func (r *Plugin) safeSendSchemaResponse(requestID string, response *plugin_pb.ConfigSchemaResponse) {
|
|
r.pendingSchemaMu.Lock()
|
|
ch := r.pendingSchema[requestID]
|
|
r.pendingSchemaMu.Unlock()
|
|
safeSendCh(ch, response, r.shutdownCh)
|
|
}
|
|
|
|
func safeSendCh[T any](ch chan T, val T, shutdownCh <-chan struct{}) {
|
|
if ch == nil {
|
|
return
|
|
}
|
|
defer func() {
|
|
recover()
|
|
}()
|
|
select {
|
|
case ch <- val:
|
|
case <-shutdownCh:
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) ensureJobTypeConfigFromDescriptor(jobType string, descriptor *plugin_pb.JobTypeDescriptor) error {
|
|
if descriptor == nil || strings.TrimSpace(jobType) == "" {
|
|
return nil
|
|
}
|
|
|
|
workerDefaults := CloneConfigValueMap(descriptor.WorkerDefaultValues)
|
|
if len(workerDefaults) == 0 && descriptor.WorkerConfigForm != nil {
|
|
workerDefaults = CloneConfigValueMap(descriptor.WorkerConfigForm.DefaultValues)
|
|
}
|
|
|
|
adminDefaults := map[string]*plugin_pb.ConfigValue{}
|
|
if descriptor.AdminConfigForm != nil {
|
|
adminDefaults = CloneConfigValueMap(descriptor.AdminConfigForm.DefaultValues)
|
|
}
|
|
|
|
adminRuntime := &plugin_pb.AdminRuntimeConfig{}
|
|
if descriptor.AdminRuntimeDefaults != nil {
|
|
defaults := descriptor.AdminRuntimeDefaults
|
|
adminRuntime = &plugin_pb.AdminRuntimeConfig{
|
|
Enabled: defaults.Enabled,
|
|
DetectionIntervalMinutes: defaults.DetectionIntervalMinutes,
|
|
DetectionTimeoutSeconds: defaults.DetectionTimeoutSeconds,
|
|
MaxJobsPerDetection: defaults.MaxJobsPerDetection,
|
|
GlobalExecutionConcurrency: defaults.GlobalExecutionConcurrency,
|
|
PerWorkerExecutionConcurrency: defaults.PerWorkerExecutionConcurrency,
|
|
RetryLimit: defaults.RetryLimit,
|
|
RetryBackoffSeconds: defaults.RetryBackoffSeconds,
|
|
JobTypeMaxRuntimeSeconds: defaults.JobTypeMaxRuntimeSeconds,
|
|
ExecutionTimeoutSeconds: defaults.ExecutionTimeoutSeconds,
|
|
}
|
|
}
|
|
|
|
cfg := &plugin_pb.PersistedJobTypeConfig{
|
|
JobType: jobType,
|
|
DescriptorVersion: descriptor.DescriptorVersion,
|
|
AdminConfigValues: adminDefaults,
|
|
WorkerConfigValues: workerDefaults,
|
|
AdminRuntime: adminRuntime,
|
|
UpdatedAt: timestamppb.Now(),
|
|
UpdatedBy: "plugin",
|
|
}
|
|
|
|
// Check existence first to avoid calling configDefaultsProvider unnecessarily
|
|
// (e.g., it may make a blocking gRPC call to fetch master config).
|
|
existing, err := r.store.LoadJobTypeConfig(jobType)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if existing != nil {
|
|
return nil
|
|
}
|
|
|
|
if r.configDefaultsProvider != nil {
|
|
cfg = r.configDefaultsProvider(cfg)
|
|
}
|
|
|
|
_, err = r.store.SaveJobTypeConfigIfNotExists(cfg)
|
|
return err
|
|
}
|
|
|
|
func (r *Plugin) handleDetectionProposals(workerID string, message *plugin_pb.DetectionProposals) {
|
|
if message == nil || message.RequestId == "" {
|
|
return
|
|
}
|
|
|
|
r.pendingDetectionMu.Lock()
|
|
state := r.pendingDetection[message.RequestId]
|
|
if state != nil {
|
|
state.proposals = append(state.proposals, cloneJobProposals(message.Proposals)...)
|
|
}
|
|
r.pendingDetectionMu.Unlock()
|
|
if state == nil {
|
|
return
|
|
}
|
|
|
|
resolvedWorkerID := strings.TrimSpace(workerID)
|
|
if resolvedWorkerID == "" {
|
|
resolvedWorkerID = state.workerID
|
|
}
|
|
resolvedJobType := strings.TrimSpace(message.JobType)
|
|
if resolvedJobType == "" {
|
|
resolvedJobType = state.jobType
|
|
}
|
|
if resolvedJobType == "" {
|
|
resolvedJobType = "unknown"
|
|
}
|
|
|
|
r.appendActivity(JobActivity{
|
|
JobType: resolvedJobType,
|
|
RequestID: message.RequestId,
|
|
WorkerID: resolvedWorkerID,
|
|
Source: "detector",
|
|
Stage: "proposals_batch",
|
|
Message: fmt.Sprintf("received %d proposal(s)", len(message.Proposals)),
|
|
OccurredAt: timeToPtr(time.Now().UTC()),
|
|
Details: map[string]interface{}{
|
|
"batch_size": len(message.Proposals),
|
|
"has_more": message.HasMore,
|
|
},
|
|
})
|
|
|
|
for _, proposal := range message.Proposals {
|
|
if proposal == nil {
|
|
continue
|
|
}
|
|
details := map[string]interface{}{
|
|
"proposal_id": proposal.ProposalId,
|
|
"dedupe_key": proposal.DedupeKey,
|
|
"priority": proposal.Priority.String(),
|
|
"summary": proposal.Summary,
|
|
"detail": proposal.Detail,
|
|
"labels": proposal.Labels,
|
|
}
|
|
if params := configValueMapToPlain(proposal.Parameters); len(params) > 0 {
|
|
details["parameters"] = params
|
|
}
|
|
|
|
messageText := strings.TrimSpace(proposal.Summary)
|
|
if messageText == "" {
|
|
messageText = fmt.Sprintf("proposal %s", strings.TrimSpace(proposal.ProposalId))
|
|
}
|
|
if messageText == "" {
|
|
messageText = "proposal generated"
|
|
}
|
|
|
|
r.appendActivity(JobActivity{
|
|
JobType: resolvedJobType,
|
|
RequestID: message.RequestId,
|
|
WorkerID: resolvedWorkerID,
|
|
Source: "detector",
|
|
Stage: "proposal",
|
|
Message: messageText,
|
|
OccurredAt: timeToPtr(time.Now().UTC()),
|
|
Details: details,
|
|
})
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) handleDetectionComplete(workerID string, message *plugin_pb.DetectionComplete) {
|
|
if message == nil {
|
|
return
|
|
}
|
|
if !message.Success {
|
|
glog.Warningf("Plugin detection failed job_type=%s: %s", message.JobType, message.ErrorMessage)
|
|
}
|
|
if message.RequestId == "" {
|
|
return
|
|
}
|
|
|
|
r.pendingDetectionMu.Lock()
|
|
state := r.pendingDetection[message.RequestId]
|
|
r.pendingDetectionMu.Unlock()
|
|
if state == nil {
|
|
return
|
|
}
|
|
|
|
resolvedWorkerID := strings.TrimSpace(workerID)
|
|
if resolvedWorkerID == "" {
|
|
resolvedWorkerID = state.workerID
|
|
}
|
|
resolvedJobType := strings.TrimSpace(message.JobType)
|
|
if resolvedJobType == "" {
|
|
resolvedJobType = state.jobType
|
|
}
|
|
if resolvedJobType == "" {
|
|
resolvedJobType = "unknown"
|
|
}
|
|
|
|
stage := "completed"
|
|
messageText := "detection completed"
|
|
if !message.Success {
|
|
stage = "failed"
|
|
messageText = strings.TrimSpace(message.ErrorMessage)
|
|
if messageText == "" {
|
|
messageText = "detection failed"
|
|
}
|
|
} else if message.SkippedReason != "" {
|
|
stage = "skipped"
|
|
messageText = message.SkippedReason
|
|
}
|
|
details := map[string]interface{}{
|
|
"success": message.Success,
|
|
"total_proposals": message.TotalProposals,
|
|
}
|
|
if message.SkippedReason != "" {
|
|
details["skipped_reason"] = message.SkippedReason
|
|
}
|
|
if message.EligibleAfter != nil {
|
|
details["eligible_after"] = message.EligibleAfter.AsTime().UTC().Format(time.RFC3339)
|
|
}
|
|
r.appendActivity(JobActivity{
|
|
JobType: resolvedJobType,
|
|
RequestID: message.RequestId,
|
|
WorkerID: resolvedWorkerID,
|
|
Source: "detector",
|
|
Stage: stage,
|
|
Message: messageText,
|
|
OccurredAt: timeToPtr(time.Now().UTC()),
|
|
Details: details,
|
|
})
|
|
|
|
r.safeSendDetectionComplete(message.RequestId, message)
|
|
}
|
|
|
|
func (r *Plugin) safeSendDetectionComplete(requestID string, message *plugin_pb.DetectionComplete) {
|
|
r.pendingDetectionMu.Lock()
|
|
state, found := r.pendingDetection[requestID]
|
|
var ch chan *plugin_pb.DetectionComplete
|
|
if found && state != nil {
|
|
ch = state.complete
|
|
}
|
|
r.pendingDetectionMu.Unlock()
|
|
|
|
if ch != nil {
|
|
safeSendCh(ch, message, r.shutdownCh)
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) handleJobCompleted(completed *plugin_pb.JobCompleted) {
|
|
if completed == nil || completed.JobType == "" {
|
|
return
|
|
}
|
|
|
|
if completed.RequestId != "" {
|
|
r.safeSendJobCompleted(completed.RequestId, completed)
|
|
}
|
|
|
|
tracked := r.trackExecutionCompletion(completed)
|
|
workerID := ""
|
|
if tracked != nil && tracked.WorkerID != "" {
|
|
workerID = tracked.WorkerID
|
|
}
|
|
|
|
r.trackWorkerActivities(completed.JobType, completed.JobId, completed.RequestId, workerID, completed.Activities)
|
|
|
|
record := &JobRunRecord{
|
|
RunID: completed.RequestId,
|
|
JobID: completed.JobId,
|
|
JobType: completed.JobType,
|
|
WorkerID: "",
|
|
Outcome: RunOutcomeError,
|
|
CompletedAt: timeToPtr(time.Now().UTC()),
|
|
}
|
|
if completed.CompletedAt != nil {
|
|
record.CompletedAt = timeToPtr(completed.CompletedAt.AsTime().UTC())
|
|
}
|
|
if completed.Success {
|
|
record.Outcome = RunOutcomeSuccess
|
|
record.Message = "completed"
|
|
if completed.Result != nil && completed.Result.Summary != "" {
|
|
record.Message = completed.Result.Summary
|
|
}
|
|
} else {
|
|
record.Outcome = RunOutcomeError
|
|
record.Message = completed.ErrorMessage
|
|
}
|
|
|
|
if tracked != nil {
|
|
if workerID != "" {
|
|
record.WorkerID = workerID
|
|
}
|
|
if tracked.CreatedAt != nil && record.CompletedAt != nil && record.CompletedAt.After(*tracked.CreatedAt) {
|
|
record.DurationMs = int64(record.CompletedAt.Sub(*tracked.CreatedAt) / time.Millisecond)
|
|
}
|
|
}
|
|
|
|
if err := r.store.AppendRunRecord(completed.JobType, record); err != nil {
|
|
glog.Warningf("Plugin failed to append run record for %s: %v", completed.JobType, err)
|
|
}
|
|
}
|
|
|
|
func (r *Plugin) safeSendJobCompleted(requestID string, completed *plugin_pb.JobCompleted) {
|
|
r.pendingExecutionMu.Lock()
|
|
ch := r.pendingExecution[requestID]
|
|
r.pendingExecutionMu.Unlock()
|
|
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()
|
|
r.publishSessionLocked(session)
|
|
}
|
|
|
|
// 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()
|
|
current, exists := r.sessions[workerID]
|
|
superseded := exists && current != session
|
|
if exists && !superseded {
|
|
delete(r.sessions, workerID)
|
|
}
|
|
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)
|
|
}
|
|
|
|
func newRequestID(prefix string) (string, error) {
|
|
buf := make([]byte, 8)
|
|
if _, err := rand.Read(buf); err != nil {
|
|
return "", fmt.Errorf("generate request id: %w", err)
|
|
}
|
|
if prefix == "" {
|
|
prefix = "req"
|
|
}
|
|
return fmt.Sprintf("%s-%d-%s", prefix, time.Now().UnixNano(), hex.EncodeToString(buf)), nil
|
|
}
|
|
|
|
func (r *Plugin) loadJobTypeConfigPayload(jobType string) (
|
|
*plugin_pb.AdminRuntimeConfig,
|
|
map[string]*plugin_pb.ConfigValue,
|
|
map[string]*plugin_pb.ConfigValue,
|
|
error,
|
|
) {
|
|
config, err := r.store.LoadJobTypeConfig(jobType)
|
|
if err != nil {
|
|
return nil, nil, nil, err
|
|
}
|
|
|
|
if config == nil {
|
|
return &plugin_pb.AdminRuntimeConfig{}, map[string]*plugin_pb.ConfigValue{}, map[string]*plugin_pb.ConfigValue{}, nil
|
|
}
|
|
|
|
adminRuntime := config.AdminRuntime
|
|
if adminRuntime == nil {
|
|
adminRuntime = &plugin_pb.AdminRuntimeConfig{}
|
|
}
|
|
return adminRuntime, CloneConfigValueMap(config.AdminConfigValues), CloneConfigValueMap(config.WorkerConfigValues), nil
|
|
}
|
|
|
|
func cloneJobProposals(in []*plugin_pb.JobProposal) []*plugin_pb.JobProposal {
|
|
if len(in) == 0 {
|
|
return nil
|
|
}
|
|
out := make([]*plugin_pb.JobProposal, 0, len(in))
|
|
for _, proposal := range in {
|
|
if proposal == nil {
|
|
continue
|
|
}
|
|
out = append(out, proto.Clone(proposal).(*plugin_pb.JobProposal))
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (r *Plugin) loadLastSuccessfulRun(jobType string) *timestamppb.Timestamp {
|
|
history, err := r.store.LoadRunHistory(jobType)
|
|
if err != nil {
|
|
glog.Warningf("Plugin failed to load run history for %s: %v", jobType, err)
|
|
return nil
|
|
}
|
|
if history == nil || len(history.SuccessfulRuns) == 0 {
|
|
return nil
|
|
}
|
|
|
|
var latest time.Time
|
|
for i := range history.SuccessfulRuns {
|
|
completedAt := history.SuccessfulRuns[i].CompletedAt
|
|
if completedAt == nil || completedAt.IsZero() {
|
|
continue
|
|
}
|
|
if latest.IsZero() || completedAt.After(latest) {
|
|
latest = *completedAt
|
|
}
|
|
}
|
|
if latest.IsZero() {
|
|
return nil
|
|
}
|
|
return timestamppb.New(latest.UTC())
|
|
}
|
|
|
|
func (r *Plugin) loadLastCompletedRun(jobType string) *timestamppb.Timestamp {
|
|
history, err := r.store.LoadRunHistory(jobType)
|
|
if err != nil {
|
|
glog.Warningf("Plugin failed to load run history for %s: %v", jobType, err)
|
|
return nil
|
|
}
|
|
if history == nil {
|
|
return nil
|
|
}
|
|
|
|
var latest time.Time
|
|
for i := range history.SuccessfulRuns {
|
|
completedAt := history.SuccessfulRuns[i].CompletedAt
|
|
if completedAt == nil || completedAt.IsZero() {
|
|
continue
|
|
}
|
|
if latest.IsZero() || completedAt.After(latest) {
|
|
latest = *completedAt
|
|
}
|
|
}
|
|
for i := range history.ErrorRuns {
|
|
completedAt := history.ErrorRuns[i].CompletedAt
|
|
if completedAt == nil || completedAt.IsZero() {
|
|
continue
|
|
}
|
|
if latest.IsZero() || completedAt.After(latest) {
|
|
latest = *completedAt
|
|
}
|
|
}
|
|
if latest.IsZero() {
|
|
return nil
|
|
}
|
|
return timestamppb.New(latest.UTC())
|
|
}
|
|
|
|
func CloneConfigValueMap(in map[string]*plugin_pb.ConfigValue) map[string]*plugin_pb.ConfigValue {
|
|
if len(in) == 0 {
|
|
return map[string]*plugin_pb.ConfigValue{}
|
|
}
|
|
out := make(map[string]*plugin_pb.ConfigValue, len(in))
|
|
for key, value := range in {
|
|
if value == nil {
|
|
continue
|
|
}
|
|
out[key] = proto.Clone(value).(*plugin_pb.ConfigValue)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (s *streamSession) close() {
|
|
s.closeOnce.Do(func() {
|
|
close(s.done)
|
|
})
|
|
}
|
|
|
|
// WorkerConnectForTest simulates a worker connecting (test helper).
|
|
func (r *Plugin) WorkerConnectForTest(hello *plugin_pb.WorkerHello) {
|
|
r.registry.UpsertFromHello(hello)
|
|
}
|
|
|
|
// 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)
|
|
}
|