Some AsyncMemoize fixes (dotnet#18074)

* deflake tests

* ilverify

* fix race

* ilverify

* restore accidental change

src/Compiler/Facilities/AsyncMemoize.fs

@@ -1,12 +1,10 @@
 namespace Internal.Utilities.Collections
 
-open System
 open System.Diagnostics
 open System.IO
 open System.Threading
 open System.Threading.Tasks
 open System.Runtime.CompilerServices
-open System.Runtime.ExceptionServices
 
 open FSharp.Compiler.DiagnosticsLogger
 open Internal.Utilities.Library
@@ -34,44 +32,44 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
 
     let afterRequest () =
         match state with
-        | Running(computation, work, _, _) when work.IsCompleted ->
-            state <-
-                try
-                    Completed work.Result
-                with
-                | exn ->
-                    if cacheException then Faulted exn else Initial computation
         | Running(c, work, cts, count) ->
             state <- Running(c, work, cts, count - 1)
             if cancelUnawaited then cancelIfUnawaited () 
-        | _ -> () // Nothing more to do if another request already transitioned the state.
+        | _ -> () // Nothing more to do if state already transitioned.
 
     let detachable (work: Task<'t>) =
         async {
             try
                 let! ct = Async.CancellationToken
-                let options = TaskContinuationOptions.ExecuteSynchronously
-                try
-                    return!
-                        // Using ContinueWith with a CancellationToken allows detaching from the running 'work' task.
-                        // This ensures the lazy 'work' and its awaiting requests can be independently managed 
-                        // by separate CancellationTokenSources, enabling individual cancellation.
-                        // Essentially, if this async computation is canceled, it won't wait for the 'work' to complete
-                        // but will immediately proceed to the finally block.
-                        work.ContinueWith((fun (t: Task<_>) -> t.Result), ct, options, TaskScheduler.Current)
-                        |> Async.AwaitTask
-                // Cancellation check before entering the `with` ensures TaskCanceledException coming from the ContinueWith task will never be raised here.
-                // The cancellation continuation will always be called in case of cancellation.
-                with exn -> return raise exn
-            finally
-                lock stateUpdateSync afterRequest
+                return!
+                    // Using ContinueWith with a CancellationToken allows detaching from the running 'work' task.
+                    // This ensures the lazy 'work' and its awaiting requests can be independently managed 
+                    // by separate CancellationTokenSources, enabling individual cancellation.
+                    // Essentially, if this async computation is canceled, it won't wait for the 'work' to complete
+                    // but will immediately proceed to the finally block.
+                    work.ContinueWith((fun (t: Task<_>) -> t.Result), ct)
+                    |> Async.AwaitTask
+            // Cancellation check before entering the `with` ensures TaskCanceledException coming from the ContinueWith task will never be raised here.
+            // 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
         | 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)
             detachable work
         | Running (c, work, cts, count) ->
@@ -86,7 +84,13 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
     new (computation, ?cancelUnawaited: bool, ?cacheException) =
         AsyncLazy(Initial computation, defaultArg cancelUnawaited true, defaultArg cacheException false)
 
-    member _.Request() = lock stateUpdateSync request
+    member _.Request() =
+        async { 
+            try
+                return! lock stateUpdateSync request
+            finally
+                lock stateUpdateSync afterRequest
+        }
 
     member _.CancelIfUnawaited() = lock stateUpdateSync cancelIfUnawaited
 

tests/FSharp.Compiler.ComponentTests/CompilerService/AsyncMemoize.fs

@@ -64,6 +64,8 @@ let internal wrapKey key =
 let assertTaskCanceled (task: Task<_>) =
     Assert.ThrowsAnyAsync<OperationCanceledException>(fun () -> task).Result |> ignore
 
+let awaitHandle h = h |> Async.AwaitWaitHandle |> Async.Ignore
+
 [<Fact>]
 let ``Basics``() =
     let computation key = async {
@@ -101,13 +103,13 @@ let ``Basics``() =
 let ``We can disconnect a request from a running job`` () =
 
     let cts = new CancellationTokenSource()
-    let canFinish = new ManualResetEventSlim(false)
+    let canFinish = new ManualResetEvent(false)
 
     let computation = async {
-        canFinish.Wait()
+        do! awaitHandle canFinish
     }
 
-    let memoize = AsyncMemoize<_, int, _>(cancelUnawaitedJobs = false, cancelDuplicateRunningJobs = true)
+    let memoize = AsyncMemoize<_, int, _>(cancelUnawaitedJobs = false)
     let events = observe memoize
 
     let key = 1
@@ -156,10 +158,10 @@ let ``We can cancel a job`` () =
 
 [<Fact>]
 let ``Job is restarted if first requestor cancels`` () =
-    let jobCanComplete = new ManualResetEventSlim(false)
+    let jobCanComplete = new ManualResetEvent(false)
 
     let computation key = async {
-        jobCanComplete.Wait()
+        do! awaitHandle jobCanComplete
         return key * 2
     }
 
@@ -170,11 +172,13 @@ let ``Job is restarted if first requestor cancels`` () =
 
     let key = 1
 
-    let _task1 = Async.StartAsTask( memoize.Get(wrapKey key, computation key), cancellationToken = cts1.Token)
+    let task1 = Async.StartAsTask( memoize.Get(wrapKey key, computation key), cancellationToken = cts1.Token)
 
     waitUntil events (received Started)
     cts1.Cancel()
 
+    assertTaskCanceled task1
+
     waitUntil events (received Canceled)
 
     let task2 = Async.StartAsTask( memoize.Get(wrapKey key, computation key))
@@ -197,10 +201,10 @@ let ``Job is restarted if first requestor cancels`` () =
 [<Fact>]
 let ``Job keeps running if only one requestor cancels`` () =
 
-    let jobCanComplete = new ManualResetEventSlim(false)
+    let jobCanComplete = new ManualResetEvent(false)
 
     let computation key = async {
-        jobCanComplete.Wait()
+        do! awaitHandle jobCanComplete
         return key * 2
     }
 
@@ -362,10 +366,10 @@ let ``Cancel running jobs with the same key`` () =
 
     let events = observe cache
 
-    let jobCanContinue = new ManualResetEventSlim(false)
+    let jobCanContinue = new ManualResetEvent(false)
 
     let work = async {
-        jobCanContinue.Wait()
+        do! awaitHandle jobCanContinue
     }
 
     let key version =
@@ -392,7 +396,8 @@ let ``Cancel running jobs with the same key`` () =
     let current = eventsWhen events (received Requested)
     Assert.Equal(0, current |> countOf Canceled)
 
-    // new request should cancel the unobserved jobs
+    waitUntil events (countOf Canceled >> (=) 10)
+
     waitUntil events (received Started)
 
     jobCanContinue.Set() |> ignore

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-791::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-787::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-504::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-504::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-504::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-504::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-504::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-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]: : 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.

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-791::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-787::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-504::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-504::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-504::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-504::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-504::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-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>.$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.