Merge branch 'TDE_REL_17_STABLE' into backup_new_wal_key

Author: dAdAbird

ci_scripts/suppressions/lsan.supp

@@ -59,3 +59,8 @@ leak:GenerateRecoveryConfig
 # Returns strdup'd string which never gets freed in frontend tools. Trying to 
 # free it leads to comiler complains that it discards 'const' qualifier.
 leak:get_progname
+
+# A buch of not freed allocations in putVariableInt right before exit(1). And 
+# malloced `threads` ("Thread state") never gets released (allocated in main(),
+# once, and not in loop).
+leak:bin/pgbench/pgbench.c

contrib/pg_tde/src/access/pg_tde_xlog_keys.c

@@ -16,6 +16,7 @@
 #include "common/pg_tde_utils.h"
 #include "encryption/enc_aes.h"
 #include "encryption/enc_tde.h"
+#include "utils/palloc.h"
 
 #ifdef FRONTEND
 #include "pg_tde_fe.h"
@@ -43,7 +44,9 @@ typedef struct WalKeyFileEntry
 } WalKeyFileEntry;
 
 static WALKeyCacheRec *tde_wal_key_cache = NULL;
+static WALKeyCacheRec *tde_wal_prealloc_record = NULL;
 static WALKeyCacheRec *tde_wal_key_last_rec = NULL;
+static WalEncryptionKey *tde_wal_prealloc_key = NULL;
 
 static WALKeyCacheRec *pg_tde_add_wal_key_to_cache(WalEncryptionKey *cached_key);
 static WalEncryptionKey *pg_tde_decrypt_wal_key(const TDEPrincipalKey *principal_key, WalKeyFileEntry *entry);
@@ -326,6 +329,34 @@ pg_tde_fetch_wal_keys(WalLocation start)
 	return return_wal_rec;
 }
 
+/*
+ * In special cases, we have to add one more record to the WAL key cache during write (in the critical section, when we can't allocate).
+ * This method is a helper to deal with that: when adding a single key, we potentially allocate 2 records.
+ * These variables help preallocate them, so in the critical section we can just use the already allocated objects, and update the cache with them.
+ *
+ * While this is somewhat a hack, it is also simple, safe, reliable, and way easier to implement than to refactor the cache or the decryption/encryption loop.
+ */
+void
+pg_tde_wal_cache_extra_palloc(void)
+{
+#ifndef FRONTEND
+	MemoryContext oldCtx;
+
+	oldCtx = MemoryContextSwitchTo(TopMemoryContext);
+#endif
+	if (tde_wal_prealloc_record == NULL)
+	{
+		tde_wal_prealloc_record = palloc0_object(WALKeyCacheRec);
+	}
+	if (tde_wal_prealloc_key == NULL)
+	{
+		tde_wal_prealloc_key = palloc0_object(WalEncryptionKey);
+	}
+#ifndef FRONTEND
+	MemoryContextSwitchTo(oldCtx);
+#endif
+}
+
 static WALKeyCacheRec *
 pg_tde_add_wal_key_to_cache(WalEncryptionKey *key)
 {
@@ -335,7 +366,8 @@ pg_tde_add_wal_key_to_cache(WalEncryptionKey *key)
 
 	oldCtx = MemoryContextSwitchTo(TopMemoryContext);
 #endif
-	wal_rec = palloc0_object(WALKeyCacheRec);
+	wal_rec = tde_wal_prealloc_record == NULL ? palloc0_object(WALKeyCacheRec) : tde_wal_prealloc_record;
+	tde_wal_prealloc_record = NULL;
 #ifndef FRONTEND
 	MemoryContextSwitchTo(oldCtx);
 #endif
@@ -553,7 +585,9 @@ pg_tde_write_wal_key_file_entry(const WalEncryptionKey *rel_key_data,
 static WalEncryptionKey *
 pg_tde_decrypt_wal_key(const TDEPrincipalKey *principal_key, WalKeyFileEntry *entry)
 {
-	WalEncryptionKey *key = palloc_object(WalEncryptionKey);
+	WalEncryptionKey *key = tde_wal_prealloc_key == NULL ? palloc_object(WalEncryptionKey) : tde_wal_prealloc_key;
+
+	tde_wal_prealloc_key = NULL;
 
 	Assert(principal_key);
 

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(),
@@ -248,6 +250,7 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)
 
 		/* cache is empty, prefetch keys from disk */
 		pg_tde_fetch_wal_keys(start);
+		pg_tde_wal_cache_extra_palloc();
 	}
 
 	if (key)
@@ -284,8 +287,8 @@ TDEXLogWriteEncryptedPagesOldKeys(int fd, const void *buf, size_t count, off_t o
 	memcpy(enc_buff, buf, count);
 
 	/*
-	 * This method potentially allocates, but only in very early execution
-	 * Shouldn't happen in a write, where we are in a critical section
+	 * This method potentially allocates, but only in very early execution Can
+	 * happen during a write, but we have one more cache entry preallocated.
 	 */
 	TDEXLogCryptBuffer(buf, enc_buff, count, offset, tli, segno, segSize);
 
@@ -342,7 +345,7 @@ tde_ensure_xlog_key_location(WalLocation loc)
 	lastKeyUsable = (TDEXLogGetEncKeyLsn() != 0);
 	pg_read_barrier();
 
-	if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && !lastKeyUsable)
+	if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && !lastKeyUsable && afterWriteKey)
 	{
 		WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();
 
@@ -452,10 +455,8 @@ TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
 		WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();
 		WalLocation write_loc = {.tli = TDEXLogGetEncKeyTli(),.lsn = write_key_lsn};
 
-		Assert(last_key);
-
 		/* write has generated a new key, need to fetch it */
-		if (wal_location_cmp(last_key->start, write_loc) < 0)
+		if (last_key != NULL && wal_location_cmp(last_key->start, write_loc) < 0)
 		{
 			pg_tde_fetch_wal_keys(write_loc);
 

contrib/pg_tde/src/bin/pg_tde_archive_decrypt.c

@@ -212,6 +212,8 @@ main(int argc, char *argv[])
 	if (system(command) != 0)
 		pg_fatal("ARCHIVE-COMMAND \"%s\" failed: %m", command);
 
+	free(command);
+
 	if (issegment)
 	{
 		if (unlink(tmppath) < 0)

contrib/pg_tde/src/bin/pg_tde_restore_encrypt.c

@@ -205,6 +205,8 @@ main(int argc, char *argv[])
 	if (system(command) != 0)
 		pg_fatal("RESTORE-COMMAND \"%s\" failed: %m", command);
 
+	free(command);
+
 	if (issegment)
 	{
 		write_encrypted_segment(targetpath, sourcename, tmppath);

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);
 }
 
 /*

contrib/pg_tde/src/include/access/pg_tde_xlog_keys.h

@@ -82,5 +82,6 @@ extern WalEncryptionKey *pg_tde_read_last_wal_key(void);
 extern void pg_tde_save_server_key(const TDEPrincipalKey *principal_key, bool write_xlog);
 extern void pg_tde_save_server_key_redo(const TDESignedPrincipalKeyInfo *signed_key_info);
 extern void pg_tde_wal_last_key_set_location(WalLocation loc);
+extern void pg_tde_wal_cache_extra_palloc(void);
 
 #endif							/* PG_TDE_XLOG_KEYS_H */

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