From 557767c4480e81be17c637d9bdef90c6f1a5c3f3 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Fri, 5 Jul 2019 12:30:24 +1200
Subject: [PATCH 2/2] Simple module for waiting for other sessions.

---
 contrib/poke/Makefile      | 18 ++++++++
 contrib/poke/poke--1.0.sql | 12 +++++
 contrib/poke/poke.c        | 93 ++++++++++++++++++++++++++++++++++++++
 contrib/poke/poke.control  |  5 ++
 4 files changed, 128 insertions(+)
 create mode 100644 contrib/poke/Makefile
 create mode 100644 contrib/poke/poke--1.0.sql
 create mode 100644 contrib/poke/poke.c
 create mode 100644 contrib/poke/poke.control

diff --git a/contrib/poke/Makefile b/contrib/poke/Makefile
new file mode 100644
index 0000000000..c56227a754
--- /dev/null
+++ b/contrib/poke/Makefile
@@ -0,0 +1,18 @@
+# contrib/poke/Makefile
+
+MODULES = poke
+
+EXTENSION = poke
+DATA = poke--1.0.sql
+PGFILEDESC = "Simple waiting mechanism."
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/poke
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/contrib/poke/poke--1.0.sql b/contrib/poke/poke--1.0.sql
new file mode 100644
index 0000000000..bb2273e517
--- /dev/null
+++ b/contrib/poke/poke--1.0.sql
@@ -0,0 +1,12 @@
+/* contrib/poke/poke--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION poke" to load this file. \quit
+
+CREATE PROCEDURE poke()
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
+
+CREATE FUNCTION wait_for_poke(milliseconds int DEFAULT -1) RETURNS boolean
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
diff --git a/contrib/poke/poke.c b/contrib/poke/poke.c
new file mode 100644
index 0000000000..f707e140dc
--- /dev/null
+++ b/contrib/poke/poke.c
@@ -0,0 +1,93 @@
+#include "postgres.h"
+
+#include "access/xact.h"
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "pgstat.h"
+#include "storage/condition_variable.h"
+#include "storage/ipc.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(poke);
+PG_FUNCTION_INFO_V1(wait_for_poke);
+extern void _PG_init(void);
+
+static ConditionVariable *my_cv;
+static bool poke_on_commit;
+static shmem_startup_hook_type prev_shmem_startup_hook;
+
+static void
+poke_xact_callback(XactEvent event, void *data)
+{
+	switch (event)
+	{
+	case XACT_EVENT_PRE_COMMIT:
+		/* ignore */
+		break;
+	case XACT_EVENT_COMMIT:
+		if (poke_on_commit)
+			ConditionVariableBroadcast(my_cv);
+		/* fall through */
+	default:
+		poke_on_commit = false;
+	}
+}
+
+static void
+poke_shmem_init(void)
+{
+	bool	found;
+
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	my_cv = ShmemInitStruct("poke", sizeof(ConditionVariable), &found);
+	if (!found)
+		ConditionVariableInit(my_cv);
+	LWLockRelease(AddinShmemInitLock);
+
+	RegisterXactCallback(poke_xact_callback, NULL);
+}
+
+static void
+poke_shmem_startup_hook(void)
+{
+	poke_shmem_init();
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+}
+
+void
+_PG_init(void)
+{
+	if (process_shared_preload_libraries_in_progress)
+	{
+		/* If preloading, then politely request the space. */
+		prev_shmem_startup_hook = shmem_startup_hook;
+		shmem_startup_hook = poke_shmem_startup_hook;
+		RequestAddinShmemSpace(sizeof(ConditionVariable));
+		return;
+	}
+
+	poke_shmem_init();
+}
+
+Datum
+poke(PG_FUNCTION_ARGS)
+{
+	poke_on_commit = true;
+
+	PG_RETURN_VOID();
+}
+
+Datum
+wait_for_poke(PG_FUNCTION_ARGS)
+{
+	int		timeout = PG_GETARG_INT32(0);
+	bool	result;
+
+	ConditionVariablePrepareToSleep(my_cv);
+	result = ConditionVariableTimedSleep(my_cv, timeout, PG_WAIT_EXTENSION);
+	ConditionVariableCancelSleep();
+
+	PG_RETURN_BOOL(result);
+}
diff --git a/contrib/poke/poke.control b/contrib/poke/poke.control
new file mode 100644
index 0000000000..f4cc50f13b
--- /dev/null
+++ b/contrib/poke/poke.control
@@ -0,0 +1,5 @@
+# poke extension
+comment = 'simple IPC between sessions'
+default_version = '1.0'
+module_pathname = '$libdir/poke'
+relocatable = true
-- 
2.21.0

