NonBlockingPipedOutputStream.close() is not idempotent: concurrent close silently drops the pending buffer

[polyglotAI-bot]

Description

com.clickhouse.data.stream.NonBlockingPipedOutputStream.close() is not thread-safe and not idempotent. When two threads close the same stream at the same time, both can pass the if (closed) guard. Two things then happen:

1. postCloseAction runs twice.
2. Both threads flush the same pending ByteBuffer object, and the second one sets its limit to 0 before the reader polls it. The reader then gets an immediate EOF, so the data still pending in the buffer is lost silently - no exception, no short-read signal, just fewer bytes.

This is the sibling of #3055 (the BlockingPipedOutputStream case). It is not the same failure: AdaptiveQueue.add() does not block, so there is no Close stream timed out error here. Instead the failure is silent data loss, and it hits the default implementation: ClickHouseDataConfig.DEFAULT_USE_BLOCKING_QUEUE is false, so ClickHouseDataStreamFactory.createPipedOutputStream(...) returns NonBlockingPipedOutputStream unless use_blocking_queue=true is set.

Steps to reproduce

1. Build clickhouse-data on main.
2. Create a NonBlockingPipedOutputStream, write fewer bytes than the buffer size (so the payload is still in the in-memory buffer).
3. Close it from two threads that meet on a barrier, then read everything through getInputStream() and compare the byte count.
4. Repeat 20000 times.

Result on an idle 4-vCPU container (no CPU contention needed):

iterations                 = 20000
postCloseAction ran twice  = 66
reader lost data (short)   = 16     <- reader saw 0 of 16 bytes
reader saw extra data      = 0
reader correct             = 19984
read errors                = 0

Sample lines from the run:

iter 126: reader saw 0 of 16 bytes (postCloseAction ran 2x)
iter 182: reader saw 0 of 16 bytes (postCloseAction ran 2x)
iter 208: reader saw 0 of 16 bytes (postCloseAction ran 2x)

Error Log or Exception StackTrace

(none - this is the point: the loss is silent, the reader gets a clean EOF)

Expected Behaviour

close() must be idempotent, as Closeable#close requires:

If the stream is already closed then invoking this method has no effect.

A second, concurrent close() must return without touching the buffer or the queue. The first close() must still flush the pending buffer completely and enqueue exactly one EMPTY_BUFFER, and postCloseAction must run exactly once.

Root cause

NonBlockingPipedOutputStream.close() (clickhouse-data/src/main/java/com/clickhouse/data/stream/NonBlockingPipedOutputStream.java:106-121):

public void close() throws IOException {
    if (closed) {                 // <-- check
        return;
    }
    try {
        if (buffer.position() > 0) {
            updateBuffer(false);  // flushes the shared buffer, does NOT replace it
        }
    } finally {
        queue.add(ClickHouseByteBuffer.EMPTY_BUFFER);
        closed = true;            // <-- set, only at the very end
        ClickHouseDataStreamFactory.handleCustomAction(postCloseAction);
    }
}

closed is volatile, but check-and-set is not atomic and closed is assigned last, so the whole body is a race window.

The data loss comes from updateBuffer(boolean) (:40-52) being run twice on the same object. Called from close() with allocateNewBuffer == false, it never replaces the buffer field:

ByteBuffer b = buffer;
if (b.hasRemaining()) {
    ((Buffer) b).limit(b.position());
}
((Buffer) b).rewind();
updateBuffer(b);            // offers b to the queue

Thread A: limit = position (= n), rewind() sets position = 0, offers b.
Thread B: b.hasRemaining() is now true (position 0 < limit n), so it sets limit = position = 0, rewinds, and offers the same b again.

The queue holds the buffer object itself, so B's limit(0) mutates the buffer A already handed over. NonBlockingInputStream.updateBuffer() polls it, sees remaining() == 0, and AbstractByteBufferInputStream treats that as end of stream. The pending payload is gone with no error.

The concurrent-close shape is the same one recorded in #3055: a writer thread closes the stream when it finishes (ClickHousePipedOutputStream.writeAsync closes it in try-with-resources) while the calling thread also closes it through try-with-resources.

Suggested fix

Same shape as the fix proposed for #3055: make the close path atomic and mark the stream closed before doing any work.

• Make close() synchronized with the closed re-check inside the lock, or move closed to an AtomicBoolean and enter the body only if compareAndSet(false, true) wins.
• Verified locally: adding synchronized to close() and changing nothing else takes the same 20000-iteration harness to postCloseAction ran twice = 0, reader lost data = 0, reader correct = 20000.
• Contrast cases that must keep their current behavior: the first close() still flushes the pending buffer and enqueues exactly one EMPTY_BUFFER so the reader still sees EOF; flush()/write() after close must still fail through ensureOpen(); and updateBuffer(true) (from flush()) must keep allocating a fresh buffer.

Because both piped streams share this pattern, fixing them together may be preferable. Both classes are @Deprecated, so a maintainer may want to weigh that - but NonBlockingPipedOutputStream is still the default on the v1 code path, and the failure mode here is lost data rather than a test timeout.

Code Example

final AtomicInteger postClose = new AtomicInteger();
final NonBlockingPipedOutputStream out = new NonBlockingPipedOutputStream(
        1024, 16, 2000, CapacityPolicy.fixedCapacity(16), postClose::incrementAndGet);
out.write(payload);                       // 16 bytes, still in the in-memory buffer

final CyclicBarrier barrier = new CyclicBarrier(2);
Runnable closer = () -> { barrier.await(); out.close(); };   // exceptions elided
Thread t1 = new Thread(closer), t2 = new Thread(closer);
t1.start(); t2.start(); t1.join(); t2.join();

int n = 0;
try (ClickHouseInputStream in = out.getInputStream()) {
    byte[] tmp = new byte[64];
    int r;
    while ((r = in.read(tmp)) > 0) { n += r; }
}
// occasionally: n == 0 while payload.length == 16, and postClose.get() == 2

Configuration

Client Configuration

// defaults; use_blocking_queue is false by default, which selects NonBlockingPipedOutputStream

Environment

• Cloud
• Client version: main @ 601ade16
• Language version: OpenJDK 17.0.19 (Temurin)
• OS: Ubuntu 24.04 container, 4 vCPU

ClickHouse Server

• ClickHouse Server version: not involved - clickhouse-data unit-level reproduction, no server interaction
• ClickHouse Server non-default settings, if any: n/a
• CREATE TABLE statements for tables involved: n/a
• Sample data for all these tables: n/a

---

Found by automated analysis of the client while working on #3055, which noted this sibling class was worth checking. Verified by running the code, not by inspection: the numbers above come from a 20000-iteration harness on main, and from the same harness after the suggested one-word change.