refactor backchannel

node/zombienet-backchannel/Cargo.toml

@@ -13,6 +13,7 @@ tokio = { version = "1.0.0", default-features = false, features = ["macros", "ne
 url = "2.0.0"
 tokio-tungstenite = "0.16"
 futures-util = "0.3.18"
+lazy_static = "1.4.0"
 parity-scale-codec = { version = "2.3.1", features = ["derive"] }
 reqwest = "0.11"
 thiserror = "1.0.30"

node/zombienet-backchannel/src/errors.rs

@@ -22,6 +22,12 @@ pub enum BackchannelError {
 	#[error("Error connecting websocket server")]
 	CantConnectToWS,
 
+	#[error("Backchannel not initialized yet")]
+	Uninitialized,
+
+	#[error("Backchannel already initialized")]
+	AlreadyInitialized,
+
 	#[error("Error sending new value to backchannel")]
 	SendItemFail,
 }

node/zombienet-backchannel/src/lib.rs

@@ -20,9 +20,11 @@
 //! implemented as a `backchannel`.
 
 use futures_util::{stream::SplitSink, SinkExt, StreamExt};
+use lazy_static::lazy_static;
 use parity_scale_codec as codec;
 use serde::{Deserialize, Serialize};
 use std::env;
+use std::sync::Mutex;
 use tokio::{net::TcpStream, sync::broadcast};
 use tokio_tungstenite::{
 	connect_async, tungstenite::protocol::Message, MaybeTlsStream, WebSocketStream,
@@ -31,10 +33,14 @@ use tokio_tungstenite::{
 mod errors;
 use errors::BackchannelError;
 
+lazy_static! {
+	pub static ref ZOMBIENET_BACKCHANNEL: Mutex<Option<ZombienetBackchannel>> = Mutex::new(None);
+}
+
 #[derive(Debug)]
 pub struct ZombienetBackchannel {
 	broadcast_tx: broadcast::Sender<BackchannelItem>,
-	ws_tx: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
+	ws_tx: broadcast::Sender<BackchannelItem>,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -43,65 +49,106 @@ pub struct BackchannelItem {
 	value: String,
 }
 
-pub const ZOMBIENET: &str = "🧟ZOMBIENET🧟";
+pub struct Broadcaster;
 
-impl ZombienetBackchannel {
-	pub async fn init() -> Result<ZombienetBackchannel, BackchannelError> {
-		let backchannel_host =
-			env::var("BACKCHANNEL_HOST").unwrap_or_else(|_| "backchannel".to_string());
-		let backchannel_port = env::var("BACKCHANNEL_PORT").unwrap_or_else(|_| "3000".to_string());
-		let ws_url = format!("ws://{}:{}/ws", backchannel_host, backchannel_port);
-		tracing::debug!(target = ZOMBIENET, "Connecting to : {}", &ws_url);
-		let (ws_stream, _) =
-			connect_async(ws_url).await.map_err(|_| BackchannelError::CantConnectToWS)?;
-		let (write, mut read) = ws_stream.split();
-
-		let (tx, _rx) = broadcast::channel(256);
-
-		let tx1 = tx.clone();
-		tokio::spawn(async move {
-			while let Some(Ok(Message::Text(text))) = read.next().await {
-				println!("es {}", &text);
-
-				match serde_json::from_str::<BackchannelItem>(&text) {
-					Ok(backchannel_item) =>
-						if tx1.send(backchannel_item).is_err() {
-							tracing::error!("Error sending through the channel");
-							return
-						},
-					Err(_) => {
-						tracing::error!("Invalid payload received");
-					},
-				}
-			}
-		});
-
-		Ok(ZombienetBackchannel { broadcast_tx: tx, ws_tx: write })
-	}
+pub const ZOMBIENET: &str = "🧟ZOMBIENET🧟";
 
-	pub fn subscribe(&self) -> broadcast::Receiver<BackchannelItem> {
-		self.broadcast_tx.subscribe()
+impl Broadcaster {
+	pub fn subscribe(&self) -> Result<broadcast::Receiver<BackchannelItem>, BackchannelError> {
+		let mut zombienet_bkc = ZOMBIENET_BACKCHANNEL.lock().unwrap();
+		let sender = zombienet_bkc.as_mut().unwrap().broadcast_tx.clone();
+		if zombienet_bkc.is_some() {
+			Ok(sender.subscribe())
+		} else {
+			Err(BackchannelError::Uninitialized)
+		}
 	}
 
 	pub async fn send(
 		&mut self,
 		key: &'static str,
 		val: impl codec::Encode,
 	) -> Result<(), BackchannelError> {
+		let mut zombienet_bkc = ZOMBIENET_BACKCHANNEL.lock().unwrap();
+		if zombienet_bkc.is_none() {
+			return Err(BackchannelError::Uninitialized);
+		}
 		let encoded = val.encode();
 		let backchannel_item = BackchannelItem {
 			key: key.to_string(),
 			value: String::from_utf8_lossy(&encoded).to_string(),
 		};
 
-		self.ws_tx
-			.send(Message::Text(serde_json::to_string(&backchannel_item).unwrap()))
-			.await
-			.map_err(|e| {
-				tracing::error!(target = ZOMBIENET, "Error sending new item: {}", e);
-				BackchannelError::SendItemFail
-			})?;
+		let sender = zombienet_bkc.as_mut().unwrap().ws_tx.clone();
+		sender.send(backchannel_item).map_err(|e| {
+			tracing::error!(target = ZOMBIENET, "Error sending new item: {}", e);
+			BackchannelError::SendItemFail
+		})?;
 
 		Ok(())
 	}
 }
+
+impl ZombienetBackchannel {
+	pub async fn init() -> Result<(), BackchannelError> {
+		let mut zombienet_bkc = ZOMBIENET_BACKCHANNEL.lock().unwrap();
+		if zombienet_bkc.is_none() {
+			let backchannel_host =
+				env::var("BACKCHANNEL_HOST").unwrap_or_else(|_| "backchannel".to_string());
+			let backchannel_port =
+				env::var("BACKCHANNEL_PORT").unwrap_or_else(|_| "3000".to_string());
+			let ws_url = format!("ws://{}:{}/ws", backchannel_host, backchannel_port);
+			tracing::debug!(target = ZOMBIENET, "Connecting to : {}", &ws_url);
+			let (ws_stream, _) =
+				connect_async(ws_url).await.map_err(|_| BackchannelError::CantConnectToWS)?;
+			let (mut write, mut read) = ws_stream.split();
+
+			let (tx, _rx) = broadcast::channel(256);
+			let (tx_relay, mut rx_relay) = broadcast::channel::<BackchannelItem>(256);
+
+			// receive from the ws and send to all subcribers
+			let tx1 = tx.clone();
+			tokio::spawn(async move {
+				while let Some(Ok(Message::Text(text))) = read.next().await {
+					match serde_json::from_str::<BackchannelItem>(&text) {
+						Ok(backchannel_item) => {
+							if tx1.send(backchannel_item).is_err() {
+								tracing::error!("Error sending through the channel");
+								return;
+							}
+						}
+						Err(_) => {
+							tracing::error!("Invalid payload received");
+						}
+					}
+				}
+			});
+
+			// receive from subscribers and relay to ws
+			tokio::spawn(async move {
+				while let Ok(item) = rx_relay.recv().await {
+					if write
+						.send(Message::Text(serde_json::to_string(&item).unwrap()))
+						.await
+						.is_err()
+					{
+						tracing::error!("Error sending through ws");
+					}
+				}
+			});
+
+			*zombienet_bkc = Some(ZombienetBackchannel { broadcast_tx: tx, ws_tx: tx_relay });
+			return Ok(());
+		}
+
+		Err(BackchannelError::AlreadyInitialized)
+	}
+
+	pub fn broadcaster() -> Result<Broadcaster, BackchannelError> {
+		if ZOMBIENET_BACKCHANNEL.lock().unwrap().is_some() {
+			Ok(Broadcaster {})
+		} else {
+			Err(BackchannelError::Uninitialized)
+		}
+	}
+}