[core] Give io context concurrency hint (#53642)

Most of our io_contexts are only run on a single thread. asio has a
concurrency hint option to improve performance in this case, so I'm
adding an optional parameter to the constructor if the io context is
guaranteed to run on a single thread. Also serves as a source of truth
for whether an io_context is only running from a single thread or
multiple threads.

This is for sure a premature optimization, but shouldn't hurt either
way.

What the boost io_context documentation says about the `1` option (we're
on 1.81).
https://www.boost.org/doc/libs/1_81_0/doc/html/boost_asio/overview/core/concurrency_hint.html

> The implementation assumes that the io_context will be run from a
single thread, and applies several optimisations based on this
assumption.
For example, when a handler is posted from within another handler, the
new handler is added to a fast thread-local queue (with the consequence
that the new handler is held back until the currently executing handler
finishes).
The io_context still provides full thread safety, and distinct I/O
objects may be used from any thread.

---------

Signed-off-by: dayshah <dhyey2019@gmail.com>

Author: dayshah

python/ray/includes/global_state_accessor.pxd

@@ -100,7 +100,7 @@ cdef extern from * namespace "ray::gcs" nogil:
       RAY_CHECK(absl::Base64Unescape(config, &config_list));
       RayConfig::instance().initialize(config_list);
 
-      instrumented_io_context io_service;
+      instrumented_io_context io_service{/*enable_lag_probe=*/false, /*running_on_single_thread=*/true};
 
       auto redis_client = std::make_shared<RedisClient>(options);
       auto status = redis_client->Connect(io_service);

src/ray/common/asio/asio_util.h

@@ -61,7 +61,7 @@ class InstrumentedIOContextWithThread {
    */
   explicit InstrumentedIOContextWithThread(const std::string &thread_name,
                                            bool enable_lag_probe = false)
-      : io_service_(enable_lag_probe),
+      : io_service_(enable_lag_probe, /*running_on_single_thread=*/true),
         work_(io_service_.get_executor()),
         thread_name_(thread_name) {
     io_thread_ = std::thread([this] {
@@ -91,7 +91,8 @@ class InstrumentedIOContextWithThread {
   }
 
  private:
-  instrumented_io_context io_service_;
+  instrumented_io_context io_service_{/*enable_lag_probe=*/false,
+                                      /*running_on_single_thread=*/true};
   boost::asio::executor_work_guard<boost::asio::io_context::executor_type>
       work_;  // to keep io_service_ running
   std::thread io_thread_;

src/ray/common/asio/instrumented_io_context.cc

@@ -73,8 +73,11 @@ void ScheduleLagProbe(instrumented_io_context &io_context) {
 }
 }  // namespace
 
-instrumented_io_context::instrumented_io_context(bool enable_lag_probe)
-    : event_stats_(std::make_shared<EventTracker>()) {
+instrumented_io_context::instrumented_io_context(bool enable_lag_probe,
+                                                 bool running_on_single_thread)
+    : boost::asio::io_context(
+          running_on_single_thread ? 1 : BOOST_ASIO_CONCURRENCY_HINT_DEFAULT),
+      event_stats_(std::make_shared<EventTracker>()) {
   if (enable_lag_probe) {
     ScheduleLagProbe(*this);
   }
@@ -98,23 +101,24 @@ void instrumented_io_context::post(std::function<void()> handler,
   }
 
   if (delay_us == 0) {
-    boost::asio::io_context::post(std::move(handler));
+    boost::asio::post(*this, std::move(handler));
   } else {
     execute_after(*this, std::move(handler), std::chrono::microseconds(delay_us));
   }
 }
 
 void instrumented_io_context::dispatch(std::function<void()> handler, std::string name) {
   if (!RayConfig::instance().event_stats()) {
-    return boost::asio::io_context::post(std::move(handler));
+    return boost::asio::post(*this, std::move(handler));
   }
   auto stats_handle = event_stats_->RecordStart(std::move(name));
   // References are only invalidated upon deletion of the corresponding item from the
   // table, which we won't do until this io_context is deleted. Provided that
   // GuardedHandlerStats synchronizes internal access, we can concurrently write to the
   // handler stats it->second from multiple threads without acquiring a table-level
   // readers lock in the callback.
-  boost::asio::io_context::dispatch(
+  boost::asio::dispatch(
+      *this,
       [handler = std::move(handler), stats_handle = std::move(stats_handle)]() mutable {
         EventTracker::RecordExecution(handler, std::move(stats_handle));
       });

src/ray/common/asio/instrumented_io_context.h

@@ -34,7 +34,8 @@ class instrumented_io_context : public boost::asio::io_context {
   /// \param enable_lag_probe If true, and if related Ray configs are set, schedule a
   /// probe to measure the event loop lag. After a probe is done, it schedules another one
   /// so a io_context.run() call will never return.
-  explicit instrumented_io_context(bool enable_lag_probe = false);
+  explicit instrumented_io_context(bool enable_lag_probe = false,
+                                   bool running_on_single_thread = false);
 
   /// A proxy post function that collects count, queueing, and execution statistics for
   /// the given handler.

src/ray/common/asio/io_service_pool.cc

@@ -24,7 +24,8 @@ IOServicePool::~IOServicePool() {}
 
 void IOServicePool::Run() {
   for (size_t i = 0; i < io_service_num_; ++i) {
-    io_services_.emplace_back(std::make_unique<instrumented_io_context>());
+    io_services_.emplace_back(std::make_unique<instrumented_io_context>(
+        /*enable_lag_probe=*/false, /*running_on_single_thread=*/true));
     instrumented_io_context &io_service = *io_services_[i];
     threads_.emplace_back([&io_service] {
       boost::asio::executor_work_guard<boost::asio::io_context::executor_type> work(

src/ray/common/client_connection.cc

@@ -193,13 +193,13 @@ void ServerConnection::ReadBufferAsync(
   if (RayConfig::instance().event_stats()) {
     auto &io_context =
         static_cast<instrumented_io_context &>(socket_.get_executor().context());
-    const auto stats_handle =
+    auto stats_handle =
         io_context.stats().RecordStart("ServerConnection.async_read.ReadBufferAsync");
     boost::asio::async_read(
         socket_,
         buffer,
         [handler, stats_handle = std::move(stats_handle)](
-            const boost::system::error_code &ec, size_t bytes_transferred) {
+            const boost::system::error_code &ec, size_t bytes_transferred) mutable {
           EventTracker::RecordExecution(
               [handler, ec]() { handler(boost_to_ray_status(ec)); },
               std::move(stats_handle));
@@ -427,13 +427,13 @@ void ClientConnection::ProcessMessages() {
     auto this_ptr = shared_ClientConnection_from_this();
     auto &io_context = static_cast<instrumented_io_context &>(
         ServerConnection::socket_.get_executor().context());
-    const auto stats_handle = io_context.stats().RecordStart(
+    auto stats_handle = io_context.stats().RecordStart(
         "ClientConnection.async_read.ProcessMessageHeader");
     boost::asio::async_read(
         ServerConnection::socket_,
         header,
         [this, this_ptr, stats_handle = std::move(stats_handle)](
-            const boost::system::error_code &ec, size_t bytes_transferred) {
+            const boost::system::error_code &ec, size_t bytes_transferred) mutable {
           EventTracker::RecordExecution(
               [this, this_ptr, ec]() { ProcessMessageHeader(ec); },
               std::move(stats_handle));
@@ -468,13 +468,13 @@ void ClientConnection::ProcessMessageHeader(const boost::system::error_code &err
     auto this_ptr = shared_ClientConnection_from_this();
     auto &io_context = static_cast<instrumented_io_context &>(
         ServerConnection::socket_.get_executor().context());
-    const auto stats_handle =
+    auto stats_handle =
         io_context.stats().RecordStart("ClientConnection.async_read.ProcessMessage");
     boost::asio::async_read(
         ServerConnection::socket_,
         boost::asio::buffer(read_message_),
         [this, this_ptr, stats_handle = std::move(stats_handle)](
-            const boost::system::error_code &ec, size_t bytes_transferred) {
+            const boost::system::error_code &ec, size_t bytes_transferred) mutable {
           EventTracker::RecordExecution([this, this_ptr, ec]() { ProcessMessage(ec); },
                                         std::move(stats_handle));
         });

src/ray/common/file_system_monitor.h

@@ -68,7 +68,8 @@ class FileSystemMonitor {
   const std::vector<std::string> paths_;
   const double capacity_threshold_;
   std::atomic<bool> over_capacity_;
-  instrumented_io_context io_context_;
+  instrumented_io_context io_context_{/*enable_lag_probe=*/false,
+                                      /*running_on_single_thread=*/true};
   std::thread monitor_thread_;
   std::shared_ptr<PeriodicalRunner> runner_;
 };

src/ray/core_worker/core_worker.h

@@ -1745,7 +1745,8 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   bool initialized_ ABSL_GUARDED_BY(initialize_mutex_) = false;
 
   /// Event loop where the IO events are handled. e.g. async GCS operations.
-  instrumented_io_context io_service_;
+  instrumented_io_context io_service_{/*enable_lag_probe=*/false,
+                                      /*running_on_single_thread=*/true};
 
   /// Keeps the io_service_ alive.
   boost::asio::executor_work_guard<boost::asio::io_context::executor_type> io_work_;

src/ray/core_worker/core_worker_process.cc

@@ -264,7 +264,8 @@ void CoreWorkerProcessImpl::InitializeSystemConfig() {
   // the system config in the constructor of `CoreWorkerProcessImpl`.
   std::promise<std::string> promise;
   std::thread thread([&] {
-    instrumented_io_context io_service;
+    instrumented_io_context io_service{/*enable_lag_probe=*/false,
+                                       /*running_on_single_thread=*/true};
     boost::asio::executor_work_guard<boost::asio::io_context::executor_type> work(
         io_service.get_executor());
     rpc::ClientCallManager client_call_manager(io_service, /*record_stats=*/false);

src/ray/core_worker/task_event_buffer.h

@@ -436,7 +436,8 @@ class TaskEventBufferImpl : public TaskEventBuffer {
   absl::Mutex profile_mutex_;
 
   /// IO service event loop owned by TaskEventBuffer.
-  instrumented_io_context io_service_;
+  instrumented_io_context io_service_{/*enable_lag_probe=*/false,
+                                      /*running_on_single_thread=*/true};
 
   /// Work guard to prevent the io_context from exiting when no work.
   boost::asio::executor_work_guard<boost::asio::io_context::executor_type> work_guard_;