From b709662c25e13896387d5e7966119e425d4808ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anders=20=C3=85strand?= <anders.astrand@percona.com>
Date: Mon, 18 Aug 2025 17:48:46 +0200
Subject: [PATCH] Hold required lock when initializing shmem

According to the documentation, each backend is supposed to hold
AddinShmemInitLock when calling ShmemInitStruct. We only did that for
half of our calls before this patch.
---
 contrib/pg_tde/src/access/pg_tde_xlog_smgr.c   | 2 ++
 contrib/pg_tde/src/catalog/tde_principal_key.c | 4 +---
 contrib/pg_tde/src/pg_tde.c                    | 4 ++++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c b/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
index abf86c1ccba..01b368d55e0 100644
--- a/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
+++ b/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
@@ -158,6 +158,8 @@ TDEXLogShmemInit(void)
 {
 	bool		foundBuf;
 
+	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
+
 	EncryptionState = (EncryptionStateData *)
 		ShmemInitStruct("TDE XLog Encryption State",
 						TDEXLogEncryptStateSize(),
diff --git a/contrib/pg_tde/src/catalog/tde_principal_key.c b/contrib/pg_tde/src/catalog/tde_principal_key.c
index e44cabfea44..abb160f47e2 100644
--- a/contrib/pg_tde/src/catalog/tde_principal_key.c
+++ b/contrib/pg_tde/src/catalog/tde_principal_key.c
@@ -127,7 +127,7 @@ PrincipalKeyShmemInit(void)
 	char	   *free_start;
 	Size		required_shmem_size = PrincipalKeyShmemSize();
 
-	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	Assert(LWLockHeldByMeInMode(AddinShmemInitLock, LW_EXCLUSIVE));
 
 	/* Create or attach to the shared memory state */
 	ereport(NOTICE, errmsg("PrincipalKeyShmemInit: requested %ld bytes", required_shmem_size));
@@ -175,8 +175,6 @@ PrincipalKeyShmemInit(void)
 
 		dshash_detach(dsh);
 	}
-
-	LWLockRelease(AddinShmemInitLock);
 }
 
 /*
diff --git a/contrib/pg_tde/src/pg_tde.c b/contrib/pg_tde/src/pg_tde.c
index 977ba0ccd35..e6f32f099bb 100644
--- a/contrib/pg_tde/src/pg_tde.c
+++ b/contrib/pg_tde/src/pg_tde.c
@@ -64,11 +64,15 @@ tde_shmem_startup(void)
 	if (prev_shmem_startup_hook)
 		prev_shmem_startup_hook();
 
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+
 	KeyProviderShmemInit();
 	PrincipalKeyShmemInit();
 	TDEXLogShmemInit();
 	TDEXLogSmgrInit();
 	TDEXLogSmgrInitWrite(EncryptXLog);
+
+	LWLockRelease(AddinShmemInitLock);
 }
 
 void