Add sequence number to gloo::transport::tcp::Address

Summary: Pull Request resolved: pytorch#239

Differential Revision: D18909096

Test Plan: Imported from OSS

Pulled By: pietern

fbshipit-source-id: f0e22f24d7cdd0120f8265268a7ad0e20bcd1109

Author: pietern

gloo/rendezvous/context.h

@@ -14,6 +14,7 @@
 #include "gloo/common/error.h"
 #include "gloo/context.h"
 #include "gloo/rendezvous/store.h"
+#include "gloo/transport/address.h"
 #include "gloo/transport/device.h"
 
 namespace gloo {
@@ -38,8 +39,8 @@ class Context : public ::gloo::Context {
 
 class ContextFactory {
  public:
-  // Assume a pair's address is no bigger than 128 bytes
-  static constexpr auto kMaxAddressSize = 128;
+  static constexpr auto kMaxAddressSize =
+      ::gloo::transport::Address::kMaxByteSize;
 
   explicit ContextFactory(std::shared_ptr<::gloo::Context> backingContext);
 

gloo/transport/address.h

@@ -16,6 +16,9 @@ namespace transport {
 
 class Address {
  public:
+  // Upper bound for an address' byte representation.
+  static constexpr auto kMaxByteSize = 192;
+
   virtual ~Address() = 0;
 
   virtual std::string str() const = 0;

gloo/transport/tcp/address.cc

@@ -17,47 +17,53 @@ namespace gloo {
 namespace transport {
 namespace tcp {
 
-Address::Address(const struct sockaddr_storage& ss) {
-  ss_ = ss;
+Address::Address(struct sockaddr_storage ss, sequence_number_t seq) {
+  impl_.ss = std::move(ss);
+  impl_.seq = seq;
 }
 
 Address::Address(const struct sockaddr* addr, size_t addrlen) {
-  memcpy(&ss_, addr, addrlen);
+  memcpy(&impl_.ss, addr, addrlen);
 }
 
 Address::Address(const std::vector<char>& bytes) {
-  GLOO_ENFORCE_EQ(sizeof(ss_), bytes.size());
-  memcpy(&ss_, bytes.data(), sizeof(ss_));
+  GLOO_ENFORCE_EQ(sizeof(impl_), bytes.size());
+  memcpy(&impl_, bytes.data(), sizeof(impl_));
 }
 
 std::vector<char> Address::bytes() const {
-  std::vector<char> bytes(sizeof(ss_));
-  memcpy(bytes.data(), &ss_, sizeof(ss_));
+  std::vector<char> bytes(sizeof(impl_));
+  memcpy(bytes.data(), &impl_, sizeof(impl_));
   return bytes;
 }
 
 std::string Address::str() const {
-  char str[INET6_ADDRSTRLEN + 8];
+  char str[INET6_ADDRSTRLEN + 128];
   int port = 0;
 
   str[0] = '[';
-  if (ss_.ss_family == AF_INET) {
-    struct sockaddr_in* in = (struct sockaddr_in*)&ss_;
+  if (impl_.ss.ss_family == AF_INET) {
+    struct sockaddr_in* in = (struct sockaddr_in*)&impl_.ss;
     inet_ntop(AF_INET, &in->sin_addr, str + 1, sizeof(str) - 1);
     port = in->sin_port;
-  } else if (ss_.ss_family == AF_INET6) {
-    struct sockaddr_in6* in6 = (struct sockaddr_in6*)&ss_;
+  } else if (impl_.ss.ss_family == AF_INET6) {
+    struct sockaddr_in6* in6 = (struct sockaddr_in6*)&impl_.ss;
     inet_ntop(AF_INET6, &in6->sin6_addr, str + 1, sizeof(str) - 1);
     port = in6->sin6_port;
   } else {
     snprintf(str + 1, sizeof(str) - 1, "none");
   }
 
-  auto len = strlen(str);
+  size_t len = strlen(str);
   if (port > 0) {
-    snprintf(str + len, sizeof(str) - len, "]:%d", port);
+    len += snprintf(str + len, sizeof(str) - len, "]:%d", port);
   } else {
-    snprintf(str + len, sizeof(str) - len, "]");
+    len += snprintf(str + len, sizeof(str) - len, "]");
+  }
+
+  // Append sequence number if one is set.
+  if (impl_.seq != kSequenceNumberUnset) {
+    len += snprintf(str + len, sizeof(str) - len, "$%ld", impl_.seq);
   }
 
   return str;

gloo/transport/tcp/address.h

@@ -8,38 +8,63 @@
 
 #pragma once
 
-#include <string>
-
 #include <sys/socket.h>
+#include <unistd.h>
 
 #include "gloo/transport/address.h"
 
 namespace gloo {
 namespace transport {
 namespace tcp {
 
-// Forward declaration
-class Pair;
+using sequence_number_t = ssize_t;
 
 class Address : public ::gloo::transport::Address {
  public:
+  static constexpr sequence_number_t kSequenceNumberUnset = -1;
+
   Address() {}
-  explicit Address(const struct sockaddr_storage&);
+
+  explicit Address(struct sockaddr_storage ss, sequence_number_t seq = -1);
+
   explicit Address(const struct sockaddr* addr, size_t addrlen);
+
   explicit Address(const std::vector<char>&);
-  virtual ~Address() {}
 
   virtual std::vector<char> bytes() const override;
+
   virtual std::string str() const override;
 
+  const struct sockaddr_storage& getSockaddr() const {
+    return impl_.ss;
+  }
+
+  sequence_number_t getSeq() const {
+    return impl_.seq;
+  }
+
   static Address fromSockName(int fd);
+
   static Address fromPeerName(int fd);
 
  protected:
-  struct sockaddr_storage ss_;
+  // Encapsulate fields such that it is trivially copyable. This class
+  // is not trivially copyable itself.
+  struct Impl {
+    // IP address of the listening socket.
+    struct sockaddr_storage ss;
+
+    // Sequence number of this address.
+    // If this is equal to -1, the address is assumed to
+    // represent the listening socket of a device. The sequence number
+    // must be set before it can be used by a pair.
+    sequence_number_t seq{kSequenceNumberUnset};
+  };
+
+  static_assert(std::is_trivially_copyable<Impl>::value, "!");
+  static_assert(sizeof(Impl) <= kMaxByteSize, "!");
 
-  // Pair can access ss_ directly
-  friend class Pair;
+  Impl impl_;
 };
 
 } // namespace tcp

gloo/transport/tcp/pair.cc

@@ -202,22 +202,25 @@ void Pair::connect(const Address& peer) {
 
   peer_ = peer;
 
+  const auto& selfAddr = self_.getSockaddr();
+  const auto& peerAddr = peer_.getSockaddr();
+
   // Addresses have to have same family
-  if (self_.ss_.ss_family != peer_.ss_.ss_family) {
+  if (selfAddr.ss_family != peerAddr.ss_family) {
     GLOO_THROW_INVALID_OPERATION_EXCEPTION("address family mismatch");
   }
 
-  if (self_.ss_.ss_family == AF_INET) {
-    struct sockaddr_in* sa = (struct sockaddr_in*)&self_.ss_;
-    struct sockaddr_in* sb = (struct sockaddr_in*)&peer_.ss_;
+  if (selfAddr.ss_family == AF_INET) {
+    struct sockaddr_in* sa = (struct sockaddr_in*)&selfAddr;
+    struct sockaddr_in* sb = (struct sockaddr_in*)&peerAddr;
     addrlen = sizeof(struct sockaddr_in);
     rv = memcmp(&sa->sin_addr, &sb->sin_addr, sizeof(struct in_addr));
     if (rv == 0) {
       rv = sa->sin_port - sb->sin_port;
     }
-  } else if (peer_.ss_.ss_family == AF_INET6) {
-    struct sockaddr_in6* sa = (struct sockaddr_in6*)&self_.ss_;
-    struct sockaddr_in6* sb = (struct sockaddr_in6*)&peer_.ss_;
+  } else if (peerAddr.ss_family == AF_INET6) {
+    struct sockaddr_in6* sa = (struct sockaddr_in6*)&selfAddr;
+    struct sockaddr_in6* sb = (struct sockaddr_in6*)&peerAddr;
     addrlen = sizeof(struct sockaddr_in6);
     rv = memcmp(&sa->sin6_addr, &sb->sin6_addr, sizeof(struct in6_addr));
     if (rv == 0) {
@@ -243,7 +246,7 @@ void Pair::connect(const Address& peer) {
   ::close(fd_);
 
   // Create new socket to connect to peer.
-  fd_ = socket(peer_.ss_.ss_family, SOCK_STREAM | SOCK_NONBLOCK, 0);
+  fd_ = socket(peerAddr.ss_family, SOCK_STREAM | SOCK_NONBLOCK, 0);
   if (fd_ == -1) {
     signalAndThrowException(GLOO_ERROR_MSG("socket: ", strerror(errno)));
   }
@@ -258,7 +261,7 @@ void Pair::connect(const Address& peer) {
   }
 
   // Connect to peer
-  rv = ::connect(fd_, (struct sockaddr*)&peer_.ss_, addrlen);
+  rv = ::connect(fd_, (struct sockaddr*)&peerAddr, addrlen);
   if (rv == -1 && errno != EINPROGRESS) {
     ::close(fd_);
     fd_ = FD_INVALID;