improved reconnect

esp32/src/main.rs

@@ -1,9 +1,12 @@
 #![no_std]
 #![no_main]
 
+extern crate alloc;
+
 use core::str::FromStr;
 
 use embassy_executor::Spawner;
+use embassy_futures::select::{Either, select};
 use embassy_net::tcp::{TcpReader, TcpWriter};
 use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
 use embassy_sync::mutex::Mutex;
@@ -40,6 +43,14 @@ enum QuestionType {
     Numeric { min: i32, max: i32 },
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum TcpDisconnect {
+    ReadError,
+    ReadEof,
+    WriteError,
+    Cancelled,
+}
+
 struct QuestionData {
     text: OwnedStr<256>,
     q_type: QuestionType,
@@ -70,13 +81,28 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
     loop {}
 }
 
-pub async fn tcp_read_loop(mut read: TcpReader<'_>) {
+pub async fn tcp_read_loop(
+    mut read: TcpReader<'_>,
+    cancel: &Signal<CriticalSectionRawMutex, ()>,
+) -> Result<(), TcpDisconnect> {
     let mut buf = [0u8; 1024];
     let mut generation = 1;
 
-    while let Ok(len) = read.read(&mut buf).await {
+    loop {
+        let read_fut = read.read(&mut buf);
+        let cancel_fut = cancel.wait();
+        let len = match select(read_fut, cancel_fut).await {
+            Either::First(Ok(len)) => len,
+            Either::First(Err(_)) => {
+                cancel.signal(());
+                return Err(TcpDisconnect::ReadError);
+            }
+            Either::Second(()) => return Err(TcpDisconnect::Cancelled),
+        };
+
         if len == 0 {
-            break;
+            cancel.signal(());
+            return Err(TcpDisconnect::ReadEof);
         }
         let Ok(str) = core::str::from_utf8(&buf[..len]) else {
             continue;
@@ -159,15 +185,26 @@ pub async fn tcp_read_loop(mut read: TcpReader<'_>) {
     }
 }
 
-pub async fn tcp_write_loop(mut write: TcpWriter<'_>) {
+pub async fn tcp_write_loop(
+    mut write: TcpWriter<'_>,
+    cancel: &Signal<CriticalSectionRawMutex, ()>,
+) -> Result<(), TcpDisconnect> {
     let mut buffer = buffer::WriteBuffer::<256>::new();
     loop {
-        let data = input::INPUT.receive().await;
+        let input_fut = input::INPUT.receive();
+        let cancel_fut = cancel.wait();
+        let data = match select(input_fut, cancel_fut).await {
+            Either::First(data) => data,
+            Either::Second(()) => return Err(TcpDisconnect::Cancelled),
+        };
         println!("button={}", data);
         let angle = *ANGLE.lock().await;
         core::writeln!(buffer, "button={} angle={}", data, angle).ok();
         println!("write: {}", &buffer);
-        write.write(&buffer).await.ok();
+        if write.write(&buffer).await.is_err() {
+            cancel.signal(());
+            return Err(TcpDisconnect::WriteError);
+        }
         buffer.clear();
     }
 }

esp32/src/net.rs

@@ -1,8 +1,11 @@
 use core::net::SocketAddrV4;
 use core::str::FromStr;
 
+use alloc::format;
 use embassy_futures::join::join;
 use embassy_net::tcp::TcpSocket;
+use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
+use embassy_sync::signal::Signal;
 use embassy_time::{Duration, Timer};
 use esp_hal::peripherals::WIFI;
 use esp_hal::rng::Rng;
@@ -10,8 +13,9 @@ use esp_println::println;
 use esp_radio::wifi::sta::StationConfig;
 use esp_radio::wifi::{Config, ControllerConfig, scan::ScanConfig};
 
+use crate::screen::overwrite_lcd;
+use crate::{TcpDisconnect, buffer::wait_for_config, tcp_read_loop, tcp_write_loop};
 use crate::{WIFI_NETWORK, WIFI_PASSWORD};
-use crate::{buffer::wait_for_config, tcp_read_loop, tcp_write_loop};
 
 pub struct NetworkConfig<'a> {
     pub wifi: WIFI<'a>,
@@ -65,6 +69,8 @@ pub async fn network_setup_task(
     spawner.spawn(net_task(runner).expect("spawn net task"));
 
     loop {
+        loop {
+            overwrite_lcd("Connecting", "").await;
             println!("Scan");
             let scan_config = ScanConfig::default().with_max(10);
             let result = controller.scan_async(&scan_config).await.unwrap();
@@ -86,12 +92,14 @@ pub async fn network_setup_task(
                     .with_password(WIFI_PASSWORD.into()),
             )) {
                 println!("Connection error: {:?}", e);
+                overwrite_lcd("Connection error", &format!("{}", e)).await;
                 Timer::after_millis(1000).await;
                 continue;
             }
 
             if let Err(e) = controller.connect_async().await {
                 println!("Connect error: {:?}", e);
+                overwrite_lcd("Connection error", &format!("{}", e)).await;
                 Timer::after_millis(1000).await;
                 continue;
             }
@@ -121,11 +129,31 @@ pub async fn network_setup_task(
                 .await
             {
                 println!("tcp connect error: {:?}", e);
+                overwrite_lcd("TCP error", &format!("{}", e)).await;
                 Timer::after_millis(1000).await;
                 continue;
             }
+            overwrite_lcd("Connected", "").await;
 
+            let cancel = Signal::<CriticalSectionRawMutex, ()>::new();
             let (read, write) = socket.split();
-        join(tcp_read_loop(read), tcp_write_loop(write)).await;
+            let (read_result, write_result) =
+                join(tcp_read_loop(read, &cancel), tcp_write_loop(write, &cancel)).await;
+            let disconnect = read_result.err().or(write_result.err());
+
+            if let Some(reason) = disconnect {
+                println!("tcp disconnected: {:?}", reason);
+            } else {
+                println!("tcp loops ended");
+            }
+
+            if !stack.is_config_up() {
+                println!("wifi down, reconnecting wifi");
+                break;
+            }
+
+            overwrite_lcd("Connection closed", "").await;
+            Timer::after_millis(500).await;
+        }
     }
 }