Add jsonbRegisterIterator()

Nikita Glukhov · Nikita Glukhov · commit 64c19e473486 · 2022-07-27T01:39:15.000+03:00
diff --git a/src/backend/access/common/detoast.c b/src/backend/access/common/detoast.c
@@ -42,6 +42,7 @@ create_detoast_iterator(struct varlena *attr)
 		iter = (DetoastIterator) palloc0(sizeof(DetoastIteratorData));
 		iter->done = false;
 		iter->nrefs = 1;
+		iter->gen.free = (void (*)(void *)) free_detoast_iterator;
 
 		/* This is an externally stored datum --- initialize fetch datum iterator */
 		iter->fetch_datum_iterator = fetch_iter = create_fetch_datum_iterator(attr);
@@ -86,6 +87,7 @@ create_detoast_iterator(struct varlena *attr)
 		iter = (DetoastIterator) palloc0(sizeof(DetoastIteratorData));
 		iter->done = false;
 		iter->nrefs = 1;
+		iter->gen.free = (void (*)(void *)) free_detoast_iterator;
 
 		iter->fetch_datum_iterator = palloc0(sizeof(*iter->fetch_datum_iterator));
 		iter->fetch_datum_iterator->buf = buf = create_toast_buffer(VARSIZE_ANY(attr), true);
diff --git a/src/backend/utils/adt/jsonb_util.c b/src/backend/utils/adt/jsonb_util.c
@@ -3022,14 +3022,31 @@ jsonbFreeIterators(void)
 	if (jsonb_detoast_iterators)
 	{
 		foreach(lc, jsonb_detoast_iterators->iterators)
-			free_detoast_iterator(lfirst(lc));
+		{
+			GenericDetoastIterator iter = lfirst(lc);
+
+			iter->free(iter);
+		}
 
 		pfree(jsonb_detoast_iterators);
 		jsonb_detoast_iterators = NULL;
 	}
 #endif
 }
 
+MemoryContext
+jsonbGetIteratorContext(void)
+{
+	return jsonb_detoast_iterators ? jsonb_detoast_iterators->mcxt : NULL;
+}
+
+void
+jsonbRegisterIterator(GenericDetoastIterator iter)
+{
+	if (jsonb_detoast_iterators)
+		jsonb_detoast_iterators->iterators = lappend(jsonb_detoast_iterators->iterators, iter);
+}
+
 static void
 jsonbzFree(JsonContainer *jc)
 {
@@ -3059,7 +3076,7 @@ jsonbzInit(JsonContainerData *jc, Datum value)
 #ifdef JSONB_FREE_ITERATORS
 	if (jsonb_detoast_iterators)
 	{
-		jsonb_detoast_iterators->iterators = lappend(jsonb_detoast_iterators->iterators, iter);
+		jsonbRegisterIterator(&iter->gen);
 		MemoryContextSwitchTo(oldcxt);
 	}
 #endif
@@ -3152,7 +3169,7 @@ DatumGetJsonbPC(Datum datum, Json *tmp, bool copy)
 # ifdef JSONB_FREE_ITERATORS
 	if (jsonb_detoast_iterators)
 	{
-		jsonb_detoast_iterators->iterators = lappend(jsonb_detoast_iterators->iterators, iter);
+		jsonbRegisterIterator(&iter->gen);
 		MemoryContextSwitchTo(oldcxt);
 	}
 # endif
diff --git a/src/include/access/detoast.h b/src/include/access/detoast.h
@@ -93,8 +93,14 @@ typedef struct FetchDatumIteratorData
 
 typedef struct FetchDatumIteratorData *FetchDatumIterator;
 
+typedef struct GenericDetoastIteratorData
+{
+	void (*free)(void *iter);
+} GenericDetoastIteratorData, *GenericDetoastIterator;
+
 typedef struct DetoastIteratorData
 {
+	GenericDetoastIteratorData gen;
 	ToastBuffer 		*buf;
 	FetchDatumIterator	fetch_datum_iterator;
 	int					nrefs;
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
@@ -15,6 +15,7 @@
 #include "lib/stringinfo.h"
 #include "utils/array.h"
 #include "utils/numeric.h"
+#include "access/detoast.h"
 
 /* Tokens used when sequentially processing a jsonb value */
 typedef enum
@@ -256,5 +257,7 @@ extern bool jsonb_partial_detoast;			/* GUC */
 
 extern void jsonbInitIterators(void);
 extern void jsonbFreeIterators(void);
+extern MemoryContext jsonbGetIteratorContext(void);
+extern void jsonbRegisterIterator(GenericDetoastIterator iter);
 
 #endif							/* __JSONB_H__ */
