Starting point for thorough testing of AsyncLocal impact

src/Compiler/Driver/ParseAndCheckInputs.fs

@@ -42,6 +42,7 @@ open FSharp.Compiler.TypedTree
 open FSharp.Compiler.TypedTreeOps
 open FSharp.Compiler.TypedTreeBasics
 open FSharp.Compiler.TcGlobals
+open FSharp.Compiler.BuildGraph
 
 let CanonicalizeFilename fileName =
     let basic = FileSystemUtils.fileNameOfPath fileName
@@ -811,23 +812,40 @@ let UseMultipleDiagnosticLoggers (inputs, diagnosticsLogger, eagerFormat) f =
             logger.CommitDelayedDiagnostics diagnosticsLogger
 
 let ParseInputFilesInParallel (tcConfig: TcConfig, lexResourceManager, sourceFiles, delayLogger: DiagnosticsLogger, retryLocked) =
+    node {
+        use _ = UseDiagnosticsLogger delayLogger
 
-    let isLastCompiland, isExe = sourceFiles |> tcConfig.ComputeCanContainEntryPoint
+        let isLastCompiland, isExe = sourceFiles |> tcConfig.ComputeCanContainEntryPoint
 
-    for fileName in sourceFiles do
-        checkInputFile tcConfig fileName
+        for fileName in sourceFiles do
+            checkInputFile tcConfig fileName
 
-    let sourceFiles = List.zip sourceFiles isLastCompiland
+        let sourceFiles = List.zip sourceFiles isLastCompiland
 
-    UseMultipleDiagnosticLoggers (sourceFiles, delayLogger, None) (fun sourceFilesWithDelayLoggers ->
-        sourceFilesWithDelayLoggers
-        |> ListParallel.map (fun ((fileName, isLastCompiland), delayLogger) ->
-            let directoryName = Path.GetDirectoryName fileName
+        let! result =
+            sourceFiles
+            |> Seq.map (fun (fileName, isLastCompiland) ->
+                node {
+                    let directoryName = Path.GetDirectoryName fileName
 
-            let input =
-                parseInputFileAux (tcConfig, lexResourceManager, fileName, (isLastCompiland, isExe), delayLogger, retryLocked)
+                    let input =
+                        parseInputFileAux (
+                            tcConfig,
+                            lexResourceManager,
+                            fileName,
+                            (isLastCompiland, isExe),
+                            DiagnosticsThreadStatics.DiagnosticsLogger,
+                            retryLocked
+                        )
 
-            (input, directoryName)))
+                    return (input, directoryName)
+                })
+            |> NodeCode.Parallel
+
+        return result |> Array.toList
+    }
+    |> Async.AwaitNodeCode
+    |> Async.RunImmediate
 
 let ParseInputFilesSequential (tcConfig: TcConfig, lexResourceManager, sourceFiles, diagnosticsLogger: DiagnosticsLogger, retryLocked) =
     let isLastCompiland, isExe = sourceFiles |> tcConfig.ComputeCanContainEntryPoint

src/Compiler/Facilities/AsyncMemoize.fs

@@ -354,15 +354,11 @@ 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 |> Async.AwaitNodeCode
+                                                post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
+                                                return ()
                                             with
                                             | TaskCancelled _ ->
                                                 Interlocked.Increment &cancel_exception_subsequent |> ignore
@@ -500,17 +496,12 @@ 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
-
                                 log (Started, key)
 
-                                try
-                                    let! result = computation |> Async.AwaitNodeCode
-                                    post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
-                                    return result
-                                finally
-                                    DiagnosticsThreadStatics.DiagnosticsLogger <- currentLogger
+                                use _ = UseDiagnosticsLogger cachingLogger
+                                let! result = computation |> Async.AwaitNodeCode
+                                post (key, (JobCompleted(result, cachingLogger.CapturedDiagnostics)))
+                                return result
                             },
                             cancellationToken = linkedCtSource.Token
                         )

src/Compiler/Facilities/BuildGraph.fs

@@ -9,29 +9,28 @@ open System.Diagnostics
 open System.Globalization
 open FSharp.Compiler.DiagnosticsLogger
 open Internal.Utilities.Library
+open Internal.Utilities.Library.Cancellable
 
 [<NoEquality; NoComparison>]
 type NodeCode<'T> = Node of Async<'T>
 
 let wrapThreadStaticInfo computation =
     async {
-        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
-        let phase = DiagnosticsThreadStatics.BuildPhase
+        let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLoggerNC
+        let phase = DiagnosticsThreadStatics.BuildPhaseNC
 
         try
             return! computation
         finally
-            DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-            DiagnosticsThreadStatics.BuildPhase <- phase
+            DiagnosticsThreadStatics.DiagnosticsLoggerNC <- diagnosticsLogger
+            DiagnosticsThreadStatics.BuildPhaseNC <- phase
     }
 
 let unwrapNode (Node(computation)) = computation
 
 type Async<'T> with
 
-    static member AwaitNodeCode(node: NodeCode<'T>) =
-        match node with
-        | Node(computation) -> wrapThreadStaticInfo computation
+    static member AwaitNodeCode(node: NodeCode<'T>) = unwrapNode node
 
 [<Sealed>]
 type NodeCodeBuilder() =
@@ -95,27 +94,8 @@ type NodeCodeBuilder() =
     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
-
-                try
-                    return! binder value |> Async.AwaitNodeCode
-                finally
-                    (value :> IDisposable).Dispose()
-            }
-        )
-
-    [<DebuggerHidden; DebuggerStepThrough>]
-    member _.Using(value: IDisposable, binder: IDisposable -> NodeCode<'U>) =
-        Node(
-            async {
-                use _ = value
-                return! binder value |> Async.AwaitNodeCode
-            }
-        )
+    member this.Using(resource: ('T :> IDisposable), binder: ('T :> IDisposable) -> NodeCode<'U>) =
+        async.Using(resource, binder >> unwrapNode) |> Node
 
 let node = NodeCodeBuilder()
 
@@ -125,49 +105,49 @@ 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.DiagnosticsLoggerNC
+        let phase = DiagnosticsThreadStatics.BuildPhaseNC
 
         try
             try
                 let work =
                     async {
-                        DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                        DiagnosticsThreadStatics.BuildPhase <- phase
-                        return! computation |> Async.AwaitNodeCode
+                        DiagnosticsThreadStatics.DiagnosticsLoggerNC <- diagnosticsLogger
+                        DiagnosticsThreadStatics.BuildPhaseNC <- phase
+                        return! unwrapNode computation
                     }
 
                 Async.StartImmediateAsTask(work, cancellationToken = ct).Result
             finally
-                DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                DiagnosticsThreadStatics.BuildPhase <- phase
+                DiagnosticsThreadStatics.DiagnosticsLoggerNC <- diagnosticsLogger
+                DiagnosticsThreadStatics.BuildPhaseNC <- 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.DiagnosticsLoggerNC
+        let phase = DiagnosticsThreadStatics.BuildPhaseNC
 
         try
             let work =
                 async {
-                    DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-                    DiagnosticsThreadStatics.BuildPhase <- phase
-                    return! computation |> Async.AwaitNodeCode
+                    DiagnosticsThreadStatics.DiagnosticsLoggerNC <- diagnosticsLogger
+                    DiagnosticsThreadStatics.BuildPhaseNC <- phase
+                    return! unwrapNode computation
                 }
 
             Async.StartAsTask(work, cancellationToken = defaultArg ct CancellationToken.None)
         finally
-            DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
-            DiagnosticsThreadStatics.BuildPhase <- phase
+            DiagnosticsThreadStatics.DiagnosticsLoggerNC <- diagnosticsLogger
+            DiagnosticsThreadStatics.BuildPhaseNC <- phase
 
     static member CancellationToken = cancellationToken
 
     static member FromCancellable(computation: Cancellable<'T>) =
-        Node(wrapThreadStaticInfo (Cancellable.toAsync computation))
+        Node(Cancellable.toAsync computation)
 
     static member AwaitAsync(computation: Async<'T>) = Node(wrapThreadStaticInfo computation)
 
@@ -176,6 +156,12 @@ type NodeCode private () =
 
     static member AwaitTask(task: Task) =
         Node(wrapThreadStaticInfo (Async.AwaitTask task))
+
+    static member AwaitTaskWithoutWrapping(task: Task<'T>) =
+        Node((Async.AwaitTask task))
+
+    static member AwaitTaskWithoutWrapping(task: Task) =
+        Node((Async.AwaitTask task))
 
     static member AwaitWaitHandle_ForTesting(waitHandle: WaitHandle) =
         Node(wrapThreadStaticInfo (Async.AwaitWaitHandle(waitHandle)))
@@ -196,26 +182,21 @@ type NodeCode private () =
 
     static member Parallel(computations: NodeCode<'T> seq) =
         node {
-            let concurrentLogging = new CaptureDiagnosticsConcurrently()
             let phase = DiagnosticsThreadStatics.BuildPhase
-            // Why does it return just IDisposable?
-            use _ = concurrentLogging
+            use concurrentLogging = new CaptureDiagnosticsConcurrently()
 
             let injectLogger i computation =
                 let logger = concurrentLogging.GetLoggerForTask($"NodeCode.Parallel {i}")
 
                 async {
-                    DiagnosticsThreadStatics.DiagnosticsLogger <- logger
-                    DiagnosticsThreadStatics.BuildPhase <- phase
-                    return! unwrapNode computation
+                    SetThreadDiagnosticsLoggerNoUnwind logger
+                    SetThreadBuildPhaseNoUnwind phase
+
+                    return! computation |> unwrapNode
                 }
 
-            return!
-                computations
-                |> Seq.mapi injectLogger
-                |> Async.Parallel
-                |> wrapThreadStaticInfo
-                |> Node
+            return! computations |> Seq.mapi injectLogger |> Async.Parallel |> wrapThreadStaticInfo |> Node
+
         }
 
 [<RequireQualifiedAccess>]
@@ -256,6 +237,7 @@ type GraphNode<'T> private (computation: NodeCode<'T>, cachedResult: ValueOption
             cachedResultNode
         else
             node {
+
                 Interlocked.Increment(&requestCount) |> ignore
 
                 try
@@ -278,18 +260,17 @@ type GraphNode<'T> private (computation: NodeCode<'T>, cachedResult: ValueOption
                                      ||| TaskContinuationOptions.NotOnFaulted
                                      ||| TaskContinuationOptions.ExecuteSynchronously)
                                 )
-                            |> NodeCode.AwaitTask
+                            |> NodeCode.AwaitTaskWithoutWrapping
 
                         match cachedResult with
                         | ValueSome value -> return value
                         | _ ->
                             let tcs = TaskCompletionSource<'T>()
-                            let (Node(p)) = computation
 
                             Async.StartWithContinuations(
                                 async {
                                     Thread.CurrentThread.CurrentUICulture <- GraphNode.culture
-                                    return! p
+                                    return! computation |> unwrapNode |> wrapThreadStaticInfo
                                 },
                                 (fun res ->
                                     cachedResult <- ValueSome res
@@ -301,7 +282,7 @@ type GraphNode<'T> private (computation: NodeCode<'T>, cachedResult: ValueOption
                                 ct
                             )
 
-                            return! tcs.Task |> NodeCode.AwaitTask
+                            return! tcs.Task |> NodeCode.AwaitTaskWithoutWrapping
                     finally
                         if taken then
                             semaphore.Release() |> ignore

src/Compiler/Facilities/BuildGraph.fsi

@@ -44,11 +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>
+    member Using: ('T :> IDisposable) * (('T :> IDisposable) -> NodeCode<'U>) -> NodeCode<'U>
 
-    /// A generic 'use' that disposes of the IDisposable at the end of the computation.
-    member Using: IDisposable * (IDisposable -> NodeCode<'T>) -> NodeCode<'T>
 
 /// Specifies code that can be run as part of the build graph.
 val node: NodeCodeBuilder