auto expand & shrink wbuf in RedisConn

Author: KernelMaker

pink/include/pink_define.h

@@ -76,6 +76,7 @@ enum RetCode {
  * define the redis protocol
  */
 #define REDIS_MAX_MESSAGE 67108864  // 64MB
+#define DEFAULT_WBUF_SIZE 262144 // 256KB
 #define REDIS_IOBUF_LEN 16384
 #define REDIS_REQ_INLINE 1
 #define REDIS_REQ_MULTIBULK 2

pink/include/redis_conn.h

@@ -25,7 +25,7 @@ class RedisConn: public PinkConn {
   virtual ~RedisConn();
   void ResetClient();
 
-  bool ExpandWbuf();
+  bool ExpandWbufTo(uint32_t new_size);
   uint32_t wbuf_size_;
 
   virtual ReadStatus GetRequest();

pink/src/redis_conn.cc

@@ -133,7 +133,7 @@ RedisConn::RedisConn(const int fd,
                      const std::string &ip_port,
                      ServerThread *thread)
     : PinkConn(fd, ip_port, thread),
-      wbuf_size_(REDIS_MAX_MESSAGE),
+      wbuf_size_(DEFAULT_WBUF_SIZE),
       wbuf_len_(0),
       last_read_pos_(-1),
       next_parse_pos_(0),
@@ -144,7 +144,7 @@ RedisConn::RedisConn(const int fd,
       is_overtake_(false),
       wbuf_pos_(0) {
   rbuf_ = reinterpret_cast<char*>(malloc(sizeof(char) * REDIS_MAX_MESSAGE));
-  wbuf_ = reinterpret_cast<char*>(malloc(sizeof(char) * REDIS_MAX_MESSAGE));
+  wbuf_ = reinterpret_cast<char*>(malloc(sizeof(char) * DEFAULT_WBUF_SIZE));
 }
 
 RedisConn::~RedisConn() {
@@ -282,14 +282,19 @@ void RedisConn::ResetClient() {
     bulk_len_ = -1;
 }
 
-bool RedisConn::ExpandWbuf() {
-  if (wbuf_size_ >= (uint32_t)(REDIS_MAX_MESSAGE * (uint32_t)32)) {
+bool RedisConn::ExpandWbufTo(uint32_t new_size) {
+  if (new_size <= wbuf_size_) {
+    return true;
+  }
+  void* new_wbuf = realloc(wbuf_, new_size);
+  if (new_wbuf != nullptr) {
+    wbuf_ = reinterpret_cast<char*>(new_wbuf);
+    wbuf_size_ = new_size;
+    return true;
+  } else {
     wbuf_pos_ = 0;
     return false;
   }
-  wbuf_size_ *= 2;
-  wbuf_ = reinterpret_cast<char*>(realloc(wbuf_, wbuf_size_));
-  return true;
 }
 
 ReadStatus RedisConn::ProcessInputBuffer() {
@@ -379,6 +384,12 @@ WriteStatus RedisConn::SendReply() {
     if (wbuf_pos_ == wbuf_len_) {
       wbuf_len_ = 0;
       wbuf_pos_ = 0;
+      if (wbuf_size_ > DEFAULT_WBUF_SIZE) {
+        free(wbuf_);
+        wbuf_ = reinterpret_cast<char*>(malloc(sizeof(char) *
+              DEFAULT_WBUF_SIZE));
+        wbuf_size_ = DEFAULT_WBUF_SIZE;
+      }
     }
   }
   if (nwritten == -1) {