Separate thread pool for I/O, add more tests for streaming, along with other minor changes for code clean up

Author: zhicwu

clickhouse-client/src/main/java/com/clickhouse/client/AbstractClient.java

@@ -170,9 +170,8 @@ public void init(ClickHouseConfig config) {
             this.config = config;
             if (this.executor == null) { // only initialize once
                 int threads = config.getMaxThreadsPerClient();
-                this.executor = threads <= 0 ? ClickHouseClient.getExecutorService()
-                        : ClickHouseUtils.newThreadPool(getClass().getSimpleName(), threads,
-                                config.getMaxQueuedRequests());
+                this.executor = threads < 1 ? ClickHouseClient.getExecutorService()
+                        : ClickHouseUtils.newThreadPool(this, threads, config.getMaxQueuedRequests());
             }
 
             initialized = true;
@@ -196,27 +195,16 @@ public final void close() {
         try {
             server = null;
 
-            if (executor != null) {
-                executor.shutdown();
-            }
-
             if (connection != null) {
                 closeConnection(connection, false);
+                connection = null;
             }
 
-            // shutdown* won't shutdown commonPool, so awaitTermination will always time out
-            // on the other hand, for a client-specific thread pool, we'd better shut it
-            // down for real
-            if (executor != null && config.getMaxThreadsPerClient() > 0
-                    && !executor.awaitTermination(config.getConnectionTimeout(), TimeUnit.MILLISECONDS)) {
-                executor.shutdownNow();
+            // avoid shutting down shared thread pool
+            if (executor != null && config.getMaxThreadsPerClient() > 0 && !executor.isTerminated()) {
+                executor.shutdown();
             }
-
             executor = null;
-            connection = null;
-        } catch (InterruptedException e) {
-            log.warn("Got interrupted when closing client", e);
-            Thread.currentThread().interrupt();
         } catch (Exception e) {
             log.warn("Exception occurred when closing client", e);
         } finally {
@@ -226,7 +214,7 @@ public final void close() {
                     closeConnection(connection, true);
                 }
 
-                if (executor != null) {
+                if (executor != null && config.getMaxThreadsPerClient() > 0) {
                     executor.shutdownNow();
                 }
             } finally {

clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseClientBuilder.java

@@ -6,7 +6,6 @@
 import java.util.Objects;
 import java.util.ServiceLoader;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.ForkJoinPool;
 
 import com.clickhouse.client.config.ClickHouseOption;
 import com.clickhouse.client.config.ClickHouseDefaults;
@@ -23,23 +22,17 @@ public class ClickHouseClientBuilder {
     static {
         int maxThreads = (int) ClickHouseDefaults.MAX_THREADS.getEffectiveDefaultValue();
         int maxRequests = (int) ClickHouseDefaults.MAX_REQUESTS.getEffectiveDefaultValue();
+        long keepAliveTimeoutMs = (long) ClickHouseDefaults.THREAD_KEEPALIVE_TIMEOUT.getEffectiveDefaultValue();
 
-        if (maxThreads <= 0 && maxRequests <= 0) {
-            // java -XX:+UnlockDiagnosticVMOptions -XX:NativeMemoryTracking=summary
-            // -XX:+PrintNMTStatistics -version
-            defaultExecutor = ForkJoinPool.commonPool();
-        } else {
-            if (maxThreads <= 0) {
-                maxThreads = Runtime.getRuntime().availableProcessors();
-            }
-
-            if (maxRequests <= 0) {
-                maxRequests = 0;
-            }
-
-            defaultExecutor = ClickHouseUtils.newThreadPool(ClickHouseClient.class.getSimpleName(), maxThreads,
-                    maxRequests);
+        if (maxThreads <= 0) {
+            maxThreads = Runtime.getRuntime().availableProcessors();
+        }
+        if (maxRequests <= 0) {
+            maxRequests = 0;
         }
+
+        defaultExecutor = ClickHouseUtils.newThreadPool(ClickHouseClient.class.getSimpleName(), maxThreads,
+                maxThreads * 2, maxRequests, keepAliveTimeoutMs);
     }
 
     protected ClickHouseConfig config;

clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseException.java

@@ -37,7 +37,7 @@ private static String buildErrorMessage(int code, String message, ClickHouseNode
         }
 
         if (server != null) {
-            builder.append(" on server ").append(server);
+            builder.append(", server ").append(server);
         }
 
         return builder.toString();

clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseInputStream.java

@@ -11,13 +11,17 @@
 import java.util.concurrent.TimeUnit;
 
 /**
- * Extended input stream.
+ * Extended input stream for read optimization.
  */
 public abstract class ClickHouseInputStream extends InputStream {
+    /**
+     * Empty byte array.
+     */
+    public static final byte[] EMPTY_BYTES = new byte[0];
     /**
      * Empty and read-only byte buffer.
      */
-    public static final ByteBuffer EMPTY = ByteBuffer.wrap(new byte[0]).asReadOnlyBuffer();
+    public static final ByteBuffer EMPTY_BUFFER = ByteBuffer.wrap(EMPTY_BYTES).asReadOnlyBuffer();
 
     static final class BlockingInputStream extends ClickHouseInputStream {
         private final BlockingQueue<ByteBuffer> queue;
@@ -28,7 +32,7 @@ static final class BlockingInputStream extends ClickHouseInputStream {
         private boolean closed;
 
         BlockingInputStream(BlockingQueue<ByteBuffer> queue, int timeout) {
-            this.queue = queue;
+            this.queue = ClickHouseChecker.nonNull(queue, "Queue");
             this.timeout = timeout;
 
             this.buffer = null;
@@ -37,10 +41,12 @@ static final class BlockingInputStream extends ClickHouseInputStream {
 
         private void ensureOpen() throws IOException {
             if (closed) {
-                throw new IOException("Stream has been closed");
+                throw new IOException(
+                        ClickHouseUtils.format("Blocking stream(queue: %d, buffer: %d) has been closed",
+                                queue.size(), buffer != null ? buffer.remaining() : 0));
             }
 
-            if (buffer == null || (buffer != EMPTY && !buffer.hasRemaining())) {
+            if (buffer == null || (buffer != EMPTY_BUFFER && !buffer.hasRemaining())) {
                 updateBuffer();
             }
         }
@@ -65,9 +71,11 @@ private int updateBuffer() throws IOException {
 
         @Override
         public int available() throws IOException {
-            ensureOpen();
+            if (closed || buffer == EMPTY_BUFFER) {
+                return 0;
+            }
 
-            return buffer.remaining();
+            return (buffer == null || !buffer.hasRemaining()) ? updateBuffer() : buffer.remaining();
         }
 
         @Override
@@ -87,7 +95,7 @@ public void close() throws IOException {
         public byte readByte() throws IOException {
             ensureOpen();
 
-            if (buffer == EMPTY) {
+            if (buffer == EMPTY_BUFFER) {
                 close();
                 throw new EOFException();
             }
@@ -99,7 +107,7 @@ public byte readByte() throws IOException {
         public int read() throws IOException {
             ensureOpen();
 
-            if (buffer == EMPTY) {
+            if (buffer == EMPTY_BUFFER) {
                 return -1;
             }
 
@@ -112,7 +120,7 @@ public int read(byte[] b, int off, int len) throws IOException {
 
             int counter = 0;
             while (len > 0) {
-                if (buffer == EMPTY) {
+                if (buffer == EMPTY_BUFFER) {
                     return counter > 0 ? counter : -1;
                 }
 
@@ -162,7 +170,7 @@ public long skip(long n) throws IOException {
             // peforms better but this is a bit tricky
             if (n == Long.MAX_VALUE) {
                 long counter = buffer.remaining();
-                while (buffer != EMPTY && buffer.limit() > 0) {
+                while (buffer != EMPTY_BUFFER && buffer.limit() > 0) {
                     counter += buffer.limit();
                     updateBuffer();
                 }
@@ -184,20 +192,32 @@ static final class WrappedInputStream extends ClickHouseInputStream {
             closed = false;
         }
 
+        private void ensureOpen() throws IOException {
+            if (closed) {
+                throw new IOException(ClickHouseUtils.format("Wrapped stream(%s) has been closed", in));
+            }
+        }
+
         @Override
         public int available() throws IOException {
             return !closed ? in.available() : 0;
         }
 
         @Override
         public byte readByte() throws IOException {
+            ensureOpen();
+
             int v = in.read();
-            if (v == -1) {
-                close();
-                throw new EOFException();
+            if (v != -1) {
+                return (byte) v;
             }
 
-            return (byte) v;
+            try {
+                close();
+            } catch (IOException e) {
+                // ignore
+            }
+            throw new EOFException();
         }
 
         @Override
@@ -207,25 +227,30 @@ public boolean isClosed() {
 
         @Override
         public void close() throws IOException {
-            try {
-                in.close();
-            } finally {
-                closed = true;
+            if (!closed) {
+                try {
+                    in.close();
+                } finally {
+                    closed = true;
+                }
             }
         }
 
         @Override
         public int read() throws IOException {
+            ensureOpen();
             return in.read();
         }
 
         @Override
         public int read(byte[] b, int off, int len) throws IOException {
+            ensureOpen();
             return in.read(b, off, len);
         }
 
         @Override
         public long skip(long n) throws IOException {
+            ensureOpen();
             return in.skip(n);
         }
     }
@@ -238,21 +263,23 @@ public long skip(long n) throws IOException {
      * @return wrapped input
      */
     public static ClickHouseInputStream of(BlockingQueue<ByteBuffer> queue, int timeout) {
-        return new BlockingInputStream(ClickHouseChecker.nonNull(queue, "queue"), timeout);
+        return new BlockingInputStream(queue, timeout);
     }
 
     /**
      * Wraps the given input stream.
      *
      * @param input non-null input stream
-     * @return wrapped input
+     * @return wrapped input, or the same input if it's instance of
+     *         {@link ClickHouseInputStream}
      */
     public static ClickHouseInputStream of(InputStream input) {
         return input instanceof ClickHouseInputStream ? (ClickHouseInputStream) input : new WrappedInputStream(input);
     }
 
     /**
-     * Reads an unsigned byte from the input stream.
+     * Reads an unsigned byte from the input stream. Unlike {@link #read()}, it will
+     * throw {@link IOException} if the input stream has been closed.
      *
      * @return unsigned byte
      * @throws IOException when failed to read value from input stream or reached
@@ -263,8 +290,10 @@ public int readUnsignedByte() throws IOException {
     }
 
     /**
-     * Reads one single byte from the input stream. It's supposed to be faster than
-     * {@link #read()}.
+     * Reads one single byte from the input stream. Unlike {@link #read()}, it will
+     * throw {@link IOException} if the input stream has been closed. In general,
+     * this method should be faster than {@link #read()}, especially when it's an
+     * input stream backed by byte[] or {@link java.nio.ByteBuffer}.
      *
      * @return byte value if present
      * @throws IOException when failed to read value from input stream or reached
@@ -274,14 +303,19 @@ public int readUnsignedByte() throws IOException {
 
     /**
      * Reads {@code length} bytes from the input stream. It behaves in the same
-     * way as {@link java.io.DataInput#readFully(byte[])}.
+     * way as {@link java.io.DataInput#readFully(byte[])}, and it will throw
+     * {@link IOException} when the input stream has been closed.
      *
      * @param length number of bytes to read
      * @return byte array and its length should be {@code length}
      * @throws IOException when failed to read value from input stream, not able to
      *                     retrieve all bytes, or reached end of the stream
      */
     public byte[] readBytes(int length) throws IOException {
+        if (length <= 0) {
+            return EMPTY_BYTES;
+        }
+
         byte[] bytes = new byte[length];
 
         for (int l = length, c = 0, n = 0; l > 0; l -= n) {

clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseThreadFactory.java

@@ -0,0 +1,49 @@
+package com.clickhouse.client;
+
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class ClickHouseThreadFactory implements ThreadFactory {
+    private final boolean daemon;
+    private final int priority;
+
+    private final ThreadGroup group;
+    private final String namePrefix;
+    private final AtomicInteger threadNumber;
+
+    public ClickHouseThreadFactory(Object owner) {
+        this(owner, false, Thread.NORM_PRIORITY);
+    }
+
+    public ClickHouseThreadFactory(Object owner, boolean daemon, int priority) {
+        String prefix = null;
+        if (owner instanceof String) {
+            prefix = ((String) owner).trim();
+        } else if (owner != null) {
+            prefix = new StringBuilder().append(owner.getClass().getSimpleName()).append('@').append(owner.hashCode())
+                    .toString();
+        }
+        this.daemon = daemon;
+        this.priority = ClickHouseChecker.between(priority, "Priority", Thread.MIN_PRIORITY, Thread.MAX_PRIORITY);
+
+        SecurityManager s = System.getSecurityManager();
+        group = s != null ? s.getThreadGroup() : Thread.currentThread().getThreadGroup();
+        namePrefix = !ClickHouseChecker.isNullOrBlank(prefix) ? prefix
+                : new StringBuilder().append(getClass().getSimpleName()).append('@').append(hashCode())
+                        .append('-').toString();
+        threadNumber = new AtomicInteger(1);
+    }
+
+    @Override
+    public Thread newThread(Runnable r) {
+        Thread t = new Thread(group, r, namePrefix + threadNumber.getAndIncrement(), 0);
+        if (daemon != t.isDaemon()) {
+            t.setDaemon(daemon);
+        }
+        if (priority != t.getPriority()) {
+            t.setPriority(priority);
+        }
+        // t.setUncaughtExceptionHandler(null);
+        return t;
+    }
+}