Re: Memory context can be its own parent and child in replication command - Mailing list pgsql-hackers

I wrote:
> ... An obvious candidate is subscription/t/100_bugs.pl, but
> the test failed there for lack of access to the test_decoding plugin.
> Perhaps we could use another plugin, but I didn't try hard.

Ah, I realized we could just use pgoutput, since the test doesn't
actually do anything that would exercise the plugin.  So here's
a v4 with the test restored.

            regards, tom lane

From cd8ff75cf45cb9515e9bf3ad1696f80901b2035f Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Sun, 20 Apr 2025 15:24:14 -0400
Subject: [PATCH v4] Use the same cmd_context throughout a walsender's
 lifetime.

exec_replication_command created a cmd_context to work in and
then deleted it on exit.  This is pretty dangerous because
some replication commands start/finish transactions.  In the
wake of commit 1afe31f03, that could lead to restoring a
CurrentMemoryContext that's already been deleted (and maybe
even reassigned, leading to hilarity such as a memory context
that is its own parent).

To fix, let's make the cmd_context persist across
exec_replication_command calls; instead of deleting it, we'll just
reset it each time.  In this way it retains the same identity and
there's no problem if transaction abort restores it as the working
context.  It probably even saves a few microseconds to do this.

This fix also ensures that exec_replication_command returns to the
caller (PostgresMain) with the same context active that had been
when it was called (probably MessageContext).  The previous
coding could get that wrong too.

Reported-by: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Author: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Reviewed-by: Tom Lane <tgl@sss.pgh.pa.us>
Discussion: https://postgr.es/m/CAO6_XqoJA7-_G6t7Uqe5nWF3nj+QBGn4F6Ptp=rUGDr0zo+KvA@mail.gmail.com
---
 src/backend/replication/walsender.c | 42 ++++++++++++++++++++++-------
 src/test/subscription/t/100_bugs.pl | 18 +++++++++++++
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c
index 987a7336ec8..9fa8beb6103 100644
--- a/src/backend/replication/walsender.c
+++ b/src/backend/replication/walsender.c
@@ -1973,8 +1973,10 @@ exec_replication_command(const char *cmd_string)
     int            parse_rc;
     Node       *cmd_node;
     const char *cmdtag;
-    MemoryContext cmd_context;
-    MemoryContext old_context;
+    MemoryContext old_context = CurrentMemoryContext;
+
+    /* We save and re-use the cmd_context across calls */
+    static MemoryContext cmd_context = NULL;

     /*
      * If WAL sender has been told that shutdown is getting close, switch its
@@ -2003,11 +2005,30 @@ exec_replication_command(const char *cmd_string)

     /*
      * Prepare to parse and execute the command.
+     *
+     * Because replication command execution can involve beginning or ending
+     * transactions, we need a working context that will survive that, so we
+     * make it a child of TopMemoryContext.  That in turn creates a hazard of
+     * long-lived memory leaks if we lose track of the working context.  We
+     * deal with that by creating it only once per walsender, and resetting it
+     * for each new command.  (Normally this reset is a no-op, but if the
+     * prior exec_replication_command call failed with an error, it won't be.)
+     *
+     * This is subtler than it looks.  The transactions we manage can extend
+     * across replication commands, indeed SnapBuildClearExportedSnapshot
+     * might have just ended one.  Because transaction exit will revert to the
+     * memory context that was current at transaction start, we need to be
+     * sure that that context is still valid.  That motivates re-using the
+     * same cmd_context rather than making a new one each time.
      */
-    cmd_context = AllocSetContextCreate(CurrentMemoryContext,
-                                        "Replication command context",
-                                        ALLOCSET_DEFAULT_SIZES);
-    old_context = MemoryContextSwitchTo(cmd_context);
+    if (cmd_context == NULL)
+        cmd_context = AllocSetContextCreate(TopMemoryContext,
+                                            "Replication command context",
+                                            ALLOCSET_DEFAULT_SIZES);
+    else
+        MemoryContextReset(cmd_context);
+
+    MemoryContextSwitchTo(cmd_context);

     replication_scanner_init(cmd_string, &scanner);

@@ -2020,7 +2041,7 @@ exec_replication_command(const char *cmd_string)
         replication_scanner_finish(scanner);

         MemoryContextSwitchTo(old_context);
-        MemoryContextDelete(cmd_context);
+        MemoryContextReset(cmd_context);

         /* XXX this is a pretty random place to make this check */
         if (MyDatabaseId == InvalidOid)
@@ -2180,9 +2201,12 @@ exec_replication_command(const char *cmd_string)
                  cmd_node->type);
     }

-    /* done */
+    /*
+     * Done.  Revert to caller's memory context, and clean out the cmd_context
+     * to recover memory right away.
+     */
     MemoryContextSwitchTo(old_context);
-    MemoryContextDelete(cmd_context);
+    MemoryContextReset(cmd_context);

     /*
      * We need not update ps display or pg_stat_activity, because PostgresMain
diff --git a/src/test/subscription/t/100_bugs.pl b/src/test/subscription/t/100_bugs.pl
index 83120f1cb6f..9c9701d58a8 100644
--- a/src/test/subscription/t/100_bugs.pl
+++ b/src/test/subscription/t/100_bugs.pl
@@ -477,6 +477,24 @@ $result =
 is( $result, qq(2|f
 3|t), 'check replicated update on subscriber');

+# Test create and immediate drop of replication slot via replication commands
+# (this exposed a memory-management bug in v18)
+my $publisher_host = $node_publisher->host;
+my $publisher_port = $node_publisher->port;
+my $connstr_db =
+  "host=$publisher_host port=$publisher_port replication=database dbname=postgres";
+
+is( $node_publisher->psql(
+        'postgres',
+        qq[
+        CREATE_REPLICATION_SLOT "test_slot" LOGICAL pgoutput ( SNAPSHOT "export");
+        DROP_REPLICATION_SLOT "test_slot";
+    ],
+        timeout => $PostgreSQL::Test::Utils::timeout_default,
+        extra_params => [ '-d', $connstr_db ]),
+    0,
+    'create and immediate drop of replication slot');
+
 $node_publisher->stop('fast');
 $node_subscriber->stop('fast');

--
2.43.5