From 1b80675206fe69a92aff527e08d53bd7f7110c99 Mon Sep 17 00:00:00 2001 From: Ion Date: Fri, 15 Nov 2019 16:04:16 -0800 Subject: [PATCH] Scheduling ids (#6137) --- BUILD.bazel | 11 ++++ src/ray/common/scheduling/scheduling_ids.cc | 53 +++++++++++++++++++ src/ray/common/scheduling/scheduling_ids.h | 45 ++++++++++++++++ src/ray/common/scheduling/scheduling_test.cc | 54 ++++++++++++++++++++ 4 files changed, 163 insertions(+) create mode 100644 src/ray/common/scheduling/scheduling_ids.cc create mode 100644 src/ray/common/scheduling/scheduling_ids.h create mode 100644 src/ray/common/scheduling/scheduling_test.cc diff --git a/BUILD.bazel b/BUILD.bazel index 7e87dd1ce..3a635a948 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -218,6 +218,7 @@ cc_library( ":ray_util", "@boost//:asio", "@com_github_grpc_grpc//:grpc++", + "@com_google_googletest//:gtest", "@plasma//:plasma_client", ], ) @@ -434,6 +435,16 @@ cc_test( ], ) +cc_test( + name = "scheduling_test", + srcs = ["src/ray/common/scheduling/scheduling_test.cc"], + copts = COPTS, + deps = [ + ":ray_common", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "lineage_cache_test", srcs = ["src/ray/raylet/lineage_cache_test.cc"], diff --git a/src/ray/common/scheduling/scheduling_ids.cc b/src/ray/common/scheduling/scheduling_ids.cc new file mode 100644 index 000000000..b8591c55d --- /dev/null +++ b/src/ray/common/scheduling/scheduling_ids.cc @@ -0,0 +1,53 @@ +#include "scheduling_ids.h" +using namespace std; + +int64_t StringIdMap::get(const string sid) { + auto it = string_to_int_.find(sid); + if (it == string_to_int_.end()) { + return -1; + } else { + return it->second; + } +}; + +int64_t StringIdMap::insert(const string sid, bool test) { + auto sit = string_to_int_.find(sid); + if (sit == string_to_int_.end()) { + int64_t id = test ? hasher_(sid) % 10 : hasher_(sid); + for (int i = 0; true; i++) { + auto it = int_to_string_.find(id); + if (it == int_to_string_.end()) { + /// No hash collision, so associated sid with id. + string_to_int_.insert(make_pair(sid, id)); + int_to_string_.insert(make_pair(id, sid)); + break; + } + id = test ? hasher_(sid + to_string(i)) % 10 : hasher_(sid + to_string(i)); + } + return id; + } else { + return sit->second; + } +}; + +void StringIdMap::remove(const string sid) { + auto sit = string_to_int_.find(sid); + if (sit != string_to_int_.end()) { + uint64_t id = string_to_int_[sid]; + string_to_int_.erase(sit); + auto it = int_to_string_.find(id); + int_to_string_.erase(it); + } +}; + +void StringIdMap::remove(int64_t id) { + auto it = int_to_string_.find(id); + if (it != int_to_string_.end()) { + string sid = int_to_string_[id]; + int_to_string_.erase(it); + auto sit = string_to_int_.find(sid); + string_to_int_.erase(sit); + } +}; + +int64_t StringIdMap::count() { return string_to_int_.size(); } diff --git a/src/ray/common/scheduling/scheduling_ids.h b/src/ray/common/scheduling/scheduling_ids.h new file mode 100644 index 000000000..337bdb564 --- /dev/null +++ b/src/ray/common/scheduling/scheduling_ids.h @@ -0,0 +1,45 @@ +#ifndef RAY_COMMON_SCHEDULING_SCHEDULING_IDS_H +#define RAY_COMMON_SCHEDULING_SCHEDULING_IDS_H + +#include +#include + +/// Class to map string IDs to unique integer IDs and back. +class StringIdMap { + std::unordered_map string_to_int_; + std::unordered_map int_to_string_; + std::hash hasher_; + + public: + StringIdMap(){}; + ~StringIdMap(){}; + + /// Get integer ID associated with an existing string ID. + /// + /// \param String ID. + /// \return The integer ID associated with the given string ID. + int64_t get(const std::string sid); + + /// Insert a string ID and get the associated integer ID. + /// + /// \param String ID to be inserted. + /// \param test: if "true" it specifies that the range of + /// IDs is limited to 0..10 for testing purposes. + /// \return The integer ID associated with string ID sid. + int64_t insert(const std::string sid, bool test = false); + + /// Delete an ID identified by its string format. + /// + /// \param ID to be deleted. + void remove(const std::string sid); + + /// Delete an ID identified by its integer format. + /// + /// \param ID to be deleted. + void remove(int64_t id); + + /// Get number of identifiers. + int64_t count(); +}; + +#endif // RAY_COMMON_SCHEDULING_SCHEDULING_IDS_H diff --git a/src/ray/common/scheduling/scheduling_test.cc b/src/ray/common/scheduling/scheduling_test.cc new file mode 100644 index 000000000..386520b3c --- /dev/null +++ b/src/ray/common/scheduling/scheduling_test.cc @@ -0,0 +1,54 @@ +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +#include +#include + +#include "ray/common/scheduling/scheduling_ids.h" +using namespace std; + +namespace ray { + +class SchedulingTest : public ::testing::Test { + public: + void SetUp() {} + + void Shutdown() {} +}; + +TEST_F(SchedulingTest, SchedulingIdTest) { + StringIdMap ids; + hash hasher; + int num = 10; // should be greater than 10. + + for (int i = 0; i < num; i++) { + ids.insert(to_string(i)); + } + ASSERT_EQ(ids.count(), num); + + ids.remove(to_string(1)); + ASSERT_EQ(ids.count(), num - 1); + + ids.remove(hasher(to_string(2))); + ASSERT_EQ(ids.count(), num - 2); + + ASSERT_TRUE(ids.get(to_string(3)) == static_cast(hasher(to_string(3)))); + + ASSERT_TRUE(ids.get(to_string(100)) == -1); + + /// Test for handling collision. + StringIdMap short_ids; + for (int i = 0; i < 10; i++) { + /// "true" reduces the range of IDs to [0..9] + int64_t id = short_ids.insert(to_string(i), true); + ASSERT_TRUE(id < 10); + } + ASSERT_EQ(short_ids.count(), 10); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}