Make citus_create_restore_point MX-safe by blocking 2PC commit decisions (#8352)

DESCRIPTION: Make citus_create_restore_point MX-safe by blocking 2PC
commit decisions

Problem:
--------
In coordinator-only mode, citus_create_restore_point() creates
consistent restore points by blocking distributed writes at the
coordinator level, which is safe because all distributed transactions
are coordinated through the coordinator.

However, in MX mode (multi-writer), any worker with metadata can
initiate distributed transactions. The existing implementation only
blocks writes at the coordinator, allowing metadata workers to continue
making 2PC commit decisions. This can result in an inconsistent cluster
state where restore points on different nodes represent different
transaction visibility.

Solution:
---------
Block distributed transaction commit decisions cluster-wide by acquiring
ExclusiveLock on pg_dist_transaction on all metadata nodes (coordinator
and MX workers). Additionally, on the coordinator only, lock
pg_dist_node and pg_dist_partition to prevent topology and schema
changes.

This selective locking strategy is based on the MX mode architecture:
- DDL operations (topology changes, table creation) can ONLY be executed
through the coordinator node, even in MX mode
- MX workers can only initiate distributed DML transactions
(INSERT/UPDATE/ DELETE) that use 2PC
- Therefore, locking pg_dist_transaction on remote metadata nodes is
sufficient to block all distributed writes they can perform, while the
coordinator's additional locks on pg_dist_node and pg_dist_partition
provide cluster-wide protection against DDL changes

The implementation:
-------------------
1. Opens connections to all nodes (metadata and non-metadata workers)
2. Begins coordinated transactions on all remote connections
3. Acquires ExclusiveLock on pg_dist_node, pg_dist_partition, and
pg_dist_transaction locally on the coordinator via LockRelationOid()
4. Acquires ExclusiveLock on pg_dist_transaction on all remote metadata
nodes via SQL LOCK TABLE command (executed in parallel)
5. Creates restore points on all nodes in parallel (both metadata and
non-metadata nodes need WAL restore points)
6. Closes remote connections, which releases locks via implicit ROLLBACK

Key Insight - Why No Transaction Drainage Is Needed: 
----------------------------------------------------- 
The commit decision in Citus 2PC occurs when LogTransactionRecord()
writes to pg_dist_transaction (using RowExclusiveLock for the insert),
which happens BEFORE the writer's local commit (in the PRE_COMMIT
callback).

By holding ExclusiveLock on pg_dist_transaction:
- Transactions that have already recorded their commit decision (already
inserted their row) will complete normally
- Transactions that haven't recorded their commit decision yet will
block on the ExclusiveLock (which conflicts with the RowExclusiveLock
needed for inserts), preventing them from proceeding

This creates a clean cut point for consistency without requiring us to
drain in-flight transactions. The restore point captures the exact state
of committed transactions across the cluster.

Recovery Correctness:
---------------------
The maintenance daemon's recovery logic relies on the presence of
pg_dist_transaction records to determine whether to COMMIT PREPARED or
ROLLBACK PREPARED. Our blocking ensures that:
- Prepared transactions WITH commit records will be committed on
recovery
- Prepared transactions WITHOUT commit records will be rolled back on
recovery

Since we create restore points while holding these locks, all nodes
capture the same set of commit decisions, ensuring cluster-wide
consistency.

Backward Compatibility:
-----------------------
- Return type unchanged: still returns coordinator LSN (pg_lsn)
- Coordinator-only mode: unchanged behavior
- MX mode: automatic detection and enhanced safety (transparent)
- No SQL function signature changes required

Author: codeforall

.github/workflows/build_and_test.yml

@@ -144,7 +144,7 @@ jobs:
           ${{ needs.params.outputs.pg17_version }},
           ${{ needs.params.outputs.pg18_version }}
         ]
-      make_targets: '["check-split", "check-multi", "check-multi-1", "check-multi-mx", "check-vanilla", "check-isolation", "check-operations", "check-follower-cluster", "check-add-backup-node", "check-columnar", "check-columnar-isolation", "check-enterprise", "check-enterprise-isolation", "check-enterprise-isolation-logicalrep-1", "check-enterprise-isolation-logicalrep-2", "check-enterprise-isolation-logicalrep-3"]'
+      make_targets: '["check-split", "check-multi", "check-multi-1", "check-multi-mx", "check-vanilla", "check-isolation", "check-operations", "check-follower-cluster", "check-add-backup-node", "check-columnar", "check-columnar-isolation", "check-enterprise", "check-enterprise-isolation", "check-enterprise-isolation-logicalrep-1", "check-enterprise-isolation-logicalrep-2", "check-enterprise-isolation-logicalrep-3", "check-tap"]'
       image_suffix: ${{ needs.params.outputs.image_suffix }}
       image_name: ${{ needs.params.outputs.test_image_name }}
     secrets:

.gitignore

@@ -40,6 +40,8 @@ lib*.pc
 /compile_commands.json
 /src/backend/distributed/cdc/build-cdc-*/*
 /src/test/cdc/tmp_check/*
+/src/test/tap/tmp_check/*
+/src/test/tap/log/*
 
 
 # temporary files vim creates

src/backend/distributed/operations/citus_create_restore_point.c

@@ -31,22 +31,54 @@
 
 #define CREATE_RESTORE_POINT_COMMAND "SELECT pg_catalog.pg_create_restore_point($1::text)"
 
+/*
+ * BLOCK_DISTRIBUTED_WRITES_COMMAND acquires ExclusiveLock on:
+ * 1. pg_dist_transaction - blocks 2PC commit decisions
+ * 2. pg_dist_partition - blocks DDL operations on distributed tables
+ *
+ * This ensures both DML (via 2PC) and DDL are blocked on metadata nodes.
+ */
+#define BLOCK_DISTRIBUTED_WRITES_COMMAND \
+		"LOCK TABLE pg_catalog.pg_dist_transaction IN EXCLUSIVE MODE; " \
+		"LOCK TABLE pg_catalog.pg_dist_partition IN EXCLUSIVE MODE"
 
 /* local functions forward declarations */
 static List * OpenConnectionsToAllWorkerNodes(LOCKMODE lockMode);
 static void BlockDistributedTransactions(void);
 static void CreateRemoteRestorePoints(char *restoreName, List *connectionList);
+static void BlockDistributedTransactionsOnAllMetadataNodes(List *connectionList);
 
 
 /* exports for SQL callable functions */
 PG_FUNCTION_INFO_V1(citus_create_restore_point);
 
 
 /*
- * citus_create_restore_point blocks writes to distributed tables and then
- * runs pg_create_restore_point on all nodes. This creates a consistent
- * restore point under the assumption that there are no other writers
- * than the coordinator.
+ * citus_create_restore_point creates a cluster-consistent restore point
+ * across all nodes in the Citus cluster.
+ *
+ * In coordinator-only mode, this function blocks new distributed writes
+ * at the coordinator and creates restore points on all worker nodes.
+ *
+ * In MX mode (multi-writer), this function blocks both DML and DDL
+ * operations on all metadata nodes by acquiring ExclusiveLock on:
+ *   - pg_dist_transaction: blocks 2PC commit decisions (DML)
+ *   - pg_dist_partition: blocks DDL on distributed tables
+ *
+ * This prevents new distributed transactions from recording commit decisions
+ * and blocks schema changes, ensuring all restore points represent the same
+ * consistent cluster state.
+ *
+ * The function returns the LSN of the restore point on the coordinator,
+ * maintaining backward compatibility with the original implementation.
+ *
+ * Key insight: We do NOT need to drain in-flight transactions. The commit
+ * decision in Citus 2PC happens when LogTransactionRecord() writes to
+ * pg_dist_transaction, which occurs BEFORE the writer's local commit.
+ * By blocking writes to pg_dist_transaction, we prevent commit decisions
+ * from being made. Transactions that have already recorded their commit
+ * decision will complete normally, while those that haven't will
+ * be blocked. This creates a clean cut point for consistency.
  */
 Datum
 citus_create_restore_point(PG_FUNCTION_ARGS)
@@ -88,22 +120,56 @@ citus_create_restore_point(PG_FUNCTION_ARGS)
 	 * ShareLock prevents new nodes being added, rendering connectionList incomplete
 	 */
 	List *connectionList = OpenConnectionsToAllWorkerNodes(ShareLock);
+	XLogRecPtr localRestorePoint = InvalidXLogRecPtr;
 
-	/*
-	 * Send a BEGIN to bust through pgbouncer. We won't actually commit since
-	 * that takes time. Instead we just close the connections and roll back,
-	 * which doesn't undo pg_create_restore_point.
-	 */
-	RemoteTransactionListBegin(connectionList);
+	PG_TRY();
+	{
+		/*
+		 * Send a BEGIN to bust through pgbouncer. We won't actually commit since
+		 * that takes time. Instead we just close the connections and roll back,
+		 * which doesn't undo pg_create_restore_point.
+		 */
+		RemoteTransactionListBegin(connectionList);
+
+		/* DANGER: finish as quickly as possible after this */
+		BlockDistributedTransactions();
 
-	/* DANGER: finish as quickly as possible after this */
-	BlockDistributedTransactions();
+		BlockDistributedTransactionsOnAllMetadataNodes(connectionList);
 
-	/* do local restore point first to bail out early if something goes wrong */
-	XLogRecPtr localRestorePoint = XLogRestorePoint(restoreNameString);
+		/* do local restore point first to bail out early if something goes wrong */
+		localRestorePoint = XLogRestorePoint(restoreNameString);
 
-	/* run pg_create_restore_point on all nodes */
-	CreateRemoteRestorePoints(restoreNameString, connectionList);
+		/* run pg_create_restore_point on all nodes */
+		CreateRemoteRestorePoints(restoreNameString, connectionList);
+
+		/* close connections to all nodes and
+		 * all locks gets released as part of the transaction rollback
+		 */
+		MultiConnection *conn = NULL;
+		foreach_declared_ptr(conn, connectionList)
+		{
+			ForgetResults(conn);
+			CloseConnection(conn);
+		}
+		connectionList = NIL;
+	}
+	PG_CATCH();
+	{
+		/*
+		 * On error, ensure we clean up connections and release locks.
+		 * Rolling back the metadata node transactions releases the
+		 * ExclusiveLocks on pg_dist_transaction cluster-wide.
+		 */
+		MultiConnection *conn = NULL;
+		foreach_declared_ptr(conn, connectionList)
+		{
+			ForgetResults(conn);
+			CloseConnection(conn);
+		}
+		connectionList = NIL;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
 
 	PG_RETURN_LSN(localRestorePoint);
 }
@@ -152,6 +218,89 @@ BlockDistributedTransactions(void)
 }
 
 
+/*
+ * BlockDistributedTransactionsOnAllMetadataNodes blocks distributed transactions
+ * on all metadata nodes by executing pg_lock_table remotely.
+ *
+ * This is the MX-mode equivalent of BlockDistributedTransactions(), extended
+ * to all nodes capable of initiating distributed transactions. We must hold
+ * these locks across the cluster to prevent commit decisions from being made
+ * on any node.
+ *
+ * The function expects that connections are already in a transaction block
+ * (BEGIN has been sent). The locks will be held until the transaction is
+ * rolled back or committed.
+ */
+static void
+BlockDistributedTransactionsOnAllMetadataNodes(List *connectionList)
+{
+	/*
+	 * Send LOCK TABLE commands to all metadata nodes in parallel. We use
+	 * standard SQL LOCK TABLE syntax to acquire ExclusiveLock on catalog
+	 * tables, mirroring what BlockDistributedTransactions() does on the
+	 * coordinator via LockRelationOid().
+	 *
+	 * The BLOCK_DISTRIBUTED_WRITES_COMMAND acquires:
+	 * 1. ExclusiveLock on pg_dist_transaction (blocks 2PC commit decisions)
+	 * 2. ExclusiveLock on pg_dist_partition (blocks DDL on distributed tables)
+	 *
+	 * Note: Unlike the local coordinator lock which also locks pg_dist_node,
+	 * we don't lock pg_dist_node on remote nodes because node management
+	 * operations (adding/removing nodes) are still coordinator-only.
+	 *
+	 * These locks naturally serialize concurrent restore point operations
+	 * cluster-wide, so no additional advisory lock is needed.
+	 */
+
+	/* Build list of remote metadata node connections */
+	List *metadataConnectionList = NIL;
+	MultiConnection *connection = NULL;
+	foreach_declared_ptr(connection, connectionList)
+	{
+		WorkerNode *workerNode = FindWorkerNode(connection->hostname, connection->port);
+		bool isRemoteMetadataNode = workerNode != NULL &&
+									NodeIsPrimaryAndRemote(workerNode);
+
+		if (isRemoteMetadataNode)
+		{
+			metadataConnectionList = lappend(metadataConnectionList, connection);
+		}
+	}
+
+	/* Send lock commands in parallel to all remote metadata nodes */
+	foreach_declared_ptr(connection, metadataConnectionList)
+	{
+		/*
+		 * We could use ExecuteCriticalRemoteCommand instead, but it would
+		 * not allow us to execute the commands in parallel. So for sake of
+		 * performance, we use SendRemoteCommand and send lock commands in parallel
+		 * to all metadata nodes, and later wait for all lock acquisitions to complete.
+		 */
+		int querySent = SendRemoteCommand(connection, BLOCK_DISTRIBUTED_WRITES_COMMAND);
+		if (querySent == 0)
+		{
+			ReportConnectionError(connection, ERROR);
+		}
+	}
+
+	/*
+	 * Wait for all lock acquisitions to complete. If any node fails to
+	 * acquire locks (e.g., due to a conflicting lock), this will error out.
+	 */
+	foreach_declared_ptr(connection, metadataConnectionList)
+	{
+		PGresult *result = GetRemoteCommandResult(connection, true);
+		if (!IsResponseOK(result))
+		{
+			ReportResultError(connection, result, ERROR);
+		}
+
+		PQclear(result);
+		ForgetResults(connection);
+	}
+}
+
+
 /*
  * CreateRemoteRestorePoints creates a restore point via each of the
  * connections in the list in parallel.
@@ -186,6 +335,5 @@ CreateRemoteRestorePoints(char *restoreName, List *connectionList)
 		PQclear(result);
 
 		ForgetResults(connection);
-		CloseConnection(connection);
 	}
 }

src/test/regress/Makefile

@@ -52,7 +52,7 @@ vanilla_diffs_file = $(citus_abs_srcdir)/pg_vanilla_outputs/$(MAJORVERSION)/regr
 # intermediate, for muscle memory backward compatibility.
 check: check-full check-enterprise-full
 # check-full triggers all tests that ought to be run routinely
-check-full: check-multi check-multi-mx check-multi-1 check-operations check-add-backup-node check-follower-cluster check-isolation check-failure check-split check-vanilla check-columnar check-columnar-isolation check-pg-upgrade check-arbitrary-configs check-citus-upgrade check-citus-upgrade-mixed check-citus-upgrade-local check-citus-upgrade-mixed-local check-pytest check-query-generator
+check-full: check-multi check-multi-mx check-multi-1 check-operations check-add-backup-node check-follower-cluster check-isolation check-failure check-split check-vanilla check-columnar check-columnar-isolation check-pg-upgrade check-arbitrary-configs check-citus-upgrade check-citus-upgrade-mixed check-citus-upgrade-local check-citus-upgrade-mixed-local check-pytest check-query-generator check-tap
 # check-enterprise-full triggers all enterprise specific tests
 check-enterprise-full: check-enterprise check-enterprise-isolation check-enterprise-failure check-enterprise-isolation-logicalrep-1 check-enterprise-isolation-logicalrep-2 check-enterprise-isolation-logicalrep-3
 
@@ -221,6 +221,9 @@ check-multi-mx: all
 	$(pg_regress_multi_check) --load-extension=citus \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/multi_mx_schedule $(EXTRA_TESTS)
 
+check-tap:
+	$(MAKE) -C $(citus_top_srcdir)/src/test/tap installcheck
+
 check-follower-cluster: all
 	$(pg_regress_multi_check) --load-extension=citus --follower-cluster \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/multi_follower_schedule $(EXTRA_TESTS)

src/test/regress/expected/multi_utilities.out

@@ -301,6 +301,10 @@ NOTICE:  issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;SELECT assign_
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 NOTICE:  issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;SELECT assign_distributed_transaction_id(xx, xx, 'xxxxxxx');
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
+NOTICE:  issuing LOCK TABLE pg_catalog.pg_dist_transaction IN EXCLUSIVE MODE; LOCK TABLE pg_catalog.pg_dist_partition IN EXCLUSIVE MODE
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
+NOTICE:  issuing LOCK TABLE pg_catalog.pg_dist_transaction IN EXCLUSIVE MODE; LOCK TABLE pg_catalog.pg_dist_partition IN EXCLUSIVE MODE
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 NOTICE:  issuing SELECT pg_catalog.pg_create_restore_point($1::text)
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 NOTICE:  issuing SELECT pg_catalog.pg_create_restore_point($1::text)

src/test/tap/Makefile

@@ -0,0 +1,35 @@
+subdir = src/test/tap
+citus_top_builddir = ../../..
+include $(citus_top_builddir)/Makefile.global
+
+pg_version = $(shell $(PG_CONFIG) --version 2>/dev/null)
+pg_whole_version = $(shell echo "$(pg_version)"| sed -e 's/^PostgreSQL \([0-9]*\)\(\.[0-9]*\)\{0,1\}\(.*\)/\1\2/')
+pg_major_version = $(shell echo "$(pg_whole_version)"| sed -e 's/^\([0-9]\{2\}\)\(.*\)/\1/')
+export pg_major_version
+
+test_path = t/*.pl
+
+ifeq ($(enable_tap_tests),yes)
+
+define citus_prove_installcheck
+rm -rf '$(CURDIR)'/tmp_check
+$(MKDIR_P) '$(CURDIR)'/tmp_check
+cd $(srcdir) && \
+TESTDIR='$(CURDIR)' \
+PATH="$(bindir):$$PATH" \
+PGPORT='6$(DEF_PGPORT)' \
+top_builddir='$(CURDIR)/$(top_builddir)' \
+PG_REGRESS='$(pgxsdir)/src/test/regress/pg_regress' \
+TEMP_CONFIG='$(CURDIR)'/postgresql.conf \
+$(PROVE) $(PG_PROVE_FLAGS) $(PROVE_FLAGS) $(if $(PROVE_TESTS),$(PROVE_TESTS),$(test_path))
+endef
+
+else
+citus_prove_installcheck = @echo "TAP tests not enabled when postgres was compiled"
+endif
+
+installcheck:
+	$(citus_prove_installcheck)
+
+clean distclean maintainer-clean:
+	rm -rf tmp_check

src/test/tap/citus_helpers.pm

@@ -0,0 +1,81 @@
+use strict;
+use warnings;
+
+my $pg_major_version =  int($ENV{'pg_major_version'} || 0);
+if ($pg_major_version >= 15) {
+    eval "use PostgreSQL::Test::Cluster";
+    eval "use PostgreSQL::Test::Utils";
+} else {
+    eval "use PostgresNode";
+}
+
+our $NODE_TYPE_COORDINATOR = 1;
+our $NODE_TYPE_WORKER = 2;
+
+sub create_node {
+    my ($name,$node_type,$host, $port, $config) = @_;
+    my $node;
+    if ($pg_major_version >= 15) {
+        $PostgreSQL::Test::Cluster::use_unix_sockets = 0;
+        $PostgreSQL::Test::Cluster::use_tcp = 1;
+        $PostgreSQL::Test::Cluster::test_pghost = 'localhost';
+        my %params = ( host => 'localhost' );
+        $params{port} = $port if defined $port;
+        $node = PostgreSQL::Test::Cluster->new($name, %params);
+    } else {
+        $PostgresNode::use_tcp = 1;
+        $PostgresNode::test_pghost = '127.0.0.1';
+        my %params = ( host => 'localhost' );
+        $params{port} = $port if defined $port;
+        $node = get_new_node($name, %params);
+    }
+    $port++ if defined $port;
+
+    my $citus_config_options = "
+max_connections = 100
+max_wal_senders = 100
+max_replication_slots = 100
+log_statement = 'all'
+ssl = off
+logging_collector = on
+log_directory = 'log'
+log_filename = 'postgresql.log'
+";
+
+    if ($config) {
+        $citus_config_options .= $config;
+    }
+
+    $node->init(allows_streaming => 'logical');
+    if ($node_type == $NODE_TYPE_COORDINATOR || $node_type == $NODE_TYPE_WORKER) {
+        $node->append_conf("postgresql.conf", "shared_preload_libraries = 'citus'\n" . $citus_config_options);
+        $node->start();
+        if (-f $node->data_dir . '/server.key') {
+            chmod 0600, $node->data_dir . '/server.key';
+        }
+        $node->safe_psql('postgres', "CREATE EXTENSION citus;");
+    }
+
+    return $node;
+}
+
+sub create_citus_cluster {
+    my ($num_workers,$host,$port,$citus_config) = @_;
+    my @workers = ();
+    my $node_coordinator;
+    $node_coordinator = create_node('coordinator', $NODE_TYPE_COORDINATOR,$host, $port, $citus_config);
+    my $coord_host = $node_coordinator->host();
+    my $coord_port = $node_coordinator->port();
+    $node_coordinator->safe_psql('postgres',"SELECT pg_catalog.citus_set_coordinator_host('$coord_host', $coord_port);");
+    for (my $i = 0; $i < $num_workers; $i++) {
+        $port = $port ? $port + 1 : undef;
+        my $node_worker = create_node("worker$i", $NODE_TYPE_WORKER,"localhost", $port, $citus_config);
+        my $node_worker_host = $node_worker->host();
+        my $node_worker_port = $node_worker->port();
+        $node_coordinator->safe_psql('postgres',"SELECT pg_catalog.citus_add_node('$node_worker_host', $node_worker_port);");
+        push @workers, $node_worker;
+    }
+    return $node_coordinator, @workers;
+}
+
+1;

src/test/tap/postgresql.conf

@@ -0,0 +1,7 @@
+# Minimal config for TAP tests
+shared_preload_libraries = 'citus'
+max_wal_senders = 10
+max_replication_slots = 10
+max_connections = 100
+log_statement = 'all'
+ssl = off