Fix/extension negotiation race

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

@@ -354,6 +354,16 @@ impl ChannelManager {
         Ok(channel_manager)
     }
 
+    /// Sets the negotiated extensions.
+    ///
+    /// This is used after upstream connection setup to store the extensions
+    /// that were successfully negotiated with the upstream server.
+    pub fn set_negotiated_extensions(&self, extensions: Vec<u16>) {
+        self.channel_manager_data.super_safe_lock(|data| {
+            data.negotiated_extensions = extensions;
+        });
+    }
+
     // Bootstraps a group channel with the given parameters.
     // Returns a `GroupChannel` if successful, otherwise returns `None`.
     //

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

@@ -17,6 +17,7 @@ use std::{
     marker::PhantomData,
 };
 use stratum_apps::{
+    extensions_negotiation::ExtensionNegotiationError,
     network_helpers,
     stratum_core::{
         binary_sv2, bitcoin,
@@ -224,9 +225,11 @@ pub enum JDCErrorKind {
     RequiredExtensionsNotSupported(Vec<u16>),
     /// Server requires extensions that the translator doesn't support
     ServerRequiresUnsupportedExtensions(Vec<u16>),
-    /// BitcoinCoreSv2TDP cancellation token activated
-    BitcoinCoreSv2TDPCancellationTokenActivated,
-    /// Failed to create BitcoinCoreSv2TDP tokio runtime
+    /// Extension negotiation timed out waiting for response
+    ExtensionNegotiationTimeout,
+    /// BitcoinCoreSv2 cancellation token activated
+    BitcoinCoreSv2CancellationTokenActivated,
+    /// Failed to create BitcoinCoreSv2 tokio runtime
     FailedToCreateBitcoinCoreTokioRuntime,
     /// Failed to send CoinbaseOutputConstraints message
     FailedToSendCoinbaseOutputConstraints,
@@ -368,8 +371,11 @@ impl fmt::Display for JDCErrorKind {
             ServerRequiresUnsupportedExtensions(extensions) => {
                 write!(f, "Server requires extensions that the translator doesn't support: {extensions:?}")
             }
-            BitcoinCoreSv2TDPCancellationTokenActivated => {
-                write!(f, "BitcoinCoreSv2TDP cancellation token activated")
+            ExtensionNegotiationTimeout => {
+                write!(f, "Extension negotiation timed out waiting for response")
+            }
+            BitcoinCoreSv2CancellationTokenActivated => {
+                write!(f, "BitcoinCoreSv2 cancellation token activated")
             }
             FailedToCreateBitcoinCoreTokioRuntime => {
                 write!(f, "Failed to create BitcoinCoreSv2TDP tokio runtime")
@@ -504,6 +510,20 @@ impl From<GroupChannelError> for JDCErrorKind {
     }
 }
 
+impl From<ExtensionNegotiationError> for JDCErrorKind {
+    fn from(e: ExtensionNegotiationError) -> Self {
+        match e {
+            ExtensionNegotiationError::SendError => JDCErrorKind::ChannelErrorSender,
+            ExtensionNegotiationError::ReceiveError(_) => JDCErrorKind::UnexpectedMessage(0, 0),
+            ExtensionNegotiationError::Timeout => JDCErrorKind::ExtensionNegotiationTimeout,
+            ExtensionNegotiationError::UnexpectedMessage(ext, msg) => {
+                JDCErrorKind::UnexpectedMessage(ext, msg as u8)
+            }
+            ExtensionNegotiationError::HandlerError(_) => JDCErrorKind::UnexpectedMessage(0, 0),
+        }
+    }
+}
+
 impl<Owner> HandlerErrorType for JDCError<Owner> {
     fn parse_error(error: ParserError) -> Self {
         Self {

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

@@ -171,7 +171,7 @@ impl JobDeclaratorClient {
             });
         }
 
-        let channel_manager_clone = channel_manager.clone();
+        let mut channel_manager_clone = channel_manager.clone();
         let mut bitcoin_core_sv2_join_handle: Option<JoinHandle<()>> = None;
 
         match self.config.template_provider_type().clone() {
@@ -291,17 +291,24 @@ impl JobDeclaratorClient {
             .await
         {
             Ok((upstream, job_declarator)) => {
-                upstream
+                // Start upstream and wait for extension negotiation to complete.
+                let negotiated_extensions = upstream
                     .start(
                         self.config.min_supported_version(),
                         self.config.max_supported_version(),
                         self.cancellation_token.clone(),
                         fallback_coordinator.clone(),
                         status_sender.clone(),
                         task_manager.clone(),
+                        &mut channel_manager_clone,
                     )
                     .await;
 
+                info!(
+                    "Upstream extension negotiation complete. Negotiated extensions: {:?}",
+                    negotiated_extensions
+                );
+
                 job_declarator
                     .start(
                         self.cancellation_token.clone(),
@@ -459,6 +466,7 @@ impl JobDeclaratorClient {
                                                 fallback_coordinator.clone(),
                                                 status_sender.clone(),
                                                 task_manager.clone(),
+                                                &mut channel_manager,
                                             )
                                             .await;
 

miner-apps/jd-client/src/lib/template_receiver/bitcoin_core.rs

@@ -39,7 +39,7 @@ pub async fn connect_to_bitcoin_core(
                 let status_sender = StatusSender::TemplateReceiver(status_sender_clone);
                 handle_error(
                     &status_sender,
-                    JDCError::<error::TemplateProvider>::shutdown(JDCErrorKind::BitcoinCoreSv2TDPCancellationTokenActivated),
+                    JDCError::<error::TemplateProvider>::shutdown(JDCErrorKind::BitcoinCoreSv2CancellationTokenActivated),
                 )
                 .await;
             }

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

@@ -6,6 +6,7 @@
 //! Responsibilities:
 //! - Establish a TCP + Noise encrypted connection to upstream
 //! - Perform `SetupConnection` handshake
+//! - Negotiate extensions synchronously before returning
 //! - Forward SV2 mining messages between upstream and channel manager
 //! - Handle common messages from upstream
 
@@ -15,12 +16,10 @@ use async_channel::{unbounded, Receiver, Sender};
 use bitcoin_core_sv2::template_distribution_protocol::CancellationToken;
 use stratum_apps::{
     custom_mutex::Mutex,
+    extensions_negotiation::negotiate_extensions,
     fallback_coordinator::FallbackCoordinator,
     network_helpers::{connect_with_noise, resolve_host},
-    stratum_core::{
-        binary_sv2::Seq064K, extensions_sv2::RequestExtensions, framing_sv2,
-        handlers_sv2::HandleCommonMessagesFromServerAsync, parsers_sv2::AnyMessage,
-    },
+    stratum_core::{framing_sv2, handlers_sv2::HandleCommonMessagesFromServerAsync},
     task_manager::TaskManager,
     utils::{
         protocol_message_type::{protocol_message_type, MessageType},
@@ -31,6 +30,7 @@ use tokio::net::TcpStream;
 use tracing::{debug, error, info, warn};
 
 use crate::{
+    channel_manager::ChannelManager,
     error::{self, JDCError, JDCErrorKind, JDCResult},
     io_task::spawn_io_tasks,
     status::{handle_error, Status, StatusSender},
@@ -151,11 +151,18 @@ impl Upstream {
     /// Perform `SetupConnection` handshake with upstream.
     ///
     /// Sends [`SetupConnection`] and awaits response.
+    /// If required extensions are configured, negotiates them synchronously
+    /// before returning.
+    ///
+    /// # Returns
+    /// * `Ok(Vec<u16>)` - The list of negotiated extensions (empty if none were requested)
+    /// * `Err(JDCError)` - Error during handshake or extension negotiation
     pub async fn setup_connection(
         &mut self,
         min_version: u16,
         max_version: u16,
-    ) -> JDCResult<(), error::Upstream> {
+        channel_manager: &mut ChannelManager,
+    ) -> JDCResult<Vec<u16>, error::Upstream> {
         info!("Upstream: initiating SV2 handshake...");
         let setup_connection =
             get_setup_connection_message(min_version, max_version, &self.address)
@@ -197,63 +204,49 @@ impl Upstream {
             .await?;
 
         // Send RequestExtensions after successful SetupConnection if there are required extensions
+        // and wait for the response before returning
         if !self.required_extensions.is_empty() {
-            self.send_request_extensions().await?;
+            let negotiated = self.negotiate_extensions(channel_manager).await?;
+            return Ok(negotiated);
         }
 
-        Ok(())
+        Ok(vec![])
     }
 
-    /// Send `RequestExtensions` message to upstream.
-    /// The supported extensions are stored for potential retry if the server requires additional
-    /// extensions.
-    async fn send_request_extensions(&mut self) -> JDCResult<(), error::Upstream> {
-        info!(
-            "Sending RequestExtensions to upstream with required extensions: {:?}",
-            self.required_extensions
-        );
-        if self.required_extensions.is_empty() {
-            return Ok(());
-        }
-
-        let requested_extensions =
-            Seq064K::new(self.required_extensions.clone()).map_err(JDCError::shutdown)?;
-
-        let request_extensions = RequestExtensions {
-            request_id: 0,
-            requested_extensions,
-        };
-
-        info!(
-            "Sending RequestExtensions to upstream with required extensions: {:?}",
-            self.required_extensions
-        );
-
-        let sv2_frame: Sv2Frame = AnyMessage::Extensions(request_extensions.into_static().into())
-            .try_into()
-            .map_err(JDCError::shutdown)?;
-
-        self.upstream_channel
-            .upstream_sender
-            .send(sv2_frame)
-            .await
-            .map_err(|e| {
-                error!(?e, "Failed to send RequestExtensions to upstream");
-                JDCError::fallback(JDCErrorKind::ChannelErrorSender)
-            })?;
-
-        info!("Sent RequestExtensions to upstream");
-        Ok(())
+    /// Sends RequestExtensions and waits for the response.
+    ///
+    /// Delegates to the shared [`stratum_apps::extensions_negotiation::negotiate_extensions`] function.
+    ///
+    /// # Returns
+    /// * `Ok(Vec<u16>)` - The list of successfully negotiated extensions
+    /// * `Err(JDCError)` - Extension negotiation failed
+    async fn negotiate_extensions(
+        &mut self,
+        channel_manager: &mut ChannelManager,
+    ) -> JDCResult<Vec<u16>, error::Upstream> {
+        negotiate_extensions(
+            self.required_extensions.clone(),
+            self.upstream_channel.upstream_sender.clone(),
+            self.upstream_channel.upstream_receiver.clone(),
+            self.upstream_channel.channel_manager_receiver.clone(),
+            channel_manager,
+        )
+        .await
+        .map_err(|e| JDCError::fallback(JDCErrorKind::from(e)))
     }
 
     /// Start unified upstream loop.
     ///
     /// Responsibilities:
-    /// - Run `setup_connection`
+    /// - Run `setup_connection` (including extension negotiation)
     /// - Handle messages from upstream (pool) and channel manager
     /// - React to shutdown signals
     ///
-    /// This function spawns an async task and returns immediately.
+    /// This function spawns an async task and returns the negotiated extensions.
+    ///
+    /// # Returns
+    /// * `Vec<u16>` - The list of negotiated extensions (empty if none were requested or setup
+    ///   failed)
     #[allow(clippy::too_many_arguments)]
     pub async fn start(
         mut self,
@@ -263,13 +256,26 @@ impl Upstream {
         fallback_coordinator: FallbackCoordinator,
         status_sender: Sender<Status>,
         task_manager: Arc<TaskManager>,
-    ) {
+        channel_manager: &mut ChannelManager,
+    ) -> Vec<u16> {
         let status_sender = StatusSender::Upstream(status_sender);
 
-        if let Err(e) = self.setup_connection(min_version, max_version).await {
-            error!(error = ?e, "Upstream: connection setup failed.");
-            return;
-        }
+        let negotiated_extensions = match self
+            .setup_connection(min_version, max_version, channel_manager)
+            .await
+        {
+            Ok(extensions) => {
+                info!(
+                    "Upstream: extension negotiation complete. Extensions: {:?}",
+                    extensions
+                );
+                extensions
+            }
+            Err(e) => {
+                error!(error = ?e, "Upstream: connection setup failed.");
+                return vec![];
+            }
+        };
 
         task_manager.spawn(async move {
             // we just spawned a new task that's relevant to fallback coordination
@@ -315,6 +321,8 @@ impl Upstream {
             // signal fallback coordinator that this task has completed its cleanup
             fallback_handler.done();
         });
+
+        negotiated_extensions
     }
 
     // Handle incoming frames from upstream (pool).

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

@@ -15,6 +15,7 @@ use std::{
     sync::PoisonError,
 };
 use stratum_apps::{
+    extensions_negotiation::ExtensionNegotiationError,
     stratum_core::{
         binary_sv2,
         channels_sv2::client::error::GroupChannelError,
@@ -176,6 +177,8 @@ pub enum TproxyErrorKind {
     RequiredExtensionsNotSupported(Vec<u16>),
     /// Server requires extensions that the translator doesn't support
     ServerRequiresUnsupportedExtensions(Vec<u16>),
+    /// Extension negotiation timed out waiting for response
+    ExtensionNegotiationTimeout,
     /// Represents a generic channel send failure, described by a string.
     General(String),
     /// Error bubbling up from translator-core library
@@ -256,6 +259,9 @@ impl fmt::Display for TproxyErrorKind {
                     extensions
                 )
             }
+            ExtensionNegotiationTimeout => {
+                write!(f, "Extension negotiation timed out waiting for response")
+            }
             SV1Error => write!(f, "Sv1 error"),
             TranslatorCore(ref e) => write!(f, "Translator core error: {e:?}"),
             NetworkHelpersError(ref e) => write!(f, "Network helpers error: {e:?}"),
@@ -396,6 +402,20 @@ impl HandlerErrorType for TproxyErrorKind {
     }
 }
 
+impl From<ExtensionNegotiationError> for TproxyErrorKind {
+    fn from(e: ExtensionNegotiationError) -> Self {
+        match e {
+            ExtensionNegotiationError::SendError => TproxyErrorKind::ChannelErrorSender,
+            ExtensionNegotiationError::ReceiveError(_) => TproxyErrorKind::UnexpectedMessage(0, 0),
+            ExtensionNegotiationError::Timeout => TproxyErrorKind::ExtensionNegotiationTimeout,
+            ExtensionNegotiationError::UnexpectedMessage(ext, msg) => {
+                TproxyErrorKind::UnexpectedMessage(ext, msg as u8)
+            }
+            ExtensionNegotiationError::HandlerError(_) => TproxyErrorKind::UnexpectedMessage(0, 0),
+        }
+    }
+}
+
 impl<Owner> HandlerErrorType for TproxyError<Owner> {
     fn parse_error(error: ParserError) -> Self {
         Self {