From 48ab1acca2b92cd4266bda611e35120c3f70b701 Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Fri, 28 Aug 2020 22:25:38 +0900
Subject: [PATCH v35 01/10] Introduce transaction manager for foreign
 transactions.

The global transaciton manager manages the transactions initiated on
the foreign server. This commit also adds both
CommitForeignTransaction and RollbackForeignTransaction FDW APIs
supporing only one-phase commit. FDW that implements these APIs can be
managed by the global transaciton manager. So FDW is able to control
its transaction using the foreign transaction manager, not using
XactCallback.

Co-authored-by: Masahiko Sawada, Ashutosh Bapat
---
 src/backend/access/transam/Makefile  |   1 +
 src/backend/access/transam/fdwxact.c | 234 +++++++++++++++++++++++++++
 src/backend/access/transam/xact.c    |   8 +
 src/backend/foreign/foreign.c        |   4 +
 src/include/access/fdwxact.h         |  34 ++++
 src/include/foreign/fdwapi.h         |  12 ++
 6 files changed, 293 insertions(+)
 create mode 100644 src/backend/access/transam/fdwxact.c
 create mode 100644 src/include/access/fdwxact.h

diff --git a/src/backend/access/transam/Makefile b/src/backend/access/transam/Makefile
index 595e02de72..b05a88549d 100644
--- a/src/backend/access/transam/Makefile
+++ b/src/backend/access/transam/Makefile
@@ -15,6 +15,7 @@ include $(top_builddir)/src/Makefile.global
 OBJS = \
 	clog.o \
 	commit_ts.o \
+	fdwxact.o \
 	generic_xlog.o \
 	multixact.o \
 	parallel.o \
diff --git a/src/backend/access/transam/fdwxact.c b/src/backend/access/transam/fdwxact.c
new file mode 100644
index 0000000000..7da90eae13
--- /dev/null
+++ b/src/backend/access/transam/fdwxact.c
@@ -0,0 +1,234 @@
+/*-------------------------------------------------------------------------
+ *
+ * fdwxact.c
+ *		PostgreSQL global transaction manager for foreign servers.
+ *
+ * This module contains the code for managing transactions started on foreign
+ * servers.
+ *
+ * An FDW that implements both commit and rollback APIs can request to register
+ * the foreign transaction by FdwXactRegisterXact() to participate it to a
+ * group of distributed tranasction.  The registered foreign transactions are
+ * identified by user mapping OID.  On commit and rollback, the global
+ * transaction manager calls corresponding FDW API to end the foreign
+ * tranasctions.
+ *
+ * Portions Copyright (c) 2021, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *	  src/backend/access/transam/fdwxact.c
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "access/fdwxact.h"
+#include "access/xact.h"
+#include "access/xlog.h"
+#include "catalog/pg_user_mapping.h"
+#include "foreign/fdwapi.h"
+#include "foreign/foreign.h"
+#include "utils/memutils.h"
+#include "utils/syscache.h"
+
+/* Check the FdwXactParticipant is capable of two-phase commit  */
+#define ServerSupportTransactionCallback(fdw_part) \
+	(((FdwXactParticipant *)(fdw_part))->commit_foreign_xact_fn != NULL)
+
+/* Check the current transaction has at least one fdwxact participant */
+#define HasFdwXactParticipant() \
+	(FdwXactParticipants != NULL && \
+	 hash_get_num_entries(FdwXactParticipants) > 0)
+
+/*
+ * Structure to bundle the foreign transaction participant.	 This struct
+ * needs to live until the end of transaction, it's allocated in the
+ * TopTransactionContext.
+ *
+ * Participants are identified by user mapping OID, rather than pair of
+ * user OID and server OID. See README.fdwxact for the discussion.
+ */
+typedef Oid FdwXactPartKey;
+typedef struct FdwXactParticipant
+{
+	FdwXactPartKey key; /* user mapping OID, hash key (must be first) */
+
+	ForeignServer *server;
+	UserMapping *usermapping;
+
+	/* Callbacks for foreign transaction */
+	CommitForeignTransaction_function commit_foreign_xact_fn;
+	RollbackForeignTransaction_function rollback_foreign_xact_fn;
+} FdwXactParticipant;
+
+/*
+ * Foreign transactions involved in the transaction.  A member of
+ * participants must support both commit and rollback APIs.
+ */
+static HTAB *FdwXactParticipants = NULL;
+
+/* Initial size of the hash table */
+#define FDWXACT_HASH_SIZE	64
+
+static void FdwXactParticipantEndTransaction(FdwXactParticipant *fdw_part,
+											 bool is_commit, bool is_parallel_worker);
+static void RemoveFdwParticipant(FdwXactPartKey umid);
+
+/*
+ * Register the given foreign transaction identified by the given user
+ * mapping OID as a participant of the transaction.
+ */
+void
+FdwXactRegisterXact(UserMapping *usermapping)
+{
+	FdwXactParticipant	*fdw_part;
+	FdwRoutine	*routine;
+	FdwXactPartKey	key;
+	MemoryContext old_ctx;
+	bool	found;
+
+	Assert(IsTransactionState());
+
+	if (FdwXactParticipants == NULL)
+	{
+		HASHCTL	ctl;
+
+		ctl.keysize = sizeof(FdwXactPartKey);
+		ctl.entrysize = sizeof(FdwXactParticipant);
+
+		/* Assume create the hash table in TopMemoryContext */
+		FdwXactParticipants = hash_create("fdw xact participants",
+										  FDWXACT_HASH_SIZE,
+										  &ctl, HASH_ELEM | HASH_BLOBS);
+	}
+
+	key = usermapping->umid;
+	fdw_part = hash_search(FdwXactParticipants, (void *) &key, HASH_ENTER, &found);
+
+	/* Already registered */
+	if (found)
+		return;
+
+	/*
+	 * The participant information needs to live until the end of the transaction
+	 * where syscache is not available, so we save them in TopMemoryContext.
+	 */
+	old_ctx = MemoryContextSwitchTo(TopMemoryContext);
+
+	fdw_part->usermapping = GetUserMapping(usermapping->userid, usermapping->serverid);
+	fdw_part->server = GetForeignServer(usermapping->serverid);
+
+	/*
+	 * Foreign server managed by the transaction manager must implement
+	 * transaction callbacks.
+	 */
+	routine = GetFdwRoutineByServerId(usermapping->serverid);
+	if (!routine->CommitForeignTransaction)
+		ereport(ERROR,
+				(errmsg("cannot register foreign server not supporting transaction callback")));
+
+	fdw_part->commit_foreign_xact_fn = routine->CommitForeignTransaction;
+	fdw_part->rollback_foreign_xact_fn = routine->RollbackForeignTransaction;
+
+	MemoryContextSwitchTo(old_ctx);
+	pfree(routine);
+}
+
+/* Remove the foreign transaction from FdwXactParticipants */
+void
+FdwXactUnregisterXact(UserMapping *usermapping)
+{
+	Assert(IsTransactionState());
+	RemoveFdwParticipant(usermapping->umid);
+}
+
+/*
+ * Remove an FdwXactParticipant entry identified by the given user mapping id
+ * from the hash table, and free the resouce if found.
+ */
+static void
+RemoveFdwParticipant(FdwXactPartKey key)
+{
+	bool found;
+	FdwXactParticipant *fdw_part;
+
+	fdw_part = hash_search(FdwXactParticipants, (void *) &key, HASH_REMOVE,
+							&found);
+
+	if (found)
+	{
+		pfree(fdw_part->server);
+		pfree(fdw_part->usermapping);
+	}
+}
+
+/*
+ * Commit or rollback all foreign transactions.
+ */
+void
+AtEOXact_FdwXact(bool is_commit, bool is_parallel_worker)
+{
+	FdwXactParticipant *fdw_part;
+	HASH_SEQ_STATUS scan;
+
+	/* If there are no foreign servers involved, we have no business here */
+	if (!HasFdwXactParticipant())
+		return;
+
+	hash_seq_init(&scan, FdwXactParticipants);
+	while ((fdw_part = (FdwXactParticipant *) hash_seq_search(&scan)))
+	{
+		Assert(ServerSupportTransactionCallback(fdw_part));
+
+		/* Commit or rollback foreign transaction */
+		FdwXactParticipantEndTransaction(fdw_part, is_commit, is_parallel_worker);
+
+		/* Successfully finished foreign transaction, remove the entry  */
+		RemoveFdwParticipant(fdw_part->key);
+	}
+
+	Assert(!HasFdwXactParticipant());
+}
+
+/*
+ * The routine for committing or rolling back the given transaction participant.
+ */
+static void
+FdwXactParticipantEndTransaction(FdwXactParticipant *fdw_part, bool is_commit,
+								 bool is_parallel_worker)
+{
+	FdwXactInfo finfo;
+
+	Assert(ServerSupportTransactionCallback(fdw_part));
+
+	finfo.server = fdw_part->server;
+	finfo.usermapping = fdw_part->usermapping;
+	finfo.flags = FDWXACT_FLAG_ONEPHASE |
+		((is_parallel_worker) ? FDWXACT_FLAG_PARALLEL_WORKER : 0);
+
+	if (is_commit)
+	{
+		fdw_part->commit_foreign_xact_fn(&finfo);
+		elog(DEBUG1, "successfully committed the foreign transaction for user mapping %u",
+			 fdw_part->key);
+	}
+	else
+	{
+		fdw_part->rollback_foreign_xact_fn(&finfo);
+		elog(DEBUG1, "successfully rolled back the foreign transaction for user mapping %u",
+			 fdw_part->key);
+	}
+}
+
+/*
+ * This function is called at PREPARE TRANSACTION.  Since we don't support
+ * preparing foreign transactions yet, raise an error if the local transaction
+ * has any foreign transaction.
+ */
+void
+AtPrepare_FdwXact(void)
+{
+	if (HasFdwXactParticipant())
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot PREPARE a transaction that has operated on foreign tables")));
+}
diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c
index 17fbc41bbb..afe40429eb 100644
--- a/src/backend/access/transam/xact.c
+++ b/src/backend/access/transam/xact.c
@@ -21,6 +21,7 @@
 #include <unistd.h>
 
 #include "access/commit_ts.h"
+#include "access/fdwxact.h"
 #include "access/multixact.h"
 #include "access/parallel.h"
 #include "access/subtrans.h"
@@ -2125,6 +2126,9 @@ CommitTransaction(void)
 	CallXactCallbacks(is_parallel_worker ? XACT_EVENT_PARALLEL_PRE_COMMIT
 					  : XACT_EVENT_PRE_COMMIT);
 
+	/* Call foreign transaction callbacks at pre-commit phase, if any */
+	AtEOXact_FdwXact(true, is_parallel_worker);
+
 	/* If we might have parallel workers, clean them up now. */
 	if (IsInParallelMode())
 		AtEOXact_Parallel(true);
@@ -2369,6 +2373,9 @@ PrepareTransaction(void)
 	 * the transaction-abort path.
 	 */
 
+	/* Process foreign trasactions */
+	AtPrepare_FdwXact();
+
 	/* Shut down the deferred-trigger manager */
 	AfterTriggerEndXact(true);
 
@@ -2705,6 +2712,7 @@ AbortTransaction(void)
 	AtAbort_Notify();
 	AtEOXact_RelationMap(false, is_parallel_worker);
 	AtAbort_Twophase();
+	AtEOXact_FdwXact(false, is_parallel_worker);
 
 	/*
 	 * Advertise the fact that we aborted in pg_xact (assuming that we got as
diff --git a/src/backend/foreign/foreign.c b/src/backend/foreign/foreign.c
index 5564dc3a1e..f8eb4fa215 100644
--- a/src/backend/foreign/foreign.c
+++ b/src/backend/foreign/foreign.c
@@ -328,6 +328,10 @@ GetFdwRoutine(Oid fdwhandler)
 		elog(ERROR, "foreign-data wrapper handler function %u did not return an FdwRoutine struct",
 			 fdwhandler);
 
+	/* The FDW must support both or nothing */
+	Assert((routine->CommitForeignTransaction && routine->RollbackForeignTransaction) ||
+		   (!routine->CommitForeignTransaction && !routine->RollbackForeignTransaction));
+
 	return routine;
 }
 
diff --git a/src/include/access/fdwxact.h b/src/include/access/fdwxact.h
new file mode 100644
index 0000000000..0699011af5
--- /dev/null
+++ b/src/include/access/fdwxact.h
@@ -0,0 +1,34 @@
+/*
+ * fdwxact.h
+ *
+ * PostgreSQL global transaction manager
+ *
+ * Portions Copyright (c) 2021, PostgreSQL Global Development Group
+ *
+ * src/include/access/fdwxact.h
+ */
+#ifndef FDWXACT_H
+#define FDWXACT_H
+
+#include "access/xact.h"
+#include "foreign/foreign.h"
+
+/* Flag passed to FDW transaction management APIs */
+#define FDWXACT_FLAG_ONEPHASE		0x01	/* transaction can commit/rollback
+											 * without preparation */
+#define FDWXACT_FLAG_PARALLEL_WORKER	0x02	/* is parallel worker? */
+
+/* State data for foreign transaction resolution, passed to FDW callbacks */
+typedef struct FdwXactInfo
+{
+	ForeignServer	*server;
+	UserMapping		*usermapping;
+
+	int	flags;			/* OR of FDWXACT_FLAG_xx flags */
+} FdwXactInfo;
+
+/* Function declarations */
+extern void AtEOXact_FdwXact(bool is_commit, bool is_parallel_worker);
+extern void AtPrepare_FdwXact(void);
+
+#endif /* FDWXACT_H */
diff --git a/src/include/foreign/fdwapi.h b/src/include/foreign/fdwapi.h
index 248f78da45..55c040e0c3 100644
--- a/src/include/foreign/fdwapi.h
+++ b/src/include/foreign/fdwapi.h
@@ -13,6 +13,7 @@
 #define FDWAPI_H
 
 #include "access/parallel.h"
+#include "access/fdwxact.h"
 #include "nodes/execnodes.h"
 #include "nodes/pathnodes.h"
 
@@ -178,6 +179,9 @@ typedef List *(*ReparameterizeForeignPathByChild_function) (PlannerInfo *root,
 															List *fdw_private,
 															RelOptInfo *child_rel);
 
+typedef void (*CommitForeignTransaction_function) (FdwXactInfo *finfo);
+typedef void (*RollbackForeignTransaction_function) (FdwXactInfo *finfo);
+
 /*
  * FdwRoutine is the struct returned by a foreign-data wrapper's handler
  * function.  It provides pointers to the callback functions needed by the
@@ -256,6 +260,10 @@ typedef struct FdwRoutine
 
 	/* Support functions for path reparameterization. */
 	ReparameterizeForeignPathByChild_function ReparameterizeForeignPathByChild;
+
+	/* Support functions for transaction management */
+	CommitForeignTransaction_function CommitForeignTransaction;
+	RollbackForeignTransaction_function RollbackForeignTransaction;
 } FdwRoutine;
 
 
@@ -269,4 +277,8 @@ extern bool IsImportableForeignTable(const char *tablename,
 									 ImportForeignSchemaStmt *stmt);
 extern Path *GetExistingLocalJoinPath(RelOptInfo *joinrel);
 
+/* Functions in fdwxact/fdwxact.c */
+extern void FdwXactRegisterXact(UserMapping *usermapping);
+extern void FdwXactUnregisterXact(UserMapping *usermapping);
+
 #endif							/* FDWAPI_H */
-- 
2.27.0

