Speed up writeVInt

benchmarks/README.md

@@ -78,7 +78,6 @@ cd fcml*
 make
 cd example/hsdis
 make
-cp .libs/libhsdis.so.0.0.0
 sudo cp .libs/libhsdis.so.0.0.0 /usr/lib/jvm/java-14-adoptopenjdk/lib/hsdis-amd64.so
 ```
 

benchmarks/src/main/java/org/elasticsearch/benchmark/search/aggregations/bucket/terms/StringTermsSerializationBenchmark.java

@@ -0,0 +1,92 @@
+/*
+ * 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.benchmark.search.aggregations.bucket.terms;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.io.stream.DelayableWriteable;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.BucketOrder;
+import org.elasticsearch.search.aggregations.InternalAggregation;
+import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+@Fork(2)
+@Warmup(iterations = 10)
+@Measurement(iterations = 5)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Benchmark)
+public class StringTermsSerializationBenchmark {
+    private static final NamedWriteableRegistry REGISTRY = new NamedWriteableRegistry(
+        List.of(new NamedWriteableRegistry.Entry(InternalAggregation.class, StringTerms.NAME, StringTerms::new))
+    );
+    @Param(value = { "1000" })
+    private int buckets;
+
+    private DelayableWriteable<InternalAggregations> results;
+
+    @Setup
+    public void initResults() {
+        results = DelayableWriteable.referencing(InternalAggregations.from(List.of(newTerms(true))));
+    }
+
+    private StringTerms newTerms(boolean withNested) {
+        List<StringTerms.Bucket> resultBuckets = new ArrayList<>(buckets);
+        for (int i = 0; i < buckets; i++) {
+            InternalAggregations inner = withNested ? InternalAggregations.from(List.of(newTerms(false))) : InternalAggregations.EMPTY;
+            resultBuckets.add(new StringTerms.Bucket(new BytesRef("test" + i), i, inner, false, 0, DocValueFormat.RAW));
+        }
+        return new StringTerms(
+            "test",
+            BucketOrder.key(true),
+            BucketOrder.key(true),
+            buckets,
+            1,
+            null,
+            DocValueFormat.RAW,
+            buckets,
+            false,
+            100000,
+            resultBuckets,
+            0
+        );
+    }
+
+    @Benchmark
+    public DelayableWriteable<InternalAggregations> serialize() {
+        return results.asSerialized(InternalAggregations::readFrom, REGISTRY);
+    }
+}

server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java

@@ -218,14 +218,85 @@ public void writeInt(int i) throws IOException {
      * using {@link #writeInt}
      */
     public void writeVInt(int i) throws IOException {
-        final byte[] buffer = scratch.get();
-        int index = 0;
-        while ((i & ~0x7F) != 0) {
-            buffer[index++] = ((byte) ((i & 0x7f) | 0x80));
-            i >>>= 7;
+        /*
+         * Pick the number of bytes that we need based on the value and then
+         * encode the int, unrolling the loops by hand. This allows writing
+         * small numbers to use `writeByte` which is simple and fast. The
+         * unrolling saves a few comparisons and bitwise operations. All
+         * together this saves quite a bit of time compared to a naive
+         * implementation.
+         */
+        switch (Integer.numberOfLeadingZeros(i)) {
+            case 32:
+            case 31:
+            case 30:
+            case 29:
+            case 28:
+            case 27:
+            case 26:
+            case 25:
+                writeByte((byte) i);
+                return;
+            case 24:
+            case 23:
+            case 22:
+            case 21:
+            case 20:
+            case 19:
+            case 18:
+                byte[] buffer = scratch.get();
+                buffer[0] = (byte) (i & 0x7f | 0x80);
+                buffer[1] = (byte) (i >>> 7);
+                assert buffer[1] <= 0x7f;
+                writeBytes(buffer, 0, 2);
+                return;
+            case 17:
+            case 16:
+            case 15:
+            case 14:
+            case 13:
+            case 12:
+            case 11:
+                buffer = scratch.get();
+                buffer[0] = (byte) (i & 0x7f | 0x80);
+                buffer[1] = (byte) ((i >>> 7) & 0x7f | 0x80);
+                buffer[2] = (byte) (i >>> 14);
+                assert buffer[2] <= 0x7f;
+                writeBytes(buffer, 0, 3);
+                return;
+            case 10:
+            case 9:
+            case 8:
+            case 7:
+            case 6:
+            case 5:
+            case 4:
+                buffer = scratch.get();
+                buffer[0] = (byte) (i & 0x7f | 0x80);
+                buffer[1] = (byte) ((i >>> 7) & 0x7f | 0x80);
+                buffer[2] = (byte) ((i >>> 14) & 0x7f | 0x80);
+                buffer[3] = (byte) (i >>> 21);
+                assert buffer[3] <= 0x7f;
+                writeBytes(buffer, 0, 4);
+                return;
+            case 3:
+            case 2:
+            case 1:
+            case 0:
+                buffer = scratch.get();
+                buffer[0] = (byte) (i & 0x7f | 0x80);
+                buffer[1] = (byte) ((i >>> 7) & 0x7f | 0x80);
+                buffer[2] = (byte) ((i >>> 14) & 0x7f | 0x80);
+                buffer[3] = (byte) ((i >>> 21) & 0x7f | 0x80);
+                buffer[4] = (byte) (i >>> 28);
+                assert buffer[4] <= 0x7f;
+                writeBytes(buffer, 0, 5);
+                return;
+            default:
+                throw new UnsupportedOperationException(
+                    "Can't encode [" + i + "]. Missing case for [" + Integer.numberOfLeadingZeros(i) + "]?"
+                );
         }
-        buffer[index++] = ((byte) i);
-        writeBytes(buffer, 0, index);
     }
 
     /**

server/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.common.io.stream;
 
+import com.carrotsearch.randomizedtesting.annotations.Repeat;
+
 import org.apache.lucene.store.AlreadyClosedException;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.Constants;
@@ -61,7 +63,7 @@
 import static org.hamcrest.Matchers.sameInstance;
 
 /**
- * Tests for {@link BytesStreamOutput} paging behaviour.
+ * Tests for {@link StreamOutput}.
  */
 public class BytesStreamsTests extends ESTestCase {
     public void testEmpty() throws Exception {
@@ -823,10 +825,21 @@ public void testReadNegativeArraySize() throws IOException {
         }
     }
 
+    @Repeat(iterations=1000)
     public void testVInt() throws IOException {
         final int value = randomInt();
         BytesStreamOutput output = new BytesStreamOutput();
         output.writeVInt(value);
+
+        BytesStreamOutput simple = new BytesStreamOutput();
+        int i = value;
+        while ((i & ~0x7F) != 0) {
+            simple.writeByte(((byte) ((i & 0x7f) | 0x80)));
+            i >>>= 7;
+        }
+        simple.writeByte((byte) i);
+        assertEquals(simple.bytes().toBytesRef().toString(), output.bytes().toBytesRef().toString());
+
         StreamInput input = output.bytes().streamInput();
         assertEquals(value, input.readVInt());
     }