Issue #5499 - Reduce buffer allocations and copying from ByteAccumulator

jetty-io/src/main/java/org/eclipse/jetty/io/ByteBufferAccumulator.java

@@ -0,0 +1,115 @@
+//
+//  ========================================================================
+//  Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
+//  ------------------------------------------------------------------------
+//  All rights reserved. This program and the accompanying materials
+//  are made available under the terms of the Eclipse Public License v1.0
+//  and Apache License v2.0 which accompanies this distribution.
+//
+//      The Eclipse Public License is available at
+//      http://www.eclipse.org/legal/epl-v10.html
+//
+//      The Apache License v2.0 is available at
+//      http://www.opensource.org/licenses/apache2.0.php
+//
+//  You may elect to redistribute this code under either of these licenses.
+//  ========================================================================
+//
+
+package org.eclipse.jetty.io;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.eclipse.jetty.util.BufferUtil;
+
+public class ByteBufferAccumulator implements AutoCloseable
+{
+    private static final int MIN_SPACE = 8;
+    private static final int DEFAULT_BUFFER_SIZE = 1024;
+
+    private final List<ByteBuffer> _buffers = new ArrayList<>();
+    private final ByteBufferPool _bufferPool;
+
+    public ByteBufferAccumulator()
+    {
+        this(null);
+    }
+
+    public ByteBufferAccumulator(ByteBufferPool bufferPool)
+    {
+        _bufferPool = (bufferPool == null) ? new NullByteBufferPool() : bufferPool;
+    }
+
+    public int getLength()
+    {
+        int length = 0;
+        for (ByteBuffer buffer : _buffers)
+            length += buffer.remaining();
+        return length;
+    }
+
+    public ByteBuffer ensureBuffer()
+    {
+        return ensureBuffer(DEFAULT_BUFFER_SIZE);
+    }
+
+    public ByteBuffer ensureBuffer(int minAllocationSize)
+    {
+        ByteBuffer buffer = _buffers.isEmpty() ? BufferUtil.EMPTY_BUFFER : _buffers.get(_buffers.size() - 1);
+        if (BufferUtil.space(buffer) <= MIN_SPACE)
+        {
+            buffer = _bufferPool.acquire(Math.max(DEFAULT_BUFFER_SIZE, minAllocationSize), false);
+            _buffers.add(buffer);
+        }
+
+        return buffer;
+    }
+
+    public void copyBytes(byte[] buf, int offset, int length)
+    {
+        copyBuffer(BufferUtil.toBuffer(buf, offset, length));
+    }
+
+    public void copyBuffer(ByteBuffer buffer)
+    {
+        while (buffer.hasRemaining())
+        {
+            ByteBuffer b = ensureBuffer(buffer.remaining());
+            int pos = BufferUtil.flipToFill(b);
+            BufferUtil.put(buffer, b);
+            BufferUtil.flipToFlush(b, pos);
+        }
+    }
+
+    public void writeTo(ByteBuffer buffer)
+    {
+        int pos = BufferUtil.flipToFill(buffer);
+        for (ByteBuffer bb : _buffers)
+        {
+            buffer.put(bb);
+        }
+        BufferUtil.flipToFlush(buffer, pos);
+    }
+
+    public void writeTo(OutputStream out) throws IOException
+    {
+        for (ByteBuffer bb : _buffers)
+        {
+            BufferUtil.writeTo(bb, out);
+        }
+    }
+
+    @Override
+    public void close()
+    {
+        for (ByteBuffer buffer : _buffers)
+        {
+            _bufferPool.release(buffer);
+        }
+        _buffers.clear();
+    }
+}

jetty-io/src/main/java/org/eclipse/jetty/io/ByteBufferOutputStream2.java

@@ -0,0 +1,149 @@
+//
+//  ========================================================================
+//  Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
+//  ------------------------------------------------------------------------
+//  All rights reserved. This program and the accompanying materials
+//  are made available under the terms of the Eclipse Public License v1.0
+//  and Apache License v2.0 which accompanies this distribution.
+//
+//      The Eclipse Public License is available at
+//      http://www.eclipse.org/legal/epl-v10.html
+//
+//      The Apache License v2.0 is available at
+//      http://www.opensource.org/licenses/apache2.0.php
+//
+//  You may elect to redistribute this code under either of these licenses.
+//  ========================================================================
+//
+
+package org.eclipse.jetty.io;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+
+import org.eclipse.jetty.util.BufferUtil;
+
+/**
+ * This class implements an output stream in which the data is written into a list of ByteBuffer,
+ * the buffer list automatically grows as data is written to it, the buffers are taken from the
+ * supplied {@link ByteBufferPool} or freshly allocated if one is not supplied.
+ *
+ * Designed to mimic {@link java.io.ByteArrayOutputStream} but with better memory usage, and less copying.
+ */
+public class ByteBufferOutputStream2 extends OutputStream
+{
+    private final ByteBufferAccumulator _accumulator;
+    private final ByteBufferPool _bufferPool;
+    private ByteBuffer _combinedByteBuffer;
+    private int _size = 0;
+
+    public ByteBufferOutputStream2()
+    {
+        this(null);
+    }
+
+    public ByteBufferOutputStream2(ByteBufferPool bufferPool)
+    {
+        _bufferPool = (bufferPool == null) ? new NullByteBufferPool() : bufferPool;
+        _accumulator = new ByteBufferAccumulator(bufferPool);
+    }
+
+    /**
+     * Get an aggregated content written to the OutputStream in a ByteBuffer.
+     * @return the content in a ByteBuffer.
+     */
+    public ByteBuffer toByteBuffer()
+    {
+        int length = _accumulator.getLength();
+        if (length == 0)
+            return BufferUtil.EMPTY_BUFFER;
+
+        if (_combinedByteBuffer != null && length == _combinedByteBuffer.remaining())
+            return _combinedByteBuffer;
+
+        ByteBuffer buffer = _bufferPool.acquire(_size, false);
+        _accumulator.writeTo(buffer);
+        if (_combinedByteBuffer != null)
+        {
+            _bufferPool.release(_combinedByteBuffer);
+            _combinedByteBuffer = buffer;
+        }
+
+        return buffer;
+    }
+
+    /**
+     * Get an aggregated content written to the OutputStream in a byte array.
+     * @return the content in a byte array.
+     */
+    public byte[] toByteArray()
+    {
+        int length = _accumulator.getLength();
+        if (length == 0)
+            return new byte[0];
+
+        byte[] bytes = new byte[_size];
+        ByteBuffer buffer = BufferUtil.toBuffer(bytes);
+        _accumulator.writeTo(buffer);
+        return bytes;
+    }
+
+    public int size()
+    {
+        return _size;
+    }
+
+    @Override
+    public void write(int b)
+    {
+        write(new byte[]{(byte)b}, 0, 1);
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len)
+    {
+        releaseAggregateBuffer();
+        _accumulator.copyBytes(b, off, len);
+    }
+
+    public void write(ByteBuffer buffer)
+    {
+        releaseAggregateBuffer();
+        _accumulator.copyBuffer(buffer);
+    }
+
+    public void writeTo(ByteBuffer buffer)
+    {
+        _accumulator.writeTo(buffer);
+    }
+
+    public void writeTo(OutputStream out) throws IOException
+    {
+        _accumulator.writeTo(out);
+    }
+
+    private void releaseAggregateBuffer()
+    {
+        if (_combinedByteBuffer != null)
+        {
+            _bufferPool.release(_combinedByteBuffer);
+            _combinedByteBuffer = null;
+        }
+    }
+
+    @Override
+    public void close()
+    {
+        releaseAggregateBuffer();
+        _accumulator.close();
+        _size = 0;
+    }
+
+    @Override
+    public synchronized String toString()
+    {
+        return String.format("%s@%x{size=%d, bufferPool=%s, byteAccumulator=%s}", getClass().getSimpleName(),
+            hashCode(), _size, _bufferPool, _accumulator);
+    }
+}

jetty-io/src/main/java/org/eclipse/jetty/io/NullByteBufferPool.java

@@ -0,0 +1,41 @@
+//
+//  ========================================================================
+//  Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
+//  ------------------------------------------------------------------------
+//  All rights reserved. This program and the accompanying materials
+//  are made available under the terms of the Eclipse Public License v1.0
+//  and Apache License v2.0 which accompanies this distribution.
+//
+//      The Eclipse Public License is available at
+//      http://www.eclipse.org/legal/epl-v10.html
+//
+//      The Apache License v2.0 is available at
+//      http://www.opensource.org/licenses/apache2.0.php
+//
+//  You may elect to redistribute this code under either of these licenses.
+//  ========================================================================
+//
+
+package org.eclipse.jetty.io;
+
+import java.nio.ByteBuffer;
+
+import org.eclipse.jetty.util.BufferUtil;
+
+public class NullByteBufferPool implements ByteBufferPool
+{
+    @Override
+    public ByteBuffer acquire(int size, boolean direct)
+    {
+        if (direct)
+            return BufferUtil.allocateDirect(size);
+        else
+            return BufferUtil.allocate(size);
+    }
+
+    @Override
+    public void release(ByteBuffer buffer)
+    {
+        BufferUtil.clear(buffer);
+    }
+}