feat(js/ts/python/beam): support Async.TryCancelled

src/Fable.Transforms/Python/Fable2Python.Transforms.fs

@@ -1718,7 +1718,12 @@ let transformSet (com: IPythonCompiler) ctx range fableExpr typ (value: Fable.Ex
 let transformBindingExprBody (com: IPythonCompiler) (ctx: Context) (var: Fable.Ident) (value: Fable.Expr) =
     match value with
     | Function(args, body) ->
-        let name = Some var.Name
+        // decision: mutable callbacks keep their Callable annotation instead of inheriting a def's parameter names
+        let name =
+            if var.IsMutable then
+                None
+            else
+                Some var.Name
 
         let args', body', returnType, typeParams =
             Annotation.transformFunctionWithAnnotations com ctx name args body

src/fable-library-beam/fable_async.erl

@@ -9,6 +9,7 @@
     parallel/1,
     sequential/1,
     catch_async/1,
+    try_cancelled/2,
     ignore/1,
     from_continuations/1,
     start_as_task/1,
@@ -43,6 +44,7 @@
 -spec parallel(list() | reference()) -> async(reference()).
 -spec sequential(list()) -> async(list()).
 -spec catch_async(async(term())) -> async(term()).
+-spec try_cancelled(async(term()), fun((term()) -> term())) -> async(term()).
 -spec ignore(async(term())) -> async(ok).
 -spec from_continuations(fun()) -> async(term()).
 -spec start_as_task(async(term())) -> term().
@@ -479,6 +481,21 @@ catch_async(Computation) ->
         end
     end.
 
+%% TryCancelled: observe cancellation without changing its terminal outcome.
+try_cancelled(Computation, Compensation) ->
+    fun(Ctx) ->
+        OnCancel = maps:get(on_cancel, Ctx),
+        Ctx2 = Ctx#{
+            on_cancel => fun(Error) ->
+                %% decision: preserve the original cancellation even when compensation throws, matching .NET.
+                try Compensation(Error) catch _:_ -> ok end,
+                OnCancel(Error)
+            end
+        },
+        %% invariant: install the handler without adding a cancellation checkpoint, including for nested handlers.
+        Computation(Ctx2)
+    end.
+
 %% Ignore: discard result
 ignore(Computation) ->
     fable_async_builder:bind(Computation, fun(_) -> fable_async_builder:return(ok) end).

src/fable-library-py/fable_library/async_.py

@@ -215,6 +215,23 @@ def _arrow21(__unit: Unit = UNIT) -> Async[Array[T]]:
     return delay(delayed)
 
 
+def try_cancelled[T](computation: Async[T], compensation: Callable[[OperationCanceledError], None]) -> Async[T]:
+    def cont(ctx: IAsyncContext[T]) -> None:
+        def on_cancel(error: OperationCanceledError) -> None:
+            # decision: preserves the original cancellation when compensation throws, matching .NET
+            try:
+                compensation(error)
+            except Exception:
+                pass
+            ctx.trampoline.run(lambda: ctx.on_cancel(error))
+
+        ctx_ = IAsyncContext.create(ctx.trampoline, ctx.cancel_token, ctx.on_success, ctx.on_error, on_cancel)
+        # invariant: installs the handler without adding a cancellation checkpoint, including for nested handlers
+        computation(ctx_)
+
+    return cont
+
+
 def catch_async[T](work: Async[T]) -> Async[FSharpChoice_2[T, Exception]]:
     def cont(ctx: IAsyncContext[FSharpChoice_2[T, Exception]]) -> None:
         def on_success(x: T):
@@ -528,4 +545,5 @@ def exception_handler(loop: asyncio.AbstractEventLoop, context: dict[str, Any])
     "start_child",
     "start_immediate",
     "start_with_continuations",
+    "try_cancelled",
 ]

src/fable-library-ts/Async.ts

@@ -127,6 +127,24 @@ export function catchAsync<T>(work: Async<T>) {
   });
 }
 
+export function tryCancelled<T>(computation: Async<T>, compensation: (error: OperationCanceledException) => void): Async<T> {
+  return (ctx) => {
+    // invariant: install the handler without adding a cancellation checkpoint, including for nested handlers
+    computation({
+      ...ctx,
+      onCancel: (error) => {
+        // decision: preserve the original cancellation even when compensation throws, matching .NET
+        try { compensation(error); } catch { /* Cancellation takes precedence. */ }
+        if (ctx.trampoline.incrementAndCheck()) {
+          ctx.trampoline.hijack(() => ctx.onCancel(error));
+        } else {
+          ctx.onCancel(error);
+        }
+      },
+    });
+  };
+}
+
 export function fromContinuations<T>(f: (conts: Continuations<T>) => void) {
   return (ctx: IAsyncContext<T>) => {
     let settled = false;

tests/Beam/AsyncTests.fs

@@ -5,8 +5,160 @@ open Fable.Tests.Util
 open Util.Testing
 open System.Threading
 
+let private runOffThreadPool action =
+#if FABLE_COMPILER
+    action ()
+#else
+    // .NET reschedules deferred FromContinuations callbacks invoked on a pool thread.
+    // Use a dedicated thread so the probe can observe completion and thrown callbacks.
+    let mutable failure = None
+    let thread = Thread(ThreadStart(fun () ->
+        try action ()
+        with error -> failure <- Some error))
+    thread.Start()
+    thread.Join()
+    failure |> Option.iter raise
+#endif
+
+let private observeTryCancelled token (work: Async<int>) =
+    let calls = ResizeArray<string>()
+    Async.StartWithContinuations(
+        work,
+        (fun value -> calls.Add(string value)),
+        (fun error -> calls.Add("error: " + error.Message)),
+        (fun error -> calls.Add("cancel: " + error.Message)),
+        token)
+    calls
+
 let private raiseCaughtCancellation (error: exn) : unit = raise error
 
+[<Fact>]
+let ``test TryCancelled preserves success and ordinary errors`` () =
+    let mutable compensations = 0
+    let wrap work = Async.TryCancelled(work, fun _ -> compensations <- compensations + 1)
+    equal [|"42"|] ((observeTryCancelled CancellationToken.None (wrap (async.Return 42))).ToArray())
+    for error in [Exception "body"; OperationCanceledException "thrown"] do
+        let work: Async<int> = async { return raise error }
+        equal [|"error: " + error.Message|] ((observeTryCancelled CancellationToken.None (wrap work)).ToArray())
+    let mutable completed = false
+    Async.StartWithContinuations(
+        Async.TryCancelled(async.Return (), fun _ -> compensations <- compensations + 1),
+        (fun () -> completed <- true), raise, raise)
+    equal true completed
+    equal 0 compensations
+
+[<Fact>]
+let ``test TryCancelled preserves cancellation identity and supports reuse`` () =
+    for throws in [false; true] do
+        let original = OperationCanceledException "original"
+        let calls = ResizeArray<string>()
+        let mutable handlerSawOriginal = true
+        let body: Async<int> = Async.FromContinuations(fun (_, _, cancel) -> cancel original)
+        let work = Async.TryCancelled(body, fun error ->
+            handlerSawOriginal <- handlerSawOriginal && obj.ReferenceEquals(original, error)
+            calls.Add "compensation"
+            if throws then failwith "compensation defect")
+        for _ in 1 .. 2 do
+            Async.StartWithContinuations(work,
+                (fun _ -> failwith "unexpected success"),
+                (fun _ -> failwith "unexpected error"),
+                (fun error ->
+                    equal true (obj.ReferenceEquals(original, error))
+                    calls.Add "cancel"))
+        equal [|"compensation"; "cancel"; "compensation"; "cancel"|] (calls.ToArray())
+        equal true handlerSawOriginal
+
+[<Fact>]
+let ``test TryCancelled installs nested handlers before cancellation`` () =
+    use source = new CancellationTokenSource()
+    source.Cancel()
+    let mutable entered = false
+    let calls = ResizeArray<string>()
+    let body = async { entered <- true; return 42 }
+    let inner = Async.TryCancelled(body, fun _ -> calls.Add "inner")
+    let outer = Async.TryCancelled(inner, fun _ -> calls.Add "outer")
+    Async.StartWithContinuations(outer,
+        (fun _ -> failwith "unexpected success"), raise,
+        (fun _ -> calls.Add "cancel"), source.Token)
+    equal false entered
+    equal [|"inner"; "outer"; "cancel"|] (calls.ToArray())
+
+[<Fact>]
+let ``test TryCancelled handles deferred cancellation`` () =
+    runOffThreadPool (fun () ->
+        let original = OperationCanceledException "deferred"
+        let mutable cancel = fun (_: OperationCanceledException) -> failwith "not started"
+        let calls = ResizeArray<string>()
+        let body: Async<int> = Async.FromContinuations(fun (_, _, continuation) -> cancel <- continuation)
+        let work = Async.TryCancelled(body, fun _ -> calls.Add "compensation")
+        Async.StartWithContinuations(work,
+            (fun _ -> failwith "unexpected success"), raise,
+            (fun error -> equal true (obj.ReferenceEquals(original, error)); calls.Add "cancel"))
+        equal 0 calls.Count
+        cancel original
+        equal [|"compensation"; "cancel"|] (calls.ToArray()))
+
+[<Fact>]
+let ``test TryCancelled preserves finalizer order when handlers throw`` () =
+    let calls = ResizeArray<string>()
+    let body: Async<int> = Async.FromContinuations(fun (_, _, cancel) -> cancel (OperationCanceledException "original"))
+    let finalized = async.TryFinally(body, fun () -> calls.Add "finally")
+    let inner = Async.TryCancelled(finalized, fun _ -> calls.Add "inner"; failwith "inner defect")
+    let outer = Async.TryCancelled(inner, fun _ -> calls.Add "outer"; failwith "outer defect")
+    equal [|"cancel: original"|] ((observeTryCancelled CancellationToken.None outer).ToArray())
+    equal [|"finally"; "inner"; "outer"|] (calls.ToArray())
+
+[<Fact>]
+let ``test TryCancelled propagates downstream errors once`` () =
+    runOffThreadPool (fun () ->
+        for channel in ["success"; "error"; "cancel"] do
+            let mutable compensations = 0
+            let calls = ResizeArray<string>()
+            let mutable complete = fun () -> failwith "not started"
+            let body: Async<int> = Async.FromContinuations(fun (ok, error, cancel) ->
+                complete <- fun () ->
+                    match channel with
+                    | "success" -> ok 42
+                    | "error" -> error (Exception "body")
+                    | _ -> cancel (OperationCanceledException "original"))
+            let work = Async.TryCancelled(body, fun _ -> compensations <- compensations + 1)
+            let terminal name = calls.Add name; failwith "downstream"
+            Async.StartWithContinuations(work,
+                (fun _ -> terminal "success"),
+                (fun _ -> terminal "error"),
+                (fun _ -> terminal "cancel"))
+            let caught =
+                try
+                    complete ()
+                    "no exception"
+                with error -> error.Message
+            equal "downstream" caught
+            equal [|channel|] (calls.ToArray())
+            equal (if channel = "cancel" then 1 else 0) compensations)
+
+[<Fact>]
+let ``test TryCancelled preserves bounded trampoline execution`` () =
+    async {
+        let original = OperationCanceledException "deep cancellation"
+        let mutable compensations = 0
+        let rec loop remaining =
+            if remaining = 0 then
+                Async.FromContinuations(fun (_, _, cancel) -> cancel original)
+            else
+                Async.TryCancelled(async {
+                    do! async.Return ()
+                    return! loop (remaining - 1)
+                }, fun _ -> compensations <- compensations + 1)
+        do! Async.FromContinuations(fun (ok, error, _) ->
+            Async.StartWithContinuations(loop 10000,
+                (fun () -> error (Exception "unexpected success")), error,
+                (fun cancelled ->
+                    if obj.ReferenceEquals(original, cancelled) then ok ()
+                    else error (Exception "cancellation identity changed"))))
+        equal 10000 compensations
+    }
+    |> Async.RunSynchronously
+
 [<Fact>]
 let ``test RunSynchronously optional timeout and cancellation token`` () =
     let work = async { return 42 }