feat(rust): add possibility to change socket buffer size

Author: SanjoDeundiak

implementations/rust/ockam/ockam_command/src/environment/static/env_info.txt

@@ -54,10 +54,12 @@ UDP
 TCP
 - OCKAM_PRIVILEGED: if variable is set, all TCP Inlets/Outlets will use eBPF (overrides `--privileged` argument for `ockam tcp-inlet create` and `ockam tcp-outlet create`). WARNING: This flag value should be equal on both ends of a portal (inlet and outlet)
 - OCKAM_TCP_PORTAL_PAYLOAD_LENGTH: size of the buffer into which TCP Portal reads the TCP stream. Default value: `128 * 1024`
+- OCKAM_TCP_PORTAL_SOCKET_LENGTH: // FIXME
 - OCKAM_TCP_PORTAL_SKIP_HANDSHAKE: skip Portal handshake for lower latency, but also lower throughput. WARNING: This flag value should be equal on both ends of a portal (inlet and outlet)
 - OCKAM_TCP_PORTAL_ENABLE_NAGLE: enable Nagle's algorithm for Portal TCP streams for potentially higher throughput, but higher latency
 - OCKAM_TCP_PORTAL_ENABLE_MPTCP: // FIXME
 - OCKAM_ENABLE_MPTCP: // FIXME
+- OCKAM_TCP_SOCKET_LENGTH: // FIXME
 
 Devs Usage
 - OCKAM: a `string` that defines the path to the ockam binary to use.

implementations/rust/ockam/ockam_transport_core/src/error.rs

@@ -82,6 +82,8 @@ pub enum TransportError {
     InvalidOckamPortalPacket(String),
     /// Connection timeout
     ConnectionTimeout,
+    /// Invalid socket options
+    SockOpt(String),
 }
 
 impl ockam_core::compat::error::Error for TransportError {}
@@ -133,6 +135,7 @@ impl core::fmt::Display for TransportError {
             ),
             Self::InvalidOckamPortalPacket(e) => write!(f, "invalid OckamPortalPacket: {}", e),
             Self::ConnectionTimeout => write!(f, "connection timed out"),
+            Self::SockOpt(e) => write!(f, "socket options failed: {}", e),
         }
     }
 }
@@ -175,6 +178,7 @@ impl From<TransportError> for Error {
             IdentifierChanged => Kind::Conflict,
             InvalidOckamPortalPacket(_) => Kind::Invalid,
             ConnectionTimeout => Kind::Io,
+            SockOpt(_) => Kind::Io,
         };
 
         Error::new(Origin::Transport, kind, err)

implementations/rust/ockam/ockam_transport_tcp/src/options.rs

@@ -1,6 +1,7 @@
 use crate::workers::Addresses;
 use ockam_core::compat::sync::Arc;
 use ockam_core::compat::time::Duration;
+use ockam_core::env::get_env;
 use ockam_core::flow_control::{FlowControlId, FlowControlOutgoingAccessControl, FlowControls};
 use ockam_core::{Address, OutgoingAccessControl};
 
@@ -11,17 +12,20 @@ pub struct TcpConnectionOptions {
     pub(super) consumer: Vec<FlowControlId>,
     pub(crate) flow_control_id: FlowControlId,
     pub(crate) enable_mptcp: bool,
+    pub(crate) buffer_size: Option<usize>,
 }
 
 impl TcpConnectionOptions {
     #[allow(clippy::new_without_default)]
     /// Mark this Tcp Receiver as a Producer with a random [`FlowControlId`]
     pub fn new() -> Self {
+        let buffer_size = get_env("OCKAM_TCP_SOCKET_LENGTH").ok().flatten();
         Self {
             timeout: None,
             consumer: vec![],
             flow_control_id: FlowControls::generate_flow_control_id(),
             enable_mptcp: false,
+            buffer_size,
         }
     }
 
@@ -65,6 +69,12 @@ impl TcpConnectionOptions {
         self.enable_mptcp = true;
         self
     }
+
+    /// Set socket buffer size
+    pub fn set_buffer_size(mut self, buffer_size: Option<usize>) -> Self {
+        self.buffer_size = buffer_size;
+        self
+    }
 }
 
 impl TcpConnectionOptions {
@@ -98,6 +108,7 @@ impl TcpConnectionOptions {
 pub struct TcpListenerOptions {
     pub(crate) flow_control_id: FlowControlId,
     pub(crate) enable_mptcp: bool,
+    pub(crate) buffer_size: Option<usize>,
 }
 
 impl TcpListenerOptions {
@@ -106,9 +117,11 @@ impl TcpListenerOptions {
     /// with Spawner's [`FlowControlId`]
     #[allow(clippy::new_without_default)]
     pub fn new() -> Self {
+        let buffer_size = get_env("OCKAM_TCP_SOCKET_LENGTH").ok().flatten();
         Self {
             flow_control_id: FlowControls::generate_flow_control_id(),
             enable_mptcp: false,
+            buffer_size,
         }
     }
 
@@ -128,6 +141,12 @@ impl TcpListenerOptions {
         self.enable_mptcp = true;
         self
     }
+
+    /// Set socket buffer size
+    pub fn set_buffer_size(mut self, buffer_size: Option<usize>) -> Self {
+        self.buffer_size = buffer_size;
+        self
+    }
 }
 
 impl TcpListenerOptions {

implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_listener.rs

@@ -1,17 +1,20 @@
 use crate::portal::addresses::{Addresses, PortalType};
 use crate::portal::tls_certificate::TlsCertificateProvider;
 use crate::portal::{InletSharedState, ReadHalfMaybeTls, WriteHalfMaybeTls};
+use crate::transport::set_socket_buffer_size;
 use crate::{portal::TcpPortalWorker, TcpInlet, TcpInletOptions, TcpRegistry};
 use log::warn;
 use ockam_core::compat::net::SocketAddr;
 use ockam_core::compat::sync::{Arc, RwLock as SyncRwLock};
+use ockam_core::env::get_env;
 use ockam_core::errcode::{Kind, Origin};
 use ockam_core::{async_trait, compat::boxed::Box, Result};
 use ockam_core::{Address, Processor, Route};
 use ockam_node::Context;
 use ockam_transport_core::{HostnamePort, TransportError};
 use rustls::pki_types::CertificateDer;
 use std::io::BufReader;
+use std::os::fd::AsRawFd;
 use std::time::Duration;
 use tokio::net::TcpListener;
 use tokio::time::Instant;
@@ -64,6 +67,11 @@ impl TcpInletListenProcessor {
                 return Err(TransportError::from(err))?;
             }
         };
+
+        if let Ok(Some(buffer_size)) = get_env::<usize>("OCKAM_TCP_PORTAL_SOCKET_LENGTH") {
+            set_socket_buffer_size(inner.as_raw_fd(), buffer_size)?;
+        }
+
         let socket_addr = inner.local_addr().map_err(TransportError::from)?;
         let inlet_shared_state =
             InletSharedState::create(ctx, outlet_listener_route, options.is_paused)?;

implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_worker.rs

@@ -5,6 +5,7 @@ use crate::portal::portal_worker::WriteHalfMaybeTls::{WriteHalfNoTls, WriteHalfW
 use crate::transport::{connect_tcp, connect_tls};
 use crate::{portal::TcpPortalRecvProcessor, PortalInternalMessage, PortalMessage, TcpRegistry};
 use ockam_core::compat::{boxed::Box, sync::Arc};
+use ockam_core::env::get_env;
 use ockam_core::{
     async_trait, AllowAll, AllowOnwardAddress, AllowSourceAddress, Decodable, DenyAll,
     IncomingAccessControl, LocalInfoIdentifier, Mailbox, Mailboxes, OutgoingAccessControl,
@@ -469,10 +470,18 @@ impl TcpPortalWorker {
     }
 
     async fn connect(&mut self) -> Result<()> {
+        let buffer_size = get_env::<usize>("OCKAM_TCP_PORTAL_SOCKET_LENGTH")
+            .ok()
+            .flatten();
         if self.is_tls {
             debug!(portal_type = %self.portal_type, sender_internal = %self.addresses.sender_internal, "connect to {} via TLS", &self.hostname_port);
-            let (rx, tx) =
-                connect_tls(&self.hostname_port, self.enable_mptcp, self.enable_nagle).await?;
+            let (rx, tx) = connect_tls(
+                &self.hostname_port,
+                self.enable_mptcp,
+                self.enable_nagle,
+                buffer_size,
+            )
+            .await?;
             self.write_half = Some(WriteHalfWithTls(tx));
             self.read_half = Some(ReadHalfWithTls(rx));
         } else {
@@ -482,6 +491,7 @@ impl TcpPortalWorker {
                 self.enable_mptcp,
                 self.enable_nagle,
                 None,
+                buffer_size,
             )
             .await?
             .into_split();

implementations/rust/ockam/ockam_transport_tcp/src/transport/common.rs

@@ -5,6 +5,7 @@ use ockam_core::{Error, Result};
 use ockam_transport_core::{HostnamePort, TransportError};
 use socket2::{SockRef, TcpKeepalive};
 use std::net::SocketAddr;
+use std::os::fd::{AsRawFd, RawFd};
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::io::{ReadHalf, WriteHalf};
@@ -14,47 +15,104 @@ use tokio_rustls::rustls::{ClientConfig, RootCertStore};
 use tokio_rustls::{TlsConnector, TlsStream};
 use tracing::{debug, instrument};
 
-pub(crate) async fn bind_tcp_listener(at: SocketAddr, enable_mptcp: bool) -> Result<TcpListener> {
-    if !enable_mptcp {
-        return Ok(TcpListener::bind(&at).await.map_err(TransportError::from)?);
+pub fn set_socket_buffer_size(socket: RawFd, buffer_size: usize) -> Result<()> {
+    let buffer_size = buffer_size as nix::libc::c_int;
+
+    let res = unsafe {
+        #[allow(trivial_casts)]
+        nix::libc::setsockopt(
+            socket.as_raw_fd(),
+            nix::libc::SOL_SOCKET,
+            nix::libc::SO_RCVBUF,
+            (&buffer_size as *const nix::libc::c_int) as *const nix::libc::c_void,
+            size_of::<nix::libc::c_int>() as nix::libc::socklen_t,
+        )
+    };
+
+    assert!(res >= 0,);
+
+    if res < 0 {
+        return Err(TransportError::SockOpt("TCP portal receive buffer error".to_string()).into());
     }
 
-    let socket = bind_mptcp(at).await.map_err(TransportError::from)?;
+    let res = unsafe {
+        #[allow(trivial_casts)]
+        nix::libc::setsockopt(
+            socket.as_raw_fd(),
+            nix::libc::SOL_SOCKET,
+            nix::libc::SO_SNDBUF,
+            (&buffer_size as *const nix::libc::c_int) as *const nix::libc::c_void,
+            size_of::<nix::libc::c_int>() as nix::libc::socklen_t,
+        )
+    };
+
+    if res < 0 {
+        return Err(TransportError::SockOpt("TCP portal send buffer error".to_string()).into());
+    }
 
-    Ok(socket)
+    Ok(())
 }
 
-async fn create_tcp_stream(to: &HostnamePort, enable_mptcp: bool) -> Result<TcpStream> {
-    if !enable_mptcp {
-        return Ok(TcpStream::connect(to.to_string())
-            .await
-            .map_err(TransportError::from)?);
+pub(crate) async fn bind_tcp_listener(
+    at: SocketAddr,
+    enable_mptcp: bool,
+    buffer_size: Option<usize>,
+) -> Result<TcpListener> {
+    let listener = if !enable_mptcp {
+        TcpListener::bind(&at).await.map_err(TransportError::from)?
+    } else {
+        bind_mptcp(at).await.map_err(TransportError::from)?
+    };
+
+    if let Some(buffer_size) = buffer_size {
+        set_socket_buffer_size(listener.as_raw_fd(), buffer_size)?;
     }
 
-    // TODO: Add timeout
-    let socket = connect_mptcp(to.to_string())
-        .await
-        .map_err(TransportError::from)?;
+    Ok(listener)
+}
+
+async fn create_tcp_stream(
+    to: &HostnamePort,
+    enable_mptcp: bool,
+    buffer_size: Option<usize>,
+) -> Result<TcpStream> {
+    let stream = if !enable_mptcp {
+        TcpStream::connect(to.to_string())
+            .await
+            .map_err(TransportError::from)?
+    } else {
+        // TODO: Add timeout
+        connect_mptcp(to.to_string())
+            .await
+            .map_err(TransportError::from)?
+    };
+
+    if let Some(buffer_size) = buffer_size {
+        set_socket_buffer_size(stream.as_raw_fd(), buffer_size)?;
+    }
 
-    Ok(socket)
+    Ok(stream)
 }
 
 async fn create_tcp_stream_timeout(
     to: &HostnamePort,
     enable_mptcp: bool,
     timeout: Option<Duration>,
+    buffer_size: Option<usize>,
 ) -> Result<TcpStream> {
     match timeout {
         Some(timeout) => {
-            match tokio::time::timeout(timeout, create_tcp_stream(to, enable_mptcp)).await {
+            match tokio::time::timeout(timeout, create_tcp_stream(to, enable_mptcp, buffer_size))
+                .await
+            {
                 Ok(result) => result,
                 Err(_) => {
                     debug!(addr = %to, timeout = %timeout.as_secs(),  "Timeout");
                     Err(TransportError::ConnectionTimeout)?
                 }
             }
         }
-        None => create_tcp_stream(to, enable_mptcp).await,
+        None => create_tcp_stream(to, enable_mptcp, buffer_size).await,
     }
 }
 
@@ -64,10 +122,11 @@ pub(crate) async fn connect_tcp(
     enable_mptcp: bool,
     enable_nagle: bool,
     timeout: Option<Duration>,
+    buffer_size: Option<usize>,
 ) -> Result<TcpStream> {
     debug!(addr = %to, "Connecting");
 
-    let result = create_tcp_stream_timeout(to, enable_mptcp, timeout).await;
+    let result = create_tcp_stream_timeout(to, enable_mptcp, timeout, buffer_size).await;
 
     let connection = match result {
         Ok(c) => {
@@ -109,14 +168,15 @@ pub(crate) async fn connect_tls(
     to: &HostnamePort,
     enable_mptcp: bool,
     enable_nagle: bool,
+    buffer_size: Option<usize>,
 ) -> Result<(
     ReadHalf<TlsStream<TcpStream>>,
     WriteHalf<TlsStream<TcpStream>>,
 )> {
     debug!(to = %to, "Trying to connect using TLS");
 
     // create a tcp stream
-    let connection = connect_tcp(to, enable_mptcp, enable_nagle, None).await?;
+    let connection = connect_tcp(to, enable_mptcp, enable_nagle, None, buffer_size).await?;
 
     // create a TLS connector
     let tls_connector = create_tls_connector().await?;

implementations/rust/ockam/ockam_transport_tcp/src/transport/connection.rs

@@ -112,7 +112,7 @@ impl TcpTransport {
         debug!("Connecting to {}", peer.clone());
 
         let (read_half, write_half) =
-            connect_tcp(&peer, options.enable_mptcp, false, options.timeout)
+            connect_tcp(&peer, options.enable_mptcp, false, options.timeout, None)
                 .await?
                 .into_split();
         let socket = read_half

implementations/rust/ockam/ockam_transport_tcp/src/workers/listener.rs

@@ -29,7 +29,7 @@ impl TcpListenProcessor {
         options: TcpListenerOptions,
     ) -> Result<(SocketAddr, Address)> {
         debug!("Binding TcpListener to {}", addr);
-        let inner = bind_tcp_listener(addr, options.enable_mptcp).await?;
+        let inner = bind_tcp_listener(addr, options.enable_mptcp, None).await?;
         let saddr = inner.local_addr().map_err(TransportError::from)?;
 
         let address = Address::random_tagged("TcpListenProcessor");