Implement flush_all feature

mckelvin · mckelvin · commit ae279e6b0475 · 2019-07-31T22:54:32.000+08:00
diff --git a/.gitignore b/.gitignore
@@ -14,6 +14,7 @@ venv/
 bench_venv/
 _client.cpp
 _client.so
+_client.*.so
 .tox/
 
 .cache/
diff --git a/ext/gtest/CMakeLists.txt.in b/ext/gtest/CMakeLists.txt.in
@@ -5,7 +5,7 @@ include(ExternalProject)
 
 ExternalProject_Add(googletest
     GIT_REPOSITORY https://github.com/google/googletest.git
-    GIT_TAG master
+    GIT_TAG release-1.8.0
     SOURCE_DIR "${CMAKE_BINARY_DIR}/googletest-src"
     BINARY_DIR "${CMAKE_BINARY_DIR}/googletest-build"
     CMAKE_ARGS -DCMAKE_ARCHIVE_OUTPUT_DIRECTORY_DEBUG:PATH=DebugLibs
diff --git a/include/Client.h b/include/Client.h
@@ -50,6 +50,7 @@ DECL_RETRIEVAL_CMD(gets)
   err_code_t version(broadcast_result_t** results, size_t* nHosts);
   err_code_t quit();
   err_code_t stats(broadcast_result_t** results, size_t* nHosts);
+  err_code_t flushAll(broadcast_result_t** results, size_t* nHosts);
 
   // touch
   err_code_t touch(const char* const* keys, const size_t* keyLens,
@@ -65,18 +66,23 @@ DECL_RETRIEVAL_CMD(gets)
            const bool noreply,
            unsigned_result_t** result, size_t* nResults);
 
+  inline void toggleFlushAllFeature(bool enabled) {
+    m_flushAllEnabled = enabled;
+  }
   void _sleep(uint32_t seconds); // check GIL in Python
 
  protected:
   void collectRetrievalResult(retrieval_result_t*** results, size_t* nResults);
   void collectMessageResult(message_result_t*** results, size_t* nResults);
-  void collectBroadcastResult(broadcast_result_t** results, size_t* nHosts);
+  void collectBroadcastResult(broadcast_result_t** results, size_t* nHosts, bool isFlushAll=false);
   void collectUnsignedResult(unsigned_result_t** results, size_t* nResults);
 
   std::vector<retrieval_result_t*> m_outRetrievalResultPtrs;
   std::vector<message_result_t*> m_outMessageResultPtrs;
   std::vector<broadcast_result_t> m_outBroadcastResultPtrs;
   std::vector<unsigned_result_t*> m_outUnsignedResultPtrs;
+
+  bool m_flushAllEnabled;
 };
 
 } // namespace mc
diff --git a/include/ConnectionPool.h b/include/ConnectionPool.h
@@ -41,7 +41,7 @@ class ConnectionPool {
 
   void collectRetrievalResult(std::vector<retrieval_result_t*>& results);
   void collectMessageResult(std::vector<message_result_t*>& results);
-  void collectBroadcastResult(std::vector<broadcast_result_t>& results);
+  void collectBroadcastResult(std::vector<broadcast_result_t>& results, bool isFlushAll=false);
   void collectUnsignedResult(std::vector<unsigned_result_t*>& results);
   void reset();
   void setPollTimeout(int timeout);
diff --git a/include/Export.h b/include/Export.h
@@ -44,14 +44,6 @@ typedef uint32_t flags_t;
 typedef uint64_t cas_unique_t;
 
 
-typedef struct {
-  char* host;
-  char** lines;
-  size_t* line_lens;
-  size_t len;
-} broadcast_result_t;
-
-
 typedef struct {
   char* key; // 8B
   char* data_block; // 8B
@@ -63,7 +55,8 @@ typedef struct {
 
 
 enum message_result_type {
-  MSG_EXISTS,
+  MSG_LIBMC_INVALID = -1,
+  MSG_EXISTS = 0,
   MSG_OK,
   MSG_STORED,
   MSG_NOT_STORED,
@@ -85,3 +78,16 @@ typedef struct {
   size_t key_len;
   uint64_t value;
 } unsigned_result_t;
+
+
+// For flush_all command, we need to specify
+// {host} and {msg_type},
+// for other broadcast commands, we need to specify
+// all fields except {msg_type}
+typedef struct {
+  char* host;
+  char** lines;
+  size_t* line_lens;
+  size_t len;
+  enum message_result_type msg_type;  // for flush_all command
+} broadcast_result_t;
diff --git a/include/Keywords.h b/include/Keywords.h
@@ -31,6 +31,7 @@ static const char k_NOREPLY[] = " noreply";
 
 static const char kVERSION[] = "version";
 static const char kSTATS[] = "stats";
+static const char kFLUSHALL[] = "flush_all";
 
 static const char kQUIT[] = "quit";
 
diff --git a/include/c_client.h b/include/c_client.h
@@ -63,6 +63,8 @@ extern "C" {
   void client_destroy_unsigned_result(void* client);
 
   err_code_t client_stats(void* client, broadcast_result_t** results, size_t* n_servers);
+  void client_toggle_flush_all_feature(void* client, bool enabled);
+  err_code_t client_flush_all(void* client, broadcast_result_t** results, size_t* n_servers);
   err_code_t client_quit(void* client);
 
   const char* err_code_to_string(err_code_t err);
diff --git a/libmc/__init__.py b/libmc/__init__.py
@@ -28,10 +28,10 @@
 )
 
 __VERSION__ = "1.2.0"
-__version__ = "v1.2.1"
+__version__ = "v1.2.1-4-gc3c6500"
 __author__ = "mckelvin"
 __email__ = "mckelvin@users.noreply.github.com"
-__date__ = "Wed Mar 13 14:40:14 2019 +0800"
+__date__ = "Tue Jul 30 14:06:44 2019 +0800"
 
 
 class Client(PyClient):
diff --git a/libmc/_client.pyx b/libmc/_client.pyx
@@ -72,13 +72,8 @@ cdef extern from "Export.h":
         uint32_t bytes
         cas_unique_t cas_unique
 
-    ctypedef struct broadcast_result_t:
-        char* host
-        char** lines
-        size_t* line_lens
-        size_t len
-
     ctypedef enum message_result_type:
+        MSG_LIBMC_INVALID
         MSG_EXISTS
         MSG_OK
         MSG_STORED
@@ -109,6 +104,13 @@ cdef extern from "Export.h":
         size_t key_len
         uint64_t value
 
+    ctypedef struct broadcast_result_t:
+        char* host
+        char** lines
+        size_t* line_lens
+        size_t len
+        message_result_type msg_type;
+
 
 cdef extern from "Client.h" namespace "douban::mc":
     cdef cppclass Client:
@@ -189,6 +191,8 @@ cdef extern from "Client.h" namespace "douban::mc":
         err_code_t version(broadcast_result_t** results, size_t* nHosts) nogil
         err_code_t quit() nogil
         err_code_t stats(broadcast_result_t** results, size_t* nHosts) nogil
+        err_code_t flushAll(broadcast_result_t** results, size_t* nHosts) nogil
+        void toggleFlushAllFeature(bool_t enabled)
         void destroyBroadcastResult() nogil
 
         err_code_t incr(
@@ -952,6 +956,30 @@ cdef class PyClient:
             self._imp.destroyBroadcastResult()
         return rv
 
+    def toggle_flush_all_feature(self, enabled):
+        self._imp.toggleFlushAllFeature(enabled)
+
+    def flush_all(self):
+        self._record_thread_ident()
+        cdef broadcast_result_t* rst = NULL
+        cdef size_t n = 0
+        with nogil:
+            self.last_error = self._imp.flushAll(&rst, &n)
+
+        rv = []
+        for i in range(n):
+            if rst[i].msg_type == MSG_OK:
+                rv.append(rst[i].host)
+
+        with nogil:
+            self._imp.destroyBroadcastResult()
+        if self.last_error == RET_PROGRAMMING_ERR:
+            raise RuntimeError(
+                "Please call toggle_flush_all_feature(true) first "
+                "to enable the flush_all feature."
+            )
+        return rv
+
     def quit(self):
         self._record_thread_ident()
         with nogil:
diff --git a/misc/update_version.sh b/misc/update_version.sh
@@ -7,8 +7,8 @@ for VERSIONING_FILE in $VERSIONING_FILES
 do
     TMPFILE=$VERSIONING_FILE".2"
     cat $VERSIONING_FILE | \
-        python $VERSIONING_SCRIPT --clean | \
-        python $VERSIONING_SCRIPT > $TMPFILE
+        python2 $VERSIONING_SCRIPT --clean | \
+        python2 $VERSIONING_SCRIPT > $TMPFILE
     mv $TMPFILE $VERSIONING_FILE
     git add $VERSIONING_FILE
 done
diff --git a/misc/versioning.py b/misc/versioning.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python2
 # from: https://gist.githubusercontent.com/pkrusche/7369262/raw/5bf2dc8afb88d3fdde7be6d16ee4290db6735f37/versioning.py
 
 """ Git Versioning Script
diff --git a/setup.py b/setup.py
@@ -2,6 +2,7 @@
 import os
 import re
 import sys
+import shlex
 import pkg_resources
 from glob import glob
 from setuptools import setup, Extension
@@ -56,7 +57,7 @@ class PyTest(TestCommand):
 
     def initialize_options(self):
         TestCommand.initialize_options(self)
-        self.pytest_args = ['tests']
+        self.pytest_args = 'tests'
 
     def finalize_options(self):
         TestCommand.finalize_options(self)
@@ -66,7 +67,7 @@ def finalize_options(self):
     def run_tests(self):
         #import here, cause outside the eggs aren't loaded
         import pytest
-        errno = pytest.main(self.pytest_args)
+        errno = pytest.main(shlex.split(self.pytest_args))
         sys.exit(errno)
 
 setup(
diff --git a/src/Client.cpp b/src/Client.cpp
@@ -8,7 +8,7 @@
 namespace douban {
 namespace mc {
 
-Client::Client() {
+Client::Client(): m_flushAllEnabled(false) {
 }
 
 
@@ -125,10 +125,10 @@ err_code_t Client::_delete(const char* const* keys, const size_t* keyLens,
 }
 
 
-void Client::collectBroadcastResult(broadcast_result_t** results, size_t* nHosts) {
+void Client::collectBroadcastResult(broadcast_result_t** results, size_t* nHosts, bool isFlushAll) {
   assert(m_outBroadcastResultPtrs.empty());
   *nHosts = m_nConns;
-  ConnectionPool::collectBroadcastResult(m_outBroadcastResultPtrs);
+  ConnectionPool::collectBroadcastResult(m_outBroadcastResultPtrs, isFlushAll);
   *results = &m_outBroadcastResultPtrs.front();
 }
 
@@ -166,6 +166,19 @@ err_code_t Client::stats(broadcast_result_t** results, size_t* nHosts) {
   return rv;
 }
 
+err_code_t Client::flushAll(broadcast_result_t** results, size_t* nHosts) {
+  if (!m_flushAllEnabled) {
+    *results = NULL;
+    *nHosts = 0;
+    return RET_PROGRAMMING_ERR;
+  }
+
+  broadcastCommand(keywords::kFLUSHALL, 9);
+  err_code_t rv = waitPoll();
+  collectBroadcastResult(results, nHosts, true);
+  return rv;
+}
+
 
 err_code_t Client::touch(const char* const* keys, const size_t* keyLens,
                    const exptime_t exptime, const bool noreply, size_t nItems,
diff --git a/src/ConnectionPool.cpp b/src/ConnectionPool.cpp
@@ -582,27 +582,39 @@ void ConnectionPool::collectMessageResult(std::vector<message_result_t*>& result
 }
 
 
-void ConnectionPool::collectBroadcastResult(std::vector<broadcast_result_t>& results) {
+void ConnectionPool::collectBroadcastResult(std::vector<broadcast_result_t>& results, bool isFlushAll) {
   results.resize(m_nConns);
   for (size_t i = 0; i < m_nConns; ++i) {
     Connection* conn = m_conns + i;
     broadcast_result_t* conn_result = &results[i];
     conn_result->host = const_cast<char*>(conn->name());
-    types::LineResultList* rst = conn->getLineResults();
-    conn_result->len = rst->size();
+    conn_result->lines = NULL;
+    conn_result->line_lens = NULL;
+    conn_result->len = 0;
+    conn_result->msg_type = MSG_LIBMC_INVALID;
+
+    if (isFlushAll) {
+      types::MessageResultList* rst = conn->getMessageResults();
+      if (rst->size() == 1) {
+        conn_result->msg_type = (rst->front()).type_;
+      } else {
+        conn_result->msg_type = MSG_LIBMC_INVALID;
+      }
+    } else {
+      types::LineResultList* rst = conn->getLineResults();
+      conn_result->len = rst->size();
 
-    if (conn_result->len == 0) {
-      conn_result->lines = NULL;
-      conn_result->line_lens = NULL;
-      continue;
-    }
-    conn_result->lines = new char*[conn_result->len];
-    conn_result->line_lens = new size_t[conn_result->len];
+      if (conn_result->len == 0) {
+        continue;
+      }
+      conn_result->lines = new char*[conn_result->len];
+      conn_result->line_lens = new size_t[conn_result->len];
 
-    int j = 0;
-    for (types::LineResultList::iterator it2 = rst->begin(); it2 != rst->end(); ++it2, ++j) {
-      types::LineResult* r1 = &(*it2);
-      conn_result->lines[j] = r1->inner(conn_result->line_lens[j]);
+      int j = 0;
+      for (types::LineResultList::iterator it2 = rst->begin(); it2 != rst->end(); ++it2, ++j) {
+        types::LineResult* r1 = &(*it2);
+        conn_result->lines[j] = r1->inner(conn_result->line_lens[j]);
+      }
     }
   }
 }
diff --git a/src/Parser.cpp b/src/Parser.cpp
@@ -35,13 +35,20 @@ void PacketParser::setMode(ParserMode md) {
 
 void PacketParser::processMessageResult(enum message_result_type tp) {
   m_messageResults.push_back(message_result_t());
-
   message_result_t* inner_rst = &m_messageResults.back();
-  struct iovec iov = m_requestKeys[m_requestKeyIdx];
-  ++m_requestKeyIdx;
   inner_rst->type_ = tp;
-  inner_rst->key = static_cast<char*>(iov.iov_base);
-  inner_rst->key_len = iov.iov_len;
+
+  // "OK\r\n" is not a key related response code,
+  // so we don't need to fill the key.
+  if (tp == MSG_OK) {
+    inner_rst->key = NULL;
+    inner_rst->key_len = 0;
+  } else {
+    struct iovec iov = m_requestKeys[m_requestKeyIdx];
+    ++m_requestKeyIdx;
+    inner_rst->key = static_cast<char*>(iov.iov_base);
+    inner_rst->key_len = iov.iov_len;
+  }
 }
 
 
@@ -338,6 +345,7 @@ int PacketParser::start_state(err_code_t& err) {
         // OK
         EXPECT_BYTES("OK\r\n", 4);
         processMessageResult(MSG_OK);
+        m_state = FSM_END;
       }
       break;
     case 'S':
diff --git a/src/c_client.cpp b/src/c_client.cpp
@@ -142,11 +142,20 @@ err_code_t client_stats(void* client, broadcast_result_t** results, size_t* n_se
   return c->stats(results, n_servers);
 }
 
+void client_toggle_flush_all_feature(void* client, bool enabled) {
+  douban::mc::Client* c = static_cast<Client*>(client);
+  return c->toggleFlushAllFeature(enabled);
+}
+err_code_t client_flush_all(void* client, broadcast_result_t** results, size_t* n_servers) {
+  douban::mc::Client* c = static_cast<Client*>(client);
+  return c->flushAll(results, n_servers);
+}
+
 err_code_t client_quit(void* client) {
   douban::mc::Client* c = static_cast<Client*>(client);
   return c->quit();
 }
 
 const char* err_code_to_string(err_code_t err) {
   return douban::mc::errCodeToString(err);
-}
+}
diff --git a/src/golibmc.go b/src/golibmc.go
diff --git a/src/golibmc_test.go b/src/golibmc_test.go
diff --git a/src/version.go b/src/version.go
diff --git a/tests/test_client.cpp b/tests/test_client.cpp
diff --git a/tests/test_client.py b/tests/test_client.py
