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:
Efremenko Arhip
2026-07-23 17:36:26 +03:00
parent 4fc3c69fdf
commit 3b41455b20
23 changed files with 332 additions and 32 deletions
@@ -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