From 2c4514a2c095d223440c39f91aca778703d81f93 Mon Sep 17 00:00:00 2001
From: dHannasch <David.A.Hannasch@gmail.com>
Date: Mon, 23 Nov 2020 21:39:50 -0700
Subject: [PATCH] [minor] Refactor to expose RedisContext::PingPort (#12022)

---
 src/ray/gcs/redis_context.cc | 28 +++++++++++++++++++++++-----
 src/ray/gcs/redis_context.h  |  7 +++++++
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/ray/gcs/redis_context.cc b/src/ray/gcs/redis_context.cc
index 8a4a19a97..90c028912 100644
--- a/src/ray/gcs/redis_context.cc
+++ b/src/ray/gcs/redis_context.cc
@@ -284,6 +284,12 @@ void SetDisconnectCallback(RedisAsyncContext *redis_async_context) {
                                   RedisAsyncContextDisconnectCallback);
 }
 
+void FreeRedisContext(redisContext *context) { redisFree(context); }
+
+void FreeRedisContext(redisAsyncContext *context) {}
+
+void FreeRedisContext(RedisAsyncContext *context) {}
+
 template <typename RedisContext, typename RedisConnectFunction>
 Status ConnectWithoutRetries(const std::string &address, int port,
                              const RedisConnectFunction &connect_function,
@@ -291,17 +297,23 @@ Status ConnectWithoutRetries(const std::string &address, int port,
   // This currently returns the errorMessage in two different ways,
   // as an output parameter and in the Status::RedisError,
   // because we're not sure whether we'll want to change what this returns.
-  *context = connect_function(address.c_str(), port);
-  if (*context == nullptr || (*context)->err) {
+  RedisContext *newContext = connect_function(address.c_str(), port);
+  if (newContext == nullptr || (newContext)->err) {
     std::ostringstream oss(errorMessage);
-    if (*context == nullptr) {
+    if (newContext == nullptr) {
       oss << "Could not allocate Redis context.";
-    } else if ((*context)->err) {
+    } else if (newContext->err) {
       oss << "Could not establish connection to Redis " << address << ":" << port
-          << " (context.err = " << (*context)->err << ")";
+          << " (context.err = " << newContext->err << ")";
     }
     return Status::RedisError(errorMessage);
   }
+  if (context != nullptr) {
+    // Don't crash if the RedisContext** is null.
+    *context = newContext;
+  } else {
+    FreeRedisContext(newContext);
+  }
   return Status::OK();
 }
 
@@ -342,6 +354,12 @@ Status ConnectWithRetries(const std::string &address, int port,
   return Status::OK();
 }
 
+Status RedisContext::PingPort(const std::string &address, int port) {
+  std::string errorMessage;
+  return ConnectWithoutRetries(address, port, redisConnect,
+                               static_cast<redisContext **>(nullptr), errorMessage);
+}
+
 Status RedisContext::Connect(const std::string &address, int port, bool sharding,
                              const std::string &password = "") {
   RAY_CHECK(!context_);
diff --git a/src/ray/gcs/redis_context.h b/src/ray/gcs/redis_context.h
index c8bcb04a1..7af0eea3d 100644
--- a/src/ray/gcs/redis_context.h
+++ b/src/ray/gcs/redis_context.h
@@ -170,6 +170,13 @@ class RedisContext {
 
   ~RedisContext();
 
+  /// Test whether the address and port has a reachable Redis service.
+  ///
+  /// \param address IP address to test.
+  /// \param port port number to test.
+  /// \return The Status that we would get if we Connected.
+  Status PingPort(const std::string &address, int port);
+
   Status Connect(const std::string &address, int port, bool sharding,
                  const std::string &password);