Add fixed capacity policy for testing

Author: zhicwu

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

@@ -162,7 +162,7 @@ public ClickHousePipedOutputStream createPipedOutputStream(ClickHouseConfig conf
         }
         return blocking
                 ? new BlockingPipedOutputStream(bufferSize, queue, timeout, postCloseAction)
-                : new NonBlockingPipedOutputStream(bufferSize, queue, timeout, postCloseAction);
+                : new NonBlockingPipedOutputStream(bufferSize, queue, timeout, null, postCloseAction);
     }
 
     public ClickHousePipedOutputStream createPipedOutputStream(int writeBufferSize, int queueSize, int timeout,

clickhouse-client/src/main/java/com/clickhouse/client/stream/CapacityPolicy.java

@@ -2,6 +2,19 @@
 
 @FunctionalInterface
 public interface CapacityPolicy {
+    static class FixedCapacity implements CapacityPolicy {
+        private final int capacity;
+
+        protected FixedCapacity(int capacity) {
+            this.capacity = capacity < 1 ? 0 : capacity;
+        }
+
+        @Override
+        public boolean ensureCapacity(int current) {
+            return capacity < 1 || current < capacity;
+        }
+    }
+
     static class LinearDynamicCapacity implements CapacityPolicy {
         private volatile int capacity;
         private volatile int count;
@@ -31,6 +44,10 @@ public boolean ensureCapacity(int current) {
         }
     }
 
+    static CapacityPolicy fixedCapacity(int capacity) {
+        return new FixedCapacity(capacity);
+    }
+
     static CapacityPolicy linearDynamicCapacity(int initialSize, int maxSize, int threshold) {
         return new LinearDynamicCapacity(initialSize, maxSize, threshold);
     }

clickhouse-client/src/main/java/com/clickhouse/client/stream/NonBlockingPipedOutputStream.java

@@ -80,10 +80,12 @@ private void updateBuffer(byte[] bytes, int offset, int length) throws IOExcepti
         }
     }
 
-    public NonBlockingPipedOutputStream(int bufferSize, int queueLength, int timeout, Runnable postCloseAction) {
+    public NonBlockingPipedOutputStream(int bufferSize, int queueLength, int timeout, CapacityPolicy policy,
+            Runnable postCloseAction) {
         super(postCloseAction);
 
-        this.queue = new AdaptiveQueue<>(CapacityPolicy.linearDynamicCapacity(1, queueLength, 0));
+        this.queue = new AdaptiveQueue<>(
+                policy != null ? policy : CapacityPolicy.linearDynamicCapacity(1, queueLength, 0));
 
         // may need an initialBufferSize and a monitor to update bufferSize in runtime
         this.bufferSize = ClickHouseUtils.getBufferSize(bufferSize,

clickhouse-client/src/test/java/com/clickhouse/client/ClickHouseDataStreamFactoryTest.java

@@ -18,9 +18,8 @@ public void testCreatePipedOutputStream() throws Exception {
         // read in worker thread
         for (int i = 0; i < 256; i++) {
             CompletableFuture<Integer> future;
-            try (ClickHousePipedOutputStream out = ClickHouseDataStreamFactory.getInstance().createPipedOutputStream(
-                    config,
-                    null)) {
+            try (ClickHousePipedOutputStream out = ClickHouseDataStreamFactory.getInstance()
+                    .createPipedOutputStream(config, null)) {
                 future = ClickHouseClient.submit(() -> {
                     try (ClickHouseInputStream in = out.getInputStream()) {
                         return in.read();
@@ -33,8 +32,7 @@ public void testCreatePipedOutputStream() throws Exception {
 
         // write in worker thread
         for (int i = 0; i < 256; i++) {
-            ClickHousePipedOutputStream out = ClickHouseDataStreamFactory.getInstance().createPipedOutputStream(
-                    config,
+            ClickHousePipedOutputStream out = ClickHouseDataStreamFactory.getInstance().createPipedOutputStream(config,
                     null);
             final int num = i;
             try (ClickHouseInputStream in = out.getInputStream()) {

clickhouse-client/src/test/java/com/clickhouse/client/stream/NonBlockingPipedOutputStreamTest.java

@@ -18,7 +18,8 @@
 public class NonBlockingPipedOutputStreamTest {
     @Test(groups = { "unit" })
     public void testRead() throws Exception {
-        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(4, 3, 1, null);
+        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(4, 3, 1, CapacityPolicy.fixedCapacity(3),
+                null);
         Assert.assertEquals(stream.queue.size(), 0);
         try (InputStream in = stream.getInputStream()) {
             in.read();
@@ -76,7 +77,8 @@ public void testRead() throws Exception {
 
     @Test(groups = { "unit" })
     public void testReadBytes() throws Exception {
-        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(4, 3, 1, null);
+        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(4, 3, 1, CapacityPolicy.fixedCapacity(3),
+                null);
         Assert.assertEquals(stream.queue.size(), 0);
         byte[] bytes = new byte[3];
         try (InputStream in = stream.getInputStream()) {
@@ -138,7 +140,8 @@ public void testReadBytes() throws Exception {
 
     @Test(groups = { "unit" })
     public void testWrite() throws Exception {
-        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(2, 3, 2, null);
+        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(2, 3, 2, CapacityPolicy.fixedCapacity(3),
+                null);
         Assert.assertEquals(stream.queue.size(), 0);
         try (OutputStream out = stream) {
             out.write(5);
@@ -153,7 +156,7 @@ public void testWrite() throws Exception {
             Assert.assertEquals(stream.queue.poll(), new byte[] { (byte) 7 });
         }
 
-        stream = new NonBlockingPipedOutputStream(1, 1, 2, null);
+        stream = new NonBlockingPipedOutputStream(1, 1, 2, CapacityPolicy.fixedCapacity(1), null);
         Assert.assertEquals(stream.queue.size(), 0);
         try (OutputStream out = stream) {
             out.write(5);
@@ -178,7 +181,8 @@ public void testWrite() throws Exception {
 
     @Test(groups = { "unit" })
     public void testWriteBytes() throws Exception {
-        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(2, 3, 2, null);
+        NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(2, 3, 2, CapacityPolicy.fixedCapacity(3),
+                null);
         Assert.assertEquals(stream.queue.size(), 0);
         try (OutputStream out = stream) {
             out.write(new byte[] { (byte) 9, (byte) 10 });
@@ -204,12 +208,12 @@ public void testWriteBytes() throws Exception {
 
     @Test(groups = { "unit" })
     public void testPipedStream() throws Exception {
-        final int timeout = 60000;
+        final int timeout = 10000;
         ExecutorService executor = Executors.newFixedThreadPool(2);
         for (int bufferSize = -1; bufferSize < 10; bufferSize++) {
             for (int queueLength = -1; queueLength < 10; queueLength++) {
                 final NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(bufferSize, queueLength,
-                        timeout, null);
+                        timeout, CapacityPolicy.fixedCapacity(queueLength), null);
                 try (InputStream in = stream.getInputStream(); OutputStream out = stream) {
                     int count = 10000;
                     CountDownLatch latch = new CountDownLatch(count + 1);