Fix AsyncLazy cancellation (#18211)

* fix AsyncLazy cancellation

* baselines

Author: majocha

src/Compiler/Facilities/AsyncMemoize.fs

@@ -18,6 +18,7 @@ type AsyncLazyState<'t> =
 /// Represents a computation that will execute only once but can be requested by multiple clients.
 /// It keeps track of the number of requests. When all clients cancel their requests, the underlying computation will also cancel and can be restarted.
 /// If cancelUnawaited is set to false, the computation will run to completion even when all requests are canceled.
+/// When cacheException is false, subsequent requests will restart the computation after an exceptional result.
 type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool, cacheException: bool) =
 
     let stateUpdateSync = obj()
@@ -33,6 +34,7 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
     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.
+            // This means single execution is not strictly enforced.
             cts.Cancel()
             Initial computation
         | state -> state
@@ -45,38 +47,27 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
         async {
             try
                 let! ct = Async.CancellationToken
-                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
+                // Using ContinueWith with a CancellationToken allows detaching from the running 'work' task.
+                // If the current async workflow is canceled, the 'work' task will continue running independently.
+                do!  work.ContinueWith(ignore<Task<'t>>, ct) |> Async.AwaitTask
+            with :? TaskCanceledException -> ()
+            // If we're here it means there was no cancellation and the 'work' task has completed.
+            return! work |> Async.AwaitTask
         }
 
+    let onComplete (t: Task<'t>) =
+        updateState <| function
+            | Running (computation, _, _, _) ->
+                try Completed t.Result with exn -> if cacheException then Faulted exn else Initial computation
+            | state -> state
+        t.Result
+
     let request = function
         | Initial computation ->
             let cts = new CancellationTokenSource()
-            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)
+            let work =
+                Async.StartAsTask(computation, cancellationToken = cts.Token)
+                    .ContinueWith(onComplete, TaskContinuationOptions.NotOnCanceled)
             Running (computation, work, cts, 1),
             detachable work
         | Running (c, work, cts, count) ->
@@ -89,10 +80,10 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
 
     // computation will deallocate after state transition to Completed ot Faulted.
     new (computation, ?cancelUnawaited: bool, ?cacheException) =
-        AsyncLazy(Initial computation, defaultArg cancelUnawaited true, defaultArg cacheException false)
+        AsyncLazy(Initial computation, defaultArg cancelUnawaited true, defaultArg cacheException true)
 
     member _.Request() =
-        async { 
+        async {
             try
                 return! withStateUpdate request
             finally
@@ -103,7 +94,7 @@ type AsyncLazy<'t> private (initial: AsyncLazyState<'t>, cancelUnawaited: bool,
 
     member _.State = state
 
-    member this.TryResult =
+    member _.TryResult =
         match state with
         | Completed result -> Some result
         | _ -> None
@@ -164,12 +155,7 @@ type private KeyData<'TKey, 'TVersion> =
 type Job<'t> = AsyncLazy<Result<'t, exn> * CapturingDiagnosticsLogger>
 
 [<DebuggerDisplay("{DebuggerDisplay}")>]
-type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'TVersion: equality
-#if !NO_CHECKNULLS
-    and 'TKey:not null
-    and 'TVersion:not null
-#endif
-    >
+type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'TVersion: equality and 'TKey:not null and 'TVersion:not null>
     (?keepStrongly, ?keepWeakly, ?name: string, ?cancelUnawaitedJobs: bool, ?cancelDuplicateRunningJobs: bool) =
 
     let name = defaultArg name "N/A"
@@ -211,8 +197,7 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
             }
 
         let wrappedComputation =
-            async {
-                use! _handler = Async.OnCancel (fun () -> log Canceled key)
+            Async.TryCancelled( async {
                 let sw = Stopwatch.StartNew()
                 log Started key
                 let logger = CapturingDiagnosticsLogger "cache"
@@ -226,14 +211,14 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
                 | Choice2Of2 exn ->
                     log Failed key
                     return Result.Error exn, logger
-            }
+            }, fun _ -> log Canceled key)
 
         let getOrAdd () =
             let cached, otherVersions = cache.GetAll(key.Key, key.Version)
 
             let countHit v = Interlocked.Increment &hits |> ignore; v
             let cacheSetNewJob () =
-                let job = Job(wrappedComputation, cancelUnawaited = cancelUnawaitedJobs)
+                let job = Job(wrappedComputation, cancelUnawaited = cancelUnawaitedJobs, cacheException = false)
                 cache.Set(key.Key, key.Version, key.Label, job)
                 job
 

src/Compiler/Facilities/AsyncMemoize.fsi

@@ -42,12 +42,7 @@ type Extensions =
 ///
 /// Strongly holds at most one result per key.
 /// </summary>
-type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'TVersion: equality
-#if !NO_CHECKNULLS
-    and 'TKey:not null
-    and 'TVersion:not null
-#endif
-    > =
+type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'TVersion: equality and 'TKey:not null and 'TVersion:not null> =
 
     /// <param name="keepStrongly">Maximum number of strongly held results to keep in the cache</param>
     /// <param name="keepWeakly">Maximum number of weakly held results to keep in the cache</param>

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-790::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-796::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-502::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-502::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-502::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-502::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-502::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-508::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-508::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-508::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-508::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-508::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-790::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-796::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-502::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-502::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-502::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-502::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-502::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-508::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-508::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-508::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-508::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-508::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.