From dbf578c5006bb8bc5f5c6884ad7a3f2463c47a45 Mon Sep 17 00:00:00 2001 From: Efremenko Arhip Date: Thu, 23 Jul 2026 14:06:19 +0300 Subject: [PATCH] feat(coordinator): artifact storage foundation (CTX-05, part 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce durable, coordinator-owned artifacts — the model the master plan requires instead of trusting worker-supplied result URIs. - migration 0003: artifacts table + artifact_kind enum, plus nullable input/result_artifact_id columns on jobs and tasks. - domain.Artifact with NewArtifact/SetContent; storage keys derive from a fresh UUID, never from a client filename (no path traversal). - BlobStore and ArtifactRepository ports. - blob.FSStore: filesystem blob storage that streams while hashing (SHA-256), fsyncs, and atomically renames into place — a failed upload leaves neither a committed artifact nor a staging file. Covered by unit tests. - ArtifactRepo (squirrel) and COORDINATOR_STORAGE_DIR config. HTTP upload/download handlers and the switch of result submission to artifact_id come in the next parts. --- coordinator/internal/domain/artifact.go | 69 +++++++++++ coordinator/internal/domain/errors.go | 17 +-- coordinator/internal/infra/config.go | 3 + coordinator/internal/storage/blob/store.go | 116 ++++++++++++++++++ .../internal/storage/blob/store_test.go | 115 +++++++++++++++++ .../storage/postgres/artifact_repo.go | 72 +++++++++++ coordinator/internal/usecase/ports.go | 17 +++ .../migrations/0003_artifacts.down.sql | 11 ++ coordinator/migrations/0003_artifacts.up.sql | 31 +++++ 9 files changed, 443 insertions(+), 8 deletions(-) create mode 100644 coordinator/internal/domain/artifact.go create mode 100644 coordinator/internal/storage/blob/store.go create mode 100644 coordinator/internal/storage/blob/store_test.go create mode 100644 coordinator/internal/storage/postgres/artifact_repo.go create mode 100644 coordinator/migrations/0003_artifacts.down.sql create mode 100644 coordinator/migrations/0003_artifacts.up.sql diff --git a/coordinator/internal/domain/artifact.go b/coordinator/internal/domain/artifact.go new file mode 100644 index 0000000..43bc219 --- /dev/null +++ b/coordinator/internal/domain/artifact.go @@ -0,0 +1,69 @@ +package domain + +import ( + "time" + + "github.com/google/uuid" +) + +type ArtifactKind string + +const ( + ArtifactInput ArtifactKind = "input" + ArtifactShard ArtifactKind = "shard" + ArtifactPartialResult ArtifactKind = "partial_result" + ArtifactFinalResult ArtifactKind = "final_result" + ArtifactLog ArtifactKind = "log" +) + +// Artifact is a durable file the coordinator owns, described by its metadata. +// The bytes live in blob storage under StorageKey; this struct is what the +// database persists and what every other layer reasons about. +type Artifact struct { + ID uuid.UUID + JobID uuid.UUID + TaskID *uuid.UUID // nil for a job-level input + Kind ArtifactKind + Filename string + StorageKey string + ContentType string + SizeBytes int64 + SHA256 string + CreatedAt time.Time +} + +// NewArtifact begins an artifact record. Size and checksum are unknown until the +// bytes have been streamed to storage, so they are filled in later by SetContent. +// +// StorageKey is derived from a fresh UUID, never from the client-supplied +// filename — that is what stops a "../../etc/passwd" filename from escaping the +// storage directory. +func NewArtifact(jobID uuid.UUID, taskID *uuid.UUID, kind ArtifactKind, + filename, contentType string, now time.Time) (*Artifact, error) { + + if filename == "" || kind == "" { + return nil, ErrInvalidInput + } + if contentType == "" { + contentType = "application/octet-stream" + } + id := uuid.New() + return &Artifact{ + ID: id, + JobID: jobID, + TaskID: taskID, + Kind: kind, + Filename: filename, + StorageKey: id.String(), + ContentType: contentType, + CreatedAt: now, + }, nil +} + +// SetContent records the size and checksum measured while streaming the bytes +// into storage. Both are computed by the coordinator, never trusted from the +// client — the whole point of owning the artifact. +func (a *Artifact) SetContent(sha256 string, size int64) { + a.SHA256 = sha256 + a.SizeBytes = size +} diff --git a/coordinator/internal/domain/errors.go b/coordinator/internal/domain/errors.go index 6603844..d64b1f6 100644 --- a/coordinator/internal/domain/errors.go +++ b/coordinator/internal/domain/errors.go @@ -8,12 +8,13 @@ import "errors" // // 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") - 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") + ErrJobNotFound = errors.New("job not found") + ErrTaskNotFound = errors.New("task not found") + ErrWorkerNotFound = errors.New("worker not found") + ErrArtifactNotFound = errors.New("artifact not found") + ErrLeaseConflict = errors.New("task leased to another worker") + ErrStaleAttempt = errors.New("attempt does not match lease") + ErrResultConflict = errors.New("different result already recorded") + ErrInvalidInput = errors.New("invalid input") + ErrTaskNotLeased = errors.New("task is not currently leased") ) diff --git a/coordinator/internal/infra/config.go b/coordinator/internal/infra/config.go index ce8bc12..b7aa308 100644 --- a/coordinator/internal/infra/config.go +++ b/coordinator/internal/infra/config.go @@ -29,6 +29,8 @@ type Config struct { LogLevel string // Path to a rotated log file. Empty logs to stdout only. LogFile string + // Directory where artifact bytes are stored. + StorageDir string // Connection pool upper bound. DBMaxConns int32 @@ -73,6 +75,7 @@ func LoadConfig() (Config, error) { Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")), LogLevel: getEnv("LOG_LEVEL", "info"), LogFile: os.Getenv("LOG_FILE"), + StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"), DBMaxConns: 10, DBConnectTimeout: 30 * time.Second, RequestTimeout: 15 * time.Second, diff --git a/coordinator/internal/storage/blob/store.go b/coordinator/internal/storage/blob/store.go new file mode 100644 index 0000000..b7efafd --- /dev/null +++ b/coordinator/internal/storage/blob/store.go @@ -0,0 +1,116 @@ +// 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, 0o755); 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 + } + f, err := os.Open(filepath.Join(s.dir, key)) + if err != nil { + return nil, err + } + return f, nil +} + +// checkKey rejects anything that could escape the storage directory. Keys are +// coordinator-generated UUIDs, so this is defence in depth, not the only guard. +func checkKey(key string) error { + if key == "" || strings.ContainsAny(key, `/\`) || strings.Contains(key, "..") { + return fmt.Errorf("invalid storage key %q", key) + } + return nil +} + +// ctxReader aborts a copy when the request context is cancelled, so a stalled +// or disconnected upload does not tie up a file handle indefinitely. +type ctxReader struct { + ctx context.Context + r io.Reader +} + +func (c *ctxReader) Read(p []byte) (int, error) { + if err := c.ctx.Err(); err != nil { + return 0, err + } + return c.r.Read(p) +} diff --git a/coordinator/internal/storage/blob/store_test.go b/coordinator/internal/storage/blob/store_test.go new file mode 100644 index 0000000..5f57294 --- /dev/null +++ b/coordinator/internal/storage/blob/store_test.go @@ -0,0 +1,115 @@ +package blob + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "io" + "os" + "path/filepath" + "strings" + "testing" +) + +func newStore(t *testing.T) *FSStore { + t.Helper() + s, err := NewFSStore(t.TempDir()) + if err != nil { + t.Fatalf("NewFSStore: %v", err) + } + return s +} + +func TestPutComputesChecksumAndSize(t *testing.T) { + s := newStore(t) + data := bytes.Repeat([]byte("chembl-row\n"), 10000) // ~110 KB, streamed + + sum, size, err := s.Put(context.Background(), "key-1", bytes.NewReader(data)) + if err != nil { + t.Fatalf("Put: %v", err) + } + + want := sha256.Sum256(data) + if sum != hex.EncodeToString(want[:]) { + t.Errorf("sha256 = %s, want %s", sum, hex.EncodeToString(want[:])) + } + if size != int64(len(data)) { + t.Errorf("size = %d, want %d", size, len(data)) + } +} + +func TestPutThenOpenRoundTrips(t *testing.T) { + s := newStore(t) + data := []byte("partial result csv\n1,2,3\n") + + if _, _, err := s.Put(context.Background(), "key-2", bytes.NewReader(data)); err != nil { + t.Fatalf("Put: %v", err) + } + + rc, err := s.Open(context.Background(), "key-2") + if err != nil { + t.Fatalf("Open: %v", err) + } + defer rc.Close() + + got, _ := io.ReadAll(rc) + if !bytes.Equal(got, data) { + t.Errorf("round-trip mismatch: got %q", got) + } +} + +func TestPutLeavesNoStagingFileBehind(t *testing.T) { + s := newStore(t) + if _, _, err := s.Put(context.Background(), "key-3", strings.NewReader("x")); err != nil { + t.Fatalf("Put: %v", err) + } + + entries, _ := os.ReadDir(s.staging) + if len(entries) != 0 { + t.Errorf("staging dir not empty after a successful put: %v", entries) + } +} + +func TestPutFailureLeavesNoArtifactOrStaging(t *testing.T) { + s := newStore(t) + // A reader that errors partway through simulates a dropped upload. + r := io.MultiReader(strings.NewReader("half"), &erroringReader{}) + + if _, _, err := s.Put(context.Background(), "key-4", r); err == nil { + t.Fatal("expected an error from a failing reader") + } + + if _, err := os.Stat(filepath.Join(s.dir, "key-4")); !os.IsNotExist(err) { + t.Error("a failed put must not leave a committed artifact") + } + if entries, _ := os.ReadDir(s.staging); len(entries) != 0 { + t.Errorf("a failed put must not leave staging files: %v", entries) + } +} + +func TestPutRejectsUnsafeKeys(t *testing.T) { + s := newStore(t) + for _, key := range []string{"", "../escape", "a/b", `a\b`, "with..dots"} { + if _, _, err := s.Put(context.Background(), key, strings.NewReader("x")); err == nil { + t.Errorf("key %q should have been rejected", key) + } + } +} + +func TestPutHonoursContextCancellation(t *testing.T) { + s := newStore(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() // already cancelled before the copy starts + + if _, _, err := s.Put(ctx, "key-5", strings.NewReader("data")); err == nil { + t.Fatal("expected cancellation to abort the put") + } + if _, err := os.Stat(filepath.Join(s.dir, "key-5")); !os.IsNotExist(err) { + t.Error("a cancelled put must not leave an artifact") + } +} + +type erroringReader struct{} + +func (*erroringReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } diff --git a/coordinator/internal/storage/postgres/artifact_repo.go b/coordinator/internal/storage/postgres/artifact_repo.go new file mode 100644 index 0000000..771b5ee --- /dev/null +++ b/coordinator/internal/storage/postgres/artifact_repo.go @@ -0,0 +1,72 @@ +package postgres + +import ( + "context" + "errors" + "fmt" + + sq "github.com/Masterminds/squirrel" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// ArtifactRepo implements usecase.ArtifactRepository. +type ArtifactRepo struct { + pool *pgxpool.Pool +} + +func NewArtifactRepo(pool *pgxpool.Pool) *ArtifactRepo { + return &ArtifactRepo{pool: pool} +} + +var _ usecase.ArtifactRepository = (*ArtifactRepo)(nil) + +var artifactColumns = []string{ + "id", "job_id", "task_id", "kind", "filename", "storage_key", + "content_type", "size_bytes", "sha256", "created_at", +} + +func (r *ArtifactRepo) Insert(ctx context.Context, a *domain.Artifact) error { + sql, args, err := psql.Insert("artifacts"). + Columns(artifactColumns...). + Values(a.ID, a.JobID, a.TaskID, string(a.Kind), a.Filename, a.StorageKey, + a.ContentType, a.SizeBytes, a.SHA256, a.CreatedAt). + ToSql() + if err != nil { + return err + } + if _, err := conn(ctx, r.pool).Exec(ctx, sql, args...); err != nil { + return fmt.Errorf("insert artifact: %w", err) + } + return nil +} + +func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error) { + sql, args, err := psql.Select(artifactColumns...). + From("artifacts"). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return nil, err + } + + var ( + a domain.Artifact + kind string + ) + err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( + &a.ID, &a.JobID, &a.TaskID, &kind, &a.Filename, &a.StorageKey, + &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return nil, domain.ErrArtifactNotFound + } + if err != nil { + return nil, fmt.Errorf("get artifact: %w", err) + } + a.Kind = domain.ArtifactKind(kind) + return &a, nil +} diff --git a/coordinator/internal/usecase/ports.go b/coordinator/internal/usecase/ports.go index 183141b..747bcef 100644 --- a/coordinator/internal/usecase/ports.go +++ b/coordinator/internal/usecase/ports.go @@ -9,6 +9,7 @@ package usecase import ( "context" "errors" + "io" "time" "github.com/google/uuid" @@ -68,6 +69,22 @@ type WorkerRepository interface { Get(ctx context.Context, id uuid.UUID) (*domain.Worker, error) } +// ArtifactRepository persists artifact metadata. The bytes live in a BlobStore; +// this keeps only the record that points at them. +type ArtifactRepository interface { + Insert(ctx context.Context, a *domain.Artifact) error + Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error) +} + +// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams +// in both directions so a large shard never has to sit in memory, and reports +// the checksum and size it measured while writing — the coordinator's own +// numbers, not the client's claim. +type BlobStore interface { + Put(ctx context.Context, key string, r io.Reader) (sha256 string, size int64, err error) + Open(ctx context.Context, key string) (io.ReadCloser, error) +} + // 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. diff --git a/coordinator/migrations/0003_artifacts.down.sql b/coordinator/migrations/0003_artifacts.down.sql new file mode 100644 index 0000000..54900e0 --- /dev/null +++ b/coordinator/migrations/0003_artifacts.down.sql @@ -0,0 +1,11 @@ +BEGIN; + +ALTER TABLE tasks DROP COLUMN IF EXISTS input_artifact_id; +ALTER TABLE tasks DROP COLUMN IF EXISTS result_artifact_id; +ALTER TABLE jobs DROP COLUMN IF EXISTS input_artifact_id; +ALTER TABLE jobs DROP COLUMN IF EXISTS result_artifact_id; + +DROP TABLE IF EXISTS artifacts; +DROP TYPE IF EXISTS artifact_kind; + +COMMIT; diff --git a/coordinator/migrations/0003_artifacts.up.sql b/coordinator/migrations/0003_artifacts.up.sql new file mode 100644 index 0000000..103b646 --- /dev/null +++ b/coordinator/migrations/0003_artifacts.up.sql @@ -0,0 +1,31 @@ +BEGIN; + +CREATE TYPE artifact_kind AS ENUM ('input','shard','partial_result','final_result','log'); + +-- A durable file the coordinator owns: input, shard, partial/final result, log. +-- The database is the source of truth; files are found through this metadata, +-- never by scanning directories. +CREATE TABLE artifacts ( + id uuid PRIMARY KEY, + job_id uuid NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, + task_id uuid REFERENCES tasks(id) ON DELETE CASCADE, -- null for job-level inputs + kind artifact_kind NOT NULL, + filename text NOT NULL, + storage_key text NOT NULL UNIQUE, -- coordinator-generated, never a client path + content_type text NOT NULL DEFAULT 'application/octet-stream', + size_bytes bigint NOT NULL CHECK (size_bytes >= 0), + sha256 text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now() +); + +CREATE INDEX ix_artifacts_job ON artifacts (job_id); +CREATE INDEX ix_artifacts_task ON artifacts (task_id); + +-- Jobs and tasks reference their artifacts. Nullable during the transition from +-- URI-based inputs/results to artifact-based ones. +ALTER TABLE jobs ADD COLUMN input_artifact_id uuid REFERENCES artifacts(id); +ALTER TABLE jobs ADD COLUMN result_artifact_id uuid REFERENCES artifacts(id); +ALTER TABLE tasks ADD COLUMN input_artifact_id uuid REFERENCES artifacts(id); +ALTER TABLE tasks ADD COLUMN result_artifact_id uuid REFERENCES artifacts(id); + +COMMIT;