1 files changed, 77 insertions, 13 deletions

diff --git a/tensorflow/stream_executor/stream_test.cc b/tensorflow/stream_executor/stream_test.cc
index 47dd675834..cfc051fd09 100644
--- a/tensorflow/stream_executor/stream_test.cc
+++ b/tensorflow/stream_executor/stream_test.cc
@@ -95,18 +95,18 @@ TEST_F(StreamTest, TwoSubStreams) {
   EXPECT_NE(sub_stream3, sub_stream4);
 }
 
-TEST_F(StreamTest, FailedSubStreamNotReused) {
+TEST_F(StreamTest, FailedSubStreamBeforeReturnNotReused) {
   std::unique_ptr<StreamExecutor> executor = NewStreamExecutor();
   Stream stream(executor.get());
   stream.Init();
   EXPECT_TRUE(stream.ok());
 
-  // Get a sub-stream.
+  // Get sub_stream1.
   Stream* sub_stream1 = stream.GetOrCreateSubStream();
   EXPECT_TRUE(sub_stream1->ok());
 
-  // Force an error on the stream; here we call a method that requires
-  // DNN support, which we know the Host platform doesn't support.
+  // Force an error on sub_stream1; here we call a method that requires DNN
+  // support, which we know the Host platform doesn't support.
   sub_stream1->ThenDepthConcatenate({}, {}, nullptr);
   EXPECT_FALSE(sub_stream1->ok());
 
@@ -115,20 +115,84 @@ TEST_F(StreamTest, FailedSubStreamNotReused) {
   Stream* sub_stream2 = stream.GetOrCreateSubStream();
   EXPECT_TRUE(sub_stream2->ok());
 
-  // The underlying streams should be different. They would have been
-  // the same, but since we forced an error on sub_stream1, it will
-  // not be re-used. Sadly we can't just check:
+  // The underlying sub_streams should be different. They would have been the
+  // same, but since we forced an error on sub_stream1, it will not be
+  // re-used. Sadly we can't just check:
   //   EXPECT_NE(sub_stream1, sub_stream2);
   //
-  // The above should hold logically, but it may fail if the new
-  // stream instance allocated for sub_stream2 happens to reside in
-  // the same memory address as sub_stream1.
+  // The above should hold logically, but it may fail if the new Stream instance
+  // allocated for sub_stream2 happens to reside in the same memory address as
+  // sub_stream1.
   //
   // The check that sub_stream2->ok() serves as a good-enough check.
 
-  // Return sub_stream2 and get sub_stream3. The previous error on
-  // sub_stream1 has no effect on these streams, and they are the
-  // same.
+  // Return sub_stream2 and get sub_stream3. The previous error on sub_stream1
+  // has no effect on these streams, and they are the same.
+  stream.ReturnSubStream(sub_stream2);
+  Stream* sub_stream3 = stream.GetOrCreateSubStream();
+  EXPECT_TRUE(sub_stream3->ok());
+  EXPECT_EQ(sub_stream2, sub_stream3);
+}
+
+TEST_F(StreamTest, FailedSubStreamAfterReturnNotReused) {
+  std::unique_ptr<StreamExecutor> executor = NewStreamExecutor();
+  Stream stream(executor.get());
+  stream.Init();
+  EXPECT_TRUE(stream.ok());
+
+  // Get and return sub_stream1.
+  Stream* sub_stream1 = stream.GetOrCreateSubStream();
+  EXPECT_TRUE(sub_stream1->ok());
+  stream.ReturnSubStream(sub_stream1);
+
+  // Force an error on sub_stream1; here we call a method that requires DNN
+  // support, which we know the Host platform doesn't support.
+  //
+  // It is a bit weird to use sub_stream1 after it has already been returned. By
+  // doing this, we're simulating an asynchronous error that occurs during
+  // execution of the sub_stream, that occurs after the sub_stream is returned.
+  //
+  // E.g. the following is a common pattern of usage, where the execution of the
+  // operations enqueued onto the sub streams may occur after the streams have
+  // already been returned.
+  //
+  //   void EnqueueOnSubStreams(Stream* stream) {
+  //     Stream* sub_stream1 = stream.GetOrCreateSubStream();
+  //     Stream* sub_stream2 = stream.GetOrCreateSubStream();
+  //     // ... enqueue some operations on the sub streams ...
+  //     stream.ThenWaitFor(sub_stream1).ThenWaitFor(sub_stream2);
+  //     stream.ReturnSubStream(sub_stream1);
+  //     stream.ReturnSubStream(sub_stream2);
+  //   }
+  //
+  //   Stream* main_stream = ...;
+  //   EnqueueOnSubStreams(main_stream);
+  //   main_stream.BlockHostUntilDone();
+  //
+  // TODO(b/112196569): The semantics of failed sub-streams is error-prone;
+  // GetOrCreateSubStream can still return a sub-stream that has not encountered
+  // an error yet, but will encounter one in the future, based on previously
+  // enqueued operations.
+  sub_stream1->ThenDepthConcatenate({}, {}, nullptr);
+  EXPECT_FALSE(sub_stream1->ok());
+
+  // Get and return sub_stream2.
+  Stream* sub_stream2 = stream.GetOrCreateSubStream();
+  EXPECT_TRUE(sub_stream2->ok());
+
+  // The underlying streams should be different. They would have been the same,
+  // but since we forced an error on sub_stream1, it will not be re-used. Sadly
+  // we can't just check:
+  //   EXPECT_NE(sub_stream1, sub_stream2);
+  //
+  // The above should hold logically, but it may fail if the new stream instance
+  // allocated for sub_stream2 happens to reside in the same memory address as
+  // sub_stream1.
+  //
+  // The check that sub_stream2->ok() serves as a good-enough check.
+
+  // Return sub_stream2 and get sub_stream3. The previous error on sub_stream1
+  // has no effect on these streams, and they are the same.
   stream.ReturnSubStream(sub_stream2);
   Stream* sub_stream3 = stream.GetOrCreateSubStream();
   EXPECT_TRUE(sub_stream3->ok());