Skip to content

Commit 9873cfd

Browse files
brandurbrandur
committed
Cleaner use of internal plugins + plugin initialization
1 parent 4c5579e commit 9873cfd

8 files changed

Lines changed: 103 additions & 43 deletions

File tree

‎client.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import (
2525
"github.com/riverqueue/river/internal/notifylimiter"
2626
"github.com/riverqueue/river/internal/pluginlookup"
2727
"github.com/riverqueue/river/internal/rivercommon"
28-
"github.com/riverqueue/river/internal/rivermiddleware"
28+
"github.com/riverqueue/river/internal/riverplugin"
2929
"github.com/riverqueue/river/internal/workunit"
3030
"github.com/riverqueue/river/riverdriver"
3131
"github.com/riverqueue/river/rivershared/baseservice"
@@ -854,10 +854,10 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
854854
}
855855
}
856856

857-
var (
858-
configuredMiddleware = middlewareFromConfig(config)
859-
allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...)
860-
allPlugins = pluginlookup.NormalizePlugins(config.Hooks, allMiddleware, config.Plugins)
857+
allPlugins := pluginlookup.NormalizePlugins(
858+
config.Hooks,
859+
middlewareFromConfig(config),
860+
append(riverplugin.DefaultPlugins(), config.Plugins...),
861861
)
862862

863863
for _, plugin := range allPlugins {

‎internal/pluginlookup/plugin_lookup.go‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -89,26 +89,30 @@ func NewPluginLookup(plugins []rivertype.Plugin) PluginLookupInterface {
8989
// NormalizePlugins converts hook, middleware, and plugin registrations into a
9090
// single plugin slice while preserving legacy hook and middleware registrations
9191
// that don't yet opt into Plugin.
92+
93+
// Plugins passed explicitly are ordered before hooks and middleware, so
94+
// Config.Plugins entries take precedence over plugins bridged from Config.Hooks
95+
// or Config.Middleware. Relative order is preserved within each input slice.
9296
func NormalizePlugins(hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Plugin {
9397
normalizedPlugins := make([]rivertype.Plugin, 0, len(hooks)+len(middlewares)+len(plugins))
9498

99+
normalizedPlugins = append(normalizedPlugins, plugins...)
100+
95101
for _, hook := range hooks {
96102
if plugin, ok := hook.(rivertype.Plugin); ok {
97103
normalizedPlugins = append(normalizedPlugins, plugin)
98-
continue
104+
} else {
105+
normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook))
99106
}
100-
101-
normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook))
102107
}
108+
103109
for _, middleware := range middlewares {
104110
if plugin, ok := middleware.(rivertype.Plugin); ok {
105111
normalizedPlugins = append(normalizedPlugins, plugin)
106-
continue
112+
} else {
113+
normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware))
107114
}
108-
109-
normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware))
110115
}
111-
normalizedPlugins = append(normalizedPlugins, plugins...)
112116

113117
return normalizedPlugins
114118
}

‎internal/pluginlookup/plugin_lookup_test.go‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,36 @@ func TestJobPluginLookup(t *testing.T) {
110110
})
111111
}
112112

113+
func TestNormalizePlugins(t *testing.T) {
114+
t.Parallel()
115+
116+
t.Run("PluginsPrecedeHooksAndMiddleware", func(t *testing.T) {
117+
t.Parallel()
118+
119+
hookPlugin := &testHookMiddlewarePlugin{}
120+
middlewarePlugin := &testHookMiddlewarePlugin{}
121+
plugin := &testHookMiddlewarePlugin{}
122+
123+
lookup, isPluginLookup := NewPluginLookup(NormalizePlugins(
124+
[]rivertype.Hook{hookPlugin},
125+
[]rivertype.Middleware{middlewarePlugin},
126+
[]rivertype.Plugin{plugin},
127+
)).(*pluginLookup)
128+
require.True(t, isPluginLookup)
129+
130+
require.Equal(t, []rivertype.Plugin{
131+
plugin,
132+
hookPlugin,
133+
middlewarePlugin,
134+
}, lookup.ByKind(HookKindInsertBegin))
135+
require.Equal(t, []rivertype.Plugin{
136+
plugin,
137+
hookPlugin,
138+
middlewarePlugin,
139+
}, lookup.ByKind(MiddlewareKindJobInsert))
140+
})
141+
}
142+
113143
func TestPluginLookup(t *testing.T) {
114144
t.Parallel()
115145

@@ -307,6 +337,29 @@ func (t *testHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype
307337
return nil
308338
}
309339

340+
//
341+
// testHookMiddlewarePlugin
342+
//
343+
344+
var (
345+
_ rivertype.HookInsertBegin = &testHookMiddlewarePlugin{}
346+
_ rivertype.JobInsertMiddleware = &testHookMiddlewarePlugin{}
347+
)
348+
349+
type testHookMiddlewarePlugin struct{}
350+
351+
func (t *testHookMiddlewarePlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error {
352+
return nil
353+
}
354+
355+
func (t *testHookMiddlewarePlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
356+
return doInner(ctx)
357+
}
358+
359+
func (t *testHookMiddlewarePlugin) IsHook() bool { return true }
360+
func (t *testHookMiddlewarePlugin) IsMiddleware() bool { return true }
361+
func (t *testHookMiddlewarePlugin) IsPlugin() bool { return true }
362+
310363
//
311364
// testHookWorkBegin
312365
//

internal/rivermiddleware/middleware.go renamed to internal/riverplugin/plugin.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package rivermiddleware
1+
package riverplugin
22

33
import (
44
"context"
@@ -13,10 +13,10 @@ import (
1313
"github.com/riverqueue/river/rivertype"
1414
)
1515

16-
// DefaultMiddleware returns the default middleware that River applies to all
17-
// jobs. This includes internal middleware like the resumable step middleware.
18-
func DefaultMiddleware() []rivertype.Middleware {
19-
return []rivertype.Middleware{&ResumableMiddleware{}}
16+
// DefaultPlugins returns the default plugins that River applies to all jobs.
17+
// This includes internal middleware like the resumable step middleware.
18+
func DefaultPlugins() []rivertype.Plugin {
19+
return []rivertype.Plugin{&ResumableMiddleware{}}
2020
}
2121

2222
// ResumableMiddleware is internal middleware that enables resumable step

‎resumable.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import (
66
"errors"
77
"fmt"
88

9-
"github.com/riverqueue/river/internal/rivermiddleware"
9+
"github.com/riverqueue/river/internal/riverplugin"
1010
)
1111

1212
var (
@@ -150,7 +150,7 @@ func ResumableStepCursor[TCursor any](ctx context.Context, name string, opts *St
150150
delete(state.Cursors, name)
151151
}
152152

153-
func mustResumableState(ctx context.Context) *rivermiddleware.ResumableState {
153+
func mustResumableState(ctx context.Context) *riverplugin.ResumableState {
154154
state, ok := resumableStateFromContext(ctx)
155155
if !ok {
156156
panic(errResumableStepNotInWorker)
@@ -159,7 +159,7 @@ func mustResumableState(ctx context.Context) *rivermiddleware.ResumableState {
159159
return state
160160
}
161161

162-
func registerResumableStepName(state *rivermiddleware.ResumableState, name string) bool {
162+
func registerResumableStepName(state *riverplugin.ResumableState, name string) bool {
163163
if _, ok := state.AllStepNames[name]; ok {
164164
state.Err = fmt.Errorf("river: duplicate resumable step name %q", name)
165165
return false
@@ -169,13 +169,13 @@ func registerResumableStepName(state *rivermiddleware.ResumableState, name strin
169169
return true
170170
}
171171

172-
func resumableStateFromContext(ctx context.Context) (*rivermiddleware.ResumableState, bool) {
173-
state := ctx.Value(rivermiddleware.ResumableContextKey{})
172+
func resumableStateFromContext(ctx context.Context) (*riverplugin.ResumableState, bool) {
173+
state := ctx.Value(riverplugin.ResumableContextKey{})
174174
if state == nil {
175175
return nil, false
176176
}
177177

178-
typedState, ok := state.(*rivermiddleware.ResumableState)
178+
typedState, ok := state.(*riverplugin.ResumableState)
179179
if !ok || typedState == nil {
180180
return nil, false
181181
}

‎resumable_step_tx_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111
"github.com/riverqueue/river/internal/execution"
1212
"github.com/riverqueue/river/internal/jobexecutor"
1313
"github.com/riverqueue/river/internal/rivercommon"
14-
"github.com/riverqueue/river/internal/rivermiddleware"
14+
"github.com/riverqueue/river/internal/riverplugin"
1515
"github.com/riverqueue/river/riverdbtest"
1616
"github.com/riverqueue/river/riverdriver"
1717
"github.com/riverqueue/river/riverdriver/riverpgxv5"
@@ -47,7 +47,7 @@ func TestResumableSetStepTx(t *testing.T) {
4747
require.NoError(t, err)
4848
ctx = context.WithValue(ctx, rivercommon.ContextKeyClient{}, client)
4949
ctx = context.WithValue(ctx, jobexecutor.ContextKeyMetadataUpdates, make(map[string]any))
50-
ctx = context.WithValue(ctx, rivermiddleware.ResumableContextKey{}, &rivermiddleware.ResumableState{
50+
ctx = context.WithValue(ctx, riverplugin.ResumableContextKey{}, &riverplugin.ResumableState{
5151
Cursors: make(map[string]json.RawMessage),
5252
StepName: stepName,
5353
})

‎resumable_test.go‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import (
1010

1111
"github.com/riverqueue/river/internal/jobexecutor"
1212
"github.com/riverqueue/river/internal/rivercommon"
13-
"github.com/riverqueue/river/internal/rivermiddleware"
13+
"github.com/riverqueue/river/internal/riverplugin"
1414
"github.com/riverqueue/river/rivertype"
1515
)
1616

@@ -32,7 +32,7 @@ func TestResumableStep(t *testing.T) {
3232
ctx, _, job := setup(t, `{}`)
3333

3434
var ran []string
35-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
35+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
3636
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
3737
ran = append(ran, "first")
3838
return nil
@@ -54,7 +54,7 @@ func TestResumableStep(t *testing.T) {
5454
ctx, _, job := setup(t, `{"river:resumable_step":"step2"}`)
5555

5656
var ran []string
57-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
57+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
5858
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
5959
ran = append(ran, "first")
6060
return nil
@@ -88,7 +88,7 @@ func TestResumableStep(t *testing.T) {
8888
ctx, metadataUpdates, job := setup(t, `{}`)
8989

9090
var ran []string
91-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
91+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
9292
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
9393
ran = append(ran, "step1")
9494
return nil
@@ -111,7 +111,7 @@ func TestResumableStep(t *testing.T) {
111111
ctx, metadataUpdates, job = setup(t, `{"river:resumable_step":"step1"}`)
112112
ran = nil
113113

114-
err = (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
114+
err = (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
115115
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
116116
ran = append(ran, "step1")
117117
return nil
@@ -140,7 +140,7 @@ func TestResumableStep(t *testing.T) {
140140
defer cancel()
141141

142142
var ran []string
143-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
143+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
144144
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
145145
ran = append(ran, "step1")
146146
cancel()
@@ -185,7 +185,7 @@ func TestResumableStepCursor(t *testing.T) {
185185
ctx, _, job := setup(t, `{}`)
186186

187187
var ran []string
188-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
188+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
189189
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
190190
ran = append(ran, "step")
191191
return nil
@@ -210,7 +210,7 @@ func TestResumableStepCursor(t *testing.T) {
210210
cursorResult resumableCursor
211211
ran []int
212212
)
213-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
213+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
214214
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
215215
ran = append(ran, 1)
216216
return nil
@@ -240,7 +240,7 @@ func TestResumableStepCursor(t *testing.T) {
240240
cursorResult = resumableCursor{}
241241
ran = nil
242242

243-
err = (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
243+
err = (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
244244
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
245245
ran = append(ran, 1)
246246
return nil
@@ -271,7 +271,7 @@ func TestResumableStepCursor(t *testing.T) {
271271
ctx, metadataUpdates, job := setup(t, `{"river:resumable_cursor":{"step2":{"id":42}},"river:resumable_step":"step1"}`)
272272

273273
var ran []int
274-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
274+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
275275
ResumableStep(ctx, "step1", nil, func(ctx context.Context) error {
276276
ran = append(ran, 1)
277277
return nil
@@ -306,7 +306,7 @@ func TestResumableStepCursor(t *testing.T) {
306306

307307
ctx, _, _ := setup(t, `{}`)
308308

309-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, &rivertype.JobRow{Metadata: []byte(`{}`)}, func(ctx context.Context) error {
309+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, &rivertype.JobRow{Metadata: []byte(`{}`)}, func(ctx context.Context) error {
310310
return ResumableSetCursor(ctx, 1)
311311
})
312312
require.ErrorIs(t, err, errResumableCursorNotInStep)
@@ -321,7 +321,7 @@ func TestResumableStepCursor(t *testing.T) {
321321

322322
ctx, metadataUpdates, job := setup(t, `{}`)
323323

324-
err := (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
324+
err := (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
325325
ResumableStepCursor(ctx, "step1", nil, func(ctx context.Context, cursor int) error {
326326
require.Zero(t, cursor)
327327
require.NoError(t, ResumableSetCursor(ctx, 123))
@@ -343,7 +343,7 @@ func TestResumableStepCursor(t *testing.T) {
343343

344344
ctx, metadataUpdates, job = setup(t, `{"river:resumable_cursor":{"step1":123,"step2":{"id":"abc"}},"river:resumable_step":"step1"}`)
345345

346-
err = (&rivermiddleware.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
346+
err = (&riverplugin.ResumableMiddleware{}).Work(ctx, job, func(ctx context.Context) error {
347347
ResumableStepCursor(ctx, "step1", nil, func(ctx context.Context, cursor int) error {
348348
require.Equal(t, 123, cursor)
349349
return nil

‎rivertest/worker.go‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212
"github.com/riverqueue/river/internal/jobexecutor"
1313
"github.com/riverqueue/river/internal/maintenance"
1414
"github.com/riverqueue/river/internal/pluginlookup"
15-
"github.com/riverqueue/river/internal/rivermiddleware"
15+
"github.com/riverqueue/river/internal/riverplugin"
1616
"github.com/riverqueue/river/riverdriver"
1717
"github.com/riverqueue/river/rivershared/baseservice"
1818
"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
146146
}
147147
completer := jobcompleter.NewInlineCompleter(archetype, w.config.Schema, exec, w.client.Pilot(), subscribeCh)
148148

149-
var (
150-
configuredMiddleware = middlewareFromConfig(w.config)
151-
allMiddleware = append(rivermiddleware.DefaultMiddleware(), configuredMiddleware...)
152-
allPlugins = pluginlookup.NormalizePlugins(w.config.Hooks, allMiddleware, w.config.Plugins) //nolint:staticcheck // Bridge legacy Hooks into the test worker's unified plugin path.
149+
allPlugins := pluginlookup.NormalizePlugins(
150+
w.config.Hooks, //nolint:staticcheck // Bridge legacy Hooks into the test worker's unified plugin path.
151+
middlewareFromConfig(w.config),
152+
append(
153+
riverplugin.DefaultPlugins(),
154+
w.config.Plugins...,
155+
),
153156
)
154157

155158
for _, plugin := range allPlugins {

0 commit comments

Comments
 (0)