From b65b69ba154c51eef9236bc56d5acca4ef008c52 Mon Sep 17 00:00:00 2001
From: Colin Bradley <colin@colinbradley.co.uk>
Date: Sat, 22 Mar 2025 18:14:28 +0000
Subject: [PATCH] SimpleRetryHandler - sleep now uses CancellationToken. closes
 #1031

---
 .../RetryHandlers/SimpleRetryHandlerTest.cs   | 45 +++++++++++++++++++
 .../RetryHandlers/SimpleRetryHandler.cs       | 15 ++++---
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/SpotifyAPI.Web.Tests/RetryHandlers/SimpleRetryHandlerTest.cs b/SpotifyAPI.Web.Tests/RetryHandlers/SimpleRetryHandlerTest.cs
index 14fdb1c7f..07943ac79 100644
--- a/SpotifyAPI.Web.Tests/RetryHandlers/SimpleRetryHandlerTest.cs
+++ b/SpotifyAPI.Web.Tests/RetryHandlers/SimpleRetryHandlerTest.cs
@@ -182,6 +182,51 @@ public async Task HandleRetry_DirectSuccess()
       setup.Sleep.Verify(s => s(TimeSpan.FromMilliseconds(50)), Times.Exactly(0));
     }
 
+    [Test]
+    public async Task HandleRetry_CancellationTokenHonoredInSleep()
+    {
+      var setup = new Setup();
+      setup.Response.SetupGet(r => r.StatusCode).Returns(HttpStatusCode.TooManyRequests);
+      setup.Response.SetupGet(r => r.Headers).Returns(new Dictionary<string, string> {
+        { "Retry-After", "5" }
+      });
+
+      var retryCalled = 0;
+      setup.Retry = (request, ct) =>
+      {
+        retryCalled++;
+        return Task.FromResult(setup.Response.Object);
+      };
+
+      var handler = new SimpleRetryHandler()
+      {
+        TooManyRequestsConsumesARetry = true,
+        RetryTimes = 0,
+        RetryAfter = TimeSpan.FromSeconds(1)
+      };
+
+      var cancellationTokenSource = new CancellationTokenSource();
+
+      var startTime = DateTimeOffset.UtcNow;
+      try
+      {
+        var attemptTask = handler.HandleRetry(setup.Request.Object, setup.Response.Object, setup.Retry, cancellationTokenSource.Token);
+
+        // Wait enough that we're probably in the sleep
+        await Task.Delay(TimeSpan.FromMilliseconds(100));
+
+        cancellationTokenSource.Cancel();
+
+        await attemptTask;
+      }
+      catch (OperationCanceledException)
+      {
+        //Expected
+      }
+
+      Assert.That(DateTimeOffset.UtcNow - startTime, Is.LessThan(TimeSpan.FromSeconds(4)));
+    }
+
     private class Setup
     {
       public Mock<Func<TimeSpan, Task>> Sleep { get; set; } = new Mock<Func<TimeSpan, Task>>();
diff --git a/SpotifyAPI.Web/RetryHandlers/SimpleRetryHandler.cs b/SpotifyAPI.Web/RetryHandlers/SimpleRetryHandler.cs
index 96f849c31..c6cdc2636 100644
--- a/SpotifyAPI.Web/RetryHandlers/SimpleRetryHandler.cs
+++ b/SpotifyAPI.Web/RetryHandlers/SimpleRetryHandler.cs
@@ -10,10 +10,10 @@ namespace SpotifyAPI.Web
 {
   public class SimpleRetryHandler : IRetryHandler
   {
-    private readonly Func<TimeSpan, Task> _sleep;
+    private readonly Func<TimeSpan, CancellationToken, Task> _sleep;
 
     /// <summary>
-    ///     Specifies after how many miliseconds should a failed request be retried.
+    ///     Specifies after how many milliseconds should a failed request be retried.
     /// </summary>
     public TimeSpan RetryAfter { get; set; }
 
@@ -38,10 +38,11 @@ public class SimpleRetryHandler : IRetryHandler
     ///   the Retry-After header
     /// </summary>
     /// <returns></returns>
-    public SimpleRetryHandler() : this(Task.Delay) { }
-    public SimpleRetryHandler(Func<TimeSpan, Task> sleep)
+    public SimpleRetryHandler() : this(sleepWithCancel: Task.Delay) { }
+    public SimpleRetryHandler(Func<TimeSpan, Task> sleep) : this((t, _) => sleep(t)) { }
+    public SimpleRetryHandler(Func<TimeSpan, CancellationToken, Task> sleepWithCancel)
     {
-      _sleep = sleep;
+      _sleep = sleepWithCancel;
       RetryAfter = TimeSpan.FromMilliseconds(50);
       RetryTimes = 10;
       TooManyRequestsConsumesARetry = false;
@@ -88,7 +89,7 @@ private async Task<IResponse> HandleRetryInternally(
       var secondsToWait = ParseTooManyRetries(response);
       if (secondsToWait != null && (!TooManyRequestsConsumesARetry || triesLeft > 0))
       {
-        await _sleep(secondsToWait.Value).ConfigureAwait(false);
+        await _sleep(secondsToWait.Value, cancel).ConfigureAwait(false);
         response = await retry(request, cancel).ConfigureAwait(false);
         var newTriesLeft = TooManyRequestsConsumesARetry ? triesLeft - 1 : triesLeft;
         return await HandleRetryInternally(request, response, retry, newTriesLeft, cancel).ConfigureAwait(false);
@@ -96,7 +97,7 @@ private async Task<IResponse> HandleRetryInternally(
 
       while (RetryErrorCodes.Contains(response.StatusCode) && triesLeft > 0)
       {
-        await _sleep(RetryAfter).ConfigureAwait(false);
+        await _sleep(RetryAfter, cancel).ConfigureAwait(false);
         response = await retry(request, cancel).ConfigureAwait(false);
         return await HandleRetryInternally(request, response, retry, triesLeft - 1, cancel).ConfigureAwait(false);
       }