diff --git a/src/common/io.cc b/src/common/io.cc index 9a7f95507..1999b7054 100644 --- a/src/common/io.cc +++ b/src/common/io.cc @@ -251,7 +251,7 @@ int write_bytes(int fd, uint8_t *cursor, size_t length) { return 0; } -int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) { +int do_write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) { int64_t version = RayConfig::instance().ray_protocol_version(); int closed; closed = write_bytes(fd, (uint8_t *) &version, sizeof(version)); @@ -273,6 +273,19 @@ int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) { return 0; } +int write_message(int fd, + int64_t type, + int64_t length, + uint8_t *bytes, + std::mutex *mutex) { + if (mutex != NULL) { + std::unique_lock guard(*mutex); + return do_write_message(fd, type, length, bytes); + } else { + return do_write_message(fd, type, length, bytes); + } +} + int read_bytes(int fd, uint8_t *cursor, size_t length) { ssize_t nbytes = 0; /* Termination condition: EOF or read 'length' bytes total. */ @@ -388,8 +401,8 @@ disconnected: void write_log_message(int fd, const char *message) { /* Account for the \0 at the end of the string. */ - write_message(fd, static_cast(CommonMessageType::LOG_MESSAGE), - strlen(message) + 1, (uint8_t *) message); + do_write_message(fd, static_cast(CommonMessageType::LOG_MESSAGE), + strlen(message) + 1, (uint8_t *) message); } char *read_log_message(int fd) { diff --git a/src/common/io.h b/src/common/io.h index 8b4a18b43..3f976445a 100644 --- a/src/common/io.h +++ b/src/common/io.h @@ -4,6 +4,7 @@ #include #include +#include #include struct aeEventLoop; @@ -121,10 +122,16 @@ int accept_client(int socket_fd); * @param type The type of the message to send. * @param length The size in bytes of the bytes parameter. * @param bytes The address of the message to send. + * @param mutex If not NULL, the whole write operation will be locked + * with this mutex, otherwise do nothing. * @return int Whether there was an error while writing. 0 corresponds to * success and -1 corresponds to an error (errno will be set). */ -int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes); +int write_message(int fd, + int64_t type, + int64_t length, + uint8_t *bytes, + std::mutex *mutex = NULL); /** * Read a sequence of bytes written by write_message from a file descriptor. diff --git a/src/local_scheduler/local_scheduler_client.cc b/src/local_scheduler/local_scheduler_client.cc index a9b1113f8..fc187c870 100644 --- a/src/local_scheduler/local_scheduler_client.cc +++ b/src/local_scheduler/local_scheduler_client.cc @@ -31,7 +31,7 @@ LocalSchedulerConnection *LocalSchedulerConnection_init( /* Register the process ID with the local scheduler. */ int success = write_message( result->conn, static_cast(MessageType::RegisterClientRequest), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &result->write_mutex); RAY_CHECK(success == 0) << "Unable to register worker with local scheduler"; return result; @@ -47,7 +47,7 @@ void local_scheduler_disconnect_client(LocalSchedulerConnection *conn) { auto message = ray::local_scheduler::protocol::CreateDisconnectClient(fbb); fbb.Finish(message); write_message(conn->conn, static_cast(MessageType::DisconnectClient), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } void local_scheduler_log_event(LocalSchedulerConnection *conn, @@ -63,7 +63,7 @@ void local_scheduler_log_event(LocalSchedulerConnection *conn, fbb, key_string, value_string, timestamp); fbb.Finish(message); write_message(conn->conn, static_cast(MessageType::EventLogMessage), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } void local_scheduler_submit(LocalSchedulerConnection *conn, @@ -78,7 +78,7 @@ void local_scheduler_submit(LocalSchedulerConnection *conn, fbb, execution_dependencies, task_spec); fbb.Finish(message); write_message(conn->conn, static_cast(MessageType::SubmitTask), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } void local_scheduler_submit_raylet( @@ -91,19 +91,22 @@ void local_scheduler_submit_raylet( fbb, execution_dependencies_message, task_spec.ToFlatbuffer(fbb)); fbb.Finish(message); write_message(conn->conn, static_cast(MessageType::SubmitTask), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn, int64_t *task_size) { - write_message(conn->conn, static_cast(MessageType::GetTask), 0, - NULL); int64_t type; int64_t reply_size; uint8_t *reply; - /* Receive a task from the local scheduler. This will block until the local - * scheduler gives this client a task. */ - read_message(conn->conn, &type, &reply_size, &reply); + { + std::unique_lock guard(conn->mutex); + write_message(conn->conn, static_cast(MessageType::GetTask), 0, + NULL, &conn->write_mutex); + /* Receive a task from the local scheduler. This will block until the local + * scheduler gives this client a task. */ + read_message(conn->conn, &type, &reply_size, &reply); + } if (type == static_cast(CommonMessageType::DISCONNECT_CLIENT)) { RAY_LOG(DEBUG) << "Exiting because local scheduler closed connection."; exit(1); @@ -139,14 +142,17 @@ TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn, // the raylet and non-raylet code paths. TaskSpec *local_scheduler_get_task_raylet(LocalSchedulerConnection *conn, int64_t *task_size) { - write_message(conn->conn, static_cast(MessageType::GetTask), 0, - NULL); int64_t type; int64_t reply_size; uint8_t *reply; - // Receive a task from the local scheduler. This will block until the local - // scheduler gives this client a task. - read_message(conn->conn, &type, &reply_size, &reply); + { + std::unique_lock guard(conn->mutex); + write_message(conn->conn, static_cast(MessageType::GetTask), 0, + NULL, &conn->write_mutex); + // Receive a task from the local scheduler. This will block until the local + // scheduler gives this client a task. + read_message(conn->conn, &type, &reply_size, &reply); + } if (type == static_cast(CommonMessageType::DISCONNECT_CLIENT)) { RAY_LOG(DEBUG) << "Exiting because local scheduler closed connection."; exit(1); @@ -197,7 +203,7 @@ TaskSpec *local_scheduler_get_task_raylet(LocalSchedulerConnection *conn, void local_scheduler_task_done(LocalSchedulerConnection *conn) { write_message(conn->conn, static_cast(MessageType::TaskDone), 0, - NULL); + NULL, &conn->write_mutex); } void local_scheduler_reconstruct_objects( @@ -211,18 +217,18 @@ void local_scheduler_reconstruct_objects( fbb.Finish(message); write_message(conn->conn, static_cast(MessageType::ReconstructObjects), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); /* TODO(swang): Propagate the error. */ } void local_scheduler_log_message(LocalSchedulerConnection *conn) { write_message(conn->conn, static_cast(MessageType::EventLogMessage), - 0, NULL); + 0, NULL, &conn->write_mutex); } void local_scheduler_notify_unblocked(LocalSchedulerConnection *conn) { write_message(conn->conn, static_cast(MessageType::NotifyUnblocked), - 0, NULL); + 0, NULL, &conn->write_mutex); } void local_scheduler_put_object(LocalSchedulerConnection *conn, @@ -234,7 +240,7 @@ void local_scheduler_put_object(LocalSchedulerConnection *conn, fbb.Finish(message); write_message(conn->conn, static_cast(MessageType::PutObject), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } const std::vector local_scheduler_get_actor_frontier( @@ -244,13 +250,16 @@ const std::vector local_scheduler_get_actor_frontier( auto message = ray::local_scheduler::protocol::CreateGetActorFrontierRequest( fbb, to_flatbuf(fbb, actor_id)); fbb.Finish(message); - write_message(conn->conn, - static_cast(MessageType::GetActorFrontierRequest), - fbb.GetSize(), fbb.GetBufferPointer()); - int64_t type; std::vector reply; - read_vector(conn->conn, &type, reply); + { + std::unique_lock guard(conn->mutex); + write_message(conn->conn, + static_cast(MessageType::GetActorFrontierRequest), + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); + + read_vector(conn->conn, &type, reply); + } if (static_cast(type) == CommonMessageType::DISCONNECT_CLIENT) { RAY_LOG(DEBUG) << "Exiting because local scheduler closed connection."; @@ -264,7 +273,8 @@ const std::vector local_scheduler_get_actor_frontier( void local_scheduler_set_actor_frontier(LocalSchedulerConnection *conn, const std::vector &frontier) { write_message(conn->conn, static_cast(MessageType::SetActorFrontier), - frontier.size(), const_cast(frontier.data())); + frontier.size(), const_cast(frontier.data()), + &conn->write_mutex); } std::pair, std::vector> local_scheduler_wait( @@ -279,14 +289,17 @@ std::pair, std::vector> local_scheduler_wait( fbb, to_flatbuf(fbb, object_ids), num_returns, timeout_milliseconds, wait_local); fbb.Finish(message); - write_message(conn->conn, - static_cast(ray::protocol::MessageType::WaitRequest), - fbb.GetSize(), fbb.GetBufferPointer()); - // Read result. int64_t type; int64_t reply_size; uint8_t *reply; - read_message(conn->conn, &type, &reply_size, &reply); + { + std::unique_lock guard(conn->mutex); + write_message(conn->conn, + static_cast(ray::protocol::MessageType::WaitRequest), + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); + // Read result. + read_message(conn->conn, &type, &reply_size, &reply); + } RAY_CHECK(static_cast(type) == ray::protocol::MessageType::WaitReply); auto reply_message = flatbuffers::GetRoot(reply); @@ -320,7 +333,7 @@ void local_scheduler_push_error(LocalSchedulerConnection *conn, write_message(conn->conn, static_cast( ray::protocol::MessageType::PushErrorRequest), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } void local_scheduler_push_profile_events( @@ -334,5 +347,5 @@ void local_scheduler_push_profile_events( write_message(conn->conn, static_cast( ray::protocol::MessageType::PushProfileEventsRequest), - fbb.GetSize(), fbb.GetBufferPointer()); + fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex); } diff --git a/src/local_scheduler/local_scheduler_client.h b/src/local_scheduler/local_scheduler_client.h index e00342ce1..0769439f7 100644 --- a/src/local_scheduler/local_scheduler_client.h +++ b/src/local_scheduler/local_scheduler_client.h @@ -1,6 +1,8 @@ #ifndef LOCAL_SCHEDULER_CLIENT_H #define LOCAL_SCHEDULER_CLIENT_H +#include + #include "common/task.h" #include "local_scheduler_shared.h" #include "ray/raylet/task_spec.h" @@ -19,6 +21,10 @@ struct LocalSchedulerConnection { /// of that resource allocated for this worker. std::unordered_map>> resource_ids_; + /// A mutex to protect stateful operations of the local scheduler client. + std::mutex mutex; + /// A mutext to protect write operations of the local scheduler client. + std::mutex write_mutex; }; /**