[#25773] YSQL: fix INSERT ON CONFLICT batching with SPI

Summary:
INSERT ON CONFLICT batching was introduced by commit
a000dbb.  Since then, commit
a120c25 allowed batching for INSERT ON
CONFLICT with RETURNING clause.  That breaks cases where the INSERT ON
CONFLICT query is sent using SPI_execute_with_args with a positive
tcount.  That tcount can cause ExecutePlan to stop calling the insert's
ExecModifyTable partway and then go to ExecEndModifyTable.  In that
case, ExecEndModifyTable is called with a partially processed insert on
conflict state, which is unexpected.  Modify expectations to allow this
and free up such state.

I could not find a way to get ExecutePlan to be called with numberTuples
> 0 besides via SPI_execute_with_args, so for testing, create a C
function yb_run_spi in regress.c that calls SPI_execute_with_args on the
given query and returns the number of rows returned.

As part of this, add checks to make sure that slots are dropped before
clearing them from the C++ maps.

Jira: DB-15059

Test Plan:
On Almalinux 8:

    #!/usr/bin/env bash
    set -euo pipefail
    ./yb_build.sh fastdebug --gcc11
    find java/yb-pgsql/src/test/java/org/yb/pgsql -name 'TestPgRegressInsertOnConflict*' \
    | grep -oE 'TestPgRegress\w+' \
    | while read -r testname; do
      ./yb_build.sh fastdebug --gcc11 --java-test "$testname" --sj
    done

Close: #25773

Reviewers: kramanathan

Reviewed By: kramanathan

Subscribers: yql

Differential Revision: https://phorge.dev.yugabyte.com/D41471

src/postgres/src/backend/executor/nodeModifyTable.c

@@ -262,6 +262,7 @@ static TupleTableSlot *YbExecInsertAct(ModifyTableContext *context,
 									   ResultRelInfo **insert_destrel);
 
 static void YbFreeInsertOnConflictBatchState(YbInsertOnConflictBatchState *state);
+static void YbDropInsertOnConflictReadSlots(YbInsertOnConflictBatchState *state);
 
 static Bitmapset *YbFetchColumnsMarkedForUpdate(ModifyTableContext *context,
 												ResultRelInfo *resultRelInfo);
@@ -5279,15 +5280,7 @@ ExecEndModifyTable(ModifyTableState *node)
 	 * CONFLICT batching happens.
 	 */
 	if (node->yb_ioc_state)
-	{
-		/*
-		 * First, free up the key cache (and intents cache if possible)
-		 * corresponding to this state.
-		 */
-		YBCPgClearInsertOnConflictCache(node->yb_ioc_state);
-
 		YbFreeInsertOnConflictBatchState(node->yb_ioc_state);
-	}
 
 	/*
 	 * Allow any FDWs to shut down
@@ -5389,12 +5382,58 @@ static void
 YbFreeInsertOnConflictBatchState(YbInsertOnConflictBatchState *state)
 {
 	Assert(state);
-	Assert(state->pending_relids == NIL);
+
+	/*
+	 * Drop input slots and read slots (this should only be the case when
+	 * ExecutePlan numberTuples is nonzero and reached).
+	 */
+	list_free(state->pending_relids);
+	for (int i = 0; i < state->num_slots; ++i)
+	{
+		if (state->slots[i])
+		{
+			ExecDropSingleTupleTableSlot(state->slots[i]);
+			Assert(state->planSlots[i]);
+			ExecDropSingleTupleTableSlot(state->planSlots[i]);
+		}
+	}
 	pfree(state->slots);
 	pfree(state->planSlots);
+	YbDropInsertOnConflictReadSlots(state);
+
+	/*
+	 * Remove map from maps.
+	 */
+	YBCPgClearInsertOnConflictCache(state);
+
+	/*
+	 * Free state.
+	 */
 	pfree(state);
 }
 
+static void
+YbDropInsertOnConflictReadSlots(YbInsertOnConflictBatchState *state)
+{
+	/*
+	 * Drop the slots stored by the batch-read operation one by one. Doing
+	 * this here (rather than in pggate) keeps all the slot management code
+	 * in one place: slots are allocated in execIndexing and deallocated in
+	 * this function.
+	 */
+	YbcPgInsertOnConflictKeyInfo info = {NULL};
+	const uint64_t key_count = YBCPgGetInsertOnConflictKeyCount(state);
+
+	for (uint64_t i = 0; i < key_count; i++)
+	{
+		HandleYBStatus(YBCPgDeleteNextInsertOnConflictKey(state,
+														  &info));
+		if (info.slot)
+			ExecDropSingleTupleTableSlot(info.slot);
+	}
+	Assert(YBCPgGetInsertOnConflictKeyCount(state) == 0);
+}
+
 /*
  * Parts copied from FDW batching.
  */
@@ -5560,24 +5599,7 @@ YbFlushSlotsFromBatch(ModifyTableContext *context,
 		state->flush_idx = 0;
 	}
 
-	/*
-	 * Drop the slots stored by the batch-read operation one by one. Doing
-	 * this here (rather than in pggate) keeps all the slot management code
-	 * in one place: slots are allocated in execIndexing and deallocated in
-	 * this function.
-	 */
-	YbcPgInsertOnConflictKeyInfo info = {NULL};
-	const uint64_t key_count = YBCPgGetInsertOnConflictKeyCount(state);
-
-	for (uint64_t i = 0; i < key_count; i++)
-	{
-		HandleYBStatus(YBCPgDeleteNextInsertOnConflictKey(state,
-														  &info));
-		if (info.slot)
-			ExecDropSingleTupleTableSlot(info.slot);
-	}
-	Assert(YBCPgGetInsertOnConflictKeyCount(state) == 0);
-
+	YbDropInsertOnConflictReadSlots(state);
 	state->num_slots = 0;
 	Assert(!state->flush_idx);
 	Assert(state->pending_relids == NIL);

src/postgres/src/test/regress/expected/yb_create_function.out

@@ -16,3 +16,7 @@ CREATE FUNCTION bigname_out(bigname)
    AS :'regresslib'
    LANGUAGE C STRICT IMMUTABLE;
 NOTICE:  argument type bigname is only a shell
+CREATE FUNCTION yb_run_spi(text, int)
+   RETURNS int8
+   AS :'regresslib'
+   LANGUAGE C STRICT IMMUTABLE;

src/postgres/src/test/regress/expected/yb_insert_on_conflict.out

@@ -175,6 +175,79 @@ ERROR:  invalid reference to FROM-clause entry for table "excluded"
 LINE 1: ...NFLICT (a) DO UPDATE SET b = EXCLUDED.a RETURNING EXCLUDED.b...
                                                              ^
 HINT:  There is an entry for table "excluded", but it cannot be referenced from this part of the query.
+--- SPI
+BEGIN;
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (1) ON CONFLICT DO NOTHING RETURNING a
+$$, 1);
+ yb_run_spi
+------------
+          0
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    10 |   1 |  10
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (null) ON CONFLICT DO NOTHING RETURNING a
+$$, 1);
+ yb_run_spi
+------------
+          1
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    11 |   1 |  10
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (generate_series(-10, 20)) ON CONFLICT DO NOTHING RETURNING b
+$$, 100);
+ yb_run_spi
+------------
+         21
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    32 | -10 |  20
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (generate_series(-20, 30)) ON CONFLICT DO NOTHING RETURNING b, a
+$$, 15);
+ yb_run_spi
+------------
+         15
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    47 | -20 |  25
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (generate_series(-30, 40)) ON CONFLICT (a) DO UPDATE SET b = EXCLUDED.a RETURNING (a + b)
+$$, 45);
+ yb_run_spi
+------------
+         45
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    57 | -30 |  25
+(1 row)
+
+ABORT;
 --- Multiple arbiter indexes
 CREATE UNIQUE INDEX NONCONCURRENTLY br_idx ON ab_tab (b ASC);
 -- No constraint specification.

src/postgres/src/test/regress/expected/yb_insert_on_conflict_1.out

@@ -175,6 +175,79 @@ ERROR:  invalid reference to FROM-clause entry for table "excluded"
 LINE 1: ...NFLICT (a) DO UPDATE SET b = EXCLUDED.a RETURNING EXCLUDED.b...
                                                              ^
 HINT:  There is an entry for table "excluded", but it cannot be referenced from this part of the query.
+--- SPI
+BEGIN;
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (1) ON CONFLICT DO NOTHING RETURNING a
+$$, 1);
+ yb_run_spi
+------------
+          0
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    10 |   1 |  10
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (null) ON CONFLICT DO NOTHING RETURNING a
+$$, 1);
+ yb_run_spi
+------------
+          1
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    11 |   1 |  10
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (generate_series(-10, 20)) ON CONFLICT DO NOTHING RETURNING b
+$$, 100);
+ yb_run_spi
+------------
+         21
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    32 | -10 |  20
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (generate_series(-20, 30)) ON CONFLICT DO NOTHING RETURNING b, a
+$$, 15);
+ yb_run_spi
+------------
+         15
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    47 | -20 |  25
+(1 row)
+
+SELECT yb_run_spi($$
+    INSERT INTO ab_tab VALUES (generate_series(-30, 40)) ON CONFLICT (a) DO UPDATE SET b = EXCLUDED.a RETURNING (a + b)
+$$, 45);
+ yb_run_spi
+------------
+         45
+(1 row)
+
+SELECT count(*), min(a), max(a) FROM ab_tab;
+ count | min | max
+-------+-----+-----
+    57 | -30 |  25
+(1 row)
+
+ABORT;
 --- Multiple arbiter indexes
 CREATE UNIQUE INDEX NONCONCURRENTLY br_idx ON ab_tab (b ASC);
 -- No constraint specification.