serve.zone/remoteingress
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

BREAKING CHANGE(remoteingress-core): add cancellation tokens and cooperative shutdown; switch event channels to bounded mpsc and improve cleanup

Browse Source
This commit is contained in:
Jürgen Kunz
2026-02-19 08:45:32 +00:00
parent a3af2487b7
commit 4e511b3350
6 changed files with 339 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

179
rust/crates/remoteingress-core/src/edge.rs
View File

@@ -6,6 +6,7 @@ use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio_rustls::TlsConnector;
use tokio_util::sync::CancellationToken;
use serde::{Deserialize, Serialize};
use remoteingress_protocol::*;
@@ -69,8 +70,8 @@ pub struct EdgeStatus {
/// The tunnel edge that listens for client connections and multiplexes them to the hub.
pub struct TunnelEdge {
config: RwLock<EdgeConfig>,
event_tx: mpsc::UnboundedSender<EdgeEvent>,
event_rx: Mutex<Option<mpsc::UnboundedReceiver<EdgeEvent>>>,
event_tx: mpsc::Sender<EdgeEvent>,
event_rx: Mutex<Option<mpsc::Receiver<EdgeEvent>>>,
shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
running: RwLock<bool>,
connected: Arc<RwLock<bool>>,
@@ -78,11 +79,12 @@ pub struct TunnelEdge {
active_streams: Arc<AtomicU32>,
next_stream_id: Arc<AtomicU32>,
listen_ports: Arc<RwLock<Vec<u16>>>,
cancel_token: CancellationToken,
}
impl TunnelEdge {
pub fn new(config: EdgeConfig) -> Self {
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (event_tx, event_rx) = mpsc::channel(1024);
Self {
config: RwLock::new(config),
event_tx,
@@ -94,11 +96,12 @@ impl TunnelEdge {
active_streams: Arc::new(AtomicU32::new(0)),
next_stream_id: Arc::new(AtomicU32::new(1)),
listen_ports: Arc::new(RwLock::new(Vec::new())),
cancel_token: CancellationToken::new(),
}
}
/// Take the event receiver (can only be called once).
pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<EdgeEvent>> {
pub async fn take_event_rx(&self) -> Option<mpsc::Receiver<EdgeEvent>> {
self.event_rx.lock().await.take()
}
@@ -126,6 +129,7 @@ impl TunnelEdge {
let next_stream_id = self.next_stream_id.clone();
let event_tx = self.event_tx.clone();
let listen_ports = self.listen_ports.clone();
let cancel_token = self.cancel_token.clone();
tokio::spawn(async move {
edge_main_loop(
@@ -137,6 +141,7 @@ impl TunnelEdge {
event_tx,
listen_ports,
shutdown_rx,
cancel_token,
)
.await;
});
@@ -146,6 +151,7 @@ impl TunnelEdge {
/// Stop the edge.
pub async fn stop(&self) {
self.cancel_token.cancel();
if let Some(tx) = self.shutdown_tx.lock().await.take() {
let _ = tx.send(()).await;
}
@@ -155,20 +161,30 @@ impl TunnelEdge {
}
}
impl Drop for TunnelEdge {
fn drop(&mut self) {
self.cancel_token.cancel();
}
}
async fn edge_main_loop(
config: EdgeConfig,
connected: Arc<RwLock<bool>>,
public_ip: Arc<RwLock<Option<String>>>,
active_streams: Arc<AtomicU32>,
next_stream_id: Arc<AtomicU32>,
event_tx: mpsc::UnboundedSender<EdgeEvent>,
event_tx: mpsc::Sender<EdgeEvent>,
listen_ports: Arc<RwLock<Vec<u16>>>,
mut shutdown_rx: mpsc::Receiver<()>,
cancel_token: CancellationToken,
) {
let mut backoff_ms: u64 = 1000;
let max_backoff_ms: u64 = 30000;
loop {
// Create a per-connection child token
let connection_token = cancel_token.child_token();
// Try to connect to hub
let result = connect_to_hub_and_run(
&config,
@@ -179,12 +195,18 @@ async fn edge_main_loop(
&event_tx,
&listen_ports,
&mut shutdown_rx,
&connection_token,
)
.await;
// Cancel connection token to kill all orphaned tasks from this cycle
connection_token.cancel();
*connected.write().await = false;
let _ = event_tx.send(EdgeEvent::TunnelDisconnected);
let _ = event_tx.try_send(EdgeEvent::TunnelDisconnected);
active_streams.store(0, Ordering::Relaxed);
// Reset stream ID counter for next connection cycle
next_stream_id.store(1, Ordering::Relaxed);
listen_ports.write().await.clear();
match result {
@@ -193,6 +215,7 @@ async fn edge_main_loop(
log::info!("Reconnecting in {}ms...", backoff_ms);
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {}
_ = cancel_token.cancelled() => break,
_ = shutdown_rx.recv() => break,
}
backoff_ms = (backoff_ms * 2).min(max_backoff_ms);
@@ -212,9 +235,10 @@ async fn connect_to_hub_and_run(
public_ip: &Arc<RwLock<Option<String>>>,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
event_tx: &mpsc::UnboundedSender<EdgeEvent>,
event_tx: &mpsc::Sender<EdgeEvent>,
listen_ports: &Arc<RwLock<Vec<u16>>>,
shutdown_rx: &mut mpsc::Receiver<()>,
connection_token: &CancellationToken,
) -> EdgeLoopResult {
// Build TLS connector that skips cert verification (auth is via secret)
let tls_config = rustls::ClientConfig::builder()
@@ -282,12 +306,12 @@ async fn connect_to_hub_and_run(
);
*connected.write().await = true;
let _ = event_tx.send(EdgeEvent::TunnelConnected);
let _ = event_tx.try_send(EdgeEvent::TunnelConnected);
log::info!("Connected to hub at {}", addr);
// Store initial ports and emit event
*listen_ports.write().await = handshake.listen_ports.clone();
let _ = event_tx.send(EdgeEvent::PortsAssigned {
let _ = event_tx.try_send(EdgeEvent::PortsAssigned {
listen_ports: handshake.listen_ports.clone(),
});
@@ -295,17 +319,26 @@ async fn connect_to_hub_and_run(
let stun_interval = handshake.stun_interval_secs;
let public_ip_clone = public_ip.clone();
let event_tx_clone = event_tx.clone();
let stun_token = connection_token.clone();
let stun_handle = tokio::spawn(async move {
loop {
if let Some(ip) = crate::stun::discover_public_ip().await {
let mut pip = public_ip_clone.write().await;
let changed = pip.as_ref() != Some(&ip);
*pip = Some(ip.clone());
if changed {
let _ = event_tx_clone.send(EdgeEvent::PublicIpDiscovered { ip });
tokio::select! {
ip_result = crate::stun::discover_public_ip() => {
if let Some(ip) = ip_result {
let mut pip = public_ip_clone.write().await;
let changed = pip.as_ref() != Some(&ip);
*pip = Some(ip.clone());
if changed {
let _ = event_tx_clone.try_send(EdgeEvent::PublicIpDiscovered { ip });
}
}
}
_ = stun_token.cancelled() => break,
}
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(stun_interval)) => {}
_ = stun_token.cancelled() => break,
}
tokio::time::sleep(std::time::Duration::from_secs(stun_interval)).await;
}
});
@@ -326,6 +359,7 @@ async fn connect_to_hub_and_run(
active_streams,
next_stream_id,
&config.edge_id,
connection_token,
);
// Read frames from hub
@@ -350,7 +384,7 @@ async fn connect_to_hub_and_run(
if let Ok(update) = serde_json::from_slice::<ConfigUpdate>(&frame.payload) {
log::info!("Config update from hub: ports {:?}", update.listen_ports);
*listen_ports.write().await = update.listen_ports.clone();
let _ = event_tx.send(EdgeEvent::PortsUpdated {
let _ = event_tx.try_send(EdgeEvent::PortsUpdated {
listen_ports: update.listen_ports.clone(),
});
apply_port_config(
@@ -361,6 +395,7 @@ async fn connect_to_hub_and_run(
active_streams,
next_stream_id,
&config.edge_id,
connection_token,
);
}
}
@@ -379,13 +414,18 @@ async fn connect_to_hub_and_run(
}
}
}
_ = connection_token.cancelled() => {
log::info!("Connection cancelled");
break EdgeLoopResult::Shutdown;
}
_ = shutdown_rx.recv() => {
break EdgeLoopResult::Shutdown;
}
}
};
// Cleanup
// Cancel connection token to propagate to all child tasks BEFORE aborting
connection_token.cancel();
stun_handle.abort();
for (_, h) in port_listeners.drain() {
h.abort();
@@ -403,6 +443,7 @@ fn apply_port_config(
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
edge_id: &str,
connection_token: &CancellationToken,
) {
let new_set: std::collections::HashSet<u16> = new_ports.iter().copied().collect();
let old_set: std::collections::HashSet<u16> = port_listeners.keys().copied().collect();
@@ -422,6 +463,7 @@ fn apply_port_config(
let active_streams = active_streams.clone();
let next_stream_id = next_stream_id.clone();
let edge_id = edge_id.to_string();
let port_token = connection_token.child_token();
let handle = tokio::spawn(async move {
let listener = match TcpListener::bind(("0.0.0.0", port)).await {
@@ -434,32 +476,42 @@ fn apply_port_config(
log::info!("Listening on port {}", port);
loop {
match listener.accept().await {
Ok((client_stream, client_addr)) => {
let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed);
let tunnel_writer = tunnel_writer.clone();
let client_writers = client_writers.clone();
let active_streams = active_streams.clone();
let edge_id = edge_id.clone();
tokio::select! {
accept_result = listener.accept() => {
match accept_result {
Ok((client_stream, client_addr)) => {
let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed);
let tunnel_writer = tunnel_writer.clone();
let client_writers = client_writers.clone();
let active_streams = active_streams.clone();
let edge_id = edge_id.clone();
let client_token = port_token.child_token();
active_streams.fetch_add(1, Ordering::Relaxed);
active_streams.fetch_add(1, Ordering::Relaxed);
tokio::spawn(async move {
handle_client_connection(
client_stream,
client_addr,
stream_id,
port,
&edge_id,
tunnel_writer,
client_writers,
)
.await;
active_streams.fetch_sub(1, Ordering::Relaxed);
});
tokio::spawn(async move {
handle_client_connection(
client_stream,
client_addr,
stream_id,
port,
&edge_id,
tunnel_writer,
client_writers,
client_token,
)
.await;
active_streams.fetch_sub(1, Ordering::Relaxed);
});
}
Err(e) => {
log::error!("Accept error on port {}: {}", port, e);
}
}
}
Err(e) => {
log::error!("Accept error on port {}: {}", port, e);
_ = port_token.cancelled() => {
log::info!("Port {} listener cancelled", port);
break;
}
}
}
@@ -476,6 +528,7 @@ async fn handle_client_connection(
edge_id: &str,
tunnel_writer: Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>,
client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
client_token: CancellationToken,
) {
let client_ip = client_addr.ip().to_string();
let client_port = client_addr.port();
@@ -503,10 +556,21 @@ async fn handle_client_connection(
let (mut client_read, mut client_write) = client_stream.into_split();
// Task: hub -> client
let hub_to_client_token = client_token.clone();
let hub_to_client = tokio::spawn(async move {
while let Some(data) = back_rx.recv().await {
if client_write.write_all(&data).await.is_err() {
break;
loop {
tokio::select! {
data = back_rx.recv() => {
match data {
Some(data) => {
if client_write.write_all(&data).await.is_err() {
break;
}
}
None => break,
}
}
_ = hub_to_client_token.cancelled() => break,
}
}
let _ = client_write.shutdown().await;
@@ -515,22 +579,27 @@ async fn handle_client_connection(
// Task: client -> hub
let mut buf = vec![0u8; 32768];
loop {
match client_read.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]);
let mut w = tunnel_writer.lock().await;
if w.write_all(&data_frame).await.is_err() {
break;
tokio::select! {
read_result = client_read.read(&mut buf) => {
match read_result {
Ok(0) => break,
Ok(n) => {
let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]);
let mut w = tunnel_writer.lock().await;
if w.write_all(&data_frame).await.is_err() {
break;
}
}
Err(_) => break,
}
}
Err(_) => break,
_ = client_token.cancelled() => break,
}
}
// Send CLOSE frame
let close_frame = encode_frame(stream_id, FRAME_CLOSE, &[]);
{
// Send CLOSE frame (only if not cancelled)
if !client_token.is_cancelled() {
let close_frame = encode_frame(stream_id, FRAME_CLOSE, &[]);
let mut w = tunnel_writer.lock().await;
let _ = w.write_all(&close_frame).await;
}