Streamline adaptive executor for single task queries

[colm-mchugh]

DESCRIPTION: New scan path for fast path queries that is a streamlined version of adaptive executor specialized for one task.

Introduces SingleTaskAdaptiveExecutor, a streamlined execution path for single-shard fast-path router queries that bypasses the full adaptive executor's WaitEventSet machinery, pool management, and slow-start logic.

Eligibility: fastPathRouterPlan && citus.enable_single_task_execution GUC (default on) && not multi-row INSERT && no dependent jobs && single-replicated table. EXPLAIN ANALYZE falls back to the regular AdaptiveExecutor at execution time.

Uses simple WaitLatchOrSocket polling on a single connection instead of the multi-connection WaitEventSet infrastructure.

Multi-replicated tables (shard_replication_factor > 1) are excluded because SingleTaskAdaptiveExecutor only uses the first placement, which would silently skip replica writes. This covers modifications to reference tables.

New GUC: citus.enable_single_task_execution (PGC_USERSET, default true)
New executor type: MULTI_EXECUTOR_SINGLE_TASK
New scan name: "Citus Single Task"

[codecov]

Codecov Report

❌ Patch coverage is 84.77612% with 51 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.71%. Comparing base (f91592f) to head (812fb54).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8607      +/-   ##
==========================================
- Coverage   88.73%   88.71%   -0.02%     
==========================================
  Files         288      288              
  Lines       64428    64725     +297     
  Branches     8113     8169      +56     
==========================================
+ Hits        57167    57424     +257     
- Misses       4917     4941      +24     
- Partials     2344     2360      +16     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[sfc-gh-mslot]

Doesn't this create a wait event set in every call? If so, what's the gain here? Seems like a lot of code to maintain.

[colm-mchugh]

Yes, it creates a WaitEventSet per call. This path doesn't fire when results are buffered on entering the loop. But for large result sets it could accumulate overhead so will likely change to create lazily, when the socket is not immediately readable.

what's the gain here?

The gain comes from bypassing the pool/session/connection-state-machine infrastructure, not the wait mechanics. Running pgbench on CPU-bound workload shows TPS gain of 15% (SELECT), 12% (INSERT) and 12% (UPDATE) so for an in-memory OLTP workload there's a substantial gain.

[Copilot]

Pull request overview

This PR introduces a new streamlined execution path (“SingleTaskAdaptiveExecutor”) for eligible single-shard fast-path router queries, aiming to reduce overhead by bypassing the full adaptive executor machinery while preserving EXPLAIN visibility and maintaining correctness via eligibility checks and fallbacks.

Changes:

• Add a new executor type (MULTI_EXECUTOR_SINGLE_TASK) plus custom scan plumbing to run a dedicated single-task execution path for eligible fast-path router queries.
• Introduce a new GUC (citus.enable_single_task_execution, default on) and update plan selection logic to choose the single-task executor when safe.
• Expand regression coverage with a new single_task_execution test, update multiple expected files for the new scan name, and add normalization to reduce output flakiness.

Reviewed changes

Copilot reviewed 43 out of 45 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/backend/distributed/executor/adaptive_executor.c	Implements SingleTaskExecutor() and supporting logic for single-connection polling and result ingestion.
src/backend/distributed/executor/citus_custom_scan.c	Registers the new custom scan methods and routes execution to SingleTaskExecutor (or falls back for EXPLAIN ANALYZE).
src/backend/distributed/executor/multi_server_executor.c	Chooses MULTI_EXECUTOR_SINGLE_TASK for eligible fast-path router plans (with replication/eligibility checks).
src/backend/distributed/planner/distributed_planner.c	Wires MULTI_EXECUTOR_SINGLE_TASK to the new custom scan methods during plan finalization.
src/backend/distributed/shared_library_init.c	Adds the citus.enable_single_task_execution GUC definition.
src/include/distributed/adaptive_executor.h	Exposes the SingleTaskExecutor(CitusScanState *) entry point.
src/include/distributed/citus_custom_scan.h	Declares SingleTaskExecutorCustomScanMethods.
src/include/distributed/multi_server_executor.h	Adds MULTI_EXECUTOR_SINGLE_TASK and the EnableSingleTaskExecution GUC variable.
src/test/regress/bin/normalize.sed	Normalizes repartition-related NOTICE output (hash boundaries, suffixes, shard id ranges) to reduce flakiness.
src/test/regress/citus_tests/run_test.py	Adjusts test deps (adds multi_transaction_recovery to minimal_schedule).
src/test/regress/multi_1_schedule	Schedules the new single_task_execution test in the non-parallelized block with other router prepared-statement tests.
src/test/regress/sql/single_task_execution.sql	New regression test validating explain output and result equivalence with the GUC toggled.
src/test/regress/sql/multi_transaction_recovery.sql	Makes the test more robust (cleanup + dynamic group IDs) and disables single-task exec for predictable 2PC behavior.
src/test/regress/sql/multi_router_planner_fast_path.sql	Forces custom plans for prepared statements to avoid flaky debug output; resets afterward.
src/test/regress/sql/local_shard_execution.sql	Makes colocations explicit; disables single-task for predictable error messages; wraps flaky error cases in DO/EXCEPTION checks.
src/test/regress/sql/failure_single_select.sql	Disables single-task exec for predictable error behavior in failure tests.
src/test/regress/sql/failure_savepoints.sql	Disables single-task exec for predictable error behavior in failure tests.
src/test/regress/sql/failure_multi_dml.sql	Disables single-task exec for predictable error behavior; ensures worker-side script also disables it.
src/test/regress/sql/failure_connection_establishment.sql	Disables single-task exec for predictable error behavior in connection-establishment tests.
src/test/regress/expected/single_task_execution.out	New expected output file for the new regression test.
src/test/regress/expected/upgrade_basic_after.out	Updates expected EXPLAIN scan name for single-task-eligible plans.
src/test/regress/expected/upgrade_basic_after_0.out	Adds an additional expected output variant for upgrade coverage.
src/test/regress/expected/sqlancer_failures.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/multi_router_planner_fast_path.out	Updates expected debug output + plan_cache_mode SET/RESET lines.
src/test/regress/expected/multi_transaction_recovery.out	Updates expected output for cleanup, dynamic group IDs, and the single-task disable.
src/test/regress/expected/local_shard_execution_replicated.out	Updates expected NOTICE output after normalization changes.
src/test/regress/expected/failure_single_select.out	Mirrors the SQL change disabling single-task exec.
src/test/regress/expected/failure_savepoints.out	Mirrors the SQL change disabling single-task exec.
src/test/regress/expected/failure_multi_dml.out	Mirrors the SQL change disabling single-task exec; expected worker script output changes.
src/test/regress/expected/failure_connection_establishment.out	Mirrors the SQL change disabling single-task exec.
src/test/regress/expected/multi_mx_insert_select_repartition.out	Updates expected NOTICE output after normalization changes.
src/test/regress/expected/multi_mx_explain.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/multi_mx_explain_0.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/multi_explain.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/multi_explain_0.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/multi_tenant_isolation.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/multi_tenant_isolation_nonblocking.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/pg13.out	Updates expected EXPLAIN/ANALYZE output formatting and scan name for single-task-eligible plan(s).
src/test/regress/expected/pg15.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/pg18.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable; aligns formatting for pg_get_loaded_modules output.
src/test/regress/expected/multi_data_types.out	Updates expected EXPLAIN/ANALYZE scan name + spacing for single-task-eligible DML.
src/test/regress/expected/drop_column_partitioned_table.out	Updates expected EXPLAIN scan name to “Citus Single Task” where applicable.
src/test/regress/expected/coordinator_shouldhaveshards.out	Updates expected NOTICE output after normalization changes.
src/test/regress/expected/single_node.out	Updates expected NOTICE output after normalization changes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[colm-mchugh]

The !OidIsValid(relationId) check covers SELECT queries, which always have targetRelationId = InvalidOid. Reading from a single placement is safe regardless of replication factor. For DML, targetRelationId is always set, so the SingleReplicatedTable() check correctly routes multi-replicated tables to the adaptive executor for write replication.

[onurctirtir]

This review was produced by Claude Opus 4.8 (model ID claude-opus-4.8), driven by @onurctirtir.

Every finding below was reproduced empirically on an isolated 2-worker Citus cluster built from this PR branch (colm/1-task-AE-take-2).

Methodology: read the STE path against the README + the adaptive executor it forks from, then tried to break it — focusing on the control-flow the new SingleTaskExecutorStart fast path skips relative to the full adaptive executor.

[codeforall]

No correctness concerns from me. Just a few nits, all take-it-or-leave-it:

• The STE path sends only TaskQueryStringAtIndex(task, 0) with no loop over task->queryCount. It's safe today because every multi-query shape (multi-row INSERT, modifying CTEs, etc.) is excluded upstream, and I couldn't get a queryCount > 1 task onto the path in testing , but nothing enforces it, and if a future fast-path change ever slipped one through, the extra queries would be silently dropped. A one-line Assert(task->queryCount == 1) in SingleTaskExecutorStart would lock the invariant cheaply.

• Two small test additions would be nice since I had to build a cluster to confirm the behavior: a prepared-statement param case (the deferred-pruning path Copilot asked about) and an EXPLAIN ANALYZE counter-is-exactly-one assertion. Both already work; it'd just pin them.

• A tiny thing, the refactor dropped the old XXX: SendRemoteCommandParams … isolation_select_vs_all.spec breadcrumb in SendNextQuery, might be worth keeping in SendQueryStringToConnection.

Nothing blocking here, LGTM. Thanks for the thorough comments and the strong test coverage, made this an easy read.

[Copilot]

Pull request overview

Copilot reviewed 43 out of 45 changed files in this pull request and generated 2 comments.