Sasedev/khadhroony-solana-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
888ce3f35f12491c6258ea4e7bae6711be2f4cb8
khadhroony-solana-project/crates/ksp-onchain-transport-lib/src/ws_session.rs
SinuS Von SifriduS 6721a4142a v0.3.14-pre.003
2026-09-11 21:54:19 +02:00

2471 lines
117 KiB
Rust
Raw Blame History

// file: crates/ksp-onchain-transport-lib/src/ws_session.rs
// version: 15
use futures_util::SinkExt; // rust-rules: trait-import
use futures_util::StreamExt; // rust-rules: trait-import
static NEXT_WS_SESSION_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
const HELIUS_WS_HEARTBEAT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
type WsPhysicalStream = tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
/// Cloneable latest-value observer for one physical WebSocket session snapshot.
///
/// The observer exposes only the already-safe [`crate::WsSessionSnapshot`] projection and keeps the internal Tokio watch channel private. Cloning it does
/// not create another socket, actor, reconnect loop or subscription registry.
#[derive(Clone)]
pub struct WsSessionSnapshotSource {
receiver: tokio::sync::watch::Receiver<crate::WsSessionSnapshot>,
}
impl crate::WsSessionSnapshotSource {
fn new(receiver: tokio::sync::watch::Receiver<crate::WsSessionSnapshot>) -> Self {
return Self { receiver };
}
/// Returns the current safe physical-session snapshot without waiting for another actor transition.
#[must_use]
pub fn current(&self) -> crate::WsSessionSnapshot {
return (*self.receiver.borrow()).clone();
}
/// Waits for one newer safe physical-session snapshot.
///
/// `None` means the owning WebSocket actor dropped the latest-value publisher and no further snapshot can arrive.
pub async fn wait_for_change(&mut self) -> std::option::Option<crate::WsSessionSnapshot> {
return match self.receiver.changed().await {
std::result::Result::Ok(()) => std::option::Option::Some((*self.receiver.borrow_and_update()).clone()),
std::result::Result::Err(_) => std::option::Option::None,
};
}
}
impl std::fmt::Debug for crate::WsSessionSnapshotSource {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
return formatter.debug_struct("WsSessionSnapshotSource").field("current", &self.current()).finish();
}
}
/// Shareable handle for one explicitly created physical WebSocket session.
///
/// The handle never exposes the sensitive endpoint URL or the underlying socket. All socket I/O is owned by one internal actor task and all caller
/// interaction is serialized through bounded channels.
#[derive(Clone)]
pub struct WsSession {
id: crate::WsSessionId,
command_tx: tokio::sync::mpsc::Sender<WsSessionCommand>,
shutdown_tx: tokio::sync::watch::Sender<std::option::Option<tokio::time::Instant>>,
snapshot_rx: tokio::sync::watch::Receiver<crate::WsSessionSnapshot>,
command_timeout: std::time::Duration,
close_timeout: std::time::Duration,
notification_queue_capacity: usize,
}
impl WsSession {
/// Opens one physical standard Solana WebSocket connection for the supplied endpoint settings.
///
/// This historical constructor remains standard-only after provider-specific protocol kinds are added. Calling this function twice with the same endpoint
/// creates two independent physical sessions. Provider-specific callers must use their typed protocol facade instead of obtaining a generic handle.
pub async fn connect(endpoint: crate::WsEndpointSettings) -> ksp_core_lib::Result<Self> {
return Self::connect_for_protocol(endpoint, crate::WsProtocolKind::SolanaStandard).await;
}
/// Opens one physical WebSocket connection after validating the typed facade protocol.
pub(crate) async fn connect_for_protocol(endpoint: crate::WsEndpointSettings, expected_protocol: crate::WsProtocolKind) -> ksp_core_lib::Result<Self> {
let validation = crate::WsTransportSettings::new(std::vec![endpoint.clone()]).validate();
if let std::result::Result::Err(error) = validation {
return std::result::Result::Err(error);
}
if endpoint.protocol() != expected_protocol {
return std::result::Result::Err(
ksp_core_lib::Error::new(crate::ERROR_CODE_INVALID_SETTINGS, "WebSocket session constructor does not accept the endpoint protocol")
.with_context("field", "ws_endpoints.protocol")
.with_context("expected_protocol", expected_protocol.as_str())
.with_context("actual_protocol", endpoint.protocol().as_str()),
);
}
return Self::connect_physical(endpoint).await;
}
async fn connect_physical(endpoint: crate::WsEndpointSettings) -> ksp_core_lib::Result<Self> {
let id_result = next_session_id();
let id = match id_result {
std::result::Result::Ok(id) => id,
std::result::Result::Err(error) => return std::result::Result::Err(error),
};
let initial_snapshot = crate::WsSessionSnapshot::new(
id,
endpoint.name(),
endpoint.provider().clone(),
endpoint.cluster().clone(),
endpoint.protocol(),
crate::WsSessionState::Connecting,
0,
0,
0,
std::vec::Vec::new(),
);
let (command_tx, command_rx) = tokio::sync::mpsc::channel(endpoint.session().command_queue_capacity());
let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(std::option::Option::None::<tokio::time::Instant>);
let (snapshot_tx, snapshot_rx) = tokio::sync::watch::channel(initial_snapshot);
let (startup_tx, startup_rx) = tokio::sync::oneshot::channel();
let command_timeout = endpoint.session().command_timeout();
let close_timeout = endpoint.session().close_timeout();
let notification_queue_capacity = endpoint.session().notification_queue_capacity();
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
provider = endpoint.provider().as_str(),
cluster = endpoint.cluster().as_str(),
protocol = endpoint.protocol().as_str(),
"starting physical WebSocket session actor"
);
let join_handle = tokio::spawn(run_ws_session_actor(id, endpoint, command_rx, shutdown_rx, snapshot_tx, startup_tx));
let startup_wait = tokio::time::timeout(command_timeout, startup_rx).await;
return match startup_wait {
std::result::Result::Ok(std::result::Result::Ok(std::result::Result::Ok(()))) => std::result::Result::Ok(Self {
id,
command_tx,
shutdown_tx,
snapshot_rx,
command_timeout,
close_timeout,
notification_queue_capacity,
}),
std::result::Result::Ok(std::result::Result::Ok(std::result::Result::Err(error))) => {
join_handle.abort();
std::result::Result::Err(error)
},
std::result::Result::Ok(std::result::Result::Err(_)) => {
join_handle.abort();
std::result::Result::Err(ws_session_closed_error(id, "WebSocket session actor ended during startup"))
},
std::result::Result::Err(_) => {
join_handle.abort();
std::result::Result::Err(ws_timeout_error(id, "WebSocket handshake exceeded the configured command timeout"))
},
};
}
/// Returns the stable local session identity.
#[must_use]
pub const fn id(&self) -> crate::WsSessionId {
return self.id;
}
/// Returns the latest safe runtime snapshot published by the actor.
#[must_use]
pub fn snapshot(&self) -> crate::WsSessionSnapshot {
return self.snapshot_rx.borrow().clone();
}
/// Returns a cloneable latest-value observer for safe physical-session snapshots.
#[must_use]
pub fn snapshot_source(&self) -> crate::WsSessionSnapshotSource {
return crate::WsSessionSnapshotSource::new(self.snapshot_rx.clone());
}
/// Returns the latest observable physical-session state.
#[must_use]
pub fn state(&self) -> crate::WsSessionState {
return self.snapshot_rx.borrow().state();
}
/// Explicitly closes this physical session under the configured close timeout.
///
/// Shutdown is session-wide: calling this method through any clone requests the actor to enter `Closing`, cancels pending JSON-RPC requests, sends a
/// best-effort WebSocket Close frame and waits for the safe lifecycle snapshot to reach `Closed`. The shutdown signal is independent from the bounded
/// command queue so a saturated command queue cannot prevent close from being requested.
pub async fn close(&self) -> ksp_core_lib::Result<()> {
match self.state() {
crate::WsSessionState::Closed => return std::result::Result::Ok(()),
crate::WsSessionState::Failed => {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session is already in terminal failed state"));
},
_ => {},
}
let deadline = tokio::time::Instant::now() + self.close_timeout;
self.shutdown_tx.send_replace(std::option::Option::Some(deadline));
let mut snapshot_rx = self.snapshot_rx.clone();
loop {
match snapshot_rx.borrow().state() {
crate::WsSessionState::Closed => return std::result::Result::Ok(()),
crate::WsSessionState::Failed => {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session failed while explicit shutdown was in progress"));
},
_ => {},
}
let changed = tokio::time::timeout_at(deadline, snapshot_rx.changed()).await;
match changed {
std::result::Result::Ok(std::result::Result::Ok(())) => {},
std::result::Result::Ok(std::result::Result::Err(_)) => {
return match snapshot_rx.borrow().state() {
crate::WsSessionState::Closed => std::result::Result::Ok(()),
_ => std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session actor ended before publishing Closed state")),
};
},
std::result::Result::Err(_) => {
if snapshot_rx.borrow().state() == crate::WsSessionState::Closed {
return std::result::Result::Ok(());
}
return std::result::Result::Err(ws_timeout_error(self.id, "WebSocket explicit shutdown exceeded the configured close timeout"));
},
}
}
}
/// Creates one crate-internal typed standard Solana subscription through the actor-owned registry.
///
/// Public typed standard wrappers consume this constructor while it remains crate-private, preventing a raw provider-extension subscription API from
/// becoming part of the stable KSP surface.
pub(crate) async fn subscribe_typed<T, F>(
&self,
kind: crate::WsSubscriptionKind,
params: std::vec::Vec<serde_json::Value>,
decoder: F,
) -> ksp_core_lib::Result<crate::WsSubscription<T>>
where
T: std::marker::Send + 'static,
F: Fn(serde_json::Value) -> ksp_core_lib::Result<T> + std::marker::Send + std::marker::Sync + 'static,
{
return self.subscribe_typed_with_completion(kind, params, decoder, |_| return false).await;
}
/// Creates one crate-internal typed subscription whose decoder can identify a delivered terminal notification.
pub(crate) async fn subscribe_typed_with_completion<T, F, C>(
&self,
kind: crate::WsSubscriptionKind,
params: std::vec::Vec<serde_json::Value>,
decoder: F,
is_terminal: C,
) -> ksp_core_lib::Result<crate::WsSubscription<T>>
where
T: std::marker::Send + 'static,
F: Fn(serde_json::Value) -> ksp_core_lib::Result<T> + std::marker::Send + std::marker::Sync + 'static,
C: Fn(&T) -> bool + std::marker::Send + std::marker::Sync + 'static,
{
if self.state() != crate::WsSessionState::Active {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session is not active"));
}
kind.warn_if_unstable();
let (dispatcher, notification_rx) = crate::typed_notification_channel_with_completion(self.notification_queue_capacity, decoder, is_terminal);
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
let command = WsSessionCommand::Subscribe { kind, params, dispatcher, response_tx };
let send_wait = tokio::time::timeout(self.command_timeout, self.command_tx.send(command)).await;
match send_wait {
std::result::Result::Ok(std::result::Result::Ok(())) => {},
std::result::Result::Ok(std::result::Result::Err(_)) => {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session command channel is closed"));
},
std::result::Result::Err(_) => {
return std::result::Result::Err(ws_timeout_error(self.id, "WebSocket subscribe command queue remained unavailable until timeout"));
},
}
let registration = match response_rx.await {
std::result::Result::Ok(std::result::Result::Ok(registration)) => registration,
std::result::Result::Ok(std::result::Result::Err(error)) => return std::result::Result::Err(error),
std::result::Result::Err(_) => {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session ended before subscribe completion"));
},
};
return std::result::Result::Ok(crate::WsSubscription::new(self.id, registration, notification_rx, self.command_tx.clone(), self.command_timeout));
}
/// Executes one internal JSON-RPC request through the actor-owned socket.
///
/// This remains crate-private in `0.2.7`; standard subscriptions consume it without exposing a public raw provider-extension escape hatch.
#[allow(dead_code)] // Staged in pre.004 and consumed by the subscription engine starting in pre.006.
pub(crate) async fn execute_json_rpc(&self, method: &'static str, params: std::vec::Vec<serde_json::Value>) -> ksp_core_lib::Result<serde_json::Value> {
if self.state() != crate::WsSessionState::Active {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session is not active"));
}
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
let command = WsSessionCommand::ExecuteJsonRpc { method, params, response_tx };
let send_wait = tokio::time::timeout(self.command_timeout, self.command_tx.send(command)).await;
match send_wait {
std::result::Result::Ok(std::result::Result::Ok(())) => {},
std::result::Result::Ok(std::result::Result::Err(_)) => {
return std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session command channel is closed"));
},
std::result::Result::Err(_) => {
return std::result::Result::Err(ws_timeout_error(self.id, "WebSocket session command queue remained unavailable until timeout"));
},
}
return match response_rx.await {
std::result::Result::Ok(result) => result,
std::result::Result::Err(_) => {
std::result::Result::Err(ws_session_closed_error(self.id, "WebSocket session ended before the JSON-RPC response was delivered"))
},
};
}
}
impl std::fmt::Debug for WsSession {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
return formatter.debug_struct("WsSession").field("id", &self.id).field("snapshot", &self.snapshot()).finish();
}
}
/// Internal bounded commands accepted by the actor-owned WebSocket session runtime.
pub(crate) enum WsSessionCommand {
ExecuteJsonRpc {
method: &'static str,
params: std::vec::Vec<serde_json::Value>,
response_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<serde_json::Value>>,
},
Subscribe {
kind: crate::WsSubscriptionKind,
params: std::vec::Vec<serde_json::Value>,
dispatcher: crate::WsNotificationDispatcher,
response_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<crate::WsSubscriptionRegistration>>,
},
Unsubscribe {
subscription_id: crate::WsSubscriptionId,
response_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<bool>>,
},
}
enum PendingWsResponse {
Raw(tokio::sync::oneshot::Sender<ksp_core_lib::Result<serde_json::Value>>),
Subscribe {
subscription_id: crate::WsSubscriptionId,
response_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<crate::WsSubscriptionRegistration>>,
},
Resubscribe {
subscription_id: crate::WsSubscriptionId,
kind: crate::WsSubscriptionKind,
},
Unsubscribe {
subscription_id: crate::WsSubscriptionId,
response_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<bool>>,
},
}
struct PendingWsRequest {
method: &'static str,
deadline: tokio::time::Instant,
response: PendingWsResponse,
}
enum WsActorIoOutcome {
Continue,
RemoteClosed,
ShutdownRequested { deadline: tokio::time::Instant },
BestEffortUnsubscribe { kind: crate::WsSubscriptionKind, remote_id: u64 },
Failed { code: ksp_core_lib::ErrorCode, pending_message: &'static str },
}
#[derive(Clone, Copy, Default)]
struct WsRuntimeCounters {
continuity_gap_count: u64,
overflow_count: u64,
}
enum WsReconnectOutcome {
Connected { websocket: std::boxed::Box<WsPhysicalStream> },
ShutdownRequested { deadline: tokio::time::Instant },
HandlesDropped,
Exhausted,
}
enum WsReconnectControlOutcome {
Continue,
ShutdownRequested { deadline: tokio::time::Instant },
HandlesDropped,
}
enum WsConnectAttemptOutcome {
Connected { websocket: std::boxed::Box<WsPhysicalStream>, handshake_status: u16 },
Retry,
ShutdownRequested { deadline: tokio::time::Instant },
HandlesDropped,
}
async fn run_ws_session_actor(
id: crate::WsSessionId,
endpoint: crate::WsEndpointSettings,
mut command_rx: tokio::sync::mpsc::Receiver<WsSessionCommand>,
mut shutdown_rx: tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
snapshot_tx: tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
startup_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<()>>,
) {
ksp_logging_lib::trace!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
provider = endpoint.provider().as_str(),
cluster = endpoint.cluster().as_str(),
"opening physical WebSocket connection"
);
let connect_wait = tokio::time::timeout(
endpoint.session().command_timeout(),
tokio_tungstenite::connect_async_with_config(endpoint.url().as_str(), std::option::Option::Some(websocket_config(&endpoint)), false),
)
.await;
let (mut websocket, handshake_status) = match connect_wait {
std::result::Result::Ok(std::result::Result::Ok((websocket, response))) => (websocket, response.status().as_u16()),
std::result::Result::Ok(std::result::Result::Err(_)) => {
publish_snapshot(&snapshot_tx, id, &endpoint, crate::WsSessionState::Failed, 0, WsRuntimeCounters::default(), &std::collections::BTreeMap::new());
let error = ws_connection_error(id, &endpoint, "WebSocket connection or handshake failed");
let _ = startup_tx.send(std::result::Result::Err(error));
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
provider = endpoint.provider().as_str(),
cluster = endpoint.cluster().as_str(),
"physical WebSocket connection failed"
);
return;
},
std::result::Result::Err(_) => {
publish_snapshot(&snapshot_tx, id, &endpoint, crate::WsSessionState::Failed, 0, WsRuntimeCounters::default(), &std::collections::BTreeMap::new());
let error = ws_timeout_error(id, "WebSocket connection handshake timed out");
let _ = startup_tx.send(std::result::Result::Err(error));
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"physical WebSocket connection handshake timed out"
);
return;
},
};
let mut continuity_gap_count = 0_u64;
let mut overflow_count = 0_u64;
publish_snapshot(
&snapshot_tx,
id,
&endpoint,
crate::WsSessionState::Active,
0,
WsRuntimeCounters { continuity_gap_count, overflow_count },
&std::collections::BTreeMap::new(),
);
let _ = startup_tx.send(std::result::Result::Ok(()));
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
handshake_status,
"physical WebSocket session is active"
);
let mut next_request_id = 1_u64;
let mut next_subscription_id = 1_u64;
let mut pending = std::collections::BTreeMap::<u64, PendingWsRequest>::new();
let mut subscriptions = std::collections::BTreeMap::<u64, crate::WsSubscriptionRuntime>::new();
let mut remote_to_local = std::collections::BTreeMap::<u64, crate::WsSubscriptionId>::new();
let heartbeat_enabled = helius_heartbeat_enabled(endpoint.protocol());
let mut heartbeat_deadline = next_helius_heartbeat_deadline();
loop {
prune_cancelled_pending(id, &mut pending, &mut subscriptions, &mut remote_to_local);
let timeout_deadline = next_pending_deadline(&pending);
let outcome = tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(&shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
WsActorIoOutcome::ShutdownRequested { deadline }
},
maybe_command = command_rx.recv() => {
let command = match maybe_command {
std::option::Option::Some(command) => command,
std::option::Option::None => {
let deadline = tokio::time::Instant::now() + endpoint.session().close_timeout();
ksp_logging_lib::trace!(target: crate::TRACING_TARGET, session_id = id.get(), "all WebSocket session handles dropped; closing actor");
close_session_actor(
id,
&endpoint,
&snapshot_tx,
&mut websocket,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
deadline,
)
.await;
return;
},
};
handle_session_command(
id,
&endpoint,
&mut websocket,
&mut pending,
&mut next_request_id,
&mut next_subscription_id,
&mut subscriptions,
&mut remote_to_local,
&mut shutdown_rx,
command,
)
.await
},
maybe_message = websocket.next() => {
handle_socket_message(
id,
&endpoint,
maybe_message,
&mut websocket,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
&mut overflow_count,
&mut shutdown_rx,
)
.await
},
() = tokio::time::sleep_until(heartbeat_deadline), if heartbeat_enabled => {
let heartbeat = send_helius_heartbeat(id, &endpoint, &mut websocket, &mut shutdown_rx).await;
if matches!(&heartbeat, WsActorIoOutcome::Continue) {
heartbeat_deadline = next_helius_heartbeat_deadline();
}
heartbeat
},
() = tokio::time::sleep_until(timeout_deadline) => {
expire_pending_requests(id, &mut pending, &mut subscriptions, &mut remote_to_local);
WsActorIoOutcome::Continue
},
};
match outcome {
WsActorIoOutcome::Continue => {
publish_snapshot(
&snapshot_tx,
id,
&endpoint,
crate::WsSessionState::Active,
pending.len(),
WsRuntimeCounters { continuity_gap_count, overflow_count },
&subscriptions,
);
},
WsActorIoOutcome::BestEffortUnsubscribe { kind, remote_id } => {
let cleanup = best_effort_remote_unsubscribe(id, &endpoint, &mut websocket, &mut shutdown_rx, kind, remote_id).await;
if let WsActorIoOutcome::ShutdownRequested { deadline } = cleanup {
close_session_actor(
id,
&endpoint,
&snapshot_tx,
&mut websocket,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
deadline,
)
.await;
return;
}
publish_snapshot(
&snapshot_tx,
id,
&endpoint,
crate::WsSessionState::Active,
pending.len(),
WsRuntimeCounters { continuity_gap_count, overflow_count },
&subscriptions,
);
},
WsActorIoOutcome::ShutdownRequested { deadline } => {
close_session_actor(
id,
&endpoint,
&snapshot_tx,
&mut websocket,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
deadline,
)
.await;
return;
},
WsActorIoOutcome::RemoteClosed => {
let recovery = recover_websocket_session(
id,
&endpoint,
&mut command_rx,
&mut shutdown_rx,
&snapshot_tx,
&mut pending,
&mut next_request_id,
&mut subscriptions,
&mut remote_to_local,
&mut continuity_gap_count,
&mut overflow_count,
crate::ERROR_CODE_WS_CONNECTION_FAILED,
"Remote peer closed WebSocket session before pending response delivery",
)
.await;
match recovery {
WsReconnectOutcome::Connected { websocket: replacement } => {
websocket = *replacement;
if heartbeat_enabled {
heartbeat_deadline = next_helius_heartbeat_deadline();
}
},
WsReconnectOutcome::ShutdownRequested { deadline } => {
finish_disconnected_shutdown(
id,
&endpoint,
&snapshot_tx,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
deadline,
);
return;
},
WsReconnectOutcome::HandlesDropped => {
finish_disconnected_close(
id,
&endpoint,
&snapshot_tx,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
);
return;
},
WsReconnectOutcome::Exhausted => {
finish_reconnect_exhaustion(
id,
&endpoint,
&snapshot_tx,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
);
return;
},
}
},
WsActorIoOutcome::Failed { code, pending_message } => {
let recovery = recover_websocket_session(
id,
&endpoint,
&mut command_rx,
&mut shutdown_rx,
&snapshot_tx,
&mut pending,
&mut next_request_id,
&mut subscriptions,
&mut remote_to_local,
&mut continuity_gap_count,
&mut overflow_count,
code,
pending_message,
)
.await;
match recovery {
WsReconnectOutcome::Connected { websocket: replacement } => {
websocket = *replacement;
if heartbeat_enabled {
heartbeat_deadline = next_helius_heartbeat_deadline();
}
},
WsReconnectOutcome::ShutdownRequested { deadline } => {
finish_disconnected_shutdown(
id,
&endpoint,
&snapshot_tx,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
deadline,
);
return;
},
WsReconnectOutcome::HandlesDropped => {
finish_disconnected_close(
id,
&endpoint,
&snapshot_tx,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
);
return;
},
WsReconnectOutcome::Exhausted => {
finish_reconnect_exhaustion(
id,
&endpoint,
&snapshot_tx,
&mut pending,
&mut subscriptions,
&mut remote_to_local,
continuity_gap_count,
overflow_count,
);
return;
},
}
},
}
}
}
const fn helius_heartbeat_enabled(protocol: crate::WsProtocolKind) -> bool {
return matches!(protocol, crate::WsProtocolKind::HeliusLaserStream);
}
fn next_helius_heartbeat_deadline() -> tokio::time::Instant {
return tokio::time::Instant::now() + HELIUS_WS_HEARTBEAT_INTERVAL;
}
async fn send_helius_heartbeat<S>(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
websocket: &mut tokio_tungstenite::WebSocketStream<S>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
) -> WsActorIoOutcome
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + std::marker::Unpin,
{
let message = tokio_tungstenite::tungstenite::Message::Ping(std::vec::Vec::new().into());
let send_result = tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
return WsActorIoOutcome::ShutdownRequested { deadline };
},
send_result = websocket.send(message) => send_result,
() = tokio::time::sleep(endpoint.session().command_timeout()) => {
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"Helius WebSocket heartbeat Ping write timed out"
);
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection failed while writing Helius heartbeat Ping",
};
},
};
return match send_result {
std::result::Result::Ok(()) => {
ksp_logging_lib::trace!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"sent Helius WebSocket heartbeat Ping control frame"
);
WsActorIoOutcome::Continue
},
std::result::Result::Err(_) => {
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"Helius WebSocket heartbeat Ping write failed"
);
WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection failed while writing Helius heartbeat Ping",
}
},
};
}
fn websocket_config(endpoint: &crate::WsEndpointSettings) -> tokio_tungstenite::tungstenite::protocol::WebSocketConfig {
return tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default()
.write_buffer_size(0)
.max_write_buffer_size(endpoint.session().max_write_buffer_size_bytes())
.max_message_size(std::option::Option::Some(endpoint.session().max_message_size_bytes()))
.max_frame_size(std::option::Option::Some(endpoint.session().max_frame_size_bytes()));
}
#[allow(clippy::too_many_arguments)]
async fn recover_websocket_session(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
command_rx: &mut tokio::sync::mpsc::Receiver<WsSessionCommand>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
next_request_id: &mut u64,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
continuity_gap_count: &mut u64,
overflow_count: &mut u64,
code: ksp_core_lib::ErrorCode,
pending_message: &'static str,
) -> WsReconnectOutcome {
fail_pending_for_reconnect(pending, id, code, pending_message, subscriptions, remote_to_local);
prepare_subscriptions_for_reconnect(endpoint.session().resubscribe(), subscriptions, remote_to_local, code);
*continuity_gap_count = (*continuity_gap_count).saturating_add(1);
let max_retries = endpoint.session().reconnect().max_retries();
if max_retries == 0 {
return WsReconnectOutcome::Exhausted;
}
let mut attempt = 1_u32;
while attempt <= max_retries {
publish_snapshot(
snapshot_tx,
id,
endpoint,
crate::WsSessionState::Reconnecting { attempt },
pending.len(),
WsRuntimeCounters { continuity_gap_count: *continuity_gap_count, overflow_count: *overflow_count },
subscriptions,
);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
reconnect_attempt = attempt,
reconnect_max_retries = max_retries,
continuity_gap_count = *continuity_gap_count,
"physical WebSocket session is reconnecting"
);
let backoff = reconnect_backoff(endpoint.session().reconnect(), attempt);
let backoff_outcome = wait_reconnect_backoff(id, endpoint, command_rx, shutdown_rx, subscriptions, remote_to_local, backoff).await;
match backoff_outcome {
WsReconnectControlOutcome::Continue => {},
WsReconnectControlOutcome::ShutdownRequested { deadline } => return WsReconnectOutcome::ShutdownRequested { deadline },
WsReconnectControlOutcome::HandlesDropped => return WsReconnectOutcome::HandlesDropped,
}
let connection = connect_replacement_websocket(id, endpoint, command_rx, shutdown_rx, subscriptions, remote_to_local).await;
let mut websocket = match connection {
WsConnectAttemptOutcome::Connected { websocket, handshake_status } => {
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
reconnect_attempt = attempt,
handshake_status,
"replacement physical WebSocket connection established"
);
*websocket
},
WsConnectAttemptOutcome::Retry => {
attempt = attempt.saturating_add(1);
continue;
},
WsConnectAttemptOutcome::ShutdownRequested { deadline } => return WsReconnectOutcome::ShutdownRequested { deadline },
WsConnectAttemptOutcome::HandlesDropped => return WsReconnectOutcome::HandlesDropped,
};
let restore = restore_subscriptions_after_reconnect(
id,
endpoint,
command_rx,
shutdown_rx,
snapshot_tx,
&mut websocket,
pending,
next_request_id,
subscriptions,
remote_to_local,
*continuity_gap_count,
overflow_count,
attempt,
)
.await;
match restore {
WsActorIoOutcome::Continue => {
publish_snapshot(
snapshot_tx,
id,
endpoint,
crate::WsSessionState::Active,
pending.len(),
WsRuntimeCounters { continuity_gap_count: *continuity_gap_count, overflow_count: *overflow_count },
subscriptions,
);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
reconnect_attempt = attempt,
continuity_gap_count = *continuity_gap_count,
"physical WebSocket reconnect completed and retry budget reset"
);
return WsReconnectOutcome::Connected { websocket: std::boxed::Box::new(websocket) };
},
WsActorIoOutcome::ShutdownRequested { deadline } => return WsReconnectOutcome::ShutdownRequested { deadline },
WsActorIoOutcome::RemoteClosed => {
fail_pending_for_reconnect(
pending,
id,
crate::ERROR_CODE_WS_CONNECTION_FAILED,
"Replacement WebSocket closed during subscription restoration",
subscriptions,
remote_to_local,
);
prepare_subscriptions_for_reconnect(endpoint.session().resubscribe(), subscriptions, remote_to_local, crate::ERROR_CODE_WS_CONNECTION_FAILED);
},
WsActorIoOutcome::Failed { code: restore_code, pending_message: restore_message } => {
fail_pending_for_reconnect(pending, id, restore_code, restore_message, subscriptions, remote_to_local);
prepare_subscriptions_for_reconnect(endpoint.session().resubscribe(), subscriptions, remote_to_local, restore_code);
},
WsActorIoOutcome::BestEffortUnsubscribe { .. } => {},
}
attempt = attempt.saturating_add(1);
}
return WsReconnectOutcome::Exhausted;
}
#[allow(clippy::too_many_arguments)]
async fn restore_subscriptions_after_reconnect(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
command_rx: &mut tokio::sync::mpsc::Receiver<WsSessionCommand>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
websocket: &mut WsPhysicalStream,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
next_request_id: &mut u64,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
continuity_gap_count: u64,
overflow_count: &mut u64,
attempt: u32,
) -> WsActorIoOutcome {
let restore_ids = subscriptions
.values()
.filter(|runtime| return runtime.state == crate::WsSubscriptionState::Resubscribing)
.map(|runtime| return runtime.id)
.collect::<std::vec::Vec<_>>();
for subscription_id in restore_ids {
let (kind, params) = match subscriptions.get(&subscription_id.get()) {
std::option::Option::Some(runtime) if runtime.state == crate::WsSubscriptionState::Resubscribing => (runtime.kind, runtime.params.clone()),
_ => continue,
};
let method = kind.subscribe_method();
let write_result = write_json_rpc_request(id, endpoint, websocket, pending.len(), next_request_id, shutdown_rx, method, params).await;
let (request_id, deadline) = match write_result {
std::result::Result::Ok(result) => result,
std::result::Result::Err((error, WsActorIoOutcome::Continue)) => {
let error_code = error.code();
fail_local_subscription(subscriptions, remote_to_local, subscription_id, error_code);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
error_code = error_code.code(),
"logical WebSocket resubscribe could not be queued"
);
continue;
},
std::result::Result::Err((_, outcome)) => return outcome,
};
pending.insert(request_id, PendingWsRequest { method, deadline, response: PendingWsResponse::Resubscribe { subscription_id, kind } });
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
reconnect_attempt = attempt,
"requested deterministic WebSocket resubscribe"
);
while pending.contains_key(&request_id) {
let timeout_deadline = next_pending_deadline(pending);
let outcome = tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
WsActorIoOutcome::ShutdownRequested { deadline }
},
maybe_command = command_rx.recv() => {
match maybe_command {
std::option::Option::Some(command) => {
handle_reconnecting_command_with_socket(
id,
endpoint,
websocket,
subscriptions,
remote_to_local,
shutdown_rx,
command,
)
.await
},
std::option::Option::None => {
return WsActorIoOutcome::ShutdownRequested {
deadline: tokio::time::Instant::now() + endpoint.session().close_timeout(),
};
},
}
},
maybe_message = websocket.next() => {
handle_socket_message(id, endpoint, maybe_message, websocket, pending, subscriptions, remote_to_local, overflow_count, shutdown_rx).await
},
() = tokio::time::sleep_until(timeout_deadline) => {
expire_pending_requests(id, pending, subscriptions, remote_to_local);
WsActorIoOutcome::Continue
},
};
match outcome {
WsActorIoOutcome::Continue => {},
WsActorIoOutcome::BestEffortUnsubscribe { kind: stale_kind, remote_id } => {
let cleanup = best_effort_remote_unsubscribe(id, endpoint, websocket, shutdown_rx, stale_kind, remote_id).await;
if let WsActorIoOutcome::ShutdownRequested { deadline } = cleanup {
return WsActorIoOutcome::ShutdownRequested { deadline };
}
},
_ => return outcome,
}
publish_snapshot(
snapshot_tx,
id,
endpoint,
crate::WsSessionState::Reconnecting { attempt },
pending.len(),
WsRuntimeCounters { continuity_gap_count, overflow_count: *overflow_count },
subscriptions,
);
}
}
return WsActorIoOutcome::Continue;
}
async fn handle_reconnecting_command_with_socket(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
websocket: &mut WsPhysicalStream,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
command: WsSessionCommand,
) -> WsActorIoOutcome {
return match command {
WsSessionCommand::Unsubscribe { subscription_id, response_tx } => {
let cleanup = subscriptions.get(&subscription_id.get()).and_then(|runtime| {
return runtime.remote_id.map(|remote_id| return (runtime.kind, remote_id));
});
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
if let std::option::Option::Some((kind, remote_id)) = cleanup {
let cleanup_outcome = best_effort_remote_unsubscribe(id, endpoint, websocket, shutdown_rx, kind, remote_id).await;
if let WsActorIoOutcome::ShutdownRequested { deadline } = cleanup_outcome {
let _ = response_tx.send(std::result::Result::Ok(false));
return WsActorIoOutcome::ShutdownRequested { deadline };
}
}
let _ = response_tx.send(std::result::Result::Ok(false));
WsActorIoOutcome::Continue
},
WsSessionCommand::ExecuteJsonRpc { response_tx, .. } => {
let _ = response_tx.send(std::result::Result::Err(ws_connection_error(id, endpoint, "WebSocket session is reconnecting")));
WsActorIoOutcome::Continue
},
WsSessionCommand::Subscribe { response_tx, .. } => {
let _ = response_tx.send(std::result::Result::Err(ws_connection_error(id, endpoint, "WebSocket session is reconnecting")));
WsActorIoOutcome::Continue
},
};
}
async fn wait_reconnect_backoff(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
command_rx: &mut tokio::sync::mpsc::Receiver<WsSessionCommand>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
backoff: std::time::Duration,
) -> WsReconnectControlOutcome {
let sleep = tokio::time::sleep(backoff);
tokio::pin!(sleep);
loop {
tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
return WsReconnectControlOutcome::ShutdownRequested { deadline };
},
maybe_command = command_rx.recv() => {
match maybe_command {
std::option::Option::Some(command) => {
handle_reconnecting_command_without_socket(id, endpoint, subscriptions, remote_to_local, command);
},
std::option::Option::None => return WsReconnectControlOutcome::HandlesDropped,
}
},
() = &mut sleep => return WsReconnectControlOutcome::Continue,
}
}
}
async fn connect_replacement_websocket(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
command_rx: &mut tokio::sync::mpsc::Receiver<WsSessionCommand>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
) -> WsConnectAttemptOutcome {
let connect = tokio_tungstenite::connect_async_with_config(endpoint.url().as_str(), std::option::Option::Some(websocket_config(endpoint)), false);
let timeout = tokio::time::sleep(endpoint.session().command_timeout());
tokio::pin!(connect);
tokio::pin!(timeout);
loop {
tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
return WsConnectAttemptOutcome::ShutdownRequested { deadline };
},
maybe_command = command_rx.recv() => {
match maybe_command {
std::option::Option::Some(command) => {
handle_reconnecting_command_without_socket(id, endpoint, subscriptions, remote_to_local, command);
},
std::option::Option::None => return WsConnectAttemptOutcome::HandlesDropped,
}
},
result = &mut connect => {
return match result {
std::result::Result::Ok((websocket, response)) => WsConnectAttemptOutcome::Connected {
websocket: std::boxed::Box::new(websocket),
handshake_status: response.status().as_u16(),
},
std::result::Result::Err(_) => {
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"replacement physical WebSocket connection attempt failed"
);
WsConnectAttemptOutcome::Retry
},
};
},
() = &mut timeout => {
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"replacement physical WebSocket connection attempt timed out"
);
return WsConnectAttemptOutcome::Retry;
},
}
}
}
fn handle_reconnecting_command_without_socket(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
command: WsSessionCommand,
) {
match command {
WsSessionCommand::Unsubscribe { subscription_id, response_tx } => {
handle_reconnecting_unsubscribe(id, subscriptions, remote_to_local, subscription_id, response_tx);
},
WsSessionCommand::ExecuteJsonRpc { response_tx, .. } => {
let _ = response_tx.send(std::result::Result::Err(ws_connection_error(id, endpoint, "WebSocket session is reconnecting")));
},
WsSessionCommand::Subscribe { response_tx, .. } => {
let _ = response_tx.send(std::result::Result::Err(ws_connection_error(id, endpoint, "WebSocket session is reconnecting")));
},
}
}
fn handle_reconnecting_unsubscribe(
id: crate::WsSessionId,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
subscription_id: crate::WsSubscriptionId,
response_tx: tokio::sync::oneshot::Sender<ksp_core_lib::Result<bool>>,
) {
let exists = subscriptions.contains_key(&subscription_id.get());
if exists {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
"local WebSocket unsubscribe won while the physical session was reconnecting"
);
}
let _ = response_tx.send(std::result::Result::Ok(false));
}
fn prepare_subscriptions_for_reconnect(
policy: crate::WsResubscribePolicy,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
failure_code: ksp_core_lib::ErrorCode,
) {
let terminal_failure_code = if policy == crate::WsResubscribePolicy::Never { crate::ERROR_CODE_WS_CONNECTION_FAILED } else { failure_code };
remote_to_local.clear();
let mut failed = std::vec::Vec::new();
let mut already_terminal = std::vec::Vec::new();
for runtime in subscriptions.values_mut() {
runtime.remote_id = std::option::Option::None;
match runtime.state {
crate::WsSubscriptionState::Active | crate::WsSubscriptionState::Resubscribing => {
if policy == crate::WsResubscribePolicy::ActiveSubscriptions {
runtime.set_state(crate::WsSubscriptionState::Resubscribing);
} else {
failed.push(runtime.id);
}
},
crate::WsSubscriptionState::Requested | crate::WsSubscriptionState::Cancelling => failed.push(runtime.id),
crate::WsSubscriptionState::Closed | crate::WsSubscriptionState::Failed => already_terminal.push(runtime.id),
}
}
for subscription_id in failed {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, terminal_failure_code);
}
for subscription_id in already_terminal {
subscriptions.remove(&subscription_id.get());
}
}
fn reconnect_backoff(settings: &crate::WsReconnectSettings, attempt: u32) -> std::time::Duration {
let mut backoff = settings.initial_backoff();
let mut step = 1_u32;
while step < attempt {
backoff = std::cmp::min(backoff.saturating_mul(2), settings.max_backoff());
step = step.saturating_add(1);
}
return std::cmp::min(backoff, settings.max_backoff());
}
fn fail_pending_for_reconnect(
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
id: crate::WsSessionId,
code: ksp_core_lib::ErrorCode,
message: &'static str,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
) {
let requests = std::mem::take(pending);
for (request_id, request) in requests {
let error = ksp_core_lib::Error::new(code, message)
.with_context("session_id", id.get().to_string())
.with_context("request_id", request_id.to_string())
.with_context("method", request.method);
match request.response {
PendingWsResponse::Raw(response_tx) => {
let _ = response_tx.send(std::result::Result::Err(error));
},
PendingWsResponse::Subscribe { subscription_id, response_tx } => {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, code);
let _ = response_tx.send(std::result::Result::Err(error));
},
PendingWsResponse::Resubscribe { .. } => {},
PendingWsResponse::Unsubscribe { subscription_id, response_tx } => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let _ = response_tx.send(std::result::Result::Err(error));
},
}
}
}
async fn best_effort_remote_unsubscribe(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
websocket: &mut WsPhysicalStream,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
kind: crate::WsSubscriptionKind,
remote_id: u64,
) -> WsActorIoOutcome {
let request = match crate::JsonRpcRequest::new(0, kind.unsubscribe_method(), std::vec![serde_json::Value::from(remote_id)]) {
std::result::Result::Ok(request) => request,
std::result::Result::Err(_) => return WsActorIoOutcome::Continue,
};
let payload = match request.to_json_string() {
std::result::Result::Ok(payload) => payload,
std::result::Result::Err(_) => return WsActorIoOutcome::Continue,
};
let message = tokio_tungstenite::tungstenite::Message::Text(payload.into());
let send = tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
return WsActorIoOutcome::ShutdownRequested { deadline };
},
send = websocket.send(message) => send,
() = tokio::time::sleep(endpoint.session().command_timeout()) => return WsActorIoOutcome::Continue,
};
if send.is_ok() {
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_kind = kind.as_str(),
"sent best-effort remote WebSocket unsubscribe"
);
}
return WsActorIoOutcome::Continue;
}
fn finish_disconnected_shutdown(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
continuity_gap_count: u64,
overflow_count: u64,
_deadline: tokio::time::Instant,
) {
publish_snapshot(
snapshot_tx,
id,
endpoint,
crate::WsSessionState::Closing,
pending.len(),
WsRuntimeCounters { continuity_gap_count, overflow_count },
subscriptions,
);
fail_all_pending(pending, id, crate::ERROR_CODE_WS_SESSION_CLOSED, "WebSocket session shutdown cancelled pending request");
terminate_all_subscriptions(subscriptions, remote_to_local, crate::WsSubscriptionState::Closed);
publish_snapshot(snapshot_tx, id, endpoint, crate::WsSessionState::Closed, 0, WsRuntimeCounters { continuity_gap_count, overflow_count }, subscriptions);
}
fn finish_disconnected_close(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
continuity_gap_count: u64,
overflow_count: u64,
) {
fail_all_pending(pending, id, crate::ERROR_CODE_WS_SESSION_CLOSED, "WebSocket session handles were dropped during reconnect");
terminate_all_subscriptions(subscriptions, remote_to_local, crate::WsSubscriptionState::Closed);
publish_snapshot(snapshot_tx, id, endpoint, crate::WsSessionState::Closed, 0, WsRuntimeCounters { continuity_gap_count, overflow_count }, subscriptions);
}
fn finish_reconnect_exhaustion(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
continuity_gap_count: u64,
overflow_count: u64,
) {
fail_all_pending(pending, id, crate::ERROR_CODE_WS_CONNECTION_FAILED, "WebSocket reconnect budget was exhausted");
fail_all_subscriptions(subscriptions, remote_to_local, crate::ERROR_CODE_WS_CONNECTION_FAILED);
publish_snapshot(snapshot_tx, id, endpoint, crate::WsSessionState::Failed, 0, WsRuntimeCounters { continuity_gap_count, overflow_count }, subscriptions);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
reconnect_max_retries = endpoint.session().reconnect().max_retries(),
continuity_gap_count,
"physical WebSocket reconnect budget exhausted"
);
}
#[allow(clippy::too_many_arguments)]
async fn handle_session_command<S>(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
websocket: &mut tokio_tungstenite::WebSocketStream<S>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
next_request_id: &mut u64,
next_subscription_id: &mut u64,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
command: WsSessionCommand,
) -> WsActorIoOutcome
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + std::marker::Unpin,
{
return match command {
WsSessionCommand::ExecuteJsonRpc { method, params, response_tx } => {
let write_result = write_json_rpc_request(id, endpoint, websocket, pending.len(), next_request_id, shutdown_rx, method, params).await;
match write_result {
std::result::Result::Ok((request_id, deadline)) => {
pending.insert(request_id, PendingWsRequest { method, deadline, response: PendingWsResponse::Raw(response_tx) });
WsActorIoOutcome::Continue
},
std::result::Result::Err((error, outcome)) => {
let _ = response_tx.send(std::result::Result::Err(error));
outcome
},
}
},
WsSessionCommand::Subscribe { kind, params, dispatcher, response_tx } => {
if subscriptions.len() >= endpoint.session().max_active_subscriptions() {
let error = ksp_core_lib::Error::new(crate::ERROR_CODE_WS_BACKPRESSURE_OVERFLOW, "WebSocket active subscription capacity is exhausted")
.with_context("session_id", id.get().to_string())
.with_context("subscription_kind", kind.as_str());
let _ = response_tx.send(std::result::Result::Err(error));
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_kind = kind.as_str(),
active_subscription_count = subscriptions.len(),
max_active_subscriptions = endpoint.session().max_active_subscriptions(),
"rejected WebSocket subscription because active capacity is exhausted"
);
return WsActorIoOutcome::Continue;
}
let subscription_id_result = next_local_subscription_id(next_subscription_id);
let subscription_id = match subscription_id_result {
std::result::Result::Ok(subscription_id) => subscription_id,
std::result::Result::Err(error) => {
let _ = response_tx.send(std::result::Result::Err(error));
return WsActorIoOutcome::Continue;
},
};
let (state_tx, _) = tokio::sync::watch::channel(crate::WsSubscriptionState::Requested);
let (terminal_error_tx, _) = tokio::sync::watch::channel(std::option::Option::None::<ksp_core_lib::ErrorCode>);
let resubscribe_params = params.clone();
subscriptions.insert(
subscription_id.get(),
crate::WsSubscriptionRuntime {
id: subscription_id,
kind,
state: crate::WsSubscriptionState::Requested,
params: resubscribe_params,
remote_id: std::option::Option::None,
state_tx,
terminal_error_tx,
dispatcher,
},
);
let method = kind.subscribe_method();
let write_result = write_json_rpc_request(id, endpoint, websocket, pending.len(), next_request_id, shutdown_rx, method, params).await;
match write_result {
std::result::Result::Ok((request_id, deadline)) => {
pending.insert(request_id, PendingWsRequest { method, deadline, response: PendingWsResponse::Subscribe { subscription_id, response_tx } });
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
"requested logical WebSocket subscription"
);
WsActorIoOutcome::Continue
},
std::result::Result::Err((error, outcome)) => {
let error_code = error.code();
fail_local_subscription(subscriptions, remote_to_local, subscription_id, error_code);
let _ = response_tx.send(std::result::Result::Err(error));
outcome
},
}
},
WsSessionCommand::Unsubscribe { subscription_id, response_tx } => {
let remote_id = match subscriptions.get_mut(&subscription_id.get()) {
std::option::Option::Some(runtime) => {
runtime.set_state(crate::WsSubscriptionState::Cancelling);
runtime.remote_id.take()
},
std::option::Option::None => {
let _ = response_tx.send(std::result::Result::Ok(false));
return WsActorIoOutcome::Continue;
},
};
let remote_id = match remote_id {
std::option::Option::Some(remote_id) => remote_id,
std::option::Option::None => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let _ = response_tx.send(std::result::Result::Ok(false));
return WsActorIoOutcome::Continue;
},
};
remote_to_local.remove(&remote_id);
let kind = match subscriptions.get(&subscription_id.get()) {
std::option::Option::Some(runtime) => runtime.kind,
std::option::Option::None => {
let _ = response_tx.send(std::result::Result::Ok(false));
return WsActorIoOutcome::Continue;
},
};
let method = kind.unsubscribe_method();
let params = std::vec![serde_json::Value::from(remote_id)];
let write_result = write_json_rpc_request(id, endpoint, websocket, pending.len(), next_request_id, shutdown_rx, method, params).await;
match write_result {
std::result::Result::Ok((request_id, deadline)) => {
pending
.insert(request_id, PendingWsRequest { method, deadline, response: PendingWsResponse::Unsubscribe { subscription_id, response_tx } });
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
"requested logical WebSocket unsubscribe"
);
WsActorIoOutcome::Continue
},
std::result::Result::Err((error, outcome)) => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let _ = response_tx.send(std::result::Result::Err(error));
outcome
},
}
},
};
}
#[allow(clippy::too_many_arguments)]
async fn write_json_rpc_request<S>(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
websocket: &mut tokio_tungstenite::WebSocketStream<S>,
pending_request_count: usize,
next_request_id: &mut u64,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
method: &'static str,
params: std::vec::Vec<serde_json::Value>,
) -> std::result::Result<(u64, tokio::time::Instant), (ksp_core_lib::Error, WsActorIoOutcome)>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + std::marker::Unpin,
{
if pending_request_count >= endpoint.session().max_pending_requests() {
let error = ksp_core_lib::Error::new(crate::ERROR_CODE_WS_BACKPRESSURE_OVERFLOW, "WebSocket pending JSON-RPC request capacity is exhausted")
.with_context("session_id", id.get().to_string())
.with_context("method", method);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
pending_request_count,
max_pending_requests = endpoint.session().max_pending_requests(),
"rejected WebSocket JSON-RPC request because pending capacity is exhausted"
);
return std::result::Result::Err((error, WsActorIoOutcome::Continue));
}
let request_id = *next_request_id;
*next_request_id = match request_id.checked_add(1) {
std::option::Option::Some(value) => value,
std::option::Option::None => {
let error = ksp_core_lib::Error::new(crate::ERROR_CODE_WS_PROTOCOL_ERROR, "WebSocket JSON-RPC request identifier space is exhausted")
.with_context("session_id", id.get().to_string());
return std::result::Result::Err((error, WsActorIoOutcome::Continue));
},
};
let request = match crate::JsonRpcRequest::new(request_id, method, params) {
std::result::Result::Ok(request) => request,
std::result::Result::Err(error) => return std::result::Result::Err((error, WsActorIoOutcome::Continue)),
};
let payload = match request.to_json_string() {
std::result::Result::Ok(payload) => payload,
std::result::Result::Err(error) => return std::result::Result::Err((error, WsActorIoOutcome::Continue)),
};
let payload_size = payload.len();
if payload_size > endpoint.session().max_message_size_bytes() || payload_size > endpoint.session().max_frame_size_bytes() {
let error = ksp_core_lib::Error::new(crate::ERROR_CODE_INVALID_RPC_PARAMETERS, "WebSocket JSON-RPC request exceeds the configured outbound size bound")
.with_context("session_id", id.get().to_string())
.with_context("method", method)
.with_context("payload_size_bytes", payload_size.to_string())
.with_context("max_message_size_bytes", endpoint.session().max_message_size_bytes().to_string())
.with_context("max_frame_size_bytes", endpoint.session().max_frame_size_bytes().to_string());
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
method,
payload_size_bytes = payload_size,
max_message_size_bytes = endpoint.session().max_message_size_bytes(),
max_frame_size_bytes = endpoint.session().max_frame_size_bytes(),
"rejected oversized WebSocket JSON-RPC request before socket write"
);
return std::result::Result::Err((error, WsActorIoOutcome::Continue));
}
ksp_logging_lib::trace!(
target: crate::TRACING_TARGET,
session_id = id.get(),
request_id,
method,
pending_request_count,
"sending WebSocket JSON-RPC request"
);
let message = tokio_tungstenite::tungstenite::Message::Text(payload.into());
let send_result = tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
let error = ws_session_closed_error(id, "WebSocket JSON-RPC request cancelled by session shutdown").with_context("method", method);
return std::result::Result::Err((error, WsActorIoOutcome::ShutdownRequested { deadline }));
},
send_result = websocket.send(message) => send_result,
() = tokio::time::sleep(endpoint.session().command_timeout()) => {
let error = ws_timeout_error(id, "WebSocket JSON-RPC request write exceeded the configured command timeout").with_context("method", method);
ksp_logging_lib::warn!(target: crate::TRACING_TARGET, session_id = id.get(), request_id, method, "WebSocket JSON-RPC request write timed out");
return std::result::Result::Err((
error,
WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection failed while a request write timed out",
},
));
},
};
match send_result {
std::result::Result::Ok(()) => {},
std::result::Result::Err(tokio_tungstenite::tungstenite::Error::WriteBufferFull(_))
| std::result::Result::Err(tokio_tungstenite::tungstenite::Error::Capacity(_)) => {
let error = ksp_core_lib::Error::new(crate::ERROR_CODE_WS_BACKPRESSURE_OVERFLOW, "WebSocket write capacity is exhausted")
.with_context("session_id", id.get().to_string())
.with_context("method", method);
return std::result::Result::Err((error, WsActorIoOutcome::Continue));
},
std::result::Result::Err(_) => {
let error = ws_connection_error(id, endpoint, "WebSocket connection failed while writing a JSON-RPC request").with_context("method", method);
return std::result::Result::Err((
error,
WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection failed while writing a request",
},
));
},
}
let deadline = tokio::time::Instant::now() + endpoint.session().command_timeout();
return std::result::Result::Ok((request_id, deadline));
}
async fn handle_socket_message<S>(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
maybe_message: std::option::Option<std::result::Result<tokio_tungstenite::tungstenite::Message, tokio_tungstenite::tungstenite::Error>>,
websocket: &mut tokio_tungstenite::WebSocketStream<S>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
overflow_count: &mut u64,
shutdown_rx: &mut tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
) -> WsActorIoOutcome
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + std::marker::Unpin,
{
let message = match maybe_message {
std::option::Option::Some(std::result::Result::Ok(message)) => message,
std::option::Option::Some(std::result::Result::Err(error)) => {
let code = classify_websocket_read_error(&error);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
protocol_error = code == crate::ERROR_CODE_WS_PROTOCOL_ERROR,
"physical WebSocket read failed"
);
return WsActorIoOutcome::Failed { code, pending_message: "WebSocket read failed before pending response delivery" };
},
std::option::Option::None => {
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"physical WebSocket stream ended without a Close frame"
);
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection ended before pending response delivery",
};
},
};
return match message {
tokio_tungstenite::tungstenite::Message::Text(text) => handle_text_message(id, text.as_str(), pending, subscriptions, remote_to_local, overflow_count),
tokio_tungstenite::tungstenite::Message::Binary(_) => {
ksp_logging_lib::warn!(target: crate::TRACING_TARGET, session_id = id.get(), "received unexpected binary WebSocket message");
WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket wire data violated the expected JSON text protocol",
}
},
tokio_tungstenite::tungstenite::Message::Ping(_) => {
ksp_logging_lib::trace!(target: crate::TRACING_TARGET, session_id = id.get(), "received WebSocket ping control frame; flushing automatic Pong");
let flush_result = tokio::select! {
biased;
shutdown_changed = shutdown_rx.changed() => {
let deadline = resolve_shutdown_deadline(shutdown_rx, shutdown_changed, endpoint.session().close_timeout());
return WsActorIoOutcome::ShutdownRequested { deadline };
},
flush_result = websocket.flush() => flush_result,
() = tokio::time::sleep(endpoint.session().command_timeout()) => {
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection failed while flushing automatic Pong",
};
},
};
return match flush_result {
std::result::Result::Ok(()) => WsActorIoOutcome::Continue,
std::result::Result::Err(_) => WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_CONNECTION_FAILED,
pending_message: "WebSocket connection failed while flushing automatic Pong",
},
};
},
tokio_tungstenite::tungstenite::Message::Pong(_) => {
ksp_logging_lib::trace!(target: crate::TRACING_TARGET, session_id = id.get(), "received WebSocket pong control frame");
WsActorIoOutcome::Continue
},
tokio_tungstenite::tungstenite::Message::Close(_) => {
ksp_logging_lib::debug!(target: crate::TRACING_TARGET, session_id = id.get(), "remote peer closed physical WebSocket session");
let _ = tokio::time::timeout(endpoint.session().close_timeout(), websocket.flush()).await;
WsActorIoOutcome::RemoteClosed
},
tokio_tungstenite::tungstenite::Message::Frame(_) => {
ksp_logging_lib::trace!(target: crate::TRACING_TARGET, session_id = id.get(), "ignored internal WebSocket frame event");
WsActorIoOutcome::Continue
},
};
}
fn handle_text_message(
id: crate::WsSessionId,
text: &str,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
overflow_count: &mut u64,
) -> WsActorIoOutcome {
let decoded = serde_json::from_str::<serde_json::Value>(text);
let value = match decoded {
std::result::Result::Ok(value) => value,
std::result::Result::Err(_) => {
ksp_logging_lib::warn!(target: crate::TRACING_TARGET, session_id = id.get(), "received malformed JSON WebSocket message");
return WsActorIoOutcome::Failed { code: crate::ERROR_CODE_WS_PROTOCOL_ERROR, pending_message: "WebSocket JSON payload was malformed" };
},
};
let object = match value.as_object() {
std::option::Option::Some(object) => object,
std::option::Option::None => {
ksp_logging_lib::warn!(target: crate::TRACING_TARGET, session_id = id.get(), "received non-object JSON WebSocket message");
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket JSON-RPC payload was not an object",
};
},
};
if !object.contains_key("id") {
return handle_subscription_notification(id, object, subscriptions, remote_to_local, overflow_count);
}
let response_id = match object.get("id").and_then(serde_json::Value::as_u64) {
std::option::Option::Some(response_id) => response_id,
std::option::Option::None => {
ksp_logging_lib::warn!(target: crate::TRACING_TARGET, session_id = id.get(), "received WebSocket JSON-RPC response with invalid id");
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket JSON-RPC response contained an invalid id",
};
},
};
let pending_request = match pending.remove(&response_id) {
std::option::Option::Some(pending_request) => pending_request,
std::option::Option::None => {
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
response_id,
"ignored unknown or stale WebSocket JSON-RPC response id"
);
return WsActorIoOutcome::Continue;
},
};
let parsed = crate::parse_json_rpc_response_value(value, response_id);
return dispatch_pending_response(id, response_id, pending_request, parsed, subscriptions, remote_to_local);
}
fn dispatch_pending_response(
id: crate::WsSessionId,
response_id: u64,
pending_request: PendingWsRequest,
parsed: ksp_core_lib::Result<crate::JsonRpcResponse>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
) -> WsActorIoOutcome {
let response = match parsed {
std::result::Result::Ok(response) => response,
std::result::Result::Err(_) => {
fail_pending_response(
pending_request.response,
id,
crate::ERROR_CODE_WS_PROTOCOL_ERROR,
"WebSocket JSON-RPC response violated protocol invariants",
);
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket JSON-RPC response violated protocol invariants",
};
},
};
let rpc_result = response.into_result();
match pending_request.response {
PendingWsResponse::Raw(response_tx) => {
let result = match rpc_result {
std::result::Result::Ok(value) => std::result::Result::Ok(value),
std::result::Result::Err(error) => std::result::Result::Err(error.with_context("method", pending_request.method)),
};
let _ = response_tx.send(result);
},
PendingWsResponse::Subscribe { subscription_id, response_tx } => match rpc_result {
std::result::Result::Ok(value) => {
let remote_id = match value.as_u64() {
std::option::Option::Some(remote_id) => remote_id,
std::option::Option::None => {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, crate::ERROR_CODE_WS_PROTOCOL_ERROR);
let error = ksp_core_lib::Error::new(
crate::ERROR_CODE_WS_PROTOCOL_ERROR,
"WebSocket subscribe response did not contain a numeric remote subscription id",
)
.with_context("session_id", id.get().to_string())
.with_context("subscription_id", subscription_id.get().to_string());
let _ = response_tx.send(std::result::Result::Err(error));
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket subscribe response violated protocol invariants",
};
},
};
if remote_to_local.contains_key(&remote_id) {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, crate::ERROR_CODE_WS_PROTOCOL_ERROR);
let error = ksp_core_lib::Error::new(crate::ERROR_CODE_WS_PROTOCOL_ERROR, "WebSocket endpoint reused an active remote subscription id")
.with_context("session_id", id.get().to_string())
.with_context("subscription_id", subscription_id.get().to_string());
let _ = response_tx.send(std::result::Result::Err(error));
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket endpoint reused an active remote subscription id",
};
}
let registration = match subscriptions.get_mut(&subscription_id.get()) {
std::option::Option::Some(runtime) => {
runtime.remote_id = std::option::Option::Some(remote_id);
runtime.set_state(crate::WsSubscriptionState::Active);
remote_to_local.insert(remote_id, subscription_id);
crate::WsSubscriptionRegistration::new(
subscription_id,
runtime.kind,
runtime.state_tx.subscribe(),
runtime.terminal_error_tx.subscribe(),
)
},
std::option::Option::None => {
let error = ksp_core_lib::Error::new(
crate::ERROR_CODE_WS_PROTOCOL_ERROR,
"WebSocket local subscription disappeared before subscribe acknowledgement",
)
.with_context("session_id", id.get().to_string())
.with_context("subscription_id", subscription_id.get().to_string());
let _ = response_tx.send(std::result::Result::Err(error));
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket local subscription registry lost a pending entry",
};
},
};
let _ = response_tx.send(std::result::Result::Ok(registration));
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
"logical WebSocket subscription is active"
);
},
std::result::Result::Err(error) => {
let error_code = error.code();
fail_local_subscription(subscriptions, remote_to_local, subscription_id, error_code);
let _ = response_tx.send(std::result::Result::Err(error.with_context("method", pending_request.method)));
},
},
PendingWsResponse::Resubscribe { subscription_id, kind } => match rpc_result {
std::result::Result::Ok(value) => {
let remote_id = match value.as_u64() {
std::option::Option::Some(remote_id) => remote_id,
std::option::Option::None => {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, crate::ERROR_CODE_WS_PROTOCOL_ERROR);
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket resubscribe response violated protocol invariants",
};
},
};
if remote_to_local.contains_key(&remote_id) {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, crate::ERROR_CODE_WS_PROTOCOL_ERROR);
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket endpoint reused an active remote subscription id during resubscribe",
};
}
let can_bind =
subscriptions.get(&subscription_id.get()).is_some_and(|runtime| return runtime.state == crate::WsSubscriptionState::Resubscribing);
if !can_bind {
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
"late WebSocket resubscribe acknowledgement lost to local cancellation"
);
return WsActorIoOutcome::BestEffortUnsubscribe { kind, remote_id };
}
if let std::option::Option::Some(runtime) = subscriptions.get_mut(&subscription_id.get()) {
runtime.remote_id = std::option::Option::Some(remote_id);
runtime.set_state(crate::WsSubscriptionState::Active);
}
remote_to_local.insert(remote_id, subscription_id);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
"logical WebSocket subscription restored with a new remote binding"
);
},
std::result::Result::Err(error) => {
let error_code = error.code();
fail_local_subscription(subscriptions, remote_to_local, subscription_id, error_code);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
subscription_kind = kind.as_str(),
error_code = error_code.code(),
"remote WebSocket resubscribe returned an application error"
);
},
},
PendingWsResponse::Unsubscribe { subscription_id, response_tx } => match rpc_result {
std::result::Result::Ok(value) => {
let removed = match value.as_bool() {
std::option::Option::Some(removed) => removed,
std::option::Option::None => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let error =
ksp_core_lib::Error::new(crate::ERROR_CODE_WS_PROTOCOL_ERROR, "WebSocket unsubscribe response did not contain a boolean result")
.with_context("session_id", id.get().to_string())
.with_context("subscription_id", subscription_id.get().to_string());
let _ = response_tx.send(std::result::Result::Err(error));
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket unsubscribe response violated protocol invariants",
};
},
};
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let _ = response_tx.send(std::result::Result::Ok(removed));
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = subscription_id.get(),
remote_removed = removed,
"logical WebSocket subscription is closed"
);
},
std::result::Result::Err(error) => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let _ = response_tx.send(std::result::Result::Err(error.with_context("method", pending_request.method)));
},
},
}
ksp_logging_lib::trace!(
target: crate::TRACING_TARGET,
session_id = id.get(),
response_id,
method = pending_request.method,
"dispatched WebSocket JSON-RPC response"
);
return WsActorIoOutcome::Continue;
}
fn handle_subscription_notification(
id: crate::WsSessionId,
object: &serde_json::Map<std::string::String, serde_json::Value>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
overflow_count: &mut u64,
) -> WsActorIoOutcome {
if object.get("jsonrpc").and_then(serde_json::Value::as_str) != std::option::Option::Some("2.0") {
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket notification used an invalid JSON-RPC version",
};
}
let method = match object.get("method").and_then(serde_json::Value::as_str) {
std::option::Option::Some(method) => method,
std::option::Option::None => {
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket notification omitted its method",
};
},
};
let params = match object.get("params").and_then(serde_json::Value::as_object) {
std::option::Option::Some(params) => params,
std::option::Option::None => {
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket notification omitted its params object",
};
},
};
let remote_id = match params.get("subscription").and_then(serde_json::Value::as_u64) {
std::option::Option::Some(remote_id) => remote_id,
std::option::Option::None => {
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket notification contained an invalid remote subscription id",
};
},
};
let result = match params.get("result") {
std::option::Option::Some(result) => result.clone(),
std::option::Option::None => {
return WsActorIoOutcome::Failed {
code: crate::ERROR_CODE_WS_PROTOCOL_ERROR,
pending_message: "WebSocket notification omitted its result payload",
};
},
};
let local_id = match remote_to_local.get(&remote_id).copied() {
std::option::Option::Some(local_id) => local_id,
std::option::Option::None => {
ksp_logging_lib::debug!(target: crate::TRACING_TARGET, session_id = id.get(), "ignored notification for unknown or stale remote subscription id");
return WsActorIoOutcome::Continue;
},
};
let runtime = match subscriptions.get(&local_id.get()) {
std::option::Option::Some(runtime) => runtime,
std::option::Option::None => {
remote_to_local.remove(&remote_id);
ksp_logging_lib::debug!(target: crate::TRACING_TARGET, session_id = id.get(), "discarded stale remote-to-local WebSocket subscription mapping");
return WsActorIoOutcome::Continue;
},
};
if method != runtime.kind.notification_method() {
let subscription_kind = runtime.kind;
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = local_id.get(),
subscription_kind = subscription_kind.as_str(),
"WebSocket notification method mismatched the registered subscription family"
);
fail_local_subscription(subscriptions, remote_to_local, local_id, crate::ERROR_CODE_WS_PROTOCOL_ERROR);
return WsActorIoOutcome::BestEffortUnsubscribe { kind: subscription_kind, remote_id };
}
let subscription_kind = runtime.kind;
let dispatch = (runtime.dispatcher)(result);
return match dispatch {
crate::WsNotificationDispatchOutcome::Delivered => WsActorIoOutcome::Continue,
crate::WsNotificationDispatchOutcome::DeliveredTerminal => {
close_local_subscription(subscriptions, remote_to_local, local_id, crate::WsSubscriptionState::Closed);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = local_id.get(),
subscription_kind = subscription_kind.as_str(),
"logical WebSocket subscription observed its server-terminal notification"
);
WsActorIoOutcome::Continue
},
crate::WsNotificationDispatchOutcome::ReceiverClosed => {
close_local_subscription(subscriptions, remote_to_local, local_id, crate::WsSubscriptionState::Closed);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = local_id.get(),
subscription_kind = subscription_kind.as_str(),
"logical WebSocket subscription receiver was dropped; scheduling remote cleanup"
);
WsActorIoOutcome::BestEffortUnsubscribe { kind: subscription_kind, remote_id }
},
crate::WsNotificationDispatchOutcome::QueueFull => {
*overflow_count = (*overflow_count).saturating_add(1);
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = local_id.get(),
subscription_kind = subscription_kind.as_str(),
overflow_count = *overflow_count,
"bounded WebSocket subscription notification queue overflowed; failing only the slow subscription"
);
fail_local_subscription(subscriptions, remote_to_local, local_id, crate::ERROR_CODE_WS_BACKPRESSURE_OVERFLOW);
WsActorIoOutcome::BestEffortUnsubscribe { kind: subscription_kind, remote_id }
},
crate::WsNotificationDispatchOutcome::DecodeFailed { code } => {
ksp_logging_lib::warn!(
target: crate::TRACING_TARGET,
session_id = id.get(),
subscription_id = local_id.get(),
subscription_kind = subscription_kind.as_str(),
error_code = code.code(),
"typed WebSocket notification decoding failed for one subscription"
);
fail_local_subscription(subscriptions, remote_to_local, local_id, code);
WsActorIoOutcome::BestEffortUnsubscribe { kind: subscription_kind, remote_id }
},
};
}
async fn close_session_actor<S>(
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
websocket: &mut tokio_tungstenite::WebSocketStream<S>,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
continuity_gap_count: u64,
overflow_count: u64,
deadline: tokio::time::Instant,
) where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + std::marker::Unpin,
{
publish_snapshot(
snapshot_tx,
id,
endpoint,
crate::WsSessionState::Closing,
pending.len(),
WsRuntimeCounters { continuity_gap_count, overflow_count },
subscriptions,
);
fail_all_pending(pending, id, crate::ERROR_CODE_WS_SESSION_CLOSED, "WebSocket session shutdown cancelled the pending request");
terminate_all_subscriptions(subscriptions, remote_to_local, crate::WsSubscriptionState::Closed);
publish_snapshot(snapshot_tx, id, endpoint, crate::WsSessionState::Closing, 0, WsRuntimeCounters { continuity_gap_count, overflow_count }, subscriptions);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
endpoint_name = endpoint.name(),
"closing physical WebSocket session"
);
let now = tokio::time::Instant::now();
let remaining = deadline.saturating_duration_since(now);
let io_deadline = now + (remaining / 2);
let close_wait = tokio::time::timeout_at(io_deadline, websocket.send(tokio_tungstenite::tungstenite::Message::Close(std::option::Option::None))).await;
match close_wait {
std::result::Result::Ok(std::result::Result::Ok(())) => {
ksp_logging_lib::trace!(target: crate::TRACING_TARGET, session_id = id.get(), "sent WebSocket Close control frame");
},
std::result::Result::Ok(std::result::Result::Err(_)) => {
ksp_logging_lib::debug!(target: crate::TRACING_TARGET, session_id = id.get(), "best-effort WebSocket Close frame could not be sent");
},
std::result::Result::Err(_) => {
ksp_logging_lib::debug!(target: crate::TRACING_TARGET, session_id = id.get(), "best-effort WebSocket Close frame reached close deadline");
},
}
publish_snapshot(snapshot_tx, id, endpoint, crate::WsSessionState::Closed, 0, WsRuntimeCounters { continuity_gap_count, overflow_count }, subscriptions);
ksp_logging_lib::debug!(target: crate::TRACING_TARGET, session_id = id.get(), endpoint_name = endpoint.name(), "physical WebSocket session is closed");
}
fn classify_websocket_read_error(error: &tokio_tungstenite::tungstenite::Error) -> ksp_core_lib::ErrorCode {
return match error {
tokio_tungstenite::tungstenite::Error::Capacity(_)
| tokio_tungstenite::tungstenite::Error::Protocol(_)
| tokio_tungstenite::tungstenite::Error::Utf8(_)
| tokio_tungstenite::tungstenite::Error::AttackAttempt => crate::ERROR_CODE_WS_PROTOCOL_ERROR,
_ => crate::ERROR_CODE_WS_CONNECTION_FAILED,
};
}
fn resolve_shutdown_deadline(
shutdown_rx: &tokio::sync::watch::Receiver<std::option::Option<tokio::time::Instant>>,
shutdown_changed: std::result::Result<(), tokio::sync::watch::error::RecvError>,
close_timeout: std::time::Duration,
) -> tokio::time::Instant {
if shutdown_changed.is_ok() {
let requested = shutdown_rx.borrow().to_owned();
if let std::option::Option::Some(deadline) = requested {
return deadline;
}
}
return tokio::time::Instant::now() + close_timeout;
}
fn next_pending_deadline(pending: &std::collections::BTreeMap<u64, PendingWsRequest>) -> tokio::time::Instant {
let mut earliest = std::option::Option::None::<tokio::time::Instant>;
for request in pending.values() {
earliest = match earliest {
std::option::Option::Some(current) if current <= request.deadline => std::option::Option::Some(current),
_ => std::option::Option::Some(request.deadline),
};
}
return match earliest {
std::option::Option::Some(deadline) => deadline,
std::option::Option::None => tokio::time::Instant::now() + std::time::Duration::from_secs(86_400),
};
}
fn expire_pending_requests(
id: crate::WsSessionId,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
) {
let now = tokio::time::Instant::now();
let mut expired_ids = std::vec::Vec::new();
for (request_id, request) in pending.iter() {
if request.deadline <= now {
expired_ids.push(*request_id);
}
}
for request_id in expired_ids {
if let std::option::Option::Some(request) = pending.remove(&request_id) {
expire_pending_response(id, request.response, subscriptions, remote_to_local, request.method);
ksp_logging_lib::debug!(
target: crate::TRACING_TARGET,
session_id = id.get(),
request_id,
method = request.method,
"expired pending WebSocket JSON-RPC request"
);
}
}
}
fn expire_pending_response(
id: crate::WsSessionId,
response: PendingWsResponse,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
method: &'static str,
) {
let error = ws_timeout_error(id, "WebSocket JSON-RPC request timed out while awaiting the remote response").with_context("method", method);
match response {
PendingWsResponse::Raw(response_tx) => {
let _ = response_tx.send(std::result::Result::Err(error));
},
PendingWsResponse::Subscribe { subscription_id, response_tx } => {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, crate::ERROR_CODE_TIMEOUT);
let _ = response_tx.send(std::result::Result::Err(error));
},
PendingWsResponse::Resubscribe { subscription_id, .. } => {
fail_local_subscription(subscriptions, remote_to_local, subscription_id, crate::ERROR_CODE_TIMEOUT);
},
PendingWsResponse::Unsubscribe { subscription_id, response_tx } => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
let _ = response_tx.send(std::result::Result::Err(error));
},
}
}
fn prune_cancelled_pending(
id: crate::WsSessionId,
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
) {
let mut cancelled_ids = std::vec::Vec::new();
for (request_id, request) in pending.iter() {
if pending_response_is_closed(&request.response) {
cancelled_ids.push(*request_id);
}
}
for request_id in cancelled_ids {
if let std::option::Option::Some(request) = pending.remove(&request_id) {
match request.response {
PendingWsResponse::Raw(_) => {},
PendingWsResponse::Subscribe { subscription_id, .. } => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
},
PendingWsResponse::Resubscribe { .. } => {},
PendingWsResponse::Unsubscribe { subscription_id, .. } => {
close_local_subscription(subscriptions, remote_to_local, subscription_id, crate::WsSubscriptionState::Closed);
},
}
ksp_logging_lib::trace!(
target: crate::TRACING_TARGET,
session_id = id.get(),
request_id,
method = request.method,
"removed abandoned WebSocket JSON-RPC request after caller cancellation"
);
}
}
}
fn pending_response_is_closed(response: &PendingWsResponse) -> bool {
return match response {
PendingWsResponse::Raw(response_tx) => response_tx.is_closed(),
PendingWsResponse::Subscribe { response_tx, .. } => response_tx.is_closed(),
PendingWsResponse::Resubscribe { .. } => false,
PendingWsResponse::Unsubscribe { response_tx, .. } => response_tx.is_closed(),
};
}
fn fail_all_pending(
pending: &mut std::collections::BTreeMap<u64, PendingWsRequest>,
id: crate::WsSessionId,
code: ksp_core_lib::ErrorCode,
message: &'static str,
) {
let requests = std::mem::take(pending);
for (request_id, request) in requests {
let error = ksp_core_lib::Error::new(code, message)
.with_context("session_id", id.get().to_string())
.with_context("request_id", request_id.to_string())
.with_context("method", request.method);
fail_pending_response_with_error(request.response, error);
}
}
fn fail_pending_response(response: PendingWsResponse, id: crate::WsSessionId, code: ksp_core_lib::ErrorCode, message: &'static str) {
let error = ksp_core_lib::Error::new(code, message).with_context("session_id", id.get().to_string());
fail_pending_response_with_error(response, error);
}
fn fail_pending_response_with_error(response: PendingWsResponse, error: ksp_core_lib::Error) {
match response {
PendingWsResponse::Raw(response_tx) => {
let _ = response_tx.send(std::result::Result::Err(error));
},
PendingWsResponse::Subscribe { response_tx, .. } => {
let _ = response_tx.send(std::result::Result::Err(error));
},
PendingWsResponse::Resubscribe { .. } => {},
PendingWsResponse::Unsubscribe { response_tx, .. } => {
let _ = response_tx.send(std::result::Result::Err(error));
},
}
}
fn terminate_all_subscriptions(
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
terminal_state: crate::WsSubscriptionState,
) {
for runtime in subscriptions.values_mut() {
runtime.remote_id = std::option::Option::None;
runtime.set_state(terminal_state);
}
remote_to_local.clear();
}
fn fail_all_subscriptions(
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
error_code: ksp_core_lib::ErrorCode,
) {
for runtime in subscriptions.values_mut() {
runtime.remote_id = std::option::Option::None;
runtime.fail_with_code(error_code);
}
remote_to_local.clear();
}
fn fail_local_subscription(
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
subscription_id: crate::WsSubscriptionId,
error_code: ksp_core_lib::ErrorCode,
) {
if let std::option::Option::Some(mut runtime) = subscriptions.remove(&subscription_id.get()) {
if let std::option::Option::Some(remote_id) = runtime.remote_id.take() {
remote_to_local.remove(&remote_id);
}
runtime.fail_with_code(error_code);
}
}
fn close_local_subscription(
subscriptions: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
remote_to_local: &mut std::collections::BTreeMap<u64, crate::WsSubscriptionId>,
subscription_id: crate::WsSubscriptionId,
terminal_state: crate::WsSubscriptionState,
) {
if let std::option::Option::Some(mut runtime) = subscriptions.remove(&subscription_id.get()) {
if let std::option::Option::Some(remote_id) = runtime.remote_id.take() {
remote_to_local.remove(&remote_id);
}
runtime.set_state(terminal_state);
}
}
fn publish_snapshot(
snapshot_tx: &tokio::sync::watch::Sender<crate::WsSessionSnapshot>,
id: crate::WsSessionId,
endpoint: &crate::WsEndpointSettings,
state: crate::WsSessionState,
pending_request_count: usize,
counters: WsRuntimeCounters,
subscriptions: &std::collections::BTreeMap<u64, crate::WsSubscriptionRuntime>,
) {
let subscription_snapshots = subscriptions.values().map(crate::WsSubscriptionRuntime::snapshot).collect();
let snapshot = crate::WsSessionSnapshot::new(
id,
endpoint.name(),
endpoint.provider().clone(),
endpoint.cluster().clone(),
endpoint.protocol(),
state,
pending_request_count,
counters.continuity_gap_count,
counters.overflow_count,
subscription_snapshots,
);
snapshot_tx.send_replace(snapshot);
}
fn next_local_subscription_id(next_subscription_id: &mut u64) -> ksp_core_lib::Result<crate::WsSubscriptionId> {
let value = *next_subscription_id;
*next_subscription_id = match value.checked_add(1) {
std::option::Option::Some(next) => next,
std::option::Option::None => {
return std::result::Result::Err(ksp_core_lib::Error::new(
crate::ERROR_CODE_WS_PROTOCOL_ERROR,
"WebSocket local subscription identity space is exhausted",
));
},
};
return match std::num::NonZeroU64::new(value) {
std::option::Option::Some(value) => std::result::Result::Ok(crate::WsSubscriptionId::new(value)),
std::option::Option::None => std::result::Result::Err(ksp_core_lib::Error::new(
crate::ERROR_CODE_WS_PROTOCOL_ERROR,
"WebSocket local subscription identity generator produced zero",
)),
};
}
fn next_session_id() -> ksp_core_lib::Result<crate::WsSessionId> {
let update_result = NEXT_WS_SESSION_ID.fetch_update(std::sync::atomic::Ordering::Relaxed, std::sync::atomic::Ordering::Relaxed, |current| {
return current.checked_add(1);
});
let value = match update_result {
std::result::Result::Ok(value) => value,
std::result::Result::Err(_) => {
return std::result::Result::Err(ksp_core_lib::Error::new(crate::ERROR_CODE_WS_PROTOCOL_ERROR, "WebSocket session identity space is exhausted"));
},
};
return match std::num::NonZeroU64::new(value) {
std::option::Option::Some(value) => std::result::Result::Ok(crate::WsSessionId::new(value)),
std::option::Option::None => {
std::result::Result::Err(ksp_core_lib::Error::new(crate::ERROR_CODE_WS_PROTOCOL_ERROR, "WebSocket session identity generator produced zero"))
},
};
}
fn ws_connection_error(id: crate::WsSessionId, endpoint: &crate::WsEndpointSettings, message: &'static str) -> ksp_core_lib::Error {
return ksp_core_lib::Error::new(crate::ERROR_CODE_WS_CONNECTION_FAILED, message)
.with_context("session_id", id.get().to_string())
.with_context("endpoint_name", endpoint.name())
.with_context("provider", endpoint.provider().as_str())
.with_context("cluster", endpoint.cluster().as_str());
}
fn ws_session_closed_error(id: crate::WsSessionId, message: &'static str) -> ksp_core_lib::Error {
return ksp_core_lib::Error::new(crate::ERROR_CODE_WS_SESSION_CLOSED, message).with_context("session_id", id.get().to_string());
}
fn ws_timeout_error(id: crate::WsSessionId, message: &'static str) -> ksp_core_lib::Error {
return ksp_core_lib::Error::new(crate::ERROR_CODE_TIMEOUT, message).with_context("session_id", id.get().to_string());
}
#[cfg(test)]
#[path = "../unit_tests/ws_session.rs"]
mod tests;
Reference in New Issue View Git Blame Copy Permalink