diff --git a/pkg/repository/mq.go b/pkg/repository/mq.go index aab0233304..a5cd5b2445 100644 --- a/pkg/repository/mq.go +++ b/pkg/repository/mq.go @@ -7,7 +7,9 @@ import ( "fmt" "github.com/google/uuid" + "github.com/jackc/pgerrcode" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgtype" "github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1" @@ -70,12 +72,11 @@ func (m *messageQueueRepository) Notify(ctx context.Context, name string, payloa // If the wrapped message exceeds this, fall back to database storage if len(wrappedPayload) > 8000 { // An auto-deleted queue can be reaped by CleanupMessageQueue in the window - // between the producer's 15s existence-cache hit and this insert, so ensure - // the parent row exists in the same statement (FK can't fail). This covers - // EXCLUSIVE auto-deleted queues too — the dispatcher queue (expirable ⇒ - // autoDeleted, exclusive) and controller consumer queues — which the reaper - // deletes regardless of exclusivity; gating on `!exclusive` left them - // exposed to the exact 23503 this fix exists to kill. + // between the producer's 15s existence-cache hit and this insert, so route + // through the self-healing path. This covers EXCLUSIVE auto-deleted queues + // too — the dispatcher queue (expirable ⇒ autoDeleted, exclusive) and + // controller consumer queues — which the reaper deletes regardless of + // exclusivity. if autoDeleted { return m.AddMessageEnsuringQueue(ctx, name, []byte(payload), durable, autoDeleted, exclusive) } @@ -94,10 +95,31 @@ func (m *messageQueueRepository) AddMessage(ctx context.Context, queue string, p } func (m *messageQueueRepository) AddMessageEnsuringQueue(ctx context.Context, queue string, payload []byte, durable, autoDeleted, exclusive bool) error { - if !autoDeleted { - return m.AddMessage(ctx, queue, payload) + // The plain insert's FK check takes only a KEY SHARE lock on the parent + // MessageQueue row, so concurrent publishers to the same queue do not + // serialize on it (an unconditional ON CONFLICT DO UPDATE here takes a + // row-exclusive tuple lock per publish, convoying every publisher on one + // row). It fails with a foreign-key violation exactly when the parent has + // been reaped, which is the only case that needs the self-heal. + err := m.AddMessage(ctx, queue, payload) + + if err == nil || !autoDeleted || !isForeignKeyViolation(err) { + return err } + // The parent queue was reaped by CleanupMessageQueue (auto-deleted queues + // are reap-eligible after 1h idle). Healing is expected at most once per + // queue per reap: the producer's queue-bind cache expires every 15s and + // each re-bind refreshes lastActive, so an actively-published queue can + // never accrue the 1h of staleness the reaper requires (15s ≪ 1h). A + // sustained heal rate therefore means something else is deleting + // MessageQueue rows out-of-band. + m.l.Warn().Ctx(ctx).Str("queue", queue).Msg("parent queue row missing on publish; self-healing by recreating it") + + // Recreate the queue and insert the item in a single atomic statement. + // The upsert sets lastActive = NOW(), so the healed queue is not + // reap-eligible for another hour and subsequent publishes take the plain + // insert above. return m.queries.AddMessageEnsuringQueue(ctx, m.pool, sqlcv1.AddMessageEnsuringQueueParams{ Queueid: queue, Payload: payload, @@ -107,6 +129,11 @@ func (m *messageQueueRepository) AddMessageEnsuringQueue(ctx context.Context, qu }) } +func isForeignKeyViolation(err error) bool { + var pgErr *pgconn.PgError + return errors.As(err, &pgErr) && pgErr.Code == pgerrcode.ForeignKeyViolation +} + func (m *messageQueueRepository) BindQueue(ctx context.Context, queue string, durable, autoDeleted, exclusive bool, exclusiveConsumer *string) error { // if exclusive, but no consumer, return error if exclusive && exclusiveConsumer == nil { diff --git a/pkg/repository/mq_selfheal_test.go b/pkg/repository/mq_selfheal_test.go index 8528a7f24b..aeaff76bdc 100644 --- a/pkg/repository/mq_selfheal_test.go +++ b/pkg/repository/mq_selfheal_test.go @@ -7,6 +7,8 @@ import ( "testing" "time" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/rs/zerolog" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -184,3 +186,119 @@ func TestBindRefreshesLastActive(t *testing.T) { require.NoError(t, err) assert.Equal(t, 1, exists, "a recently rebound queue must survive cleanup") } + +func mqRepoForTest(pool *pgxpool.Pool) *messageQueueRepository { + l := zerolog.Nop() + + return &messageQueueRepository{ + sharedRepository: &sharedRepository{ + pool: pool, + queries: sqlcv1.New(), + l: &l, + }, + } +} + +func TestAddMessageEnsuringQueueFastPathSkipsParentUpsert(t *testing.T) { + pool, cleanup := setupPostgresWithMigration(t) + defer cleanup() + + ctx := context.Background() + repo := mqRepoForTest(pool) + q := sqlcv1.New() + const name = "mq-optimistic-fast-path-queue" + + _, err := q.UpsertMessageQueue(ctx, pool, sqlcv1.UpsertMessageQueueParams{ + Name: name, + Durable: false, + Autodeleted: true, + Exclusive: false, + }) + require.NoError(t, err) + + // Backdate lastActive so any parent-row write by the publish is detectable. + _, err = pool.Exec(ctx, + `UPDATE "MessageQueue" SET "lastActive" = NOW() - INTERVAL '30 minutes' WHERE "name" = $1`, + name, + ) + require.NoError(t, err) + + err = repo.AddMessageEnsuringQueue(ctx, name, []byte(`{"hello":"world"}`), false, true, false) + require.NoError(t, err) + + var lastActive time.Time + err = pool.QueryRow(ctx, `SELECT "lastActive" FROM "MessageQueue" WHERE "name" = $1`, name).Scan(&lastActive) + require.NoError(t, err) + assert.Less(t, time.Since(lastActive), 45*time.Minute) + assert.Greater(t, time.Since(lastActive), 15*time.Minute, + "a publish to an existing queue must not rewrite the parent row (no ON CONFLICT DO UPDATE on the hot path)") + + var items int + err = pool.QueryRow(ctx, `SELECT count(*) FROM "MessageQueueItem" WHERE "queueId" = $1`, name).Scan(&items) + require.NoError(t, err) + assert.Equal(t, 1, items) +} + +func TestAddMessageEnsuringQueueHealsReapedParentOnFKViolation(t *testing.T) { + pool, cleanup := setupPostgresWithMigration(t) + defer cleanup() + + ctx := context.Background() + repo := mqRepoForTest(pool) + q := sqlcv1.New() + const name = "mq-optimistic-heal-queue" + + _, err := q.UpsertMessageQueue(ctx, pool, sqlcv1.UpsertMessageQueueParams{ + Name: name, + Durable: true, + Autodeleted: true, + Exclusive: true, + }) + require.NoError(t, err) + + err = repo.AddMessageEnsuringQueue(ctx, name, []byte(`{"seq":1}`), true, true, true) + require.NoError(t, err) + + // Reap the queue mid-stream, exactly as CleanupMessageQueue does in the race. + _, err = pool.Exec(ctx, + `UPDATE "MessageQueue" SET "lastActive" = NOW() - INTERVAL '2 hours' WHERE "name" = $1`, + name, + ) + require.NoError(t, err) + require.NoError(t, q.CleanupMessageQueue(ctx, pool)) + + var reaped int + err = pool.QueryRow(ctx, `SELECT count(*) FROM "MessageQueue" WHERE "name" = $1`, name).Scan(&reaped) + require.NoError(t, err) + require.Equal(t, 0, reaped, "queue should have been reaped, setting up the race") + + err = repo.AddMessageEnsuringQueue(ctx, name, []byte(`{"seq":2}`), true, true, true) + require.NoError(t, err, "a publish against a reaped queue must self-heal, not surface 23503") + + var exclusive bool + var lastActive time.Time + err = pool.QueryRow(ctx, + `SELECT "exclusive", "lastActive" FROM "MessageQueue" WHERE "name" = $1`, name, + ).Scan(&exclusive, &lastActive) + require.NoError(t, err) + assert.True(t, exclusive, "the healed queue must preserve the supplied bind attributes") + assert.WithinDuration(t, time.Now(), lastActive, time.Minute, + "the heal must refresh lastActive so the queue is not immediately reap-eligible again") + + var items int + err = pool.QueryRow(ctx, `SELECT count(*) FROM "MessageQueueItem" WHERE "queueId" = $1`, name).Scan(&items) + require.NoError(t, err) + assert.Equal(t, 1, items, "the post-reap message must be stored against the recreated queue") +} + +func TestAddMessageEnsuringQueueNonAutoDeletedPropagatesFKViolation(t *testing.T) { + pool, cleanup := setupPostgresWithMigration(t) + defer cleanup() + + ctx := context.Background() + repo := mqRepoForTest(pool) + + err := repo.AddMessageEnsuringQueue(ctx, "mq-optimistic-missing-static-queue", []byte(`{}`), true, false, false) + require.Error(t, err, "non-auto-deleted queues are never reaped, so a missing parent is a real bug and must not be masked by a self-heal") + assert.True(t, isForeignKeyViolation(err)) +}