Revert "Replacing notify by cancellation token to correctly shutdown …

…tasks" This reverts commit 945dab5.
blockworks-foundation · Mar 28, 2024 · 3773319 · 3773319
1 parent 945dab5
commit 3773319
Show file tree

Hide file tree

Showing 9 changed files with 24 additions and 31 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -61,7 +61,6 @@ lazy_static = "1.4.0"
 dotenv = "0.15.0"
 async-channel = "1.8.0"
 jemallocator = "0.5"
-tokio-util = "0.7.10"
 
 quinn = "0.10.2"
 quinn-proto = "0.10.5"

diff --git a/cluster-endpoints/Cargo.toml b/cluster-endpoints/Cargo.toml
@@ -45,4 +45,3 @@ itertools = {workspace = true}
 prometheus = { workspace = true }
 lazy_static = { workspace = true }
 tonic-health = { workspace = true }
-tokio-util = { workspace = true }
diff --git a/cluster-endpoints/src/grpc_multiplex.rs b/cluster-endpoints/src/grpc_multiplex.rs
@@ -16,9 +16,10 @@ use solana_lite_rpc_core::AnyhowJoinHandle;
 use solana_sdk::clock::Slot;
 use solana_sdk::commitment_config::CommitmentConfig;
 use std::collections::{BTreeSet, HashMap, HashSet};
+use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::broadcast::Receiver;
-use tokio_util::sync::CancellationToken;
+use tokio::sync::Notify;
 use yellowstone_grpc_proto::geyser::subscribe_update::UpdateOneof;
 use yellowstone_grpc_proto::geyser::SubscribeUpdate;
 
@@ -54,7 +55,7 @@ impl FromYellowstoneExtractor for BlockMetaHashExtractor {
 fn create_grpc_multiplex_processed_block_stream(
     grpc_sources: &Vec<GrpcSourceConfig>,
     processed_block_sender: async_channel::Sender<ProducedBlock>,
-    exit_notfier: CancellationToken,
+    exit_notfier: Arc<Notify>,
 ) -> Vec<AnyhowJoinHandle> {
     let commitment_config = CommitmentConfig::processed();
 
@@ -138,7 +139,7 @@ pub fn create_grpc_multiplex_blocks_subscription(
                 let (processed_block_sender, processed_block_reciever) =
                     async_channel::unbounded::<ProducedBlock>();
 
-                let exit_notify = CancellationToken::new();
+                let exit_notify = Arc::new(Notify::new());
                 let processed_blocks_tasks = create_grpc_multiplex_processed_block_stream(
                     &grpc_sources,
                     processed_block_sender,
@@ -244,7 +245,7 @@ pub fn create_grpc_multiplex_blocks_subscription(
                         }
                     }
                 }
-                exit_notify.cancel();
+                exit_notify.notify_waiters();
                 futures::future::join_all(processed_blocks_tasks).await;
             }
         })

diff --git a/cluster-endpoints/src/grpc_subscription.rs b/cluster-endpoints/src/grpc_subscription.rs
@@ -33,7 +33,7 @@ use solana_sdk::{
 };
 use solana_transaction_status::{Reward, RewardType};
 use std::{collections::HashMap, sync::Arc};
-use tokio_util::sync::CancellationToken;
+use tokio::sync::Notify;
 use yellowstone_grpc_client::GeyserGrpcClient;
 use yellowstone_grpc_proto::geyser::{SubscribeRequestFilterSlots, SubscribeUpdateSlot};
 
@@ -298,7 +298,7 @@ pub fn create_block_processing_task(
     grpc_x_token: Option<String>,
     block_sx: async_channel::Sender<SubscribeUpdateBlock>,
     commitment_level: CommitmentLevel,
-    exit_notfier: CancellationToken,
+    exit_notfier: Arc<Notify>,
 ) -> AnyhowJoinHandle {
     tokio::spawn(async move {
         loop {
@@ -361,7 +361,7 @@ pub fn create_block_processing_task(
                             }
                         };
                     },
-                    _ = exit_notfier.cancelled() => {
+                    _ = exit_notfier.notified() => {
                         break;
                     }
                 }

diff --git a/services/Cargo.toml b/services/Cargo.toml
@@ -39,7 +39,6 @@ quinn = { workspace = true }
 chrono = { workspace = true }
 rustls = { workspace = true }
 solana-lite-rpc-core = { workspace = true }
-tokio-util = { workspace = true }
 
 [dev-dependencies]
 tracing = { workspace = true }

diff --git a/services/src/quic_connection.rs b/services/src/quic_connection.rs
@@ -14,8 +14,7 @@ use std::{
         Arc,
     },
 };
-use tokio::sync::{OwnedSemaphorePermit, RwLock, Semaphore};
-use tokio_util::sync::CancellationToken;
+use tokio::sync::{Notify, OwnedSemaphorePermit, RwLock, Semaphore};
 
 pub type EndpointPool = RotatingQueue<Endpoint>;
 
@@ -41,7 +40,7 @@ pub struct QuicConnection {
     identity: Pubkey,
     socket_address: SocketAddr,
     connection_params: QuicConnectionParameters,
-    exit_notify: CancellationToken,
+    exit_notify: Arc<Notify>,
     timeout_counters: Arc<AtomicU64>,
     has_connected_once: Arc<AtomicBool>,
 }
@@ -52,7 +51,7 @@ impl QuicConnection {
         endpoint: Endpoint,
         socket_address: SocketAddr,
         connection_params: QuicConnectionParameters,
-        exit_notify: CancellationToken,
+        exit_notify: Arc<Notify>,
     ) -> Self {
         Self {
             connection: Arc::new(RwLock::new(None)),
@@ -135,7 +134,7 @@ impl QuicConnection {
                 conn = self.get_connection() => {
                     conn
                 },
-                _ = exit_notify.cancelled() => {
+                _ = exit_notify.notified() => {
                     break;
                 }
             };
@@ -150,7 +149,7 @@ impl QuicConnection {
                     ) => {
                         res
                     },
-                    _ = exit_notify.cancelled() => {
+                    _ = exit_notify.notified() => {
                         break;
                     }
                 };
@@ -165,7 +164,7 @@ impl QuicConnection {
                             ) => {
                                 res
                             },
-                            _ = exit_notify.cancelled() => {
+                            _ = exit_notify.notified() => {
                                 break;
                             }
                         };
@@ -248,7 +247,7 @@ impl QuicConnectionPool {
         endpoints: EndpointPool,
         socket_address: SocketAddr,
         connection_parameters: QuicConnectionParameters,
-        exit_notify: CancellationToken,
+        exit_notify: Arc<Notify>,
         nb_connection: usize,
         max_number_of_unistream_connection: usize,
     ) -> Self {

diff --git a/services/src/quic_connection_utils.rs b/services/src/quic_connection_utils.rs
@@ -14,8 +14,7 @@ use std::{
     sync::Arc,
     time::Duration,
 };
-use tokio::time::timeout;
-use tokio_util::sync::CancellationToken;
+use tokio::{sync::Notify, time::timeout};
 
 lazy_static::lazy_static! {
     static ref NB_QUIC_0RTT_ATTEMPTED: GenericGauge<prometheus::core::AtomicI64> =
@@ -220,7 +219,7 @@ impl QuicConnectionUtils {
         addr: SocketAddr,
         connection_timeout: Duration,
         connection_retry_count: usize,
-        exit_notified: CancellationToken,
+        exit_notified: Arc<Notify>,
     ) -> Option<Connection> {
         for _ in 0..connection_retry_count {
             let conn = if already_connected {
@@ -229,7 +228,7 @@ impl QuicConnectionUtils {
                     res = Self::make_connection_0rtt(endpoint.clone(), addr, connection_timeout) => {
                         res
                     },
-                    _ = exit_notified.cancelled() => {
+                    _ = exit_notified.notified() => {
                         break;
                     }
                 }
@@ -239,7 +238,7 @@ impl QuicConnectionUtils {
                     res = Self::make_connection(endpoint.clone(), addr, connection_timeout) => {
                         res
                     },
-                    _ = exit_notified.cancelled() => {
+                    _ = exit_notified.notified() => {
                         break;
                     }
                 }

diff --git a/services/src/tpu_utils/tpu_connection_manager.rs b/services/src/tpu_utils/tpu_connection_manager.rs
@@ -18,7 +18,6 @@ use tokio::sync::{
     broadcast::{Receiver, Sender},
     Notify,
 };
-use tokio_util::sync::CancellationToken;
 
 use crate::{
     quic_connection::{PooledConnection, QuicConnectionPool},
@@ -50,7 +49,7 @@ struct ActiveConnection {
     tpu_address: SocketAddr,
     data_cache: DataCache,
     connection_parameters: QuicConnectionParameters,
-    exit_notifier: CancellationToken,
+    exit_notifier: Arc<Notify>,
 }
 
 impl ActiveConnection {
@@ -67,7 +66,7 @@ impl ActiveConnection {
             identity,
             data_cache,
             connection_parameters,
-            exit_notifier: CancellationToken::new(),
+            exit_notifier: Arc::new(Notify::new()),
         }
     }
 
@@ -117,7 +116,7 @@ impl ActiveConnection {
                         tx = transaction_reciever.recv() => {
                             tx
                         },
-                        _ = exit_notifier.cancelled() => {
+                        _ = exit_notifier.notified() => {
                             break;
                         }
                     };
@@ -210,7 +209,7 @@ impl ActiveConnection {
                         });
                     }
                 },
-                _ = exit_notifier.cancelled() => {
+                _ = exit_notifier.notified() => {
                     break 'main_loop;
                 }
             }
@@ -288,7 +287,7 @@ impl TpuConnectionManager {
             if !connections_to_keep.contains_key(key) {
                 trace!("removing a connection for {}", key.to_string());
                 // ignore error for exit channel
-                value.exit_notifier.cancel();
+                value.exit_notifier.notify_waiters();
                 false
             } else {
                 true