Use AsyncLocal to keep diagnostics instead of ThreadStatic

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

@@ -23,3 +23,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. ([PR #16602](https://github.com/dotnet/fsharp/pull/16602))

src/Compiler/Driver/CompilerImports.fs

@@ -2238,13 +2238,15 @@ and [<Sealed>] TcImports
                 | ParallelReferenceResolution.On -> NodeCode.Parallel
                 | ParallelReferenceResolution.Off -> NodeCode.Sequential
 
+            let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
             let! results =
                 nms
                 |> List.map (fun nm ->
                     node {
                         try
                             return! tcImports.TryRegisterAndPrepareToImportReferencedDll(ctok, nm)
                         with e ->
+                            use _ = UseDiagnosticsLogger diagnosticsLogger
                             errorR (Error(FSComp.SR.buildProblemReadingAssembly (nm.resolvedPath, e.Message), nm.originalReference.Range))
                             return None
                     })

src/Compiler/Facilities/BuildGraph.fs

@@ -7,22 +7,23 @@ open System.Threading
 open System.Threading.Tasks
 open System.Diagnostics
 open System.Globalization
-open FSharp.Compiler.DiagnosticsLogger
+//open FSharp.Compiler.DiagnosticsLogger
 open Internal.Utilities.Library
 
 [<NoEquality; NoComparison>]
 type NodeCode<'T> = Node of Async<'T>
 
 let wrapThreadStaticInfo computation =
     async {
-        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-        let phase = DiagnosticsThreadStatics.BuildPhase
+        //let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+        //let phase = DiagnosticsThreadStatics.BuildPhase
 
         try
             return! computation
         finally
-            DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-            DiagnosticsThreadStatics.BuildPhase <- phase
+            //DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+            //DiagnosticsThreadStatics.BuildPhase <- phase
+            ()
     }
 
 type Async<'T> with
@@ -92,19 +93,19 @@ type NodeCodeBuilder() =
     [<DebuggerHidden; DebuggerStepThrough>]
     member _.Combine(Node(p1): NodeCode<unit>, Node(p2): NodeCode<'T>) : NodeCode<'T> = Node(async.Combine(p1, p2))
 
-    [<DebuggerHidden; DebuggerStepThrough>]
-    member _.Using(value: CompilationGlobalsScope, binder: CompilationGlobalsScope -> NodeCode<'U>) =
-        Node(
-            async {
-                DiagnosticsThreadStatics.DiagnosticsLogger <- value.DiagnosticsLogger
-                DiagnosticsThreadStatics.BuildPhase <- value.BuildPhase
+    //[<DebuggerHidden; DebuggerStepThrough>]
+    //member _.Using(value: CompilationGlobalsScope, binder: CompilationGlobalsScope -> NodeCode<'U>) =
+    //    Node(
+    //        async {
+    //            DiagnosticsThreadStatics.DiagnosticsLogger <- value.DiagnosticsLogger
+    //            DiagnosticsThreadStatics.BuildPhase <- value.BuildPhase
 
-                try
-                    return! binder value |> Async.AwaitNodeCode
-                finally
-                    (value :> IDisposable).Dispose()
-            }
-        )
+    //            try
+    //                return! binder value |> Async.AwaitNodeCode
+    //            finally
+    //                (value :> IDisposable).Dispose()
+    //        }
+    //    )
 
     [<DebuggerHidden; DebuggerStepThrough>]
     member _.Using(value: IDisposable, binder: IDisposable -> NodeCode<'U>) =
@@ -123,44 +124,46 @@ type NodeCode private () =
     static let cancellationToken = Node(wrapThreadStaticInfo Async.CancellationToken)
 
     static member RunImmediate(computation: NodeCode<'T>, ct: CancellationToken) =
-        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-        let phase = DiagnosticsThreadStatics.BuildPhase
+        //let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+        //let phase = DiagnosticsThreadStatics.BuildPhase
 
         try
             try
                 let work =
                     async {
-                        DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                        DiagnosticsThreadStatics.BuildPhase <- phase
+                        //DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+                        //DiagnosticsThreadStatics.BuildPhase <- phase
                         return! computation |> Async.AwaitNodeCode
                     }
 
                 Async.StartImmediateAsTask(work, cancellationToken = ct).Result
             finally
-                DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                DiagnosticsThreadStatics.BuildPhase <- phase
+                //DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+                //DiagnosticsThreadStatics.BuildPhase <- phase
+                ()
         with :? AggregateException as ex when ex.InnerExceptions.Count = 1 ->
             raise (ex.InnerExceptions[0])
 
     static member RunImmediateWithoutCancellation(computation: NodeCode<'T>) =
         NodeCode.RunImmediate(computation, CancellationToken.None)
 
     static member StartAsTask_ForTesting(computation: NodeCode<'T>, ?ct: CancellationToken) =
-        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-        let phase = DiagnosticsThreadStatics.BuildPhase
+        //let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+        //let phase = DiagnosticsThreadStatics.BuildPhase
 
         try
             let work =
                 async {
-                    DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                    DiagnosticsThreadStatics.BuildPhase <- phase
+                    //DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+                    //DiagnosticsThreadStatics.BuildPhase <- phase
                     return! computation |> Async.AwaitNodeCode
                 }
 
             Async.StartAsTask(work, cancellationToken = defaultArg ct CancellationToken.None)
         finally
-            DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-            DiagnosticsThreadStatics.BuildPhase <- phase
+            //DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+            //DiagnosticsThreadStatics.BuildPhase <- phase
+            ()
 
     static member CancellationToken = cancellationToken
 
@@ -193,14 +196,14 @@ type NodeCode private () =
         }
 
     static member Parallel(computations: NodeCode<'T> seq) =
-        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-        let phase = DiagnosticsThreadStatics.BuildPhase
+        //let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+        //let phase = DiagnosticsThreadStatics.BuildPhase
 
         computations
         |> Seq.map (fun (Node x) ->
             async {
-                DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                DiagnosticsThreadStatics.BuildPhase <- phase
+                //DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+                //DiagnosticsThreadStatics.BuildPhase <- phase
                 return! x
             })
         |> Async.Parallel

src/Compiler/Facilities/BuildGraph.fsi

@@ -44,8 +44,8 @@ type NodeCodeBuilder =
 
     member Combine: x1: NodeCode<unit> * x2: NodeCode<'T> -> NodeCode<'T>
 
-    /// A limited form 'use' for establishing the compilation globals.
-    member Using: CompilationGlobalsScope * (CompilationGlobalsScope -> NodeCode<'T>) -> NodeCode<'T>
+    ///// A limited form 'use' for establishing the compilation globals.
+    //member Using: CompilationGlobalsScope * (CompilationGlobalsScope -> NodeCode<'T>) -> NodeCode<'T>
 
     /// A generic 'use' that disposes of the IDisposable at the end of the computation.
     member Using: IDisposable * (IDisposable -> NodeCode<'T>) -> NodeCode<'T>

src/Compiler/Facilities/DiagnosticsLogger.fs

@@ -13,6 +13,7 @@ open System.Reflection
 open System.Threading
 open Internal.Utilities.Library
 open Internal.Utilities.Library.Extras
+open System.Threading
 
 /// Represents the style being used to format errors
 [<RequireQualifiedAccess>]
@@ -375,28 +376,25 @@ type CapturingDiagnosticsLogger(nm, ?eagerFormat) =
         errors |> Array.iter diagnosticsLogger.DiagnosticSink
 
 /// Type holds thread-static globals for use by the compiler.
-type internal DiagnosticsThreadStatics =
-    [<ThreadStatic; DefaultValue>]
-    static val mutable private buildPhase: BuildPhase
-
-    [<ThreadStatic; DefaultValue>]
-    static val mutable private diagnosticsLogger: DiagnosticsLogger
+type DiagnosticsThreadStatics =
+    static let buildPhase = new AsyncLocal<BuildPhase>()
+    static let diagnosticsLogger = new AsyncLocal<DiagnosticsLogger>()
 
-    static member BuildPhaseUnchecked = DiagnosticsThreadStatics.buildPhase
+    static let EnsureCreated (h: AsyncLocal<_>) d =
+        if box h.Value |> isNull then
+            h.Value <- d
 
     static member BuildPhase
         with get () =
-            match box DiagnosticsThreadStatics.buildPhase with
-            | Null -> BuildPhase.DefaultPhase
-            | _ -> DiagnosticsThreadStatics.buildPhase
-        and set v = DiagnosticsThreadStatics.buildPhase <- v
+            EnsureCreated buildPhase BuildPhase.DefaultPhase
+            buildPhase.Value
+        and set v = buildPhase.Value <- v
 
     static member DiagnosticsLogger
         with get () =
-            match box DiagnosticsThreadStatics.diagnosticsLogger with
-            | Null -> AssertFalseDiagnosticsLogger
-            | _ -> DiagnosticsThreadStatics.diagnosticsLogger
-        and set v = DiagnosticsThreadStatics.diagnosticsLogger <- v
+            EnsureCreated diagnosticsLogger AssertFalseDiagnosticsLogger
+            diagnosticsLogger.Value
+        and set v = diagnosticsLogger.Value <- v
 
 [<AutoOpen>]
 module DiagnosticsLoggerExtensions =
@@ -498,7 +496,7 @@ module DiagnosticsLoggerExtensions =
 
 /// NOTE: The change will be undone when the returned "unwind" object disposes
 let UseBuildPhase (phase: BuildPhase) =
-    let oldBuildPhase = DiagnosticsThreadStatics.BuildPhaseUnchecked
+    let oldBuildPhase = DiagnosticsThreadStatics.BuildPhase
     DiagnosticsThreadStatics.BuildPhase <- phase
 
     { new IDisposable with

src/Compiler/Facilities/DiagnosticsLogger.fsi

@@ -232,8 +232,6 @@ type DiagnosticsThreadStatics =
 
     static member BuildPhase: BuildPhase with get, set
 
-    static member BuildPhaseUnchecked: BuildPhase
-
     static member DiagnosticsLogger: DiagnosticsLogger with get, set
 
 [<AutoOpen>]

tests/FSharp.Compiler.ComponentTests/FSharpChecker/TransparentCompiler.fs

@@ -626,7 +626,7 @@ let fuzzingTest seed (project: SyntheticProject) = task {
             ()
         with
             | e ->
-                let _log = log.Values |> Seq.collect id |> Seq.sortBy p13 |> Seq.toArray
+                // let _log = log.Values |> Seq.collect id |> Seq.sortBy p13 |> Seq.toArray
                 failwith $"Seed: {seed}\nException: %A{e}"
     }
     let log = log.Values |> Seq.collect id |> Seq.sortBy p13 |> Seq.toArray