feat(coordinator): artifact storage foundation (CTX-05, part 1)
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.
This commit is contained in:
Efremenko Arhip
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,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
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
Reference in New Issue
Block a user