Forward-port 0.6.16 to main, step 2: batch inserts, new lanes and review fixes

.github/workflows/ci.yml

@@ -259,7 +259,7 @@ jobs:
     # doctests run as `workspace`. Shard membership lives in
     # scripts/ci-test-shard.sh; the lint job runs its `check` mode so renames
     # can't silently drop coverage.
-    timeout-minutes: 20
+    timeout-minutes: 30
     strategy:
       fail-fast: false
       matrix:
@@ -477,7 +477,7 @@ jobs:
       - name: Create venv and install deps
         run: |
           uv venv .venv --python "${{ matrix.python }}"
-          uv pip install maturin pytest pytest-asyncio asyncpg django greenlet psycopg psycopg-binary sqlalchemy
+          uv pip install maturin pytest pytest-asyncio asyncpg django greenlet "psycopg[binary]" sqlalchemy
       - name: Build extension (maturin develop)
         run: .venv/bin/maturin develop
       - name: Wait for Postgres

CHANGELOG.md

@@ -15,6 +15,9 @@ Notable changes between releases. Detailed migration notes for storage transitio
 ### Forward-ported from 0.6.16
 
 - Queue-storage rescue drains a large backlog of abandoned claims in bounded time. A killed runtime leaves every attempt it claimed open, and until they close, queue prune skips the queue slots they reference and claim prune skips their claim slots, so neither ring can wrap. A rescue pass takes 500-attempt batches, each in its own transaction, until a batch comes back short, `rescue_pass_budget` (default 1 s, at most half the rescue interval) is spent, or it has rescued `rescue_pass_max_rows` (default 5,000), and while a backlog remains the next pass runs a second later. `ClientBuilder::rescue_pass_budget` and `rescue_pass_max_rows` (Python `rescue_pass_budget_ms` and `rescue_pass_max_rows`) set both; zero for either keeps one batch per interval. Deadline and callback rescue act on their expired deadlines; stale-heartbeat rescue takes more than one batch per pass only once runtime snapshots show the claims' owners gone. Rescue's cursor scans seek to their cursor through the index, including the compact deadline-batch scan that 0.6 does not have, and rescue transactions run without JIT. `awa.maintenance.rescue.pass_size` (by `awa.rescue.kind` and `awa.rescue.pass_end`) and `awa.maintenance.rescue.oldest_attempt_age` show a drain in progress. Heartbeats no longer queue behind a completion: an attempt whose closer holds its lock is skipped, lease and progress writes skip rows a finalizer has locked and retry them on the next pass, and a finished attempt stops flushing progress once its handler returns. A point read of a failed job that starts while queue prune moves its row out of the slot being pruned is covered by a test; on this branch the read already locks the ring before it takes its snapshot. See [draining a backlog of abandoned claims](docs/configuration.md#rescue-backlog).
+- Claim prune no longer treats a partial GIN index on `lease_claim_batches.job_ids` as the claim batch index (migration v056). It recognises only a valid, non-partial GIN index on `job_ids` alone. When another index or relation already has the index's name, `ensure_claim_batch_job_index` leaves it in place, returns false and raises a warning naming it, instead of reporting the slot indexed; the slot keeps scanning until an operator drops or renames it. Refreshing the function as a migration role that owns it but may only `SET ROLE` to the table owner (PostgreSQL 16+ membership `WITH INHERIT FALSE`) no longer leaves `EXECUTE` with `PUBLIC`: the refresh resets `EXECUTE` after it hands the function to the table owner, acting as the new owner when it cannot as itself. Migrations run with the owner as the effective role, as documented, were not affected. v056 replaces functions only and takes no table locks. See [v056](docs/upgrade-0.6-to-0.7.md#v056-claim-batch-index-recognition).
+- Creating a new queue lane through `QueueStorage` (its batch enqueues, and workers promoting scheduled or retried jobs) no longer deadlocks with a concurrent `insert`, `insert_many` or `insert_many_copy` that is also the first to use the same `(queue, priority)` (#692). `QueueStorage` took the lane's provisioning lock before writing its `queue_lanes` row, while SQL inserts write the row first and take the lock in the head-table trigger that creates the lane's sequences, so each transaction waited on the other and one failed with `40P01`. `QueueStorage` now writes the row first too, and so does the `provision_lane` function it calls (migration v058, [v058](docs/upgrade-0.6-to-0.7.md#v058-lane-row-before-lane-lock)). Lanes that already existed were never affected. `QueueStorage` processes on 0.6.15 or earlier, or on an earlier 0.7 build, keep the old order until they are upgraded.
+- `insert_many_copy` and `insert_many_skip_unique_conflicts` no longer deadlock with per-row producers (generic `insert_many`, single-job inserts, binaries before v055) that share their unique keys or lanes (migration v057). With the v055 function, when a per-row transaction held a lane and then inserted a key a waiting batch had claimed, PostgreSQL aborted one of the two with `deadlock detected` (`40P01`). The batch function now waits at most 500 ms (or half of its session's `deadlock_timeout`, if lower) while it holds claims or lanes: on a longer wait it releases them and tries again, so the per-row transaction proceeds and the batch then treats the keys it committed as already claimed. The same applies to a transaction that enqueues lanes in separate statements in the opposite order to a batch. This covers per-row sessions whose `deadlock_timeout` is above 500 ms, as PostgreSQL's default of 1 s is, and batches whose unique claims insert within that budget (10,000 to 20,000 keys on a fast server, fewer on a slow or busy one); larger batches behave as before. The protection lasts for 2 s of retries; then the batch makes one last attempt with ordinary lock waits. A per-row transaction that holds a lane for longer than that before it inserts one of the batch's keys can therefore still deadlock with it. A cycle through a lock the batch's own transaction took in an earlier statement also ends in `deadlock detected`, after about 3 to 4 s instead of 1 to 2 s, and usually aborts the batch's transaction. The caller's `lock_timeout` applies to each lock wait, as before. Results and the lane lock window are unchanged. v057 is one function replacement with no table locks, and producers that already call the function get the fix as soon as it commits. See [v057](docs/upgrade-0.6-to-0.7.md#v057-batch-inserts-give-way-to-per-row-producers).
 
 ### Added
 

awa-model/migrations/v023_install_queue_storage_substrate.sql

@@ -325,13 +325,22 @@ BEGIN
         SET search_path = pg_catalog, pg_temp
         AS $func$
         BEGIN
-            -- Validates the lane and takes its lock before any row is written.
-            PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard);
+            IF p_queue IS NULL
+               OR p_priority IS NULL
+               OR p_enqueue_shard IS NULL OR p_enqueue_shard NOT BETWEEN 0 AND 63 THEN
+                RAISE EXCEPTION 'invalid queue lane: queue %%, priority %%, enqueue shard %%',
+                    quote_nullable(p_queue), p_priority, p_enqueue_shard
+                    USING ERRCODE = '22023';
+            END IF;
 
+            -- The lane row comes before the lane lock, the order every
+            -- producer takes them in.
             INSERT INTO %1$I.queue_lanes (queue, priority)
             VALUES (p_queue, p_priority)
             ON CONFLICT (queue, priority) DO NOTHING;
 
+            PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard);
+
             -- The head triggers fire before ON CONFLICT and set the lane's
             -- sequence, so rows that already exist are skipped outright.
             INSERT INTO %1$I.queue_enqueue_heads (queue, priority, enqueue_shard)
@@ -798,7 +807,9 @@ REVOKE EXECUTE ON FUNCTION awa.install_queue_storage_job_output(TEXT) FROM PUBLI
 -- index on a populated partition holds SHARE, which blocks claim inserts, for
 -- the whole build (minutes per million batches), so the index is only ever
 -- created on an empty partition: by the installer, or by claim prune right
--- after it truncates a slot.
+-- after it truncates a slot. A different relation that already has the
+-- index's name is never dropped: the function warns, naming it, and the slot
+-- keeps scanning until an operator removes or renames it.
 --
 -- fastupdate is off: job ids arrive in ascending order, so direct inserts
 -- append to the right edge of the entry tree and cost no more than the
@@ -822,6 +833,7 @@ DECLARE
     v_function REGPROCEDURE;
     v_function_owner OID;
     v_grantee TEXT;
+    v_previous_role TEXT;
     v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000
                             THEN 'SET' ELSE 'MEMBER' END;
 BEGIN
@@ -873,19 +885,44 @@ BEGIN
                     USING ERRCODE = '22023';
             END IF;
 
+            -- Only a valid, non-partial GIN index on job_ids alone, with the
+            -- default array_ops opclass, serves the by-job lookups.
             IF EXISTS (
                 SELECT 1
                 FROM pg_index AS idx
                 JOIN pg_class AS index_class ON index_class.oid = idx.indexrelid
                 JOIN pg_am AS method ON method.oid = index_class.relam
+                JOIN pg_attribute AS job_ids
+                  ON job_ids.attrelid = idx.indrelid
+                 AND job_ids.attname = 'job_ids'
+                 AND NOT job_ids.attisdropped
+                JOIN pg_opclass AS opclass ON opclass.oid = idx.indclass[0]
                 WHERE idx.indrelid = to_regclass(format('%%I.%%I', %1$L, v_child))
                   AND method.amname = 'gin'
                   AND idx.indisvalid
-                  AND pg_get_indexdef(idx.indexrelid) LIKE '%%USING gin (job_ids)%%'
+                  AND idx.indpred IS NULL
+                  AND idx.indexprs IS NULL
+                  AND idx.indnatts = 1
+                  AND idx.indkey[0] = job_ids.attnum
+                  AND opclass.opcname = 'array_ops'
+                  AND opclass.opcnamespace = 'pg_catalog'::regnamespace
             ) THEN
                 RETURN TRUE;
             END IF;
 
+            -- Something else already has the index's name: a partial or
+            -- otherwise different index, or another relation. Leave it alone
+            -- and keep scanning; the operator decides what to do with it.
+            IF to_regclass(format('%%I.%%I', %1$L, v_index)) IS NOT NULL THEN
+                RAISE WARNING '%%.%% is not a GIN index on job_ids, so claim slot %% of schema %% keeps scanning for job ids',
+                    quote_ident(%1$L), quote_ident(v_index), p_slot, %1$L
+                    USING HINT = format(
+                        'Drop or rename %%s.%%s; claim prune builds the job_ids index the next time it empties claim slot %%s.',
+                        quote_ident(%1$L), quote_ident(v_index), p_slot
+                    );
+                RETURN FALSE;
+            END IF;
+
             -- NOWAIT: a partition someone is writing to is not empty for long
             -- enough to index, so give up rather than queue behind claims.
             -- Leaving the block through the exception rolls back its
@@ -899,11 +936,11 @@ BEGIN
                     RAISE EXCEPTION USING ERRCODE = 'object_not_in_prerequisite_state';
                 END IF;
                 EXECUTE format(
-                    'CREATE INDEX IF NOT EXISTS %%I ON %%I.%%I USING GIN (job_ids) WITH (fastupdate = off)',
+                    'CREATE INDEX %%I ON %%I.%%I USING GIN (job_ids) WITH (fastupdate = off)',
                     v_index, %1$L, v_child
                 );
             EXCEPTION
-                WHEN lock_not_available OR object_not_in_prerequisite_state THEN
+                WHEN lock_not_available OR object_not_in_prerequisite_state OR duplicate_table THEN
                     RETURN FALSE;
             END;
             RETURN TRUE;
@@ -928,6 +965,7 @@ BEGIN
                 v_function,
                 pg_get_userbyid(v_table_owner)
             );
+            v_function_owner := v_table_owner;
         ELSE
             RAISE WARNING '% is owned by %, not the owner % of %.lease_claim_batches, so it cannot create the job_ids index',
                 v_function, pg_get_userbyid(v_function_owner), pg_get_userbyid(v_table_owner), p_schema
@@ -939,6 +977,14 @@ BEGIN
         END IF;
     END IF;
 
+    -- Reset EXECUTE as the function's owner, after any transfer, so a former
+    -- owner that holds TRUNCATE keeps it. A role that may only SET ROLE to the
+    -- new owner cannot grant or revoke as itself, so it acts as the owner for
+    -- the reset and then restores its previous role.
+    IF NOT pg_has_role(current_user, v_function_owner, 'USAGE') THEN
+        v_previous_role := current_setting('role');
+        PERFORM set_config('role', pg_get_userbyid(v_function_owner), true);
+    END IF;
     FOR v_grantee IN
         SELECT DISTINCT CASE WHEN acl.grantee = 0 THEN 'PUBLIC'
                              ELSE quote_ident(pg_get_userbyid(acl.grantee)) END
@@ -958,6 +1004,9 @@ BEGIN
     LOOP
         EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO %s', v_function, v_grantee);
     END LOOP;
+    IF v_previous_role IS NOT NULL THEN
+        PERFORM set_config('role', v_previous_role, true);
+    END IF;
 END;
 $refresh$;
 

awa-model/migrations/v056_claim_batch_job_index_shape.sql

@@ -0,0 +1,57 @@
+-- v056: recognise only a full job_ids GIN index as the claim batch index.
+-- Tracking: https://github.com/hardbyte/awa/pull/701
+-- <schema>.ensure_claim_batch_job_index(slot) from v054 reported a claim
+-- partition indexed when it had any valid GIN index whose definition named
+-- job_ids, a partial one included, and its CREATE INDEX IF NOT EXISTS took a
+-- different relation that already had the index's name for the index. The
+-- function now accepts only a valid, non-partial GIN index on job_ids alone
+-- with the default opclass. When something else has the index's name it
+-- raises a warning naming it, returns false and leaves it in place.
+--
+-- awa.refresh_claim_batch_job_index_helper now resets EXECUTE on the function
+-- after it hands the function to the owner of lease_claim_batches, acting as
+-- that owner when the caller may only SET ROLE to it, so EXECUTE ends with
+-- exactly the roles that hold TRUNCATE on the table.
+--
+-- This file only reinstalls the function in each queue-storage schema; it
+-- builds and drops no index. `awa migrate` and
+-- `awa migrate --sql --from 55 --to 56` run V23_UP first, which defines the
+-- helper.
+--
+-- N-1 binaries call the function with the same argument and result. An index
+-- that v054's function built is a full job_ids index and is still recognised.
+--
+-- Live load: function replacement only. It takes no lock on any awa or
+-- queue-storage table. Run it as the owner of the queue-storage tables, or as
+-- a migrator that can SET ROLE to it, so the function is owned by the table
+-- owner and can create the index.
+
+DO $$
+DECLARE
+    v_schema TEXT;
+BEGIN
+    FOR v_schema IN
+        SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered
+    LOOP
+        IF NOT EXISTS (
+            SELECT 1 FROM pg_class AS batches
+            WHERE batches.oid = to_regclass(format('%I.lease_claim_batches', v_schema))
+              AND batches.relkind = 'p'
+        ) THEN
+            CONTINUE;
+        END IF;
+
+        -- Serialize with a concurrent prepare_schema() of the same schema. This
+        -- is an advisory lock, not a table lock: claims and completions do not
+        -- take it.
+        PERFORM pg_advisory_xact_lock(
+            hashtextextended(format('awa.queue_storage.install:%I', v_schema), 0)
+        );
+        PERFORM awa.refresh_claim_batch_job_index_helper(v_schema);
+    END LOOP;
+END;
+$$;
+
+INSERT INTO awa.schema_version (version, description)
+VALUES (56, 'Recognise only a full job_ids index as the claim batch index')
+ON CONFLICT (version) DO NOTHING;