From 1798deae94b5788ba2a4aab306138f4c7bcc2ef2 Mon Sep 17 00:00:00 2001 From: "Siyuan (Ryans) Zhuang" Date: Fri, 10 Jul 2020 09:22:29 -0700 Subject: [PATCH] [Core] Plasma RAII support (#9370) --- src/ray/object_manager/plasma/client.cc | 26 ++-- src/ray/object_manager/plasma/plasma.cc | 14 ++ src/ray/object_manager/plasma/plasma.h | 18 ++- src/ray/object_manager/plasma/protocol.cc | 148 ++++++++++++---------- src/ray/object_manager/plasma/protocol.h | 64 +++++----- src/ray/object_manager/plasma/store.cc | 94 +++++++------- src/ray/object_manager/plasma/store.h | 26 ++-- 7 files changed, 212 insertions(+), 178 deletions(-) diff --git a/src/ray/object_manager/plasma/client.cc b/src/ray/object_manager/plasma/client.cc index fcf3fa3bc..08dbd30fa 100644 --- a/src/ray/object_manager/plasma/client.cc +++ b/src/ray/object_manager/plasma/client.cc @@ -337,7 +337,7 @@ class PlasmaClient::Impl : public std::enable_shared_from_this store_conn_; /// Table of dlmalloc buffer files that have been memory mapped so far. This /// is a hash table mapping a file descriptor to a struct containing the /// address of the corresponding memory-mapped file. @@ -364,7 +364,7 @@ class PlasmaClient::Impl : public std::enable_shared_from_thisRelease(object_id_)); } -PlasmaClient::Impl::Impl() : store_conn_(0), store_capacity_(0) { +PlasmaClient::Impl::Impl() : store_capacity_(0) { #ifdef PLASMA_CUDA auto maybe_manager = CudaDeviceManager::Instance(); DCHECK_OK(maybe_manager.status()); @@ -406,7 +406,7 @@ bool PlasmaClient::Impl::IsInUse(const ObjectID& object_id) { int PlasmaClient::Impl::GetStoreFd(int store_fd) { auto entry = mmap_table_.find(store_fd); if (entry == mmap_table_.end()) { - int fd = recv_fd(store_conn_); + int fd = recv_fd(store_conn_->fd); RAY_CHECK(fd >= 0) << "recv not successful"; return fd; } else { @@ -716,7 +716,7 @@ Status PlasmaClient::Impl::Release(const ObjectID& object_id) { std::lock_guard guard(client_mutex_); // If the client is already disconnected, ignore release requests. - if (store_conn_ < 0) { + if (!store_conn_) { return Status::OK(); } auto object_entry = objects_in_use_.find(object_id); @@ -907,8 +907,7 @@ Status PlasmaClient::Impl::Abort(const ObjectID& object_id) { std::vector buffer; ObjectID id; - MessageType type; - RAY_RETURN_NOT_OK(ReadMessage(store_conn_, &type, &buffer)); + RAY_RETURN_NOT_OK(PlasmaReceive(store_conn_, MessageType::PlasmaAbortReply, &buffer)); return ReadAbortReply(buffer.data(), buffer.size(), &id); } @@ -944,8 +943,7 @@ Status PlasmaClient::Impl::Evict(int64_t num_bytes, int64_t& num_bytes_evicted) RAY_RETURN_NOT_OK(SendEvictRequest(store_conn_, num_bytes)); // Wait for a response with the number of bytes actually evicted. std::vector buffer; - MessageType type; - RAY_RETURN_NOT_OK(ReadMessage(store_conn_, &type, &buffer)); + RAY_RETURN_NOT_OK(PlasmaReceive(store_conn_, MessageType::PlasmaEvictReply, &buffer)); return ReadEvictReply(buffer.data(), buffer.size(), num_bytes_evicted); } @@ -954,8 +952,7 @@ Status PlasmaClient::Impl::Refresh(const std::vector& object_ids) { RAY_RETURN_NOT_OK(SendRefreshLRURequest(store_conn_, object_ids)); std::vector buffer; - MessageType type; - RAY_RETURN_NOT_OK(ReadMessage(store_conn_, &type, &buffer)); + RAY_RETURN_NOT_OK(PlasmaReceive(store_conn_, MessageType::PlasmaRefreshLRUReply, &buffer)); return ReadRefreshLRUReply(buffer.data(), buffer.size()); } @@ -1002,7 +999,7 @@ Status PlasmaClient::Impl::Subscribe(int* fd) { RAY_RETURN_NOT_OK(SendSubscribeRequest(store_conn_)); // Send the file descriptor that the Plasma store should use to push // notifications about sealed objects to this client. - RAY_CHECK(send_fd(store_conn_, sock[1]) >= 0); + RAY_CHECK(send_fd(store_conn_->fd, sock[1]) >= 0); close(sock[1]); // Return the file descriptor that the client should use to read notifications // about sealed objects. @@ -1068,7 +1065,9 @@ Status PlasmaClient::Impl::Connect(const std::string& store_socket_name, int release_delay, int num_retries) { std::lock_guard guard(client_mutex_); - RAY_RETURN_NOT_OK(ConnectIpcSocketRetry(store_socket_name, num_retries, -1, &store_conn_)); + int fd = -1; + RAY_RETURN_NOT_OK(ConnectIpcSocketRetry(store_socket_name, num_retries, -1, &fd)); + store_conn_.reset(new StoreConn(fd)); if (manager_socket_name != "") { return Status::NotImplemented("plasma manager is no longer supported"); } @@ -1102,8 +1101,7 @@ Status PlasmaClient::Impl::Disconnect() { // Close the connections to Plasma. The Plasma store will release the objects // that were in use by us when handling the SIGPIPE. - close(store_conn_); - store_conn_ = -1; + store_conn_.reset(); return Status::OK(); } diff --git a/src/ray/object_manager/plasma/plasma.cc b/src/ray/object_manager/plasma/plasma.cc index 55afa22e0..5a20cecab 100644 --- a/src/ray/object_manager/plasma/plasma.cc +++ b/src/ray/object_manager/plasma/plasma.cc @@ -29,6 +29,20 @@ namespace fb = ray::object_manager::protocol; namespace plasma { +Client::~Client() { close(fd); } + +std::ostream &operator<<(std::ostream &os, const std::shared_ptr &client) { + os << client->fd; + return os; +} + +StoreConn::~StoreConn() { close(fd); } + +std::ostream &operator<<(std::ostream &os, const std::shared_ptr &store_conn) { + os << store_conn->fd; + return os; +} + ObjectTableEntry::ObjectTableEntry() : pointer(nullptr), ref_count(0) {} ObjectTableEntry::~ObjectTableEntry() { pointer = nullptr; } diff --git a/src/ray/object_manager/plasma/plasma.h b/src/ray/object_manager/plasma/plasma.h index c2c9f4e4b..428259786 100644 --- a/src/ray/object_manager/plasma/plasma.h +++ b/src/ray/object_manager/plasma/plasma.h @@ -70,7 +70,9 @@ constexpr int64_t kBlockSize = 64; /// Contains all information that is associated with a Plasma store client. struct Client { - explicit Client(int fd); + explicit Client(int fd) : fd(fd), notification_fd(-1) {} + + ~Client(); /// The file descriptor used to communicate with the client. int fd; @@ -88,6 +90,20 @@ struct Client { std::string name = "anonymous_client"; }; +std::ostream &operator<<(std::ostream &os, const std::shared_ptr &client); + +/// Connection to Plasma Store. +struct StoreConn { + explicit StoreConn(int fd) : fd(fd) {} + + ~StoreConn(); + + /// The file descriptor used to communicate with the store. + int fd; +}; + +std::ostream &operator<<(std::ostream &os, const std::shared_ptr &store_conn); + // TODO(pcm): Replace this by the flatbuffers message PlasmaObjectSpec. struct PlasmaObject { #ifdef PLASMA_CUDA diff --git a/src/ray/object_manager/plasma/protocol.cc b/src/ray/object_manager/plasma/protocol.cc index 1bf9cf97a..bcc9e553e 100644 --- a/src/ray/object_manager/plasma/protocol.cc +++ b/src/ray/object_manager/plasma/protocol.cc @@ -69,9 +69,12 @@ flatbuffers::Offset> ToFlatbuffer( return fbb->CreateVector(arrow::util::MakeNonNull(data.data()), data.size()); } -Status PlasmaReceive(int sock, MessageType message_type, std::vector* buffer) { +Status PlasmaReceive(const std::shared_ptr &store_conn, MessageType message_type, std::vector* buffer) { + if (!store_conn) { + return Status::IOError("Connection is closed."); + } MessageType type; - RAY_RETURN_NOT_OK(ReadMessage(sock, &type, buffer)); + RAY_RETURN_NOT_OK(ReadMessage(store_conn->fd, &type, buffer)); RAY_CHECK(type == message_type) << "type = " << static_cast(type) << ", message_type = " << static_cast(message_type); @@ -101,10 +104,23 @@ void ConvertToVector(const FlatbufferVectorPointer fbvector, std::vector* out } template -Status PlasmaSend(int sock, MessageType message_type, flatbuffers::FlatBufferBuilder* fbb, +Status PlasmaSend(const std::shared_ptr &store_conn, MessageType message_type, flatbuffers::FlatBufferBuilder* fbb, const Message& message) { + if (!store_conn) { + return Status::IOError("Connection is closed."); + } fbb->Finish(message); - return WriteMessage(sock, message_type, fbb->GetSize(), fbb->GetBufferPointer()); + return WriteMessage(store_conn->fd, message_type, fbb->GetSize(), fbb->GetBufferPointer()); +} + +template +Status PlasmaSend(const std::shared_ptr &client, MessageType message_type, flatbuffers::FlatBufferBuilder* fbb, + const Message& message) { + if (!client) { + return Status::IOError("Connection is closed."); + } + fbb->Finish(message); + return WriteMessage(client->fd, message_type, fbb->GetSize(), fbb->GetBufferPointer()); } Status PlasmaErrorStatus(fb::PlasmaError plasma_error) { @@ -125,12 +141,12 @@ Status PlasmaErrorStatus(fb::PlasmaError plasma_error) { // Set options messages. -Status SendSetOptionsRequest(int sock, const std::string& client_name, +Status SendSetOptionsRequest(const std::shared_ptr &store_conn, const std::string& client_name, int64_t output_memory_limit) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaSetOptionsRequest(fbb, fbb.CreateString(client_name), output_memory_limit); - return PlasmaSend(sock, MessageType::PlasmaSetOptionsRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaSetOptionsRequest, &fbb, message); } Status ReadSetOptionsRequest(uint8_t* data, size_t size, std::string* client_name, @@ -143,10 +159,10 @@ Status ReadSetOptionsRequest(uint8_t* data, size_t size, std::string* client_nam return Status::OK(); } -Status SendSetOptionsReply(int sock, PlasmaError error) { +Status SendSetOptionsReply(const std::shared_ptr &client, PlasmaError error) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaSetOptionsReply(fbb, error); - return PlasmaSend(sock, MessageType::PlasmaSetOptionsReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaSetOptionsReply, &fbb, message); } Status ReadSetOptionsReply(uint8_t* data, size_t size) { @@ -158,16 +174,16 @@ Status ReadSetOptionsReply(uint8_t* data, size_t size) { // Get debug string messages. -Status SendGetDebugStringRequest(int sock) { +Status SendGetDebugStringRequest(const std::shared_ptr &store_conn) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaGetDebugStringRequest(fbb); - return PlasmaSend(sock, MessageType::PlasmaGetDebugStringRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaGetDebugStringRequest, &fbb, message); } -Status SendGetDebugStringReply(int sock, const std::string& debug_string) { +Status SendGetDebugStringReply(const std::shared_ptr &client, const std::string& debug_string) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaGetDebugStringReply(fbb, fbb.CreateString(debug_string)); - return PlasmaSend(sock, MessageType::PlasmaGetDebugStringReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaGetDebugStringReply, &fbb, message); } Status ReadGetDebugStringReply(uint8_t* data, size_t size, std::string* debug_string) { @@ -180,13 +196,13 @@ Status ReadGetDebugStringReply(uint8_t* data, size_t size, std::string* debug_st // Create messages. -Status SendCreateRequest(int sock, ObjectID object_id, bool evict_if_full, +Status SendCreateRequest(const std::shared_ptr &store_conn, ObjectID object_id, bool evict_if_full, int64_t data_size, int64_t metadata_size, int device_num) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaCreateRequest(fbb, fbb.CreateString(object_id.Binary()), evict_if_full, data_size, metadata_size, device_num); - return PlasmaSend(sock, MessageType::PlasmaCreateRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaCreateRequest, &fbb, message); } Status ReadCreateRequest(uint8_t* data, size_t size, ObjectID* object_id, @@ -203,7 +219,7 @@ Status ReadCreateRequest(uint8_t* data, size_t size, ObjectID* object_id, return Status::OK(); } -Status SendCreateReply(int sock, ObjectID object_id, PlasmaObject* object, +Status SendCreateReply(const std::shared_ptr &client, ObjectID object_id, PlasmaObject* object, PlasmaError error_code, int64_t mmap_size) { flatbuffers::FlatBufferBuilder fbb; PlasmaObjectSpec plasma_object(object->store_fd, object->data_offset, object->data_size, @@ -233,7 +249,7 @@ Status SendCreateReply(int sock, ObjectID object_id, PlasmaObject* object, #endif } auto message = crb.Finish(); - return PlasmaSend(sock, MessageType::PlasmaCreateReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaCreateReply, &fbb, message); } Status ReadCreateReply(uint8_t* data, size_t size, ObjectID* object_id, @@ -262,13 +278,13 @@ Status ReadCreateReply(uint8_t* data, size_t size, ObjectID* object_id, return PlasmaErrorStatus(message->error()); } -Status SendCreateAndSealRequest(int sock, const ObjectID& object_id, bool evict_if_full, +Status SendCreateAndSealRequest(const std::shared_ptr &store_conn, const ObjectID& object_id, bool evict_if_full, const std::string& data, const std::string& metadata) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaCreateAndSealRequest( fbb, fbb.CreateString(object_id.Binary()), evict_if_full, fbb.CreateString(data), fbb.CreateString(metadata)); - return PlasmaSend(sock, MessageType::PlasmaCreateAndSealRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaCreateAndSealRequest, &fbb, message); } Status ReadCreateAndSealRequest(uint8_t* data, size_t size, ObjectID* object_id, @@ -285,7 +301,7 @@ Status ReadCreateAndSealRequest(uint8_t* data, size_t size, ObjectID* object_id, return Status::OK(); } -Status SendCreateAndSealBatchRequest(int sock, const std::vector& object_ids, +Status SendCreateAndSealBatchRequest(const std::shared_ptr &store_conn, const std::vector& object_ids, bool evict_if_full, const std::vector& data, const std::vector& metadata) { @@ -295,7 +311,7 @@ Status SendCreateAndSealBatchRequest(int sock, const std::vector& obje fbb, ToFlatbuffer(&fbb, object_ids.data(), object_ids.size()), evict_if_full, ToFlatbuffer(&fbb, data), ToFlatbuffer(&fbb, metadata)); - return PlasmaSend(sock, MessageType::PlasmaCreateAndSealBatchRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaCreateAndSealBatchRequest, &fbb, message); } Status ReadCreateAndSealBatchRequest(uint8_t* data, size_t size, @@ -322,10 +338,10 @@ Status ReadCreateAndSealBatchRequest(uint8_t* data, size_t size, return Status::OK(); } -Status SendCreateAndSealReply(int sock, PlasmaError error) { +Status SendCreateAndSealReply(const std::shared_ptr &client, PlasmaError error) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaCreateAndSealReply(fbb, static_cast(error)); - return PlasmaSend(sock, MessageType::PlasmaCreateAndSealReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaCreateAndSealReply, &fbb, message); } Status ReadCreateAndSealReply(uint8_t* data, size_t size) { @@ -335,11 +351,11 @@ Status ReadCreateAndSealReply(uint8_t* data, size_t size) { return PlasmaErrorStatus(message->error()); } -Status SendCreateAndSealBatchReply(int sock, PlasmaError error) { +Status SendCreateAndSealBatchReply(const std::shared_ptr &client, PlasmaError error) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaCreateAndSealBatchReply(fbb, static_cast(error)); - return PlasmaSend(sock, MessageType::PlasmaCreateAndSealBatchReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaCreateAndSealBatchReply, &fbb, message); } Status ReadCreateAndSealBatchReply(uint8_t* data, size_t size) { @@ -349,10 +365,10 @@ Status ReadCreateAndSealBatchReply(uint8_t* data, size_t size) { return PlasmaErrorStatus(message->error()); } -Status SendAbortRequest(int sock, ObjectID object_id) { +Status SendAbortRequest(const std::shared_ptr &store_conn, ObjectID object_id) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaAbortRequest(fbb, fbb.CreateString(object_id.Binary())); - return PlasmaSend(sock, MessageType::PlasmaAbortRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaAbortRequest, &fbb, message); } Status ReadAbortRequest(uint8_t* data, size_t size, ObjectID* object_id) { @@ -363,10 +379,10 @@ Status ReadAbortRequest(uint8_t* data, size_t size, ObjectID* object_id) { return Status::OK(); } -Status SendAbortReply(int sock, ObjectID object_id) { +Status SendAbortReply(const std::shared_ptr &client, ObjectID object_id) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaAbortReply(fbb, fbb.CreateString(object_id.Binary())); - return PlasmaSend(sock, MessageType::PlasmaAbortReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaAbortReply, &fbb, message); } Status ReadAbortReply(uint8_t* data, size_t size, ObjectID* object_id) { @@ -379,10 +395,10 @@ Status ReadAbortReply(uint8_t* data, size_t size, ObjectID* object_id) { // Seal messages. -Status SendSealRequest(int sock, ObjectID object_id) { +Status SendSealRequest(const std::shared_ptr &store_conn, ObjectID object_id) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaSealRequest(fbb, fbb.CreateString(object_id.Binary())); - return PlasmaSend(sock, MessageType::PlasmaSealRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaSealRequest, &fbb, message); } Status ReadSealRequest(uint8_t* data, size_t size, ObjectID* object_id) { @@ -393,11 +409,11 @@ Status ReadSealRequest(uint8_t* data, size_t size, ObjectID* object_id) { return Status::OK(); } -Status SendSealReply(int sock, ObjectID object_id, PlasmaError error) { +Status SendSealReply(const std::shared_ptr &client, ObjectID object_id, PlasmaError error) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaSealReply(fbb, fbb.CreateString(object_id.Binary()), error); - return PlasmaSend(sock, MessageType::PlasmaSealReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaSealReply, &fbb, message); } Status ReadSealReply(uint8_t* data, size_t size, ObjectID* object_id) { @@ -410,11 +426,11 @@ Status ReadSealReply(uint8_t* data, size_t size, ObjectID* object_id) { // Release messages. -Status SendReleaseRequest(int sock, ObjectID object_id) { +Status SendReleaseRequest(const std::shared_ptr &store_conn, ObjectID object_id) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaReleaseRequest(fbb, fbb.CreateString(object_id.Binary())); - return PlasmaSend(sock, MessageType::PlasmaReleaseRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaReleaseRequest, &fbb, message); } Status ReadReleaseRequest(uint8_t* data, size_t size, ObjectID* object_id) { @@ -425,11 +441,11 @@ Status ReadReleaseRequest(uint8_t* data, size_t size, ObjectID* object_id) { return Status::OK(); } -Status SendReleaseReply(int sock, ObjectID object_id, PlasmaError error) { +Status SendReleaseReply(const std::shared_ptr &client, ObjectID object_id, PlasmaError error) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaReleaseReply(fbb, fbb.CreateString(object_id.Binary()), error); - return PlasmaSend(sock, MessageType::PlasmaReleaseReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaReleaseReply, &fbb, message); } Status ReadReleaseReply(uint8_t* data, size_t size, ObjectID* object_id) { @@ -442,12 +458,12 @@ Status ReadReleaseReply(uint8_t* data, size_t size, ObjectID* object_id) { // Delete objects messages. -Status SendDeleteRequest(int sock, const std::vector& object_ids) { +Status SendDeleteRequest(const std::shared_ptr &store_conn, const std::vector& object_ids) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaDeleteRequest( fbb, static_cast(object_ids.size()), ToFlatbuffer(&fbb, &object_ids[0], object_ids.size())); - return PlasmaSend(sock, MessageType::PlasmaDeleteRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaDeleteRequest, &fbb, message); } Status ReadDeleteRequest(uint8_t* data, size_t size, std::vector* object_ids) { @@ -463,7 +479,7 @@ Status ReadDeleteRequest(uint8_t* data, size_t size, std::vector* obje return Status::OK(); } -Status SendDeleteReply(int sock, const std::vector& object_ids, +Status SendDeleteReply(const std::shared_ptr &client, const std::vector& object_ids, const std::vector& errors) { RAY_DCHECK(object_ids.size() == errors.size()); flatbuffers::FlatBufferBuilder fbb; @@ -473,7 +489,7 @@ Status SendDeleteReply(int sock, const std::vector& object_ids, fbb.CreateVector( arrow::util::MakeNonNull(reinterpret_cast(errors.data())), object_ids.size())); - return PlasmaSend(sock, MessageType::PlasmaDeleteReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaDeleteReply, &fbb, message); } Status ReadDeleteReply(uint8_t* data, size_t size, std::vector* object_ids, @@ -496,11 +512,11 @@ Status ReadDeleteReply(uint8_t* data, size_t size, std::vector* object // Contains messages. -Status SendContainsRequest(int sock, ObjectID object_id) { +Status SendContainsRequest(const std::shared_ptr &store_conn, ObjectID object_id) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaContainsRequest(fbb, fbb.CreateString(object_id.Binary())); - return PlasmaSend(sock, MessageType::PlasmaContainsRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaContainsRequest, &fbb, message); } Status ReadContainsRequest(uint8_t* data, size_t size, ObjectID* object_id) { @@ -511,11 +527,11 @@ Status ReadContainsRequest(uint8_t* data, size_t size, ObjectID* object_id) { return Status::OK(); } -Status SendContainsReply(int sock, ObjectID object_id, bool has_object) { +Status SendContainsReply(const std::shared_ptr &client, ObjectID object_id, bool has_object) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaContainsReply(fbb, fbb.CreateString(object_id.Binary()), has_object); - return PlasmaSend(sock, MessageType::PlasmaContainsReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaContainsReply, &fbb, message); } Status ReadContainsReply(uint8_t* data, size_t size, ObjectID* object_id, @@ -530,18 +546,18 @@ Status ReadContainsReply(uint8_t* data, size_t size, ObjectID* object_id, // Connect messages. -Status SendConnectRequest(int sock) { +Status SendConnectRequest(const std::shared_ptr &store_conn) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaConnectRequest(fbb); - return PlasmaSend(sock, MessageType::PlasmaConnectRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaConnectRequest, &fbb, message); } Status ReadConnectRequest(uint8_t* data) { return Status::OK(); } -Status SendConnectReply(int sock, int64_t memory_capacity) { +Status SendConnectReply(const std::shared_ptr &client, int64_t memory_capacity) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaConnectReply(fbb, memory_capacity); - return PlasmaSend(sock, MessageType::PlasmaConnectReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaConnectReply, &fbb, message); } Status ReadConnectReply(uint8_t* data, size_t size, int64_t* memory_capacity) { @@ -554,10 +570,10 @@ Status ReadConnectReply(uint8_t* data, size_t size, int64_t* memory_capacity) { // Evict messages. -Status SendEvictRequest(int sock, int64_t num_bytes) { +Status SendEvictRequest(const std::shared_ptr &store_conn, int64_t num_bytes) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaEvictRequest(fbb, num_bytes); - return PlasmaSend(sock, MessageType::PlasmaEvictRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaEvictRequest, &fbb, message); } Status ReadEvictRequest(uint8_t* data, size_t size, int64_t* num_bytes) { @@ -568,10 +584,10 @@ Status ReadEvictRequest(uint8_t* data, size_t size, int64_t* num_bytes) { return Status::OK(); } -Status SendEvictReply(int sock, int64_t num_bytes) { +Status SendEvictReply(const std::shared_ptr &client, int64_t num_bytes) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaEvictReply(fbb, num_bytes); - return PlasmaSend(sock, MessageType::PlasmaEvictReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaEvictReply, &fbb, message); } Status ReadEvictReply(uint8_t* data, size_t size, int64_t& num_bytes) { @@ -584,12 +600,12 @@ Status ReadEvictReply(uint8_t* data, size_t size, int64_t& num_bytes) { // Get messages. -Status SendGetRequest(int sock, const ObjectID* object_ids, int64_t num_objects, +Status SendGetRequest(const std::shared_ptr &store_conn, const ObjectID* object_ids, int64_t num_objects, int64_t timeout_ms) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaGetRequest( fbb, ToFlatbuffer(&fbb, object_ids, num_objects), timeout_ms); - return PlasmaSend(sock, MessageType::PlasmaGetRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaGetRequest, &fbb, message); } Status ReadGetRequest(uint8_t* data, size_t size, std::vector& object_ids, @@ -605,7 +621,7 @@ Status ReadGetRequest(uint8_t* data, size_t size, std::vector& object_ return Status::OK(); } -Status SendGetReply(int sock, ObjectID object_ids[], +Status SendGetReply(const std::shared_ptr &client, ObjectID object_ids[], std::unordered_map& plasma_objects, int64_t num_objects, const std::vector& store_fds, const std::vector& mmap_sizes) { @@ -633,7 +649,7 @@ Status SendGetReply(int sock, ObjectID object_ids[], fbb.CreateVector(arrow::util::MakeNonNull(store_fds.data()), store_fds.size()), fbb.CreateVector(arrow::util::MakeNonNull(mmap_sizes.data()), mmap_sizes.size()), fbb.CreateVector(arrow::util::MakeNonNull(handles.data()), handles.size())); - return PlasmaSend(sock, MessageType::PlasmaGetReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaGetReply, &fbb, message); } Status ReadGetReply(uint8_t* data, size_t size, ObjectID object_ids[], @@ -675,20 +691,20 @@ Status ReadGetReply(uint8_t* data, size_t size, ObjectID object_ids[], // Subscribe messages. -Status SendSubscribeRequest(int sock) { +Status SendSubscribeRequest(const std::shared_ptr &store_conn) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaSubscribeRequest(fbb); - return PlasmaSend(sock, MessageType::PlasmaSubscribeRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaSubscribeRequest, &fbb, message); } // Data messages. -Status SendDataRequest(int sock, ObjectID object_id, const char* address, int port) { +Status SendDataRequest(const std::shared_ptr &store_conn, ObjectID object_id, const char* address, int port) { flatbuffers::FlatBufferBuilder fbb; auto addr = fbb.CreateString(address, strlen(address)); auto message = fb::CreatePlasmaDataRequest(fbb, fbb.CreateString(object_id.Binary()), addr, port); - return PlasmaSend(sock, MessageType::PlasmaDataRequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaDataRequest, &fbb, message); } Status ReadDataRequest(uint8_t* data, size_t size, ObjectID* object_id, char** address, @@ -707,12 +723,12 @@ Status ReadDataRequest(uint8_t* data, size_t size, ObjectID* object_id, char** a return Status::OK(); } -Status SendDataReply(int sock, ObjectID object_id, int64_t object_size, +Status SendDataReply(const std::shared_ptr &client, ObjectID object_id, int64_t object_size, int64_t metadata_size) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaDataReply(fbb, fbb.CreateString(object_id.Binary()), object_size, metadata_size); - return PlasmaSend(sock, MessageType::PlasmaDataReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaDataReply, &fbb, message); } Status ReadDataReply(uint8_t* data, size_t size, ObjectID* object_id, @@ -728,13 +744,13 @@ Status ReadDataReply(uint8_t* data, size_t size, ObjectID* object_id, // RefreshLRU messages. -Status SendRefreshLRURequest(int sock, const std::vector& object_ids) { +Status SendRefreshLRURequest(const std::shared_ptr &store_conn, const std::vector& object_ids) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaRefreshLRURequest( fbb, ToFlatbuffer(&fbb, object_ids.data(), object_ids.size())); - return PlasmaSend(sock, MessageType::PlasmaRefreshLRURequest, &fbb, message); + return PlasmaSend(store_conn, MessageType::PlasmaRefreshLRURequest, &fbb, message); } Status ReadRefreshLRURequest(uint8_t* data, size_t size, @@ -749,10 +765,10 @@ Status ReadRefreshLRURequest(uint8_t* data, size_t size, return Status::OK(); } -Status SendRefreshLRUReply(int sock) { +Status SendRefreshLRUReply(const std::shared_ptr &client) { flatbuffers::FlatBufferBuilder fbb; auto message = fb::CreatePlasmaRefreshLRUReply(fbb); - return PlasmaSend(sock, MessageType::PlasmaRefreshLRUReply, &fbb, message); + return PlasmaSend(client, MessageType::PlasmaRefreshLRUReply, &fbb, message); } Status ReadRefreshLRUReply(uint8_t* data, size_t size) { diff --git a/src/ray/object_manager/plasma/protocol.h b/src/ray/object_manager/plasma/protocol.h index a2924f88f..2289d26bb 100644 --- a/src/ray/object_manager/plasma/protocol.h +++ b/src/ray/object_manager/plasma/protocol.h @@ -52,51 +52,51 @@ flatbuffers::Offset> ToFlatbuffer( /* Plasma receive message. */ -Status PlasmaReceive(int sock, MessageType message_type, std::vector* buffer); +Status PlasmaReceive(const std::shared_ptr &store_conn, MessageType message_type, std::vector* buffer); /* Set options messages. */ -Status SendSetOptionsRequest(int sock, const std::string& client_name, +Status SendSetOptionsRequest(const std::shared_ptr &store_conn, const std::string& client_name, int64_t output_memory_limit); Status ReadSetOptionsRequest(uint8_t* data, size_t size, std::string* client_name, int64_t* output_memory_quota); -Status SendSetOptionsReply(int sock, PlasmaError error); +Status SendSetOptionsReply(const std::shared_ptr &client, PlasmaError error); Status ReadSetOptionsReply(uint8_t* data, size_t size); /* Debug string messages. */ -Status SendGetDebugStringRequest(int sock); +Status SendGetDebugStringRequest(const std::shared_ptr &store_conn); -Status SendGetDebugStringReply(int sock, const std::string& debug_string); +Status SendGetDebugStringReply(const std::shared_ptr &client, const std::string& debug_string); Status ReadGetDebugStringReply(uint8_t* data, size_t size, std::string* debug_string); /* Plasma Create message functions. */ -Status SendCreateRequest(int sock, ObjectID object_id, bool evict_if_full, +Status SendCreateRequest(const std::shared_ptr &store_conn, ObjectID object_id, bool evict_if_full, int64_t data_size, int64_t metadata_size, int device_num); Status ReadCreateRequest(uint8_t* data, size_t size, ObjectID* object_id, bool* evict_if_full, int64_t* data_size, int64_t* metadata_size, int* device_num); -Status SendCreateReply(int sock, ObjectID object_id, PlasmaObject* object, +Status SendCreateReply(const std::shared_ptr &client, ObjectID object_id, PlasmaObject* object, PlasmaError error, int64_t mmap_size); Status ReadCreateReply(uint8_t* data, size_t size, ObjectID* object_id, PlasmaObject* object, int* store_fd, int64_t* mmap_size); -Status SendCreateAndSealRequest(int sock, const ObjectID& object_id, bool evict_if_full, +Status SendCreateAndSealRequest(const std::shared_ptr &store_conn, const ObjectID& object_id, bool evict_if_full, const std::string& data, const std::string& metadata); Status ReadCreateAndSealRequest(uint8_t* data, size_t size, ObjectID* object_id, bool* evict_if_full, std::string* object_data, std::string* metadata); -Status SendCreateAndSealBatchRequest(int sock, const std::vector& object_ids, +Status SendCreateAndSealBatchRequest(const std::shared_ptr &store_conn, const std::vector& object_ids, bool evict_if_full, const std::vector& data, const std::vector& metadata); @@ -107,41 +107,41 @@ Status ReadCreateAndSealBatchRequest(uint8_t* data, size_t size, std::vector* object_data, std::vector* metadata); -Status SendCreateAndSealReply(int sock, PlasmaError error); +Status SendCreateAndSealReply(const std::shared_ptr &client, PlasmaError error); Status ReadCreateAndSealReply(uint8_t* data, size_t size); -Status SendCreateAndSealBatchReply(int sock, PlasmaError error); +Status SendCreateAndSealBatchReply(const std::shared_ptr &client, PlasmaError error); Status ReadCreateAndSealBatchReply(uint8_t* data, size_t size); -Status SendAbortRequest(int sock, ObjectID object_id); +Status SendAbortRequest(const std::shared_ptr &store_conn, ObjectID object_id); Status ReadAbortRequest(uint8_t* data, size_t size, ObjectID* object_id); -Status SendAbortReply(int sock, ObjectID object_id); +Status SendAbortReply(const std::shared_ptr &client, ObjectID object_id); Status ReadAbortReply(uint8_t* data, size_t size, ObjectID* object_id); /* Plasma Seal message functions. */ -Status SendSealRequest(int sock, ObjectID object_id); +Status SendSealRequest(const std::shared_ptr &store_conn, ObjectID object_id); Status ReadSealRequest(uint8_t* data, size_t size, ObjectID* object_id); -Status SendSealReply(int sock, ObjectID object_id, PlasmaError error); +Status SendSealReply(const std::shared_ptr &client, ObjectID object_id, PlasmaError error); Status ReadSealReply(uint8_t* data, size_t size, ObjectID* object_id); /* Plasma Get message functions. */ -Status SendGetRequest(int sock, const ObjectID* object_ids, int64_t num_objects, +Status SendGetRequest(const std::shared_ptr &store_conn, const ObjectID* object_ids, int64_t num_objects, int64_t timeout_ms); Status ReadGetRequest(uint8_t* data, size_t size, std::vector& object_ids, int64_t* timeout_ms); -Status SendGetReply(int sock, ObjectID object_ids[], +Status SendGetReply(const std::shared_ptr &client, ObjectID object_ids[], std::unordered_map& plasma_objects, int64_t num_objects, const std::vector& store_fds, const std::vector& mmap_sizes); @@ -152,21 +152,21 @@ Status ReadGetReply(uint8_t* data, size_t size, ObjectID object_ids[], /* Plasma Release message functions. */ -Status SendReleaseRequest(int sock, ObjectID object_id); +Status SendReleaseRequest(const std::shared_ptr &store_conn, ObjectID object_id); Status ReadReleaseRequest(uint8_t* data, size_t size, ObjectID* object_id); -Status SendReleaseReply(int sock, ObjectID object_id, PlasmaError error); +Status SendReleaseReply(const std::shared_ptr &client, ObjectID object_id, PlasmaError error); Status ReadReleaseReply(uint8_t* data, size_t size, ObjectID* object_id); /* Plasma Delete objects message functions. */ -Status SendDeleteRequest(int sock, const std::vector& object_ids); +Status SendDeleteRequest(const std::shared_ptr &store_conn, const std::vector& object_ids); Status ReadDeleteRequest(uint8_t* data, size_t size, std::vector* object_ids); -Status SendDeleteReply(int sock, const std::vector& object_ids, +Status SendDeleteReply(const std::shared_ptr &client, const std::vector& object_ids, const std::vector& errors); Status ReadDeleteReply(uint8_t* data, size_t size, std::vector* object_ids, @@ -174,47 +174,47 @@ Status ReadDeleteReply(uint8_t* data, size_t size, std::vector* object /* Plasma Contains message functions. */ -Status SendContainsRequest(int sock, ObjectID object_id); +Status SendContainsRequest(const std::shared_ptr &store_conn, ObjectID object_id); Status ReadContainsRequest(uint8_t* data, size_t size, ObjectID* object_id); -Status SendContainsReply(int sock, ObjectID object_id, bool has_object); +Status SendContainsReply(const std::shared_ptr &client, ObjectID object_id, bool has_object); Status ReadContainsReply(uint8_t* data, size_t size, ObjectID* object_id, bool* has_object); /* Plasma Connect message functions. */ -Status SendConnectRequest(int sock); +Status SendConnectRequest(const std::shared_ptr &store_conn); Status ReadConnectRequest(uint8_t* data, size_t size); -Status SendConnectReply(int sock, int64_t memory_capacity); +Status SendConnectReply(const std::shared_ptr &client, int64_t memory_capacity); Status ReadConnectReply(uint8_t* data, size_t size, int64_t* memory_capacity); /* Plasma Evict message functions (no reply so far). */ -Status SendEvictRequest(int sock, int64_t num_bytes); +Status SendEvictRequest(const std::shared_ptr &store_conn, int64_t num_bytes); Status ReadEvictRequest(uint8_t* data, size_t size, int64_t* num_bytes); -Status SendEvictReply(int sock, int64_t num_bytes); +Status SendEvictReply(const std::shared_ptr &client, int64_t num_bytes); Status ReadEvictReply(uint8_t* data, size_t size, int64_t& num_bytes); /* Plasma Subscribe message functions. */ -Status SendSubscribeRequest(int sock); +Status SendSubscribeRequest(const std::shared_ptr &store_conn); /* Data messages. */ -Status SendDataRequest(int sock, ObjectID object_id, const char* address, int port); +Status SendDataRequest(const std::shared_ptr &store_conn, ObjectID object_id, const char* address, int port); Status ReadDataRequest(uint8_t* data, size_t size, ObjectID* object_id, char** address, int* port); -Status SendDataReply(int sock, ObjectID object_id, int64_t object_size, +Status SendDataReply(const std::shared_ptr &client, ObjectID object_id, int64_t object_size, int64_t metadata_size); Status ReadDataReply(uint8_t* data, size_t size, ObjectID* object_id, @@ -222,12 +222,12 @@ Status ReadDataReply(uint8_t* data, size_t size, ObjectID* object_id, /* Plasma refresh LRU cache functions. */ -Status SendRefreshLRURequest(int sock, const std::vector& object_ids); +Status SendRefreshLRURequest(const std::shared_ptr &store_conn, const std::vector& object_ids); Status ReadRefreshLRURequest(uint8_t* data, size_t size, std::vector* object_ids); -Status SendRefreshLRUReply(int sock); +Status SendRefreshLRUReply(const std::shared_ptr &client); Status ReadRefreshLRUReply(uint8_t* data, size_t size); diff --git a/src/ray/object_manager/plasma/store.cc b/src/ray/object_manager/plasma/store.cc index 6eb384db7..a8a9c286b 100644 --- a/src/ray/object_manager/plasma/store.cc +++ b/src/ray/object_manager/plasma/store.cc @@ -62,9 +62,9 @@ namespace fb = plasma::flatbuf; namespace plasma { struct GetRequest { - GetRequest(Client* client, const std::vector& object_ids); + GetRequest(const std::shared_ptr &client, const std::vector& object_ids); /// The client that called get. - Client* client; + std::shared_ptr client; /// The ID of the timer that will time out and cause this wait to return to /// the client if it hasn't already returned. int64_t timer; @@ -80,7 +80,7 @@ struct GetRequest { int64_t num_satisfied; }; -GetRequest::GetRequest(Client* client, const std::vector& object_ids) +GetRequest::GetRequest(const std::shared_ptr &client, const std::vector& object_ids) : client(client), timer(-1), object_ids(object_ids.begin(), object_ids.end()), @@ -90,8 +90,6 @@ GetRequest::GetRequest(Client* client, const std::vector& object_ids) num_objects_to_wait_for = unique_ids.size(); } -Client::Client(int fd) : fd(fd), notification_fd(-1) {} - PlasmaStore::PlasmaStore(EventLoop* loop, std::string directory, bool hugepages_enabled, const std::string& socket_name, std::shared_ptr external_store) @@ -115,7 +113,7 @@ const PlasmaStoreInfo* PlasmaStore::GetPlasmaStoreInfo() { return &store_info_; // If this client is not already using the object, add the client to the // object's list of clients, otherwise do nothing. void PlasmaStore::AddToClientObjectIds(const ObjectID& object_id, ObjectTableEntry* entry, - Client* client) { + const std::shared_ptr &client) { // Check if this client is already using the object. if (client->object_ids.find(object_id) != client->object_ids.end()) { return; @@ -135,12 +133,12 @@ void PlasmaStore::AddToClientObjectIds(const ObjectID& object_id, ObjectTableEnt // Allocate memory uint8_t* PlasmaStore::AllocateMemory(size_t size, bool evict_if_full, int* fd, - int64_t* map_size, ptrdiff_t* offset, Client* client, + int64_t* map_size, ptrdiff_t* offset, const std::shared_ptr &client, bool is_create) { // First free up space from the client's LRU queue if quota enforcement is on. if (evict_if_full) { std::vector client_objects_to_evict; - bool quota_ok = eviction_policy_.EnforcePerClientQuota(client, size, is_create, + bool quota_ok = eviction_policy_.EnforcePerClientQuota(client.get(), size, is_create, &client_objects_to_evict); if (!quota_ok) { return nullptr; @@ -205,7 +203,7 @@ Status PlasmaStore::FreeCudaMemory(int device_num, int64_t size, uint8_t* pointe // Create a new object buffer in the hash table. PlasmaError PlasmaStore::CreateObject(const ObjectID& object_id, bool evict_if_full, int64_t data_size, int64_t metadata_size, - int device_num, Client* client, + int device_num, const std::shared_ptr &client, PlasmaObject* result) { RAY_LOG(DEBUG) << "creating object " << object_id.Hex(); @@ -275,7 +273,7 @@ PlasmaError PlasmaStore::CreateObject(const ObjectID& object_id, bool evict_if_f // Notify the eviction policy that this object was created. This must be done // immediately before the call to AddToClientObjectIds so that the // eviction policy does not have an opportunity to evict the object. - eviction_policy_.ObjectCreated(object_id, client, true); + eviction_policy_.ObjectCreated(object_id, client.get(), true); // Record that this client is using this object. AddToClientObjectIds(object_id, store_info_.objects[object_id].get(), client); return PlasmaError::OK; @@ -324,7 +322,7 @@ void PlasmaStore::RemoveGetRequest(GetRequest* get_request) { delete get_request; } -void PlasmaStore::RemoveGetRequestsForClient(Client* client) { +void PlasmaStore::RemoveGetRequestsForClient(const std::shared_ptr &client) { std::unordered_set get_requests_to_remove; for (auto const& pair : object_get_requests_) { for (GetRequest* get_request : pair.second) { @@ -358,7 +356,7 @@ void PlasmaStore::ReturnFromGet(GetRequest* get_req) { } // Send the get reply to the client. - Status s = SendGetReply(get_req->client->fd, &get_req->object_ids[0], get_req->objects, + Status s = SendGetReply(get_req->client, &get_req->object_ids[0], get_req->objects, get_req->object_ids.size(), store_fds, mmap_sizes); WarnIfSigpipe(s.ok() ? 0 : -1, get_req->client->fd); // If we successfully sent the get reply message to the client, then also send @@ -425,7 +423,7 @@ void PlasmaStore::UpdateObjectGetRequests(const ObjectID& object_id) { } } -void PlasmaStore::ProcessGetRequest(Client* client, +void PlasmaStore::ProcessGetRequest(const std::shared_ptr &client, const std::vector& object_ids, int64_t timeout_ms) { // Create a get request for this object. @@ -453,7 +451,7 @@ void PlasmaStore::ProcessGetRequest(Client* client, if (entry->pointer) { entry->state = ObjectState::PLASMA_CREATED; entry->create_time = std::time(nullptr); - eviction_policy_.ObjectCreated(object_id, client, false); + eviction_policy_.ObjectCreated(object_id, client.get(), false); AddToClientObjectIds(object_id, store_info_.objects[object_id].get(), client); evicted_ids.push_back(object_id); evicted_entries.push_back(entry); @@ -513,7 +511,7 @@ void PlasmaStore::ProcessGetRequest(Client* client, } int PlasmaStore::RemoveFromClientObjectIds(const ObjectID& object_id, - ObjectTableEntry* entry, Client* client) { + ObjectTableEntry* entry, const std::shared_ptr &client) { auto it = client->object_ids.find(object_id); if (it != client->object_ids.end()) { client->object_ids.erase(it); @@ -554,7 +552,7 @@ void PlasmaStore::EraseFromObjectTable(const ObjectID& object_id) { store_info_.objects.erase(object_id); } -void PlasmaStore::ReleaseObject(const ObjectID& object_id, Client* client) { +void PlasmaStore::ReleaseObject(const ObjectID& object_id, const std::shared_ptr &client) { auto entry = GetObjectTableEntry(&store_info_, object_id); RAY_CHECK(entry != nullptr); // Remove the client from the object's array of clients. @@ -597,7 +595,7 @@ void PlasmaStore::SealObjects(const std::vector& object_ids) { } } -int PlasmaStore::AbortObject(const ObjectID& object_id, Client* client) { +int PlasmaStore::AbortObject(const ObjectID& object_id, const std::shared_ptr &client) { auto entry = GetObjectTableEntry(&store_info_, object_id); RAY_CHECK(entry != nullptr) << "To abort an object it must be in the object table."; RAY_CHECK(entry->state != ObjectState::PLASMA_SEALED) @@ -700,9 +698,7 @@ void PlasmaStore::EvictObjects(const std::vector& object_ids) { void PlasmaStore::ConnectClient(int listener_sock) { int client_fd = AcceptClient(listener_sock); - - Client* client = new Client(client_fd); - connected_clients_[client_fd] = std::unique_ptr(client); + auto client = std::make_shared(client_fd); // Add a callback to handle events on this socket. // TODO(pcm): Check return value. @@ -712,20 +708,15 @@ void PlasmaStore::ConnectClient(int listener_sock) { RAY_LOG(FATAL) << "Failed to process file event: " << s; } }); - RAY_LOG(DEBUG) << "New connection with fd " << client_fd; + RAY_LOG(DEBUG) << "New connection with fd " << client; } -void PlasmaStore::DisconnectClient(int client_fd) { +void PlasmaStore::DisconnectClient(const std::shared_ptr &client) { + int client_fd = client->fd; RAY_CHECK(client_fd > 0); - auto it = connected_clients_.find(client_fd); - RAY_CHECK(it != connected_clients_.end()); - loop_->RemoveFileEvent(client_fd); - // Close the socket. - close(client_fd); - RAY_LOG(DEBUG) << "Disconnecting client on fd " << client_fd; + RAY_LOG(DEBUG) << "Disconnecting client on fd " << client; // Release all the objects that the client was using. - auto client = it->second.get(); - eviction_policy_.ClientDisconnected(client); + eviction_policy_.ClientDisconnected(client.get()); std::unordered_map sealed_objects; for (const auto& object_id : client->object_ids) { auto it = store_info_.objects.find(object_id); @@ -763,7 +754,8 @@ void PlasmaStore::DisconnectClient(int client_fd) { client->notification_fd = -1; } - connected_clients_.erase(it); + // We lose the last borrower of the Client instance here. + loop_->RemoveFileEvent(client_fd); } /// Send notifications about sealed objects to the subscribers. This is called @@ -866,7 +858,7 @@ void PlasmaStore::PushNotification(ObjectInfoT* object_info, int client_fd) { } // Subscribe to notifications about sealed objects. -void PlasmaStore::SubscribeToUpdates(Client* client) { +void PlasmaStore::SubscribeToUpdates(const std::shared_ptr &client) { RAY_LOG(DEBUG) << "subscribing to updates on fd " << client->fd; if (client->notification_fd > 0) { // This client has already subscribed. Return. @@ -879,7 +871,7 @@ void PlasmaStore::SubscribeToUpdates(Client* client) { if (fd < 0) { // This may mean that the client died before sending the file descriptor. RAY_LOG(WARNING) << "Failed to receive file descriptor from client on fd " - << client->fd << "."; + << client << "."; return; } @@ -899,7 +891,7 @@ void PlasmaStore::SubscribeToUpdates(Client* client) { } } -Status PlasmaStore::ProcessMessage(Client* client) { +Status PlasmaStore::ProcessMessage(const std::shared_ptr &client) { fb::MessageType type; Status s = ReadMessage(client->fd, &type, &input_buffer_); RAY_CHECK(s.ok() || s.IsIOError()); @@ -925,7 +917,7 @@ Status PlasmaStore::ProcessMessage(Client* client) { mmap_size = GetMmapSize(object.store_fd); } HANDLE_SIGPIPE( - SendCreateReply(client->fd, object_id, &object, error_code, mmap_size), + SendCreateReply(client, object_id, &object, error_code, mmap_size), client->fd); // Only send the file descriptor if it hasn't been sent (see analogous // logic in GetStoreFd in client.cc). Similar in ReturnFromGet. @@ -963,7 +955,7 @@ Status PlasmaStore::ProcessMessage(Client* client) { } // Reply to the client. - HANDLE_SIGPIPE(SendCreateAndSealReply(client->fd, error_code), client->fd); + HANDLE_SIGPIPE(SendCreateAndSealReply(client, error_code), client->fd); } break; case fb::MessageType::PlasmaCreateAndSealBatchRequest: { bool evict_if_full; @@ -1014,14 +1006,14 @@ Status PlasmaStore::ProcessMessage(Client* client) { } } - HANDLE_SIGPIPE(SendCreateAndSealBatchReply(client->fd, error_code), client->fd); + HANDLE_SIGPIPE(SendCreateAndSealBatchReply(client, error_code), client->fd); } break; case fb::MessageType::PlasmaAbortRequest: { RAY_RETURN_NOT_OK(ReadAbortRequest(input, input_size, &object_id)); RAY_CHECK(AbortObject(object_id, client) == 1) << "To abort an object, the only " "client currently using it " "must be the creator."; - HANDLE_SIGPIPE(SendAbortReply(client->fd, object_id), client->fd); + HANDLE_SIGPIPE(SendAbortReply(client, object_id), client->fd); } break; case fb::MessageType::PlasmaGetRequest: { std::vector object_ids_to_get; @@ -1041,20 +1033,20 @@ Status PlasmaStore::ProcessMessage(Client* client) { for (auto& object_id : object_ids) { error_codes.push_back(DeleteObject(object_id)); } - HANDLE_SIGPIPE(SendDeleteReply(client->fd, object_ids, error_codes), client->fd); + HANDLE_SIGPIPE(SendDeleteReply(client, object_ids, error_codes), client->fd); } break; case fb::MessageType::PlasmaContainsRequest: { RAY_RETURN_NOT_OK(ReadContainsRequest(input, input_size, &object_id)); if (ContainsObject(object_id) == ObjectStatus::OBJECT_FOUND) { - HANDLE_SIGPIPE(SendContainsReply(client->fd, object_id, 1), client->fd); + HANDLE_SIGPIPE(SendContainsReply(client, object_id, 1), client->fd); } else { - HANDLE_SIGPIPE(SendContainsReply(client->fd, object_id, 0), client->fd); + HANDLE_SIGPIPE(SendContainsReply(client, object_id, 0), client->fd); } } break; case fb::MessageType::PlasmaSealRequest: { RAY_RETURN_NOT_OK(ReadSealRequest(input, input_size, &object_id)); SealObjects({object_id}); - HANDLE_SIGPIPE(SendSealReply(client->fd, object_id, PlasmaError::OK), client->fd); + HANDLE_SIGPIPE(SendSealReply(client, object_id, PlasmaError::OK), client->fd); } break; case fb::MessageType::PlasmaEvictRequest: { // This code path should only be used for testing. @@ -1064,24 +1056,24 @@ Status PlasmaStore::ProcessMessage(Client* client) { int64_t num_bytes_evicted = eviction_policy_.ChooseObjectsToEvict(num_bytes, &objects_to_evict); EvictObjects(objects_to_evict); - HANDLE_SIGPIPE(SendEvictReply(client->fd, num_bytes_evicted), client->fd); + HANDLE_SIGPIPE(SendEvictReply(client, num_bytes_evicted), client->fd); } break; case fb::MessageType::PlasmaRefreshLRURequest: { std::vector object_ids; RAY_RETURN_NOT_OK(ReadRefreshLRURequest(input, input_size, &object_ids)); eviction_policy_.RefreshObjects(object_ids); - HANDLE_SIGPIPE(SendRefreshLRUReply(client->fd), client->fd); + HANDLE_SIGPIPE(SendRefreshLRUReply(client), client->fd); } break; case fb::MessageType::PlasmaSubscribeRequest: SubscribeToUpdates(client); break; case fb::MessageType::PlasmaConnectRequest: { - HANDLE_SIGPIPE(SendConnectReply(client->fd, PlasmaAllocator::GetFootprintLimit()), + HANDLE_SIGPIPE(SendConnectReply(client, PlasmaAllocator::GetFootprintLimit()), client->fd); } break; case fb::MessageType::PlasmaDisconnectClient: - RAY_LOG(DEBUG) << "Disconnecting client on fd " << client->fd; - DisconnectClient(client->fd); + RAY_LOG(DEBUG) << "Disconnecting client on fd " << client; + DisconnectClient(client); break; case fb::MessageType::PlasmaSetOptionsRequest: { std::string client_name; @@ -1089,13 +1081,13 @@ Status PlasmaStore::ProcessMessage(Client* client) { RAY_RETURN_NOT_OK( ReadSetOptionsRequest(input, input_size, &client_name, &output_memory_quota)); client->name = client_name; - bool success = eviction_policy_.SetClientQuota(client, output_memory_quota); - HANDLE_SIGPIPE(SendSetOptionsReply(client->fd, success ? PlasmaError::OK - : PlasmaError::OutOfMemory), + bool success = eviction_policy_.SetClientQuota(client.get(), output_memory_quota); + HANDLE_SIGPIPE(SendSetOptionsReply(client, success ? PlasmaError::OK + : PlasmaError::OutOfMemory), client->fd); } break; case fb::MessageType::PlasmaGetDebugStringRequest: { - HANDLE_SIGPIPE(SendGetDebugStringReply(client->fd, eviction_policy_.DebugString()), + HANDLE_SIGPIPE(SendGetDebugStringReply(client, eviction_policy_.DebugString()), client->fd); } break; default: diff --git a/src/ray/object_manager/plasma/store.h b/src/ray/object_manager/plasma/store.h index 128af541d..cf037ecb8 100644 --- a/src/ray/object_manager/plasma/store.h +++ b/src/ray/object_manager/plasma/store.h @@ -94,7 +94,7 @@ class PlasmaStore { /// plasma_release. PlasmaError CreateObject(const ObjectID& object_id, bool evict_if_full, int64_t data_size, int64_t metadata_size, int device_num, - Client* client, PlasmaObject* result); + const std::shared_ptr &client, PlasmaObject* result); /// Abort a created but unsealed object. If the client is not the /// creator, then the abort will fail. @@ -103,7 +103,7 @@ class PlasmaStore { /// \param client The client who created the object. If this does not /// match the creator of the object, then the abort will fail. /// \return 1 if the abort succeeds, else 0. - int AbortObject(const ObjectID& object_id, Client* client); + int AbortObject(const ObjectID& object_id, const std::shared_ptr &client); /// Delete a specific object by object_id that have been created in the hash table. /// @@ -130,7 +130,7 @@ class PlasmaStore { /// \param client The client making this request. /// \param object_ids Object IDs of the objects to be gotten. /// \param timeout_ms The timeout for the get request in milliseconds. - void ProcessGetRequest(Client* client, const std::vector& object_ids, + void ProcessGetRequest(const std::shared_ptr &client, const std::vector& object_ids, int64_t timeout_ms); /// Seal a vector of objects. The objects are now immutable and can be accessed with @@ -150,12 +150,12 @@ class PlasmaStore { /// /// \param object_id The object ID of the object that is being released. /// \param client The client making this request. - void ReleaseObject(const ObjectID& object_id, Client* client); + void ReleaseObject(const ObjectID& object_id, const std::shared_ptr &client); /// Subscribe a file descriptor to updates about new sealed objects. /// /// \param client The client making this request. - void SubscribeToUpdates(Client* client); + void SubscribeToUpdates(const std::shared_ptr &client); /// Connect a new client to the PlasmaStore. /// @@ -164,12 +164,12 @@ class PlasmaStore { /// Disconnect a client from the PlasmaStore. /// - /// \param client_fd The client file descriptor that is disconnected. - void DisconnectClient(int client_fd); + /// \param client The client that is disconnected. + void DisconnectClient(const std::shared_ptr &client); NotificationMap::iterator SendNotifications(NotificationMap::iterator it); - Status ProcessMessage(Client* client); + Status ProcessMessage(const std::shared_ptr &client); void SetNotificationListener( const std::shared_ptr ¬ification_listener) { @@ -196,7 +196,7 @@ class PlasmaStore { void PushNotification(ObjectInfoT* object_notification, int client_fd); void AddToClientObjectIds(const ObjectID& object_id, ObjectTableEntry* entry, - Client* client); + const std::shared_ptr &client); /// Remove a GetRequest and clean up the relevant data structures. /// @@ -206,19 +206,19 @@ class PlasmaStore { /// Remove all of the GetRequests for a given client. /// /// \param client The client whose GetRequests should be removed. - void RemoveGetRequestsForClient(Client* client); + void RemoveGetRequestsForClient(const std::shared_ptr &client); void ReturnFromGet(GetRequest* get_req); void UpdateObjectGetRequests(const ObjectID& object_id); int RemoveFromClientObjectIds(const ObjectID& object_id, ObjectTableEntry* entry, - Client* client); + const std::shared_ptr &client); void EraseFromObjectTable(const ObjectID& object_id); uint8_t* AllocateMemory(size_t size, bool evict_if_full, int* fd, int64_t* map_size, - ptrdiff_t* offset, Client* client, bool is_create); + ptrdiff_t* offset, const std::shared_ptr &client, bool is_create); #ifdef PLASMA_CUDA Status AllocateCudaMemory(int device_num, int64_t size, uint8_t** out_pointer, std::shared_ptr* out_ipc_handle); @@ -246,8 +246,6 @@ class PlasmaStore { /// reorganize the code slightly. NotificationMap pending_notifications_; - std::unordered_map> connected_clients_; - std::unordered_set deletion_cache_; /// Manages worker threads for handling asynchronous/multi-threaded requests