Skip to content

Commit

Permalink
chore: improve interface for redis parsing
Browse files Browse the repository at this point in the history 
eliminate cases where we return INPUT_PENDING but do not consume the whole string.
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]>
  • Loading branch information
@romange
romange committed Nov 16, 2024
1 parent 8e3b8cc commit 4201591
Show file tree
Hide file tree
Showing 4 changed files with 125 additions and 48 deletions.
2 changes: 1 addition & 1 deletion src/facade/dragonfly_connection.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1124,7 +1124,7 @@ Connection::ParserStatus Connection::ParseRedis() {
DispatchSingle(has_more, dispatch_sync, dispatch_async);
}
io_buf_.ConsumeInput(consumed);
} while (RedisParser::OK == result && !reply_builder_->GetError());
} while (RedisParser::OK == result && io_buf_.InputLen() > 0 && !reply_builder_->GetError());

parser_error_ = result;
if (result == RedisParser::OK)
Expand Down
127 changes: 93 additions & 34 deletions src/facade/redis_parser.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,20 @@ namespace facade {
using namespace std;

auto RedisParser::Parse(Buffer str, uint32_t* consumed, RespExpr::Vec* res) -> Result {
DCHECK(!str.empty());
*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_)
Expand All @@ -38,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());
Expand All @@ -65,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());

Expand Down Expand Up @@ -212,19 +217,36 @@ auto RedisParser::ParseInline(Buffer str) -> ResultConsumed {
auto RedisParser::ParseLen(Buffer str, int64_t* res) -> ResultConsumed {
DCHECK(!str.empty());

char* s = reinterpret_cast<char*>(str.data() + 1);
char* pos = reinterpret_cast<char*>(memchr(s, '\n', str.size() - 1));
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') {
return {BAD_ARRAYLEN, 0};
}

bool success = absl::SimpleAtoi(std::string_view{s, size_t(pos - s - 1)}, res);
return ResultConsumed{success ? OK : BAD_ARRAYLEN, (pos - s) + 2};
// 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);
return ResultConsumed{success ? OK : BAD_ARRAYLEN, consumed};
}

auto RedisParser::ConsumeArrayLen(Buffer str) -> ResultConsumed {
Expand Down Expand Up @@ -288,7 +310,15 @@ 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 (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 == '$') {
int64_t len;
Expand Down Expand Up @@ -321,12 +351,22 @@ auto RedisParser::ParseArg(Buffer str) -> ResultConsumed {
}

if (c == '_') { // Resp3 NIL
// TODO: Do we need to validate that str[1:2] == "\r\n"?
// '_','\r','\n'
DCHECK_GE(small_len_ + str.size(), 3u);
DCHECK_LT(small_len_, 3);

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};
}

cached_expr_->emplace_back(RespExpr::NIL);
cached_expr_->back().u = Buffer{};
HandleFinishArg();
return {OK, 3}; // // '_','\r','\n'
return {OK, consumed};
}

if (c == '*') {
Expand Down Expand Up @@ -390,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_) {
Expand All @@ -415,26 +475,24 @@ auto RedisParser::ConsumeBulk(Buffer str) -> ResultConsumed {
return {INPUT_PENDING, consumed};
}

if (str.size() >= 32) {
DCHECK(bulk_len_);
size_t len = std::min<size_t>(str.size(), bulk_len_);
DCHECK(bulk_len_);
size_t len = std::min<size_t>(str.size(), bulk_len_);

if (is_broken_token_) {
memcpy(bulk_str.end(), str.data(), len);
bulk_str = Buffer{bulk_str.data(), bulk_str.size() + len};
DVLOG(1) << "Extending bulk stash to size " << bulk_str.size();
} else {
DVLOG(1) << "New bulk stash size " << bulk_len_;
vector<uint8_t> nb(bulk_len_);
memcpy(nb.data(), str.data(), len);
bulk_str = Buffer{nb.data(), len};
buf_stash_.emplace_back(std::move(nb));
is_broken_token_ = true;
cached_expr_->back().has_support = true;
}
consumed = len;
bulk_len_ -= len;
if (is_broken_token_) {
memcpy(bulk_str.end(), str.data(), len);
bulk_str = Buffer{bulk_str.data(), bulk_str.size() + len};
DVLOG(1) << "Extending bulk stash to size " << bulk_str.size();
} else {
DVLOG(1) << "New bulk stash size " << bulk_len_;
vector<uint8_t> nb(bulk_len_);
memcpy(nb.data(), str.data(), len);
bulk_str = Buffer{nb.data(), len};
buf_stash_.emplace_back(std::move(nb));
is_broken_token_ = true;
cached_expr_->back().has_support = true;
}
consumed = len;
bulk_len_ -= len;

return {INPUT_PENDING, consumed};
}
Expand All @@ -457,6 +515,7 @@ void RedisParser::HandleFinishArg() {
}
cached_expr_ = parse_stack_.back().second;
}
small_len_ = 0;
}

void RedisParser::ExtendLastString(Buffer str) {
Expand Down
4 changes: 2 additions & 2 deletions src/facade/redis_parser.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
*/

Expand Down Expand Up @@ -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;
Expand All @@ -114,6 +113,7 @@ class RedisParser {

using Blob = std::vector<uint8_t>;
std::vector<Blob> buf_stash_;
char small_buf_[32];
};

} // namespace facade
40 changes: 29 additions & 11 deletions src/facade/redis_parser_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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_);
Expand All @@ -125,9 +125,9 @@ TEST_F(RedisParserTest, Multi2) {

TEST_F(RedisParserTest, Multi3) {
const char kFirst[] = "*3\r\n$3\r\nSET\r\n$16\r\nkey:";
const char kSecond[] = "key:000002273458\r\n$3\r\nVXK";
const char kSecond[] = "000002273458\r\n$3\r\nVXK";
ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(kFirst));
ASSERT_EQ(strlen(kFirst) - 4, consumed_);
ASSERT_EQ(strlen(kFirst), consumed_);
ASSERT_EQ(RedisParser::INPUT_PENDING, Parse(kSecond));
ASSERT_EQ(strlen(kSecond), consumed_);
ASSERT_EQ(RedisParser::OK, Parse("\r\n*3\r\n$3\r\nSET"));
Expand All @@ -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)));

Expand All @@ -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) {
Expand All @@ -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"));
}

Expand Down Expand Up @@ -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

0 comments on commit 4201591

Please sign in to comment.