From dc8d30c68454ad4d36797d350ff76daa119e296e Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Mon, 24 Oct 2022 18:04:22 +0900
Subject: [PATCH] cleanup delete, purge and recover methods

---
 src/ai/backend/manager/api/vfolder.py  |  3 ++-
 src/ai/backend/storage/abc.py          |  7 ++++--
 src/ai/backend/storage/vfs/__init__.py | 30 +++++++++++++++++++++-----
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/src/ai/backend/manager/api/vfolder.py b/src/ai/backend/manager/api/vfolder.py
index 54cd1674e27..f7aac1189dd 100644
--- a/src/ai/backend/manager/api/vfolder.py
+++ b/src/ai/backend/manager/api/vfolder.py
@@ -767,6 +767,7 @@ async def get_info(request: web.Request, row: VFolderRow) -> web.Response:
         "is_owner": is_owner,
         "permission": permission,
         "usage_mode": row["usage_mode"],
+        "status": row["status"],
         "cloneable": row["cloneable"],
         "max_size": row["max_size"],
     }
@@ -2006,7 +2007,7 @@ async def purge(request: web.Request) -> web.Response:
 
 @auth_required
 @server_status_required(ALL_ALLOWED)
-async def recover(request: web.Request, params: Any, row: VFolderRow) -> web.Response:
+async def recover(request: web.Request) -> web.Response:
     await ensure_vfolder_status(
         request, VFolderAccessStatus.RECOVERABLE, folder_name=request.match_info["name"]
     )
diff --git a/src/ai/backend/storage/abc.py b/src/ai/backend/storage/abc.py
index ef759d021e9..24c5f92444c 100644
--- a/src/ai/backend/storage/abc.py
+++ b/src/ai/backend/storage/abc.py
@@ -40,11 +40,14 @@ async def init(self) -> None:
     async def shutdown(self) -> None:
         pass
 
-    def mangle_vfpath(self, vfid: UUID) -> Path:
+    def mangle_rel_path(self, vfid: UUID) -> Path:
         prefix1 = vfid.hex[0:2]
         prefix2 = vfid.hex[2:4]
         rest = vfid.hex[4:]
-        return Path(self.mount_path, prefix1, prefix2, rest)
+        return Path(prefix1, prefix2, rest)
+
+    def mangle_vfpath(self, vfid: UUID) -> Path:
+        return Path(self.mount_path, self.mangle_rel_path(vfid))
 
     def sanitize_vfpath(
         self,
diff --git a/src/ai/backend/storage/vfs/__init__.py b/src/ai/backend/storage/vfs/__init__.py
index f1298aff4b6..460b3db192f 100644
--- a/src/ai/backend/storage/vfs/__init__.py
+++ b/src/ai/backend/storage/vfs/__init__.py
@@ -48,6 +48,9 @@ async def run(cmd: Sequence[Union[str, Path]]) -> str:
 
 
 class BaseVolume(AbstractVolume):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        Path.mkdir(self.trash_path, parents=True, exist_ok=True)
 
     # ------ volume operations -------
 
@@ -75,18 +78,30 @@ async def create_vfolder(
             lambda: vfpath.mkdir(0o755, parents=True, exist_ok=exist_ok),
         )
 
+    @staticmethod
+    def _clean_empty_parents(path: Path) -> None:
+        # remove intermediate prefix directories if they become empty
+        if not os.listdir(path.parent):
+            path.parent.rmdir()
+        if not os.listdir(path.parent.parent):
+            path.parent.parent.rmdir()
+
     async def delete_vfolder(self, vfid: UUID) -> None:
         vfpath = self.mangle_vfpath(vfid)
-        dst = self.trash_path / vfpath
+        dst = self.trash_path / self.mangle_rel_path(vfid)
         loop = asyncio.get_running_loop()
 
+        def _delete_vfolder():
+            shutil.move(vfpath, dst)
+            self._clean_empty_parents(vfpath)
+
         await loop.run_in_executor(
             None,
-            lambda: shutil.move(vfpath, dst),
+            _delete_vfolder,
         )
 
     async def purge_vfolder(self, vfid: UUID) -> None:
-        vfpath = self.trash_path / self.mangle_vfpath(vfid)
+        vfpath = self.trash_path / self.mangle_rel_path(vfid)
         loop = asyncio.get_running_loop()
 
         def _purge_vfolder():
@@ -94,17 +109,22 @@ def _purge_vfolder():
                 shutil.rmtree(vfpath)
             except FileNotFoundError:
                 pass
+            self._clean_empty_parents(vfpath)
 
         await loop.run_in_executor(None, _purge_vfolder)
 
     async def recover_vfolder(self, vfid: UUID) -> None:
-        src = self.trash_path / self.mangle_vfpath(vfid)
+        src = self.trash_path / self.mangle_rel_path(vfid)
         dst = self.mangle_vfpath(vfid)
         loop = asyncio.get_running_loop()
 
+        def _recover_vfolder():
+            shutil.move(src, dst)
+            self._clean_empty_parents(src)
+
         await loop.run_in_executor(
             None,
-            lambda: shutil.move(src, dst),
+            _recover_vfolder,
         )
 
     async def clone_vfolder(