fix(python): share per-object locks and align StartChild timeouts

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

@@ -32,6 +32,7 @@
 )
 from .event import IEvent_2
 from .protocols import IEnumerable_1
+from .system import TimeoutException__ctor
 from .task import TaskCompletionSource
 from .time_span import TimeSpan, to_milliseconds
 from .types import UNIT, Unit
@@ -291,23 +292,29 @@ def cancel(_: OperationCanceledError) -> None:
     return tcs.get_task()
 
 
-def start_child[T](computation: Async[T], ms: int | TimeSpan | None = None) -> Async[Async[T]]:
-    if ms is not None:
-        # Race the computation against a timeout: whichever settles first wins.
-        # asyncio.gather (the previous implementation via parallel2) waited for BOTH to settle,
-        # which meant the timeout always fired even when the computation finished first.
-        task = start_as_task(computation)
+def _start_child_timeout[T](computation: Async[T], ms: int | TimeSpan) -> Async[Async[T]]:
+    # Return the child result when ready, or raise TimeoutException when the deadline expires.
+    task = asyncio.ensure_future(start_as_task(computation))
 
-        async def with_timeout() -> T:
-            try:
-                return await asyncio.wait_for(task, timeout=to_milliseconds(ms) / 1000.0)
-            except TimeoutError:
-                raise TimeoutError()
+    async def with_timeout() -> T:
+        try:
+            return await asyncio.wait_for(task, timeout=to_milliseconds(ms) / 1000.0)
+        except TimeoutError:
+            # invariant: a completed child retains its exception; only deadline cancellation is translated
+            if not task.cancelled():
+                raise
+            raise TimeoutException__ctor()
 
-        def cont(ctx: IAsyncContext[Async[T]]) -> None:
-            protected_return(await_task(with_timeout()))(ctx)
+    def cont(ctx: IAsyncContext[Async[T]]) -> None:
+        protected_return(await_task(with_timeout()))(ctx)
 
-        return protected_cont(cont)
+    return protected_cont(cont)
+
+
+def start_child[T](computation: Async[T], ms: int | TimeSpan | None = None) -> Async[Async[T]]:
+    # decision: share TimeoutException across callers; its native base preserves Python TimeoutError handlers
+    if ms is not None:
+        return _start_child_timeout(computation, ms)
 
     task = start_as_task(computation)
 

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

@@ -29,6 +29,7 @@
     overload,
 )
 from urllib.parse import quote, unquote
+from weakref import WeakValueDictionary
 
 from .array_ import Array
 from .bases import (
@@ -351,16 +352,21 @@ def assert_not_equal[T](actual: T, expected: T, msg: str | None = None) -> None:
         raise Exception(msg or f"Expected not equal to: {expected} - Actual: {actual}")
 
 
-MAX_LOCKS = 1024
+# decision: weak values give each active object its own lock without retaining idle locks or their objects
+_locks: WeakValueDictionary[int, RLock] = WeakValueDictionary()
+_lock_registry_guard = RLock()
 
 
-def lock[T](lock_obj: Any, fn: Callable[[], T]) -> T:
-    @functools.lru_cache(maxsize=MAX_LOCKS)
-    def get_lock(n: int) -> RLock:
-        return RLock()
+def lock[T](lock_obj: object, fn: Callable[[], T]) -> T:
+    key = id(lock_obj)
+    with _lock_registry_guard:
+        object_lock = _locks.get(key)
+        if object_lock is None:
+            object_lock = RLock()
+            _locks[key] = object_lock
 
-    lock = get_lock(id(lock_obj))
-    with lock:
+    # invariant: each holder and waiter keeps a strong lock reference outside the registry guard
+    with object_lock:
         return fn()
 
 

tests/Python/TestAsync.fs

@@ -467,7 +467,7 @@ let ``test Async.StartChild applies timeout`` () =
         let mutable x = ""
 
         let task = async {
-            x <- x + "A"
+            // A busy .NET runner can time out before this child is scheduled.
             do! Async.Sleep 1_000
             x <- x + "X" // Never hit
         }
@@ -482,8 +482,8 @@ let ``test Async.StartChild applies timeout`` () =
 
         x <- x + "C"
 
-        equal x "ABC"
-    } |> Async.StartImmediate
+        equal "BC" x
+    } |> Async.RunSynchronously
 
 [<Fact>]
 let ``test Async.StartChild with timeout completes when computation finishes before timeout`` () = // See #4481
@@ -496,7 +496,74 @@ let ``test Async.StartChild with timeout completes when computation finishes bef
         with
             | :? TimeoutException ->
                 failwith "should not time out"
-    } |> Async.StartImmediate
+    } |> Async.RunSynchronously
+
+let private assertChildFailure startChild (error: exn) =
+    async {
+        let failing: Async<int> = async {
+            do! Async.Sleep 10
+            return raise error
+        }
+        let! child = startChild(failing, 60_000)
+        let! result = Async.Catch child
+        match result with
+        | Choice2Of2 caught -> obj.ReferenceEquals(error, caught) |> equal true
+        | _ -> failwith "Expected child failure"
+    } |> Async.RunSynchronously
+
+[<Fact>]
+let ``test Async.StartChild preserves child TimeoutException`` () =
+    assertChildFailure
+        (fun (computation, timeout) -> Async.StartChild(computation, timeout))
+        (TimeoutException("child timeout"))
+
+let private assertChildTimeout startChild assertError =
+    async {
+        let pending: Async<int> = Async.FromContinuations(ignore)
+        let! child = startChild(pending, 0)
+        let! result = Async.Catch child
+        match result with
+        | Choice2Of2 ex -> assertError ex
+        | _ -> failwith "Expected child timeout"
+    } |> Async.RunSynchronously
+
+[<Fact>]
+let ``test Async.StartChild zero timeout raises TimeoutException`` () =
+    assertChildTimeout
+        (fun (computation, timeout) -> Async.StartChild(computation, timeout))
+        (fun ex -> ex :? TimeoutException |> equal true)
+
+#if FABLE_COMPILER
+module private NativeTimeout =
+    open Fable.Core
+    open Fable.Core.PyInterop
+
+    let startChild: Func<Async<int>, int, Async<Async<int>>> = import "start_child" "fable_library.async_"
+    let timeoutError: obj = import "TimeoutError" "builtins"
+
+    [<Emit("TimeoutError($0)")>]
+    let createTimeoutError (message: string) : exn = nativeOnly
+
+[<Fact>]
+let ``test Python start_child deadline matches both timeout types`` () =
+    assertChildTimeout
+        (fun (computation, timeout) -> NativeTimeout.startChild.Invoke(computation, timeout))
+        (fun ex ->
+            ex :? TimeoutException |> equal true
+            Fable.Core.PyInterop.pyInstanceof ex NativeTimeout.timeoutError |> equal true)
+
+[<Fact>]
+let ``test Async.StartChild preserves child Python TimeoutError`` () =
+    assertChildFailure
+        (fun (computation, timeout) -> Async.StartChild(computation, timeout))
+        (NativeTimeout.createTimeoutError "child timeout")
+
+[<Fact>]
+let ``test Python start_child preserves child TimeoutError`` () =
+    assertChildFailure
+        (fun (computation, timeout) -> NativeTimeout.startChild.Invoke(computation, timeout))
+        (NativeTimeout.createTimeoutError "child timeout")
+#endif
 
 [<Fact>]
 let ``test Unit arguments are erased`` () = // See #1832

tests/Python/TestMisc.fs

@@ -1288,7 +1288,7 @@
     | 2 -> x <- "2"
     | 3 | 4 -> x <- "3" // Multiple cases are allowed
     // | 5 | 6 as j -> x <- string j // This prevents the optimization
     | 4 -> x <- "4" // Unreachable cases are removed
     | _ -> x <- "?"
     equal "3" x
 
@@ -1677,3 +1677,111 @@
         with ex ->
             (fun (_: int) -> ex.Message)  // lambda-with-arg captures ex
     sub 42 |> equal "boom"
+
+[<Fact>]
+let ``test lock is reentrant and returns the body value`` () =
+    let key = obj()
+    lock key (fun () -> lock key (fun () -> 42)) |> equal 42
+
+[<Fact>]
+let ``test lock releases after a body exception`` () =
+    let key = obj()
+    let mutable error = ""
+    try
+        lock key (fun () -> failwith "body")
+    with ex ->
+        error <- ex.Message
+    equal "body" error
+    lock key (fun () -> 42) |> equal 42
+
+#if FABLE_COMPILER
+module private LockRuntime =
+    open Fable.Core.PyInterop
+
+    type Threads =
+        abstract trace: string array
+        abstract errors: string array
+        abstract first_entered: bool
+        abstract second_attempted: bool
+        abstract second_blocked: bool
+        abstract threads_finished: bool
+        abstract record: string -> unit
+        abstract wait_release: unit -> bool
+        abstract run: Action * Action -> unit
+
+    type Storage =
+        abstract stable: bool
+        abstract count: int
+
+    let makeThreads: Func<Threads> = import "make_threads" "./py/lock_runtime.py"
+    let makeStorage: Func<Storage> = import "make_storage" "./py/lock_runtime.py"
+    let makeKey: Func<obj> = import "make_key" "./py/lock_runtime.py"
+    let makeCollidingKeys: Func<obj array> = import "make_colliding_keys" "./py/lock_runtime.py"
+
+[<Fact>]
+let ``test real threads share a lock and release after an exception`` () =
+    let key = LockRuntime.makeKey.Invoke()
+    let threads = LockRuntime.makeThreads.Invoke()
+    let mutable error = ""
+    let first = Action(fun () ->
+        try
+            lock key (fun () ->
+                threads.record "first-enter"
+                threads.wait_release() |> equal true
+                threads.record "first-exit"
+                failwith "body")
+        with ex ->
+            error <- ex.Message)
+    let second = Action(fun () ->
+        lock key (fun () -> threads.record "second-enter")
+        threads.record "second-exit")
+    threads.run(first, second)
+    equal true threads.first_entered
+    equal true threads.second_attempted
+    equal true threads.second_blocked
+    equal true threads.threads_finished
+    equal [||] threads.errors
+    equal [| "first-enter"; "first-exit"; "second-enter"; "second-exit" |] threads.trace
+    equal "body" error
+    lock key (fun () -> 42) |> equal 42
+
+[<Fact>]
+let ``test independent objects do not share a lock`` () =
+    let keys = LockRuntime.makeCollidingKeys.Invoke()
+    let threads = LockRuntime.makeThreads.Invoke()
+    let first = Action(fun () ->
+        lock keys.[0] (fun () ->
+            threads.record "first-enter"
+            threads.wait_release() |> equal true
+            threads.record "first-exit"))
+    let second = Action(fun () ->
+        lock keys.[1] (fun () -> threads.record "second-enter")
+        threads.record "second-exit")
+    threads.run(first, second)
+    equal true threads.first_entered
+    equal true threads.second_attempted
+    equal false threads.second_blocked
+    equal true threads.threads_finished
+    equal [||] threads.errors
+    equal 4 threads.trace.Length
+
+[<Fact>]
+let ``test lock storage releases idle entries across new objects`` () =
+    let iterations = 2048
+    let storage = LockRuntime.makeStorage.Invoke()
+    let before = storage.count
+    lock (LockRuntime.makeKey.Invoke()) (fun () -> equal (before + 1) storage.count)
+    for _ in 1 .. iterations do
+        lock (LockRuntime.makeKey.Invoke()) (fun () -> 1) |> equal 1
+    equal true storage.stable
+    equal before storage.count
+
+[<Fact>]
+let ``test lock storage releases idle entries after exceptions`` () =
+    let storage = LockRuntime.makeStorage.Invoke()
+    let before = storage.count
+    try
+        lock (LockRuntime.makeKey.Invoke()) (fun () -> failwith "body")
+    with _ -> ()
+    equal before storage.count
+#endif