From 32cabb6032bbda2e8ffb5ad1ca11395bef178b6e Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Tue, 26 Nov 2024 18:26:59 +0000
Subject: [PATCH 01/13] add shared lru cache

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/BUILD                    |   9 ++
 src/ray/util/shared_lru.h             | 135 ++++++++++++++++++++++++++
 src/ray/util/tests/BUILD              |  12 +++
 src/ray/util/tests/shared_lru_test.cc |  41 ++++++++
 4 files changed, 197 insertions(+)
 create mode 100644 src/ray/util/shared_lru.h
 create mode 100644 src/ray/util/tests/shared_lru_test.cc

diff --git a/src/ray/util/BUILD b/src/ray/util/BUILD
index 23d9f1e90150..0a65db3726e1 100644
--- a/src/ray/util/BUILD
+++ b/src/ray/util/BUILD
@@ -55,3 +55,12 @@ cc_library(
     srcs = ["thread_checker.cc"],
     visibility = ["//visibility:public"],
 )
+
+cc_library(
+    name = "shared_lru",
+    hdrs = ["shared_lru.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
new file mode 100644
index 000000000000..1ff1cc15385a
--- /dev/null
+++ b/src/ray/util/shared_lru.h
@@ -0,0 +1,135 @@
+// SharedLruCache is a LRU cache, with all entries shared, which means a single entry
+// could be accessed by multiple getters. When `Get`, a copy of the value is returned, so
+// for heavy-loaded types it's suggested to wrap with `std::shared_ptr<>`.
+//
+// Example usage:
+// SharedLruCache<std::string, std::string> cache{cap};
+// // Put a key-value pair into cache.
+// cache.Put("key", "val");
+//
+// // Get a key-value pair from cache.
+// auto val = ache.Get("key");
+// // Check and consume `val`.
+//
+// TODO(hjiang):
+// 1. Add template arguments for key hash and key equal, to pass into absl::flat_hash_map.
+// 2. Provide a key hash wrapper to save a copy.
+// 3. flat hash map supports heterogeneous lookup, expose `KeyLike` templated interface.
+
+#pragma once
+
+#include <cstdint>
+#include <list>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+
+namespace ray::utils::container {
+
+template <typename Key, typename Val>
+class SharedLruCache final {
+ public:
+  using key_type = Key;
+  using mapped_type = Val;
+
+  // A `max_entries` of 0 means that there is no limit on the number of entries
+  // in the cache.
+  explicit SharedLruCache(size_t max_entries) : max_entries_(max_entries) {}
+
+  ~SharedLruCache() = default;
+
+  // Insert `value` with key `key`. This will replace any previous entry with
+  // the same key.
+  void Put(Key key, Val value) {
+    std::lock_guard lck(mu_);
+    lru_list_.emplace_front(key);
+    Entry new_entry{std::move(value), lru_list_.begin()};
+    cache_[key] = std::move(new_entry);
+
+    if (max_entries_ > 0 && lru_list_.size() > max_entries_) {
+      const auto &stale_key = lru_list_.back();
+      cache_.erase(stale_key);
+      lru_list_.pop_back();
+    }
+  }
+
+  // Delete the entry with key `key`. Return true if the entry was found for
+  // `key`, false if the entry was not found. In both cases, there is no entry
+  // with key `key` existed after the call.
+  bool Delete(const Key &key) {
+    std::lock_guard lck(mu_);
+    auto it = cache_.find(key);
+    if (it == cache_.end()) {
+      return false;
+    }
+    cache_.erase(it);
+    return true;
+  }
+
+  // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
+  // If found, return a copy for the value.
+  std::optional<Val> Get(const Key &key) {
+    std::lock_guard lck(mu_);
+    const auto cache_iter = cache_.find(key);
+    if (cache_iter == cache_.end()) {
+      return std::nullopt;
+    }
+    Val value = std::move(cache_iter->second.value);
+    lru_list_.erase(cache_iter->second.lru_iterator);
+    cache_.erase(cache_iter);
+
+    // Re-insert into the cache, no need to check capacity.
+    lru_list_.emplace_front(key);
+    Entry new_entry{value, lru_list_.begin()};
+    cache_[key] = std::move(new_entry);
+
+    return value;
+  }
+
+  // Clear the cache.
+  void Clear() {
+    std::lock_guard lck(mu_);
+    cache_.clear();
+    lru_list_.clear();
+  }
+
+  // Accessors for cache parameters.
+  size_t max_entries() const { return max_entries_; }
+
+ private:
+  struct Entry {
+    // The entry's value.
+    Val value;
+
+    // A list iterator pointing to the entry's position in the LRU list.
+    typename std::list<Key>::iterator lru_iterator;
+  };
+
+  using EntryMap = absl::flat_hash_map<Key, Entry>;
+
+  // The maximum number of entries in the cache. A value of 0 means there is no
+  // limit on entry count.
+  const size_t max_entries_;
+
+  // Guards access to the cache and the LRU list.
+  std::mutex mu_;
+
+  // All keys are stored as refernce (`std::reference_wrapper`), and the
+  // ownership lies in `lru_list_`.
+  EntryMap cache_;
+
+  // The LRU list of entries. The front of the list identifies the most
+  // recently accessed entry.
+  std::list<Key> lru_list_;
+};
+
+// Same interfaces as `SharedLruCache`, but all cached values are
+// `const`-specified to avoid concurrent updates.
+template <typename K, typename V>
+using SharedLruConstCache = SharedLruCache<K, const V>;
+
+}  // namespace ray::utils::container
diff --git a/src/ray/util/tests/BUILD b/src/ray/util/tests/BUILD
index 2941d105cf91..b85c01f28ebf 100644
--- a/src/ray/util/tests/BUILD
+++ b/src/ray/util/tests/BUILD
@@ -194,3 +194,15 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+cc_test(
+    name = "shared_lru_test",
+    srcs = ["shared_lru_test.cc"],
+    deps = [
+        "//src/ray/util:shared_lru",
+        "@com_google_googletest//:gtest_main",
+    ],
+    size = "small",
+    copts = COPTS,
+    tags = ["team:core"],
+)
diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
new file mode 100644
index 000000000000..49787fe3af60
--- /dev/null
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -0,0 +1,41 @@
+#include "src/ray/util/shared_lru.h"
+
+#include <gtest/gtest.h>
+
+#include <string>
+
+namespace ray::utils::container {
+
+namespace {
+constexpr size_t kTestCacheSz = 1;
+}  // namespace
+
+TEST(SharedLruCache, PutAndGet) {
+  SharedLruCache<std::string, std::string> cache{kTestCacheSz};
+
+  // No value initially.
+  auto val = cache.Get("1");
+  EXPECT_FALSE(val.has_value());
+
+  // Check put and get.
+  cache.Put("1", "1");
+  val = cache.Get("1");
+  EXPECT_TRUE(val.has_value());
+  EXPECT_EQ(*val, "1");
+
+  // Check key eviction.
+  cache.Put("2", "2");
+  val = cache.Get("1");
+  EXPECT_FALSE(val.has_value());
+  val = cache.Get("2");
+  EXPECT_TRUE(val.has_value());
+  EXPECT_EQ(*val, "2");
+
+  // Check deletion.
+  EXPECT_FALSE(cache.Delete("1"));
+  EXPECT_TRUE(cache.Delete("2"));
+  val = cache.Get("2");
+  EXPECT_FALSE(val.has_value());
+}
+
+}  // namespace ray::utils::container

From d0265fdbdf3836c61aea9904fdfed6b3f47d8260 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Tue, 26 Nov 2024 22:21:12 +0000
Subject: [PATCH 02/13] move key at insertion

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index 1ff1cc15385a..9c9b31bf966e 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -48,7 +48,7 @@ class SharedLruCache final {
     std::lock_guard lck(mu_);
     lru_list_.emplace_front(key);
     Entry new_entry{std::move(value), lru_list_.begin()};
-    cache_[key] = std::move(new_entry);
+    cache_[std::move(key)] = std::move(new_entry);
 
     if (max_entries_ > 0 && lru_list_.size() > max_entries_) {
       const auto &stale_key = lru_list_.back();

From bb8f7d885de871ea2fb0f7013a9acce5cfdc7fef Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Tue, 26 Nov 2024 22:27:52 +0000
Subject: [PATCH 03/13] fix lru list removal

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/BUILD        | 1 +
 src/ray/util/shared_lru.h | 1 +
 2 files changed, 2 insertions(+)

diff --git a/src/ray/util/BUILD b/src/ray/util/BUILD
index 0a65db3726e1..83604a428f3a 100644
--- a/src/ray/util/BUILD
+++ b/src/ray/util/BUILD
@@ -62,5 +62,6 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "@com_google_absl//absl/container:flat_hash_map",
+        ":util",
     ],
 )
diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index 9c9b31bf966e..f62eac29251c 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -66,6 +66,7 @@ class SharedLruCache final {
     if (it == cache_.end()) {
       return false;
     }
+    lru_list_.erase(it->second.lru_iterator);
     cache_.erase(it);
     return true;
   }

From 29380a117b1a5cb9ebba440864f63cbddce10066 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Tue, 26 Nov 2024 22:33:58 +0000
Subject: [PATCH 04/13] split thread safe and no lock version

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/BUILD                    |  1 -
 src/ray/util/shared_lru.h             | 66 ++++++++++++++++++++++++---
 src/ray/util/tests/shared_lru_test.cc |  2 +-
 3 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/src/ray/util/BUILD b/src/ray/util/BUILD
index 83604a428f3a..0a65db3726e1 100644
--- a/src/ray/util/BUILD
+++ b/src/ray/util/BUILD
@@ -62,6 +62,5 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "@com_google_absl//absl/container:flat_hash_map",
-        ":util",
     ],
 )
diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index f62eac29251c..f487d24b4201 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -40,12 +40,14 @@ class SharedLruCache final {
   // in the cache.
   explicit SharedLruCache(size_t max_entries) : max_entries_(max_entries) {}
 
+  SharedLruCache(const SharedLruCache &) = delete;
+  SharedLruCache &operator=(const SharedLruCache &) = delete;
+
   ~SharedLruCache() = default;
 
   // Insert `value` with key `key`. This will replace any previous entry with
   // the same key.
   void Put(Key key, Val value) {
-    std::lock_guard lck(mu_);
     lru_list_.emplace_front(key);
     Entry new_entry{std::move(value), lru_list_.begin()};
     cache_[std::move(key)] = std::move(new_entry);
@@ -61,7 +63,6 @@ class SharedLruCache final {
   // `key`, false if the entry was not found. In both cases, there is no entry
   // with key `key` existed after the call.
   bool Delete(const Key &key) {
-    std::lock_guard lck(mu_);
     auto it = cache_.find(key);
     if (it == cache_.end()) {
       return false;
@@ -74,7 +75,6 @@ class SharedLruCache final {
   // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
   // If found, return a copy for the value.
   std::optional<Val> Get(const Key &key) {
-    std::lock_guard lck(mu_);
     const auto cache_iter = cache_.find(key);
     if (cache_iter == cache_.end()) {
       return std::nullopt;
@@ -93,7 +93,6 @@ class SharedLruCache final {
 
   // Clear the cache.
   void Clear() {
-    std::lock_guard lck(mu_);
     cache_.clear();
     lru_list_.clear();
   }
@@ -116,9 +115,6 @@ class SharedLruCache final {
   // limit on entry count.
   const size_t max_entries_;
 
-  // Guards access to the cache and the LRU list.
-  std::mutex mu_;
-
   // All keys are stored as refernce (`std::reference_wrapper`), and the
   // ownership lies in `lru_list_`.
   EntryMap cache_;
@@ -133,4 +129,60 @@ class SharedLruCache final {
 template <typename K, typename V>
 using SharedLruConstCache = SharedLruCache<K, const V>;
 
+template <typename Key, typename Val>
+class ThreadSafeSharedLruCache final {
+ public:
+  using key_type = Key;
+  using mapped_type = Val;
+
+  // A `max_entries` of 0 means that there is no limit on the number of entries
+  // in the cache.
+  explicit ThreadSafeSharedLruCache(size_t max_entries) : cache_(max_entries) {}
+
+  ThreadSafeSharedLruCache(const ThreadSafeSharedLruCache &) = delete;
+  ThreadSafeSharedLruCache &operator=(const ThreadSafeSharedLruCache &) = delete;
+
+  ~ThreadSafeSharedLruCache() = default;
+
+  // Insert `value` with key `key`. This will replace any previous entry with
+  // the same key.
+  void Put(Key key, Val value) {
+    std::lock_guard lck(mtx_);
+    cache_.Put(std::move(key), std::move(value));
+  }
+
+  // Delete the entry with key `key`. Return true if the entry was found for
+  // `key`, false if the entry was not found. In both cases, there is no entry
+  // with key `key` existed after the call.
+  bool Delete(const Key &key) {
+    std::lock_guard lck(mtx_);
+    return cache_.Delete(key);
+  }
+
+  // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
+  // If found, return a copy for the value.
+  std::optional<Val> Get(const Key &key) {
+    std::lock_guard lck(mtx_);
+    return cache_.Get(key);
+  }
+
+  // Clear the cache.
+  void Clear() {
+    std::lock_guard lck(mtx_);
+    cache_.Clear();
+  }
+
+  // Accessors for cache parameters.
+  size_t max_entries() const { return cache_.max_entries(); }
+
+ private:
+  std::mutex mtx_;
+  SharedLruCache<Key, Val> cache_;
+};
+
+// Same interfaces as `SharedLruCache`, but all cached values are
+// `const`-specified to avoid concurrent updates.
+template <typename K, typename V>
+using ThreadSafeSharedLruConstCache = ThreadSafeSharedLruCache<K, const V>;
+
 }  // namespace ray::utils::container
diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
index 49787fe3af60..3ecddb620146 100644
--- a/src/ray/util/tests/shared_lru_test.cc
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -11,7 +11,7 @@ constexpr size_t kTestCacheSz = 1;
 }  // namespace
 
 TEST(SharedLruCache, PutAndGet) {
-  SharedLruCache<std::string, std::string> cache{kTestCacheSz};
+  ThreadSafeSharedLruCache<std::string, std::string> cache{kTestCacheSz};
 
   // No value initially.
   auto val = cache.Get("1");

From 095e70472c73957da6ec7bbd7cf12fdb387ab963 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Tue, 26 Nov 2024 22:37:38 +0000
Subject: [PATCH 05/13] doc on thread-safe

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index f487d24b4201..169a5bd85d86 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -129,6 +129,7 @@ class SharedLruCache final {
 template <typename K, typename V>
 using SharedLruConstCache = SharedLruCache<K, const V>;
 
+// Same interface and functionality as `SharedLruCache`, but thread-safe version.
 template <typename Key, typename Val>
 class ThreadSafeSharedLruCache final {
  public:

From 73e2fae846881066ec0111863047ad264170d573 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Wed, 27 Nov 2024 00:09:36 +0000
Subject: [PATCH 06/13] copyright

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h             | 14 ++++++++++++++
 src/ray/util/tests/shared_lru_test.cc | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index 169a5bd85d86..070a7b484926 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -1,3 +1,17 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
 // SharedLruCache is a LRU cache, with all entries shared, which means a single entry
 // could be accessed by multiple getters. When `Get`, a copy of the value is returned, so
 // for heavy-loaded types it's suggested to wrap with `std::shared_ptr<>`.
diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
index 3ecddb620146..8e1db5bd733a 100644
--- a/src/ray/util/tests/shared_lru_test.cc
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -1,3 +1,17 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 #include "src/ray/util/shared_lru.h"
 
 #include <gtest/gtest.h>

From 2e9b7a9f24fa05a0924f67f7fc81a3bd5142a645 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Wed, 27 Nov 2024 01:35:04 +0000
Subject: [PATCH 07/13] some updates for lru cache

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h             | 63 ++++++++++++++++++---------
 src/ray/util/tests/shared_lru_test.cc | 20 ++++-----
 2 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index 070a7b484926..eee8745ccb30 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -25,10 +25,9 @@
 // auto val = ache.Get("key");
 // // Check and consume `val`.
 //
-// TODO(hjiang):
-// 1. Add template arguments for key hash and key equal, to pass into absl::flat_hash_map.
-// 2. Provide a key hash wrapper to save a copy.
-// 3. flat hash map supports heterogeneous lookup, expose `KeyLike` templated interface.
+// TODO(hjiang): Write a wrapper around KeyHash and KeyEq, which takes
+// std::reference_wrapper<Key>, so we could store keys only in std::list, and reference in
+// absl::flat_hash_map.
 
 #pragma once
 
@@ -41,14 +40,22 @@
 #include <utility>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/internal/hash_function_defaults.h"
 
 namespace ray::utils::container {
 
-template <typename Key, typename Val>
+// TODO(hjiang): `absl::container_internal::hash_default_` has been promoted to stable
+// public interface in the latest abseil, update after we bump up abseil version.
+template <typename Key,
+          typename Val,
+          typename KeyHash = absl::container_internal::hash_default_hash<Key>,
+          typename KeyEq = absl::container_internal::hash_default_eq<Key>>
 class SharedLruCache final {
  public:
   using key_type = Key;
   using mapped_type = Val;
+  using hasher = KeyHash;
+  using key_equal = KeyEq;
 
   // A `max_entries` of 0 means that there is no limit on the number of entries
   // in the cache.
@@ -76,7 +83,8 @@ class SharedLruCache final {
   // Delete the entry with key `key`. Return true if the entry was found for
   // `key`, false if the entry was not found. In both cases, there is no entry
   // with key `key` existed after the call.
-  bool Delete(const Key &key) {
+  template <typename KeyLike>
+  bool Delete(KeyLike &&key) {
     auto it = cache_.find(key);
     if (it == cache_.end()) {
       return false;
@@ -88,7 +96,8 @@ class SharedLruCache final {
 
   // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
   // If found, return a copy for the value.
-  std::optional<Val> Get(const Key &key) {
+  template <typename KeyLike>
+  std::optional<Val> Get(KeyLike &&key) {
     const auto cache_iter = cache_.find(key);
     if (cache_iter == cache_.end()) {
       return std::nullopt;
@@ -98,9 +107,9 @@ class SharedLruCache final {
     cache_.erase(cache_iter);
 
     // Re-insert into the cache, no need to check capacity.
-    lru_list_.emplace_front(key);
+    lru_list_.emplace_front(static_cast<Key>(key));
     Entry new_entry{value, lru_list_.begin()};
-    cache_[key] = std::move(new_entry);
+    cache_[static_cast<Key>(std::forward<KeyLike>(key))] = std::move(new_entry);
 
     return value;
   }
@@ -123,7 +132,7 @@ class SharedLruCache final {
     typename std::list<Key>::iterator lru_iterator;
   };
 
-  using EntryMap = absl::flat_hash_map<Key, Entry>;
+  using EntryMap = absl::flat_hash_map<Key, Entry, KeyHash, KeyEq>;
 
   // The maximum number of entries in the cache. A value of 0 means there is no
   // limit on entry count.
@@ -140,15 +149,23 @@ class SharedLruCache final {
 
 // Same interfaces as `SharedLruCache`, but all cached values are
 // `const`-specified to avoid concurrent updates.
-template <typename K, typename V>
-using SharedLruConstCache = SharedLruCache<K, const V>;
+template <typename K,
+          typename V,
+          typename KeyHash = absl::container_internal::hash_default_hash<K>,
+          typename KeyEq = absl::container_internal::hash_default_eq<K>>
+using SharedLruConstCache = SharedLruCache<K, const V, KeyHash, KeyEq>;
 
 // Same interface and functionality as `SharedLruCache`, but thread-safe version.
-template <typename Key, typename Val>
+template <typename Key,
+          typename Val,
+          typename KeyHash = absl::container_internal::hash_default_hash<Key>,
+          typename KeyEq = absl::container_internal::hash_default_eq<Key>>
 class ThreadSafeSharedLruCache final {
  public:
   using key_type = Key;
   using mapped_type = Val;
+  using hasher = KeyHash;
+  using key_equal = KeyEq;
 
   // A `max_entries` of 0 means that there is no limit on the number of entries
   // in the cache.
@@ -169,16 +186,18 @@ class ThreadSafeSharedLruCache final {
   // Delete the entry with key `key`. Return true if the entry was found for
   // `key`, false if the entry was not found. In both cases, there is no entry
   // with key `key` existed after the call.
-  bool Delete(const Key &key) {
+  template <typename KeyLike>
+  bool Delete(KeyLike &&key) {
     std::lock_guard lck(mtx_);
-    return cache_.Delete(key);
+    return cache_.Delete(std::forward<KeyLike>(key));
   }
 
   // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
   // If found, return a copy for the value.
-  std::optional<Val> Get(const Key &key) {
+  template <typename KeyLike>
+  std::optional<Val> Get(KeyLike &&key) {
     std::lock_guard lck(mtx_);
-    return cache_.Get(key);
+    return cache_.Get(std::forward<KeyLike>(key));
   }
 
   // Clear the cache.
@@ -192,12 +211,16 @@ class ThreadSafeSharedLruCache final {
 
  private:
   std::mutex mtx_;
-  SharedLruCache<Key, Val> cache_;
+  SharedLruCache<Key, Val, KeyHash, KeyEq> cache_;
 };
 
 // Same interfaces as `SharedLruCache`, but all cached values are
 // `const`-specified to avoid concurrent updates.
-template <typename K, typename V>
-using ThreadSafeSharedLruConstCache = ThreadSafeSharedLruCache<K, const V>;
+template <typename K,
+          typename V,
+          typename KeyHash = absl::container_internal::hash_default_hash<K>,
+          typename KeyEq = absl::container_internal::hash_default_eq<K>>
+using ThreadSafeSharedLruConstCache =
+    ThreadSafeSharedLruCache<K, const V, KeyHash, KeyEq>;
 
 }  // namespace ray::utils::container
diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
index 8e1db5bd733a..78066ed51824 100644
--- a/src/ray/util/tests/shared_lru_test.cc
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -24,31 +24,31 @@ namespace {
 constexpr size_t kTestCacheSz = 1;
 }  // namespace
 
-TEST(SharedLruCache, PutAndGet) {
+TEST(SharedLruCacheTest, PutAndGet) {
   ThreadSafeSharedLruCache<std::string, std::string> cache{kTestCacheSz};
 
   // No value initially.
-  auto val = cache.Get("1");
+  auto val = cache.Get(std::string_view{"1"});
   EXPECT_FALSE(val.has_value());
 
   // Check put and get.
   cache.Put("1", "1");
-  val = cache.Get("1");
+  val = cache.Get(std::string_view{"1"});
   EXPECT_TRUE(val.has_value());
-  EXPECT_EQ(*val, "1");
+  EXPECT_EQ(*val, std::string_view{"1"});
 
   // Check key eviction.
   cache.Put("2", "2");
-  val = cache.Get("1");
+  val = cache.Get(std::string_view{"1"});
   EXPECT_FALSE(val.has_value());
-  val = cache.Get("2");
+  val = cache.Get(std::string_view{"2"});
   EXPECT_TRUE(val.has_value());
-  EXPECT_EQ(*val, "2");
+  EXPECT_EQ(*val, std::string_view{"2"});
 
   // Check deletion.
-  EXPECT_FALSE(cache.Delete("1"));
-  EXPECT_TRUE(cache.Delete("2"));
-  val = cache.Get("2");
+  EXPECT_FALSE(cache.Delete(std::string_view{"1"}));
+  EXPECT_TRUE(cache.Delete(std::string_view{"2"}));
+  val = cache.Get(std::string_view{"2"});
   EXPECT_FALSE(val.has_value());
 }
 

From 05b67f7e177f1c079b00886740f62b7729db4cde Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Wed, 27 Nov 2024 01:48:34 +0000
Subject: [PATCH 08/13] fix typo

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index eee8745ccb30..7060e752fbc4 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -22,7 +22,7 @@
 // cache.Put("key", "val");
 //
 // // Get a key-value pair from cache.
-// auto val = ache.Get("key");
+// auto val = cache.Get("key");
 // // Check and consume `val`.
 //
 // TODO(hjiang): Write a wrapper around KeyHash and KeyEq, which takes
@@ -138,8 +138,7 @@ class SharedLruCache final {
   // limit on entry count.
   const size_t max_entries_;
 
-  // All keys are stored as refernce (`std::reference_wrapper`), and the
-  // ownership lies in `lru_list_`.
+  // Key-value pairs.
   EntryMap cache_;
 
   // The LRU list of entries. The front of the list identifies the most

From e1cc2c2181225fc2dfd937fb05b76c3278172c0c Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Thu, 28 Nov 2024 03:13:56 +0000
Subject: [PATCH 09/13] fix conflict

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h             | 18 ++++++++++--------
 src/ray/util/tests/shared_lru_test.cc |  3 +--
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index 93cd6b8fcf8c..5bc84cb7164e 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -26,9 +26,9 @@
 // auto val = cache.Get("key");
 // // Check and consume `val`.
 //
-// TODO(hjiang): 
-// 1. Write a wrapper around KeyHash and KeyEq, which takes std::reference_wrapper<Key>, so we could store keys only in std::list, and reference in
-// absl::flat_hash_map.
+// TODO(hjiang):
+// 1. Write a wrapper around KeyHash and KeyEq, which takes std::reference_wrapper<Key>,
+// so we could store keys only in std::list, and reference in absl::flat_hash_map.
 // 2. Add a `GetOrCreate` interface, which takes factory function to creation value.
 // 3. For thread-safe cache, add a sharded container wrapper to reduce lock contention.
 
@@ -38,13 +38,12 @@
 #include <list>
 #include <memory>
 #include <mutex>
-#include <optional>
 #include <string>
 #include <utility>
 
-#include "src/ray/util/logging.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/internal/hash_function_defaults.h"
+#include "src/ray/util/logging.h"
 
 namespace ray::utils::container {
 
@@ -111,7 +110,7 @@ class SharedLruCache final {
   // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
   // If found, return a copy for the value.
   template <typename KeyLike>
-  std::optional<Val> Get(KeyLike &&key) {
+  std::shared_ptr<Val> Get(KeyLike &&key) {
     const auto cache_iter = cache_.find(key);
     if (cache_iter == cache_.end()) {
       return nullptr;
@@ -161,7 +160,10 @@ template <typename K,
 using SharedLruConstCache = SharedLruCache<K, const V, KeyHash, KeyEq>;
 
 // Same interface and functionality as `SharedLruCache`, but thread-safe version.
-template <typename Key, typename Val>
+template <typename Key,
+          typename Val,
+          typename KeyHash = absl::container_internal::hash_default_hash<Key>,
+          typename KeyEq = absl::container_internal::hash_default_eq<Key>>
 class ThreadSafeSharedLruCache final {
  public:
   using key_type = Key;
@@ -197,7 +199,7 @@ class ThreadSafeSharedLruCache final {
   // Look up the entry with key `key`. Return std::nullopt if key doesn't exist.
   // If found, return a copy for the value.
   template <typename KeyLike>
-  std::optional<Val> Get(KeyLike &&key) {
+  std::shared_ptr<Val> Get(KeyLike &&key) {
     std::lock_guard lck(mu_);
     return cache_.Get(std::forward<KeyLike>(key));
   }
diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
index 5e0f50da7819..ef5cda30319f 100644
--- a/src/ray/util/tests/shared_lru_test.cc
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -42,13 +42,12 @@ TEST(SharedLruCache, PutAndGet) {
   cache.Put("2", std::make_shared<std::string>("2"));
   val = cache.Get(std::string_view{"1"});
   EXPECT_EQ(val, nullptr);
-  val = cache.Get(std::string_view{"1"});
+  val = cache.Get(std::string_view{"2"});
   EXPECT_NE(val, nullptr);
   EXPECT_EQ(*val, "2");
 
   // Check deletion.
   EXPECT_FALSE(cache.Delete(std::string_view{"1"}));
-  EXPECT_TRUE(cache.Delete(std::string_view{"1"}));
   val = cache.Get(std::string_view{"1"});
   EXPECT_EQ(val, nullptr);
 }

From 1ef094aed23f4f5fb1d88615fa19c8b37ce592f7 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Thu, 28 Nov 2024 11:32:30 +0000
Subject: [PATCH 10/13] key hash and equal test

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/tests/shared_lru_test.cc | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
index ef5cda30319f..05e892447041 100644
--- a/src/ray/util/tests/shared_lru_test.cc
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -23,6 +23,19 @@ namespace ray::utils::container {
 
 namespace {
 constexpr size_t kTestCacheSz = 1;
+
+class TestClass {
+ public:
+  TestClass(std::string data) : data_(std::move(data)) {}
+  bool operator==(const TestClass &rhs) const { return data_ == rhs.data_; }
+  template <typename H>
+  friend H AbslHashValue(H h, const TestClass &obj) {
+    return H::combine(std::move(h), obj.data_);
+  }
+
+ private:
+  std::string data_;
+};
 }  // namespace
 
 TEST(SharedLruCache, PutAndGet) {
@@ -72,4 +85,13 @@ TEST(SharedLruConstCache, TypeAliasAssertion) {
       std::is_same_v<SharedLruConstCache<int, int>, SharedLruCache<int, const int>>);
 }
 
+TEST(SharedLruConstCache, CustomizedKey) {
+  TestClass obj1{"hello"};
+  TestClass obj2{"hello"};
+  SharedLruCache<TestClass, std::string> cache{2};
+  cache.Put(obj1, std::make_shared<std::string>("val"));
+  auto val = cache.Get(obj2);
+  EXPECT_EQ(*val, "val");
+}
+
 }  // namespace ray::utils::container

From ad8ac27f7b507772b5420bd6e89cd1ef296430a7 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Mon, 2 Dec 2024 21:23:59 +0000
Subject: [PATCH 11/13] update test class name

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/tests/shared_lru_test.cc | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/ray/util/tests/shared_lru_test.cc b/src/ray/util/tests/shared_lru_test.cc
index 05e892447041..673395a82f3b 100644
--- a/src/ray/util/tests/shared_lru_test.cc
+++ b/src/ray/util/tests/shared_lru_test.cc
@@ -24,12 +24,12 @@ namespace ray::utils::container {
 namespace {
 constexpr size_t kTestCacheSz = 1;
 
-class TestClass {
+class TestClassWithHashAndEq {
  public:
-  TestClass(std::string data) : data_(std::move(data)) {}
-  bool operator==(const TestClass &rhs) const { return data_ == rhs.data_; }
+  TestClassWithHashAndEq(std::string data) : data_(std::move(data)) {}
+  bool operator==(const TestClassWithHashAndEq &rhs) const { return data_ == rhs.data_; }
   template <typename H>
-  friend H AbslHashValue(H h, const TestClass &obj) {
+  friend H AbslHashValue(H h, const TestClassWithHashAndEq &obj) {
     return H::combine(std::move(h), obj.data_);
   }
 
@@ -86,9 +86,9 @@ TEST(SharedLruConstCache, TypeAliasAssertion) {
 }
 
 TEST(SharedLruConstCache, CustomizedKey) {
-  TestClass obj1{"hello"};
-  TestClass obj2{"hello"};
-  SharedLruCache<TestClass, std::string> cache{2};
+  TestClassWithHashAndEq obj1{"hello"};
+  TestClassWithHashAndEq obj2{"hello"};
+  SharedLruCache<TestClassWithHashAndEq, std::string> cache{2};
   cache.Put(obj1, std::make_shared<std::string>("val"));
   auto val = cache.Get(obj2);
   EXPECT_EQ(*val, "val");

From 357c3cda4b52bf6ca13806d82ac8483847da1018 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Mon, 2 Dec 2024 22:08:54 +0000
Subject: [PATCH 12/13] remove customized hash and eq

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h | 35 ++++++++---------------------------
 1 file changed, 8 insertions(+), 27 deletions(-)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index 5bc84cb7164e..e57dd0b9cb16 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -47,18 +47,11 @@
 
 namespace ray::utils::container {
 
-// TODO(hjiang): `absl::container_internal::hash_default_` has been promoted to stable
-// public interface in the latest abseil, update after we bump up abseil version.
-template <typename Key,
-          typename Val,
-          typename KeyHash = absl::container_internal::hash_default_hash<Key>,
-          typename KeyEq = absl::container_internal::hash_default_eq<Key>>
+template <typename Key, typename Val>
 class SharedLruCache final {
  public:
   using key_type = Key;
   using mapped_type = Val;
-  using hasher = KeyHash;
-  using key_equal = KeyEq;
 
   // A `max_entries` of 0 means that there is no limit on the number of entries
   // in the cache.
@@ -137,7 +130,7 @@ class SharedLruCache final {
     typename std::list<Key>::iterator lru_iterator;
   };
 
-  using EntryMap = absl::flat_hash_map<Key, Entry, KeyHash, KeyEq>;
+  using EntryMap = absl::flat_hash_map<Key, Entry>;
 
   // The maximum number of entries in the cache. A value of 0 means there is no
   // limit on entry count.
@@ -153,23 +146,15 @@ class SharedLruCache final {
 
 // Same interfaces as `SharedLruCache`, but all cached values are
 // `const`-specified to avoid concurrent updates.
-template <typename K,
-          typename V,
-          typename KeyHash = absl::container_internal::hash_default_hash<K>,
-          typename KeyEq = absl::container_internal::hash_default_eq<K>>
-using SharedLruConstCache = SharedLruCache<K, const V, KeyHash, KeyEq>;
+template <typename K, typename V>
+using SharedLruConstCache = SharedLruCache<K, const V>;
 
 // Same interface and functionality as `SharedLruCache`, but thread-safe version.
-template <typename Key,
-          typename Val,
-          typename KeyHash = absl::container_internal::hash_default_hash<Key>,
-          typename KeyEq = absl::container_internal::hash_default_eq<Key>>
+template <typename Key, typename Val>
 class ThreadSafeSharedLruCache final {
  public:
   using key_type = Key;
   using mapped_type = Val;
-  using hasher = KeyHash;
-  using key_equal = KeyEq;
 
   // A `max_entries` of 0 means that there is no limit on the number of entries
   // in the cache.
@@ -215,16 +200,12 @@ class ThreadSafeSharedLruCache final {
 
  private:
   std::mutex mu_;
-  SharedLruCache<Key, Val, KeyHash, KeyEq> cache_;
+  SharedLruCache<Key, Val> cache_;
 };
 
 // Same interfaces as `SharedLruCache`, but all cached values are
 // `const`-specified to avoid concurrent updates.
-template <typename K,
-          typename V,
-          typename KeyHash = absl::container_internal::hash_default_hash<K>,
-          typename KeyEq = absl::container_internal::hash_default_eq<K>>
-using ThreadSafeSharedLruConstCache =
-    ThreadSafeSharedLruCache<K, const V, KeyHash, KeyEq>;
+template <typename K, typename V>
+using ThreadSafeSharedLruConstCache = ThreadSafeSharedLruCache<K, const V>;
 
 }  // namespace ray::utils::container

From 1af2ff0e05a438d5deaaafcb6ec97c62668642f1 Mon Sep 17 00:00:00 2001
From: hjiang <dentinyhao@gmail.com>
Date: Mon, 2 Dec 2024 22:58:34 +0000
Subject: [PATCH 13/13] remove useless header

Signed-off-by: hjiang <dentinyhao@gmail.com>
---
 src/ray/util/shared_lru.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/ray/util/shared_lru.h b/src/ray/util/shared_lru.h
index e57dd0b9cb16..39c43231ac4d 100644
--- a/src/ray/util/shared_lru.h
+++ b/src/ray/util/shared_lru.h
@@ -42,7 +42,6 @@
 #include <utility>
 
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/internal/hash_function_defaults.h"
 #include "src/ray/util/logging.h"
 
 namespace ray::utils::container {