rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replace supervisor with agentrun worker service

Browse Source 
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>
This commit is contained in:
Bryan Frimin
2026-06-07 09:00:01 +02:00
parent 400800fd41
commit 3dfc833671
27 changed files with 1076 additions and 1260 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
pkg/agentrun/export_test.go Normal file
View File

@@ -0,0 +1,24 @@
// 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
// ShutdownBroadcast returns the channel that closes once the worker has
// broadcast graceful shutdown to all in-flight runs. It is compiled only
// in test builds so external tests can synchronize tool release with
// shutdown propagation without leaking a test-only method into the
// worker's public API.
func (w *Worker) ShutdownBroadcast() <-chan struct{} {
return w.handler.shutdownCh
}

375
pkg/agentrun/handler.go Normal file
View File

@@ -0,0 +1,375 @@
// 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"
"encoding/json"
"errors"
"fmt"
"sync"
"time"
"unicode/utf8"
"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"
"go.probo.inc/probo/pkg/llm"
)
type handler struct {
pg *pg.Client
store *coredata.PGCheckpointer
registry agent.AgentRegistry
logger *log.Logger
leaseDuration time.Duration
shutdownCh chan struct{}
shutdownOnce sync.Once
}
var (
_ worker.Handler[coredata.AgentRun] = (*handler)(nil)
_ worker.StaleRecoverer = (*handler)(nil)
)
// Claim loads the next pending agent run, marks it RUNNING with a lease
// owned by this worker, and returns the row. When no work is available it
// returns worker.ErrNoTask so the kit can back off until the next tick.
func (h *handler) Claim(ctx context.Context) (coredata.AgentRun, error) {
var (
run = coredata.AgentRun{}
now = time.Now()
leaseExpiresAt = now.Add(h.leaseDuration)
)
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := run.LoadNextPendingForUpdateSkipLocked(ctx, tx); err != nil {
return fmt.Errorf("cannot load next pending agent run: %w", err)
}
run.Status = coredata.AgentRunStatusRunning
run.StartedAt = &now
run.LeaseExpiresAt = &leaseExpiresAt
run.LeaseGeneration++
run.UpdatedAt = now
if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return fmt.Errorf("cannot update agent run: %w", err)
}
return nil
},
); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.AgentRun{}, worker.ErrNoTask
}
return coredata.AgentRun{}, err
}
return run, nil
}
// Process executes a single agent run. It spawns a heartbeat goroutine
// that renews the lease while the run is active, and a forwarder
// goroutine that converts the handler-level shutdown broadcast into a
// per-run ctx cancellation so the agent loop checkpoints cleanly at
// its next turn boundary.
//
// The returned error mirrors the run outcome so the worker kit's
// task metrics and OTel span status reflect actual agent failures.
// nil is returned for both successful runs and graceful exits
// (lease loss, infrastructure suspension) where the row state is
// already consistent.
func (h *handler) Process(ctx context.Context, run coredata.AgentRun) error {
runCtx, cancelRun := context.WithCancelCause(ctx)
defer cancelRun(nil)
leaseGeneration := run.LeaseGeneration
forwarderDone := make(chan struct{})
defer close(forwarderDone)
go func() {
select {
case <-h.shutdownCh:
cancelRun(agent.ErrSuspendForCheckpoint)
case <-forwarderDone:
}
}()
heartbeatCtx, cancelHeartbeat := context.WithCancel(ctx)
defer cancelHeartbeat()
go h.heartbeatLease(heartbeatCtx, run.ID.String(), leaseGeneration, cancelRun)
return h.executeRun(runCtx, &run, leaseGeneration)
}
// RecoverStale resets agent runs whose worker lease has expired back to
// PENDING so a fresh worker can pick them up on the next cycle.
func (h *handler) RecoverStale(ctx context.Context) error {
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleAgentRuns(ctx, conn)
},
); err != nil {
return fmt.Errorf("cannot reset stale agent runs: %w", err)
}
return nil
}
// signalShutdown closes the handler-level shutdown broadcast channel. All
// in-flight Process forwarder goroutines observe the close and propagate
// it to their per-run agent stop channels, letting agents checkpoint at
// the next turn boundary before Process returns.
func (h *handler) signalShutdown() {
h.shutdownOnce.Do(func() { close(h.shutdownCh) })
}
func (h *handler) heartbeatLease(
ctx context.Context,
runID string,
leaseGeneration int64,
cancelRun context.CancelCauseFunc,
) {
ticker := time.NewTicker(h.leaseDuration / 3)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
expiresAt := time.Now().Add(h.leaseDuration)
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
rowsAffected, err := coredata.HeartbeatAgentRunLease(
ctx,
conn,
runID,
leaseGeneration,
expiresAt,
)
if err != nil {
return err
}
if rowsAffected == 0 {
return ErrLeaseLost
}
return nil
},
); err != nil {
h.logger.ErrorCtx(ctx, "cannot heartbeat agent run lease", log.Error(err))
if errors.Is(err, ErrLeaseLost) {
cancelRun(ErrLeaseLost)
} else {
cancelRun(fmt.Errorf("%w: %w", ErrHeartbeatFailed, err))
}
return
}
}
}
}
const (
// errorMessageMaxLen caps the error string persisted to the
// agent_runs.error_message column. Raw tool or LLM errors can embed
// URLs with credentials, response snippets containing PII, or partial
// records from failed DB lookups; the full context is logged while
// only a truncated summary is stored for caller-visible state.
errorMessageMaxLen = 512
)
func sanitizeError(err error) string {
msg := err.Error()
if len(msg) <= errorMessageMaxLen {
return msg
}
cut := errorMessageMaxLen
for cut > 0 && !utf8.RuneStart(msg[cut]) {
cut--
}
return msg[:cut] + "…"
}
type leasedCheckpointer struct {
store *coredata.PGCheckpointer
leaseGeneration int64
}
func (s leasedCheckpointer) Save(ctx context.Context, runID string, cp *agent.Checkpoint) error {
return s.store.SaveForLease(ctx, runID, cp, s.leaseGeneration)
}
func (s leasedCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) {
return s.store.Load(ctx, runID)
}
func (h *handler) executeRun(
ctx context.Context,
run *coredata.AgentRun,
leaseGeneration int64,
) error {
runID := run.ID.String()
checkpointer := leasedCheckpointer{
store: h.store,
leaseGeneration: leaseGeneration,
}
var (
result *agent.Result
runErr error
)
if run.Checkpoint != nil {
h.logger.InfoCtx(ctx, "resuming agent run", log.String("run_id", runID))
result, runErr = agent.Restore(ctx, checkpointer, runID, h.registry)
} else {
h.logger.InfoCtx(ctx, "starting agent run", log.String("run_id", runID))
a, err := h.registry.Agent(run.StartAgentName)
if err != nil {
runErr = fmt.Errorf("cannot resolve agent %q: %w", run.StartAgentName, err)
} else {
var inputMsgs []llm.Message
if err := json.Unmarshal(run.InputMessages, &inputMsgs); err != nil {
runErr = fmt.Errorf("cannot unmarshal input messages: %w", err)
} else {
result, runErr = a.Run(
ctx,
inputMsgs,
agent.WithCheckpointer(checkpointer, runID),
)
}
}
}
// Heartbeat loss: another worker may have taken over. Do not commit
// any status — stale recovery will handle the row. Surface the cause
// so the worker kit logs and traces a failure for this attempt.
if cause := context.Cause(ctx); errors.Is(cause, ErrLeaseLost) || errors.Is(cause, ErrHeartbeatFailed) {
h.logger.WarnCtx(
context.WithoutCancel(ctx),
"agent run stopped after heartbeat failure; leaving status for stale recovery",
log.String("run_id", runID),
log.Error(cause),
)
return cause
}
// Infrastructure-triggered suspension (graceful shutdown): leave the
// row as RUNNING so stale recovery resets it to PENDING on restart.
// The checkpoint was already saved by coreLoop before returning
// SuspendedError, so Restore will pick up where it left off. This
// is not a failure from the worker kit's perspective.
if runErr != nil {
if _, ok := errors.AsType[*agent.SuspendedError](runErr); ok {
h.logger.InfoCtx(
context.WithoutCancel(ctx),
"agent run suspended by infrastructure; leaving for stale recovery",
log.String("run_id", runID),
)
return nil
}
}
now := time.Now()
run.UpdatedAt = now
run.StartedAt = nil
run.LeaseExpiresAt = nil
if runErr == nil {
run.Status = coredata.AgentRunStatusCompleted
if result != nil {
data, err := json.Marshal(result)
if err != nil {
h.logger.ErrorCtx(ctx, "cannot marshal agent run result", log.Error(err))
runErr = fmt.Errorf("cannot marshal agent run result: %w", err)
} else {
run.Result = data
}
}
}
if runErr != nil {
run.Status = coredata.AgentRunStatusFailed
run.Result = nil
h.logger.ErrorCtx(
context.WithoutCancel(ctx),
"agent run failed",
log.String("run_id", runID),
log.Error(runErr),
)
msg := sanitizeError(runErr)
run.ErrorMessage = &msg
}
commitCtx := context.WithoutCancel(ctx)
if err := h.pg.WithTx(
commitCtx,
func(ctx context.Context, tx pg.Tx) error {
rowsAffected, err := coredata.CommitAgentRunResult(ctx, tx, run, leaseGeneration)
if err != nil {
return err
}
if rowsAffected == 0 {
return ErrLeaseLost
}
if run.Status == coredata.AgentRunStatusCompleted {
if err := run.ClearCheckpoint(ctx, tx, coredata.NewNoScope()); err != nil {
return err
}
}
return nil
},
); err != nil {
if errors.Is(err, ErrLeaseLost) {
h.logger.WarnCtx(
commitCtx,
"agent run lost lease before commit; discarding stale completion",
log.String("run_id", runID),
)
return nil
}
h.logger.ErrorCtx(commitCtx, "cannot commit agent run status", log.Error(err))
return fmt.Errorf("cannot commit agent run status: %w", err)
}
return runErr
}

111
pkg/agentrun/internal_test.go Normal file
View File

@@ -0,0 +1,111 @@
// 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 (
"errors"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestSanitizeError(t *testing.T) {
t.Parallel()
t.Run(
"short message unchanged",
func(t *testing.T) {
t.Parallel()
err := errors.New("short")
assert.Equal(t, "short", sanitizeError(err))
},
)
t.Run(
"boundary length unchanged",
func(t *testing.T) {
t.Parallel()
msg := strings.Repeat("a", errorMessageMaxLen)
assert.Equal(t, msg, sanitizeError(errors.New(msg)))
},
)
t.Run(
"long utf8 message is rune safe and suffixed",
func(t *testing.T) {
t.Parallel()
msg := strings.Repeat("é", errorMessageMaxLen)
sanitized := sanitizeError(errors.New(msg))
assert.True(t, strings.HasSuffix(sanitized, "…"))
assert.True(t, len(sanitized) <= errorMessageMaxLen+len("…"))
assert.True(t, strings.HasPrefix(msg, strings.TrimSuffix(sanitized, "…")))
},
)
}
func TestWorkerOptions(t *testing.T) {
t.Parallel()
t.Run(
"interval updates only when positive",
func(t *testing.T) {
t.Parallel()
cfg := workerConfig{interval: 3 * time.Second}
WithWorkerInterval(0)(&cfg)
assert.Equal(t, 3*time.Second, cfg.interval)
WithWorkerInterval(7 * time.Second)(&cfg)
assert.Equal(t, 7*time.Second, cfg.interval)
},
)
t.Run(
"lease duration updates only when positive",
func(t *testing.T) {
t.Parallel()
cfg := workerConfig{leaseDuration: 5 * time.Second}
WithWorkerLeaseDuration(-1)(&cfg)
assert.Equal(t, 5*time.Second, cfg.leaseDuration)
WithWorkerLeaseDuration(12 * time.Second)(&cfg)
assert.Equal(t, 12*time.Second, cfg.leaseDuration)
},
)
t.Run(
"max concurrency updates only when positive",
func(t *testing.T) {
t.Parallel()
cfg := workerConfig{maxConcurrency: 2}
WithWorkerMaxConcurrency(0)(&cfg)
assert.Equal(t, 2, cfg.maxConcurrency)
WithWorkerMaxConcurrency(9)(&cfg)
assert.Equal(t, 9, cfg.maxConcurrency)
},
)
}

114
pkg/agentrun/service.go Normal file
View File

@@ -0,0 +1,114 @@
// 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"
"fmt"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/page"
)
type Service struct {
pg *pg.Client
}
func NewService(pgClient *pg.Client) *Service {
return &Service{pg: pgClient}
}
func (s *Service) Get(
ctx context.Context,
scope coredata.Scoper,
agentRunID gid.GID,
) (*coredata.AgentRun, error) {
run := &coredata.AgentRun{}
err := s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := run.LoadByID(ctx, conn, scope, agentRunID); err != nil {
return fmt.Errorf("cannot load agent run: %w", err)
}
return nil
},
)
if err != nil {
return nil, err
}
return run, nil
}
func (s *Service) ListForOrganizationID(
ctx context.Context,
scope coredata.Scoper,
organizationID gid.GID,
cursor *page.Cursor[coredata.AgentRunOrderField],
) (*page.Page[*coredata.AgentRun, coredata.AgentRunOrderField], error) {
var runs coredata.AgentRuns
err := s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
organization := &coredata.Organization{}
if err := organization.LoadByID(ctx, conn, scope, organizationID); err != nil {
return fmt.Errorf("cannot load organization: %w", err)
}
if err := runs.LoadByOrganizationID(ctx, conn, scope, organization.ID, cursor); err != nil {
return fmt.Errorf("cannot load agent runs: %w", err)
}
return nil
},
)
if err != nil {
return nil, err
}
return page.NewPage(runs, cursor), nil
}
func (s *Service) CountForOrganizationID(
ctx context.Context,
scope coredata.Scoper,
organizationID gid.GID,
) (int, error) {
var count int
err := s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) (err error) {
runs := &coredata.AgentRuns{}
count, err = runs.CountByOrganizationID(ctx, conn, scope, organizationID)
if err != nil {
return fmt.Errorf("cannot count agent runs: %w", err)
}
return nil
},
)
if err != nil {
return 0, err
}
return count, nil
}

95
pkg/agentrun/service_test.go Normal file
View File

@@ -0,0 +1,95 @@
// 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_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.probo.inc/probo/pkg/agentrun"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/page"
)
func TestService_Get(t *testing.T) {
client := pgClient(t)
svc := agentrun.NewService(client)
run := insertPendingRun(
t,
client,
"service-get-agent",
nil,
)
got, err := svc.Get(context.Background(), coredata.NewNoScope(), run.ID)
require.NoError(t, err)
require.NotNil(t, got)
assert.Equal(t, run.ID, got.ID)
missingID := gid.New(run.ID.TenantID(), coredata.AgentRunEntityType)
_, err = svc.Get(context.Background(), coredata.NewNoScope(), missingID)
require.Error(t, err)
assert.ErrorIs(t, err, coredata.ErrResourceNotFound)
}
func TestService_ListForOrganizationID(t *testing.T) {
client := pgClient(t)
svc := agentrun.NewService(client)
orgID := insertTestOrganization(t, client)
runA := insertPendingRunInOrg(t, client, orgID, "service-list-agent-a", nil)
runB := insertPendingRunInOrg(t, client, orgID, "service-list-agent-b", nil)
cursor := page.NewCursor(
10,
nil,
page.Head,
page.OrderBy[coredata.AgentRunOrderField]{
Field: coredata.AgentRunOrderFieldCreatedAt,
Direction: page.OrderDirectionDesc,
},
)
got, err := svc.ListForOrganizationID(context.Background(), coredata.NewNoScope(), orgID, cursor)
require.NoError(t, err)
require.NotNil(t, got)
ids := make(map[gid.GID]bool)
for _, run := range got.Data {
ids[run.ID] = true
}
assert.True(t, ids[runA.ID])
assert.True(t, ids[runB.ID])
}
func TestService_CountForOrganizationID(t *testing.T) {
client := pgClient(t)
svc := agentrun.NewService(client)
orgID := insertTestOrganization(t, client)
_ = insertPendingRunInOrg(t, client, orgID, "service-count-agent-a", nil)
_ = insertPendingRunInOrg(t, client, orgID, "service-count-agent-b", nil)
_ = insertPendingRunInOrg(t, client, orgID, "service-count-agent-c", nil)
count, err := svc.CountForOrganizationID(context.Background(), coredata.NewNoScope(), orgID)
require.NoError(t, err)
assert.Equal(t, 3, count)
}

121
pkg/agentrun/worker.go Normal file
View File

@@ -0,0 +1,121 @@
// 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)
}