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