387 lines
20 KiB
Rust
387 lines
20 KiB
Rust
// file: kb-onchain-transport/src/ws_client.rs
|
|
// version: 8
|
|
|
|
//! Standard Solana WebSocket client helpers.
|
|
|
|
use futures_util::SinkExt; // rust-rules: trait-import
|
|
use futures_util::StreamExt; // rust-rules: trait-import
|
|
|
|
/// Snapshot of one pooled WebSocket endpoint.
|
|
#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct WsPoolClientSnapshot {
|
|
/// Logical endpoint name.
|
|
pub endpoint_name: std::string::String,
|
|
/// Provider name.
|
|
pub provider: std::string::String,
|
|
/// Endpoint URL.
|
|
pub endpoint_url: std::string::String,
|
|
/// Supported roles.
|
|
pub roles: std::vec::Vec<crate::EndpointRoleSnapshot>,
|
|
/// Status string.
|
|
pub status: std::string::String,
|
|
}
|
|
|
|
/// Standard Solana WebSocket client bound to one configured endpoint.
|
|
#[derive(Clone, Debug)]
|
|
pub struct WsClient {
|
|
endpoint: kb_config::WsEndpointConfig,
|
|
next_request_id: std::sync::Arc<std::sync::atomic::AtomicU64>,
|
|
}
|
|
|
|
impl crate::WsClient {
|
|
/// Creates a new WebSocket client bound to one endpoint.
|
|
pub fn new(endpoint: kb_config::WsEndpointConfig) -> kb_core::Result<Self> {
|
|
if !endpoint.enabled {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "create_ws_client", endpoint_name = %endpoint.name, provider = %endpoint.provider, error_code = "ws_endpoint_disabled", "cannot create WebSocket client for disabled endpoint");
|
|
return std::result::Result::Err(kb_core::Error::config(format!(
|
|
"ws endpoint '{}' is disabled",
|
|
endpoint.name
|
|
)));
|
|
}
|
|
tracing::debug!(target: crate::TRACING_TARGET, action = "create_ws_client", endpoint_name = %endpoint.name, provider = %endpoint.provider, role_count = endpoint.roles.len(), "WebSocket client created");
|
|
return std::result::Result::Ok(Self {
|
|
endpoint,
|
|
next_request_id: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(1)),
|
|
});
|
|
}
|
|
|
|
/// Returns the endpoint name.
|
|
pub fn endpoint_name(&self) -> &str {
|
|
return self.endpoint.name.as_str();
|
|
}
|
|
|
|
/// Returns the provider name.
|
|
pub fn provider(&self) -> &str {
|
|
return self.endpoint.provider.as_str();
|
|
}
|
|
|
|
/// Returns the endpoint URL.
|
|
pub fn endpoint_url(&self) -> &str {
|
|
return self.endpoint.url.as_str();
|
|
}
|
|
|
|
/// Returns the endpoint configuration.
|
|
pub fn endpoint_config(&self) -> &kb_config::WsEndpointConfig {
|
|
return &self.endpoint;
|
|
}
|
|
|
|
/// Returns true when this endpoint supports the required role and request kind.
|
|
pub fn can_handle(&self, required_role: &str, request_kind: &str) -> bool {
|
|
if !self.endpoint.enabled {
|
|
return false;
|
|
}
|
|
for role in &self.endpoint.roles {
|
|
if crate::role_matches(role, required_role, request_kind) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// Returns a serializable endpoint snapshot.
|
|
pub fn snapshot(&self) -> crate::WsPoolClientSnapshot {
|
|
let mut roles = std::vec::Vec::new();
|
|
for role in &self.endpoint.roles {
|
|
roles.push(crate::EndpointRoleSnapshot::from_config(role));
|
|
}
|
|
return crate::WsPoolClientSnapshot {
|
|
endpoint_name: self.endpoint.name.clone(),
|
|
provider: self.endpoint.provider.clone(),
|
|
endpoint_url: self.endpoint.url.clone(),
|
|
roles,
|
|
status: "idle".to_string(),
|
|
};
|
|
}
|
|
|
|
/// Builds a JSON-RPC request with a generated id.
|
|
pub fn build_json_rpc_request(
|
|
&self,
|
|
method: std::string::String,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
) -> crate::JsonRpcRequest {
|
|
let request_id = self.next_request_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
|
return crate::JsonRpcRequest::new_with_u64_id(request_id, method, params);
|
|
}
|
|
|
|
/// Builds a subscribe request for one explicitly registered standard subscription.
|
|
pub fn build_standard_subscribe_request(
|
|
&self,
|
|
subscription: &crate::StandardWsSubscriptionSpec,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
) -> crate::JsonRpcRequest {
|
|
return self.build_json_rpc_request(subscription.subscribe_method.to_string(), params);
|
|
}
|
|
|
|
/// Builds an unsubscribe request for one explicitly registered standard subscription.
|
|
pub fn build_standard_unsubscribe_request(
|
|
&self,
|
|
subscription: &crate::StandardWsSubscriptionSpec,
|
|
subscription_id: u64,
|
|
) -> crate::JsonRpcRequest {
|
|
return self.build_json_rpc_request(
|
|
subscription.unsubscribe_method.to_string(),
|
|
std::vec![serde_json::Value::from(subscription_id)],
|
|
);
|
|
}
|
|
|
|
/// Connects, sends one JSON-RPC request, waits for one response and closes.
|
|
pub async fn execute_json_rpc_once(
|
|
&self,
|
|
method: std::string::String,
|
|
params: std::vec::Vec<serde_json::Value>,
|
|
) -> kb_core::Result<crate::JsonRpcResponse> {
|
|
let parameter_count = params.len();
|
|
let request = self.build_json_rpc_request(method.clone(), params);
|
|
tracing::debug!(target: crate::TRACING_TARGET, action = "execute_ws_json_rpc_once", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, parameter_count, "start one-shot WebSocket JSON-RPC request");
|
|
let request_text = match request.to_json_string() {
|
|
std::result::Result::Ok(text) => text,
|
|
std::result::Result::Err(error) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "serialize_ws_json_rpc_request", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket JSON-RPC request serialization failed");
|
|
return std::result::Result::Err(error);
|
|
},
|
|
};
|
|
let connect_timeout = std::time::Duration::from_millis(self.endpoint.connect_timeout_ms);
|
|
let connect_future = tokio_tungstenite::connect_async(self.endpoint.url.as_str());
|
|
let connect_timeout_result = tokio::time::timeout(connect_timeout, connect_future).await;
|
|
let connect_result = match connect_timeout_result {
|
|
std::result::Result::Ok(result) => result,
|
|
std::result::Result::Err(_) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "connect_ws_endpoint", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, timeout_ms = self.endpoint.connect_timeout_ms, error_code = "ws_connect_timeout", "WebSocket endpoint connection timed out");
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket connect timed out for endpoint '{}'",
|
|
self.endpoint.name
|
|
)));
|
|
},
|
|
};
|
|
let (mut stream, _response) = match connect_result {
|
|
std::result::Result::Ok(pair) => pair,
|
|
std::result::Result::Err(error) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "connect_ws_endpoint", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket endpoint connection failed");
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"cannot connect websocket endpoint '{}': {error}",
|
|
self.endpoint.name
|
|
)));
|
|
},
|
|
};
|
|
let send_result = stream
|
|
.send(tokio_tungstenite::tungstenite::Message::Text(request_text.into()))
|
|
.await;
|
|
if let std::result::Result::Err(error) = send_result {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "send_ws_json_rpc_request", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket JSON-RPC request send failed");
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"cannot send websocket request '{}' to endpoint '{}': {error}",
|
|
method, self.endpoint.name
|
|
)));
|
|
}
|
|
let response_timeout = std::time::Duration::from_millis(self.endpoint.request_timeout_ms);
|
|
let next_timeout_result = tokio::time::timeout(response_timeout, stream.next()).await;
|
|
let next_result = match next_timeout_result {
|
|
std::result::Result::Ok(result) => result,
|
|
std::result::Result::Err(_) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, timeout_ms = self.endpoint.request_timeout_ms, error_code = "ws_response_timeout", "WebSocket JSON-RPC response timed out");
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket response timed out for endpoint '{}'",
|
|
self.endpoint.name
|
|
)));
|
|
},
|
|
};
|
|
let message = match next_result {
|
|
std::option::Option::Some(result) => match result {
|
|
std::result::Result::Ok(message) => message,
|
|
std::result::Result::Err(error) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket JSON-RPC response read failed");
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket read failed for endpoint '{}': {error}",
|
|
self.endpoint.name
|
|
)));
|
|
},
|
|
},
|
|
std::option::Option::None => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error_code = "ws_closed_before_response", "WebSocket endpoint closed before response");
|
|
return std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket endpoint '{}' closed before response",
|
|
self.endpoint.name
|
|
)));
|
|
},
|
|
};
|
|
let close_result = stream
|
|
.send(tokio_tungstenite::tungstenite::Message::Close(std::option::Option::None))
|
|
.await;
|
|
if let std::result::Result::Err(error) = close_result {
|
|
tracing::debug!(target: crate::TRACING_TARGET, action = "close_ws_endpoint", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket close send failed");
|
|
}
|
|
return match message {
|
|
tokio_tungstenite::tungstenite::Message::Text(text) => {
|
|
let parse_result = crate::parse_json_rpc_text(text.as_str());
|
|
match parse_result {
|
|
std::result::Result::Ok(response) => {
|
|
if let crate::JsonRpcResponse::Error(error_response) = &response {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "execute_ws_json_rpc_once", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, response_kind = response.kind_name(), rpc_error_code = error_response.error.code, rpc_error_message = %error_response.error.message, "WebSocket JSON-RPC endpoint returned an RPC error");
|
|
} else {
|
|
tracing::debug!(target: crate::TRACING_TARGET, action = "execute_ws_json_rpc_once", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, response_kind = response.kind_name(), "one-shot WebSocket JSON-RPC request completed");
|
|
}
|
|
std::result::Result::Ok(response)
|
|
},
|
|
std::result::Result::Err(error) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "parse_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, response_byte_length = text.len(), error = %error, "WebSocket JSON-RPC response parsing failed");
|
|
std::result::Result::Err(error)
|
|
},
|
|
}
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Binary(binary) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "binary", response_byte_length = binary.len(), "WebSocket endpoint returned binary data before JSON response");
|
|
std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket endpoint '{}' returned binary message with {} bytes",
|
|
self.endpoint.name,
|
|
binary.len()
|
|
)))
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Ping(_) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "ping", "WebSocket endpoint returned ping before JSON response");
|
|
std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket endpoint '{}' returned ping before json response",
|
|
self.endpoint.name
|
|
)))
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Pong(_) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "pong", "WebSocket endpoint returned pong before JSON response");
|
|
std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket endpoint '{}' returned pong before json response",
|
|
self.endpoint.name
|
|
)))
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Close(_) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "close", "WebSocket endpoint closed before JSON response");
|
|
std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket endpoint '{}' closed before json response",
|
|
self.endpoint.name
|
|
)))
|
|
},
|
|
tokio_tungstenite::tungstenite::Message::Frame(_) => {
|
|
tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "frame", "WebSocket endpoint returned raw frame before JSON response");
|
|
std::result::Result::Err(kb_core::Error::ws(format!(
|
|
"websocket endpoint '{}' returned raw frame before json response",
|
|
self.endpoint.name
|
|
)))
|
|
},
|
|
};
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
fn role_config(
|
|
role: &str,
|
|
request_kinds: std::vec::Vec<std::string::String>,
|
|
) -> kb_config::EndpointRoleConfig {
|
|
return kb_config::EndpointRoleConfig {
|
|
role: role.to_string(),
|
|
enabled: true,
|
|
request_kinds,
|
|
priority: 1,
|
|
requests_per_second: 10,
|
|
burst_capacity: 10,
|
|
max_concurrent_requests: 4,
|
|
max_subscriptions: 16,
|
|
pause_after_rate_limit_ms: 1500,
|
|
};
|
|
}
|
|
|
|
fn endpoint(enabled: bool) -> kb_config::WsEndpointConfig {
|
|
return kb_config::WsEndpointConfig {
|
|
name: "ws_a".to_string(),
|
|
enabled,
|
|
provider: "test".to_string(),
|
|
cluster: "devnet".to_string(),
|
|
url: "wss://example.invalid".to_string(),
|
|
connect_timeout_ms: 100,
|
|
request_timeout_ms: 100,
|
|
unsubscribe_timeout_ms: 100,
|
|
write_channel_capacity: 8,
|
|
event_channel_capacity: 16,
|
|
auto_reconnect: false,
|
|
roles: std::vec![
|
|
role_config("slot_notifications", std::vec!["slot_subscribe".to_string()]),
|
|
role_config("program_subscribe", std::vec!["program_subscribe".to_string()]),
|
|
role_config("ws_any", std::vec!["*".to_string()]),
|
|
],
|
|
};
|
|
}
|
|
|
|
#[test]
|
|
fn new_rejects_disabled_endpoint() {
|
|
let result = crate::WsClient::new(endpoint(false));
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn can_handle_matches_exact_role_and_kind() {
|
|
let client = match crate::WsClient::new(endpoint(true)) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client creation failed: {error}"),
|
|
};
|
|
assert!(client.can_handle("slot_notifications", "slot_subscribe"));
|
|
assert!(!client.can_handle("slot_notifications", "root_subscribe"));
|
|
}
|
|
|
|
#[test]
|
|
fn can_handle_matches_wildcard_kind() {
|
|
let client = match crate::WsClient::new(endpoint(true)) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client creation failed: {error}"),
|
|
};
|
|
assert!(client.can_handle("ws_any", "logs_subscribe_mentions"));
|
|
}
|
|
|
|
#[test]
|
|
fn snapshot_preserves_endpoint_metadata() {
|
|
let client = match crate::WsClient::new(endpoint(true)) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client creation failed: {error}"),
|
|
};
|
|
let snapshot = client.snapshot();
|
|
assert_eq!(snapshot.endpoint_name, "ws_a");
|
|
assert_eq!(snapshot.provider, "test");
|
|
assert_eq!(snapshot.endpoint_url, "wss://example.invalid");
|
|
assert_eq!(snapshot.roles.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn standard_subscription_builders_preserve_exact_methods_and_subscription_id() {
|
|
let client = match crate::WsClient::new(endpoint(true)) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client creation failed: {error}"),
|
|
};
|
|
let subscription = match crate::standard_ws_subscription("slotSubscribe") {
|
|
std::option::Option::Some(value) => value,
|
|
std::option::Option::None => panic!("slot subscription missing"),
|
|
};
|
|
let subscribe = client.build_standard_subscribe_request(
|
|
subscription,
|
|
std::vec![serde_json::json!({"commitment": "confirmed"})],
|
|
);
|
|
let unsubscribe = client.build_standard_unsubscribe_request(subscription, 42);
|
|
assert_eq!(subscribe.method, "slotSubscribe");
|
|
assert_eq!(subscribe.params.len(), 1);
|
|
assert_eq!(unsubscribe.method, "slotUnsubscribe");
|
|
assert_eq!(unsubscribe.params, std::vec![serde_json::Value::from(42)]);
|
|
assert_eq!(subscribe.id, serde_json::Value::from(1));
|
|
assert_eq!(unsubscribe.id, serde_json::Value::from(2));
|
|
}
|
|
|
|
#[test]
|
|
fn build_json_rpc_request_increments_ids() {
|
|
let client = match crate::WsClient::new(endpoint(true)) {
|
|
std::result::Result::Ok(client) => client,
|
|
std::result::Result::Err(error) => panic!("client creation failed: {error}"),
|
|
};
|
|
let first =
|
|
client.build_json_rpc_request("slotSubscribe".to_string(), std::vec::Vec::new());
|
|
let second =
|
|
client.build_json_rpc_request("rootSubscribe".to_string(), std::vec::Vec::new());
|
|
assert_eq!(first.id, serde_json::Value::from(1));
|
|
assert_eq!(second.id, serde_json::Value::from(2));
|
|
}
|
|
}
|