Yet another NodeCode problem affecting AsyncMemoize

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

@@ -430,6 +430,162 @@ let ``Cancel running jobs with the same key`` cancelDuplicate expectFinished =
 type DummyException(msg) =
     inherit Exception(msg)
 
+[<Fact>]
+let ``Preserve thread static diagnostics with inner task`` () = 
+
+    let seed = System.Random().Next()
+
+    let rng = System.Random seed
+
+    let job1Cache = AsyncMemoize()
+    let job2Cache = AsyncMemoize()
+
+    let taskJob = task {
+        do! Task.Delay 10
+        let ex = DummyException("job1 error")
+        DiagnosticsThreadStatics.DiagnosticsLogger.ErrorR(ex)
+    }
+
+    let job1 (input: string) = node {
+        let! _ = Async.Sleep (rng.Next(1, 30)) |> NodeCode.AwaitAsync
+        do! taskJob |> NodeCode.AwaitTask
+        return Ok input
+    }
+
+    let job2 (input: int) = node {
+
+        DiagnosticsThreadStatics.DiagnosticsLogger.Warning(DummyException("job2 error 1"))
+
+        let! _ = Async.Sleep (rng.Next(1, 30)) |> NodeCode.AwaitAsync
+
+        let key = { new ICacheKey<_, _> with
+                        member _.GetKey() = "job1"
+                        member _.GetVersion() = input
+                        member _.GetLabel() = "job1" }
+
+        let! result = job1Cache.Get(key, job1 "${input}" )
+
+        DiagnosticsThreadStatics.DiagnosticsLogger.Warning(DummyException("job2 error 2"))
+
+        return input, result
+
+    }
+
+    let tasks = seq {
+        for i in 1 .. 100 do
+
+            task {
+                let diagnosticsLogger =
+                    CompilationDiagnosticLogger($"Testing task {i}", FSharpDiagnosticOptions.Default)
+
+                use _ = new CompilationGlobalsScope(diagnosticsLogger, BuildPhase.Optimize)
+
+                DiagnosticsThreadStatics.DiagnosticsLogger.Warning(DummyException("task error"))
+
+
+                let key = { new ICacheKey<_, _> with
+                                member _.GetKey() = "job2"
+                                member _.GetVersion() = rng.Next(1, 10)
+                                member _.GetLabel() = "job2" }
+
+                let! result = job2Cache.Get(key, job2 (i % 10)) |> Async.AwaitNodeCode
+
+                let diagnostics = diagnosticsLogger.GetDiagnostics()
+
+                //Assert.Equal(3, diagnostics.Length)
+
+                return result, diagnostics
+            }
+    }
+
+    let results = (Task.WhenAll tasks).Result
+
+    let _diagnosticCounts = results |> Seq.map snd |> Seq.map Array.length |> Seq.groupBy id |> Seq.map (fun (k, v) -> k, v |> Seq.length) |> Seq.sortBy fst |> Seq.toList
+
+    //Assert.Equal<(int * int) list>([4, 100], diagnosticCounts)
+
+    let diagnosticMessages = results |> Seq.map snd |> Seq.map (Array.map (fun (d, _) -> d.Exception.Message) >> Array.toList) |> Set
+
+    Assert.Equal<Set<_>>(Set [["task error"; "job2 error 1"; "job1 error"; "job2 error 2"; ]], diagnosticMessages)
+
+[<Fact>]
+let ``Preserve thread static diagnostics with async computation`` () = 
+
+    let seed = System.Random().Next()
+
+    let rng = System.Random seed
+
+    let job1Cache = AsyncMemoize()
+    let job2Cache = AsyncMemoize()
+
+    let asyncJob = async {
+        do! Async.Sleep 10
+        let ex = DummyException("job1 error")
+        DiagnosticsThreadStatics.DiagnosticsLogger.ErrorR(ex)
+    }
+
+    let job1 (input: string) = node {
+        let! _ = Async.Sleep (rng.Next(1, 30)) |> NodeCode.AwaitAsync
+        do! asyncJob |> NodeCode.AwaitAsync
+        return Ok input
+    }
+
+    let job2 (input: int) = node {
+
+        DiagnosticsThreadStatics.DiagnosticsLogger.Warning(DummyException("job2 error 1"))
+
+        let! _ = Async.Sleep (rng.Next(1, 30)) |> NodeCode.AwaitAsync
+
+        let key = { new ICacheKey<_, _> with
+                        member _.GetKey() = "job1"
+                        member _.GetVersion() = input
+                        member _.GetLabel() = "job1" }
+
+        let! result = job1Cache.Get(key, job1 "${input}" )
+
+        DiagnosticsThreadStatics.DiagnosticsLogger.Warning(DummyException("job2 error 2"))
+
+        return input, result
+
+    }
+
+    let tasks = seq {
+        for i in 1 .. 100 do
+
+            task {
+                let diagnosticsLogger =
+                    CompilationDiagnosticLogger($"Testing task {i}", FSharpDiagnosticOptions.Default)
+
+                use _ = new CompilationGlobalsScope(diagnosticsLogger, BuildPhase.Optimize)
+
+                DiagnosticsThreadStatics.DiagnosticsLogger.Warning(DummyException("task error"))
+
+
+                let key = { new ICacheKey<_, _> with
+                                member _.GetKey() = "job2"
+                                member _.GetVersion() = rng.Next(1, 10)
+                                member _.GetLabel() = "job2" }
+
+                let! result = job2Cache.Get(key, job2 (i % 10)) |> Async.AwaitNodeCode
+
+                let diagnostics = diagnosticsLogger.GetDiagnostics()
+
+                //Assert.Equal(3, diagnostics.Length)
+
+                return result, diagnostics
+            }
+    }
+
+    let results = (Task.WhenAll tasks).Result
+
+    let _diagnosticCounts = results |> Seq.map snd |> Seq.map Array.length |> Seq.groupBy id |> Seq.map (fun (k, v) -> k, v |> Seq.length) |> Seq.sortBy fst |> Seq.toList
+
+    //Assert.Equal<(int * int) list>([4, 100], diagnosticCounts)
+
+    let diagnosticMessages = results |> Seq.map snd |> Seq.map (Array.map (fun (d, _) -> d.Exception.Message) >> Array.toList) |> Set
+
+    Assert.Equal<Set<_>>(Set [["task error"; "job2 error 1"; "job1 error"; "job2 error 2"; ]], diagnosticMessages)
+
 [<Fact>]
 let ``Preserve thread static diagnostics`` () = 
 

tests/FSharp.Compiler.UnitTests/BuildGraphTests.fs

@@ -237,23 +237,29 @@ module BuildGraphTests =
 
 
     [<Fact>]
-    let internal ``NodeCode preserves DiagnosticsThreadStatics`` () =
+    let ``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 asyncTask expectedPhase = async {
+            do! Async.Sleep 10
+            Assert.Equal(expectedPhase, DiagnosticsThreadStatics.BuildPhase)
         }
-
+
+        let rec job phase i = node {  
+            for i in 1 .. 10 do
+                Assert.Equal(phase, DiagnosticsThreadStatics.BuildPhase)
+                do! asyncTask phase |> NodeCode.AwaitAsync
+        }
+
         let work (phase: BuildPhase) =
             node {
-                use _ = new CompilationGlobalsScope(DiscardErrorsLogger, phase)
+                DiagnosticsThreadStatics.BuildPhase <- phase
                 let! _ = Seq.init 8 (job phase) |> NodeCode.Parallel
                 Assert.Equal(phase, DiagnosticsThreadStatics.BuildPhase)
             }
-    
+
         let phases = [|
             BuildPhase.DefaultPhase
             BuildPhase.Compile
@@ -267,9 +273,10 @@ module BuildGraphTests =
             BuildPhase.Output
             BuildPhase.Interactive
         |]
-    
+
         let pickRandomPhase _ = phases[random phases.Length]
-        Seq.init 100 pickRandomPhase
+
+        Seq.init 10 pickRandomPhase
         |> Seq.map (work >> Async.AwaitNodeCode)
         |> Async.Parallel
         |> Async.RunSynchronously