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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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;
|
||||
Reference in New Issue
Block a user