NodeCode.Parallel keeps ThreadStatics (#16576)

dotnet · Jan 25, 2024 · 4c0a592 · 4c0a592
1 parent 544a1c9
commit 4c0a592
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 1 deletion.
diff --git a/src/Compiler/Facilities/BuildGraph.fs b/src/Compiler/Facilities/BuildGraph.fs
@@ -193,7 +193,19 @@ type NodeCode private () =
         }
 
     static member Parallel(computations: NodeCode<'T> seq) =
-        computations |> Seq.map (fun (Node x) -> x) |> Async.Parallel |> Node
+        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
+        let phase = DiagnosticsThreadStatics.BuildPhase
+
+        computations
+        |> Seq.map (fun (Node x) ->
+            async {
+                DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
+                DiagnosticsThreadStatics.BuildPhase <- phase
+                return! x
+            })
+        |> Async.Parallel
+        |> wrapThreadStaticInfo
+        |> Node
 
 [<RequireQualifiedAccess>]
 module GraphNode =

diff --git a/tests/FSharp.Compiler.UnitTests/BuildGraphTests.fs b/tests/FSharp.Compiler.UnitTests/BuildGraphTests.fs
@@ -8,6 +8,7 @@ open Xunit
 open FSharp.Test
 open FSharp.Test.Compiler
 open FSharp.Compiler.BuildGraph
+open FSharp.Compiler.DiagnosticsLogger
 open Internal.Utilities.Library
 
 module BuildGraphTests =
@@ -233,3 +234,42 @@ module BuildGraphTests =
 
         Assert.shouldBeTrue graphNode.HasValue
         Assert.shouldBe (ValueSome 1) (graphNode.TryPeekValue())
+
+
+    [<Fact>]
+    let internal ``NodeCode preserves DiagnosticsThreadStatics`` () =
+        let random =
+            let rng = Random()
+            fun n -> rng.Next n
+
+        let job phase _ = node {
+            do! random 10 |> Async.Sleep |> NodeCode.AwaitAsync
+            Assert.Equal(phase, DiagnosticsThreadStatics.BuildPhase)
+        }
+
+        let work (phase: BuildPhase) =
+            node {
+                use _ = new CompilationGlobalsScope(DiscardErrorsLogger, phase)
+                let! _ = Seq.init 8 (job phase) |> NodeCode.Parallel
+                Assert.Equal(phase, DiagnosticsThreadStatics.BuildPhase)
+            }
+
+        let phases = [|
+            BuildPhase.DefaultPhase
+            BuildPhase.Compile
+            BuildPhase.Parameter
+            BuildPhase.Parse
+            BuildPhase.TypeCheck
+            BuildPhase.CodeGen
+            BuildPhase.Optimize
+            BuildPhase.IlxGen
+            BuildPhase.IlGen
+            BuildPhase.Output
+            BuildPhase.Interactive
+        |]
+
+        let pickRandomPhase _ = phases[random phases.Length]
+        Seq.init 100 pickRandomPhase
+        |> Seq.map (work >> Async.AwaitNodeCode)
+        |> Async.Parallel
+        |> Async.RunSynchronously