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

CHANGELOG.md

@@ -240,6 +240,11 @@
   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)
+- **[data]** Fixed `NonBlockingPipedOutputStream.close()` not being idempotent under concurrency. Two threads could
+  both flush and mutate the same pending buffer before the reader consumed it, silently replacing the payload with
+  an empty buffer and running the post-close action twice. Exactly one caller now flushes the pending data, enqueues
+  the end-of-stream marker, and runs the post-close action; concurrent or repeated `close()` calls return immediately.
+  (https://github.com/ClickHouse/clickhouse-java/issues/3057)
 - **[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/NonBlockingPipedOutputStream.java

@@ -4,6 +4,7 @@
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.clickhouse.data.ClickHouseByteBuffer;
 import com.clickhouse.data.ClickHouseChecker;
@@ -34,6 +35,7 @@ public class NonBlockingPipedOutputStream extends ClickHousePipedOutputStream {
     protected final int bufferSize;
     protected final CompletableFuture<Void> future;
     protected final long timeout;
+    private final AtomicBoolean closing = new AtomicBoolean(false);
 
     protected ByteBuffer buffer;
 
@@ -104,7 +106,7 @@ public ClickHouseInputStream getInputStream(Runnable postCloseAction) {
 
     @Override
     public void close() throws IOException {
-        if (closed) {
+        if (closed || !closing.compareAndSet(false, true)) {
             return;
         }
 

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

@@ -9,7 +9,10 @@
 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;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import com.clickhouse.data.ClickHouseByteBuffer;
@@ -208,6 +211,61 @@ public void testWriteBytes() throws IOException {
         }
     }
 
+    @Test(groups = { "unit" })
+    public void testConcurrentClose() throws Exception {
+        final long timeout = 5000L;
+        final AtomicBoolean acceptBuffer = new AtomicBoolean(false);
+        final CountDownLatch firstBufferOffer = new CountDownLatch(1);
+        final AtomicInteger closeCount = new AtomicInteger(0);
+        final CapacityPolicy policy = current -> {
+            firstBufferOffer.countDown();
+            return acceptBuffer.get();
+        };
+        final NonBlockingPipedOutputStream stream = new NonBlockingPipedOutputStream(4, 2, timeout * 4L,
+                policy, (Runnable) closeCount::incrementAndGet);
+        final byte[] expected = new byte[] { (byte) 1, (byte) 2, (byte) 3 };
+        stream.write(expected);
+
+        final ExecutorService executor = Executors.newFixedThreadPool(2);
+        try {
+            Future<?> firstClose = executor.submit(() -> {
+                stream.close();
+                return null;
+            });
+            Assert.assertTrue(firstBufferOffer.await(timeout, TimeUnit.MILLISECONDS),
+                    "First close did not try to flush the pending buffer");
+
+            Future<?> concurrentClose = executor.submit(() -> {
+                stream.close();
+                return null;
+            });
+            boolean concurrentCloseReturned = true;
+            try {
+                concurrentClose.get(timeout, TimeUnit.MILLISECONDS);
+            } catch (TimeoutException e) {
+                concurrentCloseReturned = false;
+            }
+
+            acceptBuffer.set(true);
+            firstClose.get(timeout, TimeUnit.MILLISECONDS);
+            concurrentClose.get(timeout, TimeUnit.MILLISECONDS);
+
+            Assert.assertTrue(concurrentCloseReturned, "Concurrent close should return without waiting");
+            Assert.assertEquals(closeCount.get(), 1, "Post close action should have been executed exactly once");
+            try (InputStream in = stream.getInputStream()) {
+                byte[] actual = new byte[expected.length];
+                Assert.assertEquals(in.read(actual), expected.length);
+                Assert.assertEquals(actual, expected);
+                Assert.assertEquals(in.read(), -1);
+            }
+        } finally {
+            acceptBuffer.set(true);
+            executor.shutdownNow();
+            Assert.assertTrue(executor.awaitTermination(timeout, TimeUnit.MILLISECONDS),
+                    "Concurrent close executor did not terminate");
+        }
+    }
+
     @Test(groups = { "unit" })
     public void testPipedStream() throws InterruptedException, IOException {
         final int timeout = 10000;