use non-cancellable pause helper everywhere

src/internal/coroutine/coroutine.mbt

@@ -95,6 +95,24 @@ pub async fn pause() -> Unit noraise {
   }
 }
 
+///|
+/// Helper for a non-cancellable `pause`, more efficient than `protect_from_cancel`
+pub async fn pause_nocancel() -> Unit noraise + nocancel {
+  guard! scheduler.curr_coro is Some(coro)
+  let shielded = coro.shielded
+  coro.shielded = true
+  defer {
+    coro.shielded = shielded
+  }
+  let result = async_suspend() <| cont => {
+    guard! coro.state is Running
+    coro.state = Suspend(cont)
+    coro.ready = true
+    scheduler.run_later.push_back(coro)
+  }
+  guard! result is Continue
+}
+
 ///|
 pub async fn suspend_check_cancel() -> SuspendResult noraise {
   guard! scheduler.curr_coro is Some(coro)
@@ -211,22 +229,3 @@ pub async fn[X] protect_from_cancel(f : async () -> X) -> X nocancel {
 pub async fn[X] handle_cancellation(f : async () -> X) -> X? {
   capture_cancellation(handler=() => None, () => Some(f()))
 }
-
-///|
-/// Yield after an I/O operation has made progress. Preserve cancellation for
-/// the caller's next cancellation point without losing the completed result.
-pub async fn pause_after_io() -> Unit noraise + nocancel {
-  guard! scheduler.curr_coro is Some(coro)
-  let shielded = coro.shielded
-  coro.shielded = true
-  defer {
-    coro.shielded = shielded
-  }
-  let result = async_suspend() <| cont => {
-    guard! coro.state is Running
-    coro.state = Suspend(cont)
-    coro.ready = true
-    scheduler.run_later.push_back(coro)
-  }
-  guard! result is Continue
-}

src/internal/coroutine/pause_wbtest.mbt

@@ -13,18 +13,18 @@
 // limitations under the License.
 
 ///|
-test "IO pause preserves round robin scheduling" {
+test "non-cancellable pause preserves round robin scheduling" {
   let log : Array[Int] = []
   let a = spawn(() => {
     for _ in 0..<3 {
       log.push(1)
-      pause_after_io()
+      pause_nocancel()
     }
   })
   let b = spawn(() => {
     for _ in 0..<3 {
       log.push(2)
-      pause_after_io()
+      pause_nocancel()
     }
   })
   while has_immediately_ready_task() {
@@ -36,64 +36,60 @@ test "IO pause preserves round robin scheduling" {
 }
 
 ///|
-test "IO pause delivers completed IO then preserves cancellation" {
+test "non-cancellable pause completes then preserves cancellation" {
   let log : Array[String] = []
-  let a = spawn(() => {
+  let _ = spawn(() => {
+    let self = current_coroutine()
     log.push("started")
-    pause_after_io()
-    log.push("delivered")
-    assert_false(current_coroutine().shielded)
-    check_cancellation()
-    log.push("unreachable")
+    self.cancel()
+    pause_nocancel()
+    log.push("completed")
+    assert_false(self.shielded)
+    assert_true(self.cancelled)
   })
-  a.cancel()
   while has_immediately_ready_task() {
     reschedule()
   }
-  guard! a.unwrap() is Cancelled
-  assert_eq(log, ["started", "delivered"])
+  json_inspect(log, content=["started", "completed"])
 }
 
 ///|
-test "IO pause restores an enclosing cancellation shield" {
+test "non-cancellable pause restores an enclosing cancellation shield" {
   let log : Array[String] = []
-  let a = spawn(() => {
+  let _ = spawn(() => {
+    let self = current_coroutine()
     protect_from_cancel(() => {
-      current_coroutine().cancel()
-      pause_after_io()
-      assert_true(current_coroutine().shielded)
-      assert_true(current_coroutine().cancelled)
-      log.push("delivered")
+      self.cancel()
+      pause_nocancel()
+      log.push("completed")
+      assert_true(self.shielded)
+      assert_true(self.cancelled)
     })
-    assert_false(current_coroutine().shielded)
-    check_cancellation()
-    log.push("unreachable")
+    assert_false(self.shielded)
+    assert_true(self.cancelled)
   })
   while has_immediately_ready_task() {
     reschedule()
   }
-  guard! a.unwrap() is Cancelled
-  assert_eq(log, ["delivered"])
+  assert_eq(log, ["completed"])
 }
 
 ///|
-test "IO pause delivers completed IO when cancelled during suspension" {
+test "non-cancellable pause cancelled during suspension" {
   let log : Array[String] = []
   let a = spawn(() => {
     log.push("started")
-    pause_after_io()
-    log.push("delivered")
+    pause_nocancel()
+    log.push("completed")
     check_cancellation()
     log.push("unreachable")
   })
-  let b = spawn(() => {
-    assert_eq(log, ["started"])
-    a.cancel()
-  })
+  reschedule()
+  assert_eq(log, ["started"])
+  a.cancel()
   while has_immediately_ready_task() {
     reschedule()
   }
   guard! a.unwrap() is Cancelled
-  guard! b.unwrap() is Continue
-  assert_eq(log, ["started", "delivered"])
+  assert_eq(log, ["started", "completed"])
 }

src/internal/coroutine/pkg.generated.mbti

@@ -20,7 +20,7 @@ pub fn is_being_cancelled() -> Bool
 
 pub async fn pause() -> Unit noraise
 
-pub async fn pause_after_io() -> Unit noraise + nocancel
+pub async fn pause_nocancel() -> Unit noraise + nocancel
 
 pub async fn[X] protect_from_cancel(async () -> X) -> X nocancel
 

src/internal/event_loop/fs.mbt

@@ -536,7 +536,7 @@ pub async fn IoHandle::read_dir_changes(
   if n >= 0 {
     // The `pause` here prevent a busy loop from exhausting the whole program,
     // and give chance for other tasks to execute.
-    @coroutine.protect_from_cancel(() => @coroutine.pause())
+    @coroutine.pause_nocancel()
     n
   } else if @os_error.is_nonblocking_io_error() {
     handle.wait_read_result(result, context~)

src/internal/event_loop/io_unix.mbt

@@ -77,7 +77,7 @@ async fn IoHandle::read_via_event_bus_unix(
   if ret >= 0 {
     // The `pause` here prevent a busy read loop from exhausting the whole program
     // and give chance for other tasks to execute.
-    @coroutine.pause_after_io()
+    @coroutine.pause_nocancel()
     return ret
   }
   let ret = if @os_error.is_nonblocking_io_error() {
@@ -138,7 +138,7 @@ async fn IoHandle::write_via_event_bus_unix(
   if ret >= 0 {
     // The `pause` here prevent a busy write loop from exhausting the whole program,
     // and give chance for other tasks to execute.
-    @coroutine.pause_after_io()
+    @coroutine.pause_nocancel()
     return ret
   }
   let ret = if @os_error.is_nonblocking_io_error() {

src/internal/event_loop/io_windows.mbt

@@ -124,7 +124,7 @@ async fn IoHandle::generic_read_windows(
   let n = if n >= 0 {
     // The `pause` here prevent a busy read loop from exhausting the whole program,
     // and give chance for other tasks to execute.
-    @coroutine.protect_from_cancel(() => @coroutine.pause())
+    @coroutine.pause_nocancel()
     // It is important to note that the logic here is correct
     // only if `FILE_SKIP_COMPLETION_PORT_ON_SUCCESS` is set
     // via `SetFileCompletionNotificationModes`.
@@ -203,7 +203,7 @@ async fn IoHandle::generic_write_windows(
   if n >= 0 {
     // The `pause` here prevent a busy write loop from exhausting the whole program,
     // and give chance for other tasks to execute.
-    @coroutine.protect_from_cancel(() => @coroutine.pause())
+    @coroutine.pause_nocancel()
     // It is important to note that the logic here is correct
     // only if `FILE_SKIP_COMPLETION_PORT_ON_SUCCESS` is set
     // via `SetFileCompletionNotificationModes`.

src/internal/event_loop/network_unix.mbt

@@ -60,7 +60,7 @@ async fn IoHandle::recvfrom_unix(
   if ret >= 0 {
     // The `pause` here prevent a busy loop from exhausting the whole program,
     // and give chance for other tasks to execute.
-    @coroutine.protect_from_cancel(() => @coroutine.pause())
+    @coroutine.pause_nocancel()
     return ret
   }
   let ret = if @os_error.is_nonblocking_io_error() {
@@ -125,7 +125,7 @@ async fn IoHandle::sendto_unix(
 ) -> Int {
   let ret = sendto_unix_ffi(handle.fd, buf, offset~, len~, addr~)
   if ret >= 0 {
-    @coroutine.protect_from_cancel(() => @coroutine.pause())
+    @coroutine.pause_nocancel()
     return ret
   }
   let ret = if @os_error.is_nonblocking_io_error() {
@@ -242,7 +242,7 @@ pub async fn IoHandle::accept_unix(
   let conn = if @fd_util.fd_is_valid(conn) {
     // The `pause` here prevent a busy `accept` loop from exhausting the whole program,
     // and give chance for other tasks to execute.
-    @coroutine.protect_from_cancel(() => @coroutine.pause())
+    @coroutine.pause_nocancel()
     conn
   } else if @os_error.is_nonblocking_io_error() {
     handle.wait_read()

src/io/pipe.mbt

@@ -123,7 +123,7 @@ pub impl Reader for PipeRead with fn _direct_read(self, dst, offset~, max_len~)
       // The pause here is not cancellable,
       // because otherwise the writer cannot have correct knowledge on what is written,
       // making writing not cancellation safe.
-      @coroutine.protect_from_cancel(() => @coroutine.pause())
+      @coroutine.pause_nocancel()
       len
     }
   }
@@ -186,7 +186,7 @@ pub impl Writer for PipeWrite with fn write_once(self, buf, offset~, len~) {
       // The pause here is not cancellable,
       // because otherwise the writer cannot have correct knowledge on what is written,
       // making writing not cancellation safe.
-      @coroutine.protect_from_cancel(() => @coroutine.pause())
+      @coroutine.pause_nocancel()
       len
     }
   }