[RFC PATCH proxmox-backup 1/4] tape: drive: add a pool to use several drives for one task

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

Restoring a media set loads one tape after the other into a single
drive, even if the library has more drives that are idle. In
preparation for using them in parallel, add a pool of drives that runs
a list of jobs, each needing one media, on all of its drives.

With a single drive, the jobs run in the given order in the calling
thread, exactly like the existing loops that load one media after the
other.

With several drives, each drive runs in its own thread. A drive takes
the job for the media already loaded in it, else the next one whose
media is not loaded in another drive of the pool, as loading such a
media fails until that drive moves on. Changer operations are
serialized, so two drives do not pick the same free slot when
unloading a media. The first error stops all drives and is returned
once all of them finished.

As the lines of all drives are interleaved in the task log, their
messages start with the name of the drive, and the jobs get this
prefix for their own messages. With a single drive, the prefix is
empty, so the log stays as it is.

Signed-off-by: Jakob-Tobias Winter <tobias@stuttgart.de>
---
 src/tape/drive/mod.rs  |  41 ++++-
 src/tape/drive/pool.rs | 358 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 393 insertions(+), 6 deletions(-)
 create mode 100644 src/tape/drive/pool.rs

diff --git a/src/tape/drive/mod.rs b/src/tape/drive/mod.rs
index fa84fd54a..4df30be72 100644
--- a/src/tape/drive/mod.rs
+++ b/src/tape/drive/mod.rs
@@ -1,6 +1,7 @@
 //! Tape drivers
 
 use std::path::PathBuf;
+use std::sync::{Mutex, PoisonError};
 
 use anyhow::{Error, bail, format_err};
 use nix::fcntl::OFlag;
@@ -39,6 +40,9 @@ mod virtual_tape;
 mod lto;
 pub use lto::*;
 
+mod pool;
+pub use pool::*;
+
 /// Tape driver interface
 pub trait TapeDriver {
     /// Flush all data to the tape
@@ -378,11 +382,28 @@ pub fn request_and_load_media(
     drive: &str,
     label: &MediaLabel,
     notification_mode: &TapeNotificationMode,
+) -> Result<(Box<dyn TapeDriver>, MediaId), Error> {
+    request_and_load_media_locked(worker, config, drive, label, notification_mode, None, "")
+}
+
+/// Like [`request_and_load_media`], but holds `changer_lock` while the changer moves media and
+/// starts each log message with `log_prefix`.
+///
+/// This serializes the changer operations of several drives used by the same task, and tells
+/// their messages apart in the task log.
+pub fn request_and_load_media_locked(
+    worker: &dyn WorkerTaskContext,
+    config: &SectionConfigData,
+    drive: &str,
+    label: &MediaLabel,
+    notification_mode: &TapeNotificationMode,
+    changer_lock: Option<&Mutex<()>>,
+    log_prefix: &str,
 ) -> Result<(Box<dyn TapeDriver>, MediaId), Error> {
     let check_label = |handle: &mut dyn TapeDriver, uuid: &proxmox_uuid::Uuid| {
         if let Ok((Some(media_id), _)) = handle.read_label() {
             info!(
-                "found media label {} ({})",
+                "{log_prefix}found media label {} ({})",
                 media_id.label.label_text, media_id.label.uuid,
             );
 
@@ -422,7 +443,7 @@ pub fn request_and_load_media(
                                                         new: TapeRequestError|
                      -> Result<(), Error> {
                         if new != *old {
-                            info!("{new}");
+                            info!("{log_prefix}{new}");
                             let (device_type, device) = if let Some(changer) = changer {
                                 ("changer", changer.as_str())
                             } else {
@@ -430,7 +451,8 @@ pub fn request_and_load_media(
                             };
 
                             info!(
-                                "Please insert media '{label_text}' into {device_type} '{device}'"
+                                "{log_prefix}Please insert media '{label_text}' into {device_type} \
+                                '{device}'"
                             );
                             send_load_media_notification(
                                 notification_mode,
@@ -455,13 +477,20 @@ pub fn request_and_load_media(
                                 std::thread::sleep(std::time::Duration::from_millis(100));
                             }
                         } else if drive_config.changer.is_none() {
-                            info!("Checking for media '{label_text}' in drive '{drive}'");
+                            info!(
+                                "{log_prefix}Checking for media '{label_text}' in drive '{drive}'"
+                            );
                         } else {
-                            info!("trying to load media '{label_text}' into drive '{drive}'");
+                            info!(
+                                "{log_prefix}trying to load media '{label_text}' into drive \
+                                '{drive}'"
+                            );
                         }
 
                         if drive_config.changer.is_some() {
                             let mut changer = MtxMediaChanger::with_drive_config(&drive_config)?;
+                            let _changer_guard = changer_lock
+                                .map(|lock| lock.lock().unwrap_or_else(PoisonError::into_inner));
                             if let Err(err) = changer.load_media(&label_text) {
                                 update_and_log_request_error(
                                     &mut last_error,
@@ -485,7 +514,7 @@ pub fn request_and_load_media(
                         let request_error = match handle.read_label() {
                             Ok((Some(media_id), _)) if media_id.label.uuid == label.uuid => {
                                 info!(
-                                    "found media label {} ({})",
+                                    "{log_prefix}found media label {} ({})",
                                     media_id.label.label_text,
                                     media_id.label.uuid.to_string(),
                                 );
diff --git a/src/tape/drive/pool.rs b/src/tape/drive/pool.rs
new file mode 100644
index 000000000..3ee7a2984
--- /dev/null
+++ b/src/tape/drive/pool.rs
@@ -0,0 +1,358 @@
+//! Use several tape drives for one task
+
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
+
+use anyhow::{Error, bail};
+use tracing::{info, warn};
+
+use proxmox_log::LogContext;
+use proxmox_section_config::SectionConfigData;
+use proxmox_worker_task::WorkerTaskContext;
+
+use pbs_api_types::LtoTapeDrive;
+use pbs_tape::{ElementStatus, MtxStatus};
+
+use crate::tape::{
+    MediaId, TapeNotificationMode,
+    changer::{MediaChange, MtxMediaChanger},
+    file_formats::MediaLabel,
+};
+
+use super::{TapeDriver, request_and_load_media, request_and_load_media_locked};
+
+fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
+    mutex.lock().unwrap_or_else(PoisonError::into_inner)
+}
+
+/// Task context for the jobs of one drive, also stops when a job on another drive failed.
+struct DriveJobContext {
+    worker: Arc<dyn WorkerTaskContext>,
+    stop: Arc<AtomicBool>,
+}
+
+impl WorkerTaskContext for DriveJobContext {
+    fn abort_requested(&self) -> bool {
+        self.stop.load(Ordering::SeqCst) || self.worker.abort_requested()
+    }
+
+    fn shutdown_requested(&self) -> bool {
+        self.worker.shutdown_requested()
+    }
+}
+
+/// Stops the other drives if a drive thread panics.
+struct StopOnPanic<'a>(&'a AtomicBool);
+
+impl Drop for StopOnPanic<'_> {
+    fn drop(&mut self) {
+        if std::thread::panicking() {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+}
+
+/// A media loaded into a drive of a [`DrivePool`].
+pub struct LoadedMedia {
+    /// The task context the job has to use for abort checks.
+    pub worker: Arc<dyn WorkerTaskContext>,
+    pub drive: Box<dyn TapeDriver>,
+    pub media_id: MediaId,
+    /// The start of the job's log messages, names the drive if the pool has several.
+    pub log_prefix: String,
+}
+
+/// A set of drives that work on the media of one task in parallel.
+pub struct DrivePool<'a> {
+    worker: Arc<dyn WorkerTaskContext>,
+    config: &'a SectionConfigData,
+    drives: &'a [String],
+    notification_mode: &'a TapeNotificationMode,
+    /// Serializes the changer operations of all drives.
+    changer_lock: Mutex<()>,
+}
+
+impl<'a> DrivePool<'a> {
+    /// Create a pool of `drives`, which the caller has to lock.
+    pub fn new(
+        worker: Arc<dyn WorkerTaskContext>,
+        config: &'a SectionConfigData,
+        drives: &'a [String],
+        notification_mode: &'a TapeNotificationMode,
+    ) -> Self {
+        Self {
+            worker,
+            config,
+            drives,
+            notification_mode,
+            changer_lock: Mutex::new(()),
+        }
+    }
+
+    /// Load the media of each job into a drive and run `func` for it.
+    ///
+    /// With a single drive, the jobs run in the given order in the calling thread. With several
+    /// drives, each drive runs in its own thread and takes the next job whose media is not loaded
+    /// in another drive of the pool, preferring the media loaded in itself. The first error stops
+    /// the jobs on all drives and is returned once all of them finished.
+    ///
+    /// The results are returned in the order the jobs finished.
+    pub fn run<T, R, F>(&self, jobs: Vec<(MediaLabel, T)>, func: F) -> Result<Vec<R>, Error>
+    where
+        T: Send,
+        R: Send,
+        F: Fn(LoadedMedia, T) -> Result<R, Error> + Sync,
+    {
+        if let [drive_name] = self.drives {
+            let mut results = Vec::with_capacity(jobs.len());
+            for (label, job) in jobs {
+                let (drive, media_id) = request_and_load_media(
+                    &*self.worker,
+                    self.config,
+                    drive_name,
+                    &label,
+                    self.notification_mode,
+                )?;
+                let media = LoadedMedia {
+                    worker: Arc::clone(&self.worker),
+                    drive,
+                    media_id,
+                    log_prefix: String::new(),
+                };
+                results.push(func(media, job)?);
+            }
+            return Ok(results);
+        }
+
+        let queue = Mutex::new(VecDeque::from(jobs));
+        let results = Mutex::new(Vec::new());
+        let first_error = Mutex::new(None);
+        let stop = Arc::new(AtomicBool::new(false));
+        let log_context = LogContext::current();
+
+        std::thread::scope(|scope| {
+            for drive_name in self.drives {
+                let worker: Arc<dyn WorkerTaskContext> = Arc::new(DriveJobContext {
+                    worker: Arc::clone(&self.worker),
+                    stop: Arc::clone(&stop),
+                });
+                let (queue, results, first_error, stop, func) =
+                    (&queue, &results, &first_error, &stop, &func);
+                let log_context = log_context.clone();
+
+                scope.spawn(move || {
+                    let run = || {
+                        let _stop_on_panic = StopOnPanic(stop);
+                        let res = self.drive_loop(drive_name, worker, queue, results, func);
+                        if let Err(err) = res {
+                            let mut first_error = lock(first_error);
+                            if first_error.is_none() {
+                                if !stop.load(Ordering::SeqCst) && !self.worker.abort_requested() {
+                                    warn!(
+                                        "drive '{drive_name}': failed, stopping all drives - {err}"
+                                    );
+                                }
+                                *first_error = Some(err);
+                            }
+                            stop.store(true, Ordering::SeqCst);
+                        }
+                    };
+                    match log_context {
+                        Some(log_context) => log_context.sync_scope(run),
+                        None => run(),
+                    }
+                });
+            }
+        });
+
+        if let Some(err) = lock(&first_error).take() {
+            return Err(err);
+        }
+
+        let left = lock(&queue);
+        if !left.is_empty() {
+            // only possible without drives, a drive always takes the media loaded in itself
+            let labels: Vec<&str> = left
+                .iter()
+                .map(|(label, _)| label.label_text.as_str())
+                .collect();
+            bail!("unable to process media {}", labels.join(", "));
+        }
+
+        Ok(results.into_inner().unwrap_or_else(PoisonError::into_inner))
+    }
+
+    fn drive_loop<T, R, F>(
+        &self,
+        drive_name: &str,
+        worker: Arc<dyn WorkerTaskContext>,
+        queue: &Mutex<VecDeque<(MediaLabel, T)>>,
+        results: &Mutex<Vec<R>>,
+        func: &F,
+    ) -> Result<(), Error>
+    where
+        F: Fn(LoadedMedia, T) -> Result<R, Error>,
+    {
+        // the lines of all drives are interleaved in the task log
+        let log_prefix = format!("drive '{drive_name}': ");
+
+        loop {
+            worker.check_abort()?;
+
+            let (label, job) = {
+                let mut queue = lock(queue);
+                if queue.is_empty() {
+                    return Ok(());
+                }
+
+                let (own, busy) = self.loaded_media(drive_name, &log_prefix);
+                match next_job(&queue, own.as_deref(), &busy) {
+                    Some(pos) => queue.remove(pos).expect("position is valid"),
+                    None => return Ok(()),
+                }
+            };
+
+            info!("{log_prefix}using media '{}'", label.label_text);
+
+            let (drive, media_id) = request_and_load_media_locked(
+                &*worker,
+                self.config,
+                drive_name,
+                &label,
+                self.notification_mode,
+                Some(&self.changer_lock),
+                &log_prefix,
+            )?;
+            let media = LoadedMedia {
+                worker: Arc::clone(&worker),
+                drive,
+                media_id,
+                log_prefix: log_prefix.clone(),
+            };
+            let result = func(media, job)?;
+            lock(results).push(result);
+
+            info!("{log_prefix}done with media '{}'", label.label_text);
+        }
+    }
+
+    /// The label texts of the media loaded in `drive_name` and in the other drives of this pool.
+    ///
+    /// Only drives in a changer can report this, errors are ignored (loading a media that is in
+    /// another drive then fails and is retried by [`request_and_load_media_locked`]).
+    fn loaded_media(
+        &self,
+        drive_name: &str,
+        log_prefix: &str,
+    ) -> (Option<String>, HashSet<String>) {
+        let mut own = None;
+        let mut others = HashSet::new();
+        let mut status_list: HashMap<String, Option<MtxStatus>> = HashMap::new();
+
+        let _changer_guard = lock(&self.changer_lock);
+
+        for name in self.drives {
+            let Ok(drive) = self.config.lookup::<LtoTapeDrive>("lto", name) else {
+                continue;
+            };
+            let Some(changer) = drive.changer.clone() else {
+                continue;
+            };
+
+            let status = status_list.entry(changer).or_insert_with_key(|changer| {
+                match MtxMediaChanger::with_drive_config(&drive).and_then(|mut c| c.status()) {
+                    Ok(status) => Some(status),
+                    Err(err) => {
+                        warn!("{log_prefix}unable to get status of changer '{changer}' - {err}");
+                        None
+                    }
+                }
+            });
+
+            let drive_status = status.as_ref().and_then(|status| {
+                status
+                    .drives
+                    .get(drive.changer_drivenum.unwrap_or(0) as usize)
+            });
+            if let Some(drive_status) = drive_status {
+                if let ElementStatus::VolumeTag(ref tag) = drive_status.status {
+                    if name == drive_name {
+                        own = Some(tag.clone());
+                    } else {
+                        others.insert(tag.clone());
+                    }
+                }
+            }
+        }
+
+        (own, others)
+    }
+}
+
+/// Choose the next job for a drive: the media loaded in it, else the first one in order that is
+/// not loaded in another drive (loading that would fail, the other drive will take it).
+fn next_job<T>(
+    queue: &VecDeque<(MediaLabel, T)>,
+    own: Option<&str>,
+    busy: &HashSet<String>,
+) -> Option<usize> {
+    if let Some(own) = own {
+        if let Some(pos) = queue.iter().position(|(label, _)| label.label_text == own) {
+            return Some(pos);
+        }
+    }
+    queue
+        .iter()
+        .position(|(label, _)| !busy.contains(&label.label_text))
+}
+
+#[cfg(test)]
+mod test {
+    use std::collections::{HashSet, VecDeque};
+
+    use proxmox_uuid::Uuid;
+
+    use super::next_job;
+    use crate::tape::file_formats::MediaLabel;
+
+    fn queue(labels: &[&str]) -> VecDeque<(MediaLabel, ())> {
+        labels
+            .iter()
+            .map(|label_text| {
+                let label = MediaLabel {
+                    label_text: label_text.to_string(),
+                    uuid: Uuid::generate(),
+                    ctime: 0,
+                    pool: None,
+                };
+                (label, ())
+            })
+            .collect()
+    }
+
+    #[test]
+    fn next_job_in_order() {
+        let queue = queue(&["A", "B", "C"]);
+        assert_eq!(next_job(&queue, None, &HashSet::new()), Some(0));
+        // the media in this drive is not in the queue
+        assert_eq!(next_job(&queue, Some("X"), &HashSet::new()), Some(0));
+    }
+
+    #[test]
+    fn next_job_prefers_own_media() {
+        let queue = queue(&["A", "B", "C"]);
+        assert_eq!(next_job(&queue, Some("C"), &HashSet::new()), Some(2));
+    }
+
+    #[test]
+    fn next_job_skips_media_in_other_drives() {
+        let queue = queue(&["A", "B", "C"]);
+        let busy = HashSet::from(["A".to_string(), "B".to_string()]);
+        assert_eq!(next_job(&queue, None, &busy), Some(2));
+
+        let busy = HashSet::from(["A".to_string(), "B".to_string(), "C".to_string()]);
+        assert_eq!(next_job(&queue, None, &busy), None);
+        assert_eq!(next_job(&queue, Some("B"), &busy), Some(1));
+    }
+}
-- 
2.53.0