Inject worker metadata to pilot + testable job args

Here, add a mechanism that has the client inject worker metadata into
pilots. This gives pilots information about what jobs/workers are
registered, and helps in cases where say we want pilots to be able to
examine things like `Hooks()` implementations on each worker.

I also ended up adding a function for `river.AddWorkerArgs` that lets a
concrete args object be injected into the registration process rather
than one that's implicitly allocated. I'm a little partial on this one
because it'd be better if it wasn't necessary, but it cleans up tests
somewhat that are trying to test specific things on args like `Hooks()`,
and make tests on features like that _possible_ when there's no
workaround available like storing information that'll be persisted to a
job's metadata (there's otherwise no good way to carry them out, with
the only way I can think of a global variable)

Author: brandur

client.go

@@ -776,10 +776,29 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		pluginDriver.PluginInit(archetype)
 		client.pilot = pluginDriver.PluginPilot()
 	}
+
+	var workerMetadata []*rivertype.WorkerMetadata
+	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,
+				Kind:        kind,
+			})
+		}
+	}
+
 	if client.pilot == nil {
 		client.pilot = &riverpilot.StandardPilot{}
 	}
-	client.pilot.PilotInit(archetype)
+	client.pilot.PilotInit(archetype, &riverpilot.PilotInitParams{
+		WorkerMetadata: workerMetadata,
+	})
 	pluginPilot, _ := client.pilot.(pilotPlugin)
 
 	if withBaseService, ok := config.RetryPolicy.(baseservice.WithBaseService); ok {

client_test.go

@@ -24,7 +24,6 @@ import (
 	"github.com/tidwall/sjson"
 
 	"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"
@@ -840,38 +839,47 @@ func Test_Client_Common(t *testing.T) {
 
 		_, bundle := setup(t)
 
-		type JobArgs struct {
-			testutil.JobArgsReflectKind[JobArgs]
-			hookEmbed[metadataHookInsertBegin]
-		}
+		var hookInsertBeginCalled atomic.Bool
 
-		AddWorker(bundle.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error {
+		jobArgs := JobArgsWithHooksFunc(func() []rivertype.Hook {
+			return []rivertype.Hook{
+				HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error {
+					hookInsertBeginCalled.Store(true)
+					return nil
+				}),
+			}
+		})
+
+		AddWorkerArgs(bundle.config.Workers, jobArgs, WorkFunc(func(ctx context.Context, job *Job[JobArgsWithHooksFunc]) error {
 			return nil
 		}))
 
 		client, err := NewClient(riverpgxv5.New(bundle.dbPool), bundle.config)
 		require.NoError(t, err)
 
-		insertRes, err := client.Insert(ctx, JobArgs{}, nil)
+		_, err = client.Insert(ctx, jobArgs, nil)
 		require.NoError(t, err)
 
-		var metadataMap map[string]any
-		err = json.Unmarshal(insertRes.Job.Metadata, &metadataMap)
-		require.NoError(t, err)
-		require.Equal(t, metadataHookCalled, metadataMap[metadataHookInsertBeginKey])
+		require.True(t, hookInsertBeginCalled.Load())
 	})
 
 	t.Run("WithWorkBeginHookOnJobArgs", func(t *testing.T) { //nolint:dupl
 		t.Parallel()
 
 		_, bundle := setup(t)
 
-		type JobArgs struct {
-			testutil.JobArgsReflectKind[JobArgs]
-			hookEmbed[metadataHookWorkBegin]
-		}
+		var hookWorkBeginCalled atomic.Bool
 
-		AddWorker(bundle.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error {
+		jobArgs := JobArgsWithHooksFunc(func() []rivertype.Hook {
+			return []rivertype.Hook{
+				HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error {
+					hookWorkBeginCalled.Store(true)
+					return nil
+				}),
+			}
+		})
+
+		AddWorkerArgs(bundle.config.Workers, jobArgs, WorkFunc(func(ctx context.Context, job *Job[JobArgsWithHooksFunc]) error {
 			return nil
 		}))
 
@@ -881,30 +889,33 @@ func Test_Client_Common(t *testing.T) {
 		subscribeChan := subscribe(t, client)
 		startClient(ctx, t, client)
 
-		insertRes, err := client.Insert(ctx, JobArgs{}, nil)
+		insertRes, err := client.Insert(ctx, jobArgs, nil)
 		require.NoError(t, err)
 
 		event := riversharedtest.WaitOrTimeout(t, subscribeChan)
 		require.Equal(t, EventKindJobCompleted, event.Kind)
 		require.Equal(t, insertRes.Job.ID, event.Job.ID)
 
-		var metadataMap map[string]any
-		err = json.Unmarshal(event.Job.Metadata, &metadataMap)
-		require.NoError(t, err)
-		require.Equal(t, metadataHookCalled, metadataMap[metadataHookWorkBeginKey])
+		require.True(t, hookWorkBeginCalled.Load())
 	})
 
 	t.Run("WithWorkEndHookOnJobArgs", func(t *testing.T) { //nolint:dupl
 		t.Parallel()
 
 		_, bundle := setup(t)
 
-		type JobArgs struct {
-			testutil.JobArgsReflectKind[JobArgs]
-			hookEmbed[metadataHookWorkEnd]
-		}
+		var hookWorkEndCalled atomic.Bool
 
-		AddWorker(bundle.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error {
+		jobArgs := JobArgsWithHooksFunc(func() []rivertype.Hook {
+			return []rivertype.Hook{
+				HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error {
+					hookWorkEndCalled.Store(true)
+					return nil
+				}),
+			}
+		})
+
+		AddWorkerArgs(bundle.config.Workers, jobArgs, WorkFunc(func(ctx context.Context, job *Job[JobArgsWithHooksFunc]) error {
 			return nil
 		}))
 
@@ -914,17 +925,14 @@ func Test_Client_Common(t *testing.T) {
 		subscribeChan := subscribe(t, client)
 		startClient(ctx, t, client)
 
-		insertRes, err := client.Insert(ctx, JobArgs{}, nil)
+		insertRes, err := client.Insert(ctx, jobArgs, nil)
 		require.NoError(t, err)
 
 		event := riversharedtest.WaitOrTimeout(t, subscribeChan)
 		require.Equal(t, EventKindJobCompleted, event.Kind)
 		require.Equal(t, insertRes.Job.ID, event.Job.ID)
 
-		var metadataMap map[string]any
-		err = json.Unmarshal(event.Job.Metadata, &metadataMap)
-		require.NoError(t, err)
-		require.Equal(t, metadataHookCalled, metadataMap[metadataHookWorkEndKey])
+		require.True(t, hookWorkEndCalled.Load())
 	})
 
 	t.Run("WithGlobalWorkerMiddleware", func(t *testing.T) {
@@ -1384,73 +1392,6 @@ func Test_Client_Common(t *testing.T) {
 	})
 }
 
-// hookEmbed can be embedded on a JobArgs to add a hook to it in such a way that
-// it can be encapsulated within a test case.
-type hookEmbed[T rivertype.Hook] struct{}
-
-func (f hookEmbed[T]) Hooks() []rivertype.Hook {
-	var hook T
-	return []rivertype.Hook{hook}
-}
-
-const (
-	metadataHookCalled         = "called"
-	metadataHookInsertBeginKey = "insert_begin"
-	metadataHookWorkBeginKey   = "work_begin"
-	metadataHookWorkEndKey     = "work_end"
-)
-
-type metadataHookInsertBegin struct{ rivertype.Hook }
-
-func (metadataHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error {
-	var metadataMap map[string]any
-	if err := json.Unmarshal(params.Metadata, &metadataMap); err != nil {
-		return err
-	}
-
-	metadataMap[metadataHookInsertBeginKey] = metadataHookCalled
-
-	var err error
-	params.Metadata, err = json.Marshal(metadataMap)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
-type metadataHookWorkBegin struct{ rivertype.Hook }
-
-func (metadataHookWorkBegin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error {
-	metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx)
-	if !hasMetadataUpdates {
-		panic("expected to be called from within job executor")
-	}
-
-	metadataUpdates[metadataHookWorkBeginKey] = metadataHookCalled
-
-	return nil
-}
-
-type metadataHookWorkEnd struct{ rivertype.Hook }
-
-func (metadataHookWorkEnd) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error {
-	metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx)
-	if !hasMetadataUpdates {
-		panic("expected to be called from within job executor")
-	}
-
-	metadataUpdates[metadataHookWorkEndKey] = metadataHookCalled
-
-	return err
-}
-
-var (
-	_ rivertype.HookInsertBegin = metadataHookInsertBegin{}
-	_ rivertype.HookWorkBegin   = metadataHookWorkBegin{}
-	_ rivertype.HookWorkEnd     = metadataHookWorkEnd{}
-)
-
 type workerWithMiddleware[T JobArgs] struct {
 	WorkerDefaults[T]
 
@@ -8036,3 +7977,15 @@ func NewDriverWithoutListenNotify(dbPool *pgxpool.Pool) *DriverWithoutListenNoti
 
 func (d *DriverWithoutListenNotify) SupportsListener() bool     { return false }
 func (d *DriverWithoutListenNotify) SupportsListenNotify() bool { return false }
+
+type JobArgsWithHooksFunc func() []rivertype.Hook
+
+func (JobArgsWithHooksFunc) Kind() string { return "job_args_with_hooks" }
+
+func (f JobArgsWithHooksFunc) Hooks() []rivertype.Hook {
+	return f()
+}
+
+func (JobArgsWithHooksFunc) MarshalJSON() ([]byte, error) { return []byte("{}"), nil }
+
+func (JobArgsWithHooksFunc) UnmarshalJSON([]byte) error { return nil }

rivershared/riverpilot/pilot.go

@@ -36,7 +36,7 @@ type Pilot interface {
 
 	JobSetStateIfRunningMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error)
 
-	PilotInit(archetype *baseservice.Archetype)
+	PilotInit(archetype *baseservice.Archetype, params *PilotInitParams)
 
 	// ProducerInit is called when a producer is started. It should return the ID
 	// of the new producer, a new state object that will be used to track the
@@ -50,6 +50,17 @@ type Pilot interface {
 	QueueMetadataChanged(ctx context.Context, exec riverdriver.Executor, params *QueueMetadataChangedParams) error
 }
 
+// PilotInitParams are parameters for initializing a pilot.
+//
+// API is not stable. DO NOT USE.
+type PilotInitParams struct {
+	// WorkerMetadata is metadata about registered workers as received from the
+	// client's worker bundle. Only available when a client will work jobs (i.e.
+	// has Workers configured), so while it's safe to assume the presence of
+	// this value in places like maintenance services, it's not in all contexts.
+	WorkerMetadata []*rivertype.WorkerMetadata
+}
+
 // PilotPeriodicJob contains pilot functions related to periodic jobs. This is
 // extracted as its own interface so there's less surface area to mock in places
 // like the periodic job enqueuer where that's needed.

rivershared/riverpilot/standard.go

@@ -48,7 +48,7 @@ func (p *StandardPilot) PeriodicJobUpsertMany(ctx context.Context, exec riverdri
 	return nil, nil
 }
 
-func (p *StandardPilot) PilotInit(archetype *baseservice.Archetype) {
+func (p *StandardPilot) PilotInit(archetype *baseservice.Archetype, params *PilotInitParams) {
 	// No-op
 }
 

rivertype/river_type.go

@@ -494,3 +494,13 @@ func UniqueOptsByStateDefault() []JobState {
 		JobStateScheduled,
 	}
 }
+
+// WorkerMetadata is metadata about workers registerd with a client.
+type WorkerMetadata struct {
+	// JobArgHooks are job args specific hooks returned from a JobArgsWithHooks
+	// implementation.
+	JobArgHooks []Hook
+
+	// Kind is the kind returned from job args and recognized by worker to work.
+	Kind string
+}

worker.go

@@ -105,6 +105,17 @@ func AddWorker[T JobArgs](workers *Workers, worker Worker[T]) {
 	}
 }
 
+// AddWorkerArgs is the same as AddWorker except that it lets args be passed
+// explicitly rather than being instantiated implicitly. We don't know of any
+// use for this function beyond exercising some args-related edge cases in tests
+// are are difficult/impossible to exercise otherwise, and its use should be
+// considered internal only.
+func AddWorkerArgs[T JobArgs](workers *Workers, jobArgs T, worker Worker[T]) {
+	if err := workers.add(jobArgs, &workUnitFactoryWrapper[T]{worker: worker}); err != nil {
+		panic(err)
+	}
+}
+
 // AddWorkerSafely registers a worker on the provided Workers bundle. Unlike AddWorker,
 // AddWorkerSafely does not panic and instead returns an error if the worker
 // is already registered or if its configuration is invalid.

worker_test.go

@@ -45,6 +45,16 @@ func TestWork(t *testing.T) {
 			return nil
 		}))
 	})
+
+	type OtherJobArgs struct {
+		testutil.JobArgsReflectKind[OtherJobArgs]
+	}
+
+	var jobArgs OtherJobArgs
+	AddWorkerArgs(workers, jobArgs, WorkFunc(func(ctx context.Context, job *Job[OtherJobArgs]) error {
+		return nil
+	}))
+	require.Contains(t, workers.workersMap, (OtherJobArgs{}).Kind())
 }
 
 type configurableArgs struct {