1.x: fix unsubscription and producer issues in sample(other)

Author: akarnokd

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

@@ -16,9 +16,9 @@
 package rx.internal.operators;
 
 import java.util.concurrent.atomic.AtomicReference;
-import rx.Observable;
+
+import rx.*;
 import rx.Observable.Operator;
-import rx.Subscriber;
 import rx.observers.SerializedSubscriber;
 
 /**
@@ -44,7 +44,9 @@ public Subscriber<? super T> call(Subscriber<? super T> child) {
 
         final AtomicReference<Object> value = new AtomicReference<Object>(EMPTY_TOKEN);
 
-        Subscriber<U> samplerSub = new Subscriber<U>(child) {
+        final AtomicReference<Subscription> main = new AtomicReference<Subscription>();
+        
+        final Subscriber<U> samplerSub = new Subscriber<U>() {
             @Override
             public void onNext(U t) {
                 Object localValue = value.getAndSet(EMPTY_TOKEN);
@@ -57,16 +59,18 @@ public void onNext(U t) {
 
             @Override
             public void onError(Throwable e) {
+                // no need to null check, main is assigned before any of the two gets subscribed
+                main.get().unsubscribe();
                 s.onError(e);
-                unsubscribe();
             }
 
             @Override
             public void onCompleted() {
+                // no need to null check, main is assigned before any of the two gets subscribed
+                main.get().unsubscribe();
+                // onNext(null); // emit the very last value?
                 s.onCompleted();
-                unsubscribe();
             }
-            
         };
 
         Subscriber<T> result = new Subscriber<T>() {
@@ -77,18 +81,22 @@ public void onNext(T t) {
 
             @Override
             public void onError(Throwable e) {
+                samplerSub.unsubscribe();
                 s.onError(e);
-                unsubscribe();
             }
 
             @Override
             public void onCompleted() {
+                samplerSub.unsubscribe();
+                // samplerSub.onNext(null); // emit the very last value?
                 s.onCompleted();
-                unsubscribe();
             }
         };
 
+        main.lazySet(result);
+        
         child.add(result);
+        child.add(samplerSub);
 
         sampler.unsafeSubscribe(samplerSub);
 

src/test/java/rx/internal/operators/OperatorSampleTest.java

@@ -19,13 +19,15 @@
 import static org.mockito.Mockito.*;
 
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.*;
 import org.mockito.InOrder;
 
 import rx.*;
 import rx.Observable.OnSubscribe;
 import rx.functions.*;
+import rx.observers.TestSubscriber;
 import rx.schedulers.TestScheduler;
 import rx.subjects.PublishSubject;
 
@@ -312,4 +314,81 @@ public void call(Long t) {
 
         Assert.assertEquals(Long.MAX_VALUE, requested[0]);
     }
+    
+    @Test
+    public void dontUnsubscribeChild1() {
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        
+        PublishSubject<Integer> source = PublishSubject.create();
+        
+        PublishSubject<Integer> sampler = PublishSubject.create();
+        
+        source.sample(sampler).unsafeSubscribe(ts);
+        
+        source.onCompleted();
+        
+        Assert.assertFalse("Source has subscribers?", source.hasObservers());
+        Assert.assertFalse("Sampler has subscribers?", sampler.hasObservers());
+        
+        Assert.assertFalse("TS unsubscribed?", ts.isUnsubscribed());
+    }
+
+    @Test
+    public void dontUnsubscribeChild2() {
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        
+        PublishSubject<Integer> source = PublishSubject.create();
+        
+        PublishSubject<Integer> sampler = PublishSubject.create();
+        
+        source.sample(sampler).unsafeSubscribe(ts);
+        
+        sampler.onCompleted();
+        
+        Assert.assertFalse("Source has subscribers?", source.hasObservers());
+        Assert.assertFalse("Sampler has subscribers?", sampler.hasObservers());
+        
+        Assert.assertFalse("TS unsubscribed?", ts.isUnsubscribed());
+    }
+    
+    @Test
+    public void neverSetProducer() {
+        Observable<Integer> neverBackpressure = Observable.create(new OnSubscribe<Integer>() {
+            @Override
+            public void call(Subscriber<? super Integer> t) {
+                t.setProducer(new Producer() {
+                    @Override
+                    public void request(long n) {
+                        // irrelevant in this test
+                    }
+                });
+            }
+        });
+        
+        final AtomicInteger count = new AtomicInteger();
+        
+        neverBackpressure.sample(neverBackpressure).unsafeSubscribe(new Subscriber<Integer>() {
+            @Override
+            public void onNext(Integer t) {
+                // irrelevant
+            }
+            
+            @Override
+            public void onError(Throwable e) {
+                // irrelevant
+            }
+            
+            @Override
+            public void onCompleted() {
+                // irrelevant
+            }
+            
+            @Override
+            public void setProducer(Producer p) {
+                count.incrementAndGet();
+            }
+        });
+        
+        Assert.assertEquals(0, count.get());
+    }
 }