Use an isolated clone of the Session in workunit handlers to shield against cancellation

src/python/pants/engine/internals/native_engine.pyi

@@ -120,6 +120,7 @@ def session_get_observation_histograms(scheduler: PyScheduler, session: PySessio
 def session_record_test_observation(
     scheduler: PyScheduler, session: PySession, value: int
 ) -> None: ...
+def session_isolated_shallow_clone(session: PySession) -> PySession: ...
 def all_counter_names() -> list[str]: ...
 def graph_len(scheduler: PyScheduler) -> int: ...
 def graph_visualize(scheduler: PyScheduler, session: PySession, path: str) -> None: ...

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:

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,

src/rust/engine/src/externs/interface.rs

@@ -318,6 +318,12 @@ py_module_initializer!(native_engine, |py, m| {
       session_record_test_observation(a: PyScheduler, b: PySession, c: u64)
     ),
   )?;
+  m.add(
+    py,
+    "session_isolated_shallow_clone",
+    py_fn!(py, session_isolated_shallow_clone(a: PySession)),
+  )?;
+
   m.add(py, "all_counter_names", py_fn!(py, all_counter_names()))?;
 
   m.add(
@@ -1389,6 +1395,12 @@ fn session_record_test_observation(
   })
 }
 
+fn session_isolated_shallow_clone(py: Python, session_ptr: PySession) -> CPyResult<PySession> {
+  with_session(py, session_ptr, |session| {
+    PySession::create_instance(py, session.isolated_shallow_clone())
+  })
+}
+
 fn validate_reachability(py: Python, scheduler_ptr: PyScheduler) -> PyUnitResult {
   with_scheduler(py, scheduler_ptr, |scheduler| {
     scheduler