From a470e472450264ddc3941611f457c270d0a816da Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Wed, 29 Jul 2026 20:04:58 -0500 Subject: [PATCH] make producer race test deterministic The producer concurrency regression processes 10,000 jobs and relies on finishing the last one within a fixed deadline. Race-enabled CI can exercise the intended overlap many times and still fail solely because it misses that throughput deadline. Replace the workload with a coordinated three-job scenario that blocks fetches after capacity is calculated. Complete jobs while fetches are pending and read capacity concurrently across removals and additions. This keeps unsafe active-job accounting visible to the race detector. Assert each capacity transition and verify that a queued third job runs after a slot reopens. Run the scenario through both poll-only and notifier producer setups, using the pilot seam so no production-only test signal is needed. Keep reusable test plumbing generic with a pre-fetch pilot hook while the scenario owns its fetch schedule. Use zero-based job indices and single-owner gate cleanup. Stream capacity transitions in order, with cancellation-aware sends, so failures cannot strand the observer. Document the gated sequence so completion during a pending fetch, the retained fetch limit, refill timing, and observer lifetime are explicit. --- producer_test.go | 317 ++++++++++++++++++++++++++++------------------- 1 file changed, 190 insertions(+), 127 deletions(-) diff --git a/producer_test.go b/producer_test.go index a23262f1..1f0629c3 100644 --- a/producer_test.go +++ b/producer_test.go @@ -5,13 +5,13 @@ import ( "encoding/json" "fmt" "slices" + "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" "github.com/riverqueue/river/internal/jobcompleter" - "github.com/riverqueue/river/internal/maintenance" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" @@ -33,6 +33,27 @@ import ( const testClientID = "test-client-id" +// beforeJobGetAvailablePilot calls a hook before delegating JobGetAvailable to +// the wrapped pilot. +type beforeJobGetAvailablePilot struct { + riverpilot.Pilot + + beforeJobGetAvailableFunc func(params *riverdriver.JobGetAvailableParams) +} + +func (p *beforeJobGetAvailablePilot) JobGetAvailable( + ctx context.Context, + exec riverdriver.Executor, + state riverpilot.ProducerState, + params *riverdriver.JobGetAvailableParams, +) ([]*rivertype.JobRow, error) { + if p.beforeJobGetAvailableFunc != nil { + p.beforeJobGetAvailableFunc(params) + } + + return p.Pilot.JobGetAvailable(ctx, exec, state, params) +} + type countingPluginLookup struct { pluginlookup.PluginLookupInterface @@ -44,132 +65,6 @@ func (l *countingPluginLookup) ByKind(kind pluginlookup.PluginKind) []any { return l.PluginLookupInterface.ByKind(kind) } -func Test_Producer_CanSafelyCompleteJobsWhileFetchingNewOnes(t *testing.T) { - // We have encountered previous data races with the list of active jobs on - // Producer because we need to know the count of active jobs in order to - // determine how many we can fetch for the next batch, while we're managing - // the map of active jobs in a different goroutine. - // - // This test attempts to exercise that race condition so that the race - // detector can tell us if we're protected against it. - t.Parallel() - - ctx := context.Background() - require := require.New(t) - dbPool := riversharedtest.DBPool(ctx, t) - - const maxJobCount = 10000 - // This doesn't strictly mean that there are no more jobs left to process, - // merely that the final job we inserted is now being processed, which is - // close enough for our purposes here. - lastJobRun := make(chan struct{}) - - archetype := riversharedtest.BaseServiceArchetype(t) - - var ( - driver = riverpgxv5.New(dbPool) - schema = riverdbtest.TestSchema(ctx, t, driver, nil) - config = newTestConfig(t, schema) - exec = driver.GetExecutor() - listener = driver.GetListener(&riverdriver.GetListenenerParams{Schema: schema}) - pilot = &riverpilot.StandardPilot{} - ) - - subscribeChan := make(chan []jobcompleter.CompleterJobUpdated, 100) - t.Cleanup(riverinternaltest.DiscardContinuously(subscribeChan)) - - completer := jobcompleter.NewInlineCompleter(archetype, schema, exec, &riverpilot.StandardPilot{}, subscribeChan) - t.Cleanup(completer.Stop) - - type WithJobNumArgs struct { - testutil.JobArgsReflectKind[WithJobNumArgs] - - JobNum int `json:"job_num"` - } - - workers := NewWorkers() - AddWorker(workers, WorkFunc(func(ctx context.Context, job *Job[WithJobNumArgs]) error { - var jobArgs WithJobNumArgs - require.NoError(json.Unmarshal(job.EncodedArgs, &jobArgs)) - - if jobArgs.JobNum == maxJobCount-1 { - select { - case <-ctx.Done(): - case lastJobRun <- struct{}{}: - } - } - return nil - })) - - notifier := notifier.New(archetype, listener) - - producer := newProducer(archetype, exec, pilot, &producerConfig{ - ClientID: testClientID, - Completer: completer, - ErrorHandler: newTestErrorHandler(), - // Fetch constantly to more aggressively trigger the potential data race: - FetchCooldown: time.Millisecond, - FetchPollInterval: time.Millisecond, - PluginLookupByJob: pluginlookup.NewJobPluginLookup(), - PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), - JobTimeout: JobTimeoutDefault, - MaxWorkers: 1000, - Notifier: notifier, - Queue: rivercommon.QueueDefault, - QueuePollInterval: queuePollIntervalDefault, - QueueReportInterval: queueReportIntervalDefault, - RetryPolicy: &DefaultClientRetryPolicy{}, - SchedulerInterval: maintenance.JobSchedulerIntervalDefault, - Schema: schema, - StaleProducerRetentionPeriod: time.Minute, - Workers: workers, - }) - - params := make([]*riverdriver.JobInsertFastParams, maxJobCount) - for i := range params { - insertParams, err := insertParamsFromConfigArgsAndOptions(archetype, config, WithJobNumArgs{JobNum: i}, nil) - require.NoError(err) - - params[i] = (*riverdriver.JobInsertFastParams)(insertParams) - } - - ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) - t.Cleanup(cancel) - - go func() { - // The producer should never exceed its MaxWorkerCount. If it does, panic so - // we can get a trace. - for { - select { - case <-ctx.Done(): - return - default: - } - numActiveJobs := producer.numJobsActive.Load() - if int(numActiveJobs) > producer.config.MaxWorkers { - panic(fmt.Sprintf("producer exceeded MaxWorkerCount=%d, actual count=%d", producer.config.MaxWorkers, numActiveJobs)) - } - } - }() - - _, err := exec.JobInsertFastMany(ctx, &riverdriver.JobInsertFastManyParams{ - Jobs: params, - Schema: producer.config.Schema, - }) - require.NoError(err) - - require.NoError(producer.StartWorkContext(ctx, ctx)) - t.Cleanup(producer.Stop) - - select { - case <-lastJobRun: - t.Logf("Last job reported in; cancelling context") - cancel() - case <-ctx.Done(): - t.Error("timed out waiting for last job to run") - } -} - func TestProducer_MetricEmitHook(t *testing.T) { t.Parallel() @@ -608,6 +503,174 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRetryable, update.Job.State) }) + t.Run("CompletesJobWhileFetchingNewOnes", func(t *testing.T) { + t.Parallel() + + // Exercise the case where a job finishes while the producer is fetching + // more jobs. One of two worker slots is occupied, so the next fetch + // computes a limit of one before the test holds it pending. + // Completing the active job must be processed without waiting for the + // fetch, reopening both slots before the fetched job occupies one. + // A third queued job verifies that the producer continues fetching and + // runs the remaining work. + // + // The test also reads capacity throughout these active-job changes, + // exercising the shared-state boundary so the race detector can catch + // unsynchronized access. + + producer, bundle := setup(t) + producer.config.FetchCooldown = time.Millisecond + producer.config.FetchPollInterval = time.Hour + producer.config.MaxWorkers = 2 + + const numJobs = 3 + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + + JobIndex int `json:"job_index"` + } + + capacityReadCtx, stopCapacityRead := context.WithCancel(ctx) + + var ( + capacityChanged = make(chan int) + capacityReadDone = make(chan struct{}) + fetchCount atomic.Int32 + jobStarted = make(chan int, numJobs) + releaseJob = make([]chan struct{}, numJobs) + secondFetchStarted = make(chan int, 1) + thirdFetchStarted = make(chan int, 1) + unblockSecondFetch = make(chan struct{}, 1) + unblockThirdFetch = make(chan struct{}, 1) + ) + + for jobIndex := range numJobs { + releaseJob[jobIndex] = make(chan struct{}, 1) + } + + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + jobIndex := job.Args.JobIndex + if jobIndex < 0 || jobIndex >= numJobs { + return fmt.Errorf("unexpected job index: %d", jobIndex) + } + + jobStarted <- jobIndex + <-releaseJob[jobIndex] + return nil + })) + + producer.pilot = &beforeJobGetAvailablePilot{ + Pilot: producer.pilot, + beforeJobGetAvailableFunc: func(params *riverdriver.JobGetAvailableParams) { + switch fetchCount.Add(1) { + case 2: + secondFetchStarted <- params.MaxToLock + <-unblockSecondFetch + case 3: + thirdFetchStarted <- params.MaxToLock + <-unblockThirdFetch + } + }, + } + + mustInsert(ctx, t, producer, bundle, &JobArgs{JobIndex: 0}) + + startProducer(t, ctx, ctx, producer) + t.Cleanup(func() { + for _, release := range releaseJob { + close(release) + } + close(unblockSecondFetch) + close(unblockThirdFetch) + stopCapacityRead() + }) + + require.Equal(t, 0, riversharedtest.WaitOrTimeout(t, jobStarted)) + + mustInsert(ctx, t, producer, bundle, &JobArgs{JobIndex: 1}) + mustInsert(ctx, t, producer, bundle, &JobArgs{JobIndex: 2}) + producer.TriggerJobFetch() + + secondFetchLimit := riversharedtest.WaitOrTimeout(t, secondFetchStarted) + require.Equal(t, 1, secondFetchLimit) + + go func() { + defer close(capacityReadDone) + + lastCapacity := -1 + for { + capacity := producer.maxJobsToFetch() + if capacity != lastCapacity { + select { + case capacityChanged <- capacity: + lastCapacity = capacity + case <-capacityReadCtx.Done(): + return + } + } + + select { + case <-capacityReadCtx.Done(): + return + default: + } + } + }() + t.Cleanup(func() { + stopCapacityRead() + riversharedtest.WaitOrTimeout(t, capacityReadDone) + }) + + // Keep reading maxJobsToFetch between transitions so reads can overlap + // active-job changes. Wait for each observed transition before causing + // the next. + waitForCapacity := func(expected int) { + t.Helper() + + require.Equal(t, expected, riversharedtest.WaitOrTimeout(t, capacityChanged)) + require.Equal(t, producer.config.MaxWorkers-expected, int(producer.numJobsActive.Load())) + require.Equal(t, expected, producer.maxJobsToFetch()) + } + + // The second fetch is blocked with a limit of one while job 0 occupies a + // slot. Complete job 0 and verify that the producer processes its result, + // reopening both slots before the fetch resumes. + waitForCapacity(1) + releaseJob[0] <- struct{}{} + + waitForCapacity(2) + unblockSecondFetch <- struct{}{} + + // The fetch keeps its original limit of one, so it starts only job 1 and + // leaves job 2 queued. The third fetch may calculate its limit while job 1 + // is still active (one slot) or after its completion is processed (two + // slots). + require.Equal(t, 1, riversharedtest.WaitOrTimeout(t, jobStarted)) + waitForCapacity(1) + + releaseJob[1] <- struct{}{} + waitForCapacity(2) + + thirdFetchLimit := riversharedtest.WaitOrTimeout(t, thirdFetchStarted) + require.Contains(t, []int{1, 2}, thirdFetchLimit) + unblockThirdFetch <- struct{}{} + + // After job 2 starts, wait for the reader to observe the final transition + // back to capacity one. Stop it before releasing job 2, then require all + // three jobs to complete. + require.Equal(t, 2, riversharedtest.WaitOrTimeout(t, jobStarted)) + waitForCapacity(1) + stopCapacityRead() + riversharedtest.WaitOrTimeout(t, capacityReadDone) + + releaseJob[2] <- struct{}{} + updates := riversharedtest.WaitOrTimeoutN(t, bundle.jobUpdates, numJobs) + for _, update := range updates { + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + } + }) + t.Run("MaxWorkers", func(t *testing.T) { t.Parallel()