Ordered execution of tasks per actor handle (#5664)

BUILD.bazel

@@ -416,6 +416,16 @@ cc_binary(
     ],
 )
 
+cc_test(
+    name = "scheduling_queue_test",
+    srcs = ["src/ray/core_worker/test/scheduling_queue_test.cc"],
+    copts = COPTS,
+    deps = [
+        ":core_worker_lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "lineage_cache_test",
     srcs = ["src/ray/raylet/lineage_cache_test.cc"],

src/ray/core_worker/test/core_worker_test.cc

@@ -698,6 +698,7 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) {
                     raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, "",
                     nullptr);
   std::unique_ptr<ActorHandle> actor_handle;
+  std::vector<ObjectID> object_ids;
 
   // Test creating actor.
   uint8_t array[] = {1, 2, 3};
@@ -715,11 +716,14 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) {
                                           30 * 1000 /* 30s */));
   // Test submitting some tasks with by-value args for that actor.
   int64_t start_ms = current_time_ms();
-  const int num_tasks = 10000;
+  const int num_tasks = 100000;
   RAY_LOG(INFO) << "start submitting " << num_tasks << " tasks";
   for (int i = 0; i < num_tasks; i++) {
     // Create arguments with PassByValue.
     std::vector<TaskArg> args;
+    int64_t array[] = {SHOULD_CHECK_MESSAGE_ORDER, i};
+    auto buffer = std::make_shared<LocalMemoryBuffer>(reinterpret_cast<uint8_t *>(array),
+                                                      sizeof(array));
     args.emplace_back(TaskArg::PassByValue(std::make_shared<RayObject>(buffer, nullptr)));
 
     TaskOptions options{1, resources};
@@ -729,9 +733,18 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) {
     RAY_CHECK_OK(
         driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids));
     ASSERT_EQ(return_ids.size(), 1);
+    object_ids.emplace_back(return_ids[0]);
   }
   RAY_LOG(INFO) << "finish submitting " << num_tasks << " tasks"
                 << ", which takes " << current_time_ms() - start_ms << " ms";
+
+  for (const auto &object_id : object_ids) {
+    std::vector<std::shared_ptr<RayObject>> results;
+    RAY_CHECK_OK(driver.Objects().Get({object_id}, -1, &results));
+    ASSERT_EQ(results.size(), 1);
+  }
+  RAY_LOG(INFO) << "finish executing " << num_tasks << " tasks"
+                << ", which takes " << current_time_ms() - start_ms << " ms";
 }
 
 TEST_F(ZeroNodeTest, TestWorkerContext) {

src/ray/core_worker/test/mock_worker.cc

@@ -44,6 +44,16 @@ class MockWorker {
       auto &data = arg->GetData();
       buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size());
     }
+    if (buffer.size() >= 8) {
+      auto int_arr = reinterpret_cast<int64_t *>(buffer.data());
+      if (int_arr[0] == SHOULD_CHECK_MESSAGE_ORDER) {
+        auto seq_no = int_arr[1];
+        if (seq_no > 0) {
+          RAY_CHECK(seq_no == prev_seq_no_ + 1) << seq_no << " vs " << prev_seq_no_;
+        }
+        prev_seq_no_ = seq_no;
+      }
+    }
     auto memory_buffer =
         std::make_shared<LocalMemoryBuffer>(buffer.data(), buffer.size(), true);
 
@@ -56,6 +66,7 @@ class MockWorker {
   }
 
   CoreWorker worker_;
+  int64_t prev_seq_no_ = 0;
 };
 
 }  // namespace ray

src/ray/core_worker/test/scheduling_queue_test.cc

@@ -0,0 +1,81 @@
+#include <thread>
+#include "gtest/gtest.h"
+
+#include "ray/core_worker/transport/direct_actor_transport.h"
+
+namespace ray {
+
+TEST(SchedulingQueueTest, TestInOrder) {
+  boost::asio::io_service io_service;
+  SchedulingQueue queue(io_service, 0);
+  int n_ok = 0;
+  int n_rej = 0;
+  auto fn_ok = [&n_ok]() { n_ok++; };
+  auto fn_rej = [&n_rej]() { n_rej++; };
+  queue.Add(0, -1, fn_ok, fn_rej);
+  queue.Add(1, -1, fn_ok, fn_rej);
+  queue.Add(2, -1, fn_ok, fn_rej);
+  queue.Add(3, -1, fn_ok, fn_rej);
+  io_service.run();
+  ASSERT_EQ(n_ok, 4);
+  ASSERT_EQ(n_rej, 0);
+}
+
+TEST(SchedulingQueueTest, TestOutOfOrder) {
+  boost::asio::io_service io_service;
+  SchedulingQueue queue(io_service, 0);
+  int n_ok = 0;
+  int n_rej = 0;
+  auto fn_ok = [&n_ok]() { n_ok++; };
+  auto fn_rej = [&n_rej]() { n_rej++; };
+  queue.Add(2, -1, fn_ok, fn_rej);
+  queue.Add(0, -1, fn_ok, fn_rej);
+  queue.Add(3, -1, fn_ok, fn_rej);
+  queue.Add(1, -1, fn_ok, fn_rej);
+  io_service.run();
+  ASSERT_EQ(n_ok, 4);
+  ASSERT_EQ(n_rej, 0);
+}
+
+TEST(SchedulingQueueTest, TestDepWaitTimeout) {
+  boost::asio::io_service io_service;
+  SchedulingQueue queue(io_service, 0);
+  int n_ok = 0;
+  int n_rej = 0;
+  auto fn_ok = [&n_ok]() { n_ok++; };
+  auto fn_rej = [&n_rej]() { n_rej++; };
+  queue.Add(2, -1, fn_ok, fn_rej);
+  queue.Add(0, -1, fn_ok, fn_rej);
+  queue.Add(3, -1, fn_ok, fn_rej);
+  ASSERT_EQ(n_ok, 1);
+  ASSERT_EQ(n_rej, 0);
+  io_service.run();  // immediately triggers timeout
+  ASSERT_EQ(n_ok, 1);
+  ASSERT_EQ(n_rej, 2);
+  queue.Add(4, -1, fn_ok, fn_rej);
+  queue.Add(5, -1, fn_ok, fn_rej);
+  ASSERT_EQ(n_ok, 3);
+  ASSERT_EQ(n_rej, 2);
+}
+
+TEST(SchedulingQueueTest, TestSkipAlreadyProcessedByClient) {
+  boost::asio::io_service io_service;
+  SchedulingQueue queue(io_service, 0);
+  int n_ok = 0;
+  int n_rej = 0;
+  auto fn_ok = [&n_ok]() { n_ok++; };
+  auto fn_rej = [&n_rej]() { n_rej++; };
+  queue.Add(2, 2, fn_ok, fn_rej);
+  queue.Add(3, 2, fn_ok, fn_rej);
+  queue.Add(1, 2, fn_ok, fn_rej);
+  io_service.run();
+  ASSERT_EQ(n_ok, 1);
+  ASSERT_EQ(n_rej, 2);
+}
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -1,4 +1,3 @@
-
 #include "ray/core_worker/transport/direct_actor_transport.h"
 #include "ray/common/task/task.h"
 
@@ -68,7 +67,7 @@ Status CoreWorkerDirectActorTaskSubmitter::SubmitTask(
 
     // Submit request.
     auto &client = rpc_clients_[actor_id];
-    PushTask(*client, *request, actor_id, task_id, num_returns);
+    PushTask(*client, std::move(request), actor_id, task_id, num_returns);
     return Status::OK();
   } else {
     // Actor is dead, treat the task as failure.
@@ -138,32 +137,31 @@ Status CoreWorkerDirectActorTaskSubmitter::SubscribeActorUpdates(
 
 void CoreWorkerDirectActorTaskSubmitter::ConnectAndSendPendingTasks(
     const ActorID &actor_id, std::string ip_address, int port) {
-  std::unique_ptr<rpc::DirectActorClient> grpc_client(
-      new rpc::DirectActorClient(ip_address, port, client_call_manager_));
+  std::shared_ptr<rpc::DirectActorClient> grpc_client =
+      rpc::DirectActorClient::make(ip_address, port, client_call_manager_);
   RAY_CHECK(rpc_clients_.emplace(actor_id, std::move(grpc_client)).second);
 
   // Submit all pending requests.
   auto &client = rpc_clients_[actor_id];
   auto &requests = pending_requests_[actor_id];
   while (!requests.empty()) {
-    const auto &request = *requests.front();
-    PushTask(*client, request, actor_id,
-             TaskID::FromBinary(request.task_spec().task_id()),
-             request.task_spec().num_returns());
+    auto request = std::move(requests.front());
+    auto num_returns = request->task_spec().num_returns();
+    auto task_id = TaskID::FromBinary(request->task_spec().task_id());
+    PushTask(*client, std::move(request), actor_id, task_id, num_returns);
     requests.pop_front();
   }
 }
 
-void CoreWorkerDirectActorTaskSubmitter::PushTask(rpc::DirectActorClient &client,
-                                                  const rpc::PushTaskRequest &request,
-                                                  const ActorID &actor_id,
-                                                  const TaskID &task_id,
-                                                  int num_returns) {
+void CoreWorkerDirectActorTaskSubmitter::PushTask(
+    rpc::DirectActorClient &client, std::unique_ptr<rpc::PushTaskRequest> request,
+    const ActorID &actor_id, const TaskID &task_id, int num_returns) {
   RAY_LOG(DEBUG) << "Pushing task " << task_id << " to actor " << actor_id;
   waiting_reply_tasks_[actor_id].insert(std::make_pair(task_id, num_returns));
-  auto status =
-      client.PushTask(request, [this, actor_id, task_id, num_returns](
-                                   Status status, const rpc::PushTaskReply &reply) {
+
+  auto status = client.PushTask(
+      std::move(request), [this, actor_id, task_id, num_returns](
+                              Status status, const rpc::PushTaskReply &reply) {
         {
           std::unique_lock<std::mutex> guard(mutex_);
           waiting_reply_tasks_[actor_id].erase(task_id);
@@ -230,6 +228,7 @@ CoreWorkerDirectActorTaskReceiver::CoreWorkerDirectActorTaskReceiver(
     boost::asio::io_service &io_service, rpc::GrpcServer &server,
     const TaskHandler &task_handler)
     : worker_context_(worker_context),
+      io_service_(io_service),
       object_interface_(object_interface),
       task_service_(io_service, *this),
       task_handler_(task_handler) {
@@ -253,33 +252,47 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask(
     return;
   }
 
-  auto num_returns = task_spec.NumReturns();
-  RAY_CHECK(task_spec.IsActorCreationTask() || task_spec.IsActorTask());
-  RAY_CHECK(num_returns > 0);
-  // Decrease to account for the dummy object id.
-  num_returns--;
-
-  std::vector<std::shared_ptr<RayObject>> results;
-  auto status = task_handler_(task_spec, &results);
-  RAY_CHECK(results.size() == num_returns) << results.size() << "  " << num_returns;
-
-  for (size_t i = 0; i < results.size(); i++) {
-    auto return_object = (*reply).add_return_objects();
-    ObjectID id = ObjectID::ForTaskReturn(
-        task_spec.TaskId(), /*index=*/i + 1,
-        /*transport_type=*/static_cast<int>(TaskTransportType::DIRECT_ACTOR));
-    return_object->set_object_id(id.Binary());
-    const auto &result = results[i];
-    if (result->HasData()) {
-      return_object->set_data(result->GetData()->Data(), result->GetData()->Size());
-    }
-    if (result->HasMetadata()) {
-      return_object->set_metadata(result->GetMetadata()->Data(),
-                                  result->GetMetadata()->Size());
-    }
+  auto it = scheduling_queue_.find(task_spec.ActorHandleId());
+  if (it == scheduling_queue_.end()) {
+    auto result = scheduling_queue_.emplace(
+        task_spec.ActorHandleId(),
+        std::unique_ptr<SchedulingQueue>(new SchedulingQueue(io_service_)));
+    it = result.first;
   }
+  it->second->Add(
+      request.sequence_number(), request.client_processed_up_to(),
+      [this, reply, send_reply_callback, task_spec]() {
+        auto num_returns = task_spec.NumReturns();
+        RAY_CHECK(task_spec.IsActorCreationTask() || task_spec.IsActorTask());
+        RAY_CHECK(num_returns > 0);
+        // Decrease to account for the dummy object id.
+        num_returns--;
 
-  send_reply_callback(status, nullptr, nullptr);
+        std::vector<std::shared_ptr<RayObject>> results;
+        auto status = task_handler_(task_spec, &results);
+        RAY_CHECK(results.size() == num_returns) << results.size() << "  " << num_returns;
+
+        for (size_t i = 0; i < results.size(); i++) {
+          auto return_object = (*reply).add_return_objects();
+          ObjectID id = ObjectID::ForTaskReturn(
+              task_spec.TaskId(), /*index=*/i + 1,
+              /*transport_type=*/static_cast<int>(TaskTransportType::DIRECT_ACTOR));
+          return_object->set_object_id(id.Binary());
+          const auto &result = results[i];
+          if (result->GetData() != nullptr) {
+            return_object->set_data(result->GetData()->Data(), result->GetData()->Size());
+          }
+          if (result->GetMetadata() != nullptr) {
+            return_object->set_metadata(result->GetMetadata()->Data(),
+                                        result->GetMetadata()->Size());
+          }
+        }
+
+        send_reply_callback(status, nullptr, nullptr);
+      },
+      [this, send_reply_callback]() {
+        send_reply_callback(Status::Invalid("client cancelled rpc"), nullptr, nullptr);
+      });
 }
 
 }  // namespace ray

src/ray/core_worker/transport/direct_actor_transport.h

@@ -3,7 +3,9 @@
 
 #include <list>
 #include <set>
+#include <utility>
 
+#include "ray/common/id.h"
 #include "ray/core_worker/object_interface.h"
 #include "ray/core_worker/transport/transport.h"
 #include "ray/gcs/redis_gcs_client.h"
@@ -12,6 +14,9 @@
 
 namespace ray {
 
+/// The max time to wait for out-of-order tasks.
+const int kMaxReorderWaitSeconds = 30;
+
 /// In direct actor call task submitter and receiver, a task is directly submitted
 /// to the actor that will execute it.
 
@@ -53,8 +58,9 @@ class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter {
   /// \param[in] task_id The ID of a task.
   /// \param[in] num_returns Number of return objects.
   /// \return Void.
-  void PushTask(rpc::DirectActorClient &client, const rpc::PushTaskRequest &request,
-                const ActorID &actor_id, const TaskID &task_id, int num_returns);
+  void PushTask(rpc::DirectActorClient &client,
+                std::unique_ptr<rpc::PushTaskRequest> request, const ActorID &actor_id,
+                const TaskID &task_id, int num_returns);
 
   /// Treat a task as failed.
   ///
@@ -98,7 +104,11 @@ class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter {
   std::unordered_map<ActorID, ActorStateData> actor_states_;
 
   /// Map from actor id to rpc client. This only includes actors that we send tasks to.
-  std::unordered_map<ActorID, std::unique_ptr<rpc::DirectActorClient>> rpc_clients_;
+  /// We use shared_ptr to enable shared_from_this for pending client callbacks.
+  ///
+  /// TODO(zhijunfu): this will be moved into `actor_states_` later when we can
+  /// subscribe updates for a specific actor.
+  std::unordered_map<ActorID, std::shared_ptr<rpc::DirectActorClient>> rpc_clients_;
 
   /// Map from actor id to the actor's pending requests.
   std::unordered_map<ActorID, std::list<std::unique_ptr<rpc::PushTaskRequest>>>
@@ -116,6 +126,78 @@ class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter {
   friend class CoreWorkerTest;
 };
 
+/// Used to ensure serial order of task execution per actor handle.
+/// See direct_actor.proto for a description of the ordering protocol.
+class SchedulingQueue {
+ public:
+  SchedulingQueue(boost::asio::io_service &io_service,
+                  int64_t reorder_wait_seconds = kMaxReorderWaitSeconds)
+      : wait_timer_(io_service), reorder_wait_seconds_(reorder_wait_seconds) {}
+
+  void Add(int64_t seq_no, int64_t client_processed_up_to,
+           std::function<void()> accept_request, std::function<void()> reject_request) {
+    if (client_processed_up_to >= next_seq_no_) {
+      RAY_LOG(DEBUG) << "client skipping requests " << next_seq_no_ << " to "
+                     << client_processed_up_to;
+      next_seq_no_ = client_processed_up_to + 1;
+    }
+    pending_tasks_[seq_no] = make_pair(accept_request, reject_request);
+
+    // Reject any stale requests that the client doesn't need any longer.
+    while (!pending_tasks_.empty() && pending_tasks_.begin()->first < next_seq_no_) {
+      auto head = pending_tasks_.begin();
+      head->second.second();  // reject_request
+      pending_tasks_.erase(head);
+    }
+
+    // Process as many in-order requests as we can.
+    while (!pending_tasks_.empty() && pending_tasks_.begin()->first == next_seq_no_) {
+      auto head = pending_tasks_.begin();
+      head->second.first();  // accept_request
+      pending_tasks_.erase(head);
+      next_seq_no_++;
+    }
+
+    // Set a timeout on the queued tasks to avoid an infinite wait on failure.
+    wait_timer_.expires_from_now(boost::posix_time::seconds(reorder_wait_seconds_));
+    if (!pending_tasks_.empty()) {
+      RAY_LOG(DEBUG) << "waiting for " << next_seq_no_ << " queue size "
+                     << pending_tasks_.size();
+      wait_timer_.async_wait([this](const boost::system::error_code &error) {
+        if (error == boost::asio::error::operation_aborted) {
+          return;  // time deadline was adjusted
+        }
+        OnDependencyWaitTimeout();
+      });
+    }
+  }
+
+ private:
+  /// Called when we time out waiting for a task dependency to show up.
+  void OnDependencyWaitTimeout() {
+    RAY_LOG(ERROR) << "timed out waiting for " << next_seq_no_
+                   << ", cancelling all queued tasks";
+    while (!pending_tasks_.empty()) {
+      auto head = pending_tasks_.begin();
+      head->second.second();  // reject_request
+      pending_tasks_.erase(head);
+      next_seq_no_ = std::max(next_seq_no_, head->first + 1);
+    }
+  }
+
+  /// Max time in seconds to wait for dependencies to show up.
+  const int64_t reorder_wait_seconds_ = 0;
+  /// Sorted map of (accept, rej) task callbacks keyed by their sequence number.
+  std::map<int64_t, std::pair<std::function<void()>, std::function<void()>>>
+      pending_tasks_;
+  /// The next sequence number we are waiting for to arrive.
+  int64_t next_seq_no_ = 0;
+  /// Timer for waiting on dependencies.
+  boost::asio::deadline_timer wait_timer_;
+
+  friend class SchedulingQueueTest;
+};
+
 class CoreWorkerDirectActorTaskReceiver : public CoreWorkerTaskReceiver,
                                           public rpc::DirectActorHandler {
  public:
@@ -138,12 +220,17 @@ class CoreWorkerDirectActorTaskReceiver : public CoreWorkerTaskReceiver,
  private:
   // Worker context.
   WorkerContext &worker_context_;
+  /// The IO event loop.
+  boost::asio::io_service &io_service_;
   // Object interface.
   CoreWorkerObjectInterface &object_interface_;
   /// The rpc service for `DirectActorService`.
   rpc::DirectActorGrpcService task_service_;
   /// The callback function to process a task.
   TaskHandler task_handler_;
+  /// Queue of pending requests per actor handle.
+  /// TODO(ekl) GC these queues once the handle is no longer active.
+  std::unordered_map<ActorHandleID, std::unique_ptr<SchedulingQueue>> scheduling_queue_;
 };
 
 }  // namespace ray

src/ray/protobuf/direct_actor.proto

@@ -16,6 +16,16 @@ message ReturnObject {
 message PushTaskRequest {
   // The task to be pushed.
   TaskSpec task_spec = 1;
+  // The sequence number of the task for this client. This must increase
+  // sequentially starting from zero for each actor handle. The server
+  // will guarantee tasks execute in this sequence, waiting for any
+  // out-of-order request messages to arrive as necessary.
+  int64 sequence_number = 2;
+  // The max sequence number the client has processed responses for. This
+  // is a performance optimization that allows the client to tell the server
+  // to cancel any PushTaskRequests with seqno <= this value, rather than
+  // waiting for the server to time out waiting for missing messages.
+  int64 client_processed_up_to = 3;
 }
 
 message PushTaskReply {

src/ray/rpc/worker/direct_actor_client.h

@@ -1,12 +1,16 @@
 #ifndef RAY_RPC_DIRECT_ACTOR_CLIENT_H
 #define RAY_RPC_DIRECT_ACTOR_CLIENT_H
 
+#include <deque>
+#include <memory>
+#include <mutex>
 #include <thread>
 
 #include <grpcpp/grpcpp.h>
 
 #include "ray/common/status.h"
 #include "ray/rpc/client_call.h"
+#include "ray/rpc/worker/direct_actor_common.h"
 #include "ray/util/logging.h"
 #include "src/ray/protobuf/direct_actor.grpc.pb.h"
 #include "src/ray/protobuf/direct_actor.pb.h"
@@ -15,8 +19,76 @@ namespace ray {
 namespace rpc {
 
 /// Client used for communicating with a direct actor server.
-class DirectActorClient {
+class DirectActorClient : public std::enable_shared_from_this<DirectActorClient> {
  public:
+  /// Constructor.
+  ///
+  /// \param[in] address Address of the direct actor server.
+  /// \param[in] port Port of the direct actor server.
+  /// \param[in] client_call_manager The `ClientCallManager` used for managing requests.
+  static std::shared_ptr<DirectActorClient> make(const std::string &address,
+                                                 const int port,
+                                                 ClientCallManager &client_call_manager) {
+    auto instance = new DirectActorClient(address, port, client_call_manager);
+    return std::shared_ptr<DirectActorClient>(instance);
+  }
+
+  /// Push a task.
+  ///
+  /// \param[in] request The request message.
+  /// \param[in] callback The callback function that handles reply.
+  /// \return if the rpc call succeeds
+  ray::Status PushTask(std::unique_ptr<PushTaskRequest> request,
+                       const ClientCallback<PushTaskReply> &callback) {
+    request->set_sequence_number(next_seq_no_++);
+    send_queue_.push_back(std::make_pair(std::move(request), callback));
+    SendRequests();
+    return ray::Status::OK();
+  }
+
+  /// Send as many pending tasks as possible. This method is thread-safe.
+  ///
+  /// The client will guarantee no more than kMaxBytesInFlight bytes of RPCs are being
+  /// sent at once. This prevents the server scheduling queue from being overwhelmed.
+  /// See direct_actor.proto for a description of the ordering protocol.
+  void SendRequests() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto this_ptr = this->shared_from_this();
+
+    while (!send_queue_.empty() && rpc_bytes_in_flight_ < kMaxBytesInFlight) {
+      auto pair = std::move(*send_queue_.begin());
+      send_queue_.pop_front();
+
+      auto request = std::move(pair.first);
+      auto callback = pair.second;
+      int64_t task_size = RequestSizeInBytes(*request);
+      int64_t seq_no = request->sequence_number();
+      request->set_client_processed_up_to(max_finished_seq_no_);
+      rpc_bytes_in_flight_ += task_size;
+
+      client_call_manager_.CreateCall<DirectActorService, PushTaskRequest, PushTaskReply>(
+          *stub_, &DirectActorService::Stub::PrepareAsyncPushTask, *request,
+          [this, this_ptr, seq_no, task_size, callback](Status status,
+                                                        const rpc::PushTaskReply &reply) {
+            {
+              std::lock_guard<std::mutex> lock(mutex_);
+              if (seq_no > max_finished_seq_no_) {
+                max_finished_seq_no_ = seq_no;
+              }
+              rpc_bytes_in_flight_ -= task_size;
+              RAY_CHECK(rpc_bytes_in_flight_ >= 0);
+            }
+            SendRequests();
+            callback(status, reply);
+          });
+    }
+
+    if (!send_queue_.empty()) {
+      RAY_LOG(DEBUG) << "client send queue size " << send_queue_.size();
+    }
+  }
+
+ private:
   /// Constructor.
   ///
   /// \param[in] address Address of the direct actor server.
@@ -30,26 +102,27 @@ class DirectActorClient {
     stub_ = DirectActorService::NewStub(channel);
   };
 
-  /// Push a task.
-  ///
-  /// \param[in] request The request message.
-  /// \param[in] callback The callback function that handles reply.
-  /// \return if the rpc call succeeds
-  ray::Status PushTask(const PushTaskRequest &request,
-                       const ClientCallback<PushTaskReply> &callback) {
-    auto call = client_call_manager_
-                    .CreateCall<DirectActorService, PushTaskRequest, PushTaskReply>(
-                        *stub_, &DirectActorService::Stub::PrepareAsyncPushTask, request,
-                        callback);
-    return call->GetStatus();
-  }
+  /// Protects against unsafe concurrent access from the callback thread.
+  std::mutex mutex_;
 
- private:
   /// The gRPC-generated stub.
   std::unique_ptr<DirectActorService::Stub> stub_;
 
   /// The `ClientCallManager` used for managing requests.
   ClientCallManager &client_call_manager_;
+
+  /// Queue of requests to send.
+  std::deque<std::pair<std::unique_ptr<PushTaskRequest>, ClientCallback<PushTaskReply>>>
+      send_queue_;
+
+  /// The next sequence number to assign to a task for this server.
+  int64_t next_seq_no_ = 0;
+
+  /// The number of bytes currently in flight.
+  int64_t rpc_bytes_in_flight_ = 0;
+
+  /// The max sequence number we have processed responses for.
+  int64_t max_finished_seq_no_ = -1;
 };
 
 }  // namespace rpc

src/ray/rpc/worker/direct_actor_common.h

@@ -0,0 +1,28 @@
+#ifndef RAY_RPC_DIRECT_ACTOR_COMMON_H
+#define RAY_RPC_DIRECT_ACTOR_COMMON_H
+
+#include "src/ray/protobuf/direct_actor.grpc.pb.h"
+#include "src/ray/protobuf/direct_actor.pb.h"
+
+namespace ray {
+namespace rpc {
+
+/// The maximum number of requests in flight per client.
+const int64_t kMaxBytesInFlight = 16 * 1024 * 1024;
+
+/// The base size in bytes per request.
+const int64_t kBaseRequestSize = 1024;
+
+/// Get the estimated size in bytes of the given task.
+const static int64_t RequestSizeInBytes(const PushTaskRequest &request) {
+  int64_t size = kBaseRequestSize;
+  for (auto &arg : request.task_spec().args()) {
+    size += arg.data().size();
+  }
+  return size;
+}
+
+}  // namespace rpc
+}  // namespace ray
+
+#endif  // RAY_RPC_DIRECT_ACTOR_COMMON_H

src/ray/util/test_util.h

@@ -7,6 +7,9 @@
 
 namespace ray {
 
+// Magic argument to signal to mock_worker we should check message order.
+int64_t SHOULD_CHECK_MESSAGE_ORDER = 123450000;
+
 /// Wait until the condition is met, or timeout is reached.
 ///
 /// \param[in] condition The condition to wait for.