Proper streaming ByteBody implementation

servlet-engine/src/main/java/io/micronaut/servlet/engine/body/AbstractServletByteBody.java → servlet-core/src/main/java/io/micronaut/servlet/http/body/AbstractServletByteBody.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package io.micronaut.servlet.engine.body;
+package io.micronaut.servlet.http.body;
 
 import io.micronaut.core.annotation.Internal;
 import io.micronaut.http.body.CloseableByteBody;
@@ -25,7 +25,7 @@
  * @since 4.9.0
  */
 @Internal
-public abstract class AbstractServletByteBody implements CloseableByteBody {
+abstract class AbstractServletByteBody implements CloseableByteBody {
     static void failClaim() {
         throw new IllegalStateException("Request body has already been claimed: Two conflicting sites are trying to access the request body. If this is intentional, the first user must ByteBody#split the body. To find out where the body was claimed, turn on TRACE logging for io.micronaut.http.server.netty.body.NettyByteBody.");
     }

servlet-engine/src/main/java/io/micronaut/servlet/engine/body/AvailableByteArrayBody.java → servlet-core/src/main/java/io/micronaut/servlet/http/body/AvailableByteArrayBody.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package io.micronaut.servlet.engine.body;
+package io.micronaut.servlet.http.body;
 
 import io.micronaut.core.annotation.Internal;
 import io.micronaut.core.annotation.NonNull;
@@ -27,6 +27,8 @@
 
 /**
  * {@link io.micronaut.http.body.AvailableByteBody} implementation based on a byte array.
+ * <p>
+ * Note: While internal, this is also used from the AWS and GCP modules.
  *
  * @author Jonas Konrad
  * @since 4.9.0

servlet-core/src/main/java/io/micronaut/servlet/http/body/ByteQueue.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2017-2024 original authors
+ *
+ * Licensed 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
+ *
+ * https://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 io.micronaut.servlet.http.body;
+
+import io.micronaut.core.annotation.Internal;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayDeque;
+import java.util.Arrays;
+import java.util.Queue;
+
+/**
+ * Non-thread-safe queue for bytes.
+ *
+ * @author Jonas Konrad
+ * @since 4.9.0
+ */
+@Internal
+final class ByteQueue {
+    // not the most efficient implementation, but the most readable.
+
+    private final Queue<ByteBuffer> queue = new ArrayDeque<>();
+
+    /**
+     * Add a copy of the given array to this queue.
+     *
+     * @param arr The input array
+     * @param off The offset of the section to add
+     * @param len The length of the section to add
+     */
+    public void addCopy(byte[] arr, int off, int len) {
+        add(Arrays.copyOfRange(arr, off, off + len));
+    }
+
+    private void add(byte[] arr) {
+        if (arr.length == 0) {
+            return;
+        }
+        queue.add(ByteBuffer.wrap(arr));
+    }
+
+    public boolean isEmpty() {
+        return queue.isEmpty();
+    }
+
+    public int take(byte[] arr, int off, int len) {
+        ByteBuffer peek = queue.peek();
+        if (peek == null) {
+            throw new IllegalStateException("Queue is empty");
+        }
+        int n = Math.min(len, peek.remaining());
+        peek.get(arr, off, n);
+        if (peek.remaining() == 0) {
+            queue.poll();
+        }
+        return n;
+    }
+
+    public void clear() {
+        queue.clear();
+    }
+}

servlet-core/src/main/java/io/micronaut/servlet/http/body/ExtendedInputStream.java

@@ -0,0 +1,142 @@
+/*
+ * Copyright 2017-2024 original authors
+ *
+ * Licensed 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
+ *
+ * https://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 io.micronaut.servlet.http.body;
+
+import io.micronaut.core.annotation.Internal;
+import io.micronaut.core.annotation.Nullable;
+import io.micronaut.http.body.ByteBody;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Arrays;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * Extended InputStream API for better backpressure/cancellation handling.
+ *
+ * @author Jonas Konrad
+ * @since 4.9.0
+ */
+@Internal
+abstract class ExtendedInputStream extends InputStream {
+    private static final int CHUNK_SIZE = 8192;
+    private static final Logger LOG = LoggerFactory.getLogger(ExtendedInputStream.class);
+
+    static ExtendedInputStream wrap(InputStream inputStream) {
+        return new Wrapper(inputStream);
+    }
+
+    @Override
+    public int read() throws IOException {
+        byte[] arr1 = new byte[1];
+        int n = read(arr1);
+        if (n == -1) {
+            return -1;
+        } else if (n == 0) {
+            throw new IllegalStateException("Read 0 bytes");
+        } else {
+            return arr1[0] & 0xff;
+        }
+    }
+
+    @Override
+    public abstract int read(byte[] b, int off, int len) throws IOException;
+
+    /**
+     * Read some data into a new byte array. The array may be of any size. This is usually the same
+     * as allocating a new array, calling {@link #read(byte[])}, and then truncating the array, but
+     * may be optimized in some implementations.
+     */
+    @Nullable
+    public byte[] readSome() throws IOException {
+        byte[] arr = new byte[CHUNK_SIZE];
+        int n = read(arr);
+        if (n == -1) {
+            return null;
+        } else if (n == arr.length) {
+            return arr;
+        } else {
+            return Arrays.copyOf(arr, n);
+        }
+    }
+
+    @Override
+    public void close() {
+        allowDiscard();
+        cancelInput();
+    }
+
+    /**
+     * Allow discarding the input of this stream. See {@link ByteBody#allowDiscard()}.
+     */
+    public abstract void allowDiscard();
+
+    /**
+     * Cancel any further upstream input. This also removes any backpressure that this stream
+     * may apply on its upstream.
+     */
+    public abstract void cancelInput();
+
+    private static final class Wrapper extends ExtendedInputStream {
+        private final Lock lock = new ReentrantLock();
+        private final InputStream delegate;
+        private boolean discarded;
+
+        Wrapper(InputStream delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public int read(byte[] b, int off, int len) throws IOException {
+            lock.lock();
+            try {
+                if (discarded) {
+                    throw ByteBody.BodyDiscardedException.create();
+                }
+                return delegate.read(b, off, len);
+            } finally {
+                lock.unlock();
+            }
+        }
+
+        @Override
+        public void close() {
+            try {
+                delegate.close();
+            } catch (IOException e) {
+                LOG.debug("Failed to close request stream", e);
+            }
+        }
+
+        @Override
+        public void allowDiscard() {
+            lock.lock();
+            try {
+                discarded = true;
+                close();
+            } finally {
+                lock.unlock();
+            }
+        }
+
+        @Override
+        public void cancelInput() {
+        }
+    }
+}