Merge pull request #966: [direct-io-kafka] idle watermark robustness

Author: je-ik

direct/io-kafka/src/main/java/cz/o2/proxima/direct/io/kafka/KafkaLogReader.java

@@ -402,6 +402,7 @@ private void submitConsumerWithObserver(
               new KafkaThroughputLimiter(maxBytesPerSec);
           final Map<TopicPartition, Long> polledOffsets = new HashMap<>();
           final Map<TopicPartition, Integer> emptyPolls = new HashMap<>();
+          final Map<TopicPartition, Long> endOffsets = new HashMap<>();
 
           handle.set(
               createObserveHandle(shutdown, seekOffsets, consumer, readyLatch, completedLatch));
@@ -414,6 +415,7 @@ private void submitConsumerWithObserver(
                   topicPartitionToId,
                   emptyPolls,
                   polledOffsets,
+                  endOffsets,
                   watermarkEstimator);
 
           try (KafkaConsumer<Object, Object> kafka =
@@ -423,12 +425,9 @@ private void submitConsumerWithObserver(
 
             // we need to poll first to initialize kafka assignments and rebalance listener
             ConsumerRecords<Object, Object> poll;
-            Map<TopicPartition, Long> endOffsets;
 
             do {
               poll = kafka.poll(pollDuration);
-              endOffsets = stopAtCurrent ? findNonEmptyEndOffsets(kafka) : null;
-
               if (log.isDebugEnabled()) {
                 log.debug(
                     "End offsets of current assignment {}: {}", kafka.assignment(), endOffsets);
@@ -498,7 +497,12 @@ private void submitConsumerWithObserver(
               terminateIfConsumed(stopAtCurrent, kafka, endOffsets, polledOffsets, completed);
 
               progressWatermarkOnEmptyPartitions(
-                  consumer, emptyPolls, topicPartitionToId, watermarkEstimator);
+                  consumer,
+                  emptyPolls,
+                  endOffsets,
+                  polledOffsets,
+                  topicPartitionToId,
+                  watermarkEstimator);
               throughputLimiter.sleepToLimitThroughput(bytesPolled, pollTimeMs);
               long startTime = System.currentTimeMillis();
               poll = kafka.poll(pollDuration);
@@ -548,13 +552,19 @@ private void submitConsumerWithObserver(
   private void progressWatermarkOnEmptyPartitions(
       ElementConsumer<Object, Object> consumer,
       Map<TopicPartition, Integer> emptyPolls,
+      Map<TopicPartition, Long> endOffsets,
+      Map<TopicPartition, Long> polledOffsets,
       Map<TopicPartition, Integer> topicPartitionToId,
       AtomicReference<PartitionedWatermarkEstimator> watermarkEstimator) {
 
     int partitions = emptyPolls.size();
     List<TopicPartition> idlingPartitions =
         emptyPolls.entrySet().stream()
             .filter(e -> e.getValue() >= partitions)
+            .filter(
+                e ->
+                    MoreObjects.firstNonNull(polledOffsets.get(e.getKey()), -1L) + 1
+                        >= MoreObjects.firstNonNull(endOffsets.get(e.getKey()), 0L))
             .map(Entry::getKey)
             .collect(Collectors.toList());
 
@@ -624,7 +634,7 @@ private void logConsumerWatermark(
 
     if (log.isDebugEnabled()) {
       log.debug(
-          "Current watermark of consumer name {} with offsets {} " + "on {} poll'd records is {}",
+          "Current watermark of consumer name {} with offsets {} on {} poll'd records is {}",
           name,
           offsets,
           polledCount,
@@ -723,14 +733,6 @@ public void waitUntilReady() throws InterruptedException {
     };
   }
 
-  private Map<TopicPartition, Long> findNonEmptyEndOffsets(final KafkaConsumer<?, ?> kafka) {
-    Set<TopicPartition> assignment = kafka.assignment();
-    Map<TopicPartition, Long> beginning = kafka.beginningOffsets(assignment);
-    return kafka.endOffsets(assignment).entrySet().stream()
-        .filter(entry -> beginning.get(entry.getKey()) < entry.getValue())
-        .collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
-  }
-
   private KafkaConsumer<Object, Object> createConsumer() {
     return createConsumer(UUID.randomUUID().toString(), null, null, Position.NEWEST);
   }
@@ -883,6 +885,7 @@ private ConsumerRebalanceListener listener(
       Map<TopicPartition, Integer> topicPartitionToId,
       Map<TopicPartition, Integer> emptyPolls,
       Map<TopicPartition, Long> polledOffsets,
+      Map<TopicPartition, Long> endOffsets,
       AtomicReference<PartitionedWatermarkEstimator> watermarkEstimator) {
 
     return new ConsumerRebalanceListener() {
@@ -901,6 +904,9 @@ public void onPartitionsAssigned(Collection<TopicPartition> parts) {
         topicPartitionToId.clear();
         AtomicInteger id = new AtomicInteger();
         emptyPolls.clear();
+        polledOffsets.clear();
+        endOffsets.clear();
+        Optional.ofNullable(kafka.get()).ifPresent(k -> endOffsets.putAll(k.endOffsets(parts)));
 
         currentlyAssigned.forEach(p -> topicPartitionToId.put(p, id.getAndIncrement()));
 
@@ -921,13 +927,14 @@ public void onPartitionsAssigned(Collection<TopicPartition> parts) {
                       name != null
                           ? getCommittedTopicOffsets(currentlyAssigned, c)
                           : getCurrentTopicOffsets(currentlyAssigned, c);
-                  polledOffsets.clear();
-                  newOffsets.forEach(
-                      o ->
-                          polledOffsets.put(
-                              new TopicPartition(
-                                  o.getPartition().getTopic(), o.getPartition().getPartition()),
-                              o.getOffset() - 1));
+                  newOffsets.stream()
+                      .filter(o -> o.getOffset() > 0)
+                      .forEach(
+                          o ->
+                              polledOffsets.put(
+                                  new TopicPartition(
+                                      o.getPartition().getTopic(), o.getPartition().getPartition()),
+                                  o.getOffset() - 1));
                   consumer.onAssign(c, newOffsets);
                 });
       }

direct/io-kafka/src/test/java/cz/o2/proxima/direct/io/kafka/LocalKafkaCommitLogDescriptor.java

@@ -337,7 +337,7 @@ private Map<TopicPartition, Long> getEndOffsets(String name, Collection<TopicPar
     private <K, V> ConsumerRecords<K, V> pollConsumer(
         ConsumerGroup group,
         long period,
-        ConsumerId consumerId,
+        final ConsumerId consumerId,
         ElementSerializer<K, V> serializer,
         @Nullable ConsumerRebalanceListener listener)
         throws InterruptedException {
@@ -388,7 +388,6 @@ private <K, V> ConsumerRecords<K, V> pollConsumer(
               Optional.ofNullable(offsets.get(partition))
                   .orElse(getCommittedOffset(name, part.getId()));
           log.trace("Partition {} has last {}, reading from {}", partition, last, off);
-
           while (off < last && maxToPoll-- > 0) {
             if (off >= 0) {
               records.add(toConsumerRecord(partitionData.get(off), serializer, part.getId(), off));

direct/io-kafka/src/test/java/cz/o2/proxima/direct/io/kafka/LocalKafkaCommitLogDescriptorTest.java

@@ -79,6 +79,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
@@ -1124,15 +1125,16 @@ void testPollFromNConsumersMovesWatermark(
     final long now = System.currentTimeMillis();
     CountDownLatch latch = new CountDownLatch(numObservers);
     Map<CommitLogObserver, Long> observerWatermarks = new ConcurrentHashMap<>();
-    AtomicInteger readyObservers = new AtomicInteger();
+    Map<Integer, Integer> idles = new ConcurrentHashMap<>();
     for (int i = 0; i < numObservers; i++) {
+      int id = i;
       reader
           .observe(
               "test-" + expectMoved,
               Position.OLDEST,
               new CommitLogObserver() {
 
-                boolean confimedLatch = false;
+                boolean confirmedLatch = false;
 
                 @Override
                 public boolean onNext(StreamElement element, OnNextContext context) {
@@ -1147,28 +1149,24 @@ public boolean onError(Throwable error) {
 
                 @Override
                 public void onIdle(OnIdleContext context) {
-                  if (readyObservers.get() == numObservers) {
-                    Long oldWatermark = observerWatermarks.put(this, context.getWatermark());
-                    boolean hasWatermark =
-                        MoreObjects.firstNonNull(oldWatermark, Long.MIN_VALUE) > 0;
-                    if ((!expectMoved || !hasWatermark && context.getWatermark() > 0)
-                        && !confimedLatch) {
-
-                      confimedLatch = true;
+                  idles.compute(id, (k, v) -> MoreObjects.firstNonNull(v, 0) + 1);
+                  if (context.getWatermark() > 0) {
+                    observerWatermarks.put(this, context.getWatermark());
+                    if ((!expectMoved || context.getWatermark() > 0) && !confirmedLatch) {
+                      confirmedLatch = true;
                       latch.countDown();
                     }
                   }
                 }
               })
           .waitUntilReady();
-      readyObservers.incrementAndGet();
     }
 
     assertTrue(
         String.format(
-            "Timeout, readyObservers = %d, observerWatermarks == %s, numObservers = %d",
-            readyObservers.get(), observerWatermarks, numObservers),
-        latch.await(10, TimeUnit.SECONDS));
+            "Timeout, observerWatermarks = %s, numObservers = %d",
+            observerWatermarks, numObservers),
+        latch.await(30, TimeUnit.SECONDS));
 
     assertEquals(numObservers, observerWatermarks.size());
     long watermark = observerWatermarks.values().stream().min(Long::compare).orElse(Long.MIN_VALUE);
@@ -1511,10 +1509,11 @@ public boolean onNext(StreamElement element, OnNextContext confirm) {
               }
             });
     reader.observe("test", observer);
+    AtomicBoolean finished = new AtomicBoolean();
     Executors.newCachedThreadPool()
         .execute(
             () -> {
-              while (true) {
+              while (!finished.get()) {
                 try {
                   TimeUnit.MILLISECONDS.sleep(100);
                 } catch (InterruptedException ex) {
@@ -1525,6 +1524,7 @@ public boolean onNext(StreamElement element, OnNextContext confirm) {
             });
     latch.await();
     assertEquals(3, restarts.get());
+    finished.set(true);
   }
 
   @Test(timeout = 10000)