WIP Free jsonb resources in user functions

Author: Nikita Glukhov

src/backend/utils/adt/jsonb.c

@@ -141,11 +141,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);
@@ -159,13 +162,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

@@ -149,10 +149,18 @@ json_exists_all(PG_FUNCTION_ARGS)
 static bool
 json_contains_internal(Json *val, Json *tmpl)
 {
+	bool		res;
+
+	jsonbInitIterators();
+
 	if (JB_ROOT_IS_OBJECT(val) != JB_ROOT_IS_OBJECT(tmpl))
-		return false;
+		res = false;
+	else
+		res = JsonbDeepContains(JsonRoot(val), JsonRoot(tmpl));
 
-	return JsonbDeepContains(JsonRoot(val), JsonRoot(tmpl));
+	jsonbFreeIterators();
+
+	return res;
 }
 
 Datum

src/backend/utils/adt/jsonb_util.c

@@ -3188,6 +3188,8 @@ jsonbzIteratorInit(JsonContainer *jc)
 	return jsonbIteratorInit(jc, jbc, cjb);
 }
 
+List **jsonb_detoast_iterators;
+
 static void
 #ifndef JSONB_DETOAST_ITERATOR
 jsonbzInitFromCompresedDatum(JsonContainerData *jc, CompressedDatum *cd, JsonbContainerHeader *header)
@@ -3212,6 +3214,12 @@ jsonbzInitFromDetoastIterator(JsonContainerData *jc, DetoastIterator iter, Jsonb
 	cjb->iter = iter;
 	cjb->offset = offsetof(Jsonb, 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 if (!header)
@@ -3221,14 +3229,36 @@ jsonbzInitFromDetoastIterator(JsonContainerData *jc, DetoastIterator iter, Jsonb
 #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 = jsonbzGetCompressedJsonb(jc);
 
 #ifdef JSONB_DETOAST_ITERATOR
-	if (cjb->iter)
-		free_detoast_iterator(cjb->iter);
+//	if (cjb->iter)
+//		free_detoast_iterator(cjb->iter);
 #endif
 }
 
@@ -3825,6 +3855,8 @@ jsonb_toaster_cmp(Relation rel, JsonContainer *new_jc, JsonContainer *old_jc, ch
 Datum
 jsonb_toaster(Relation rel, Datum new_val, Datum old_val, int max_size, char cmethod)
 {
+	jsonbInitIterators();
+
 	Json	   *new_js = new_val != (Datum) 0 ? DatumGetJsonbPC(new_val, NULL, false) : NULL;
 	Json	   *old_js = old_val != (Datum) 0 ? DatumGetJsonbPC(old_val, NULL, false) : NULL;
 	Datum		res = (Datum) 0;
@@ -3852,6 +3884,8 @@ jsonb_toaster(Relation rel, Datum new_val, Datum old_val, int max_size, char cme
 			jsonb_toaster_delete_recursive(rel, JsonRoot(old_js), false);
 	}
 
+	jsonbFreeIterators();
+
 	return res;
 }
 

src/backend/utils/adt/jsonfuncs.c

@@ -634,6 +634,8 @@ json_object_field_internal(Json *jb, text *key)
 Datum
 jsonb_object_field(PG_FUNCTION_ARGS)
 {
+	jsonbInitIterators();
+
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	JsonValue  *res = json_object_field_internal(jb,
 												 PG_GETARG_TEXT_PP(1));
@@ -644,11 +646,13 @@ jsonb_object_field(PG_FUNCTION_ARGS)
 
 		//if (res && res->type == jbvBinary)
 		//	JsonContainerFree(res->val.binary.data);
+		jsonbFreeIterators();
 		PG_FREE_IF_COPY_JSONB(jb, 0);
 		PG_RETURN_DATUM(r);
 	}
 	else
 	{
+		jsonbFreeIterators();
 		PG_FREE_IF_COPY_JSONB(jb, 0);
 		PG_RETURN_NULL();
 	}
@@ -669,6 +673,8 @@ json_object_field(PG_FUNCTION_ARGS)
 Datum
 jsonb_object_field_text(PG_FUNCTION_ARGS)
 {
+	jsonbInitIterators();
+
 	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
 	JsonValue  *res = json_object_field_internal(jb, PG_GETARG_TEXT_PP(1));
 
@@ -678,11 +684,13 @@ jsonb_object_field_text(PG_FUNCTION_ARGS)
 
 		//if (res && res->type == jbvBinary)
 		//	JsonContainerFree(res->val.binary.data);
+		jsonbFreeIterators();
 		PG_FREE_IF_COPY_JSONB(jb, 0);
 		PG_RETURN_TEXT_P(r);
 	}
 	else
 	{
+		jsonbFreeIterators();
 		PG_FREE_IF_COPY_JSONB(jb, 0);
 		PG_RETURN_NULL();
 	}
@@ -726,11 +734,17 @@ json_array_element_internal(Json *jb, int element)
 Datum
 jsonb_array_element(PG_FUNCTION_ARGS)
 {
-	JsonValue  *res = json_array_element_internal(PG_GETARG_JSONB_P(0),
+	jsonbInitIterators();
+
+	JsonValue  *res = json_array_element_internal(PG_GETARG_JSONB_PC(0),
 												  PG_GETARG_INT32(1));
 
+	Datum		r = res ? JsonValueGetJsonbDatum(res) : (Datum) 0;
+
+	jsonbFreeIterators();
+
 	if (res)
-		PG_RETURN_JSONB_P(JsonbValueToJsonb(res));
+		PG_RETURN_DATUM(r);
 	else
 		PG_RETURN_NULL();
 }
@@ -750,11 +764,15 @@ json_array_element(PG_FUNCTION_ARGS)
 Datum
 jsonb_array_element_text(PG_FUNCTION_ARGS)
 {
-	JsonValue  *res = json_array_element_internal(PG_GETARG_JSONB_P(0),
+	jsonbInitIterators();
+	JsonValue  *res = json_array_element_internal(PG_GETARG_JSONB_PC(0),
 												  PG_GETARG_INT32(1));
+	text	   *r = res && res->type ? JsonbValueAsText(res) : NULL;
 
-	if (res && res->type != jbvNull)
-		PG_RETURN_TEXT_P(JsonbValueAsText(res));
+	jsonbFreeIterators();
+
+	if (r)
+		PG_RETURN_TEXT_P(r);
 	else
 		PG_RETURN_NULL();
 }
@@ -1168,7 +1186,13 @@ jsonb_each(PG_FUNCTION_ARGS)
 Datum
 jsonb_each_text(PG_FUNCTION_ARGS)
 {
-	return each_worker_json(fcinfo, "jsonb_each_text", true, true);
+	Datum		res;
+
+	jsonbInitIterators();
+	res = each_worker_json(fcinfo, "jsonb_each_text", true, true);
+	jsonbFreeIterators();
+
+	return res;
 }
 
 Datum
@@ -1318,7 +1342,13 @@ jsonb_array_elements(PG_FUNCTION_ARGS)
 Datum
 jsonb_array_elements_text(PG_FUNCTION_ARGS)
 {
-	return elements_worker_json(fcinfo, "jsonb_array_elements_text", true, true);
+	Datum res;
+
+	jsonbInitIterators();
+	res = elements_worker_json(fcinfo, "jsonb_array_elements_text", true, true);
+	jsonbFreeIterators();
+
+	PG_RETURN_DATUM(res);
 }
 
 Datum
@@ -2810,7 +2840,7 @@ json_concat_internal(Json *jb1, Json *jb2, bool is_jsonb)
 
 	Assert(res != NULL);
 
-	return JsonbValueToJsonb(res);
+	return JsonValueToJson(res);
 }
 
 /*
@@ -2821,17 +2851,35 @@ json_concat_internal(Json *jb1, Json *jb2, bool is_jsonb)
 Datum
 jsonb_concat(PG_FUNCTION_ARGS)
 {
-	PG_RETURN_JSONB_P(json_concat_internal(PG_GETARG_JSONB_P(0),
-										   PG_GETARG_JSONB_P(1),
-										   true));
+	jsonbInitIterators();
+
+	Json	   *js1 = PG_GETARG_JSONB_P(0);
+	Json	   *js2 = PG_GETARG_JSONB_P(1);
+	Json	   *js = json_concat_internal(js1, js2, true);
+	Datum		res = JsonbPGetDatum(js); //JsonValueGetJsonbDatum(res);
+
+	jsonbFreeIterators();
+	PG_FREE_IF_COPY_JSONB(js1, 0);
+	PG_FREE_IF_COPY_JSONB(js2, 0);
+
+	PG_RETURN_DATUM(res);
 }
 
 Datum
 json_concat(PG_FUNCTION_ARGS)
 {
-	PG_RETURN_JSONT_P(json_concat_internal(PG_GETARG_JSONT_P(0),
-										   PG_GETARG_JSONT_P(1),
-										   false));
+	jsonbInitIterators();
+
+	Json	   *js1 = PG_GETARG_JSONT_P(0);
+	Json	   *js2 = PG_GETARG_JSONT_P(1);
+	Json	   *js = json_concat_internal(js1, js2, false);
+	Datum		res = JsontPGetDatum(js);
+
+	jsonbFreeIterators();
+	PG_FREE_IF_COPY_JSONB(js1, 0);
+	PG_FREE_IF_COPY_JSONB(js2, 0);
+
+	PG_RETURN_DATUM(res);
 }
 
 static Json *
@@ -3139,10 +3187,21 @@ json_set_internal(Json *in, ArrayType *path, Json *newjsonb, bool create)
 Datum
 jsonb_set(PG_FUNCTION_ARGS)
 {
-	PG_RETURN_JSONB_P(json_set_internal(PG_GETARG_JSONB_P(0),
+	jsonbInitIterators();
+
+	Json	   *jb1 = PG_GETARG_JSONB_P(0);
+	Json	   *jb2 = PG_GETARG_JSONB_P(2);
+	Json	   *res = json_set_internal(jb1,
 										PG_GETARG_ARRAYTYPE_P(1),
-										PG_GETARG_JSONB_P(2),
-										PG_GETARG_BOOL(3)));
+										jb2,
+										PG_GETARG_BOOL(3));
+	Datum		r = JsonbPGetDatum(res);
+
+	jsonbFreeIterators();
+	PG_FREE_IF_COPY_JSONB(jb1, 0);
+	PG_FREE_IF_COPY_JSONB(jb2, 2);
+
+	PG_RETURN_DATUM(r);
 }
 
 Datum

src/backend/utils/adt/jsonpath_exec.c

@@ -459,6 +459,8 @@ jsonb_path_query_internal(FunctionCallInfo fcinfo, bool tz, bool is_jsonb)
 		funcctx = SRF_FIRSTCALL_INIT();
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
+		jsonbInitIterators();
+
 		jb = is_jsonb ? PG_GETARG_JSONB_P_COPY(0) : PG_GETARG_JSONT_P_COPY(0);
 		jp = PG_GETARG_JSONPATH_P_COPY(1);
 		vars = is_jsonb ? PG_GETARG_JSONB_P_COPY(2) : PG_GETARG_JSONT_P_COPY(2);
@@ -477,7 +479,10 @@ jsonb_path_query_internal(FunctionCallInfo fcinfo, bool tz, bool is_jsonb)
 	c = list_head(found);
 
 	if (c == NULL)
+	{
+		jsonbFreeIterators();
 		SRF_RETURN_DONE(funcctx);
+	}
 
 	v = lfirst(c);
 	funcctx->user_fctx = list_delete_first(found);
@@ -490,7 +495,11 @@ jsonb_path_query_internal(FunctionCallInfo fcinfo, bool tz, bool is_jsonb)
 Datum
 jsonb_path_query(PG_FUNCTION_ARGS)
 {
-	return jsonb_path_query_internal(fcinfo, false, true);
+	Datum res;
+
+	res = jsonb_path_query_internal(fcinfo, false, true);
+
+	return res;
 }
 
 Datum

src/include/access/detoast.h

@@ -216,8 +216,12 @@ detoast_iterate(DetoastIterator detoast_iter, const char *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
 	}
 }
 

src/include/utils/jsonb.h

@@ -437,4 +437,7 @@ extern Datum jsonb_toaster(Relation rel, Datum new_val, Datum old_val,
 						   int max_size, char cmethod);	/* FIXME */
 extern bool JsonbHasExternal(Datum jb);
 
+extern void jsonbInitIterators(void);
+extern void jsonbFreeIterators(void);
+
 #endif							/* __JSONB_H__ */