chore: import upstream snapshot with attribution

tensorflow/c/experimental/filesystem/plugins/gcs/BUILD

@@ -0,0 +1,139 @@
+# Experimental gcs filesystem plugin.
+load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+# Filesystem implementation for GCS environments
+tf_cc_shared_object(
+    name = "gcs_filesystem",
+    framework_so = [],
+    linkstatic = False,
+    per_os_targets = 1,
+    visibility = ["//visibility:public"],
+    deps = [":gcs_filesystem_impl"],
+)
+
+# The real implementation of the filesystem.
+cc_library(
+    name = "gcs_filesystem_impl",
+    srcs = ["gcs_filesystem.cc"],
+    hdrs = ["gcs_filesystem.h"],
+    copts = select({
+        "//conditions:default": [],
+        "//tensorflow:windows": get_win_copts(),
+    }),
+    deps = [
+        ":expiring_lru_cache",
+        ":gcs_helper",
+        ":ram_file_block_cache",
+        "//tensorflow/c:env",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c/experimental/filesystem:filesystem_interface",
+        "@com_github_googlecloudplatform_google_cloud_cpp//:storage_client",
+        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud:google_cloud_cpp_common",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+cc_library(
+    name = "gcs_helper",
+    srcs = ["gcs_helper.cc"],
+    hdrs = ["gcs_helper.h"],
+    linkstatic = 1,
+    deps = [
+        "//tensorflow/c:env",
+    ],
+)
+
+cc_library(
+    name = "cleanup",
+    hdrs = ["cleanup.h"],
+)
+
+cc_library(
+    name = "ram_file_block_cache",
+    srcs = ["ram_file_block_cache.cc"],
+    hdrs = ["ram_file_block_cache.h"],
+    deps = [
+        ":cleanup",
+        "//tensorflow/c:env",
+        "//tensorflow/c:tf_status",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+    ],
+)
+
+tf_cc_test(
+    name = "ram_file_block_cache_test",
+    size = "small",
+    srcs = ["ram_file_block_cache_test.cc"],
+    deps = [
+        ":ram_file_block_cache",
+        "//tensorflow/c:tf_status_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform/cloud:now_seconds_env",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+tf_cc_test(
+    name = "gcs_filesystem_test",
+    srcs = [
+        "gcs_filesystem_test.cc",
+    ],
+    tags = [
+        "manual",
+        "notap",
+    ],
+    deps = [
+        ":gcs_filesystem_impl",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/core/platform:stacktrace_handler",
+        "//tensorflow/core/platform:test",
+        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud:google_cloud_cpp_common",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "expiring_lru_cache",
+    hdrs = ["expiring_lru_cache.h"],
+    deps = [
+        "//tensorflow/c:env",
+        "//tensorflow/c:tf_status",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+tf_cc_test(
+    name = "expiring_lru_cache_test",
+    size = "small",
+    srcs = ["expiring_lru_cache_test.cc"],
+    deps = [
+        ":expiring_lru_cache",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/c:tf_status_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform/cloud:now_seconds_env",
+    ],
+)

tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h

@@ -0,0 +1,109 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+// MakeCleanup(f) returns an RAII cleanup object that calls 'f' in its
+// destructor. The easiest way to use MakeCleanup is with a lambda argument,
+// capturing the return value in an 'auto' local variable. Most users will not
+// need more sophisticated syntax than that.
+//
+// Example:
+//   void func() {
+//     FILE* fp = fopen("data.txt", "r");
+//     if (fp == nullptr) return;
+//     auto fp_cleaner = gtl::MakeCleanup([fp] { fclose(fp); });
+//     // No matter what, fclose(fp) will happen.
+//     DataObject d;
+//     while (ReadDataObject(fp, &d)) {
+//       if (d.IsBad()) {
+//         LOG(ERROR) << "Bad Data";
+//         return;
+//       }
+//       PushGoodData(d);
+//     }
+//   }
+//
+// You can use Cleanup<F> directly, instead of using MakeCleanup and auto,
+// but there's rarely a reason to do that.
+//
+// You can call 'release()' on a Cleanup object to cancel the cleanup.
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
+
+#include <type_traits>
+#include <utility>
+
+namespace tf_gcs_filesystem {
+
+// A move-only RAII object that calls a stored cleanup functor when
+// destroyed. Cleanup<F> is the return type of gtl::MakeCleanup(F).
+template <typename F>
+class Cleanup {
+ public:
+  Cleanup() : released_(true), f_() {}
+
+  template <typename G>
+  explicit Cleanup(G&& f)          // NOLINT
+      : f_(std::forward<G>(f)) {}  // NOLINT(build/c++11)
+
+  Cleanup(Cleanup&& src)  // NOLINT
+      : released_(src.is_released()), f_(src.release()) {}
+
+  // Implicitly move-constructible from any compatible Cleanup<G>.
+  // The source will be released as if src.release() were called.
+  // A moved-from Cleanup can be safely destroyed or reassigned.
+  template <typename G>
+  Cleanup(Cleanup<G>&& src)  // NOLINT
+      : released_(src.is_released()), f_(src.release()) {}
+
+  // Assignment to a Cleanup object behaves like destroying it
+  // and making a new one in its place, analogous to unique_ptr
+  // semantics.
+  Cleanup& operator=(Cleanup&& src) {  // NOLINT
+    if (!released_) f_();
+    released_ = src.released_;
+    f_ = src.release();
+    return *this;
+  }
+
+  ~Cleanup() {
+    if (!released_) f_();
+  }
+
+  // Releases the cleanup function instead of running it.
+  // Hint: use c.release()() to run early.
+  F release() {
+    released_ = true;
+    return std::move(f_);
+  }
+
+  bool is_released() const { return released_; }
+
+ private:
+  static_assert(!std::is_reference<F>::value, "F must not be a reference");
+
+  bool released_ = false;
+  F f_;
+};
+
+template <int&... ExplicitParameterBarrier, typename F,
+          typename DecayF = typename std::decay<F>::type>
+Cleanup<DecayF> MakeCleanup(F&& f) {
+  return Cleanup<DecayF>(std::forward<F>(f));
+}
+
+}  // namespace tf_gcs_filesystem
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_

tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h

@@ -0,0 +1,191 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
+
+#include <functional>
+#include <list>
+#include <map>
+#include <memory>
+#include <string>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/c/env.h"
+#include "tensorflow/c/tf_status.h"
+
+namespace tf_gcs_filesystem {
+
+/// \brief An LRU cache of string keys and arbitrary values, with configurable
+/// max item age (in seconds) and max entries.
+///
+/// This class is thread safe.
+template <typename T>
+class ExpiringLRUCache {
+ public:
+  /// A `max_age` of 0 means that nothing is cached. A `max_entries` of 0 means
+  /// that there is no limit on the number of entries in the cache (however, if
+  /// `max_age` is also 0, the cache will not be populated).
+  ExpiringLRUCache(uint64_t max_age, size_t max_entries,
+                   std::function<uint64_t()> timer_seconds = TF_NowSeconds)
+      : max_age_(max_age),
+        max_entries_(max_entries),
+        timer_seconds_(timer_seconds) {}
+
+  /// Insert `value` with key `key`. This will replace any previous entry with
+  /// the same key.
+  void Insert(const std::string& key, const T& value) {
+    if (max_age_ == 0) {
+      return;
+    }
+    absl::MutexLock lock(mu_);
+    InsertLocked(key, 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 std::string& key) {
+    absl::MutexLock lock(mu_);
+    return DeleteLocked(key);
+  }
+
+  /// Look up the entry with key `key` and copy it to `value` if found. Returns
+  /// true if an entry was found for `key`, and its timestamp is not more than
+  /// max_age_ seconds in the past.
+  bool Lookup(const std::string& key, T* value) {
+    if (max_age_ == 0) {
+      return false;
+    }
+    absl::MutexLock lock(mu_);
+    return LookupLocked(key, value);
+  }
+
+  typedef std::function<void(const std::string&, T*, TF_Status*)> ComputeFunc;
+
+  /// Look up the entry with key `key` and copy it to `value` if found. If not
+  /// found, call `compute_func`. If `compute_func` set `status` to `TF_OK`,
+  /// store a copy of the output parameter in the cache, and another copy in
+  /// `value`.
+  void LookupOrCompute(const std::string& key, T* value,
+                       const ComputeFunc& compute_func, TF_Status* status) {
+    if (max_age_ == 0) {
+      return compute_func(key, value, status);
+    }
+
+    // Note: we hold onto mu_ for the rest of this function. In practice, this
+    // is okay, as stat requests are typically fast, and concurrent requests are
+    // often for the same file. Future work can split this up into one lock per
+    // key if this proves to be a significant performance bottleneck.
+    absl::MutexLock lock(mu_);
+    if (LookupLocked(key, value)) {
+      return TF_SetStatus(status, TF_OK, "");
+    }
+    compute_func(key, value, status);
+    if (TF_GetCode(status) == TF_OK) {
+      InsertLocked(key, *value);
+    }
+  }
+
+  /// Clear the cache.
+  void Clear() {
+    absl::MutexLock lock(mu_);
+    cache_.clear();
+    lru_list_.clear();
+  }
+
+  /// Accessors for cache parameters.
+  uint64_t max_age() const { return max_age_; }
+  size_t max_entries() const { return max_entries_; }
+
+ private:
+  struct Entry {
+    /// The timestamp (seconds) at which the entry was added to the cache.
+    uint64_t timestamp;
+
+    /// The entry's value.
+    T value;
+
+    /// A list iterator pointing to the entry's position in the LRU list.
+    std::list<std::string>::iterator lru_iterator;
+  };
+
+  bool LookupLocked(const std::string& key, T* value)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    auto it = cache_.find(key);
+    if (it == cache_.end()) {
+      return false;
+    }
+    lru_list_.erase(it->second.lru_iterator);
+    if (timer_seconds_() - it->second.timestamp > max_age_) {
+      cache_.erase(it);
+      return false;
+    }
+    *value = it->second.value;
+    lru_list_.push_front(it->first);
+    it->second.lru_iterator = lru_list_.begin();
+    return true;
+  }
+
+  void InsertLocked(const std::string& key, const T& value)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    lru_list_.push_front(key);
+    Entry entry{timer_seconds_(), value, lru_list_.begin()};
+    auto insert = cache_.insert(std::make_pair(key, entry));
+    if (!insert.second) {
+      lru_list_.erase(insert.first->second.lru_iterator);
+      insert.first->second = entry;
+    } else if (max_entries_ > 0 && cache_.size() > max_entries_) {
+      cache_.erase(lru_list_.back());
+      lru_list_.pop_back();
+    }
+  }
+
+  bool DeleteLocked(const std::string& key) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    auto it = cache_.find(key);
+    if (it == cache_.end()) {
+      return false;
+    }
+    lru_list_.erase(it->second.lru_iterator);
+    cache_.erase(it);
+    return true;
+  }
+
+  /// The maximum age of entries in the cache, in seconds. A value of 0 means
+  /// that no entry is ever placed in the cache.
+  const uint64_t max_age_;
+
+  /// 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_;
+
+  /// The callback to read timestamps.
+  std::function<uint64_t()> timer_seconds_;
+
+  /// Guards access to the cache and the LRU list.
+  absl::Mutex mu_;
+
+  /// The cache (a map from string key to Entry).
+  std::map<std::string, Entry> cache_ ABSL_GUARDED_BY(mu_);
+
+  /// The LRU list of entries. The front of the list identifies the most
+  /// recently accessed entry.
+  std::list<std::string> lru_list_ ABSL_GUARDED_BY(mu_);
+};
+
+}  // namespace tf_gcs_filesystem
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_

tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache_test.cc

@@ -0,0 +1,216 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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 "tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/cloud/now_seconds_env.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(ExpiringLRUCacheTest, MaxAge) {
+  const std::string key = "a";
+  std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache(
+      1, 0, [&env]() { return env->NowSeconds(); });
+  env->SetNowSeconds(1);
+  // Verify that replacement of an existing element works, and updates the
+  // timestamp of the entry.
+  cache.Insert(key, 41);
+  env->SetNowSeconds(2);
+  cache.Insert(key, 42);
+  // 1 second after the most recent insertion, the entry is still valid.
+  env->SetNowSeconds(3);
+  int value = 0;
+  EXPECT_TRUE(cache.Lookup(key, &value));
+  EXPECT_EQ(value, 42);
+  // 2 seconds after the most recent insertion, the entry is no longer valid.
+  env->SetNowSeconds(4);
+  EXPECT_FALSE(cache.Lookup(key, &value));
+  // Re-insert the entry.
+  cache.Insert(key, 43);
+  EXPECT_TRUE(cache.Lookup(key, &value));
+  EXPECT_EQ(value, 43);
+  // The entry is valid 1 second after the insertion...
+  env->SetNowSeconds(5);
+  value = 0;
+  EXPECT_TRUE(cache.Lookup(key, &value));
+  EXPECT_EQ(value, 43);
+  // ...but is no longer valid 2 seconds after the insertion.
+  env->SetNowSeconds(6);
+  EXPECT_FALSE(cache.Lookup(key, &value));
+}
+
+TEST(ExpiringLRUCacheTest, MaxEntries) {
+  // max_age of 0 means nothing will be cached.
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache1(0, 4);
+  cache1.Insert("a", 1);
+  int value = 0;
+  EXPECT_FALSE(cache1.Lookup("a", &value));
+  // Now set max_age = 1 and verify the LRU eviction logic.
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache2(1, 4);
+  cache2.Insert("a", 1);
+  cache2.Insert("b", 2);
+  cache2.Insert("c", 3);
+  cache2.Insert("d", 4);
+  EXPECT_TRUE(cache2.Lookup("a", &value));
+  EXPECT_EQ(value, 1);
+  EXPECT_TRUE(cache2.Lookup("b", &value));
+  EXPECT_EQ(value, 2);
+  EXPECT_TRUE(cache2.Lookup("c", &value));
+  EXPECT_EQ(value, 3);
+  EXPECT_TRUE(cache2.Lookup("d", &value));
+  EXPECT_EQ(value, 4);
+  // Insertion of "e" causes "a" to be evicted, but the other entries are still
+  // there.
+  cache2.Insert("e", 5);
+  EXPECT_FALSE(cache2.Lookup("a", &value));
+  EXPECT_TRUE(cache2.Lookup("b", &value));
+  EXPECT_EQ(value, 2);
+  EXPECT_TRUE(cache2.Lookup("c", &value));
+  EXPECT_EQ(value, 3);
+  EXPECT_TRUE(cache2.Lookup("d", &value));
+  EXPECT_EQ(value, 4);
+  EXPECT_TRUE(cache2.Lookup("e", &value));
+  EXPECT_EQ(value, 5);
+}
+
+TEST(ExpiringLRUCacheTest, LookupOrCompute) {
+  // max_age of 0 means we should always compute.
+  uint64_t num_compute_calls = 0;
+  tf_gcs_filesystem::ExpiringLRUCache<int>::ComputeFunc compute_func =
+      [&num_compute_calls](const std::string& key, int* value,
+                           TF_Status* status) {
+        *value = num_compute_calls;
+        num_compute_calls++;
+        return TF_SetStatus(status, TF_OK, "");
+      };
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache1(0, 4);
+
+  int value = -1;
+  TF_Status status;
+  cache1.LookupOrCompute("a", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 0);
+  EXPECT_EQ(num_compute_calls, 1);
+  // re-read the same value, expect another lookup
+  cache1.LookupOrCompute("a", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 1);
+  EXPECT_EQ(num_compute_calls, 2);
+
+  // Define a new cache with max_age > 0 and verify correct behavior.
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache2(2, 4);
+  num_compute_calls = 0;
+  value = -1;
+
+  // Read our first value
+  cache2.LookupOrCompute("a", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 0);
+  EXPECT_EQ(num_compute_calls, 1);
+  // Re-read, exepct no additional function compute_func calls.
+  cache2.LookupOrCompute("a", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 0);
+  EXPECT_EQ(num_compute_calls, 1);
+
+  // Read a sequence of additional values, eventually evicting "a".
+  cache2.LookupOrCompute("b", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 1);
+  EXPECT_EQ(num_compute_calls, 2);
+  cache2.LookupOrCompute("c", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 2);
+  EXPECT_EQ(num_compute_calls, 3);
+  cache2.LookupOrCompute("d", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 3);
+  EXPECT_EQ(num_compute_calls, 4);
+  cache2.LookupOrCompute("e", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 4);
+  EXPECT_EQ(num_compute_calls, 5);
+  // Verify the other values remain in the cache.
+  cache2.LookupOrCompute("b", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 1);
+  EXPECT_EQ(num_compute_calls, 5);
+  cache2.LookupOrCompute("c", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 2);
+  EXPECT_EQ(num_compute_calls, 5);
+  cache2.LookupOrCompute("d", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 3);
+  EXPECT_EQ(num_compute_calls, 5);
+
+  // Re-read "a", ensure it is re-computed.
+  cache2.LookupOrCompute("a", &value, compute_func, &status);
+  TF_EXPECT_OK(status.status);
+  EXPECT_EQ(value, 5);
+  EXPECT_EQ(num_compute_calls, 6);
+}
+
+TEST(ExpiringLRUCacheTest, Clear) {
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache(1, 4);
+  cache.Insert("a", 1);
+  cache.Insert("b", 2);
+  cache.Insert("c", 3);
+  cache.Insert("d", 4);
+  int value = 0;
+  EXPECT_TRUE(cache.Lookup("a", &value));
+  EXPECT_EQ(value, 1);
+  EXPECT_TRUE(cache.Lookup("b", &value));
+  EXPECT_EQ(value, 2);
+  EXPECT_TRUE(cache.Lookup("c", &value));
+  EXPECT_EQ(value, 3);
+  EXPECT_TRUE(cache.Lookup("d", &value));
+  EXPECT_EQ(value, 4);
+  cache.Clear();
+  EXPECT_FALSE(cache.Lookup("a", &value));
+  EXPECT_FALSE(cache.Lookup("b", &value));
+  EXPECT_FALSE(cache.Lookup("c", &value));
+  EXPECT_FALSE(cache.Lookup("d", &value));
+}
+
+TEST(ExpiringLRUCacheTest, Delete) {
+  // Insert an entry.
+  tf_gcs_filesystem::ExpiringLRUCache<int> cache(1, 4);
+  cache.Insert("a", 1);
+  int value = 0;
+  EXPECT_TRUE(cache.Lookup("a", &value));
+  EXPECT_EQ(value, 1);
+
+  // Delete the entry.
+  EXPECT_TRUE(cache.Delete("a"));
+  EXPECT_FALSE(cache.Lookup("a", &value));
+
+  // Try deleting the entry again.
+  EXPECT_FALSE(cache.Delete("a"));
+  EXPECT_FALSE(cache.Lookup("a", &value));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.h

@@ -0,0 +1,117 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+ 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.
+ ==============================================================================*/
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/synchronization/mutex.h"
+#include "google/cloud/storage/client.h"
+#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
+#include "tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h"
+#include "tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
+#include "tensorflow/c/tf_status.h"
+
+void ParseGCSPath(const std::string& fname, bool object_empty_ok,
+                  std::string* bucket, std::string* object, TF_Status* status);
+
+namespace tf_random_access_file {
+void Cleanup(TF_RandomAccessFile* file);
+int64_t Read(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
+             char* buffer, TF_Status* status);
+}  // namespace tf_random_access_file
+
+namespace tf_writable_file {
+void Cleanup(TF_WritableFile* file);
+void Append(const TF_WritableFile* file, const char* buffer, size_t n,
+            TF_Status* status);
+int64_t Tell(const TF_WritableFile* file, TF_Status* status);
+void Flush(const TF_WritableFile* file, TF_Status* status);
+void Sync(const TF_WritableFile* file, TF_Status* status);
+void Close(const TF_WritableFile* file, TF_Status* status);
+}  // namespace tf_writable_file
+
+namespace tf_read_only_memory_region {
+void Cleanup(TF_ReadOnlyMemoryRegion* region);
+const void* Data(const TF_ReadOnlyMemoryRegion* region);
+uint64_t Length(const TF_ReadOnlyMemoryRegion* region);
+}  // namespace tf_read_only_memory_region
+
+namespace tf_gcs_filesystem {
+typedef struct GcsFileStat {
+  TF_FileStatistics base;
+  int64_t generation_number;
+} GcsFileStat;
+
+typedef struct GCSFile {
+  google::cloud::storage::Client gcs_client;  // owned
+  bool compose;
+  absl::Mutex block_cache_lock;
+  std::shared_ptr<RamFileBlockCache> file_block_cache
+      ABSL_GUARDED_BY(block_cache_lock);
+  uint64_t block_size;  // Reads smaller than block_size will trigger a read
+                        // of block_size.
+  std::unique_ptr<ExpiringLRUCache<GcsFileStat>> stat_cache;
+  GCSFile(google::cloud::storage::Client&& gcs_client);
+  // This constructor is used for testing purpose only.
+  GCSFile(google::cloud::storage::Client&& gcs_client, bool compose,
+          uint64_t block_size, size_t max_bytes, uint64_t max_staleness,
+          uint64_t stat_cache_max_age, size_t stat_cache_max_entries);
+} GCSFile;
+
+// This function is used to initialize a filesystem without the need of setting
+// manually environement variables.
+void InitTest(TF_Filesystem* filesystem, bool compose, uint64_t block_size,
+              size_t max_bytes, uint64_t max_staleness,
+              uint64_t stat_cache_max_age, size_t stat_cache_max_entries,
+              TF_Status* status);
+
+void Init(TF_Filesystem* filesystem, TF_Status* status);
+void Cleanup(TF_Filesystem* filesystem);
+void NewRandomAccessFile(const TF_Filesystem* filesystem, const char* path,
+                         TF_RandomAccessFile* file, TF_Status* status);
+void NewWritableFile(const TF_Filesystem* filesystem, const char* path,
+                     TF_WritableFile* file, TF_Status* status);
+void NewAppendableFile(const TF_Filesystem* filesystem, const char* path,
+                       TF_WritableFile* file, TF_Status* status);
+void NewReadOnlyMemoryRegionFromFile(const TF_Filesystem* filesystem,
+                                     const char* path,
+                                     TF_ReadOnlyMemoryRegion* region,
+                                     TF_Status* status);
+int64_t GetFileSize(const TF_Filesystem* filesystem, const char* path,
+                    TF_Status* status);
+void PathExists(const TF_Filesystem* filesystem, const char* path,
+                TF_Status* status);
+void CreateDir(const TF_Filesystem* filesystem, const char* path,
+               TF_Status* status);
+int GetChildren(const TF_Filesystem* filesystem, const char* path,
+                char*** entries, TF_Status* status);
+void DeleteFile(const TF_Filesystem* filesystem, const char* path,
+                TF_Status* status);
+void Stat(const TF_Filesystem* filesystem, const char* path,
+          TF_FileStatistics* stats, TF_Status* status);
+void DeleteDir(const TF_Filesystem* filesystem, const char* path,
+               TF_Status* status);
+void CopyFile(const TF_Filesystem* filesystem, const char* src, const char* dst,
+              TF_Status* status);
+void RenameFile(const TF_Filesystem* filesystem, const char* src,
+                const char* dst, TF_Status* status);
+}  // namespace tf_gcs_filesystem
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_

tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc

@@ -0,0 +1,705 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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 "tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <iostream>
+#include <iterator>
+#include <memory>
+#include <random>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/strings/match.h"
+#include "absl/strings/string_view.h"
+#include "third_party/cloud_cpp/google/cloud/common_options.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/stacktrace_handler.h"
+#include "tensorflow/core/platform/test.h"
+
+#define ASSERT_TF_OK(x) ASSERT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)
+#define EXPECT_TF_OK(x) EXPECT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)
+
+static const char* content = "abcdefghijklmnopqrstuvwxyz1234567890";
+// We will work with content_view instead of content.
+static const absl::string_view content_view = content;
+
+namespace gcs = google::cloud::storage;
+
+static std::string InitializeTmpDir() {
+  // This env should be something like `gs://bucket/path`
+  const char* test_dir = getenv("GCS_TEST_TMPDIR");
+  if (test_dir != nullptr) {
+    std::string bucket, object;
+    TF_Status* status = TF_NewStatus();
+    ParseGCSPath(test_dir, true, &bucket, &object, status);
+    if (TF_GetCode(status) != TF_OK) {
+      TF_DeleteStatus(status);
+      return "";
+    }
+    TF_DeleteStatus(status);
+
+    // We add a random value into `test_dir` to ensures that two consecutive
+    // runs are unlikely to clash.
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<> distribution;
+    std::string rng_val = std::to_string(distribution(gen));
+    return tensorflow::io::JoinPath(std::string(test_dir), rng_val);
+  } else {
+    return "";
+  }
+}
+
+static std::string* GetTmpDir() {
+  static std::string tmp_dir = InitializeTmpDir();
+  if (tmp_dir == "")
+    return nullptr;
+  else
+    return &tmp_dir;
+}
+
+namespace tensorflow {
+namespace {
+
+// TODO(vnvo2409): Refactor `gcs_filesystem_test` to remove unnecessary tests
+// after porting all tests from
+// `//tensorflow/core/platform/cloud:gcs_file_system_test`.
+class GCSFilesystemTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    root_dir_ = io::JoinPath(
+        *GetTmpDir(),
+        ::testing::UnitTest::GetInstance()->current_test_info()->name());
+    status_ = TF_NewStatus();
+    filesystem_ = new TF_Filesystem;
+    filesystem_->plugin_filesystem = nullptr;
+    // Because different tests requires different setup for filesystem. We
+    // initialize filesystem in each testcase.
+  }
+  void TearDown() override {
+    TF_DeleteStatus(status_);
+    if (filesystem_->plugin_filesystem != nullptr)
+      tf_gcs_filesystem::Cleanup(filesystem_);
+    delete filesystem_;
+  }
+
+  std::string GetURIForPath(absl::string_view path) {
+    const std::string translated_name =
+        tensorflow::io::JoinPath(root_dir_, path);
+    return translated_name;
+  }
+
+  std::unique_ptr<TF_WritableFile, void (*)(TF_WritableFile* file)>
+  GetWriter() {
+    std::unique_ptr<TF_WritableFile, void (*)(TF_WritableFile * file)> writer(
+        new TF_WritableFile, [](TF_WritableFile* file) {
+          if (file != nullptr) {
+            if (file->plugin_file != nullptr) tf_writable_file::Cleanup(file);
+            delete file;
+          }
+        });
+    writer->plugin_file = nullptr;
+    return writer;
+  }
+
+  std::unique_ptr<TF_RandomAccessFile, void (*)(TF_RandomAccessFile* file)>
+  GetReader() {
+    std::unique_ptr<TF_RandomAccessFile, void (*)(TF_RandomAccessFile * file)>
+        reader(new TF_RandomAccessFile, [](TF_RandomAccessFile* file) {
+          if (file != nullptr) {
+            if (file->plugin_file != nullptr)
+              tf_random_access_file::Cleanup(file);
+            delete file;
+          }
+        });
+    reader->plugin_file = nullptr;
+    return reader;
+  }
+
+  void WriteString(const std::string& path, const std::string& content) {
+    auto writer = GetWriter();
+    tf_gcs_filesystem::NewWritableFile(filesystem_, path.c_str(), writer.get(),
+                                       status_);
+    if (TF_GetCode(status_) != TF_OK) return;
+    tf_writable_file::Append(writer.get(), content.c_str(), content.length(),
+                             status_);
+    if (TF_GetCode(status_) != TF_OK) return;
+    tf_writable_file::Close(writer.get(), status_);
+    if (TF_GetCode(status_) != TF_OK) return;
+  }
+
+  std::string ReadAll(const std::string& path) {
+    auto reader = GetReader();
+    tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(),
+                                           reader.get(), status_);
+    if (TF_GetCode(status_) != TF_OK) return "";
+
+    auto file_size =
+        tf_gcs_filesystem::GetFileSize(filesystem_, path.c_str(), status_);
+    if (TF_GetCode(status_) != TF_OK) return "";
+
+    std::string content;
+    content.resize(file_size);
+    auto read = tf_random_access_file::Read(reader.get(), 0, file_size,
+                                            &content[0], status_);
+    if (TF_GetCode(status_) != TF_OK) return "";
+    if (read >= 0) content.resize(read);
+    if (file_size != content.size())
+      TF_SetStatus(
+          status_, TF_DATA_LOSS,
+          std::string("expected " + std::to_string(file_size) + " got " +
+                      std::to_string(content.size()) + " bytes")
+              .c_str());
+    return content;
+  }
+
+ protected:
+  TF_Filesystem* filesystem_;
+  TF_Status* status_;
+
+ private:
+  std::string root_dir_;
+};
+
+::testing::AssertionResult WriteToServer(const std::string& path, size_t offset,
+                                         size_t length, gcs::Client* gcs_client,
+                                         TF_Status* status) {
+  std::string bucket, object;
+  ParseGCSPath(path, false, &bucket, &object, status);
+  if (TF_GetCode(status) != TF_OK)
+    return ::testing::AssertionFailure() << TF_Message(status);
+
+  auto writer = gcs_client->WriteObject(bucket, object);
+  writer.write(content + offset, length);
+  writer.Close();
+  if (writer.metadata()) {
+    return ::testing::AssertionSuccess();
+  } else {
+    return ::testing::AssertionFailure()
+           << writer.metadata().status().message();
+  }
+}
+
+::testing::AssertionResult InsertObject(const std::string& path,
+                                        const std::string& content,
+                                        gcs::Client* gcs_client,
+                                        TF_Status* status) {
+  std::string bucket, object;
+  ParseGCSPath(path, false, &bucket, &object, status);
+  if (TF_GetCode(status) != TF_OK)
+    return ::testing::AssertionFailure() << TF_Message(status);
+  auto metadata = gcs_client->InsertObject(bucket, object, content);
+  if (metadata)
+    return ::testing::AssertionSuccess();
+  else
+    return ::testing::AssertionFailure() << metadata.status().message();
+}
+
+::testing::AssertionResult CompareSubString(int64_t offset, size_t length,
+                                            absl::string_view result,
+                                            size_t read) {
+  // Result isn't a null-terminated string so we have to wrap it inside a
+  // `string_view`
+  if (length == read && content_view.substr(offset, length) ==
+                            absl::string_view(result).substr(0, read))
+    return ::testing::AssertionSuccess();
+  else
+    return ::testing::AssertionFailure()
+           << "Result: " << absl::string_view(result).substr(0, read)
+           << " Read: " << read;
+}
+
+::testing::AssertionResult CompareWithServer(const std::string& path,
+                                             size_t offset, size_t length,
+                                             gcs::Client* gcs_client,
+                                             TF_Status* status) {
+  std::string bucket, object;
+  ParseGCSPath(path, false, &bucket, &object, status);
+  if (TF_GetCode(status) != TF_OK)
+    return ::testing::AssertionFailure() << TF_Message(status);
+
+  auto reader = gcs_client->ReadObject(bucket, object);
+  if (!reader) {
+    return ::testing::AssertionFailure() << reader.status().message();
+  } else {
+    std::string content{std::istreambuf_iterator<char>{reader}, {}};
+    return CompareSubString(offset, length, content, content.length());
+  }
+}
+
+TEST_F(GCSFilesystemTest, ParseGCSPath) {
+  std::string bucket, object;
+  ParseGCSPath("gs://bucket/path/to/object", false, &bucket, &object, status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_EQ(bucket, "bucket");
+  ASSERT_EQ(object, "path/to/object");
+
+  ParseGCSPath("gs://bucket/", true, &bucket, &object, status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_EQ(bucket, "bucket");
+
+  ParseGCSPath("bucket/path/to/object", false, &bucket, &object, status_);
+  ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
+
+  // bucket name must end with "/"
+  ParseGCSPath("gs://bucket", true, &bucket, &object, status_);
+  ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
+
+  ParseGCSPath("gs://bucket/", false, &bucket, &object, status_);
+  ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
+}
+
+TEST_F(GCSFilesystemTest, RandomAccessFile) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string filepath = GetURIForPath("a_file");
+  TF_RandomAccessFile* file = new TF_RandomAccessFile;
+  tf_gcs_filesystem::NewRandomAccessFile(filesystem_, filepath.c_str(), file,
+                                         status_);
+  ASSERT_TF_OK(status_);
+  char* result = new char[content_view.length()];
+  int64_t read = tf_random_access_file::Read(file, 0, 1, result, status_);
+  ASSERT_EQ(read, -1) << "Read: " << read;
+  ASSERT_EQ(TF_GetCode(status_), TF_NOT_FOUND) << TF_Message(status_);
+  TF_SetStatus(status_, TF_OK, "");
+
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(WriteToServer(filepath, 0, content_view.length(),
+                            &gcs_file->gcs_client, status_));
+
+  read = tf_random_access_file::Read(file, 0, content_view.length(), result,
+                                     status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_TRUE(CompareSubString(0, content_view.length(), result, read));
+
+  read = tf_random_access_file::Read(file, 0, 4, result, status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_TRUE(CompareSubString(0, 4, result, read));
+
+  read = tf_random_access_file::Read(file, content_view.length() - 2, 4, result,
+                                     status_);
+  ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
+  ASSERT_TRUE(CompareSubString(content_view.length() - 2, 2, result, read));
+
+  delete[] result;
+  tf_random_access_file::Cleanup(file);
+  delete file;
+}
+
+TEST_F(GCSFilesystemTest, WritableFile) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string filepath = GetURIForPath("a_file");
+  TF_WritableFile* file = new TF_WritableFile;
+  tf_gcs_filesystem::NewWritableFile(filesystem_, filepath.c_str(), file,
+                                     status_);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Append(file, content, 4, status_);
+  ASSERT_TF_OK(status_);
+  auto length = tf_writable_file::Tell(file, status_);
+  ASSERT_EQ(length, 4);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Flush(file, status_);
+  ASSERT_TF_OK(status_);
+
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(
+      CompareWithServer(filepath, 0, 4, &gcs_file->gcs_client, status_));
+
+  tf_writable_file::Append(file, content + 4, 4, status_);
+  ASSERT_TF_OK(status_);
+  length = tf_writable_file::Tell(file, status_);
+  ASSERT_EQ(length, 8);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Flush(file, status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_TRUE(
+      CompareWithServer(filepath, 0, 8, &gcs_file->gcs_client, status_));
+
+  tf_writable_file::Close(file, status_);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Cleanup(file);
+
+  // Testing for compose objects
+  gcs_file->compose = true;
+  filepath = GetURIForPath("b_file");
+  tf_gcs_filesystem::NewWritableFile(filesystem_, filepath.c_str(), file,
+                                     status_);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Append(file, content, 4, status_);
+  ASSERT_TF_OK(status_);
+  length = tf_writable_file::Tell(file, status_);
+  ASSERT_EQ(length, 4);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Flush(file, status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_TRUE(
+      CompareWithServer(filepath, 0, 4, &gcs_file->gcs_client, status_));
+
+  tf_writable_file::Append(file, content + 4, 4, status_);
+  ASSERT_TF_OK(status_);
+  length = tf_writable_file::Tell(file, status_);
+  ASSERT_EQ(length, 8);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Flush(file, status_);
+  ASSERT_TF_OK(status_);
+  ASSERT_TRUE(
+      CompareWithServer(filepath, 0, 8, &gcs_file->gcs_client, status_));
+
+  tf_writable_file::Close(file, status_);
+  ASSERT_TF_OK(status_);
+  tf_writable_file::Cleanup(file);
+  delete file;
+}
+
+TEST_F(GCSFilesystemTest, ReadOnlyMemoryRegion) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string path = GetURIForPath("a_file");
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(WriteToServer(path, 0, 0, &gcs_file->gcs_client, status_));
+  TF_ReadOnlyMemoryRegion* region = new TF_ReadOnlyMemoryRegion;
+  tf_gcs_filesystem::NewReadOnlyMemoryRegionFromFile(filesystem_, path.c_str(),
+                                                     region, status_);
+  ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT) << TF_Message(status_);
+
+  TF_SetStatus(status_, TF_OK, "");
+  ASSERT_TRUE(WriteToServer(path, 0, content_view.length(),
+                            &gcs_file->gcs_client, status_));
+  tf_gcs_filesystem::NewReadOnlyMemoryRegionFromFile(filesystem_, path.c_str(),
+                                                     region, status_);
+  ASSERT_TF_OK(status_);
+  auto length = tf_read_only_memory_region::Length(region);
+  ASSERT_EQ(length, content_view.length());
+  auto data =
+      static_cast<const char*>(tf_read_only_memory_region::Data(region));
+  ASSERT_TRUE(CompareSubString(0, content_view.length(), data, length));
+
+  tf_read_only_memory_region::Cleanup(region);
+  delete region;
+}
+
+TEST_F(GCSFilesystemTest, PathExists) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string path = GetURIForPath("PathExists");
+  tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
+  EXPECT_EQ(TF_NOT_FOUND, TF_GetCode(status_)) << TF_Message(status_);
+  TF_SetStatus(status_, TF_OK, "");
+  WriteString(path, "test");
+  ASSERT_TF_OK(status_);
+  tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
+  EXPECT_TF_OK(status_);
+}
+
+TEST_F(GCSFilesystemTest, GetChildren) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string base = GetURIForPath("GetChildren");
+  tf_gcs_filesystem::CreateDir(filesystem_, base.c_str(), status_);
+  EXPECT_TF_OK(status_);
+
+  const std::string file = io::JoinPath(base, "TestFile.csv");
+  WriteString(file, "test");
+  EXPECT_TF_OK(status_);
+
+  const std::string subdir = io::JoinPath(base, "SubDir");
+  tf_gcs_filesystem::CreateDir(filesystem_, subdir.c_str(), status_);
+  EXPECT_TF_OK(status_);
+  const std::string subfile = io::JoinPath(subdir, "TestSubFile.csv");
+  WriteString(subfile, "test");
+  EXPECT_TF_OK(status_);
+
+  char** entries;
+  auto num_entries = tf_gcs_filesystem::GetChildren(filesystem_, base.c_str(),
+                                                    &entries, status_);
+  EXPECT_TF_OK(status_);
+
+  std::vector<std::string> childrens;
+  for (int i = 0; i < num_entries; ++i) {
+    childrens.push_back(entries[i]);
+  }
+  std::sort(childrens.begin(), childrens.end());
+  EXPECT_EQ(std::vector<std::string>({"SubDir/", "TestFile.csv"}), childrens);
+}
+
+TEST_F(GCSFilesystemTest, DeleteFile) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string path = GetURIForPath("DeleteFile");
+  WriteString(path, "test");
+  ASSERT_TF_OK(status_);
+  tf_gcs_filesystem::DeleteFile(filesystem_, path.c_str(), status_);
+  EXPECT_TF_OK(status_);
+  tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
+  EXPECT_EQ(TF_GetCode(status_), TF_NOT_FOUND);
+}
+
+TEST_F(GCSFilesystemTest, CreateDir) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string dir = GetURIForPath("CreateDir");
+  tf_gcs_filesystem::CreateDir(filesystem_, dir.c_str(), status_);
+  EXPECT_TF_OK(status_);
+
+  TF_FileStatistics stat;
+  tf_gcs_filesystem::Stat(filesystem_, dir.c_str(), &stat, status_);
+  EXPECT_TF_OK(status_);
+  EXPECT_TRUE(stat.is_directory);
+}
+
+TEST_F(GCSFilesystemTest, DeleteDir) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string dir = GetURIForPath("DeleteDir");
+  const std::string file = io::JoinPath(dir, "DeleteDirFile.csv");
+  WriteString(file, "test");
+  ASSERT_TF_OK(status_);
+  tf_gcs_filesystem::DeleteDir(filesystem_, dir.c_str(), status_);
+  EXPECT_EQ(TF_GetCode(status_), TF_FAILED_PRECONDITION);
+
+  TF_SetStatus(status_, TF_OK, "");
+  tf_gcs_filesystem::DeleteFile(filesystem_, file.c_str(), status_);
+  EXPECT_TF_OK(status_);
+  tf_gcs_filesystem::DeleteDir(filesystem_, dir.c_str(), status_);
+  EXPECT_TF_OK(status_);
+  TF_FileStatistics stat;
+  tf_gcs_filesystem::Stat(filesystem_, dir.c_str(), &stat, status_);
+  EXPECT_EQ(TF_GetCode(status_), TF_NOT_FOUND) << TF_Message(status_);
+}
+
+TEST_F(GCSFilesystemTest, StatFile) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string path = GetURIForPath("StatFile");
+  WriteString(path, "test");
+  ASSERT_TF_OK(status_);
+
+  TF_FileStatistics stat;
+  tf_gcs_filesystem::Stat(filesystem_, path.c_str(), &stat, status_);
+  EXPECT_TF_OK(status_);
+  EXPECT_EQ(4, stat.length);
+  EXPECT_FALSE(stat.is_directory);
+}
+
+TEST_F(GCSFilesystemTest, RenameFile) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string src = GetURIForPath("RenameFileSrc");
+  const std::string dst = GetURIForPath("RenameFileDst");
+  WriteString(src, "test");
+  ASSERT_TF_OK(status_);
+
+  tf_gcs_filesystem::RenameFile(filesystem_, src.c_str(), dst.c_str(), status_);
+  EXPECT_TF_OK(status_);
+  auto result = ReadAll(dst);
+  EXPECT_TF_OK(status_);
+  EXPECT_EQ("test", result);
+}
+
+TEST_F(GCSFilesystemTest, RenameFileOverwrite) {
+  tf_gcs_filesystem::Init(filesystem_, status_);
+  ASSERT_TF_OK(status_);
+  const std::string src = GetURIForPath("RenameFileOverwriteSrc");
+  const std::string dst = GetURIForPath("RenameFileOverwriteDst");
+
+  WriteString(src, "test_old");
+  ASSERT_TF_OK(status_);
+  WriteString(dst, "test_new");
+  ASSERT_TF_OK(status_);
+
+  tf_gcs_filesystem::PathExists(filesystem_, dst.c_str(), status_);
+  EXPECT_TF_OK(status_);
+  tf_gcs_filesystem::RenameFile(filesystem_, src.c_str(), dst.c_str(), status_);
+  EXPECT_TF_OK(status_);
+
+  auto result = ReadAll(dst);
+  EXPECT_TF_OK(status_);
+  EXPECT_EQ("test_old", result);
+}
+
+// These tests below are ported from
+// `//tensorflow/core/platform/cloud:gcs_file_system_test`
+TEST_F(GCSFilesystemTest, NewRandomAccessFile_NoBlockCache) {
+  tf_gcs_filesystem::InitTest(filesystem_, false, 0, 0, 0, 0, 0, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string path = GetURIForPath("a_file");
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
+
+  TF_RandomAccessFile* file = new TF_RandomAccessFile;
+  tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
+                                         status_);
+  ASSERT_TF_OK(status_);
+
+  std::string result;
+  result.resize(6);
+  int64_t read = tf_random_access_file::Read(file, 0, 6, &result[0], status_);
+  ASSERT_EQ(read, 6) << "Read: " << read << "\n";
+  ASSERT_TF_OK(status_);
+  ASSERT_EQ(result, "012345") << "Result: " << result << "\n";
+
+  read = tf_random_access_file::Read(file, 6, 6, &result[0], status_);
+  ASSERT_EQ(read, 4) << "Read: " << read << "\n";
+  ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
+  result.resize(read);
+  ASSERT_EQ(result, "6789") << "Result: " << result << "\n";
+}
+
+TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered) {
+  tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string path = GetURIForPath("a_file");
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
+
+  TF_RandomAccessFile* file = new TF_RandomAccessFile;
+  tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
+                                         status_);
+  ASSERT_TF_OK(status_);
+
+  std::string result;
+  result.resize(6);
+  int64_t read = tf_random_access_file::Read(file, 0, 6, &result[0], status_);
+  ASSERT_EQ(read, 6) << "Read: " << read << "\n";
+  ASSERT_TF_OK(status_);
+  ASSERT_EQ(result, "012345") << "Result: " << result << "\n";
+
+  read = tf_random_access_file::Read(file, 6, 6, &result[0], status_);
+  ASSERT_EQ(read, 4) << "Read: " << read << "\n";
+  ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
+  result.resize(read);
+  ASSERT_EQ(result, "6789") << "Result: " << result << "\n";
+}
+
+TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered_ReadAtEOF) {
+  tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string path = GetURIForPath("a_file");
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
+
+  TF_RandomAccessFile* file = new TF_RandomAccessFile;
+  tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
+                                         status_);
+  ASSERT_TF_OK(status_);
+
+  std::string result;
+  result.resize(10);
+  int64_t read = tf_random_access_file::Read(file, 0, result.length(),
+                                             &result[0], status_);
+  ASSERT_EQ(read, 10) << "Read: " << read << "\n";
+  ASSERT_TF_OK(status_);
+  ASSERT_EQ(result, "0123456789") << "Result: " << result << "\n";
+
+  read = tf_random_access_file::Read(file, result.length(), result.length(),
+                                     &result[0], status_);
+  ASSERT_EQ(read, 0) << "Read: " << read << "\n";
+  ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
+  result.resize(read);
+  ASSERT_EQ(result, "") << "Result: " << result << "\n";
+}
+
+TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered_CachedOutOfRange) {
+  tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  std::string path = GetURIForPath("a_file");
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  ASSERT_TRUE(InsertObject(path, "012345678", &gcs_file->gcs_client, status_));
+
+  TF_RandomAccessFile* file = new TF_RandomAccessFile;
+  tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
+                                         status_);
+  ASSERT_TF_OK(status_);
+
+  std::string result;
+  result.resize(5);
+  int64_t read = tf_random_access_file::Read(file, 0, result.length(),
+                                             &result[0], status_);
+  ASSERT_EQ(read, 5) << "Read: " << read << "\n";
+  ASSERT_TF_OK(status_);
+  ASSERT_EQ(result, "01234") << "Result: " << result << "\n";
+
+  read = tf_random_access_file::Read(file, 4, result.length(), &result[0],
+                                     status_);
+  ASSERT_EQ(read, 5) << "Read: " << read << "\n";
+  ASSERT_TF_OK(status_);
+  result.resize(read);
+  ASSERT_EQ(result, "45678") << "Result: " << result << "\n";
+
+  read = tf_random_access_file::Read(file, 5, result.length(), &result[0],
+                                     status_);
+  ASSERT_EQ(read, 4) << "Read: " << read << "\n";
+  ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
+  result.resize(read);
+  ASSERT_EQ(result, "5678") << "Result: " << result << "\n";
+}
+
+TEST_F(GCSFilesystemTest, TestGetStorageClientOptions) {
+  tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
+  ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
+                        << TF_Message(status_);
+  auto gcs_file =
+      static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
+  // Get the raw connection
+  auto connection = gcs_file->gcs_client.raw_client();
+  // Get the options from the connection
+  auto options = connection->options();
+
+  // Access the UserAgentProductsOptions which contain information about
+  // the UserAgent
+  auto user_agent_products =
+      options.get<google::cloud::UserAgentProductsOption>();
+  bool found_x_goog_api_client = false;
+  for (const auto& product : user_agent_products) {
+    if (absl::StrContains(product, "TensorFlow-C-API")) {
+      found_x_goog_api_client = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(found_x_goog_api_client);
+}
+
+}  // namespace
+}  // namespace tensorflow
+
+GTEST_API_ int main(int argc, char** argv) {
+  tensorflow::testing::InstallStacktraceHandler();
+  if (!GetTmpDir()) {
+    std::cerr << "Could not read GCS_TEST_TMPDIR env";
+    return -1;
+  }
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc

@@ -0,0 +1,42 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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 "tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
+
+#include <stdio.h>
+
+#include <cstdio>
+#include <fstream>
+#include <ios>
+#include <string>
+#include <utility>
+
+TempFile::TempFile(const std::string& temp_file_name, std::ios::openmode mode)
+    : std::fstream(temp_file_name, mode), name_(temp_file_name) {}
+
+TempFile::TempFile(TempFile&& rhs)
+    : std::fstream(std::move(rhs)), name_(std::move(rhs.name_)) {}
+
+TempFile::~TempFile() {
+  std::fstream::close();
+  std::remove(name_.c_str());
+}
+
+const std::string TempFile::getName() const { return name_; }
+
+bool TempFile::truncate() {
+  std::fstream::close();
+  std::fstream::open(name_, std::ios::binary | std::ios::out);
+  return std::fstream::is_open();
+}

tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h

@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
+
+#include <fstream>
+#include <ios>
+#include <string>
+
+class TempFile : public std::fstream {
+ public:
+  // We should specify openmode each time we call TempFile.
+  TempFile(const std::string& temp_file_name, std::ios::openmode mode);
+  TempFile(TempFile&& rhs);
+  ~TempFile() override;
+  const std::string getName() const;
+  bool truncate();
+
+ private:
+  const std::string name_;
+};
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_

tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc

@@ -0,0 +1,316 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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 "tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
+
+#include <cstdint>
+#include <cstring>
+#include <iostream>
+#include <iterator>
+#include <limits>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/synchronization/mutex.h"
+#include "absl/time/time.h"
+#include "tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h"
+
+namespace tf_gcs_filesystem {
+
+bool RamFileBlockCache::BlockNotStale(const std::shared_ptr<Block>& block) {
+  absl::MutexLock l(block->mu);
+  if (block->state != FetchState::FINISHED) {
+    return true;  // No need to check for staleness.
+  }
+  if (max_staleness_ == 0) return true;  // Not enforcing staleness.
+  return timer_seconds_() - block->timestamp <= max_staleness_;
+}
+
+std::shared_ptr<RamFileBlockCache::Block> RamFileBlockCache::Lookup(
+    const Key& key) {
+  absl::MutexLock lock(mu_);
+  auto entry = block_map_.find(key);
+  if (entry != block_map_.end()) {
+    if (BlockNotStale(entry->second)) {
+      return entry->second;
+    } else {
+      // Remove the stale block and continue.
+      RemoveFile_Locked(key.first);
+    }
+  }
+
+  // Insert a new empty block, setting the bookkeeping to sentinel values
+  // in order to update them as appropriate.
+  auto new_entry = std::make_shared<Block>();
+  lru_list_.push_front(key);
+  lra_list_.push_front(key);
+  new_entry->lru_iterator = lru_list_.begin();
+  new_entry->lra_iterator = lra_list_.begin();
+  new_entry->timestamp = timer_seconds_();
+  block_map_.emplace(std::make_pair(key, new_entry));
+  return new_entry;
+}
+
+// Remove blocks from the cache until we do not exceed our maximum size.
+void RamFileBlockCache::Trim() {
+  while (!lru_list_.empty() && cache_size_ > max_bytes_) {
+    RemoveBlock(block_map_.find(lru_list_.back()));
+  }
+}
+
+/// Move the block to the front of the LRU list if it isn't already there.
+void RamFileBlockCache::UpdateLRU(const Key& key,
+                                  const std::shared_ptr<Block>& block,
+                                  TF_Status* status) {
+  absl::MutexLock lock(mu_);
+  if (block->timestamp == 0) {
+    // The block was evicted from another thread. Allow it to remain evicted.
+    return TF_SetStatus(status, TF_OK, "");
+  }
+  if (block->lru_iterator != lru_list_.begin()) {
+    lru_list_.erase(block->lru_iterator);
+    lru_list_.push_front(key);
+    block->lru_iterator = lru_list_.begin();
+  }
+
+  // Check for inconsistent state. If there is a block later in the same file
+  // in the cache, and our current block is not block size, this likely means
+  // we have inconsistent state within the cache. Note: it's possible some
+  // incomplete reads may still go undetected.
+  if (block->data.size() < block_size_) {
+    Key fmax = std::make_pair(key.first, std::numeric_limits<size_t>::max());
+    auto fcmp = block_map_.upper_bound(fmax);
+    if (fcmp != block_map_.begin() && key < (--fcmp)->first) {
+      return TF_SetStatus(status, TF_INTERNAL,
+                          "Block cache contents are inconsistent.");
+    }
+  }
+
+  Trim();
+
+  return TF_SetStatus(status, TF_OK, "");
+}
+
+void RamFileBlockCache::MaybeFetch(const Key& key,
+                                   const std::shared_ptr<Block>& block,
+                                   TF_Status* status) {
+  bool downloaded_block = false;
+  auto reconcile_state = MakeCleanup([this, &downloaded_block, &key, &block] {
+    // Perform this action in a cleanup callback to avoid locking mu_ after
+    // locking block->mu.
+    if (downloaded_block) {
+      absl::MutexLock l(mu_);
+      // Do not update state if the block is already to be evicted.
+      if (block->timestamp != 0) {
+        // Use capacity() instead of size() to account for all  memory
+        // used by the cache.
+        cache_size_ += block->data.capacity();
+        // Put to beginning of LRA list.
+        lra_list_.erase(block->lra_iterator);
+        lra_list_.push_front(key);
+        block->lra_iterator = lra_list_.begin();
+        block->timestamp = timer_seconds_();
+      }
+    }
+  });
+  // Loop until either block content is successfully fetched, or our request
+  // encounters an error.
+  absl::MutexLock l(block->mu);
+  TF_SetStatus(status, TF_OK, "");
+  while (true) {
+    switch (block->state) {
+      case FetchState::ERROR:
+        // TF_FALLTHROUGH_INTENDED
+      case FetchState::CREATED:
+        block->state = FetchState::FETCHING;
+        block->mu.unlock();  // Release the lock while making the API call.
+        block->data.clear();
+        block->data.resize(block_size_, 0);
+        int64_t bytes_transferred;
+        bytes_transferred = block_fetcher_(key.first, key.second, block_size_,
+                                           block->data.data(), status);
+        block->mu.lock();  // Reacquire the lock immediately afterwards
+        if (TF_GetCode(status) == TF_OK) {
+          block->data.resize(bytes_transferred, 0);
+          // Shrink the data capacity to the actual size used.
+          // NOLINTNEXTLINE: shrink_to_fit() may not shrink the capacity.
+          std::vector<char>(block->data).swap(block->data);
+          downloaded_block = true;
+          block->state = FetchState::FINISHED;
+        } else {
+          block->state = FetchState::ERROR;
+        }
+        block->cond_var.SignalAll();
+        return;
+      case FetchState::FETCHING:
+        block->cond_var.WaitWithTimeout(&block->mu, absl::Minutes(1));
+        if (block->state == FetchState::FINISHED) {
+          return TF_SetStatus(status, TF_OK, "");
+        }
+        // Re-loop in case of errors.
+        break;
+      case FetchState::FINISHED:
+        return TF_SetStatus(status, TF_OK, "");
+    }
+  }
+  return TF_SetStatus(
+      status, TF_INTERNAL,
+      "Control flow should never reach the end of RamFileBlockCache::Fetch.");
+}
+
+int64_t RamFileBlockCache::Read(const std::string& filename, size_t offset,
+                                size_t n, char* buffer, TF_Status* status) {
+  if (n == 0) {
+    TF_SetStatus(status, TF_OK, "");
+    return 0;
+  }
+  if (!IsCacheEnabled() || (n > max_bytes_)) {
+    // The cache is effectively disabled, so we pass the read through to the
+    // fetcher without breaking it up into blocks.
+    return block_fetcher_(filename, offset, n, buffer, status);
+  }
+  // Calculate the block-aligned start and end of the read.
+  size_t start = block_size_ * (offset / block_size_);
+  size_t finish = block_size_ * ((offset + n) / block_size_);
+  if (finish < offset + n) {
+    finish += block_size_;
+  }
+  size_t total_bytes_transferred = 0;
+  // Now iterate through the blocks, reading them one at a time.
+  for (size_t pos = start; pos < finish; pos += block_size_) {
+    Key key = std::make_pair(filename, pos);
+    // Look up the block, fetching and inserting it if necessary, and update the
+    // LRU iterator for the key and block.
+    std::shared_ptr<Block> block = Lookup(key);
+    if (!block) {
+      std::cerr << "No block for key " << key.first << "@" << key.second;
+      abort();
+    }
+    MaybeFetch(key, block, status);
+    if (TF_GetCode(status) != TF_OK) return -1;
+    UpdateLRU(key, block, status);
+    if (TF_GetCode(status) != TF_OK) return -1;
+    // Copy the relevant portion of the block into the result buffer.
+    const auto& data = block->data;
+    if (offset >= pos + data.size()) {
+      // The requested offset is at or beyond the end of the file. This can
+      // happen if `offset` is not block-aligned, and the read returns the last
+      // block in the file, which does not extend all the way out to `offset`.
+      std::stringstream os;
+      os << "EOF at offset " << offset << " in file " << filename
+         << " at position " << pos << " with data size " << data.size();
+      TF_SetStatus(status, TF_OUT_OF_RANGE, std::move(os).str().c_str());
+      return total_bytes_transferred;
+    }
+    auto begin = data.begin();
+    if (offset > pos) {
+      // The block begins before the slice we're reading.
+      begin += offset - pos;
+    }
+    auto end = data.end();
+    if (pos + data.size() > offset + n) {
+      // The block extends past the end of the slice we're reading.
+      end -= (pos + data.size()) - (offset + n);
+    }
+    if (begin < end) {
+      size_t bytes_to_copy = end - begin;
+      memcpy(&buffer[total_bytes_transferred], &*begin, bytes_to_copy);
+      total_bytes_transferred += bytes_to_copy;
+    }
+    if (data.size() < block_size_) {
+      // The block was a partial block and thus signals EOF at its upper bound.
+      break;
+    }
+  }
+  TF_SetStatus(status, TF_OK, "");
+  return total_bytes_transferred;
+}
+
+bool RamFileBlockCache::ValidateAndUpdateFileSignature(
+    const std::string& filename, int64_t file_signature) {
+  absl::MutexLock lock(mu_);
+  auto it = file_signature_map_.find(filename);
+  if (it != file_signature_map_.end()) {
+    if (it->second == file_signature) {
+      return true;
+    }
+    // Remove the file from cache if the signatures don't match.
+    RemoveFile_Locked(filename);
+    it->second = file_signature;
+    return false;
+  }
+  file_signature_map_[filename] = file_signature;
+  return true;
+}
+
+size_t RamFileBlockCache::CacheSize() const {
+  absl::MutexLock lock(mu_);
+  return cache_size_;
+}
+
+void RamFileBlockCache::Prune() {
+  while (!stop_pruning_thread_.WaitForNotificationWithTimeout(
+      absl::Microseconds(1000000))) {
+    absl::MutexLock lock(mu_);
+    uint64_t now = timer_seconds_();
+    while (!lra_list_.empty()) {
+      auto it = block_map_.find(lra_list_.back());
+      if (now - it->second->timestamp <= max_staleness_) {
+        // The oldest block is not yet expired. Come back later.
+        break;
+      }
+      // We need to make a copy of the filename here, since it could otherwise
+      // be used within RemoveFile_Locked after `it` is deleted.
+      RemoveFile_Locked(std::string(it->first.first));
+    }
+  }
+}
+
+void RamFileBlockCache::Flush() {
+  absl::MutexLock lock(mu_);
+  block_map_.clear();
+  lru_list_.clear();
+  lra_list_.clear();
+  cache_size_ = 0;
+}
+
+void RamFileBlockCache::RemoveFile(const std::string& filename) {
+  absl::MutexLock lock(mu_);
+  RemoveFile_Locked(filename);
+}
+
+void RamFileBlockCache::RemoveFile_Locked(const std::string& filename) {
+  Key begin = std::make_pair(filename, 0);
+  auto it = block_map_.lower_bound(begin);
+  while (it != block_map_.end() && it->first.first == filename) {
+    auto next = std::next(it);
+    RemoveBlock(it);
+    it = next;
+  }
+}
+
+void RamFileBlockCache::RemoveBlock(BlockMap::iterator entry) {
+  // This signals that the block is removed, and should not be inadvertently
+  // reinserted into the cache in UpdateLRU.
+  entry->second->timestamp = 0;
+  lru_list_.erase(entry->second->lru_iterator);
+  lra_list_.erase(entry->second->lra_iterator);
+  cache_size_ -= entry->second->data.capacity();
+  block_map_.erase(entry);
+}
+
+}  // namespace tf_gcs_filesystem

tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h

@@ -0,0 +1,270 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <iostream>
+#include <list>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/log/log.h"
+#include "absl/strings/str_format.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/synchronization/notification.h"
+#include "tensorflow/c/env.h"
+#include "tensorflow/c/tf_status.h"
+
+namespace tf_gcs_filesystem {
+
+/// \brief An LRU block cache of file contents, keyed by {filename, offset}.
+///
+/// This class should be shared by read-only random access files on a remote
+/// filesystem (e.g. GCS).
+class RamFileBlockCache {
+ public:
+  /// The callback executed when a block is not found in the cache, and needs to
+  /// be fetched from the backing filesystem. This callback is provided when the
+  /// cache is constructed. It returns total bytes read ( -1 in case of errors
+  /// ). The `status` should be `TF_OK` as long as the read from the remote
+  /// filesystem succeeded (similar to the semantics of the read(2) system
+  /// call).
+  typedef std::function<int64_t(const std::string& filename, size_t offset,
+                                size_t buffer_size, char* buffer,
+                                TF_Status* status)>
+      BlockFetcher;
+
+  RamFileBlockCache(size_t block_size, size_t max_bytes, uint64_t max_staleness,
+                    BlockFetcher block_fetcher,
+                    std::function<uint64_t()> timer_seconds = TF_NowSeconds)
+      : block_size_(block_size),
+        max_bytes_(max_bytes),
+        max_staleness_(max_staleness),
+        block_fetcher_(block_fetcher),
+        timer_seconds_(timer_seconds),
+        pruning_thread_(nullptr,
+                        [](TF_Thread* thread) { TF_JoinThread(thread); }) {
+    if (max_staleness_ > 0) {
+      TF_ThreadOptions thread_options;
+      TF_DefaultThreadOptions(&thread_options);
+      pruning_thread_.reset(
+          TF_StartThread(&thread_options, "TF_prune_FBC", PruneThread, this));
+    }
+    VLOG(1) << absl::StrFormat("GCS file block cache is %s.\n",
+                               (IsCacheEnabled() ? "enabled" : "disabled"));
+  }
+
+  ~RamFileBlockCache() {
+    if (pruning_thread_) {
+      stop_pruning_thread_.Notify();
+      // Destroying pruning_thread_ will block until Prune() receives the above
+      // notification and returns.
+      pruning_thread_.reset();
+    }
+  }
+
+  /// Read `n` bytes from `filename` starting at `offset` into `buffer`. It
+  /// returns total bytes read ( -1 in case of errors ). This method will set
+  /// `status` to:
+  ///
+  /// 1) The error from the remote filesystem, if the read from the remote
+  ///    filesystem failed.
+  /// 2) `TF_FAILED_PRECONDITION` if the read from the remote filesystem
+  /// succeeded,
+  ///    but the read returned a partial block, and the LRU cache contained a
+  ///    block at a higher offset (indicating that the partial block should have
+  ///    been a full block).
+  /// 3) `TF_OUT_OF_RANGE` if the read from the remote filesystem succeeded, but
+  ///    the file contents do not extend past `offset` and thus nothing was
+  ///    placed in `out`.
+  /// 4) `TF_OK` otherwise (i.e. the read succeeded, and at least one byte was
+  /// placed
+  ///    in `buffer`).
+  ///
+  /// Caller is responsible for allocating memory for `buffer`.
+  /// `buffer` will be left unchanged in case of errors.
+  int64_t Read(const std::string& filename, size_t offset, size_t n,
+               char* buffer, TF_Status* status);
+
+  // Validate the given file signature with the existing file signature in the
+  // cache. Returns true if the signature doesn't change or the file doesn't
+  // exist before. If the signature changes, update the existing signature with
+  // the new one and remove the file from cache.
+  bool ValidateAndUpdateFileSignature(const std::string& filename,
+                                      int64_t file_signature)
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  /// Remove all cached blocks for `filename`.
+  void RemoveFile(const std::string& filename) ABSL_LOCKS_EXCLUDED(mu_);
+
+  /// Remove all cached data.
+  void Flush() ABSL_LOCKS_EXCLUDED(mu_);
+
+  /// Accessors for cache parameters.
+  size_t block_size() const { return block_size_; }
+  size_t max_bytes() const { return max_bytes_; }
+  uint64_t max_staleness() const { return max_staleness_; }
+
+  /// The current size (in bytes) of the cache.
+  size_t CacheSize() const ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Returns true if the cache is enabled. If false, the BlockFetcher callback
+  // is always executed during Read.
+  bool IsCacheEnabled() const { return block_size_ > 0 && max_bytes_ > 0; }
+
+  // We can not pass a lambda with capture as a function pointer to
+  // `TF_StartThread`, so we have to wrap `Prune` inside a static function.
+  static void PruneThread(void* param) {
+    auto ram_file_block_cache = static_cast<RamFileBlockCache*>(param);
+    ram_file_block_cache->Prune();
+  }
+
+ private:
+  /// The size of the blocks stored in the LRU cache, as well as the size of the
+  /// reads from the underlying filesystem.
+  const size_t block_size_;
+  /// The maximum number of bytes (sum of block sizes) allowed in the LRU cache.
+  const size_t max_bytes_;
+  /// The maximum staleness of any block in the LRU cache, in seconds.
+  const uint64_t max_staleness_;
+  /// The callback to read a block from the underlying filesystem.
+  const BlockFetcher block_fetcher_;
+  /// The callback to read timestamps.
+  const std::function<uint64_t()> timer_seconds_;
+
+  /// \brief The key type for the file block cache.
+  ///
+  /// The file block cache key is a {filename, offset} pair.
+  typedef std::pair<std::string, size_t> Key;
+
+  /// \brief The state of a block.
+  ///
+  /// A block begins in the CREATED stage. The first thread will attempt to read
+  /// the block from the filesystem, transitioning the state of the block to
+  /// FETCHING. After completing, if the read was successful the state should
+  /// be FINISHED. Otherwise the state should be ERROR. A subsequent read can
+  /// re-fetch the block if the state is ERROR.
+  enum class FetchState {
+    CREATED,
+    FETCHING,
+    FINISHED,
+    ERROR,
+  };
+
+  /// \brief A block of a file.
+  ///
+  /// A file block consists of the block data, the block's current position in
+  /// the LRU cache, the timestamp (seconds since epoch) at which the block
+  /// was cached, a coordination lock, and state & condition variables.
+  ///
+  /// Thread safety:
+  /// The iterator and timestamp fields should only be accessed while holding
+  /// the block-cache-wide mu_ instance variable. The state variable should only
+  /// be accessed while holding the Block's mu lock. The data vector should only
+  /// be accessed after state == FINISHED, and it should never be modified.
+  ///
+  /// In order to prevent deadlocks, never grab the block-cache-wide mu_ lock
+  /// AFTER grabbing any block's mu lock. It is safe to grab mu without locking
+  /// mu_.
+  struct Block {
+    /// The block data.
+    std::vector<char> data;
+    /// A list iterator pointing to the block's position in the LRU list.
+    std::list<Key>::iterator lru_iterator;
+    /// A list iterator pointing to the block's position in the LRA list.
+    std::list<Key>::iterator lra_iterator;
+    /// The timestamp (seconds since epoch) at which the block was cached.
+    uint64_t timestamp;
+    /// Mutex to guard state variable
+    absl::Mutex mu;
+    /// The state of the block.
+    FetchState state ABSL_GUARDED_BY(mu) = FetchState::CREATED;
+    /// Wait on cond_var if state is FETCHING.
+    absl::CondVar cond_var;
+  };
+
+  /// \brief The block map type for the file block cache.
+  ///
+  /// The block map is an ordered map from Key to Block.
+  typedef std::map<Key, std::shared_ptr<Block>> BlockMap;
+
+  /// Prune the cache by removing files with expired blocks.
+  void Prune() ABSL_LOCKS_EXCLUDED(mu_);
+
+  bool BlockNotStale(const std::shared_ptr<Block>& block)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  /// Look up a Key in the block cache.
+  std::shared_ptr<Block> Lookup(const Key& key) ABSL_LOCKS_EXCLUDED(mu_);
+
+  void MaybeFetch(const Key& key, const std::shared_ptr<Block>& block,
+                  TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
+
+  /// Trim the block cache to make room for another entry.
+  void Trim() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  /// Update the LRU iterator for the block at `key`.
+  void UpdateLRU(const Key& key, const std::shared_ptr<Block>& block,
+                 TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
+
+  /// Remove all blocks of a file, with mu_ already held.
+  void RemoveFile_Locked(const std::string& filename)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  /// Remove the block `entry` from the block map and LRU list, and update the
+  /// cache size accordingly.
+  void RemoveBlock(BlockMap::iterator entry) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  /// The cache pruning thread that removes files with expired blocks.
+  std::unique_ptr<TF_Thread, std::function<void(TF_Thread*)>> pruning_thread_;
+
+  /// Notification for stopping the cache pruning thread.
+  absl::Notification stop_pruning_thread_;
+
+  /// Guards access to the block map, LRU list, and cached byte count.
+  mutable absl::Mutex mu_;
+
+  /// The block map (map from Key to Block).
+  BlockMap block_map_ ABSL_GUARDED_BY(mu_);
+
+  /// The LRU list of block keys. The front of the list identifies the most
+  /// recently accessed block.
+  std::list<Key> lru_list_ ABSL_GUARDED_BY(mu_);
+
+  /// The LRA (least recently added) list of block keys. The front of the list
+  /// identifies the most recently added block.
+  ///
+  /// Note: blocks are added to lra_list_ only after they have successfully been
+  /// fetched from the underlying block store.
+  std::list<Key> lra_list_ ABSL_GUARDED_BY(mu_);
+
+  /// The combined number of bytes in all of the cached blocks.
+  size_t cache_size_ ABSL_GUARDED_BY(mu_) = 0;
+
+  // A filename->file_signature map.
+  std::map<std::string, int64_t> file_signature_map_ ABSL_GUARDED_BY(mu_);
+};
+
+}  // namespace tf_gcs_filesystem
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_

tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc

@@ -0,0 +1,614 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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 "tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <list>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/ascii.h"
+#include "absl/synchronization/blocking_counter.h"
+#include "absl/synchronization/notification.h"
+#include "absl/time/time.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/cloud/now_seconds_env.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+absl::Status ReadCache(tf_gcs_filesystem::RamFileBlockCache* cache,
+                       const std::string& filename, size_t offset, size_t n,
+                       std::vector<char>* out) {
+  out->clear();
+  out->resize(n, 0);
+  TF_Status status;
+  auto bytes_transferred =
+      cache->Read(filename, offset, n, out->data(), &status);
+  if (bytes_transferred >= 0) {
+    EXPECT_LE(bytes_transferred, n);
+    out->resize(bytes_transferred, n);
+  }
+  return status.status;
+}
+
+TEST(RamFileBlockCacheTest, IsCacheEnabled) {
+  auto fetcher = [](const std::string& filename, size_t offset, size_t n,
+                    char* buffer, TF_Status* status) -> int64_t {
+    // Do nothing.
+    TF_SetStatus(status, TF_OK, "");
+    return 0;
+  };
+  tf_gcs_filesystem::RamFileBlockCache cache1(0, 0, 0, fetcher);
+  tf_gcs_filesystem::RamFileBlockCache cache2(16, 0, 0, fetcher);
+  tf_gcs_filesystem::RamFileBlockCache cache3(0, 32, 0, fetcher);
+  tf_gcs_filesystem::RamFileBlockCache cache4(16, 32, 0, fetcher);
+
+  EXPECT_FALSE(cache1.IsCacheEnabled());
+  EXPECT_FALSE(cache2.IsCacheEnabled());
+  EXPECT_FALSE(cache3.IsCacheEnabled());
+  EXPECT_TRUE(cache4.IsCacheEnabled());
+}
+
+TEST(RamFileBlockCacheTest, ValidateAndUpdateFileSignature) {
+  int calls = 0;
+  auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
+                          char* buffer, TF_Status* status) -> int64_t {
+    calls++;
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  std::string filename = "file";
+  tf_gcs_filesystem::RamFileBlockCache cache(16, 32, 0, fetcher);
+  std::vector<char> out;
+
+  // First read.
+  EXPECT_TRUE(cache.ValidateAndUpdateFileSignature(filename, 123));
+  TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
+  EXPECT_EQ(calls, 1);
+
+  // Second read. Hit cache.
+  EXPECT_TRUE(cache.ValidateAndUpdateFileSignature(filename, 123));
+  TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
+  EXPECT_EQ(calls, 1);
+
+  // Third read. File signatures are different.
+  EXPECT_FALSE(cache.ValidateAndUpdateFileSignature(filename, 321));
+  TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
+  EXPECT_EQ(calls, 2);
+}
+
+TEST(RamFileBlockCacheTest, PassThrough) {
+  const std::string want_filename = "foo/bar";
+  const size_t want_offset = 42;
+  const size_t want_n = 1024;
+  int calls = 0;
+  auto fetcher = [&calls, want_filename, want_offset, want_n](
+                     const std::string& got_filename, size_t got_offset,
+                     size_t got_n, char* buffer, TF_Status* status) -> int64_t {
+    EXPECT_EQ(got_filename, want_filename);
+    EXPECT_EQ(got_offset, want_offset);
+    EXPECT_EQ(got_n, want_n);
+    calls++;
+    memset(buffer, 'x', got_n);
+    TF_SetStatus(status, TF_OK, "");
+    return got_n;
+  };
+  // If block_size, max_bytes, or both are zero, or want_n is larger than
+  // max_bytes the cache is a pass-through.
+  tf_gcs_filesystem::RamFileBlockCache cache1(1, 0, 0, fetcher);
+  tf_gcs_filesystem::RamFileBlockCache cache2(0, 1, 0, fetcher);
+  tf_gcs_filesystem::RamFileBlockCache cache3(0, 0, 0, fetcher);
+  tf_gcs_filesystem::RamFileBlockCache cache4(1000, 1000, 0, fetcher);
+  std::vector<char> out;
+  TF_EXPECT_OK(ReadCache(&cache1, want_filename, want_offset, want_n, &out));
+  EXPECT_EQ(calls, 1);
+  TF_EXPECT_OK(ReadCache(&cache2, want_filename, want_offset, want_n, &out));
+  EXPECT_EQ(calls, 2);
+  TF_EXPECT_OK(ReadCache(&cache3, want_filename, want_offset, want_n, &out));
+  EXPECT_EQ(calls, 3);
+  TF_EXPECT_OK(ReadCache(&cache4, want_filename, want_offset, want_n, &out));
+  EXPECT_EQ(calls, 4);
+}
+
+TEST(RamFileBlockCacheTest, BlockAlignment) {
+  // Initialize a 256-byte buffer.  This is the file underlying the reads we'll
+  // do in this test.
+  const size_t size = 256;
+  std::vector<char> buf;
+  buf.reserve(size);
+  for (int i = 0; i < size; i++) {
+    buf.push_back(i);
+  }
+  // The fetcher just fetches slices of the buffer.
+  auto fetcher = [&buf](const std::string& filename, size_t offset, size_t n,
+                        char* buffer, TF_Status* status) -> int64_t {
+    int64_t bytes_transferred;
+    if (offset < buf.size()) {
+      size_t bytes_to_copy = std::min<size_t>(buf.size() - offset, n);
+      memcpy(buffer, buf.data() + offset, bytes_to_copy);
+      bytes_transferred = bytes_to_copy;
+    } else {
+      bytes_transferred = 0;
+    }
+    TF_SetStatus(status, TF_OK, "");
+    return bytes_transferred;
+  };
+  for (size_t block_size = 2; block_size <= 4; block_size++) {
+    // Make a cache of N-byte block size (1 block) and verify that reads of
+    // varying offsets and lengths return correct data.
+    tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
+                                               fetcher);
+    for (size_t offset = 0; offset < 10; offset++) {
+      for (size_t n = block_size - 2; n <= block_size + 2; n++) {
+        std::vector<char> got;
+        TF_EXPECT_OK(ReadCache(&cache, "", offset, n, &got));
+        // Verify the size of the read.
+        if (offset + n <= size) {
+          // Expect a full read.
+          EXPECT_EQ(got.size(), n) << "block size = " << block_size
+                                   << ", offset = " << offset << ", n = " << n;
+        } else {
+          // Expect a partial read.
+          EXPECT_EQ(got.size(), size - offset)
+              << "block size = " << block_size << ", offset = " << offset
+              << ", n = " << n;
+        }
+        // Verify the contents of the read.
+        std::vector<char>::const_iterator begin = buf.begin() + offset;
+        std::vector<char>::const_iterator end =
+            offset + n > buf.size() ? buf.end() : begin + n;
+        std::vector<char> want(begin, end);
+        EXPECT_EQ(got, want) << "block size = " << block_size
+                             << ", offset = " << offset << ", n = " << n;
+      }
+    }
+  }
+}
+
+TEST(RamFileBlockCacheTest, CacheHits) {
+  const size_t block_size = 16;
+  std::set<size_t> calls;
+  auto fetcher = [&calls, block_size](const std::string& filename,
+                                      size_t offset, size_t n, char* buffer,
+                                      TF_Status* status) -> int64_t {
+    EXPECT_EQ(n, block_size);
+    EXPECT_EQ(offset % block_size, 0);
+    EXPECT_EQ(calls.find(offset), calls.end()) << "at offset " << offset;
+    calls.insert(offset);
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  const uint32_t block_count = 256;
+  tf_gcs_filesystem::RamFileBlockCache cache(
+      block_size, block_count * block_size, 0, fetcher);
+  std::vector<char> out;
+  out.resize(block_count, 0);
+  // The cache has space for `block_count` blocks. The loop with i = 0 should
+  // fill the cache, and the loop with i = 1 should be all cache hits. The
+  // fetcher checks that it is called once and only once for each offset (to
+  // fetch the corresponding block).
+  for (int i = 0; i < 2; i++) {
+    for (int j = 0; j < block_count; j++) {
+      TF_EXPECT_OK(ReadCache(&cache, "", block_size * j, block_size, &out));
+    }
+  }
+}
+
+TEST(RamFileBlockCacheTest, OutOfRange) {
+  // Tests reads of a 24-byte file with block size 16.
+  const size_t block_size = 16;
+  const size_t file_size = 24;
+  bool first_block = false;
+  bool second_block = false;
+  auto fetcher = [block_size, file_size, &first_block, &second_block](
+                     const std::string& filename, size_t offset, size_t n,
+                     char* buffer, TF_Status* status) -> int64_t {
+    EXPECT_EQ(n, block_size);
+    EXPECT_EQ(offset % block_size, 0);
+    size_t bytes_to_copy = 0;
+    if (offset == 0) {
+      // The first block (16 bytes) of the file.
+      memset(buffer, 'x', n);
+      bytes_to_copy = n;
+      first_block = true;
+    } else if (offset == block_size) {
+      // The second block (8 bytes) of the file.
+      bytes_to_copy = file_size - block_size;
+      memset(buffer, 'x', bytes_to_copy);
+      second_block = true;
+    }
+    TF_SetStatus(status, TF_OK, "");
+    return bytes_to_copy;
+  };
+  tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
+                                             fetcher);
+  std::vector<char> out;
+  // Reading the first 16 bytes should be fine.
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, block_size, &out));
+  EXPECT_TRUE(first_block);
+  EXPECT_EQ(out.size(), block_size);
+  // Reading at offset file_size + 4 will read the second block (since the read
+  // at file_size + 4 = 28 will be aligned to an offset of 16) but will return
+  // OutOfRange because the offset is past the end of the 24-byte file.
+  absl::Status status = ReadCache(&cache, "", file_size + 4, 4, &out);
+  EXPECT_EQ(status.code(), error::OUT_OF_RANGE);
+  EXPECT_TRUE(second_block);
+  // Reading the second full block will return 8 bytes, from a cache hit.
+  second_block = false;
+  TF_EXPECT_OK(ReadCache(&cache, "", block_size, block_size, &out));
+  EXPECT_FALSE(second_block);
+  EXPECT_EQ(out.size(), file_size - block_size);
+}
+
+TEST(RamFileBlockCacheTest, Inconsistent) {
+  // Tests the detection of interrupted reads leading to partially filled blocks
+  // where we expected complete blocks.
+  const size_t block_size = 16;
+  // This fetcher returns OK but only fills in one byte for any offset.
+  auto fetcher = [block_size](const std::string& filename, size_t offset,
+                              size_t n, char* buffer,
+                              TF_Status* status) -> int64_t {
+    EXPECT_EQ(n, block_size);
+    EXPECT_EQ(offset % block_size, 0);
+    EXPECT_GE(n, 1);
+    memset(buffer, 'x', 1);
+    TF_SetStatus(status, TF_OK, "");
+    return 1;
+  };
+  tf_gcs_filesystem::RamFileBlockCache cache(block_size, 2 * block_size, 0,
+                                             fetcher);
+  std::vector<char> out;
+  // Read the second block; this should yield an OK status and a single byte.
+  TF_EXPECT_OK(ReadCache(&cache, "", block_size, block_size, &out));
+  EXPECT_EQ(out.size(), 1);
+  // Now read the first block; this should yield an INTERNAL error because we
+  // had already cached a partial block at a later position.
+  absl::Status status = ReadCache(&cache, "", 0, block_size, &out);
+  EXPECT_EQ(status.code(), error::INTERNAL);
+}
+
+TEST(RamFileBlockCacheTest, LRU) {
+  const size_t block_size = 16;
+  std::list<size_t> calls;
+  auto fetcher = [&calls, block_size](const std::string& filename,
+                                      size_t offset, size_t n, char* buffer,
+                                      TF_Status* status) -> int64_t {
+    EXPECT_EQ(n, block_size);
+    EXPECT_FALSE(calls.empty()) << "at offset = " << offset;
+    if (!calls.empty()) {
+      EXPECT_EQ(offset, calls.front());
+      calls.pop_front();
+    }
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  const uint32_t block_count = 2;
+  tf_gcs_filesystem::RamFileBlockCache cache(
+      block_size, block_count * block_size, 0, fetcher);
+  std::vector<char> out;
+  // Read blocks from the cache, and verify the LRU behavior based on the
+  // fetcher calls that the cache makes.
+  calls.push_back(0);
+  // Cache miss - drains an element from `calls`.
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
+  // Cache hit - does not drain an element from `calls`.
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
+  calls.push_back(block_size);
+  // Cache miss followed by cache hit.
+  TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
+  TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
+  calls.push_back(2 * block_size);
+  // Cache miss followed by cache hit.  Causes eviction of LRU element.
+  TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
+  TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
+  // LRU element was at offset 0.  Cache miss.
+  calls.push_back(0);
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
+  // Element at 2 * block_size is still in cache, and this read should update
+  // its position in the LRU list so it doesn't get evicted by the next read.
+  TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
+  // Element at block_size was evicted.  Reading this element will also cause
+  // the LRU element (at 0) to be evicted.
+  calls.push_back(block_size);
+  TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
+  // Element at 0 was evicted again.
+  calls.push_back(0);
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
+}
+
+TEST(RamFileBlockCacheTest, MaxStaleness) {
+  int calls = 0;
+  auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
+                          char* buffer, TF_Status* status) -> int64_t {
+    calls++;
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  std::vector<char> out;
+  std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
+  // Create a cache with max staleness of 2 seconds, and verify that it works as
+  // expected.
+  tf_gcs_filesystem::RamFileBlockCache cache1(
+      8, 16, 2 /* max staleness */, fetcher,
+      [&env]() { return env->NowSeconds(); });
+  // Execute the first read to load the block.
+  TF_EXPECT_OK(ReadCache(&cache1, "", 0, 1, &out));
+  EXPECT_EQ(calls, 1);
+  // Now advance the clock one second at a time and redo the read. The call
+  // count should advance every 3 seconds (i.e. every time the staleness is
+  // greater than 2).
+  for (int i = 1; i <= 10; i++) {
+    env->SetNowSeconds(i + 1);
+    TF_EXPECT_OK(ReadCache(&cache1, "", 0, 1, &out));
+    EXPECT_EQ(calls, 1 + i / 3);
+  }
+  // Now create a cache with max staleness of 0, and verify that it also works
+  // as expected.
+  calls = 0;
+  env->SetNowSeconds(0);
+  tf_gcs_filesystem::RamFileBlockCache cache2(
+      8, 16, 0 /* max staleness */, fetcher,
+      [&env]() { return env->NowSeconds(); });
+  // Execute the first read to load the block.
+  TF_EXPECT_OK(ReadCache(&cache2, "", 0, 1, &out));
+  EXPECT_EQ(calls, 1);
+  // Advance the clock by a huge amount and verify that the cached block is
+  // used to satisfy the read.
+  env->SetNowSeconds(365 * 24 * 60 * 60);  // ~1 year, just for fun.
+  TF_EXPECT_OK(ReadCache(&cache2, "", 0, 1, &out));
+  EXPECT_EQ(calls, 1);
+}
+
+TEST(RamFileBlockCacheTest, RemoveFile) {
+  int calls = 0;
+  auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
+                          char* buffer, TF_Status* status) -> int64_t {
+    calls++;
+    char c = (filename == "a") ? 'a' : (filename == "b") ? 'b' : 'x';
+    if (offset > 0) {
+      // The first block is lower case and all subsequent blocks are upper case.
+      c = absl::ascii_toupper(c);
+    }
+    memset(buffer, c, n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  // This cache has space for 4 blocks; we'll read from two files.
+  const size_t n = 3;
+  tf_gcs_filesystem::RamFileBlockCache cache(8, 32, 0, fetcher);
+  std::vector<char> out;
+  std::vector<char> a(n, 'a');
+  std::vector<char> b(n, 'b');
+  std::vector<char> A(n, 'A');
+  std::vector<char> B(n, 'B');
+  // Fill the cache.
+  TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
+  EXPECT_EQ(out, a);
+  EXPECT_EQ(calls, 1);
+  TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
+  EXPECT_EQ(out, A);
+  EXPECT_EQ(calls, 2);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
+  EXPECT_EQ(out, b);
+  EXPECT_EQ(calls, 3);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
+  EXPECT_EQ(out, B);
+  EXPECT_EQ(calls, 4);
+  // All four blocks should be in the cache now.
+  TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
+  EXPECT_EQ(out, a);
+  TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
+  EXPECT_EQ(out, A);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
+  EXPECT_EQ(out, b);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
+  EXPECT_EQ(out, B);
+  EXPECT_EQ(calls, 4);
+  // Remove the blocks from "a".
+  cache.RemoveFile("a");
+  // Both blocks from "b" should still be there.
+  TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
+  EXPECT_EQ(out, b);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
+  EXPECT_EQ(out, B);
+  EXPECT_EQ(calls, 4);
+  // The blocks from "a" should not be there.
+  TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
+  EXPECT_EQ(out, a);
+  EXPECT_EQ(calls, 5);
+  TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
+  EXPECT_EQ(out, A);
+  EXPECT_EQ(calls, 6);
+}
+
+TEST(RamFileBlockCacheTest, Prune) {
+  int calls = 0;
+  auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
+                          char* buffer, TF_Status* status) -> int64_t {
+    calls++;
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  std::vector<char> out;
+  // Our fake environment is initialized with the current timestamp.
+  std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
+  uint64_t now = Env::Default()->NowSeconds();
+  env->SetNowSeconds(now);
+  tf_gcs_filesystem::RamFileBlockCache cache(
+      8, 32, 1 /* max staleness */, fetcher,
+      [&env]() { return env->NowSeconds(); });
+  // Read three blocks into the cache, and advance the timestamp by one second
+  // with each read. Start with a block of "a" at the current timestamp `now`.
+  TF_EXPECT_OK(ReadCache(&cache, "a", 0, 1, &out));
+  // Now load a block of a different file "b" at timestamp `now` + 1
+  env->SetNowSeconds(now + 1);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
+  // Now load a different block of file "a" at timestamp `now` + 1. When the
+  // first block of "a" expires, this block should also be removed because it
+  // also belongs to file "a".
+  TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));
+  // Ensure that all blocks are in the cache (i.e. reads are cache hits).
+  EXPECT_EQ(cache.CacheSize(), 24);
+  EXPECT_EQ(calls, 3);
+  TF_EXPECT_OK(ReadCache(&cache, "a", 0, 1, &out));
+  TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
+  TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));
+  EXPECT_EQ(calls, 3);
+  // Advance the fake timestamp so that "a" becomes stale via its first block.
+  env->SetNowSeconds(now + 2);
+  // The pruning thread periodically compares env->NowSeconds() with the oldest
+  // block's timestamp to see if it should evict any files. At the current fake
+  // timestamp of `now` + 2, file "a" is stale because its first block is stale,
+  // but file "b" is not stale yet. Thus, once the pruning thread wakes up (in
+  // one second of wall time), it should remove "a" and leave "b" alone.
+  uint64_t start = Env::Default()->NowSeconds();
+  do {
+    Env::Default()->SleepForMicroseconds(100000);
+  } while (cache.CacheSize() == 24 && Env::Default()->NowSeconds() - start < 3);
+  // There should be one block left in the cache, and it should be the first
+  // block of "b".
+  EXPECT_EQ(cache.CacheSize(), 8);
+  TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
+  EXPECT_EQ(calls, 3);
+  // Advance the fake time to `now` + 3, at which point "b" becomes stale.
+  env->SetNowSeconds(now + 3);
+  // Wait for the pruner to remove "b".
+  start = Env::Default()->NowSeconds();
+  do {
+    Env::Default()->SleepForMicroseconds(100000);
+  } while (cache.CacheSize() == 8 && Env::Default()->NowSeconds() - start < 3);
+  // The cache should now be empty.
+  EXPECT_EQ(cache.CacheSize(), 0);
+}
+
+TEST(RamFileBlockCacheTest, ParallelReads) {
+  // This fetcher won't respond until either `callers` threads are calling it
+  // concurrently (at which point it will respond with success to all callers),
+  // or 10 seconds have elapsed (at which point it will respond with an error).
+  const int callers = 4;
+  absl::BlockingCounter counter(callers);
+  absl::Notification notification;
+  auto fetcher = [&counter, &notification](
+                     const std::string& filename, size_t offset, size_t n,
+                     char* buffer, TF_Status* status) -> int64_t {
+    if (counter.DecrementCount()) {
+      notification.Notify();
+      // This call to `Wait()` is not expected to block. Calling `Wait()` here
+      // allows us to satisfy `BlockingCounter`'s requirement: "When `Wait()`
+      // returns, it is legal to destroy the `BlockingCounter`.".
+      counter.Wait();
+    }
+    if (!notification.WaitForNotificationWithTimeout(absl::Seconds(10))) {
+      // This avoids having the test time out, which is harder to debug.
+      TF_SetStatus(status, TF_FAILED_PRECONDITION,
+                   "desired concurrency not reached");
+      return -1;
+    }
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  const int block_size = 8;
+  tf_gcs_filesystem::RamFileBlockCache cache(
+      block_size, 2 * callers * block_size, 0, fetcher);
+  std::vector<std::unique_ptr<Thread>> threads;
+  threads.reserve(callers);
+  for (int i = 0; i < callers; i++) {
+    threads.emplace_back(
+        Env::Default()->StartThread({}, "caller", [block_size, &cache, i]() {
+          std::vector<char> out;
+          TF_EXPECT_OK(
+              ReadCache(&cache, "a", i * block_size, block_size, &out));
+          std::vector<char> x(block_size, 'x');
+          EXPECT_EQ(out, x);
+        }));
+  }
+  // The `threads` destructor blocks until the threads can be joined, once their
+  // respective reads finish (which happens once they are all concurrently being
+  // executed, or 10 seconds have passed).
+}
+
+TEST(RamFileBlockCacheTest, CoalesceConcurrentReads) {
+  // Concurrent reads to the same file blocks should be de-duplicated.
+  const size_t block_size = 16;
+  int num_requests = 0;
+  absl::Notification notification;
+  auto fetcher = [&num_requests, &notification, block_size](
+                     const std::string& filename, size_t offset, size_t n,
+                     char* buffer, TF_Status* status) -> int64_t {
+    EXPECT_EQ(n, block_size);
+    EXPECT_EQ(offset, 0);
+    num_requests++;
+    memset(buffer, 'x', n);
+    notification.Notify();
+    // Wait for other thread to issue read.
+    Env::Default()->SleepForMicroseconds(100000);  // 0.1 secs
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
+                                             fetcher);
+  // Fork off thread for parallel read.
+  std::unique_ptr<Thread> concurrent(
+      Env::Default()->StartThread({}, "concurrent", [block_size, &cache] {
+        std::vector<char> out;
+        TF_EXPECT_OK(ReadCache(&cache, "", 0, block_size / 2, &out));
+        EXPECT_EQ(out.size(), block_size / 2);
+      }));
+  notification.WaitForNotification();
+  std::vector<char> out;
+  TF_EXPECT_OK(ReadCache(&cache, "", block_size / 2, block_size / 2, &out));
+  EXPECT_EQ(out.size(), block_size / 2);
+
+  EXPECT_EQ(1, num_requests);
+}
+
+TEST(RamFileBlockCacheTest, Flush) {
+  int calls = 0;
+  auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
+                          char* buffer, TF_Status* status) -> int64_t {
+    calls++;
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    return n;
+  };
+  tf_gcs_filesystem::RamFileBlockCache cache(16, 32, 0, fetcher);
+  std::vector<char> out;
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
+  EXPECT_EQ(calls, 1);
+  cache.Flush();
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
+  EXPECT_EQ(calls, 2);
+}
+
+}  // namespace
+}  // namespace tensorflow