diff --git a/client.go b/client.go index 68f7e51c..06661802 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 ( @@ -823,11 +824,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) @@ -895,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) @@ -2301,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 @@ -2316,6 +2319,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 +2341,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 { @@ -2890,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, } @@ -2923,9 +2940,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 +2976,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..85ed3959 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,10 @@ type pilotSpy struct { jobInsertManyCalls atomic.Int64 jobRetryCalls atomic.Int64 pilotInitCalls atomic.Int64 + pilotInitParams *riverpilot.PilotInitParams + producerInitErr error + producerInitID int64 + producerShutdownErr error testSignals pilotSpyTestSignals } @@ -106,12 +111,16 @@ 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) } 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 +131,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) } @@ -325,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() @@ -334,4 +347,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/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/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/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/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.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 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/river_driver_interface.go b/riverdriver/river_driver_interface.go index 20b92c0c..ee75571f 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) @@ -814,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/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/internal/dbsqlc/river_queue.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_queue.sql.go index b33808c8..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 } @@ -249,7 +259,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 211a4374..ffab8520 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" @@ -888,9 +889,14 @@ 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), "{}"), + Metadata: cmp.Or(string(rivercommon.QueueMetadataForUserWrite(params.Metadata)), "{}"), Name: params.Name, Now: params.Now, PausedAt: params.PausedAt, @@ -905,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 { @@ -976,7 +983,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, }) @@ -1250,7 +1257,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/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/riverdrivertest/queue.go b/riverdriver/riverdrivertest/queue.go index 5c1bad68..a68b2fbd 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() @@ -125,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() @@ -504,6 +552,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/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/internal/dbsqlc/river_queue.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_queue.sql index abc45ed4..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 *; @@ -72,7 +81,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..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 } @@ -238,7 +248,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 41b510eb..9bf1b0ce 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" @@ -873,9 +874,14 @@ 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, + Metadata: rivercommon.QueueMetadataForUserWrite(params.Metadata), Name: params.Name, Now: params.Now, PausedAt: params.PausedAt, @@ -890,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 { @@ -961,7 +968,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, }) @@ -1299,7 +1306,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..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 ) @@ -73,7 +84,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..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 } @@ -257,7 +269,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 c7df4964..2d6a914a 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -1153,9 +1153,13 @@ 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("{}")), + Metadata: sliceutil.FirstNonEmpty(rivercommon.QueueMetadataForUserWrite(params.Metadata), []byte("{}")), Name: params.Name, Now: timeStringNullable(params.Now), PausedAt: timeStringNullable(params.PausedAt), @@ -1170,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 { @@ -1249,7 +1254,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, }) @@ -1657,7 +1662,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/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