mirror of
https://github.com/acaloiaro/neoq
synced 2026-07-21 18:29:08 +00:00
Refactor: separate neoq core from backends
- API now supports user-supplied backends when calling neoq.New() - Separated postgres backend code from core API
This commit is contained in:
parent
f1f5ea2988
commit
9dff9f241c
7 changed files with 977 additions and 865 deletions
|
|
@ -42,7 +42,7 @@ Queue Handlers are simple Go functions that accept a `Context` parameter.
|
|||
**Example**: Add a listener on the `hello_world` queue
|
||||
|
||||
```go
|
||||
nq, _ := neoq.New("postgres://username:password@127.0.0.1:5432/neoq?sslmode=disable")
|
||||
nq, _ := neoq.New(neoq.ConnectionString("postgres://postgres:postgres@localhost:5432/neoq?sslmode=disable"))
|
||||
nq.Listen("hello_world", neoq.NewHandler(func(ctx context.Context) (err error) {
|
||||
j, err := neoq.JobFromContext(ctx)
|
||||
log.Println("got job id:", j.ID, "messsage:", j.Payload["message"])
|
||||
|
|
@ -55,7 +55,7 @@ nq.Listen("hello_world", neoq.NewHandler(func(ctx context.Context) (err error) {
|
|||
**Example**: Add a "Hello World" job to the `hello_world` queue
|
||||
|
||||
```go
|
||||
nq, _ := neoq.New("postgres://username:password@127.0.0.1:5432/neoq?sslmode=disable")
|
||||
nq, _ := neoq.New(neoq.ConnectionString("postgres://postgres:postgres@localhost:5432/neoq?sslmode=disable"))
|
||||
jid, _ := nq.Enqueue(neoq.Job{
|
||||
Queue: "hello_world",
|
||||
Payload: map[string]interface{}{
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import (
|
|||
|
||||
func main() {
|
||||
const queue = "foobar"
|
||||
nq, _ := neoq.New("postgres://postgres:postgres@127.0.0.1:5432/neoq?sslmode=disable", neoq.TransactionTimeoutOpt(1000))
|
||||
nq, _ := neoq.New(neoq.PgTransactionTimeoutOpt(1000))
|
||||
|
||||
// Add a job that will execute 1 hour from now
|
||||
jobID, err := nq.Enqueue(neoq.Job{
|
||||
|
|
|
|||
|
|
@ -11,11 +11,11 @@ import (
|
|||
func main() {
|
||||
var err error
|
||||
const queue = "foobar"
|
||||
//
|
||||
nq, _ := neoq.New(
|
||||
"postgres://postgres:postgres@127.0.0.1:5432/neoq?sslmode=disable",
|
||||
neoq.TransactionTimeoutOpt(1000), // transactions may be idle up to one second
|
||||
)
|
||||
|
||||
// by default neoq connects to a local postgres server using: [neoq.DefaultPgConnectionString]
|
||||
// connection strings can be set explicitly as follows:
|
||||
// neoq.New(neoq.ConnectionString("postgres://username:passsword@hostname/database"))
|
||||
nq, _ := neoq.New()
|
||||
|
||||
// we use a done channel here to make sure that our test doesn't exit before the job finishes running
|
||||
// this is probably not a pattern you want to use in production jobs and you see it here only for testing reasons
|
||||
|
|
|
|||
|
|
@ -9,7 +9,10 @@ import (
|
|||
)
|
||||
|
||||
func main() {
|
||||
nq, _ := neoq.New("postgres://postgres:postgres@127.0.0.1:5432/neoq?sslmode=disable", neoq.TransactionTimeoutOpt(1000))
|
||||
// by default neoq connects to a local postgres server using: [neoq.DefaultPgConnectionString]
|
||||
// connection strings can be set explicitly as follows:
|
||||
// neoq.New(neoq.ConnectionString("postgres://username:passsword@hostname/database"))
|
||||
nq, _ := neoq.New(neoq.PgTransactionTimeoutOpt(1000))
|
||||
// run a job periodically
|
||||
handler := neoq.NewHandler(func(ctx context.Context) (err error) {
|
||||
log.Println("running periodic job")
|
||||
|
|
|
|||
18
neoq_test.go
18
neoq_test.go
|
|
@ -4,15 +4,19 @@ import (
|
|||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/exp/slog"
|
||||
)
|
||||
|
||||
func TestWorkerListenConn(t *testing.T) {
|
||||
const queue = "foobar"
|
||||
nq, err := New("postgres://postgres:postgres@127.0.0.1:5432/neoq?sslmode=disable")
|
||||
pgBackend, err := NewPgBackend("postgres://postgres:postgres@127.0.0.1:5432/neoq?sslmode=disable")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
nq, err := New(Backend(pgBackend))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
|
@ -23,7 +27,7 @@ func TestWorkerListenConn(t *testing.T) {
|
|||
handler := NewHandler(func(ctx context.Context) (err error) {
|
||||
var j *Job
|
||||
j, err = JobFromContext(ctx)
|
||||
slog.Info("got job", "id", j.ID, "messsage", j.Payload["message"])
|
||||
log.Println("queue:", j.Queue, "got job", "id:", j.ID, "messsage:", j.Payload["message"])
|
||||
done <- true
|
||||
return
|
||||
})
|
||||
|
|
@ -85,7 +89,7 @@ func TestWorkerListenConn(t *testing.T) {
|
|||
|
||||
func TestWorkerListenCron(t *testing.T) {
|
||||
const cron = "* * * * * *"
|
||||
nq, err := New("postgres://postgres:postgres@127.0.0.1:5432/neoq?sslmode=disable")
|
||||
nq, err := New()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
|
@ -93,10 +97,11 @@ func TestWorkerListenCron(t *testing.T) {
|
|||
jobRan := false
|
||||
var done = make(chan bool)
|
||||
handler := NewHandler(func(ctx context.Context) (err error) {
|
||||
slog.Info("got periodic job")
|
||||
log.Println("got periodic job")
|
||||
done <- true
|
||||
return
|
||||
})
|
||||
|
||||
handler = handler.
|
||||
WithOption(HandlerDeadlineOpt(time.Duration(500 * time.Millisecond))).
|
||||
WithOption(HandlerConcurrencyOpt(1))
|
||||
|
|
@ -123,5 +128,4 @@ func TestWorkerListenCron(t *testing.T) {
|
|||
if !jobRan {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
838
postgres_backend.go
Normal file
838
postgres_backend.go
Normal file
|
|
@ -0,0 +1,838 @@
|
|||
package neoq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/md5"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/iancoleman/strcase"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"github.com/jsuar/go-cron-descriptor/pkg/crondescriptor"
|
||||
"github.com/pkg/errors"
|
||||
"github.com/robfig/cron"
|
||||
"golang.org/x/exp/slog"
|
||||
)
|
||||
|
||||
// NeoqPg is a Postgres-based Neoq backend
|
||||
type NeoqPg struct {
|
||||
config *neoqPgConfig
|
||||
cron *cron.Cron
|
||||
listenConn *pgx.Conn
|
||||
pool *pgxpool.Pool
|
||||
handlers map[string]Handler // a map of queue names to queue handlers
|
||||
mu *sync.Mutex // mutext to protect mutating state on a pgWorker
|
||||
futureJobs map[int64]time.Time // map of future job IDs to their due time
|
||||
logger Logger
|
||||
}
|
||||
|
||||
// If the database does not yet exist, Neoq will attempt to create the database and related tables by default.
|
||||
//
|
||||
// Connection strings may be a URL or DSN-style connection string to neoq's database. The connection string supports multiple
|
||||
// options detailed below.
|
||||
//
|
||||
// options:
|
||||
// pool_max_conns: integer greater than 0
|
||||
// pool_min_conns: integer 0 or greater
|
||||
// pool_max_conn_lifetime: duration string
|
||||
// pool_max_conn_idle_time: duration string
|
||||
// pool_health_check_period: duration string
|
||||
// pool_max_conn_lifetime_jitter: duration string
|
||||
//
|
||||
// # Example DSN
|
||||
// user=worker password=secret host=workerdb.example.com port=5432 dbname=mydb sslmode=verify-ca pool_max_conns=10
|
||||
//
|
||||
// # Example URL
|
||||
// postgres://worker:secret@workerdb.example.com:5432/mydb?sslmode=verify-ca&pool_max_conns=10
|
||||
//
|
||||
// Available ConfigOption
|
||||
// - WithTransactionTimeout(timeout int): configure the idle_in_transaction_timeout for the worker's database
|
||||
// connection(s)
|
||||
func NewPgBackend(connectString string, opts ...ConfigOption) (n Neoq, err error) {
|
||||
w := NeoqPg{
|
||||
mu: &sync.Mutex{},
|
||||
config: &neoqPgConfig{connectString: connectString},
|
||||
handlers: make(map[string]Handler),
|
||||
futureJobs: make(map[int64]time.Time),
|
||||
logger: slog.New(slog.NewTextHandler(os.Stdout)),
|
||||
cron: cron.New(),
|
||||
}
|
||||
w.cron.Start()
|
||||
|
||||
// Set all options
|
||||
for _, opt := range opts {
|
||||
opt(&w)
|
||||
}
|
||||
|
||||
err = w.initializeDB()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
if w.pool == nil {
|
||||
var poolConfig *pgxpool.Config
|
||||
poolConfig, err = pgxpool.ParseConfig(w.config.connectString)
|
||||
if err != nil || w.config.connectString == "" {
|
||||
return nil, errors.New("invalid connecton string: see documentation for valid connection strings")
|
||||
}
|
||||
|
||||
// ensure that workers don't consume connections with idle transactions
|
||||
poolConfig.AfterConnect = func(ctx context.Context, conn *pgx.Conn) (err error) {
|
||||
var query string
|
||||
if w.config.idleTxTimeout > 0 {
|
||||
query = fmt.Sprintf("SET idle_in_transaction_session_timeout = '%dms'", w.config.idleTxTimeout)
|
||||
} else {
|
||||
// there is no limit to the amount of time a worker's transactions may be idle
|
||||
query = "SET idle_in_transaction_session_timeout = 0"
|
||||
}
|
||||
_, err = conn.Exec(ctx, query)
|
||||
return
|
||||
}
|
||||
|
||||
w.pool, err = pgxpool.NewWithConfig(context.Background(), poolConfig)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
n = w
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// neoqPgConfig configures NeoqPg's general behavior
|
||||
type neoqPgConfig struct {
|
||||
idleTxTimeout int // time a worker's transaction may be idle before its connection is closed
|
||||
connectString string // postgres connect string / DSN
|
||||
}
|
||||
|
||||
// initializeDB initializes the tables, types, and indices necessary to operate Neoq
|
||||
func (w NeoqPg) initializeDB() (err error) {
|
||||
var pgxCfg *pgx.ConnConfig
|
||||
var tx pgx.Tx
|
||||
ctx := context.Background()
|
||||
pgxCfg, err = pgx.ParseConfig(w.config.connectString)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
connectStr := fmt.Sprintf("postgres://%s:%s@%s", pgxCfg.User, pgxCfg.Password, pgxCfg.Host)
|
||||
conn, err := pgx.Connect(context.Background(), connectStr)
|
||||
if err != nil {
|
||||
w.logger.Error("unableto connect to database", err)
|
||||
return
|
||||
}
|
||||
defer conn.Close(ctx)
|
||||
|
||||
var dbExists bool
|
||||
dbExistsQ := fmt.Sprintf(`SELECT EXISTS (SELECT datname FROM pg_catalog.pg_database WHERE datname = '%s');`, pgxCfg.Database)
|
||||
rows, err := conn.Query(context.Background(), dbExistsQ)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "unable to determne if jobs table exists: %v", err)
|
||||
return
|
||||
}
|
||||
for rows.Next() {
|
||||
err = rows.Scan(&dbExists)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "unable to determine if jobs table exists: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
conn.Close(ctx)
|
||||
conn, err = pgx.Connect(context.Background(), connectStr)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Unable to connect to database: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer conn.Close(ctx)
|
||||
|
||||
if !dbExists {
|
||||
createDBQ := fmt.Sprintf("CREATE DATABASE %s", pgxCfg.Database)
|
||||
_, err := conn.Exec(ctx, createDBQ)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "unable to create neoq database: %v", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
conn, err = pgx.ConnectConfig(context.Background(), pgxCfg)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
tx, err = conn.Begin(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer tx.Rollback(ctx) // rollback has no effect if the transaction has been committed
|
||||
|
||||
jobsTableExistsQ := `SELECT EXISTS (SELECT FROM
|
||||
pg_tables
|
||||
WHERE
|
||||
schemaname = 'public' AND
|
||||
tablename = 'neoq_jobs'
|
||||
);`
|
||||
rows, err = tx.Query(context.Background(), jobsTableExistsQ)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "unable to determne if jobs table exists: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
var tablesInitialized bool
|
||||
for rows.Next() {
|
||||
err = rows.Scan(&tablesInitialized)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "unable to determine if jobs table exists: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
if !tablesInitialized {
|
||||
createTablesQ := `
|
||||
CREATE TYPE job_status AS ENUM (
|
||||
'new',
|
||||
'processed',
|
||||
'failed'
|
||||
);
|
||||
|
||||
CREATE TABLE neoq_dead_jobs (
|
||||
id SERIAL NOT NULL,
|
||||
fingerprint text NOT NULL,
|
||||
queue text NOT NULL,
|
||||
status job_status NOT NULL default 'failed',
|
||||
payload jsonb,
|
||||
retries integer,
|
||||
max_retries integer,
|
||||
created_at timestamp with time zone DEFAULT now(),
|
||||
error text
|
||||
);
|
||||
|
||||
CREATE TABLE neoq_jobs (
|
||||
id SERIAL NOT NULL,
|
||||
fingerprint text NOT NULL,
|
||||
queue text NOT NULL,
|
||||
status job_status NOT NULL default 'new',
|
||||
payload jsonb,
|
||||
retries integer default 0,
|
||||
max_retries integer default 23,
|
||||
run_after timestamp with time zone DEFAULT now(),
|
||||
ran_at timestamp with time zone,
|
||||
created_at timestamp with time zone DEFAULT now(),
|
||||
error text
|
||||
);
|
||||
|
||||
CREATE INDEX neoq_job_fetcher_idx ON neoq_jobs (id, status, run_after);
|
||||
CREATE INDEX neoq_jobs_fetcher_idx ON neoq_jobs (queue, status, run_after);
|
||||
CREATE INDEX neoq_jobs_fingerprint_idx ON neoq_jobs (fingerprint, status);
|
||||
`
|
||||
|
||||
_, err = tx.Exec(ctx, createTablesQ)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "unable to create job status enum: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
err = tx.Commit(ctx)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Enqueue adds jobs to the specified queue
|
||||
func (w NeoqPg) Enqueue(job Job) (jobID int64, err error) {
|
||||
ctx := context.Background()
|
||||
conn, err := w.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Release()
|
||||
|
||||
tx, err := conn.Begin(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Rollback is safe to call even if the tx is already closed, so if
|
||||
// the tx commits successfully, this is a no-op
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
// Make sure RunAfter is set to a non-zero value if not provided by the caller
|
||||
// if already set, schedule the future job
|
||||
now := time.Now()
|
||||
if job.RunAfter.IsZero() {
|
||||
job.RunAfter = now
|
||||
}
|
||||
|
||||
if job.Queue == "" {
|
||||
err = errors.New("this job does not specify a Queue. Please specify a queue")
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
jobID, err = w.enqueueJob(ctx, tx, job)
|
||||
|
||||
if err != nil || jobID == DuplicateJobID {
|
||||
return
|
||||
}
|
||||
|
||||
// notify listeners that a new job has arrived if it's not a future job
|
||||
if job.RunAfter == now {
|
||||
_, err = tx.Exec(ctx, fmt.Sprintf("NOTIFY %s, '%d'", job.Queue, jobID))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
} else {
|
||||
w.mu.Lock()
|
||||
w.futureJobs[jobID] = job.RunAfter
|
||||
w.mu.Unlock()
|
||||
}
|
||||
|
||||
err = tx.Commit(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Listen sets the queue handler function for the specified queue.
|
||||
//
|
||||
// Neoq will not process any queues until Listen() is called at least once.
|
||||
func (w NeoqPg) Listen(queue string, h Handler) (err error) {
|
||||
w.mu.Lock()
|
||||
w.handlers[queue] = h
|
||||
w.mu.Unlock()
|
||||
|
||||
err = w.start(queue)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// ListenCron listens for jobs on a cron schedule and handles them with the provided handler
|
||||
//
|
||||
// See: https://pkg.go.dev/github.com/robfig/cron?#hdr-CRON_Expression_Format for details on the cron spec format
|
||||
func (w NeoqPg) ListenCron(cronSpec string, h Handler) (err error) {
|
||||
cd, err := crondescriptor.NewCronDescriptor(cronSpec)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cdStr, err := cd.GetDescription(crondescriptor.Full)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
queue := stripNonAlphanum(strcase.ToSnake(*cdStr))
|
||||
w.cron.AddFunc(cronSpec, func() {
|
||||
w.Enqueue(Job{Queue: queue})
|
||||
})
|
||||
|
||||
err = w.Listen(queue, h)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (w NeoqPg) Shutdown() (err error) {
|
||||
w.pool.Close()
|
||||
w.cron.Stop()
|
||||
|
||||
err = w.listenConn.Close(context.Background())
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (w NeoqPg) WithConfigOpt(opt ConfigOption) Neoq {
|
||||
opt(&w)
|
||||
return &w
|
||||
}
|
||||
|
||||
// enqueueJob adds jobs to the queue, returning the job ID
|
||||
//
|
||||
// Jobs that are not already fingerprinted are fingerprinted before being added
|
||||
// Duplicate jobs are not added to the queue. Any two unprocessed jobs with the same fingerprint are duplicates
|
||||
func (w NeoqPg) enqueueJob(ctx context.Context, tx pgx.Tx, j Job) (jobID int64, err error) {
|
||||
|
||||
// fingerprint the job if it hasn't been fingerprinted already
|
||||
// a fingerprint is the md5 sum of: queue + payload
|
||||
if j.Fingerprint == "" {
|
||||
var js []byte
|
||||
js, err = json.Marshal(j.Payload)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
h := md5.New()
|
||||
io.WriteString(h, j.Queue)
|
||||
io.WriteString(h, string(js))
|
||||
j.Fingerprint = fmt.Sprintf("%x", h.Sum(nil))
|
||||
}
|
||||
|
||||
var rowCount int64
|
||||
countRow := tx.QueryRow(ctx, `SELECT COUNT(*) as row_count
|
||||
FROM neoq_jobs
|
||||
WHERE fingerprint = $1
|
||||
AND status NOT IN ('processed')`, j.Fingerprint)
|
||||
err = countRow.Scan(&rowCount)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// this is a duplicate job; skip it
|
||||
if rowCount > 0 {
|
||||
return DuplicateJobID, nil
|
||||
}
|
||||
|
||||
err = tx.QueryRow(ctx, `INSERT INTO neoq_jobs(queue, fingerprint, payload, run_after)
|
||||
VALUES ($1, $2, $3, $4) RETURNING id`,
|
||||
j.Queue, j.Fingerprint, j.Payload, j.RunAfter).Scan(&jobID)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// moveToDeadQueue moves jobs from the pending queue to the dead queue
|
||||
func (w NeoqPg) moveToDeadQueue(ctx context.Context, tx pgx.Tx, j *Job, jobErr error) (err error) {
|
||||
_, err = tx.Exec(ctx, "DELETE FROM neoq_jobs WHERE id = $1", j.ID)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
_, err = tx.Exec(ctx, `INSERT INTO neoq_dead_jobs(id, queue, fingerprint, payload, retries, max_retries, error)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)`,
|
||||
j.ID, j.Queue, j.Fingerprint, j.Payload, j.Retries, j.MaxRetries, jobErr.Error())
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// updateJob updates the status of jobs with: status, run time, error messages, and retries
|
||||
// if the retry count exceeds the maximum number of retries for the job, move the job to the dead jobs queue
|
||||
//
|
||||
// if `tx`'s underlying connection dies, it results in this function throwing an error and job status inaccurately
|
||||
// reflecting the status of the job and its number of retries. TODO: consider solutions to this problem, e.g. acquiring
|
||||
// a new connection in the event of connection failure
|
||||
func (w NeoqPg) updateJob(ctx context.Context, jobErr error) (err error) {
|
||||
status := JobStatusProcessed
|
||||
errMsg := ""
|
||||
|
||||
if jobErr != nil {
|
||||
status = JobStatusFailed
|
||||
errMsg = jobErr.Error()
|
||||
}
|
||||
|
||||
var job *Job
|
||||
if job, err = JobFromContext(ctx); err != nil {
|
||||
return errors.New("unable to get job from context")
|
||||
}
|
||||
|
||||
var tx pgx.Tx
|
||||
if tx, err = txFromContext(ctx); err != nil {
|
||||
return errors.New("unable to get transaction from context")
|
||||
}
|
||||
|
||||
if job.Retries >= job.MaxRetries {
|
||||
err = w.moveToDeadQueue(ctx, tx, job, jobErr)
|
||||
return
|
||||
}
|
||||
|
||||
var runAfter time.Time
|
||||
if job.Retries > 0 && status == JobStatusFailed {
|
||||
runAfter = calculateBackoff(job.Retries)
|
||||
qstr := "UPDATE neoq_jobs SET ran_at = $1, error = $2, status = $3, retries = $4, run_after = $5 WHERE id = $6"
|
||||
_, err = tx.Exec(ctx, qstr, time.Now(), errMsg, status, job.Retries, runAfter, job.ID)
|
||||
} else if job.Retries > 0 && status != JobStatusFailed {
|
||||
qstr := "UPDATE neoq_jobs SET ran_at = $1, error = $2, status = $3, retries = $4 WHERE id = $5"
|
||||
_, err = tx.Exec(ctx, qstr, time.Now(), errMsg, status, job.Retries, job.ID)
|
||||
} else {
|
||||
qstr := "UPDATE neoq_jobs SET ran_at = $1, error = $2, status = $3 WHERE id = $4"
|
||||
_, err = tx.Exec(ctx, qstr, time.Now(), errMsg, status, job.ID)
|
||||
}
|
||||
|
||||
if err == nil && time.Until(runAfter) > 0 {
|
||||
w.mu.Lock()
|
||||
w.futureJobs[job.ID] = runAfter
|
||||
w.mu.Unlock()
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// start starts a queue listener, processes pending job, and fires up goroutines to process future jobs
|
||||
func (w NeoqPg) start(queue string) (err error) {
|
||||
|
||||
var handler Handler
|
||||
var ok bool
|
||||
if handler, ok = w.handlers[queue]; !ok {
|
||||
return fmt.Errorf("no handler for queue: %s", queue)
|
||||
}
|
||||
ctx := context.Background()
|
||||
conn, err := w.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// use a single connection to listen for jobs on all queues
|
||||
// TODO: Give more thought to the implications of hijacking a conn from the pool here
|
||||
// should this connect not come from the pool, to avoid tainting it with connections that don't have an idle in
|
||||
// transaction time out set?
|
||||
w.mu.Lock()
|
||||
if w.listenConn == nil {
|
||||
w.listenConn = conn.Hijack()
|
||||
}
|
||||
w.mu.Unlock()
|
||||
|
||||
listenJobChan := w.listen(ctx, queue) // listen for 'new' jobs
|
||||
pendingJobsChan := w.pendingJobs(ctx, queue) // process overdue jobs *at startup*
|
||||
|
||||
// process all future jobs and retries
|
||||
// TODO there are no illusions about the current future jobs system being robust
|
||||
// the current implementation is a brute force proof of concept that can certainly be improved upon
|
||||
go func() { w.scheduleFutureJobs(ctx, queue) }()
|
||||
|
||||
for i := 0; i < handler.concurrency; i++ {
|
||||
go func() {
|
||||
var jobID int64
|
||||
|
||||
for {
|
||||
select {
|
||||
case jobID = <-listenJobChan:
|
||||
err = w.handleJob(ctx, jobID, handler)
|
||||
case jobID = <-pendingJobsChan:
|
||||
err = w.handleJob(ctx, jobID, handler)
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if err.Error() == "no rows in result set" {
|
||||
err = nil
|
||||
} else {
|
||||
w.logger.Error("error handling job", err, "job_id", jobID)
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// removeFutureJob removes a future job from the in-memory list of jobs that will execute in the future
|
||||
func (w NeoqPg) removeFutureJob(jobID int64) {
|
||||
if _, ok := w.futureJobs[jobID]; ok {
|
||||
w.mu.Lock()
|
||||
delete(w.futureJobs, jobID)
|
||||
w.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// initFutureJobs is intended to be run once to initialize the list of future jobs that must be monitored for
|
||||
// execution. it should be run only during system startup.
|
||||
func (w NeoqPg) initFutureJobs(ctx context.Context, queue string) {
|
||||
rows, err := w.pool.Query(context.Background(), FutureJobQuery, queue)
|
||||
if err != nil {
|
||||
w.logger.Error("error fetching future jobs list", err)
|
||||
return
|
||||
}
|
||||
|
||||
var id int64
|
||||
var runAfter time.Time
|
||||
pgx.ForEachRow(rows, []any{&id, &runAfter}, func() error {
|
||||
w.mu.Lock()
|
||||
w.futureJobs[id] = runAfter
|
||||
w.mu.Unlock()
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// scheduleFutureJobs announces future jobs using NOTIFY on an interval
|
||||
func (w NeoqPg) scheduleFutureJobs(ctx context.Context, queue string) {
|
||||
w.initFutureJobs(ctx, queue)
|
||||
|
||||
// check for new future jobs on an interval
|
||||
// TODO make this time configurable
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
|
||||
for {
|
||||
// loop over list of future jobs, scheduling goroutines to wait for jobs that are due within the next 30 seconds
|
||||
// TODO: Make 30 seconds configurable
|
||||
for jobID, runAfter := range w.futureJobs {
|
||||
at := time.Until(runAfter)
|
||||
if at <= time.Duration(30*time.Second) {
|
||||
w.removeFutureJob(jobID)
|
||||
go func(jid int64) {
|
||||
scheduleCh := time.After(at)
|
||||
<-scheduleCh
|
||||
w.announceJob(ctx, queue, jid)
|
||||
}(jobID)
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ticker.C:
|
||||
continue
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// announceJob announces jobs to queue listeners.
|
||||
//
|
||||
// Announced jobs are executed by the first worker to respond to the announcement.
|
||||
func (w NeoqPg) announceJob(ctx context.Context, queue string, jobID int64) {
|
||||
conn, err := w.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Release()
|
||||
|
||||
tx, err := conn.Begin(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Rollback is safe to call even if the tx is already closed, so if
|
||||
// the tx commits successfully, this is a no-op
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
// notify listeners that a job is ready to run
|
||||
_, err = tx.Exec(ctx, fmt.Sprintf("NOTIFY %s, '%d'", queue, jobID))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
err = tx.Commit(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (w NeoqPg) pendingJobs(ctx context.Context, queue string) (jobsCh chan int64) {
|
||||
jobsCh = make(chan int64)
|
||||
|
||||
// TODO Consider refactoring to use pgxpool.AcquireFunc()
|
||||
conn, err := w.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
w.logger.Error("failed to acquire database connection to listen for pending queue items", err)
|
||||
return
|
||||
}
|
||||
|
||||
go func(ctx context.Context) {
|
||||
defer conn.Release()
|
||||
|
||||
for {
|
||||
jobID, err := w.getPendingJobID(ctx, conn, queue)
|
||||
if err != nil {
|
||||
if err.Error() != "no rows in result set" {
|
||||
w.logger.Error("failed to fetch pending job", err, "job_id", jobID)
|
||||
} else {
|
||||
// done fetching pending jobs
|
||||
break
|
||||
}
|
||||
} else {
|
||||
jobsCh <- jobID
|
||||
}
|
||||
}
|
||||
|
||||
}(ctx)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// handleJob is the workhorse of Neoq
|
||||
// it receives pending, periodic, and retry job ids asynchronously
|
||||
// 1. handleJob first creates a transactions inside of which a row lock is acquired for the job to be processed.
|
||||
// 2. handleJob secondly calls the handler on the job, and finally updates the job's status
|
||||
func (w NeoqPg) handleJob(ctx context.Context, jobID int64, handler Handler) (err error) {
|
||||
var tx pgx.Tx
|
||||
conn, err := w.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Release()
|
||||
|
||||
tx, err = conn.Begin(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer tx.Rollback(ctx) // rollback has no effect if the transaction has been committed
|
||||
|
||||
ctxv := handlerCtxVars{tx: tx}
|
||||
var job Job
|
||||
job, err = w.getPendingJob(ctx, tx, jobID)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
ctxv.job = &job
|
||||
ctx = withHandlerContext(ctx, ctxv)
|
||||
|
||||
// check if the job is being retried and increment retry count accordingly
|
||||
if job.Status != JobStatusNew {
|
||||
job.Retries = job.Retries + 1
|
||||
}
|
||||
|
||||
// execute the queue handler of this job
|
||||
handlerErr := w.execHandler(ctx, handler)
|
||||
err = w.updateJob(ctx, handlerErr)
|
||||
if err != nil {
|
||||
err = errors.Wrap(err, "error updating job status")
|
||||
return
|
||||
}
|
||||
|
||||
err = tx.Commit(ctx)
|
||||
if err != nil {
|
||||
w.logger.Error("unable to commit job transaction. retrying this job may dupliate work", err, "job_id", job.ID)
|
||||
err = errors.Wrap(err, "unable to commit job transaction. retrying this job may dupliate work")
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// exechandler executes handler functions with a concrete time deadline
|
||||
func (w NeoqPg) execHandler(ctx context.Context, handler Handler) (err error) {
|
||||
deadlineCtx, cancel := context.WithDeadline(ctx, time.Now().Add(handler.deadline))
|
||||
defer cancel()
|
||||
|
||||
var done = make(chan bool)
|
||||
go func(ctx context.Context) {
|
||||
err = handler.handle(ctx)
|
||||
done <- true
|
||||
}(ctx)
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case <-deadlineCtx.Done():
|
||||
err = fmt.Errorf("job exceeded its %s deadline", handler.deadline)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// listen uses Postgres LISTEN to listen for jobs on a queue
|
||||
// TODO: There is currently no handling for listener disconnects. This will lead to jobs not getting processed until the
|
||||
// worker is restarted. Implement disconnect handling.
|
||||
func (w NeoqPg) listen(ctx context.Context, queue string) (c chan int64) {
|
||||
var err error
|
||||
c = make(chan int64)
|
||||
|
||||
// set this connection's idle in transaction timeout to infinite so it is not intermittently disconnected
|
||||
_, err = w.listenConn.Exec(ctx, fmt.Sprintf("SET idle_in_transaction_session_timeout = '0'; LISTEN %s", queue))
|
||||
if err != nil {
|
||||
msg := "unable to create database connection for listener"
|
||||
err = errors.Wrap(err, msg)
|
||||
w.logger.Error(msg, err)
|
||||
return
|
||||
}
|
||||
|
||||
go func(ctx context.Context) {
|
||||
for {
|
||||
notification, waitErr := w.listenConn.WaitForNotification(ctx)
|
||||
if waitErr != nil {
|
||||
w.logger.Error("failed to wait for notification", waitErr)
|
||||
time.Sleep(1 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
var jobID int64
|
||||
if jobID, err = strconv.ParseInt(notification.Payload, 0, 64); err != nil {
|
||||
w.logger.Error("unable to fetch job", err)
|
||||
continue
|
||||
}
|
||||
|
||||
c <- jobID
|
||||
}
|
||||
}(ctx)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (w NeoqPg) getPendingJob(ctx context.Context, tx pgx.Tx, jobID int64) (job Job, err error) {
|
||||
row, err := tx.Query(ctx, PendingJobQuery, jobID)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
var j *Job
|
||||
j, err = pgx.CollectOneRow(row, pgx.RowToAddrOfStructByName[Job])
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return *j, err
|
||||
}
|
||||
|
||||
func (w NeoqPg) getPendingJobID(ctx context.Context, conn *pgxpool.Conn, queue string) (jobID int64, err error) {
|
||||
err = conn.QueryRow(ctx, PendingJobIDQuery, queue).Scan(&jobID)
|
||||
return
|
||||
}
|
||||
|
||||
// PgTransactionTimeoutOpt sets the time that NeoqPg's transactions may be idle before its underlying connection is
|
||||
// closed
|
||||
// The timeout is the number of milliseconds that a transaction may sit idle before postgres terminates the
|
||||
// transaction's underlying connection. The timeout should be longer than your longest job takes to complete. If set
|
||||
// too short, job state will become unpredictable, e.g. retry counts may become incorrect.
|
||||
//
|
||||
// PgTransactionTimeoutOpt is best set when calling neoq.New() rather than after creation using WithConfigOpt() because this
|
||||
// setting results in the creation of a new database connection pool.
|
||||
func PgTransactionTimeoutOpt(txTimeout int) ConfigOption {
|
||||
return func(n Neoq) {
|
||||
var ok bool
|
||||
var npg NeoqPg
|
||||
|
||||
if npg, ok = n.(NeoqPg); !ok {
|
||||
return
|
||||
}
|
||||
|
||||
npg.config.idleTxTimeout = txTimeout
|
||||
|
||||
// if the worker already has a pool configured, then setting the transaction timeout prints a warning and is a
|
||||
// no-op
|
||||
if npg.pool != nil {
|
||||
log.Println("create a new Neoq instance to set transaction timeout")
|
||||
return
|
||||
}
|
||||
|
||||
var err error
|
||||
var poolConfig *pgxpool.Config
|
||||
|
||||
poolConfig, err = pgxpool.ParseConfig(npg.config.connectString)
|
||||
if err != nil || npg.config.connectString == "" {
|
||||
return
|
||||
}
|
||||
|
||||
// ensure that workers don't consume connections with idle transactions
|
||||
poolConfig.AfterConnect = func(ctx context.Context, conn *pgx.Conn) (err error) {
|
||||
var query string
|
||||
if npg.config.idleTxTimeout > 0 {
|
||||
query = fmt.Sprintf("SET idle_in_transaction_session_timeout = '%dms'", npg.config.idleTxTimeout)
|
||||
} else {
|
||||
// there is no limit to the amount of time a worker's transactions may be idle
|
||||
query = "SET idle_in_transaction_session_timeout = 0"
|
||||
}
|
||||
_, err = conn.Exec(ctx, query)
|
||||
return
|
||||
}
|
||||
|
||||
npg.pool, err = pgxpool.NewWithConfig(context.Background(), poolConfig)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue