chore: Introduce small buffer in redis parser

This is needed in order to eliminate cases where we return INPUT_PENDING but do not consume the whole string by rejecting just several bytes. This should simplify buffer management for the caller, so that if they pass a string that did not result in complete parsed request, at least the whole string is consumed and can be discarded. Signed-off-by: Roman Gershman <[email protected]>
dragonflydb · Nov 29, 2024 · b788c86 · b788c86
1 parent 68b7baf
commit b788c86
Show file tree

Hide file tree

Showing 3 changed files with 90 additions and 29 deletions.
diff --git a/src/facade/redis_parser.cc b/src/facade/redis_parser.cc
@@ -17,12 +17,16 @@ auto RedisParser::Parse(Buffer str, uint32_t* consumed, RespExpr::Vec* res) -> R
   *consumed = 0;
   res->clear();
 
-  if (str.size() < 2) {
+  if (str.size() + small_len_ < 2) {
+    memcpy(small_buf_ + small_len_, str.data(), str.size());
+    small_len_ += str.size();
+    *consumed = str.size();
+
     return INPUT_PENDING;
   }
 
   if (state_ == CMD_COMPLETE_S) {
-    InitStart(str[0], res);
+    InitStart(small_len_ > 0 ? small_buf_[0] : str[0], res);
   } else {
     // We continue parsing in the middle.
     if (!cached_expr_)
@@ -39,11 +43,8 @@ auto RedisParser::Parse(Buffer str, uint32_t* consumed, RespExpr::Vec* res) -> R
         resultc = ConsumeArrayLen(str);
         break;
       case PARSE_ARG_S:
-        if (str.size() == 0 || (str.size() < 4 && str[0] != '_')) {
-          resultc.first = INPUT_PENDING;
-        } else {
-          resultc = ParseArg(str);
-        }
+        DCHECK(!str.empty());
+        resultc = ParseArg(str);
         break;
       case INLINE_S:
         DCHECK(parse_stack_.empty());
@@ -66,13 +67,16 @@ auto RedisParser::Parse(Buffer str, uint32_t* consumed, RespExpr::Vec* res) -> R
     }
 
     if (resultc.first == INPUT_PENDING) {
+      DCHECK(str.empty());
       StashState(res);
     }
     return resultc.first;
   }
 
   if (resultc.first == OK) {
     DCHECK(cached_expr_);
+    DCHECK_EQ(0, small_len_);
+
     if (res != cached_expr_) {
       DCHECK(!stash_.empty());
 
@@ -213,13 +217,27 @@ auto RedisParser::ParseInline(Buffer str) -> ResultConsumed {
 auto RedisParser::ParseLen(Buffer str, int64_t* res) -> ResultConsumed {
   DCHECK(!str.empty());
 
-  DCHECK(str[0] == '$' || str[0] == '*' || str[0] == '%' || str[0] == '~');
+  DCHECK(small_len_ > 0 || str[0] == '$' || str[0] == '*' || str[0] == '%' || str[0] == '~');
 
   const char* s = reinterpret_cast<const char*>(str.data());
+  unsigned consumed = 0;
   const char* pos = reinterpret_cast<const char*>(memchr(s, '\n', str.size()));
   if (!pos) {
-    Result r = str.size() < 32 ? INPUT_PENDING : BAD_ARRAYLEN;
-    return {r, 0};
+    if (str.size() + small_len_ < sizeof(small_buf_)) {
+      memcpy(small_buf_ + small_len_, str.data(), str.size());
+      small_len_ += str.size();
+      return {INPUT_PENDING, str.size()};
+    }
+    return ResultConsumed{BAD_ARRAYLEN, 0};
+  }
+
+  consumed = pos - s + 1;
+  if (small_len_ > 0) {
+    memcpy(small_buf_ + small_len_, str.data(), consumed);
+    small_len_ += consumed;
+    s = small_buf_;
+    pos = small_buf_ + small_len_ - 1;
+    small_len_ = 0;
   }
 
   if (pos[-1] != '\r') {
@@ -228,8 +246,6 @@ auto RedisParser::ParseLen(Buffer str, int64_t* res) -> ResultConsumed {
 
   // Skip the first character and 2 last ones (\r\n).
   bool success = absl::SimpleAtoi(std::string_view{s + 1, size_t(pos - 1 - s)}, res);
-  unsigned consumed = pos - s + 1;
-
   return ResultConsumed{success ? OK : BAD_ARRAYLEN, consumed};
 }
 
@@ -295,11 +311,13 @@ auto RedisParser::ConsumeArrayLen(Buffer str) -> ResultConsumed {
 auto RedisParser::ParseArg(Buffer str) -> ResultConsumed {
   DCHECK(!str.empty());
 
-  char c = str[0];
+  char c = small_len_ > 0 ? small_buf_[0] : str[0];
   unsigned min_len = 3 + int(c != '_');
 
-  if (str.size() < min_len) {
-    return {INPUT_PENDING, 0};
+  if (small_len_ + str.size() < min_len) {
+    memcpy(small_buf_ + small_len_, str.data(), str.size());
+    small_len_ += str.size();
+    return {INPUT_PENDING, str.size()};
   }
 
   if (c == '$') {
@@ -334,10 +352,14 @@ auto RedisParser::ParseArg(Buffer str) -> ResultConsumed {
 
   if (c == '_') {  // Resp3 NIL
     // '_','\r','\n'
-    DCHECK_GE(str.size(), 3u);
+    DCHECK_GE(small_len_ + str.size(), 3u);
+    DCHECK_LT(small_len_, 3);
 
-    unsigned consumed = 3;
-    if (str[1] != '\r' || str[2] != '\n') {
+    unsigned consumed = 3 - small_len_;
+    for (unsigned i = 0; i < consumed; ++i) {
+      small_buf_[small_len_ + i] = str[i];
+    }
+    if (small_buf_[1] != '\r' || small_buf_[2] != '\n') {
       return {BAD_STRING, 0};
     }
 
@@ -408,6 +430,26 @@ auto RedisParser::ConsumeBulk(Buffer str) -> ResultConsumed {
 
   uint32_t consumed = 0;
 
+  if (small_len_ > 0) {
+    DCHECK(!is_broken_token_);
+    DCHECK_EQ(bulk_len_, 0u);
+
+    if (bulk_len_ == 0) {
+      DCHECK_EQ(small_len_, 1);
+      DCHECK_GE(str.size(), 1u);
+      if (small_buf_[0] != '\r' || str[0] != '\n') {
+        return {BAD_STRING, 0};
+      }
+      consumed = bulk_len_ + 2;
+      small_len_ = 0;
+      HandleFinishArg();
+
+      return {OK, 1};
+    }
+  }
+
+  DCHECK_EQ(small_len_, 0);
+
   if (str.size() >= bulk_len_) {
     consumed = bulk_len_;
     if (bulk_len_) {
@@ -473,6 +515,7 @@ void RedisParser::HandleFinishArg() {
     }
     cached_expr_ = parse_stack_.back().second;
   }
+  small_len_ = 0;
 }
 
 void RedisParser::ExtendLastString(Buffer str) {

diff --git a/src/facade/redis_parser.h b/src/facade/redis_parser.h
@@ -47,8 +47,6 @@ class RedisParser {
    * part of str because parser caches the intermediate state internally according to 'consumed'
    * result.
    *
-   * Note: A parser does not always guarantee progress, i.e. if a small buffer was passed it may
-   * returns INPUT_PENDING with consumed == 0.
    *
    */
 
@@ -99,6 +97,7 @@ class RedisParser {
   State state_ = CMD_COMPLETE_S;
   bool is_broken_token_ = false;  // true, if a token (inline or bulk) is broken during the parsing.
   bool server_mode_ = true;
+  uint8_t small_len_ = 0;
 
   uint32_t bulk_len_ = 0;
   uint32_t last_stashed_level_ = 0, last_stashed_index_ = 0;
@@ -114,6 +113,7 @@ class RedisParser {
 
   using Blob = std::vector<uint8_t>;
   std::vector<Blob> buf_stash_;
+  char small_buf_[32];
 };
 
 }  // namespace facade
diff --git a/src/facade/redis_parser_test.cc b/src/facade/redis_parser_test.cc
@@ -107,10 +107,10 @@ TEST_F(RedisParserTest, Multi1) {
 
 TEST_F(RedisParserTest, Multi2) {
   ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("*1\r\n$"));
-  EXPECT_EQ(4, consumed_);
+  EXPECT_EQ(5, consumed_);
 
-  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("$4\r\nMSET"));
-  EXPECT_EQ(8, consumed_);
+  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("4\r\nMSET"));
+  EXPECT_EQ(7, consumed_);
 
   ASSERT_EQ(RedisParser::OK, Parse("\r\n*2\r\n"));
   EXPECT_EQ(2, consumed_);
@@ -138,6 +138,7 @@ TEST_F(RedisParserTest, Multi3) {
 TEST_F(RedisParserTest, ClientMode) {
   parser_.SetClientMode();
 
+#if 0
   ASSERT_EQ(RedisParser::OK, Parse(":-1\r\n"));
   EXPECT_THAT(args_, ElementsAre(IntArg(-1)));
 
@@ -146,6 +147,16 @@ TEST_F(RedisParserTest, ClientMode) {
 
   ASSERT_EQ(RedisParser::OK, Parse("-ERR foo bar\r\n"));
   EXPECT_THAT(args_, ElementsAre(ErrArg("ERR foo")));
+
+  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("_"));
+  EXPECT_EQ(1, consumed_);
+  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("\r"));
+  EXPECT_EQ(1, consumed_);
+  ASSERT_EQ(RedisParser::OK, Parse("\n"));
+  EXPECT_EQ(1, consumed_);
+  EXPECT_THAT(args_, ElementsAre(ArgType(RespExpr::NIL)));
+#endif
+  ASSERT_EQ(RedisParser::OK, Parse("*2\r\n_\r\n_\r\n"));
 }
 
 TEST_F(RedisParserTest, Hierarchy) {
@@ -171,25 +182,25 @@ TEST_F(RedisParserTest, Empty) {
 
 TEST_F(RedisParserTest, LargeBulk) {
   std::string_view prefix("*1\r\n$1024\r\n");
+  string half(512, 'a');
 
   ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(prefix));
   ASSERT_EQ(prefix.size(), consumed_);
   ASSERT_GE(parser_.parselen_hint(), 1024);
 
-  string half(512, 'a');
   ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(half));
   ASSERT_EQ(512, consumed_);
   ASSERT_GE(parser_.parselen_hint(), 512);
   ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(half));
   ASSERT_EQ(512, consumed_);
   ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("\r"));
-  ASSERT_EQ(0, consumed_);
-  ASSERT_EQ(RedisParser::OK, Parse("\r\n"));
-  ASSERT_EQ(2, consumed_);
+  ASSERT_EQ(1, consumed_);
+  ASSERT_EQ(RedisParser::OK, Parse("\n"));
+  EXPECT_EQ(1, consumed_);
 
   string part1 = absl::StrCat(prefix, half);
-  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(part1));
-  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(half));
+  EXPECT_EQ(RedisParser::INPUT_PENDING, Parse(part1));
+  EXPECT_EQ(RedisParser::INPUT_PENDING, Parse(half));
   ASSERT_EQ(RedisParser::OK, Parse("\r\n"));
 }
 
@@ -231,4 +242,11 @@ TEST_F(RedisParserTest, UsedMemory) {
   EXPECT_GT(dfly::HeapSize(stash), 30000);
 }
 
+TEST_F(RedisParserTest, Eol) {
+  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("*1\r"));
+  EXPECT_EQ(3, consumed_);
+  ASSERT_EQ(RedisParser::INPUT_PENDING, Parse("\n$5\r\n"));
+  EXPECT_EQ(5, consumed_);
+}
+
 }  // namespace facade