Unsubscribe on retry using SerialSubscription

Author: benjchristensen

rxjava-core/src/main/java/rx/operators/OperatorRetry.java

@@ -32,16 +32,14 @@
  */
 
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
 
 import rx.Observable;
 import rx.Observable.Operator;
-import rx.Scheduler;
 import rx.Scheduler.Inner;
 import rx.Subscriber;
-import rx.Subscription;
 import rx.functions.Action1;
 import rx.schedulers.Schedulers;
+import rx.subscriptions.SerialSubscription;
 
 public class OperatorRetry<T> implements Operator<T, Observable<T>> {
 
@@ -58,63 +56,59 @@ public OperatorRetry() {
     }
 
     @Override
-    public Subscriber<? super Observable<T>> call(final Subscriber<? super T> s) {
-        return new Subscriber<Observable<T>>(s) {
+    public Subscriber<? super Observable<T>> call(final Subscriber<? super T> child) {
+        final SerialSubscription serialSubscription = new SerialSubscription();
+        // add serialSubscription so it gets unsubscribed if child is unsubscribed
+        child.add(serialSubscription);
+        return new Subscriber<Observable<T>>(child) {
             final AtomicInteger attempts = new AtomicInteger(0);
-            
+
             @Override
             public void onCompleted() {
                 // ignore as we expect a single nested Observable<T>
             }
 
             @Override
             public void onError(Throwable e) {
-                s.onError(e);
+                child.onError(e);
             }
 
             @Override
             public void onNext(final Observable<T> o) {
-
-                final AtomicReference<Subscription> retrySub=new AtomicReference<Subscription>();
-
                 Schedulers.trampoline().schedule(new Action1<Inner>() {
 
                     @Override
                     public void call(final Inner inner) {
                         final Action1<Inner> _self = this;
                         attempts.incrementAndGet();
-                        retrySub.set(o.unsafeSubscribe(new Subscriber<T>(s) {
+
+                        Subscriber<T> subscriber = new Subscriber<T>(child) {
 
                             @Override
                             public void onCompleted() {
-                                s.onCompleted();
+                                child.onCompleted();
                             }
 
                             @Override
                             public void onError(Throwable e) {
                                 if ((retryCount == INFINITE_RETRY || attempts.get() <= retryCount) && !inner.isUnsubscribed()) {
                                     // retry again
-                                    inner.schedule(new Action1<Inner>() {
-                                        @Override
-                                        public void call(Inner inner)
-                                        {
-                                            // Remove the failed subscription first
-                                            retrySub.get().unsubscribe();
-                                            _self.call(inner);
-                                        }
-                                    });
+                                    inner.schedule(_self);
                                 } else {
                                     // give up and pass the failure
-                                    s.onError(e);
+                                    child.onError(e);
                                 }
                             }
 
                             @Override
                             public void onNext(T v) {
-                                s.onNext(v);
+                                child.onNext(v);
                             }
 
-                        }));
+                        };
+                        // register this Subscription (and unsubscribe previous if exists) 
+                        serialSubscription.set(subscriber);
+                        o.unsafeSubscribe(subscriber);
                     }
                 });
             }