Add WaitForState

4 files changed, 57 insertions, 2 deletions

diff --git a/include/grpc++/channel_interface.h b/include/grpc++/channel_interface.h
index 65161275c6..335b6ccaae 100644
--- a/include/grpc++/channel_interface.h
+++ b/include/grpc++/channel_interface.h
@@ -69,10 +69,15 @@ class ChannelInterface : public CallHook,
                                    gpr_timespec deadline,
                                    CompletionQueue* cq, void* tag) = 0;
 
-  // Blocking wait for channel state change or deadline expires.
+  // Blocking wait for channel state change or deadline expiration.
   // GetState needs to called to get the current state.
   virtual bool WaitForStateChange(grpc_connectivity_state last_observed,
                                   gpr_timespec deadline) = 0;
+
+  // Blocking wait for target state or deadline expriration.
+  virtual bool WaitForState(grpc_connectivity_state target_state,
+                            gpr_timespec deadline) = 0;
+
 #ifndef GRPC_CXX0X_NO_CHRONO
   virtual void NotifyOnStateChange(
       grpc_connectivity_state last_observed,
@@ -81,6 +86,9 @@ class ChannelInterface : public CallHook,
   virtual bool WaitForStateChange(
       grpc_connectivity_state last_observed,
       const std::chrono::system_clock::time_point& deadline) = 0;
+  virtual bool WaitForState(
+      grpc_connectivity_state target_state,
+      const std::chrono::system_clock::time_point& deadline) = 0;
 #endif  // !GRPC_CXX0X_NO_CHRONO
 
 };
diff --git a/src/cpp/client/channel.cc b/src/cpp/client/channel.cc
index 1bdb9ae0de..ccd30c0f46 100644
--- a/src/cpp/client/channel.cc
+++ b/src/cpp/client/channel.cc
@@ -136,6 +136,30 @@ bool WaitForStateChangeShared(grpc_channel* channel,
   return ok;
 }
 
+template <typename T>
+bool WaitForStateShared(grpc_channel* channel,
+                        grpc_connectivity_state target_state,
+                        const T& deadline) {
+  grpc_connectivity_state current_state =
+      grpc_channel_check_connectivity_state(channel, 0);
+  if (current_state == target_state) {
+    return true;
+  }
+  TimePoint<T> deadline_tp(deadline);
+  CompletionQueue cq;
+  bool ok = false;
+  void* tag = NULL;
+  while (current_state != target_state) {
+    NotifyOnStateChangeShared(channel, current_state, deadline_tp.raw_time(),
+                              &cq, NULL);
+    cq.Next(&tag, &ok);
+    if (!ok) {
+      return false;
+    }
+    current_state = grpc_channel_check_connectivity_state(channel, 0);
+  }
+  return true;
+}
 }  // namespace
 
 void Channel::NotifyOnStateChange(grpc_connectivity_state last_observed,
@@ -149,6 +173,11 @@ bool Channel::WaitForStateChange(grpc_connectivity_state last_observed,
   return WaitForStateChangeShared(c_channel_, last_observed, deadline);
 }
 
+bool Channel::WaitForState(grpc_connectivity_state target_state,
+                           gpr_timespec deadline) {
+  return WaitForStateShared(c_channel_, target_state, deadline);
+}
+
 #ifndef GRPC_CXX0X_NO_CHRONO
 void Channel::NotifyOnStateChange(
       grpc_connectivity_state last_observed,
@@ -162,5 +191,12 @@ bool Channel::WaitForStateChange(
       const std::chrono::system_clock::time_point& deadline) {
   return WaitForStateChangeShared(c_channel_, last_observed, deadline);
 }
+
+bool Channel::WaitForState(
+    grpc_connectivity_state target_state,
+    const std::chrono::system_clock::time_point& deadline) {
+  return WaitForStateShared(c_channel_, target_state, deadline);
+}
+
 #endif  // !GRPC_CXX0X_NO_CHRONO
 }  // namespace grpc
diff --git a/src/cpp/client/channel.h b/src/cpp/client/channel.h
index fa3aedc9eb..4dc6723778 100644
--- a/src/cpp/client/channel.h
+++ b/src/cpp/client/channel.h
@@ -71,15 +71,24 @@ class Channel GRPC_FINAL : public GrpcLibrary, public ChannelInterface {
   bool WaitForStateChange(grpc_connectivity_state last_observed,
                           gpr_timespec deadline) GRPC_OVERRIDE;
 
+  bool WaitForState(grpc_connectivity_state target_state,
+                    gpr_timespec deadline) GRPC_OVERRIDE;
+
 #ifndef GRPC_CXX0X_NO_CHRONO
   void NotifyOnStateChange(
       grpc_connectivity_state last_observed,
       const std::chrono::system_clock::time_point& deadline,
       CompletionQueue* cq, void* tag) GRPC_OVERRIDE;
+
   bool WaitForStateChange(
       grpc_connectivity_state last_observed,
       const std::chrono::system_clock::time_point& deadline) GRPC_OVERRIDE;
+
+  bool WaitForState(grpc_connectivity_state target_state,
+                    const std::chrono::system_clock::time_point& deadline)
+      GRPC_OVERRIDE;
 #endif  // !GRPC_CXX0X_NO_CHRONO
+
  private:
   const grpc::string host_;
   grpc_channel* const c_channel_;  // owned
diff --git a/test/cpp/end2end/end2end_test.cc b/test/cpp/end2end/end2end_test.cc
index 8963382a87..12ac25c6df 100644
--- a/test/cpp/end2end/end2end_test.cc
+++ b/test/cpp/end2end/end2end_test.cc
@@ -876,10 +876,10 @@ TEST_F(End2endTest, ChannelState) {
   // Start IDLE
   EXPECT_EQ(GRPC_CHANNEL_IDLE, channel_->GetState(false));
 
+  // Did not ask to connect, no state change.
   CompletionQueue cq;
   std::chrono::system_clock::time_point deadline =
       std::chrono::system_clock::now() + std::chrono::milliseconds(10);
-  // No state change.
   channel_->NotifyOnStateChange(GRPC_CHANNEL_IDLE, deadline, &cq, NULL);
   void* tag;
   bool ok = true;
@@ -890,6 +890,8 @@ TEST_F(End2endTest, ChannelState) {
   EXPECT_TRUE(channel_->WaitForStateChange(
       GRPC_CHANNEL_IDLE, gpr_inf_future(GPR_CLOCK_REALTIME)));
   EXPECT_EQ(GRPC_CHANNEL_CONNECTING, channel_->GetState(false));
+  EXPECT_TRUE(channel_->WaitForState(GRPC_CHANNEL_READY,
+                                     gpr_inf_future(GPR_CLOCK_REALTIME)));
 }
 
 }  // namespace testing