From 3ba8eb10b948bbffed96c06f7c7bdede56e819b1 Mon Sep 17 00:00:00 2001
From: skr <shikharid@gmail.com>
Date: Fri, 23 Feb 2024 09:41:33 +0530
Subject: [PATCH] Use pre-allocated State's in `RedisStateMachine`, avoiding
 need for State object allocs #2610

* adds gc and thrpt profiling in RedisStateMachine benchmark
* fixes a stale benchmark which caused compilation errors ClusterDistributionChannelWriterBenchmark

Original pull request: #2768
---
 .../core/protocol/RedisStateMachine.java      | 53 ++++++++++---------
 ...terDistributionChannelWriterBenchmark.java |  3 +-
 .../jmh/io/lettuce/core/protocol/JmhMain.java | 17 ++++--
 3 files changed, 42 insertions(+), 31 deletions(-)

diff --git a/src/main/java/io/lettuce/core/protocol/RedisStateMachine.java b/src/main/java/io/lettuce/core/protocol/RedisStateMachine.java
index 1d446b8117..ac292f0c49 100644
--- a/src/main/java/io/lettuce/core/protocol/RedisStateMachine.java
+++ b/src/main/java/io/lettuce/core/protocol/RedisStateMachine.java
@@ -37,6 +37,7 @@
  * @author Will Glozer
  * @author Mark Paluch
  * @author Helly Guo
+ * @author shikharid
  */
 public class RedisStateMachine {
 
@@ -225,9 +226,27 @@ public String toString() {
             return sb.toString();
         }
 
+        void reset() {
+            this.type = null;
+            this.count = NOT_FOUND;
+        }
+
+        /**
+         * Pre-allocates a State array of given len
+         *
+         * @param len len of the states array to be created
+         * @return array of State's of len size
+         */
+        private static State[] createStates(int len) {
+            final State[] stack = new State[len];
+            for (int i = 0;i < len; ++i) {
+                stack[i] = new State();
+            }
+            return stack;
+        }
     }
 
-    private final State[] stack = new State[32];
+    private final State[] stack = State.createStates(32);
 
     private final boolean debugEnabled = logger.isDebugEnabled();
 
@@ -290,7 +309,7 @@ public boolean decode(ByteBuf buffer, CommandOutput<?, ?, ?> output, Consumer<Ex
         buffer.touch("RedisStateMachine.decode(…)");
 
         if (isEmpty(stack)) {
-            add(stack, new State());
+            addHead(stack);
         }
 
         if (output == null) {
@@ -549,7 +568,7 @@ static State.Result returnDependStateCount(RedisStateMachine rsm, State state) {
         }
 
         state.count--;
-        rsm.addFirst(rsm.stack, new State());
+        rsm.addHead(rsm.stack);
 
         return State.Result.CONTINUE_LOOP;
     }
@@ -587,7 +606,9 @@ private static State.Result handleAttribute(RedisStateMachine rsm, State state,
      * Reset the state machine.
      */
     public void reset() {
-        Arrays.fill(stack, null);
+        for (State state : stack) {
+            state.reset();
+        }
         stackElements = 0;
     }
 
@@ -689,18 +710,16 @@ private ByteBuffer readBytes0(ByteBuf buffer, int count) {
      * @param stack
      */
     private void remove(State[] stack) {
-        stack[stackElements - 1] = null;
-        stackElements--;
+        stack[--stackElements].reset();
     }
 
     /**
      * Add the element to the stack to be the new head element.
      *
      * @param stack
-     * @param state
      */
-    private void addFirst(State[] stack, State state) {
-        stack[stackElements++] = state;
+    private void addHead(State[] stack) {
+        ++stackElements;
     }
 
     /**
@@ -713,22 +732,6 @@ private State peek(State[] stack) {
         return stack[stackElements - 1];
     }
 
-    /**
-     * Add a state as tail element. This method shifts the whole stack if the stack is not empty.
-     *
-     * @param stack
-     * @param state
-     */
-    private void add(State[] stack, State state) {
-
-        if (stackElements != 0) {
-            System.arraycopy(stack, 0, stack, 1, stackElements);
-        }
-
-        stack[0] = state;
-        stackElements++;
-    }
-
     /**
      * @param stack
      * @return number of stack elements.
diff --git a/src/test/jmh/io/lettuce/core/cluster/ClusterDistributionChannelWriterBenchmark.java b/src/test/jmh/io/lettuce/core/cluster/ClusterDistributionChannelWriterBenchmark.java
index f2d36a27d6..b7cb44e7bb 100644
--- a/src/test/jmh/io/lettuce/core/cluster/ClusterDistributionChannelWriterBenchmark.java
+++ b/src/test/jmh/io/lettuce/core/cluster/ClusterDistributionChannelWriterBenchmark.java
@@ -22,6 +22,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import io.lettuce.core.protocol.ConnectionIntent;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.Setup;
@@ -69,7 +70,7 @@ public class ClusterDistributionChannelWriterBenchmark {
     @Setup
     public void setup() {
 
-        writer = new ClusterDistributionChannelWriter(CLIENT_OPTIONS, EMPTY_WRITER, ClusterEventListener.NO_OP);
+        writer = new ClusterDistributionChannelWriter(EMPTY_WRITER, CLIENT_OPTIONS, ClusterEventListener.NO_OP);
 
         Partitions partitions = new Partitions();
 
diff --git a/src/test/jmh/io/lettuce/core/protocol/JmhMain.java b/src/test/jmh/io/lettuce/core/protocol/JmhMain.java
index b20117baf6..ff475d746b 100644
--- a/src/test/jmh/io/lettuce/core/protocol/JmhMain.java
+++ b/src/test/jmh/io/lettuce/core/protocol/JmhMain.java
@@ -35,9 +35,9 @@ public static void main(String... args) throws RunnerException {
 
         // run selectively
         // runCommandBenchmark();
-        runCommandHandlerBenchmark();
+        // runCommandHandlerBenchmark();
         // runRedisEndpointBenchmark();
-        // runRedisStateMachineBenchmark();
+        runRedisStateMachineBenchmark();
         // runCommandEncoderBenchmark();
 
         // or all
@@ -83,14 +83,21 @@ private static void runCommandEncoderBenchmark() throws RunnerException {
 
     private static void runRedisStateMachineBenchmark() throws RunnerException {
 
-        new Runner(prepareOptions().mode(Mode.AverageTime).timeUnit(TimeUnit.NANOSECONDS)
+        // measures AverageTime in ns (time/op)
+        new Runner(prepareRSMOptions().mode(Mode.AverageTime).timeUnit(TimeUnit.NANOSECONDS)
+                .include(".*RedisStateMachineBenchmark.*").build()).run();
+
+        // measures Throughput (ops/sec)
+        new Runner(prepareRSMOptions().mode(Mode.Throughput).timeUnit(TimeUnit.SECONDS)
                 .include(".*RedisStateMachineBenchmark.*").build()).run();
-        // new
-        // Runner(prepareOptions().mode(Mode.Throughput).timeUnit(TimeUnit.SECONDS).include(".*CommandHandlerBenchmark.*").build()).run();
     }
 
     private static ChainedOptionsBuilder prepareOptions() {
         return new OptionsBuilder().forks(1).warmupIterations(5).threads(1).measurementIterations(5)
                 .timeout(TimeValue.seconds(2));
     }
+
+    private static ChainedOptionsBuilder prepareRSMOptions() {
+        return prepareOptions().addProfiler("gc");
+    }
 }