From aafe485f21c5b6582aa5ef7a9f979e94713dbd1b Mon Sep 17 00:00:00 2001
From: Jamie Gaskins <jgaskins@hey.com>
Date: Sun, 22 Dec 2024 23:21:01 -0600
Subject: [PATCH] Optimize simple strings for `+OK\r\n` responses

Previously, `redis.set "foo", "bar"` allocated 16 bytes of heap memory
for the "OK" response from the server. With this commit, it doesn't
allocate any heap memory at all.
---
 spec/parser_spec.cr |  9 +++++++--
 spec/redis_spec.cr  |  2 +-
 src/parser.cr       | 17 ++++++++++++++++-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/spec/parser_spec.cr b/spec/parser_spec.cr
index 3d6571d..8a49512 100644
--- a/spec/parser_spec.cr
+++ b/spec/parser_spec.cr
@@ -12,8 +12,13 @@ module Redis
     end
 
     it "reads simple strings" do
-      io = IO::Memory.new("+OK\r\n")
-      Parser.new(io).read.should eq "OK"
+      io = IO::Memory.new("+OK\r\n+QUEUED\r\n+OK\r\n+QUEUED\r\n")
+      parser = Parser.new(io)
+
+      2.times do
+        parser.read.should eq "OK"
+        parser.read.should eq "QUEUED"
+      end
     end
 
     it "reads bulk strings" do
diff --git a/spec/redis_spec.cr b/spec/redis_spec.cr
index 04b93a3..6b8be0e 100644
--- a/spec/redis_spec.cr
+++ b/spec/redis_spec.cr
@@ -27,7 +27,7 @@ end
 describe Redis::Client do
   test "can set, get, and delete keys" do
     redis.get(random_key).should eq nil
-    redis.set(key, "hello")
+    redis.set(key, "hello").should eq "OK"
     redis.get(key).should eq "hello"
     redis.del(key).should eq 1
     redis.del(key).should eq 0
diff --git a/src/parser.cr b/src/parser.cr
index 2cc8921..4c1c746 100644
--- a/src/parser.cr
+++ b/src/parser.cr
@@ -48,7 +48,22 @@ module Redis
           value
         end
       when '+'
-        @io.read_line
+        # Most of the time, RESP simple strings are just "OK", so we can
+        # optimize for that case to avoid heap allocations. If it is *not* the
+        # "OK" string, this does an extra heap allocation, but that seems like
+        # a decent tradeoff considering the vast majority of times a simple
+        # string will be returned from the server is from a SET call.
+        buffer = uninitialized UInt8[4] # "OK\r\n"
+        slice = buffer.to_slice
+        read = @io.read slice
+        result = if read == 4 && slice == "OK\r\n".to_slice
+          "OK"
+        else
+          String.build do |str|
+            str.write slice[0...read]
+            str << @io.read_line
+          end
+        end
       when '-'
         type, message = @io.read_line.split(' ', 2)
         ERROR_MAP[type].new("#{type} #{message}")