diff --git a/coordinator/.env.example b/coordinator/.env.example index 1e8f90f..d20a617 100644 --- a/coordinator/.env.example +++ b/coordinator/.env.example @@ -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 diff --git a/coordinator/cmd/coordinator/main.go b/coordinator/cmd/coordinator/main.go index 767ebeb..dd11a56 100644 --- a/coordinator/cmd/coordinator/main.go +++ b/coordinator/cmd/coordinator/main.go @@ -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 diff --git a/coordinator/internal/domain/task.go b/coordinator/internal/domain/task.go index 8908e44..4fadca2 100644 --- a/coordinator/internal/domain/task.go +++ b/coordinator/internal/domain/task.go @@ -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 diff --git a/coordinator/internal/domain/task_test.go b/coordinator/internal/domain/task_test.go index 7ccaccf..a8fa352 100644 --- a/coordinator/internal/domain/task_test.go +++ b/coordinator/internal/domain/task_test.go @@ -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) diff --git a/coordinator/internal/infra/config.go b/coordinator/internal/infra/config.go index b7aa308..8703a27 100644 --- a/coordinator/internal/infra/config.go +++ b/coordinator/internal/infra/config.go @@ -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 == "" { diff --git a/coordinator/internal/infra/server.go b/coordinator/internal/infra/server.go index c017726..9858c19 100644 --- a/coordinator/internal/infra/server.go +++ b/coordinator/internal/infra/server.go @@ -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) } } } diff --git a/coordinator/internal/memstore/memstore.go b/coordinator/internal/memstore/memstore.go index 71243ed..7916951 100644 --- a/coordinator/internal/memstore/memstore.go +++ b/coordinator/internal/memstore/memstore.go @@ -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 { diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go index 6ee522e..fc917d6 100644 --- a/coordinator/internal/storage/postgres/integration_test.go +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -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() diff --git a/coordinator/internal/storage/postgres/task_repo.go b/coordinator/internal/storage/postgres/task_repo.go index 827ab77..2a4480f 100644 --- a/coordinator/internal/storage/postgres/task_repo.go +++ b/coordinator/internal/storage/postgres/task_repo.go @@ -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 diff --git a/coordinator/internal/storage/postgres/worker_repo.go b/coordinator/internal/storage/postgres/worker_repo.go index 034a97f..1df545e 100644 --- a/coordinator/internal/storage/postgres/worker_repo.go +++ b/coordinator/internal/storage/postgres/worker_repo.go @@ -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 diff --git a/coordinator/internal/transport/http/errors.go b/coordinator/internal/transport/http/errors.go index 04e5ae7..f808ad7 100644 --- a/coordinator/internal/transport/http/errors.go +++ b/coordinator/internal/transport/http/errors.go @@ -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() diff --git a/coordinator/internal/transport/http/handlers.go b/coordinator/internal/transport/http/handlers.go index 1165105..cdd9385 100644 --- a/coordinator/internal/transport/http/handlers.go +++ b/coordinator/internal/transport/http/handlers.go @@ -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, diff --git a/coordinator/internal/transport/http/server.go b/coordinator/internal/transport/http/server.go index 109392c..cf7c304 100644 --- a/coordinator/internal/transport/http/server.go +++ b/coordinator/internal/transport/http/server.go @@ -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, } } diff --git a/coordinator/internal/transport/http/server_test.go b/coordinator/internal/transport/http/server_test.go index 89aef8a..0f00024 100644 --- a/coordinator/internal/transport/http/server_test.go +++ b/coordinator/internal/transport/http/server_test.go @@ -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} diff --git a/coordinator/internal/usecase/job.go b/coordinator/internal/usecase/job.go index 6552565..28e6c6e 100644 --- a/coordinator/internal/usecase/job.go +++ b/coordinator/internal/usecase/job.go @@ -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 diff --git a/coordinator/internal/usecase/ports.go b/coordinator/internal/usecase/ports.go index 1269e42..a85481b 100644 --- a/coordinator/internal/usecase/ports.go +++ b/coordinator/internal/usecase/ports.go @@ -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; diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go index 0451187..0c2e4f0 100644 --- a/coordinator/internal/usecase/task.go +++ b/coordinator/internal/usecase/task.go @@ -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 } diff --git a/coordinator/internal/usecase/usecase_test.go b/coordinator/internal/usecase/usecase_test.go index 9c3b458..254a8ef 100644 --- a/coordinator/internal/usecase/usecase_test.go +++ b/coordinator/internal/usecase/usecase_test.go @@ -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) { diff --git a/coordinator/internal/usecase/worker.go b/coordinator/internal/usecase/worker.go index 915879b..c8ccab2 100644 --- a/coordinator/internal/usecase/worker.go +++ b/coordinator/internal/usecase/worker.go @@ -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)) +} diff --git a/coordinator/migrations/0006_task_running_enum.down.sql b/coordinator/migrations/0006_task_running_enum.down.sql new file mode 100644 index 0000000..5144db0 --- /dev/null +++ b/coordinator/migrations/0006_task_running_enum.down.sql @@ -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; diff --git a/coordinator/migrations/0006_task_running_enum.up.sql b/coordinator/migrations/0006_task_running_enum.up.sql new file mode 100644 index 0000000..4216518 --- /dev/null +++ b/coordinator/migrations/0006_task_running_enum.up.sql @@ -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'; diff --git a/coordinator/migrations/0007_task_running_lease.down.sql b/coordinator/migrations/0007_task_running_lease.down.sql new file mode 100644 index 0000000..fdf4795 --- /dev/null +++ b/coordinator/migrations/0007_task_running_lease.down.sql @@ -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; diff --git a/coordinator/migrations/0007_task_running_lease.up.sql b/coordinator/migrations/0007_task_running_lease.up.sql new file mode 100644 index 0000000..5cf0cfa --- /dev/null +++ b/coordinator/migrations/0007_task_running_lease.up.sql @@ -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;