From 4312f30f34c1ffd87e3c97ab26d528858deb33b6 Mon Sep 17 00:00:00 2001
From: Irene Zhang <irene.zhang@microsoft.com>
Date: Sat, 25 May 2024 12:32:40 -0700
Subject: [PATCH] [catnap] Enhancement: Yield to kernel

---
 src/rust/catnap/linux/transport.rs | 11 +++++++----
 src/rust/catnap/mod.rs             |  2 ++
 src/rust/catnap/win/overlapped.rs  | 15 +++++++++++++--
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/rust/catnap/linux/transport.rs b/src/rust/catnap/linux/transport.rs
index de1d89c7a..48d598812 100644
--- a/src/rust/catnap/linux/transport.rs
+++ b/src/rust/catnap/linux/transport.rs
@@ -14,9 +14,12 @@ mod socket;
 //======================================================================================================================
 
 use crate::{
-    catnap::transport::socket::{
-        SharedSocketData,
-        SocketData,
+    catnap::{
+        transport::socket::{
+            SharedSocketData,
+            SocketData,
+        },
+        YIELD_TIMEOUT_MS,
     },
     demikernel::config::Config,
     expect_ok,
@@ -173,7 +176,7 @@ impl SharedCatnapTransport {
                     self.epoll_fd,
                     events.as_mut_ptr() as *mut libc::epoll_event,
                     EPOLL_BATCH_SIZE as i32,
-                    0,
+                    YIELD_TIMEOUT_MS as i32,
                 )
             } {
                 result if result >= 0 => {
diff --git a/src/rust/catnap/mod.rs b/src/rust/catnap/mod.rs
index e5ae4b276..e3c6115ab 100644
--- a/src/rust/catnap/mod.rs
+++ b/src/rust/catnap/mod.rs
@@ -4,3 +4,5 @@
 #[cfg_attr(target_os = "linux", path = "linux/transport.rs")]
 #[cfg_attr(target_os = "windows", path = "win/transport.rs")]
 pub mod transport;
+
+const YIELD_TIMEOUT_MS: u32 = 1;
diff --git a/src/rust/catnap/win/overlapped.rs b/src/rust/catnap/win/overlapped.rs
index d7f38486e..6b9bef72b 100644
--- a/src/rust/catnap/win/overlapped.rs
+++ b/src/rust/catnap/win/overlapped.rs
@@ -31,7 +31,10 @@ use windows::Win32::{
 };
 
 use crate::{
-    catnap::transport::error::translate_ntstatus,
+    catnap::{
+        transport::error::translate_ntstatus,
+        YIELD_TIMEOUT_MS,
+    },
     collections::pin_slab::PinSlab,
     expect_some,
     runtime::{
@@ -245,7 +248,15 @@ impl<S: Unpin> IoCompletionPort<S> {
 
         loop {
             let mut dequeued: u32 = 0;
-            match unsafe { GetQueuedCompletionStatusEx(self.iocp, entries.as_mut_slice(), &mut dequeued, 0, FALSE) } {
+            match unsafe {
+                GetQueuedCompletionStatusEx(
+                    self.iocp,
+                    entries.as_mut_slice(),
+                    &mut dequeued,
+                    YIELD_TIMEOUT_MS,
+                    FALSE,
+                )
+            } {
                 Ok(()) => {
                     for i in 0..dequeued {
                         self.process_overlapped(&entries[i as usize]);