diff --git a/src/ray/common/common_protocol.h b/src/ray/common/common_protocol.h index bc3d9b646..792bb173b 100644 --- a/src/ray/common/common_protocol.h +++ b/src/ray/common/common_protocol.h @@ -104,15 +104,13 @@ string_vec_to_flatbuf(flatbuffers::FlatBufferBuilder &fbb, template flatbuffers::Offset to_flatbuf(flatbuffers::FlatBufferBuilder &fbb, ID id) { - return fbb.CreateString(reinterpret_cast(id.data()), sizeof(ID)); + return fbb.CreateString(reinterpret_cast(id.data()), id.size()); } template ID from_flatbuf(const flatbuffers::String &string) { - ID id; - RAY_CHECK(string.size() == sizeof(ID)); - memcpy(id.mutable_data(), string.data(), sizeof(ID)); - return id; + RAY_CHECK(string.size() == ID::size()); + return ID::from_binary(string.str()); } template @@ -129,12 +127,13 @@ template const std::vector ids_from_flatbuf(const flatbuffers::String &string) { const auto &ids = string_from_flatbuf(string); std::vector ret; - RAY_CHECK(ids.size() % kUniqueIDSize == 0); - auto count = ids.size() / kUniqueIDSize; + size_t id_size = ID::size(); + RAY_CHECK(ids.size() % id_size == 0); + auto count = ids.size() / id_size; for (size_t i = 0; i < count; ++i) { - auto pos = static_cast(kUniqueIDSize * i); - const auto &id = ids.substr(pos, kUniqueIDSize); + auto pos = static_cast(id_size * i); + const auto &id = ids.substr(pos, id_size); ret.push_back(ID::from_binary(id)); } diff --git a/src/ray/id.cc b/src/ray/id.cc index a9d9c5a7e..8d72cef8b 100644 --- a/src/ray/id.cc +++ b/src/ray/id.cc @@ -24,18 +24,23 @@ std::mt19937 RandomlySeededMersenneTwister() { return seeded_engine; } +uint64_t MurmurHash64A(const void *key, int len, unsigned int seed); + UniqueID::UniqueID() { // Set the ID to nil. std::fill_n(id_, kUniqueIDSize, 255); } +UniqueID::UniqueID(const std::string &binary) { + std::memcpy(&id_, binary.data(), kUniqueIDSize); +} + UniqueID::UniqueID(const plasma::UniqueID &from) { std::memcpy(&id_, from.data(), kUniqueIDSize); } UniqueID UniqueID::from_random() { - UniqueID id; - uint8_t *data = id.mutable_data(); + std::string data(kUniqueIDSize, 0); // NOTE(pcm): The right way to do this is to have one std::mt19937 per // thread (using the thread_local keyword), but that's not supported on // older versions of macOS (see https://stackoverflow.com/a/29929949) @@ -46,14 +51,10 @@ UniqueID UniqueID::from_random() { for (int i = 0; i < kUniqueIDSize; i++) { data[i] = static_cast(dist(generator)); } - return id; + return UniqueID::from_binary(data); } -UniqueID UniqueID::from_binary(const std::string &binary) { - UniqueID id; - std::memcpy(&id, binary.data(), sizeof(id)); - return id; -} +UniqueID UniqueID::from_binary(const std::string &binary) { return UniqueID(binary); } const UniqueID &UniqueID::nil() { static const UniqueID nil_id; @@ -72,9 +73,7 @@ bool UniqueID::is_nil() const { const uint8_t *UniqueID::data() const { return id_; } -uint8_t *UniqueID::mutable_data() { return id_; } - -size_t UniqueID::size() const { return kUniqueIDSize; } +size_t UniqueID::size() { return kUniqueIDSize; } std::string UniqueID::binary() const { return std::string(reinterpret_cast(id_), kUniqueIDSize); @@ -152,7 +151,14 @@ uint64_t MurmurHash64A(const void *key, int len, unsigned int seed) { return h; } -size_t UniqueID::hash() const { return MurmurHash64A(&id_[0], kUniqueIDSize, 0); } +size_t UniqueID::hash() const { + // Note(ashione): hash code lazy calculation(it's invoked every time if hash code is + // default value 0) + if (!hash_) { + hash_ = MurmurHash64A(&id_[0], kUniqueIDSize, 0); + } + return hash_; +} std::ostream &operator<<(std::ostream &os, const UniqueID &id) { if (id.is_nil()) { @@ -206,19 +212,15 @@ const TaskID GenerateTaskId(const DriverID &driver_id, const TaskID &parent_task // Compute hashes. SHA256_CTX ctx; sha256_init(&ctx); - sha256_update(&ctx, (BYTE *)&driver_id, sizeof(driver_id)); - sha256_update(&ctx, (BYTE *)&parent_task_id, sizeof(parent_task_id)); - sha256_update(&ctx, (BYTE *)&parent_task_counter, sizeof(parent_task_counter)); + sha256_update(&ctx, reinterpret_cast(driver_id.data()), driver_id.size()); + sha256_update(&ctx, reinterpret_cast(parent_task_id.data()), + parent_task_id.size()); + sha256_update(&ctx, (const BYTE *)&parent_task_counter, sizeof(parent_task_counter)); // Compute the final task ID from the hash. BYTE buff[DIGEST_SIZE]; sha256_final(&ctx, buff); - TaskID task_id; - RAY_DCHECK(sizeof(task_id) <= DIGEST_SIZE); - memcpy(&task_id, buff, sizeof(task_id)); - task_id = FinishTaskId(task_id); - - return task_id; + return FinishTaskId(TaskID::from_binary(std::string(buff, buff + kUniqueIDSize))); } int64_t ComputeObjectIndex(const ObjectID &object_id) { diff --git a/src/ray/id.h b/src/ray/id.h index 35c67b220..9467c1a3f 100644 --- a/src/ray/id.h +++ b/src/ray/id.h @@ -24,33 +24,39 @@ class RAY_EXPORT UniqueID { bool operator==(const UniqueID &rhs) const; bool operator!=(const UniqueID &rhs) const; const uint8_t *data() const; - uint8_t *mutable_data(); - size_t size() const; + static size_t size(); std::string binary() const; std::string hex() const; plasma::UniqueID to_plasma_id() const; + private: + UniqueID(const std::string &binary); + protected: uint8_t id_[kUniqueIDSize]; + mutable size_t hash_ = 0; }; static_assert(std::is_standard_layout::value, "UniqueID must be standard"); std::ostream &operator<<(std::ostream &os, const UniqueID &id); -#define DEFINE_UNIQUE_ID(type) \ - class RAY_EXPORT type : public UniqueID { \ - public: \ - explicit type(const UniqueID &from) { \ - std::memcpy(&id_, from.data(), kUniqueIDSize); \ - } \ - type() : UniqueID() {} \ - static type from_random() { return type(UniqueID::from_random()); } \ - static type from_binary(const std::string &binary) { return type(binary); } \ - static type nil() { return type(UniqueID::nil()); } \ - \ - private: \ - type(const std::string &binary) { std::memcpy(id_, binary.data(), kUniqueIDSize); } \ +#define DEFINE_UNIQUE_ID(type) \ + class RAY_EXPORT type : public UniqueID { \ + public: \ + explicit type(const UniqueID &from) { \ + std::memcpy(&id_, from.data(), kUniqueIDSize); \ + } \ + type() : UniqueID() {} \ + static type from_random() { return type(UniqueID::from_random()); } \ + static type from_binary(const std::string &binary) { return type(binary); } \ + static type nil() { return type(UniqueID::nil()); } \ + static size_t size() { return kUniqueIDSize; } \ + \ + private: \ + explicit type(const std::string &binary) { \ + std::memcpy(&id_, binary.data(), kUniqueIDSize); \ + } \ }; #include "id_def.h" diff --git a/src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.cc b/src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.cc index fb7e7b373..eb9d2f0e5 100644 --- a/src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.cc +++ b/src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.cc @@ -11,8 +11,11 @@ class UniqueIdFromJByteArray { public: const ID &GetId() const { return id; } - UniqueIdFromJByteArray(JNIEnv *env, const jbyteArray& bytes) { - env->GetByteArrayRegion(bytes, 0, kUniqueIDSize, reinterpret_cast(id.mutable_data())); + UniqueIdFromJByteArray(JNIEnv *env, const jbyteArray &bytes) { + std::string id_str(ID::size(), 0); + env->GetByteArrayRegion(bytes, 0, ID::size(), + reinterpret_cast(&id_str.front())); + id = ID::from_binary(id_str); } private: @@ -228,11 +231,12 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeGenerateTaskId( TaskID task_id = ray::GenerateTaskId(driver_id.GetId(), parent_task_id.GetId(), parent_task_counter); - jbyteArray result = env->NewByteArray(sizeof(TaskID)); + jbyteArray result = env->NewByteArray(task_id.size()); if (nullptr == result) { return nullptr; } - env->SetByteArrayRegion(result, 0, sizeof(TaskID), reinterpret_cast(&task_id)); + env->SetByteArrayRegion(result, 0, task_id.size(), + reinterpret_cast(task_id.data())); return result; } @@ -276,9 +280,9 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativePrepareCheckpoint(JNIEnv *env if (ThrowRayExceptionIfNotOK(env, status)) { return nullptr; } - jbyteArray result = env->NewByteArray(sizeof(ActorCheckpointID)); - env->SetByteArrayRegion(result, 0, sizeof(ActorCheckpointID), - reinterpret_cast(&checkpoint_id)); + jbyteArray result = env->NewByteArray(checkpoint_id.size()); + env->SetByteArrayRegion(result, 0, checkpoint_id.size(), + reinterpret_cast(checkpoint_id.data())); return result; }