Home | History | Annotate | Download | only in util 
      1 /* Copyright 2015 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 #include "tensorflow/core/util/tensor_slice_reader_cache.h"
     17 
     18 #include <utility>
     19 
     20 #include "tensorflow/core/lib/gtl/stl_util.h"
     21 #include "tensorflow/core/platform/logging.h"
     22 
     23 namespace tensorflow {
     24 
     25 namespace checkpoint {
     26 
     27 TensorSliceReaderCacheWrapper::TensorSliceReaderCacheWrapper() {}
     28 TensorSliceReaderCacheWrapper::~TensorSliceReaderCacheWrapper() {
     29   delete cache_;
     30   cache_ = nullptr;
     31 }
     32 
     33 const TensorSliceReader* TensorSliceReaderCacheWrapper::GetReader(
     34     const string& filepattern,
     35     TensorSliceReader::OpenTableFunction open_function,
     36     int preferred_shard) const {
     37   mutex_lock l(mu_);
     38   if (!cache_) {
     39     cache_ = new TensorSliceReaderCache;
     40   }
     41   return cache_->GetReader(filepattern, std::move(open_function),
     42                            preferred_shard);
     43 }
     44 
     45 TensorSliceReaderCache::TensorSliceReaderCache() {}
     46 
     47 TensorSliceReaderCache::~TensorSliceReaderCache() {
     48   for (auto pair : readers_) {
     49     delete pair.second.second;
     50   }
     51 }
     52 
     53 const TensorSliceReader* TensorSliceReaderCache::GetReader(
     54     const string& filepattern,
     55     TensorSliceReader::OpenTableFunction open_function, int preferred_shard) {
     56   mutex_lock l(mu_);
     57 
     58 #if defined(__GXX_RTTI) || defined(_CPPRTTI)
     59   // Get the function pointer from the open_function value.
     60   TensorSliceReaderCache::OpenFuncType* func_ptr =
     61       open_function.target<TensorSliceReaderCache::OpenFuncType>();
     62 #else   // __GXX_RTTI
     63   // When RTTI is disabled, we will hard-code func_ptr to be zero,
     64   // since we cannot figure out the target type for open_function.
     65   // TODO(jiayq): find a more elegant way to possibly enable cache again.
     66   TensorSliceReaderCache::OpenFuncType* func_ptr = nullptr;
     67 #endif  // _GXX_RTTI
     68 
     69   if (!func_ptr) {
     70     // We could not get the pointer, no caching is possible.
     71     LOG(WARNING) << "Caching disabled because the open function is a lambda or "
     72                     "RTTI is not enabled in this build.";
     73     return nullptr;
     74   }
     75 
     76   // Wait if another thread is already trying to open the same files.
     77   while (still_opening_.find(filepattern) != still_opening_.end()) {
     78     cv_.wait(l);
     79   }
     80 
     81   TensorSliceReader* reader = nullptr;
     82   if (readers_.find(filepattern) == readers_.end()) {
     83     VLOG(1) << "Creating new TensorSliceReader for " << filepattern;
     84     still_opening_.insert(filepattern);
     85     // Release the lock temporary as constructing TensorSliceReader is
     86     // expensive.
     87     mu_.unlock();
     88     TensorSliceReader* tmp_reader(
     89         new TensorSliceReader(filepattern, open_function, preferred_shard));
     90     // Acquire the lock again.
     91     mu_.lock();
     92     if (tmp_reader->status().ok()) {
     93       reader = tmp_reader;
     94       readers_[filepattern] = std::make_pair(*func_ptr, reader);
     95     } else {
     96       delete tmp_reader;
     97     }
     98     CHECK_EQ(size_t{1}, still_opening_.erase(filepattern));
     99     VLOG(1) << "Cached TensorSliceReader for " << filepattern << ": " << reader;
    100   } else {
    101     auto cached_val = readers_[filepattern];
    102     if (cached_val.first == *func_ptr) {
    103       reader = cached_val.second;
    104       VLOG(1) << "Using cached TensorSliceReader for " << filepattern << ": "
    105               << reader;
    106     } else {
    107       LOG(WARNING) << "Caching disabled because the checkpoint file "
    108                    << "is being opened with two different open functions: "
    109                    << filepattern;
    110     }
    111   }
    112 
    113   cv_.notify_all();
    114   return reader;
    115 }
    116 
    117 }  // namespace checkpoint
    118 
    119 }  // namespace tensorflow
    120