ITWorx Pulse release export
bd774932d5
Public source validation / validate (push) Failing after 3m8s
273 lines
10 KiB
Go
273 lines
10 KiB
Go
package workerruntime
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
// ErrLeaseLost reports that the job run row this worker owned was completed or
|
|
// reclaimed elsewhere, so this worker's result must not overwrite it.
|
|
var ErrLeaseLost = errors.New("worker job lease was lost")
|
|
|
|
// Lease identifies one claimed scheduling window.
|
|
type Lease struct {
|
|
ID string
|
|
JobType string
|
|
JobKey string
|
|
ScheduledAt time.Time
|
|
Owner string
|
|
}
|
|
|
|
// LeaseStore coordinates duplicate workers through the database.
|
|
//
|
|
// Acquire must be atomic per (jobType, jobKey, scheduledAt): exactly one caller
|
|
// may hold a window at a time, an already completed window must never be handed
|
|
// out again, and a window whose lease has expired must be reclaimable so a
|
|
// crashed worker cannot block a job forever.
|
|
type LeaseStore interface {
|
|
Acquire(ctx context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error)
|
|
Complete(ctx context.Context, lease Lease, status, errorCode string, counts map[string]int64) error
|
|
}
|
|
|
|
// PostgresLeaseStore implements LeaseStore on the existing job_runs table.
|
|
//
|
|
// It follows the same pattern as alertworker.PostgresLeaseStore, which is
|
|
// hardwired to the alert evaluation job type; this one is parameterized by job
|
|
// type and additionally records result counts, so every worker job shares one
|
|
// coordination and reporting mechanism instead of inventing its own.
|
|
type PostgresLeaseStore struct {
|
|
Pool *pgxpool.Pool
|
|
}
|
|
|
|
// Acquire claims one scheduling window, reclaiming an expired lease if needed.
|
|
func (s PostgresLeaseStore) Acquire(ctx context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error) {
|
|
if s.Pool == nil {
|
|
return Lease{}, false, fmt.Errorf("%w: lease store has no database pool", ErrInvalidConfig)
|
|
}
|
|
if strings.TrimSpace(jobType) == "" || strings.TrimSpace(jobKey) == "" || strings.TrimSpace(owner) == "" || ttl <= 0 {
|
|
return Lease{}, false, fmt.Errorf("%w: lease identity is invalid", ErrInvalidConfig)
|
|
}
|
|
scheduled := scheduledAt.UTC()
|
|
moment := now.UTC()
|
|
leaseUntil := moment.Add(ttl)
|
|
tx, err := s.Pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return Lease{}, false, fmt.Errorf("begin worker lease: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
leaseID := newID()
|
|
var insertedID string
|
|
err = tx.QueryRow(ctx, `INSERT INTO job_runs (id,job_type,job_key,scheduled_at,started_at,status,lease_owner,lease_until) VALUES ($1,$2,$3,$4,$5,'running',$6,$7) ON CONFLICT (job_type,job_key,scheduled_at) DO NOTHING RETURNING id`,
|
|
leaseID, jobType, jobKey, scheduled, moment, owner, leaseUntil).Scan(&insertedID)
|
|
if err == nil {
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return Lease{}, false, fmt.Errorf("commit worker lease: %w", err)
|
|
}
|
|
return Lease{ID: insertedID, JobType: jobType, JobKey: jobKey, ScheduledAt: scheduled, Owner: owner}, true, nil
|
|
}
|
|
if !errors.Is(err, pgx.ErrNoRows) {
|
|
return Lease{}, false, fmt.Errorf("insert worker lease: %w", err)
|
|
}
|
|
var status string
|
|
var existingUntil *time.Time
|
|
if err := tx.QueryRow(ctx, `SELECT status,lease_until FROM job_runs WHERE job_type=$1 AND job_key=$2 AND scheduled_at=$3 FOR UPDATE`, jobType, jobKey, scheduled).Scan(&status, &existingUntil); err != nil {
|
|
return Lease{}, false, fmt.Errorf("read worker lease: %w", err)
|
|
}
|
|
if status != "running" || (existingUntil != nil && existingUntil.After(moment)) {
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return Lease{}, false, fmt.Errorf("commit worker lease contention: %w", err)
|
|
}
|
|
return Lease{}, false, nil
|
|
}
|
|
tag, err := tx.Exec(ctx, `UPDATE job_runs SET started_at=$1,completed_at=NULL,error_code=NULL,lease_owner=$2,lease_until=$3 WHERE job_type=$4 AND job_key=$5 AND scheduled_at=$6 AND status='running' AND (lease_until IS NULL OR lease_until <= $7)`,
|
|
moment, owner, leaseUntil, jobType, jobKey, scheduled, moment)
|
|
if err != nil {
|
|
return Lease{}, false, fmt.Errorf("reclaim worker lease: %w", err)
|
|
}
|
|
if tag.RowsAffected() != 1 {
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return Lease{}, false, fmt.Errorf("commit worker lease reclaim contention: %w", err)
|
|
}
|
|
return Lease{}, false, nil
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return Lease{}, false, fmt.Errorf("commit worker lease reclaim: %w", err)
|
|
}
|
|
return Lease{ID: leaseID, JobType: jobType, JobKey: jobKey, ScheduledAt: scheduled, Owner: owner}, true, nil
|
|
}
|
|
|
|
// Complete records the outcome of a claimed window. The update is scoped to the
|
|
// owning worker, so a run whose lease already expired cannot overwrite the
|
|
// result of the worker that took the window over.
|
|
func (s PostgresLeaseStore) Complete(ctx context.Context, lease Lease, status, errorCode string, counts map[string]int64) error {
|
|
if s.Pool == nil {
|
|
return fmt.Errorf("%w: lease store has no database pool", ErrInvalidConfig)
|
|
}
|
|
if !recordedStatus(status) {
|
|
return fmt.Errorf("%w: worker job status %q is invalid", ErrInvalidConfig, status)
|
|
}
|
|
payload, err := encodeCounts(status, errorCode, counts)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// error_code describes a failure; the reason for a skipped or disabled run
|
|
// lives in counts so a non-failure is never read back as an error.
|
|
failureCode := ""
|
|
if status == StatusFailed {
|
|
failureCode = boundedCode(errorCode)
|
|
}
|
|
tag, err := s.Pool.Exec(ctx, `UPDATE job_runs SET status=$1,completed_at=now(),counts=$2::jsonb,error_code=NULLIF($3,''),lease_owner=NULL,lease_until=NULL WHERE job_type=$4 AND job_key=$5 AND scheduled_at=$6 AND lease_owner=$7 AND status='running'`,
|
|
status, payload, failureCode, lease.JobType, lease.JobKey, lease.ScheduledAt.UTC(), lease.Owner)
|
|
if err != nil {
|
|
return fmt.Errorf("complete worker job: %w", err)
|
|
}
|
|
if tag.RowsAffected() != 1 {
|
|
return ErrLeaseLost
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// recordedStatus bounds what may be written to job_runs.status. Skipped and
|
|
// disabled are kept distinct from completed on purpose: a run that did nothing
|
|
// must never be read back as a successful observation.
|
|
func recordedStatus(status string) bool {
|
|
switch status {
|
|
case StatusCompleted, StatusFailed, StatusSkipped, StatusDisabled:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func encodeCounts(status, reason string, counts map[string]int64) ([]byte, error) {
|
|
document := make(map[string]any, len(counts)+2)
|
|
for key, value := range counts {
|
|
if len(document) >= 20 {
|
|
break
|
|
}
|
|
document[boundedCode(key)] = value
|
|
}
|
|
document["status"] = status
|
|
if reason != "" {
|
|
document["reason"] = boundedCode(reason)
|
|
}
|
|
payload, err := json.Marshal(document)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("encode worker job counts: %w", err)
|
|
}
|
|
return payload, nil
|
|
}
|
|
|
|
func boundedCode(value string) string {
|
|
value = strings.TrimSpace(value)
|
|
if len(value) > 160 {
|
|
return value[:160]
|
|
}
|
|
return value
|
|
}
|
|
|
|
// MemoryLeaseStore is an in-memory LeaseStore with the same semantics, used by
|
|
// tests and by a worker running without a database.
|
|
type MemoryLeaseStore struct {
|
|
mu sync.Mutex
|
|
records map[string]memoryLease
|
|
}
|
|
|
|
type memoryLease struct {
|
|
lease Lease
|
|
status string
|
|
leaseUntil time.Time
|
|
counts map[string]int64
|
|
reason string
|
|
}
|
|
|
|
// NewMemoryLeaseStore returns an empty in-memory lease store.
|
|
func NewMemoryLeaseStore() *MemoryLeaseStore {
|
|
return &MemoryLeaseStore{records: make(map[string]memoryLease)}
|
|
}
|
|
|
|
// Acquire claims a window unless it is completed or still leased.
|
|
func (s *MemoryLeaseStore) Acquire(_ context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error) {
|
|
if strings.TrimSpace(jobType) == "" || strings.TrimSpace(jobKey) == "" || strings.TrimSpace(owner) == "" || ttl <= 0 {
|
|
return Lease{}, false, fmt.Errorf("%w: lease identity is invalid", ErrInvalidConfig)
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.records == nil {
|
|
s.records = make(map[string]memoryLease)
|
|
}
|
|
key := memoryLeaseKey(jobType, jobKey, scheduledAt)
|
|
if record, exists := s.records[key]; exists {
|
|
if record.status != StatusRunning || record.leaseUntil.After(now.UTC()) {
|
|
return Lease{}, false, nil
|
|
}
|
|
}
|
|
lease := Lease{ID: newID(), JobType: jobType, JobKey: jobKey, ScheduledAt: scheduledAt.UTC(), Owner: owner}
|
|
s.records[key] = memoryLease{lease: lease, status: StatusRunning, leaseUntil: now.UTC().Add(ttl)}
|
|
return lease, true, nil
|
|
}
|
|
|
|
// Complete records the outcome for a window this owner still holds.
|
|
func (s *MemoryLeaseStore) Complete(_ context.Context, lease Lease, status, errorCode string, counts map[string]int64) error {
|
|
if !recordedStatus(status) {
|
|
return fmt.Errorf("%w: worker job status %q is invalid", ErrInvalidConfig, status)
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
key := memoryLeaseKey(lease.JobType, lease.JobKey, lease.ScheduledAt)
|
|
record, exists := s.records[key]
|
|
if !exists || record.lease.ID != lease.ID || record.lease.Owner != lease.Owner || record.status != StatusRunning {
|
|
return ErrLeaseLost
|
|
}
|
|
record.status = status
|
|
record.reason = errorCode
|
|
record.counts = counts
|
|
record.leaseUntil = time.Time{}
|
|
s.records[key] = record
|
|
return nil
|
|
}
|
|
|
|
// Status reports the recorded status of one window, or an empty string when the
|
|
// window was never claimed.
|
|
func (s *MemoryLeaseStore) Status(jobType, jobKey string, scheduledAt time.Time) string {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
record, exists := s.records[memoryLeaseKey(jobType, jobKey, scheduledAt)]
|
|
if !exists {
|
|
return ""
|
|
}
|
|
return record.status
|
|
}
|
|
|
|
// Windows reports how many distinct scheduling windows were claimed.
|
|
func (s *MemoryLeaseStore) Windows() int {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return len(s.records)
|
|
}
|
|
|
|
func memoryLeaseKey(jobType, jobKey string, scheduledAt time.Time) string {
|
|
return jobType + "|" + jobKey + "|" + scheduledAt.UTC().Format(time.RFC3339Nano)
|
|
}
|
|
|
|
func newID() string {
|
|
raw := make([]byte, 16)
|
|
if _, err := rand.Read(raw); err != nil {
|
|
return fmt.Sprintf("00000000-0000-4000-8000-%012d", time.Now().UnixNano()%1_000_000_000_000)
|
|
}
|
|
raw[6] = (raw[6] & 0x0f) | 0x40
|
|
raw[8] = (raw[8] & 0x3f) | 0x80
|
|
encoded := hex.EncodeToString(raw)
|
|
return encoded[:8] + "-" + encoded[8:12] + "-" + encoded[12:16] + "-" + encoded[16:20] + "-" + encoded[20:]
|
|
}
|