[Libguestfs] [PATCH 1/2] mllib: add an hook to cleanup directories on exit

[Pino Toscano]

Much similar to unlink_on_exit, but recursively cleaning directories.
---
 mllib/common_utils.ml | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/mllib/common_utils.ml b/mllib/common_utils.ml
index 3943417..d02a2d3 100644
--- a/mllib/common_utils.ml
+++ b/mllib/common_utils.ml
@@ -386,6 +386,41 @@ let unlink_on_exit =
       registered_handlers := true
     )
 
+(* Remove a temporary directory on exit. *)
+let rmdir_on_exit =
+  let dirs = ref [] in
+  let registered_handlers = ref false in
+
+  let rec unlink_dirs () =
+    let rec recursive_rmdir fn =
+      match (Unix.lstat fn).Unix.st_kind with
+      | Unix.S_DIR ->
+        let names = Array.map (fun d -> fn // d) (Sys.readdir fn) in
+        Array.iter recursive_rmdir names;
+        Unix.rmdir fn
+      | Unix.S_REG
+      | Unix.S_CHR
+      | Unix.S_BLK
+      | Unix.S_LNK
+      | Unix.S_FIFO
+      | Unix.S_SOCK ->
+        Unix.unlink fn
+    in
+    List.iter (
+      fun dir -> try recursive_rmdir dir with _ -> ()
+    ) !dirs
+  and register_handlers () =
+    (* Remove on exit. *)
+    at_exit unlink_dirs
+  in
+
+  fun dir ->
+    dirs := dir :: !dirs;
+    if not !registered_handlers then (
+      register_handlers ();
+      registered_handlers := true
+    )
+
 (* Using the libguestfs API, recursively remove only files from the
  * given directory.  Useful for cleaning /var/cache etc in sysprep
  * without removing the actual directory structure.  Also if 'dir' is
-- 
1.8.3.1