From acb974d96ebd9c611424bc15c14ac89c4dc9e940 Mon Sep 17 00:00:00 2001
From: Andreas Karlsson <andreas.karlsson@percona.com>
Date: Fri, 2 May 2025 19:43:58 +0200
Subject: [PATCH] Do not use our shmem utils in the key provider code

Since the only thing our key provider code does with shared memory is
look up a LWLock tranche it is quite a waste of lines of codes to use
our own layer on top of the PostgreSQL shared memory.
---
 contrib/pg_tde/src/catalog/tde_keyring.c      | 40 +++----------------
 .../pg_tde/src/include/catalog/tde_keyring.h  |  2 +-
 contrib/pg_tde/src/pg_tde.c                   |  2 +-
 3 files changed, 7 insertions(+), 37 deletions(-)

diff --git a/contrib/pg_tde/src/catalog/tde_keyring.c b/contrib/pg_tde/src/catalog/tde_keyring.c
index c5aee34e377..bab3f19efe9 100644
--- a/contrib/pg_tde/src/catalog/tde_keyring.c
+++ b/contrib/pg_tde/src/catalog/tde_keyring.c
@@ -87,12 +87,10 @@ PG_FUNCTION_INFO_V1(pg_tde_list_all_global_key_providers);
 
 static const char *get_keyring_provider_typename(ProviderType p_type);
 static List *GetAllKeyringProviders(Oid dbOid);
-static Size initialize_shared_state(void *start_address);
 static Datum pg_tde_add_key_provider_internal(PG_FUNCTION_ARGS, Oid dbOid);
 static Datum pg_tde_change_key_provider_internal(PG_FUNCTION_ARGS, Oid dbOid);
 static Datum pg_tde_delete_key_provider_internal(PG_FUNCTION_ARGS, Oid dbOid);
 static Datum pg_tde_list_all_key_providers_internal(PG_FUNCTION_ARGS, const char *fname, Oid dbOid);
-static Size required_shared_mem_size(void);
 static List *scan_key_provider_file(ProviderScanType scanType, void *scanKey, Oid dbOid);
 static void save_new_key_provider_info(KeyringProviderRecord *provider, Oid databaseId);
 static void modify_key_provider_info(KeyringProviderRecord *provider, Oid databaseId);
@@ -101,47 +99,19 @@ static void check_provider_record(KeyringProviderRecord *provider_record);
 
 #define PG_TDE_LIST_PROVIDERS_COLS 4
 
-typedef struct TdeKeyProviderInfoSharedState
-{
-	LWLockPadded *Locks;
-} TdeKeyProviderInfoSharedState;
-
-TdeKeyProviderInfoSharedState *sharedPrincipalKeyState = NULL;	/* Lives in shared state */
-
-static const TDEShmemSetupRoutine key_provider_info_shmem_routine = {
-	.init_shared_state = initialize_shared_state,
-	.init_dsa_area_objects = NULL,
-	.required_shared_mem_size = required_shared_mem_size,
-	.shmem_kill = NULL
-};
-
-static Size
-required_shared_mem_size(void)
-{
-	return MAXALIGN(sizeof(TdeKeyProviderInfoSharedState));
-}
-
-static Size
-initialize_shared_state(void *start_address)
-{
-	sharedPrincipalKeyState = (TdeKeyProviderInfoSharedState *) start_address;
-	sharedPrincipalKeyState->Locks = GetNamedLWLockTranche(TDE_TRANCHE_NAME);
-
-	return sizeof(TdeKeyProviderInfoSharedState);
-}
+static LWLockPadded *tdeLocks = NULL;	/* Lives in shared state */
 
 static inline LWLock *
 tde_provider_info_lock(void)
 {
-	Assert(sharedPrincipalKeyState);
-	return &sharedPrincipalKeyState->Locks[TDE_LWLOCK_PI_FILES].lock;
+	Assert(tdeLocks);
+	return &tdeLocks[TDE_LWLOCK_PI_FILES].lock;
 }
 
 void
-InitializeKeyProviderInfo(void)
+KeyProviderShmemInit(void)
 {
-	ereport(LOG, errmsg("initializing TDE key provider info"));
-	RegisterShmemRequest(&key_provider_info_shmem_routine);
+	tdeLocks = GetNamedLWLockTranche(TDE_TRANCHE_NAME);
 }
 
 void
diff --git a/contrib/pg_tde/src/include/catalog/tde_keyring.h b/contrib/pg_tde/src/include/catalog/tde_keyring.h
index 11c24027747..5dd09f3e2c1 100644
--- a/contrib/pg_tde/src/include/catalog/tde_keyring.h
+++ b/contrib/pg_tde/src/include/catalog/tde_keyring.h
@@ -33,7 +33,7 @@ typedef struct KeyringProviderRecordInFile
 extern GenericKeyring *GetKeyProviderByName(const char *provider_name, Oid dbOid);
 extern GenericKeyring *GetKeyProviderByID(int provider_id, Oid dbOid);
 extern ProviderType get_keyring_provider_from_typename(char *provider_type);
-extern void InitializeKeyProviderInfo(void);
+extern void KeyProviderShmemInit(void);
 extern void key_provider_startup_cleanup(Oid databaseId);
 extern bool get_keyring_info_file_record_by_name(char *provider_name,
 												 Oid database_id,
diff --git a/contrib/pg_tde/src/pg_tde.c b/contrib/pg_tde/src/pg_tde.c
index f38ad795609..4157b72d6f7 100644
--- a/contrib/pg_tde/src/pg_tde.c
+++ b/contrib/pg_tde/src/pg_tde.c
@@ -74,6 +74,7 @@ tde_shmem_startup(void)
 		prev_shmem_startup_hook();
 
 	TdeShmemInit();
+	InitializeKeyProviderInfo();
 	TDEXLogShmemInit();
 	TDEXLogSmgrInit();
 }
@@ -100,7 +101,6 @@ _PG_init(void)
 	TdeGucInit();
 	TdeEventCaptureInit();
 	InitializePrincipalKeyInfo();
-	InitializeKeyProviderInfo();
 	InstallFileKeyring();
 	InstallVaultV2Keyring();
 	InstallKmipKeyring();