p2p: fix hanging shutdown

4 files changed, 50 insertions, 19 deletions

diff --git a/contrib/epee/include/net/abstract_tcp_server2.h b/contrib/epee/include/net/abstract_tcp_server2.h
index 33e50cd35..fe201f155 100644
--- a/contrib/epee/include/net/abstract_tcp_server2.h
+++ b/contrib/epee/include/net/abstract_tcp_server2.h
@@ -128,7 +128,7 @@ namespace net_utils
 
     void start_handshake();
     void start_read();
-    void finish_read(size_t bytes_transferred);
+    void handle_read(size_t bytes_transferred);
     void start_write();
     void start_shutdown();
     void cancel_socket();
@@ -326,7 +326,7 @@ namespace net_utils
     //----------------- i_service_endpoint ---------------------
     virtual bool do_send(byte_slice message); ///< (see do_send from i_service_endpoint)
     virtual bool send_done();
-    virtual bool close();
+    virtual bool close(const bool wait_for_shutdown);
     virtual bool call_run_once_service_io();
     virtual bool request_callback();
     virtual io_context_t& get_io_context();
@@ -379,7 +379,7 @@ namespace net_utils
     bool timed_wait_server_stop(uint64_t wait_mseconds);
 
     /// Stop the server.
-    void send_stop_signal();
+    void send_stop_signal(std::function<void()> close_all_connections = [](){});
 
     bool is_stop_signal_sent() const noexcept { return m_stop_signal_sent; };
 
diff --git a/contrib/epee/include/net/abstract_tcp_server2.inl b/contrib/epee/include/net/abstract_tcp_server2.inl
index 814ad91b1..60cc1b55f 100644
--- a/contrib/epee/include/net/abstract_tcp_server2.inl
+++ b/contrib/epee/include/net/abstract_tcp_server2.inl
@@ -34,6 +34,7 @@
 #include <boost/asio/post.hpp>
 #include <boost/foreach.hpp>
 #include <boost/uuid/random_generator.hpp>
+#include <boost/uuid/uuid_io.hpp>
 #include <boost/chrono.hpp>
 #include <boost/utility/value_init.hpp>
 #include <boost/asio/bind_executor.hpp>
@@ -245,7 +246,7 @@ namespace net_utils
               )
             ) {
               m_state.ssl.enabled = false;
-              finish_read(bytes_transferred);
+              handle_read(bytes_transferred);
             }
             else {
               m_state.ssl.detected = true;
@@ -383,7 +384,7 @@ namespace net_utils
           m_conn_context.m_recv_cnt += bytes_transferred;
           start_timer(get_timeout_from_bytes_read(bytes_transferred), true);
         }
-        finish_read(bytes_transferred);
+        handle_read(bytes_transferred);
       }
     };
     if (!m_state.ssl.enabled)
@@ -410,7 +411,7 @@ namespace net_utils
   }
 
   template<typename T>
-  void connection<T>::finish_read(size_t bytes_transferred)
+  void connection<T>::handle_read(size_t bytes_transferred)
   {
     // Post handle_recv to a separate `strand_`, distinct from `m_strand`
     // which is listening for reads/writes. This avoids a circular dep.
@@ -737,6 +738,7 @@ namespace net_utils
     }
     else
       m_state.status = status_t::WASTED;
+    m_state.condition.notify_all();
   }
 
   template<typename T>
@@ -795,6 +797,7 @@ namespace net_utils
       m_state.socket.connected = false;
     }
     m_state.status = status_t::WASTED;
+    m_state.condition.notify_all();
   }
 
   template<typename T>
@@ -1113,7 +1116,7 @@ namespace net_utils
   template<typename T>
   bool connection<T>::cancel()
   {
-    return close();
+    return close(false);
   }
 
   template<typename T>
@@ -1129,12 +1132,34 @@ namespace net_utils
   }
 
   template<typename T>
-  bool connection<T>::close()
+  bool connection<T>::close(const bool wait_for_shutdown)
   {
     std::lock_guard<std::mutex> guard(m_state.lock);
     if (m_state.status != status_t::RUNNING)
       return false;
     terminate_async();
+
+    // Sometimes we do *not* want to wait for the connection to shut down because for example handle_recv might try to
+    // close the connection when handling a request. But handle_recv can't complete the shutdown sequence because
+    // handle_read is set to true. So, in that case we call terminate_async and return here.
+    if (!wait_for_shutdown)
+      return true;
+
+    // Sometimes we *do* want to wait for the connection to shut down for example when stopping the server. When
+    // stopping the server, we don't want the io_context to stop before the shutdown sequence completes, since we
+    // execute terminate inside m_strand. So we wait for the connection's shutdown sequence to complete before stopping
+    // the io_context.
+    MDEBUG("Waiting for connection " << m_conn_context.m_connection_id << " to shutdown, current state: " << m_state.status);
+    m_state.condition.wait(
+      m_state.lock,
+      [this]{
+        return (
+          m_state.status == status_t::TERMINATED || m_state.status == status_t::WASTED
+        );
+      }
+    );
+    MDEBUG("Shut down connection " << m_conn_context.m_connection_id);
+
     return true;
   }
 
@@ -1541,7 +1566,7 @@ namespace net_utils
   }
   //---------------------------------------------------------------------------------
   template<class t_protocol_handler>
-  void boosted_tcp_server<t_protocol_handler>::send_stop_signal()
+  void boosted_tcp_server<t_protocol_handler>::send_stop_signal(std::function<void()> close_all_connections)
   {
     m_stop_signal_sent = true;
     typename connection<t_protocol_handler>::shared_state *state = static_cast<typename connection<t_protocol_handler>::shared_state*>(m_state.get());
@@ -1554,7 +1579,13 @@ namespace net_utils
     }
     connections_.clear();
     connections_mutex.unlock();
+
+    // Since we shut down connections in the strand, we want to make sure to complete the shutdown sequence before
+    // stopping the io_context. We let the caller handle closing because the caller is the one keeping track of all
+    // connections (connections_ is only a subset of all connections).
+    close_all_connections();
     io_context_.stop();
+    MDEBUG("Done with send_stop_signal");
     CATCH_ENTRY_L0("boosted_tcp_server<t_protocol_handler>::send_stop_signal()", void());
   }
   //---------------------------------------------------------------------------------
diff --git a/contrib/epee/include/net/levin_protocol_handler_async.h b/contrib/epee/include/net/levin_protocol_handler_async.h
index 756560c2a..73d46fd8b 100644
--- a/contrib/epee/include/net/levin_protocol_handler_async.h
+++ b/contrib/epee/include/net/levin_protocol_handler_async.h
@@ -107,7 +107,7 @@ public:
   int invoke_async(int command, message_writer in_msg, boost::uuids::uuid connection_id, const callback_t &cb, size_t timeout = LEVIN_DEFAULT_TIMEOUT_PRECONFIGURED);
 
   int send(epee::byte_slice message, const boost::uuids::uuid& connection_id);
-  bool close(boost::uuids::uuid connection_id);
+  bool close(boost::uuids::uuid connection_id, const bool wait_for_shutdown);
   bool update_connection_context(const t_connection_context& contxt);
   bool request_callback(boost::uuids::uuid connection_id);
   template<class callback_t>
@@ -121,7 +121,7 @@ public:
 
   async_protocol_handler_config():m_pcommands_handler(NULL), m_pcommands_handler_destroy(NULL), m_initial_max_packet_size(LEVIN_INITIAL_MAX_PACKET_SIZE), m_max_packet_size(LEVIN_DEFAULT_MAX_PACKET_SIZE), m_invoke_timeout(LEVIN_DEFAULT_TIMEOUT_PRECONFIGURED)
   {}
-  ~async_protocol_handler_config() { set_handler(NULL, NULL); }
+  virtual ~async_protocol_handler_config() { set_handler(NULL, NULL); }
   void del_out_connections(size_t count);
   void del_in_connections(size_t count);
 };
@@ -214,7 +214,7 @@ public:
           MINFO(con.get_context_ref() << "Timeout on invoke operation happened, command: " << command << " timeout: " << timeout);
           epee::span<const uint8_t> fake;
           cb(LEVIN_ERROR_CONNECTION_TIMEDOUT, fake, con.get_context_ref());
-          con.close();
+          con.close(false);
           con.finish_outer_call();
         });
         m_timer_started = true;
@@ -278,7 +278,7 @@ public:
           MINFO(con.get_context_ref() << "Timeout on invoke operation happened, command: " << command << " timeout: " << timeout);
           epee::span<const uint8_t> fake;
           cb(LEVIN_ERROR_CONNECTION_TIMEDOUT, fake, con.get_context_ref());
-          con.close();
+          con.close(false);
           con.finish_outer_call();
         });
       }
@@ -379,11 +379,11 @@ public:
     return true;
   }
   
-  bool close()
+  bool close(const bool wait_for_shutdown)
   {    
     ++m_close_called;
 
-    m_pservice_endpoint->close();
+    m_pservice_endpoint->close(wait_for_shutdown);
     return true;
   }
 
@@ -792,7 +792,7 @@ void async_protocol_handler_config<t_connection_context>::delete_connections(siz
   CRITICAL_REGION_END();
 
   for (size_t i = 0; i < connections.size() && i < count; ++i)
-    connections[i]->close();
+    connections[i]->close(false);
 }
 //------------------------------------------------------------------------------------------
 template<class t_connection_context>
@@ -935,14 +935,14 @@ int async_protocol_handler_config<t_connection_context>::send(byte_slice message
 }
 //------------------------------------------------------------------------------------------
 template<class t_connection_context>
-bool async_protocol_handler_config<t_connection_context>::close(boost::uuids::uuid connection_id)
+bool async_protocol_handler_config<t_connection_context>::close(boost::uuids::uuid connection_id, const bool wait_for_shutdown)
 {
   async_protocol_handler<t_connection_context>* aph = nullptr;
   if (find_and_lock_connection(connection_id, aph) != LEVIN_OK)
     return false;
   auto scope_exit_handler = misc_utils::create_scope_leave_handler(
     boost::bind(&async_protocol_handler<t_connection_context>::finish_outer_call, aph));
-  if (!aph->close())
+  if (!aph->close(wait_for_shutdown))
     return false;
   CRITICAL_REGION_LOCAL(m_connects_lock);
   m_connects.erase(connection_id);
diff --git a/contrib/epee/include/net/net_utils_base.h b/contrib/epee/include/net/net_utils_base.h
index 2cf6e795d..88b50f602 100644
--- a/contrib/epee/include/net/net_utils_base.h
+++ b/contrib/epee/include/net/net_utils_base.h
@@ -441,7 +441,7 @@ namespace net_utils
 	struct i_service_endpoint
 	{
 		virtual bool do_send(byte_slice message)=0;
-    virtual bool close()=0;
+    virtual bool close(const bool wait_for_shutdown)=0;
     virtual bool send_done()=0;
     virtual bool call_run_once_service_io()=0;
     virtual bool request_callback()=0;