use std::time::Duration;
use crate::{
backoff::{HermesExponentialBackoff, HermesExponentialBackoffBuilder},
resilient_ws_connection::HermesResilientWSConnection,
ws_connection::{
HermesClientMessage, HermesClientMessageSubscribe, HermesClientMessageUnsubscribe,
HermesServerMessage,
},
CHANNEL_CAPACITY,
};
use anyhow::{bail, Result};
use backoff::ExponentialBackoff;
use tokio::sync::mpsc::{self, error::TrySendError};
use tracing::{error, warn};
use ttl_cache::TtlCache;
use url::Url;
const DEDUP_CACHE_SIZE: usize = 100_000;
const DEDUP_TTL: Duration = Duration::from_secs(10);
const DEFAULT_ENDPOINTS: [&str; 1] = ["wss://hermes.pyth.network/ws"];
const DEFAULT_NUM_CONNECTIONS: usize = 3;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
pub struct HermesClient {
endpoints: Vec<Url>,
num_connections: usize,
ws_connections: Vec<HermesResilientWSConnection>,
backoff: ExponentialBackoff,
timeout: Duration,
channel_capacity: usize,
access_token: Option<String>,
}
impl HermesClient {
pub fn new(
endpoints: Vec<Url>,
num_connections: usize,
backoff: HermesExponentialBackoff,
timeout: Duration,
channel_capacity: usize,
access_token: Option<String>,
) -> Result<Self> {
if endpoints.is_empty() {
bail!("At least one endpoint must be provided");
}
Ok(Self {
endpoints,
num_connections,
ws_connections: Vec::with_capacity(num_connections),
backoff: backoff.into(),
timeout,
channel_capacity,
access_token,
})
}
pub async fn start(&mut self) -> Result<mpsc::Receiver<HermesServerMessage>> {
let (sender, receiver) = mpsc::channel::<HermesServerMessage>(self.channel_capacity);
let (ws_connection_sender, mut ws_connection_receiver) =
mpsc::channel::<HermesServerMessage>(CHANNEL_CAPACITY);
for i in 0..self.num_connections {
let mut endpoint = self.endpoints[i % self.endpoints.len()].clone();
if let Some(ref token) = self.access_token {
endpoint
.query_pairs_mut()
.append_pair("ACCESS_TOKEN", token);
}
let connection = HermesResilientWSConnection::new(
endpoint,
self.backoff.clone(),
self.timeout,
ws_connection_sender.clone(),
);
self.ws_connections.push(connection);
}
let mut seen_updates = TtlCache::new(DEDUP_CACHE_SIZE);
tokio::spawn(async move {
while let Some(response) = ws_connection_receiver.recv().await {
let cache_key = response.cache_key();
if seen_updates.contains_key(&cache_key) {
continue;
}
seen_updates.insert(cache_key, response.clone(), DEDUP_TTL);
match sender.try_send(response) {
Ok(_) => (),
Err(TrySendError::Full(r)) => {
warn!("Sender channel is full, responses will be delayed");
if sender.send(r).await.is_err() {
error!("Sender channel is closed, stopping client");
}
}
Err(TrySendError::Closed(_)) => {
error!("Sender channel is closed, stopping client");
}
}
}
});
Ok(receiver)
}
pub async fn subscribe(
&mut self,
subscribe_request: HermesClientMessageSubscribe,
) -> Result<()> {
for connection in &mut self.ws_connections {
connection
.send_request(HermesClientMessage::Subscribe(subscribe_request.clone()))
.await?;
}
Ok(())
}
pub async fn unsubscribe(
&mut self,
unsubscribe_request: HermesClientMessageUnsubscribe,
) -> Result<()> {
for connection in &mut self.ws_connections {
connection
.send_request(HermesClientMessage::Unsubscribe(
unsubscribe_request.clone(),
))
.await?;
}
Ok(())
}
}
pub struct HermesClientBuilder {
endpoints: Vec<Url>,
num_connections: usize,
backoff: HermesExponentialBackoff,
timeout: Duration,
channel_capacity: usize,
access_token: Option<String>,
}
impl Default for HermesClientBuilder {
fn default() -> Self {
Self {
endpoints: DEFAULT_ENDPOINTS
.iter()
.map(|&s| s.parse().unwrap())
.collect(),
num_connections: DEFAULT_NUM_CONNECTIONS,
backoff: HermesExponentialBackoffBuilder::default().build(),
timeout: DEFAULT_TIMEOUT,
channel_capacity: CHANNEL_CAPACITY,
access_token: None,
}
}
}
impl HermesClientBuilder {
pub fn with_endpoints(mut self, endpoints: Vec<Url>) -> Self {
self.endpoints = endpoints;
self
}
pub fn with_num_connections(mut self, num_connections: usize) -> Self {
self.num_connections = num_connections;
self
}
pub fn with_backoff(mut self, backoff: HermesExponentialBackoff) -> Self {
self.backoff = backoff;
self
}
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
pub fn with_channel_capacity(mut self, channel_capacity: usize) -> Self {
self.channel_capacity = channel_capacity;
self
}
pub fn with_access_token(mut self, access_token: String) -> Self {
self.access_token = Some(access_token);
self
}
pub fn build(self) -> Result<HermesClient> {
HermesClient::new(
self.endpoints,
self.num_connections,
self.backoff,
self.timeout,
self.channel_capacity,
self.access_token,
)
}
}