Replace UNIX sockets with TCP sockets in Ray on Windows (#6823)

* Replace UNIX sockets with TCP sockets in Ray

BUILD.bazel

@@ -650,6 +650,17 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "url_test",
+    srcs = ["src/ray/util/url_test.cc"],
+    copts = COPTS,
+    deps = [
+        ":raylet_lib",
+        "@boost//:asio",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "stats_test",
     srcs = ["src/ray/stats/stats_test.cc"],
@@ -766,6 +777,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":sha256",
+        "@boost//:asio",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/time",
         "@plasma//:plasma_client",

src/ray/common/client_connection.cc

@@ -9,28 +9,6 @@
 
 namespace ray {
 
-ray::Status TcpConnect(boost::asio::ip::tcp::socket &socket,
-                       const std::string &ip_address_string, int port) {
-  // Disable Nagle's algorithm, which caused transfer delays of 10s of ms in
-  // certain cases.
-  socket.open(boost::asio::ip::tcp::v4());
-  boost::asio::ip::tcp::no_delay option(true);
-  socket.set_option(option);
-
-  boost::asio::ip::address ip_address =
-      boost::asio::ip::address::from_string(ip_address_string);
-  boost::asio::ip::tcp::endpoint endpoint(ip_address, port);
-  boost::system::error_code error;
-  socket.connect(endpoint, error);
-  const auto status = boost_to_ray_status(error);
-  if (!status.ok()) {
-    // Close the socket if the connect failed.
-    boost::system::error_code close_error;
-    socket.close(close_error);
-  }
-  return status;
-}
-
 template <class T>
 std::shared_ptr<ServerConnection<T>> ServerConnection<T>::Create(
     boost::asio::basic_stream_socket<T> &&socket) {
@@ -361,9 +339,12 @@ std::string ServerConnection<T>::DebugString() const {
   return result.str();
 }
 
+#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
+// We compile conditionally to prevent duplicate explicit instantiation error
 template class ServerConnection<local_stream_protocol>;
-template class ServerConnection<remote_stream_protocol>;
 template class ClientConnection<local_stream_protocol>;
+#endif
+template class ServerConnection<remote_stream_protocol>;
 template class ClientConnection<remote_stream_protocol>;
 
 }  // namespace ray

src/ray/common/client_connection.h

@@ -13,15 +13,6 @@
 
 namespace ray {
 
-/// Connect a TCP socket.
-///
-/// \param socket The socket to connect.
-/// \param ip_address The IP address to connect to.
-/// \param port The port to connect to.
-/// \return Status.
-ray::Status TcpConnect(boost::asio::ip::tcp::socket &socket,
-                       const std::string &ip_address, int port);
-
 /// \typename ServerConnection
 ///
 /// A generic type representing a client connection to a server. This typename
@@ -217,7 +208,7 @@ typedef
 #if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
     boost::asio::local::stream_protocol
 #else
-    boost::asio::generic::stream_protocol
+    boost::asio::ip::tcp
 #endif
         local_stream_protocol;
 

src/ray/core_worker/core_worker.cc

@@ -150,7 +150,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
       node_ip_address, node_manager_port, *client_call_manager_);
   ClientID local_raylet_id;
   local_raylet_client_ = std::shared_ptr<raylet::RayletClient>(new raylet::RayletClient(
-      std::move(grpc_client), raylet_socket, worker_context_.GetWorkerID(),
+      io_service_, std::move(grpc_client), raylet_socket, worker_context_.GetWorkerID(),
       (worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(),
       language_, &local_raylet_id, core_worker_server_.GetPort()));
   connected_ = true;

src/ray/object_manager/object_store_notification_manager.cc

@@ -22,14 +22,33 @@ ObjectStoreNotificationManager::ObjectStoreNotificationManager(
       socket_(io_service) {
   RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket_name.c_str(), "", 0, 300));
 
-  RAY_ARROW_CHECK_OK(store_client_.Subscribe(&c_socket_));
+  int c_socket;  // TODO(mehrdadn): This should be type SOCKET for Windows
+  RAY_ARROW_CHECK_OK(store_client_.Subscribe(&c_socket));
   boost::system::error_code ec;
-#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
-  local_stream_protocol sp;
-#else  // TODO(mehrdadn): HACK: FIXME: This is just to get things compiling!
-  local_stream_protocol sp(AF_UNIX, 0);
+#ifdef _WIN32
+  WSAPROTOCOL_INFO pi;
+  size_t n = sizeof(pi);
+  char *p = reinterpret_cast<char *>(&pi);
+  const int level = SOL_SOCKET;
+  const int opt = SO_PROTOCOL_INFO;
+  if (boost::asio::detail::socket_ops::getsockopt(c_socket, 0, level, opt, p, &n, ec) !=
+      boost::asio::detail::socket_error_retval) {
+    switch (pi.iAddressFamily) {
+    case AF_INET:
+      socket_.assign(local_stream_protocol::v4(), c_socket, ec);
+      break;
+    case AF_INET6:
+      socket_.assign(local_stream_protocol::v6(), c_socket, ec);
+      break;
+    default:
+      ec = boost::system::errc::make_error_code(
+          boost::system::errc::address_family_not_supported);
+      break;
+    }
+  }
+#else
+  socket_.assign(local_stream_protocol(), c_socket, ec);
 #endif
-  socket_.assign(sp, c_socket_, ec);
   assert(!ec.value());
   NotificationWait();
 }

src/ray/object_manager/object_store_notification_manager.h

@@ -66,7 +66,6 @@ class ObjectStoreNotificationManager {
   std::vector<std::function<void(const ray::ObjectID &)>> rem_handlers_;
 
   plasma::PlasmaClient store_client_;
-  int c_socket_;
   int64_t length_;
   int64_t num_adds_processed_;
   int64_t num_removes_processed_;

src/ray/raylet/raylet.cc

@@ -6,6 +6,7 @@
 #include <iostream>
 
 #include "ray/common/status.h"
+#include "ray/util/url.h"
 
 namespace {
 
@@ -53,13 +54,13 @@ Raylet::Raylet(boost::asio::io_service &main_service, const std::string &socket_
       node_manager_(main_service, self_node_id_, node_manager_config, object_manager_,
                     gcs_client_, object_directory_),
       socket_name_(socket_name),
-      acceptor_(main_service, local_stream_protocol::endpoint(
+      acceptor_(main_service,
 #if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
-                                  socket_name
-#else  // TODO(mehrdadn): HACK: FIXME: This is just to get things compiling!
-                                  socket_name.data(), socket_name.size()
+                local_stream_protocol::endpoint(socket_name)
+#else
+                parse_ip_tcp_endpoint(socket_name)
 #endif
-                                  )),
+                    ),
       socket_(main_service) {
   self_node_info_.set_node_id(self_node_id_.Binary());
   self_node_info_.set_state(GcsNodeInfo::ALIVE);

src/ray/raylet/raylet_client.cc

@@ -10,92 +10,34 @@
 #include <sys/ioctl.h>
 #include <sys/socket.h>
 #include <sys/types.h>
-#include <sys/un.h>
 
 #include "ray/common/common_protocol.h"
 #include "ray/common/ray_config.h"
 #include "ray/common/task/task_spec.h"
 #include "ray/raylet/format/node_manager_generated.h"
 #include "ray/util/logging.h"
+#include "ray/util/url.h"
 
 using MessageType = ray::protocol::MessageType;
 
-// TODO(rkn): The io methods below should be removed.
-bool connect_ipc_sock(Socket &sock, const std::string &socket_pathname) {
-  struct sockaddr_un socket_address;
-
-  sock.reset(socket(AF_UNIX, SOCK_STREAM, 0));
-  if (sock.get() < 0) {
-    RAY_LOG(ERROR) << "socket() failed for pathname " << socket_pathname;
-    return false;
-  }
-
-  memset(&socket_address, 0, sizeof(socket_address));
-  socket_address.sun_family = AF_UNIX;
-  if (socket_pathname.length() + 1 > sizeof(socket_address.sun_path)) {
-    RAY_LOG(ERROR) << "Socket pathname is too long.";
-    return false;
-  }
-  strncpy(socket_address.sun_path, socket_pathname.c_str(), socket_pathname.length() + 1);
-
-  if (connect(sock.get(), (struct sockaddr *)&socket_address, sizeof(socket_address)) !=
-      0) {
-    return false;
-  }
-  return true;
+static int read_bytes(local_stream_protocol::socket &conn, void *cursor, size_t length) {
+  boost::system::error_code ec;
+  size_t nread = boost::asio::read(conn, boost::asio::buffer(cursor, length), ec);
+  return nread == length ? 0 : -1;
 }
 
-int read_bytes(Socket &conn, uint8_t *cursor, size_t length) {
-  ssize_t nbytes = 0;
-  // Termination condition: EOF or read 'length' bytes total.
-  size_t bytesleft = length;
-  size_t offset = 0;
-  while (bytesleft > 0) {
-    nbytes = read(conn.get(), cursor + offset, bytesleft);
-    if (nbytes < 0) {
-      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
-        continue;
-      }
-      return -1;  // Errno will be set.
-    } else if (0 == nbytes) {
-      // Encountered early EOF.
-      return -1;
-    }
-    RAY_CHECK(nbytes > 0);
-    bytesleft -= nbytes;
-    offset += nbytes;
-  }
-  return 0;
-}
-
-int write_bytes(Socket &conn, uint8_t *cursor, size_t length) {
-  ssize_t nbytes = 0;
-  size_t bytesleft = length;
-  size_t offset = 0;
-  while (bytesleft > 0) {
-    // While we haven't written the whole message, write to the file
-    // descriptor, advance the cursor, and decrease the amount left to write.
-    nbytes = write(conn.get(), cursor + offset, bytesleft);
-    if (nbytes < 0) {
-      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
-        continue;
-      }
-      return -1;  // Errno will be set.
-    } else if (0 == nbytes) {
-      // Encountered early EOF.
-      return -1;
-    }
-    RAY_CHECK(nbytes > 0);
-    bytesleft -= nbytes;
-    offset += nbytes;
-  }
-  return 0;
+static int write_bytes(local_stream_protocol::socket &conn, void *cursor, size_t length) {
+  boost::system::error_code ec;
+  size_t nread = boost::asio::write(conn, boost::asio::buffer(cursor, length), ec);
+  return nread == length ? 0 : -1;
 }
 
 namespace ray {
 
-raylet::RayletConnection::RayletConnection(const std::string &raylet_socket,
-                                           int num_retries, int64_t timeout) {
+raylet::RayletConnection::RayletConnection(boost::asio::io_service &io_service,
+                                           const std::string &raylet_socket,
+                                           int num_retries, int64_t timeout)
+    : conn_(io_service) {
   // Pick the default values if the user did not specify.
   if (num_retries < 0) {
     num_retries = RayConfig::instance().num_connect_attempts();
@@ -104,10 +46,14 @@ raylet::RayletConnection::RayletConnection(const std::string &raylet_socket,
     timeout = RayConfig::instance().connect_timeout_milliseconds();
   }
   RAY_CHECK(!raylet_socket.empty());
-  bool connected = false;
+  boost::system::error_code ec;
   for (int num_attempts = 0; num_attempts < num_retries; ++num_attempts) {
-    connected = connect_ipc_sock(conn_, raylet_socket);
-    if (connected) {
+#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
+    local_stream_protocol::endpoint endpoint(raylet_socket);
+#else
+    local_stream_protocol::endpoint endpoint = parse_ip_tcp_endpoint(raylet_socket);
+#endif
+    if (!conn_.connect(endpoint, ec)) {
       break;
     }
     if (num_attempts > 0) {
@@ -119,7 +65,7 @@ raylet::RayletConnection::RayletConnection(const std::string &raylet_socket,
     usleep(timeout * 1000);
   }
   // If we could not connect to the socket, exit.
-  if (!connected) {
+  if (ec) {
     RAY_LOG(FATAL) << "Could not connect to socket " << raylet_socket;
   }
 }
@@ -143,12 +89,12 @@ Status raylet::RayletConnection::ReadMessage(MessageType type,
   int64_t cookie;
   int64_t type_field;
   int64_t length;
-  int closed = read_bytes(conn_, (uint8_t *)&cookie, sizeof(cookie));
+  int closed = read_bytes(conn_, &cookie, sizeof(cookie));
   if (closed) goto disconnected;
   RAY_CHECK(cookie == RayConfig::instance().ray_cookie());
-  closed = read_bytes(conn_, (uint8_t *)&type_field, sizeof(type_field));
+  closed = read_bytes(conn_, &type_field, sizeof(type_field));
   if (closed) goto disconnected;
-  closed = read_bytes(conn_, (uint8_t *)&length, sizeof(length));
+  closed = read_bytes(conn_, &length, sizeof(length));
   if (closed) goto disconnected;
   message = std::unique_ptr<uint8_t[]>(new uint8_t[length]);
   closed = read_bytes(conn_, message.get(), length);
@@ -182,11 +128,11 @@ Status raylet::RayletConnection::WriteMessage(MessageType type,
   int64_t type_field = static_cast<int64_t>(type);
   auto io_error = Status::IOError("[RayletClient] Connection closed unexpectedly.");
   int closed;
-  closed = write_bytes(conn_, (uint8_t *)&cookie, sizeof(cookie));
+  closed = write_bytes(conn_, &cookie, sizeof(cookie));
   if (closed) return io_error;
-  closed = write_bytes(conn_, (uint8_t *)&type_field, sizeof(type_field));
+  closed = write_bytes(conn_, &type_field, sizeof(type_field));
   if (closed) return io_error;
-  closed = write_bytes(conn_, (uint8_t *)&length, sizeof(length));
+  closed = write_bytes(conn_, &length, sizeof(length));
   if (closed) return io_error;
   closed = write_bytes(conn_, bytes, length * sizeof(char));
   if (closed) return io_error;
@@ -207,13 +153,14 @@ raylet::RayletClient::RayletClient(
     : grpc_client_(std::move(grpc_client)) {}
 
 raylet::RayletClient::RayletClient(
+    boost::asio::io_service &io_service,
     std::shared_ptr<rpc::NodeManagerWorkerClient> grpc_client,
     const std::string &raylet_socket, const WorkerID &worker_id, bool is_worker,
     const JobID &job_id, const Language &language, ClientID *raylet_id, int port)
     : grpc_client_(std::move(grpc_client)), worker_id_(worker_id), job_id_(job_id) {
   // For C++14, we could use std::make_unique
   conn_ = std::unique_ptr<raylet::RayletConnection>(
-      new raylet::RayletConnection(raylet_socket, -1, -1));
+      new raylet::RayletConnection(io_service, raylet_socket, -1, -1));
 
   flatbuffers::FlatBufferBuilder fbb;
   auto message = protocol::CreateRegisterClientRequest(

src/ray/raylet/raylet_client.h

@@ -27,9 +27,16 @@ using ray::rpc::ProfileTableData;
 using MessageType = ray::protocol::MessageType;
 using ResourceMappingType =
     std::unordered_map<std::string, std::vector<std::pair<int64_t, double>>>;
-using Socket = boost::asio::detail::socket_holder;
 using WaitResultPair = std::pair<std::vector<ObjectID>, std::vector<ObjectID>>;
 
+typedef
+#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS)
+    boost::asio::local::stream_protocol
+#else
+    boost::asio::ip::tcp
+#endif
+        local_stream_protocol;
+
 namespace ray {
 
 /// Interface for leasing workers. Abstract for testing.
@@ -66,7 +73,8 @@ class RayletConnection {
   /// \param job_id The ID of the driver. This is non-nil if the client is a
   ///        driver.
   /// \return The connection information.
-  RayletConnection(const std::string &raylet_socket, int num_retries, int64_t timeout);
+  RayletConnection(boost::asio::io_service &io_service, const std::string &raylet_socket,
+                   int num_retries, int64_t timeout);
 
   /// Notify the raylet that this client is disconnecting gracefully. This
   /// is used by actors to exit gracefully so that the raylet doesn't
@@ -86,7 +94,7 @@ class RayletConnection {
 
  private:
   /// The Unix domain socket that connects to raylet.
-  Socket conn_;
+  local_stream_protocol::socket conn_;
   /// A mutex to protect stateful operations of the raylet client.
   std::mutex mutex_;
   /// A mutex to protect write operations of the raylet client.
@@ -107,7 +115,8 @@ class RayletClient : public WorkerLeaseInterface {
   /// \param raylet_id This will be populated with the local raylet's ClientID.
   /// \param port The port that the worker will listen on for gRPC requests, if
   /// any.
-  RayletClient(std::shared_ptr<ray::rpc::NodeManagerWorkerClient> grpc_client,
+  RayletClient(boost::asio::io_service &io_service,
+               std::shared_ptr<ray::rpc::NodeManagerWorkerClient> grpc_client,
                const std::string &raylet_socket, const WorkerID &worker_id,
                bool is_worker, const JobID &job_id, const Language &language,
                ClientID *raylet_id, int port = -1);

src/ray/util/url.cc

@@ -0,0 +1,65 @@
+#include "ray/util/url.h"
+
+#include <stdlib.h>
+
+#include <algorithm>
+#include <boost/asio/ip/address.hpp>
+#include <boost/asio/ip/tcp.hpp>
+#include <string>
+
+#include "ray/util/logging.h"
+
+boost::asio::ip::tcp::endpoint parse_ip_tcp_endpoint(const std::string &endpoint,
+                                                     int default_port) {
+  const std::string scheme_sep = "://";
+  size_t scheme_begin = 0, scheme_end = endpoint.find(scheme_sep, scheme_begin);
+  size_t host_begin;
+  if (scheme_end < endpoint.size()) {
+    host_begin = scheme_end + scheme_sep.size();
+  } else {
+    scheme_end = scheme_begin;
+    host_begin = scheme_end;
+  }
+  std::string scheme = endpoint.substr(scheme_begin, scheme_end - scheme_begin);
+  RAY_CHECK(scheme_end == host_begin || scheme == "tcp");
+  size_t port_end = endpoint.find('/', host_begin);
+  if (port_end >= endpoint.size()) {
+    port_end = endpoint.size();
+  }
+  size_t host_end, port_begin;
+  if (endpoint.find('[', host_begin) == host_begin) {
+    // IPv6 with brackets (optional ports)
+    ++host_begin;
+    host_end = endpoint.find(']', host_begin);
+    if (host_end < port_end) {
+      port_begin = endpoint.find(':', host_end + 1);
+      if (port_begin < port_end) {
+        ++port_begin;
+      } else {
+        port_begin = port_end;
+      }
+    } else {
+      host_end = port_end;
+      port_begin = host_end;
+    }
+  } else if (std::count(endpoint.begin() + static_cast<ptrdiff_t>(host_begin),
+                        endpoint.begin() + static_cast<ptrdiff_t>(port_end), ':') > 1) {
+    // IPv6 without brackets (no ports)
+    port_begin = port_end;
+    host_end = port_begin;
+  } else {
+    // IPv4
+    host_end = endpoint.find(':', host_begin);
+    if (host_end < port_end) {
+      port_begin = host_end + 1;
+    } else {
+      host_end = port_end;
+      port_begin = host_end;
+    }
+  }
+  std::string host = endpoint.substr(host_begin, host_end - host_begin);
+  std::string port_str = endpoint.substr(port_begin, port_end - port_begin);
+  boost::asio::ip::address address = boost::asio::ip::make_address(host);
+  int port = port_str.empty() ? default_port : atoi(port_str.c_str());
+  return boost::asio::ip::tcp::endpoint(address, port);
+}

src/ray/util/url.h

@@ -0,0 +1,10 @@
+#ifndef RAY_UTIL_URL_H
+#define RAY_UTIL_URL_H
+
+#include <boost/asio/ip/tcp.hpp>
+
+// Parses the endpoint (host + port number) of a URL.
+boost::asio::ip::tcp::endpoint parse_ip_tcp_endpoint(const std::string &endpoint,
+                                                     int default_port = 0);
+
+#endif

src/ray/util/url_test.cc

@@ -0,0 +1,44 @@
+#include "ray/util/url.h"
+
+#include <sstream>
+
+#include "gtest/gtest.h"
+
+namespace ray {
+
+template <class T>
+static std::string to_str(const T &obj) {
+  std::ostringstream ss;
+  ss << obj;
+  return ss.str();
+}
+
+TEST(UrlTest, UrlIpTcpParseTest) {
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://[::1]:1/", 0)) == "[::1]:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://[::1]/", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://[::1]:1", 0)) == "[::1]:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://[::1]", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://::1/", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://::1", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://127.0.0.1:1/", 0)) == "127.0.0.1:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://127.0.0.1/", 0)) == "127.0.0.1:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://127.0.0.1:1", 0)) == "127.0.0.1:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("tcp://127.0.0.1", 0)) == "127.0.0.1:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("[::1]:1/", 0)) == "[::1]:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("[::1]/", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("[::1]:1", 0)) == "[::1]:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("[::1]", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("::1/", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("::1", 0)) == "[::1]:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("127.0.0.1:1/", 0)) == "127.0.0.1:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("127.0.0.1/", 0)) == "127.0.0.1:0");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("127.0.0.1:1", 0)) == "127.0.0.1:1");
+  ASSERT_TRUE(to_str(parse_ip_tcp_endpoint("127.0.0.1", 0)) == "127.0.0.1:0");
+}
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}