|
| 1 | +package river_test |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "errors" |
| 6 | + "fmt" |
| 7 | + "log/slog" |
| 8 | + "os" |
| 9 | + |
| 10 | + "github.com/jackc/pgx/v5/pgxpool" |
| 11 | + |
| 12 | + "github.com/riverqueue/river" |
| 13 | + "github.com/riverqueue/river/riverdbtest" |
| 14 | + "github.com/riverqueue/river/riverdriver/riverpgxv5" |
| 15 | + "github.com/riverqueue/river/rivershared/riversharedtest" |
| 16 | + "github.com/riverqueue/river/rivershared/util/slogutil" |
| 17 | + "github.com/riverqueue/river/rivershared/util/testutil" |
| 18 | +) |
| 19 | + |
| 20 | +// UploadCancelArgs represents work on a user-initiated upload that the |
| 21 | +// worker may discover to be unrecoverable. A separate "uploads" row is |
| 22 | +// created by the web app in the same transaction that enqueues this job; |
| 23 | +// the job keeps the external status in sync with River's view of |
| 24 | +// the job's outcome. |
| 25 | +type UploadCancelArgs struct { |
| 26 | + UploadID int64 `json:"upload_id"` |
| 27 | + // Valid is used to simulate validation outcomes. When false, the |
| 28 | + // worker treats the upload as permanently unprocessable and cancels |
| 29 | + // the job rather than retrying. |
| 30 | + Valid bool `json:"valid"` |
| 31 | +} |
| 32 | + |
| 33 | +func (UploadCancelArgs) Kind() string { return "example_upload_cancel_worker" } |
| 34 | + |
| 35 | +// UploadCancelWorker processes an upload. If the upload is permanently |
| 36 | +// invalid, the worker uses JobCancelTx to record the failure: the job will |
| 37 | +// not be retried regardless of MaxAttempts. The external "uploads" row's |
| 38 | +// status is promoted to "cancelled" in the same transaction so the web |
| 39 | +// app sees a consistent final state. |
| 40 | +type UploadCancelWorker struct { |
| 41 | + river.WorkerDefaults[UploadCancelArgs] |
| 42 | + |
| 43 | + dbPool *pgxpool.Pool |
| 44 | + schema string |
| 45 | +} |
| 46 | + |
| 47 | +func (w *UploadCancelWorker) Work(ctx context.Context, job *river.Job[UploadCancelArgs]) error { |
| 48 | + tx, err := w.dbPool.Begin(ctx) |
| 49 | + if err != nil { |
| 50 | + // return err #1: could not begin a tx. Nothing has been done yet; |
| 51 | + // River will retry per its policy. |
| 52 | + return err |
| 53 | + } |
| 54 | + defer tx.Rollback(ctx) |
| 55 | + |
| 56 | + // Increment the external "attempts" counter so the web app can see |
| 57 | + // that the worker ran. |
| 58 | + if _, err := tx.Exec(ctx, |
| 59 | + fmt.Sprintf(`UPDATE %q.uploads SET attempts = attempts + 1 WHERE id = $1`, w.schema), |
| 60 | + job.Args.UploadID, |
| 61 | + ); err != nil { |
| 62 | + // return err #2: external UPDATE failed; tx rolls back. River retries. |
| 63 | + return err |
| 64 | + } |
| 65 | + |
| 66 | + // Perform validation. If the upload is unrecoverably invalid there's |
| 67 | + // no point retrying - cancel the job instead. |
| 68 | + if !job.Args.Valid { |
| 69 | + validationErr := errors.New("upload rejected: invalid content") |
| 70 | + |
| 71 | + if _, err := tx.Exec(ctx, |
| 72 | + fmt.Sprintf(`UPDATE %q.uploads SET status = 'cancelled', last_error = $1 WHERE id = $2`, w.schema), |
| 73 | + validationErr.Error(), job.Args.UploadID, |
| 74 | + ); err != nil { |
| 75 | + // return err #3: the external "mark cancelled" UPDATE failed. |
| 76 | + // Tx rolls back with everything in it. River will retry the |
| 77 | + // job normally - which on the next run will hit the same |
| 78 | + // validation error and (if we then reach JobCancelTx |
| 79 | + // successfully) cancel for real. |
| 80 | + return err |
| 81 | + } |
| 82 | + |
| 83 | + if _, err := river.JobCancelTx[*riverpgxv5.Driver](ctx, tx, job, validationErr); err != nil { |
| 84 | + // return err #4: JobCancelTx itself errored. Tx rolls back. |
| 85 | + return err |
| 86 | + } |
| 87 | + |
| 88 | + if err := tx.Commit(ctx); err != nil { |
| 89 | + // return err #5: commit failed. Nothing persisted. River retries. |
| 90 | + return err |
| 91 | + } |
| 92 | + |
| 93 | + // return validationErr #6: the cancellation is recorded and the |
| 94 | + // external row is in sync. Return the error up so that any |
| 95 | + // registered HookWorkEnd hooks (and ErrorHandler) still see it |
| 96 | + // and can log/trace it or add metadata via river.RecordOutput. |
| 97 | + // Metadata merges on the executor's subsequent UPDATE are *not* |
| 98 | + // guarded by state, so those land asynchronously after our tx |
| 99 | + // commits. Hooks and handlers cannot steer the job's outcome at |
| 100 | + // this point: the state change and any new AttemptError are |
| 101 | + // guarded by state = 'running' in the SQL and won't apply now |
| 102 | + // that JobCancelTx has moved the row out of that state. |
| 103 | + return validationErr |
| 104 | + } |
| 105 | + |
| 106 | + // Normal success path omitted for brevity; a real worker would do its |
| 107 | + // work here, optionally call JobCompleteTx to finalize the job within |
| 108 | + // this same tx, then commit. |
| 109 | + return errors.New("success path not shown in this example") |
| 110 | +} |
| 111 | + |
| 112 | +// Example_cancelJobWithinTx demonstrates how to transactionally cancel a |
| 113 | +// job when the worker determines that further retries are pointless, |
| 114 | +// while keeping an external status row in sync. |
| 115 | +func Example_cancelJobWithinTx() { |
| 116 | + ctx := context.Background() |
| 117 | + |
| 118 | + dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL()) |
| 119 | + if err != nil { |
| 120 | + panic(err) |
| 121 | + } |
| 122 | + defer dbPool.Close() |
| 123 | + |
| 124 | + schema := riverdbtest.TestSchema(ctx, testutil.PanicTB(), riverpgxv5.New(dbPool), nil) |
| 125 | + |
| 126 | + if _, err := dbPool.Exec(ctx, fmt.Sprintf(` |
| 127 | + CREATE TABLE %q.uploads ( |
| 128 | + id bigint PRIMARY KEY, |
| 129 | + status text NOT NULL, |
| 130 | + attempts int NOT NULL DEFAULT 0, |
| 131 | + last_error text |
| 132 | + )`, schema)); err != nil { |
| 133 | + panic(err) |
| 134 | + } |
| 135 | + if _, err := dbPool.Exec(ctx, fmt.Sprintf(`INSERT INTO %q.uploads (id, status) VALUES (99, 'working')`, schema)); err != nil { |
| 136 | + panic(err) |
| 137 | + } |
| 138 | + |
| 139 | + workers := river.NewWorkers() |
| 140 | + river.AddWorker(workers, &UploadCancelWorker{dbPool: dbPool, schema: schema}) |
| 141 | + |
| 142 | + riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{ |
| 143 | + Logger: slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelWarn, ReplaceAttr: slogutil.NoLevelTime})), |
| 144 | + Queues: map[string]river.QueueConfig{ |
| 145 | + river.QueueDefault: {MaxWorkers: 100}, |
| 146 | + }, |
| 147 | + Schema: schema, |
| 148 | + TestOnly: true, |
| 149 | + Workers: workers, |
| 150 | + }) |
| 151 | + if err != nil { |
| 152 | + panic(err) |
| 153 | + } |
| 154 | + |
| 155 | + // Subscribe to the cancelled-event stream so the example can wait for |
| 156 | + // the job to be fully processed. |
| 157 | + subscribeChan, subscribeCancel := riverClient.Subscribe(river.EventKindJobCancelled) |
| 158 | + defer subscribeCancel() |
| 159 | + |
| 160 | + if err := riverClient.Start(ctx); err != nil { |
| 161 | + panic(err) |
| 162 | + } |
| 163 | + |
| 164 | + // MaxAttempts=5 to make clear the job stops on the first run because |
| 165 | + // of the cancel - attempts remaining doesn't matter. |
| 166 | + if _, err := riverClient.Insert(ctx, UploadCancelArgs{UploadID: 99, Valid: false}, &river.InsertOpts{MaxAttempts: 5}); err != nil { |
| 167 | + panic(err) |
| 168 | + } |
| 169 | + |
| 170 | + riversharedtest.WaitOrTimeoutN(testutil.PanicTB(), subscribeChan, 1) |
| 171 | + |
| 172 | + if err := riverClient.Stop(ctx); err != nil { |
| 173 | + panic(err) |
| 174 | + } |
| 175 | + |
| 176 | + var ( |
| 177 | + status string |
| 178 | + attempts int |
| 179 | + lastError *string |
| 180 | + ) |
| 181 | + if err := dbPool.QueryRow(ctx, |
| 182 | + fmt.Sprintf(`SELECT status, attempts, last_error FROM %q.uploads WHERE id = 99`, schema), |
| 183 | + ).Scan(&status, &attempts, &lastError); err != nil { |
| 184 | + panic(err) |
| 185 | + } |
| 186 | + lastErrTxt := "" |
| 187 | + if lastError != nil { |
| 188 | + lastErrTxt = *lastError |
| 189 | + } |
| 190 | + fmt.Printf("upload 99: status=%q attempts=%d last_error=%q\n", status, attempts, lastErrTxt) |
| 191 | + |
| 192 | + // Output: |
| 193 | + // upload 99: status="cancelled" attempts=1 last_error="upload rejected: invalid content" |
| 194 | +} |
0 commit comments