Yet Another Fix for RetryWhen

Sources/Common/DemandBuffer.swift

@@ -33,7 +33,17 @@ class DemandBuffer<S: Subscriber> {
     init(subscriber: S) {
         self.subscriber = subscriber
     }
-
+
+    /// Signal to the buffer that the upstream has changed
+    /// - Returns: A demand that has already been requested from the previous upstream, but has not yet been processed
+    func attachToNewUpstream() -> Subscribers.Demand {
+        lock.lock()
+        defer { lock.unlock() }
+
+        demandState.sent = demandState.requested
+        return demandState.requested - demandState.processed
+    }
+
     /// Buffer an upstream value to later be forwarded to
     /// the downstream subscriber, once it demands it
     ///

Sources/Common/Sink.swift

@@ -62,7 +62,13 @@ class Sink<Upstream: Publisher, Downstream: Subscriber>: Subscriber {
     }
 
     func receive(subscription: Subscription) {
-        upstreamSubscription = subscription
+        defer { upstreamSubscription = subscription }
+
+        if let upstreamSubscription {
+            upstreamSubscription.cancel()
+            let newDemand = buffer.attachToNewUpstream()
+            subscription.requestIfNeeded(newDemand)
+        }
     }
 
     func receive(_ input: Upstream.Output) -> Subscribers.Demand {

Sources/Operators/RetryWhen.swift

@@ -89,6 +89,9 @@ extension Publishers.RetryWhen {
         }
 
         func cancel() {
+            cancellable?.cancel()
+            cancellable = nil
+            sink?.cancelUpstream()
             sink = nil
         }
     }

Tests/RetryWhenTests.swift

@@ -71,7 +71,93 @@ class RetryWhenTests: XCTestCase {
         XCTAssertEqual(completion, .finished)
         XCTAssertEqual(times, 2)
     }
+
+    func testSuccessfulRetryWithManyRetries() {
+        var times = 0
+        var retriesCount = 0
+        var subscriptionsCount = 0
+        var cancelCount = 0
+        var resultOutput: [Int] = []
+        var completion: Subscribers.Completion<RetryWhenTests.MyError>?
+
+        subscription = Deferred(createPublisher: { () -> AnyPublisher<Int, MyError> in
+            defer { times += 1 }
+            if times == 0 {
+                return Fail<Int, MyError>(error: MyError.someError).eraseToAnyPublisher()
+            } else {
+                return Just(times)
+                    .setFailureType(to: MyError.self)
+                    .handleEvents(
+                        receiveSubscription: { _ in subscriptionsCount += 1 },
+                        receiveCancel: { cancelCount += 1 }
+                    )
+                    .eraseToAnyPublisher()
+            }
+        })
+        .retryWhen { error in
+            return error
+                .handleEvents(receiveOutput: { _ in
+                    retriesCount += 1
+                })
+                .flatMapLatest { _ in [1, 2].publisher }
+        }
+        .sink(
+            receiveCompletion: { completion = $0 },
+            receiveValue: { resultOutput.append($0) }
+        )
+
+        XCTAssertEqual(resultOutput, [2])
+        XCTAssertEqual(completion, .finished)
+        XCTAssertEqual(times, 3)
+        XCTAssertEqual(retriesCount, 1)
+        XCTAssertEqual(subscriptionsCount, 2)
+        XCTAssertEqual(cancelCount, 1)
+    }
+
+    func testSuccessfulRetryWithCustomDemand() {
+        var times = 0
+        var retriesCount = 0
+        var resultOutput: [Int] = []
+        var completion: Subscribers.Completion<RetryWhenTests.MyError>?
+
+        AnyPublisher<Int, MyError>.init { subscriber in
+            defer { times += 1 }
+            if times < 2 {
+                subscriber.send(times)
+                subscriber.send(completion: .failure(MyError.someError))
+            }
+            else {
+                subscriber.send(times)
+                subscriber.send(completion: .finished)
+            }
+
+            return AnyCancellable({})
+        }
+        .retryWhen { error in
+            error
+                .handleEvents(receiveOutput: { _ in retriesCount += 1})
+                .map { _ in }
+        }
+        .subscribe(
+            AnySubscriber(
+                receiveSubscription: { subscription in
+                    self.subscription = AnyCancellable { subscription.cancel() }
+                    subscription.request(.max(3))
+                },
+                receiveValue: {
+                    resultOutput.append($0)
+                    return .none
+                },
+                receiveCompletion: { completion = $0 }
+            )
+        )
 
+        XCTAssertEqual(resultOutput, [0, 1, 2])
+        XCTAssertEqual(completion, .finished)
+        XCTAssertEqual(times, 3)
+        XCTAssertEqual(retriesCount, 2)
+    }
+
     func testRetryFailure() {
         var expectedOutput: Int?