From d790be2124509e0ba3099a751918dec4ca657921 Mon Sep 17 00:00:00 2001
From: Smilencer <527646889@qq.com>
Date: Fri, 24 Apr 2020 20:09:30 +0800
Subject: [PATCH] feat(util): add restore_read function in rpc_message (#442)

---
 include/dsn/cpp/rpc_stream.h       |  2 +-
 include/dsn/tool-api/rpc_message.h |  6 ++++++
 src/core/core/rpc_message.cpp      |  7 +++++++
 src/core/tests/rpc_message.cpp     | 12 ++++++++++++
 4 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/include/dsn/cpp/rpc_stream.h b/include/dsn/cpp/rpc_stream.h
index e7480acf25..ed9e81dddd 100644
--- a/include/dsn/cpp/rpc_stream.h
+++ b/include/dsn/cpp/rpc_stream.h
@@ -121,4 +121,4 @@ class rpc_write_stream : public binary_writer
     int _last_write_next_total_size;
 };
 typedef ::dsn::ref_ptr<rpc_write_stream> rpc_write_stream_ptr;
-}
+} // namespace dsn
diff --git a/include/dsn/tool-api/rpc_message.h b/include/dsn/tool-api/rpc_message.h
index d4c83774e3..f4cc86a796 100644
--- a/include/dsn/tool-api/rpc_message.h
+++ b/include/dsn/tool-api/rpc_message.h
@@ -209,6 +209,12 @@ class message_ex : public ref_counter,
     size_t body_size() { return (size_t)header->body_length; }
     DSN_API void *rw_ptr(size_t offset_begin);
 
+    // rpc_read_stream can read a msg many times by restore()
+    // rpc_read_stream stream1(msg)
+    // msg->restore_read()
+    // rpc_read_stream stream2(msg)
+    DSN_API void restore_read();
+
     bool is_backup_request() const { return header->context.u.is_backup_request; }
 
 private:
diff --git a/src/core/core/rpc_message.cpp b/src/core/core/rpc_message.cpp
index 1f48e8feff..6f09c19b41 100644
--- a/src/core/core/rpc_message.cpp
+++ b/src/core/core/rpc_message.cpp
@@ -520,6 +520,13 @@ void message_ex::read_commit(size_t size)
     this->_rw_committed = true;
 }
 
+void message_ex::restore_read()
+{
+    _rw_index = -1;
+    _rw_committed = true;
+    _rw_offset = 0;
+}
+
 void *message_ex::rw_ptr(size_t offset_begin)
 {
     // printf("%p %s\n", this, __FUNCTION__);
diff --git a/src/core/tests/rpc_message.cpp b/src/core/tests/rpc_message.cpp
index eabef10fe6..14bd15d304 100644
--- a/src/core/tests/rpc_message.cpp
+++ b/src/core/tests/rpc_message.cpp
@@ -36,6 +36,7 @@
 #include <dsn/utility/crc.h>
 #include <dsn/utility/transient_memory.h>
 #include <dsn/tool-api/rpc_message.h>
+#include <core/core/message_utils.cpp>
 #include <gtest/gtest.h>
 
 using namespace ::dsn;
@@ -183,3 +184,14 @@ TEST(core, message_ex)
         request->release_ref();
     }
 }
+
+TEST(rpc_message, restore_read)
+{
+    using namespace dsn;
+    configuration_query_by_index_request request, result;
+    message_ptr msg = from_thrift_request_to_received_message(request, RPC_CODE_FOR_TEST);
+    for (int i = 0; i < 10; i++) {
+        unmarshall(msg, result);
+        msg->restore_read();
+    }
+}