grpc: reestablish connection and start stream (#31)

Shikugawa · web-flow · commit 754841ef0b16 · 2020-12-10T17:52:01.000+08:00
diff --git a/cpp2sky/internal/async_client.h b/cpp2sky/internal/async_client.h
@@ -18,6 +18,7 @@
 #include <grpcpp/grpcpp.h>
 
 #include <memory>
+#include <queue>
 
 using google::protobuf::Message;
 
@@ -59,6 +60,27 @@ class AsyncClient {
    * Peer address of current gRPC client..
    */
   virtual std::string peerAddress() = 0;
+
+  /**
+   * Drain pending messages
+   */
+  virtual void drainPendingMessages(
+      std::queue<RequestType>& pending_messages) = 0;
+
+  /**
+   * Reset stream if it is living.
+   */
+  virtual void resetStream() = 0;
+
+  /**
+   * Start stream if there is no living stream.
+   */
+  virtual void startStream() = 0;
+
+  /**
+   * The number of drained pending messages.
+   */
+  virtual size_t numOfMessages() = 0;
 };
 
 enum class Operation : uint8_t {
@@ -105,7 +127,8 @@ class AsyncStreamFactory {
    * Create async stream entity
    */
   virtual AsyncStreamPtr<RequestType> create(
-      AsyncClient<RequestType, ResponseType>* client) = 0;
+      AsyncClient<RequestType, ResponseType>* client,
+      std::queue<RequestType>& drained_messages) = 0;
 };
 
 template <class RequestType, class ResponseType>
diff --git a/source/grpc_async_client_impl.cc b/source/grpc_async_client_impl.cc
@@ -43,14 +43,23 @@ GrpcAsyncSegmentReporterClient::GrpcAsyncSegmentReporterClient(
     AsyncStreamFactory<TracerRequestType, TracerResponseType>& factory,
     std::shared_ptr<grpc::ChannelCredentials> cred, std::string address,
     std::string token)
-    : token_(token), address_(address), factory_(factory), cq_(cq) {
-  stub_ = std::make_unique<TracerStubImpl>(grpc::CreateChannel(address, cred));
-  stream_ = factory_.create(this);
-  stream_->startStream();
+    : token_(token),
+      factory_(factory),
+      cq_(cq),
+      channel_(grpc::CreateChannel(address, cred)) {
+  stub_ = std::make_unique<TracerStubImpl>(channel_);
+  startStream();
 }
 
 void GrpcAsyncSegmentReporterClient::sendMessage(TracerRequestType message) {
-  GPR_ASSERT(stream_ != nullptr);
+  if (!stream_) {
+    drained_messages_.emplace(message);
+    gpr_log(GPR_INFO,
+            "No active stream, inserted message into draining message queue. "
+            "pending message size: %ld",
+            drained_messages_.size());
+    return;
+  }
   stream_->sendMessage(message);
 }
 
@@ -64,16 +73,45 @@ GrpcAsyncSegmentReporterClient::createWriter(grpc::ClientContext* ctx,
   return stub_->createWriter(ctx, response, cq_, tag);
 }
 
+void GrpcAsyncSegmentReporterClient::startStream() {
+  resetStream();
+
+  // Try to establish connection.
+  channel_->GetState(true);
+  stream_ = factory_.create(this, drained_messages_);
+  stream_->startStream();
+}
+
+void GrpcAsyncSegmentReporterClient::drainPendingMessages(
+    std::queue<TracerRequestType>& pending_messages) {
+  const auto pending_messages_size = pending_messages.size();
+  while (!pending_messages.empty()) {
+    auto msg = pending_messages.front();
+    pending_messages.pop();
+    drained_messages_.emplace(msg);
+  }
+  gpr_log(GPR_INFO, "%ld pending messages drained.", pending_messages_size);
+}
+
 GrpcAsyncSegmentReporterStream::GrpcAsyncSegmentReporterStream(
-    AsyncClient<TracerRequestType, TracerResponseType>* client)
-    : client_(client) {}
+    AsyncClient<TracerRequestType, TracerResponseType>* client,
+    std::queue<TracerRequestType>& drained_messages)
+    : client_(client) {
+  const auto drained_messages_size = drained_messages.size();
+  while (!drained_messages.empty()) {
+    auto msg = drained_messages.front();
+    pending_messages_.emplace(msg);
+    drained_messages.pop();
+  }
+  gpr_log(GPR_INFO, "%ld drained messages inserted into pending messages.",
+          drained_messages_size);
+}
 
 GrpcAsyncSegmentReporterStream::~GrpcAsyncSegmentReporterStream() {
   {
     std::unique_lock<std::mutex> lck_(mux_);
     cond_.wait(lck_, [this] { return pending_messages_.empty(); });
   }
-
   ctx_.TryCancel();
   request_writer_->Finish(&status_, toTag(&finish_));
 }
@@ -132,6 +170,7 @@ bool GrpcAsyncSegmentReporterStream::handleOperation(Operation incoming_op) {
   } else if (state_ == Operation::Finished) {
     gpr_log(GPR_INFO, "Stream closed with http status: %d",
             grpcStatusToGenericHttpStatus(status_.error_code()));
+    client_->drainPendingMessages(pending_messages_);
     if (!status_.ok()) {
       gpr_log(GPR_ERROR, "%s", status_.error_message().c_str());
     }
@@ -141,11 +180,13 @@ bool GrpcAsyncSegmentReporterStream::handleOperation(Operation incoming_op) {
 }
 
 AsyncStreamPtr<TracerRequestType> GrpcAsyncSegmentReporterStreamFactory::create(
-    AsyncClient<TracerRequestType, TracerResponseType>* client) {
+    AsyncClient<TracerRequestType, TracerResponseType>* client,
+    std::queue<TracerRequestType>& drained_messages) {
   if (client == nullptr) {
     return nullptr;
   }
-  return std::make_shared<GrpcAsyncSegmentReporterStream>(client);
+  return std::make_shared<GrpcAsyncSegmentReporterStream>(client,
+                                                          drained_messages);
 }
 
 }  // namespace cpp2sky
diff --git a/source/grpc_async_client_impl.h b/source/grpc_async_client_impl.h
@@ -44,6 +44,8 @@ class TracerStubImpl final
   std::unique_ptr<TraceSegmentReportService::Stub> stub_;
 };
 
+class GrpcAsyncSegmentReporterStream;
+
 class GrpcAsyncSegmentReporterClient final
     : public AsyncClient<TracerRequestType, TracerResponseType> {
  public:
@@ -59,18 +61,28 @@ class GrpcAsyncSegmentReporterClient final
   std::unique_ptr<grpc::ClientAsyncWriter<TracerRequestType>> createWriter(
       grpc::ClientContext* ctx, TracerResponseType* response,
       void* tag) override;
+  void drainPendingMessages(
+      std::queue<TracerRequestType>& pending_messages) override;
+  void resetStream() override {
+    if (stream_) {
+      stream_.reset();
+      stream_ = nullptr;
+    }
+  }
+  void startStream() override;
+  size_t numOfMessages() override { return drained_messages_.size(); }
 
  private:
   std::string token_;
   std::string address_;
   AsyncStreamFactory<TracerRequestType, TracerResponseType>& factory_;
   TracerStubPtr<TracerRequestType, TracerResponseType> stub_;
   grpc::CompletionQueue* cq_;
+  std::shared_ptr<grpc::Channel> channel_;
   AsyncStreamPtr<TracerRequestType> stream_;
+  std::queue<TracerRequestType> drained_messages_;
 };
 
-class GrpcAsyncSegmentReporterStream;
-
 struct TaggedStream {
   Operation operation;
   GrpcAsyncSegmentReporterStream* stream;
@@ -83,7 +95,8 @@ class GrpcAsyncSegmentReporterStream final
     : public AsyncStream<TracerRequestType> {
  public:
   GrpcAsyncSegmentReporterStream(
-      AsyncClient<TracerRequestType, TracerResponseType>* client);
+      AsyncClient<TracerRequestType, TracerResponseType>* client,
+      std::queue<TracerRequestType>& drained_messages);
   ~GrpcAsyncSegmentReporterStream() override;
 
   // AsyncStream
@@ -115,7 +128,8 @@ class GrpcAsyncSegmentReporterStreamFactory final
  public:
   // AsyncStreamFactory
   AsyncStreamPtr<TracerRequestType> create(
-      AsyncClient<TracerRequestType, TracerResponseType>* client) override;
+      AsyncClient<TracerRequestType, TracerResponseType>* client,
+      std::queue<TracerRequestType>& drained_messages) override;
 };
 
 }  // namespace cpp2sky
diff --git a/source/tracer_impl.cc b/source/tracer_impl.cc
@@ -22,12 +22,12 @@ TracerImpl::TracerImpl(TracerConfig& config,
                        std::shared_ptr<grpc::ChannelCredentials> cred,
                        GrpcAsyncSegmentReporterStreamFactory& factory)
     : th_([this] { this->run(); }) {
-  client_ = new GrpcAsyncSegmentReporterClient(
+  client_ = std::make_unique<GrpcAsyncSegmentReporterClient>(
       &cq_, factory, cred, config.address(), config.token());
 }
 
 TracerImpl::~TracerImpl() {
-  delete client_;
+  client_.reset();
   th_.join();
   cq_.Shutdown();
 }
@@ -50,10 +50,11 @@ void TracerImpl::run() {
     }
     TaggedStream* t_stream = deTag(got_tag);
     if (!ok) {
+      client_->resetStream();
       continue;
     }
     if (!t_stream->stream->handleOperation(t_stream->operation)) {
-      return;
+      client_->startStream();
     }
   }
 }
diff --git a/source/tracer_impl.h b/source/tracer_impl.h
@@ -21,6 +21,9 @@
 
 namespace cpp2sky {
 
+using TracerRequestType = SegmentObject;
+using TracerResponseType = Commands;
+
 class TracerImpl : public Tracer {
  public:
   TracerImpl(TracerConfig& config,
@@ -33,7 +36,7 @@ class TracerImpl : public Tracer {
  private:
   void run();
 
-  GrpcAsyncSegmentReporterClient* client_;
+  AsyncClientPtr<TracerRequestType, TracerResponseType> client_;
   grpc::CompletionQueue cq_;
   std::thread th_;
 };
diff --git a/test/grpc_async_client_test.cc b/test/grpc_async_client_test.cc
@@ -28,7 +28,7 @@ using testing::_;
 class GrpcAsyncSegmentReporterClientTest : public testing::Test {
  public:
   GrpcAsyncSegmentReporterClientTest() {
-    EXPECT_CALL(factory_, create(_));
+    EXPECT_CALL(factory_, create(_, _));
     EXPECT_CALL(*stream_, startStream());
     client_ = std::make_unique<GrpcAsyncSegmentReporterClient>(
         &cq_, factory_, grpc::InsecureChannelCredentials(), address_, token_);
@@ -51,4 +51,14 @@ TEST_F(GrpcAsyncSegmentReporterClientTest, SendMessageTest) {
   client_->sendMessage(fake_message);
 }
 
+TEST_F(GrpcAsyncSegmentReporterClientTest, MessageDrainTest) {
+  std::queue<TracerRequestType> fake_pending_messages;
+  for (int i = 0; i < 3; ++i) {
+    fake_pending_messages.emplace(SegmentObject());
+  }
+  client_->drainPendingMessages(fake_pending_messages);
+  EXPECT_EQ(fake_pending_messages.size(), 0);
+  EXPECT_EQ(client_->numOfMessages(), 3);
+}
+
 }  // namespace cpp2sky
diff --git a/test/mocks.h b/test/mocks.h
@@ -38,15 +38,20 @@ class MockAsyncStream : public AsyncStream<RequestType> {
   MOCK_METHOD(void, sendMessage, (RequestType));
   MOCK_METHOD(std::string, peerAddress, ());
   MOCK_METHOD(bool, handleOperation, (Operation));
+  MOCK_METHOD(size_t, numOfMessages, ());
 };
 
 template <class RequestType, class ResponseType>
 class MockAsyncClient : public AsyncClient<RequestType, ResponseType> {
  public:
-  MOCK_METHOD(void, sendMessage, (Message&));
+  MOCK_METHOD(void, sendMessage, (RequestType));
   MOCK_METHOD(std::unique_ptr<grpc::ClientAsyncWriter<RequestType>>,
               createWriter, (grpc::ClientContext*, ResponseType*, void*));
   MOCK_METHOD(std::string, peerAddress, ());
+  MOCK_METHOD(void, drainPendingMessages, (std::queue<RequestType>&));
+  MOCK_METHOD(void, resetStream, ());
+  MOCK_METHOD(void, startStream, ());
+  MOCK_METHOD(size_t, numOfMessages, ());
 };
 
 template <class RequestType, class ResponseType>
@@ -55,9 +60,10 @@ class MockAsyncStreamFactory
  public:
   using AsyncClientType = AsyncClient<RequestType, ResponseType>;
   MockAsyncStreamFactory(AsyncStreamPtr<RequestType> stream) : stream_(stream) {
-    ON_CALL(*this, create(_)).WillByDefault(Return(stream_));
+    ON_CALL(*this, create(_, _)).WillByDefault(Return(stream_));
   }
-  MOCK_METHOD(AsyncStreamPtr<RequestType>, create, (AsyncClientType*));
+  MOCK_METHOD(AsyncStreamPtr<RequestType>, create,
+              (AsyncClientType*, std::queue<RequestType>&));
 
  private:
   AsyncStreamPtr<RequestType> stream_;

Original file line number	Diff line number	Diff line change
`@@ -22,12 +22,12 @@ TracerImpl::TracerImpl(TracerConfig& config,`
`22`	`22`	`std::shared_ptr<grpc::ChannelCredentials> cred,`
`23`	`23`	`GrpcAsyncSegmentReporterStreamFactory& factory)`
`24`	`24`	`: th_([this] { this->run(); }) {`
`25`		`- client_ = new GrpcAsyncSegmentReporterClient(`
	`25`	`+ client_ = std::make_unique<GrpcAsyncSegmentReporterClient>(`
`26`	`26`	`&cq_, factory, cred, config.address(), config.token());`
`27`	`27`	`}`
`28`	`28`
`29`	`29`	`TracerImpl::~TracerImpl() {`
`30`		`- delete client_;`
	`30`	`+ client_.reset();`
`31`	`31`	`th_.join();`
`32`	`32`	`cq_.Shutdown();`
`33`	`33`	`}`
`@@ -50,10 +50,11 @@ void TracerImpl::run() {`
`50`	`50`	`}`
`51`	`51`	`TaggedStream* t_stream = deTag(got_tag);`
`52`	`52`	`if (!ok) {`
	`53`	`+ client_->resetStream();`
`53`	`54`	`continue;`
`54`	`55`	`}`
`55`	`56`	`if (!t_stream->stream->handleOperation(t_stream->operation)) {`
`56`		`- return;`
	`57`	`+ client_->startStream();`
`57`	`58`	`}`
`58`	`59`	`}`
`59`	`60`	`}`