Refactor MultipleAssignment

- simplified state machine

Author: benjchristensen

rxjava-core/src/main/java/rx/subscriptions/MultipleAssignmentSubscription.java

@@ -25,43 +25,77 @@
  * 
  * @see <a href="http://msdn.microsoft.com/en-us/library/system.reactive.disposables.multipleassignmentdisposable">Rx.Net equivalent MultipleAssignmentDisposable</a>
  */
-public class MultipleAssignmentSubscription implements Subscription {
-    private AtomicReference<Subscription> reference = new AtomicReference<Subscription>();
-    /** Sentinel for the unsubscribed state. */
-    private static final Subscription UNSUBSCRIBED_SENTINEL = new Subscription() {
-        @Override
-        public void unsubscribe() {
+public final class MultipleAssignmentSubscription implements Subscription {
+
+    private final AtomicReference<State> state = new AtomicReference<State>(new State(false, Subscriptions.empty()));
+
+    private static final class State {
+        final boolean isUnsubscribed;
+        final Subscription subscription;
+
+        State(boolean u, Subscription s) {
+            this.isUnsubscribed = u;
+            this.subscription = s;
+        }
+
+        State unsubscribe() {
+            return new State(true, subscription);
         }
-    };
+
+        State set(Subscription s) {
+            return new State(isUnsubscribed, s);
+        }
+
+    }
 
     public boolean isUnsubscribed() {
-        return reference.get() == UNSUBSCRIBED_SENTINEL;
+        return state.get().isUnsubscribed;
     }
 
     @Override
     public void unsubscribe() {
-        Subscription s = reference.getAndSet(UNSUBSCRIBED_SENTINEL);
-        if (s != null) {
-            s.unsubscribe();
-        }
+        State oldState;
+        State newState;
+        do {
+            oldState = state.get();
+            if (oldState.isUnsubscribed) {
+                return;
+            } else {
+                newState = oldState.unsubscribe();
+            }
+        } while (!state.compareAndSet(oldState, newState));
+        oldState.subscription.unsubscribe();
     }
 
+    @Deprecated
     public void setSubscription(Subscription s) {
+        set(s);
+    }
+
+    public void set(Subscription s) {
+        if (s == null) {
+            throw new IllegalArgumentException("Subscription can not be null");
+        }
+        State oldState;
+        State newState;
         do {
-            Subscription r = reference.get();
-            if (r == UNSUBSCRIBED_SENTINEL) {
+            oldState = state.get();
+            if (oldState.isUnsubscribed) {
                 s.unsubscribe();
                 return;
+            } else {
+                newState = oldState.set(s);
             }
-            if (reference.compareAndSet(r, s)) {
-                break;
-            }
-        } while (true);
+        } while (!state.compareAndSet(oldState, newState));
     }
 
+    @Deprecated
     public Subscription getSubscription() {
-        Subscription s = reference.get();
-        return s != UNSUBSCRIBED_SENTINEL ? s : Subscriptions.empty();
+        return get();
+    }
+
+    public Subscription get() {
+        return state.get().subscription;
     }
 
 }