Scheduling ids (#6137)

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"],

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(); }

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 <string>
+#include <unordered_map>
+
+/// Class to map string IDs to unique integer IDs and back.
+class StringIdMap {
+  std::unordered_map<std::string, int64_t> string_to_int_;
+  std::unordered_map<int64_t, std::string> int_to_string_;
+  std::hash<std::string> 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

src/ray/common/scheduling/scheduling_test.cc

@@ -0,0 +1,54 @@
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include <string>
+#include <thread>
+
+#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<string> 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<int64_t>(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();
+}