package-updater/internal/master/store/jobs.go

package store

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"

	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/protocol"
)

const lostJobMessage = "lost: the agent did not know this job when it reconnected"

// failActiveJobsSQL fails every queued or running job of agent $1 with
// message $2 at time $3.
const failActiveJobsSQL = `
UPDATE update_jobs SET status = 'failed', error = $2, finished_at = $3
WHERE agent_id = $1 AND status IN ('queued', 'running')
RETURNING id`

const jobSelectSQL = `
SELECT j.id, j.service_id, s.kind, s.compose_project, s.compose_service, s.swarm_service, s.container_name,
       j.agent_id, COALESCE(a.name, ''), j.from_ref, j.to_ref, j.trigger, j.requested_by, j.status,
       %s, j.output_truncated, j.error, j.created_at, j.started_at, j.finished_at
FROM update_jobs j
JOIN services s ON s.id = j.service_id
LEFT JOIN agents a ON a.id = j.agent_id`

func scanJob(row pgx.Row) (model.Job, error) {
	var j model.Job
	var svc model.Service
	var kind, trigger, status string
	err := row.Scan(&j.ID, &j.ServiceID, &kind, &svc.ComposeProject, &svc.ComposeService, &svc.SwarmService, &svc.ContainerName,
		&j.AgentID, &j.AgentName, &j.FromRef, &j.ToRef, &trigger, &j.RequestedBy, &status,
		&j.Output, &j.OutputTruncated, &j.Error, &j.CreatedAt, &j.StartedAt, &j.FinishedAt)
	svc.Kind = protocol.Kind(kind)
	j.ServiceName = svc.DisplayName()
	j.Trigger = model.JobTrigger(trigger)
	j.Status = model.JobStatus(status)
	return j, err
}

func (s *Store) getJob(ctx context.Context, id int64) (model.Job, error) {
	j, err := scanJob(s.pool.QueryRow(ctx, fmt.Sprintf(jobSelectSQL, "j.output")+` WHERE j.id = $1`, id))
	if errors.Is(err, pgx.ErrNoRows) {
		return model.Job{}, model.ErrNotFound
	}
	if err != nil {
		return model.Job{}, fmt.Errorf("get job: %w", err)
	}
	return j, nil
}

// CreateJob inserts a queued job. If the service already has a queued or
// running job, it returns that job and model.ErrJobActive (detected via the
// unique index, so it holds under concurrent calls). A job for a revoked
// agent is refused with model.ErrAgentOffline: the agent's revoke fails its
// jobs once, so a job created afterwards would never be settled and would
// block the service. The share lock on the agent row makes a concurrent
// revoke wait for this insert, so its job cleanup sees the new row.
func (s *Store) CreateJob(ctx context.Context, j model.NewJob) (model.Job, error) {
	// The active job can finish between the unique violation and the lookup
	// below; then the insert is simply tried once more.
	for attempt := 0; attempt < 2; attempt++ {
		var id int64
		err := s.pool.QueryRow(ctx, `
			INSERT INTO update_jobs (service_id, agent_id, from_ref, to_ref, trigger, requested_by, status, created_at)
			SELECT $1::bigint, $2::bigint, $3::text, $4::text, $5::text, $6::text, 'queued', $7::timestamptz
			WHERE $2::bigint IS NULL
			   OR EXISTS (SELECT 1 FROM agents WHERE id = $2 AND revoked_at IS NULL FOR SHARE)
			RETURNING id`,
			j.ServiceID, j.AgentID, j.FromRef, j.ToRef, string(j.Trigger), j.RequestedBy, j.CreatedAt).Scan(&id)
		if errors.Is(err, pgx.ErrNoRows) {
			return model.Job{}, model.ErrAgentOffline
		}
		var pgErr *pgconn.PgError
		if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "update_jobs_one_active" {
			active, err := scanJob(s.pool.QueryRow(ctx, fmt.Sprintf(jobSelectSQL, "j.output")+
				` WHERE j.service_id = $1 AND j.status IN ('queued', 'running')`, j.ServiceID))
			if errors.Is(err, pgx.ErrNoRows) {
				continue
			}
			if err != nil {
				return model.Job{}, fmt.Errorf("find active job: %w", err)
			}
			return active, model.ErrJobActive
		}
		if err != nil {
			return model.Job{}, fmt.Errorf("create job: %w", err)
		}
		return s.getJob(ctx, id)
	}
	return model.Job{}, fmt.Errorf("create job: active job vanished twice")
}

// FailJob sets a queued or running job to failed with msg and finished_at =
// at. Returns false (no error) if the job is no longer active.
func (s *Store) FailJob(ctx context.Context, id int64, msg string, at time.Time) (bool, error) {
	tag, err := s.pool.Exec(ctx, `
		UPDATE update_jobs SET status = 'failed', error = $2, finished_at = $3
		WHERE id = $1 AND status IN ('queued', 'running')`, id, msg, at)
	if err != nil {
		return false, fmt.Errorf("fail job: %w", err)
	}
	return tag.RowsAffected() == 1, nil
}

// GetJob returns one job with its output. Returns model.ErrNotFound if missing.
func (s *Store) GetJob(ctx context.Context, id int64) (model.Job, error) {
	return s.getJob(ctx, id)
}

// ListJobs returns jobs newest first, without Output.
func (s *Store) ListJobs(ctx context.Context, f model.JobFilter) ([]model.Job, error) {
	limit := f.Limit
	if limit <= 0 {
		limit = 100
	}
	rows, err := s.pool.Query(ctx, fmt.Sprintf(jobSelectSQL, "''")+`
		WHERE ($1 = 0 OR j.service_id = $1)
		ORDER BY j.id DESC
		LIMIT $2`, f.ServiceID, limit)
	if err != nil {
		return nil, fmt.Errorf("list jobs: %w", err)
	}
	defer rows.Close()
	var out []model.Job
	for rows.Next() {
		j, err := scanJob(rows)
		if err != nil {
			return nil, fmt.Errorf("scan job: %w", err)
		}
		out = append(out, j)
	}
	return out, rows.Err()
}

// FailLostJobs fails every queued or running job sent to agentID whose ID is
// not in known, with a "lost" error. Returns the IDs it failed.
func (s *Store) FailLostJobs(ctx context.Context, agentID int64, known []int64, at time.Time) ([]int64, error) {
	if known == nil {
		known = []int64{}
	}
	rows, err := s.pool.Query(ctx, `
		UPDATE update_jobs SET status = 'failed', error = $3, finished_at = $4
		WHERE agent_id = $1 AND status IN ('queued', 'running') AND NOT (id = ANY($2))
		RETURNING id`, agentID, known, lostJobMessage, at)
	if err != nil {
		return nil, fmt.Errorf("fail lost jobs: %w", err)
	}
	defer rows.Close()
	var ids []int64
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			return nil, fmt.Errorf("scan lost job: %w", err)
		}
		ids = append(ids, id)
	}
	return ids, rows.Err()
}

// JobProgress appends chunk to the output of a queued or running job sent to
// agentID, keeping the last protocol.MaxJobOutput characters, and moves a
// queued job to running (started_at = at if unset). Returns false if no such
// active job belongs to agentID.
func (s *Store) JobProgress(ctx context.Context, agentID, jobID int64, chunk string, at time.Time) (bool, error) {
	tag, err := s.pool.Exec(ctx, `
		UPDATE update_jobs SET
			status           = 'running',
			started_at       = COALESCE(started_at, $4),
			output_truncated = output_truncated OR length(output) + length($3::text) > $5,
			output           = right(output || $3::text, $5)
		WHERE id = $1 AND agent_id = $2 AND status IN ('queued', 'running')`,
		jobID, agentID, chunk, at, protocol.MaxJobOutput)
	if err != nil {
		return false, fmt.Errorf("job progress: %w", err)
	}
	return tag.RowsAffected() == 1, nil
}

// FinishJob records the agent's result for a queued or running job sent to
// agentID: status, error, output (replacing what progress assembled),
// started_at (the agent's, or the existing one if the agent sent none),
// finished_at; to_ref = f.Ref when set. On success it also sets the
// service's desired_ref = f.Ref and clears its available ref when that is
// the ref the job was created for (ignoring any @digest) and was reported no
// later than the job was created (that version is now running), in the same
// transaction. Returns false if no
// such active job belongs to agentID.
func (s *Store) FinishJob(ctx context.Context, agentID, jobID int64, f model.JobFinish) (bool, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return false, fmt.Errorf("finish job: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	// The agent's clock wins when it sent a start time; otherwise the time the
	// first progress arrived stays.
	var startedAt *time.Time
	if !f.StartedAt.IsZero() {
		startedAt = &f.StartedAt
	}
	var serviceID int64
	var createdAt time.Time
	var requestedRef string
	// requested is the job's to_ref as it was created, before f.Ref overwrites it.
	err = tx.QueryRow(ctx, `
		UPDATE update_jobs j SET
			status = $3, error = $4, output = $5, output_truncated = $6,
			to_ref = CASE WHEN $7::text <> '' THEN $7::text ELSE j.to_ref END,
			started_at = COALESCE($8, j.started_at), finished_at = $9
		FROM (SELECT id, to_ref FROM update_jobs WHERE id = $1) requested
		WHERE j.id = requested.id AND j.agent_id = $2 AND j.status IN ('queued', 'running')
		RETURNING j.service_id, j.created_at, requested.to_ref`,
		jobID, agentID, string(f.Status), f.Error, f.Output, f.OutputTruncated, f.Ref, startedAt, f.FinishedAt).Scan(&serviceID, &createdAt, &requestedRef)
	if errors.Is(err, pgx.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("finish job: %w", err)
	}
	if f.Status == model.JobSucceeded {
		requestedRef, _, _ = strings.Cut(requestedRef, "@")
		if _, err := tx.Exec(ctx, `
			UPDATE services SET desired_ref = $2,
				available_ref    = CASE WHEN available_ref = $4 AND available_at <= $3 THEN '' ELSE available_ref END,
				available_source = CASE WHEN available_ref = $4 AND available_at <= $3 THEN '' ELSE available_source END,
				available_at     = CASE WHEN available_ref = $4 AND available_at <= $3 THEN NULL ELSE available_at END
			WHERE id = $1`, serviceID, f.Ref, createdAt, requestedRef); err != nil {
			return false, fmt.Errorf("finish job: desired ref: %w", err)
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return false, fmt.Errorf("finish job: commit: %w", err)
	}
	return true, nil
}