[pbs-devel] [PATCH proxmox-backup v2 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>
---
 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           | 86 ++++++++++++++++++++++++++------
 4 files changed, 84 insertions(+), 22 deletions(-)

diff --git a/src/api2/node/disks/directory.rs b/src/api2/node/disks/directory.rs
index 123a8d7b..cb2799fa 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(&disk).smart(false).query()?;
 
     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..2bd3f859 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, DisksUsageQuery, 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 DisksUsageQuery::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(&disk).query()?;
 
     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..a2d6ff0f 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::DisksUsageQuery::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 ea4c687a..a1436fc3 100644
--- a/src/tools/disks/mod.rs
+++ b/src/tools/disks/mod.rs
@@ -781,19 +781,77 @@ 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 DisksUsageQuery {
+    disks: Option<Vec<String>>,
+    smart: bool,
+    partitions: bool,
+}
+
+impl DisksUsageQuery {
+    pub fn new() -> DisksUsageQuery {
+        DisksUsageQuery {
+            disks: None,
+            smart: true,
+            partitions: false,
+        }
+    }
+
+    pub fn disks<'a>(&'a mut self, disks: Vec<String>) -> &'a mut DisksUsageQuery {
+        self.disks = Some(disks);
+        self
+    }
+
+    pub fn smart<'a>(&'a mut self, smart: bool) -> &'a mut DisksUsageQuery {
+        self.smart = smart;
+        self
+    }
+
+    pub fn partitions<'a>(&'a mut self, partitions: bool) -> &'a mut DisksUsageQuery {
+        self.partitions = partitions;
+        self
+    }
+
+    pub fn query(&self) -> Result<HashMap<String, DiskUsageInfo>, Error> {
+        get_disks(self.disks.clone(), !self.smart, self.partitions)
+    }
+}
+
+pub struct DiskUsageQuery {
+    disk: String,
+    smart: bool,
+    partitions: bool,
+}
+
+impl DiskUsageQuery {
+    pub fn new(disk: &str) -> DiskUsageQuery {
+        DiskUsageQuery {
+            disk: String::from(disk),
+            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<DiskUsageInfo, Error> {
+        let mut map = DisksUsageQuery::new()
+            .disks(vec![self.disk.clone()])
+            .smart(self.smart)
+            .partitions(self.partitions)
+            .query()?;
+        if let Some(info) = map.remove(&self.disk) {
+            Ok(info)
+        } else {
+            bail!("failed to get disk usage info - internal error"); // should not happen
+        }
     }
 }
 
@@ -856,7 +914,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