Repository: pantsbuild/pants

Issue:
Async streaming workunit handlers are cancelled when the client goes away
As described (but incorrectly diagnosed: see below) in https://github.com/pantsbuild/pants/pull/11687#issuecomment-797872718, when a streaming workunit handler completes in the background, any requests it makes to the engine after the client has gone away will fail.

The warning at the top of [the error message](https://gist.github.com/stuhood/a498acc9d9a19cb9587370ac51ec0647) is (mostly) a red herring: although we do fail to start a UI after the console has gone away (and would probably fail in the middle of the run if it were to go away), the actual error is rendered by the stacktrace: a `KeyboardInterrupt`.

This is because when a client disconnects we set `Session.cancelled`, which will cause all further calls to `Scheduler.execute` to immediately fail with a KeyboardInterrupt here: https://github.com/pantsbuild/pants/blob/030c99f729eb1bb1ddaa7257ade1d40f8012a648/src/rust/engine/src/scheduler.rs#L380-L383

----

The SWH needs access to the metrics (the `WorkunitStore`, `workunit_metadata_map`, etc) of the `Session` that it was started for, but it needs to be disconnected from the original `Session` in terms of cancellation. Fixing this might involve a special-cased variant of clone like `Session::clone_with_shared_metrics`.


Discussion:

Resolving patch:
diff --git a/src/python/pants/engine/internals/scheduler.py b/src/python/pants/engine/internals/scheduler.py
--- a/src/python/pants/engine/internals/scheduler.py
+++ b/src/python/pants/engine/internals/scheduler.py
@@ -130,6 +130,7 @@ def __init__(
         """
         self.include_trace_on_error = include_trace_on_error
         self._visualize_to_dir = visualize_to_dir
+        self._visualize_run_count = 0
         # Validate and register all provided and intrinsic tasks.
         rule_index = RuleIndex.create(rules)
         tasks = register_rules(rule_index, union_membership)
@@ -311,7 +312,6 @@ class SchedulerSession:
     def __init__(self, scheduler: Scheduler, session: PySession) -> None:
         self._scheduler = scheduler
         self._py_session = session
-        self._run_count = 0
 
     @property
     def scheduler(self) -> Scheduler:
@@ -325,6 +325,11 @@ def py_scheduler(self) -> PyScheduler:
     def py_session(self) -> PySession:
         return self._py_session
 
+    def isolated_shallow_clone(self) -> SchedulerSession:
+        return SchedulerSession(
+            self._scheduler, native_engine.session_isolated_shallow_clone(self._py_session)
+        )
+
     def poll_workunits(self, max_log_verbosity: LogLevel) -> PolledWorkunits:
         result = native_engine.session_poll_workunits(
             self.py_scheduler, self.py_session, max_log_verbosity.level
@@ -401,8 +406,9 @@ def metrics(self) -> dict[str, int]:
 
     def _maybe_visualize(self) -> None:
         if self._scheduler.visualize_to_dir is not None:
-            name = f"graph.{self._run_count:03d}.dot"
-            self._run_count += 1
+            # TODO: This increment-and-get is racey.
+            name = f"graph.{self._scheduler._visualize_run_count:03d}.dot"
+            self._scheduler._visualize_run_count += 1
             self.visualize_graph_to_file(os.path.join(self._scheduler.visualize_to_dir, name))
 
     def teardown_dynamic_ui(self) -> None:
diff --git a/src/python/pants/engine/streaming_workunit_handler.py b/src/python/pants/engine/streaming_workunit_handler.py
--- a/src/python/pants/engine/streaming_workunit_handler.py
+++ b/src/python/pants/engine/streaming_workunit_handler.py
@@ -185,6 +185,7 @@ def __init__(
         pantsd: bool,
         max_workunit_verbosity: LogLevel = LogLevel.TRACE,
     ) -> None:
+        scheduler = scheduler.isolated_shallow_clone()
         self.callbacks = callbacks
         self.context = StreamingWorkunitContext(
             _scheduler=scheduler,


Test patch: