gloo/transport/tcp/loop: better memory management and shutdown (#422)

Summary:
Pull Request resolved: #422

This overhauls how we schedule async operations on the tcp Loop.

These changes greatly simplifies the memory management and simplifies the shutdown logic since memory ownership is better tied to epoll usage and relies less on shared_ptr copies everywhere to work around it.

This makes a few major changes:

* adds a new `registerDescriptor` method to Loop that provides a shared_ptr. This method will track the lifetime of the provided Handler and will release it when the epoll handle is updated/deleted.
* Make `Handler::handleEvents` signature provide a `Loop&` reference so all objects don't need to hold a shared_ptr to the loop.
* Update `helpers.h` to not manage it's own lifetime (instead using registerDescriptor shared_ptr support) and also not hold a shared_ptr to `Loop`.
* Introduces a new `shutdown()` method that allows the loop to gracefully shutdown all operations prior to destruction. This simplifies teardown on error since there a few places where handlers (Device) are managed with raw pointers. This isn't used outside of Loop currently but will be in a follow up PR.

Differential Revision: D72080276

Author: d4l3k

gloo/test/tcp_test.cc

@@ -17,15 +17,15 @@ TEST(TcpTest, ConnectTimeout) {
   // Use bad address
   auto remote = Address("::1", 10);
   auto timeout = std::chrono::milliseconds(100);
-  auto fn = [&](std::shared_ptr<Socket>, const Error& e) {
+  auto fn = [&](Loop&, std::shared_ptr<Socket>, const Error& e) {
     std::lock_guard<std::mutex> lock(m);
     done = true;
     cv.notify_all();
 
     EXPECT_TRUE(e);
     EXPECT_TRUE(dynamic_cast<const TimeoutError*>(&e));
   };
-  connectLoop(loop, remote, 0, 5, timeout, std::move(fn));
+  connectLoop(*loop, remote, 0, 5, timeout, std::move(fn));
 
   std::unique_lock<std::mutex> lock(m);
   cv.wait(lock, [&] { return done; });

gloo/transport/tcp/device.cc

@@ -341,11 +341,12 @@ void Device::connectAsInitiator(
     const int size,
     std::chrono::milliseconds timeout,
     connect_callback_t fn) {
-  auto writeSeq = [loop = loop_, seq = remote.getSeq()](
-                      std::shared_ptr<Socket> socket, connect_callback_t fn) {
-    // Write sequence number for peer to new socket.
-    write<sequence_number_t>(loop, std::move(socket), seq, std::move(fn));
-  };
+  auto writeSeq =
+      [seq = remote.getSeq()](
+          Loop& loop, std::shared_ptr<Socket> socket, connect_callback_t fn) {
+        // Write sequence number for peer to new socket.
+        write<sequence_number_t>(loop, std::move(socket), seq, std::move(fn));
+      };
 
   if (disableConnectionRetries()) {
     const auto& sockaddr = remote.getSockaddr();
@@ -356,22 +357,22 @@ void Device::connectAsInitiator(
     socket->noDelay(true);
     socket->connect(sockaddr);
 
-    writeSeq(std::move(socket), std::move(fn));
+    writeSeq(*loop_, std::move(socket), std::move(fn));
   } else {
     connectLoop(
-        loop_,
+        *loop_,
         remote,
         rank,
         size,
         timeout,
-        [loop = loop_, fn = std::move(fn), writeSeq = std::move(writeSeq)](
-            std::shared_ptr<Socket> socket, const Error& error) {
+        [fn = std::move(fn), writeSeq = std::move(writeSeq)](
+            Loop& loop, std::shared_ptr<Socket> socket, const Error& error) {
           if (error) {
             fn(socket, error);
             return;
           }
 
-          writeSeq(std::move(socket), std::move(fn));
+          writeSeq(loop, std::move(socket), fn);
         });
   }
 }

gloo/transport/tcp/helpers.cc

@@ -5,15 +5,15 @@ namespace transport {
 namespace tcp {
 
 void connectLoop(
-    std::shared_ptr<Loop> loop,
+    Loop& loop,
     const Address& remote,
     const int rank,
     const int size,
     std::chrono::milliseconds timeout,
     typename ConnectOperation::callback_t fn) {
   auto x = std::make_shared<ConnectOperation>(
-      std::move(loop), remote, rank, size, timeout, std::move(fn));
-  x->run();
+      remote, rank, size, timeout, std::move(fn));
+  x->run(loop);
 }
 
 } // namespace tcp

gloo/transport/tcp/helpers.h

@@ -26,8 +26,6 @@ namespace tcp {
 // socket specified at construction. Upon completion or error, the
 // callback is called. Its lifetime is coupled with completion of the
 // operation, so the called doesn't need to hold on to the instance.
-// It does so by storing a shared_ptr to itself (effectively a leak)
-// until the event loop calls back.
 template <typename T>
 class ReadValueOperation final
     : public Handler,
@@ -36,29 +34,15 @@ class ReadValueOperation final
   using callback_t =
       std::function<void(std::shared_ptr<Socket>, const Error& error, T&& t)>;
 
-  ReadValueOperation(
-      std::shared_ptr<Loop> loop,
-      std::shared_ptr<Socket> socket,
-      callback_t fn)
-      : loop_(std::move(loop)),
-        socket_(std::move(socket)),
-        fn_(std::move(fn)) {}
+  ReadValueOperation(std::shared_ptr<Socket> socket, callback_t fn)
+      : socket_(std::move(socket)), fn_(std::move(fn)) {}
 
-  void run() {
-    // 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);
+  void run(Loop& loop) {
+    loop.registerDescriptor(
+        socket_->fd(), EPOLLIN | EPOLLONESHOT, this->shared_from_this());
   }
 
-  void handleEvents(int events) override {
-    // Move leaked shared_ptr to the stack so that this object
-    // destroys itself once this function returns.
-    auto self = std::move(this->leak_);
-
+  void handleEvents(Loop&, int /*events*/) override {
     // Read T.
     auto rv = socket_->read(&t_, sizeof(t_));
     if (rv == -1) {
@@ -80,30 +64,26 @@ class ReadValueOperation final
   }
 
  private:
-  std::shared_ptr<Loop> loop_;
   std::shared_ptr<Socket> socket_;
   callback_t fn_;
-  std::shared_ptr<ReadValueOperation<T>> leak_;
 
   T t_;
 };
 
 template <typename T>
 void read(
-    std::shared_ptr<Loop> loop,
+    Loop& loop,
     std::shared_ptr<Socket> socket,
     typename ReadValueOperation<T>::callback_t fn) {
-  auto x = std::make_shared<ReadValueOperation<T>>(
-      std::move(loop), std::move(socket), std::move(fn));
-  x->run();
+  auto x =
+      std::make_shared<ReadValueOperation<T>>(std::move(socket), std::move(fn));
+  x->run(loop);
 }
 
 // WriteValueOperation asynchronously writes a value of type T to the
 // socket specified at construction. Upon completion or error, the
 // callback is called. Its lifetime is coupled with completion of the
 // operation, so the called doesn't need to hold on to the instance.
-// It does so by storing a shared_ptr to itself (effectively a leak)
-// until the event loop calls back.
 template <typename T>
 class WriteValueOperation final
     : public Handler,
@@ -112,31 +92,15 @@ class WriteValueOperation final
   using callback_t =
       std::function<void(std::shared_ptr<Socket>, const Error& error)>;
 
-  WriteValueOperation(
-      std::shared_ptr<Loop> loop,
-      std::shared_ptr<Socket> socket,
-      T t,
-      callback_t fn)
-      : loop_(std::move(loop)),
-        socket_(std::move(socket)),
-        fn_(std::move(fn)),
-        t_(std::move(t)) {}
-
-  void run() {
-    // 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(), EPOLLOUT | EPOLLONESHOT, this);
-  }
+  WriteValueOperation(std::shared_ptr<Socket> socket, T t, callback_t fn)
+      : socket_(std::move(socket)), fn_(std::move(fn)), t_(std::move(t)) {}
 
-  void handleEvents(int events) override {
-    // Move leaked shared_ptr to the stack so that this object
-    // destroys itself once this function returns.
-    auto leak = std::move(this->leak_);
+  void run(Loop& loop) {
+    loop.registerDescriptor(
+        socket_->fd(), EPOLLOUT | EPOLLONESHOT, this->shared_from_this());
+  }
 
+  void handleEvents(Loop&, int /*events*/) override {
     // Write T.
     auto rv = socket_->write(&t_, sizeof(t_));
     if (rv == -1) {
@@ -154,33 +118,30 @@ class WriteValueOperation final
   }
 
  private:
-  std::shared_ptr<Loop> loop_;
   std::shared_ptr<Socket> socket_;
   callback_t fn_;
-  std::shared_ptr<WriteValueOperation<T>> leak_;
 
   T t_;
 };
 
 template <typename T>
 void write(
-    std::shared_ptr<Loop> loop,
+    Loop& loop,
     std::shared_ptr<Socket> socket,
     T t,
     typename WriteValueOperation<T>::callback_t fn) {
   auto x = std::make_shared<WriteValueOperation<T>>(
-      std::move(loop), std::move(socket), std::move(t), std::move(fn));
-  x->run();
+      std::move(socket), std::move(t), std::move(fn));
+  x->run(loop);
 }
 
 class ConnectOperation final
     : public Handler,
       public std::enable_shared_from_this<ConnectOperation> {
  public:
-  using callback_t =
-      std::function<void(std::shared_ptr<Socket>, const Error& error)>;
+  using callback_t = std::function<
+      void(Loop& loop, std::shared_ptr<Socket>, const Error& error)>;
   ConnectOperation(
-      std::shared_ptr<Loop> loop,
       const Address& remote,
       const int rank,
       const int size,
@@ -190,15 +151,9 @@ class ConnectOperation final
         rank_(rank),
         size_(size),
         deadline_(std::chrono::steady_clock::now() + timeout),
-        loop_(std::move(loop)),
         fn_(std::move(fn)) {}
 
-  void run() {
-    // 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();
-
+  void run(Loop& loop) {
     const auto& sockaddr = remote_.getSockaddr();
 
     // Create new socket to connect to peer.
@@ -207,29 +162,26 @@ class ConnectOperation final
     socket_->noDelay(true);
     socket_->connect(sockaddr);
 
-    // Register with loop only after we've leaked the shared_ptr,
-    // because we unleak it when the event loop thread calls.
     // Register for EPOLLOUT, because we want to be notified when
     // the connect completes. EPOLLERR is also necessary because
     // connect() can fail.
-    if (auto loop = loop_.lock()) {
-      loop->registerDescriptor(
-          socket_->fd(), EPOLLOUT | EPOLLERR | EPOLLONESHOT, this);
-    } else {
-      fn_(socket_, LoopError("loop is gone"));
-    }
+    loop.registerDescriptor(
+        socket_->fd(),
+        EPOLLOUT | EPOLLERR | EPOLLONESHOT,
+        this->shared_from_this());
   }
 
-  void handleEvents(int events) override {
-    // Move leaked shared_ptr to the stack so that this object
-    // destroys itself once this function returns.
-    auto leak = std::move(this->leak_);
+  void handleEvents(Loop& loop, int /*events*/) override {
+    // Hold a reference to this object to keep it alive until the
+    // callback is called.
+    auto leak = shared_from_this();
+    loop.unregisterDescriptor(socket_->fd(), this);
 
     int result;
     socklen_t result_len = sizeof(result);
     if (getsockopt(socket_->fd(), SOL_SOCKET, SO_ERROR, &result, &result_len) <
         0) {
-      fn_(socket_, SystemError("getsockopt", errno, remote_));
+      fn_(loop, socket_, SystemError("getsockopt", errno, remote_));
       return;
     }
     if (result != 0) {
@@ -248,16 +200,18 @@ class ConnectOperation final
           socket_->sockName().str(),
       };
       DebugLogger::log(debugData);
+
       // check deadline
       if (willRetry) {
-        run();
+        run(loop);
       } else {
-        fn_(socket_, TimeoutError("timed out connecting: " + e.what()));
+        fn_(loop, socket_, TimeoutError("timed out connecting: " + e.what()));
       }
+
       return;
     }
 
-    fn_(socket_, Error::kSuccess);
+    fn_(loop, socket_, Error::kSuccess);
   }
 
  private:
@@ -269,16 +223,12 @@ class ConnectOperation final
 
   int retry_{0};
 
-  // We use a weak_ptr to the loop to avoid a reference cycle when an error
-  // occurs.
-  std::weak_ptr<Loop> loop_;
   std::shared_ptr<Socket> socket_;
   callback_t fn_;
-  std::shared_ptr<ConnectOperation> leak_;
 };
 
 void connectLoop(
-    std::shared_ptr<Loop> loop,
+    Loop& loop,
     const Address& remote,
     const int rank,
     const int size,

gloo/transport/tcp/listener.cc

@@ -40,7 +40,7 @@ Listener::~Listener() {
   }
 }
 
-void Listener::handleEvents(int /* unused */) {
+void Listener::handleEvents(Loop& loop, int /* unused */) {
   std::lock_guard<std::mutex> guard(mutex_);
 
   for (;;) {
@@ -59,7 +59,7 @@ void Listener::handleEvents(int /* unused */) {
 
     // Read sequence number.
     read<sequence_number_t>(
-        loop_,
+        loop,
         sock,
         [this](
             std::shared_ptr<Socket> socket,

gloo/transport/tcp/listener.h

@@ -38,7 +38,7 @@ class Listener final : public Handler {
 
   ~Listener() override;
 
-  void handleEvents(int events) override;
+  void handleEvents(Loop& loop, int events) override;
 
   Address nextAddress();