refactor(cleanup): add logs and use new type annotation

openwpm · Oct 21, 2023 · 288100f · 288100f
1 parent 62e7c82
commit 288100f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/openwpm/browser_manager.py b/openwpm/browser_manager.py
@@ -549,8 +549,8 @@ def kill_browser_manager(self):
         if self.browser_manager is not None and self.browser_manager.pid is not None:
             self.logger.debug(
                 "BROWSER %i: Attempting to kill BrowserManager with pid %i. "
-                "Browser PID: %s"
-                % (self.browser_id, self.browser_manager.pid, self.geckodriver_pid)
+                "Browser PID: %s",
+                (self.browser_id, self.browser_manager.pid, self.geckodriver_pid),
             )
             try:
                 os.kill(self.browser_manager.pid, signal.SIGKILL)

diff --git a/openwpm/storage/storage_controller.py b/openwpm/storage/storage_controller.py
@@ -72,11 +72,11 @@ def __init__(
         self._shutdown_flag = False
         self._relaxed = False
         self.logger = logging.getLogger("openwpm")
-        self.store_record_tasks: DefaultDict[VisitId, List[Task[None]]] = defaultdict(
+        self.store_record_tasks: DefaultDict[VisitId, list[Task[None]]] = defaultdict(
             list
         )
         """Contains all store_record tasks for a given visit_id"""
-        self.finalize_tasks: List[Tuple[VisitId, Optional[Task[None]], bool]] = []
+        self.finalize_tasks: list[tuple[VisitId, Optional[Task[None]], bool]] = []
         """Contains all information required for update_completion_queue to work
             Tuple structure is: VisitId, optional completion token, success
         """
@@ -251,6 +251,7 @@ async def update_status_queue(self) -> NoReturn:
             )
 
     async def shutdown(self, completion_queue_task: Task[None]) -> None:
+        self.logger.info("Entering self.shutdown")
         completion_tokens = {}
         visit_ids = list(self.store_record_tasks.keys())
         for visit_id in visit_ids:
@@ -264,6 +265,7 @@ async def shutdown(self, completion_queue_task: Task[None]) -> None:
             self.completion_queue.put((visit_id, False))
 
         await self.structured_storage.shutdown()
+        self.logger.info("structured_storage is shut down")
 
         if self.unstructured_storage is not None:
             await self.unstructured_storage.flush_cache()
@@ -345,13 +347,21 @@ async def _run(self) -> None:
         update_completion_queue = asyncio.create_task(
             self.update_completion_queue(), name="CompletionQueueFeeder"
         )
-        # Blocks until we should shutdown
+        # Blocks until we should shut down
         await self.should_shutdown()
-
+        self.logger.info(f"Closing Server")
         server.close()
+        self.logger.info("Closed Server")
+        self.logger.info("Cancelling status_queue_update")
         status_queue_update.cancel()
+        self.logger.info("Cancelled status_queue_update")
+        self.logger.info("Cancelling timeout_check")
         timeout_check.cancel()
+        self.logger.info("Cancelled timeout_check")
+        self.logger.info("Starting wait_closed")
         await server.wait_closed()
+        self.logger.info("Completed wait_closed")
+
         await self.shutdown(update_completion_queue)
 
     def run(self) -> None: