Allow TaskExecutor to be used in async tests (#3178)

# Description

Since the `TaskExecutor` currently requires a `Weak<Runtime>`, it's impossible to use it in an async test where the `Runtime` is created outside our scope. Whilst we *could* create a new `Runtime` instance inside the async test, dropping that `Runtime` would cause a panic (you can't drop a `Runtime` in an async context).

To address this issue, this PR creates the `enum Handle`, which supports either:

- A `Weak<Runtime>` (for use in our production code)
- A `Handle` to a runtime (for use in testing)

In theory, there should be no change to the behaviour of our production code (beyond some slightly different descriptions in HTTP 500 errors), or even our tests. If there is no change, you might ask *"why bother?"*. There are two PRs (#3070 and #3175) that are waiting on these fixes to introduce some new tests. Since we've added the EL to the `BeaconChain` (for the merge), we are now doing more async stuff in tests.

I've also added a `RuntimeExecutor` to the `BeaconChainTestHarness`. Whilst that's not immediately useful, it will become useful in the near future with all the new async testing.

beacon_node/beacon_chain/Cargo.toml

@@ -61,6 +61,7 @@ execution_layer = { path = "../execution_layer" }
 sensitive_url = { path = "../../common/sensitive_url" }
 superstruct = "0.5.0"
 hex = "0.4.2"
+exit-future = "0.2.0"
 
 [[test]]
 name = "beacon_chain_tests"

beacon_node/beacon_chain/src/builder.rs

@@ -27,7 +27,7 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 use std::time::Duration;
 use store::{Error as StoreError, HotColdDB, ItemStore, KeyValueStoreOp};
-use task_executor::ShutdownReason;
+use task_executor::{ShutdownReason, TaskExecutor};
 use types::{
     BeaconBlock, BeaconState, ChainSpec, Checkpoint, EthSpec, Graffiti, Hash256, PublicKeyBytes,
     Signature, SignedBeaconBlock, Slot,
@@ -91,6 +91,7 @@ pub struct BeaconChainBuilder<T: BeaconChainTypes> {
     // Pending I/O batch that is constructed during building and should be executed atomically
     // alongside `PersistedBeaconChain` storage when `BeaconChainBuilder::build` is called.
     pending_io_batch: Vec<KeyValueStoreOp>,
+    task_executor: Option<TaskExecutor>,
 }
 
 impl<TSlotClock, TEth1Backend, TEthSpec, THotStore, TColdStore>
@@ -129,6 +130,7 @@ where
             slasher: None,
             validator_monitor: None,
             pending_io_batch: vec![],
+            task_executor: None,
         }
     }
 
@@ -182,6 +184,13 @@ where
         self.log = Some(log);
         self
     }
+
+    /// Sets the task executor.
+    pub fn task_executor(mut self, task_executor: TaskExecutor) -> Self {
+        self.task_executor = Some(task_executor);
+        self
+    }
+
     /// Attempt to load an existing eth1 cache from the builder's `Store`.
     pub fn get_persisted_eth1_backend(&self) -> Result<Option<SszEth1>, String> {
         let store = self
@@ -919,6 +928,7 @@ mod test {
     use std::time::Duration;
     use store::config::StoreConfig;
     use store::{HotColdDB, MemoryStore};
+    use task_executor::test_utils::TestRuntime;
     use types::{EthSpec, MinimalEthSpec, Slot};
 
     type TestEthSpec = MinimalEthSpec;
@@ -952,10 +962,12 @@ mod test {
         .expect("should create interop genesis state");
 
         let (shutdown_tx, _) = futures::channel::mpsc::channel(1);
+        let runtime = TestRuntime::default();
 
         let chain = BeaconChainBuilder::new(MinimalEthSpec)
             .logger(log.clone())
             .store(Arc::new(store))
+            .task_executor(runtime.task_executor.clone())
             .genesis_state(genesis_state)
             .expect("should build state using recent genesis")
             .dummy_eth1_backend()

beacon_node/beacon_chain/src/test_utils.rs

@@ -12,15 +12,12 @@ use crate::{
 };
 use bls::get_withdrawal_credentials;
 use execution_layer::{
-    test_utils::{
-        ExecutionBlockGenerator, ExecutionLayerRuntime, MockExecutionLayer, DEFAULT_TERMINAL_BLOCK,
-    },
+    test_utils::{ExecutionBlockGenerator, MockExecutionLayer, DEFAULT_TERMINAL_BLOCK},
     ExecutionLayer,
 };
 use futures::channel::mpsc::Receiver;
 pub use genesis::{interop_genesis_state, DEFAULT_ETH1_BLOCK_HASH};
 use int_to_bytes::int_to_bytes32;
-use logging::test_logger;
 use merkle_proof::MerkleTree;
 use parking_lot::Mutex;
 use parking_lot::RwLockWriteGuard;
@@ -41,7 +38,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
 use store::{config::StoreConfig, HotColdDB, ItemStore, LevelDB, MemoryStore};
-use task_executor::ShutdownReason;
+use task_executor::{test_utils::TestRuntime, ShutdownReason};
 use tree_hash::TreeHash;
 use types::sync_selection_proof::SyncSelectionProof;
 pub use types::test_utils::generate_deterministic_keypairs;
@@ -151,8 +148,8 @@ pub struct Builder<T: BeaconChainTypes> {
     initial_mutator: Option<BoxedMutator<T::EthSpec, T::HotStore, T::ColdStore>>,
     store_mutator: Option<BoxedMutator<T::EthSpec, T::HotStore, T::ColdStore>>,
     execution_layer: Option<ExecutionLayer>,
-    execution_layer_runtime: Option<ExecutionLayerRuntime>,
     mock_execution_layer: Option<MockExecutionLayer<T::EthSpec>>,
+    runtime: TestRuntime,
     log: Logger,
 }
 
@@ -255,6 +252,9 @@ where
     Cold: ItemStore<E>,
 {
     pub fn new(eth_spec_instance: E) -> Self {
+        let runtime = TestRuntime::default();
+        let log = runtime.log.clone();
+
         Self {
             eth_spec_instance,
             spec: None,
@@ -266,8 +266,8 @@ where
             store_mutator: None,
             execution_layer: None,
             mock_execution_layer: None,
-            execution_layer_runtime: None,
-            log: test_logger(),
+            runtime,
+            log,
         }
     }
 
@@ -330,8 +330,6 @@ where
             "execution layer already defined"
         );
 
-        let el_runtime = ExecutionLayerRuntime::default();
-
         let urls: Vec<SensitiveUrl> = urls
             .iter()
             .map(|s| SensitiveUrl::parse(*s))
@@ -346,19 +344,19 @@ where
         };
         let execution_layer = ExecutionLayer::from_config(
             config,
-            el_runtime.task_executor.clone(),
-            el_runtime.log.clone(),
+            self.runtime.task_executor.clone(),
+            self.log.clone(),
         )
         .unwrap();
 
         self.execution_layer = Some(execution_layer);
-        self.execution_layer_runtime = Some(el_runtime);
         self
     }
 
     pub fn mock_execution_layer(mut self) -> Self {
         let spec = self.spec.clone().expect("cannot build without spec");
         let mock = MockExecutionLayer::new(
+            self.runtime.task_executor.clone(),
             spec.terminal_total_difficulty,
             DEFAULT_TERMINAL_BLOCK,
             spec.terminal_block_hash,
@@ -383,7 +381,7 @@ where
     pub fn build(self) -> BeaconChainHarness<BaseHarnessType<E, Hot, Cold>> {
         let (shutdown_tx, shutdown_receiver) = futures::channel::mpsc::channel(1);
 
-        let log = test_logger();
+        let log = self.log;
         let spec = self.spec.expect("cannot build without spec");
         let seconds_per_slot = spec.seconds_per_slot;
         let validator_keypairs = self
@@ -395,6 +393,7 @@ where
             .custom_spec(spec)
             .store(self.store.expect("cannot build without store"))
             .store_migrator_config(MigratorConfig::default().blocking())
+            .task_executor(self.runtime.task_executor.clone())
             .execution_layer(self.execution_layer)
             .dummy_eth1_backend()
             .expect("should build dummy backend")
@@ -434,8 +433,8 @@ where
             chain: Arc::new(chain),
             validator_keypairs,
             shutdown_receiver: Arc::new(Mutex::new(shutdown_receiver)),
+            runtime: self.runtime,
             mock_execution_layer: self.mock_execution_layer,
-            execution_layer_runtime: self.execution_layer_runtime,
             rng: make_rng(),
         }
     }
@@ -451,9 +450,9 @@ pub struct BeaconChainHarness<T: BeaconChainTypes> {
     pub chain: Arc<BeaconChain<T>>,
     pub spec: ChainSpec,
     pub shutdown_receiver: Arc<Mutex<Receiver<ShutdownReason>>>,
+    pub runtime: TestRuntime,
 
     pub mock_execution_layer: Option<MockExecutionLayer<T::EthSpec>>,
-    pub execution_layer_runtime: Option<ExecutionLayerRuntime>,
 
     pub rng: Mutex<StdRng>,
 }

beacon_node/client/src/builder.rs

@@ -166,6 +166,7 @@ where
         let builder = BeaconChainBuilder::new(eth_spec_instance)
             .logger(context.log().clone())
             .store(store)
+            .task_executor(context.executor.clone())
             .custom_spec(spec.clone())
             .chain_config(chain_config)
             .graffiti(graffiti)

beacon_node/execution_layer/src/lib.rs

@@ -304,11 +304,7 @@ impl ExecutionLayer {
         T: Fn(&'a Self) -> U,
         U: Future<Output = Result<V, Error>>,
     {
-        let runtime = self
-            .executor()
-            .runtime()
-            .upgrade()
-            .ok_or(Error::ShuttingDown)?;
+        let runtime = self.executor().handle().ok_or(Error::ShuttingDown)?;
         // TODO(merge): respect the shutdown signal.
         runtime.block_on(generate_future(self))
     }
@@ -322,11 +318,7 @@ impl ExecutionLayer {
         T: Fn(&'a Self) -> U,
         U: Future<Output = V>,
     {
-        let runtime = self
-            .executor()
-            .runtime()
-            .upgrade()
-            .ok_or(Error::ShuttingDown)?;
+        let runtime = self.executor().handle().ok_or(Error::ShuttingDown)?;
         // TODO(merge): respect the shutdown signal.
         Ok(runtime.block_on(generate_future(self)))
     }
@@ -1263,13 +1255,15 @@ impl ExecutionLayer {
 mod test {
     use super::*;
     use crate::test_utils::MockExecutionLayer as GenericMockExecutionLayer;
+    use task_executor::test_utils::TestRuntime;
     use types::MainnetEthSpec;
 
     type MockExecutionLayer = GenericMockExecutionLayer<MainnetEthSpec>;
 
     #[tokio::test]
     async fn produce_three_valid_pos_execution_blocks() {
-        MockExecutionLayer::default_params()
+        let runtime = TestRuntime::default();
+        MockExecutionLayer::default_params(runtime.task_executor.clone())
             .move_to_terminal_block()
             .produce_valid_execution_payload_on_head()
             .await
@@ -1281,7 +1275,8 @@ mod test {
 
     #[tokio::test]
     async fn finds_valid_terminal_block_hash() {
-        MockExecutionLayer::default_params()
+        let runtime = TestRuntime::default();
+        MockExecutionLayer::default_params(runtime.task_executor.clone())
             .move_to_block_prior_to_terminal_block()
             .with_terminal_block(|spec, el, _| async move {
                 el.engines().upcheck_not_synced(Logging::Disabled).await;
@@ -1300,7 +1295,8 @@ mod test {
 
     #[tokio::test]
     async fn verifies_valid_terminal_block_hash() {
-        MockExecutionLayer::default_params()
+        let runtime = TestRuntime::default();
+        MockExecutionLayer::default_params(runtime.task_executor.clone())
             .move_to_terminal_block()
             .with_terminal_block(|spec, el, terminal_block| async move {
                 el.engines().upcheck_not_synced(Logging::Disabled).await;
@@ -1316,7 +1312,8 @@ mod test {
 
     #[tokio::test]
     async fn rejects_invalid_terminal_block_hash() {
-        MockExecutionLayer::default_params()
+        let runtime = TestRuntime::default();
+        MockExecutionLayer::default_params(runtime.task_executor.clone())
             .move_to_terminal_block()
             .with_terminal_block(|spec, el, terminal_block| async move {
                 el.engines().upcheck_not_synced(Logging::Disabled).await;
@@ -1334,7 +1331,8 @@ mod test {
 
     #[tokio::test]
     async fn rejects_unknown_terminal_block_hash() {
-        MockExecutionLayer::default_params()
+        let runtime = TestRuntime::default();
+        MockExecutionLayer::default_params(runtime.task_executor.clone())
             .move_to_terminal_block()
             .with_terminal_block(|spec, el, _| async move {
                 el.engines().upcheck_not_synced(Logging::Disabled).await;

beacon_node/execution_layer/src/test_utils/mock_execution_layer.rs

@@ -2,61 +2,22 @@ use crate::{
     test_utils::{MockServer, DEFAULT_TERMINAL_BLOCK, DEFAULT_TERMINAL_DIFFICULTY, JWT_SECRET},
     Config, *,
 };
-use environment::null_logger;
 use sensitive_url::SensitiveUrl;
-use std::sync::Arc;
 use task_executor::TaskExecutor;
 use tempfile::NamedTempFile;
 use types::{Address, ChainSpec, Epoch, EthSpec, FullPayload, Hash256, Uint256};
 
-pub struct ExecutionLayerRuntime {
-    pub runtime: Option<Arc<tokio::runtime::Runtime>>,
-    pub _runtime_shutdown: exit_future::Signal,
-    pub task_executor: TaskExecutor,
-    pub log: Logger,
-}
-
-impl Default for ExecutionLayerRuntime {
-    fn default() -> Self {
-        let runtime = Arc::new(
-            tokio::runtime::Builder::new_multi_thread()
-                .enable_all()
-                .build()
-                .unwrap(),
-        );
-        let (runtime_shutdown, exit) = exit_future::signal();
-        let (shutdown_tx, _) = futures::channel::mpsc::channel(1);
-        let log = null_logger().unwrap();
-        let task_executor =
-            TaskExecutor::new(Arc::downgrade(&runtime), exit, log.clone(), shutdown_tx);
-
-        Self {
-            runtime: Some(runtime),
-            _runtime_shutdown: runtime_shutdown,
-            task_executor,
-            log,
-        }
-    }
-}
-
-impl Drop for ExecutionLayerRuntime {
-    fn drop(&mut self) {
-        if let Some(runtime) = self.runtime.take() {
-            Arc::try_unwrap(runtime).unwrap().shutdown_background()
-        }
-    }
-}
-
 pub struct MockExecutionLayer<T: EthSpec> {
     pub server: MockServer<T>,
     pub el: ExecutionLayer,
-    pub el_runtime: ExecutionLayerRuntime,
+    pub executor: TaskExecutor,
     pub spec: ChainSpec,
 }
 
 impl<T: EthSpec> MockExecutionLayer<T> {
-    pub fn default_params() -> Self {
+    pub fn default_params(executor: TaskExecutor) -> Self {
         Self::new(
+            executor,
             DEFAULT_TERMINAL_DIFFICULTY.into(),
             DEFAULT_TERMINAL_BLOCK,
             ExecutionBlockHash::zero(),
@@ -65,13 +26,13 @@ impl<T: EthSpec> MockExecutionLayer<T> {
     }
 
     pub fn new(
+        executor: TaskExecutor,
         terminal_total_difficulty: Uint256,
         terminal_block: u64,
         terminal_block_hash: ExecutionBlockHash,
         terminal_block_hash_activation_epoch: Epoch,
     ) -> Self {
-        let el_runtime = ExecutionLayerRuntime::default();
-        let handle = el_runtime.runtime.as_ref().unwrap().handle();
+        let handle = executor.handle().unwrap();
 
         let mut spec = T::default_spec();
         spec.terminal_total_difficulty = terminal_total_difficulty;
@@ -79,7 +40,7 @@ impl<T: EthSpec> MockExecutionLayer<T> {
         spec.terminal_block_hash_activation_epoch = terminal_block_hash_activation_epoch;
 
         let server = MockServer::new(
-            handle,
+            &handle,
             terminal_total_difficulty,
             terminal_block,
             terminal_block_hash,
@@ -97,17 +58,13 @@ impl<T: EthSpec> MockExecutionLayer<T> {
             suggested_fee_recipient: Some(Address::repeat_byte(42)),
             ..Default::default()
         };
-        let el = ExecutionLayer::from_config(
-            config,
-            el_runtime.task_executor.clone(),
-            el_runtime.log.clone(),
-        )
-        .unwrap();
+        let el =
+            ExecutionLayer::from_config(config, executor.clone(), executor.log().clone()).unwrap();
 
         Self {
             server,
             el,
-            el_runtime,
+            executor,
             spec,
         }
     }

beacon_node/execution_layer/src/test_utils/mod.rs

@@ -22,7 +22,7 @@ use types::{EthSpec, ExecutionBlockHash, Uint256};
 use warp::{http::StatusCode, Filter, Rejection};
 
 pub use execution_block_generator::{generate_pow_block, ExecutionBlockGenerator};
-pub use mock_execution_layer::{ExecutionLayerRuntime, MockExecutionLayer};
+pub use mock_execution_layer::MockExecutionLayer;
 
 pub const DEFAULT_TERMINAL_DIFFICULTY: u64 = 6400;
 pub const DEFAULT_TERMINAL_BLOCK: u64 = 64;

beacon_node/http_api/Cargo.toml

@@ -30,6 +30,7 @@ futures = "0.3.8"
 execution_layer = {path = "../execution_layer"}
 parking_lot = "0.12.0"
 safe_arith = {path = "../../consensus/safe_arith"}
+task_executor = { path = "../../common/task_executor" }
 
 
 [dev-dependencies]

beacon_node/http_api/tests/tests.rs

@@ -20,6 +20,7 @@ use slot_clock::SlotClock;
 use state_processing::per_slot_processing;
 use std::convert::TryInto;
 use std::sync::Arc;
+use task_executor::test_utils::TestRuntime;
 use tokio::sync::{mpsc, oneshot};
 use tokio::time::Duration;
 use tree_hash::TreeHash;
@@ -63,6 +64,7 @@ struct ApiTester {
     network_rx: mpsc::UnboundedReceiver<NetworkMessage<E>>,
     local_enr: Enr,
     external_peer_id: PeerId,
+    _runtime: TestRuntime,
 }
 
 impl ApiTester {
@@ -185,7 +187,7 @@ impl ApiTester {
             external_peer_id,
         } = create_api_server(chain.clone(), log).await;
 
-        tokio::spawn(server);
+        harness.runtime.task_executor.spawn(server, "api_server");
 
         let client = BeaconNodeHttpClient::new(
             SensitiveUrl::parse(&format!(
@@ -212,6 +214,7 @@ impl ApiTester {
             network_rx,
             local_enr,
             external_peer_id,
+            _runtime: harness.runtime,
         }
     }
 
@@ -263,7 +266,7 @@ impl ApiTester {
             external_peer_id,
         } = create_api_server(chain.clone(), log).await;
 
-        tokio::spawn(server);
+        harness.runtime.task_executor.spawn(server, "api_server");
 
         let client = BeaconNodeHttpClient::new(
             SensitiveUrl::parse(&format!(
@@ -290,6 +293,7 @@ impl ApiTester {
             network_rx,
             local_enr,
             external_peer_id,
+            _runtime: harness.runtime,
         }
     }
 

beacon_node/network/src/beacon_processor/tests.rs

@@ -20,7 +20,7 @@ use std::cmp;
 use std::iter::Iterator;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::runtime::Runtime;
+use tokio::runtime::Handle;
 use tokio::sync::mpsc;
 use types::{
     Attestation, AttesterSlashing, EthSpec, MainnetEthSpec, ProposerSlashing, SignedBeaconBlock,
@@ -324,20 +324,19 @@ impl TestRig {
             .unwrap();
     }
 
-    fn runtime(&mut self) -> Arc<Runtime> {
+    fn handle(&mut self) -> Handle {
         self.environment
             .as_mut()
             .unwrap()
             .core_context()
             .executor
-            .runtime()
-            .upgrade()
+            .handle()
             .unwrap()
     }
 
     /// Assert that the `BeaconProcessor` doesn't produce any events in the given `duration`.
     pub fn assert_no_events_for(&mut self, duration: Duration) {
-        self.runtime().block_on(async {
+        self.handle().block_on(async {
             tokio::select! {
                 _ = tokio::time::sleep(duration) => (),
                 event = self.work_journal_rx.recv() => panic!(
@@ -360,7 +359,7 @@ impl TestRig {
             .iter()
             .all(|ev| ev != &WORKER_FREED && ev != &NOTHING_TO_DO));
 
-        let (events, worker_freed_remaining) = self.runtime().block_on(async {
+        let (events, worker_freed_remaining) = self.handle().block_on(async {
             let mut events = Vec::with_capacity(expected.len());
             let mut worker_freed_remaining = expected.len();
 
@@ -415,7 +414,7 @@ impl TestRig {
     /// We won't attempt to listen for any more than `expected.len()` events. As such, it makes sense
     /// to use the `NOTHING_TO_DO` event to ensure that execution has completed.
     pub fn assert_event_journal_with_timeout(&mut self, expected: &[&str], timeout: Duration) {
-        let events = self.runtime().block_on(async {
+        let events = self.handle().block_on(async {
             let mut events = Vec::with_capacity(expected.len());
 
             let drain_future = async {