[runtime] Feature: memory segments and batching #1550

iyzhang · 2025-05-28T17:27:55Z

This PR closes #88

Copilot

Pull Request Overview

This PR implements memory segments and batching support across the networking stack by replacing single-buffer operations with batch processing using ArrayVec. Key changes include:

Updating TCP, UDP, and lower-layer protocols to process batches of packets.
Replacing the old RECEIVE_BATCH_SIZE constant with MAX_BATCH_SIZE_NUM_PACKETS.
Updating related runtime and OS modules (demikernel, catnap, catnip, etc.) to support the new batching logic.

Reviewed Changes

Copilot reviewed 34 out of 34 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/inetstack/protocols/layer4/tcp/peer.rs	Updated push/pop functions to accept and return ArrayVec batches.
src/inetstack/protocols/layer4/tcp/established/sender.rs	Modified push function to iterate over a batch of buffers.
src/inetstack/protocols/layer4/tcp/established/receiver.rs	Adjusted pop function logic for handling batched packet retrieval.
src/inetstack/protocols/layer4/established/mod.rs	Updated push/pop interfaces to work with packet batches.
src/inetstack/protocols/layer4/mod.rs	Replaced single-buffer push/pop with batch-based operations for UDP.
src/inetstack/protocols/layer3/mod.rs, layer2/mod.rs	Updated receiver and transmitter to use MAX_BATCH_SIZE_NUM_PACKETS.
src/inetstack/consts.rs	Deprecated RECEIVE_BATCH_SIZE in favor of MAX_BATCH_SIZE_NUM_PACKETS.
src/demikernel/libos/network/queue.rs, libos.rs	Adapted network queue operations to work with buffer batches.
src/catpowder and src/catnip runtimes, catnap/linux modules	Updated physical layer and transport implementations to support batching.

src/catnap/linux/transport.rs

anandbonde · 2025-06-24T02:38:15Z

src/inetstack/protocols/layer2/mod.rs

+        for mut pkt in pkts.iter_mut() {
+            eth2_header.serialize_and_attach(&mut pkt);
+        }


pkt is already a mutable reference, so there is no need to take a reference again. You can just use the code as below.

Suggested change

for mut pkt in pkts.iter_mut() {

eth2_header.serialize_and_attach(&mut pkt);

}

for pkt in pkts.iter_mut() {

eth2_header.serialize_and_attach(pkt);

}

anandbonde · 2025-06-24T02:40:26Z

src/runtime/memory/mod.rs

+        let cause: String = format!("cannot allocate a zero element scatter-gather array");
+        error!("into_sgarray(): {}", cause);
+        return Err(Fail::new(libc::EINVAL, &cause));


A string literal can be used here instead of format! to make it more efficient and readable, also there will be no need to take the reference to cause with that change. The block can be written like this:

Suggested change

let cause: String = format!("cannot allocate a zero element scatter-gather array");

error!("into_sgarray(): {}", cause);

return Err(Fail::new(libc::EINVAL, &cause));

let cause: &'static str = format!("cannot allocate a zero element scatter-gather array");

error!("into_sgarray(): {}", cause);

return Err(Fail::new(libc::EINVAL, cause));

Note: this applies to any other places as well in this change which use string literals in format.

anandbonde · 2025-06-24T02:43:49Z

src/inetstack/protocols/layer4/tcp/established/receiver.rs

+                    let remaining_buf: DemiBuffer = buf.split_front(*size)?;
+                    self.pop_queue.push_front(remaining_buf);
+                }
+                *size = *size - buf.len();


This can be shortened to

Suggested change

*size = *size - buf.len();

*size -= buf.len();

anandbonde · 2025-06-24T02:46:04Z

src/inetstack/protocols/layer4/tcp/established/receiver.rs

+            }
+            bufs.push(buf);
+            match size {
+                Some(size) if size == 0 => break,


You can simply use this representation instead to avoid specifying the condition

Suggested change

Some(size) if size == 0 => break,

Some(0) => break,

anandbonde · 2025-06-24T02:57:35Z

src/inetstack/mod.rs

        addr: Option<SocketAddr>,
    ) -> Result<(), Fail> {
-        self.layer4_endpoint.push(sd, buf, addr).await
+        debug_assert!(DEMI_SGARRAY_MAXLEN >= MAX_BATCH_SIZE_NUM_PACKETS);


Can this be done outside the critical path (even though it's only compiled in for debug mode)? Both operands are consts, so this should be checked just once, if possible.

yes, let me check if there is a compile assert that I can use

anandbonde · 2025-06-24T03:03:41Z

src/inetstack/mod.rs

        addr: Option<SocketAddr>,
    ) -> Result<(), Fail> {
-        self.layer4_endpoint.push(sd, buf, addr).await
+        debug_assert!(DEMI_SGARRAY_MAXLEN >= MAX_BATCH_SIZE_NUM_PACKETS);
+        let push_bufs: ArrayVec<DemiBuffer, MAX_BATCH_SIZE_NUM_PACKETS> = bufs.into();


This .into() is not needed because the types to and from are the same. Then, the question arises, do we need push_bufs at all? Or can we just pass in bufs into layer4_endpoint.push()?

anandbonde · 2025-06-24T03:13:38Z

src/catnap/linux/transport.rs

-    ) -> Result<(Option<SocketAddr>, DemiBuffer), Fail> {
-        self.data_from_sd(sd).pop(size).await
+    ) -> Result<(Option<SocketAddr>, ArrayVec<DemiBuffer, DEMI_SGARRAY_MAXLEN>), Fail> {
+        let total_size: usize = 0;


What's the role of total_size in the while loop? It's value never changes.

anandbonde · 2025-06-24T03:31:50Z

src/demikernel/libos/network/libos.rs

-            let cause: String = String::from("zero-length buffer");
+        let bufs: ArrayVec<DemiBuffer, DEMI_SGARRAY_MAXLEN> = clone_sgarray(sga)?;
+        if bufs.is_empty() {
+            let cause: String = String::from("zero-length list of buffers");


String::from can be replaced with string literals and references to cause can be removed while passing into other blocks.

Note: This comment applies to all other places in this change where String::from is used.

anandbonde · 2025-06-24T03:34:35Z

src/inetstack/consts.rs

-/// TODO: This Should be Generic
-pub const RECEIVE_BATCH_SIZE: usize = 4;
+/// Max batch size of packets for both transmit and receive up and down the stack.
+pub const MAX_BATCH_SIZE_NUM_PACKETS: usize = 20;


Why is this number bumped up 5x?

I bumped it to match the new sgaarray max so that we could potentially get up to one sgarray worth of buffers in a single poll. Otherwise, we will have to poll 4 times to fill a single batch that we return to the application. The max sgarray size was chosen because Kyle said that their RPCs grow up to 20 MTUs in size.

That makes sense! It would be good to log this thougth process in form of comment near the const.

anandbonde · 2025-06-24T03:38:20Z

src/inetstack/protocols/layer4/mod.rs

+            Socket::Tcp(socket) => self.tcp.push(socket, bufs).await,
+            Socket::Udp(socket) => {
+                for buf in bufs {
+                    trace!("pushing");


If we need a trace here, then we should add slightly more information in this log statement.

anandbonde · 2025-06-24T04:01:36Z

src/runtime/types/memory.rs

@@ -16,6 +18,8 @@ pub const DEMI_SGARRAY_MAXLEN: usize = 1;
 #[repr(C, packed)]
 #[derive(Copy, Clone)]
 pub struct demi_sgaseg_t {
+    /// Reserved for DemiBuffer metadata.
+    pub sgaseg_md: *mut libc::c_void,


We should evaluate renaming these fields to avoid redundant prefix sgaseg. That will also gives us more space to write meaningful names.

anandbonde · 2025-06-24T04:09:10Z

src/runtime/types/ops.rs

@@ -97,7 +97,7 @@ mod test {
        // Size of a demi_qresult_t structure.
        crate::ensure_eq!(
            mem::size_of::<demi_qresult_t>(),
-            QR_OPCODE_SIZE + QR_QD_SIZE + QR_QT_SIZE + QR_RET_SIZE + QR_VALUE_SIZE
+            QR_OPCODE_SIZE + QR_QD_SIZE + QR_QT_SIZE + QR_RET_SIZE + QR_VALUE_SIZE + 4,


It would be good to give a name to this magic const 4.

anandbonde

The overall approach looks good! Left a few comments.

iyzhang requested review from anandbonde, kyleholohan and Copilot May 28, 2025 17:27

iyzhang self-assigned this May 28, 2025

Copilot AI reviewed May 28, 2025

View reviewed changes

src/catnap/linux/transport.rs Show resolved Hide resolved

iyzhang marked this pull request as draft May 28, 2025 18:20

API: Implement multi segment DemiBuffer push and pop

fc32d02

iyzhang force-pushed the feature-memory-add-segments branch from 74af05a to fc32d02 Compare June 23, 2025 18:47

iyzhang marked this pull request as ready for review June 23, 2025 18:49

anandbonde reviewed Jun 24, 2025

View reviewed changes

anandbonde requested changes Jun 24, 2025

View reviewed changes

[runtime] Feature: memory segments and batching #1550

Are you sure you want to change the base?

[runtime] Feature: memory segments and batching #1550

Conversation

iyzhang commented May 28, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

anandbonde Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anandbonde Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anandbonde Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anandbonde left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

anandbonde Jun 24, 2025 •

edited

Loading

anandbonde Jun 24, 2025 •

edited

Loading

anandbonde Jun 24, 2025 •

edited

Loading