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,67 +69,91 @@ pub async fn network_setup_task(
     spawner.spawn(net_task(runner).expect("spawn net task"));
 
     loop {
-        println!("Scan");
+        loop {
-        let scan_config = ScanConfig::default().with_max(10);
+            overwrite_lcd("Connecting", "").await;
-        let result = controller.scan_async(&scan_config).await.unwrap();
+            println!("Scan");
-        for ap in result {
+            let scan_config = ScanConfig::default().with_max(10);
-            println!(
+            let result = controller.scan_async(&scan_config).await.unwrap();
-                "- {} (CH {} Strength {})",
+            for ap in result {
-                ap.ssid.as_str(),
+                println!(
-                ap.channel,
+                    "- {} (CH {} Strength {})",
-                ap.signal_strength
+                    ap.ssid.as_str(),
-            );
+                    ap.channel,
+                    ap.signal_strength
+                );
+            }
+
+            println!("Scan ended");
+
+            if let Err(e) = controller.set_config(&Config::Station(
+                StationConfig::default()
+                    // .with_auth_method(esp_radio::wifi::AuthenticationMethod::Wpa2Wpa3Personal)
+                    .with_ssid(WIFI_NETWORK)
+                    .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;
+            }
+
+            println!("Configured");
+
+            stack.wait_config_up().await;
+
+            break;
         }
 
-        println!("Scan ended");
+        println!("Waiting for DHCP");
 
-        if let Err(e) = controller.set_config(&Config::Station(
+        let cfg = wait_for_config(stack).await;
-            StationConfig::default()
+        let local_addr = cfg.address.address();
-                // .with_auth_method(esp_radio::wifi::AuthenticationMethod::Wpa2Wpa3Personal)
+        println!("IP address: {:?}", local_addr);
-                .with_ssid(WIFI_NETWORK)
+
-                .with_password(WIFI_PASSWORD.into()),
+        loop {
-        )) {
+            let host_addr = embassy_net::Ipv4Address::from_str(config.target_ip).unwrap();
-            println!("Connection error: {:?}", e);
+            let mut rx_buffer = [0; 4096];
-            Timer::after_millis(1000).await;
+            let mut tx_buffer = [0; 4096];
-            continue;
+            let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
+            socket.set_timeout(Some(Duration::from_secs(10)));
+            socket.set_keep_alive(Some(Duration::from_secs(1)));
+            if let Err(e) = socket
+                .connect(SocketAddrV4::new(host_addr, config.target_port))
+                .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();
+            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;
         }
-
-        if let Err(e) = controller.connect_async().await {
-            println!("Connect error: {:?}", e);
-            Timer::after_millis(1000).await;
-            continue;
-        }
-
-        println!("Configured");
-
-        stack.wait_config_up().await;
-
-        break;
-    }
-
-    println!("Waiting for DHCP");
-
-    let cfg = wait_for_config(stack).await;
-    let local_addr = cfg.address.address();
-    println!("IP address: {:?}", local_addr);
-
-    loop {
-        let host_addr = embassy_net::Ipv4Address::from_str(config.target_ip).unwrap();
-        let mut rx_buffer = [0; 4096];
-        let mut tx_buffer = [0; 4096];
-        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
-        socket.set_timeout(Some(Duration::from_secs(10)));
-        socket.set_keep_alive(Some(Duration::from_secs(1)));
-        if let Err(e) = socket
-            .connect(SocketAddrV4::new(host_addr, config.target_port))
-            .await
-        {
-            println!("tcp connect error: {:?}", e);
-            Timer::after_millis(1000).await;
-            continue;
-        }
-
-        let (read, write) = socket.split();
-        join(tcp_read_loop(read), tcp_write_loop(write)).await;
     }
 }