add queue-state count regression benchmark

This adds a benchmark on top of #1203 to make the queue-state count
query regression easy to reproduce and discuss. It compares the current
`JobCountByQueueAndState` implementation against the legacy query shape
on the same migrated `river_job` schema.

The benchmark stays lightweight by default, but can be scaled locally
with `RIVER_BENCH_QUEUE_STATE_COUNT_NUM_JOBS` to reproduce the planner
regression with a couple hundred thousand rows and quantify the gap.

Author: bgentry

riverdriver/riverpgxv5/queue_state_count_benchmark_test.go

@@ -0,0 +1,209 @@
+package riverpgxv5_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strconv"
+	"testing"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/stretchr/testify/require"
+
+	"github.com/riverqueue/river/riverdbtest"
+	"github.com/riverqueue/river/riverdriver"
+	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+)
+
+const (
+	queueStateCountBenchmarkNumJobsDefault = 20000
+	queueStateCountBenchmarkNumJobsEnvVar  = "RIVER_BENCH_QUEUE_STATE_COUNT_NUM_JOBS"
+)
+
+func BenchmarkJobCountByQueueAndState(b *testing.B) {
+	ctx := context.Background()
+
+	dbPool := riversharedtest.DBPool(ctx, b)
+	driver := riverpgxv5.New(dbPool)
+	schema := riverdbtest.TestSchema(ctx, b, driver, nil)
+	numJobs := queueStateCountBenchmarkNumJobs(b)
+
+	seedQueueStateCountBenchmarkData(ctx, b, dbPool, schema, numJobs)
+
+	queueNamesTwo := queueStateCountBenchmarkQueueNames(2)
+	queueNamesTen := queueStateCountBenchmarkQueueNames(10)
+
+	for _, benchmarkCase := range []struct {
+		name       string
+		queueNames []string
+	}{
+		{name: "TwoQueues", queueNames: queueNamesTwo},
+		{name: "TenQueues", queueNames: queueNamesTen},
+	} {
+		b.Run("Current/"+benchmarkCase.name, func(b *testing.B) {
+			b.ReportAllocs()
+
+			params := &riverdriver.JobCountByQueueAndStateParams{
+				QueueNames: benchmarkCase.queueNames,
+				Schema:     schema,
+			}
+
+			b.ResetTimer()
+			for range b.N {
+				results, err := driver.GetExecutor().JobCountByQueueAndState(ctx, params)
+				require.NoError(b, err)
+				require.NotEmpty(b, results)
+			}
+		})
+
+		b.Run("Legacy/"+benchmarkCase.name, func(b *testing.B) {
+			b.ReportAllocs()
+
+			query := legacyJobCountByQueueAndStateQuery(schema)
+
+			b.ResetTimer()
+			for range b.N {
+				rows, err := dbPool.Query(ctx, query, benchmarkCase.queueNames)
+				require.NoError(b, err)
+
+				var numRows int
+				for rows.Next() {
+					var (
+						countAvailable int64
+						countRunning   int64
+						queue          string
+					)
+
+					require.NoError(b, rows.Scan(&queue, &countAvailable, &countRunning))
+					numRows++
+				}
+
+				rows.Close()
+				require.NoError(b, rows.Err())
+				require.Equal(b, len(benchmarkCase.queueNames), numRows)
+			}
+		})
+	}
+}
+
+func legacyJobCountByQueueAndStateQuery(schema string) string {
+	return fmt.Sprintf(`
+WITH all_queues AS (
+    SELECT DISTINCT unnest($1::text[])::text AS queue
+),
+
+running_job_counts AS (
+    SELECT
+        queue,
+        COUNT(*) AS count
+    FROM %s.river_job
+    WHERE queue = ANY($1::text[])
+        AND state = 'running'
+    GROUP BY queue
+),
+
+available_job_counts AS (
+    SELECT
+        queue,
+        COUNT(*) AS count
+    FROM %s.river_job
+    WHERE queue = ANY($1::text[])
+        AND state = 'available'
+    GROUP BY queue
+)
+
+SELECT
+    all_queues.queue,
+    COALESCE(available_job_counts.count, 0) AS count_available,
+    COALESCE(running_job_counts.count, 0) AS count_running
+FROM
+    all_queues
+LEFT JOIN
+    running_job_counts ON all_queues.queue = running_job_counts.queue
+LEFT JOIN
+    available_job_counts ON all_queues.queue = available_job_counts.queue
+ORDER BY all_queues.queue ASC
+`, schema, schema)
+}
+
+func queueStateCountBenchmarkNumJobs(b *testing.B) int {
+	b.Helper()
+
+	numJobs := queueStateCountBenchmarkNumJobsDefault
+	if numJobsEnv := os.Getenv(queueStateCountBenchmarkNumJobsEnvVar); numJobsEnv != "" {
+		parsedNumJobs, err := strconv.Atoi(numJobsEnv)
+		require.NoError(b, err)
+		require.Greater(b, parsedNumJobs, 0)
+
+		numJobs = parsedNumJobs
+	}
+
+	return numJobs
+}
+
+func queueStateCountBenchmarkQueueNames(numQueues int) []string {
+	queueNames := make([]string, numQueues)
+	for i := range numQueues {
+		queueNames[i] = fmt.Sprintf("queue_%03d", i+1)
+	}
+
+	return queueNames
+}
+
+func seedQueueStateCountBenchmarkData(ctx context.Context, b *testing.B, dbPool *pgxpool.Pool, schema string, numJobs int) {
+	b.Helper()
+
+	query := fmt.Sprintf(`
+WITH generated_jobs AS (
+    SELECT
+        CASE gs %% 8
+            WHEN 0 THEN 'running'
+            WHEN 1 THEN 'available'
+            WHEN 2 THEN 'completed'
+            WHEN 3 THEN 'cancelled'
+            WHEN 4 THEN 'discarded'
+            WHEN 5 THEN 'retryable'
+            WHEN 6 THEN 'scheduled'
+            ELSE 'pending'
+        END AS state,
+        now() - ((gs %% 100000)::text || ' seconds')::interval AS scheduled_at,
+        'queue_' || lpad(((gs %% 100) + 1)::text, 3, '0') AS queue
+    FROM generate_series(1, %d) AS gs
+)
+INSERT INTO %s.river_job (
+    args,
+    finalized_at,
+    kind,
+    max_attempts,
+    metadata,
+    queue,
+    scheduled_at,
+    state
+)
+SELECT
+    '{}'::jsonb,
+    CASE
+        WHEN state IN ('cancelled', 'completed', 'discarded') THEN scheduled_at + interval '1 second'
+        ELSE NULL
+    END AS finalized_at,
+    'benchmark',
+    25,
+    '{}'::jsonb,
+    queue,
+    scheduled_at,
+    state::%s.river_job_state
+FROM generated_jobs;
+
+ANALYZE %s.river_job;
+`, numJobs, schema, schema, schema)
+
+	_, err := dbPool.Exec(ctx, query)
+	require.NoError(b, err)
+
+	row := dbPool.QueryRow(ctx, "SELECT count(*) FROM "+schema+".river_job")
+
+	var numRows int
+	require.NoError(b, row.Scan(&numRows))
+	require.Equal(b, numJobs, numRows)
+}