From 09b7bcd6c8e3fbc9438c6edf1aac75a55b3909be Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Thu, 24 Oct 2024 17:34:57 -0700
Subject: [PATCH v3 3/4] Support shared itereation on TidStore.

Author:
Reviewed-by:
Discussion: https://postgr.es/m/
Backpatch-through:
---
 src/backend/access/common/tidstore.c | 59 ++++++++++++++++++++++++++++
 src/include/access/tidstore.h        |  3 ++
 2 files changed, 62 insertions(+)

diff --git a/src/backend/access/common/tidstore.c b/src/backend/access/common/tidstore.c
index a7179759d6..637d26012d 100644
--- a/src/backend/access/common/tidstore.c
+++ b/src/backend/access/common/tidstore.c
@@ -483,6 +483,7 @@ TidStoreBeginIterate(TidStore *ts)
 	iter = palloc0(sizeof(TidStoreIter));
 	iter->ts = ts;
 
+	/* begin iteration on the radix tree */
 	if (TidStoreIsShared(ts))
 		iter->tree_iter.shared = shared_ts_begin_iterate(ts->tree.shared);
 	else
@@ -533,6 +534,56 @@ TidStoreEndIterate(TidStoreIter *iter)
 	pfree(iter);
 }
 
+/*
+ * Prepare to iterate through a shared TidStore in shared mode. This function
+ * is aimed to start the iteration on the given TidStore with parallel workers.
+ *
+ * The TidStoreIter struct is created in the caller's memory context, and it
+ * will be freed in TidStoreEndIterate.
+ *
+ * The caller is responsible for locking TidStore until the iteration is
+ * finished.
+ */
+TidStoreIter *
+TidStoreBeginIterateShared(TidStore *ts)
+{
+	TidStoreIter *iter;
+
+	if (!TidStoreIsShared(ts))
+		elog(ERROR, "cannot begin shared iteration on local TidStore");
+
+	iter = palloc0(sizeof(TidStoreIter));
+	iter->ts = ts;
+
+	/* begin the shared iteration on radix tree */
+	iter->tree_iter.shared =
+		(shared_ts_iter *) shared_ts_begin_iterate_shared(ts->tree.shared);
+
+	return iter;
+}
+
+/*
+ * Attach to the shared TidStore iterator. 'iter_handle' is the dsa_pointer
+ * returned by TidStoreGetSharedIterHandle(). The returned object is allocated
+ * in backend-local memory using CurrentMemoryContext.
+ */
+TidStoreIter *
+TidStoreAttachIterateShared(TidStore *ts, dsa_pointer iter_handle)
+{
+	TidStoreIter *iter;
+
+	Assert(TidStoreIsShared(ts));
+
+	iter = palloc0(sizeof(TidStoreIter));
+	iter->ts = ts;
+
+	/* Attach to the shared iterator */
+	iter->tree_iter.shared = shared_ts_attach_iterate_shared(ts->tree.shared,
+															 iter_handle);
+
+	return iter;
+}
+
 /*
  * Return the memory usage of TidStore.
  */
@@ -564,6 +615,14 @@ TidStoreGetHandle(TidStore *ts)
 	return (dsa_pointer) shared_ts_get_handle(ts->tree.shared);
 }
 
+dsa_pointer
+TidStoreGetSharedIterHandle(TidStoreIter *iter)
+{
+	Assert(TidStoreIsShared(iter->ts));
+
+	return (dsa_pointer) shared_ts_get_iter_handle(iter->tree_iter.shared);
+}
+
 /*
  * Given a TidStoreIterResult returned by TidStoreIterateNext(), extract the
  * offset numbers.  Returns the number of offsets filled in, if <=
diff --git a/src/include/access/tidstore.h b/src/include/access/tidstore.h
index d95cabd7b5..0c79a101fd 100644
--- a/src/include/access/tidstore.h
+++ b/src/include/access/tidstore.h
@@ -37,6 +37,9 @@ extern void TidStoreDetach(TidStore *ts);
 extern void TidStoreLockExclusive(TidStore *ts);
 extern void TidStoreLockShare(TidStore *ts);
 extern void TidStoreUnlock(TidStore *ts);
+extern TidStoreIter *TidStoreBeginIterateShared(TidStore *ts);
+extern TidStoreIter *TidStoreAttachIterateShared(TidStore *ts, dsa_pointer iter_handle);
+extern dsa_pointer TidStoreGetSharedIterHandle(TidStoreIter *iter);
 extern void TidStoreDestroy(TidStore *ts);
 extern void TidStoreSetBlockOffsets(TidStore *ts, BlockNumber blkno, OffsetNumber *offsets,
 									int num_offsets);
-- 
2.43.5