diff --git a/go.mod b/go.mod index cc390ce..ce126c9 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,7 @@ require ( github.com/containerd/log v0.1.0 // indirect github.com/containerd/platforms v0.2.1 // indirect github.com/cpuguy83/dockercfg v0.3.2 // indirect - github.com/davecgh/go-spew v1.1.1 // indirect + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/distribution/reference v0.6.0 // indirect github.com/docker/go-connections v0.6.0 // indirect @@ -31,7 +31,9 @@ require ( github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/go-ole/go-ole v1.2.6 // indirect + github.com/golang-migrate/migrate/v4 v4.19.1 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect @@ -50,7 +52,7 @@ require ( github.com/moby/term v0.5.2 // indirect github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/image-spec v1.1.1 // indirect - github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect github.com/shirou/gopsutil/v4 v4.26.3 // indirect github.com/sirupsen/logrus v1.9.4 // indirect @@ -60,7 +62,7 @@ require ( github.com/tklauser/numcpus v0.11.0 // indirect github.com/yusufpapurcu/wmi v1.2.4 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 // indirect go.opentelemetry.io/otel v1.41.0 // indirect go.opentelemetry.io/otel/metric v1.41.0 // indirect go.opentelemetry.io/otel/trace v1.41.0 // indirect diff --git a/go.sum b/go.sum index 034343b..70a62e9 100644 --- a/go.sum +++ b/go.sum @@ -27,8 +27,8 @@ github.com/cpuguy83/dockercfg v0.3.2/go.mod h1:sugsbF4//dDlL/i+S+rtpIWp+5h0BHJHf github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= 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/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= @@ -50,11 +50,15 @@ github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/golang-migrate/migrate/v4 v4.19.1 h1:OCyb44lFuQfYXYLx1SCxPZQGU7mcaZ7gH9yH4jSFbBA= +github.com/golang-migrate/migrate/v4 v4.19.1/go.mod h1:CTcgfjxhaUtsLipnLoQRWCrjYXycRz/g5+RWDuYgPrE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= 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/pgerrcode v0.0.0-20220416144525-469b46aa5efa h1:s+4MhCQ6YrzisK6hFJUX53drDT4UsSW3DEhKn0ifuHw= +github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds= 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-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -101,8 +105,9 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= -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/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt9k/+g42oCprj/FisM4qX9L3sZB3upGN2ZU= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= github.com/redis/go-redis/v9 v9.18.0 h1:pMkxYPkEbMPwRdenAzUNyFNrDgHx9U+DrBabWNfSRQs= @@ -136,16 +141,16 @@ github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0= github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 h1:F7Jx+6hwnZ41NSFTO5q4LYDtJRXBf2PD0rNBkeB/lus= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0/go.mod h1:UHB22Z8QsdRDrnAtX4PntOl36ajSxcdUMt1sF7Y6E7Q= go.opentelemetry.io/otel v1.41.0 h1:YlEwVsGAlCvczDILpUXpIpPSL/VPugt7zHThEMLce1c= go.opentelemetry.io/otel v1.41.0/go.mod h1:Yt4UwgEKeT05QbLwbyHXEwhnjxNO6D8L5PQP51/46dE= go.opentelemetry.io/otel/metric v1.41.0 h1:rFnDcs4gRzBcsO9tS8LCpgR0dxg4aaxWlJxCno7JlTQ= go.opentelemetry.io/otel/metric v1.41.0/go.mod h1:xPvCwd9pU0VN8tPZYzDZV/BMj9CM9vs00GuBjeKhJps= -go.opentelemetry.io/otel/sdk v1.35.0 h1:iPctf8iprVySXSKJffSS79eOjl9pvxV9ZqOWT0QejKY= -go.opentelemetry.io/otel/sdk v1.35.0/go.mod h1:+ga1bZliga3DxJ3CQGg3updiaAJoNECOgJREo9KHGQg= -go.opentelemetry.io/otel/sdk/metric v1.35.0 h1:1RriWBmCKgkeHEhM7a2uMjMUfP7MsOF5JpUCaEqEI9o= -go.opentelemetry.io/otel/sdk/metric v1.35.0/go.mod h1:is6XYCUMpcKi+ZsOvfluY5YstFnhW0BidkR+gL+qN+w= +go.opentelemetry.io/otel/sdk v1.36.0 h1:b6SYIuLRs88ztox4EyrvRti80uXIFy+Sqzoh9kFULbs= +go.opentelemetry.io/otel/sdk v1.36.0/go.mod h1:+lC+mTgD+MUWfjJubi2vvXWcVxyr9rmlshZni72pXeY= +go.opentelemetry.io/otel/sdk/metric v1.36.0 h1:r0ntwwGosWGaa0CrSt8cuNuTcccMXERFwHX4dThiPis= +go.opentelemetry.io/otel/sdk/metric v1.36.0/go.mod h1:qTNOhFDfKRwX0yXOqJYegL5WRaW376QbB7P4Pb0qva4= go.opentelemetry.io/otel/trace v1.41.0 h1:Vbk2co6bhj8L59ZJ6/xFTskY+tGAbOnCtQGVVa9TIN0= go.opentelemetry.io/otel/trace v1.41.0/go.mod h1:U1NU4ULCoxeDKc09yCWdWe+3QoyweJcISEVa1RBzOis= go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= diff --git a/internal/store/migrate.go b/internal/store/migrate.go new file mode 100644 index 0000000..444c207 --- /dev/null +++ b/internal/store/migrate.go @@ -0,0 +1,45 @@ +package store + +import ( + "context" + "embed" + "fmt" + + "github.com/golang-migrate/migrate/v4" + pgxdriver "github.com/golang-migrate/migrate/v4/database/pgx/v5" + "github.com/golang-migrate/migrate/v4/source/iofs" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/stdlib" +) + +//go:embed migrations/*.sql +var migrationsFS embed.FS + +func Migrate(ctx context.Context, dsn string) error { + cfg, err := pgx.ParseConfig(dsn) + if err != nil { + return fmt.Errorf("parse dsn: %w", err) + } + db := stdlib.OpenDB(*cfg) + defer func() { _ = db.Close() }() + + driver, err := pgxdriver.WithInstance(db, &pgxdriver.Config{}) + if err != nil { + return fmt.Errorf("driver: %w", err) + } + + src, err := iofs.New(migrationsFS, "migrations") + if err != nil { + return fmt.Errorf("source: %w", err) + } + + m, err := migrate.NewWithInstance("iofs", src, "pgx", driver) + if err != nil { + return fmt.Errorf("migrate: %w", err) + } + + if err := m.Up(); err != nil && err != migrate.ErrNoChange { + return fmt.Errorf("up: %w", err) + } + return nil +} diff --git a/internal/store/migrate_integration_test.go b/internal/store/migrate_integration_test.go new file mode 100644 index 0000000..56bd5ae --- /dev/null +++ b/internal/store/migrate_integration_test.go @@ -0,0 +1,30 @@ +//go:build integration + +package store_test + +import ( + "context" + "testing" + + "github.com/smallchungus/disttaskqueue/internal/store" + "github.com/smallchungus/disttaskqueue/internal/testutil" +) + +func TestMigrate_AppliesInitialSchema(t *testing.T) { + pool := testutil.StartPostgres(t) + + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatalf("migrate: %v", err) + } + + var n int + err := pool.QueryRow(context.Background(), + `SELECT count(*) FROM information_schema.tables WHERE table_schema='public' AND table_name IN ('pipeline_jobs','job_status_history')`, + ).Scan(&n) + if err != nil { + t.Fatal(err) + } + if n != 2 { + t.Fatalf("expected 2 tables, got %d", n) + } +} diff --git a/internal/store/migrations/0001_init.down.sql b/internal/store/migrations/0001_init.down.sql new file mode 100644 index 0000000..bb675b6 --- /dev/null +++ b/internal/store/migrations/0001_init.down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS job_status_history; +DROP TABLE IF EXISTS pipeline_jobs; diff --git a/internal/store/migrations/0001_init.up.sql b/internal/store/migrations/0001_init.up.sql new file mode 100644 index 0000000..38578c8 --- /dev/null +++ b/internal/store/migrations/0001_init.up.sql @@ -0,0 +1,31 @@ +CREATE TABLE pipeline_jobs ( + id UUID PRIMARY KEY, + stage TEXT NOT NULL, + status TEXT NOT NULL, + payload JSONB NOT NULL DEFAULT '{}'::jsonb, + worker_id TEXT, + attempts INT NOT NULL DEFAULT 0, + max_attempts INT NOT NULL DEFAULT 8, + last_error TEXT, + next_run_at TIMESTAMPTZ NOT NULL DEFAULT now(), + claimed_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX pipeline_jobs_stage_status_idx ON pipeline_jobs (stage, status); +CREATE INDEX pipeline_jobs_queued_next_run_idx ON pipeline_jobs (status, next_run_at) WHERE status = 'queued'; + +CREATE TABLE job_status_history ( + id BIGSERIAL PRIMARY KEY, + job_id UUID NOT NULL REFERENCES pipeline_jobs(id) ON DELETE CASCADE, + stage TEXT NOT NULL, + from_status TEXT, + to_status TEXT NOT NULL, + worker_id TEXT, + error TEXT, + at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX job_status_history_job_at_idx ON job_status_history (job_id, at); diff --git a/internal/store/store.go b/internal/store/store.go new file mode 100644 index 0000000..024e7b0 --- /dev/null +++ b/internal/store/store.go @@ -0,0 +1,121 @@ +package store + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +type Store struct { + pool *pgxpool.Pool +} + +func New(pool *pgxpool.Pool) *Store { + return &Store{pool: pool} +} + +var ErrJobNotFound = errors.New("job not found") + +func (s *Store) EnqueueJob(ctx context.Context, nj NewJob) (Job, error) { + id := uuid.New() + payload := nj.Payload + if payload == nil { + payload = json.RawMessage(`{}`) + } + + const q = ` + INSERT INTO pipeline_jobs (id, stage, status, payload) + VALUES ($1, $2, $3, $4) + RETURNING id, stage, status, payload, worker_id, attempts, max_attempts, + last_error, next_run_at, claimed_at, completed_at, created_at, updated_at` + + row := s.pool.QueryRow(ctx, q, id, nj.Stage, StatusQueued, payload) + var j Job + if err := row.Scan(&j.ID, &j.Stage, &j.Status, &j.Payload, &j.WorkerID, &j.Attempts, + &j.MaxAttempts, &j.LastError, &j.NextRunAt, &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt); err != nil { + return Job{}, fmt.Errorf("enqueue: %w", err) + } + return j, nil +} + +var ErrJobNotClaimable = errors.New("job not claimable") + +func (s *Store) ClaimJob(ctx context.Context, id uuid.UUID, workerID string) error { + const q = ` + UPDATE pipeline_jobs + SET status = $1, worker_id = $2, claimed_at = now(), updated_at = now() + WHERE id = $3 AND status = $4 + RETURNING id` + + var got uuid.UUID + err := s.pool.QueryRow(ctx, q, StatusRunning, workerID, id, StatusQueued).Scan(&got) + if errors.Is(err, pgx.ErrNoRows) { + return ErrJobNotClaimable + } + if err != nil { + return fmt.Errorf("claim: %w", err) + } + return nil +} + +func (s *Store) MarkDone(ctx context.Context, id uuid.UUID) error { + const q = ` + UPDATE pipeline_jobs + SET status = $1, completed_at = now(), updated_at = now() + WHERE id = $2` + + tag, err := s.pool.Exec(ctx, q, StatusDone, id) + if err != nil { + return fmt.Errorf("mark done: %w", err) + } + if tag.RowsAffected() == 0 { + return ErrJobNotFound + } + return nil +} + +func (s *Store) MarkFailed(ctx context.Context, id uuid.UUID, errMsg string, nextRunAt time.Time) error { + const q = ` + UPDATE pipeline_jobs + SET attempts = attempts + 1, + status = CASE WHEN attempts + 1 >= max_attempts THEN $1::text ELSE $2::text END, + worker_id = NULL, + last_error = $3, + next_run_at = $4, + updated_at = now() + WHERE id = $5` + + tag, err := s.pool.Exec(ctx, q, StatusDead, StatusQueued, errMsg, nextRunAt, id) + if err != nil { + return fmt.Errorf("mark failed: %w", err) + } + if tag.RowsAffected() == 0 { + return ErrJobNotFound + } + return nil +} + +func (s *Store) GetJob(ctx context.Context, id uuid.UUID) (Job, error) { + const q = ` + SELECT id, stage, status, payload, worker_id, attempts, max_attempts, + last_error, next_run_at, claimed_at, completed_at, created_at, updated_at + FROM pipeline_jobs WHERE id = $1` + + row := s.pool.QueryRow(ctx, q, id) + var j Job + err := row.Scan(&j.ID, &j.Stage, &j.Status, &j.Payload, &j.WorkerID, &j.Attempts, + &j.MaxAttempts, &j.LastError, &j.NextRunAt, &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return Job{}, ErrJobNotFound + } + if err != nil { + return Job{}, fmt.Errorf("get: %w", err) + } + return j, nil +} diff --git a/internal/store/store_integration_test.go b/internal/store/store_integration_test.go new file mode 100644 index 0000000..7587aef --- /dev/null +++ b/internal/store/store_integration_test.go @@ -0,0 +1,162 @@ +//go:build integration + +package store_test + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/smallchungus/disttaskqueue/internal/store" + "github.com/smallchungus/disttaskqueue/internal/testutil" +) + +func newStore(t *testing.T) *store.Store { + t.Helper() + pool := testutil.StartPostgres(t) + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatalf("migrate: %v", err) + } + return store.New(pool) +} + +func TestEnqueueJob_PersistsRow(t *testing.T) { + s := newStore(t) + ctx := context.Background() + + j, err := s.EnqueueJob(ctx, store.NewJob{Stage: "test", Payload: json.RawMessage(`{"k":"v"}`)}) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + if j.ID.String() == "" || j.Stage != "test" || j.Status != store.StatusQueued { + t.Fatalf("unexpected job: %+v", j) + } + if string(j.Payload) != `{"k": "v"}` && string(j.Payload) != `{"k":"v"}` { + t.Fatalf("payload roundtrip: got %s", string(j.Payload)) + } +} + +func TestGetJob_RoundTripsAllFields(t *testing.T) { + s := newStore(t) + ctx := context.Background() + + enq, err := s.EnqueueJob(ctx, store.NewJob{Stage: "test", Payload: json.RawMessage(`{"x":1}`)}) + if err != nil { + t.Fatal(err) + } + + got, err := s.GetJob(ctx, enq.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.ID != enq.ID || got.Stage != enq.Stage || got.Status != enq.Status { + t.Fatalf("mismatch: enq=%+v got=%+v", enq, got) + } +} + +func TestGetJob_ReturnsErrNotFoundForUnknownID(t *testing.T) { + s := newStore(t) + _, err := s.GetJob(context.Background(), uuid.New()) + if !errors.Is(err, store.ErrJobNotFound) { + t.Fatalf("got %v, want ErrJobNotFound", err) + } +} + +func TestClaimJob_TransitionsQueuedToRunning(t *testing.T) { + s := newStore(t) + ctx := context.Background() + j, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "test"}) + + if err := s.ClaimJob(ctx, j.ID, "worker-1"); err != nil { + t.Fatalf("claim: %v", err) + } + + got, _ := s.GetJob(ctx, j.ID) + if got.Status != store.StatusRunning { + t.Fatalf("status: %s", got.Status) + } + if got.WorkerID == nil || *got.WorkerID != "worker-1" { + t.Fatalf("worker_id: %v", got.WorkerID) + } + if got.ClaimedAt == nil { + t.Fatalf("claimed_at: nil") + } +} + +func TestClaimJob_FailsWhenNotQueued(t *testing.T) { + s := newStore(t) + ctx := context.Background() + j, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "test"}) + if err := s.ClaimJob(ctx, j.ID, "worker-1"); err != nil { + t.Fatal(err) + } + + err := s.ClaimJob(ctx, j.ID, "worker-2") + if !errors.Is(err, store.ErrJobNotClaimable) { + t.Fatalf("got %v, want ErrJobNotClaimable", err) + } +} + +func TestMarkDone_SetsTerminalState(t *testing.T) { + s := newStore(t) + ctx := context.Background() + j, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "test"}) + _ = s.ClaimJob(ctx, j.ID, "worker-1") + + if err := s.MarkDone(ctx, j.ID); err != nil { + t.Fatalf("done: %v", err) + } + + got, _ := s.GetJob(ctx, j.ID) + if got.Status != store.StatusDone { + t.Fatalf("status: %s", got.Status) + } + if got.CompletedAt == nil { + t.Fatalf("completed_at: nil") + } +} + +func TestMarkFailed_RequeuesWithIncrementedAttempts(t *testing.T) { + s := newStore(t) + ctx := context.Background() + j, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "test"}) + _ = s.ClaimJob(ctx, j.ID, "worker-1") + + nextRun := time.Now().Add(2 * time.Second) + if err := s.MarkFailed(ctx, j.ID, "transient boom", nextRun); err != nil { + t.Fatalf("fail: %v", err) + } + + got, _ := s.GetJob(ctx, j.ID) + if got.Status != store.StatusQueued { + t.Fatalf("status: %s, want queued", got.Status) + } + if got.Attempts != 1 { + t.Fatalf("attempts: %d, want 1", got.Attempts) + } + if got.LastError == nil || *got.LastError != "transient boom" { + t.Fatalf("last_error: %v", got.LastError) + } + if got.WorkerID != nil { + t.Fatalf("worker_id: %v, want nil after fail", got.WorkerID) + } +} + +func TestMarkFailed_MarksDeadAtMaxAttempts(t *testing.T) { + s := newStore(t) + ctx := context.Background() + j, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "test"}) + + for i := 0; i < 8; i++ { + _ = s.ClaimJob(ctx, j.ID, "worker-1") + _ = s.MarkFailed(ctx, j.ID, "still broken", time.Now()) + } + + got, _ := s.GetJob(ctx, j.ID) + if got.Status != store.StatusDead { + t.Fatalf("status: %s, want dead", got.Status) + } +} diff --git a/internal/store/types.go b/internal/store/types.go new file mode 100644 index 0000000..0eb3ccd --- /dev/null +++ b/internal/store/types.go @@ -0,0 +1,38 @@ +package store + +import ( + "encoding/json" + "time" + + "github.com/google/uuid" +) + +type JobStatus string + +const ( + StatusQueued JobStatus = "queued" + StatusRunning JobStatus = "running" + StatusDone JobStatus = "done" + StatusDead JobStatus = "dead" +) + +type Job struct { + ID uuid.UUID + Stage string + Status JobStatus + Payload json.RawMessage + WorkerID *string + Attempts int + MaxAttempts int + LastError *string + NextRunAt time.Time + ClaimedAt *time.Time + CompletedAt *time.Time + CreatedAt time.Time + UpdatedAt time.Time +} + +type NewJob struct { + Stage string + Payload json.RawMessage +}