all: Rename graph::task_spawn to tokio and move TEST_RUNTIME there

Author: lutter

graph/src/lib.rs

@@ -38,8 +38,10 @@ pub mod env;
 pub mod ipfs;
 
 /// Wrapper for spawning tasks that abort on panic, which is our default.
-mod task_spawn;
-pub use task_spawn::{
+mod tokio;
+#[cfg(debug_assertions)]
+pub use tokio::TEST_RUNTIME;
+pub use tokio::{
     block_on, spawn, spawn_allow_panic, spawn_blocking, spawn_blocking_allow_panic, spawn_thread,
 };
 

graph/src/log/elastic.rs

@@ -202,7 +202,7 @@ impl ElasticDrain {
         let mut interval = tokio::time::interval(self.config.flush_interval);
         let max_retries = self.config.max_retries;
 
-        crate::task_spawn::spawn(async move {
+        crate::tokio::spawn(async move {
             loop {
                 interval.tick().await;
 

graph/src/task_spawn.rs renamed to graph/src/tokio.rs

@@ -1,20 +1,49 @@
-//! The functions in this module should be used to execute futures, serving as a facade to the
-//! underlying executor implementation which currently is tokio. This serves a few purposes:
-//! - Avoid depending directly on tokio APIs, making upgrades or a potential switch easier.
-//! - Reflect our chosen default semantics of aborting on task panic, offering `*_allow_panic`
-//!   functions to opt out of that.
-//! - Reflect that historically we've used blocking futures due to making DB calls directly within
-//!   futures. This point should go away once https://github.com/graphprotocol/graph-node/issues/905
-//!   is resolved. Then the blocking flavors should no longer accept futures but closures.
+//! Helpers for dealing with certain aspects of tokio.
 //!
-//! These should not be called from within executors other than tokio, particularly the blocking
-//! functions will panic in that case. We should generally avoid mixing executors whenever possible.
+//! This module sets up a runtime on which all tests should run, as well as
+//! providing some functions for spawning tasks with our desired semantics.
+//!
+//! The functions in this module should be used to execute futures, serving
+//! as a facade to the underlying executor implementation which currently is
+//! tokio. This serves a few purposes:
+//! - Avoid depending directly on tokio APIs, making upgrades or a potential
+//!   switch easier.
+//! - Reflect our chosen default semantics of aborting on task panic,
+//!   offering `*_allow_panic` functions to opt out of that.
+//! - Reflect that historically we've used blocking futures due to making DB
+//!   calls directly within futures. This point should go away once
+//!   https://github.com/graphprotocol/graph-node/issues/905 is resolved.
+//!   Then the blocking flavors should no longer accept futures but
+//!   closures.
+//!
+//! These should not be called from within executors other than tokio,
+//! particularly the blocking functions will panic in that case. We should
+//! generally avoid mixing executors whenever possible.
 
 use futures03::future::{FutureExt, TryFutureExt};
 use std::future::Future as Future03;
 use std::panic::AssertUnwindSafe;
 use tokio::task::JoinHandle;
 
+#[cfg(debug_assertions)]
+use tokio::runtime::{Builder, Runtime};
+
+#[cfg(debug_assertions)]
+lazy_static::lazy_static! {
+    /// The one true runtime for all tests. Tests should use the
+    /// `graph::test` macro to make sure they are using this runtime, the
+    /// same way they would use `#[tokio::test]`.
+    ///
+    /// We need to make sure we use a single runtime because if there are
+    /// multiple runtimes involved, the task that diesel_async spawns to
+    /// drive database connections (see `drive_connection` in the
+    /// `diesel_async` crate) may end up on a different runtime than the one
+    /// the test is using, leading to that task getting dropped, and the
+    /// test using a connection receiving a `Connection closed` error.
+    pub static ref TEST_RUNTIME: Runtime =
+        Builder::new_multi_thread().enable_all().build().unwrap();
+}
+
 fn abort_on_panic<T: Send + 'static>(
     f: impl Future03<Output = T> + Send + 'static,
 ) -> impl Future03<Output = T> {

server/graphman/tests/util/mod.rs

@@ -4,18 +4,14 @@ pub mod server;
 use std::future::Future;
 use std::sync::Mutex;
 
+use graph::TEST_RUNTIME;
 use lazy_static::lazy_static;
 use test_store::store::remove_subgraphs;
 use test_store::store::PRIMARY_POOL;
-use tokio::runtime::Builder;
-use tokio::runtime::Runtime;
 
 lazy_static! {
     // Used to make sure tests will run sequentially.
     static ref SEQ_MUX: Mutex<()> = Mutex::new(());
-
-    // One runtime helps share the same server between the tests.
-    static ref RUNTIME: Runtime = Builder::new_current_thread().enable_all().build().unwrap();
 }
 
 pub fn run_test<T, F>(test: T)
@@ -25,7 +21,7 @@ where
 {
     let _lock = SEQ_MUX.lock().unwrap_or_else(|err| err.into_inner());
 
-    RUNTIME.block_on(async {
+    TEST_RUNTIME.block_on(async {
         cleanup_graphman_command_executions_table().await;
         remove_subgraphs().await;
 

store/test-store/src/store.rs

@@ -37,7 +37,6 @@ use std::collections::BTreeSet;
 use std::collections::HashMap;
 use std::time::Instant;
 use std::{marker::PhantomData, sync::Mutex};
-use tokio::runtime::{Builder, Runtime};
 use web3::types::H256;
 
 pub const NETWORK_NAME: &str = "fake_network";
@@ -54,8 +53,6 @@ lazy_static! {
         None => Logger::root(slog::Discard, o!()),
     };
     static ref SEQ_LOCK: Mutex<()> = Mutex::new(());
-    pub static ref STORE_RUNTIME: Runtime =
-        Builder::new_multi_thread().enable_all().build().unwrap();
     pub static ref METRICS_REGISTRY: Arc<MetricsRegistry> = Arc::new(MetricsRegistry::mock());
     pub static ref LOAD_MANAGER: Arc<LoadManager> = Arc::new(LoadManager::new(
         &LOGGER,
@@ -119,7 +116,7 @@ where
         Err(err) => err.into_inner(),
     };
 
-    STORE_RUNTIME.handle().block_on(async {
+    graph::TEST_RUNTIME.handle().block_on(async {
         let store = STORE.clone();
         test(store).await
     })
@@ -630,7 +627,7 @@ fn build_store() -> (Arc<Store>, ConnectionPool, Config, Arc<SubscriptionManager
         .unwrap_or_else(|_| panic!("config is not valid (file={:?})", &opt.config));
     let registry = Arc::new(MetricsRegistry::mock());
     std::thread::spawn(move || {
-        STORE_RUNTIME.handle().block_on(async {
+        graph::TEST_RUNTIME.handle().block_on(async {
             let builder = StoreBuilder::new(&LOGGER, &NODE_ID, &config, None, registry).await;
             let subscription_manager = builder.subscription_manager();
             let primary_pool = builder.primary_pool();