Add ZFS subdirs

lib/store_spec.ml

@@ -5,7 +5,7 @@ open Sexplib.Conv
 
 type t = [
   | `Btrfs of string  (* Path *)
-  | `Zfs of string    (* Path with pool at end *)
+  | `Zfs of (string option * string * bool)    (* Path with pool at end, optional sub-directory path and whether the pool should be in the file paths *)
   | `Rsync of (string * Rsync_store.mode)  (* Path for the root of the store *)
   | `Xfs of string    (* Path *)
   | `Docker of string (* Path *)
@@ -15,15 +15,21 @@ let is_absolute path = not (Filename.is_relative path)
 
 let of_string s =
   match Astring.String.cut s ~sep:":" with
-  | Some ("zfs", pool) -> Ok (`Zfs pool)
+  | Some ("zfs", pool) -> (
+    (* zfs:<prefix/pool>:<|subdir|> *)
+    match Astring.String.cut pool ~sep:":" with
+    | Some (pool, subdir) -> Ok (`Zfs (Some subdir, pool))
+    | None -> Ok (`Zfs (None, pool))
+  )
   | Some ("btrfs", path) when is_absolute path -> Ok (`Btrfs path)
   | Some ("rsync", path) when is_absolute path -> Ok (`Rsync path)
   | Some ("xfs", path) when is_absolute path -> Ok (`Xfs path)
   | Some ("docker", path) -> Ok (`Docker path)
   | _ -> Error (`Msg "Store must start with zfs:, btrfs:/, rsync:/ or xfs:/")
 
 let pp f = function
-  | `Zfs path -> Fmt.pf f "zfs:%s" path
+  | `Zfs (None, path) -> Fmt.pf f "zfs:%s" path
+  | `Zfs (Some subdir, path) -> Fmt.pf f "zfs:%s:%s" path subdir
   | `Btrfs path -> Fmt.pf f "btrfs:%s" path
   | `Rsync path -> Fmt.pf f "rsync:%s" path
   | `Xfs path -> Fmt.pf f "xfs:%s" path
@@ -35,8 +41,8 @@ let to_store = function
   | `Btrfs path ->
     `Native, Btrfs_store.create path >|= fun store ->
     Store ((module Btrfs_store), store)
-  | `Zfs path ->
-    `Native, Zfs_store.create ~path >|= fun store ->
+  | `Zfs (subdir, path, path_without_pool) ->
+    `Native, Zfs_store.create ~path_with_pool:(not path_without_pool) ?subdir ~path >|= fun store ->
     Store ((module Zfs_store), store)
   | `Rsync (path, rsync_mode) ->
     `Native, Rsync_store.create ~path ~mode:rsync_mode () >|= fun store ->
@@ -74,19 +80,25 @@ let rsync_mode_opt =
       ~docs:"RSYNC STORE"
       ["rsync-mode"]
 
+let zfs_path_without_pool =
+  Arg.flag @@
+  Arg.info ~doc:"Whether the ZFS file paths should include the pool name at the start or not"
+  ~docv:"ZFS_PATH_WITHOUT_POOL"
+  [ "zfs-path-without-pool" ]
+
 let rsync_mode =
   Arg.value @@ rsync_mode_opt
 
 (** Transform a [store] and [rsync-mode] into a validated combination.
 
     For example an rsync store must supply an rsync-mode.
  *)
-let of_t store rsync_mode =
+let of_t store rsync_mode zfs_path_without_pool =
   match store, rsync_mode with
   | Some (`Rsync path), Some rsync_mode -> `Rsync (path, rsync_mode)
   | Some (`Rsync _path), None -> failwith "Store rsync:/ must supply an rsync-mode"
   | Some (`Btrfs path), None -> (`Btrfs path)
-  | Some (`Zfs path), None -> (`Zfs path)
+  | Some (`Zfs (path, subdir)), None -> (`Zfs (path, subdir, zfs_path_without_pool))
   | Some (`Xfs path), None -> (`Xfs path)
   | Some (`Docker path), None -> (`Docker path)
   | _, _ -> failwith "Store type required must be one of btrfs:/path, rsync:/path, xfs:/path, zfs:pool or docker:path for the OBuilder cache."
@@ -95,7 +107,8 @@ let of_t store rsync_mode =
 let v =
   Term.(const of_t
         $ Arg.value @@ store ["store"]
-        $ Arg.value @@ rsync_mode_opt)
+        $ Arg.value @@ rsync_mode_opt
+        $ Arg.value @@ zfs_path_without_pool)
 
 (** Parse cli arguments for t and initialise a [store]. *)
 let cmdliner =

lib/zfs_store.ml

@@ -33,6 +33,8 @@ type cache = {
 
 type t = {
   pool : string;
+  path_with_pool : bool;
+  subdir : string option; (* This allows the destination datasets to be store under <prefix>/<pool>/<subdir>/result... *)
   prefix : string; (* To be prepended to `pool` to give the full path to the pool *)
   caches : (string, cache) Hashtbl.t;
   mutable next : int;
@@ -46,7 +48,7 @@ module Dataset : sig
   val state : dataset
   val cache_tmp_group : dataset
   val groups : dataset list
-
+  val subdir : t -> dataset option 
   val result : S.id -> dataset
   val cache : string -> dataset
   val cache_tmp : int -> string -> dataset
@@ -56,6 +58,7 @@ module Dataset : sig
 
   val exists : ?snapshot:string -> t -> dataset -> bool Lwt.t
   val if_missing : ?snapshot:string -> t -> dataset -> (unit -> unit Lwt.t) -> unit Lwt.t
+  val if_missing_subdir : t -> (string -> unit Lwt.t) -> unit Lwt.t
 end = struct
   type dataset = string
 
@@ -64,33 +67,62 @@ end = struct
   let cache_group = "cache"
   let cache_tmp_group = "cache-tmp"
 
+  let subdir t = t.subdir
+
   let groups = [state; result_group; cache_group; cache_tmp_group]
 
   let result id = "result/" ^ id
   let cache name = "cache/" ^ Escape.cache name
   let cache_tmp i name = strf "cache-tmp/%d-%s" i (Escape.cache name)
 
+  let dataset pool subdir ds = match subdir with
+    | None -> strf "%s/%s" pool ds
+    | Some subdir -> strf "%s/%s/%s" pool subdir ds
+
+  let dataset_no_pool subdir ds = match subdir with
+    | None -> ds
+    | Some subdir -> strf "%s/%s" subdir ds
+
   let full_name ?snapshot ?subvolume t ds =
     match snapshot, subvolume with
-    | None, None -> strf "%s/%s" t.pool ds
-    | Some snapshot, None -> strf "%s/%s@%s" t.pool ds snapshot
-    | None, Some subvolume -> strf "%s/%s/%s" t.pool ds subvolume
-    | Some snapshot, Some subvolume -> strf "%s/%s/%s@%s" t.pool ds subvolume snapshot
+    | None, None -> dataset t.pool t.subdir ds 
+    | Some snapshot, None -> 
+        strf "%s@%s" (dataset t.pool t.subdir ds) snapshot
+    | None, Some subvolume -> 
+        strf "%s/%s" (dataset t.pool t.subdir ds) subvolume
+    | Some snapshot, Some subvolume -> 
+        strf "%s/%s@%s" (dataset t.pool t.subdir ds) subvolume snapshot
 
   let path ?snapshot t ds =
     match snapshot with
-    | None -> strf "%s%s/%s" t.prefix t.pool ds
-    | Some snapshot -> strf "%s%s/%s/.zfs/snapshot/%s" t.prefix t.pool ds snapshot
-
-  let exists ?snapshot t ds =
-    Lwt_process.pread ("", [| "zfs"; "list"; "-p"; "-H"; full_name t ds ?snapshot |]) >>= function
+    | None ->
+        if t.path_with_pool then strf "%s%s" t.prefix (dataset t.pool t.subdir ds)
+        else strf "%s%s" t.prefix (dataset_no_pool t.subdir ds)
+    | Some snapshot -> 
+        if t.path_with_pool then strf "%s%s/.zfs/snapshot/%s" t.prefix (dataset t.pool t.subdir ds) snapshot
+        else strf "%s%s/.zfs/snapshot/%s" t.prefix (dataset_no_pool t.subdir ds) snapshot
+
+  let exists_raw raw =
+    Lwt_process.pread ("", [| "zfs"; "list"; "-p"; "-H"; raw |]) >>= function
     | "" -> Lwt.return false
     | _ -> Lwt.return true
 
+  let exists ?snapshot t ds =
+    exists_raw (full_name t ds ?snapshot)
+
   let if_missing ?snapshot t ds fn =
     exists ?snapshot t ds >>= function
     | true -> Lwt.return_unit
     | false -> fn ()
+
+  let if_missing_subdir t fn =
+    match t.subdir with 
+    | None -> Lwt.return_unit
+    | Some dir ->
+      let path = t.pool ^ "/" ^ dir in
+      exists_raw path >>= function
+      | true -> Lwt.return_unit
+      | false -> fn path
 end
 
 let user = `Unix { Obuilder_spec.uid = Unix.getuid (); gid = Unix.getgid () }
@@ -103,6 +135,9 @@ module Zfs = struct
   let create t ds =
     Os.sudo ["zfs"; "create"; "--"; Dataset.full_name t ds]
 
+  let create_raw t ds =
+    Os.sudo ["zfs"; "create"; "--"; ds ]
+
   let destroy t ds mode =
     let opts =
       match mode with
@@ -190,11 +225,12 @@ let prefix_and_pool path =
    | Some prefix -> (prefix, pool)
    | None -> failwith ("Failed to get preffix from: " ^ path)
 
-let create ~path =
+let create ?(path_with_pool=true) ?subdir ~path =
   let prefix, pool = prefix_and_pool path in
-  let t = { pool; prefix; caches = Hashtbl.create 10; next = 0 } in
+  let t = { pool; path_with_pool; subdir; prefix; caches = Hashtbl.create 10; next = 0 } in
   (* Ensure any left-over temporary datasets are removed before we start. *)
   delete_if_exists t (Dataset.cache_tmp_group) `And_snapshots_and_clones >>= fun () ->
+  Dataset.if_missing_subdir t (fun subdir -> Zfs.create_raw t subdir) >>= fun () -> 
   Dataset.groups |> Lwt_list.iter_s (fun group ->
       Dataset.if_missing t group (fun () -> Zfs.create t group) >>= fun () ->
       Zfs.chown ~user t group

lib/zfs_store.mli

@@ -2,7 +2,7 @@
 
 include S.STORE
 
-val create : path:string -> t Lwt.t
+val create : ?path_with_pool:bool -> ?subdir:string -> path:string -> t Lwt.t
 (** [create ~path] creates a new zfs store in a pool mounted at [path].
     The pool name is [Filename.basename path]. If only a poolname is passed
-    such as [tank] the path is inferred as [/tank].*)
+    such as [tank] the path is inferred as [/tank]. [subdir] is then added to the end of this if it is passed. *)