WIP Free jsonb resources in user functions

Author: Nikita Glukhov

src/backend/utils/adt/jsonb.c

@@ -108,11 +108,14 @@ jsonb_recv(PG_FUNCTION_ARGS)
 Datum
 jsonb_out(PG_FUNCTION_ARGS)
 {
+	jsonbInitIterators();
+
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	char	   *out;
 
 	out = JsonToCString(JsonbRoot(jb), NULL);
 
+	jsonbFreeIterators();
 	PG_FREE_IF_COPY_JSONB(jb, 0);
 
 	PG_RETURN_CSTRING(out);
@@ -126,13 +129,17 @@ jsonb_out(PG_FUNCTION_ARGS)
 Datum
 jsonb_send(PG_FUNCTION_ARGS)
 {
+	jsonbInitIterators();
+
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	StringInfoData buf;
 	StringInfo	jtext = makeStringInfo();
 	int			version = 1;
 
 	(void) JsonToCString(JsonbRoot(jb), jtext);
 
+	jsonbFreeIterators();
+
 	pq_begintypsend(&buf);
 	pq_sendint8(&buf, version);
 	pq_sendtext(&buf, jtext->data, jtext->len);

src/backend/utils/adt/jsonb_op.c

@@ -112,11 +112,18 @@ jsonb_contains(PG_FUNCTION_ARGS)
 {
 	Jsonb	   *val = PG_GETARG_JSONB_P(0);
 	Jsonb	   *tmpl = PG_GETARG_JSONB_P(1);
+	bool		res;
+
+	jsonbInitIterators();
 
 	if (JB_ROOT_IS_OBJECT(val) != JB_ROOT_IS_OBJECT(tmpl))
-		PG_RETURN_BOOL(false);
+		res = false;
+	else
+		res = JsonbDeepContains(JsonRoot(val), JsonRoot(tmpl));
 
-	PG_RETURN_BOOL(JsonbDeepContains(JsonRoot(val), JsonRoot(tmpl)));
+	jsonbFreeIterators();
+
+	PG_RETURN_BOOL(res);
 }
 
 Datum

src/backend/utils/adt/jsonb_util.c

@@ -2961,6 +2961,8 @@ jsonbzIteratorInit(JsonContainer *jc)
 	return jsonbIteratorInitExt(jc, jbc, cjb);
 }
 
+List **jsonb_detoast_iterators;
+
 static void
 #ifndef JSONB_DETOAST_ITERATOR
 jsonbzInitFromCompresedDatum(JsonContainerData *jc, CompressedDatum *cd)
@@ -2985,6 +2987,12 @@ jsonbzInitFromDetoastIterator(JsonContainerData *jc, DetoastIterator iter)
 	cjb->iter = iter;
 	cjb->offset = offsetof(JsonbDatum, root);
 
+#define JSONB_FREE_ITERATORS
+#ifdef JSONB_FREE_ITERATORS
+	if (jsonb_detoast_iterators)
+		*jsonb_detoast_iterators = lappend(*jsonb_detoast_iterators, iter);
+#endif
+
 	if (!jsonb_partial_decompression)
 		PG_DETOAST_ITERATE(iter, iter->buf->capacity);
 	else
@@ -2994,14 +3002,36 @@ jsonbzInitFromDetoastIterator(JsonContainerData *jc, DetoastIterator iter)
 #endif
 }
 
+void
+jsonbInitIterators(void)
+{
+#ifdef JSONB_FREE_ITERATORS
+	jsonb_detoast_iterators = palloc0(sizeof(*jsonb_detoast_iterators));
+#endif
+}
+
+void
+jsonbFreeIterators(void)
+{
+#ifdef JSONB_FREE_ITERATORS
+	ListCell *lc;
+
+	if (jsonb_detoast_iterators)
+		foreach(lc, *jsonb_detoast_iterators)
+			free_detoast_iterator(lfirst(lc));
+
+	jsonb_detoast_iterators = NULL;
+#endif
+}
+
 static void
 jsonbzFree(JsonContainer *jc)
 {
 	CompressedJsonb *cjb = (CompressedJsonb *) &jc->_data;
 
 #ifdef JSONB_DETOAST_ITERATOR
-	if (cjb->iter)
-		free_detoast_iterator(cjb->iter);
+//	if (cjb->iter)
+//		free_detoast_iterator(cjb->iter);
 #endif
 }
 

src/backend/utils/adt/jsonfuncs.c

@@ -822,17 +822,24 @@ jsonb_object_field(PG_FUNCTION_ARGS)
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	JsonbValue *v;
 	JsonbValue	vbuf;
+	Datum		res;
 
 	if (!JB_ROOT_IS_OBJECT(jb))
 		PG_RETURN_NULL();
 
+	jsonbInitIterators();
+
 	v = getKeyJsonValueFromContainer(&jb->root,
 									 VARDATA_ANY(key),
 									 VARSIZE_ANY_EXHDR(key),
 									 &vbuf);
 
-	if (v != NULL)
-		PG_RETURN_JSONB_VALUE_P(v);
+	res = v ? JsonValueToJsonbDatum(v) : (Datum) 0;
+
+	jsonbFreeIterators();
+
+	if (res != (Datum) 0)
+		PG_RETURN_DATUM(res);
 
 	PG_RETURN_NULL();
 }
@@ -860,17 +867,24 @@ jsonb_object_field_text(PG_FUNCTION_ARGS)
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	JsonbValue *v;
 	JsonbValue	vbuf;
+	text	   *res;
 
 	if (!JB_ROOT_IS_OBJECT(jb))
 		PG_RETURN_NULL();
 
+	jsonbInitIterators();
+
 	v = getKeyJsonValueFromContainer(&jb->root,
 									 VARDATA_ANY(key),
 									 VARSIZE_ANY_EXHDR(key),
 									 &vbuf);
 
-	if (v != NULL && v->type != jbvNull)
-		PG_RETURN_TEXT_P(JsonbValueAsText(v));
+	res = v != NULL && v->type != jbvNull ? JsonbValueAsText(v) : NULL;
+
+	jsonbFreeIterators();
+
+	if (res)
+		PG_RETURN_TEXT_P(res);
 
 	PG_RETURN_NULL();
 }
@@ -896,6 +910,7 @@ jsonb_array_element(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	int			element = PG_GETARG_INT32(1);
 	JsonbValue *v;
+	Datum		res;
 
 	if (!JB_ROOT_IS_ARRAY(jb))
 		PG_RETURN_NULL();
@@ -911,9 +926,15 @@ jsonb_array_element(PG_FUNCTION_ARGS)
 			element += nelements;
 	}
 
+	jsonbInitIterators();
+
 	v = getIthJsonbValueFromContainer(JsonbRoot(jb), element);
-	if (v != NULL)
-		PG_RETURN_JSONB_VALUE_P(v);
+	res = v != NULL ? JsonValueToJsonbDatum(v) : (Datum) 0;
+
+	jsonbFreeIterators();
+
+	if (res != (Datum) 0)
+		PG_RETURN_DATUM(res);
 
 	PG_RETURN_NULL();
 }
@@ -939,6 +960,7 @@ jsonb_array_element_text(PG_FUNCTION_ARGS)
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	int			element = PG_GETARG_INT32(1);
 	JsonbValue *v;
+	text	   *res;
 
 	if (!JB_ROOT_IS_ARRAY(jb))
 		PG_RETURN_NULL();
@@ -954,10 +976,15 @@ jsonb_array_element_text(PG_FUNCTION_ARGS)
 			element += nelements;
 	}
 
+	jsonbInitIterators();
+
 	v = getIthJsonbValueFromContainer(JsonbRoot(jb), element);
+	res = v != NULL && v->type != jbvNull ? JsonbValueAsText(v) : NULL;
+
+	jsonbFreeIterators();
 
-	if (v != NULL && v->type != jbvNull)
-		PG_RETURN_TEXT_P(JsonbValueAsText(v));
+	if (res)
+		PG_RETURN_TEXT_P(res);
 
 	PG_RETURN_NULL();
 }
@@ -1457,8 +1484,12 @@ get_jsonb_path_all(FunctionCallInfo fcinfo, bool as_text)
 
 	deconstruct_array_builtin(path, TEXTOID, &pathtext, &pathnulls, &npath);
 
+	jsonbInitIterators();
+
 	res = jsonb_get_element(jb, pathtext, npath, &isnull, as_text);
 
+	jsonbFreeIterators();
+
 	if (isnull)
 		PG_RETURN_NULL();
 	else
@@ -1886,7 +1917,13 @@ json_each(PG_FUNCTION_ARGS)
 Datum
 jsonb_each(PG_FUNCTION_ARGS)
 {
-	return each_worker_jsonb(fcinfo, "jsonb_each", false);
+	Datum		res;
+
+	jsonbInitIterators();
+	res = each_worker_jsonb(fcinfo, "jsonb_each", false);
+	jsonbFreeIterators();
+
+	PG_RETURN_DATUM(res);
 }
 
 Datum
@@ -2128,13 +2165,25 @@ each_scalar(void *state, char *token, JsonTokenType tokentype)
 Datum
 jsonb_array_elements(PG_FUNCTION_ARGS)
 {
-	return elements_worker_jsonb(fcinfo, "jsonb_array_elements", false);
+	Datum res;
+
+	jsonbInitIterators();
+	res = elements_worker_jsonb(fcinfo, "jsonb_array_elements", false);
+	jsonbFreeIterators();
+
+	PG_RETURN_DATUM(res);
 }
 
 Datum
 jsonb_array_elements_text(PG_FUNCTION_ARGS)
 {
-	return elements_worker_jsonb(fcinfo, "jsonb_array_elements_text", true);
+	Datum res;
+
+	jsonbInitIterators();
+	res = elements_worker_jsonb(fcinfo, "jsonb_array_elements_text", true);
+	jsonbFreeIterators();
+
+	PG_RETURN_DATUM(res);
 }
 
 static Datum

src/backend/utils/adt/jsonpath_exec.c

@@ -484,6 +484,8 @@ jsonb_path_query_internal(FunctionCallInfo fcinfo, bool tz)
 		funcctx = SRF_FIRSTCALL_INIT();
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
+		jsonbInitIterators();
+
 		jb = PG_GETARG_JSONB_P_COPY(0);
 		jp = PG_GETARG_JSONPATH_P_COPY(1);
 		vars = PG_GETARG_JSONB_P_COPY(2);
@@ -503,7 +505,10 @@ jsonb_path_query_internal(FunctionCallInfo fcinfo, bool tz)
 	c = list_head(found);
 
 	if (c == NULL)
+	{
+		jsonbFreeIterators();
 		SRF_RETURN_DONE(funcctx);
+	}
 
 	v = lfirst(c);
 	funcctx->user_fctx = list_delete_first(found);

src/include/access/detoast.h

@@ -174,7 +174,15 @@ detoast_iterate(DetoastIterator detoast_iter, const char *destend)
 								 destend);
 
 	if (detoast_iter->buf->limit == detoast_iter->buf->capacity)
+	{
 		detoast_iter->done = true;
+#if 0
+		if (detoast_iter->buf == fetch_iter->buf)
+			fetch_iter->buf = NULL;
+		free_fetch_datum_iterator(fetch_iter);
+		detoast_iter->fetch_datum_iterator = NULL;
+#endif
+	}
 }
 
 #endif							/* FRONTEND */

src/include/utils/jsonb.h

@@ -254,4 +254,7 @@ extern bool jsonb_sort_field_values;		/* GUC */
 extern bool jsonb_partial_decompression;	/* GUC */
 extern bool jsonb_partial_detoast;			/* GUC */
 
+extern void jsonbInitIterators(void);
+extern void jsonbFreeIterators(void);
+
 #endif							/* __JSONB_H__ */