Skip to content

Inject worker metadata to pilot + testable job args #996

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
brandur merged 1 commit into from
Aug 3, 2025
Merged

Inject worker metadata to pilot + testable job args #996

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
149 changes: 51 additions & 98 deletions client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 {

@brandur brandur Aug 2, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See here and the red below for how AddWorkerArgs helps to remove boilerplate and bring implementation closer to the specific test case, but the more critical piece is that unless a hook has some "hack" workaround available like being able to add information to metadata (like this test case was doing), I couldn't find any other way to test hooks that are specific to only some job args.

The trouble is that AddWorker always implicitly constructs a JobArgs like:

func AddWorkerSafely[T JobArgs](workers *Workers, worker Worker[T]) error {
	var jobArgs T
	return workers.add(jobArgs, &workUnitFactoryWrapper[T]{worker: worker})
}

Which makes attaching a specific Hooks() implementation to it for a single test case impossible.

By bringing in something like AddWorkerArgs which is implemented like:

func AddWorkerArgs[T JobArgs](workers *Workers, jobArgs T, worker Worker[T]) {
	if err := workers.add(jobArgs, &workUnitFactoryWrapper[T]{worker: worker}); err != nil {
		panic(err)
	}
}

It lets us do what you see above, with a closure hook implementation that can set a local variable and be tested against.

The only other way I could think to do it is to have a job args specific to every test case and which sets a global variable, which sucks. So the only options I'm able to come up with are:

  1. AddWorkerArgs as you see here.
  2. Job args for each test case that set a global (or use a workaround like middleware if that's available).
  3. Don't test the hooks on specific job args case.

@brandur brandur Aug 2, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I should note also that I'm not married to this at all, but for the life of me could not find a better way.

@bgentry bgentry Aug 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't have a better idea unfortunately, annoying problem to have to hack around 😕

@brandur brandur Aug 3, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

K thx. And yeah :( We'll have a bit of time ahead of release probably. I'll keep mulling it over.

It'd be nice if Go allowed you to mark some API is "pretty internal, you probably don't want to use this man" or something of that nature even and at least hide it from Godoc or something. Couldn't think of any tricks here though unfortunately.

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
}))

Expand All @@ -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
}))

Expand All @@ -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) {
Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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 }
13 changes: 12 additions & 1 deletion rivershared/riverpilot/pilot.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion rivershared/riverpilot/standard.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
10 changes: 10 additions & 0 deletions rivertype/river_type.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
11 changes: 11 additions & 0 deletions worker.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions worker_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading