From 9b88de523dbede9888b695e660992f803c0110ee Mon Sep 17 00:00:00 2001 From: Tom Lane Date: Fri, 23 May 2025 20:36:08 -0400 Subject: [PATCH v1 1/2] Fix memory leakage in postgres_fdw's DirectModify code path. postgres_fdw tries to use PG_TRY blocks to ensure that it will eventually free the PGresult created by the remote modify command. However, it's fundamentally impossible for this scheme to work reliably when there's RETURNING data, because the query could fail in between invocations of postgres_fdw's DirectModify methods. There is at least one instance of exactly this situation in the regression tests, and the ensuing session-lifespan leak is visible under Valgrind. We can improve matters by using a memory context reset callback attached to the ExecutorState context. That ensures that the PGresult will be freed when the ExecutorState context is torn down, even if control never reaches postgresEndDirectModify. Also, by switching to this approach, we can eliminate some PG_TRY overhead since it's no longer necessary to be so cautious about errors. I added a MemoryContextUnregisterResetCallback function to mcxt.c, which we hadn't needed before. This fix could have been made without that, but the logic is simpler this way. Also I think we might be able to use the same idea to make some other postgres_fdw code less fragile, and having an unregister capability will help there. --- contrib/postgres_fdw/postgres_fdw.c | 53 ++++++++++++++--------------- src/backend/utils/mmgr/mcxt.c | 39 +++++++++++++++++++-- src/include/utils/palloc.h | 2 ++ 3 files changed, 64 insertions(+), 30 deletions(-) diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c index 331f3fc088d..bd3f3e669fb 100644 --- a/contrib/postgres_fdw/postgres_fdw.c +++ b/contrib/postgres_fdw/postgres_fdw.c @@ -240,6 +240,7 @@ 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 */ @@ -2817,7 +2818,13 @@ postgresEndDirectModify(ForeignScanState *node) return; /* Release PGresult */ - PQclear(dmstate->result); + if (dmstate->result) + { + MemoryContextUnregisterResetCallback(GetMemoryChunkContext(dmstate), + &dmstate->result_cb); + PQclear(dmstate->result); + dmstate->result = NULL; + } /* Release remote connection */ ReleaseConnection(dmstate->conn); @@ -4591,13 +4598,19 @@ execute_dml_stmt(ForeignScanState *node) /* * Get the result, and check for success. * - * We don't use a PG_TRY block here, so be careful not to throw error - * without releasing the PGresult. + * 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. */ + Assert(dmstate->result == NULL); dmstate->result = pgfdw_get_result(dmstate->conn); + dmstate->result_cb.func = (MemoryContextCallbackFunction) PQclear; + dmstate->result_cb.arg = dmstate->result; + MemoryContextRegisterResetCallback(GetMemoryChunkContext(dmstate), + &dmstate->result_cb); + if (PQresultStatus(dmstate->result) != (dmstate->has_returning ? PGRES_TUPLES_OK : PGRES_COMMAND_OK)) - pgfdw_report_error(ERROR, dmstate->result, dmstate->conn, true, + pgfdw_report_error(ERROR, dmstate->result, dmstate->conn, false, dmstate->query); /* Get the number of rows affected. */ @@ -4641,30 +4654,16 @@ get_returning_data(ForeignScanState *node) } else { - /* - * On error, be sure to release the PGresult on the way out. Callers - * do not have PG_TRY blocks to ensure this happens. - */ - PG_TRY(); - { - HeapTuple newtup; - - newtup = make_tuple_from_result_row(dmstate->result, - dmstate->next_tuple, - dmstate->rel, - dmstate->attinmeta, - dmstate->retrieved_attrs, - node, - dmstate->temp_cxt); - ExecStoreHeapTuple(newtup, slot, false); - } - PG_CATCH(); - { - PQclear(dmstate->result); - PG_RE_THROW(); - } - PG_END_TRY(); + HeapTuple newtup; + newtup = make_tuple_from_result_row(dmstate->result, + dmstate->next_tuple, + dmstate->rel, + dmstate->attinmeta, + dmstate->retrieved_attrs, + node, + dmstate->temp_cxt); + ExecStoreHeapTuple(newtup, slot, false); /* Get the updated/deleted tuple. */ if (dmstate->rel) resultSlot = slot; diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c index 15fa4d0a55e..ce01dce9861 100644 --- a/src/backend/utils/mmgr/mcxt.c +++ b/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 - * want it to not do anything anymore, adjust the state pointed to by its - * "arg" to indicate that. + * Note that callers can assume this cannot fail. */ 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. diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h index e1b42267b22..039b9cba61a 100644 --- a/src/include/utils/palloc.h +++ b/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 -- 2.43.5