content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF 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.apache.kafka.jmh.util;
  
  import org.apache.kafka.common.utils.ByteUtils;
  
  import org.openjdk.jmh.annotations.Benchmark;
  import org.openjdk.jmh.annotations.CompilerControl;
  import org.openjdk.jmh.annotations.Fork;
  import org.openjdk.jmh.annotations.Level;
  import org.openjdk.jmh.annotations.Measurement;
  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 org.openjdk.jmh.infra.Blackhole;
  import org.openjdk.jmh.runner.Runner;
  import org.openjdk.jmh.runner.RunnerException;
  import org.openjdk.jmh.runner.options.Options;
  import org.openjdk.jmh.runner.options.OptionsBuilder;
  
  import java.nio.ByteBuffer;
  import java.security.SecureRandom;
  import java.util.concurrent.TimeUnit;
  
  /**
   * This benchmark calculates the empirical evidence of different implementation for encoding/decoding a protobuf
   * <a href="https://protobuf.dev/programming-guides/encoding/#varints">VarInt</a> and VarLong.
   *
   * The benchmark uses JMH and calculates results for different sizes of variable length integer. We expect most of the
   * usage in Kafka code base to be 1 or 2 byte integers.
   */
  @OutputTimeUnit(TimeUnit.SECONDS)
  @Fork(3)
  @Warmup(iterations = 3, time = 1)
  @Measurement(iterations = 5, time = 1)
  public class ByteUtilsBenchmark {
      private static final int DATA_SET_SAMPLE_SIZE = 16384;
  
      @State(Scope.Benchmark)
      public static class BaseBenchmarkState {
          private ByteBuffer testBuffer;
          private SecureRandom random;
          @Setup(Level.Trial)
          public void setUpBenchmarkLevel() {
              // Initialize the random number generator with a seed so that for each benchmark it produces the same sequence
              // of random numbers. Note that it is important to initialize it again with the seed before every benchmark.
              random = new SecureRandom();
              random.setSeed(133713371337L);
          }
  
          @Setup(Level.Invocation)
          public void setUpInvocationBuffer() {
              testBuffer = ByteBuffer.allocate(10);
          }
  
          /**
           * Generates a random int64 number which occupies exactly bytesSet in the variable length encoding for int64
           *
           * Upper bound is set by finding the largest number that can be represented by N bits. This number is found
           * by bit shifting by N and subtracting decimal 1. For example, for 2 bytes = 16 bits, we calculate the
           * upper bound by:
           * 1. 0001 0000 0000 0000 0000 // bit shift by 16 = 65536 decimal
           * 2. Subtract 1 from 65536 = 65535
           * 3. 65535 is the upper bound which is represented in binary as 1111 1111 1111 1111 i.e. largest number
           *    that could be represented by 2 bytes.
           *
           * Hence, range of random number of different byte length is:
           * 1 byte - [1, 255)
           * 2 byte - [256, 65535)
           * ... and so on.
           */
          long generateRandomLongWithExactBytesSet(int bytesSet) {
              long lowerBound = 1L << ((bytesSet - 1) * 8);
              long upperBound = (1L << (bytesSet * 8)) - 1;
              if (lowerBound >= upperBound) {
                  throw new IllegalArgumentException();
              }
              return lowerBound +
                  random.longs(lowerBound, upperBound).findFirst()
                      .orElseThrow(() -> new IllegalStateException("Unable to create a random long in the range=[" + lowerBound + ", " + upperBound + "]"));
          }
  
          /**
           * Generates a random int32 number which occupies exactly bytesSet in the variable length encoding for int32
           *
           * @see {@link #generateRandomLongWithExactBytesSet(int)} for implementation details.
           *
           */
          int generateRandomIntWithExactBytesSet(int bytesSet) {
              int lowerBound = 1 << ((bytesSet - 1) * 8);
              int upperBound = (1 << (bytesSet * 8)) - 1;
              if (lowerBound >= upperBound) {
                  throw new IllegalArgumentException();
              }
              return lowerBound + random.nextInt(upperBound - lowerBound);
          }
  
          public ByteBuffer getTestBuffer() {
              return testBuffer;
          }
      }
  
      @State(Scope.Benchmark)
      public static class IterationStateForLong extends BaseBenchmarkState {
          @Param({"1", "3", "5", "7"})
          int numNonZeroBytes;
  
          private long[] randomLongs;
  
          @Setup(Level.Iteration)
          public void setup() {
              randomLongs = new long[DATA_SET_SAMPLE_SIZE];
              for (int i = 0; i < DATA_SET_SAMPLE_SIZE; i++) {
                  this.randomLongs[i] = generateRandomLongWithExactBytesSet(numNonZeroBytes);
              }
          }
  
          public long[] getRandomValues() {
              return randomLongs;
          }
      }
  
      @State(Scope.Benchmark)
      public static class IterationStateForInt extends BaseBenchmarkState {
          @Param({"1", "2", "3"})
          int numNonZeroBytes;
          private int[] randomInts;
  
          @Setup(Level.Iteration)
          public void setup() {
              randomInts = new int[DATA_SET_SAMPLE_SIZE];
              for (int i = 0; i < DATA_SET_SAMPLE_SIZE; i++) {
                  this.randomInts[i] = generateRandomIntWithExactBytesSet(numNonZeroBytes);
              }
          }
  
          public int[] getRandomValues() {
              return randomInts;
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedReadVarintNetty(IterationStateForInt state, Blackhole bk) {
          for (int randomValue : state.getRandomValues()) {
              ByteUtils.writeUnsignedVarint(randomValue, state.getTestBuffer());
              // prepare for reading
              state.getTestBuffer().flip();
              bk.consume(ByteUtilsBenchmark.readUnsignedVarintNetty(state.getTestBuffer()));
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedReadVarintLegacy(IterationStateForInt state, Blackhole bk) {
          for (int randomValue : state.getRandomValues()) {
              ByteUtils.writeUnsignedVarint(randomValue, state.getTestBuffer());
              // prepare for reading
              state.getTestBuffer().flip();
              bk.consume(ByteUtilsBenchmark.readUnsignedVarintLegacy(state.getTestBuffer()));
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedReadVarintProtobuf(IterationStateForInt state, Blackhole bk) {
          for (int randomValue : state.getRandomValues()) {
              ByteUtils.writeUnsignedVarint(randomValue, state.getTestBuffer());
              // prepare for reading
              state.getTestBuffer().flip();
              bk.consume(ByteUtilsBenchmark.readUnsignedVarintProtoBuf(state.getTestBuffer()));
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedReadVarlongUnrolled(IterationStateForLong state, Blackhole bk) {
          for (long randomValue : state.getRandomValues()) {
              ByteUtils.writeUnsignedVarlong(randomValue, state.getTestBuffer());
              // prepare for reading
              state.getTestBuffer().flip();
              bk.consume(ByteUtilsBenchmark.readUnsignedVarlongNetty(state.getTestBuffer()));
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedReadVarlongLegacy(IterationStateForLong state, Blackhole bk) {
          for (long randomValue : state.getRandomValues()) {
              ByteUtils.writeUnsignedVarlong(randomValue, state.getTestBuffer());
              // prepare for reading
              state.getTestBuffer().flip();
              bk.consume(ByteUtilsBenchmark.readUnsignedVarlongLegacy(state.getTestBuffer()));
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedWriteVarintUnrolled(IterationStateForInt state) {
          for (int randomValue : state.getRandomValues()) {
              ByteUtilsBenchmark.writeUnsignedVarintUnrolled(randomValue, state.getTestBuffer());
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedWriteVarintLegacy(IterationStateForInt state) {
          for (int randomValue : state.getRandomValues()) {
              ByteUtilsBenchmark.writeUnsignedVarintLegacy(randomValue, state.getTestBuffer());
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedWriteVarlongUnrolled(IterationStateForLong state) {
          for (long randomValue : state.getRandomValues()) {
              ByteUtilsBenchmark.writeUnsignedVarlongUnrolled(randomValue, state.getTestBuffer());
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testUnsignedWriteVarlongLegacy(IterationStateForLong state) {
          for (long randomValue : state.getRandomValues()) {
              ByteUtilsBenchmark.writeUnsignedVarlongLegacy(randomValue, state.getTestBuffer());
              state.getTestBuffer().clear();
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testSizeOfUnsignedVarint(IterationStateForInt state, Blackhole bk) {
          for (int randomValue : state.getRandomValues()) {
              bk.consume(ByteUtils.sizeOfUnsignedVarint(randomValue));
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testSizeOfUnsignedVarintSimple(IterationStateForInt state, Blackhole bk) {
          for (int randomValue : state.getRandomValues()) {
              int value = randomValue;
              int bytes = 1;
              while ((value & 0xffffff80) != 0L) {
                  bytes += 1;
                  value >>>= 7;
              }
              bk.consume(bytes);
          }
      }
  
      @Benchmark
      @CompilerControl(CompilerControl.Mode.DONT_INLINE)
      public void testSizeOfUnsignedVarlong(IterationStateForLong state, Blackhole bk) {
          for (long randomValue : state.getRandomValues()) {
              bk.consume(ByteUtils.sizeOfUnsignedVarlong(randomValue));
          }
      }
  
      public static void main(String[] args) throws RunnerException {
          Options opt = new OptionsBuilder()
                  .include(ByteUtilsBenchmark.class.getSimpleName())
                  .forks(2)
                  .build();
  
          new Runner(opt).run();
      }
  
  
      /* Implementations */
  
      /*
       * Implementation in Trunk as of Apr 2023 / v3.4
       */
      private static int readUnsignedVarintLegacy(ByteBuffer buffer) {
          int value = 0;
          int i = 0;
          int b;
          while (((b = buffer.get()) & 0x80) != 0) {
              value |= (b & 0x7f) << i;
              i += 7;
              if (i > 28)
                  throw new IllegalArgumentException();
          }
          value |= b << i;
          return value;
      }
  
      /*
       * Implementation in Trunk as of Apr 2023 / v3.4
       */
      private static long readUnsignedVarlongLegacy(ByteBuffer buffer)  {
          long value = 0L;
          int i = 0;
          long b;
          while (((b = buffer.get()) & 0x80) != 0) {
              value |= (b & 0x7f) << i;
              i += 7;
              if (i > 63)
                  throw new IllegalArgumentException();
          }
          value |= b << i;
          return value;
      }
  
      /**
       * Implementation copied from Protobuf's implementation.
       * see: https://github.com/protocolbuffers/protobuf/blob/f1c7820c9bd0e31f8b7d091092851441ad2716b6/java/core/src/main/java/com/google/protobuf/CodedInputStream.java#L1048
       */
      private static int readUnsignedVarintProtoBuf(ByteBuffer buf) {
          fastpath:
          {
              int tempPos = buf.position();
              int limit = buf.limit();
  
              if (limit == tempPos) {
                  break fastpath;
              }
  
              final byte[] buffer = buf.array();
              int x;
              if ((x = buffer[tempPos++]) >= 0) {
                  buf.position(tempPos);
                  return x;
              } else if (limit - tempPos < 9) {
                  break fastpath;
              } else if ((x ^= buffer[tempPos++] << 7) < 0) {
                  x ^= ~0 << 7;
              } else if ((x ^= buffer[tempPos++] << 14) >= 0) {
                  x ^= (~0 << 7) ^ (~0 << 14);
              } else if ((x ^= buffer[tempPos++] << 21) < 0) {
                  x ^= (~0 << 7) ^ (~0 << 14) ^ (~0 << 21);
              } else {
                  int y = buffer[tempPos++];
                  x ^= y << 28;
                  x ^= (~0 << 7) ^ (~0 << 14) ^ (~0 << 21) ^ (~0 << 28);
                  if (y < 0
                      && buffer[tempPos++] < 0
                      && buffer[tempPos++] < 0
                      && buffer[tempPos++] < 0
                      && buffer[tempPos++] < 0
                      && buffer[tempPos++] < 0) {
                      break fastpath; // Will throw malformedVarint()
                  }
              }
              buf.position(tempPos);
              return x;
          }
          return readUnsignedVarintLegacy(buf);
      }
  
      /**
       * Implementation copied from Netty
       * see: https://github.com/netty/netty/blob/59aa6e635b9996cf21cd946e64353270679adc73/codec/src/main/java/io/netty/handler/codec/protobuf/ProtobufVarint32FrameDecoder.java#L73
       */
      private static int readUnsignedVarintNetty(ByteBuffer buffer) {
          byte tmp = buffer.get();
          if (tmp >= 0) {
              return tmp;
          } else {
              int result = tmp & 127;
              if ((tmp = buffer.get()) >= 0) {
                  result |= tmp << 7;
              } else {
                  result |= (tmp & 127) << 7;
                  if ((tmp = buffer.get()) >= 0) {
                      result |= tmp << 14;
                  } else {
                      result |= (tmp & 127) << 14;
                      if ((tmp = buffer.get()) >= 0) {
                          result |= tmp << 21;
                      } else {
                          result |= (tmp & 127) << 21;
                          result |= (tmp = buffer.get()) << 28;
                          if (tmp < 0) {
                              throw new IllegalArgumentException();
                          }
                      }
                  }
              }
              return result;
          }
      }
  
      /**
       * Implementation extended from Int implementation from Netty
       * see: https://github.com/netty/netty/blob/59aa6e635b9996cf21cd946e64353270679adc73/codec/src/main/java/io/netty/handler/codec/protobuf/ProtobufVarint32FrameDecoder.java#L73
       */
      private static long readUnsignedVarlongNetty(ByteBuffer buffer) {
          byte tmp = buffer.get();
          if (tmp >= 0) {
              return tmp;
          } else {
              long result = tmp & 0x7f;
              if ((tmp = buffer.get()) >= 0) {
                  result |= tmp << 7;
              } else {
                  result |= (tmp & 0x7f) << 7;
                  if ((tmp = buffer.get()) >= 0) {
                      result |= tmp << 14;
                  } else {
                      result |= (tmp & 0x7f) << 14;
                      if ((tmp = buffer.get()) >= 0) {
                          result |= tmp << 21;
                      } else {
                          result |= (tmp & 0x7f) << 21;
                          if ((tmp = buffer.get()) >= 0) {
                              result |= (long) tmp << 28;
                          } else {
                              result |= (long) (tmp & 0x7f) << 28;
                              result = innerUnsignedReadVarLong(buffer, result);
                          }
                      }
                  }
              }
              return result;
          }
      }
  
      private static long innerUnsignedReadVarLong(ByteBuffer buffer, long result) {
          byte tmp;
          if ((tmp = buffer.get()) >= 0) {
              result |= (long) tmp << 35;
          } else {
              result |= (long) (tmp & 0x7f) << 35;
              if ((tmp = buffer.get()) >= 0) {
                  result |= (long) tmp << 42;
              } else {
                  result |= (long) (tmp & 0x7f) << 42;
                  if ((tmp = buffer.get()) >= 0) {
                      result |= (long) tmp << 49;
                  } else {
                      result |= (long) (tmp & 0x7f) << 49;
                      if ((tmp = buffer.get()) >= 0) {
                          result |= (long) tmp << 56;
                      } else {
                          result |= (long) (tmp & 0x7f) << 56;
                          result |= (long) (tmp = buffer.get()) << 63;
                          if (tmp < 0) {
                              throw new IllegalArgumentException();
                          }
                      }
                  }
              }
          }
          return result;
      }
  
      /*
       * Implementation in Trunk as of Apr 2023 / v3.4
       */
      private static void writeUnsignedVarintLegacy(int value, ByteBuffer buffer) {
          while ((value & 0xffffff80) != 0L) {
              byte b = (byte) ((value & 0x7f) | 0x80);
              buffer.put(b);
              value >>>= 7;
          }
          buffer.put((byte) value);
      }
  
      /*
       * Implementation in Trunk as of Apr 2023 / v3.4
       */
      private static void writeUnsignedVarlongLegacy(long v, ByteBuffer buffer) {
          while ((v & 0xffffffffffffff80L) != 0L) {
              byte b = (byte) ((v & 0x7f) | 0x80);
              buffer.put(b);
              v >>>= 7;
          }
          buffer.put((byte) v);
      }
  
      /*
       * Implementation extended for Long from the Int implementation at https://github.com/astei/varint-writing-showdown/tree/dev (MIT License)
       * see: https://github.com/astei/varint-writing-showdown/blob/6b1a4baec4b1f0ce65fa40cf0b282ec775fdf43e/src/jmh/java/me/steinborn/varintshowdown/res/SmartNoDataDependencyUnrolledVarIntWriter.java#L8
       */
      private static void writeUnsignedVarlongUnrolled(long value, ByteBuffer buffer) {
          if ((value & (0xFFFFFFFF << 7)) == 0) {
              buffer.put((byte) value);
          } else {
              buffer.put((byte) (value & 0x7F | 0x80));
              if ((value & (0xFFFFFFFF << 14)) == 0) {
                  buffer.put((byte) (value >>> 7));
              } else {
                  buffer.put((byte) ((value >>> 7) & 0x7F | 0x80));
                  if ((value & (0xFFFFFFFF << 21)) == 0) {
                      buffer.put((byte) (value >>> 14));
                  } else {
                      buffer.put((byte) ((value >>> 14) & 0x7F | 0x80));
                      if ((value & (0xFFFFFFFF << 28)) == 0) {
                          buffer.put((byte) (value >>> 21));
                      } else {
                          buffer.put((byte) ((value >>> 21) & 0x7F | 0x80));
                          innerWriteUnsignedVarlongUnrolled(value, buffer);
                      }
                  }
              }
          }
      }
  
      private static void innerWriteUnsignedVarlongUnrolled(long value, ByteBuffer buffer) {
          if ((value & (0xFFFFFFFFFFFFFFFFL << 35)) == 0) {
              buffer.put((byte) (value >>> 28));
          } else {
              buffer.put((byte) ((value >>> 28) & 0x7FL | 0x80));
              if ((value & (0xFFFFFFFFFFFFFFFFL << 42)) == 0) {
                  buffer.put((byte) (value >>> 35));
              } else {
                  buffer.put((byte) ((value >>> 35) & 0x7FL | 0x80));
                  if ((value & (0xFFFFFFFFFFFFFFFFL << 49)) == 0) {
                      buffer.put((byte) (value >>> 42));
                  } else {
                      buffer.put((byte) ((value >>> 42) & 0x7FL | 0x80));
                      if ((value & (0xFFFFFFFFFFFFFFFFL << 56)) == 0) {
                          buffer.put((byte) (value >>> 49));
                      } else {
                          buffer.put((byte) ((value >>> 49) & 0x7FL | 0x80));
                          if ((value & (0xFFFFFFFFFFFFFFFFL << 63)) == 0) {
                              buffer.put((byte) (value >>> 56));
                          } else {
                              buffer.put((byte) ((value >>> 56) & 0x7FL | 0x80));
                              buffer.put((byte) (value >>> 63));
                          }
                      }
                  }
              }
          }
      }
  
      /*
       * Implementation copied from https://github.com/astei/varint-writing-showdown/tree/dev (MIT License)
       * see: https://github.com/astei/varint-writing-showdown/blob/6b1a4baec4b1f0ce65fa40cf0b282ec775fdf43e/src/jmh/java/me/steinborn/varintshowdown/res/SmartNoDataDependencyUnrolledVarIntWriter.java#L8
       */
      private static void writeUnsignedVarintUnrolled(int value, ByteBuffer buffer) {
          if ((value & (0xFFFFFFFF << 7)) == 0) {
              buffer.put((byte) value);
          } else {
              buffer.put((byte) (value & 0x7F | 0x80));
              if ((value & (0xFFFFFFFF << 14)) == 0) {
                  buffer.put((byte) ((value >>> 7) & 0xFF));
              } else {
                  buffer.put((byte) ((value >>> 7) & 0x7F | 0x80));
                  if ((value & (0xFFFFFFFF << 21)) == 0) {
                      buffer.put((byte) ((value >>> 14) & 0xFF));
                  } else {
                      buffer.put((byte) ((value >>> 14) & 0x7F | 0x80));
                      if ((value & (0xFFFFFFFF << 28)) == 0) {
                          buffer.put((byte) ((value >>> 21) & 0xFF));
                      } else {
                          buffer.put((byte) ((value >>> 21) & 0x7F | 0x80));
                          buffer.put((byte) ((value >>> 28) & 0xFF));
                      }
                  }
              }
          }
      }
  }
  