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.

common/task_executor/Cargo.toml

@@ -5,9 +5,10 @@ authors = ["Sigma Prime <contact@sigmaprime.io>"]
 edition = "2021"
 
 [dependencies]
-tokio = { version = "1.14.0", features = ["rt"] }
+tokio = { version = "1.14.0", features = ["rt-multi-thread"] }
 slog = "2.5.2"
 futures = "0.3.7"
 exit-future = "0.2.0"
 lazy_static = "1.4.0"
 lighthouse_metrics = { path = "../lighthouse_metrics" }
+sloggers = { version = "2.1.1", features = ["json"] }

common/task_executor/src/lib.rs

@@ -1,10 +1,11 @@
 mod metrics;
+pub mod test_utils;
 
 use futures::channel::mpsc::Sender;
 use futures::prelude::*;
 use slog::{crit, debug, o, trace};
 use std::sync::Weak;
-use tokio::runtime::Runtime;
+use tokio::runtime::{Handle, Runtime};
 
 /// Provides a reason when Lighthouse is shut down.
 #[derive(Copy, Clone, Debug, PartialEq)]
@@ -24,11 +25,51 @@ impl ShutdownReason {
     }
 }
 
+/// Provides a `Handle` by either:
+///
+/// 1. Holding a `Weak<Runtime>` and calling `Runtime::handle`.
+/// 2. Directly holding a `Handle` and cloning it.
+///
+/// This enum allows the `TaskExecutor` to work in production where a `Weak<Runtime>` is directly
+/// accessible and in testing where the `Runtime` is hidden outside our scope.
+#[derive(Clone)]
+pub enum HandleProvider {
+    Runtime(Weak<Runtime>),
+    Handle(Handle),
+}
+
+impl From<Handle> for HandleProvider {
+    fn from(handle: Handle) -> Self {
+        HandleProvider::Handle(handle)
+    }
+}
+
+impl From<Weak<Runtime>> for HandleProvider {
+    fn from(weak_runtime: Weak<Runtime>) -> Self {
+        HandleProvider::Runtime(weak_runtime)
+    }
+}
+
+impl HandleProvider {
+    /// Returns a `Handle` to a `Runtime`.
+    ///
+    /// May return `None` if the weak reference to the `Runtime` has been dropped (this generally
+    /// means Lighthouse is shutting down).
+    pub fn handle(&self) -> Option<Handle> {
+        match self {
+            HandleProvider::Runtime(weak_runtime) => weak_runtime
+                .upgrade()
+                .map(|runtime| runtime.handle().clone()),
+            HandleProvider::Handle(handle) => Some(handle.clone()),
+        }
+    }
+}
+
 /// A wrapper over a runtime handle which can spawn async and blocking tasks.
 #[derive(Clone)]
 pub struct TaskExecutor {
     /// The handle to the runtime on which tasks are spawned
-    runtime: Weak<Runtime>,
+    handle_provider: HandleProvider,
     /// The receiver exit future which on receiving shuts down the task
     exit: exit_future::Exit,
     /// Sender given to tasks, so that if they encounter a state in which execution cannot
@@ -43,16 +84,19 @@ pub struct TaskExecutor {
 impl TaskExecutor {
     /// Create a new task executor.
     ///
-    /// Note: this function is mainly useful in tests. A `TaskExecutor` should be normally obtained from
-    /// a [`RuntimeContext`](struct.RuntimeContext.html)
-    pub fn new(
-        runtime: Weak<Runtime>,
+    /// ## Note
+    ///
+    /// This function should only be used during testing. In production, prefer to obtain an
+    /// instance of `Self` via a `environment::RuntimeContext` (see the `lighthouse/environment`
+    /// crate).
+    pub fn new<T: Into<HandleProvider>>(
+        handle: T,
         exit: exit_future::Exit,
         log: slog::Logger,
         signal_tx: Sender<ShutdownReason>,
     ) -> Self {
         Self {
-            runtime,
+            handle_provider: handle.into(),
             exit,
             signal_tx,
             log,
@@ -62,7 +106,7 @@ impl TaskExecutor {
     /// Clones the task executor adding a service name.
     pub fn clone_with_name(&self, service_name: String) -> Self {
         TaskExecutor {
-            runtime: self.runtime.clone(),
+            handle_provider: self.handle_provider.clone(),
             exit: self.exit.clone(),
             signal_tx: self.signal_tx.clone(),
             log: self.log.new(o!("service" => service_name)),
@@ -94,8 +138,8 @@ impl TaskExecutor {
         let mut shutdown_sender = self.shutdown_sender();
         let log = self.log.clone();
 
-        if let Some(runtime) = self.runtime.upgrade() {
-            runtime.spawn(async move {
+        if let Some(handle) = self.handle() {
+            handle.spawn(async move {
                 let timer = metrics::start_timer_vec(&metrics::TASKS_HISTOGRAM, &[name]);
                 if let Err(join_error) = task_handle.await {
                     if let Ok(panic) = join_error.try_into_panic() {
@@ -160,8 +204,8 @@ impl TaskExecutor {
             });
 
             int_gauge.inc();
-            if let Some(runtime) = self.runtime.upgrade() {
-                runtime.spawn(future);
+            if let Some(handle) = self.handle() {
+                handle.spawn(future);
             } else {
                 debug!(self.log, "Couldn't spawn task. Runtime shutting down");
             }
@@ -211,8 +255,8 @@ impl TaskExecutor {
             });
 
             int_gauge.inc();
-            if let Some(runtime) = self.runtime.upgrade() {
-                Some(runtime.spawn(future))
+            if let Some(handle) = self.handle() {
+                Some(handle.spawn(future))
             } else {
                 debug!(self.log, "Couldn't spawn task. Runtime shutting down");
                 None
@@ -242,8 +286,8 @@ impl TaskExecutor {
         let timer = metrics::start_timer_vec(&metrics::BLOCKING_TASKS_HISTOGRAM, &[name]);
         metrics::inc_gauge_vec(&metrics::BLOCKING_TASKS_COUNT, &[name]);
 
-        let join_handle = if let Some(runtime) = self.runtime.upgrade() {
-            runtime.spawn_blocking(task)
+        let join_handle = if let Some(handle) = self.handle() {
+            handle.spawn_blocking(task)
         } else {
             debug!(self.log, "Couldn't spawn task. Runtime shutting down");
             return None;
@@ -268,8 +312,9 @@ impl TaskExecutor {
         Some(future)
     }
 
-    pub fn runtime(&self) -> Weak<Runtime> {
-        self.runtime.clone()
+    /// Returns a `Handle` to the current runtime.
+    pub fn handle(&self) -> Option<Handle> {
+        self.handle_provider.handle()
     }
 
     /// Returns a copy of the `exit_future::Exit`.

common/task_executor/src/test_utils.rs

@@ -0,0 +1,68 @@
+use crate::TaskExecutor;
+use slog::Logger;
+use sloggers::{null::NullLoggerBuilder, Build};
+use std::sync::Arc;
+use tokio::runtime;
+
+/// Whilst the `TestRuntime` is not necessarily useful in itself, it provides the necessary
+/// components for creating a `TaskExecutor` during tests.
+///
+/// May create its own runtime or use an existing one.
+///
+/// ## Warning
+///
+/// This struct should never be used in production, only testing.
+pub struct TestRuntime {
+    runtime: Option<Arc<tokio::runtime::Runtime>>,
+    _runtime_shutdown: exit_future::Signal,
+    pub task_executor: TaskExecutor,
+    pub log: Logger,
+}
+
+impl Default for TestRuntime {
+    /// If called *inside* an existing runtime, instantiates `Self` using a handle to that runtime. If
+    /// called *outside* any existing runtime, create a new `Runtime` and keep it alive until the
+    /// `Self` is dropped.
+    fn default() -> Self {
+        let (runtime_shutdown, exit) = exit_future::signal();
+        let (shutdown_tx, _) = futures::channel::mpsc::channel(1);
+        let log = null_logger().unwrap();
+
+        let (runtime, handle) = if let Ok(handle) = runtime::Handle::try_current() {
+            (None, handle)
+        } else {
+            let runtime = Arc::new(
+                runtime::Builder::new_multi_thread()
+                    .enable_all()
+                    .build()
+                    .unwrap(),
+            );
+            let handle = runtime.handle().clone();
+            (Some(runtime), handle)
+        };
+
+        let task_executor = TaskExecutor::new(handle, exit, log.clone(), shutdown_tx);
+
+        Self {
+            runtime,
+            _runtime_shutdown: runtime_shutdown,
+            task_executor,
+            log,
+        }
+    }
+}
+
+impl Drop for TestRuntime {
+    fn drop(&mut self) {
+        if let Some(runtime) = self.runtime.take() {
+            Arc::try_unwrap(runtime).unwrap().shutdown_background()
+        }
+    }
+}
+
+pub fn null_logger() -> Result<Logger, String> {
+    let log_builder = NullLoggerBuilder;
+    log_builder
+        .build()
+        .map_err(|e| format!("Failed to start null logger: {:?}", e))
+}