diff --git a/CHANGELOG.md b/CHANGELOG.md index baa537f3..5ced190a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- Both hooks and middleware should now be registered under a general `Config.Plugins` instead, a change that simpifies things somewhat, and better handles cases where a hook or middleware might be _both_ a hook and a middleware as opposed to only one or the other. `Config.Hooks` and `Config.Middleware` are still available for use, but considered deprecated in favor of the more general `Config.Plugins`. [PR #1284](https://github.com/riverqueue/river/pull/1284). +- Hooks passed to `Config.Hooks` that also implement middleware now have their middleware functionality activate automatically. The converse is also true for middleware sent to `Config.Middleware` that implement hooks. [PR #1284](https://github.com/riverqueue/river/pull/1284). + ## [0.40.0] - 2026-07-02 ⚠️ Version 0.40.0 contains a new database migration, version 7, that rolls up some database cleanups and a few SQLite features: diff --git a/client.go b/client.go index 68f7e51c..1dab4695 100644 --- a/client.go +++ b/client.go @@ -17,16 +17,15 @@ import ( "github.com/riverqueue/river/internal/dblist" "github.com/riverqueue/river/internal/dbunique" - "github.com/riverqueue/river/internal/hooklookup" "github.com/riverqueue/river/internal/jobcompleter" "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/leadership" "github.com/riverqueue/river/internal/maintenance" - "github.com/riverqueue/river/internal/middlewarelookup" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/notifylimiter" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" - "github.com/riverqueue/river/internal/rivermiddleware" + "github.com/riverqueue/river/internal/riverplugin" "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" @@ -241,6 +240,11 @@ type Config struct { // work hook runs and the insertion hooks on either side of it are skipped. // // Jobs may have their own specific hooks by implementing JobArgsWithHooks. + // + // If a type in Hooks also implements rivertype.Middleware, it will be + // installed as middleware too. + // + // Deprecated: Use Plugins instead. Hooks []rivertype.Hook // Logger is the structured logger to use for logging purposes. If none is @@ -274,8 +278,27 @@ type Config struct { // middlewares will run one after another, and the work middleware between // them will not run. When a job is worked, the work middleware runs and the // insertion middlewares on either side of it are skipped. + // + // If a type in Middleware also implements rivertype.Hook, it will be + // installed as a hook too. + // + // Deprecated: Use Plugins instead. Middleware []rivertype.Middleware + // Plugins contains extensions installed globally as hooks, middleware, or + // both. + // + // A type qualifies as a plugin by implementing [rivertype.Plugin]. Most + // existing hook and middleware implementations already do this by embedding + // HookDefaults or MiddlewareDefaults. If a type participates on both sides, + // it may embed PluginDefaults, or embed both HookDefaults and + // MiddlewareDefaults directly and define its own IsPlugin method, then + // implement any operation-specific hook or middleware interfaces it needs. + // + // Hooks and Middleware are still supported for backward compatibility, but + // Plugins is the preferred place to register new extensions. + Plugins []rivertype.Plugin + // PeriodicJobs are a set of periodic jobs to run at the specified intervals // in the client. PeriodicJobs []*PeriodicJob @@ -500,6 +523,7 @@ func (c *Config) WithDefaults() *Config { MaxAttempts: cmp.Or(c.MaxAttempts, MaxAttemptsDefault), Middleware: c.Middleware, PeriodicJobs: c.PeriodicJobs, + Plugins: c.Plugins, PollOnly: c.PollOnly, Queues: c.Queues, ReindexerIndexNames: reindexerIndexNames, @@ -520,6 +544,35 @@ func (c *Config) WithDefaults() *Config { } } +func middlewareFromConfig(config *Config) []rivertype.Middleware { + middleware := make([]rivertype.Middleware, 0, + len(config.Middleware)+len(config.JobInsertMiddleware)+len(config.WorkerMiddleware)) + middleware = append(middleware, config.Middleware...) + + for _, jobInsertMiddleware := range config.JobInsertMiddleware { + middleware = append(middleware, jobInsertMiddleware) + } + +outerLoop: + for _, workerMiddleware := range config.WorkerMiddleware { + // Don't add the middleware if it also implements JobInsertMiddleware + // and the instance has been added to config.JobInsertMiddleware. This + // is a hedge to make sure we don't accidentally double add middleware + // as we've converted over to the unified config.Middleware setting. + if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddleware.(rivertype.JobInsertMiddleware); ok { + for _, jobInsertMiddleware := range config.JobInsertMiddleware { + if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddleware { + continue outerLoop + } + } + } + + middleware = append(middleware, workerMiddleware) + } + + return middleware +} + func (c *Config) validate() error { if c.CancelledJobRetentionPeriod < -1 { return errors.New("CancelledJobRetentionPeriod time cannot be less than zero, except for -1 (infinite)") @@ -678,28 +731,27 @@ type Client[TTx any] struct { baseService baseservice.BaseService baseStartStop startstop.BaseStartStop - clientNotifyBundle *ClientNotifyBundle[TTx] - completer jobcompleter.JobCompleter - config *Config - driver riverdriver.Driver[TTx] - elector *leadership.Elector - hookLookupByJob *hooklookup.JobHookLookup - hookLookupGlobal hooklookup.HookLookupInterface - insertNotifyLimiter *notifylimiter.Limiter - middlewareLookupGlobal middlewarelookup.MiddlewareLookupInterface - notifier *notifier.Notifier // may be nil in poll-only mode - periodicJobs *PeriodicJobBundle - pilot riverpilot.Pilot - producersByQueueName map[string]*producer - producersMu sync.RWMutex - queueMaintainer *maintenance.QueueMaintainer - queueMaintainerLeader *maintenance.QueueMaintainerLeader - queues *QueueBundle - services []startstop.Service - stopped <-chan struct{} - stuckJobCount atomic.Int32 - subscriptionManager *subscriptionManager - testSignals clientTestSignals + clientNotifyBundle *ClientNotifyBundle[TTx] + completer jobcompleter.JobCompleter + config *Config + driver riverdriver.Driver[TTx] + elector *leadership.Elector + pluginLookupByJob *pluginlookup.JobPluginLookup + pluginLookupGlobal pluginlookup.PluginLookupInterface + insertNotifyLimiter *notifylimiter.Limiter + notifier *notifier.Notifier // may be nil in poll-only mode + periodicJobs *PeriodicJobBundle + pilot riverpilot.Pilot + producersByQueueName map[string]*producer + producersMu sync.RWMutex + queueMaintainer *maintenance.QueueMaintainer + queueMaintainerLeader *maintenance.QueueMaintainerLeader + queues *QueueBundle + services []startstop.Service + stopped <-chan struct{} + stuckJobCount atomic.Int32 + subscriptionManager *subscriptionManager + testSignals clientTestSignals // workCancel cancels the context used for all work goroutines. Normal Stop // does not cancel that context. @@ -802,8 +854,14 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client } } - for _, hook := range config.Hooks { - if withBaseService, ok := hook.(baseservice.WithBaseService); ok { + allPlugins := pluginlookup.NormalizePlugins( + config.Hooks, + middlewareFromConfig(config), + append(riverplugin.DefaultPlugins(), config.Plugins...), + ) + + for _, plugin := range allPlugins { + if withBaseService, ok := plugin.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } } @@ -815,8 +873,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client }, config: config, driver: driver, - hookLookupByJob: hooklookup.NewJobHookLookup(), - hookLookupGlobal: hooklookup.NewHookLookup(config.Hooks), + pluginLookupByJob: pluginlookup.NewJobPluginLookup(), + pluginLookupGlobal: pluginlookup.NewPluginLookup(allPlugins), producersByQueueName: make(map[string]*producer), testSignals: clientTestSignals{}, workCancel: func(cause error) {}, // replaced on start, but here in case StopAndCancel is called before start up @@ -834,41 +892,6 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client client.baseService.Name = "Client" // Have to correct the name because base service isn't embedded like it usually is client.insertNotifyLimiter = notifylimiter.NewLimiter(archetype, config.FetchCooldown) - // Validation ensures that config.JobInsertMiddleware/WorkerMiddleware or - // the more abstract config.Middleware for middleware are set, but not both, - // so in practice we never append all three of these to each other. - { - middleware := rivermiddleware.DefaultMiddleware() - middleware = append(middleware, config.Middleware...) - for _, jobInsertMiddleware := range config.JobInsertMiddleware { - middleware = append(middleware, jobInsertMiddleware) - } - outerLoop: - for _, workerMiddleware := range config.WorkerMiddleware { - // Don't add the middleware if it also implements JobInsertMiddleware - // and the instance has been added to config.JobInsertMiddleware. This - // is a hedge to make sure we don't accidentally double add middleware - // as we've converted over to the unified config.Middleware setting. - if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddleware.(rivertype.JobInsertMiddleware); ok { - for _, jobInsertMiddleware := range config.JobInsertMiddleware { - if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddleware { - continue outerLoop - } - } - } - - middleware = append(middleware, workerMiddleware) - } - - for _, middleware := range middleware { - if withBaseService, ok := middleware.(baseservice.WithBaseService); ok { - baseservice.Init(archetype, withBaseService) - } - } - - client.middlewareLookupGlobal = middlewarelookup.NewMiddlewareLookup(middleware) - } - pluginDriver, _ := driver.(driverPlugin[TTx]) if pluginDriver != nil { pluginDriver.PluginInit(archetype) @@ -996,7 +1019,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client { periodicJobEnqueuer, err := maintenance.NewPeriodicJobEnqueuer(archetype, &maintenance.PeriodicJobEnqueuerConfig{ AdvisoryLockPrefix: config.AdvisoryLockPrefix, - HookLookupGlobal: client.hookLookupGlobal, + PluginLookupGlobal: client.pluginLookupGlobal, Insert: client.insertMany, Pilot: client.pilot, Schema: config.Schema, @@ -1976,8 +1999,8 @@ func (c *Client[TTx]) insertManyShared( doInner := func(ctx context.Context) ([]*rivertype.JobInsertResult, error) { for _, params := range insertParams { for _, hook := range append( - c.hookLookupGlobal.ByHookKind(hooklookup.HookKindInsertBegin), - c.hookLookupByJob.ByJobArgs(params.Args).ByHookKind(hooklookup.HookKindInsertBegin)..., + c.pluginLookupGlobal.ByKind(pluginlookup.HookKindInsertBegin), + c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.HookKindInsertBegin)..., ) { if err := hook.(rivertype.HookInsertBegin).InsertBegin(ctx, params); err != nil { //nolint:forcetypeassert return nil, err @@ -2008,7 +2031,7 @@ func (c *Client[TTx]) insertManyShared( return insertResults, nil } - jobInsertMiddleware := c.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert) + jobInsertMiddleware := c.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert) if len(jobInsertMiddleware) > 0 { // Wrap middlewares in reverse order so the one defined first is wrapped // as the outermost function and is first to receive the operation. @@ -2286,14 +2309,13 @@ func (c *Client[TTx]) producerAdd(queueName string, queueConfig QueueConfig) (*p ErrorHandler: c.config.ErrorHandler, FetchCooldown: cmp.Or(queueConfig.FetchCooldown, c.config.FetchCooldown), FetchPollInterval: cmp.Or(queueConfig.FetchPollInterval, c.config.FetchPollInterval), - HookLookupByJob: c.hookLookupByJob, - HookLookupGlobal: c.hookLookupGlobal, + PluginLookupByJob: c.pluginLookupByJob, + PluginLookupGlobal: c.pluginLookupGlobal, JobStuckHandler: c.config.JobStuckHandler, JobStuckCount: &c.stuckJobCount, JobStuckThreshold: c.config.JobStuckThreshold, JobTimeout: c.config.JobTimeout, MaxWorkers: queueConfig.MaxWorkers, - MiddlewareLookupGlobal: c.middlewareLookupGlobal, Notifier: c.notifier, Queue: queueName, QueueEventCallback: c.subscriptionManager.distributeQueueEvent, diff --git a/client_test.go b/client_test.go index d3c666fe..1b2ca255 100644 --- a/client_test.go +++ b/client_test.go @@ -26,8 +26,8 @@ import ( "github.com/riverqueue/river/internal/dbunique" "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/maintenance" - "github.com/riverqueue/river/internal/middlewarelookup" "github.com/riverqueue/river/internal/notifier" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/internal/riverinternaltest/retrypolicytest" @@ -70,6 +70,82 @@ type noOpWorker struct { func (w *noOpWorker) Work(ctx context.Context, job *Job[noOpArgs]) error { return nil } +var ( + _ rivertype.HookInsertBegin = &hookMiddlewareEmbeddedDefaultsPlugin{} + _ rivertype.JobInsertMiddleware = &hookMiddlewareEmbeddedDefaultsPlugin{} + _ rivertype.Plugin = &hookMiddlewareEmbeddedDefaultsPlugin{} + + _ rivertype.HookInsertBegin = &hookMiddlewarePlugin{} + _ rivertype.JobInsertMiddleware = &hookMiddlewarePlugin{} + _ rivertype.Plugin = &hookMiddlewarePlugin{} + + _ rivertype.HookInsertBegin = hookMiddlewareValuePlugin{} + _ rivertype.JobInsertMiddleware = hookMiddlewareValuePlugin{} + _ rivertype.Plugin = hookMiddlewareValuePlugin{} +) + +type hookMiddlewareEmbeddedDefaultsPlugin struct { + HookDefaults + MiddlewareDefaults + + insertBeginCount int + insertManyCount int +} + +func (p *hookMiddlewareEmbeddedDefaultsPlugin) IsPlugin() bool { return true } + +func (p *hookMiddlewareEmbeddedDefaultsPlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + p.insertBeginCount++ + return nil +} + +func (p *hookMiddlewareEmbeddedDefaultsPlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + p.insertManyCount++ + return doInner(ctx) +} + +type hookMiddlewarePlugin struct { + PluginDefaults + + insertBeginCount int + insertManyCount int +} + +func (p *hookMiddlewarePlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + p.insertBeginCount++ + return nil +} + +func (p *hookMiddlewarePlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + p.insertManyCount++ + return doInner(ctx) +} + +type hookMiddlewareValuePlugin struct { + counts *hookMiddlewareValuePluginCounts +} + +func (p hookMiddlewareValuePlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + p.counts.insertBeginCount++ + return nil +} + +func (p hookMiddlewareValuePlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + p.counts.insertManyCount++ + return doInner(ctx) +} + +func (p hookMiddlewareValuePlugin) IsHook() bool { return true } + +func (p hookMiddlewareValuePlugin) IsPlugin() bool { return true } + +func (p hookMiddlewareValuePlugin) IsMiddleware() bool { return true } + +type hookMiddlewareValuePluginCounts struct { + insertBeginCount int + insertManyCount int +} + type periodicJobArgs struct{} func (periodicJobArgs) Kind() string { return "periodic_job" } @@ -8204,6 +8280,121 @@ func Test_NewClient_Overrides(t *testing.T) { require.Len(t, client.config.WorkerMiddleware, 1) } +func Test_NewClient_PluginsAndHybrids(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + config *Config + dbPool *pgxpool.Pool + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + dbPool := riversharedtest.DBPool(ctx, t) + driver := riverpgxv5.New(dbPool) + schema := riverdbtest.TestSchema(ctx, t, driver, nil) + + return &testBundle{ + config: newTestConfig(t, schema), + dbPool: dbPool, + } + } + + insertAndRequireCounts := func(t *testing.T, bundle *testBundle, plugin *hookMiddlewarePlugin, expectedCount int) { + t.Helper() + + client := newTestClient(t, bundle.dbPool, bundle.config) + + _, err := client.Insert(ctx, noOpArgs{}, nil) + require.NoError(t, err) + + require.Equal(t, expectedCount, plugin.insertBeginCount) + require.Equal(t, expectedCount, plugin.insertManyCount) + } + + t.Run("DuplicatesAcrossHooksMiddlewareAndPluginsRunMultipleTimes", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + plugin := &hookMiddlewarePlugin{} + bundle.config.Hooks = []rivertype.Hook{plugin} + bundle.config.Middleware = []rivertype.Middleware{plugin} + bundle.config.Plugins = []rivertype.Plugin{plugin} + + insertAndRequireCounts(t, bundle, plugin, 3) + }) + + t.Run("HookAlsoRegisteredAsMiddleware", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + plugin := &hookMiddlewarePlugin{} + bundle.config.Hooks = []rivertype.Hook{plugin} + + insertAndRequireCounts(t, bundle, plugin, 1) + }) + + t.Run("MiddlewareAlsoRegisteredAsHook", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + plugin := &hookMiddlewarePlugin{} + bundle.config.Middleware = []rivertype.Middleware{plugin} + + insertAndRequireCounts(t, bundle, plugin, 1) + }) + + t.Run("PluginRegisteredAsHookAndMiddleware", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + plugin := &hookMiddlewarePlugin{} + bundle.config.Plugins = []rivertype.Plugin{plugin} + + insertAndRequireCounts(t, bundle, plugin, 1) + }) + + t.Run("PluginRegisteredWithEmbeddedDefaults", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + plugin := &hookMiddlewareEmbeddedDefaultsPlugin{} + bundle.config.Plugins = []rivertype.Plugin{plugin} + + client := newTestClient(t, bundle.dbPool, bundle.config) + + _, err := client.Insert(ctx, noOpArgs{}, nil) + require.NoError(t, err) + + require.Equal(t, 1, plugin.insertBeginCount) + require.Equal(t, 1, plugin.insertManyCount) + }) + + t.Run("SeparateEqualValueInstancesRunSeparately", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + counts := &hookMiddlewareValuePluginCounts{} + bundle.config.Hooks = []rivertype.Hook{ + hookMiddlewareValuePlugin{counts: counts}, + } + bundle.config.Middleware = []rivertype.Middleware{ + hookMiddlewareValuePlugin{counts: counts}, + } + + client := newTestClient(t, bundle.dbPool, bundle.config) + + _, err := client.Insert(ctx, noOpArgs{}, nil) + require.NoError(t, err) + + require.Equal(t, 2, counts.insertBeginCount) + require.Equal(t, 2, counts.insertManyCount) + }) +} + func Test_NewClient_ReindexerIndexNamesExplicitEmptyOverride(t *testing.T) { t.Parallel() @@ -8388,8 +8579,8 @@ func Test_NewClient_Validations(t *testing.T) { config.Middleware = []rivertype.Middleware{&overridableJobMiddleware{}} }, validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper - require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 1) - require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2) + require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert), 1) + require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker), 2) }, }, { @@ -8399,8 +8590,8 @@ func Test_NewClient_Validations(t *testing.T) { config.WorkerMiddleware = []rivertype.WorkerMiddleware{&overridableJobMiddleware{}} }, validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper - require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 2) - require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 3) + require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert), 2) + require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker), 3) }, }, { @@ -8411,8 +8602,8 @@ func Test_NewClient_Validations(t *testing.T) { config.WorkerMiddleware = []rivertype.WorkerMiddleware{middleware} }, validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper - require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 1) - require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2) + require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert), 1) + require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker), 2) }, }, { diff --git a/example_global_hooks_test.go b/example_global_hooks_test.go index edc2dce2..9f0b2b76 100644 --- a/example_global_hooks_test.go +++ b/example_global_hooks_test.go @@ -65,7 +65,7 @@ func Example_globalHooks() { //nolint:dupl riverClient, err := river.NewClient(riverpgxv5.New(dbPool), initTestConfig(ctx, dbPool, &river.Config{ // Order is significant. See output below. - Hooks: []rivertype.Hook{ + Plugins: []rivertype.Plugin{ &BothInsertAndWorkBeginHook{}, &InsertBeginHook{}, &WorkBeginHook{}, diff --git a/example_global_middleware_test.go b/example_global_middleware_test.go index 29750727..fb594457 100644 --- a/example_global_middleware_test.go +++ b/example_global_middleware_test.go @@ -65,7 +65,7 @@ func Example_globalMiddleware() { //nolint:dupl riverClient, err := river.NewClient(riverpgxv5.New(dbPool), initTestConfig(ctx, dbPool, &river.Config{ // Order is significant. See output below. - Middleware: []rivertype.Middleware{ + Plugins: []rivertype.Plugin{ &JobBothInsertAndWorkMiddleware{}, &JobInsertMiddleware{}, &WorkerMiddleware{}, diff --git a/hook_defaults_funcs.go b/hook_defaults_funcs.go index 93a7f4fe..c768b6b8 100644 --- a/hook_defaults_funcs.go +++ b/hook_defaults_funcs.go @@ -7,12 +7,14 @@ import ( ) // HookDefaults should be embedded on any hooks implementation. It helps -// identify a struct as hooks, and guarantee forward compatibility in case -// additions are necessary to the rivertype.Hook interface. +// identify a struct as hooks and a plugin, and guarantee forward compatibility +// in case additions are necessary to the rivertype.Hook interface. type HookDefaults struct{} func (d *HookDefaults) IsHook() bool { return true } +func (d *HookDefaults) IsPlugin() bool { return true } + // HookInsertBeginFunc is a convenience helper for implementing // rivertype.HookInsertBegin using a simple function instead of a struct. type HookInsertBeginFunc func(ctx context.Context, params *rivertype.JobInsertParams) error @@ -23,12 +25,16 @@ func (f HookInsertBeginFunc) InsertBegin(ctx context.Context, params *rivertype. func (f HookInsertBeginFunc) IsHook() bool { return true } +func (f HookInsertBeginFunc) IsPlugin() bool { return true } + // HookPeriodicJobsStartFunc is a convenience helper for implementing // rivertype.HookPeriodicJobsStart using a simple function instead of a struct. type HookPeriodicJobsStartFunc func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error func (f HookPeriodicJobsStartFunc) IsHook() bool { return true } +func (f HookPeriodicJobsStartFunc) IsPlugin() bool { return true } + func (f HookPeriodicJobsStartFunc) Start(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return f(ctx, params) } @@ -37,18 +43,22 @@ func (f HookPeriodicJobsStartFunc) Start(ctx context.Context, params *rivertype. // rivertype.HookWorkBegin using a simple function instead of a struct. type HookWorkBeginFunc func(ctx context.Context, job *rivertype.JobRow) error +func (f HookWorkBeginFunc) IsHook() bool { return true } + +func (f HookWorkBeginFunc) IsPlugin() bool { return true } + func (f HookWorkBeginFunc) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { return f(ctx, job) } -func (f HookWorkBeginFunc) IsHook() bool { return true } - // HookWorkEndFunc is a convenience helper for implementing // rivertype.HookWorkEnd using a simple function instead of a struct. type HookWorkEndFunc func(ctx context.Context, job *rivertype.JobRow, err error) error +func (f HookWorkEndFunc) IsHook() bool { return true } + +func (f HookWorkEndFunc) IsPlugin() bool { return true } + func (f HookWorkEndFunc) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error { return f(ctx, job, err) } - -func (f HookWorkEndFunc) IsHook() bool { return true } diff --git a/hook_defaults_funcs_test.go b/hook_defaults_funcs_test.go index 43a20297..52007e2f 100644 --- a/hook_defaults_funcs_test.go +++ b/hook_defaults_funcs_test.go @@ -10,10 +10,17 @@ import ( var ( _ rivertype.Hook = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil }) _ rivertype.HookInsertBegin = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil }) + _ rivertype.Plugin = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil }) _ rivertype.Hook = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil }) _ rivertype.HookPeriodicJobsStart = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil }) + _ rivertype.Plugin = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil }) _ rivertype.Hook = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil }) _ rivertype.HookWorkBegin = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil }) + _ rivertype.Plugin = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil }) + + _ rivertype.Hook = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return nil }) + _ rivertype.HookWorkEnd = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return nil }) + _ rivertype.Plugin = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return nil }) ) diff --git a/internal/hooklookup/hook_lookup.go b/internal/hooklookup/hook_lookup.go deleted file mode 100644 index 420debb2..00000000 --- a/internal/hooklookup/hook_lookup.go +++ /dev/null @@ -1,166 +0,0 @@ -package hooklookup - -import ( - "sync" - - "github.com/riverqueue/river/rivertype" -) - -// -// HookKind -// - -type HookKind string - -const ( - HookKindInsertBegin HookKind = "insert_begin" - HookKindPeriodicJobsStart HookKind = "periodic_job_start" - HookKindWorkBegin HookKind = "work_begin" - HookKindWorkEnd HookKind = "work_end" -) - -// -// HookLookupInterface -// - -// HookLookupInterface is an interface to look up hooks by hook kind. It's -// commonly implemented by HookLookup, but may also be EmptyHookLookup as a -// memory allocation optimization for bundles where no hooks are present. -type HookLookupInterface interface { - ByHookKind(kind HookKind) []rivertype.Hook -} - -// NewHookLookup returns a new hook lookup interface based on the given hooks -// that satisfies HookLookupInterface. This is often hookLookup, but may be -// emptyHookLookup as an optimization for the common case of an empty hook -// bundle. -func NewHookLookup(hooks []rivertype.Hook) HookLookupInterface { - if len(hooks) < 1 { - return &emptyHookLookup{} - } - - return &hookLookup{ - hooks: hooks, - hooksByKind: make(map[HookKind][]rivertype.Hook), - mu: &sync.RWMutex{}, - } -} - -// -// hookLookup -// - -// hookLookup looks up and caches hooks based on a HookKind, saving work when -// looking up hooks for specific operations, a common operation that gets -// repeated over and over again. This struct may be used as a lookup for -// globally installed hooks or hooks for specific job kinds through the use of -// JobHookLookup. -type hookLookup struct { - hooks []rivertype.Hook - hooksByKind map[HookKind][]rivertype.Hook - mu *sync.RWMutex -} - -func (c *hookLookup) ByHookKind(kind HookKind) []rivertype.Hook { - c.mu.RLock() - cache, ok := c.hooksByKind[kind] - c.mu.RUnlock() - if ok { - return cache - } - - c.mu.Lock() - defer c.mu.Unlock() - - // Even if this ends up being empty, make sure there's an entry for the next - // time the cache gets invoked for this kind. - c.hooksByKind[kind] = nil - - // Rely on exhaustlint to find any missing hook kinds here. - switch kind { - case HookKindInsertBegin: - for _, hook := range c.hooks { - if typedHook, ok := hook.(rivertype.HookInsertBegin); ok { - c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook) - } - } - case HookKindPeriodicJobsStart: - for _, hook := range c.hooks { - if typedHook, ok := hook.(rivertype.HookPeriodicJobsStart); ok { - c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook) - } - } - case HookKindWorkBegin: - for _, hook := range c.hooks { - if typedHook, ok := hook.(rivertype.HookWorkBegin); ok { - c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook) - } - } - case HookKindWorkEnd: - for _, hook := range c.hooks { - if typedHook, ok := hook.(rivertype.HookWorkEnd); ok { - c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook) - } - } - } - - return c.hooksByKind[kind] -} - -// -// emptyHookLookup -// - -// emptyHookLookup is an empty version of HookLookup that's zero allocation. For -// most applications, most job args won't have hooks, so this prevents us from -// allocating dozens/hundreds of small HookLookup objects that go unused. -type emptyHookLookup struct{} - -func (c *emptyHookLookup) ByHookKind(kind HookKind) []rivertype.Hook { return nil } - -// -// JobHookLookup -// - -type JobHookLookup struct { - hookLookupByKind map[string]HookLookupInterface - mu sync.RWMutex -} - -func NewJobHookLookup() *JobHookLookup { - return &JobHookLookup{ - hookLookupByKind: make(map[string]HookLookupInterface), - } -} - -// ByJobArgs returns a HookLookupInterface by job args, which is a HookLookup if -// the job args had specific hooks (i.e. implements JobArgsWithHooks and returns -// a non-empty set of hooks), or an EmptyHashLookup otherwise. -func (c *JobHookLookup) ByJobArgs(args rivertype.JobArgs) HookLookupInterface { - kind := args.Kind() - - c.mu.RLock() - lookup, ok := c.hookLookupByKind[kind] - c.mu.RUnlock() - if ok { - return lookup - } - - c.mu.Lock() - defer c.mu.Unlock() - - var hooks []rivertype.Hook - if argsWithHooks, ok := args.(jobArgsWithHooks); ok { - hooks = argsWithHooks.Hooks() - } - - lookup = NewHookLookup(hooks) - c.hookLookupByKind[kind] = lookup - return lookup -} - -// Same as river.JobArgsWithHooks, but duplicated here so that can still live in -// the top level package. -type jobArgsWithHooks interface { - Hooks() []rivertype.Hook -} diff --git a/internal/hooklookup/hook_lookup_test.go b/internal/hooklookup/hook_lookup_test.go deleted file mode 100644 index 4390ca6c..00000000 --- a/internal/hooklookup/hook_lookup_test.go +++ /dev/null @@ -1,266 +0,0 @@ -package hooklookup - -import ( - "context" - "sync" - "testing" - - "github.com/stretchr/testify/require" - - "github.com/riverqueue/river/rivertype" -) - -func TestHookLookup(t *testing.T) { - t.Parallel() - - type testBundle struct{} - - setup := func(t *testing.T) (*hookLookup, *testBundle) { //nolint:unparam - t.Helper() - - return NewHookLookup([]rivertype.Hook{ //nolint:forcetypeassert - &testHookInsertAndWorkBegin{}, - &testHookInsertBegin{}, - &testHookWorkBegin{}, - &testHookWorkEnd{}, - }).(*hookLookup), &testBundle{} - } - - t.Run("LooksUpHooks", func(t *testing.T) { - t.Parallel() - - hookLookup, _ := setup(t) - - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookInsertBegin{}, - }, hookLookup.ByHookKind(HookKindInsertBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookWorkBegin{}, - }, hookLookup.ByHookKind(HookKindWorkBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookWorkEnd{}, - }, hookLookup.ByHookKind(HookKindWorkEnd)) - - require.Len(t, hookLookup.hooksByKind, 3) - - // Repeat lookups to make sure we get the same result. - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookInsertBegin{}, - }, hookLookup.ByHookKind(HookKindInsertBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookWorkBegin{}, - }, hookLookup.ByHookKind(HookKindWorkBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookWorkEnd{}, - }, hookLookup.ByHookKind(HookKindWorkEnd)) - }) - - t.Run("Stress", func(t *testing.T) { - t.Parallel() - - hookLookup, _ := setup(t) - - var wg sync.WaitGroup - - parallelLookupLoop := func(kind HookKind) { - wg.Go(func() { - for range 50 { - hookLookup.ByHookKind(kind) - } - }) - } - - parallelLookupLoop(HookKindInsertBegin) - parallelLookupLoop(HookKindWorkBegin) - parallelLookupLoop(HookKindInsertBegin) - parallelLookupLoop(HookKindWorkBegin) - - wg.Wait() - }) -} - -func TestEmptyHookLookup(t *testing.T) { - t.Parallel() - - type testBundle struct{} - - setup := func(t *testing.T) (*emptyHookLookup, *testBundle) { - t.Helper() - - return NewHookLookup(nil).(*emptyHookLookup), &testBundle{} //nolint:forcetypeassert - } - - t.Run("AlwaysReturnsNil", func(t *testing.T) { - t.Parallel() - - hookLookup, _ := setup(t) - - require.Nil(t, hookLookup.ByHookKind(HookKindInsertBegin)) - require.Nil(t, hookLookup.ByHookKind(HookKindWorkBegin)) - }) -} - -func TestJobHookLookup(t *testing.T) { - t.Parallel() - - type testBundle struct{} - - setup := func(t *testing.T) (*JobHookLookup, *testBundle) { //nolint:unparam - t.Helper() - - return NewJobHookLookup(), &testBundle{} - } - - t.Run("LooksUpHooks", func(t *testing.T) { - t.Parallel() - - jobHookLookup, _ := setup(t) - - require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindInsertBegin)) - require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkBegin)) - require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkEnd)) - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookInsertBegin{}, - }, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindInsertBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookWorkBegin{}, - }, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookWorkEnd{}, - }, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkEnd)) - - require.Len(t, jobHookLookup.hookLookupByKind, 2) - - // Repeat lookups to make sure we get the same result. - require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindInsertBegin)) - require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkBegin)) - require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkEnd)) - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookInsertBegin{}, - }, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindInsertBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookWorkBegin{}, - }, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkBegin)) - require.Equal(t, []rivertype.Hook{ - &testHookWorkEnd{}, - }, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkEnd)) - }) - - t.Run("Stress", func(t *testing.T) { - t.Parallel() - - jobHookLookup, _ := setup(t) - - var wg sync.WaitGroup - - parallelLookupLoop := func(args rivertype.JobArgs) { - wg.Go(func() { - for range 50 { - jobHookLookup.ByJobArgs(args) - } - }) - } - - parallelLookupLoop(&jobArgsNoHooks{}) - parallelLookupLoop(&jobArgsWithCustomHooks{}) - parallelLookupLoop(&jobArgsNoHooks{}) - parallelLookupLoop(&jobArgsWithCustomHooks{}) - - wg.Wait() - }) -} - -// -// jobArgsNoHooks -// - -var _ rivertype.JobArgs = &jobArgsNoHooks{} - -type jobArgsNoHooks struct{} - -func (jobArgsNoHooks) Kind() string { return "no_hooks" } - -// -// jobArgsWithHooks -// - -var ( - _ rivertype.JobArgs = &jobArgsWithCustomHooks{} - _ jobArgsWithHooks = &jobArgsWithCustomHooks{} -) - -type jobArgsWithCustomHooks struct{} - -func (jobArgsWithCustomHooks) Hooks() []rivertype.Hook { - return []rivertype.Hook{ - &testHookInsertAndWorkBegin{}, - &testHookInsertBegin{}, - &testHookWorkBegin{}, - &testHookWorkEnd{}, - } -} - -func (jobArgsWithCustomHooks) Kind() string { return "with_custom_hooks" } - -// -// testHookInsertAndWorkBegin -// - -var ( - _ rivertype.HookInsertBegin = &testHookInsertAndWorkBegin{} - _ rivertype.HookWorkBegin = &testHookInsertAndWorkBegin{} -) - -type testHookInsertAndWorkBegin struct{ rivertype.Hook } - -func (t *testHookInsertAndWorkBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { - return nil -} - -func (t *testHookInsertAndWorkBegin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { - return nil -} - -// -// testHookInsertBegin -// - -var _ rivertype.HookInsertBegin = &testHookInsertBegin{} - -type testHookInsertBegin struct{ rivertype.Hook } - -func (t *testHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { - return nil -} - -// -// testHookWorkBegin -// - -var _ rivertype.HookWorkBegin = &testHookWorkBegin{} - -type testHookWorkBegin struct{ rivertype.Hook } - -func (t *testHookWorkBegin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { - return nil -} - -// -// testHookWorkEnd -// - -var _ rivertype.HookWorkEnd = &testHookWorkEnd{} - -type testHookWorkEnd struct{ rivertype.Hook } - -func (t *testHookWorkEnd) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error { - return nil -} diff --git a/internal/jobexecutor/job_executor.go b/internal/jobexecutor/job_executor.go index 6c5f2eba..328dc313 100644 --- a/internal/jobexecutor/job_executor.go +++ b/internal/jobexecutor/job_executor.go @@ -15,10 +15,9 @@ import ( "github.com/tidwall/gjson" "github.com/riverqueue/river/internal/execution" - "github.com/riverqueue/river/internal/hooklookup" "github.com/riverqueue/river/internal/jobcompleter" "github.com/riverqueue/river/internal/jobstats" - "github.com/riverqueue/river/internal/middlewarelookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" @@ -111,10 +110,9 @@ type JobExecutor struct { ClientRetryPolicy ClientRetryPolicy DefaultClientRetryPolicy ClientRetryPolicy ErrorHandler ErrorHandler - HookLookupByJob *hooklookup.JobHookLookup - HookLookupGlobal hooklookup.HookLookupInterface + PluginLookupByJob *pluginlookup.JobPluginLookup + PluginLookupGlobal pluginlookup.PluginLookupInterface JobRow *rivertype.JobRow - MiddlewareLookupGlobal middlewarelookup.MiddlewareLookupInterface ProducerCallbacks struct { JobDone func(jobRow *rivertype.JobRow) Stuck func(ctx context.Context, jobRow *rivertype.JobRow) @@ -220,8 +218,8 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { doInner := execution.Func(func(ctx context.Context) error { { for _, hook := range append( - e.HookLookupGlobal.ByHookKind(hooklookup.HookKindWorkBegin), - e.WorkUnit.HookLookup(e.HookLookupByJob).ByHookKind(hooklookup.HookKindWorkBegin)..., + e.PluginLookupGlobal.ByKind(pluginlookup.HookKindWorkBegin), + e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.HookKindWorkBegin)..., ) { if err := hook.(rivertype.HookWorkBegin).WorkBegin(ctx, e.JobRow); err != nil { //nolint:forcetypeassert return err @@ -248,8 +246,8 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { { for _, hook := range append( - e.HookLookupGlobal.ByHookKind(hooklookup.HookKindWorkEnd), - e.WorkUnit.HookLookup(e.HookLookupByJob).ByHookKind(hooklookup.HookKindWorkEnd)..., + e.PluginLookupGlobal.ByKind(pluginlookup.HookKindWorkEnd), + e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.HookKindWorkEnd)..., ) { err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, e.JobRow, err) //nolint:forcetypeassert } @@ -258,8 +256,13 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { return err }) + globalMiddleware := make([]rivertype.Middleware, 0, len(e.PluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker))) + for _, plugin := range e.PluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker) { + globalMiddleware = append(globalMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert + } + executeFunc := execution.MiddlewareChain( - e.MiddlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), + globalMiddleware, e.WorkUnit.Middleware(), doInner, e.JobRow, diff --git a/internal/jobexecutor/job_executor_test.go b/internal/jobexecutor/job_executor_test.go index c149dac3..62dadc1b 100644 --- a/internal/jobexecutor/job_executor_test.go +++ b/internal/jobexecutor/job_executor_test.go @@ -10,9 +10,8 @@ import ( "github.com/stretchr/testify/require" - "github.com/riverqueue/river/internal/hooklookup" "github.com/riverqueue/river/internal/jobcompleter" - "github.com/riverqueue/river/internal/middlewarelookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/internal/riverinternaltest/retrypolicytest" @@ -38,8 +37,8 @@ type customizableWorkUnit struct { work func() error } -func (w *customizableWorkUnit) HookLookup(lookup *hooklookup.JobHookLookup) hooklookup.HookLookupInterface { - return hooklookup.NewHookLookup(nil) +func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { + return pluginlookup.NewPluginLookup(nil) } func (w *customizableWorkUnit) Middleware() []rivertype.WorkerMiddleware { @@ -189,10 +188,9 @@ func TestJobExecutor_Execute(t *testing.T) { Completer: bundle.completer, DefaultClientRetryPolicy: &retrypolicytest.RetryPolicyNoJitter{}, ErrorHandler: bundle.errorHandler, - HookLookupByJob: hooklookup.NewJobHookLookup(), - HookLookupGlobal: hooklookup.NewHookLookup(nil), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobRow: bundle.jobRow, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(nil), ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) Stuck func(ctx context.Context, jobRow *rivertype.JobRow) @@ -909,7 +907,7 @@ func TestJobExecutor_Execute(t *testing.T) { executor, bundle := setup(t) // Add a middleware so we can verify it's in the trace too: - executor.MiddlewareLookupGlobal = middlewarelookup.NewMiddlewareLookup([]rivertype.Middleware{ + executor.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ &testMiddleware{ work: func(ctx context.Context, job *rivertype.JobRow, next func(context.Context) error) error { return next(ctx) @@ -1080,7 +1078,7 @@ func TestJobExecutor_Execute(t *testing.T) { workBeginCalled bool workEndCalled bool ) - executor.HookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{ + executor.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { workBeginCalled = true return nil @@ -1110,7 +1108,7 @@ func TestJobExecutor_Execute(t *testing.T) { workEnd1Called bool workEnd2Called bool ) - executor.HookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{ + executor.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { workEnd1Called = true require.EqualError(t, err, "job error") @@ -1144,7 +1142,7 @@ func TestJobExecutor_Execute(t *testing.T) { workEnd1Called bool workEnd2Called bool ) - executor.HookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{ + executor.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { workEnd1Called = true require.EqualError(t, err, "job error") @@ -1181,7 +1179,8 @@ func (f HookWorkBeginFunc) WorkBegin(ctx context.Context, job *rivertype.JobRow) return f(ctx, job) } -func (f HookWorkBeginFunc) IsHook() bool { return true } +func (f HookWorkBeginFunc) IsHook() bool { return true } +func (f HookWorkBeginFunc) IsPlugin() bool { return true } type HookWorkEndFunc func(ctx context.Context, job *rivertype.JobRow, err error) error @@ -1189,13 +1188,15 @@ func (f HookWorkEndFunc) WorkEnd(ctx context.Context, job *rivertype.JobRow, err return f(ctx, job, err) } -func (f HookWorkEndFunc) IsHook() bool { return true } +func (f HookWorkEndFunc) IsHook() bool { return true } +func (f HookWorkEndFunc) IsPlugin() bool { return true } type testMiddleware struct { work func(ctx context.Context, job *rivertype.JobRow, next func(context.Context) error) error } func (m *testMiddleware) IsMiddleware() bool { return true } +func (m *testMiddleware) IsPlugin() bool { return true } func (m *testMiddleware) Work(ctx context.Context, job *rivertype.JobRow, next func(context.Context) error) error { return m.work(ctx, job, next) diff --git a/internal/maintenance/job_rescuer_test.go b/internal/maintenance/job_rescuer_test.go index 508824ee..bc4e457e 100644 --- a/internal/maintenance/job_rescuer_test.go +++ b/internal/maintenance/job_rescuer_test.go @@ -9,7 +9,7 @@ import ( "github.com/stretchr/testify/require" - "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/riverdbtest" "github.com/riverqueue/river/riverdriver" @@ -40,7 +40,7 @@ type callbackWorkUnit struct { timeout time.Duration // defaults to 0, which signals default timeout } -func (w *callbackWorkUnit) HookLookup(cache *hooklookup.JobHookLookup) hooklookup.HookLookupInterface { +func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { return nil } func (w *callbackWorkUnit) Middleware() []rivertype.WorkerMiddleware { return nil } diff --git a/internal/maintenance/periodic_job_enqueuer.go b/internal/maintenance/periodic_job_enqueuer.go index 907fa50c..b82b6cf9 100644 --- a/internal/maintenance/periodic_job_enqueuer.go +++ b/internal/maintenance/periodic_job_enqueuer.go @@ -10,7 +10,7 @@ import ( "github.com/tidwall/sjson" - "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" @@ -89,7 +89,7 @@ type InsertFunc func(ctx context.Context, tx riverdriver.ExecutorTx, insertParam type PeriodicJobEnqueuerConfig struct { AdvisoryLockPrefix int32 - HookLookupGlobal hooklookup.HookLookupInterface + PluginLookupGlobal pluginlookup.PluginLookupInterface // Insert is the function to call to insert jobs into the database. Insert InsertFunc @@ -150,9 +150,9 @@ func NewPeriodicJobEnqueuer(archetype *baseservice.Archetype, config *PeriodicJo nextHandle++ } - hookLookupGlobal := config.HookLookupGlobal - if hookLookupGlobal == nil { - hookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{}) + pluginLookupGlobal := config.PluginLookupGlobal + if pluginLookupGlobal == nil { + pluginLookupGlobal = pluginlookup.NewPluginLookup(nil) } pilot := config.Pilot @@ -163,7 +163,7 @@ func NewPeriodicJobEnqueuer(archetype *baseservice.Archetype, config *PeriodicJo svc := baseservice.Init(archetype, &PeriodicJobEnqueuer{ Config: (&PeriodicJobEnqueuerConfig{ AdvisoryLockPrefix: config.AdvisoryLockPrefix, - HookLookupGlobal: hookLookupGlobal, + PluginLookupGlobal: pluginLookupGlobal, Insert: config.Insert, PeriodicJobs: config.PeriodicJobs, Pilot: pilot, @@ -331,7 +331,7 @@ func (s *PeriodicJobEnqueuer) Start(ctx context.Context) error { return err } - for _, hook := range s.Config.HookLookupGlobal.ByHookKind(hooklookup.HookKindPeriodicJobsStart) { + for _, hook := range s.Config.PluginLookupGlobal.ByKind(pluginlookup.HookKindPeriodicJobsStart) { if err := hook.(rivertype.HookPeriodicJobsStart).Start(ctx, &rivertype.HookPeriodicJobsStartParams{ //nolint:forcetypeassert DurableJobs: sliceutil.Map(initialPeriodicJobs, func(job *riverpilot.PeriodicJob) *rivertype.DurablePeriodicJob { return (*rivertype.DurablePeriodicJob)(job) diff --git a/internal/maintenance/periodic_job_enqueuer_test.go b/internal/maintenance/periodic_job_enqueuer_test.go index 197926d7..3ffcd3ea 100644 --- a/internal/maintenance/periodic_job_enqueuer_test.go +++ b/internal/maintenance/periodic_job_enqueuer_test.go @@ -15,7 +15,7 @@ import ( "github.com/tidwall/gjson" "github.com/riverqueue/river/internal/dbunique" - "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/riverdbtest" "github.com/riverqueue/river/riverdriver" @@ -1055,7 +1055,7 @@ func TestPeriodicJobEnqueuer(t *testing.T) { require.NoError(t, err) var periodicJobsStartHookCalled bool - svc.Config.HookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{ + svc.Config.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { periodicJobsStartHookCalled = true @@ -1117,7 +1117,7 @@ func TestPeriodicJobEnqueuer(t *testing.T) { } var periodicJobsStartHookCalled bool - svc.Config.HookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{ + svc.Config.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { periodicJobsStartHookCalled = true @@ -1182,7 +1182,7 @@ func TestPeriodicJobEnqueuer(t *testing.T) { svc, _ := setup(t) - svc.Config.HookLookupGlobal = hooklookup.NewHookLookup([]rivertype.Hook{ + svc.Config.PluginLookupGlobal = pluginlookup.NewPluginLookup([]rivertype.Plugin{ HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return errors.New("hook start error") }), @@ -1231,7 +1231,8 @@ func (p *PilotPeriodicJobMock) PeriodicJobUpsertMany(ctx context.Context, exec r type HookPeriodicJobsStartFunc func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error -func (f HookPeriodicJobsStartFunc) IsHook() bool { return true } +func (f HookPeriodicJobsStartFunc) IsHook() bool { return true } +func (f HookPeriodicJobsStartFunc) IsPlugin() bool { return true } func (f HookPeriodicJobsStartFunc) Start(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return f(ctx, params) diff --git a/internal/middlewarelookup/middleware_lookup.go b/internal/middlewarelookup/middleware_lookup.go deleted file mode 100644 index 4b45265f..00000000 --- a/internal/middlewarelookup/middleware_lookup.go +++ /dev/null @@ -1,108 +0,0 @@ -package middlewarelookup - -import ( - "sync" - - "github.com/riverqueue/river/rivertype" -) - -// -// MiddlewareKind -// - -type MiddlewareKind string - -const ( - MiddlewareKindJobInsert MiddlewareKind = "job_insert" - MiddlewareKindWorker MiddlewareKind = "worker" -) - -// -// MiddlewareLookupInterface -// - -// MiddlewareLookupInterface is an interface to look up middlewares by -// middleware kind. It's commonly implemented by MiddlewareLookup, but may also -// be EmptyMiddlewareLookup as a memory allocation optimization for bundles -// where no middlewares are present. -type MiddlewareLookupInterface interface { - ByMiddlewareKind(kind MiddlewareKind) []rivertype.Middleware -} - -// NewMiddlewareLookup returns a new middleware lookup interface based on the given middlewares -// that satisfies MiddlewareLookupInterface. This is often middlewareLookup, but may be -// emptyMiddlewareLookup as an optimization for the common case of an empty middleware -// bundle. -func NewMiddlewareLookup(middlewares []rivertype.Middleware) MiddlewareLookupInterface { - if len(middlewares) < 1 { - return &emptyMiddlewareLookup{} - } - - return &middlewareLookup{ - middlewares: middlewares, - middlewaresByKind: make(map[MiddlewareKind][]rivertype.Middleware), - mu: &sync.RWMutex{}, - } -} - -// -// middlewareLookup -// - -// middlewareLookup looks up and caches middlewares based on a MiddlewareKind, saving work when -// looking up middlewares for specific operations, a common operation that gets -// repeated over and over again. This struct may be used as a lookup for -// globally installed middlewares or middlewares for specific job kinds through the use of -// JobMiddlewareLookup. -type middlewareLookup struct { - middlewares []rivertype.Middleware - middlewaresByKind map[MiddlewareKind][]rivertype.Middleware - mu *sync.RWMutex -} - -func (c *middlewareLookup) ByMiddlewareKind(kind MiddlewareKind) []rivertype.Middleware { - c.mu.RLock() - cache, ok := c.middlewaresByKind[kind] - c.mu.RUnlock() - if ok { - return cache - } - - c.mu.Lock() - defer c.mu.Unlock() - - // Even if this ends up being empty, make sure there's an entry for the next - // time the cache gets invoked for this kind. - c.middlewaresByKind[kind] = nil - - // Rely on exhaustlint to find any missing middleware kinds here. - switch kind { - case MiddlewareKindJobInsert: - for _, middleware := range c.middlewares { - if typedMiddleware, ok := middleware.(rivertype.JobInsertMiddleware); ok { - c.middlewaresByKind[kind] = append(c.middlewaresByKind[kind], typedMiddleware) - } - } - case MiddlewareKindWorker: - for _, middleware := range c.middlewares { - if typedMiddleware, ok := middleware.(rivertype.WorkerMiddleware); ok { - c.middlewaresByKind[kind] = append(c.middlewaresByKind[kind], typedMiddleware) - } - } - } - - return c.middlewaresByKind[kind] -} - -// -// emptyMiddlewareLookup -// - -// emptyMiddlewareLookup is an empty version of MiddlewareLookup that's zero allocation. For -// most applications, most job args won't have middlewares, so this prevents us from -// allocating dozens/hundreds of small MiddlewareLookup objects that go unused. -type emptyMiddlewareLookup struct{} - -func (c *emptyMiddlewareLookup) ByMiddlewareKind(kind MiddlewareKind) []rivertype.Middleware { - return nil -} diff --git a/internal/middlewarelookup/middleware_lookup_test.go b/internal/middlewarelookup/middleware_lookup_test.go deleted file mode 100644 index 9b5f9a68..00000000 --- a/internal/middlewarelookup/middleware_lookup_test.go +++ /dev/null @@ -1,141 +0,0 @@ -package middlewarelookup - -import ( - "context" - "sync" - "testing" - - "github.com/stretchr/testify/require" - - "github.com/riverqueue/river/rivertype" -) - -func TestMiddlewareLookup(t *testing.T) { - t.Parallel() - - type testBundle struct{} - - setup := func(t *testing.T) (*middlewareLookup, *testBundle) { //nolint:unparam - t.Helper() - - return NewMiddlewareLookup([]rivertype.Middleware{ //nolint:forcetypeassert - &testMiddlewareJobInsertAndWorker{}, - &testMiddlewareJobInsert{}, - &testMiddlewareWorker{}, - }).(*middlewareLookup), &testBundle{} - } - - t.Run("LooksUpMiddleware", func(t *testing.T) { - t.Parallel() - - middlewareLookup, _ := setup(t) - - require.Equal(t, []rivertype.Middleware{ - &testMiddlewareJobInsertAndWorker{}, - &testMiddlewareJobInsert{}, - }, middlewareLookup.ByMiddlewareKind(MiddlewareKindJobInsert)) - require.Equal(t, []rivertype.Middleware{ - &testMiddlewareJobInsertAndWorker{}, - &testMiddlewareWorker{}, - }, middlewareLookup.ByMiddlewareKind(MiddlewareKindWorker)) - - require.Len(t, middlewareLookup.middlewaresByKind, 2) - - // Repeat lookups to make sure we get the same result. - require.Equal(t, []rivertype.Middleware{ - &testMiddlewareJobInsertAndWorker{}, - &testMiddlewareJobInsert{}, - }, middlewareLookup.ByMiddlewareKind(MiddlewareKindJobInsert)) - require.Equal(t, []rivertype.Middleware{ - &testMiddlewareJobInsertAndWorker{}, - &testMiddlewareWorker{}, - }, middlewareLookup.ByMiddlewareKind(MiddlewareKindWorker)) - }) - - t.Run("Stress", func(t *testing.T) { - t.Parallel() - - middlewareLookup, _ := setup(t) - - var wg sync.WaitGroup - - parallelLookupLoop := func(kind MiddlewareKind) { - wg.Go(func() { - for range 50 { - middlewareLookup.ByMiddlewareKind(kind) - } - }) - } - - parallelLookupLoop(MiddlewareKindJobInsert) - parallelLookupLoop(MiddlewareKindWorker) - parallelLookupLoop(MiddlewareKindJobInsert) - parallelLookupLoop(MiddlewareKindWorker) - - wg.Wait() - }) -} - -func TestEmptyMiddlewareLookup(t *testing.T) { - t.Parallel() - - type testBundle struct{} - - setup := func(t *testing.T) (*emptyMiddlewareLookup, *testBundle) { - t.Helper() - - return NewMiddlewareLookup(nil).(*emptyMiddlewareLookup), &testBundle{} //nolint:forcetypeassert - } - - t.Run("AlwaysReturnsNil", func(t *testing.T) { - t.Parallel() - - middlewareLookup, _ := setup(t) - - require.Nil(t, middlewareLookup.ByMiddlewareKind(MiddlewareKindJobInsert)) - require.Nil(t, middlewareLookup.ByMiddlewareKind(MiddlewareKindWorker)) - }) -} - -// -// testMiddlewareInsertAndWorkBegin -// - -var ( - _ rivertype.JobInsertMiddleware = &testMiddlewareJobInsertAndWorker{} - _ rivertype.WorkerMiddleware = &testMiddlewareJobInsertAndWorker{} -) - -type testMiddlewareJobInsertAndWorker struct{ rivertype.Middleware } - -func (t *testMiddlewareJobInsertAndWorker) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { - return doInner(ctx) -} - -func (t *testMiddlewareJobInsertAndWorker) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { - return doInner(ctx) -} - -// -// testMiddlewareJobInsert -// - -var _ rivertype.JobInsertMiddleware = &testMiddlewareJobInsert{} - -type testMiddlewareJobInsert struct{ rivertype.Middleware } - -func (t *testMiddlewareJobInsert) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { - return doInner(ctx) -} - -// -// testMiddlewareWorker -// - -var _ rivertype.WorkerMiddleware = &testMiddlewareWorker{} - -type testMiddlewareWorker struct{ rivertype.Middleware } - -func (t *testMiddlewareWorker) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { - return doInner(ctx) -} diff --git a/internal/pluginconfig/plugin_config.go b/internal/pluginconfig/plugin_config.go new file mode 100644 index 00000000..4c306902 --- /dev/null +++ b/internal/pluginconfig/plugin_config.go @@ -0,0 +1,43 @@ +package pluginconfig + +import "github.com/riverqueue/river/rivertype" + +// Hooks returns the effective hook list from configured hooks and middleware. +// Explicit hooks are preserved first, followed by middleware that also +// implement hooks. +func Hooks(hooks []rivertype.Hook, middleware []rivertype.Middleware) []rivertype.Hook { + effectiveHooks := make([]rivertype.Hook, 0, len(hooks)+len(middleware)) + + effectiveHooks = append(effectiveHooks, hooks...) + + for _, middlewareItem := range middleware { + hook, ok := middlewareItem.(rivertype.Hook) + if !ok { + continue + } + + effectiveHooks = append(effectiveHooks, hook) + } + + return effectiveHooks +} + +// Middleware returns the effective middleware list from configured hooks and +// middleware. Explicit middleware are preserved first, followed by hooks that +// also implement middleware. +func Middleware(hooks []rivertype.Hook, middleware []rivertype.Middleware) []rivertype.Middleware { + effectiveMiddleware := make([]rivertype.Middleware, 0, len(hooks)+len(middleware)) + + effectiveMiddleware = append(effectiveMiddleware, middleware...) + + for _, hook := range hooks { + middlewareItem, ok := hook.(rivertype.Middleware) + if !ok { + continue + } + + effectiveMiddleware = append(effectiveMiddleware, middlewareItem) + } + + return effectiveMiddleware +} diff --git a/internal/pluginlookup/plugin_lookup.go b/internal/pluginlookup/plugin_lookup.go new file mode 100644 index 00000000..de0c6111 --- /dev/null +++ b/internal/pluginlookup/plugin_lookup.go @@ -0,0 +1,280 @@ +package pluginlookup + +import ( + "context" + "sync" + + "github.com/riverqueue/river/rivertype" +) + +// +// PluginKind +// + +type PluginKind string + +const ( + PluginKindInsertBegin PluginKind = "insert_begin" + PluginKindJobInsert PluginKind = "job_insert" + PluginKindPeriodicJobsStart PluginKind = "periodic_job_start" + PluginKindWorkBegin PluginKind = "work_begin" + PluginKindWorkEnd PluginKind = "work_end" + PluginKindWorker PluginKind = "worker" +) + +type ( + HookKind = PluginKind + MiddlewareKind = PluginKind +) + +const ( + HookKindInsertBegin HookKind = PluginKindInsertBegin + HookKindPeriodicJobsStart HookKind = PluginKindPeriodicJobsStart + HookKindWorkBegin HookKind = PluginKindWorkBegin + HookKindWorkEnd HookKind = PluginKindWorkEnd + MiddlewareKindJobInsert MiddlewareKind = PluginKindJobInsert + MiddlewareKindWorker MiddlewareKind = PluginKindWorker +) + +// +// PluginLookupInterface +// + +// PluginLookupInterface looks up plugins by kind. It's commonly implemented by +// PluginLookup, but may also be EmptyPluginLookup as a memory allocation +// optimization for bundles where no plugins are present. +type PluginLookupInterface interface { + ByKind(kind PluginKind) []rivertype.Plugin +} + +// NewPluginLookup returns a new plugin lookup interface based on the given +// plugins that satisfies PluginLookupInterface. This is often pluginLookup, +// but may be emptyPluginLookup as an optimization for the common case of an +// empty plugin bundle. +func NewPluginLookup(plugins []rivertype.Plugin) PluginLookupInterface { + if len(plugins) < 1 { + return &emptyPluginLookup{} + } + + pluginsByKind := make(map[PluginKind][]rivertype.Plugin) + + for _, plugin := range plugins { + if plugin == nil { + continue + } + + if _, ok := plugin.(rivertype.HookInsertBegin); ok { + pluginsByKind[PluginKindInsertBegin] = append(pluginsByKind[PluginKindInsertBegin], plugin) + } + if _, ok := plugin.(rivertype.HookPeriodicJobsStart); ok { + pluginsByKind[PluginKindPeriodicJobsStart] = append(pluginsByKind[PluginKindPeriodicJobsStart], plugin) + } + if _, ok := plugin.(rivertype.HookWorkBegin); ok { + pluginsByKind[PluginKindWorkBegin] = append(pluginsByKind[PluginKindWorkBegin], plugin) + } + if _, ok := plugin.(rivertype.HookWorkEnd); ok { + pluginsByKind[PluginKindWorkEnd] = append(pluginsByKind[PluginKindWorkEnd], plugin) + } + if _, ok := plugin.(rivertype.JobInsertMiddleware); ok { + pluginsByKind[PluginKindJobInsert] = append(pluginsByKind[PluginKindJobInsert], plugin) + } + if _, ok := plugin.(rivertype.WorkerMiddleware); ok { + pluginsByKind[PluginKindWorker] = append(pluginsByKind[PluginKindWorker], plugin) + } + } + + return &pluginLookup{pluginsByKind: pluginsByKind} +} + +// NormalizePlugins converts hook, middleware, and plugin registrations into a +// single plugin slice while preserving legacy hook and middleware registrations +// that don't yet opt into Plugin. + +// Plugins passed explicitly are ordered before hooks and middleware, so +// Config.Plugins entries take precedence over plugins bridged from Config.Hooks +// or Config.Middleware. Relative order is preserved within each input slice. +func NormalizePlugins(hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Plugin { + normalizedPlugins := make([]rivertype.Plugin, 0, len(hooks)+len(middlewares)+len(plugins)) + + normalizedPlugins = append(normalizedPlugins, plugins...) + + for _, hook := range hooks { + if plugin, ok := hook.(rivertype.Plugin); ok { + normalizedPlugins = append(normalizedPlugins, plugin) + } else { + normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook)) + } + } + + for _, middleware := range middlewares { + if plugin, ok := middleware.(rivertype.Plugin); ok { + normalizedPlugins = append(normalizedPlugins, plugin) + } else { + normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware)) + } + } + + return normalizedPlugins +} + +// +// pluginLookup +// + +// pluginLookup looks up and caches plugins based on their kind, saving work +// when looking up plugin bundles for specific operations, a common operation +// that gets repeated over and over again. This struct may be used as a lookup +// for globally installed plugins or plugins for specific job kinds through the +// use of JobPluginLookup. +type pluginLookup struct { + pluginsByKind map[PluginKind][]rivertype.Plugin +} + +func (c *pluginLookup) ByKind(kind PluginKind) []rivertype.Plugin { + return c.pluginsByKind[kind] +} + +// +// emptyPluginLookup +// + +// emptyPluginLookup is an empty version of PluginLookup that's zero +// allocation. For most applications, most job args won't have plugins, so this +// prevents us from allocating dozens/hundreds of small PluginLookup objects +// that go unused. +type emptyPluginLookup struct{} + +func (c *emptyPluginLookup) ByKind(kind PluginKind) []rivertype.Plugin { + return nil +} + +// +// hookPlugin +// + +// hookPlugin wraps a legacy hook so it can participate in plugin lookup. +type hookPlugin struct { + hook rivertype.Hook +} + +func newHookPlugin(hook rivertype.Hook) rivertype.Plugin { + return &hookPlugin{hook: hook} +} + +func (p *hookPlugin) IsPlugin() bool { return true } +func (p *hookPlugin) IsHook() bool { return true } + +func (p *hookPlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + hook, ok := p.hook.(rivertype.HookInsertBegin) + if !ok { + return nil + } + return hook.InsertBegin(ctx, params) +} + +func (p *hookPlugin) Start(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { + hook, ok := p.hook.(rivertype.HookPeriodicJobsStart) + if !ok { + return nil + } + return hook.Start(ctx, params) +} + +func (p *hookPlugin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { + hook, ok := p.hook.(rivertype.HookWorkBegin) + if !ok { + return nil + } + return hook.WorkBegin(ctx, job) +} + +func (p *hookPlugin) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error { + hook, ok := p.hook.(rivertype.HookWorkEnd) + if !ok { + return err + } + return hook.WorkEnd(ctx, job, err) +} + +// +// middlewarePlugin +// + +// middlewarePlugin wraps a legacy middleware so it can participate in plugin +// lookup. +type middlewarePlugin struct { + middleware rivertype.Middleware +} + +func newMiddlewarePlugin(middleware rivertype.Middleware) rivertype.Plugin { + return &middlewarePlugin{middleware: middleware} +} + +func (p *middlewarePlugin) IsPlugin() bool { return true } +func (p *middlewarePlugin) IsMiddleware() bool { + return true +} + +func (p *middlewarePlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + middleware, ok := p.middleware.(rivertype.JobInsertMiddleware) + if !ok { + return doInner(ctx) + } + return middleware.InsertMany(ctx, manyParams, doInner) +} + +func (p *middlewarePlugin) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { + middleware, ok := p.middleware.(rivertype.WorkerMiddleware) + if !ok { + return doInner(ctx) + } + return middleware.Work(ctx, job, doInner) +} + +// +// JobPluginLookup +// + +type JobPluginLookup struct { + pluginLookupByKind map[string]PluginLookupInterface + mu sync.RWMutex +} + +func NewJobPluginLookup() *JobPluginLookup { + return &JobPluginLookup{ + pluginLookupByKind: make(map[string]PluginLookupInterface), + } +} + +// ByJobArgs returns a PluginLookupInterface by job args, which is a +// PluginLookup if the job args had specific hooks (i.e. implements +// JobArgsWithHooks and returns a non-empty set of hooks), or an +// EmptyPluginLookup otherwise. +func (c *JobPluginLookup) ByJobArgs(args rivertype.JobArgs) PluginLookupInterface { + kind := args.Kind() + + c.mu.RLock() + lookup, ok := c.pluginLookupByKind[kind] + c.mu.RUnlock() + if ok { + return lookup + } + + c.mu.Lock() + defer c.mu.Unlock() + + var hooks []rivertype.Hook + if argsWithHooks, ok := args.(jobArgsWithHooks); ok { + hooks = argsWithHooks.Hooks() + } + + lookup = NewPluginLookup(NormalizePlugins(hooks, nil, nil)) + c.pluginLookupByKind[kind] = lookup + return lookup +} + +// Same as river.JobArgsWithHooks, but duplicated here so that can still live in +// the top level package. +type jobArgsWithHooks interface { + Hooks() []rivertype.Hook +} diff --git a/internal/pluginlookup/plugin_lookup_test.go b/internal/pluginlookup/plugin_lookup_test.go new file mode 100644 index 00000000..a5d0c665 --- /dev/null +++ b/internal/pluginlookup/plugin_lookup_test.go @@ -0,0 +1,471 @@ +package pluginlookup + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/rivertype" +) + +func TestEmptyPluginLookup(t *testing.T) { + t.Parallel() + + type testBundle struct{} + + setup := func(t *testing.T) (*emptyPluginLookup, *testBundle) { + t.Helper() + + lookup, isEmptyLookup := NewPluginLookup(nil).(*emptyPluginLookup) + require.True(t, isEmptyLookup) + + return lookup, &testBundle{} + } + + t.Run("AlwaysReturnsNil", func(t *testing.T) { + t.Parallel() + + pluginLookup, _ := setup(t) + + require.Nil(t, pluginLookup.ByKind(HookKindInsertBegin)) + require.Nil(t, pluginLookup.ByKind(HookKindWorkBegin)) + require.Nil(t, pluginLookup.ByKind(MiddlewareKindJobInsert)) + require.Nil(t, pluginLookup.ByKind(MiddlewareKindWorker)) + }) +} + +func TestJobPluginLookup(t *testing.T) { + t.Parallel() + + type testBundle struct{} + + setup := func(t *testing.T) (*JobPluginLookup, *testBundle) { //nolint:unparam + t.Helper() + + return NewJobPluginLookup(), &testBundle{} + } + + t.Run("LooksUpHooks", func(t *testing.T) { + t.Parallel() + + jobPluginLookup, _ := setup(t) + + require.Nil(t, jobPluginLookup.ByJobArgs(&jobArgsNoHooks{}).ByKind(HookKindInsertBegin)) + require.Nil(t, jobPluginLookup.ByJobArgs(&jobArgsNoHooks{}).ByKind(HookKindWorkBegin)) + require.Nil(t, jobPluginLookup.ByJobArgs(&jobArgsNoHooks{}).ByKind(HookKindWorkEnd)) + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookInsertBegin{}, + }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(HookKindInsertBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookWorkBegin{}, + }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(HookKindWorkBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookWorkEnd{}, + }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(HookKindWorkEnd)) + + require.Len(t, jobPluginLookup.pluginLookupByKind, 2) + + // Repeat lookups to make sure we get the same result. + require.Nil(t, jobPluginLookup.ByJobArgs(&jobArgsNoHooks{}).ByKind(HookKindInsertBegin)) + require.Nil(t, jobPluginLookup.ByJobArgs(&jobArgsNoHooks{}).ByKind(HookKindWorkBegin)) + require.Nil(t, jobPluginLookup.ByJobArgs(&jobArgsNoHooks{}).ByKind(HookKindWorkEnd)) + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookInsertBegin{}, + }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(HookKindInsertBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookWorkBegin{}, + }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(HookKindWorkBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookWorkEnd{}, + }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(HookKindWorkEnd)) + }) + + t.Run("Stress", func(t *testing.T) { + t.Parallel() + + jobPluginLookup, _ := setup(t) + + var wg sync.WaitGroup + + parallelLookupLoop := func(args rivertype.JobArgs) { + wg.Go(func() { + for range 50 { + jobPluginLookup.ByJobArgs(args) + } + }) + } + + parallelLookupLoop(&jobArgsNoHooks{}) + parallelLookupLoop(&jobArgsWithCustomHooks{}) + parallelLookupLoop(&jobArgsNoHooks{}) + parallelLookupLoop(&jobArgsWithCustomHooks{}) + + wg.Wait() + }) +} + +func TestNormalizePlugins(t *testing.T) { + t.Parallel() + + t.Run("PluginsPrecedeHooksAndMiddleware", func(t *testing.T) { + t.Parallel() + + hookPlugin := &testHookMiddlewarePlugin{} + middlewarePlugin := &testHookMiddlewarePlugin{} + plugin := &testHookMiddlewarePlugin{} + + lookup, isPluginLookup := NewPluginLookup(NormalizePlugins( + []rivertype.Hook{hookPlugin}, + []rivertype.Middleware{middlewarePlugin}, + []rivertype.Plugin{plugin}, + )).(*pluginLookup) + require.True(t, isPluginLookup) + + require.Equal(t, []rivertype.Plugin{ + plugin, + hookPlugin, + middlewarePlugin, + }, lookup.ByKind(HookKindInsertBegin)) + require.Equal(t, []rivertype.Plugin{ + plugin, + hookPlugin, + middlewarePlugin, + }, lookup.ByKind(MiddlewareKindJobInsert)) + }) +} + +func TestPluginLookup(t *testing.T) { + t.Parallel() + + type testBundle struct{} + + setup := func(t *testing.T) (*pluginLookup, *testBundle) { //nolint:unparam + t.Helper() + + lookup, isPluginLookup := NewPluginLookup([]rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookInsertBegin{}, + &testHookWorkBegin{}, + &testHookWorkEnd{}, + &testMiddlewareJobInsertAndWorker{}, + &testMiddlewareJobInsert{}, + &testMiddlewareWorker{}, + }).(*pluginLookup) + require.True(t, isPluginLookup) + + return lookup, &testBundle{} + } + + t.Run("LooksUpPlugins", func(t *testing.T) { + t.Parallel() + + pluginLookup, _ := setup(t) + + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookInsertBegin{}, + }, pluginLookup.ByKind(HookKindInsertBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookWorkBegin{}, + }, pluginLookup.ByKind(HookKindWorkBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookWorkEnd{}, + }, pluginLookup.ByKind(HookKindWorkEnd)) + + require.Equal(t, []rivertype.Plugin{ + &testMiddlewareJobInsertAndWorker{}, + &testMiddlewareJobInsert{}, + }, pluginLookup.ByKind(MiddlewareKindJobInsert)) + require.Equal(t, []rivertype.Plugin{ + &testMiddlewareJobInsertAndWorker{}, + &testMiddlewareWorker{}, + }, pluginLookup.ByKind(MiddlewareKindWorker)) + + require.Len(t, pluginLookup.pluginsByKind, 5) + + // Repeat lookups to make sure we get the same result. + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookInsertBegin{}, + }, pluginLookup.ByKind(HookKindInsertBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookInsertAndWorkBegin{}, + &testHookWorkBegin{}, + }, pluginLookup.ByKind(HookKindWorkBegin)) + require.Equal(t, []rivertype.Plugin{ + &testHookWorkEnd{}, + }, pluginLookup.ByKind(HookKindWorkEnd)) + require.Equal(t, []rivertype.Plugin{ + &testMiddlewareJobInsertAndWorker{}, + &testMiddlewareJobInsert{}, + }, pluginLookup.ByKind(MiddlewareKindJobInsert)) + require.Equal(t, []rivertype.Plugin{ + &testMiddlewareJobInsertAndWorker{}, + &testMiddlewareWorker{}, + }, pluginLookup.ByKind(MiddlewareKindWorker)) + }) + + t.Run("Stress", func(t *testing.T) { + t.Parallel() + + pluginLookup, _ := setup(t) + + var wg sync.WaitGroup + + parallelLookupLoop := func(kind HookKind) { + wg.Go(func() { + for range 50 { + pluginLookup.ByKind(kind) + } + }) + } + + parallelLookupLoop(HookKindInsertBegin) + parallelLookupLoop(HookKindWorkBegin) + parallelLookupLoop(HookKindInsertBegin) + parallelLookupLoop(HookKindWorkBegin) + + wg.Wait() + }) + + t.Run("LooksUpLegacyHooksAndMiddleware", func(t *testing.T) { + t.Parallel() + + legacyHook := &testLegacyHookInsertBegin{} + legacyMiddleware := &testLegacyMiddlewareJobInsert{} + + lookup, isPluginLookup := NewPluginLookup(NormalizePlugins( + []rivertype.Hook{legacyHook}, + []rivertype.Middleware{legacyMiddleware}, + nil, + )).(*pluginLookup) + require.True(t, isPluginLookup) + + hookPlugins := lookup.ByKind(HookKindInsertBegin) + require.Len(t, hookPlugins, 1) + + hook, isHookInsertBegin := hookPlugins[0].(rivertype.HookInsertBegin) + require.True(t, isHookInsertBegin) + require.NoError(t, hook.InsertBegin(context.Background(), &rivertype.JobInsertParams{})) + require.True(t, legacyHook.insertBeginCalled) + + middlewarePlugins := lookup.ByKind(MiddlewareKindJobInsert) + require.Len(t, middlewarePlugins, 1) + + middleware, isJobInsertMiddleware := middlewarePlugins[0].(rivertype.JobInsertMiddleware) + require.True(t, isJobInsertMiddleware) + _, err := middleware.InsertMany(context.Background(), []*rivertype.JobInsertParams{{}}, func(context.Context) ([]*rivertype.JobInsertResult, error) { + legacyMiddleware.doInnerCalled = true + return nil, nil + }) + require.NoError(t, err) + require.True(t, legacyMiddleware.doInnerCalled) + }) +} + +// +// jobArgsNoHooks +// + +var _ rivertype.JobArgs = &jobArgsNoHooks{} + +type jobArgsNoHooks struct{} + +func (jobArgsNoHooks) Kind() string { return "no_hooks" } + +// +// jobArgsWithHooks +// + +var ( + _ rivertype.JobArgs = &jobArgsWithCustomHooks{} + _ jobArgsWithHooks = &jobArgsWithCustomHooks{} +) + +type jobArgsWithCustomHooks struct{} + +func (jobArgsWithCustomHooks) Hooks() []rivertype.Hook { + return []rivertype.Hook{ + &testHookInsertAndWorkBegin{}, + &testHookInsertBegin{}, + &testHookWorkBegin{}, + &testHookWorkEnd{}, + } +} + +func (jobArgsWithCustomHooks) Kind() string { return "with_custom_hooks" } + +// +// testHookInsertAndWorkBegin +// + +var ( + _ rivertype.HookInsertBegin = &testHookInsertAndWorkBegin{} + _ rivertype.HookWorkBegin = &testHookInsertAndWorkBegin{} +) + +type testHookInsertAndWorkBegin struct{ rivertype.Hook } + +func (t *testHookInsertAndWorkBegin) IsPlugin() bool { return true } + +func (t *testHookInsertAndWorkBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + return nil +} + +func (t *testHookInsertAndWorkBegin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { + return nil +} + +// +// testHookInsertBegin +// + +var _ rivertype.HookInsertBegin = &testHookInsertBegin{} + +type testHookInsertBegin struct{ rivertype.Hook } + +func (t *testHookInsertBegin) IsPlugin() bool { return true } + +func (t *testHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + return nil +} + +// +// testHookMiddlewarePlugin +// + +var ( + _ rivertype.HookInsertBegin = &testHookMiddlewarePlugin{} + _ rivertype.JobInsertMiddleware = &testHookMiddlewarePlugin{} +) + +type testHookMiddlewarePlugin struct{} + +func (t *testHookMiddlewarePlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + return nil +} + +func (t *testHookMiddlewarePlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return doInner(ctx) +} + +func (t *testHookMiddlewarePlugin) IsHook() bool { return true } +func (t *testHookMiddlewarePlugin) IsMiddleware() bool { return true } +func (t *testHookMiddlewarePlugin) IsPlugin() bool { return true } + +// +// testHookWorkBegin +// + +var _ rivertype.HookWorkBegin = &testHookWorkBegin{} + +type testHookWorkBegin struct{ rivertype.Hook } + +func (t *testHookWorkBegin) IsPlugin() bool { return true } + +func (t *testHookWorkBegin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { + return nil +} + +// +// testHookWorkEnd +// + +var _ rivertype.HookWorkEnd = &testHookWorkEnd{} + +type testHookWorkEnd struct{ rivertype.Hook } + +func (t *testHookWorkEnd) IsPlugin() bool { return true } + +func (t *testHookWorkEnd) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error { + return nil +} + +// +// testMiddlewareInsertAndWorkBegin +// + +var ( + _ rivertype.JobInsertMiddleware = &testMiddlewareJobInsertAndWorker{} + _ rivertype.WorkerMiddleware = &testMiddlewareJobInsertAndWorker{} +) + +type testMiddlewareJobInsertAndWorker struct{ rivertype.Middleware } + +func (t *testMiddlewareJobInsertAndWorker) IsPlugin() bool { return true } + +func (t *testMiddlewareJobInsertAndWorker) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return doInner(ctx) +} + +func (t *testMiddlewareJobInsertAndWorker) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { + return doInner(ctx) +} + +// +// testMiddlewareJobInsert +// + +var _ rivertype.JobInsertMiddleware = &testMiddlewareJobInsert{} + +type testMiddlewareJobInsert struct{ rivertype.Middleware } + +func (t *testMiddlewareJobInsert) IsPlugin() bool { return true } + +func (t *testMiddlewareJobInsert) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return doInner(ctx) +} + +// +// testMiddlewareWorker +// + +var _ rivertype.WorkerMiddleware = &testMiddlewareWorker{} + +type testMiddlewareWorker struct{ rivertype.Middleware } + +func (t *testMiddlewareWorker) IsPlugin() bool { return true } + +func (t *testMiddlewareWorker) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { + return doInner(ctx) +} + +// +// testLegacyHookInsertBegin +// + +var _ rivertype.HookInsertBegin = &testLegacyHookInsertBegin{} + +type testLegacyHookInsertBegin struct { + insertBeginCalled bool +} + +func (t *testLegacyHookInsertBegin) IsHook() bool { return true } + +func (t *testLegacyHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + t.insertBeginCalled = true + return nil +} + +// +// testLegacyMiddlewareJobInsert +// + +var _ rivertype.JobInsertMiddleware = &testLegacyMiddlewareJobInsert{} + +type testLegacyMiddlewareJobInsert struct { + doInnerCalled bool +} + +func (t *testLegacyMiddlewareJobInsert) IsMiddleware() bool { return true } + +func (t *testLegacyMiddlewareJobInsert) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return doInner(ctx) +} diff --git a/internal/rivermiddleware/middleware.go b/internal/riverplugin/plugin.go similarity index 89% rename from internal/rivermiddleware/middleware.go rename to internal/riverplugin/plugin.go index 81a6caa2..9cce6d39 100644 --- a/internal/rivermiddleware/middleware.go +++ b/internal/riverplugin/plugin.go @@ -1,4 +1,4 @@ -package rivermiddleware +package riverplugin import ( "context" @@ -13,10 +13,10 @@ import ( "github.com/riverqueue/river/rivertype" ) -// DefaultMiddleware returns the default middleware that River applies to all -// jobs. This includes internal middleware like the resumable step middleware. -func DefaultMiddleware() []rivertype.Middleware { - return []rivertype.Middleware{&ResumableMiddleware{}} +// DefaultPlugins returns the default plugins that River applies to all jobs. +// This includes internal middleware like the resumable step middleware. +func DefaultPlugins() []rivertype.Plugin { + return []rivertype.Plugin{&ResumableMiddleware{}} } // ResumableMiddleware is internal middleware that enables resumable step @@ -26,6 +26,7 @@ func DefaultMiddleware() []rivertype.Middleware { type ResumableMiddleware struct{} func (*ResumableMiddleware) IsMiddleware() bool { return true } +func (*ResumableMiddleware) IsPlugin() bool { return true } func (*ResumableMiddleware) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) diff --git a/internal/workunit/work_unit.go b/internal/workunit/work_unit.go index dbc05345..3d6ff458 100644 --- a/internal/workunit/work_unit.go +++ b/internal/workunit/work_unit.go @@ -4,7 +4,7 @@ import ( "context" "time" - "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/rivertype" ) @@ -16,10 +16,10 @@ import ( // // Implemented by river.wrapperWorkUnit. type WorkUnit interface { - // HookLookup procures the a hook lookup bundle for the wrapped job using + // PluginLookup procures the a hook lookup bundle for the wrapped job using // the given job hook lookup bundle. Hooks are looked up by job args and // otherwise not available to jobexecutor. - HookLookup(lookup *hooklookup.JobHookLookup) hooklookup.HookLookupInterface + PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface Middleware() []rivertype.WorkerMiddleware NextRetry() time.Time diff --git a/middleware_defaults.go b/middleware_defaults.go index 2645e0e6..3633449b 100644 --- a/middleware_defaults.go +++ b/middleware_defaults.go @@ -7,12 +7,15 @@ import ( ) // MiddlewareDefaults should be embedded on any middleware implementation. It -// helps identify a struct as middleware, and guarantees forward compatibility in -// case additions are necessary to the rivertype.Middleware interface. +// helps identify a struct as middleware and a plugin, and guarantees forward +// compatibility in case additions are necessary to the rivertype.Middleware +// interface. type MiddlewareDefaults struct{} func (d *MiddlewareDefaults) IsMiddleware() bool { return true } +func (d *MiddlewareDefaults) IsPlugin() bool { return true } + // JobInsertMiddlewareDefaults is an embeddable struct that provides default // implementations for the rivertype.JobInsertMiddleware. Use of this struct is // recommended in case rivertype.JobInsertMiddleware is expanded in the future @@ -35,6 +38,8 @@ func (f JobInsertMiddlewareFunc) InsertMany(ctx context.Context, manyParams []*r func (f JobInsertMiddlewareFunc) IsMiddleware() bool { return true } +func (f JobInsertMiddlewareFunc) IsPlugin() bool { return true } + // WorkerInsertMiddlewareDefaults is an embeddable struct that provides default // implementations for the rivertype.WorkerMiddleware. Use of this struct is // recommended in case rivertype.WorkerMiddleware is expanded in the future so @@ -51,8 +56,10 @@ func (d *WorkerMiddlewareDefaults) Work(ctx context.Context, job *rivertype.JobR // rivertype.WorkerMiddleware using a simple function instead of a struct. type WorkerMiddlewareFunc func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error +func (f WorkerMiddlewareFunc) IsMiddleware() bool { return true } + +func (f WorkerMiddlewareFunc) IsPlugin() bool { return true } + func (f WorkerMiddlewareFunc) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { return f(ctx, job, doInner) } - -func (f WorkerMiddlewareFunc) IsMiddleware() bool { return true } diff --git a/middleware_defaults_test.go b/middleware_defaults_test.go index 8d469a74..657b0077 100644 --- a/middleware_defaults_test.go +++ b/middleware_defaults_test.go @@ -9,6 +9,8 @@ import ( var ( _ rivertype.JobInsertMiddleware = &JobInsertMiddlewareDefaults{} _ rivertype.WorkerMiddleware = &WorkerMiddlewareDefaults{} + _ rivertype.Plugin = &JobInsertMiddlewareDefaults{} + _ rivertype.Plugin = &WorkerMiddlewareDefaults{} _ rivertype.JobInsertMiddleware = JobInsertMiddlewareFunc(func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { return doInner(ctx) @@ -16,6 +18,9 @@ var ( _ rivertype.Middleware = JobInsertMiddlewareFunc(func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { return doInner(ctx) }) + _ rivertype.Plugin = JobInsertMiddlewareFunc(func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + return doInner(ctx) + }) _ rivertype.Middleware = WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { return doInner(ctx) @@ -23,4 +28,7 @@ var ( _ rivertype.WorkerMiddleware = WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { return doInner(ctx) }) + _ rivertype.Plugin = WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { + return doInner(ctx) + }) ) diff --git a/plugin_defaults.go b/plugin_defaults.go new file mode 100644 index 00000000..2930e479 --- /dev/null +++ b/plugin_defaults.go @@ -0,0 +1,12 @@ +package river + +// PluginDefaults should be embedded on plugin implementations. It helps +// identify a struct as both hook and middleware, and guarantees forward +// compatibility in case additions are necessary to the rivertype.Hook or +// rivertype.Middleware interfaces. +type PluginDefaults struct { + HookDefaults + MiddlewareDefaults +} + +func (d *PluginDefaults) IsPlugin() bool { return true } diff --git a/plugin_defaults_test.go b/plugin_defaults_test.go new file mode 100644 index 00000000..64f15598 --- /dev/null +++ b/plugin_defaults_test.go @@ -0,0 +1,11 @@ +package river + +import "github.com/riverqueue/river/rivertype" + +var ( + _ rivertype.Hook = &PluginDefaults{} + _ rivertype.Middleware = &PluginDefaults{} + _ rivertype.Plugin = &PluginDefaults{} + _ rivertype.Plugin = &HookDefaults{} + _ rivertype.Plugin = &MiddlewareDefaults{} +) diff --git a/producer.go b/producer.go index 07554ac9..b6b0267b 100644 --- a/producer.go +++ b/producer.go @@ -13,11 +13,10 @@ import ( "sync/atomic" "time" - "github.com/riverqueue/river/internal/hooklookup" "github.com/riverqueue/river/internal/jobcompleter" "github.com/riverqueue/river/internal/jobexecutor" - "github.com/riverqueue/river/internal/middlewarelookup" "github.com/riverqueue/river/internal/notifier" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/util/chanutil" "github.com/riverqueue/river/internal/workunit" @@ -82,14 +81,13 @@ type producerConfig struct { // LISTEN/NOTIFY, but this provides a fallback. FetchPollInterval time.Duration - HookLookupByJob *hooklookup.JobHookLookup - HookLookupGlobal hooklookup.HookLookupInterface - JobStuckHandler JobStuckHandler - JobStuckCount *atomic.Int32 - JobStuckThreshold time.Duration - JobTimeout time.Duration - MaxWorkers int - MiddlewareLookupGlobal middlewarelookup.MiddlewareLookupInterface + PluginLookupByJob *pluginlookup.JobPluginLookup + PluginLookupGlobal pluginlookup.PluginLookupInterface + JobStuckHandler JobStuckHandler + JobStuckCount *atomic.Int32 + JobStuckThreshold time.Duration + JobTimeout time.Duration + MaxWorkers int // Notifier is a notifier for subscribing to new job inserts and job // control. If nil, the producer will operate in poll-only mode. @@ -874,9 +872,8 @@ func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype. Completer: p.completer, DefaultClientRetryPolicy: &DefaultClientRetryPolicy{}, ErrorHandler: p.errorHandler, - HookLookupByJob: p.config.HookLookupByJob, - HookLookupGlobal: p.config.HookLookupGlobal, - MiddlewareLookupGlobal: p.config.MiddlewareLookupGlobal, + PluginLookupByJob: p.config.PluginLookupByJob, + PluginLookupGlobal: p.config.PluginLookupGlobal, JobRow: job, ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) diff --git a/producer_test.go b/producer_test.go index afac80ae..69a2f8bf 100644 --- a/producer_test.go +++ b/producer_test.go @@ -10,11 +10,10 @@ import ( "github.com/stretchr/testify/require" - "github.com/riverqueue/river/internal/hooklookup" "github.com/riverqueue/river/internal/jobcompleter" "github.com/riverqueue/river/internal/maintenance" - "github.com/riverqueue/river/internal/middlewarelookup" "github.com/riverqueue/river/internal/notifier" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/internal/riverinternaltest/sharedtx" @@ -100,11 +99,10 @@ func Test_Producer_CanSafelyCompleteJobsWhileFetchingNewOnes(t *testing.T) { // Fetch constantly to more aggressively trigger the potential data race: FetchCooldown: time.Millisecond, FetchPollInterval: time.Millisecond, - HookLookupByJob: hooklookup.NewJobHookLookup(), - HookLookupGlobal: hooklookup.NewHookLookup(nil), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobTimeout: JobTimeoutDefault, MaxWorkers: 1000, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(nil), Notifier: notifier, Queue: rivercommon.QueueDefault, QueuePollInterval: queuePollIntervalDefault, @@ -196,11 +194,10 @@ func TestProducer_PollOnly(t *testing.T) { ErrorHandler: newTestErrorHandler(), FetchCooldown: FetchCooldownDefault, FetchPollInterval: 50 * time.Millisecond, // more aggressive than normal because we have no notifier - HookLookupByJob: hooklookup.NewJobHookLookup(), - HookLookupGlobal: hooklookup.NewHookLookup(nil), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobTimeout: JobTimeoutDefault, MaxWorkers: 1_000, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(nil), Notifier: nil, // no notifier Queue: queueName, QueuePollInterval: queuePollIntervalDefault, @@ -250,11 +247,10 @@ func TestProducer_WithNotifier(t *testing.T) { ErrorHandler: newTestErrorHandler(), FetchCooldown: FetchCooldownDefault, FetchPollInterval: 50 * time.Millisecond, // more aggressive than normal so in case we miss the event, tests still pass quickly - HookLookupByJob: hooklookup.NewJobHookLookup(), - HookLookupGlobal: hooklookup.NewHookLookup(nil), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobTimeout: JobTimeoutDefault, MaxWorkers: 1_000, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(nil), Notifier: notifier, Queue: queueName, QueuePollInterval: queuePollIntervalDefault, diff --git a/resumable.go b/resumable.go index 447cff63..272b3d9c 100644 --- a/resumable.go +++ b/resumable.go @@ -6,7 +6,7 @@ import ( "errors" "fmt" - "github.com/riverqueue/river/internal/rivermiddleware" + "github.com/riverqueue/river/internal/riverplugin" ) var ( @@ -150,7 +150,7 @@ func ResumableStepCursor[TCursor any](ctx context.Context, name string, opts *St delete(state.Cursors, name) } -func mustResumableState(ctx context.Context) *rivermiddleware.ResumableState { +func mustResumableState(ctx context.Context) *riverplugin.ResumableState { state, ok := resumableStateFromContext(ctx) if !ok { panic(errResumableStepNotInWorker) @@ -159,7 +159,7 @@ func mustResumableState(ctx context.Context) *rivermiddleware.ResumableState { return state } -func registerResumableStepName(state *rivermiddleware.ResumableState, name string) bool { +func registerResumableStepName(state *riverplugin.ResumableState, name string) bool { if _, ok := state.AllStepNames[name]; ok { state.Err = fmt.Errorf("river: duplicate resumable step name %q", name) return false @@ -169,13 +169,13 @@ func registerResumableStepName(state *rivermiddleware.ResumableState, name strin return true } -func resumableStateFromContext(ctx context.Context) (*rivermiddleware.ResumableState, bool) { - state := ctx.Value(rivermiddleware.ResumableContextKey{}) +func resumableStateFromContext(ctx context.Context) (*riverplugin.ResumableState, bool) { + state := ctx.Value(riverplugin.ResumableContextKey{}) if state == nil { return nil, false } - typedState, ok := state.(*rivermiddleware.ResumableState) + typedState, ok := state.(*riverplugin.ResumableState) if !ok || typedState == nil { return nil, false } diff --git a/resumable_step_tx_test.go b/resumable_step_tx_test.go index 3202759f..24bc89a2 100644 --- a/resumable_step_tx_test.go +++ b/resumable_step_tx_test.go @@ -11,7 +11,7 @@ import ( "github.com/riverqueue/river/internal/execution" "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/rivercommon" - "github.com/riverqueue/river/internal/rivermiddleware" + "github.com/riverqueue/river/internal/riverplugin" "github.com/riverqueue/river/riverdbtest" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riverpgxv5" @@ -47,7 +47,7 @@ func TestResumableSetStepTx(t *testing.T) { require.NoError(t, err) ctx = context.WithValue(ctx, rivercommon.ContextKeyClient{}, client) ctx = context.WithValue(ctx, jobexecutor.ContextKeyMetadataUpdates, make(map[string]any)) - ctx = context.WithValue(ctx, rivermiddleware.ResumableContextKey{}, &rivermiddleware.ResumableState{ + ctx = context.WithValue(ctx, riverplugin.ResumableContextKey{}, &riverplugin.ResumableState{ Cursors: make(map[string]json.RawMessage), StepName: stepName, }) diff --git a/resumable_test.go b/resumable_test.go index 8c17b336..ce59bf9b 100644 --- a/resumable_test.go +++ b/resumable_test.go @@ -10,7 +10,7 @@ import ( "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/rivercommon" - "github.com/riverqueue/river/internal/rivermiddleware" + "github.com/riverqueue/river/internal/riverplugin" "github.com/riverqueue/river/rivertype" ) @@ -32,7 +32,7 @@ func TestResumableStep(t *testing.T) { ctx, _, job := setup(t, `{}`) var ran []string - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, "first") return nil @@ -54,7 +54,7 @@ func TestResumableStep(t *testing.T) { ctx, _, job := setup(t, `{"river:resumable_step":"step2"}`) var ran []string - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, "first") return nil @@ -88,7 +88,7 @@ func TestResumableStep(t *testing.T) { ctx, metadataUpdates, job := setup(t, `{}`) var ran []string - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, "step1") return nil @@ -111,7 +111,7 @@ func TestResumableStep(t *testing.T) { ctx, metadataUpdates, job = setup(t, `{"river:resumable_step":"step1"}`) ran = nil - err = (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err = (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, "step1") return nil @@ -140,7 +140,7 @@ func TestResumableStep(t *testing.T) { defer cancel() var ran []string - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, "step1") cancel() @@ -185,7 +185,7 @@ func TestResumableStepCursor(t *testing.T) { ctx, _, job := setup(t, `{}`) var ran []string - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, "step") return nil @@ -210,7 +210,7 @@ func TestResumableStepCursor(t *testing.T) { cursorResult resumableCursor ran []int ) - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, 1) return nil @@ -240,7 +240,7 @@ func TestResumableStepCursor(t *testing.T) { cursorResult = resumableCursor{} ran = nil - err = (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err = (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, 1) return nil @@ -271,7 +271,7 @@ func TestResumableStepCursor(t *testing.T) { ctx, metadataUpdates, job := setup(t, `{"river:resumable_cursor":{"step2":{"id":42}},"river:resumable_step":"step1"}`) var ran []int - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStep(ctx, "step1", nil, func(ctx context.Context) error { ran = append(ran, 1) return nil @@ -306,7 +306,7 @@ func TestResumableStepCursor(t *testing.T) { ctx, _, _ := setup(t, `{}`) - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, &rivertype.JobRow{Metadata: []byte(`{}`)}, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, &rivertype.JobRow{Metadata: []byte(`{}`)}, func(ctx context.Context) error { return ResumableSetCursor(ctx, 1) }) require.ErrorIs(t, err, errResumableCursorNotInStep) @@ -321,7 +321,7 @@ func TestResumableStepCursor(t *testing.T) { ctx, metadataUpdates, job := setup(t, `{}`) - err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStepCursor(ctx, "step1", nil, func(ctx context.Context, cursor int) error { require.Zero(t, cursor) require.NoError(t, ResumableSetCursor(ctx, 123)) @@ -343,7 +343,7 @@ func TestResumableStepCursor(t *testing.T) { ctx, metadataUpdates, job = setup(t, `{"river:resumable_cursor":{"step1":123,"step2":{"id":"abc"}},"river:resumable_step":"step1"}`) - err = (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { + err = (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error { ResumableStepCursor(ctx, "step1", nil, func(ctx context.Context, cursor int) error { require.Equal(t, 123, cursor) return nil diff --git a/riverlog/river_log.go b/riverlog/river_log.go index 9d2811b3..307af99f 100644 --- a/riverlog/river_log.go +++ b/riverlog/river_log.go @@ -69,6 +69,8 @@ type Middleware struct { newSlogHandler func(w io.Writer) slog.Handler } +func (m *Middleware) IsPlugin() bool { return true } + // MiddlewareConfig is configuration for Middleware. type MiddlewareConfig struct { // MaxSizeBytes is the maximum size of log data that'll be persisted in diff --git a/rivertest/work_unit_wrapper.go b/rivertest/work_unit_wrapper.go index 1dd3f293..e36e4979 100644 --- a/rivertest/work_unit_wrapper.go +++ b/rivertest/work_unit_wrapper.go @@ -6,7 +6,7 @@ import ( "time" "github.com/riverqueue/river" - "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/rivertype" ) @@ -45,7 +45,7 @@ type wrapperWorkUnit[T river.JobArgs] struct { worker river.Worker[T] } -func (w *wrapperWorkUnit[T]) HookLookup(lookup *hooklookup.JobHookLookup) hooklookup.HookLookupInterface { +func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { var job T return lookup.ByJobArgs(job) } diff --git a/rivertest/worker.go b/rivertest/worker.go index da4f13f5..9377eb5f 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -8,12 +8,11 @@ import ( "github.com/riverqueue/river" "github.com/riverqueue/river/internal/execution" - "github.com/riverqueue/river/internal/hooklookup" "github.com/riverqueue/river/internal/jobcompleter" "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/maintenance" - "github.com/riverqueue/river/internal/middlewarelookup" - "github.com/riverqueue/river/internal/rivermiddleware" + "github.com/riverqueue/river/internal/pluginlookup" + "github.com/riverqueue/river/internal/riverplugin" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/riversharedtest" @@ -147,13 +146,17 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job } completer := jobcompleter.NewInlineCompleter(archetype, w.config.Schema, exec, w.client.Pilot(), subscribeCh) - for _, hook := range w.config.Hooks { - if withBaseService, ok := hook.(baseservice.WithBaseService); ok { - baseservice.Init(archetype, withBaseService) - } - } - for _, middleware := range w.config.Middleware { - if withBaseService, ok := middleware.(baseservice.WithBaseService); ok { + allPlugins := pluginlookup.NormalizePlugins( + w.config.Hooks, //nolint:staticcheck // Bridge legacy Hooks into the test worker's unified plugin path. + middlewareFromConfig(w.config), + append( + riverplugin.DefaultPlugins(), + w.config.Plugins..., + ), + ) + + for _, plugin := range allPlugins { + if withBaseService, ok := plugin.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } } @@ -205,10 +208,9 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job return nil }, }, - HookLookupGlobal: hooklookup.NewHookLookup(w.config.Hooks), - HookLookupByJob: hooklookup.NewJobHookLookup(), - JobRow: job, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(append(rivermiddleware.DefaultMiddleware(), w.config.Middleware...)), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupGlobal: pluginlookup.NewPluginLookup(allPlugins), + JobRow: job, ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) Stuck func(ctx context.Context, jobRow *rivertype.JobRow) @@ -240,6 +242,36 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job panic("unreachable") } +//nolint:staticcheck +func middlewareFromConfig(config *river.Config) []rivertype.Middleware { + middleware := make([]rivertype.Middleware, 0, + len(config.Middleware)+len(config.JobInsertMiddleware)+len(config.WorkerMiddleware)) + middleware = append(middleware, config.Middleware...) + + for _, jobInsertMiddleware := range config.JobInsertMiddleware { + middleware = append(middleware, jobInsertMiddleware) + } + +outerLoop: + for _, workerMiddleware := range config.WorkerMiddleware { + // Don't add the middleware if it also implements JobInsertMiddleware + // and the instance has been added to config.JobInsertMiddleware. This + // is a hedge to make sure we don't accidentally double add middleware + // as we've converted over to the unified config.Middleware setting. + if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddleware.(rivertype.JobInsertMiddleware); ok { + for _, jobInsertMiddleware := range config.JobInsertMiddleware { + if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddleware { + continue outerLoop + } + } + } + + middleware = append(middleware, workerMiddleware) + } + + return middleware +} + // WorkResult is the result of working a job in the test Worker. type WorkResult struct { // EventKind is the kind of event that occurred following execution. diff --git a/rivertest/worker_test.go b/rivertest/worker_test.go index ced315c5..5024d975 100644 --- a/rivertest/worker_test.go +++ b/rivertest/worker_test.go @@ -389,7 +389,7 @@ func TestWorker_Work(t *testing.T) { return nil } - bundle.config.Hooks = []rivertype.Hook{ + bundle.config.Plugins = []rivertype.Plugin{ hookWithBaseService, river.HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { hookCalled = true @@ -430,9 +430,9 @@ func TestWorker_Work(t *testing.T) { return doInner(ctx) } - bundle.config.Middleware = []rivertype.Middleware{ + bundle.config.Plugins = []rivertype.Plugin{ middlewareWithBaseService, - river.WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error { + river.WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { middlewareCalled = true return doInner(ctx) }), diff --git a/rivertype/river_type.go b/rivertype/river_type.go index 13d07f63..d9795d97 100644 --- a/rivertype/river_type.go +++ b/rivertype/river_type.go @@ -447,6 +447,17 @@ type Middleware interface { IsMiddleware() bool } +// Plugin is a generic extension installed globally. +// +// Plugin structs should embed river.PluginDefaults, or embed either +// river.HookDefaults or river.MiddlewareDefaults, then implement any +// operation-specific hook or middleware interfaces they need. If a type embeds +// both defaults directly, it must also implement IsPlugin itself to avoid the +// embedded method-name collision. +type Plugin interface { + IsPlugin() bool +} + // JobInsertMiddleware provides an interface for middleware that integrations // can use to encapsulate common logic around job insertion. // diff --git a/work_unit_wrapper.go b/work_unit_wrapper.go index d0fcabc2..e05b5749 100644 --- a/work_unit_wrapper.go +++ b/work_unit_wrapper.go @@ -5,7 +5,7 @@ import ( "encoding/json" "time" - "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/rivertype" ) @@ -26,7 +26,7 @@ type wrapperWorkUnit[T JobArgs] struct { worker Worker[T] } -func (w *wrapperWorkUnit[T]) HookLookup(lookup *hooklookup.JobHookLookup) hooklookup.HookLookupInterface { +func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { var job T return lookup.ByJobArgs(job) }