diff --git a/BUILD.bazel b/BUILD.bazel index 2d99efd65..d8f7c0462 100644 --- a/BUILD.bazel +++ b/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", diff --git a/src/ray/common/client_connection.cc b/src/ray/common/client_connection.cc index 7f7e8a476..d92ddd3fb 100644 --- a/src/ray/common/client_connection.cc +++ b/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 std::shared_ptr> ServerConnection::Create( boost::asio::basic_stream_socket &&socket) { @@ -361,9 +339,12 @@ std::string ServerConnection::DebugString() const { return result.str(); } +#if defined(BOOST_ASIO_HAS_LOCAL_SOCKETS) +// We compile conditionally to prevent duplicate explicit instantiation error template class ServerConnection; -template class ServerConnection; template class ClientConnection; +#endif +template class ServerConnection; template class ClientConnection; } // namespace ray diff --git a/src/ray/common/client_connection.h b/src/ray/common/client_connection.h index 4c9c6292a..01cbb70c8 100644 --- a/src/ray/common/client_connection.h +++ b/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; diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 5db41b211..08925312c 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/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(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; diff --git a/src/ray/object_manager/object_store_notification_manager.cc b/src/ray/object_manager/object_store_notification_manager.cc index e2fe67623..2141656f8 100644 --- a/src/ray/object_manager/object_store_notification_manager.cc +++ b/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(&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(); } diff --git a/src/ray/object_manager/object_store_notification_manager.h b/src/ray/object_manager/object_store_notification_manager.h index 35ed0f1ee..6d9d3b315 100644 --- a/src/ray/object_manager/object_store_notification_manager.h +++ b/src/ray/object_manager/object_store_notification_manager.h @@ -66,7 +66,6 @@ class ObjectStoreNotificationManager { std::vector> rem_handlers_; plasma::PlasmaClient store_client_; - int c_socket_; int64_t length_; int64_t num_adds_processed_; int64_t num_removes_processed_; diff --git a/src/ray/raylet/raylet.cc b/src/ray/raylet/raylet.cc index cd67c8ce4..5575c70b0 100644 --- a/src/ray/raylet/raylet.cc +++ b/src/ray/raylet/raylet.cc @@ -6,6 +6,7 @@ #include #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); diff --git a/src/ray/raylet/raylet_client.cc b/src/ray/raylet/raylet_client.cc index 176c72fae..6bb07befb 100644 --- a/src/ray/raylet/raylet_client.cc +++ b/src/ray/raylet/raylet_client.cc @@ -10,92 +10,34 @@ #include #include #include -#include #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(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(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 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( - new raylet::RayletConnection(raylet_socket, -1, -1)); + new raylet::RayletConnection(io_service, raylet_socket, -1, -1)); flatbuffers::FlatBufferBuilder fbb; auto message = protocol::CreateRegisterClientRequest( diff --git a/src/ray/raylet/raylet_client.h b/src/ray/raylet/raylet_client.h index e5618c42f..4f4e81f75 100644 --- a/src/ray/raylet/raylet_client.h +++ b/src/ray/raylet/raylet_client.h @@ -27,9 +27,16 @@ using ray::rpc::ProfileTableData; using MessageType = ray::protocol::MessageType; using ResourceMappingType = std::unordered_map>>; -using Socket = boost::asio::detail::socket_holder; using WaitResultPair = std::pair, std::vector>; +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 grpc_client, + RayletClient(boost::asio::io_service &io_service, + std::shared_ptr 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); diff --git a/src/ray/util/url.cc b/src/ray/util/url.cc new file mode 100644 index 000000000..159111f62 --- /dev/null +++ b/src/ray/util/url.cc @@ -0,0 +1,65 @@ +#include "ray/util/url.h" + +#include + +#include +#include +#include +#include + +#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(host_begin), + endpoint.begin() + static_cast(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); +} diff --git a/src/ray/util/url.h b/src/ray/util/url.h new file mode 100644 index 000000000..ae53df030 --- /dev/null +++ b/src/ray/util/url.h @@ -0,0 +1,10 @@ +#ifndef RAY_UTIL_URL_H +#define RAY_UTIL_URL_H + +#include + +// 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 diff --git a/src/ray/util/url_test.cc b/src/ray/util/url_test.cc new file mode 100644 index 000000000..0825865ea --- /dev/null +++ b/src/ray/util/url_test.cc @@ -0,0 +1,44 @@ +#include "ray/util/url.h" + +#include + +#include "gtest/gtest.h" + +namespace ray { + +template +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(); +}