1673 lines
71 KiB
Rust
1673 lines
71 KiB
Rust
// file: kb-onchain-transport/src/ws_session.rs
|
|
// version: 5
|
|
|
|
//! Persistent multiplexed WebSocket session with bounded reconnect and resubscription.
|
|
|
|
use futures_util::SinkExt; // rust-rules: trait-import
|
|
use futures_util::StreamExt; // rust-rules: trait-import
|
|
|
|
/// Bounded reconnect policy for a persistent WebSocket session.
|
|
#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct WsReconnectPolicy {
|
|
/// Whether reconnect and resubscription are enabled.
|
|
pub enabled: bool,
|
|
/// Maximum connection attempts after one transport failure.
|
|
pub max_attempts: u32,
|
|
/// Delay before the first reconnect attempt.
|
|
pub initial_delay_ms: u64,
|
|
/// Maximum exponential-backoff delay.
|
|
pub max_delay_ms: u64,
|
|
}
|
|
|
|
impl crate::WsReconnectPolicy {
|
|
/// Returns a policy that never reconnects.
|
|
pub const fn disabled() -> Self {
|
|
return Self {
|
|
enabled: false,
|
|
max_attempts: 0,
|
|
initial_delay_ms: 0,
|
|
max_delay_ms: 0,
|
|
};
|
|
}
|
|
|
|
/// Creates a bounded reconnect policy.
|
|
pub fn bounded(
|
|
max_attempts: u32,
|
|
initial_delay_ms: u64,
|
|
max_delay_ms: u64,
|
|
) -> kb_core::Result<Self> {
|
|
if max_attempts == 0 {
|
|
return std::result::Result::Err(kb_core::Error::config(
|
|
"WebSocket reconnect max attempts must be greater than zero",
|
|
));
|
|
}
|
|
if initial_delay_ms == 0 {
|
|
return std::result::Result::Err(kb_core::Error::config(
|
|
"WebSocket reconnect initial delay must be greater than zero",
|
|
));
|
|
}
|
|
if max_delay_ms < initial_delay_ms {
|
|
return std::result::Result::Err(kb_core::Error::config(
|
|
"WebSocket reconnect maximum delay must not be smaller than the initial delay",
|
|
));
|
|
}
|
|
return std::result::Result::Ok(Self {
|
|
enabled: true,
|
|
max_attempts,
|
|
initial_delay_ms,
|
|
max_delay_ms,
|
|
});
|
|
}
|
|
|
|
fn delay_for_attempt(&self, attempt: u32) -> std::time::Duration {
|
|
let exponent = attempt.saturating_sub(1).min(31);
|
|
let multiplier = match 1_u64.checked_shl(exponent) {
|
|
std::option::Option::Some(multiplier) => multiplier,
|
|
std::option::Option::None => u64::MAX,
|
|
};
|
|
let delay = self.initial_delay_ms.saturating_mul(multiplier).min(self.max_delay_ms);
|
|
return std::time::Duration::from_millis(delay);
|
|
}
|
|
}
|
|
|
|
/// Observable lifecycle state of a persistent WebSocket session.
|
|
#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum WsSessionState {
|
|
/// Socket is connected and accepts commands.
|
|
Connected,
|
|
/// Socket failed and a bounded reconnect is running.
|
|
Reconnecting,
|
|
/// Session is permanently disconnected.
|
|
Disconnected,
|
|
}
|
|
|
|
/// One active standard subscription in a persistent session.
|
|
#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct WsSubscriptionSnapshot {
|
|
/// Stable local identifier retained across reconnects.
|
|
pub local_subscription_id: u64,
|
|
/// Current server identifier, absent while resubscription is pending.
|
|
pub remote_subscription_id: std::option::Option<u64>,
|
|
/// Exact subscribe method.
|
|
pub subscribe_method: std::string::String,
|
|
/// Exact unsubscribe method.
|
|
pub unsubscribe_method: std::string::String,
|
|
/// Exact notification method.
|
|
pub notification_method: std::string::String,
|
|
}
|
|
|
|
/// Current snapshot of one persistent WebSocket session.
|
|
#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct WsSessionSnapshot {
|
|
/// Endpoint name.
|
|
pub endpoint_name: std::string::String,
|
|
/// Provider name.
|
|
pub provider: std::string::String,
|
|
/// Endpoint URL.
|
|
pub endpoint_url: std::string::String,
|
|
/// Current lifecycle state.
|
|
pub state: crate::WsSessionState,
|
|
/// Number of successful reconnects since session creation.
|
|
pub reconnect_count: u32,
|
|
/// Effective standard WebSocket capabilities after runtime probing.
|
|
pub capabilities: crate::StandardWsCapabilities,
|
|
/// Active subscriptions ordered by stable local identifier.
|
|
pub subscriptions: std::vec::Vec<crate::WsSubscriptionSnapshot>,
|
|
}
|
|
|
|
/// Successful subscription acknowledgement.
|
|
#[derive(Clone, Debug, PartialEq)]
|
|
pub struct WsSubscriptionAck {
|
|
/// JSON-RPC response received from the node.
|
|
pub response: crate::JsonRpcResponse,
|
|
/// Registered subscription snapshot.
|
|
pub subscription: crate::WsSubscriptionSnapshot,
|
|
}
|
|
|
|
/// Successful explicit unsubscribe acknowledgement.
|
|
#[derive(Clone, Debug, PartialEq)]
|
|
pub struct WsUnsubscribeAck {
|
|
/// JSON-RPC response received from the node.
|
|
pub response: crate::JsonRpcResponse,
|
|
/// Removed subscription snapshot.
|
|
pub subscription: crate::WsSubscriptionSnapshot,
|
|
}
|
|
|
|
/// Event emitted by a persistent WebSocket session.
|
|
#[derive(Clone, Debug, PartialEq)]
|
|
pub enum WsSessionEvent {
|
|
/// Initial socket connection is ready.
|
|
Connected,
|
|
/// Reconnect attempt started.
|
|
Reconnecting {
|
|
/// One-based attempt number.
|
|
attempt: u32,
|
|
/// Configured attempt limit.
|
|
maximum_attempts: u32,
|
|
},
|
|
/// Socket reconnected and active subscriptions are being restored.
|
|
Reconnected {
|
|
/// Successful reconnect count.
|
|
reconnect_count: u32,
|
|
},
|
|
/// New subscription was acknowledged.
|
|
SubscriptionAdded(crate::WsSubscriptionSnapshot),
|
|
/// Existing subscription received a new remote identifier after reconnect.
|
|
SubscriptionRemapped {
|
|
/// Stable local identifier.
|
|
local_subscription_id: u64,
|
|
/// Previous remote identifier.
|
|
previous_remote_subscription_id: std::option::Option<u64>,
|
|
/// New remote identifier.
|
|
remote_subscription_id: u64,
|
|
},
|
|
/// Subscription was removed explicitly or completed by a one-shot server contract.
|
|
SubscriptionRemoved(crate::WsSubscriptionSnapshot),
|
|
/// Typed standard notification.
|
|
Notification {
|
|
/// Subscription that received the notification.
|
|
subscription: crate::WsSubscriptionSnapshot,
|
|
/// Typed notification payload.
|
|
notification: std::boxed::Box<crate::StandardWsNotification>,
|
|
/// Original JSON-RPC notification envelope.
|
|
raw: crate::JsonRpcNotification,
|
|
},
|
|
/// Protocol or transport diagnostic.
|
|
Diagnostic {
|
|
/// Stable error family or diagnostic code.
|
|
code: std::string::String,
|
|
/// Human-readable diagnostic.
|
|
message: std::string::String,
|
|
},
|
|
/// Session ended and will not reconnect.
|
|
Disconnected,
|
|
}
|
|
|
|
/// Persistent multiplexed standard Solana WebSocket session.
|
|
pub struct WsSession {
|
|
endpoint: kb_config::WsEndpointConfig,
|
|
capabilities: std::sync::Arc<tokio::sync::RwLock<crate::StandardWsCapabilities>>,
|
|
control_sender: tokio::sync::mpsc::Sender<WsSessionCommand>,
|
|
event_sender: tokio::sync::broadcast::Sender<crate::WsSessionEvent>,
|
|
snapshot: std::sync::Arc<tokio::sync::RwLock<crate::WsSessionSnapshot>>,
|
|
}
|
|
|
|
impl std::fmt::Debug for crate::WsSession {
|
|
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
return formatter
|
|
.debug_struct("WsSession")
|
|
.field("endpoint_name", &self.endpoint.name)
|
|
.field("provider", &self.endpoint.provider)
|
|
.field("endpoint_url", &self.endpoint.url)
|
|
.finish_non_exhaustive();
|
|
}
|
|
}
|
|
|
|
impl crate::WsSession {
|
|
/// Connects one persistent session and starts its multiplexing task.
|
|
///
|
|
/// Enabled unstable methods are probed by their first real subscription attempt. When the
|
|
/// node explicitly rejects one as absent or disabled, that capability is disabled for the
|
|
/// remaining lifetime of this session.
|
|
pub async fn connect(
|
|
client: crate::WsClient,
|
|
capabilities: crate::StandardWsCapabilities,
|
|
reconnect_policy: crate::WsReconnectPolicy,
|
|
) -> kb_core::Result<Self> {
|
|
let stream = match connect_stream(&client).await {
|
|
std::result::Result::Ok(stream) => stream,
|
|
std::result::Result::Err(error) => return std::result::Result::Err(error),
|
|
};
|
|
let endpoint = client.endpoint_config().clone();
|
|
let channel_capacity = match usize::try_from(endpoint.write_channel_capacity) {
|
|
std::result::Result::Ok(capacity) if capacity > 0 => capacity,
|
|
_ => 1,
|
|
};
|
|
let event_capacity = match usize::try_from(endpoint.event_channel_capacity) {
|
|
std::result::Result::Ok(capacity) if capacity > 0 => capacity,
|
|
_ => 1,
|
|
};
|
|
let (control_sender, control_receiver) = tokio::sync::mpsc::channel(channel_capacity);
|
|
let (event_sender, _event_receiver) = tokio::sync::broadcast::channel(event_capacity);
|
|
let initial_capabilities = capabilities;
|
|
let capabilities = std::sync::Arc::new(tokio::sync::RwLock::new(initial_capabilities));
|
|
let snapshot = std::sync::Arc::new(tokio::sync::RwLock::new(crate::WsSessionSnapshot {
|
|
endpoint_name: endpoint.name.clone(),
|
|
provider: endpoint.provider.clone(),
|
|
endpoint_url: endpoint.url.clone(),
|
|
state: crate::WsSessionState::Connected,
|
|
reconnect_count: 0,
|
|
capabilities: initial_capabilities,
|
|
subscriptions: std::vec::Vec::new(),
|
|
}));
|
|
let runtime = WsSessionRuntime {
|
|
client,
|
|
reconnect_policy,
|
|
capabilities: capabilities.clone(),
|
|
control_receiver,
|
|
event_sender: event_sender.clone(),
|
|
snapshot: snapshot.clone(),
|
|
active: std::collections::BTreeMap::new(),
|
|
remote_to_local: std::collections::BTreeMap::new(),
|
|
pending: std::collections::BTreeMap::new(),
|
|
next_local_subscription_id: 1,
|
|
reconnect_count: 0,
|
|
};
|
|
tokio::spawn(async move {
|
|
runtime.run(stream).await;
|
|
});
|
|
let _send_result = event_sender.send(crate::WsSessionEvent::Connected);
|
|
tracing::info!(target: crate::TRACING_TARGET, action = "connect_ws_session", endpoint_name = %endpoint.name, provider = %endpoint.provider, "persistent WebSocket session connected");
|
|
return std::result::Result::Ok(Self {
|
|
endpoint,
|
|
capabilities,
|
|
control_sender,
|
|
event_sender,
|
|
snapshot,
|
|
});
|
|
}
|
|
|
|
/// Returns a receiver for session events.
|
|
pub fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver<crate::WsSessionEvent> {
|
|
return self.event_sender.subscribe();
|
|
}
|
|
|
|
/// Returns the current session snapshot.
|
|
pub async fn snapshot(&self) -> crate::WsSessionSnapshot {
|
|
return self.snapshot.read().await.clone();
|
|
}
|
|
|
|
/// Subscribes one typed standard request on the existing socket.
|
|
///
|
|
/// For an enabled unstable method, the first request also verifies actual node support.
|
|
pub async fn subscribe(
|
|
&self,
|
|
request: crate::StandardWsRequest,
|
|
) -> kb_core::Result<crate::WsSubscriptionAck> {
|
|
let capabilities = *self.capabilities.read().await;
|
|
let params = match request.params(capabilities) {
|
|
std::result::Result::Ok(params) => params,
|
|
std::result::Result::Err(error) => return std::result::Result::Err(error),
|
|
};
|
|
let (response_sender, response_receiver) = tokio::sync::oneshot::channel();
|
|
let command = WsSessionCommand::Subscribe { request, params, response_sender };
|
|
let send_result = self.control_sender.send(command).await;
|
|
if send_result.is_err() {
|
|
return std::result::Result::Err(kb_core::Error::not_connected(format!(
|
|
"WebSocket session '{}' is not accepting subscriptions",
|
|
self.endpoint.name
|
|
)));
|
|
}
|
|
return match response_receiver.await {
|
|
std::result::Result::Ok(result) => result,
|
|
std::result::Result::Err(error) => {
|
|
std::result::Result::Err(kb_core::Error::not_connected(format!(
|
|
"WebSocket subscription response channel closed for endpoint '{}': {error}",
|
|
self.endpoint.name
|
|
)))
|
|
},
|
|
};
|
|
}
|
|
|
|
/// Explicitly unsubscribes one current remote subscription identifier.
|
|
pub async fn unsubscribe(
|
|
&self,
|
|
remote_subscription_id: u64,
|
|
) -> kb_core::Result<WsUnsubscribeAck> {
|
|
let (response_sender, response_receiver) = tokio::sync::oneshot::channel();
|
|
let command = WsSessionCommand::Unsubscribe { remote_subscription_id, response_sender };
|
|
let send_result = self.control_sender.send(command).await;
|
|
if send_result.is_err() {
|
|
return std::result::Result::Err(kb_core::Error::not_connected(format!(
|
|
"WebSocket session '{}' is not accepting unsubscribe commands",
|
|
self.endpoint.name
|
|
)));
|
|
}
|
|
return match response_receiver.await {
|
|
std::result::Result::Ok(result) => result,
|
|
std::result::Result::Err(error) => {
|
|
std::result::Result::Err(kb_core::Error::not_connected(format!(
|
|
"WebSocket unsubscribe response channel closed for endpoint '{}': {error}",
|
|
self.endpoint.name
|
|
)))
|
|
},
|
|
};
|
|
}
|
|
|
|
/// Unsubscribes active subscriptions, closes the socket and stops the session task.
|
|
pub async fn disconnect(&self) -> kb_core::Result<()> {
|
|
let (response_sender, response_receiver) = tokio::sync::oneshot::channel();
|
|
let send_result =
|
|
self.control_sender.send(WsSessionCommand::Disconnect { response_sender }).await;
|
|
if send_result.is_err() {
|
|
return std::result::Result::Ok(());
|
|
}
|
|
return match response_receiver.await {
|
|
std::result::Result::Ok(result) => result,
|
|
std::result::Result::Err(_) => std::result::Result::Ok(()),
|
|
};
|
|
}
|
|
}
|
|
|
|
type WsStream =
|
|
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
|
|
|
|
enum WsSessionCommand {
|
|
Subscribe {
|
|
request: crate::StandardWsRequest,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<crate::WsSubscriptionAck>>,
|
|
},
|
|
Unsubscribe {
|
|
remote_subscription_id: u64,
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<WsUnsubscribeAck>>,
|
|
},
|
|
Disconnect {
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<()>>,
|
|
},
|
|
}
|
|
|
|
struct ActiveWsSubscription {
|
|
request: crate::StandardWsRequest,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
snapshot: crate::WsSubscriptionSnapshot,
|
|
}
|
|
|
|
enum PendingWsRequest {
|
|
Subscribe {
|
|
local_subscription_id: u64,
|
|
request: crate::StandardWsRequest,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
deadline: tokio::time::Instant,
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<crate::WsSubscriptionAck>>,
|
|
},
|
|
Unsubscribe {
|
|
local_subscription_id: u64,
|
|
remote_subscription_id: u64,
|
|
deadline: tokio::time::Instant,
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<WsUnsubscribeAck>>,
|
|
},
|
|
Resubscribe {
|
|
local_subscription_id: u64,
|
|
previous_remote_subscription_id: std::option::Option<u64>,
|
|
deadline: tokio::time::Instant,
|
|
},
|
|
}
|
|
|
|
struct WsSessionRuntime {
|
|
client: crate::WsClient,
|
|
reconnect_policy: crate::WsReconnectPolicy,
|
|
capabilities: std::sync::Arc<tokio::sync::RwLock<crate::StandardWsCapabilities>>,
|
|
control_receiver: tokio::sync::mpsc::Receiver<WsSessionCommand>,
|
|
event_sender: tokio::sync::broadcast::Sender<crate::WsSessionEvent>,
|
|
snapshot: std::sync::Arc<tokio::sync::RwLock<crate::WsSessionSnapshot>>,
|
|
active: std::collections::BTreeMap<u64, ActiveWsSubscription>,
|
|
remote_to_local: std::collections::BTreeMap<u64, u64>,
|
|
pending: std::collections::BTreeMap<u64, PendingWsRequest>,
|
|
next_local_subscription_id: u64,
|
|
reconnect_count: u32,
|
|
}
|
|
|
|
impl WsSessionRuntime {
|
|
async fn run(mut self, mut stream: WsStream) {
|
|
let mut timeout_tick = tokio::time::interval(std::time::Duration::from_millis(50));
|
|
timeout_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
loop {
|
|
tokio::select! {
|
|
command = self.control_receiver.recv() => {
|
|
let should_continue = match command {
|
|
std::option::Option::Some(command) => self.handle_command(&mut stream, command).await,
|
|
std::option::Option::None => false,
|
|
};
|
|
if !should_continue {
|
|
break;
|
|
}
|
|
},
|
|
message = stream.next() => {
|
|
let transport_ok = self.handle_stream_message(&mut stream, message).await;
|
|
if !transport_ok {
|
|
let reconnect = self.reconnect().await;
|
|
match reconnect {
|
|
std::option::Option::Some(reconnected_stream) => {
|
|
stream = reconnected_stream;
|
|
},
|
|
std::option::Option::None => break,
|
|
}
|
|
}
|
|
},
|
|
_ = timeout_tick.tick() => {
|
|
self.expire_pending().await;
|
|
},
|
|
}
|
|
}
|
|
self.finish_disconnected().await;
|
|
}
|
|
|
|
async fn handle_command(&mut self, stream: &mut WsStream, command: WsSessionCommand) -> bool {
|
|
return match command {
|
|
WsSessionCommand::Subscribe { request, params, response_sender } => {
|
|
self.send_subscribe(stream, request, params, response_sender).await;
|
|
true
|
|
},
|
|
WsSessionCommand::Unsubscribe { remote_subscription_id, response_sender } => {
|
|
self.send_unsubscribe(stream, remote_subscription_id, response_sender).await;
|
|
true
|
|
},
|
|
WsSessionCommand::Disconnect { response_sender } => {
|
|
let result = self.close_stream(stream).await;
|
|
let _send_result = response_sender.send(result);
|
|
false
|
|
},
|
|
};
|
|
}
|
|
|
|
async fn send_subscribe(
|
|
&mut self,
|
|
stream: &mut WsStream,
|
|
request: crate::StandardWsRequest,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<crate::WsSubscriptionAck>>,
|
|
) {
|
|
let local_subscription_id = self.next_local_subscription_id;
|
|
self.next_local_subscription_id = self.next_local_subscription_id.saturating_add(1);
|
|
let json_request = self
|
|
.client
|
|
.build_json_rpc_request(request.subscribe_method().to_string(), params.clone());
|
|
let request_id = match json_request.id.as_u64() {
|
|
std::option::Option::Some(request_id) => request_id,
|
|
std::option::Option::None => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(
|
|
kb_core::Error::invalid_state("numeric WebSocket request id is required"),
|
|
));
|
|
return;
|
|
},
|
|
};
|
|
let text = match json_request.to_json_string() {
|
|
std::result::Result::Ok(text) => text,
|
|
std::result::Result::Err(error) => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(error));
|
|
return;
|
|
},
|
|
};
|
|
let send_result =
|
|
stream.send(tokio_tungstenite::tungstenite::Message::Text(text.into())).await;
|
|
if let std::result::Result::Err(error) = send_result {
|
|
let _send_result = response_sender.send(std::result::Result::Err(kb_core::Error::ws(
|
|
format!("cannot send WebSocket subscribe request: {error}"),
|
|
)));
|
|
return;
|
|
}
|
|
self.pending.insert(
|
|
request_id,
|
|
PendingWsRequest::Subscribe {
|
|
local_subscription_id,
|
|
request,
|
|
params,
|
|
deadline: tokio::time::Instant::now()
|
|
+ std::time::Duration::from_millis(
|
|
self.client.endpoint_config().request_timeout_ms,
|
|
),
|
|
response_sender,
|
|
},
|
|
);
|
|
}
|
|
|
|
async fn send_unsubscribe(
|
|
&mut self,
|
|
stream: &mut WsStream,
|
|
remote_subscription_id: u64,
|
|
response_sender: tokio::sync::oneshot::Sender<kb_core::Result<WsUnsubscribeAck>>,
|
|
) {
|
|
let local_subscription_id = match self.remote_to_local.get(&remote_subscription_id) {
|
|
std::option::Option::Some(local_subscription_id) => *local_subscription_id,
|
|
std::option::Option::None => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(
|
|
kb_core::Error::invalid_state(format!(
|
|
"unknown remote WebSocket subscription id {remote_subscription_id}"
|
|
)),
|
|
));
|
|
return;
|
|
},
|
|
};
|
|
let subscription = match self.active.get(&local_subscription_id) {
|
|
std::option::Option::Some(subscription) => subscription,
|
|
std::option::Option::None => {
|
|
let _send_result =
|
|
response_sender.send(std::result::Result::Err(kb_core::Error::invalid_state(
|
|
format!("missing local WebSocket subscription id {local_subscription_id}"),
|
|
)));
|
|
return;
|
|
},
|
|
};
|
|
let json_request = self.client.build_json_rpc_request(
|
|
subscription.snapshot.unsubscribe_method.clone(),
|
|
std::vec![serde_json::Value::from(remote_subscription_id)],
|
|
);
|
|
let request_id = match json_request.id.as_u64() {
|
|
std::option::Option::Some(request_id) => request_id,
|
|
std::option::Option::None => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(
|
|
kb_core::Error::invalid_state("numeric WebSocket request id is required"),
|
|
));
|
|
return;
|
|
},
|
|
};
|
|
let text = match json_request.to_json_string() {
|
|
std::result::Result::Ok(text) => text,
|
|
std::result::Result::Err(error) => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(error));
|
|
return;
|
|
},
|
|
};
|
|
let send_result =
|
|
stream.send(tokio_tungstenite::tungstenite::Message::Text(text.into())).await;
|
|
if let std::result::Result::Err(error) = send_result {
|
|
let _send_result = response_sender.send(std::result::Result::Err(kb_core::Error::ws(
|
|
format!("cannot send WebSocket unsubscribe request: {error}"),
|
|
)));
|
|
return;
|
|
}
|
|
self.pending.insert(
|
|
request_id,
|
|
PendingWsRequest::Unsubscribe {
|
|
local_subscription_id,
|
|
remote_subscription_id,
|
|
deadline: tokio::time::Instant::now()
|
|
+ std::time::Duration::from_millis(
|
|
self.client.endpoint_config().unsubscribe_timeout_ms,
|
|
),
|
|
response_sender,
|
|
},
|
|
);
|
|
}
|
|
|
|
async fn handle_stream_message(
|
|
&mut self,
|
|
stream: &mut WsStream,
|
|
message: std::option::Option<
|
|
std::result::Result<
|
|
tokio_tungstenite::tungstenite::Message,
|
|
tokio_tungstenite::tungstenite::Error,
|
|
>,
|
|
>,
|
|
) -> bool {
|
|
let message = match message {
|
|
std::option::Option::Some(std::result::Result::Ok(message)) => message,
|
|
std::option::Option::Some(std::result::Result::Err(error)) => {
|
|
self.emit_diagnostic("ws_read", format!("WebSocket read failed: {error}"));
|
|
return false;
|
|
},
|
|
std::option::Option::None => {
|
|
self.emit_diagnostic("ws_closed", "WebSocket stream ended".to_string());
|
|
return false;
|
|
},
|
|
};
|
|
return match message {
|
|
tokio_tungstenite::tungstenite::Message::Text(text) => {
|
|
self.handle_text(text.as_str()).await;
|
|
true
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Ping(payload) => {
|
|
let pong_result =
|
|
stream.send(tokio_tungstenite::tungstenite::Message::Pong(payload)).await;
|
|
if let std::result::Result::Err(error) = pong_result {
|
|
self.emit_diagnostic("ws_pong", format!("WebSocket pong failed: {error}"));
|
|
return false;
|
|
}
|
|
true
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Close(_) => false,
|
|
tokio_tungstenite::tungstenite::Message::Binary(_)
|
|
| tokio_tungstenite::tungstenite::Message::Pong(_)
|
|
| tokio_tungstenite::tungstenite::Message::Frame(_) => true,
|
|
};
|
|
}
|
|
|
|
async fn handle_text(&mut self, text: &str) {
|
|
let response = match crate::parse_json_rpc_text(text) {
|
|
std::result::Result::Ok(response) => response,
|
|
std::result::Result::Err(error) => {
|
|
self.emit_error(error);
|
|
return;
|
|
},
|
|
};
|
|
match response {
|
|
crate::JsonRpcResponse::Success(success) => {
|
|
self.handle_success(success).await;
|
|
},
|
|
crate::JsonRpcResponse::Error(error_response) => {
|
|
self.handle_error_response(error_response).await;
|
|
},
|
|
crate::JsonRpcResponse::Notification(notification) => {
|
|
self.handle_notification(notification).await;
|
|
},
|
|
}
|
|
}
|
|
|
|
async fn handle_success(&mut self, success: crate::JsonRpcSuccessResponse) {
|
|
let request_id = match success.id.as_u64() {
|
|
std::option::Option::Some(request_id) => request_id,
|
|
std::option::Option::None => {
|
|
self.emit_diagnostic(
|
|
"ws_response_id",
|
|
"WebSocket success response has no numeric request id".to_string(),
|
|
);
|
|
return;
|
|
},
|
|
};
|
|
let pending = match self.pending.remove(&request_id) {
|
|
std::option::Option::Some(pending) => pending,
|
|
std::option::Option::None => {
|
|
self.emit_diagnostic(
|
|
"ws_unmatched_response",
|
|
format!("unmatched WebSocket response id {request_id}"),
|
|
);
|
|
return;
|
|
},
|
|
};
|
|
let response = crate::JsonRpcResponse::Success(success.clone());
|
|
match pending {
|
|
PendingWsRequest::Subscribe {
|
|
local_subscription_id,
|
|
request,
|
|
params,
|
|
deadline: _,
|
|
response_sender,
|
|
} => {
|
|
let remote_subscription_id = match success.result.as_u64() {
|
|
std::option::Option::Some(value) => value,
|
|
std::option::Option::None => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(
|
|
kb_core::Error::json("subscribe result must be a numeric id"),
|
|
));
|
|
return;
|
|
},
|
|
};
|
|
let specification = request.specification();
|
|
let snapshot = crate::WsSubscriptionSnapshot {
|
|
local_subscription_id,
|
|
remote_subscription_id: std::option::Option::Some(remote_subscription_id),
|
|
subscribe_method: specification.subscribe_method.to_string(),
|
|
unsubscribe_method: specification.unsubscribe_method.to_string(),
|
|
notification_method: specification.notification_method.to_string(),
|
|
};
|
|
self.remote_to_local.insert(remote_subscription_id, local_subscription_id);
|
|
self.active.insert(
|
|
local_subscription_id,
|
|
ActiveWsSubscription {
|
|
request,
|
|
params,
|
|
snapshot: snapshot.clone(),
|
|
},
|
|
);
|
|
self.refresh_snapshot().await;
|
|
let _event_result = self
|
|
.event_sender
|
|
.send(crate::WsSessionEvent::SubscriptionAdded(snapshot.clone()));
|
|
let _send_result =
|
|
response_sender.send(std::result::Result::Ok(crate::WsSubscriptionAck {
|
|
response,
|
|
subscription: snapshot,
|
|
}));
|
|
},
|
|
PendingWsRequest::Unsubscribe {
|
|
local_subscription_id,
|
|
remote_subscription_id,
|
|
deadline: _,
|
|
response_sender,
|
|
} => {
|
|
if success.result.as_bool() != std::option::Option::Some(true) {
|
|
let _send_result = response_sender.send(std::result::Result::Err(
|
|
kb_core::Error::json("unsubscribe result must be true"),
|
|
));
|
|
return;
|
|
}
|
|
self.remote_to_local.remove(&remote_subscription_id);
|
|
let removed = self.active.remove(&local_subscription_id);
|
|
let snapshot = match removed {
|
|
std::option::Option::Some(subscription) => subscription.snapshot,
|
|
std::option::Option::None => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(
|
|
kb_core::Error::invalid_state(
|
|
"acknowledged subscription disappeared before removal",
|
|
),
|
|
));
|
|
return;
|
|
},
|
|
};
|
|
self.refresh_snapshot().await;
|
|
let _event_result = self
|
|
.event_sender
|
|
.send(crate::WsSessionEvent::SubscriptionRemoved(snapshot.clone()));
|
|
let _send_result =
|
|
response_sender.send(std::result::Result::Ok(WsUnsubscribeAck {
|
|
response,
|
|
subscription: snapshot,
|
|
}));
|
|
},
|
|
PendingWsRequest::Resubscribe {
|
|
local_subscription_id,
|
|
previous_remote_subscription_id,
|
|
deadline: _,
|
|
} => {
|
|
let remote_subscription_id = match success.result.as_u64() {
|
|
std::option::Option::Some(value) => value,
|
|
std::option::Option::None => {
|
|
self.emit_diagnostic(
|
|
"ws_resubscribe_result",
|
|
"resubscribe result must be a numeric id".to_string(),
|
|
);
|
|
return;
|
|
},
|
|
};
|
|
let remapped = match self.active.get_mut(&local_subscription_id) {
|
|
std::option::Option::Some(active) => {
|
|
active.snapshot.remote_subscription_id =
|
|
std::option::Option::Some(remote_subscription_id);
|
|
true
|
|
},
|
|
std::option::Option::None => false,
|
|
};
|
|
if remapped {
|
|
self.remote_to_local.insert(remote_subscription_id, local_subscription_id);
|
|
let _event_result =
|
|
self.event_sender.send(crate::WsSessionEvent::SubscriptionRemapped {
|
|
local_subscription_id,
|
|
previous_remote_subscription_id,
|
|
remote_subscription_id,
|
|
});
|
|
self.refresh_snapshot().await;
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
async fn handle_error_response(&mut self, error_response: crate::JsonRpcErrorResponse) {
|
|
let request_id = error_response.id.as_u64();
|
|
let error = kb_core::Error::ws(format!(
|
|
"WebSocket JSON-RPC error {}: {}",
|
|
error_response.error.code, error_response.error.message
|
|
));
|
|
if let std::option::Option::Some(request_id) = request_id {
|
|
if let std::option::Option::Some(pending) = self.pending.remove(&request_id) {
|
|
match pending {
|
|
PendingWsRequest::Subscribe { request, response_sender, .. } => {
|
|
self.disable_unsupported_capability(&request, &error_response).await;
|
|
let _send_result = response_sender.send(std::result::Result::Err(error));
|
|
},
|
|
PendingWsRequest::Unsubscribe { response_sender, .. } => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(error));
|
|
},
|
|
PendingWsRequest::Resubscribe { local_subscription_id, .. } => {
|
|
let request = self
|
|
.active
|
|
.get(&local_subscription_id)
|
|
.map(|active| return active.request.clone());
|
|
let disabled = match request {
|
|
std::option::Option::Some(request) => {
|
|
self.disable_unsupported_capability(&request, &error_response).await
|
|
},
|
|
std::option::Option::None => false,
|
|
};
|
|
if disabled {
|
|
let removed = self.active.remove(&local_subscription_id);
|
|
if let std::option::Option::Some(removed) = removed {
|
|
self.refresh_snapshot().await;
|
|
let _event_result = self.event_sender.send(
|
|
crate::WsSessionEvent::SubscriptionRemoved(removed.snapshot),
|
|
);
|
|
}
|
|
} else {
|
|
self.emit_error(error);
|
|
}
|
|
},
|
|
}
|
|
return;
|
|
}
|
|
}
|
|
self.emit_error(error);
|
|
}
|
|
|
|
async fn disable_unsupported_capability(
|
|
&mut self,
|
|
request: &crate::StandardWsRequest,
|
|
error_response: &crate::JsonRpcErrorResponse,
|
|
) -> bool {
|
|
let method = request.subscribe_method();
|
|
if !crate::StandardWsCapabilities::is_unstable_method(method)
|
|
|| !is_method_unavailable_error(error_response)
|
|
{
|
|
return false;
|
|
}
|
|
let disabled = {
|
|
let mut capabilities = self.capabilities.write().await;
|
|
capabilities.disable_method(method)
|
|
};
|
|
if disabled {
|
|
self.refresh_snapshot().await;
|
|
self.emit_diagnostic(
|
|
"ws_capability_disabled",
|
|
format!(
|
|
"endpoint '{}' rejected unstable method '{method}'; capability disabled for this session: JSON-RPC {} {}",
|
|
self.client.endpoint_name(),
|
|
error_response.error.code,
|
|
error_response.error.message
|
|
),
|
|
);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
async fn handle_notification(&mut self, notification: crate::JsonRpcNotification) {
|
|
let remote_subscription_id = notification.params.subscription;
|
|
let local_subscription_id = match self.remote_to_local.get(&remote_subscription_id) {
|
|
std::option::Option::Some(value) => *value,
|
|
std::option::Option::None => {
|
|
self.emit_diagnostic(
|
|
"ws_unknown_subscription",
|
|
format!(
|
|
"notification references unknown remote subscription id {remote_subscription_id}"
|
|
),
|
|
);
|
|
return;
|
|
},
|
|
};
|
|
let subscription = match self.active.get(&local_subscription_id) {
|
|
std::option::Option::Some(active) => active.snapshot.clone(),
|
|
std::option::Option::None => return,
|
|
};
|
|
if notification.method != subscription.notification_method {
|
|
self.emit_diagnostic(
|
|
"ws_notification_method",
|
|
format!(
|
|
"notification method '{}' does not match expected '{}'",
|
|
notification.method, subscription.notification_method
|
|
),
|
|
);
|
|
return;
|
|
}
|
|
let typed = match crate::adapt_standard_ws_notification(¬ification) {
|
|
std::result::Result::Ok(typed) => typed,
|
|
std::result::Result::Err(error) => {
|
|
self.emit_error(error);
|
|
return;
|
|
},
|
|
};
|
|
let terminal_signature = matches!(
|
|
&typed,
|
|
crate::StandardWsNotification::Signature(crate::WsSignatureNotification {
|
|
result: crate::RpcResponse {
|
|
value: crate::WsSignatureValue::Status { .. },
|
|
..
|
|
},
|
|
..
|
|
})
|
|
);
|
|
let _event_result = self.event_sender.send(crate::WsSessionEvent::Notification {
|
|
subscription: subscription.clone(),
|
|
notification: std::boxed::Box::new(typed),
|
|
raw: notification,
|
|
});
|
|
if terminal_signature {
|
|
self.remote_to_local.remove(&remote_subscription_id);
|
|
self.active.remove(&local_subscription_id);
|
|
self.refresh_snapshot().await;
|
|
let _event_result =
|
|
self.event_sender.send(crate::WsSessionEvent::SubscriptionRemoved(subscription));
|
|
}
|
|
}
|
|
|
|
async fn expire_pending(&mut self) {
|
|
let now = tokio::time::Instant::now();
|
|
let mut expired = std::vec::Vec::new();
|
|
for (request_id, pending) in &self.pending {
|
|
let deadline = match pending {
|
|
PendingWsRequest::Subscribe { deadline, .. }
|
|
| PendingWsRequest::Unsubscribe { deadline, .. }
|
|
| PendingWsRequest::Resubscribe { deadline, .. } => *deadline,
|
|
};
|
|
if deadline <= now {
|
|
expired.push(*request_id);
|
|
}
|
|
}
|
|
return for request_id in expired {
|
|
let pending = self.pending.remove(&request_id);
|
|
if let std::option::Option::Some(pending) = pending {
|
|
let error =
|
|
kb_core::Error::ws(format!("WebSocket request id {request_id} timed out"));
|
|
match pending {
|
|
PendingWsRequest::Subscribe { response_sender, .. } => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(error));
|
|
},
|
|
PendingWsRequest::Unsubscribe { response_sender, .. } => {
|
|
let _send_result = response_sender.send(std::result::Result::Err(error));
|
|
},
|
|
PendingWsRequest::Resubscribe { .. } => {
|
|
self.emit_error(error);
|
|
},
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
async fn reconnect(&mut self) -> std::option::Option<WsStream> {
|
|
self.fail_pending(kb_core::Error::ws(
|
|
"WebSocket transport failed before request completion",
|
|
));
|
|
if !self.reconnect_policy.enabled {
|
|
return std::option::Option::None;
|
|
}
|
|
let mut previous_remote_ids = std::collections::BTreeMap::new();
|
|
self.remote_to_local.clear();
|
|
for (local_subscription_id, active) in &mut self.active {
|
|
previous_remote_ids
|
|
.insert(*local_subscription_id, active.snapshot.remote_subscription_id);
|
|
active.snapshot.remote_subscription_id = std::option::Option::None;
|
|
}
|
|
self.set_state(crate::WsSessionState::Reconnecting).await;
|
|
let mut attempt = 1_u32;
|
|
while attempt <= self.reconnect_policy.max_attempts {
|
|
let _event_result = self.event_sender.send(crate::WsSessionEvent::Reconnecting {
|
|
attempt,
|
|
maximum_attempts: self.reconnect_policy.max_attempts,
|
|
});
|
|
tokio::time::sleep(self.reconnect_policy.delay_for_attempt(attempt)).await;
|
|
let stream_result = connect_stream(&self.client).await;
|
|
match stream_result {
|
|
std::result::Result::Ok(mut stream) => {
|
|
self.reconnect_count = self.reconnect_count.saturating_add(1);
|
|
self.set_state(crate::WsSessionState::Connected).await;
|
|
let _event_result =
|
|
self.event_sender.send(crate::WsSessionEvent::Reconnected {
|
|
reconnect_count: self.reconnect_count,
|
|
});
|
|
let restore_result =
|
|
self.restore_subscriptions(&mut stream, &previous_remote_ids).await;
|
|
if let std::result::Result::Err(error) = restore_result {
|
|
self.emit_error(error);
|
|
attempt = attempt.saturating_add(1);
|
|
continue;
|
|
}
|
|
return std::option::Option::Some(stream);
|
|
},
|
|
std::result::Result::Err(error) => self.emit_error(error),
|
|
}
|
|
attempt = attempt.saturating_add(1);
|
|
}
|
|
return std::option::Option::None;
|
|
}
|
|
|
|
async fn restore_subscriptions(
|
|
&mut self,
|
|
stream: &mut WsStream,
|
|
previous_remote_ids: &std::collections::BTreeMap<u64, std::option::Option<u64>>,
|
|
) -> kb_core::Result<()> {
|
|
let local_ids = self.active.keys().copied().collect::<std::vec::Vec<_>>();
|
|
for local_subscription_id in local_ids {
|
|
let (subscribe_method, params, previous_remote_subscription_id) =
|
|
match self.active.get(&local_subscription_id) {
|
|
std::option::Option::Some(active) => (
|
|
active.request.subscribe_method().to_string(),
|
|
active.params.clone(),
|
|
previous_remote_ids.get(&local_subscription_id).copied().flatten(),
|
|
),
|
|
std::option::Option::None => continue,
|
|
};
|
|
let json_request = self.client.build_json_rpc_request(subscribe_method, params);
|
|
let request_id = match json_request.id.as_u64() {
|
|
std::option::Option::Some(value) => value,
|
|
std::option::Option::None => {
|
|
return std::result::Result::Err(kb_core::Error::invalid_state(
|
|
"numeric WebSocket request id is required",
|
|
));
|
|
},
|
|
};
|
|
let text = match json_request.to_json_string() {
|
|
std::result::Result::Ok(text) => text,
|
|
std::result::Result::Err(error) => return std::result::Result::Err(error),
|
|
};
|
|
let send_result =
|
|
stream.send(tokio_tungstenite::tungstenite::Message::Text(text.into())).await;
|
|
if let std::result::Result::Err(error) = send_result {
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"cannot restore WebSocket subscription: {error}"
|
|
)));
|
|
}
|
|
self.pending.insert(
|
|
request_id,
|
|
PendingWsRequest::Resubscribe {
|
|
local_subscription_id,
|
|
previous_remote_subscription_id,
|
|
deadline: tokio::time::Instant::now()
|
|
+ std::time::Duration::from_millis(
|
|
self.client.endpoint_config().request_timeout_ms,
|
|
),
|
|
},
|
|
);
|
|
}
|
|
return std::result::Result::Ok(());
|
|
}
|
|
|
|
fn fail_pending(&mut self, error: kb_core::Error) {
|
|
let pending = std::mem::take(&mut self.pending);
|
|
for (_request_id, request) in pending {
|
|
match request {
|
|
PendingWsRequest::Subscribe { response_sender, .. } => {
|
|
let _send_result =
|
|
response_sender.send(std::result::Result::Err(error.clone()));
|
|
},
|
|
PendingWsRequest::Unsubscribe { response_sender, .. } => {
|
|
let _send_result =
|
|
response_sender.send(std::result::Result::Err(error.clone()));
|
|
},
|
|
PendingWsRequest::Resubscribe { .. } => {},
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn close_stream(&mut self, stream: &mut WsStream) -> kb_core::Result<()> {
|
|
let subscriptions = self
|
|
.active
|
|
.values()
|
|
.filter_map(|active| {
|
|
return active.snapshot.remote_subscription_id.map(|remote_id| {
|
|
return (active.snapshot.unsubscribe_method.clone(), remote_id);
|
|
});
|
|
})
|
|
.collect::<std::vec::Vec<_>>();
|
|
for (method, remote_id) in subscriptions {
|
|
let request = self
|
|
.client
|
|
.build_json_rpc_request(method, std::vec![serde_json::Value::from(remote_id)]);
|
|
let text = match request.to_json_string() {
|
|
std::result::Result::Ok(text) => text,
|
|
std::result::Result::Err(error) => {
|
|
self.emit_error(error);
|
|
continue;
|
|
},
|
|
};
|
|
let send_result =
|
|
stream.send(tokio_tungstenite::tungstenite::Message::Text(text.into())).await;
|
|
if let std::result::Result::Err(error) = send_result {
|
|
self.emit_diagnostic(
|
|
"ws_disconnect_unsubscribe",
|
|
format!("cannot send disconnect unsubscribe: {error}"),
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
let close_result = stream
|
|
.send(tokio_tungstenite::tungstenite::Message::Close(std::option::Option::None))
|
|
.await;
|
|
return match close_result {
|
|
std::result::Result::Ok(()) => std::result::Result::Ok(()),
|
|
std::result::Result::Err(error) => std::result::Result::Err(kb_core::Error::ws(
|
|
format!("cannot close WebSocket session: {error}"),
|
|
)),
|
|
};
|
|
}
|
|
|
|
async fn refresh_snapshot(&self) {
|
|
let capabilities = *self.capabilities.read().await;
|
|
let mut snapshot = self.snapshot.write().await;
|
|
snapshot.reconnect_count = self.reconnect_count;
|
|
snapshot.capabilities = capabilities;
|
|
snapshot.subscriptions =
|
|
self.active.values().map(|active| return active.snapshot.clone()).collect();
|
|
}
|
|
|
|
async fn set_state(&self, state: crate::WsSessionState) {
|
|
let capabilities = *self.capabilities.read().await;
|
|
let mut snapshot = self.snapshot.write().await;
|
|
snapshot.state = state;
|
|
snapshot.reconnect_count = self.reconnect_count;
|
|
snapshot.capabilities = capabilities;
|
|
snapshot.subscriptions =
|
|
self.active.values().map(|active| return active.snapshot.clone()).collect();
|
|
}
|
|
|
|
async fn finish_disconnected(&mut self) {
|
|
self.fail_pending(kb_core::Error::not_connected(
|
|
"persistent WebSocket session disconnected",
|
|
));
|
|
self.set_state(crate::WsSessionState::Disconnected).await;
|
|
let _event_result = self.event_sender.send(crate::WsSessionEvent::Disconnected);
|
|
tracing::info!(target: crate::TRACING_TARGET, action = "disconnect_ws_session", endpoint_name = %self.client.endpoint_name(), provider = %self.client.provider(), "persistent WebSocket session disconnected");
|
|
}
|
|
|
|
fn emit_error(&self, error: kb_core::Error) {
|
|
self.emit_diagnostic(error.code(), error.message().to_string());
|
|
}
|
|
|
|
fn emit_diagnostic(&self, code: &str, message: std::string::String) {
|
|
tracing::warn!(target: crate::TRACING_TARGET, action = "ws_session_diagnostic", endpoint_name = %self.client.endpoint_name(), provider = %self.client.provider(), diagnostic_code = code, diagnostic = %message, "persistent WebSocket session diagnostic");
|
|
let _event_result = self
|
|
.event_sender
|
|
.send(crate::WsSessionEvent::Diagnostic { code: code.to_string(), message });
|
|
}
|
|
}
|
|
|
|
fn is_method_unavailable_error(error_response: &crate::JsonRpcErrorResponse) -> bool {
|
|
if error_response.error.code == -32601 {
|
|
return true;
|
|
}
|
|
let message = error_response.error.message.to_ascii_lowercase();
|
|
return message.contains("method not found")
|
|
|| message.contains("method is not available")
|
|
|| message.contains("method not available")
|
|
|| message.contains("method is not enabled")
|
|
|| message.contains("method not enabled")
|
|
|| message.contains("unsupported method")
|
|
|| message.contains("unstable method");
|
|
}
|
|
|
|
async fn connect_stream(client: &crate::WsClient) -> kb_core::Result<WsStream> {
|
|
let connect_future = tokio_tungstenite::connect_async(client.endpoint_url());
|
|
let timeout = std::time::Duration::from_millis(client.endpoint_config().connect_timeout_ms);
|
|
let timeout_result = tokio::time::timeout(timeout, connect_future).await;
|
|
let connect_result = match timeout_result {
|
|
std::result::Result::Ok(connect_result) => connect_result,
|
|
std::result::Result::Err(_) => {
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"WebSocket connect timed out for endpoint '{}'",
|
|
client.endpoint_name()
|
|
)));
|
|
},
|
|
};
|
|
return match connect_result {
|
|
std::result::Result::Ok((stream, _response)) => std::result::Result::Ok(stream),
|
|
std::result::Result::Err(error) => std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"cannot connect WebSocket endpoint '{}': {error}",
|
|
client.endpoint_name()
|
|
))),
|
|
};
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use futures_util::SinkExt; // rust-rules: trait-import
|
|
use futures_util::StreamExt; // rust-rules: trait-import
|
|
|
|
fn endpoint(url: std::string::String) -> kb_config::WsEndpointConfig {
|
|
return kb_config::WsEndpointConfig {
|
|
name: "local-test".to_string(),
|
|
enabled: true,
|
|
provider: "offline".to_string(),
|
|
cluster: "localnet".to_string(),
|
|
url,
|
|
connect_timeout_ms: 1_000,
|
|
request_timeout_ms: 1_000,
|
|
unsubscribe_timeout_ms: 1_000,
|
|
write_channel_capacity: 16,
|
|
event_channel_capacity: 32,
|
|
auto_reconnect: true,
|
|
roles: std::vec![kb_config::EndpointRoleConfig {
|
|
role: "slot_notifications".to_string(),
|
|
enabled: true,
|
|
request_kinds: std::vec!["slot_subscribe".to_string()],
|
|
priority: 1,
|
|
requests_per_second: 100,
|
|
burst_capacity: 100,
|
|
max_concurrent_requests: 8,
|
|
max_subscriptions: 32,
|
|
pause_after_rate_limit_ms: 10,
|
|
}],
|
|
};
|
|
}
|
|
|
|
async fn next_text(
|
|
stream: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>,
|
|
) -> serde_json::Value {
|
|
let message = match stream.next().await {
|
|
std::option::Option::Some(std::result::Result::Ok(message)) => message,
|
|
std::option::Option::Some(std::result::Result::Err(error)) => {
|
|
panic!("server read failed: {error}")
|
|
},
|
|
std::option::Option::None => panic!("server stream ended"),
|
|
};
|
|
let text = match message {
|
|
tokio_tungstenite::tungstenite::Message::Text(text) => text,
|
|
other => panic!("expected text, got {other:?}"),
|
|
};
|
|
return match serde_json::from_str(text.as_str()) {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("server json failed: {error}"),
|
|
};
|
|
}
|
|
|
|
async fn send_json(
|
|
stream: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>,
|
|
value: serde_json::Value,
|
|
) {
|
|
let text = match serde_json::to_string(&value) {
|
|
std::result::Result::Ok(text) => text,
|
|
std::result::Result::Err(error) => panic!("server serialization failed: {error}"),
|
|
};
|
|
if let std::result::Result::Err(error) =
|
|
stream.send(tokio_tungstenite::tungstenite::Message::Text(text.into())).await
|
|
{
|
|
panic!("server send failed: {error}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn reconnect_policy_is_bounded_and_validated() {
|
|
assert!(crate::WsReconnectPolicy::bounded(0, 100, 1000).is_err());
|
|
assert!(crate::WsReconnectPolicy::bounded(3, 0, 1000).is_err());
|
|
assert!(crate::WsReconnectPolicy::bounded(3, 1000, 100).is_err());
|
|
let policy = match crate::WsReconnectPolicy::bounded(3, 100, 250) {
|
|
std::result::Result::Ok(policy) => policy,
|
|
std::result::Result::Err(error) => panic!("policy failed: {error}"),
|
|
};
|
|
assert_eq!(policy.delay_for_attempt(1), std::time::Duration::from_millis(100));
|
|
assert_eq!(policy.delay_for_attempt(2), std::time::Duration::from_millis(200));
|
|
assert_eq!(policy.delay_for_attempt(3), std::time::Duration::from_millis(250));
|
|
}
|
|
|
|
#[test]
|
|
fn disabled_policy_never_reconnects() {
|
|
assert_eq!(
|
|
crate::WsReconnectPolicy::disabled(),
|
|
crate::WsReconnectPolicy {
|
|
enabled: false,
|
|
max_attempts: 0,
|
|
initial_delay_ms: 0,
|
|
max_delay_ms: 0,
|
|
}
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn persistent_session_multiplexes_notification_and_explicit_unsubscribe() {
|
|
let listener = match tokio::net::TcpListener::bind("127.0.0.1:0").await {
|
|
std::result::Result::Ok(listener) => listener,
|
|
std::result::Result::Err(error) => panic!("listener failed: {error}"),
|
|
};
|
|
let address = match listener.local_addr() {
|
|
std::result::Result::Ok(address) => address,
|
|
std::result::Result::Err(error) => panic!("local address failed: {error}"),
|
|
};
|
|
let server = tokio::spawn(async move {
|
|
let (socket, _peer) = match listener.accept().await {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("accept failed: {error}"),
|
|
};
|
|
let mut stream = match tokio_tungstenite::accept_async(socket).await {
|
|
std::result::Result::Ok(stream) => stream,
|
|
std::result::Result::Err(error) => panic!("handshake failed: {error}"),
|
|
};
|
|
let subscribe = next_text(&mut stream).await;
|
|
assert_eq!(subscribe["method"], "slotSubscribe");
|
|
send_json(
|
|
&mut stream,
|
|
serde_json::json!({ "jsonrpc": "2.0", "result": 42, "id": subscribe["id"].clone() }),
|
|
)
|
|
.await;
|
|
send_json(
|
|
&mut stream,
|
|
serde_json::json!({
|
|
"jsonrpc": "2.0",
|
|
"method": "slotNotification",
|
|
"params": {
|
|
"result": { "slot": 8, "parent": 7, "root": 6 },
|
|
"subscription": 42
|
|
}
|
|
}),
|
|
)
|
|
.await;
|
|
let unsubscribe = next_text(&mut stream).await;
|
|
assert_eq!(unsubscribe["method"], "slotUnsubscribe");
|
|
assert_eq!(unsubscribe["params"][0], 42);
|
|
send_json(
|
|
&mut stream,
|
|
serde_json::json!({ "jsonrpc": "2.0", "result": true, "id": unsubscribe["id"].clone() }),
|
|
)
|
|
.await;
|
|
let close = match stream.next().await {
|
|
std::option::Option::Some(std::result::Result::Ok(message)) => message,
|
|
other => panic!("close read failed: {other:?}"),
|
|
};
|
|
assert!(matches!(close, tokio_tungstenite::tungstenite::Message::Close(_)));
|
|
});
|
|
let client = match crate::WsClient::new(endpoint(format!("ws://{address}"))) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client failed: {error}"),
|
|
};
|
|
let session = match crate::WsSession::connect(
|
|
client,
|
|
crate::StandardWsCapabilities::default(),
|
|
crate::WsReconnectPolicy::disabled(),
|
|
)
|
|
.await
|
|
{
|
|
std::result::Result::Ok(session) => session,
|
|
std::result::Result::Err(error) => panic!("session failed: {error}"),
|
|
};
|
|
let mut events = session.subscribe_events();
|
|
let acknowledgement = match session
|
|
.subscribe(crate::StandardWsRequest::Slot(crate::SlotSubscribeRequest))
|
|
.await
|
|
{
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("subscribe failed: {error}"),
|
|
};
|
|
assert_eq!(acknowledgement.subscription.remote_subscription_id, Some(42));
|
|
let notification = tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
|
loop {
|
|
let event = match events.recv().await {
|
|
std::result::Result::Ok(event) => event,
|
|
std::result::Result::Err(error) => panic!("event failed: {error}"),
|
|
};
|
|
if let crate::WsSessionEvent::Notification { notification, .. } = event {
|
|
return notification;
|
|
}
|
|
}
|
|
})
|
|
.await;
|
|
let notification = match notification {
|
|
std::result::Result::Ok(notification) => notification,
|
|
std::result::Result::Err(_) => panic!("notification timed out"),
|
|
};
|
|
assert!(matches!(notification.as_ref(), crate::StandardWsNotification::Slot(_)));
|
|
let unsubscribe = match session.unsubscribe(42).await {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("unsubscribe failed: {error}"),
|
|
};
|
|
assert_eq!(unsubscribe.subscription.remote_subscription_id, Some(42));
|
|
assert!(session.disconnect().await.is_ok());
|
|
if let std::result::Result::Err(error) = server.await {
|
|
panic!("server task failed: {error}");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn unsupported_unstable_method_is_disabled_after_first_server_rejection() {
|
|
let listener = match tokio::net::TcpListener::bind("127.0.0.1:0").await {
|
|
std::result::Result::Ok(listener) => listener,
|
|
std::result::Result::Err(error) => panic!("listener failed: {error}"),
|
|
};
|
|
let address = match listener.local_addr() {
|
|
std::result::Result::Ok(address) => address,
|
|
std::result::Result::Err(error) => panic!("local address failed: {error}"),
|
|
};
|
|
let server = tokio::spawn(async move {
|
|
let (socket, _peer) = match listener.accept().await {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("accept failed: {error}"),
|
|
};
|
|
let mut stream = match tokio_tungstenite::accept_async(socket).await {
|
|
std::result::Result::Ok(stream) => stream,
|
|
std::result::Result::Err(error) => panic!("handshake failed: {error}"),
|
|
};
|
|
let subscribe = next_text(&mut stream).await;
|
|
assert_eq!(subscribe["method"], "blockSubscribe");
|
|
send_json(
|
|
&mut stream,
|
|
serde_json::json!({
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32601,
|
|
"message": "Method not found"
|
|
},
|
|
"id": subscribe["id"].clone()
|
|
}),
|
|
)
|
|
.await;
|
|
let close = match stream.next().await {
|
|
std::option::Option::Some(std::result::Result::Ok(message)) => message,
|
|
other => panic!("close read failed: {other:?}"),
|
|
};
|
|
assert!(matches!(close, tokio_tungstenite::tungstenite::Message::Close(_)));
|
|
});
|
|
let client = match crate::WsClient::new(endpoint(format!("ws://{address}"))) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client failed: {error}"),
|
|
};
|
|
let session = match crate::WsSession::connect(
|
|
client,
|
|
crate::StandardWsCapabilities {
|
|
block_subscribe: true,
|
|
slots_updates_subscribe: false,
|
|
vote_subscribe: false,
|
|
},
|
|
crate::WsReconnectPolicy::disabled(),
|
|
)
|
|
.await
|
|
{
|
|
std::result::Result::Ok(session) => session,
|
|
std::result::Result::Err(error) => panic!("session failed: {error}"),
|
|
};
|
|
let request = crate::StandardWsRequest::Block(crate::BlockSubscribeRequest {
|
|
filter: crate::WsBlockFilter::All,
|
|
config: std::option::Option::None,
|
|
});
|
|
assert!(session.subscribe(request.clone()).await.is_err());
|
|
let snapshot = session.snapshot().await;
|
|
assert!(!snapshot.capabilities.block_subscribe);
|
|
assert!(session.subscribe(request).await.is_err());
|
|
assert!(session.disconnect().await.is_ok());
|
|
if let std::result::Result::Err(error) = server.await {
|
|
panic!("server task failed: {error}");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn terminal_signature_notification_removes_one_shot_subscription() {
|
|
let listener = match tokio::net::TcpListener::bind("127.0.0.1:0").await {
|
|
std::result::Result::Ok(listener) => listener,
|
|
std::result::Result::Err(error) => panic!("listener failed: {error}"),
|
|
};
|
|
let address = match listener.local_addr() {
|
|
std::result::Result::Ok(address) => address,
|
|
std::result::Result::Err(error) => panic!("local address failed: {error}"),
|
|
};
|
|
let server = tokio::spawn(async move {
|
|
let (socket, _peer) = match listener.accept().await {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("accept failed: {error}"),
|
|
};
|
|
let mut stream = match tokio_tungstenite::accept_async(socket).await {
|
|
std::result::Result::Ok(stream) => stream,
|
|
std::result::Result::Err(error) => panic!("handshake failed: {error}"),
|
|
};
|
|
let subscribe = next_text(&mut stream).await;
|
|
assert_eq!(subscribe["method"], "signatureSubscribe");
|
|
send_json(
|
|
&mut stream,
|
|
serde_json::json!({
|
|
"jsonrpc": "2.0",
|
|
"result": 77,
|
|
"id": subscribe["id"].clone()
|
|
}),
|
|
)
|
|
.await;
|
|
send_json(
|
|
&mut stream,
|
|
serde_json::json!({
|
|
"jsonrpc": "2.0",
|
|
"method": "signatureNotification",
|
|
"params": {
|
|
"result": {
|
|
"context": { "slot": 9 },
|
|
"value": { "err": null }
|
|
},
|
|
"subscription": 77
|
|
}
|
|
}),
|
|
)
|
|
.await;
|
|
let close = match stream.next().await {
|
|
std::option::Option::Some(std::result::Result::Ok(message)) => message,
|
|
other => panic!("close read failed: {other:?}"),
|
|
};
|
|
assert!(matches!(close, tokio_tungstenite::tungstenite::Message::Close(_)));
|
|
});
|
|
let client = match crate::WsClient::new(endpoint(format!("ws://{address}"))) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client failed: {error}"),
|
|
};
|
|
let session = match crate::WsSession::connect(
|
|
client,
|
|
crate::StandardWsCapabilities::default(),
|
|
crate::WsReconnectPolicy::disabled(),
|
|
)
|
|
.await
|
|
{
|
|
std::result::Result::Ok(session) => session,
|
|
std::result::Result::Err(error) => panic!("session failed: {error}"),
|
|
};
|
|
let mut events = session.subscribe_events();
|
|
let acknowledgement = match session
|
|
.subscribe(crate::StandardWsRequest::Signature(crate::SignatureSubscribeRequest {
|
|
signature: "1111111111111111111111111111111111111111111111111111111111111111"
|
|
.to_string(),
|
|
config: std::option::Option::None,
|
|
}))
|
|
.await
|
|
{
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("subscribe failed: {error}"),
|
|
};
|
|
assert_eq!(acknowledgement.subscription.remote_subscription_id, Some(77));
|
|
let removed = tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
|
loop {
|
|
let event = match events.recv().await {
|
|
std::result::Result::Ok(event) => event,
|
|
std::result::Result::Err(error) => panic!("event failed: {error}"),
|
|
};
|
|
if let crate::WsSessionEvent::SubscriptionRemoved(subscription) = event {
|
|
return subscription;
|
|
}
|
|
}
|
|
})
|
|
.await;
|
|
let removed = match removed {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(_) => panic!("terminal removal timed out"),
|
|
};
|
|
assert_eq!(removed.remote_subscription_id, Some(77));
|
|
assert!(session.snapshot().await.subscriptions.is_empty());
|
|
assert!(session.disconnect().await.is_ok());
|
|
if let std::result::Result::Err(error) = server.await {
|
|
panic!("server task failed: {error}");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn reconnect_restores_subscription_with_new_remote_id() {
|
|
let listener = match tokio::net::TcpListener::bind("127.0.0.1:0").await {
|
|
std::result::Result::Ok(listener) => listener,
|
|
std::result::Result::Err(error) => panic!("listener failed: {error}"),
|
|
};
|
|
let address = match listener.local_addr() {
|
|
std::result::Result::Ok(address) => address,
|
|
std::result::Result::Err(error) => panic!("local address failed: {error}"),
|
|
};
|
|
let server = tokio::spawn(async move {
|
|
let (first_socket, _peer) = match listener.accept().await {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("first accept failed: {error}"),
|
|
};
|
|
let mut first = match tokio_tungstenite::accept_async(first_socket).await {
|
|
std::result::Result::Ok(stream) => stream,
|
|
std::result::Result::Err(error) => panic!("first handshake failed: {error}"),
|
|
};
|
|
let subscribe = next_text(&mut first).await;
|
|
send_json(
|
|
&mut first,
|
|
serde_json::json!({ "jsonrpc": "2.0", "result": 10, "id": subscribe["id"].clone() }),
|
|
)
|
|
.await;
|
|
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
|
if let std::result::Result::Err(error) =
|
|
first.send(tokio_tungstenite::tungstenite::Message::Close(None)).await
|
|
{
|
|
panic!("first close failed: {error}");
|
|
}
|
|
let (second_socket, _peer) = match listener.accept().await {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("second accept failed: {error}"),
|
|
};
|
|
let mut second = match tokio_tungstenite::accept_async(second_socket).await {
|
|
std::result::Result::Ok(stream) => stream,
|
|
std::result::Result::Err(error) => panic!("second handshake failed: {error}"),
|
|
};
|
|
let resubscribe = next_text(&mut second).await;
|
|
assert_eq!(resubscribe["method"], "slotSubscribe");
|
|
send_json(
|
|
&mut second,
|
|
serde_json::json!({ "jsonrpc": "2.0", "result": 20, "id": resubscribe["id"].clone() }),
|
|
)
|
|
.await;
|
|
let disconnect_unsubscribe = next_text(&mut second).await;
|
|
assert_eq!(disconnect_unsubscribe["method"], "slotUnsubscribe");
|
|
assert_eq!(disconnect_unsubscribe["params"][0], 20);
|
|
let close = match second.next().await {
|
|
std::option::Option::Some(std::result::Result::Ok(message)) => message,
|
|
other => panic!("second close read failed: {other:?}"),
|
|
};
|
|
assert!(matches!(close, tokio_tungstenite::tungstenite::Message::Close(_)));
|
|
});
|
|
let client = match crate::WsClient::new(endpoint(format!("ws://{address}"))) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client failed: {error}"),
|
|
};
|
|
let policy = match crate::WsReconnectPolicy::bounded(2, 10, 20) {
|
|
std::result::Result::Ok(policy) => policy,
|
|
std::result::Result::Err(error) => panic!("policy failed: {error}"),
|
|
};
|
|
let session = match crate::WsSession::connect(
|
|
client,
|
|
crate::StandardWsCapabilities::default(),
|
|
policy,
|
|
)
|
|
.await
|
|
{
|
|
std::result::Result::Ok(session) => session,
|
|
std::result::Result::Err(error) => panic!("session failed: {error}"),
|
|
};
|
|
let mut events = session.subscribe_events();
|
|
let acknowledgement = match session
|
|
.subscribe(crate::StandardWsRequest::Slot(crate::SlotSubscribeRequest))
|
|
.await
|
|
{
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(error) => panic!("subscribe failed: {error}"),
|
|
};
|
|
assert_eq!(acknowledgement.subscription.remote_subscription_id, Some(10));
|
|
let remapped = tokio::time::timeout(std::time::Duration::from_secs(3), async {
|
|
loop {
|
|
let event = match events.recv().await {
|
|
std::result::Result::Ok(event) => event,
|
|
std::result::Result::Err(error) => panic!("event failed: {error}"),
|
|
};
|
|
if let crate::WsSessionEvent::SubscriptionRemapped {
|
|
previous_remote_subscription_id,
|
|
remote_subscription_id,
|
|
..
|
|
} = event
|
|
{
|
|
return (previous_remote_subscription_id, remote_subscription_id);
|
|
}
|
|
}
|
|
})
|
|
.await;
|
|
let remapped = match remapped {
|
|
std::result::Result::Ok(value) => value,
|
|
std::result::Result::Err(_) => panic!("remap timed out"),
|
|
};
|
|
assert_eq!(remapped, (Some(10), 20));
|
|
let snapshot = session.snapshot().await;
|
|
assert_eq!(snapshot.reconnect_count, 1);
|
|
assert_eq!(snapshot.subscriptions[0].remote_subscription_id, Some(20));
|
|
assert!(session.disconnect().await.is_ok());
|
|
if let std::result::Result::Err(error) = server.await {
|
|
panic!("server task failed: {error}");
|
|
}
|
|
}
|
|
}
|