Add BlockHostUntilDoneWithStatus, which returns Status rather than bool.

Also fixed a deadlock in Stream::BlockHostUntilDone.  The problem with the
original code was that it grabbed mu_ before looping over substreams, and would
call CheckError with mu_ still held.  But CheckError will attempt to lock mu_ in
the failure case, which would deadlock.

PiperOrigin-RevId: 178191634

1 files changed, 15 insertions, 8 deletions

diff --git a/tensorflow/stream_executor/stream.cc b/tensorflow/stream_executor/stream.cc
index 22fd6bce78..de65038d17 100644
--- a/tensorflow/stream_executor/stream.cc
+++ b/tensorflow/stream_executor/stream.cc
@@ -5055,22 +5055,24 @@ Stream &Stream::ThenEnqueueOnBackgroundThread(
   });
 }
 
-bool Stream::BlockHostUntilDone() {
+port::Status Stream::BlockHostUntilDoneWithStatus() {
   VLOG_CALL();
 
   if (!ok()) {
-    LOG(INFO)
-        << "stream " << this
-        << " did not block host until done; was already in an error state";
-    return false;
+    port::Status status = port::Status(
+        port::error::INTERNAL,
+        "stream did not block host until done; was already in an error state");
+    LOG(INFO) << status << " " << this;
+    return status;
   }
 
+  port::Status first_error;
   {
     // Wait until all active sub-streams have done their tasks.
     mutex_lock lock{mu_};
     for (auto &stream : sub_streams_) {
       if (!stream.second) {
-        CheckError(stream.first->BlockHostUntilDone());
+        first_error.Update(stream.first->BlockHostUntilDoneWithStatus());
         // Set this sub-stream as available.
         stream.second = true;
       }
@@ -5079,8 +5081,13 @@ bool Stream::BlockHostUntilDone() {
 
   temporary_memory_manager_.DeallocateFinalizedTemporaries();
 
-  CheckError(parent_->BlockHostUntilDone(this));
-  return ok();
+  first_error.Update(parent_->BlockHostUntilDoneWithStatus(this));
+  CheckError(first_error.ok());
+  return first_error;
+}
+
+bool Stream::BlockHostUntilDone() {
+  return BlockHostUntilDoneWithStatus().ok();
 }
 
 }  // namespace gputools