diff --git a/.github/workflows/coordinator.yml b/.github/workflows/coordinator.yml new file mode 100644 index 0000000..11a5607 --- /dev/null +++ b/.github/workflows/coordinator.yml @@ -0,0 +1,66 @@ +name: coordinator + +on: + push: + paths: + - "coordinator/**" + - ".github/workflows/coordinator.yml" + pull_request: + paths: + - "coordinator/**" + - ".github/workflows/coordinator.yml" + +defaults: + run: + working-directory: coordinator + +jobs: + test: + runs-on: ubuntu-latest + + services: + postgres: + image: postgres:16-alpine + env: + POSTGRES_USER: scimesh + POSTGRES_PASSWORD: scimesh + POSTGRES_DB: scimesh + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U scimesh" + --health-interval 5s + --health-timeout 3s + --health-retries 10 + + env: + TEST_DATABASE_URL: postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version-file: coordinator/go.mod + cache-dependency-path: coordinator/go.sum + + - name: go vet + run: go vet ./... + + - name: gofmt + run: test -z "$(gofmt -l .)" || (gofmt -l . && exit 1) + + - name: unit tests (race) + run: go test -race ./... + + - name: lint + run: go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.12.2 run --build-tags=integration ./... + + - name: install migrate CLI + run: go install -tags 'postgres' github.com/golang-migrate/migrate/v4/cmd/migrate@v4.17.1 + + - name: apply migrations + run: migrate -path migrations -database "$TEST_DATABASE_URL" up + + - name: integration tests + run: go test -tags=integration ./internal/storage/postgres/ -v diff --git a/coordinator/.dockerignore b/coordinator/.dockerignore new file mode 100644 index 0000000..a4d0d5c --- /dev/null +++ b/coordinator/.dockerignore @@ -0,0 +1,14 @@ +# Keep the build context small and never bake secrets or local state into an image. +.env +.git +.gitignore +*.md +Makefile +docker-compose.yml +Dockerfile +.dockerignore + +# Local build artifacts +/coordinator +/bin/ +*.out diff --git a/coordinator/.env.example b/coordinator/.env.example new file mode 100644 index 0000000..d20a617 --- /dev/null +++ b/coordinator/.env.example @@ -0,0 +1,28 @@ +# Copy to .env and adjust. All settings are read from the environment. + +COORDINATOR_ADDR=:8080 +DATABASE_URL=postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable + +# Shared bearer token every worker must present. Leave empty to disable auth (dev only). +WORKER_AUTH_TOKEN=change-me + +# Logging. LOG_LEVEL: debug|info|warn|error. LOG_FILE empty = stdout only; +# set a path to also write a size-rotated file (kept across restarts). +LOG_LEVEL=info +# LOG_FILE=./logs/coordinator.log + +# 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 +# How long to keep retrying the initial DB connection while Postgres boots. +DB_CONNECT_TIMEOUT=30s +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/.gitignore b/coordinator/.gitignore new file mode 100644 index 0000000..c8c4f7b --- /dev/null +++ b/coordinator/.gitignore @@ -0,0 +1,6 @@ +/coordinator +/bin/ +.env +*.out +/logs/ +/data/ diff --git a/coordinator/.golangci.yml b/coordinator/.golangci.yml new file mode 100644 index 0000000..428bfed --- /dev/null +++ b/coordinator/.golangci.yml @@ -0,0 +1,54 @@ +version: "2" + +run: + timeout: 3m + +linters: + # "standard" = errcheck, govet, ineffassign, staticcheck, unused. + default: standard + enable: + # Catches `err == ErrFoo` where errors.Is is required. Directly relevant + # here: domain exposes sentinel errors that use cases may wrap with %w. + - errorlint + # Returning nil after checking a non-nil error — a silent bug factory. + - nilerr + # http.Get/Do without a context: every outbound call must be cancellable. + - noctx + # Unclosed response bodies leak connections. + - bodyclose + # Common security mistakes (weak crypto, unhandled file perms). + - gosec + # Style and naming consistency. + - revive + - misspell + - unconvert + + settings: + errcheck: + # Deferred Close/Rollback are intentionally ignored in a few places + # (rollback after commit is a documented no-op). + check-type-assertions: true + revive: + rules: + - name: exported + disabled: true # internal packages need no exported-symbol comments + gosec: + excludes: + - G404 # math/rand is fine for jitter; nothing here is security-sensitive + + exclusions: + rules: + # Tests may skip error checks and use long literals freely. + - path: _test\.go + linters: + - errcheck + - gosec + +formatters: + enable: + - gofmt + - goimports + settings: + goimports: + local-prefixes: + - github.com/emil28092005/SciMesh/coordinator diff --git a/coordinator/ARCHITECTURE.md b/coordinator/ARCHITECTURE.md new file mode 100644 index 0000000..77b6976 --- /dev/null +++ b/coordinator/ARCHITECTURE.md @@ -0,0 +1,144 @@ +# Архитектура координатора + +Карта кода. Читать сверху вниз: сначала «где что лежит», потом «как проходит +запрос», в конце — «куда добавлять новое». + +--- + +## 1. Четыре слоя + +``` + infra конфиг, пул БД, часы, HTTP-сервер, reaper ← драйверы + transport HTTP-хендлеры ← входящее: кто зовёт нас + storage репозитории на SQL ← исходящее: кого зовём мы + usecase операции + ПОРТЫ (интерфейсы) ← прикладные правила + domain Task, Job и их инварианты ← бизнес-правила + + ┌── transport ──┐ + domain ◄── usecase ◄┤ ├◄── infra + └── storage ────┘ +``` + +`transport` и `storage` — один и тот же слой (в книгах он зовётся «адаптеры»), +просто разделённый по направлению: транспорт принимает запросы снаружи, storage +обращается наружу сам. Так путь к файлу говорит о его роли, а не о категории. + +**Единственное правило:** зависимости идут только внутрь. `domain` не импортирует +ничего из проекта. `usecase` видит только `domain`. `transport` и `storage` не +знают друг о друге. + +Проверить в любой момент: + +```sh +go list -f '{{range .Imports}}{{.}}{{"\n"}}{{end}}' ./internal/domain | grep internal +# пусто = правило соблюдено +``` + +--- + +## 2. Где что лежит + +| Файл | Что внутри | Строк | +| --- | --- | --- | +| `domain/task.go` | `Task` и **все** переходы состояний: аренда, завершение, провал, истечение | ~245 | +| `domain/job.go` | `Job`, разбиение на чанки, вывод статуса из счётчиков задач | ~107 | +| `domain/errors.go` | Нарушения бизнес-правил (`ErrLeaseConflict`, `ErrStaleAttempt`, …) | ~18 | +| `usecase/ports.go` | **Порты**: `TaskRepository`, `JobRepository`, `TxManager`, `Clock` | ~79 | +| `usecase/task.go` | Операции над задачей: claim, renew, complete, fail, expire | ~200 | +| `usecase/job.go` | Операции над job: create, status, results, stitch | ~180 | +| `usecase/dto.go` | Входные структуры юзкейсов | ~51 | +| `transport/http/server.go` | Роутер и сборка middleware | ~60 | +| `transport/http/handlers.go` | По хендлеру на эндпоинт | ~180 | +| `transport/http/dto.go` | JSON-форматы запросов и ответов | ~118 | +| `transport/http/middleware.go` | request-ID, access-лог, bearer-авторизация | ~103 | +| `transport/http/errors.go` | Маппинг доменных ошибок в HTTP-коды | ~55 | +| `storage/postgres/task_repo.go` | SQL по задачам, включая атомарный claim | ~109 | +| `storage/postgres/job_repo.go` | SQL по job'ам | ~39 | +| `storage/postgres/tx.go` | `TxManager`: транзакция через контекст | ~65 | +| `infra/*.go` | Конфиг, пул, часы, сервер, reaper | ~240 | +| `cmd/coordinator/main.go` | **Composition root** — единственное место со всеми конкретными типами | ~73 | + +--- + +## 3. Трасса запроса: `POST /tasks/claim` + +Как воркер получает задачу. Четыре остановки, по одной на слой: + +``` + ① transport/http/handlers.go → handleClaim + разбирает JSON, отдаёт usecase.ClaimTaskInput + │ + ▼ + ② usecase/task.go → ClaimTask.Execute + сначала подчищает протухшие аренды, потом просит одну задачу + через ПОРТ TaskRepository (реализацию не знает) + │ + ▼ + ③ usecase/ports.go → TaskRepository.ClaimNext + контракт: «атомарно выдай одну задачу» + │ + ▼ + ④ storage/postgres/task_repo.go → claimNextSQL + SELECT ... FOR UPDATE SKIP LOCKED + UPDATE одним запросом +``` + +Обратно поднимается `*domain.Task`, юзкейс сужает его до `domain.ClaimedTask` +(воркеру не отдаём `version`, `max_attempts` и чужие ошибки), хендлер +превращает в JSON. Пустая очередь — это `nil, nil` на шаге ② и `204` на ①. + +**Трасса `POST /tasks/{id}/result`** такая же, но с одним отличием: решение +принимает **сущность**, а не юзкейс. + +``` + handlers.go → CompleteTask.Execute → tx.WithinTx( + GetForUpdate → task.CompleteWith(...) ←── ЗДЕСЬ правила + │ (чужая аренда? устаревший + Update ←─────────────┘ attempt? повтор того же + syncJobStatus манифеста?) + ) +``` + +--- + +## 4. Куда добавлять новое + +| Хочу… | Правлю | +| --- | --- | +| новое бизнес-правило (когда задачу можно повторить) | `domain/task.go` + тест рядом | +| новую операцию (отменить job) | `usecase/job.go` + порт в `ports.go`, если нужен новый запрос к БД | +| новый HTTP-эндпоинт | `transport/http/handlers.go` + маршрут в `server.go` + DTO в `dto.go` | +| новый SQL-запрос | `storage/postgres/*_repo.go` | +| новую настройку | `infra/config.go` + `.env.example` | +| поменять код ответа на ошибку | `transport/http/errors.go` | + +**Правило при сомнении:** если код можно описать фразой «когда X, то Y» без +упоминания HTTP, SQL и конфигов — это `domain`. Если он оркеструет несколько +шагов и транзакцию — `usecase`. Если знает про JSON — `transport`, про SQL — `storage`. + +--- + +## 5. Три вещи, которые надо понять один раз + +**Порты объявляет потребитель.** `TaskRepository` описан в `usecase/ports.go`, а +реализован в `storage/postgres`. Поэтому `usecase` не импортирует `storage` — +стрелка зависимости смотрит внутрь, хотя вызов на рантайме идёт наружу. + +**Транзакция едет в контексте.** `TxManager.WithinTx` кладёт `pgx.Tx` в контекст +по неэкспортируемому ключу; репозитории достают её через `conn(ctx, pool)`. +Благодаря этому юзкейс говорит «сделай это атомарно», ни разу не упомянув pgx. + +**Атомарный claim нельзя разложить на шаги.** `ClaimNext` — один SQL-запрос, +потому что `SELECT` + отдельный `UPDATE` вернул бы гонку, при которой одну +задачу выдают двум воркерам. Поэтому `ClaimTask.Execute` выглядит тонким: там +нечего оркестровать, вся гарантия — внутри запроса. + +--- + +## 6. Что уже работает, а что заглушка + +Работает: слои и проводка, роутинг, авторизация, access-лог, маппинг ошибок, +транзакции, graceful shutdown, миграции, **весь domain с 12 юнит-тестами без БД**. + +Заглушки (`ErrNotImplemented` → HTTP 501): методы репозиториев. SQL для двух +главных операций уже написан в `task_repo.go` — `claimNextSQL` и +`expireLeasesSQL`, осталось их подключить. diff --git a/coordinator/Dockerfile b/coordinator/Dockerfile new file mode 100644 index 0000000..fc28a60 --- /dev/null +++ b/coordinator/Dockerfile @@ -0,0 +1,52 @@ +# syntax=docker/dockerfile:1 +# +# Requires BuildKit (the RUN --mount cache lines below). Docker 23+ enables it +# by default when the buildx plugin is present; install `docker-buildx` if a +# build fails with "the --mount option requires BuildKit". + +# --- build stage ---------------------------------------------------------- +FROM golang:1.24-alpine AS build + +WORKDIR /src + +# Copy manifests first: this layer stays cached until dependencies actually +# change, so editing Go sources does not re-download the module graph. +COPY go.mod go.sum ./ +RUN --mount=type=cache,target=/go/pkg/mod go mod download + +COPY . . + +# The cache mounts persist the module cache and the compiler's build cache +# *across* builds, so a rebuild after a code edit recompiles only what changed +# instead of the whole dependency tree. +# +# CGO_ENABLED=0 produces a fully static binary, so the runtime image needs no +# libc. -trimpath strips local paths; -s -w drop the symbol table and DWARF. +RUN --mount=type=cache,target=/go/pkg/mod \ + --mount=type=cache,target=/root/.cache/go-build \ + CGO_ENABLED=0 GOOS=linux go build \ + -trimpath -ldflags="-s -w" \ + -o /out/coordinator ./cmd/coordinator + +# --- runtime stage -------------------------------------------------------- +FROM alpine:3.20 + +# ca-certificates for outbound TLS; wget backs the container healthcheck. +RUN apk add --no-cache ca-certificates wget \ + && adduser -D -H -u 10001 coordinator \ + # Pre-create the storage and log dirs owned by the non-root user. A named + # volume mounted here inherits this ownership from the image, so the process + # can write to it — a host bind mount, owned by root, cannot. + && mkdir -p /var/lib/scimesh/artifacts /var/log/scimesh \ + && chown -R coordinator:coordinator /var/lib/scimesh /var/log/scimesh + +COPY --from=build /out/coordinator /usr/local/bin/coordinator + +# Never run as root: a compromised process should not own the container. +USER coordinator + +EXPOSE 8080 + +# Exec form, not shell: the binary becomes PID 1 and receives SIGTERM directly, +# which is what its graceful shutdown depends on. +ENTRYPOINT ["/usr/local/bin/coordinator"] diff --git a/coordinator/Makefile b/coordinator/Makefile new file mode 100644 index 0000000..bdc1313 --- /dev/null +++ b/coordinator/Makefile @@ -0,0 +1,85 @@ +.PHONY: build run test test-integration vet lint tidy check migrate-up migrate-down up down down-clean logs ps rebuild psql smoke + +# --- build / run --------------------------------------------------------- +build: + go build ./... + +run: + go run ./cmd/coordinator + +test: + go test ./... + +# Needs a running PostgreSQL; the spec forbids mocks for these guarantees. +# make test-integration TEST_DATABASE_URL='postgres://...' +test-integration: + TEST_DATABASE_URL="$(TEST_DATABASE_URL)" go test -tags=integration ./... -v + +vet: + go vet ./... + +# One command that runs everything: unit tests + vet + lint, then brings up the +# stack and runs the integration suite and the end-to-end smoke test. +# Needs Docker. Hand this to a reviewer. +check: vet lint + go test -race ./... + docker compose up -d --build + @echo "waiting for the coordinator to be ready..." + @sleep 6 + TEST_DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable' \ + go test -tags=integration ./internal/storage/postgres/ -v + ./scripts/smoke.sh + @echo "\nall checks passed ✓" + +# Runs golangci-lint without installing it system-wide. Install it for speed: +# pacman -S golangci-lint (Arch) +LINT_VERSION := v2.12.2 +lint: + @command -v golangci-lint >/dev/null 2>&1 \ + && golangci-lint run --build-tags=integration ./... \ + || go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@$(LINT_VERSION) run --build-tags=integration ./... + +tidy: + go mod tidy + +# --- migrations ---------------------------------------------------------- +# Requires the golang-migrate CLI: +# go install -tags 'postgres' github.com/golang-migrate/migrate/v4/cmd/migrate@latest +# DATABASE_URL must be set, e.g.: +# export DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable' +migrate-up: + migrate -path migrations -database "$(DATABASE_URL)" up + +migrate-down: + migrate -path migrations -database "$(DATABASE_URL)" down 1 + +# --- docker -------------------------------------------------------------- +# `up` starts Postgres, applies migrations, then launches the coordinator. +up: + docker compose up -d --build + +down: + docker compose down + +# Also drops the database volume — use when the schema is beyond repair. +down-clean: + docker compose down -v + +logs: + docker compose logs -f coordinator + +ps: + docker compose ps + +rebuild: + docker compose up -d --build --force-recreate coordinator + +psql: + docker compose exec postgres psql -U scimesh -d scimesh + +# --- api ------------------------------------------------------------------ +# Exercises every endpoint against a running coordinator; exits non-zero on the +# first unexpected status. See also api/requests.http for clicking through them +# one at a time in an editor. +smoke: + ./scripts/smoke.sh diff --git a/coordinator/README.md b/coordinator/README.md new file mode 100644 index 0000000..c190806 --- /dev/null +++ b/coordinator/README.md @@ -0,0 +1,171 @@ +# SciMesh Coordinator + +Durable task-queue server for SciMesh, in Go on PostgreSQL. It owns all database +access; workers talk to it only over HTTP and never receive DB credentials. + +Built as a **modular monolith following Clean Architecture** — one binary, four +layers, dependencies pointing strictly inward. See +`docs/database-integration-task.md` and `docs/worker-daemon-task.md` in the repo +root for the full contract. + +## Layers + +``` + infra config, pgxpool, http.Server, clock ← frameworks & drivers + transport http handlers ← inbound: who calls us + storage sql repositories ← outbound: who we call + usecase business operations + PORTS ← application rules + domain Task, Job + their invariants ← enterprise rules + + ┌── transport ──┐ + domain ◄── usecase ◄┤ ├◄── infra + └── storage ────┘ +``` + +`transport` and `storage` are one layer — the "interface adapters" ring — split +by direction rather than by category, so a file's path tells you its role. + +The rule that matters: **source dependencies point only inward**. `domain` +imports nothing from this module; `usecase` sees only `domain`; `transport` and +`storage` know nothing of each other. Verify it at any time with: + +```sh +go list -f '{{range .Imports}}{{.}}{{"\n"}}{{end}}' ./internal/domain | grep internal # must be empty +``` + +## Layout + +``` +coordinator/ + cmd/coordinator/main.go # composition root: the only place with concrete types + internal/ + domain/ # entities + rules, no I/O + task.go Task, lease/complete/fail/expire transitions + job.go Job, chunk fan-out, status derivation + errors.go business-rule violations + usecase/ # one type per operation, dependencies injected + ports.go TaskRepository, JobRepository, TxManager, Clock + dto.go use-case boundary inputs + task.go claim, renew, complete, fail, expire + job.go create, status, results, stitch + transport/http/ # routing, DTOs, middleware, error mapping + storage/postgres/ # SQL behind the ports; TxManager via context + infra/ # config.go db.go clock.go server.go + migrations/ # golang-migrate SQL, run as an explicit command +``` + +A full map — file-by-file table, a request traced through every layer, and a +"where do I add X" guide — lives in [ARCHITECTURE.md](ARCHITECTURE.md). + +## Quickstart + +### With Docker (nothing to install but Docker) + +```sh +make up # Postgres → migrations → coordinator +curl localhost:8080/health +make logs # follow the coordinator +make down # stop (add down-clean to drop the DB volume) +``` + +`up` starts three services in order: Postgres waits until `pg_isready` passes, a +one-shot `migrate` container applies the schema and exits, and only then does the +coordinator start — so it never queries a database that has no tables. + +> **Needs BuildKit.** The Dockerfile uses `RUN --mount=type=cache` to reuse the +> Go module and compiler caches between builds. If the build fails with +> *"the --mount option requires BuildKit"*, install the buildx plugin — +> `pacman -S docker-buildx` on Arch, `apt install docker-buildx-plugin` on Debian. + +### Locally, against your own Postgres + +```sh +cp .env.example .env # then edit DATABASE_URL / WORKER_AUTH_TOKEN + # it is loaded automatically — no export needed + +make tidy # fetch deps (needs network once) +make migrate-up # apply schema (needs the migrate CLI) +make run # start the server +``` + +## Configuration + +Settings come from the environment. A `.env` file is loaded at startup via +`godotenv` as a local-dev convenience (override its path with `ENV_FILE`): + +- a missing `.env` is not an error — production injects real env vars; +- **real environment variables always win** over the file, so an orchestrator's + values are never shadowed by a stale `.env` baked into an image. + +See `.env.example`; only `DATABASE_URL` is required. + +## Endpoints + +| Method | Path | Purpose | +| ------ | ---------------------------------- | --------------------------------------------- | +| POST | `/workers/register` | Register a worker, get its id | +| POST | `/jobs` | Create job + tasks from chunk URIs | +| POST | `/jobs/upload` | Upload a dataset; coordinator chunks it | +| GET | `/jobs/{job_id}` | Aggregate job progress | +| POST | `/tasks/claim` | Atomically lease one task (`204` if none) | +| GET | `/tasks/{task_id}/input` | Download the task's input shard | +| POST | `/tasks/{task_id}/heartbeat` | Renew the caller's lease (→ `running`) | +| PUT | `/tasks/{task_id}/artifacts/{name}`| Upload a partial-result artifact | +| POST | `/tasks/{task_id}/result` | Complete with an artifact id (idempotent) | +| POST | `/tasks/{task_id}/failure` | Record failure / retryable state | +| GET | `/artifacts/{artifact_id}/download`| Download an artifact by id | +| GET | `/health` | Readiness incl. database (unauthenticated) | + +The full contract is in [`docs/api-contract.md`](../docs/api-contract.md) and +[`docs/openapi.yaml`](../docs/openapi.yaml); a worker-author guide is in +[`docs/building-workers.md`](../docs/building-workers.md). + +## Poking the API + +Two ways, both checked in: + +```sh +make smoke # every endpoint, asserted; non-zero exit on failure +``` + +`api/requests.http` runs the same calls one at a time from an editor with a REST +client (VSCodium/VS Code "REST Client", JetBrains HTTP Client). Later requests +reuse ids captured from earlier responses, so it doubles as API documentation. + +## Status + +Works end to end: a worker registers, a dataset is uploaded and chunked into +shard tasks (or a job is created from chunk URIs), tasks are leased one at a +time, downloaded, heartbeated (`leased → running`), completed via uploaded +result artifacts, and reflected in job progress. A reaper reclaims expired +leases and marks silent workers offline. + +Done: schema + migrations, atomic claim (`FOR UPDATE SKIP LOCKED`), optimistic +concurrency, result/failure paths, lease expiry, worker registry + liveness, +artifact storage, dataset upload + chunking, request-size limits. + +Still stubbed: `StitchJob.Execute` — merging per-chunk top-k into the final CSV +is workload semantics that belongs to the Python side (reducer). + +## Tests + +Unit tests need **no database** — domain rules, use-case orchestration (over +in-memory `internal/memstore`), and HTTP handlers (via `httptest`): + +```sh +make test # go test ./... +make vet +make lint +go test -race ./... +``` + +Integration tests run against a **real PostgreSQL** (the spec forbids mocks +here — they verify `FOR UPDATE SKIP LOCKED`, optimistic concurrency, rollback): + +```sh +docker compose up -d +make test-integration TEST_DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable' +``` + +CI (`.github/workflows/coordinator.yml`) runs vet, gofmt, race tests, lint, and +the integration suite against a Postgres service on every push and PR. diff --git a/coordinator/api/requests.http b/coordinator/api/requests.http new file mode 100644 index 0000000..6c50a59 --- /dev/null +++ b/coordinator/api/requests.http @@ -0,0 +1,234 @@ +# SciMesh Coordinator — API requests +# +# Runnable from any editor with a REST client (VSCodium/VS Code "REST Client", +# JetBrains HTTP Client). Click "Send Request" above each block, top to bottom: +# later requests reuse ids captured from earlier responses. +# +# Start the stack first: docker compose up -d + +@host = http://localhost:8080 +@token = change-me +@worker = worker-1 + +### Readiness — the only unauthenticated endpoint (probes the database) +GET {{host}}/health + +### Auth check — no token must be rejected with 401 +POST {{host}}/tasks/claim +Content-Type: application/json + +{ "worker_id": "{{worker}}" } + +### 0. Register a worker (201) +# @name register +POST {{host}}/workers/register +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "name": "lab-worker-01", + "capabilities": ["similarity_search"], + "cpu_count": 8, + "memory_mb": 16384 +} + +@workerId = {{register.response.body.worker_id}} + +### 0b. Upload a dataset — the coordinator splits it into shard tasks (201) +# Text fields first, the file part last (it is streamed, not buffered). +# @name uploadJob +POST {{host}}/jobs/upload +Authorization: Bearer {{token}} +Content-Type: multipart/form-data; boundary=----scimesh + +------scimesh +Content-Disposition: form-data; name="workload" + +similarity_search +------scimesh +Content-Disposition: form-data; name="parameters" + +{"top_k":10} +------scimesh +Content-Disposition: form-data; name="chunk_rows" + +2 +------scimesh +Content-Disposition: form-data; name="file"; filename="chembl.tsv" +Content-Type: text/tab-separated-values + +id smiles +A CC +B CCC +C CCCC +D CCCCC +------scimesh-- + +### Download a task's input shard (200) — taskId must be a shard task from an +### uploaded job (claim one first; its input.uri is /tasks/{id}/input). +GET {{host}}/tasks/{{taskId}}/input +Authorization: Bearer {{token}} + +### 1. Create a job and its chunks (201) +# The coordinator splits the submission into one task per chunk, transactionally. +# @name createJob +POST {{host}}/jobs +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "workload": "similarity_search", + "input_uri": "s3://chembl/full.sdf", + "parameters": { "top_k": 10 }, + "chunks": [ + { "chunk_index": 0, "input_uri": "s3://chembl/shard-0.sdf", "input_sha256": "aaa", "max_attempts": 3 }, + { "chunk_index": 1, "input_uri": "s3://chembl/shard-1.sdf", "input_sha256": "bbb", "max_attempts": 3 }, + { "chunk_index": 2, "input_uri": "s3://chembl/shard-2.sdf", "input_sha256": "ccc", "max_attempts": 3 } + ] +} + +@jobId = {{createJob.response.body.id}} + +### 2. Claim a task (200, or 204 when the queue is empty) +# Each call leases a different task; run it repeatedly to see chunk_index advance. +# @name claim +POST {{host}}/tasks/claim +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "{{worker}}", + "capabilities": ["similarity_search"], + "max_concurrency": 1 +} + +@taskId = {{claim.response.body.task_id}} +@attempt = {{claim.response.body.attempt}} + +### 3. Heartbeat — renew the lease while the task is still running (200) +POST {{host}}/tasks/{{taskId}}/heartbeat +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "{{worker}}", + "attempt": {{attempt}} +} + +### 3a. Upload a partial-result artifact (200) — while the task is leased +# Identity travels in headers per the contract; the body is streamed as-is. +# @name uploadArtifact +PUT {{host}}/tasks/{{taskId}}/artifacts/result.csv +Authorization: Bearer {{token}} +Content-Type: text/csv +X-Worker-ID: {{worker}} +X-Task-Attempt: {{attempt}} + +query,match,score +CHEMBL25,CHEMBL139,0.87 + +@artifactId = {{uploadArtifact.response.body.artifact_id}} + +### 3b. Download the artifact by id (200) +GET {{host}}/artifacts/{{artifactId}}/download +Authorization: Bearer {{token}} + +### 3c. Upload a second artifact — used by the conflict check below (200) +# @name uploadArtifact2 +PUT {{host}}/tasks/{{taskId}}/artifacts/secondary.csv +Authorization: Bearer {{token}} +Content-Type: text/csv +X-Worker-ID: {{worker}} +X-Task-Attempt: {{attempt}} + +query,match,score +CHEMBL25,CHEMBL521,0.42 + +@artifactId2 = {{uploadArtifact2.response.body.artifact_id}} + +### 4. Submit the result, referencing the uploaded artifact (200) +POST {{host}}/tasks/{{taskId}}/result +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "{{worker}}", + "attempt": {{attempt}}, + "result": { "artifact_id": "{{artifactId}}", "content_type": "text/csv" }, + "metrics": { "elapsed_ms": 1234, "candidates": 50000 } +} + +### 4a. Replay the same result — must be idempotent (200, not 409) +POST {{host}}/tasks/{{taskId}}/result +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "{{worker}}", + "attempt": {{attempt}}, + "result": { "artifact_id": "{{artifactId}}" } +} + +### 4b. A different artifact for the same task — conflict (409) +POST {{host}}/tasks/{{taskId}}/result +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "{{worker}}", + "attempt": {{attempt}}, + "result": { "artifact_id": "{{artifactId2}}" } +} + +### 4c. Another worker submitting for this task — conflict (409) +POST {{host}}/tasks/{{taskId}}/result +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "impostor", + "attempt": {{attempt}}, + "result": { "artifact_id": "{{artifactId}}" } +} + +### 5. Report a failure instead (200) +# retryable=true returns the task to the queue while attempts remain; +# retryable=false fails it terminally. +POST {{host}}/tasks/{{taskId}}/failure +Authorization: Bearer {{token}} +Content-Type: application/json + +{ + "worker_id": "{{worker}}", + "attempt": {{attempt}}, + "error_code": "download_failed", + "error_message": "checksum mismatch on shard", + "retryable": true +} + +### 6. Job progress (200) +GET {{host}}/jobs/{{jobId}} +Authorization: Bearer {{token}} + +### --- error cases ------------------------------------------------------- + +### Malformed UUID in the path (400) +POST {{host}}/tasks/not-a-uuid/result +Authorization: Bearer {{token}} +Content-Type: application/json + +{ "worker_id": "{{worker}}", "attempt": 1, "result_uri": "s3://x", "result_sha256": "x" } + +### Unknown field in the body (400) — a misspelled key must not pass silently +POST {{host}}/tasks/claim +Authorization: Bearer {{token}} +Content-Type: application/json + +{ "worker_ID": "{{worker}}" } + +### Unknown job (404) +GET {{host}}/jobs/00000000-0000-0000-0000-000000000000 +Authorization: Bearer {{token}} + +### Stitching is not implemented yet (501) +# Any endpoint whose use case is still a stub answers 501. diff --git a/coordinator/cmd/coordinator/main.go b/coordinator/cmd/coordinator/main.go new file mode 100644 index 0000000..dd11a56 --- /dev/null +++ b/coordinator/cmd/coordinator/main.go @@ -0,0 +1,125 @@ +package main + +import ( + "context" + "log/slog" + "os" + "os/signal" + "sync" + "syscall" + + "github.com/emil28092005/SciMesh/coordinator/internal/infra" + "github.com/emil28092005/SciMesh/coordinator/internal/storage/blob" + "github.com/emil28092005/SciMesh/coordinator/internal/storage/postgres" + httptransport "github.com/emil28092005/SciMesh/coordinator/internal/transport/http" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +func main() { + // All work happens in run() so its defers (pool.Close, log flush, signal + // stop) still execute: os.Exit skips deferred calls entirely. + if err := run(); err != nil { + os.Exit(1) + } +} + +func run() error { + // Bootstrap logger, used only until config says where logs should go. It + // writes to stderr so it never contaminates the configured stdout stream. + boot := slog.New(slog.NewJSONHandler(os.Stderr, nil)) + + cfg, err := infra.LoadConfig() + if err != nil { + boot.Error("load config", "err", err) + return err + } + + // The real logger: stdout plus an optional rotated file (LOG_FILE). + log, logCloser, err := infra.NewLogger(cfg) + if err != nil { + boot.Error("init logger", "err", err) + return err + } + defer func() { _ = logCloser.Close() }() + + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + pool, err := infra.NewPool(ctx, cfg, log) + if err != nil { + log.Error("connect database", "err", err) + return err + } + defer pool.Close() + + blobStore, err := blob.NewFSStore(cfg.StorageDir) + if err != nil { + log.Error("init blob storage", "err", err) + return err + } + + var ( + clk = infra.NewClock() + tx = postgres.NewTxManager(pool) + taskRepo = postgres.NewTaskRepo(pool) + jobRepo = postgres.NewJobRepo(pool) + workerRepo = postgres.NewWorkerRepo(pool) + artifactRepo = postgres.NewArtifactRepo(pool) + ) + + useCases := httptransport.UseCases{ + RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk), + 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, 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), + UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, clk), + DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore), + GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore), + } + + // 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 + 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, 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 + // LIFO, so the deferred stop() would fire *after* the wait below). + // + // 1. stop() cancel the context, telling the reaper to finish + // 2. wg.Wait() let it return from its current tick + // 3. deferred pool.Close() closes an idle pool, not a busy one + // + // Calling stop() here also covers the path where RunServer failed on its + // own: the context would never be cancelled otherwise and wg.Wait() + // would block forever. + stop() + wg.Wait() + log.Info("shutdown complete") + + return err +} diff --git a/coordinator/docker-compose.yml b/coordinator/docker-compose.yml new file mode 100644 index 0000000..901ac55 --- /dev/null +++ b/coordinator/docker-compose.yml @@ -0,0 +1,81 @@ +name: scimesh + +services: + postgres: + image: postgres:16-alpine + environment: + POSTGRES_USER: ${POSTGRES_USER:-scimesh} + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-scimesh} + POSTGRES_DB: ${POSTGRES_DB:-scimesh} + ports: + - "${POSTGRES_PORT:-5432}:5432" + volumes: + - pgdata:/var/lib/postgresql/data + healthcheck: + # Everything else waits on this, so the check must prove the server + # accepts queries — not merely that the port is open. + test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER:-scimesh} -d ${POSTGRES_DB:-scimesh}"] + interval: 5s + timeout: 3s + retries: 10 + start_period: 5s + + # One-shot: applies migrations, then exits. Schema changes stay an explicit + # deployment step — the coordinator binary never migrates on startup. + migrate: + image: migrate/migrate:v4.17.1 + depends_on: + postgres: + condition: service_healthy + volumes: + - ./migrations:/migrations:ro + command: + - -path=/migrations + - -database=postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres:5432/${POSTGRES_DB:-scimesh}?sslmode=disable + - up + restart: on-failure + + coordinator: + build: + context: . + depends_on: + postgres: + condition: service_healthy + # Start only once the schema exists, otherwise the first query fails. + migrate: + condition: service_completed_successfully + environment: + COORDINATOR_ADDR: ":8080" + # Host is the service name: compose resolves it on the project network. + DATABASE_URL: postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres:5432/${POSTGRES_DB:-scimesh}?sslmode=disable + WORKER_AUTH_TOKEN: ${WORKER_AUTH_TOKEN:-dev-token} + DB_MAX_CONNS: "10" + REQUEST_TIMEOUT: "15s" + LEASE_DURATION: "2m" + REAPER_INTERVAL: "30s" + LOG_LEVEL: ${LOG_LEVEL:-info} + # Logs are teed to stdout (docker logs) and this rotated file on a named + # volume, so they survive a rebuild. + LOG_FILE: /var/log/scimesh/coordinator.log + # Artifact bytes live on a named volume, durable across rebuilds. + COORDINATOR_STORAGE_DIR: /var/lib/scimesh/artifacts + ports: + - "${COORDINATOR_PORT:-8080}:8080" + # Named volumes (not host bind mounts): they inherit the image's directory + # ownership, so the non-root process can write to them. A bind mount would + # be root-owned and unwritable by uid 10001. + volumes: + - coordinator_logs:/var/log/scimesh + - coordinator_data:/var/lib/scimesh/artifacts + healthcheck: + test: ["CMD", "wget", "-qO-", "http://127.0.0.1:8080/health"] + interval: 10s + timeout: 3s + retries: 3 + start_period: 5s + restart: unless-stopped + +volumes: + pgdata: + coordinator_logs: + coordinator_data: diff --git a/coordinator/go.mod b/coordinator/go.mod new file mode 100644 index 0000000..4fc8bcc --- /dev/null +++ b/coordinator/go.mod @@ -0,0 +1,23 @@ +module github.com/emil28092005/SciMesh/coordinator + +go 1.22 + +require ( + github.com/Masterminds/squirrel v1.5.4 + github.com/cenkalti/backoff/v4 v4.3.0 + github.com/google/uuid v1.6.0 + github.com/jackc/pgx/v5 v5.6.0 + github.com/joho/godotenv v1.5.1 + gopkg.in/natefinch/lumberjack.v2 v2.2.1 +) + +require ( + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect + github.com/jackc/puddle/v2 v2.2.1 // indirect + github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect + github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect + golang.org/x/crypto v0.17.0 // indirect + golang.org/x/sync v0.1.0 // indirect + golang.org/x/text v0.14.0 // indirect +) diff --git a/coordinator/go.sum b/coordinator/go.sum new file mode 100644 index 0000000..d9c880a --- /dev/null +++ b/coordinator/go.sum @@ -0,0 +1,43 @@ +github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8afzqM= +github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10= +github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8= +github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk= +github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.6.0 h1:SWJzexBzPL5jb0GEsrPMLIsi/3jOo7RHlzTjcAeDrPY= +github.com/jackc/pgx/v5 v5.6.0/go.mod h1:DNZ/vlrUnhWCoFGxHAG8U2ljioxukquj7utPDgtQdTw= +github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk= +github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= +github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= +github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 h1:SOEGU9fKiNWd/HOJuq6+3iTQz8KNCLtVX6idSoTLdUw= +github.com/lann/builder v0.0.0-20180802200727-47ae307949d0/go.mod h1:dXGbAdH5GtBTC4WfIxhKZfyBF/HBFgRZSWwZ9g/He9o= +github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 h1:P6pPBnrTSX3DEVR4fDembhRWSsG5rVo6hYhAB/ADZrk= +github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0/go.mod h1:vmVJ0l/dxyfGW6FmdpVm2joNMFikkuWg0EoCKLGUMNw= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k= +golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4= +golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o= +golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= +golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc= +gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/coordinator/internal/chunk/tsv.go b/coordinator/internal/chunk/tsv.go new file mode 100644 index 0000000..6b968eb --- /dev/null +++ b/coordinator/internal/chunk/tsv.go @@ -0,0 +1,92 @@ +// Package chunk splits a tabular input into deterministic shards. It is generic +// row splitting only — no workload semantics (SMILES, top-k) live here. +package chunk + +import ( + "bufio" + "bytes" + "fmt" + "io" +) + +// ErrNoRows is returned when the input has a header but no data rows: a job with +// zero tasks could never complete, so it is rejected at the source. +var ErrNoRows = fmt.Errorf("input has no data rows") + +// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at +// most rowsPerShard data rows. Every shard repeats the header, so a worker can +// parse its shard in isolation. emit is called once per shard, in order, with a +// reader over that shard's bytes; the reader is valid only for the duration of +// the call. +// +// Splitting is deterministic: the same input and rowsPerShard always produce the +// same shards, byte for byte — which is what lets chunk_index refer to a stable +// piece and makes a re-run reproducible. +// +// Only one shard is buffered at a time, so memory is bounded by shard size (a +// worker-sized slice of the data), not by the size of the whole dataset. +func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error { + if rowsPerShard <= 0 { + return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard) + } + + sc := bufio.NewScanner(r) + // Allow long lines: a SMILES row can be far wider than bufio's 64 KB default. + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + + if !sc.Scan() { + if err := sc.Err(); err != nil { + return fmt.Errorf("read header: %w", err) + } + return ErrNoRows // completely empty input + } + header := append([]byte(nil), sc.Bytes()...) + + var ( + buf bytes.Buffer + rows int + index int + ) + + // flush emits the buffered shard and resets for the next one. + flush := func() error { + if err := emit(index, bytes.NewReader(buf.Bytes())); err != nil { + return err + } + index++ + buf.Reset() + rows = 0 + return nil + } + + for sc.Scan() { + if rows == 0 { + buf.Write(header) + buf.WriteByte('\n') + } + buf.Write(sc.Bytes()) + buf.WriteByte('\n') + rows++ + + if rows == rowsPerShard { + if err := flush(); err != nil { + return err + } + } + } + if err := sc.Err(); err != nil { + return fmt.Errorf("read rows: %w", err) + } + + // A partial final shard still has to go out. + if rows > 0 { + if err := flush(); err != nil { + return err + } + } + + if index == 0 { + return ErrNoRows // header only, no data + } + return nil +} diff --git a/coordinator/internal/chunk/tsv_test.go b/coordinator/internal/chunk/tsv_test.go new file mode 100644 index 0000000..ba12917 --- /dev/null +++ b/coordinator/internal/chunk/tsv_test.go @@ -0,0 +1,117 @@ +package chunk + +import ( + "bytes" + "errors" + "fmt" + "io" + "strings" + "testing" +) + +// collect runs SplitTSV and returns every shard as a string. +func collect(t *testing.T, input string, rowsPerShard int) []string { + t.Helper() + var shards []string + err := SplitTSV(strings.NewReader(input), rowsPerShard, func(index int, shard io.Reader) error { + b, _ := io.ReadAll(shard) + if index != len(shards) { + t.Fatalf("emit index = %d, want %d (out of order)", index, len(shards)) + } + shards = append(shards, string(b)) + return nil + }) + if err != nil { + t.Fatalf("SplitTSV: %v", err) + } + return shards +} + +func TestSplitCountsShardsAndRepeatsHeader(t *testing.T) { + input := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" + shards := collect(t, input, 2) + + if len(shards) != 3 { // 5 rows / 2 per shard = ceil = 3 + t.Fatalf("got %d shards, want 3", len(shards)) + } + for i, s := range shards { + if !strings.HasPrefix(s, "id\tsmiles\n") { + t.Errorf("shard %d missing header: %q", i, s) + } + } + if shards[0] != "id\tsmiles\nA\tCC\nB\tCCC\n" { + t.Errorf("shard 0 = %q", shards[0]) + } + if shards[2] != "id\tsmiles\nE\tCCCCCC\n" { // partial final shard + t.Errorf("shard 2 = %q", shards[2]) + } +} + +func TestSplitExactMultipleHasNoEmptyTrailingShard(t *testing.T) { + input := "h\nr1\nr2\nr3\nr4\n" + shards := collect(t, input, 2) + if len(shards) != 2 { // exactly 4/2, no empty third shard + t.Fatalf("got %d shards, want 2", len(shards)) + } +} + +func TestSplitIsDeterministic(t *testing.T) { + input := "h\n" + strings.Repeat("row\n", 100) + a := collect(t, input, 7) + b := collect(t, input, 7) + if fmt.Sprint(a) != fmt.Sprint(b) { + t.Error("two runs produced different shards") + } +} + +func TestSplitRejectsHeaderOnly(t *testing.T) { + err := SplitTSV(strings.NewReader("id\tsmiles\n"), 10, func(int, io.Reader) error { return nil }) + if !errors.Is(err, ErrNoRows) { + t.Errorf("err = %v, want ErrNoRows", err) + } +} + +func TestSplitRejectsEmptyInput(t *testing.T) { + err := SplitTSV(strings.NewReader(""), 10, func(int, io.Reader) error { return nil }) + if !errors.Is(err, ErrNoRows) { + t.Errorf("err = %v, want ErrNoRows", err) + } +} + +func TestSplitRejectsNonPositiveSize(t *testing.T) { + err := SplitTSV(strings.NewReader("h\nr\n"), 0, func(int, io.Reader) error { return nil }) + if err == nil { + t.Error("expected an error for rowsPerShard = 0") + } +} + +func TestSplitPropagatesEmitError(t *testing.T) { + boom := errors.New("boom") + err := SplitTSV(strings.NewReader("h\nr1\nr2\n"), 1, func(int, io.Reader) error { return boom }) + if !errors.Is(err, boom) { + t.Errorf("err = %v, want boom", err) + } +} + +func TestSplitSingleShardWhenSizeExceedsRows(t *testing.T) { + shards := collect(t, "h\nr1\nr2\n", 100) + if len(shards) != 1 { + t.Fatalf("got %d shards, want 1", len(shards)) + } + if shards[0] != "h\nr1\nr2\n" { + t.Errorf("shard 0 = %q", shards[0]) + } +} + +// The scanned bytes are reused by bufio; the shard buffer must copy them, or a +// later row would corrupt an earlier one. This guards that copy. +func TestSplitDoesNotAliasScannerBuffer(t *testing.T) { + var got bytes.Buffer + _ = SplitTSV(strings.NewReader("h\naaaa\nbbbb\n"), 2, func(_ int, shard io.Reader) error { + _, _ = io.Copy(&got, shard) + return nil + }) + if want := "h\naaaa\nbbbb\n"; got.String() != want { + t.Errorf("got %q, want %q", got.String(), want) + } +} diff --git a/coordinator/internal/domain/artifact.go b/coordinator/internal/domain/artifact.go new file mode 100644 index 0000000..43bc219 --- /dev/null +++ b/coordinator/internal/domain/artifact.go @@ -0,0 +1,69 @@ +package domain + +import ( + "time" + + "github.com/google/uuid" +) + +type ArtifactKind string + +const ( + ArtifactInput ArtifactKind = "input" + ArtifactShard ArtifactKind = "shard" + ArtifactPartialResult ArtifactKind = "partial_result" + ArtifactFinalResult ArtifactKind = "final_result" + ArtifactLog ArtifactKind = "log" +) + +// Artifact is a durable file the coordinator owns, described by its metadata. +// The bytes live in blob storage under StorageKey; this struct is what the +// database persists and what every other layer reasons about. +type Artifact struct { + ID uuid.UUID + JobID uuid.UUID + TaskID *uuid.UUID // nil for a job-level input + Kind ArtifactKind + Filename string + StorageKey string + ContentType string + SizeBytes int64 + SHA256 string + CreatedAt time.Time +} + +// NewArtifact begins an artifact record. Size and checksum are unknown until the +// bytes have been streamed to storage, so they are filled in later by SetContent. +// +// StorageKey is derived from a fresh UUID, never from the client-supplied +// filename — that is what stops a "../../etc/passwd" filename from escaping the +// storage directory. +func NewArtifact(jobID uuid.UUID, taskID *uuid.UUID, kind ArtifactKind, + filename, contentType string, now time.Time) (*Artifact, error) { + + if filename == "" || kind == "" { + return nil, ErrInvalidInput + } + if contentType == "" { + contentType = "application/octet-stream" + } + id := uuid.New() + return &Artifact{ + ID: id, + JobID: jobID, + TaskID: taskID, + Kind: kind, + Filename: filename, + StorageKey: id.String(), + ContentType: contentType, + CreatedAt: now, + }, nil +} + +// SetContent records the size and checksum measured while streaming the bytes +// into storage. Both are computed by the coordinator, never trusted from the +// client — the whole point of owning the artifact. +func (a *Artifact) SetContent(sha256 string, size int64) { + a.SHA256 = sha256 + a.SizeBytes = size +} diff --git a/coordinator/internal/domain/artifact_test.go b/coordinator/internal/domain/artifact_test.go new file mode 100644 index 0000000..da4bd22 --- /dev/null +++ b/coordinator/internal/domain/artifact_test.go @@ -0,0 +1,55 @@ +package domain + +import ( + "errors" + "testing" + + "github.com/google/uuid" +) + +func TestNewArtifact(t *testing.T) { + jobID := uuid.New() + taskID := uuid.New() + a, err := NewArtifact(jobID, &taskID, ArtifactPartialResult, "result.csv", "text/csv", testNow) + if err != nil { + t.Fatal(err) + } + if a.JobID != jobID || a.TaskID == nil || *a.TaskID != taskID { + t.Error("ownership not recorded") + } + // Storage key is derived from the artifact id, never the filename — no path + // traversal from a hostile "../.." name. + if a.StorageKey != a.ID.String() { + t.Errorf("storage key = %q, want the artifact id", a.StorageKey) + } + if a.SizeBytes != 0 || a.SHA256 != "" { + t.Error("size and checksum are unknown until SetContent") + } +} + +func TestNewArtifactDefaultsContentType(t *testing.T) { + a, err := NewArtifact(uuid.New(), nil, ArtifactInput, "data", "", testNow) + if err != nil { + t.Fatal(err) + } + if a.ContentType != "application/octet-stream" { + t.Errorf("content type = %q, want the default", a.ContentType) + } +} + +func TestNewArtifactRejectsBadInput(t *testing.T) { + if _, err := NewArtifact(uuid.New(), nil, ArtifactInput, "", "text/csv", testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("empty filename: err = %v, want ErrInvalidInput", err) + } + if _, err := NewArtifact(uuid.New(), nil, "", "f", "text/csv", testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("empty kind: err = %v, want ErrInvalidInput", err) + } +} + +func TestArtifactSetContent(t *testing.T) { + a, _ := NewArtifact(uuid.New(), nil, ArtifactShard, "shard-0.tsv", "text/csv", testNow) + a.SetContent("deadbeef", 42) + if a.SHA256 != "deadbeef" || a.SizeBytes != 42 { + t.Error("SetContent must record checksum and size") + } +} diff --git a/coordinator/internal/domain/errors.go b/coordinator/internal/domain/errors.go new file mode 100644 index 0000000..d64b1f6 --- /dev/null +++ b/coordinator/internal/domain/errors.go @@ -0,0 +1,20 @@ +package domain + +import "errors" + +// Business-rule violations. They live in the innermost layer because they +// describe what the rules are, not how a transport reports them: the HTTP +// adapter maps these to status codes, and nothing here knows 409 exists. +// +// Always compare with errors.Is — outer layers may wrap these with %w. +var ( + ErrJobNotFound = errors.New("job not found") + ErrTaskNotFound = errors.New("task not found") + ErrWorkerNotFound = errors.New("worker not found") + ErrArtifactNotFound = errors.New("artifact not found") + ErrLeaseConflict = errors.New("task leased to another worker") + ErrStaleAttempt = errors.New("attempt does not match lease") + ErrResultConflict = errors.New("different result already recorded") + ErrInvalidInput = errors.New("invalid input") + ErrTaskNotLeased = errors.New("task is not currently leased") +) diff --git a/coordinator/internal/domain/job.go b/coordinator/internal/domain/job.go new file mode 100644 index 0000000..0aa4827 --- /dev/null +++ b/coordinator/internal/domain/job.go @@ -0,0 +1,125 @@ +package domain + +import ( + "time" + + "github.com/google/uuid" +) + +type JobStatus string + +const ( + JobPending JobStatus = "pending" + JobRunning JobStatus = "running" + JobCompleted JobStatus = "completed" + JobFailed JobStatus = "failed" + JobCancelled JobStatus = "cancelled" +) + +// Job is one user submission that fans out into one or more tasks. +type Job struct { + ID uuid.UUID + Workload string + InputURI string // external input URI; empty for uploaded datasets + InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions + Parameters map[string]any + Status JobStatus + CreatedAt time.Time + CompletedAt *time.Time +} + +// NewUploadedJob builds a job whose input was uploaded to the coordinator. The +// job's id is generated here so the input artifact can reference it; the reverse +// link (jobs.input_artifact_id) is left unset — the input is found via the +// artifact's job_id — which also sidesteps the circular job↔artifact FK. +func NewUploadedJob(workload string, params map[string]any, now time.Time) (*Job, error) { + if workload == "" { + return nil, ErrInvalidInput + } + return &Job{ + ID: uuid.New(), + Workload: workload, + Parameters: params, + Status: JobPending, + CreatedAt: now, + }, nil +} + +// ChunkSpec describes one piece a job is split into. Callers build these from +// whatever chunking strategy the workload uses; the domain only validates them. +type ChunkSpec struct { + ChunkIndex int + Workload string // empty inherits the job's workload + InputURI string + InputSHA256 string + Parameters map[string]any + MaxAttempts int +} + +// NewJobWithTasks builds a job together with all of its tasks, validating the +// set as a whole. Returning both from one constructor keeps the invariant +// visible: a job without tasks, or with duplicate chunk indexes, cannot exist. +func NewJobWithTasks(workload, inputURI string, params map[string]any, + chunks []ChunkSpec, now time.Time) (*Job, []*Task, error) { + + if workload == "" || inputURI == "" || len(chunks) == 0 { + return nil, nil, ErrInvalidInput + } + + job := &Job{ + ID: uuid.New(), + Workload: workload, + InputURI: inputURI, + Parameters: params, + Status: JobPending, + CreatedAt: now, + } + + seen := make(map[int]struct{}, len(chunks)) + tasks := make([]*Task, 0, len(chunks)) + for _, c := range chunks { + if _, dup := seen[c.ChunkIndex]; dup { + return nil, nil, ErrInvalidInput // unique (job_id, chunk_index) + } + seen[c.ChunkIndex] = struct{}{} + + w := c.Workload + if w == "" { + w = workload + } + task, err := NewTask(job.ID, c.ChunkIndex, w, c.InputURI, c.InputSHA256, + c.Parameters, c.MaxAttempts, now) + if err != nil { + return nil, nil, err + } + tasks = append(tasks, task) + } + return job, tasks, nil +} + +// JobProgress is the aggregate view of a job and the state of its tasks. +type JobProgress struct { + Job Job + Total int + Pending int + Leased int + Done int + Failed int +} + +// DeriveStatus computes what the job's status should be from its task counts, +// so the rule lives here rather than in a SQL trigger or a handler. +func (p JobProgress) DeriveStatus() JobStatus { + switch { + case p.Total == 0: + return JobPending + case p.Done == p.Total: + return JobCompleted + case p.Failed > 0 && p.Done+p.Failed == p.Total: + return JobFailed + case p.Leased > 0 || p.Done > 0 || p.Failed > 0: + return JobRunning + default: + return JobPending + } +} diff --git a/coordinator/internal/domain/job_test.go b/coordinator/internal/domain/job_test.go new file mode 100644 index 0000000..0b58902 --- /dev/null +++ b/coordinator/internal/domain/job_test.go @@ -0,0 +1,140 @@ +package domain + +import ( + "errors" + "testing" + + "github.com/google/uuid" +) + +func TestNewJobWithTasksBuildsBoth(t *testing.T) { + job, tasks, err := NewJobWithTasks("similarity_search", "s3://in", nil, []ChunkSpec{ + {ChunkIndex: 0, InputURI: "s3://c0", InputSHA256: "a"}, + {ChunkIndex: 1, InputURI: "s3://c1", InputSHA256: "b"}, + }, testNow) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(tasks) != 2 { + t.Fatalf("got %d tasks, want 2", len(tasks)) + } + for _, tk := range tasks { + if tk.JobID != job.ID { + t.Error("task not linked to job") + } + if tk.Workload != "similarity_search" { + t.Error("task should inherit the job workload") + } + } + if job.Status != JobPending { + t.Errorf("status = %q, want pending", job.Status) + } +} + +func TestNewJobWithTasksRejectsBadInput(t *testing.T) { + good := []ChunkSpec{{ChunkIndex: 0, InputURI: "s3://c0", InputSHA256: "a"}} + cases := map[string]struct { + workload string + inputURI string + chunks []ChunkSpec + }{ + "empty workload": {"", "s3://in", good}, + "empty input": {"w", "", good}, + "no chunks": {"w", "s3://in", nil}, + "duplicate index": {"w", "s3://in", []ChunkSpec{ + {ChunkIndex: 0, InputURI: "a", InputSHA256: "x"}, + {ChunkIndex: 0, InputURI: "b", InputSHA256: "y"}, + }}, + } + for name, c := range cases { + t.Run(name, func(t *testing.T) { + if _, _, err := NewJobWithTasks(c.workload, c.inputURI, nil, c.chunks, testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("err = %v, want ErrInvalidInput", err) + } + }) + } +} + +func TestNewJobWithTasksInheritsAndOverridesWorkload(t *testing.T) { + _, tasks, err := NewJobWithTasks("base", "s3://in", nil, []ChunkSpec{ + {ChunkIndex: 0, InputURI: "a", InputSHA256: "x"}, + {ChunkIndex: 1, InputURI: "b", InputSHA256: "y", Workload: "special"}, + }, testNow) + if err != nil { + t.Fatal(err) + } + if tasks[0].Workload != "base" || tasks[1].Workload != "special" { + t.Errorf("workloads = %q, %q", tasks[0].Workload, tasks[1].Workload) + } +} + +func TestDeriveStatus(t *testing.T) { + cases := []struct { + name string + p JobProgress + want JobStatus + }{ + {"empty", JobProgress{Total: 0}, JobPending}, + {"all pending", JobProgress{Total: 3, Pending: 3}, JobPending}, + {"one leased", JobProgress{Total: 3, Pending: 2, Leased: 1}, JobRunning}, + {"partly done", JobProgress{Total: 3, Pending: 1, Done: 2}, JobRunning}, + {"all done", JobProgress{Total: 3, Done: 3}, JobCompleted}, + {"done and failed", JobProgress{Total: 3, Done: 2, Failed: 1}, JobFailed}, + {"failed but work remains", JobProgress{Total: 3, Pending: 1, Failed: 2}, JobRunning}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := c.p.DeriveStatus(); got != c.want { + t.Errorf("DeriveStatus() = %q, want %q", got, c.want) + } + }) + } +} + +func TestNewUploadedJob(t *testing.T) { + job, err := NewUploadedJob("w", map[string]any{"k": 1}, testNow) + if err != nil { + t.Fatal(err) + } + if job.Status != JobPending || job.InputURI != "" { + t.Error("uploaded job should be pending with no input URI") + } + if _, err := NewUploadedJob("", nil, testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("empty workload: err = %v, want ErrInvalidInput", err) + } +} + +func TestNewShardTask(t *testing.T) { + art := uuid.New() + task, err := NewShardTask(uuid.New(), 2, "w", art, "sha", nil, 0, testNow) + if err != nil { + t.Fatal(err) + } + if task.InputArtifactID == nil || *task.InputArtifactID != art { + t.Error("shard task must reference its input artifact") + } + if task.InputURI != "" { + t.Error("shard task must not carry a URI") + } + if task.MaxAttempts != DefaultMaxAttempts { + t.Errorf("maxAttempts = %d, want default %d", task.MaxAttempts, DefaultMaxAttempts) + } + + bad := []struct { + name string + art uuid.UUID + sha string + idx int + }{ + {"nil artifact", uuid.Nil, "sha", 0}, + {"empty sha", art, "", 0}, + {"negative index", art, "sha", -1}, + } + for _, c := range bad { + t.Run(c.name, func(t *testing.T) { + if _, err := NewShardTask(uuid.New(), c.idx, "w", c.art, c.sha, nil, 0, testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("err = %v, want ErrInvalidInput", err) + } + }) + } +} diff --git a/coordinator/internal/domain/task.go b/coordinator/internal/domain/task.go new file mode 100644 index 0000000..4fadca2 --- /dev/null +++ b/coordinator/internal/domain/task.go @@ -0,0 +1,281 @@ +// Package domain holds SciMesh's entities and the rules that govern them. It +// is the innermost layer: it imports nothing from this module and knows nothing +// about HTTP, SQL, or configuration. Every state transition a task can undergo +// is a method here, so the rules are unit-testable without a database. +package domain + +import ( + "time" + + "github.com/google/uuid" +) + +type TaskStatus string + +const ( + TaskPending TaskStatus = "pending" + TaskLeased TaskStatus = "leased" + TaskRunning TaskStatus = "running" + TaskCompleted TaskStatus = "completed" + TaskFailed TaskStatus = "failed" + TaskCancelled TaskStatus = "cancelled" +) + +// ErrCodeLeaseExpired marks tasks failed by the reaper rather than by a worker. +const ErrCodeLeaseExpired = "lease_expired" + +// Task is one independently executable chunk of a job. +// +// Nullable columns are pointers so "no lease" stays distinguishable from +// "lease owned by the empty string" — a plain string cannot express both. +type Task struct { + ID uuid.UUID + JobID uuid.UUID + ChunkIndex int + Workload string + InputURI string // external input URI; empty for uploaded shards + InputArtifactID *uuid.UUID // coordinator-stored shard; nil for URI inputs + InputSHA256 string + Parameters map[string]any + Status TaskStatus + Attempt int + MaxAttempts int + LeaseOwner *string + LeaseExpiresAt *time.Time + ResultArtifactID *uuid.UUID + Metrics map[string]any + ErrorCode *string + ErrorMessage *string + CreatedAt time.Time + StartedAt *time.Time + CompletedAt *time.Time + Version int +} + +// NewTask builds a pending task. maxAttempts <= 0 falls back to the default. +func NewTask(jobID uuid.UUID, chunkIndex int, workload, inputURI, inputSHA256 string, + params map[string]any, maxAttempts int, now time.Time) (*Task, error) { + + if inputURI == "" { + return nil, ErrInvalidInput + } + if inputSHA256 == "" { + return nil, ErrInvalidInput // checksum is mandatory: workers verify inputs + } + if chunkIndex < 0 { + return nil, ErrInvalidInput + } + if maxAttempts <= 0 { + maxAttempts = DefaultMaxAttempts + } + return &Task{ + ID: uuid.New(), + JobID: jobID, + ChunkIndex: chunkIndex, + Workload: workload, + InputURI: inputURI, + InputSHA256: inputSHA256, + Parameters: params, + Status: TaskPending, + Attempt: 0, + MaxAttempts: maxAttempts, + CreatedAt: now, + }, nil +} + +// NewShardTask builds a pending task whose input is a coordinator-stored shard +// artifact rather than an external URI. The worker fetches it from the +// coordinator, so no InputURI is set — inputSHA256 is the shard's checksum. +func NewShardTask(jobID uuid.UUID, chunkIndex int, workload string, inputArtifactID uuid.UUID, + inputSHA256 string, params map[string]any, maxAttempts int, now time.Time) (*Task, error) { + + if inputArtifactID == uuid.Nil || inputSHA256 == "" || chunkIndex < 0 { + return nil, ErrInvalidInput + } + if maxAttempts <= 0 { + maxAttempts = DefaultMaxAttempts + } + return &Task{ + ID: uuid.New(), + JobID: jobID, + ChunkIndex: chunkIndex, + Workload: workload, + InputArtifactID: &inputArtifactID, + InputSHA256: inputSHA256, + Parameters: params, + Status: TaskPending, + Attempt: 0, + MaxAttempts: maxAttempts, + CreatedAt: now, + }, nil +} + +// DefaultMaxAttempts applies when a task does not specify its own ceiling. +const DefaultMaxAttempts = 3 + +// CanRetry reports whether any attempts remain. +func (t *Task) CanRetry() bool { return t.Attempt < t.MaxAttempts } + +// IsLeaseHeldBy reports whether worker currently holds this task at attempt. +func (t *Task) IsLeaseHeldBy(worker string, attempt int) bool { + return t.LeaseOwner != nil && *t.LeaseOwner == worker && t.Attempt == attempt +} + +// AsClaimed projects the task into the trimmed view handed to a worker: +// everything needed to execute, nothing it has no business seeing. +func (t *Task) AsClaimed() ClaimedTask { + ct := ClaimedTask{ + TaskID: t.ID, + JobID: t.JobID, + ChunkIndex: t.ChunkIndex, + Workload: t.Workload, + InputURI: t.InputURI, + InputArtifactID: t.InputArtifactID, + InputSHA256: t.InputSHA256, + Parameters: t.Parameters, + Attempt: t.Attempt, + } + if t.LeaseOwner != nil { + ct.LeaseOwner = *t.LeaseOwner + } + if t.LeaseExpiresAt != nil { + ct.LeaseExpiresAt = *t.LeaseExpiresAt + } + return ct +} + +// 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 { + // 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 { + return ErrLeaseConflict + } + if t.Attempt != attempt { + return ErrStaleAttempt + } + return nil +} + +// 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 +} + +// CompleteWith records a successful result. +// +// Idempotency comes first deliberately: a worker whose network dropped will +// retry the same manifest, and that must succeed rather than trip the lease +// check on a task the coordinator already finished. A *different* manifest for +// an already-completed task is a genuine conflict. +func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any, + worker string, attempt int, now time.Time) error { + + if resultArtifactID == uuid.Nil { + return ErrInvalidInput + } + + if t.Status == TaskCompleted { + if t.Attempt == attempt && t.ResultArtifactID != nil && *t.ResultArtifactID == resultArtifactID { + return nil // same attempt, same artifact — replay of a successful call + } + return ErrResultConflict + } + + if err := t.verifyLease(worker, attempt); err != nil { + return err + } + + t.Status = TaskCompleted + t.ResultArtifactID = &resultArtifactID + t.Metrics = metrics + t.CompletedAt = &now + t.LeaseOwner = nil + t.LeaseExpiresAt = nil + t.ErrorCode = nil + t.ErrorMessage = nil + t.Version++ + return nil +} + +// Fail records a worker-reported failure. A retryable failure with attempts +// left returns the task to the queue; otherwise it terminates as failed. +func (t *Task) Fail(worker string, attempt int, code, message string, retryable bool, now time.Time) error { + if err := t.verifyLease(worker, attempt); err != nil { + return err + } + t.ErrorCode = &code + t.ErrorMessage = &message + t.LeaseOwner = nil + t.LeaseExpiresAt = nil + t.Version++ + + if retryable && t.CanRetry() { + t.Status = TaskPending + return nil + } + t.Status = TaskFailed + t.CompletedAt = &now + return nil +} + +// 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) { + // 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 + t.LeaseExpiresAt = nil + t.Version++ + + if t.CanRetry() { + t.Status = TaskPending + return + } + code, msg := ErrCodeLeaseExpired, "lease expired after the final attempt" + t.ErrorCode = &code + t.ErrorMessage = &msg + t.Status = TaskFailed + t.CompletedAt = &now +} + +// ClaimedTask is the worker-facing projection of a leased task. Input is either +// an external URI or a coordinator-stored shard (InputArtifactID set); the +// transport turns the latter into a coordinator download URL. +type ClaimedTask struct { + TaskID uuid.UUID + JobID uuid.UUID + ChunkIndex int + Workload string + InputURI string + InputArtifactID *uuid.UUID + InputSHA256 string + Parameters map[string]any + Attempt int + LeaseOwner string + LeaseExpiresAt time.Time +} + +// ResultManifest is a completed task's output, ordered for the stitcher. It +// points at the coordinator-owned result artifact rather than a worker URI. +type ResultManifest struct { + TaskID uuid.UUID + ChunkIndex int + ResultArtifactID uuid.UUID + Metrics map[string]any +} diff --git a/coordinator/internal/domain/task_test.go b/coordinator/internal/domain/task_test.go new file mode 100644 index 0000000..a8fa352 --- /dev/null +++ b/coordinator/internal/domain/task_test.go @@ -0,0 +1,221 @@ +package domain + +import ( + "errors" + "testing" + "time" + + "github.com/google/uuid" +) + +var ( + testNow = time.Date(2026, 7, 21, 12, 0, 0, 0, time.UTC) + testLater = testNow.Add(time.Hour) + testWorker = "worker-1" + testResult = uuid.New() + testResultAlt = uuid.New() +) + +// leasedTask builds a task already leased to testWorker at the given attempt. +func leasedTask(attempt, maxAttempts int) *Task { + owner := testWorker + expires := testLater + return &Task{ + ID: uuid.New(), + JobID: uuid.New(), + Status: TaskLeased, + Attempt: attempt, + MaxAttempts: maxAttempts, + LeaseOwner: &owner, + LeaseExpiresAt: &expires, + } +} + +func TestCompleteWithRecordsResult(t *testing.T) { + task := leasedTask(1, 3) + + if err := task.CompleteWith(testResult, nil, testWorker, 1, testNow); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if task.Status != TaskCompleted { + t.Errorf("status = %q, want completed", task.Status) + } + if task.LeaseOwner != nil || task.LeaseExpiresAt != nil { + t.Error("lease must be released on completion") + } + if task.CompletedAt == nil || !task.CompletedAt.Equal(testNow) { + t.Error("completed_at must be stamped") + } +} + +// A worker whose network dropped retries the same manifest; that must succeed +// rather than fail on the lease it has already given up. +func TestCompleteWithIsIdempotentForSameManifest(t *testing.T) { + task := leasedTask(1, 3) + if err := task.CompleteWith(testResult, nil, testWorker, 1, testNow); err != nil { + t.Fatalf("first call: %v", err) + } + versionAfterFirst := task.Version + + if err := task.CompleteWith(testResult, nil, testWorker, 1, testLater); err != nil { + t.Fatalf("replay must be idempotent, got %v", err) + } + if task.Version != versionAfterFirst { + t.Error("replay must not mutate the task") + } +} + +func TestCompleteWithRejectsDifferentManifest(t *testing.T) { + task := leasedTask(1, 3) + if err := task.CompleteWith(testResult, nil, testWorker, 1, testNow); err != nil { + t.Fatalf("first call: %v", err) + } + + err := task.CompleteWith(testResultAlt, nil, testWorker, 1, testLater) + if !errors.Is(err, ErrResultConflict) { + t.Errorf("err = %v, want ErrResultConflict", err) + } +} + +func TestCompleteWithRejectsForeignWorker(t *testing.T) { + task := leasedTask(1, 3) + + err := task.CompleteWith(testResult, nil, "worker-2", 1, testNow) + if !errors.Is(err, ErrLeaseConflict) { + t.Errorf("err = %v, want ErrLeaseConflict", err) + } +} + +func TestCompleteWithRejectsStaleAttempt(t *testing.T) { + task := leasedTask(2, 3) // task is on attempt 2 + + err := task.CompleteWith(testResult, nil, testWorker, 1, testNow) // worker thinks it is 1 + if !errors.Is(err, ErrStaleAttempt) { + t.Errorf("err = %v, want ErrStaleAttempt", err) + } +} + +func TestFailRequeuesWhileAttemptsRemain(t *testing.T) { + task := leasedTask(1, 3) + + if err := task.Fail(testWorker, 1, "boom", "exploded", true, testNow); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if task.Status != TaskPending { + t.Errorf("status = %q, want pending", task.Status) + } + if task.LeaseOwner != nil { + t.Error("lease must be released so another worker can claim it") + } +} + +func TestFailTerminatesOnFinalAttempt(t *testing.T) { + task := leasedTask(3, 3) // no attempts left + + if err := task.Fail(testWorker, 3, "boom", "exploded", true, testNow); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if task.Status != TaskFailed { + t.Errorf("status = %q, want failed", task.Status) + } +} + +func TestFailIsTerminalWhenNotRetryable(t *testing.T) { + task := leasedTask(1, 3) // attempts remain, but the error is fatal + + if err := task.Fail(testWorker, 1, "bad_input", "checksum mismatch", false, testNow); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if task.Status != TaskFailed { + t.Errorf("status = %q, want failed", task.Status) + } +} + +// This is the MVP acceptance criterion: a dead worker must not strand its task. +func TestExpireLeaseRequeuesWhileAttemptsRemain(t *testing.T) { + task := leasedTask(1, 3) + + task.ExpireLease(testNow) + + if task.Status != TaskPending { + t.Errorf("status = %q, want pending", task.Status) + } + if task.LeaseOwner != nil || task.LeaseExpiresAt != nil { + t.Error("expired lease must be cleared") + } +} + +func TestExpireLeaseFailsAfterFinalAttempt(t *testing.T) { + task := leasedTask(3, 3) + + task.ExpireLease(testNow) + + if task.Status != TaskFailed { + t.Errorf("status = %q, want failed", task.Status) + } + if task.ErrorCode == nil || *task.ErrorCode != ErrCodeLeaseExpired { + t.Error("expected a lease_expired error code") + } +} + +func TestExpireLeaseIgnoresUnleasedTasks(t *testing.T) { + task := &Task{Status: TaskCompleted, Attempt: 1, MaxAttempts: 3} + + task.ExpireLease(testNow) + + if task.Status != TaskCompleted { + t.Errorf("status = %q, completed tasks must be untouched", task.Status) + } +} + +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) + + if err := task.RenewLease(testWorker, 1, until); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !task.LeaseExpiresAt.Equal(until) { + t.Error("lease must be extended") + } + + if err := task.RenewLease("worker-2", 1, until); !errors.Is(err, ErrLeaseConflict) { + t.Errorf("err = %v, want ErrLeaseConflict", err) + } +} diff --git a/coordinator/internal/domain/worker.go b/coordinator/internal/domain/worker.go new file mode 100644 index 0000000..e7e0bf9 --- /dev/null +++ b/coordinator/internal/domain/worker.go @@ -0,0 +1,45 @@ +package domain + +import ( + "time" + + "github.com/google/uuid" +) + +type WorkerStatus string + +const ( + WorkerOnline WorkerStatus = "online" + WorkerBusy WorkerStatus = "busy" + WorkerOffline WorkerStatus = "offline" +) + +// Worker is a registered process/machine allowed to claim tasks. Its +// capabilities are the allowlisted workload names it can run; the coordinator +// never hands it a task outside that set. +type Worker struct { + ID uuid.UUID + Name string + Capabilities []string + Status WorkerStatus + LastHeartbeatAt time.Time + CreatedAt time.Time + UpdatedAt time.Time +} + +// NewWorker registers a worker. A worker with no capabilities could never be +// handed a task, so an empty set is rejected rather than silently stored. +func NewWorker(name string, capabilities []string, now time.Time) (*Worker, error) { + if len(capabilities) == 0 { + return nil, ErrInvalidInput + } + return &Worker{ + ID: uuid.New(), + Name: name, + Capabilities: capabilities, + Status: WorkerOnline, + LastHeartbeatAt: now, + CreatedAt: now, + UpdatedAt: now, + }, nil +} diff --git a/coordinator/internal/domain/worker_test.go b/coordinator/internal/domain/worker_test.go new file mode 100644 index 0000000..d39716e --- /dev/null +++ b/coordinator/internal/domain/worker_test.go @@ -0,0 +1,31 @@ +package domain + +import ( + "errors" + "testing" +) + +func TestNewWorker(t *testing.T) { + w, err := NewWorker("lab-01", []string{"similarity_search"}, testNow) + if err != nil { + t.Fatal(err) + } + if w.Status != WorkerOnline { + t.Errorf("status = %q, want online", w.Status) + } + if w.ID.String() == "" { + t.Error("worker must get an id") + } + if !w.LastHeartbeatAt.Equal(testNow) || !w.CreatedAt.Equal(testNow) { + t.Error("timestamps must be stamped") + } +} + +func TestNewWorkerRejectsNoCapabilities(t *testing.T) { + if _, err := NewWorker("lab-01", nil, testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("err = %v, want ErrInvalidInput", err) + } + if _, err := NewWorker("lab-01", []string{}, testNow); !errors.Is(err, ErrInvalidInput) { + t.Errorf("empty slice: err = %v, want ErrInvalidInput", err) + } +} diff --git a/coordinator/internal/infra/clock.go b/coordinator/internal/infra/clock.go new file mode 100644 index 0000000..eedbcde --- /dev/null +++ b/coordinator/internal/infra/clock.go @@ -0,0 +1,13 @@ +// Clock: the real implementation of the usecase.Clock port. It lives out here +// because reading the system clock is infrastructure; tests substitute a fixed one. +package infra + +import "time" + +type System struct{} + +func NewClock() System { return System{} } + +// Now returns UTC so every timestamp the coordinator writes is comparable +// regardless of the host's timezone. +func (System) Now() time.Time { return time.Now().UTC() } diff --git a/coordinator/internal/infra/config.go b/coordinator/internal/infra/config.go new file mode 100644 index 0000000..8703a27 --- /dev/null +++ b/coordinator/internal/infra/config.go @@ -0,0 +1,184 @@ +// Config: coordinator settings, read only from the environment, so the same +// binary behaves identically in CI, local, and prod. +package infra + +import ( + "errors" + "fmt" + "io/fs" + "math" + "os" + "strconv" + "time" + + "github.com/joho/godotenv" +) + +// defaultEnvFile is loaded by Load unless ENV_FILE points elsewhere. +const defaultEnvFile = ".env" + +type Config struct { + // HTTP listen address, e.g. ":8080". + Addr string + // PostgreSQL connection string (pgx format / libpq URL). + DatabaseURL string + // Shared bearer token workers must present. Empty disables auth (dev only). + Token string + + // Minimum log level: debug, info, warn, error. + LogLevel string + // Path to a rotated log file. Empty logs to stdout only. + 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 + // How long to keep retrying the initial database connection at startup + // before giving up. Covers a Postgres container that is still booting. + DBConnectTimeout time.Duration + // Per-request context timeout applied to handlers and DB calls. + RequestTimeout time.Duration + + // Suggested heartbeat cadence returned to workers on registration. + HeartbeatInterval time.Duration + // Default lease length handed out on claim. + LeaseDuration time.Duration + // Default attempt ceiling for newly created tasks. + 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 +// or malformed, so a misconfigured process never limps along half-wired. +// +// A .env file (path overridable via ENV_FILE) is loaded first as a local-dev +// convenience. It only fills variables the environment does not already define. +func LoadConfig() (Config, error) { + envFile := os.Getenv("ENV_FILE") + if envFile == "" { + envFile = defaultEnvFile + } + // godotenv.Load never overwrites variables already present in the + // environment, so an orchestrator's values always beat the file. A missing + // file is expected in production, where env vars are injected directly. + if err := godotenv.Load(envFile); err != nil && !errors.Is(err, fs.ErrNotExist) { + return Config{}, fmt.Errorf("load env file %q: %w", envFile, err) + } + + cfg := Config{ + Addr: getEnv("COORDINATOR_ADDR", ":8080"), + DatabaseURL: os.Getenv("DATABASE_URL"), + // COORDINATOR_TOKEN is the contract name; WORKER_AUTH_TOKEN is the + // former name, still honoured so existing .env files keep working. + Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")), + 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, + HeartbeatInterval: 15 * time.Second, + LeaseDuration: 2 * time.Minute, + DefaultMaxAttempts: 3, + ReaperInterval: 30 * time.Second, + WorkerOfflineAfter: 1 * time.Minute, + } + + if cfg.DatabaseURL == "" { + return Config{}, fmt.Errorf("DATABASE_URL is required") + } + + var err error + if cfg.DBMaxConns, err = getEnvInt32("DB_MAX_CONNS", cfg.DBMaxConns); err != nil { + return Config{}, err + } + 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 + } + if cfg.HeartbeatInterval, err = getEnvDuration("HEARTBEAT_INTERVAL", cfg.HeartbeatInterval); err != nil { + return Config{}, err + } + if cfg.LeaseDuration, err = getEnvDuration("LEASE_DURATION", cfg.LeaseDuration); err != nil { + return Config{}, err + } + 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 + } + + return cfg, nil +} + +func getEnv(key, def string) string { + if v := os.Getenv(key); v != "" { + return v + } + return def +} + +func getEnvInt(key string, def int) (int, error) { + v := os.Getenv(key) + if v == "" { + return def, nil + } + n, err := strconv.Atoi(v) + if err != nil { + return 0, fmt.Errorf("%s: %w", key, err) + } + return n, nil +} + +func getEnvInt32(key string, def int32) (int32, error) { + n, err := getEnvInt(key, int(def)) + if err != nil { + return 0, err + } + // On 64-bit builds int is wider than int32, so an oversized value would + // wrap silently — DB_MAX_CONNS=2147483648 becoming a negative pool size. + if n < math.MinInt32 || n > math.MaxInt32 { + return 0, fmt.Errorf("%s: %d is out of range for int32", key, n) + } + 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 == "" { + return def, nil + } + d, err := time.ParseDuration(v) + if err != nil { + return 0, fmt.Errorf("%s: %w", key, err) + } + return d, nil +} diff --git a/coordinator/internal/infra/db.go b/coordinator/internal/infra/db.go new file mode 100644 index 0000000..d1135cd --- /dev/null +++ b/coordinator/internal/infra/db.go @@ -0,0 +1,65 @@ +// DB: the PostgreSQL connection pool. +package infra + +import ( + "context" + "log/slog" + "time" + + "github.com/cenkalti/backoff/v4" + "github.com/jackc/pgx/v5/pgxpool" +) + +// NewPool builds the single shared pool. The caller owns its lifetime and must +// Close() it on shutdown. +func NewPool(ctx context.Context, cfg Config, log *slog.Logger) (*pgxpool.Pool, error) { + poolCfg, err := pgxpool.ParseConfig(cfg.DatabaseURL) + if err != nil { + return nil, err + } + poolCfg.MaxConns = cfg.DBMaxConns + + pool, err := pgxpool.NewWithConfig(ctx, poolCfg) + if err != nil { + return nil, err + } + // pgxpool.New is lazy, so a ping is needed to actually reach the server. + // It is retried because at startup — especially under docker-compose, where + // the coordinator can boot before Postgres is accepting connections — a + // service should wait for its database rather than crash-loop. + if err := pingWithRetry(ctx, pool, cfg.DBConnectTimeout, log); err != nil { + pool.Close() + return nil, err + } + return pool, nil +} + +// pingWithRetry waits for the database to accept connections, backing off +// between attempts until the budget elapses or ctx is cancelled. +// +// Unlike the transaction retry in storage/postgres, this retries *any* ping +// error: at startup a "connection refused" is the expected, retryable state, +// not an anomaly. +func pingWithRetry(ctx context.Context, pool *pgxpool.Pool, budget time.Duration, log *slog.Logger) error { + b := backoff.NewExponentialBackOff() + b.InitialInterval = 200 * time.Millisecond + b.MaxInterval = 3 * time.Second + b.MaxElapsedTime = budget + + attempt := 0 + return backoff.RetryNotify( + func() error { + // A bounded per-attempt timeout so one hung dial cannot eat the + // whole budget in a single try. + pingCtx, cancel := context.WithTimeout(ctx, 3*time.Second) + defer cancel() + return pool.Ping(pingCtx) + }, + backoff.WithContext(b, ctx), + func(err error, next time.Duration) { + attempt++ + log.Warn("database not ready, retrying", + "attempt", attempt, "retry_in", next.String(), "err", err) + }, + ) +} diff --git a/coordinator/internal/infra/logging.go b/coordinator/internal/infra/logging.go new file mode 100644 index 0000000..b5f49e3 --- /dev/null +++ b/coordinator/internal/infra/logging.go @@ -0,0 +1,65 @@ +package infra + +import ( + "fmt" + "io" + "log/slog" + "os" + "path/filepath" + "strings" + + "gopkg.in/natefinch/lumberjack.v2" +) + +// NewLogger builds the process logger. +// +// It always writes JSON to stdout, so `docker logs` and any 12-factor log +// collector keep working. When LogFile is set it *also* writes to a +// size-rotated file, so logs survive a container rebuild instead of vanishing +// with the previous stdout stream. Rotation is delegated to lumberjack rather +// than hand-rolled. +// +// The returned Closer flushes and closes the file; call it on shutdown. +func NewLogger(cfg Config) (*slog.Logger, io.Closer, error) { + opts := &slog.HandlerOptions{Level: parseLevel(cfg.LogLevel)} + + var ( + out io.Writer = os.Stdout + closer io.Closer = noopCloser{} + ) + + if cfg.LogFile != "" { + if err := os.MkdirAll(filepath.Dir(cfg.LogFile), 0o750); err != nil { + return nil, nil, fmt.Errorf("create log directory: %w", err) + } + rotator := &lumberjack.Logger{ + Filename: cfg.LogFile, + MaxSize: 50, // megabytes before a rotation + MaxBackups: 5, // keep this many rotated files + MaxAge: 30, // days + Compress: true, + } + // Tee to both: the console stays live while the file is the durable copy. + out = io.MultiWriter(os.Stdout, rotator) + closer = rotator + } + + return slog.New(slog.NewJSONHandler(out, opts)), closer, nil +} + +func parseLevel(s string) slog.Level { + switch strings.ToLower(strings.TrimSpace(s)) { + case "debug": + return slog.LevelDebug + case "warn", "warning": + return slog.LevelWarn + case "error": + return slog.LevelError + default: + return slog.LevelInfo + } +} + +type noopCloser struct{} + +func (noopCloser) Close() error { return nil } diff --git a/coordinator/internal/infra/server.go b/coordinator/internal/infra/server.go new file mode 100644 index 0000000..9858c19 --- /dev/null +++ b/coordinator/internal/infra/server.go @@ -0,0 +1,74 @@ +// Server: the HTTP listener and the background lease reaper, both shut down +// cleanly on a signal. +package infra + +import ( + "context" + "errors" + "log/slog" + "net/http" + "time" +) + +const shutdownGrace = 15 * time.Second + +// Run serves handler until ctx is cancelled, then drains in-flight requests. +func RunServer(ctx context.Context, log *slog.Logger, addr string, handler http.Handler) error { + srv := &http.Server{ + Addr: addr, + Handler: handler, + ReadHeaderTimeout: 5 * time.Second, + } + + // Buffered so this goroutine can exit even when nobody reads the channel + // (the ctx.Done branch below) — an unbuffered send would leak it forever. + errCh := make(chan error, 1) + go func() { + log.Info("coordinator listening", "addr", addr) + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + errCh <- err + } + }() + + select { + case err := <-errCh: + return err + case <-ctx.Done(): + log.Info("shutdown signal received") + } + + // A fresh context: ctx is already cancelled, and reusing it would abort the + // very requests we are trying to let finish. + shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownGrace) + defer cancel() + return srv.Shutdown(shutdownCtx) +} + +// RunReaper periodically reclaims tasks whose lease elapsed, so a worker that +// died without a heartbeat cannot strand its task in 'leased' forever. +// 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() + + for { + select { + case <-ctx.Done(): + return + case <-t.C: + n, err := fn(ctx) + if err != nil { + log.Debug(name+" skipped", "err", err) + continue + } + if n > 0 { + log.Info(name, "count", n) + } + } + } +} diff --git a/coordinator/internal/memstore/memstore.go b/coordinator/internal/memstore/memstore.go new file mode 100644 index 0000000..7916951 --- /dev/null +++ b/coordinator/internal/memstore/memstore.go @@ -0,0 +1,344 @@ +// Package memstore holds in-memory implementations of the usecase ports for +// tests: they exercise use-case orchestration without a database or filesystem. +// The real invariants that depend on Postgres (SKIP LOCKED, row locking) are +// covered separately by the integration tests. +package memstore + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "io" + "sort" + "sync" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// Clock returns a fixed, advanceable time. +type Clock struct{ t time.Time } + +func NewClock(t time.Time) *Clock { return &Clock{t: t} } +func (c *Clock) Now() time.Time { return c.t } +func (c *Clock) Advance(d time.Duration) { c.t = c.t.Add(d) } + +// Tx is a no-op transaction manager: the in-memory stores need no atomicity to +// be observed, so it simply runs the function. +type Tx struct{} + +func (Tx) WithinTx(ctx context.Context, fn func(ctx context.Context) error) error { return fn(ctx) } + +// --- TaskRepo ------------------------------------------------------------ + +type TaskRepo struct { + mu sync.Mutex + tasks map[uuid.UUID]*domain.Task +} + +func NewTaskRepo() *TaskRepo { return &TaskRepo{tasks: map[uuid.UUID]*domain.Task{}} } + +var _ usecase.TaskRepository = (*TaskRepo)(nil) + +// clone returns a copy so a caller's mutations do not touch stored state until +// Update — mirroring how a repository hands back detached entities. +func clone(t *domain.Task) *domain.Task { cp := *t; return &cp } + +func (r *TaskRepo) put(t *domain.Task) { + r.mu.Lock() + defer r.mu.Unlock() + r.tasks[t.ID] = clone(t) +} + +func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domain.Task, error) { + r.mu.Lock() + defer r.mu.Unlock() + + var cands []*domain.Task + for _, t := range r.tasks { + if t.Status != domain.TaskPending || t.Attempt >= t.MaxAttempts { + continue + } + if len(f.Workloads) > 0 && !contains(f.Workloads, t.Workload) { + continue + } + cands = append(cands, t) + } + if len(cands) == 0 { + return nil, nil + } + sort.Slice(cands, func(i, j int) bool { + if cands[i].CreatedAt.Equal(cands[j].CreatedAt) { + return cands[i].ChunkIndex < cands[j].ChunkIndex + } + return cands[i].CreatedAt.Before(cands[j].CreatedAt) + }) + + t := cands[0] + t.Status = domain.TaskLeased + t.Attempt++ + owner := f.Owner + t.LeaseOwner = &owner + t.LeaseExpiresAt = &f.LeaseUntil + if t.StartedAt == nil { + t.StartedAt = &f.Now + } + t.Version++ + return clone(t), nil +} + +func (r *TaskRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Task, error) { + r.mu.Lock() + defer r.mu.Unlock() + t, ok := r.tasks[id] + if !ok { + return nil, domain.ErrTaskNotFound + } + return clone(t), nil +} + +func (r *TaskRepo) GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task, error) { + return r.Get(ctx, id) +} + +func (r *TaskRepo) Update(ctx context.Context, t *domain.Task) error { + r.mu.Lock() + defer r.mu.Unlock() + stored, ok := r.tasks[t.ID] + if !ok || stored.Version != t.Version-1 { + return domain.ErrLeaseConflict // vanished or advanced under us + } + r.tasks[t.ID] = clone(t) + return nil +} + +func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error { + for _, t := range tasks { + r.put(t) + } + return nil +} + +func (r *TaskRepo) ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domain.Task, error) { + r.mu.Lock() + defer r.mu.Unlock() + var out []*domain.Task + for _, t := range r.tasks { + if t.JobID == jobID && t.Status == domain.TaskCompleted { + out = append(out, clone(t)) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex }) + return out, nil +} + +func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) { + r.mu.Lock() + defer r.mu.Unlock() + counts := map[domain.TaskStatus]int{} + for _, t := range r.tasks { + if t.JobID == jobID { + counts[t.Status]++ + } + } + return counts, nil +} + +func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + var n int64 + for _, t := range r.tasks { + if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) { + t.ExpireLease(now) + n++ + } + } + return n, nil +} + +// --- JobRepo ------------------------------------------------------------- + +type JobRepo struct { + mu sync.Mutex + jobs map[uuid.UUID]*domain.Job +} + +func NewJobRepo() *JobRepo { return &JobRepo{jobs: map[uuid.UUID]*domain.Job{}} } + +var _ usecase.JobRepository = (*JobRepo)(nil) + +func (r *JobRepo) Insert(ctx context.Context, j *domain.Job) error { + r.mu.Lock() + defer r.mu.Unlock() + cp := *j + r.jobs[j.ID] = &cp + return nil +} + +func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) { + r.mu.Lock() + defer r.mu.Unlock() + j, ok := r.jobs[id] + if !ok { + return nil, domain.ErrJobNotFound + } + cp := *j + return &cp, nil +} + +func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, status domain.JobStatus, completedAt *time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + j, ok := r.jobs[id] + if !ok { + return domain.ErrJobNotFound + } + j.Status = status + j.CompletedAt = completedAt + return nil +} + +// --- WorkerRepo ---------------------------------------------------------- + +type WorkerRepo struct { + mu sync.Mutex + workers map[uuid.UUID]*domain.Worker +} + +func NewWorkerRepo() *WorkerRepo { return &WorkerRepo{workers: map[uuid.UUID]*domain.Worker{}} } + +var _ usecase.WorkerRepository = (*WorkerRepo)(nil) + +func (r *WorkerRepo) Insert(ctx context.Context, w *domain.Worker) error { + r.mu.Lock() + defer r.mu.Unlock() + cp := *w + r.workers[w.ID] = &cp + return nil +} + +func (r *WorkerRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Worker, error) { + r.mu.Lock() + defer r.mu.Unlock() + w, ok := r.workers[id] + if !ok { + return nil, domain.ErrWorkerNotFound + } + cp := *w + 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 { + mu sync.Mutex + arts map[uuid.UUID]*domain.Artifact +} + +func NewArtifactRepo() *ArtifactRepo { return &ArtifactRepo{arts: map[uuid.UUID]*domain.Artifact{}} } + +var _ usecase.ArtifactRepository = (*ArtifactRepo)(nil) + +func (r *ArtifactRepo) Insert(ctx context.Context, a *domain.Artifact) error { + r.mu.Lock() + defer r.mu.Unlock() + cp := *a + r.arts[a.ID] = &cp + return nil +} + +func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error) { + r.mu.Lock() + defer r.mu.Unlock() + a, ok := r.arts[id] + if !ok { + return nil, domain.ErrArtifactNotFound + } + cp := *a + return &cp, nil +} + +// --- BlobStore ----------------------------------------------------------- + +type BlobStore struct { + mu sync.Mutex + blobs map[string][]byte +} + +func NewBlobStore() *BlobStore { return &BlobStore{blobs: map[string][]byte{}} } + +var _ usecase.BlobStore = (*BlobStore)(nil) + +func (b *BlobStore) Put(ctx context.Context, key string, r io.Reader) (string, int64, error) { + data, err := io.ReadAll(r) + if err != nil { + return "", 0, err + } + sum := sha256.Sum256(data) + b.mu.Lock() + b.blobs[key] = data + b.mu.Unlock() + return hex.EncodeToString(sum[:]), int64(len(data)), nil +} + +func (b *BlobStore) Open(ctx context.Context, key string) (io.ReadCloser, error) { + b.mu.Lock() + defer b.mu.Unlock() + data, ok := b.blobs[key] + if !ok { + return nil, domain.ErrArtifactNotFound + } + return io.NopCloser(bytes.NewReader(data)), nil +} + +func (b *BlobStore) Delete(ctx context.Context, key string) error { + b.mu.Lock() + defer b.mu.Unlock() + delete(b.blobs, key) + return nil +} + +// Has reports whether a blob exists — handy for asserting cleanup in tests. +func (b *BlobStore) Has(key string) bool { + b.mu.Lock() + defer b.mu.Unlock() + _, ok := b.blobs[key] + return ok +} + +func contains(ss []string, s string) bool { + for _, x := range ss { + if x == s { + return true + } + } + return false +} diff --git a/coordinator/internal/storage/blob/store.go b/coordinator/internal/storage/blob/store.go new file mode 100644 index 0000000..db98600 --- /dev/null +++ b/coordinator/internal/storage/blob/store.go @@ -0,0 +1,130 @@ +// Package blob stores artifact bytes on the local filesystem. It implements +// usecase.BlobStore; no other layer knows where or how the bytes are kept. +package blob + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "os" + "path/filepath" + "strings" + + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// FSStore keeps each artifact as one file under dir, named by its storage key. +type FSStore struct { + dir string + staging string +} + +var _ usecase.BlobStore = (*FSStore)(nil) + +// NewFSStore prepares the storage and staging directories. Staging lives inside +// dir so a finished file can be renamed into place on the same filesystem — +// rename is only atomic within one filesystem. +func NewFSStore(dir string) (*FSStore, error) { + staging := filepath.Join(dir, ".staging") + if err := os.MkdirAll(staging, 0o750); err != nil { + return nil, fmt.Errorf("create blob dirs: %w", err) + } + return &FSStore{dir: dir, staging: staging}, nil +} + +// Put streams r to a staging file while hashing it, then atomically renames it +// into place. A caller that dies mid-upload leaves at most a staging temp file, +// never a half-written artifact that looks complete. +func (s *FSStore) Put(ctx context.Context, key string, r io.Reader) (string, int64, error) { + if err := checkKey(key); err != nil { + return "", 0, err + } + + tmp, err := os.CreateTemp(s.staging, key+"-*") + if err != nil { + return "", 0, fmt.Errorf("create staging file: %w", err) + } + tmpName := tmp.Name() + // On any failure past this point, do not leave the temp file behind. + defer func() { + if tmpName != "" { + _ = os.Remove(tmpName) + } + }() + + h := sha256.New() + // Tee the stream: one copy to disk, one to the hasher, in a single pass so + // the bytes are never held in memory or read twice. + size, err := io.Copy(io.MultiWriter(tmp, h), &ctxReader{ctx: ctx, r: r}) + if err != nil { + _ = tmp.Close() + return "", 0, fmt.Errorf("write artifact: %w", err) + } + // fsync before rename so a crash cannot leave a renamed-but-empty file. + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return "", 0, fmt.Errorf("sync artifact: %w", err) + } + if err := tmp.Close(); err != nil { + return "", 0, fmt.Errorf("close artifact: %w", err) + } + + final := filepath.Join(s.dir, key) + if err := os.Rename(tmpName, final); err != nil { + return "", 0, fmt.Errorf("commit artifact: %w", err) + } + tmpName = "" // committed — the deferred cleanup must not delete it now + + return hex.EncodeToString(h.Sum(nil)), size, nil +} + +// Open returns the artifact bytes for streaming to a client. The caller closes. +func (s *FSStore) Open(ctx context.Context, key string) (io.ReadCloser, error) { + if err := checkKey(key); err != nil { + return nil, err + } + // checkKey has rejected any traversal, so the joined path stays under s.dir. + f, err := os.Open(filepath.Join(s.dir, key)) //nolint:gosec // key validated by checkKey + + if err != nil { + return nil, err + } + return f, nil +} + +// Delete removes a stored blob. Absence is not an error: cleaning up after a +// failed metadata insert must be idempotent. +func (s *FSStore) Delete(ctx context.Context, key string) error { + if err := checkKey(key); err != nil { + return err + } + if err := os.Remove(filepath.Join(s.dir, key)); err != nil && !os.IsNotExist(err) { + return err + } + return nil +} + +// checkKey rejects anything that could escape the storage directory. Keys are +// coordinator-generated UUIDs, so this is defence in depth, not the only guard. +func checkKey(key string) error { + if key == "" || strings.ContainsAny(key, `/\`) || strings.Contains(key, "..") { + return fmt.Errorf("invalid storage key %q", key) + } + return nil +} + +// ctxReader aborts a copy when the request context is cancelled, so a stalled +// or disconnected upload does not tie up a file handle indefinitely. +type ctxReader struct { + ctx context.Context + r io.Reader +} + +func (c *ctxReader) Read(p []byte) (int, error) { + if err := c.ctx.Err(); err != nil { + return 0, err + } + return c.r.Read(p) +} diff --git a/coordinator/internal/storage/blob/store_test.go b/coordinator/internal/storage/blob/store_test.go new file mode 100644 index 0000000..5f57294 --- /dev/null +++ b/coordinator/internal/storage/blob/store_test.go @@ -0,0 +1,115 @@ +package blob + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "io" + "os" + "path/filepath" + "strings" + "testing" +) + +func newStore(t *testing.T) *FSStore { + t.Helper() + s, err := NewFSStore(t.TempDir()) + if err != nil { + t.Fatalf("NewFSStore: %v", err) + } + return s +} + +func TestPutComputesChecksumAndSize(t *testing.T) { + s := newStore(t) + data := bytes.Repeat([]byte("chembl-row\n"), 10000) // ~110 KB, streamed + + sum, size, err := s.Put(context.Background(), "key-1", bytes.NewReader(data)) + if err != nil { + t.Fatalf("Put: %v", err) + } + + want := sha256.Sum256(data) + if sum != hex.EncodeToString(want[:]) { + t.Errorf("sha256 = %s, want %s", sum, hex.EncodeToString(want[:])) + } + if size != int64(len(data)) { + t.Errorf("size = %d, want %d", size, len(data)) + } +} + +func TestPutThenOpenRoundTrips(t *testing.T) { + s := newStore(t) + data := []byte("partial result csv\n1,2,3\n") + + if _, _, err := s.Put(context.Background(), "key-2", bytes.NewReader(data)); err != nil { + t.Fatalf("Put: %v", err) + } + + rc, err := s.Open(context.Background(), "key-2") + if err != nil { + t.Fatalf("Open: %v", err) + } + defer rc.Close() + + got, _ := io.ReadAll(rc) + if !bytes.Equal(got, data) { + t.Errorf("round-trip mismatch: got %q", got) + } +} + +func TestPutLeavesNoStagingFileBehind(t *testing.T) { + s := newStore(t) + if _, _, err := s.Put(context.Background(), "key-3", strings.NewReader("x")); err != nil { + t.Fatalf("Put: %v", err) + } + + entries, _ := os.ReadDir(s.staging) + if len(entries) != 0 { + t.Errorf("staging dir not empty after a successful put: %v", entries) + } +} + +func TestPutFailureLeavesNoArtifactOrStaging(t *testing.T) { + s := newStore(t) + // A reader that errors partway through simulates a dropped upload. + r := io.MultiReader(strings.NewReader("half"), &erroringReader{}) + + if _, _, err := s.Put(context.Background(), "key-4", r); err == nil { + t.Fatal("expected an error from a failing reader") + } + + if _, err := os.Stat(filepath.Join(s.dir, "key-4")); !os.IsNotExist(err) { + t.Error("a failed put must not leave a committed artifact") + } + if entries, _ := os.ReadDir(s.staging); len(entries) != 0 { + t.Errorf("a failed put must not leave staging files: %v", entries) + } +} + +func TestPutRejectsUnsafeKeys(t *testing.T) { + s := newStore(t) + for _, key := range []string{"", "../escape", "a/b", `a\b`, "with..dots"} { + if _, _, err := s.Put(context.Background(), key, strings.NewReader("x")); err == nil { + t.Errorf("key %q should have been rejected", key) + } + } +} + +func TestPutHonoursContextCancellation(t *testing.T) { + s := newStore(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() // already cancelled before the copy starts + + if _, _, err := s.Put(ctx, "key-5", strings.NewReader("data")); err == nil { + t.Fatal("expected cancellation to abort the put") + } + if _, err := os.Stat(filepath.Join(s.dir, "key-5")); !os.IsNotExist(err) { + t.Error("a cancelled put must not leave an artifact") + } +} + +type erroringReader struct{} + +func (*erroringReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } diff --git a/coordinator/internal/storage/postgres/artifact_repo.go b/coordinator/internal/storage/postgres/artifact_repo.go new file mode 100644 index 0000000..771b5ee --- /dev/null +++ b/coordinator/internal/storage/postgres/artifact_repo.go @@ -0,0 +1,72 @@ +package postgres + +import ( + "context" + "errors" + "fmt" + + sq "github.com/Masterminds/squirrel" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// ArtifactRepo implements usecase.ArtifactRepository. +type ArtifactRepo struct { + pool *pgxpool.Pool +} + +func NewArtifactRepo(pool *pgxpool.Pool) *ArtifactRepo { + return &ArtifactRepo{pool: pool} +} + +var _ usecase.ArtifactRepository = (*ArtifactRepo)(nil) + +var artifactColumns = []string{ + "id", "job_id", "task_id", "kind", "filename", "storage_key", + "content_type", "size_bytes", "sha256", "created_at", +} + +func (r *ArtifactRepo) Insert(ctx context.Context, a *domain.Artifact) error { + sql, args, err := psql.Insert("artifacts"). + Columns(artifactColumns...). + Values(a.ID, a.JobID, a.TaskID, string(a.Kind), a.Filename, a.StorageKey, + a.ContentType, a.SizeBytes, a.SHA256, a.CreatedAt). + ToSql() + if err != nil { + return err + } + if _, err := conn(ctx, r.pool).Exec(ctx, sql, args...); err != nil { + return fmt.Errorf("insert artifact: %w", err) + } + return nil +} + +func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error) { + sql, args, err := psql.Select(artifactColumns...). + From("artifacts"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + + var ( + a domain.Artifact + kind string + ) + err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( + &a.ID, &a.JobID, &a.TaskID, &kind, &a.Filename, &a.StorageKey, + &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrArtifactNotFound + } + if err != nil { + return nil, fmt.Errorf("get artifact: %w", err) + } + a.Kind = domain.ArtifactKind(kind) + return &a, nil +} diff --git a/coordinator/internal/storage/postgres/builder.go b/coordinator/internal/storage/postgres/builder.go new file mode 100644 index 0000000..0025347 --- /dev/null +++ b/coordinator/internal/storage/postgres/builder.go @@ -0,0 +1,11 @@ +package postgres + +import sq "github.com/Masterminds/squirrel" + +// psql is the shared statement builder, fixed to PostgreSQL $N placeholders so +// no call site repeats PlaceholderFormat(sq.Dollar). +// +// Not everything goes through it. Two genuinely set-based statements stay as +// raw SQL — claimNext (a FOR UPDATE SKIP LOCKED CTE) and expireLeases (CASE +// logic in the SET) — because a builder would obscure them, not clarify them. +var psql = sq.StatementBuilder.PlaceholderFormat(sq.Dollar) diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go new file mode 100644 index 0000000..fc917d6 --- /dev/null +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -0,0 +1,516 @@ +//go:build integration + +// Integration tests run against a real PostgreSQL instance supplied through +// TEST_DATABASE_URL. The spec forbids mocks or SQLite here: the guarantees +// being verified — FOR UPDATE SKIP LOCKED, optimistic concurrency, transaction +// rollback — are properties of Postgres, not of our Go code. +// +// docker compose up -d +// TEST_DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable' \ +// go test -tags=integration ./internal/storage/postgres/ -v +package postgres + +import ( + "context" + "errors" + "fmt" + "os" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +func testPool(t *testing.T) *pgxpool.Pool { + t.Helper() + url := os.Getenv("TEST_DATABASE_URL") + if url == "" { + t.Skip("TEST_DATABASE_URL is not set") + } + pool, err := pgxpool.New(context.Background(), url) + if err != nil { + t.Fatalf("connect: %v", err) + } + t.Cleanup(pool.Close) + return pool +} + +// seedJob creates a job with n pending tasks and removes them afterwards, so +// tests stay independent of each other and of leftovers from earlier runs. +func seedJob(t *testing.T, pool *pgxpool.Pool, n int) (*domain.Job, []*domain.Task) { + t.Helper() + ctx := context.Background() + + chunks := make([]domain.ChunkSpec, 0, n) + for i := 0; i < n; i++ { + chunks = append(chunks, domain.ChunkSpec{ + ChunkIndex: i, + InputURI: fmt.Sprintf("s3://chunk-%d", i), + InputSHA256: fmt.Sprintf("sha-%d", i), + }) + } + job, tasks, err := domain.NewJobWithTasks("similarity_search", "s3://ds", nil, chunks, time.Now().UTC()) + if err != nil { + t.Fatalf("build job: %v", err) + } + + jobs, taskRepo, tx := NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool) + err = tx.WithinTx(ctx, func(ctx context.Context) error { + if err := jobs.Insert(ctx, job); err != nil { + return err + } + return taskRepo.InsertBatch(ctx, tasks) + }) + if err != nil { + t.Fatalf("seed: %v", err) + } + + t.Cleanup(func() { + // ON DELETE CASCADE removes the tasks with it. + _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, job.ID) + }) + return job, tasks +} + +func TestCreateJobPersistsEveryTask(t *testing.T) { + pool := testPool(t) + job, _ := seedJob(t, pool, 3) + + counts, err := NewTaskRepo(pool).CountByStatus(context.Background(), job.ID) + if err != nil { + t.Fatalf("count: %v", err) + } + if counts[domain.TaskPending] != 3 { + t.Errorf("pending = %d, want 3", counts[domain.TaskPending]) + } +} + +// A job must land whole or not at all: a half-created job leaves chunks no +// worker could ever complete. +func TestCreateJobRollsBackOnFailure(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + chunks := []domain.ChunkSpec{{ChunkIndex: 0, InputURI: "s3://c0", InputSHA256: "sha0"}} + job, tasks, err := domain.NewJobWithTasks("similarity_search", "s3://ds", nil, chunks, time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + + jobs, taskRepo, tx := NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool) + boom := errors.New("boom") + err = tx.WithinTx(ctx, func(ctx context.Context) error { + if err := jobs.Insert(ctx, job); err != nil { + return err + } + if err := taskRepo.InsertBatch(ctx, tasks); err != nil { + return err + } + return boom // fail after both writes + }) + if !errors.Is(err, boom) { + t.Fatalf("err = %v, want boom", err) + } + + if _, err := jobs.Get(ctx, job.ID); !errors.Is(err, domain.ErrJobNotFound) { + t.Errorf("job survived the rollback: %v", err) + } +} + +// The acceptance criterion: N workers claiming at once must each get a +// different task, and no task may be handed out twice. +func TestConcurrentClaimGivesEachTaskToExactlyOneWorker(t *testing.T) { + pool := testPool(t) + const tasks = 8 + job, _ := seedJob(t, pool, tasks) + + repo := NewTaskRepo(pool) + now := time.Now().UTC() + + var ( + mu sync.Mutex + claimed = make(map[uuid.UUID]string) + wg sync.WaitGroup + ) + // More workers than tasks, so the surplus must come back empty rather than + // steal an already-leased row. + for i := 0; i < tasks*2; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{ + Owner: fmt.Sprintf("worker-%d", n), + Now: now, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Errorf("claim: %v", err) + return + } + if task == nil || task.JobID != job.ID { + return // empty queue, or a task from another test's job + } + mu.Lock() + defer mu.Unlock() + if prev, dup := claimed[task.ID]; dup { + t.Errorf("task %s handed to both %s and worker-%d", task.ID, prev, n) + } + claimed[task.ID] = fmt.Sprintf("worker-%d", n) + }(i) + } + wg.Wait() + + if len(claimed) != tasks { + t.Errorf("claimed %d tasks, want %d", len(claimed), tasks) + } +} + +func TestClaimNextReturnsNilOnEmptyQueue(t *testing.T) { + pool := testPool(t) + now := time.Now().UTC() + + // Drain everything first, then ask once more. + repo := NewTaskRepo(pool) + for { + task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{ + Owner: "drainer", Now: now, LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Fatalf("drain: %v", err) + } + if task == nil { + break + } + } + + task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{ + Owner: "worker-1", Now: now, LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Fatalf("claim: %v", err) + } + if task != nil { + t.Errorf("expected nil on an empty queue, got %s", task.ID) + } +} + +func TestUpdateRejectsStaleVersion(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 1) + + repo, tx := NewTaskRepo(pool), NewTxManager(pool) + now := time.Now().UTC() + + task, err := repo.ClaimNext(ctx, usecase.ClaimFilter{ + Owner: "worker-1", Now: now, LeaseUntil: now.Add(time.Minute), + }) + if err != nil || task == nil || task.JobID != job.ID { + t.Skipf("could not claim this job's task (got %v, %v)", task, err) + } + + // A stale copy: same row, but the version it remembers is behind. + stale := *task + stale.Version = task.Version // pretend the caller mutated it once + + err = tx.WithinTx(ctx, func(ctx context.Context) error { + fresh, err := repo.GetForUpdate(ctx, task.ID) + if err != nil { + return err + } + if err := fresh.RenewLease("worker-1", fresh.Attempt, now.Add(2*time.Minute)); err != nil { + return err + } + return repo.Update(ctx, fresh) + }) + if err != nil { + t.Fatalf("legitimate update failed: %v", err) + } + + // Now the stale copy's version is behind by one; its write must be refused. + stale.Version++ // as a domain method would have done + if err := repo.Update(ctx, &stale); !errors.Is(err, domain.ErrLeaseConflict) { + t.Errorf("stale update err = %v, want ErrLeaseConflict", err) + } +} + +func TestListCompletedIsOrderedByChunkIndex(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, tasks := seedJob(t, pool, 4) + + repo, artifacts, tx := NewTaskRepo(pool), NewArtifactRepo(pool), NewTxManager(pool) + now := time.Now().UTC() + + // Complete them out of order to prove the ordering comes from SQL. + for _, i := range []int{2, 0, 3, 1} { + task := tasks[i] + err := tx.WithinTx(ctx, func(ctx context.Context) error { + // A completed task must reference a real result artifact (FK + check). + taskID := task.ID + art, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, + fmt.Sprintf("result-%d.csv", task.ChunkIndex), "text/csv", now) + if err != nil { + return err + } + art.SetContent(fmt.Sprintf("rsha-%d", task.ChunkIndex), 1) + if err := artifacts.Insert(ctx, art); err != nil { + return err + } + + fresh, err := repo.GetForUpdate(ctx, task.ID) + if err != nil { + return err + } + owner := "worker-1" + fresh.Status = domain.TaskLeased + fresh.LeaseOwner = &owner + expires := now.Add(time.Minute) + fresh.LeaseExpiresAt = &expires + if err := fresh.CompleteWith(art.ID, nil, owner, fresh.Attempt, now); err != nil { + return err + } + return repo.Update(ctx, fresh) + }) + if err != nil { + t.Fatalf("complete chunk %d: %v", i, err) + } + } + + done, err := repo.ListCompleted(ctx, job.ID) + if err != nil { + t.Fatalf("list: %v", err) + } + if len(done) != 4 { + t.Fatalf("got %d completed, want 4", len(done)) + } + for i, task := range done { + if task.ChunkIndex != i { + t.Errorf("position %d holds chunk_index %d — order is not deterministic", i, task.ChunkIndex) + } + } +} + +// A worker whose network dropped resends the same manifest. That must succeed: +// the entity is unchanged, so nothing is written, and the optimistic-concurrency +// guard must not turn the replay into a conflict. +func TestCompleteTaskReplayIsIdempotent(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 1) + + tasks, jobs, artifacts, tx := NewTaskRepo(pool), NewJobRepo(pool), NewArtifactRepo(pool), NewTxManager(pool) + clk := fixedClock{now: time.Now().UTC()} + uc := usecase.NewCompleteTask(tasks, jobs, artifacts, tx, clk) + + claimed, err := tasks.ClaimNext(ctx, usecase.ClaimFilter{ + Owner: "worker-1", Now: clk.now, LeaseUntil: clk.now.Add(time.Minute), + }) + if err != nil || claimed == nil || claimed.JobID != job.ID { + t.Skipf("could not claim this job's task (got %v, %v)", claimed, err) + } + + // A partial-result artifact the coordinator stored for this task. + art := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult) + + in := usecase.CompleteTaskInput{ + TaskID: claimed.ID, WorkerID: "worker-1", Attempt: claimed.Attempt, + ResultArtifactID: art.ID, + } + if _, err := uc.Execute(ctx, in); err != nil { + t.Fatalf("first submission: %v", err) + } + if _, err := uc.Execute(ctx, in); err != nil { + t.Errorf("replay must be idempotent, got %v", err) + } + + // A different result artifact for the same task is a genuine conflict. + art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult) + other := in + other.ResultArtifactID = art2.ID + if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) { + t.Errorf("err = %v, want ErrResultConflict", err) + } +} + +type fixedClock struct{ now time.Time } + +func (c fixedClock) Now() time.Time { return c.now } + +// seedArtifact inserts an artifact and returns it, cleaned up with its job. +func seedArtifact(t *testing.T, pool *pgxpool.Pool, jobID uuid.UUID, taskID *uuid.UUID, kind domain.ArtifactKind) *domain.Artifact { + t.Helper() + art, err := domain.NewArtifact(jobID, taskID, kind, "f.csv", "text/csv", time.Now().UTC()) + if err != nil { + t.Fatalf("build artifact: %v", err) + } + art.SetContent(fmt.Sprintf("sha-%s", art.ID), 3) + if err := NewArtifactRepo(pool).Insert(context.Background(), art); err != nil { + t.Fatalf("insert artifact: %v", err) + } + return art +} + +func TestWorkerRepoRoundTrip(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + repo := NewWorkerRepo(pool) + + w, err := domain.NewWorker("lab-int", []string{"similarity_search", "similarity_graph"}, time.Now().UTC()) + 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) }) + + got, err := repo.Get(ctx, w.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.Status != domain.WorkerOnline || len(got.Capabilities) != 2 { + t.Errorf("round-trip mismatch: %+v", got) + } + // capabilities must survive the jsonb round-trip. + if got.Capabilities[0] != "similarity_search" { + t.Errorf("capabilities = %v", got.Capabilities) + } + + if _, err := repo.Get(ctx, uuid.New()); !errors.Is(err, domain.ErrWorkerNotFound) { + t.Errorf("missing worker err = %v, want ErrWorkerNotFound", err) + } +} + +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() + job, _ := seedJob(t, pool, 1) + + art := seedArtifact(t, pool, job.ID, nil, domain.ArtifactInput) + got, err := NewArtifactRepo(pool).Get(ctx, art.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.Kind != domain.ArtifactInput || got.StorageKey != art.StorageKey || got.SizeBytes != 3 { + t.Errorf("round-trip mismatch: %+v", got) + } + if _, err := NewArtifactRepo(pool).Get(ctx, uuid.New()); !errors.Is(err, domain.ErrArtifactNotFound) { + t.Errorf("missing artifact err = %v, want ErrArtifactNotFound", err) + } +} + +// A shard task stores its input as an artifact and no URI: this exercises the +// nullable input_uri column, the input_artifact_id round-trip, and the +// ck_tasks_has_input check that requires one or the other. +func TestShardTaskRoundTrip(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + job, err := domain.NewUploadedJob("similarity_search", nil, time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + jobs, taskRepo, tx := NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool) + if err := jobs.Insert(ctx, job); err != nil { + t.Fatalf("insert job: %v", err) + } + t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, job.ID) }) + + shard := seedArtifact(t, pool, job.ID, nil, domain.ArtifactShard) + task, err := domain.NewShardTask(job.ID, 0, "similarity_search", shard.ID, shard.SHA256, nil, 0, time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + if err := tx.WithinTx(ctx, func(ctx context.Context) error { + return taskRepo.InsertBatch(ctx, []*domain.Task{task}) + }); err != nil { + t.Fatalf("insert shard task: %v", err) + } + + got, err := taskRepo.Get(ctx, task.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.InputArtifactID == nil || *got.InputArtifactID != shard.ID { + t.Errorf("input_artifact_id did not round-trip: %v", got.InputArtifactID) + } + if got.InputURI != "" { + t.Errorf("shard task input_uri = %q, want empty (NULL)", got.InputURI) + } +} + +func TestExpireLeasesRequeuesElapsedTasks(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 1) + + repo := NewTaskRepo(pool) + past := time.Now().UTC().Add(-time.Hour) + + // Lease it with an expiry already in the past. + task, err := repo.ClaimNext(ctx, usecase.ClaimFilter{ + Owner: "dead-worker", Now: past, LeaseUntil: past.Add(time.Minute), + }) + if err != nil || task == nil || task.JobID != job.ID { + t.Skipf("could not claim this job's task (got %v, %v)", task, err) + } + + if _, err := repo.ExpireLeases(ctx, time.Now().UTC()); err != nil { + t.Fatalf("expire: %v", err) + } + + counts, err := repo.CountByStatus(ctx, job.ID) + if err != nil { + t.Fatalf("count: %v", err) + } + if counts[domain.TaskPending] != 1 { + t.Errorf("pending = %d, want 1 — a dead worker must not strand its task", counts[domain.TaskPending]) + } +} diff --git a/coordinator/internal/storage/postgres/job_repo.go b/coordinator/internal/storage/postgres/job_repo.go new file mode 100644 index 0000000..f2a7e2e --- /dev/null +++ b/coordinator/internal/storage/postgres/job_repo.go @@ -0,0 +1,91 @@ +package postgres + +import ( + "context" + "errors" + "time" + + sq "github.com/Masterminds/squirrel" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// JobRepo implements usecase.JobRepository. +type JobRepo struct { + pool *pgxpool.Pool +} + +func NewJobRepo(pool *pgxpool.Pool) *JobRepo { + return &JobRepo{pool: pool} +} + +var _ usecase.JobRepository = (*JobRepo)(nil) + +var jobColumns = []string{"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at"} + +// Insert runs inside the caller's transaction, alongside the job's tasks — that +// is what makes "all tasks or none" hold. +func (r *JobRepo) Insert(ctx context.Context, j *domain.Job) error { + sql, args, err := psql.Insert("jobs"). + Columns("id", "workload", "input_uri", "parameters", "status", "created_at"). + Values(j.ID, j.Workload, j.InputURI, jsonbOrEmpty(j.Parameters), string(j.Status), j.CreatedAt). + ToSql() + if err != nil { + return err + } + _, err = conn(ctx, r.pool).Exec(ctx, sql, args...) + return err +} + +func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) { + sql, args, err := psql.Select(jobColumns...). + From("jobs"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + + var ( + j domain.Job + status string + ) + err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( + &j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrJobNotFound + } + if err != nil { + return nil, err + } + j.Status = domain.JobStatus(status) + return &j, nil +} + +func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, + status domain.JobStatus, completedAt *time.Time) error { + + sql, args, err := psql.Update("jobs"). + SetMap(map[string]any{ + "status": string(status), + "completed_at": completedAt, + }). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return err + } + + tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return domain.ErrJobNotFound + } + return nil +} diff --git a/coordinator/internal/storage/postgres/retry.go b/coordinator/internal/storage/postgres/retry.go new file mode 100644 index 0000000..8f0ae86 --- /dev/null +++ b/coordinator/internal/storage/postgres/retry.go @@ -0,0 +1,82 @@ +package postgres + +import ( + "context" + "errors" + "time" + + "github.com/cenkalti/backoff/v4" + "github.com/jackc/pgx/v5/pgconn" +) + +// Transient PostgreSQL failures. Under concurrent claiming these are expected +// rather than exceptional: two coordinators touching neighbouring rows can +// deadlock or fail to serialize, and the correct response is to try again. +const ( + codeSerializationFailure = "40001" + codeDeadlockDetected = "40P01" + codeTooManyConnections = "53300" + codeCannotConnectNow = "57P03" +) + +// Retry budget: short and bounded. A worker polling for tasks would rather get +// a fast error and poll again than have its request hang for half a minute. +const ( + retryInitialInterval = 50 * time.Millisecond + retryMaxInterval = 1 * time.Second + retryMaxElapsedTime = 5 * time.Second +) + +// isTransient reports whether err is worth retrying. +// +// The default is *not* to retry: a constraint violation or a syntax error will +// fail identically every time, and retrying it only multiplies the damage. +func isTransient(err error) bool { + if err == nil { + return false + } + // A cancelled caller does not want another attempt. + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) { + switch pgErr.Code { + case codeSerializationFailure, codeDeadlockDetected, + codeTooManyConnections, codeCannotConnectNow: + return true + default: + return false + } + } + + // Connection-level trouble (dropped socket, closed pool). pgconn knows + // whether the query could have been executed before the failure — retrying + // a maybe-executed write would risk duplicating it. + return pgconn.SafeToRetry(err) +} + +// withRetry runs op, retrying only transient database failures with +// exponential backoff and jitter, and giving up as soon as ctx is done. +// +// Jitter matters here: without it, several coordinators that collide once will +// retry in lockstep and collide again at exactly the same moment. +func withRetry(ctx context.Context, op func(context.Context) error) error { + b := backoff.NewExponentialBackOff() + b.InitialInterval = retryInitialInterval + b.MaxInterval = retryMaxInterval + b.MaxElapsedTime = retryMaxElapsedTime + // RandomizationFactor defaults to 0.5, which is the jitter. + + return backoff.Retry(func() error { + err := op(ctx) + if err == nil { + return nil + } + if !isTransient(err) { + return backoff.Permanent(err) // stop now, do not burn the budget + } + return err + }, backoff.WithContext(b, ctx)) +} diff --git a/coordinator/internal/storage/postgres/retry_test.go b/coordinator/internal/storage/postgres/retry_test.go new file mode 100644 index 0000000..bdf77ec --- /dev/null +++ b/coordinator/internal/storage/postgres/retry_test.go @@ -0,0 +1,98 @@ +package postgres + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgconn" +) + +func TestIsTransient(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"serialization failure", &pgconn.PgError{Code: codeSerializationFailure}, true}, + {"deadlock", &pgconn.PgError{Code: codeDeadlockDetected}, true}, + {"too many connections", &pgconn.PgError{Code: codeTooManyConnections}, true}, + // A unique-violation repeats identically forever — retrying is pointless. + {"unique violation", &pgconn.PgError{Code: "23505"}, false}, + {"syntax error", &pgconn.PgError{Code: "42601"}, false}, + {"context cancelled", context.Canceled, false}, + {"deadline exceeded", context.DeadlineExceeded, false}, + {"unknown error", errors.New("boom"), false}, + // Wrapping must not hide the cause: errors.As walks the chain. + {"wrapped deadlock", errors2Wrap(&pgconn.PgError{Code: codeDeadlockDetected}), true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isTransient(tt.err); got != tt.want { + t.Errorf("isTransient(%v) = %v, want %v", tt.err, got, tt.want) + } + }) + } +} + +func errors2Wrap(err error) error { + return errors.Join(errors.New("query failed"), err) +} + +func TestWithRetrySucceedsAfterTransientFailures(t *testing.T) { + calls := 0 + err := withRetry(context.Background(), func(context.Context) error { + calls++ + if calls < 3 { + return &pgconn.PgError{Code: codeSerializationFailure} + } + return nil + }) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if calls != 3 { + t.Errorf("calls = %d, want 3", calls) + } +} + +func TestWithRetryStopsOnPermanentError(t *testing.T) { + permanent := &pgconn.PgError{Code: "23505"} // unique violation + calls := 0 + + err := withRetry(context.Background(), func(context.Context) error { + calls++ + return permanent + }) + + if !errors.Is(err, permanent) { + t.Errorf("err = %v, want the original error", err) + } + if calls != 1 { + t.Errorf("calls = %d, want 1 — a permanent error must not be retried", calls) + } +} + +func TestWithRetryHonoursContextCancellation(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + calls := 0 + start := time.Now() + err := withRetry(ctx, func(context.Context) error { + calls++ + return &pgconn.PgError{Code: codeDeadlockDetected} + }) + + if err == nil { + t.Fatal("expected an error once the context expired") + } + // Must abort at the deadline, not run the full 5s retry budget. + if elapsed := time.Since(start); elapsed > time.Second { + t.Errorf("took %v, expected to stop at the context deadline", elapsed) + } +} diff --git a/coordinator/internal/storage/postgres/task_repo.go b/coordinator/internal/storage/postgres/task_repo.go new file mode 100644 index 0000000..2a4480f --- /dev/null +++ b/coordinator/internal/storage/postgres/task_repo.go @@ -0,0 +1,330 @@ +package postgres + +import ( + "context" + "errors" + "strings" + "time" + + sq "github.com/Masterminds/squirrel" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// TaskRepo implements usecase.TaskRepository. +type TaskRepo struct { + pool *pgxpool.Pool +} + +func NewTaskRepo(pool *pgxpool.Pool) *TaskRepo { + return &TaskRepo{pool: pool} +} + +var _ usecase.TaskRepository = (*TaskRepo)(nil) + +// taskColumns is the single source of truth for the shape scanTask expects. +// Every query that returns a task selects exactly this list, in this order — +// three hand-written column lists would drift apart within a week. +var taskColumns = []string{ + "id", "job_id", "chunk_index", "workload", "input_uri", "input_artifact_id", "input_sha256", + "parameters", "status", "attempt", "max_attempts", "lease_owner", "lease_expires_at", + "result_artifact_id", "metrics", "error_code", "error_message", + "created_at", "started_at", "completed_at", "version", +} + +// taskColumnList is the same set as a comma string, for the raw claim query's +// RETURNING clause, which the builder does not touch. +var taskColumnList = strings.Join(taskColumns, ", ") + +// scanTask maps one row onto an entity. +// +// status is read into a plain string rather than domain.TaskStatus: pgx does +// not know the task_status enum, and going through string keeps the driver out +// of the domain's type system. +func scanTask(row pgx.Row) (*domain.Task, error) { + var ( + t domain.Task + status string + // input_uri is nullable now (uploaded shards have none), so it cannot + // scan straight into a string; NULL becomes the empty InputURI. + inputURI *string + ) + err := row.Scan( + &t.ID, &t.JobID, &t.ChunkIndex, &t.Workload, &inputURI, &t.InputArtifactID, &t.InputSHA256, + &t.Parameters, &status, &t.Attempt, &t.MaxAttempts, &t.LeaseOwner, &t.LeaseExpiresAt, + &t.ResultArtifactID, &t.Metrics, &t.ErrorCode, &t.ErrorMessage, + &t.CreatedAt, &t.StartedAt, &t.CompletedAt, &t.Version, + ) + if err != nil { + return nil, err + } + if inputURI != nil { + t.InputURI = *inputURI + } + t.Status = domain.TaskStatus(status) + return &t, nil +} + +// claimNextSQL leases one task in a single statement. +// +// Left as raw SQL on purpose: it is a data-modifying CTE with FOR UPDATE SKIP +// LOCKED, which no query builder expresses — and which is the whole point. +// SKIP LOCKED is what makes concurrent coordinators safe: each process locks a +// different candidate row instead of queueing on the same one, so no task is +// ever handed to two workers and no claim blocks behind another. Splitting this +// into SELECT + UPDATE would reintroduce exactly that race. +var claimNextSQL = ` +WITH candidate AS ( + SELECT id AS cid + FROM tasks + WHERE status = 'pending' + AND attempt < max_attempts + AND (cardinality($1::text[]) = 0 OR workload = ANY($1)) + ORDER BY created_at, chunk_index + FOR UPDATE SKIP LOCKED + LIMIT 1 +) +UPDATE tasks +SET status = 'leased', + attempt = attempt + 1, + lease_owner = $2, + lease_expires_at = $3, + started_at = COALESCE(started_at, $4), + version = version + 1 +FROM candidate +WHERE tasks.id = candidate.cid +RETURNING ` + taskColumnList + +// ClaimNext atomically leases the next eligible task. +func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domain.Task, error) { + workloads := f.Workloads + if workloads == nil { + workloads = []string{} // NULL would make the cardinality() guard fail + } + + var task *domain.Task + err := withRetry(ctx, func(ctx context.Context) error { + row := conn(ctx, r.pool).QueryRow(ctx, claimNextSQL, workloads, f.Owner, f.LeaseUntil, f.Now) + t, err := scanTask(row) + if errors.Is(err, pgx.ErrNoRows) { + task = nil + return nil // an empty queue is a normal state, not a failure + } + if err != nil { + return err + } + task = t + return nil + }) + if err != nil { + return nil, err + } + return task, nil +} + +// Get reads a task without locking its row. +func (r *TaskRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Task, error) { + sql, args, err := psql.Select(taskColumns...). + From("tasks"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + t, err := scanTask(conn(ctx, r.pool).QueryRow(ctx, sql, args...)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrTaskNotFound + } + if err != nil { + return nil, err + } + return t, nil +} + +// GetForUpdate reads a task and holds its row lock until the caller's +// transaction ends, so read-modify-write use cases cannot interleave. +func (r *TaskRepo) GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task, error) { + sql, args, err := psql.Select(taskColumns...). + From("tasks"). + Where(sq.Eq{"id": id}). + Suffix("FOR UPDATE"). + ToSql() + if err != nil { + return nil, err + } + t, err := scanTask(conn(ctx, r.pool).QueryRow(ctx, sql, args...)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrTaskNotFound + } + if err != nil { + return nil, err + } + return t, nil +} + +// Update writes the mutated entity back under optimistic concurrency. The entity +// has already incremented its Version in memory, so the new value goes into SET +// while the WHERE guard matches against the previous one (Version-1). +func (r *TaskRepo) Update(ctx context.Context, t *domain.Task) error { + sql, args, err := psql.Update("tasks"). + SetMap(map[string]any{ + "status": string(t.Status), + "attempt": t.Attempt, + "lease_owner": t.LeaseOwner, + "lease_expires_at": t.LeaseExpiresAt, + "result_artifact_id": t.ResultArtifactID, + "metrics": t.Metrics, + "error_code": t.ErrorCode, + "error_message": t.ErrorMessage, + "started_at": t.StartedAt, + "completed_at": t.CompletedAt, + "version": t.Version, + }). + Where(sq.Eq{"id": t.ID, "version": t.Version - 1}). + ToSql() + if err != nil { + return err + } + + tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + // Either the row vanished or someone else advanced its version while we + // held a stale copy. Both mean this write must not land. + return domain.ErrLeaseConflict + } + return nil +} + +// InsertBatch writes every task in one round trip. It runs inside the caller's +// transaction, which is what makes "all tasks or none" hold. +func (r *TaskRepo) InsertBatch(ctx context.Context, tasks []*domain.Task) error { + if len(tasks) == 0 { + return nil + } + + batch := &pgx.Batch{} + for _, t := range tasks { + sql, args, err := psql.Insert("tasks"). + Columns("id", "job_id", "chunk_index", "workload", "input_uri", "input_artifact_id", + "input_sha256", "parameters", "status", "attempt", "max_attempts", "created_at", "version"). + // input_uri is stored NULL (not "") when empty, so the ck_tasks_has_input + // check actually bites: a task with neither a URI nor an artifact fails. + Values(t.ID, t.JobID, t.ChunkIndex, t.Workload, nullIfEmpty(t.InputURI), t.InputArtifactID, + t.InputSHA256, jsonbOrEmpty(t.Parameters), string(t.Status), t.Attempt, t.MaxAttempts, t.CreatedAt, t.Version). + ToSql() + if err != nil { + return err + } + batch.Queue(sql, args...) + } + + results := conn(ctx, r.pool).SendBatch(ctx, batch) + for range tasks { + if _, err := results.Exec(); err != nil { + _ = results.Close() + return err + } + } + return results.Close() +} + +// ListCompleted returns results in chunk order, which the stitcher relies on: +// a non-deterministic order would make the merged output depend on which worker +// happened to finish first. +func (r *TaskRepo) ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domain.Task, error) { + sql, args, err := psql.Select(taskColumns...). + From("tasks"). + Where(sq.Eq{"job_id": jobID, "status": "completed"}). + OrderBy("chunk_index"). + ToSql() + if err != nil { + return nil, err + } + + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + var tasks []*domain.Task + for rows.Next() { + t, err := scanTask(rows) + if err != nil { + return nil, err + } + tasks = append(tasks, t) + } + return tasks, rows.Err() +} + +func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) { + sql, args, err := psql.Select("status", "count(*)"). + From("tasks"). + Where(sq.Eq{"job_id": jobID}). + GroupBy("status"). + ToSql() + if err != nil { + return nil, err + } + + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + counts := make(map[domain.TaskStatus]int) + for rows.Next() { + var ( + status string + n int + ) + if err := rows.Scan(&status, &n); err != nil { + return nil, err + } + counts[domain.TaskStatus(status)] = n + } + return counts, rows.Err() +} + +// expireLeasesSQL applies the lease-expiry rule set-based, mirroring +// domain.Task.ExpireLease: requeue while attempts remain, otherwise fail. +// +// Left as raw SQL: the branching lives in CASE expressions inside the SET, which +// a builder cannot express more clearly than this. It is one statement rather +// than a load-decide-save loop because several coordinators run it concurrently; +// an atomic UPDATE makes the duplicate work harmless — the loser updates zero rows. +var expireLeasesSQL = ` +UPDATE tasks +SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_status + ELSE 'failed'::task_status END, + lease_owner = NULL, + lease_expires_at = NULL, + error_code = CASE WHEN attempt >= max_attempts THEN $2 ELSE error_code END, + error_message = CASE WHEN attempt >= max_attempts + THEN 'lease expired after the final attempt' + ELSE error_message END, + completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END, + version = version + 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 + err := withRetry(ctx, func(ctx context.Context) error { + tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired) + if err != nil { + return err + } + affected = tag.RowsAffected() + return nil + }) + return affected, err +} diff --git a/coordinator/internal/storage/postgres/tx.go b/coordinator/internal/storage/postgres/tx.go new file mode 100644 index 0000000..38adaa7 --- /dev/null +++ b/coordinator/internal/storage/postgres/tx.go @@ -0,0 +1,102 @@ +// Package postgres implements the usecase repository ports on PostgreSQL. +// SQL and pgx types never escape this package. +package postgres + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" +) + +// querier is satisfied by both *pgxpool.Pool and pgx.Tx, letting every +// repository method run identically inside or outside a transaction. +type querier interface { + Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) + QueryRow(ctx context.Context, sql string, args ...any) pgx.Row + Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) + SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults +} + +// txKey is an unexported struct type, so no other package can collide with it +// or reach the transaction we stash in the context. +type txKey struct{} + +// TxManager implements usecase.TxManager. +type TxManager struct { + pool *pgxpool.Pool +} + +func NewTxManager(pool *pgxpool.Pool) *TxManager { + return &TxManager{pool: pool} +} + +// WithinTx runs fn inside one transaction, committing on success and rolling +// back on any error or panic. +// +// The transaction travels in the context rather than in fn's signature, which +// is what lets the usecase layer express "do these repository calls atomically" +// without its port ever mentioning pgx. +// Retrying happens here, around the whole transaction, and deliberately not +// inside the repositories. Once Postgres aborts a transaction with a +// serialization failure or deadlock, every further statement in it fails too — +// replaying a single query would accomplish nothing. The unit of retry is +// Begin → fn → Commit. +// +// This is safe because fn re-reads its rows (via GetForUpdate) on each attempt, +// so a retry starts from the current state rather than stale entities. +func (m *TxManager) WithinTx(ctx context.Context, fn func(ctx context.Context) error) error { + if _, ok := ctx.Value(txKey{}).(pgx.Tx); ok { + // Already inside a transaction — join it. Retrying here would be wrong + // twice over: the outer transaction owns the retry, and re-running fn + // alone cannot undo what the outer one already wrote. + return fn(ctx) + } + + return withRetry(ctx, func(ctx context.Context) error { + return m.runTx(ctx, fn) + }) +} + +func (m *TxManager) runTx(ctx context.Context, fn func(ctx context.Context) error) error { + tx, err := m.pool.Begin(ctx) + if err != nil { + return err + } + // Rollback after a successful Commit is a no-op, so this defer is safe and + // also covers the panic path. + defer func() { _ = tx.Rollback(ctx) }() + + if err := fn(context.WithValue(ctx, txKey{}, tx)); err != nil { + return err + } + return tx.Commit(ctx) +} + +// jsonbOrEmpty keeps a nil map from reaching a NOT NULL jsonb column. pgx +// encodes a nil map as SQL NULL rather than omitting the column, so the +// DEFAULT '{}' never gets a chance to apply. +func jsonbOrEmpty(m map[string]any) map[string]any { + if m == nil { + return map[string]any{} + } + return m +} + +// nullIfEmpty maps "" to a SQL NULL, so an absent optional string is stored as +// NULL rather than an empty string that would defeat a NOT-NULL-or check. +func nullIfEmpty(s string) any { + if s == "" { + return nil + } + return s +} + +// conn returns the transaction bound to ctx, or the pool when there is none. +func conn(ctx context.Context, pool *pgxpool.Pool) querier { + if tx, ok := ctx.Value(txKey{}).(pgx.Tx); ok { + return tx + } + return pool +} diff --git a/coordinator/internal/storage/postgres/worker_repo.go b/coordinator/internal/storage/postgres/worker_repo.go new file mode 100644 index 0000000..1df545e --- /dev/null +++ b/coordinator/internal/storage/postgres/worker_repo.go @@ -0,0 +1,105 @@ +package postgres + +import ( + "context" + "errors" + "fmt" + "time" + + sq "github.com/Masterminds/squirrel" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// WorkerRepo implements usecase.WorkerRepository. +type WorkerRepo struct { + pool *pgxpool.Pool +} + +func NewWorkerRepo(pool *pgxpool.Pool) *WorkerRepo { + return &WorkerRepo{pool: pool} +} + +var workerColumns = []string{"id", "name", "capabilities", "status", "last_heartbeat_at", "created_at", "updated_at"} + +func (r *WorkerRepo) Insert(ctx context.Context, w *domain.Worker) error { + sql, args, err := psql.Insert("workers"). + Columns(workerColumns...). + // capabilities is a jsonb column; pgx marshals the []string to a JSON array. + Values(w.ID, w.Name, w.Capabilities, string(w.Status), + w.LastHeartbeatAt, w.CreatedAt, w.UpdatedAt). + ToSql() + if err != nil { + return err + } + if _, err := conn(ctx, r.pool).Exec(ctx, sql, args...); err != nil { + return fmt.Errorf("insert worker: %w", err) + } + return nil +} + +func (r *WorkerRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Worker, error) { + sql, args, err := psql.Select(workerColumns...). + From("workers"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + + w, err := scanWorker(conn(ctx, r.pool).QueryRow(ctx, sql, args...)) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrWorkerNotFound + } + if err != nil { + return nil, fmt.Errorf("get worker: %w", 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 + status string + ) + if err := row.Scan(&w.ID, &w.Name, &w.Capabilities, &status, + &w.LastHeartbeatAt, &w.CreatedAt, &w.UpdatedAt); err != nil { + return nil, err + } + w.Status = domain.WorkerStatus(status) + return &w, nil +} diff --git a/coordinator/internal/transport/http/dto.go b/coordinator/internal/transport/http/dto.go new file mode 100644 index 0000000..c652b25 --- /dev/null +++ b/coordinator/internal/transport/http/dto.go @@ -0,0 +1,164 @@ +package http + +import ( + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// Wire formats. Keeping them separate from domain entities means the API +// contract can evolve without reshaping the database, and nothing internal +// (version counters, other workers' errors) leaks by accident. + +type createJobRequest struct { + Workload string `json:"workload"` + InputURI string `json:"input_uri"` + Parameters map[string]any `json:"parameters"` + Chunks []chunkDTO `json:"chunks"` +} + +type chunkDTO struct { + ChunkIndex int `json:"chunk_index"` + Workload string `json:"workload"` + InputURI string `json:"input_uri"` + InputSHA256 string `json:"input_sha256"` + Parameters map[string]any `json:"parameters"` + MaxAttempts int `json:"max_attempts"` +} + +type registerRequest struct { + Name string `json:"name"` + Capabilities []string `json:"capabilities"` + // Accepted per the contract for forward compatibility; not yet persisted. + CPUCount int `json:"cpu_count"` + MemoryMB int `json:"memory_mb"` +} + +type registerResponse struct { + WorkerID uuid.UUID `json:"worker_id"` + HeartbeatIntervalSeconds int `json:"heartbeat_interval_seconds"` +} + +type claimRequest struct { + WorkerID string `json:"worker_id"` + Capabilities []string `json:"capabilities"` + // Accepted per the contract; the coordinator leases one task per call. + MaxConcurrency int `json:"max_concurrency"` +} + +type heartbeatRequest struct { + WorkerID string `json:"worker_id"` + Attempt int `json:"attempt"` +} + +type resultRequest struct { + WorkerID string `json:"worker_id"` + Attempt int `json:"attempt"` + Result resultManifest `json:"result"` + Metrics map[string]any `json:"metrics"` +} + +// resultManifest references the artifact the worker already uploaded. sha256 and +// content_type are accepted for the worker's own cross-checking; the coordinator +// trusts its own stored metadata, not these. +type resultManifest struct { + ArtifactID uuid.UUID `json:"artifact_id"` + SHA256 string `json:"sha256"` + ContentType string `json:"content_type"` +} + +type failureRequest struct { + WorkerID string `json:"worker_id"` + Attempt int `json:"attempt"` + ErrorCode string `json:"error_code"` + ErrorMessage string `json:"error_message"` + Retryable bool `json:"retryable"` +} + +type jobResponse struct { + ID uuid.UUID `json:"id"` + Status string `json:"status"` +} + +type taskResponse struct { + ID uuid.UUID `json:"id"` + JobID uuid.UUID `json:"job_id"` + Status string `json:"status"` +} + +type inputRef struct { + URI string `json:"uri"` + SHA256 string `json:"sha256"` +} + +type claimedTaskResponse struct { + TaskID uuid.UUID `json:"task_id"` + JobID uuid.UUID `json:"job_id"` + ChunkIndex int `json:"chunk_index"` + Workload string `json:"workload"` + Input inputRef `json:"input"` + Parameters map[string]any `json:"parameters"` + Attempt int `json:"attempt"` + LeaseExpiresAt time.Time `json:"lease_expires_at"` +} + +type uploadJobResponse struct { + JobID uuid.UUID `json:"job_id"` + TaskCount int `json:"task_count"` + InputArtifactID uuid.UUID `json:"input_artifact_id"` +} + +type jobProgressResponse struct { + ID uuid.UUID `json:"id"` + Status string `json:"status"` + Total int `json:"total"` + Pending int `json:"pending"` + Leased int `json:"leased"` + Done int `json:"completed"` + Failed int `json:"failed"` +} + +type uploadArtifactResponse struct { + ArtifactID uuid.UUID `json:"artifact_id"` + URI string `json:"uri"` + SHA256 string `json:"sha256"` + SizeBytes int64 `json:"size_bytes"` +} + +type errorResponse struct { + Error string `json:"error"` + RequestID string `json:"request_id,omitempty"` +} + +func toClaimedTaskResponse(c domain.ClaimedTask) claimedTaskResponse { + // A shard's input lives in the coordinator; hand the worker a URL to fetch + // it from. A URI-based task keeps its external URI. + uri := c.InputURI + if c.InputArtifactID != nil { + uri = "/tasks/" + c.TaskID.String() + "/input" + } + return claimedTaskResponse{ + TaskID: c.TaskID, + JobID: c.JobID, + ChunkIndex: c.ChunkIndex, + Workload: c.Workload, + Input: inputRef{URI: uri, SHA256: c.InputSHA256}, + Parameters: c.Parameters, + Attempt: c.Attempt, + LeaseExpiresAt: c.LeaseExpiresAt, + } +} + +func toJobProgressResponse(p domain.JobProgress) jobProgressResponse { + return jobProgressResponse{ + ID: p.Job.ID, + Status: string(p.DeriveStatus()), + Total: p.Total, + Pending: p.Pending, + Leased: p.Leased, + Done: p.Done, + Failed: p.Failed, + } +} diff --git a/coordinator/internal/transport/http/errors.go b/coordinator/internal/transport/http/errors.go new file mode 100644 index 0000000..f808ad7 --- /dev/null +++ b/coordinator/internal/transport/http/errors.go @@ -0,0 +1,68 @@ +package http + +import ( + "encoding/json" + "errors" + "net/http" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +func writeJSON(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = 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(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() + return dec.Decode(dst) +} + +// writeError translates domain errors into status codes. This mapping is the +// only place in the codebase that knows HTTP status codes exist — the inner +// layers speak only in business terms. +func (s *Server) writeError(w http.ResponseWriter, r *http.Request, err error) { + reqID := requestIDFrom(r.Context()) + + status := http.StatusInternalServerError + switch { + case errors.Is(err, domain.ErrInvalidInput): + status = http.StatusBadRequest + case errors.Is(err, domain.ErrJobNotFound), errors.Is(err, domain.ErrTaskNotFound), + errors.Is(err, domain.ErrWorkerNotFound), errors.Is(err, domain.ErrArtifactNotFound): + status = http.StatusNotFound + case errors.Is(err, domain.ErrLeaseConflict), + errors.Is(err, domain.ErrStaleAttempt), + errors.Is(err, domain.ErrResultConflict), + errors.Is(err, domain.ErrTaskNotLeased): + status = http.StatusConflict + case errors.Is(err, usecase.ErrNotImplemented): + status = http.StatusNotImplemented + } + + // 501 says "this endpoint has no implementation yet" — that leaks nothing and + // is far more useful than a generic failure, which sent one debugging session + // hunting a database problem that did not exist. + if status == http.StatusNotImplemented { + writeJSON(w, status, errorResponse{Error: "not implemented", RequestID: reqID}) + return + } + + if status >= 500 { + // Never echo an internal error: it can carry table names, query + // fragments, and values. The request ID is the bridge to the logs. + s.log.Error("request failed", "request_id", reqID, "path", r.URL.Path, "err", err) + writeJSON(w, status, errorResponse{Error: "internal error", RequestID: reqID}) + return + } + writeJSON(w, status, errorResponse{Error: err.Error(), RequestID: reqID}) +} diff --git a/coordinator/internal/transport/http/handlers.go b/coordinator/internal/transport/http/handlers.go new file mode 100644 index 0000000..cdd9385 --- /dev/null +++ b/coordinator/internal/transport/http/handlers.go @@ -0,0 +1,367 @@ +package http + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "strings" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// Every handler follows the same shape: decode, map to a use-case input, +// execute, translate. Anything resembling a rule belongs one layer inward. + +func (s *Server) handleCreateJob(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + var req createJobRequest + if err := decodeJSON(r, &req); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + in := usecase.CreateJobInput{ + Workload: req.Workload, + InputURI: req.InputURI, + Parameters: req.Parameters, + } + for _, c := range req.Chunks { + in.Chunks = append(in.Chunks, usecase.ChunkInput(c)) + } + + job, err := s.uc.CreateJob.Execute(ctx, in) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusCreated, jobResponse{ID: job.ID, Status: string(job.Status)}) +} + +func (s *Server) handleRegister(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + var req registerRequest + if err := decodeJSON(r, &req); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + worker, err := s.uc.RegisterWorker.Execute(ctx, usecase.RegisterWorkerInput{ + Name: req.Name, + Capabilities: req.Capabilities, + }) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusCreated, registerResponse{ + WorkerID: worker.ID, + HeartbeatIntervalSeconds: int(s.heartbeatInterval.Seconds()), + }) +} + +func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + var req claimRequest + if err := decodeJSON(r, &req); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{ + WorkerID: req.WorkerID, + Workloads: req.Capabilities, + }) + if err != nil { + s.writeError(w, r, err) + return + } + if claimed == nil { + w.WriteHeader(http.StatusNoContent) // empty queue, not an error + return + } + writeJSON(w, http.StatusOK, toClaimedTaskResponse(*claimed)) +} + +func (s *Server) handleHeartbeat(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + taskID, ok := s.pathUUID(w, r, "task_id") + if !ok { + return + } + var req heartbeatRequest + if err := decodeJSON(r, &req); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + claimed, err := s.uc.RenewLease.Execute(ctx, usecase.RenewLeaseInput{ + TaskID: taskID, + WorkerID: req.WorkerID, + Attempt: req.Attempt, + }) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusOK, toClaimedTaskResponse(*claimed)) +} + +func (s *Server) handleResult(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + taskID, ok := s.pathUUID(w, r, "task_id") + if !ok { + return + } + var req resultRequest + if err := decodeJSON(r, &req); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + task, err := s.uc.CompleteTask.Execute(ctx, usecase.CompleteTaskInput{ + TaskID: taskID, + WorkerID: req.WorkerID, + Attempt: req.Attempt, + ResultArtifactID: req.Result.ArtifactID, + Metrics: req.Metrics, + }) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusOK, taskResponse{ID: task.ID, JobID: task.JobID, Status: string(task.Status)}) +} + +func (s *Server) handleFailure(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + taskID, ok := s.pathUUID(w, r, "task_id") + if !ok { + return + } + var req failureRequest + if err := decodeJSON(r, &req); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + task, err := s.uc.FailTask.Execute(ctx, usecase.FailTaskInput{ + TaskID: taskID, + WorkerID: req.WorkerID, + Attempt: req.Attempt, + ErrorCode: req.ErrorCode, + ErrorMessage: req.ErrorMessage, + Retryable: req.Retryable, + }) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusOK, taskResponse{ID: task.ID, JobID: task.JobID, Status: string(task.Status)}) +} + +// defaultChunkRows is the shard size used when a request omits chunk_rows. +const defaultChunkRows = 1000 + +// handleUploadDataset accepts a multipart submission — the dataset file plus the +// workload/parameters/chunk_rows fields — and hands the file, streamed, to the +// 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) + return + } + + var ( + workload string + params map[string]any + rows = defaultChunkRows + result usecase.SubmitDatasetResult + gotDataset bool + ) + + for { + part, err := mr.NextPart() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + + switch part.FormName() { + case "workload": + b, _ := io.ReadAll(io.LimitReader(part, 1<<10)) + workload = strings.TrimSpace(string(b)) + case "parameters": + b, _ := io.ReadAll(io.LimitReader(part, 1<<16)) + if len(b) > 0 { + if err := json.Unmarshal(b, ¶ms); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + } + case "chunk_rows": + b, _ := io.ReadAll(io.LimitReader(part, 32)) + if n, err := strconv.Atoi(strings.TrimSpace(string(b))); err == nil { + rows = n + } + case "file", "dataset": + filename := part.FileName() + if filename == "" { + filename = "dataset" + } + result, err = s.uc.SubmitDataset.Execute(r.Context(), usecase.SubmitDatasetInput{ + Workload: workload, + Parameters: params, + RowsPerShard: rows, + Filename: filename, + ContentType: part.Header.Get("Content-Type"), + Body: part, + }) + if err != nil { + s.writeError(w, r, err) + return + } + gotDataset = true + } + _ = part.Close() + } + + if !gotDataset { + s.writeError(w, r, domain.ErrInvalidInput) // no file part + return + } + writeJSON(w, http.StatusCreated, uploadJobResponse{ + JobID: result.JobID, + TaskCount: result.TaskCount, + InputArtifactID: result.InputArtifactID, + }) +} + +// handleGetTaskInput streams a task's input shard back to the worker. +func (s *Server) handleGetTaskInput(w http.ResponseWriter, r *http.Request) { + taskID, ok := s.pathUUID(w, r, "task_id") + if !ok { + return + } + art, body, err := s.uc.GetTaskInput.Execute(r.Context(), taskID) + if err != nil { + s.writeError(w, r, err) + return + } + defer func() { _ = body.Close() }() + + w.Header().Set("Content-Type", art.ContentType) + w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10)) + w.Header().Set("X-Checksum-SHA256", art.SHA256) + _, _ = io.Copy(w, body) +} + +// handleUploadArtifact streams a worker's partial result into blob storage. It +// deliberately does not use the short request timeout — a large shard upload +// would trip it — and reads identity from headers per the contract (§5.5). +func (s *Server) handleUploadArtifact(w http.ResponseWriter, r *http.Request) { + taskID, ok := s.pathUUID(w, r, "task_id") + if !ok { + return + } + attempt, err := strconv.Atoi(r.Header.Get("X-Task-Attempt")) + if err != nil { + 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, + WorkerID: r.Header.Get("X-Worker-ID"), + Attempt: attempt, + Filename: r.PathValue("filename"), + ContentType: r.Header.Get("Content-Type"), + Body: r.Body, + }) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusOK, uploadArtifactResponse{ + ArtifactID: art.ID, + URI: "/artifacts/" + art.ID.String() + "/download", + SHA256: art.SHA256, + SizeBytes: art.SizeBytes, + }) +} + +// handleDownloadArtifact streams an artifact's bytes back to the caller. +func (s *Server) handleDownloadArtifact(w http.ResponseWriter, r *http.Request) { + artifactID, ok := s.pathUUID(w, r, "artifact_id") + if !ok { + return + } + art, body, err := s.uc.DownloadArtifact.Execute(r.Context(), artifactID) + if err != nil { + s.writeError(w, r, err) + return + } + defer func() { _ = body.Close() }() + + w.Header().Set("Content-Type", art.ContentType) + w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10)) + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", art.Filename)) + w.Header().Set("X-Checksum-SHA256", art.SHA256) + _, _ = io.Copy(w, body) +} + +func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + + jobID, ok := s.pathUUID(w, r, "job_id") + if !ok { + return + } + progress, err := s.uc.GetJobStatus.Execute(ctx, jobID) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusOK, toJobProgressResponse(progress)) +} + +// --- helpers --- + +func (s *Server) reqCtx(r *http.Request) (context.Context, context.CancelFunc) { + return context.WithTimeout(r.Context(), s.requestTimeout) +} + +func (s *Server) pathUUID(w http.ResponseWriter, r *http.Request, name string) (uuid.UUID, bool) { + id, err := uuid.Parse(r.PathValue(name)) + if err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return uuid.Nil, false + } + return id, true +} diff --git a/coordinator/internal/transport/http/middleware.go b/coordinator/internal/transport/http/middleware.go new file mode 100644 index 0000000..feb3f4f --- /dev/null +++ b/coordinator/internal/transport/http/middleware.go @@ -0,0 +1,103 @@ +package http + +import ( + "context" + "crypto/rand" + "crypto/subtle" + "encoding/hex" + "log/slog" + "net/http" + "strings" + "time" +) + +type ctxKey string + +const requestIDKey ctxKey = "request_id" + +// withRequestID stamps every request with an ID for correlated logs and error +// bodies. It wraps the auth middleware rather than the other way round, so even +// a rejected request carries an ID the caller can quote in a bug report. +func withRequestID(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id := newRequestID() + w.Header().Set("X-Request-ID", id) + next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), requestIDKey, id))) + }) +} + +func requestIDFrom(ctx context.Context) string { + if v, ok := ctx.Value(requestIDKey).(string); ok { + return v + } + return "" +} + +func newRequestID() string { + var b [8]byte + _, _ = rand.Read(b[:]) + return hex.EncodeToString(b[:]) +} + +// withAuth enforces the shared bearer token every worker presents. +// An empty token disables the check (local development only). +func withAuth(token string) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if token == "" { + next.ServeHTTP(w, r) + return + } + presented := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + // Constant-time compare: a byte-by-byte early exit would let an + // attacker recover the token by timing responses. + if subtle.ConstantTimeCompare([]byte(presented), []byte(token)) != 1 { + w.Header().Set("WWW-Authenticate", "Bearer") + writeJSON(w, http.StatusUnauthorized, errorResponse{ + Error: "unauthorized", + RequestID: requestIDFrom(r.Context()), + }) + return + } + next.ServeHTTP(w, r) + }) + } +} + +// statusRecorder captures the status code for the access log. +type statusRecorder struct { + http.ResponseWriter + status int +} + +func (s *statusRecorder) WriteHeader(code int) { + s.status = code + s.ResponseWriter.WriteHeader(code) +} + +// withAccessLog records one structured line per request — the minimum needed to +// debug a distributed system after the fact. +func withAccessLog(log *slog.Logger) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + start := time.Now() + rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK} + next.ServeHTTP(rec, r) + log.Info("request", + "request_id", requestIDFrom(r.Context()), + "method", r.Method, + "path", r.URL.Path, + "status", rec.status, + "duration_ms", time.Since(start).Milliseconds(), + ) + }) + } +} + +// chain applies middleware so that the first argument is the outermost layer. +func chain(h http.Handler, mw ...func(http.Handler) http.Handler) http.Handler { + for i := len(mw) - 1; i >= 0; i-- { + h = mw[i](h) + } + return h +} diff --git a/coordinator/internal/transport/http/server.go b/coordinator/internal/transport/http/server.go new file mode 100644 index 0000000..cf7c304 --- /dev/null +++ b/coordinator/internal/transport/http/server.go @@ -0,0 +1,93 @@ +// Package http adapts the use-case layer to HTTP. Handlers decode requests, +// map them onto use-case inputs, and translate results and errors back — no +// business rules live here. +package http + +import ( + "context" + "log/slog" + "net/http" + "time" + + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// UseCases collects everything the transport needs. Depending on concrete +// use-case types (not one fat interface) keeps each handler's dependency +// explicit and the wiring visible in the composition root. +type UseCases struct { + RegisterWorker *usecase.RegisterWorker + CreateJob *usecase.CreateJob + SubmitDataset *usecase.SubmitDataset + ClaimTask *usecase.ClaimTask + RenewLease *usecase.RenewLease + CompleteTask *usecase.CompleteTask + FailTask *usecase.FailTask + GetJobStatus *usecase.GetJobStatus + UploadArtifact *usecase.UploadArtifact + DownloadArtifact *usecase.DownloadArtifact + GetTaskInput *usecase.GetTaskInput +} + +type Server struct { + uc UseCases + 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, + maxUploadBytes int64, ready func(context.Context) error) *Server { + return &Server{ + uc: uc, + log: log, + requestTimeout: requestTimeout, + heartbeatInterval: heartbeatInterval, + maxUploadBytes: maxUploadBytes, + ready: ready, + } +} + +// Handler builds the router. Go 1.22's ServeMux matches on method and path +// wildcards, so no third-party router is needed. +func (s *Server) Handler(token string) http.Handler { + protected := http.NewServeMux() + protected.HandleFunc("POST /workers/register", s.handleRegister) + protected.HandleFunc("POST /jobs", s.handleCreateJob) + protected.HandleFunc("POST /jobs/upload", s.handleUploadDataset) + protected.HandleFunc("GET /jobs/{job_id}", s.handleGetJob) + protected.HandleFunc("POST /tasks/claim", s.handleClaim) + protected.HandleFunc("GET /tasks/{task_id}/input", s.handleGetTaskInput) + protected.HandleFunc("POST /tasks/{task_id}/heartbeat", s.handleHeartbeat) + protected.HandleFunc("POST /tasks/{task_id}/result", s.handleResult) + protected.HandleFunc("POST /tasks/{task_id}/failure", s.handleFailure) + protected.HandleFunc("PUT /tasks/{task_id}/artifacts/{filename}", s.handleUploadArtifact) + protected.HandleFunc("GET /artifacts/{artifact_id}/download", s.handleDownloadArtifact) + + mux := http.NewServeMux() + mux.HandleFunc("GET /health", s.handleHealth) + mux.Handle("/", chain(protected, + withRequestID, // outermost: every response gets an ID, + withAccessLog(s.log), // including the 401s below + withAuth(token), + )) + return mux +} + +// handleHealth reports readiness. It probes the database so an orchestrator +// learns the difference between "process is up" and "process can serve". +func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) { + if s.ready != nil { + ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second) + defer cancel() + if err := s.ready(ctx); err != nil { + writeJSON(w, http.StatusServiceUnavailable, map[string]string{"status": "unavailable"}) + return + } + } + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) +} diff --git a/coordinator/internal/transport/http/server_test.go b/coordinator/internal/transport/http/server_test.go new file mode 100644 index 0000000..0f00024 --- /dev/null +++ b/coordinator/internal/transport/http/server_test.go @@ -0,0 +1,310 @@ +package http_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "mime/multipart" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/emil28092005/SciMesh/coordinator/internal/memstore" + coordhttp "github.com/emil28092005/SciMesh/coordinator/internal/transport/http" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +const token = "secret" + +type env struct { + ts *httptest.Server + blobs *memstore.BlobStore +} + +func newEnv(t *testing.T, ready func(context.Context) error) *env { + t.Helper() + tasks := memstore.NewTaskRepo() + jobs := memstore.NewJobRepo() + work := memstore.NewWorkerRepo() + arts := memstore.NewArtifactRepo() + blobs := memstore.NewBlobStore() + clk := memstore.NewClock(time.Date(2026, 7, 21, 12, 0, 0, 0, time.UTC)) + tx := memstore.Tx{} + lease := 2 * time.Minute + + uc := coordhttp.UseCases{ + RegisterWorker: usecase.NewRegisterWorker(work, clk), + 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, work, tx, clk, lease), + CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk), + FailTask: usecase.NewFailTask(tasks, jobs, tx, clk), + GetJobStatus: usecase.NewGetJobStatus(jobs, tasks), + UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk), + 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, 1<<30, ready) + ts := httptest.NewServer(srv.Handler(token)) + t.Cleanup(ts.Close) + return &env{ts: ts, blobs: blobs} +} + +func healthy(context.Context) error { return nil } + +// do sends an authenticated JSON request and returns status + decoded body. +func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) { + t.Helper() + req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer "+token) + if body != "" { + req.Header.Set("Content-Type", "application/json") + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("%s %s: %v", method, path, err) + } + defer resp.Body.Close() + var m map[string]any + b, _ := io.ReadAll(resp.Body) + _ = json.Unmarshal(b, &m) + return resp.StatusCode, m +} + +// get issues an unauthenticated GET and returns the response, failing on error. +func (e *env) get(t *testing.T, path string) *http.Response { + t.Helper() + req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+path, nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("GET %s: %v", path, err) + } + return resp +} + +func TestHealthOK(t *testing.T) { + e := newEnv(t, healthy) + resp := e.get(t, "/health") // unauthenticated + defer resp.Body.Close() + if resp.StatusCode != 200 { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +func TestHealthUnavailableWhenDBDown(t *testing.T) { + e := newEnv(t, func(context.Context) error { return context.DeadlineExceeded }) + resp := e.get(t, "/health") + defer resp.Body.Close() + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("status = %d, want 503", resp.StatusCode) + } +} + +func TestAuthRequired(t *testing.T) { + e := newEnv(t, healthy) + send := func(authz string) int { + req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/tasks/claim", + strings.NewReader(`{"worker_id":"w1"}`)) + req.Header.Set("Content-Type", "application/json") + if authz != "" { + req.Header.Set("Authorization", authz) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("claim: %v", err) + } + defer resp.Body.Close() + return resp.StatusCode + } + if code := send(""); code != 401 { + t.Errorf("no token: status = %d, want 401", code) + } + if code := send("Bearer nope"); code != 401 { + t.Errorf("wrong token: status = %d, want 401", code) + } +} + +func TestRegisterWorker(t *testing.T) { + e := newEnv(t, healthy) + code, body := e.do(t, "POST", "/workers/register", `{"name":"lab","capabilities":["w"]}`) + if code != 201 { + t.Fatalf("status = %d, want 201", code) + } + if body["worker_id"] == nil || body["heartbeat_interval_seconds"] == nil { + t.Errorf("missing fields in %v", body) + } +} + +func TestRegisterRejectsNoCapabilities(t *testing.T) { + e := newEnv(t, healthy) + if code, _ := e.do(t, "POST", "/workers/register", `{"name":"lab"}`); code != 400 { + t.Errorf("status = %d, want 400", code) + } +} + +func TestFullLifecycle(t *testing.T) { + e := newEnv(t, healthy) + + // Create a one-chunk job. + code, job := e.do(t, "POST", "/jobs", `{ + "workload":"w","input_uri":"s3://in", + "chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"}]}`) + if code != 201 { + t.Fatalf("create job: %d", code) + } + jobID := job["id"].(string) + + // Claim it. + code, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`) + if code != 200 { + t.Fatalf("claim: %d", code) + } + taskID := claim["task_id"].(string) + attempt := int(claim["attempt"].(float64)) + + // Heartbeat. + if code, _ := e.do(t, "POST", "/tasks/"+taskID+"/heartbeat", + `{"worker_id":"w1","attempt":`+itoa(attempt)+`}`); code != 200 { + t.Fatalf("heartbeat: %d", code) + } + + // Upload a result artifact (PUT, headers carry identity). + artID := e.putArtifact(t, taskID, "w1", attempt, "q,m\nA,B\n") + + // Submit the result by artifact id. + if code, _ := e.do(t, "POST", "/tasks/"+taskID+"/result", + `{"worker_id":"w1","attempt":`+itoa(attempt)+`,"result":{"artifact_id":"`+artID+`"}}`); code != 200 { + t.Fatalf("result: %d", code) + } + + // Job is now completed. + code, prog := e.do(t, "GET", "/jobs/"+jobID, "") + if code != 200 || prog["status"] != "completed" { + t.Errorf("job status = %v (code %d), want completed", prog["status"], code) + } +} + +func TestForeignArtifactResultConflict(t *testing.T) { + e := newEnv(t, healthy) + e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in", + "chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"}, + {"chunk_index":1,"input_uri":"s3://c1","input_sha256":"sha"}]}`) + + _, cA := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`) + _, cB := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`) + taskA, attA := cA["task_id"].(string), int(cA["attempt"].(float64)) + taskB, attB := cB["task_id"].(string), int(cB["attempt"].(float64)) + artA := e.putArtifact(t, taskA, "w1", attA, "data") + + // Complete taskB with taskA's artifact → 409. + if code, _ := e.do(t, "POST", "/tasks/"+taskB+"/result", + `{"worker_id":"w1","attempt":`+itoa(attB)+`,"result":{"artifact_id":"`+artA+`"}}`); code != 409 { + t.Errorf("cross-task result: status = %d, want 409", code) + } +} + +func TestUploadDatasetChunksAndServesInput(t *testing.T) { + e := newEnv(t, healthy) + tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" + + code, body := e.uploadDataset(t, "w", 2, tsv) + if code != 201 { + t.Fatalf("upload: status = %d", code) + } + if int(body["task_count"].(float64)) != 3 { + t.Fatalf("task_count = %v, want 3", body["task_count"]) + } + + // Claim a shard, follow its input.uri, and pull the shard bytes. + _, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`) + input := claim["input"].(map[string]any) + uri := input["uri"].(string) + if !strings.HasPrefix(uri, "/tasks/") || !strings.HasSuffix(uri, "/input") { + t.Fatalf("input.uri = %q", uri) + } + req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+uri, nil) + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("get input: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + t.Fatalf("get input: status = %d", resp.StatusCode) + } + shard, _ := io.ReadAll(resp.Body) + if !strings.HasPrefix(string(shard), "id\tsmiles\n") { + t.Errorf("shard missing header: %q", shard) + } +} + +func TestErrorMappings(t *testing.T) { + e := newEnv(t, healthy) + zero := "00000000-0000-0000-0000-000000000000" + + if code, _ := e.do(t, "GET", "/jobs/"+zero, ""); code != 404 { + t.Errorf("unknown job: %d, want 404", code) + } + if code, _ := e.do(t, "POST", "/tasks/not-a-uuid/heartbeat", `{"worker_id":"w1","attempt":1}`); code != 400 { + t.Errorf("malformed uuid: %d, want 400", code) + } + if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","totally_unknown":1}`); code != 400 { + t.Errorf("unknown field: %d, want 400", code) + } +} + +// --- helpers ------------------------------------------------------------- + +func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data string) string { + t.Helper() + req, _ := http.NewRequestWithContext(context.Background(), "PUT", + e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "text/csv") + req.Header.Set("X-Worker-ID", worker) + req.Header.Set("X-Task-Attempt", itoa(attempt)) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + t.Fatalf("put artifact: status = %d", resp.StatusCode) + } + var m map[string]any + b, _ := io.ReadAll(resp.Body) + _ = json.Unmarshal(b, &m) + return m["artifact_id"].(string) +} + +func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string) (int, map[string]any) { + t.Helper() + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + _ = mw.WriteField("workload", workload) + _ = mw.WriteField("chunk_rows", itoa(rows)) + fw, _ := mw.CreateFormFile("file", "chembl.tsv") + _, _ = io.Copy(fw, strings.NewReader(tsv)) + _ = mw.Close() + + req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", mw.FormDataContentType()) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + var m map[string]any + b, _ := io.ReadAll(resp.Body) + _ = json.Unmarshal(b, &m) + return resp.StatusCode, m +} + +func itoa(n int) string { return strconv.Itoa(n) } diff --git a/coordinator/internal/usecase/artifact.go b/coordinator/internal/usecase/artifact.go new file mode 100644 index 0000000..a515d09 --- /dev/null +++ b/coordinator/internal/usecase/artifact.go @@ -0,0 +1,80 @@ +package usecase + +import ( + "context" + "io" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// UploadArtifact stores a worker's partial-result bytes and records the metadata. +type UploadArtifact struct { + tasks TaskRepository + artifacts ArtifactRepository + blobs BlobStore + clk Clock +} + +func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository, + blobs BlobStore, clk Clock) *UploadArtifact { + return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk} +} + +func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) { + task, err := uc.tasks.Get(ctx, in.TaskID) + if err != nil { + return nil, err + } + // Only the worker holding the current lease at this attempt may upload the + // task's output — the coordinator never trusts an ownership claim on faith. + if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt) { + return nil, domain.ErrLeaseConflict + } + + taskID := task.ID + art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult, + in.Filename, in.ContentType, uc.clk.Now()) + if err != nil { + return nil, err + } + + // Stream to storage first: size and checksum are measured here, by us, not + // taken from the worker. A large shard never sits in memory. + sum, size, err := uc.blobs.Put(ctx, art.StorageKey, in.Body) + if err != nil { + return nil, err + } + art.SetContent(sum, size) + + // Persist the record. If that fails the blob would be an orphan, so remove it. + if err := uc.artifacts.Insert(ctx, art); err != nil { + _ = uc.blobs.Delete(ctx, art.StorageKey) + return nil, err + } + return art, nil +} + +// DownloadArtifact returns an artifact's metadata together with a reader over +// its bytes. The caller must close the reader. +type DownloadArtifact struct { + artifacts ArtifactRepository + blobs BlobStore +} + +func NewDownloadArtifact(artifacts ArtifactRepository, blobs BlobStore) *DownloadArtifact { + return &DownloadArtifact{artifacts: artifacts, blobs: blobs} +} + +func (uc *DownloadArtifact) Execute(ctx context.Context, id uuid.UUID) (*domain.Artifact, io.ReadCloser, error) { + a, err := uc.artifacts.Get(ctx, id) + if err != nil { + return nil, nil, err + } + rc, err := uc.blobs.Open(ctx, a.StorageKey) + if err != nil { + return nil, nil, err + } + return a, rc, nil +} diff --git a/coordinator/internal/usecase/dto.go b/coordinator/internal/usecase/dto.go new file mode 100644 index 0000000..6491cbf --- /dev/null +++ b/coordinator/internal/usecase/dto.go @@ -0,0 +1,83 @@ +package usecase + +import ( + "io" + + "github.com/google/uuid" +) + +// Use-case boundary types. Adapters map their wire formats onto these, so the +// HTTP shape can change without touching business code. + +type CreateJobInput struct { + Workload string + InputURI string + Parameters map[string]any + Chunks []ChunkInput +} + +type ChunkInput struct { + ChunkIndex int + Workload string + InputURI string + InputSHA256 string + Parameters map[string]any + MaxAttempts int +} + +type RegisterWorkerInput struct { + Name string + Capabilities []string +} + +type ClaimTaskInput struct { + WorkerID string + Workloads []string +} + +type RenewLeaseInput struct { + TaskID uuid.UUID + WorkerID string + Attempt int +} + +type CompleteTaskInput struct { + TaskID uuid.UUID + WorkerID string + Attempt int + ResultArtifactID uuid.UUID + Metrics map[string]any +} + +type SubmitDatasetInput struct { + Workload string + Parameters map[string]any + RowsPerShard int + Filename string + ContentType string + Body io.Reader +} + +type SubmitDatasetResult struct { + JobID uuid.UUID + TaskCount int + InputArtifactID uuid.UUID +} + +type UploadArtifactInput struct { + TaskID uuid.UUID + WorkerID string + Attempt int + Filename string + ContentType string + Body io.Reader +} + +type FailTaskInput struct { + TaskID uuid.UUID + WorkerID string + Attempt int + ErrorCode string + ErrorMessage string + Retryable bool +} diff --git a/coordinator/internal/usecase/job.go b/coordinator/internal/usecase/job.go new file mode 100644 index 0000000..28e6c6e --- /dev/null +++ b/coordinator/internal/usecase/job.go @@ -0,0 +1,174 @@ +package usecase + +import ( + "context" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// Job operations: the submitter-facing lifecycle of a whole submission. +// +// CreateJob register a job and fan it out into tasks +// GetJobStatus aggregate progress +// ListResults completed manifests, ordered for the stitcher +// StitchJob merge partial results into the final artifact + +// --- CreateJob ----------------------------------------------------------- + +type CreateJob struct { + jobs JobRepository + tasks TaskRepository + tx TxManager + clock Clock +} + +func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CreateJob { + return &CreateJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock} +} + +// Execute builds the job and its tasks, then writes them in one transaction. +// The all-or-none guarantee comes from TxManager: a half-created job would +// leave chunks no worker could ever complete. +func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) { + chunks := make([]domain.ChunkSpec, 0, len(in.Chunks)) + for _, c := range in.Chunks { + chunks = append(chunks, domain.ChunkSpec(c)) + } + + job, tasks, err := domain.NewJobWithTasks(in.Workload, in.InputURI, in.Parameters, chunks, uc.clock.Now()) + if err != nil { + return nil, err + } + + err = uc.tx.WithinTx(ctx, func(ctx context.Context) error { + if err := uc.jobs.Insert(ctx, job); err != nil { + return err + } + return uc.tasks.InsertBatch(ctx, tasks) + }) + if err != nil { + return nil, err + } + return job, nil +} + +// --- GetJobStatus -------------------------------------------------------- + +type GetJobStatus struct { + jobs JobRepository + tasks TaskRepository +} + +func NewGetJobStatus(jobs JobRepository, tasks TaskRepository) *GetJobStatus { + return &GetJobStatus{jobs: jobs, tasks: tasks} +} + +func (uc *GetJobStatus) Execute(ctx context.Context, jobID uuid.UUID) (domain.JobProgress, error) { + job, err := uc.jobs.Get(ctx, jobID) + if err != nil { + return domain.JobProgress{}, err + } + counts, err := uc.tasks.CountByStatus(ctx, jobID) + if err != nil { + return domain.JobProgress{}, err + } + return progressFrom(*job, counts), nil +} + +// --- ListResults --------------------------------------------------------- + +type ListResults struct { + tasks TaskRepository +} + +func NewListResults(tasks TaskRepository) *ListResults { + return &ListResults{tasks: tasks} +} + +// Execute preserves chunk_index order: the stitcher merges these into one +// artifact, and a non-deterministic order would make the final result depend on +// which worker happened to finish first. +func (uc *ListResults) Execute(ctx context.Context, jobID uuid.UUID) ([]domain.ResultManifest, error) { + tasks, err := uc.tasks.ListCompleted(ctx, jobID) + if err != nil { + return nil, err + } + + manifests := make([]domain.ResultManifest, 0, len(tasks)) + for _, t := range tasks { + if t.ResultArtifactID == nil { + continue // a completed task always references its result; skip defensively + } + manifests = append(manifests, domain.ResultManifest{ + TaskID: t.ID, + ChunkIndex: t.ChunkIndex, + ResultArtifactID: *t.ResultArtifactID, + Metrics: t.Metrics, + }) + } + return manifests, nil +} + +// --- StitchJob ----------------------------------------------------------- + +// StitchJob merges every chunk's partial result into the job's final artifact. +// For similarity search that means concatenating each worker's local top-k, +// sorting by similarity, and keeping the global top-k — the distributed result +// must match what a single local run would produce. +type StitchJob struct { + results *ListResults +} + +func NewStitchJob(results *ListResults) *StitchJob { + return &StitchJob{results: results} +} + +// Execute returns the URI of the assembled artifact. +// +// TODO(phase 6): fetch each manifest's CSV, merge, and persist the result. +func (uc *StitchJob) Execute(ctx context.Context, jobID uuid.UUID) (string, error) { + if _, err := uc.results.Execute(ctx, jobID); err != nil { + return "", err + } + return "", ErrNotImplemented +} + +// --- shared helpers ------------------------------------------------------ + +// progressFrom turns a status histogram into the domain's progress view. +func progressFrom(job domain.Job, counts map[domain.TaskStatus]int) domain.JobProgress { + p := domain.JobProgress{ + Job: job, + Pending: counts[domain.TaskPending], + // 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 + } + return p +} + +// syncJobStatus recomputes a job's status from its task counts and persists it. +// Shared by CompleteTask and FailTask so both close a job by the same rule — +// the rule itself lives in domain.JobProgress.DeriveStatus. +func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository, + jobID uuid.UUID, now time.Time) error { + + counts, err := tasks.CountByStatus(ctx, jobID) + if err != nil { + return err + } + status := progressFrom(domain.Job{}, counts).DeriveStatus() + + var completedAt *time.Time + if status == domain.JobCompleted || status == domain.JobFailed { + completedAt = &now + } + return jobs.UpdateStatus(ctx, jobID, status, completedAt) +} diff --git a/coordinator/internal/usecase/ports.go b/coordinator/internal/usecase/ports.go new file mode 100644 index 0000000..a85481b --- /dev/null +++ b/coordinator/internal/usecase/ports.go @@ -0,0 +1,117 @@ +// Package usecase holds the application's business operations. Each use case is +// a small type with its dependencies injected and a single Execute method. +// +// The interfaces below are *ports*: they are declared here, by the consumer, +// and implemented further out in storage/postgres. That is what keeps the +// dependency rule intact — usecase never imports storage or transport. +package usecase + +import ( + "context" + "errors" + "io" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// ClaimFilter narrows which task a worker may be handed. +type ClaimFilter struct { + Workloads []string // workloads this worker can execute + Owner string // worker ID taking the lease + Now time.Time + LeaseUntil time.Time +} + +// TaskRepository persists tasks. +// +// ClaimNext is deliberately coarse: leasing must be a single atomic statement +// (SELECT ... FOR UPDATE SKIP LOCKED + UPDATE), so it cannot be decomposed into +// Get+Update without losing the guarantee that one task goes to one worker. +type TaskRepository interface { + // ClaimNext atomically leases one matching pending task. + // Returns (nil, nil) when nothing is available. + ClaimNext(ctx context.Context, f ClaimFilter) (*domain.Task, error) + + // Get reads a task without locking. Use it for read-only checks (e.g. + // verifying lease ownership before a long upload) where holding a row lock + // across the operation would be wrong. + Get(ctx context.Context, id uuid.UUID) (*domain.Task, error) + + // GetForUpdate reads a task and locks its row for the enclosing + // transaction, so read-modify-write use cases stay serialized. + GetForUpdate(ctx context.Context, id uuid.UUID) (*domain.Task, error) + + // Update persists a mutated task, honouring its Version for optimistic + // concurrency. + Update(ctx context.Context, t *domain.Task) error + + InsertBatch(ctx context.Context, tasks []*domain.Task) error + + // ListCompleted returns completed tasks ordered by chunk_index. + ListCompleted(ctx context.Context, jobID uuid.UUID) ([]*domain.Task, error) + + // CountByStatus aggregates a job's tasks for progress reporting. + CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) + + // ExpireLeases applies the lease-expiry rule to every elapsed task and + // reports how many were affected. + ExpireLeases(ctx context.Context, now time.Time) (int64, error) +} + +// JobRepository persists jobs. +type JobRepository interface { + Insert(ctx context.Context, j *domain.Job) error + Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) + UpdateStatus(ctx context.Context, id uuid.UUID, status domain.JobStatus, completedAt *time.Time) error +} + +// WorkerRepository persists the worker registry. +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; +// this keeps only the record that points at them. +type ArtifactRepository interface { + Insert(ctx context.Context, a *domain.Artifact) error + Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error) +} + +// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams +// in both directions so a large shard never has to sit in memory, and reports +// the checksum and size it measured while writing — the coordinator's own +// numbers, not the client's claim. +type BlobStore interface { + Put(ctx context.Context, key string, r io.Reader) (sha256 string, size int64, err error) + Open(ctx context.Context, key string) (io.ReadCloser, error) + // Delete removes a stored blob. Used to clean up after a metadata insert + // fails, so a committed blob never outlives its (absent) record. + Delete(ctx context.Context, key string) error +} + +// TxManager runs a function inside one database transaction. The transaction +// travels in the context, so repositories pick it up without this port ever +// mentioning pgx. +type TxManager interface { + WithinTx(ctx context.Context, fn func(ctx context.Context) error) error +} + +// Clock supplies the current time. Injecting it keeps lease and expiry rules +// testable without sleeping or freezing the system clock. +type Clock interface { + Now() time.Time +} + +// ErrNotImplemented marks scaffold code with no body yet. Unlike the errors in +// domain, it describes the state of this codebase, not a business rule. +var ErrNotImplemented = errors.New("not implemented") diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go new file mode 100644 index 0000000..0c2e4f0 --- /dev/null +++ b/coordinator/internal/usecase/task.go @@ -0,0 +1,246 @@ +package usecase + +import ( + "context" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// Task operations: the worker-facing lifecycle of a single chunk. +// +// ClaimTask lease the next available task +// RenewLease extend a held lease (heartbeat) +// CompleteTask record a successful result +// FailTask record a failure +// ExpireLeases reclaim leases that elapsed without a heartbeat + +// --- ClaimTask ----------------------------------------------------------- + +type ClaimTask struct { + tasks TaskRepository + clock Clock + leaseDuration time.Duration +} + +func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask { + return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration} +} + +// Execute reclaims elapsed leases first, then hands out one task. +// +// Sweeping before claiming matters: otherwise a task abandoned by a dead worker +// stays invisible until the reaper's next tick, and a waiting worker is told the +// queue is empty while work sits idle. +// +// This use case is thin by design — the atomicity that makes claiming correct +// lives in one SQL statement behind ClaimNext, and splitting it across the layer +// boundary would break it. +func (uc *ClaimTask) Execute(ctx context.Context, in ClaimTaskInput) (*domain.ClaimedTask, error) { + if in.WorkerID == "" { + return nil, domain.ErrInvalidInput + } + now := uc.clock.Now() + + if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil { + return nil, err + } + + task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{ + Workloads: in.Workloads, + Owner: in.WorkerID, + Now: now, + LeaseUntil: now.Add(uc.leaseDuration), + }) + if err != nil { + return nil, err + } + if task == nil { + return nil, nil // empty queue is a normal state, not an error + } + + claimed := task.AsClaimed() + return &claimed, nil +} + +// --- RenewLease ---------------------------------------------------------- + +type RenewLease struct { + tasks TaskRepository + workers WorkerRepository + tx TxManager + clock Clock + leaseDuration time.Duration +} + +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 +// locked: two concurrent heartbeats must not interleave into a lost update. +// Whether the caller may renew at all is decided by the entity, not here. +func (uc *RenewLease) Execute(ctx context.Context, in RenewLeaseInput) (*domain.ClaimedTask, error) { + var claimed domain.ClaimedTask + + err := uc.tx.WithinTx(ctx, func(ctx context.Context) error { + task, err := uc.tasks.GetForUpdate(ctx, in.TaskID) + if err != nil { + return err + } + if err := task.RenewLease(in.WorkerID, in.Attempt, uc.clock.Now().Add(uc.leaseDuration)); err != nil { + return err + } + if err := uc.tasks.Update(ctx, task); err != nil { + return err + } + claimed = task.AsClaimed() + return nil + }) + 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 +} + +// --- CompleteTask -------------------------------------------------------- + +type CompleteTask struct { + tasks TaskRepository + jobs JobRepository + artifacts ArtifactRepository + tx TxManager + clock Clock +} + +func NewCompleteTask(tasks TaskRepository, jobs JobRepository, artifacts ArtifactRepository, + tx TxManager, clock Clock) *CompleteTask { + return &CompleteTask{tasks: tasks, jobs: jobs, artifacts: artifacts, tx: tx, clock: clock} +} + +// Execute applies the result and, when that was the job's last outstanding +// task, closes the job in the same transaction — so a caller who sees a +// completed task never observes its job still marked running. +// +// Lease ownership, staleness, and idempotent replays are all decided by +// Task.CompleteWith; this use case only orchestrates. +func (uc *CompleteTask) Execute(ctx context.Context, in CompleteTaskInput) (*domain.Task, error) { + var out *domain.Task + + err := uc.tx.WithinTx(ctx, func(ctx context.Context) error { + task, err := uc.tasks.GetForUpdate(ctx, in.TaskID) + if err != nil { + return err + } + // Rule 10: never trust a worker-supplied artifact reference. The result + // must be an artifact the coordinator itself stored for *this* task. + if err := uc.verifyResultArtifact(ctx, in.TaskID, in.ResultArtifactID); err != nil { + return err + } + now := uc.clock.Now() + before := task.Version + if err := task.CompleteWith(in.ResultArtifactID, in.Metrics, + in.WorkerID, in.Attempt, now); err != nil { + return err + } + out = task + + // A replay of an already-recorded result leaves the entity untouched. + // Writing anyway would fail the optimistic-concurrency guard (the stored + // version already equals ours) and turn an idempotent call into a 409. + if task.Version == before { + return nil + } + + if err := uc.tasks.Update(ctx, task); err != nil { + return err + } + return syncJobStatus(ctx, uc.jobs, uc.tasks, task.JobID, now) + }) + if err != nil { + return nil, err + } + return out, nil +} + +// verifyResultArtifact enforces that the referenced artifact was stored by the +// coordinator for this exact task. It stops a worker from completing task B with +// an artifact it uploaded for task A, and from naming an id that isn't a result. +func (uc *CompleteTask) verifyResultArtifact(ctx context.Context, taskID, artifactID uuid.UUID) error { + art, err := uc.artifacts.Get(ctx, artifactID) + if err != nil { + return err + } + if art.TaskID == nil || *art.TaskID != taskID || art.Kind != domain.ArtifactPartialResult { + return domain.ErrResultConflict + } + return nil +} + +// --- FailTask ------------------------------------------------------------ + +type FailTask struct { + tasks TaskRepository + jobs JobRepository + tx TxManager + clock Clock +} + +func NewFailTask(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *FailTask { + return &FailTask{tasks: tasks, jobs: jobs, tx: tx, clock: clock} +} + +// Execute delegates the requeue-or-terminate decision to Task.Fail, then keeps +// the parent job's status consistent in the same transaction. +func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task, error) { + var out *domain.Task + + err := uc.tx.WithinTx(ctx, func(ctx context.Context) error { + task, err := uc.tasks.GetForUpdate(ctx, in.TaskID) + if err != nil { + return err + } + now := uc.clock.Now() + if err := task.Fail(in.WorkerID, in.Attempt, in.ErrorCode, in.ErrorMessage, in.Retryable, now); err != nil { + return err + } + if err := uc.tasks.Update(ctx, task); err != nil { + return err + } + out = task + return syncJobStatus(ctx, uc.jobs, uc.tasks, task.JobID, now) + }) + if err != nil { + return nil, err + } + return out, nil +} + +// --- ExpireLeases -------------------------------------------------------- + +type ExpireLeases struct { + tasks TaskRepository + clock Clock +} + +func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases { + return &ExpireLeases{tasks: tasks, clock: clock} +} + +// Execute reports how many tasks were reclaimed. +// +// The sweep is one set-based statement rather than a load-decide-save loop: +// several coordinators run it concurrently, and a single atomic UPDATE makes +// the duplicate work harmless — the loser simply updates 0 rows. +func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) { + return uc.tasks.ExpireLeases(ctx, uc.clock.Now()) +} diff --git a/coordinator/internal/usecase/upload.go b/coordinator/internal/usecase/upload.go new file mode 100644 index 0000000..f04f582 --- /dev/null +++ b/coordinator/internal/usecase/upload.go @@ -0,0 +1,151 @@ +package usecase + +import ( + "context" + "fmt" + "io" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/chunk" + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// SubmitDataset accepts an uploaded dataset, splits it into shard artifacts, and +// creates the job with one task per shard — the coordinator-side counterpart of +// a client submitting pre-chunked URIs. +type SubmitDataset struct { + blobs BlobStore + artifacts ArtifactRepository + jobs JobRepository + tasks TaskRepository + tx TxManager + clk Clock +} + +func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository, + tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset { + return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk} +} + +func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) { + now := uc.clk.Now() + + job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now) + if err != nil { + return SubmitDatasetResult{}, err + } + + // Everything written to blob storage, so a failed transaction can undo it. + var putKeys []string + cleanup := func() { + for _, k := range putKeys { + _ = uc.blobs.Delete(ctx, k) + } + } + + // 1. Stream the upload into the input artifact; we measure size and sha256. + input, err := domain.NewArtifact(job.ID, nil, domain.ArtifactInput, in.Filename, in.ContentType, now) + if err != nil { + return SubmitDatasetResult{}, err + } + sum, size, err := uc.blobs.Put(ctx, input.StorageKey, in.Body) + if err != nil { + return SubmitDatasetResult{}, err + } + putKeys = append(putKeys, input.StorageKey) + input.SetContent(sum, size) + + // 2. Re-open the stored input and split it into shard artifacts + tasks. + shards := []*domain.Artifact{} + tasks := []*domain.Task{} + rc, err := uc.blobs.Open(ctx, input.StorageKey) + if err != nil { + cleanup() + return SubmitDatasetResult{}, err + } + splitErr := chunk.SplitTSV(rc, in.RowsPerShard, func(index int, shard io.Reader) error { + art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard, + fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now) + if err != nil { + return err + } + ssum, ssize, err := uc.blobs.Put(ctx, art.StorageKey, shard) + if err != nil { + return err + } + putKeys = append(putKeys, art.StorageKey) + art.SetContent(ssum, ssize) + + task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, 0, now) + if err != nil { + return err + } + shards = append(shards, art) + tasks = append(tasks, task) + return nil + }) + _ = rc.Close() + if splitErr != nil { + cleanup() + return SubmitDatasetResult{}, splitErr + } + + // 3. Persist job + all artifacts + all tasks atomically. + err = uc.tx.WithinTx(ctx, func(ctx context.Context) error { + if err := uc.jobs.Insert(ctx, job); err != nil { + return err + } + if err := uc.artifacts.Insert(ctx, input); err != nil { + return err + } + for _, a := range shards { + if err := uc.artifacts.Insert(ctx, a); err != nil { + return err + } + } + return uc.tasks.InsertBatch(ctx, tasks) + }) + if err != nil { + cleanup() + return SubmitDatasetResult{}, err + } + + return SubmitDatasetResult{ + JobID: job.ID, + TaskCount: len(tasks), + InputArtifactID: input.ID, + }, nil +} + +// GetTaskInput resolves a task's input shard and opens it for streaming. The +// caller closes the reader. +type GetTaskInput struct { + tasks TaskRepository + artifacts ArtifactRepository + blobs BlobStore +} + +func NewGetTaskInput(tasks TaskRepository, artifacts ArtifactRepository, blobs BlobStore) *GetTaskInput { + return &GetTaskInput{tasks: tasks, artifacts: artifacts, blobs: blobs} +} + +func (uc *GetTaskInput) Execute(ctx context.Context, taskID uuid.UUID) (*domain.Artifact, io.ReadCloser, error) { + task, err := uc.tasks.Get(ctx, taskID) + if err != nil { + return nil, nil, err + } + if task.InputArtifactID == nil { + // A URI-based task keeps its input outside the coordinator. + return nil, nil, domain.ErrArtifactNotFound + } + art, err := uc.artifacts.Get(ctx, *task.InputArtifactID) + if err != nil { + return nil, nil, err + } + rc, err := uc.blobs.Open(ctx, art.StorageKey) + if err != nil { + return nil, nil, err + } + return art, rc, nil +} diff --git a/coordinator/internal/usecase/usecase_test.go b/coordinator/internal/usecase/usecase_test.go new file mode 100644 index 0000000..254a8ef --- /dev/null +++ b/coordinator/internal/usecase/usecase_test.go @@ -0,0 +1,441 @@ +package usecase_test + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/memstore" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +var ctx = context.Background() + +const lease = 2 * time.Minute + +// harness wires every use case to in-memory stores so orchestration can be +// tested without a database. +type harness struct { + tasks *memstore.TaskRepo + jobs *memstore.JobRepo + work *memstore.WorkerRepo + arts *memstore.ArtifactRepo + blobs *memstore.BlobStore + clk *memstore.Clock + + createJob *usecase.CreateJob + submit *usecase.SubmitDataset + claim *usecase.ClaimTask + renew *usecase.RenewLease + complete *usecase.CompleteTask + fail *usecase.FailTask + status *usecase.GetJobStatus + results *usecase.ListResults + register *usecase.RegisterWorker + uploadArt *usecase.UploadArtifact + downloadArt *usecase.DownloadArtifact + getInput *usecase.GetTaskInput + expire *usecase.ExpireLeases +} + +func newHarness() *harness { + h := &harness{ + tasks: memstore.NewTaskRepo(), + jobs: memstore.NewJobRepo(), + work: memstore.NewWorkerRepo(), + arts: memstore.NewArtifactRepo(), + blobs: memstore.NewBlobStore(), + clk: memstore.NewClock(time.Date(2026, 7, 21, 12, 0, 0, 0, time.UTC)), + } + tx := memstore.Tx{} + 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, 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) + h.results = usecase.NewListResults(h.tasks) + h.register = usecase.NewRegisterWorker(h.work, h.clk) + h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, h.clk) + h.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs) + h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs) + h.expire = usecase.NewExpireLeases(h.tasks, h.clk) + return h +} + +// seedJob creates a URI-chunked job with n chunks and returns its id. +func (h *harness) seedJob(t *testing.T, workload string, n int) uuid.UUID { + t.Helper() + in := usecase.CreateJobInput{Workload: workload, InputURI: "s3://in"} + for i := 0; i < n; i++ { + in.Chunks = append(in.Chunks, usecase.ChunkInput{ + ChunkIndex: i, InputURI: fmt.Sprintf("s3://c%d", i), InputSHA256: "sha", + }) + } + job, err := h.createJob.Execute(ctx, in) + if err != nil { + t.Fatalf("seedJob: %v", err) + } + return job.ID +} + +// leaseOne claims a single task for worker and returns its id and attempt. +func (h *harness) leaseOne(t *testing.T, worker, workload string) (uuid.UUID, int) { + t.Helper() + c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: worker, Workloads: []string{workload}}) + if err != nil || c == nil { + t.Fatalf("leaseOne: claim returned (%v, %v)", c, err) + } + return c.TaskID, c.Attempt +} + +// uploadResult stores a partial-result artifact for a leased task. +func (h *harness) uploadResult(t *testing.T, taskID uuid.UUID, worker string, attempt int) uuid.UUID { + t.Helper() + art, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{ + TaskID: taskID, WorkerID: worker, Attempt: attempt, + Filename: "r.csv", ContentType: "text/csv", Body: strings.NewReader("q,m\nA,B\n"), + }) + if err != nil { + t.Fatalf("uploadResult: %v", err) + } + return art.ID +} + +// --- ClaimTask ----------------------------------------------------------- + +func TestClaimLeasesAndAdvancesAttempt(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + + c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}}) + if err != nil || c == nil { + t.Fatalf("claim = (%v, %v)", c, err) + } + if c.Attempt != 1 || c.LeaseOwner != "w1" { + t.Errorf("attempt=%d owner=%q, want 1/w1", c.Attempt, c.LeaseOwner) + } +} + +func TestClaimEmptyQueueReturnsNil(t *testing.T) { + h := newHarness() + c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}}) + if err != nil || c != nil { + t.Errorf("claim on empty queue = (%v, %v), want (nil, nil)", c, err) + } +} + +func TestClaimRequiresWorkerID(t *testing.T) { + h := newHarness() + if _, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{}); !errors.Is(err, domain.ErrInvalidInput) { + t.Errorf("err = %v, want ErrInvalidInput", err) + } +} + +func TestClaimSweepsExpiredLeaseFirst(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + // w1 leases it, then goes silent past the lease. + taskID, _ := h.leaseOne(t, "w1", "w") + h.clk.Advance(lease + time.Minute) + + // w2 claims: the sweep requeues the dead lease, so w2 gets the same task at attempt 2. + c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w2", Workloads: []string{"w"}}) + if err != nil || c == nil { + t.Fatalf("claim = (%v, %v)", c, err) + } + if c.TaskID != taskID || c.Attempt != 2 || c.LeaseOwner != "w2" { + t.Errorf("got task=%v attempt=%d owner=%q", c.TaskID, c.Attempt, c.LeaseOwner) + } +} + +// --- RenewLease ---------------------------------------------------------- + +func TestRenewExtendsForHolder(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + + c, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt}) + if err != nil { + t.Fatalf("renew: %v", err) + } + if !c.LeaseExpiresAt.Equal(h.clk.Now().Add(lease)) { + t.Error("lease not extended to now+lease") + } +} + +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) + taskID, attempt := h.leaseOne(t, "w1", "w") + + _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: taskID, WorkerID: "intruder", Attempt: attempt}) + if !errors.Is(err, domain.ErrLeaseConflict) { + t.Errorf("err = %v, want ErrLeaseConflict", err) + } +} + +// --- CompleteTask -------------------------------------------------------- + +func TestCompleteHappyPathClosesJob(t *testing.T) { + h := newHarness() + jobID := h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + 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: %v", err) + } + + prog, _ := h.status.Execute(ctx, jobID) + if prog.DeriveStatus() != domain.JobCompleted { + t.Errorf("job status = %q, want completed", prog.DeriveStatus()) + } +} + +func TestCompleteRejectsForeignArtifact(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 2) + // Lease two tasks; upload an artifact for taskA, try to complete taskB with it. + taskA, attA := h.leaseOne(t, "w1", "w") + taskB, attB := h.leaseOne(t, "w1", "w") + artA := h.uploadResult(t, taskA, "w1", attA) + + _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{ + TaskID: taskB, WorkerID: "w1", Attempt: attB, ResultArtifactID: artA, + }) + if !errors.Is(err, domain.ErrResultConflict) { + t.Errorf("cross-task artifact: err = %v, want ErrResultConflict", err) + } +} + +func TestCompleteIsIdempotentOnReplay(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + artID := h.uploadResult(t, taskID, "w1", attempt) + in := usecase.CompleteTaskInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt, ResultArtifactID: artID} + + if _, err := h.complete.Execute(ctx, in); err != nil { + t.Fatalf("first complete: %v", err) + } + if _, err := h.complete.Execute(ctx, in); err != nil { + t.Errorf("replay must be idempotent, got %v", err) + } +} + +// --- FailTask ------------------------------------------------------------ + +func TestFailRequeuesWhileAttemptsRemain(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + + task, err := h.fail.Execute(ctx, usecase.FailTaskInput{ + TaskID: taskID, WorkerID: "w1", Attempt: attempt, + ErrorCode: "boom", ErrorMessage: "exploded", Retryable: true, + }) + if err != nil { + t.Fatalf("fail: %v", err) + } + if task.Status != domain.TaskPending { + t.Errorf("status = %q, want pending (requeued)", task.Status) + } + // It should be claimable again. + if c, _ := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w2", Workloads: []string{"w"}}); c == nil { + t.Error("requeued task should be claimable") + } +} + +// --- CreateJob / status -------------------------------------------------- + +func TestCreateJobFansOutIntoTasks(t *testing.T) { + h := newHarness() + jobID := h.seedJob(t, "w", 3) + prog, err := h.status.Execute(ctx, jobID) + if err != nil { + t.Fatalf("status: %v", err) + } + if prog.Total != 3 || prog.Pending != 3 { + t.Errorf("progress total=%d pending=%d, want 3/3", prog.Total, prog.Pending) + } +} + +// --- RegisterWorker ------------------------------------------------------ + +func TestRegisterWorkerPersists(t *testing.T) { + h := newHarness() + w, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{Name: "lab", Capabilities: []string{"w"}}) + if err != nil { + t.Fatalf("register: %v", err) + } + got, err := h.work.Get(ctx, w.ID) + if err != nil || got.Status != domain.WorkerOnline { + t.Errorf("worker not stored online: %v %v", got, err) + } +} + +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) { + t.Errorf("err = %v, want ErrInvalidInput", err) + } +} + +// --- UploadArtifact ------------------------------------------------------ + +func TestUploadArtifactRejectsForeignWorker(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + + _, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{ + TaskID: taskID, WorkerID: "intruder", Attempt: attempt, + Filename: "r.csv", ContentType: "text/csv", Body: strings.NewReader("x"), + }) + if !errors.Is(err, domain.ErrLeaseConflict) { + t.Errorf("err = %v, want ErrLeaseConflict", err) + } +} + +func TestDownloadArtifactRoundTrips(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attempt := h.leaseOne(t, "w1", "w") + artID := h.uploadResult(t, taskID, "w1", attempt) + + art, rc, err := h.downloadArt.Execute(ctx, artID) + if err != nil { + t.Fatalf("download: %v", err) + } + defer rc.Close() + if art.Kind != domain.ArtifactPartialResult { + t.Errorf("kind = %q", art.Kind) + } +} + +// --- SubmitDataset / GetTaskInput --------------------------------------- + +func TestSubmitDatasetChunksAndServesInput(t *testing.T) { + h := newHarness() + tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" + + res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{ + Workload: "w", RowsPerShard: 2, Filename: "chembl.tsv", + ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv), + }) + if err != nil { + t.Fatalf("submit: %v", err) + } + if res.TaskCount != 3 { // 5 rows / 2 + t.Fatalf("task_count = %d, want 3", res.TaskCount) + } + + // The job now has three claimable shard tasks; each serves its own input. + prog, _ := h.status.Execute(ctx, res.JobID) + if prog.Total != 3 { + t.Errorf("job total = %d, want 3", prog.Total) + } + + c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}}) + if err != nil || c == nil { + t.Fatalf("claim shard: %v", err) + } + if c.InputArtifactID == nil { + t.Fatal("shard task must reference an input artifact") + } + art, rc, err := h.getInput.Execute(ctx, c.TaskID) + if err != nil { + t.Fatalf("get input: %v", err) + } + defer rc.Close() + if art.Kind != domain.ArtifactShard { + t.Errorf("input kind = %q, want shard", art.Kind) + } +} + +func TestGetTaskInputMissingForURITask(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) // URI-based task, no coordinator-stored input + taskID, _ := h.leaseOne(t, "w1", "w") + + if _, _, err := h.getInput.Execute(ctx, taskID); !errors.Is(err, domain.ErrArtifactNotFound) { + t.Errorf("err = %v, want ErrArtifactNotFound", err) + } +} + +// --- ExpireLeases -------------------------------------------------------- + +func TestExpireLeasesReclaims(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + h.leaseOne(t, "w1", "w") + h.clk.Advance(lease + time.Minute) + + n, err := h.expire.Execute(ctx) + if err != nil || n != 1 { + t.Errorf("expire = (%d, %v), want (1, nil)", n, err) + } +} diff --git a/coordinator/internal/usecase/worker.go b/coordinator/internal/usecase/worker.go new file mode 100644 index 0000000..c8ccab2 --- /dev/null +++ b/coordinator/internal/usecase/worker.go @@ -0,0 +1,45 @@ +package usecase + +import ( + "context" + "time" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" +) + +// RegisterWorker records a worker in the registry and hands back its identity. +type RegisterWorker struct { + workers WorkerRepository + clk Clock +} + +func NewRegisterWorker(workers WorkerRepository, clk Clock) *RegisterWorker { + return &RegisterWorker{workers: workers, clk: clk} +} + +func (uc *RegisterWorker) Execute(ctx context.Context, in RegisterWorkerInput) (*domain.Worker, error) { + w, err := domain.NewWorker(in.Name, in.Capabilities, uc.clk.Now()) + if err != nil { + return nil, err + } + if err := uc.workers.Insert(ctx, w); err != nil { + return nil, err + } + 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/0001_init.down.sql b/coordinator/migrations/0001_init.down.sql new file mode 100644 index 0000000..78df969 --- /dev/null +++ b/coordinator/migrations/0001_init.down.sql @@ -0,0 +1,8 @@ +BEGIN; + +DROP TABLE IF EXISTS tasks; +DROP TABLE IF EXISTS jobs; +DROP TYPE IF EXISTS task_status; +DROP TYPE IF EXISTS job_status; + +COMMIT; diff --git a/coordinator/migrations/0001_init.up.sql b/coordinator/migrations/0001_init.up.sql new file mode 100644 index 0000000..c84c469 --- /dev/null +++ b/coordinator/migrations/0001_init.up.sql @@ -0,0 +1,58 @@ +BEGIN; + +CREATE TYPE job_status AS ENUM ('pending','running','completed','failed','cancelled'); +CREATE TYPE task_status AS ENUM ('pending','leased','completed','failed','cancelled'); + +-- One user submission, possibly split into several tasks. +CREATE TABLE jobs ( + id uuid PRIMARY KEY, + workload text NOT NULL, + input_uri text NOT NULL, + parameters jsonb NOT NULL DEFAULT '{}'::jsonb, + status job_status NOT NULL DEFAULT 'pending', + created_at timestamptz NOT NULL DEFAULT now(), + completed_at timestamptz +); + +-- One independently executable chunk. +CREATE TABLE tasks ( + id uuid PRIMARY KEY, + job_id uuid NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, + chunk_index integer NOT NULL, + workload text NOT NULL, + input_uri text NOT NULL, + input_sha256 text NOT NULL, + parameters jsonb NOT NULL DEFAULT '{}'::jsonb, + status task_status NOT NULL DEFAULT 'pending', + attempt integer NOT NULL DEFAULT 0, + max_attempts integer NOT NULL DEFAULT 3, + lease_owner text, + lease_expires_at timestamptz, + result_uri text, + result_sha256 text, + metrics jsonb, + error_code text, + error_message text, + created_at timestamptz NOT NULL DEFAULT now(), + started_at timestamptz, + completed_at timestamptz, + version integer NOT NULL DEFAULT 0, + + CONSTRAINT uq_tasks_job_chunk UNIQUE (job_id, chunk_index), + CONSTRAINT ck_tasks_attempt CHECK (attempt >= 0), + CONSTRAINT ck_tasks_max_attempts CHECK (max_attempts > 0), + -- A completed task must carry its result manifest. + CONSTRAINT ck_tasks_completed_result CHECK ( + status <> 'completed' OR (result_uri IS NOT NULL AND result_sha256 IS NOT NULL) + ), + -- A leased task must carry its lease. + CONSTRAINT ck_tasks_leased_owner CHECK ( + status <> 'leased' OR (lease_owner IS NOT NULL AND lease_expires_at IS NOT NULL) + ) +); + +-- Claim path: find the oldest pending task fast. +CREATE INDEX ix_tasks_claim ON tasks (status, lease_expires_at, created_at); +CREATE INDEX ix_tasks_job ON tasks (job_id); + +COMMIT; diff --git a/coordinator/migrations/0002_workers.down.sql b/coordinator/migrations/0002_workers.down.sql new file mode 100644 index 0000000..c01563d --- /dev/null +++ b/coordinator/migrations/0002_workers.down.sql @@ -0,0 +1,6 @@ +BEGIN; + +DROP TABLE IF EXISTS workers; +DROP TYPE IF EXISTS worker_status; + +COMMIT; diff --git a/coordinator/migrations/0002_workers.up.sql b/coordinator/migrations/0002_workers.up.sql new file mode 100644 index 0000000..25547fd --- /dev/null +++ b/coordinator/migrations/0002_workers.up.sql @@ -0,0 +1,20 @@ +BEGIN; + +CREATE TYPE worker_status AS ENUM ('online','busy','offline'); + +-- A registered process/machine that can claim tasks. Registration returns the +-- id; liveness is tracked by last_heartbeat_at. +CREATE TABLE workers ( + id uuid PRIMARY KEY, + name text NOT NULL DEFAULT '', + capabilities jsonb NOT NULL DEFAULT '[]'::jsonb, + status worker_status NOT NULL DEFAULT 'online', + last_heartbeat_at timestamptz NOT NULL DEFAULT now(), + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now() +); + +-- Liveness sweep: find workers that have gone quiet. +CREATE INDEX ix_workers_liveness ON workers (status, last_heartbeat_at); + +COMMIT; diff --git a/coordinator/migrations/0003_artifacts.down.sql b/coordinator/migrations/0003_artifacts.down.sql new file mode 100644 index 0000000..54900e0 --- /dev/null +++ b/coordinator/migrations/0003_artifacts.down.sql @@ -0,0 +1,11 @@ +BEGIN; + +ALTER TABLE tasks DROP COLUMN IF EXISTS input_artifact_id; +ALTER TABLE tasks DROP COLUMN IF EXISTS result_artifact_id; +ALTER TABLE jobs DROP COLUMN IF EXISTS input_artifact_id; +ALTER TABLE jobs DROP COLUMN IF EXISTS result_artifact_id; + +DROP TABLE IF EXISTS artifacts; +DROP TYPE IF EXISTS artifact_kind; + +COMMIT; diff --git a/coordinator/migrations/0003_artifacts.up.sql b/coordinator/migrations/0003_artifacts.up.sql new file mode 100644 index 0000000..103b646 --- /dev/null +++ b/coordinator/migrations/0003_artifacts.up.sql @@ -0,0 +1,31 @@ +BEGIN; + +CREATE TYPE artifact_kind AS ENUM ('input','shard','partial_result','final_result','log'); + +-- A durable file the coordinator owns: input, shard, partial/final result, log. +-- The database is the source of truth; files are found through this metadata, +-- never by scanning directories. +CREATE TABLE artifacts ( + id uuid PRIMARY KEY, + job_id uuid NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, + task_id uuid REFERENCES tasks(id) ON DELETE CASCADE, -- null for job-level inputs + kind artifact_kind NOT NULL, + filename text NOT NULL, + storage_key text NOT NULL UNIQUE, -- coordinator-generated, never a client path + content_type text NOT NULL DEFAULT 'application/octet-stream', + size_bytes bigint NOT NULL CHECK (size_bytes >= 0), + sha256 text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now() +); + +CREATE INDEX ix_artifacts_job ON artifacts (job_id); +CREATE INDEX ix_artifacts_task ON artifacts (task_id); + +-- Jobs and tasks reference their artifacts. Nullable during the transition from +-- URI-based inputs/results to artifact-based ones. +ALTER TABLE jobs ADD COLUMN input_artifact_id uuid REFERENCES artifacts(id); +ALTER TABLE jobs ADD COLUMN result_artifact_id uuid REFERENCES artifacts(id); +ALTER TABLE tasks ADD COLUMN input_artifact_id uuid REFERENCES artifacts(id); +ALTER TABLE tasks ADD COLUMN result_artifact_id uuid REFERENCES artifacts(id); + +COMMIT; diff --git a/coordinator/migrations/0004_result_artifact.down.sql b/coordinator/migrations/0004_result_artifact.down.sql new file mode 100644 index 0000000..acbbdb3 --- /dev/null +++ b/coordinator/migrations/0004_result_artifact.down.sql @@ -0,0 +1,11 @@ +BEGIN; + +ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_completed_result; +ALTER TABLE tasks ADD COLUMN result_uri text; +ALTER TABLE tasks ADD COLUMN result_sha256 text; + +ALTER TABLE tasks ADD CONSTRAINT ck_tasks_completed_result CHECK ( + status <> 'completed' OR (result_uri IS NOT NULL AND result_sha256 IS NOT NULL) +); + +COMMIT; diff --git a/coordinator/migrations/0004_result_artifact.up.sql b/coordinator/migrations/0004_result_artifact.up.sql new file mode 100644 index 0000000..e062e1a --- /dev/null +++ b/coordinator/migrations/0004_result_artifact.up.sql @@ -0,0 +1,14 @@ +BEGIN; + +-- Results are now coordinator-owned artifacts, not worker-supplied URIs. +-- Drop the URI-based completion guard and columns, and require a completed task +-- to reference its result artifact instead (PLAN.md §6.2). +ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_completed_result; +ALTER TABLE tasks DROP COLUMN IF EXISTS result_uri; +ALTER TABLE tasks DROP COLUMN IF EXISTS result_sha256; + +ALTER TABLE tasks ADD CONSTRAINT ck_tasks_completed_result CHECK ( + status <> 'completed' OR result_artifact_id IS NOT NULL +); + +COMMIT; diff --git a/coordinator/migrations/0005_uploaded_input.down.sql b/coordinator/migrations/0005_uploaded_input.down.sql new file mode 100644 index 0000000..87731e2 --- /dev/null +++ b/coordinator/migrations/0005_uploaded_input.down.sql @@ -0,0 +1,9 @@ +BEGIN; + +ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_has_input; + +-- Restoring NOT NULL requires the columns to be populated; safe on a fresh DB. +ALTER TABLE tasks ALTER COLUMN input_uri SET NOT NULL; +ALTER TABLE jobs ALTER COLUMN input_uri SET NOT NULL; + +COMMIT; diff --git a/coordinator/migrations/0005_uploaded_input.up.sql b/coordinator/migrations/0005_uploaded_input.up.sql new file mode 100644 index 0000000..04d79fb --- /dev/null +++ b/coordinator/migrations/0005_uploaded_input.up.sql @@ -0,0 +1,13 @@ +BEGIN; + +-- Inputs can now arrive as uploaded artifacts (POST /jobs/upload), not only as +-- external URIs. Relax the URI requirement and require every task to have an +-- input one way or the other. +ALTER TABLE jobs ALTER COLUMN input_uri DROP NOT NULL; +ALTER TABLE tasks ALTER COLUMN input_uri DROP NOT NULL; + +ALTER TABLE tasks ADD CONSTRAINT ck_tasks_has_input CHECK ( + input_uri IS NOT NULL OR input_artifact_id IS NOT NULL +); + +COMMIT; 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; diff --git a/coordinator/scripts/smoke.sh b/coordinator/scripts/smoke.sh new file mode 100755 index 0000000..2d9a1f2 --- /dev/null +++ b/coordinator/scripts/smoke.sh @@ -0,0 +1,204 @@ +#!/usr/bin/env bash +# +# End-to-end smoke test against a running coordinator. +# +# ./scripts/smoke.sh # localhost:8080, token from .env +# HOST=http://1.2.3.4:8080 TOKEN=x ./scripts/smoke.sh +# +# Exits non-zero on the first unexpected status, so it is usable in CI. + +set -uo pipefail + +HOST="${HOST:-http://localhost:8080}" +TOKEN="${TOKEN:-$(grep -s '^WORKER_AUTH_TOKEN=' .env | cut -d= -f2- || echo change-me)}" + +pass=0 +fail=0 + +# check