Introduce resizable inbound byte buffer

core/src/main/java/org/elasticsearch/common/bytes/ByteBufferReference.java

@@ -0,0 +1,89 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch 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.elasticsearch.common.bytes;
+
+import org.apache.lucene.util.BytesRef;
+
+import java.nio.ByteBuffer;
+
+/**
+ * This is a {@link BytesReference} backed by a {@link ByteBuffer}. The byte buffer can either be a heap or
+ * direct byte buffer. The reference is composed of the space between the {@link ByteBuffer#position} and
+ * {@link ByteBuffer#limit} at construction time. If the position or limit of the underlying byte buffer is
+ * changed, those changes will not be reflected in this reference. However, modifying the limit or position
+ * of the underlying byte buffer is not recommended as those can be used during {@link ByteBuffer#get()}
+ * bounds checks. Use {@link ByteBuffer#duplicate()} at creation time if you plan on modifying the markers of
+ * the underlying byte buffer. Any changes to the underlying data in the byte buffer will be reflected.
+ */
+public class ByteBufferReference extends BytesReference {
+
+    private final ByteBuffer buffer;
+    private final int offset;
+    private final int length;
+
+    public ByteBufferReference(ByteBuffer buffer) {
+        this.buffer = buffer;
+        this.offset = buffer.position();
+        this.length = buffer.remaining();
+    }
+
+    @Override
+    public byte get(int index) {
+        return buffer.get(index + offset);
+    }
+
+    @Override
+    public int length() {
+        return length;
+    }
+
+    @Override
+    public BytesReference slice(int from, int length) {
+        if (from < 0 || (from + length) > this.length) {
+            throw new IndexOutOfBoundsException("can't slice a buffer with length [" + this.length + "], with slice parameters from ["
+                + from + "], length [" + length + "]");
+        }
+        ByteBuffer newByteBuffer = buffer.duplicate();
+        newByteBuffer.position(offset + from);
+        newByteBuffer.limit(offset + from + length);
+        return new ByteBufferReference(newByteBuffer);
+    }
+
+    /**
+     * This will return a bytes ref composed of the bytes. If this is a direct byte buffer, the bytes will
+     * have to be copied.
+     *
+     * @return the bytes ref
+     */
+    @Override
+    public BytesRef toBytesRef() {
+        if (buffer.hasArray()) {
+            return new BytesRef(buffer.array(), buffer.arrayOffset() + offset, length);
+        }
+        final byte[] copy = new byte[length];
+        buffer.get(copy, offset, length);
+        return new BytesRef(copy);
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        return buffer.capacity();
+    }
+}

core/src/test/java/org/elasticsearch/common/bytes/ByteBufferReferenceTests.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch 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.elasticsearch.common.bytes;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class ByteBufferReferenceTests extends AbstractBytesReferenceTestCase {
+
+    private void initializeBytes(byte[] bytes) {
+        for (int i = 0 ; i < bytes.length; ++i) {
+            bytes[i] = (byte) i;
+        }
+    }
+
+    @Override
+    protected BytesReference newBytesReference(int length) throws IOException {
+        return newBytesReferenceWithOffsetOfZero(length);
+    }
+
+    @Override
+    protected BytesReference newBytesReferenceWithOffsetOfZero(int length) throws IOException {
+        byte[] bytes = new byte[length];
+        initializeBytes(bytes);
+        return new ByteBufferReference(ByteBuffer.wrap(bytes));
+    }
+}

test/framework/src/main/java/org/elasticsearch/transport/nio/InboundChannelBuffer.java

@@ -0,0 +1,204 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch 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.elasticsearch.transport.nio;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayDeque;
+import java.util.Iterator;
+import java.util.function.Supplier;
+
+/**
+ * This is a channel byte buffer composed internally of 16kb pages. When an entire message has been read
+ * and consumed, the {@link #release(long)} method releases the bytes from the head of the buffer and closes
+ * the pages internally. If more space is needed at the end of the buffer {@link #ensureCapacity(long)} can
+ * be called and the buffer will expand using the supplier provided.
+ */
+public final class InboundChannelBuffer {
+
+    private static final int PAGE_SIZE = 1 << 14;
+    private static final int PAGE_MASK = PAGE_SIZE - 1;
+    private static final int PAGE_SHIFT = Integer.numberOfTrailingZeros(PAGE_SIZE);
+    private static final ByteBuffer[] EMPTY_BYTE_BUFFER_ARRAY = new ByteBuffer[0];
+
+
+    private final ArrayDeque<ByteBuffer> pages;
+    private final Supplier<ByteBuffer> pageSupplier;
+
+    private long capacity = 0;
+    private long internalIndex = 0;
+    // The offset is an int as it is the offset of where the bytes begin in the first buffer
+    private int offset = 0;
+
+    public InboundChannelBuffer() {
+        this(() -> ByteBuffer.wrap(new byte[PAGE_SIZE]));
+    }
+
+    private InboundChannelBuffer(Supplier<ByteBuffer> pageSupplier) {
+        this.pageSupplier = pageSupplier;
+        this.pages = new ArrayDeque<>();
+        this.capacity = PAGE_SIZE * pages.size();
+        ensureCapacity(PAGE_SIZE);
+    }
+
+    public void ensureCapacity(long requiredCapacity) {
+        if (capacity < requiredCapacity) {
+            int numPages = numPages(requiredCapacity + offset);
+            int pagesToAdd = numPages - pages.size();
+            for (int i = 0; i < pagesToAdd; i++) {
+                pages.addLast(pageSupplier.get());
+            }
+            capacity += pagesToAdd * PAGE_SIZE;
+        }
+    }
+
+    /**
+     * This method will release bytes from the head of this buffer. If you release bytes past the current
+     * index the index is truncated to zero.
+     *
+     * @param bytesToRelease number of bytes to drop
+     */
+    public void release(long bytesToRelease) {
+        if (bytesToRelease > capacity) {
+            throw new IllegalArgumentException("Releasing more bytes [" + bytesToRelease + "] than buffer capacity [" + capacity + "].");
+        }
+
+        int pagesToRelease = pageIndex(offset + bytesToRelease);
+        for (int i = 0; i < pagesToRelease; i++) {
+            pages.removeFirst();
+        }
+        capacity -= bytesToRelease;
+        internalIndex = Math.max(internalIndex - bytesToRelease, 0);
+        offset = indexInPage(bytesToRelease + offset);
+    }
+
+    /**
+     * This method will return an array of {@link ByteBuffer} representing the bytes from the beginning of
+     * this buffer up through the index argument that was passed. The buffers will be duplicates of the
+     * internal buffers, so any modifications to the markers {@link ByteBuffer#position()},
+     * {@link ByteBuffer#limit()}, etc will not modify the this class.
+     *
+     * @param to the index to slice up to
+     * @return the byte buffers
+     */
+    public ByteBuffer[] sliceBuffersTo(long to) {
+        if (to > capacity) {
+            throw new IndexOutOfBoundsException("can't slice a channel buffer with capacity [" + capacity +
+                "], with slice parameters to [" + to + "]");
+        } else if (to == 0) {
+            return EMPTY_BYTE_BUFFER_ARRAY;
+        }
+        long indexWithOffset = to + offset;
+        int pageCount = pageIndex(indexWithOffset);
+        int finalLimit = indexInPage(indexWithOffset);
+        if (finalLimit != 0) {
+            pageCount += 1;
+        }
+
+        ByteBuffer[] buffers = new ByteBuffer[pageCount];
+        Iterator<ByteBuffer> pageIterator = pages.iterator();
+        ByteBuffer firstBuffer = pageIterator.next().duplicate();
+        firstBuffer.position(firstBuffer.position() + offset);
+        buffers[0] = firstBuffer;
+        for (int i = 1; i < buffers.length; i++) {
+            buffers[i] = pageIterator.next().duplicate();
+        }
+        if (finalLimit != 0) {
+            buffers[buffers.length - 1].limit(finalLimit);
+        }
+
+        return buffers;
+    }
+
+    /**
+     * This method will return an array of {@link ByteBuffer} representing the bytes from the index passed
+     * through the end of this buffer. The buffers will be duplicates of the internal buffers, so any
+     * modifications to the markers {@link ByteBuffer#position()}, {@link ByteBuffer#limit()}, etc will not
+     * modify the this class.
+     *
+     * @param from the index to slice from
+     * @return the byte buffers
+     */
+    public ByteBuffer[] sliceBuffersFrom(long from) {
+        if (from > capacity) {
+            throw new IndexOutOfBoundsException("can't slice a channel buffer with capacity [" + capacity +
+                "], with slice parameters from [" + from + "]");
+        } else if (from == capacity) {
+            return EMPTY_BYTE_BUFFER_ARRAY;
+        }
+        long indexWithOffset = from + offset;
+
+        int pageIndex = pageIndex(indexWithOffset);
+        int indexInPage = indexInPage(indexWithOffset);
+
+        ByteBuffer[] buffers = new ByteBuffer[pages.size() - pageIndex];
+        Iterator<ByteBuffer> pageIterator = pages.descendingIterator();
+        for (int i = buffers.length - 1; i > 0; --i) {
+            buffers[i] = pageIterator.next().duplicate();
+        }
+        ByteBuffer firstPostIndexBuffer = pageIterator.next().duplicate();
+        firstPostIndexBuffer.position(firstPostIndexBuffer.position() + indexInPage);
+        buffers[0] = firstPostIndexBuffer;
+
+        return buffers;
+    }
+
+    public void incrementIndex(long delta) {
+        if (delta < 0) {
+            throw new IllegalArgumentException("Cannot increment an index with a negative delta [" + delta + "]");
+        }
+
+        long newIndex = delta + internalIndex;
+        if (newIndex > capacity) {
+            throw new IllegalArgumentException("Cannot increment an index [" + internalIndex + "] with a delta [" + delta +
+                "] that will result in a new index [" + newIndex + "] that is greater than the capacity [" + capacity + "].");
+        }
+        internalIndex = newIndex;
+    }
+
+    public long getIndex() {
+        return internalIndex;
+    }
+
+    public long getCapacity() {
+        return capacity;
+    }
+
+    public long getRemaining() {
+        long remaining = capacity - internalIndex;
+        assert remaining >= 0 : "The remaining [" + remaining + "] number of bytes should not be less than zero.";
+        return remaining;
+    }
+
+    private int numPages(long capacity) {
+        final long numPages = (capacity + PAGE_MASK) >>> PAGE_SHIFT;
+        if (numPages > Integer.MAX_VALUE) {
+            throw new IllegalArgumentException("pageSize=" + (PAGE_MASK + 1) + " is too small for such as capacity: " + capacity);
+        }
+        return (int) numPages;
+    }
+
+    private int pageIndex(long index) {
+        return (int) (index >>> PAGE_SHIFT);
+    }
+
+    private int indexInPage(long index) {
+        return (int) (index & PAGE_MASK);
+    }
+}