Skip to content

Commit 5c11a7b

Browse files
Ruiyu Zhufacebook-github-bot
Ruiyu Zhu
authored and
facebook-github-bot
committed
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
1 parent 1d58088 commit 5c11a7b

3 files changed

Lines changed: 114 additions & 63 deletions

File tree

‎fbpcf/engine/communication/test/PartyCommunicationAgentTest.cpp‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
#include "fbpcf/engine/communication/InMemoryPartyCommunicationAgentHost.h"
2323
#include "fbpcf/engine/communication/SocketPartyCommunicationAgentFactory.h"
2424
#include "fbpcf/engine/communication/test/AgentFactoryCreationHelper.h"
25+
#include "fbpcf/engine/communication/test/SocketInTestHelper.h"
2526
#include "fbpcf/engine/communication/test/TlsCommunicationUtils.h"
2627

2728
namespace fbpcf::engine::communication {
@@ -102,16 +103,7 @@ TEST(InMemoryPartyCommunicationAgentTest, testSendAndReceive) {
102103
TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithTls) {
103104
auto createdDir = setUpTlsFiles();
104105

105-
std::random_device rd;
106-
std::default_random_engine defEngine(rd());
107-
std::uniform_int_distribution<int> intDistro(10000, 25000);
108-
109-
/*
110-
* We use random ports rather than constants because, during
111-
* stress runs, we get errors when trying to bind to the
112-
* same port multiple times.
113-
*/
114-
auto port01 = intDistro(defEngine);
106+
auto port01 = SocketInTestHelper::findNextOpenPort(5000);
115107
auto port02 = port01 + 4;
116108
auto port12 = port01 + 8;
117109

@@ -151,11 +143,7 @@ TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithTls) {
151143
}
152144

153145
TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithoutTls) {
154-
std::random_device rd;
155-
std::default_random_engine defEngine(rd());
156-
std::uniform_int_distribution<int> intDistro(10000, 25000);
157-
158-
auto port01 = intDistro(defEngine);
146+
auto port01 = SocketInTestHelper::findNextOpenPort(5000);
159147
auto port02 = port01 + 4;
160148
auto port12 = port01 + 8;
161149

@@ -194,11 +182,23 @@ TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithoutTls) {
194182
}
195183

196184
TEST(SocketPartyCommunicationAgentTest, testSendAndReceiveWithJammedPort) {
197-
std::random_device rd;
198-
std::default_random_engine defEngine(rd());
199-
std::uniform_int_distribution<int> intDistro(10000, 25000);
185+
{
186+
// jam port 5000
187+
struct sockaddr_in servAddr;
188+
189+
memset((char*)&servAddr, 0, sizeof(servAddr));
190+
servAddr.sin_family = AF_INET;
191+
servAddr.sin_addr.s_addr = INADDR_ANY;
192+
193+
servAddr.sin_port = htons(5000);
194+
auto sockfd1 = socket(AF_INET, SOCK_STREAM, 0);
195+
::bind(sockfd1, (struct sockaddr*)&servAddr, sizeof(struct sockaddr_in));
196+
}
197+
198+
auto port01 = SocketInTestHelper::findNextOpenPort(5000);
199+
// make sure a port larger than 5000 is found.
200+
ASSERT_GE(port01, 5000);
200201

201-
auto port01 = intDistro(defEngine);
202202
auto port02 = port01 + 4;
203203
auto port12 = port01 + 8;
204204

‎fbpcf/engine/communication/test/SocketInTestHelper.h‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
*
4+
* This source code is licensed under the MIT license found in the
5+
* LICENSE file in the root directory of this source tree.
6+
*/
7+
8+
#pragma once
9+
10+
#include <stdexcept>
11+
#include <string>
12+
13+
#include <fcntl.h>
14+
15+
#include <arpa/inet.h>
16+
#include <netinet/in.h>
17+
#include <sys/socket.h>
18+
19+
namespace fbpcf::engine::communication {
20+
21+
class SocketInTestHelper {
22+
public:
23+
static int findNextOpenPort(int portNo) {
24+
while (!isPortOpen(portNo)) {
25+
portNo++;
26+
if (portNo > 65535) {
27+
throw std::runtime_error("Can't find a open port!");
28+
}
29+
}
30+
return portNo;
31+
}
32+
33+
private:
34+
static bool isPortOpen(int port) {
35+
auto sockfd = socket(AF_INET, SOCK_STREAM, 0);
36+
if (sockfd < 0) {
37+
throw std::runtime_error("error opening socket");
38+
}
39+
int enable = 1;
40+
41+
if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)) <
42+
0) {
43+
throw std::runtime_error("setsockopt(SO_REUSEADDR) failed");
44+
}
45+
46+
struct sockaddr_in servAddr;
47+
48+
memset((char*)&servAddr, 0, sizeof(servAddr));
49+
servAddr.sin_family = AF_INET;
50+
servAddr.sin_addr.s_addr = INADDR_ANY;
51+
servAddr.sin_port = htons(port);
52+
53+
// if binding to the original port number fails, try to grab a free port
54+
// instead
55+
if (::bind(
56+
sockfd, (struct sockaddr*)&servAddr, sizeof(struct sockaddr_in)) <
57+
0) {
58+
return false;
59+
} else {
60+
close(sockfd);
61+
return true;
62+
}
63+
}
64+
};
65+
66+
} // namespace fbpcf::engine::communication

‎fbpcf/engine/util/test/benchmarks/BenchmarkHelper.h‎

Lines changed: 29 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include "fbpcf/engine/communication/IPartyCommunicationAgent.h"
1919
#include "fbpcf/engine/communication/IPartyCommunicationAgentFactory.h"
2020
#include "fbpcf/engine/communication/SocketPartyCommunicationAgentFactory.h"
21+
#include "fbpcf/engine/communication/test/SocketInTestHelper.h"
2122
#include "folly/Random.h"
2223
#include "folly/logging/xlog.h"
2324

@@ -39,51 +40,35 @@ getSocketAgents() {
3940
std::mt19937_64 e(rd());
4041
std::uniform_int_distribution<int> intDistro(10000, 25000);
4142

42-
// Since we're using random port numbers, there's a chance we'll encounter a
43-
// bind error. Add some retries to reduce the flakiness.
44-
auto retries = 5;
45-
while (retries--) {
46-
try {
47-
auto port = intDistro(e);
48-
std::map<
49-
int,
50-
communication::SocketPartyCommunicationAgentFactory::PartyInfo>
51-
partyInfo0 = {{1, {"127.0.0.1", port}}};
52-
std::map<
53-
int,
54-
communication::SocketPartyCommunicationAgentFactory::PartyInfo>
55-
partyInfo1 = {{0, {"127.0.0.1", port}}};
56-
57-
auto factory1Future = std::async([&partyInfo1]() {
58-
return std::make_unique<
59-
communication::SocketPartyCommunicationAgentFactory>(1, partyInfo1);
60-
});
61-
auto factory0 =
62-
std::make_unique<communication::SocketPartyCommunicationAgentFactory>(
63-
0, partyInfo0);
64-
auto factory1 = factory1Future.get();
65-
66-
auto task =
67-
[](std::unique_ptr<communication::IPartyCommunicationAgentFactory>
68-
factory,
69-
int myId) {
70-
return factory->create(1 - myId, "benchmark_traffic");
71-
};
72-
73-
auto createSocketAgent0 = std::async(task, std::move(factory0), 0);
74-
auto createSocketAgent1 = std::async(task, std::move(factory1), 1);
75-
76-
auto agent0 = createSocketAgent0.get();
77-
auto agent1 = createSocketAgent1.get();
78-
79-
return {std::move(agent0), std::move(agent1)};
80-
} catch (...) {
81-
XLOG(INFO) << "Failed to create socket agents. " << retries
82-
<< " retries remaining.";
83-
}
84-
}
43+
auto port = communication::SocketInTestHelper::findNextOpenPort(5000);
44+
std::map<int, communication::SocketPartyCommunicationAgentFactory::PartyInfo>
45+
partyInfo0 = {{1, {"127.0.0.1", port}}};
46+
std::map<int, communication::SocketPartyCommunicationAgentFactory::PartyInfo>
47+
partyInfo1 = {{0, {"127.0.0.1", port}}};
48+
49+
auto factory1Future = std::async([&partyInfo1]() {
50+
return std::make_unique<
51+
communication::SocketPartyCommunicationAgentFactory>(
52+
1, partyInfo1, "party_1_unit_test_traffic");
53+
});
54+
auto factory0 =
55+
std::make_unique<communication::SocketPartyCommunicationAgentFactory>(
56+
0, partyInfo0, "party_0_unit_test_traffic");
57+
auto factory1 = factory1Future.get();
58+
59+
auto task = [](std::unique_ptr<communication::IPartyCommunicationAgentFactory>
60+
factory,
61+
int myId) {
62+
return factory->create(1 - myId, "benchmark_traffic");
63+
};
64+
65+
auto createSocketAgent0 = std::async(task, std::move(factory0), 0);
66+
auto createSocketAgent1 = std::async(task, std::move(factory1), 1);
67+
68+
auto agent0 = createSocketAgent0.get();
69+
auto agent1 = createSocketAgent1.get();
8570

86-
throw std::runtime_error("Failed to create socket agents. Out of retries.");
71+
return {std::move(agent0), std::move(agent1)};
8772
}
8873

8974
// A wrapper class around the existing SocketPartyCommunicationAgentFactory,

0 commit comments

Comments
 (0)