[pbs-devel] [PATCH proxmox-backup v3 2/3] disks: use builder pattern for querying disk usage

[Hannes Laimer <h.laimer@proxmox.com>]

Signed-off-by: Hannes Laimer <h.laimer@proxmox.com>
---
v2, thanks @Wolfgang Bumiller <w.bumiller@proxmox.com>
 * use just one struct for querying
 * remove not needed lifetimes

the two 'find' functions arguments are for filtering the data and the builder
functions are for defining what data the (maybe filtered) resulting
elements should contain. 

src/api2/node/disks/directory.rs |  6 ++--
 src/api2/node/disks/mod.rs       | 12 +++++---
 src/api2/node/disks/zfs.rs       |  2 +-
 src/tools/disks/mod.rs           | 53 +++++++++++++++++++++++---------
 4 files changed, 51 insertions(+), 22 deletions(-)

diff --git a/src/api2/node/disks/directory.rs b/src/api2/node/disks/directory.rs
index 123a8d7b..cada95cd 100644
--- a/src/api2/node/disks/directory.rs
+++ b/src/api2/node/disks/directory.rs
@@ -13,8 +13,8 @@ use pbs_api_types::{
 };
 
 use crate::tools::disks::{
-    create_file_system, create_single_linux_partition, get_disk_usage_info, get_fs_uuid,
-    DiskManage, DiskUsageType, FileSystemType,
+    create_file_system, create_single_linux_partition, get_fs_uuid, DiskManage, DiskUsageQuery,
+    DiskUsageType, FileSystemType,
 };
 use crate::tools::systemd::{self, types::*};
 
@@ -147,7 +147,7 @@ pub fn create_datastore_disk(
 
     let auth_id = rpcenv.get_auth_id().unwrap();
 
-    let info = get_disk_usage_info(&disk, true, false)?;
+    let info = DiskUsageQuery::new().smart(false).find(&disk)?;
 
     if info.used != DiskUsageType::Unused {
         bail!("disk '{}' is already in use.", disk);
diff --git a/src/api2/node/disks/mod.rs b/src/api2/node/disks/mod.rs
index 478829fb..c16be46a 100644
--- a/src/api2/node/disks/mod.rs
+++ b/src/api2/node/disks/mod.rs
@@ -12,8 +12,8 @@ use pbs_api_types::{
 };
 
 use crate::tools::disks::{
-    get_disk_usage_info, get_disks, get_smart_data, inititialize_gpt_disk, DiskManage,
-    DiskUsageInfo, DiskUsageType, SmartData,
+    get_smart_data, inititialize_gpt_disk, DiskManage, DiskUsageInfo, DiskUsageQuery,
+    DiskUsageType, SmartData,
 };
 use proxmox_rest_server::WorkerTask;
 
@@ -64,7 +64,11 @@ pub fn list_disks(
 ) -> Result<Vec<DiskUsageInfo>, Error> {
     let mut list = Vec::new();
 
-    for (_, info) in get_disks(None, skipsmart, include_partitions)? {
+    for (_, info) in DiskUsageQuery::new()
+        .smart(!skipsmart)
+        .partitions(include_partitions)
+        .query()?
+    {
         if let Some(ref usage_type) = usage_type {
             if info.used == *usage_type {
                 list.push(info);
@@ -147,7 +151,7 @@ pub fn initialize_disk(
 
     let auth_id = rpcenv.get_auth_id().unwrap();
 
-    let info = get_disk_usage_info(&disk, true, false)?;
+    let info = DiskUsageQuery::new().find(&disk)?;
 
     if info.used != DiskUsageType::Unused {
         bail!("disk '{}' is already in use.", disk);
diff --git a/src/api2/node/disks/zfs.rs b/src/api2/node/disks/zfs.rs
index 5cb23e70..3efc8a05 100644
--- a/src/api2/node/disks/zfs.rs
+++ b/src/api2/node/disks/zfs.rs
@@ -174,7 +174,7 @@ pub fn create_zpool(
         .map(|v| v.as_str().unwrap().to_string())
         .collect();
 
-    let disk_map = crate::tools::disks::get_disks(None, true, false)?;
+    let disk_map = crate::tools::disks::DiskUsageQuery::new().query()?;
     for disk in devices.iter() {
         match disk_map.get(disk) {
             Some(info) => {
diff --git a/src/tools/disks/mod.rs b/src/tools/disks/mod.rs
index 61e0f17a..4b2d638b 100644
--- a/src/tools/disks/mod.rs
+++ b/src/tools/disks/mod.rs
@@ -763,19 +763,44 @@ fn scan_partitions(
     Ok(used)
 }
 
-/// Get disk usage information for a single disk
-pub fn get_disk_usage_info(
-    disk: &str,
-    no_smart: bool,
-    include_partitions: bool,
-) -> Result<DiskUsageInfo, Error> {
-    let mut filter = Vec::new();
-    filter.push(disk.to_string());
-    let mut map = get_disks(Some(filter), no_smart, include_partitions)?;
-    if let Some(info) = map.remove(disk) {
-        Ok(info)
-    } else {
-        bail!("failed to get disk usage info - internal error"); // should not happen
+pub struct DiskUsageQuery {
+    smart: bool,
+    partitions: bool,
+}
+
+impl DiskUsageQuery {
+    pub fn new() -> DiskUsageQuery {
+        DiskUsageQuery {
+            smart: true,
+            partitions: false,
+        }
+    }
+
+    pub fn smart<'a>(&'a mut self, smart: bool) -> &'a mut DiskUsageQuery {
+        self.smart = smart;
+        self
+    }
+
+    pub fn partitions<'a>(&'a mut self, partitions: bool) -> &'a mut DiskUsageQuery {
+        self.partitions = partitions;
+        self
+    }
+
+    pub fn query(&self) -> Result<HashMap<String, DiskUsageInfo>, Error> {
+        get_disks(None, !self.smart, self.partitions)
+    }
+
+    pub fn find(&self, disk: &str) -> Result<DiskUsageInfo, Error> {
+        let mut map = get_disks(Some(vec![disk.to_string()]), !self.smart, self.partitions)?;
+        if let Some(info) = map.remove(disk) {
+            Ok(info)
+        } else {
+            bail!("failed to get disk usage info - internal error"); // should not happen
+        }
+    }
+
+    pub fn find_all(&self, disks: Vec<String>) -> Result<HashMap<String, DiskUsageInfo>, Error> {
+        get_disks(Some(disks), !self.smart, self.partitions)
     }
 }
 
@@ -838,7 +863,7 @@ fn get_partitions_info(
 }
 
 /// Get disk usage information for multiple disks
-pub fn get_disks(
+fn get_disks(
     // filter - list of device names (without leading /dev)
     disks: Option<Vec<String>>,
     // do no include data from smartctl
-- 
2.30.2