remove broadcast channel from tproxy

Author: Shourya742

miner-apps/jd-client/src/lib/mod.rs

@@ -16,6 +16,7 @@ use stratum_apps::{
     tp_type::TemplateProviderType,
     utils::types::{Sv2Frame, GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS},
 };
+use tokio::sync::Notify;
 use tracing::{debug, error, info, warn};
 
 use crate::{

miner-apps/translator/src/lib/error.rs

@@ -29,7 +29,6 @@ use stratum_apps::{
         MessageType,
     },
 };
-use tokio::sync::broadcast;
 
 pub type TproxyResult<T, Owner> = Result<T, TproxyError<Owner>>;
 
@@ -158,10 +157,6 @@ pub enum TproxyErrorKind {
     ChannelErrorReceiver(async_channel::RecvError),
     /// Channel sender error
     ChannelErrorSender,
-    /// Broadcast channel receiver error
-    BroadcastChannelErrorReceiver(broadcast::error::RecvError),
-    /// Tokio channel receiver error
-    TokioChannelErrorRecv(tokio::sync::broadcast::error::RecvError),
     /// Error converting SetDifficulty to Message
     SetDifficultyToMessage(SetDifficulty),
     /// Received an unexpected message type
@@ -223,11 +218,7 @@ impl fmt::Display for TproxyErrorKind {
             ParseInt(ref e) => write!(f, "Bad convert from `String` to `int`: `{e:?}`"),
             PoisonLock => write!(f, "Poison Lock error"),
             ChannelErrorReceiver(ref e) => write!(f, "Channel receive error: `{e:?}`"),
-            BroadcastChannelErrorReceiver(ref e) => {
-                write!(f, "Broadcast channel receive error: {e:?}")
-            }
             ChannelErrorSender => write!(f, "Sender error"),
-            TokioChannelErrorRecv(ref e) => write!(f, "Channel receive error: `{e:?}`"),
             SetDifficultyToMessage(ref e) => {
                 write!(f, "Error converting SetDifficulty to Message: `{e:?}`")
             }
@@ -333,12 +324,6 @@ impl From<async_channel::RecvError> for TproxyErrorKind {
     }
 }
 
-impl From<tokio::sync::broadcast::error::RecvError> for TproxyErrorKind {
-    fn from(e: tokio::sync::broadcast::error::RecvError) -> Self {
-        TproxyErrorKind::TokioChannelErrorRecv(e)
-    }
-}
-
 //*** LOCK ERRORS ***
 impl<T> From<PoisonError<T>> for TproxyErrorKind {
     fn from(_e: PoisonError<T>) -> Self {

miner-apps/translator/src/lib/sv1/downstream/channel.rs

@@ -1,9 +1,5 @@
 use async_channel::{Receiver, Sender};
-use stratum_apps::{
-    stratum_core::sv1_api::json_rpc,
-    utils::types::{ChannelId, DownstreamId},
-};
-use tokio::sync::broadcast;
+use stratum_apps::{stratum_core::sv1_api::json_rpc, utils::types::DownstreamId};
 use tokio_util::sync::CancellationToken;
 use tracing::debug;
 
@@ -12,8 +8,7 @@ pub struct DownstreamChannelState {
     pub downstream_sv1_sender: Sender<json_rpc::Message>,
     pub downstream_sv1_receiver: Receiver<json_rpc::Message>,
     pub sv1_server_sender: Sender<(DownstreamId, json_rpc::Message)>,
-    pub sv1_server_broadcast:
-        broadcast::Sender<(ChannelId, Option<DownstreamId>, json_rpc::Message)>, /* channel_id, optional downstream_id, message */
+    pub sv1_server_receiver: Receiver<json_rpc::Message>,
     /// Per-connection cancellation token (child of the global token).
     /// Cancelled when this downstream's task loop exits, causing
     /// the associated SV1 I/O task to shut down.
@@ -26,17 +21,13 @@ impl DownstreamChannelState {
         downstream_sv1_sender: Sender<json_rpc::Message>,
         downstream_sv1_receiver: Receiver<json_rpc::Message>,
         sv1_server_sender: Sender<(DownstreamId, json_rpc::Message)>,
-        sv1_server_broadcast: broadcast::Sender<(
-            ChannelId,
-            Option<DownstreamId>,
-            json_rpc::Message,
-        )>,
+        sv1_server_receiver: Receiver<json_rpc::Message>,
         connection_token: CancellationToken,
     ) -> Self {
         Self {
             downstream_sv1_receiver,
             downstream_sv1_sender,
-            sv1_server_broadcast,
+            sv1_server_receiver,
             sv1_server_sender,
             connection_token,
         }

miner-apps/translator/src/lib/sv1/downstream/downstream.rs

@@ -2,7 +2,6 @@ use crate::{
     error::{self, TproxyError, TproxyErrorKind, TproxyResult},
     status::{handle_error, StatusSender},
     sv1::downstream::{channel::DownstreamChannelState, data::DownstreamData},
-    utils::AGGREGATED_CHANNEL_ID,
 };
 use async_channel::{Receiver, Sender};
 use std::{
@@ -23,9 +22,8 @@ use stratum_apps::{
         },
     },
     task_manager::TaskManager,
-    utils::types::{ChannelId, DownstreamId, Hashrate},
+    utils::types::{DownstreamId, Hashrate},
 };
-use tokio::sync::broadcast;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, warn};
 
@@ -62,11 +60,7 @@ impl Downstream {
         downstream_sv1_sender: Sender<json_rpc::Message>,
         downstream_sv1_receiver: Receiver<json_rpc::Message>,
         sv1_server_sender: Sender<(DownstreamId, json_rpc::Message)>,
-        sv1_server_broadcast: broadcast::Sender<(
-            ChannelId,
-            Option<DownstreamId>,
-            json_rpc::Message,
-        )>,
+        sv1_server_receiver: Receiver<json_rpc::Message>,
         target: Target,
         hashrate: Option<Hashrate>,
         connection_token: CancellationToken,
@@ -76,7 +70,7 @@ impl Downstream {
             downstream_sv1_sender,
             downstream_sv1_receiver,
             sv1_server_sender,
-            sv1_server_broadcast,
+            sv1_server_receiver,
             connection_token,
         );
         Self {
@@ -106,10 +100,6 @@ impl Downstream {
         status_sender: StatusSender,
         task_manager: Arc<TaskManager>,
     ) {
-        let mut sv1_server_receiver = self
-            .downstream_channel_state
-            .sv1_server_broadcast
-            .subscribe();
         let downstream_id = self.downstream_id;
         task_manager.spawn(async move {
             // we just spawned a new task that's relevant to fallback coordination
@@ -141,7 +131,7 @@ impl Downstream {
                     }
 
                     // Handle server -> downstream message
-                    res = self.handle_sv1_server_message(&mut sv1_server_receiver) => {
+                    res = self.handle_sv1_server_message() => {
                         if let Err(e) = res {
                             error!("Downstream {downstream_id}: error in server message handler: {e:?}");
                             if handle_error(&status_sender, e).await {
@@ -181,25 +171,16 @@ impl Downstream {
     ///   complete
     /// - On handshake completion: sends cached messages in correct order (set_difficulty first,
     ///   then notify)
-    pub async fn handle_sv1_server_message(
-        &self,
-        sv1_server_receiver: &mut broadcast::Receiver<(
-            ChannelId,
-            Option<DownstreamId>,
-            json_rpc::Message,
-        )>,
-    ) -> TproxyResult<(), error::Downstream> {
-        match sv1_server_receiver.recv().await {
-            Ok((channel_id, downstream_id, message)) => {
-                let my_channel_id = self.downstream_data.super_safe_lock(|d| d.channel_id);
-                let my_downstream_id = self.downstream_id;
+    pub async fn handle_sv1_server_message(&self) -> TproxyResult<(), error::Downstream> {
+        match self
+            .downstream_channel_state
+            .sv1_server_receiver
+            .recv()
+            .await
+        {
+            Ok(message) => {
+                let downstream_id = self.downstream_id;
                 let handshake_complete = self.sv1_handshake_complete.load(Ordering::SeqCst);
-                let id_matches = (my_channel_id == Some(channel_id)
-                    || channel_id == AGGREGATED_CHANNEL_ID)
-                    && (downstream_id.is_none() || downstream_id == Some(my_downstream_id));
-                if !id_matches {
-                    return Ok(()); // Message not intended for this downstream
-                }
 
                 // Check if this is a queued message response
                 let is_queued_sv1_handshake_response = self
@@ -267,7 +248,7 @@ impl Downstream {
                                                 "Down: Failed to send mining.set_difficulty to downstream: {:?}",
                                                 e
                                             );
-                                            TproxyError::disconnect(TproxyErrorKind::ChannelErrorSender, downstream_id.unwrap_or(0))
+                                            TproxyError::disconnect(TproxyErrorKind::ChannelErrorSender, downstream_id)
                                         })?;
                                 }
 
@@ -279,7 +260,7 @@ impl Downstream {
                                         .await
                                         .map_err(|e| {
                                             error!("Down: Failed to send mining.notify to downstream: {:?}", e);
-                                            TproxyError::disconnect(TproxyErrorKind::ChannelErrorSender, downstream_id.unwrap_or(0))
+                                            TproxyError::disconnect(TproxyErrorKind::ChannelErrorSender, downstream_id)
                                         })?;
                                 }
                                 return Ok(());
@@ -297,7 +278,7 @@ impl Downstream {
                                         );
                                         TproxyError::disconnect(
                                             TproxyErrorKind::ChannelErrorSender,
-                                            downstream_id.unwrap_or(0),
+                                            downstream_id,
                                         )
                                     })?;
                             }
@@ -338,7 +319,7 @@ impl Downstream {
                             error!("Down: Failed to send queued message to downstream: {:?}", e);
                             TproxyError::disconnect(
                                 TproxyErrorKind::ChannelErrorSender,
-                                downstream_id.unwrap_or(0),
+                                downstream_id,
                             )
                         })?;
                 } else {
@@ -348,12 +329,11 @@ impl Downstream {
                 }
             }
             Err(e) => {
-                let downstream_id = self.downstream_id;
                 error!(
                     "Sv1 message handler error for downstream {}: {:?}",
-                    downstream_id, e
+                    self.downstream_id, e
                 );
-                return Err(TproxyError::disconnect(e, downstream_id));
+                return Err(TproxyError::disconnect(e, self.downstream_id));
             }
         }
 

miner-apps/translator/src/lib/sv1/sv1_server/channel.rs

@@ -1,16 +1,13 @@
 use async_channel::{unbounded, Receiver, Sender};
+use dashmap::DashMap;
+use std::sync::Arc;
 use stratum_apps::stratum_core::parsers_sv2::{Mining, Tlv};
 
-use stratum_apps::{
-    stratum_core::sv1_api::json_rpc,
-    utils::types::{ChannelId, DownstreamId},
-};
-use tokio::sync::broadcast;
+use stratum_apps::{stratum_core::sv1_api::json_rpc, utils::types::DownstreamId};
 
 #[derive(Clone)]
 pub struct Sv1ServerChannelState {
-    pub sv1_server_to_downstream_sender:
-        broadcast::Sender<(ChannelId, Option<DownstreamId>, json_rpc::Message)>,
+    pub sv1_server_to_downstream_sender: Arc<DashMap<DownstreamId, Sender<json_rpc::Message>>>,
     pub downstream_to_sv1_server_sender: Sender<(DownstreamId, json_rpc::Message)>,
     pub downstream_to_sv1_server_receiver: Receiver<(DownstreamId, json_rpc::Message)>,
     pub channel_manager_receiver: Receiver<(Mining<'static>, Option<Vec<Tlv>>)>,
@@ -23,11 +20,10 @@ impl Sv1ServerChannelState {
         channel_manager_receiver: Receiver<(Mining<'static>, Option<Vec<Tlv>>)>,
         channel_manager_sender: Sender<(Mining<'static>, Option<Vec<Tlv>>)>,
     ) -> Self {
-        let (sv1_server_to_downstream_sender, _) = broadcast::channel(1000);
         let (downstream_to_sv1_server_sender, downstream_to_sv1_server_receiver) = unbounded();
 
         Self {
-            sv1_server_to_downstream_sender,
+            sv1_server_to_downstream_sender: Arc::new(DashMap::new()),
             downstream_to_sv1_server_receiver,
             downstream_to_sv1_server_sender,
             channel_manager_receiver,

miner-apps/translator/src/lib/sv1/sv1_server/difficulty_manager.rs

@@ -150,24 +150,26 @@ impl Sv1Server {
         }
 
         // Process immediate set_difficulty updates (for new_target >= upstream_target)
-        for (channel_id, downstream_id, target) in immediate_updates {
+        for (_channel_id, downstream_id, target) in immediate_updates {
             // Send set_difficulty message immediately
             if let Ok(set_difficulty_msg) = build_sv1_set_difficulty_from_sv2_target(target) {
-                if let Err(e) = self
+                let ds_id = downstream_id.unwrap_or(0);
+                if let Some(sender) = self
                     .sv1_server_channel_state
                     .sv1_server_to_downstream_sender
-                    .send((channel_id, downstream_id, set_difficulty_msg))
+                    .get(&ds_id)
                 {
-                    error!(
-                        "Failed to send immediate SetDifficulty message to downstream {}: {:?}",
-                        downstream_id.unwrap_or(0),
-                        e
-                    );
-                } else {
-                    trace!(
-                        "Sent immediate SetDifficulty to downstream {} (new_target >= upstream_target)",
-                        downstream_id.unwrap_or(0)
-                    );
+                    if let Err(e) = sender.value().try_send(set_difficulty_msg) {
+                        error!(
+                            "Failed to send immediate SetDifficulty message to downstream {}: {:?}",
+                            ds_id, e
+                        );
+                    } else {
+                        trace!(
+                            "Sent immediate SetDifficulty to downstream {} (new_target >= upstream_target)",
+                            ds_id
+                        );
+                    }
                 }
             }
         }
@@ -426,7 +428,7 @@ impl Sv1Server {
                 .get(&update.downstream_id)
                 .and_then(|ds| ds.downstream_data.super_safe_lock(|d| d.channel_id));
 
-            let Some(channel_id) = channel_id else {
+            let Some(_channel_id) = channel_id else {
                 trace!(
                     "Skipping SetDifficulty for downstream {}: no channel_id yet",
                     update.downstream_id
@@ -446,17 +448,19 @@ impl Sv1Server {
                     }
                 };
 
-            if let Err(e) = self
+            if let Some(sender) = self
                 .sv1_server_channel_state
                 .sv1_server_to_downstream_sender
-                .send((channel_id, Some(update.downstream_id), set_difficulty_msg))
+                .get(&update.downstream_id)
             {
-                error!(
-                    "Failed to send SetDifficulty to downstream {}: {:?}",
-                    update.downstream_id, e
-                );
-            } else {
-                trace!("Sent SetDifficulty to downstream {}", update.downstream_id);
+                if let Err(e) = sender.value().try_send(set_difficulty_msg) {
+                    error!(
+                        "Failed to send SetDifficulty to downstream {}: {:?}",
+                        update.downstream_id, e
+                    );
+                } else {
+                    trace!("Sent SetDifficulty to downstream {}", update.downstream_id);
+                }
             }
         }
     }