Add detach and destroy feature to dsm_registry

Author: swchangdev

src/backend/storage/ipc/dsm_registry.c

@@ -47,6 +47,12 @@ typedef struct DSMRegistryEntry
 	size_t		size;
 } DSMRegistryEntry;
 
+typedef struct DSMDetachCallbackContext
+{
+	void (*on_detach_callback) (void *);
+	void *arg;
+} DSMDetachCallbackContext;
+
 static const dshash_parameters dsh_params = {
 	offsetof(DSMRegistryEntry, handle),
 	sizeof(DSMRegistryEntry),
@@ -121,6 +127,36 @@ init_dsm_registry(void)
 	LWLockRelease(DSMRegistryLock);
 }
 
+static void
+call_on_detach_callback(dsm_segment *seg, Datum arg)
+{
+	char *ptr = DatumGetPointer(arg);
+	DSMDetachCallbackContext *cb_ctx = (DSMDetachCallbackContext *)ptr;
+	cb_ctx->on_detach_callback(cb_ctx->arg);
+}
+
+static void
+detach_dsm_segment(dsm_handle handle, void (*detach_cb) (void *), void *arg)
+{
+	dsm_segment *seg = dsm_find_mapping(handle);
+
+	/* Detach the segment only if it is already attached */
+	if (seg != NULL)
+	{
+		if (detach_cb != NULL)
+		{
+			DSMDetachCallbackContext *cb_ctx = palloc(sizeof(DSMDetachCallbackContext));
+			cb_ctx->on_detach_callback = detach_cb;
+			cb_ctx->arg = arg;
+
+			on_dsm_detach(seg, call_on_detach_callback, PointerGetDatum(cb_ctx));
+		}
+
+		dsm_unpin_mapping(seg);
+		dsm_detach(seg);
+	}
+}
+
 /*
  * Initialize or attach a named DSM segment.
  *
@@ -172,6 +208,8 @@ GetNamedDSMSegment(const char *name, size_t size,
 	}
 	else if (entry->size != size)
 	{
+		dshash_release_lock(dsm_registry_table, entry);
+		MemoryContextSwitchTo(oldcontext);
 		ereport(ERROR,
 				(errmsg("requested DSM segment size does not match size of "
 						"existing segment")));
@@ -185,8 +223,11 @@ GetNamedDSMSegment(const char *name, size_t size,
 		{
 			seg = dsm_attach(entry->handle);
 			if (seg == NULL)
+			{
+				dshash_release_lock(dsm_registry_table, entry);
+				MemoryContextSwitchTo(oldcontext);
 				elog(ERROR, "could not map dynamic shared memory segment");
-
+			}
 			dsm_pin_mapping(seg);
 		}
 
@@ -198,3 +239,127 @@ GetNamedDSMSegment(const char *name, size_t size,
 
 	return ret;
 }
+
+/*
+ * Detach a named DSM segment
+ *
+ * This routine detaches the DSM segment. If the DSM segment was not attached
+ * by this process, then the routine just returns. on_detach_callback is passed
+ * on to dsm_segment by calling on_dsm_detach for the corresponding dsm_segment
+ * before actually detaching.
+ */
+void
+DetachNamedDSMSegment(const char *name, size_t size,
+					  void (*on_detach_callback) (void *), void *arg)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+
+	if (!name || *name == '\0')
+		ereport(ERROR,
+				(errmsg("DSM segment name cannot be empty")));
+
+	if (strlen(name) >= offsetof(DSMRegistryEntry, handle))
+		ereport(ERROR,
+				(errmsg("DSM segment name too long")));
+
+	if (size == 0)
+		ereport(ERROR,
+				(errmsg("DSM segment size must be nonzero")));
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	entry = dshash_find(dsm_registry_table, name, true);
+
+	if (entry == NULL)
+	{
+		MemoryContextSwitchTo(oldcontext);
+		ereport(ERROR,
+				(errmsg("cannot detach a DSM segment that does not exist")));
+	}
+
+	if (entry->size != size)
+	{
+		dshash_release_lock(dsm_registry_table, entry);
+		MemoryContextSwitchTo(oldcontext);
+		ereport(ERROR,
+				(errmsg("requested DSM segment size does not match size of "
+						"existing segment")));
+	}
+
+	detach_dsm_segment(entry->handle, on_detach_callback, arg);
+
+	dshash_release_lock(dsm_registry_table, entry);
+	MemoryContextSwitchTo(oldcontext);
+}
+
+/*
+ * Attempt to destroy a named DSM segment
+ *
+ * This routine attempts to destroy the DSM segment. We unpin the dsm_segment
+ * and delete the entry from dsm_registry_table. This may not destroy the
+ * dsm_segment instantly, but it would die out once all the other processes
+ * attached to this dsm_segment either exit or manually detach from the
+ * dsm_segment.
+ *
+ * Because we deleted the key from dsm_registry_table, calling
+ * GetNamedDSMSegment with the same key would result into creating a new
+ * dsm_segment instead of retrieving the old unpinned dsm_segment.
+ */
+void
+DestroyNamedDSMSegment(const char *name, size_t size,
+					   void (*on_detach_callback) (void *), void *arg)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+
+	if (!name || *name == '\0')
+		ereport(ERROR,
+				(errmsg("DSM segment name cannot be empty")));
+
+	if (strlen(name) >= offsetof(DSMRegistryEntry, handle))
+		ereport(ERROR,
+				(errmsg("DSM segment name too long")));
+
+	if (size == 0)
+		ereport(ERROR,
+				(errmsg("DSM segment size must be nonzero")));
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	entry = dshash_find(dsm_registry_table, name, true);
+
+	if (entry == NULL)
+	{
+		MemoryContextSwitchTo(oldcontext);
+		ereport(ERROR,
+				(errmsg("cannot destroy a DSM segment that does not exist")));
+	}
+
+	if (entry->size != size)
+	{
+		dshash_release_lock(dsm_registry_table, entry);
+		MemoryContextSwitchTo(oldcontext);
+		ereport(ERROR,
+				(errmsg("requested DSM segment size does not match size of "
+						"existing segment")));
+	}
+
+	detach_dsm_segment(entry->handle, on_detach_callback, arg);
+	dsm_unpin_segment(entry->handle);
+
+	/* dshash_delete_entry calls LWLockRelease internally. We shouldn't
+	 * release lock twice */
+	dshash_delete_entry(dsm_registry_table, entry);
+	dshash_delete_key(dsm_registry_table, name);
+
+	MemoryContextSwitchTo(oldcontext);
+}

src/include/storage/dsm_registry.h

@@ -17,6 +17,14 @@ extern void *GetNamedDSMSegment(const char *name, size_t size,
 								void (*init_callback) (void *ptr),
 								bool *found);
 
+extern void DetachNamedDSMSegment(const char *name, size_t size,
+								  void (*on_detach_callback) (void *),
+								  void *arg);
+
+extern void DestroyNamedDSMSegment(const char *name, size_t size,
+								   void (*on_detach_callback) (void *),
+								   void *arg);
+
 extern Size DSMRegistryShmemSize(void);
 extern void DSMRegistryShmemInit(void);
 

src/test/modules/test_dsm_registry/expected/test_dsm_registry.out

@@ -12,3 +12,15 @@ SELECT get_val_in_shmem();
              1236
 (1 row)
 
+SELECT detach_from_tdr_segment();
+ detach_from_tdr_segment 
+-------------------------
+ t
+(1 row)
+
+SELECT destroy_tdr_segment();
+ destroy_tdr_segment 
+---------------------
+ 
+(1 row)
+

src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql

@@ -2,3 +2,5 @@ CREATE EXTENSION test_dsm_registry;
 SELECT set_val_in_shmem(1236);
 \c
 SELECT get_val_in_shmem();
+SELECT detach_from_tdr_segment();
+SELECT destroy_tdr_segment();

src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql

@@ -8,3 +8,9 @@ CREATE FUNCTION set_val_in_shmem(val INT) RETURNS VOID
 
 CREATE FUNCTION get_val_in_shmem() RETURNS INT
 	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION detach_from_tdr_segment() RETURNS BOOL
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION destroy_tdr_segment() RETURNS VOID
+	AS 'MODULE_PATHNAME' LANGUAGE C;

src/test/modules/test_dsm_registry/test_dsm_registry.c

@@ -35,6 +35,12 @@ tdr_init_shmem(void *ptr)
 	state->val = 0;
 }
 
+static void
+reset_tdr_state_to_null(void *arg)
+{
+	tdr_state = NULL;
+}
+
 static void
 tdr_attach_shmem(void)
 {
@@ -74,3 +80,25 @@ get_val_in_shmem(PG_FUNCTION_ARGS)
 
 	PG_RETURN_UINT32(ret);
 }
+
+PG_FUNCTION_INFO_V1(detach_from_tdr_segment);
+Datum
+detach_from_tdr_segment(PG_FUNCTION_ARGS)
+{
+	DetachNamedDSMSegment("test_dsm_registry",
+						  sizeof(TestDSMRegistryStruct),
+						  reset_tdr_state_to_null, NULL);
+
+	PG_RETURN_BOOL(tdr_state == NULL);
+}
+
+PG_FUNCTION_INFO_V1(destroy_tdr_segment);
+Datum
+destroy_tdr_segment(PG_FUNCTION_ARGS)
+{
+	DestroyNamedDSMSegment("test_dsm_registry",
+						   sizeof(TestDSMRegistryStruct),
+						   NULL, NULL);
+
+	PG_RETURN_VOID();
+}