JDBC Client: fix a bug where TrinoResultSet cannot be properly closed

client/trino-jdbc/src/main/java/io/trino/jdbc/TrinoResultSet.java

@@ -13,6 +13,7 @@
  */
 package io.trino.jdbc;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.Streams;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
@@ -29,9 +30,9 @@
 import java.util.Optional;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
 import java.util.concurrent.Semaphore;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
@@ -144,27 +145,39 @@ private static <T> Iterator<T> flatten(Iterator<Iterable<T>> iterator, long maxR
         return stream.iterator();
     }
 
-    private static class AsyncIterator<T>
+    @VisibleForTesting
+    static class AsyncIterator<T>
             extends AbstractIterator<T>
     {
         private static final int MAX_QUEUED_ROWS = 50_000;
         private static final ExecutorService executorService = newCachedThreadPool(
                 new ThreadFactoryBuilder().setNameFormat("Trino JDBC worker-%s").setDaemon(true).build());
 
         private final StatementClient client;
-        private final BlockingQueue<T> rowQueue = new ArrayBlockingQueue<>(MAX_QUEUED_ROWS);
+        private final BlockingQueue<T> rowQueue;
         // Semaphore to indicate that some data is ready.
         // Each permit represents a row of data (or that the underlying iterator is exhausted).
         private final Semaphore semaphore = new Semaphore(0);
-        private final CompletableFuture<Void> future;
+        private final Future<?> future;
+        private volatile boolean cancelled;
+        private volatile boolean finished;
 
         public AsyncIterator(Iterator<T> dataIterator, StatementClient client)
+        {
+            this(dataIterator, client, Optional.empty());
+        }
+
+        @VisibleForTesting
+        AsyncIterator(Iterator<T> dataIterator, StatementClient client, Optional<BlockingQueue<T>> queue)
         {
             requireNonNull(dataIterator, "dataIterator is null");
             this.client = client;
-            this.future = CompletableFuture.runAsync(() -> {
+            this.rowQueue = queue.orElseGet(() -> new ArrayBlockingQueue<>(MAX_QUEUED_ROWS));
+            this.cancelled = false;
+            this.finished = false;
+            this.future = executorService.submit(() -> {
                 try {
-                    while (dataIterator.hasNext()) {
+                    while (!cancelled && dataIterator.hasNext()) {
                         rowQueue.put(dataIterator.next());
                         semaphore.release();
                     }
@@ -174,22 +187,46 @@ public AsyncIterator(Iterator<T> dataIterator, StatementClient client)
                 }
                 finally {
                     semaphore.release();
+                    finished = true;
                 }
-            }, executorService);
+            });
         }
 
         public void cancel()
         {
+            cancelled = true;
             future.cancel(true);
+            cleanup();
         }
 
         public void interrupt(InterruptedException e)
         {
-            client.close();
+            cleanup();
             Thread.currentThread().interrupt();
             throw new RuntimeException(new SQLException("ResultSet thread was interrupted", e));
         }
 
+        private void cleanup()
+        {
+            // When thread interruption is mis-handled by underlying implementation of `client`, the thread which
+            // is working for `future` may be blocked by `rowQueue.put` (`rowQueue` is full) and will never finish
+            // its work. It is necessary to close `client` and drain `rowQueue` to avoid such leaks.
+            client.close();
+            rowQueue.clear();
+        }
+
+        @VisibleForTesting
+        Future<?> getFuture()
+        {
+            return future;
+        }
+
+        @VisibleForTesting
+        boolean isBackgroundThreadFinished()
+        {
+            return finished;
+        }
+
         @Override
         protected T computeNext()
         {

client/trino-jdbc/src/test/java/io/trino/jdbc/TestJdbcResultSet.java

@@ -24,6 +24,9 @@
 
 import static java.lang.String.format;
 
+/**
+ * An integration test for JDBC client interacting with Trino server.
+ */
 public class TestJdbcResultSet
         extends BaseTestJdbcResultSet
 {

client/trino-jdbc/src/test/java/io/trino/jdbc/TestTrinoResultSet.java

@@ -0,0 +1,272 @@
+/*
+ * Licensed 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 io.trino.jdbc;
+
+import com.google.common.collect.ImmutableList;
+import io.trino.client.ClientSelectedRole;
+import io.trino.client.QueryData;
+import io.trino.client.QueryStatusInfo;
+import io.trino.client.StatementClient;
+import io.trino.client.StatementStats;
+import org.testng.annotations.Test;
+
+import java.time.ZoneId;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.testng.Assert.assertTrue;
+
+/**
+ * A unit test for {@link TrinoResultSet}.
+ *
+ * @see TestJdbcResultSet for an integration test.
+ */
+public class TestTrinoResultSet
+{
+    @Test(timeOut = 10000)
+    public void testIteratorCancelWhenQueueNotFull()
+            throws Exception
+    {
+        AtomicReference<Thread> thread = new AtomicReference<>();
+        CountDownLatch interruptedButSwallowedLatch = new CountDownLatch(1);
+        MockAsyncIterator<Iterable<List<Object>>> iterator = new MockAsyncIterator<>(
+                new Iterator<Iterable<List<Object>>>()
+                {
+                    @Override
+                    public boolean hasNext()
+                    {
+                        return true;
+                    }
+
+                    @Override
+                    public Iterable<List<Object>> next()
+                    {
+                        thread.compareAndSet(null, Thread.currentThread());
+                        try {
+                            TimeUnit.MILLISECONDS.sleep(1000);
+                        }
+                        catch (InterruptedException e) {
+                            interruptedButSwallowedLatch.countDown();
+                        }
+                        return ImmutableList.of((ImmutableList.of(new Object())));
+                    }
+                },
+                new ArrayBlockingQueue<>(100));
+
+        while (thread.get() == null || thread.get().getState() != Thread.State.TIMED_WAITING) {
+            // wait for thread being waiting
+        }
+        iterator.cancel();
+        while (!iterator.getFuture().isDone() || !iterator.isBackgroundThreadFinished()) {
+            TimeUnit.MILLISECONDS.sleep(10);
+        }
+        boolean interruptedButSwallowed = interruptedButSwallowedLatch.await(5000, TimeUnit.MILLISECONDS);
+        assertTrue(interruptedButSwallowed);
+    }
+
+    @Test(timeOut = 10000)
+    public void testIteratorCancelWhenQueueIsFull()
+            throws Exception
+    {
+        BlockingQueue<Iterable<List<Object>>> queue = new ArrayBlockingQueue<>(1);
+        queue.put(ImmutableList.of());
+        // queue is full at the beginning
+        AtomicReference<Thread> thread = new AtomicReference<>();
+        MockAsyncIterator<Iterable<List<Object>>> iterator = new MockAsyncIterator<>(
+                new Iterator<Iterable<List<Object>>>()
+                {
+                    @Override
+                    public boolean hasNext()
+                    {
+                        return true;
+                    }
+
+                    @Override
+                    public Iterable<List<Object>> next()
+                    {
+                        thread.compareAndSet(null, Thread.currentThread());
+                        return ImmutableList.of((ImmutableList.of(new Object())));
+                    }
+                },
+                queue);
+
+        while (thread.get() == null || thread.get().getState() != Thread.State.WAITING) {
+            // wait for thread being waiting (for queue being not full)
+            TimeUnit.MILLISECONDS.sleep(10);
+        }
+        iterator.cancel();
+        while (!iterator.isBackgroundThreadFinished()) {
+            TimeUnit.MILLISECONDS.sleep(10);
+        }
+    }
+
+    private static class MockAsyncIterator<T>
+            extends TrinoResultSet.AsyncIterator<T>
+    {
+        public MockAsyncIterator(Iterator<T> dataIterator, BlockingQueue<T> queue)
+        {
+            super(
+                    dataIterator,
+                    new StatementClient()
+                    {
+                        @Override
+                        public String getQuery()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public ZoneId getTimeZone()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public boolean isRunning()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public boolean isClientAborted()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public boolean isClientError()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public boolean isFinished()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public StatementStats getStats()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public QueryStatusInfo currentStatusInfo()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public QueryData currentData()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public QueryStatusInfo finalStatusInfo()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Optional<String> getSetCatalog()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Optional<String> getSetSchema()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Optional<String> getSetPath()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Map<String, String> getSetSessionProperties()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Set<String> getResetSessionProperties()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Map<String, ClientSelectedRole> getSetRoles()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Map<String, String> getAddedPreparedStatements()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public Set<String> getDeallocatedPreparedStatements()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public String getStartedTransactionId()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public boolean isClearTransactionId()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public boolean advance()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public void cancelLeafStage()
+                        {
+                            throw new UnsupportedOperationException();
+                        }
+
+                        @Override
+                        public void close()
+                        {
+                            // do nothing
+                        }
+                    },
+                    Optional.of(queue));
+        }
+    }
+}