Sorted merge: Fix latent bugs in SortedMergeAdapterRescan

[neildsh]

Fix two latent bugs in SortedMergeAdapterRescan

Summary

This PR fixes two correctness/robustness bugs in the sorted-merge rescan callback that surfaced during a pen test of #8529. Both bugs live in code that is unreachable in production because the materialize_finished_plan wrapper added in #8529 absorbs every rescan that PostgreSQL would otherwise forward to the sorted-merge CustomScan. The fixes harden the defensive rescan path so it produces correct results — and does not crash — if the absorber is ever weakened, removed, or bypassed by a future PG release.

Background

SortedMergeAdapter streams globally-sorted tuples by running a binary min-heap over per-task tuplestores (sorted_merge.c). SortedMergeAdapterNext() has two branches:

• First call (!initialized): rewind every per-task store, seed the heap with the first tuple from each, then return the heap's current winner.
• Subsequent calls (initialized): advance the previous winner's store, update the heap, then return the new winner.

SortedMergeAdapterRescan() is the rescan callback installed on the sorted-merge CustomScan. PostgreSQL invokes it from ExecutorRewind() on the only SQL shape that reaches it: a SCROLL CURSOR WITH HOLD whose holdable portal is being persisted by PersistHoldablePortal() at COMMIT time. In production, the Material node inserted by materialize_finished_plan() in distributed_planner.c absorbs that rescan, so the sorted-merge callback never runs.

Bug 1 — off-by-one tuple drop after rescan (correctness)

Where: src/backend/distributed/executor/sorted_merge.c::SortedMergeAdapterRescan

The original implementation pre-seeded the heap with the first tuple from each store and then set initialized = true:

for (int i = 0; i < adapter->nstores; i++)
{
    tuplestore_rescan(adapter->perTaskStores[i]);
    if (tuplestore_gettupleslot(adapter->perTaskStores[i], true, false,
                                adapter->mergeCtx.slots[i]))
    {
        binaryheap_add_unordered(adapter->heap, Int32GetDatum(i));
    }
}
binaryheap_build(adapter->heap);

adapter->exhausted = binaryheap_empty(adapter->heap);
adapter->initialized = true;            /* ← problem */

Because the adapter is now marked as initialized, the next SortedMergeAdapterNext() call enters the else branch, treats the freshly-seeded heap winner as a previous winner, advances that store and only then returns the new winner. The first globally-minimal tuple after every rescan is silently dropped.

Reproduction (with the Material absorber temporarily removed)

CREATE TABLE t (id int);
SELECT create_distributed_table('t', 'id');
INSERT INTO t SELECT g FROM generate_series(1, 20) g;

BEGIN;
DECLARE c SCROLL CURSOR WITH HOLD FOR
    SELECT id FROM t ORDER BY id;
FETCH 3 FROM c;   -- 1, 2, 3
COMMIT;           -- PersistHoldablePortal → ExecutorRewind → rescan
FETCH ALL FROM c; -- expected: 4..20  (17 rows)
                  -- actual:   5..20  (16 rows; row 4 dropped)
CLOSE c;

Fix

Don't seed the heap in Rescan. Mark the adapter as uninitialized and let the next Next() call do its normal first-call seed via the existing if (!adapter->initialized) branch, which already rewinds each store, seeds the heap, and returns the winner instead of advancing past it:

void
SortedMergeAdapterRescan(SortedMergeAdapter *adapter)
{
    binaryheap_reset(adapter->heap);

    /* Let the next Next() reseed.  Zero-store adapters stay exhausted. */
    adapter->initialized = (adapter->nstores == 0);
    adapter->exhausted   = (adapter->nstores == 0);
}

Bug 2 — NULL-deref crash when rescan fires before any FETCH

Where: src/backend/distributed/executor/citus_custom_scan.c::SortedMergeReScan

The sorted-merge CustomScan runs AdaptiveExecutor() (and thus CreatePerTaskDispatchDests(), which installs scanState->mergeAdapter) lazily — only on the first tuple fetched from the scan. So a SCROLL CURSOR WITH HOLD that is declared, never fetched from, and then committed will reach SortedMergeReScan() with scanState->mergeAdapter == NULL.

The previous code dereferenced it unconditionally:

static void
SortedMergeReScan(CustomScanState *node)
{
    elog(WARNING, "unexpected sorted merge rescan");
    CitusReScanCommon(node);
    CitusScanState *scanState = (CitusScanState *) node;
    SortedMergeAdapterRescan(scanState->mergeAdapter);  /* SEGV when NULL */
}

Reproduction (with the Material absorber temporarily removed)

BEGIN;
DECLARE c SCROLL CURSOR WITH HOLD FOR
    SELECT id FROM t ORDER BY id;
FETCH 0 FROM c;        -- no FETCH ⇒ AdaptiveExecutor never runs
COMMIT;                -- PersistHoldablePortal → rescan → NULL deref ⇒ SIGSEGV

Fix

Guard the adapter pointer before calling rescan, mirroring the existing CitusReScan() pattern that already guards scanState->tuplestorestate:

static void
SortedMergeReScan(CustomScanState *node)
{
    elog(WARNING, "unexpected sorted merge rescan");
    CitusReScanCommon(node);
    CitusScanState *scanState = (CitusScanState *) node;

    /*
     * The adapter is created lazily by AdaptiveExecutor on the first scan
     * fetch.  If rescan fires before any tuple was fetched (e.g. DECLARE
     * followed by COMMIT without any FETCH), there is nothing to rewind.
     */
    if (scanState->mergeAdapter != NULL)
    {
        SortedMergeAdapterRescan(scanState->mergeAdapter);
    }
}

Test coverage

Added seven new sub-tests (N2..N7) to the existing Category N (SCROLL CURSOR WITH HOLD rescan) section of src/test/regress/sql/multi_orderby_pushdown.sql, alongside the pre-existing N1 test. These document the exact SQL shapes that reproduced the bugs and act as regression guards:

Test	Shape	Bug it would surface (if Material absorber is removed)
N2	WITH HOLD + FETCH 3 + COMMIT + FETCH ALL, no LIMIT	Bug 1 — row at rescan boundary dropped
N3	WITH HOLD + FETCH 0 + COMMIT + FETCH ALL	Bug 2 — NULL adapter ⇒ SIGSEGV
N4	WITH HOLD + multi-key ORDER BY id%3, id + rescan	Multi-key comparator correctness across rescan
N5	WITH HOLD + ORDER BY id DESC + rescan	Inverted comparator correctness across rescan
N6	WITH HOLD + repeated single-row fetches after rescan	holdStore sequencing post-rescan
N7	WITH HOLD over empty result + rescan	Empty-store rescan: no phantoms, no heap underflow

In production (with materialize_finished_plan in place) all of these queries already returned correct results because the Material node absorbed the rescan. They are still kept in the suite as documentation and to catch any future regression where the absorber is changed, weakened, or bypassed by a new PG release. With the absorber temporarily removed during pen-testing, N2/N4/N5/N6 reproduce the row-drop bug and N3 reproduces the crash; with the fixes applied (and the absorber still removed) all seven produce correct results.

Impact

• Correctness: No behaviour change visible to users of the unmodified PR. Both bugs were latent under the materialize_finished_plan absorber, so no released sorted-merge plan can hit them today.
• Defensive hardening: The rescan callback is now self-consistent and crash-safe. A future PG-side change (e.g. the planner choosing not to wrap, or a new portal type that bypasses the wrapper) would no longer silently produce wrong results or crash.
• No performance change. The rescan callback's hot work is moved from Rescan() into the next Next() call (where the existing uninitialized-path code already lives). Zero extra work either way.

Files changed

• src/backend/distributed/executor/sorted_merge.c — Fix Bug 1 in SortedMergeAdapterRescan.
• src/backend/distributed/executor/citus_custom_scan.c — Fix Bug 2 in SortedMergeReScan.
• src/test/regress/sql/multi_orderby_pushdown.sql — Add N2..N7 tests under the existing Category N.
• src/test/regress/expected/multi_orderby_pushdown.out — Expected output for the new tests.

Verification

• make -C src/test/regress check-base EXTRA_TESTS='multi_orderby_pushdown' passes (15/15 tests, including the new N2..N7).
• Bug repro under a manually patched build that removes the materialize_finished_plan call (used only locally for verification — not part of this PR):
  • Before fix: N2 returns 16 rows instead of 17 (row 4 missing); N3 SIGSEGVs in SortedMergeAdapterRescan on a NULL adapter.
  • After fix: N2..N7 all return correct results.