wip

Author: lvboudre

apps/jet/src/tpu_sink.rs

@@ -7,13 +7,13 @@ use {
         task::{Context, Poll, ready},
     },
     yellowstone_jet_tpu_client::yellowstone_grpc::sender::{
-        PollTpuSender, SendErrorKind, ShieldBlockList, YellowstoneTpuSender,
+        PollYellowstoneTpuSender, SendErrorKind, ShieldBlockList, YellowstoneTpuSender,
     },
     yellowstone_shield_store::PolicyStore,
 };
 
 pub struct TpuSink<St> {
-    tpu_sender: PollTpuSender,
+    tpu_sender: PollYellowstoneTpuSender,
     source: St,
     shield_store: Option<PolicyStore>,
 }
@@ -25,7 +25,7 @@ impl<St> TpuSink<St> {
         shield_store: Option<PolicyStore>,
     ) -> Self {
         Self {
-            tpu_sender: PollTpuSender::new(tpu_sender),
+            tpu_sender: PollYellowstoneTpuSender::new(tpu_sender),
             source,
             shield_store,
         }

crates/tpu-client/src/core.rs

@@ -43,6 +43,7 @@ use crate::prom;
 use {
     crate::config::{TpuOverrideInfo, TpuPortKind, TpuSenderConfig},
     bytes::Bytes,
+    core::fmt,
     derive_more::Display,
     futures::task::AtomicWaker,
     quinn::{
@@ -58,11 +59,12 @@ use {
     solana_tls_utils::{QuicClientCertificate, SkipServerVerification, new_dummy_x509_certificate},
     std::{
         collections::{BTreeMap, HashMap, HashSet, VecDeque},
+        future,
         net::{IpAddr, Ipv4Addr, SocketAddr},
         num::NonZeroUsize,
         pin::Pin,
         sync::{Arc, Mutex as StdMutex, atomic::AtomicBool},
-        task::Poll,
+        task::{Poll, ready},
         time::{Duration, Instant},
     },
     tokio::{
@@ -74,6 +76,7 @@ use {
         task::{self, Id, JoinError, JoinHandle, JoinSet},
         time::{Sleep, interval},
     },
+    tokio_util::sync::PollSender,
 };
 
 /// This has been copy-pasted from `solana_streamer::nonblocking::quic::ALPN_TPU_PROTOCOL_ID`
@@ -149,10 +152,25 @@ struct MultiStepIdentitySynchronizationCommand {
 /// Command to control driver behavior.
 ///
 enum DriverCommand {
-    UpdateIdenttiy(UpdateIdentityCommand),
+    UpdateIdentity(UpdateIdentityCommand),
     MultiStepIdentitySynchronization(MultiStepIdentitySynchronizationCommand),
 }
 
+impl fmt::Debug for DriverCommand {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            DriverCommand::UpdateIdentity(cmd) => f
+                .debug_struct("UpdateIdentity")
+                .field("new_identity", &cmd.new_identity.pubkey())
+                .finish(),
+            DriverCommand::MultiStepIdentitySynchronization(cmd) => f
+                .debug_struct("MultiStepIdentitySynchronization")
+                .field("new_identity", &cmd.new_identity.pubkey())
+                .finish(),
+        }
+    }
+}
+
 enum DriverTaskMeta {
     DropAllWorkers,
 }
@@ -2696,7 +2714,7 @@ where
 
     async fn handle_cnc(&mut self, command: DriverCommand) {
         match command {
-            DriverCommand::UpdateIdenttiy(cmd) => {
+            DriverCommand::UpdateIdentity(cmd) => {
                 let UpdateIdentityCommand {
                     new_identity,
                     callback,
@@ -3382,7 +3400,7 @@ impl TpuSenderDriverSpawner {
         TpuSenderSessionContext {
             driver_tx_sink: tx_inlet,
             identity_updater: TpuSenderIdentityUpdater {
-                cnc_tx: driver_cnc_tx,
+                cnc_tx: PollSender::new(driver_cnc_tx),
             },
             driver_join_handle: jh,
         }
@@ -3392,11 +3410,12 @@ impl TpuSenderDriverSpawner {
 ///
 /// Handle to update the identity used by the TPU sender driver.
 ///
+#[derive(Clone)]
 pub struct TpuSenderIdentityUpdater {
     ///
     /// Command-and-control channel to send command to the QUIC driver
     ///  
-    cnc_tx: mpsc::Sender<DriverCommand>,
+    cnc_tx: PollSender<DriverCommand>,
 }
 
 ///
@@ -3406,7 +3425,7 @@ impl TpuSenderIdentityUpdater {
     ///
     /// Changes the configured identity in the QUIC driver
     ///
-    pub async fn update_identity(&mut self, identity: Keypair) {
+    pub fn update_identity(&mut self, identity: Keypair) -> UpdateIdentity {
         let shared = UpdateIdentityInner {
             set: AtomicBool::new(false),
             waker: AtomicWaker::new(),
@@ -3416,15 +3435,12 @@ impl TpuSenderIdentityUpdater {
             new_identity: identity,
             callback: Arc::clone(&shared),
         };
-        self.cnc_tx
-            .send(DriverCommand::UpdateIdenttiy(cmd))
-            .await
-            .expect("disconnected");
-        let update_identity = UpdateIdentity {
+        let cnc_tx = self.cnc_tx.clone();
+
+        UpdateIdentity {
             inner: shared,
-            _this: self,
-        };
-        update_identity.await
+            state: UpdateIdentityState::Init { cnc_tx, cmd },
+        }
     }
 
     ///
@@ -3446,10 +3462,13 @@ impl TpuSenderIdentityUpdater {
             new_identity: identity,
             barrier,
         };
-        self.cnc_tx
-            .send(DriverCommand::MultiStepIdentitySynchronization(cmd))
+        let mut cnc_tx = self.cnc_tx.clone();
+        future::poll_fn(|cx| cnc_tx.poll_reserve(cx))
             .await
             .expect("disconnected");
+        cnc_tx
+            .send_item(DriverCommand::MultiStepIdentitySynchronization(cmd))
+            .expect("disconnected");
     }
 }
 
@@ -3461,39 +3480,106 @@ struct UpdateIdentityInner {
     waker: AtomicWaker,
 }
 
+enum UpdateIdentityState {
+    Init {
+        cnc_tx: PollSender<DriverCommand>,
+        cmd: UpdateIdentityCommand,
+    },
+    Closed,
+    WaitingForCompletion,
+}
+
 ///
 /// Future that waits for the identity update to complete.
 /// This future is used to ensure that the identity update is completed before proceeding.
 ///
-pub struct UpdateIdentity<'a> {
+pub struct UpdateIdentity {
     inner: Arc<UpdateIdentityInner>,
-    _this: &'a TpuSenderIdentityUpdater, /* phantom data to prevent two threads from updating the identity at the same time */
+    state: UpdateIdentityState,
+}
+
+impl UpdateIdentity {
+    fn take_state(&mut self) -> UpdateIdentityState {
+        std::mem::replace(&mut self.state, UpdateIdentityState::Closed)
+    }
 }
 
-impl Future for UpdateIdentity<'_> {
-    type Output = ();
+#[derive(Debug, thiserror::Error)]
+#[error("tpu sender runtime closed")]
+pub struct UpdateIdentityError;
+
+impl Future for UpdateIdentity {
+    type Output = Result<(), UpdateIdentityError>;
 
     fn poll(
         self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Self::Output> {
-        // quick check to avoid registration if already done.
-        if self.inner.set.load(std::sync::atomic::Ordering::Relaxed) {
-            return Poll::Ready(());
-        }
-
-        self.inner.waker.register(cx.waker());
+        let this = self.get_mut();
+
+        let (result, next_state) = match this.take_state() {
+            UpdateIdentityState::Init { mut cnc_tx, cmd } => {
+                match ready!(cnc_tx.poll_reserve(cx)) {
+                    Ok(()) => {
+                        match cnc_tx.send_item(DriverCommand::UpdateIdentity(cmd)) {
+                            Ok(()) => {
+                                this.inner.waker.register(cx.waker());
+                                // Between the time we send and the registration, the driver may have already completed the update and set the atomic bool.
+                                // So we need to immediately wake the waker, so that the next poll will check the atomic bool and return ready if the update is already complete.
+                                cx.waker().wake_by_ref();
+                                (None, UpdateIdentityState::WaitingForCompletion)
+                            }
+                            Err(_) => (Some(Err(UpdateIdentityError)), UpdateIdentityState::Closed),
+                        }
+                    }
+                    Err(_) => (Some(Err(UpdateIdentityError)), UpdateIdentityState::Closed),
+                }
+            }
+            UpdateIdentityState::WaitingForCompletion => {
+                if this.inner.set.load(std::sync::atomic::Ordering::Relaxed) {
+                    (Some(Ok(())), UpdateIdentityState::Closed)
+                } else {
+                    (None, UpdateIdentityState::WaitingForCompletion)
+                }
+            }
+            UpdateIdentityState::Closed => {
+                panic!("UpdateIdentity polled after completion");
+            }
+        };
 
-        // Need to check condition **after** `register` to avoid a race
-        // condition that would result in lost notifications.
-        if self.inner.set.load(std::sync::atomic::Ordering::Relaxed) {
-            Poll::Ready(())
+        this.state = next_state;
+        if let Some(result) = result {
+            return Poll::Ready(result);
         } else {
-            Poll::Pending
+            return Poll::Pending;
         }
     }
 }
 
+// impl Future for UpdateIdentity<'_> {
+//     type Output = ();
+
+//     fn poll(
+//         self: std::pin::Pin<&mut Self>,
+//         cx: &mut std::task::Context<'_>,
+//     ) -> std::task::Poll<Self::Output> {
+//         // quick check to avoid registration if already done.
+//         if self.inner.set.load(std::sync::atomic::Ordering::Relaxed) {
+//             return Poll::Ready(());
+//         }
+
+//         self.inner.waker.register(cx.waker());
+
+//         // Need to check condition **after** `register` to avoid a race
+//         // condition that would result in lost notifications.
+//         if self.inner.set.load(std::sync::atomic::Ordering::Relaxed) {
+//             Poll::Ready(())
+//         } else {
+//             Poll::Pending
+//         }
+//     }
+// }
+
 pub const fn module_path_for_test() -> &'static str {
     module_path!()
 }
@@ -3508,15 +3594,18 @@ mod test {
         solana_pubkey::Pubkey,
         std::time::Duration,
         tokio::sync::mpsc,
+        tokio_util::sync::PollSender,
     };
 
     #[tokio::test]
     async fn test_update_identity_fut() {
         let (cnc_tx, mut cnc_rx) = mpsc::channel(10);
-        let mut updater = TpuSenderIdentityUpdater { cnc_tx };
+        let mut updater = TpuSenderIdentityUpdater {
+            cnc_tx: PollSender::new(cnc_tx),
+        };
 
         let jh = tokio::spawn(async move {
-            let DriverCommand::UpdateIdenttiy(UpdateIdentityCommand {
+            let DriverCommand::UpdateIdentity(UpdateIdentityCommand {
                 new_identity,
                 callback,
             }) = cnc_rx.recv().await.unwrap()
@@ -3533,7 +3622,10 @@ mod test {
         });
 
         let identity = Keypair::new();
-        updater.update_identity(identity.insecure_clone()).await;
+        updater
+            .update_identity(identity.insecure_clone())
+            .await
+            .unwrap();
 
         let actual = jh.await.unwrap();
         assert_eq!(actual, identity)