[proxima-direct-io-kafka] #350 additional fixes of watermarks

direct/io-kafka/build.gradle

@@ -44,6 +44,7 @@ dependencies {
   testImplementation "org.apache.kafka:kafka_2.13:${kafka_version}"
   testImplementation "org.apache.kafka:kafka-streams:${kafka_version}"
   testImplementation "org.apache.kafka:kafka-streams-test-utils:${kafka_version}"
+  testImplementation libraries.log4j_core
   compileOnly libraries.auto_service_annotations
   compileAnnotationProcessor libraries.lombok
   annotationProcessor libraries.auto_service

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

@@ -38,6 +38,7 @@
 import cz.o2.proxima.internal.com.google.common.annotations.VisibleForTesting;
 import cz.o2.proxima.internal.com.google.common.base.MoreObjects;
 import cz.o2.proxima.internal.com.google.common.base.Preconditions;
+import cz.o2.proxima.internal.com.google.common.collect.Sets;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -430,33 +431,34 @@ private void submitConsumerWithObserver(
                 && !shutdown.get()
                 && !Thread.currentThread().isInterrupted());
 
-            Set<TopicPartition> assignment = kafka.assignment();
-            if (!assignment.isEmpty()) {
-              listener.onPartitionsRevoked(assignment);
-              listener.onPartitionsAssigned(assignment);
+            {
+              Set<TopicPartition> assignment = kafka.assignment();
+              if (!assignment.isEmpty()) {
+                listener.onPartitionsRevoked(assignment);
+                listener.onPartitionsAssigned(assignment);
+              }
             }
 
             readyLatch.countDown();
 
             AtomicReference<Throwable> error = new AtomicReference<>();
             long pollTimeMs = 0L;
 
-            Map<TopicPartition, Long> polledOffsets = new HashMap<>(assignment.size());
+            Map<TopicPartition, Long> polledOffsets = new HashMap<>();
             do {
-              if (poll.isEmpty()) {
-                Optional.ofNullable(watermarkEstimator.get()).ifPresent(consumer::onIdle);
-              }
               logConsumerWatermark(name, offsets, watermarkEstimator, poll.count());
               poll =
                   seekToNewOffsetsIfNeeded(
                       seekOffsets, consumer, watermarkEstimator, kafka, poll, polledOffsets);
               long bytesPolled = 0L;
-              Set<TopicPartition> emptyPartitions = new HashSet<>(assignment);
+              Set<TopicPartition> currentAssignment = kafka.assignment();
+              Set<TopicPartition> nonEmptyPartitions = new HashSet<>(currentAssignment.size());
               for (ConsumerRecord<Object, Object> r : poll) {
                 bytesPolled += r.serializedKeySize() + r.serializedValueSize();
                 TopicPartition tp = new TopicPartition(r.topic(), r.partition());
                 // offsets should be increasing in polled records
                 polledOffsets.put(tp, r.offset());
+                nonEmptyPartitions.add(tp);
                 preWrite.accept(tp, r);
                 StreamElement ingest = serializer.read(r, getEntityDescriptor());
                 if (ingest != null) {
@@ -487,12 +489,16 @@ private void submitConsumerWithObserver(
                   }
                 }
               }
-              if (!emptyPartitions.isEmpty()) {
+              if (!Sets.difference(currentAssignment, nonEmptyPartitions).isEmpty()
+                  || currentAssignment.isEmpty()) {
                 increaseWatermarkOnEmptyPolls(
-                    kafka.endOffsets(kafka.assignment()),
+                    currentAssignment.isEmpty()
+                        ? Collections.emptyMap()
+                        : kafka.endOffsets(currentAssignment),
                     polledOffsets,
                     topicPartitionToId,
-                    watermarkEstimator);
+                    watermarkEstimator,
+                    consumer);
               }
               if (!flushCommits(kafka, consumer)) {
                 handleRebalanceInOffsetCommit(kafka, listener);
@@ -661,15 +667,21 @@ private void increaseWatermarkOnEmptyPolls(
       Map<TopicPartition, Long> endOffsets,
       Map<TopicPartition, Long> polledOffsets,
       Map<TopicPartition, Integer> topicPartitionToId,
-      AtomicReference<PartitionedWatermarkEstimator> watermarkEstimator) {
+      AtomicReference<PartitionedWatermarkEstimator> watermarkEstimator,
+      ElementConsumer<?, ?> consumer) {
 
+    AtomicInteger emptyPartitions = new AtomicInteger();
     endOffsets.forEach(
         (tp, endOffset) -> {
-          long polledOffset = MoreObjects.firstNonNull(polledOffsets.get(tp), 0L);
+          long polledOffset = MoreObjects.firstNonNull(polledOffsets.get(tp), -1L);
           if (endOffset <= polledOffset + 1) {
+            emptyPartitions.incrementAndGet();
             watermarkEstimator.get().idle(topicPartitionToId.get(tp));
           }
         });
+    if (emptyPartitions.get() == endOffsets.size()) {
+      Optional.ofNullable(watermarkEstimator.get()).ifPresent(consumer::onIdle);
+    }
   }
 
   private ObserveHandle createObserveHandle(

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

@@ -1996,14 +1996,15 @@ public void testCachedView() throws InterruptedException {
     final LocalKafkaWriter writer = accessor.newWriter();
     final CachedView view = Optionals.get(accessor.getCachedView(context()));
     final AtomicReference<CountDownLatch> latch = new AtomicReference<>(new CountDownLatch(1));
+    long now = System.currentTimeMillis();
     StreamElement update =
         StreamElement.upsert(
             entity,
             attr,
             UUID.randomUUID().toString(),
             "key",
             attr.getName(),
-            System.currentTimeMillis(),
+            now,
             new byte[] {1, 2});
 
     writer.write(
@@ -2016,15 +2017,14 @@ public void testCachedView() throws InterruptedException {
     latch.set(new CountDownLatch(1));
     view.assign(IntStream.range(0, 3).mapToObj(this::getPartition).collect(Collectors.toList()));
     assertArrayEquals(new byte[] {1, 2}, Optionals.get(view.get("key", attr)).getValue());
-    long now = System.currentTimeMillis();
     update =
         StreamElement.upsert(
             entity,
             attr,
             UUID.randomUUID().toString(),
             "key",
             attr.getName(),
-            now,
+            now + 1,
             new byte[] {1, 2, 3});
     assertEquals(Watermarks.MIN_WATERMARK, getMinWatermark(view.getRunningHandle().get()));
     writer.write(