From 98f63d123cc8c6017db1eec48a608af2cf9f568d Mon Sep 17 00:00:00 2001
From: Nikolay Kasyanov <136644+nikolaykasyanov@users.noreply.github.com>
Date: Fri, 4 Oct 2024 20:08:10 +0200
Subject: [PATCH] Don't retain receiver of Completable.andThen beyond its
 completion (#2604)

---
 .../Completable+AndThen.swift                 | 16 ++---
 Sources/AllTestz/main.swift                   |  2 +
 Tests/RxSwiftTests/Completable+AndThen.swift  | 67 +++++++++++++++++++
 3 files changed, 77 insertions(+), 8 deletions(-)

diff --git a/RxSwift/Traits/PrimitiveSequence/Completable+AndThen.swift b/RxSwift/Traits/PrimitiveSequence/Completable+AndThen.swift
index f63c3f17f..dd1dc9723 100644
--- a/RxSwift/Traits/PrimitiveSequence/Completable+AndThen.swift
+++ b/RxSwift/Traits/PrimitiveSequence/Completable+AndThen.swift
@@ -70,8 +70,8 @@ final private class ConcatCompletable<Element>: Producer<Element> {
     }
 
     override func run<Observer: ObserverType>(_ observer: Observer, cancel: Cancelable) -> (sink: Disposable, subscription: Disposable) where Observer.Element == Element {
-        let sink = ConcatCompletableSink(parent: self, observer: observer, cancel: cancel)
-        let subscription = sink.run()
+        let sink = ConcatCompletableSink(second: second, observer: observer, cancel: cancel)
+        let subscription = sink.run(completable: completable)
         return (sink: sink, subscription: subscription)
     }
 }
@@ -82,11 +82,11 @@ final private class ConcatCompletableSink<Observer: ObserverType>
     typealias Element = Never
     typealias Parent = ConcatCompletable<Observer.Element>
 
-    private let parent: Parent
+    private let second: Observable<Observer.Element>
     private let subscription = SerialDisposable()
     
-    init(parent: Parent, observer: Observer, cancel: Cancelable) {
-        self.parent = parent
+    init(second: Observable<Observer.Element>, observer: Observer, cancel: Cancelable) {
+        self.second = second
         super.init(observer: observer, cancel: cancel)
     }
 
@@ -99,14 +99,14 @@ final private class ConcatCompletableSink<Observer: ObserverType>
             break
         case .completed:
             let otherSink = ConcatCompletableSinkOther(parent: self)
-            self.subscription.disposable = self.parent.second.subscribe(otherSink)
+            self.subscription.disposable = self.second.subscribe(otherSink)
         }
     }
 
-    func run() -> Disposable {
+    func run(completable: Observable<Never>) -> Disposable {
         let subscription = SingleAssignmentDisposable()
         self.subscription.disposable = subscription
-        subscription.setDisposable(self.parent.completable.subscribe(self))
+        subscription.setDisposable(completable.subscribe(self))
         return self.subscription
     }
 }
diff --git a/Sources/AllTestz/main.swift b/Sources/AllTestz/main.swift
index 4a0dc0a4b..ab665d5fd 100644
--- a/Sources/AllTestz/main.swift
+++ b/Sources/AllTestz/main.swift
@@ -105,6 +105,8 @@ final class CompletableAndThenTest_ : CompletableAndThenTest, RxTestCase {
     ("testCompletableCompleted_CompletableCompleted", CompletableAndThenTest.testCompletableCompleted_CompletableCompleted),
     ("testCompletableError_CompletableCompleted", CompletableAndThenTest.testCompletableError_CompletableCompleted),
     ("testCompletableCompleted_CompletableError", CompletableAndThenTest.testCompletableCompleted_CompletableError),
+    ("testCompletable_FirstCompletableNotRetainedBeyondCompletion", CompletableAndThenTest.testCompletable_FirstCompletableNotRetainedBeyondCompletion),
+    ("testCompletable_FirstCompletableNotRetainedBeyondFailure", CompletableAndThenTest.testCompletable_FirstCompletableNotRetainedBeyondFailure),
     ("testCompletableEmpty_SingleCompleted", CompletableAndThenTest.testCompletableEmpty_SingleCompleted),
     ("testCompletableCompleted_SingleNormal", CompletableAndThenTest.testCompletableCompleted_SingleNormal),
     ("testCompletableError_SingleNormal", CompletableAndThenTest.testCompletableError_SingleNormal),
diff --git a/Tests/RxSwiftTests/Completable+AndThen.swift b/Tests/RxSwiftTests/Completable+AndThen.swift
index 2ea8c26b2..5ae78cf02 100644
--- a/Tests/RxSwiftTests/Completable+AndThen.swift
+++ b/Tests/RxSwiftTests/Completable+AndThen.swift
@@ -120,6 +120,70 @@ extension CompletableAndThenTest {
             ])
     }
 
+    func testCompletable_FirstCompletableNotRetainedBeyondCompletion() {
+        let scheduler = TestScheduler(initialClock: 0)
+
+        let x: TestableObservable<Never> = scheduler.createHotObservable([
+            .completed(210),
+        ])
+
+        var object = Optional.some(TestObject())
+
+        var completable = x.asCompletable()
+            .do(onCompleted: { [object] in
+                _ = object
+            })
+
+        let disposable = completable
+            .andThen(.never())
+            .subscribe()
+
+        defer {
+            disposable.dispose()
+        }
+
+        // completable has completed by now
+        scheduler.advanceTo(300)
+
+        weak var weakObject = object
+        object = nil
+        completable = .never()
+
+        XCTAssertNil(weakObject)
+    }
+
+    func testCompletable_FirstCompletableNotRetainedBeyondFailure() {
+        let scheduler = TestScheduler(initialClock: 0)
+
+        let x: TestableObservable<Never> = scheduler.createHotObservable([
+            .error(210, TestError.dummyError),
+        ])
+
+        var object = Optional.some(TestObject())
+
+        var completable = x.asCompletable()
+            .do(onCompleted: { [object] in
+                _ = object
+            })
+
+        let disposable = completable
+            .andThen(.never())
+            .subscribe()
+
+        defer {
+            disposable.dispose()
+        }
+
+        // completable has terminated with error by now
+        scheduler.advanceTo(300)
+
+        weak var weakObject = object
+        object = nil
+        completable = .never()
+
+        XCTAssertNil(weakObject)
+    }
+
     #if TRACE_RESOURCES
         func testAndThenCompletableReleasesResourcesOnComplete() {
             _ = Completable.empty().andThen(Completable.empty()).subscribe()
@@ -575,3 +639,6 @@ extension CompletableAndThenTest {
         }
     #endif
 }
+
+private class TestObject {
+}