gloo: async

Differential Revision: D69698406

Author: d4l3k

gloo/common/store.h

@@ -9,12 +9,13 @@
 #pragma once
 
 #include <chrono>
+#include <memory>
 #include <string>
 #include <vector>
 
 namespace gloo {
 
-class IStore {
+class IStore : public std::enable_shared_from_this<IStore> {
  public:
   virtual ~IStore() = default;
 

gloo/common/utils.cc

@@ -31,9 +31,10 @@ std::string getHostname() {
 }
 
 bool useRankAsSeqNumber() {
-  const auto& res = getenv("GLOO_ENABLE_RANK_AS_SEQUENCE_NUMBER");
-  return res != nullptr &&
-      (std::string(res) == "True" || std::string(res) == "1");
+  // const auto& res = getenv("GLOO_ENABLE_RANK_AS_SEQUENCE_NUMBER");
+  // return res != nullptr && (std::string(res) == "True" || std::string(res) ==
+  // "1");
+  return true;
 }
 
 bool isStoreExtendedApiEnabled() {

gloo/test/base_test.h

@@ -106,7 +106,7 @@ class BaseTest : public ::testing::Test {
       std::function<void(std::shared_ptr<Context>)> fn,
       int base = 2) {
     Barrier barrier(size);
-    ::gloo::rendezvous::HashStore store;
+    auto store = std::make_shared<::gloo::rendezvous::HashStore>();
 
     spawnThreads(size, [&](int rank) {
       auto context =
@@ -118,7 +118,7 @@ class BaseTest : public ::testing::Test {
       if (!device) {
         return;
       }
-      context->connectFullMesh(store, device);
+      context->connectFullMesh(*store, device);
 
       try {
         fn(context);

gloo/transport/tcp/context.cc

@@ -24,10 +24,10 @@ namespace gloo {
 namespace transport {
 namespace tcp {
 
-constexpr int kDefaultBatchSize = 128;
-
 Context::Context(std::shared_ptr<Device> device, int rank, int size)
-    : ::gloo::transport::Context(rank, size), device_(std::move(device)) {}
+    : ::gloo::transport::Context(rank, size), device_(std::move(device)) {
+  connecting_.resize(size);
+}
 
 Context::~Context() {
   // Pairs refer to device by raw pointer.
@@ -50,8 +50,6 @@ void Context::createAndConnectAllPairs(IStore& store) {
   // it's not super straightforward so left for folks having more bandwidth
   // later on.
 
-  int localRank = 0;
-  bool localRankSet = false;
   auto localHostName = getHostname();
   bool useRankAsSeqNum = useRankAsSeqNumber();
 
@@ -60,13 +58,15 @@ void Context::createAndConnectAllPairs(IStore& store) {
   // it's just to keep the later seq num matching logic simple
   std::vector<ssize_t> pairIdentifiers;
   for (int i = 0; i < size; i++) {
-    const auto& pair = createPair(i, useRankAsSeqNum);
-    if (!useRankAsSeqNum) {
-      // Need to preserve the order of the pair identifiers if we are not using
-      // the rank as seq number
-      pairIdentifiers.emplace_back(
-          static_cast<Pair*>(pair.get())->address().getSeq());
-    }
+    // const auto& pair =
+    createPair(i, useRankAsSeqNum);
+    // if (!useRankAsSeqNum) {
+    //   // Need to preserve the order of the pair identifiers if we are not
+    //   using
+    //   // the rank as seq number
+    //   pairIdentifiers.emplace_back(
+    //       static_cast<Pair*>(pair.get())->address().getSeq());
+    // }
   }
 
   // Obtain the pair object for this rank
@@ -88,68 +88,96 @@ void Context::createAndConnectAllPairs(IStore& store) {
       localHostName, deviceAddress.bytes(), std::move(pairIdentifiers));
   store.set(std::to_string(rank), currentRankInfo.bytes());
 
-  std::vector<std::vector<char>> remoteRankInfos;
-  int key = 0;
-  if (isStoreExtendedApiEnabled() && store.has_v2_support()) {
-    auto sizeRemaining = size;
-    while (sizeRemaining > 0) {
-      const auto batchKeys = std::min(kDefaultBatchSize, sizeRemaining);
-      std::vector<std::string> keys(batchKeys);
-      std::generate_n(
-          keys.begin(), batchKeys, [&] { return std::to_string(key++); });
-      const auto& batchRemoteInfos = store.multi_get(keys);
-      remoteRankInfos.insert(
-          remoteRankInfos.end(),
-          batchRemoteInfos.begin(),
-          batchRemoteInfos.end());
-      sizeRemaining -= batchKeys;
+  store_ = store.shared_from_this();
+
+  // we don't use local rank so why even set it??
+
+  /*
+    std::vector<std::vector<char>> remoteRankInfos;
+    int key = 0;
+    if (isStoreExtendedApiEnabled() && store.has_v2_support()) {
+      auto sizeRemaining = size;
+      while (sizeRemaining > 0) {
+        const auto batchKeys = std::min(kDefaultBatchSize, sizeRemaining);
+        std::vector<std::string> keys(batchKeys);
+        std::generate_n(
+            keys.begin(), batchKeys, [&] { return std::to_string(key++); });
+        const auto& batchRemoteInfos = store.multi_get(keys);
+        remoteRankInfos.insert(
+            remoteRankInfos.end(),
+            batchRemoteInfos.begin(),
+            batchRemoteInfos.end());
+        sizeRemaining -= batchKeys;
+      }
+    } else {
+      std::generate_n(std::back_inserter(remoteRankInfos), size, [&] {
+        const auto& keyStr = std::to_string(key++);
+        store.wait({keyStr.c_str()}, getTimeout());
+        return store.get(keyStr);
+      });
     }
-  } else {
-    std::generate_n(std::back_inserter(remoteRankInfos), size, [&] {
-      const auto& keyStr = std::to_string(key++);
-      store.wait({keyStr.c_str()}, getTimeout());
-      return store.get(keyStr);
-    });
-  }
 
-  // Connect every pair
-  for (int i = 0; i < size; i++) {
-    if (i == rank) {
-      // at this point we have enumerated all the ranks located on this host
-      // up to the current rank, so the current `localRank` number is
-      // what we'll set to the pairs.
-      localRankSet = true;
-      // We are not going to connect self.
-      continue;
-    }
+    // Connect every pair
+    for (int i = 0; i < size; i++) {
+      if (i == rank) {
+        // at this point we have enumerated all the ranks located on this host
+        // up to the current rank, so the current `localRank` number is
+        // what we'll set to the pairs.
+        localRankSet = true;
+        // We are not going to connect self.
+        continue;
+      }
 
-    Rank remoteRankInfo(remoteRankInfos[i]);
+      Rank remoteRankInfo(remoteRankInfos[i]);
 
-    if (!localRankSet && remoteRankInfo.hostname == localHostName) {
-      ++localRank;
-    }
+      if (!localRankSet && remoteRankInfo.hostname == localHostName) {
+        ++localRank;
+      }
 
-    const auto& pair = getPair(i);
-    auto remoteDeviceAddr = Address(remoteRankInfo.addressBytes).getSockaddr();
-    auto remoteAddr = Address(
-        remoteDeviceAddr,
-        useRankAsSeqNum ? (sequence_number_t)rank
-                        : remoteRankInfo.pairIdentifiers[rank]);
-    pair->connect(remoteAddr.bytes());
-  }
+      const auto& pair = getPair(i);
+      auto remoteDeviceAddr =
+    Address(remoteRankInfo.addressBytes).getSockaddr(); auto remoteAddr =
+    Address( remoteDeviceAddr, useRankAsSeqNum ? (sequence_number_t)rank :
+    remoteRankInfo.pairIdentifiers[rank]); pair->connect(remoteAddr.bytes());
+    }
 
-  // Set the local rank info for all mesh pairs involving current rank
-  for (int i = 0; i < size; i++) {
-    if (i == rank) {
-      continue;
+    // Set the local rank info for all mesh pairs involving current rank
+    for (int i = 0; i < size; i++) {
+      if (i == rank) {
+        continue;
+      }
+      const auto& pair = getPair(i);
+      pair->setLocalRank(localRank);
     }
-    const auto& pair = getPair(i);
-    pair->setLocalRank(localRank);
-  }
+    */
 
   printConnectivityInfo();
 }
 
+std::unique_ptr<transport::Pair>& Context::getPair(int rank) {
+  auto& pair = pairs_[rank];
+
+  // don't connect to self
+  if (rank == this->rank) {
+    return pair;
+  }
+
+  if (!connecting_[rank]) {
+    connecting_[rank] = true;
+
+    const auto& keyStr = std::to_string(rank);
+    store_->wait({keyStr.c_str()}, getTimeout());
+    auto remoteRankInfoBytes = store_->get(keyStr);
+
+    Rank remoteRankInfo(remoteRankInfoBytes);
+
+    auto remoteDeviceAddr = Address(remoteRankInfo.addressBytes).getSockaddr();
+    auto remoteAddr = Address(remoteDeviceAddr, this->rank);
+    pair->connect(remoteAddr.bytes());
+  }
+  return pair;
+}
+
 std::unique_ptr<transport::Pair>& Context::createPair(int rank) {
   pairs_[rank] = std::unique_ptr<transport::Pair>(
       new tcp::Pair(this, device_.get(), rank, getTimeout(), false));

gloo/transport/tcp/context.h

@@ -43,12 +43,16 @@ class Context : public ::gloo::transport::Context,
       int rank,
       bool useRankAsSeqNumber);
 
+  virtual std::unique_ptr<transport::Pair>& getPair(int rank) override;
+
   std::unique_ptr<transport::UnboundBuffer> createUnboundBuffer(
       void* ptr,
       size_t size) override;
 
  protected:
   std::shared_ptr<Device> device_;
+  std::shared_ptr<IStore> store_{nullptr};
+  std::vector<bool> connecting_;
 
   using pendingRecvTuple = std::tuple<
       WeakNonOwningPtr<UnboundBuffer>,

gloo/transport/tcp/device.cc

@@ -13,6 +13,7 @@
 #include <netinet/in.h>
 #include <string.h>
 #include <array>
+#include <iostream>
 
 #include "gloo/common/error.h"
 #include "gloo/common/linux.h"
@@ -217,7 +218,9 @@ Device::Device(const struct attr& attr)
       interfaceSpeedMbps_(getInterfaceSpeedByName(interfaceName_)),
       pciBusID_(interfaceToBusID(interfaceName_)) {}
 
-Device::~Device() {}
+Device::~Device() {
+  loop_->shutdown();
+}
 
 std::string Device::str() const {
   std::stringstream ss;
@@ -328,6 +331,10 @@ void Device::connectAsListener(
   listener_->waitForConnection(local.getSeq(), std::move(fn));
 }
 
+void Device::cancelConnect(const Address& local) {
+  listener_->cancelConnect(local.getSeq());
+}
+
 // Connecting as initiator is active.
 //
 // The connect callback is fired when the connection to the other side

gloo/transport/tcp/device.h

@@ -101,6 +101,8 @@ class Device : public ::gloo::transport::Device,
       std::chrono::milliseconds timeout,
       connect_callback_t fn);
 
+  void cancelConnect(const Address& local);
+
   void connectAsListener(
       const Address& local,
       std::chrono::milliseconds timeout,

gloo/transport/tcp/helpers.h

@@ -45,13 +45,19 @@ class ReadValueOperation final
         fn_(std::move(fn)) {}
 
   void run() {
+    auto loop = loop_.lock();
+    if (!loop) {
+      return;
+    }
+
     // Cannot initialize leak until after the object has been
     // constructed, because the std::make_shared initialization
     // doesn't run after construction of the underlying object.
     leak_ = this->shared_from_this();
+
     // Register with loop only after we've leaked the shared_ptr,
     // because we unleak it when the event loop thread calls.
-    loop_->registerDescriptor(socket_->fd(), EPOLLIN | EPOLLONESHOT, this);
+    loop->registerDescriptor(socket_->fd(), EPOLLIN | EPOLLONESHOT, this);
   }
 
   void handleEvents(int events) override {
@@ -80,7 +86,7 @@ class ReadValueOperation final
   }
 
  private:
-  std::shared_ptr<Loop> loop_;
+  std::weak_ptr<Loop> loop_;
   std::shared_ptr<Socket> socket_;
   callback_t fn_;
   std::shared_ptr<ReadValueOperation<T>> leak_;

gloo/transport/tcp/listener.cc

@@ -35,6 +35,7 @@ Listener::Listener(std::shared_ptr<Loop> loop, const attr& attr)
 }
 
 Listener::~Listener() {
+  *closed_ = true;
   if (listener_) {
     loop_->unregisterDescriptor(listener_->fd(), this);
   }
@@ -61,7 +62,7 @@ void Listener::handleEvents(int /* unused */) {
     read<sequence_number_t>(
         loop_,
         sock,
-        [this](
+        [this, closed = closed_](
             std::shared_ptr<Socket> socket,
             const Error& error,
             sequence_number_t&& seq) {
@@ -72,6 +73,10 @@ void Listener::handleEvents(int /* unused */) {
             return;
           }
 
+          if (*closed) {
+            return;
+          }
+
           haveConnection(std::move(socket), seq);
         });
   }
@@ -108,6 +113,19 @@ void Listener::waitForConnection(sequence_number_t seq, connect_callback_t fn) {
   loop_->defer([fn, socket]() { fn(socket, Error::kSuccess); });
 }
 
+void Listener::cancelConnect(sequence_number_t seq) {
+  std::unique_lock<std::mutex> lock(mutex_);
+
+  // If we don't yet have a callback for this sequence number, do nothing.
+  auto it = seqToCallback_.find(seq);
+  if (it == seqToCallback_.end()) {
+    return;
+  }
+
+  // If we already have a callback for this sequence number, cancel it.
+  seqToCallback_.erase(it);
+}
+
 void Listener::haveConnection(
     std::shared_ptr<Socket> socket,
     sequence_number_t seq) {

gloo/transport/tcp/listener.h

@@ -49,10 +49,14 @@ class Listener final : public Handler {
   // even if the connection is already available.
   void waitForConnection(sequence_number_t seq, connect_callback_t fn);
 
+  void cancelConnect(sequence_number_t seq);
+
  private:
   std::mutex mutex_;
   std::shared_ptr<Loop> loop_;
   std::shared_ptr<Socket> listener_;
+  std::shared_ptr<std::atomic<bool>> closed_{
+      std::make_shared<std::atomic<bool>>(false)};
 
   // Address of this listener and the sequence number for the next
   // connection. Sequence numbers are written by a peer right after