Add WaitForState

2 files changed, 45 insertions, 0 deletions

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