From 9130ba73e6bc420100b1dfb1bae2adade3c4e767 Mon Sep 17 00:00:00 2001
From: Vijay Vasudevan <vrv@google.com>
Date: Fri, 12 Jun 2020 13:57:03 -0700
Subject: [PATCH] GCS File system: Remove tempfile in destructor regardless of
 status.

Otherwise this implementation leaves tempfiles around on failures, filling
up the temp filesystem: the Sync() call fails and Close() returns before the tempfile can be removed on destruction.

By removing the tempfile in the destructor, this mirrors the creation of the
tempfile in the constructor, and ensures it only happens once.

PiperOrigin-RevId: 316171772
Change-Id: I7a6ab9c417cd32b0e1fb46499b80f23bb476b6b6
---
 .../core/platform/cloud/gcs_file_system.cc    | 13 ++++--
 .../platform/cloud/gcs_file_system_test.cc    | 43 +++++++++++++------
 2 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index 92210498b01..5d395f3d821 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -416,7 +416,10 @@ class GcsWritableFile : public WritableFile {
                   std::ofstream::binary | std::ofstream::app);
   }
 
-  ~GcsWritableFile() override { Close().IgnoreError(); }
+  ~GcsWritableFile() override {
+    Close().IgnoreError();
+    std::remove(tmp_content_filename_.c_str());
+  }
 
   Status Append(StringPiece data) override {
     TF_RETURN_IF_ERROR(CheckWritable());
@@ -431,9 +434,11 @@ class GcsWritableFile : public WritableFile {
 
   Status Close() override {
     if (outfile_.is_open()) {
-      TF_RETURN_IF_ERROR(Sync());
-      outfile_.close();
-      std::remove(tmp_content_filename_.c_str());
+      Status sync_status = Sync();
+      if (sync_status.ok()) {
+        outfile_.close();
+      }
+      return sync_status;
     }
     return Status::OK();
   }
diff --git a/tensorflow/core/platform/cloud/gcs_file_system_test.cc b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
index 14af9f979e6..544ddc32043 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system_test.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
@@ -1225,6 +1225,11 @@ TEST(GcsFileSystemTest, NewWritableFile_ResumeUploadAllAttemptsFail) {
 }
 
 TEST(GcsFileSystemTest, NewWritableFile_UploadReturns410) {
+  std::vector<string> results;
+  TF_EXPECT_OK(
+      Env::Default()->GetMatchingPaths("/tmp/tmp_file_tensorflow*", &results));
+  const int64 tmp_files_before = results.size();
+
   std::vector<HttpRequest*> requests(
       {new FakeHttpRequest(
            "Uri: https://www.googleapis.com/upload/storage/v1/b/bucket/o?"
@@ -1268,21 +1273,31 @@ TEST(GcsFileSystemTest, NewWritableFile_UploadReturns410) {
                    kTestTimeoutConfig, *kAllowedLocationsDefault,
                    nullptr /* gcs additional header */);
 
-  std::unique_ptr<WritableFile> file;
-  TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
+  {
+    std::unique_ptr<WritableFile> file;
+    TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
 
-  TF_EXPECT_OK(file->Append("content1,"));
-  TF_EXPECT_OK(file->Append("content2"));
-  const auto& status = file->Close();
-  EXPECT_EQ(errors::Code::UNAVAILABLE, status.code());
-  EXPECT_TRUE(
-      absl::StrContains(status.error_message(),
-                        "Upload to gs://bucket/path/writeable.txt failed, "
-                        "caused by: Not found: important HTTP error 410"))
-      << status;
-  EXPECT_TRUE(absl::StrContains(
-      status.error_message(), "when uploading gs://bucket/path/writeable.txt"))
-      << status;
+    TF_EXPECT_OK(file->Append("content1,"));
+    TF_EXPECT_OK(file->Append("content2"));
+    const auto& status = file->Close();
+    EXPECT_EQ(errors::Code::UNAVAILABLE, status.code());
+    EXPECT_TRUE(
+        absl::StrContains(status.error_message(),
+                          "Upload to gs://bucket/path/writeable.txt failed, "
+                          "caused by: Not found: important HTTP error 410"))
+        << status;
+    EXPECT_TRUE(
+        absl::StrContains(status.error_message(),
+                          "when uploading gs://bucket/path/writeable.txt"))
+        << status;
+  }
+
+  // Check that no new tempfiles were left over after failure and destruction
+  // of the file.
+  results.clear();
+  TF_EXPECT_OK(
+      Env::Default()->GetMatchingPaths("/tmp/tmp_file_tensorflow*", &results));
+  EXPECT_EQ(tmp_files_before, results.size());
 }
 
 TEST(GcsFileSystemTest, NewWritableFile_NoObjectName) {