diff --git a/src/ray/common/client_connection.cc b/src/ray/common/client_connection.cc index 1dc7c88d4..1ae225443 100644 --- a/src/ray/common/client_connection.cc +++ b/src/ray/common/client_connection.cc @@ -112,7 +112,7 @@ void ServerConnection::WriteMessageAsync( auto size = async_write_queue_.size(); auto size_is_power_of_two = (size & (size - 1)) == 0; - if (size > 100 && size_is_power_of_two) { + if (size > 1000 && size_is_power_of_two) { RAY_LOG(WARNING) << "ServerConnection has " << size << " buffered async writes"; } diff --git a/src/ray/common/client_connection.h b/src/ray/common/client_connection.h index 86f0f20c1..d4ca993d2 100644 --- a/src/ray/common/client_connection.h +++ b/src/ray/common/client_connection.h @@ -1,7 +1,7 @@ #ifndef RAY_COMMON_CLIENT_CONNECTION_H #define RAY_COMMON_CLIENT_CONNECTION_H -#include +#include #include #include @@ -94,7 +94,7 @@ class ServerConnection : public std::enable_shared_from_this const int async_write_max_messages_; /// List of pending messages to write. - std::list> async_write_queue_; + std::deque> async_write_queue_; /// Whether we are in the middle of an async write. bool async_write_in_flight_; diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index 4f4807275..6f5e25d0b 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -12,7 +12,7 @@ namespace { // A helper function to get a worker from a list. std::shared_ptr GetWorker( - const std::list> &worker_pool, + const std::unordered_set> &worker_pool, const std::shared_ptr &connection) { for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) { if ((*it)->Connection() == connection) { @@ -24,15 +24,9 @@ std::shared_ptr GetWorker( // A helper function to remove a worker from a list. Returns true if the worker // was found and removed. -bool RemoveWorker(std::list> &worker_pool, +bool RemoveWorker(std::unordered_set> &worker_pool, const std::shared_ptr &worker) { - for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) { - if (*it == worker) { - worker_pool.erase(it); - return true; - } - } - return false; + return worker_pool.erase(worker) > 0; } } // namespace @@ -152,7 +146,7 @@ void WorkerPool::RegisterWorker(std::shared_ptr worker) { auto pid = worker->Pid(); RAY_LOG(DEBUG) << "Registering worker with pid " << pid; auto &state = GetStateForLanguage(worker->GetLanguage()); - state.registered_workers.push_back(std::move(worker)); + state.registered_workers.insert(std::move(worker)); auto it = starting_worker_processes_.find(pid); RAY_CHECK(it != starting_worker_processes_.end()); @@ -165,7 +159,7 @@ void WorkerPool::RegisterWorker(std::shared_ptr worker) { void WorkerPool::RegisterDriver(std::shared_ptr driver) { RAY_CHECK(!driver->GetAssignedTaskId().is_nil()); auto &state = GetStateForLanguage(driver->GetLanguage()); - state.registered_drivers.push_back(driver); + state.registered_drivers.insert(std::move(driver)); } std::shared_ptr WorkerPool::GetRegisteredWorker( @@ -197,7 +191,7 @@ void WorkerPool::PushWorker(std::shared_ptr worker) { auto &state = GetStateForLanguage(worker->GetLanguage()); // Add the worker to the idle pool. if (worker->GetActorId().is_nil()) { - state.idle.push_back(std::move(worker)); + state.idle.insert(std::move(worker)); } else { state.idle_actor[worker->GetActorId()] = std::move(worker); } @@ -209,8 +203,8 @@ std::shared_ptr WorkerPool::PopWorker(const TaskSpecification &task_spec std::shared_ptr worker = nullptr; if (actor_id.is_nil()) { if (!state.idle.empty()) { - worker = std::move(state.idle.back()); - state.idle.pop_back(); + worker = std::move(*state.idle.begin()); + state.idle.erase(state.idle.begin()); } } else { auto actor_entry = state.idle_actor.find(actor_id); diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h index 78e0642ab..32360b769 100644 --- a/src/ray/raylet/worker_pool.h +++ b/src/ray/raylet/worker_pool.h @@ -2,9 +2,9 @@ #define RAY_RAYLET_WORKER_POOL_H #include -#include #include #include +#include #include "ray/common/client_connection.h" #include "ray/gcs/format/util.h" @@ -136,15 +136,14 @@ class WorkerPool { /// The commands and arguments used to start the worker process std::vector worker_command; /// The pool of idle non-actor workers. - std::list> idle; + std::unordered_set> idle; /// The pool of idle actor workers. std::unordered_map> idle_actor; /// All workers that have registered and are still connected, including both /// idle and executing. - // TODO(swang): Make this a map to make GetRegisteredWorker faster. - std::list> registered_workers; + std::unordered_set> registered_workers; /// All drivers that have registered and are still connected. - std::list> registered_drivers; + std::unordered_set> registered_drivers; }; /// A helper function that returns the reference of the pool state