|
| 1 | +// Copyright 2026 The Go Authors. All rights reserved. |
| 2 | +// Use of this source code is governed by a BSD-style |
| 3 | +// license that can be found in the LICENSE file. |
| 4 | + |
| 5 | +// Package pgqueue provides a Postgres-backed queue implementation for |
| 6 | +// scheduling and processing fetch actions. It supports multiple concurrent |
| 7 | +// workers (processes or goroutines) |
| 8 | +package pgqueue |
| 9 | + |
| 10 | +import ( |
| 11 | + "context" |
| 12 | + "database/sql" |
| 13 | + "errors" |
| 14 | + "fmt" |
| 15 | + "sync" |
| 16 | + "time" |
| 17 | + |
| 18 | + "golang.org/x/pkgsite/internal/database" |
| 19 | + "golang.org/x/pkgsite/internal/log" |
| 20 | + "golang.org/x/pkgsite/internal/queue" |
| 21 | +) |
| 22 | + |
| 23 | +// The frequency at which we poll for work. |
| 24 | +const pollInterval = 5 * time.Second |
| 25 | + |
| 26 | +// ProcessFunc is the function signature for processing dequeued work. |
| 27 | +type ProcessFunc func(ctx context.Context, modulePath, version string) (int, error) |
| 28 | + |
| 29 | +// Queue implements the Queue interface backed by a Postgres table. It is safe |
| 30 | +// for concurrent use by multiple goroutines and processes. |
| 31 | +type Queue struct { |
| 32 | + db *database.DB |
| 33 | +} |
| 34 | + |
| 35 | +// New creates the queue_tasks table if it doesn't exist and returns a Queue. |
| 36 | +func New(ctx context.Context, db *database.DB) (*Queue, error) { |
| 37 | + // TODO(jbarkhuysen): If we find it onerous to do table updates over time, we |
| 38 | + // may want to consider alternatives to doing this here. |
| 39 | + if _, err := db.Exec(ctx, createTableQuery); err != nil { |
| 40 | + return nil, fmt.Errorf("pgqueue.New: creating table: %w", err) |
| 41 | + } |
| 42 | + return &Queue{db: db}, nil |
| 43 | +} |
| 44 | + |
| 45 | +const createTableQuery = ` |
| 46 | +CREATE TABLE IF NOT EXISTS queue_tasks ( |
| 47 | + id BIGSERIAL PRIMARY KEY, |
| 48 | + task_name TEXT UNIQUE NOT NULL, |
| 49 | + module_path TEXT NOT NULL, |
| 50 | + version TEXT NOT NULL, |
| 51 | + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), |
| 52 | + started_at TIMESTAMPTZ |
| 53 | +); |
| 54 | +CREATE INDEX IF NOT EXISTS idx_queue_tasks_started_created |
| 55 | +ON queue_tasks (started_at, created_at);` |
| 56 | + |
| 57 | +// ScheduleFetch inserts a task into queue_tasks. It returns (true, nil) if the |
| 58 | +// task was inserted, or (false, nil) if it was a duplicate. |
| 59 | +func (q *Queue) ScheduleFetch(ctx context.Context, modulePath, version string, opts *queue.Options) (bool, error) { |
| 60 | + taskName := modulePath + "@" + version |
| 61 | + if opts != nil && opts.Suffix != "" { |
| 62 | + taskName += "-" + opts.Suffix |
| 63 | + } |
| 64 | + n, err := q.db.Exec(ctx, |
| 65 | + `INSERT INTO queue_tasks (task_name, module_path, version) VALUES ($1, $2, $3) ON CONFLICT (task_name) DO NOTHING`, |
| 66 | + taskName, modulePath, version) |
| 67 | + if err != nil { |
| 68 | + return false, fmt.Errorf("pgqueue.ScheduleFetch(%q, %q): %w", modulePath, version, err) |
| 69 | + } |
| 70 | + return n == 1, nil |
| 71 | +} |
| 72 | + |
| 73 | +// Poll starts background polling for work. It spawns the given number of worker |
| 74 | +// goroutines, each of which periodically claims a task, runs processFunc, and |
| 75 | +// deletes the task on completion. It blocks until ctx is cancelled. |
| 76 | +func (q *Queue) Poll(ctx context.Context, workers int, processFunc ProcessFunc) { |
| 77 | + wg := sync.WaitGroup{} |
| 78 | + for range workers { |
| 79 | + wg.Go(func() { |
| 80 | + // Periodically claim work. |
| 81 | + ticker := time.NewTicker(pollInterval) |
| 82 | + defer ticker.Stop() |
| 83 | + for { |
| 84 | + select { |
| 85 | + case <-ctx.Done(): |
| 86 | + return |
| 87 | + case <-ticker.C: |
| 88 | + q.claimAndProcess(ctx, processFunc) |
| 89 | + } |
| 90 | + } |
| 91 | + }) |
| 92 | + } |
| 93 | + wg.Wait() |
| 94 | +} |
| 95 | + |
| 96 | +// TODO(jbarkhuysen): 5m stall timeout is baked in; we might want to make it |
| 97 | +// variable in the future. |
| 98 | +const dequeueQuery = ` |
| 99 | +WITH next_task AS ( |
| 100 | + SELECT id |
| 101 | + FROM queue_tasks |
| 102 | + WHERE started_at IS NULL |
| 103 | + OR started_at + INTERVAL '5 minutes' < NOW() |
| 104 | + ORDER BY created_at ASC |
| 105 | + LIMIT 1 |
| 106 | + FOR UPDATE SKIP LOCKED |
| 107 | +) |
| 108 | +UPDATE queue_tasks |
| 109 | +SET started_at = NOW() |
| 110 | +WHERE id = (SELECT id FROM next_task) |
| 111 | +RETURNING id, module_path, version, started_at` |
| 112 | + |
| 113 | +func (q *Queue) claimAndProcess(ctx context.Context, processFunc ProcessFunc) { |
| 114 | + var id int64 |
| 115 | + var modulePath, version string |
| 116 | + var startedAt time.Time |
| 117 | + err := q.db.QueryRow(ctx, dequeueQuery).Scan(&id, &modulePath, &version, &startedAt) |
| 118 | + if errors.Is(err, sql.ErrNoRows) { |
| 119 | + return // There's no work: no-op. |
| 120 | + } |
| 121 | + if err != nil { |
| 122 | + log.Errorf(ctx, "pgqueue: dequeue: %v", err) |
| 123 | + return |
| 124 | + } |
| 125 | + |
| 126 | + log.Infof(ctx, "pgqueue: processing %s@%s (task %d)", modulePath, version, id) |
| 127 | + code, err := processFunc(ctx, modulePath, version) |
| 128 | + if err != nil { |
| 129 | + log.Errorf(ctx, "pgqueue: processing %s@%s: status=%d err=%v", modulePath, version, code, err) |
| 130 | + // This still gets removed (delete below) so that we don't endlessly |
| 131 | + // fail the same work item. |
| 132 | + } |
| 133 | + |
| 134 | + // Use a background context for cleanup so the delete succeeds even if |
| 135 | + // the poll context has been cancelled. |
| 136 | + delCtx := context.Background() |
| 137 | + if _, err := q.db.Exec(delCtx, `DELETE FROM queue_tasks WHERE id = $1 AND started_at = $2`, id, startedAt); err != nil { |
| 138 | + log.Errorf(delCtx, "pgqueue: deleting task %d: %v", id, err) |
| 139 | + } |
| 140 | +} |
0 commit comments