1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_PLATFORM_CLOUD_EXPIRING_LRU_CACHE_H_ 17 #define TENSORFLOW_CORE_PLATFORM_CLOUD_EXPIRING_LRU_CACHE_H_ 18 19 #include <list> 20 #include <map> 21 #include <memory> 22 #include <string> 23 #include "tensorflow/core/platform/env.h" 24 #include "tensorflow/core/platform/mutex.h" 25 #include "tensorflow/core/platform/thread_annotations.h" 26 #include "tensorflow/core/platform/types.h" 27 28 namespace tensorflow { 29 30 /// \brief An LRU cache of string keys and arbitrary values, with configurable 31 /// max item age (in seconds) and max entries. 32 /// 33 /// This class is thread safe. 34 template <typename T> 35 class ExpiringLRUCache { 36 public: 37 /// A `max_age` of 0 means that nothing is cached. A `max_entries` of 0 means 38 /// that there is no limit on the number of entries in the cache (however, if 39 /// `max_age` is also 0, the cache will not be populated). 40 ExpiringLRUCache(uint64 max_age, size_t max_entries, 41 Env* env = Env::Default()) 42 : max_age_(max_age), max_entries_(max_entries), env_(env) {} 43 44 /// Insert `value` with key `key`. This will replace any previous entry with 45 /// the same key. 46 void Insert(const string& key, const T& value) { 47 if (max_age_ == 0) { 48 return; 49 } 50 mutex_lock lock(mu_); 51 InsertLocked(key, value); 52 } 53 54 /// Look up the entry with key `key` and copy it to `value` if found. Returns 55 /// true if an entry was found for `key`, and its timestamp is not more than 56 /// max_age_ seconds in the past. 57 bool Lookup(const string& key, T* value) { 58 if (max_age_ == 0) { 59 return false; 60 } 61 mutex_lock lock(mu_); 62 return LookupLocked(key, value); 63 } 64 65 typedef std::function<Status(const string&, T*)> ComputeFunc; 66 67 /// Look up the entry with key `key` and copy it to `value` if found. If not 68 /// found, call `compute_func`. If `compute_func` returns successfully, store 69 /// a copy of the output parameter in the cache, and another copy in `value`. 70 Status LookupOrCompute(const string& key, T* value, 71 const ComputeFunc& compute_func) { 72 if (max_age_ == 0) { 73 return compute_func(key, value); 74 } 75 76 // Note: we hold onto mu_ for the rest of this function. In practice, this 77 // is okay, as stat requests are typically fast, and concurrent requests are 78 // often for the same file. Future work can split this up into one lock per 79 // key if this proves to be a significant performance bottleneck. 80 mutex_lock lock(mu_); 81 if (LookupLocked(key, value)) { 82 return Status::OK(); 83 } 84 Status s = compute_func(key, value); 85 if (s.ok()) { 86 InsertLocked(key, *value); 87 } 88 return s; 89 } 90 91 /// Clear the cache. 92 void Clear() { 93 mutex_lock lock(mu_); 94 cache_.clear(); 95 lru_list_.clear(); 96 } 97 98 /// Accessors for cache parameters. 99 uint64 max_age() const { return max_age_; } 100 size_t max_entries() const { return max_entries_; } 101 102 private: 103 struct Entry { 104 /// The timestamp (seconds) at which the entry was added to the cache. 105 uint64 timestamp; 106 107 /// The entry's value. 108 T value; 109 110 /// A list iterator pointing to the entry's position in the LRU list. 111 std::list<string>::iterator lru_iterator; 112 }; 113 114 bool LookupLocked(const string& key, T* value) EXCLUSIVE_LOCKS_REQUIRED(mu_) { 115 auto it = cache_.find(key); 116 if (it == cache_.end()) { 117 return false; 118 } 119 lru_list_.erase(it->second.lru_iterator); 120 if (env_->NowSeconds() - it->second.timestamp > max_age_) { 121 cache_.erase(it); 122 return false; 123 } 124 *value = it->second.value; 125 lru_list_.push_front(it->first); 126 it->second.lru_iterator = lru_list_.begin(); 127 return true; 128 } 129 130 void InsertLocked(const string& key, const T& value) 131 EXCLUSIVE_LOCKS_REQUIRED(mu_) { 132 lru_list_.push_front(key); 133 Entry entry{env_->NowSeconds(), value, lru_list_.begin()}; 134 auto insert = cache_.insert(std::make_pair(key, entry)); 135 if (!insert.second) { 136 lru_list_.erase(insert.first->second.lru_iterator); 137 insert.first->second = entry; 138 } else if (max_entries_ > 0 && cache_.size() > max_entries_) { 139 cache_.erase(lru_list_.back()); 140 lru_list_.pop_back(); 141 } 142 } 143 144 /// The maximum age of entries in the cache, in seconds. A value of 0 means 145 /// that no entry is ever placed in the cache. 146 const uint64 max_age_; 147 148 /// The maximum number of entries in the cache. A value of 0 means there is no 149 /// limit on entry count. 150 const size_t max_entries_; 151 152 /// The Env from which we read timestamps. 153 Env* const env_; // not owned 154 155 /// Guards access to the cache and the LRU list. 156 mutex mu_; 157 158 /// The cache (a map from string key to Entry). 159 std::map<string, Entry> cache_ GUARDED_BY(mu_); 160 161 /// The LRU list of entries. The front of the list identifies the most 162 /// recently accessed entry. 163 std::list<string> lru_list_ GUARDED_BY(mu_); 164 }; 165 166 } // namespace tensorflow 167 168 #endif // TENSORFLOW_CORE_PLATFORM_CLOUD_EXPIRING_LRU_CACHE_H_ 169
