Preserve Redis line scan progress across reads

src/brpc/details/redis_line_scanner.h

@@ -0,0 +1,72 @@
+// 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.
+
+
+#ifndef BRPC_DETAILS_REDIS_LINE_SCANNER_H
+#define BRPC_DETAILS_REDIS_LINE_SCANNER_H
+
+#include "butil/iobuf.h"
+
+namespace brpc {
+namespace detail {
+
+// Input may be appended between calls, but its existing prefix must not change.
+// Keep the backing-block position too: forwarding a fresh iterator to a saved
+// byte offset would still repeatedly walk all earlier blocks.
+struct RedisLineScanState {
+    uint32_t block_index;
+    uint32_t block_offset;
+    uint32_t scanned;
+    bool previous_cr;
+
+    void Reset() {
+        block_index = block_offset = scanned = 0;
+        previous_cr = false;
+    }
+
+    size_t FindCRLF(const butil::IOBuf& buf, size_t limit) {
+        if (scanned > buf.size()) {
+            Reset();
+        }
+        while (block_index < buf.backing_block_num() && scanned < limit) {
+            const butil::StringPiece block = buf.backing_block(block_index);
+            while (block_offset < block.size() && scanned < limit) {
+                const char c = block[block_offset++];
+                ++scanned;
+                if (previous_cr && c == '\n') {
+                    return scanned - 2;
+                }
+                previous_cr = (c == '\r');
+            }
+            if (block_offset < block.size() ||
+                block_index + 1 >= buf.backing_block_num()) {
+                // The last block may grow on the next append.
+                break;
+            }
+            ++block_index;
+            block_offset = 0;
+        }
+        return static_cast<size_t>(-1);
+    }
+};
+
+static_assert(sizeof(RedisLineScanState) <= 16,
+              "Line scan state must fit in the RedisReply union");
+
+}  // namespace detail
+}  // namespace brpc
+#endif  // BRPC_DETAILS_REDIS_LINE_SCANNER_H

src/brpc/redis_command.cpp

@@ -29,22 +29,6 @@ namespace {
 const size_t CTX_WIDTH = 5;
 const size_t CRLF_NOT_FOUND = (size_t)-1;
 
-size_t FindCRLF(const butil::IOBuf& buf, size_t max_scan_size) {
-    butil::IOBufBytesIterator it(buf);
-    bool prev_cr = false;
-    size_t pos = 0;
-    while (it && pos < max_scan_size) {
-        const char c = static_cast<char>(*it);
-        if (prev_cr && c == '\n') {
-            return pos - 1;
-        }
-        prev_cr = (c == '\r');
-        ++it;
-        ++pos;
-    }
-    return CRLF_NOT_FOUND;
-}
-
 } // namespace
 
 namespace brpc {
@@ -454,7 +438,9 @@ butil::Status RedisCommandByComponents(butil::IOBuf* output,
 RedisCommandParser::RedisCommandParser()
     : _parsing_array(false)
     , _length(0)
-    , _index(0) {}
+    , _index(0) {
+    _line_scan.Reset();
+}
 
 size_t RedisCommandParser::ParsedArgsSize() {
     return _args.size();
@@ -514,7 +500,7 @@ RedisCommandConsumeState RedisCommandParser::ConsumeImpl(butil::IOBuf& buf,
         const size_t max_inline_size =
             FLAGS_redis_max_allocation_size > 0 ?
             static_cast<size_t>(FLAGS_redis_max_allocation_size) : 0;
-        const size_t crlf_pos = FindCRLF(buf, max_inline_size + 2);
+        const size_t crlf_pos = _line_scan.FindCRLF(buf, max_inline_size + 2);
         if (crlf_pos == CRLF_NOT_FOUND) {
             if (buf.size() <= max_inline_size) {
                 *err = PARSE_ERROR_NOT_ENOUGH_DATA;
@@ -696,6 +682,7 @@ RedisCommandConsumeState RedisCommandParser::ConsumeImpl(butil::IOBuf& buf,
 }
 
 void RedisCommandParser::Reset() {
+    _line_scan.Reset();
     _parsing_array = false;
     _length = 0;
     _index = 0;

src/brpc/redis_command.h

@@ -26,6 +26,7 @@
 #include "butil/status.h"
 #include "butil/arena.h"
 #include "brpc/parse_result.h"
+#include "brpc/details/redis_line_scanner.h"
 
 namespace brpc {
 
@@ -73,6 +74,7 @@ class RedisCommandParser {
                                          butil::Arena* arena,
                                          ParseError* err);
 
+    detail::RedisLineScanState _line_scan;
     bool _parsing_array;            // if the parser has met array indicator '*'
     int _length;                    // array length
     int _index;                     // current parsing array index

src/brpc/redis_reply.cpp

@@ -150,8 +150,15 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& buf, int depth,
     switch (fc) {
     case '-':   // Error          "-<message>\r\n"
     case '+': { // Simple String  "+<string>\r\n"
-        butil::IOBuf str;
-        if (buf.cut_until(&str, "\r\n") != 0) {
+        if (_length != kScanningLine) {
+            _type = REDIS_REPLY_NIL;
+            _length = kScanningLine;
+            _data.line = {};
+        }
+        const size_t max_scan = FLAGS_redis_max_allocation_size >= 0
+            ? static_cast<size_t>(FLAGS_redis_max_allocation_size) + 3 : 0;
+        const size_t crlf_pos = _data.line.FindCRLF(buf, max_scan);
+        if (crlf_pos == static_cast<size_t>(-1)) {
             const size_t len = buf.size();
             if (len > std::numeric_limits<uint32_t>::max()) {
                 LOG(ERROR) << "simple string is too long! max length=2^32-1,"
@@ -172,6 +179,9 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& buf, int depth,
             }
             return PARSE_ERROR_NOT_ENOUGH_DATA;
         }
+        butil::IOBuf str;
+        buf.cutn(&str, crlf_pos);
+        buf.pop_front(2);
         const size_t len = str.size() - 1;
         if (FLAGS_redis_max_allocation_size < 0 ||
             len > (size_t)FLAGS_redis_max_allocation_size) {

src/brpc/redis_reply.h

@@ -23,6 +23,7 @@
 #include "butil/strings/string_piece.h"   // butil::StringPiece
 #include "butil/arena.h"                  // butil::Arena
 #include "butil/logging.h"                // CHECK
+#include "brpc/details/redis_line_scanner.h"
 #include "parse_result.h"                 // ParseError
 
 
@@ -141,6 +142,7 @@ class RedisReply {
 
 private:
     static const int npos;
+    static const int kScanningLine = -2;
 
     // RedisReply does not own the memory of fields, copying must be done
     // by calling CopyFrom[Different|Same]Arena.
@@ -152,7 +154,8 @@ class RedisReply {
     void SetStringImpl(const butil::StringPiece& str, RedisReplyType type);
 
     RedisReplyType _type;
-    int _length;  // length of short_str/long_str, count of replies
+    // Length of strings/arrays, or kScanningLine during a partial simple line.
+    int _length;
     union {
         int64_t integer;
         char short_str[16];
@@ -165,6 +168,8 @@ class RedisReply {
             uint32_t used_bytes;
             RedisReply* replies;
         } array;
+        // Reuse the union while scanning so the reply does not grow.
+        detail::RedisLineScanState line;
         uint64_t padding[2]; // For swapping, must cover all bytes.
     } _data;
     butil::Arena* _arena;

test/brpc_redis_unittest.cpp

@@ -139,6 +139,75 @@ class RedisTest : public testing::Test {
     void TearDown() {}
 };
 
+TEST(RedisLineScanTest, fragmented_lines_preserve_input_until_complete) {
+    butil::Arena arena;
+    brpc::RedisCommandParser parser;
+    std::vector<butil::StringPiece> args;
+    butil::IOBuf command;
+    const std::string text = "GET payload\r\n";
+    for (size_t i = 0; i < text.size(); ++i) {
+        command.append(&text[i], 1);
+        if (i + 1 < text.size()) {
+            const std::string saved = command.to_string();
+            EXPECT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA,
+                      parser.Consume(command, &args, &arena));
+            EXPECT_EQ(saved, command.to_string());
+        } else {
+            ASSERT_EQ(brpc::PARSE_OK, parser.Consume(command, &args, &arena));
+        }
+    }
+    ASSERT_EQ(2u, args.size());
+    EXPECT_EQ("get", args[0]);
+    EXPECT_EQ("payload", args[1]);
+    EXPECT_TRUE(command.empty());
+    command.append("PING\r\n");
+    ASSERT_EQ(brpc::PARSE_OK, parser.Consume(command, &args, &arena));
+    ASSERT_EQ(1u, args.size());
+    EXPECT_EQ("ping", args[0]);
+
+    for (const std::string& line : {std::string("+payload\r\n"),
+                                    std::string("-error\r\n"),
+                                    std::string("*2\r\n+one\r\n-two\r\n")}) {
+        brpc::RedisReply reply(&arena);
+        butil::IOBuf buf;
+        for (size_t i = 0; i < line.size(); ++i) {
+            buf.append(&line[i], 1);
+            const brpc::ParseError result = reply.ConsumePartialIOBuf(buf);
+            EXPECT_EQ(i + 1 == line.size() ? brpc::PARSE_OK
+                                          : brpc::PARSE_ERROR_NOT_ENOUGH_DATA,
+                      result);
+        }
+        EXPECT_TRUE(buf.empty());
+        if (line[0] == '*') {
+            ASSERT_EQ(2u, reply.size());
+            EXPECT_EQ("one", reply[0].data());
+            EXPECT_STREQ("two", reply[1].error_message());
+        } else if (line[0] == '-') {
+            EXPECT_EQ(line.substr(1, line.size() - 3), reply.error_message());
+        } else {
+            EXPECT_EQ(line.substr(1, line.size() - 3), reply.data());
+        }
+    }
+}
+
+TEST(RedisLineScanTest, block_boundaries_limits_and_resume) {
+    brpc::detail::RedisLineScanState scanner;
+    scanner.Reset();
+    butil::IOBuf buf;
+    buf.append("first\r");
+    EXPECT_EQ(static_cast<size_t>(-1), scanner.FindCRLF(buf, 100));
+    EXPECT_EQ(static_cast<size_t>(-1), scanner.FindCRLF(buf, 100));
+    char* tail = new char[2]{'\n', 'x'};
+    buf.append_user_data(tail, 2, [](void* p) { delete[] static_cast<char*>(p); });
+    EXPECT_EQ(5u, scanner.FindCRLF(buf, 100));
+    scanner.Reset();
+    EXPECT_EQ(static_cast<size_t>(-1), scanner.FindCRLF(buf, 6));
+    EXPECT_EQ(5u, scanner.FindCRLF(buf, 7));
+    scanner.Reset();
+    EXPECT_EQ(static_cast<size_t>(-1), scanner.FindCRLF(buf, 3));
+    EXPECT_EQ(5u, scanner.FindCRLF(buf, 7));
+}
+
 TEST(RedisCommandFormatTest, wide_numeric_conversion) {
     butil::IOBuf buf;
     ASSERT_TRUE(brpc::RedisCommandFormat(&buf, "SET key %100d tail %d", 7, 9).ok());