PublishSubject Add Before onSubscribe

Proposed change to register with state before onSubscribe so that "doOnSubscribe" side-effects are safe.

Author: benjchristensen

src/main/java/io/reactivex/subjects/PublishSubject.java

@@ -183,26 +183,29 @@ static final class State<T> extends AtomicReference<Object> implements Publisher
                 AtomicReferenceFieldUpdater.newUpdater(State.class, PublishSubscriber[].class, "subscribers");
 
         @Override
-        public void subscribe(Subscriber<? super T> t) {
-            PublishSubscriber<T> ps = new PublishSubscriber<>(t, this);
-            t.onSubscribe(ps);
-            if (ps.cancelled == 0) {
-                if (add(ps)) {
-                    // if cancellation happened while a successful add, the remove() didn't work
-                    // so we need to do it again
-                    if (ps.cancelled != 0) {
-                        remove(ps);
-                    }
-                } else {
-                    Object o = get();
-                    if (o == COMPLETE) {
-                        ps.onComplete();
-                    } else {
-                        ps.onError((Throwable)o);
-                    }
-                }
-            }
-        }
+		public void subscribe(Subscriber<? super T> t) {
+			PublishSubscriber<T> ps = new PublishSubscriber<>(t, this);
+			if (ps.cancelled == 0) {
+				boolean added = add(ps);
+				if (added) {
+					// if cancellation happened while a successful add, the remove() didn't work
+					// so we need to do it again
+					if (ps.cancelled != 0) {
+						remove(ps);
+					}
+				}
+				t.onSubscribe(ps);
+				if (!added) {
+					Object o = get();
+					if (o == COMPLETE) {
+						ps.onComplete();
+					} else {
+						ps.onError((Throwable) o);
+					}
+					return;
+				}
+			}
+		}
 
         /**
          * @return the array of currently subscribed subscribers

src/test/java/io/reactivex/subjects/PublishSubjectTest.java

@@ -18,7 +18,10 @@
 import static org.mockito.Mockito.*;
 
 import java.util.ArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.*;
 
 import org.junit.Test;
@@ -426,4 +429,49 @@ public void testCurrentStateMethodsError() {
         assertFalse(as.hasComplete());
         assertTrue(as.getThrowable() instanceof TestException);
     }
+    
+    @Test
+    public void testRegisteredBeforeOnSubscribe() throws InterruptedException {
+    	PublishSubject<Object> as = PublishSubject.create();
+    	
+    	Runnable sideEffect = () -> {
+    		as.onNext(1);
+    		as.onComplete();
+    	};
+    	
+    	final AtomicReference<Object> next = new AtomicReference<Object>();
+    	final CountDownLatch latch = new CountDownLatch(1);
+    	
+    	as.subscribe(new Subscriber<Object>() {
+
+			@Override
+			public void onSubscribe(Subscription s) {
+				s.request(Long.MAX_VALUE);
+				// we are subscribed, so expect that we can now trigger side-effects that cause data to flow
+				sideEffect.run();
+			}
+
+			@Override
+			public void onNext(Object t) {
+				next.set(t);
+			}
+
+			@Override
+			public void onError(Throwable t) {
+				latch.countDown();
+			}
+
+			@Override
+			public void onComplete() {
+				latch.countDown();
+			}
+    		
+    	});
+    	
+    	if(!latch.await(500, TimeUnit.MILLISECONDS)) {
+    		fail("Did not receive events");
+    	} 
+    	
+    	assertEquals(1, next.get());
+    }
 }