package store import ( "context" "errors" "fmt" "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, 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 err = tx.QueryRow(ctx, ` UPDATE update_jobs SET status = $3, error = $4, output = $5, output_truncated = $6, to_ref = CASE WHEN $7::text <> '' THEN $7::text ELSE to_ref END, started_at = COALESCE($8, started_at), finished_at = $9 WHERE id = $1 AND agent_id = $2 AND status IN ('queued', 'running') RETURNING service_id`, jobID, agentID, string(f.Status), f.Error, f.Output, f.OutputTruncated, f.Ref, startedAt, f.FinishedAt).Scan(&serviceID) 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 { if _, err := tx.Exec(ctx, `UPDATE services SET desired_ref = $2 WHERE id = $1`, serviceID, f.Ref); 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 }