Adds a C++ test utility that makes it possible to run multi-process tests.

This will be necessary for tests the distributed runtime (issue #23).
Change: 115339579

3 files changed, 113 insertions, 1 deletions

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 0ec4878990..a38e32f410 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -63,6 +63,7 @@ load(
     "tf_proto_library",
     "tf_additional_lib_srcs",
     "tf_additional_stream_executor_srcs",
+    "tf_additional_test_deps",
     "tf_additional_test_srcs",
     "tf_kernel_tests_linkstatic",
 )
@@ -157,7 +158,7 @@ cc_library(
         ":lib",
         ":lib_internal",
         "//tensorflow/core/platform/default/build_config:gtest",
-    ],
+    ] + tf_additional_test_deps(),
 )
 
 tf_cuda_library(
diff --git a/tensorflow/core/platform/posix/test.cc b/tensorflow/core/platform/posix/test.cc
new file mode 100644
index 0000000000..2fa2d36ceb
--- /dev/null
+++ b/tensorflow/core/platform/posix/test.cc
@@ -0,0 +1,79 @@
+/* Copyright 2016 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/platform/test.h"
+
+#include <signal.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+namespace tensorflow {
+namespace testing {
+
+namespace {
+class PosixSubProcess : public SubProcess {
+ public:
+  PosixSubProcess(const std::vector<string>& argv) : argv_(argv), pid_(0) {}
+
+  ~PosixSubProcess() override {}
+
+  bool Start() override {
+    if (pid_ != 0) {
+      LOG(ERROR) << "Tried to start process multiple times.";
+      return false;
+    }
+    pid_ = fork();
+    if (pid_ == 0) {
+      // We are in the child process.
+      const char* path = argv_[0].c_str();
+      const char** argv = new const char*[argv_.size()];
+      for (int i = 1; i < argv_.size(); ++i) {
+        argv[i - 1] = argv_[i].c_str();
+      }
+      argv[argv_.size() - 1] = nullptr;
+      execv(path, (char* const*)argv);
+      // Never executes.
+      return true;
+    } else if (pid_ < 0) {
+      LOG(ERROR) << "Failed to fork process.";
+      return false;
+    } else {
+      // We are in the parent process and fork() was successful.
+      // TODO(mrry): Consider collecting stderr from the child.
+      return true;
+    }
+  }
+
+  bool Kill(int signal) override {
+    if (pid_ == 0) {
+      LOG(ERROR) << "Tried to kill process before starting it.";
+      return false;
+    }
+    return kill(pid_, signal) == 0;
+  }
+
+ private:
+  const std::vector<string> argv_;
+  pid_t pid_;
+  TF_DISALLOW_COPY_AND_ASSIGN(PosixSubProcess);
+};
+}  // namespace
+
+std::unique_ptr<SubProcess> CreateSubProcess(const std::vector<string>& argv) {
+  return std::unique_ptr<SubProcess>(new PosixSubProcess(argv));
+}
+
+}  // namespace testing
+}  // namespace tensorflow
diff --git a/tensorflow/core/platform/test.h b/tensorflow/core/platform/test.h
index 98be407679..4b9cc1278d 100644
--- a/tensorflow/core/platform/test.h
+++ b/tensorflow/core/platform/test.h
@@ -16,6 +16,10 @@ limitations under the License.
 #ifndef TENSORFLOW_PLATFORM_TEST_H_
 #define TENSORFLOW_PLATFORM_TEST_H_
 
+#include <memory>
+#include <vector>
+
+#include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/platform.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -38,6 +42,34 @@ string SrcDir();
 // Returns the same value for the lifetime of the process.
 int RandomSeed();
 
+// Supports spawning and killing child processes, for use in
+// multi-process testing.
+class SubProcess {
+ public:
+  virtual ~SubProcess() {}
+
+  // Starts the subprocess. Returns true on success, otherwise false.
+  // NOTE: This method is not thread-safe.
+  virtual bool Start() = 0;
+
+  // Kills the subprocess with the given signal number. Returns true
+  // on success, otherwise false.
+  // NOTE: This method is not thread-safe.
+  virtual bool Kill(int signal) = 0;
+
+ protected:
+  SubProcess() {}
+
+ private:
+  TF_DISALLOW_COPY_AND_ASSIGN(SubProcess);
+};
+
+// Returns an object that represents a child process that will be
+// launched with the given command-line arguments `argv`. The process
+// must be explicitly started by calling the Start() method on the
+// returned object.
+std::unique_ptr<SubProcess> CreateSubProcess(const std::vector<string>& argv);
+
 }  // namespace testing
 }  // namespace tensorflow