@@ -13,7 +13,7 @@ use std::{result, thread};
 use vm_memory::bitmap::BitmapSlice;
 use vm_memory::{ReadVolatile, VolatileMemoryError, VolatileSlice, WriteVolatile};
 
-use crate::protocol::{Request, Response};
+use crate::protocol::Response;
 
 /// The `KeepAliveStream` is a stream that is intended to be used for the main
 /// connection of live migrations. If the `KeepAliveStream` does not read or
@@ -51,16 +51,14 @@ enum KeepAliveStreamAnswer {
 
 struct KeepAliveWorker<S: Read + Write + AsFd> {
     stream: S,
-    /// Is this running on the sender or receiver side?
-    is_sender: bool,
 }
 
 impl<S> KeepAliveWorker<S>
 where
     S: Read + Write + AsFd,
 {
-    pub fn new(stream: S, is_sender: bool) -> Self {
-        Self { stream, is_sender }
+    pub fn new(stream: S) -> Self {
+        Self { stream }
     }
 
     pub fn read(&mut self, mut buf: Vec<u8>, len: usize) -> io::Result<(Vec<u8>, usize)> {
@@ -74,8 +72,11 @@ where
 
     pub fn write(&mut self, buf: Vec<u8>, len: usize) -> io::Result<(Vec<u8>, usize)> {
         debug_assert!(len <= buf.len());
-        let n = Write::write(&mut self.stream, &buf[..len])?;
-        Ok((buf, n))
+        // If only 6 of 16 response bytes are written, the caller's retry could
+        // follow a keepalive: 6 response bytes, keepalive, then 10 bytes.
+        // Finish this write before the timer can run again.
+        Write::write_all(&mut self.stream, &buf[..len])?;
+        Ok((buf, len))
     }
 
     pub fn flush(&mut self) -> io::Result<()> {
@@ -103,7 +104,6 @@ impl KeepAliveStream {
     pub fn new<T: Read + Write + Send + AsFd + 'static>(
         stream: T,
         timeout: Duration,
-        is_sender: bool,
     ) -> result::Result<Self, io::Error> {
         let fd = stream.as_fd().try_clone_to_owned()?;
 
@@ -114,7 +114,7 @@ impl KeepAliveStream {
         let thread = thread::Builder::new()
             .name("migration_keep_alive_thread".to_string())
             .spawn(move || {
-                let mut worker = KeepAliveWorker::new(stream, is_sender);
+                let mut worker = KeepAliveWorker::new(stream);
                 loop {
                     // The idea is to always send a keep alive message when this times out.
                     match message_rx.recv_timeout(timeout) {
@@ -148,13 +148,8 @@ impl KeepAliveStream {
                             KeepAliveStreamMessage::Disconnect => break,
                         },
                         Err(RecvTimeoutError::Timeout) => {
-                            if worker.is_sender {
-                                let keep_alive = Request::keep_alive();
-                                let _ = keep_alive.write_to(&mut worker.stream);
-                            } else {
-                                let keep_alive = Response::keep_alive();
-                                let _ = keep_alive.write_to(&mut worker.stream);
-                            }
+                            let keep_alive = Response::keep_alive();
+                            let _ = keep_alive.write_to(&mut worker.stream);
                         }
                         Err(RecvTimeoutError::Disconnected) => break,
                     }
 
@@ -18,6 +18,7 @@ use crate::protocol::MemoryRangeTable;
 
 mod bitpos_iterator;
 mod context;
+// TODO(keepalive): remove this module when gardenlinux-release-26-09-18 compatibility ends.
 pub mod keep_alive_stream;
 pub mod progress;
 pub mod protocol;
 
@@ -127,10 +127,7 @@ pub enum Command {
     /// Finalizes the migration without resuming the VM on the destination.
     /// Sent when the source VM was paused at migration time.
     CompletePaused = 9,
-    // We introduced this with discriminant eight but in the meantime,
-    // upstream introduced a new command with discriminant 8. For
-    // migration-compatibility we stick to this temporarily, until we have
-    // a solution for the discriminant collision.
+    // TODO(keepalive): remove this command when gardenlinux-release-26-09-18 compatibility ends.
     KeepAlive = 8,
 }
 
@@ -185,10 +182,6 @@ impl Request {
         Self::new(Command::Abandon, 0)
     }
 
-    pub fn keep_alive() -> Self {
-        Self::new(Command::KeepAlive, 0)
-    }
-
     pub fn command(&self) -> Command {
         self.command
     }
@@ -216,6 +209,7 @@ impl Request {
                 .map_err(|error| MigratableError::DeserializeError(anyhow!("{error:?}")))?;
 
             // If we read a keep alive message, we throw it away and keep reading.
+            // TODO(keepalive): remove this branch when gardenlinux-release-26-09-18 compatibility ends.
             if request.command() == Command::KeepAlive {
                 *request = Request::default();
                 continue;
@@ -237,6 +231,7 @@ pub enum Status {
     Invalid,
     Ok,
     Error,
+    // TODO(keepalive): remove this status when gardenlinux-release-26-09-18 compatibility ends.
     KeepAlive,
 }
 
@@ -265,6 +260,7 @@ impl Response {
         Self::new(Status::Error, 0)
     }
 
+    // TODO(keepalive): remove this constructor when gardenlinux-release-26-09-18 compatibility ends.
     pub fn keep_alive() -> Self {
         Self::new(Status::KeepAlive, 0)
     }
@@ -296,6 +292,7 @@ impl Response {
                 .map_err(|error| MigratableError::DeserializeError(anyhow!("{error:?}")))?;
 
             // If we read a keep alive message, we throw it away and keep reading.
+            // TODO(keepalive): remove this branch when gardenlinux-release-26-09-18 compatibility ends.
             if response.status() == Status::KeepAlive {
                 *response = Response::default();
                 continue;
 
@@ -86,6 +86,7 @@ serde_json = { workspace = true }
 serial_buffer = { path = "../serial_buffer" }
 sha2 = { workspace = true }
 signal-hook = { workspace = true }
+socket2 = { version = "0.6.5", features = ["all"] }
 thiserror = { workspace = true }
 tracer = { path = "../tracer" }
 uuid = { workspace = true }
 
@@ -1734,17 +1734,10 @@ impl Vmm {
         let mut ctx = OngoingMigrationContext::new();
 
         // Set up the socket connection
-        let mut socket = if send_data_migration.local {
-            transport::send_migration_socket(
-                &send_data_migration.destination_url,
-                send_data_migration.tls_dir.as_deref(),
-            )?
-        } else {
-            transport::send_migration_socket_with_keep_alive(
-                &send_data_migration.destination_url,
-                send_data_migration.tls_dir.as_deref(),
-            )?
-        };
+        let mut socket = transport::send_migration_socket(
+            &send_data_migration.destination_url,
+            send_data_migration.tls_dir.as_deref(),
+        )?;
 
         // Start the migration
         transport::send_request_expect_ok(
 
@@ -20,6 +20,7 @@ use std::time::Duration;
 use anyhow::{Context, anyhow};
 use log::{debug, error, info, warn};
 use serde_json;
+use socket2::{SockRef, TcpKeepalive};
 use vm_memory::bitmap::BitmapSlice;
 use vm_memory::{
     Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, ReadVolatile, VolatileMemoryError,
@@ -39,33 +40,10 @@ use crate::{GuestMemoryMmap, VmMigrationConfig};
 /// receiver side.
 pub(crate) const MAX_MIGRATION_CONNECTIONS: u32 = 128;
 
-/// The time a writer may block on a socket until it throws an error.
-///
-/// Also the interval at which the [`KeepAliveStream`] sends keep alive messages.
-///
-/// # Relation with [`MIGRATION_READ_TIMEOUT_DURATION`]
-///
-/// This timeout has to be smaller than [`MIGRATION_READ_TIMEOUT_DURATION`],
-/// otherwise spurious timeouts may happen.
-const MIGRATION_WRITE_TIMEOUT_DURATION: Duration = Duration::from_secs(5);
-
-/// The time a reader may block on a socket until it throws an error.
-///
-/// # Relation with [`MIGRATION_WRITE_TIMEOUT_DURATION`]
-///
-/// This timeout has to be larger than [`MIGRATION_WRITE_TIMEOUT_DURATION`],
-/// otherwise spurious timeouts may happen.
-const MIGRATION_READ_TIMEOUT_DURATION: Duration = {
-    let migration_read_timeout_duration = Duration::from_secs(10);
-
-    // This timeout has to be larger than [`MIGRATION_WRITE_TIMEOUT_DURATION`],
-    // otherwise spurious timeouts may happen.
-    assert!(
-        MIGRATION_WRITE_TIMEOUT_DURATION.as_millis() < migration_read_timeout_duration.as_millis(),
-        "MIGRATION_WRITE_TIMEOUT_DURATION must be smaller than MIGRATION_READ_TIMEOUT_DURATION",
-    );
-    migration_read_timeout_duration
-};
+/// Interval at which the [`KeepAliveStream`] sends keep alive messages.
+/// This must remain shorter than the old sender's 10-second read timeout.
+// TODO(keepalive): remove this interval when gardenlinux-release-26-09-18 compatibility ends.
+const MIGRATION_KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(5);
 
 /// The timeout of the migration-receiver.
 ///
@@ -76,16 +54,6 @@ const MIGRATION_READ_TIMEOUT_DURATION: Duration = {
 /// if the management software doesn't kill the process correctly.
 const MIGRATION_ACCEPT_TIMEOUT_DURATION: Duration = Duration::from_secs(60);
 
-fn set_migration_socket_timeouts(socket: &TcpStream) -> anyhow::Result<()> {
-    socket
-        .set_read_timeout(Some(MIGRATION_READ_TIMEOUT_DURATION))
-        .context("Error setting read timeout on TCP socket")?;
-    socket
-        .set_write_timeout(Some(MIGRATION_WRITE_TIMEOUT_DURATION))
-        .context("Error setting write timeout on TCP socket")?;
-    Ok(())
-}
-
 /// Transport-agnostic listener used to receive connections.
 #[derive(Debug)]
 pub(crate) enum ReceiveListener {
@@ -98,6 +66,7 @@ impl ReceiveListener {
     /// Block until a connection is accepted.
     pub(crate) fn accept(
         &mut self,
+        // TODO(keepalive): remove this parameter when gardenlinux-release-26-09-18 compatibility ends.
         main_connection: bool,
     ) -> Result<SocketStream, MigratableError> {
         match self {
@@ -109,11 +78,16 @@ impl ReceiveListener {
                 let (socket, _) = accept_with_timeout(listener, MIGRATION_ACCEPT_TIMEOUT_DURATION)
                     .context("Failed to accept TCP migration connection")
                     .map_err(MigratableError::MigrateReceive)?;
-                set_migration_socket_timeouts(&socket).map_err(MigratableError::MigrateReceive)?;
+
+                set_tcp_keepalive_and_user_timeout(&socket)
+                    .context("Failed to set socket options")
+                    .map_err(MigratableError::MigrateReceive)?;
 
                 let socket = SocketStream::Tcp(socket);
+
+                // TODO(keepalive): remove this wrapper when gardenlinux-release-26-09-18 compatibility ends.
                 if main_connection {
-                    KeepAliveStream::new(socket, MIGRATION_WRITE_TIMEOUT_DURATION, false)
+                    KeepAliveStream::new(socket, MIGRATION_KEEP_ALIVE_INTERVAL)
                         .map(SocketStream::KeepAlive)
                         .context("Error creating keep-alive migration stream")
                         .map_err(MigratableError::MigrateReceive)
@@ -134,16 +108,20 @@ impl ReceiveListener {
                 let (socket, _) = accept_with_timeout(listener, MIGRATION_ACCEPT_TIMEOUT_DURATION)
                     .context("Failed to accept TCP connection")
                     .map_err(MigratableError::MigrateReceive)?;
-                set_migration_socket_timeouts(&socket).map_err(MigratableError::MigrateReceive)?;
+
+                set_tcp_keepalive_and_user_timeout(&socket)
+                    .context("Failed to set socket options")
+                    .map_err(MigratableError::MigrateReceive)?;
 
                 let socket = TlsStream::new_server(socket, config)
                     .map(Box::new)
                     .map(SocketStream::Tls)
                     .context("Failed to accept TLS migration connection")
                     .map_err(MigratableError::MigrateReceive)?;
 
+                // TODO(keepalive): remove this wrapper when gardenlinux-release-26-09-18 compatibility ends.
                 if main_connection {
-                    KeepAliveStream::new(socket, MIGRATION_WRITE_TIMEOUT_DURATION, false)
+                    KeepAliveStream::new(socket, MIGRATION_KEEP_ALIVE_INTERVAL)
                         .map(SocketStream::KeepAlive)
                         .context("Error creating keep-alive migration stream")
                         .map_err(MigratableError::MigrateReceive)
@@ -228,6 +206,7 @@ pub(crate) enum SocketStream {
     Unix(UnixStream),
     Tcp(TcpStream),
     Tls(Box<TlsStream>),
+    // TODO(keepalive): remove this variant when gardenlinux-release-26-09-18 compatibility ends.
     KeepAlive(KeepAliveStream),
 }
 
@@ -282,6 +261,7 @@ impl ReadVolatile for SocketStream {
             SocketStream::Unix(s) => s.read_volatile(buf),
             SocketStream::Tcp(s) => s.read_volatile(buf),
             SocketStream::Tls(s) => s.read_volatile(buf),
+
             SocketStream::KeepAlive(s) => s.read_volatile(buf),
         }
     }
@@ -940,6 +920,33 @@ fn tcp_address_to_server_name(address: &str) -> Result<&str, anyhow::Error> {
     }
 }
 
+/// Enables `SO_KEEPALIVE` and `TCP_USER_TIMEOUT` for the `tcp_stream`'s socket.
+///
+/// The set options target a failure detection time of two to three minutes.
+fn set_tcp_keepalive_and_user_timeout(tcp_stream: &TcpStream) -> io::Result<()> {
+    /// [`TcpKeepalive`] config for migration TCP sockets.
+    ///
+    /// After 60 seconds, the kernel starts sending keepalive probes.
+    /// Every 15 seconds another keepalive probe is sent.
+    /// If 4 probes are sent without response, the connection is dropped.
+    const MIGRATION_TCP_KEEPALIVE: TcpKeepalive = TcpKeepalive::new()
+        .with_retries(4)
+        .with_time(Duration::from_secs(60))
+        .with_interval(Duration::from_secs(15));
+
+    let socket_ref = SockRef::from(&tcp_stream);
+
+    // `TCP_USER_TIMEOUT` overrides `SO_KEEPALIVE` partially as per documentation:
+    // "[...] TCP_USER_TIMEOUT will override keepalive to determine when
+    // to close a connection due to keepalive failure."
+    // https://man7.org/linux/man-pages/man7/tcp.7.html
+    //TODO: Look into setting the keepalive via `std` once
+    // https://github.com/rust-lang/rust/issues/155889 is stabilized.
+    socket_ref.set_tcp_keepalive(&MIGRATION_TCP_KEEPALIVE)?;
+
+    socket_ref.set_tcp_user_timeout(Some(Duration::from_secs(120)))
+}
+
 /// Connect to a migration endpoint and return the established stream.
 pub(crate) fn send_migration_socket(
     destination_url: &str,
@@ -951,7 +958,10 @@ pub(crate) fn send_migration_socket(
         let socket = TcpStream::connect(address).map_err(|e| {
             MigratableError::MigrateSend(anyhow!("Error connecting to TCP socket: {e}"))
         })?;
-        set_migration_socket_timeouts(&socket).map_err(MigratableError::MigrateSend)?;
+
+        set_tcp_keepalive_and_user_timeout(&socket)
+            .context("Failed to set socket options")
+            .map_err(MigratableError::MigrateSend)?;
 
         if let Some(tls_dir) = tls_dir {
             let server_name = tcp_address_to_server_name(address)
@@ -977,23 +987,6 @@ pub(crate) fn send_migration_socket(
     }
 }
 
-/// Connect to the main migration endpoint and keep the connection active while
-/// memory is transferred over additional streams.
-pub(crate) fn send_migration_socket_with_keep_alive(
-    destination_url: &str,
-    tls_dir: Option<&Path>,
-) -> Result<SocketStream, MigratableError> {
-    match send_migration_socket(destination_url, tls_dir)? {
-        socket @ (SocketStream::Tcp(_) | SocketStream::Tls(_)) => {
-            KeepAliveStream::new(socket, MIGRATION_WRITE_TIMEOUT_DURATION, true)
-                .map(SocketStream::KeepAlive)
-                .context("Error creating keep-alive migration stream")
-                .map_err(MigratableError::MigrateSend)
-        }
-        socket => Ok(socket),
-    }
-}
-
 /// Bind a migration listener for the receiver side.
 pub(crate) fn receive_migration_listener(
     receiver_url: &str,