OperatorMulticast.connect(connection) should not return null

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

@@ -41,11 +41,13 @@ public final class OperatorMulticast<T, R> extends ConnectableObservable<R> {
     final Observable<? extends T> source;
     final Object guard;
     final Func0<? extends Subject<? super T, ? extends R>> subjectFactory;
-    private final AtomicReference<Subject<? super T, ? extends R>> connectedSubject;
-    private final List<Subscriber<? super R>> waitingForConnect;
+    final AtomicReference<Subject<? super T, ? extends R>> connectedSubject;
+    final List<Subscriber<? super R>> waitingForConnect;
 
     /** Guarded by guard. */
-    Subscriber<T> subscription;
+    private Subscriber<T> subscription;
+    // wraps subscription above for unsubscription using guard
+    private Subscription guardedSubscription;
 
     public OperatorMulticast(Observable<? extends T> source, final Func0<? extends Subject<? super T, ? extends R>> subjectFactory) {
         this(new Object(), new AtomicReference<Subject<? super T, ? extends R>>(), new ArrayList<Subscriber<? super R>>(), source, subjectFactory);
@@ -77,15 +79,13 @@ public void call(Subscriber<? super R> subscriber) {
     public void connect(Action1<? super Subscription> connection) {
         // each time we connect we create a new Subject and Subscription
 
-        boolean shouldSubscribe = false;
-
         // subscription is the state of whether we are connected or not
         synchronized (guard) {
             if (subscription != null) {
-                // already connected, return as there is nothing to do
+                // already connected
+                connection.call(guardedSubscription);
                 return;
             } else {
-                shouldSubscribe = true;
                 // we aren't connected, so let's create a new Subject and connect
                 final Subject<? super T, ? extends R> subject = subjectFactory.call();
                 // create new Subscriber that will pass-thru to the subject we just created
@@ -106,6 +106,26 @@ public void onNext(T args) {
                         subject.onNext(args);
                     }
                 };
+                final AtomicReference<Subscription> gs = new AtomicReference<Subscription>();
+                gs.set(Subscriptions.create(new Action0() {
+                    @Override
+                    public void call() {
+                        Subscription s;
+                        synchronized (guard) {
+                            if ( guardedSubscription == gs.get()) {
+                                s = subscription;
+                                subscription = null;
+                                guardedSubscription = null;
+                                connectedSubject.set(null);
+                            } else 
+                                return;
+                        }
+                        if (s != null) {
+                            s.unsubscribe();
+                        }
+                    }
+                }));
+                guardedSubscription = gs.get();
 
                 // register any subscribers that are waiting with this new subject
                 for(Subscriber<? super R> s : waitingForConnect) {
@@ -116,34 +136,20 @@ public void onNext(T args) {
                 // record the Subject so OnSubscribe can see it
                 connectedSubject.set(subject);
             }
+
         }
 
         // in the lock above we determined we should subscribe, do it now outside the lock
-        if (shouldSubscribe) {
-            // register a subscription that will shut this down
-            connection.call(Subscriptions.create(new Action0() {
-                @Override
-                public void call() {
-                    Subscription s;
-                    synchronized (guard) {
-                        s = subscription;
-                        subscription = null;
-                        connectedSubject.set(null);
-                    }
-                    if (s != null) {
-                        s.unsubscribe();
-                    }
-                }
-            }));
+        // register a subscription that will shut this down
+        connection.call(guardedSubscription);
 
-            // now that everything is hooked up let's subscribe
-            // as long as the subscription is not null (which can happen if already unsubscribed)
-            boolean subscriptionIsNull;
-            synchronized(guard) {
-                subscriptionIsNull = subscription == null;
-            }
-            if (!subscriptionIsNull)
-                source.subscribe(subscription);
+        // now that everything is hooked up let's subscribe
+        // as long as the subscription is not null (which can happen if already unsubscribed)
+        Subscriber<T> sub; 
+        synchronized (guard) {
+            sub = subscription;
         }
+        if (sub != null)
+            source.subscribe(sub);
     }
 }

src/test/java/rx/internal/operators/OnSubscribeMulticastTest.java → src/test/java/rx/internal/operators/OperatorMulticastTest.java

@@ -15,11 +15,13 @@
  */
 package rx.internal.operators;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import org.junit.Assert;
 import org.junit.Test;
 
 import rx.Observer;
@@ -29,7 +31,7 @@
 import rx.subjects.PublishSubject;
 import rx.subjects.Subject;
 
-public class OnSubscribeMulticastTest {
+public class OperatorMulticastTest {
 
     @Test
     public void testMulticast() {
@@ -70,15 +72,17 @@ public void testMulticastConnectTwice() {
 
         source.onNext("one");
 
-        multicasted.connect();
-        multicasted.connect();
-
+        Subscription sub = multicasted.connect();
+        Subscription sub2 = multicasted.connect();
+        
         source.onNext("two");
         source.onCompleted();
 
         verify(observer, never()).onNext("one");
         verify(observer, times(1)).onNext("two");
         verify(observer, times(1)).onCompleted();
+
+        assertEquals(sub, sub2);
 
     }