Fixed synchronous ConnectableObservable.connect problem

rxjava-core/src/main/java/rx/observables/ConnectableObservable.java

@@ -18,6 +18,7 @@
 import rx.Observable;
 import rx.Subscriber;
 import rx.Subscription;
+import rx.functions.Action1;
 import rx.operators.OperatorRefCount;
 
 /**
@@ -44,10 +45,28 @@ protected ConnectableObservable(OnSubscribe<T> onSubscribe) {
     /**
      * Call a ConnectableObservable's connect() method to instruct it to begin emitting the
      * items from its underlying {@link Observable} to its {@link Subscriber}s.
+     * <p>To disconnect from a synchronous source, use the {@link #connect(rx.functions.Action1)}
+     * method.
      * @return the subscription representing the connection
      */
-    public abstract Subscription connect();
-
+    public final Subscription connect() {
+        final Subscription[] out = new Subscription[1];
+        connect(new Action1<Subscription>() {
+            @Override
+            public void call(Subscription t1) {
+                out[0] = t1;
+            }
+        });
+        return out[0];
+    }
+    /**
+     * Call a ConnectableObservable's connect() method to instruct it to begin emitting the
+     * items from its underlying {@link Observable} to its {@link Subscriber}s.
+     * @param connection the action that receives the connection subscription
+     * before the subscription to source happens allowing the caller
+     * to synchronously disconnect a synchronous source.
+     */
+    public abstract void connect(Action1<? super Subscription> connection);
     /**
      * Returns an observable sequence that stays connected to the source as long
      * as there is at least one subscription to the observable sequence.

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

@@ -19,6 +19,7 @@
 import rx.Subscriber;
 import rx.Subscription;
 import rx.functions.Action0;
+import rx.functions.Action1;
 import rx.observables.ConnectableObservable;
 import rx.subjects.Subject;
 import rx.subscriptions.Subscriptions;
@@ -48,7 +49,20 @@ public void call(Subscriber<? super R> subscriber) {
     }
 
     @Override
-    public Subscription connect() {
+    public void connect(Action1<? super Subscription> connection) {
+        connection.call(Subscriptions.create(new Action0() {
+            @Override
+            public void call() {
+                Subscription s;
+                synchronized (guard) {
+                    s = subscription;
+                    subscription = null;
+                }
+                if (s != null) {
+                    s.unsubscribe();
+                }
+            }
+        }));
         Subscriber<T> s = null;
         synchronized (guard) {
             if (subscription == null) {
@@ -74,20 +88,5 @@ public void onNext(T args) {
         if (s != null) {
             source.unsafeSubscribe(s);
         }
-
-        return Subscriptions.create(new Action0() {
-            @Override
-            public void call() {
-                Subscription s;
-                synchronized (guard) {
-                    s = subscription;
-                    subscription = null;
-                }
-                if (s != null) {
-                    s.unsubscribe();
-                }
-            }
-        });
     }
-
 }

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

@@ -18,12 +18,13 @@
 import rx.Observable;
 import rx.Observable.OnSubscribe;
 import rx.Subscriber;
+import rx.Subscription;
+import rx.functions.Action1;
 import rx.functions.Func0;
 import rx.functions.Func1;
 import rx.observables.ConnectableObservable;
 import rx.observers.SafeSubscriber;
 import rx.subjects.Subject;
-import rx.subscriptions.CompositeSubscription;
 
 /**
  * Returns an observable sequence that contains the elements of a sequence
@@ -63,14 +64,16 @@ public void call(Subscriber<? super TResult> child) {
             return;
         }
 
-        CompositeSubscription csub = new CompositeSubscription();
-        child.add(csub);
-
-        SafeSubscriber<TResult> s = new SafeSubscriber<TResult>(child);
+        final SafeSubscriber<TResult> s = new SafeSubscriber<TResult>(child);
 
         observable.unsafeSubscribe(s);
 
-        csub.add(connectable.connect());
+        connectable.connect(new Action1<Subscription>() {
+            @Override
+            public void call(Subscription t1) {
+                s.add(t1);
+            }
+        });
     }
 
 }

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

@@ -24,13 +24,17 @@
 
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.Assert;
+import static org.junit.Assert.assertEquals;
 import org.junit.Test;
 import org.mockito.InOrder;
 
 import rx.Observable;
 import rx.Observer;
+import rx.functions.Action0;
+import rx.functions.Action1;
 import rx.functions.Func1;
 import rx.observables.ConnectableObservable;
 import rx.operators.OperatorReplay.VirtualBoundedList;
@@ -470,4 +474,50 @@ public void testWindowedReplayError() {
             verify(observer1, never()).onCompleted();
         }
     }
+    @Test
+    public void testSynchronousDisconnect() {
+        final AtomicInteger effectCounter = new AtomicInteger();
+        Observable<Integer> source = Observable.from(1, 2, 3, 4)
+        .doOnNext(new Action1<Integer>() {
+            @Override
+            public void call(Integer v) {
+                effectCounter.incrementAndGet();
+                System.out.println("Sideeffect #" + v);
+            }
+        });
+
+        Observable<Integer> result = source.replay(
+        new Func1<Observable<Integer>, Observable<Integer>>() {
+            @Override
+            public Observable<Integer> call(Observable<Integer> o) {
+                return o.take(2);
+            }
+        });
+
+        for (int i = 1; i < 3; i++) {
+            effectCounter.set(0);
+            System.out.printf("- %d -%n", i);
+            result.subscribe(new Action1<Integer>() {
+
+                @Override
+                public void call(Integer t1) {
+                    System.out.println(t1);
+                }
+
+            }, new Action1<Throwable>() {
+
+                @Override
+                public void call(Throwable t1) {
+                    t1.printStackTrace();
+                }
+            }, 
+            new Action0() {
+                @Override
+                public void call() {
+                    System.out.println("Done");
+                }
+            });
+            assertEquals(2, effectCounter.get());
+        }
+    }
 }