Move the FileSystemRegistry inside Env instead of it being global and shared

by all Envs.
Change: 119342847

tensorflow/core/platform/env.cc

@@ -21,11 +21,50 @@ limitations under the License.
 
 namespace tensorflow {
 
-Env::~Env() {}
+class FileSystemRegistryImpl : public FileSystemRegistry {
+ public:
+  void Register(const string& scheme, Factory factory) override;
+  FileSystem* Lookup(const string& scheme) override;
+  Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes) override;
+
+ private:
+  mutable mutex mu_;
+  mutable std::unordered_map<string, FileSystem*> registry_ GUARDED_BY(mu_);
+};
+
+void FileSystemRegistryImpl::Register(const string& scheme,
+                                      FileSystemRegistry::Factory factory) {
+  mutex_lock lock(mu_);
+  QCHECK(!gtl::FindOrNull(registry_, scheme)) << "File factory for " << scheme
+                                              << " already registered";
+  registry_[scheme] = factory();
+}
+
+FileSystem* FileSystemRegistryImpl::Lookup(const string& scheme) {
+  mutex_lock lock(mu_);
+  auto fs_ptr = gtl::FindOrNull(registry_, scheme);
+  if (!fs_ptr) {
+    return nullptr;
+  }
+  return *fs_ptr;
+}
+
+Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes(
+    std::vector<string>* schemes) {
+  mutex_lock lock(mu_);
+  for (auto const e : registry_) {
+    schemes->push_back(e.first);
+  }
+  return Status::OK();
+}
+
+Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {}
+
+Env::~Env() { delete file_system_registry_; }
 
 Status Env::GetFileSystemForFile(const string& fname, FileSystem** result) {
   string scheme = GetSchemeFromURI(fname);
-  FileSystem* file_system = GlobalFileSystemRegistry()->Lookup(scheme);
+  FileSystem* file_system = file_system_registry_->Lookup(scheme);
   if (!file_system) {
     return errors::Unimplemented("File system scheme ", scheme,
                                  " not implemented");
@@ -34,6 +73,15 @@ Status Env::GetFileSystemForFile(const string& fname, FileSystem** result) {
   return Status::OK();
 }
 
+Status Env::GetRegisteredFileSystemSchemes(std::vector<string>* schemes) {
+  return file_system_registry_->GetRegisteredFileSystemSchemes(schemes);
+}
+
+void Env::RegisterFileSystem(const string& scheme,
+                             FileSystemRegistry::Factory factory) {
+  file_system_registry_->Register(scheme, factory);
+}
+
 Status Env::NewRandomAccessFile(const string& fname,
                                 RandomAccessFile** result) {
   FileSystem* fs;

tensorflow/core/platform/env.h

@@ -44,7 +44,7 @@ struct ThreadOptions;
 /// multiple threads without any external synchronization.
 class Env {
  public:
-  Env() {}
+  Env();
   virtual ~Env();
 
   /// \brief Returns a default environment suitable for the current operating
@@ -61,6 +61,13 @@ class Env {
   /// for the file system related (non-virtual) functions that follow.
   virtual Status GetFileSystemForFile(const string& fname, FileSystem** result);
 
+  /// \brief Returns the file system schemes registered for this Env.
+  virtual Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes);
+
+  // \brief Register a file system for a scheme.
+  virtual void RegisterFileSystem(const string& scheme,
+                                  FileSystemRegistry::Factory factory);
+
   /// \brief Creates a brand new random access read-only file with the
   /// specified name.
 
@@ -184,6 +191,8 @@ class Env {
   /// No copying allowed
   Env(const Env&);
   void operator=(const Env&);
+
+  FileSystemRegistry* file_system_registry_;
 };
 
 /// \brief An implementation of Env that forwards all calls to another Env.
@@ -204,6 +213,15 @@ class EnvWrapper : public Env {
     return target_->GetFileSystemForFile(fname, result);
   }
 
+  Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes) override {
+    return target_->GetRegisteredFileSystemSchemes(schemes);
+  }
+
+  void RegisterFileSystem(const string& scheme,
+                          FileSystemRegistry::Factory factory) override {
+    target_->RegisterFileSystem(scheme, factory);
+  }
+
   uint64 NowMicros() override { return target_->NowMicros(); }
   void SleepForMicroseconds(int micros) override {
     target_->SleepForMicroseconds(micros);
@@ -267,6 +285,32 @@ Status WriteStringToFile(Env* env, const string& fname,
 Status ReadBinaryProto(Env* env, const string& fname,
                        ::tensorflow::protobuf::MessageLite* proto);
 
+namespace register_file_system {
+
+template <typename Factory>
+struct Register {
+  Register(Env* env, const string& scheme) {
+    env->RegisterFileSystem(scheme,
+                            []() -> FileSystem* { return new Factory; });
+  }
+};
+
+}  // namespace register_file_system
+
 }  // namespace tensorflow
 
+// Register a FileSystem implementation for a scheme. Files with names that have
+// "scheme://" prefixes are routed to use this implementation.
+#define REGISTER_FILE_SYSTEM_ENV(env, scheme, factory) \
+  REGISTER_FILE_SYSTEM_UNIQ_HELPER(__COUNTER__, env, scheme, factory)
+#define REGISTER_FILE_SYSTEM_UNIQ_HELPER(ctr, env, scheme, factory) \
+  REGISTER_FILE_SYSTEM_UNIQ(ctr, env, scheme, factory)
+#define REGISTER_FILE_SYSTEM_UNIQ(ctr, env, scheme, factory)   \
+  static ::tensorflow::register_file_system::Register<factory> \
+      register_ff##ctr TF_ATTRIBUTE_UNUSED =                   \
+          ::tensorflow::register_file_system::Register<factory>(env, scheme)
+
+#define REGISTER_FILE_SYSTEM(scheme, factory) \
+  REGISTER_FILE_SYSTEM_ENV(Env::Default(), scheme, factory);
+
 #endif  // TENSORFLOW_CORE_PLATFORM_ENV_H_

tensorflow/core/platform/file_system.cc

@@ -33,38 +33,6 @@ WritableFile::~WritableFile() {}
 
 FileSystemRegistry::~FileSystemRegistry() {}
 
-class FileSystemRegistryImpl : public FileSystemRegistry {
- public:
-  void Register(const string& scheme, Factory factory) override;
-  FileSystem* Lookup(const string& scheme) override;
-
- private:
-  mutable mutex mu_;
-  mutable std::unordered_map<string, FileSystem*> registry_ GUARDED_BY(mu_);
-};
-
-FileSystemRegistry* GlobalFileSystemRegistry() {
-  static FileSystemRegistry* registry = new FileSystemRegistryImpl;
-  return registry;
-}
-
-void FileSystemRegistryImpl::Register(const string& scheme,
-                                      FileSystemRegistry::Factory factory) {
-  mutex_lock lock(mu_);
-  QCHECK(!gtl::FindOrNull(registry_, scheme)) << "File factory for " << scheme
-                                              << " already registered";
-  registry_[scheme] = factory();
-}
-
-FileSystem* FileSystemRegistryImpl::Lookup(const string& scheme) {
-  mutex_lock lock(mu_);
-  auto fs_ptr = gtl::FindOrNull(registry_, scheme);
-  if (!fs_ptr) {
-    return nullptr;
-  }
-  return *fs_ptr;
-}
-
 string GetSchemeFromURI(const string& name) {
   auto colon_loc = name.find(":");
   // Make sure scheme matches [a-zA-Z][0-9a-zA-Z.]*

tensorflow/core/platform/file_system.h

@@ -205,22 +205,10 @@ class FileSystemRegistry {
   virtual ~FileSystemRegistry();
   virtual void Register(const string& scheme, Factory factory) = 0;
   virtual FileSystem* Lookup(const string& scheme) = 0;
+  virtual Status GetRegisteredFileSystemSchemes(
+      std::vector<string>* schemes) = 0;
 };
 
-FileSystemRegistry* GlobalFileSystemRegistry();
-
-namespace register_file_system {
-
-template <typename Factory>
-struct Register {
-  Register(const string& scheme) {
-    ::tensorflow::GlobalFileSystemRegistry()->Register(
-        scheme, []() -> FileSystem* { return new Factory; });
-  }
-};
-
-}  // namespace register_file_system
-
 // Given URI of the form [scheme://]<filename>, return 'scheme'.
 string GetSchemeFromURI(const string& name);
 
@@ -229,15 +217,4 @@ string GetNameFromURI(const string& name);
 
 }  // namespace tensorflow
 
-// Register a FileSystem implementation for a scheme. Files with names that have
-// "scheme://" prefixes are routed to use this implementation.
-#define REGISTER_FILE_SYSTEM(scheme, factory) \
-  REGISTER_FILE_SYSTEM_UNIQ_HELPER(__COUNTER__, scheme, factory)
-#define REGISTER_FILE_SYSTEM_UNIQ_HELPER(ctr, scheme, factory) \
-  REGISTER_FILE_SYSTEM_UNIQ(ctr, scheme, factory)
-#define REGISTER_FILE_SYSTEM_UNIQ(ctr, scheme, factory)        \
-  static ::tensorflow::register_file_system::Register<factory> \
-      register_ff##ctr TF_ATTRIBUTE_UNUSED =                   \
-          ::tensorflow::register_file_system::Register<factory>(scheme)
-
 #endif  // TENSORFLOW_CORE_PLATFORM_FILE_SYSTEM_H_

tensorflow/python/framework/test_file_system.cc

@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/env.h"
 
 namespace tensorflow {