Add a helper to capture a free port for tests. (#252)

Summary: Pull Request resolved: #252 We recently have seen a failed unit test caused by an occupied random port. Our system always expects a pre-assigned free port. To meet this requirement, we add a helper function to find this free port. Note: this is a test-only helper, shouldn't be used in production environment. Reviewed By: adshastri Differential Revision: D37505830 fbshipit-source-id: a23afbcadaafe59c5105589e1119f6c9a6b3b0e2
facebookresearch · Jun 30, 2022 · 5c11a7b · 5c11a7b
1 parent 1d58088
commit 5c11a7b
Show file tree

Hide file tree

Showing 3 changed files with 114 additions and 63 deletions.
diff --git a/fbpcf/engine/communication/test/PartyCommunicationAgentTest.cpp b/fbpcf/engine/communication/test/PartyCommunicationAgentTest.cpp
@@ -22,6 +22,7 @@
 #include "fbpcf/engine/communication/InMemoryPartyCommunicationAgentHost.h"
 #include "fbpcf/engine/communication/SocketPartyCommunicationAgentFactory.h"
 #include "fbpcf/engine/communication/test/AgentFactoryCreationHelper.h"
+#include "fbpcf/engine/communication/test/SocketInTestHelper.h"
 #include "fbpcf/engine/communication/test/TlsCommunicationUtils.h"
 
 namespace fbpcf::engine::communication {
@@ -102,16 +103,7 @@ TEST(InMemoryPartyCommunicationAgentTest, testSendAndReceive) {
 TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithTls) {
   auto createdDir = setUpTlsFiles();
 
-  std::random_device rd;
-  std::default_random_engine defEngine(rd());
-  std::uniform_int_distribution<int> intDistro(10000, 25000);
-
-  /*
-   * We use random ports rather than constants because, during
-   * stress runs, we get errors when trying to bind to the
-   * same port multiple times.
-   */
-  auto port01 = intDistro(defEngine);
+  auto port01 = SocketInTestHelper::findNextOpenPort(5000);
   auto port02 = port01 + 4;
   auto port12 = port01 + 8;
 
@@ -151,11 +143,7 @@ TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithTls) {
 }
 
 TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithoutTls) {
-  std::random_device rd;
-  std::default_random_engine defEngine(rd());
-  std::uniform_int_distribution<int> intDistro(10000, 25000);
-
-  auto port01 = intDistro(defEngine);
+  auto port01 = SocketInTestHelper::findNextOpenPort(5000);
   auto port02 = port01 + 4;
   auto port12 = port01 + 8;
 
@@ -194,11 +182,23 @@ TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithoutTls) {
 }
 
 TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithJammedPort) {
-  std::random_device rd;
-  std::default_random_engine defEngine(rd());
-  std::uniform_int_distribution<int> intDistro(10000, 25000);
+  {
+    // jam port 5000
+    struct sockaddr_in servAddr;
+
+    memset((char*)&servAddr, 0, sizeof(servAddr));
+    servAddr.sin_family = AF_INET;
+    servAddr.sin_addr.s_addr = INADDR_ANY;
+
+    servAddr.sin_port = htons(5000);
+    auto sockfd1 = socket(AF_INET, SOCK_STREAM, 0);
+    ::bind(sockfd1, (struct sockaddr*)&servAddr, sizeof(struct sockaddr_in));
+  }
+
+  auto port01 = SocketInTestHelper::findNextOpenPort(5000);
+  // make sure a port larger than 5000 is found.
+  ASSERT_GE(port01, 5000);
 
-  auto port01 = intDistro(defEngine);
   auto port02 = port01 + 4;
   auto port12 = port01 + 8;
 

diff --git a/fbpcf/engine/communication/test/SocketInTestHelper.h b/fbpcf/engine/communication/test/SocketInTestHelper.h
@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <stdexcept>
+#include <string>
+
+#include <fcntl.h>
+
+#include <arpa/inet.h>
+#include <netinet/in.h>
+#include <sys/socket.h>
+
+namespace fbpcf::engine::communication {
+
+class SocketInTestHelper {
+ public:
+  static int findNextOpenPort(int portNo) {
+    while (!isPortOpen(portNo)) {
+      portNo++;
+      if (portNo > 65535) {
+        throw std::runtime_error("Can't find a open port!");
+      }
+    }
+    return portNo;
+  }
+
+ private:
+  static bool isPortOpen(int port) {
+    auto sockfd = socket(AF_INET, SOCK_STREAM, 0);
+    if (sockfd < 0) {
+      throw std::runtime_error("error opening socket");
+    }
+    int enable = 1;
+
+    if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)) <
+        0) {
+      throw std::runtime_error("setsockopt(SO_REUSEADDR) failed");
+    }
+
+    struct sockaddr_in servAddr;
+
+    memset((char*)&servAddr, 0, sizeof(servAddr));
+    servAddr.sin_family = AF_INET;
+    servAddr.sin_addr.s_addr = INADDR_ANY;
+    servAddr.sin_port = htons(port);
+
+    // if binding to the original port number fails, try to grab a free port
+    // instead
+    if (::bind(
+            sockfd, (struct sockaddr*)&servAddr, sizeof(struct sockaddr_in)) <
+        0) {
+      return false;
+    } else {
+      close(sockfd);
+      return true;
+    }
+  }
+};
+
+} // namespace fbpcf::engine::communication
diff --git a/fbpcf/engine/util/test/benchmarks/BenchmarkHelper.h b/fbpcf/engine/util/test/benchmarks/BenchmarkHelper.h
@@ -18,6 +18,7 @@
 #include "fbpcf/engine/communication/IPartyCommunicationAgent.h"
 #include "fbpcf/engine/communication/IPartyCommunicationAgentFactory.h"
 #include "fbpcf/engine/communication/SocketPartyCommunicationAgentFactory.h"
+#include "fbpcf/engine/communication/test/SocketInTestHelper.h"
 #include "folly/Random.h"
 #include "folly/logging/xlog.h"
 
@@ -39,51 +40,35 @@ getSocketAgents() {
   std::mt19937_64 e(rd());
   std::uniform_int_distribution<int> intDistro(10000, 25000);
 
-  // Since we're using random port numbers, there's a chance we'll encounter a
-  // bind error. Add some retries to reduce the flakiness.
-  auto retries = 5;
-  while (retries--) {
-    try {
-      auto port = intDistro(e);
-      std::map<
-          int,
-          communication::SocketPartyCommunicationAgentFactory::PartyInfo>
-          partyInfo0 = {{1, {"127.0.0.1", port}}};
-      std::map<
-          int,
-          communication::SocketPartyCommunicationAgentFactory::PartyInfo>
-          partyInfo1 = {{0, {"127.0.0.1", port}}};
-
-      auto factory1Future = std::async([&partyInfo1]() {
-        return std::make_unique<
-            communication::SocketPartyCommunicationAgentFactory>(1, partyInfo1);
-      });
-      auto factory0 =
-          std::make_unique<communication::SocketPartyCommunicationAgentFactory>(
-              0, partyInfo0);
-      auto factory1 = factory1Future.get();
-
-      auto task =
-          [](std::unique_ptr<communication::IPartyCommunicationAgentFactory>
-                 factory,
-             int myId) {
-            return factory->create(1 - myId, "benchmark_traffic");
-          };
-
-      auto createSocketAgent0 = std::async(task, std::move(factory0), 0);
-      auto createSocketAgent1 = std::async(task, std::move(factory1), 1);
-
-      auto agent0 = createSocketAgent0.get();
-      auto agent1 = createSocketAgent1.get();
-
-      return {std::move(agent0), std::move(agent1)};
-    } catch (...) {
-      XLOG(INFO) << "Failed to create socket agents. " << retries
-                 << " retries remaining.";
-    }
-  }
+  auto port = communication::SocketInTestHelper::findNextOpenPort(5000);
+  std::map<int, communication::SocketPartyCommunicationAgentFactory::PartyInfo>
+      partyInfo0 = {{1, {"127.0.0.1", port}}};
+  std::map<int, communication::SocketPartyCommunicationAgentFactory::PartyInfo>
+      partyInfo1 = {{0, {"127.0.0.1", port}}};
+
+  auto factory1Future = std::async([&partyInfo1]() {
+    return std::make_unique<
+        communication::SocketPartyCommunicationAgentFactory>(
+        1, partyInfo1, "party_1_unit_test_traffic");
+  });
+  auto factory0 =
+      std::make_unique<communication::SocketPartyCommunicationAgentFactory>(
+          0, partyInfo0, "party_0_unit_test_traffic");
+  auto factory1 = factory1Future.get();
+
+  auto task = [](std::unique_ptr<communication::IPartyCommunicationAgentFactory>
+                     factory,
+                 int myId) {
+    return factory->create(1 - myId, "benchmark_traffic");
+  };
+
+  auto createSocketAgent0 = std::async(task, std::move(factory0), 0);
+  auto createSocketAgent1 = std::async(task, std::move(factory1), 1);
+
+  auto agent0 = createSocketAgent0.get();
+  auto agent1 = createSocketAgent1.get();
 
-  throw std::runtime_error("Failed to create socket agents. Out of retries.");
+  return {std::move(agent0), std::move(agent1)};
 }
 
 // A wrapper class around the existing SocketPartyCommunicationAgentFactory,