Fixes a data race condition in tfdbg that could lead to using more

than the storage limit.

PiperOrigin-RevId: 210930360

2 files changed, 11 insertions, 1 deletions

diff --git a/tensorflow/core/debug/debug_io_utils.cc b/tensorflow/core/debug/debug_io_utils.cc
index 09c2b58168..38863db1cc 100644
--- a/tensorflow/core/debug/debug_io_utils.cc
+++ b/tensorflow/core/debug/debug_io_utils.cc
@@ -690,6 +690,8 @@ const uint64 DebugFileIO::defaultGlobalDiskBytesLimit = 107374182400L;
 uint64 DebugFileIO::globalDiskBytesLimit = 0;
 uint64 DebugFileIO::diskBytesUsed = 0;
 
+mutex DebugFileIO::bytes_mu(LINKER_INITIALIZED);
+
 bool DebugFileIO::requestDiskByteUsage(uint64 bytes) {
   if (globalDiskBytesLimit == 0) {
     const char* env_tfdbg_disk_bytes_limit = getenv("TFDBG_DISK_BYTES_LIMIT");
@@ -705,6 +707,7 @@ bool DebugFileIO::requestDiskByteUsage(uint64 bytes) {
   if (bytes == 0) {
     return true;
   }
+  mutex_lock l(bytes_mu);
   if (diskBytesUsed + bytes < globalDiskBytesLimit) {
     diskBytesUsed += bytes;
     return true;
@@ -713,7 +716,10 @@ bool DebugFileIO::requestDiskByteUsage(uint64 bytes) {
   }
 }
 
-void DebugFileIO::resetDiskByteUsage() { diskBytesUsed = 0; }
+void DebugFileIO::resetDiskByteUsage() {
+  mutex_lock l(bytes_mu);
+  diskBytesUsed = 0;
+}
 
 #ifndef PLATFORM_WINDOWS
 DebugGrpcChannel::DebugGrpcChannel(const string& server_stream_addr)
diff --git a/tensorflow/core/debug/debug_io_utils.h b/tensorflow/core/debug/debug_io_utils.h
index 56f8b74e18..5390ce408a 100644
--- a/tensorflow/core/debug/debug_io_utils.h
+++ b/tensorflow/core/debug/debug_io_utils.h
@@ -225,7 +225,11 @@ class DebugFileIO {
   // fixed.
   static Status RecursiveCreateDir(Env* env, const string& dir);
 
+  // Tracks how much disk has been used so far.
   static uint64 diskBytesUsed;
+  // Mutex for thread-safe access to diskBytesUsed.
+  static mutex bytes_mu;
+  // Default limit for the disk space.
   static const uint64 defaultGlobalDiskBytesLimit;
 
   friend class DiskUsageLimitTest;