KAFKA-16100: Add timeout to all the CompletableApplicationEvents

clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java

@@ -929,13 +929,14 @@ public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition
                 return Collections.emptyMap();
             }
 
+            final Timer timer = time.timer(timeout);
             final FetchCommittedOffsetsEvent event = new FetchCommittedOffsetsEvent(
                 partitions,
-                timeout.toMillis());
+                timer);
             wakeupTrigger.setActiveTask(event.future());
             try {
                 final Map<TopicPartition, OffsetAndMetadata> committedOffsets = applicationEventHandler.addAndGet(event,
-                    time.timer(timeout));
+                    timer);
                 committedOffsets.forEach(this::updateLastSeenEpochIfNewer);
                 return committedOffsets;
             } catch (TimeoutException e) {
@@ -981,11 +982,12 @@ public List<PartitionInfo> partitionsFor(String topic, Duration timeout) {
                 throw new TimeoutException();
             }
 
-            final TopicMetadataEvent topicMetadataEvent = new TopicMetadataEvent(topic, timeout.toMillis());
+            final Timer timer = time.timer(timeout);
+            final TopicMetadataEvent topicMetadataEvent = new TopicMetadataEvent(topic, timer);
             wakeupTrigger.setActiveTask(topicMetadataEvent.future());
             try {
                 Map<String, List<PartitionInfo>> topicMetadata =
-                        applicationEventHandler.addAndGet(topicMetadataEvent, time.timer(timeout));
+                        applicationEventHandler.addAndGet(topicMetadataEvent, timer);
 
                 return topicMetadata.getOrDefault(topic, Collections.emptyList());
             } finally {
@@ -1009,10 +1011,11 @@ public Map<String, List<PartitionInfo>> listTopics(Duration timeout) {
                 throw new TimeoutException();
             }
 
-            final AllTopicsMetadataEvent topicMetadataEvent = new AllTopicsMetadataEvent(timeout.toMillis());
+            final Timer timer = time.timer(timeout);
+            final AllTopicsMetadataEvent topicMetadataEvent = new AllTopicsMetadataEvent(timer);
             wakeupTrigger.setActiveTask(topicMetadataEvent.future());
             try {
-                return applicationEventHandler.addAndGet(topicMetadataEvent, time.timer(timeout));
+                return applicationEventHandler.addAndGet(topicMetadataEvent, timer);
             } finally {
                 wakeupTrigger.clearTask();
             }
@@ -1080,16 +1083,18 @@ public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(Map<TopicPartitio
             if (timestampsToSearch.isEmpty()) {
                 return Collections.emptyMap();
             }
+            final Timer timer = time.timer(timeout);
             final ListOffsetsEvent listOffsetsEvent = new ListOffsetsEvent(
                 timestampsToSearch,
-                true);
+                true,
+                timer);
 
             // If timeout is set to zero return empty immediately; otherwise try to get the results
             // and throw timeout exception if it cannot complete in time.
             if (timeout.toMillis() == 0L)
                 return listOffsetsEvent.emptyResult();
 
-            return applicationEventHandler.addAndGet(listOffsetsEvent, time.timer(timeout));
+            return applicationEventHandler.addAndGet(listOffsetsEvent, timer);
         } finally {
             release();
         }
@@ -1130,12 +1135,14 @@ private Map<TopicPartition, Long> beginningOrEndOffset(Collection<TopicPartition
             Map<TopicPartition, Long> timestampToSearch = partitions
                 .stream()
                 .collect(Collectors.toMap(Function.identity(), tp -> timestamp));
+            Timer timer = time.timer(timeout);
             ListOffsetsEvent listOffsetsEvent = new ListOffsetsEvent(
                 timestampToSearch,
-                false);
+                false,
+                timer);
             Map<TopicPartition, OffsetAndTimestamp> offsetAndTimestampMap = applicationEventHandler.addAndGet(
                 listOffsetsEvent,
-                time.timer(timeout));
+                timer);
             return offsetAndTimestampMap
                 .entrySet()
                 .stream()
@@ -1265,7 +1272,7 @@ void prepareShutdown(final Timer timer, final AtomicReference<Throwable> firstEx
         completeQuietly(
             () -> {
                 maybeRevokePartitions();
-                applicationEventHandler.addAndGet(new LeaveOnCloseEvent(), timer);
+                applicationEventHandler.addAndGet(new LeaveOnCloseEvent(timer), timer);
             },
             "Failed to send leaveGroup heartbeat with a timeout(ms)=" + timer.timeoutMs(), firstException);
     }
@@ -1342,7 +1349,7 @@ public void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets, Duration
         long commitStart = time.nanoseconds();
         try {
             Timer requestTimer = time.timer(timeout.toMillis());
-            SyncCommitEvent syncCommitEvent = new SyncCommitEvent(offsets, timeout.toMillis());
+            SyncCommitEvent syncCommitEvent = new SyncCommitEvent(offsets, requestTimer);
             CompletableFuture<Void> commitFuture = commit(syncCommitEvent);
             wakeupTrigger.setActiveTask(commitFuture);
             ConsumerUtils.getResult(commitFuture, requestTimer);
@@ -1456,10 +1463,10 @@ public void unsubscribe() {
         try {
             fetchBuffer.retainAll(Collections.emptySet());
             if (groupMetadata.get().isPresent()) {
-                UnsubscribeEvent unsubscribeEvent = new UnsubscribeEvent();
+                Timer timer = time.timer(Long.MAX_VALUE);
+                UnsubscribeEvent unsubscribeEvent = new UnsubscribeEvent(timer);
                 applicationEventHandler.add(unsubscribeEvent);
                 log.info("Unsubscribing all topics or patterns and assigned partitions");
-                Timer timer = time.timer(Long.MAX_VALUE);
 
                 try {
                     processBackgroundEvents(backgroundEventProcessor, unsubscribeEvent.future(), timer);
@@ -1570,7 +1577,7 @@ private boolean updateFetchPositions(final Timer timer) {
             // Validate positions using the partition leader end offsets, to detect if any partition
             // has been truncated due to a leader change. This will trigger an OffsetForLeaderEpoch
             // request, retrieve the partition end offsets, and validate the current position against it.
-            applicationEventHandler.addAndGet(new ValidatePositionsEvent(), timer);
+            applicationEventHandler.addAndGet(new ValidatePositionsEvent(timer), timer);
 
             cachedSubscriptionHasAllFetchPositions = subscriptions.hasAllFetchPositions();
             if (cachedSubscriptionHasAllFetchPositions) return true;
@@ -1593,7 +1600,7 @@ private boolean updateFetchPositions(final Timer timer) {
             // which are awaiting reset. This will trigger a ListOffset request, retrieve the
             // partition offsets according to the strategy (ex. earliest, latest), and update the
             // positions.
-            applicationEventHandler.addAndGet(new ResetPositionsEvent(), timer);
+            applicationEventHandler.addAndGet(new ResetPositionsEvent(timer), timer);
             return true;
         } catch (TimeoutException e) {
             return false;
@@ -1626,7 +1633,7 @@ private boolean initWithCommittedOffsetsIfNeeded(Timer timer) {
             final FetchCommittedOffsetsEvent event =
                 new FetchCommittedOffsetsEvent(
                     initializingPartitions,
-                    timer.remainingMs());
+                    timer);
             final Map<TopicPartition, OffsetAndMetadata> offsets = applicationEventHandler.addAndGet(event, timer);
             refreshCommittedOffsets(offsets, metadata, subscriptions);
             return true;

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AbstractTopicMetadataEvent.java

@@ -17,25 +17,14 @@
 package org.apache.kafka.clients.consumer.internals.events;
 
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.List;
 import java.util.Map;
 
 public abstract class AbstractTopicMetadataEvent extends CompletableApplicationEvent<Map<String, List<PartitionInfo>>> {
 
-    private final long timeoutMs;
-
-    protected AbstractTopicMetadataEvent(final Type type, final long timeoutMs) {
-        super(type);
-        this.timeoutMs = timeoutMs;
-    }
-
-    public long timeoutMs() {
-        return timeoutMs;
-    }
-
-    @Override
-    public String toStringBase() {
-        return super.toStringBase() + ", timeoutMs=" + timeoutMs;
+    protected AbstractTopicMetadataEvent(final Type type, final Timer timer) {
+        super(type, timer);
     }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AllTopicsMetadataEvent.java

@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 public class AllTopicsMetadataEvent extends AbstractTopicMetadataEvent {
 
-    public AllTopicsMetadataEvent(final long timeoutMs) {
-        super(Type.ALL_TOPICS_METADATA, timeoutMs);
+    public AllTopicsMetadataEvent(final Timer timer) {
+        super(Type.ALL_TOPICS_METADATA, timer);
     }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java

@@ -165,8 +165,7 @@ private void process(final SyncCommitEvent event) {
         }
 
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
-        long expirationTimeoutMs = getExpirationTimeForTimeout(event.retryTimeoutMs());
-        CompletableFuture<Void> future = manager.commitSync(event.offsets(), expirationTimeoutMs);
+        CompletableFuture<Void> future = manager.commitSync(event.offsets(), event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
@@ -177,8 +176,7 @@ private void process(final FetchCommittedOffsetsEvent event) {
             return;
         }
         CommitRequestManager manager = requestManagers.commitRequestManager.get();
-        long expirationTimeMs = getExpirationTimeForTimeout(event.timeout());
-        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future = manager.fetchOffsets(event.partitions(), expirationTimeMs);
+        CompletableFuture<Map<TopicPartition, OffsetAndMetadata>> future = manager.fetchOffsets(event.partitions(), event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
@@ -250,16 +248,14 @@ private void process(final ValidatePositionsEvent event) {
     }
 
     private void process(final TopicMetadataEvent event) {
-        final long expirationTimeMs = getExpirationTimeForTimeout(event.timeoutMs());
         final CompletableFuture<Map<String, List<PartitionInfo>>> future =
-                requestManagers.topicMetadataRequestManager.requestTopicMetadata(event.topic(), expirationTimeMs);
+                requestManagers.topicMetadataRequestManager.requestTopicMetadata(event.topic(), event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
     private void process(final AllTopicsMetadataEvent event) {
-        final long expirationTimeMs = getExpirationTimeForTimeout(event.timeoutMs());
         final CompletableFuture<Map<String, List<PartitionInfo>>> future =
-                requestManagers.topicMetadataRequestManager.requestAllTopicsMetadata(expirationTimeMs);
+                requestManagers.topicMetadataRequestManager.requestAllTopicsMetadata(event.deadlineMs());
         future.whenComplete(complete(event.future()));
     }
 
@@ -296,19 +292,6 @@ private void process(final LeaveOnCloseEvent event) {
         future.whenComplete(complete(event.future()));
     }
 
-    /**
-     * @return Expiration time in milliseconds calculated with the current time plus the given
-     * timeout. Returns Long.MAX_VALUE if the expiration overflows it.
-     * Visible for testing.
-     */
-    long getExpirationTimeForTimeout(final long timeoutMs) {
-        long expiration = System.currentTimeMillis() + timeoutMs;
-        if (expiration < 0) {
-            return Long.MAX_VALUE;
-        }
-        return expiration;
-    }
-
     private <T> BiConsumer<? super T, ? super Throwable> complete(final CompletableFuture<T> b) {
         return (value, exception) -> {
             if (exception != null)

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/AsyncCommitEvent.java

@@ -27,6 +27,6 @@
 public class AsyncCommitEvent extends CommitEvent {
 
     public AsyncCommitEvent(final Map<TopicPartition, OffsetAndMetadata> offsets) {
-        super(Type.COMMIT_ASYNC, offsets);
+        super(Type.COMMIT_ASYNC, offsets, Long.MAX_VALUE);
     }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CommitEvent.java

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.Collections;
 import java.util.Map;
@@ -29,15 +30,28 @@ public abstract class CommitEvent extends CompletableApplicationEvent<Void> {
      */
     private final Map<TopicPartition, OffsetAndMetadata> offsets;
 
-    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets) {
-        super(type);
-        this.offsets = Collections.unmodifiableMap(offsets);
+    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets, final Timer timer) {
+        super(type, timer);
+        this.offsets = validate(offsets);
+    }
+
+    protected CommitEvent(final Type type, final Map<TopicPartition, OffsetAndMetadata> offsets, final long deadlineMs) {
+        super(type, deadlineMs);
+        this.offsets = validate(offsets);
+    }
 
+    /**
+     * Validates the offsets are not negative and then returns the given offset map as
+     * {@link Collections#unmodifiableMap(Map) as unmodifiable}.
+     */
+    private static Map<TopicPartition, OffsetAndMetadata> validate(final Map<TopicPartition, OffsetAndMetadata> offsets) {
         for (OffsetAndMetadata offsetAndMetadata : offsets.values()) {
             if (offsetAndMetadata.offset() < 0) {
                 throw new IllegalArgumentException("Invalid offset: " + offsetAndMetadata.offset());
             }
         }
+
+        return Collections.unmodifiableMap(offsets);
     }
 
     public Map<TopicPartition, OffsetAndMetadata> offsets() {

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/CompletableApplicationEvent.java

@@ -16,6 +16,9 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -27,19 +30,32 @@
 public abstract class CompletableApplicationEvent<T> extends ApplicationEvent implements CompletableEvent<T> {
 
     private final CompletableFuture<T> future;
+    private final long deadlineMs;
 
-    protected CompletableApplicationEvent(final Type type) {
+    protected CompletableApplicationEvent(final Type type, final Timer timer) {
         super(type);
         this.future = new CompletableFuture<>();
+        Objects.requireNonNull(timer);
+        this.deadlineMs = timer.remainingMs() + timer.currentTimeMs();
+    }
+
+    protected CompletableApplicationEvent(final Type type, final long deadlineMs) {
+        super(type);
+        this.future = new CompletableFuture<>();
+        this.deadlineMs = deadlineMs;
     }
 
     @Override
     public CompletableFuture<T> future() {
         return future;
     }
 
+    public long deadlineMs() {
+        return deadlineMs;
+    }
+
     @Override
     protected String toStringBase() {
-        return super.toStringBase() + ", future=" + future;
+        return super.toStringBase() + ", future=" + future + ", deadlineMs=" + deadlineMs;
     }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/FetchCommittedOffsetsEvent.java

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Timer;
 
 import java.util.Collections;
 import java.util.Map;
@@ -30,27 +31,17 @@ public class FetchCommittedOffsetsEvent extends CompletableApplicationEvent<Map<
      */
     private final Set<TopicPartition> partitions;
 
-    /**
-     * Time until which the request will be retried if it fails with a retriable error.
-     */
-    private final long timeoutMs;
-
-    public FetchCommittedOffsetsEvent(final Set<TopicPartition> partitions, final long timeoutMs) {
-        super(Type.FETCH_COMMITTED_OFFSETS);
+    public FetchCommittedOffsetsEvent(final Set<TopicPartition> partitions, final Timer timer) {
+        super(Type.FETCH_COMMITTED_OFFSETS, timer);
         this.partitions = Collections.unmodifiableSet(partitions);
-        this.timeoutMs = timeoutMs;
     }
 
     public Set<TopicPartition> partitions() {
         return partitions;
     }
 
-    public long timeout() {
-        return timeoutMs;
-    }
-
     @Override
     public String toStringBase() {
-        return super.toStringBase() + ", partitions=" + partitions + ", partitions=" + partitions;
+        return super.toStringBase() + ", partitions=" + partitions;
     }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/LeaveOnCloseEvent.java

@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
+import org.apache.kafka.common.utils.Timer;
+
 public class LeaveOnCloseEvent extends CompletableApplicationEvent<Void> {
 
-    public LeaveOnCloseEvent() {
-        super(Type.LEAVE_ON_CLOSE);
+    public LeaveOnCloseEvent(final Timer timer) {
+        super(Type.LEAVE_ON_CLOSE, timer);
     }
 }