|
@@ -1,30 +1,62 @@
|
|
|
//! Relayers
|
|
|
//!
|
|
|
//! This is the main entry point to the client library.
|
|
|
-use crate::{client::ActiveSubscription, Client, Error};
|
|
|
+use crate::{client::ActiveSubscription as ClientActiveSubscription, Client, Error};
|
|
|
use futures::future::join_all;
|
|
|
use nostr_rs_types::{
|
|
|
client::{self, subscribe},
|
|
|
- types::SubscriptionId,
|
|
|
Response,
|
|
|
};
|
|
|
-use std::{collections::HashMap, sync::Arc};
|
|
|
+use std::{
|
|
|
+ collections::HashMap,
|
|
|
+ sync::{
|
|
|
+ atomic::{AtomicUsize, Ordering},
|
|
|
+ Arc,
|
|
|
+ },
|
|
|
+};
|
|
|
use tokio::sync::{mpsc, RwLock};
|
|
|
use url::Url;
|
|
|
|
|
|
-type Subscriptions =
|
|
|
- Arc<RwLock<HashMap<SubscriptionId, (subscribe::Subscribe, Vec<ActiveSubscription>)>>>;
|
|
|
+pub mod subscription;
|
|
|
+
|
|
|
+pub(crate) type AllClients =
|
|
|
+ Arc<RwLock<HashMap<Url, (Arc<AtomicUsize>, (ClientActiveSubscription, Client))>>>;
|
|
|
+
|
|
|
/// Clients
|
|
|
///
|
|
|
/// This is a set of outgoing connections to relayers. This struct can connect
|
|
|
/// async to N relayers offering a simple API to talk to all of them at the same
|
|
|
/// time, and to receive messages
|
|
|
-#[derive(Debug)]
|
|
|
pub struct Pool {
|
|
|
- clients: RwLock<HashMap<Url, Client>>,
|
|
|
+ clients: AllClients,
|
|
|
sender: mpsc::Sender<(Response, Url)>,
|
|
|
receiver: Option<mpsc::Receiver<(Response, Url)>>,
|
|
|
- subscriptions: Subscriptions,
|
|
|
+ subscription_manager: Arc<subscription::Manager>,
|
|
|
+}
|
|
|
+
|
|
|
+/// Active client
|
|
|
+///
|
|
|
+/// For each connection on the pool this object will be returned. When dropped,
|
|
|
+/// that connection is also dropped from the connection pool.
|
|
|
+pub struct ActiveClient {
|
|
|
+ client_id: Url,
|
|
|
+ counter: Arc<AtomicUsize>,
|
|
|
+ all_clients: AllClients,
|
|
|
+}
|
|
|
+
|
|
|
+impl Drop for ActiveClient {
|
|
|
+ fn drop(&mut self) {
|
|
|
+ let counter = self.counter.fetch_sub(1, Ordering::SeqCst);
|
|
|
+ if counter == 0 {
|
|
|
+ let all_clients = self.all_clients.clone();
|
|
|
+ let client_id = self.client_id.clone();
|
|
|
+ tokio::spawn(async move {
|
|
|
+ // remove the client from the pool, when it goes out of scope
|
|
|
+ // it will be disconnected
|
|
|
+ all_clients.write().await.remove(&client_id);
|
|
|
+ });
|
|
|
+ }
|
|
|
+ }
|
|
|
}
|
|
|
|
|
|
/// Default channel buffer size for the pool
|
|
@@ -36,44 +68,25 @@ impl Default for Pool {
|
|
|
Self {
|
|
|
clients: Default::default(),
|
|
|
receiver: Some(receiver),
|
|
|
- subscriptions: Default::default(),
|
|
|
+ subscription_manager: Default::default(),
|
|
|
sender,
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
-/// Return a subscription that will be removed when dropped
|
|
|
-#[derive(Debug)]
|
|
|
-pub struct PoolSubscription {
|
|
|
- subscription_id: SubscriptionId,
|
|
|
- subscriptions: Subscriptions,
|
|
|
-}
|
|
|
-
|
|
|
-impl Drop for PoolSubscription {
|
|
|
- fn drop(&mut self) {
|
|
|
- let subscriptions = self.subscriptions.clone();
|
|
|
- let subscription_id = self.subscription_id.clone();
|
|
|
- tokio::spawn(async move {
|
|
|
- subscriptions.write().await.remove(&subscription_id);
|
|
|
- });
|
|
|
- }
|
|
|
-}
|
|
|
-
|
|
|
impl Pool {
|
|
|
/// Creates a new instance with a list of urls
|
|
|
- pub fn new_with_clients(clients: Vec<Url>) -> Self {
|
|
|
- let (sender, receiver) = mpsc::channel(DEFAULT_CHANNEL_BUFFER_SIZE);
|
|
|
- let clients = clients
|
|
|
- .into_iter()
|
|
|
- .map(|url| (url.clone(), Client::new(sender.clone(), url)))
|
|
|
- .collect::<HashMap<_, _>>();
|
|
|
-
|
|
|
- Self {
|
|
|
- clients: RwLock::new(clients),
|
|
|
- subscriptions: Default::default(),
|
|
|
- receiver: Some(receiver),
|
|
|
- sender,
|
|
|
- }
|
|
|
+ pub fn new_with_clients(clients: Vec<Url>) -> Result<(Self, Vec<ActiveClient>), Error> {
|
|
|
+ let pool = Self::default();
|
|
|
+ let connect = clients.into_iter().map(|url| pool.connect_to(url));
|
|
|
+
|
|
|
+ futures::executor::block_on(async {
|
|
|
+ futures::future::join_all(connect)
|
|
|
+ .await
|
|
|
+ .into_iter()
|
|
|
+ .collect::<Result<Vec<_>, _>>()
|
|
|
+ })
|
|
|
+ .map(|clients| (pool, clients))
|
|
|
}
|
|
|
|
|
|
/// Splits the pool removing the receiver to be used in a different context
|
|
@@ -93,38 +106,15 @@ impl Pool {
|
|
|
|
|
|
/// Returns the number of active subscriptions
|
|
|
pub async fn active_subscriptions(&self) -> usize {
|
|
|
- self.subscriptions.read().await.keys().len()
|
|
|
+ self.subscription_manager.total_subscribers()
|
|
|
}
|
|
|
|
|
|
/// Subscribes to all the connected relayers
|
|
|
pub async fn subscribe(
|
|
|
&self,
|
|
|
subscription: subscribe::Subscribe,
|
|
|
- ) -> Result<PoolSubscription, Error> {
|
|
|
- let clients = self.clients.read().await;
|
|
|
-
|
|
|
- let wait_all = clients
|
|
|
- .values()
|
|
|
- .map(|sender| sender.subscribe(subscription.clone()))
|
|
|
- .collect::<Vec<_>>();
|
|
|
-
|
|
|
- let subscription_id = subscription.subscription_id.clone();
|
|
|
-
|
|
|
- self.subscriptions.write().await.insert(
|
|
|
- subscription_id.clone(),
|
|
|
- (
|
|
|
- subscription,
|
|
|
- join_all(wait_all)
|
|
|
- .await
|
|
|
- .into_iter()
|
|
|
- .collect::<Result<Vec<_>, _>>()?,
|
|
|
- ),
|
|
|
- );
|
|
|
-
|
|
|
- Ok(PoolSubscription {
|
|
|
- subscription_id,
|
|
|
- subscriptions: self.subscriptions.clone(),
|
|
|
- })
|
|
|
+ ) -> subscription::ActiveSubscription {
|
|
|
+ self.subscription_manager.subcribe(subscription, None).await
|
|
|
}
|
|
|
|
|
|
/// Sends a request to all the connected relayers
|
|
@@ -133,7 +123,7 @@ impl Pool {
|
|
|
join_all(
|
|
|
clients
|
|
|
.values()
|
|
|
- .map(|sender| sender.post(request.clone()))
|
|
|
+ .map(|(_, (_, sender))| sender.post(request.clone()))
|
|
|
.collect::<Vec<_>>(),
|
|
|
)
|
|
|
.await;
|
|
@@ -145,7 +135,7 @@ impl Pool {
|
|
|
.read()
|
|
|
.await
|
|
|
.iter()
|
|
|
- .filter(|(_, client)| client.is_connected())
|
|
|
+ .filter(|(_, (_, (_, client)))| client.is_connected())
|
|
|
.collect::<Vec<_>>()
|
|
|
.len()
|
|
|
}
|
|
@@ -154,22 +144,43 @@ impl Pool {
|
|
|
///
|
|
|
/// This function will open a connection at most once, if a connection
|
|
|
/// already exists false will be returned
|
|
|
- pub async fn connect_to(&self, url: Url) {
|
|
|
+ pub async fn connect_to(&self, url: Url) -> Result<ActiveClient, Error> {
|
|
|
let mut clients = self.clients.write().await;
|
|
|
- let mut subscriptions = self.subscriptions.write().await;
|
|
|
|
|
|
- if !clients.contains_key(&url) {
|
|
|
+ let ref_id = if let Some((id, _)) = clients.get(&url) {
|
|
|
+ id.fetch_add(1, Ordering::SeqCst);
|
|
|
+ id.clone()
|
|
|
+ } else {
|
|
|
log::warn!("Connecting to {}", url);
|
|
|
- let client = Client::new(self.sender.clone(), url.clone());
|
|
|
-
|
|
|
- for (filter, sub) in subscriptions.values_mut() {
|
|
|
- let _ = client.subscribe(filter.clone()).await.map(|subscription| {
|
|
|
- sub.push(subscription);
|
|
|
- });
|
|
|
- }
|
|
|
-
|
|
|
- clients.insert(url, client);
|
|
|
- }
|
|
|
+ let subscription_manager = self.subscription_manager.clone();
|
|
|
+ let client = Client::new(
|
|
|
+ self.sender.clone(),
|
|
|
+ url.clone(),
|
|
|
+ move |response, url, return_to| {
|
|
|
+ let subscription_manager = subscription_manager.clone();
|
|
|
+ Box::pin(async move {
|
|
|
+ subscription_manager
|
|
|
+ .process_message(response, url, return_to)
|
|
|
+ .await
|
|
|
+ })
|
|
|
+ },
|
|
|
+ );
|
|
|
+
|
|
|
+ // subscribe to all events
|
|
|
+ let meta_subscription = client
|
|
|
+ .subscribe(subscribe::Subscribe::to_all_events())
|
|
|
+ .await?;
|
|
|
+
|
|
|
+ let ref_id: Arc<AtomicUsize> = Arc::new(1.into());
|
|
|
+ clients.insert(url.clone(), (ref_id.clone(), (meta_subscription, client)));
|
|
|
+ ref_id
|
|
|
+ };
|
|
|
+
|
|
|
+ Ok(ActiveClient {
|
|
|
+ client_id: url,
|
|
|
+ counter: ref_id,
|
|
|
+ all_clients: self.clients.clone(),
|
|
|
+ })
|
|
|
}
|
|
|
}
|
|
|
|
|
@@ -199,10 +210,7 @@ mod test {
|
|
|
#[tokio::test]
|
|
|
async fn droppable_subscription() {
|
|
|
let client_pool = Pool::default();
|
|
|
- let subscription = client_pool
|
|
|
- .subscribe(Default::default())
|
|
|
- .await
|
|
|
- .expect("valid subscription");
|
|
|
+ let subscription = client_pool.subscribe(Default::default()).await;
|
|
|
|
|
|
assert_eq!(client_pool.active_subscriptions().await, 1);
|
|
|
drop(subscription);
|
|
@@ -213,7 +221,7 @@ mod test {
|
|
|
#[tokio::test]
|
|
|
async fn connect_to_dummy_server() {
|
|
|
let (addr, stopper) = dummy_server(0).await;
|
|
|
- let client_pool = Pool::new_with_clients(vec![addr]);
|
|
|
+ let (client_pool, _connections) = Pool::new_with_clients(vec![addr]).expect("valid pool");
|
|
|
|
|
|
assert_eq!(0, client_pool.check_active_connections().await);
|
|
|
|
|
@@ -230,13 +238,11 @@ mod test {
|
|
|
#[tokio::test]
|
|
|
async fn two_clients_communication() {
|
|
|
let (addr, _) = dummy_server(0).await;
|
|
|
- let mut client_pool1 = Pool::new_with_clients(vec![addr.clone()]);
|
|
|
- let client_pool2 = Pool::new_with_clients(vec![addr]);
|
|
|
+ let (mut client_pool1, _c1) =
|
|
|
+ Pool::new_with_clients(vec![addr.clone()]).expect("valid pool");
|
|
|
+ let (client_pool2, _c2) = Pool::new_with_clients(vec![addr]).expect("valid pool");
|
|
|
|
|
|
- let _sub1 = client_pool1
|
|
|
- .subscribe(Default::default())
|
|
|
- .await
|
|
|
- .expect("valid subscription");
|
|
|
+ let _sub1 = client_pool1.subscribe(Default::default()).await;
|
|
|
|
|
|
sleep(Duration::from_millis(10)).await;
|
|
|
|
|
@@ -270,13 +276,11 @@ mod test {
|
|
|
#[tokio::test]
|
|
|
async fn reconnect_and_resubscribe() {
|
|
|
let (addr, stopper) = dummy_server(0).await;
|
|
|
- let mut client_pool1 = Pool::new_with_clients(vec![addr.clone()]);
|
|
|
- let client_pool2 = Pool::new_with_clients(vec![addr.clone()]);
|
|
|
+ let (mut client_pool1, _c1) =
|
|
|
+ Pool::new_with_clients(vec![addr.clone()]).expect("valid pool");
|
|
|
+ let (client_pool2, _c2) = Pool::new_with_clients(vec![addr.clone()]).expect("valid pool");
|
|
|
|
|
|
- let _sub1 = client_pool1
|
|
|
- .subscribe(Default::default())
|
|
|
- .await
|
|
|
- .expect("valid subscription");
|
|
|
+ let _sub1 = client_pool1.subscribe(Default::default()).await;
|
|
|
|
|
|
sleep(Duration::from_millis(10)).await;
|
|
|
|
|
@@ -353,13 +357,11 @@ mod test {
|
|
|
async fn connect_multiple_servers() {
|
|
|
let (addr1, _) = dummy_server(0).await;
|
|
|
let (addr2, _) = dummy_server(0).await;
|
|
|
- let mut client_pool1 = Pool::new_with_clients(vec![addr1.clone(), addr2]);
|
|
|
- let client_pool2 = Pool::new_with_clients(vec![addr1]);
|
|
|
+ let (mut client_pool1, _c1) =
|
|
|
+ Pool::new_with_clients(vec![addr1.clone(), addr2]).expect("valid pool");
|
|
|
+ let (client_pool2, _c2) = Pool::new_with_clients(vec![addr1]).expect("valid pool");
|
|
|
|
|
|
- let _sub1 = client_pool1
|
|
|
- .subscribe(Default::default())
|
|
|
- .await
|
|
|
- .expect("valid subscription");
|
|
|
+ let _sub1 = client_pool1.subscribe(Default::default()).await;
|
|
|
|
|
|
sleep(Duration::from_millis(10)).await;
|
|
|
|