From 6d60ff7874e16e3cb6b06038645a28b33d85c307 Mon Sep 17 00:00:00 2001 From: Brandur Date: Fri, 12 Jun 2026 07:37:44 -0500 Subject: [PATCH 1/6] Cross-compatible hook/middleware + "plugins" which are allowed to be hook and middleware This one's largely aimed at extending a few parts of `otelriver` to be able to emit some additional useful metrics like time that it takes to lock jobs, number of jobs locked per batch, or any other arbitrary metrics we want to emit down the road. I previously had something similar back in #1203, but here we extract an isolated piece of it. This change does two things: * If either a hook sent to `Config.Hooks` implements a middleware or a middleware sent to `Config.Middleware` implements a hook, activate its alternate side as well. * Establish a new `Config.Plugins` that acts as a more generalized place where a hook/middleware can go. We define a plugin as this type: type Plugin interface { Hook Middleware } The reason for the first point is better backward compatibility. Notably, if I add a hook to `otelriver.Middleware`, I want it to be able to still work even if the user doesn't explicitly movie it from `Config.Middleware` to `Config.Plugins`. --- client.go | 86 ++++++++---- client_test.go | 186 +++++++++++++++++++++++++ internal/pluginconfig/plugin_config.go | 51 +++++++ plugin_defaults.go | 10 ++ plugin_defaults_test.go | 9 ++ rivertest/worker.go | 12 +- rivertype/river_type.go | 11 ++ 7 files changed, 336 insertions(+), 29 deletions(-) create mode 100644 internal/pluginconfig/plugin_config.go create mode 100644 plugin_defaults.go create mode 100644 plugin_defaults_test.go diff --git a/client.go b/client.go index 68f7e51c..7f1d5bb8 100644 --- a/client.go +++ b/client.go @@ -25,6 +25,7 @@ import ( "github.com/riverqueue/river/internal/middlewarelookup" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/notifylimiter" + "github.com/riverqueue/river/internal/pluginconfig" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/rivermiddleware" "github.com/riverqueue/river/internal/workunit" @@ -241,6 +242,9 @@ 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. Hooks []rivertype.Hook // Logger is the structured logger to use for logging purposes. If none is @@ -274,8 +278,25 @@ 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. Middleware []rivertype.Middleware + // Plugins contains extensions installed globally as both hooks and + // middleware. + // + // A plugin must implement both rivertype.Hook and rivertype.Middleware. It + // may embed PluginDefaults, or embed both HookDefaults and + // MiddlewareDefaults directly, then implement any operation-specific hook or + // middleware interfaces it needs. + // + // Hooks and Middleware are still supported. If a type in Hooks also + // implements middleware, or a type in Middleware also implements hooks, River + // will install it on both sides automatically. The Plugins list exists as a + // generic place for new extensions to be registered. + Plugins []rivertype.Plugin + // PeriodicJobs are a set of periodic jobs to run at the specified intervals // in the client. PeriodicJobs []*PeriodicJob @@ -500,6 +521,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, @@ -802,7 +824,11 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client } } - for _, hook := range config.Hooks { + configuredMiddleware := middlewareFromConfig(config) + effectiveHooks := pluginconfig.Hooks(config.Hooks, configuredMiddleware, config.Plugins) + effectiveMiddleware := pluginconfig.Middleware(config.Hooks, configuredMiddleware, config.Plugins) + + for _, hook := range effectiveHooks { if withBaseService, ok := hook.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } @@ -816,7 +842,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client config: config, driver: driver, hookLookupByJob: hooklookup.NewJobHookLookup(), - hookLookupGlobal: hooklookup.NewHookLookup(config.Hooks), + hookLookupGlobal: hooklookup.NewHookLookup(effectiveHooks), 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,31 +860,12 @@ 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. + // effectiveMiddleware contains configured middleware, hook/middleware + // hybrids, and plugins. Default middleware stays first so user middleware + // wraps inside River's internal defaults like before. { 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) - } + middleware = append(middleware, effectiveMiddleware...) for _, middleware := range middleware { if withBaseService, ok := middleware.(baseservice.WithBaseService); ok { @@ -1068,6 +1075,35 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client return client, nil } +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 +} + // Start starts the client's job fetching and working loops. Once this is called, // the client will run in a background goroutine until stopped. All jobs are // run with a context inheriting from the provided context, but with a timeout diff --git a/client_test.go b/client_test.go index d3c666fe..3ee9f0d1 100644 --- a/client_test.go +++ b/client_test.go @@ -70,6 +70,77 @@ 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{} +) + +type hookMiddlewareEmbeddedDefaultsPlugin struct { + HookDefaults + MiddlewareDefaults + + insertBeginCount int + insertManyCount int +} + +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) IsMiddleware() bool { return true } + +type hookMiddlewareValuePluginCounts struct { + insertBeginCount int + insertManyCount int +} + type periodicJobArgs struct{} func (periodicJobArgs) Kind() string { return "periodic_job" } @@ -8204,6 +8275,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() diff --git a/internal/pluginconfig/plugin_config.go b/internal/pluginconfig/plugin_config.go new file mode 100644 index 00000000..1377970e --- /dev/null +++ b/internal/pluginconfig/plugin_config.go @@ -0,0 +1,51 @@ +package pluginconfig + +import "github.com/riverqueue/river/rivertype" + +// Hooks returns the effective hook list from configured hooks, middleware, and +// plugins. Explicit hooks are preserved first, followed by middleware that also +// implement hooks, then plugins. +func Hooks(hooks []rivertype.Hook, middleware []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Hook { + effectiveHooks := make([]rivertype.Hook, 0, len(hooks)+len(middleware)+len(plugins)) + + effectiveHooks = append(effectiveHooks, hooks...) + + for _, middlewareItem := range middleware { + hook, ok := middlewareItem.(rivertype.Hook) + if !ok { + continue + } + + effectiveHooks = append(effectiveHooks, hook) + } + + for _, plugin := range plugins { + effectiveHooks = append(effectiveHooks, plugin) + } + + return effectiveHooks +} + +// Middleware returns the effective middleware list from configured hooks, +// middleware, and plugins. Explicit middleware are preserved first, followed by +// hooks that also implement middleware, then plugins. +func Middleware(hooks []rivertype.Hook, middleware []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Middleware { + effectiveMiddleware := make([]rivertype.Middleware, 0, len(hooks)+len(middleware)+len(plugins)) + + effectiveMiddleware = append(effectiveMiddleware, middleware...) + + for _, hook := range hooks { + middlewareItem, ok := hook.(rivertype.Middleware) + if !ok { + continue + } + + effectiveMiddleware = append(effectiveMiddleware, middlewareItem) + } + + for _, plugin := range plugins { + effectiveMiddleware = append(effectiveMiddleware, plugin) + } + + return effectiveMiddleware +} diff --git a/plugin_defaults.go b/plugin_defaults.go new file mode 100644 index 00000000..7c9596b7 --- /dev/null +++ b/plugin_defaults.go @@ -0,0 +1,10 @@ +package river + +// PluginDefaults should be embedded on plugin implementations. It helps +// identify a struct as both hooks and middleware, and guarantees forward +// compatibility in case additions are necessary to the rivertype.Hook or +// rivertype.Middleware interfaces. +type PluginDefaults struct { + HookDefaults + MiddlewareDefaults +} diff --git a/plugin_defaults_test.go b/plugin_defaults_test.go new file mode 100644 index 00000000..e995398f --- /dev/null +++ b/plugin_defaults_test.go @@ -0,0 +1,9 @@ +package river + +import "github.com/riverqueue/river/rivertype" + +var ( + _ rivertype.Hook = &PluginDefaults{} + _ rivertype.Middleware = &PluginDefaults{} + _ rivertype.Plugin = &PluginDefaults{} +) diff --git a/rivertest/worker.go b/rivertest/worker.go index da4f13f5..5280d442 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -13,6 +13,7 @@ import ( "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/maintenance" "github.com/riverqueue/river/internal/middlewarelookup" + "github.com/riverqueue/river/internal/pluginconfig" "github.com/riverqueue/river/internal/rivermiddleware" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" @@ -147,12 +148,15 @@ 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 { + effectiveHooks := pluginconfig.Hooks(w.config.Hooks, w.config.Middleware, w.config.Plugins) + effectiveMiddleware := pluginconfig.Middleware(w.config.Hooks, w.config.Middleware, w.config.Plugins) + + for _, hook := range effectiveHooks { if withBaseService, ok := hook.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } } - for _, middleware := range w.config.Middleware { + for _, middleware := range effectiveMiddleware { if withBaseService, ok := middleware.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } @@ -205,10 +209,10 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job return nil }, }, - HookLookupGlobal: hooklookup.NewHookLookup(w.config.Hooks), + HookLookupGlobal: hooklookup.NewHookLookup(effectiveHooks), HookLookupByJob: hooklookup.NewJobHookLookup(), JobRow: job, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(append(rivermiddleware.DefaultMiddleware(), w.config.Middleware...)), + MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(append(rivermiddleware.DefaultMiddleware(), effectiveMiddleware...)), ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) Stuck func(ctx context.Context, jobRow *rivertype.JobRow) diff --git a/rivertype/river_type.go b/rivertype/river_type.go index 13d07f63..f350e877 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 as both a hook and +// middleware. +// +// Plugin structs should embed river.PluginDefaults, or embed both +// river.HookDefaults and river.MiddlewareDefaults directly, then implement any +// operation-specific hook or middleware interfaces they need. +type Plugin interface { + Hook + Middleware +} + // JobInsertMiddleware provides an interface for middleware that integrations // can use to encapsulate common logic around job insertion. // From 000c1521166caf7c02a421d6bcdd0e079591e41d Mon Sep 17 00:00:00 2001 From: Brandur Date: Tue, 7 Jul 2026 18:35:01 -0500 Subject: [PATCH 2/6] Make `Plugin` a marker interface Make hooks and middleware opt into `Plugin` through their default embedders and function helpers. Keep `Config.Plugins` as the unified registration path. --- client.go | 136 +++++++++++++++++-------- client_test.go | 5 + example_global_hooks_test.go | 2 +- example_global_middleware_test.go | 2 +- hook_defaults_funcs.go | 22 ++-- internal/pluginconfig/plugin_config.go | 28 ++--- middleware_defaults.go | 15 ++- plugin_defaults.go | 4 +- plugin_defaults_test.go | 2 + rivertest/worker.go | 85 +++++++++++++++- rivertest/worker_test.go | 6 +- rivertype/river_type.go | 14 +-- 12 files changed, 234 insertions(+), 87 deletions(-) diff --git a/client.go b/client.go index 7f1d5bb8..614de243 100644 --- a/client.go +++ b/client.go @@ -25,7 +25,6 @@ import ( "github.com/riverqueue/river/internal/middlewarelookup" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/notifylimiter" - "github.com/riverqueue/river/internal/pluginconfig" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/rivermiddleware" "github.com/riverqueue/river/internal/workunit" @@ -245,6 +244,8 @@ type Config struct { // // 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 @@ -281,20 +282,22 @@ type Config struct { // // 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 both hooks and - // middleware. + // Plugins contains extensions installed globally as hooks, middleware, or + // both. // - // A plugin must implement both rivertype.Hook and rivertype.Middleware. It - // may embed PluginDefaults, or embed both HookDefaults and - // MiddlewareDefaults directly, then implement any operation-specific hook or - // middleware interfaces it needs. + // 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. If a type in Hooks also - // implements middleware, or a type in Middleware also implements hooks, River - // will install it on both sides automatically. The Plugins list exists as a - // generic place for new extensions to be registered. + // 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 @@ -542,6 +545,83 @@ func (c *Config) WithDefaults() *Config { } } +func effectiveHooksForConfig(config *Config) []rivertype.Hook { + configuredMiddleware := middlewareFromConfig(config) + hooks := make([]rivertype.Hook, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) + hooks = append(hooks, config.Hooks...) + + for _, middleware := range configuredMiddleware { + if hook, ok := middleware.(rivertype.Hook); ok { + hooks = append(hooks, hook) + } + } + + for _, plugin := range config.Plugins { + if plugin == nil { + continue + } + + if hook, ok := plugin.(rivertype.Hook); ok { + hooks = append(hooks, hook) + } + } + + return hooks +} + +func effectiveMiddlewareForConfig(config *Config) []rivertype.Middleware { + configuredMiddleware := middlewareFromConfig(config) + middleware := make([]rivertype.Middleware, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) + middleware = append(middleware, configuredMiddleware...) + + for _, hook := range config.Hooks { + if middlewareItem, ok := hook.(rivertype.Middleware); ok { + middleware = append(middleware, middlewareItem) + } + } + + for _, plugin := range config.Plugins { + if plugin == nil { + continue + } + + if middlewareItem, ok := plugin.(rivertype.Middleware); ok { + middleware = append(middleware, middlewareItem) + } + } + + return middleware +} + +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)") @@ -824,9 +904,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client } } - configuredMiddleware := middlewareFromConfig(config) - effectiveHooks := pluginconfig.Hooks(config.Hooks, configuredMiddleware, config.Plugins) - effectiveMiddleware := pluginconfig.Middleware(config.Hooks, configuredMiddleware, config.Plugins) + effectiveHooks := effectiveHooksForConfig(config) + effectiveMiddleware := effectiveMiddlewareForConfig(config) for _, hook := range effectiveHooks { if withBaseService, ok := hook.(baseservice.WithBaseService); ok { @@ -1075,35 +1154,6 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client return client, nil } -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 -} - // Start starts the client's job fetching and working loops. Once this is called, // the client will run in a background goroutine until stopped. All jobs are // run with a context inheriting from the provided context, but with a timeout diff --git a/client_test.go b/client_test.go index 3ee9f0d1..4f19d444 100644 --- a/client_test.go +++ b/client_test.go @@ -81,6 +81,7 @@ var ( _ rivertype.HookInsertBegin = hookMiddlewareValuePlugin{} _ rivertype.JobInsertMiddleware = hookMiddlewareValuePlugin{} + _ rivertype.Plugin = hookMiddlewareValuePlugin{} ) type hookMiddlewareEmbeddedDefaultsPlugin struct { @@ -91,6 +92,8 @@ type hookMiddlewareEmbeddedDefaultsPlugin struct { insertManyCount int } +func (p *hookMiddlewareEmbeddedDefaultsPlugin) IsPlugin() bool { return true } + func (p *hookMiddlewareEmbeddedDefaultsPlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { p.insertBeginCount++ return nil @@ -134,6 +137,8 @@ func (p hookMiddlewareValuePlugin) InsertMany(ctx context.Context, manyParams [] func (p hookMiddlewareValuePlugin) IsHook() bool { return true } +func (p hookMiddlewareValuePlugin) IsPlugin() bool { return true } + func (p hookMiddlewareValuePlugin) IsMiddleware() bool { return true } type hookMiddlewareValuePluginCounts struct { 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/internal/pluginconfig/plugin_config.go b/internal/pluginconfig/plugin_config.go index 1377970e..4c306902 100644 --- a/internal/pluginconfig/plugin_config.go +++ b/internal/pluginconfig/plugin_config.go @@ -2,11 +2,11 @@ package pluginconfig import "github.com/riverqueue/river/rivertype" -// Hooks returns the effective hook list from configured hooks, middleware, and -// plugins. Explicit hooks are preserved first, followed by middleware that also -// implement hooks, then plugins. -func Hooks(hooks []rivertype.Hook, middleware []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Hook { - effectiveHooks := make([]rivertype.Hook, 0, len(hooks)+len(middleware)+len(plugins)) +// 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...) @@ -19,18 +19,14 @@ func Hooks(hooks []rivertype.Hook, middleware []rivertype.Middleware, plugins [] effectiveHooks = append(effectiveHooks, hook) } - for _, plugin := range plugins { - effectiveHooks = append(effectiveHooks, plugin) - } - return effectiveHooks } -// Middleware returns the effective middleware list from configured hooks, -// middleware, and plugins. Explicit middleware are preserved first, followed by -// hooks that also implement middleware, then plugins. -func Middleware(hooks []rivertype.Hook, middleware []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Middleware { - effectiveMiddleware := make([]rivertype.Middleware, 0, len(hooks)+len(middleware)+len(plugins)) +// 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...) @@ -43,9 +39,5 @@ func Middleware(hooks []rivertype.Hook, middleware []rivertype.Middleware, plugi effectiveMiddleware = append(effectiveMiddleware, middlewareItem) } - for _, plugin := range plugins { - effectiveMiddleware = append(effectiveMiddleware, plugin) - } - return effectiveMiddleware } 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/plugin_defaults.go b/plugin_defaults.go index 7c9596b7..2930e479 100644 --- a/plugin_defaults.go +++ b/plugin_defaults.go @@ -1,10 +1,12 @@ package river // PluginDefaults should be embedded on plugin implementations. It helps -// identify a struct as both hooks and middleware, and guarantees forward +// 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 index e995398f..64f15598 100644 --- a/plugin_defaults_test.go +++ b/plugin_defaults_test.go @@ -6,4 +6,6 @@ var ( _ rivertype.Hook = &PluginDefaults{} _ rivertype.Middleware = &PluginDefaults{} _ rivertype.Plugin = &PluginDefaults{} + _ rivertype.Plugin = &HookDefaults{} + _ rivertype.Plugin = &MiddlewareDefaults{} ) diff --git a/rivertest/worker.go b/rivertest/worker.go index 5280d442..eb96fdc6 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -13,7 +13,6 @@ import ( "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/maintenance" "github.com/riverqueue/river/internal/middlewarelookup" - "github.com/riverqueue/river/internal/pluginconfig" "github.com/riverqueue/river/internal/rivermiddleware" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" @@ -148,8 +147,8 @@ 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) - effectiveHooks := pluginconfig.Hooks(w.config.Hooks, w.config.Middleware, w.config.Plugins) - effectiveMiddleware := pluginconfig.Middleware(w.config.Hooks, w.config.Middleware, w.config.Plugins) + effectiveHooks := effectiveHooksForConfig(w.config) + effectiveMiddleware := effectiveMiddlewareForConfig(w.config) for _, hook := range effectiveHooks { if withBaseService, ok := hook.(baseservice.WithBaseService); ok { @@ -244,6 +243,86 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job panic("unreachable") } +//nolint:staticcheck +func effectiveHooksForConfig(config *river.Config) []rivertype.Hook { + configuredMiddleware := middlewareFromConfig(config) + hooks := make([]rivertype.Hook, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) + hooks = append(hooks, config.Hooks...) + + for _, middleware := range configuredMiddleware { + if hook, ok := middleware.(rivertype.Hook); ok { + hooks = append(hooks, hook) + } + } + + for _, plugin := range config.Plugins { + if plugin == nil { + continue + } + + if hook, ok := plugin.(rivertype.Hook); ok { + hooks = append(hooks, hook) + } + } + + return hooks +} + +//nolint:staticcheck +func effectiveMiddlewareForConfig(config *river.Config) []rivertype.Middleware { + configuredMiddleware := middlewareFromConfig(config) + middleware := make([]rivertype.Middleware, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) + middleware = append(middleware, configuredMiddleware...) + + for _, hook := range config.Hooks { + if middlewareItem, ok := hook.(rivertype.Middleware); ok { + middleware = append(middleware, middlewareItem) + } + } + + for _, plugin := range config.Plugins { + if plugin == nil { + continue + } + + if middlewareItem, ok := plugin.(rivertype.Middleware); ok { + middleware = append(middleware, middlewareItem) + } + } + + return middleware +} + +//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 f350e877..d9795d97 100644 --- a/rivertype/river_type.go +++ b/rivertype/river_type.go @@ -447,15 +447,15 @@ type Middleware interface { IsMiddleware() bool } -// Plugin is a generic extension installed globally as both a hook and -// middleware. +// Plugin is a generic extension installed globally. // -// Plugin structs should embed river.PluginDefaults, or embed both -// river.HookDefaults and river.MiddlewareDefaults directly, then implement any -// operation-specific hook or middleware interfaces they need. +// 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 { - Hook - Middleware + IsPlugin() bool } // JobInsertMiddleware provides an interface for middleware that integrations From 0e837361332ec9adf415cdb8c59360f4d38cb49a Mon Sep 17 00:00:00 2001 From: Brandur Date: Wed, 8 Jul 2026 12:28:33 -0500 Subject: [PATCH 3/6] Consolidate hook and middleware lookup into `pluginlookup` Merge the split lookup packages into one precomputed plugin lookup. Build the kind buckets up front and use the unified path throughout client startup and job execution. --- CHANGELOG.md | 5 + client.go | 138 ++----- client_test.go | 14 +- hook_defaults_funcs_test.go | 7 + internal/hooklookup/hook_lookup.go | 166 --------- internal/hooklookup/hook_lookup_test.go | 266 ------------- internal/jobexecutor/job_executor.go | 23 +- internal/jobexecutor/job_executor_test.go | 27 +- internal/maintenance/job_rescuer_test.go | 4 +- internal/maintenance/periodic_job_enqueuer.go | 14 +- .../maintenance/periodic_job_enqueuer_test.go | 11 +- .../middlewarelookup/middleware_lookup.go | 108 ------ .../middleware_lookup_test.go | 141 ------- internal/pluginlookup/plugin_lookup.go | 190 ++++++++++ internal/pluginlookup/plugin_lookup_test.go | 351 ++++++++++++++++++ internal/rivermiddleware/middleware.go | 1 + internal/workunit/work_unit.go | 6 +- middleware_defaults_test.go | 8 + producer.go | 23 +- producer_test.go | 18 +- riverlog/river_log.go | 2 + rivertest/work_unit_wrapper.go | 4 +- rivertest/worker.go | 10 +- work_unit_wrapper.go | 4 +- 24 files changed, 678 insertions(+), 863 deletions(-) delete mode 100644 internal/hooklookup/hook_lookup.go delete mode 100644 internal/hooklookup/hook_lookup_test.go delete mode 100644 internal/middlewarelookup/middleware_lookup.go delete mode 100644 internal/middlewarelookup/middleware_lookup_test.go create mode 100644 internal/pluginlookup/plugin_lookup.go create mode 100644 internal/pluginlookup/plugin_lookup_test.go 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 614de243..a7d849d6 100644 --- a/client.go +++ b/client.go @@ -17,14 +17,13 @@ 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/workunit" @@ -545,54 +544,6 @@ func (c *Config) WithDefaults() *Config { } } -func effectiveHooksForConfig(config *Config) []rivertype.Hook { - configuredMiddleware := middlewareFromConfig(config) - hooks := make([]rivertype.Hook, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) - hooks = append(hooks, config.Hooks...) - - for _, middleware := range configuredMiddleware { - if hook, ok := middleware.(rivertype.Hook); ok { - hooks = append(hooks, hook) - } - } - - for _, plugin := range config.Plugins { - if plugin == nil { - continue - } - - if hook, ok := plugin.(rivertype.Hook); ok { - hooks = append(hooks, hook) - } - } - - return hooks -} - -func effectiveMiddlewareForConfig(config *Config) []rivertype.Middleware { - configuredMiddleware := middlewareFromConfig(config) - middleware := make([]rivertype.Middleware, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) - middleware = append(middleware, configuredMiddleware...) - - for _, hook := range config.Hooks { - if middlewareItem, ok := hook.(rivertype.Middleware); ok { - middleware = append(middleware, middlewareItem) - } - } - - for _, plugin := range config.Plugins { - if plugin == nil { - continue - } - - if middlewareItem, ok := plugin.(rivertype.Middleware); ok { - middleware = append(middleware, middlewareItem) - } - } - - return middleware -} - func middlewareFromConfig(config *Config) []rivertype.Middleware { middleware := make([]rivertype.Middleware, 0, len(config.Middleware)+len(config.JobInsertMiddleware)+len(config.WorkerMiddleware)) @@ -780,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. @@ -904,11 +854,14 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client } } - effectiveHooks := effectiveHooksForConfig(config) - effectiveMiddleware := effectiveMiddlewareForConfig(config) + var ( + configuredMiddleware = middlewareFromConfig(config) + allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...) + allPlugins = pluginlookup.NormalizePlugins(config.Hooks, allMiddleware, config.Plugins) + ) - for _, hook := range effectiveHooks { - if withBaseService, ok := hook.(baseservice.WithBaseService); ok { + for _, plugin := range allPlugins { + if withBaseService, ok := plugin.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } } @@ -920,8 +873,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client }, config: config, driver: driver, - hookLookupByJob: hooklookup.NewJobHookLookup(), - hookLookupGlobal: hooklookup.NewHookLookup(effectiveHooks), + 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 @@ -939,22 +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) - // effectiveMiddleware contains configured middleware, hook/middleware - // hybrids, and plugins. Default middleware stays first so user middleware - // wraps inside River's internal defaults like before. - { - middleware := rivermiddleware.DefaultMiddleware() - middleware = append(middleware, effectiveMiddleware...) - - 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) @@ -1082,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, @@ -2062,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 @@ -2094,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. @@ -2372,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 4f19d444..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" @@ -8579,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) }, }, { @@ -8590,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) }, }, { @@ -8602,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/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/pluginlookup/plugin_lookup.go b/internal/pluginlookup/plugin_lookup.go new file mode 100644 index 00000000..c8f2e39a --- /dev/null +++ b/internal/pluginlookup/plugin_lookup.go @@ -0,0 +1,190 @@ +package pluginlookup + +import ( + "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. Hook and middleware registrations must already satisfy +// rivertype.Plugin to participate. +func NormalizePlugins(hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Plugin { + normalizedPlugins := make([]rivertype.Plugin, 0, len(hooks)+len(middlewares)+len(plugins)) + + for _, hook := range hooks { + if plugin, ok := hook.(rivertype.Plugin); ok { + normalizedPlugins = append(normalizedPlugins, plugin) + } + } + for _, middleware := range middlewares { + if plugin, ok := middleware.(rivertype.Plugin); ok { + normalizedPlugins = append(normalizedPlugins, plugin) + } + } + for _, plugin := range plugins { + if plugin != nil { + normalizedPlugins = append(normalizedPlugins, plugin) + } + } + + 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 +} + +// +// 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..d0d8d34e --- /dev/null +++ b/internal/pluginlookup/plugin_lookup_test.go @@ -0,0 +1,351 @@ +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 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() + }) +} + +// +// 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 +} + +// +// 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) +} diff --git a/internal/rivermiddleware/middleware.go b/internal/rivermiddleware/middleware.go index 81a6caa2..aa1a32f2 100644 --- a/internal/rivermiddleware/middleware.go +++ b/internal/rivermiddleware/middleware.go @@ -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_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/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/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 eb96fdc6..c1f26d48 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -8,11 +8,10 @@ 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/pluginlookup" "github.com/riverqueue/river/internal/rivermiddleware" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" @@ -208,10 +207,9 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job return nil }, }, - HookLookupGlobal: hooklookup.NewHookLookup(effectiveHooks), - HookLookupByJob: hooklookup.NewJobHookLookup(), - JobRow: job, - MiddlewareLookupGlobal: middlewarelookup.NewMiddlewareLookup(append(rivermiddleware.DefaultMiddleware(), effectiveMiddleware...)), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupGlobal: pluginlookup.NewPluginLookup(pluginlookup.NormalizePlugins(effectiveHooks, append(rivermiddleware.DefaultMiddleware(), effectiveMiddleware...), nil)), + JobRow: job, ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) Stuck func(ctx context.Context, jobRow *rivertype.JobRow) 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) } From 50adaf0e77548fccbe040a2780dd60133d1a4473 Mon Sep 17 00:00:00 2001 From: Brandur Date: Wed, 8 Jul 2026 17:52:33 -0500 Subject: [PATCH 4/6] Bring `rivertest` more in line with unified hook/middleware path from client --- rivertest/worker.go | 70 ++++++--------------------------------------- 1 file changed, 9 insertions(+), 61 deletions(-) diff --git a/rivertest/worker.go b/rivertest/worker.go index c1f26d48..19e534e4 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -146,16 +146,14 @@ 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) - effectiveHooks := effectiveHooksForConfig(w.config) - effectiveMiddleware := effectiveMiddlewareForConfig(w.config) - - for _, hook := range effectiveHooks { - if withBaseService, ok := hook.(baseservice.WithBaseService); ok { - baseservice.Init(archetype, withBaseService) - } - } - for _, middleware := range effectiveMiddleware { - if withBaseService, ok := middleware.(baseservice.WithBaseService); ok { + var ( + configuredMiddleware = middlewareFromConfig(w.config) + allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...) + allPlugins = pluginlookup.NormalizePlugins(w.config.Hooks, allMiddleware, w.config.Plugins) //nolint:staticcheck // Bridge legacy Hooks into the test worker's unified plugin path. + ) + + for _, plugin := range allPlugins { + if withBaseService, ok := plugin.(baseservice.WithBaseService); ok { baseservice.Init(archetype, withBaseService) } } @@ -208,7 +206,7 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job }, }, PluginLookupByJob: pluginlookup.NewJobPluginLookup(), - PluginLookupGlobal: pluginlookup.NewPluginLookup(pluginlookup.NormalizePlugins(effectiveHooks, append(rivermiddleware.DefaultMiddleware(), effectiveMiddleware...), nil)), + PluginLookupGlobal: pluginlookup.NewPluginLookup(allPlugins), JobRow: job, ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) @@ -241,56 +239,6 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job panic("unreachable") } -//nolint:staticcheck -func effectiveHooksForConfig(config *river.Config) []rivertype.Hook { - configuredMiddleware := middlewareFromConfig(config) - hooks := make([]rivertype.Hook, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) - hooks = append(hooks, config.Hooks...) - - for _, middleware := range configuredMiddleware { - if hook, ok := middleware.(rivertype.Hook); ok { - hooks = append(hooks, hook) - } - } - - for _, plugin := range config.Plugins { - if plugin == nil { - continue - } - - if hook, ok := plugin.(rivertype.Hook); ok { - hooks = append(hooks, hook) - } - } - - return hooks -} - -//nolint:staticcheck -func effectiveMiddlewareForConfig(config *river.Config) []rivertype.Middleware { - configuredMiddleware := middlewareFromConfig(config) - middleware := make([]rivertype.Middleware, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins)) - middleware = append(middleware, configuredMiddleware...) - - for _, hook := range config.Hooks { - if middlewareItem, ok := hook.(rivertype.Middleware); ok { - middleware = append(middleware, middlewareItem) - } - } - - for _, plugin := range config.Plugins { - if plugin == nil { - continue - } - - if middlewareItem, ok := plugin.(rivertype.Middleware); ok { - middleware = append(middleware, middlewareItem) - } - } - - return middleware -} - //nolint:staticcheck func middlewareFromConfig(config *river.Config) []rivertype.Middleware { middleware := make([]rivertype.Middleware, 0, From 4c5579ec498f5175f850832c65f5838f15713201 Mon Sep 17 00:00:00 2001 From: Brandur Date: Thu, 9 Jul 2026 11:14:20 -0500 Subject: [PATCH 5/6] Code review feedback to allow legacy hooks/middlewares to work unchanged --- internal/pluginlookup/plugin_lookup.go | 100 ++++++++++++++++++-- internal/pluginlookup/plugin_lookup_test.go | 67 +++++++++++++ 2 files changed, 160 insertions(+), 7 deletions(-) diff --git a/internal/pluginlookup/plugin_lookup.go b/internal/pluginlookup/plugin_lookup.go index c8f2e39a..e9523b16 100644 --- a/internal/pluginlookup/plugin_lookup.go +++ b/internal/pluginlookup/plugin_lookup.go @@ -1,6 +1,7 @@ package pluginlookup import ( + "context" "sync" "github.com/riverqueue/river/rivertype" @@ -86,26 +87,28 @@ func NewPluginLookup(plugins []rivertype.Plugin) PluginLookupInterface { } // NormalizePlugins converts hook, middleware, and plugin registrations into a -// single plugin slice. Hook and middleware registrations must already satisfy -// rivertype.Plugin to participate. +// single plugin slice while preserving legacy hook and middleware registrations +// that don't yet opt into Plugin. func NormalizePlugins(hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Plugin { normalizedPlugins := make([]rivertype.Plugin, 0, len(hooks)+len(middlewares)+len(plugins)) for _, hook := range hooks { if plugin, ok := hook.(rivertype.Plugin); ok { normalizedPlugins = append(normalizedPlugins, plugin) + continue } + + normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook)) } for _, middleware := range middlewares { if plugin, ok := middleware.(rivertype.Plugin); ok { normalizedPlugins = append(normalizedPlugins, plugin) + continue } + + normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware)) } - for _, plugin := range plugins { - if plugin != nil { - normalizedPlugins = append(normalizedPlugins, plugin) - } - } + normalizedPlugins = append(normalizedPlugins, plugins...) return normalizedPlugins } @@ -141,6 +144,89 @@ 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 // diff --git a/internal/pluginlookup/plugin_lookup_test.go b/internal/pluginlookup/plugin_lookup_test.go index d0d8d34e..890f3344 100644 --- a/internal/pluginlookup/plugin_lookup_test.go +++ b/internal/pluginlookup/plugin_lookup_test.go @@ -204,6 +204,40 @@ func TestPluginLookup(t *testing.T) { 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) + }) } // @@ -349,3 +383,36 @@ 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) +} From 9873cfd54acff07d29af2c4934d69bac99d79e99 Mon Sep 17 00:00:00 2001 From: Brandur Date: Fri, 10 Jul 2026 13:23:35 -0500 Subject: [PATCH 6/6] Cleaner use of internal plugins + plugin initialization --- client.go | 10 ++-- internal/pluginlookup/plugin_lookup.go | 18 ++++--- internal/pluginlookup/plugin_lookup_test.go | 53 +++++++++++++++++++ .../middleware.go => riverplugin/plugin.go} | 10 ++-- resumable.go | 12 ++--- resumable_step_tx_test.go | 4 +- resumable_test.go | 26 ++++----- rivertest/worker.go | 13 +++-- 8 files changed, 103 insertions(+), 43 deletions(-) rename internal/{rivermiddleware/middleware.go => riverplugin/plugin.go} (91%) diff --git a/client.go b/client.go index a7d849d6..1dab4695 100644 --- a/client.go +++ b/client.go @@ -25,7 +25,7 @@ import ( "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" @@ -854,10 +854,10 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client } } - var ( - configuredMiddleware = middlewareFromConfig(config) - allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...) - allPlugins = pluginlookup.NormalizePlugins(config.Hooks, allMiddleware, config.Plugins) + allPlugins := pluginlookup.NormalizePlugins( + config.Hooks, + middlewareFromConfig(config), + append(riverplugin.DefaultPlugins(), config.Plugins...), ) for _, plugin := range allPlugins { diff --git a/internal/pluginlookup/plugin_lookup.go b/internal/pluginlookup/plugin_lookup.go index e9523b16..de0c6111 100644 --- a/internal/pluginlookup/plugin_lookup.go +++ b/internal/pluginlookup/plugin_lookup.go @@ -89,26 +89,30 @@ func NewPluginLookup(plugins []rivertype.Plugin) PluginLookupInterface { // 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) - continue + } else { + normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook)) } - - normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook)) } + for _, middleware := range middlewares { if plugin, ok := middleware.(rivertype.Plugin); ok { normalizedPlugins = append(normalizedPlugins, plugin) - continue + } else { + normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware)) } - - normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware)) } - normalizedPlugins = append(normalizedPlugins, plugins...) return normalizedPlugins } diff --git a/internal/pluginlookup/plugin_lookup_test.go b/internal/pluginlookup/plugin_lookup_test.go index 890f3344..a5d0c665 100644 --- a/internal/pluginlookup/plugin_lookup_test.go +++ b/internal/pluginlookup/plugin_lookup_test.go @@ -110,6 +110,36 @@ func TestJobPluginLookup(t *testing.T) { }) } +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() @@ -307,6 +337,29 @@ func (t *testHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype 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 // diff --git a/internal/rivermiddleware/middleware.go b/internal/riverplugin/plugin.go similarity index 91% rename from internal/rivermiddleware/middleware.go rename to internal/riverplugin/plugin.go index aa1a32f2..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 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/rivertest/worker.go b/rivertest/worker.go index 19e534e4..9377eb5f 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -12,7 +12,7 @@ import ( "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/maintenance" "github.com/riverqueue/river/internal/pluginlookup" - "github.com/riverqueue/river/internal/rivermiddleware" + "github.com/riverqueue/river/internal/riverplugin" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/riversharedtest" @@ -146,10 +146,13 @@ 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) - var ( - configuredMiddleware = middlewareFromConfig(w.config) - allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...) - allPlugins = pluginlookup.NormalizePlugins(w.config.Hooks, allMiddleware, w.config.Plugins) //nolint:staticcheck // Bridge legacy Hooks into the test worker's unified plugin path. + 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 {