WIP: Delete old key in the same scan as we add the new one

Author: jeltz

contrib/pg_tde/src/access/pg_tde_tdemap.c

@@ -72,6 +72,7 @@ static void pg_tde_sign_principal_key_info(TDESignedPrincipalKeyInfo *signed_key
 static void pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path);
 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 void pg_tde_replace_key_map_entry(const RelFileLocator *rlocator, const InternalKey *rel_key_data, TDEPrincipalKey *principal_key);
 
 void
 pg_tde_save_smgr_key(RelFileLocator rel, const InternalKey *rel_key_data)
@@ -88,7 +89,7 @@ pg_tde_save_smgr_key(RelFileLocator rel, const InternalKey *rel_key_data)
 				errhint("Use pg_tde_set_key_using_database_key_provider() or pg_tde_set_key_using_global_key_provider() to configure one."));
 	}
 
-	pg_tde_write_key_map_entry(&rel, rel_key_data, principal_key);
+	pg_tde_replace_key_map_entry(&rel, rel_key_data, principal_key);
 	LWLockRelease(lock_pk);
 }
 
@@ -362,6 +363,65 @@ pg_tde_delete_principal_key(Oid dbOid)
 	durable_unlink(path, ERROR);
 }
 
+/*
+ * The caller must hold an exclusive lock on the key file to avoid
+ * concurrent in place updates leading to data conflicts.
+ */
+void
+pg_tde_replace_key_map_entry(const RelFileLocator *rlocator, const InternalKey *rel_key_data, TDEPrincipalKey *principal_key)
+{
+	char		db_map_path[MAXPGPATH];
+	int			map_fd;
+	off_t		curr_pos = 0;
+	off_t		write_pos = 0;
+	TDEMapEntry write_map_entry;
+	TDESignedPrincipalKeyInfo signed_key_Info;
+
+	Assert(rlocator);
+
+	pg_tde_set_db_file_path(rlocator->dbOid, db_map_path);
+
+	pg_tde_sign_principal_key_info(&signed_key_Info, principal_key);
+
+	/* Open and validate file for basic correctness. */
+	map_fd = pg_tde_open_file_write(db_map_path, &signed_key_Info, false, &curr_pos);
+
+	/*
+	 * Read until we find an empty slot. Otherwise, read until end. This seems
+	 * to be less frequent than vacuum. So let's keep this function here
+	 * rather than overloading the vacuum process.
+	 */
+	while (1)
+	{
+		TDEMapEntry read_map_entry;
+		off_t		prev_pos = curr_pos;
+
+		if (!pg_tde_read_one_map_entry(map_fd, &read_map_entry, &curr_pos))
+		{
+			if (write_pos == 0)
+				write_pos = prev_pos;
+			break;
+		}
+
+		if (read_map_entry.spcOid == rlocator->spcOid && read_map_entry.relNumber == rlocator->relNumber)
+		{
+			write_pos = prev_pos;
+			break;
+		}
+
+		if (write_pos == 0 && read_map_entry.type == MAP_ENTRY_EMPTY)
+			write_pos = prev_pos;
+	}
+
+	/* Initialize map entry and encrypt key */
+	pg_tde_initialize_map_entry(&write_map_entry, principal_key, rlocator, rel_key_data);
+
+	/* Write the given entry at curr_pos; i.e. the free entry. */
+	pg_tde_write_one_map_entry(map_fd, &write_map_entry, &write_pos, db_map_path);
+
+	CloseTransientFile(map_fd);
+}
+
 #endif							/* !FRONTEND */
 
 /*

contrib/pg_tde/src/smgr/pg_tde_smgr.c

@@ -102,8 +102,6 @@ tde_smgr_create_key_redo(const RelFileLocator *rlocator)
 {
 	InternalKey key;
 
-	pg_tde_free_key_map_entry(*rlocator);
-
 	pg_tde_generate_internal_key(&key, TDE_KEY_TYPE_SMGR);
 
 	pg_tde_save_smgr_key(*rlocator, &key);
@@ -373,18 +371,18 @@ tde_mdcreate(RelFileLocator relold, SMgrRelation reln, ForkNumber forknum, bool
 	if (forknum != MAIN_FORKNUM)
 		return;
 
-	/*
-	 * If we have a key for this relation already, we need to remove it. This
-	 * can happen if OID is re-used after a crash left a key for a
-	 * non-existing relation in the key file.
-	 */
-	if (!RelFileLocatorBackendIsTemp(reln->smgr_rlocator))
-		pg_tde_free_key_map_entry(reln->smgr_rlocator.locator);
-
 	if (!tde_smgr_should_encrypt(&reln->smgr_rlocator, &relold))
 	{
+		/*
+		 * If we have a key for this relation already, we need to remove it.
+		 * This can happen if OID is re-used after a crash left a key for a
+		 * non-existing relation in the key file.
+		 */
 		if (!RelFileLocatorBackendIsTemp(reln->smgr_rlocator))
+		{
+			pg_tde_free_key_map_entry(reln->smgr_rlocator.locator);
 			tde_smgr_log_delete_key(&reln->smgr_rlocator);
+		}
 
 		tdereln->encryption_status = RELATION_NOT_ENCRYPTED;
 		return;