From a7f61935a41d9c09526852d47b059e2d4cbad958 Mon Sep 17 00:00:00 2001 From: flcrom <161213233+flcrom@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:08:09 +0530 Subject: [PATCH 1/3] fix: leave worker registry unchanged on alias collision --- worker.go | 29 +++++++++++++++++++---------- worker_test.go | 20 ++++++++++++++++++++ 2 files changed, 39 insertions(+), 10 deletions(-) diff --git a/worker.go b/worker.go index f833b760d..a9d1370cb 100644 --- a/worker.go +++ b/worker.go @@ -174,20 +174,29 @@ func (w Workers) add(jobArgs JobArgs, workUnitFactory workunit.WorkUnitFactory) workUnitFactory: workUnitFactory, } - kind := jobArgs.Kind() - if err := checkRegistered(kind); err != nil { - return err - } - w.workersMap[kind] = workerInfo + // Collect the primary kind and any aliases first so that all of them can + // be validated before anything is written to the map. This way, a failed + // registration leaves no partial state behind. + kinds := []string{jobArgs.Kind()} // Jobs can register an alternate kind to make renaming easier. if jobArgsWithKindAliases, ok := jobArgs.(JobArgsWithKindAliases); ok { - for _, kind := range jobArgsWithKindAliases.KindAliases() { - if err := checkRegistered(kind); err != nil { - return err - } - w.workersMap[kind] = workerInfo + kinds = append(kinds, jobArgsWithKindAliases.KindAliases()...) + } + + seen := make(map[string]struct{}, len(kinds)) + for _, kind := range kinds { + if err := checkRegistered(kind); err != nil { + return err } + if _, ok := seen[kind]; ok { + return fmt.Errorf("worker for kind %q is already registered", kind) + } + seen[kind] = struct{}{} + } + + for _, kind := range kinds { + w.workersMap[kind] = workerInfo } return nil diff --git a/worker_test.go b/worker_test.go index 8654b21ba..5fa58121e 100644 --- a/worker_test.go +++ b/worker_test.go @@ -13,6 +13,19 @@ import ( "github.com/riverqueue/river/rivershared/util/testutil" ) +func TestAddWorkerSafelyAliasCollisionDoesNotRegisterKinds(t *testing.T) { + t.Parallel() + + workers := NewWorkers() + + require.NoError(t, AddWorkerSafely(workers, WorkFunc(func(context.Context, *Job[withKindAliasesArgs]) error { return nil }))) + before := len(workers.workersMap) + require.Error(t, AddWorkerSafely(workers, WorkFunc(func(context.Context, *Job[lateAliasCollisionArgs]) error { return nil }))) + require.Len(t, workers.workersMap, before) + require.NotContains(t, workers.workersMap, "candidate_primary") + require.NotContains(t, workers.workersMap, "candidate_early_alias") +} + func TestWork(t *testing.T) { t.Parallel() @@ -74,6 +87,13 @@ func (w *configurableWorker) Work(ctx context.Context, job *Job[configurableArgs return nil } +type lateAliasCollisionArgs struct{} + +func (lateAliasCollisionArgs) Kind() string { return "candidate_primary" } +func (lateAliasCollisionArgs) KindAliases() []string { + return []string{"candidate_early_alias", "with_kind_alternate_alternate"} +} + type withKindAliasesArgs struct{} func (a withKindAliasesArgs) Kind() string { return "with_kind_alternate" } From 3f38b47e71939feb73ba7448956beaf6ab1648c9 Mon Sep 17 00:00:00 2001 From: flcrom <161213233+flcrom@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:36:31 +0530 Subject: [PATCH 2/3] docs: note atomic worker registration fix in changelog --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 05277e667..9649508b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440). + ## [0.49.0] - 2026-10-05 ### Changed From 804c3143e77934f36bc6c2b807a6d9a8da26f347 Mon Sep 17 00:00:00 2001 From: Brandur Date: Fri, 9 Oct 2026 23:45:55 -0500 Subject: [PATCH 3/3] Tweak new implementation to be a little more succinct Produces a net reduction in LOCs (-15 lines) for `Workers.add`. --- worker.go | 37 +++++++++++-------------------------- 1 file changed, 11 insertions(+), 26 deletions(-) diff --git a/worker.go b/worker.go index a9d1370cb..0c1ede252 100644 --- a/worker.go +++ b/worker.go @@ -162,41 +162,26 @@ func NewWorkers() *Workers { } func (w Workers) add(jobArgs JobArgs, workUnitFactory workunit.WorkUnitFactory) error { - checkRegistered := func(kind string) error { - if _, ok := w.workersMap[kind]; ok { - return fmt.Errorf("worker for kind %q is already registered", kind) - } - return nil - } - - workerInfo := workerInfo{ - jobArgs: jobArgs, - workUnitFactory: workUnitFactory, - } - - // Collect the primary kind and any aliases first so that all of them can - // be validated before anything is written to the map. This way, a failed - // registration leaves no partial state behind. kinds := []string{jobArgs.Kind()} - - // Jobs can register an alternate kind to make renaming easier. - if jobArgsWithKindAliases, ok := jobArgs.(JobArgsWithKindAliases); ok { - kinds = append(kinds, jobArgsWithKindAliases.KindAliases()...) + if args, ok := jobArgs.(JobArgsWithKindAliases); ok { + kinds = append(kinds, args.KindAliases()...) } - seen := make(map[string]struct{}, len(kinds)) + // Validate all kinds before changing the registry. + seen := make(map[string]bool, len(kinds)) for _, kind := range kinds { - if err := checkRegistered(kind); err != nil { - return err - } - if _, ok := seen[kind]; ok { + if _, ok := w.workersMap[kind]; ok || seen[kind] { return fmt.Errorf("worker for kind %q is already registered", kind) } - seen[kind] = struct{}{} + seen[kind] = true } + info := workerInfo{ + jobArgs: jobArgs, + workUnitFactory: workUnitFactory, + } for _, kind := range kinds { - w.workersMap[kind] = workerInfo + w.workersMap[kind] = info } return nil