From 5b80d4e8dd77b0f8163e17e71d7634df5d0f340d Mon Sep 17 00:00:00 2001
From: Jason Petersen <jason@citusdata.com>
Date: Wed, 28 Sep 2016 10:49:28 -0600
Subject: [PATCH] Directly register multi-shard callbacks in PG_init

I had changed these callbacks to use the same method I chose for the
router executor (for consistency), but as that method is flawed, we now
want to ensure we directly register them from PG_init as well.
---
 src/backend/distributed/shared_library_init.c |  2 +-
 .../transaction/multi_shard_transaction.c     | 24 ++-----------------
 .../distributed/multi_shard_transaction.h     |  2 +-
 3 files changed, 4 insertions(+), 24 deletions(-)

diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c
index c52437dbb..98c88c184 100644
--- a/src/backend/distributed/shared_library_init.c
+++ b/src/backend/distributed/shared_library_init.c
@@ -155,7 +155,7 @@ _PG_init(void)
 
 	/* initialize transaction callbacks */
 	RegisterRouterExecutorXactCallbacks();
-	InstallMultiShardXactShmemHook();
+	RegisterShardPlacementXactCallbacks();
 }
 
 
diff --git a/src/backend/distributed/transaction/multi_shard_transaction.c b/src/backend/distributed/transaction/multi_shard_transaction.c
index 92b66f786..712beaa5e 100644
--- a/src/backend/distributed/transaction/multi_shard_transaction.c
+++ b/src/backend/distributed/transaction/multi_shard_transaction.c
@@ -27,11 +27,9 @@
 
 /* Global variables used in commit handler */
 static HTAB *shardConnectionHash = NULL;
-static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 static bool subXactAbortAttempted = false;
 
 /* functions needed by callbacks and hooks */
-static void RegisterShardPlacementXactCallbacks(void);
 static void CompleteShardPlacementTransactions(XactEvent event, void *arg);
 static void MultiShardSubXactCallback(SubXactEvent event, SubTransactionId subId,
 									  SubTransactionId parentSubid, void *arg);
@@ -229,32 +227,14 @@ ConnectionList(HTAB *connectionHash)
 
 
 /*
- * InstallMultiShardXactShmemHook simply installs a hook (intended to be called
- * once during backend startup), which will itself register all the transaction
- * callbacks needed by multi-shard transaction logic.
+ * RegisterShardPlacementXactCallbacks registers transaction callbacks needed
+ * for multi-shard transactions.
  */
 void
-InstallMultiShardXactShmemHook(void)
-{
-	prev_shmem_startup_hook = shmem_startup_hook;
-	shmem_startup_hook = RegisterShardPlacementXactCallbacks;
-}
-
-
-/*
- * RegisterShardPlacementXactCallbacks registers transaction callbacks needed
- * for multi-shard transactions before calling previous shmem startup hooks.
- */
-static void
 RegisterShardPlacementXactCallbacks(void)
 {
 	RegisterXactCallback(CompleteShardPlacementTransactions, NULL);
 	RegisterSubXactCallback(MultiShardSubXactCallback, NULL);
-
-	if (prev_shmem_startup_hook != NULL)
-	{
-		prev_shmem_startup_hook();
-	}
 }
 
 
diff --git a/src/include/distributed/multi_shard_transaction.h b/src/include/distributed/multi_shard_transaction.h
index 4b68378d7..4b5ec354a 100644
--- a/src/include/distributed/multi_shard_transaction.h
+++ b/src/include/distributed/multi_shard_transaction.h
@@ -33,7 +33,7 @@ extern ShardConnections * GetShardHashConnections(HTAB *connectionHash, int64 sh
 												  bool *connectionsFound);
 extern List * ConnectionList(HTAB *connectionHash);
 extern void CloseConnections(List *connectionList);
-extern void InstallMultiShardXactShmemHook(void);
+extern void RegisterShardPlacementXactCallbacks(void);
 
 
 #endif /* MULTI_SHARD_TRANSACTION_H */