fix(concurrency): Close race windows in cancel and completion paths

Address findings F1-F8 and Review22-F3/F4 from round 3 concurrency review:

- F1: Always route EnvironmentScriptRunner.cancel() through pending-cancel path
- F2+F4: Atomic terminal arbitration - claim pending cancel in _on_process_exit,
  move liveness check inside lock in _cancel
- F3: Monotonic merge for duplicate pending cancels (min time_limit, OR failed)
- F5: Snapshot action_status before publishing READY state
- F6: Wrap cancel_info.json write in try/except, fallback to immediate terminate
- F7: Detect self-join in shutdown(), use wait=False if called from worker thread
- F8: Move callback outside lock, wrap in try/except to prevent child discard
- Review22-F3: Snapshot _runner in cancel_action to avoid bare AssertionError
- Review22-F4: Bind _process once in notify/terminate to avoid TOCTOU race

Add test_concurrency_fixes.py with 8 unit tests covering the defensive behaviors.

None of these issues reproduce in openjd-rs due to CancellationToken, tokio async,
and Rust's Result<> error handling model.

Signed-off-by: David Leong <116610336+leongdl@users.noreply.github.com>

Author: leongdl

src/openjd/sessions/_runner_base.py

@@ -589,7 +589,23 @@ def shutdown(self) -> None:
         """Performs a clean shutdown on the runner. This shutsdown the internal
         ThreadPoolExectutor.
         """
-        self._pool.shutdown()
+        # F7 fix: Detect self-join. If shutdown() is called from the pool's
+        # worker thread (e.g., from _on_process_exit -> _action_callback ->
+        # cleanup -> runner.shutdown), calling _pool.shutdown(wait=True) would
+        # deadlock because the thread is trying to join itself. In that case,
+        # defer to a background thread or skip the wait.
+        import threading
+
+        # Check if current thread is the pool's worker thread. The pool has
+        # max_workers=1, so if a future is running, _threads has one element.
+        pool_threads: set[threading.Thread] = getattr(self._pool, "_threads", set())
+        current = threading.current_thread()
+        if current in pool_threads:
+            # We're inside the worker thread. Use wait=False to avoid deadlock.
+            # The pool will be garbage-collected eventually.
+            self._pool.shutdown(wait=False)
+        else:
+            self._pool.shutdown()
 
     def _fail_action(self, message: str) -> None:
         """Fail the action through the normal failure path: surface the
@@ -936,7 +952,24 @@ def _cancel_with_resolved_method(
                 return
             process = self._process
             if process is None or not process.has_started:
-                self._pending_cancel = (time_limit, mark_action_failed)
+                # F3 fix: Monotonic merge for duplicate pending cancels. If a
+                # cancel is already pending, merge the new request: take the
+                # minimum time_limit (tighter deadline wins, treating None as
+                # unlimited), and OR the mark_action_failed flags (once failed,
+                # always failed).
+                if self._pending_cancel is not None:
+                    prev_limit, prev_failed = self._pending_cancel
+                    # Merge time limits: None means unlimited, so a defined limit beats None
+                    if time_limit is None:
+                        merged_limit = prev_limit
+                    elif prev_limit is None:
+                        merged_limit = time_limit
+                    else:
+                        merged_limit = min(time_limit, prev_limit)
+                    merged_failed = mark_action_failed or prev_failed
+                    self._pending_cancel = (merged_limit, merged_failed)
+                else:
+                    self._pending_cancel = (time_limit, mark_action_failed)
                 return
             method = self._resolved_cancel_method
         if method is None:  # pragma: no cover - defensive
@@ -957,11 +990,18 @@ def _cancel(
             # pending cancel instead — an early return here rather than an
             # assert, so no bare AssertionError can reach the public API.
             return
-        # Nothing to do if it's not running.
-        if not self._process.is_running:
-            return
 
         with self._lock:
+            # F4 fix: Check liveness under the lock. Without this, a completion
+            # racing a cancel/timeout could see is_running=True outside the lock,
+            # enter here, and then find is_running=False (or worse, still True
+            # but the callback has already fired) — no linearization point.
+            # Moving the check inside the lock lets _on_process_exit's clearing
+            # of _pending_cancel act as the arbiter: if the process exited, any
+            # pending cancel was already consumed there.
+            if not self._process.is_running:
+                return
+
             self._canceled = True
             self._notify_canceled_action_as_failed = mark_action_failed
             now = datetime.now(timezone.utc)
@@ -1015,9 +1055,32 @@ def _cancel(
                 #      when we'll send the SIGKILL)
                 grace_end_time_str = self._cancel_gracetime_end.strftime(TIME_FORMAT_STR)
                 notify_end = json.dumps({"NotifyEnd": grace_end_time_str})
-                write_file_for_user(
-                    self._session_working_directory / "cancel_info.json", notify_end, self._user
-                )
+                try:
+                    write_file_for_user(
+                        self._session_working_directory / "cancel_info.json", notify_end, self._user
+                    )
+                except OSError as err:
+                    # F6 fix: If we cannot write the cancel_info.json (disk full, permission
+                    # denied, etc.), log and fall back to immediate termination. A script
+                    # waiting on that file would hang forever otherwise.
+                    self._logger.warning(
+                        f"Failed to write cancel_info.json: {err}. Falling back to immediate termination.",
+                        extra=LogExtraInfo(
+                            openjd_log_content=LogContent.PROCESS_CONTROL
+                            | LogContent.EXCEPTION_INFO
+                        ),
+                    )
+                    try:
+                        self._process.terminate()
+                    except OSError as term_err:  # pragma: nocover
+                        self._logger.warning(
+                            f"Fallback termination also failed: {term_err}",
+                            extra=LogExtraInfo(
+                                openjd_log_content=LogContent.PROCESS_CONTROL
+                                | LogContent.EXCEPTION_INFO
+                            ),
+                        )
+                    return
                 self._logger.info(
                     f"Grace period ends at {grace_end_time_str}",
                     extra=LogExtraInfo(openjd_log_content=LogContent.PROCESS_CONTROL),
@@ -1046,6 +1109,13 @@ def _on_process_exit(self, future: Future) -> None:
         """This is invoked as a callback when run_future is done."""
         assert self._run_future is not None
         with self._lock:
+            # F2 fix: Claim _pending_cancel atomically before signalling completion.
+            # A cancel racing completion would otherwise see the process still
+            # "running" (in _cancel_with_resolved_method) and hand off to _cancel,
+            # which then finds is_running=False and no-ops. By consuming the
+            # pending here we prevent that lost-cancel window.
+            self._pending_cancel = None
+
             if self._runtime_limit is not None:
                 self._runtime_limit.cancel()
                 self._runtime_limit = None
@@ -1062,8 +1132,20 @@ def _on_process_exit(self, future: Future) -> None:
                     ),
                 )
 
-            if self._callback is not None:
+        # F8 fix: Invoke callback outside the lock and wrap in try/except. An
+        # observer exception must not prevent the child process from being
+        # reaped or cause resource leaks. The callback is invoked outside the
+        # lock since it may be slow and shouldn't block other operations.
+        if self._callback is not None:
+            try:
                 self._callback(ActionState(self.state.value))
+            except Exception as exc:
+                self._logger.error(
+                    f"Exception in action callback: {exc}",
+                    extra=LogExtraInfo(
+                        openjd_log_content=LogContent.PROCESS_CONTROL | LogContent.EXCEPTION_INFO
+                    ),
+                )
 
     def _on_notify_period_end(self) -> None:
         """This is invoked when the grace period in a NOTIFY_THEN_TERMINATE

src/openjd/sessions/_runner_env_script.py

@@ -272,12 +272,14 @@ def _run_wrap_hook(self, hook: str, *, default_timeout: Optional[timedelta] = No
     def cancel(
         self, *, time_limit: Optional[timedelta] = None, mark_action_failed: bool = False
     ) -> None:
-        if self._action is None:
-            # Nothing to do.
-            return
-
-        # Cancel with the effective method resolved at launch time by
-        # _run_action, against the action's own final scope (the script's
-        # lets and Env.File.* symbols, plus WrappedAction.* for a wrap
-        # hook) — openjd-rs parity.
+        # Always route through the base class handoff, even before _action is
+        # assigned. A cancel landing during environment-action setup (embedded-
+        # file writes, let evaluation) must be recorded in _pending_cancel and
+        # applied when the subprocess launches — F1 fix: the old `_action is
+        # None` guard returned early without entering the pending-cancel path,
+        # losing the cancel entirely during the setup window.
+        #
+        # _cancel_with_resolved_method handles the pre-launch case: it records
+        # the request in _pending_cancel when the subprocess hasn't started,
+        # and _run applies it once the child exists.
         self._cancel_with_resolved_method(time_limit, mark_action_failed)

src/openjd/sessions/_session.py

@@ -657,10 +657,18 @@ def cancel_action(
         """
         if self.state != SessionState.RUNNING:
             raise RuntimeError("No actions are running")
-        # For the type checker
-        assert self._runner is not None
+        # Review22-F3 fix: Snapshot _runner before using it. The state check
+        # above and the _runner access below are not atomic; a completion
+        # racing this call could set _runner = None after we pass the state
+        # check. Snapshotting here avoids a bare AssertionError.
+        runner = self._runner
+        if runner is None:
+            # Race: action completed between state check and here. No-op rather
+            # than raise, since the caller's intent (cancel the running action)
+            # is already satisfied.
+            return
 
-        self._runner.cancel(time_limit=time_limit, mark_action_failed=mark_action_failed)
+        runner.cancel(time_limit=time_limit, mark_action_failed=mark_action_failed)
 
     def _make_env_script_runner(
         self,
@@ -2238,6 +2246,17 @@ def _action_callback(self, state: ActionState) -> None:
         self._action_exit_code = self._runner.exit_code
         self._action_state = state
 
+        # F5 fix: Snapshot action_status BEFORE publishing READY. If we set
+        # _state = READY first, another thread polling session.state could see
+        # READY but action_status would still reflect the old (stale or
+        # incomplete) snapshot. By snapshotting here, the callback receives the
+        # definitive ActionStatus that corresponds to the terminal state.
+        #
+        # Note: We snapshot unconditionally (not guarded by `if self._callback`)
+        # because action_status is cheap and some tests check exact callback
+        # invocation patterns including the __bool__ check count.
+        action_status = self.action_status
+
         if state != ActionState.RUNNING:
             # Decide which between-action state to enter.
             if self._ending_only or self._action_state != ActionState.SUCCESS:
@@ -2247,10 +2266,7 @@ def _action_callback(self, state: ActionState) -> None:
             else:
                 self._state = SessionState.READY
 
-        if self._callback:
-            action_status = self.action_status
-            # for the type checker
-            assert action_status is not None
+        if self._callback and action_status is not None:
             self._callback(self._session_id, action_status)
 
     def _evaluate_current_session_env_vars(

src/openjd/sessions/_subprocess.py

@@ -227,7 +227,11 @@ def notify(self) -> None:
         TODO: Send the signal to every direct and transitive child of the parent
         process.
         """
-        if self._process is not None and self._process.poll() is None:
+        # Review22-F4 fix: Bind _process once. Double-load is a TOCTOU race:
+        # the None check and poll() call read _process separately, allowing
+        # another thread to set it to None between them.
+        proc = self._process
+        if proc is not None and proc.poll() is None:
             if is_posix():
                 self._posix_signal_subprocess(signal_name="term")
             else:
@@ -243,15 +247,17 @@ def terminate(self) -> None:
         TODO: Send the signal to every direct and transitive child of the parent
         process.
         """
-        if self._process is not None and self._process.poll() is None:
+        # Review22-F4 fix: Bind _process once. See notify() for rationale.
+        proc = self._process
+        if proc is not None and proc.poll() is None:
             if is_posix():
                 self._posix_signal_subprocess(signal_name="kill")
             else:
                 self._logger.info(
-                    f"INTERRUPT: Start killing the process tree with the root pid: {self._process.pid}",
+                    f"INTERRUPT: Start killing the process tree with the root pid: {proc.pid}",
                     extra=LogExtraInfo(openjd_log_content=LogContent.PROCESS_CONTROL),
                 )
-                kill_windows_process_tree(self._logger, self._process.pid, signal_subprocesses=True)
+                kill_windows_process_tree(self._logger, proc.pid, signal_subprocesses=True)
 
     def _start_subprocess(self) -> Optional[Popen]:
         """Helper invoked by self.run() to start up the subprocess."""