percona
diff --git a/‎contrib/pg_tde/src/access/pg_tde_tdemap.c
Lines changed: 5 additions & 47 deletions b/‎contrib/pg_tde/src/access/pg_tde_tdemap.c
Lines changed: 5 additions & 47 deletions
diff --git a/‎contrib/pg_tde/src/catalog/tde_principal_key.c
Lines changed: 2 additions & 2 deletions b/‎contrib/pg_tde/src/catalog/tde_principal_key.c
Lines changed: 2 additions & 2 deletions
diff --git a/‎contrib/pg_tde/src/pg_tde_event_capture.c
Lines changed: 5 additions & 1 deletion b/‎contrib/pg_tde/src/pg_tde_event_capture.c
Lines changed: 5 additions & 1 deletion
@@ -524,37 +524,6 @@ pg_tde_perform_rotate_key(TDEPrincipalKey *principal_key, TDEPrincipalKey *new_p
 	}
 }
 
-/*
- * Checks if the key map file for the given database has any non-empty entries.
- *
- * The caller must hold the lock on the files before calling this function.
-*/
-bool
-pg_tde_is_key_map_empty(Oid dbOid)
-{
-	char		path[MAXPGPATH];
-	off_t		pos;
-	int			fd;
-	TDEMapEntry map_entry;
-
-	pg_tde_set_db_file_path(dbOid, path);
-
-	fd = pg_tde_open_file_read(path, false, &pos);
-
-	while (pg_tde_read_one_map_entry(fd, &map_entry, &pos))
-	{
-		if (map_entry.type != MAP_ENTRY_EMPTY)
-		{
-			close(fd);
-			return false;
-		}
-	}
-
-	close(fd);
-	return true;
-}
-
-
 void
 pg_tde_delete_principal_key_redo(Oid dbOid)
 {
@@ -578,20 +547,10 @@ pg_tde_delete_principal_key(Oid dbOid, bool write_xlog)
 	char		path[MAXPGPATH];
 
 	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
+	Assert(pg_tde_count_relations(dbOid) == 0);
 
 	pg_tde_set_db_file_path(dbOid, path);
 
-	/*
-	 * Check that there aro no local keys in the map file, i.e. nothing is
-	 * encrypted with the principal key.
-	 */
-	if (!pg_tde_is_key_map_empty(dbOid))
-	{
-		ereport(ERROR,
-				errcode(ERRCODE_INTERNAL_ERROR),
-				errmsg("could not delete key, key map file is not empty"));
-	}
-
 	if (write_xlog)
 	{
 		XLogBeginInsert();
@@ -742,20 +701,21 @@ pg_tde_find_map_entry(const RelFileLocator *rlocator, TDEMapEntryType key_type,
  * positive is more harmful this might not be.
  *
  * Works even if the database has no map file.
+ * The caller must hold a shared or exclusive lock on the
+ * tde_lwlock_enc_keys() lock.
  */
 int
 pg_tde_count_relations(Oid dbOid)
 {
 	char		db_map_path[MAXPGPATH];
-	LWLock	   *lock_pk = tde_lwlock_enc_keys();
 	File		map_fd;
 	off_t		curr_pos = 0;
 	TDEMapEntry map_entry;
 	int			count = 0;
 
-	pg_tde_set_db_file_path(dbOid, db_map_path);
+	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_SHARED) || LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
 
-	LWLockAcquire(lock_pk, LW_SHARED);
+	pg_tde_set_db_file_path(dbOid, db_map_path);
 
 	map_fd = pg_tde_open_file_read(db_map_path, true, &curr_pos);
 	if (map_fd < 0)
@@ -769,8 +729,6 @@ pg_tde_count_relations(Oid dbOid)
 
 	close(map_fd);
 
-	LWLockRelease(lock_pk);
-
 	return count;
 }
 
 
@@ -608,7 +608,7 @@ pg_tde_delete_key(PG_FUNCTION_ARGS)
 	 * If database has something encryted, we can try to fallback to the
 	 * default principal key
 	 */
-	if (!pg_tde_is_key_map_empty(MyDatabaseId))
+	if (pg_tde_count_relations(MyDatabaseId) != 0)
 	{
 		default_principal_key = GetPrincipalKeyNoDefault(DEFAULT_DATA_TDE_OID, LW_EXCLUSIVE);
 		if (default_principal_key == NULL)
@@ -684,7 +684,7 @@ pg_tde_delete_default_key(PG_FUNCTION_ARGS)
 			 * If database key map is empty we should remove it befroe
 			 * deleting default principal key
 			 */
-			if (pg_tde_is_key_map_empty(dbOid))
+			if (pg_tde_count_relations(dbOid) == 0)
 			{
 				pg_tde_delete_principal_key(dbOid, true);
 				clear_principal_key_cache(dbOid);
 
@@ -633,7 +633,11 @@ pg_tde_proccess_utility(PlannedStmt *pstmt,
 
 					if (dbOid != InvalidOid)
 					{
-						int			count = pg_tde_count_relations(dbOid);
+						int count;
+
+						LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
+						count = pg_tde_count_relations(dbOid);
+						LWLockRelease(tde_lwlock_enc_keys());
 
 						if (count > 0)
 							ereport(ERROR,