From 049d0f2caa3f5713caaf7f3582b1c66c0b1ae456 Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Tue, 27 Aug 2024 17:55:06 +0800
Subject: [PATCH] work in progress mirror artiq mainline

---
 src/runtime/src/comms.rs | 102 +++++++++++++++++++++++++++------------
 1 file changed, 72 insertions(+), 30 deletions(-)

diff --git a/src/runtime/src/comms.rs b/src/runtime/src/comms.rs
index 33cefc7..cf54bab 100644
--- a/src/runtime/src/comms.rs
+++ b/src/runtime/src/comms.rs
@@ -818,6 +818,59 @@ pub fn main(timer: GlobalTimer, cfg: Config) {
         let aux_mutex = aux_mutex.clone();
         let routing_table = drtio_routing_table.clone();
 
+        // copy same idea as artiq mainline
+        // loop
+        // check if socket.is_active() -> accept and run kernel task
+        // check if task is terminated -> run idle -> terminate idle on restart_idle
+        // yield and reloop (annoying thing is if async behaves as expected)
+        
+        // start listening on a socket and return the socket
+        let listener = TcpListener::new(10_000);
+        listener.listen(1381);
+
+        loop {
+            if listener.can_accept() { // ie can_accept (socket in state SYN-RECEIVED)
+                let mut stream = listener.accept().await;
+                if connection.try_wait().is_none() {
+                    // there is an existing connection
+                    terminate.signal();
+                    connection.async_wait().await;
+                }
+
+                let _ = terminate.try_wait();
+                task::spawn(async move {
+                    select_biased! {
+                        _ = (async {
+
+                        }) => (),
+                        _ = terminate.async_wait().fuse() => ()
+                    }
+                    let _ = stream.flush().await;
+                    let _ = stream.abort().await;
+                    connection.signal();
+                });
+            }
+
+            if connection.try_wait().is_some() {
+                let _ = terminate.try_wait(); // ensure terminate and restart_idle are always set
+                let _ = restart_idle.try_wait();
+                task::spawn(async move {
+                    select_biased! {
+                        _ = (async {
+                            // needs to ensure that if there is no idle_kernel, enter idle loop and
+                            // keep yielding
+                        }) => (),
+                        _ = terminate.async_wait().fuse() => (),
+                        _ = restart_idle.async_wait().fuse() => ()
+                    }
+                    connection.signal();
+                });
+            }
+
+            // yield and rerun
+            task::r#yield().await;
+        }
+
         select_biased! {
             _ = (async {
                     loop {
@@ -830,33 +883,25 @@ pub fn main(timer: GlobalTimer, cfg: Config) {
                         let running = running.clone();
                         let repeat_lock = repeat_lock.clone();
 
-                        info!("awaiting connection");
-                        if let Ok(mut stream) = TcpStream::accept(1381, 0x10_000, 0x10_000).await {
-                            info!("accepted connection");
-                            task::spawn(async move {
-                                info!("awaiting lock");
-                                let _rep_lock = repeat_lock.async_lock().await;
-                                terminate.signal();
-                                let _lock = running.async_lock().await;
-                                drop(_rep_lock);
-                                info!("lock achieved");
-                                let routing_table = routing_table.borrow();
-                                select_biased! {
-                                    _ = (async {
-                                        info!("pre handling connection");
-                                        let _ = handle_connection(&mut stream, control.clone(), &up_destinations, &aux_mutex, &routing_table, timer)
-                                            .await
-                                            .map_err(|e| warn!("connection terminated: {}", e));
-                                    }).fuse() => (),
-                                    _ = terminate.async_wait().fuse() => ()
-                                }
-                                drop(_lock);
-                                let _ = stream.flush().await;
-                                let _ = stream.abort().await;
-                            });
-                        } else {
-                            info!("failed to accept connection");
-                        }
+                        let mut stream = TcpStream::accept(1381, 0x10_000, 0x10_000).await.unwrap();
+                        task::spawn(async move {
+                            let _rep_lock = repeat_lock.async_lock().await;
+                            terminate.signal();
+                            let _lock = running.async_lock().await;
+                            drop(_rep_lock);
+                            let routing_table = routing_table.borrow();
+                            select_biased! {
+                                _ = (async {
+                                    let _ = handle_connection(&mut stream, control.clone(), &up_destinations, &aux_mutex, &routing_table, timer)
+                                        .await
+                                        .map_err(|e| warn!("connection terminated: {}", e));
+                                }).fuse() => (),
+                                _ = terminate.async_wait().fuse() => ()
+                            }
+                            drop(_lock);
+                            let _ = stream.flush().await;
+                            let _ = stream.abort().await;
+                        });
                     }
             }).fuse() => (),
             _ = (async {
@@ -873,9 +918,7 @@ pub fn main(timer: GlobalTimer, cfg: Config) {
                         
                         let _rep_lock = repeat_lock.async_lock().await;
                         drop(_rep_lock);
-                        info!("getting lock idle kernel");
                         let _lock = running.async_lock().await;
-                        info!("locked idle kernel");
                         let routing_table = routing_table.borrow();
                         select_biased! {
                                 _ = (async {
@@ -886,7 +929,6 @@ pub fn main(timer: GlobalTimer, cfg: Config) {
                                 _ = terminate.async_wait().fuse() => (),
                                 _ = restart_idle.async_wait().fuse() => ()
                         }
-                        info!("supposed to drop idle kernel");
                         drop(_lock);
                     }
             }).fuse() => (),