perf(msgqueue): optimistic insert + heal-on-23503 — drop the per-publish parent-row lock

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,22 @@ 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 only takes a KEY SHARE lock on the parent MessageQueue
+	// row (via the FK check), so concurrent publishers to the same queue do not
+	// serialize on it. 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, and the producer's existence cache can
+	// be stale). Recreate it 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 +120,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 {

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))
+}