[RFC PATCH proxmox-backup 2/4] tape: restore: use a drive pool for the selective restore

[Jakob-Tobias Winter <tobias@stuttgart.de>]

In preparation for restoring from several drives in parallel, load the
media of the first two phases of a selective restore through a drive
pool, and start the messages logged while reading a media with the
prefix the pool provides. The pool only contains the given drive for
now, so the media are loaded in the same order as before and the
prefix is empty.

No functional change intended.

Signed-off-by: Jakob-Tobias Winter <tobias@stuttgart.de>
---
 src/api2/tape/restore.rs | 116 ++++++++++++++++++++++++---------------
 1 file changed, 72 insertions(+), 44 deletions(-)

diff --git a/src/api2/tape/restore.rs b/src/api2/tape/restore.rs
index 6544aa726..ee8906f85 100644
--- a/src/api2/tape/restore.rs
+++ b/src/api2/tape/restore.rs
@@ -40,7 +40,9 @@ use pbs_tape::{
 use crate::backup::check_ns_modification_privs;
 use crate::tape::{
     Inventory, MediaCatalog, MediaId, MediaSet, MediaSetCatalog, TAPE_STATUS_DIR,
-    drive::{TapeDriver, lock_tape_device, request_and_load_media, set_tape_device_state},
+    drive::{
+        DrivePool, TapeDriver, lock_tape_device, request_and_load_media, set_tape_device_state,
+    },
     file_formats::{
         CatalogArchiveHeader, ChunkArchiveDecoder, ChunkArchiveHeader,
         PROXMOX_BACKUP_CATALOG_ARCHIVE_MAGIC_1_0, PROXMOX_BACKUP_CATALOG_ARCHIVE_MAGIC_1_1,
@@ -418,7 +420,7 @@ pub fn restore(
                     inventory,
                     media_set_uuid,
                     drive_config,
-                    &drive,
+                    std::slice::from_ref(&drive),
                     store_map,
                     restore_owner,
                     &notification_mode,
@@ -628,7 +630,7 @@ fn restore_list_worker(
     inventory: Inventory,
     media_set_uuid: Uuid,
     drive_config: SectionConfigData,
-    drive_name: &str,
+    drives: &[String],
     store_map: DataStoreMap,
     restore_owner: &Authid,
     notification_mode: &TapeNotificationMode,
@@ -636,6 +638,7 @@ fn restore_list_worker(
     auth_id: &Authid,
 ) -> Result<(), Error> {
     let catalog = get_media_set_catalog(&inventory, &media_set_uuid)?;
+    let drive_pool = DrivePool::new(worker, &drive_config, drives, notification_mode);
 
     let mut datastore_locks = Vec::new();
     let mut snapshot_file_hash: BTreeMap<Uuid, Vec<(u64, String)>> = BTreeMap::new();
@@ -760,30 +763,39 @@ fn restore_list_worker(
 
         info!("Phase 1: temporarily restore snapshots to temp dir");
         log_required_tapes(&inventory, snapshot_file_hash.keys());
-        let mut datastore_chunk_map: HashMap<String, HashSet<[u8; 32]>> = HashMap::new();
-        let mut tmp_paths = Vec::new();
-        for (media_uuid, file_list) in snapshot_file_hash.iter_mut() {
-            let media_id = inventory.lookup_media(media_uuid).unwrap();
-            let (drive, info) = request_and_load_media(
-                &worker,
-                &drive_config,
-                drive_name,
-                &media_id.label,
-                notification_mode,
-            )?;
+        let jobs = snapshot_file_hash
+            .into_iter()
+            .map(|(media_uuid, file_list)| {
+                let media_id = inventory.lookup_media(&media_uuid).unwrap();
+                (media_id.label.clone(), file_list)
+            })
+            .collect();
+        let results = drive_pool.run(jobs, |media, mut file_list| {
             file_list.sort_unstable();
-
+            let mut chunk_map = HashMap::new();
             let tmp_path = restore_snapshots_to_tmpdir(
-                worker.clone(),
+                media.worker,
                 &store_map,
-                file_list,
-                drive,
-                &info,
+                &file_list,
+                media.drive,
+                &media.media_id,
                 &media_set_uuid,
-                &mut datastore_chunk_map,
+                &mut chunk_map,
+                &media.log_prefix,
             )
             .map_err(|err| format_err!("could not restore snapshots to tmpdir: {}", err))?;
+            Ok((tmp_path, chunk_map))
+        })?;
+        let mut datastore_chunk_map: HashMap<String, HashSet<[u8; 32]>> = HashMap::new();
+        let mut tmp_paths = Vec::new();
+        for (tmp_path, chunk_map) in results {
             tmp_paths.extend(tmp_path);
+            for (source_datastore, chunks) in chunk_map {
+                datastore_chunk_map
+                    .entry(source_datastore)
+                    .or_default()
+                    .extend(chunks);
+            }
         }
 
         // sorted media_uuid => (sorted file_num => (set of digests)))
@@ -816,17 +828,23 @@ fn restore_list_worker(
             info!("All chunks are already present, skip phase 2...");
         }
 
-        for (media_uuid, file_chunk_map) in media_file_chunk_map.iter_mut() {
-            let media_id = inventory.lookup_media(media_uuid).unwrap();
-            let (mut drive, _info) = request_and_load_media(
-                &worker,
-                &drive_config,
-                drive_name,
-                &media_id.label,
-                notification_mode,
-            )?;
-            restore_file_chunk_map(worker.clone(), &mut drive, &store_map, file_chunk_map)?;
-        }
+        let jobs = media_file_chunk_map
+            .into_iter()
+            .map(|(media_uuid, file_chunk_map)| {
+                let media_id = inventory.lookup_media(&media_uuid).unwrap();
+                (media_id.label.clone(), file_chunk_map)
+            })
+            .collect();
+        drive_pool.run(jobs, |media, mut file_chunk_map| {
+            let mut drive = media.drive;
+            restore_file_chunk_map(
+                media.worker,
+                &mut drive,
+                &store_map,
+                &mut file_chunk_map,
+                &media.log_prefix,
+            )
+        })?;
 
         info!("Phase 3: copy snapshots from temp dir to datastores");
         let mut errors = false;
@@ -997,14 +1015,16 @@ fn snapshot_tmpdir(
     Ok(path)
 }
 
+#[allow(clippy::too_many_arguments)]
 fn restore_snapshots_to_tmpdir(
-    worker: Arc<WorkerTask>,
+    worker: Arc<dyn WorkerTaskContext>,
     store_map: &DataStoreMap,
     file_list: &[(u64, String)],
     mut drive: Box<dyn TapeDriver>,
     media_id: &MediaId,
     media_set_uuid: &Uuid,
     chunks_list: &mut HashMap<String, HashSet<[u8; 32]>>,
+    log_prefix: &str,
 ) -> Result<Vec<PathBuf>, Error> {
     let mut tmp_paths = Vec::new();
     match media_id.media_set_label {
@@ -1030,10 +1050,10 @@ fn restore_snapshots_to_tmpdir(
     for (file_num, snapshot_name) in file_list {
         let current_file_number = drive.current_file_number()?;
         if current_file_number != *file_num {
-            info!("was at file {current_file_number}, moving to {file_num}");
+            info!("{log_prefix}was at file {current_file_number}, moving to {file_num}");
             drive.move_to_file(*file_num)?;
             let current_file_number = drive.current_file_number()?;
-            info!("now at file {current_file_number}");
+            info!("{log_prefix}now at file {current_file_number}");
         }
         let mut reader = drive.read_next_file()?;
 
@@ -1062,7 +1082,9 @@ fn restore_snapshots_to_tmpdir(
                     bail!("unexpected snapshot name (expected {snapshot_name}, got {snapshot})");
                 }
 
-                info!("File {file_num}: snapshot archive {source_datastore}:{snapshot}",);
+                info!(
+                    "{log_prefix}File {file_num}: snapshot archive {source_datastore}:{snapshot}",
+                );
 
                 let mut decoder =
                     pxar::decoder::sync::Decoder::from_std(pxar::PxarVariant::Unified(reader))?;
@@ -1070,7 +1092,10 @@ fn restore_snapshots_to_tmpdir(
                 let target_datastore = match store_map.target_store(&source_datastore) {
                     Some(datastore) => datastore,
                     None => {
-                        warn!("could not find target datastore for {source_datastore}:{snapshot}",);
+                        warn!(
+                            "{log_prefix}could not find target datastore for \
+                            {source_datastore}:{snapshot}",
+                        );
                         continue;
                     }
                 };
@@ -1115,18 +1140,19 @@ fn restore_snapshots_to_tmpdir(
 }
 
 fn restore_file_chunk_map(
-    worker: Arc<WorkerTask>,
+    worker: Arc<dyn WorkerTaskContext>,
     drive: &mut Box<dyn TapeDriver>,
     store_map: &DataStoreMap,
     file_chunk_map: &mut BTreeMap<u64, HashSet<[u8; 32]>>,
+    log_prefix: &str,
 ) -> Result<(), Error> {
     for (nr, chunk_map) in file_chunk_map.iter_mut() {
         let current_file_number = drive.current_file_number()?;
         if current_file_number != *nr {
-            info!("was at file {current_file_number}, moving to {nr}");
+            info!("{log_prefix}was at file {current_file_number}, moving to {nr}");
             drive.move_to_file(*nr)?;
             let current_file_number = drive.current_file_number()?;
-            info!("now at file {current_file_number}");
+            info!("{log_prefix}now at file {current_file_number}");
         }
         let mut reader = drive.read_next_file()?;
         let header: MediaContentHeader = unsafe { reader.read_le_value()? };
@@ -1143,7 +1169,7 @@ fn restore_file_chunk_map(
 
                 let source_datastore = archive_header.store;
 
-                info!("File {nr}: chunk archive for datastore '{source_datastore}'",);
+                info!("{log_prefix}File {nr}: chunk archive for datastore '{source_datastore}'",);
 
                 let datastore = store_map.target_store(&source_datastore).ok_or_else(|| {
                     format_err!("unexpected chunk archive for store: {source_datastore}")
@@ -1154,8 +1180,9 @@ fn restore_file_chunk_map(
                     reader,
                     datastore.clone(),
                     chunk_map,
+                    log_prefix,
                 )?;
-                info!("restored {count} chunks");
+                info!("{log_prefix}restored {count} chunks");
             }
             _ => bail!("unexpected content magic {:?}", header.content_magic),
         }
@@ -1165,10 +1192,11 @@ fn restore_file_chunk_map(
 }
 
 fn restore_partial_chunk_archive<'a>(
-    worker: Arc<WorkerTask>,
+    worker: Arc<dyn WorkerTaskContext>,
     reader: Box<dyn 'a + TapeRead>,
     datastore: Arc<DataStore>,
     chunk_list: &mut HashSet<[u8; 32]>,
+    log_prefix: &str,
 ) -> Result<usize, Error> {
     let mut decoder = ChunkArchiveDecoder::new(reader);
 
@@ -1217,7 +1245,7 @@ fn restore_partial_chunk_archive<'a>(
     let elapsed = start_time.elapsed()?.as_secs_f64();
     let bytes = bytes.load(std::sync::atomic::Ordering::SeqCst) as f64;
     info!(
-        "restored {} ({:.2}/s)",
+        "{log_prefix}restored {} ({:.2}/s)",
         HumanByte::new_decimal(bytes),
         HumanByte::new_decimal(bytes / elapsed),
     );
@@ -1675,7 +1703,7 @@ fn restore_snapshot_archive<'a>(
 }
 
 fn try_restore_snapshot_archive<R: pxar::decoder::SeqRead>(
-    worker: Arc<WorkerTask>,
+    worker: Arc<dyn WorkerTaskContext>,
     decoder: &mut pxar::decoder::sync::Decoder<R>,
     snapshot_path: &Path,
 ) -> Result<BackupManifest, Error> {
-- 
2.53.0