diff --git a/backend/windmill-trigger-websocket/src/listener.rs b/backend/windmill-trigger-websocket/src/listener.rs index d46faf1b1a..e34fe6e4f6 100644 --- a/backend/windmill-trigger-websocket/src/listener.rs +++ b/backend/windmill-trigger-websocket/src/listener.rs @@ -11,7 +11,7 @@ use serde::Deserialize; use serde_json::value::RawValue; use std::{borrow::Cow, collections::HashMap, sync::Arc}; use tokio::{net::TcpStream, sync::RwLock}; -use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::{tungstenite, tungstenite::Message, MaybeTlsStream, WebSocketStream}; use windmill_common::{ error::{to_anyhow, Error, Result}, jobs::JobTriggerKind, @@ -22,13 +22,30 @@ use windmill_common::{ }; use windmill_queue::PushArgsOwned; use windmill_trigger::filter::{check_filters, Filter}; -use windmill_trigger::listener::ListeningTrigger; +use windmill_trigger::listener::{update_rw_lock, ListeningTrigger}; use windmill_trigger::trigger_helpers::{ trigger_runnable, trigger_runnable_and_wait_for_raw_result, trigger_runnable_and_wait_for_raw_result_with_error_ctx, TriggerJobArgs, }; use windmill_trigger::Listener; +const MAX_CONNECT_ATTEMPTS: u32 = 5; + +/// Whether a failed WebSocket connect is worth retrying: network-level IO +/// errors and HTTP 5xx/429 handshake responses are typically transient (edge +/// proxies like Cloudflare return 502/520 sporadically), while other errors +/// (bad URL, protocol mismatch, other 4xx) point at configuration and would +/// fail identically on every attempt. +fn is_transient_connect_error(err: &tungstenite::Error) -> bool { + match err { + tungstenite::Error::Io(_) => true, + tungstenite::Error::Http(resp) => { + resp.status().is_server_error() || resp.status().as_u16() == 429 + } + _ => false, + } +} + async fn send_initial_messages( listening_trigger: &ListeningTrigger, writer: &mut SplitSink>, Message>, @@ -176,12 +193,50 @@ impl Listener for WebsocketTrigger { validate_websocket_url_for_ssrf(&connect_url).await?; - let connection = connect_async_with_proxy(&*connect_url) - .await - .map(|conn| Some(conn)) - .map_err(|err| to_anyhow(err).into()); - - connection + // Gateway endpoints are often fronted by an edge proxy (e.g. Cloudflare) + // that sporadically answers the upgrade request with a transient 5xx + // instead of `101 Switching Protocols`, and a `get_consumer` error + // disables the trigger until a human re-enables it — so retry transient + // failures with backoff before giving up. The caller runs `loop_ping` + // concurrently so `last_server_ping` stays alive across the sleeps, and + // killpill cancels this future between awaits. + let mut attempt = 0; + loop { + attempt += 1; + match connect_async_with_proxy(&*connect_url).await { + Ok(conn) => return Ok(Some(conn)), + // Only retry in trigger mode: a failed connect there disables the + // trigger until a human re-enables it, while capture mode is an + // interactive test where instant feedback beats resilience (an + // `Io` error can also be a permanent misconfiguration, e.g. a + // typo'd host, which should surface immediately when iterating). + Err(err) + if listening_trigger.trigger_mode + && attempt < MAX_CONNECT_ATTEMPTS + && is_transient_connect_error(&err) => + { + let delay_secs = 1u64 << attempt; + tracing::warn!( + "Transient error connecting to WebSocket for trigger {} (attempt {}/{}), retrying in {}s: {}", + listening_trigger.path, + attempt, + MAX_CONNECT_ATTEMPTS, + delay_secs, + err + ); + update_rw_lock( + err_message.clone(), + Some(format!( + "Connection attempt {}/{} failed ({}), retrying in {}s...", + attempt, MAX_CONNECT_ATTEMPTS, err, delay_secs + )), + ) + .await; + tokio::time::sleep(std::time::Duration::from_secs(delay_secs)).await; + } + Err(err) => return Err(to_anyhow(err).into()), + } + } } async fn consume( &self, diff --git a/backend/windmill-trigger-websocket/tests/connect_retry.rs b/backend/windmill-trigger-websocket/tests/connect_retry.rs new file mode 100644 index 0000000000..2d21512692 --- /dev/null +++ b/backend/windmill-trigger-websocket/tests/connect_retry.rs @@ -0,0 +1,153 @@ +//! End-to-end tests for the websocket trigger connect-retry behavior: a mock +//! TCP server rejects the websocket upgrade with a configurable HTTP status a +//! number of times before completing a real handshake, and the tests assert +//! which failures `get_consumer` retries. +//! +//! This lives in an integration-test binary (own process) because it sets +//! ALLOW_PRIVATE_WEBSOCKET_URLS — the mock server listens on 127.0.0.1, which +//! the SSRF check blocks — and that process-global env var must not leak into +//! the crate's unit tests, which assert loopback URLs are rejected. + +use std::sync::{ + atomic::{AtomicU32, Ordering}, + Arc, +}; + +use tokio::{ + io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, + net::TcpListener, + sync::{broadcast, RwLock}, +}; +use windmill_trigger::{listener::ListeningTrigger, Listener}; +use windmill_trigger_websocket::{ + WebsocketConfig, WebsocketTrigger, ALLOW_PRIVATE_WEBSOCKET_URLS_ENV, +}; + +/// Mock server: rejects the first `failures` upgrade requests with +/// `status_line` and closes, then completes real websocket handshakes and +/// parks the connection open. Returns the bound address and the +/// connection-attempt counter. +async fn mock_ws_server( + failures: u32, + status_line: &'static str, +) -> (std::net::SocketAddr, Arc) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let attempts = Arc::new(AtomicU32::new(0)); + let served = attempts.clone(); + tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let n = served.fetch_add(1, Ordering::SeqCst); + if n < failures { + // Drain the request head, then reject the upgrade. + let mut reader = BufReader::new(&mut socket); + let mut line = String::new(); + loop { + line.clear(); + let read = reader.read_line(&mut line).await.unwrap_or(0); + if read == 0 || line == "\r\n" { + break; + } + } + socket + .write_all( + format!("HTTP/1.1 {status_line}\r\nContent-Length: 0\r\n\r\n").as_bytes(), + ) + .await + .ok(); + } else if let Ok(ws) = tokio_tungstenite::accept_async(socket).await { + tokio::spawn(async move { + let _open = ws; + std::future::pending::<()>().await + }); + } + } + }); + (addr, attempts) +} + +fn trigger(url: String, trigger_mode: bool) -> ListeningTrigger { + ListeningTrigger { + path: "f/test/ws".to_string(), + is_flow: false, + workspace_id: "test".to_string(), + edited_by: "test".to_string(), + permissioned_as: "u/test".to_string(), + trigger_config: WebsocketConfig { + url, + filters: vec![], + filter_logic: "and".to_string(), + initial_messages: None, + url_runnable_args: None, + can_return_message: false, + can_return_error_result: false, + heartbeat: None, + }, + script_path: "f/test/script".to_string(), + trigger_mode, + error_handling: None, + suspended_mode: false, + } +} + +async fn get_consumer_result( + lt: &ListeningTrigger, + err_message: Arc>>, +) -> windmill_common::error::Result::Consumer>> { + // The static-URL path of `get_consumer` never touches the DB; a lazy pool + // satisfies the signature without a running postgres. + let db: windmill_common::DB = + sqlx::Pool::connect_lazy("postgres://unused:unused@127.0.0.1:1/unused").unwrap(); + let (_killpill_tx, killpill_rx) = broadcast::channel::<()>(1); + WebsocketTrigger + .get_consumer(&db, lt, err_message, killpill_rx) + .await +} + +#[tokio::test] +async fn transient_502s_are_retried_until_the_upgrade_succeeds() { + std::env::set_var(ALLOW_PRIVATE_WEBSOCKET_URLS_ENV, "true"); + let (addr, attempts) = mock_ws_server(2, "502 Bad Gateway").await; + let lt = trigger(format!("ws://{addr}"), true); + let err_message = Arc::new(RwLock::new(None)); + + let consumer = get_consumer_result(<, err_message.clone()) + .await + .expect("connect should succeed after retries"); + + assert!(consumer.is_some(), "expected an established connection"); + assert_eq!(attempts.load(Ordering::SeqCst), 3); + // Retry progress was reported through the shared status lock. + let status = err_message.read().await.clone().unwrap(); + assert!(status.contains("attempt 2/5"), "got status: {status}"); + assert!(status.contains("502"), "got status: {status}"); +} + +#[tokio::test] +async fn non_transient_http_errors_fail_on_the_first_attempt() { + std::env::set_var(ALLOW_PRIVATE_WEBSOCKET_URLS_ENV, "true"); + let (addr, attempts) = mock_ws_server(u32::MAX, "404 Not Found").await; + let lt = trigger(format!("ws://{addr}"), true); + + let err = get_consumer_result(<, Arc::new(RwLock::new(None))) + .await + .expect_err("a 404 upgrade response should not be retried"); + + assert!(err.to_string().contains("404"), "got error: {err}"); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn capture_mode_fails_fast_even_on_transient_errors() { + std::env::set_var(ALLOW_PRIVATE_WEBSOCKET_URLS_ENV, "true"); + let (addr, attempts) = mock_ws_server(u32::MAX, "502 Bad Gateway").await; + let lt = trigger(format!("ws://{addr}"), false); + + let err = get_consumer_result(<, Arc::new(RwLock::new(None))) + .await + .expect_err("capture mode should surface the first failure immediately"); + + assert!(err.to_string().contains("502"), "got error: {err}"); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +}