diff --git a/src/backend/access/index/indexam.c b/src/backend/access/index/indexam.c
index aefdd2916d..cc43300072 100644
--- a/src/backend/access/index/indexam.c
+++ b/src/backend/access/index/indexam.c
@@ -480,8 +480,7 @@ index_beginscan_parallel(Relation heaprel, Relation indexrel, int nkeys,
 	IndexScanDesc scan;
 
 	Assert(RelationGetRelid(heaprel) == pscan->ps_relid);
-	snapshot = RestoreSnapshot(pscan->ps_snapshot_data);
-	RegisterSnapshot(snapshot);
+	snapshot = RegisterSnapshot(RestoreSnapshot(pscan->ps_snapshot_data));
 	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 12adf59085..c4b0e221da 100644
--- a/src/backend/access/table/tableam.c
+++ b/src/backend/access/table/tableam.c
@@ -107,8 +107,7 @@ table_scan_update_snapshot(TableScanDesc scan, Snapshot snapshot)
 {
 	Assert(IsMVCCSnapshot(snapshot));
 
-	RegisterSnapshot(snapshot);
-	scan->rs_snapshot = snapshot;
+	scan->rs_snapshot = RegisterSnapshot(snapshot);
 	scan->rs_flags |= SO_TEMP_SNAPSHOT;
 }
 
@@ -167,7 +166,7 @@ table_beginscan_parallel(Relation relation, ParallelTableScanDesc parallel_scan)
 		/* Snapshot was serialized -- restore it */
 		snapshot = RestoreSnapshot((char *) parallel_scan +
 								   parallel_scan->phs_snapshot_off);
-		RegisterSnapshot(snapshot);
+		snapshot = RegisterSnapshot(snapshot);
 		flags |= SO_TEMP_SNAPSHOT;
 	}
 	else