package distworker import ( "context" "crypto/ed25519" "encoding/base64" "encoding/json" "fmt" "log" "os" "path/filepath" "sync" "sync/atomic" "time" "github.com/Jeffail/tunny" "github.com/gorilla/websocket" "gorm.io/datatypes" "rocketgit.ru/rsmon/worker/app/models" "rocketgit.ru/rsmon/worker/internal/checkexec" "rocketgit.ru/rsmon/worker/internal/wire" ) const ( heartbeatInterval = 10 * time.Second // minQueueCapacity is the lower bound for the bounded job/result // channels so that a small pool still has some backpressure headroom. minQueueCapacity = 16 malformedTaskEnvelopeError = "malformed_task_envelope" finalDrainLimit = 32 finalDrainTimeout = 5 * time.Second ) // jobPool is the minimal interface the Runner needs from a worker pool. It // exists so tests can observe SetSize calls without wrapping the tunny // concrete type. type jobPool interface { Process(payload interface{}) interface{} SetSize(n int) Close() } // resultEnvelope pairs a job with its reports so the writer can include // context (e.g., the check kind) when logging. type resultEnvelope struct { job wire.CheckJob reports []wire.CheckResultReport } type metricEnvelope struct { generation uint64 report wire.ServerMetricReport } // Runner manages the worker execution loop. // // Concurrency design: // // - jobQueue and results are bounded channels sized relative to // maxConcurrency. The websocket read loop enqueues into jobQueue, // providing backpressure to the control plane when the worker is // saturated. // - A fixed number of dispatcher goroutines (maxConcurrency) drain the // jobQueue and call pool.Process. The number of in-flight executions // is bounded by the tunny.Pool size, which the control plane // configures via the "init" / "config" message (clamped to // maxConcurrency). // - Each websocket connection owns a single writer goroutine that // serializes result and heartbeat writes through a mutex. This keeps // websocket writes thread-safe and removes the per-task goroutine // that previously blocked pool execution while holding the mutex. type Runner struct { config *Config client *Client pool jobPool concurrency int64 maxConcurrency int jobQueue chan wire.CheckJob results chan resultEnvelope notifyQueue chan wire.NotificationTask notifyResults chan notifyResultEnvelope metricResults chan metricEnvelope stopCh chan struct{} wg sync.WaitGroup controlWG sync.WaitGroup backgroundWG sync.WaitGroup lifecycleMu sync.Mutex controlCtx context.Context controlCancel context.CancelFunc backgroundCtx context.Context backgroundCancel context.CancelFunc rotationCtx context.Context rotationCancel context.CancelFunc started atomic.Bool queueDepth int64 activeCount int64 notifyDepth int64 notifyActive int64 // credentialsMu guards credentials during init/config refresh. credentialsMu sync.RWMutex credentials *wire.NotificationCredentials // systemContactsMu guards system contacts pushed via init/config. // Workers notify these contacts directly when the main API is // unreachable (see docs/distributed/notifications-from-worker.md // "System Selfcheck"). systemContactsMu sync.RWMutex systemContacts []wire.SystemContact // urlMu guards the worker URL pushed via init/config refresh. The // URL is what other workers and the main app dial to reach this // worker (it can differ from the bind host:port because of reverse // proxies / Traefik with HTTPS). See docs/worker-http-settings.md. urlMu sync.RWMutex url string // peersMu guards the peer list pushed via init/config. The // control plane builds the list from worker_nodes (excluding // this worker) and refreshes it on every 5m config push; see // docs/distributed/worker-to-worker-raft.md §10 for the // intended use. The selfcheck module reads the list to drive // the peer poller and the consensus. peersMu sync.RWMutex peers []wire.PeerInfo // peerCache stores the latest observation per peer. Written by // the peer poller (every peerPollInterval) and read by the // selfcheck consensus helper. Constructed in NewRunner so tests // can drive it without Start(). peerCache *peerCache // masterStatusMu guards the local "master API up/down" snapshot // that the selfcheck writes after each probe and the // /api/peer/status handler reads. observedAt is the wall-clock // time of the most recent local probe; up is the verdict of // that probe. The zero time means "no probe has run yet". masterStatusMu sync.RWMutex masterStatusUp *bool masterStatusAt time.Time // executor is the function the pool runs for each job. It is a // field so tests can swap it for a deterministic stub without // touching the websocket plumbing. executor func(payload interface{}) interface{} // notificationExecutor is a test seam for the notification task // boundary. Production uses ExecuteNotification's normal delivery path. notificationExecutor func(context.Context, models.Task) wire.NotificationResultReport // resultsBuf holds the most recent result rows. The webapp // reads from it via RecentResults(n). Cleared by Stop so the // ring does not leak between worker runs. resultsBuf *resultBuffer // notificationsBuf mirrors resultsBuf for emitted // notifications. Phase 1 only fills this from selfcheck alerts. notificationsBuf *notificationBuffer // lastHeartbeatAt tracks the most recent successful heartbeat // write, so the webapp can render "last ack" without polling. lastHeartbeatMu sync.RWMutex lastHeartbeatAt time.Time // tokenRotatedAt records the wall-clock time of the last token // rotation, so the settings page can render "last rotated". tokenRotatedMu sync.RWMutex tokenRotatedAt time.Time // workerID / regionCode / workerVersion / workerCaps are // captured from the most recent init/config websocket message // so the webapp can render them read-only. workerIDMu sync.RWMutex workerID string regionMu sync.RWMutex regionCode string versionMu sync.RWMutex workerVersion string capsMu sync.RWMutex workerCaps []string identityMu sync.RWMutex state workerState serverID atomic.Int64 metricGeneration atomic.Uint64 nextMetricGeneration atomic.Uint64 // clientMu guards the control-plane client and its active websocket. // Rotation swaps the client and closes only controlConn, leaving the // runner's global stop signal and local subsystems untouched. clientMu sync.Mutex controlConn *websocket.Conn controlWriteMu *sync.Mutex reconnectCh chan struct{} rotationMu sync.Mutex leasesMu sync.Mutex leases map[string]string // outbox holds messages removed from a per-connection writer queue that // could not be written before its websocket closed. outboxMu sync.Mutex outbox []wire.WorkerMessage outboxWake chan struct{} // beforeControlWrite is a test seam used to hold a dequeued message while // a connection rotates. beforeControlWrite func() // beforeTokenCommit lets tests make shutdown win between a successful HTTP // rotation response and the lifecycle-protected in-memory commit. beforeTokenCommit func() } // NewRunner creates a new worker runner. Config is taken by pointer to // keep the parameter cheap as the struct grows (it now carries the // HTTP listener settings on top of the control-plane connection // fields). func NewRunner(cfg *Config) *Runner { maxConc := cfg.MaxConcurrency if maxConc <= 0 { maxConc = DefaultMaxConcurrency } controlCtx, controlCancel := context.WithCancel(context.Background()) backgroundCtx, backgroundCancel := context.WithCancel(context.Background()) rotationCtx, rotationCancel := context.WithCancel(context.Background()) return &Runner{ config: cfg, maxConcurrency: maxConc, stopCh: make(chan struct{}), reconnectCh: make(chan struct{}, 1), outboxWake: make(chan struct{}, 1), controlCtx: controlCtx, controlCancel: controlCancel, backgroundCtx: backgroundCtx, backgroundCancel: backgroundCancel, rotationCtx: rotationCtx, rotationCancel: rotationCancel, resultsBuf: newResultBuffer(), notificationsBuf: newNotificationBuffer(), peerCache: newPeerCache(), leases: make(map[string]string), } } // Start begins the worker execution func (r *Runner) Start() error { log.Println("worker: starting...") r.lifecycleMu.Lock() if r.config.URL == "" || (r.config.Token == "" && r.config.BootstrapToken == "") { r.lifecycleMu.Unlock() return fmt.Errorf("RSMON_URL and RSMON_TOKEN must be set") } if !r.started.CompareAndSwap(false, true) { r.lifecycleMu.Unlock() return fmt.Errorf("worker: runner already started") } if r.stopped() { r.lifecycleMu.Unlock() return fmt.Errorf("worker: runner stopped") } if r.config.BootstrapToken != "" && (r.config.StateFile == "" || !filepath.IsAbs(r.config.StateFile)) { r.lifecycleMu.Unlock() return fmt.Errorf("bootstrap requires an absolute RSMON_STATE_FILE") } state, err := loadWorkerState(r.config.StateFile) if err != nil { r.lifecycleMu.Unlock() return fmt.Errorf("worker state: %w", err) } if state.Token != "" { r.config.Token = state.Token } if r.config.Token == "" { client := NewClient(r.config.URL, "") boot, err := client.Bootstrap(context.Background(), r.config.WorkerID, r.config.BootstrapToken) if err != nil { r.lifecycleMu.Unlock() return fmt.Errorf("worker bootstrap: %w", err) } state = workerState{Token: boot.AuthToken, WorkerID: boot.WorkerID, VerificationKey: boot.ConfigVerificationKey, SigningKeyID: boot.SigningKeyID} if err := saveWorkerState(r.config.StateFile, state); err != nil { r.lifecycleMu.Unlock() return fmt.Errorf("worker state: %w", err) } r.config.Token = boot.AuthToken r.config.BootstrapToken = "" } r.state = state r.clientMu.Lock() r.client = NewClient(r.config.URL, r.config.Token) r.clientMu.Unlock() queueCap := r.queueCapacity() r.jobQueue = make(chan wire.CheckJob, queueCap) r.results = make(chan resultEnvelope, queueCap) r.notifyQueue = make(chan wire.NotificationTask, queueCap) r.notifyResults = make(chan notifyResultEnvelope, queueCap) r.metricResults = make(chan metricEnvelope, queueCap) if r.executor == nil { r.executor = r.defaultExecuteJob } atomic.StoreInt64(&r.concurrency, 1) r.pool = tunny.NewFunc(int(atomic.LoadInt64(&r.concurrency)), r.executor) // Start a fixed pool of dispatchers. The tunny.Pool size ultimately // limits the number of concurrent job executions; extra dispatchers // simply wait inside pool.Process when the pool is saturated. for i := 0; i < r.maxConcurrency; i++ { r.wg.Add(1) go r.dispatcher() } // Notification dispatchers share the pool's overall concurrency budget: // we run maxConcurrency notification dispatchers and let the runner's // NotificationMethods capability gate keep them quiet when the worker // is not authorized for the method. for i := 0; i < r.maxConcurrency; i++ { r.wg.Add(1) go r.notifyDispatcher() } // Start websocket task loop. r.controlWG.Add(1) go func() { defer r.controlWG.Done() r.websocketLoop() }() r.backgroundWG.Add(1) go func() { defer r.backgroundWG.Done() r.serverMetricLoop(r.backgroundCtx) }() // Start periodic selfcheck loop. This probes the main API and, on // sustained unreachability, notifies system contacts directly via // the cached credentials. Lifetimes of selfcheck goroutines are // bound to stopCh (and the explicit cancel, kept for symmetry). r.backgroundWG.Add(1) go func() { defer r.backgroundWG.Done() r.startSelfcheck(r.backgroundCtx) }() // Start the peer poller. It refreshes r.peerCache with each // peer's latest /api/peer/status verdict. The selfcheck // consumes the cache to drive consensus. The poller is bound // to selfcheckCtx so it shuts down together with the // selfcheck loop on Stop. r.backgroundWG.Add(1) go func() { defer r.backgroundWG.Done() r.peerPollerLoop(r.backgroundCtx) }() r.lifecycleMu.Unlock() // Wait for stop signal <-r.stopCh log.Println("worker: shutting down...") if r.backgroundCancel != nil { r.backgroundCancel() } r.controlWG.Wait() r.backgroundWG.Wait() // Dispatchers exit via stopCh. Do not close jobQueue here: the websocket // reader can still be unwinding and may otherwise race with a send. r.wg.Wait() // Pool is safe to close only after all dispatchers have returned, // otherwise an in-flight pool.Process would panic. r.pool.Close() return nil } // Stop gracefully stops the worker func (r *Runner) Stop() { r.lifecycleMu.Lock() r.drainFinalResults() select { case <-r.stopCh: // already closed default: close(r.stopCh) } if r.controlCancel != nil { r.controlCancel() } if r.backgroundCancel != nil { r.backgroundCancel() } if r.rotationCancel != nil { r.rotationCancel() } r.clientMu.Lock() conn := r.controlConn r.clientMu.Unlock() if conn != nil { _ = conn.Close() } r.lifecycleMu.Unlock() // Cancel and close before waiting for rotation. A rotation can be waiting // on controlWriteMu while a blocked writer needs that close to return. r.rotationMu.Lock() r.rotationMu.Unlock() r.controlWG.Wait() r.backgroundWG.Wait() r.wg.Wait() } // Enqueue submits a job to the worker pool. It returns false if the runner // has been stopped. The call blocks while the bounded jobQueue is full, // providing natural backpressure to the caller. func (r *Runner) Enqueue(job wire.CheckJob) bool { //nolint:lll,gocritic // wire.CheckJob is ~104B; keep by-value to avoid forcing callers to take an address if r.jobQueue == nil { return false } if r.stopped() { return false } select { case r.jobQueue <- job: atomic.AddInt64(&r.queueDepth, 1) return true case <-r.stopCh: return false } } // Concurrency returns the current tunny.Pool size. func (r *Runner) Concurrency() int { return int(atomic.LoadInt64(&r.concurrency)) } // EnqueueNotification submits a notification task to the worker pool. It // returns false if the runner has been stopped. The call blocks while the // bounded notifyQueue is full, providing backpressure to the control plane. // //nolint:gocritic // wire payload is shared with the dispatcher; keep by-value func (r *Runner) EnqueueNotification(task wire.NotificationTask) bool { if r.notifyQueue == nil { return false } if r.stopped() { return false } select { case r.notifyQueue <- task: atomic.AddInt64(&r.notifyDepth, 1) return true case <-r.stopCh: return false } } // notifyDispatcher is the per-notification-task execution loop. It mirrors // dispatcher() but routes to ExecuteNotification + the notification_result // writer rather than checkexec + check_result_report. The same // maxConcurrency budget caps total in-flight work across both kinds. func (r *Runner) notifyDispatcher() { defer r.wg.Done() for { select { case task := <-r.notifyQueue: atomic.AddInt64(&r.notifyDepth, -1) atomic.AddInt64(&r.notifyActive, 1) r.executeAndForwardNotification(task) atomic.AddInt64(&r.notifyActive, -1) case <-r.stopCh: return } } } // executeAndForwardNotification converts the wire NotificationTask to a Task // shape, runs the executor, and pushes the result into notifyResults. The // writer goroutine picks it up and serializes the websocket write. func (r *Runner) executeAndForwardNotification(task wire.NotificationTask) { //nolint:gocritic // wire payload is shared started := time.Now() deadline := started.Add(models.DefaultNotificationExecutionTimeout) var taskDeadline *time.Time if task.Deadline != nil { parsedDeadline, err := time.Parse(time.RFC3339Nano, *task.Deadline) if err != nil { r.completeNotification(task, notificationPermanentReport(task, "invalid notification deadline", started)) return } taskDeadline = &parsedDeadline if !parsedDeadline.After(started) { r.completeNotification(task, notificationPermanentReport(task, "notification deadline expired", started)) return } if parsedDeadline.Before(deadline) { deadline = parsedDeadline } } ctx, cancel := context.WithDeadline(context.Background(), deadline) defer cancel() payload, _ := json.Marshal(task) dbTask := models.Task{JobID: task.JobID, LeaseToken: task.LeaseToken, Payload: payload} if taskDeadline != nil { dbTask.Deadline = taskDeadline } if task.MessageID != 0 { msgID := task.MessageID dbTask.MessageID = &msgID } if len(task.EventIDs) > 0 { ev := task.EventIDs[0] _ = ev } report := r.ExecuteNotification(ctx, dbTask) r.forwardNotificationResult(task, report) } func notificationPermanentReport(task wire.NotificationTask, message string, started time.Time) wire.NotificationResultReport { return wire.NotificationResultReport{ JobID: task.JobID, LeaseToken: task.LeaseToken, MessageID: task.MessageID, Status: wire.NotificationResultPermanent, Error: stringPtr(message), DurationMs: int(time.Since(started) / time.Millisecond), } } // completeNotification records the terminal local outcome before forwarding // the result. Rejected tasks use this path because they do not enter the executor. func (r *Runner) completeNotification(task wire.NotificationTask, report wire.NotificationResultReport) { r.recordDelegatedNotification(report.JobID, task.Method, report) r.forwardNotificationResult(task, report) } func (r *Runner) forwardNotificationResult(task wire.NotificationTask, report wire.NotificationResultReport) { if r.notifyResults == nil || r.stopped() { return } select { case r.notifyResults <- notifyResultEnvelope{task: task, report: report}: case <-r.stopCh: } } // MaxConcurrency returns the upper bound for pool size and dispatchers. func (r *Runner) MaxConcurrency() int { return r.maxConcurrency } // QueueCapacity returns the bounded buffer size used for jobQueue and // results. Exposed for tests and observability. func (r *Runner) QueueCapacity() int { return r.queueCapacity() } // ActiveCount returns only running check executions. QueueDepth reports the // disjoint pending-check count used with it in heartbeat capacity accounting. func (r *Runner) ActiveCount() int { return int(atomic.LoadInt64(&r.activeCount)) } // QueueDepth returns the current number of pending jobs in jobQueue. func (r *Runner) QueueDepth() int { return int(atomic.LoadInt64(&r.queueDepth)) } // ActiveNotifications returns the number of queued and in-flight deliveries. func (r *Runner) ActiveNotifications() int { return int(atomic.LoadInt64(&r.notifyDepth) + atomic.LoadInt64(&r.notifyActive)) } // NotificationQueueDepth returns deliveries waiting for a notification worker. func (r *Runner) NotificationQueueDepth() int { return int(atomic.LoadInt64(&r.notifyDepth)) } // queueCapacity returns the bounded buffer size for the job/result // channels. The capacity is derived from maxConcurrency so that // backpressure scales with the configured pool. func (r *Runner) queueCapacity() int { qcap := 2 * r.maxConcurrency if qcap < minQueueCapacity { qcap = minQueueCapacity } return qcap } func (r *Runner) dispatcher() { defer r.wg.Done() for { select { case job := <-r.jobQueue: atomic.AddInt64(&r.queueDepth, -1) atomic.AddInt64(&r.activeCount, 1) r.executeAndForward(job) atomic.AddInt64(&r.activeCount, -1) case <-r.stopCh: return } } } func (r *Runner) executeAndForward(job wire.CheckJob) { //nolint:lll,gocritic // see Enqueue; CheckJob is forwarded into tunny.Pool as interface{} result := r.pool.Process(job) reports, ok := result.([]wire.CheckResultReport) if !ok || len(reports) == 0 { return } env := resultEnvelope{job: job, reports: reports} // Prefer an early exit when the runner is stopping so we never // block on a full results channel. select { case <-r.stopCh: return default: } select { case r.results <- env: case <-r.stopCh: } } func (r *Runner) websocketLoop() { for { select { case <-r.stopCh: return default: } if err := r.runWebsocket(); err != nil { log.Println("worker: websocket error:", err) } select { case <-time.After(3 * time.Second): case <-r.reconnectCh: case <-r.stopCh: return } } } func (r *Runner) runWebsocket() error { var writeMu sync.Mutex r.clientMu.Lock() client := r.client r.clientMu.Unlock() if client == nil { return fmt.Errorf("worker: client not initialized") } conn, err := client.WorkerSocketContext(r.controlCtx) if err != nil { return err } // A rotation can complete while the websocket dial is in flight. Do not // install a connection authenticated with the superseded token. r.clientMu.Lock() if r.client != client || r.stopped() { r.clientMu.Unlock() _ = conn.Close() return nil } r.controlConn = conn r.controlWriteMu = &writeMu metricGeneration := r.nextMetricGeneration.Add(1) r.metricGeneration.Store(metricGeneration) r.clientMu.Unlock() defer func() { r.clientMu.Lock() if r.controlConn == conn { r.controlConn = nil r.controlWriteMu = nil r.metricGeneration.CompareAndSwap(metricGeneration, 0) } r.clientMu.Unlock() _ = conn.Close() }() log.Println("worker: websocket connected") done := make(chan struct{}) var doneOnce sync.Once closeDone := func() { doneOnce.Do(func() { close(done) }) } var connectionWG sync.WaitGroup defer func() { closeDone() _ = conn.Close() connectionWG.Wait() }() // Heartbeat goroutine — shares writeMu with the writer. connectionWG.Add(1) go func() { defer connectionWG.Done() r.heartbeat(conn, &writeMu, done) }() // Single writer goroutine for this connection: serializes result // and heartbeat writes through writeMu so websocket.WriteJSON is // never called concurrently. The dispatcher loop feeds it via the // bounded results channel. connectionWG.Add(1) go func() { defer connectionWG.Done() r.writer(conn, &writeMu, done, metricGeneration) }() for { var msg wire.WorkerMessage if err := conn.ReadJSON(&msg); err != nil { closeDone() return err } if (msg.Kind == "init" || msg.Kind == "config") && msg.Init != nil { if err := r.applyInit(msg.Init); err != nil { log.Printf("worker: rejected config: %v", err) continue } if msg.Init.RotateToken != "" && msg.Init.RotationID != "" { r.applyRotation(conn, &writeMu, msg.Init.RotateToken, msg.Init.RotationID) } continue } if msg.Kind == "stale_lease_ack" && msg.StaleLeaseAck != nil { r.leasesMu.Lock() if r.leases[msg.StaleLeaseAck.JobID] == msg.StaleLeaseAck.LeaseToken { delete(r.leases, msg.StaleLeaseAck.JobID) } r.leasesMu.Unlock() continue } if msg.Kind != "task" { continue } if !r.enqueueTaskMessage(msg) { closeDone() return nil } } } // enqueueTaskMessage prefers the v2 envelope over sibling legacy fields. An // envelope is accepted only when it selects exactly one matching payload with // the same non-empty outer and inner job IDs. Invalid envelopes never fall // back to a sibling legacy payload, which could otherwise execute a task the // control plane did not intend to send. func (r *Runner) enqueueTaskMessage(msg wire.WorkerMessage) bool { //nolint:gocritic // wire envelope is the dispatcher boundary if msg.TaskEnvelope != nil { return r.enqueueTaskEnvelope(msg.TaskEnvelope) } switch { case msg.NotificationTask != nil: if msg.NotificationTask.JobID == "" || msg.NotificationTask.LeaseToken == "" { return true } if !r.taskInScope(msg.NotificationTask.AccountID) { return r.enqueueFailedNotification(*msg.NotificationTask, "account_scope_mismatch") } log.Printf("worker: received websocket notification task %s method=%s", msg.NotificationTask.JobID, msg.NotificationTask.Method) return r.EnqueueNotification(*msg.NotificationTask) case msg.Task != nil: if msg.Task.JobID == "" || msg.Task.LeaseToken == "" { return true } log.Printf("worker: received websocket task %s", msg.Task.JobID) if !r.taskInScope(msg.Task.AccountID) { return r.enqueueFailedCheck(*msg.Task, "account_scope_mismatch") } if !checkexec.SupportsKind(msg.Task.Kind) { return r.enqueueUnsupportedCheck(*msg.Task) } return r.Enqueue(*msg.Task) default: return true } } func (r *Runner) enqueueTaskEnvelope(envelope *wire.TaskEnvelope) bool { if (envelope.Job == nil) == (envelope.Notify == nil) { return true } if envelope.Job != nil { if !matchingEnvelopeJobID(envelope.JobID, envelope.Job.JobID) || envelope.Job.LeaseToken == "" { return true } if !r.taskInScope(envelope.Job.AccountID) { return r.enqueueFailedCheck(*envelope.Job, "account_scope_mismatch") } if envelope.Type != wire.TaskTypeCheck { return r.enqueueFailedCheck(*envelope.Job, malformedTaskEnvelopeError) } if !checkexec.SupportsKind(envelope.Job.Kind) { return r.enqueueUnsupportedCheck(*envelope.Job) } r.rememberLease(envelope.Job.JobID, envelope.Job.LeaseToken) return r.Enqueue(*envelope.Job) } if !matchingEnvelopeJobID(envelope.JobID, envelope.Notify.JobID) || envelope.Notify.LeaseToken == "" { return true } if envelope.Type != wire.TaskTypeNotification { return r.enqueueFailedNotification(*envelope.Notify, malformedTaskEnvelopeError) } if !r.taskInScope(envelope.Notify.AccountID) { return r.enqueueFailedNotification(*envelope.Notify, "account_scope_mismatch") } r.rememberLease(envelope.Notify.JobID, envelope.Notify.LeaseToken) return r.EnqueueNotification(*envelope.Notify) } func (r *Runner) rememberLease(jobID, token string) { if jobID == "" || token == "" { return } r.leasesMu.Lock() if r.leases == nil { r.leases = make(map[string]string) } r.leases[jobID] = token r.leasesMu.Unlock() } func (r *Runner) forgetLease(jobID string) { r.leasesMu.Lock() delete(r.leases, jobID) r.leasesMu.Unlock() } func (r *Runner) drainFinalResults() { // Shutdown keeps the existing connection just long enough to report a // bounded set of terminal frames and leases that cannot complete. r.clientMu.Lock() conn, writeMu := r.controlConn, r.controlWriteMu r.clientMu.Unlock() if conn == nil || writeMu == nil { return } deadline := time.Now().Add(finalDrainTimeout) count := 0 r.leasesMu.Lock() leases := make(map[string]string, len(r.leases)) for id, token := range r.leases { leases[id] = token } r.leasesMu.Unlock() for count < finalDrainLimit && time.Now().Before(deadline) { select { case env := <-r.results: for i := range env.reports { env.reports[i].LeaseToken = env.job.LeaseToken if r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "result", Result: &env.reports[i]}) != nil { return } r.forgetLease(env.job.JobID) count++ if count >= finalDrainLimit { return } } case env := <-r.notifyResults: if r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "result", NotificationResult: &env.report}) != nil { return } r.forgetLease(env.task.JobID) count++ default: for id, token := range leases { if count >= finalDrainLimit || time.Now().After(deadline) { return } r.leasesMu.Lock() _, pending := r.leases[id] r.leasesMu.Unlock() if pending { if r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "stale_lease", StaleLease: &wire.StaleLeaseReport{JobID: id, LeaseToken: token}}) != nil { return } count++ } } return } } } func (r *Runner) taskInScope(accountID int64) bool { if accountID == 0 { return true } // legacy task payload r.identityMu.RLock() defer r.identityMu.RUnlock() return r.state.AccountID == 0 || r.state.AccountID == accountID } func matchingEnvelopeJobID(outer, inner string) bool { return outer != "" && outer == inner } func (r *Runner) enqueueUnsupportedCheck(job wire.CheckJob) bool { return r.enqueueFailedCheck(job, "unsupported_kind: "+job.Kind) } func (r *Runner) enqueueFailedCheck(job wire.CheckJob, errorCode string) bool { if r.results == nil || r.stopped() { return false } report := wire.CheckResultReport{ JobID: job.JobID, CheckID: job.CheckID, MonitorID: job.MonitorID, State: "FAIL", Error: stringPtr(errorCode), DurationMs: 0, } select { case r.results <- resultEnvelope{job: job, reports: []wire.CheckResultReport{report}}: return true case <-r.stopCh: return false } } func (r *Runner) enqueueFailedNotification(task wire.NotificationTask, errorCode string) bool { if r.stopped() { return false } r.completeNotification(task, notificationPermanentReport(task, errorCode, time.Now())) return true } func (r *Runner) stopped() bool { select { case <-r.stopCh: return true default: return false } } func (r *Runner) heartbeat(conn *websocket.Conn, writeMu *sync.Mutex, done <-chan struct{}) { ticker := time.NewTicker(heartbeatInterval) defer ticker.Stop() for { select { case <-ticker.C: writeMu.Lock() err := conn.WriteJSON(wire.WorkerMessage{ Kind: "heartbeat", Heartbeat: &wire.HeartbeatRequest{ ActiveChecks: r.ActiveCount(), QueueDepth: r.QueueDepth(), ActiveNotifications: r.ActiveNotifications(), NotificationQueueDepth: r.NotificationQueueDepth(), }, }) writeMu.Unlock() if err != nil { return } r.touchHeartbeat(time.Now().UTC()) case <-done: return case <-r.stopCh: return } } } // touchHeartbeat records the most recent successful heartbeat write. // Called from the heartbeat goroutine; the webapp reads it via // LastHeartbeatAck. func (r *Runner) touchHeartbeat(at time.Time) { if r == nil { return } r.lastHeartbeatMu.Lock() r.lastHeartbeatAt = at r.lastHeartbeatMu.Unlock() } // LastHeartbeatAck returns the wall-clock time of the most recent // successful heartbeat. Returns the zero time if no heartbeat has // been written yet. func (r *Runner) LastHeartbeatAck() time.Time { if r == nil { return time.Time{} } r.lastHeartbeatMu.RLock() defer r.lastHeartbeatMu.RUnlock() return r.lastHeartbeatAt } func (r *Runner) writer(conn *websocket.Conn, writeMu *sync.Mutex, done <-chan struct{}, metricGeneration uint64) { for { if message, ok := r.takeOutbox(); ok { if err := r.writeControlMessage(conn, writeMu, message); err != nil { r.requeueOutbox(message) _ = conn.Close() return } continue } select { case env, ok := <-r.results: if !ok { return } for i := range env.reports { report := &env.reports[i] report.LeaseToken = env.job.LeaseToken message := wire.WorkerMessage{Kind: "result", Result: report} if err := r.writeControlMessage(conn, writeMu, message); err != nil { remaining := []wire.WorkerMessage{message} for j := i + 1; j < len(env.reports); j++ { next := env.reports[j] next.LeaseToken = env.job.LeaseToken remaining = append(remaining, wire.WorkerMessage{Kind: "result", Result: &next}) } r.requeueOutbox(remaining...) log.Printf( "worker: failed to report result job=%s check=%d kind=%s state=%s: %v", report.JobID, report.CheckID, env.job.Kind, report.State, err, ) _ = conn.Close() return } log.Printf( "worker: completed job=%s check=%d kind=%s state=%s reported=true", report.JobID, report.CheckID, env.job.Kind, report.State, ) r.resultsBuf.add(resultRowFromReport(&env, report, time.Now().UTC())) r.forgetLease(env.job.JobID) } case env, ok := <-r.notifyResults: if !ok { return } message := wire.WorkerMessage{Kind: "result", NotificationResult: &env.report} if err := r.writeControlMessage(conn, writeMu, message); err != nil { r.requeueOutbox(message) log.Printf( "worker: failed to report notification result job=%s method=%s status=%s: %v", env.report.JobID, env.task.Method, env.report.Status, err, ) _ = conn.Close() return } log.Printf( "worker: completed notification job=%s method=%s status=%s message=%d duration_ms=%d", env.report.JobID, env.task.Method, env.report.Status, env.task.MessageID, env.report.DurationMs, ) r.forgetLease(env.task.JobID) case metric := <-r.metricResults: if metric.generation != metricGeneration { continue } message := wire.WorkerMessage{Kind: "result", ServerMetric: &metric.report} if err := r.writeControlMessage(conn, writeMu, message); err != nil { // Metrics are periodic snapshots without a lease or idempotency key. // Dropping a failed snapshot avoids duplicate control-plane inserts; // the next collection tick supplies a fresh replacement. _ = conn.Close() return } case <-r.outboxWake: case <-done: return case <-r.stopCh: return } } } func (r *Runner) takeOutbox() (wire.WorkerMessage, bool) { r.outboxMu.Lock() defer r.outboxMu.Unlock() if len(r.outbox) == 0 { return wire.WorkerMessage{}, false } message := r.outbox[0] r.outbox = r.outbox[1:] return message, true } func (r *Runner) requeueOutbox(messages ...wire.WorkerMessage) { if len(messages) == 0 { return } r.outboxMu.Lock() r.outbox = append(messages, r.outbox...) r.outboxMu.Unlock() select { case r.outboxWake <- struct{}{}: default: } } func (r *Runner) writeControlMessage(conn *websocket.Conn, writeMu *sync.Mutex, message wire.WorkerMessage) error { writeMu.Lock() defer writeMu.Unlock() if r.beforeControlWrite != nil { r.beforeControlWrite() } return conn.WriteJSON(message) } func (r *Runner) applyInit(init *wire.WorkerInit) error { if err := r.verifyInit(init); err != nil { return err } if init.Concurrency > 0 && init.Concurrency != r.Concurrency() { size := init.Concurrency if size > r.maxConcurrency { size = r.maxConcurrency } if size < 1 { size = 1 } atomic.StoreInt64(&r.concurrency, int64(size)) r.pool.SetSize(size) } if len(init.LLMs) > 0 { llm := init.LLMs[0] setEnvIfNotEmpty("LLM_URL", llm.URL) setEnvIfNotEmpty("LLM_MODEL", llm.Model) setEnvIfNotEmpty("LLM_APIKEY", llm.APIKey) setEnvIfNotEmpty("LLM_KIND", llm.Kind) } r.credentialsMu.Lock() r.credentials = init.Credentials r.credentialsMu.Unlock() r.systemContactsMu.Lock() r.systemContacts = init.SystemContacts r.systemContactsMu.Unlock() r.urlMu.Lock() accepted, field := init.PublicURL, "public_url" if accepted == "" { accepted, field = init.URL, "url" // legacy wire field (bounded migration) } // Validate the control-plane-supplied endpoint before accepting it. // On an unusable value keep the previous accepted URL (never regress // to a blank or garbage endpoint) and surface the rejection. A valid // empty value still clears the stored URL. var err error if accepted != "" { err = ValidateAdvertisedURL(accepted) } if err != nil { log.Printf("worker: ignoring invalid advertised URL from control plane (%s=%q): %v", field, accepted, err) } else { r.url = accepted } r.urlMu.Unlock() // Peers is the slice of other workers this node can reach for // cross-worker confirmation. Refresh the cache so a removed // peer's stale observation is dropped immediately rather than // lingering until peerStatusMaxAge. r.peersMu.Lock() r.peers = append([]wire.PeerInfo(nil), init.Peers...) r.peersMu.Unlock() if r.peerCache != nil { r.peerCache.resetFor(init.Peers) } r.workerIDMu.Lock() r.workerID = init.WorkerID r.workerIDMu.Unlock() r.regionMu.Lock() r.regionCode = init.RegionCode r.regionMu.Unlock() r.versionMu.Lock() r.workerVersion = init.Version r.versionMu.Unlock() r.capsMu.Lock() r.workerCaps = append([]string(nil), init.Capabilities...) r.capsMu.Unlock() if init.ServerID == nil { r.serverID.Store(0) } else { r.serverID.Store(*init.ServerID) } smtpCount, tgCount := 0, 0 if init.Credentials != nil { smtpCount = len(init.Credentials.SMTP) tgCount = len(init.Credentials.Telegram) } log.Printf( "worker: %s received worker_id=%s region=%s version=%s capabilities=%v "+ "concurrency=%d llms=%d credentials_smtp=%d credentials_telegram=%d system_contacts=%d peers=%d", "config", init.WorkerID, init.RegionCode, init.Version, init.Capabilities, init.Concurrency, len(init.LLMs), smtpCount, tgCount, len(init.SystemContacts), len(init.Peers), ) return nil } func (r *Runner) verifyInit(init *wire.WorkerInit) error { // Legacy control planes do not sign config. Do not turn existing workers // into an outage; once bootstrap pins a key, signed config is mandatory. r.identityMu.Lock() defer r.identityMu.Unlock() if r.state.VerificationKey == "" { return nil } if init.Signature == "" || init.SigningKeyID == "" || init.SigningKeyID != r.state.SigningKeyID || init.AccountID == nil || init.WorkerID != r.state.WorkerID { return fmt.Errorf("unsigned or mismatched config") } if r.state.AccountID != 0 && *init.AccountID != r.state.AccountID { return fmt.Errorf("account scope changed") } if init.ConfigVersion < r.state.ConfigVersion { return fmt.Errorf("config version downgrade") } expires, err := time.Parse(time.RFC3339Nano, init.ExpiresAt) if err != nil || !expires.After(time.Now()) { return fmt.Errorf("config expired") } key, err := base64.StdEncoding.DecodeString(r.state.VerificationKey) if err != nil || len(key) != ed25519.PublicKeySize { return fmt.Errorf("invalid pinned verification key") } sig, err := base64.StdEncoding.DecodeString(init.Signature) if err != nil { return fmt.Errorf("invalid config signature") } signature := init.Signature init.Signature = "" body, err := json.Marshal(init) init.Signature = signature if err != nil || !ed25519.Verify(ed25519.PublicKey(key), body, sig) { return fmt.Errorf("invalid config signature") } // Issued/expiry timestamps are signed freshness metadata, regenerated on // reconnect; they do not change immutable configuration content. canonical := *init canonical.Signature, canonical.IssuedAt, canonical.ExpiresAt = "", "", "" canonicalBody, err := json.Marshal(&canonical) if err != nil { return err } payload := base64.StdEncoding.EncodeToString(canonicalBody) if init.ConfigVersion == r.state.ConfigVersion && r.state.ConfigPayload != "" && r.state.ConfigPayload != payload { return fmt.Errorf("config version replay content mismatch") } r.state.AccountID, r.state.ConfigVersion, r.state.ConfigPayload = *init.AccountID, init.ConfigVersion, payload if err := saveWorkerState(r.config.StateFile, r.state); err != nil { return err } return nil } func (r *Runner) applyRotation(conn *websocket.Conn, writeMu *sync.Mutex, token, rotationID string) { if token == "" || token == r.config.Token { return } r.clientMu.Lock() r.config.Token = token r.client = NewClient(r.config.URL, token) r.state.Token = token r.clientMu.Unlock() if err := saveWorkerState(r.config.StateFile, r.state); err != nil { log.Printf("worker: rotation state: %v", err) return } _ = r.writeControlMessage(conn, writeMu, wire.WorkerMessage{Kind: "rotation_ack", RotationAck: &wire.RotationAck{RotationID: rotationID}}) _ = conn.Close() } // Credentials returns a snapshot of the current notification credentials. // Safe for concurrent use. Returns nil if no credentials have been pushed. func (r *Runner) Credentials() *wire.NotificationCredentials { r.credentialsMu.RLock() defer r.credentialsMu.RUnlock() return r.credentials } // SystemContacts returns a snapshot of the cached system contacts. // Safe for concurrent use. func (r *Runner) SystemContacts() []wire.SystemContact { r.systemContactsMu.RLock() defer r.systemContactsMu.RUnlock() out := make([]wire.SystemContact, len(r.systemContacts)) copy(out, r.systemContacts) return out } // URL returns the publicly-advertised URL the main app and peer workers // should dial to reach this worker. It is set from the init/config // websocket message and may be empty if the main app did not push one // (e.g. legacy worker). Safe for concurrent use. func (r *Runner) URL() string { r.urlMu.RLock() defer r.urlMu.RUnlock() return r.url } // Peers returns a snapshot of the cached peer list pushed by the // control plane via WorkerInit.Peers. The current worker is excluded // upstream so a worker never dials itself. Returns nil when no init // has landed or when the control plane ships an empty list (e.g. a // single-worker install); callers must handle that case explicitly so // they fall back to the single-node selfcheck verdict. Safe for // concurrent use. func (r *Runner) Peers() []wire.PeerInfo { r.peersMu.RLock() defer r.peersMu.RUnlock() if len(r.peers) == 0 { return nil } out := make([]wire.PeerInfo, len(r.peers)) copy(out, r.peers) return out } // SetMasterStatus records the most recent local selfcheck verdict. // Called by the selfcheck loop after every probe so the // /api/peer/status HTTP handler can answer with the same value the // consensus uses. Passing up == nil resets the snapshot to "unknown" // (no probe yet) and is used by tests. func (r *Runner) SetMasterStatus(up *bool, observedAt time.Time) { if r == nil { return } r.masterStatusMu.Lock() r.masterStatusUp = up r.masterStatusAt = observedAt r.masterStatusMu.Unlock() } // MasterStatus returns the most recent local selfcheck verdict and // the wall-clock time it was produced. up == nil means the local // selfcheck has not produced a verdict yet (very first tick or the // runner has not been started). Safe for concurrent use. func (r *Runner) MasterStatus() (up *bool, observedAt time.Time) { if r == nil { return nil, time.Time{} } r.masterStatusMu.RLock() defer r.masterStatusMu.RUnlock() return r.masterStatusUp, r.masterStatusAt } func setEnvIfNotEmpty(key, value string) { if value != "" { _ = os.Setenv(key, value) } } // defaultExecuteJob executes a single check job. func (r *Runner) defaultExecuteJob(payload interface{}) interface{} { job := payload.(wire.CheckJob) log.Printf("worker: executing job %s (check %d, kind %s)", job.JobID, job.CheckID, job.Kind) monitor := &models.Monitor{Host: job.Host} monitor.ID = job.MonitorID check := models.Check{ Kind: job.Kind, URL: job.URL, Settings: datatypes.JSON(job.Settings), } check.ID = job.CheckID check.MonitorID = job.MonitorID results := checkexec.Execute(monitor, []models.Check{check}) if len(results) == 0 { log.Printf("worker: no results for job %s", job.JobID) return []wire.CheckResultReport{} } reports := make([]wire.CheckResultReport, 0, len(results)) for _, result := range results { report := wire.CheckResultReport{ JobID: job.JobID, CheckID: job.CheckID, MonitorID: job.MonitorID, State: result.Result.State, DurationMs: result.Result.Duration.Milliseconds(), Warnings: result.Result.Warnings, Infos: result.Result.Infos, Metrics: result.Metrics, } if result.Result.Error != nil { errStr := result.Result.Error.Error() report.Error = &errStr } if result.Result.Expires != nil { exp := result.Result.Expires.Format(time.RFC3339) report.ExpiresAt = &exp } reports = append(reports, report) } return reports } // RecentResults returns the most recent n result rows produced by // this worker, in chronological order. n <= 0 returns an empty slice. // Safe to call before Start (returns nil). func (r *Runner) RecentResults(n int) []ResultRow { if r == nil || r.resultsBuf == nil { return nil } return r.resultsBuf.snapshot(n) } // RecentNotifications returns the most recent n notification rows. func (r *Runner) RecentNotifications(n int) []NotificationRow { if r == nil || r.notificationsBuf == nil { return nil } return r.notificationsBuf.snapshot(n) } // RecordNotification appends one row to the notification ring buffer. // Called from selfcheck.sendSystemAlert. Safe before Start. func (r *Runner) RecordNotification(n *NotificationRow) { if r == nil || r.notificationsBuf == nil { return } if n.At.IsZero() { n.At = time.Now().UTC() } r.notificationsBuf.add(n) } func (r *Runner) recordDelegatedNotification(jobID, method string, report wire.NotificationResultReport) { r.RecordNotification(&NotificationRow{ JobID: jobID, Method: method, Status: report.Status, DurationMs: report.DurationMs, At: time.Now().UTC(), }) } // Token returns the current bearer token. The webapp settings page // masks this for display. func (r *Runner) Token() string { if r == nil || r.config == nil { return "" } r.clientMu.Lock() defer r.clientMu.Unlock() return r.config.Token } // TokenRotatedAt returns the wall-clock time of the most recent // successful token rotation. Zero before the first rotation. func (r *Runner) TokenRotatedAt() time.Time { if r == nil { return time.Time{} } r.tokenRotatedMu.RLock() defer r.tokenRotatedMu.RUnlock() return r.tokenRotatedAt } // WorkerID returns the worker_id pushed by the main app via the // init/config message. Empty before init lands. func (r *Runner) WorkerID() string { if r == nil { return "" } // The init payload is captured via the credentials/url caches; we // also expose the worker_id through the URL setter below. r.workerIDMu.RLock() defer r.workerIDMu.RUnlock() return r.workerID } // RegionCode returns the region_code from the init payload. func (r *Runner) RegionCode() string { if r == nil { return "" } r.regionMu.RLock() defer r.regionMu.RUnlock() return r.regionCode } // WorkerVersion returns the worker_version from the init payload. func (r *Runner) WorkerVersion() string { if r == nil { return "" } r.versionMu.RLock() defer r.versionMu.RUnlock() return r.workerVersion } // WorkerCapabilities returns the capabilities from the init payload. func (r *Runner) WorkerCapabilities() []string { if r == nil { return nil } r.capsMu.RLock() defer r.capsMu.RUnlock() return append([]string(nil), r.workerCaps...) } // HTTPConfig returns the worker's local HTTP listener settings. Used // by the webapp's settings page to render the configured bind // address and URL. func (r *Runner) HTTPConfig() HTTPConfig { if r == nil || r.config == nil { return HTTPConfig{} } return r.config.HTTP } // RotateToken asks the main app to issue a fresh bearer token, // updates the in-memory config + client, and closes the current // websocket so the reconnect loop picks up the new token. // // Returns the new token string. On any failure the old token and // client are kept untouched. func (r *Runner) RotateToken(ctx context.Context) (string, error) { if r == nil || r.config == nil { return "", fmt.Errorf("worker: runner not initialized") } // The endpoint authenticates with the current token, so concurrent // rotations must not mint replacements from the same stale client. r.rotationMu.Lock() defer r.rotationMu.Unlock() if r.stopped() { return "", fmt.Errorf("worker: runner stopped") } requestCtx, cancelRequest := context.WithCancel(ctx) rotationDone := make(chan struct{}) defer func() { close(rotationDone) cancelRequest() }() go func() { select { case <-r.rotationCtx.Done(): cancelRequest() case <-rotationDone: } }() r.clientMu.Lock() client := r.client r.clientMu.Unlock() if client == nil { return "", fmt.Errorf("worker: client not yet started") } newToken, err := client.RotateToken(requestCtx) if err != nil { return "", err } if r.beforeTokenCommit != nil { r.beforeTokenCommit() } // Stop and the post-HTTP commit share lifecycleMu. Once Stop has closed // stopCh, this rotation cannot install or report a replacement token. r.lifecycleMu.Lock() r.clientMu.Lock() if r.stopped() { r.clientMu.Unlock() r.lifecycleMu.Unlock() return "", fmt.Errorf("worker: runner stopped") } if newToken == "" || newToken == r.config.Token { r.clientMu.Unlock() r.lifecycleMu.Unlock() return "", fmt.Errorf("worker: rotate-token returned unchanged or empty token") } // Swap client and detach only the active control-plane connection. stopCh // remains exclusively owned by Stop, so dispatch, metrics, selfcheck, peer // polling, and services started alongside the runner keep running. r.config.Token = newToken r.client = NewClient(r.config.URL, newToken) conn := r.controlConn writeMu := r.controlWriteMu r.clientMu.Unlock() r.lifecycleMu.Unlock() // Stamp the rotation time so the settings page can show it. r.tokenRotatedMu.Lock() r.tokenRotatedAt = time.Now().UTC() r.tokenRotatedMu.Unlock() // Interrupt the current connection before waiting for its writer. A stalled // WriteJSON holds writeMu; closing the socket makes that write fail so the // writer can requeue its dequeued envelope and release the mutex. if conn != nil { _ = conn.Close() } if writeMu != nil { writeMu.Lock() writeMu.Unlock() } r.lifecycleMu.Lock() stopped := r.stopped() r.lifecycleMu.Unlock() if stopped { return "", fmt.Errorf("worker: runner stopped") } select { case r.reconnectCh <- struct{}{}: default: } return newToken, nil }