Make unsubscribe + reference modification an atomic operation

Author: jloisel

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

@@ -19,6 +19,9 @@
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import rx.Subscription;
 
@@ -29,25 +32,43 @@
  * @see <a href="http://msdn.microsoft.com/en-us/library/system.reactive.disposables.serialdisposable(v=vs.103).aspx">Rx.Net equivalent SerialDisposable</a>
  */
 public class SerialSubscription implements Subscription {
-	private final AtomicBoolean unsubscribed = new AtomicBoolean();
-	private final AtomicReference<Subscription> reference = new AtomicReference<Subscription>(empty());
-
-	@Override
-	public void unsubscribe() {
-		if (unsubscribed.compareAndSet(false, true)) {
-			reference.getAndSet(empty()).unsubscribe();
-		}
-	}
-
-	public void setSubscription(final Subscription subscription) {
-		if (unsubscribed.get()) {
-			subscription.unsubscribe();
-		} else {
-			reference.getAndSet(subscription == null ? empty() : subscription).unsubscribe();
-		}
-	}
-
-	public Subscription getSubscription() {
-		return reference.get();
-	}
+    private final AtomicBoolean unsubscribed = new AtomicBoolean();
+    private final AtomicReference<Subscription> reference = new AtomicReference<Subscription>(empty());
+    private final Lock read;
+    private final Lock write;
+
+    public SerialSubscription() {
+        final ReadWriteLock lock = new ReentrantReadWriteLock();
+        read = lock.readLock();
+        write = lock.writeLock();
+    }
+
+    @Override
+    public void unsubscribe() {
+        write.lock();
+        try {
+            if (unsubscribed.compareAndSet(false, true)) {
+                reference.getAndSet(empty()).unsubscribe();
+            }
+        } finally {
+            write.unlock();
+        }
+    }
+
+    public void setSubscription(final Subscription subscription) {
+        read.lock();
+        try {
+            if (unsubscribed.get()) {
+                subscription.unsubscribe();
+            } else {
+                reference.getAndSet(subscription).unsubscribe();
+            }
+        } finally {
+            read.unlock();
+        }
+    }
+
+    public Subscription getSubscription() {
+        return reference.get();
+    }
 }