0006-WIP-Make-RestoreSnapshot-register-the-snapshot-with-.patch

From f542826b420d41d18dc82d068f774559a4d8f649 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Mon, 31 Mar 2025 19:54:39 +0300
Subject: [PATCH 6/9] WIP: Make RestoreSnapshot register the snapshot with
 current resowner

This simplifies the next commit
---
 src/backend/access/index/indexam.c    | 1 -
 src/backend/access/table/tableam.c    | 2 --
 src/backend/access/transam/parallel.c | 4 ++++
 src/backend/utils/time/snapmgr.c      | 9 +++++++--
 4 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/backend/access/index/indexam.c b/src/backend/access/index/indexam.c
index 7ef5031d716..3828622665c 100644
--- a/src/backend/access/index/indexam.c
+++ b/src/backend/access/index/indexam.c
@@ -603,7 +603,6 @@ index_beginscan_parallel(Relation heaprel, Relation indexrel,
 	Assert(RelFileLocatorEquals(indexrel->rd_locator, pscan->ps_indexlocator));
 
 	snapshot = (Snapshot) RestoreSnapshot(pscan->ps_snapshot_data);
-	snapshot = RegisterSnapshot(snapshot);
 	scan = index_beginscan_internal(indexrel, nkeys, norderbys, snapshot,
 									pscan, true);
 
diff --git a/src/backend/access/table/tableam.c b/src/backend/access/table/tableam.c
index fde28accfd3..c8db2918f40 100644
--- a/src/backend/access/table/tableam.c
+++ b/src/backend/access/table/tableam.c
@@ -175,7 +175,6 @@ table_beginscan_parallel(Relation relation, ParallelTableScanDesc pscan)
 	{
 		/* Snapshot was serialized -- restore it */
 		snapshot = (Snapshot) RestoreSnapshot((char *) pscan + pscan->phs_snapshot_off);
-		snapshot = RegisterSnapshot(snapshot);
 		flags |= SO_TEMP_SNAPSHOT;
 	}
 	else
@@ -205,7 +204,6 @@ table_beginscan_parallel_tidrange(Relation relation,
 	{
 		/* Snapshot was serialized -- restore it */
 		snapshot = (Snapshot) RestoreSnapshot((char *) pscan + pscan->phs_snapshot_off);
-		RegisterSnapshot(snapshot);
 		flags |= SO_TEMP_SNAPSHOT;
 	}
 	else
diff --git a/src/backend/access/transam/parallel.c b/src/backend/access/transam/parallel.c
index 1fd2146358d..0bd3e65f849 100644
--- a/src/backend/access/transam/parallel.c
+++ b/src/backend/access/transam/parallel.c
@@ -1506,6 +1506,10 @@ ParallelWorkerMain(Datum main_arg)
 							   fps->parallel_leader_pgproc);
 	PushActiveSnapshot(asnapshot);
 
+	UnregisterSnapshot(asnapshot);
+	if (tsnapshot != asnapshot)
+		UnregisterSnapshot(tsnapshot);
+
 	/*
 	 * We've changed which tuples we can see, and must therefore invalidate
 	 * system caches.
diff --git a/src/backend/utils/time/snapmgr.c b/src/backend/utils/time/snapmgr.c
index 0ab81736742..e9a42437828 100644
--- a/src/backend/utils/time/snapmgr.c
+++ b/src/backend/utils/time/snapmgr.c
@@ -518,7 +518,6 @@ SetTransactionSnapshot(MVCCSnapshot sourcesnap, VirtualTransactionId *sourcevxid
 	/* Better do this to ensure following Assert succeeds. */
 	InvalidateCatalogSnapshot();
 
-	Assert(pairingheap_is_empty(&RegisteredSnapshots));
 	Assert(!FirstXactSnapshotRegistered);
 	Assert(!HistoricSnapshotActive());
 
@@ -1831,7 +1830,7 @@ SerializeSnapshot(MVCCSnapshot snapshot, char *start_address)
  *		Restore a serialized snapshot from the specified address.
  *
  * The copy is palloc'd in TopTransactionContext and has initial refcounts set
- * to 0.  The returned snapshot has the copied flag set.
+ * to 0.  The returned snapshot is registered with the current resource owner.
  */
 MVCCSnapshot
 RestoreSnapshot(char *start_address)
@@ -1888,6 +1887,12 @@ RestoreSnapshot(char *start_address)
 	snapshot->copied = true;
 	snapshot->valid = true;
 
+	/* and tell resowner.c about it, just like RegisterSnapshot() */
+	ResourceOwnerEnlarge(CurrentResourceOwner);
+	snapshot->regd_count++;
+	ResourceOwnerRememberSnapshot(CurrentResourceOwner, (Snapshot) snapshot);
+	pairingheap_add(&RegisteredSnapshots, &snapshot->ph_node);
+
 	return snapshot;
 }
 
-- 
2.47.3