diff --git a/.github/workflows/pipeline.yaml b/.github/workflows/pipeline.yaml index e784906..0065e86 100644 --- a/.github/workflows/pipeline.yaml +++ b/.github/workflows/pipeline.yaml @@ -1,5 +1,5 @@ name: Build -on: [push] +on: [push, pull_request] jobs: nightly-clippy: @@ -7,10 +7,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v2 + uses: actions/checkout@v6 - name: Cache Cargo registry - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: ~/.cargo/registry key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }} @@ -18,7 +18,7 @@ jobs: ${{ runner.os }}-cargo-registry- - name: Cache Cargo build - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: target key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }} @@ -26,11 +26,9 @@ jobs: ${{ runner.os }}-cargo-build- - name: Set up Rust Nightly - uses: actions-rs/toolchain@v1 + uses: dtolnay/rust-toolchain@nightly with: - toolchain: nightly components: clippy - override: true - name: Run Clippy run: cargo clippy -- -D warnings @@ -40,10 +38,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v2 + uses: actions/checkout@v6 - name: Cache Cargo registry - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: ~/.cargo/registry key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }} @@ -51,7 +49,7 @@ jobs: ${{ runner.os }}-cargo-registry- - name: Cache Cargo build - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: target key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }} @@ -59,10 +57,7 @@ jobs: ${{ runner.os }}-cargo-build- - name: Set up Rust Nightly - uses: actions-rs/toolchain@v1 - with: - toolchain: nightly - override: true + uses: dtolnay/rust-toolchain@nightly - name: Run Cargo Tests run: cargo test --all-features @@ -72,10 +67,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v2 + uses: actions/checkout@v6 - name: Cache Cargo registry - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: ~/.cargo/registry key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }} @@ -83,7 +78,7 @@ jobs: ${{ runner.os }}-cargo-registry- - name: Cache Cargo build - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: target key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }} @@ -91,11 +86,9 @@ jobs: ${{ runner.os }}-cargo-build- - name: Set up Rust Stable - uses: actions-rs/toolchain@v1 + uses: dtolnay/rust-toolchain@stable with: - toolchain: stable components: clippy - override: true - name: Run Clippy run: cargo clippy -- -D warnings @@ -105,10 +98,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v2 + uses: actions/checkout@v6 - name: Cache Cargo registry - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: ~/.cargo/registry key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }} @@ -116,7 +109,7 @@ jobs: ${{ runner.os }}-cargo-registry- - name: Cache Cargo build - uses: actions/cache@v3 + uses: actions/cache@v5 with: path: target key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }} @@ -124,10 +117,7 @@ jobs: ${{ runner.os }}-cargo-build- - name: Set up Rust Stable - uses: actions-rs/toolchain@v1 - with: - toolchain: stable - override: true + uses: dtolnay/rust-toolchain@stable - name: Run Cargo Tests run: cargo test --all-features diff --git a/CHANGES.md b/CHANGES.md new file mode 100644 index 0000000..eb82da9 --- /dev/null +++ b/CHANGES.md @@ -0,0 +1,21 @@ + + + +## 0.2.0 - Unreleased + +### Added + +* Reconnect with callback installed via `WStompConfig::on_reconnect`. + The reconnect loop is then spawned by `WStompConfig::build_and_connect`, + which in the reconnecting state returns a `WStompReconnectHandle` + instead of a `WStompClient`. + +## 0.1.0 - 2025-11-18 + +* Initial release +* Connects to STOMP servers over WebSocket using [`awc`](https://crates.io/crates/awc). +* Handles all STOMP protocol encoding and decoding via [`async-stomp`](https://crates.io/crates/async-stomp). +* Manages WebSocket ping/pong heartbeats automatically in a background task. +* Provides a simple `tokio::mpsc` channel-based API ([`WStompClient`]) for sending and receiving STOMP frames. +* Connection helpers for various authentication methods +* Optional `rustls` feature for SSL connections, with helpers that force HTTP/1.1 for compatibility with servers like SockJS. diff --git a/README.md b/README.md index bebfefb..8076c2e 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ This crate provides a simple client to connect to a STOMP-enabled WebSocket serv * [`connect`]: Anonymous connection. * [`connect_with_pass`]: Login and passcode authentication. * [`connect_with_token`]: Authentication using an authorization token header. + * [`WStompConfig::auth_token_fn`]: Dynamic token provider — refreshes credentials on every (re)connect. * Optional `rustls` feature for SSL connections, with helpers that force HTTP/1.1 for compatibility with servers like SockJS. ## Installation @@ -42,7 +43,7 @@ wstomp = { version = "0.1.0", features = ["rustls"] } ## Usage -Here is a basic example of connecting, subscribing to a topic, and receiving messages. +Here is a basic example of connecting, subscribing to a topic and receiving messages. ```rust,no_run use wstomp::{ @@ -182,31 +183,42 @@ async fn main() { ### Auto-reconnect -Use [`WStompConfig::build_and_connect_with_reconnection_cb`] method to automatically perform a full reconnect upon errors. +Install a callback via [`WStompConfig::on_reconnect`] and call `build_and_connect` to automatically perform a full reconnect upon errors. The callback receives the result of each connection attempt and returns a [`ReconnectControl`] value that drives the loop. ```rust,no_run -use wstomp::{WStompClient, WStompConfig, WStompConnectError}; +use wstomp::{ReconnectControl, WStompClient, WStompConfig, WStompConnectError, WStompReconnectHandle}; #[actix_rt::main] async fn main() { let url = "wss://secure-server.com/ws"; let session_token = "session_token"; - let cb = { - move |wstomp_client_res: Result| { - async move { - // Unwrap wstomp client here or react to an error. - // Upon an error you can return from the callback to make wstomp library a re-connection attempt + let cb = move |wstomp_client_res: Result| { + async move { + match wstomp_client_res { + Ok(client) => { + // Use the client; when the session ends, Continue triggers a reconnect. + ReconnectControl::Continue + } + Err(err) => { + eprintln!("Connection failed: {}", err); + // Return Stop to give up, or Continue to keep retrying with backoff. + ReconnectControl::Continue + } } } }; - let res = WStompConfig::new(url) + // Returns immediately with a handle; the reconnection loop runs in the background. + let _handle: WStompReconnectHandle = WStompConfig::new(url) .ssl() .auth_token(session_token) - .build_and_connect_with_reconnection_cb(cb); + .on_reconnect(cb) + .build_and_connect() + .expect("Failed to start reconnection loop"); - // ... do different stuff here, but don't exit immediately as this will terminate wstomp loop. + // Keep _handle alive for as long as the loop should run. + // Dropping it (or calling _handle.abort()) stops the loop. } ``` diff --git a/src/client.rs b/src/client.rs index 5a09ffe..724b725 100644 --- a/src/client.rs +++ b/src/client.rs @@ -33,7 +33,8 @@ impl WStompClient { /// It also manages websocket ping-pong heartbeat. /// /// NOTE: This method does not perform automatic reconnection. - /// Use [WStompConfig::build_and_connect_with_reconnection_cb] to auto-reconnect. + /// Install a callback via [WStompConfig::on_reconnect] and call + /// `build_and_connect` to auto-reconnect. pub fn from_framed(ws_framed: Framed) -> Self { // Channel for you to send STOMP frames to the handler task let (app_tx, app_rx) = mpsc::channel::>(100); diff --git a/src/config.rs b/src/config.rs index 32a6db1..eb34cbc 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,17 +1,36 @@ -use std::sync::Arc; +use std::{future::Future, pin::Pin, sync::Arc}; #[cfg(feature = "rustls")] use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer}; -pub struct WStompConfig { - url: U, - opts: WStompConfigOpts, +/// Closure type for dynamic auth token retrieval. Invoked on every +/// (re)connection attempt so the token can be refreshed out-of-band by the +/// application (e.g. via a background token-refresh task) without recreating +/// the wstomp client. +pub type AuthTokenFn = + Arc Pin> + Send>> + Send + Sync>; + +/// Marker for a [`WStompConfig`] without a reconnection callback. In this +/// state [`WStompConfig::build_and_connect`] is a one-shot async method that +/// resolves to a [`crate::WStompClient`]. +pub struct NoReconnect; + +/// Marker for a [`WStompConfig`] with a reconnection callback installed via +/// [`WStompConfig::on_reconnect`]. In this state +/// [`WStompConfig::build_and_connect`] spawns the reconnect loop and returns +/// a [`crate::WStompReconnectHandle`]. +pub struct Reconnecting(pub(crate) F); + +pub struct WStompConfig { + pub(crate) url: U, + pub(crate) opts: WStompConfigOpts, + pub(crate) reconnect: R, } #[derive(Clone)] pub struct WStompConfigOpts { pub ssl: bool, - pub auth_token: Option, + pub auth_token: Option, pub login: Option, pub passcode: Option, #[cfg(feature = "rustls")] @@ -54,14 +73,17 @@ impl Default for WStompConfigOpts { } } -impl WStompConfig { +impl WStompConfig { pub fn new(url: U) -> Self { Self { url, opts: WStompConfigOpts::default(), + reconnect: NoReconnect, } } +} +impl WStompConfig { /// Get url to which this config is assigned to use. pub fn get_url(&self) -> &U { &self.url @@ -88,9 +110,29 @@ impl WStompConfig { self } - /// Sets the authentication token for the connection. + /// Sets a static authentication token for the connection. The token is + /// fixed for the lifetime of the client — use [Self::auth_token_fn] when + /// the token needs to be refreshed between reconnects. pub fn auth_token(mut self, auth_token: impl Into) -> Self { - self.opts.auth_token = Some(auth_token.into()); + let token = auth_token.into(); + self.opts.auth_token = Some(Arc::new(move || { + let token = token.clone(); + Box::pin(async move { Some(token) }) + })); + self + } + + /// Sets a dynamic authentication token provider. `f` is invoked on every + /// (re)connection attempt, so returning a fresh token allows the app to + /// refresh credentials without tearing down the wstomp client. + /// + /// Returning `None` from `f` omits the `Authorization` header entirely. + pub fn auth_token_fn(mut self, f: F) -> Self + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.opts.auth_token = Some(Arc::new(move || Box::pin(f()))); self } @@ -145,8 +187,8 @@ impl WStompConfig { self } - /// If [Self::build_and_connect_with_reconnection_cb] method is used, - /// sets the initial retry interval in seconds. + /// When [`Self::on_reconnect`] is used to install a reconnection + /// callback, sets the initial retry interval in seconds. /// /// Example: Start retrying after 3 seconds. pub fn retry_initial_interval(mut self, seconds: u64) -> Self { @@ -154,8 +196,8 @@ impl WStompConfig { self } - /// If [Self::build_and_connect_with_reconnection_cb] method is used, - /// sets the maximum retry interval in seconds. + /// When [`Self::on_reconnect`] is used to install a reconnection + /// callback, sets the maximum retry interval in seconds. /// /// Example: Cap the wait time at 30 seconds. pub fn retry_max_interval(mut self, seconds: u64) -> Self { @@ -163,8 +205,8 @@ impl WStompConfig { self } - /// If [Self::build_and_connect_with_reconnection_cb] method is used, - /// sets the multiplier for the backoff. + /// When [`Self::on_reconnect`] is used to install a reconnection + /// callback, sets the multiplier for the backoff. /// /// Example: 2.0 doubles the wait time after every failure. pub fn retry_multiplier(mut self, multiplier: f64) -> Self { @@ -172,8 +214,9 @@ impl WStompConfig { self } - /// If [Self::build_and_connect_with_reconnection_cb] method is used, - /// sets a maximum total time to try reconnecting before giving up. + /// When [`Self::on_reconnect`] is used to install a reconnection + /// callback, sets a maximum total time to try reconnecting before giving + /// up. /// /// Defaults to no limit if method not invoked. pub fn retry_max_elapsed_time(mut self, seconds: u64) -> Self { diff --git a/src/connect.rs b/src/connect.rs index b3b9f2e..9d8b3af 100644 --- a/src/connect.rs +++ b/src/connect.rs @@ -5,10 +5,56 @@ use awc::{ ws::WebsocketsRequest, }; use backoff::{ExponentialBackoffBuilder, backoff::Backoff}; -use std::time::Duration; +use std::time::{Duration, Instant}; use tokio::time::sleep; -use crate::{WStompClient, WStompConfig, WStompConnectError, config::WStompConfigOpts}; +use crate::{ + WStompClient, WStompConfig, WStompConnectError, + config::{NoReconnect, Reconnecting, WStompConfigOpts}, +}; + +/// Signal returned by the reconnection callback to control what the +/// reconnection loop does next. +#[derive(Debug, Clone)] +pub enum ReconnectControl { + /// Reconnect using the normal exponential-backoff schedule. Short-lived + /// sessions (closed within one `retry_initial_interval`) are treated as + /// failures for backoff purposes so a subscribe error doesn't cause a + /// tight reconnect loop. + Continue, + /// Sleep for the given duration before the next reconnect attempt, then + /// reset the backoff. Use this when the application already knows how + /// long the server needs to cool down (e.g. rate-limit headers). + DelayThen(Duration), + /// Stop the reconnection loop cleanly. + Stop, +} + +/// Handle to a running reconnection loop. Dropping it aborts the loop, as +/// does calling [`Self::abort`] explicitly. +pub struct WStompReconnectHandle { + join: actix_rt::task::JoinHandle<()>, +} + +impl WStompReconnectHandle { + /// Abort the reconnection loop. Safe to call more than once; subsequent + /// calls are no-ops. + pub fn abort(&self) { + self.join.abort(); + } + + /// Returns whether the reconnection loop has finished (either aborted or + /// ran to completion via [`ReconnectControl::Stop`]). + pub fn is_finished(&self) -> bool { + self.join.is_finished() + } +} + +impl Drop for WStompReconnectHandle { + fn drop(&mut self) { + self.join.abort(); + } +} /// Connect to STOMP server without additional parameters /// @@ -75,7 +121,27 @@ impl StompConnect for WebsocketsRequest { } } -impl WStompConfig +impl WStompConfig { + /// Install a reconnection callback. Switches the config into the + /// reconnecting state, where [`Self::build_and_connect`] spawns a + /// reconnect loop and hands the resulting client (or connection error) + /// to `cb` on every attempt. The callback's return value + /// ([`ReconnectControl`]) drives the loop. + pub fn on_reconnect(self, cb: F) -> WStompConfig> + where + F: Fn(Result) -> R + 'static, + R: Future, + { + let (url, opts) = self.into_parts(); + WStompConfig { + url, + opts, + reconnect: Reconnecting(cb), + } + } +} + +impl WStompConfig where Uri: TryFrom, >::Error: Into, @@ -89,20 +155,34 @@ where WStompConnectError::WsClientError(WsClientError::from(err)) })?; - inner_connect(uri, opts).await + let auth_token = match &opts.auth_token { + Some(f) => f().await, + None => None, + }; + inner_connect(uri, opts, auth_token).await } +} - /// Build the client and spawns connect procedure with reconnection mechanism. - /// The result from the connection procedure and all subsequent reconnection attempts is passed into the callback. - pub fn build_and_connect_with_reconnection_cb( - self, - cb: F, - ) -> Result<(), WStompConnectError> - where - F: Fn(Result) -> R + 'static, - R: Future, - { - let (url, opts) = self.into_parts(); +impl WStompConfig> +where + Uri: TryFrom, + >::Error: Into, + F: Fn(Result) -> R + 'static, + R: Future, +{ + /// Build the client and spawn the reconnection loop. The returned + /// [`WStompReconnectHandle`] aborts the loop on drop. + /// + /// The callback installed via [`WStompConfig::on_reconnect`] is invoked + /// after every connection attempt with either the newly connected + /// `WStompClient` or the error from the failed attempt. Its return value + /// ([`ReconnectControl`]) controls what the loop does next. + pub fn build_and_connect(self) -> Result { + let WStompConfig { + url, + opts, + reconnect: Reconnecting(cb), + } = self; let uri = Uri::try_from(url).map_err(|e| { let err: HttpError = e.into(); @@ -116,24 +196,49 @@ where .with_max_elapsed_time(opts.retry_max_elapsed_time.map(Duration::from_secs)) .build(); - actix_rt::spawn(async move { + // A session that closes faster than retry_initial_interval is treated + // as a failure for backoff purposes — prevents tight reconnect loops + // when e.g. the STOMP subscribe fails right after CONNECT succeeds. + let min_session_duration = Duration::from_secs(opts.retry_initial_interval); + + let join = actix_rt::spawn(async move { loop { - let tx = inner_connect(uri.clone(), opts.clone()).await; - - if tx.is_ok() { - backoff.reset(); - } else if let Some(duration) = backoff.next_backoff() { - sleep(duration).await; - } else { - cb(Err(WStompConnectError::ReconnectionLimit)).await; - break; - } + let auth_token = match &opts.auth_token { + Some(f) => f().await, + None => None, + }; + + let connect_start = Instant::now(); + let tx = inner_connect(uri.clone(), opts.clone(), auth_token).await; + let was_ok = tx.is_ok(); - cb(tx).await; + let control = cb(tx).await; + let session_duration = connect_start.elapsed(); + + match control { + ReconnectControl::Stop => break, + ReconnectControl::DelayThen(d) => { + sleep(d).await; + backoff.reset(); + } + ReconnectControl::Continue => { + let short_or_failed = !was_ok || session_duration < min_session_duration; + if short_or_failed { + if let Some(duration) = backoff.next_backoff() { + sleep(duration).await; + } else { + let _ = cb(Err(WStompConnectError::ReconnectionLimit)).await; + break; + } + } else { + backoff.reset(); + } + } + } } }); - Ok(()) + Ok(WStompReconnectHandle { join }) } } @@ -144,6 +249,7 @@ pub(crate) fn headers_for_token(auth_token: impl Into) -> Vec<(String, S async fn inner_connect( uri: Uri, opts: WStompConfigOpts, + auth_token: Option, ) -> Result { let client = if let Some(client) = opts.client { client @@ -165,7 +271,7 @@ async fn inner_connect( let mut headers = opts.additional_headers; - if let Some(auth_token) = opts.auth_token { + if let Some(auth_token) = auth_token { headers.extend(headers_for_token(auth_token)); } diff --git a/src/lib.rs b/src/lib.rs index 2be8667..ae54cc0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,13 +1,16 @@ #![doc = pretty_readme::docify!("README.md", "https://docs.rs/wstomp/latest/wstomp/", "./")] mod config; -pub use config::WStompConfig; +pub use config::{AuthTokenFn, WStompConfig}; mod client; pub use client::{WStompClient, WStompReceiver, WStompSender}; mod connect; -pub use connect::{StompConnect, connect, connect_with_pass, connect_with_token}; +pub use connect::{ + ReconnectControl, StompConnect, WStompReconnectHandle, connect, connect_with_pass, + connect_with_token, +}; mod stomp_handler; diff --git a/src/wstomp_event.rs b/src/wstomp_event.rs index 29b7409..2746739 100644 --- a/src/wstomp_event.rs +++ b/src/wstomp_event.rs @@ -54,7 +54,12 @@ pub enum WStompError { impl std::fmt::Display for WStompConnectError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Self::WsClientError(err) => write!(f, "WebSocket receive error: {}", err), + Self::WsClientError(WsClientError::InvalidResponseStatus(status)) => write!( + f, + "WebSocket connection error: server rejected upgrade with HTTP {}", + status + ), + Self::WsClientError(err) => write!(f, "WebSocket connection error: {}", err), Self::ConnectMessageFailed(msg) => write!(f, "WebSocket receive error: {}", msg), Self::ReconnectionLimit => write!(f, "Reconnection retry limit reached"), }