package repository import ( "context" "encoding/json" "errors" "fmt" "hash/fnv" "os" "strconv" "strings" "time" "evobgp/internal/store" "github.com/google/uuid" "github.com/jackc/pgx/v5" ) // jobEnqueueGlobalLock serializes the global pending-job count. It is not held // while idempotency is resolved; that uses jobEnqueueTenantLock. const jobEnqueueGlobalLock int64 = 38042341 const jobColumns = `id::text, tenant_id::text, kind, status, idempotency_key, module_id::text, meta_json, created_at, started_at, finished_at, error_message, COALESCE(owner_id,''), fencing_token, attempt, cancel_requested, lease_until, request_sequence, COALESCE(request_hash,'')` func scanJob(row pgx.Row) (store.JobRecord, error) { var j store.JobRecord var meta []byte err := row.Scan(&j.ID, &j.TenantID, &j.Kind, &j.Status, &j.IdempotencyKey, &j.ModuleID, &meta, &j.CreatedAt, &j.StartedAt, &j.FinishedAt, &j.Error, &j.OwnerID, &j.FencingToken, &j.Attempt, &j.CancelRequested, &j.LeaseUntil, &j.RequestSequence, &j.RequestHash) if err != nil { return j, err } if err = json.Unmarshal(meta, &j.Meta); err != nil { return j, fmt.Errorf("repository: decode job metadata: %w", err) } return j, nil } func jobLimit(env string, fallback int) int { n, e := strconv.Atoi(os.Getenv(env)) if e == nil && n > 0 { return n } return fallback } func rollbackWithTimeout(tx pgx.Tx) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = tx.Rollback(ctx) } func (p *Postgres) CountJobsByStatus(ctx context.Context) (map[string]int64, error) { rows, err := p.pool.Query(ctx, `SELECT status,count(*) FROM job_audit GROUP BY status`) if err != nil { return nil, err } defer rows.Close() out := map[string]int64{} for rows.Next() { var status string var n int64 if err := rows.Scan(&status, &n); err != nil { return nil, err } out[status] = n } return out, rows.Err() } func (p *Postgres) FindJobByIdempotencyKey(ctx context.Context, tenant, key string) (store.JobRecord, error) { j, err := scanJob(p.pool.QueryRow(ctx, `SELECT `+jobColumns+` FROM job_audit WHERE tenant_id=$1::uuid AND idempotency_key=$2`, tenant, key)) if errors.Is(err, pgx.ErrNoRows) { return j, store.ErrJobNotFound } return j, err } // EnqueueJob inserts exactly one immutable request and returns the committed row. func (p *Postgres) EnqueueJob(ctx context.Context, in store.JobEnqueue) (store.JobRecord, bool, error) { h, err := store.JobRequestHash(in) if err != nil { return store.JobRecord{}, false, fmt.Errorf("repository: hash job: %w", err) } meta, err := json.Marshal(in.Meta) if err != nil { return store.JobRecord{}, false, fmt.Errorf("repository: encode job: %w", err) } tx, err := p.pool.Begin(ctx) if err != nil { return store.JobRecord{}, false, err } defer rollbackWithTimeout(tx) if err = lockJobEnqueue(ctx, tx, jobEnqueueTenantLock(in.TenantID)); err != nil { return store.JobRecord{}, false, err } if in.IdempotencyKey != nil && *in.IdempotencyKey != "" { var existingHash string j, e := scanJob(tx.QueryRow(ctx, `SELECT `+jobColumns+` FROM job_audit WHERE tenant_id=$1::uuid AND idempotency_key=$2 FOR UPDATE`, in.TenantID, *in.IdempotencyKey)) if e == nil { if err = tx.QueryRow(ctx, `SELECT COALESCE(request_hash,'') FROM job_audit WHERE id=$1::uuid`, j.ID).Scan(&existingHash); err != nil { return store.JobRecord{}, false, err } if existingHash != h { return store.JobRecord{}, false, store.ErrJobConflict } if err = tx.Commit(ctx); err != nil { return store.JobRecord{}, false, err } return j, false, nil } if !errors.Is(e, pgx.ErrNoRows) { return store.JobRecord{}, false, e } } var tenantN int if err = tx.QueryRow(ctx, `SELECT count(*)::int FROM job_audit WHERE status IN ('queued','running') AND tenant_id=$1::uuid`, in.TenantID).Scan(&tenantN); err != nil { return store.JobRecord{}, false, err } if tenantN >= jobLimit("EVOBGP_JOB_MAX_TENANT_PENDING", 1000) { return store.JobRecord{}, false, store.ErrJobQueueFull } if err = lockJobEnqueue(ctx, tx, jobEnqueueGlobalLock); err != nil { return store.JobRecord{}, false, err } var global int if err = tx.QueryRow(ctx, `SELECT count(*)::int FROM job_audit WHERE status IN ('queued','running')`).Scan(&global); err != nil { return store.JobRecord{}, false, err } if global >= jobLimit("EVOBGP_JOB_MAX_PENDING", 10000) { return store.JobRecord{}, false, store.ErrJobQueueFull } id := uuid.NewString() var idem, mod any if in.IdempotencyKey != nil && *in.IdempotencyKey != "" { idem = *in.IdempotencyKey } if in.ModuleID != nil && *in.ModuleID != "" { mod = *in.ModuleID } j, err := scanJob(tx.QueryRow(ctx, `INSERT INTO job_audit(id,tenant_id,kind,status,idempotency_key,module_id,meta_json,request_hash,request_json) VALUES($1::uuid,$2::uuid,$3,'queued',$4,$5::uuid,$6::jsonb,$7,$6::jsonb) RETURNING `+jobColumns, id, in.TenantID, in.Kind, idem, mod, meta, h)) if err != nil { return store.JobRecord{}, false, err } if err = tx.Commit(ctx); err != nil { return store.JobRecord{}, false, err } return j, true, nil } // ClaimJobs atomically acquires available rows with a fresh fencing token. func (p *Postgres) ClaimJobs(ctx context.Context, owner string, limit int, lease time.Duration) ([]store.JobRecord, error) { if limit <= 0 { return nil, nil } if limit > 64 { limit = 64 } if lease <= 0 { lease = time.Minute } rows, err := p.pool.Query(ctx, `WITH picked AS ( SELECT id AS picked_id FROM job_audit WHERE ((status='queued' AND available_at<=now()) OR (status='running' AND lease_until=3`); err != nil { return fmt.Errorf("repository: reap exhausted jobs: %w", err) } return nil } // HeartbeatJob extends a live lease and reports a cross-replica cancel request. func (p *Postgres) HeartbeatJob(ctx context.Context, id, owner string, token int64, lease time.Duration) (bool, error) { if lease <= 0 { lease = time.Minute } var cancelled bool err := p.pool.QueryRow(ctx, `UPDATE job_audit SET lease_until=clock_timestamp()+$4::bigint * interval '1 millisecond' WHERE id=$1::uuid AND owner_id=$2 AND fencing_token=$3 AND status='running' AND lease_until>clock_timestamp() RETURNING cancel_requested`, id, owner, token, lease.Milliseconds()).Scan(&cancelled) if errors.Is(err, pgx.ErrNoRows) { return false, store.ErrJobLeaseLost } return cancelled, err } // CheckpointJob persists mutable progress without altering the immutable request. func (p *Postgres) CheckpointJob(ctx context.Context, j store.JobRecord) error { meta, err := json.Marshal(j.Meta) if err != nil { return err } tag, err := p.pool.Exec(ctx, `UPDATE job_audit SET meta_json=$4::jsonb WHERE id=$1::uuid AND owner_id=$2 AND fencing_token=$3 AND status='running' AND lease_until>clock_timestamp()`, j.ID, j.OwnerID, j.FencingToken, meta) if err != nil { return err } if tag.RowsAffected() != 1 { return store.ErrJobLeaseLost } return nil } // FinishJob commits terminal status only for the current live lease owner. func (p *Postgres) FinishJob(ctx context.Context, j store.JobRecord) error { meta, err := json.Marshal(j.Meta) if err != nil { return err } tag, err := p.pool.Exec(ctx, `UPDATE job_audit SET status=CASE WHEN cancel_requested THEN 'cancelled' ELSE $4 END,error_message=CASE WHEN cancel_requested THEN NULL ELSE $5 END,finished_at=now(),meta_json=$6::jsonb,owner_id=NULL,lease_until=NULL WHERE id=$1::uuid AND owner_id=$2 AND fencing_token=$3 AND status='running' AND lease_until>clock_timestamp()`, j.ID, j.OwnerID, j.FencingToken, j.Status, j.Error, meta) if err != nil { return err } if tag.RowsAffected() != 1 { return store.ErrJobLeaseLost } return nil } // ReleaseJob returns an owned running job to the queue during graceful shutdown. func (p *Postgres) ReleaseJob(ctx context.Context, id, owner string, token int64) error { tag, err := p.pool.Exec(ctx, `UPDATE job_audit SET status=CASE WHEN cancel_requested THEN 'cancelled' ELSE 'queued' END,finished_at=CASE WHEN cancel_requested THEN now() ELSE NULL END,owner_id=NULL,lease_until=NULL,available_at=now(),attempt=GREATEST(attempt-1,0) WHERE id=$1::uuid AND owner_id=$2 AND fencing_token=$3 AND status='running'`, id, owner, token) if err != nil { return err } if tag.RowsAffected() != 1 { return store.ErrJobLeaseLost } return nil } // GetJob returns a tenant-scoped job from shared storage. func (p *Postgres) GetJob(ctx context.Context, tenant, id string) (store.JobRecord, error) { if _, err := uuid.Parse(id); err != nil { return store.JobRecord{}, store.ErrJobNotFound } j, err := scanJob(p.pool.QueryRow(ctx, `SELECT `+jobColumns+` FROM job_audit WHERE id=$1::uuid AND tenant_id=$2::uuid`, id, tenant)) if errors.Is(err, pgx.ErrNoRows) { return j, store.ErrJobNotFound } return j, err } // ListJobs returns a bounded shared page in creation order. func (p *Postgres) ListJobs(ctx context.Context, tenant, status, kind, cursor string, limit int) ([]store.JobRecord, string, bool, error) { if limit <= 0 || limit > 100 { limit = 50 } off, _ := strconv.Atoi(cursor) if off < 0 { off = 0 } rows, err := p.pool.Query(ctx, `SELECT `+jobColumns+` FROM job_audit WHERE tenant_id=$1::uuid AND ($2='' OR status=$2) AND ($3='' OR kind=$3) ORDER BY created_at DESC,request_sequence DESC,id DESC OFFSET $4 LIMIT $5`, tenant, status, kind, off, limit+1) if err != nil { return nil, "", false, err } defer rows.Close() out := make([]store.JobRecord, 0, limit+1) for rows.Next() { j, e := scanJob(rows) if e != nil { return nil, "", false, e } out = append(out, j) } if err = rows.Err(); err != nil { return nil, "", false, err } more := len(out) > limit if more { out = out[:limit] } next := "" if more { next = strconv.Itoa(off + limit) } return out, next, more, nil } // CancelJob atomically cancels queued work or signals its active owner. func (p *Postgres) CancelJob(ctx context.Context, tenant, id string) (store.JobRecord, error) { j, err := scanJob(p.pool.QueryRow(ctx, `UPDATE job_audit SET status=CASE WHEN status='queued' THEN 'cancelled' ELSE status END,finished_at=CASE WHEN status='queued' THEN now() ELSE finished_at END,cancel_requested=CASE WHEN status='running' THEN true ELSE cancel_requested END WHERE id=$1::uuid AND tenant_id=$2::uuid RETURNING `+jobColumns, id, tenant)) if errors.Is(err, pgx.ErrNoRows) { return j, store.ErrJobNotFound } return j, err } // RetryJob explicitly schedules a failed operation under its original ID. func (p *Postgres) RetryJob(ctx context.Context, tenant, id string) (store.JobRecord, error) { tx, err := p.pool.Begin(ctx) if err != nil { return store.JobRecord{}, err } defer rollbackWithTimeout(tx) if err = lockJobEnqueue(ctx, tx, jobEnqueueTenantLock(tenant)); err != nil { return store.JobRecord{}, err } old, err := scanJob(tx.QueryRow(ctx, `SELECT `+jobColumns+` FROM job_audit WHERE id=$1::uuid AND tenant_id=$2::uuid FOR UPDATE`, id, tenant)) if errors.Is(err, pgx.ErrNoRows) { return store.JobRecord{}, store.ErrJobNotFound } if err != nil { return store.JobRecord{}, err } if old.Status != "failed" { return store.JobRecord{}, store.ErrJobConflict } var pending int if err = tx.QueryRow(ctx, `SELECT count(*)::int FROM job_audit WHERE status IN ('queued','running') AND tenant_id=$1::uuid`, tenant).Scan(&pending); err != nil { return store.JobRecord{}, err } if pending >= jobLimit("EVOBGP_JOB_MAX_TENANT_PENDING", 1000) { return store.JobRecord{}, store.ErrJobQueueFull } if err = lockJobEnqueue(ctx, tx, jobEnqueueGlobalLock); err != nil { return store.JobRecord{}, err } var global int if err = tx.QueryRow(ctx, `SELECT count(*)::int FROM job_audit WHERE status IN ('queued','running')`).Scan(&global); err != nil { return store.JobRecord{}, err } if global >= jobLimit("EVOBGP_JOB_MAX_PENDING", 10000) { return store.JobRecord{}, store.ErrJobQueueFull } j, err := scanJob(tx.QueryRow(ctx, `UPDATE job_audit SET status='queued',available_at=now(),started_at=NULL,finished_at=NULL,error_message=NULL,cancel_requested=false,attempt=0,owner_id=NULL,lease_until=NULL,meta_json=COALESCE(request_json,meta_json) WHERE id=$1::uuid AND tenant_id=$2::uuid AND status='failed' RETURNING `+jobColumns, id, tenant)) if errors.Is(err, pgx.ErrNoRows) { return j, store.ErrJobConflict } if err != nil { return j, err } return j, tx.Commit(ctx) } func lockJobEnqueue(ctx context.Context, tx pgx.Tx, key int64) error { if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1::bigint)`, key); err != nil { return fmt.Errorf("repository: lock job enqueue: %w", err) } return nil } // jobEnqueueTenantLock is the per-tenant advisory key for idempotency and the tenant pending limit. func jobEnqueueTenantLock(tenantID string) int64 { h := fnv.New64a() _, _ = h.Write([]byte("evobgp-job-enqueue-v1\x00")) _, _ = h.Write([]byte(strings.TrimSpace(tenantID))) key := int64(h.Sum64()) if key == jobEnqueueGlobalLock { key ^= 0x5a17 } return key }