Merge pull request #2567 from akarnokd/RefCountDisconnectOnTerminalEvent

RefCount: disconnect all if upstream terminates

Author: benjchristensen

src/main/java/rx/internal/operators/OnSubscribeRefCount.java

@@ -80,11 +80,8 @@ public void call(final Subscriber<? super T> subscriber) {
             }
         } else {
             try {
-                // handle unsubscribing from the base subscription
-                subscriber.add(disconnect());
-
                 // ready to subscribe to source so do it
-                source.unsafeSubscribe(subscriber);
+                doSubscribe(subscriber, baseSubscription);
             } finally {
                 // release the read lock
                 lock.unlock();
@@ -101,12 +98,8 @@ public void call(Subscription subscription) {
 
                 try {
                     baseSubscription.add(subscription);
-
-                    // handle unsubscribing from the base subscription
-                    subscriber.add(disconnect());
-
                     // ready to subscribe to source so do it
-                    source.unsafeSubscribe(subscriber);
+                    doSubscribe(subscriber, baseSubscription);
                 } finally {
                     // release the write lock
                     lock.unlock();
@@ -115,18 +108,54 @@ public void call(Subscription subscription) {
             }
         };
     }
+    
+    void doSubscribe(final Subscriber<? super T> subscriber, final CompositeSubscription currentBase) {
+        // handle unsubscribing from the base subscription
+        subscriber.add(disconnect(currentBase));
+        
+        source.unsafeSubscribe(new Subscriber<T>(subscriber) {
+            @Override
+            public void onError(Throwable e) {
+                cleanup();
+                subscriber.onError(e);
+            }
+            @Override
+            public void onNext(T t) {
+                subscriber.onNext(t);
+            }
+            @Override
+            public void onCompleted() {
+                cleanup();
+                subscriber.onCompleted();
+            }
+            void cleanup() {
+                lock.lock();
+                try {
+                    if (baseSubscription == currentBase) {
+                        baseSubscription.unsubscribe();
+                        baseSubscription = new CompositeSubscription();
+                        subscriptionCount.set(0);
+                    }
+                } finally {
+                    lock.unlock();
+                }
+            }
+        });
+    }
 
-    private Subscription disconnect() {
+    private Subscription disconnect(final CompositeSubscription current) {
         return Subscriptions.create(new Action0() {
             @Override
             public void call() {
                 lock.lock();
                 try {
-                    if (subscriptionCount.decrementAndGet() == 0) {
-                        baseSubscription.unsubscribe();
-                        // need a new baseSubscription because once
-                        // unsubscribed stays that way
-                        baseSubscription = new CompositeSubscription();
+                    if (baseSubscription == current) {
+                        if (subscriptionCount.decrementAndGet() == 0) {
+                            baseSubscription.unsubscribe();
+                            // need a new baseSubscription because once
+                            // unsubscribed stays that way
+                            baseSubscription = new CompositeSubscription();
+                        }
                     }
                 } finally {
                     lock.unlock();

src/test/java/rx/internal/operators/OnSubscribeRefCountTest.java

@@ -15,39 +15,24 @@
  */
 package rx.internal.operators;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.inOrder;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
+import static org.mockito.Mockito.*;
+
+import java.util.*;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
 
-import org.junit.Before;
-import org.junit.Test;
-import org.mockito.InOrder;
-import org.mockito.MockitoAnnotations;
+import org.junit.*;
+import org.mockito.*;
 
-import rx.Observable;
+import rx.*;
 import rx.Observable.OnSubscribe;
+import rx.Observable;
 import rx.Observer;
-import rx.Subscriber;
-import rx.Subscription;
-import rx.functions.Action0;
-import rx.functions.Action1;
-import rx.functions.Func2;
-import rx.observers.Subscribers;
-import rx.observers.TestSubscriber;
-import rx.schedulers.Schedulers;
-import rx.schedulers.TestScheduler;
+import rx.functions.*;
+import rx.observers.*;
+import rx.schedulers.*;
 import rx.subjects.ReplaySubject;
 import rx.subscriptions.Subscriptions;
 
@@ -285,6 +270,13 @@ public void call() {
         s.assertNoErrors();
     }
 
+    @Test
+    public void testConnectUnsubscribeRaceConditionLoop() throws InterruptedException {
+        for (int i = 0; i < 1000; i++) {
+            testConnectUnsubscribeRaceCondition();
+        }
+    }
+    
     @Test
     public void testConnectUnsubscribeRaceCondition() throws InterruptedException {
         final AtomicInteger subUnsubCount = new AtomicInteger();
@@ -310,12 +302,14 @@ public void call() {
                 });
 
         TestSubscriber<Long> s = new TestSubscriber<Long>();
-        o.publish().refCount().subscribeOn(Schedulers.newThread()).subscribe(s);
+        
+        o.publish().refCount().subscribeOn(Schedulers.computation()).subscribe(s);
         System.out.println("send unsubscribe");
         // now immediately unsubscribe while subscribeOn is racing to subscribe
         s.unsubscribe();
         // this generally will mean it won't even subscribe as it is already unsubscribed by the time connect() gets scheduled
-
+        // give time to the counter to update
+        Thread.sleep(1);
         // either we subscribed and then unsubscribed, or we didn't ever even subscribe
         assertEquals(0, subUnsubCount.get());
 
@@ -532,4 +526,72 @@ public Integer call(Integer t1, Integer t2) {
         ts2.assertReceivedOnNext(Arrays.asList(30));
     }
 
+    @Test(timeout = 10000)
+    public void testUpstreamErrorAllowsRetry() throws InterruptedException {
+        final AtomicInteger intervalSubscribed = new AtomicInteger();
+        Observable<String> interval =
+                Observable.interval(200,TimeUnit.MILLISECONDS)
+                        .doOnSubscribe(
+                                new Action0() {
+                                    @Override
+                                    public void call() {
+                                        System.out.println("Subscribing to interval " + intervalSubscribed.incrementAndGet());
+                                    }
+                                }
+                         )
+                        .flatMap(new Func1<Long, Observable<String>>() {
+                            @Override
+                            public Observable<String> call(Long t1) {
+                                return Observable.defer(new Func0<Observable<String>>() {
+                                    @Override
+                                    public Observable<String> call() {
+                                        return Observable.<String>error(new Exception("Some exception"));
+                                    }
+                                });
+                            }
+                        })
+                        .onErrorResumeNext(new Func1<Throwable, Observable<String>>() {
+                            @Override
+                            public Observable<String> call(Throwable t1) {
+                                return Observable.error(t1);
+                            }
+                        })
+                        .publish()
+                        .refCount();
+
+        interval
+                .doOnError(new Action1<Throwable>() {
+                    @Override
+                    public void call(Throwable t1) {
+                        System.out.println("Subscriber 1 onError: " + t1);
+                    }
+                })
+                .retry(5)
+                .subscribe(new Action1<String>() {
+                    @Override
+                    public void call(String t1) {
+                        System.out.println("Subscriber 1: " + t1);
+                    }
+                });
+        Thread.sleep(100);
+        interval
+        .doOnError(new Action1<Throwable>() {
+            @Override
+            public void call(Throwable t1) {
+                System.out.println("Subscriber 2 onError: " + t1);
+            }
+        })
+        .retry(5)
+                .subscribe(new Action1<String>() {
+                    @Override
+                    public void call(String t1) {
+                        System.out.println("Subscriber 2: " + t1);
+                    }
+                });
+        
+        Thread.sleep(1300);
+        
+        System.out.println(intervalSubscribed.get());
+        assertEquals(6, intervalSubscribed.get());
+    }
 }