diff --git a/client.go b/client.go index 814712a7..0c68b500 100644 --- a/client.go +++ b/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 { diff --git a/client_test.go b/client_test.go index 74d2add0..5fbc4198 100644 --- a/client_test.go +++ b/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,25 +839,28 @@ 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 @@ -866,12 +868,18 @@ func Test_Client_Common(t *testing.T) { _, 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,17 +889,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[metadataHookWorkBeginKey]) + require.True(t, hookWorkBeginCalled.Load()) }) t.Run("WithWorkEndHookOnJobArgs", func(t *testing.T) { //nolint:dupl @@ -899,12 +904,18 @@ func Test_Client_Common(t *testing.T) { _, 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 } diff --git a/rivershared/riverpilot/pilot.go b/rivershared/riverpilot/pilot.go index c2edc97c..98e1a217 100644 --- a/rivershared/riverpilot/pilot.go +++ b/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. diff --git a/rivershared/riverpilot/standard.go b/rivershared/riverpilot/standard.go index 3db36b8e..88c3acf7 100644 --- a/rivershared/riverpilot/standard.go +++ b/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 } diff --git a/rivertype/river_type.go b/rivertype/river_type.go index 0519ac65..79f9852b 100644 --- a/rivertype/river_type.go +++ b/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 +} diff --git a/worker.go b/worker.go index ff33a832..22f83b21 100644 --- a/worker.go +++ b/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. diff --git a/worker_test.go b/worker_test.go index da54f438..8654b21b 100644 --- a/worker_test.go +++ b/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 {