Implement resizing of zpool

rust/Cargo.lock

@@ -270,6 +270,18 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "disko-zpool-expand"
+version = "1.0.0"
+dependencies = [
+ "anyhow",
+ "clap",
+ "common",
+ "hakari",
+ "log",
+ "serde",
+]
+
 [[package]]
 name = "displaydoc"
 version = "0.2.5"

rust/lib/common/src/proc.rs

@@ -308,21 +308,67 @@ impl Command {
         cmd.status()
             .with_context(|| format!("Could not spawn command: {}", command_to_string(self)))
     }
+
+    pub fn spawn_into_parts(&mut self) -> anyhow::Result<(String, String, ExitStatus)> {
+        let mut cmd = self.as_command();
+        let mut child = cmd
+            .stderr(self.stderr.as_std_stdio())
+            .stdin(self.stdin.as_std_stdio())
+            .stdout(Stdio::piped())
+            .spawn()
+            .with_context(|| format!("Could not spawn command: {}", command_to_string(self)))?;
+        let mut stdin = Stdin::Null;
+        std::mem::swap(&mut self.stdin, &mut stdin);
+        let join_handle = if let Some(data) = stdin.into_data() {
+            let mut stdin_pipe = child.stdin.take().expect("Child has no stdin");
+            Some(std::thread::spawn(move || {
+                stdin_pipe
+                    .write_all(data.as_slice())
+                    .expect("Could not write to child");
+            }))
+        } else {
+            None
+        };
+        let output: Result<(Vec<u8>, Vec<u8>, ExitStatus), anyhow::Error> =
+            wait_with_output_into_parts(child, || command_to_string(self));
+        if let Some(join_handle) = join_handle {
+            join_handle
+                .join()
+                .map_err(|e| anyhow::format_err!("Thread sending stdin panicked: {e:?}"))?;
+        }
+        let (stdout, stderr, status) = output?;
+        let stdout =
+            String::from_utf8(stdout).context("Could not read stdout of command as UTF-8")?;
+        let stderr =
+            String::from_utf8(stderr).context("Could not read stderr of command as UTF-8")?;
+        Ok((stdout, stderr, status))
+    }
 }
 
-fn wait_with_output(child: Child, cmd_str: impl Fn() -> String) -> anyhow::Result<Vec<u8>> {
-    let output = child
-        .wait_with_output()
-        .with_context(|| format!("Could not wait for output of commnad: {}", cmd_str()))?;
-    if !output.status.success() {
+fn wait_with_output(child: Child, cmd_str: impl Fn() -> String + Clone) -> anyhow::Result<Vec<u8>> {
+    let (stdout, stderr, status) = wait_with_output_into_parts(child, cmd_str.clone())?;
+    if !status.success() {
         return Err(anyhow::format_err!(
             "Command {}, exited unexpectedly: {:?}. With stderr: {}",
             cmd_str(),
-            output.status,
-            String::from_utf8_lossy(&output.stderr),
+            status,
+            String::from_utf8_lossy(&stderr),
         ));
     }
-    Ok(output.stdout)
+    Ok(stdout)
+}
+
+fn wait_with_output_into_parts(
+    child: Child,
+    cmd_str: impl Fn() -> String,
+) -> anyhow::Result<(Vec<u8>, Vec<u8>, ExitStatus)> {
+    let output = child
+        .wait_with_output()
+        .with_context(|| format!("Could not wait for output of commnad: {}", cmd_str()))?;
+    let status = output.status;
+    let stdout = output.stdout;
+    let stderr = output.stderr;
+    Ok((stdout, stderr, status))
 }
 
 #[cfg(test)]

rust/program/disko-zpool-expand/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "disko-zpool-expand"
+edition = "2024"
+version = "1.0.0"
+metadata.crane.name = "disko-zpool-expand"
+
+[dependencies]
+anyhow = { workspace = true }
+clap = { workspace = true }
+common = { path = "../../lib/common" }
+log = { workspace = true }
+serde = { workspace = true }
+hakari = { version = "0.1", path = "../../lib/hakari" }

rust/program/disko-zpool-expand/src/main.rs

@@ -0,0 +1,117 @@
+use serde::Deserialize;
+use std::{collections::BTreeMap, path::PathBuf};
+
+use anyhow::Context as _;
+use clap::{Parser, Subcommand};
+
+fn main() {
+    common::entrypoint(program);
+}
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+pub struct Args {
+    #[command(subcommand)]
+    pub command: Commands,
+}
+
+#[derive(Debug, Subcommand)]
+pub enum Commands {
+    /// Expands the partitions based on a zpool and brings the pool up to the new size.
+    #[command(name = "expand-zpool")]
+    ExpandZpool(ExpandZpool),
+}
+
+#[derive(Debug, Clone, clap::Args)]
+pub struct ExpandZpool {
+    /// Name of the pool to expand
+    pool_name: String,
+}
+
+fn program() -> anyhow::Result<()> {
+    let args = Args::parse();
+    match args.command {
+        Commands::ExpandZpool(pool) => expand_zpool(pool),
+    }
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatus {
+    pools: BTreeMap<String, ZpoolStatusPool>,
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatusPool {
+    state: Option<ZpoolState>,
+    vdevs: BTreeMap<String, ZpoolStatusVdev>,
+}
+
+#[derive(Clone, Copy, Deserialize, PartialEq)]
+enum ZpoolState {
+    #[serde(rename = "ONLINE")]
+    Online,
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatusVdev {
+    vdevs: BTreeMap<String, ZpoolStatusVdevVdev>,
+}
+#[derive(Deserialize)]
+struct ZpoolStatusVdevVdev {
+    path: PathBuf,
+}
+
+fn expand_zpool(p: ExpandZpool) -> anyhow::Result<()> {
+    let mut proc = common::proc::Command::new("zpool");
+    proc.args(["status", "--json", &p.pool_name]);
+    let result: ZpoolStatus = proc
+        .try_spawn_to_json()
+        .context("Could not get zpool status")?;
+
+    let pool = result
+        .pools
+        .get(&p.pool_name)
+        .context("Could not find requested pool in status output")?;
+
+    if !pool
+        .state
+        .as_ref()
+        .is_some_and(|st| *st == ZpoolState::Online)
+    {
+        return Err(anyhow::format_err!("Zpool {} is not online", p.pool_name));
+    }
+
+    for vdev in pool.vdevs.values() {
+        for vdev in vdev.vdevs.values() {
+            let partition_dev = vdev.path.display().to_string();
+            let Some(dev) = partition_dev.strip_suffix("-part1") else {
+                return Err(anyhow::format_err!(
+                    "Expected vdev path {} to end with -part1",
+                    vdev.path.display()
+                ));
+            };
+            let mut proc = common::proc::Command::new("growpart");
+            proc.args([dev, "1"]);
+            let (stdout, _stderr, status) = proc.spawn_into_parts()?;
+            if !status.success() && !stdout.starts_with("NOCHANGE: ") {
+                return Err(anyhow::format_err!(
+                    "Could not resize partitin for {}, err: {stdout}",
+                    vdev.path.display()
+                ));
+            }
+            // let name = partition_dev
+            //     .split("/")
+            //     .last()
+            //     .expect("Should always have at least one element");
+            let mut proc = common::proc::Command::new("zpool");
+            proc.args(["online", "-e", &p.pool_name, &partition_dev]);
+            proc.try_spawn_to_string().with_context(|| {
+                format!(
+                    "Could not bring zpool {} online with expand flag",
+                    p.pool_name
+                )
+            })?;
+        }
+    }
+    Ok(())
+}