fix(Caching): Correct Cancellation Handling

CHANGELOG.md

@@ -19,6 +19,8 @@ The `Unreleased` section name is replaced by the expected version of next releas
 ### Removed
 ### Fixed
 
+- `Caching`: Store Load `Task` Cancellation (e.g., as triggered by the Cosmos SDK under rate limiting) leads to perpetual `TaskCancellation` exceptions for cached requests (`LoadOption.AnyCachedValue`, `LoadOption.AllowStale`) [#451](https://github.com/jet/equinox/pull/451)
+
 <a name="4.0.0"></a>
 ## [4.0.0] - 2024-03-20
 

src/Equinox/LazyTask.fs

@@ -23,8 +23,8 @@ type
     member _.TryAwaitValid() = task {
         let t = workflow.Value
 
-        // Determines if the last attempt completed, but failed; For TMI see https://stackoverflow.com/a/33946166/11635
-        if t.IsFaulted then return ValueNone
+        // Determines if the last attempt completed, but failed, or was cancelled (e.g. due to timeout); For TMI see https://stackoverflow.com/a/33946166/11635
+        if t.IsFaulted || t.IsCanceled then return ValueNone
         else
             let! res = t
             return ValueSome res }

tests/Equinox.Core.Tests/LazyTaskTests.fs

@@ -19,3 +19,42 @@ let [<Fact>] ``LazyTask.Empty is a true singleton, does not allocate`` () =
     let i1 = LazyTask<int>.Empty
     let i2 = LazyTask<int>.Empty
     test <@ obj.ReferenceEquals(i1, i2) @>
+
+[<Theory; InlineData false; InlineData true>]
+let ``LazyTask TryAwaitValid fault handling`` immediately = async {
+    let expected = if immediately then "bad beginning" else "bad ending"
+    let cell = LazyTask(fun () -> task {  if immediately then failwith "bad beginning"
+                                          do! Task.Delay 10
+                                          failwith "bad ending" })
+    let res = cell.TryAwaitValid()
+    // We've not awaited it yet, so nothing bad yet
+    test <@ not res.IsFaulted @>
+    let! res = res |> Async.AwaitTaskCorrect |> Async.Catch
+    // Depending on whether there's a continuation in the task, we'll see different outcomes
+    if immediately then Choice1Of2 ValueNone =! res
+    else test <@ match res with Choice2Of2 e -> e.Message = expected | _ -> false @>
+
+    let! res2 = cell.TryAwaitValid() |> Async.AwaitTaskCorrect
+    // next attempt does not propagate the fault
+    res2 =! ValueOption.None }
+
+[<Theory; InlineData false; InlineData true>]
+let ``LazyTask TryAwaitValid cancellation`` immediately = async {
+    let cts = new System.Threading.CancellationTokenSource()
+    let cell: LazyTask<unit> =
+        if immediately then
+            cts.Cancel()
+            LazyTask(fun () -> Task.FromCanceled<unit> cts.Token)
+        else
+            LazyTask(fun () -> task { do! Task.Delay(10000, cts.Token)
+                                      return () })
+    if immediately then
+        let! res = cell.TryAwaitValid() |> Async.AwaitTaskCorrect
+        res =! ValueNone
+    else
+        cts.CancelAfter 100
+        let! res = cell.TryAwaitValid() |> Async.AwaitTaskCorrect |> Async.Catch
+        test <@ match res with Choice2Of2 (:? TaskCanceledException) -> true | _ -> false @>
+    // Next time, we know it's faulted
+    let! res = cell.TryAwaitValid() |> Async.AwaitTaskCorrect
+    res =! ValueOption.None }