From 58012054c96bb7c65e85a0c4357904fce0ccc6c4 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <github-tech@jeltef.nl>
Date: Fri, 6 Sep 2019 12:44:42 +0200
Subject: [PATCH] Add an extra advisory lock tag class

---
 src/include/distributed/resource_lock.h | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/include/distributed/resource_lock.h b/src/include/distributed/resource_lock.h
index dec6ff426..3c15702f9 100644
--- a/src/include/distributed/resource_lock.h
+++ b/src/include/distributed/resource_lock.h
@@ -35,7 +35,8 @@ typedef enum AdvisoryLocktagClass
 	/* Citus lock types */
 	ADV_LOCKTAG_CLASS_CITUS_SHARD_METADATA = 4,
 	ADV_LOCKTAG_CLASS_CITUS_SHARD = 5,
-	ADV_LOCKTAG_CLASS_CITUS_JOB = 6
+	ADV_LOCKTAG_CLASS_CITUS_JOB = 6,
+	ADV_LOCKTAG_CLASS_CITUS_REBALANCE_COLOCATION = 7
 } AdvisoryLocktagClass;
 
 
@@ -63,6 +64,16 @@ typedef enum AdvisoryLocktagClass
 						 (uint32) (jobid), \
 						 ADV_LOCKTAG_CLASS_CITUS_JOB)
 
+/* reuse advisory lock, but with different, unused field 4 (7)
+ * Also it has the the database hardcoded to MyDatabaseId, to ensure the locks
+ * are local to each database */
+#define SET_LOCKTAG_REBALANCE_COLOCATION(tag, colocationOrTableId) \
+	SET_LOCKTAG_ADVISORY(tag, \
+						 MyDatabaseId, \
+						 (uint32) ((colocationOrTableId) >> 32), \
+						 (uint32) (colocationOrTableId), \
+						 ADV_LOCKTAG_CLASS_CITUS_REBALANCE_COLOCATION)
+
 
 /* Lock shard/relation metadata for safe modifications */
 extern void LockShardDistributionMetadata(int64 shardId, LOCKMODE lockMode);