fix: Use a good index for dequeue

[tomasfarias]

Our current index is bad.

I created a database and table with 100 million jobs and ran the dequeue tx query as a prepared statement:

test_database_1# EXPLAIN ANALYZE EXECUTE dequeue_tx('default', 'test');
                                                                            QUERY PLAN                                        
                                    
══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════
════════════════════════════════════
 Update on job_queue  (cost=4999952.58..4999960.63 rows=1 width=86) (actual time=70519.543..70519.546 rows=1 loops=1)
   CTE available_in_queue
     ->  Limit  (cost=4999952.00..4999952.01 rows=1 width=26) (actual time=70497.417..70497.418 rows=1 loops=1)
           ->  LockRows  (cost=4999952.00..6249827.50 rows=99990040 width=26) (actual time=70393.454..70393.455 rows=1 loops=1
)
                 ->  Sort  (cost=4999952.00..5249927.10 rows=99990040 width=26) (actual time=70393.328..70393.328 rows=1 loops
=1)
                       Sort Key: job_queue_1.attempt, job_queue_1.scheduled_at
                       Sort Method: external merge  Disk: 3913976kB
                       ->  Seq Scan on job_queue job_queue_1  (cost=0.00..4500001.80 rows=99990040 width=26) (actual time=0.01
5..10763.967 rows=99999999 loops=1)
                             Filter: ((status = 'available'::job_status) AND (queue = 'default'::text) AND (scheduled_at <= no
w()))
                             Rows Removed by Filter: 1
   ->  Nested Loop  (cost=0.57..8.61 rows=1 width=86) (actual time=70506.223..70506.225 rows=1 loops=1)
         ->  CTE Scan on available_in_queue  (cost=0.00..0.02 rows=1 width=40) (actual time=70497.438..70497.439 rows=1 loops=
1)
         ->  Index Scan using job_queue_pkey on job_queue  (cost=0.57..8.59 rows=1 width=31) (actual time=8.773..8.773 rows=1 
loops=1)
               Index Cond: (id = available_in_queue.id)
 Planning Time: 0.522 ms
 JIT:
   Functions: 15
   Options: Inlining true, Optimization true, Expressions true, Deforming true
   Timing: Generation 0.954 ms, Inlining 32.689 ms, Optimization 49.676 ms, Emission 34.811 ms, Total 118.131 ms
 Execution Time: 71108.776 ms

Notice the following:

• Execution time is through the roof!
• The main performance hits are on the CTE: notice the UPDATE query uses the pkey index and selects only one row (which should be fast according to the query plan cost).
• In the CTE:
  • Postgres is sorting on attempt, scheduled_at (due to ORDER BY clause) which our current index cannot help with as it's defined on scheduled_at, attempt (notice the order is reversed).
  • After sorting, pg scans the entire thing looking for available jobs that match the queue and scheduled_at, just to find a single row!

Solution: A partial index (after all, we only dequeue available jobs, no need to index other values) with the correct column order:

test_database_1# CREATE INDEX idx_queue_pg_is_quite_fast_if_you_help_it ON job_queue(queue, attempt, scheduled_at) where status = 'available' :: job_status;
CREATE INDEX
test_database_1# EXPLAIN ANALYZE EXECUTE dequeue_tx('default', 'test');
                                                                                             QUERY PLAN                       
                                                                       
══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════
═══════════════════════════════════════════════════════════════════════
 Update on job_queue  (cost=2.64..10.69 rows=1 width=86) (actual time=8.824..8.826 rows=1 loops=1)
   CTE available_in_queue
     ->  Limit  (cost=0.57..2.08 rows=1 width=26) (actual time=0.056..0.057 rows=1 loops=1)
           ->  LockRows  (cost=0.57..150551172.94 rows=99990000 width=26) (actual time=0.055..0.056 rows=1 loops=1)
                 ->  Index Scan using idx_queue_pg_is_quite_fast_if_you_help_it on job_queue job_queue_1  (cost=0.57..14955127
2.94 rows=99990000 width=26) (actual time=0.040..0.040 rows=1 loops=1)
                       Index Cond: ((queue = 'default'::text) AND (scheduled_at <= now()))
                       Filter: (status = 'available'::job_status)
   ->  Nested Loop  (cost=0.57..8.61 rows=1 width=86) (actual time=8.771..8.772 rows=1 loops=1)
         ->  CTE Scan on available_in_queue  (cost=0.00..0.02 rows=1 width=40) (actual time=0.062..0.063 rows=1 loops=1)
         ->  Index Scan using job_queue_pkey on job_queue  (cost=0.57..8.59 rows=1 width=31) (actual time=8.705..8.705 rows=1 
loops=1)
               Index Cond: (id = available_in_queue.id)
 Planning Time: 0.746 ms
 Execution Time: 8.885 ms
(13 rows)

Down to 8ms in a 100 million rows:

test_database_1# select count(*) from job_queue;
   count   
═══════════
 100000000
(1 row)

Notice the following:

• We are now scanning the index.
• UPDATE still uses pkey index.

TODO for later: We should prepare the queries as statements to save parsing time. AFAIK sqlx does not do this for us (it has a prepare method but I think it's only used to check types, not the actual PREPARE command in Postgres)

[tomasfarias]

Probably the 3 lines with most perf impact I've written in a while.

[tomasfarias]

For the sake of completeness, here's the script I used to generate the test data (warning: it took 10 minutes on my machine).

CREATE DATABASE test_database_1;
\c test_database_1;

CREATE TYPE job_status AS ENUM(
    'available',
    'completed',
    'failed',
    'running'
);

CREATE TABLE job_queue(
    id BIGSERIAL PRIMARY KEY,
    attempt INT NOT NULL DEFAULT 0,
    attempted_at TIMESTAMPTZ DEFAULT NULL,
    attempted_by TEXT [] DEFAULT ARRAY [] :: TEXT [],
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    errors JSONB [],
    max_attempts INT NOT NULL DEFAULT 1,
    metadata JSONB,
    last_attempt_finished_at TIMESTAMPTZ DEFAULT NULL,
    parameters JSONB,
    queue TEXT NOT NULL DEFAULT 'default' :: text,
    scheduled_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    status job_status NOT NULL DEFAULT 'available' :: job_status,
    target TEXT NOT NULL
);


INSERT INTO job_queue
    (attempt, created_at, scheduled_at, max_attempts, metadata, parameters, queue, status, target)
SELECT
    0,
    NOW() - random() * INTERVAL '1 years',
    NOW() - random() * INTERVAL '1 years',
    3,
    '{"team_id": 1, "meta": {}}'::JSONB,
    '{"url": "test.com", "method": "POST"}'::JSONB,
    'default',
    'available'::job_status,
    'test.com'
FROM generate_series(1, 100_000_000);

VACUUM FULL job_queue;

[xvello]

Let's merge as is as sqlx-migrate currently does not support concurrent index operations. We'll block inserts for a couple seconds but plugin-server will retry safely.
We'll need to invest into supporting concurrent index operations for the near future though, as we want the ability to tune these indexes as we iterate on the code.

[tomasfarias]

This will likely block inserts for a bit due to a limitation on sqlx. We are hoping is fine.

The issue is that sqlx runs all migrations in transactions, but CREATE INDEX CONCURRENTLY has to run outside a transaction. So, while using sqlx as our migration tool we are forced to do CREATE INDEX (non-concurrently). We will be looking into alternatives later.

See https://redirect.github.com/launchbadge/sqlx/issues/767 for more details.

[bretthoerner]

Thanks for checking the indexes, this is great.