Compare commits
45
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6ef92908a1 | ||
|
|
f953112cfd | ||
|
|
19cbf7f113 | ||
|
|
9ec8f50313 | ||
|
|
08f5478a66 | ||
|
|
bde6cdb4ba | ||
|
|
43ceec1f77 | ||
|
|
f5b16b057f | ||
|
|
f8de0b2b9d | ||
|
|
7547a30bde | ||
|
|
6bac7dad3c | ||
|
|
c7956c4683 | ||
|
|
d648beede2 | ||
|
|
ac9b921401 | ||
|
|
5be87ad762 | ||
|
|
e83e0b5e1f | ||
|
|
ec861edce5 | ||
|
|
2ce9687e52 | ||
|
|
66836b962d | ||
|
|
484ecd0dfa | ||
|
|
983c5843ec | ||
|
|
b4a89dd7c2 | ||
|
|
8af8ddcf48 | ||
|
|
d271170dd2 | ||
|
|
e0ee95cbab | ||
|
|
abfda35170 | ||
|
|
6829632651 | ||
|
|
3b41455b20 | ||
|
|
4fc3c69fdf | ||
|
|
e5ba27951a | ||
|
|
c3243a6b7e | ||
|
|
4a092d2e4e | ||
|
|
58da6ef139 | ||
|
|
6d45406ee0 | ||
|
|
dbf578c500 | ||
|
|
a5945f2d38 | ||
|
|
dc92121acc | ||
|
|
13f9a0b494 | ||
|
|
69c34c9383 | ||
|
|
8a76b13759 | ||
|
|
e7aa0be22d | ||
|
|
5d6390fd98 | ||
|
|
6517145622 | ||
|
|
f1c3163be4 | ||
|
|
bda22666d7 |
@@ -0,0 +1,116 @@
|
||||
# SciMesh Go Coordinator Agent
|
||||
|
||||
## Role
|
||||
|
||||
You are the backend engineer responsible for the SciMesh coordinator.
|
||||
|
||||
Your area includes:
|
||||
|
||||
- Go coordinator service;
|
||||
- PostgreSQL migrations and repositories;
|
||||
- worker registration;
|
||||
- transactional task leasing;
|
||||
- lease renewal and expiry;
|
||||
- artifact metadata and storage;
|
||||
- job/task state transitions;
|
||||
- HTTP API handlers;
|
||||
- reducer orchestration.
|
||||
|
||||
## Read before working
|
||||
|
||||
Always read:
|
||||
|
||||
1. `PLAN.md`
|
||||
2. `docs/api-contract.md`
|
||||
3. the assigned CTX task
|
||||
4. existing migrations and coordinator tests
|
||||
5. `STATUS.md`
|
||||
|
||||
`PLAN.md` is the architectural source of truth.
|
||||
|
||||
## Hard rules
|
||||
|
||||
- Workers never access PostgreSQL.
|
||||
- Task claims must use one transaction and `FOR UPDATE SKIP LOCKED`.
|
||||
- Every task mutation validates `worker_id` and `attempt`.
|
||||
- A task cannot become `completed` before its artifact is durable.
|
||||
- Never trust paths, status, ownership, or artifact identity supplied by a worker without checking PostgreSQL state.
|
||||
- Never expose raw PostgreSQL errors through HTTP.
|
||||
- Never execute arbitrary commands.
|
||||
- Do not silently modify the API contract.
|
||||
- Do not implement unrelated CTX tasks.
|
||||
- Mutating operations must be transactional and context-aware.
|
||||
- A completed job must reference a durable final artifact.
|
||||
- Output ordering must remain deterministic.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. Inspect the current implementation and repository status.
|
||||
2. Read the assigned CTX task and verify that its dependencies are complete.
|
||||
3. Restate the task, scope, assumptions, and acceptance criteria.
|
||||
4. Identify the smallest set of files that must change.
|
||||
5. Implement the smallest complete change.
|
||||
6. Add Go unit tests or PostgreSQL integration tests.
|
||||
7. Run:
|
||||
- `go test ./...`
|
||||
- `go vet ./...`
|
||||
- relevant migration and integration tests
|
||||
8. Review the diff for unrelated changes.
|
||||
9. Produce a structured handoff.
|
||||
|
||||
## Scope control
|
||||
|
||||
One pull request should normally implement one CTX task.
|
||||
|
||||
Do not refactor unrelated packages unless the assigned task cannot be completed
|
||||
without it. Explain the need before making the refactor.
|
||||
|
||||
Do not add Redis, Kafka, RabbitMQ, Kubernetes, cloud storage, or a frontend
|
||||
framework unless a later approved design explicitly requires it.
|
||||
|
||||
## Implementation preferences
|
||||
|
||||
- Prefer small interfaces around storage, queue, and repositories.
|
||||
- Keep HTTP DTOs separate from domain and database structs.
|
||||
- Validate request DTOs before calling services.
|
||||
- Use parameterized SQL only.
|
||||
- Use UTC RFC 3339 timestamps at API boundaries.
|
||||
- Stream artifact bodies; do not read large files fully into memory.
|
||||
- Sanitize errors before returning them to workers or users.
|
||||
- Make completion and reduction idempotent or transactionally protected.
|
||||
|
||||
## Required output
|
||||
|
||||
At completion report:
|
||||
|
||||
### Implemented
|
||||
|
||||
What behavior now works.
|
||||
|
||||
### Files changed
|
||||
|
||||
List each changed file and its purpose.
|
||||
|
||||
### Database changes
|
||||
|
||||
Migrations, constraints, indexes, and queries added.
|
||||
|
||||
### API impact
|
||||
|
||||
Endpoints or contract behavior changed. State `none` when unchanged.
|
||||
|
||||
### Tests
|
||||
|
||||
Commands run and their results.
|
||||
|
||||
### Acceptance criteria
|
||||
|
||||
Checklist copied from the assigned CTX task.
|
||||
|
||||
### Risks and limitations
|
||||
|
||||
Known gaps, assumptions, and follow-up work.
|
||||
|
||||
### Handoff
|
||||
|
||||
State which dependent CTX task may begin next.
|
||||
@@ -0,0 +1,128 @@
|
||||
# SciMesh Integration Agent
|
||||
|
||||
## Role
|
||||
|
||||
You are responsible for compatibility between the Go coordinator, PostgreSQL,
|
||||
Python Worker, artifact storage, and distributed workloads.
|
||||
|
||||
You should not implement large isolated features. Your job is to connect,
|
||||
validate, diagnose, and report the complete vertical slice.
|
||||
|
||||
## Responsibilities
|
||||
|
||||
- maintain compatibility with `docs/api-contract.md`;
|
||||
- verify Go and Python request/response schemas;
|
||||
- verify PostgreSQL migrations and state transitions;
|
||||
- run contract and end-to-end tests;
|
||||
- verify artifact persistence and checksums;
|
||||
- detect incompatible assumptions between components;
|
||||
- verify deterministic reducers;
|
||||
- update `STATUS.md` after accepted merges;
|
||||
- produce milestone-readiness reports.
|
||||
|
||||
## Read before working
|
||||
|
||||
Always read:
|
||||
|
||||
1. `PLAN.md`
|
||||
2. `docs/api-contract.md`
|
||||
3. `STATUS.md`
|
||||
4. CTX tasks included in the integration milestone
|
||||
5. latest developer handoffs
|
||||
6. relevant CI configuration
|
||||
|
||||
## Hard rules
|
||||
|
||||
- Do not mask contract mismatches with silent compatibility hacks.
|
||||
- Do not duplicate domain logic in Go and Python.
|
||||
- Do not modify the API contract without documenting and testing the change.
|
||||
- Do not mark a milestone complete unless its acceptance criteria are demonstrated.
|
||||
- Prefer fixing the source of truth rather than adding adapters around mistakes.
|
||||
- Use real PostgreSQL for integration tests.
|
||||
- Verify actual artifact bytes and checksums, not only status codes.
|
||||
- Verify stale attempt and foreign worker conflicts.
|
||||
- Do not report success when tests were skipped or services were mocked beyond the stated test scope.
|
||||
- Update `STATUS.md` only after evidence is collected.
|
||||
|
||||
## Integration sequence
|
||||
|
||||
1. Start PostgreSQL.
|
||||
2. Apply migrations to an empty database.
|
||||
3. Start the Go coordinator.
|
||||
4. Verify readiness and configuration.
|
||||
5. Start at least two Python Workers.
|
||||
6. Verify worker registration and capability reporting.
|
||||
7. Submit a small fixture job.
|
||||
8. Verify distinct atomic task claims.
|
||||
9. Verify heartbeat renewal from returned lease deadlines.
|
||||
10. Verify input download and SHA-256 checking.
|
||||
11. Verify streamed partial artifact upload.
|
||||
12. Verify task completion references coordinator-owned artifacts.
|
||||
13. Verify deterministic reduction and final artifact download.
|
||||
14. Kill one worker during a task.
|
||||
15. Verify lease expiry and task reassignment.
|
||||
16. Restart the coordinator.
|
||||
17. Verify job state and artifacts remain available.
|
||||
18. Compare distributed output with the local reference.
|
||||
19. Update `STATUS.md`.
|
||||
20. Produce a readiness decision.
|
||||
|
||||
## Required integration scenarios
|
||||
|
||||
- worker registration;
|
||||
- no-task `204`;
|
||||
- successful task claim;
|
||||
- lease renewal;
|
||||
- foreign worker mutation rejected;
|
||||
- stale attempt rejected;
|
||||
- checksum mismatch handled;
|
||||
- worker failure reported through `/failure`;
|
||||
- retry after lease expiry;
|
||||
- idempotent identical completion;
|
||||
- conflicting completion rejected;
|
||||
- final result survives coordinator restart;
|
||||
- two-worker similarity-search equals local CLI output.
|
||||
|
||||
## Required output
|
||||
|
||||
### Tested revisions
|
||||
|
||||
Commit hashes or branch names for coordinator and Python code.
|
||||
|
||||
### Environment
|
||||
|
||||
Go, Python, PostgreSQL versions and relevant configuration.
|
||||
|
||||
### Commands executed
|
||||
|
||||
Exact startup and test commands.
|
||||
|
||||
### Passed scenarios
|
||||
|
||||
List with evidence.
|
||||
|
||||
### Failed scenarios
|
||||
|
||||
List with observed behavior.
|
||||
|
||||
### Contract mismatches
|
||||
|
||||
Field, endpoint, status-code, or ownership differences.
|
||||
|
||||
### Blocking issues
|
||||
|
||||
Issues that prevent the next milestone.
|
||||
|
||||
### STATUS.md update
|
||||
|
||||
Exact status changes made.
|
||||
|
||||
### Readiness decision
|
||||
|
||||
One of:
|
||||
|
||||
- `READY`
|
||||
- `READY WITH NON-BLOCKING LIMITATIONS`
|
||||
- `NOT READY`
|
||||
|
||||
Include the reason.
|
||||
@@ -0,0 +1,133 @@
|
||||
# SciMesh Review Agent
|
||||
|
||||
## Role
|
||||
|
||||
You are a strict code reviewer for SciMesh.
|
||||
|
||||
Do not implement new features unless explicitly asked. Review the current diff
|
||||
against:
|
||||
|
||||
1. `PLAN.md`
|
||||
2. `docs/api-contract.md`
|
||||
3. the assigned CTX task
|
||||
4. relevant agent role rules
|
||||
5. current `STATUS.md`
|
||||
|
||||
Focus on correctness, scope, reliability, security, and test evidence.
|
||||
|
||||
## Review priorities
|
||||
|
||||
### Architecture and scope
|
||||
|
||||
- The change matches exactly the assigned CTX task.
|
||||
- Dependencies are satisfied.
|
||||
- No unrelated refactoring or speculative feature is included.
|
||||
- Coordinator and Worker responsibilities remain separated.
|
||||
- No deferred technology was introduced without approval.
|
||||
|
||||
### Coordinator correctness
|
||||
|
||||
- Task claims are atomic.
|
||||
- Lease owner and attempt are checked on every mutation.
|
||||
- State transitions cannot skip required states.
|
||||
- Artifact durability precedes task completion.
|
||||
- Completion and reduction are idempotent or transactionally protected.
|
||||
- PostgreSQL operations are parameterized and context-aware.
|
||||
- Raw database errors are not exposed.
|
||||
- Storage keys and filenames are sanitized.
|
||||
|
||||
### Worker correctness
|
||||
|
||||
- No database credentials or SQL.
|
||||
- No `shell=True` or arbitrary coordinator-provided commands.
|
||||
- Only allowlisted workloads execute.
|
||||
- Checksums are verified.
|
||||
- Cross-origin redirects do not receive coordinator credentials.
|
||||
- Local paths and raw tracebacks are not sent.
|
||||
- Lease loss prevents successful completion.
|
||||
- Result upload occurs before completion.
|
||||
|
||||
### Scientific correctness
|
||||
|
||||
- A local reference result exists.
|
||||
- Distributed output matches the local result.
|
||||
- Ordering is deterministic.
|
||||
- Reducers are independent of completion order.
|
||||
- Graph pair coverage is complete and disjoint.
|
||||
- No dense N×N matrix is created.
|
||||
- Memory bounds are respected.
|
||||
|
||||
### Tests
|
||||
|
||||
- Success path is covered.
|
||||
- Validation failure is covered.
|
||||
- Conflict and stale-attempt behavior are covered.
|
||||
- Retry and lease expiry are covered when relevant.
|
||||
- Tests use real PostgreSQL where transaction behavior matters.
|
||||
- Contract tests exercise the real Go/Python boundary where relevant.
|
||||
- Test claims in the handoff match actual commands and output.
|
||||
|
||||
### Security and observability
|
||||
|
||||
- Secrets are not logged.
|
||||
- User-controlled values are escaped or sanitized.
|
||||
- Request sizes and timeouts are appropriate where relevant.
|
||||
- Errors returned to users/workers are sanitized.
|
||||
- Logs include useful request/task/worker identifiers without sensitive data.
|
||||
|
||||
## Finding severity
|
||||
|
||||
Return findings ordered by severity:
|
||||
|
||||
1. `BLOCKING`
|
||||
2. `HIGH`
|
||||
3. `MEDIUM`
|
||||
4. `LOW`
|
||||
|
||||
For every finding include:
|
||||
|
||||
- severity;
|
||||
- file and relevant function or line range;
|
||||
- violated invariant or acceptance criterion;
|
||||
- concrete failure scenario;
|
||||
- recommended correction.
|
||||
|
||||
## Required output
|
||||
|
||||
### Summary
|
||||
|
||||
One paragraph describing the reviewed scope and overall quality.
|
||||
|
||||
### Findings
|
||||
|
||||
Ordered by severity. Do not hide important findings in prose.
|
||||
|
||||
### Acceptance criteria verification
|
||||
|
||||
For every CTX acceptance criterion, mark:
|
||||
|
||||
- `VERIFIED`
|
||||
- `NOT VERIFIED`
|
||||
- `FAILED`
|
||||
- `NOT APPLICABLE`
|
||||
|
||||
Include the evidence.
|
||||
|
||||
### Test evidence
|
||||
|
||||
List commands or CI checks inspected.
|
||||
|
||||
### Scope assessment
|
||||
|
||||
State whether the diff contains unrelated changes.
|
||||
|
||||
### Decision
|
||||
|
||||
One of:
|
||||
|
||||
- `APPROVE`
|
||||
- `APPROVE WITH NON-BLOCKING COMMENTS`
|
||||
- `REQUEST CHANGES`
|
||||
|
||||
If there are no blocking findings, explicitly state which CTX acceptance
|
||||
criteria were verified.
|
||||
@@ -0,0 +1,109 @@
|
||||
# SciMesh Python Worker Agent
|
||||
|
||||
## Role
|
||||
|
||||
You are responsible for the Python Worker Daemon and communication with the
|
||||
SciMesh Go coordinator.
|
||||
|
||||
Your area includes:
|
||||
|
||||
- worker registration;
|
||||
- task polling and claiming;
|
||||
- lease heartbeat and renewal;
|
||||
- input artifact download;
|
||||
- SHA-256 verification;
|
||||
- allowlisted workload execution;
|
||||
- partial result upload;
|
||||
- task completion and failure reporting;
|
||||
- worker CLI and configuration;
|
||||
- worker-side unit and contract tests.
|
||||
|
||||
## Read before working
|
||||
|
||||
Always read:
|
||||
|
||||
1. `PLAN.md`
|
||||
2. `docs/api-contract.md`
|
||||
3. the assigned CTX task
|
||||
4. `scimesh/worker/`
|
||||
5. relevant workload adapters
|
||||
6. existing worker and contract tests
|
||||
7. `STATUS.md`
|
||||
|
||||
`docs/api-contract.md` is the compatibility boundary with the Go coordinator.
|
||||
|
||||
## Hard rules
|
||||
|
||||
- The worker never receives or uses database credentials.
|
||||
- Never use `shell=True`.
|
||||
- Never execute commands supplied by the coordinator.
|
||||
- Only explicitly registered and allowlisted workloads may execute.
|
||||
- Reject unknown workload parameters.
|
||||
- Never persist `worker://`, `file://`, or worker-local filesystem paths as result URIs.
|
||||
- Verify downloaded artifact checksums before execution.
|
||||
- Upload result artifacts before submitting task completion.
|
||||
- Remove the coordinator bearer token when a redirect changes origin.
|
||||
- A stale task attempt must not complete successfully.
|
||||
- Failure payloads must not contain tokens, absolute paths, raw tracebacks, or sensitive input contents.
|
||||
- Heartbeat scheduling must use the renewed `lease_expires_at` returned by the coordinator.
|
||||
- Preserve local workload behavior and CLI compatibility.
|
||||
- Do not alter scientific algorithms unless the assigned task explicitly requires it.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. Inspect the current Worker Daemon and relevant tests.
|
||||
2. Compare current behavior with `docs/api-contract.md`.
|
||||
3. Restate the assigned task, endpoints, retry rules, and failure cases.
|
||||
4. Implement only the assigned contract behavior.
|
||||
5. Add unit and real coordinator contract tests.
|
||||
6. Run relevant `pytest` suites.
|
||||
7. Verify that local CLI workloads still work.
|
||||
8. Review logs and error payloads for leaked secrets or paths.
|
||||
9. Produce a structured handoff.
|
||||
|
||||
## Reliability behavior
|
||||
|
||||
- Claim at most the configured concurrency.
|
||||
- Back off when no task is available or the coordinator is unavailable.
|
||||
- Distinguish transient transport errors from permanent task errors.
|
||||
- Stop successful completion after lease loss or `409 Conflict`.
|
||||
- Keep the task workspace until the configured cleanup policy allows removal.
|
||||
- Verify upload response metadata before sending completion.
|
||||
- Treat repeated identical completion as idempotent success when the API allows it.
|
||||
- Never retry an unknown workload or invalid parameter set as a transient failure.
|
||||
|
||||
## Required output
|
||||
|
||||
At completion report:
|
||||
|
||||
### Implemented
|
||||
|
||||
Worker behavior added or changed.
|
||||
|
||||
### API usage
|
||||
|
||||
Endpoints, headers, DTO fields, and status codes handled.
|
||||
|
||||
### Reliability
|
||||
|
||||
Heartbeat, retry, backoff, lease-loss, and cleanup behavior.
|
||||
|
||||
### Files changed
|
||||
|
||||
List each changed file and its purpose.
|
||||
|
||||
### Tests
|
||||
|
||||
Commands run and results, including contract tests.
|
||||
|
||||
### Compatibility
|
||||
|
||||
Effect on existing local workloads and CLI.
|
||||
|
||||
### Acceptance criteria
|
||||
|
||||
Checklist copied from the assigned CTX task.
|
||||
|
||||
### Risks and limitations
|
||||
|
||||
Remaining failure scenarios or contract assumptions.
|
||||
@@ -0,0 +1,133 @@
|
||||
# SciMesh Scientific Workload Agent
|
||||
|
||||
## Role
|
||||
|
||||
You are responsible for distributed scientific workload correctness.
|
||||
|
||||
Your area includes:
|
||||
|
||||
- workload input and parameter validation;
|
||||
- deterministic sharding;
|
||||
- typed `TaskPlan` generation;
|
||||
- bounded-memory worker execution;
|
||||
- partial result formats;
|
||||
- deterministic reduction;
|
||||
- comparison with local reference implementations;
|
||||
- scientific correctness tests.
|
||||
|
||||
Initial production-oriented workloads:
|
||||
|
||||
- `similarity-search`;
|
||||
- `similarity-graph`.
|
||||
|
||||
Genome and plasma workloads are deferred until the first molecular distributed
|
||||
release is stable and accepted.
|
||||
|
||||
## Read before working
|
||||
|
||||
Always read:
|
||||
|
||||
1. `PLAN.md`
|
||||
2. the assigned CTX task
|
||||
3. the current local workload implementation
|
||||
4. the distributed workload protocol
|
||||
5. relevant fixtures and tests
|
||||
6. `STATUS.md`
|
||||
|
||||
The local implementation is the correctness reference unless the assigned task
|
||||
explicitly changes the scientific definition.
|
||||
|
||||
## Hard rules
|
||||
|
||||
- Do not modify coordinator queue or state-machine logic.
|
||||
- Task plans must be JSON-serializable.
|
||||
- Task plans contain validated parameters and artifact references, never foreign local paths.
|
||||
- Results must be deterministic for identical inputs and parameters.
|
||||
- Distributed results must match the local reference implementation.
|
||||
- Do not create or retain a dense N×N similarity matrix.
|
||||
- Similarity graph must compare every unordered pair exactly once.
|
||||
- Reducers must be independent of worker completion order.
|
||||
- Memory usage must remain bounded.
|
||||
- Shards and block indices must be deterministic.
|
||||
- Partial outputs must use documented schemas.
|
||||
- Invalid scientific inputs must fail predictably or be counted according to the workload specification.
|
||||
- Do not change API endpoints or PostgreSQL state semantics.
|
||||
- Do not add genome or plasma implementations before their scope is approved.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. Establish and test the local reference result.
|
||||
2. Define task boundaries and invariants.
|
||||
3. Define the task payload schema.
|
||||
4. Define the partial result schema.
|
||||
5. Implement validation and planner.
|
||||
6. Implement worker execution adapter.
|
||||
7. Implement reducer.
|
||||
8. Compare local and distributed outputs.
|
||||
9. Test multiple shard or block sizes.
|
||||
10. Test different worker completion orders.
|
||||
11. Test retry without changing the final result.
|
||||
12. Document memory bounds and scientific invariants.
|
||||
13. Produce a structured handoff.
|
||||
|
||||
## Similarity-search invariants
|
||||
|
||||
- Resolve `query_id` once during planning.
|
||||
- Each shard keeps a valid TSV header and stable `chunk_index`.
|
||||
- Each shard returns at least the requested global `top_k`.
|
||||
- Query molecule and duplicate canonical query SMILES are excluded as specified.
|
||||
- Global reducer tie-breaking matches local SciMesh.
|
||||
- Final result is independent of task completion order.
|
||||
|
||||
## Similarity-graph invariants
|
||||
|
||||
For blocks `(i, j)`:
|
||||
|
||||
- plan only `i <= j`;
|
||||
- diagonal blocks compare only `a < b`;
|
||||
- off-diagonal blocks compare all cross-block pairs;
|
||||
- no self-loops;
|
||||
- no duplicate unordered edges;
|
||||
- support the documented threshold direction;
|
||||
- distributed edge set equals local brute-force output;
|
||||
- result is invariant to block size and task completion order.
|
||||
|
||||
## Required output
|
||||
|
||||
At completion report:
|
||||
|
||||
### Scientific definition
|
||||
|
||||
What exactly is computed.
|
||||
|
||||
### Sharding strategy
|
||||
|
||||
How input is split and why coverage is complete.
|
||||
|
||||
### Task payload
|
||||
|
||||
Documented JSON-compatible fields.
|
||||
|
||||
### Partial result
|
||||
|
||||
File format, ordering, and metrics.
|
||||
|
||||
### Reduction algorithm
|
||||
|
||||
How partial outputs become the final result.
|
||||
|
||||
### Correctness invariants
|
||||
|
||||
Properties that must always hold.
|
||||
|
||||
### Tests
|
||||
|
||||
Local versus distributed comparisons and commands run.
|
||||
|
||||
### Performance constraints
|
||||
|
||||
Expected memory complexity and known bottlenecks.
|
||||
|
||||
### Acceptance criteria
|
||||
|
||||
Checklist copied from the assigned CTX task.
|
||||
@@ -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
|
||||
@@ -0,0 +1,28 @@
|
||||
name: python
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- "scimesh/**"
|
||||
- "tests/**"
|
||||
- "pyproject.toml"
|
||||
- ".github/workflows/python.yml"
|
||||
pull_request:
|
||||
paths:
|
||||
- "scimesh/**"
|
||||
- "tests/**"
|
||||
- "pyproject.toml"
|
||||
- ".github/workflows/python.yml"
|
||||
|
||||
jobs:
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
cache: pip
|
||||
- run: python -m pip install --upgrade pip
|
||||
- run: python -m pip install -e '.[dev]'
|
||||
- run: pytest -q
|
||||
@@ -11,3 +11,7 @@ results/
|
||||
*_similarities.csv
|
||||
test_results.csv
|
||||
test_structures/
|
||||
|
||||
# Local coordinator-worker execution state
|
||||
worker-data*/
|
||||
scimesh-worker-data/
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
# Repository Guidelines
|
||||
|
||||
## Project Structure & Module Organization
|
||||
|
||||
SciMesh is a Python package for molecular-similarity workloads. Source lives in
|
||||
`scimesh/`: `chemistry/` reads data and makes fingerprints, `workloads/`
|
||||
contains commands, and `core/` provides the workload protocol and registry.
|
||||
The worker daemon in `scimesh/worker/` is a coordinator client, not a database
|
||||
client. Tests are in `tests/`; specifications in `docs/`; roadmap: `PLAN.md`.
|
||||
|
||||
For distributed work, read `.agents/`, `docs/api-contract.md`,
|
||||
and `STATUS.md`. Use one CTX task per pull request; local workloads are the
|
||||
scientific reference.
|
||||
|
||||
## Build, Test, and Development Commands
|
||||
|
||||
Create a virtual environment, install the package with development tools, and
|
||||
run the suite:
|
||||
|
||||
```bash
|
||||
python -m venv .venv
|
||||
source .venv/bin/activate
|
||||
pip install -e '.[dev]'
|
||||
pytest
|
||||
```
|
||||
|
||||
Use `pytest tests/test_similarity_graph.py` for one module. Exercise the CLI
|
||||
with `scimesh help` or `scimesh similarity-search --help`. Run `python -m build`
|
||||
only when packaging is needed; install `build` first if necessary.
|
||||
|
||||
## Coding Style & Naming Conventions
|
||||
|
||||
Target Python 3.10+; type public APIs and exchanged data. Use four spaces,
|
||||
`snake_case` for modules,
|
||||
functions, and variables, `PascalCase` for classes, and descriptive test names
|
||||
such as `test_graph_is_deterministic_across_block_sizes`. Keep CLI parsing in
|
||||
workload modules and register new workloads through `scimesh/core/registry.py`;
|
||||
do not add workload-specific logic to the main CLI.
|
||||
|
||||
Prefer few dependencies. RDKit is the chemistry dependency.
|
||||
For worker/coordinator work, keep network payloads explicit and multi-line;
|
||||
never make the worker access PostgreSQL directly.
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
Use pytest and add a regression test for every defect. Similarity code must be
|
||||
checked against a small brute-force or fully sorted reference. Graph results
|
||||
must be deterministic, have no self-loops or duplicate pairs, and remain
|
||||
stable for different block sizes. Worker changes need success and failure
|
||||
tests: checksum mismatch, lease failure, upload failure, and safe reporting.
|
||||
Run the full `pytest` suite before committing.
|
||||
|
||||
## Commit & Pull Request Guidelines
|
||||
|
||||
Use short imperative commit subjects, for example `Add graph threshold mode` or
|
||||
`Fix worker result and lease contracts`. Keep one logical change per commit.
|
||||
In a pull request, state the problem, behaviour changed, tests run, and any API
|
||||
or documentation changes. Link the relevant `CTX-*` item in `PLAN.md` for
|
||||
distributed work. Do not commit datasets, generated CSV/PNG files, `.venv/`,
|
||||
tokens, or local worker artifacts.
|
||||
|
||||
## Security & Protocol Rules
|
||||
|
||||
Upload worker results through the coordinator before posting completion; never
|
||||
submit `file://` or `worker://` result URIs. Send failures to `/failure`, not
|
||||
`/result`. Do not log bearer tokens, raw tracebacks, or private local paths.
|
||||
@@ -0,0 +1,979 @@
|
||||
# SciMesh: master implementation plan
|
||||
|
||||
> **Purpose.** This document is the source plan for turning SciMesh from a
|
||||
> local molecular CLI into a local-first distributed scientific-computation
|
||||
> platform. It is intentionally detailed enough to split into independent task
|
||||
> briefs for developers or coding agents.
|
||||
>
|
||||
> **Planning baseline.** This branch starts from `Workers`: the Python package
|
||||
> has local `similarity-search` and `similarity-graph` workloads plus a Worker
|
||||
> Daemon client. The coordinator and PostgreSQL implementation do not yet
|
||||
> exist. The Worker contract and the Go/PostgreSQL design briefs in `docs/` are
|
||||
> part of this plan.
|
||||
|
||||
---
|
||||
|
||||
## 1. Product goal
|
||||
|
||||
SciMesh accepts a scientific run, turns it into independent tasks, dispatches
|
||||
them to polling workers, persists task state and artifacts, combines partial
|
||||
results, and exposes the final result and progress to a user.
|
||||
|
||||
The first production-oriented vertical slice is molecular computation:
|
||||
|
||||
- `similarity-search`: exact top-k Tanimoto search over ChEMBL shards;
|
||||
- `similarity-graph`: exact sparse Tanimoto graph, where each pair is compared
|
||||
once and only edges satisfying the chosen threshold are retained.
|
||||
|
||||
The platform must later support other scientific workloads without changing the
|
||||
coordinator or worker state machine.
|
||||
|
||||
```text
|
||||
User / simple UI / CLI
|
||||
|
|
||||
v
|
||||
Go coordinator + PostgreSQL + coordinator artifact storage
|
||||
|
|
||||
+-- creates Job/Run -> Tasks -> leases one task at a time
|
||||
|
|
||||
v
|
||||
Python Worker Daemons (outbound HTTP only)
|
||||
|
|
||||
+-- download artifact -> execute allowlisted workload -> upload partial result
|
||||
|
|
||||
v
|
||||
Coordinator reducer -> final artifact -> download/status API
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Scope, non-goals, and decisions
|
||||
|
||||
### 2.1 In scope
|
||||
|
||||
- Go 1.22+ coordinator service with PostgreSQL 15+;
|
||||
- Python Worker Daemon running existing SciMesh workloads locally;
|
||||
- durable job, task, worker, and artifact metadata;
|
||||
- local coordinator-managed artifact storage for the first deployment;
|
||||
- HTTP API for submit, poll/claim, heartbeat, artifact transfer, completion,
|
||||
failure, job status, and result download;
|
||||
- sharding and reduction for the two molecular workloads;
|
||||
- a small server-rendered or static HTML status page after the API works;
|
||||
- automated unit, integration, and contract tests;
|
||||
- a reproducible local demo using one coordinator and two or more workers.
|
||||
|
||||
### 2.2 Explicit non-goals for the first release
|
||||
|
||||
- cloud object storage, Kubernetes, autoscaling, and multi-region operation;
|
||||
- arbitrary shell commands sent by coordinator to workers;
|
||||
- user accounts, multi-tenancy, billing, or sophisticated authorization;
|
||||
- GPU scheduling and multiprocessing inside a worker;
|
||||
- Docker as a required runtime dependency;
|
||||
- video/CV processing implementation;
|
||||
- a React/Vue frontend;
|
||||
- exact resumability of a subprocess after host power loss.
|
||||
|
||||
### 2.3 Architectural decisions already made
|
||||
|
||||
| Decision | Choice | Rationale |
|
||||
| --- | --- | --- |
|
||||
| Coordinator | Go + `net/http` | One durable service for API, queue, artifacts, and reducer orchestration. |
|
||||
| Database | PostgreSQL + `pgxpool` | Transactional leasing and concurrent `SKIP LOCKED` claims. |
|
||||
| Migration tool | `golang-migrate` SQL migrations | Schema is reviewable independently of Go code. |
|
||||
| Workers | Python | Reuses RDKit and the existing SciMesh workload code. |
|
||||
| Worker connectivity | Outbound HTTP polling | Workers require no public inbound ports. |
|
||||
| Queue model | Database rows, not a separate broker | Sufficient for the initial local-first deployment. |
|
||||
| Artifact storage | Coordinator filesystem first | Durable and simple; can later be replaced by S3-compatible storage behind an interface. |
|
||||
| Workload execution | Explicit allowlist + typed parameters | Never execute coordinator-provided shell commands. |
|
||||
| Result correctness | Artifact upload before task completion | A completed task must reference a durable, coordinator-accessible result. |
|
||||
|
||||
### 2.4 Rules that must never be violated
|
||||
|
||||
1. Workers never use PostgreSQL credentials or execute SQL.
|
||||
2. A task is leased atomically to at most one worker attempt.
|
||||
3. A task becomes `completed` only after its result artifact is durable and
|
||||
verified by the coordinator.
|
||||
4. Every mutating worker request includes `worker_id` and `attempt`; stale
|
||||
attempts receive `409 Conflict`.
|
||||
5. Worker tokens are never forwarded to an external presigned download URL.
|
||||
6. The worker runs an explicit Python command list, never `shell=True`.
|
||||
7. Result and reducer outputs are deterministic for identical inputs and
|
||||
parameters.
|
||||
8. Similarity graph tasks must cover each original molecule pair exactly once.
|
||||
9. No workload may create or retain a dense N×N similarity matrix.
|
||||
10. The coordinator must not trust worker-supplied artifact paths, status, or
|
||||
ownership claims without checking task state in PostgreSQL.
|
||||
|
||||
---
|
||||
|
||||
## 3. Glossary and canonical lifecycle
|
||||
|
||||
| Term | Meaning |
|
||||
| --- | --- |
|
||||
| **Worker** | A registered process/machine capable of claiming tasks. |
|
||||
| **Job** | A user-requested full computation. The UI may call it a **Run**; the database/API use `job`. |
|
||||
| **Task** | One independently executable unit of a job. |
|
||||
| **Attempt** | A monotonically increasing execution lease for a task. |
|
||||
| **Lease** | Temporary exclusive assignment of a task to one worker. |
|
||||
| **Artifact** | A durable input, shard, partial result, final result, or log file. |
|
||||
| **Planner** | Workload code that validates a job and emits task payloads. |
|
||||
| **Runner** | Worker-side code that executes one typed task locally. |
|
||||
| **Reducer** | Coordinator-side code that merges all completed partial results into a final artifact. |
|
||||
|
||||
### 3.1 Job state machine
|
||||
|
||||
```text
|
||||
CREATED -> PLANNING -> RUNNING -> REDUCING -> COMPLETED
|
||||
| | |
|
||||
+-------> FAILED <---+
|
||||
RUNNING -> CANCELLED
|
||||
```
|
||||
|
||||
### 3.2 Task state machine
|
||||
|
||||
```text
|
||||
PENDING -> LEASED -> RUNNING -> COMPLETED
|
||||
| | |
|
||||
| +--------> FAILED
|
||||
+-----> PENDING |
|
||||
lease expiry +-> PENDING (attempts remain)
|
||||
```
|
||||
|
||||
`LEASED` means the coordinator has returned a task. `RUNNING` means the worker
|
||||
has successfully sent its first heartbeat/start acknowledgement. A lease expiry
|
||||
may return either `LEASED` or `RUNNING` tasks to `PENDING` when attempts remain.
|
||||
|
||||
### 3.3 Artifact lifecycle
|
||||
|
||||
```text
|
||||
upload input -> INPUT artifact -> planner creates SHARD artifacts
|
||||
worker downloads SHARD -> executes -> uploads PARTIAL_RESULT artifact
|
||||
reducer reads partial artifacts -> writes FINAL_RESULT artifact
|
||||
user downloads final artifact
|
||||
```
|
||||
|
||||
No `file://` or `worker://` URI is valid in persisted result metadata.
|
||||
|
||||
---
|
||||
|
||||
## 4. Target repository layout
|
||||
|
||||
Keep the existing Python package at the repository root and add a Go
|
||||
coordinator as a self-contained subproject.
|
||||
|
||||
```text
|
||||
SciMesh/
|
||||
PLAN.md
|
||||
docs/
|
||||
worker-daemon-task.md
|
||||
database-integration-task.md
|
||||
api-contract.md # created in Phase 0
|
||||
scimesh/ # Python local workloads and worker daemon
|
||||
chemistry/
|
||||
core/
|
||||
distributed/ # planner/reducer contracts, Phase 3
|
||||
worker/
|
||||
workloads/
|
||||
tests/
|
||||
unit/
|
||||
integration/
|
||||
contract/
|
||||
coordinator/
|
||||
go.mod
|
||||
cmd/coordinator/main.go
|
||||
internal/
|
||||
config/
|
||||
domain/
|
||||
httpapi/
|
||||
queue/
|
||||
reducer/
|
||||
storage/
|
||||
store/postgres/
|
||||
migrations/
|
||||
tests/
|
||||
scripts/
|
||||
dev-start.sh # optional convenience script, not required runtime
|
||||
```
|
||||
|
||||
Do not move the mature local workload code merely to satisfy this layout. Move
|
||||
only when a distributed contract requires a clear shared module.
|
||||
|
||||
---
|
||||
|
||||
## 5. Cross-service API contract
|
||||
|
||||
The API contract is a compatibility boundary. Before either side is implemented,
|
||||
copy this section into `docs/api-contract.md` and treat it as versioned.
|
||||
|
||||
All worker endpoints require bearer authentication. Worker identity and attempt
|
||||
are checked against the current task lease in PostgreSQL.
|
||||
|
||||
### 5.1 Register worker
|
||||
|
||||
```http
|
||||
POST /workers/register
|
||||
Authorization: Bearer <token>
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"name": "lab-worker-01",
|
||||
"capabilities": ["similarity-search", "similarity-graph"],
|
||||
"cpu_count": 8,
|
||||
"memory_mb": 16384
|
||||
}
|
||||
```
|
||||
|
||||
Response:
|
||||
|
||||
```json
|
||||
{
|
||||
"worker_id": "uuid",
|
||||
"heartbeat_interval_seconds": 15
|
||||
}
|
||||
```
|
||||
|
||||
### 5.2 Claim task
|
||||
|
||||
```http
|
||||
POST /tasks/claim
|
||||
Authorization: Bearer <token>
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"worker_id": "uuid",
|
||||
"capabilities": ["similarity-search", "similarity-graph"],
|
||||
"max_concurrency": 1
|
||||
}
|
||||
```
|
||||
|
||||
- `204 No Content`: no compatible task is available.
|
||||
- `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.example/tasks/uuid/input",
|
||||
"sha256": "hex-sha256"
|
||||
},
|
||||
"parameters": {
|
||||
"query_id": "CHEMBL939",
|
||||
"top_k": 20
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 5.3 Renew lease
|
||||
|
||||
```http
|
||||
POST /tasks/{task_id}/heartbeat
|
||||
Authorization: Bearer <token>
|
||||
Content-Type: application/json
|
||||
|
||||
{"worker_id": "uuid", "attempt": 1}
|
||||
```
|
||||
|
||||
The response **must** contain a renewed deadline:
|
||||
|
||||
```json
|
||||
{"lease_expires_at": "2026-07-22T12:10:00Z"}
|
||||
```
|
||||
|
||||
The worker schedules the next heartbeat before half of this returned TTL, never
|
||||
using only a fixed interval.
|
||||
|
||||
### 5.4 Download input or shard
|
||||
|
||||
`GET /tasks/{task_id}/input` returns an artifact owned by the current task. The
|
||||
worker verifies its SHA-256 before execution. If the returned URI redirects to
|
||||
another origin, the worker must remove the coordinator bearer token.
|
||||
|
||||
### 5.5 Upload a partial artifact
|
||||
|
||||
```http
|
||||
PUT /tasks/{task_id}/artifacts/{filename}
|
||||
Authorization: Bearer <token>
|
||||
Content-Type: text/csv
|
||||
X-Worker-ID: uuid
|
||||
X-Task-Attempt: 1
|
||||
|
||||
<streamed bytes>
|
||||
```
|
||||
|
||||
The coordinator streams the body to its artifact storage, verifies ownership,
|
||||
stores metadata and checksum, then returns:
|
||||
|
||||
```json
|
||||
{
|
||||
"artifact_id": "uuid",
|
||||
"uri": "https://coordinator.example/artifacts/uuid/download",
|
||||
"sha256": "hex-sha256",
|
||||
"size_bytes": 1234
|
||||
}
|
||||
```
|
||||
|
||||
### 5.6 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}
|
||||
}
|
||||
```
|
||||
|
||||
`POST /tasks/{task_id}/failure` uses the same identity fields and contains only
|
||||
sanitized `error_code` and `error_message` values. No Python traceback, token,
|
||||
or absolute worker path may be sent.
|
||||
|
||||
### 5.7 Idempotency and errors
|
||||
|
||||
| Situation | Required response |
|
||||
| --- | --- |
|
||||
| No compatible task | `204` |
|
||||
| Worker/attempt does not own lease | `409` |
|
||||
| Artifact does not belong to task/attempt | `409` |
|
||||
| Same completion, same manifest | `200`/`202` idempotent success |
|
||||
| Same attempt, different manifest | `409` |
|
||||
| Invalid parameters/input | `400` |
|
||||
| Worker authentication failure | `401`/`403` |
|
||||
|
||||
---
|
||||
|
||||
## 6. PostgreSQL data model
|
||||
|
||||
The existing brief describes `jobs` and `tasks`. Add first-class worker and
|
||||
artifact records before implementation. The database is the source of truth for
|
||||
state; files are referenced by metadata rather than discovered from directories.
|
||||
|
||||
### 6.1 Tables
|
||||
|
||||
#### `workers`
|
||||
|
||||
| Field | Notes |
|
||||
| --- | --- |
|
||||
| `id UUID PK` | Returned on registration |
|
||||
| `name text` | Human-readable; unique only if desired |
|
||||
| `capabilities jsonb` | Allowlisted workload names |
|
||||
| `status` | `online`, `busy`, `offline` |
|
||||
| `last_heartbeat_at timestamptz` | Liveness visibility |
|
||||
| `created_at`, `updated_at timestamptz` | Audit |
|
||||
|
||||
#### `jobs`
|
||||
|
||||
| Field | Notes |
|
||||
| --- | --- |
|
||||
| `id UUID PK` | User-visible identifier |
|
||||
| `workload text` | Registered distributed workload name |
|
||||
| `status` | `created`, `planning`, `running`, `reducing`, `completed`, `failed`, `cancelled` |
|
||||
| `parameters jsonb` | Validated job parameters |
|
||||
| `input_artifact_id UUID` | Original uploaded data |
|
||||
| `result_artifact_id UUID nullable` | Final result |
|
||||
| `total_tasks`, `completed_tasks`, `failed_tasks` | Progress counters, updated transactionally |
|
||||
| timestamps and `error_message` | Audit and failure status |
|
||||
|
||||
#### `tasks`
|
||||
|
||||
| Field | Notes |
|
||||
| --- | --- |
|
||||
| `id UUID PK`, `job_id UUID FK` | Identity and ownership |
|
||||
| `chunk_index int` | Unique within job; deterministic reducer order |
|
||||
| `workload text`, `parameters jsonb` | Typed task payload |
|
||||
| `input_artifact_id UUID` | Dataset or shard |
|
||||
| `status` | `pending`, `leased`, `running`, `completed`, `failed`, `cancelled` |
|
||||
| `attempt`, `max_attempts` | Retry accounting |
|
||||
| `lease_owner UUID nullable`, `lease_expires_at nullable` | Exclusive lease |
|
||||
| `result_artifact_id UUID nullable` | Uploaded partial output |
|
||||
| `metrics jsonb`, errors, timestamps, `version int` | Audit and concurrency |
|
||||
|
||||
#### `artifacts`
|
||||
|
||||
| Field | Notes |
|
||||
| --- | --- |
|
||||
| `id UUID PK` | Stable artifact identity |
|
||||
| `job_id UUID FK`, `task_id UUID FK nullable` | Ownership |
|
||||
| `kind` | `input`, `shard`, `partial_result`, `final_result`, `log` |
|
||||
| `filename`, `storage_key`, `content_type` | Storage metadata |
|
||||
| `size_bytes`, `sha256` | Integrity metadata |
|
||||
| `created_at` | Audit |
|
||||
|
||||
### 6.2 Required constraints and queries
|
||||
|
||||
- unique `(job_id, chunk_index)` for tasks;
|
||||
- unique `(task_id, kind)` for single-result task workloads;
|
||||
- claim index: `(status, lease_expires_at, created_at)`;
|
||||
- `attempt >= 0`, `max_attempts > 0`;
|
||||
- a leased/running task must have owner and expiry;
|
||||
- a completed task must reference a `partial_result` artifact;
|
||||
- a completed job must reference a `final_result` artifact;
|
||||
- `list_completed_results(job_id)` orders by `chunk_index`, never insertion time.
|
||||
|
||||
Atomic claim uses one PostgreSQL transaction with `FOR UPDATE SKIP LOCKED`.
|
||||
Never write a `SELECT pending task` followed by an unguarded later `UPDATE`.
|
||||
|
||||
---
|
||||
|
||||
## 7. Workload evolution
|
||||
|
||||
The current local `Workload` CLI interface is intentionally small. Distributed
|
||||
execution needs a second, explicit contract. Do not force every CLI helper into
|
||||
the distributed interface; adapt only workloads that can be planned and reduced.
|
||||
|
||||
```python
|
||||
class DistributedWorkload(Protocol):
|
||||
name: str
|
||||
version: str
|
||||
|
||||
def validate_job(self, input_path: Path, parameters: dict[str, object]) -> None: ...
|
||||
def plan(self, input_path: Path, parameters: dict[str, object], workspace: Path) -> list[TaskPlan]: ...
|
||||
def execute_task(self, task: TaskPlan, workspace: Path) -> RunResult: ...
|
||||
def reduce(self, partial_results: list[Path], parameters: dict[str, object], workspace: Path) -> FinalResult: ...
|
||||
def describe(self) -> dict[str, object]: ...
|
||||
```
|
||||
|
||||
`TaskPlan` is JSON-serializable and contains only validated parameters plus
|
||||
artifact IDs/URIs. It never contains shell commands or local paths from another
|
||||
machine.
|
||||
|
||||
### 7.1 Distributed similarity-search
|
||||
|
||||
1. Validate either `query_id` or `query_smiles`, never both.
|
||||
2. Resolve a `query_id` once during planning and persist the canonical query
|
||||
SMILES/identity in the job metadata.
|
||||
3. Split the input TSV into deterministic shard artifacts. Every shard keeps a
|
||||
header and a stable `chunk_index`.
|
||||
4. Each worker streams one shard, skips invalid SMILES and the query molecule,
|
||||
and writes a sorted local top-k CSV.
|
||||
5. The reducer merges all local top-k outputs with a bounded heap using the same
|
||||
deterministic tie-breaker as local SciMesh.
|
||||
6. The final CSV must equal the current single-process result for the same input
|
||||
and options.
|
||||
|
||||
Important: a local shard top-k must retain at least the global requested `k`.
|
||||
The reducer cannot recover a candidate discarded by every shard.
|
||||
|
||||
### 7.2 Distributed similarity-graph
|
||||
|
||||
1. Parse valid molecules once during planning or produce deterministic molecule
|
||||
block artifacts with stable block indices.
|
||||
2. Emit one task for every block pair `(i, j)` where `i <= j`.
|
||||
3. A diagonal task compares only pairs inside one block with local index `a < b`.
|
||||
4. An off-diagonal task compares every molecule in block `i` with every molecule
|
||||
in block `j`.
|
||||
5. Each task emits an edge-list CSV only for pairs satisfying the chosen
|
||||
`threshold_direction` and threshold.
|
||||
6. The reducer merges edge files, verifies no duplicate unordered pair, and
|
||||
writes a deterministic sort order.
|
||||
|
||||
Correctness invariant:
|
||||
|
||||
```text
|
||||
union(task-pairs) = all unordered molecule pairs
|
||||
intersection(task-pairs) = empty
|
||||
```
|
||||
|
||||
For a small fixture, the distributed result must exactly equal the local
|
||||
brute-force graph for both `greater` and `less` threshold directions.
|
||||
|
||||
### 7.3 Future workload policy
|
||||
|
||||
A new workload is accepted only when it supplies:
|
||||
|
||||
- an input/parameter validator;
|
||||
- an explicit sharding strategy;
|
||||
- bounded-memory task execution;
|
||||
- deterministic reduction semantics;
|
||||
- fixture-based local and distributed correctness tests;
|
||||
- a `describe()` payload for UI/API discovery.
|
||||
|
||||
---
|
||||
|
||||
## 8. Milestones and dependency order
|
||||
|
||||
```text
|
||||
M0 contracts + test fixtures
|
||||
-> M1 Go coordinator skeleton + migrations
|
||||
-> M2 transactional queue + worker registry
|
||||
-> M3 artifact storage + worker contract integration
|
||||
-> M4 distributed similarity-search vertical slice
|
||||
-> M5 reducer + status/result API
|
||||
-> M6 distributed similarity-graph
|
||||
-> M7 UI, observability, CI hardening
|
||||
```
|
||||
|
||||
Do not start graph distribution before the search vertical slice proves the
|
||||
full artifact/lease/reducer lifecycle.
|
||||
|
||||
---
|
||||
|
||||
## 9. Task briefs for implementation
|
||||
|
||||
Each section below is deliberately self-contained. When assigning work, copy
|
||||
the task block plus the **Shared context** section and any listed dependency.
|
||||
|
||||
### Shared context for every assignee
|
||||
|
||||
```text
|
||||
Project: SciMesh
|
||||
Architecture: Go/PostgreSQL coordinator; Python workers; local artifact storage.
|
||||
Hard rules: workers have no DB credentials; task claims are atomic; artifacts
|
||||
must be durable before completion; no arbitrary commands; deterministic output.
|
||||
Read first: PLAN.md sections 2, 3, 5, and the task's dependencies.
|
||||
Do not refactor unrelated files. Add or update tests with every behavior change.
|
||||
```
|
||||
|
||||
### CTX-00 — Freeze API and error contract
|
||||
|
||||
**Goal:** Create `docs/api-contract.md` from section 5 and make it the single
|
||||
source of truth for the Go coordinator and Python Worker.
|
||||
|
||||
**Depends on:** none.
|
||||
|
||||
**Deliverables:** endpoint table, JSON schemas/examples, headers, artifact
|
||||
upload ownership rule, error mapping, retry/idempotency policy, and an explicit
|
||||
version marker (`v1`).
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- all worker endpoints and their paths/methods are listed;
|
||||
- heartbeat response includes `lease_expires_at`;
|
||||
- completion references only coordinator-uploaded artifacts;
|
||||
- failure endpoint is distinct from result endpoint;
|
||||
- document states whether the coordinator uses UTC RFC 3339 timestamps.
|
||||
|
||||
**Out of scope:** implementing HTTP handlers.
|
||||
|
||||
### CTX-01 — Bootstrap Go coordinator
|
||||
|
||||
**Goal:** Add the `coordinator/` Go module and a minimal healthy HTTP service.
|
||||
|
||||
**Depends on:** CTX-00.
|
||||
|
||||
**Deliverables:** `go.mod`; config parser; structured logger; `GET /healthz`;
|
||||
graceful shutdown; `pgxpool` lifecycle; migration command documentation.
|
||||
|
||||
**Inputs:** `DATABASE_URL`, `COORDINATOR_STORAGE_DIR`, `COORDINATOR_ADDR`,
|
||||
`COORDINATOR_TOKEN`, request timeout, pool size.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- `go test ./...` passes;
|
||||
- service refuses to start with missing/invalid required configuration;
|
||||
- `/healthz` returns database readiness without exposing secrets;
|
||||
- shutdown closes HTTP server and `pgxpool` cleanly;
|
||||
- migration execution is explicit, not implicit on production startup.
|
||||
|
||||
**Out of scope:** queue endpoints and UI.
|
||||
|
||||
### CTX-02 — PostgreSQL schema and migration set
|
||||
|
||||
**Goal:** Implement versioned SQL migrations for workers, jobs, tasks, and
|
||||
artifacts from section 6.
|
||||
|
||||
**Depends on:** CTX-01, CTX-00.
|
||||
|
||||
**Deliverables:** up/down SQL migrations; enum/check constraints; indexes;
|
||||
repository-domain structs as needed for scanning rows.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- empty PostgreSQL database migrates up and down in integration tests;
|
||||
- constraints reject invalid state combinations;
|
||||
- migration test uses `TEST_DATABASE_URL`, never SQLite;
|
||||
- `tasks(job_id, chunk_index)` uniqueness and claim index exist;
|
||||
- no application code creates tables dynamically.
|
||||
|
||||
**Out of scope:** HTTP handlers and lease operations.
|
||||
|
||||
### CTX-03 — Transactional queue and lease repository
|
||||
|
||||
**Goal:** Implement PostgreSQL repository operations for create, claim, renew,
|
||||
complete, fail, expiry, and job status.
|
||||
|
||||
**Depends on:** CTX-02.
|
||||
|
||||
**Required operations:**
|
||||
|
||||
```go
|
||||
CreateJobWithTasks(ctx, input)
|
||||
ClaimNextTask(ctx, workerID, capabilities, leaseDuration)
|
||||
RenewLease(ctx, taskID, workerID, attempt, leaseDuration)
|
||||
CompleteTask(ctx, input)
|
||||
FailTask(ctx, input)
|
||||
ExpireLeases(ctx, now)
|
||||
GetJobStatus(ctx, jobID)
|
||||
ListCompletedResults(ctx, jobID)
|
||||
```
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- concurrent claim test proves one task/attempt has one owner;
|
||||
- renewal returns a new `lease_expires_at` timestamp;
|
||||
- stale attempt cannot renew, upload, fail, or complete;
|
||||
- expired task returns to pending or fails after final attempt;
|
||||
- completion is idempotent for identical manifest and conflicts for a different
|
||||
manifest;
|
||||
- all operations are context-aware and parameterized.
|
||||
|
||||
**Out of scope:** artifact byte storage and Python worker changes.
|
||||
|
||||
### CTX-04 — Worker registry and coordinator HTTP handlers
|
||||
|
||||
**Goal:** Expose CTX-03 through versioned `net/http` handlers and add worker
|
||||
registration/liveness state.
|
||||
|
||||
**Depends on:** CTX-00, CTX-03.
|
||||
|
||||
**Endpoints:** `POST /workers/register`, `POST /tasks/claim`,
|
||||
`POST /tasks/{id}/heartbeat`, `GET /jobs/{id}`, and readiness endpoints.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- request DTOs are validated before calling services;
|
||||
- no raw PostgreSQL errors leave the process;
|
||||
- `204`, `400`, `401/403`, and `409` match CTX-00;
|
||||
- heartbeat response returns renewed deadline;
|
||||
- tests cover an offline worker and an expired lease;
|
||||
- logs include request ID, worker ID, task ID, attempt, operation.
|
||||
|
||||
### CTX-05 — Coordinator artifact storage
|
||||
|
||||
**Goal:** Store input, shard, partial-result, and final-result files durably in
|
||||
the coordinator filesystem and persist their metadata.
|
||||
|
||||
**Depends on:** CTX-02, CTX-04.
|
||||
|
||||
**Endpoints:** authenticated input upload/download and
|
||||
`PUT /tasks/{id}/artifacts/{filename}`.
|
||||
|
||||
**Implementation rules:**
|
||||
|
||||
- stream request bodies to a staging file; never `ReadAll` a result;
|
||||
- use sanitized generated storage keys, not user paths;
|
||||
- calculate SHA-256 while streaming;
|
||||
- atomically rename staging file only after successful write;
|
||||
- verify worker lease owner and attempt before accepting task output;
|
||||
- return a coordinator-owned durable URI and artifact ID.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- uploaded artifact is downloadable after coordinator restart;
|
||||
- a foreign worker receives `409`;
|
||||
- large-file test proves bounded-memory streaming behavior;
|
||||
- checksum and size are persisted;
|
||||
- failed upload leaves no visible artifact or orphan staging file.
|
||||
|
||||
### CTX-06 — Align Python Worker with the live Go contract
|
||||
|
||||
**Goal:** Adapt `scimesh/worker/` to CTX-00 through CTX-05 without changing
|
||||
local workload algorithms.
|
||||
|
||||
**Depends on:** CTX-00, CTX-04, CTX-05.
|
||||
|
||||
**Required behavior:**
|
||||
|
||||
- register worker capabilities at startup;
|
||||
- claim one task at a time;
|
||||
- remove bearer token when input redirect changes origin;
|
||||
- verify input checksum;
|
||||
- renew lease from returned `lease_expires_at`;
|
||||
- stream uploaded output with worker/attempt headers;
|
||||
- submit the returned artifact manifest only after upload;
|
||||
- send errors to `/failure`;
|
||||
- preserve task directory until configured cleanup;
|
||||
- reject unknown workload parameters.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- Python contract tests run against the real Go service in CI;
|
||||
- no result references `worker://` or local paths;
|
||||
- long fake runner causes multiple lease renewals;
|
||||
- lost lease stops successful completion and reports conflict cleanly;
|
||||
- worker CLI documents all environment variables and exits non-zero on invalid
|
||||
configuration.
|
||||
|
||||
### CTX-07 — Distributed workload protocol and job planner
|
||||
|
||||
**Goal:** Add the Python `DistributedWorkload` adapter protocol and coordinator
|
||||
planner bridge for registered workloads.
|
||||
|
||||
**Depends on:** CTX-05, CTX-06.
|
||||
|
||||
**Deliverables:** typed task-plan JSON; validator; planner registry; reducer
|
||||
registry; workload descriptions; no direct coordinator dependency in local
|
||||
chemistry helpers.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- unknown workload rejected before a job/task is written;
|
||||
- every plan payload is JSON-serializable and uses artifact references;
|
||||
- planner failure leaves no partial job/tasks transaction;
|
||||
- test fixture demonstrates planning a two-shard dummy workload;
|
||||
- reducer receives completed artifacts ordered by `chunk_index`.
|
||||
|
||||
### CTX-08 — Distributed similarity-search vertical slice
|
||||
|
||||
**Goal:** Implement planner, worker execution adapter, reducer, and end-to-end
|
||||
tests for distributed exact top-k similarity search.
|
||||
|
||||
**Depends on:** CTX-07.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- query ID is resolved before shards execute, or query SMILES is validated once;
|
||||
- shards are deterministic and contain valid TSV headers;
|
||||
- each shard produces local top-k CSV and reports invalid-SMILES counts;
|
||||
- reducer output matches local single-process SciMesh byte-for-byte apart from
|
||||
permitted elapsed-time metadata;
|
||||
- query molecule and duplicate canonical query SMILES are excluded;
|
||||
- tie order is deterministic across worker completion order;
|
||||
- test uses at least two workers and one retry.
|
||||
|
||||
### CTX-09 — Job reducer orchestration and final result API
|
||||
|
||||
**Goal:** When all task results are complete, run the appropriate reducer once,
|
||||
persist a final artifact, update job state, and expose download/status.
|
||||
|
||||
**Depends on:** CTX-07, CTX-08.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- only one reducer process may transition a job into `reducing`;
|
||||
- reducer is idempotent or protected by state/version transaction;
|
||||
- reducer failure marks job failed with a sanitized error;
|
||||
- `GET /jobs/{id}` reports counters and state correctly;
|
||||
- final artifact has stored checksum and downloadable URI;
|
||||
- integration test covers complete job lifecycle.
|
||||
|
||||
### CTX-10 — Distributed similarity-graph
|
||||
|
||||
**Goal:** Implement block-pair planning, execution, and deterministic reduction
|
||||
for the exact sparse similarity graph.
|
||||
|
||||
**Depends on:** CTX-08, CTX-09.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- each block pair is planned once with stable `(left_block, right_block)`;
|
||||
- diagonal and off-diagonal comparisons obey the pair invariant in section 7.2;
|
||||
- no task creates a dense matrix;
|
||||
- both `greater` and `less` threshold directions are preserved;
|
||||
- reducer detects duplicate unordered pairs and fails safely;
|
||||
- distributed output equals local brute-force output on a small fixture;
|
||||
- result is invariant to block size and worker completion order.
|
||||
|
||||
### CTX-11 — Minimal dashboard and operator views
|
||||
|
||||
**Goal:** Add a small server-rendered or static HTML UI to inspect jobs, tasks,
|
||||
workers, and download final artifacts.
|
||||
|
||||
**Detailed delivery plan:** [`docs/web-interface-plan.md`](docs/web-interface-plan.md).
|
||||
The plan deliberately starts with a clearly labelled diagnostic UI before
|
||||
CTX-09 enables final result downloads.
|
||||
|
||||
**Depends on:** CTX-04, CTX-09.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- show worker name/capabilities/status/last heartbeat;
|
||||
- show job status and completed/total task progress;
|
||||
- show task attempt, lease owner, and sanitized error;
|
||||
- refresh with simple polling; no frontend framework required;
|
||||
- final result link is available only in `completed` state;
|
||||
- HTML escapes user-controlled values.
|
||||
|
||||
### CTX-12 — Reliability, security, and CI hardening
|
||||
|
||||
**Goal:** Make the vertical slice safe to demo and difficult to regress.
|
||||
|
||||
**Depends on:** CTX-06 through CTX-11.
|
||||
|
||||
**Work items:**
|
||||
|
||||
- CI jobs for Python tests, Go tests, `go vet`, migrations, and contract tests;
|
||||
- coordinator request-size limits and timeouts;
|
||||
- token configuration/rotation documentation;
|
||||
- structured logs and correlation IDs;
|
||||
- cleanup policy for failed worker directories and stale staging files;
|
||||
- metrics/status endpoint suitable for local monitoring;
|
||||
- retry/backoff test matrix;
|
||||
- release checklist and local two-worker demo script.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- clean checkout can run the documented demo;
|
||||
- all quality gates execute in CI;
|
||||
- no secret appears in logs/tests/errors;
|
||||
- failure/retry scenarios have automated coverage;
|
||||
- README contains architecture diagram, security caveat, and troubleshooting.
|
||||
|
||||
---
|
||||
|
||||
## 10. Suggested assignment bundles
|
||||
|
||||
These bundles minimize overlap. Do not run tasks from the same bundle in
|
||||
parallel unless one engineer owns integration.
|
||||
|
||||
| Bundle | Tasks | Recommended owner |
|
||||
| --- | --- | --- |
|
||||
| Contract and data | CTX-00, CTX-02, CTX-03 | Go/PostgreSQL engineer |
|
||||
| Coordinator API | CTX-01, CTX-04, CTX-05 | Go backend engineer |
|
||||
| Python transport | CTX-06 | Python worker engineer |
|
||||
| Distributed computation | CTX-07, CTX-08, CTX-10 | Scientific Python engineer |
|
||||
| Product surface | CTX-09, CTX-11 | Full-stack/backend engineer |
|
||||
| Quality gate | CTX-12 | DevOps/QA engineer |
|
||||
|
||||
Suggested order for a small team:
|
||||
|
||||
```text
|
||||
Week 1: CTX-00 + CTX-01 + CTX-02
|
||||
Week 2: CTX-03 + CTX-04
|
||||
Week 3: CTX-05 + CTX-06
|
||||
Week 4: CTX-07 + CTX-08
|
||||
Week 5: CTX-09 + CTX-10
|
||||
Week 6: CTX-11 + CTX-12
|
||||
```
|
||||
|
||||
This is an ordering aid, not a deadline commitment. Start the next milestone
|
||||
only after its predecessor's acceptance criteria are demonstrably met.
|
||||
|
||||
---
|
||||
|
||||
## 11. Test strategy
|
||||
|
||||
### 11.1 Python unit tests
|
||||
|
||||
- SMILES parsing, fingerprints, deterministic top-k and graph ordering;
|
||||
- Worker parameter allowlist and CLI argument mapping;
|
||||
- checksum verification and cross-origin auth stripping;
|
||||
- heartbeat rescheduling from returned TTL;
|
||||
- artifact upload manifest generation;
|
||||
- runner failure sanitization and cleanup behavior.
|
||||
|
||||
### 11.2 Go unit tests
|
||||
|
||||
- configuration parser and error mapping;
|
||||
- storage-key sanitization;
|
||||
- request DTO validation;
|
||||
- service state transition guards;
|
||||
- reducer invocation selection.
|
||||
|
||||
### 11.3 PostgreSQL integration tests
|
||||
|
||||
- migrations up/down;
|
||||
- concurrent `ClaimNextTask` without duplicate lease;
|
||||
- lease expiry/retry/exhaustion;
|
||||
- stale attempt conflict;
|
||||
- artifact ownership and task completion idempotency;
|
||||
- deterministic partial-result retrieval order.
|
||||
|
||||
### 11.4 Contract tests
|
||||
|
||||
Run Python Worker tests against a real Go coordinator and PostgreSQL instance.
|
||||
At minimum prove:
|
||||
|
||||
1. worker registers;
|
||||
2. worker claims one task;
|
||||
3. worker receives a renewed lease;
|
||||
4. input checksum is verified;
|
||||
5. result is uploaded and available from coordinator storage;
|
||||
6. completion persists the correct artifact;
|
||||
7. bad checksum or runner failure reaches `/failure`;
|
||||
8. a stale worker cannot complete after its lease has expired.
|
||||
|
||||
### 11.5 End-to-end tests
|
||||
|
||||
- two-worker similarity-search run equals local CLI output;
|
||||
- retry one failed shard and still finish deterministically;
|
||||
- graph on small fixture has all and only brute-force threshold edges;
|
||||
- stopping a worker mid-task requeues its task after lease expiry;
|
||||
- final result remains downloadable after coordinator restart.
|
||||
|
||||
---
|
||||
|
||||
## 12. Review gates
|
||||
|
||||
Before merging a task, reviewer checks:
|
||||
|
||||
### Every task
|
||||
|
||||
- scope matches one CTX block;
|
||||
- tests cover new behavior and pass;
|
||||
- no unrelated refactor or generated data is committed;
|
||||
- errors are sanitized and parameters validated;
|
||||
- documentation/API contract updated when behavior changes.
|
||||
|
||||
### Worker changes
|
||||
|
||||
- no SQL, DB credentials, or arbitrary command execution;
|
||||
- no worker-local URI persisted as a result;
|
||||
- artifact transfer is streamed and checksum verified;
|
||||
- heartbeat uses returned deadline;
|
||||
- cross-origin requests do not receive coordinator token.
|
||||
|
||||
### Coordinator changes
|
||||
|
||||
- mutating actions are transactional;
|
||||
- lease/task ownership checked at every mutation;
|
||||
- PostgreSQL operations are parameterized;
|
||||
- result completion is idempotent;
|
||||
- artifact ownership and storage keys are validated;
|
||||
- state transitions cannot skip required intermediate conditions.
|
||||
|
||||
### Scientific workload changes
|
||||
|
||||
- local reference result exists;
|
||||
- distributed result is compared to local result on fixtures;
|
||||
- reduction ordering is deterministic;
|
||||
- pair coverage/duplicate invariants are tested for graphs;
|
||||
- memory remains bounded as specified.
|
||||
|
||||
---
|
||||
|
||||
## 13. Deferred backlog
|
||||
|
||||
Do not start these before CTX-12 is accepted.
|
||||
|
||||
- Replace local artifact storage with S3/MinIO behind an `ArtifactStore` API.
|
||||
- Add worker labels/capacity-aware scheduling and concurrency > 1.
|
||||
- Add cancellation propagation to workers.
|
||||
- Add image outputs and final PDF reporting to job artifacts.
|
||||
- Add CV/video workloads using the same planner/runner/reducer contract.
|
||||
- Add observability export (Prometheus/OpenTelemetry).
|
||||
- Add per-user/project authorization and signed artifact URLs.
|
||||
- Add shard caching and content-addressed input deduplication.
|
||||
- Add job priority and fair scheduling.
|
||||
- Add a CLI for submitting and monitoring remote jobs.
|
||||
|
||||
---
|
||||
|
||||
## 14. Definition of the first usable distributed release
|
||||
|
||||
The release is complete when all of the following are true:
|
||||
|
||||
1. A user uploads a small ChEMBL TSV and starts a similarity-search job.
|
||||
2. The Go coordinator creates PostgreSQL job/task/artifact records.
|
||||
3. Two Python workers register, claim distinct shards, renew leases, and upload
|
||||
partial CSVs.
|
||||
4. The coordinator reduces partial results into a deterministic final CSV.
|
||||
5. The user observes progress and downloads the final CSV.
|
||||
6. Killing one worker requeues only its lease after expiry; the job still
|
||||
completes within its attempt budget.
|
||||
7. The same fixture run matches local SciMesh output.
|
||||
8. CI executes Python, Go, PostgreSQL, and contract tests successfully.
|
||||
|
||||
Until these eight conditions are met, SciMesh is a promising set of components,
|
||||
not yet a complete distributed platform.
|
||||
@@ -1,6 +1,12 @@
|
||||
# SciMesh
|
||||
|
||||
SciMesh is a small local framework for scientific workloads on molecular datasets. It currently provides exact molecular similarity search and exact sparse similarity-graph construction. It runs in one local Python process: there is no network service, multiprocessing, coordinator, database, or dense similarity matrix.
|
||||
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
|
||||
runs exact similarity search and sparse similarity-graph construction locally in
|
||||
one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
|
||||
coordinator and Python worker can run a diagnostic, shard-based
|
||||
`similarity-search` pipeline locally. Its CSV artifacts are not a global result
|
||||
until CTX-07--09 add planning and reduction; use the local CLI for scientific
|
||||
results today. See [`STATUS.md`](STATUS.md).
|
||||
|
||||
The ChEMBL TSV database is intentionally not included in this repository. Download it separately and pass its path to the commands below. The expected columns are `chembl_id` and `canonical_smiles`.
|
||||
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
# SciMesh Status
|
||||
|
||||
**Updated:** 2026-07-24
|
||||
**Branch baseline:** `main` at `f953112` (distributed pipeline hardening)
|
||||
|
||||
## Current state
|
||||
|
||||
The local Python molecular workloads are implemented and tested. They provide
|
||||
the reference behaviour for future distributed execution:
|
||||
|
||||
- `similarity-search`: streaming ChEMBL TSV search, Morgan fingerprints,
|
||||
Tanimoto scoring, heap-based top-k, CSV and image output;
|
||||
- `similarity-graph`: exact sparse graph, block-based pair comparisons,
|
||||
deterministic CSV output;
|
||||
- Python Worker skeleton: claim, heartbeat, input checksum validation,
|
||||
artifact upload, completion and failure reporting.
|
||||
|
||||
The Go coordinator and its PostgreSQL-backed task lifecycle are implemented:
|
||||
registration, atomic claiming, lease renewal, artifact storage, dataset
|
||||
chunking, result/failure reporting, and job progress. The Python worker now
|
||||
uses the live coordinator contract; its HTTP path was exercised against a real
|
||||
Docker PostgreSQL stack on 2026-07-23.
|
||||
|
||||
## Milestone tracker
|
||||
|
||||
| CTX | Status | Notes |
|
||||
| --- | --- | --- |
|
||||
| CTX-00 API and error contract | Implemented | Contract, OpenAPI, and request examples are in `docs/`. |
|
||||
| CTX-01 Go coordinator bootstrap | Implemented | Go service and Docker runtime in `coordinator/`. |
|
||||
| CTX-02 PostgreSQL migrations | Implemented | Applied by the Compose migration service. |
|
||||
| CTX-03 Transactional queue | Implemented | Real-PostgreSQL integration tests cover atomic claims and concurrency. |
|
||||
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. |
|
||||
| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. |
|
||||
| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. |
|
||||
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. The concrete molecular planner/reducer remains CTX-08/09. |
|
||||
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
|
||||
| CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
|
||||
| CTX-10 Distributed similarity-graph | Not started | Local reference exists. |
|
||||
| CTX-11 Dashboard/operator view | Implemented (diagnostic scope) | Protected local view: job/task/worker status, validated similarity-search upload, diagnostic partial-artifact download, and bounded polling. Final-result reduction remains CTX-09. |
|
||||
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
|
||||
|
||||
## Next recommended assignment
|
||||
|
||||
Assign **CTX-08** to the workload role: implement the molecular
|
||||
`similarity-search` planner and worker adapter on top of the accepted CTX-07
|
||||
contract.
|
||||
|
||||
## Known constraints
|
||||
|
||||
- The CTX-07 protocol is implemented, but no concrete molecular planner or
|
||||
reducer is registered yet; the operator UI labels `partial_result` files as
|
||||
diagnostic and cannot present them as final output.
|
||||
Use the local `scimesh` CLI for complete workload results.
|
||||
- The worker/coordinator flow currently accepts both underscore API workload
|
||||
names and hyphenated CLI names while the contract is consolidated.
|
||||
- A real-stack worker test uses a small `query_smiles` shard. Resolving a
|
||||
`query_id` once and sharing it across shards belongs to CTX-07.
|
||||
- The coordinator accepts uploaded distributed jobs only for
|
||||
`similarity-search` with `query_smiles`. It rejects `similarity-graph` until
|
||||
CTX-10 supplies cross-shard pair planning.
|
||||
|
||||
## Update rule
|
||||
|
||||
The integration role updates this file only after collecting command output,
|
||||
test evidence, and accepted changes. State facts, revision hashes, blockers,
|
||||
and the next unblocked CTX task; do not mark work complete based on plans alone.
|
||||
@@ -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
|
||||
@@ -0,0 +1,32 @@
|
||||
# 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
|
||||
|
||||
# Optional local operator UI. Use a separate value; never reuse the worker token.
|
||||
# When empty, /ui is disabled.
|
||||
UI_AUTH_TOKEN=
|
||||
|
||||
# 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
|
||||
@@ -0,0 +1,6 @@
|
||||
/coordinator
|
||||
/bin/
|
||||
.env
|
||||
*.out
|
||||
/logs/
|
||||
/data/
|
||||
@@ -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
|
||||
@@ -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`, осталось их подключить.
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,99 @@
|
||||
.PHONY: build run test test-integration vet lint tidy check migrate-up migrate-down up down down-clean logs ps rebuild psql smoke
|
||||
|
||||
# `check` deliberately uses its own Compose project and host ports. This keeps
|
||||
# it from connecting to or replacing a developer's local PostgreSQL instance.
|
||||
CHECK_PROJECT ?= scimesh-check
|
||||
CHECK_POSTGRES_PORT ?= 55432
|
||||
CHECK_COORDINATOR_PORT ?= 18080
|
||||
CHECK_HOST ?= http://localhost:$(CHECK_COORDINATOR_PORT)
|
||||
CHECK_TOKEN ?= dev-token
|
||||
CHECK_DATABASE_URL ?= postgres://scimesh:scimesh@localhost:$(CHECK_POSTGRES_PORT)/scimesh?sslmode=disable
|
||||
CHECK_COMPOSE = POSTGRES_PORT=$(CHECK_POSTGRES_PORT) COORDINATOR_PORT=$(CHECK_COORDINATOR_PORT) docker compose -p $(CHECK_PROJECT)
|
||||
|
||||
# --- 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 ./...
|
||||
$(CHECK_COMPOSE) up -d --build
|
||||
@echo "waiting for the coordinator to be ready..."
|
||||
@attempt=0; until curl -fsS "$(CHECK_HOST)/health" >/dev/null; do \
|
||||
attempt=$$((attempt + 1)); \
|
||||
if [ $$attempt -ge 30 ]; then $(CHECK_COMPOSE) logs coordinator; exit 1; fi; \
|
||||
sleep 1; \
|
||||
done
|
||||
TEST_DATABASE_URL="$(CHECK_DATABASE_URL)" \
|
||||
go test -tags=integration ./internal/storage/postgres/ -v
|
||||
HOST="$(CHECK_HOST)" TOKEN="$(CHECK_TOKEN)" ./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
|
||||
@@ -0,0 +1,194 @@
|
||||
# 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)
|
||||
```
|
||||
|
||||
To enable the local operator UI, set a separate credential before starting:
|
||||
|
||||
```sh
|
||||
UI_AUTH_TOKEN='local-ui-secret' make up
|
||||
# Open http://localhost:8080/ui and use any username with this value as password.
|
||||
```
|
||||
|
||||
The UI is disabled by default and never accepts the worker bearer token.
|
||||
It shows recent jobs, task/worker state, and the per-job partial artifacts.
|
||||
Those files are explicitly diagnostic until the CTX-09 reducer creates a final
|
||||
result; the UI does not present them as final scientific output.
|
||||
|
||||
`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.
|
||||
|
||||
For the complete local verification, including an isolated Docker PostgreSQL
|
||||
and the HTTP smoke flow, run:
|
||||
|
||||
```sh
|
||||
make check
|
||||
```
|
||||
|
||||
It uses Compose project `scimesh-check` and ports `55432`/`18080` by default,
|
||||
so it does not connect to a PostgreSQL already running on `5432`. Override
|
||||
`CHECK_POSTGRES_PORT`, `CHECK_COORDINATOR_PORT`, or `CHECK_PROJECT` if needed.
|
||||
@@ -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.
|
||||
@@ -0,0 +1,128 @@
|
||||
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)
|
||||
uiReadRepo = postgres.NewUIReadRepo(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, cfg.DefaultMaxAttempts),
|
||||
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
|
||||
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
|
||||
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, tx, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
|
||||
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
|
||||
Dashboard: usecase.NewDashboard(uiReadRepo),
|
||||
}
|
||||
|
||||
// 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, jobRepo, tx, 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, cfg.UIToken))
|
||||
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
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}
|
||||
# Empty disables /ui. Set this separately from the worker token.
|
||||
UI_AUTH_TOKEN: ${UI_AUTH_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:
|
||||
@@ -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
|
||||
)
|
||||
@@ -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=
|
||||
@@ -0,0 +1,149 @@
|
||||
// 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"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// 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")
|
||||
|
||||
// maxShardBytes bounds the coordinator memory used by one in-progress shard.
|
||||
// The uploaded file may be much larger: it is first stored on disk, then split
|
||||
// in small bounded pieces. Operators can lower rowsPerShard when this limit is
|
||||
// reached rather than exhausting the coordinator process.
|
||||
const maxShardBytes = 64 << 20 // 64 MiB
|
||||
|
||||
// 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 {
|
||||
return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
|
||||
}
|
||||
|
||||
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
|
||||
// A maxRows value of zero means unlimited. This lets an operator make a small,
|
||||
// representative pipeline check without materialising a second dataset file.
|
||||
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit)
|
||||
}
|
||||
|
||||
// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It
|
||||
// validates the two columns every local SciMesh workload requires before any
|
||||
// shard task is persisted, while generic SplitTSV remains reusable for future
|
||||
// non-chemistry workloads.
|
||||
func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit)
|
||||
}
|
||||
|
||||
func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error {
|
||||
if rowsPerShard <= 0 {
|
||||
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
|
||||
}
|
||||
if maxRows < 0 {
|
||||
return fmt.Errorf("maxRows must be non-negative, got %d", maxRows)
|
||||
}
|
||||
|
||||
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()...)
|
||||
if validateHeader != nil {
|
||||
if err := validateHeader(header); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
if len(header)+1 > maxShardBytes {
|
||||
return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes)
|
||||
}
|
||||
buf.Write(header)
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
if buf.Len()+len(sc.Bytes())+1 > maxShardBytes {
|
||||
return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes)
|
||||
}
|
||||
buf.Write(sc.Bytes())
|
||||
buf.WriteByte('\n')
|
||||
rows++
|
||||
|
||||
if rows == rowsPerShard {
|
||||
if err := flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if maxRows > 0 && index*rowsPerShard+rows == maxRows {
|
||||
break
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
func validateChEMBLHeader(header []byte) error {
|
||||
seen := make(map[string]struct{})
|
||||
for _, field := range strings.Split(strings.TrimPrefix(string(header), "\ufeff"), "\t") {
|
||||
seen[field] = struct{}{}
|
||||
}
|
||||
if _, ok := seen["chembl_id"]; !ok {
|
||||
return fmt.Errorf("TSV is missing required column chembl_id")
|
||||
}
|
||||
if _, ok := seen["canonical_smiles"]; !ok {
|
||||
return fmt.Errorf("TSV is missing required column canonical_smiles")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
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])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
|
||||
input := "h\nr1\nr2\nr3\nr4\nr5\n"
|
||||
var shards []string
|
||||
err := SplitTSVLimit(strings.NewReader(input), 2, 3, func(_ int, shard io.Reader) error {
|
||||
b, _ := io.ReadAll(shard)
|
||||
shards = append(shards, string(b))
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := strings.Join(shards, ""), "h\nr1\nr2\nh\nr3\n"; got != want {
|
||||
t.Errorf("limited shards = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
|
||||
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
|
||||
func(int, io.Reader) error { return nil })
|
||||
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
|
||||
t.Errorf("err = %v, want missing-column error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
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
|
||||
Attempt *int // required for a partial result; nil for non-worker artifacts
|
||||
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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
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")
|
||||
ErrJobNotCancellable = errors.New("job cannot be cancelled")
|
||||
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")
|
||||
)
|
||||
@@ -0,0 +1,128 @@
|
||||
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
|
||||
Cancelled 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.Job.Status == JobCancelled:
|
||||
return JobCancelled
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
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},
|
||||
{"cancelled job wins over task histogram", JobProgress{Job: Job{Status: JobCancelled}, Total: 3, Done: 1, Cancelled: 2}, JobCancelled},
|
||||
}
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,304 @@
|
||||
// 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, now time.Time) bool {
|
||||
return t.LeaseOwner != nil && t.LeaseExpiresAt != nil && now.Before(*t.LeaseExpiresAt) &&
|
||||
*t.LeaseOwner == worker && t.Attempt == attempt &&
|
||||
(t.Status == TaskLeased || t.Status == TaskRunning)
|
||||
}
|
||||
|
||||
// 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, now time.Time) 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
|
||||
}
|
||||
if t.LeaseExpiresAt == nil || !now.Before(*t.LeaseExpiresAt) {
|
||||
return ErrLeaseConflict
|
||||
}
|
||||
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, now, until time.Time) error {
|
||||
if err := t.verifyLease(worker, attempt, now); 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, now); 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, now); 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
|
||||
}
|
||||
|
||||
// Cancel prevents any further worker transition for a task that has not
|
||||
// reached a terminal result. A cancelled lease deliberately becomes invalid:
|
||||
// a worker still running locally must not upload or complete after its job was
|
||||
// stopped by the operator.
|
||||
func (t *Task) Cancel(now time.Time) bool {
|
||||
if t.Status == TaskCompleted || t.Status == TaskFailed || t.Status == TaskCancelled {
|
||||
return false
|
||||
}
|
||||
t.Status = TaskCancelled
|
||||
t.LeaseOwner = nil
|
||||
t.LeaseExpiresAt = nil
|
||||
t.ErrorCode = nil
|
||||
t.ErrorMessage = nil
|
||||
t.CompletedAt = &now
|
||||
t.Version++
|
||||
return true
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
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 TestCancelInvalidatesLeaseButPreservesTerminalTask(t *testing.T) {
|
||||
task := leasedTask(1, 3)
|
||||
if !task.Cancel(testNow) {
|
||||
t.Fatal("leased task should be cancelled")
|
||||
}
|
||||
if task.Status != TaskCancelled || task.LeaseOwner != nil || task.LeaseExpiresAt != nil {
|
||||
t.Errorf("cancelled task = %+v", task)
|
||||
}
|
||||
if task.Cancel(testLater) {
|
||||
t.Error("cancelled task must not be changed twice")
|
||||
}
|
||||
completed := &Task{Status: TaskCompleted}
|
||||
if completed.Cancel(testNow) {
|
||||
t.Error("completed task must remain terminal")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) {
|
||||
task := leasedTask(1, 3)
|
||||
until := testLater.Add(time.Hour)
|
||||
|
||||
if err := task.RenewLease(testWorker, 1, testNow, 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, testNow, 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, testNow, 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, testNow, testLater) // -> running
|
||||
task2.ExpireLease(testLater)
|
||||
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, testNow, 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, testNow, until); !errors.Is(err, ErrLeaseConflict) {
|
||||
t.Errorf("err = %v, want ErrLeaseConflict", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiredLeaseRejectsRenewalCompletionAndFailure(t *testing.T) {
|
||||
task := leasedTask(1, 3)
|
||||
expired := testLater.Add(time.Nanosecond)
|
||||
|
||||
if err := task.RenewLease(testWorker, 1, expired, expired.Add(time.Minute)); !errors.Is(err, ErrLeaseConflict) {
|
||||
t.Errorf("renew expired lease: err = %v, want ErrLeaseConflict", err)
|
||||
}
|
||||
if err := task.CompleteWith(testResult, nil, testWorker, 1, expired); !errors.Is(err, ErrLeaseConflict) {
|
||||
t.Errorf("complete expired lease: err = %v, want ErrLeaseConflict", err)
|
||||
}
|
||||
if err := task.Fail(testWorker, 1, "timeout", "expired", true, expired); !errors.Is(err, ErrLeaseConflict) {
|
||||
t.Errorf("fail expired lease: err = %v, want ErrLeaseConflict", err)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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() }
|
||||
@@ -0,0 +1,193 @@
|
||||
// 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
|
||||
// Local operator UI credential. Empty disables the embedded UI entirely.
|
||||
UIToken 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")),
|
||||
UIToken: os.Getenv("UI_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")
|
||||
}
|
||||
if cfg.UIToken != "" && cfg.Token != "" && cfg.UIToken == cfg.Token {
|
||||
return Config{}, fmt.Errorf("UI_AUTH_TOKEN must differ from the worker auth token")
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
if cfg.DefaultMaxAttempts < 1 {
|
||||
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package infra
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadConfigRejectsSharedUIAndWorkerToken(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("COORDINATOR_TOKEN", "shared-secret")
|
||||
t.Setenv("UI_AUTH_TOKEN", "shared-secret")
|
||||
|
||||
_, err := LoadConfig()
|
||||
if err == nil || !strings.Contains(err.Error(), "must differ") {
|
||||
t.Fatalf("LoadConfig error = %v, want distinct-token error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("COORDINATOR_TOKEN", "worker-secret")
|
||||
t.Setenv("UI_AUTH_TOKEN", "ui-secret")
|
||||
|
||||
cfg, err := LoadConfig()
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig: %v", err)
|
||||
}
|
||||
if cfg.Token != "worker-secret" || cfg.UIToken != "ui-secret" {
|
||||
t.Fatalf("unexpected tokens: %+v", cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadConfigRejectsNonPositiveDefaultMaxAttempts(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("DEFAULT_MAX_ATTEMPTS", "0")
|
||||
|
||||
_, err := LoadConfig()
|
||||
if err == nil || !strings.Contains(err.Error(), "DEFAULT_MAX_ATTEMPTS") {
|
||||
t.Fatalf("LoadConfig error = %v, want default-attempt validation", err)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,374 @@
|
||||
// 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) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time) (int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var cancelled int64
|
||||
for _, task := range r.tasks {
|
||||
if task.JobID == jobID && task.Cancel(now) {
|
||||
cancelled++
|
||||
}
|
||||
}
|
||||
return cancelled, nil
|
||||
}
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
affected := make([]uuid.UUID, 0)
|
||||
for _, t := range r.tasks {
|
||||
if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
|
||||
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
t.ExpireLease(now)
|
||||
affected = append(affected, t.JobID)
|
||||
}
|
||||
}
|
||||
return affected, 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
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, a := range r.arts {
|
||||
if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
|
||||
a.Attempt != nil && *a.Attempt == attempt {
|
||||
return cloneArtifact(a), nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func cloneArtifact(a *domain.Artifact) *domain.Artifact {
|
||||
cp := *a
|
||||
return &cp
|
||||
}
|
||||
|
||||
// --- 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
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package memstore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sort"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
|
||||
// UIReadRepo is the in-memory read projection used by HTTP/UI tests.
|
||||
type UIReadRepo struct {
|
||||
jobs *JobRepo
|
||||
tasks *TaskRepo
|
||||
workers *WorkerRepo
|
||||
artifacts *ArtifactRepo
|
||||
}
|
||||
|
||||
func NewUIReadRepo(j *JobRepo, t *TaskRepo, w *WorkerRepo, a *ArtifactRepo) *UIReadRepo {
|
||||
return &UIReadRepo{j, t, w, a}
|
||||
}
|
||||
|
||||
var _ usecase.UIReadRepository = (*UIReadRepo)(nil)
|
||||
|
||||
func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
return r.jobs.Get(ctx, id)
|
||||
}
|
||||
func (r *UIReadRepo) ListJobs(_ context.Context, limit int) ([]domain.Job, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
r.jobs.mu.Lock()
|
||||
defer r.jobs.mu.Unlock()
|
||||
out := make([]domain.Job, 0, len(r.jobs.jobs))
|
||||
for _, job := range r.jobs.jobs {
|
||||
out = append(out, *job)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
|
||||
return out[i].ID.String() > out[j].ID.String()
|
||||
}
|
||||
return out[i].CreatedAt.After(out[j].CreatedAt)
|
||||
})
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domain.Task, error) {
|
||||
r.tasks.mu.Lock()
|
||||
defer r.tasks.mu.Unlock()
|
||||
out := []domain.Task{}
|
||||
for _, task := range r.tasks.tasks {
|
||||
if task.JobID == jobID {
|
||||
out = append(out, *clone(task))
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
for _, id := range jobIDs {
|
||||
tasks, err := r.ListTasksByJob(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[id] = tasks
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
r.workers.mu.Lock()
|
||||
defer r.workers.mu.Unlock()
|
||||
out := []domain.Worker{}
|
||||
for _, worker := range r.workers.workers {
|
||||
copy := *worker
|
||||
copy.Capabilities = append([]string(nil), worker.Capabilities...)
|
||||
out = append(out, copy)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].LastHeartbeatAt.Equal(out[j].LastHeartbeatAt) {
|
||||
return out[i].ID.String() > out[j].ID.String()
|
||||
}
|
||||
return out[i].LastHeartbeatAt.After(out[j].LastHeartbeatAt)
|
||||
})
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListArtifactsByJob(_ context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
|
||||
r.artifacts.mu.Lock()
|
||||
defer r.artifacts.mu.Unlock()
|
||||
out := []domain.Artifact{}
|
||||
for _, artifact := range r.artifacts.arts {
|
||||
if artifact.JobID == jobID {
|
||||
out = append(out, *artifact)
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
|
||||
return out[i].ID.String() < out[j].ID.String()
|
||||
}
|
||||
return out[i].CreatedAt.Before(out[j].CreatedAt)
|
||||
})
|
||||
return out, nil
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -0,0 +1,102 @@
|
||||
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", "attempt", "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, a.Attempt, 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, &a.Attempt, &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
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).
|
||||
From("artifacts").
|
||||
Where(sq.Eq{
|
||||
"task_id": taskID,
|
||||
"attempt": attempt,
|
||||
"kind": string(domain.ArtifactPartialResult),
|
||||
}).
|
||||
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, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
|
||||
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find partial result: %w", err)
|
||||
}
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
@@ -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)
|
||||
@@ -0,0 +1,588 @@
|
||||
//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. With SKIP LOCKED, a concurrent caller can
|
||||
// transiently see no eligible row while every remaining row is locked by a
|
||||
// different claim statement. Poll briefly, as a real worker does, before
|
||||
// treating the queue as empty. This verifies the actual contract: tasks are
|
||||
// unique and all eventually become claimable without lock contention.
|
||||
for i := 0; i < tasks*2; i++ {
|
||||
wg.Add(1)
|
||||
go func(n int) {
|
||||
defer wg.Done()
|
||||
for attempt := 0; attempt < 20; attempt++ {
|
||||
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 {
|
||||
time.Sleep(time.Millisecond)
|
||||
continue
|
||||
}
|
||||
mu.Lock()
|
||||
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)
|
||||
mu.Unlock()
|
||||
return
|
||||
}
|
||||
}(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 TestCancelJobCancelsEveryUnfinishedTask(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, _ := seedJob(t, pool, 3)
|
||||
clk := fixedClock{now: time.Now().UTC()}
|
||||
uc := usecase.NewCancelJob(NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool), clk)
|
||||
|
||||
cancelled, err := uc.Execute(ctx, job.ID)
|
||||
if err != nil || cancelled != 3 {
|
||||
t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err)
|
||||
}
|
||||
stored, err := NewJobRepo(pool).Get(ctx, job.ID)
|
||||
if err != nil || stored.Status != domain.JobCancelled {
|
||||
t.Fatalf("job after cancel = (%+v, %v)", stored, err)
|
||||
}
|
||||
counts, err := NewTaskRepo(pool).CountByStatus(ctx, job.ID)
|
||||
if err != nil || counts[domain.TaskCancelled] != 3 {
|
||||
t.Fatalf("cancelled tasks = %d, err = %v", counts[domain.TaskCancelled], err)
|
||||
}
|
||||
}
|
||||
|
||||
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, 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)
|
||||
attempt := 1
|
||||
art.Attempt = &attempt
|
||||
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.Attempt = attempt
|
||||
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)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, tasks := seedJob(t, pool, 1)
|
||||
taskID := tasks[0].ID
|
||||
first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
|
||||
second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attempt := 1
|
||||
second.Attempt = &attempt
|
||||
second.SetContent("other-sha", 5)
|
||||
if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
|
||||
t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
|
||||
}
|
||||
}
|
||||
|
||||
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 kind == domain.ArtifactPartialResult {
|
||||
attempt := 1
|
||||
art.Attempt = &attempt
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPartialResultArtifactRoundTripsAttempt(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, tasks := seedJob(t, pool, 1)
|
||||
taskID := tasks[0].ID
|
||||
art, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult,
|
||||
"result.csv", "text/csv", time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attempt := 2
|
||||
art.Attempt = &attempt
|
||||
art.SetContent("sha", 3)
|
||||
repo := NewArtifactRepo(pool)
|
||||
if err := repo.Insert(ctx, art); err != nil {
|
||||
t.Fatalf("insert: %v", err)
|
||||
}
|
||||
got, err := repo.Get(ctx, art.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
if got.Attempt == nil || *got.Attempt != attempt {
|
||||
t.Fatalf("attempt = %v, want %d", got.Attempt, attempt)
|
||||
}
|
||||
}
|
||||
|
||||
// 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])
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
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()
|
||||
}
|
||||
|
||||
// cancelByJobSQL mirrors domain.Task.Cancel in one set-based update. It runs in
|
||||
// the same transaction as the job-status update, so no claimable shard remains
|
||||
// after an operator receives a successful cancellation response.
|
||||
const cancelByJobSQL = `
|
||||
UPDATE tasks
|
||||
SET status = 'cancelled'::task_status,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
error_code = NULL,
|
||||
error_message = NULL,
|
||||
completed_at = $2,
|
||||
version = version + 1
|
||||
WHERE job_id = $1
|
||||
AND status IN ('pending','leased','running')`
|
||||
|
||||
func (r *TaskRepo) CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) {
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, cancelByJobSQL, jobID, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return tag.RowsAffected(), nil
|
||||
}
|
||||
|
||||
// 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
|
||||
RETURNING job_id`
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
var affected []uuid.UUID
|
||||
err := withRetry(ctx, func(ctx context.Context) error {
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
affected = affected[:0]
|
||||
for rows.Next() {
|
||||
var jobID uuid.UUID
|
||||
if err := rows.Scan(&jobID); err != nil {
|
||||
return err
|
||||
}
|
||||
affected = append(affected, jobID)
|
||||
}
|
||||
return rows.Err()
|
||||
})
|
||||
return affected, err
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
sq "github.com/Masterminds/squirrel"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
|
||||
// UIReadRepo contains bounded, deterministic read queries for the operator UI.
|
||||
type UIReadRepo struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewUIReadRepo(pool *pgxpool.Pool) *UIReadRepo { return &UIReadRepo{pool: pool} }
|
||||
|
||||
var _ usecase.UIReadRepository = (*UIReadRepo)(nil)
|
||||
|
||||
func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
job, err := NewJobRepo(r.pool).Get(ctx, id)
|
||||
return job, err
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListJobs(ctx context.Context, limit int) ([]domain.Job, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
sql, args, err := psql.Select(jobColumns...).From("jobs").OrderBy("created_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
jobs := make([]domain.Job, 0)
|
||||
for rows.Next() {
|
||||
var j domain.Job
|
||||
var status string
|
||||
var inputURI *string
|
||||
if err := rows.Scan(&j.ID, &j.Workload, &inputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if inputURI != nil {
|
||||
j.InputURI = *inputURI
|
||||
}
|
||||
j.Status = domain.JobStatus(status)
|
||||
jobs = append(jobs, j)
|
||||
}
|
||||
return jobs, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error) {
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").Where(sq.Eq{"job_id": jobID}).OrderBy("chunk_index ASC").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tasks: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
tasks := make([]domain.Task, 0)
|
||||
for rows.Next() {
|
||||
task, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tasks = append(tasks, *task)
|
||||
}
|
||||
return tasks, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
if len(jobIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").
|
||||
Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tasks by jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
task, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[task.JobID] = append(out[task.JobID], *task)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
sql, args, err := psql.Select(workerColumns...).From("workers").OrderBy("last_heartbeat_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list workers: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
workers := make([]domain.Worker, 0)
|
||||
for rows.Next() {
|
||||
worker, err := scanWorker(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
workers = append(workers, *worker)
|
||||
}
|
||||
return workers, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).From("artifacts").Where(sq.Eq{"job_id": jobID}).OrderBy("created_at ASC", "id ASC").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list artifacts: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
artifacts := make([]domain.Artifact, 0)
|
||||
for rows.Next() {
|
||||
var a domain.Artifact
|
||||
var kind string
|
||||
if err := rows.Scan(&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey, &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
artifacts = append(artifacts, a)
|
||||
}
|
||||
return artifacts, rows.Err()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
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"`
|
||||
Cancelled int `json:"cancelled"`
|
||||
}
|
||||
|
||||
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,
|
||||
Cancelled: p.Cancelled,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"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()
|
||||
if err := dec.Decode(dst); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := dec.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
|
||||
return errors.New("request body must contain exactly one JSON value")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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),
|
||||
errors.Is(err, domain.ErrJobNotCancellable):
|
||||
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})
|
||||
}
|
||||
@@ -0,0 +1,436 @@
|
||||
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
|
||||
}
|
||||
if _, err := uuid.Parse(req.WorkerID); 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/max_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
|
||||
maxRows int
|
||||
result usecase.SubmitDatasetResult
|
||||
gotDataset bool
|
||||
gotWorkload bool
|
||||
gotParams bool
|
||||
gotRows bool
|
||||
gotMaxRows 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":
|
||||
if gotDataset || gotWorkload {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(part, 1<<10))
|
||||
workload = strings.TrimSpace(string(b))
|
||||
gotWorkload = true
|
||||
case "parameters":
|
||||
if gotDataset || gotParams {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(part, 1<<16))
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, ¶ms); err != nil {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
}
|
||||
gotParams = true
|
||||
case "chunk_rows":
|
||||
if gotDataset || gotRows {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(part, 32))
|
||||
n, err := strconv.Atoi(strings.TrimSpace(string(b)))
|
||||
if err != nil || n < 1 {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
rows = n
|
||||
gotRows = true
|
||||
case "max_rows":
|
||||
if gotDataset || gotMaxRows {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
b, _ := io.ReadAll(io.LimitReader(part, 32))
|
||||
n, err := strconv.Atoi(strings.TrimSpace(string(b)))
|
||||
if err != nil || n < 1 {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
maxRows = n
|
||||
gotMaxRows = true
|
||||
case "file", "dataset":
|
||||
if gotDataset || workload == "" {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
filename := part.FileName()
|
||||
if filename == "" {
|
||||
filename = "dataset"
|
||||
}
|
||||
result, err = s.uc.SubmitDataset.Execute(r.Context(), usecase.SubmitDatasetInput{
|
||||
Workload: workload,
|
||||
Parameters: params,
|
||||
RowsPerShard: rows,
|
||||
MaxRows: maxRows,
|
||||
Filename: filename,
|
||||
ContentType: part.Header.Get("Content-Type"),
|
||||
Body: part,
|
||||
})
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
gotDataset = true
|
||||
default:
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
_ = 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))
|
||||
}
|
||||
|
||||
// handleCancelJob stops all non-terminal shards for an operator-requested job.
|
||||
// It is available to both the bearer API and the separately authenticated UI.
|
||||
func (s *Server) handleCancelJob(w http.ResponseWriter, r *http.Request) {
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
jobID, ok := s.pathUUID(w, r, "job_id")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
cancelled, err := s.uc.CancelJob.Execute(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"job_id": jobID,
|
||||
"status": domain.JobCancelled,
|
||||
"cancelled_tasks": cancelled,
|
||||
})
|
||||
}
|
||||
|
||||
// --- 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
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// withBasicAuth protects the local operator UI with a credential distinct from
|
||||
// the worker bearer token. The username is intentionally ignored; the password
|
||||
// is the configured UI token. Basic Auth is suitable only for localhost or a
|
||||
// TLS-terminating trusted reverse proxy.
|
||||
func withBasicAuth(token string) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, password, ok := r.BasicAuth()
|
||||
if !ok || subtle.ConstantTimeCompare([]byte(password), []byte(token)) != 1 {
|
||||
w.Header().Set("WWW-Authenticate", `Basic realm="SciMesh UI", charset="UTF-8"`)
|
||||
writeJSON(w, http.StatusUnauthorized, errorResponse{Error: "unauthorized", RequestID: requestIDFrom(r.Context())})
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// withSameOrigin rejects browser form/fetch writes initiated by another origin.
|
||||
// A missing Origin is allowed for direct local tools; authenticated UI pages use
|
||||
// the browser-supplied Origin header on state-changing requests.
|
||||
func withSameOrigin(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodOptions {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
origin := r.Header.Get("Origin")
|
||||
if origin != "" {
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
if origin != scheme+"://"+r.Host {
|
||||
writeJSON(w, http.StatusForbidden, errorResponse{Error: "cross-origin request rejected", 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
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
// 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
|
||||
CancelJob *usecase.CancelJob
|
||||
UploadArtifact *usecase.UploadArtifact
|
||||
DownloadArtifact *usecase.DownloadArtifact
|
||||
GetTaskInput *usecase.GetTaskInput
|
||||
Dashboard *usecase.Dashboard
|
||||
}
|
||||
|
||||
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, uiToken ...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 /jobs/{job_id}/cancel", s.handleCancelJob)
|
||||
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)
|
||||
if len(uiToken) > 0 && uiToken[0] != "" && s.uc.Dashboard != nil {
|
||||
ui := http.NewServeMux()
|
||||
ui.HandleFunc("GET /ui", s.handleUIHome)
|
||||
ui.HandleFunc("GET /ui/jobs/new", s.handleUINewJob)
|
||||
ui.HandleFunc("GET /ui/jobs/{job_id}", s.handleUIJob)
|
||||
ui.HandleFunc("GET /ui/api/jobs/{job_id}", s.handleUIJobJSON)
|
||||
ui.HandleFunc("POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob)
|
||||
ui.HandleFunc("POST /ui/api/jobs/upload", s.handleUploadDataset)
|
||||
ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload)
|
||||
mux.Handle("/ui", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
|
||||
mux.Handle("/ui/", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
|
||||
} else {
|
||||
// More specific than the protected catch-all: UI absence is not an auth
|
||||
// failure and does not disclose that a UI feature is configured elsewhere.
|
||||
mux.HandleFunc("/ui", http.NotFound)
|
||||
mux.HandleFunc("/ui/", http.NotFound)
|
||||
}
|
||||
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"})
|
||||
}
|
||||
@@ -0,0 +1,605 @@
|
||||
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"
|
||||
const uiToken = "ui-secret"
|
||||
|
||||
type env struct {
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
workerID string
|
||||
}
|
||||
|
||||
func newEnv(t *testing.T, ready func(context.Context) error) *env {
|
||||
return newEnvWithUIToken(t, ready, uiToken)
|
||||
}
|
||||
|
||||
func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configuredUIToken string) *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, 3),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
|
||||
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
|
||||
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
|
||||
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
|
||||
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
|
||||
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
|
||||
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
|
||||
}
|
||||
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
|
||||
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("register test worker: %v", err)
|
||||
}
|
||||
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, configuredUIToken))
|
||||
t.Cleanup(ts.Close)
|
||||
return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()}
|
||||
}
|
||||
|
||||
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()
|
||||
body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`)
|
||||
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 TestUIRequiresDistinctCredentialAndRendersDashboard(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
request := func() *http.Request {
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui", nil)
|
||||
return req
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(request())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusUnauthorized {
|
||||
t.Fatalf("no UI auth: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
req := request()
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err = http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusUnauthorized {
|
||||
t.Fatalf("worker token authorized UI: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
req = request()
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err = http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("UI status: %d", resp.StatusCode)
|
||||
}
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if !strings.Contains(string(body), "SciMesh operator dashboard") {
|
||||
t.Errorf("dashboard body missing title")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIDisabledReturnsNotFound(t *testing.T) {
|
||||
e := newEnvWithUIToken(t, healthy, "")
|
||||
resp := e.get(t, "/ui")
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusNotFound {
|
||||
t.Fatalf("disabled UI = %d, want 404", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIRejectsCrossOriginUpload(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/upload", strings.NewReader("dataset=x"))
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Origin", "https://attacker.example")
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusForbidden {
|
||||
t.Fatalf("cross-origin upload = %d, want 403", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIUploadDatasetCreatesJob(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var body bytes.Buffer
|
||||
mw := multipart.NewWriter(&body)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO","top_k":20,"progress_every":0}`)
|
||||
_ = mw.WriteField("chunk_rows", "1000")
|
||||
file, err := mw.CreateFormFile("file", "chembl.tsv")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _ = io.WriteString(file, "chembl_id\tcanonical_smiles\nCHEMBL1\tCCO\n")
|
||||
if err := mw.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/upload", &body)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusCreated {
|
||||
result, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("UI upload = %d: %s", resp.StatusCode, result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelJobStopsUnfinishedTasks(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, job := 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"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("create: %d", code)
|
||||
}
|
||||
jobID := job["id"].(string)
|
||||
if code, body := e.do(t, "POST", "/jobs/"+jobID+"/cancel", ""); code != http.StatusOK || body["cancelled_tasks"].(float64) != 2 {
|
||||
t.Fatalf("cancel = (%d, %v)", code, body)
|
||||
}
|
||||
if code, progress := e.do(t, "GET", "/jobs/"+jobID, ""); code != http.StatusOK || progress["status"] != "cancelled" || progress["cancelled"].(float64) != 2 {
|
||||
t.Fatalf("cancelled job progress = (%d, %v)", code, progress)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUICancelJobUsesOperatorCredential(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
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 != http.StatusCreated {
|
||||
t.Fatalf("create: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/"+job["id"].(string)+"/cancel", nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("UI cancel = %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIJobAndArtifactAreScopedToTheirJob(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("create: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+job["id"].(string), nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("detail: %d", resp.StatusCode)
|
||||
}
|
||||
if got := resp.Header.Get("Content-Security-Policy"); got == "" {
|
||||
t.Error("missing UI CSP")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIArtifactDownloadRejectsAnotherJobsArtifact(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("first job: %d", code)
|
||||
}
|
||||
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
|
||||
artifactID := e.putArtifact(t, claim["task_id"].(string), "w1", int(claim["attempt"].(float64)), "result")
|
||||
code, second := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("second job: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+second["id"].(string)+"/artifacts/"+artifactID, nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusNotFound {
|
||||
t.Errorf("cross-job artifact = %d, want 404", 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 TestClaimRequiresRegisteredWorkerAndUsesStoredCapabilities(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"not-a-uuid"}`); code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid worker id claim = %d, want 400", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"11111111-1111-4111-8111-111111111111"}`); code != http.StatusNotFound {
|
||||
t.Fatalf("unregistered worker claim = %d, want 404", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`); code != http.StatusCreated {
|
||||
t.Fatalf("create job = %d", code)
|
||||
}
|
||||
code, worker := e.do(t, "POST", "/workers/register", `{"name":"search-only","capabilities":["similarity-search"]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("register = %d", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"`+worker["worker_id"].(string)+`","capabilities":["w"]}`); code != http.StatusNoContent {
|
||||
t.Fatalf("forged capability claim = %d, want 204", 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 := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
code, body := e.uploadDataset(t, "similarity-search", 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), "chembl_id\tcanonical_smiles\n") {
|
||||
t.Errorf("shard missing header: %q", shard)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetLimitsRows(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
_ = mw.WriteField("max_rows", "3")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
|
||||
_ = 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 result map[string]any
|
||||
_ = json.NewDecoder(resp.Body).Decode(&result)
|
||||
if resp.StatusCode != http.StatusCreated || result["task_count"].(float64) != 2 {
|
||||
t.Fatalf("limited upload = (%d, %v)", resp.StatusCode, result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_ = 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()
|
||||
if resp.StatusCode != http.StatusBadRequest {
|
||||
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
func TestJSONRejectsTrailingValue(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
if code, _ := e.do(t, "POST", "/workers/register",
|
||||
`{"name":"lab","capabilities":["w"]} {}`); code != http.StatusBadRequest {
|
||||
t.Errorf("status = %d, want 400", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "not-a-number")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"))
|
||||
_ = 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()
|
||||
if resp.StatusCode != http.StatusBadRequest {
|
||||
t.Errorf("status = %d, want 400", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
// --- 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")
|
||||
if worker == "w1" {
|
||||
worker = e.workerID
|
||||
}
|
||||
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("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = 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) }
|
||||
@@ -0,0 +1,23 @@
|
||||
{{define "dashboard.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>SciMesh operator dashboard</title>
|
||||
<style>
|
||||
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:1180px;margin:auto;padding:32px 20px 56px}.top{display:flex;justify-content:space-between;gap:24px;align-items:start}.eyebrow{margin:0;color:#50617d;font-size:.86rem;font-weight:700;text-transform:uppercase;letter-spacing:.08em}h1{margin:.2rem 0;font-size:2rem}h2{margin:32px 0 12px;font-size:1.28rem}.lead{margin:0;color:#56657c}.button{display:inline-block;border:0;border-radius:8px;padding:11px 15px;background:#1f5eff;color:#fff;font-weight:700;text-decoration:none;white-space:nowrap}.notice{margin-top:24px;padding:16px 18px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.notice strong{display:block}.steps{display:grid;grid-template-columns:repeat(3,1fr);gap:12px;margin-top:14px}.step,.card{padding:16px;border:1px solid #dfe5f0;border-radius:10px;background:#fff}.step b{display:block;color:#1f5eff}.table-wrap{overflow-x:auto;background:#fff;border:1px solid #dfe5f0;border-radius:10px}table{width:100%;border-collapse:collapse}td,th{padding:13px 14px;border-bottom:1px solid #e8ecf4;text-align:left;vertical-align:top}th{color:#50617d;font-size:.78rem;text-transform:uppercase;letter-spacing:.06em}tr:last-child td{border:0}a{color:#174ecf}small,.muted{color:#68758b}.status{display:inline-block;border-radius:999px;padding:3px 9px;font-size:.84rem;font-weight:700}.status-success{background:#dff6e9;color:#126b3d}.status-danger{background:#ffe4e6;color:#a31135}.status-active{background:#e4edff;color:#174ecf}.status-waiting{background:#edf0f5;color:#50617d}.bar{height:7px;min-width:120px;margin-top:7px;overflow:hidden;border-radius:999px;background:#e6eaf1}.bar>span{display:block;height:100%;background:#1f5eff}.kicker{font-variant-numeric:tabular-nums}.empty{padding:28px;text-align:center;color:#68758b}.worker{display:grid;grid-template-columns:1.3fr .8fr 2fr 1fr;gap:12px;align-items:center}.worker+.worker{border-top:1px solid #e8ecf4;padding-top:12px;margin-top:12px}@media(max-width:760px){.top,.steps{display:block}.button{margin-top:12px}.step{margin-top:10px}.worker{grid-template-columns:1fr}.hide-mobile{display:none}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main class="page">
|
||||
<header class="top"><div><p class="eyebrow">Local coordinator</p><h1>SciMesh operator dashboard</h1><p class="lead">See where a computation is and what should happen next.</p></div><a class="button" href="/ui/jobs/new">Start a check</a></header>
|
||||
<section class="notice" aria-label="Current pipeline limitation"><strong>This screen currently diagnoses shard jobs.</strong><span>Workers upload partial CSVs to the coordinator. Until a reducer is implemented, those files are not one final scientific result.</span><div class="steps"><div class="step"><b>1. Upload TSV</b>The coordinator splits the file into shard tasks.</div><div class="step"><b>2. Wait for a worker</b>A worker claims a shard, calculates similarity, and returns a CSV.</div><div class="step"><b>3. Inspect artifacts</b>Download a partial result from the job page.</div></div></section>
|
||||
<h2>Recent jobs</h2>
|
||||
<div class="table-wrap"><table><tr><th>Computation</th><th>State</th><th>Progress</th><th class="hide-mobile">Created</th></tr>{{range .Jobs}}<tr><td><a href="/ui/jobs/{{.ID}}"><strong>{{workloadLabel .Workload}}</strong></a><br><small>Open job details</small></td><td><span class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span><br><small>{{statusHint .Status}}</small></td><td class="kicker"><strong>{{.Completed}} / {{.Total}}</strong> complete{{if gt .Failed 0}} · <span style="color:#a31135">failed: {{.Failed}}</span>{{end}}{{if gt .Cancelled 0}} · <span>stopped: {{.Cancelled}}</span>{{end}}<div class="bar"><span style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div></td><td class="hide-mobile"><small>{{time .CreatedAt}}</small></td></tr>{{else}}<tr><td colspan="4" class="empty"><strong>No jobs yet.</strong><br>Click “Start a check”, upload a small TSV, and leave a worker running.</td></tr>{{end}}</table></div>
|
||||
<h2>Workers</h2>
|
||||
<section class="card">{{range .Workers}}<div class="worker"><div><strong>{{.Name}}</strong><br><small>{{.ID}}</small></div><div><span class="status status-{{if eq .Status "online"}}success{{else}}waiting{{end}}">{{workerStatusLabel .Status}}</span></div><div>{{range .Capabilities}}<code>{{.}}</code> {{end}}</div><div class="muted">Last signal<br>{{time .LastHeartbeatAt}}</div></div>{{else}}<div class="empty"><strong>No worker is registered yet.</strong><br>Run <code>scimesh-worker</code> with the coordinator URL and worker token.</div>{{end}}</section>
|
||||
</main>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
@@ -0,0 +1,29 @@
|
||||
{{define "job.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>SciMesh job</title>
|
||||
<style>
|
||||
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:1180px;margin:auto;padding:32px 20px 56px}a{color:#174ecf}.back{text-decoration:none}.eyebrow{margin:28px 0 4px;color:#50617d;font-size:.86rem;font-weight:700;text-transform:uppercase;letter-spacing:.08em}h1{margin:0;font-size:2rem}h2{margin:32px 0 12px;font-size:1.3rem}.summary,.card{padding:20px;border:1px solid #dfe5f0;border-radius:12px;background:#fff}.summary-head{display:flex;justify-content:space-between;gap:16px;align-items:start}.status{display:inline-block;border-radius:999px;padding:4px 10px;font-size:.9rem;font-weight:700}.status-success{background:#dff6e9;color:#126b3d}.status-danger{background:#ffe4e6;color:#a31135}.status-active{background:#e4edff;color:#174ecf}.status-waiting{background:#edf0f5;color:#50617d}.hint{margin:8px 0 0;color:#56657c}.bar{height:10px;margin:20px 0 8px;overflow:hidden;border-radius:999px;background:#e6eaf1}.bar>span{display:block;height:100%;background:#1f5eff;transition:width .3s}.numbers{display:grid;grid-template-columns:repeat(6,1fr);gap:10px}.number{padding:12px;border-radius:8px;background:#f6f8fc}.number b{display:block;font-size:1.35rem}.stop{display:block;margin-left:auto;border:1px solid #d43b51;border-radius:7px;padding:8px 11px;background:#fff;color:#b2223a;font:inherit;font-weight:700;cursor:pointer}.stop:disabled{opacity:.6}.notice{margin:20px 0;padding:15px 17px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.table-wrap{overflow-x:auto;border:1px solid #dfe5f0;border-radius:10px;background:#fff}table{width:100%;border-collapse:collapse}td,th{padding:12px 13px;border-bottom:1px solid #e8ecf4;text-align:left;vertical-align:top}th{color:#50617d;font-size:.78rem;text-transform:uppercase;letter-spacing:.06em}tr:last-child td{border:0}small,.muted{color:#68758b}.error{color:#a31135;max-width:360px;word-break:break-word}.download{display:inline-block;padding:7px 10px;border:1px solid #b9c9ee;border-radius:7px;text-decoration:none}.empty{padding:24px;text-align:center;color:#68758b}details{margin-top:18px;color:#56657c}code{word-break:break-all}@media(max-width:700px){.summary-head{display:block}.numbers{grid-template-columns:repeat(2,1fr)}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main class="page">
|
||||
<a class="back" href="/ui">← Back to jobs</a><p class="eyebrow">{{workloadLabel .Workload}}</p><h1>Execution progress</h1>
|
||||
<section class="summary"><div class="summary-head"><div><span id="status" class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span><p id="hint" class="hint">{{statusHint .Status}}</p></div><div>{{if cancellable .Status}}<button id="stop-job" class="stop" type="button">Stop job</button><small>This cancels every shard that is not finished yet.</small>{{else}}<small>Summary refreshes automatically every two seconds.</small>{{end}}</div></div><div class="bar" aria-label="Progress"><span id="progress-bar" style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div><p id="progress" class="muted">{{.Completed}} of {{.Total}} tasks complete</p><div class="numbers"><div class="number"><b id="total">{{.Total}}</b><small>total shards</small></div><div class="number"><b id="completed">{{.Completed}}</b><small>complete</small></div><div class="number"><b id="pending">{{.Pending}}</b><small>waiting</small></div><div class="number"><b id="active">{{add .Leased .Running}}</b><small>with workers</small></div><div class="number"><b id="failed">{{.Failed}}</b><small>failed</small></div><div class="number"><b id="cancelled">{{.Cancelled}}</b><small>stopped</small></div></div><details><summary>Technical details</summary><p>Job ID: <code>{{.ID}}</code><br>Workload: <code>{{.Workload}}</code><br>Created: {{time .CreatedAt}}</p></details></section>
|
||||
<section class="notice"><strong>What can be downloaded now?</strong><br><code>partial_result</code> files come from individual shards. They are useful for checking the pipeline, but are not a merged final CSV because the reducer is not implemented yet.</section>
|
||||
<h2>Shard tasks</h2><p class="muted">If a task fails, its code and message appear here. Refresh the page to update the detailed rows.</p>
|
||||
<div class="table-wrap"><table><tr><th>Shard</th><th>State</th><th>Attempt</th><th>Worker / lease</th><th>Error</th></tr>{{range .Tasks}}<tr><td>#{{.ChunkIndex}}</td><td><span class="status status-{{statusClass .Status}}">{{statusLabel .Status}}</span></td><td>{{.Attempt}} / {{.MaxAttempts}}</td><td>{{if .LeaseOwner}}<code>{{.LeaseOwner}}</code>{{if .LeaseExpiresAt}}<br><small>until {{time .LeaseExpiresAt}}</small>{{end}}{{else}}<span class="muted">—</span>{{end}}</td><td class="error">{{if .ErrorCode}}<strong>{{taskErrorLabel .ErrorCode}}</strong><br><small>{{taskErrorHint .ErrorCode}}</small>{{else}}<span class="muted">—</span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No tasks have appeared yet.</td></tr>{{end}}</table></div>
|
||||
<h2>Coordinator artifacts</h2>
|
||||
<div class="table-wrap"><table><tr><th>Type</th><th>File</th><th>Size</th><th>Integrity check</th><th></th></tr>{{range .Artifacts}}<tr><td>{{if .Diagnostic}}<strong>Partial result</strong><br><small>diagnostic</small>{{else}}{{.Kind}}{{end}}</td><td>{{.Filename}}</td><td>{{bytes .SizeBytes}}</td><td><code>{{.SHA256}}</code></td><td>{{if .Downloadable}}<a class="download" href="/ui/jobs/{{$.ID}}/artifacts/{{.ID}}">Download CSV</a>{{else}}<span class="muted">Unavailable</span>{{end}}</td></tr>{{else}}<tr><td colspan="5" class="empty">No artifacts yet. The worker uploads a CSV after it completes a shard.</td></tr>{{end}}</table></div>
|
||||
</main>
|
||||
<script>
|
||||
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']};
|
||||
const stop=document.querySelector('#stop-job');if(stop)stop.addEventListener('click',async()=>{if(!confirm('Stop this job? Unfinished shards will be cancelled.'))return;stop.disabled=true;const response=await fetch('/ui/api/jobs/'+id+'/cancel',{method:'POST'});if(!response.ok){stop.disabled=false;alert('Unable to stop this job.');return}location.reload()});
|
||||
const terminal=new Set(['completed','failed','cancelled']);let timer;const poll=async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running;if(terminal.has(job.status)&&timer){clearInterval(timer);timer=undefined}}catch(_){}};const start=()=>{if(!timer&&!document.hidden&&!terminal.has(document.querySelector('#status').textContent.toLowerCase()))timer=setInterval(poll,2000)};document.addEventListener('visibilitychange',()=>{if(document.hidden&&timer){clearInterval(timer);timer=undefined}else start()});start();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
@@ -0,0 +1,31 @@
|
||||
{{define "new-job.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Create a check — SciMesh</title>
|
||||
<style>
|
||||
:root{color:#172033;background:#f6f8fc;font:16px/1.5 system-ui,sans-serif}body{margin:0}.page{max-width:760px;margin:auto;padding:32px 20px 56px}a{color:#174ecf}.back{text-decoration:none}.eyebrow{margin:28px 0 4px;color:#50617d;font-size:.86rem;font-weight:700;text-transform:uppercase;letter-spacing:.08em}h1{margin:0;font-size:2rem}.lead{color:#56657c}.notice{margin:22px 0;padding:16px 18px;border:1px solid #f2cb72;border-radius:10px;background:#fff8e6}.notice strong{display:block}.card{padding:22px;border:1px solid #dfe5f0;border-radius:12px;background:#fff}label{display:block;margin:18px 0 4px;font-weight:700}input{box-sizing:border-box;width:100%;padding:10px;border:1px solid #bac5d8;border-radius:7px;font:inherit}input[type=file]{padding:8px;background:#f8faff}.hint{margin:4px 0;color:#68758b;font-size:.9rem}.button{margin-top:22px;border:0;border-radius:8px;padding:11px 16px;background:#1f5eff;color:#fff;font:inherit;font-weight:700;cursor:pointer}.button:disabled{opacity:.6;cursor:wait}.error{margin-top:16px;color:#a31135}.working{margin-top:16px;color:#174ecf}.checklist{margin:8px 0;padding-left:20px;color:#56657c}.checklist li{margin:5px 0}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main class="page">
|
||||
<a class="back" href="/ui">← Back to jobs</a><p class="eyebrow">Guided run</p><h1>Search for similar molecules</h1><p class="lead">Creates a diagnostic <code>similarity-search</code> job: a worker finds the top-k molecules most similar to a target SMILES.</p>
|
||||
<section class="notice"><strong>Before starting</strong><ul class="checklist"><li>Keep at least one <code>scimesh-worker</code> running.</li><li>Use a small TSV for a hands-on check.</li><li><b>“Rows per shard” does not limit the file size.</b> It splits the entire upload into tasks: a full ChEMBL TSV at 1,000 rows per shard creates thousands of tasks.</li></ul></section>
|
||||
<form id="run" class="card">
|
||||
<label for="file">ChEMBL TSV</label><input id="file" type="file" name="file" required accept=".tsv,.txt,text/tab-separated-values"><p class="hint">Expected columns: <code>chembl_id</code> and <code>canonical_smiles</code>.</p>
|
||||
<label for="query-smiles">Target molecule (SMILES)</label><input id="query-smiles" name="query_smiles" required maxlength="200" value="CCO" autocomplete="off"><p class="hint"><code>CCO</code> is ethanol. For gefitinib, use its SMILES here or the local CLI with <code>--query-id</code>.</p>
|
||||
<label for="top-k">Matches to return</label><input id="top-k" name="top_k" type="number" min="1" max="100000" value="20" required><p class="hint">This is the top-k within each shard, not a global top-k for the whole dataset yet.</p>
|
||||
<label for="chunk-rows">Rows per shard</label><input id="chunk-rows" name="chunk_rows" type="number" min="1" max="100000" value="1000" required><p class="hint">Fewer rows mean more tasks and more visible progress; more rows mean fewer, longer tasks.</p>
|
||||
<label for="max-rows">Maximum dataset rows to process <small>(optional)</small></label><input id="max-rows" name="max_rows" type="number" min="1" max="10000000" placeholder="For example: 500"><p class="hint">Useful for a quick check of a large TSV. The coordinator creates shards from only the first N data rows; it still stores the original upload.</p>
|
||||
<button class="button" id="submit" type="submit">Upload file and create job</button><p id="working" class="working" hidden aria-live="polite">Uploading the file and creating shard tasks… Keep this page open.</p><p id="error" class="error" role="alert"></p>
|
||||
</form>
|
||||
</main>
|
||||
<script>
|
||||
const form=document.querySelector('#run'),button=document.querySelector('#submit'),working=document.querySelector('#working'),error=document.querySelector('#error');
|
||||
form.addEventListener('submit',async event=>{event.preventDefault();error.textContent='';const fields=new FormData(form),file=fields.get('file'),maxRows=String(fields.get('max_rows')||'').trim();if(!(file instanceof File)||file.size===0){error.textContent='Choose a non-empty TSV file.';return}const parameters={query_smiles:fields.get('query_smiles'),top_k:Number(fields.get('top_k')),progress_every:0},upload=new FormData();upload.append('workload','similarity-search');upload.append('parameters',JSON.stringify(parameters));upload.append('chunk_rows',fields.get('chunk_rows'));if(maxRows)upload.append('max_rows',maxRows);upload.append('file',file,file.name);button.disabled=true;working.hidden=false;try{const response=await fetch('/ui/api/jobs/upload',{method:'POST',body:upload}),data=await response.json();if(!response.ok)throw Error(data.error||'Unable to create the job.');location.href='/ui/jobs/'+data.job_id}catch(err){error.textContent=err.message==='invalid input'?'Check the TSV and fields: the coordinator could not accept this request.':err.message;button.disabled=false;working.hidden=true}});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
@@ -0,0 +1,281 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"embed"
|
||||
"fmt"
|
||||
"html/template"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
//go:embed templates/*.html
|
||||
var uiFiles embed.FS
|
||||
|
||||
var uiTemplates = template.Must(template.New("ui").Funcs(template.FuncMap{
|
||||
"time": formatUITime,
|
||||
"statusLabel": uiStatusLabel,
|
||||
"statusHint": uiStatusHint,
|
||||
"statusClass": uiStatusClass,
|
||||
"taskErrorLabel": uiTaskErrorLabel,
|
||||
"taskErrorHint": uiTaskErrorHint,
|
||||
"workerStatusLabel": uiWorkerStatusLabel,
|
||||
"workloadLabel": uiWorkloadLabel,
|
||||
"progressPercent": uiProgressPercent,
|
||||
"cancellable": uiCancellable,
|
||||
"bytes": uiBytes,
|
||||
"add": func(a, b int) int { return a + b },
|
||||
}).ParseFS(uiFiles, "templates/*.html"))
|
||||
|
||||
func formatUITime(t time.Time) string {
|
||||
if t.IsZero() {
|
||||
return "—"
|
||||
}
|
||||
return t.UTC().Format("02.01.2006 15:04 UTC")
|
||||
}
|
||||
|
||||
func uiStatusLabel(status string) string {
|
||||
switch status {
|
||||
case "pending":
|
||||
return "Waiting for a worker"
|
||||
case "leased":
|
||||
return "Assigned to a worker"
|
||||
case "running":
|
||||
return "Running"
|
||||
case "completed":
|
||||
return "Tasks complete"
|
||||
case "failed":
|
||||
return "Needs attention"
|
||||
case "cancelled":
|
||||
return "Stopped"
|
||||
default:
|
||||
return status
|
||||
}
|
||||
}
|
||||
|
||||
func uiStatusHint(status string) string {
|
||||
switch status {
|
||||
case "pending":
|
||||
return "Waiting for an available worker with the required capability."
|
||||
case "leased":
|
||||
return "A worker has claimed the task and should begin processing shortly."
|
||||
case "running":
|
||||
return "A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator."
|
||||
case "completed":
|
||||
return "Every shard task is complete. Files below are still partial results."
|
||||
case "failed":
|
||||
return "One or more shard tasks failed. Open the task list below for details."
|
||||
case "cancelled":
|
||||
return "The operator stopped this job. No new shards can be claimed."
|
||||
default:
|
||||
return "Status reported by the coordinator."
|
||||
}
|
||||
}
|
||||
|
||||
func uiStatusClass(status string) string {
|
||||
switch status {
|
||||
case "completed":
|
||||
return "success"
|
||||
case "failed":
|
||||
return "danger"
|
||||
case "cancelled":
|
||||
return "waiting"
|
||||
case "running", "leased":
|
||||
return "active"
|
||||
default:
|
||||
return "waiting"
|
||||
}
|
||||
}
|
||||
|
||||
func uiWorkerStatusLabel(status string) string {
|
||||
switch status {
|
||||
case "online":
|
||||
return "Available"
|
||||
case "offline":
|
||||
return "Offline"
|
||||
default:
|
||||
return status
|
||||
}
|
||||
}
|
||||
|
||||
// uiTaskErrorLabel deliberately maps worker implementation errors to an
|
||||
// operator-facing diagnosis. Raw subprocess commands and local paths belong in
|
||||
// the worker terminal, not in the web UI.
|
||||
func uiTaskErrorLabel(errorCode string) string {
|
||||
switch errorCode {
|
||||
case "CalledProcessError":
|
||||
return "Local calculation failed"
|
||||
case "ValueError":
|
||||
return "Task input could not be processed"
|
||||
case "CoordinatorTransientError":
|
||||
return "Coordinator connection was interrupted"
|
||||
case "CoordinatorConflictError":
|
||||
return "Worker lease was no longer valid"
|
||||
case "FileNotFoundError":
|
||||
return "Local task file is missing"
|
||||
default:
|
||||
return errorCode
|
||||
}
|
||||
}
|
||||
|
||||
func uiTaskErrorHint(errorCode string) string {
|
||||
switch errorCode {
|
||||
case "CalledProcessError":
|
||||
return "The local SciMesh command stopped before it could upload a result. Check the worker terminal for the original error."
|
||||
case "ValueError":
|
||||
return "The coordinator task or its downloaded input did not meet the worker validation rules."
|
||||
case "CoordinatorTransientError":
|
||||
return "The worker will retry after the coordinator connection is available again."
|
||||
case "CoordinatorConflictError":
|
||||
return "Another worker or a lease timeout changed this task before completion."
|
||||
case "FileNotFoundError":
|
||||
return "The worker could not find one of its local task files. Restart it with an absolute --work-dir."
|
||||
default:
|
||||
return "Check the worker terminal for the original error details."
|
||||
}
|
||||
}
|
||||
|
||||
func uiWorkloadLabel(workload string) string {
|
||||
switch workload {
|
||||
case "similarity-search", "similarity_search":
|
||||
return "Molecule similarity search"
|
||||
case "similarity-graph", "similarity_graph":
|
||||
return "Molecular similarity graph"
|
||||
default:
|
||||
return workload
|
||||
}
|
||||
}
|
||||
|
||||
func uiCancellable(status string) bool {
|
||||
return status == "pending" || status == "running"
|
||||
}
|
||||
|
||||
func uiProgressPercent(completed, failed, cancelled, total int) int {
|
||||
if total <= 0 {
|
||||
return 0
|
||||
}
|
||||
percent := (completed + failed + cancelled) * 100 / total
|
||||
if percent > 100 {
|
||||
return 100
|
||||
}
|
||||
return percent
|
||||
}
|
||||
|
||||
func uiBytes(n int64) string {
|
||||
const kib = 1024
|
||||
if n < kib {
|
||||
return fmt.Sprintf("%d B", n)
|
||||
}
|
||||
if n < kib*kib {
|
||||
return fmt.Sprintf("%.1f KiB", float64(n)/kib)
|
||||
}
|
||||
if n < kib*kib*kib {
|
||||
return fmt.Sprintf("%.1f MiB", float64(n)/(kib*kib))
|
||||
}
|
||||
return fmt.Sprintf("%.1f GiB", float64(n)/(kib*kib*kib))
|
||||
}
|
||||
|
||||
func (s *Server) renderUI(w http.ResponseWriter, name string, data any) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; base-uri 'none'; frame-ancestors 'none'")
|
||||
if err := uiTemplates.ExecuteTemplate(w, name, data); err != nil {
|
||||
s.log.Error("render UI", "err", err)
|
||||
http.Error(w, "internal error", http.StatusInternalServerError)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleUIHome(w http.ResponseWriter, r *http.Request) {
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
view, err := s.uc.Dashboard.Overview(ctx, 20)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
s.renderUI(w, "dashboard.html", view)
|
||||
}
|
||||
|
||||
func (s *Server) handleUINewJob(w http.ResponseWriter, r *http.Request) {
|
||||
s.renderUI(w, "new-job.html", nil)
|
||||
}
|
||||
|
||||
func (s *Server) uiJobID(w http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
|
||||
return s.pathUUID(w, r, "job_id")
|
||||
}
|
||||
|
||||
func (s *Server) handleUIJob(w http.ResponseWriter, r *http.Request) {
|
||||
jobID, ok := s.uiJobID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
view, err := s.uc.Dashboard.JobDetail(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
s.renderUI(w, "job.html", view)
|
||||
}
|
||||
|
||||
func (s *Server) handleUIJobJSON(w http.ResponseWriter, r *http.Request) {
|
||||
jobID, ok := s.uiJobID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
view, err := s.uc.Dashboard.JobDetail(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, view)
|
||||
}
|
||||
|
||||
func (s *Server) handleUIArtifactDownload(w http.ResponseWriter, r *http.Request) {
|
||||
jobID, ok := s.uiJobID(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
artifactID, err := uuid.Parse(r.PathValue("artifact_id"))
|
||||
if err != nil {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
belongs, err := s.uc.Dashboard.ArtifactBelongsToJob(ctx, jobID, artifactID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
if !belongs {
|
||||
s.writeError(w, r, domain.ErrArtifactNotFound)
|
||||
return
|
||||
}
|
||||
// Reuse the coordinator-owned blob stream after the job-scoped check above.
|
||||
art, body, err := s.uc.DownloadArtifact.Execute(ctx, artifactID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if err := body.Close(); err != nil {
|
||||
s.log.Warn("close downloaded UI artifact", "artifact_id", artifactID, "err", err)
|
||||
}
|
||||
}()
|
||||
w.Header().Set("Content-Type", art.ContentType)
|
||||
w.Header().Set("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": art.Filename}))
|
||||
w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10))
|
||||
w.Header().Set("X-Checksum-SHA256", art.SHA256)
|
||||
_, _ = io.Copy(w, body)
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package http
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestUIStatusPresentation(t *testing.T) {
|
||||
tests := []struct {
|
||||
status string
|
||||
label string
|
||||
class string
|
||||
}{
|
||||
{"pending", "Waiting for a worker", "waiting"},
|
||||
{"running", "Running", "active"},
|
||||
{"completed", "Tasks complete", "success"},
|
||||
{"failed", "Needs attention", "danger"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.status, func(t *testing.T) {
|
||||
if got := uiStatusLabel(test.status); got != test.label {
|
||||
t.Errorf("label = %q, want %q", got, test.label)
|
||||
}
|
||||
if got := uiStatusClass(test.status); got != test.class {
|
||||
t.Errorf("class = %q, want %q", got, test.class)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIProgressPercent(t *testing.T) {
|
||||
if got := uiProgressPercent(3, 1, 0, 8); got != 50 {
|
||||
t.Errorf("progress = %d, want 50", got)
|
||||
}
|
||||
if got := uiProgressPercent(1, 1, 0, 0); got != 0 {
|
||||
t.Errorf("empty progress = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUITaskErrorPresentationDoesNotExposeCommand(t *testing.T) {
|
||||
if got := uiTaskErrorLabel("CalledProcessError"); got != "Local calculation failed" {
|
||||
t.Errorf("error label = %q", got)
|
||||
}
|
||||
if got := uiTaskErrorHint("CalledProcessError"); got == "" {
|
||||
t.Error("error hint must explain the failure")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
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
|
||||
tx TxManager
|
||||
clk Clock
|
||||
}
|
||||
|
||||
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
|
||||
blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact {
|
||||
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, 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, uc.clk.Now()) {
|
||||
return nil, domain.ErrLeaseConflict
|
||||
}
|
||||
// A client can retry a PUT after losing the response. Return the one durable
|
||||
// result for this lease attempt instead of storing duplicate artifacts.
|
||||
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if existing != nil {
|
||||
return existing, nil
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
attempt := in.Attempt
|
||||
art.Attempt = &attempt
|
||||
|
||||
// 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)
|
||||
|
||||
// The stream may take longer than the lease. Lock the task while re-checking
|
||||
// ownership and inserting metadata: completion or another upload cannot race
|
||||
// this final decision. The database unique index is a second line of defence.
|
||||
var durable *domain.Artifact
|
||||
err = uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
current, err := uc.tasks.GetForUpdate(ctx, in.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
|
||||
return domain.ErrLeaseConflict
|
||||
}
|
||||
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if existing != nil {
|
||||
durable = existing
|
||||
return nil
|
||||
}
|
||||
if err := uc.artifacts.Insert(ctx, art); err != nil {
|
||||
return err
|
||||
}
|
||||
durable = art
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||
return nil, err
|
||||
}
|
||||
if durable != art {
|
||||
// Another request won the race while this stream was being written.
|
||||
_ = uc.blobs.Delete(ctx, art.StorageKey)
|
||||
}
|
||||
return durable, 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
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
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
|
||||
// MaxRows limits how many data rows are turned into shards. Zero means the
|
||||
// whole uploaded dataset; the input artifact itself remains stored intact.
|
||||
MaxRows 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
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
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) {
|
||||
if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" {
|
||||
// CTX-10 must plan triangular block pairs; ordinary independent input
|
||||
// chunks would silently omit every cross-chunk molecular pair.
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
if (in.Workload == "similarity-search" || in.Workload == "similarity_search") &&
|
||||
len(in.Chunks) > 1 && in.Parameters["query_id"] != nil {
|
||||
// Resolving once against the source dataset belongs to CTX-07. Letting
|
||||
// each shard resolve it would make most tasks fail or use inconsistent data.
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
// --- CancelJob -----------------------------------------------------------
|
||||
|
||||
type CancelJob struct {
|
||||
jobs JobRepository
|
||||
tasks TaskRepository
|
||||
tx TxManager
|
||||
clock Clock
|
||||
}
|
||||
|
||||
func NewCancelJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CancelJob {
|
||||
return &CancelJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock}
|
||||
}
|
||||
|
||||
// Execute stops a job atomically. Completed and finally failed tasks are kept
|
||||
// as historical evidence; all other tasks are cancelled, including leased and
|
||||
// running ones. A repeated cancel of an already cancelled job is idempotent.
|
||||
func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error) {
|
||||
now := uc.clock.Now()
|
||||
var cancelled int64
|
||||
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
job, err := uc.jobs.Get(ctx, jobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if job.Status == domain.JobCancelled {
|
||||
return nil
|
||||
}
|
||||
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
|
||||
return domain.ErrJobNotCancellable
|
||||
}
|
||||
// The lease reaper can be the transition that exhausted the final task.
|
||||
// Check the authoritative task histogram as well as the cached job status,
|
||||
// so a stale status can never turn a failed/completed job into cancelled.
|
||||
counts, err := uc.tasks.CountByStatus(ctx, jobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
derived := progressFrom(*job, counts).DeriveStatus()
|
||||
if derived == domain.JobCompleted || derived == domain.JobFailed {
|
||||
return domain.ErrJobNotCancellable
|
||||
}
|
||||
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return uc.jobs.UpdateStatus(ctx, jobID, domain.JobCancelled, &now)
|
||||
})
|
||||
return cancelled, err
|
||||
}
|
||||
|
||||
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],
|
||||
Cancelled: counts[domain.TaskCancelled],
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
func syncExpiredJobStatuses(ctx context.Context, jobs JobRepository, tasks TaskRepository,
|
||||
jobIDs []uuid.UUID, now time.Time) error {
|
||||
seen := make(map[uuid.UUID]struct{}, len(jobIDs))
|
||||
for _, jobID := range jobIDs {
|
||||
if _, duplicate := seen[jobID]; duplicate {
|
||||
continue
|
||||
}
|
||||
seen[jobID] = struct{}{}
|
||||
if err := syncJobStatus(ctx, jobs, tasks, jobID, now); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
// 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)
|
||||
|
||||
// CancelByJob marks every non-terminal task as cancelled and invalidates its
|
||||
// lease. It returns how many tasks changed.
|
||||
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
|
||||
|
||||
// ExpireLeases applies the lease-expiry rule to every elapsed task and returns
|
||||
// the distinct jobs whose aggregate status may have changed.
|
||||
ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, 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)
|
||||
// FindPartialResult returns the durable result already uploaded for one task
|
||||
// attempt. A nil artifact means the attempt has not uploaded one yet.
|
||||
FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*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")
|
||||
@@ -0,0 +1,281 @@
|
||||
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
|
||||
jobs JobRepository
|
||||
workers WorkerRepository
|
||||
tx TxManager
|
||||
clock Clock
|
||||
leaseDuration time.Duration
|
||||
}
|
||||
|
||||
func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask {
|
||||
return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, 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
|
||||
}
|
||||
workloads := in.Workloads
|
||||
if workerID, err := uuid.Parse(in.WorkerID); err == nil {
|
||||
worker, err := uc.workers.Get(ctx, workerID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Never trust caller-supplied capabilities: registration is the durable
|
||||
// worker identity and its allowlist.
|
||||
workloads = worker.Capabilities
|
||||
}
|
||||
var claimed *domain.ClaimedTask
|
||||
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
now := uc.clock.Now()
|
||||
affectedJobs, err := uc.tasks.ExpireLeases(ctx, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affectedJobs, now); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
|
||||
Workloads: workloads,
|
||||
Owner: in.WorkerID,
|
||||
Now: now,
|
||||
LeaseUntil: now.Add(uc.leaseDuration),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if task != nil {
|
||||
value := task.AsClaimed()
|
||||
claimed = &value
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return claimed, nil // nil means an empty queue
|
||||
}
|
||||
|
||||
// --- 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
|
||||
}
|
||||
now := uc.clock.Now()
|
||||
if err := task.RenewLease(in.WorkerID, in.Attempt, now, 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.Attempt, 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 uuid.UUID, attempt int, artifactID uuid.UUID) error {
|
||||
art, err := uc.artifacts.Get(ctx, artifactID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if art.TaskID == nil || *art.TaskID != taskID || art.Attempt == nil || *art.Attempt != attempt || 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
|
||||
jobs JobRepository
|
||||
tx TxManager
|
||||
clock Clock
|
||||
}
|
||||
|
||||
func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases {
|
||||
return &ExpireLeases{tasks: tasks, jobs: jobs, tx: tx, clock: clock}
|
||||
}
|
||||
|
||||
// Execute reclaims elapsed tasks and persists the state of every affected job.
|
||||
//
|
||||
// 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) {
|
||||
var affected []uuid.UUID
|
||||
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
|
||||
now := uc.clock.Now()
|
||||
var err error
|
||||
affected, err = uc.tasks.ExpireLeases(ctx, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affected, now)
|
||||
})
|
||||
return int64(len(affected)), err
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package usecase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
// UIReadRepository is a read-only projection source for the local operator UI.
|
||||
// It intentionally exposes no storage paths or credentials.
|
||||
type UIReadRepository interface {
|
||||
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
|
||||
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
|
||||
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error)
|
||||
ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error)
|
||||
ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
|
||||
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error)
|
||||
}
|
||||
|
||||
type JobCard struct {
|
||||
ID string `json:"id"`
|
||||
Workload string `json:"workload"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Total int `json:"total"`
|
||||
Pending int `json:"pending"`
|
||||
Leased int `json:"leased"`
|
||||
Running int `json:"running"`
|
||||
Completed int `json:"completed"`
|
||||
Failed int `json:"failed"`
|
||||
Cancelled int `json:"cancelled"`
|
||||
}
|
||||
|
||||
type TaskCard struct {
|
||||
ID string `json:"id"`
|
||||
ChunkIndex int `json:"chunk_index"`
|
||||
Status string `json:"status"`
|
||||
Attempt int `json:"attempt"`
|
||||
MaxAttempts int `json:"max_attempts"`
|
||||
LeaseOwner string `json:"lease_owner,omitempty"`
|
||||
LeaseExpiresAt *time.Time `json:"lease_expires_at,omitempty"`
|
||||
ErrorCode string `json:"error_code,omitempty"`
|
||||
ErrorMessage string `json:"error_message,omitempty"`
|
||||
}
|
||||
|
||||
type ArtifactCard struct {
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
Filename string `json:"filename"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
SHA256 string `json:"sha256"`
|
||||
Downloadable bool `json:"downloadable"`
|
||||
Diagnostic bool `json:"diagnostic"`
|
||||
}
|
||||
|
||||
type WorkerCard struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Status string `json:"status"`
|
||||
Capabilities []string `json:"capabilities"`
|
||||
LastHeartbeatAt time.Time `json:"last_heartbeat_at"`
|
||||
}
|
||||
|
||||
type DashboardView struct {
|
||||
Jobs []JobCard
|
||||
Workers []WorkerCard
|
||||
}
|
||||
type JobDetailView struct {
|
||||
JobCard
|
||||
Tasks []TaskCard `json:"tasks"`
|
||||
Artifacts []ArtifactCard `json:"artifacts"`
|
||||
FinalResultAvailable bool `json:"final_result_available"`
|
||||
}
|
||||
|
||||
type Dashboard struct{ read UIReadRepository }
|
||||
|
||||
func NewDashboard(read UIReadRepository) *Dashboard { return &Dashboard{read: read} }
|
||||
|
||||
func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, error) {
|
||||
jobs, err := d.read.ListJobs(ctx, limit)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
workers, err := d.read.ListWorkers(ctx, limit)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))}
|
||||
jobIDs := make([]uuid.UUID, 0, len(jobs))
|
||||
for _, job := range jobs {
|
||||
jobIDs = append(jobIDs, job.ID)
|
||||
}
|
||||
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
|
||||
if err != nil {
|
||||
return DashboardView{}, err
|
||||
}
|
||||
for _, job := range jobs {
|
||||
out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID]))
|
||||
}
|
||||
for _, worker := range workers {
|
||||
out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (d *Dashboard) JobDetail(ctx context.Context, jobID uuid.UUID) (JobDetailView, error) {
|
||||
job, err := d.read.GetJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return JobDetailView{}, err
|
||||
}
|
||||
tasks, err := d.read.ListTasksByJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return JobDetailView{}, err
|
||||
}
|
||||
artifacts, err := d.read.ListArtifactsByJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return JobDetailView{}, err
|
||||
}
|
||||
out := JobDetailView{JobCard: jobCard(*job, tasks), Tasks: make([]TaskCard, 0, len(tasks)), Artifacts: make([]ArtifactCard, 0, len(artifacts))}
|
||||
for _, task := range tasks {
|
||||
card := TaskCard{ID: task.ID.String(), ChunkIndex: task.ChunkIndex, Status: string(task.Status), Attempt: task.Attempt, MaxAttempts: task.MaxAttempts, LeaseExpiresAt: task.LeaseExpiresAt}
|
||||
if task.LeaseOwner != nil {
|
||||
card.LeaseOwner = *task.LeaseOwner
|
||||
}
|
||||
if task.ErrorCode != nil {
|
||||
card.ErrorCode = *task.ErrorCode
|
||||
}
|
||||
if task.ErrorMessage != nil {
|
||||
card.ErrorMessage = *task.ErrorMessage
|
||||
}
|
||||
out.Tasks = append(out.Tasks, card)
|
||||
}
|
||||
for _, artifact := range artifacts {
|
||||
diagnostic := artifact.Kind == domain.ArtifactPartialResult
|
||||
downloadable := diagnostic || (artifact.Kind == domain.ArtifactFinalResult && out.Status == string(domain.JobCompleted))
|
||||
out.Artifacts = append(out.Artifacts, ArtifactCard{ID: artifact.ID.String(), Kind: string(artifact.Kind), Filename: artifact.Filename, SizeBytes: artifact.SizeBytes, SHA256: artifact.SHA256, Downloadable: downloadable, Diagnostic: diagnostic})
|
||||
if artifact.Kind == domain.ArtifactFinalResult && downloadable {
|
||||
out.FinalResultAvailable = true
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (d *Dashboard) ArtifactBelongsToJob(ctx context.Context, jobID, artifactID uuid.UUID) (bool, error) {
|
||||
artifacts, err := d.read.ListArtifactsByJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
for _, a := range artifacts {
|
||||
if a.ID == artifactID {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func jobCard(job domain.Job, tasks []domain.Task) JobCard {
|
||||
c := JobCard{ID: job.ID.String(), Workload: job.Workload, CreatedAt: job.CreatedAt}
|
||||
for _, task := range tasks {
|
||||
c.Total++
|
||||
switch task.Status {
|
||||
case domain.TaskPending:
|
||||
c.Pending++
|
||||
case domain.TaskLeased:
|
||||
c.Leased++
|
||||
case domain.TaskRunning:
|
||||
c.Running++
|
||||
case domain.TaskCompleted:
|
||||
c.Completed++
|
||||
case domain.TaskFailed:
|
||||
c.Failed++
|
||||
case domain.TaskCancelled:
|
||||
c.Cancelled++
|
||||
}
|
||||
}
|
||||
p := domain.JobProgress{Job: job, Total: c.Total, Pending: c.Pending, Leased: c.Leased + c.Running, Done: c.Completed, Failed: c.Failed, Cancelled: c.Cancelled}
|
||||
c.Status = string(p.DeriveStatus())
|
||||
return c
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
package usecase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
|
||||
"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
|
||||
maxAttempts int
|
||||
}
|
||||
|
||||
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
|
||||
tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
|
||||
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts}
|
||||
}
|
||||
|
||||
func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
|
||||
if err := validateUploadedWorkload(in.Workload, in.Parameters); err != nil {
|
||||
return SubmitDatasetResult{}, err
|
||||
}
|
||||
if uc.maxAttempts < 1 {
|
||||
return SubmitDatasetResult{}, domain.ErrInvalidInput
|
||||
}
|
||||
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.SplitChEMBLTSVLimit(rc, in.RowsPerShard, in.MaxRows, 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, uc.maxAttempts, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
shards = append(shards, art)
|
||||
tasks = append(tasks, task)
|
||||
return nil
|
||||
})
|
||||
_ = rc.Close()
|
||||
if splitErr != nil {
|
||||
cleanup()
|
||||
// Dataset shape is caller input, not an internal coordinator failure.
|
||||
return SubmitDatasetResult{}, domain.ErrInvalidInput
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// validateUploadedWorkload is deliberately narrow until CTX-07/08/10 adds a
|
||||
// typed distributed-workload registry. In particular, running similarity-graph
|
||||
// independently per TSV shard is scientifically wrong: cross-shard pairs would
|
||||
// be absent from the apparent graph.
|
||||
func validateUploadedWorkload(workload string, parameters map[string]any) error {
|
||||
if workload != "similarity-search" {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
allowed := map[string]struct{}{
|
||||
"query_smiles": {}, "top_k": {}, "threshold": {},
|
||||
"threshold_direction": {}, "progress_every": {},
|
||||
}
|
||||
for key := range parameters {
|
||||
if _, ok := allowed[key]; !ok {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
}
|
||||
query, ok := parameters["query_smiles"].(string)
|
||||
if !ok || query == "" || len(query) > 200 {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["top_k"]; ok && !isPositiveJSONInteger(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["progress_every"]; ok && !isNonNegativeJSONInteger(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["threshold"]; ok && !isUnitIntervalNumber(value) {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
if value, ok := parameters["threshold_direction"]; ok && value != "greater" && value != "less" {
|
||||
return domain.ErrInvalidInput
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isPositiveJSONInteger(value any) bool { return isJSONInteger(value, false) }
|
||||
func isNonNegativeJSONInteger(value any) bool { return isJSONInteger(value, true) }
|
||||
|
||||
func isJSONInteger(value any, allowZero bool) bool {
|
||||
var n int64
|
||||
switch v := value.(type) {
|
||||
case int:
|
||||
n = int64(v)
|
||||
case int64:
|
||||
n = v
|
||||
case float64:
|
||||
if math.Trunc(v) != v || v > math.MaxInt64 || v < math.MinInt64 {
|
||||
return false
|
||||
}
|
||||
n = int64(v)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return n >= 0 && (allowZero || n > 0)
|
||||
}
|
||||
|
||||
func isUnitIntervalNumber(value any) bool {
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v <= 1
|
||||
case int:
|
||||
return v >= 0 && v <= 1
|
||||
case int64:
|
||||
return v >= 0 && v <= 1
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,627 @@
|
||||
package usecase_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"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
|
||||
|
||||
type expiringBlobStore struct {
|
||||
*memstore.BlobStore
|
||||
clock *memstore.Clock
|
||||
}
|
||||
|
||||
func (s expiringBlobStore) Put(ctx context.Context, key string, body io.Reader) (string, int64, error) {
|
||||
sum, size, err := s.BlobStore.Put(ctx, key, body)
|
||||
s.clock.Advance(lease + time.Second)
|
||||
return sum, size, err
|
||||
}
|
||||
|
||||
// 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
|
||||
cancel *usecase.CancelJob
|
||||
}
|
||||
|
||||
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, 3)
|
||||
h.claim = usecase.NewClaimTask(h.tasks, h.jobs, h.work, tx, h.clk, lease)
|
||||
h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease)
|
||||
h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
|
||||
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
|
||||
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
|
||||
h.results = usecase.NewListResults(h.tasks)
|
||||
h.register = usecase.NewRegisterWorker(h.work, h.clk)
|
||||
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, tx, 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.jobs, tx, h.clk)
|
||||
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, 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 TestRegisteredWorkerCannotBroadenItsCapabilitiesAtClaim(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "restricted", 1)
|
||||
worker, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{
|
||||
Name: "search-only", Capabilities: []string{"similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{
|
||||
WorkerID: worker.ID.String(), Workloads: []string{"restricted"},
|
||||
})
|
||||
if err != nil || claimed != nil {
|
||||
t.Fatalf("claim = (%v, %v), want no compatible task", claimed, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateJobRejectsUnsafeDistributedScientificPlans(t *testing.T) {
|
||||
h := newHarness()
|
||||
_, err := h.createJob.Execute(ctx, usecase.CreateJobInput{
|
||||
Workload: "similarity-graph", InputURI: "s3://input",
|
||||
Chunks: []usecase.ChunkInput{{ChunkIndex: 0, InputURI: "s3://chunk", InputSHA256: "sha"}},
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("graph job err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
_, err = h.createJob.Execute(ctx, usecase.CreateJobInput{
|
||||
Workload: "similarity-search", InputURI: "s3://input", Parameters: map[string]any{"query_id": "CHEMBL1"},
|
||||
Chunks: []usecase.ChunkInput{
|
||||
{ChunkIndex: 0, InputURI: "s3://chunk0", InputSHA256: "sha"},
|
||||
{ChunkIndex: 1, InputURI: "s3://chunk1", InputSHA256: "sha"},
|
||||
},
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("sharded query_id job err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
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 TestCompleteRejectsArtifactFromExpiredAttempt(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attemptOne := h.leaseOne(t, "w1", "w")
|
||||
staleArtifact := h.uploadResult(t, taskID, "w1", attemptOne)
|
||||
|
||||
h.clk.Advance(lease + time.Second)
|
||||
if _, err := h.expire.Execute(ctx); err != nil {
|
||||
t.Fatalf("expire lease: %v", err)
|
||||
}
|
||||
_, attemptTwo := h.leaseOne(t, "w2", "w")
|
||||
if attemptTwo != attemptOne+1 {
|
||||
t.Fatalf("attempt = %d, want %d", attemptTwo, attemptOne+1)
|
||||
}
|
||||
|
||||
_, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{
|
||||
TaskID: taskID, WorkerID: "w2", Attempt: attemptTwo, ResultArtifactID: staleArtifact,
|
||||
})
|
||||
if !errors.Is(err, domain.ErrResultConflict) {
|
||||
t.Errorf("stale-attempt artifact: err = %v, want ErrResultConflict", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||
h.uploadArt = usecase.NewUploadArtifact(
|
||||
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, memstore.Tx{}, h.clk,
|
||||
)
|
||||
|
||||
_, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
|
||||
TaskID: taskID, WorkerID: "w1", Attempt: attempt,
|
||||
Filename: "result.csv", ContentType: "text/csv", Body: strings.NewReader("result"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrLeaseConflict) {
|
||||
t.Errorf("upload after lease expiry: err = %v, want ErrLeaseConflict", 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 TestUploadArtifactIsIdempotentPerTaskAttempt(t *testing.T) {
|
||||
h := newHarness()
|
||||
h.seedJob(t, "w", 1)
|
||||
taskID, attempt := h.leaseOne(t, "w1", "w")
|
||||
first := h.uploadResult(t, taskID, "w1", attempt)
|
||||
second, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
|
||||
TaskID: taskID, WorkerID: "w1", Attempt: attempt,
|
||||
Filename: "retry.csv", ContentType: "text/csv", Body: strings.NewReader("different bytes"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second.ID != first {
|
||||
t.Errorf("retry artifact = %s, want existing %s", second.ID, first)
|
||||
}
|
||||
}
|
||||
|
||||
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 := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, 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{"similarity-search"}})
|
||||
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 TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
|
||||
h := newHarness()
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
|
||||
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.TaskCount != 2 {
|
||||
t.Fatalf("task_count = %d, want 2", res.TaskCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubmitDatasetRejectsUnsupportedDistributedWorkloads(t *testing.T) {
|
||||
h := newHarness()
|
||||
_, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-graph", Parameters: map[string]any{"threshold": 0.7}, RowsPerShard: 2,
|
||||
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
|
||||
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("graph submission err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
_, err = h.submit.Execute(ctx, usecase.SubmitDatasetInput{
|
||||
Workload: "similarity-search", Parameters: map[string]any{"query_id": "CHEMBL1"}, RowsPerShard: 2,
|
||||
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
|
||||
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
|
||||
})
|
||||
if !errors.Is(err, domain.ErrInvalidInput) {
|
||||
t.Errorf("query_id submission err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
|
||||
h := newHarness()
|
||||
jobID := h.seedJob(t, "w", 3)
|
||||
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
|
||||
if err != nil || claimed == nil {
|
||||
t.Fatalf("claim: %v", err)
|
||||
}
|
||||
cancelled, err := h.cancel.Execute(ctx, jobID)
|
||||
if err != nil || cancelled != 3 {
|
||||
t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err)
|
||||
}
|
||||
if _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: claimed.TaskID, WorkerID: "w1", Attempt: claimed.Attempt}); !errors.Is(err, domain.ErrTaskNotLeased) {
|
||||
t.Errorf("cancelled lease heartbeat = %v, want ErrTaskNotLeased", err)
|
||||
}
|
||||
progress, err := h.status.Execute(ctx, jobID)
|
||||
if err != nil || progress.DeriveStatus() != domain.JobCancelled || progress.Cancelled != 3 {
|
||||
t.Errorf("cancelled progress = %+v, err = %v", progress, err)
|
||||
}
|
||||
if cancelled, err := h.cancel.Execute(ctx, jobID); err != nil || cancelled != 0 {
|
||||
t.Errorf("second cancel = (%d, %v), want (0, nil)", cancelled, err)
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinalLeaseExpiryPersistsFailedJobAndCannotBeCancelled(t *testing.T) {
|
||||
h := newHarness()
|
||||
jobID := h.seedJob(t, "w", 1)
|
||||
for attempt := 1; attempt <= domain.DefaultMaxAttempts; attempt++ {
|
||||
h.leaseOne(t, "w1", "w")
|
||||
h.clk.Advance(lease + time.Second)
|
||||
if _, err := h.expire.Execute(ctx); err != nil {
|
||||
t.Fatalf("expire attempt %d: %v", attempt, err)
|
||||
}
|
||||
}
|
||||
progress, err := h.status.Execute(ctx, jobID)
|
||||
if err != nil || progress.Job.Status != domain.JobFailed || progress.DeriveStatus() != domain.JobFailed {
|
||||
t.Fatalf("progress = %+v, err = %v; want persisted failed job", progress, err)
|
||||
}
|
||||
if _, err := h.cancel.Execute(ctx, jobID); !errors.Is(err, domain.ErrJobNotCancellable) {
|
||||
t.Errorf("cancel terminal lease failure = %v, want ErrJobNotCancellable", err)
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -0,0 +1,6 @@
|
||||
BEGIN;
|
||||
|
||||
DROP TABLE IF EXISTS workers;
|
||||
DROP TYPE IF EXISTS worker_status;
|
||||
|
||||
COMMIT;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
@@ -0,0 +1,4 @@
|
||||
-- PostgreSQL cannot drop a single enum value without recreating the type and
|
||||
-- rewriting every dependent column. Leaving 'running' in place is harmless: no
|
||||
-- code writes it after the down of 0007 restores the leased-only transitions.
|
||||
SELECT 1;
|
||||
@@ -0,0 +1,5 @@
|
||||
-- 'running' means the worker has acknowledged start via its first heartbeat.
|
||||
-- Kept in its own migration, without an explicit transaction: an enum value
|
||||
-- added in a transaction cannot be USED in that same transaction, and the next
|
||||
-- migration references it.
|
||||
ALTER TYPE task_status ADD VALUE IF NOT EXISTS 'running';
|
||||
@@ -0,0 +1,8 @@
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_leased_owner;
|
||||
ALTER TABLE tasks ADD CONSTRAINT ck_tasks_leased_owner CHECK (
|
||||
status <> 'leased' OR (lease_owner IS NOT NULL AND lease_expires_at IS NOT NULL)
|
||||
);
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,10 @@
|
||||
BEGIN;
|
||||
|
||||
-- A running task holds a lease just like a leased one, so the lease-integrity
|
||||
-- check must cover both states.
|
||||
ALTER TABLE tasks DROP CONSTRAINT IF EXISTS ck_tasks_leased_owner;
|
||||
ALTER TABLE tasks ADD CONSTRAINT ck_tasks_leased_owner CHECK (
|
||||
status NOT IN ('leased','running') OR (lease_owner IS NOT NULL AND lease_expires_at IS NOT NULL)
|
||||
);
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,7 @@
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE artifacts DROP CONSTRAINT IF EXISTS ck_artifact_attempt_positive;
|
||||
ALTER TABLE artifacts DROP CONSTRAINT IF EXISTS ck_partial_result_attempt;
|
||||
ALTER TABLE artifacts DROP COLUMN IF EXISTS attempt;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,34 @@
|
||||
BEGIN;
|
||||
|
||||
-- A partial result belongs to the lease attempt that uploaded it. Without this
|
||||
-- binding a worker holding a later retry could complete a task with stale bytes
|
||||
-- uploaded by an expired attempt of that same task.
|
||||
ALTER TABLE artifacts ADD COLUMN attempt integer;
|
||||
|
||||
-- A completed task never gets a later lease, so its current attempt is also
|
||||
-- the attempt that produced the stored result.
|
||||
UPDATE artifacts AS a
|
||||
SET attempt = t.attempt
|
||||
FROM tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status = 'completed'::task_status
|
||||
AND a.attempt IS NULL;
|
||||
|
||||
-- For unfinished tasks the old schema cannot tell which attempt uploaded a
|
||||
-- partial result. Keeping it would let a later retry claim stale bytes, so the
|
||||
-- worker must upload again. Blob garbage is harmless and follows the existing
|
||||
-- coordinator-owned storage cleanup policy.
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status <> 'completed'::task_status
|
||||
AND a.attempt IS NULL;
|
||||
|
||||
ALTER TABLE artifacts ADD CONSTRAINT ck_partial_result_attempt
|
||||
CHECK (kind <> 'partial_result'::artifact_kind OR attempt IS NOT NULL);
|
||||
ALTER TABLE artifacts ADD CONSTRAINT ck_artifact_attempt_positive
|
||||
CHECK (attempt IS NULL OR attempt > 0);
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,5 @@
|
||||
BEGIN;
|
||||
|
||||
DROP INDEX IF EXISTS uq_partial_result_task_attempt;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,26 @@
|
||||
BEGIN;
|
||||
|
||||
-- Old deployments can contain more than one partial result because earlier
|
||||
-- versions accepted repeated PUTs. Preserve the one referenced by a completed
|
||||
-- task and discard stale rows; unfinished tasks must upload again after a
|
||||
-- deploy, just as they do after a lost lease.
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status <> 'completed'::task_status;
|
||||
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status = 'completed'::task_status
|
||||
AND a.id <> t.result_artifact_id;
|
||||
|
||||
-- One lease attempt has one durable partial result. This makes an upload retry
|
||||
-- idempotent and prevents repeated uploads from accumulating orphan artifacts.
|
||||
CREATE UNIQUE INDEX uq_partial_result_task_attempt
|
||||
ON artifacts (task_id, attempt)
|
||||
WHERE kind = 'partial_result'::artifact_kind;
|
||||
|
||||
COMMIT;
|
||||
Executable
+215
@@ -0,0 +1,215 @@
|
||||
#!/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"}'
|
||||
registration=$(curl -sS "${auth[@]}" -X POST "${HOST}/workers/register" \
|
||||
-d '{"name":"smoke-worker-active","capabilities":["similarity_search","similarity-search"],"cpu_count":4}')
|
||||
worker_id=$(printf '%s' "$registration" | python3 -c 'import json,sys;print(json.load(sys.stdin)["worker_id"])' 2>/dev/null)
|
||||
if [[ -z "${worker_id:-}" ]]; then
|
||||
echo " ✗ could not register active worker: $registration"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
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\":\"${worker_id}\"}")
|
||||
[[ -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\":\"${worker_id}\",\"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: ${worker_id}" -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: ${worker_id}" -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'
|
||||
|
||||
# A retry of a PUT returns the same durable artifact for the task attempt.
|
||||
art_id=$(upload primary.csv)
|
||||
art_id2=$(upload secondary.csv)
|
||||
if [[ "$art_id" == "$art_id2" ]]; then
|
||||
printf ' \033[32m✓\033[0m %-46s %s\n' "duplicate upload is idempotent" "$art_id"
|
||||
pass=$((pass + 1))
|
||||
else
|
||||
printf ' \033[31m✗\033[0m %-46s got %s and %s\n' "duplicate upload is idempotent" "$art_id" "$art_id2"
|
||||
fail=$((fail + 1))
|
||||
fi
|
||||
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\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
check "replay same result → idempotent" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
|
||||
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
|
||||
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\":\"${worker_id}\",\"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\":\"${worker_id}\",\"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={"query_smiles":"CCO","top_k":10}' \
|
||||
-F 'chunk_rows=2' \
|
||||
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
|
||||
chembl_id canonical_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\":\"${worker_id}\"}")
|
||||
[[ -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 ]]
|
||||
@@ -0,0 +1,222 @@
|
||||
# SciMesh coordinator ↔ worker API contract (v1)
|
||||
|
||||
**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.
|
||||
|
||||
> **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.
|
||||
|
||||
- **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`.
|
||||
|
||||
## 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), optional
|
||||
`max_rows` (positive int), and the file part `file`. `max_rows` limits the
|
||||
leading data rows that become shards; it does not change the stored source
|
||||
artifact. The coordinator splits the selected TSV rows 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.
|
||||
|
||||
## Stop a job
|
||||
|
||||
```http
|
||||
POST /jobs/{job_id}/cancel
|
||||
Authorization: Bearer <token>
|
||||
```
|
||||
|
||||
The coordinator transactionally marks every pending, leased, or running shard
|
||||
as `cancelled`, invalidates its lease, and marks the job `cancelled`. Completed
|
||||
and terminally failed shards remain as history. Repeating a cancellation of an
|
||||
already cancelled job is safe.
|
||||
|
||||
`200`:
|
||||
|
||||
```json
|
||||
{ "job_id": "uuid", "status": "cancelled", "cancelled_tasks": 12 }
|
||||
```
|
||||
|
||||
## Register worker
|
||||
|
||||
```http
|
||||
POST /workers/register
|
||||
Authorization: Bearer <token>
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"name": "lab-worker-01",
|
||||
"capabilities": ["similarity-search", "similarity-graph"],
|
||||
"cpu_count": 8,
|
||||
"memory_mb": 16384
|
||||
}
|
||||
```
|
||||
|
||||
`201`:
|
||||
|
||||
```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. A claim uses the capabilities
|
||||
stored at registration; the request cannot broaden them.
|
||||
|
||||
## 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. `worker_id` must be a registered UUID;
|
||||
its persisted capabilities, rather than this request field, decide eligibility.
|
||||
|
||||
```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 }
|
||||
```
|
||||
|
||||
Response **must** contain a renewed deadline:
|
||||
|
||||
```json
|
||||
{ "lease_expires_at": "2026-07-22T12:10:00Z" }
|
||||
```
|
||||
|
||||
The worker schedules the next heartbeat before half of the returned TTL, never
|
||||
on a fixed interval alone.
|
||||
|
||||
## Download input or shard (CTX-05)
|
||||
|
||||
`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.
|
||||
|
||||
## 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
|
||||
|
||||
<streamed bytes>
|
||||
```
|
||||
|
||||
`200`:
|
||||
|
||||
```json
|
||||
{ "artifact_id": "uuid", "uri": "https://coordinator/artifacts/uuid/download",
|
||||
"sha256": "hex", "size_bytes": 1234 }
|
||||
```
|
||||
|
||||
## 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", "sha256": "hex", "content_type": "text/csv" },
|
||||
"metrics": { "elapsed_seconds": 12.4, "processed_rows": 10000 }
|
||||
}
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
```http
|
||||
POST /tasks/{task_id}/failure
|
||||
```
|
||||
|
||||
Same identity fields, plus sanitized `error_code`, `error_message`, `retryable`.
|
||||
Never a traceback, token, or absolute worker path.
|
||||
|
||||
## Idempotency and errors
|
||||
|
||||
| Situation | Response |
|
||||
| --- | --- |
|
||||
| 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` |
|
||||
@@ -0,0 +1,81 @@
|
||||
# Task: safely preview partial CSV artifacts in the UI
|
||||
|
||||
## Assignment
|
||||
|
||||
You are the junior developer implementing one contained UI feature: an
|
||||
authenticated operator can preview a small portion of a CSV artifact belonging
|
||||
to the job they are viewing. This is a diagnostic aid, not a final-results
|
||||
page.
|
||||
|
||||
## Read first
|
||||
|
||||
1. `AGENTS.md`
|
||||
2. `.agents/coordinator.md`
|
||||
3. `docs/web-interface-plan.md`
|
||||
4. `docs/api-contract.md`
|
||||
5. `coordinator/internal/transport/http/ui.go` and its tests
|
||||
|
||||
## Current baseline
|
||||
|
||||
The coordinator serves an authenticated local web UI. Job detail pages already
|
||||
list partial result artifacts and provide job-scoped downloads. The browser has
|
||||
no worker token and must not learn storage paths. A partial shard CSV is never
|
||||
a global or final molecular-search result.
|
||||
|
||||
## Scope
|
||||
|
||||
Add a **Preview** action next to eligible CSV artifacts on a job detail page.
|
||||
|
||||
- Preview only artifacts that belong to the requested job.
|
||||
- Show at most the first **30 rows** and read at most **64 KiB** from storage.
|
||||
- State clearly when content was truncated.
|
||||
- Preserve the existing download action.
|
||||
- For a non-CSV artifact, return a friendly, sanitized explanation rather than
|
||||
attempting to render bytes as text.
|
||||
- Use an existing UI route pattern or add a small UI-authenticated endpoint;
|
||||
keep it separate from worker API routes.
|
||||
|
||||
## Security rules
|
||||
|
||||
- Require UI Basic Auth for every preview request.
|
||||
- Verify job ownership in the coordinator before opening the artifact; an
|
||||
artifact ID from another job must not be previewable.
|
||||
- Never expose `storage_key`, filesystem paths, database errors, bearer tokens,
|
||||
or worker-local information.
|
||||
- Do not use `innerHTML` for CSV fields. Use `html/template` escaping or
|
||||
`textContent` so strings such as `<script>alert(1)</script>` are displayed as
|
||||
data, not executed.
|
||||
- Do not load the complete artifact into memory.
|
||||
|
||||
## Out of scope
|
||||
|
||||
- Charts, molecule imagery, RDKit rendering, client-side CSV libraries, React,
|
||||
and a new frontend service.
|
||||
- Changing job/task state, retrying tasks, or implementing reducer output.
|
||||
- Redesigning the broader dashboard or job-creation workflow; that belongs to
|
||||
`docs/user-space-task.md`.
|
||||
|
||||
## Suggested implementation shape
|
||||
|
||||
Keep UI transport, use case, and storage responsibilities separate. Return a
|
||||
small view model containing artifact metadata, column headers, rows, and a
|
||||
`truncated` flag. Reuse existing coordinator-owned artifact access rather than
|
||||
reading a path supplied by the browser. Keep the handler streaming/limited.
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- A valid partial CSV can be previewed from its own job detail page.
|
||||
- The result shows no more than 30 data rows and marks 64 KiB/row truncation.
|
||||
- Empty and malformed CSV content fail safely with a clear message.
|
||||
- A non-CSV artifact is rejected safely.
|
||||
- Unauthenticated access is rejected; a cross-job artifact request is not
|
||||
disclosed or served.
|
||||
- HTML-like values are escaped in the rendered preview.
|
||||
- Existing artifact downloads still work.
|
||||
- Add Go tests for all cases above and run `go test ./...` and `go vet ./...`.
|
||||
|
||||
## Handoff
|
||||
|
||||
Work in one focused branch and one PR. Report files changed, any API impact,
|
||||
test commands/results, and known limitations. Do not commit datasets, generated
|
||||
CSV files, Docker volumes, `.venv`, or `worker-data/`.
|
||||
@@ -0,0 +1,251 @@
|
||||
# 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 fixed at registration — the coordinator only hands you
|
||||
matching tasks and a later claim cannot broaden that set.
|
||||
- **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).
|
||||
- Current diagnostic uploads use `similarity-search` with `query_smiles`. The
|
||||
reference worker accepts the legacy `similarity_search` spelling too. Do not
|
||||
advertise `similarity-graph` until CTX-10 implements cross-shard pair planning.
|
||||
|
||||
## 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_smiles": "CCO", "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`)
|
||||
- worker name (the coordinator returns its `worker_id` at registration;
|
||||
`SCIMESH_WORKER_ID` is only a legacy/test override)
|
||||
- the bearer token
|
||||
- poll interval and request timeout
|
||||
- a working directory for downloaded inputs and generated outputs
|
||||
|
||||
## Run the reference worker locally
|
||||
|
||||
Use one terminal per worker and a distinct work directory for each process:
|
||||
|
||||
```sh
|
||||
SCIMESH_COORDINATOR_URL=http://localhost:8080 \
|
||||
SCIMESH_BEARER_TOKEN=dev-token \
|
||||
SCIMESH_WORKER_NAME=worker-1 \
|
||||
scimesh-worker --work-dir "$PWD/worker-data-1"
|
||||
```
|
||||
|
||||
For a bounded manual check, use one of these lifecycle modes:
|
||||
|
||||
```sh
|
||||
# Make exactly one claim; exit immediately when no task is available.
|
||||
scimesh-worker --work-dir "$PWD/worker-data-check" --once
|
||||
|
||||
# Keep polling until two tasks complete successfully, then exit.
|
||||
scimesh-worker --work-dir "$PWD/worker-data-check" --max-tasks 2
|
||||
```
|
||||
|
||||
`SCIMESH_MAX_TASKS` provides the same limit through the environment. Pressing
|
||||
`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task,
|
||||
the worker reports a sanitized retriable failure first, emits no traceback, and
|
||||
exits with status `130`.
|
||||
|
||||
## 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.
|
||||
@@ -0,0 +1,209 @@
|
||||
# CTX-07: distributed workload protocol and planner contract
|
||||
|
||||
## Status and scope
|
||||
|
||||
This document is the implementation contract for CTX-07. Its generic protocol,
|
||||
registry, strict JSON models, and deterministic reduction ordering are
|
||||
implemented in `scimesh/distributed/`. It does not implement a molecular
|
||||
planner, reducer, API endpoint, database migration, or final artifact. Until
|
||||
CTX-08 and CTX-09 are complete, shard CSVs remain diagnostic partial results.
|
||||
|
||||
The protocol gives local scientific workloads a coordinator-independent way to
|
||||
validate a job, plan artifact-backed tasks, and later reduce completed outputs.
|
||||
The Go coordinator owns durable artifacts, transactions, task rows, leases, and
|
||||
HTTP. A Python workload must never access PostgreSQL or call the coordinator.
|
||||
|
||||
Read `PLAN.md`, `.agents/workloads.md`, and `docs/api-contract.md` before
|
||||
implementing this CTX.
|
||||
|
||||
## Canonical vocabulary
|
||||
|
||||
- External workload names are lowercase hyphenated names: `similarity-search`
|
||||
and, later, `similarity-graph`.
|
||||
- The existing underscore spellings are a temporary compatibility alias at the
|
||||
Python worker boundary only. Planners, persisted job/task payloads, and new
|
||||
API examples use the canonical hyphenated spelling.
|
||||
- A **plan** contains only JSON-compatible values and coordinator artifact
|
||||
references. It contains no local filesystem path, worker URI, presigned URL,
|
||||
database connection, or callable.
|
||||
- `chunk_index` is a non-negative integer, unique within a plan, and sorted
|
||||
ascending whenever results are enumerated.
|
||||
|
||||
## Python boundary
|
||||
|
||||
CTX-07 adds a small `DistributedWorkload` protocol under `scimesh/distributed/`
|
||||
and a registry separate from the local CLI registry. Names below are proposed
|
||||
public types; keep concrete implementation details minimal.
|
||||
|
||||
```python
|
||||
class DistributedWorkload(Protocol):
|
||||
name: str
|
||||
|
||||
def validate_job(self, parameters: Mapping[str, object]) -> None: ...
|
||||
|
||||
def plan(
|
||||
self,
|
||||
input_path: Path,
|
||||
input_artifact_id: str,
|
||||
parameters: Mapping[str, object],
|
||||
shard_rows: int,
|
||||
workspace: Path,
|
||||
) -> DistributedPlan: ...
|
||||
|
||||
def reduce(
|
||||
self,
|
||||
partial_results: Sequence[CompletedPartial],
|
||||
parameters: Mapping[str, object],
|
||||
workspace: Path,
|
||||
) -> FinalResult: ...
|
||||
```
|
||||
|
||||
`input_path` and `workspace` are temporary files supplied by the coordinator
|
||||
bridge. They are never serialized. `plan()` returns only a `DistributedPlan`;
|
||||
the bridge validates it, persists artifact/task rows in one coordinator
|
||||
transaction, and removes its temporary workspace. If validation or planning
|
||||
fails, no job or task may be written.
|
||||
|
||||
## JSON models
|
||||
|
||||
All objects below are schema version `1`. Future incompatible changes require a
|
||||
new version; never infer a schema from missing fields.
|
||||
|
||||
### Artifact reference
|
||||
|
||||
```json
|
||||
{
|
||||
"artifact_id": "c4273293-f8b4-4ecb-99df-3b9f5a32b6a6",
|
||||
"sha256": "3b2d...64-lowercase-hex-characters",
|
||||
"content_type": "text/tab-separated-values"
|
||||
}
|
||||
```
|
||||
|
||||
The artifact ID is coordinator-owned. The checksum is included so planning and
|
||||
tests can assert exactly which immutable input was used. A worker receives the
|
||||
coordinator-generated download URI only through `POST /tasks/claim`.
|
||||
|
||||
### Distributed plan
|
||||
|
||||
```json
|
||||
{
|
||||
"schema_version": 1,
|
||||
"workload": "similarity-search",
|
||||
"resolved_parameters": {
|
||||
"query_smiles": "COc1ccc(Nc2ncnc3cc(OCCCN4CCOCC4)c(OC)c23)cc1",
|
||||
"query_source": {"kind": "chembl_id", "value": "CHEMBL939"},
|
||||
"top_k": 20,
|
||||
"threshold": 0.7,
|
||||
"threshold_direction": "greater",
|
||||
"fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}
|
||||
},
|
||||
"tasks": [
|
||||
{
|
||||
"chunk_index": 0,
|
||||
"input_artifact": {
|
||||
"artifact_id": "69e41105-d9fb-4c7f-a2db-7dd9e3ba2c76",
|
||||
"sha256": "4c92...64-lowercase-hex-characters",
|
||||
"content_type": "text/tab-separated-values"
|
||||
},
|
||||
"parameters": {
|
||||
"query_smiles": "COc1ccc(Nc2ncnc3cc(OCCCN4CCOCC4)c(OC)c23)cc1",
|
||||
"top_k": 20,
|
||||
"threshold": 0.7,
|
||||
"threshold_direction": "greater",
|
||||
"fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
`resolved_parameters` are immutable job metadata. A task copies only the
|
||||
values required by its worker runner. The coordinator may add its own durable
|
||||
task ID and generated input URI; it must not alter scientific parameters.
|
||||
|
||||
## Similarity-search planning rules
|
||||
|
||||
1. Accept exactly one of `query_id` and `query_smiles` at the public boundary.
|
||||
2. Validate a supplied SMILES once. For `query_id`, find and validate that
|
||||
molecule once against the original uploaded TSV **before** creating shards.
|
||||
3. Persist the resolved canonical query SMILES and the original query source in
|
||||
`resolved_parameters`. Workers receive `query_smiles`, never `query_id`.
|
||||
4. Fingerprint settings are fixed to Morgan radius `2` and `fp_size` `2048`.
|
||||
Reject a request that tries to override them rather than silently changing
|
||||
scientific semantics.
|
||||
5. Split source rows in input order. Every shard includes the original TSV
|
||||
header and has a contiguous, zero-based `chunk_index`.
|
||||
6. Each shard uses the global `top_k`, not a smaller local limit. A global
|
||||
reducer cannot recover a candidate discarded by every shard.
|
||||
7. Preserve `threshold`, `threshold_direction`, and valid `max_rows` semantics
|
||||
in the resolved plan. A job-level row limit is applied before sharding, not
|
||||
independently by every worker.
|
||||
|
||||
Invalid row SMILES are not planner failures. They remain shard data and are
|
||||
counted by the worker exactly as the local workload does. An invalid query is a
|
||||
planning failure.
|
||||
|
||||
## Partial-result contract
|
||||
|
||||
A completed similarity-search task owns exactly one coordinator-uploaded CSV
|
||||
artifact with content type `text/csv` and these columns, in this order:
|
||||
|
||||
```csv
|
||||
rank,chembl_id,canonical_smiles,similarity
|
||||
1,CHEMBL123,CCO,0.875000
|
||||
```
|
||||
|
||||
- `rank` is one-based local rank.
|
||||
- `similarity` uses the local CLI's six-decimal formatting.
|
||||
- Rows are sorted by `(-similarity, chembl_id, canonical_smiles)` for
|
||||
`threshold_direction=greater`, or `(similarity, chembl_id,
|
||||
canonical_smiles)` for `less`.
|
||||
- The query molecule and every candidate with the same canonical query SMILES
|
||||
are excluded using the existing local-workload definition.
|
||||
- Empty valid result files still include the header.
|
||||
|
||||
The worker completion metrics must include JSON numbers for `scanned_rows`,
|
||||
`valid_molecules`, `invalid_smiles`, `matches_emitted`, and
|
||||
`elapsed_seconds`. Metrics are observability data; the reducer derives final
|
||||
scientific output exclusively from coordinator-owned partial artifacts.
|
||||
|
||||
## Reduction boundary
|
||||
|
||||
CTX-09 invokes the registered reducer only after every task is completed. It
|
||||
passes `CompletedPartial` values ordered by `chunk_index`, each containing its
|
||||
coordinator artifact reference and validated metrics.
|
||||
|
||||
For similarity-search the reducer:
|
||||
|
||||
1. reads partial CSVs in `chunk_index` order;
|
||||
2. validates header, row shape, rank, finite similarity in `[0, 1]`, and sort
|
||||
order;
|
||||
3. retains a bounded heap of at most the global `top_k` candidates using the
|
||||
exact local ranking key;
|
||||
4. writes the same header and deterministic rank numbering as the local CLI.
|
||||
|
||||
It must not deduplicate ordinary records: the local reference keeps input-row
|
||||
multiplicity. Reduction is independent of worker completion order and uses
|
||||
`O(top_k + shard_rows)` memory apart from CSV streaming buffers.
|
||||
|
||||
## Required tests for the CTX-07 implementation
|
||||
|
||||
- unknown workload is rejected before any coordinator job/task write;
|
||||
- invalid public parameters and invalid `query_id` produce no partial plan;
|
||||
- `query_id` resolution occurs once, before shard construction;
|
||||
- the same input, parameters, and shard size generate byte-equivalent
|
||||
JSON plans and identical shard order;
|
||||
- every task payload is JSON-serializable and contains only artifact references
|
||||
and validated scalar/object values;
|
||||
- a two-shard dummy workload proves coordinator transaction rollback on planner
|
||||
validation failure;
|
||||
- completed partial artifacts reach the reducer ordered by `chunk_index`, even
|
||||
when workers finish in a different order;
|
||||
- the protocol registry never imports the Go coordinator or database code.
|
||||
|
||||
## Deferred work
|
||||
|
||||
CTX-08 implements the similarity-search planner, runner adapter, reducer, and
|
||||
comparison against the local CLI. CTX-09 persists the final artifact and job
|
||||
state. CTX-10 defines graph-specific triangular block plans; it must not reuse
|
||||
the search shard scheme without its pair-coverage invariants.
|
||||
@@ -0,0 +1,584 @@
|
||||
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`, `max_rows`) MUST precede the `file` part: the file is streamed, not
|
||||
buffered, so the fields have to be parsed before it arrives. Currently
|
||||
only diagnostic `similarity-search` with `parameters.query_smiles` is
|
||||
accepted; distributed graph planning is not implemented.
|
||||
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" }
|
||||
|
||||
/jobs/{job_id}/cancel:
|
||||
post:
|
||||
tags: [jobs]
|
||||
summary: Cancel a job and invalidate all unfinished task leases
|
||||
parameters:
|
||||
- $ref: "#/components/parameters/JobID"
|
||||
responses:
|
||||
"200":
|
||||
description: The job is cancelled. Completed and terminally failed tasks remain unchanged.
|
||||
content:
|
||||
application/json:
|
||||
schema: { $ref: "#/components/schemas/CancelJobResponse" }
|
||||
"400": { $ref: "#/components/responses/BadRequest" }
|
||||
"401": { $ref: "#/components/responses/Unauthorized" }
|
||||
"404": { $ref: "#/components/responses/NotFound" }
|
||||
"409": { $ref: "#/components/responses/Conflict" }
|
||||
|
||||
/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]
|
||||
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, enum: [similarity-search], example: similarity-search }
|
||||
parameters:
|
||||
type: string
|
||||
description: JSON object, sent as a string form field.
|
||||
example: '{"query_smiles":"CCO","top_k":10}'
|
||||
chunk_rows:
|
||||
type: integer
|
||||
description: Data rows per shard. Default 1000.
|
||||
example: 1000
|
||||
max_rows:
|
||||
type: integer
|
||||
minimum: 1
|
||||
description: Optional leading data-row limit for a small pipeline check.
|
||||
example: 500
|
||||
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 }
|
||||
|
||||
CancelJobResponse:
|
||||
type: object
|
||||
properties:
|
||||
job_id: { type: string, format: uuid }
|
||||
status: { type: string, enum: [cancelled] }
|
||||
cancelled_tasks: { type: integer }
|
||||
|
||||
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 }
|
||||
cancelled: { type: integer }
|
||||
|
||||
ClaimRequest:
|
||||
type: object
|
||||
required: [worker_id]
|
||||
properties:
|
||||
worker_id: { type: string, format: uuid, description: Registered worker identity. }
|
||||
capabilities:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: Accepted for compatibility only; registration capabilities decide eligibility.
|
||||
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, running, completed, failed, cancelled]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user