more msd async

2 files changed, 40 insertions, 20 deletions

diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index ddb90a79..53082cc6 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -431,10 +431,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                 with Inotify() as inotify:
                     for path in [
-                        *self.__storage.get_watchable_paths(),
+                        *(await self.__storage.get_watchable_paths()),
                         *self.__drive.get_watchable_paths(),
                     ]:
-                        inotify.watch(path, InotifyMask.ALL_MODIFY_EVENTS)
+                        await inotify.watch(path, InotifyMask.ALL_MODIFY_EVENTS)
 
                     # После установки вотчеров еще раз проверяем стейт, чтобы ничего не потерять
                     await self.__reload_state()
@@ -471,7 +471,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     await self.__storage.remount_rw(False)
                     await self.__setup_initial()
 
-                storage_state = self.__get_storage_state()
+                storage_state = await self.__get_storage_state()
 
             except Exception:
                 logger.exception("Error while reloading MSD state; switching to offline")
@@ -514,8 +514,8 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     # =====
 
-    def __get_storage_state(self) -> _StorageState:
-        images = self.__storage.get_images()
+    async def __get_storage_state(self) -> _StorageState:
+        images = await self.__storage.get_images()
         space = self.__storage.get_space(fatal=True)
         assert space
         return _StorageState(
diff --git a/kvmd/plugins/msd/otg/storage.py b/kvmd/plugins/msd/otg/storage.py
index 4cca878b..c6a3e34b 100644
--- a/kvmd/plugins/msd/otg/storage.py
+++ b/kvmd/plugins/msd/otg/storage.py
@@ -21,6 +21,7 @@
 
 
 import os
+import operator
 import dataclasses
 
 from typing import Generator
@@ -31,6 +32,7 @@ import aiofiles.os
 
 from ....logging import get_logger
 
+from .... import aiotools
 from .... import aiohelpers
 
 from .. import MsdError
@@ -130,26 +132,42 @@ class Storage:
         self.__path = path
         self.__remount_cmd = remount_cmd
 
-    def get_watchable_paths(self) -> list[str]:
-        paths: list[str] = []
-        for (root_path, dirs, _) in os.walk(self.__path):
-            dirs[:] = list(self.__filtered(dirs))
-            paths.append(root_path)
-        return paths
+    async def get_watchable_paths(self) -> list[str]:
+        return (await aiotools.run_async(self.__get_watchable_paths))
 
-    def get_images(self) -> dict[str, Image]:
+    async def get_images(self) -> dict[str, Image]:
+        return (await aiotools.run_async(self.__get_images))
+
+    def __get_watchable_paths(self) -> list[str]:
+        return list(map(operator.itemgetter(0), self.__walk(with_files=False)))
+
+    def __get_images(self) -> dict[str, Image]:
         images: dict[str, Image] = {}
-        for (root_path, dirs, files) in os.walk(self.__path):
-            dirs[:] = list(self.__filtered(dirs))
-            for file in self.__filtered(files):
-                name = os.path.relpath(os.path.join(root_path, file), self.__path)
+        for (_, files) in self.__walk(with_files=True):
+            for path in files:
+                name = os.path.relpath(path, self.__path)
                 images[name] = self.get_image_by_name(name)
         return images
 
-    def __filtered(self, items: list[str]) -> Generator[str, None, None]:
-        for item in sorted(map(str.strip, items)):
-            if not item.startswith(".") and item != "lost+found":
-                yield item
+    def __walk(self, with_files: bool, root_path: (str | None)=None) -> Generator[tuple[str, list[str]], None, None]:
+        if root_path is None:
+            root_path = self.__path
+        files: list[str] = []
+        with os.scandir(root_path) as dir_iter:
+            for item in sorted(dir_iter, key=operator.attrgetter("name")):
+                if item.name.startswith(".") or item.name == "lost+found":
+                    continue
+                try:
+                    if item.is_dir(follow_symlinks=False):
+                        item.stat()  # Проверяем, не сдохла ли смонтированная NFS
+                        yield from self.__walk(with_files, item.path)
+                    elif with_files and item.is_file(follow_symlinks=False):
+                        files.append(item.path)
+                except Exception:
+                    pass
+        yield (root_path, files)
+
+    # =====
 
     def get_image_by_name(self, name: str) -> Image:
         assert name
@@ -170,6 +188,8 @@ class Storage:
         assert path
         return Image(name, path, (self if in_storage else None))
 
+    # =====
+
     def get_space(self, fatal: bool) -> (StorageSpace | None):
         try:
             st = os.statvfs(self.__path)