Simplify Subscription/Observer State

Author: benjchristensen

rxjava-core/src/main/java/rx/subjects/SubjectSubscriptionManager.java

@@ -23,10 +23,8 @@
 import rx.Observable.OnSubscribe;
 import rx.Observer;
 import rx.Subscriber;
-import rx.Subscription;
 import rx.functions.Action0;
 import rx.functions.Action1;
-import rx.operators.SafeObservableSubscription;
 import rx.subscriptions.Subscriptions;
 
 /* package */class SubjectSubscriptionManager<T> {
@@ -39,10 +37,11 @@
      *            Always runs at the beginning of 'subscribe' regardless of terminal state.
      * @param onTerminated
      *            Only runs if Subject is in terminal state and the Observer ends up not being registered.
-     * @param onUnsubscribe called after the child subscription is removed from the state
+     * @param onUnsubscribe
+     *            called after the child subscription is removed from the state
      * @return
      */
-    public OnSubscribe<T> getOnSubscribeFunc(final Action1<SubjectObserver<? super T>> onSubscribe, 
+    public OnSubscribe<T> getOnSubscribeFunc(final Action1<SubjectObserver<? super T>> onSubscribe,
             final Action1<SubjectObserver<? super T>> onTerminated,
             final Action1<SubjectObserver<? super T>> onUnsubscribe) {
         return new OnSubscribe<T>() {
@@ -73,10 +72,9 @@ public void call(Subscriber<? super T> actualOperator) {
                         }
                         break;
                     } else {
-                        final SafeObservableSubscription subscription = new SafeObservableSubscription();
-                        actualOperator.add(subscription); // add to parent if the Subject itself is unsubscribed
                         addedObserver = true;
-                        subscription.wrap(Subscriptions.create(new Action0() {
+                        // add to parent if the Subject itself is unsubscribed
+                        actualOperator.add(Subscriptions.create(new Action0() {
 
                             @Override
                             public void call() {
@@ -85,19 +83,19 @@ public void call() {
                                 do {
                                     current = state.get();
                                     // on unsubscribe remove it from the map of outbound observers to notify
-                                    newState = current.removeObserver(subscription);
+                                    newState = current.removeObserver(observer);
                                 } while (!state.compareAndSet(current, newState));
                                 if (onUnsubscribe != null) {
                                     onUnsubscribe.call(observer);
                                 }
                             }
                         }));
-                        if (subscription.isUnsubscribed()) {
-                            addedObserver = false;
+                        if (actualOperator.isUnsubscribed()) {
+                            // we've been unsubscribed while working so return and do nothing
                             return;
                         }
                         // on subscribe add it to the map of outbound observers to notify
-                        newState = current.addObserver(subscription, observer);
+                        newState = current.addObserver(observer);
                     }
                 } while (!state.compareAndSet(current, newState));
 
@@ -126,7 +124,7 @@ protected Collection<SubjectObserver<? super T>> terminate(Action1<Collection<Su
             }
         } while (!state.compareAndSet(current, newState));
 
-        Collection<SubjectObserver<? super T>> observerCollection = (Collection)Arrays.asList(newState.observers);
+        Collection<SubjectObserver<? super T>> observerCollection = (Collection) Arrays.asList(newState.observers);
         /*
          * if we get here then we won setting the state to terminated
          * and have a deterministic set of Observers to emit to (concurrent subscribes
@@ -158,97 +156,79 @@ public SubjectObserver<Object>[] rawSnapshot() {
     protected static class State<T> {
         final boolean terminated;
         final CountDownLatch terminationLatch;
-        final Subscription[] subscriptions;
         final SubjectObserver[] observers;
         // to avoid lots of empty arrays
-        final Subscription[] EMPTY_S = new Subscription[0];
-        // to avoid lots of empty arrays
         final SubjectObserver[] EMPTY_O = new SubjectObserver[0];
 
-        private State(boolean isTerminated, CountDownLatch terminationLatch,
-                Subscription[] subscriptions, SubjectObserver[] observers) {
+        private State(boolean isTerminated, CountDownLatch terminationLatch, SubjectObserver[] observers) {
             this.terminationLatch = terminationLatch;
             this.terminated = isTerminated;
-            this.subscriptions = subscriptions;
             this.observers = observers;
         }
 
         State() {
             this.terminated = false;
             this.terminationLatch = null;
-            this.subscriptions = EMPTY_S;
             this.observers = EMPTY_O;
         }
 
         public State<T> terminate() {
             if (terminated) {
                 throw new IllegalStateException("Already terminated.");
             }
-            return new State<T>(true, new CountDownLatch(1), subscriptions, observers);
+            return new State<T>(true, new CountDownLatch(1), observers);
         }
 
-        public State<T> addObserver(Subscription s, SubjectObserver<? super T> observer) {
+        public State<T> addObserver(SubjectObserver<? super T> observer) {
             int n = this.observers.length;
 
-            Subscription[] newsubscriptions = Arrays.copyOf(this.subscriptions, n + 1);
             SubjectObserver[] newobservers = Arrays.copyOf(this.observers, n + 1);
 
-            newsubscriptions[n] = s;
             newobservers[n] = observer;
 
-            return createNewWith(newsubscriptions, newobservers);
+            return createNewWith(newobservers);
         }
 
-        private State<T> createNewWith(Subscription[] newsubscriptions, SubjectObserver[] newobservers) {
-            return new State<T>(terminated, terminationLatch, newsubscriptions, newobservers);
+        private State<T> createNewWith(SubjectObserver[] newobservers) {
+            return new State<T>(terminated, terminationLatch, newobservers);
         }
 
-        public State<T> removeObserver(Subscription s) {
+        public State<T> removeObserver(SubjectObserver<? super T> o) {
             // we are empty, nothing to remove
             if (this.observers.length == 0) {
                 return this;
-            } else
-            if (this.observers.length == 1) {
-                if (this.subscriptions[0].equals(s)) {
-                    return createNewWith(EMPTY_S, EMPTY_O);
-                }
-                return this;
             }
+
             int n = this.observers.length - 1;
             int copied = 0;
-            Subscription[] newsubscriptions = new Subscription[n];
             SubjectObserver[] newobservers = new SubjectObserver[n];
 
-            for (int i = 0; i < this.subscriptions.length; i++) {
-                Subscription s0 = this.subscriptions[i];
-                if (!s0.equals(s)) {
+            for (int i = 0; i < this.observers.length; i++) {
+                SubjectObserver s0 = this.observers[i];
+                if (!s0.equals(o)) {
                     if (copied == n) {
                         // if s was not found till the end of the iteration
                         // we return ourselves since no modification should
                         // have happened
                         return this;
                     }
-                    newsubscriptions[copied] = s0;
-                    newobservers[copied] = this.observers[i];
+                    newobservers[copied] = s0;
                     copied++;
                 }
             }
 
             if (copied == 0) {
-                return createNewWith(EMPTY_S, EMPTY_O);
+                return createNewWith(EMPTY_O);
             }
             // if somehow copied less than expected, truncate the arrays
             // if s is unique, this should never happen
             if (copied < n) {
-                Subscription[] newsubscriptions2 = new Subscription[copied];
-                System.arraycopy(newsubscriptions, 0, newsubscriptions2, 0, copied);
-                
                 SubjectObserver[] newobservers2 = new SubjectObserver[copied];
                 System.arraycopy(newobservers, 0, newobservers2, 0, copied);
 
-                return createNewWith(newsubscriptions2, newobservers2);
+                return createNewWith(newobservers2);
             }
-            return createNewWith(newsubscriptions, newobservers);
+            return createNewWith(newobservers);
         }
     }