feat: add scheduler support for timeout management in Multiselect (#312)

* feat: add scheduler support for timeout management in Multiselect

* fix: pass scheduler to Multiselect for improved timeout management

Author: kamilsa

include/libp2p/protocol_muxer/multiselect.hpp

@@ -11,6 +11,10 @@
 
 #include "protocol_muxer.hpp"
 
+namespace libp2p::basic {
+  class Scheduler;
+}
+
 namespace libp2p::protocol_muxer::multiselect {
 
   class MultiselectInstance;
@@ -20,6 +24,8 @@ namespace libp2p::protocol_muxer::multiselect {
    public:
     using Instance = std::shared_ptr<MultiselectInstance>;
 
+    explicit Multiselect(std::shared_ptr<basic::Scheduler> scheduler);
+
     ~Multiselect() override = default;
 
     /// Implements ProtocolMuxer API
@@ -45,6 +51,9 @@ namespace libp2p::protocol_muxer::multiselect {
     /// Returns instance either from cache or creates a new one
     Instance getInstance();
 
+    /// Scheduler for timeout management
+    std::shared_ptr<basic::Scheduler> scheduler_;
+
     /// Active instances, keep them here to hold shared ptrs alive
     std::unordered_set<Instance> active_instances_;
 

include/libp2p/protocol_muxer/multiselect/multiselect_instance.hpp

@@ -6,13 +6,18 @@
 
 #pragma once
 
+#include <libp2p/basic/scheduler.hpp>
 #include <libp2p/protocol_muxer/multiselect.hpp>
 #include "parser.hpp"
 
 namespace soralog {
   class Logger;
 }
 
+namespace libp2p::basic {
+  class Scheduler;
+}
+
 namespace libp2p::protocol_muxer::multiselect {
 
   class Multiselect;
@@ -21,7 +26,8 @@ namespace libp2p::protocol_muxer::multiselect {
   class MultiselectInstance
       : public std::enable_shared_from_this<MultiselectInstance> {
    public:
-    explicit MultiselectInstance(Multiselect &owner);
+    explicit MultiselectInstance(Multiselect &owner,
+                                 std::shared_ptr<basic::Scheduler> scheduler);
 
     /// Implements ProtocolMuxer API
     void selectOneOf(std::span<const peer::ProtocolName> protocols,
@@ -74,6 +80,9 @@ namespace libp2p::protocol_muxer::multiselect {
     /// Handles "na" reply, client-specific
     MaybeResult handleNA();
 
+    /// Handles timeout when negotiation takes too long
+    void onTimeout();
+
     /// Owner of this object, needed for reuse of instances
     Multiselect &owner_;
 
@@ -125,6 +134,12 @@ namespace libp2p::protocol_muxer::multiselect {
 
     /// Cache: serialized NA response
     boost::optional<Packet> na_response_;
+
+    /// Scheduler for timeout handling
+    std::shared_ptr<basic::Scheduler> scheduler_;
+
+    /// Timeout handle for negotiation timeout
+    basic::Scheduler::Handle timeout_handle_;
   };
 
 }  // namespace libp2p::protocol_muxer::multiselect

src/protocol_muxer/multiselect.cpp

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <libp2p/basic/scheduler.hpp>
 #include <libp2p/log/logger.hpp>
 #include <libp2p/protocol_muxer/multiselect/multiselect_instance.hpp>
 #include <libp2p/protocol_muxer/multiselect/simple_stream_negotiate.hpp>
@@ -21,6 +22,9 @@ namespace libp2p::protocol_muxer::multiselect {
     constexpr size_t kMaxCacheSize = 8;
   }  // namespace
 
+  Multiselect::Multiselect(std::shared_ptr<basic::Scheduler> scheduler)
+      : scheduler_(std::move(scheduler)) {}
+
   void Multiselect::selectOneOf(std::span<const peer::ProtocolName> protocols,
                                 std::shared_ptr<basic::ReadWriter> connection,
                                 bool is_initiator,
@@ -62,7 +66,7 @@ namespace libp2p::protocol_muxer::multiselect {
   Multiselect::Instance Multiselect::getInstance() {
     Instance instance;
     if (cache_.empty()) {
-      instance = std::make_shared<MultiselectInstance>(*this);
+      instance = std::make_shared<MultiselectInstance>(*this, scheduler_);
     } else {
       SL_TRACE(log(),
                "cache: {}->{}, active {}->{}",

src/protocol_muxer/multiselect/multiselect_instance.cpp

@@ -7,7 +7,9 @@
 #include <libp2p/protocol_muxer/multiselect/multiselect_instance.hpp>
 
 #include <cctype>
+#include <span>
 
+#include <libp2p/basic/scheduler.hpp>
 #include <libp2p/basic/write_return_size.hpp>
 #include <libp2p/common/trace.hpp>
 #include <libp2p/protocol_muxer/multiselect/serializing.hpp>
@@ -20,10 +22,14 @@ namespace libp2p::protocol_muxer::multiselect {
       static log::Logger logger = log::createLogger("Multiselect");
       return logger;
     }
+
+    /// Timeout for protocol negotiation (5 seconds)
+    constexpr std::chrono::milliseconds kNegotiationTimeout{5000};
   }  // namespace
 
-  MultiselectInstance::MultiselectInstance(Multiselect &owner)
-      : owner_(owner) {}
+  MultiselectInstance::MultiselectInstance(
+      Multiselect &owner, std::shared_ptr<basic::Scheduler> scheduler)
+      : owner_(owner), scheduler_(std::move(scheduler)) {}
 
   void MultiselectInstance::selectOneOf(
       std::span<const peer::ProtocolName> protocols,
@@ -62,6 +68,15 @@ namespace libp2p::protocol_muxer::multiselect {
     write_queue_.clear();
     is_writing_ = false;
 
+    // Schedule timeout for negotiation
+    timeout_handle_ = scheduler_->scheduleWithHandle(
+        [wptr = std::weak_ptr<MultiselectInstance>(shared_from_this())]() {
+          if (auto self = wptr.lock()) {
+            self->onTimeout();
+          }
+        },
+        kNegotiationTimeout);
+
     if (is_initiator_) {
       std::ignore = sendProposal();
     } else if (negotiate_multiselect) {
@@ -166,6 +181,10 @@ namespace libp2p::protocol_muxer::multiselect {
 
   void MultiselectInstance::close(outcome::result<std::string> result) {
     closed_ = true;
+
+    // Cancel timeout if it's still active
+    timeout_handle_.reset();
+
     ++current_round_;
     write_queue_.clear();
     Multiselect::ProtocolHandlerFunc callback;
@@ -338,4 +357,11 @@ namespace libp2p::protocol_muxer::multiselect {
     return MaybeResult(ProtocolMuxer::Error::PROTOCOL_VIOLATION);
   }
 
+  void MultiselectInstance::onTimeout() {
+    SL_DEBUG(log(),
+             "Protocol negotiation timeout after {}ms",
+             kNegotiationTimeout.count());
+    close(ProtocolMuxer::Error::NEGOTIATION_FAILED);
+  }
+
 }  // namespace libp2p::protocol_muxer::multiselect

test/acceptance/p2p/host/peer/test_peer.cpp

@@ -136,7 +136,7 @@ Peer::sptr<host::BasicHost> Peer::makeHost(const crypto::KeyPair &keyPair) {
       std::make_shared<peer::IdentityManagerImpl>(keyPair, key_marshaller);
 
   auto multiselect =
-      std::make_shared<protocol_muxer::multiselect::Multiselect>();
+      std::make_shared<protocol_muxer::multiselect::Multiselect>(scheduler_);
 
   auto router = std::make_shared<network::RouterImpl>();