Revert "Remove dedicated SSL network write buffer (#41283)"

This reverts commit f65a86c.

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

@@ -25,8 +25,6 @@
 
 public class FlushOperation {
 
-    private static final ByteBuffer[] EMPTY_ARRAY = new ByteBuffer[0];
-
     private final BiConsumer<Void, Exception> listener;
     private final ByteBuffer[] buffers;
     private final int[] offsets;
@@ -63,38 +61,19 @@ public void incrementIndex(int delta) {
     }
 
     public ByteBuffer[] getBuffersToWrite() {
-        return getBuffersToWrite(length);
-    }
-
-    public ByteBuffer[] getBuffersToWrite(int maxBytes) {
         final int index = Arrays.binarySearch(offsets, internalIndex);
-        final int offsetIndex = index < 0 ? (-(index + 1)) - 1 : index;
-        final int finalIndex = Arrays.binarySearch(offsets, Math.min(internalIndex + maxBytes, length));
-        final int finalOffsetIndex = finalIndex < 0 ? (-(finalIndex + 1)) - 1 : finalIndex;
+        int offsetIndex = index < 0 ? (-(index + 1)) - 1 : index;
 
-        int nBuffers = (finalOffsetIndex - offsetIndex) + 1;
+        ByteBuffer[] postIndexBuffers = new ByteBuffer[buffers.length - offsetIndex];
 
-        int firstBufferPosition = internalIndex - offsets[offsetIndex];
         ByteBuffer firstBuffer = buffers[offsetIndex].duplicate();
-        firstBuffer.position(firstBufferPosition);
-        if (nBuffers == 1 && firstBuffer.remaining() == 0) {
-            return EMPTY_ARRAY;
-        }
-
-        ByteBuffer[] postIndexBuffers = new ByteBuffer[nBuffers];
+        firstBuffer.position(internalIndex - offsets[offsetIndex]);
         postIndexBuffers[0] = firstBuffer;
-        int finalOffset = offsetIndex + nBuffers;
-        int nBytes = firstBuffer.remaining();
         int j = 1;
-        for (int i = (offsetIndex + 1); i < finalOffset; ++i) {
-            ByteBuffer buffer = buffers[i].duplicate();
-            nBytes += buffer.remaining();
-            postIndexBuffers[j++] = buffer;
+        for (int i = (offsetIndex + 1); i < buffers.length; ++i) {
+            postIndexBuffers[j++] = buffers[i].duplicate();
         }
 
-        int excessBytes = Math.max(0, nBytes - maxBytes);
-        ByteBuffer lastBuffer = postIndexBuffers[postIndexBuffers.length - 1];
-        lastBuffer.limit(lastBuffer.limit() - excessBytes);
         return postIndexBuffers;
     }
 }

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

@@ -27,7 +27,7 @@ public class FlushReadyWrite extends FlushOperation implements WriteOperation {
     private final SocketChannelContext channelContext;
     private final ByteBuffer[] buffers;
 
-    public FlushReadyWrite(SocketChannelContext channelContext, ByteBuffer[] buffers, BiConsumer<Void, Exception> listener) {
+    FlushReadyWrite(SocketChannelContext channelContext, ByteBuffer[] buffers, BiConsumer<Void, Exception> listener) {
         super(buffers, listener);
         this.channelContext = channelContext;
         this.buffers = buffers;

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;
@@ -139,11 +140,11 @@ public ByteBuffer[] sliceBuffersTo(long to) {
 
         ByteBuffer[] buffers = new ByteBuffer[pageCount];
         Iterator<Page> pageIterator = pages.iterator();
-        ByteBuffer firstBuffer = pageIterator.next().byteBuffer().duplicate();
+        ByteBuffer firstBuffer = pageIterator.next().byteBuffer.duplicate();
         firstBuffer.position(firstBuffer.position() + offset);
         buffers[0] = firstBuffer;
         for (int i = 1; i < buffers.length; i++) {
-            buffers[i] = pageIterator.next().byteBuffer().duplicate();
+            buffers[i] = pageIterator.next().byteBuffer.duplicate();
         }
         if (finalLimit != 0) {
             buffers[buffers.length - 1].limit(finalLimit);
@@ -179,14 +180,14 @@ public Page[] sliceAndRetainPagesTo(long to) {
         Page[] pages = new Page[pageCount];
         Iterator<Page> pageIterator = this.pages.iterator();
         Page firstPage = pageIterator.next().duplicate();
-        ByteBuffer firstBuffer = firstPage.byteBuffer();
+        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);
+            pages[pages.length - 1].byteBuffer.limit(finalLimit);
         }
 
         return pages;
@@ -216,9 +217,9 @@ public ByteBuffer[] sliceBuffersFrom(long from) {
         ByteBuffer[] buffers = new ByteBuffer[pages.size() - pageIndex];
         Iterator<Page> pageIterator = pages.descendingIterator();
         for (int i = buffers.length - 1; i > 0; --i) {
-            buffers[i] = pageIterator.next().byteBuffer().duplicate();
+            buffers[i] = pageIterator.next().byteBuffer.duplicate();
         }
-        ByteBuffer firstPostIndexBuffer = pageIterator.next().byteBuffer().duplicate();
+        ByteBuffer firstPostIndexBuffer = pageIterator.next().byteBuffer.duplicate();
         firstPostIndexBuffer.position(firstPostIndexBuffer.position() + indexInPage);
         buffers[0] = firstPostIndexBuffer;
 
@@ -267,4 +268,53 @@ private int pageIndex(long index) {
     private int indexInPage(long index) {
         return (int) (index & PAGE_MASK);
     }
+
+    public static class Page implements AutoCloseable {
+
+        private final ByteBuffer byteBuffer;
+        // 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.refCountedCloseable = refCountedCloseable;
+        }
+
+        private Page duplicate() {
+            refCountedCloseable.incRef();
+            return new Page(byteBuffer.duplicate(), refCountedCloseable);
+        }
+
+        public ByteBuffer getByteBuffer() {
+            return byteBuffer;
+        }
+
+        @Override
+        public void close() {
+            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/main/java/org/elasticsearch/nio/SocketChannelContext.java

@@ -325,7 +325,7 @@ protected int flushToChannel(FlushOperation flushOperation) throws IOException {
             ioBuffer.clear();
             ioBuffer.limit(Math.min(WRITE_LIMIT, ioBuffer.limit()));
             int j = 0;
-            ByteBuffer[] buffers = flushOperation.getBuffersToWrite(WRITE_LIMIT);
+            ByteBuffer[] buffers = flushOperation.getBuffersToWrite();
             while (j < buffers.length && ioBuffer.remaining() > 0) {
                 ByteBuffer buffer = buffers[j++];
                 copyBytes(buffer, ioBuffer);

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

@@ -31,7 +31,6 @@
 import java.util.function.Consumer;
 
 import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -169,7 +168,7 @@ public void testQueuedWriteIsFlushedInFlushCall() throws Exception {
 
         assertTrue(context.readyForFlush());
 
-        when(flushOperation.getBuffersToWrite(anyInt())).thenReturn(buffers);
+        when(flushOperation.getBuffersToWrite()).thenReturn(buffers);
         when(flushOperation.isFullyFlushed()).thenReturn(false, true);
         when(flushOperation.getListener()).thenReturn(listener);
         context.flushChannel();
@@ -188,7 +187,7 @@ public void testPartialFlush() throws IOException {
         assertTrue(context.readyForFlush());
 
         when(flushOperation.isFullyFlushed()).thenReturn(false);
-        when(flushOperation.getBuffersToWrite(anyInt())).thenReturn(new ByteBuffer[] {ByteBuffer.allocate(3)});
+        when(flushOperation.getBuffersToWrite()).thenReturn(new ByteBuffer[] {ByteBuffer.allocate(3)});
         context.flushChannel();
 
         verify(listener, times(0)).accept(null, null);
@@ -202,8 +201,8 @@ public void testMultipleWritesPartialFlushes() throws IOException {
         BiConsumer<Void, Exception> listener2 = mock(BiConsumer.class);
         FlushReadyWrite flushOperation1 = mock(FlushReadyWrite.class);
         FlushReadyWrite flushOperation2 = mock(FlushReadyWrite.class);
-        when(flushOperation1.getBuffersToWrite(anyInt())).thenReturn(new ByteBuffer[] {ByteBuffer.allocate(3)});
-        when(flushOperation2.getBuffersToWrite(anyInt())).thenReturn(new ByteBuffer[] {ByteBuffer.allocate(3)});
+        when(flushOperation1.getBuffersToWrite()).thenReturn(new ByteBuffer[] {ByteBuffer.allocate(3)});
+        when(flushOperation2.getBuffersToWrite()).thenReturn(new ByteBuffer[] {ByteBuffer.allocate(3)});
         when(flushOperation1.getListener()).thenReturn(listener);
         when(flushOperation2.getListener()).thenReturn(listener2);
 
@@ -238,7 +237,7 @@ public void testWhenIOExceptionThrownListenerIsCalled() throws IOException {
         assertTrue(context.readyForFlush());
 
         IOException exception = new IOException();
-        when(flushOperation.getBuffersToWrite(anyInt())).thenReturn(buffers);
+        when(flushOperation.getBuffersToWrite()).thenReturn(buffers);
         when(rawChannel.write(any(ByteBuffer.class))).thenThrow(exception);
         when(flushOperation.getListener()).thenReturn(listener);
         expectThrows(IOException.class, () -> context.flushChannel());
@@ -253,7 +252,7 @@ public void testWriteIOExceptionMeansChannelReadyToClose() throws IOException {
         context.queueWriteOperation(flushOperation);
 
         IOException exception = new IOException();
-        when(flushOperation.getBuffersToWrite(anyInt())).thenReturn(buffers);
+        when(flushOperation.getBuffersToWrite()).thenReturn(buffers);
         when(rawChannel.write(any(ByteBuffer.class))).thenThrow(exception);
 
         assertFalse(context.selectorShouldClose());

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

@@ -65,45 +65,29 @@ public void testMultipleFlushesWithCompositeBuffer() throws IOException {
         ByteBuffer[] byteBuffers = writeOp.getBuffersToWrite();
         assertEquals(3, byteBuffers.length);
         assertEquals(5, byteBuffers[0].remaining());
-        ByteBuffer[] byteBuffersWithLimit = writeOp.getBuffersToWrite(10);
-        assertEquals(2, byteBuffersWithLimit.length);
-        assertEquals(5, byteBuffersWithLimit[0].remaining());
-        assertEquals(5, byteBuffersWithLimit[1].remaining());
 
         writeOp.incrementIndex(5);
         assertFalse(writeOp.isFullyFlushed());
         byteBuffers = writeOp.getBuffersToWrite();
         assertEquals(2, byteBuffers.length);
         assertEquals(15, byteBuffers[0].remaining());
-        assertEquals(3, byteBuffers[1].remaining());
-        byteBuffersWithLimit = writeOp.getBuffersToWrite(10);
-        assertEquals(1, byteBuffersWithLimit.length);
-        assertEquals(10, byteBuffersWithLimit[0].remaining());
 
         writeOp.incrementIndex(2);
         assertFalse(writeOp.isFullyFlushed());
         byteBuffers = writeOp.getBuffersToWrite();
         assertEquals(2, byteBuffers.length);
         assertEquals(13, byteBuffers[0].remaining());
-        assertEquals(3, byteBuffers[1].remaining());
-        byteBuffersWithLimit = writeOp.getBuffersToWrite(10);
-        assertEquals(1, byteBuffersWithLimit.length);
-        assertEquals(10, byteBuffersWithLimit[0].remaining());
 
         writeOp.incrementIndex(15);
         assertFalse(writeOp.isFullyFlushed());
         byteBuffers = writeOp.getBuffersToWrite();
         assertEquals(1, byteBuffers.length);
         assertEquals(1, byteBuffers[0].remaining());
-        byteBuffersWithLimit = writeOp.getBuffersToWrite(10);
-        assertEquals(1, byteBuffersWithLimit.length);
-        assertEquals(1, byteBuffersWithLimit[0].remaining());
 
         writeOp.incrementIndex(1);
         assertTrue(writeOp.isFullyFlushed());
         byteBuffers = writeOp.getBuffersToWrite();
-        assertEquals(0, byteBuffers.length);
-        byteBuffersWithLimit = writeOp.getBuffersToWrite(10);
-        assertEquals(0, byteBuffersWithLimit.length);
+        assertEquals(1, byteBuffers.length);
+        assertEquals(0, byteBuffers[0].remaining());
     }
 }