Grpc tcp connection limiter

Cargo.toml

@@ -145,7 +145,7 @@ parking_lot = "0.12"
 smallvec = { version = "1.10.0", features = ["serde"] }
 borsh = { version = "0.9.1", features = ["rc"] } # please keep this fixed
 async-std = { version = "1.12.0", features = ['attributes'] }
-async-channel = "1.8.0"
+async-channel = "2.0.0"
 clap = { version = "4.0.23", features = ["derive", "string", "cargo"] }
 derive_more = { version = "0.99" }
 log = "0.4"
@@ -207,16 +207,16 @@ pbkdf2 = { version = "0.12.1" }
 
 
 # workflow dependencies
-workflow-d3 = { version = "0.7.0" }
-workflow-nw = { version = "0.7.0" }
-workflow-log = { version = "0.7.0" }
-workflow-core = { version = "0.7.0" }
-workflow-wasm = { version = "0.7.0" }
-workflow-dom = { version = "0.7.0" }
-workflow-rpc = { version = "0.7.0" }
-workflow-node = { version = "0.7.0" }
-workflow-store = { version = "0.7.0" }
-workflow-terminal = { version = "0.7.0" }
+workflow-d3 = { version = "0.8.0" }
+workflow-nw = { version = "0.8.0" }
+workflow-log = { version = "0.8.0" }
+workflow-core = { version = "0.8.0" }
+workflow-wasm = { version = "0.8.0" }
+workflow-dom = { version = "0.8.0" }
+workflow-rpc = { version = "0.8.0" }
+workflow-node = { version = "0.8.0" }
+workflow-store = { version = "0.8.0" }
+workflow-terminal = { version = "0.8.0" }
 nw-sys = "0.1.5"
 
 # if below is enabled, this means that there is an ongoing work

consensus/Cargo.toml

@@ -38,7 +38,7 @@ kaspa-consensusmanager.workspace = true
 once_cell.workspace = true
 
 crossbeam-channel = "0.5"
-async-channel = "1.8.0"
+async-channel.workspace = true
 secp256k1 = { version = "0.24", features = ["global-context", "rand-std"] }
 
 rand = { workspace = true, features = ["small_rng"] }

kaspad/src/args.rs

@@ -65,6 +65,8 @@ pub struct Args {
  pub prealloc_amount: u64,
 
  pub disable_upnp: bool,
+
+ pub max_tcp_connections: Option<u16>,
 }
 
 impl Default for Args {
@@ -111,6 +113,7 @@ impl Default for Args {
  prealloc_amount: 1_000_000,
 
  disable_upnp: false,
+ max_tcp_connections: None,
  }
  }
 }
@@ -316,7 +319,14 @@ pub fn cli() -> Command {
  .require_equals(true)
  .value_parser(clap::value_parser!(u64))
  .help("Interval in seconds for performance metrics collection."),
- )
+ ).arg(
+ Arg::new("max-tcp-conns")
+ .long("max-tcp-conns")
+ .value_name("max-tcp-conns")
+ .require_equals(true)
+ .value_parser(clap::value_parser!(u16))
+ .help("Max number of tcp connections."),
+ )
  .arg(arg!(--"disable-upnp" "Disable upnp"));
 
  #[cfg(feature = "devnet-prealloc")]
@@ -377,6 +387,7 @@ pub fn parse_args() -> Args {
  #[cfg(feature = "devnet-prealloc")]
  prealloc_amount: m.get_one::<u64>("prealloc-amount").cloned().unwrap_or(defaults.prealloc_amount),
  disable_upnp: m.get_one::<bool>("disable-upnp").cloned().unwrap_or(defaults.disable_upnp),
+ max_tcp_connections: m.get_one::<u16>("max-tcp-conns").cloned(),
  }
 }
 

kaspad/src/daemon.rs

@@ -27,6 +27,8 @@ use kaspa_mining::{
 use kaspa_p2p_flows::{flow_context::FlowContext, service::P2pService};
 
 use kaspa_perf_monitor::builder::Builder as PerfMonitorBuilder;
+use kaspa_utils::fd_budget::acquire_guard;
+use kaspa_utils::tcp_limiter::Limit;
 use kaspa_utxoindex::{api::UtxoIndexProxy, UtxoIndex};
 use kaspa_wrpc_server::service::{Options as WrpcServerOptions, ServerCounters as WrpcServerCounters, WrpcEncoding, WrpcService};
 
@@ -177,6 +179,18 @@ pub fn create_core(args: Args, fd_total_budget: i32) -> (Arc<Core>, Arc<RpcCoreS
 pub fn create_core_with_runtime(runtime: &Runtime, args: &Args, fd_total_budget: i32) -> (Arc<Core>, Arc<RpcCoreService>) {
  let network = args.network();
  let mut fd_remaining = fd_total_budget;
+ let tcp_limit = if let Some(tcp_limit) = args.max_tcp_connections {
+ tcp_limit as i32
+ } else {
+ args.rpc_max_clients as i32 - args.inbound_limit as i32 - args.outbound_target as i32
+ };
+ let Ok(_tcp_limit_guard) = acquire_guard(tcp_limit) else {
+ println!("Oops! Looks like we've hit a system limit. You have a couple of options:");
+ println!("1. Increase the file limit for this process. If you're on Linux, you can do this with the 'ulimit' command.");
+ println!("2. Reduce the TCP connection limit by passing the `max-tcp-conns` argument when starting this program.");
+ exit(1);
+ };
+ fd_remaining -= tcp_limit;
  let utxo_files_limit = if args.utxoindex {
  let utxo_files_limit = fd_remaining * 10 / 100;
  fd_remaining -= utxo_files_limit;
@@ -190,6 +204,7 @@ pub fn create_core_with_runtime(runtime: &Runtime, args: &Args, fd_total_budget:
  exit(1);
  }
 
+ let tcp_limit = Some(Arc::new(Limit::new(tcp_limit)));
  let config = Arc::new(
  ConfigBuilder::new(network.into())
  .adjust_perf_params_to_consensus_params()
@@ -368,6 +383,7 @@ do you confirm? (answer y/n or pass --yes to the Kaspad command line to confirm
  args.inbound_limit,
  dns_seeders,
  config.default_p2p_port(),
+ tcp_limit.clone(),
  ));
 
  let rpc_core_service = Arc::new(RpcCoreService::new(
@@ -384,7 +400,7 @@ do you confirm? (answer y/n or pass --yes to the Kaspad command line to confirm
  wrpc_json_counters.clone(),
  perf_monitor.clone(),
  ));
- let grpc_service = Arc::new(GrpcService::new(grpc_server_addr, rpc_core_service.clone(), args.rpc_max_clients));
+ let grpc_service = Arc::new(GrpcService::new(grpc_server_addr, rpc_core_service.clone(), args.rpc_max_clients, tcp_limit.clone()));
 
  // Create an async runtime and register the top-level async services
  let async_runtime = Arc::new(AsyncRuntime::new(args.async_threads));
@@ -419,6 +435,7 @@ do you confirm? (answer y/n or pass --yes to the Kaspad command line to confirm
  WrpcServerOptions {
  listen_address: listen_address.to_address(&network.network_type, &encoding).to_string(), // TODO: use a normalized ContextualNetAddress instead of a String
  verbose: args.wrpc_verbose,
+ tcp_limit: tcp_limit.clone(),
  ..WrpcServerOptions::default()
  },
  ))

kaspad/src/main.rs

@@ -17,8 +17,8 @@ pub fn main() {
  let _profiler = dhat::Profiler::builder().file_name("kaspad-heap.json").build();
 
  let args = parse_args();
- let fd_total_budget = fd_budget::limit() - args.rpc_max_clients as i32 - args.inbound_limit as i32 - args.outbound_target as i32;
- let (core, _) = create_core(args, fd_total_budget);
+
+ let (core, _) = create_core(args, fd_budget::limit());
 
  // Bind the keyboard signal to the core
  Arc::new(Signals::new(&core)).init();

protocol/flows/src/service.rs

@@ -7,6 +7,7 @@ use kaspa_core::{
  trace,
 };
 use kaspa_p2p_lib::Adaptor;
+use kaspa_utils::tcp_limiter::Limit;
 use kaspa_utils::triggers::SingleTrigger;
 
 use crate::flow_context::FlowContext;
@@ -23,6 +24,7 @@ pub struct P2pService {
  dns_seeders: &'static [&'static str],
  default_port: u16,
  shutdown: SingleTrigger,
+ tcp_limit: Option<Arc<Limit>>,
 }
 
 impl P2pService {
@@ -35,6 +37,7 @@ impl P2pService {
  inbound_limit: usize,
  dns_seeders: &'static [&'static str],
  default_port: u16,
+ tcp_limit: Option<Arc<Limit>>,
  ) -> Self {
  Self {
  flow_context,
@@ -46,6 +49,7 @@ impl P2pService {
  inbound_limit,
  dns_seeders,
  default_port,
+ tcp_limit,
  }
  }
 }
@@ -61,7 +65,9 @@ impl AsyncService for P2pService {
  // Prepare a shutdown signal receiver
  let shutdown_signal = self.shutdown.listener.clone();
 
- let p2p_adaptor = Adaptor::bidirectional(self.listen, self.flow_context.hub().clone(), self.flow_context.clone()).unwrap();
+ let p2p_adaptor =
+ Adaptor::bidirectional(self.listen, self.flow_context.hub().clone(), self.flow_context.clone(), self.tcp_limit.clone())
+ .unwrap();
  let connection_manager = ConnectionManager::new(
  p2p_adaptor.clone(),
  self.outbound_target,

protocol/p2p/src/bin/server.rs

@@ -10,7 +10,7 @@ async fn main() {
  // [0] - init p2p-adaptor - server side
  let ip_port = NetAddress::from_str("[::1]:50051").unwrap();
  let initializer = Arc::new(EchoFlowInitializer::new());
- let adaptor = kaspa_p2p_lib::Adaptor::bidirectional(ip_port, kaspa_p2p_lib::Hub::new(), initializer).unwrap();
+ let adaptor = kaspa_p2p_lib::Adaptor::bidirectional(ip_port, kaspa_p2p_lib::Hub::new(), initializer, None).unwrap();
  // [1] - connect to a few peers
  let ip_port = String::from("[::1]:16111");
  for i in 0..1 {

protocol/p2p/src/core/adaptor.rs

@@ -3,6 +3,7 @@ use crate::core::hub::Hub;
 use crate::ConnectionError;
 use crate::{core::connection_handler::ConnectionHandler, Router};
 use kaspa_utils::networking::NetAddress;
+use kaspa_utils::tcp_limiter::Limit;
 use std::ops::Deref;
 use std::sync::Arc;
 use std::time::Duration;
@@ -58,17 +59,19 @@ impl Adaptor {
  serve_address: NetAddress,
  hub: Hub,
  initializer: Arc<dyn ConnectionInitializer>,
+ tcp_limit: Option<Arc<Limit>>,
  ) -> Result<Arc<Self>, ConnectionError> {
  let (hub_sender, hub_receiver) = mpsc_channel(Self::hub_channel_size());
  let connection_handler = ConnectionHandler::new(hub_sender, initializer.clone());
- let server_termination = connection_handler.serve(serve_address)?;
+ let server_termination = connection_handler.serve(serve_address, tcp_limit)?;
  let adaptor = Arc::new(Adaptor::new(Some(server_termination), connection_handler, hub));
  adaptor.hub.clone().start_event_loop(hub_receiver, initializer);
  Ok(adaptor)
  }
 
  /// Connect to a new peer (no retries)
  pub async fn connect_peer(&self, peer_address: String) -> Result<PeerKey, ConnectionError> {
+ // todo pass tcp limit
  self.connection_handler.connect_with_retry(peer_address, 1, Default::default()).await.map(|r| r.key())
  }
 
@@ -79,6 +82,7 @@ impl Adaptor {
  retry_attempts: u8,
  retry_interval: Duration,
  ) -> Result<PeerKey, ConnectionError> {
+ // todo pass tcp limit
  self.connection_handler.connect_with_retry(peer_address, retry_attempts, retry_interval).await.map(|r| r.key())
  }
 

protocol/p2p/src/core/connection_handler.rs

@@ -7,18 +7,19 @@ use crate::{ConnectionInitializer, Router};
 use futures::FutureExt;
 use kaspa_core::{debug, info};
 use kaspa_utils::networking::NetAddress;
+use kaspa_utils::tcp_limiter::{Limit, Wrapper};
 use std::net::ToSocketAddrs;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::time::Duration;
 use thiserror::Error;
+use tokio::net::TcpListener;
 use tokio::sync::mpsc::{channel as mpsc_channel, Sender as MpscSender};
 use tokio::sync::oneshot::{channel as oneshot_channel, Sender as OneshotSender};
-use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
 use tokio_stream::StreamExt;
 use tonic::transport::{Error as TonicError, Server as TonicServer};
 use tonic::{Request, Response, Status as TonicStatus, Streaming};
-
 #[derive(Error, Debug)]
 pub enum ConnectionError {
  #[error("missing socket address")]
@@ -54,7 +55,11 @@ impl ConnectionHandler {
  }
 
  /// Launches a P2P server listener loop
- pub(crate) fn serve(&self, serve_address: NetAddress) -> Result<OneshotSender<()>, ConnectionError> {
+ pub(crate) fn serve(
+ &self,
+ serve_address: NetAddress,
+ tcp_limit: Option<Arc<Limit>>,
+ ) -> Result<OneshotSender<()>, ConnectionError> {
  let (termination_sender, termination_receiver) = oneshot_channel::<()>();
  let connection_handler = self.clone();
  info!("P2P Server starting on: {}", serve_address);
@@ -64,11 +69,21 @@ impl ConnectionHandler {
  .send_compressed(tonic::codec::CompressionEncoding::Gzip)
  .max_decoding_message_size(P2P_MAX_MESSAGE_SIZE);
 
- // TODO: check whether we should set tcp_keepalive
- let serve_result = TonicServer::builder()
- .add_service(proto_server)
- .serve_with_shutdown(serve_address.into(), termination_receiver.map(drop))
- .await;
+ let builder = TonicServer::builder().add_service(proto_server);
+ let serve_result = if let Some(limit) = tcp_limit {
+ let listener = TcpListener::bind(serve_address.to_string()).await.unwrap();
+ let tcp_stream = TcpListenerStream::new(listener).filter_map(|tcp_stream| match tcp_stream {
+ Ok(tcp_stream) => Wrapper::new(tcp_stream, limit.clone()).map(Ok),
+ Err(e) => Some(Err(e)),
+ });
+ builder
+ // TODO: check whether we should set tcp_keepalive
+ .serve_with_incoming_shutdown(tcp_stream, termination_receiver.map(drop))
+ .await
+ } else {
+ // TODO: check whether we should set tcp_keepalive
+ builder.serve_with_shutdown(serve_address.into(), termination_receiver.map(drop)).await
+ };
 
  match serve_result {
  Ok(_) => info!("P2P Server stopped: {}", serve_address),
@@ -128,6 +143,7 @@ impl ConnectionHandler {
  retry_attempts: u8,
  retry_interval: Duration,
  ) -> Result<Arc<Router>, ConnectionError> {
+ // todo consider tcp limit
  let mut counter = 0;
  loop {
  counter += 1;

protocol/p2p/src/echo.rs

@@ -157,10 +157,10 @@ mod tests {
  kaspa_core::log::try_init_logger("debug");
 
  let address1 = NetAddress::from_str("[::1]:50053").unwrap();
- let adaptor1 = Adaptor::bidirectional(address1, Hub::new(), Arc::new(EchoFlowInitializer::new())).unwrap();
+ let adaptor1 = Adaptor::bidirectional(address1, Hub::new(), Arc::new(EchoFlowInitializer::new()), None).unwrap();
 
  let address2 = NetAddress::from_str("[::1]:50054").unwrap();
- let adaptor2 = Adaptor::bidirectional(address2, Hub::new(), Arc::new(EchoFlowInitializer::new())).unwrap();
+ let adaptor2 = Adaptor::bidirectional(address2, Hub::new(), Arc::new(EchoFlowInitializer::new()), None).unwrap();
 
  // Initiate the connection from `adaptor1` (outbound) to `adaptor2` (inbound)
  let peer2_id = adaptor1

rpc/grpc/server/src/adaptor.rs

@@ -3,6 +3,7 @@ use kaspa_core::debug;
 use kaspa_notify::notifier::Notifier;
 use kaspa_rpc_core::{api::rpc::DynRpcService, notify::connection::ChannelConnection, Notification, RpcResult};
 use kaspa_utils::networking::NetAddress;
+use kaspa_utils::tcp_limiter::Limit;
 use std::{ops::Deref, sync::Arc};
 use tokio::sync::{mpsc::channel as mpsc_channel, oneshot::Sender as OneshotSender};
 
@@ -35,10 +36,11 @@ impl Adaptor {
  manager: Manager,
  core_service: DynRpcService,
  core_notifier: Arc<Notifier<Notification, ChannelConnection>>,
+ tcp_limit: Option<Arc<Limit>>,
  ) -> Arc<Self> {
  let (manager_sender, manager_receiver) = mpsc_channel(Self::manager_channel_size());
  let connection_handler = ConnectionHandler::new(manager_sender, core_service.clone(), core_notifier);
- let server_termination = connection_handler.serve(serve_address);
+ let server_termination = connection_handler.serve(serve_address, tcp_limit);
  let adaptor = Arc::new(Adaptor::new(Some(server_termination), connection_handler, manager, serve_address));
  adaptor.manager.clone().start_event_loop(manager_receiver);
  adaptor.start();