From 33d81809a7e53f70cd51730d5b20497ef1c7072a Mon Sep 17 00:00:00 2001
From: mulhern <amulhern@redhat.com>
Date: Tue, 6 Aug 2024 22:29:28 -0400
Subject: [PATCH] Roughly abstract the calculations that differ

Signed-off-by: mulhern <amulhern@redhat.com>
---
 src/bin/utils/predict_usage.rs | 123 +++++++++++++++++++++++----------
 1 file changed, 85 insertions(+), 38 deletions(-)

diff --git a/src/bin/utils/predict_usage.rs b/src/bin/utils/predict_usage.rs
index 492319b453..21a3a7515b 100644
--- a/src/bin/utils/predict_usage.rs
+++ b/src/bin/utils/predict_usage.rs
@@ -156,6 +156,85 @@ pub fn predict_filesystem_usage(
     Ok(())
 }
 
+struct V1usage {
+    crypt_metadata_alloc: Sectors,
+    stratis_metadata_alloc: Sectors,
+}
+
+impl V1usage {
+    fn collect_usage_params(crypt_meta_alloc: Sectors) -> Self {
+        V1usage {
+            crypt_metadata_alloc: crypt_meta_alloc,
+            stratis_metadata_alloc: BDA::default().extended_size().sectors(),
+        }
+    }
+
+    fn predict_pool_usage(
+        &self,
+        encrypted: bool,
+        device_sizes: Vec<Sectors>,
+    ) -> Result<Sectors, Box<dyn Error>> {
+        let crypt_metadata_size = if encrypted {
+            self.crypt_metadata_alloc
+        } else {
+            Sectors(0)
+        };
+
+        let stratis_avail_sizes = device_sizes
+            .iter()
+            .map(|&s| {
+                (*s).checked_sub(*crypt_metadata_size)
+                    .and_then(|r| r.checked_sub(*self.stratis_metadata_alloc))
+                    .map(Sectors)
+                    .ok_or_else(|| {
+                        Box::new(PredictError(
+                            "Some device sizes too small for encryption metadata.".into(),
+                        ))
+                    })
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        Ok(stratis_avail_sizes.iter().cloned().sum::<Sectors>())
+    }
+}
+
+struct V2usage {
+    crypt_metadata_alloc: Sectors,
+    integrity_metadata_alloc: Sectors,
+    raid_metadata_alloc: Sectors,
+    stratis_metadata_alloc: Sectors,
+}
+
+impl V2usage {
+    fn collect_usage_params(crypt_meta_alloc: Sectors) -> Self {
+        V2usage {
+            crypt_metadata_alloc: crypt_meta_alloc,
+            integrity_metadata_alloc: Sectors(0),
+            raid_metadata_alloc: Sectors(0),
+            stratis_metadata_alloc: BDA::default().extended_size().sectors(),
+        }
+    }
+
+    fn predict_pool_usage(&self, device_sizes: Vec<Sectors>) -> Result<Sectors, Box<dyn Error>> {
+        let stratis_avail_sizes = device_sizes
+            .iter()
+            .map(|&s| {
+                (*s).checked_sub(*self.integrity_metadata_alloc)
+                    .and_then(|r| r.checked_sub(*self.raid_metadata_alloc))
+                    .and_then(|r| r.checked_sub(*self.stratis_metadata_alloc))
+                    .map(Sectors)
+                    .ok_or_else(|| {
+                        Box::new(PredictError(
+                            "Some device sizes too small for DM metadata.".into(),
+                        ))
+                    })
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        Ok(stratis_avail_sizes.iter().cloned().sum::<Sectors>() - self.crypt_metadata_alloc)
+    }
+}
+
 // Predict usage for a newly created pool given information about whether
 // or not the pool is encrypted, a list of device sizes, and an optional list
 // of filesystem sizes.
@@ -170,52 +249,20 @@ pub fn predict_pool_usage(
         .map(|sizes| get_filesystem_prediction(overprovisioned, sizes))
         .transpose()?;
 
-    let crypt_metadata_size = if encrypted {
-        crypt_metadata_size()
-    } else {
-        Bytes(0)
-    };
-
+    let crypt_metadata_size = crypt_metadata_size();
     let crypt_metadata_size_sectors = crypt_metadata_size.sectors();
-
     // verify that crypt metadata size is divisible by sector size
     assert_eq!(crypt_metadata_size_sectors.bytes(), crypt_metadata_size);
 
     let device_sizes = device_sizes.iter().map(|s| s.sectors()).collect::<Vec<_>>();
     let total_size: Sectors = device_sizes.iter().cloned().sum();
 
-    let stratis_device_sizes = if metadata_version == 1u8 {
-        device_sizes
-            .iter()
-            .map(|&s| {
-                (*s).checked_sub(*crypt_metadata_size_sectors)
-                    .map(Sectors)
-                    .ok_or_else(|| {
-                        Box::new(PredictError(
-                            "Some device sizes too small for encryption metadata.".into(),
-                        ))
-                    })
-            })
-            .collect::<Result<Vec<_>, _>>()?
+    let non_metadata_size = if metadata_version == 1 {
+        V1usage::collect_usage_params(crypt_metadata_size_sectors)
+            .predict_pool_usage(encrypted, device_sizes)
     } else {
-        device_sizes
-    };
-
-    let stratis_metadata_size = BDA::default().extended_size().sectors();
-    let stratis_avail_sizes = stratis_device_sizes
-        .iter()
-        .map(|&s| {
-            (*s).checked_sub(*stratis_metadata_size)
-                .map(Sectors)
-                .ok_or_else(|| {
-                    Box::new(PredictError(
-                        "Some device sizes too small for Stratis metadata.".into(),
-                    ))
-                })
-        })
-        .collect::<Result<Vec<_>, _>>()?;
-
-    let non_metadata_size: Sectors = stratis_avail_sizes.iter().cloned().sum();
+        V2usage::collect_usage_params(crypt_metadata_size_sectors).predict_pool_usage(device_sizes)
+    }?;
 
     let size_params = ThinPoolSizeParams::new(non_metadata_size)?;
     let total_non_data = 2usize * size_params.meta_size() + size_params.mdv_size();