From b936bde9369ccc2f80e6b5413a60b7d1eaf269fd Mon Sep 17 00:00:00 2001
From: Andres Freund <andres@anarazel.de>
Date: Thu, 17 Aug 2017 01:35:09 -0700
Subject: [PATCH] Take AccessShareLock on the extension prior to running
 deadlock detection

---
 src/backend/distributed/utils/maintenanced.c | 48 ++++++++++++++++++--
 1 file changed, 43 insertions(+), 5 deletions(-)

diff --git a/src/backend/distributed/utils/maintenanced.c b/src/backend/distributed/utils/maintenanced.c
index 4b943c18c..e190982be 100644
--- a/src/backend/distributed/utils/maintenanced.c
+++ b/src/backend/distributed/utils/maintenanced.c
@@ -21,6 +21,8 @@
 #include "pgstat.h"
 
 #include "access/xact.h"
+#include "catalog/pg_extension.h"
+#include "commands/extension.h"
 #include "libpq/pqsignal.h"
 #include "distributed/distributed_deadlock_detection.h"
 #include "distributed/maintenanced.h"
@@ -29,6 +31,7 @@
 #include "storage/ipc.h"
 #include "storage/proc.h"
 #include "storage/latch.h"
+#include "storage/lmgr.h"
 #include "storage/lwlock.h"
 #include "tcop/tcopprot.h"
 
@@ -85,6 +88,8 @@ static void MaintenanceDaemonSigHupHandler(SIGNAL_ARGS);
 static size_t MaintenanceDaemonShmemSize(void);
 static void MaintenanceDaemonShmemInit(void);
 static void MaintenanceDaemonErrorContext(void *arg);
+static bool LockCitusExtension(void);
+
 
 /*
  * InitializeMaintenanceDaemon, called at server start, is responsible for
@@ -272,11 +277,12 @@ CitusMaintenanceDaemonMain(Datum main_arg)
 		{
 			StartTransactionCommand();
 
-			/*
-			 * We don't want to run the deadlock checks if there exists
-			 * any version mistmatch.
-			 */
-			if (CheckCitusVersion(DEBUG1))
+			if (!LockCitusExtension())
+			{
+				ereport(DEBUG1, (errmsg("could not lock the citus extension, "
+										"skipping deadlock detection")));
+			}
+			else if (CheckCitusVersion(DEBUG1) && CitusHasBeenLoaded())
 			{
 				foundDeadlock = CheckForDistributedDeadlocks();
 			}
@@ -472,3 +478,35 @@ MaintenanceDaemonErrorContext(void *arg)
 	errcontext("Citus maintenance daemon for database %u user %u",
 			   myDbData->databaseOid, myDbData->userOid);
 }
+
+
+/*
+ * LockCitusExtension acquires a lock on the Citus extension or returns
+ * false if the extension does not exist or is being dropped.
+ */
+static bool
+LockCitusExtension(void)
+{
+	Oid recheckExtensionOid = InvalidOid;
+
+	Oid extensionOid = get_extension_oid("citus", true);
+	if (extensionOid == InvalidOid)
+	{
+		/* citus extension does not exist */
+		return false;
+	}
+
+	LockDatabaseObject(ExtensionRelationId, extensionOid, 0, AccessShareLock);
+
+	/*
+	 * The extension may have been dropped and possibly recreated prior to
+	 * obtaining a lock. Check whether we still get the expected OID.
+	 */
+	recheckExtensionOid = get_extension_oid("citus", true);
+	if (recheckExtensionOid != extensionOid)
+	{
+		return false;
+	}
+
+	return true;
+}