From f7464a6c05421c50cf3ed7645c724832482063d9 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 20 Sep 2019 15:13:50 -0700 Subject: [PATCH] Write metadata to a temp file before atomically renaming it. PiperOrigin-RevId: 270360252 --- .../kernels/data/experimental/snapshot_dataset_op.cc | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/data/experimental/snapshot_dataset_op.cc b/tensorflow/core/kernels/data/experimental/snapshot_dataset_op.cc index 336c602dfca..d10bfe2881f 100644 --- a/tensorflow/core/kernels/data/experimental/snapshot_dataset_op.cc +++ b/tensorflow/core/kernels/data/experimental/snapshot_dataset_op.cc @@ -224,18 +224,19 @@ Status WriteMetadataFile(const string& hash_dir, string metadata_filename = absl::StrCat(hash_dir, "/", kSnapshotFilename); TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(hash_dir)); - Status exists = Env::Default()->FileExists(metadata_filename); - if (exists.ok()) { - TF_RETURN_IF_ERROR(Env::Default()->DeleteFile(metadata_filename)); - } + std::string tmp_filename = + absl::StrCat(metadata_filename, "-tmp-", random::New64()); std::unique_ptr file; - TF_RETURN_IF_ERROR(Env::Default()->NewWritableFile(metadata_filename, &file)); + TF_RETURN_IF_ERROR(Env::Default()->NewWritableFile(tmp_filename, &file)); auto writer = absl::make_unique(file.get()); TF_RETURN_IF_ERROR(writer->WriteRecord(metadata.SerializeAsString())); TF_RETURN_IF_ERROR(writer->Close()); + TF_RETURN_IF_ERROR( + Env::Default()->RenameFile(tmp_filename, metadata_filename)); + return Status::OK(); }