1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include "tensorflow/core/platform/net.h" 17 18 #include <cstdlib> 19 #include <unordered_set> 20 21 #include <sys/types.h> 22 #include <winsock2.h> 23 24 #include "tensorflow/core/platform/logging.h" 25 #include "tensorflow/core/platform/windows/error.h" 26 27 #undef ERROR 28 29 #pragma comment(lib, "Ws2_32.lib") 30 31 namespace tensorflow { 32 namespace internal { 33 34 namespace { 35 36 bool IsPortAvailable(int* port, bool is_tcp) { 37 const int protocol = is_tcp ? IPPROTO_TCP : 0; 38 SOCKET sock = socket(AF_INET, is_tcp ? SOCK_STREAM : SOCK_DGRAM, protocol); 39 40 struct sockaddr_in addr; 41 int addr_len = static_cast<int>(sizeof(addr)); 42 int actual_port; 43 44 CHECK_GE(*port, 0); 45 CHECK_LE(*port, 65535); 46 if (sock == INVALID_SOCKET) { 47 LOG(ERROR) << "socket() failed: " 48 << GetWindowsErrorMessage(WSAGetLastError()); 49 return false; 50 } 51 52 // SO_REUSEADDR lets us start up a server immediately after it exits. 53 const int one = 1; 54 int result = setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, 55 reinterpret_cast<const char*>(&one), sizeof(one)); 56 if (result == SOCKET_ERROR) { 57 LOG(ERROR) << "setsockopt() failed: " 58 << GetWindowsErrorMessage(WSAGetLastError()); 59 closesocket(sock); 60 return false; 61 } 62 63 // Try binding to port. 64 addr.sin_family = AF_INET; 65 addr.sin_addr.s_addr = INADDR_ANY; 66 addr.sin_port = htons((uint16_t)*port); 67 result = bind(sock, (struct sockaddr*)&addr, sizeof(addr)); 68 if (result == SOCKET_ERROR) { 69 LOG(WARNING) << "bind(port=" << *port 70 << ") failed: " << GetWindowsErrorMessage(WSAGetLastError()); 71 closesocket(sock); 72 return false; 73 } 74 75 // Get the bound port number. 76 result = getsockname(sock, (struct sockaddr*)&addr, &addr_len); 77 if (result == SOCKET_ERROR) { 78 LOG(WARNING) << "getsockname() failed: " 79 << GetWindowsErrorMessage(WSAGetLastError()); 80 closesocket(sock); 81 return false; 82 } 83 84 CHECK_LE(addr_len, sizeof(addr)); 85 actual_port = ntohs(addr.sin_port); 86 CHECK_GT(actual_port, 0); 87 if (*port == 0) { 88 *port = actual_port; 89 } else { 90 CHECK_EQ(*port, actual_port); 91 } 92 93 closesocket(sock); 94 return true; 95 } 96 97 const int kNumRandomPortsToPick = 100; 98 const int kMaximumTrials = 1000; 99 100 } // namespace 101 102 int PickUnusedPortOrDie() { 103 WSADATA wsaData; 104 if (WSAStartup(MAKEWORD(2, 2), &wsaData) != NO_ERROR) { 105 LOG(ERROR) << "Error at WSAStartup()"; 106 return false; 107 } 108 109 static std::unordered_set<int> chosen_ports; 110 111 // Type of port to first pick in the next iteration. 112 bool is_tcp = true; 113 int trial = 0; 114 while (true) { 115 int port; 116 trial++; 117 CHECK_LE(trial, kMaximumTrials) 118 << "Failed to pick an unused port for testing."; 119 if (trial == 1) { 120 port = GetCurrentProcessId() % (65536 - 30000) + 30000; 121 } else if (trial <= kNumRandomPortsToPick) { 122 port = rand() % (65536 - 30000) + 30000; 123 } else { 124 port = 0; 125 } 126 127 if (chosen_ports.find(port) != chosen_ports.end()) { 128 continue; 129 } 130 if (!IsPortAvailable(&port, is_tcp)) { 131 continue; 132 } 133 134 CHECK_GT(port, 0); 135 if (!IsPortAvailable(&port, !is_tcp)) { 136 is_tcp = !is_tcp; 137 continue; 138 } 139 140 chosen_ports.insert(port); 141 WSACleanup(); 142 return port; 143 } 144 145 return 0; 146 } 147 148 } // namespace internal 149 } // namespace tensorflow 150
