Add byte array pooling to nio http transport

server/src/main/java/org/elasticsearch/common/util/concurrent/AbstractRefCounted.java → libs/core/src/main/java/org/elasticsearch/common/util/concurrent/AbstractRefCounted.java

@@ -19,8 +19,6 @@
 
 package org.elasticsearch.common.util.concurrent;
 
-import org.apache.lucene.store.AlreadyClosedException;
-
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -68,7 +66,7 @@ public final void decRef() {
     }
 
     protected void alreadyClosed() {
-        throw new AlreadyClosedException(name + " is already closed can't increment refCount current count [" + refCount.get() + "]");
+        throw new IllegalStateException(name + " is already closed can't increment refCount current count [" + refCount.get() + "]");
     }
 
     /**

server/src/main/java/org/elasticsearch/common/util/concurrent/RefCounted.java → libs/core/src/main/java/org/elasticsearch/common/util/concurrent/RefCounted.java

@@ -44,7 +44,7 @@ public interface RefCounted {
      *
      * @see #decRef
      * @see #tryIncRef()
-     * @throws org.apache.lucene.store.AlreadyClosedException iff the reference counter can not be incremented.
+     * @throws IllegalStateException iff the reference counter can not be incremented.
      */
     void incRef();
 

server/src/test/java/org/elasticsearch/common/util/concurrent/RefCountedTests.java → libs/core/src/test/java/org/elasticsearch/common/util/concurrent/RefCountedTests.java

@@ -18,7 +18,6 @@
  */
 package org.elasticsearch.common.util.concurrent;
 
-import org.apache.lucene.store.AlreadyClosedException;
 import org.elasticsearch.test.ESTestCase;
 import org.hamcrest.Matchers;
 
@@ -70,14 +69,14 @@ public void testRefCount() throws IOException {
         try {
             counted.incRef();
             fail(" expected exception");
-        } catch (AlreadyClosedException ex) {
+        } catch (IllegalStateException ex) {
             assertThat(ex.getMessage(), equalTo("test is already closed can't increment refCount current count [0]"));
         }
 
         try {
             counted.ensureOpen();
             fail(" expected exception");
-        } catch (AlreadyClosedException ex) {
+        } catch (IllegalStateException ex) {
             assertThat(ex.getMessage(), equalTo("closed"));
         }
     }
@@ -116,7 +115,7 @@ public void run() {
         try {
             counted.ensureOpen();
             fail("expected to be closed");
-        } catch (AlreadyClosedException ex) {
+        } catch (IllegalStateException ex) {
             assertThat(ex.getMessage(), equalTo("closed"));
         }
         assertThat(counted.refCount(), is(0));
@@ -140,7 +139,7 @@ protected void closeInternal() {
         public void ensureOpen() {
             if (closed.get()) {
                 assert this.refCount() == 0;
-                throw new AlreadyClosedException("closed");
+                throw new IllegalStateException("closed");
             }
         }
     }

libs/nio/src/main/java/org/elasticsearch/nio/InboundChannelBuffer.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.nio;
 
+import org.elasticsearch.common.util.concurrent.AbstractRefCounted;
 import org.elasticsearch.nio.utils.ExceptionsHelper;
 
 import java.nio.ByteBuffer;
@@ -41,6 +42,7 @@ public final class InboundChannelBuffer implements AutoCloseable {
     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 static final Page[] EMPTY_BYTE_PAGE_ARRAY = new Page[0];
 
 
     private final ArrayDeque<Page> pages;
@@ -152,6 +154,46 @@ public ByteBuffer[] sliceBuffersTo(long to) {
         return buffers;
     }
 
+    /**
+     * This method will return an array of {@link Page} representing the bytes from the beginning of
+     * this buffer up through the index argument that was passed. The pages and buffers will be duplicates of
+     * the internal components, so any modifications to the markers {@link ByteBuffer#position()},
+     * {@link ByteBuffer#limit()}, etc will not modify the this class. Additionally, this will internally
+     * retain the underlying pages, so the pages returned by this method must be closed.
+     *
+     * @param to the index to slice up to
+     * @return the pages
+     */
+    public Page[] sliceAndRetainPagesTo(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_PAGE_ARRAY;
+        }
+        long indexWithOffset = to + offset;
+        int pageCount = pageIndex(indexWithOffset);
+        int finalLimit = indexInPage(indexWithOffset);
+        if (finalLimit != 0) {
+            pageCount += 1;
+        }
+
+        Page[] pages = new Page[pageCount];
+        Iterator<Page> pageIterator = this.pages.iterator();
+        Page firstPage = pageIterator.next().duplicate();
+        ByteBuffer firstBuffer = firstPage.byteBuffer;
+        firstBuffer.position(firstBuffer.position() + offset);
+        pages[0] = firstPage;
+        for (int i = 1; i < pages.length; i++) {
+            pages[i] = pageIterator.next().duplicate();
+        }
+        if (finalLimit != 0) {
+            pages[pages.length - 1].byteBuffer.limit(finalLimit);
+        }
+
+        return pages;
+    }
+
     /**
      * 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
@@ -231,16 +273,49 @@ private int indexInPage(long index) {
     public static class Page implements AutoCloseable {
 
         private final ByteBuffer byteBuffer;
-        private final Runnable closeable;
+        // This is reference counted as some implementations want to retain the byte pages by calling
+        // sliceAndRetainPagesTo. With reference counting we can increment the reference count, return the
+        // pages, and safely close them when this channel buffer is done with them. The reference count
+        // would be 1 at that point, meaning that the pages will remain until the implementation closes
+        // theirs.
+        private final RefCountedCloseable refCountedCloseable;
 
         public Page(ByteBuffer byteBuffer, Runnable closeable) {
+            this(byteBuffer, new RefCountedCloseable(closeable));
+        }
+
+        private Page(ByteBuffer byteBuffer, RefCountedCloseable refCountedCloseable) {
             this.byteBuffer = byteBuffer;
-            this.closeable = closeable;
+            this.refCountedCloseable = refCountedCloseable;
+        }
+
+        private Page duplicate() {
+            refCountedCloseable.incRef();
+            return new Page(byteBuffer.duplicate(), refCountedCloseable);
+        }
+
+        public ByteBuffer getByteBuffer() {
+            return byteBuffer;
         }
 
         @Override
         public void close() {
-            closeable.run();
+            refCountedCloseable.decRef();
+        }
+
+        private static class RefCountedCloseable extends AbstractRefCounted {
+
+            private final Runnable closeable;
+
+            private RefCountedCloseable(Runnable closeable) {
+                super("byte array page");
+                this.closeable = closeable;
+            }
+
+            @Override
+            protected void closeInternal() {
+                closeable.run();
+            }
         }
     }
 }

libs/nio/src/test/java/org/elasticsearch/nio/InboundChannelBufferTests.java

@@ -31,7 +31,8 @@ public class InboundChannelBufferTests extends ESTestCase {
 
     private static final int PAGE_SIZE = BigArrays.PAGE_SIZE_IN_BYTES;
     private final Supplier<InboundChannelBuffer.Page> defaultPageSupplier = () ->
-        new InboundChannelBuffer.Page(ByteBuffer.allocate(BigArrays.BYTE_PAGE_SIZE), () -> {});
+        new InboundChannelBuffer.Page(ByteBuffer.allocate(BigArrays.BYTE_PAGE_SIZE), () -> {
+        });
 
     public void testNewBufferHasSinglePage() {
         InboundChannelBuffer channelBuffer = new InboundChannelBuffer(defaultPageSupplier);
@@ -167,6 +168,49 @@ public void testClose() {
         expectThrows(IllegalStateException.class, () -> channelBuffer.ensureCapacity(1));
     }
 
+    public void testCloseRetainedPages() {
+        ConcurrentLinkedQueue<AtomicBoolean> queue = new ConcurrentLinkedQueue<>();
+        Supplier<InboundChannelBuffer.Page> supplier = () -> {
+            AtomicBoolean atomicBoolean = new AtomicBoolean();
+            queue.add(atomicBoolean);
+            return new InboundChannelBuffer.Page(ByteBuffer.allocate(PAGE_SIZE), () -> atomicBoolean.set(true));
+        };
+        InboundChannelBuffer channelBuffer = new InboundChannelBuffer(supplier);
+        channelBuffer.ensureCapacity(PAGE_SIZE * 4);
+
+        assertEquals(4, queue.size());
+
+        for (AtomicBoolean closedRef : queue) {
+            assertFalse(closedRef.get());
+        }
+
+        InboundChannelBuffer.Page[] pages = channelBuffer.sliceAndRetainPagesTo(PAGE_SIZE * 2);
+
+        pages[1].close();
+
+        for (AtomicBoolean closedRef : queue) {
+            assertFalse(closedRef.get());
+        }
+
+        channelBuffer.close();
+
+        int i = 0;
+        for (AtomicBoolean closedRef : queue) {
+            if (i < 1) {
+                assertFalse(closedRef.get());
+            } else {
+                assertTrue(closedRef.get());
+            }
+            ++i;
+        }
+
+        pages[0].close();
+
+        for (AtomicBoolean closedRef : queue) {
+            assertTrue(closedRef.get());
+        }
+    }
+
     public void testAccessByteBuffers() {
         InboundChannelBuffer channelBuffer = new InboundChannelBuffer(defaultPageSupplier);
 

modules/transport-netty4/src/main/java/org/elasticsearch/transport/Netty4Plugin.java

@@ -87,8 +87,8 @@ public Map<String, Supplier<Transport>> getTransports(Settings settings, ThreadP
 
     @Override
     public Map<String, Supplier<HttpServerTransport>> getHttpTransports(Settings settings, ThreadPool threadPool, BigArrays bigArrays,
+                                                                        PageCacheRecycler pageCacheRecycler,
                                                                         CircuitBreakerService circuitBreakerService,
-                                                                        NamedWriteableRegistry namedWriteableRegistry,
                                                                         NamedXContentRegistry xContentRegistry,
                                                                         NetworkService networkService,
                                                                         HttpServerTransport.Dispatcher dispatcher) {

plugins/transport-nio/src/main/java/org/elasticsearch/http/nio/HttpReadWriteHandler.java

@@ -78,7 +78,7 @@ public class HttpReadWriteHandler implements ReadWriteHandler {
 
     @Override
     public int consumeReads(InboundChannelBuffer channelBuffer) throws IOException {
-        int bytesConsumed = adaptor.read(channelBuffer.sliceBuffersTo(channelBuffer.getIndex()));
+        int bytesConsumed = adaptor.read(channelBuffer.sliceAndRetainPagesTo(channelBuffer.getIndex()));
         Object message;
         while ((message = adaptor.pollInboundMessage()) != null) {
             handleRequest(message);

plugins/transport-nio/src/main/java/org/elasticsearch/http/nio/NettyAdaptor.java

@@ -29,6 +29,7 @@
 import io.netty.channel.embedded.EmbeddedChannel;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.nio.FlushOperation;
+import org.elasticsearch.nio.InboundChannelBuffer;
 import org.elasticsearch.nio.WriteOperation;
 
 import java.nio.ByteBuffer;
@@ -97,6 +98,13 @@ public int read(ByteBuffer[] buffers) {
         return byteBuf.readerIndex() - initialReaderIndex;
     }
 
+    public int read(InboundChannelBuffer.Page[] pages) {
+        ByteBuf byteBuf = PagedByteBuf.byteBufFromPages(pages);
+        int readableBytes = byteBuf.readableBytes();
+        nettyChannel.writeInbound(byteBuf);
+        return readableBytes;
+    }
+
     public Object pollInboundMessage() {
         return nettyChannel.readInbound();
     }

plugins/transport-nio/src/main/java/org/elasticsearch/http/nio/NioHttpServerTransport.java

@@ -32,12 +32,14 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.network.NetworkAddress;
 import org.elasticsearch.common.network.NetworkService;
+import org.elasticsearch.common.recycler.Recycler;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.NetworkExceptionHelper;
 import org.elasticsearch.common.transport.TransportAddress;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.http.AbstractHttpServerTransport;
@@ -63,6 +65,7 @@
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
 import java.util.ArrayList;
@@ -103,6 +106,8 @@ public class NioHttpServerTransport extends AbstractHttpServerTransport {
             (s) -> Integer.toString(EsExecutors.numberOfProcessors(s) * 2),
             (s) -> Setting.parseInt(s, 1, "http.nio.worker_count"), Setting.Property.NodeScope);
 
+    private final PageCacheRecycler pageCacheRecycler;
+
     private final boolean tcpNoDelay;
     private final boolean tcpKeepAlive;
     private final boolean reuseAddress;
@@ -115,9 +120,11 @@ public class NioHttpServerTransport extends AbstractHttpServerTransport {
     private HttpChannelFactory channelFactory;
     private final NioCorsConfig corsConfig;
 
-    public NioHttpServerTransport(Settings settings, NetworkService networkService, BigArrays bigArrays, ThreadPool threadPool,
-                                  NamedXContentRegistry xContentRegistry, HttpServerTransport.Dispatcher dispatcher) {
+    public NioHttpServerTransport(Settings settings, NetworkService networkService, BigArrays bigArrays,
+                                  PageCacheRecycler pageCacheRecycler, ThreadPool threadPool, NamedXContentRegistry xContentRegistry,
+                                  HttpServerTransport.Dispatcher dispatcher) {
         super(settings, networkService, bigArrays, threadPool, xContentRegistry, dispatcher);
+        this.pageCacheRecycler = pageCacheRecycler;
 
         ByteSizeValue maxChunkSize = SETTING_HTTP_MAX_CHUNK_SIZE.get(settings);
         ByteSizeValue maxHeaderSize = SETTING_HTTP_MAX_HEADER_SIZE.get(settings);
@@ -329,11 +336,15 @@ private HttpChannelFactory() {
         @Override
         public NioHttpChannel createChannel(NioSelector selector, SocketChannel channel) throws IOException {
             NioHttpChannel nioChannel = new NioHttpChannel(channel);
+            java.util.function.Supplier<InboundChannelBuffer.Page> pageSupplier = () -> {
+                Recycler.V<byte[]> bytes = pageCacheRecycler.bytePage(false);
+                return new InboundChannelBuffer.Page(ByteBuffer.wrap(bytes.v()), bytes::close);
+            };
             HttpReadWriteHandler httpReadWritePipeline = new HttpReadWriteHandler(nioChannel,NioHttpServerTransport.this,
                 handlingSettings, corsConfig);
             Consumer<Exception> exceptionHandler = (e) -> exceptionCaught(nioChannel, e);
             SocketChannelContext context = new BytesChannelContext(nioChannel, selector, exceptionHandler, httpReadWritePipeline,
-                InboundChannelBuffer.allocatingInstance());
+                new InboundChannelBuffer(pageSupplier));
             nioChannel.setContext(context);
             return nioChannel;
         }