Fix broken pipe crash in task sessions

NEWS.md

@@ -2,6 +2,8 @@
 
 - Retire persistent `callr` workers when cancelling tasks so a late interrupt
   cannot leak into the next diagnostics run.
+- Prevent a dead persistent worker from crashing the language server when a
+  task is dispatched to its closed command pipe.
 - Add preferred quick fixes for common `lintr` diagnostics and a conflict-aware
   `source.fixAll` action, while making multi-line `nolint` actions apply to every
   affected line.

R/task.R

@@ -86,15 +86,26 @@ Task <- R6::R6Class("Task",
                 TRUE
             }
         },
+        fail = function(error) {
+            if (!private$cancelled && !is.null(private$error)) {
+                private$error(error)
+            }
+            invisible(NULL)
+        },
         kill = function() {
             private$cancelled <- TRUE
+            retired_session <- NULL
             if (!is.null(private$session)) {
                 if (!identical(Sys.getenv("R_COVR"), "true")) {
                     # An interrupt can arrive after this call has completed and
                     # interrupt the next task on the persistent session. Retire
                     # the worker instead so cancellation cannot leak across tasks.
-                    private$session$kill(
-                        grace = 0, close_connections = FALSE)
+                    retired_session <- private$session
+                    tryCatch(
+                        private$session$kill(
+                            grace = 0, close_connections = FALSE),
+                        error = function(e) NULL
+                    )
                 }
             } else if (!is.null(private$process) && private$process$is_alive()) {
                 if (identical(Sys.getenv("R_COVR"), "true")) {
@@ -103,6 +114,7 @@ Task <- R6::R6Class("Task",
                     private$process$kill()
                 }
             }
+            invisible(retired_session)
         }
     )
 )
@@ -120,6 +132,20 @@ TaskManager <- R6::R6Class("TaskManager",
         min_idle_sessions = NULL,
         cancelled_tasks = NULL,
         stopping = FALSE,
+        remove_session = function(session) {
+            keep <- !vapply(
+                private$sessions, identical, logical(1L), y = session)
+            private$sessions <- private$sessions[keep]
+        },
+        retire_session = function(session) {
+            private$remove_session(session)
+            if (!identical(Sys.getenv("R_COVR"), "true")) {
+                tryCatch(
+                    session$kill(grace = 0),
+                    error = function(e) NULL
+                )
+            }
+        },
         create_session = function() {
             session <- callr::r_session$new(
                 options = callr::r_session_options(
@@ -243,7 +269,13 @@ TaskManager <- R6::R6Class("TaskManager",
             }
             if (private$running_tasks$has(id)) {
                 old_task <- private$running_tasks$pop(id)
-                old_task$kill()
+                retired_session <- old_task$kill()
+                if (!is.null(retired_session)) {
+                    # A killed session can still report a buffered result and
+                    # transition back to "idle". Remove it before that stale
+                    # state makes it eligible for another task.
+                    private$remove_session(retired_session)
+                }
                 private$cancelled_tasks <- append(
                     private$cancelled_tasks, old_task)
             }
@@ -286,7 +318,20 @@ TaskManager <- R6::R6Class("TaskManager",
 
                 task <- private$pending_tasks$pop(id)
                 private$running_tasks$set(id, task)
-                task$start(session)
+                start_error <- tryCatch(
+                    {
+                        task$start(session)
+                        NULL
+                    },
+                    error = function(e) e
+                )
+                if (!is.null(start_error)) {
+                    private$running_tasks$remove(id)
+                    if (!is.null(session)) {
+                        private$retire_session(session)
+                    }
+                    task$fail(start_error)
+                }
             }
         },
         check_tasks = function() {

tests/testthat/test-task.R

@@ -184,6 +184,69 @@ test_that("Task cancellation retires its persistent session", {
     expect_false(interrupted)
 })
 
+test_that("TaskManager removes a cancelled session from the pool", {
+    withr::local_envvar(R_COVR = "false")
+    state <- "idle"
+    session <- list(
+        call = function(...) state <<- "busy",
+        get_state = function() state,
+        read = function() {
+            if (state != "busy") return(NULL)
+            state <<- "idle"
+            list(code = 200L, error = NULL, result = NULL)
+        },
+        kill = function(...) NULL,
+        close = function(...) NULL
+    )
+    tm <- TaskManager$new(
+        "cancelled session", use_session = TRUE, min_idle_sessions = 0,
+        max_running_tasks = 1L, cpu_load = 1
+    )
+    private <- tm$.__enclos_env__$private
+    private$sessions <- list(session)
+    tm$add_task("doc", create_task(function() NULL, list()))
+    tm$run_tasks()
+
+    tm$add_task("doc", create_task(function() NULL, list()))
+    tm$check_tasks()
+
+    expect_equal(state, "idle")
+    expect_length(private$sessions, 0L)
+    expect_true(private$pending_tasks$has("doc"))
+    tm$stop()
+})
+
+test_that("TaskManager handles errors while dispatching to a session", {
+    withr::local_envvar(R_COVR = "false")
+    killed <- FALSE
+    session <- list(
+        call = function(...) stop("broken pipe"),
+        get_state = function() "idle",
+        read = function() NULL,
+        kill = function(...) killed <<- TRUE,
+        close = function(...) NULL
+    )
+    tm <- TaskManager$new(
+        "dispatch error", use_session = TRUE, min_idle_sessions = 0,
+        max_running_tasks = 1L, cpu_load = 1
+    )
+    private <- tm$.__enclos_env__$private
+    private$sessions <- list(session)
+    task_error <- NULL
+    tm$add_task("doc", create_task(
+        function() NULL, list(),
+        error = function(e) task_error <<- e
+    ))
+
+    expect_silent(tm$run_tasks())
+
+    expect_match(conditionMessage(task_error), "broken pipe")
+    expect_true(killed)
+    expect_length(private$sessions, 0L)
+    expect_false(private$running_tasks$has("doc"))
+    tm$stop()
+})
+
 test_that("TaskManager does not overprovision while a session starts", {
     tm <- TaskManager$new(
         "starting", use_session = TRUE, min_idle_sessions = 0,