refactor: update TaskManager to implement CommandExecutionContext

Author: vringar

openwpm/task_manager.py

@@ -21,13 +21,16 @@
 from .browser_manager import BrowserManagerHandle
 from .command_sequence import CommandSequence
 from .errors import CommandExecutionError
+from .failure_tracker import FailureTracker
 from .js_instrumentation import clean_js_instrumentation_settings
 from .mp_logger import MPLogger
 from .storage.storage_controller import DataSocket, StorageControllerHandle
 from .storage.storage_providers import (
     StructuredStorageProvider,
+    TableName,
     UnstructuredStorageProvider,
 )
+from .types import VisitId
 from .utilities.multiprocess_utils import kill_process_and_children
 from .utilities.platform_utils import get_configuration_string, get_version
 from .utilities.storage_watchdog import StorageLogger
@@ -101,11 +104,7 @@ def __init__(
 
         # Flow control
         self.closing = False
-        self.failure_status: Optional[Dict[str, Any]] = None
-        self.threadlock = threading.Lock()
-        self.failure_count = 0
-
-        self.failure_limit = manager_params.failure_limit
+        self.failure_tracker = FailureTracker(manager_params.failure_limit)
         # Start logging server thread
         self.logging_server = MPLogger(
             self.manager_params.log_path,
@@ -331,6 +330,18 @@ def _shutdown_manager(
         if hasattr(self, "callback_thread"):
             self.callback_thread.join()
 
+    # CommandExecutionContext protocol implementation
+
+    def store_record(
+        self, table: TableName, visit_id: VisitId, data: Dict[str, Any]
+    ) -> None:
+        """Send a record to the StorageController via DataSocket."""
+        self.sock.store_record(table, visit_id, data)
+
+    def finalize_visit_id(self, visit_id: VisitId, success: bool) -> None:
+        """Signal that all data for a visit_id has been sent."""
+        self.sock.finalize_visit_id(visit_id, success)
+
     def _check_failure_status(self) -> None:
         """Check the status of command failures. Raise exceptions as necessary
 
@@ -340,25 +351,27 @@ def _check_failure_status(self) -> None:
         appropriate steps are taken to gracefully close the infrastructure
         """
         self.logger.debug("Checking command failure status indicator...")
-        if not self.failure_status:
+        if not self.failure_tracker.has_critical_failure:
             return
 
+        failure_status = self.failure_tracker.critical_failure
+        assert failure_status is not None
         self.logger.debug("TaskManager failure status set, halting command execution.")
         self._shutdown_manager()
-        if self.failure_status["ErrorType"] == "ExceedCommandFailureLimit":
+        if failure_status["ErrorType"] == "ExceedCommandFailureLimit":
             raise CommandExecutionError(
                 "TaskManager exceeded maximum consecutive command "
                 "execution failures.",
-                self.failure_status["CommandSequence"],
+                failure_status["CommandSequence"],
             )
-        elif self.failure_status["ErrorType"] == "ExceedLaunchFailureLimit":
+        elif failure_status["ErrorType"] == "ExceedLaunchFailureLimit":
             raise CommandExecutionError(
                 "TaskManager failed to launch browser within allowable "
                 "failure limit.",
-                self.failure_status["CommandSequence"],
+                failure_status["CommandSequence"],
             )
-        if self.failure_status["ErrorType"] == "CriticalChildException":
-            _, exc, tb = pickle.loads(self.failure_status["Exception"])
+        if failure_status["ErrorType"] == "CriticalChildException":
+            _, exc, tb = pickle.loads(failure_status["Exception"])
             raise exc.with_traceback(tb)
 
     # CRAWLER COMMAND CODE