AsyncLazy bugfix (#18101)

majocha · web-flow · commit b204b964dab7 · 2024-12-06T10:45:15.000+01:00
diff --git a/src/Compiler/Facilities/AsyncMemoize.fs b/src/Compiler/Facilities/AsyncMemoize.fs
@@ -22,20 +22,24 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
 
     let stateUpdateSync = obj()
     let mutable state = initial
-
-    let cancelIfUnawaited () =
-        match state with
-        | Running(computation, _, cts, 0) ->
+    // This should remain the only function that mutates the state.
+    let withStateUpdate f =
+        lock stateUpdateSync <| fun () ->
+            let next, result = f state
+            state <- next
+            result
+    let updateState f = withStateUpdate <| fun prev -> f prev, ()
+
+    let cancelIfUnawaited cancelUnawaited = function
+        | Running(computation, _, cts, 0) when cancelUnawaited ->
+            // To keep state updates fast we don't actually wait for the work to cancel.
             cts.Cancel()
-            state <- Initial computation
-        | _ -> ()
+            Initial computation
+        | state -> state
 
-    let afterRequest () =
-        match state with
-        | Running(c, work, cts, count) ->
-            state <- Running(c, work, cts, count - 1)
-            if cancelUnawaited then cancelIfUnawaited () 
-        | _ -> () // Nothing more to do if state already transitioned.
+    let afterRequest = function
+        | Running(c, work, cts, count) -> Running(c, work, cts, count - 1) |> cancelIfUnawaited cancelUnawaited
+        | state -> state // Nothing more to do if state already transitioned.
 
     let detachable (work: Task<'t>) =
         async {
@@ -53,32 +57,35 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
             // The cancellation continuation will always be called in case of cancellation.
             with exn -> return raise exn
         }
-    
-    let workCompleted computation (work: Task<_>) =
-        lock stateUpdateSync <| fun () ->
-            state <-
-                try
-                    Completed work.Result
-                with
-                | exn ->
-                    if cacheException then Faulted exn else Initial computation
 
-    let request () =
-        match state with
+    let request = function
         | Initial computation ->
             let cts = new CancellationTokenSource()
-            let work = Async.StartAsTask(computation, cancellationToken = cts.Token)
-            // Ensure state is updated even when not awaited.
-            work.ContinueWith(workCompleted computation, TaskContinuationOptions.NotOnCanceled) |> ignore
-            state <- Running (computation, work, cts, 1)
+            let work = Async.StartAsTask( async {
+                try
+                    let! result = computation
+                    // If associated cts is signalled it means this work item was abandoned
+                    // and it should not alter the state.
+                    updateState <| function
+                        | state when cts.IsCancellationRequested -> state
+                        | _ -> Completed result
+                    return result
+                with
+                | exn ->
+                    updateState <| function
+                        | state when cts.IsCancellationRequested -> state
+                        | _ -> if cacheException then Faulted exn else Initial computation
+                    return raise exn
+            }, cancellationToken = cts.Token)
+            Running (computation, work, cts, 1),
             detachable work
         | Running (c, work, cts, count) ->
-            state <- Running (c, work, cts, count + 1)
+            Running (c, work, cts, count + 1),
             detachable work
-        | Completed result ->
-            async { return result }
-        | Faulted exn ->
-            async { return raise exn }
+        | Completed result as state ->
+            state, async { return result }
+        | Faulted exn as state ->
+            state, async { return raise exn }
 
     // computation will deallocate after state transition to Completed ot Faulted.
     new (computation, ?cancelUnawaited: bool, ?cacheException) =
@@ -87,12 +94,12 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
     member _.Request() =
         async { 
             try
-                return! lock stateUpdateSync request
+                return! withStateUpdate request
             finally
-                lock stateUpdateSync afterRequest
+                updateState afterRequest
         }
 
-    member _.CancelIfUnawaited() = lock stateUpdateSync cancelIfUnawaited
+    member _.CancelIfUnawaited() = updateState (cancelIfUnawaited true)
 
     member _.State = state
 
diff --git a/tests/FSharp.Compiler.ComponentTests/CompilerService/AsyncMemoize.fs b/tests/FSharp.Compiler.ComponentTests/CompilerService/AsyncMemoize.fs
@@ -197,6 +197,40 @@ let ``Job is restarted if first requestor cancels`` () =
         Started, key
         Finished, key ]
 
+[<Fact>]
+let ``Job is actually cancelled and restarted`` () =
+    let jobCanComplete = new ManualResetEvent(false)
+    let mutable finishedCount = 0
+
+    let computation = async {
+        do! awaitHandle jobCanComplete
+        Interlocked.Increment &finishedCount |> ignore
+        return 42
+    }
+
+    let memoize = AsyncMemoize<_, int, _>()
+    let events = observe memoize
+
+    let key = wrapKey 1
+
+    for i in 1 .. 10 do
+        use cts = new CancellationTokenSource()
+        let task = Async.StartAsTask( memoize.Get(key, computation), cancellationToken = cts.Token)
+        waitUntil events (received Started)
+        cts.Cancel()
+        assertTaskCanceled task
+        waitUntil events (received Canceled)
+        Assert.Equal(1, memoize.Count)
+
+    let _task2 = Async.StartAsTask( memoize.Get(key, computation))
+
+    waitUntil events (received Started)
+
+    jobCanComplete.Set() |> ignore
+
+    waitUntil events (received Finished)
+
+    Assert.Equal(1, finishedCount)
 
 [<Fact>]
 let ``Job keeps running if only one requestor cancels`` () =
diff --git a/tests/ILVerify/ilverify_FSharp.Compiler.Service_Debug_net9.0.bsl b/tests/ILVerify/ilverify_FSharp.Compiler.Service_Debug_net9.0.bsl
@@ -21,14 +21,14 @@
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.CodeAnalysis.Hosted.CompilerHelpers::fscCompile([FSharp.Compiler.Service]FSharp.Compiler.CodeAnalysis.LegacyReferenceResolver, string, string[])][offset 0x00000082][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.CodeAnalysis.Hosted.CompilerHelpers::fscCompile([FSharp.Compiler.Service]FSharp.Compiler.CodeAnalysis.LegacyReferenceResolver, string, string[])][offset 0x0000008B][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+MagicAssemblyResolution::ResolveAssemblyCore([FSharp.Compiler.Service]Internal.Utilities.Library.CompilationThreadToken, [FSharp.Compiler.Service]FSharp.Compiler.Text.Range, [FSharp.Compiler.Service]FSharp.Compiler.CompilerConfig+TcConfigBuilder, [FSharp.Compiler.Service]FSharp.Compiler.CompilerImports+TcImports, [FSharp.Compiler.Service]FSharp.Compiler.Interactive.Shell+FsiDynamicCompiler, [FSharp.Compiler.Service]FSharp.Compiler.Interactive.Shell+FsiConsoleOutput, string)][offset 0x00000015][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+clo@3510-788::Invoke([S.P.CoreLib]System.Tuple`3<char[],int32,int32>)][offset 0x000001E5][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+clo@3510-789::Invoke([S.P.CoreLib]System.Tuple`3<char[],int32,int32>)][offset 0x000001E5][found Char] Unexpected type on the stack.
 [IL]: Error [UnmanagedPointer]: : FSharp.Compiler.Interactive.Shell+Utilities+pointerToNativeInt@110::Invoke(object)][offset 0x00000007] Unmanaged pointers are not a verifiable type.
 [IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$FSharpCheckerResults+dataTipOfReferences@2205::Invoke([FSharp.Core]Microsoft.FSharp.Core.Unit)][offset 0x00000084][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000032][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000003B][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000082][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000008B][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000094][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000032][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000003B][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000082][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000008B][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000094][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.StaticLinking+TypeForwarding::followTypeForwardForILTypeRef([FSharp.Compiler.Service]FSharp.Compiler.AbstractIL.IL+ILTypeRef)][offset 0x00000010][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.CompilerOptions::getCompilerOption([FSharp.Compiler.Service]FSharp.Compiler.CompilerOptions+CompilerOption, [FSharp.Core]Microsoft.FSharp.Core.FSharpOption`1<int32>)][offset 0x000000E6][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.CompilerOptions::AddPathMapping([FSharp.Compiler.Service]FSharp.Compiler.CompilerConfig+TcConfigBuilder, string)][offset 0x0000000B][found Char] Unexpected type on the stack.
diff --git a/tests/ILVerify/ilverify_FSharp.Compiler.Service_Debug_netstandard2.0.bsl b/tests/ILVerify/ilverify_FSharp.Compiler.Service_Debug_netstandard2.0.bsl
@@ -28,18 +28,18 @@
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.CodeAnalysis.Hosted.CompilerHelpers::fscCompile([FSharp.Compiler.Service]FSharp.Compiler.CodeAnalysis.LegacyReferenceResolver, string, string[])][offset 0x0000008B][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+FsiStdinSyphon::GetLine(string, int32)][offset 0x00000039][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+MagicAssemblyResolution::ResolveAssemblyCore([FSharp.Compiler.Service]Internal.Utilities.Library.CompilationThreadToken, [FSharp.Compiler.Service]FSharp.Compiler.Text.Range, [FSharp.Compiler.Service]FSharp.Compiler.CompilerConfig+TcConfigBuilder, [FSharp.Compiler.Service]FSharp.Compiler.CompilerImports+TcImports, [FSharp.Compiler.Service]FSharp.Compiler.Interactive.Shell+FsiDynamicCompiler, [FSharp.Compiler.Service]FSharp.Compiler.Interactive.Shell+FsiConsoleOutput, string)][offset 0x00000015][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+clo@3510-788::Invoke([S.P.CoreLib]System.Tuple`3<char[],int32,int32>)][offset 0x000001E5][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+clo@3510-789::Invoke([S.P.CoreLib]System.Tuple`3<char[],int32,int32>)][offset 0x000001E5][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.Interactive.Shell+FsiInteractionProcessor::CompletionsForPartialLID([FSharp.Compiler.Service]FSharp.Compiler.Interactive.Shell+FsiDynamicCompilerState, string)][offset 0x0000001B][found Char] Unexpected type on the stack.
 [IL]: Error [UnmanagedPointer]: : FSharp.Compiler.Interactive.Shell+Utilities+pointerToNativeInt@110::Invoke(object)][offset 0x00000007] Unmanaged pointers are not a verifiable type.
 [IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$FSharpCheckerResults+dataTipOfReferences@2205::Invoke([FSharp.Core]Microsoft.FSharp.Core.Unit)][offset 0x00000084][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.EditorServices.AssemblyContent+traverseMemberFunctionAndValues@176::Invoke([FSharp.Compiler.Service]FSharp.Compiler.Symbols.FSharpMemberOrFunctionOrValue)][offset 0x00000059][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.EditorServices.AssemblyContent+traverseEntity@218::GenerateNext([S.P.CoreLib]System.Collections.Generic.IEnumerable`1<FSharp.Compiler.EditorServices.AssemblySymbol>&)][offset 0x000000DA][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.EditorServices.ParsedInput+visitor@1423-6::VisitExpr([FSharp.Core]Microsoft.FSharp.Collections.FSharpList`1<FSharp.Compiler.Syntax.SyntaxNode>, [FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<FSharp.Compiler.Syntax.SynExpr,Microsoft.FSharp.Core.FSharpOption`1<FSharp.Compiler.EditorServices.CompletionContext>>, [FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<FSharp.Compiler.Syntax.SynExpr,Microsoft.FSharp.Core.FSharpOption`1<FSharp.Compiler.EditorServices.CompletionContext>>, [FSharp.Compiler.Service]FSharp.Compiler.Syntax.SynExpr)][offset 0x00000605][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000032][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000003B][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000082][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000008B][found Char] Unexpected type on the stack.
-[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-500::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000094][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000032][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000003B][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000082][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x0000008B][found Char] Unexpected type on the stack.
+[IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$ServiceLexing+clo@921-501::Invoke([FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<System.Tuple`3<FSharp.Compiler.Parser+token,int32,int32>,Microsoft.FSharp.Core.Unit>)][offset 0x00000094][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : <StartupCode$FSharp-Compiler-Service>.$Symbols+fullName@2490-1::Invoke([FSharp.Core]Microsoft.FSharp.Core.Unit)][offset 0x00000015][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.CreateILModule+MainModuleBuilder::ConvertProductVersionToILVersionInfo(string)][offset 0x00000011][found Char] Unexpected type on the stack.
 [IL]: Error [StackUnexpected]: : FSharp.Compiler.StaticLinking+TypeForwarding::followTypeForwardForILTypeRef([FSharp.Compiler.Service]FSharp.Compiler.AbstractIL.IL+ILTypeRef)][offset 0x00000010][found Char] Unexpected type on the stack.
