Remove our own shared memory tooling

Since it just has one user now let's remove our own shared memory code
and remove a lot of generic code that we do not need.

Author: jeltz

contrib/pg_tde/src/catalog/tde_principal_key.c

@@ -60,7 +60,6 @@ typedef struct TdePrincipalKeySharedState
 	LWLockPadded *Locks;
 	dshash_table_handle hashHandle;
 	void	   *rawDsaArea;		/* DSA area pointer */
-
 } TdePrincipalKeySharedState;
 
 typedef struct TdePrincipalKeylocalState
@@ -84,10 +83,6 @@ static dshash_parameters principal_key_dsh_params = {
 static TdePrincipalKeylocalState principalKeyLocalState;
 
 static void principal_key_info_attach_shmem(void);
-static Size initialize_shared_state(void *start_address);
-static void initialize_objects_in_dsa_area(dsa_area *dsa, void *raw_dsa_area);
-static Size required_shared_mem_size(void);
-static void shared_memory_shutdown(int code, Datum arg);
 static void clear_principal_key_cache(Oid databaseId);
 static inline dshash_table *get_principal_key_hash(void);
 static TDEPrincipalKey *get_principal_key_from_cache(Oid dbOid);
@@ -111,32 +106,6 @@ PG_FUNCTION_INFO_V1(pg_tde_set_server_key_using_global_key_provider);
 
 static void pg_tde_set_principal_key_internal(Oid providerOid, Oid dbOid, const char *principal_key_name, const char *provider_name, bool ensure_new_key);
 
-static const TDEShmemSetupRoutine principal_key_info_shmem_routine = {
-	.init_shared_state = initialize_shared_state,
-	.init_dsa_area_objects = initialize_objects_in_dsa_area,
-	.required_shared_mem_size = required_shared_mem_size,
-	.shmem_kill = shared_memory_shutdown
-};
-
-void
-InitializePrincipalKeyInfo(void)
-{
-	ereport(LOG, errmsg("Initializing TDE principal key info"));
-	RegisterShmemRequest(&principal_key_info_shmem_routine);
-}
-
-/*
- * Lock to guard internal/principal key. Usually, this lock has to be held until
- * the caller fetches an internal_key or rotates the principal.
- */
-LWLock *
-tde_lwlock_enc_keys(void)
-{
-	Assert(principalKeyLocalState.sharedPrincipalKeyState);
-
-	return &principalKeyLocalState.sharedPrincipalKeyState->Locks[TDE_LWLOCK_ENC_KEY].lock;
-}
-
 /*
  * Request some pages so we can fit the DSA header, empty hash table plus some
  * extra. Additional memory to grow the hash map will be allocated as needed
@@ -147,50 +116,84 @@ tde_lwlock_enc_keys(void)
  */
 #define CACHE_DSA_INITIAL_SIZE (4096 * 64)
 
-static Size
-required_shared_mem_size(void)
+Size
+PrincipalKeyShmemSize(void)
 {
 	Size		sz = CACHE_DSA_INITIAL_SIZE;
 
 	sz = add_size(sz, sizeof(TdePrincipalKeySharedState));
 	return MAXALIGN(sz);
 }
 
-/*
- * Initialize the shared area for Principal key info.
- * This includes locks and cache area for principal key info
- */
-
-static Size
-initialize_shared_state(void *start_address)
+void
+PrincipalKeyShmemInit(void)
 {
-	TdePrincipalKeySharedState *sharedState = (TdePrincipalKeySharedState *) start_address;
+	bool		found;
+	char	   *free_start;
+	Size		required_shmem_size = PrincipalKeyShmemSize();
 
-	ereport(LOG, errmsg("initializing shared state for principal key"));
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
 
-	sharedState->Locks = GetNamedLWLockTranche(TDE_TRANCHE_NAME);
+	/* Create or attach to the shared memory state */
+	ereport(NOTICE, errmsg("PrincipalKeyShmemInit: requested %ld bytes", required_shmem_size));
+	free_start = ShmemInitStruct("pg_tde", required_shmem_size, &found);
 
-	principalKeyLocalState.sharedPrincipalKeyState = sharedState;
-	principalKeyLocalState.sharedHash = NULL;
+	if (!found)
+	{
+		TdePrincipalKeySharedState *sharedState;
+		Size		sz;
+		Size		dsa_area_size;
+		dsa_area   *dsa;
+		dshash_table *dsh;
 
-	return sizeof(TdePrincipalKeySharedState);
-}
+		/* Now place shared state structure */
+		sharedState = (TdePrincipalKeySharedState *) free_start;
+		sz = MAXALIGN(sizeof(TdePrincipalKeySharedState));
+		free_start += sz;
+		Assert(sz <= required_shmem_size);
 
-static void
-initialize_objects_in_dsa_area(dsa_area *dsa, void *raw_dsa_area)
-{
-	dshash_table *dsh;
-	TdePrincipalKeySharedState *sharedState = principalKeyLocalState.sharedPrincipalKeyState;
+		/* Create DSA area */
+		dsa_area_size = required_shmem_size - sz;
+		Assert(dsa_area_size > 0);
+
+		ereport(LOG, errmsg("creating DSA area of size %lu", dsa_area_size));
+
+		dsa = dsa_create_in_place(free_start,
+								  dsa_area_size,
+								  LWLockNewTrancheId(), 0);
+		dsa_pin(dsa);
+
+		/* Limit area size during population to get a nice error */
+		dsa_set_size_limit(dsa, dsa_area_size);
+
+		principal_key_dsh_params.tranche_id = LWLockNewTrancheId();
+		dsh = dshash_create(dsa, &principal_key_dsh_params, NULL);
+
+		dsa_set_size_limit(dsa, -1);
 
-	ereport(LOG, errmsg("initializing dsa area objects for principal key"));
+		sharedState->Locks = GetNamedLWLockTranche(TDE_TRANCHE_NAME);
+		sharedState->hashHandle = dshash_get_hash_table_handle(dsh);
+		sharedState->rawDsaArea = free_start;
 
-	Assert(sharedState != NULL);
+		principalKeyLocalState.sharedPrincipalKeyState = sharedState;
+		principalKeyLocalState.sharedHash = NULL;
 
-	sharedState->rawDsaArea = raw_dsa_area;
-	principal_key_dsh_params.tranche_id = LWLockNewTrancheId();
-	dsh = dshash_create(dsa, &principal_key_dsh_params, NULL);
-	sharedState->hashHandle = dshash_get_hash_table_handle(dsh);
-	dshash_detach(dsh);
+		dshash_detach(dsh);
+	}
+
+	LWLockRelease(AddinShmemInitLock);
+}
+
+/*
+ * Lock to guard internal/principal key. Usually, this lock has to be held until
+ * the caller fetches an internal_key or rotates the principal.
+ */
+LWLock *
+tde_lwlock_enc_keys(void)
+{
+	Assert(principalKeyLocalState.sharedPrincipalKeyState);
+
+	return &principalKeyLocalState.sharedPrincipalKeyState->Locks[TDE_LWLOCK_ENC_KEY].lock;
 }
 
 /*
@@ -217,12 +220,6 @@ principal_key_info_attach_shmem(void)
 	MemoryContextSwitchTo(oldcontext);
 }
 
-static void
-shared_memory_shutdown(int code, Datum arg)
-{
-	principalKeyLocalState.sharedPrincipalKeyState = NULL;
-}
-
 void
 set_principal_key_with_keyring(const char *key_name, const char *provider_name,
 							   Oid providerOid, Oid dbOid, bool ensure_new_key)

contrib/pg_tde/src/common/pg_tde_shmem.c

@@ -11,122 +11,10 @@
 
 #include "postgres.h"
 
-#include "lib/dshash.h"
-#include "nodes/pg_list.h"
-#include "storage/ipc.h"
-#include "storage/lwlock.h"
-#include "storage/shmem.h"
-
 #include "common/pg_tde_shmem.h"
 
-static void tde_shmem_shutdown(int code, Datum arg);
-
-List	   *registeredShmemRequests = NIL;
-bool		shmemInited = false;
-
-void
-RegisterShmemRequest(const TDEShmemSetupRoutine *routine)
-{
-	Assert(shmemInited == false);
-	registeredShmemRequests = lappend(registeredShmemRequests, (void *) routine);
-}
-
-Size
-TdeRequiredSharedMemorySize(void)
-{
-	Size		sz = 0;
-	ListCell   *lc;
-
-	foreach(lc, registeredShmemRequests)
-	{
-		TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc);
-
-		if (routine->required_shared_mem_size)
-			sz = add_size(sz, routine->required_shared_mem_size());
-	}
-	return MAXALIGN(sz);
-}
-
 int
 TdeRequiredLocksCount(void)
 {
 	return TDE_LWLOCK_COUNT;
 }
-
-void
-TdeShmemInit(void)
-{
-	bool		found;
-	char	   *free_start;
-	Size		required_shmem_size = TdeRequiredSharedMemorySize();
-
-	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-	/* Create or attach to the shared memory state */
-	ereport(NOTICE, errmsg("TdeShmemInit: requested %ld bytes", required_shmem_size));
-	free_start = ShmemInitStruct("pg_tde", required_shmem_size, &found);
-
-	if (!found)
-	{
-		/* First time through ... */
-		dsa_area   *dsa;
-		ListCell   *lc;
-		Size		used_size = 0;
-		Size		dsa_area_size;
-
-		/* Now place all shared state structures */
-		foreach(lc, registeredShmemRequests)
-		{
-			Size		sz = 0;
-			TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc);
-
-			if (routine->init_shared_state)
-			{
-				sz = routine->init_shared_state(free_start);
-				used_size += MAXALIGN(sz);
-				free_start += MAXALIGN(sz);
-				Assert(used_size <= required_shmem_size);
-			}
-		}
-		/* Create DSA area */
-		dsa_area_size = required_shmem_size - used_size;
-		Assert(dsa_area_size > 0);
-
-		ereport(LOG, errmsg("creating DSA area of size %lu", dsa_area_size));
-		dsa = dsa_create_in_place(free_start,
-								  dsa_area_size,
-								  LWLockNewTrancheId(), 0);
-		dsa_pin(dsa);
-		dsa_set_size_limit(dsa, dsa_area_size);
-
-		/* Initialize all DSA area objects */
-		foreach(lc, registeredShmemRequests)
-		{
-			TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc);
-
-			if (routine->init_dsa_area_objects)
-				routine->init_dsa_area_objects(dsa, free_start);
-		}
-		ereport(LOG, errmsg("setting no limit to DSA area of size %lu", dsa_area_size));
-
-		dsa_set_size_limit(dsa, -1);	/* Let it grow outside the shared
-										 * memory */
-
-		shmemInited = true;
-	}
-	LWLockRelease(AddinShmemInitLock);
-	on_shmem_exit(tde_shmem_shutdown, (Datum) 0);
-}
-
-static void
-tde_shmem_shutdown(int code, Datum arg)
-{
-	ListCell   *lc;
-
-	foreach(lc, registeredShmemRequests)
-	{
-		TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc);
-
-		if (routine->shmem_kill)
-			routine->shmem_kill(code, arg);
-	}
-}

contrib/pg_tde/src/include/catalog/tde_principal_key.h

@@ -42,7 +42,8 @@ typedef struct XLogPrincipalKeyRotate
 
 #define SizeoOfXLogPrincipalKeyRotate	offsetof(XLogPrincipalKeyRotate, buff)
 
-extern void InitializePrincipalKeyInfo(void);
+extern void PrincipalKeyShmemInit(void);
+extern Size PrincipalKeyShmemSize(void);
 
 #ifndef FRONTEND
 extern void principal_key_startup_cleanup(Oid databaseId);

contrib/pg_tde/src/include/common/pg_tde_shmem.h

@@ -8,9 +8,6 @@
 #ifndef PG_TDE_SHMEM_H
 #define PG_TDE_SHMEM_H
 
-#include "postgres.h"
-#include "utils/dsa.h"
-
 #define TDE_TRANCHE_NAME "pg_tde_tranche"
 
 typedef enum
@@ -22,39 +19,6 @@ typedef enum
 	TDE_LWLOCK_COUNT
 }			TDELockTypes;
 
-typedef struct TDEShmemSetupRoutine
-{
-	/*
-	 * init_shared_state gets called at the time of extension load you can
-	 * initialize the data structures required to be placed in shared memory
-	 * in this callback The callback must return the size of the shared memory
-	 * area acquired. The argument to the function is the start of the shared
-	 * memory address that can be used to store the shared data structures.
-	 */
-	Size		(*init_shared_state) (void *raw_dsa_area);
-
-	/*
-	 * shmem_startup gets called at the time of postmaster shutdown
-	 */
-	void		(*shmem_kill) (int code, Datum arg);
-
-	/*
-	 * The callback must return the size of the shared memory acquired.
-	 */
-	Size		(*required_shared_mem_size) (void);
-
-	/*
-	 * Gets called after all shared memory structures are initialized and here
-	 * you can create shared memory hash tables or any other shared objects
-	 * that needs to live in DSA area.
-	 */
-	void		(*init_dsa_area_objects) (dsa_area *dsa, void *raw_dsa_area);
-} TDEShmemSetupRoutine;
-
-/* Interface to register the shared memory requests */
-extern void RegisterShmemRequest(const TDEShmemSetupRoutine *routine);
-extern void TdeShmemInit(void);
-extern Size TdeRequiredSharedMemorySize(void);
 extern int	TdeRequiredLocksCount(void);
 
 #endif							/* PG_TDE_SHMEM_H */

contrib/pg_tde/src/pg_tde.c

@@ -55,13 +55,15 @@ PG_FUNCTION_INFO_V1(pg_tdeam_handler);
 static void
 tde_shmem_request(void)
 {
-	Size		sz = TdeRequiredSharedMemorySize();
+	Size		sz = 0;
 	int			required_locks = TdeRequiredLocksCount();
 
+	sz = add_size(sz, PrincipalKeyShmemSize());
 	sz = add_size(sz, TDEXLogEncryptStateSize());
 
 	if (prev_shmem_request_hook)
 		prev_shmem_request_hook();
+
 	RequestAddinShmemSpace(sz);
 	RequestNamedLWLockTranche(TDE_TRANCHE_NAME, required_locks);
 	ereport(LOG, errmsg("tde_shmem_request: requested %ld bytes", sz));
@@ -73,8 +75,8 @@ tde_shmem_startup(void)
 	if (prev_shmem_startup_hook)
 		prev_shmem_startup_hook();
 
-	TdeShmemInit();
-	InitializeKeyProviderInfo();
+	KeyProviderShmemInit();
+	PrincipalKeyShmemInit();
 	TDEXLogShmemInit();
 	TDEXLogSmgrInit();
 }
@@ -100,7 +102,6 @@ _PG_init(void)
 	AesInit();
 	TdeGucInit();
 	TdeEventCaptureInit();
-	InitializePrincipalKeyInfo();
 	InstallFileKeyring();
 	InstallVaultV2Keyring();
 	InstallKmipKeyring();