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

[Dominik Csapak <d.csapak@proxmox.com>]

most of it looks good to me, some comments inline

On 11/6/26 2:39 PM, Jakob-Tobias Winter wrote:
> 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();

i noticed that you did this here and...

> +        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();

... here instead of in the closure.

is there a special reason for that?
i think the code would look a bit better if we just give the drive pool 
the map directly, and pass the uuid in the closure instead.

each closure can do the lookup itself then (like it was before
just in a closure instead of the loop body)

> +        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> {