text: // Copyright 2017 The Ray Authors.
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //  http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  
  #include "ray/raylet_client/raylet_client.h"
  
  #include "ray/common/client_connection.h"
  #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/util.h"
  
  using MessageType = ray::protocol::MessageType;
  
  namespace {
  
  flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<ray::protocol::Address>>>
  AddressesToFlatbuffer(flatbuffers::FlatBufferBuilder &fbb,
                        const std::vector<ray::rpc::Address> &addresses) {
    std::vector<flatbuffers::Offset<ray::protocol::Address>> address_vec;
    address_vec.reserve(addresses.size());
    for (const auto &addr : addresses) {
      auto fbb_addr = ray::protocol::CreateAddress(
          fbb, fbb.CreateString(addr.raylet_id()), fbb.CreateString(addr.ip_address()),
          addr.port(), fbb.CreateString(addr.worker_id()));
      address_vec.push_back(fbb_addr);
    }
    return fbb.CreateVector(address_vec);
  }
  
  }  // namespace
  
  namespace ray {
  
  raylet::RayletConnection::RayletConnection(boost::asio::io_service &io_service,
                                             const std::string &raylet_socket,
                                             int num_retries, int64_t timeout) {
    local_stream_socket socket(io_service);
    Status s = ConnectSocketRetry(socket, raylet_socket, num_retries, timeout);
    // If we could not connect to the socket, exit.
    if (!s.ok()) {
      RAY_LOG(FATAL) << "Could not connect to socket " << raylet_socket;
    }
    conn_ = ServerConnection::Create(std::move(socket));
  }
  
  Status raylet::RayletConnection::WriteMessage(MessageType type,
                                                flatbuffers::FlatBufferBuilder *fbb) {
    std::unique_lock<std::mutex> guard(write_mutex_);
    int64_t length = fbb ? fbb->GetSize() : 0;
    uint8_t *bytes = fbb ? fbb->GetBufferPointer() : nullptr;
    return conn_->WriteMessage(static_cast<int64_t>(type), length, bytes);
  }
  
  Status raylet::RayletConnection::AtomicRequestReply(MessageType request_type,
                                                      MessageType reply_type,
                                                      std::vector<uint8_t> *reply_message,
                                                      flatbuffers::FlatBufferBuilder *fbb) {
    std::unique_lock<std::mutex> guard(mutex_);
    RAY_RETURN_NOT_OK(WriteMessage(request_type, fbb));
    return conn_->ReadMessage(static_cast<int64_t>(reply_type), reply_message);
  }
  
  raylet::RayletClient::RayletClient(
      std::shared_ptr<rpc::NodeManagerWorkerClient> grpc_client)
      : grpc_client_(std::move(grpc_client)) {}
  
  raylet::RayletClient::RayletClient(
      boost::asio::io_service &io_service,
      std::shared_ptr<ray::rpc::NodeManagerWorkerClient> grpc_client,
      const std::string &raylet_socket, const WorkerID &worker_id,
      rpc::WorkerType worker_type, const JobID &job_id, const Language &language,
      const std::string &ip_address, ClientID *raylet_id, int *port,
      std::unordered_map<std::string, std::string> *system_config,
      const std::string &job_config)
      : grpc_client_(std::move(grpc_client)),
        worker_id_(worker_id),
        job_id_(job_id),
        job_config_(job_config) {
    // For C++14, we could use std::make_unique
    conn_ = std::unique_ptr<raylet::RayletConnection>(
        new raylet::RayletConnection(io_service, raylet_socket, -1, -1));
  
    flatbuffers::FlatBufferBuilder fbb;
    // TODO(suquark): Use `WorkerType` in `common.proto` without converting to int.
    auto message = protocol::CreateRegisterClientRequest(
        fbb, static_cast<int>(worker_type), to_flatbuf(fbb, worker_id), getpid(),
        to_flatbuf(fbb, job_id), language, fbb.CreateString(ip_address), /*port=*/0,
        fbb.CreateString(job_config_));
    fbb.Finish(message);
    // Register the process ID with the raylet.
    // NOTE(swang): If raylet exits and we are registered as a worker, we will get killed.
    std::vector<uint8_t> reply;
    auto status = conn_->AtomicRequestReply(MessageType::RegisterClientRequest,
                                            MessageType::RegisterClientReply, &reply, &fbb);
    RAY_CHECK_OK_PREPEND(status, "[RayletClient] Unable to register worker with raylet.");
    auto reply_message = flatbuffers::GetRoot<protocol::RegisterClientReply>(reply.data());
    *raylet_id = ClientID::FromBinary(reply_message->raylet_id()->str());
    *port = reply_message->port();
  
    RAY_CHECK(system_config);
    auto keys = reply_message->system_config_keys();
    auto values = reply_message->system_config_values();
    RAY_CHECK(keys->size() == values->size());
    for (size_t i = 0; i < keys->size(); i++) {
      system_config->emplace(keys->Get(i)->str(), values->Get(i)->str());
    }
  }
  
  Status raylet::RayletClient::Disconnect() {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateDisconnectClient(fbb);
    fbb.Finish(message);
    auto status = conn_->WriteMessage(MessageType::IntentionalDisconnectClient, &fbb);
    // Don't be too strict for disconnection errors.
    // Just create logs and prevent it from crash.
    if (!status.ok()) {
      RAY_LOG(ERROR) << status.ToString()
                     << " [RayletClient] Failed to disconnect from raylet.";
    }
    return Status::OK();
  }
  
  Status raylet::RayletClient::AnnounceWorkerPort(int port) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateAnnounceWorkerPort(fbb, port);
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::AnnounceWorkerPort, &fbb);
  }
  
  Status raylet::RayletClient::SubmitTask(const TaskSpecification &task_spec) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message =
        protocol::CreateSubmitTaskRequest(fbb, fbb.CreateString(task_spec.Serialize()));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::SubmitTask, &fbb);
  }
  
  Status raylet::RayletClient::TaskDone() {
    return conn_->WriteMessage(MessageType::TaskDone);
  }
  
  Status raylet::RayletClient::FetchOrReconstruct(
      const std::vector<ObjectID> &object_ids,
      const std::vector<rpc::Address> &owner_addresses, bool fetch_only,
      bool mark_worker_blocked, const TaskID &current_task_id) {
    RAY_CHECK(object_ids.size() == owner_addresses.size());
    flatbuffers::FlatBufferBuilder fbb;
    auto object_ids_message = to_flatbuf(fbb, object_ids);
    auto message = protocol::CreateFetchOrReconstruct(
        fbb, object_ids_message, AddressesToFlatbuffer(fbb, owner_addresses), fetch_only,
        mark_worker_blocked, to_flatbuf(fbb, current_task_id));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::FetchOrReconstruct, &fbb);
  }
  
  Status raylet::RayletClient::NotifyUnblocked(const TaskID &current_task_id) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateNotifyUnblocked(fbb, to_flatbuf(fbb, current_task_id));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::NotifyUnblocked, &fbb);
  }
  
  Status raylet::RayletClient::NotifyDirectCallTaskBlocked() {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateNotifyDirectCallTaskBlocked(fbb);
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::NotifyDirectCallTaskBlocked, &fbb);
  }
  
  Status raylet::RayletClient::NotifyDirectCallTaskUnblocked() {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateNotifyDirectCallTaskUnblocked(fbb);
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::NotifyDirectCallTaskUnblocked, &fbb);
  }
  
  Status raylet::RayletClient::Wait(const std::vector<ObjectID> &object_ids,
                                    const std::vector<rpc::Address> &owner_addresses,
                                    int num_returns, int64_t timeout_milliseconds,
                                    bool wait_local, bool mark_worker_blocked,
                                    const TaskID &current_task_id, WaitResultPair *result) {
    // Write request.
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateWaitRequest(
        fbb, to_flatbuf(fbb, object_ids), AddressesToFlatbuffer(fbb, owner_addresses),
        num_returns, timeout_milliseconds, wait_local, mark_worker_blocked,
        to_flatbuf(fbb, current_task_id));
    fbb.Finish(message);
    std::vector<uint8_t> reply;
    RAY_RETURN_NOT_OK(conn_->AtomicRequestReply(MessageType::WaitRequest,
                                                MessageType::WaitReply, &reply, &fbb));
    // Parse the flatbuffer object.
    auto reply_message = flatbuffers::GetRoot<protocol::WaitReply>(reply.data());
    auto found = reply_message->found();
    for (size_t i = 0; i < found->size(); i++) {
      ObjectID object_id = ObjectID::FromBinary(found->Get(i)->str());
      result->first.push_back(object_id);
    }
    auto remaining = reply_message->remaining();
    for (size_t i = 0; i < remaining->size(); i++) {
      ObjectID object_id = ObjectID::FromBinary(remaining->Get(i)->str());
      result->second.push_back(object_id);
    }
    return Status::OK();
  }
  
  Status raylet::RayletClient::WaitForDirectActorCallArgs(
      const std::vector<rpc::ObjectReference> &references, int64_t tag) {
    flatbuffers::FlatBufferBuilder fbb;
    std::vector<ObjectID> object_ids;
    std::vector<rpc::Address> owner_addresses;
    for (const auto &ref : references) {
      object_ids.push_back(ObjectID::FromBinary(ref.object_id()));
      owner_addresses.push_back(ref.owner_address());
    }
    auto message = protocol::CreateWaitForDirectActorCallArgsRequest(
        fbb, to_flatbuf(fbb, object_ids), AddressesToFlatbuffer(fbb, owner_addresses), tag);
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::WaitForDirectActorCallArgsRequest, &fbb);
  }
  
  Status raylet::RayletClient::PushError(const JobID &job_id, const std::string &type,
                                         const std::string &error_message,
                                         double timestamp) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreatePushErrorRequest(
        fbb, to_flatbuf(fbb, job_id), fbb.CreateString(type),
        fbb.CreateString(error_message), timestamp);
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::PushErrorRequest, &fbb);
  }
  
  Status raylet::RayletClient::PushProfileEvents(const ProfileTableData &profile_events) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = fbb.CreateString(profile_events.SerializeAsString());
    fbb.Finish(message);
  
    auto status = conn_->WriteMessage(MessageType::PushProfileEventsRequest, &fbb);
    // Don't be too strict for profile errors. Just create logs and prevent it from crash.
    if (!status.ok()) {
      RAY_LOG(ERROR) << status.ToString()
                     << " [RayletClient] Failed to push profile events.";
    }
    return Status::OK();
  }
  
  Status raylet::RayletClient::FreeObjects(const std::vector<ObjectID> &object_ids,
                                           bool local_only, bool delete_creating_tasks) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateFreeObjectsRequest(
        fbb, local_only, delete_creating_tasks, to_flatbuf(fbb, object_ids));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::FreeObjectsInObjectStoreRequest, &fbb);
  }
  
  Status raylet::RayletClient::PrepareActorCheckpoint(const ActorID &actor_id,
                                                      ActorCheckpointID *checkpoint_id) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message =
        protocol::CreatePrepareActorCheckpointRequest(fbb, to_flatbuf(fbb, actor_id));
    fbb.Finish(message);
  
    std::vector<uint8_t> reply;
    RAY_RETURN_NOT_OK(conn_->AtomicRequestReply(MessageType::PrepareActorCheckpointRequest,
                                                MessageType::PrepareActorCheckpointReply,
                                                &reply, &fbb));
    auto reply_message =
        flatbuffers::GetRoot<protocol::PrepareActorCheckpointReply>(reply.data());
    *checkpoint_id = ActorCheckpointID::FromBinary(reply_message->checkpoint_id()->str());
    return Status::OK();
  }
  
  Status raylet::RayletClient::NotifyActorResumedFromCheckpoint(
      const ActorID &actor_id, const ActorCheckpointID &checkpoint_id) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateNotifyActorResumedFromCheckpoint(
        fbb, to_flatbuf(fbb, actor_id), to_flatbuf(fbb, checkpoint_id));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::NotifyActorResumedFromCheckpoint, &fbb);
  }
  
  Status raylet::RayletClient::SetResource(const std::string &resource_name,
                                           const double capacity,
                                           const ClientID &client_Id) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateSetResourceRequest(fbb, fbb.CreateString(resource_name),
                                                      capacity, to_flatbuf(fbb, client_Id));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::SetResourceRequest, &fbb);
  }
  
  void raylet::RayletClient::RequestWorkerLease(
      const TaskSpecification &resource_spec,
      const rpc::ClientCallback<rpc::RequestWorkerLeaseReply> &callback) {
    rpc::RequestWorkerLeaseRequest request;
    request.mutable_resource_spec()->CopyFrom(resource_spec.GetMessage());
    grpc_client_->RequestWorkerLease(request, callback);
  }
  
  /// Spill objects to external storage.
  void raylet::RayletClient::RequestObjectSpillage(
      const ObjectID &object_id,
      const rpc::ClientCallback<rpc::RequestObjectSpillageReply> &callback) {
    rpc::RequestObjectSpillageRequest request;
    request.set_object_id(object_id.Binary());
    grpc_client_->RequestObjectSpillage(request, callback);
  }
  
  /// Restore spilled objects from external storage.
  /// \param object_ids The IDs of objects to be restored.
  Status raylet::RayletClient::ForceRestoreSpilledObjects(
      const std::vector<ObjectID> &object_ids) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message =
        protocol::CreateForceRestoreSpilledObjectsRequest(fbb, to_flatbuf(fbb, object_ids));
    fbb.Finish(message);
    std::vector<uint8_t> reply;
    RAY_RETURN_NOT_OK(conn_->AtomicRequestReply(
        MessageType::ForceRestoreSpilledObjectsRequest,
        MessageType::ForceRestoreSpilledObjectsReply, &reply, &fbb));
    RAY_UNUSED(
        flatbuffers::GetRoot<protocol::ForceRestoreSpilledObjectsReply>(reply.data()));
    return Status::OK();
  }
  
  Status raylet::RayletClient::ReturnWorker(int worker_port, const WorkerID &worker_id,
                                            bool disconnect_worker) {
    rpc::ReturnWorkerRequest request;
    request.set_worker_port(worker_port);
    request.set_worker_id(worker_id.Binary());
    request.set_disconnect_worker(disconnect_worker);
    grpc_client_->ReturnWorker(
        request, [](const Status &status, const rpc::ReturnWorkerReply &reply) {
          if (!status.ok()) {
            RAY_LOG(INFO) << "Error returning worker: " << status;
          }
        });
    return Status::OK();
  }
  
  void raylet::RayletClient::ReleaseUnusedWorkers(
      const std::vector<WorkerID> &workers_in_use,
      const rpc::ClientCallback<rpc::ReleaseUnusedWorkersReply> &callback) {
    rpc::ReleaseUnusedWorkersRequest request;
    for (auto &worker_id : workers_in_use) {
      request.add_worker_ids_in_use(worker_id.Binary());
    }
    grpc_client_->ReleaseUnusedWorkers(
        request,
        [callback](const Status &status, const rpc::ReleaseUnusedWorkersReply &reply) {
          if (!status.ok()) {
            RAY_LOG(WARNING)
                << "Error releasing workers from raylet, the raylet may have died:"
                << status;
          }
          callback(status, reply);
        });
  }
  
  void raylet::RayletClient::CancelWorkerLease(
      const TaskID &task_id,
      const rpc::ClientCallback<rpc::CancelWorkerLeaseReply> &callback) {
    rpc::CancelWorkerLeaseRequest request;
    request.set_task_id(task_id.Binary());
    grpc_client_->CancelWorkerLease(request, callback);
  }
  
  void raylet::RayletClient::PrepareBundleResources(
      const BundleSpecification &bundle_spec,
      const ray::rpc::ClientCallback<ray::rpc::PrepareBundleResourcesReply> &callback) {
    rpc::PrepareBundleResourcesRequest request;
    request.mutable_bundle_spec()->CopyFrom(bundle_spec.GetMessage());
    grpc_client_->PrepareBundleResources(request, callback);
  }
  
  void raylet::RayletClient::CommitBundleResources(
      const BundleSpecification &bundle_spec,
      const ray::rpc::ClientCallback<ray::rpc::CommitBundleResourcesReply> &callback) {
    rpc::CommitBundleResourcesRequest request;
    request.mutable_bundle_spec()->CopyFrom(bundle_spec.GetMessage());
    grpc_client_->CommitBundleResources(request, callback);
  }
  
  void raylet::RayletClient::CancelResourceReserve(
      BundleSpecification &bundle_spec,
      const ray::rpc::ClientCallback<ray::rpc::CancelResourceReserveReply> &callback) {
    rpc::CancelResourceReserveRequest request;
    request.mutable_bundle_spec()->CopyFrom(bundle_spec.GetMessage());
    grpc_client_->CancelResourceReserve(request, callback);
  }
  
  void raylet::RayletClient::PinObjectIDs(
      const rpc::Address &caller_address, const std::vector<ObjectID> &object_ids,
      const rpc::ClientCallback<rpc::PinObjectIDsReply> &callback) {
    rpc::PinObjectIDsRequest request;
    request.mutable_owner_address()->CopyFrom(caller_address);
    for (const ObjectID &object_id : object_ids) {
      request.add_object_ids(object_id.Binary());
    }
    grpc_client_->PinObjectIDs(request, callback);
  }
  
  void raylet::RayletClient::GlobalGC(
      const rpc::ClientCallback<rpc::GlobalGCReply> &callback) {
    rpc::GlobalGCRequest request;
    grpc_client_->GlobalGC(request, callback);
  }
  
  Status raylet::RayletClient::SubscribeToPlasma(const ObjectID &object_id) {
    flatbuffers::FlatBufferBuilder fbb;
    auto message = protocol::CreateSubscribePlasmaReady(fbb, to_flatbuf(fbb, object_id));
    fbb.Finish(message);
    return conn_->WriteMessage(MessageType::SubscribePlasmaReady, &fbb);
  }
  
  }  // namespace ray
  