SAMZA-1471: SystemConsumers should not poll ssp that hit end of stream

When SystemConsumers poll from SSPs that have hit end of stream, obviously there will be no data return and the poll will exhaust the timeout. This would cause performance issue.

Author: Hai Lu <[email protected]>

Reviewers: Xinyu Liu <[email protected]>

Closes apache#342 from lhaiesp/master

Author: lhaiesp

samza-core/src/main/scala/org/apache/samza/system/SystemConsumers.scala

@@ -28,6 +28,7 @@ import org.apache.samza.util.{Logging, TimerUtils}
 import org.apache.samza.system.chooser.MessageChooser
 import org.apache.samza.SamzaException
 import java.util.ArrayDeque
+import java.util.Collections
 import java.util.HashSet
 import java.util.HashMap
 import java.util.Queue
@@ -186,7 +187,7 @@ class SystemConsumers (
     if (IncomingMessageEnvelope.END_OF_STREAM_OFFSET.equals(offset)) {
       info("Stream : %s is already at end of stream" format (systemStreamPartition))
       endOfStreamSSPs.add(systemStreamPartition)
-      return;
+      return
     }
 
     metrics.registerSystemStreamPartition(systemStreamPartition)
@@ -258,7 +259,14 @@ class SystemConsumers (
 
     trace("Getting fetch map for system: %s" format systemName)
 
-    val systemFetchSet = emptySystemStreamPartitionsBySystem.get(systemName)
+    val systemFetchSet : util.Set[SystemStreamPartition] =
+      if (emptySystemStreamPartitionsBySystem.containsKey(systemName)) {
+        val sspToFetch = new util.HashSet(emptySystemStreamPartitionsBySystem.get(systemName))
+        sspToFetch.removeAll(endOfStreamSSPs)
+        sspToFetch
+      } else {
+        Collections.emptySet()
+      }
 
     // Poll when at least one SSP in this system needs more messages.
 
@@ -293,7 +301,7 @@ class SystemConsumers (
         }
       }
     } else {
-      trace("Skipping polling for %s. Already have messages available for all registered SystemStreamPartitions." format (systemName))
+      trace("Skipping polling for %s. Already have messages available for all registered SystemStreamPartitions." format systemName)
     }
   }
 

samza-core/src/test/scala/org/apache/samza/system/TestSystemConsumers.scala

@@ -19,6 +19,9 @@
 
 package org.apache.samza.system
 
+import java.util
+import java.util.Collections
+
 import org.junit.Assert._
 import org.junit.Test
 import org.apache.samza.Partition
@@ -285,10 +288,62 @@ class TestSystemConsumers {
 
   }
 
+  @Test
+  def testSystemConsumersShouldNotPollEndOfStreamSSPs {
+    val system = "test-system"
+    val stream = "some-stream"
+    val systemStreamPartition1 = new SystemStreamPartition(system, stream, new Partition(1))
+    val systemStreamPartition2 = new SystemStreamPartition(system, stream, new Partition(2))
+    val normalEnvelope = new IncomingMessageEnvelope(systemStreamPartition1, "1", "k", "v")
+    val endOfStreamEnvelope = IncomingMessageEnvelope.buildEndOfStreamEnvelope(systemStreamPartition2)
+    val consumer = new CustomPollResponseSystemConsumer(normalEnvelope)
+    val consumers = new SystemConsumers(new MockMessageChooser, Map(system -> consumer),
+      new SerdeManager, new SystemConsumersMetrics,
+      SystemConsumers.DEFAULT_NO_NEW_MESSAGES_TIMEOUT,
+      SystemConsumers.DEFAULT_DROP_SERIALIZATION_ERROR,
+      SystemConsumers.DEFAULT_POLL_INTERVAL_MS, clock = () => 0)
+
+    consumers.register(systemStreamPartition1, "0")
+    consumers.register(systemStreamPartition2, "0")
+    consumers.start
+
+    // Start should trigger a poll to the consumer.
+    assertEquals(1, consumer.polls)
+    assertEquals(2, consumer.lastPoll.size())
+
+    // Tell the consumer to start returning messages when polled.
+    val nextResponse = Map[SystemStreamPartition, java.util.List[IncomingMessageEnvelope]](
+      systemStreamPartition1 -> Collections.singletonList(normalEnvelope),
+      systemStreamPartition2 -> Collections.singletonList(endOfStreamEnvelope)
+    )
+    consumer.setNextResponse(nextResponse)
+
+    // Choose to trigger a refresh with data.
+    assertNull(consumers.choose())
+
+    // Choose should have triggered a second poll, since no messages are available.
+    assertEquals(2, consumer.polls)
+    assertEquals(2, consumer.lastPoll.size())
+
+    // Choose a few times and let chooser handle the end of stream message
+    assertNotNull(consumers.choose())
+    assertNotNull(consumers.choose())
+    consumers.tryUpdate(systemStreamPartition1)
+    // Now assuming that chooser has processed end of stream message,
+    // tryUpdate shouldn't add ssp back to emptySystemStreamPartitionsBySystem
+    consumers.tryUpdate(systemStreamPartition2)
+    assertNull(consumers.choose())
+    assertEquals(3, consumer.polls)
+    // SystemConsumers should poll only one partition: ssp1
+    assertEquals(1, consumer.lastPoll.size())
+    assertTrue(consumer.lastPoll.contains(systemStreamPartition1))
+  }
+
   /**
    * A simple MockSystemConsumer that keeps track of what was polled, and lets
    * you define how many envelopes to return in the poll response. You can
    * supply the envelope to use for poll responses through the constructor.
+   * You can also directly set the next response by calling setNextResponse
    */
   private class CustomPollResponseSystemConsumer(envelope: IncomingMessageEnvelope) extends SystemConsumer {
     var polls = 0
@@ -299,7 +354,7 @@ class TestSystemConsumers {
     def register(systemStreamPartition: SystemStreamPartition, offset: String) {}
     def poll(systemStreamPartitions: java.util.Set[SystemStreamPartition], timeout: Long) = {
       polls += 1
-      lastPoll = systemStreamPartitions
+      lastPoll = new util.HashSet[SystemStreamPartition](systemStreamPartitions)
       pollResponse.asJava
     }
     def setResponseSizes(numEnvelopes: Int) {
@@ -308,6 +363,9 @@ class TestSystemConsumers {
       pollResponse = Map(envelope.getSystemStreamPartition -> q)
       pollResponse = Map[SystemStreamPartition, java.util.List[IncomingMessageEnvelope]]()
     }
+    def setNextResponse(nextResponse: Map[SystemStreamPartition, java.util.List[IncomingMessageEnvelope]]) {
+      pollResponse = nextResponse
+    }
   }
 
   /**