Add ability for channels to pause/resume processing of incoming messages

nano/core_test/network.cpp

@@ -607,7 +607,7 @@ TEST (tcp_listener, tcp_listener_timeout_node_id_handshake)
 	ASSERT_TRUE (cookie);
 	nano::node_id_handshake::query_payload query{ *cookie };
 	nano::node_id_handshake node_id_handshake{ nano::dev::network_params.network, query };
-	auto channel = std::make_shared<nano::transport::channel_tcp> (*node0, socket);
+	auto channel = std::make_shared<nano::transport::channel_tcp> (*node0, socket, std::weak_ptr<nano::transport::tcp_server>{});
 	socket->async_connect (node0->tcp_listener.endpoint (), [&node_id_handshake, channel] (boost::system::error_code const & ec) {
 		ASSERT_FALSE (ec);
 		channel->send (node_id_handshake, [] (boost::system::error_code const & ec, size_t size_a) {
@@ -1125,3 +1125,31 @@ TEST (network, purge_dead_channel_remote)
 	};
 	ASSERT_TIMELY (5s, !channel_exists (node2, channel));
 }
+
+TEST (network, tcp_server_throttle)
+{
+	nano::test::system system;
+
+	auto config = system.default_config ();
+	config.message_processor.threads = 0; // Disable message processing
+	auto & node1 = *system.add_node (config);
+	// We need a second node to hijack one of the channels
+	auto & node2 = *system.add_node (config);
+
+	auto channel = node2.network.find_node_id (node1.node_id.pub);
+
+	// Up to (queue size + 1) messages should be received and deserialized, with the last one causing the channel to be throttled
+	nano::keepalive msg{ nano::dev::network_params.network };
+	for (int n = 0; n < config.message_processor.max_queue + 1; ++n)
+	{
+		channel->send (msg);
+	}
+	ASSERT_TIMELY_EQ (5s, node1.stats.count (nano::stat::type::tcp_server_message, nano::stat::detail::keepalive), config.message_processor.max_queue + 1);
+
+	// No messages should be processed
+	ASSERT_ALWAYS_EQ (1s, node1.stats.count (nano::stat::type::message, nano::stat::detail::keepalive), 0);
+
+	// Any additional messages should not be eagerly deserialized from the socket
+	channel->send (msg);
+	ASSERT_ALWAYS_EQ (1s, node1.stats.count (nano::stat::type::tcp_server_message, nano::stat::detail::keepalive), config.message_processor.max_queue + 1);
+}

nano/core_test/request_aggregator.cpp

@@ -37,7 +37,7 @@ TEST (request_aggregator, one)
 	std::vector<std::pair<nano::block_hash, nano::root>> request;
 	request.emplace_back (send1->hash (), send1->root ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	node.aggregator.request (request, dummy_channel);
 	ASSERT_TIMELY (3s, node.aggregator.empty ());
 	// Not yet in the ledger
@@ -103,7 +103,7 @@ TEST (request_aggregator, one_update)
 	std::vector<std::pair<nano::block_hash, nano::root>> request;
 	request.emplace_back (send2->hash (), send2->root ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	node.aggregator.request (request, dummy_channel);
 	request.clear ();
 	request.emplace_back (receive1->hash (), receive1->root ());
@@ -169,7 +169,7 @@ TEST (request_aggregator, two)
 	request.emplace_back (send2->hash (), send2->root ());
 	request.emplace_back (receive1->hash (), receive1->root ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	// Process both blocks
 	node.aggregator.request (request, dummy_channel);
 	// One vote should be generated for both blocks
@@ -288,7 +288,7 @@ TEST (request_aggregator, split)
 	ASSERT_TIMELY_EQ (5s, max_vbh + 2, node.ledger.cemented_count ());
 	ASSERT_EQ (max_vbh + 1, request.size ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	node.aggregator.request (request, dummy_channel);
 	// In the ledger but no vote generated yet
 	ASSERT_TIMELY_EQ (3s, 2, node.stats.count (nano::stat::type::requests, nano::stat::detail::requests_generated_votes));
@@ -327,7 +327,7 @@ TEST (request_aggregator, channel_max_queue)
 	std::vector<std::pair<nano::block_hash, nano::root>> request;
 	request.emplace_back (send1->hash (), send1->root ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	node.aggregator.request (request, dummy_channel);
 	node.aggregator.request (request, dummy_channel);
 	ASSERT_LT (0, node.stats.count (nano::stat::type::aggregator, nano::stat::detail::aggregator_dropped));
@@ -356,7 +356,7 @@ TEST (request_aggregator, DISABLED_unique)
 	std::vector<std::pair<nano::block_hash, nano::root>> request;
 	request.emplace_back (send1->hash (), send1->root ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	node.aggregator.request (request, dummy_channel);
 	node.aggregator.request (request, dummy_channel);
 	node.aggregator.request (request, dummy_channel);
@@ -401,7 +401,7 @@ TEST (request_aggregator, cannot_vote)
 	// Incorrect hash, correct root
 	request.emplace_back (1, send2->root ());
 	auto client = std::make_shared<nano::transport::socket> (node);
-	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client);
+	std::shared_ptr<nano::transport::channel> dummy_channel = std::make_shared<nano::transport::channel_tcp> (node, client, std::weak_ptr<nano::transport::tcp_server>{});
 	node.aggregator.request (request, dummy_channel);
 	ASSERT_TIMELY (3s, node.aggregator.empty ());
 	ASSERT_EQ (1, node.stats.count (nano::stat::type::aggregator, nano::stat::detail::aggregator_accepted));

nano/core_test/socket.cpp

@@ -422,7 +422,7 @@ TEST (socket, drop_policy)
 		});
 
 		auto client = std::make_shared<nano::transport::socket> (*node);
-		auto channel = std::make_shared<nano::transport::channel_tcp> (*node, client);
+		auto channel = std::make_shared<nano::transport::channel_tcp> (*node, client, std::weak_ptr<nano::transport::tcp_server>{});
 
 		std::atomic completed_writes{ 0 };
 

nano/lib/logging_enums.hpp

@@ -158,6 +158,10 @@ enum class detail
 	frontier_req,
 	bulk_pull_account,
 
+	// message_processor
+	channel_pause,
+	channel_resume,
+
 	_last // Must be the last enum
 };
 
@@ -216,4 +220,4 @@ struct magic_enum::customize::enum_range<nano::log::detail>
 {
 	static constexpr int min = 0;
 	static constexpr int max = 512;
-};
+};

nano/lib/stats_enums.hpp

@@ -22,6 +22,7 @@ enum class type
 	bootstrap,
 	network,
 	tcp_server,
+	tcp_server_message,
 	vote,
 	vote_processor,
 	vote_processor_tier,

nano/node/bootstrap/bootstrap_connections.cpp

@@ -160,7 +160,7 @@ void nano::bootstrap_connections::connect_client (nano::tcp_endpoint const & end
 		{
 			this_l->node.logger.debug (nano::log::type::bootstrap, "Connection established to: {}", nano::util::to_str (endpoint_a));
 
-			auto client (std::make_shared<nano::bootstrap_client> (this_l->node.shared (), std::make_shared<nano::transport::channel_tcp> (*this_l->node.shared (), socket), socket));
+			auto client (std::make_shared<nano::bootstrap_client> (this_l->node.shared (), std::make_shared<nano::transport::channel_tcp> (*this_l->node.shared (), socket, std::weak_ptr<nano::transport::tcp_server>{}), socket));
 			this_l->pool_connection (client, true, push_front);
 		}
 		else

nano/node/fair_queue.hpp

@@ -161,6 +161,11 @@ class fair_queue final
 		{
 			return requests.size ();
 		}
+
+		bool full () const
+		{
+			return requests.size () == max_size;
+		}
 	};
 
 public:
@@ -192,6 +197,12 @@ class fair_queue final
 		return total_size;
 	};
 
+	bool full (origin_type source) const
+	{
+		auto it = queues.find (source);
+		return it == queues.end () ? false : it->second.full ();
+	}
+
 	bool empty () const
 	{
 		return size () == 0;

nano/node/message_processor.cpp

@@ -76,7 +76,7 @@ void nano::message_processor::stop ()
 	threads.clear ();
 }
 
-bool nano::message_processor::put (std::unique_ptr<nano::message> message, std::shared_ptr<nano::transport::channel> const & channel)
+void nano::message_processor::put (std::unique_ptr<nano::message> message, std::shared_ptr<nano::transport::channel> const & channel)
 {
 	release_assert (message != nullptr);
 	release_assert (channel != nullptr);
@@ -87,6 +87,12 @@ bool nano::message_processor::put (std::unique_ptr<nano::message> message, std::
 	{
 		nano::lock_guard<nano::mutex> guard{ mutex };
 		added = queue.push ({ std::move (message), channel }, { nano::no_value{}, channel });
+		if (queue.full ({ nano::no_value{}, channel }))
+		{
+			node.logger.trace (nano::log::type::message_processor, nano::log::detail::channel_pause,
+			nano::log::arg{ "channel", channel });
+			channel->pause ();
+		}
 	}
 	if (added)
 	{
@@ -100,7 +106,6 @@ bool nano::message_processor::put (std::unique_ptr<nano::message> message, std::
 		stats.inc (nano::stat::type::message_processor, nano::stat::detail::overfill);
 		stats.inc (nano::stat::type::message_processor_overfill, to_stat_detail (type));
 	}
-	return added;
 }
 
 void nano::message_processor::run ()
@@ -143,6 +148,12 @@ void nano::message_processor::run_batch (nano::unique_lock<nano::mutex> & lock)
 	{
 		auto const & [message, channel] = entry;
 		release_assert (message != nullptr);
+		auto resumed = channel->resume_maybe ();
+		if (resumed)
+		{
+			node.logger.trace (nano::log::type::message_processor, nano::log::detail::channel_resume,
+			nano::log::arg{ "channel", channel });
+		}
 		process (*message, channel);
 	}
 
@@ -313,4 +324,4 @@ nano::error nano::message_processor_config::deserialize (nano::tomlconfig & toml
 	toml.get ("max_queue", max_queue);
 
 	return toml.get_error ();
-}
+}

nano/node/message_processor.hpp

@@ -33,7 +33,7 @@ class message_processor final
 	void start ();
 	void stop ();
 
-	bool put (std::unique_ptr<nano::message>, std::shared_ptr<nano::transport::channel> const &);
+	void put (std::unique_ptr<nano::message>, std::shared_ptr<nano::transport::channel> const &);
 	void process (nano::message const &, std::shared_ptr<nano::transport::channel> const &);
 
 	std::unique_ptr<container_info_component> collect_container_info (std::string const & name);
@@ -57,4 +57,4 @@ class message_processor final
 	nano::condition_variable condition;
 	std::vector<std::thread> threads;
 };
-}
+}

nano/node/transport/channel.hpp

@@ -127,6 +127,16 @@ class channel
 		network_version = network_version_a;
 	}
 
+	void pause ()
+	{
+		paused = true;
+	}
+
+	virtual bool resume_maybe ()
+	{
+		return paused.exchange (false);
+	}
+
 	nano::endpoint get_peering_endpoint () const;
 	void set_peering_endpoint (nano::endpoint endpoint);
 
@@ -139,11 +149,12 @@ class channel
 	boost::optional<nano::account> node_id{ boost::none };
 	std::atomic<uint8_t> network_version{ 0 };
 	std::optional<nano::endpoint> peering_endpoint{};
+	std::atomic<bool> paused{ false };
 
 protected:
 	nano::node & node;
 
 public: // Logging
 	virtual void operator() (nano::object_stream &) const;
 };
-}
+}

nano/node/transport/tcp.cpp

@@ -11,9 +11,10 @@
  * channel_tcp
  */
 
-nano::transport::channel_tcp::channel_tcp (nano::node & node_a, std::weak_ptr<nano::transport::socket> socket_a) :
-	channel (node_a),
-	socket (std::move (socket_a))
+nano::transport::channel_tcp::channel_tcp (nano::node & node_a, std::weak_ptr<nano::transport::socket> socket_a, std::weak_ptr<nano::transport::tcp_server> server) :
+	channel{ node_a },
+	socket{ std::move (socket_a) },
+	server{ server }
 {
 }
 
@@ -103,6 +104,19 @@ void nano::transport::channel_tcp::operator() (nano::object_stream & obs) const
 	obs.write ("socket", socket);
 }
 
+bool nano::transport::channel_tcp::resume_maybe ()
+{
+	bool result;
+	if ((result = channel::resume_maybe ()))
+	{
+		if (auto server = this->server.lock ())
+		{
+			server->receive_message ();
+		}
+	}
+	return result;
+}
+
 /*
  * tcp_channels
  */
@@ -219,7 +233,7 @@ std::shared_ptr<nano::transport::channel_tcp> nano::transport::tcp_channels::cre
 	fmt::streamed (socket->remote_endpoint ()),
 	node_id.to_node_id ());
 
-	auto channel = std::make_shared<nano::transport::channel_tcp> (node, socket);
+	auto channel = std::make_shared<nano::transport::channel_tcp> (node, socket, server);
 	channel->update_endpoints ();
 	channel->set_node_id (node_id);
 

nano/node/transport/tcp.hpp

@@ -31,7 +31,7 @@ namespace transport
 		friend class nano::transport::tcp_channels;
 
 	public:
-		channel_tcp (nano::node &, std::weak_ptr<nano::transport::socket>);
+		channel_tcp (nano::node &, std::weak_ptr<nano::transport::socket>, std::weak_ptr<nano::transport::tcp_server> server);
 		~channel_tcp () override;
 
 		void update_endpoints ();
@@ -90,8 +90,11 @@ namespace transport
 			}
 		}
 
+		bool resume_maybe () override;
+
 	public:
 		std::weak_ptr<nano::transport::socket> socket;
+		std::weak_ptr<nano::transport::tcp_server> server;
 
 	private:
 		nano::endpoint endpoint;

nano/node/transport/tcp_server.cpp

@@ -161,7 +161,7 @@ auto nano::transport::tcp_server::process_message (std::unique_ptr<nano::message
 		return process_result::abort;
 	}
 
-	node->stats.inc (nano::stat::type::tcp_server, to_stat_detail (message->type ()), nano::stat::dir::in);
+	node->stats.inc (nano::stat::type::tcp_server_message, to_stat_detail (message->type ()), nano::stat::dir::in);
 
 	debug_assert (is_undefined_connection () || is_realtime_connection () || is_bootstrap_connection ());
 
@@ -254,8 +254,7 @@ void nano::transport::tcp_server::queue_realtime (std::unique_ptr<nano::message>
 
 	channel->set_last_packet_received (std::chrono::steady_clock::now ());
 
-	bool added = node->message_processor.put (std::move (message), channel);
-	// TODO: Throttle if not added
+	node->message_processor.put (std::move (message), channel);
 }
 
 auto nano::transport::tcp_server::process_handshake (nano::node_id_handshake const & message) -> handshake_status

nano/node/transport/tcp_server.hpp

@@ -134,6 +134,7 @@ class tcp_server final : public std::enable_shared_from_this<tcp_server>
 		std::shared_ptr<tcp_server> server;
 	};
 
+	friend class channel_tcp;
 	friend class handshake_message_visitor;
 };
 }