Emil_Shanaty/SciMesh
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'feat/coordinator'

Browse Source 
# Conflicts:
#	docs/api-contract.md
This commit is contained in:
Emil
2026-07-23 21:17:53 +03:00
parent abfda35170 8af8ddcf48
commit b4a89dd7c2
75 changed files with 8574 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/coordinator.yml
+66
View File
@@ -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
coordinator/.dockerignore
+14
View File
@@ -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
coordinator/.env.example
+28
View File
@@ -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
coordinator/.gitignore
+6
View File
@@ -0,0 +1,6 @@
/coordinator
/bin/
.env
*.out
/logs/
/data/
coordinator/.golangci.yml
+54
View File
@@ -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
coordinator/ARCHITECTURE.md
+144
View File
@@ -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`, осталось их подключить.
coordinator/Dockerfile
+52
View File
@@ -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"]
coordinator/Makefile
+85
View File
@@ -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
coordinator/README.md
+171
View File
@@ -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.
coordinator/api/requests.http
+234
View File
@@ -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.
coordinator/cmd/coordinator/main.go
+125
View File
@@ -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
}
coordinator/docker-compose.yml
+81
View File
@@ -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:
coordinator/go.mod
+23
View File
@@ -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
)
coordinator/go.sum
+43
View File
@@ -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=
coordinator/internal/chunk/tsv.go
+92
View File
@@ -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
}
coordinator/internal/chunk/tsv_test.go
+117
View File
@@ -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)
}
}
coordinator/internal/domain/artifact.go
+69
View File
@@ -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
}
coordinator/internal/domain/artifact_test.go
+55
View File
@@ -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")
}
}
coordinator/internal/domain/errors.go
+20
View File
@@ -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")
)
coordinator/internal/domain/job.go
+125
View File
@@ -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
}
}
coordinator/internal/domain/job_test.go
+140
View File
@@ -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)
}
})
}
}
coordinator/internal/domain/task.go
+281
View File
@@ -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
}
coordinator/internal/domain/task_test.go
+221
View File
@@ -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)
}
}
coordinator/internal/domain/worker.go
+45
View File
@@ -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
}
coordinator/internal/domain/worker_test.go
+31
View File
@@ -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)
}
}
coordinator/internal/infra/clock.go
+13
View File
@@ -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() }
coordinator/internal/infra/config.go
+184
View File
@@ -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
}
coordinator/internal/infra/db.go
+65
View File
@@ -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)
},
)
}
coordinator/internal/infra/logging.go
+65
View File
@@ -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 }
coordinator/internal/infra/server.go
+74
View File
@@ -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)
}
}
}
}
coordinator/internal/memstore/memstore.go
+344
View File
@@ -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
}
coordinator/internal/storage/blob/store.go
+130
View File
@@ -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)
}
coordinator/internal/storage/blob/store_test.go
+115
View File
@@ -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 }
coordinator/internal/storage/postgres/artifact_repo.go
+72
View File
@@ -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
}
coordinator/internal/storage/postgres/builder.go
+11
View File
@@ -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)
coordinator/internal/storage/postgres/integration_test.go
+516
View File
@@ -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])
}
}
coordinator/internal/storage/postgres/job_repo.go
+91
View File
@@ -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
}
coordinator/internal/storage/postgres/retry.go
+82
View File
@@ -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))
}
coordinator/internal/storage/postgres/retry_test.go
+98
View File
@@ -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)
}
}
coordinator/internal/storage/postgres/task_repo.go
+330
View File
@@ -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
}
coordinator/internal/storage/postgres/tx.go
+102
View File
@@ -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
}
coordinator/internal/storage/postgres/worker_repo.go
+105
View File
@@ -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
}
coordinator/internal/transport/http/dto.go
+164
View File
@@ -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,
}
}
coordinator/internal/transport/http/errors.go
+68
View File
@@ -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})
}
coordinator/internal/transport/http/handlers.go
+367
View File
@@ -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, &params); 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
}
coordinator/internal/transport/http/middleware.go
+103
View File
@@ -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
}
coordinator/internal/transport/http/server.go
+93
View File
@@ -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"})
}
coordinator/internal/transport/http/server_test.go
+310
View File
@@ -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) }
coordinator/internal/usecase/artifact.go
+80
View File
@@ -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
}
coordinator/internal/usecase/dto.go
+83
View File
@@ -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
}
coordinator/internal/usecase/job.go
+174
View File
@@ -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)
}
coordinator/internal/usecase/ports.go
+117
View File
@@ -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")
coordinator/internal/usecase/task.go
+246
View File
@@ -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())
}
coordinator/internal/usecase/upload.go
+151
View File
@@ -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
}
coordinator/internal/usecase/usecase_test.go
+441
View File
@@ -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)
}
}
coordinator/internal/usecase/worker.go
+45
View File
@@ -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))
}
coordinator/migrations/0001_init.down.sql
+8
View File
@@ -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;
coordinator/migrations/0001_init.up.sql
+58
View File
@@ -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;
coordinator/migrations/0002_workers.down.sql
+6
View File
@@ -0,0 +1,6 @@
BEGIN;
DROP TABLE IF EXISTS workers;
DROP TYPE IF EXISTS worker_status;
COMMIT;
coordinator/migrations/0002_workers.up.sql
+20
View File
@@ -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;
coordinator/migrations/0003_artifacts.down.sql
+11
View File
@@ -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;
coordinator/migrations/0003_artifacts.up.sql
+31
View File
@@ -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;
coordinator/migrations/0004_result_artifact.down.sql
+11
View File
@@ -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;
coordinator/migrations/0004_result_artifact.up.sql
+14
View File
@@ -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;
coordinator/migrations/0005_uploaded_input.down.sql
+9
View File
@@ -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;
coordinator/migrations/0005_uploaded_input.up.sql
+13
View File
@@ -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;
coordinator/migrations/0006_task_running_enum.down.sql
+4
View File
@@ -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;
coordinator/migrations/0006_task_running_enum.up.sql
+5
View File
@@ -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';
coordinator/migrations/0007_task_running_lease.down.sql
+8
View File
@@ -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;
coordinator/migrations/0007_task_running_lease.up.sql
+10
View File
@@ -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;
coordinator/scripts/smoke.sh
Executable
+204
View File
@@ -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 <label> <expected-status> <curl args...>
check() {
local label="$1" want="$2"
shift 2
local body status
body=$(curl -sS -w '\n%{http_code}' "$@" 2>&1)
status=$(printf '%s' "$body" | tail -n1)
if [[ "$status" == "$want" ]]; then
printf ' \033[32m✓\033[0m %-46s %s\n' "$label" "$status"
pass=$((pass + 1))
else
printf ' \033[31m✗\033[0m %-46s got %s, want %s\n' "$label" "$status" "$want"
printf ' %s\n' "$(printf '%s' "$body" | head -n-1)"
fail=$((fail + 1))
fi
}
json() { printf '%s' "$1" | head -n-1; }
auth=(-H "Authorization: Bearer ${TOKEN}" -H 'Content-Type: application/json')
echo "coordinator: ${HOST}"
echo
echo "health & auth"
check "GET /health" 200 "${HOST}/health"
check "claim without a token → 401" 401 -X POST "${HOST}/tasks/claim" \
-H 'Content-Type: application/json' -d '{"worker_id":"w1"}'
echo
echo "worker registry"
check "register worker" 201 -X POST "${HOST}/workers/register" "${auth[@]}" \
-d '{"name":"smoke-worker","capabilities":["similarity_search"],"cpu_count":4,"memory_mb":8192}'
check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \
-d '{"name":"bad"}'
echo
echo "job lifecycle"
job=$(curl -sS "${auth[@]}" -X POST "${HOST}/jobs" -d '{
"workload":"similarity_search","input_uri":"s3://chembl","parameters":{"top_k":10},
"chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"aaa"},
{"chunk_index":1,"input_uri":"s3://c1","input_sha256":"bbb"}]}')
job_id=$(printf '%s' "$job" | python3 -c 'import json,sys;print(json.load(sys.stdin)["id"])' 2>/dev/null)
if [[ -z "${job_id:-}" ]]; then
echo " ✗ could not create a job: $job"
exit 1
fi
printf ' \033[32m✓\033[0m %-46s %s\n' "POST /jobs" "$job_id"
pass=$((pass + 1))
# The database may hold pending tasks from earlier runs, so claim until we have
# both of *our* chunks rather than assuming the queue starts empty. The attempt
# number comes from the response too: a task requeued by an expired lease is
# handed out with attempt 2 or 3, and hard-coding 1 would fail the lease check.
declare -A our_chunks
task_id=""
attempt=""
for _ in $(seq 1 40); do
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d '{"worker_id":"w1"}')
[[ -z "$claim" ]] && break # 204: queue drained
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
python3 -c 'import json,sys;d=json.load(sys.stdin);print(d["job_id"],d["task_id"],d["chunk_index"],d["attempt"])' 2>/dev/null)
[[ "$c_job" != "$job_id" ]] && continue # someone else's leftover task
our_chunks["$c_chunk"]=1
if [[ -z "$task_id" ]]; then
task_id="$c_task"
attempt="$c_attempt"
fi
[[ "${#our_chunks[@]}" -eq 2 ]] && break
done
if [[ "${#our_chunks[@]}" -eq 2 ]]; then
printf ' \033[32m✓\033[0m %-46s chunks %s\n' "POST /tasks/claim × 2 (distinct)" "${!our_chunks[*]}"
pass=$((pass + 1))
else
printf ' \033[31m✗\033[0m %-46s got %d distinct chunks, want 2\n' "claim" "${#our_chunks[@]}"
fail=$((fail + 1))
exit 1
fi
check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt}}"
# --- artifacts + result (uploads happen while the task is still leased) ---
bearer=(-H "Authorization: Bearer ${TOKEN}")
# upload <filename> -> prints the artifact_id
upload() {
curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
--data-binary $'query,match,score\nA,B,0.9\n' |
python3 -c 'import json,sys;print(json.load(sys.stdin)["artifact_id"])' 2>/dev/null
}
check "upload artifact" 200 -X PUT "${HOST}/tasks/${task_id}/artifacts/result.csv" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
--data-binary $'query,match,score\nA,B,0.9\n'
check "foreign worker upload → 409" 409 -X PUT "${HOST}/tasks/${task_id}/artifacts/x.csv" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \
--data-binary 'x'
# Two result artifacts, uploaded now while the lease is held: one to complete
# with, a second to prove a different manifest is rejected after completion.
art_id=$(upload primary.csv)
art_id2=$(upload secondary.csv)
check "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}"
check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "replay same result → idempotent" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "different result → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id2}\"}}"
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
echo
echo "input validation"
check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \
-d '{"worker_id":"w1","attempt":1,"result":{"artifact_id":"00000000-0000-0000-0000-000000000000"}}'
# Note: Go's encoding/json matches field names case-insensitively, so
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
# trips DisallowUnknownFields.
check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \
-d '{"worker_id":"w1","totally_unknown":1}'
check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}"
echo
echo "dataset upload → chunking"
# Upload a 5-row TSV split at 2 rows/shard → expect 3 shard tasks. The text
# fields precede the file part, which the coordinator streams.
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
-F 'workload=similarity_search' \
-F 'parameters={"top_k":10}' \
-F 'chunk_rows=2' \
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
id smiles
A CC
B CCC
C CCCC
D CCCCC
E CCCCCC
TSV
)
up_job=$(printf '%s' "$up" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
up_count=$(printf '%s' "$up" | python3 -c 'import json,sys;print(json.load(sys.stdin)["task_count"])' 2>/dev/null)
if [[ "$up_count" == "3" ]]; then
printf ' \033[32m✓\033[0m %-46s task_count=3\n' "POST /jobs/upload (5 rows / 2)"
pass=$((pass + 1))
else
printf ' \033[31m✗\033[0m %-46s got task_count=%s, want 3\n' "POST /jobs/upload" "${up_count:-?}"
printf ' %s\n' "$up"
fail=$((fail + 1))
fi
# Claim one of this job's shard tasks and pull its input shard from the coordinator.
up_input=""
for _ in $(seq 1 30); do
c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \
-d '{"worker_id":"up-w","capabilities":["similarity_search"]}')
[[ -z "$c" ]] && break
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
[[ "$cj" != "$up_job" ]] && continue
up_input=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["input"]["uri"])' 2>/dev/null)
break
done
if [[ "$up_input" == /tasks/*/input ]]; then
printf ' \033[32m✓\033[0m %-46s %s\n' "claim → input.uri points at coordinator" "$up_input"
pass=$((pass + 1))
else
printf ' \033[31m✗\033[0m %-46s got %q\n' "claim shard input.uri" "$up_input"
fail=$((fail + 1))
fi
check "download shard input" 200 "${HOST}${up_input}" "${bearer[@]}"
echo
curl -sS "${HOST}/jobs/${job_id}" "${auth[@]}"
echo
printf '\n%d passed, %d failed\n' "$pass" "$fail"
[[ "$fail" -eq 0 ]]
docs/api-contract.md
+145 -96
View File
@@ -1,86 +1,150 @@
# SciMesh Coordinator API Contract
# SciMesh coordinator ↔ worker API contract (v1)
**Status:** draft, version 1. This document is the compatibility boundary
between the Go coordinator and the Python Worker. Change it only in the same
pull request as both implementation and contract tests.
**Status marker:** `v1`. This document is the single source of truth for the Go
coordinator and the Python Worker Daemon. It is derived from `PLAN.md` §5 and
must be updated in the same change as any behaviour it describes.
## General rules
> **Machine-readable:** [`openapi.yaml`](openapi.yaml) is the OpenAPI 3.0 mirror
> of this document — feed it to `openapi-python-client` or `datamodel-code-generator`
> to generate the Python client/models. This markdown stays the human-readable
> source; keep the two in sync.
- All worker endpoints require `Authorization: Bearer <token>`.
- Times use UTC RFC 3339, for example `2026-07-23T12:05:00Z`.
- JSON requests and responses use `application/json`.
- `worker_id` and `attempt` identify a lease. The coordinator validates them
transactionally on every task mutation.
- A task becomes `completed` only after a coordinator-owned artifact is durable.
- Identical repeated completion is successful; a different result for the same
attempt is a conflict.
- **Auth:** every endpoint except readiness requires `Authorization: Bearer <token>`.
- **Identity:** every mutating worker request carries `worker_id` and `attempt`;
they are checked against the current task lease in PostgreSQL. A stale attempt
gets `409`.
- **Timestamps:** UTC, RFC 3339 (e.g. `2026-07-22T12:05:00Z`).
- **Unknown JSON fields are rejected** with `400`.
## Worker registration
## Implementation status
| Endpoint | Contract | Coordinator |
| --- | --- | --- |
| `GET /health` | readiness incl. DB | ✅ done |
| `POST /workers/register` | register + capabilities | ✅ done |
| `POST /tasks/claim` | atomic lease | ✅ done |
| `POST /tasks/{id}/heartbeat` | renew lease | ✅ done |
| `POST /tasks/{id}/result` | complete | ✅ done, references `artifact_id` |
| `POST /tasks/{id}/failure` | fail | ✅ done |
| `GET /jobs/{id}` | progress | ✅ done |
| `PUT /tasks/{id}/artifacts/{name}` | upload partial | ✅ done |
| `GET /artifacts/{id}/download` | download by id | ✅ done |
| `POST /jobs/upload` | upload dataset, coordinator chunks it | ✅ done |
| `GET /tasks/{id}/input` | download shard | ✅ done |
---
## Readiness
```http
GET /health
```
`200 {"status":"ok"}` when the database is reachable; `503 {"status":"unavailable"}`
otherwise. Unauthenticated.
## Submit a dataset (submitter-side)
```http
POST /jobs/upload
Authorization: Bearer <token>
Content-Type: multipart/form-data
```
Fields, in order (text fields first, file last — the file is streamed):
`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), and the file
part `file`. The coordinator stores the input, splits the TSV into shard
artifacts (header repeated per shard), and creates one task per shard.
`201`:
```json
{ "job_id": "uuid", "task_count": 3, "input_artifact_id": "uuid" }
```
Each resulting task's claim response carries `input.uri = /tasks/{id}/input`,
served by §5.4.
## Register worker
```http
POST /workers/register
Authorization: Bearer <token>
Content-Type: application/json
{"name":"lab-worker-01","capabilities":["similarity-search"],"cpu_count":8,"memory_mb":16384}
```
Returns `200 OK`:
```json
{"worker_id":"uuid","heartbeat_interval_seconds":15}
```
## Task lifecycle
### Claim
```http
POST /tasks/claim
{"worker_id":"uuid","capabilities":["similarity-search"],"max_concurrency":1}
```
Returns `204 No Content` when no compatible task exists. A successful atomic
claim returns `200 OK`:
```json
{
"task_id":"uuid",
"attempt":1,
"lease_expires_at":"2026-07-23T12:05:00Z",
"workload":"similarity-search",
"input":{"uri":"https://coordinator.example/tasks/uuid/input","sha256":"hex-sha256"},
"parameters":{"query_id":"CHEMBL939","top_k":20}
"name": "lab-worker-01",
"capabilities": ["similarity-search", "similarity-graph"],
"cpu_count": 8,
"memory_mb": 16384
}
```
The claim is one PostgreSQL transaction using `FOR UPDATE SKIP LOCKED`.
`201`:
### Heartbeat
```json
{ "worker_id": "uuid", "heartbeat_interval_seconds": 15 }
```
`cpu_count`/`memory_mb` are accepted for forward compatibility and not yet
persisted. `capabilities` must be non-empty (an allowlisted workload set).
## Claim task
```http
POST /tasks/claim
Authorization: Bearer <token>
Content-Type: application/json
{ "worker_id": "uuid", "capabilities": ["similarity-search"], "max_concurrency": 1 }
```
- `204 No Content`: no compatible task.
- `200 OK`: a task is leased atomically.
```json
{
"task_id": "uuid",
"attempt": 1,
"lease_expires_at": "2026-07-22T12:05:00Z",
"workload": "similarity-search",
"input": { "uri": "https://coordinator/tasks/uuid/input", "sha256": "hex" },
"parameters": { "query_id": "CHEMBL939", "top_k": 20 }
}
```
`max_concurrency` is accepted; the coordinator leases one task per call for now.
## Renew lease (heartbeat)
```http
POST /tasks/{task_id}/heartbeat
Authorization: Bearer <token>
Content-Type: application/json
{"worker_id":"uuid","attempt":1}
{ "worker_id": "uuid", "attempt": 1 }
```
Returns `200 OK` and the renewed deadline:
Response **must** contain a renewed deadline:
```json
{"lease_expires_at":"2026-07-23T12:10:00Z"}
{ "lease_expires_at": "2026-07-22T12:10:00Z" }
```
The Worker schedules its next heartbeat before half of the returned TTL.
The worker schedules the next heartbeat before half of the returned TTL, never
on a fixed interval alone.
### Input download
## Download input or shard (CTX-05)
`GET /tasks/{task_id}/input` returns the claimed task input. The Worker verifies
its SHA-256 before execution. On a redirect to another origin, it removes the
coordinator bearer token.
`GET /tasks/{task_id}/input` returns the artifact owned by the current task. The
worker verifies its SHA-256 before execution. If the URI redirects to another
origin, the worker removes the coordinator bearer token.
## Artifact upload
## Upload a partial artifact (CTX-05)
```http
PUT /tasks/{task_id}/artifacts/{filename}
Authorization: Bearer <token>
Content-Type: text/csv
X-Worker-ID: uuid
X-Task-Attempt: 1
@@ -88,64 +152,49 @@ X-Task-Attempt: 1
<streamed bytes>
```
The coordinator streams the body to storage, checks lease ownership, records
the checksum and returns `201 Created`:
`200`:
```json
{
"artifact_id":"uuid",
"uri":"https://coordinator.example/artifacts/uuid/download",
"sha256":"hex-sha256",
"size_bytes":1234
}
{ "artifact_id": "uuid", "uri": "https://coordinator/artifacts/uuid/download",
"sha256": "hex", "size_bytes": 1234 }
```
The returned URI is the only URI the Worker may send in task completion.
`worker://` and `file://` are invalid.
## Completion and failure
## Complete or fail task
```http
POST /tasks/{task_id}/result
Authorization: Bearer <token>
Content-Type: application/json
{
"worker_id":"uuid",
"attempt":1,
"result":{
"artifact_id":"uuid",
"uri":"https://coordinator.example/artifacts/uuid/download",
"sha256":"hex-sha256",
"content_type":"text/csv"
},
"metrics":{"elapsed_seconds":12.4,"processed_rows":10000}
"worker_id": "uuid",
"attempt": 1,
"result": { "artifact_id": "uuid", "sha256": "hex", "content_type": "text/csv" },
"metrics": { "elapsed_seconds": 12.4, "processed_rows": 10000 }
}
```
The coordinator returns `200`, `201`, or `202` for a valid completion. It must
verify that the artifact belongs to that task and attempt before completing it.
The worker uploads its partial result first (§5.5), then completes with that
`artifact_id`. The coordinator verifies the artifact was stored for this exact
task before accepting it — a worker cannot complete one task with another task's
artifact. No worker-supplied URI is ever persisted.
Use `POST /tasks/{task_id}/failure` only for a failed attempt:
```json
{"worker_id":"uuid","attempt":1,"error_code":"ValueError","error_message":"input checksum mismatch"}
```http
POST /tasks/{task_id}/failure
```
Messages are sanitised: no token, traceback, absolute local path, or raw input.
Same identity fields, plus sanitized `error_code`, `error_message`, `retryable`.
Never a traceback, token, or absolute worker path.
## Error responses
## Idempotency and errors
| Situation | Response |
| --- | --- |
| Invalid JSON, field, or parameter | `400 Bad Request` |
| Missing or invalid authentication | `401 Unauthorized` / `403 Forbidden` |
| Worker/attempt does not own an active lease | `409 Conflict` |
| Artifact does not belong to the task/attempt | `409 Conflict` |
| Same attempt, different completion manifest | `409 Conflict` |
| Unexpected coordinator failure | `500` without internal details |
## Compatibility tests
Contract tests must cover: registration, `204` claim, successful claim,
heartbeat renewal, foreign worker and stale attempt conflicts, streamed upload,
checksum mismatch, success after upload, failure through `/failure`, and
idempotent completion.
| No compatible task | `204` |
| Worker/attempt does not own lease | `409` |
| Artifact does not belong to task/attempt | `409` |
| Same completion, same manifest | `200` idempotent |
| Same attempt, different manifest | `409` |
| Invalid parameters/input | `400` |
| Auth failure | `401` |
| Unknown job/task | `404` |
docs/building-workers.md
+221
View File
@@ -0,0 +1,221 @@
# Building a SciMesh worker
A worker is a process that pulls tasks from the coordinator, runs them, and
returns results. It talks to the coordinator **only over HTTP** — it never sees
the database, and it needs no inbound port (all requests are outbound). This
guide is what you need to implement one (the reference is a Python daemon, but
nothing here is Python-specific).
**Read alongside:**
[`api-contract.md`](api-contract.md) (the contract in prose) and
[`openapi.yaml`](openapi.yaml) (machine-readable — generate a typed client from
it, see the bottom).
---
## The one loop
A worker is essentially this loop:
```text
register once
loop forever:
task = POST /tasks/claim
if no task (204): sleep, continue
download the task's input, verify its checksum
run the workload ── while running, POST heartbeat before the lease expires
upload the result artifact (PUT)
POST /tasks/{id}/result with the artifact id
on any failure: POST /tasks/{id}/failure
```
Everything below fills in the details.
## 0. Auth
Every request except `GET /health` carries a shared bearer token:
```
Authorization: Bearer <COORDINATOR_TOKEN>
```
The token is handed to you out of band (env var / secret) — the same string the
coordinator was started with. Never log it, never send it in an error body.
## 1. Register (once, at startup)
```http
POST /workers/register
{ "name": "lab-worker-01", "capabilities": ["similarity_search"] }
```
Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
- `capabilities` are the workload names you can run — the coordinator only hands
you matching tasks.
- **Keep `worker_id`**. Use it as your identity in every later call. Using the
registered UUID is what lets the coordinator track your liveness (it marks
workers offline after they go silent).
## 2. Claim a task
```http
POST /tasks/claim
{ "worker_id": "<uuid>", "capabilities": ["similarity_search"] }
```
- `200` → a leased task (below).
- `204` → nothing to do; back off a little and poll again.
```json
{
"task_id": "<uuid>",
"attempt": 1,
"lease_expires_at": "2026-07-22T12:05:00Z",
"workload": "similarity_search",
"input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" },
"parameters": { "query_id": "CHEMBL939", "top_k": 20 }
}
```
**`attempt` matters.** Every later call for this task must echo the exact
`attempt` you were handed. A task requeued after a lost lease comes back with a
higher attempt; an old attempt is rejected with `409`.
## 3. Download the input, verify it
```http
GET {input.uri} # e.g. GET /tasks/<uuid>/input
```
Stream it to disk and **check the SHA-256 against `input.sha256`** before
running. A mismatch means a corrupt shard — fail the task with a clear code,
don't process garbage.
> If `input.uri` ever redirects to another host (object storage), **strip the
> `Authorization` header** on the redirect — never send the coordinator token to
> a third party.
## 4. Run — and heartbeat while you run
Long tasks must prove they are alive, or the coordinator's reaper reclaims the
lease and hands the task to someone else.
```http
POST /tasks/{task_id}/heartbeat
{ "worker_id": "<uuid>", "attempt": 1 }
```
Response: `{ "lease_expires_at": "<new deadline>" }`.
- Schedule the next heartbeat at **less than half** the remaining TTL — don't
rely on a fixed interval. If `lease_expires_at` is 2 minutes out, heartbeat
every ~45s.
- The first heartbeat also moves the task from `leased` to `running` on the
server; you don't have to do anything special for that.
If you miss the deadline, your lease expires: a later `heartbeat`/`result` will
come back `409`, and the task is already back in the queue.
## 5. Upload the result artifact
The coordinator owns results — you upload the bytes, it stores them and computes
the checksum. Identity travels in **headers** here, not the body:
```http
PUT /tasks/{task_id}/artifacts/result.csv
Content-Type: text/csv
X-Worker-ID: <uuid>
X-Task-Attempt: 1
<streamed result bytes>
```
Response: `{ "artifact_id": "<uuid>", "uri": "...", "sha256": "<hex>", "size_bytes": 1234 }`.
Keep the returned `artifact_id`.
## 6. Complete the task
```http
POST /tasks/{task_id}/result
{ "worker_id": "<uuid>", "attempt": 1,
"result": { "artifact_id": "<uuid>" },
"metrics": { "elapsed_seconds": 12.4, "processed_rows": 10000 } }
```
- Reference the `artifact_id` you just uploaded **for this task**. The
coordinator verifies it belongs to this task; another task's artifact → `409`.
- **Idempotent:** if your network dropped and you retry the same `artifact_id`,
you get `200` again, not a conflict. Safe to retry.
## 7. …or fail it
```http
POST /tasks/{task_id}/failure
{ "worker_id": "<uuid>", "attempt": 1,
"error_code": "download_failed", "error_message": "checksum mismatch",
"retryable": true }
```
- `retryable: true` → the task returns to the queue while attempts remain (a new
worker gets it at a higher `attempt`).
- `retryable: false` → it fails terminally.
- Send only a short, sanitized `error_code`/`error_message`. **Never** a Python
traceback, a token, or an absolute local path.
---
## Status-code cheat sheet
| Code | Meaning for the worker |
| --- | --- |
| `204` | claim: queue empty — back off and retry |
| `400` | your request is malformed (bad UUID, unknown field) |
| `401` | bad/missing token |
| `404` | task/job/artifact doesn't exist |
| `409` | you don't hold the lease, or your `attempt` is stale, or a different result was already recorded — **stop working on this task**, it's no longer yours |
A `409` is normal, not a crash: it means the coordinator gave the task to
someone else (usually because your lease expired). Log it and move on to the
next claim.
## Config the worker should expose
Per the worker contract, at minimum:
- `SCIMESH_COORDINATOR_URL` (e.g. `http://coordinator:8080`)
- `SCIMESH_WORKER_ID` (or derive from hostname)
- the bearer token
- poll interval and request timeout
- a working directory for downloaded inputs and generated outputs
## Generate a client from the spec
Instead of hand-writing request code, generate it:
```sh
# typed async client
openapi-python-client generate --path docs/openapi.yaml
# or just the Pydantic models
datamodel-codegen --input docs/openapi.yaml --output scimesh_models.py
```
## Try the endpoints by hand first
`coordinator/api/requests.http` walks the whole flow one request at a time
(register → claim → heartbeat → upload → result), and
`coordinator/scripts/smoke.sh` runs it end to end. Read those to see real
request/response bodies before writing code.
## The rules you must not break
1. Never touch the database — HTTP only.
2. Every mutating call carries `worker_id` **and** `attempt`.
3. Verify the input checksum before running.
4. Upload the result artifact **before** calling `/result`.
5. Never persist a `worker://` or local path as a result — the coordinator owns
artifacts.
6. Strip the bearer token on any cross-origin redirect.
7. Sanitize error output — no tracebacks, tokens, or absolute paths.
docs/openapi.yaml
+552
View File
@@ -0,0 +1,552 @@
openapi: 3.0.3
info:
title: SciMesh Coordinator API
version: 1.0.0
description: >
Durable task-queue server for SciMesh. Workers register, claim tasks one at a
time, heartbeat, upload partial-result artifacts, and complete or fail tasks.
Submitters create jobs — either with pre-chunked input URIs or by uploading a
dataset the coordinator chunks itself.
Machine-readable mirror of `docs/api-contract.md` (v1). All timestamps are
UTC, RFC 3339. Every endpoint except `GET /health` requires a bearer token.
Unknown JSON fields are rejected with 400.
servers:
- url: "{scheme}://{host}"
variables:
scheme:
default: http
enum: [http, https]
host:
default: localhost:8080
security:
- bearerAuth: []
tags:
- name: health
- name: workers
- name: jobs
- name: tasks
- name: artifacts
paths:
/health:
get:
tags: [health]
summary: Readiness (probes the database)
security: []
responses:
"200":
description: The coordinator and its database are ready.
content:
application/json:
schema: { $ref: "#/components/schemas/Health" }
"503":
description: The database is unreachable.
content:
application/json:
schema: { $ref: "#/components/schemas/Health" }
/workers/register:
post:
tags: [workers]
summary: Register a worker
requestBody:
required: true
content:
application/json:
schema: { $ref: "#/components/schemas/RegisterRequest" }
responses:
"201":
description: Registered.
content:
application/json:
schema: { $ref: "#/components/schemas/RegisterResponse" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
/jobs:
post:
tags: [jobs]
summary: Create a job from pre-chunked input URIs
description: >
The submitter supplies each chunk's input URI and checksum. To have the
coordinator split a dataset instead, use `POST /jobs/upload`.
requestBody:
required: true
content:
application/json:
schema: { $ref: "#/components/schemas/CreateJobRequest" }
responses:
"201":
description: Job and its tasks were created transactionally.
content:
application/json:
schema: { $ref: "#/components/schemas/JobCreated" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
/jobs/upload:
post:
tags: [jobs]
summary: Upload a dataset; the coordinator chunks it into shard tasks
description: >
multipart/form-data. The text fields (`workload`, `parameters`,
`chunk_rows`) MUST precede the `file` part: the file is streamed, not
buffered, so the fields have to be parsed before it arrives.
requestBody:
required: true
content:
multipart/form-data:
schema: { $ref: "#/components/schemas/UploadJobForm" }
encoding:
file:
contentType: text/tab-separated-values
responses:
"201":
description: Job, input artifact, shard artifacts, and shard tasks created.
content:
application/json:
schema: { $ref: "#/components/schemas/UploadJobResponse" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
/jobs/{job_id}:
get:
tags: [jobs]
summary: Aggregate job progress
parameters:
- $ref: "#/components/parameters/JobID"
responses:
"200":
description: Progress counts and derived status.
content:
application/json:
schema: { $ref: "#/components/schemas/JobProgress" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
/tasks/claim:
post:
tags: [tasks]
summary: Atomically lease one task
requestBody:
required: true
content:
application/json:
schema: { $ref: "#/components/schemas/ClaimRequest" }
responses:
"200":
description: A task was leased.
content:
application/json:
schema: { $ref: "#/components/schemas/ClaimedTask" }
"204":
description: No compatible task is available.
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
/tasks/{task_id}/heartbeat:
post:
tags: [tasks]
summary: Renew the caller's lease
description: >
The response carries a renewed `lease_expires_at`. Schedule the next
heartbeat before half of the remaining TTL, never on a fixed interval alone.
parameters:
- $ref: "#/components/parameters/TaskID"
requestBody:
required: true
content:
application/json:
schema: { $ref: "#/components/schemas/IdentityRequest" }
responses:
"200":
description: Lease renewed.
content:
application/json:
schema: { $ref: "#/components/schemas/ClaimedTask" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
"409": { $ref: "#/components/responses/Conflict" }
/tasks/{task_id}/input:
get:
tags: [tasks]
summary: Download the task's input shard
description: >
Streams the shard bytes for an uploaded-dataset task. The worker verifies
the `X-Checksum-SHA256` header (also delivered as `input.sha256` on claim)
before executing. URI-based tasks have no coordinator-stored input and
return 404.
parameters:
- $ref: "#/components/parameters/TaskID"
responses:
"200":
description: The shard bytes.
headers:
X-Checksum-SHA256:
schema: { type: string }
description: SHA-256 of the shard.
content:
application/octet-stream:
schema: { type: string, format: binary }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
/tasks/{task_id}/artifacts/{filename}:
put:
tags: [tasks, artifacts]
summary: Upload a partial-result artifact
description: >
Streams the body into blob storage. Identity travels in headers, not the
body. The coordinator measures the size and SHA-256 itself and returns them.
parameters:
- $ref: "#/components/parameters/TaskID"
- name: filename
in: path
required: true
schema: { type: string }
- name: X-Worker-ID
in: header
required: true
schema: { type: string }
- name: X-Task-Attempt
in: header
required: true
schema: { type: integer }
requestBody:
required: true
content:
application/octet-stream:
schema: { type: string, format: binary }
text/csv:
schema: { type: string, format: binary }
responses:
"200":
description: Artifact stored.
content:
application/json:
schema: { $ref: "#/components/schemas/ArtifactUploaded" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
"409": { $ref: "#/components/responses/Conflict" }
/tasks/{task_id}/result:
post:
tags: [tasks]
summary: Complete a task with an uploaded result artifact
description: >
References an artifact previously uploaded for THIS task. The coordinator
verifies ownership before accepting it. Idempotent: replaying the same
artifact_id succeeds; a different one for a completed task is a 409.
parameters:
- $ref: "#/components/parameters/TaskID"
requestBody:
required: true
content:
application/json:
schema: { $ref: "#/components/schemas/ResultRequest" }
responses:
"200":
description: Recorded.
content:
application/json:
schema: { $ref: "#/components/schemas/TaskState" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
"409": { $ref: "#/components/responses/Conflict" }
/tasks/{task_id}/failure:
post:
tags: [tasks]
summary: Report a task failure
description: >
`retryable: true` returns the task to the queue while attempts remain;
otherwise it fails terminally. Send only sanitized error fields — never a
traceback, token, or absolute worker path.
parameters:
- $ref: "#/components/parameters/TaskID"
requestBody:
required: true
content:
application/json:
schema: { $ref: "#/components/schemas/FailureRequest" }
responses:
"200":
description: Recorded.
content:
application/json:
schema: { $ref: "#/components/schemas/TaskState" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
"409": { $ref: "#/components/responses/Conflict" }
/artifacts/{artifact_id}/download:
get:
tags: [artifacts]
summary: Download an artifact by id
parameters:
- name: artifact_id
in: path
required: true
schema: { type: string, format: uuid }
responses:
"200":
description: The artifact bytes.
headers:
X-Checksum-SHA256:
schema: { type: string }
content:
application/octet-stream:
schema: { type: string, format: binary }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
components:
securitySchemes:
bearerAuth:
type: http
scheme: bearer
parameters:
JobID:
name: job_id
in: path
required: true
schema: { type: string, format: uuid }
TaskID:
name: task_id
in: path
required: true
schema: { type: string, format: uuid }
responses:
BadRequest:
description: Invalid input.
content:
application/json:
schema: { $ref: "#/components/schemas/Error" }
Unauthorized:
description: Missing or invalid bearer token.
content:
application/json:
schema: { $ref: "#/components/schemas/Error" }
NotFound:
description: The referenced job, task, or artifact does not exist.
content:
application/json:
schema: { $ref: "#/components/schemas/Error" }
Conflict:
description: Lease not held, stale attempt, or a different result already recorded.
content:
application/json:
schema: { $ref: "#/components/schemas/Error" }
schemas:
Health:
type: object
properties:
status: { type: string, example: ok }
Error:
type: object
properties:
error: { type: string, example: "invalid input" }
request_id: { type: string, description: Correlates with the server logs. }
RegisterRequest:
type: object
required: [capabilities]
properties:
name: { type: string, example: lab-worker-01 }
capabilities:
type: array
minItems: 1
items: { type: string }
example: [similarity_search, similarity_graph]
cpu_count:
type: integer
description: Accepted for forward compatibility; not yet persisted.
memory_mb:
type: integer
description: Accepted for forward compatibility; not yet persisted.
RegisterResponse:
type: object
properties:
worker_id: { type: string, format: uuid }
heartbeat_interval_seconds: { type: integer, example: 15 }
ChunkSpec:
type: object
required: [chunk_index, input_uri, input_sha256]
properties:
chunk_index: { type: integer }
workload:
type: string
description: Empty inherits the job's workload.
input_uri: { type: string }
input_sha256: { type: string }
parameters: { type: object, additionalProperties: true }
max_attempts: { type: integer }
CreateJobRequest:
type: object
required: [workload, input_uri, chunks]
properties:
workload: { type: string, example: similarity_search }
input_uri: { type: string }
parameters: { type: object, additionalProperties: true }
chunks:
type: array
minItems: 1
items: { $ref: "#/components/schemas/ChunkSpec" }
JobCreated:
type: object
properties:
id: { type: string, format: uuid }
status: { $ref: "#/components/schemas/JobStatus" }
UploadJobForm:
type: object
required: [workload, file]
properties:
workload: { type: string, example: similarity_search }
parameters:
type: string
description: JSON object, sent as a string form field.
example: '{"top_k":10}'
chunk_rows:
type: integer
description: Data rows per shard. Default 1000.
example: 1000
file:
type: string
format: binary
description: The dataset (TSV; header repeated into each shard).
UploadJobResponse:
type: object
properties:
job_id: { type: string, format: uuid }
task_count: { type: integer, example: 3 }
input_artifact_id: { type: string, format: uuid }
JobProgress:
type: object
properties:
id: { type: string, format: uuid }
status: { $ref: "#/components/schemas/JobStatus" }
total: { type: integer }
pending: { type: integer }
leased: { type: integer }
completed: { type: integer }
failed: { type: integer }
ClaimRequest:
type: object
required: [worker_id]
properties:
worker_id: { type: string }
capabilities:
type: array
items: { type: string }
description: Workloads this worker can run. Empty means "any".
max_concurrency:
type: integer
description: Accepted; the coordinator leases one task per call.
InputRef:
type: object
properties:
uri:
type: string
description: >
For an uploaded shard, a coordinator path `/tasks/{id}/input`. For a
URI-based task, the external input URI.
sha256: { type: string }
ClaimedTask:
type: object
properties:
task_id: { type: string, format: uuid }
job_id: { type: string, format: uuid }
chunk_index: { type: integer }
workload: { type: string }
input: { $ref: "#/components/schemas/InputRef" }
parameters: { type: object, additionalProperties: true }
attempt: { type: integer }
lease_expires_at: { type: string, format: date-time }
IdentityRequest:
type: object
required: [worker_id, attempt]
properties:
worker_id: { type: string }
attempt: { type: integer }
ResultManifest:
type: object
required: [artifact_id]
properties:
artifact_id:
type: string
format: uuid
description: An artifact previously uploaded for this task.
sha256:
type: string
description: Accepted for the worker's own cross-check; the coordinator trusts its stored metadata.
content_type: { type: string }
ResultRequest:
type: object
required: [worker_id, attempt, result]
properties:
worker_id: { type: string }
attempt: { type: integer }
result: { $ref: "#/components/schemas/ResultManifest" }
metrics: { type: object, additionalProperties: true }
FailureRequest:
type: object
required: [worker_id, attempt, error_code]
properties:
worker_id: { type: string }
attempt: { type: integer }
error_code: { type: string, example: download_failed }
error_message: { type: string }
retryable: { type: boolean }
ArtifactUploaded:
type: object
properties:
artifact_id: { type: string, format: uuid }
uri:
type: string
description: Coordinator download path, `/artifacts/{id}/download`.
sha256: { type: string }
size_bytes: { type: integer, format: int64 }
TaskState:
type: object
properties:
id: { type: string, format: uuid }
job_id: { type: string, format: uuid }
status: { $ref: "#/components/schemas/TaskStatus" }
JobStatus:
type: string
enum: [pending, running, completed, failed, cancelled]
TaskStatus:
type: string
enum: [pending, leased, completed, failed, cancelled]
docs/worker-agent-brief.md
+83
View File
@@ -0,0 +1,83 @@
# Brief: build the SciMesh worker against the coordinator
You are implementing the **worker side**. The **coordinator** (Go/PostgreSQL) is
already built, tested, and running on branch `feat/coordinator`. This brief tells
you what exists, where the contract is, and what to deliver.
## What the coordinator already does (done — do not reimplement)
A durable task-queue server. Over HTTP only (workers never touch the database):
- **Worker registry** — `POST /workers/register` returns a `worker_id`; the
coordinator tracks liveness and marks silent workers offline.
- **Jobs** — created from chunk URIs (`POST /jobs`) or by uploading a dataset
(`POST /jobs/upload`), which the coordinator splits into shard tasks itself.
- **Queue** — atomic claim (`FOR UPDATE SKIP LOCKED`), leases, heartbeats
(`leased → running`), a reaper that requeues expired leases, retry budget.
- **Artifacts** — the worker uploads a partial result (`PUT`), the coordinator
stores it (streamed, checksummed) and owns it; completion references an
`artifact_id`, not a worker URI.
- **Input delivery** — `GET /tasks/{id}/input` streams a task's shard.
Full endpoint list and status: `coordinator/README.md`.
## The contract (read these first)
| File | What it is |
| --- | --- |
| `docs/openapi.yaml` | OpenAPI 3.0 — **generate your client from this** |
| `docs/building-workers.md` | step-by-step guide: the claim→heartbeat→upload→complete loop, auth, lease semantics, status codes, and the rules you must not break |
| `docs/api-contract.md` | the same contract in prose |
| `coordinator/api/requests.http` | real request/response examples for every endpoint |
Generate a typed client instead of hand-writing HTTP:
```sh
openapi-python-client generate --path docs/openapi.yaml
# or just models:
datamodel-codegen --input docs/openapi.yaml --output scimesh_models.py
```
## Run the coordinator locally to develop against it
```sh
cd coordinator && docker compose up -d # listens on :8080, migrations auto-applied
make smoke # exercises the whole flow (should pass)
```
Auth: every request except `GET /health` needs `Authorization: Bearer <token>`
(the compose default is `dev-token`; check `coordinator/.env.example`).
## Your deliverable (CTX-06 in PLAN.md)
A worker daemon that:
1. registers at startup and reuses its `worker_id`;
2. claims one task at a time; backs off on `204`;
3. downloads the input via `input.uri` and **verifies its `sha256`** before running;
4. heartbeats before half the lease TTL elapses;
5. uploads the result artifact, then completes the task with that `artifact_id`;
6. reports failures to `/failure` with sanitized error fields;
7. is configured by env: coordinator URL, worker id, token, poll interval, work dir.
## Acceptance criteria
- A worker registers, claims a shard, downloads and checksum-verifies its input,
heartbeats through a long run, uploads a result, and completes it — end to end
against the real coordinator.
- A lost lease (missed heartbeats) surfaces as a clean `409` and the worker moves
on rather than crashing.
- No result ever references a `worker://` or local path — only uploaded artifacts.
- Contract tests run the worker against the real Go coordinator + Postgres in CI.
## Rules you must not break
1. HTTP only — never the database.
2. Every mutating call carries `worker_id` **and** `attempt`; a stale attempt is `409`.
3. Verify the input checksum before executing.
4. Upload the result artifact **before** calling `/result`.
5. Strip the bearer token on any cross-origin redirect.
6. Sanitize errors — never send a traceback, token, or absolute path.
Anything about the coordinator's behavior that isn't clear here is answered by
`docs/openapi.yaml` (authoritative shapes) and `docs/building-workers.md`.