Adds a C++ test utility for picking an unused port.

The PickUnusedPortOrDie implementation is based on a simplified
version of `grpc_pick_unused_port_or_die()` in gRPC. This utility will
be necessary for tests of the distributed runtime (issue #23).
Change: 115345502

2 files changed, 104 insertions, 0 deletions

diff --git a/tensorflow/core/platform/posix/test.cc b/tensorflow/core/platform/posix/test.cc
index 2fa2d36ceb..366aa0fd53 100644
--- a/tensorflow/core/platform/posix/test.cc
+++ b/tensorflow/core/platform/posix/test.cc
@@ -15,7 +15,11 @@ limitations under the License.
 
 #include "tensorflow/core/platform/test.h"
 
+#include <unordered_set>
+
+#include <netinet/in.h>
 #include <signal.h>
+#include <sys/socket.h>
 #include <sys/types.h>
 #include <unistd.h>
 
@@ -75,5 +79,101 @@ std::unique_ptr<SubProcess> CreateSubProcess(const std::vector<string>& argv) {
   return std::unique_ptr<SubProcess>(new PosixSubProcess(argv));
 }
 
+namespace {
+bool IsPortAvailable(int* port, bool is_tcp) {
+  const int protocol = is_tcp ? IPPROTO_TCP : 0;
+  const int fd = socket(AF_INET, is_tcp ? SOCK_STREAM : SOCK_DGRAM, protocol);
+
+  struct sockaddr_in addr;
+  socklen_t addr_len = sizeof(addr);
+  int actual_port;
+
+  CHECK_GE(*port, 0);
+  CHECK_LE(*port, 65535);
+  if (fd < 0) {
+    LOG(ERROR) << "socket() failed: " << strerror(errno);
+    return false;
+  }
+
+  // SO_REUSEADDR lets us start up a server immediately after it exists.
+  int one = 1;
+  if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) < 0) {
+    LOG(ERROR) << "setsockopt() failed: " << strerror(errno);
+    close(fd);
+    return false;
+  }
+
+  // Try binding to port.
+  addr.sin_family = AF_INET;
+  addr.sin_addr.s_addr = INADDR_ANY;
+  addr.sin_port = htons((uint16_t)*port);
+  if (bind(fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
+    LOG(WARNING) << "bind(port=" << *port << ") failed: " << strerror(errno);
+    close(fd);
+    return false;
+  }
+
+  // Get the bound port number.
+  if (getsockname(fd, (struct sockaddr*)&addr, &addr_len) < 0) {
+    LOG(WARNING) << "getsockname() failed: " << strerror(errno);
+    close(fd);
+    return false;
+  }
+  CHECK_LE(addr_len, sizeof(addr));
+  actual_port = ntohs(addr.sin_port);
+  CHECK_GT(actual_port, 0);
+  if (*port == 0) {
+    *port = actual_port;
+  } else {
+    CHECK_EQ(*port, actual_port);
+  }
+  close(fd);
+  return true;
+}
+
+const int kNumRandomPortsToPick = 100;
+const int kMaximumTrials = 1000;
+
+}  // namespace
+
+int PickUnusedPortOrDie() {
+  static std::unordered_set<int> chosen_ports;
+
+  // Type of port to first pick in the next iteration.
+  bool is_tcp = true;
+  int trial = 0;
+  while (true) {
+    int port;
+    trial++;
+    CHECK_LE(trial, kMaximumTrials)
+        << "Failed to pick an unused port for testing.";
+    if (trial == 1) {
+      port = getpid() % (65536 - 30000) + 30000;
+    } else if (trial <= kNumRandomPortsToPick) {
+      port = rand() % (65536 - 30000) + 30000;
+    } else {
+      port = 0;
+    }
+
+    if (chosen_ports.find(port) != chosen_ports.end()) {
+      continue;
+    }
+    if (!IsPortAvailable(&port, is_tcp)) {
+      continue;
+    }
+
+    CHECK_GT(port, 0);
+    if (!IsPortAvailable(&port, !is_tcp)) {
+      is_tcp = !is_tcp;
+      continue;
+    }
+
+    chosen_ports.insert(port);
+    return port;
+  }
+
+  return 0;
+}
+
 }  // namespace testing
 }  // namespace tensorflow
diff --git a/tensorflow/core/platform/test.h b/tensorflow/core/platform/test.h
index 4b9cc1278d..73f2f1e272 100644
--- a/tensorflow/core/platform/test.h
+++ b/tensorflow/core/platform/test.h
@@ -70,6 +70,10 @@ class SubProcess {
 // returned object.
 std::unique_ptr<SubProcess> CreateSubProcess(const std::vector<string>& argv);
 
+// Returns an unused port number, for use in multi-process testing.
+// NOTE: This function is not thread-safe.
+int PickUnusedPortOrDie();
+
 }  // namespace testing
 }  // namespace tensorflow