Improve AsyncMemoize tests (#16580)

dotnet · Jan 24, 2024 · 61a5792 · 61a5792
1 parent 2caef37
commit 61a5792
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 14 deletions.
diff --git a/src/Compiler/Facilities/AsyncMemoize.fs b/src/Compiler/Facilities/AsyncMemoize.fs
@@ -74,7 +74,9 @@ type internal Job<'TValue> =
         | Failed(_, ex) -> $"Failed {ex}"
 
 type internal JobEvent =
+    | Requested
     | Started
+    | Restarted
     | Finished
     | Canceled
     | Evicted
@@ -245,7 +247,7 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
 
                 let cached, otherVersions = cache.GetAll(key.Key, key.Version)
 
-                return
+                let result =
                     match msg, cached with
                     | GetOrCompute _, Some(Completed(result, diags)) ->
                         Interlocked.Increment &hits |> ignore
@@ -298,6 +300,9 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
                                 cts.Cancel())
 
                         New cts.Token
+
+                log (Requested, key)
+                return result
             })
 
     let internalError key message =
@@ -346,7 +351,7 @@ type internal AsyncMemoize<'TKey, 'TVersion, 'TValue when 'TKey: equality and 'T
 
                                             try
                                                 // TODO: Should unify starting and restarting
-                                                log (Started, key)
+                                                log (Restarted, key)
                                                 Interlocked.Increment &restarted |> ignore
                                                 System.Diagnostics.Trace.TraceInformation $"{name} Restarted {key.Label}"
                                                 let currentLogger = DiagnosticsThreadStatics.DiagnosticsLogger

diff --git a/src/Compiler/Facilities/AsyncMemoize.fsi b/src/Compiler/Facilities/AsyncMemoize.fsi
@@ -12,7 +12,9 @@ module internal Utils =
     val (|TaskCancelled|_|): ex: exn -> TaskCanceledException option
 
 type internal JobEvent =
+    | Requested
     | Started
+    | Restarted
     | Finished
     | Canceled
     | Evicted

diff --git a/tests/FSharp.Compiler.ComponentTests/CompilerService/AsyncMemoize.fs b/tests/FSharp.Compiler.ComponentTests/CompilerService/AsyncMemoize.fs
@@ -61,7 +61,7 @@ let ``Basics``() =
     let groups = eventLog |> Seq.groupBy snd |> Seq.toList
     Assert.Equal(3, groups.Length)
     for key, events in groups do
-        Assert.Equal<Set<(JobEvent * int)>>(Set [ Started, key; Finished, key ], Set events)
+        Assert.Equal<Set<(JobEvent * int)>>(Set [ Requested, key; Started, key; Finished, key ], Set events)
 
 [<Fact>]
 let ``We can cancel a job`` () =
@@ -96,9 +96,14 @@ let ``We can cancel a job`` () =
         waitFor jobStarted
         jobStarted.Reset() |> ignore
 
+        let jobRequested = new ManualResetEvent(false)
+        memoize.OnEvent(fun (e, _) -> if e = Requested then jobRequested.Set() |> ignore)
+
         let _task2 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation ignore), ct = cts2.Token)
         let _task3 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation ignore), ct = cts3.Token)
 
+        waitFor jobRequested
+
         cts1.Cancel()
         cts2.Cancel()
 
@@ -108,7 +113,14 @@ let ``We can cancel a job`` () =
 
         waitFor jobCanceled
 
-        Assert.Equal<(JobEvent * int) array>([| Started, key; Started, key; Canceled, key |], eventLog |> Seq.toArray )
+        Assert.Equal<(JobEvent * int) array>([|
+            Requested, key
+            Started, key
+            Requested, key
+            Requested, key
+            Restarted, key
+            Canceled, key
+        |], eventLog |> Seq.toArray )
     }
 
 [<Fact>]
@@ -139,22 +151,31 @@ let ``Job is restarted if first requestor cancels`` () =
         waitFor jobStarted
         jobStarted.Reset() |> ignore
 
+        let jobRequested = new ManualResetEvent(false)
+        memoize.OnEvent(fun (e, _) -> if e = Requested then jobRequested.Set() |> ignore)
+
         let _task2 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation key), ct = cts2.Token)
         let _task3 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation key), ct = cts3.Token)
 
+        waitFor jobRequested
+
         cts1.Cancel()
 
         waitFor jobStarted
 
-        cts3.Cancel()
-
         jobCanComplete.Set() |> ignore
 
         let! result = _task2
         Assert.Equal(2, result)
 
         let orderedLog = eventLog |> Seq.rev |> Seq.toList
-        let expected = [ Started, key; Started, key; Finished, key ]
+        let expected = [
+            Requested, key
+            Started, key
+            Requested, key
+            Requested, key
+            Restarted, key
+            Finished, key ]
 
         Assert.Equal<_ list>(expected, orderedLog)
     }
@@ -184,15 +205,20 @@ let ``Job is restarted if first requestor cancels but keeps running if second re
 
         let _task1 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation key), ct = cts1.Token)
 
-        jobStarted.WaitOne() |> ignore
+        waitFor jobStarted
         jobStarted.Reset() |> ignore
 
+        let jobRequested = new ManualResetEvent(false)
+        memoize.OnEvent(fun (e, _) -> if e = Requested then jobRequested.Set() |> ignore)
+
         let _task2 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation key), ct = cts2.Token)
         let _task3 = NodeCode.StartAsTask_ForTesting( memoize.Get'(key, computation key), ct = cts3.Token)
 
+        waitFor jobRequested
+
         cts1.Cancel()
 
-        jobStarted.WaitOne() |> ignore
+        waitFor jobStarted
 
         cts2.Cancel()
 
@@ -202,7 +228,13 @@ let ``Job is restarted if first requestor cancels but keeps running if second re
         Assert.Equal(2, result)
 
         let orderedLog = eventLog |> Seq.rev |> Seq.toList
-        let expected = [ Started, key; Started, key; Finished, key ]
+        let expected = [
+            Requested, key
+            Started, key
+            Requested, key
+            Requested, key
+            Restarted, key
+            Finished, key ]
 
         Assert.Equal<_ list>(expected, orderedLog)
     }
@@ -228,7 +260,7 @@ let ``Stress test`` () =
     let keyCount = rng.Next(5, 200)
     let keys = [| 1 .. keyCount |]
 
-    let testTimeoutMs = threads * iterations * maxDuration
+    let testTimeoutMs = threads * iterations * maxDuration * 2
 
     let intenseComputation durationMs result =
         async {

diff --git a/tests/FSharp.Compiler.ComponentTests/FSharpChecker/TransparentCompiler.fs b/tests/FSharp.Compiler.ComponentTests/FSharpChecker/TransparentCompiler.fs
@@ -237,8 +237,8 @@ let ``File is not checked twice`` () =
         |> Seq.map (fun (k, g) -> k, g |> Seq.map fst |> Seq.toList)
         |> Map
 
-    Assert.Equal<JobEvent list>([Weakened; Started; Finished], intermediateTypeChecks["FileFirst.fs"])
-    Assert.Equal<JobEvent list>([Weakened; Started; Finished], intermediateTypeChecks["FileThird.fs"])
+    Assert.Equal<JobEvent list>([Weakened; Requested; Started; Finished], intermediateTypeChecks["FileFirst.fs"])
+    Assert.Equal<JobEvent list>([Weakened; Requested; Started; Finished], intermediateTypeChecks["FileThird.fs"])
 
 [<Fact>]
 let ``If a file is checked as a dependency it's not re-checked later`` () =
@@ -261,7 +261,7 @@ let ``If a file is checked as a dependency it's not re-checked later`` () =
         |> Seq.map (fun (k, g) -> k, g |> Seq.map fst |> Seq.toList)
         |> Map
 
-    Assert.Equal<JobEvent list>([Weakened; Started; Finished], intermediateTypeChecks["FileThird.fs"])
+    Assert.Equal<JobEvent list>([Weakened; Requested; Started; Finished; Requested], intermediateTypeChecks["FileThird.fs"])
 
 
 // [<Fact>] TODO: differentiate complete and minimal checking requests