AsyncLocal part 2 - replace all ThreadStatics and remove NodeCode and node CE

docs/release-notes/.FSharp.Compiler.Service/8.0.300.md

@@ -29,3 +29,4 @@
 * Reduce allocations in compiler checking via `ValueOption` usage ([PR #16323](https://github.com/dotnet/fsharp/pull/16323), [PR #16567](https://github.com/dotnet/fsharp/pull/16567))
 * Reverted [#16348](https://github.com/dotnet/fsharp/pull/16348) `ThreadStatic` `CancellationToken` changes to improve test stability and prevent potential unwanted cancellations. ([PR #16536](https://github.com/dotnet/fsharp/pull/16536))
 * Refactored parenthesization API. ([PR #16461])(https://github.com/dotnet/fsharp/pull/16461))
+* Replaced `ThreadStatic` diagnostics globals with `AsyncLocal` for better stability of Transparent Compiler. Removed `NodeCode` in favor of unflavored Async ([PR #16645](https://github.com/dotnet/fsharp/pull/16645))

docs/release-notes/.VisualStudio/17.10.md

@@ -6,3 +6,4 @@
 ### Changed
 
 * Use refactored parenthesization API in unnecessary parentheses code fix. ([PR #16461])(https://github.com/dotnet/fsharp/pull/16461))
+* Removed `NodeCode` in favor of unflavored Async with AsyncLocal state ([PR #16645](https://github.com/dotnet/fsharp/pull/16645))

src/Compiler/Driver/CompilerConfig.fs

@@ -252,7 +252,7 @@ and IProjectReference =
     abstract FileName: string
 
     /// Evaluate raw contents of the assembly file generated by the project
-    abstract EvaluateRawContents: unit -> NodeCode<ProjectAssemblyDataResult>
+    abstract EvaluateRawContents: unit -> Async<ProjectAssemblyDataResult>
 
     /// Get the logical timestamp that would be the timestamp of the assembly file generated by the project
     ///

src/Compiler/Driver/CompilerConfig.fsi

@@ -86,7 +86,7 @@ and IProjectReference =
     /// Evaluate raw contents of the assembly file generated by the project.
     /// 'None' may be returned if an in-memory view of the contents is, for some reason,
     /// not available.  In this case the on-disk view of the contents will be preferred.
-    abstract EvaluateRawContents: unit -> NodeCode<ProjectAssemblyDataResult>
+    abstract EvaluateRawContents: unit -> Async<ProjectAssemblyDataResult>
 
     /// Get the logical timestamp that would be the timestamp of the assembly file generated by the project.
     ///

src/Compiler/Driver/CompilerImports.fs

@@ -1735,7 +1735,7 @@ and [<Sealed>] TcImports
             m
         ) =
 
-        let startingErrorCount = DiagnosticsThreadStatics.DiagnosticsLogger.ErrorCount
+        let startingErrorCount = DiagnosticsAsyncState.DiagnosticsLogger.ErrorCount
 
         // Find assembly level TypeProviderAssemblyAttributes. These will point to the assemblies that
         // have class which implement ITypeProvider and which have TypeProviderAttribute on them.
@@ -1936,7 +1936,7 @@ and [<Sealed>] TcImports
                     with RecoverableException e ->
                         errorRecovery e m
 
-                if startingErrorCount < DiagnosticsThreadStatics.DiagnosticsLogger.ErrorCount then
+                if startingErrorCount < DiagnosticsAsyncState.DiagnosticsLogger.ErrorCount then
                     error (Error(FSComp.SR.etOneOrMoreErrorsSeenDuringExtensionTypeSetting (), m))
 
             providers
@@ -2158,14 +2158,14 @@ and [<Sealed>] TcImports
         (
             ctok,
             r: AssemblyResolution
-        ) : NodeCode<(_ * (unit -> AvailableImportedAssembly list)) option> =
-        node {
+        ) : Async<(_ * (unit -> AvailableImportedAssembly list)) option> =
+        async {
             CheckDisposed()
             let m = r.originalReference.Range
             let fileName = r.resolvedPath
 
             let! contentsOpt =
-                node {
+                async {
                     match r.ProjectReference with
                     | Some ilb -> return! ilb.EvaluateRawContents()
                     | None -> return ProjectAssemblyDataResult.Unavailable true
@@ -2228,27 +2228,29 @@ and [<Sealed>] TcImports
 
     // NOTE: When used in the Language Service this can cause the transitive checking of projects. Hence it must be cancellable.
     member tcImports.RegisterAndImportReferencedAssemblies(ctok, nms: AssemblyResolution list) =
-        node {
+        async {
             CheckDisposed()
 
             let tcConfig = tcConfigP.Get ctok
 
             let runMethod =
                 match tcConfig.parallelReferenceResolution with
-                | ParallelReferenceResolution.On -> NodeCode.Parallel
-                | ParallelReferenceResolution.Off -> NodeCode.Sequential
-
-            let! results =
-                nms
-                |> List.map (fun nm ->
-                    node {
-                        try
-                            return! tcImports.TryRegisterAndPrepareToImportReferencedDll(ctok, nm)
-                        with e ->
-                            errorR (Error(FSComp.SR.buildProblemReadingAssembly (nm.resolvedPath, e.Message), nm.originalReference.Range))
-                            return None
-                    })
-                |> runMethod
+                | ParallelReferenceResolution.On -> Async.Parallel
+                | ParallelReferenceResolution.Off -> Async.SequentialImmediate
+
+            let! results = async {
+                    use captureTasks = new CaptureDiagnosticsConcurrently(DiagnosticsAsyncState.DiagnosticsLogger)
+                    return! nms |> List.map (fun nm ->
+                        async {
+                            use _ = UseDiagnosticsLogger captureTasks.LoggerForTask
+                            try
+                                return! tcImports.TryRegisterAndPrepareToImportReferencedDll(ctok, nm)
+                            with e ->
+                                errorR (Error(FSComp.SR.buildProblemReadingAssembly (nm.resolvedPath, e.Message), nm.originalReference.Range))
+                                return None
+                        })
+                        |> runMethod
+                }
 
             let _dllinfos, phase2s = results |> Array.choose id |> List.ofArray |> List.unzip
             fixupOrphanCcus ()
@@ -2282,7 +2284,7 @@ and [<Sealed>] TcImports
                     ReportWarnings warns
 
                     tcImports.RegisterAndImportReferencedAssemblies(ctok, res)
-                    |> NodeCode.RunImmediateWithoutCancellation
+                    |> Async.RunImmediateWithoutCancellation
                     |> ignore
 
                     true
@@ -2383,7 +2385,7 @@ and [<Sealed>] TcImports
     // we dispose TcImports is because we need to dispose type providers, and type providers are never included in the framework DLL set.
     // If a framework set ever includes type providers, you will not have to worry about explicitly calling Dispose as the Finalizer will handle it.
     static member BuildFrameworkTcImports(tcConfigP: TcConfigProvider, frameworkDLLs, nonFrameworkDLLs) =
-        node {
+        async {
             let ctok = CompilationThreadToken()
             let tcConfig = tcConfigP.Get ctok
 
@@ -2460,7 +2462,7 @@ and [<Sealed>] TcImports
                 resolvedAssemblies |> List.choose tryFindEquivPrimaryAssembly
 
             let! fslibCcu, fsharpCoreAssemblyScopeRef =
-                node {
+                async {
                     if tcConfig.compilingFSharpCore then
                         // When compiling FSharp.Core.dll, the fslibCcu reference to FSharp.Core.dll is a delayed ccu thunk fixed up during type checking
                         return CcuThunk.CreateDelayed getFSharpCoreLibraryName, ILScopeRef.Local
@@ -2553,7 +2555,7 @@ and [<Sealed>] TcImports
             dependencyProvider
         ) =
 
-        node {
+        async {
             let ctok = CompilationThreadToken()
             let tcConfig = tcConfigP.Get ctok
 
@@ -2571,7 +2573,7 @@ and [<Sealed>] TcImports
         }
 
     static member BuildTcImports(tcConfigP: TcConfigProvider, dependencyProvider) =
-        node {
+        async {
             let ctok = CompilationThreadToken()
             let tcConfig = tcConfigP.Get ctok
 
@@ -2603,7 +2605,7 @@ let RequireReferences (ctok, tcImports: TcImports, tcEnv, thisAssemblyName, reso
 
     let ccuinfos =
         tcImports.RegisterAndImportReferencedAssemblies(ctok, resolutions)
-        |> NodeCode.RunImmediateWithoutCancellation
+        |> Async.RunImmediateWithoutCancellation
 
     let asms =
         ccuinfos

src/Compiler/Driver/CompilerImports.fsi

@@ -199,14 +199,14 @@ type TcImports =
     member internal Base: TcImports option
 
     static member BuildFrameworkTcImports:
-        TcConfigProvider * AssemblyResolution list * AssemblyResolution list -> NodeCode<TcGlobals * TcImports>
+        TcConfigProvider * AssemblyResolution list * AssemblyResolution list -> Async<TcGlobals * TcImports>
 
     static member BuildNonFrameworkTcImports:
         TcConfigProvider * TcImports * AssemblyResolution list * UnresolvedAssemblyReference list * DependencyProvider ->
-            NodeCode<TcImports>
+            Async<TcImports>
 
     static member BuildTcImports:
-        tcConfigP: TcConfigProvider * dependencyProvider: DependencyProvider -> NodeCode<TcGlobals * TcImports>
+        tcConfigP: TcConfigProvider * dependencyProvider: DependencyProvider -> Async<TcGlobals * TcImports>
 
 /// Process a group of #r in F# Interactive.
 /// Adds the reference to the tcImports and add the ccu to the type checking environment.

src/Compiler/Driver/ParseAndCheckInputs.fs

@@ -1828,7 +1828,7 @@ let CheckMultipleInputsUsingGraphMode
             |> Graph.writeMermaidToFile graphFile)
 
     let _ = ctok // TODO Use it
-    let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+    let diagnosticsLogger = DiagnosticsAsyncState.DiagnosticsLogger
 
     // In the first linear part of parallel checking, we use a 'checkForErrors' that checks either for errors
     // somewhere in the files processed prior to each one, or in the processing of this particular file.

src/Compiler/Driver/fsc.fs

@@ -614,7 +614,7 @@ let main1
     // Import basic assemblies
     let tcGlobals, frameworkTcImports =
         TcImports.BuildFrameworkTcImports(foundationalTcConfigP, sysRes, otherRes)
-        |> NodeCode.RunImmediateWithoutCancellation
+        |> Async.RunImmediateWithoutCancellation
 
     let ilSourceDocs =
         [
@@ -663,7 +663,7 @@ let main1
 
     let tcImports =
         TcImports.BuildNonFrameworkTcImports(tcConfigP, frameworkTcImports, otherRes, knownUnresolved, dependencyProvider)
-        |> NodeCode.RunImmediateWithoutCancellation
+        |> Async.RunImmediateWithoutCancellation
 
     // register tcImports to be disposed in future
     disposables.Register tcImports

src/Compiler/Facilities/AsyncMemoize.fs

@@ -50,11 +50,11 @@ type internal MemoizeReply<'TValue> =
     | New of CancellationToken
     | Existing of Task<'TValue>
 
-type internal MemoizeRequest<'TValue> = GetOrCompute of NodeCode<'TValue> * CancellationToken
+type internal MemoizeRequest<'TValue> = GetOrCompute of Async<'TValue> * CancellationToken
 
 [<DebuggerDisplay("{DebuggerDisplay}")>]
 type internal Job<'TValue> =
-    | Running of TaskCompletionSource<'TValue> * CancellationTokenSource * NodeCode<'TValue> * DateTime * ResizeArray<DiagnosticsLogger>
+    | Running of TaskCompletionSource<'TValue> * CancellationTokenSource * Async<'TValue> * DateTime * ResizeArray<DiagnosticsLogger>
     | Completed of 'TValue * (PhasedDiagnostic * FSharpDiagnosticSeverity) list
     | Canceled of DateTime
     | Failed of DateTime * exn // TODO: probably we don't need to keep this
@@ -354,15 +354,12 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
                                                 log (Restarted, key)
                                                 Interlocked.Increment &restarted |> ignore
                                                 System.Diagnostics.Trace.TraceInformation $"{name} Restarted {key.Label}"
-                                                let currentLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-                                                DiagnosticsThreadStatics.DiagnosticsLogger <- cachingLogger
-
-                                                try
-                                                    let! result = computation |> Async.AwaitNodeCode
-                                                    post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
-                                                    return ()
-                                                finally
-                                                    DiagnosticsThreadStatics.DiagnosticsLogger <- currentLogger
+                                                use _ = UseDiagnosticsLogger cachingLogger
+
+                                                let! result = computation
+                                                post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
+                                                return ()
+
                                             with
                                             | TaskCancelled _ ->
                                                 Interlocked.Increment &cancel_exception_subsequent |> ignore
@@ -481,14 +478,22 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
                 Version = key.GetVersion()
             }
 
-        node {
-            let! ct = NodeCode.CancellationToken
+        let callerDiagnosticLogger =
+            if DiagnosticsAsyncState.DiagnosticsLogger = UninitializedDiagnosticsLogger then
+                // TODO: Telemetry?
+                DiagnosticsAsyncState.DiagnosticsLogger <- DiscardErrorsLogger
+
+            DiagnosticsAsyncState.DiagnosticsLogger
+
+        // Preserve outer diagnostics scope in case the computation doesn't.
+        let computation = computation |> Async.CompilationScope
 
-            let callerDiagnosticLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+        async {
+            let! ct = Async.CancellationToken
 
             match!
                 processRequest post (key, GetOrCompute(computation, ct)) callerDiagnosticLogger
-                |> NodeCode.AwaitTask
+                |> Async.AwaitTask
             with
             | New internalCt ->
 
@@ -500,21 +505,17 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
                         Async.StartAsTask(
                             async {
                                 // TODO: Should unify starting and restarting
-                                let currentLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-                                DiagnosticsThreadStatics.DiagnosticsLogger <- cachingLogger
+                                use _ = UseDiagnosticsLogger cachingLogger
 
                                 log (Started, key)
 
-                                try
-                                    let! result = computation |> Async.AwaitNodeCode
-                                    post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
-                                    return result
-                                finally
-                                    DiagnosticsThreadStatics.DiagnosticsLogger <- currentLogger
+                                let! result = computation
+                                post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
+                                return result
                             },
                             cancellationToken = linkedCtSource.Token
                         )
-                        |> NodeCode.AwaitTask
+                        |> Async.AwaitTask
                 with
                 | TaskCancelled ex ->
                     // TODO: do we need to do anything else here? Presumably it should be done by the registration on
@@ -530,7 +531,7 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
                     post (key, (JobFailed(ex, cachingLogger.CapturedDiagnostics)))
                     return raise ex
 
-            | Existing job -> return! job |> NodeCode.AwaitTask
+            | Existing job -> return! job |> Async.AwaitTask
 
         }
 

src/Compiler/Facilities/AsyncMemoize.fsi

@@ -65,9 +65,9 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
 
     member Clear: predicate: ('TKey -> bool) -> unit
 
-    member Get: key: ICacheKey<'TKey, 'TVersion> * computation: NodeCode<'TValue> -> NodeCode<'TValue>
+    member Get: key: ICacheKey<'TKey, 'TVersion> * computation: Async<'TValue> -> Async<'TValue>
 
-    member Get': key: 'TKey * computation: NodeCode<'TValue> -> NodeCode<'TValue>
+    member Get': key: 'TKey * computation: Async<'TValue> -> Async<'TValue>
 
     member Event: IEvent<JobEvent * (string * 'TKey * 'TVersion)>