From 113662703545c6ec0d7dcc516edeb85ae73f5b77 Mon Sep 17 00:00:00 2001
From: Virgile Robles <virgile.robles@protonmail.ch>
Date: Tue, 3 Sep 2024 18:19:42 +0200
Subject: [PATCH] kv: rework example to expose pure KV

---
 README.md                 |  5 +--
 unikernel-kv/config.ml    | 73 +++++++++++++++++++++++++++++++++++++--
 unikernel-kv/unikernel.ml | 38 +++++++++-----------
 3 files changed, 87 insertions(+), 29 deletions(-)

diff --git a/README.md b/README.md
index a6d36dc..d39eef6 100644
--- a/README.md
+++ b/README.md
@@ -17,11 +17,8 @@ As a work-in-progress demo, `notafs` includes a partial implementation of the [`
 To run the unikernel demos, you'll need to pin the `notafs` library, copy the `unikernel-kv` folder out of the project (to avoid recursive issues with `opam-monorepo`), compile it for your prefered Mirage target and create a disk to use:
 
 ```shell
-# Pin the library
-$ cd notafs
-notafs/ $ opam pin add notafs . --with-version=dev
-
 # Copy the mirage-kv demo to another folder
+$ cd notafs
 notafs/ $ cp -r unikernel-kv ../unikernel-kv
 notafs/ $ cd ../unikernel-kv
 
diff --git a/unikernel-kv/config.ml b/unikernel-kv/config.ml
index 41c01ff..8b830b5 100644
--- a/unikernel-kv/config.ml
+++ b/unikernel-kv/config.ml
@@ -1,5 +1,72 @@
 open Mirage
 
-let main = main "Unikernel.Main" (block @-> job) ~packages:[ package "notafs" ]
-let img = if_impl Key.is_solo5 (block_of_file "storage") (block_of_file "/tmp/storage")
-let () = register "block_test" [ main $ img ]
+(* This needs to be included until support for notafs is merged
+   in the upstream mirage tool. *)
+open (
+  struct
+    type checksum = CHECKSUM
+
+    let checksum_t = Type.v CHECKSUM
+
+    let notafs_kv_rw_conf ~format =
+      (* TODO remove pin when notafs is published on opam *)
+      let packages =
+        [ package ~pin:"git+https://github.com/tarides/notafs.git" "notafs" ]
+      in
+      let connect _ modname = function
+        | [ _pclock_v; _checksum; block_v ] ->
+            let connect_c = Fmt.str "%s.connect %s" modname block_v in
+            let format_c = Fmt.str "%s.format %s" modname block_v in
+            let connect_format_c =
+              match format with
+              | `Never -> connect_c
+              | `Always -> format_c
+              | `If_not ->
+                  Fmt.str
+                    {ml|%s >>= function
+                        | Error `Disk_not_formatted -> %s
+                        | x -> Lwt.return x|ml}
+                    connect_c format_c
+            in
+            code ~pos:__POS__
+              {ml|(%s)
+                >|= Result.map_error (Fmt.str "notafs_kv_rw: %%a" %s.pp_error)
+                >|= Result.fold ~ok:Fun.id ~error:failwith|ml}
+              connect_format_c modname
+        | _ -> connect_err "notafs_kv_rw" 3
+      in
+      impl ~packages ~connect "Notafs.KV"
+        (pclock @-> checksum_t @-> block @-> kv_rw)
+
+    let notafs_kv_rw ?(pclock = default_posix_clock) ?(checksum = `Adler32)
+        ?(format = `If_not) block =
+      let checksum_modname =
+        match checksum with
+        | `Adler32 -> "Notafs.Adler32"
+        | `No_checksum -> "Notafs.No_checksum"
+      in
+      let checksum = impl checksum_modname checksum_t in
+      notafs_kv_rw_conf ~format $ pclock $ checksum $ block
+  end :
+    sig
+      val notafs_kv_rw :
+        ?pclock:pclock impl ->
+        ?checksum:[ `Adler32 | `No_checksum ] ->
+        ?format:[ `Always | `Never | `If_not ] ->
+        block impl ->
+        kv_rw impl
+      (** [notafs_kv_rw ~checksum ~format block] exposes a KV_RW interface from
+          a notafs block, with the given checksum mechanism. The underlying
+          block is expected to be a well-formed notafs volume if
+          [format = `Never], is always formatted (and cleared) to be one if
+          [format = `Always], or only as needed (the first time it's opened) if
+          [format = `If_not] (the default). *)
+    end)
+
+let main = main "Unikernel.Main" (kv_rw @-> job)
+
+let block =
+  if_impl Key.is_solo5 (block_of_file "storage") (block_of_file "/tmp/storage")
+
+let kv = notafs_kv_rw block
+let () = register "block_test" [ main $ kv ]
diff --git a/unikernel-kv/unikernel.ml b/unikernel-kv/unikernel.ml
index e412068..838d467 100644
--- a/unikernel-kv/unikernel.ml
+++ b/unikernel-kv/unikernel.ml
@@ -1,32 +1,26 @@
 open Lwt.Syntax
 
-module Main (Block : Mirage_block.S) = struct
-  module Kv = Notafs.KV (Pclock) (Notafs.Adler32) (Block)
-
+module Main (KV : Mirage_kv.RW) = struct
   let force lwt =
     let open Lwt.Infix in
-    lwt
-    >|= function
+    lwt >|= function
     | Ok v -> v
     | Error e ->
-      Format.printf "ERROR: %a@." Kv.pp_error e ;
-      failwith "error"
+        Logs.err (fun f -> f "Error: %a" KV.pp_write_error e);
+        failwith "fatal error"
 
-  let start block =
-    let* fs = Kv.connect block in
-    let* fs =
-      match fs with
-      | Ok fs -> Lwt.return fs
-      | Error `Disk_not_formatted ->
-        let* fs = force @@ Kv.format block in
-        let+ () = force @@ Kv.set fs (Mirage_kv.Key.v "hello") "world!" in
-        fs
-      | Error e ->
-        Format.printf "ERROR: %a@." Kv.pp_error e ;
-        failwith "unexpected error"
+  let start kv =
+    let key = Mirage_kv.Key.v "hello" in
+    let* result = KV.get kv key in
+    let* () =
+      match result with
+      | Ok contents ->
+          Logs.info (fun f -> f "Key hello contains %S" contents);
+          Lwt.return_unit
+      | Error _ ->
+          Logs.warn (fun f -> f "Key hello doesn't exist, creating it!");
+          force @@ KV.set kv key "world!"
     in
-    let* contents = force @@ Kv.get fs (Mirage_kv.Key.v "hello") in
-    Format.printf "%S@." contents ;
-    let* () = Block.disconnect block in
+    let* () = KV.disconnect kv in
     Lwt.return_unit
 end