feat(coordinator): running state, worker liveness, request-size limits
Polish pass hardening the queue and closing plan gaps. - Task state machine gains `running`: the first heartbeat moves a task from leased to running (migrations 0006/0007 add the enum value and extend the lease-integrity check). verifyLease, ExpireLease, the reaper SQL, and job progress all treat leased and running alike. - Worker liveness: a heartbeat from a registered worker (UUID worker_id) bumps its last_heartbeat_at online; a second background reaper marks workers offline after WORKER_OFFLINE_AFTER of silence (RunReaper generalized to RunPeriodic). - Request-size limits: JSON bodies capped at 1 MiB; dataset/artifact uploads capped at MAX_UPLOAD_BYTES (default 1 GiB) via http.MaxBytesReader. - Tests cover the running transition, liveness + offline reaper (unit over memstore and integration over Postgres).
This commit is contained in:
@@ -387,6 +387,47 @@ func TestWorkerRepoRoundTrip(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkerLivenessAndOfflineReaper(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
repo := NewWorkerRepo(pool)
|
||||
|
||||
w, err := domain.NewWorker("liveness", []string{"similarity_search"}, time.Now().UTC().Add(-time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := repo.Insert(ctx, w); err != nil {
|
||||
t.Fatalf("insert: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM workers WHERE id = $1`, w.ID) })
|
||||
|
||||
// A fresh heartbeat bumps it online.
|
||||
now := time.Now().UTC()
|
||||
if err := repo.Touch(ctx, w.ID, now); err != nil {
|
||||
t.Fatalf("touch: %v", err)
|
||||
}
|
||||
if got, _ := repo.Get(ctx, w.ID); got.Status != domain.WorkerOnline {
|
||||
t.Errorf("status = %q, want online after touch", got.Status)
|
||||
}
|
||||
|
||||
// Touching an unregistered id is a harmless no-op.
|
||||
if err := repo.Touch(ctx, uuid.New(), now); err != nil {
|
||||
t.Errorf("touch of unknown worker returned %v, want nil", err)
|
||||
}
|
||||
|
||||
// The reaper marks it offline once its heartbeat is older than the cutoff.
|
||||
n, err := repo.MarkStaleOffline(ctx, now.Add(time.Minute))
|
||||
if err != nil {
|
||||
t.Fatalf("mark offline: %v", err)
|
||||
}
|
||||
if n < 1 {
|
||||
t.Errorf("marked %d offline, want at least 1", n)
|
||||
}
|
||||
if got, _ := repo.Get(ctx, w.ID); got.Status != domain.WorkerOffline {
|
||||
t.Errorf("status = %q, want offline after reaper", got.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestArtifactRepoRoundTrip(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
|
||||
@@ -314,7 +314,7 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
|
||||
ELSE error_message END,
|
||||
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
|
||||
version = version + 1
|
||||
WHERE status = 'leased' AND lease_expires_at < $1`
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1`
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
var affected int64
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
sq "github.com/Masterminds/squirrel"
|
||||
"github.com/google/uuid"
|
||||
@@ -59,6 +60,37 @@ func (r *WorkerRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Worker, err
|
||||
return w, nil
|
||||
}
|
||||
|
||||
func (r *WorkerRepo) Touch(ctx context.Context, id uuid.UUID, at time.Time) error {
|
||||
sql, args, err := psql.Update("workers").
|
||||
SetMap(map[string]any{"last_heartbeat_at": at, "status": "online", "updated_at": at}).
|
||||
Where(sq.Eq{"id": id}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// A worker that never registered simply matches no row; that is not an error.
|
||||
if _, err := conn(ctx, r.pool).Exec(ctx, sql, args...); err != nil {
|
||||
return fmt.Errorf("touch worker: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *WorkerRepo) MarkStaleOffline(ctx context.Context, cutoff time.Time) (int64, error) {
|
||||
sql, args, err := psql.Update("workers").
|
||||
SetMap(map[string]any{"status": "offline", "updated_at": cutoff}).
|
||||
Where(sq.Lt{"last_heartbeat_at": cutoff}).
|
||||
Where(sq.NotEq{"status": "offline"}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("mark stale workers offline: %w", err)
|
||||
}
|
||||
return tag.RowsAffected(), nil
|
||||
}
|
||||
|
||||
func scanWorker(row pgx.Row) (*domain.Worker, error) {
|
||||
var (
|
||||
w domain.Worker
|
||||
|
||||
Reference in New Issue
Block a user