Simplify TcpTransport interface by reducing send code to a single send method (#19223)

Due to some optimization on the netty layer we had quite some code / cruft
added to the TcpTransport to allow for those optimizations. After cleaning
up BytesReference we can now move this optimization into TcpTransport and
have a simple send method on the implementation layer instead. This commit
adds a CompositeBytesReference that also allows message headers to be written
separately which simplify the header code as well since no skips are needed
anymore.

Author: s1monw

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

@@ -35,10 +35,7 @@ public final class BytesArray extends BytesReference {
     private final int length;
 
     public BytesArray(String bytes) {
-        BytesRef bytesRef = new BytesRef(bytes);
-        this.bytes = bytesRef.bytes;
-        this.offset = bytesRef.offset;
-        this.length = bytesRef.length;
+        this(new BytesRef(bytes));
     }
 
     public BytesArray(BytesRef bytesRef) {
@@ -47,21 +44,15 @@ public BytesArray(BytesRef bytesRef) {
 
     public BytesArray(BytesRef bytesRef, boolean deepCopy) {
         if (deepCopy) {
-            BytesRef copy = BytesRef.deepCopyOf(bytesRef);
-            bytes = copy.bytes;
-            offset = copy.offset;
-            length = copy.length;
-        } else {
-            bytes = bytesRef.bytes;
-            offset = bytesRef.offset;
-            length = bytesRef.length;
+            bytesRef = BytesRef.deepCopyOf(bytesRef);
         }
+        bytes = bytesRef.bytes;
+        offset = bytesRef.offset;
+        length = bytesRef.length;
     }
 
     public BytesArray(byte[] bytes) {
-        this.bytes = bytes;
-        this.offset = 0;
-        this.length = bytes.length;
+        this(bytes, 0, bytes.length);
     }
 
     public BytesArray(byte[] bytes, int offset, int length) {
@@ -105,4 +96,5 @@ public BytesRef toBytesRef() {
     public long ramBytesUsed() {
         return bytes.length;
     }
+
 }

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

@@ -24,6 +24,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.function.ToIntBiFunction;
 
@@ -52,9 +53,8 @@ public abstract class BytesReference implements Accountable, Comparable<BytesRef
     /**
      * A stream input of the bytes.
      */
-    public StreamInput streamInput() {
-        BytesRef ref = toBytesRef();
-        return StreamInput.wrap(ref.bytes, ref.offset, ref.length);
+    public StreamInput streamInput() throws IOException {
+        return new MarkSupportingStreamInputWrapper(this);
     }
 
     /**
@@ -208,4 +208,73 @@ private static void advance(final BytesRef ref, final int length) {
         ref.length -= length;
         ref.offset += length;
     }
+
+    /**
+     * Instead of adding the complexity of {@link InputStream#reset()} etc to the actual impl
+     * this wrapper builds it on top of the BytesReferenceStreamInput which is much simpler
+     * that way.
+     */
+    private static final class MarkSupportingStreamInputWrapper extends StreamInput {
+        private final BytesReference reference;
+        private BytesReferenceStreamInput input;
+        private int mark = 0;
+
+        private MarkSupportingStreamInputWrapper(BytesReference reference) throws IOException {
+            this.reference = reference;
+            this.input = new BytesReferenceStreamInput(reference.iterator(), reference.length());
+        }
+
+        @Override
+        public byte readByte() throws IOException {
+            return input.readByte();
+        }
+
+        @Override
+        public void readBytes(byte[] b, int offset, int len) throws IOException {
+            input.readBytes(b, offset, len);
+        }
+
+        @Override
+        public int read(byte[] b, int off, int len) throws IOException {
+            return input.read(b, off, len);
+        }
+
+        @Override
+        public void close() throws IOException {
+            input.close();
+        }
+
+        @Override
+        public int read() throws IOException {
+            return input.read();
+        }
+
+        @Override
+        public int available() throws IOException {
+            return input.available();
+        }
+
+        @Override
+        public void reset() throws IOException {
+            input = new BytesReferenceStreamInput(reference.iterator(), reference.length());
+            input.skip(mark);
+        }
+
+        @Override
+        public boolean markSupported() {
+            return true;
+        }
+
+        @Override
+        public void mark(int readLimit) {
+            // readLimit is optional it only guarantees that the stream remembers data upto this limit but it can remember more
+            // which we do in our case
+            this.mark = input.getOffset();
+        }
+
+        @Override
+        public long skip(long n) throws IOException {
+            return input.skip(n);
+        }
+    }
 }

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

@@ -0,0 +1,136 @@
+/*
+ * 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 org.apache.lucene.util.BytesRefIterator;
+import org.elasticsearch.common.io.stream.StreamInput;
+
+import java.io.EOFException;
+import java.io.IOException;
+
+/**
+ * A StreamInput that reads off a {@link BytesRefIterator}. This is used to provide
+ * generic stream access to {@link BytesReference} instances without materializing the
+ * underlying bytes reference.
+ */
+final class BytesReferenceStreamInput extends StreamInput {
+    private final BytesRefIterator iterator;
+    private int sliceOffset;
+    private BytesRef slice;
+    private final int length; // the total size of the stream
+    private int offset; // the current position of the stream
+
+    public BytesReferenceStreamInput(BytesRefIterator iterator, final int length) throws IOException {
+        this.iterator = iterator;
+        this.slice = iterator.next();
+        this.length = length;
+        this.offset = 0;
+        this.sliceOffset = 0;
+    }
+
+    @Override
+    public byte readByte() throws IOException {
+        if (offset >= length) {
+            throw new EOFException();
+        }
+        maybeNextSlice();
+        byte b = slice.bytes[slice.offset + (sliceOffset++)];
+        offset++;
+        return b;
+    }
+
+    private void maybeNextSlice() throws IOException {
+        while (sliceOffset == slice.length) {
+            slice = iterator.next();
+            sliceOffset = 0;
+            if (slice == null) {
+                throw new EOFException();
+            }
+        }
+    }
+
+    @Override
+    public void readBytes(byte[] b, int bOffset, int len) throws IOException {
+        if (offset + len > length) {
+            throw new IndexOutOfBoundsException("Cannot read " + len + " bytes from stream with length " + length + " at offset " + offset);
+        }
+        read(b, bOffset, len);
+    }
+
+    @Override
+    public int read() throws IOException {
+        if (offset >= length) {
+            return -1;
+        }
+        return Byte.toUnsignedInt(readByte());
+    }
+
+    @Override
+    public int read(final byte[] b, final int bOffset, final int len) throws IOException {
+        if (offset >= length) {
+            return -1;
+        }
+        final int numBytesToCopy =  Math.min(len, length - offset);
+        int remaining = numBytesToCopy; // copy the full length or the remaining part
+        int destOffset = bOffset;
+        while (remaining > 0) {
+            maybeNextSlice();
+            final int currentLen = Math.min(remaining, slice.length - sliceOffset);
+            assert currentLen > 0 : "length has to be > 0 to make progress but was: " + currentLen;
+            System.arraycopy(slice.bytes, slice.offset + sliceOffset, b, destOffset, currentLen);
+            destOffset += currentLen;
+            remaining -= currentLen;
+            sliceOffset += currentLen;
+            offset += currentLen;
+            assert remaining >= 0 : "remaining: " + remaining;
+        }
+        return numBytesToCopy;
+    }
+
+    @Override
+    public void close() throws IOException {
+        // do nothing
+    }
+
+    @Override
+    public int available() throws IOException {
+        return length - offset;
+    }
+
+    @Override
+    public long skip(long n) throws IOException {
+        final int skip = (int) Math.min(Integer.MAX_VALUE, n);
+        final int numBytesSkipped =  Math.min(skip, length - offset);
+        int remaining = numBytesSkipped;
+        while (remaining > 0) {
+            maybeNextSlice();
+            int currentLen = Math.min(remaining, slice.length - (slice.offset + sliceOffset));
+            remaining -= currentLen;
+            sliceOffset += currentLen;
+            offset += currentLen;
+            assert remaining >= 0 : "remaining: " + remaining;
+        }
+        return numBytesSkipped;
+    }
+
+    int getOffset() {
+        return offset;
+    }
+}