From 1e26df678b05c444df9afd8d082d57f497b428c8 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Mon, 17 Feb 2025 12:20:19 -0600 Subject: [PATCH 1/5] allow setting created_at on insert for clock overrides While the client code was attempting to set `created_at` when the clock was stubbed, the driver code was not actually utilizing this at all. This commit makes it possible to set the `created_at` (similar to `scheduled_at`) using a mocked clock. It also removes the ability to set `finalized_at`, which has no need to be set at insert (other than for tests utilizing `JobInsertFull`). --- .../riverdrivertest/riverdrivertest.go | 39 +++++++++++++-- .../internal/dbsqlc/river_job.sql.go | 48 +++++++++++-------- .../river_database_sql_driver.go | 14 ++++++ .../riverpgxv5/internal/dbsqlc/copyfrom.go | 4 +- .../riverpgxv5/internal/dbsqlc/river_job.sql | 4 ++ .../internal/dbsqlc/river_job.sql.go | 48 +++++++++++-------- .../internal/dbsqlc/river_job_copyfrom.sql | 4 +- .../internal/dbsqlc/river_job_copyfrom.sql.go | 2 +- riverdriver/riverpgxv5/river_pgx_v5_driver.go | 13 +++++ 9 files changed, 128 insertions(+), 48 deletions(-) diff --git a/internal/riverinternaltest/riverdrivertest/riverdrivertest.go b/internal/riverinternaltest/riverdrivertest/riverdrivertest.go index 569bef8f..aea7a295 100644 --- a/internal/riverinternaltest/riverdrivertest/riverdrivertest.go +++ b/internal/riverinternaltest/riverdrivertest/riverdrivertest.go @@ -877,6 +877,7 @@ func Exercise[TTx any](ctx context.Context, t *testing.T, insertParams := make([]*riverdriver.JobInsertFastParams, 10) for i := 0; i < len(insertParams); i++ { insertParams[i] = &riverdriver.JobInsertFastParams{ + CreatedAt: ptrutil.Ptr(now.Add(time.Duration(i) * 5 * time.Second)), EncodedArgs: []byte(`{"encoded": "args"}`), Kind: "test_kind", MaxAttempts: rivercommon.MaxAttemptsDefault, @@ -901,7 +902,7 @@ func Exercise[TTx any](ctx context.Context, t *testing.T, require.Equal(t, 0, job.Attempt) require.Nil(t, job.AttemptedAt) require.Empty(t, job.AttemptedBy) - require.WithinDuration(t, now, job.CreatedAt, 2*time.Second) + require.WithinDuration(t, now.Add(time.Duration(i)*5*time.Second), job.CreatedAt, time.Millisecond) require.JSONEq(t, `{"encoded": "args"}`, string(job.EncodedArgs)) require.Empty(t, job.Errors) require.Nil(t, job.FinalizedAt) @@ -974,6 +975,7 @@ func Exercise[TTx any](ctx context.Context, t *testing.T, insertParams := make([]*riverdriver.JobInsertFastParams, 10) for i := 0; i < len(insertParams); i++ { insertParams[i] = &riverdriver.JobInsertFastParams{ + CreatedAt: ptrutil.Ptr(now.Add(time.Duration(i) * 5 * time.Second)), EncodedArgs: []byte(`{"encoded": "args"}`), Kind: "test_kind", MaxAttempts: rivercommon.MaxAttemptsDefault, @@ -994,10 +996,10 @@ func Exercise[TTx any](ctx context.Context, t *testing.T, jobsAfter, err := exec.JobGetByKindMany(ctx, []string{"test_kind"}) require.NoError(t, err) require.Len(t, jobsAfter, len(insertParams)) - for _, job := range jobsAfter { + for i, job := range jobsAfter { require.Equal(t, 0, job.Attempt) require.Nil(t, job.AttemptedAt) - require.WithinDuration(t, time.Now().UTC(), job.CreatedAt, 2*time.Second) + require.WithinDuration(t, now.Add(time.Duration(i)*5*time.Second), job.CreatedAt, time.Millisecond) require.JSONEq(t, `{"encoded": "args"}`, string(job.EncodedArgs)) require.Empty(t, job.Errors) require.Nil(t, job.FinalizedAt) @@ -1012,6 +1014,37 @@ func Exercise[TTx any](ctx context.Context, t *testing.T, } }) + t.Run("MissingCreatedAtDefaultsToNow", func(t *testing.T) { + exec, _ := setup(ctx, t) + + insertParams := make([]*riverdriver.JobInsertFastParams, 10) + for i := 0; i < len(insertParams); i++ { + insertParams[i] = &riverdriver.JobInsertFastParams{ + CreatedAt: nil, // explicit nil + EncodedArgs: []byte(`{"encoded": "args"}`), + Kind: "test_kind", + MaxAttempts: rivercommon.MaxAttemptsDefault, + Metadata: []byte(`{"meta": "data"}`), + Priority: rivercommon.PriorityDefault, + Queue: rivercommon.QueueDefault, + ScheduledAt: ptrutil.Ptr(time.Now().UTC()), + State: rivertype.JobStateAvailable, + Tags: []string{"tag"}, + } + } + + count, err := exec.JobInsertFastManyNoReturning(ctx, insertParams) + require.NoError(t, err) + require.Len(t, insertParams, count) + + jobsAfter, err := exec.JobGetByKindMany(ctx, []string{"test_kind"}) + require.NoError(t, err) + require.Len(t, jobsAfter, len(insertParams)) + for _, job := range jobsAfter { + require.WithinDuration(t, time.Now().UTC(), job.CreatedAt, 2*time.Second) + } + }) + t.Run("MissingScheduledAtDefaultsToNow", func(t *testing.T) { exec, _ := setup(ctx, t) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 14f676e2..59b6d1f7 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -528,6 +528,7 @@ func (q *Queries) JobGetStuck(ctx context.Context, db DBTX, arg *JobGetStuckPara const jobInsertFastMany = `-- name: JobInsertFastMany :many INSERT INTO river_job( args, + created_at, kind, max_attempts, metadata, @@ -540,22 +541,23 @@ INSERT INTO river_job( unique_states ) SELECT unnest($1::jsonb[]), - unnest($2::text[]), - unnest($3::smallint[]), - unnest($4::jsonb[]), - unnest($5::smallint[]), - unnest($6::text[]), - unnest($7::timestamptz[]), + unnest($2::timestamptz[]), + unnest($3::text[]), + unnest($4::smallint[]), + unnest($5::jsonb[]), + unnest($6::smallint[]), + unnest($7::text[]), + unnest($8::timestamptz[]), -- To avoid requiring pgx users to register the OID of the river_job_state[] -- type, we cast the array to text[] and then to river_job_state. - unnest($8::text[])::river_job_state, + unnest($9::text[])::river_job_state, -- Unnest on a multi-dimensional array will fully flatten the array, so we -- encode the tag list as a comma-separated string and split it in the -- query. - string_to_array(unnest($9::text[]), ','), + string_to_array(unnest($10::text[]), ','), - unnest($10::bytea[]), - unnest($11::bit(8)[]) + unnest($11::bytea[]), + unnest($12::bit(8)[]) ON CONFLICT (unique_key) WHERE unique_key IS NOT NULL @@ -568,6 +570,7 @@ RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_a type JobInsertFastManyParams struct { Args []string + CreatedAt []time.Time Kind []string MaxAttempts []int16 Metadata []string @@ -588,6 +591,7 @@ type JobInsertFastManyRow struct { func (q *Queries) JobInsertFastMany(ctx context.Context, db DBTX, arg *JobInsertFastManyParams) ([]*JobInsertFastManyRow, error) { rows, err := db.QueryContext(ctx, jobInsertFastMany, pq.Array(arg.Args), + pq.Array(arg.CreatedAt), pq.Array(arg.Kind), pq.Array(arg.MaxAttempts), pq.Array(arg.Metadata), @@ -643,6 +647,7 @@ func (q *Queries) JobInsertFastMany(ctx context.Context, db DBTX, arg *JobInsert const jobInsertFastManyNoReturning = `-- name: JobInsertFastManyNoReturning :execrows INSERT INTO river_job( args, + created_at, kind, max_attempts, metadata, @@ -655,22 +660,23 @@ INSERT INTO river_job( unique_states ) SELECT unnest($1::jsonb[]), - unnest($2::text[]), - unnest($3::smallint[]), - unnest($4::jsonb[]), - unnest($5::smallint[]), - unnest($6::text[]), - unnest($7::timestamptz[]), - unnest($8::river_job_state[]), + unnest($2::timestamptz[]), + unnest($3::text[]), + unnest($4::smallint[]), + unnest($5::jsonb[]), + unnest($6::smallint[]), + unnest($7::text[]), + unnest($8::timestamptz[]), + unnest($9::river_job_state[]), -- lib/pq really, REALLY does not play nicely with multi-dimensional arrays, -- so instead we pack each set of tags into a string, send them through, -- then unpack them here into an array to put in each row. This isn't -- necessary in the Pgx driver where copyfrom is used instead. - string_to_array(unnest($9::text[]), ','), + string_to_array(unnest($10::text[]), ','), - unnest($10::bytea[]), - unnest($11::bit(8)[]) + unnest($11::bytea[]), + unnest($12::bit(8)[]) ON CONFLICT (unique_key) WHERE unique_key IS NOT NULL @@ -681,6 +687,7 @@ DO NOTHING type JobInsertFastManyNoReturningParams struct { Args []string + CreatedAt []time.Time Kind []string MaxAttempts []int16 Metadata []string @@ -696,6 +703,7 @@ type JobInsertFastManyNoReturningParams struct { func (q *Queries) JobInsertFastManyNoReturning(ctx context.Context, db DBTX, arg *JobInsertFastManyNoReturningParams) (int64, error) { result, err := db.ExecContext(ctx, jobInsertFastManyNoReturning, pq.Array(arg.Args), + pq.Array(arg.CreatedAt), pq.Array(arg.Kind), pq.Array(arg.MaxAttempts), pq.Array(arg.Metadata), diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index 1677bdb8..ceacaaf5 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -204,6 +204,7 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver.JobInsertFastParams) ([]*riverdriver.JobInsertFastResult, error) { insertJobsParams := &dbsqlc.JobInsertFastManyParams{ Args: make([]string, len(params)), + CreatedAt: make([]time.Time, len(params)), Kind: make([]string, len(params)), MaxAttempts: make([]int16, len(params)), Metadata: make([]string, len(params)), @@ -220,6 +221,11 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver. for i := 0; i < len(params); i++ { params := params[i] + createdAt := now + if params.CreatedAt != nil { + createdAt = *params.CreatedAt + } + scheduledAt := now if params.ScheduledAt != nil { scheduledAt = *params.ScheduledAt @@ -233,6 +239,7 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver. defaultObject := "{}" insertJobsParams.Args[i] = valutil.ValOrDefault(string(params.EncodedArgs), defaultObject) + insertJobsParams.CreatedAt[i] = createdAt insertJobsParams.Kind[i] = params.Kind insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec insertJobsParams.Metadata[i] = valutil.ValOrDefault(string(params.Metadata), defaultObject) @@ -262,6 +269,7 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver. func (e *Executor) JobInsertFastManyNoReturning(ctx context.Context, params []*riverdriver.JobInsertFastParams) (int, error) { insertJobsParams := &dbsqlc.JobInsertFastManyNoReturningParams{ Args: make([]string, len(params)), + CreatedAt: make([]time.Time, len(params)), Kind: make([]string, len(params)), MaxAttempts: make([]int16, len(params)), Metadata: make([]string, len(params)), @@ -278,6 +286,11 @@ func (e *Executor) JobInsertFastManyNoReturning(ctx context.Context, params []*r for i := 0; i < len(params); i++ { params := params[i] + createdAt := now + if params.CreatedAt != nil { + createdAt = *params.CreatedAt + } + scheduledAt := now if params.ScheduledAt != nil { scheduledAt = *params.ScheduledAt @@ -291,6 +304,7 @@ func (e *Executor) JobInsertFastManyNoReturning(ctx context.Context, params []*r defaultObject := "{}" insertJobsParams.Args[i] = valutil.ValOrDefault(string(params.EncodedArgs), defaultObject) + insertJobsParams.CreatedAt[i] = createdAt insertJobsParams.Kind[i] = params.Kind insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec insertJobsParams.Metadata[i] = valutil.ValOrDefault(string(params.Metadata), defaultObject) diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go b/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go index 8cedc3f3..724187a1 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/copyfrom.go @@ -30,7 +30,7 @@ func (r *iteratorForJobInsertFastManyCopyFrom) Next() bool { func (r iteratorForJobInsertFastManyCopyFrom) Values() ([]interface{}, error) { return []interface{}{ r.rows[0].Args, - r.rows[0].FinalizedAt, + r.rows[0].CreatedAt, r.rows[0].Kind, r.rows[0].MaxAttempts, r.rows[0].Metadata, @@ -49,5 +49,5 @@ func (r iteratorForJobInsertFastManyCopyFrom) Err() error { } func (q *Queries) JobInsertFastManyCopyFrom(ctx context.Context, db DBTX, arg []*JobInsertFastManyCopyFromParams) (int64, error) { - return db.CopyFrom(ctx, []string{"river_job"}, []string{"args", "finalized_at", "kind", "max_attempts", "metadata", "priority", "queue", "scheduled_at", "state", "tags", "unique_key", "unique_states"}, &iteratorForJobInsertFastManyCopyFrom{rows: arg}) + return db.CopyFrom(ctx, []string{"river_job"}, []string{"args", "created_at", "kind", "max_attempts", "metadata", "priority", "queue", "scheduled_at", "state", "tags", "unique_key", "unique_states"}, &iteratorForJobInsertFastManyCopyFrom{rows: arg}) } diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index 7f3ab708..09b9b154 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -197,6 +197,7 @@ LIMIT @max; -- name: JobInsertFastMany :many INSERT INTO river_job( args, + created_at, kind, max_attempts, metadata, @@ -209,6 +210,7 @@ INSERT INTO river_job( unique_states ) SELECT unnest(@args::jsonb[]), + unnest(@created_at::timestamptz[]), unnest(@kind::text[]), unnest(@max_attempts::smallint[]), unnest(@metadata::jsonb[]), @@ -237,6 +239,7 @@ RETURNING sqlc.embed(river_job), (xmax != 0) AS unique_skipped_as_duplicate; -- name: JobInsertFastManyNoReturning :execrows INSERT INTO river_job( args, + created_at, kind, max_attempts, metadata, @@ -249,6 +252,7 @@ INSERT INTO river_job( unique_states ) SELECT unnest(@args::jsonb[]), + unnest(@created_at::timestamptz[]), unnest(@kind::text[]), unnest(@max_attempts::smallint[]), unnest(@metadata::jsonb[]), diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go index d1da01a1..ad8746c0 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -515,6 +515,7 @@ func (q *Queries) JobGetStuck(ctx context.Context, db DBTX, arg *JobGetStuckPara const jobInsertFastMany = `-- name: JobInsertFastMany :many INSERT INTO river_job( args, + created_at, kind, max_attempts, metadata, @@ -527,22 +528,23 @@ INSERT INTO river_job( unique_states ) SELECT unnest($1::jsonb[]), - unnest($2::text[]), - unnest($3::smallint[]), - unnest($4::jsonb[]), - unnest($5::smallint[]), - unnest($6::text[]), - unnest($7::timestamptz[]), + unnest($2::timestamptz[]), + unnest($3::text[]), + unnest($4::smallint[]), + unnest($5::jsonb[]), + unnest($6::smallint[]), + unnest($7::text[]), + unnest($8::timestamptz[]), -- To avoid requiring pgx users to register the OID of the river_job_state[] -- type, we cast the array to text[] and then to river_job_state. - unnest($8::text[])::river_job_state, + unnest($9::text[])::river_job_state, -- Unnest on a multi-dimensional array will fully flatten the array, so we -- encode the tag list as a comma-separated string and split it in the -- query. - string_to_array(unnest($9::text[]), ','), + string_to_array(unnest($10::text[]), ','), - unnest($10::bytea[]), - unnest($11::bit(8)[]) + unnest($11::bytea[]), + unnest($12::bit(8)[]) ON CONFLICT (unique_key) WHERE unique_key IS NOT NULL @@ -555,6 +557,7 @@ RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_a type JobInsertFastManyParams struct { Args [][]byte + CreatedAt []time.Time Kind []string MaxAttempts []int16 Metadata [][]byte @@ -575,6 +578,7 @@ type JobInsertFastManyRow struct { func (q *Queries) JobInsertFastMany(ctx context.Context, db DBTX, arg *JobInsertFastManyParams) ([]*JobInsertFastManyRow, error) { rows, err := db.Query(ctx, jobInsertFastMany, arg.Args, + arg.CreatedAt, arg.Kind, arg.MaxAttempts, arg.Metadata, @@ -627,6 +631,7 @@ func (q *Queries) JobInsertFastMany(ctx context.Context, db DBTX, arg *JobInsert const jobInsertFastManyNoReturning = `-- name: JobInsertFastManyNoReturning :execrows INSERT INTO river_job( args, + created_at, kind, max_attempts, metadata, @@ -639,22 +644,23 @@ INSERT INTO river_job( unique_states ) SELECT unnest($1::jsonb[]), - unnest($2::text[]), - unnest($3::smallint[]), - unnest($4::jsonb[]), - unnest($5::smallint[]), - unnest($6::text[]), - unnest($7::timestamptz[]), - unnest($8::river_job_state[]), + unnest($2::timestamptz[]), + unnest($3::text[]), + unnest($4::smallint[]), + unnest($5::jsonb[]), + unnest($6::smallint[]), + unnest($7::text[]), + unnest($8::timestamptz[]), + unnest($9::river_job_state[]), -- lib/pq really, REALLY does not play nicely with multi-dimensional arrays, -- so instead we pack each set of tags into a string, send them through, -- then unpack them here into an array to put in each row. This isn't -- necessary in the Pgx driver where copyfrom is used instead. - string_to_array(unnest($9::text[]), ','), + string_to_array(unnest($10::text[]), ','), - unnest($10::bytea[]), - unnest($11::bit(8)[]) + unnest($11::bytea[]), + unnest($12::bit(8)[]) ON CONFLICT (unique_key) WHERE unique_key IS NOT NULL @@ -665,6 +671,7 @@ DO NOTHING type JobInsertFastManyNoReturningParams struct { Args [][]byte + CreatedAt []time.Time Kind []string MaxAttempts []int16 Metadata [][]byte @@ -680,6 +687,7 @@ type JobInsertFastManyNoReturningParams struct { func (q *Queries) JobInsertFastManyNoReturning(ctx context.Context, db DBTX, arg *JobInsertFastManyNoReturningParams) (int64, error) { result, err := db.Exec(ctx, jobInsertFastManyNoReturning, arg.Args, + arg.CreatedAt, arg.Kind, arg.MaxAttempts, arg.Metadata, diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql index 3ce273d4..54fdbeaa 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql @@ -1,7 +1,7 @@ -- name: JobInsertFastManyCopyFrom :copyfrom INSERT INTO river_job( args, - finalized_at, + created_at, kind, max_attempts, metadata, @@ -14,7 +14,7 @@ INSERT INTO river_job( unique_states ) VALUES ( @args, - @finalized_at, + @created_at, @kind, @max_attempts, @metadata, diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go index 47738d18..c44bbab8 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job_copyfrom.sql.go @@ -13,7 +13,7 @@ import ( type JobInsertFastManyCopyFromParams struct { Args []byte - FinalizedAt *time.Time + CreatedAt time.Time Kind string MaxAttempts int16 Metadata []byte diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/riverdriver/riverpgxv5/river_pgx_v5_driver.go index f7ac4250..5f304383 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -198,6 +198,7 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver.JobInsertFastParams) ([]*riverdriver.JobInsertFastResult, error) { insertJobsParams := &dbsqlc.JobInsertFastManyParams{ Args: make([][]byte, len(params)), + CreatedAt: make([]time.Time, len(params)), Kind: make([]string, len(params)), MaxAttempts: make([]int16, len(params)), Metadata: make([][]byte, len(params)), @@ -214,6 +215,11 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver. for i := 0; i < len(params); i++ { params := params[i] + createdAt := now + if params.CreatedAt != nil { + createdAt = *params.CreatedAt + } + scheduledAt := now if params.ScheduledAt != nil { scheduledAt = *params.ScheduledAt @@ -227,6 +233,7 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params []*riverdriver. defaultObject := []byte("{}") insertJobsParams.Args[i] = sliceutil.FirstNonEmpty(params.EncodedArgs, defaultObject) + insertJobsParams.CreatedAt[i] = createdAt insertJobsParams.Kind[i] = params.Kind insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec insertJobsParams.Metadata[i] = sliceutil.FirstNonEmpty(params.Metadata, defaultObject) @@ -260,6 +267,11 @@ func (e *Executor) JobInsertFastManyNoReturning(ctx context.Context, params []*r for i := 0; i < len(params); i++ { params := params[i] + createdAt := now + if params.CreatedAt != nil { + createdAt = *params.CreatedAt + } + metadata := params.Metadata if metadata == nil { metadata = []byte("{}") @@ -277,6 +289,7 @@ func (e *Executor) JobInsertFastManyNoReturning(ctx context.Context, params []*r insertJobsParams[i] = &dbsqlc.JobInsertFastManyCopyFromParams{ Args: params.EncodedArgs, + CreatedAt: createdAt, Kind: params.Kind, MaxAttempts: int16(min(params.MaxAttempts, math.MaxInt16)), //nolint:gosec Metadata: metadata, From 39b42c842dc9e021b4f1214d3ae59ede08201868 Mon Sep 17 00:00:00 2001 From: Brandur Date: Sat, 15 Feb 2025 14:46:53 -0800 Subject: [PATCH 2/5] Let test workers call `JobCompleteTx` + test work transaction variants Follows up #753 to make it possible to call `JobCompleteTx` inside test worker implementations. The test worker tries to update an equivalent job in the database's state to `running` so `JobCompleteTx` will be able to find it when called (it only considers jobs in `running` state). We also augment the API so that it's workable for users who want to use test transactions in their tests (which is hopefully the vast majority of users) by adding transaction variants `WorkTx` and `WorkJobTx`. --- rivertest/worker.go | 51 +++++++++ rivertest/worker_test.go | 235 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 270 insertions(+), 16 deletions(-) diff --git a/rivertest/worker.go b/rivertest/worker.go index 55a623c9..947904de 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -3,6 +3,7 @@ package rivertest import ( "context" "encoding/json" + "errors" "sync/atomic" "testing" "time" @@ -77,12 +78,62 @@ func (w *Worker[T, TTx]) Work(ctx context.Context, tb testing.TB, args T, opts * return w.WorkJob(ctx, tb, job) } +// WorkTx allocates a synthetic job with the provided arguments and optional +// insert options, then works it in the given transaction. +func (w *Worker[T, TTx]) WorkTx(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) error { + tb.Helper() + job := makeJobFromArgs(tb, w.config, args, opts) + return w.WorkJobTx(ctx, tb, tx, job) +} + // WorkJob works the provided job. The job must be constructed to be a realistic // job using external logic prior to calling this method. Unlike the other // variants, this method offers total control over the job's attributes. func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, job *river.Job[T]) error { tb.Helper() + var exec riverdriver.Executor + if w.client.Driver().HasPool() { + exec = w.client.Driver().GetExecutor() + } + + return w.workJobExec(ctx, tb, exec, job) +} + +// WorkJobTx works the provided job in the given transaction. The job must be +// constructed to be a realistic job using external logic prior to calling this +// method. Unlike the other variants, this method offers total control over the +// job's attributes. +func (w *Worker[T, TTx]) WorkJobTx(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { + tb.Helper() + return w.workJobExec(ctx, tb, w.client.Driver().UnwrapExecutor(tx), job) +} + +func (w *Worker[T, TTx]) workJobExec(ctx context.Context, tb testing.TB, exec riverdriver.Executor, job *river.Job[T]) error { + tb.Helper() + + // Try to transition an existing job row to running, but also tolerate the + // row not existing in the database. Most of the time you don't need to + // insert a real job row to use this function (it's only necessary if you + // want to use `JobCompleteTx`). + if exec != nil { + updatedJobRow, err := exec.JobUpdate(ctx, &riverdriver.JobUpdateParams{ + ID: job.ID, + StateDoUpdate: true, + State: rivertype.JobStateRunning, + }) + if err != nil && !errors.Is(err, rivertype.ErrNotFound) { + return err + } + if updatedJobRow != nil { + job.JobRow = updatedJobRow + } + } + + // Regardless of whether the job was actually in the database, transition it + // to running. + job.JobRow.State = rivertype.JobStateRunning + // populate river client into context: ctx = WorkContext(ctx, w.client) ctx = context.WithValue(ctx, execution.ContextKeyInsideTestWorker{}, true) diff --git a/rivertest/worker_test.go b/rivertest/worker_test.go index 66a9f7f3..9114be04 100644 --- a/rivertest/worker_test.go +++ b/rivertest/worker_test.go @@ -5,11 +5,14 @@ import ( "testing" "time" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/require" "github.com/riverqueue/river" "github.com/riverqueue/river/internal/dbunique" "github.com/riverqueue/river/internal/execution" + "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/riversharedtest" @@ -118,7 +121,7 @@ func TestWorker_Work(t *testing.T) { }) tw := NewWorker(t, driver, config, worker) - // You can also use the WorkOpts method to pass in custom insert options: + // You can also pass in custom insert options: require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test3"}, &river.InsertOpts{ MaxAttempts: 420, Metadata: []byte(`{"key": "value"}`), @@ -163,24 +166,224 @@ func TestWorker_Work(t *testing.T) { }) } +func TestWorker_WorkTx(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + client *river.Client[pgx.Tx] + driver *riverpgxv5.Driver + tx pgx.Tx + workFunc func(ctx context.Context, job *river.Job[testArgs]) error + } + + setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { + t.Helper() + + var ( + config = &river.Config{} + driver = riverpgxv5.New(nil) + ) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + + bundle := &testBundle{ + client: client, + driver: driver, + tx: riverinternaltest.TestTx(ctx, t), + workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, + } + + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + return bundle.workFunc(ctx, job) + }) + + return NewWorker(t, driver, config, worker), bundle + } + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.Equal(t, rivertype.JobStateRunning, job.State) + return nil + } + + require.NoError(t, testWorker.WorkTx(ctx, t, bundle.tx, testArgs{}, nil)) + }) +} + func TestWorker_WorkJob(t *testing.T) { t.Parallel() - config := &river.Config{} - driver := riverpgxv5.New(nil) + ctx := context.Background() + + type testBundle struct { + client *river.Client[pgx.Tx] + dbPool *pgxpool.Pool + driver *riverpgxv5.Driver + workFunc func(ctx context.Context, job *river.Job[testArgs]) error + } + + setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { + t.Helper() + + var ( + config = &river.Config{} + dbPool = riverinternaltest.TestDB(ctx, t) + driver = riverpgxv5.New(dbPool) + ) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + + bundle := &testBundle{ + client: client, + driver: driver, + dbPool: dbPool, + workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, + } + + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + return bundle.workFunc(ctx, job) + }) + + return NewWorker(t, driver, config, worker), bundle + } + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.Equal(t, []string{"worker123"}, job.JobRow.AttemptedBy) + require.Equal(t, rivertype.JobStateRunning, job.State) + return nil + } + + now := time.Now() + require.NoError(t, testWorker.WorkJob(ctx, t, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{ + AttemptedAt: &now, + AttemptedBy: []string{"worker123"}, + CreatedAt: &now, + EncodedArgs: []byte(`{"value": "test"}`), + Errors: nil, + }))) + }) + + t.Run("JobCompleteTxWithInsertedJobRow", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) - worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { - require.Equal(t, []string{"worker123"}, job.JobRow.AttemptedBy) - return nil + args := testArgs{} + insertRes, err := bundle.client.Insert(ctx, args, nil) + require.NoError(t, err) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + tx, err := bundle.dbPool.Begin(ctx) + require.NoError(t, err) + + updatedJob, err := bundle.driver.GetExecutor().JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + + _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, tx, job) + require.NoError(t, err) + + err = tx.Commit(ctx) + require.NoError(t, err) + + return nil + } + + require.NoError(t, testWorker.WorkJob(ctx, t, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + + updatedJob, err := bundle.driver.GetExecutor().JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateCompleted, updatedJob.State) + }) +} + +func TestWorker_WorkJobTx(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + client *river.Client[pgx.Tx] + driver *riverpgxv5.Driver + tx pgx.Tx + workFunc func(ctx context.Context, job *river.Job[testArgs]) error + } + + setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { + t.Helper() + + var ( + config = &river.Config{} + driver = riverpgxv5.New(nil) + ) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + + bundle := &testBundle{ + client: client, + driver: driver, + tx: riverinternaltest.TestTx(ctx, t), + workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, + } + + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + return bundle.workFunc(ctx, job) + }) + + return NewWorker(t, driver, config, worker), bundle + } + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.Equal(t, rivertype.JobStateRunning, job.State) + return nil + } + + require.NoError(t, testWorker.WorkJobTx(ctx, t, bundle.tx, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{}))) + }) + + t.Run("JobCompleteTxWithInsertedJobRow", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + args := testArgs{} + insertRes, err := bundle.client.InsertTx(ctx, bundle.tx, args, nil) + require.NoError(t, err) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + + _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, bundle.tx, job) + require.NoError(t, err) + + return nil + } + + require.NoError(t, testWorker.WorkJobTx(ctx, t, bundle.tx, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + + updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateCompleted, updatedJob.State) }) - tw := NewWorker(t, driver, config, worker) - - now := time.Now() - require.NoError(t, tw.WorkJob(context.Background(), t, makeJobFromFactoryBuild(t, testArgs{Value: "test"}, &testfactory.JobOpts{ - AttemptedAt: &now, - AttemptedBy: []string{"worker123"}, - CreatedAt: &now, - EncodedArgs: []byte(`{"value": "test"}`), - Errors: nil, - }))) } From 4c466fc98406c9a58ad9e80be4061617eb27a3b7 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 16 Feb 2025 15:33:29 -0600 Subject: [PATCH 3/5] always take tx in rivertest.Worker methods --- client.go | 18 +++++++++++++++--- rivertest/worker.go | 46 +++++++++++++++++++++++++++++++-------------- 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/client.go b/client.go index 6f6233f3..aaaebc61 100644 --- a/client.go +++ b/client.go @@ -48,6 +48,15 @@ const ( // TestConfig contains configuration specific to test environments. type TestConfig struct { + // DisableUniqueEnforcement disables the application of unique job + // constraints. This is useful for testing scenarios when testing a worker + // that typically uses uniqueness, but where enforcing uniqueness would cause + // conflicts with parallel test execution. + // + // The [rivertest.Worker] type automatically disables uniqueness enforcement + // when creating jobs. + DisableUniqueEnforcement bool + // Time is a time generator to make time stubbable in tests. Time rivertype.TimeGenerator } @@ -1253,9 +1262,12 @@ func insertParamsFromConfigArgsAndOptions(archetype *baseservice.Archetype, conf return nil, errors.New("priority must be between 1 and 4") } - uniqueOpts := insertOpts.UniqueOpts - if uniqueOpts.isEmpty() { - uniqueOpts = jobInsertOpts.UniqueOpts + var uniqueOpts UniqueOpts + if !config.Test.DisableUniqueEnforcement { + uniqueOpts = insertOpts.UniqueOpts + if uniqueOpts.isEmpty() { + uniqueOpts = jobInsertOpts.UniqueOpts + } } if err := uniqueOpts.validate(); err != nil { return nil, err diff --git a/rivertest/worker.go b/rivertest/worker.go index 947904de..88281101 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -58,6 +58,8 @@ func NewWorker[T river.JobArgs, TTx any](tb testing.TB, driver riverdriver.Drive tb.Helper() config = config.WithDefaults() + config.Test.DisableUniqueEnforcement = true + client, err := river.NewClient(driver, config) if err != nil { tb.Fatalf("failed to create client: %s", err) @@ -70,34 +72,48 @@ func NewWorker[T river.JobArgs, TTx any](tb testing.TB, driver riverdriver.Drive } } +func (w *Worker[T, TTx]) insertJob(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) *river.Job[T] { + tb.Helper() + + result, err := w.client.InsertTx(ctx, tx, args, opts) + if err != nil { + tb.Fatalf("failed to insert job: %s", err) + } + var decodedArgs T + if err := json.Unmarshal(result.Job.EncodedArgs, &decodedArgs); err != nil { + tb.Fatalf("failed to unmarshal args: %s", err) + } + + return &river.Job[T]{ + Args: decodedArgs, + JobRow: result.Job, + } +} + // Work allocates a synthetic job with the provided arguments and optional // insert options, then works it. -func (w *Worker[T, TTx]) Work(ctx context.Context, tb testing.TB, args T, opts *river.InsertOpts) error { +func (w *Worker[T, TTx]) Work(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) error { tb.Helper() - job := makeJobFromArgs(tb, w.config, args, opts) - return w.WorkJob(ctx, tb, job) + + job := w.insertJob(ctx, tb, tx, args, opts) + return w.WorkJob(ctx, tb, tx, job) } // WorkTx allocates a synthetic job with the provided arguments and optional // insert options, then works it in the given transaction. func (w *Worker[T, TTx]) WorkTx(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) error { tb.Helper() - job := makeJobFromArgs(tb, w.config, args, opts) + + job := w.insertJob(ctx, tb, tx, args, opts) return w.WorkJobTx(ctx, tb, tx, job) } // WorkJob works the provided job. The job must be constructed to be a realistic // job using external logic prior to calling this method. Unlike the other // variants, this method offers total control over the job's attributes. -func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, job *river.Job[T]) error { +func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { tb.Helper() - - var exec riverdriver.Executor - if w.client.Driver().HasPool() { - exec = w.client.Driver().GetExecutor() - } - - return w.workJobExec(ctx, tb, exec, job) + return w.workJob(ctx, tb, tx, job) } // WorkJobTx works the provided job in the given transaction. The job must be @@ -106,12 +122,14 @@ func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, job *river. // job's attributes. func (w *Worker[T, TTx]) WorkJobTx(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { tb.Helper() - return w.workJobExec(ctx, tb, w.client.Driver().UnwrapExecutor(tx), job) + return w.workJob(ctx, tb, tx, job) } -func (w *Worker[T, TTx]) workJobExec(ctx context.Context, tb testing.TB, exec riverdriver.Executor, job *river.Job[T]) error { +func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { tb.Helper() + exec := w.client.Driver().UnwrapExecutor(tx) + // Try to transition an existing job row to running, but also tolerate the // row not existing in the database. Most of the time you don't need to // insert a real job row to use this function (it's only necessary if you From 1549634e4d1aa37e5162b9545b9d64c4c26de3ae Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 16 Feb 2025 17:28:48 -0600 Subject: [PATCH 4/5] always insert jobs for testworker --- rivertest/worker.go | 90 ++++------------------------ rivertest/worker_test.go | 126 ++++++++++++++++++--------------------- 2 files changed, 72 insertions(+), 144 deletions(-) diff --git a/rivertest/worker.go b/rivertest/worker.go index 88281101..3095db98 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -9,13 +9,11 @@ import ( "time" "github.com/riverqueue/river" - "github.com/riverqueue/river/internal/dbunique" "github.com/riverqueue/river/internal/execution" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/testfactory" "github.com/riverqueue/river/rivershared/util/ptrutil" - "github.com/riverqueue/river/rivershared/util/sliceutil" "github.com/riverqueue/river/rivershared/util/valutil" "github.com/riverqueue/river/rivertype" ) @@ -135,10 +133,20 @@ func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job // insert a real job row to use this function (it's only necessary if you // want to use `JobCompleteTx`). if exec != nil { + timeGen := w.config.Test.Time + if timeGen == nil { + timeGen = &baseservice.UnStubbableTimeGenerator{} + } updatedJobRow, err := exec.JobUpdate(ctx, &riverdriver.JobUpdateParams{ - ID: job.ID, - StateDoUpdate: true, - State: rivertype.JobStateRunning, + ID: job.ID, + Attempt: job.JobRow.Attempt + 1, + AttemptDoUpdate: true, + AttemptedAt: ptrutil.Ptr(timeGen.NowUTC()), + AttemptedAtDoUpdate: true, + AttemptedBy: append(job.JobRow.AttemptedBy, "worker1"), + AttemptedByDoUpdate: true, + StateDoUpdate: true, + State: rivertype.JobStateRunning, }) if err != nil && !errors.Is(err, rivertype.ErrNotFound) { return err @@ -176,78 +184,6 @@ func nextID() int64 { return atomic.AddInt64(&idSeq, 1) } -func makeJobFromArgs[T river.JobArgs](tb testing.TB, config *river.Config, args T, opts *river.InsertOpts) *river.Job[T] { - tb.Helper() - - encodedArgs, err := json.Marshal(args) - if err != nil { - tb.Fatalf("failed to marshal args: %s", err) - } - - // Round trip the args to validate JSON marshalling/unmarshalling on the type. - var decodedArgs T - if err = json.Unmarshal(encodedArgs, &decodedArgs); err != nil { - tb.Fatalf("failed to unmarshal args: %s", err) - } - - if opts == nil { - opts = &river.InsertOpts{} - } - - // Extract any job insert opts from args type if present - var jobInsertOpts river.InsertOpts - if argsWithOpts, ok := any(args).(river.JobArgsWithInsertOpts); ok { - jobInsertOpts = argsWithOpts.InsertOpts() - } - - now := valutil.FirstNonZero(opts.ScheduledAt, jobInsertOpts.ScheduledAt, time.Now()) - - internalUniqueOpts := (*dbunique.UniqueOpts)(&opts.UniqueOpts) - insertParams := &rivertype.JobInsertParams{ - Args: args, - CreatedAt: &now, - EncodedArgs: encodedArgs, - Kind: args.Kind(), - MaxAttempts: valutil.FirstNonZero(opts.MaxAttempts, jobInsertOpts.MaxAttempts, config.MaxAttempts, river.MaxAttemptsDefault), - Metadata: sliceutil.FirstNonEmpty(opts.Metadata, jobInsertOpts.Metadata, []byte(`{}`)), - Priority: valutil.FirstNonZero(opts.Priority, jobInsertOpts.Priority, river.PriorityDefault), - Queue: valutil.FirstNonZero(opts.Queue, jobInsertOpts.Queue, river.QueueDefault), - State: rivertype.JobStateAvailable, - Tags: sliceutil.FirstNonEmpty(opts.Tags, jobInsertOpts.Tags, []string{}), - } - var uniqueKey []byte - timeGen := config.Test.Time - if timeGen == nil { - timeGen = &baseservice.UnStubbableTimeGenerator{} - } - if !internalUniqueOpts.IsEmpty() { - uniqueKey, err = dbunique.UniqueKey(valutil.ValOrDefault(config.Test.Time, timeGen), internalUniqueOpts, insertParams) - if err != nil { - tb.Fatalf("failed to create unique key: %s", err) - } - } - - return makeJobFromFactoryBuild(tb, args, &testfactory.JobOpts{ - Attempt: ptrutil.Ptr(1), - AttemptedAt: &now, - AttemptedBy: []string{"worker1"}, - CreatedAt: &now, - EncodedArgs: encodedArgs, - Errors: nil, - FinalizedAt: nil, - Kind: ptrutil.Ptr(args.Kind()), - MaxAttempts: ptrutil.Ptr(valutil.FirstNonZero(opts.MaxAttempts, jobInsertOpts.MaxAttempts, config.MaxAttempts, river.MaxAttemptsDefault)), - Metadata: sliceutil.FirstNonEmpty(opts.Metadata, jobInsertOpts.Metadata, []byte(`{}`)), - Priority: ptrutil.Ptr(valutil.FirstNonZero(opts.Priority, jobInsertOpts.Priority, river.PriorityDefault)), - Queue: ptrutil.Ptr(valutil.FirstNonZero(opts.Queue, jobInsertOpts.Queue, river.QueueDefault)), - ScheduledAt: &now, - State: ptrutil.Ptr(rivertype.JobStateRunning), - Tags: sliceutil.FirstNonEmpty(opts.Tags, jobInsertOpts.Tags, []string{}), - UniqueKey: uniqueKey, - UniqueStates: internalUniqueOpts.StateBitmask(), - }) -} - func makeJobFromFactoryBuild[T river.JobArgs](tb testing.TB, args T, jobOpts *testfactory.JobOpts) *river.Job[T] { tb.Helper() diff --git a/rivertest/worker_test.go b/rivertest/worker_test.go index 9114be04..7197828a 100644 --- a/rivertest/worker_test.go +++ b/rivertest/worker_test.go @@ -6,15 +6,12 @@ import ( "time" "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/require" "github.com/riverqueue/river" - "github.com/riverqueue/river/internal/dbunique" "github.com/riverqueue/river/internal/execution" "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/riverdriver/riverpgxv5" - "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/testfactory" "github.com/riverqueue/river/rivertype" @@ -29,12 +26,35 @@ func (testArgs) Kind() string { return "rivertest_work_test" } func TestWorker_Work(t *testing.T) { t.Parallel() - config := &river.Config{} - driver := riverpgxv5.New(nil) + ctx := context.Background() + + type testBundle struct { + config *river.Config + driver *riverpgxv5.Driver + tx pgx.Tx + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + var ( + config = &river.Config{} + driver = riverpgxv5.New(nil) + tx = riverinternaltest.TestTx(ctx, t) + ) + + return &testBundle{ + config: config, + driver: driver, + tx: tx, + } + } t.Run("WorkASimpleJob", func(t *testing.T) { t.Parallel() + bundle := setup(t) + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { require.Equal(t, testArgs{Value: "test"}, job.Args) require.Equal(t, 1, job.JobRow.Attempt) @@ -61,48 +81,37 @@ func TestWorker_Work(t *testing.T) { return nil }) - tw := NewWorker(t, driver, config, worker) - require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test"}, nil)) + tw := NewWorker(t, bundle.driver, bundle.config, worker) + require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test"}, nil)) }) t.Run("Reusable", func(t *testing.T) { t.Parallel() + bundle := setup(t) + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { return nil }) - tw := NewWorker(t, driver, config, worker) - require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test"}, nil)) - require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test2"}, nil)) + tw := NewWorker(t, bundle.driver, bundle.config, worker) + require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test"}, nil)) + require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test2"}, nil)) }) t.Run("SetsCustomInsertOpts", func(t *testing.T) { t.Parallel() - uniqueOpts := river.UniqueOpts{ByQueue: true} - hourFromNow := time.Now().Add(1 * time.Hour) - internalUniqueOpts := (*dbunique.UniqueOpts)(&uniqueOpts) - uniqueKey, err := dbunique.UniqueKey(&baseservice.UnStubbableTimeGenerator{}, internalUniqueOpts, &rivertype.JobInsertParams{ - Args: testArgs{Value: "test3"}, - CreatedAt: &hourFromNow, - EncodedArgs: []byte(`{"value": "test3"}`), - Kind: "rivertest_work_test", - MaxAttempts: 420, - Metadata: []byte(`{"key": "value"}`), - Priority: 3, - Queue: "custom_queue", - State: rivertype.JobStateAvailable, - Tags: []string{"tag1", "tag2"}, - }) - require.NoError(t, err) + bundle := setup(t) + + hourFromNow := time.Now().UTC().Add(1 * time.Hour) worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { require.Equal(t, testArgs{Value: "test3"}, job.Args) require.Equal(t, 1, job.JobRow.Attempt) require.NotNil(t, job.JobRow.AttemptedAt) - require.WithinDuration(t, hourFromNow, *job.JobRow.AttemptedAt, 2*time.Second) + require.WithinDuration(t, time.Now().UTC(), *job.JobRow.AttemptedAt, 2*time.Second) require.Equal(t, []string{"worker1"}, job.JobRow.AttemptedBy) - require.WithinDuration(t, hourFromNow, job.JobRow.CreatedAt, 2*time.Second) + require.WithinDuration(t, time.Now().UTC(), job.JobRow.CreatedAt, 2*time.Second) require.JSONEq(t, `{"value": "test3"}`, string(job.JobRow.EncodedArgs)) require.Empty(t, job.JobRow.Errors) require.Nil(t, job.JobRow.FinalizedAt) @@ -115,14 +124,13 @@ func TestWorker_Work(t *testing.T) { require.WithinDuration(t, hourFromNow, job.JobRow.ScheduledAt, 2*time.Second) require.Equal(t, rivertype.JobStateRunning, job.JobRow.State) require.Equal(t, []string{"tag1", "tag2"}, job.JobRow.Tags) - require.Equal(t, uniqueKey, job.JobRow.UniqueKey) return nil }) - tw := NewWorker(t, driver, config, worker) + tw := NewWorker(t, bundle.driver, bundle.config, worker) // You can also pass in custom insert options: - require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test3"}, &river.InsertOpts{ + require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test3"}, &river.InsertOpts{ MaxAttempts: 420, Metadata: []byte(`{"key": "value"}`), Pending: true, // ignored but added to ensure non-default behavior @@ -130,38 +138,29 @@ func TestWorker_Work(t *testing.T) { Queue: "custom_queue", ScheduledAt: hourFromNow, Tags: []string{"tag1", "tag2"}, - UniqueOpts: uniqueOpts, })) }) - t.Run("UniqueOptsByPeriodRespectsCustomStubbedTime", func(t *testing.T) { + t.Run("UniqueOptsAreIgnored", func(t *testing.T) { t.Parallel() + // UniqueOpts must be ignored because otherwise there's a likelihood of + // conflicts with parallel tests inserting jobs with the same unique key. + + bundle := setup(t) stubTime := &riversharedtest.TimeStub{} now := time.Now().UTC() stubTime.StubNowUTC(now) - config := &river.Config{ - Test: river.TestConfig{Time: stubTime}, - } - - uniqueOpts := river.UniqueOpts{ByPeriod: 1 * time.Hour} - internalUniqueOpts := (*dbunique.UniqueOpts)(&uniqueOpts) - uniqueKey, err := dbunique.UniqueKey(stubTime, internalUniqueOpts, &rivertype.JobInsertParams{ - Args: testArgs{Value: "test3"}, - CreatedAt: &now, - EncodedArgs: []byte(`{"value": "test3"}`), - Kind: "rivertest_work_test", - }) - require.NoError(t, err) + bundle.config.Test.Time = stubTime worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { - require.Equal(t, uniqueKey, job.JobRow.UniqueKey) - + require.Empty(t, job.JobRow.UniqueKey) + require.Empty(t, job.JobRow.UniqueStates) return nil }) - tw := NewWorker(t, driver, config, worker) - require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test"}, &river.InsertOpts{ - UniqueOpts: uniqueOpts, + tw := NewWorker(t, bundle.driver, bundle.config, worker) + require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test"}, &river.InsertOpts{ + UniqueOpts: river.UniqueOpts{ByPeriod: 1 * time.Hour}, })) }) } @@ -224,8 +223,8 @@ func TestWorker_WorkJob(t *testing.T) { type testBundle struct { client *river.Client[pgx.Tx] - dbPool *pgxpool.Pool driver *riverpgxv5.Driver + tx pgx.Tx workFunc func(ctx context.Context, job *river.Job[testArgs]) error } @@ -234,8 +233,7 @@ func TestWorker_WorkJob(t *testing.T) { var ( config = &river.Config{} - dbPool = riverinternaltest.TestDB(ctx, t) - driver = riverpgxv5.New(dbPool) + driver = riverpgxv5.New(nil) ) client, err := river.NewClient(driver, config) @@ -244,7 +242,7 @@ func TestWorker_WorkJob(t *testing.T) { bundle := &testBundle{ client: client, driver: driver, - dbPool: dbPool, + tx: riverinternaltest.TestTx(ctx, t), workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, } @@ -267,7 +265,7 @@ func TestWorker_WorkJob(t *testing.T) { } now := time.Now() - require.NoError(t, testWorker.WorkJob(ctx, t, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{ + require.NoError(t, testWorker.WorkJob(ctx, t, bundle.tx, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{ AttemptedAt: &now, AttemptedBy: []string{"worker123"}, CreatedAt: &now, @@ -282,29 +280,23 @@ func TestWorker_WorkJob(t *testing.T) { testWorker, bundle := setup(t) args := testArgs{} - insertRes, err := bundle.client.Insert(ctx, args, nil) + insertRes, err := bundle.client.InsertTx(ctx, bundle.tx, args, nil) require.NoError(t, err) bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { - tx, err := bundle.dbPool.Begin(ctx) - require.NoError(t, err) - - updatedJob, err := bundle.driver.GetExecutor().JobGetByID(ctx, insertRes.Job.ID) + updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) require.NoError(t, err) require.Equal(t, rivertype.JobStateRunning, updatedJob.State) - _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, tx, job) - require.NoError(t, err) - - err = tx.Commit(ctx) + _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, bundle.tx, job) require.NoError(t, err) return nil } - require.NoError(t, testWorker.WorkJob(ctx, t, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + require.NoError(t, testWorker.WorkJob(ctx, t, bundle.tx, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) - updatedJob, err := bundle.driver.GetExecutor().JobGetByID(ctx, insertRes.Job.ID) + updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) require.NoError(t, err) require.Equal(t, rivertype.JobStateCompleted, updatedJob.State) }) From d2607ac25676ff4b9feb73d385b47c7b0b619d1a Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 16 Feb 2025 18:03:53 -0600 Subject: [PATCH 5/5] use real executor for rivertest worker, alter APIs Leverage the actual job executor for the `rivertest.Worker` type so that it closely matches real job executions. Reduce the exposed APIs so it only has `Work` and `WorkJob`, both of which require a transaction (but the latter works on existing jobs). Rather than returning only an error, also return a `WorkResult` struct that contains additional metadata. Populate it with the final updated job and the event kind to indicate whether the job completed, errored, snoozed, cancelled, etc. --- CHANGELOG.md | 11 ++ rivertest/worker.go | 327 +++++++++++++++++++++++---------------- rivertest/worker_test.go | 197 ++++++++--------------- 3 files changed, 274 insertions(+), 261 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a62585b5..96fb58c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +⚠️ Version 0.18.0 has breaking changes for the `rivertest.Worker` type that was just introduced. While attempting to round out some edge cases with its design, we realized some of them simply couldn't be solved adequately without changing the overall design such that all tested jobs are inserted into the database. Given the short duration since it was released (over a weekend) it's unlikely many users have adopted it and it seemed best to rip off the bandaid to fix it before it gets widely used. + +### Changed + +- **Breaking change:** The `rivertest.Worker` type now requires all jobs to be inserted into the database. The original design allowed workers to be tested without hitting the database at all. Ultimately this design made it hard to correctly simulate features like `JobCompleteTx` and the other potential solutions seemed undesirable. + + As part of this change, the `Work` and `WorkJob` methods now take a transaction argument. The expectation is that a transaction will be opened by the caller and rolled back after test completion. Additionally, the return signature was changed to return a `WorkResult` struct alongside the error. The struct includes the post-execution job row as well as the event kind that occurred, making it easy to inspect the job's state after execution. + + Finally, the implementation was refactored so that it uses the _real_ `river.Client` insert path, and also uses the same job execution path as real execution. This minimizes the potential for differences in behavior between testing and real execution. + [PR #766](https://github.com/riverqueue/river/pull/766). + ## [0.17.0] - 2025-02-16 ### Added diff --git a/rivertest/worker.go b/rivertest/worker.go index 3095db98..752bf24f 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -4,40 +4,46 @@ import ( "context" "encoding/json" "errors" - "sync/atomic" "testing" "time" "github.com/riverqueue/river" "github.com/riverqueue/river/internal/execution" + "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/maintenance" + "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" - "github.com/riverqueue/river/rivershared/testfactory" + "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/util/ptrutil" - "github.com/riverqueue/river/rivershared/util/valutil" "github.com/riverqueue/river/rivertype" ) // Worker makes it easier to test river workers. Once built, the worker can be -// used to work any number of synthetic jobs without touching the database: +// used to insert and work any number jobs: // // worker := rivertest.NewWorker(t, driver, config, worker) -// if err := worker.Work(ctx, t, args, nil); err != nil { +// result, err := worker.Work(ctx, t, tx, args, nil) +// if err != nil { // t.Fatalf("failed to work job: %s", err) // } +// if result.Kind != river.EventKindJobCompleted { +// t.Fatalf("expected job to be completed, got %s", result.Kind) +// } +// +// An existing job (inserted using external logic) can also be worked: // -// if err := worker.Work(ctx, t, args, &river.InsertOpts{Queue: "custom_queue"}); err != nil { +// job := worker.insertJob(ctx, t, tx, args, nil) +// result, err := worker.WorkJob(ctx, t, tx, job) +// if err != nil { // t.Fatalf("failed to work job: %s", err) // } // -// In all cases the underlying [river.Worker] will be called with the synthetic -// job. The execution environment has a realistic River context with access to -// all River features, including [river.ClientFromContext] and worker -// middleware. -// -// When relying on features that require a database record (such as JobCompleteTx), -// the job must be inserted into the database first and then executed with -// WorkJob. +// In all cases the underlying [river.Worker] will be called with the job +// transitioned into a running state. The execution environment has a realistic +// River context with access to all River features, including +// [river.ClientFromContext] and worker middleware. type Worker[T river.JobArgs, TTx any] struct { client *river.Client[TTx] config *river.Config @@ -48,10 +54,13 @@ type Worker[T river.JobArgs, TTx any] struct { // The worker uses the provided driver and River config to populate default // values on test jobs and to configure the execution environment. // -// The database is not required to use this helper, and a pool-less driver is -// recommended for most usage. This enables individual test cases to run in -// parallel and with full isolation, even using a single shared `Worker` -// instance across many test cases. +// A pool-less driver is recommended for most usage, as individual job inserts +// and executions will happen within a provided transaction. This enables many +// parallel test cases to run with full isolation, each in their own +// transaction. +// +// The worker is configured to disable unique enforcement by default, as this +// would otherwise prevent conflicting jobs from being tested in parallel. func NewWorker[T river.JobArgs, TTx any](tb testing.TB, driver riverdriver.Driver[TTx], config *river.Config, worker river.Worker[T]) *Worker[T, TTx] { tb.Helper() @@ -70,159 +79,215 @@ func NewWorker[T river.JobArgs, TTx any](tb testing.TB, driver riverdriver.Drive } } -func (w *Worker[T, TTx]) insertJob(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) *river.Job[T] { +func (w *Worker[T, TTx]) insertJob(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) *rivertype.JobRow { tb.Helper() result, err := w.client.InsertTx(ctx, tx, args, opts) if err != nil { tb.Fatalf("failed to insert job: %s", err) } - var decodedArgs T - if err := json.Unmarshal(result.Job.EncodedArgs, &decodedArgs); err != nil { - tb.Fatalf("failed to unmarshal args: %s", err) - } - - return &river.Job[T]{ - Args: decodedArgs, - JobRow: result.Job, - } + return result.Job } -// Work allocates a synthetic job with the provided arguments and optional -// insert options, then works it. -func (w *Worker[T, TTx]) Work(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) error { +// Work inserts a job with the provided arguments and optional insert options, +// then works it. The transaction is used for all work-related database +// operations so that the caller can easily roll back at the end of the test to +// maintain a clean database state. +// +// The returned WorkResult contains the updated job row from the database +// _after_ it has been worked, as well as the kind of event that occurred. +// +// The returned error only reflects _real_ errors and does not include +// explicitly returned snooze or cancel errors from the worker. +func (w *Worker[T, TTx]) Work(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) (WorkResult, error) { tb.Helper() job := w.insertJob(ctx, tb, tx, args, opts) return w.WorkJob(ctx, tb, tx, job) } -// WorkTx allocates a synthetic job with the provided arguments and optional -// insert options, then works it in the given transaction. -func (w *Worker[T, TTx]) WorkTx(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) error { - tb.Helper() - - job := w.insertJob(ctx, tb, tx, args, opts) - return w.WorkJobTx(ctx, tb, tx, job) -} - -// WorkJob works the provided job. The job must be constructed to be a realistic -// job using external logic prior to calling this method. Unlike the other -// variants, this method offers total control over the job's attributes. -func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { +// WorkJob works an existing job already in the database. The job must be +// inserted using external logic prior to calling this method. The transaction +// is used for all work-related database operations so that the caller can +// easily roll back at the end of the test to maintain a clean database state. +// +// The returned WorkResult contains the updated job row from the database +// _after_ it has been worked, as well as the kind of event that occurred. +// +// The returned error only reflects _real_ errors and does not include +// explicitly returned snooze or cancel errors from the worker. +func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, tx TTx, job *rivertype.JobRow) (WorkResult, error) { tb.Helper() return w.workJob(ctx, tb, tx, job) } -// WorkJobTx works the provided job in the given transaction. The job must be -// constructed to be a realistic job using external logic prior to calling this -// method. Unlike the other variants, this method offers total control over the -// job's attributes. -func (w *Worker[T, TTx]) WorkJobTx(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { +func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job *rivertype.JobRow) (WorkResult, error) { tb.Helper() - return w.workJob(ctx, tb, tx, job) -} -func (w *Worker[T, TTx]) workJob(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { - tb.Helper() + timeGen := w.config.Test.Time + if timeGen == nil { + timeGen = &baseservice.UnStubbableTimeGenerator{} + } exec := w.client.Driver().UnwrapExecutor(tx) + subscribeCh := make(chan []jobcompleter.CompleterJobUpdated, 1) + archetype := riversharedtest.BaseServiceArchetype(tb) + if withStub, ok := timeGen.(baseservice.TimeGeneratorWithStub); ok { + archetype.Time = withStub + } else { + archetype.Time = &baseservice.TimeGeneratorWithStubWrapper{TimeGenerator: timeGen} + } + completer := jobcompleter.NewInlineCompleter(archetype, exec, w.client.Pilot(), subscribeCh) - // Try to transition an existing job row to running, but also tolerate the - // row not existing in the database. Most of the time you don't need to - // insert a real job row to use this function (it's only necessary if you - // want to use `JobCompleteTx`). - if exec != nil { - timeGen := w.config.Test.Time - if timeGen == nil { - timeGen = &baseservice.UnStubbableTimeGenerator{} - } - updatedJobRow, err := exec.JobUpdate(ctx, &riverdriver.JobUpdateParams{ - ID: job.ID, - Attempt: job.JobRow.Attempt + 1, - AttemptDoUpdate: true, - AttemptedAt: ptrutil.Ptr(timeGen.NowUTC()), - AttemptedAtDoUpdate: true, - AttemptedBy: append(job.JobRow.AttemptedBy, "worker1"), - AttemptedByDoUpdate: true, - StateDoUpdate: true, - State: rivertype.JobStateRunning, - }) - if err != nil && !errors.Is(err, rivertype.ErrNotFound) { - return err - } - if updatedJobRow != nil { - job.JobRow = updatedJobRow - } + updatedJobRow, err := exec.JobUpdate(ctx, &riverdriver.JobUpdateParams{ + ID: job.ID, + Attempt: job.Attempt + 1, + AttemptDoUpdate: true, + AttemptedAt: ptrutil.Ptr(timeGen.NowUTC()), + AttemptedAtDoUpdate: true, + AttemptedBy: append(job.AttemptedBy, w.config.ID), + AttemptedByDoUpdate: true, + StateDoUpdate: true, + State: rivertype.JobStateRunning, + }) + if err != nil && !errors.Is(err, rivertype.ErrNotFound) { + tb.Fatalf("test worker internal error: failed to update job to running state: %s", err) } + job = updatedJobRow - // Regardless of whether the job was actually in the database, transition it - // to running. - job.JobRow.State = rivertype.JobStateRunning + workUnit := (&workUnitFactoryWrapper[T]{worker: w.worker}).MakeUnit(job) // populate river client into context: ctx = WorkContext(ctx, w.client) + // TODO: remove ContextKeyInsideTestWorker ctx = context.WithValue(ctx, execution.ContextKeyInsideTestWorker{}, true) - doInner := execution.MiddlewareChain( - w.config.WorkerMiddleware, - w.worker.Middleware(job), - func(ctx context.Context) error { return w.worker.Work(ctx, job) }, - job.JobRow, - ) + // jobCancel will always be called by the executor to prevent leaks. + jobCtx, jobCancel := context.WithCancelCause(ctx) + + executionDone := make(chan struct{}) + + var resultErr error + + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{ + CancelFunc: jobCancel, + ClientJobTimeout: w.config.JobTimeout, + ClientRetryPolicy: w.config.RetryPolicy, + Completer: completer, + DefaultClientRetryPolicy: &river.DefaultClientRetryPolicy{}, + ErrorHandler: &errorHandlerWrapper{ + HandleErrorFunc: func(ctx context.Context, job *rivertype.JobRow, err error) *jobexecutor.ErrorHandlerResult { + resultErr = err + return nil + }, + HandlePanicFunc: func(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *jobexecutor.ErrorHandlerResult { + tb.Fatalf("panic: %v", panicVal) + return nil + }, + }, + InformProducerDoneFunc: func(job *rivertype.JobRow) { close(executionDone) }, + GlobalMiddleware: w.config.WorkerMiddleware, + JobRow: job, + SchedulerInterval: maintenance.JobSchedulerIntervalDefault, + WorkUnit: workUnit, + }) - jobTimeout := valutil.FirstNonZero(w.worker.Timeout(job), w.config.JobTimeout) - ctx, cancel := execution.MaybeApplyTimeout(ctx, jobTimeout) - defer cancel() + executor.Execute(jobCtx) + <-executionDone - return doInner(ctx) + select { + case completerResult := <-subscribeCh: + if len(completerResult) == 0 { + tb.Fatal("test worker internal error: empty job completion received") + } + if len(completerResult) > 1 { + tb.Fatalf("test worker internal error: received %d job completions, expected 1", len(completerResult)) + } + return completerResultToWorkResult(tb, completerResult[0]), resultErr + default: + tb.Fatal("test worker internal error: no job completions received") + } + panic("unreachable") } -var idSeq int64 = 0 //nolint:gochecknoglobals +// WorkResult is the result of working a job in the test Worker. +type WorkResult struct { + // Job is the updated job row from the database _after_ it has been worked. + Job *rivertype.JobRow -func nextID() int64 { - return atomic.AddInt64(&idSeq, 1) + // Kind is the kind of event that occurred following execution. + Kind river.EventKind } -func makeJobFromFactoryBuild[T river.JobArgs](tb testing.TB, args T, jobOpts *testfactory.JobOpts) *river.Job[T] { +func completerResultToWorkResult(tb testing.TB, completerResult jobcompleter.CompleterJobUpdated) WorkResult { tb.Helper() - jobParams := testfactory.Job_Build(tb, jobOpts) - - var errors []rivertype.AttemptError - if jobParams.Errors != nil { - errors = make([]rivertype.AttemptError, len(jobParams.Errors)) - for i, err := range jobParams.Errors { - var attemptError rivertype.AttemptError - if err := json.Unmarshal(err, &attemptError); err != nil { - tb.Fatalf("failed to unmarshal attempt error: %s", err) - } - errors[i] = attemptError - } + var kind river.EventKind + switch completerResult.Job.State { + case rivertype.JobStateCancelled: + kind = river.EventKindJobCancelled + case rivertype.JobStateCompleted: + kind = river.EventKindJobCompleted + case rivertype.JobStateScheduled: + kind = river.EventKindJobSnoozed + case rivertype.JobStateAvailable, rivertype.JobStateDiscarded, rivertype.JobStateRetryable, rivertype.JobStateRunning: + kind = river.EventKindJobFailed + case rivertype.JobStatePending: + panic("test worker internal error: completion subscriber unexpectedly received job in pending state, river bug") + default: + // linter exhaustive rule prevents this from being reached + panic("test worker internal error: unreachable state to distribute, river bug") } - now := time.Now() - - return &river.Job[T]{ - Args: args, - JobRow: &rivertype.JobRow{ - ID: nextID(), - Attempt: jobParams.Attempt, - AttemptedAt: jobParams.AttemptedAt, - AttemptedBy: jobParams.AttemptedBy, - CreatedAt: ptrutil.ValOrDefault(jobParams.CreatedAt, now), - EncodedArgs: jobParams.EncodedArgs, - Errors: errors, - FinalizedAt: jobParams.FinalizedAt, - Kind: jobParams.Kind, - MaxAttempts: jobParams.MaxAttempts, - Metadata: jobParams.Metadata, - Priority: jobParams.Priority, - Queue: jobParams.Queue, - ScheduledAt: ptrutil.ValOrDefault(jobParams.ScheduledAt, now), - State: jobParams.State, - Tags: jobParams.Tags, - UniqueKey: jobParams.UniqueKey, - }, + return WorkResult{ + Job: completerResult.Job, + Kind: kind, } } + +type errorHandlerWrapper struct { + HandleErrorFunc func(ctx context.Context, job *rivertype.JobRow, err error) *jobexecutor.ErrorHandlerResult + HandlePanicFunc func(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *jobexecutor.ErrorHandlerResult +} + +func (h *errorHandlerWrapper) HandleError(ctx context.Context, job *rivertype.JobRow, err error) *jobexecutor.ErrorHandlerResult { + return h.HandleErrorFunc(ctx, job, err) +} + +func (h *errorHandlerWrapper) HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *jobexecutor.ErrorHandlerResult { + return h.HandlePanicFunc(ctx, job, panicVal, trace) +} + +// TODO: move work_unit_wrapper.go so I don't need to copy paste it here: + +// workUnitFactoryWrapper wraps a Worker to implement workUnitFactory. +type workUnitFactoryWrapper[T river.JobArgs] struct { + worker river.Worker[T] +} + +func (w *workUnitFactoryWrapper[T]) MakeUnit(jobRow *rivertype.JobRow) workunit.WorkUnit { + return &wrapperWorkUnit[T]{jobRow: jobRow, worker: w.worker} +} + +// wrapperWorkUnit implements workUnit for a job and Worker. +type wrapperWorkUnit[T river.JobArgs] struct { + job *river.Job[T] // not set until after UnmarshalJob is invoked + jobRow *rivertype.JobRow + worker river.Worker[T] +} + +func (w *wrapperWorkUnit[T]) NextRetry() time.Time { return w.worker.NextRetry(w.job) } +func (w *wrapperWorkUnit[T]) Timeout() time.Duration { return w.worker.Timeout(w.job) } +func (w *wrapperWorkUnit[T]) Work(ctx context.Context) error { return w.worker.Work(ctx, w.job) } + +func (w *wrapperWorkUnit[T]) Middleware() []rivertype.WorkerMiddleware { + return w.worker.Middleware(w.job) +} + +func (w *wrapperWorkUnit[T]) UnmarshalJob() error { + w.job = &river.Job[T]{ + JobRow: w.jobRow, + } + + return json.Unmarshal(w.jobRow.EncodedArgs, &w.job.Args) +} diff --git a/rivertest/worker_test.go b/rivertest/worker_test.go index 7197828a..393604ed 100644 --- a/rivertest/worker_test.go +++ b/rivertest/worker_test.go @@ -13,7 +13,6 @@ import ( "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivershared/riversharedtest" - "github.com/riverqueue/river/rivershared/testfactory" "github.com/riverqueue/river/rivertype" ) @@ -38,7 +37,7 @@ func TestWorker_Work(t *testing.T) { t.Helper() var ( - config = &river.Config{} + config = &river.Config{ID: "rivertest-worker"} driver = riverpgxv5.New(nil) tx = riverinternaltest.TestTx(ctx, t) ) @@ -60,7 +59,7 @@ func TestWorker_Work(t *testing.T) { require.Equal(t, 1, job.JobRow.Attempt) require.NotNil(t, job.JobRow.AttemptedAt) require.WithinDuration(t, time.Now(), *job.JobRow.AttemptedAt, 5*time.Second) - require.Equal(t, []string{"worker1"}, job.JobRow.AttemptedBy) + require.Equal(t, []string{"rivertest-worker"}, job.JobRow.AttemptedBy) require.WithinDuration(t, time.Now(), job.JobRow.CreatedAt, 5*time.Second) require.JSONEq(t, `{"value": "test"}`, string(job.JobRow.EncodedArgs)) require.Empty(t, job.JobRow.Errors) @@ -82,7 +81,9 @@ func TestWorker_Work(t *testing.T) { return nil }) tw := NewWorker(t, bundle.driver, bundle.config, worker) - require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test"}, nil)) + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test"}, nil) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) }) t.Run("Reusable", func(t *testing.T) { @@ -94,8 +95,12 @@ func TestWorker_Work(t *testing.T) { return nil }) tw := NewWorker(t, bundle.driver, bundle.config, worker) - require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test"}, nil)) - require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test2"}, nil)) + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test"}, nil) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) + res, err = tw.Work(ctx, t, bundle.tx, testArgs{Value: "test2"}, nil) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) }) t.Run("SetsCustomInsertOpts", func(t *testing.T) { @@ -110,7 +115,7 @@ func TestWorker_Work(t *testing.T) { require.Equal(t, 1, job.JobRow.Attempt) require.NotNil(t, job.JobRow.AttemptedAt) require.WithinDuration(t, time.Now().UTC(), *job.JobRow.AttemptedAt, 2*time.Second) - require.Equal(t, []string{"worker1"}, job.JobRow.AttemptedBy) + require.Equal(t, []string{"rivertest-worker"}, job.JobRow.AttemptedBy) require.WithinDuration(t, time.Now().UTC(), job.JobRow.CreatedAt, 2*time.Second) require.JSONEq(t, `{"value": "test3"}`, string(job.JobRow.EncodedArgs)) require.Empty(t, job.JobRow.Errors) @@ -130,7 +135,7 @@ func TestWorker_Work(t *testing.T) { tw := NewWorker(t, bundle.driver, bundle.config, worker) // You can also pass in custom insert options: - require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test3"}, &river.InsertOpts{ + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test3"}, &river.InsertOpts{ MaxAttempts: 420, Metadata: []byte(`{"key": "value"}`), Pending: true, // ignored but added to ensure non-default behavior @@ -138,7 +143,9 @@ func TestWorker_Work(t *testing.T) { Queue: "custom_queue", ScheduledAt: hourFromNow, Tags: []string{"tag1", "tag2"}, - })) + }) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) }) t.Run("UniqueOptsAreIgnored", func(t *testing.T) { @@ -159,156 +166,76 @@ func TestWorker_Work(t *testing.T) { return nil }) tw := NewWorker(t, bundle.driver, bundle.config, worker) - require.NoError(t, tw.Work(context.Background(), t, bundle.tx, testArgs{Value: "test"}, &river.InsertOpts{ + + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test"}, &river.InsertOpts{ UniqueOpts: river.UniqueOpts{ByPeriod: 1 * time.Hour}, - })) + }) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) }) -} - -func TestWorker_WorkTx(t *testing.T) { - t.Parallel() - - ctx := context.Background() - - type testBundle struct { - client *river.Client[pgx.Tx] - driver *riverpgxv5.Driver - tx pgx.Tx - workFunc func(ctx context.Context, job *river.Job[testArgs]) error - } - - setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { - t.Helper() - - var ( - config = &river.Config{} - driver = riverpgxv5.New(nil) - ) - client, err := river.NewClient(driver, config) - require.NoError(t, err) + t.Run("ReturnsASnoozeEventKindWhenSnoozed", func(t *testing.T) { + t.Parallel() - bundle := &testBundle{ - client: client, - driver: driver, - tx: riverinternaltest.TestTx(ctx, t), - workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, - } + bundle := setup(t) worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { - return bundle.workFunc(ctx, job) + return river.JobSnooze(time.Hour) }) + tw := NewWorker(t, bundle.driver, bundle.config, worker) - return NewWorker(t, driver, config, worker), bundle - } - - t.Run("Success", func(t *testing.T) { - t.Parallel() - - testWorker, bundle := setup(t) - - bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { - require.Equal(t, rivertype.JobStateRunning, job.State) - return nil - } - - require.NoError(t, testWorker.WorkTx(ctx, t, bundle.tx, testArgs{}, nil)) + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test"}, nil) + require.NoError(t, err) + require.Equal(t, river.EventKindJobSnoozed, res.Kind) }) -} - -func TestWorker_WorkJob(t *testing.T) { - t.Parallel() - - ctx := context.Background() - - type testBundle struct { - client *river.Client[pgx.Tx] - driver *riverpgxv5.Driver - tx pgx.Tx - workFunc func(ctx context.Context, job *river.Job[testArgs]) error - } - - setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { - t.Helper() - - var ( - config = &river.Config{} - driver = riverpgxv5.New(nil) - ) - client, err := river.NewClient(driver, config) - require.NoError(t, err) + t.Run("ReturnsACancelEventKindWhenCancelled", func(t *testing.T) { + t.Parallel() - bundle := &testBundle{ - client: client, - driver: driver, - tx: riverinternaltest.TestTx(ctx, t), - workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, - } + bundle := setup(t) worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { - return bundle.workFunc(ctx, job) + return river.JobCancel(nil) }) + tw := NewWorker(t, bundle.driver, bundle.config, worker) - return NewWorker(t, driver, config, worker), bundle - } - - t.Run("Success", func(t *testing.T) { - t.Parallel() - - testWorker, bundle := setup(t) - - bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { - require.Equal(t, []string{"worker123"}, job.JobRow.AttemptedBy) - require.Equal(t, rivertype.JobStateRunning, job.State) - return nil - } - - now := time.Now() - require.NoError(t, testWorker.WorkJob(ctx, t, bundle.tx, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{ - AttemptedAt: &now, - AttemptedBy: []string{"worker123"}, - CreatedAt: &now, - EncodedArgs: []byte(`{"value": "test"}`), - Errors: nil, - }))) + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test"}, nil) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCancelled, res.Kind) }) - t.Run("JobCompleteTxWithInsertedJobRow", func(t *testing.T) { + t.Run("UsesACustomClockWhenProvided", func(t *testing.T) { t.Parallel() - testWorker, bundle := setup(t) - - args := testArgs{} - insertRes, err := bundle.client.InsertTx(ctx, bundle.tx, args, nil) - require.NoError(t, err) - - bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { - updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) - require.NoError(t, err) - require.Equal(t, rivertype.JobStateRunning, updatedJob.State) - - _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, bundle.tx, job) - require.NoError(t, err) + bundle := setup(t) + hourFromNow := time.Now().UTC().Add(1 * time.Hour) + timeStub := &TimeStub{} + timeStub.StubNowUTC(hourFromNow) + bundle.config.Test.Time = timeStub + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + require.WithinDuration(t, hourFromNow, *job.JobRow.AttemptedAt, time.Millisecond) + require.WithinDuration(t, hourFromNow, job.JobRow.CreatedAt, time.Millisecond) + require.WithinDuration(t, hourFromNow, job.JobRow.ScheduledAt, time.Millisecond) return nil - } - - require.NoError(t, testWorker.WorkJob(ctx, t, bundle.tx, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + }) + tw := NewWorker(t, bundle.driver, bundle.config, worker) - updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) + res, err := tw.Work(ctx, t, bundle.tx, testArgs{Value: "test"}, nil) require.NoError(t, err) - require.Equal(t, rivertype.JobStateCompleted, updatedJob.State) + require.Equal(t, river.EventKindJobCompleted, res.Kind) + require.WithinDuration(t, hourFromNow, *res.Job.FinalizedAt, time.Millisecond) }) } -func TestWorker_WorkJobTx(t *testing.T) { +func TestWorker_WorkJob(t *testing.T) { t.Parallel() ctx := context.Background() type testBundle struct { client *river.Client[pgx.Tx] + config *river.Config driver *riverpgxv5.Driver tx pgx.Tx workFunc func(ctx context.Context, job *river.Job[testArgs]) error @@ -318,7 +245,7 @@ func TestWorker_WorkJobTx(t *testing.T) { t.Helper() var ( - config = &river.Config{} + config = &river.Config{ID: "rivertest-workjob"} driver = riverpgxv5.New(nil) ) @@ -327,6 +254,7 @@ func TestWorker_WorkJobTx(t *testing.T) { bundle := &testBundle{ client: client, + config: config, driver: driver, tx: riverinternaltest.TestTx(ctx, t), workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, @@ -345,11 +273,18 @@ func TestWorker_WorkJobTx(t *testing.T) { testWorker, bundle := setup(t) bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.WithinDuration(t, time.Now(), *job.JobRow.AttemptedAt, 5*time.Second) + require.Equal(t, []string{"rivertest-workjob"}, job.JobRow.AttemptedBy) require.Equal(t, rivertype.JobStateRunning, job.State) return nil } - require.NoError(t, testWorker.WorkJobTx(ctx, t, bundle.tx, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{}))) + insertRes, err := bundle.client.InsertTx(ctx, bundle.tx, testArgs{}, nil) + require.NoError(t, err) + + res, err := testWorker.WorkJob(ctx, t, bundle.tx, insertRes.Job) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) }) t.Run("JobCompleteTxWithInsertedJobRow", func(t *testing.T) { @@ -372,7 +307,9 @@ func TestWorker_WorkJobTx(t *testing.T) { return nil } - require.NoError(t, testWorker.WorkJobTx(ctx, t, bundle.tx, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + res, err := testWorker.WorkJob(ctx, t, bundle.tx, insertRes.Job) + require.NoError(t, err) + require.Equal(t, river.EventKindJobCompleted, res.Kind) updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) require.NoError(t, err)