Improved net bot implementation

main.py

@@ -5,11 +5,16 @@
 import asyncio
 import traceback
 import signal
+from pathlib import Path
 
 import config
 import util
-import bot
+#import bot
-import logic
+#import logic
+
+from new_bot import MatrixBot
+from bot_types import MatrixBotConfig
+
 from datatypes import AppConfig
 
 async def main() -> None:
@@ -30,17 +35,16 @@ async def main() -> None:
         util.log_error("Could't load config")
         return
 
-    # start the bot
+    matrix_bot_config = MatrixBotConfig(
-    if not await bot.start(cfg):
+        matrix_homeserver_url=cfg.matrix_homeserver,
-        util.log_error("Could't start the bot")
+        matrix_username_localpart=cfg.matrix_user,
-        return
+        storage_directory=Path(cfg.store_dir)
-    await logic.setup(bot.get_client())
+    )
+    bot = MatrixBot(matrix_bot_config)
+    bot.start()
 
-    # wait for stop
     await util.get_app_stop_event().wait()
 
-    # stop
-    await logic.stop()
     await bot.stop()
 
 if __name__ == "__main__":

new_bot.py

@@ -2,6 +2,7 @@ import asyncio
 import aiofiles
 import aioconsole
 import traceback
+import logging
 import time
 import json
 import os
@@ -93,6 +94,31 @@ class MatrixBot:
         localpart = config.matrix_username_localpart
         return f"@{localpart}:{homeserver_name}"
 
+    @staticmethod
+    async def _wait_for_task_and_stop_event(payload_task: asyncio.Task, stop_wait_task: asyncio.Task, print_exc: bool = False) -> Any:
+        """Waits for `payload_task` to finish. Uses `stop_wait_task` to check if payload task
+            should be cancelled. Returns value returned by `payload_task` task. Raises exception
+            raised by `payload_task` task.
+
+            If `stop_wait_task` finishes, then `asyncio.CancelledError` is raised.
+        """
+        try:
+            done, _ = await asyncio.wait(
+                [payload_task, stop_wait_task],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+            if stop_wait_task in done:
+                payload_task.cancel()
+                await payload_task
+                raise asyncio.CancelledError()
+            return payload_task.result()
+        except asyncio.CancelledError:
+            raise
+        except:
+            if print_exc:
+                traceback.print_exc()
+            raise
+
     @staticmethod
     async def _default_password_callback() -> str:
         """Gets password from `MATRIX_PASSWORD` envvar if it is set. Asks
@@ -108,6 +134,7 @@ class MatrixBot:
         path = self._config.storage_directory / "next_batch"
         async with aiofiles.open(path, "w") as f:
             await f.write(next_batch)
+        self._logger.debug("next_batch value is written to the disk")
 
     async def _read_next_batch(self) -> str | None:
         """Reads `next_batch` value from disk. Returns None if file does not exist."""
@@ -126,6 +153,7 @@ class MatrixBot:
         }
         async with aiofiles.open(path, "w") as f:
             await f.write(json.dumps(data, indent=4))
+        self._logger.debug("Session data is writter to the disk")
 
     async def _read_session_data(self) -> dict[str, Any] | None:
         """Read session data from disk."""
@@ -143,10 +171,11 @@ class MatrixBot:
         """This callback is called when AsyncClient syncs with the server"""
         current_time = time.time()
         delta_time = current_time - self._last_next_batch_dump
+        self._last_next_batch = response.next_batch
         if delta_time >= self.NEXT_BATCH_DUMP_PERIOD:
             self._last_next_batch_dump = current_time
             try:
-                await self._write_next_batch(response.next_batch)
+                await self._write_next_batch(self._last_next_batch)
             except:
                 traceback.print_exc()
 
@@ -159,6 +188,7 @@ class MatrixBot:
 
     async def _client_login_session_data(self, session_data: dict[str, Any]) -> None:
         """Login using session data. Raises and exception on failure."""
+        self._logger.debug("Using stored session data to log in")
         # build user id
         username = self._build_matrix_username(self._config)
         self._client.restore_login(
@@ -167,25 +197,32 @@ class MatrixBot:
         )
         result = await self._client.whoami()
         if type(result) is WhoamiError:
+            self._logger.error(f"Can't log in using stored session data: '{result.message}'")
             raise RuntimeError(result.message)
         elif type(result) is not WhoamiResponse:
+            self._logger.error("Can't log in using stored session data, unknown error")
             raise RuntimeError("Unknown response for whoami request")
+        self._logger.debug("Logged in using stored session data")
 
     async def _client_login_password(self) -> None:
         """Login using password and save result to disk on success.
             Raises an exception on failure.
         """
+        self._logger.debug("Using password to log in")
         # get the password
         password = await self._cb_password()
         result = await self._client.login(password=password)
         if type(result) is LoginResponse:
+            self._logger.debug("Logged in using password")
             await self._write_session_data(
                 access_token=result.access_token,
                 device_id=result.device_id
             )
         elif type(result) is LoginError:
+            self._logger.error(f"Can't log in using password: '{result.message}'")
             raise RuntimeError(result.message)
         else:
+            self._logger.error(f"Can't log in using password, unknown error")
             raise RuntimeError("Unknown login result")
 
     async def _client_login(self) -> None:
@@ -194,31 +231,68 @@ class MatrixBot:
         session_data = await self._read_session_data()
         # session data is present, try to log in
         if session_data is not None:
+            self._logger.debug("Some session data found on the disk")
             try:
                 await self._client_login_session_data(session_data)
                 return
+            except asyncio.CancelledError:
+                raise
             except:
                 pass
         # no session data - login using password
         try:
+            self._logger.debug("No session data found on the disk OR invalid data")
             await self._client_login_password()
             return
+        except asyncio.CancelledError:
+            raise
         except:
             pass
         # can't login
+        self._logger.error("Can't log in using available methods")
         raise RuntimeError("All login methods have failed, can't continue")
 
+    async def _client_destroy(self) -> None:
+        """Gracefully destroys the client."""
+        try:
+            self._logger.debug("Closing the client")
+            await self._client.close()
+            if self._last_next_batch is not None:
+                self._logger.debug("Saving next_batch")
+                await self._write_next_batch(self._last_next_batch)
+        except:
+            traceback.print_exc()
+
     async def _background_coroutine(self) -> None:
         """This function implements bot lifecycle."""
         # we should stop when this task stops
+        self._logger.debug("_background_coroutine is started")
         stop_wait_task = asyncio.create_task(self._stop_event.wait())
         self._client = self._build_client(self._config)
         self._setup_client_callbacks()
         # perform login
+        login_task = asyncio.create_task(self._client_login())
         try:
-            await self._client_login()
+            await self._wait_for_task_and_stop_event(login_task, stop_wait_task)
+        except asyncio.CancelledError:
+            self._logger.debug("Background task is cancelled during login")
+            await self._client_destroy()
+            return
         except:
             traceback.print_exc()
+        self._logger.info("Succesfully logged in")
+        # test
+        wait_task = asyncio.create_task(asyncio.sleep(100000))
+        try:
+            self._logger.info("Bot is not implemented yet, sleeping forever")
+            await self._wait_for_task_and_stop_event(wait_task, stop_wait_task)
+        except asyncio.CancelledError:
+            self._logger.info("Background task is cancelled during eternal sleep")
+            await self._client_destroy()
+            return
+        except:
+            traceback.print_exc()
+
 
     #
     # PUBLIC
@@ -228,11 +302,15 @@ class MatrixBot:
         self._validate_bot_config(config) # may raise an Exception
         # save the config
         self._config = config
+        # create the logger
+        self._logger = logging.getLogger(self._build_matrix_username(config))
+        self._logger.setLevel(logging.DEBUG)
 
         # prepare some private data
         self._background_task: asyncio.Task | None = None
         self._client: AsyncClient = None # type: ignore
         self._last_next_batch_dump: float = 0.0
+        self._last_next_batch: str | None = None
         self._cb_password = self._default_password_callback
 
     def start(self) -> None:
@@ -249,22 +327,14 @@ class MatrixBot:
             self._background_coroutine()
         )
 
-    def stop(self) -> bool:
+    async def stop(self) -> None:
-        """Stop the bot.
+        """Stop the bot and wait for the bot stop."""
-            Signals the background task to stop. Returns True on success (the
-            signal is sent), False on failure (no background task). This
-            function won't wait for the task to stop. Use `wait_stop` to wait.
-        """
-        if self._background_task is None:
-            return False
-        self._stop_event.set()
-        return True
-
-    async def wait_stop(self) -> None:
-        """Wait for the background task to stop.
-            You must call `stop()` by yourself. This function will never return
-            otherwise. It just waits for the background task stop.
-        """
         if self._background_task is None:
             return
-        await asyncio.wait([self._background_task])
+        self._stop_event.set()
+        try:
+            await self._background_task
+        except:
+            traceback.print_exc()
+        self._stop_event = None
+        self._background_task = None