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()) max_age_(max_age)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. Insert(const string & key,const T & value)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 // Delete the entry with key `key`. Return true if the entry was found for 55 // `key`, false if the entry was not found. In both cases, there is no entry 56 // with key `key` existed after the call. Delete(const string & key)57 bool Delete(const string& key) { 58 mutex_lock lock(mu_); 59 return DeleteLocked(key); 60 } 61 62 /// Look up the entry with key `key` and copy it to `value` if found. Returns 63 /// true if an entry was found for `key`, and its timestamp is not more than 64 /// max_age_ seconds in the past. Lookup(const string & key,T * value)65 bool Lookup(const string& key, T* value) { 66 if (max_age_ == 0) { 67 return false; 68 } 69 mutex_lock lock(mu_); 70 return LookupLocked(key, value); 71 } 72 73 typedef std::function<Status(const string&, T*)> ComputeFunc; 74 75 /// Look up the entry with key `key` and copy it to `value` if found. If not 76 /// found, call `compute_func`. If `compute_func` returns successfully, store 77 /// a copy of the output parameter in the cache, and another copy in `value`. LookupOrCompute(const string & key,T * value,const ComputeFunc & compute_func)78 Status LookupOrCompute(const string& key, T* value, 79 const ComputeFunc& compute_func) { 80 if (max_age_ == 0) { 81 return compute_func(key, value); 82 } 83 84 // Note: we hold onto mu_ for the rest of this function. In practice, this 85 // is okay, as stat requests are typically fast, and concurrent requests are 86 // often for the same file. Future work can split this up into one lock per 87 // key if this proves to be a significant performance bottleneck. 88 mutex_lock lock(mu_); 89 if (LookupLocked(key, value)) { 90 return OkStatus(); 91 } 92 Status s = compute_func(key, value); 93 if (s.ok()) { 94 InsertLocked(key, *value); 95 } 96 return s; 97 } 98 99 /// Clear the cache. Clear()100 void Clear() { 101 mutex_lock lock(mu_); 102 cache_.clear(); 103 lru_list_.clear(); 104 } 105 106 /// Accessors for cache parameters. max_age()107 uint64 max_age() const { return max_age_; } max_entries()108 size_t max_entries() const { return max_entries_; } 109 110 private: 111 struct Entry { 112 /// The timestamp (seconds) at which the entry was added to the cache. 113 uint64 timestamp; 114 115 /// The entry's value. 116 T value; 117 118 /// A list iterator pointing to the entry's position in the LRU list. 119 std::list<string>::iterator lru_iterator; 120 }; 121 LookupLocked(const string & key,T * value)122 bool LookupLocked(const string& key, T* value) 123 TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { 124 auto it = cache_.find(key); 125 if (it == cache_.end()) { 126 return false; 127 } 128 lru_list_.erase(it->second.lru_iterator); 129 if (env_->NowSeconds() - it->second.timestamp > max_age_) { 130 cache_.erase(it); 131 return false; 132 } 133 *value = it->second.value; 134 lru_list_.push_front(it->first); 135 it->second.lru_iterator = lru_list_.begin(); 136 return true; 137 } 138 InsertLocked(const string & key,const T & value)139 void InsertLocked(const string& key, const T& value) 140 TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { 141 lru_list_.push_front(key); 142 Entry entry{env_->NowSeconds(), value, lru_list_.begin()}; 143 auto insert = cache_.insert(std::make_pair(key, entry)); 144 if (!insert.second) { 145 lru_list_.erase(insert.first->second.lru_iterator); 146 insert.first->second = entry; 147 } else if (max_entries_ > 0 && cache_.size() > max_entries_) { 148 cache_.erase(lru_list_.back()); 149 lru_list_.pop_back(); 150 } 151 } 152 DeleteLocked(const string & key)153 bool DeleteLocked(const string& key) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { 154 auto it = cache_.find(key); 155 if (it == cache_.end()) { 156 return false; 157 } 158 lru_list_.erase(it->second.lru_iterator); 159 cache_.erase(it); 160 return true; 161 } 162 163 /// The maximum age of entries in the cache, in seconds. A value of 0 means 164 /// that no entry is ever placed in the cache. 165 const uint64 max_age_; 166 167 /// The maximum number of entries in the cache. A value of 0 means there is no 168 /// limit on entry count. 169 const size_t max_entries_; 170 171 /// The Env from which we read timestamps. 172 Env* const env_; // not owned 173 174 /// Guards access to the cache and the LRU list. 175 mutex mu_; 176 177 /// The cache (a map from string key to Entry). 178 std::map<string, Entry> cache_ TF_GUARDED_BY(mu_); 179 180 /// The LRU list of entries. The front of the list identifies the most 181 /// recently accessed entry. 182 std::list<string> lru_list_ TF_GUARDED_BY(mu_); 183 }; 184 185 } // namespace tensorflow 186 187 #endif // TENSORFLOW_CORE_PLATFORM_CLOUD_EXPIRING_LRU_CACHE_H_ 188