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,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
|
||||
}
|
||||
Reference in New Issue
Block a user