Improve tokio task execution (#1181)

* Add logging on shutdown

* Replace tokio::spawn with handle.spawn

* Upgrade tokio

* Add a task executor

* Beacon chain tasks use task executor

* Validator client tasks use task executor

* Rename runtime_handle to executor

* Add duration histograms; minor fixes

* Cleanup

* Fix logs

* Fix tests

* Remove random file

* Get enr dependency instead of libp2p

* Address some review comments

* Libp2p takes a TaskExecutor

* Ugly fix libp2p tests

* Move TaskExecutor to own file

* Upgrade Dockerfile rust version

* Minor fixes

* Revert "Ugly fix libp2p tests"

This reverts commit 58d4bb690f.

* Pretty fix libp2p tests

* Add spawn_without_exit; change Counter to Gauge

* Tidy

* Move log from RuntimeContext to TaskExecutor

* Fix errors

* Replace histogram with int_gauge for async tasks

* Fix todo

* Fix memory leak in test by exiting all spawned tasks at the end

beacon_node/network/Cargo.toml

@@ -10,6 +10,7 @@ genesis = { path = "../genesis" }
 lazy_static = "1.4.0"
 matches = "0.1.8"
 tempfile = "3.1.0"
+exit-future = "0.2.0"
 
 [dependencies]
 beacon_chain =  { path = "../beacon_chain" }
@@ -25,7 +26,7 @@ eth2_ssz = "0.1.2"
 tree_hash = "0.1.0"
 futures = "0.3.5"
 error-chain = "0.12.2"
-tokio = { version = "0.2.20", features = ["full"] }
+tokio = { version = "0.2.21", features = ["full"] }
 parking_lot = "0.10.2"
 smallvec = "1.4.0"
 # TODO: Remove rand crate for mainnet
@@ -34,3 +35,4 @@ fnv = "1.0.6"
 rlp = "0.4.5"
 lazy_static = "1.4.0"
 lighthouse_metrics = { path = "../../common/lighthouse_metrics" }
+environment = { path = "../../lighthouse/environment" }

beacon_node/network/src/router/mod.rs

@@ -58,7 +58,7 @@ impl<T: BeaconChainTypes> Router<T> {
         beacon_chain: Arc<BeaconChain<T>>,
         network_globals: Arc<NetworkGlobals<T::EthSpec>>,
         network_send: mpsc::UnboundedSender<NetworkMessage<T::EthSpec>>,
-        runtime_handle: &tokio::runtime::Handle,
+        executor: environment::TaskExecutor,
         log: slog::Logger,
     ) -> error::Result<mpsc::UnboundedSender<RouterMessage<T::EthSpec>>> {
         let message_handler_log = log.new(o!("service"=> "router"));
@@ -68,7 +68,7 @@ impl<T: BeaconChainTypes> Router<T> {
 
         // Initialise a message instance, which itself spawns the syncing thread.
         let processor = Processor::new(
-            runtime_handle,
+            executor.clone(),
             beacon_chain,
             network_globals.clone(),
             network_send.clone(),
@@ -84,12 +84,15 @@ impl<T: BeaconChainTypes> Router<T> {
         };
 
         // spawn handler task and move the message handler instance into the spawned thread
-        runtime_handle.spawn(async move {
-            handler_recv
-                .for_each(move |msg| future::ready(handler.handle_message(msg)))
-                .await;
-            debug!(log, "Network message handler terminated.");
-        });
+        executor.spawn(
+            async move {
+                debug!(log, "Network message router started");
+                handler_recv
+                    .for_each(move |msg| future::ready(handler.handle_message(msg)))
+                    .await;
+            },
+            "router",
+        );
 
         Ok(handler_send)
     }

beacon_node/network/src/router/processor.rs

@@ -14,7 +14,7 @@ use slog::{debug, error, o, trace, warn};
 use ssz::Encode;
 use std::sync::Arc;
 use store::Store;
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::mpsc;
 use types::{
     Attestation, ChainSpec, Epoch, EthSpec, Hash256, SignedAggregateAndProof, SignedBeaconBlock,
     Slot,
@@ -33,8 +33,6 @@ pub struct Processor<T: BeaconChainTypes> {
     chain: Arc<BeaconChain<T>>,
     /// A channel to the syncing thread.
     sync_send: mpsc::UnboundedSender<SyncMessage<T::EthSpec>>,
-    /// A oneshot channel for destroying the sync thread.
-    _sync_exit: oneshot::Sender<()>,
     /// A network context to return and handle RPC requests.
     network: HandlerNetworkContext<T::EthSpec>,
     /// The `RPCHandler` logger.
@@ -44,7 +42,7 @@ pub struct Processor<T: BeaconChainTypes> {
 impl<T: BeaconChainTypes> Processor<T> {
     /// Instantiate a `Processor` instance
     pub fn new(
-        runtime_handle: &tokio::runtime::Handle,
+        executor: environment::TaskExecutor,
         beacon_chain: Arc<BeaconChain<T>>,
         network_globals: Arc<NetworkGlobals<T::EthSpec>>,
         network_send: mpsc::UnboundedSender<NetworkMessage<T::EthSpec>>,
@@ -53,8 +51,8 @@ impl<T: BeaconChainTypes> Processor<T> {
         let sync_logger = log.new(o!("service"=> "sync"));
 
         // spawn the sync thread
-        let (sync_send, _sync_exit) = crate::sync::manager::spawn(
-            runtime_handle,
+        let sync_send = crate::sync::manager::spawn(
+            executor,
             beacon_chain.clone(),
             network_globals,
             network_send.clone(),
@@ -64,7 +62,6 @@ impl<T: BeaconChainTypes> Processor<T> {
         Processor {
             chain: beacon_chain,
             sync_send,
-            _sync_exit,
             network: HandlerNetworkContext::new(network_send, log.clone()),
             log: log.clone(),
         }

beacon_node/network/src/service.rs

@@ -14,8 +14,7 @@ use rest_types::ValidatorSubscription;
 use slog::{debug, error, info, o, trace};
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::runtime::Handle;
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::mpsc;
 use tokio::time::Delay;
 use types::EthSpec;
 
@@ -53,13 +52,12 @@ impl<T: BeaconChainTypes> NetworkService<T> {
     pub fn start(
         beacon_chain: Arc<BeaconChain<T>>,
         config: &NetworkConfig,
-        runtime_handle: &Handle,
-        network_log: slog::Logger,
+        executor: environment::TaskExecutor,
     ) -> error::Result<(
         Arc<NetworkGlobals<T::EthSpec>>,
         mpsc::UnboundedSender<NetworkMessage<T::EthSpec>>,
-        oneshot::Sender<()>,
     )> {
+        let network_log = executor.log().clone();
         // build the network channel
         let (network_send, network_recv) = mpsc::unbounded_channel::<NetworkMessage<T::EthSpec>>();
         // get a reference to the beacon chain store
@@ -75,7 +73,7 @@ impl<T: BeaconChainTypes> NetworkService<T> {
 
         // launch libp2p service
         let (network_globals, mut libp2p) =
-            runtime_handle.enter(|| LibP2PService::new(config, enr_fork_id, &network_log))?;
+            LibP2PService::new(executor.clone(), config, enr_fork_id, &network_log)?;
 
         for enr in load_dht::<T::Store, T::EthSpec>(store.clone()) {
             libp2p.swarm.add_enr(enr);
@@ -88,7 +86,7 @@ impl<T: BeaconChainTypes> NetworkService<T> {
             beacon_chain.clone(),
             network_globals.clone(),
             network_send.clone(),
-            runtime_handle,
+            executor.clone(),
             network_log.clone(),
         )?;
 
@@ -111,19 +109,20 @@ impl<T: BeaconChainTypes> NetworkService<T> {
             propagation_percentage,
         };
 
-        let network_exit = runtime_handle.enter(|| spawn_service(network_service))?;
+        spawn_service(executor, network_service)?;
 
-        Ok((network_globals, network_send, network_exit))
+        Ok((network_globals, network_send))
     }
 }
 
 fn spawn_service<T: BeaconChainTypes>(
+    executor: environment::TaskExecutor,
     mut service: NetworkService<T>,
-) -> error::Result<tokio::sync::oneshot::Sender<()>> {
-    let (network_exit, mut exit_rx) = tokio::sync::oneshot::channel();
+) -> error::Result<()> {
+    let mut exit_rx = executor.exit();
 
     // spawn on the current executor
-    tokio::spawn(async move {
+    executor.spawn_without_exit(async move {
         loop {
             // build the futures to check simultaneously
             tokio::select! {
@@ -361,9 +360,9 @@ fn spawn_service<T: BeaconChainTypes>(
                 }
             }
         }
-    });
+    }, "network");
 
-    Ok(network_exit)
+    Ok(())
 }
 
 /// Returns a `Delay` that triggers shortly after the next change in the beacon chain fork version.

beacon_node/network/src/service/tests.rs

@@ -32,7 +32,9 @@ mod tests {
         let enrs = vec![enr1, enr2];
 
         let runtime = Runtime::new().unwrap();
-        let handle = runtime.handle().clone();
+
+        let (signal, exit) = exit_future::signal();
+        let executor = environment::TaskExecutor::new(runtime.handle().clone(), exit, log.clone());
 
         let mut config = NetworkConfig::default();
         config.libp2p_port = 21212;
@@ -42,8 +44,8 @@ mod tests {
             // Create a new network service which implicitly gets dropped at the
             // end of the block.
 
-            let _ =
-                NetworkService::start(beacon_chain.clone(), &config, &handle, log.clone()).unwrap();
+            let _ = NetworkService::start(beacon_chain.clone(), &config, executor).unwrap();
+            drop(signal);
         });
         runtime.shutdown_timeout(tokio::time::Duration::from_millis(300));
 

beacon_node/network/src/sync/manager.rs

@@ -48,7 +48,7 @@ use smallvec::SmallVec;
 use std::boxed::Box;
 use std::ops::Sub;
 use std::sync::Arc;
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::mpsc;
 use types::{EthSpec, Hash256, SignedBeaconBlock, Slot};
 
 /// The number of slots ahead of us that is allowed before requesting a long-range (batch)  Sync
@@ -181,17 +181,12 @@ impl SingleBlockRequest {
 /// chain. This allows the chain to be
 /// dropped during the syncing process which will gracefully end the `SyncManager`.
 pub fn spawn<T: BeaconChainTypes>(
-    runtime_handle: &tokio::runtime::Handle,
+    executor: environment::TaskExecutor,
     beacon_chain: Arc<BeaconChain<T>>,
     network_globals: Arc<NetworkGlobals<T::EthSpec>>,
     network_send: mpsc::UnboundedSender<NetworkMessage<T::EthSpec>>,
     log: slog::Logger,
-) -> (
-    mpsc::UnboundedSender<SyncMessage<T::EthSpec>>,
-    oneshot::Sender<()>,
-) {
-    // generate the exit channel
-    let (sync_exit, exit_rx) = tokio::sync::oneshot::channel();
+) -> mpsc::UnboundedSender<SyncMessage<T::EthSpec>> {
     // generate the message channel
     let (sync_send, sync_recv) = mpsc::unbounded_channel::<SyncMessage<T::EthSpec>>();
 
@@ -215,11 +210,8 @@ pub fn spawn<T: BeaconChainTypes>(
 
     // spawn the sync manager thread
     debug!(log, "Sync Manager started");
-    runtime_handle.spawn(async move {
-        futures::future::select(Box::pin(sync_manager.main()), exit_rx).await;
-        info!(log.clone(), "Sync Manager shutdown");
-    });
-    (sync_send, sync_exit)
+    executor.spawn(async move { Box::pin(sync_manager.main()).await }, "sync");
+    sync_send
 }
 
 impl<T: BeaconChainTypes> SyncManager<T> {