Improve mirroring for high stress (#2561)

* Some opts

* No blocking in the sniffer loop

* Clippy

* Changelog entry

* Docs

* Sniffer messages docs

* TcpSessionIdentifier visibility

* Sniffer submodules visibility

* TcpConnectionSniffer docs

* Update mirrord/agent/src/sniffer/messages.rs

Co-authored-by: meowjesty <[email protected]>

* Update mirrord/agent/src/sniffer.rs

Co-authored-by: meowjesty <[email protected]>

* Update mirrord/agent/src/sniffer.rs

Co-authored-by: meowjesty <[email protected]>

* CR style suggestions

* CR style suggestion

* What the hell is not working?

* Docs fixed

* Hell yeah

---------

Co-authored-by: meowjesty <[email protected]>

Cargo.toml

@@ -56,7 +56,7 @@ actix-codec = "0.5"
 bincode = { version = "2.0.0-rc.2", features = ["serde"] }
 bytes = "1"
 tokio = { version = "1" }
-tokio-stream = "0.1"
+tokio-stream = { version = "0.1", features = ["sync"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 anyhow = "1"

changelog.d/2529.fixed.md

@@ -0,0 +1 @@
+Improved agent performance when mirroring is under high load. 

mirrord/agent/src/entrypoint.rs

@@ -13,6 +13,7 @@ use client_connection::AgentTlsConnector;
 use dns::{DnsCommand, DnsWorker};
 use futures::TryFutureExt;
 use mirrord_protocol::{ClientMessage, DaemonMessage, GetEnvVarsRequest, LogMessage};
+use sniffer::tcp_capture::RawSocketTcpCapture;
 use tokio::{
     net::{TcpListener, TcpStream},
     process::Command,
@@ -35,7 +36,7 @@ use crate::{
     file::FileManager,
     outgoing::{TcpOutgoingApi, UdpOutgoingApi},
     runtime::get_container,
-    sniffer::{SnifferCommand, TcpConnectionSniffer, TcpSnifferApi},
+    sniffer::{api::TcpSnifferApi, messages::SnifferCommand, TcpConnectionSniffer},
     steal::{
         ip_tables::{
             new_iptables, IPTablesWrapper, SafeIpTables, IPTABLE_MESH, IPTABLE_MESH_ENV,
@@ -48,8 +49,7 @@ use crate::{
     *,
 };
 
-/// Size of [`mpsc`] channels connecting [`TcpStealerApi`] and [`TcpSnifferApi`] with their
-/// background tasks.
+/// Size of [`mpsc`] channels connecting [`TcpStealerApi`] with the background task.
 const CHANNEL_SIZE: usize = 1024;
 
 /// Keeps track of next client id.
@@ -201,6 +201,8 @@ struct ClientConnectionHandler {
     udp_outgoing_api: UdpOutgoingApi,
     dns_api: DnsApi,
     state: State,
+    /// Whether the client has sent us [`ClientMessage::ReadyForLogs`].
+    ready_for_logs: bool,
 }
 
 impl ClientConnectionHandler {
@@ -233,6 +235,7 @@ impl ClientConnectionHandler {
             udp_outgoing_api,
             dns_api,
             state,
+            ready_for_logs: false,
         };
 
         Ok(client_handler)
@@ -244,7 +247,7 @@ impl ClientConnectionHandler {
         connection: &mut ClientConnection,
     ) -> Option<TcpSnifferApi> {
         if let BackgroundTask::Running(sniffer_status, sniffer_sender) = task {
-            match TcpSnifferApi::new(id, sniffer_sender, sniffer_status, CHANNEL_SIZE).await {
+            match TcpSnifferApi::new(id, sniffer_sender, sniffer_status).await {
                 Ok(api) => Some(api),
                 Err(e) => {
                     let message = format!(
@@ -338,7 +341,13 @@ impl ClientConnectionHandler {
                         unreachable!()
                     }
                 }, if self.tcp_sniffer_api.is_some() => match message {
-                    Ok(message) => self.respond(DaemonMessage::Tcp(message)).await?,
+                    Ok((message, Some(log))) if self.ready_for_logs => {
+                        self.respond(DaemonMessage::LogMessage(log)).await?;
+                        self.respond(DaemonMessage::Tcp(message)).await?;
+                    }
+                    Ok((message, _)) => {
+                        self.respond(DaemonMessage::Tcp(message)).await?;
+                    },
                     Err(e) => break e,
                 },
                 message = async {
@@ -461,7 +470,9 @@ impl ClientConnectionHandler {
                 ))
                 .await?;
             }
-            ClientMessage::ReadyForLogs => {}
+            ClientMessage::ReadyForLogs => {
+                self.ready_for_logs = true;
+            }
         }
 
         Ok(true)
@@ -498,7 +509,7 @@ async fn start_agent(args: Args) -> Result<()> {
         let mesh = args.mode.mesh();
 
         let watched_task = WatchedTask::new(
-            TcpConnectionSniffer::TASK_NAME,
+            TcpConnectionSniffer::<RawSocketTcpCapture>::TASK_NAME,
             TcpConnectionSniffer::new(sniffer_command_rx, args.network_interface, mesh).and_then(
                 |sniffer| async move {
                     let res = sniffer.start(cancellation_token).await;
@@ -512,7 +523,7 @@ async fn start_agent(args: Args) -> Result<()> {
         let status = watched_task.status();
         let task = run_thread_in_namespace(
             watched_task.start(),
-            TcpConnectionSniffer::TASK_NAME.to_string(),
+            TcpConnectionSniffer::<RawSocketTcpCapture>::TASK_NAME.to_string(),
             state.container_pid(),
             "net",
         );

mirrord/agent/src/error.rs

@@ -12,7 +12,7 @@ use mirrord_protocol::{
 use thiserror::Error;
 
 use crate::{
-    client_connection::TlsSetupError, namespace::NamespaceError, sniffer::SnifferCommand,
+    client_connection::TlsSetupError, namespace::NamespaceError, sniffer::messages::SnifferCommand,
     steal::StealerCommand,
 };
 
@@ -135,6 +135,9 @@ pub(crate) enum AgentError {
     /// Child agent process spawned in `main` failed.
     #[error("Agent child process failed: {0}")]
     AgentFailed(ExitStatus),
+
+    #[error("Exhausted possible identifiers for incoming connections.")]
+    ExhaustedConnectionId,
 }
 
 pub(crate) type Result<T, E = AgentError> = std::result::Result<T, E>;

mirrord/agent/src/main.rs

@@ -1,6 +1,7 @@
 #![feature(hash_extract_if)]
 #![feature(let_chains)]
 #![feature(type_alias_impl_trait)]
+#![feature(entry_insert)]
 #![cfg_attr(target_os = "linux", feature(tcp_quickack))]
 #![feature(lazy_cell)]
 #![warn(clippy::indexing_slicing)]