ZipObservable fix

JAVA-3687

Author: rozza

driver-scala/src/main/scala/org/mongodb/scala/internal/ZipObservable.scala

@@ -17,36 +17,40 @@
 package org.mongodb.scala.internal
 
 import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.atomic.AtomicLong
 
 import org.mongodb.scala.{ Observable, Observer, Subscription }
 
-private[scala] case class ZipObservable[T, U](
-    observable1: Observable[T],
-    observable2: Observable[U]
-) extends Observable[(T, U)] {
+private[scala] case class ZipObservable[L, R](
+    leftObservable: Observable[L],
+    rightObservable: Observable[R]
+) extends Observable[(L, R)] {
 
-  def subscribe(observer: Observer[_ >: (T, U)]): Unit = {
+  def subscribe(observer: Observer[_ >: (L, R)]): Unit = {
     val helper = SubscriptionHelper(observer)
-    observable1.subscribe(SubscriptionCheckingObserver(helper.createFirstObserver))
-    observable2.subscribe(SubscriptionCheckingObserver(helper.createSecondObserver))
+    leftObservable.subscribe(SubscriptionCheckingObserver(helper.createLeftObserver))
+    rightObservable.subscribe(SubscriptionCheckingObserver(helper.createRightObserver))
   }
 
-  case class SubscriptionHelper(observer: Observer[_ >: (T, U)]) {
-    private val thisQueue: ConcurrentLinkedQueue[(Long, T)] = new ConcurrentLinkedQueue[(Long, T)]()
-    private val thatQueue: ConcurrentLinkedQueue[(Long, U)] = new ConcurrentLinkedQueue[(Long, U)]()
+  case class SubscriptionHelper(observer: Observer[_ >: (L, R)]) {
+    private val leftQueue: ConcurrentLinkedQueue[(Long, L)] = new ConcurrentLinkedQueue[(Long, L)]()
+    private val rightQueue: ConcurrentLinkedQueue[(Long, R)] = new ConcurrentLinkedQueue[(Long, R)]()
 
+    private val leftCounter: AtomicLong = new AtomicLong()
+    private val rightCounter: AtomicLong = new AtomicLong()
+    @volatile private var completedLeft: Boolean = false
+    @volatile private var completedRight: Boolean = false
     @volatile private var terminated: Boolean = false
-    @volatile private var observable1Subscription: Option[Subscription] = None
-    @volatile private var observable2Subscription: Option[Subscription] = None
+    @volatile private var leftSubscription: Option[Subscription] = None
+    @volatile private var rightSubscription: Option[Subscription] = None
 
-    def createFirstObserver: Observer[T] = createSubObserver[T](thisQueue, observer, firstSub = true)
-
-    def createSecondObserver: Observer[U] = createSubObserver[U](thatQueue, observer, firstSub = false)
+    def createLeftObserver: Observer[L] = createSubObserver[L](leftQueue, observer, isLeftSub = true)
+    def createRightObserver: Observer[R] = createSubObserver[R](rightQueue, observer, isLeftSub = false)
 
     private def createSubObserver[A](
         queue: ConcurrentLinkedQueue[(Long, A)],
-        observer: Observer[_ >: (T, U)],
-        firstSub: Boolean
+        observer: Observer[_ >: (L, R)],
+        isLeftSub: Boolean
     ): Observer[A] = {
       new Observer[A] {
         @volatile private var counter: Long = 0
@@ -56,38 +60,61 @@ private[scala] case class ZipObservable[T, U](
         }
 
         override def onSubscribe(subscription: Subscription): Unit = {
-          if (firstSub) {
-            observable1Subscription = Some(subscription)
+          if (isLeftSub) {
+            leftSubscription = Some(subscription)
           } else {
-            observable2Subscription = Some(subscription)
+            rightSubscription = Some(subscription)
           }
 
-          if (observable1Subscription.nonEmpty && observable2Subscription.nonEmpty) {
+          if (leftSubscription.nonEmpty && rightSubscription.nonEmpty) {
             observer.onSubscribe(jointSubscription)
           }
         }
 
         override def onComplete(): Unit = {
-          if (!firstSub) {
-            terminated = true
-            observer.onComplete()
-          }
+          markCompleted(isLeftSub)
+          processNext(observer)
         }
 
         override def onNext(tResult: A): Unit = {
+          if (isLeftSub) leftCounter.incrementAndGet() else rightCounter.incrementAndGet()
           counter += 1
           queue.add((counter, tResult))
-          if (!firstSub) processNext(observer)
+          processNext(observer)
         }
       }
     }
 
-    private def processNext(observer: Observer[_ >: (T, U)]): Unit = {
-      (thisQueue.peek, thatQueue.peek) match {
-        case ((k1: Long, _), (k2: Long, _)) if k1 == k2 => observer.onNext((thisQueue.poll()._2, thatQueue.poll()._2))
+    private def markCompleted(isLeftSub: Boolean): Unit = synchronized {
+      if (isLeftSub) {
+        completedLeft = true
+      } else {
+        completedRight = true
+      }
+    }
+
+    private def completed(): Unit = synchronized {
+      if (!terminated) {
+        terminated = true
+        leftSubscription.foreach(_.unsubscribe())
+        rightSubscription.foreach(_.unsubscribe())
+        observer.onComplete()
+      }
+    }
+
+    private def processNext(observer: Observer[_ >: (L, R)]): Unit = synchronized {
+      (leftQueue.peek, rightQueue.peek) match {
+        case ((k1: Long, _), (k2: Long, _)) if k1 == k2 =>
+          observer.onNext((leftQueue.poll()._2, rightQueue.poll()._2))
+          processNext(observer)
         case _ =>
-          if (!terminated && !jointSubscription.isUnsubscribed) jointSubscription.request(1) // Uneven queues request more data
-        // from downstream so to honor the original request for data.
+          if (!terminated && !jointSubscription.isUnsubscribed) {
+            if (completedLeft && rightCounter.get() >= leftCounter.get()) {
+              completed()
+            } else if (completedRight && leftCounter.get() >= rightCounter.get()) {
+              completed()
+            }
+          }
       }
     }
 
@@ -96,14 +123,14 @@ private[scala] case class ZipObservable[T, U](
       override def isUnsubscribed: Boolean = !subscribed
 
       override def request(n: Long): Unit = {
-        observable1Subscription.foreach(_.request(n))
-        observable2Subscription.foreach(_.request(n))
+        leftSubscription.foreach(_.request(n))
+        rightSubscription.foreach(_.request(n))
       }
 
       override def unsubscribe(): Unit = {
         subscribed = false
-        observable1Subscription.foreach(_.unsubscribe())
-        observable2Subscription.foreach(_.unsubscribe())
+        leftSubscription.foreach(_.unsubscribe())
+        rightSubscription.foreach(_.unsubscribe())
       }
     }
   }

driver-scala/src/test/scala/org/mongodb/scala/internal/ObservableImplementationSpec.scala

@@ -167,6 +167,18 @@ class ObservableImplementationSpec extends BaseSpec with TableDrivenPropertyChec
         observer.completed should equal(true)
       }
     }
+
+    forAll(zippedObservablesWithEmptyObservable) { (observable: Observable[(Int, Int)]) =>
+      {
+        val observer = TestObserver[(Int, Int)]()
+        observable.subscribe(observer)
+
+        observer.subscription.foreach(_.request(100))
+
+        observer.results should equal(List())
+        observer.completed should equal(true)
+      }
+    }
   }
 
   it should "error if requested amount is less than 1" in {
@@ -258,6 +270,13 @@ class ObservableImplementationSpec extends BaseSpec with TableDrivenPropertyChec
       ZipObservable[Int, Int](TestObservable[Int](), TestObservable[Int](1 to 50))
     )
 
+  private def zippedObservablesWithEmptyObservable =
+    Table[Observable[(Int, Int)]](
+      "observable",
+      ZipObservable[Int, Int](TestObservable[Int](1 to 50), TestObservable[Int](List())),
+      ZipObservable[Int, Int](TestObservable[Int](List()), TestObservable[Int](1 to 50))
+    )
+
   private def overRequestingObservables =
     Table(
       ("observable", "observer", "expected"),