From b3555db35a5e5dea9f19b4c772ba32877c67db29 Mon Sep 17 00:00:00 2001
From: Alexander Korotkov <akorotkov@postgresql.org>
Date: Fri, 29 Nov 2024 09:42:32 +0200
Subject: [PATCH] Fix assigning subtransaction logical xids

---
 expected/subtransactions.out | 23 +++++++++++++++++++++++
 sql/subtransactions.sql      | 22 ++++++++++++++++++++++
 src/transam/oxid.c           |  7 ++++++-
 src/transam/undo.c           |  2 +-
 4 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/expected/subtransactions.out b/expected/subtransactions.out
index ec1fbbbe..795dd033 100644
--- a/expected/subtransactions.out
+++ b/expected/subtransactions.out
@@ -770,6 +770,29 @@ FETCH FROM abc;
 -------
 (0 rows)
 
+ROLLBACK;
+-- Check for logical xid overflow
+BEGIN;
+DO $$
+BEGIN
+    FOR i IN 1..1000 LOOP
+        BEGIN
+	    RAISE reading_sql_data_not_permitted;
+	EXCEPTION WHEN reading_sql_data_not_permitted THEN END;
+    END LOOP;
+END;$$;
+ROLLBACK;
+BEGIN;
+DO $$
+BEGIN
+    FOR i IN 1..100000 LOOP
+        BEGIN
+	    RAISE reading_sql_data_not_permitted;
+	EXCEPTION WHEN reading_sql_data_not_permitted THEN END;
+    END LOOP;
+END;$$;
+ERROR:  not enough logical xids
+CONTEXT:  PL/pgSQL function inline_code_block line 4 during statement block entry
 ROLLBACK;
 DROP EXTENSION orioledb CASCADE;
 NOTICE:  drop cascades to table o_subtrans
diff --git a/sql/subtransactions.sql b/sql/subtransactions.sql
index db26f254..f62549b9 100644
--- a/sql/subtransactions.sql
+++ b/sql/subtransactions.sql
@@ -413,6 +413,28 @@ ROLLBACK TO s1;
 FETCH FROM abc;
 ROLLBACK;
 
+-- Check for logical xid overflow
+BEGIN;
+DO $$
+BEGIN
+    FOR i IN 1..1000 LOOP
+        BEGIN
+	    RAISE reading_sql_data_not_permitted;
+	EXCEPTION WHEN reading_sql_data_not_permitted THEN END;
+    END LOOP;
+END;$$;
+ROLLBACK;
+BEGIN;
+DO $$
+BEGIN
+    FOR i IN 1..100000 LOOP
+        BEGIN
+	    RAISE reading_sql_data_not_permitted;
+	EXCEPTION WHEN reading_sql_data_not_permitted THEN END;
+    END LOOP;
+END;$$;
+ROLLBACK;
+
 DROP EXTENSION orioledb CASCADE;
 DROP SCHEMA subtransactions CASCADE;
 RESET search_path;
diff --git a/src/transam/oxid.c b/src/transam/oxid.c
index 20816b1f..c3a4b7bf 100644
--- a/src/transam/oxid.c
+++ b/src/transam/oxid.c
@@ -260,6 +260,7 @@ acquire_logical_xid(void)
 				if (nloops > 16)
 					elog(ERROR, "not enough logical xids");
 			}
+			continue;
 		}
 		bitnum = pg_ceil_log2_32(bit);
 		Assert(bit == (1 << bitnum));
@@ -312,6 +313,10 @@ release_logical_xid(TransactionId xid)
 void
 assign_subtransaction_logical_xid(void)
 {
+	TransactionId nextLogicalXid;
+
+	nextLogicalXid = acquire_logical_xid();
+
 	if (TransactionIdIsValid(logicalXid))
 	{
 		MemoryContext mcxt = MemoryContextSwitchTo(TopMemoryContext);
@@ -324,7 +329,7 @@ assign_subtransaction_logical_xid(void)
 		MemoryContextSwitchTo(mcxt);
 	}
 
-	logicalXid = acquire_logical_xid();
+	logicalXid = nextLogicalXid;
 }
 
 /*
diff --git a/src/transam/undo.c b/src/transam/undo.c
index c257d69b..179d2e5a 100644
--- a/src/transam/undo.c
+++ b/src/transam/undo.c
@@ -1565,9 +1565,9 @@ undo_subxact_callback(SubXactEvent event, SubTransactionId mySubid,
 	{
 		case SUBXACT_EVENT_START_SUB:
 			(void) get_current_oxid();
+			add_subxact_undo_item(parentSubid);
 			prentLogicalXid = get_current_logical_xid();
 			assign_subtransaction_logical_xid();
-			add_subxact_undo_item(parentSubid);
 			add_savepoint_wal_record(parentSubid, prentLogicalXid);
 			break;
 		case SUBXACT_EVENT_COMMIT_SUB: