From eae34c6ad494ec8e96b6264a5fd4f454255f018d Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sat, 11 Jul 2026 19:35:41 -0500 Subject: [PATCH 1/5] make dynamic queue startup atomic Dynamic queue addition registers a producer before starting it. A startup failure can therefore leave the queue name occupied and, for plugin pilots, leave persisted producer state behind even though `QueueBundle.Add` returned an error. Clean up initialized producer state under a bounded timeout when startup fails. Remove successfully cleaned producers from the client registry so the queue name can be retried. If cleanup also fails, retain the producer behind a typed error so `Remove` can retry without losing track of it. Also release partially registered notifier subscriptions and cover both the successful-cleanup and cleanup-retry paths through a pilot-backed client. --- client.go | 33 +++++++++++++++++++------ client_pilot_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++ error.go | 18 ++++++++++++++ producer.go | 59 +++++++++++++++++++++++++++++++++++--------- 4 files changed, 148 insertions(+), 20 deletions(-) diff --git a/client.go b/client.go index 68f7e51c..61b5724f 100644 --- a/client.go +++ b/client.go @@ -823,11 +823,12 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client } client.queues = &QueueBundle{ - clientFetchCooldown: config.FetchCooldown, - clientFetchPollInterval: config.FetchPollInterval, - clientWillExecuteJobs: config.willExecuteJobs(), - producerAdd: client.producerAdd, - producerRemove: client.producerRemove, + clientFetchCooldown: config.FetchCooldown, + clientFetchPollInterval: config.FetchPollInterval, + clientWillExecuteJobs: config.willExecuteJobs(), + producerAdd: client.producerAdd, + producerRemove: client.producerRemove, + producerRemoveOnAddError: client.producerRemoveOnAddError, } baseservice.Init(archetype, &client.baseService) @@ -2316,6 +2317,11 @@ func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) erro if !ok { return &QueueNotFoundError{Name: queueName} } + if producer.startupCleanupPending.Load() { + if err := producer.abortStartup(ctx); err != nil { + return err + } + } shouldStop, stopped, finalizeStop := producer.StopInit() if shouldStop { @@ -2333,6 +2339,15 @@ func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) erro return nil } +func (c *Client[TTx]) producerRemoveOnAddError(queueName string, producer *producer) { + c.producersMu.Lock() + defer c.producersMu.Unlock() + + if c.producersByQueueName[queueName] == producer && !producer.startupCleanupPending.Load() { + delete(c.producersByQueueName, queueName) + } +} + var nameRegex = regexp.MustCompile(`^(?:[a-z0-9])+(?:[_|\-]?[a-z0-9]+)*$`) func validateQueueName(queueName string) error { @@ -2923,9 +2938,10 @@ type QueueBundle struct { clientWillExecuteJobs bool - fetchCtx context.Context //nolint:containedctx - producerAdd func(queueName string, queueConfig QueueConfig) (*producer, error) // add producer to associated client - producerRemove func(ctx context.Context, queueName string) error // remove producer from associated client + fetchCtx context.Context //nolint:containedctx + producerAdd func(queueName string, queueConfig QueueConfig) (*producer, error) // add producer to associated client + producerRemove func(ctx context.Context, queueName string) error // remove producer from associated client + producerRemoveOnAddError func(queueName string, producer *producer) // Mutex that's acquired when client is starting and stopping and when a // queue is being added so that we can be sure that a client is fully @@ -2958,6 +2974,7 @@ func (b *QueueBundle) Add(queueName string, queueConfig QueueConfig) error { // Start the queue if the client is already started. if b.fetchCtx != nil && b.fetchCtx.Err() == nil { if err := producer.StartWorkContext(b.fetchCtx, b.workCtx); err != nil { + b.producerRemoveOnAddError(queueName, producer) return err } } diff --git a/client_pilot_test.go b/client_pilot_test.go index 7f47cd4b..c45734d3 100644 --- a/client_pilot_test.go +++ b/client_pilot_test.go @@ -2,6 +2,7 @@ package river import ( "context" + "errors" "sync" "sync/atomic" "testing" @@ -29,6 +30,9 @@ type pilotSpy struct { jobInsertManyCalls atomic.Int64 jobRetryCalls atomic.Int64 pilotInitCalls atomic.Int64 + producerInitErr error + producerInitID int64 + producerShutdownErr error testSignals pilotSpyTestSignals } @@ -112,6 +116,9 @@ func (p *pilotSpy) PilotInit(archetype *baseservice.Archetype, params *riverpilo func (p *pilotSpy) ProducerInit(ctx context.Context, exec riverdriver.Executor, params *riverpilot.ProducerInitParams) (int64, riverpilot.ProducerState, error) { p.testSignals.ProducerInit.Signal(struct{}{}) + if p.producerInitErr != nil { + return p.producerInitID, nil, p.producerInitErr + } return p.StandardPilot.ProducerInit(ctx, exec, params) } @@ -122,6 +129,9 @@ func (p *pilotSpy) ProducerKeepAlive(ctx context.Context, exec riverdriver.Execu func (p *pilotSpy) ProducerShutdown(ctx context.Context, exec riverdriver.Executor, params *riverpilot.ProducerShutdownParams) error { p.testSignals.ProducerShutdown.Signal(struct{}{}) + if p.producerShutdownErr != nil { + return p.producerShutdownErr + } return p.StandardPilot.ProducerShutdown(ctx, exec, params) } @@ -334,4 +344,52 @@ func Test_Client_PilotUsage(t *testing.T) { pilot.testSignals.ProducerShutdown.WaitOrTimeout() }) + + t.Run("QueueAddCleansUpFailedProducerStartup", func(t *testing.T) { + t.Parallel() + + client, pilot := setup(t, nil) + pilot.testSignals.Init(t) + startClient(ctx, t, client) + + startupErr := errors.New("producer startup failed") + pilot.producerInitErr = startupErr + pilot.producerInitID = 123 + + err := client.Queues().Add("dynamic_queue", QueueConfig{MaxWorkers: 1}) + require.ErrorIs(t, err, startupErr) + pilot.testSignals.ProducerShutdown.WaitOrTimeout() + + pilot.producerInitErr = nil + require.NoError(t, client.Queues().Add("dynamic_queue", QueueConfig{MaxWorkers: 1})) + }) + + t.Run("QueueAddRetainsProducerUntilCleanupSucceeds", func(t *testing.T) { + t.Parallel() + + client, pilot := setup(t, nil) + pilot.testSignals.Init(t) + startClient(ctx, t, client) + + startupErr := errors.New("producer startup failed") + cleanupErr := errors.New("producer cleanup failed") + pilot.producerInitErr = startupErr + pilot.producerInitID = 123 + pilot.producerShutdownErr = cleanupErr + + err := client.Queues().Add("dynamic_queue", QueueConfig{MaxWorkers: 1}) + var queueAddCleanupErr *QueueAddCleanupError + require.ErrorAs(t, err, &queueAddCleanupErr) + require.ErrorIs(t, err, startupErr) + require.ErrorIs(t, err, cleanupErr) + + err = client.Queues().Add("dynamic_queue", QueueConfig{MaxWorkers: 1}) + require.ErrorIs(t, err, &QueueAlreadyAddedError{}) + + pilot.producerShutdownErr = nil + require.NoError(t, client.Queues().Remove(ctx, "dynamic_queue")) + + pilot.producerInitErr = nil + require.NoError(t, client.Queues().Add("dynamic_queue", QueueConfig{MaxWorkers: 1})) + }) } diff --git a/error.go b/error.go index 7e634d65..0eb805da 100644 --- a/error.go +++ b/error.go @@ -59,6 +59,24 @@ func (e *QueueAlreadyAddedError) Is(target error) bool { return ok } +// QueueAddCleanupError is returned when adding a queue fails after its +// producer was initialized and River cannot clean up that producer. The queue +// remains registered so that Remove can retry cleanup without losing track of +// it. +type QueueAddCleanupError struct { + AddErr error + CleanupErr error + Name string +} + +func (e *QueueAddCleanupError) Error() string { + return fmt.Sprintf("queue %q add failed: %v; cleanup failed: %v", e.Name, e.AddErr, e.CleanupErr) +} + +func (e *QueueAddCleanupError) Unwrap() []error { + return []error{e.AddErr, e.CleanupErr} +} + // QueueNotFoundError is returned when attempting to remove a queue that does // not exist on the Client. type QueueNotFoundError struct { diff --git a/producer.go b/producer.go index 07554ac9..d07d6061 100644 --- a/producer.go +++ b/producer.go @@ -223,8 +223,9 @@ type producer struct { numJobsActive atomic.Int32 numJobsStuck atomic.Int32 - numJobsRan atomic.Uint64 - paused bool + numJobsRan atomic.Uint64 + paused bool + startupCleanupPending atomic.Bool // Receives control messages from the notifier goroutine. Written by notifier // goroutine, only read from main goroutine. queueControlCh chan *controlEventPayload @@ -294,6 +295,21 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { isExpectedShutdownError := func(err error) bool { return errors.Is(err, startstop.ErrStop) || strings.HasSuffix(err.Error(), "conn closed") || fetchCtx.Err() != nil } + producerInitialized := false + finishStartupError := func(startErr error) error { + stopped() + if !producerInitialized { + return startErr + } + + p.startupCleanupPending.Store(true) + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(fetchCtx), 5*time.Second) + defer cancel() + if cleanupErr := p.abortStartup(cleanupCtx); cleanupErr != nil { + return &QueueAddCleanupError{AddErr: startErr, CleanupErr: cleanupErr, Name: p.config.Queue} + } + return startErr + } fetchedQueue, err := func() (*rivertype.Queue, error) { ctx, cancel := context.WithTimeout(fetchCtx, 10*time.Second) @@ -336,15 +352,19 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { Queue: p.config.Queue, Schema: p.config.Schema, }) + if id != 0 { + p.id.Store(id) + producerInitialized = true + } if err != nil { - stopped() - if isExpectedShutdownError(err) { + startupErr := finishStartupError(err) + var cleanupErr *QueueAddCleanupError + if isExpectedShutdownError(err) && !errors.As(startupErr, &cleanupErr) { return nil } p.Logger.ErrorContext(fetchCtx, p.Name+": Error initializing producer state", slog.String("err", err.Error())) - return err + return startupErr } - p.id.Store(id) p.fetchLimiter = chanutil.NewDebouncedChan(fetchCtx, p.config.FetchCooldown, true) @@ -369,20 +389,23 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } insertSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicInsert, handleInsertNotification) if err != nil { - stopped() - if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { + startupErr := finishStartupError(err) + var cleanupErr *QueueAddCleanupError + if (strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled)) && !errors.As(startupErr, &cleanupErr) { return nil } - return err + return startupErr } controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(workCtx)) if err != nil { - stopped() - if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { + insertSub.Unlisten(context.WithoutCancel(fetchCtx)) + startupErr := finishStartupError(err) + var cleanupErr *QueueAddCleanupError + if (strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled)) && !errors.As(startupErr, &cleanupErr) { return nil } - return err + return startupErr } } @@ -440,6 +463,18 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } +func (p *producer) abortStartup(ctx context.Context) error { + if err := p.pilot.ProducerShutdown(ctx, p.exec, &riverpilot.ProducerShutdownParams{ + ProducerID: p.id.Load(), + Queue: p.config.Queue, + Schema: p.config.Schema, + }); err != nil { + return err + } + p.startupCleanupPending.Store(false) + return nil +} + // TriggerJobFetch manually triggers the producer to perform a job fetch // (although it's debounced, so it may not happen immediately if a fetch was // performed very recently). This is used by clients using drivers that don't From 227aaa8daf0fb70e035bb29cecc47def1e80cc07 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sat, 11 Jul 2026 19:39:13 -0500 Subject: [PATCH 2/5] add shared advisory transaction locks Rate and concurrency coordination need readers that can coexist while still excluding queue-wide maintenance. Expose PostgreSQL shared advisory transaction locks through the internal driver interface and both PostgreSQL drivers. SQLite continues to return `ErrNotImplemented`. Extend the cross-driver transaction suite to verify that shared holders coexist and block an exclusive holder until both transactions finish. --- riverdriver/river_driver_interface.go | 1 + .../internal/dbsqlc/pg_misc.sql.go | 9 ++++ .../river_database_sql_driver.go | 5 ++ riverdriver/riverdrivertest/executor_tx.go | 49 +++++++++++++++++++ .../riverpgxv5/internal/dbsqlc/pg_misc.sql | 3 ++ .../riverpgxv5/internal/dbsqlc/pg_misc.sql.go | 9 ++++ riverdriver/riverpgxv5/river_pgx_v5_driver.go | 5 ++ .../riversqlite/river_sqlite_driver.go | 4 ++ 8 files changed, 85 insertions(+) diff --git a/riverdriver/river_driver_interface.go b/riverdriver/river_driver_interface.go index 20b92c0c..9bddf652 100644 --- a/riverdriver/river_driver_interface.go +++ b/riverdriver/river_driver_interface.go @@ -274,6 +274,7 @@ type Executor interface { NotifyMany(ctx context.Context, params *NotifyManyParams) error PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) + PGAdvisoryXactLockShared(ctx context.Context, key int64) (*struct{}, error) QueueCreateOrSetUpdatedAt(ctx context.Context, params *QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) QueueDeleteExpired(ctx context.Context, params *QueueDeleteExpiredParams) ([]string, error) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go index e2542e76..7598df0e 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go @@ -21,6 +21,15 @@ func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) er return err } +const pGAdvisoryXactLockShared = `-- name: PGAdvisoryXactLockShared :exec +SELECT pg_advisory_xact_lock_shared($1) +` + +func (q *Queries) PGAdvisoryXactLockShared(ctx context.Context, db DBTX, key int64) error { + _, err := db.ExecContext(ctx, pGAdvisoryXactLockShared, key) + return err +} + const pGNotifyMany = `-- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index 211a4374..6658f922 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -888,6 +888,11 @@ func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{} return &struct{}{}, interpretError(err) } +func (e *Executor) PGAdvisoryXactLockShared(ctx context.Context, key int64) (*struct{}, error) { + err := dbsqlc.New().PGAdvisoryXactLockShared(ctx, e.dbtx, key) + return &struct{}{}, interpretError(err) +} + func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ Metadata: cmp.Or(string(params.Metadata), "{}"), diff --git a/riverdriver/riverdrivertest/executor_tx.go b/riverdriver/riverdrivertest/executor_tx.go index e4313932..81f3baf8 100644 --- a/riverdriver/riverdrivertest/executor_tx.go +++ b/riverdriver/riverdrivertest/executor_tx.go @@ -214,6 +214,55 @@ func exerciseExecutorTx[TTx any](ctx context.Context, t *testing.T, require.True(t, tryAcquireLock(otherExec)) }) + t.Run("PGAdvisoryXactLockShared", func(t *testing.T) { + t.Parallel() + + { + driver, _ := driverWithSchema(ctx, t, nil) + if driver.DatabaseName() == riverdriver.DatabaseNameSQLite { + t.Logf("Skipping PGAdvisoryXactLockShared test for SQLite") + return + } + } + + lockHash := hashutil.NewAdvisoryLockHash(0) + lockHash.Write([]byte(t.Name())) + lockHash.Write([]byte(randutil.Hex(10))) + key := lockHash.Key() + + exec1 := setup(ctx, t) + tx1, err := exec1.Begin(ctx) + require.NoError(t, err) + t.Cleanup(func() { _ = tx1.Rollback(ctx) }) + + exec2 := setup(ctx, t) + tx2, err := exec2.Begin(ctx) + require.NoError(t, err) + t.Cleanup(func() { _ = tx2.Rollback(ctx) }) + + exec3 := setup(ctx, t) + tx3, err := exec3.Begin(ctx) + require.NoError(t, err) + t.Cleanup(func() { _ = tx3.Rollback(ctx) }) + + _, err = tx1.PGAdvisoryXactLockShared(ctx, key) + require.NoError(t, err) + _, err = tx2.PGAdvisoryXactLockShared(ctx, key) + require.NoError(t, err) + + tryAcquireExclusive := func() bool { + var lockAcquired bool + require.NoError(t, tx3.QueryRow(ctx, "SELECT pg_try_advisory_xact_lock($1)", key).Scan(&lockAcquired)) + return lockAcquired + } + + require.False(t, tryAcquireExclusive()) + require.NoError(t, tx1.Rollback(ctx)) + require.False(t, tryAcquireExclusive()) + require.NoError(t, tx2.Rollback(ctx)) + require.True(t, tryAcquireExclusive()) + }) + t.Run("QueryRow", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql index 19a7b99f..9e4c72d6 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql @@ -1,6 +1,9 @@ -- name: PGAdvisoryXactLock :exec SELECT pg_advisory_xact_lock(@key); +-- name: PGAdvisoryXactLockShared :exec +SELECT pg_advisory_xact_lock_shared(@key); + -- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go index 9215c089..5622fb19 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go @@ -20,6 +20,15 @@ func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) er return err } +const pGAdvisoryXactLockShared = `-- name: PGAdvisoryXactLockShared :exec +SELECT pg_advisory_xact_lock_shared($1) +` + +func (q *Queries) PGAdvisoryXactLockShared(ctx context.Context, db DBTX, key int64) error { + _, err := db.Exec(ctx, pGAdvisoryXactLockShared, key) + return err +} + const pGNotifyMany = `-- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/riverdriver/riverpgxv5/river_pgx_v5_driver.go index 41b510eb..ea4a948d 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -873,6 +873,11 @@ func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{} return &struct{}{}, interpretError(err) } +func (e *Executor) PGAdvisoryXactLockShared(ctx context.Context, key int64) (*struct{}, error) { + err := dbsqlc.New().PGAdvisoryXactLockShared(ctx, e.dbtx, key) + return &struct{}{}, interpretError(err) +} + func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ Metadata: params.Metadata, diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index c7df4964..bd6c1caa 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -1153,6 +1153,10 @@ func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{} return nil, riverdriver.ErrNotImplemented } +func (e *Executor) PGAdvisoryXactLockShared(ctx context.Context, key int64) (*struct{}, error) { + return nil, riverdriver.ErrNotImplemented +} + func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ Metadata: sliceutil.FirstNonEmpty(params.Metadata, []byte("{}")), From bf7c1dc527c787351080d2444252aa9bf2787112 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sat, 11 Jul 2026 19:39:55 -0500 Subject: [PATCH 3/5] expose producer stale horizon to pilots Pilot implementations that interpret producer state need the same liveness horizon that River uses to reap unresponsive producers. Keeping a second duration in plugins can make stale producers count toward one limiter while River has already stopped treating them as live. Name River's existing five-minute horizon once and pass it through `PilotInitParams`. Continue using the same value in producer maintenance and verify plugin pilots receive it during client initialization. --- client.go | 22 ++++++++++++---------- client_pilot_test.go | 3 +++ rivershared/riverpilot/pilot.go | 5 +++++ 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/client.go b/client.go index 61b5724f..8a82aae4 100644 --- a/client.go +++ b/client.go @@ -48,12 +48,13 @@ const ( FetchPollIntervalDefault = 1 * time.Second FetchPollIntervalMin = 1 * time.Millisecond - JobStuckThresholdDefault = 10 * time.Second - JobTimeoutDefault = 1 * time.Minute - MaxAttemptsDefault = rivercommon.MaxAttemptsDefault - PriorityDefault = rivercommon.PriorityDefault - QueueDefault = rivercommon.QueueDefault - QueueNumWorkersMax = 10_000 + JobStuckThresholdDefault = 10 * time.Second + JobTimeoutDefault = 1 * time.Minute + MaxAttemptsDefault = rivercommon.MaxAttemptsDefault + PriorityDefault = rivercommon.PriorityDefault + producerStaleRetentionPeriodDefault = 5 * time.Minute + QueueDefault = rivercommon.QueueDefault + QueueNumWorkersMax = 10_000 ) var ( @@ -896,9 +897,10 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client client.pilot = &riverpilot.StandardPilot{} } client.pilot.PilotInit(archetype, (&riverpilot.PilotInitParams{ - Insert: client.insertMany, - NotifyNonTxJobInsert: client.notifyProducerWithoutListenerJobFetch, - WorkerMetadata: workerMetadata, + Insert: client.insertMany, + NotifyNonTxJobInsert: client.notifyProducerWithoutListenerJobFetch, + ProducerStaleRetentionPeriod: producerStaleRetentionPeriodDefault, + WorkerMetadata: workerMetadata, }).Validate()) pluginPilot, _ := client.pilot.(pilotPlugin) @@ -2302,7 +2304,7 @@ func (c *Client[TTx]) producerAdd(queueName string, queueConfig QueueConfig) (*p RetryPolicy: c.config.RetryPolicy, SchedulerInterval: c.config.schedulerInterval, Schema: c.config.Schema, - StaleProducerRetentionPeriod: 5 * time.Minute, + StaleProducerRetentionPeriod: producerStaleRetentionPeriodDefault, Workers: c.config.Workers, }) c.producersByQueueName[queueName] = producer diff --git a/client_pilot_test.go b/client_pilot_test.go index c45734d3..85ed3959 100644 --- a/client_pilot_test.go +++ b/client_pilot_test.go @@ -30,6 +30,7 @@ type pilotSpy struct { jobInsertManyCalls atomic.Int64 jobRetryCalls atomic.Int64 pilotInitCalls atomic.Int64 + pilotInitParams *riverpilot.PilotInitParams producerInitErr error producerInitID int64 producerShutdownErr error @@ -110,6 +111,7 @@ func (p *pilotSpy) PeriodicJobUpsertMany(ctx context.Context, exec riverdriver.E func (p *pilotSpy) PilotInit(archetype *baseservice.Archetype, params *riverpilot.PilotInitParams) { p.pilotInitCalls.Add(1) + p.pilotInitParams = params p.testSignals.PilotInit.Signal(struct{}{}) p.StandardPilot.PilotInit(archetype, params) } @@ -335,6 +337,7 @@ func Test_Client_PilotUsage(t *testing.T) { require.NotZero(t, insertRes.Job.ID) pilot.testSignals.JobGetAvailable.WaitOrTimeout() + require.Equal(t, 5*time.Minute, pilot.pilotInitParams.ProducerStaleRetentionPeriod) pilot.testSignals.JobSetStateIfRunningMany.WaitOrTimeout() pilot.testSignals.ProducerInit.WaitOrTimeout() pilot.testSignals.ProducerKeepAlive.WaitOrTimeout() diff --git a/rivershared/riverpilot/pilot.go b/rivershared/riverpilot/pilot.go index e42ef797..a9573c22 100644 --- a/rivershared/riverpilot/pilot.go +++ b/rivershared/riverpilot/pilot.go @@ -75,6 +75,11 @@ type PilotInitParams struct { // latency between job insert and when a job is worked. NotifyNonTxJobInsert func(ctx context.Context, res []*rivertype.JobInsertResult) + // ProducerStaleRetentionPeriod is how long an unresponsive producer remains + // live before River may reap it. Pilot maintenance that interprets producer + // state must use the same horizon. + ProducerStaleRetentionPeriod time.Duration + // WorkerMetadata is metadata about registered workers as received from the // client's worker bundle. Only available when a client will work jobs (i.e. // has Workers configured), so while it's safe to assume the presence of From d19146d743ffbaf2ef1f873576189a45310c2abf Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sat, 11 Jul 2026 19:42:56 -0500 Subject: [PATCH 4/5] reserve internal queue metadata Queue metadata needs a namespace for River-managed runtime state without exposing implementation details or allowing ordinary user updates to erase them. Treat top-level `river:*` keys as reserved while preserving the existing public metadata contract. Filter reserved keys from creates, updates, reads, lists, notifications, and producer polling. Preserve them atomically in driver `QueueUpdate` queries. Wrap non-object user metadata under an internal key so arbitrary JSON values continue to round trip while runtime state remains an object alongside them. Exercise the behavior through clients and the shared PostgreSQL and SQLite driver conformance suite. --- client.go | 2 +- client_test.go | 45 +++++++- internal/rivercommon/river_common.go | 46 +++++++- internal/rivercommon/river_common_test.go | 12 +++ producer_test.go | 17 +++ .../internal/dbsqlc/river_queue.sql.go | 17 ++- .../river_database_sql_driver.go | 7 +- riverdriver/riverdrivertest/queue.go | 102 ++++++++++++++++++ .../internal/dbsqlc/river_queue.sql | 17 ++- .../internal/dbsqlc/river_queue.sql.go | 17 ++- riverdriver/riverpgxv5/river_pgx_v5_driver.go | 7 +- .../internal/dbsqlc/river_queue.sql | 20 +++- .../internal/dbsqlc/river_queue.sql.go | 20 +++- .../riversqlite/river_sqlite_driver.go | 6 +- 14 files changed, 315 insertions(+), 20 deletions(-) diff --git a/client.go b/client.go index 8a82aae4..06661802 100644 --- a/client.go +++ b/client.go @@ -2907,7 +2907,7 @@ func (c *Client[TTx]) queueUpdate(ctx context.Context, executorTx riverdriver.Ex controlEvent := &controlEventPayload{ Action: controlActionMetadataChanged, - Metadata: params.Metadata, + Metadata: queue.Metadata, Queue: queue.Name, } diff --git a/client_test.go b/client_test.go index d3c666fe..9d6409bc 100644 --- a/client_test.go +++ b/client_test.go @@ -6348,6 +6348,32 @@ func Test_Client_QueueGet(t *testing.T) { require.Nil(t, queueRes.PausedAt) }) + t.Run("HidesReservedMetadata", func(t *testing.T) { + t.Parallel() + + client, bundle := setup(t) + exec := client.driver.GetExecutor() + + queue := testfactory.Queue(ctx, t, exec, &testfactory.QueueOpts{ + Metadata: []byte(`{"foo":"bar"}`), + Schema: bundle.schema, + }) + err := exec.QueryRow(ctx, + `UPDATE `+dbutil.SafeIdentifier(bundle.schema)+`.river_queue SET metadata = metadata || '{"river:rate_limit_rollup":{"version":1}}' WHERE name = $1 RETURNING 1`, + queue.Name, + ).Scan(new(int)) + require.NoError(t, err) + + queueRes, err := client.QueueGet(ctx, queue.Name) + require.NoError(t, err) + require.JSONEq(t, `{"foo":"bar"}`, string(queueRes.Metadata)) + + listRes, err := client.QueueList(ctx, NewQueueListParams()) + require.NoError(t, err) + require.Len(t, listRes.Queues, 1) + require.JSONEq(t, `{"foo":"bar"}`, string(listRes.Queues[0].Metadata)) + }) + t.Run("ReturnsErrNotFoundIfQueueDoesNotExist", func(t *testing.T) { t.Parallel() @@ -6606,9 +6632,14 @@ func Test_Client_QueueUpdate(t *testing.T) { queue := testfactory.Queue(ctx, t, client.driver.GetExecutor(), &testfactory.QueueOpts{Schema: bundle.schema}) require.Equal(t, []byte(`{}`), queue.Metadata) + err = client.driver.GetExecutor().QueryRow(ctx, + `UPDATE `+dbutil.SafeIdentifier(bundle.schema)+`.river_queue SET metadata = metadata || '{"river:rate_limit_rollup":{"version":1}}' WHERE name = $1 RETURNING 1`, + queue.Name, + ).Scan(new(int)) + require.NoError(t, err) queue, err = client.QueueUpdate(ctx, queue.Name, &QueueUpdateParams{ - Metadata: []byte(`{"foo":"bar"}`), + Metadata: []byte(`{"foo":"bar","river:attempted_override":true}`), }) require.NoError(t, err) require.JSONEq(t, `{"foo":"bar"}`, string(queue.Metadata)) @@ -6621,6 +6652,18 @@ func Test_Client_QueueUpdate(t *testing.T) { Queue: queue.Name, }, notif.payload) + var ( + attemptedOverrideIsNull bool + rollup []byte + ) + err = client.driver.GetExecutor().QueryRow(ctx, + `SELECT metadata -> 'river:rate_limit_rollup', metadata -> 'river:attempted_override' IS NULL FROM `+dbutil.SafeIdentifier(bundle.schema)+`.river_queue WHERE name = $1`, + queue.Name, + ).Scan(&rollup, &attemptedOverrideIsNull) + require.NoError(t, err) + require.JSONEq(t, `{"version":1}`, string(rollup)) + require.True(t, attemptedOverrideIsNull) + queue, err = client.QueueUpdate(ctx, queue.Name, &QueueUpdateParams{ Metadata: []byte(`{"foo":"baz"}`), }) diff --git a/internal/rivercommon/river_common.go b/internal/rivercommon/river_common.go index 4f769b39..ae0e1ba8 100644 --- a/internal/rivercommon/river_common.go +++ b/internal/rivercommon/river_common.go @@ -1,8 +1,10 @@ package rivercommon import ( + "encoding/json" "errors" "regexp" + "strings" "time" ) @@ -12,12 +14,48 @@ import ( const ( // AllQueuesString is a special string that can be used to indicate all // queues in some operations, particularly pause and resume. - AllQueuesString = "*" - MaxAttemptsDefault = 25 - PriorityDefault = 1 - QueueDefault = "default" + AllQueuesString = "*" + MaxAttemptsDefault = 25 + PriorityDefault = 1 + QueueDefault = "default" + QueueMetadataKeyUserMetadata = "river:user_metadata" ) +// QueueMetadataForUserWrite removes reserved keys from user-supplied queue +// metadata while preserving every nonreserved value. +func QueueMetadataForUserWrite(metadata []byte) []byte { + var object map[string]json.RawMessage + if err := json.Unmarshal(metadata, &object); err != nil { + return metadata + } + + for key := range object { + if strings.HasPrefix(key, "river:") { + delete(object, key) + } + } + + metadataWithoutReserved, err := json.Marshal(object) + if err != nil { + return metadata + } + return metadataWithoutReserved +} + +// QueueMetadataWithoutReserved removes River's reserved top-level metadata +// keys before queue metadata crosses a public API boundary. A wrapped +// non-object user value is restored to its original representation. +func QueueMetadataWithoutReserved(metadata []byte) []byte { + var object map[string]json.RawMessage + if err := json.Unmarshal(metadata, &object); err != nil { + return metadata + } + if userMetadata, ok := object[QueueMetadataKeyUserMetadata]; ok { + return userMetadata + } + return QueueMetadataForUserWrite(metadata) +} + // HotOperationTimeout attempts to standardize timeouts for some "hot" // operations like locking available jobs or completing finished jobs. It's // somewhat questionable whether it makes sense to share timing on these diff --git a/internal/rivercommon/river_common_test.go b/internal/rivercommon/river_common_test.go index a492386f..0aea2595 100644 --- a/internal/rivercommon/river_common_test.go +++ b/internal/rivercommon/river_common_test.go @@ -25,3 +25,15 @@ func TestJobKindRE(t *testing.T) { require.NotRegexp(t, UserSpecifiedIDOrKindRE, "with,comma") require.NotRegexp(t, UserSpecifiedIDOrKindRE, ":no_leading_special_characters") } + +func TestQueueMetadataWithoutReserved(t *testing.T) { + t.Parallel() + + require.JSONEq(t, + `{"user":"value"}`, + string(QueueMetadataWithoutReserved([]byte(`{"river:internal":{"state":1},"user":"value"}`))), + ) + require.Equal(t, []byte(`not json`), QueueMetadataWithoutReserved([]byte(`not json`))) + require.JSONEq(t, `[1,2]`, string(QueueMetadataWithoutReserved([]byte(`{"river:user_metadata":[1,2],"river:internal":{}}`)))) + require.JSONEq(t, `{"user":"value"}`, string(QueueMetadataForUserWrite([]byte(`{"river:internal":{},"user":"value"}`)))) +} diff --git a/producer_test.go b/producer_test.go index afac80ae..70fb12d6 100644 --- a/producer_test.go +++ b/producer_test.go @@ -26,6 +26,7 @@ import ( "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/startstoptest" "github.com/riverqueue/river/rivershared/testfactory" + "github.com/riverqueue/river/rivershared/util/dbutil" "github.com/riverqueue/river/rivershared/util/ptrutil" "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/testutil" @@ -745,6 +746,22 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin producer.config.QueuePollInterval = 50 * time.Millisecond startProducer(t, ctx, ctx, producer) + if producer.config.Notifier == nil { + producer.testSignals.PolledQueueConfig.WaitOrTimeout() + + schemaAndTable := "river_queue" + if producer.config.Schema != "" { + schemaAndTable = dbutil.SafeIdentifier(producer.config.Schema) + ".river_queue" + } + err := bundle.exec.QueryRow(ctx, + `UPDATE `+schemaAndTable+` SET metadata = metadata || '{"river:rate_limit_rollup":{"version":1}}' WHERE name = $1 RETURNING 1`, + producer.config.Queue, + ).Scan(new(int)) + require.NoError(t, err) + + producer.testSignals.PolledQueueConfig.WaitOrTimeout() + producer.testSignals.MetadataChanged.RequireEmpty() + } updateMetadata := func(newMetadata []byte) { t.Helper() diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go index b33808c8..43d1a880 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go @@ -249,7 +249,22 @@ func (q *Queries) QueueResume(ctx context.Context, db DBTX, arg *QueueResumePara const queueUpdate = `-- name: QueueUpdate :one UPDATE /* TEMPLATE: schema */river_queue SET - metadata = CASE WHEN $1::boolean THEN $2::jsonb ELSE metadata END, + metadata = CASE WHEN $1::boolean THEN + CASE WHEN jsonb_typeof($2::jsonb) = 'object' THEN + $2::jsonb + ELSE + jsonb_build_object('river:user_metadata', $2::jsonb) + END || coalesce( + ( + SELECT jsonb_object_agg(key, value) + FROM jsonb_each( + CASE WHEN jsonb_typeof(river_queue.metadata) = 'object' THEN river_queue.metadata ELSE '{}'::jsonb END + ) + WHERE key LIKE 'river:%' AND key != 'river:user_metadata' + ), + '{}'::jsonb + ) + ELSE metadata END, updated_at = now() WHERE name = $3 RETURNING name, created_at, metadata, paused_at, updated_at diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index 6658f922..9129454d 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -20,6 +20,7 @@ import ( "github.com/lib/pq" + "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riverdatabasesql/internal/dbsqlc" "github.com/riverqueue/river/rivershared/sqlctemplate" @@ -895,7 +896,7 @@ func (e *Executor) PGAdvisoryXactLockShared(ctx context.Context, key int64) (*st func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ - Metadata: cmp.Or(string(params.Metadata), "{}"), + Metadata: cmp.Or(string(rivercommon.QueueMetadataForUserWrite(params.Metadata)), "{}"), Name: params.Name, Now: params.Now, PausedAt: params.PausedAt, @@ -981,7 +982,7 @@ func (e *Executor) QueueResume(ctx context.Context, params *riverdriver.QueueRes func (e *Executor) QueueUpdate(ctx context.Context, params *riverdriver.QueueUpdateParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueUpdate(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueUpdateParams{ - Metadata: string(params.Metadata), + Metadata: string(rivercommon.QueueMetadataForUserWrite(params.Metadata)), MetadataDoUpdate: params.MetadataDoUpdate, Name: params.Name, }) @@ -1255,7 +1256,7 @@ func queueFromInternal(internal *dbsqlc.RiverQueue) *rivertype.Queue { } return &rivertype.Queue{ CreatedAt: internal.CreatedAt.UTC(), - Metadata: []byte(internal.Metadata), + Metadata: rivercommon.QueueMetadataWithoutReserved([]byte(internal.Metadata)), Name: internal.Name, PausedAt: pausedAt, UpdatedAt: internal.UpdatedAt.UTC(), diff --git a/riverdriver/riverdrivertest/queue.go b/riverdriver/riverdrivertest/queue.go index 5c1bad68..09d01e92 100644 --- a/riverdriver/riverdrivertest/queue.go +++ b/riverdriver/riverdrivertest/queue.go @@ -78,6 +78,19 @@ func exerciseQueue[TTx any](ctx context.Context, t *testing.T, executorWithTx fu require.WithinDuration(t, now, *queue.PausedAt, bundle.driver.TimePrecision()) }) + t.Run("RemovesReservedMetadata", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + + queue, err := exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte(`{"foo":"bar","river:rate_limit_rollup":{"version":1}}`), + Name: "queue-with-reserved-create-metadata", + }) + require.NoError(t, err) + require.JSONEq(t, `{"foo":"bar"}`, string(queue.Metadata)) + }) + t.Run("UpdatesTheUpdatedAtOfExistingQueue", func(t *testing.T) { t.Parallel() @@ -504,6 +517,95 @@ func exerciseQueue[TTx any](ctx context.Context, t *testing.T, executorWithTx fu require.JSONEq(t, `{"baz": "qux"}`, string(updatedQueue.Metadata)) }) + t.Run("PreservesAndHidesReservedMetadata", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + + _, err := exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte(`{"foo":"bar"}`), + Name: "queue-with-reserved-update-metadata", + }) + require.NoError(t, err) + err = exec.QueryRow(ctx, `UPDATE river_queue SET metadata = '{"foo":"bar","river:rate_limit_rollup":{"version":1}}' WHERE name = 'queue-with-reserved-update-metadata' RETURNING 1`).Scan(new(int)) + require.NoError(t, err) + + updatedQueue, err := exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`{"baz":"qux","river:attempted_override":true}`), + MetadataDoUpdate: true, + Name: "queue-with-reserved-update-metadata", + }) + require.NoError(t, err) + require.JSONEq(t, `{"baz":"qux"}`, string(updatedQueue.Metadata)) + + var attemptedOverrideIsNull bool + var rollup []byte + err = exec.QueryRow(ctx, `SELECT metadata -> 'river:rate_limit_rollup', metadata -> 'river:attempted_override' IS NULL FROM river_queue WHERE name = 'queue-with-reserved-update-metadata'`).Scan(&rollup, &attemptedOverrideIsNull) + require.NoError(t, err) + require.JSONEq(t, `{"version":1}`, string(rollup)) + require.True(t, attemptedOverrideIsNull) + }) + + t.Run("PreservesArbitraryUserMetadata", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + + _, err := exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte(`[1,2]`), + Name: "queue-with-array-metadata", + }) + require.NoError(t, err) + updatedQueue, err := exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`{"foo":"bar"}`), + MetadataDoUpdate: true, + Name: "queue-with-array-metadata", + }) + require.NoError(t, err) + require.JSONEq(t, `{"foo":"bar"}`, string(updatedQueue.Metadata)) + + _, err = exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte(`{"foo":"bar"}`), + Name: "queue-with-wrapped-metadata", + }) + require.NoError(t, err) + err = exec.QueryRow(ctx, `UPDATE river_queue SET metadata = '{"river:rate_limit_rollup":{"version":1}}' WHERE name = 'queue-with-wrapped-metadata' RETURNING 1`).Scan(new(int)) + require.NoError(t, err) + updatedQueue, err = exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`[1,2]`), + MetadataDoUpdate: true, + Name: "queue-with-wrapped-metadata", + }) + require.NoError(t, err) + require.JSONEq(t, `[1,2]`, string(updatedQueue.Metadata)) + + var rollup []byte + var userMetadata []byte + err = exec.QueryRow(ctx, `SELECT metadata -> 'river:rate_limit_rollup', metadata -> 'river:user_metadata' FROM river_queue WHERE name = 'queue-with-wrapped-metadata'`).Scan(&rollup, &userMetadata) + require.NoError(t, err) + require.JSONEq(t, `{"version":1}`, string(rollup)) + require.JSONEq(t, `[1,2]`, string(userMetadata)) + }) + + t.Run("TreatsReservedPrefixAsCaseSensitive", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + + _, err := exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte(`{"River:user":"old"}`), + Name: "queue-with-mixed-case-key", + }) + require.NoError(t, err) + updatedQueue, err := exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`{"River:user":"new"}`), + MetadataDoUpdate: true, + Name: "queue-with-mixed-case-key", + }) + require.NoError(t, err) + require.JSONEq(t, `{"River:user":"new"}`, string(updatedQueue.Metadata)) + }) + t.Run("DoesNotUpdateFieldsIfDoUpdateIsFalse", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql index abc45ed4..cf3c7d05 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql @@ -72,7 +72,22 @@ WHERE CASE WHEN @name::text = '*' THEN true ELSE name = @name END; -- name: QueueUpdate :one UPDATE /* TEMPLATE: schema */river_queue SET - metadata = CASE WHEN @metadata_do_update::boolean THEN @metadata::jsonb ELSE metadata END, + metadata = CASE WHEN @metadata_do_update::boolean THEN + CASE WHEN jsonb_typeof(@metadata::jsonb) = 'object' THEN + @metadata::jsonb + ELSE + jsonb_build_object('river:user_metadata', @metadata::jsonb) + END || coalesce( + ( + SELECT jsonb_object_agg(key, value) + FROM jsonb_each( + CASE WHEN jsonb_typeof(river_queue.metadata) = 'object' THEN river_queue.metadata ELSE '{}'::jsonb END + ) + WHERE key LIKE 'river:%' AND key != 'river:user_metadata' + ), + '{}'::jsonb + ) + ELSE metadata END, updated_at = now() WHERE name = @name RETURNING *; diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go index 0988828b..ff94c499 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go @@ -238,7 +238,22 @@ func (q *Queries) QueueResume(ctx context.Context, db DBTX, arg *QueueResumePara const queueUpdate = `-- name: QueueUpdate :one UPDATE /* TEMPLATE: schema */river_queue SET - metadata = CASE WHEN $1::boolean THEN $2::jsonb ELSE metadata END, + metadata = CASE WHEN $1::boolean THEN + CASE WHEN jsonb_typeof($2::jsonb) = 'object' THEN + $2::jsonb + ELSE + jsonb_build_object('river:user_metadata', $2::jsonb) + END || coalesce( + ( + SELECT jsonb_object_agg(key, value) + FROM jsonb_each( + CASE WHEN jsonb_typeof(river_queue.metadata) = 'object' THEN river_queue.metadata ELSE '{}'::jsonb END + ) + WHERE key LIKE 'river:%' AND key != 'river:user_metadata' + ), + '{}'::jsonb + ) + ELSE metadata END, updated_at = now() WHERE name = $3 RETURNING name, created_at, metadata, paused_at, updated_at diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/riverdriver/riverpgxv5/river_pgx_v5_driver.go index ea4a948d..0b066589 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -24,6 +24,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/puddle/v2" + "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riverpgxv5/internal/dbsqlc" "github.com/riverqueue/river/rivershared/sqlctemplate" @@ -880,7 +881,7 @@ func (e *Executor) PGAdvisoryXactLockShared(ctx context.Context, key int64) (*st func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ - Metadata: params.Metadata, + Metadata: rivercommon.QueueMetadataForUserWrite(params.Metadata), Name: params.Name, Now: params.Now, PausedAt: params.PausedAt, @@ -966,7 +967,7 @@ func (e *Executor) QueueResume(ctx context.Context, params *riverdriver.QueueRes func (e *Executor) QueueUpdate(ctx context.Context, params *riverdriver.QueueUpdateParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueUpdate(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueUpdateParams{ - Metadata: params.Metadata, + Metadata: rivercommon.QueueMetadataForUserWrite(params.Metadata), MetadataDoUpdate: params.MetadataDoUpdate, Name: params.Name, }) @@ -1304,7 +1305,7 @@ func queueFromInternal(internal *dbsqlc.RiverQueue) *rivertype.Queue { } return &rivertype.Queue{ CreatedAt: internal.CreatedAt.UTC(), - Metadata: internal.Metadata, + Metadata: rivercommon.QueueMetadataWithoutReserved(internal.Metadata), Name: internal.Name, PausedAt: pausedAt, UpdatedAt: internal.UpdatedAt.UTC(), diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql index 7e769d3f..f37979e6 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql @@ -73,7 +73,25 @@ WHERE CASE WHEN cast(@name AS text) = '*' THEN true ELSE name = @name END; -- name: QueueUpdate :one UPDATE /* TEMPLATE: schema */river_queue SET - metadata = CASE WHEN cast(@metadata_do_update AS boolean) THEN jsonb(@metadata) ELSE metadata END, + metadata = CASE WHEN cast(@metadata_do_update AS boolean) THEN + jsonb_patch( + CASE WHEN json_type(jsonb(@metadata), '$') = 'object' THEN + jsonb(@metadata) + ELSE + jsonb_object('river:user_metadata', jsonb(@metadata)) + END, + CASE WHEN json_type(river_queue.metadata, '$') = 'object' THEN + coalesce( + ( + SELECT jsonb_group_object(key, jsonb(value)) + FROM json_each(river_queue.metadata) + WHERE key GLOB 'river:*' AND key != 'river:user_metadata' + ), + jsonb('{}') + ) + ELSE jsonb('{}') END + ) + ELSE metadata END, updated_at = datetime('now', 'subsec') WHERE name = @name RETURNING *; diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go index fa52751c..470d99df 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go @@ -257,7 +257,25 @@ func (q *Queries) QueueResume(ctx context.Context, db DBTX, arg *QueueResumePara const queueUpdate = `-- name: QueueUpdate :one UPDATE /* TEMPLATE: schema */river_queue SET - metadata = CASE WHEN cast(?1 AS boolean) THEN jsonb(?2) ELSE metadata END, + metadata = CASE WHEN cast(?1 AS boolean) THEN + jsonb_patch( + CASE WHEN json_type(jsonb(?2), '$') = 'object' THEN + jsonb(?2) + ELSE + jsonb_object('river:user_metadata', jsonb(?2)) + END, + CASE WHEN json_type(river_queue.metadata, '$') = 'object' THEN + coalesce( + ( + SELECT jsonb_group_object(key, jsonb(value)) + FROM json_each(river_queue.metadata) + WHERE key GLOB 'river:*' AND key != 'river:user_metadata' + ), + jsonb('{}') + ) + ELSE jsonb('{}') END + ) + ELSE metadata END, updated_at = datetime('now', 'subsec') WHERE name = ?3 RETURNING name, created_at, json(metadata), paused_at, updated_at diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index bd6c1caa..b3ba985e 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -1159,7 +1159,7 @@ func (e *Executor) PGAdvisoryXactLockShared(ctx context.Context, key int64) (*st func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverdriver.QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueCreateOrSetUpdatedAt(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueCreateOrSetUpdatedAtParams{ - Metadata: sliceutil.FirstNonEmpty(params.Metadata, []byte("{}")), + Metadata: sliceutil.FirstNonEmpty(rivercommon.QueueMetadataForUserWrite(params.Metadata), []byte("{}")), Name: params.Name, Now: timeStringNullable(params.Now), PausedAt: timeStringNullable(params.PausedAt), @@ -1253,7 +1253,7 @@ func (e *Executor) QueueResume(ctx context.Context, params *riverdriver.QueueRes func (e *Executor) QueueUpdate(ctx context.Context, params *riverdriver.QueueUpdateParams) (*rivertype.Queue, error) { queue, err := dbsqlc.New().QueueUpdate(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueUpdateParams{ - Metadata: sliceutil.FirstNonEmpty(params.Metadata, []byte("{}")), + Metadata: sliceutil.FirstNonEmpty(rivercommon.QueueMetadataForUserWrite(params.Metadata), []byte("{}")), MetadataDoUpdate: params.MetadataDoUpdate, Name: params.Name, }) @@ -1661,7 +1661,7 @@ func queueFromInternal(internal *dbsqlc.RiverQueue) *rivertype.Queue { } return &rivertype.Queue{ CreatedAt: internal.CreatedAt.UTC(), - Metadata: internal.Metadata, + Metadata: rivercommon.QueueMetadataWithoutReserved(internal.Metadata), Name: internal.Name, PausedAt: pausedAt, UpdatedAt: internal.UpdatedAt.UTC(), From 29ea9a3bf6250870eaf0de7f7bb21a8fa90c751c Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sat, 11 Jul 2026 19:43:29 -0500 Subject: [PATCH 5/5] retain queues with live internal state Queue cleanup currently deletes rows solely from `updated_at`. A queue can therefore age past the retention horizon while reserved runtime state still contains unexpired debt needed by a plugin. Pass one clock value through each cleanup pass and retain rows whose internal rate-limit rollup has a future expiry. Treat malformed runtime state conservatively so cleanup cannot destroy it. Lock selected rows with `FOR UPDATE SKIP LOCKED` to avoid racing concurrent runtime updates and to let parallel cleaners make progress safely. Cover expired, live, absent, and malformed runtime state across PostgreSQL and SQLite drivers. --- internal/maintenance/queue_cleaner.go | 4 +- internal/maintenance/queue_cleaner_test.go | 59 +++++++++++++++++++ riverdriver/river_driver_interface.go | 1 + .../internal/dbsqlc/river_queue.sql.go | 18 ++++-- .../river_database_sql_driver.go | 1 + riverdriver/riverdrivertest/queue.go | 39 +++++++++++- .../internal/dbsqlc/river_queue.sql | 13 +++- .../internal/dbsqlc/river_queue.sql.go | 18 ++++-- riverdriver/riverpgxv5/river_pgx_v5_driver.go | 1 + .../internal/dbsqlc/river_queue.sql | 11 ++++ .../internal/dbsqlc/river_queue.sql.go | 16 ++++- .../riversqlite/river_sqlite_driver.go | 1 + 12 files changed, 167 insertions(+), 15 deletions(-) diff --git a/internal/maintenance/queue_cleaner.go b/internal/maintenance/queue_cleaner.go index f64e97db..0f2ec76e 100644 --- a/internal/maintenance/queue_cleaner.go +++ b/internal/maintenance/queue_cleaner.go @@ -161,10 +161,12 @@ func (s *QueueCleaner) runOnce(ctx context.Context) (*queueCleanerRunOnceResult, ctx, cancelFunc := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault) defer cancelFunc() + now := s.Time.Now() queuesDeleted, err := s.exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{ Max: s.batchSize(), + Now: &now, Schema: s.Config.Schema, - UpdatedAtHorizon: time.Now().Add(-s.Config.RetentionPeriod), + UpdatedAtHorizon: now.Add(-s.Config.RetentionPeriod), }) if err != nil { return nil, fmt.Errorf("error deleting expired queues: %w", err) diff --git a/internal/maintenance/queue_cleaner_test.go b/internal/maintenance/queue_cleaner_test.go index 06058e95..28692ec4 100644 --- a/internal/maintenance/queue_cleaner_test.go +++ b/internal/maintenance/queue_cleaner_test.go @@ -3,6 +3,7 @@ package maintenance import ( "context" "fmt" + "sync/atomic" "testing" "time" @@ -19,6 +20,36 @@ import ( "github.com/riverqueue/river/rivertype" ) +type queueCleanerExecutorMock struct { + riverdriver.Executor + + queueDeleteExpiredFunc func(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) +} + +func (m *queueCleanerExecutorMock) QueueDeleteExpired(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) { + return m.queueDeleteExpiredFunc(ctx, params) +} + +type queueCleanerTimeGenerator struct { + now time.Time + nowCalls atomic.Int32 +} + +func (g *queueCleanerTimeGenerator) Now() time.Time { + g.nowCalls.Add(1) + return g.now +} + +func (g *queueCleanerTimeGenerator) NowOrNil() *time.Time { + now := g.Now() + return &now +} + +func (g *queueCleanerTimeGenerator) StubNow(now time.Time) time.Time { + g.now = now + return now +} + func TestQueueCleaner(t *testing.T) { t.Parallel() @@ -71,6 +102,34 @@ func TestQueueCleaner(t *testing.T) { startstoptest.Stress(ctx, t, cleaner) }) + t.Run("UsesOneNowForExpiration", func(t *testing.T) { + t.Parallel() + + cleaner, bundle := setup(t) + now := time.Date(2026, time.July, 10, 12, 0, 0, 0, time.UTC) + timeGenerator := &queueCleanerTimeGenerator{now: now} + + var capturedParams *riverdriver.QueueDeleteExpiredParams + cleaner.Time = timeGenerator + cleaner.exec = &queueCleanerExecutorMock{ + Executor: bundle.exec, + queueDeleteExpiredFunc: func(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) { + paramsCopy := *params + capturedParams = ¶msCopy + return nil, nil + }, + } + + res, err := cleaner.runOnce(ctx) + require.NoError(t, err) + require.Empty(t, res.QueuesDeleted) + require.Equal(t, int32(1), timeGenerator.nowCalls.Load()) + require.NotNil(t, capturedParams) + require.NotNil(t, capturedParams.Now) + require.Equal(t, now, *capturedParams.Now) + require.Equal(t, now.Add(-cleaner.Config.RetentionPeriod), capturedParams.UpdatedAtHorizon) + }) + t.Run("DeletesExpiredQueues", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/river_driver_interface.go b/riverdriver/river_driver_interface.go index 9bddf652..ee75571f 100644 --- a/riverdriver/river_driver_interface.go +++ b/riverdriver/river_driver_interface.go @@ -815,6 +815,7 @@ type QueueCreateOrSetUpdatedAtParams struct { type QueueDeleteExpiredParams struct { Max int + Now *time.Time Schema string UpdatedAtHorizon time.Time } diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go index 43d1a880..92db5a9f 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go @@ -60,23 +60,33 @@ func (q *Queries) QueueCreateOrSetUpdatedAt(ctx context.Context, db DBTX, arg *Q const queueDeleteExpired = `-- name: QueueDeleteExpired :many DELETE FROM /* TEMPLATE: schema */river_queue -WHERE name IN ( - SELECT name +WHERE ctid IN ( + SELECT ctid FROM /* TEMPLATE: schema */river_queue WHERE river_queue.updated_at < $1 + AND CASE + WHEN NOT (river_queue.metadata ? 'river:rate_limit_rollup') THEN true + WHEN jsonb_typeof(river_queue.metadata -> 'river:rate_limit_rollup' -> 'expires_at_unix') = 'number' + AND river_queue.metadata -> 'river:rate_limit_rollup' ->> 'expires_at_unix' ~ '^[0-9]+$' + THEN (river_queue.metadata -> 'river:rate_limit_rollup' ->> 'expires_at_unix')::numeric + <= extract(epoch FROM coalesce($2::timestamptz, now())) + ELSE false + END ORDER BY name ASC - LIMIT $2::bigint + LIMIT $3::bigint + FOR UPDATE SKIP LOCKED ) RETURNING name, created_at, metadata, paused_at, updated_at ` type QueueDeleteExpiredParams struct { UpdatedAtHorizon time.Time + Now *time.Time Max int64 } func (q *Queries) QueueDeleteExpired(ctx context.Context, db DBTX, arg *QueueDeleteExpiredParams) ([]*RiverQueue, error) { - rows, err := db.QueryContext(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Max) + rows, err := db.QueryContext(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Now, arg.Max) if err != nil { return nil, err } diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index 9129454d..ffab8520 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -911,6 +911,7 @@ func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverd func (e *Executor) QueueDeleteExpired(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) { queues, err := dbsqlc.New().QueueDeleteExpired(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueDeleteExpiredParams{ Max: int64(params.Max), + Now: params.Now, UpdatedAtHorizon: params.UpdatedAtHorizon, }) if err != nil { diff --git a/riverdriver/riverdrivertest/queue.go b/riverdriver/riverdrivertest/queue.go index 09d01e92..a68b2fbd 100644 --- a/riverdriver/riverdrivertest/queue.go +++ b/riverdriver/riverdrivertest/queue.go @@ -138,19 +138,54 @@ func exerciseQueue[TTx any](ctx context.Context, t *testing.T, executorWithTx fu _ = testfactory.Queue(ctx, t, exec, &testfactory.QueueOpts{UpdatedAt: ptrutil.Ptr(now.Add(-23 * time.Hour))}) horizon := now.Add(-24 * time.Hour) - deletedQueueNames, err := exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{Max: 2, UpdatedAtHorizon: horizon}) + deletedQueueNames, err := exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{Max: 2, Now: &now, UpdatedAtHorizon: horizon}) require.NoError(t, err) // queue2 and queue3 should be deleted, with queue4 being skipped due to max of 2: require.Equal(t, []string{queue2.Name, queue3.Name}, deletedQueueNames) // Try again, make sure queue4 gets deleted this time: - deletedQueueNames, err = exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{Max: 2, UpdatedAtHorizon: horizon}) + deletedQueueNames, err = exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{Max: 2, Now: &now, UpdatedAtHorizon: horizon}) require.NoError(t, err) require.Equal(t, []string{queue4.Name}, deletedQueueNames) }) + t.Run("QueueDeleteExpiredRetainsRuntimeState", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + + now := time.Unix(2_000_000_000, 0).UTC() + updatedAt := now.Add(-48 * time.Hour) + for _, name := range []string{"expired-rollup", "malformed-negative-rollup", "malformed-null-rollup", "malformed-rollup", "no-rollup", "unexpired-rollup"} { + _, err := exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Name: name, + UpdatedAt: &updatedAt, + }) + require.NoError(t, err) + } + + for _, query := range []string{ + `UPDATE river_queue SET metadata = '{"river:rate_limit_rollup":{"expires_at_unix":1999999999}}' WHERE name = 'expired-rollup' RETURNING 1`, + `UPDATE river_queue SET metadata = '{"river:rate_limit_rollup":{"expires_at_unix":-1}}' WHERE name = 'malformed-negative-rollup' RETURNING 1`, + `UPDATE river_queue SET metadata = '{"river:rate_limit_rollup":null}' WHERE name = 'malformed-null-rollup' RETURNING 1`, + `UPDATE river_queue SET metadata = '{"river:rate_limit_rollup":{"expires_at_unix":"invalid"}}' WHERE name = 'malformed-rollup' RETURNING 1`, + `UPDATE river_queue SET metadata = '{"river:rate_limit_rollup":{"expires_at_unix":2000003600}}' WHERE name = 'unexpired-rollup' RETURNING 1`, + } { + err := exec.QueryRow(ctx, query).Scan(new(int)) + require.NoError(t, err) + } + + deletedQueueNames, err := exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{ + Max: 10, + Now: &now, + UpdatedAtHorizon: now.Add(-24 * time.Hour), + }) + require.NoError(t, err) + require.ElementsMatch(t, []string{"expired-rollup", "no-rollup"}, deletedQueueNames) + }) + t.Run("QueueGet", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql index cf3c7d05..4cff3c0a 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql @@ -26,12 +26,21 @@ RETURNING *; -- name: QueueDeleteExpired :many DELETE FROM /* TEMPLATE: schema */river_queue -WHERE name IN ( - SELECT name +WHERE ctid IN ( + SELECT ctid FROM /* TEMPLATE: schema */river_queue WHERE river_queue.updated_at < @updated_at_horizon + AND CASE + WHEN NOT (river_queue.metadata ? 'river:rate_limit_rollup') THEN true + WHEN jsonb_typeof(river_queue.metadata -> 'river:rate_limit_rollup' -> 'expires_at_unix') = 'number' + AND river_queue.metadata -> 'river:rate_limit_rollup' ->> 'expires_at_unix' ~ '^[0-9]+$' + THEN (river_queue.metadata -> 'river:rate_limit_rollup' ->> 'expires_at_unix')::numeric + <= extract(epoch FROM coalesce(sqlc.narg('now')::timestamptz, now())) + ELSE false + END ORDER BY name ASC LIMIT @max::bigint + FOR UPDATE SKIP LOCKED ) RETURNING *; diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go index ff94c499..751420c0 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql.go @@ -58,23 +58,33 @@ func (q *Queries) QueueCreateOrSetUpdatedAt(ctx context.Context, db DBTX, arg *Q const queueDeleteExpired = `-- name: QueueDeleteExpired :many DELETE FROM /* TEMPLATE: schema */river_queue -WHERE name IN ( - SELECT name +WHERE ctid IN ( + SELECT ctid FROM /* TEMPLATE: schema */river_queue WHERE river_queue.updated_at < $1 + AND CASE + WHEN NOT (river_queue.metadata ? 'river:rate_limit_rollup') THEN true + WHEN jsonb_typeof(river_queue.metadata -> 'river:rate_limit_rollup' -> 'expires_at_unix') = 'number' + AND river_queue.metadata -> 'river:rate_limit_rollup' ->> 'expires_at_unix' ~ '^[0-9]+$' + THEN (river_queue.metadata -> 'river:rate_limit_rollup' ->> 'expires_at_unix')::numeric + <= extract(epoch FROM coalesce($2::timestamptz, now())) + ELSE false + END ORDER BY name ASC - LIMIT $2::bigint + LIMIT $3::bigint + FOR UPDATE SKIP LOCKED ) RETURNING name, created_at, metadata, paused_at, updated_at ` type QueueDeleteExpiredParams struct { UpdatedAtHorizon time.Time + Now *time.Time Max int64 } func (q *Queries) QueueDeleteExpired(ctx context.Context, db DBTX, arg *QueueDeleteExpiredParams) ([]*RiverQueue, error) { - rows, err := db.Query(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Max) + rows, err := db.Query(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Now, arg.Max) if err != nil { return nil, err } diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/riverdriver/riverpgxv5/river_pgx_v5_driver.go index 0b066589..9bf1b0ce 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -896,6 +896,7 @@ func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverd func (e *Executor) QueueDeleteExpired(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) { queues, err := dbsqlc.New().QueueDeleteExpired(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueDeleteExpiredParams{ Max: int64(params.Max), + Now: params.Now, UpdatedAtHorizon: params.UpdatedAtHorizon, }) if err != nil { diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql index f37979e6..5a406f15 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql @@ -30,6 +30,17 @@ WHERE name IN ( SELECT name FROM /* TEMPLATE: schema */river_queue WHERE river_queue.updated_at < @updated_at_horizon + AND CASE + WHEN json_type(river_queue.metadata, '$') != 'object' + OR NOT EXISTS ( + SELECT 1 FROM json_each(river_queue.metadata) WHERE key = 'river:rate_limit_rollup' + ) THEN true + WHEN json_type(river_queue.metadata, '$."river:rate_limit_rollup".expires_at_unix') = 'integer' + AND cast(json_extract(river_queue.metadata, '$."river:rate_limit_rollup".expires_at_unix') AS integer) >= 0 + THEN cast(json_extract(river_queue.metadata, '$."river:rate_limit_rollup".expires_at_unix') AS integer) + <= unixepoch(coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec'))) + ELSE false + END ORDER BY name ASC LIMIT @max ) diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go index 470d99df..1dbbdbc9 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_queue.sql.go @@ -63,19 +63,31 @@ WHERE name IN ( SELECT name FROM /* TEMPLATE: schema */river_queue WHERE river_queue.updated_at < ?1 + AND CASE + WHEN json_type(river_queue.metadata, '$') != 'object' + OR NOT EXISTS ( + SELECT 1 FROM json_each(river_queue.metadata) WHERE key = 'river:rate_limit_rollup' + ) THEN true + WHEN json_type(river_queue.metadata, '$."river:rate_limit_rollup".expires_at_unix') = 'integer' + AND cast(json_extract(river_queue.metadata, '$."river:rate_limit_rollup".expires_at_unix') AS integer) >= 0 + THEN cast(json_extract(river_queue.metadata, '$."river:rate_limit_rollup".expires_at_unix') AS integer) + <= unixepoch(coalesce(cast(?2 AS text), datetime('now', 'subsec'))) + ELSE false + END ORDER BY name ASC - LIMIT ?2 + LIMIT ?3 ) RETURNING name, created_at, json(metadata), paused_at, updated_at ` type QueueDeleteExpiredParams struct { UpdatedAtHorizon time.Time + Now *string Max int64 } func (q *Queries) QueueDeleteExpired(ctx context.Context, db DBTX, arg *QueueDeleteExpiredParams) ([]*RiverQueue, error) { - rows, err := db.QueryContext(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Max) + rows, err := db.QueryContext(ctx, queueDeleteExpired, arg.UpdatedAtHorizon, arg.Now, arg.Max) if err != nil { return nil, err } diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index b3ba985e..2d6a914a 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -1174,6 +1174,7 @@ func (e *Executor) QueueCreateOrSetUpdatedAt(ctx context.Context, params *riverd func (e *Executor) QueueDeleteExpired(ctx context.Context, params *riverdriver.QueueDeleteExpiredParams) ([]string, error) { queues, err := dbsqlc.New().QueueDeleteExpired(schemaTemplateParam(ctx, params.Schema), e.dbtx, &dbsqlc.QueueDeleteExpiredParams{ Max: int64(params.Max), + Now: timeStringNullable(params.Now), UpdatedAtHorizon: params.UpdatedAtHorizon.UTC(), }) if err != nil {