Allow on-detach callbacks for dynamic shared memory segments.

Just as backends must clean up their shared memory state (releasing
lwlocks, buffer pins, etc.) before exiting, they must also perform
any similar cleanups related to dynamic shared memory segments they
have mapped before unmapping those segments.  So add a mechanism to
ensure that.

Existing on_shmem_exit hooks include both "user level" cleanup such
as transaction abort and removal of leftover temporary relations and
also "low level" cleanup that forcibly released leftover shared
memory resources.  On-detach callbacks should run after the first
group but before the second group, so create a new before_shmem_exit
function for registering the early callbacks and keep on_shmem_exit
for the regular callbacks.  (An earlier draft of this patch added an
additional argument to on_shmem_exit, but that had a much larger
footprint and probably a substantially higher risk of breaking third
party code for no real gain.)

Patch by me, reviewed by KaiGai Kohei and Andres Freund.

src/backend/bootstrap/bootstrap.c

@@ -402,8 +402,8 @@ AuxiliaryProcessMain(int argc, char *argv[])
 		/* finish setting up bufmgr.c */
 		InitBufferPoolBackend();
 
-		/* register a shutdown callback for LWLock cleanup */
-		on_shmem_exit(ShutdownAuxiliaryProcess, 0);
+		/* register a before-shutdown callback for LWLock cleanup */
+		before_shmem_exit(ShutdownAuxiliaryProcess, 0);
 	}
 
 	/*

src/backend/catalog/namespace.c

@@ -3681,7 +3681,7 @@ AtEOXact_Namespace(bool isCommit)
 	if (myTempNamespaceSubID != InvalidSubTransactionId)
 	{
 		if (isCommit)
-			on_shmem_exit(RemoveTempRelationsCallback, 0);
+			before_shmem_exit(RemoveTempRelationsCallback, 0);
 		else
 		{
 			myTempNamespace = InvalidOid;

src/backend/commands/async.c

@@ -921,7 +921,7 @@ Exec_ListenPreCommit(void)
 	 */
 	if (!unlistenExitRegistered)
 	{
-		on_shmem_exit(Async_UnlistenOnExit, 0);
+		before_shmem_exit(Async_UnlistenOnExit, 0);
 		unlistenExitRegistered = true;
 	}
 

src/backend/storage/ipc/dsm.c

@@ -58,6 +58,14 @@
 
 #define INVALID_CONTROL_SLOT		((uint32) -1)
 
+/* Backend-local tracking for on-detach callbacks. */
+typedef struct dsm_segment_detach_callback
+{
+	on_dsm_detach_callback	function;
+	Datum					arg;
+	slist_node				node;
+} dsm_segment_detach_callback;
+
 /* Backend-local state for a dynamic shared memory segment. */
 struct dsm_segment
 {
@@ -68,6 +76,7 @@ struct dsm_segment
 	void       *impl_private;		/* Implementation-specific private data. */
 	void	   *mapped_address;		/* Mapping address, or NULL if unmapped. */
 	Size		mapped_size;		/* Size of our mapping. */
+	slist_head	on_detach;			/* On-detach callbacks. */
 };
 
 /* Shared-memory state for a dynamic shared memory segment. */
@@ -91,7 +100,6 @@ static void dsm_cleanup_for_mmap(void);
 static bool dsm_read_state_file(dsm_handle *h);
 static void dsm_write_state_file(dsm_handle h);
 static void dsm_postmaster_shutdown(int code, Datum arg);
-static void dsm_backend_shutdown(int code, Datum arg);
 static dsm_segment *dsm_create_descriptor(void);
 static bool dsm_control_segment_sane(dsm_control_header *control,
 						 Size mapped_size);
@@ -556,9 +564,6 @@ dsm_backend_startup(void)
 	}
 #endif
 
-	/* Arrange to detach segments on exit. */
-	on_shmem_exit(dsm_backend_shutdown, 0);
-
 	dsm_init_done = true;
 }
 
@@ -718,8 +723,8 @@ dsm_attach(dsm_handle h)
 /*
  * At backend shutdown time, detach any segments that are still attached.
  */
-static void
-dsm_backend_shutdown(int code, Datum arg)
+void
+dsm_backend_shutdown(void)
 {
 	while (!dlist_is_empty(&dsm_segment_list))
 	{
@@ -774,6 +779,27 @@ dsm_remap(dsm_segment *seg)
 void
 dsm_detach(dsm_segment *seg)
 {
+	/*
+	 * Invoke registered callbacks.  Just in case one of those callbacks
+	 * throws a further error that brings us back here, pop the callback
+	 * before invoking it, to avoid infinite error recursion.
+	 */
+	while (!slist_is_empty(&seg->on_detach))
+	{
+		slist_node *node;
+		dsm_segment_detach_callback *cb;
+		on_dsm_detach_callback	function;
+		Datum		arg;
+
+		node = slist_pop_head_node(&seg->on_detach);
+		cb = slist_container(dsm_segment_detach_callback, node, node);
+		function = cb->function;
+		arg = cb->arg;
+		pfree(cb);
+
+		function(seg, arg);
+	}
+
 	/*
 	 * Try to remove the mapping, if one exists.  Normally, there will be,
 	 * but maybe not, if we failed partway through a create or attach
@@ -915,6 +941,44 @@ dsm_segment_handle(dsm_segment *seg)
 	return seg->handle;
 }
 
+/*
+ * Register an on-detach callback for a dynamic shared memory segment.
+ */
+void
+on_dsm_detach(dsm_segment *seg, on_dsm_detach_callback function, Datum arg)
+{
+	dsm_segment_detach_callback *cb;
+
+	cb = MemoryContextAlloc(TopMemoryContext,
+							sizeof(dsm_segment_detach_callback));
+	cb->function = function;
+	cb->arg = arg;
+	slist_push_head(&seg->on_detach, &cb->node);
+}
+
+/*
+ * Unregister an on-detach callback for a dynamic shared memory segment.
+ */
+void
+cancel_on_dsm_detach(dsm_segment *seg, on_dsm_detach_callback function,
+					 Datum arg)
+{
+	slist_mutable_iter	iter;
+
+	slist_foreach_modify(iter, &seg->on_detach)
+	{
+		dsm_segment_detach_callback *cb;
+
+		cb = slist_container(dsm_segment_detach_callback, node, iter.cur);
+		if (cb->function == function && cb->arg == arg)
+		{
+			slist_delete_current(&iter);
+			pfree(cb);
+			break;
+		}
+	}
+}
+
 /*
  * Create a segment descriptor.
  */
@@ -937,6 +1001,8 @@ dsm_create_descriptor(void)
 	seg->resowner = CurrentResourceOwner;
 	ResourceOwnerRememberDSM(CurrentResourceOwner, seg);
 
+	slist_init(&seg->on_detach);
+
 	return seg;
 }
 

src/backend/storage/ipc/ipc.c

@@ -27,6 +27,7 @@
 #ifdef PROFILE_PID_DIR
 #include "postmaster/autovacuum.h"
 #endif
+#include "storage/dsm.h"
 #include "storage/ipc.h"
 #include "tcop/tcopprot.h"
 
@@ -64,14 +65,19 @@ static void proc_exit_prepare(int code);
 
 #define MAX_ON_EXITS 20
 
-static struct ONEXIT
+struct ONEXIT
 {
 	pg_on_exit_callback function;
 	Datum		arg;
-}	on_proc_exit_list[MAX_ON_EXITS], on_shmem_exit_list[MAX_ON_EXITS];
+};
+
+static struct ONEXIT on_proc_exit_list[MAX_ON_EXITS];
+static struct ONEXIT on_shmem_exit_list[MAX_ON_EXITS];
+static struct ONEXIT before_shmem_exit_list[MAX_ON_EXITS];
 
 static int	on_proc_exit_index,
-			on_shmem_exit_index;
+			on_shmem_exit_index,
+			before_shmem_exit_index;
 
 
 /* ----------------------------------------------------------------
@@ -202,25 +208,60 @@ proc_exit_prepare(int code)
 /* ------------------
  * Run all of the on_shmem_exit routines --- but don't actually exit.
  * This is used by the postmaster to re-initialize shared memory and
- * semaphores after a backend dies horribly.
+ * semaphores after a backend dies horribly.  As with proc_exit(), we
+ * remove each callback from the list before calling it, to avoid
+ * infinite loop in case of error.
  * ------------------
  */
 void
 shmem_exit(int code)
 {
-	elog(DEBUG3, "shmem_exit(%d): %d callbacks to make",
-		 code, on_shmem_exit_index);
+	/*
+	 * Call before_shmem_exit callbacks.
+	 *
+	 * These should be things that need most of the system to still be
+	 * up and working, such as cleanup of temp relations, which requires
+	 * catalog access; or things that need to be completed because later
+	 * cleanup steps depend on them, such as releasing lwlocks.
+	 */
+	elog(DEBUG3, "shmem_exit(%d): %d before_shmem_exit callbacks to make",
+		 code, before_shmem_exit_index);
+	while (--before_shmem_exit_index >= 0)
+		(*before_shmem_exit_list[before_shmem_exit_index].function) (code,
+					before_shmem_exit_list[before_shmem_exit_index].arg);
+	before_shmem_exit_index = 0;
 
 	/*
-	 * call all the registered callbacks.
+	 * Call dynamic shared memory callbacks.
+	 *
+	 * These serve the same purpose as late callbacks, but for dynamic shared
+	 * memory segments rather than the main shared memory segment.
+	 * dsm_backend_shutdown() has the same kind of progressive logic we use
+	 * for the main shared memory segment; namely, it unregisters each
+	 * callback before invoking it, so that we don't get stuck in an infinite
+	 * loop if one of those callbacks itself throws an ERROR or FATAL.
+	 *
+	 * Note that explicitly calling this function here is quite different
+	 * from registering it as an on_shmem_exit callback for precisely this
+	 * reason: if one dynamic shared memory callback errors out, the remaining
+	 * callbacks will still be invoked.  Thus, hard-coding this call puts it
+	 * equal footing with callbacks for the main shared memory segment.
+	 */
+	dsm_backend_shutdown();
+
+	/*
+	 * Call on_shmem_exit callbacks.
 	 *
-	 * As with proc_exit(), we remove each callback from the list before
-	 * calling it, to avoid infinite loop in case of error.
+	 * These are generally releasing low-level shared memory resources.  In
+	 * some cases, this is a backstop against the possibility that the early
+	 * callbacks might themselves fail, leading to re-entry to this routine;
+	 * in other cases, it's cleanup that only happens at process exit.
 	 */
+	elog(DEBUG3, "shmem_exit(%d): %d on_shmem_exit callbacks to make",
+		 code, on_shmem_exit_index);
 	while (--on_shmem_exit_index >= 0)
 		(*on_shmem_exit_list[on_shmem_exit_index].function) (code,
 					on_shmem_exit_list[on_shmem_exit_index].arg);
-
 	on_shmem_exit_index = 0;
 }
 
@@ -269,11 +310,40 @@ on_proc_exit(pg_on_exit_callback function, Datum arg)
 	}
 }
 
+/* ----------------------------------------------------------------
+ *		before_shmem_exit
+ *
+ *		Register early callback to perform user-level cleanup,
+ *		e.g. transaction abort, before we begin shutting down
+ *		low-level subsystems.
+ * ----------------------------------------------------------------
+ */
+void
+before_shmem_exit(pg_on_exit_callback function, Datum arg)
+{
+	if (before_shmem_exit_index >= MAX_ON_EXITS)
+		ereport(FATAL,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg_internal("out of before_shmem_exit slots")));
+
+	before_shmem_exit_list[before_shmem_exit_index].function = function;
+	before_shmem_exit_list[before_shmem_exit_index].arg = arg;
+
+	++before_shmem_exit_index;
+
+	if (!atexit_callback_setup)
+	{
+		atexit(atexit_callback);
+		atexit_callback_setup = true;
+	}
+}
+
 /* ----------------------------------------------------------------
  *		on_shmem_exit
  *
- *		this function adds a callback function to the list of
- *		functions invoked by shmem_exit().	-cim 2/6/90
+ *		Register ordinary callback to perform low-level shutdown
+ *		(e.g. releasing our PGPROC); run after before_shmem_exit
+ *		callbacks and before on_proc_exit callbacks.
  * ----------------------------------------------------------------
  */
 void
@@ -297,21 +367,22 @@ on_shmem_exit(pg_on_exit_callback function, Datum arg)
 }
 
 /* ----------------------------------------------------------------
- *		cancel_shmem_exit
+ *		cancel_before_shmem_exit
  *
- *		this function removes an entry, if present, from the list of
- *		functions to be invoked by shmem_exit().  For simplicity,
- *		only the latest entry can be removed.  (We could work harder
- *		but there is no need for current uses.)
+ *		this function removes a previously-registed before_shmem_exit
+ *		callback.  For simplicity, only the latest entry can be
+ *		removed.  (We could work harder but there is no need for
+ *		current uses.)
  * ----------------------------------------------------------------
  */
 void
-cancel_shmem_exit(pg_on_exit_callback function, Datum arg)
+cancel_before_shmem_exit(pg_on_exit_callback function, Datum arg)
 {
-	if (on_shmem_exit_index > 0 &&
-		on_shmem_exit_list[on_shmem_exit_index - 1].function == function &&
-		on_shmem_exit_list[on_shmem_exit_index - 1].arg == arg)
-		--on_shmem_exit_index;
+	if (before_shmem_exit_index > 0 &&
+		before_shmem_exit_list[before_shmem_exit_index - 1].function
+			== function &&
+		before_shmem_exit_list[before_shmem_exit_index - 1].arg == arg)
+		--before_shmem_exit_index;
 }
 
 /* ----------------------------------------------------------------
@@ -326,6 +397,7 @@ cancel_shmem_exit(pg_on_exit_callback function, Datum arg)
 void
 on_exit_reset(void)
 {
+	before_shmem_exit_index = 0;
 	on_shmem_exit_index = 0;
 	on_proc_exit_index = 0;
 }

src/backend/utils/init/postinit.c

@@ -587,15 +587,14 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 	RelationCacheInitializePhase2();
 
 	/*
-	 * Set up process-exit callback to do pre-shutdown cleanup.  This has to
-	 * be after we've initialized all the low-level modules like the buffer
-	 * manager, because during shutdown this has to run before the low-level
-	 * modules start to close down.  On the other hand, we want it in place
-	 * before we begin our first transaction --- if we fail during the
-	 * initialization transaction, as is entirely possible, we need the
-	 * AbortTransaction call to clean up.
+	 * Set up process-exit callback to do pre-shutdown cleanup.  This is the
+	 * first before_shmem_exit callback we register; thus, this will be the
+	 * last thing we do before low-level modules like the buffer manager begin
+	 * to close down.  We need to have this in place before we begin our first
+	 * transaction --- if we fail during the initialization transaction, as is
+	 * entirely possible, we need the AbortTransaction call to clean up.
 	 */
-	on_shmem_exit(ShutdownPostgres, 0);
+	before_shmem_exit(ShutdownPostgres, 0);
 
 	/* The autovacuum launcher is done here */
 	if (IsAutoVacuumLauncherProcess())

src/include/storage/dsm.h

@@ -17,8 +17,9 @@
 
 typedef struct dsm_segment dsm_segment;
 
-/* Initialization function. */
+/* Startup and shutdown functions. */
 extern void dsm_postmaster_startup(void);
+extern void dsm_backend_shutdown(void);
 
 /* Functions that create, update, or remove mappings. */
 extern dsm_segment *dsm_create(Size size);
@@ -36,4 +37,11 @@ extern void *dsm_segment_address(dsm_segment *seg);
 extern Size dsm_segment_map_length(dsm_segment *seg);
 extern dsm_handle dsm_segment_handle(dsm_segment *seg);
 
+/* Cleanup hooks. */
+typedef void (*on_dsm_detach_callback) (dsm_segment *, Datum arg);
+extern void on_dsm_detach(dsm_segment *seg,
+			  on_dsm_detach_callback function, Datum arg);
+extern void cancel_on_dsm_detach(dsm_segment *seg,
+			  on_dsm_detach_callback function, Datum arg);
+
 #endif   /* DSM_H */

src/include/storage/ipc.h

@@ -46,14 +46,14 @@ typedef void (*shmem_startup_hook_type) (void);
  */
 #define PG_ENSURE_ERROR_CLEANUP(cleanup_function, arg)	\
 	do { \
-		on_shmem_exit(cleanup_function, arg); \
+		before_shmem_exit(cleanup_function, arg); \
 		PG_TRY()
 
 #define PG_END_ENSURE_ERROR_CLEANUP(cleanup_function, arg)	\
-		cancel_shmem_exit(cleanup_function, arg); \
+		cancel_before_shmem_exit(cleanup_function, arg); \
 		PG_CATCH(); \
 		{ \
-			cancel_shmem_exit(cleanup_function, arg); \
+			cancel_before_shmem_exit(cleanup_function, arg); \
 			cleanup_function (0, arg); \
 			PG_RE_THROW(); \
 		} \
@@ -68,7 +68,8 @@ extern void proc_exit(int code) __attribute__((noreturn));
 extern void shmem_exit(int code);
 extern void on_proc_exit(pg_on_exit_callback function, Datum arg);
 extern void on_shmem_exit(pg_on_exit_callback function, Datum arg);
-extern void cancel_shmem_exit(pg_on_exit_callback function, Datum arg);
+extern void before_shmem_exit(pg_on_exit_callback function, Datum arg);
+extern void cancel_before_shmem_exit(pg_on_exit_callback function, Datum arg);
 extern void on_exit_reset(void);
 
 /* ipci.c */