Skip to content

Commit f8d3c2a

Browse files
brandurbrandur
committed
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.
1 parent 000c152 commit f8d3c2a

23 files changed

Lines changed: 673 additions & 863 deletions

‎client.go‎

Lines changed: 37 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,13 @@ import (
1717

1818
"github.com/riverqueue/river/internal/dblist"
1919
"github.com/riverqueue/river/internal/dbunique"
20-
"github.com/riverqueue/river/internal/hooklookup"
2120
"github.com/riverqueue/river/internal/jobcompleter"
2221
"github.com/riverqueue/river/internal/jobexecutor"
2322
"github.com/riverqueue/river/internal/leadership"
2423
"github.com/riverqueue/river/internal/maintenance"
25-
"github.com/riverqueue/river/internal/middlewarelookup"
2624
"github.com/riverqueue/river/internal/notifier"
2725
"github.com/riverqueue/river/internal/notifylimiter"
26+
"github.com/riverqueue/river/internal/pluginlookup"
2827
"github.com/riverqueue/river/internal/rivercommon"
2928
"github.com/riverqueue/river/internal/rivermiddleware"
3029
"github.com/riverqueue/river/internal/workunit"
@@ -545,54 +544,6 @@ func (c *Config) WithDefaults() *Config {
545544
}
546545
}
547546

548-
func effectiveHooksForConfig(config *Config) []rivertype.Hook {
549-
configuredMiddleware := middlewareFromConfig(config)
550-
hooks := make([]rivertype.Hook, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins))
551-
hooks = append(hooks, config.Hooks...)
552-
553-
for _, middleware := range configuredMiddleware {
554-
if hook, ok := middleware.(rivertype.Hook); ok {
555-
hooks = append(hooks, hook)
556-
}
557-
}
558-
559-
for _, plugin := range config.Plugins {
560-
if plugin == nil {
561-
continue
562-
}
563-
564-
if hook, ok := plugin.(rivertype.Hook); ok {
565-
hooks = append(hooks, hook)
566-
}
567-
}
568-
569-
return hooks
570-
}
571-
572-
func effectiveMiddlewareForConfig(config *Config) []rivertype.Middleware {
573-
configuredMiddleware := middlewareFromConfig(config)
574-
middleware := make([]rivertype.Middleware, 0, len(config.Hooks)+len(configuredMiddleware)+len(config.Plugins))
575-
middleware = append(middleware, configuredMiddleware...)
576-
577-
for _, hook := range config.Hooks {
578-
if middlewareItem, ok := hook.(rivertype.Middleware); ok {
579-
middleware = append(middleware, middlewareItem)
580-
}
581-
}
582-
583-
for _, plugin := range config.Plugins {
584-
if plugin == nil {
585-
continue
586-
}
587-
588-
if middlewareItem, ok := plugin.(rivertype.Middleware); ok {
589-
middleware = append(middleware, middlewareItem)
590-
}
591-
}
592-
593-
return middleware
594-
}
595-
596547
func middlewareFromConfig(config *Config) []rivertype.Middleware {
597548
middleware := make([]rivertype.Middleware, 0,
598549
len(config.Middleware)+len(config.JobInsertMiddleware)+len(config.WorkerMiddleware))
@@ -780,28 +731,27 @@ type Client[TTx any] struct {
780731
baseService baseservice.BaseService
781732
baseStartStop startstop.BaseStartStop
782733

783-
clientNotifyBundle *ClientNotifyBundle[TTx]
784-
completer jobcompleter.JobCompleter
785-
config *Config
786-
driver riverdriver.Driver[TTx]
787-
elector *leadership.Elector
788-
hookLookupByJob *hooklookup.JobHookLookup
789-
hookLookupGlobal hooklookup.HookLookupInterface
790-
insertNotifyLimiter *notifylimiter.Limiter
791-
middlewareLookupGlobal middlewarelookup.MiddlewareLookupInterface
792-
notifier *notifier.Notifier // may be nil in poll-only mode
793-
periodicJobs *PeriodicJobBundle
794-
pilot riverpilot.Pilot
795-
producersByQueueName map[string]*producer
796-
producersMu sync.RWMutex
797-
queueMaintainer *maintenance.QueueMaintainer
798-
queueMaintainerLeader *maintenance.QueueMaintainerLeader
799-
queues *QueueBundle
800-
services []startstop.Service
801-
stopped <-chan struct{}
802-
stuckJobCount atomic.Int32
803-
subscriptionManager *subscriptionManager
804-
testSignals clientTestSignals
734+
clientNotifyBundle *ClientNotifyBundle[TTx]
735+
completer jobcompleter.JobCompleter
736+
config *Config
737+
driver riverdriver.Driver[TTx]
738+
elector *leadership.Elector
739+
pluginLookupByJob *pluginlookup.JobPluginLookup
740+
pluginLookupGlobal pluginlookup.PluginLookupInterface
741+
insertNotifyLimiter *notifylimiter.Limiter
742+
notifier *notifier.Notifier // may be nil in poll-only mode
743+
periodicJobs *PeriodicJobBundle
744+
pilot riverpilot.Pilot
745+
producersByQueueName map[string]*producer
746+
producersMu sync.RWMutex
747+
queueMaintainer *maintenance.QueueMaintainer
748+
queueMaintainerLeader *maintenance.QueueMaintainerLeader
749+
queues *QueueBundle
750+
services []startstop.Service
751+
stopped <-chan struct{}
752+
stuckJobCount atomic.Int32
753+
subscriptionManager *subscriptionManager
754+
testSignals clientTestSignals
805755

806756
// workCancel cancels the context used for all work goroutines. Normal Stop
807757
// does not cancel that context.
@@ -904,11 +854,14 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
904854
}
905855
}
906856

907-
effectiveHooks := effectiveHooksForConfig(config)
908-
effectiveMiddleware := effectiveMiddlewareForConfig(config)
857+
var (
858+
configuredMiddleware = middlewareFromConfig(config)
859+
allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...)
860+
allPlugins = pluginlookup.NormalizePlugins(config.Hooks, allMiddleware, config.Plugins)
861+
)
909862

910-
for _, hook := range effectiveHooks {
911-
if withBaseService, ok := hook.(baseservice.WithBaseService); ok {
863+
for _, plugin := range allPlugins {
864+
if withBaseService, ok := plugin.(baseservice.WithBaseService); ok {
912865
baseservice.Init(archetype, withBaseService)
913866
}
914867
}
@@ -920,8 +873,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
920873
},
921874
config: config,
922875
driver: driver,
923-
hookLookupByJob: hooklookup.NewJobHookLookup(),
924-
hookLookupGlobal: hooklookup.NewHookLookup(effectiveHooks),
876+
pluginLookupByJob: pluginlookup.NewJobPluginLookup(),
877+
pluginLookupGlobal: pluginlookup.NewPluginLookup(allPlugins),
925878
producersByQueueName: make(map[string]*producer),
926879
testSignals: clientTestSignals{},
927880
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
939892
client.baseService.Name = "Client" // Have to correct the name because base service isn't embedded like it usually is
940893
client.insertNotifyLimiter = notifylimiter.NewLimiter(archetype, config.FetchCooldown)
941894

942-
// effectiveMiddleware contains configured middleware, hook/middleware
943-
// hybrids, and plugins. Default middleware stays first so user middleware
944-
// wraps inside River's internal defaults like before.
945-
{
946-
middleware := rivermiddleware.DefaultMiddleware()
947-
middleware = append(middleware, effectiveMiddleware...)
948-
949-
for _, middleware := range middleware {
950-
if withBaseService, ok := middleware.(baseservice.WithBaseService); ok {
951-
baseservice.Init(archetype, withBaseService)
952-
}
953-
}
954-
955-
client.middlewareLookupGlobal = middlewarelookup.NewMiddlewareLookup(middleware)
956-
}
957-
958895
pluginDriver, _ := driver.(driverPlugin[TTx])
959896
if pluginDriver != nil {
960897
pluginDriver.PluginInit(archetype)
@@ -1082,7 +1019,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
10821019
{
10831020
periodicJobEnqueuer, err := maintenance.NewPeriodicJobEnqueuer(archetype, &maintenance.PeriodicJobEnqueuerConfig{
10841021
AdvisoryLockPrefix: config.AdvisoryLockPrefix,
1085-
HookLookupGlobal: client.hookLookupGlobal,
1022+
PluginLookupGlobal: client.pluginLookupGlobal,
10861023
Insert: client.insertMany,
10871024
Pilot: client.pilot,
10881025
Schema: config.Schema,
@@ -2062,8 +1999,8 @@ func (c *Client[TTx]) insertManyShared(
20621999
doInner := func(ctx context.Context) ([]*rivertype.JobInsertResult, error) {
20632000
for _, params := range insertParams {
20642001
for _, hook := range append(
2065-
c.hookLookupGlobal.ByHookKind(hooklookup.HookKindInsertBegin),
2066-
c.hookLookupByJob.ByJobArgs(params.Args).ByHookKind(hooklookup.HookKindInsertBegin)...,
2002+
c.pluginLookupGlobal.ByKind(pluginlookup.HookKindInsertBegin),
2003+
c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.HookKindInsertBegin)...,
20672004
) {
20682005
if err := hook.(rivertype.HookInsertBegin).InsertBegin(ctx, params); err != nil { //nolint:forcetypeassert
20692006
return nil, err
@@ -2094,7 +2031,7 @@ func (c *Client[TTx]) insertManyShared(
20942031
return insertResults, nil
20952032
}
20962033

2097-
jobInsertMiddleware := c.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert)
2034+
jobInsertMiddleware := c.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert)
20982035
if len(jobInsertMiddleware) > 0 {
20992036
// Wrap middlewares in reverse order so the one defined first is wrapped
21002037
// 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
23722309
ErrorHandler: c.config.ErrorHandler,
23732310
FetchCooldown: cmp.Or(queueConfig.FetchCooldown, c.config.FetchCooldown),
23742311
FetchPollInterval: cmp.Or(queueConfig.FetchPollInterval, c.config.FetchPollInterval),
2375-
HookLookupByJob: c.hookLookupByJob,
2376-
HookLookupGlobal: c.hookLookupGlobal,
2312+
PluginLookupByJob: c.pluginLookupByJob,
2313+
PluginLookupGlobal: c.pluginLookupGlobal,
23772314
JobStuckHandler: c.config.JobStuckHandler,
23782315
JobStuckCount: &c.stuckJobCount,
23792316
JobStuckThreshold: c.config.JobStuckThreshold,
23802317
JobTimeout: c.config.JobTimeout,
23812318
MaxWorkers: queueConfig.MaxWorkers,
2382-
MiddlewareLookupGlobal: c.middlewareLookupGlobal,
23832319
Notifier: c.notifier,
23842320
Queue: queueName,
23852321
QueueEventCallback: c.subscriptionManager.distributeQueueEvent,

‎client_test.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ import (
2626
"github.com/riverqueue/river/internal/dbunique"
2727
"github.com/riverqueue/river/internal/jobexecutor"
2828
"github.com/riverqueue/river/internal/maintenance"
29-
"github.com/riverqueue/river/internal/middlewarelookup"
3029
"github.com/riverqueue/river/internal/notifier"
30+
"github.com/riverqueue/river/internal/pluginlookup"
3131
"github.com/riverqueue/river/internal/rivercommon"
3232
"github.com/riverqueue/river/internal/riverinternaltest"
3333
"github.com/riverqueue/river/internal/riverinternaltest/retrypolicytest"
@@ -8579,8 +8579,8 @@ func Test_NewClient_Validations(t *testing.T) {
85798579
config.Middleware = []rivertype.Middleware{&overridableJobMiddleware{}}
85808580
},
85818581
validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
8582-
require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 1)
8583-
require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2)
8582+
require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert), 1)
8583+
require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker), 2)
85848584
},
85858585
},
85868586
{
@@ -8590,8 +8590,8 @@ func Test_NewClient_Validations(t *testing.T) {
85908590
config.WorkerMiddleware = []rivertype.WorkerMiddleware{&overridableJobMiddleware{}}
85918591
},
85928592
validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
8593-
require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 2)
8594-
require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 3)
8593+
require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert), 2)
8594+
require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker), 3)
85958595
},
85968596
},
85978597
{
@@ -8602,8 +8602,8 @@ func Test_NewClient_Validations(t *testing.T) {
86028602
config.WorkerMiddleware = []rivertype.WorkerMiddleware{middleware}
86038603
},
86048604
validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
8605-
require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 1)
8606-
require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2)
8605+
require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindJobInsert), 1)
8606+
require.Len(t, client.pluginLookupGlobal.ByKind(pluginlookup.MiddlewareKindWorker), 2)
86078607
},
86088608
},
86098609
{

‎hook_defaults_funcs_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,17 @@ import (
1010
var (
1111
_ rivertype.Hook = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil })
1212
_ rivertype.HookInsertBegin = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil })
13+
_ rivertype.Plugin = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil })
1314

1415
_ rivertype.Hook = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil })
1516
_ rivertype.HookPeriodicJobsStart = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil })
17+
_ rivertype.Plugin = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil })
1618

1719
_ rivertype.Hook = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil })
1820
_ rivertype.HookWorkBegin = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil })
21+
_ rivertype.Plugin = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil })
22+
23+
_ rivertype.Hook = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return nil })
24+
_ rivertype.HookWorkEnd = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return nil })
25+
_ rivertype.Plugin = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return nil })
1926
)

0 commit comments

Comments
 (0)