Move extracting middleware config over to shared internal code

Author: brandur

client.go

@@ -23,6 +23,7 @@ import (
 	"github.com/riverqueue/river/internal/maintenance"
 	"github.com/riverqueue/river/internal/notifier"
 	"github.com/riverqueue/river/internal/notifylimiter"
+	"github.com/riverqueue/river/internal/pluginconfig"
 	"github.com/riverqueue/river/internal/pluginlookup"
 	"github.com/riverqueue/river/internal/rivercommon"
 	"github.com/riverqueue/river/internal/riverplugin"
@@ -544,35 +545,6 @@ func (c *Config) WithDefaults() *Config {
 	}
 }
 
-func middlewareFromConfig(config *Config) []rivertype.Middleware {
-	middleware := make([]rivertype.Middleware, 0,
-		len(config.Middleware)+len(config.JobInsertMiddleware)+len(config.WorkerMiddleware))
-	middleware = append(middleware, config.Middleware...)
-
-	for _, jobInsertMiddleware := range config.JobInsertMiddleware {
-		middleware = append(middleware, jobInsertMiddleware)
-	}
-
-outerLoop:
-	for _, workerMiddleware := range config.WorkerMiddleware {
-		// Don't add the middleware if it also implements JobInsertMiddleware
-		// and the instance has been added to config.JobInsertMiddleware. This
-		// is a hedge to make sure we don't accidentally double add middleware
-		// as we've converted over to the unified config.Middleware setting.
-		if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddleware.(rivertype.JobInsertMiddleware); ok {
-			for _, jobInsertMiddleware := range config.JobInsertMiddleware {
-				if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddleware {
-					continue outerLoop
-				}
-			}
-		}
-
-		middleware = append(middleware, workerMiddleware)
-	}
-
-	return middleware
-}
-
 func (c *Config) validate() error {
 	if c.CancelledJobRetentionPeriod < -1 {
 		return errors.New("CancelledJobRetentionPeriod time cannot be less than zero, except for -1 (infinite)")
@@ -855,7 +827,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 	}
 
 	var (
-		configuredMiddleware = middlewareFromConfig(config)
+		configuredMiddleware = pluginconfig.Middleware(config.Middleware, config.JobInsertMiddleware, config.WorkerMiddleware)
 		plugins              = append(riverplugin.DefaultPlugins(), config.Plugins...)
 		allPlugins           = pluginlookup.NormalizePlugins(config.Hooks, configuredMiddleware, plugins)
 	)

internal/pluginconfig/plugin_config.go

@@ -2,42 +2,31 @@ package pluginconfig
 
 import "github.com/riverqueue/river/rivertype"
 
-// Hooks returns the effective hook list from configured hooks and middleware.
-// Explicit hooks are preserved first, followed by middleware that also
-// implement hooks.
-func Hooks(hooks []rivertype.Hook, middleware []rivertype.Middleware) []rivertype.Hook {
-	effectiveHooks := make([]rivertype.Hook, 0, len(hooks)+len(middleware))
-
-	effectiveHooks = append(effectiveHooks, hooks...)
-
-	for _, middlewareItem := range middleware {
-		hook, ok := middlewareItem.(rivertype.Hook)
-		if !ok {
-			continue
-		}
-
-		effectiveHooks = append(effectiveHooks, hook)
+// Middleware combines middleware from the current and legacy configuration
+// fields into one list. Explicit middleware is preserved first, followed by
+// legacy job insert and worker middleware. A worker middleware already present
+// in the job insert list is included only once.
+func Middleware(middleware []rivertype.Middleware, jobInsertMiddleware []rivertype.JobInsertMiddleware, workerMiddleware []rivertype.WorkerMiddleware) []rivertype.Middleware {
+	allMiddleware := make([]rivertype.Middleware, 0,
+		len(middleware)+len(jobInsertMiddleware)+len(workerMiddleware))
+	allMiddleware = append(allMiddleware, middleware...)
+
+	for _, jobInsertMiddlewareItem := range jobInsertMiddleware {
+		allMiddleware = append(allMiddleware, jobInsertMiddlewareItem)
 	}
 
-	return effectiveHooks
-}
-
-// Middleware returns the effective middleware list from configured hooks and
-// middleware. Explicit middleware are preserved first, followed by hooks that
-// also implement middleware.
-func Middleware(hooks []rivertype.Hook, middleware []rivertype.Middleware) []rivertype.Middleware {
-	effectiveMiddleware := make([]rivertype.Middleware, 0, len(hooks)+len(middleware))
-
-	effectiveMiddleware = append(effectiveMiddleware, middleware...)
-
-	for _, hook := range hooks {
-		middlewareItem, ok := hook.(rivertype.Middleware)
-		if !ok {
-			continue
+outerLoop:
+	for _, workerMiddlewareItem := range workerMiddleware {
+		if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddlewareItem.(rivertype.JobInsertMiddleware); ok {
+			for _, jobInsertMiddlewareItem := range jobInsertMiddleware {
+				if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddlewareItem {
+					continue outerLoop
+				}
+			}
 		}
 
-		effectiveMiddleware = append(effectiveMiddleware, middlewareItem)
+		allMiddleware = append(allMiddleware, workerMiddlewareItem)
 	}
 
-	return effectiveMiddleware
+	return allMiddleware
 }

internal/pluginconfig/plugin_config_test.go

@@ -0,0 +1,64 @@
+package pluginconfig
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/riverqueue/river/rivertype"
+)
+
+func TestMiddleware(t *testing.T) {
+	t.Parallel()
+
+	t.Run("CombinesMiddlewareInOrder", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			middleware          = &testMiddleware{}
+			jobInsertMiddleware = &testJobInsertMiddleware{}
+			workerMiddleware    = &testWorkerMiddleware{}
+		)
+
+		require.Equal(t, []rivertype.Middleware{
+			middleware,
+			jobInsertMiddleware,
+			workerMiddleware,
+		}, Middleware(
+			[]rivertype.Middleware{middleware},
+			[]rivertype.JobInsertMiddleware{jobInsertMiddleware},
+			[]rivertype.WorkerMiddleware{workerMiddleware},
+		))
+	})
+
+	t.Run("SameSpecificMiddlewareIncludedOnce", func(t *testing.T) {
+		t.Parallel()
+
+		middleware := &testJobInsertAndWorkerMiddleware{}
+
+		require.Equal(t, []rivertype.Middleware{middleware}, Middleware(
+			nil,
+			[]rivertype.JobInsertMiddleware{middleware},
+			[]rivertype.WorkerMiddleware{middleware},
+		))
+	})
+}
+
+type testJobInsertAndWorkerMiddleware struct {
+	rivertype.JobInsertMiddleware
+	rivertype.WorkerMiddleware
+}
+
+func (*testJobInsertAndWorkerMiddleware) IsMiddleware() bool { return true }
+
+type testJobInsertMiddleware struct {
+	rivertype.JobInsertMiddleware
+}
+
+type testMiddleware struct{}
+
+func (*testMiddleware) IsMiddleware() bool { return true }
+
+type testWorkerMiddleware struct {
+	rivertype.WorkerMiddleware
+}

rivertest/worker.go

@@ -11,6 +11,7 @@ import (
 	"github.com/riverqueue/river/internal/jobcompleter"
 	"github.com/riverqueue/river/internal/jobexecutor"
 	"github.com/riverqueue/river/internal/maintenance"
+	"github.com/riverqueue/river/internal/pluginconfig"
 	"github.com/riverqueue/river/internal/pluginlookup"
 	"github.com/riverqueue/river/internal/riverplugin"
 	"github.com/riverqueue/river/riverdriver"
@@ -147,9 +148,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)
-		plugins              = append(riverplugin.DefaultPlugins(), w.config.Plugins...)
-		allPlugins           = pluginlookup.NormalizePlugins(
+		configuredMiddleware = pluginconfig.Middleware(
+			w.config.Middleware,          //nolint:staticcheck // Bridge legacy middleware into the test worker's unified plugin path.
+			w.config.JobInsertMiddleware, //nolint:staticcheck // Bridge legacy middleware into the test worker's unified plugin path.
+			w.config.WorkerMiddleware,    //nolint:staticcheck // Bridge legacy middleware into the test worker's unified plugin path.
+		)
+		plugins    = append(riverplugin.DefaultPlugins(), w.config.Plugins...)
+		allPlugins = pluginlookup.NormalizePlugins(
 			w.config.Hooks, //nolint:staticcheck // Bridge legacy Hooks into the test worker's unified plugin path.
 			configuredMiddleware,
 			plugins,
@@ -238,36 +243,6 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job
 	panic("unreachable")
 }
 
-//nolint:staticcheck
-func middlewareFromConfig(config *river.Config) []rivertype.Middleware {
-	middleware := make([]rivertype.Middleware, 0,
-		len(config.Middleware)+len(config.JobInsertMiddleware)+len(config.WorkerMiddleware))
-	middleware = append(middleware, config.Middleware...)
-
-	for _, jobInsertMiddleware := range config.JobInsertMiddleware {
-		middleware = append(middleware, jobInsertMiddleware)
-	}
-
-outerLoop:
-	for _, workerMiddleware := range config.WorkerMiddleware {
-		// Don't add the middleware if it also implements JobInsertMiddleware
-		// and the instance has been added to config.JobInsertMiddleware. This
-		// is a hedge to make sure we don't accidentally double add middleware
-		// as we've converted over to the unified config.Middleware setting.
-		if workerMiddlewareAsJobInsertMiddleware, ok := workerMiddleware.(rivertype.JobInsertMiddleware); ok {
-			for _, jobInsertMiddleware := range config.JobInsertMiddleware {
-				if workerMiddlewareAsJobInsertMiddleware == jobInsertMiddleware {
-					continue outerLoop
-				}
-			}
-		}
-
-		middleware = append(middleware, workerMiddleware)
-	}
-
-	return middleware
-}
-
 // WorkResult is the result of working a job in the test Worker.
 type WorkResult struct {
 	// EventKind is the kind of event that occurred following execution.