Fix clickhouse-data: make BlockingPipedOutputStream.close() idempotent

CHANGELOG.md

@@ -44,6 +44,13 @@
 
 ### Bug Fixes 
 
+- **[data]** Fixed `BlockingPipedOutputStream.close()` not being idempotent under concurrency: the check of the
+  `closed` flag and the closing handshake were not atomic, so two threads closing the same stream (e.g. a writer
+  thread and a try-with-resources block) could both put the end-of-stream marker into the queue, and the second one
+  failed with `Close stream timed out after <n> ms` once the reader had stopped consuming. Exactly one caller now
+  performs the handshake and runs the post-close action; a concurrent or repeated `close()` returns immediately. A
+  `close()` which fails while flushing the remaining data also marks the stream closed and runs the post-close
+  action, so the stream cannot stay half-closed. (https://github.com/ClickHouse/clickhouse-java/issues/3055)
 - **[jdbc-v2]** Fixed JDBC escape processing rewriting text inside string literals and quoted identifiers. Because
   `PreparedStatement` inlines bound parameters into the statement text, a bound value containing `{fn ` (or `{d '...'}`
   / `{ts '...'}`) was re-read as SQL syntax: the `{fn ` was removed together with the next `}` found anywhere in the

clickhouse-data/src/main/java/com/clickhouse/data/stream/BlockingPipedOutputStream.java

@@ -8,6 +8,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.clickhouse.data.ClickHouseByteBuffer;
 import com.clickhouse.data.ClickHouseChecker;
@@ -33,6 +34,7 @@ public class BlockingPipedOutputStream extends ClickHousePipedOutputStream {
     private final int bufferSize;
     private final CompletableFuture<Void> future;
     private final long timeout;
+    private final AtomicBoolean closing = new AtomicBoolean(false);
 
     private ByteBuffer buffer;
 
@@ -105,16 +107,16 @@ public ClickHouseInputStream getInputStream(Runnable postCloseAction) {
 
     @Override
     public void close() throws IOException {
-        if (closed) {
+        if (closed || !closing.compareAndSet(false, true)) {
             return;
         }
 
-        if (buffer.position() > 0) {
-            updateBuffer(false);
-        }
-
         // buffer = ClickHouseByteBuffer.EMPTY_BUFFER;
         try {
+            if (buffer.position() > 0) {
+                updateBuffer(false);
+            }
+
             if (timeout > 0L) {
                 if (!queue.offer(ClickHouseByteBuffer.EMPTY_BUFFER, timeout, TimeUnit.MILLISECONDS)) {
                     throw new IOException(ClickHouseUtils.format("Close stream timed out after %d ms", timeout));

clickhouse-data/src/test/java/com/clickhouse/data/stream/BlockingPipedOutputStreamTest.java

@@ -6,10 +6,16 @@
 import java.io.UncheckedIOException;
 import java.nio.ByteBuffer;
 import java.nio.Buffer;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -209,6 +215,74 @@ public void testWriteBytes() throws InterruptedException, IOException {
         }
     }
 
+    @Test(groups = { "unit" })
+    public void testConcurrentClose() throws Exception {
+        final int closers = 4;
+        final long timeout = 500L;
+        final AtomicInteger closeCount = new AtomicInteger(0);
+        final Collection<String> errors = new ConcurrentLinkedQueue<>();
+        final BlockingPipedOutputStream stream = new BlockingPipedOutputStream(4, 1, timeout,
+                (Runnable) closeCount::incrementAndGet);
+        // fill the only slot of the queue so that the closing handshake cannot complete
+        stream.queue.put(ByteBuffer.allocate(1));
+
+        final CyclicBarrier barrier = new CyclicBarrier(closers);
+        final ExecutorService executor = Executors.newFixedThreadPool(closers);
+        try {
+            List<Future<?>> futures = new ArrayList<>(closers);
+            for (int i = 0; i < closers; i++) {
+                futures.add(executor.submit(() -> {
+                    barrier.await();
+                    try {
+                        stream.close();
+                    } catch (IOException e) {
+                        errors.add(String.valueOf(e.getMessage()));
+                    }
+                    return null;
+                }));
+            }
+            for (Future<?> f : futures) {
+                f.get(timeout + 30000L, TimeUnit.MILLISECONDS);
+            }
+        } finally {
+            executor.shutdownNow();
+        }
+
+        Assert.assertEquals(closeCount.get(), 1, "Stream should have been closed exactly once");
+        Assert.assertEquals(errors.size(), 1, "Only the thread which closed the stream may fail");
+        Assert.assertTrue(errors.iterator().next().indexOf("Close stream timed out") == 0,
+                "Unexpected error: " + errors);
+        Assert.assertEquals(stream.queue.size(), 1, "No additional buffer should have been queued");
+
+        stream.close();
+        Assert.assertEquals(closeCount.get(), 1, "Closing a closed stream should do nothing");
+    }
+
+    @Test(groups = { "unit" })
+    public void testCloseWhenFlushingRemainingBufferFails() throws Exception {
+        final long timeout = 500L;
+        final AtomicInteger closeCount = new AtomicInteger(0);
+        final BlockingPipedOutputStream stream = new BlockingPipedOutputStream(4, 1, timeout,
+                (Runnable) closeCount::incrementAndGet);
+        stream.writeByte((byte) 1);
+        // fill the only slot of the queue so that flushing the remaining buffer fails
+        stream.queue.put(ByteBuffer.allocate(1));
+
+        try {
+            stream.close();
+            Assert.fail("Close should fail");
+        } catch (IOException e) {
+            Assert.assertTrue(e.getMessage().indexOf("Write timed out") == 0, "Unexpected error: " + e.getMessage());
+        }
+
+        Assert.assertTrue(stream.isClosed(), "Stream should have been closed");
+        Assert.assertEquals(closeCount.get(), 1, "Post close action should have been executed exactly once");
+        Assert.assertEquals(stream.queue.size(), 1, "No additional buffer should have been queued");
+
+        stream.close();
+        Assert.assertEquals(closeCount.get(), 1, "Closing a closed stream should do nothing");
+    }
+
     @Test(groups = { "unit" })
     public void testPipedStream() throws InterruptedException, IOException {
         final int timeout = 10000;