Clarification of object ownership in ::tensorflow::Env

Change: 119544956

2 files changed, 31 insertions, 12 deletions

diff --git a/tensorflow/core/platform/env.cc b/tensorflow/core/platform/env.cc
index 3e0c51d599..714b4511f8 100644
--- a/tensorflow/core/platform/env.cc
+++ b/tensorflow/core/platform/env.cc
@@ -29,30 +29,32 @@ class FileSystemRegistryImpl : public FileSystemRegistry {
 
  private:
   mutable mutex mu_;
-  mutable std::unordered_map<string, FileSystem*> registry_ GUARDED_BY(mu_);
+  mutable std::unordered_map<string, std::unique_ptr<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();
+  QCHECK(
+      registry_.emplace(string(scheme), std::unique_ptr<FileSystem>(factory()))
+          .second)
+      << "File factory for " << scheme << " already registered";
 }
 
 FileSystem* FileSystemRegistryImpl::Lookup(const string& scheme) {
   mutex_lock lock(mu_);
-  auto fs_ptr = gtl::FindOrNull(registry_, scheme);
-  if (!fs_ptr) {
+  const auto found = registry_.find(scheme);
+  if (found == registry_.end()) {
     return nullptr;
   }
-  return *fs_ptr;
+  return found->second.get();
 }
 
 Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes(
     std::vector<string>* schemes) {
   mutex_lock lock(mu_);
-  for (auto const e : registry_) {
+  for (const auto& e : registry_) {
     schemes->push_back(e.first);
   }
   return Status::OK();
@@ -60,8 +62,6 @@ Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes(
 
 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 = file_system_registry_->Lookup(scheme);
diff --git a/tensorflow/core/platform/env.h b/tensorflow/core/platform/env.h
index 6527da97a9..1abe5cd2c0 100644
--- a/tensorflow/core/platform/env.h
+++ b/tensorflow/core/platform/env.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define TENSORFLOW_CORE_PLATFORM_ENV_H_
 
 #include <stdint.h>
+#include <memory>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -45,7 +46,7 @@ struct ThreadOptions;
 class Env {
  public:
   Env();
-  virtual ~Env();
+  virtual ~Env() = default;
 
   /// \brief Returns a default environment suitable for the current operating
   /// system.
@@ -59,6 +60,8 @@ class Env {
   /// \brief Returns the FileSystem object to handle operations on the file
   /// specified by 'fname'. The FileSystem object is used as the implementation
   /// for the file system related (non-virtual) functions that follow.
+  /// Returned FileSystem object is still owned by the Env object and will
+  // (might) be destroyed when the environment is destroyed.
   virtual Status GetFileSystemForFile(const string& fname, FileSystem** result);
 
   /// \brief Returns the file system schemes registered for this Env.
@@ -77,6 +80,10 @@ class Env {
   /// status.
   ///
   /// The returned file may be concurrently accessed by multiple threads.
+  ///
+  /// The ownership of the returned RandomAccessFile is passed to the caller
+  /// and the object should be deleted when is not used. The file object
+  /// shouldn't live longer than the Env object.
   Status NewRandomAccessFile(const string& fname, RandomAccessFile** result);
 
   /// \brief Creates an object that writes to a new file with the specified
@@ -88,6 +95,10 @@ class Env {
   /// returns non-OK.
   ///
   /// The returned file will only be accessed by one thread at a time.
+  ///
+  /// The ownership of the returned WritableFile is passed to the caller
+  /// and the object should be deleted when is not used. The file object
+  /// shouldn't live longer than the Env object.
   Status NewWritableFile(const string& fname, WritableFile** result);
 
   /// \brief Creates an object that either appends to an existing file, or
@@ -98,6 +109,10 @@ class Env {
   /// non-OK.
   ///
   /// The returned file will only be accessed by one thread at a time.
+  ///
+  /// The ownership of the returned WritableFile is passed to the caller
+  /// and the object should be deleted when is not used. The file object
+  /// shouldn't live longer than the Env object.
   Status NewAppendableFile(const string& fname, WritableFile** result);
 
   /// \brief Creates a readonly region of memory with the file context.
@@ -107,6 +122,10 @@ class Env {
   /// the caller. On failure stores nullptr in *result and returns non-OK.
   ///
   /// The returned memory region can be accessed from many threads in parallel.
+  ///
+  /// The ownership of the returned ReadOnlyMemoryRegion is passed to the caller
+  /// and the object should be deleted when is not used. The memory region
+  /// object shouldn't live longer than the Env object.
   Status NewReadOnlyMemoryRegionFromFile(const string& fname,
                                          ReadOnlyMemoryRegion** result);
 
@@ -192,7 +211,7 @@ class Env {
   Env(const Env&);
   void operator=(const Env&);
 
-  FileSystemRegistry* file_system_registry_;
+  std::unique_ptr<FileSystemRegistry> file_system_registry_;
 };
 
 /// \brief An implementation of Env that forwards all calls to another Env.