From a78d3804c2c84f0a3178648138442c9b07665bda Mon Sep 17 00:00:00 2001
From: swithek <simon.devs@protonmail.com>
Date: Mon, 31 May 2021 04:50:55 +0300
Subject: [PATCH] Remove context deadline checking in backOffContext (#113)

* Remove ctx deadline check

* Add context error checking after stop
---
 context.go    | 6 +-----
 retry.go      | 4 ++++
 retry_test.go | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/context.go b/context.go
index fcff86c..4848233 100644
--- a/context.go
+++ b/context.go
@@ -57,10 +57,6 @@ func (b *backOffContext) NextBackOff() time.Duration {
 	case <-b.ctx.Done():
 		return Stop
 	default:
+		return b.BackOff.NextBackOff()
 	}
-	next := b.BackOff.NextBackOff()
-	if deadline, ok := b.ctx.Deadline(); ok && deadline.Sub(time.Now()) < next { // nolint: gosimple
-		return Stop
-	}
-	return next
 }
diff --git a/retry.go b/retry.go
index 73234a4..1ce2507 100644
--- a/retry.go
+++ b/retry.go
@@ -62,6 +62,10 @@ func RetryNotifyWithTimer(operation Operation, b BackOff, notify Notify, t Timer
 		}
 
 		if next = b.NextBackOff(); next == Stop {
+			if cerr := ctx.Err(); cerr != nil {
+				return cerr
+			}
+
 			return err
 		}
 
diff --git a/retry_test.go b/retry_test.go
index 19e93db..4d328f0 100644
--- a/retry_test.go
+++ b/retry_test.go
@@ -81,7 +81,7 @@ func TestRetryContext(t *testing.T) {
 	if err == nil {
 		t.Errorf("error is unexpectedly nil")
 	}
-	if err.Error() != "error (3)" {
+	if !errors.Is(err, context.Canceled) {
 		t.Errorf("unexpected error: %s", err.Error())
 	}
 	if i != cancelOn {