diff --git a/client.go b/client.go index bf07bfc0f..e1db8e494 100644 --- a/client.go +++ b/client.go @@ -57,6 +57,8 @@ const ( QueueNumWorkersMax = 10_000 ) +const clientInitialPingTimeout = 10 * time.Second + var ( postgresSchemaNameRE = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`) @@ -1058,9 +1060,26 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client client.testSignals.queueMaintainerLeader = &client.queueMaintainerLeader.TestSignals } + // Ping asynchronously so that constructing a client remains non-blocking. + // PostgreSQL drivers use this ping to select and cache their unique insert + // strategy. The first unique insert uses the same resolver, so it can perform + // initialization instead if it reaches the database first. + if driver.DatabaseName() == riverdriver.DatabaseNamePostgres && driver.PoolIsSet() { + go client.initialPing() + } + return client, nil } +func (c *Client[TTx]) initialPing() { + ctx, cancel := context.WithTimeout(context.Background(), clientInitialPingTimeout) + defer cancel() + + if err := c.driver.GetExecutor().Ping(ctx); err != nil { + c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": Initial database ping failed", slog.String("err", err.Error())) + } +} + // Start starts the client's job fetching and working loops. Once this is called, // the client will run in a background goroutine until stopped. All jobs are // run with a context inheriting from the provided context, but with a timeout @@ -1110,7 +1129,7 @@ func (c *Client[TTx]) Start(ctx context.Context) error { // available, the client appears to have started even though it's completely // non-functional. Here we try to make an initial assessment of health and // return quickly in case of an apparent problem. - if err := c.driver.GetExecutor().Exec(fetchCtx, "SELECT 1"); err != nil { + if err := c.driver.GetExecutor().Ping(fetchCtx); err != nil { return fmt.Errorf("error making initial connection to database: %w", err) } diff --git a/client_test.go b/client_test.go index bb7ba8150..50741fa58 100644 --- a/client_test.go +++ b/client_test.go @@ -49,6 +49,36 @@ import ( "github.com/riverqueue/river/rivertype" ) +type clientInitialPingTestDriver struct { + riverdriver.Driver[pgx.Tx] + + executor riverdriver.Executor +} + +func (d *clientInitialPingTestDriver) GetExecutor() riverdriver.Executor { return d.executor } + +type clientInitialPingTestExecutor struct { + riverdriver.Executor + + pingFunc func(context.Context) error +} + +func (e *clientInitialPingTestExecutor) Ping(ctx context.Context) error { return e.pingFunc(ctx) } + +type clientInitialPingTestSignals struct { + PingContinue testsignal.TestSignal[struct{}] + PingFinished testsignal.TestSignal[error] + PingStarted testsignal.TestSignal[struct{}] +} + +func (s *clientInitialPingTestSignals) Init(tb testing.TB) { + tb.Helper() + + s.PingContinue.Init(tb) + s.PingFinished.Init(tb) + s.PingStarted.Init(tb) +} + type clientJobCancelTestSignals struct { ContinueWork testsignal.TestSignal[struct{}] JobStarted testsignal.TestSignal[int64] @@ -8491,6 +8521,71 @@ func Test_NewClient_Defaults(t *testing.T) { require.IsType(t, &baseservice.UnStubbableTimeGenerator{}, client.baseService.Time) } +func Test_NewClient_InitialPing(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + client *Client[pgx.Tx] + signals *clientInitialPingTestSignals + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + var signals clientInitialPingTestSignals + signals.Init(t) + + baseDriver := riverpgxv5.New(riversharedtest.DBPool(ctx, t)) + schema := riverdbtest.TestSchema(ctx, t, baseDriver, nil) + baseExecutor := baseDriver.GetExecutor() + driver := &clientInitialPingTestDriver{ + Driver: baseDriver, + executor: &clientInitialPingTestExecutor{ + Executor: baseExecutor, + pingFunc: func(ctx context.Context) error { + signals.PingStarted.Signal(struct{}{}) + signals.PingContinue.WaitOrTimeout() + + err := baseExecutor.Ping(ctx) + signals.PingFinished.Signal(err) + return err + }, + }, + } + + client, err := NewClient(driver, &Config{Schema: schema}) + require.NoError(t, err) + t.Cleanup(func() { signals.PingContinue.Signal(struct{}{}) }) + + return &testBundle{client: client, signals: &signals} + } + + t.Run("InsertCanResolveBeforePing", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.signals.PingStarted.WaitOrTimeout() + + _, err := bundle.client.Insert(ctx, noOpArgs{Name: "first"}, &InsertOpts{UniqueOpts: UniqueOpts{ByArgs: true}}) + require.NoError(t, err) + + bundle.signals.PingContinue.Signal(struct{}{}) + require.NoError(t, bundle.signals.PingFinished.WaitOrTimeout()) + }) + + t.Run("PingResolvesWithoutInsert", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.signals.PingStarted.WaitOrTimeout() + + bundle.signals.PingContinue.Signal(struct{}{}) + require.NoError(t, bundle.signals.PingFinished.WaitOrTimeout()) + }) +} + func Test_NewClient_Overrides(t *testing.T) { t.Parallel() diff --git a/internal/rivercommon/river_common.go b/internal/rivercommon/river_common.go index 4f769b391..389efdc8a 100644 --- a/internal/rivercommon/river_common.go +++ b/internal/rivercommon/river_common.go @@ -45,11 +45,6 @@ const ( // MetadataKeyRescueCount records how many times the job has been rescued. MetadataKeyRescueCount = "river:rescue_count" - - // MetadataKeyUniqueNonce is a special metadata key used by the SQLite driver to - // determine whether an upsert is was skipped or not because the `(xmax != 0)` - // trick we use in Postgres doesn't work in SQLite. - MetadataKeyUniqueNonce = "river:unique_nonce" ) type ContextKeyClient struct{} diff --git a/riverdriver/river_driver_interface.go b/riverdriver/river_driver_interface.go index d4b87f315..b639e3e9f 100644 --- a/riverdriver/river_driver_interface.go +++ b/riverdriver/river_driver_interface.go @@ -289,6 +289,11 @@ type Executor interface { NotificationDeleteBefore(ctx context.Context, params *NotificationDeleteBeforeParams) (int, error) NotifyMany(ctx context.Context, params *NotifyManyParams) error + + // Ping checks that the database is reachable and performs any driver-specific + // initialization that should happen before the first operation. + Ping(ctx context.Context) error + PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) QueueCreateOrSetUpdatedAt(ctx context.Context, params *QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go index e2542e763..b4be68757 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go @@ -21,6 +21,24 @@ func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) er return err } +const pGGetProductAndVersion = `-- name: PGGetProductAndVersion :one +SELECT + version()::text AS product, + current_setting('server_version_num')::int AS version_num +` + +type PGGetProductAndVersionRow struct { + Product string + VersionNum int32 +} + +func (q *Queries) PGGetProductAndVersion(ctx context.Context, db DBTX) (*PGGetProductAndVersionRow, error) { + row := db.QueryRowContext(ctx, pGGetProductAndVersion) + var i PGGetProductAndVersionRow + err := row.Scan(&i.Product, &i.VersionNum) + return &i, err +} + const pGNotifyMany = `-- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index a197f7bb5..e7361009e 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -718,7 +718,9 @@ ON CONFLICT (unique_key) AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) -- Something needs to be updated for a row to be returned on a conflict. DO UPDATE SET kind = EXCLUDED.kind -RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, (xmax != 0) AS unique_skipped_as_duplicate +RETURNING + river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, + /* TEMPLATE_BEGIN: unique_skipped_as_duplicate */ (xmax != 0) /* TEMPLATE_END */ AS unique_skipped_as_duplicate ` type JobInsertFastManyParams struct { diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index b672c77d0..37c186e1b 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -16,6 +16,8 @@ import ( "io/fs" "math" "strings" + "sync" + "sync/atomic" "time" "github.com/jackc/pgx/v5/pgxpool" @@ -28,6 +30,7 @@ import ( "github.com/riverqueue/river/rivershared/uniquestates" "github.com/riverqueue/river/rivershared/util/dbutil" "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/savepointutil" "github.com/riverqueue/river/rivershared/util/sliceutil" "github.com/riverqueue/river/rivertype" @@ -38,9 +41,11 @@ var migrationFS embed.FS // Driver is an implementation of riverdriver.Driver for database/sql. type Driver struct { - dbPool *sql.DB - listenerDriver *riverpgxv5.Driver - replacer sqlctemplate.Replacer + dbPool *sql.DB + listenerDriver *riverpgxv5.Driver + replacer sqlctemplate.Replacer + uniqueInsertMode atomic.Uint32 + uniqueInsertModeInitMu sync.Mutex } // New returns a new database/sql River driver for use with River. @@ -402,6 +407,16 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt } func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + uniqueInsertMode, err := e.uniqueInsertMode(ctx) + if err != nil { + return nil, err + } + + var uniqueNonce string + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueNonce = randutil.Hex(8) + } + insertJobsParams := &dbsqlc.JobInsertFastManyParams{ ID: make([]int64, len(params.Jobs)), Args: make([]string, len(params.Jobs)), @@ -442,7 +457,16 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.CreatedAt[i] = createdAt insertJobsParams.Kind[i] = params.Kind insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec - insertJobsParams.Metadata[i] = cmp.Or(string(params.Metadata), "{}") + metadata := []byte(cmp.Or(string(params.Metadata), "{}")) + if uniqueNonce != "" { + var err error + metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce) + if err != nil { + return nil, err + } + } + + insertJobsParams.Metadata[i] = string(metadata) insertJobsParams.Priority[i] = int16(min(params.Priority, math.MaxInt16)) //nolint:gosec insertJobsParams.Queue[i] = params.Queue insertJobsParams.ScheduledAt[i] = scheduledAt @@ -452,6 +476,9 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.UniqueStates[i] = int32(params.UniqueStates) } + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "unique_skipped_as_duplicate": {Value: uniqueInsertMode.SQL(), Stable: true}, + }, nil) items, err := dbsqlc.New().JobInsertFastMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams) if err != nil { return nil, interpretError(err) @@ -462,7 +489,13 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo if err != nil { return nil, err } - return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: row.UniqueSkippedAsDuplicate}, nil + + uniqueSkippedAsDuplicate := row.UniqueSkippedAsDuplicate + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueSkippedAsDuplicate = riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce) + } + + return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: uniqueSkippedAsDuplicate}, nil }) } @@ -932,6 +965,11 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan }) } +func (e *Executor) Ping(ctx context.Context) error { + _, err := e.uniqueInsertMode(ctx) + return err +} + func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { err := dbsqlc.New().PGAdvisoryXactLock(ctx, e.dbtx, key) return &struct{}{}, interpretError(err) @@ -1090,6 +1128,32 @@ func (e *Executor) TableTruncate(ctx context.Context, params *riverdriver.TableT return interpretError(err) } +func (e *Executor) uniqueInsertMode(ctx context.Context) (riverdriver.UniqueInsertMode, error) { + if e.driver != nil { + if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown { + return mode, nil + } + + e.driver.uniqueInsertModeInitMu.Lock() + defer e.driver.uniqueInsertModeInitMu.Unlock() + + if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown { + return mode, nil + } + } + + productAndVersion, err := dbsqlc.New().PGGetProductAndVersion(ctx, e.dbtx) + if err != nil { + return riverdriver.UniqueInsertModeUnknown, interpretError(err) + } + + mode := riverdriver.UniqueInsertModeFromProductAndVersion(productAndVersion.Product, productAndVersion.VersionNum) + if e.driver != nil { + e.driver.uniqueInsertMode.Store(uint32(mode)) + } + return mode, nil +} + type ExecutorTx struct { Executor diff --git a/riverdriver/riverdrivertest/job_insert.go b/riverdriver/riverdrivertest/job_insert.go index 6b625a0e5..9bc8cf64a 100644 --- a/riverdriver/riverdrivertest/job_insert.go +++ b/riverdriver/riverdrivertest/job_insert.go @@ -88,7 +88,7 @@ func exerciseJobInsert[TTx any](ctx context.Context, t *testing.T, // SQLite needs to set a special metadata key to be able to // check for duplicates. Remove this for purposes of comparing // inserted metadata. - job.Metadata, err = sjson.DeleteBytes(job.Metadata, rivercommon.MetadataKeyUniqueNonce) + job.Metadata, err = sjson.DeleteBytes(job.Metadata, riverdriver.UniqueInsertMetadataKey) require.NoError(t, err) require.Equal(t, idStart+int64(i), job.ID) diff --git a/riverdriver/riverdrivertest/job_read.go b/riverdriver/riverdrivertest/job_read.go index f8fe44182..eaa090827 100644 --- a/riverdriver/riverdrivertest/job_read.go +++ b/riverdriver/riverdrivertest/job_read.go @@ -55,7 +55,7 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx for _, state := range rivertype.JobStates() { require.Contains(t, countsByState, state) - switch state { //nolint:exhaustive + switch state { case rivertype.JobStateAvailable: require.Equal(t, 2, countsByState[state]) case rivertype.JobStateCancelled: @@ -64,8 +64,10 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx require.Equal(t, 1, countsByState[state]) case rivertype.JobStateDiscarded: require.Equal(t, 1, countsByState[state]) - default: + case rivertype.JobStatePending, rivertype.JobStateRetryable, rivertype.JobStateRunning, rivertype.JobStateScheduled: require.Equal(t, 0, countsByState[state]) + default: + require.FailNow(t, "unknown job state", state) } } }) diff --git a/riverdriver/riverdrivertest/riverdrivertest.go b/riverdriver/riverdrivertest/riverdrivertest.go index aa66762c7..b85afb9a4 100644 --- a/riverdriver/riverdrivertest/riverdrivertest.go +++ b/riverdriver/riverdrivertest/riverdrivertest.go @@ -54,6 +54,13 @@ func exerciseDriverPool[TTx any](ctx context.Context, t *testing.T, ) { t.Helper() + t.Run("Ping", func(t *testing.T) { + t.Parallel() + + exec, _ := executorWithTx(ctx, t) + require.NoError(t, exec.Ping(ctx)) + }) + t.Run("PoolIsSet", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql index 19a7b99f6..02f510b2f 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql @@ -1,6 +1,11 @@ -- name: PGAdvisoryXactLock :exec SELECT pg_advisory_xact_lock(@key); +-- name: PGGetProductAndVersion :one +SELECT + version()::text AS product, + current_setting('server_version_num')::int AS version_num; + -- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go index 9215c089b..e1ef0917b 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go @@ -20,6 +20,24 @@ func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) er return err } +const pGGetProductAndVersion = `-- name: PGGetProductAndVersion :one +SELECT + version()::text AS product, + current_setting('server_version_num')::int AS version_num +` + +type PGGetProductAndVersionRow struct { + Product string + VersionNum int32 +} + +func (q *Queries) PGGetProductAndVersion(ctx context.Context, db DBTX) (*PGGetProductAndVersionRow, error) { + row := db.QueryRow(ctx, pGGetProductAndVersion) + var i PGGetProductAndVersionRow + err := row.Scan(&i.Product, &i.VersionNum) + return &i, err +} + const pGNotifyMany = `-- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index 40098ed07..ca953ae57 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -318,7 +318,9 @@ ON CONFLICT (unique_key) AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) -- Something needs to be updated for a row to be returned on a conflict. DO UPDATE SET kind = EXCLUDED.kind -RETURNING sqlc.embed(river_job), (xmax != 0) AS unique_skipped_as_duplicate; +RETURNING + sqlc.embed(river_job), + /* TEMPLATE_BEGIN: unique_skipped_as_duplicate */ (xmax != 0) /* TEMPLATE_END */ AS unique_skipped_as_duplicate; -- name: JobInsertFastManyNoReturning :execrows INSERT INTO /* TEMPLATE: schema */river_job( diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go index 7eb082c52..0043759b1 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -697,7 +697,9 @@ ON CONFLICT (unique_key) AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) -- Something needs to be updated for a row to be returned on a conflict. DO UPDATE SET kind = EXCLUDED.kind -RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, (xmax != 0) AS unique_skipped_as_duplicate +RETURNING + river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, + /* TEMPLATE_BEGIN: unique_skipped_as_duplicate */ (xmax != 0) /* TEMPLATE_END */ AS unique_skipped_as_duplicate ` type JobInsertFastManyParams struct { diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/riverdriver/riverpgxv5/river_pgx_v5_driver.go index 34954fac0..b742b8082 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -16,6 +16,7 @@ import ( "math" "strings" "sync" + "sync/atomic" "time" "github.com/jackc/pgx/v5" @@ -30,6 +31,7 @@ import ( "github.com/riverqueue/river/rivershared/uniquestates" "github.com/riverqueue/river/rivershared/util/dbutil" "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/sliceutil" "github.com/riverqueue/river/rivertype" ) @@ -39,8 +41,10 @@ var migrationFS embed.FS // Driver is an implementation of riverdriver.Driver for Pgx v5. type Driver struct { - dbPool *pgxpool.Pool - replacer sqlctemplate.Replacer + dbPool *pgxpool.Pool + replacer sqlctemplate.Replacer + uniqueInsertMode atomic.Uint32 + uniqueInsertModeInitMu sync.Mutex } // New returns a new Pgx v5 River driver for use with River. @@ -367,6 +371,16 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt } func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + uniqueInsertMode, err := e.uniqueInsertMode(ctx) + if err != nil { + return nil, err + } + + var uniqueNonce string + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueNonce = randutil.Hex(8) + } + insertJobsParams := &dbsqlc.JobInsertFastManyParams{ ID: make([]int64, len(params.Jobs)), Args: make([][]byte, len(params.Jobs)), @@ -408,7 +422,16 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo 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) + metadata := sliceutil.FirstNonEmpty(params.Metadata, defaultObject) + if uniqueNonce != "" { + var err error + metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce) + if err != nil { + return nil, err + } + } + + insertJobsParams.Metadata[i] = metadata insertJobsParams.Priority[i] = int16(min(params.Priority, math.MaxInt16)) //nolint:gosec insertJobsParams.Queue[i] = params.Queue insertJobsParams.ScheduledAt[i] = scheduledAt @@ -418,6 +441,9 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.UniqueStates[i] = int32(params.UniqueStates) } + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "unique_skipped_as_duplicate": {Value: uniqueInsertMode.SQL(), Stable: true}, + }, nil) items, err := dbsqlc.New().JobInsertFastMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams) if err != nil { return nil, interpretError(err) @@ -428,7 +454,13 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo if err != nil { return nil, err } - return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: row.UniqueSkippedAsDuplicate}, nil + + uniqueSkippedAsDuplicate := row.UniqueSkippedAsDuplicate + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueSkippedAsDuplicate = riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce) + } + + return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: uniqueSkippedAsDuplicate}, nil }) } @@ -878,6 +910,11 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan }) } +func (e *Executor) Ping(ctx context.Context) error { + _, err := e.uniqueInsertMode(ctx) + return err +} + func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { err := dbsqlc.New().PGAdvisoryXactLock(ctx, e.dbtx, key) return &struct{}{}, interpretError(err) @@ -1036,6 +1073,32 @@ func (e *Executor) TableTruncate(ctx context.Context, params *riverdriver.TableT return interpretError(err) } +func (e *Executor) uniqueInsertMode(ctx context.Context) (riverdriver.UniqueInsertMode, error) { + if e.driver != nil { + if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown { + return mode, nil + } + + e.driver.uniqueInsertModeInitMu.Lock() + defer e.driver.uniqueInsertModeInitMu.Unlock() + + if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown { + return mode, nil + } + } + + productAndVersion, err := dbsqlc.New().PGGetProductAndVersion(ctx, e.dbtx) + if err != nil { + return riverdriver.UniqueInsertModeUnknown, interpretError(err) + } + + mode := riverdriver.UniqueInsertModeFromProductAndVersion(productAndVersion.Product, productAndVersion.VersionNum) + if e.driver != nil { + e.driver.uniqueInsertMode.Store(uint32(mode)) + } + return mode, nil +} + type ExecutorTx struct { Executor diff --git a/riverdriver/riversqlite/go.mod b/riverdriver/riversqlite/go.mod index 7af5377bb..23be1c865 100644 --- a/riverdriver/riversqlite/go.mod +++ b/riverdriver/riversqlite/go.mod @@ -10,12 +10,6 @@ require ( github.com/riverqueue/river/rivershared v0.46.0 github.com/riverqueue/river/rivertype v0.46.0 github.com/stretchr/testify v1.12.1 - github.com/tidwall/gjson v1.19.0 - github.com/tidwall/sjson v1.2.5 ) -require ( - github.com/tidwall/match v1.2.0 // indirect - github.com/tidwall/pretty v1.2.1 // indirect - go.yaml.in/yaml/v3 v3.0.5 // indirect -) +require go.yaml.in/yaml/v3 v3.0.5 // indirect diff --git a/riverdriver/riversqlite/go.sum b/riverdriver/riversqlite/go.sum index f66ab1c50..1b7d4e9f5 100644 --- a/riverdriver/riversqlite/go.sum +++ b/riverdriver/riversqlite/go.sum @@ -18,17 +18,6 @@ github.com/riverqueue/river/rivertype v0.44.1 h1:G9UnsBJhlndreMPdjHSoY8wp+OTAsDT github.com/riverqueue/river/rivertype v0.44.1/go.mod h1:D1Ad+EaZiaXbQbJcJcfeicXJMBKno0n6UcfKI5Q7DIQ= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= -github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= -github.com/tidwall/gjson v1.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU= -github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc= -github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= -github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM= -github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= -github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= -github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= -github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= -github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= -github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index c778396d1..a5b1753ee 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -37,10 +37,6 @@ import ( "sync" "time" - "github.com/tidwall/gjson" - "github.com/tidwall/sjson" - - "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riversqlite/internal/dbsqlc" "github.com/riverqueue/river/rivershared/sqlctemplate" @@ -598,7 +594,7 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo return &riverdriver.JobInsertFastResult{ Job: job, - UniqueSkippedAsDuplicate: gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyUniqueNonce).Str != uniqueNonce, + UniqueSkippedAsDuplicate: riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce), }, nil }) } @@ -1186,6 +1182,10 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan return dbsqlc.New().NotificationInsertMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, notifications) } +func (e *Executor) Ping(ctx context.Context) error { + return e.Exec(ctx, "SELECT 1") +} + func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { return nil, riverdriver.ErrNotImplemented } @@ -1516,7 +1516,7 @@ func sqliteJobInsertFastManyJobsParam(jobs []*riverdriver.JobInsertFastParams, u metadata := sliceutil.FirstNonEmpty(job.Metadata, []byte("{}")) if uniqueNonce != "" { var err error - metadata, err = sjson.SetBytes(metadata, rivercommon.MetadataKeyUniqueNonce, uniqueNonce) + metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce) if err != nil { return nil, err } diff --git a/riverdriver/unique_insert.go b/riverdriver/unique_insert.go new file mode 100644 index 000000000..191e26da0 --- /dev/null +++ b/riverdriver/unique_insert.go @@ -0,0 +1,112 @@ +package riverdriver + +import ( + "encoding/json" + "fmt" + "strings" +) + +// UniqueInsertMetadataKey is a reserved job metadata key used to detect unique +// insert conflicts on databases that don't expose PostgreSQL system columns. +const UniqueInsertMetadataKey = "river:unique_nonce" + +// UniqueInsertMode is a database-specific strategy for detecting whether a +// unique insert returned a newly inserted job or an existing one. +type UniqueInsertMode uint32 + +const ( + // UniqueInsertModeUnknown indicates that a database's mode hasn't been + // detected yet. + UniqueInsertModeUnknown UniqueInsertMode = iota + + // UniqueInsertModeMetadataNonce detects conflicts by putting a nonce in the + // metadata of the proposed job and checking whether the returned job + // contains it. + UniqueInsertModeMetadataNonce + + // UniqueInsertModeReturningOld uses PostgreSQL 18's OLD row support in + // RETURNING. + UniqueInsertModeReturningOld + + // UniqueInsertModeXmax uses PostgreSQL's xmax system column. + UniqueInsertModeXmax +) + +// SQL returns the SQL expression for the mode. UniqueInsertModeMetadataNonce +// always returns false because duplicate detection is performed in Go instead. +func (m UniqueInsertMode) SQL() string { + switch m { + case UniqueInsertModeMetadataNonce: + return "false" + + case UniqueInsertModeReturningOld: + return "(OLD.id IS NOT NULL)" + + case UniqueInsertModeXmax: + return "(xmax != 0)" + + case UniqueInsertModeUnknown: + panic("unique insert mode has not been detected") + + default: + panic(fmt.Sprintf("invalid unique insert mode: %d", m)) + } +} + +// UniqueInsertMetadataIsDuplicate returns whether metadata lacks the nonce +// from a proposed insert, indicating that an existing row was returned +// instead. +func UniqueInsertMetadataIsDuplicate(metadata []byte, nonce string) bool { + var metadataMap map[string]json.RawMessage + if err := json.Unmarshal(metadata, &metadataMap); err != nil { + return true + } + + var metadataNonce string + if err := json.Unmarshal(metadataMap[UniqueInsertMetadataKey], &metadataNonce); err != nil { + return true + } + return metadataNonce != nonce +} + +// UniqueInsertMetadataWithNonce returns metadata with nonce set under +// UniqueInsertMetadataKey. +func UniqueInsertMetadataWithNonce(metadata []byte, nonce string) ([]byte, error) { + if len(metadata) == 0 { + metadata = []byte("{}") + } + + var metadataMap map[string]json.RawMessage + if err := json.Unmarshal(metadata, &metadataMap); err != nil { + return nil, fmt.Errorf("error unmarshaling job metadata: %w", err) + } + if metadataMap == nil { + metadataMap = make(map[string]json.RawMessage) + } + + nonceJSON, err := json.Marshal(nonce) + if err != nil { + return nil, fmt.Errorf("error marshaling unique insert nonce: %w", err) + } + metadataMap[UniqueInsertMetadataKey] = nonceJSON + + metadata, err = json.Marshal(metadataMap) + if err != nil { + return nil, fmt.Errorf("error marshaling job metadata: %w", err) + } + return metadata, nil +} + +// UniqueInsertModeFromProductAndVersion returns the unique insert mode +// appropriate for a database product and its PostgreSQL-compatible server +// version number. +func UniqueInsertModeFromProductAndVersion(product string, version int32) UniqueInsertMode { + productLower := strings.ToLower(product) + if strings.Contains(productLower, "-yb") || strings.Contains(productLower, "yugabyte") { + return UniqueInsertModeMetadataNonce + } + if version >= 180_000 { + return UniqueInsertModeReturningOld + } + return UniqueInsertModeXmax +} diff --git a/riverdriver/unique_insert_test.go b/riverdriver/unique_insert_test.go new file mode 100644 index 000000000..3c78de985 --- /dev/null +++ b/riverdriver/unique_insert_test.go @@ -0,0 +1,126 @@ +package riverdriver + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestUniqueInsertMetadataIsDuplicate(t *testing.T) { + t.Parallel() + + t.Run("DifferentNonce", func(t *testing.T) { + t.Parallel() + + require.True(t, UniqueInsertMetadataIsDuplicate([]byte(`{"river:unique_nonce":"old"}`), "new")) + }) + + t.Run("InvalidMetadata", func(t *testing.T) { + t.Parallel() + + require.True(t, UniqueInsertMetadataIsDuplicate([]byte(`{`), "nonce")) + }) + + t.Run("MatchingNonce", func(t *testing.T) { + t.Parallel() + + require.False(t, UniqueInsertMetadataIsDuplicate([]byte(`{"river:unique_nonce":"nonce"}`), "nonce")) + }) + + t.Run("MissingNonce", func(t *testing.T) { + t.Parallel() + + require.True(t, UniqueInsertMetadataIsDuplicate([]byte(`{"existing":123}`), "nonce")) + }) +} + +func TestUniqueInsertMetadataWithNonce(t *testing.T) { + t.Parallel() + + t.Run("EmptyMetadata", func(t *testing.T) { + t.Parallel() + + metadata, err := UniqueInsertMetadataWithNonce(nil, "nonce") + require.NoError(t, err) + require.JSONEq(t, `{"river:unique_nonce":"nonce"}`, string(metadata)) + }) + + t.Run("ExistingMetadata", func(t *testing.T) { + t.Parallel() + + metadata, err := UniqueInsertMetadataWithNonce([]byte(`{"existing":123}`), "nonce") + require.NoError(t, err) + require.JSONEq(t, `{"existing":123,"river:unique_nonce":"nonce"}`, string(metadata)) + }) + + t.Run("ExistingNonce", func(t *testing.T) { + t.Parallel() + + metadata, err := UniqueInsertMetadataWithNonce([]byte(`{"river:unique_nonce":"old"}`), "new") + require.NoError(t, err) + require.JSONEq(t, `{"river:unique_nonce":"new"}`, string(metadata)) + }) + + t.Run("InvalidMetadata", func(t *testing.T) { + t.Parallel() + + _, err := UniqueInsertMetadataWithNonce([]byte(`{`), "nonce") + require.ErrorContains(t, err, "error unmarshaling job metadata") + }) +} + +func TestUniqueInsertModeFromProductAndVersion(t *testing.T) { + t.Parallel() + + t.Run("PostgreSQL17", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeXmax, UniqueInsertModeFromProductAndVersion("PostgreSQL 17.5", 170_005)) + }) + + t.Run("PostgreSQL18", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeReturningOld, UniqueInsertModeFromProductAndVersion("PostgreSQL 18.0", 180_000)) + }) + + t.Run("YugabyteByName", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeMetadataNonce, UniqueInsertModeFromProductAndVersion("YugabyteDB", 180_000)) + }) + + t.Run("YugabytePostgreSQLVersion", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeMetadataNonce, UniqueInsertModeFromProductAndVersion("PostgreSQL 15.2-YB-2.25.1.0-b0", 150_002)) + }) +} + +func TestUniqueInsertModeSQL(t *testing.T) { + t.Parallel() + + t.Run("MetadataNonce", func(t *testing.T) { + t.Parallel() + + require.Equal(t, "false", UniqueInsertModeMetadataNonce.SQL()) + }) + + t.Run("ReturningOld", func(t *testing.T) { + t.Parallel() + + require.Equal(t, "(OLD.id IS NOT NULL)", UniqueInsertModeReturningOld.SQL()) + }) + + t.Run("Unknown", func(t *testing.T) { + t.Parallel() + + require.PanicsWithValue(t, "unique insert mode has not been detected", func() { UniqueInsertModeUnknown.SQL() }) + }) + + t.Run("Xmax", func(t *testing.T) { + t.Parallel() + + require.Equal(t, "(xmax != 0)", UniqueInsertModeXmax.SQL()) + }) +}