Fix memory leak for JSON/LIST type during Postgres table scan (#784)

Fixes #750

Author: YuweiXiao

include/pgduckdb/pg/declarations.hpp

@@ -25,6 +25,9 @@ typedef double Cardinality;
 
 typedef uintptr_t Datum;
 
+struct MemoryContextData;
+typedef MemoryContextData *MemoryContext;
+
 struct FormData_pg_attribute;
 typedef FormData_pg_attribute *Form_pg_attribute;
 

include/pgduckdb/pg/memory.hpp

@@ -0,0 +1,12 @@
+#pragma once
+
+#include "pgduckdb/pg/declarations.hpp"
+
+namespace pgduckdb::pg {
+
+MemoryContext MemoryContextCreate(MemoryContext parent, const char *name);
+MemoryContext MemoryContextSwitchTo(MemoryContext target);
+void MemoryContextReset(MemoryContext context);
+void MemoryContextDelete(MemoryContext context);
+
+} // namespace pgduckdb::pg

include/pgduckdb/scan/postgres_scan.hpp

@@ -35,6 +35,7 @@ struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState {
 	std::atomic<std::uint32_t> total_row_count;
 	std::ostringstream scan_query;
 	duckdb::shared_ptr<PostgresTableReader> table_reader_global_state;
+	MemoryContext duckdb_scan_memory_ctx;
 };
 
 // Local State

src/pg/memory.cpp

@@ -0,0 +1,32 @@
+#include "pgduckdb/pgduckdb_utils.hpp"
+
+extern "C" {
+#include "postgres.h"
+#include "utils/memutils.h"
+#include "utils/palloc.h"
+}
+
+namespace pgduckdb::pg {
+
+MemoryContext
+MemoryContextCreate(MemoryContext parent, const char *name) {
+	return PostgresFunctionGuard(::AllocSetContextCreateInternal, parent, name, ALLOCSET_DEFAULT_MINSIZE,
+	                             ALLOCSET_DEFAULT_INITSIZE, ALLOCSET_DEFAULT_MAXSIZE);
+}
+
+MemoryContext
+MemoryContextSwitchTo(MemoryContext target) {
+	return PostgresFunctionGuard(::MemoryContextSwitchTo, target);
+}
+
+void
+MemoryContextReset(MemoryContext context) {
+	PostgresFunctionGuard(::MemoryContextReset, context);
+}
+
+void
+MemoryContextDelete(MemoryContext context) {
+	PostgresFunctionGuard(::MemoryContextDelete, context);
+}
+
+} // namespace pgduckdb::pg

src/pgduckdb_types.cpp

@@ -1506,10 +1506,10 @@ AppendString(duckdb::Vector &result, Datum value, idx_t offset, bool is_bpchar)
 static void
 AppendJsonb(duckdb::Vector &result, Datum value, idx_t offset) {
 	auto jsonb = DatumGetJsonbP(value);
-	auto jsonb_str = JsonbToCString(NULL, &jsonb->root, VARSIZE(jsonb));
-	duckdb::string_t str(jsonb_str);
+	StringInfo str = PostgresFunctionGuard(makeStringInfo);
+	auto json_str = PostgresFunctionGuard(JsonbToCString, str, &jsonb->root, VARSIZE(jsonb));
 	auto data = duckdb::FlatVector::GetData<duckdb::string_t>(result);
-	data[offset] = duckdb::StringVector::AddString(result, str);
+	data[offset] = duckdb::StringVector::AddString(result, json_str, str->len);
 }
 
 static void
@@ -1820,13 +1820,13 @@ ConvertPostgresToDuckValue(Oid attr_type, Datum value, duckdb::Vector &result, i
 		int16 typlen;
 		bool typbyval;
 		char typalign;
-		get_typlenbyvalalign(elem_type, &typlen, &typbyval, &typalign);
+		PostgresFunctionGuard(get_typlenbyvalalign, elem_type, &typlen, &typbyval, &typalign);
 
 		int nelems;
 		Datum *elems;
 		bool *nulls;
 		// Deconstruct the array into Datum elements
-		deconstruct_array(array, elem_type, typlen, typbyval, typalign, &elems, &nulls, &nelems);
+		PostgresFunctionGuard(deconstruct_array, array, elem_type, typlen, typbyval, typalign, &elems, &nulls, &nelems);
 
 		if (ndims == -1) {
 			throw duckdb::InternalException("Array type has an ndims of -1, so it's actually not an array??");

src/scan/postgres_scan.cpp

@@ -4,6 +4,7 @@
 #include "pgduckdb/scan/postgres_table_reader.hpp"
 #include "pgduckdb/pgduckdb_types.hpp"
 #include "pgduckdb/pgduckdb_utils.hpp"
+#include "pgduckdb/pg/memory.hpp"
 #include "pgduckdb/pg/relations.hpp"
 
 #include "pgduckdb/pgduckdb_process_lock.hpp"
@@ -170,6 +171,7 @@ PostgresScanGlobalState::PostgresScanGlobalState(Snapshot _snapshot, Relation _r
 	ConstructTableScanQuery(input);
 	table_reader_global_state =
 	    duckdb::make_shared_ptr<PostgresTableReader>(scan_query.str().c_str(), count_tuples_only);
+	duckdb_scan_memory_ctx = pg::MemoryContextCreate(CurrentMemoryContext, "DuckdbScanContext");
 	pd_log(DEBUG2, "(DuckDB/PostgresSeqScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads());
 }
 
@@ -257,6 +259,7 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb:
 	local_state.output_vector_size = 0;
 
 	std::lock_guard<std::recursive_mutex> lock(GlobalProcessLock::GetLock());
+	MemoryContext old_context = pg::MemoryContextSwitchTo(local_state.global_state->duckdb_scan_memory_ctx);
 	for (size_t i = 0; i < STANDARD_VECTOR_SIZE; i++) {
 		TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple();
 		if (pgduckdb::TupleIsNull(slot)) {
@@ -269,6 +272,8 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb:
 		InsertTupleIntoChunk(output, local_state, slot);
 	}
 
+	pg::MemoryContextReset(local_state.global_state->duckdb_scan_memory_ctx);
+	pg::MemoryContextSwitchTo(old_context);
 	SetOutputCardinality(output, local_state);
 }