Add test for writer operation buffer accounting (#27707)

This is a follow up to #27695. This commit adds a test checking that
across multiple writes using multiple buffers, a write operation
properly keeps track of which buffers still need to be written.

test/framework/src/test/java/org/elasticsearch/transport/nio/WriteOperationTests.java

@@ -21,12 +21,15 @@ package org.elasticsearch.transport.nio;
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.bytes.CompositeBytesReference;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.transport.nio.channel.NioSocketChannel;
 import org.junit.Before;
+import org.mockito.ArgumentCaptor;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.List;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
@@ -65,4 +68,52 @@ public class WriteOperationTests extends ESTestCase {
 
         assertFalse(writeOp.isFullyFlushed());
     }
+
+    public void testMultipleFlushesWithCompositeBuffer() throws IOException {
+        BytesArray bytesReference1 = new BytesArray(new byte[10]);
+        BytesArray bytesReference2 = new BytesArray(new byte[15]);
+        BytesArray bytesReference3 = new BytesArray(new byte[3]);
+        CompositeBytesReference bytesReference = new CompositeBytesReference(bytesReference1, bytesReference2, bytesReference3);
+        WriteOperation writeOp = new WriteOperation(channel, bytesReference, listener);
+
+        ArgumentCaptor<ByteBuffer[]> buffersCaptor = ArgumentCaptor.forClass(ByteBuffer[].class);
+
+        when(channel.write(buffersCaptor.capture())).thenReturn(5)
+            .thenReturn(5)
+            .thenReturn(2)
+            .thenReturn(15)
+            .thenReturn(1);
+
+        writeOp.flush();
+        assertFalse(writeOp.isFullyFlushed());
+        writeOp.flush();
+        assertFalse(writeOp.isFullyFlushed());
+        writeOp.flush();
+        assertFalse(writeOp.isFullyFlushed());
+        writeOp.flush();
+        assertFalse(writeOp.isFullyFlushed());
+        writeOp.flush();
+        assertTrue(writeOp.isFullyFlushed());
+
+        List<ByteBuffer[]> values = buffersCaptor.getAllValues();
+        ByteBuffer[] byteBuffers = values.get(0);
+        assertEquals(3, byteBuffers.length);
+        assertEquals(10, byteBuffers[0].remaining());
+
+        byteBuffers = values.get(1);
+        assertEquals(3, byteBuffers.length);
+        assertEquals(5, byteBuffers[0].remaining());
+
+        byteBuffers = values.get(2);
+        assertEquals(2, byteBuffers.length);
+        assertEquals(15, byteBuffers[0].remaining());
+
+        byteBuffers = values.get(3);
+        assertEquals(2, byteBuffers.length);
+        assertEquals(13, byteBuffers[0].remaining());
+
+        byteBuffers = values.get(4);
+        assertEquals(1, byteBuffers.length);
+        assertEquals(1, byteBuffers[0].remaining());
+    }
 }