EvoBGP/internal/repository/job_queue.go

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<clock_timestamp() AND NOT cancel_requested)) AND attempt<3
	 ORDER BY available_at,created_at,request_sequence,id FOR UPDATE SKIP LOCKED LIMIT $1
	) UPDATE job_audit j SET status='running',owner_id=$2,lease_until=clock_timestamp()+$3::bigint * interval '1 millisecond',attempt=j.attempt+1,fencing_token=j.fencing_token+1,started_at=COALESCE(j.started_at,now()) FROM picked WHERE j.id=picked.picked_id RETURNING `+jobColumns, limit, owner, lease.Milliseconds())
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make([]store.JobRecord, 0, limit)
	for rows.Next() {
		j, e := scanJob(rows)
		if e != nil {
			return nil, e
		}
		out = append(out, j)
	}
	return out, rows.Err()
}

// ReapExpiredJobs marks orphaned cancellations and exhausted retries terminal.
// It is separate from ClaimJobs so a once-per-second claim does not rewrite job_audit.
func (p *Postgres) ReapExpiredJobs(ctx context.Context) error {
	if _, err := p.pool.Exec(ctx, `UPDATE job_audit SET status='cancelled',finished_at=now(),owner_id=NULL,lease_until=NULL WHERE status='running' AND lease_until<clock_timestamp() AND cancel_requested`); err != nil {
		return fmt.Errorf("repository: reap cancelled jobs: %w", err)
	}
	if _, err := p.pool.Exec(ctx, `UPDATE job_audit SET status='failed',finished_at=now(),error_message='retry limit reached',owner_id=NULL,lease_until=NULL WHERE ((status='running' AND lease_until<clock_timestamp()) OR status='queued') AND attempt>=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
}