Skip to content

Commit

Permalink
feat: expose connector functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@dodomorandi
dodomorandi committed Apr 24, 2024
1 parent 645d3fd commit 082d083
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 53 deletions.
114 changes: 107 additions & 7 deletions async-nats/src/connector.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use crate::tls;
use crate::AuthError;
use crate::ClientOp;
use crate::ConnectInfo;
use crate::ConnectOptions;
use crate::Event;
use crate::MaybeArc;
use crate::Protocol;
Expand All @@ -45,10 +46,11 @@ use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::time::sleep;
use tokio_rustls::rustls;

pub(crate) struct ConnectorOptions {
pub struct ConnectorOptions {
pub(crate) tls_required: bool,
pub(crate) certificates: Vec<PathBuf>,
pub(crate) client_cert: Option<PathBuf>,
Expand All @@ -68,7 +70,7 @@ pub(crate) struct ConnectorOptions {
}

/// Maintains a list of servers and establishes connections.
pub(crate) struct Connector {
pub(crate) struct Handler<const CLIENT_NODE: bool = true> {
/// A map of servers and number of connect attempts.
servers: Vec<(ServerAddr, usize)>,
options: ConnectorOptions,
Expand All @@ -88,17 +90,110 @@ pub(crate) fn reconnect_delay_callback_default(attempts: usize) -> Duration {
}
}

impl Connector {
pub(crate) fn new<A: ToServerAddrs>(
pub struct Connector<const CLIENT_NODE: bool = true> {
pub(crate) handler: Handler<CLIENT_NODE>,
pub(crate) events_rx: tokio::sync::mpsc::Receiver<Event>,
pub(crate) state_rx: tokio::sync::watch::Receiver<State>,
pub(crate) subscription_capacity: usize,
pub(crate) event_callback: Option<CallbackArg1<Event, ()>>,
pub(crate) inbox_prefix: String,
pub(crate) request_timeout: Option<Duration>,
pub(crate) retry_on_initial_connect: bool,
pub(crate) sender_capacity: usize,
pub(crate) ping_interval: Duration,
}

pub fn create<A: ToServerAddrs>(addrs: A, options: ConnectOptions) -> Result<Connector, io::Error> {
create_inner::<A, true>(addrs, options)
}

pub fn create_leaf_connector<A: ToServerAddrs>(
addrs: A,
options: ConnectOptions,
) -> Result<Connector<false>, io::Error> {
create_inner::<A, false>(addrs, options)
}

fn create_inner<A: ToServerAddrs, const CLIENT_NODE: bool>(
addrs: A,
options: ConnectOptions,
) -> Result<Connector<CLIENT_NODE>, io::Error> {
let ConnectOptions {
name,
no_echo,
max_reconnects,
connection_timeout,
auth,
tls_required,
tls_first,
certificates,
client_cert,
client_key,
tls_client_config,
ping_interval,
subscription_capacity,
sender_capacity,
event_callback,
inbox_prefix,
request_timeout,
retry_on_initial_connect,
ignore_discovered_servers,
retain_servers_order,
read_buffer_capacity,
reconnect_delay_callback,
auth_callback,
} = options;

let options = ConnectorOptions {
tls_required,
certificates,
client_cert,
client_key,
tls_client_config,
tls_first,
auth,
no_echo,
connection_timeout,
name,
ignore_discovered_servers,
retain_servers_order,
read_buffer_capacity,
reconnect_delay_callback,
auth_callback,
max_reconnects,
};

let (events_tx, events_rx) = mpsc::channel(128);
let (state_tx, state_rx) = tokio::sync::watch::channel(State::Pending);
// We're setting it to the default server payload size.
let max_payload = Arc::new(AtomicUsize::new(1024 * 1024));

let handler = Handler::<CLIENT_NODE>::new(addrs, options, events_tx, state_tx, max_payload)?;
Ok(Connector {
handler,
events_rx,
state_rx,
subscription_capacity,
inbox_prefix,
request_timeout,
event_callback,
retry_on_initial_connect,
sender_capacity,
ping_interval,
})
}

impl<const CLIENT_NODE: bool> Handler<CLIENT_NODE> {
fn new<A: ToServerAddrs>(
addrs: A,
options: ConnectorOptions,
events_tx: tokio::sync::mpsc::Sender<Event>,
state_tx: tokio::sync::watch::Sender<State>,
max_payload: Arc<AtomicUsize>,
) -> Result<Connector, io::Error> {
) -> Result<Self, io::Error> {
let servers = addrs.to_server_addrs()?.map(|addr| (addr, 0)).collect();

Ok(Connector {
Ok(Handler {
attempts: 0,
servers,
options,
Expand Down Expand Up @@ -194,12 +289,17 @@ impl Connector {
}

let tls_required = self.options.tls_required || server_addr.tls_required();
let lang = if CLIENT_NODE {
LANG.to_string()
} else {
String::new()
};
let mut connect_info = ConnectInfo {
tls_required,
name: self.options.name.clone(),
pedantic: false,
verbose: false,
lang: LANG.to_string(),
lang,
version: VERSION.to_string(),
protocol: Protocol::Dynamic,
user: self.options.auth.username.clone(),
Expand Down
72 changes: 26 additions & 46 deletions async-nats/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,7 @@
#![deny(rustdoc::invalid_rust_codeblocks)]
#![cfg_attr(docsrs, feature(doc_auto_cfg))]

use connector::Connector;
use thiserror::Error;

use futures::stream::Stream;
Expand All @@ -210,7 +211,6 @@ use std::option;
use std::pin::Pin;
use std::slice;
use std::str::{self, FromStr};
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::io::ErrorKind;
Expand All @@ -237,7 +237,6 @@ const MULTIPLEXER_SID: u64 = 0;
pub use tokio_rustls::rustls;

use connection::{Connection, State};
use connector::{Connector, ConnectorOptions};
pub use header::{HeaderMap, HeaderName, HeaderValue};
pub use subject::Subject;

Expand Down Expand Up @@ -435,7 +434,7 @@ struct Multiplexer {
/// A connection handler which facilitates communication from channels to a single shared connection.
pub(crate) struct ConnectionHandler {
connection: Connection,
connector: Connector,
connector: connector::Handler,
subscriptions: HashMap<u64, Subscription>,
multiplexer: Option<Multiplexer>,
pending_pings: usize,
Expand All @@ -449,7 +448,7 @@ pub(crate) struct ConnectionHandler {
impl ConnectionHandler {
pub(crate) fn new(
connection: Connection,
connector: Connector,
connector: connector::Handler,
info_sender: tokio::sync::watch::Sender<ServerInfo>,
ping_period: Duration,
) -> ConnectionHandler {
Expand Down Expand Up @@ -957,63 +956,44 @@ pub async fn connect_with_options<A: ToServerAddrs>(
addrs: A,
options: ConnectOptions,
) -> Result<Client, ConnectError> {
let ping_period = options.ping_interval;

let (events_tx, mut events_rx) = mpsc::channel(128);
let (state_tx, state_rx) = tokio::sync::watch::channel(State::Pending);
// We're setting it to the default server payload size.
let max_payload = Arc::new(AtomicUsize::new(1024 * 1024));

let mut connector = Connector::new(
addrs,
ConnectorOptions {
tls_required: options.tls_required,
certificates: options.certificates,
client_key: options.client_key,
client_cert: options.client_cert,
tls_client_config: options.tls_client_config,
tls_first: options.tls_first,
auth: options.auth,
no_echo: options.no_echo,
connection_timeout: options.connection_timeout,
name: options.name,
ignore_discovered_servers: options.ignore_discovered_servers,
retain_servers_order: options.retain_servers_order,
read_buffer_capacity: options.read_buffer_capacity,
reconnect_delay_callback: options.reconnect_delay_callback,
auth_callback: options.auth_callback,
max_reconnects: options.max_reconnects,
},
events_tx,
state_tx,
max_payload.clone(),
)
.map_err(ConnectError::InvalidServerAddress)?;
let Connector {
mut handler,
mut events_rx,
state_rx,
subscription_capacity,
inbox_prefix,
request_timeout,
event_callback,
retry_on_initial_connect,
sender_capacity,
ping_interval,
} = connector::create(addrs, options).map_err(ConnectError::InvalidServerAddress)?;

let mut info: ServerInfo = Default::default();
let mut connection = None;
if !options.retry_on_initial_connect {
if !retry_on_initial_connect {
debug!("retry on initial connect failure is disabled");
let (info_ok, connection_ok) = connector.try_connect().await?;
let (info_ok, connection_ok) = handler.try_connect().await?;
connection = Some(connection_ok);
info = info_ok;
}

let max_payload = Arc::clone(&handler.max_payload);
let (info_sender, info_watcher) = tokio::sync::watch::channel(info.clone());
let (sender, mut receiver) = mpsc::channel(options.sender_capacity);
let (sender, mut receiver) = mpsc::channel(sender_capacity);

let events_handler_task = task::spawn(async move {
while let Some(event) = events_rx.recv().await {
tracing::info!("event: {}", event);
if let Some(event_callback) = &options.event_callback {
if let Some(event_callback) = &event_callback {
event_callback.call(event).await;
}
}
});

let connection_handler_task = task::spawn(async move {
if connection.is_none() && options.retry_on_initial_connect {
let (info, connection_ok) = match connector.connect().await {
if connection.is_none() && retry_on_initial_connect {
let (info, connection_ok) = match handler.connect().await {
Ok((info, connection)) => (info, connection),
Err(err) => {
error!("connection closed: {}", err);
Expand All @@ -1025,17 +1005,17 @@ pub async fn connect_with_options<A: ToServerAddrs>(
}
let connection = connection.unwrap();
let mut connection_handler =
ConnectionHandler::new(connection, connector, info_sender, ping_period);
ConnectionHandler::new(connection, handler, info_sender, ping_interval);
connection_handler.process(&mut receiver).await
});

Ok(client::Builder {
info: info_watcher,
state: state_rx,
sender,
capacity: options.subscription_capacity,
inbox_prefix: options.inbox_prefix,
request_timeout: options.request_timeout,
capacity: subscription_capacity,
inbox_prefix,
request_timeout,
max_payload,
events_handler_task,
connection_handler_task,
Expand Down

0 comments on commit 082d083

Please sign in to comment.