- Makes SingleMachine::Run(...) thread-safe.

Change: 150289918

1 files changed, 21 insertions, 0 deletions

diff --git a/tensorflow/core/grappler/utils_test.cc b/tensorflow/core/grappler/utils_test.cc
index f5d6877637..766cedd65d 100644
--- a/tensorflow/core/grappler/utils_test.cc
+++ b/tensorflow/core/grappler/utils_test.cc
@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
@@ -60,6 +63,24 @@ TEST_F(UtilsTest, AddNodeNamePrefix) {
   EXPECT_EQ("OPTIMIZED-", AddPrefixToNodeName("", "OPTIMIZED"));
 }
 
+TEST_F(UtilsTest, ExecuteWithTimeout) {
+  std::unique_ptr<thread::ThreadPool> thread_pool(
+      new thread::ThreadPool(Env::Default(), "ExecuteWithTimeout", 2));
+  ASSERT_TRUE(ExecuteWithTimeout(
+      []() {  // Do nothing.
+      },
+      1 /* timeout_in_ms */, thread_pool.get()));
+  // This should time out.
+  ASSERT_FALSE(ExecuteWithTimeout([]() { sleep(1); }, 1 /* timeout_in_ms */,
+                                  thread_pool.get()));
+  // This should run till the end.
+  ASSERT_TRUE(ExecuteWithTimeout([]() { sleep(1); }, 0 /* timeout_in_ms */,
+                                 thread_pool.get()));
+
+  // Deleting before local variables go off the stack.
+  thread_pool.reset();
+}
+
 }  // namespace
 }  // namespace grappler
 }  // namespace tensorflow