PG-1617 Stop using cache for keys of temporary tables

Since are going to remove the key cache we need to make sure to figure
out another solution for the keys of temporary tables. The new solution
is to store them in a per-backend hash table from which we delete the key
when we drop the relation.

contrib/pg_tde/src/access/pg_tde_tdemap.c

@@ -20,6 +20,7 @@
 #include "access/xlog_internal.h"
 #include "access/xloginsert.h"
 #include "utils/builtins.h"
+#include "utils/hsearch.h"
 #include "miscadmin.h"
 
 #include "access/pg_tde_tdemap.h"
@@ -99,6 +100,24 @@ RelKeyCache tde_rel_key_cache = {
 	.cap = 0,
 };
 
+typedef struct
+{
+	RelFileLocator rel;
+	InternalKey key;
+} TempRelKeyEntry;
+
+#ifndef FRONTEND
+
+/* Arbitrarily picked small number of temporary relations */
+#define INIT_TEMP_RELS 16
+
+/*
+ * Each backend has a hashtable that stores the keys for all temporary tables.
+ */
+static HTAB *TempRelKeys = NULL;
+
+#endif
+
 /*
  * TODO: WAL should have its own RelKeyCache
  */
@@ -125,6 +144,7 @@ static int	pg_tde_file_header_write(const char *tde_filename, int fd, const TDES
 static void pg_tde_sign_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const TDEPrincipalKey *principal_key);
 static void pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path);
 static void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_data, TDEPrincipalKey *principal_key);
+static void pg_tde_free_key_map_entry(const RelFileLocator *rlocator);
 static int	keyrotation_init_file(const TDESignedPrincipalKeyInfo *signed_key_info, char *rotated_filename, const char *filename, off_t *curr_pos);
 static void finalize_key_rotation(const char *path_old, const char *path_new);
 static int	pg_tde_open_file_write(const char *tde_filename, const TDESignedPrincipalKeyInfo *signed_key_info, bool truncate, off_t *curr_pos);
@@ -141,11 +161,29 @@ pg_tde_create_smgr_key(const RelFileLocatorBackend *newrlocator)
 static InternalKey *
 pg_tde_create_smgr_key_temp(const RelFileLocator *newrlocator)
 {
-	InternalKey int_key;
+	TempRelKeyEntry *entry;
+	bool		found;
+
+	if (TempRelKeys == NULL)
+	{
+		HASHCTL		ctl;
+
+		ctl.keysize = sizeof(RelFileLocator);
+		ctl.entrysize = sizeof(TempRelKeyEntry);
+		TempRelKeys = hash_create("pg_tde temporary relation keys",
+								  INIT_TEMP_RELS,
+								  &ctl,
+								  HASH_ELEM | HASH_BLOBS);
+	}
+
+	entry = (TempRelKeyEntry *) hash_search(TempRelKeys,
+											newrlocator,
+											HASH_ENTER, &found);
+	Assert(!found);
 
-	pg_tde_generate_internal_key(&int_key, TDE_KEY_TYPE_SMGR);
+	pg_tde_generate_internal_key(&entry->key, TDE_KEY_TYPE_SMGR);
 
-	return pg_tde_put_key_into_cache(newrlocator, &int_key);
+	return &entry->key;
 }
 
 static InternalKey *
@@ -274,6 +312,18 @@ pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocat
 	LWLockRelease(tde_lwlock_enc_keys());
 }
 
+void
+DeleteSMGRRelationKey(RelFileLocatorBackend rel)
+{
+	if (RelFileLocatorBackendIsTemp(rel))
+	{
+		if (TempRelKeys)
+			hash_search(TempRelKeys, &rel.locator, HASH_REMOVE, NULL);
+	}
+	else
+		pg_tde_free_key_map_entry(&rel.locator);
+}
+
 /*
  * Deletes the key map file for a given database.
  */
@@ -491,7 +541,7 @@ pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_
  * This fucntion is called by the pg_tde SMGR when storage is unlinked on
  * transaction commit/abort.
  */
-void
+static void
 pg_tde_free_key_map_entry(const RelFileLocator *rlocator)
 {
 	char		db_map_path[MAXPGPATH];
@@ -1047,6 +1097,24 @@ pg_tde_get_principal_key_info(Oid dbOid)
 	return signed_key_info;
 }
 
+static InternalKey *
+pg_tde_get_temporary_rel_key(const RelFileLocator *rel)
+{
+#ifndef FRONTEND
+	TempRelKeyEntry *entry;
+
+	if (TempRelKeys == NULL)
+		return NULL;
+
+	entry = hash_search(TempRelKeys, rel, HASH_FIND, NULL);
+
+	if (entry)
+		return &entry->key;
+#endif
+
+	return NULL;
+}
+
 /*
  * Figures out whether a relation is encrypted or not, but without trying to
  * decrypt the key if it is. This also means that this function cannot push the
@@ -1062,7 +1130,7 @@ IsSMGRRelationEncrypted(RelFileLocatorBackend rel)
 	Assert(rel.locator.relNumber != InvalidRelFileNumber);
 
 	if (RelFileLocatorBackendIsTemp(rel))
-		return pg_tde_get_key_from_cache(&rel.locator, TDE_KEY_TYPE_SMGR) != NULL;
+		return pg_tde_get_temporary_rel_key(&rel.locator) != NULL;
 	else if (pg_tde_get_key_from_cache(&rel.locator, TDE_KEY_TYPE_SMGR))
 		return true;
 
@@ -1090,7 +1158,7 @@ GetSMGRRelationKey(RelFileLocatorBackend rel)
 	Assert(rel.locator.relNumber != InvalidRelFileNumber);
 
 	if (RelFileLocatorBackendIsTemp(rel))
-		return pg_tde_get_key_from_cache(&rel.locator, TDE_KEY_TYPE_SMGR);
+		return pg_tde_get_temporary_rel_key(&rel.locator);
 	else
 	{
 		InternalKey *key;

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

@@ -90,7 +90,6 @@ extern void pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path
 extern InternalKey *pg_tde_create_smgr_key(const RelFileLocatorBackend *newrlocator);
 extern void pg_tde_create_smgr_key_perm_redo(const RelFileLocator *newrlocator);
 extern void pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, uint32 flags);
-extern void pg_tde_free_key_map_entry(const RelFileLocator *rlocator);
 
 #define PG_TDE_MAP_FILENAME			"%d_keys"
 
@@ -102,6 +101,8 @@ pg_tde_set_db_file_path(Oid dbOid, char *path)
 
 extern bool IsSMGRRelationEncrypted(RelFileLocatorBackend rel);
 extern InternalKey *GetSMGRRelationKey(RelFileLocatorBackend rel);
+extern void DeleteSMGRRelationKey(RelFileLocatorBackend rel);
+
 extern int	pg_tde_count_relations(Oid dbOid);
 
 extern void pg_tde_delete_tde_files(Oid dbOid);

contrib/pg_tde/src/smgr/pg_tde_smgr.c

@@ -161,8 +161,8 @@ tde_mdunlink(RelFileLocatorBackend rlocator, ForkNumber forknum, bool isRedo)
 	 */
 	if (forknum == MAIN_FORKNUM || forknum == InvalidForkNumber)
 	{
-		if (!RelFileLocatorBackendIsTemp(rlocator) && IsSMGRRelationEncrypted(rlocator))
-			pg_tde_free_key_map_entry(&rlocator.locator);
+		if (IsSMGRRelationEncrypted(rlocator))
+			DeleteSMGRRelationKey(rlocator);
 	}
 }