use http::Uri;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio_websockets::ClientBuilder;
use crate::Config;
use crate::actor::{Actor, ActorContext, Addr};
use crate::adapters::ws_conn::WsConn;
use crate::message::Message;
use crate::tokio_time::sleep;
use async_trait::async_trait;
use log::{debug, info};
use web_time::Duration;
#[derive(Clone)]
pub struct OutgoingWebsocketManager {
config: Config,
clients: Arc<RwLock<HashMap<String, Addr>>>,
urls: Vec<String>,
}
impl OutgoingWebsocketManager {
pub fn new(config: Config, urls: Vec<String>) -> Self {
OutgoingWebsocketManager {
urls,
clients: Arc::new(RwLock::new(HashMap::new())),
config,
}
}
pub async fn connected_count(&self) -> usize {
self.clients.read().await.len()
}
pub fn urls(&self) -> &[String] {
&self.urls
}
}
#[async_trait]
impl Actor for OutgoingWebsocketManager {
async fn pre_start(&mut self, ctx: &ActorContext) {
info!("OutgoingWebsocketManager starting");
for url in self.urls.iter() {
loop {
if self.clients.read().await.contains_key(url) {
debug!("already connected to {}", url);
break;
}
debug!("attempting WebSocket connect to {}", url);
let uri = match url.parse::<Uri>() {
Ok(u) => u,
Err(_) => {
debug!("invalid URL: {}", url);
sleep(Duration::from_millis(200)).await;
continue;
}
};
let host = uri.host().unwrap_or("127.0.0.1");
let port = uri
.port_u16()
.unwrap_or(if uri.scheme_str() == Some("wss") {
443
} else {
80
});
let tcp = tokio::net::TcpStream::connect((host, port)).await;
let result = match tcp {
Ok(stream) => {
let ws_config =
tokio_websockets::Config::default().flush_threshold(usize::MAX);
ClientBuilder::from_uri(uri)
.config(ws_config)
.connect_on(stream)
.await
}
Err(e) => {
debug!("TCP connect to {}:{} failed: {}", host, port, e);
sleep(Duration::from_millis(200)).await;
continue;
}
};
if let Ok((socket, _)) = result {
let client = WsConn::new(socket, self.config.allow_public_space);
let addr = ctx.start_actor(Box::new(client));
self.clients.write().await.insert(url.clone(), addr);
debug!("connected to {}", url);
break;
}
debug!("connect to {} failed, retrying in 200ms", url);
sleep(Duration::from_millis(200)).await;
}
}
}
fn subscribe_to_everything(&self) -> bool {
true
}
async fn handle(&mut self, message: Arc<Message>, _ctx: &ActorContext) {
let snapshot: Vec<Addr> = self.clients.read().await.values().cloned().collect();
for client in snapshot {
let _ = client.send(Arc::clone(&message));
}
}
async fn stopping(&mut self, _ctx: &ActorContext) {
let count = self.clients.read().await.len();
info!(
"OutgoingWebsocketManager stopping — {} outgoing connections",
count
);
self.clients.write().await.clear();
}
}