Merge pull request #1128 from liumh8/liuminghang/fix_redis_args

fix redis args

example/redis_c++/redis_server.cpp

@@ -64,14 +64,14 @@ public:
     explicit GetCommandHandler(RedisServiceImpl* rsimpl)
         : _rsimpl(rsimpl) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                          brpc::RedisReply* output,
-                                          bool /*flush_batched*/) override {
+    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                        brpc::RedisReply* output,
+                                        bool /*flush_batched*/) override {
         if (args.size() != 2ul) {
             output->FormatError("Expect 1 arg for 'get', actually %lu", args.size()-1);
             return brpc::REDIS_CMD_HANDLED;
         }
-        const std::string key(args[1]);
+        const std::string key(args[1].data(), args[1].size());
         std::string value;
         if (_rsimpl->Get(key, &value)) {
             output->SetString(value);
@@ -90,15 +90,15 @@ public:
     explicit SetCommandHandler(RedisServiceImpl* rsimpl)
         : _rsimpl(rsimpl) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                          brpc::RedisReply* output,
-                                          bool /*flush_batched*/) override {
+    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                        brpc::RedisReply* output,
+                                        bool /*flush_batched*/) override {
         if (args.size() != 3ul) {
             output->FormatError("Expect 2 args for 'set', actually %lu", args.size()-1);
             return brpc::REDIS_CMD_HANDLED;
         }
-        const std::string key(args[1]);
-        const std::string value(args[2]);
+        const std::string key(args[1].data(), args[1].size());
+        const std::string value(args[2].data(), args[2].size());
         _rsimpl->Set(key, value);
         output->SetStatus("OK");
         return brpc::REDIS_CMD_HANDLED;

src/brpc/policy/redis_protocol.cpp

@@ -76,13 +76,13 @@ public:
 };
 
 int ConsumeCommand(RedisConnContext* ctx,
-                   const std::vector<const char*>& commands,
+                   const std::vector<butil::StringPiece>& args,
                    bool flush_batched,
                    butil::IOBufAppender* appender) {
     RedisReply output(&ctx->arena);
     RedisCommandHandlerResult result = REDIS_CMD_HANDLED;
     if (ctx->transaction_handler) {
-        result = ctx->transaction_handler->Run(commands, &output, flush_batched);
+        result = ctx->transaction_handler->Run(args, &output, flush_batched);
         if (result == REDIS_CMD_HANDLED) {
             ctx->transaction_handler.reset(NULL);
         } else if (result == REDIS_CMD_BATCHED) {
@@ -90,13 +90,13 @@ int ConsumeCommand(RedisConnContext* ctx,
             return -1;
         }
     } else {
-        RedisCommandHandler* ch = ctx->redis_service->FindCommandHandler(commands[0]);
+        RedisCommandHandler* ch = ctx->redis_service->FindCommandHandler(args[0]);
         if (!ch) {
             char buf[64];
-            snprintf(buf, sizeof(buf), "ERR unknown command `%s`", commands[0]);
+            snprintf(buf, sizeof(buf), "ERR unknown command `%s`", args[0].as_string().c_str());
             output.SetError(buf);
         } else {
-            result = ch->Run(commands, &output, flush_batched);
+            result = ch->Run(args, &output, flush_batched);
             if (result == REDIS_CMD_CONTINUE) {
                 if (ctx->batched_size != 0) {
                     LOG(ERROR) << "CONTINUE should not be returned in a batched process.";
@@ -159,26 +159,26 @@ ParseResult ParseRedisMessage(butil::IOBuf* source, Socket* socket,
             ctx = new RedisConnContext(rs);
             socket->reset_parsing_context(ctx);
         }
-        std::vector<const char*> current_commands;
+        std::vector<butil::StringPiece> current_args;
         butil::IOBufAppender appender;
         ParseError err = PARSE_OK;
 
-        err = ctx->parser.Consume(*source, &current_commands, &ctx->arena);
+        err = ctx->parser.Consume(*source, &current_args, &ctx->arena);
         if (err != PARSE_OK) {
             return MakeParseError(err);
         }
         while (true) {
-            std::vector<const char*> next_commands;
-            err = ctx->parser.Consume(*source, &next_commands, &ctx->arena);
+            std::vector<butil::StringPiece> next_args;
+            err = ctx->parser.Consume(*source, &next_args, &ctx->arena);
             if (err != PARSE_OK) {
                 break;
             }
-            if (ConsumeCommand(ctx, current_commands, false, &appender) != 0) {
+            if (ConsumeCommand(ctx, current_args, false, &appender) != 0) {
                 return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
             }
-            current_commands.swap(next_commands);
+            current_args.swap(next_args);
         }
-        if (ConsumeCommand(ctx, current_commands,
+        if (ConsumeCommand(ctx, current_args,
                     true /*must be the last message*/, &appender) != 0) {
             return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
         }

src/brpc/redis.cpp

@@ -447,9 +447,8 @@ bool RedisService::AddCommandHandler(const std::string& name, RedisCommandHandle
     return true;
 }
  
-RedisCommandHandler* RedisService::FindCommandHandler(const std::string& name) const {
-    std::string lcname = StringToLowerASCII(name);
-    auto it = _command_map.find(lcname);
+RedisCommandHandler* RedisService::FindCommandHandler(const butil::StringPiece& name) const {
+    auto it = _command_map.find(name.as_string());
     if (it != _command_map.end()) {
         return it->second;
     }

src/brpc/redis.h

@@ -22,6 +22,7 @@
 #include <google/protobuf/message.h>
 #include <unordered_map>
 #include <memory>
+#include <list>
 #include "butil/iobuf.h"
 #include "butil/strings/string_piece.h"
 #include "butil/arena.h"
@@ -224,7 +225,7 @@ public:
     bool AddCommandHandler(const std::string& name, RedisCommandHandler* handler);
 
     // This function should not be touched by user and used by brpc deverloper only.
-    RedisCommandHandler* FindCommandHandler(const std::string& name) const;
+    RedisCommandHandler* FindCommandHandler(const butil::StringPiece& name) const;
 
 private:
     typedef std::unordered_map<std::string, RedisCommandHandler*> CommandMap;
@@ -260,7 +261,7 @@ public:
     // an start marker and brpc will call MultiTransactionHandler() to new a transaction
     // handler that all the following commands are sent to this tranction handler until
     // it returns REDIS_CMD_HANDLED. Read the comment below.
-    virtual RedisCommandHandlerResult Run(const std::vector<const char*>& args,
+    virtual RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
                                           brpc::RedisReply* output,
                                           bool flush_batched) = 0;
 

src/brpc/redis_command.cpp

@@ -362,7 +362,7 @@ RedisCommandParser::RedisCommandParser()
     , _index(0) {}
 
 ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
-                                       std::vector<const char*>* commands,
+                                       std::vector<butil::StringPiece>* args,
                                        butil::Arena* arena) {
     const char* pfc = (const char*)buf.fetch1();
     if (pfc == NULL) {
@@ -389,7 +389,7 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
         LOG(ERROR) << '`' << intbuf + 1 << "' is not a valid 64-bit decimal";
         return PARSE_ERROR_ABSOLUTELY_WRONG;
     }
-    if (value <= 0) {
+    if (value < 0) {
         LOG(ERROR) << "Invalid len=" << value << " in redis command";
         return PARSE_ERROR_ABSOLUTELY_WRONG;
     }
@@ -398,8 +398,8 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
         _parsing_array = true;
         _length = value;
         _index = 0;
-        _commands.resize(value);
-        return Consume(buf, commands, arena);
+        _args.resize(value);
+        return Consume(buf, args, arena);
     }
     CHECK(_index < _length) << "a complete command has been parsed. "
             "impl of RedisCommandParser::Parse is buggy";
@@ -420,7 +420,7 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
     char* d = (char*)arena->allocate((len/8 + 1) * 8);
     buf.cutn(d, len);
     d[len] = '\0';
-    _commands[_index] = d;
+    _args[_index].set(d, len);
     if (_index == 0) {
         // convert it to lowercase when it is command name
         for (int i = 0; i < len; ++i) {
@@ -434,9 +434,9 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
         return PARSE_ERROR_ABSOLUTELY_WRONG;
     }
     if (++_index < _length) {
-        return Consume(buf, commands, arena);
+        return Consume(buf, args, arena);
     }
-    commands->swap(_commands);
+    args->swap(_args);
     Reset();
     return PARSE_OK;
 }
@@ -445,7 +445,7 @@ void RedisCommandParser::Reset() {
     _parsing_array = false;
     _length = 0;
     _index = 0;
-    _commands.clear();
+    _args.clear();
 }
 
 } // namespace brpc

src/brpc/redis_command.h

@@ -48,9 +48,9 @@ public:
     RedisCommandParser();
 
     // Parse raw message from `buf'. Return PARSE_OK and set the parsed command
-    // to `commands' and length to `len' if successful. Memory of commands are
-    // allocated in `arena'.
-    ParseError Consume(butil::IOBuf& buf, std::vector<const char*>* commands,
+    // to `args' and length to `len' if successful. Memory of args are allocated 
+    // in `arena'.
+    ParseError Consume(butil::IOBuf& buf, std::vector<butil::StringPiece>* args,
                        butil::Arena* arena);
 
 private:
@@ -60,7 +60,7 @@ private:
     bool _parsing_array;            // if the parser has met array indicator '*'
     int _length;                    // array length
     int _index;                     // current parsing array index
-    std::vector<const char*> _commands;  // parsed command string
+    std::vector<butil::StringPiece> _args;  // parsed command string
 };
 
 } // namespace brpc

test/brpc_redis_unittest.cpp

@@ -128,7 +128,7 @@ void AssertReplyEqual(const brpc::RedisReply& reply1,
         // fall through
     case brpc::REDIS_REPLY_STATUS:
         ASSERT_NE(reply1.c_str(), reply2.c_str()); // from different arena
-        ASSERT_STREQ(reply1.c_str(), reply2.c_str());
+        ASSERT_EQ(reply1.data(), reply2.data());
         break;
     case brpc::REDIS_REPLY_ERROR:
         ASSERT_NE(reply1.error_message(), reply2.error_message()); // from different arena
@@ -550,13 +550,13 @@ TEST_F(RedisTest, quote_and_escape) {
     request.Clear();
 }
 
-std::string GetCompleteCommand(const std::vector<const char*>& commands) {
+std::string GetCompleteCommand(const std::vector<butil::StringPiece>& commands) {
 	std::string res;
     for (int i = 0; i < (int)commands.size(); ++i) {
         if (i != 0) {
             res.push_back(' ');
         }
-        res.append(commands[i]);
+        res.append(commands[i].data(), commands[i].size());
     }
     return res;
 }
@@ -565,7 +565,7 @@ std::string GetCompleteCommand(const std::vector<const char*>& commands) {
 TEST_F(RedisTest, command_parser) {
     brpc::RedisCommandParser parser;
     butil::IOBuf buf;
-    std::vector<const char*> command_out;
+    std::vector<butil::StringPiece> command_out;
     butil::Arena arena;
     {
         // parse from whole command
@@ -573,7 +573,7 @@ TEST_F(RedisTest, command_parser) {
         ASSERT_TRUE(brpc::RedisCommandNoFormat(&buf, command.c_str()).ok());
         ASSERT_EQ(brpc::PARSE_OK, parser.Consume(buf, &command_out, &arena));
         ASSERT_TRUE(buf.empty());
-        ASSERT_STREQ(command.c_str(), GetCompleteCommand(command_out).c_str());
+        ASSERT_EQ(command, GetCompleteCommand(command_out));
     }
     {
         // simulate parsing from network
@@ -593,7 +593,7 @@ TEST_F(RedisTest, command_parser) {
                 }
             }
             ASSERT_TRUE(buf.empty());
-            ASSERT_STREQ(GetCompleteCommand(command_out).c_str(), "set abc def");
+            ASSERT_EQ(GetCompleteCommand(command_out), "set abc def");
         }
     }
     {
@@ -812,19 +812,19 @@ public:
     RedisServiceImpl()
         : _batch_count(0) {}
 
-    brpc::RedisCommandHandlerResult OnBatched(const std::vector<const char*> args,
+    brpc::RedisCommandHandlerResult OnBatched(const std::vector<butil::StringPiece>& args,
                    brpc::RedisReply* output, bool flush_batched) {
         if (_batched_command.empty() && flush_batched) {
-            if (strcmp(args[0], "set") == 0) {
-                DoSet(args[1], args[2], output);
-            } else if (strcmp(args[0], "get") == 0) {
-                DoGet(args[1], output);
+            if (args[0] == "set") {
+                DoSet(args[1].as_string(), args[2].as_string(), output);
+            } else if (args[0] == "get") {
+                DoGet(args[1].as_string(), output);
             }
             return brpc::REDIS_CMD_HANDLED;
         }
         std::vector<std::string> comm;
         for (int i = 0; i < (int)args.size(); ++i) {
-            comm.push_back(args[i]);
+            comm.push_back(args[i].as_string());
         }
         _batched_command.push_back(comm);
         if (flush_batched) {
@@ -869,9 +869,9 @@ public:
         : _rs(rs)
         , _batch_process(batch_process) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                          brpc::RedisReply* output,
-                                          bool flush_batched) {
+    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                        brpc::RedisReply* output,
+                                        bool flush_batched) {
         if (args.size() < 3) {
             output->SetError("ERR wrong number of arguments for 'set' command");
             return brpc::REDIS_CMD_HANDLED;
@@ -879,7 +879,7 @@ public:
         if (_batch_process) {
             return _rs->OnBatched(args, output, flush_batched);
         } else {
-            DoSet(args[1], args[2], output);
+            DoSet(args[1].as_string(), args[2].as_string(), output);
             return brpc::REDIS_CMD_HANDLED;
         }
     }
@@ -900,9 +900,9 @@ public:
         : _rs(rs)
         , _batch_process(batch_process) {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                          brpc::RedisReply* output,
-                                          bool flush_batched) {
+    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                        brpc::RedisReply* output,
+                                        bool flush_batched) {
         if (args.size() < 2) {
             output->SetError("ERR wrong number of arguments for 'get' command");
             return brpc::REDIS_CMD_HANDLED;
@@ -910,7 +910,7 @@ public:
         if (_batch_process) {
             return _rs->OnBatched(args, output, flush_batched);
         } else {
-            DoGet(args[1], output);
+            DoGet(args[1].as_string(), output);
             return brpc::REDIS_CMD_HANDLED;
         }
     }
@@ -933,17 +933,16 @@ class IncrCommandHandler : public brpc::RedisCommandHandler {
 public:
     IncrCommandHandler() {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                          brpc::RedisReply* output,
-                                          bool flush_batched) {
+    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                        brpc::RedisReply* output,
+                                        bool flush_batched) {
         if (args.size() < 2) {
             output->SetError("ERR wrong number of arguments for 'incr' command");
             return brpc::REDIS_CMD_HANDLED;
         }
-        const std::string& key = args[1];
         int64_t value;
         s_mutex.lock();
-        value = ++int_map[key];
+        value = ++int_map[args[1].as_string()];
         s_mutex.unlock();
         output->SetInteger(value);
         return brpc::REDIS_CMD_HANDLED;
@@ -994,6 +993,34 @@ TEST_F(RedisTest, server_sanity) {
     ASSERT_STREQ("value2", response.reply(5).c_str());
     ASSERT_EQ(brpc::REDIS_REPLY_ERROR, response.reply(6).type());
     ASSERT_TRUE(butil::StringPiece(response.reply(6).error_message()).starts_with("ERR unknown command"));
+
+    cntl.Reset(); 
+    request.Clear();
+    response.Clear();
+    std::string value3("value3");
+    value3.append(1, '\0');
+    value3.append(1, 'a');
+    std::vector<butil::StringPiece> pieces;
+    pieces.push_back("set");
+    pieces.push_back("key3");
+    pieces.push_back(value3);
+    ASSERT_TRUE(request.AddCommandByComponents(&pieces[0], pieces.size()));
+    ASSERT_TRUE(request.AddCommand("set key4 \"\""));
+    ASSERT_TRUE(request.AddCommand("get key3"));
+    ASSERT_TRUE(request.AddCommand("get key4"));
+    channel.CallMethod(NULL, &cntl, &request, &response, NULL);
+    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
+    ASSERT_EQ(4, response.reply_size());
+    ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(0).type());
+    ASSERT_STREQ("OK", response.reply(0).c_str());
+    ASSERT_EQ(brpc::REDIS_REPLY_STATUS, response.reply(1).type());
+    ASSERT_STREQ("OK", response.reply(1).c_str());
+    ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(2).type());
+    ASSERT_STREQ("value3", response.reply(2).c_str());
+    ASSERT_NE("value3", response.reply(2).data());
+    ASSERT_EQ(value3, response.reply(2).data());
+    ASSERT_EQ(brpc::REDIS_REPLY_STRING, response.reply(3).type());
+    ASSERT_EQ("", response.reply(3).data());
 }
 
 void* incr_thread(void* arg) {
@@ -1047,9 +1074,9 @@ class MultiCommandHandler : public brpc::RedisCommandHandler {
 public:
     MultiCommandHandler() {}
 
-    brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                          brpc::RedisReply* output,
-                                          bool flush_batched) {
+    brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                        brpc::RedisReply* output,
+                                        bool flush_batched) {
         output->SetStatus("OK");
         return brpc::REDIS_CMD_CONTINUE;
     }
@@ -1060,17 +1087,17 @@ public:
 
     class MultiTransactionHandler : public brpc::RedisCommandHandler {
     public:
-        brpc::RedisCommandHandlerResult Run(const std::vector<const char*>& args,
-                                              brpc::RedisReply* output,
-                                              bool flush_batched) {
-            if (strcmp(args[0], "multi") == 0) {
+        brpc::RedisCommandHandlerResult Run(const std::vector<butil::StringPiece>& args,
+                                            brpc::RedisReply* output,
+                                            bool flush_batched) {
+            if (args[0] == "multi") {
                 output->SetError("ERR duplicate multi");
                 return brpc::REDIS_CMD_CONTINUE;
             }
-            if (strcmp(args[0], "exec") != 0) {
+            if (args[0] != "exec") {
                 std::vector<std::string> comm;
                 for (int i = 0; i < (int)args.size(); ++i) {
-                    comm.push_back(args[i]);
+                    comm.push_back(args[i].as_string());
                 }
                 _commands.push_back(comm);
                 output->SetStatus("QUEUED");