3dfc833671
Move agent-run orchestration from the legacy supervisor path into the new agentrun worker/service package and wire it through coredata, server, policies, and GraphQL resolvers. This consolidates run lifecycle handling around lease-aware workers and aligns API surface with the new agent-run domain model so reviewers can follow one coherent execution path. Signed-off-by: Bryan Frimin <bryan@probo.com>
122 lines
2.9 KiB
Go
122 lines
2.9 KiB
Go
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
|
//
|
|
// Permission to use, copy, modify, and/or distribute this software for any
|
|
// purpose with or without fee is hereby granted, provided that the above
|
|
// copyright notice and this permission notice appear in all copies.
|
|
//
|
|
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
|
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
|
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
|
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
|
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
|
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
|
// PERFORMANCE OF THIS SOFTWARE.
|
|
|
|
package agentrun
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
"go.gearno.de/kit/log"
|
|
"go.gearno.de/kit/pg"
|
|
"go.gearno.de/kit/worker"
|
|
"go.probo.inc/probo/pkg/agent"
|
|
"go.probo.inc/probo/pkg/coredata"
|
|
)
|
|
|
|
type (
|
|
Worker struct {
|
|
handler *handler
|
|
kitWorker *worker.Worker[coredata.AgentRun]
|
|
}
|
|
|
|
WorkerOption func(*workerConfig)
|
|
|
|
workerConfig struct {
|
|
interval time.Duration
|
|
leaseDuration time.Duration
|
|
maxConcurrency int
|
|
}
|
|
)
|
|
|
|
var (
|
|
ErrHeartbeatFailed = errors.New("agent run heartbeat failed")
|
|
ErrLeaseLost = errors.New("agent run lease lost")
|
|
)
|
|
|
|
func WithWorkerInterval(d time.Duration) WorkerOption {
|
|
return func(c *workerConfig) {
|
|
if d > 0 {
|
|
c.interval = d
|
|
}
|
|
}
|
|
}
|
|
|
|
func WithWorkerLeaseDuration(d time.Duration) WorkerOption {
|
|
return func(c *workerConfig) {
|
|
if d > 0 {
|
|
c.leaseDuration = d
|
|
}
|
|
}
|
|
}
|
|
|
|
func WithWorkerMaxConcurrency(n int) WorkerOption {
|
|
return func(c *workerConfig) {
|
|
if n > 0 {
|
|
c.maxConcurrency = n
|
|
}
|
|
}
|
|
}
|
|
|
|
func NewWorker(
|
|
pgClient *pg.Client,
|
|
store *coredata.PGCheckpointer,
|
|
registry agent.AgentRegistry,
|
|
logger *log.Logger,
|
|
opts ...WorkerOption,
|
|
) *Worker {
|
|
cfg := workerConfig{
|
|
interval: 10 * time.Second,
|
|
leaseDuration: 5 * time.Minute,
|
|
maxConcurrency: 5,
|
|
}
|
|
|
|
for _, opt := range opts {
|
|
opt(&cfg)
|
|
}
|
|
|
|
h := &handler{
|
|
pg: pgClient,
|
|
store: store,
|
|
registry: registry,
|
|
logger: logger,
|
|
leaseDuration: cfg.leaseDuration,
|
|
shutdownCh: make(chan struct{}),
|
|
}
|
|
|
|
w := worker.New(
|
|
"agent-run-worker",
|
|
h,
|
|
logger,
|
|
worker.WithInterval(cfg.interval),
|
|
worker.WithMaxConcurrency(cfg.maxConcurrency),
|
|
)
|
|
|
|
return &Worker{handler: h, kitWorker: w}
|
|
}
|
|
|
|
// Run starts the worker loop. It blocks until ctx is cancelled, then
|
|
// closes the shutdown broadcast channel so in-flight Process calls can
|
|
// checkpoint and exit, and waits for all of them to drain before
|
|
// returning.
|
|
//
|
|
// signalShutdown is registered without a stop hook because it is
|
|
// idempotent (sync.Once) and we want it to fire on every ctx
|
|
// cancellation, even one that races with kitWorker.Run returning.
|
|
func (w *Worker) Run(ctx context.Context) error {
|
|
context.AfterFunc(ctx, w.handler.signalShutdown)
|
|
return w.kitWorker.Run(ctx)
|
|
}
|