unsubscribe before retry

Author: petermd

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

@@ -32,11 +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;
 
@@ -71,13 +74,16 @@ public void onError(Throwable 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();
-                        o.unsafeSubscribe(new Subscriber<T>(s) {
+                        retrySub.set(o.unsafeSubscribe(new Subscriber<T>(s) {
 
                             @Override
                             public void onCompleted() {
@@ -88,7 +94,15 @@ public void onCompleted() {
                             public void onError(Throwable e) {
                                 if ((retryCount == INFINITE_RETRY || attempts.get() <= retryCount) && !inner.isUnsubscribed()) {
                                     // retry again
-                                    inner.schedule(_self);
+                                    inner.schedule(new Action1<Inner>() {
+                                        @Override
+                                        public void call(Inner inner)
+                                        {
+                                            // Remove the failed subscription first
+                                            retrySub.get().unsubscribe();
+                                            _self.call(inner);
+                                        }
+                                    });
                                 } else {
                                     // give up and pass the failure
                                     s.onError(e);
@@ -100,7 +114,7 @@ public void onNext(T v) {
                                 s.onNext(v);
                             }
 
-                        });
+                        }));
                     }
                 });
             }

rxjava-core/src/test/java/rx/operators/OperatorRetryTest.java

@@ -19,6 +19,8 @@
 import static org.mockito.Matchers.*;
 import static org.mockito.Mockito.*;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.Test;
@@ -27,6 +29,7 @@
 import rx.Observable;
 import rx.Observable.OnSubscribeFunc;
 import rx.Observer;
+import rx.Subscriber;
 import rx.Subscription;
 import rx.functions.Action1;
 import rx.subjects.PublishSubject;
@@ -147,6 +150,93 @@ public void call(Integer n) {
         assertEquals(1, count.get());
     }
 
+    public static class SlowFuncAlwaysFails implements Observable.OnSubscribe<String> {
+
+        final AtomicInteger nextSeq=new AtomicInteger();
+        final AtomicInteger activeSubs=new AtomicInteger();
+        final AtomicInteger concurrentSubs=new AtomicInteger();
+
+        public void call(final Subscriber<? super String> s)
+        {
+            final int seq=nextSeq.incrementAndGet();
+
+            int cur=activeSubs.incrementAndGet();
+            // Track concurrent subscriptions
+            concurrentSubs.set(Math.max(cur,concurrentSubs.get()));
+
+            // Use async error
+            new Thread(new Runnable() {
+                @Override
+                public void run() {
+                    try {
+                        Thread.sleep(100);
+                    } catch (InterruptedException e) {
+                        // ignore
+                    }
+                    s.onError(new RuntimeException("Subscriber #"+seq+" fails"));
+                }
+            }).start();
+
+            // Track unsubscribes
+            s.add(new Subscription()
+            {
+                private boolean active=true;
+
+                public void unsubscribe()
+                {
+                    if (active) {
+                        activeSubs.decrementAndGet();
+                        active=false;
+                    }
+                }
+
+                public boolean isUnsubscribed()
+                {
+                    return !active;
+                }
+            });
+        }
+    }
+
+    @Test
+    public void testUnsubscribeAfterError() {
+
+        final CountDownLatch check=new CountDownLatch(1);
+        final SlowFuncAlwaysFails sf=new SlowFuncAlwaysFails();
+
+        Observable
+            .create(sf)
+            .retry(4)
+            .subscribe(
+                new Action1<String>()
+                {
+                    @Override
+                    public void call(String v)
+                    {
+                        fail("Should never happen");
+                    }
+                },
+                new Action1<Throwable>()
+                {
+                    public void call(Throwable throwable)
+                    {
+                        check.countDown();
+                    }
+                }
+            );
+
+        try
+        {
+            check.await(1, TimeUnit.SECONDS);
+        } catch (InterruptedException e)
+        {
+            fail("interrupted");
+        }
+
+        assertEquals("5 Subscribers created", 5, sf.nextSeq.get());
+        assertEquals("1 Active Subscriber", 1, sf.concurrentSubs.get());
+    }
+
     @Test
     public void testRetryAllowsSubscriptionAfterAllSubscriptionsUnsubsribed() throws InterruptedException {
         final AtomicInteger subsCount = new AtomicInteger(0);