From 2edc3189181dae4177d4703127890d1d5f8991b6 Mon Sep 17 00:00:00 2001 From: Wang Xiaofeng Date: Sun, 25 Oct 2026 23:57:16 +0800 Subject: [PATCH] Preserve Redis line scan progress across reads --- src/brpc/details/redis_line_scanner.h | 72 +++++++++++++++++++++++++++ src/brpc/redis_command.cpp | 23 ++------- src/brpc/redis_command.h | 2 + src/brpc/redis_reply.cpp | 14 +++++- src/brpc/redis_reply.h | 7 ++- test/brpc_redis_unittest.cpp | 69 +++++++++++++++++++++++++ 6 files changed, 166 insertions(+), 21 deletions(-) create mode 100644 src/brpc/details/redis_line_scanner.h diff --git a/src/brpc/details/redis_line_scanner.h b/src/brpc/details/redis_line_scanner.h new file mode 100644 index 0000000000..a574a1049d --- /dev/null +++ b/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(-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 diff --git a/src/brpc/redis_command.cpp b/src/brpc/redis_command.cpp index 5f30d04968..845493f74d 100644 --- a/src/brpc/redis_command.cpp +++ b/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(*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(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; diff --git a/src/brpc/redis_command.h b/src/brpc/redis_command.h index 25bd8859fd..fbec01e4bd 100644 --- a/src/brpc/redis_command.h +++ b/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 diff --git a/src/brpc/redis_reply.cpp b/src/brpc/redis_reply.cpp index ee77b3e6f3..fed3e71624 100644 --- a/src/brpc/redis_reply.cpp +++ b/src/brpc/redis_reply.cpp @@ -150,8 +150,15 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& buf, int depth, switch (fc) { case '-': // Error "-\r\n" case '+': { // Simple 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(FLAGS_redis_max_allocation_size) + 3 : 0; + const size_t crlf_pos = _data.line.FindCRLF(buf, max_scan); + if (crlf_pos == static_cast(-1)) { const size_t len = buf.size(); if (len > std::numeric_limits::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) { diff --git a/src/brpc/redis_reply.h b/src/brpc/redis_reply.h index 725e184ef2..598a0913ae 100644 --- a/src/brpc/redis_reply.h +++ b/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; diff --git a/test/brpc_redis_unittest.cpp b/test/brpc_redis_unittest.cpp index 7e8d0be019..0b1b67bf31 100644 --- a/test/brpc_redis_unittest.cpp +++ b/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 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(-1), scanner.FindCRLF(buf, 100)); + EXPECT_EQ(static_cast(-1), scanner.FindCRLF(buf, 100)); + char* tail = new char[2]{'\n', 'x'}; + buf.append_user_data(tail, 2, [](void* p) { delete[] static_cast(p); }); + EXPECT_EQ(5u, scanner.FindCRLF(buf, 100)); + scanner.Reset(); + EXPECT_EQ(static_cast(-1), scanner.FindCRLF(buf, 6)); + EXPECT_EQ(5u, scanner.FindCRLF(buf, 7)); + scanner.Reset(); + EXPECT_EQ(static_cast(-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());