wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-08 06:58:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Core] Plasma RAII support (#9370)

Browse Source
This commit is contained in:
Siyuan (Ryans) Zhuang
2020-07-10 09:22:29 -07:00
committed by GitHub
parent d8a0d76d02
commit 1798deae94
7 changed files with 212 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ray/object_manager/plasma/client.cc
+12 -14
View File
@@ -337,7 +337,7 @@ class PlasmaClient::Impl : public std::enable_shared_from_this<PlasmaClient::Imp
const uint8_t* metadata, int64_t metadata_size);
/// File descriptor of the Unix domain socket that connects to the store.
int store_conn_;
std::shared_ptr<StoreConn> 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_this<PlasmaClient::Imp
PlasmaBuffer::~PlasmaBuffer() { RAY_UNUSED(client_->Release(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<std::recursive_mutex> 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<uint8_t> 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<uint8_t> 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<ObjectID>& object_ids) {
RAY_RETURN_NOT_OK(SendRefreshLRURequest(store_conn_, object_ids));
std::vector<uint8_t> 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<std::recursive_mutex> 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();
}
src/ray/object_manager/plasma/plasma.cc
+14
View File
@@ -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> &client) {
os << client->fd;
return os;
}
StoreConn::~StoreConn() { close(fd); }
std::ostream &operator<<(std::ostream &os, const std::shared_ptr<StoreConn> &store_conn) {
os << store_conn->fd;
return os;
}
ObjectTableEntry::ObjectTableEntry() : pointer(nullptr), ref_count(0) {}
ObjectTableEntry::~ObjectTableEntry() { pointer = nullptr; }
src/ray/object_manager/plasma/plasma.h
+17 -1
View File
@@ -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> &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<StoreConn> &store_conn);
// TODO(pcm): Replace this by the flatbuffers message PlasmaObjectSpec.
struct PlasmaObject {
#ifdef PLASMA_CUDA
src/ray/object_manager/plasma/protocol.cc
+82 -66
View File
@@ -69,9 +69,12 @@ flatbuffers::Offset<flatbuffers::Vector<int64_t>> ToFlatbuffer(
return fbb->CreateVector(arrow::util::MakeNonNull(data.data()), data.size());
}
Status PlasmaReceive(int sock, MessageType message_type, std::vector<uint8_t>* buffer) {
Status PlasmaReceive(const std::shared_ptr<StoreConn> &store_conn, MessageType message_type, std::vector<uint8_t>* 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<int64_t>(type)
<< ", message_type = " << static_cast<int64_t>(message_type);
@@ -101,10 +104,23 @@ void ConvertToVector(const FlatbufferVectorPointer fbvector, std::vector<T>* out
}
template <typename Message>
Status PlasmaSend(int sock, MessageType message_type, flatbuffers::FlatBufferBuilder* fbb,
Status PlasmaSend(const std::shared_ptr<StoreConn> &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 <typename Message>
Status PlasmaSend(const std::shared_ptr<Client> &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<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &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<ObjectID>& object_ids,
Status SendCreateAndSealBatchRequest(const std::shared_ptr<StoreConn> &store_conn, const std::vector<ObjectID>& object_ids,
bool evict_if_full,
const std::vector<std::string>& data,
const std::vector<std::string>& metadata) {
@@ -295,7 +311,7 @@ Status SendCreateAndSealBatchRequest(int sock, const std::vector<ObjectID>& 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> &client, PlasmaError error) {
flatbuffers::FlatBufferBuilder fbb;
auto message = fb::CreatePlasmaCreateAndSealReply(fbb, static_cast<PlasmaError>(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> &client, PlasmaError error) {
flatbuffers::FlatBufferBuilder fbb;
auto message =
fb::CreatePlasmaCreateAndSealBatchReply(fbb, static_cast<PlasmaError>(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<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &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> &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<ObjectID>& object_ids) {
Status SendDeleteRequest(const std::shared_ptr<StoreConn> &store_conn, const std::vector<ObjectID>& object_ids) {
flatbuffers::FlatBufferBuilder fbb;
auto message = fb::CreatePlasmaDeleteRequest(
fbb, static_cast<int32_t>(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<ObjectID>* object_ids) {
@@ -463,7 +479,7 @@ Status ReadDeleteRequest(uint8_t* data, size_t size, std::vector<ObjectID>* obje
return Status::OK();
}
Status SendDeleteReply(int sock, const std::vector<ObjectID>& object_ids,
Status SendDeleteReply(const std::shared_ptr<Client> &client, const std::vector<ObjectID>& object_ids,
const std::vector<PlasmaError>& errors) {
RAY_DCHECK(object_ids.size() == errors.size());
flatbuffers::FlatBufferBuilder fbb;
@@ -473,7 +489,7 @@ Status SendDeleteReply(int sock, const std::vector<ObjectID>& object_ids,
fbb.CreateVector(
arrow::util::MakeNonNull(reinterpret_cast<const int32_t*>(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<ObjectID>* object_ids,
@@ -496,11 +512,11 @@ Status ReadDeleteReply(uint8_t* data, size_t size, std::vector<ObjectID>* object
// Contains messages.
Status SendContainsRequest(int sock, ObjectID object_id) {
Status SendContainsRequest(const std::shared_ptr<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &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<ObjectID>& object_ids,
@@ -605,7 +621,7 @@ Status ReadGetRequest(uint8_t* data, size_t size, std::vector<ObjectID>& object_
return Status::OK();
}
Status SendGetReply(int sock, ObjectID object_ids[],
Status SendGetReply(const std::shared_ptr<Client> &client, ObjectID object_ids[],
std::unordered_map<ObjectID, PlasmaObject>& plasma_objects,
int64_t num_objects, const std::vector<int>& store_fds,
const std::vector<int64_t>& 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<StoreConn> &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<StoreConn> &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> &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<ObjectID>& object_ids) {
Status SendRefreshLRURequest(const std::shared_ptr<StoreConn> &store_conn, const std::vector<ObjectID>& 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> &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) {
src/ray/object_manager/plasma/protocol.h
+32 -32
View File
@@ -52,51 +52,51 @@ flatbuffers::Offset<flatbuffers::Vector<int64_t>> ToFlatbuffer(
/* Plasma receive message. */
Status PlasmaReceive(int sock, MessageType message_type, std::vector<uint8_t>* buffer);
Status PlasmaReceive(const std::shared_ptr<StoreConn> &store_conn, MessageType message_type, std::vector<uint8_t>* buffer);
/* Set options messages. */
Status SendSetOptionsRequest(int sock, const std::string& client_name,
Status SendSetOptionsRequest(const std::shared_ptr<StoreConn> &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> &client, PlasmaError error);
Status ReadSetOptionsReply(uint8_t* data, size_t size);
/* Debug string messages. */
Status SendGetDebugStringRequest(int sock);
Status SendGetDebugStringRequest(const std::shared_ptr<StoreConn> &store_conn);
Status SendGetDebugStringReply(int sock, const std::string& debug_string);
Status SendGetDebugStringReply(const std::shared_ptr<Client> &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<StoreConn> &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> &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<StoreConn> &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<ObjectID>& object_ids,
Status SendCreateAndSealBatchRequest(const std::shared_ptr<StoreConn> &store_conn, const std::vector<ObjectID>& object_ids,
bool evict_if_full,
const std::vector<std::string>& data,
const std::vector<std::string>& metadata);
@@ -107,41 +107,41 @@ Status ReadCreateAndSealBatchRequest(uint8_t* data, size_t size,
std::vector<std::string>* object_data,
std::vector<std::string>* metadata);
Status SendCreateAndSealReply(int sock, PlasmaError error);
Status SendCreateAndSealReply(const std::shared_ptr<Client> &client, PlasmaError error);
Status ReadCreateAndSealReply(uint8_t* data, size_t size);
Status SendCreateAndSealBatchReply(int sock, PlasmaError error);
Status SendCreateAndSealBatchReply(const std::shared_ptr<Client> &client, PlasmaError error);
Status ReadCreateAndSealBatchReply(uint8_t* data, size_t size);
Status SendAbortRequest(int sock, ObjectID object_id);
Status SendAbortRequest(const std::shared_ptr<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &store_conn, const ObjectID* object_ids, int64_t num_objects,
int64_t timeout_ms);
Status ReadGetRequest(uint8_t* data, size_t size, std::vector<ObjectID>& object_ids,
int64_t* timeout_ms);
Status SendGetReply(int sock, ObjectID object_ids[],
Status SendGetReply(const std::shared_ptr<Client> &client, ObjectID object_ids[],
std::unordered_map<ObjectID, PlasmaObject>& plasma_objects,
int64_t num_objects, const std::vector<int>& store_fds,
const std::vector<int64_t>& 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<StoreConn> &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> &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<ObjectID>& object_ids);
Status SendDeleteRequest(const std::shared_ptr<StoreConn> &store_conn, const std::vector<ObjectID>& object_ids);
Status ReadDeleteRequest(uint8_t* data, size_t size, std::vector<ObjectID>* object_ids);
Status SendDeleteReply(int sock, const std::vector<ObjectID>& object_ids,
Status SendDeleteReply(const std::shared_ptr<Client> &client, const std::vector<ObjectID>& object_ids,
const std::vector<PlasmaError>& errors);
Status ReadDeleteReply(uint8_t* data, size_t size, std::vector<ObjectID>* object_ids,
@@ -174,47 +174,47 @@ Status ReadDeleteReply(uint8_t* data, size_t size, std::vector<ObjectID>* object
/* Plasma Contains message functions. */
Status SendContainsRequest(int sock, ObjectID object_id);
Status SendContainsRequest(const std::shared_ptr<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &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> &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<StoreConn> &store_conn);
/* Data messages. */
Status SendDataRequest(int sock, ObjectID object_id, const char* address, int port);
Status SendDataRequest(const std::shared_ptr<StoreConn> &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> &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<ObjectID>& object_ids);
Status SendRefreshLRURequest(const std::shared_ptr<StoreConn> &store_conn, const std::vector<ObjectID>& object_ids);
Status ReadRefreshLRURequest(uint8_t* data, size_t size,
std::vector<ObjectID>* object_ids);
Status SendRefreshLRUReply(int sock);
Status SendRefreshLRUReply(const std::shared_ptr<Client> &client);
Status ReadRefreshLRUReply(uint8_t* data, size_t size);
src/ray/object_manager/plasma/store.cc
+43 -51
View File
@@ -62,9 +62,9 @@ namespace fb = plasma::flatbuf;
namespace plasma {
struct GetRequest {
GetRequest(Client* client, const std::vector<ObjectID>& object_ids);
GetRequest(const std::shared_ptr<Client> &client, const std::vector<ObjectID>& object_ids);
/// The client that called get.
Client* client;
std::shared_ptr<Client> 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<ObjectID>& object_ids)
GetRequest::GetRequest(const std::shared_ptr<Client> &client, const std::vector<ObjectID>& 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<ObjectID>& 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<ExternalStore> 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> &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> &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<ObjectID> 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> &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> &client) {
std::unordered_set<GetRequest*> 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> &client,
const std::vector<ObjectID>& 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> &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> &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<ObjectID>& object_ids) {
}
}
int PlasmaStore::AbortObject(const ObjectID& object_id, Client* client) {
int PlasmaStore::AbortObject(const ObjectID& object_id, const std::shared_ptr<Client> &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<ObjectID>& 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>(client);
auto client = std::make_shared<Client>(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> &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<ObjectID, ObjectTableEntry*> 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> &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> &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<ObjectID> 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<ObjectID> 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:
src/ray/object_manager/plasma/store.h
+12 -14
View File
@@ -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> &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> &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<ObjectID>& object_ids,
void ProcessGetRequest(const std::shared_ptr<Client> &client, const std::vector<ObjectID>& 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> &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> &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> &client);
NotificationMap::iterator SendNotifications(NotificationMap::iterator it);
Status ProcessMessage(Client* client);
Status ProcessMessage(const std::shared_ptr<Client> &client);
void SetNotificationListener(
const std::shared_ptr<ray::ObjectStoreNotificationManager> &notification_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> &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> &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> &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> &client, bool is_create);
#ifdef PLASMA_CUDA
Status AllocateCudaMemory(int device_num, int64_t size, uint8_t** out_pointer,
std::shared_ptr<CudaIpcMemHandle>* out_ipc_handle);
@@ -246,8 +246,6 @@ class PlasmaStore {
/// reorganize the code slightly.
NotificationMap pending_notifications_;
std::unordered_map<int, std::unique_ptr<Client>> connected_clients_;
std::unordered_set<ObjectID> deletion_cache_;
/// Manages worker threads for handling asynchronous/multi-threaded requests