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
+4
View File
@@ -13,6 +13,8 @@ LOG_LEVEL=info
# Directory where artifact bytes are stored.
COORDINATOR_STORAGE_DIR=./data
# Upper bound on an uploaded dataset or artifact body (bytes). Default 1 GiB.
MAX_UPLOAD_BYTES=1073741824
# Optional tuning (defaults shown).
DB_MAX_CONNS=10
@@ -22,3 +24,5 @@ REQUEST_TIMEOUT=15s
LEASE_DURATION=2m
DEFAULT_MAX_ATTEMPTS=3
REAPER_INTERVAL=30s
# A worker silent longer than this is marked offline by the reaper.
WORKER_OFFLINE_AFTER=1m
+21 -10
View File
@@ -72,7 +72,7 @@ func run() error {
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk),
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration),
RenewLease: usecase.NewRenewLease(taskRepo, tx, clk, cfg.LeaseDuration),
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
@@ -81,19 +81,30 @@ func run() error {
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
}
// Background workers are tracked so shutdown can wait for them. Without
// this the process would exit while the reaper sat mid-UPDATE, and the
// deferred pool.Close() would pull connections out from under it.
// Background reapers are tracked so shutdown can wait for them. Without this
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull
// connections out from under them.
expireLeases := usecase.NewExpireLeases(taskRepo, clk)
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
infra.RunReaper(ctx, log, usecase.NewExpireLeases(taskRepo, clk), cfg.ReaperInterval)
}()
for _, r := range []struct {
name string
fn func(context.Context) (int64, error)
}{
{"reaper requeued expired leases", expireLeases.Execute},
{"reaper marked workers offline", markOffline.Execute},
} {
wg.Add(1)
go func(name string, fn func(context.Context) (int64, error)) {
defer wg.Done()
infra.RunPeriodic(ctx, log, name, cfg.ReaperInterval, fn)
}(r.name, r.fn)
}
// pool.Ping backs /health: readiness means the database answers, not just
// that the process is alive.
api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, pool.Ping)
api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, pool.Ping)
err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token))
// Shutdown order matters, and defers alone cannot express it (they run
+11 -3
View File
@@ -15,6 +15,7 @@ type TaskStatus string
const (
TaskPending TaskStatus = "pending"
TaskLeased TaskStatus = "leased"
TaskRunning TaskStatus = "running"
TaskCompleted TaskStatus = "completed"
TaskFailed TaskStatus = "failed"
TaskCancelled TaskStatus = "cancelled"
@@ -146,7 +147,9 @@ func (t *Task) AsClaimed() ClaimedTask {
// verifyLease is the guard every worker-driven transition shares: the caller
// must own the lease and reference the attempt it was granted.
func (t *Task) verifyLease(worker string, attempt int) error {
if t.Status != TaskLeased {
// A task is worker-owned while leased or running: the first heartbeat moves
// it from leased to running, but ownership rules are identical for both.
if t.Status != TaskLeased && t.Status != TaskRunning {
return ErrTaskNotLeased
}
if t.LeaseOwner == nil || *t.LeaseOwner != worker {
@@ -158,12 +161,16 @@ func (t *Task) verifyLease(worker string, attempt int) error {
return nil
}
// RenewLease extends the lease of the worker that holds it.
// RenewLease extends the lease of the worker that holds it. The first heartbeat
// also acknowledges start, moving the task from leased to running.
func (t *Task) RenewLease(worker string, attempt int, until time.Time) error {
if err := t.verifyLease(worker, attempt); err != nil {
return err
}
t.LeaseExpiresAt = &until
if t.Status == TaskLeased {
t.Status = TaskRunning
}
t.Version++
return nil
}
@@ -228,7 +235,8 @@ func (t *Task) Fail(worker string, attempt int, code, message string, retryable
// ExpireLease is applied by the reaper when a lease elapses without a
// heartbeat: requeue while attempts remain, otherwise fail terminally.
func (t *Task) ExpireLease(now time.Time) {
if t.Status != TaskLeased {
// Both a leased and a running task can go silent and must be reclaimed.
if t.Status != TaskLeased && t.Status != TaskRunning {
return
}
t.LeaseOwner = nil
+36
View File
@@ -168,6 +168,42 @@ func TestExpireLeaseIgnoresUnleasedTasks(t *testing.T) {
}
}
func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) {
task := leasedTask(1, 3)
until := testLater.Add(time.Hour)
if err := task.RenewLease(testWorker, 1, until); err != nil {
t.Fatal(err)
}
if task.Status != TaskRunning {
t.Errorf("status = %q, want running after first heartbeat", task.Status)
}
// A second heartbeat keeps it running.
if err := task.RenewLease(testWorker, 1, until); err != nil {
t.Fatal(err)
}
if task.Status != TaskRunning {
t.Errorf("status = %q, want running", task.Status)
}
}
func TestRunningTaskCanBeCompletedAndExpired(t *testing.T) {
// Complete works from running.
task := leasedTask(1, 3)
_ = task.RenewLease(testWorker, 1, testLater) // -> running
if err := task.CompleteWith(testResult, nil, testWorker, 1, testNow); err != nil {
t.Errorf("complete from running: %v", err)
}
// Expire reclaims a running task too.
task2 := leasedTask(1, 3)
_ = task2.RenewLease(testWorker, 1, testLater) // -> running
task2.ExpireLease(testNow)
if task2.Status != TaskPending {
t.Errorf("status = %q, want pending after a running lease expires", task2.Status)
}
}
func TestRenewLeaseExtendsOnlyForHolder(t *testing.T) {
task := leasedTask(1, 3)
until := testLater.Add(time.Hour)
+24
View File
@@ -31,6 +31,8 @@ type Config struct {
LogFile string
// Directory where artifact bytes are stored.
StorageDir string
// Upper bound on an uploaded dataset or artifact body, in bytes.
MaxUploadBytes int64
// Connection pool upper bound.
DBMaxConns int32
@@ -48,6 +50,8 @@ type Config struct {
DefaultMaxAttempts int
// How often the background lease-reaper runs.
ReaperInterval time.Duration
// A worker silent for longer than this is marked offline by the reaper.
WorkerOfflineAfter time.Duration
}
// Load reads the environment and fails fast on anything required-but-missing
@@ -76,6 +80,7 @@ func LoadConfig() (Config, error) {
LogLevel: getEnv("LOG_LEVEL", "info"),
LogFile: os.Getenv("LOG_FILE"),
StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"),
MaxUploadBytes: 1 << 30, // 1 GiB
DBMaxConns: 10,
DBConnectTimeout: 30 * time.Second,
RequestTimeout: 15 * time.Second,
@@ -83,6 +88,7 @@ func LoadConfig() (Config, error) {
LeaseDuration: 2 * time.Minute,
DefaultMaxAttempts: 3,
ReaperInterval: 30 * time.Second,
WorkerOfflineAfter: 1 * time.Minute,
}
if cfg.DatabaseURL == "" {
@@ -96,6 +102,9 @@ func LoadConfig() (Config, error) {
if cfg.DBConnectTimeout, err = getEnvDuration("DB_CONNECT_TIMEOUT", cfg.DBConnectTimeout); err != nil {
return Config{}, err
}
if cfg.MaxUploadBytes, err = getEnvInt64("MAX_UPLOAD_BYTES", cfg.MaxUploadBytes); err != nil {
return Config{}, err
}
if cfg.RequestTimeout, err = getEnvDuration("REQUEST_TIMEOUT", cfg.RequestTimeout); err != nil {
return Config{}, err
}
@@ -108,6 +117,9 @@ func LoadConfig() (Config, error) {
if cfg.ReaperInterval, err = getEnvDuration("REAPER_INTERVAL", cfg.ReaperInterval); err != nil {
return Config{}, err
}
if cfg.WorkerOfflineAfter, err = getEnvDuration("WORKER_OFFLINE_AFTER", cfg.WorkerOfflineAfter); err != nil {
return Config{}, err
}
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
return Config{}, err
}
@@ -147,6 +159,18 @@ func getEnvInt32(key string, def int32) (int32, error) {
return int32(n), nil
}
func getEnvInt64(key string, def int64) (int64, error) {
v := os.Getenv(key)
if v == "" {
return def, nil
}
n, err := strconv.ParseInt(v, 10, 64)
if err != nil {
return 0, fmt.Errorf("%s: %w", key, err)
}
return n, nil
}
func getEnvDuration(key string, def time.Duration) (time.Duration, error) {
v := os.Getenv(key)
if v == "" {
+10 -8
View File
@@ -8,8 +8,6 @@ import (
"log/slog"
"net/http"
"time"
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
)
const shutdownGrace = 15 * time.Second
@@ -48,7 +46,13 @@ func RunServer(ctx context.Context, log *slog.Logger, addr string, handler http.
// RunReaper periodically reclaims tasks whose lease elapsed, so a worker that
// died without a heartbeat cannot strand its task in 'leased' forever.
func RunReaper(ctx context.Context, log *slog.Logger, uc *usecase.ExpireLeases, interval time.Duration) {
// RunPeriodic invokes fn on an interval until ctx is done, logging how many rows
// each tick affected. It backs the background reapers (expired leases, offline
// workers) — each is a set-based UPDATE that is safe to run repeatedly and
// concurrently across coordinators.
func RunPeriodic(ctx context.Context, log *slog.Logger, name string, interval time.Duration,
fn func(context.Context) (int64, error)) {
t := time.NewTicker(interval)
defer t.Stop()
@@ -57,15 +61,13 @@ func RunReaper(ctx context.Context, log *slog.Logger, uc *usecase.ExpireLeases,
case <-ctx.Done():
return
case <-t.C:
n, err := uc.Execute(ctx)
n, err := fn(ctx)
if err != nil {
// Demoted to debug while the repository is still a stub;
// raise to Warn once phase 5 lands.
log.Debug("reaper skipped", "err", err)
log.Debug(name+" skipped", "err", err)
continue
}
if n > 0 {
log.Info("reaper requeued expired leases", "count", n)
log.Info(name, "count", n)
}
}
}
+23
View File
@@ -233,6 +233,29 @@ func (r *WorkerRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Worker, err
return &cp, nil
}
func (r *WorkerRepo) Touch(ctx context.Context, id uuid.UUID, at time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
if w, ok := r.workers[id]; ok {
w.LastHeartbeatAt = at
w.Status = domain.WorkerOnline
}
return nil
}
func (r *WorkerRepo) MarkStaleOffline(ctx context.Context, cutoff time.Time) (int64, error) {
r.mu.Lock()
defer r.mu.Unlock()
var n int64
for _, w := range r.workers {
if w.Status != domain.WorkerOffline && w.LastHeartbeatAt.Before(cutoff) {
w.Status = domain.WorkerOffline
n++
}
}
return n, nil
}
// --- ArtifactRepo --------------------------------------------------------
type ArtifactRepo struct {
@@ -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
@@ -15,8 +15,12 @@ func writeJSON(w http.ResponseWriter, status int, v any) {
_ = json.NewEncoder(w).Encode(v)
}
// maxJSONBody caps a JSON request body. The DTOs are tiny; anything larger is a
// mistake or an attack, and must not be read into memory unbounded.
const maxJSONBody = 1 << 20 // 1 MiB
func decodeJSON(r *http.Request, dst any) error {
dec := json.NewDecoder(r.Body)
dec := json.NewDecoder(http.MaxBytesReader(nil, r.Body, maxJSONBody))
// Reject unknown fields: silently ignoring a misspelled "worker_ID" would
// surface later as a baffling validation failure.
dec.DisallowUnknownFields()
@@ -186,6 +186,7 @@ const defaultChunkRows = 1000
// chunker. The text fields MUST precede the file part: the file is streamed, not
// buffered, so by the time it arrives the other fields are already parsed.
func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
r.Body = http.MaxBytesReader(w, r.Body, s.maxUploadBytes)
mr, err := r.MultipartReader()
if err != nil {
s.writeError(w, r, domain.ErrInvalidInput)
@@ -292,6 +293,7 @@ func (s *Server) handleUploadArtifact(w http.ResponseWriter, r *http.Request) {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
r.Body = http.MaxBytesReader(w, r.Body, s.maxUploadBytes)
art, err := s.uc.UploadArtifact.Execute(r.Context(), usecase.UploadArtifactInput{
TaskID: taskID,
@@ -34,18 +34,20 @@ type Server struct {
log *slog.Logger
requestTimeout time.Duration
heartbeatInterval time.Duration
maxUploadBytes int64
// ready probes downstream dependencies (the database) for /health. Kept as
// a func so the transport layer never imports pgx.
ready func(context.Context) error
}
func NewServer(uc UseCases, log *slog.Logger, requestTimeout, heartbeatInterval time.Duration,
ready func(context.Context) error) *Server {
maxUploadBytes int64, ready func(context.Context) error) *Server {
return &Server{
uc: uc,
log: log,
requestTimeout: requestTimeout,
heartbeatInterval: heartbeatInterval,
maxUploadBytes: maxUploadBytes,
ready: ready,
}
}
@@ -42,7 +42,7 @@ func newEnv(t *testing.T, ready func(context.Context) error) *env {
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
RenewLease: usecase.NewRenewLease(tasks, tx, clk, lease),
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
@@ -50,7 +50,7 @@ func newEnv(t *testing.T, ready func(context.Context) error) *env {
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
}
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, ready)
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready)
ts := httptest.NewServer(srv.Handler(token))
t.Cleanup(ts.Close)
return &env{ts: ts, blobs: blobs}
+4 -3
View File
@@ -143,9 +143,10 @@ func progressFrom(job domain.Job, counts map[domain.TaskStatus]int) domain.JobPr
p := domain.JobProgress{
Job: job,
Pending: counts[domain.TaskPending],
Leased: counts[domain.TaskLeased],
Done: counts[domain.TaskCompleted],
Failed: counts[domain.TaskFailed],
// Leased and running are both "in flight" for progress purposes.
Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning],
Done: counts[domain.TaskCompleted],
Failed: counts[domain.TaskFailed],
}
for _, n := range counts {
p.Total += n
+6
View File
@@ -72,6 +72,12 @@ type JobRepository interface {
type WorkerRepository interface {
Insert(ctx context.Context, w *domain.Worker) error
Get(ctx context.Context, id uuid.UUID) (*domain.Worker, error)
// Touch records liveness for a heartbeating worker, marking it online. A
// no-op for an id that is not a registered worker.
Touch(ctx context.Context, id uuid.UUID, at time.Time) error
// MarkStaleOffline flips every worker last seen before cutoff to offline and
// reports how many changed.
MarkStaleOffline(ctx context.Context, cutoff time.Time) (int64, error)
}
// ArtifactRepository persists artifact metadata. The bytes live in a BlobStore;
+10 -2
View File
@@ -69,13 +69,15 @@ func (uc *ClaimTask) Execute(ctx context.Context, in ClaimTaskInput) (*domain.Cl
type RenewLease struct {
tasks TaskRepository
workers WorkerRepository
tx TxManager
clock Clock
leaseDuration time.Duration
}
func NewRenewLease(tasks TaskRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *RenewLease {
return &RenewLease{tasks: tasks, tx: tx, clock: clock, leaseDuration: leaseDuration}
func NewRenewLease(tasks TaskRepository, workers WorkerRepository, tx TxManager,
clock Clock, leaseDuration time.Duration) *RenewLease {
return &RenewLease{tasks: tasks, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration}
}
// Execute is a read-modify-write, so it runs inside a transaction with the row
@@ -101,6 +103,12 @@ func (uc *RenewLease) Execute(ctx context.Context, in RenewLeaseInput) (*domain.
if err != nil {
return nil, err
}
// Best-effort worker liveness, outside the task transaction so it can never
// fail the heartbeat. Only registered workers (a UUID worker_id) are tracked.
if id, perr := uuid.Parse(in.WorkerID); perr == nil {
_ = uc.workers.Touch(ctx, id, uc.clock.Now())
}
return &claimed, nil
}
+53 -1
View File
@@ -57,7 +57,7 @@ func newHarness() *harness {
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk)
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk)
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease)
h.renew = usecase.NewRenewLease(h.tasks, tx, h.clk, lease)
h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease)
h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
@@ -172,6 +172,26 @@ func TestRenewExtendsForHolder(t *testing.T) {
}
}
func TestHeartbeatThenCompleteViaRunning(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "w", 1)
taskID, attempt := h.leaseOne(t, "w1", "w")
// Heartbeat moves the task to running; completion must still work from there.
if _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt}); err != nil {
t.Fatalf("heartbeat: %v", err)
}
artID := h.uploadResult(t, taskID, "w1", attempt)
if _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt, ResultArtifactID: artID,
}); err != nil {
t.Fatalf("complete after heartbeat: %v", err)
}
if prog, _ := h.status.Execute(ctx, jobID); prog.DeriveStatus() != domain.JobCompleted {
t.Errorf("job status = %q, want completed", prog.DeriveStatus())
}
}
func TestRenewRejectsForeignWorker(t *testing.T) {
h := newHarness()
h.seedJob(t, "w", 1)
@@ -285,6 +305,38 @@ func TestRegisterWorkerPersists(t *testing.T) {
}
}
func TestHeartbeatTracksWorkerLivenessAndReaperMarksOffline(t *testing.T) {
h := newHarness()
w, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{Name: "lab", Capabilities: []string{"w"}})
if err != nil {
t.Fatal(err)
}
wid := w.ID.String() // a registered worker heartbeats with its UUID
h.seedJob(t, "w", 1)
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: wid, Workloads: []string{"w"}})
if err != nil || c == nil {
t.Fatalf("claim: %v", err)
}
if _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: c.TaskID, WorkerID: wid, Attempt: c.Attempt}); err != nil {
t.Fatalf("heartbeat: %v", err)
}
if got, _ := h.work.Get(ctx, w.ID); got.Status != domain.WorkerOnline {
t.Errorf("worker status = %q, want online after heartbeat", got.Status)
}
// Go silent past the threshold; the reaper marks it offline.
offline := usecase.NewMarkWorkersOffline(h.work, h.clk, 30*time.Second)
h.clk.Advance(time.Minute)
n, err := offline.Execute(ctx)
if err != nil || n != 1 {
t.Fatalf("reaper marked %d offline (err %v), want 1", n, err)
}
if got, _ := h.work.Get(ctx, w.ID); got.Status != domain.WorkerOffline {
t.Errorf("worker status = %q, want offline after reaper", got.Status)
}
}
func TestRegisterWorkerRejectsNoCapabilities(t *testing.T) {
h := newHarness()
if _, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{Name: "lab"}); !errors.Is(err, domain.ErrInvalidInput) {
+17
View File
@@ -2,6 +2,7 @@ package usecase
import (
"context"
"time"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
)
@@ -26,3 +27,19 @@ func (uc *RegisterWorker) Execute(ctx context.Context, in RegisterWorkerInput) (
}
return w, nil
}
// MarkWorkersOffline is the liveness reaper: workers that stopped heartbeating
// longer ago than `after` are flipped to offline.
type MarkWorkersOffline struct {
workers WorkerRepository
clk Clock
after time.Duration
}
func NewMarkWorkersOffline(workers WorkerRepository, clk Clock, after time.Duration) *MarkWorkersOffline {
return &MarkWorkersOffline{workers: workers, clk: clk, after: after}
}
func (uc *MarkWorkersOffline) Execute(ctx context.Context) (int64, error) {
return uc.workers.MarkStaleOffline(ctx, uc.clk.Now().Add(-uc.after))
}
@@ -0,0 +1,4 @@
-- PostgreSQL cannot drop a single enum value without recreating the type and
-- rewriting every dependent column. Leaving 'running' in place is harmless: no
-- code writes it after the down of 0007 restores the leased-only transitions.
SELECT 1;
@@ -0,0 +1,5 @@
-- 'running' means the worker has acknowledged start via its first heartbeat.
-- Kept in its own migration, without an explicit transaction: an enum value
-- added in a transaction cannot be USED in that same transaction, and the next
-- migration references it.
ALTER TYPE task_status ADD VALUE IF NOT EXISTS 'running';
@@ -0,0 +1,8 @@
BEGIN;
ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_leased_owner;
ALTER TABLE tasks ADD CONSTRAINT ck_tasks_leased_owner CHECK (
status <> 'leased' OR (lease_owner IS NOT NULL AND lease_expires_at IS NOT NULL)
);
COMMIT;
@@ -0,0 +1,10 @@
BEGIN;
-- A running task holds a lease just like a leased one, so the lease-integrity
-- check must cover both states.
ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_leased_owner;
ALTER TABLE tasks ADD CONSTRAINT ck_tasks_leased_owner CHECK (
status NOT IN ('leased','running') OR (lease_owner IS NOT NULL AND lease_expires_at IS NOT NULL)
);
COMMIT;