Files
worker/internal/distworker/runner.go
Gleb Tv cb23f123ae
Все проверки выполнены успешно
CI / test (push) Successful in 10m15s
Docker / Build and publish worker image (push) Successful in 34m59s
feat(worker): adopt canonical public URL
2026-08-12 20:48:01 +03:00

1404 строки
41 KiB
Go

package distworker
import (
"context"
"encoding/json"
"fmt"
"log"
"os"
"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"
)
// 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
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
// 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(),
}
}
// 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.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")
}
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()
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 {
r.applyInit(msg.Init)
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
}
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 !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 envelope.Type != wire.TaskTypeCheck {
return r.enqueueFailedCheck(*envelope.Job, malformedTaskEnvelopeError)
}
if !checkexec.SupportsKind(envelope.Job.Kind) {
return r.enqueueUnsupportedCheck(*envelope.Job)
}
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)
}
return r.EnqueueNotification(*envelope.Notify)
}
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()))
}
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,
)
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) {
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),
)
}
// 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
}