vsock: flow control fix

This patch has been cherry-picked from the Firecracker tree. The
reference commit is 1cc8b8a678eb28b20f5843556bdb7fbb2dfa6284.

Fixed a logical error in the vsock flow control, that would cause credit
update packets to not be sent at the right time.

Signed-off-by: Dan Horobeanu <dhr@amazon.com>
Signed-off-by: Gabriel Ionescu <gbi@amazon.com>
Signed-off-by: Stefano Garzarella <sgarzare@redhat.com>

vm-virtio/src/vsock/csm/connection.rs

@@ -600,7 +600,9 @@ where
     /// Check if the credit information the peer has last received from us is outdated.
     ///
     fn peer_needs_credit_update(&self) -> bool {
-        (self.fwd_cnt - self.last_fwd_cnt_to_peer).0 as usize >= defs::CONN_CREDIT_UPDATE_THRESHOLD
+        let peer_seen_free_buf =
+            Wrapping(defs::CONN_TX_BUF_SIZE) - (self.fwd_cnt - self.last_fwd_cnt_to_peer);
+        peer_seen_free_buf < Wrapping(defs::CONN_CREDIT_UPDATE_THRESHOLD)
     }
 
     /// Check if we need to ask the peer for a credit update before sending any more data its
@@ -632,7 +634,7 @@ where
             .set_src_port(self.local_port)
             .set_dst_port(self.peer_port)
             .set_type(uapi::VSOCK_TYPE_STREAM)
-            .set_buf_alloc(defs::CONN_TX_BUF_SIZE as u32)
+            .set_buf_alloc(defs::CONN_TX_BUF_SIZE)
             .set_fwd_cnt(self.fwd_cnt.0)
     }
 }
@@ -1052,7 +1054,7 @@ mod tests {
         assert!(ctx.conn.has_pending_rx());
         ctx.recv();
         assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_CREDIT_UPDATE);
-        assert_eq!(ctx.pkt.buf_alloc(), csm_defs::CONN_TX_BUF_SIZE as u32);
+        assert_eq!(ctx.pkt.buf_alloc(), csm_defs::CONN_TX_BUF_SIZE);
         assert_eq!(ctx.pkt.fwd_cnt(), ctx.conn.fwd_cnt.0);
     }
 
@@ -1062,10 +1064,23 @@ mod tests {
 
         // Force a stale state, where the peer hasn't been updated on our credit situation.
         ctx.conn.last_fwd_cnt_to_peer = Wrapping(0);
-        ctx.conn.fwd_cnt = Wrapping(csm_defs::CONN_CREDIT_UPDATE_THRESHOLD as u32);
 
-        // Fake a data send from the peer, to bring us over the credit update threshold.
+        // Since a credit update token is sent when the fwd_cnt value exceeds
+        // CONN_TX_BUF_SIZE - CONN_CREDIT_UPDATE_THRESHOLD, we initialize
+        // fwd_cnt at 6 bytes below the threshold.
+        let initial_fwd_cnt =
+            csm_defs::CONN_TX_BUF_SIZE as u32 - csm_defs::CONN_CREDIT_UPDATE_THRESHOLD as u32 - 6;
+        ctx.conn.fwd_cnt = Wrapping(initial_fwd_cnt);
+
+        // Use a 4-byte packet for triggering the credit update threshold.
         let data = &[1, 2, 3, 4];
+
+        // Check that there is no pending RX.
+        ctx.init_data_pkt(data);
+        ctx.send();
+        assert!(!ctx.conn.has_pending_rx());
+
+        // Send a packet again.
         ctx.init_data_pkt(data);
         ctx.send();
 
@@ -1073,10 +1088,7 @@ mod tests {
         assert!(ctx.conn.has_pending_rx());
         ctx.recv();
         assert_eq!(ctx.pkt.op(), uapi::VSOCK_OP_CREDIT_UPDATE);
-        assert_eq!(
-            ctx.pkt.fwd_cnt() as usize,
-            csm_defs::CONN_CREDIT_UPDATE_THRESHOLD + data.len()
-        );
+        assert_eq!(ctx.pkt.fwd_cnt(), initial_fwd_cnt + data.len() as u32 * 2);
         assert_eq!(ctx.conn.fwd_cnt, ctx.conn.last_fwd_cnt_to_peer);
     }
 
@@ -1175,7 +1187,7 @@ mod tests {
         // Fill up the TX buffer.
         let data = vec![0u8; ctx.pkt.buf().unwrap().len()];
         ctx.init_data_pkt(data.as_slice());
-        for _i in 0..(csm_defs::CONN_TX_BUF_SIZE / data.len()) {
+        for _i in 0..(csm_defs::CONN_TX_BUF_SIZE / data.len() as u32) {
             ctx.send();
         }
 

vm-virtio/src/vsock/csm/mod.rs

@@ -11,11 +11,11 @@ pub use connection::VsockConnection;
 
 pub mod defs {
     /// Vsock connection TX buffer capacity.
-    pub const CONN_TX_BUF_SIZE: usize = 64 * 1024;
+    pub const CONN_TX_BUF_SIZE: u32 = 64 * 1024;
 
-    /// After the guest thinks it has filled our TX buffer up to this limit (in bytes), we'll send
-    /// them a credit update packet, to let them know we can handle more.
-    pub const CONN_CREDIT_UPDATE_THRESHOLD: usize = CONN_TX_BUF_SIZE - 4 * 4 * 1024;
+    /// When the guest thinks we have less than this amount of free buffer space,
+    /// we will send them a credit update packet.
+    pub const CONN_CREDIT_UPDATE_THRESHOLD: u32 = 4 * 1024;
 
     /// Connection request timeout, in millis.
     pub const CONN_REQUEST_TIMEOUT_MS: u64 = 2000;

vm-virtio/src/vsock/csm/txbuf.rs

@@ -25,7 +25,7 @@ pub struct TxBuf {
 impl TxBuf {
     /// Total buffer size, in bytes.
     ///
-    const SIZE: usize = defs::CONN_TX_BUF_SIZE;
+    const SIZE: usize = defs::CONN_TX_BUF_SIZE as usize;
 
     /// Ring-buffer constructor.
     ///