cli: support the stateroot option also on switch

# Background

<hash> added the `stateroot` option to the `install` subcommand

# Issue

The `stateroot` option is not available on the `switch` subcommand

# Solution

Add the `stateroot` option to the `switch` subcommand

# Implementation

* If the stateroot is different than the current, we should allow using
  the same image as the currently booted one

* Stateroot has to be explicitly created (`init_osname` binding) if it
  doesn't exist. If it does, we still call `init_osname` and simply
  ignore the error (TODO: only ignore non-already-exists errors)

* Copy `/var` from the old stateroot to the new one. I'm doing `--reflink`
  but it's still very slow

* Must use the old stateroot to find the `merge_deployment` because
  otherwise we boot without the required kargs (it manifested as a
  missing `root=UUID=...` which caused the dracut rootfs-generator to
  silently fail to create `sysroot.mount` and so `ostree-prepare-root`
  failed due to empty `/sysroot`)

Signed-off-by: Omer Tuchfeld <[email protected]>

lib/src/cli.rs

@@ -25,6 +25,7 @@ use crate::deploy::RequiredHostSpec;
 use crate::lints;
 use crate::spec::Host;
 use crate::spec::ImageReference;
+use crate::task::Task;
 use crate::utils::sigpolicy_from_opts;
 
 include!(concat!(env!("OUT_DIR"), "/version.rs"));
@@ -99,6 +100,11 @@ pub(crate) struct SwitchOpts {
 
     /// Target image to use for the next boot.
     pub(crate) target: String,
+
+    /// Make the switch into a custom stateroot. If the stateroot doesn't exist, it will be created
+    /// and `/var` of the current stateroot will be copied (copy-on-write) into it.
+    #[clap(long)]
+    pub(crate) stateroot: Option<String>,
 }
 
 /// Options controlling rollback
@@ -628,7 +634,7 @@ async fn upgrade(opts: UpgradeOpts) -> Result<()> {
             println!("No update available.")
         } else {
             let osname = booted_deployment.osname();
-            crate::deploy::stage(sysroot, &osname, &fetched, &spec).await?;
+            crate::deploy::stage(sysroot, &osname, &osname, &fetched, &spec).await?;
             changed = true;
             if let Some(prev) = booted_image.as_ref() {
                 if let Some(fetched_manifest) = fetched.get_manifest(repo)? {
@@ -664,13 +670,13 @@ async fn switch(opts: SwitchOpts) -> Result<()> {
     );
     let target = ostree_container::OstreeImageReference { sigverify, imgref };
     let target = ImageReference::from(target);
+    let root = cap_std::fs::Dir::open_ambient_dir("/", cap_std::ambient_authority())?;
 
     // If we're doing an in-place mutation, we shortcut most of the rest of the work here
     if opts.mutate_in_place {
         let deployid = {
             // Clone to pass into helper thread
             let target = target.clone();
-            let root = cap_std::fs::Dir::open_ambient_dir("/", cap_std::ambient_authority())?;
             tokio::task::spawn_blocking(move || {
                 crate::deploy::switch_origin_inplace(&root, &target)
             })
@@ -687,18 +693,52 @@ async fn switch(opts: SwitchOpts) -> Result<()> {
     let (booted_deployment, _deployments, host) =
         crate::status::get_status_require_booted(sysroot)?;
 
+    let (old_stateroot, stateroot) = {
+        let booted_osname = booted_deployment.osname();
+        let stateroot = opts
+            .stateroot
+            .as_deref()
+            .unwrap_or_else(|| booted_osname.as_str());
+
+        (booted_osname.to_owned(), stateroot.to_owned())
+    };
+
     let new_spec = {
         let mut new_spec = host.spec.clone();
         new_spec.image = Some(target.clone());
         new_spec
     };
 
-    if new_spec == host.spec {
-        println!("Image specification is unchanged.");
+    if new_spec == host.spec && old_stateroot == stateroot {
+        // TODO: Should we really be confusing users with terms like "stateroot"?
+        println!(
+            "The currently running deployment in stateroot {stateroot} is already using this image"
+        );
         return Ok(());
     }
     let new_spec = RequiredHostSpec::from_spec(&new_spec)?;
 
+    if old_stateroot != stateroot {
+        let init_result = sysroot.init_osname(&stateroot, cancellable);
+        match init_result {
+            Ok(_) => {
+                Task::new("Copying /var to new stateroot", "cp")
+                    .args([
+                        "--recursive",
+                        "--reflink=auto",
+                        "--archive",
+                        format!("/sysroot/ostree/deploy/{old_stateroot}/var").as_str(),
+                        format!("/sysroot/ostree/deploy/{stateroot}/").as_str(),
+                    ])
+                    .run()?;
+            }
+            Err(err) => {
+                // TODO: Only ignore non already-exists errors
+                println!("Ignoring error creating new stateroot: {err}");
+            }
+        }
+    }
+
     let fetched = crate::deploy::pull(repo, &target, None, opts.quiet).await?;
 
     if !opts.retain {
@@ -712,8 +752,7 @@ async fn switch(opts: SwitchOpts) -> Result<()> {
         }
     }
 
-    let stateroot = booted_deployment.osname();
-    crate::deploy::stage(sysroot, &stateroot, &fetched, &new_spec).await?;
+    crate::deploy::stage(sysroot, &old_stateroot, &stateroot, &fetched, &new_spec).await?;
 
     if opts.apply {
         crate::reboot::reboot()?;
@@ -766,7 +805,7 @@ async fn edit(opts: EditOpts) -> Result<()> {
     // TODO gc old layers here
 
     let stateroot = booted_deployment.osname();
-    crate::deploy::stage(sysroot, &stateroot, &fetched, &new_spec).await?;
+    crate::deploy::stage(sysroot, &stateroot, &stateroot, &fetched, &new_spec).await?;
 
     Ok(())
 }

lib/src/deploy.rs

@@ -373,7 +373,6 @@ async fn deploy(
     image: &ImageState,
     origin: &glib::KeyFile,
 ) -> Result<Deployment> {
-    let stateroot = Some(stateroot);
     let mut opts = ostree::SysrootDeployTreeOpts::default();
     // Compute the kernel argument overrides. In practice today this API is always expecting
     // a merge deployment. The kargs code also always looks at the booted root (which
@@ -396,7 +395,7 @@ async fn deploy(
     let cancellable = gio::Cancellable::NONE;
     return sysroot
         .stage_tree_with_options(
-            stateroot,
+            Some(stateroot),
             image.ostree_commit.as_str(),
             Some(origin),
             merge_deployment,
@@ -422,11 +421,12 @@ fn origin_from_imageref(imgref: &ImageReference) -> Result<glib::KeyFile> {
 #[context("Staging")]
 pub(crate) async fn stage(
     sysroot: &Storage,
+    merge_stateroot: &str,
     stateroot: &str,
     image: &ImageState,
     spec: &RequiredHostSpec<'_>,
 ) -> Result<()> {
-    let merge_deployment = sysroot.merge_deployment(Some(stateroot));
+    let merge_deployment = sysroot.merge_deployment(Some(merge_stateroot));
     let origin = origin_from_imageref(spec.image)?;
     let deployment = crate::deploy::deploy(
         sysroot,