diff --git a/CHANGELOG.md b/CHANGELOG.md index 4546a03d..40dbc792 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added `JobArgsWithPlugins` for installing plugins that apply only to a specific job type. [PR #1337](https://github.com/riverqueue/river/pull/1337). + ## [0.41.1] - 2026-07-29 ### Fixed diff --git a/client.go b/client.go index 34be5948..6cee0af1 100644 --- a/client.go +++ b/client.go @@ -241,7 +241,8 @@ type Config struct { // and the work hook between them will not run. When a job is worked, the // work hook runs and the insertion hooks on either side of it are skipped. // - // Jobs may have their own specific hooks by implementing JobArgsWithHooks. + // Jobs may have their own specific hooks by implementing JobArgsWithHooks or + // JobArgsWithPlugins. // // Entries in Hooks are installed only as hooks, even if they also implement // rivertype.Middleware. Use Plugins for an extension that should act as @@ -297,6 +298,8 @@ type Config struct { // // Use Hooks or Middleware when an extension should be installed only as the // corresponding kind. Use Plugins when it should be eligible as both. + // Jobs may have their own specific plugins by implementing + // JobArgsWithPlugins. Plugins []rivertype.Plugin // PeriodicJobs are a set of periodic jobs to run at the specified intervals @@ -708,7 +711,7 @@ type Client[TTx any] struct { driver riverdriver.Driver[TTx] elector *leadership.Elector pluginLookupByJob *pluginlookup.JobPluginLookup - pluginLookupGlobal pluginlookup.PluginLookupInterface + pluginLookupGlobal *pluginlookup.PluginLookup insertNotifyLimiter *notifylimiter.Limiter notifier *notifier.Notifier // may be nil in poll-only mode periodicJobs *PeriodicJobBundle @@ -832,9 +835,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client middleware = pluginconfig.CombinedMiddleware(config.Middleware, config.JobInsertMiddleware, config.WorkerMiddleware) plugins = append(riverplugin.DefaultPlugins(), config.Plugins...) ) - pluginlookup.InitBaseServices(archetype, config.Hooks) - pluginlookup.InitBaseServices(archetype, middleware) - pluginlookup.InitBaseServices(archetype, plugins) + pluginLookupByJob := pluginlookup.NewJobPluginLookup(archetype) + pluginLookupGlobal := pluginlookup.NewPluginLookupFromConfig(archetype, config.Hooks, middleware, plugins) client := &Client[TTx]{ clientNotifyBundle: &ClientNotifyBundle[TTx]{ @@ -843,8 +845,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client }, config: config, driver: driver, - pluginLookupByJob: pluginlookup.NewJobPluginLookup(), - pluginLookupGlobal: pluginlookup.NewPluginLookupFromConfig(config.Hooks, middleware, plugins), + pluginLookupByJob: pluginLookupByJob, + pluginLookupGlobal: pluginLookupGlobal, producersByQueueName: make(map[string]*producer), testSignals: clientTestSignals{}, workCancel: func(cause error) {}, // replaced on start, but here in case StopAndCancel is called before start up @@ -872,13 +874,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client if config.Workers != nil { workerMetadata = make([]*rivertype.WorkerMetadata, 0, len(config.Workers.workersMap)) for kind, workerInfo := range config.Workers.workersMap { - var hooks []rivertype.Hook - if jobArgsWithHooks, ok := workerInfo.jobArgs.(JobArgsWithHooks); ok { - hooks = jobArgsWithHooks.Hooks() - } - workerMetadata = append(workerMetadata, &rivertype.WorkerMetadata{ - JobArgHooks: hooks, + JobArgHooks: pluginLookupByJob.ByJobArgs(workerInfo.jobArgs).Hooks(), Kind: kind, }) } @@ -2004,7 +2001,20 @@ func (c *Client[TTx]) insertManyShared( return insertResults, nil } - jobInsertMiddleware := c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert) + jobInsertMiddleware := append([]any(nil), c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert)...) + jobKindsSeen := make(map[string]struct{}, len(insertParams)) + for _, params := range insertParams { + kind := params.Args.Kind() + if _, ok := jobKindsSeen[kind]; ok { + continue + } + jobKindsSeen[kind] = struct{}{} + + jobInsertMiddleware = append( + jobInsertMiddleware, + c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.PluginKindMiddlewareJobInsert)..., + ) + } 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. diff --git a/internal/jobexecutor/job_executor.go b/internal/jobexecutor/job_executor.go index 3bad60ab..c6680014 100644 --- a/internal/jobexecutor/job_executor.go +++ b/internal/jobexecutor/job_executor.go @@ -112,7 +112,7 @@ type JobExecutor struct { DefaultClientRetryPolicy ClientRetryPolicy ErrorHandler ErrorHandler PluginLookupByJob *pluginlookup.JobPluginLookup - PluginLookupGlobal pluginlookup.PluginLookupInterface + PluginLookupGlobal *pluginlookup.PluginLookup JobRow *rivertype.JobRow ProducerCallbacks struct { JobDone func(jobRow *rivertype.JobRow) @@ -217,12 +217,13 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { ) return &jobExecutorResult{Err: &rivertype.UnknownJobKindError{Kind: e.JobRow.Kind}, MetadataUpdates: metadataUpdates} } + pluginLookupByJob := e.WorkUnit.PluginLookup(e.PluginLookupByJob) doInner := execution.Func(func(ctx context.Context) error { { for _, hook := range append( e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkBegin), - e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.PluginKindHookWorkBegin)..., + pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkBegin)..., ) { if err := hook.(rivertype.HookWorkBegin).WorkBegin(ctx, e.JobRow); err != nil { //nolint:forcetypeassert return err @@ -251,7 +252,7 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { { for _, hook := range append( e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkEnd), - e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.PluginKindHookWorkEnd)..., + pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkEnd)..., ) { err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, e.JobRow, err) //nolint:forcetypeassert } @@ -260,13 +261,19 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) { return err }) - globalMiddleware := make([]rivertype.Middleware, 0, len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker))) + pluginMiddleware := make([]rivertype.Middleware, 0, + len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker))+ + len(pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker)), + ) for _, plugin := range e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker) { - globalMiddleware = append(globalMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert + pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert + } + for _, plugin := range pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker) { + pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert } executeFunc := execution.MiddlewareChain( - globalMiddleware, + pluginMiddleware, e.WorkUnit.Middleware(), doInner, e.JobRow, diff --git a/internal/jobexecutor/job_executor_test.go b/internal/jobexecutor/job_executor_test.go index f595481b..4890bc42 100644 --- a/internal/jobexecutor/job_executor_test.go +++ b/internal/jobexecutor/job_executor_test.go @@ -38,7 +38,7 @@ type customizableWorkUnit struct { work func() error } -func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { +func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup { return pluginlookup.NewPluginLookup(nil) } @@ -189,7 +189,7 @@ func TestJobExecutor_Execute(t *testing.T) { Completer: bundle.completer, DefaultClientRetryPolicy: &retrypolicytest.RetryPolicyNoJitter{}, ErrorHandler: bundle.errorHandler, - PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobRow: bundle.jobRow, ProducerCallbacks: struct { diff --git a/internal/maintenance/job_rescuer_test.go b/internal/maintenance/job_rescuer_test.go index 66d324d0..d7963c12 100644 --- a/internal/maintenance/job_rescuer_test.go +++ b/internal/maintenance/job_rescuer_test.go @@ -50,7 +50,7 @@ type callbackWorkUnit struct { unmarshalErr error } -func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { +func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup { 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 7ef38ae3..6f5a5c90 100644 --- a/internal/maintenance/periodic_job_enqueuer.go +++ b/internal/maintenance/periodic_job_enqueuer.go @@ -89,7 +89,7 @@ type InsertFunc func(ctx context.Context, tx riverdriver.ExecutorTx, insertParam type PeriodicJobEnqueuerConfig struct { AdvisoryLockPrefix int32 - PluginLookupGlobal pluginlookup.PluginLookupInterface + PluginLookupGlobal *pluginlookup.PluginLookup // Insert is the function to call to insert jobs into the database. Insert InsertFunc diff --git a/internal/pluginlookup/plugin_lookup.go b/internal/pluginlookup/plugin_lookup.go index 9829916c..5ce27534 100644 --- a/internal/pluginlookup/plugin_lookup.go +++ b/internal/pluginlookup/plugin_lookup.go @@ -23,37 +23,24 @@ const ( PluginKindMiddlewareWorker PluginKind = "middleware_worker" ) -// InitBaseServices initializes base services embedded in plugins. -func InitBaseServices[T any](archetype *baseservice.Archetype, plugins []T) { - for _, plugin := range plugins { - if withBaseService, ok := any(plugin).(baseservice.WithBaseService); ok { - baseservice.Init(archetype, withBaseService) - } - } -} - // -// PluginLookupInterface +// PluginLookup // -// 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) []any +// PluginLookup looks up plugins by kind. Its zero value is an empty lookup. +type PluginLookup struct { + hooks []rivertype.Hook + pluginsByKind map[PluginKind][]any } -// 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. Each input is considered for every hook and middleware -// kind it implements. +// NewPluginLookup returns a new plugin lookup based on the given plugins. Each +// input is considered for every hook and middleware kind it implements. // // The plugins parameter is []any rather than []rivertype.Plugin because the // lookup may contain legacy hooks and middleware that don't implement // rivertype.Plugin. Keeping their original concrete values avoids compatibility // wrappers that would need to forward every operation-specific interface. -func NewPluginLookup(plugins []any) PluginLookupInterface { +func NewPluginLookup(plugins []any) *PluginLookup { return newPluginLookup(plugins, plugins) } @@ -61,7 +48,15 @@ func NewPluginLookup(plugins []any) PluginLookupInterface { // hooks, middleware, and plugins. Explicit plugins may participate as either // hooks or middleware, while entries from the legacy Hooks and Middleware // configuration fields participate only as the kind they were configured as. -func NewPluginLookupFromConfig(hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) PluginLookupInterface { +// Base services embedded in any configured extension are initialized with +// archetype when it's non-nil. +func NewPluginLookupFromConfig(archetype *baseservice.Archetype, hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) *PluginLookup { + if archetype != nil { + initBaseServices(archetype, hooks) + initBaseServices(archetype, middlewares) + initBaseServices(archetype, plugins) + } + pluginValues := toAnySlice(plugins) hookValues := make([]any, 0, len(plugins)+len(hooks)) @@ -75,32 +70,53 @@ func NewPluginLookupFromConfig(hooks []rivertype.Hook, middlewares []rivertype.M return newPluginLookup(hookValues, middlewareValues) } -func newPluginLookup(hooks, middlewares []any) PluginLookupInterface { +func (c *PluginLookup) ByKind(kind PluginKind) []any { + return c.pluginsByKind[kind] +} + +// Hooks returns all the hooks in the lookup in configuration order. +func (c *PluginLookup) Hooks() []rivertype.Hook { + return c.hooks +} + +func initBaseServices[T any](archetype *baseservice.Archetype, plugins []T) { + for _, plugin := range plugins { + if withBaseService, ok := any(plugin).(baseservice.WithBaseService); ok { + baseservice.Init(archetype, withBaseService) + } + } +} + +func newPluginLookup(hooks, middlewares []any) *PluginLookup { + lookup := &PluginLookup{} if len(hooks) < 1 && len(middlewares) < 1 { - return &emptyPluginLookup{} + return lookup } - pluginsByKind := make(map[PluginKind][]any) + lookup.pluginsByKind = make(map[PluginKind][]any) for _, plugin := range hooks { if plugin == nil { continue } + if hook, ok := plugin.(rivertype.Hook); ok { + lookup.hooks = append(lookup.hooks, hook) + } if _, ok := plugin.(rivertype.HookInsertBegin); ok { - pluginsByKind[PluginKindHookInsertBegin] = append(pluginsByKind[PluginKindHookInsertBegin], plugin) + lookup.pluginsByKind[PluginKindHookInsertBegin] = append(lookup.pluginsByKind[PluginKindHookInsertBegin], plugin) } if _, ok := plugin.(rivertype.HookMetricEmit); ok { - pluginsByKind[PluginKindHookMetricEmit] = append(pluginsByKind[PluginKindHookMetricEmit], plugin) + lookup.pluginsByKind[PluginKindHookMetricEmit] = append(lookup.pluginsByKind[PluginKindHookMetricEmit], plugin) } if _, ok := plugin.(rivertype.HookPeriodicJobsStart); ok { - pluginsByKind[PluginKindHookPeriodicJobsStart] = append(pluginsByKind[PluginKindHookPeriodicJobsStart], plugin) + lookup.pluginsByKind[PluginKindHookPeriodicJobsStart] = append(lookup.pluginsByKind[PluginKindHookPeriodicJobsStart], plugin) } if _, ok := plugin.(rivertype.HookWorkBegin); ok { - pluginsByKind[PluginKindHookWorkBegin] = append(pluginsByKind[PluginKindHookWorkBegin], plugin) + lookup.pluginsByKind[PluginKindHookWorkBegin] = append(lookup.pluginsByKind[PluginKindHookWorkBegin], plugin) } if _, ok := plugin.(rivertype.HookWorkEnd); ok { - pluginsByKind[PluginKindHookWorkEnd] = append(pluginsByKind[PluginKindHookWorkEnd], plugin) + lookup.pluginsByKind[PluginKindHookWorkEnd] = append(lookup.pluginsByKind[PluginKindHookWorkEnd], plugin) } } @@ -110,45 +126,14 @@ func newPluginLookup(hooks, middlewares []any) PluginLookupInterface { } if _, ok := plugin.(rivertype.JobInsertMiddleware); ok { - pluginsByKind[PluginKindMiddlewareJobInsert] = append(pluginsByKind[PluginKindMiddlewareJobInsert], plugin) + lookup.pluginsByKind[PluginKindMiddlewareJobInsert] = append(lookup.pluginsByKind[PluginKindMiddlewareJobInsert], plugin) } if _, ok := plugin.(rivertype.WorkerMiddleware); ok { - pluginsByKind[PluginKindMiddlewareWorker] = append(pluginsByKind[PluginKindMiddlewareWorker], plugin) + lookup.pluginsByKind[PluginKindMiddlewareWorker] = append(lookup.pluginsByKind[PluginKindMiddlewareWorker], plugin) } } - return &pluginLookup{pluginsByKind: pluginsByKind} -} - -// -// 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][]any -} - -func (c *pluginLookup) ByKind(kind PluginKind) []any { - 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) []any { - return nil + return lookup } func toAnySlice[T any](values []T) []any { @@ -164,41 +149,50 @@ func toAnySlice[T any](values []T) []any { // type JobPluginLookup struct { - pluginLookupByKind map[string]PluginLookupInterface + archetype *baseservice.Archetype + mu sync.RWMutex + pluginLookupByKind map[string]*PluginLookup } -func NewJobPluginLookup() *JobPluginLookup { +func NewJobPluginLookup(archetype *baseservice.Archetype) *JobPluginLookup { return &JobPluginLookup{ - pluginLookupByKind: make(map[string]PluginLookupInterface), + archetype: archetype, + pluginLookupByKind: make(map[string]*PluginLookup), } } -// 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 { +// ByJobArgs returns a plugin lookup for the given job args. +func (c *JobPluginLookup) ByJobArgs(args rivertype.JobArgs) *PluginLookup { kind := args.Kind() c.mu.RLock() - lookup, ok := c.pluginLookupByKind[kind] + entry, ok := c.pluginLookupByKind[kind] c.mu.RUnlock() if ok { - return lookup + return entry } c.mu.Lock() defer c.mu.Unlock() + if entry, ok := c.pluginLookupByKind[kind]; ok { + return entry + } - var hooks []rivertype.Hook + var ( + hooks []rivertype.Hook + plugins []rivertype.Plugin + ) if argsWithHooks, ok := args.(jobArgsWithHooks); ok { hooks = argsWithHooks.Hooks() } + if argsWithPlugins, ok := args.(jobArgsWithPlugins); ok { + plugins = argsWithPlugins.Plugins() + } - lookup = newPluginLookup(toAnySlice(hooks), nil) - c.pluginLookupByKind[kind] = lookup - return lookup + entry = NewPluginLookupFromConfig(c.archetype, hooks, nil, plugins) + c.pluginLookupByKind[kind] = entry + return entry } // Same as river.JobArgsWithHooks, but duplicated here so that can still live in @@ -206,3 +200,9 @@ func (c *JobPluginLookup) ByJobArgs(args rivertype.JobArgs) PluginLookupInterfac type jobArgsWithHooks interface { Hooks() []rivertype.Hook } + +// Same as river.JobArgsWithPlugins, but duplicated here so that can still live +// in the top level package. +type jobArgsWithPlugins interface { + Plugins() []rivertype.Plugin +} diff --git a/internal/pluginlookup/plugin_lookup_test.go b/internal/pluginlookup/plugin_lookup_test.go index 1749927c..3d0c983d 100644 --- a/internal/pluginlookup/plugin_lookup_test.go +++ b/internal/pluginlookup/plugin_lookup_test.go @@ -13,28 +13,13 @@ import ( 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 := NewPluginLookup(nil) - pluginLookup, _ := setup(t) - - require.Nil(t, pluginLookup.ByKind(PluginKindHookInsertBegin)) - require.Nil(t, pluginLookup.ByKind(PluginKindHookMetricEmit)) - require.Nil(t, pluginLookup.ByKind(PluginKindHookWorkBegin)) - require.Nil(t, pluginLookup.ByKind(PluginKindMiddlewareJobInsert)) - require.Nil(t, pluginLookup.ByKind(PluginKindMiddlewareWorker)) - }) + require.Nil(t, pluginLookup.ByKind(PluginKindHookInsertBegin)) + require.Nil(t, pluginLookup.ByKind(PluginKindHookMetricEmit)) + require.Nil(t, pluginLookup.ByKind(PluginKindHookWorkBegin)) + require.Nil(t, pluginLookup.ByKind(PluginKindMiddlewareJobInsert)) + require.Nil(t, pluginLookup.ByKind(PluginKindMiddlewareWorker)) } func TestJobPluginLookup(t *testing.T) { @@ -45,7 +30,7 @@ func TestJobPluginLookup(t *testing.T) { setup := func(t *testing.T) (*JobPluginLookup, *testBundle) { //nolint:unparam t.Helper() - return NewJobPluginLookup(), &testBundle{} + return NewJobPluginLookup(nil), &testBundle{} } t.Run("LooksUpHooks", func(t *testing.T) { @@ -87,6 +72,35 @@ func TestJobPluginLookup(t *testing.T) { }, jobPluginLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByKind(PluginKindHookWorkEnd)) }) + t.Run("LooksUpPlugins", func(t *testing.T) { + t.Parallel() + + jobPluginLookup, _ := setup(t) + args := &jobArgsWithCustomHooksAndPlugins{} + + require.Equal(t, []any{ + &testHookMiddlewarePlugin{}, + &testHookInsertBegin{}, + }, jobPluginLookup.ByJobArgs(args).ByKind(PluginKindHookInsertBegin)) + require.Equal(t, []any{ + &testHookWorkEnd{}, + }, jobPluginLookup.ByJobArgs(args).ByKind(PluginKindHookWorkEnd)) + require.Equal(t, []any{ + &testHookMiddlewarePlugin{}, + &testMiddlewareJobInsertAndWorker{}, + }, jobPluginLookup.ByJobArgs(args).ByKind(PluginKindMiddlewareJobInsert)) + require.Equal(t, []any{ + &testMiddlewareJobInsertAndWorker{}, + }, jobPluginLookup.ByJobArgs(args).ByKind(PluginKindMiddlewareWorker)) + require.Equal(t, []rivertype.Hook{ + &testHookMiddlewarePlugin{}, + &testHookWorkEnd{}, + &testHookInsertBegin{}, + }, jobPluginLookup.ByJobArgs(args).Hooks()) + + require.Len(t, jobPluginLookup.pluginLookupByKind, 1) + }) + t.Run("Stress", func(t *testing.T) { t.Parallel() @@ -118,12 +132,12 @@ func TestNewPluginLookupFromConfig(t *testing.T) { middlewarePlugin := &testHookMiddlewarePlugin{} plugin := &testHookMiddlewarePlugin{} - lookup, isPluginLookup := NewPluginLookupFromConfig( + lookup := NewPluginLookupFromConfig( + nil, []rivertype.Hook{hookPlugin}, []rivertype.Middleware{middlewarePlugin}, []rivertype.Plugin{plugin}, - ).(*pluginLookup) - require.True(t, isPluginLookup) + ) require.Equal(t, []any{ plugin, @@ -140,10 +154,10 @@ func TestPluginLookup(t *testing.T) { type testBundle struct{} - setup := func(t *testing.T) (*pluginLookup, *testBundle) { //nolint:unparam + setup := func(t *testing.T) (*PluginLookup, *testBundle) { //nolint:unparam t.Helper() - lookup, isPluginLookup := NewPluginLookup([]any{ + lookup := NewPluginLookup([]any{ &testHookInsertAndWorkBegin{}, &testHookInsertBegin{}, &testHookMetricEmit{}, @@ -152,8 +166,7 @@ func TestPluginLookup(t *testing.T) { &testMiddlewareJobInsertAndWorker{}, &testMiddlewareJobInsert{}, &testMiddlewareWorker{}, - }).(*pluginLookup) - require.True(t, isPluginLookup) + }) return lookup, &testBundle{} } @@ -244,12 +257,12 @@ func TestPluginLookup(t *testing.T) { legacyHook := &testLegacyHookInsertBegin{} legacyMiddleware := &testLegacyMiddlewareJobInsert{} - lookup, isPluginLookup := NewPluginLookupFromConfig( + lookup := NewPluginLookupFromConfig( + nil, []rivertype.Hook{legacyHook}, []rivertype.Middleware{legacyMiddleware}, nil, - ).(*pluginLookup) - require.True(t, isPluginLookup) + ) hookPlugins := lookup.ByKind(PluginKindHookInsertBegin) require.Len(t, hookPlugins, 1) @@ -305,6 +318,32 @@ func (jobArgsWithCustomHooks) Hooks() []rivertype.Hook { func (jobArgsWithCustomHooks) Kind() string { return "with_custom_hooks" } +// +// jobArgsWithCustomHooksAndPlugins +// + +var ( + _ rivertype.JobArgs = &jobArgsWithCustomHooksAndPlugins{} + _ jobArgsWithHooks = &jobArgsWithCustomHooksAndPlugins{} + _ jobArgsWithPlugins = &jobArgsWithCustomHooksAndPlugins{} +) + +type jobArgsWithCustomHooksAndPlugins struct{} + +func (jobArgsWithCustomHooksAndPlugins) Hooks() []rivertype.Hook { + return []rivertype.Hook{&testHookInsertBegin{}} +} + +func (jobArgsWithCustomHooksAndPlugins) Kind() string { return "with_custom_hooks_and_plugins" } + +func (jobArgsWithCustomHooksAndPlugins) Plugins() []rivertype.Plugin { + return []rivertype.Plugin{ + &testHookMiddlewarePlugin{}, + &testMiddlewareJobInsertAndWorker{}, + &testHookWorkEnd{}, + } +} + // // testHookInsertAndWorkBegin // diff --git a/internal/workunit/work_unit.go b/internal/workunit/work_unit.go index 3d6ff458..8becb697 100644 --- a/internal/workunit/work_unit.go +++ b/internal/workunit/work_unit.go @@ -19,7 +19,7 @@ type WorkUnit interface { // 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. - PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface + PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup Middleware() []rivertype.WorkerMiddleware NextRetry() time.Time diff --git a/job.go b/job.go index e2338f84..5a04bbda 100644 --- a/job.go +++ b/job.go @@ -92,3 +92,23 @@ type JobArgsWithInsertOpts interface { // system defaults. These can also be overridden at insertion time. InsertOpts() InsertOpts } + +// JobArgsWithPlugins is an interface that job args can implement to attach +// specific plugins (i.e. other than those globally installed to a client) to +// certain kinds of jobs. +type JobArgsWithPlugins interface { + // Plugins returns specific plugins to run for this job type. Plugin hooks + // run after global hooks, and plugin middleware is combined with other + // middleware configured for the job. + // + // A plugin implementing rivertype.JobInsertMiddleware runs once around an + // insertion batch containing one or more jobs of this type. Its InsertMany + // method receives the complete batch, which may also contain other job + // types. + // + // Warning: Plugins returned should be based on the job type only and be + // invariant of the specific contents of a job. Plugins are extracted by + // instantiating a generic instance of the job even when a specific instance + // is available, so any conditional logic within will be ignored. + Plugins() []rivertype.Plugin +} diff --git a/job_args_plugins_test.go b/job_args_plugins_test.go new file mode 100644 index 00000000..55af18cd --- /dev/null +++ b/job_args_plugins_test.go @@ -0,0 +1,106 @@ +package river + +import ( + "context" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/riverdbtest" + "github.com/riverqueue/river/riverdriver/riverpgxv5" + "github.com/riverqueue/river/rivershared/baseservice" + "github.com/riverqueue/river/rivershared/riversharedtest" + "github.com/riverqueue/river/rivertype" +) + +type jobArgsPluginsTestArgs func() []rivertype.Plugin + +func (jobArgsPluginsTestArgs) Kind() string { return "job_args_plugins_test" } + +func (jobArgsPluginsTestArgs) MarshalJSON() ([]byte, error) { return []byte("{}"), nil } + +func (f jobArgsPluginsTestArgs) Plugins() []rivertype.Plugin { return f() } + +func (jobArgsPluginsTestArgs) UnmarshalJSON([]byte) error { return nil } + +type jobArgsPluginsTestPlugin struct { + baseservice.BaseService + PluginDefaults + + insertBeginCount atomic.Int32 + insertManyCount atomic.Int32 + insertManyParamCount atomic.Int32 + workBeginCount atomic.Int32 + workMiddlewareCount atomic.Int32 +} + +func (p *jobArgsPluginsTestPlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error { + p.insertBeginCount.Add(1) + return nil +} + +func (p *jobArgsPluginsTestPlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) { + p.insertManyCount.Add(1) + p.insertManyParamCount.Store(int32(len(manyParams))) //nolint:gosec // test-only count will be small + return doInner(ctx) +} + +func (p *jobArgsPluginsTestPlugin) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error { + p.workMiddlewareCount.Add(1) + return doInner(ctx) +} + +func (p *jobArgsPluginsTestPlugin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { + p.workBeginCount.Add(1) + return nil +} + +var ( + _ JobArgsWithPlugins = jobArgsPluginsTestArgs(nil) + _ rivertype.HookInsertBegin = &jobArgsPluginsTestPlugin{} + _ rivertype.JobInsertMiddleware = &jobArgsPluginsTestPlugin{} + _ rivertype.Plugin = &jobArgsPluginsTestPlugin{} + _ rivertype.WorkerMiddleware = &jobArgsPluginsTestPlugin{} +) + +func TestJobArgsWithPlugins(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dbPool := riversharedtest.DBPool(ctx, t) + driver := riverpgxv5.New(dbPool) + schema := riverdbtest.TestSchema(ctx, t, driver, nil) + plugin := &jobArgsPluginsTestPlugin{} + args := jobArgsPluginsTestArgs(func() []rivertype.Plugin { return []rivertype.Plugin{plugin} }) + + config := newTestConfig(t, schema) + AddWorkerArgs(config.Workers, args, WorkFunc(func(ctx context.Context, job *Job[jobArgsPluginsTestArgs]) error { + return nil + })) + + client, err := NewClient(driver, config) + require.NoError(t, err) + require.NotEmpty(t, plugin.Name) + + subscribeChan := subscribe(t, client) + startClient(ctx, t, client) + + insertResults, err := client.InsertMany(ctx, []InsertManyParams{ + {Args: args}, + {Args: noOpArgs{}}, + {Args: args}, + }) + require.NoError(t, err) + require.Len(t, insertResults, 3) + + events := riversharedtest.WaitOrTimeoutN(t, subscribeChan, 3) + require.Equal(t, EventKindJobCompleted, events[0].Kind) + require.Equal(t, EventKindJobCompleted, events[1].Kind) + require.Equal(t, EventKindJobCompleted, events[2].Kind) + require.Equal(t, int32(2), plugin.insertBeginCount.Load()) + require.Equal(t, int32(1), plugin.insertManyCount.Load()) + require.Equal(t, int32(3), plugin.insertManyParamCount.Load()) + require.Equal(t, int32(2), plugin.workBeginCount.Load()) + require.Equal(t, int32(2), plugin.workMiddlewareCount.Load()) +} diff --git a/producer.go b/producer.go index 7f10fd77..da11246e 100644 --- a/producer.go +++ b/producer.go @@ -84,7 +84,7 @@ type producerConfig struct { FetchPollInterval time.Duration PluginLookupByJob *pluginlookup.JobPluginLookup - PluginLookupGlobal pluginlookup.PluginLookupInterface + PluginLookupGlobal *pluginlookup.PluginLookup JobStuckHandler JobStuckHandler JobStuckCount *atomic.Int32 JobStuckThreshold time.Duration diff --git a/producer_test.go b/producer_test.go index 1f0629c3..f5d17a07 100644 --- a/producer_test.go +++ b/producer_test.go @@ -54,31 +54,19 @@ func (p *beforeJobGetAvailablePilot) JobGetAvailable( return p.Pilot.JobGetAvailable(ctx, exec, state, params) } -type countingPluginLookup struct { - pluginlookup.PluginLookupInterface - - count int -} - -func (l *countingPluginLookup) ByKind(kind pluginlookup.PluginKind) []any { - l.count++ - return l.PluginLookupInterface.ByKind(kind) -} - func TestProducer_MetricEmitHook(t *testing.T) { t.Parallel() ctx := context.Background() type testBundle struct { - archetype *baseservice.Archetype - config *Config - exec riverdriver.Executor - metrics chan *rivertype.HookMetricEmitParams - pluginLookup *countingPluginLookup - producer *producer - queue string - schema string + archetype *baseservice.Archetype + config *Config + exec riverdriver.Executor + metrics chan *rivertype.HookMetricEmitParams + producer *producer + queue string + schema string } setup := func(t *testing.T) *testBundle { @@ -104,9 +92,7 @@ func TestProducer_MetricEmitHook(t *testing.T) { paramsCopy := *params metrics <- ¶msCopy }) - pluginLookup := &countingPluginLookup{ - PluginLookupInterface: pluginlookup.NewPluginLookup([]any{metricHook}), - } + pluginLookup := pluginlookup.NewPluginLookup([]any{metricHook}) producer := newProducer(archetype, exec, pilot, &producerConfig{ ClientID: testClientID, @@ -116,7 +102,7 @@ func TestProducer_MetricEmitHook(t *testing.T) { FetchPollInterval: 50 * time.Millisecond, JobTimeout: JobTimeoutDefault, MaxWorkers: 1_000, - PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), PluginLookupGlobal: pluginLookup, Queue: queueName, QueuePollInterval: queuePollIntervalDefault, @@ -129,14 +115,13 @@ func TestProducer_MetricEmitHook(t *testing.T) { }) return &testBundle{ - archetype: archetype, - config: newTestConfig(t, schema), - exec: exec, - metrics: metrics, - pluginLookup: pluginLookup, - producer: producer, - queue: queueName, - schema: schema, + archetype: archetype, + config: newTestConfig(t, schema), + exec: exec, + metrics: metrics, + producer: producer, + queue: queueName, + schema: schema, } } @@ -168,7 +153,7 @@ func TestProducer_MetricEmitHook(t *testing.T) { fetchResult := riversharedtest.WaitOrTimeout(t, fetchResultCh) require.NoError(t, fetchResult.err) require.Len(t, fetchResult.jobs, 2) - require.Equal(t, 1, bundle.pluginLookup.count) + require.Len(t, bundle.producer.metricEmitHooks, 1) metricsByName := make(map[rivertype.MetricName]rivertype.Metric) for _, metric := range riversharedtest.WaitOrTimeoutN(t, bundle.metrics, 2) { @@ -197,7 +182,7 @@ func TestProducer_MetricEmitHook(t *testing.T) { fetchResult := riversharedtest.WaitOrTimeout(t, fetchResultCh) require.NoError(t, fetchResult.err) require.Empty(t, fetchResult.jobs) - require.Equal(t, 1, bundle.pluginLookup.count) + require.Len(t, bundle.producer.metricEmitHooks, 1) require.Empty(t, bundle.metrics) }) } @@ -237,7 +222,7 @@ func TestProducer_PollOnly(t *testing.T) { ErrorHandler: newTestErrorHandler(), FetchCooldown: FetchCooldownDefault, FetchPollInterval: 50 * time.Millisecond, // more aggressive than normal because we have no notifier - PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobTimeout: JobTimeoutDefault, MaxWorkers: 1_000, @@ -290,7 +275,7 @@ 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 - PluginLookupByJob: pluginlookup.NewJobPluginLookup(), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), JobTimeout: JobTimeoutDefault, MaxWorkers: 1_000, diff --git a/rivertest/work_unit_wrapper.go b/rivertest/work_unit_wrapper.go index e36e4979..5514fb64 100644 --- a/rivertest/work_unit_wrapper.go +++ b/rivertest/work_unit_wrapper.go @@ -45,7 +45,7 @@ type wrapperWorkUnit[T river.JobArgs] struct { worker river.Worker[T] } -func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { +func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup { var job T return lookup.ByJobArgs(job) } diff --git a/rivertest/worker.go b/rivertest/worker.go index 8f0af787..414f846b 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -153,9 +153,6 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job middleware = pluginconfig.CombinedMiddleware(w.config.Middleware, w.config.JobInsertMiddleware, w.config.WorkerMiddleware) //nolint:staticcheck plugins = append(riverplugin.DefaultPlugins(), w.config.Plugins...) ) - pluginlookup.InitBaseServices(archetype, hooks) - pluginlookup.InitBaseServices(archetype, middleware) - pluginlookup.InitBaseServices(archetype, plugins) clientRetryPolicy := w.config.RetryPolicy if _, ok := clientRetryPolicy.(*river.DefaultClientRetryPolicy); ok { clientRetryPolicy = retrypolicy.NewDefault(archetype.Time) @@ -208,8 +205,8 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job return nil }, }, - PluginLookupByJob: pluginlookup.NewJobPluginLookup(), - PluginLookupGlobal: pluginlookup.NewPluginLookupFromConfig(hooks, middleware, plugins), + PluginLookupByJob: pluginlookup.NewJobPluginLookup(archetype), + PluginLookupGlobal: pluginlookup.NewPluginLookupFromConfig(archetype, hooks, middleware, plugins), JobRow: job, ProducerCallbacks: struct { JobDone func(jobRow *rivertype.JobRow) diff --git a/rivertype/river_type.go b/rivertype/river_type.go index 64a5387e..1abdf28e 100644 --- a/rivertype/river_type.go +++ b/rivertype/river_type.go @@ -644,8 +644,8 @@ func UniqueOptsByStateDefault() []JobState { // WorkerMetadata is metadata about workers registered with a client. type WorkerMetadata struct { - // JobArgHooks are job args specific hooks returned from a JobArgsWithHooks - // implementation. + // JobArgHooks are job args specific hooks returned from JobArgsWithHooks or + // from plugins returned by JobArgsWithPlugins. JobArgHooks []Hook // Kind is the kind returned from job args and recognized by worker to work. diff --git a/work_unit_wrapper.go b/work_unit_wrapper.go index e05b5749..7c1000fb 100644 --- a/work_unit_wrapper.go +++ b/work_unit_wrapper.go @@ -26,7 +26,7 @@ type wrapperWorkUnit[T JobArgs] struct { worker Worker[T] } -func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface { +func (w *wrapperWorkUnit[T]) PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup { var job T return lookup.ByJobArgs(job) }