[TransferEngine] Add IPv6 support (#479)

* [TransferEngine] Fix compilation bug in NVLink xport

* [TransferEngine] Add IPv6 support
This commit is contained in:
Feng Ren 2025-06-12 14:35:00 +08:00 committed by GitHub
parent 20829bc4f4
commit 883a5a734e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 83 additions and 34 deletions

View File

@ -45,6 +45,7 @@ struct GlobalConfig {
int log_level = google::INFO;
bool trace = false;
int64_t slice_timeout = -1;
bool use_ipv6 = false;
};
void loadGlobalConfig(GlobalConfig &config);

View File

@ -12,11 +12,11 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include "config.h"
#include <dirent.h>
#include <unistd.h>
#include "config.h"
namespace mooncake {
void loadGlobalConfig(GlobalConfig &config) {
const char *num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX");
@ -182,14 +182,15 @@ void loadGlobalConfig(GlobalConfig &config) {
config.metacache = false;
}
const char *handshake_listen_backlog = std::getenv("MC_HANDSHAKE_LISTEN_BACKLOG");
const char *handshake_listen_backlog =
std::getenv("MC_HANDSHAKE_LISTEN_BACKLOG");
if (handshake_listen_backlog) {
int val = std::stoi(handshake_listen_backlog);
if (val > 0) {
config.handshake_listen_backlog = val;
} else {
LOG(WARNING)
<< "Ignore value from environment variable MC_HANDSHAKE_LISTEN_BACKLOG";
LOG(WARNING) << "Ignore value from environment variable "
"MC_HANDSHAKE_LISTEN_BACKLOG";
}
}
@ -223,11 +224,13 @@ void loadGlobalConfig(GlobalConfig &config) {
if (log_dir_path) {
google::InitGoogleLogging("mooncake-transfer-engine");
if (opendir(log_dir_path) == NULL) {
LOG(WARNING) << "Path [" << log_dir_path <<
"] is not a valid directory path. Still logging to stderr.";
LOG(WARNING)
<< "Path [" << log_dir_path
<< "] is not a valid directory path. Still logging to stderr.";
} else if (access(log_dir_path, W_OK) != 0) {
LOG(WARNING) << "Path [" << log_dir_path <<
"] is not a permitted directory path for the current user. \
LOG(WARNING)
<< "Path [" << log_dir_path
<< "] is not a permitted directory path for the current user. \
Still logging to stderr.";
} else {
FLAGS_log_dir = log_dir_path;
@ -235,6 +238,10 @@ void loadGlobalConfig(GlobalConfig &config) {
FLAGS_stop_logging_if_full_disk = true;
}
}
if (std::getenv("MC_USE_IPV6")) {
config.use_ipv6 = true;
}
}
std::string mtuLengthToString(ibv_mtu mtu) {

View File

@ -501,17 +501,13 @@ struct SocketHandShakePlugin : public HandShakePlugin {
return 0;
}
sockaddr_in bind_address;
int on = 1;
memset(&bind_address, 0, sizeof(sockaddr_in));
bind_address.sin_family = AF_INET;
bind_address.sin_port = htons(listen_port);
bind_address.sin_addr.s_addr = INADDR_ANY;
if (sockfd >= 0) {
listen_fd_ = sockfd;
} else {
listen_fd_ = socket(AF_INET, SOCK_STREAM, 0);
listen_fd_ = socket(globalConfig().use_ipv6 ? AF_INET6 : AF_INET,
SOCK_STREAM, 0);
if (listen_fd_ < 0) {
PLOG(ERROR) << "SocketHandShakePlugin: socket()";
return ERR_SOCKET;
@ -535,12 +531,34 @@ struct SocketHandShakePlugin : public HandShakePlugin {
return ERR_SOCKET;
}
if (bind(listen_fd_, (sockaddr *)&bind_address,
sizeof(sockaddr_in)) < 0) {
PLOG(ERROR) << "SocketHandShakePlugin: bind (port "
<< listen_port << ")";
closeListen();
return ERR_SOCKET;
if (globalConfig().use_ipv6) {
sockaddr_in6 bind_address;
memset(&bind_address, 0, sizeof(sockaddr_in6));
bind_address.sin6_family = AF_INET6;
bind_address.sin6_port = htons(listen_port);
bind_address.sin6_addr = IN6ADDR_ANY_INIT;
if (bind(listen_fd_, (sockaddr *)&bind_address,
sizeof(sockaddr_in6)) < 0) {
PLOG(ERROR) << "SocketHandShakePlugin: bind (port "
<< listen_port << ")";
closeListen();
return ERR_SOCKET;
}
} else {
sockaddr_in bind_address;
memset(&bind_address, 0, sizeof(sockaddr_in));
bind_address.sin_family = AF_INET;
bind_address.sin_port = htons(listen_port);
bind_address.sin_addr.s_addr = INADDR_ANY;
if (bind(listen_fd_, (sockaddr *)&bind_address,
sizeof(sockaddr_in)) < 0) {
PLOG(ERROR) << "SocketHandShakePlugin: bind (port "
<< listen_port << ")";
closeListen();
return ERR_SOCKET;
}
}
}
@ -857,12 +875,15 @@ std::vector<std::string> findLocalIpAddresses() {
return ips;
}
auto use_ipv6 = globalConfig().use_ipv6;
sa_family_t family = use_ipv6 ? AF_INET6 : AF_INET;
for (ifa = ifaddr; ifa != nullptr; ifa = ifa->ifa_next) {
if (ifa->ifa_addr == nullptr) {
continue;
}
if (ifa->ifa_addr->sa_family == AF_INET) {
if (ifa->ifa_addr->sa_family == family) {
if (strcmp(ifa->ifa_name, "lo") == 0) {
continue;
}
@ -875,8 +896,11 @@ std::vector<std::string> findLocalIpAddresses() {
}
char host[NI_MAXHOST];
if (getnameinfo(ifa->ifa_addr, sizeof(struct sockaddr_in), host,
NI_MAXHOST, nullptr, 0, NI_NUMERICHOST) == 0) {
if (getnameinfo(ifa->ifa_addr,
use_ipv6 ? sizeof(struct sockaddr_in6)
: sizeof(struct sockaddr_in),
host, NI_MAXHOST, nullptr, 0,
NI_NUMERICHOST) == 0) {
LOG(INFO) << "Found active interface " << ifa->ifa_name
<< " with IP " << host;
ips.push_back(host);
@ -894,9 +918,11 @@ uint16_t findAvailableTcpPort(int &sockfd) {
const int min_port = 15000;
const int max_port = 17000;
const int max_attempts = 500;
bool use_ipv6 = globalConfig().use_ipv6;
for (int attempt = 0; attempt < max_attempts; ++attempt) {
int port = min_port + rand_dist(rand_gen) % (max_port - min_port + 1);
sockfd = socket(AF_INET, SOCK_STREAM, 0);
sockfd = socket(use_ipv6 ? AF_INET6 : AF_INET, SOCK_STREAM, 0);
if (sockfd == -1) {
continue;
}
@ -918,15 +944,30 @@ uint16_t findAvailableTcpPort(int &sockfd) {
continue;
}
sockaddr_in bind_address;
memset(&bind_address, 0, sizeof(sockaddr_in));
bind_address.sin_family = AF_INET;
bind_address.sin_port = htons(port);
bind_address.sin_addr.s_addr = INADDR_ANY;
if (bind(sockfd, (sockaddr *)&bind_address, sizeof(sockaddr_in)) < 0) {
close(sockfd);
sockfd = -1;
continue;
if (use_ipv6) {
sockaddr_in6 bind_address;
memset(&bind_address, 0, sizeof(sockaddr_in6));
bind_address.sin6_family = AF_INET6;
bind_address.sin6_port = htons(port);
bind_address.sin6_addr = IN6ADDR_ANY_INIT;
if (bind(sockfd, (sockaddr *)&bind_address, sizeof(sockaddr_in6)) <
0) {
close(sockfd);
sockfd = -1;
continue;
}
} else {
sockaddr_in bind_address;
memset(&bind_address, 0, sizeof(sockaddr_in));
bind_address.sin_family = AF_INET;
bind_address.sin_port = htons(port);
bind_address.sin_addr.s_addr = INADDR_ANY;
if (bind(sockfd, (sockaddr *)&bind_address, sizeof(sockaddr_in)) <
0) {
close(sockfd);
sockfd = -1;
continue;
}
}
return port;