KAFKA-9051: Prematurely complete source offset read requests for stopped tasks

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java

@@ -49,10 +49,10 @@
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.storage.CloseableOffsetStorageReader;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.OffsetBackingStore;
-import org.apache.kafka.connect.storage.OffsetStorageReader;
 import org.apache.kafka.connect.storage.OffsetStorageReaderImpl;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
 import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -511,7 +511,7 @@ private WorkerTask buildWorkerTask(ClusterConfigState configState,
             retryWithToleranceOperator.reporters(sourceTaskReporters(id, connConfig, errorHandlingMetrics));
             TransformationChain<SourceRecord> transformationChain = new TransformationChain<>(connConfig.<SourceRecord>transformations(), retryWithToleranceOperator);
             log.info("Initializing: {}", transformationChain);
-            OffsetStorageReader offsetReader = new OffsetStorageReaderImpl(offsetBackingStore, id.connector(),
+            CloseableOffsetStorageReader offsetReader = new OffsetStorageReaderImpl(offsetBackingStore, id.connector(),
                     internalKeyConverter, internalValueConverter);
             OffsetStorageWriter offsetWriter = new OffsetStorageWriter(offsetBackingStore, id.connector(),
                     internalKeyConverter, internalValueConverter);

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTask.java

@@ -39,9 +39,9 @@
 import org.apache.kafka.connect.runtime.errors.Stage;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTask;
+import org.apache.kafka.connect.storage.CloseableOffsetStorageReader;
 import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.HeaderConverter;
-import org.apache.kafka.connect.storage.OffsetStorageReader;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
 import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -75,7 +75,7 @@ class WorkerSourceTask extends WorkerTask {
     private final HeaderConverter headerConverter;
     private final TransformationChain<SourceRecord> transformationChain;
     private KafkaProducer<byte[], byte[]> producer;
-    private final OffsetStorageReader offsetReader;
+    private final CloseableOffsetStorageReader offsetReader;
     private final OffsetStorageWriter offsetWriter;
     private final Time time;
     private final SourceTaskMetricsGroup sourceTaskMetricsGroup;
@@ -105,7 +105,7 @@ public WorkerSourceTask(ConnectorTaskId id,
                             HeaderConverter headerConverter,
                             TransformationChain<SourceRecord> transformationChain,
                             KafkaProducer<byte[], byte[]> producer,
-                            OffsetStorageReader offsetReader,
+                            CloseableOffsetStorageReader offsetReader,
                             OffsetStorageWriter offsetWriter,
                             WorkerConfig workerConfig,
                             ClusterConfigState configState,
@@ -175,6 +175,7 @@ protected void releaseResources() {
     @Override
     public void stop() {
         super.stop();
+        offsetReader.close();
         stopRequestedLatch.countDown();
         synchronized (this) {
             if (finishedStart)

connect/runtime/src/main/java/org/apache/kafka/connect/storage/CloseableOffsetStorageReader.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.storage;
+
+import java.io.Closeable;
+import java.util.Collection;
+import java.util.Map;
+
+public interface CloseableOffsetStorageReader extends Closeable, OffsetStorageReader {
+
+    /**
+     * Invoke {@link OffsetReadFuture#forceComplete()} on all outstanding offset read requests, and
+     * return an empty map and {@code null} in all future calls to {@link #offsets(Collection)} and
+     * {@link #offset(Map)}, respectively. This is useful for unblocking task threads which need to shut
+     * down but are blocked on offset reads.
+     */
+    void close();
+}

connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java

@@ -42,6 +42,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * <p>
@@ -118,17 +119,8 @@ public void stop() {
     }
 
     @Override
-    public Future<Map<ByteBuffer, ByteBuffer>> get(final Collection<ByteBuffer> keys,
-                                                   final Callback<Map<ByteBuffer, ByteBuffer>> callback) {
-        ConvertingFutureCallback<Void, Map<ByteBuffer, ByteBuffer>> future = new ConvertingFutureCallback<Void, Map<ByteBuffer, ByteBuffer>>(callback) {
-            @Override
-            public Map<ByteBuffer, ByteBuffer> convert(Void result) {
-                Map<ByteBuffer, ByteBuffer> values = new HashMap<>();
-                for (ByteBuffer key : keys)
-                    values.put(key, data.get(key));
-                return values;
-            }
-        };
+    public OffsetReadFuture get(final Collection<ByteBuffer> keys) {
+        KafkaOffsetReadFuture future = new KafkaOffsetReadFuture(keys);
         // This operation may be relatively (but not too) expensive since it always requires checking end offsets, even
         // if we've already read up to the end. However, it also should not be common (offsets should only be read when
         // resetting a task). Always requiring that we read to the end is simpler than trying to differentiate when it
@@ -231,5 +223,50 @@ public synchronized Void get(long timeout, TimeUnit unit) throws InterruptedExce
         }
     }
 
+    private class KafkaOffsetReadFuture extends ConvertingFutureCallback<Void, Map<ByteBuffer, ByteBuffer>> implements OffsetReadFuture {
+        private final Collection<ByteBuffer> keys;
+        private final AtomicBoolean completed;
+        private transient boolean forceComplete;
+
+        public KafkaOffsetReadFuture(Collection<ByteBuffer> keys) {
+            super(null);
+            this.keys = keys;
+            this.completed = new AtomicBoolean(false);
+            this.forceComplete = false;
+        }
+
+        @Override
+        public Map<ByteBuffer, ByteBuffer> convert(Void result) {
+            Map<ByteBuffer, ByteBuffer> values = new HashMap<>();
+            if (!forceComplete) {
+                for (ByteBuffer key : keys)
+                    values.put(key, data.get(key));
+            }
+            return values;
+        }
+
+        @Override
+        public void forceComplete() {
+            log.debug("Forcing completion of offset read future");
+            forceComplete = true;
+            onCompletion(null, null);
+        }
+
+        @Override
+        public Map<ByteBuffer, ByteBuffer> get() throws InterruptedException, ExecutionException {
+            log.trace("Blocking on offset read request");
+            return super.get();
+        }
+
+        @Override
+        public void onCompletion(Throwable error, Void result) {
+            if (!completed.getAndSet(true)) {
+                log.trace("Offset request completed (either normally or forcibly)");
+                super.onCompletion(error, result);
+            } else {
+                log.trace("Ignoring onCompletion invocation as offset read has already completed");
+            }
+        }
 
+    }
 }

connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryOffsetBackingStore.java

@@ -27,10 +27,12 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Implementation of OffsetBackingStore that doesn't actually persist any data. To ensure this
@@ -75,22 +77,8 @@ public void stop() {
     }
 
     @Override
-    public Future<Map<ByteBuffer, ByteBuffer>> get(
-            final Collection<ByteBuffer> keys,
-            final Callback<Map<ByteBuffer, ByteBuffer>> callback) {
-        return executor.submit(new Callable<Map<ByteBuffer, ByteBuffer>>() {
-            @Override
-            public Map<ByteBuffer, ByteBuffer> call() throws Exception {
-                Map<ByteBuffer, ByteBuffer> result = new HashMap<>();
-                for (ByteBuffer key : keys) {
-                    result.put(key, data.get(key));
-                }
-                if (callback != null)
-                    callback.onCompletion(null, result);
-                return result;
-            }
-        });
-
+    public OffsetReadFuture get(final Collection<ByteBuffer> keys) {
+        return new MemoryOffsetReadFuture(keys);
     }
 
     @Override
@@ -114,4 +102,74 @@ public Void call() throws Exception {
     protected void save() {
 
     }
+
+    private class MemoryOffsetReadFuture implements OffsetReadFuture {
+        private final Map<ByteBuffer, ByteBuffer> result;
+        private final Collection<ByteBuffer> keys;
+        private final CountDownLatch completed;
+        private final Future<?> underlyingOffsetFuture;
+
+        public MemoryOffsetReadFuture(Collection<ByteBuffer> keys) {
+            this.result = new HashMap<>();
+            this.keys = keys;
+            this.completed = new CountDownLatch(1);
+            this.underlyingOffsetFuture = executor.submit(new Runnable() {
+                @Override
+                public void run() {
+                    synchronized (this) {
+                        if (isDone()) {
+                            return;
+                        }
+                        collectResults();
+                        completed.countDown();
+                    }
+                }
+            });
+        }
+
+        @Override
+        public void forceComplete() {
+            synchronized (this) {
+                if (isDone()) {
+                    return;
+                }
+                completed.countDown();
+            }
+        }
+
+        @Override
+        public boolean cancel(boolean mayInterruptIfRunning) {
+            return underlyingOffsetFuture.cancel(mayInterruptIfRunning);
+        }
+
+        @Override
+        public boolean isCancelled() {
+            return underlyingOffsetFuture.isCancelled();
+        }
+
+        @Override
+        public boolean isDone() {
+            return completed.getCount() == 0;
+        }
+
+        @Override
+        public Map<ByteBuffer, ByteBuffer> get() throws InterruptedException {
+            completed.await();
+            return result;
+        }
+
+        @Override
+        public Map<ByteBuffer, ByteBuffer> get(long timeout, TimeUnit unit)
+            throws TimeoutException, InterruptedException {
+            if (!completed.await(timeout, unit))
+                throw new TimeoutException("Timed out waiting for future");
+            return result;
+        }
+
+        private void collectResults() {
+            for (ByteBuffer key : keys) {
+                result.put(key, data.get(key));
+            }
+        }
+    }
 }

connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetBackingStore.java

@@ -53,21 +53,17 @@ public interface OffsetBackingStore {
     /**
      * Get the values for the specified keys
      * @param keys list of keys to look up
-     * @param callback callback to invoke on completion
      * @return future for the resulting map from key to value
      */
-    Future<Map<ByteBuffer, ByteBuffer>> get(
-            Collection<ByteBuffer> keys,
-            Callback<Map<ByteBuffer, ByteBuffer>> callback);
+    OffsetReadFuture get(Collection<ByteBuffer> keys);
 
     /**
      * Set the specified keys and values.
      * @param values map from key to value
      * @param callback callback to invoke on completion
      * @return void future for the operation
      */
-    Future<Void> set(Map<ByteBuffer, ByteBuffer> values,
-                            Callback<Void> callback);
+    Future<Void> set(Map<ByteBuffer, ByteBuffer> values, Callback<Void> callback);
 
     /**
      * Configure class with the given key-value pairs

connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetReadFuture.java

@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.storage;
+
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+import java.util.concurrent.Future;
+
+public interface OffsetReadFuture extends Future<Map<ByteBuffer, ByteBuffer>> {
+
+    /**
+     * Without {@link Future#cancel(boolean) cancelling} the operation, immediately return an empty
+     * map to any threads which are waiting for this future to complete.
+     */
+    void forceComplete();
+}