From 448bf62867b9e919e5a670b6b0fc11fccb95c244 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=80=E6=8C=83?= <yizhi.fzh@alibaba-inc.com>
Date: Thu, 6 Aug 2020 21:44:23 +0800
Subject: [PATCH v1] test_module used for concurrency case simulation

---
 contrib/test_module/Makefile      | 18 ++++++++
 contrib/test_module/test_module.c | 75 +++++++++++++++++++++++++++++++
 src/backend/utils/adt/lockfuncs.c | 18 --------
 src/include/storage/lock.h        | 18 ++++++++
 4 files changed, 111 insertions(+), 18 deletions(-)
 create mode 100644 contrib/test_module/Makefile
 create mode 100644 contrib/test_module/test_module.c

diff --git a/contrib/test_module/Makefile b/contrib/test_module/Makefile
new file mode 100644
index 0000000000..5892f3995a
--- /dev/null
+++ b/contrib/test_module/Makefile
@@ -0,0 +1,18 @@
+# contrib/test_module/Makefile
+
+MODULE_big = test_module
+OBJS = \
+	$(WIN32RES) \
+	test_module.o
+PGFILEDESC = "test_module used for concurrency case simulation"
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/test_module
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/contrib/test_module/test_module.c b/contrib/test_module/test_module.c
new file mode 100644
index 0000000000..97b6c668d0
--- /dev/null
+++ b/contrib/test_module/test_module.c
@@ -0,0 +1,75 @@
+/*-------------------------------------------------------------------------
+ *
+ * test_module.c
+ *
+ *
+ * Copyright (c) 2008-2020, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *	  contrib/test_module/test_module.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "nodes/plannodes.h"
+#include "optimizer/planner.h"
+#include "storage/lock.h"
+
+
+PG_MODULE_MAGIC;
+
+/* Lock ID for planner stage */
+
+static const int64 before_planner_lock_id = 1;
+static const int64 after_planner_lock_id = 2;
+
+
+static planner_hook_type prev_planner_hook = NULL;
+
+void _PG_init(void);
+void _PG_fini(void);
+
+
+static void
+test_module_advisory_lock_unlock(int64 lock_id)
+{
+	LOCKTAG tag;
+	SET_LOCKTAG_INT64(tag, lock_id);
+	(void) LockAcquire(&tag, ExclusiveLock, true, false);
+	(void) LockRelease(&tag, ExclusiveLock, true);
+
+}
+static PlannedStmt *
+test_module_planner(Query *parse, const char *query_string, int cursorOptions,
+					ParamListInfo boundParams)
+{
+	PlannedStmt *result;
+
+	test_module_advisory_lock_unlock(before_planner_lock_id);
+	if (prev_planner_hook)
+		result = prev_planner_hook(parse, query_string, cursorOptions,
+								   boundParams);
+	else
+		result = standard_planner(parse, query_string, cursorOptions,
+								  boundParams);
+	test_module_advisory_lock_unlock(after_planner_lock_id);
+	return result;
+}
+
+void
+_PG_init(void)
+{
+	prev_planner_hook = planner_hook;
+	planner_hook = test_module_planner;
+}
+
+
+void
+_PG_fini(void)
+{
+	planner_hook = prev_planner_hook;
+}
diff --git a/src/backend/utils/adt/lockfuncs.c b/src/backend/utils/adt/lockfuncs.c
index e992d1bbfc..34ff68d647 100644
--- a/src/backend/utils/adt/lockfuncs.c
+++ b/src/backend/utils/adt/lockfuncs.c
@@ -645,24 +645,6 @@ pg_isolation_test_session_is_blocked(PG_FUNCTION_ARGS)
 }
 
 
-/*
- * Functions for manipulating advisory locks
- *
- * We make use of the locktag fields as follows:
- *
- *	field1: MyDatabaseId ... ensures locks are local to each database
- *	field2: first of 2 int4 keys, or high-order half of an int8 key
- *	field3: second of 2 int4 keys, or low-order half of an int8 key
- *	field4: 1 if using an int8 key, 2 if using 2 int4 keys
- */
-#define SET_LOCKTAG_INT64(tag, key64) \
-	SET_LOCKTAG_ADVISORY(tag, \
-						 MyDatabaseId, \
-						 (uint32) ((key64) >> 32), \
-						 (uint32) (key64), \
-						 1)
-#define SET_LOCKTAG_INT32(tag, key1, key2) \
-	SET_LOCKTAG_ADVISORY(tag, MyDatabaseId, key1, key2, 2)
 
 /*
  * pg_advisory_lock(int8) - acquire exclusive lock on an int8 key
diff --git a/src/include/storage/lock.h b/src/include/storage/lock.h
index fdabf42721..c166f829c7 100644
--- a/src/include/storage/lock.h
+++ b/src/include/storage/lock.h
@@ -265,6 +265,24 @@ typedef struct LOCKTAG
 	 (locktag).locktag_type = LOCKTAG_ADVISORY, \
 	 (locktag).locktag_lockmethodid = USER_LOCKMETHOD)
 
+/*
+ * Functions for manipulating advisory locks
+ *
+ * We make use of the locktag fields as follows:
+ *
+ *	field1: MyDatabaseId ... ensures locks are local to each database
+ *	field2: first of 2 int4 keys, or high-order half of an int8 key
+ *	field3: second of 2 int4 keys, or low-order half of an int8 key
+ *	field4: 1 if using an int8 key, 2 if using 2 int4 keys
+ */
+#define SET_LOCKTAG_INT64(tag, key64) \
+	SET_LOCKTAG_ADVISORY(tag, \
+						 MyDatabaseId, \
+						 (uint32) ((key64) >> 32), \
+						 (uint32) (key64), \
+						 1)
+#define SET_LOCKTAG_INT32(tag, key1, key2) \
+	SET_LOCKTAG_ADVISORY(tag, MyDatabaseId, key1, key2, 2)
 
 /*
  * Per-locked-object lock information:
-- 
2.21.0