use bytes::Bytes;
use helix_core::effect::TransportId;
use helix_core::ports::FrameSender;
use helix_core::PortError;
use helix_core::Tick;
use helix_driver_host::{
AsyncMetricSink, HostHeaderRegistry, HostNetworkConfig, NoopMetricSink, SharedWsClient,
TickIngressSender, WsConnectionActivation,
};
use std::sync::Arc;
use tokio::sync::mpsc;
mod ingress;
use ingress::NativeTickSender;
pub struct NativeTransport {
inner: Option<SharedWsClient>,
init_error: Option<String>,
handshake_headers: Vec<(String, String)>,
reconnect: Option<NativeReconnectFactory>,
}
#[derive(Clone)]
pub(crate) struct NativeReconnectFactory {
config: HostNetworkConfig,
headers: HostHeaderRegistry,
id: TransportId,
tick_tx: Option<NativeTickSender>,
metrics: Arc<dyn AsyncMetricSink>,
}
impl NativeReconnectFactory {
pub(crate) fn build(&self) -> Result<NativeTransport, PortError> {
let mut inner = SharedWsClient::with_registry(self.config.clone(), self.headers.clone())?
.with_metric_sink(Arc::clone(&self.metrics));
if let Some(tx) = &self.tick_tx {
inner = tx.attach(inner, self.id);
}
Ok(NativeTransport::from_shared_parts(
inner,
Some(self.clone()),
))
}
}
impl NativeTransport {
pub fn new(
url: impl Into<String>,
id: TransportId,
tick_tx: Option<mpsc::Sender<Tick>>,
) -> Self {
let config = HostNetworkConfig::new("", url.into());
match Self::with_config(config, id, tick_tx) {
Ok(transport) => transport,
Err(e) => Self {
inner: None,
init_error: Some(e.to_string()),
handshake_headers: Vec::new(),
reconnect: None,
},
}
}
pub fn new_stamped(
url: impl Into<String>,
id: TransportId,
tick_tx: Option<TickIngressSender>,
) -> Self {
let config = HostNetworkConfig::new("", url.into());
match Self::with_config_parts(config, id, tick_tx.map(NativeTickSender::Stamped)) {
Ok(transport) => transport,
Err(error) => Self {
inner: None,
init_error: Some(error.to_string()),
handshake_headers: Vec::new(),
reconnect: None,
},
}
}
pub fn with_handshake_headers(mut self, headers: Vec<(String, String)>) -> Self {
self.handshake_headers = headers;
self
}
pub fn with_metric_sink(mut self, metrics: Arc<dyn AsyncMetricSink>) -> Self {
if let Some(inner) = self.inner.take() {
self.inner = Some(inner.with_metric_sink(Arc::clone(&metrics)));
}
if let Some(factory) = &mut self.reconnect {
factory.metrics = metrics;
}
self
}
pub fn with_config(
config: HostNetworkConfig,
id: TransportId,
tick_tx: Option<mpsc::Sender<Tick>>,
) -> Result<Self, PortError> {
Self::with_config_parts(config, id, tick_tx.map(NativeTickSender::Raw))
}
fn with_config_parts(
config: HostNetworkConfig,
id: TransportId,
tick_tx: Option<NativeTickSender>,
) -> Result<Self, PortError> {
let headers = HostHeaderRegistry::default();
let mut inner = SharedWsClient::with_registry(config.clone(), headers.clone())?;
if let Some(tx) = &tick_tx {
inner = tx.attach(inner, id);
}
let reconnect = NativeReconnectFactory {
config,
headers,
id,
tick_tx,
metrics: Arc::new(NoopMetricSink),
};
Ok(Self::from_shared_parts(inner, Some(reconnect)))
}
pub fn from_shared(
inner: SharedWsClient,
id: TransportId,
tick_tx: Option<mpsc::Sender<Tick>>,
) -> Self {
let inner = match tick_tx {
Some(tx) => inner.with_inbound_tick(id, tx),
None => inner,
};
Self::from_shared_parts(inner, None)
}
fn from_shared_parts(inner: SharedWsClient, reconnect: Option<NativeReconnectFactory>) -> Self {
Self {
inner: Some(inner),
init_error: None,
handshake_headers: Vec::new(),
reconnect,
}
}
pub(crate) fn reconnect_factory(&self) -> Option<NativeReconnectFactory> {
self.reconnect.clone()
}
fn inner_mut(&mut self) -> Result<&mut SharedWsClient, PortError> {
match self.inner.as_mut() {
Some(inner) => Ok(inner),
None => Err(PortError::Transport(
self.init_error
.clone()
.unwrap_or_else(|| "transport not initialized".to_string()),
)),
}
}
fn inner(&self) -> Result<&SharedWsClient, PortError> {
match self.inner.as_ref() {
Some(inner) => Ok(inner),
None => Err(PortError::Transport(
self.init_error
.clone()
.unwrap_or_else(|| "transport not initialized".to_string()),
)),
}
}
}
impl NativeTransport {
pub async fn connect(&mut self) -> Result<WsConnectionActivation, PortError> {
if !self.handshake_headers.is_empty() {
let registry = self.inner_mut()?.headers();
registry
.replace_headers(self.handshake_headers.iter().map(|(k, v)| (k, v)))
.await?;
}
self.inner_mut()?.connect().await
}
async fn send_frame(&self, frame: Bytes) -> Result<(), PortError> {
self.inner()?.send(frame).await
}
pub async fn recv(&self) -> Result<Option<Bytes>, PortError> {
self.inner()?.recv().await
}
pub async fn close(&self) -> Result<(), PortError> {
self.inner()?.close().await
}
}
#[async_trait::async_trait]
impl FrameSender for NativeTransport {
async fn send(&self, frame: Bytes) -> Result<(), PortError> {
self.send_frame(frame).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use helix_core::effect::TransportId;
#[tokio::test]
async fn test_send_on_disconnected_returns_err() {
let (tx, _rx) = mpsc::channel(8);
let transport =
NativeTransport::new("ws://127.0.0.1:0/never", TransportId::from_raw(1), Some(tx));
let result = transport.send(Bytes::from_static(b"hello")).await;
assert!(matches!(result, Err(PortError::Transport(_))));
}
#[tokio::test]
async fn test_recv_on_disconnected_returns_none() {
let (tx, _rx) = mpsc::channel(8);
let transport =
NativeTransport::new("ws://127.0.0.1:0/never", TransportId::from_raw(2), Some(tx));
let result = transport.recv().await.expect("recv should not error");
assert!(result.is_none(), "disconnected recv should return None");
}
#[tokio::test]
async fn test_close_on_disconnected_is_ok() {
let (tx, _rx) = mpsc::channel(8);
let transport =
NativeTransport::new("ws://127.0.0.1:0/never", TransportId::from_raw(3), Some(tx));
transport.close().await.ok();
}
}