Create infrastructure to reliably prevent leakage of PGresults.

Commit 232d8caea fixed a case where postgres_fdw could lose track
of a PGresult object, resulting in a process-lifespan memory leak.
But I have little faith that there aren't other potential PGresult
leakages, now or in future, in the backend modules that use libpq.
Therefore, this patch proposes infrastructure that makes all
PGresults returned from libpq act as though they are palloc'd
in the CurrentMemoryContext (with the option to relocate them to
another context later).  This should greatly reduce the risk of
careless leaks, and it also permits removal of a bunch of code
that attempted to prevent such leaks via PG_TRY blocks.

This patch adds infrastructure that wraps each PGresult in a
"libpqsrv_PGresult" that provides a memory context reset callback
to PQclear the PGresult.  Code using this abstraction is inherently
memory-safe to the same extent as we are accustomed to in most backend
code.  Furthermore, we add some macros that automatically redirect
calls of the libpq functions concerned with PGresults to use this
infrastructure, so that almost no source-code changes are needed to
wheel this infrastructure into place in all the backend code that
uses libpq.

Perhaps in future we could create similar infrastructure for
PGconn objects, but there seems less need for that.

This patch just creates the infrastructure and makes relevant code
use it, including reverting 232d8caea in favor of this mechanism.
A good deal of follow-on simplification is possible now that we don't
have to be so cautious about freeing PGresults, but I'll put that in
a separate patch.

Author: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: Matheus Alcantara <matheusssilv97@gmail.com>
Discussion: https://postgr.es/m/2976982.1748049023@sss.pgh.pa.us

contrib/postgres_fdw/postgres_fdw.c

@@ -240,7 +240,6 @@ typedef struct PgFdwDirectModifyState
 	PGresult   *result;			/* result for query */
 	int			num_tuples;		/* # of result tuples */
 	int			next_tuple;		/* index of next one to return */
-	MemoryContextCallback result_cb;	/* ensures result will get freed */
 	Relation	resultRel;		/* relcache entry for the target relation */
 	AttrNumber *attnoMap;		/* array of attnums of input user columns */
 	AttrNumber	ctidAttno;		/* attnum of input ctid column */
@@ -2671,17 +2670,6 @@ postgresBeginDirectModify(ForeignScanState *node, int eflags)
 	dmstate = (PgFdwDirectModifyState *) palloc0(sizeof(PgFdwDirectModifyState));
 	node->fdw_state = dmstate;
 
-	/*
-	 * We use a memory context callback to ensure that the dmstate's PGresult
-	 * (if any) will be released, even if the query fails somewhere that's
-	 * outside our control.  The callback is always armed for the duration of
-	 * the query; this relies on PQclear(NULL) being a no-op.
-	 */
-	dmstate->result_cb.func = (MemoryContextCallbackFunction) PQclear;
-	dmstate->result_cb.arg = NULL;
-	MemoryContextRegisterResetCallback(CurrentMemoryContext,
-									   &dmstate->result_cb);
-
 	/*
 	 * Identify which user to do the remote access as.  This should match what
 	 * ExecCheckPermissions() does.
@@ -2829,13 +2817,7 @@ postgresEndDirectModify(ForeignScanState *node)
 		return;
 
 	/* Release PGresult */
-	if (dmstate->result)
+	PQclear(dmstate->result);
-	{
-		PQclear(dmstate->result);
-		dmstate->result = NULL;
-		/* ... and don't forget to disable the callback */
-		dmstate->result_cb.arg = NULL;
-	}
 
 	/* Release remote connection */
 	ReleaseConnection(dmstate->conn);
@@ -4615,20 +4597,20 @@ execute_dml_stmt(ForeignScanState *node)
 
 	/*
 	 * Get the result, and check for success.
-	 *
-	 * We use a memory context callback to ensure that the PGresult will be
-	 * released, even if the query fails somewhere that's outside our control.
-	 * The callback is already registered, just need to fill in its arg.
 	 */
-	Assert(dmstate->result == NULL);
 	dmstate->result = pgfdw_get_result(dmstate->conn);
-	dmstate->result_cb.arg = dmstate->result;
-
 	if (PQresultStatus(dmstate->result) !=
 		(dmstate->has_returning ? PGRES_TUPLES_OK : PGRES_COMMAND_OK))
-		pgfdw_report_error(ERROR, dmstate->result, dmstate->conn, false,
+		pgfdw_report_error(ERROR, dmstate->result, dmstate->conn, true,
 						   dmstate->query);
 
+	/*
+	 * The result potentially needs to survive across multiple executor row
+	 * cycles, so move it to the context where the dmstate is.
+	 */
+	dmstate->result = libpqsrv_PGresultSetParent(dmstate->result,
+												 GetMemoryChunkContext(dmstate));
+
 	/* Get the number of rows affected. */
 	if (dmstate->has_returning)
 		dmstate->num_tuples = PQntuples(dmstate->result);

contrib/postgres_fdw/postgres_fdw.h

@@ -15,7 +15,7 @@
 
 #include "foreign/foreign.h"
 #include "lib/stringinfo.h"
-#include "libpq-fe.h"
+#include "libpq/libpq-be-fe.h"
 #include "nodes/execnodes.h"
 #include "nodes/pathnodes.h"
 #include "utils/relcache.h"

src/backend/utils/mmgr/mcxt.c

@@ -560,9 +560,7 @@ MemoryContextDeleteChildren(MemoryContext context)
  * the specified context, since that means it will automatically be freed
  * when no longer needed.
  *
- * There is no API for deregistering a callback once registered.  If you
+ * Note that callers can assume this cannot fail.
- * want it to not do anything anymore, adjust the state pointed to by its
- * "arg" to indicate that.
  */
 void
 MemoryContextRegisterResetCallback(MemoryContext context,
@@ -577,6 +575,41 @@ MemoryContextRegisterResetCallback(MemoryContext context,
 	context->isReset = false;
 }
 
+/*
+ * MemoryContextUnregisterResetCallback
+ *		Undo the effects of MemoryContextRegisterResetCallback.
+ *
+ * This can be used if a callback's effects are no longer required
+ * at some point before the context has been reset/deleted.  It is the
+ * caller's responsibility to pfree the callback struct (if needed).
+ *
+ * An assertion failure occurs if the callback was not registered.
+ * We could alternatively define that case as a no-op, but that seems too
+ * likely to mask programming errors such as passing the wrong context.
+ */
+void
+MemoryContextUnregisterResetCallback(MemoryContext context,
+									 MemoryContextCallback *cb)
+{
+	MemoryContextCallback *prev,
+			   *cur;
+
+	Assert(MemoryContextIsValid(context));
+
+	for (prev = NULL, cur = context->reset_cbs; cur != NULL;
+		 prev = cur, cur = cur->next)
+	{
+		if (cur != cb)
+			continue;
+		if (prev)
+			prev->next = cur->next;
+		else
+			context->reset_cbs = cur->next;
+		return;
+	}
+	Assert(false);
+}
+
 /*
  * MemoryContextCallResetCallbacks
  *		Internal function to call all registered callbacks for context.

src/include/libpq/libpq-be-fe-helpers.h

@@ -30,17 +30,7 @@
 #ifndef LIBPQ_BE_FE_HELPERS_H
 #define LIBPQ_BE_FE_HELPERS_H
 
-/*
+#include "libpq/libpq-be-fe.h"
- * Despite the name, BUILDING_DLL is set only when building code directly part
- * of the backend. Which also is where libpq isn't allowed to be
- * used. Obviously this doesn't protect against libpq-fe.h getting included
- * otherwise, but perhaps still protects against a few mistakes...
- */
-#ifdef BUILDING_DLL
-#error "libpq may not be used code directly built into the backend"
-#endif
-
-#include "libpq-fe.h"
 #include "miscadmin.h"
 #include "storage/fd.h"
 #include "storage/latch.h"
@@ -462,13 +452,21 @@ exit:	;
  * This function is intended to be set via PQsetNoticeReceiver() so that
  * NOTICE, WARNING, and similar messages from the connection are reported via
  * ereport(), instead of being printed to stderr.
+ *
+ * Because this will be called from libpq with a "real" (not wrapped)
+ * PGresult, we need to temporarily ignore libpq-be-fe.h's wrapper macros
+ * for PGresult and also PQresultErrorMessage, and put back the wrappers
+ * afterwards.  That's not pretty, but there seems no better alternative.
  */
+#undef PGresult
+#undef PQresultErrorMessage
+
 static inline void
 libpqsrv_notice_receiver(void *arg, const PGresult *res)
 {
-	char	   *message;
+	const char *message;
 	int			len;
-	char	   *prefix = (char *) arg;
+	const char *prefix = (const char *) arg;
 
 	/*
 	 * Trim the trailing newline from the message text returned from
@@ -484,4 +482,7 @@ libpqsrv_notice_receiver(void *arg, const PGresult *res)
 			errmsg_internal("%s: %.*s", prefix, len, message));
 }
 
+#define PGresult libpqsrv_PGresult
+#define PQresultErrorMessage libpqsrv_PQresultErrorMessage
+
 #endif							/* LIBPQ_BE_FE_HELPERS_H */

src/include/libpq/libpq-be-fe.h

@@ -0,0 +1,259 @@
+/*-------------------------------------------------------------------------
+ *
+ * libpq-be-fe.h
+ *	  Wrapper functions for using libpq in extensions
+ *
+ * Code built directly into the backend is not allowed to link to libpq
+ * directly. Extension code is allowed to use libpq however. One of the
+ * main risks in doing so is leaking the malloc-allocated structures
+ * returned by libpq, causing a process-lifespan memory leak.
+ *
+ * This file provides wrapper objects to help in building memory-safe code.
+ * A PGresult object wrapped this way acts much as if it were palloc'd:
+ * it will go away when the specified context is reset or deleted.
+ * We might later extend the concept to other objects such as PGconns.
+ *
+ * See also the libpq-be-fe-helpers.h file, which provides additional
+ * facilities built on top of this one.
+ *
+ * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/libpq/libpq-be-fe.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef LIBPQ_BE_FE_H
+#define LIBPQ_BE_FE_H
+
+/*
+ * Despite the name, BUILDING_DLL is set only when building code directly part
+ * of the backend. Which also is where libpq isn't allowed to be
+ * used. Obviously this doesn't protect against libpq-fe.h getting included
+ * otherwise, but perhaps still protects against a few mistakes...
+ */
+#ifdef BUILDING_DLL
+#error "libpq may not be used in code directly built into the backend"
+#endif
+
+#include "libpq-fe.h"
+
+/*
+ * Memory-context-safe wrapper object for a PGresult.
+ */
+typedef struct libpqsrv_PGresult
+{
+	PGresult   *res;			/* the wrapped PGresult */
+	MemoryContext ctx;			/* the MemoryContext it's attached to */
+	MemoryContextCallback cb;	/* the callback that implements freeing */
+} libpqsrv_PGresult;
+
+
+/*
+ * Wrap the given PGresult in a libpqsrv_PGresult object, so that it will
+ * go away automatically if the current memory context is reset or deleted.
+ *
+ * To avoid potential memory leaks, backend code must always apply this
+ * immediately to the output of any PGresult-yielding libpq function.
+ */
+static inline libpqsrv_PGresult *
+libpqsrv_PQwrap(PGresult *res)
+{
+	libpqsrv_PGresult *bres;
+	MemoryContext ctx = CurrentMemoryContext;
+
+	/* We pass through a NULL result as-is, since there's nothing to free */
+	if (res == NULL)
+		return NULL;
+	/* Attempt to allocate the wrapper ... this had better not throw error */
+	bres = (libpqsrv_PGresult *)
+		MemoryContextAllocExtended(ctx,
+								   sizeof(libpqsrv_PGresult),
+								   MCXT_ALLOC_NO_OOM);
+	/* If we failed to allocate a wrapper, free the PGresult before failing */
+	if (bres == NULL)
+	{
+		PQclear(res);
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory")));
+	}
+	/* Okay, set up the wrapper */
+	bres->res = res;
+	bres->ctx = ctx;
+	bres->cb.func = (MemoryContextCallbackFunction) PQclear;
+	bres->cb.arg = res;
+	MemoryContextRegisterResetCallback(ctx, &bres->cb);
+	return bres;
+}
+
+/*
+ * Free a wrapped PGresult, after detaching it from the memory context.
+ * Like PQclear(), allow the argument to be NULL.
+ */
+static inline void
+libpqsrv_PQclear(libpqsrv_PGresult *bres)
+{
+	if (bres)
+	{
+		MemoryContextUnregisterResetCallback(bres->ctx, &bres->cb);
+		PQclear(bres->res);
+		pfree(bres);
+	}
+}
+
+/*
+ * Move a wrapped PGresult to have a different parent context.
+ */
+static inline libpqsrv_PGresult *
+libpqsrv_PGresultSetParent(libpqsrv_PGresult *bres, MemoryContext ctx)
+{
+	libpqsrv_PGresult *newres;
+
+	/* We pass through a NULL result as-is */
+	if (bres == NULL)
+		return NULL;
+	/* Make a new wrapper in the target context, raising error on OOM */
+	newres = (libpqsrv_PGresult *)
+		MemoryContextAlloc(ctx, sizeof(libpqsrv_PGresult));
+	/* Okay, set up the new wrapper */
+	newres->res = bres->res;
+	newres->ctx = ctx;
+	newres->cb.func = (MemoryContextCallbackFunction) PQclear;
+	newres->cb.arg = bres->res;
+	MemoryContextRegisterResetCallback(ctx, &newres->cb);
+	/* Disarm and delete the old wrapper */
+	MemoryContextUnregisterResetCallback(bres->ctx, &bres->cb);
+	pfree(bres);
+	return newres;
+}
+
+/*
+ * Convenience wrapper for PQgetResult.
+ *
+ * We could supply wrappers for other PGresult-returning functions too,
+ * but at present there's no need.
+ */
+static inline libpqsrv_PGresult *
+libpqsrv_PQgetResult(PGconn *conn)
+{
+	return libpqsrv_PQwrap(PQgetResult(conn));
+}
+
+/*
+ * Accessor functions for libpqsrv_PGresult.  While it's not necessary to use
+ * these, they emulate the behavior of the underlying libpq functions when
+ * passed a NULL pointer.  This is particularly important for PQresultStatus,
+ * which is often the first check on a result.
+ */
+
+static inline ExecStatusType
+libpqsrv_PQresultStatus(const libpqsrv_PGresult *res)
+{
+	if (!res)
+		return PGRES_FATAL_ERROR;
+	return PQresultStatus(res->res);
+}
+
+static inline const char *
+libpqsrv_PQresultErrorMessage(const libpqsrv_PGresult *res)
+{
+	if (!res)
+		return "";
+	return PQresultErrorMessage(res->res);
+}
+
+static inline char *
+libpqsrv_PQresultErrorField(const libpqsrv_PGresult *res, int fieldcode)
+{
+	if (!res)
+		return NULL;
+	return PQresultErrorField(res->res, fieldcode);
+}
+
+static inline char *
+libpqsrv_PQcmdStatus(const libpqsrv_PGresult *res)
+{
+	if (!res)
+		return NULL;
+	return PQcmdStatus(res->res);
+}
+
+static inline int
+libpqsrv_PQntuples(const libpqsrv_PGresult *res)
+{
+	if (!res)
+		return 0;
+	return PQntuples(res->res);
+}
+
+static inline int
+libpqsrv_PQnfields(const libpqsrv_PGresult *res)
+{
+	if (!res)
+		return 0;
+	return PQnfields(res->res);
+}
+
+static inline char *
+libpqsrv_PQgetvalue(const libpqsrv_PGresult *res, int tup_num, int field_num)
+{
+	if (!res)
+		return NULL;
+	return PQgetvalue(res->res, tup_num, field_num);
+}
+
+static inline int
+libpqsrv_PQgetlength(const libpqsrv_PGresult *res, int tup_num, int field_num)
+{
+	if (!res)
+		return 0;
+	return PQgetlength(res->res, tup_num, field_num);
+}
+
+static inline int
+libpqsrv_PQgetisnull(const libpqsrv_PGresult *res, int tup_num, int field_num)
+{
+	if (!res)
+		return 1;				/* pretend it is null */
+	return PQgetisnull(res->res, tup_num, field_num);
+}
+
+static inline char *
+libpqsrv_PQfname(const libpqsrv_PGresult *res, int field_num)
+{
+	if (!res)
+		return NULL;
+	return PQfname(res->res, field_num);
+}
+
+static inline const char *
+libpqsrv_PQcmdTuples(const libpqsrv_PGresult *res)
+{
+	if (!res)
+		return "";
+	return PQcmdTuples(res->res);
+}
+
+/*
+ * Redefine these libpq entry point names concerned with PGresults so that
+ * they will operate on libpqsrv_PGresults instead.  This avoids needing to
+ * convert a lot of pre-existing code, and reduces the notational differences
+ * between frontend and backend libpq-using code.
+ */
+#define PGresult libpqsrv_PGresult
+#define PQclear libpqsrv_PQclear
+#define PQgetResult libpqsrv_PQgetResult
+#define PQresultStatus libpqsrv_PQresultStatus
+#define PQresultErrorMessage libpqsrv_PQresultErrorMessage
+#define PQresultErrorField libpqsrv_PQresultErrorField
+#define PQcmdStatus libpqsrv_PQcmdStatus
+#define PQntuples libpqsrv_PQntuples
+#define PQnfields libpqsrv_PQnfields
+#define PQgetvalue libpqsrv_PQgetvalue
+#define PQgetlength libpqsrv_PQgetlength
+#define PQgetisnull libpqsrv_PQgetisnull
+#define PQfname libpqsrv_PQfname
+#define PQcmdTuples libpqsrv_PQcmdTuples
+
+#endif							/* LIBPQ_BE_FE_H */

src/include/utils/palloc.h

@@ -133,6 +133,8 @@ MemoryContextSwitchTo(MemoryContext context)
 /* Registration of memory context reset/delete callbacks */
 extern void MemoryContextRegisterResetCallback(MemoryContext context,
 											   MemoryContextCallback *cb);
+extern void MemoryContextUnregisterResetCallback(MemoryContext context,
+												 MemoryContextCallback *cb);
 
 /*
  * These are like standard strdup() except the copied string is

src/tools/pgindent/typedefs.list

@@ -3757,6 +3757,7 @@ leafSegmentInfo
 leaf_item
 libpq_gettext_func
 libpq_source
+libpqsrv_PGresult
 line_t
 lineno_t
 list_sort_comparator