Make CompositeInputStream safe when close() races with an in-progress read

dubbo-remoting/dubbo-remoting-http12/src/main/java/org/apache/dubbo/remoting/http12/CompositeInputStream.java

@@ -31,7 +31,9 @@ public class CompositeInputStream extends InputStream {
 
     private int readIndex = 0;
 
-    public void addInputStream(InputStream inputStream) {
+    private boolean closed = false;
+
+    public synchronized void addInputStream(InputStream inputStream) {
         inputStreams.offer(inputStream);
         try {
             totalAvailable += inputStream.available();
@@ -41,15 +43,23 @@ public void addInputStream(InputStream inputStream) {
     }
 
     @Override
-    public int read() throws IOException {
+    public synchronized int read() throws IOException {
         InputStream inputStream;
         while ((inputStream = inputStreams.peek()) != null) {
-            int available = inputStream.available();
-            if (available == 0) {
-                releaseHeadStream();
-                continue;
+            int read;
+            try {
+                if (inputStream.available() == 0) {
+                    releaseHeadStream();
+                    continue;
+                }
+                read = inputStream.read();
+            } catch (IOException e) {
+                if (closed) {
+                    // close() raced with this read: report end of stream instead of the raw exception
+                    return -1;
+                }
+                throw e;
             }
-            int read = inputStream.read();
             if (read != -1) {
                 ++readIndex;
                 releaseIfNecessary(inputStream);
@@ -61,7 +71,7 @@ public int read() throws IOException {
     }
 
     @Override
-    public int read(byte[] b, int off, int len) throws IOException {
+    public synchronized int read(byte[] b, int off, int len) throws IOException {
         if (b == null) {
             throw new NullPointerException();
         } else if (off < 0 || len < 0 || len > b.length - off) {
@@ -73,13 +83,21 @@ public int read(byte[] b, int off, int len) throws IOException {
         int total = 0;
         InputStream inputStream;
         while ((inputStream = inputStreams.peek()) != null) {
-            int available = inputStream.available();
-            if (available == 0) {
-                releaseHeadStream();
-                continue;
+            int read;
+            try {
+                int available = inputStream.available();
+                if (available == 0) {
+                    releaseHeadStream();
+                    continue;
+                }
+                read = inputStream.read(b, off + total, Math.min(len - total, available));
+            } catch (IOException e) {
+                if (closed) {
+                    // close() raced with this read: return what was already read, or end of stream
+                    break;
+                }
+                throw e;
             }
-
-            int read = inputStream.read(b, off + total, Math.min(len - total, available));
             if (read != -1) {
                 total += read;
                 readIndex += read;
@@ -97,12 +115,13 @@ public int read(byte[] b, int off, int len) throws IOException {
     }
 
     @Override
-    public int available() {
-        return totalAvailable - readIndex;
+    public synchronized int available() {
+        return closed ? 0 : totalAvailable - readIndex;
     }
 
     @Override
-    public void close() throws IOException {
+    public synchronized void close() throws IOException {
+        closed = true;
         IOException firstException = null;
         InputStream inputStream;
         while ((inputStream = inputStreams.poll()) != null) {
@@ -133,7 +152,15 @@ private void releaseHeadStream() {
     }
 
     private void releaseIfNecessary(InputStream inputStream) throws IOException {
-        int available = inputStream.available();
+        int available;
+        try {
+            available = inputStream.available();
+        } catch (IOException e) {
+            if (closed) {
+                return;
+            }
+            throw e;
+        }
         if (available == 0) {
             releaseHeadStream();
         }

dubbo-remoting/dubbo-remoting-http12/src/test/java/org/apache/dubbo/remoting/http12/CompositeInputStreamConcurrencyTest.java

@@ -0,0 +1,172 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.dubbo.remoting.http12;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+/**
+ * close() racing with an in-progress read. Most cases are simulated deterministically:
+ * the underlying stream closes the composite from inside its own read(), which is the
+ * same interleaving as a close() arriving on another thread mid-read.
+ */
+class CompositeInputStreamConcurrencyTest {
+
+    @Test
+    void readInterruptedByCloseReportsEndOfStream() throws IOException {
+        CompositeInputStream in = new CompositeInputStream();
+        in.addInputStream(new CloseDuringReadStream(in, new byte[8], true));
+
+        Assertions.assertEquals(-1, in.read(new byte[8], 0, 8));
+    }
+
+    @Test
+    void singleByteReadInterruptedByCloseReportsEndOfStream() throws IOException {
+        CompositeInputStream in = new CompositeInputStream();
+        in.addInputStream(new CloseDuringReadStream(in, new byte[8], true));
+
+        Assertions.assertEquals(-1, in.read());
+    }
+
+    @Test
+    void alreadyBufferedDataIsStillDeliveredWhenCloseHappensDuringRead() throws IOException {
+        byte[] expected = {1, 2, 3, 4, 5, 6, 7, 8};
+        CompositeInputStream in = new CompositeInputStream();
+        in.addInputStream(new CloseDuringReadStream(in, expected, false));
+
+        byte[] buf = new byte[8];
+        Assertions.assertEquals(8, in.read(buf, 0, 8));
+        Assertions.assertArrayEquals(expected, buf);
+        Assertions.assertEquals(-1, in.read(buf, 0, 8));
+    }
+
+    @Test
+    void availableReturnsZeroWhileCloseIsInProgress() throws IOException {
+        CompositeInputStream in = new CompositeInputStream();
+        AtomicInteger seenDuringClose = new AtomicInteger(-1);
+        in.addInputStream(new ByteArrayInputStream(new byte[10]) {
+            @Override
+            public void close() {
+                seenDuringClose.set(in.available());
+            }
+        });
+
+        in.close();
+
+        Assertions.assertEquals(0, seenDuringClose.get());
+    }
+
+    @Test
+    void closeWaitsForReadInProgress() throws Exception {
+        CountDownLatch insideRead = new CountDownLatch(1);
+        CountDownLatch finishRead = new CountDownLatch(1);
+        CompositeInputStream in = new CompositeInputStream();
+        in.addInputStream(new InputStream() {
+            @Override
+            public int available() {
+                return 1;
+            }
+
+            @Override
+            public int read() {
+                insideRead.countDown();
+                try {
+                    finishRead.await(5, TimeUnit.SECONDS);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                }
+                return 42;
+            }
+        });
+
+        AtomicInteger readResult = new AtomicInteger(-2);
+        Thread reader = new Thread(() -> {
+            try {
+                readResult.set(in.read());
+            } catch (IOException e) {
+                readResult.set(-3);
+            }
+        });
+        Thread closer = new Thread(() -> {
+            try {
+                in.close();
+            } catch (IOException ignored) {
+                // not under test
+            }
+        });
+
+        reader.start();
+        Assertions.assertTrue(insideRead.await(5, TimeUnit.SECONDS));
+        closer.start();
+        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+        while (closer.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) {
+            Thread.sleep(10);
+        }
+        // close() must be waiting for the in-flight read, not tearing the stream down under it
+        Assertions.assertEquals(Thread.State.BLOCKED, closer.getState());
+
+        finishRead.countDown();
+        reader.join(5000);
+        closer.join(5000);
+        Assertions.assertEquals(42, readResult.get());
+    }
+
+    /** Closes the composite from inside read(), then either fails or delivers its data. */
+    private static final class CloseDuringReadStream extends InputStream {
+        private final CompositeInputStream owner;
+        private final byte[] data;
+        private final boolean failAfterClose;
+        private int pos;
+
+        CloseDuringReadStream(CompositeInputStream owner, byte[] data, boolean failAfterClose) {
+            this.owner = owner;
+            this.data = data;
+            this.failAfterClose = failAfterClose;
+        }
+
+        @Override
+        public int available() {
+            return data.length - pos;
+        }
+
+        @Override
+        public int read() throws IOException {
+            byte[] one = new byte[1];
+            int n = read(one, 0, 1);
+            return n == -1 ? -1 : (one[0] & 0xFF);
+        }
+
+        @Override
+        public int read(byte[] b, int off, int len) throws IOException {
+            owner.close();
+            if (failAfterClose) {
+                throw new IOException("Stream already closed");
+            }
+            int n = Math.min(len, data.length - pos);
+            System.arraycopy(data, pos, b, off, n);
+            pos += n;
+            return n;
+        }
+    }
+}

dubbo-test/dubbo-test-check/src/main/java/org/apache/dubbo/test/check/registrycenter/processor/StartZookeeperWindowsProcessor.java

@@ -22,6 +22,8 @@
 import org.apache.dubbo.test.check.registrycenter.Processor;
 import org.apache.dubbo.test.check.registrycenter.context.ZookeeperWindowsContext;
 
+import java.io.IOException;
+import java.net.Socket;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.concurrent.TimeUnit;
@@ -37,6 +39,16 @@ public class StartZookeeperWindowsProcessor extends ZookeeperWindowsProcessor {
 
     private static final Logger logger = LoggerFactory.getLogger(StartZookeeperWindowsProcessor.class);
 
+    /**
+     * Maximum time to wait for each zookeeper instance to accept connections.
+     */
+    private static final long READY_TIMEOUT_MILLIS = 30_000;
+
+    /**
+     * Delay between successive readiness checks.
+     */
+    private static final long POLL_INTERVAL_MILLIS = 200;
+
     /**
      * The {@link Processor} to find the pid of zookeeper instance.
      */
@@ -70,15 +82,53 @@ protected void doProcess(ZookeeperWindowsContext context) throws DubboTestExcept
                     .toString());
             context.getExecutorService().submit(() -> executor.execute(cmdLine));
         }
-        try {
-            // TODO: Help me to optimize the ugly sleep.
-            // sleep to wait all of zookeeper instances are started successfully.
-            // The best way is to check the output log with the specified keywords,
-            // however, there maybe keep waiting for check when any exception occurred,
-            // because the output stream will be blocked to wait for continuous data without any break
-            TimeUnit.SECONDS.sleep(3);
-        } catch (InterruptedException e) {
-            // ignored
+        // Actively wait for each zookeeper instance to start accepting connections,
+        // instead of blindly sleeping for a fixed duration. This fails fast when a
+        // port never comes up, and doesn't waste time once a port is already ready.
+        waitForZookeeperReady(context.getClientPorts());
+    }
+
+    /**
+     * Blocks until every given zookeeper client port is accepting connections, or throws
+     * if any of them fails to become ready within {@link #READY_TIMEOUT_MILLIS}.
+     */
+    private void waitForZookeeperReady(int[] clientPorts) throws DubboTestException {
+        for (int clientPort : clientPorts) {
+            long deadline = System.currentTimeMillis() + READY_TIMEOUT_MILLIS;
+            boolean ready = false;
+
+            while (System.currentTimeMillis() < deadline) {
+                if (isPortOpen(clientPort)) {
+                    ready = true;
+                    break;
+                }
+
+                try {
+                    TimeUnit.MILLISECONDS.sleep(POLL_INTERVAL_MILLIS);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new DubboTestException("Interrupted while waiting for zookeeper to start", e);
+                }
+            }
+
+            if (!ready) {
+                throw new DubboTestException(String.format(
+                        "Zookeeper on port %d did not become ready within %d milliseconds",
+                        clientPort, READY_TIMEOUT_MILLIS));
+            }
+            logger.info(String.format("The zookeeper-%d is ready.", clientPort));
+        }
+    }
+
+    /**
+     * Returns true if a TCP connection to 127.0.0.1:port can be opened, meaning
+     * something (expected to be zookeeper) is already listening there.
+     */
+    private boolean isPortOpen(int port) {
+        try (Socket socket = new Socket("127.0.0.1", port)) {
+            return true;
+        } catch (IOException e) {
+            return false;
         }
     }
 }