Fail build if tokio spawn is used outside runtime

Author: arjan-bal

grpc/Cargo.toml

@@ -13,39 +13,46 @@ allowed_external_types = [
 ]
 
 [features]
-default = ["dns"]
+default = ["dns", "_runtime-tokio"]
 dns = ["dep:hickory-resolver"]
+# The following feature is used to ensure all modules use the runtime
+# abstraction instead of using tokio directly.
+# Using tower/buffer enables tokio's rt feature even though it's possible to
+# create Buffers with a user provided executor.
+_runtime-tokio = [
+    "tokio/rt",
+    "tokio/net",
+    "dep:futures",
+    "dep:socket2",
+    "dep:tower",
+]
 
 [dependencies]
 bytes = "1.10.1"
-futures = "0.3.31"
+futures = { version = "0.3.31", optional = true }
 futures-core = "0.3.31"
 futures-util = "0.3.31"
 hickory-resolver = { version = "0.25.1", optional = true }
 http = "1.1.0"
 http-body = "1.0.1"
 hyper = { version = "1.6.0", features = ["client", "http2"] }
-hyper-util = "0.1.14"
 once_cell = "1.19.0"
 parking_lot = "0.12.4"
 pin-project-lite = "0.2.16"
 rand = "0.9"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"
-socket2 = "0.5.10"
-tokio = { version = "1.37.0", features = [
-    "sync",
-    "rt",
-    "net",
-    "time",
-    "macros",
-] }
+socket2 = { version = "0.5.10", optional = true }
+tokio = { version = "1.37.0", features = ["sync", "time", "macros"] }
 tokio-stream = "0.1.17"
 tonic = { version = "0.14.0", path = "../tonic", default-features = false, features = [
     "codegen",
-    "transport",
 ] }
-tower = { version = "0.5.2", features = ["buffer", "limit", "util"] }
+tower = { version = "0.5.2", features = [
+    "limit",
+    "util",
+    "buffer",
+], optional = true }
 tower-service = "0.3.3"
 url = "2.5.0"
 

grpc/src/client/channel.rs

@@ -37,17 +37,16 @@ use std::{
 };
 
 use tokio::sync::{mpsc, oneshot, watch, Notify};
-use tokio::task::AbortHandle;
 
 use serde_json::json;
 use tonic::async_trait;
 use url::Url; // NOTE: http::Uri requires non-empty authority portion of URI
 
-use crate::credentials::Credentials;
+use crate::attributes::Attributes;
 use crate::rt;
 use crate::service::{Request, Response, Service};
-use crate::{attributes::Attributes, rt::tokio::TokioRuntime};
 use crate::{client::ConnectivityState, rt::Runtime};
+use crate::{credentials::Credentials, rt::default_runtime};
 
 use super::service_config::ServiceConfig;
 use super::transport::{TransportRegistry, GLOBAL_TRANSPORT_REGISTRY};
@@ -156,7 +155,7 @@ impl Channel {
             inner: Arc::new(PersistentChannel::new(
                 target,
                 credentials,
-                Arc::new(rt::tokio::TokioRuntime {}),
+                default_runtime(),
                 options,
             )),
         }
@@ -280,7 +279,7 @@ impl ActiveChannel {
         let resolver_opts = name_resolution::ResolverOptions {
             authority,
             work_scheduler,
-            runtime: Arc::new(TokioRuntime {}),
+            runtime: runtime.clone(),
         };
         let resolver = rb.build(&target, resolver_opts);
 
@@ -373,7 +372,7 @@ impl InternalChannelController {
         connectivity_state: Arc<Watcher<ConnectivityState>>,
         runtime: Arc<dyn Runtime>,
     ) -> Self {
-        let lb = Arc::new(GracefulSwitchBalancer::new(wqtx.clone()));
+        let lb = Arc::new(GracefulSwitchBalancer::new(wqtx.clone(), runtime.clone()));
 
         Self {
             lb,
@@ -459,6 +458,7 @@ pub(super) struct GracefulSwitchBalancer {
     policy_builder: Mutex<Option<Arc<dyn LbPolicyBuilder>>>,
     work_scheduler: WorkQueueTx,
     pending: Mutex<bool>,
+    runtime: Arc<dyn Runtime>,
 }
 
 impl WorkScheduler for GracefulSwitchBalancer {
@@ -483,12 +483,13 @@ impl WorkScheduler for GracefulSwitchBalancer {
 }
 
 impl GracefulSwitchBalancer {
-    fn new(work_scheduler: WorkQueueTx) -> Self {
+    fn new(work_scheduler: WorkQueueTx, runtime: Arc<dyn Runtime>) -> Self {
         Self {
             policy_builder: Mutex::default(),
             policy: Mutex::default(), // new(None::<Box<dyn LbPolicy>>),
             work_scheduler,
             pending: Mutex::default(),
+            runtime,
         }
     }
 
@@ -506,6 +507,7 @@ impl GracefulSwitchBalancer {
             let builder = GLOBAL_LB_REGISTRY.get_policy(policy_name).unwrap();
             let newpol = builder.build(LbPolicyOptions {
                 work_scheduler: self.clone(),
+                runtime: self.runtime.clone(),
             });
             *self.policy_builder.lock().unwrap() = Some(builder);
             *p = Some(newpol);

grpc/src/client/load_balancing/child_manager.rs

@@ -38,6 +38,7 @@ use crate::client::load_balancing::{
     WeakSubchannel, WorkScheduler,
 };
 use crate::client::name_resolution::{Address, ResolverUpdate};
+use crate::rt::Runtime;
 
 use super::{Subchannel, SubchannelState};
 
@@ -47,6 +48,7 @@ pub struct ChildManager<T> {
     children: Vec<Child<T>>,
     update_sharder: Box<dyn ResolverUpdateSharder<T>>,
     pending_work: Arc<Mutex<HashSet<usize>>>,
+    runtime: Arc<dyn Runtime>,
 }
 
 struct Child<T> {
@@ -81,12 +83,16 @@ pub trait ResolverUpdateSharder<T>: Send {
 impl<T> ChildManager<T> {
     /// Creates a new ChildManager LB policy.  shard_update is called whenever a
     /// resolver_update operation occurs.
-    pub fn new(update_sharder: Box<dyn ResolverUpdateSharder<T>>) -> Self {
+    pub fn new(
+        update_sharder: Box<dyn ResolverUpdateSharder<T>>,
+        runtime: Arc<dyn Runtime>,
+    ) -> Self {
         Self {
             update_sharder,
             subchannel_child_map: Default::default(),
             children: Default::default(),
             pending_work: Default::default(),
+            runtime,
         }
     }
 
@@ -197,6 +203,7 @@ impl<T: PartialEq + Hash + Eq + Send + Sync + 'static> LbPolicy for ChildManager
                 });
                 let policy = builder.build(LbPolicyOptions {
                     work_scheduler: work_scheduler.clone(),
+                    runtime: self.runtime.clone(),
                 });
                 let state = LbState::initial();
                 self.children.push(Child {

grpc/src/client/load_balancing/mod.rs

@@ -41,6 +41,7 @@ use tonic::{metadata::MetadataMap, Status};
 
 use crate::{
     client::channel::WorkQueueTx,
+    rt::Runtime,
     service::{Request, Response, Service},
 };
 
@@ -64,6 +65,7 @@ pub struct LbPolicyOptions {
     /// A hook into the channel's work scheduler that allows the LbPolicy to
     /// request the ability to perform operations on the ChannelController.
     pub work_scheduler: Arc<dyn WorkScheduler>,
+    pub runtime: Arc<dyn Runtime>,
 }
 
 /// Used to asynchronously request a call into the LbPolicy's work method if

grpc/src/client/load_balancing/pick_first.rs

@@ -13,6 +13,7 @@ use crate::{
         name_resolution::{Address, ResolverUpdate},
         subchannel, ConnectivityState,
     },
+    rt::Runtime,
     service::Request,
 };
 
@@ -31,6 +32,7 @@ impl LbPolicyBuilder for Builder {
             work_scheduler: options.work_scheduler,
             subchannel: None,
             next_addresses: Vec::default(),
+            runtime: options.runtime,
         })
     }
 
@@ -47,6 +49,7 @@ struct PickFirstPolicy {
     work_scheduler: Arc<dyn WorkScheduler>,
     subchannel: Option<Arc<dyn Subchannel>>,
     next_addresses: Vec<Address>,
+    runtime: Arc<dyn Runtime>,
 }
 
 impl LbPolicy for PickFirstPolicy {
@@ -73,10 +76,10 @@ impl LbPolicy for PickFirstPolicy {
         self.next_addresses = addresses;
         let work_scheduler = self.work_scheduler.clone();
         // TODO: Implement Drop that cancels this task.
-        tokio::task::spawn(async move {
+        self.runtime.spawn(Box::pin(async move {
             sleep(Duration::from_millis(200)).await;
             work_scheduler.schedule_work();
-        });
+        }));
         // TODO: return a picker that queues RPCs.
         Ok(())
     }

grpc/src/client/subchannel.rs

@@ -24,7 +24,6 @@ use std::{
 };
 use tokio::{
     sync::{mpsc, oneshot, watch, Notify},
-    task::{AbortHandle, JoinHandle},
     time::{Duration, Instant},
 };
 use tonic::async_trait;
@@ -66,7 +65,7 @@ struct InternalSubchannelReadyState {
 }
 
 struct InternalSubchannelTransientFailureState {
-    abort_handle: Option<AbortHandle>,
+    task_handle: Option<Box<dyn TaskHandle>>,
     error: String,
 }
 
@@ -168,7 +167,7 @@ impl Drop for InternalSubchannelState {
                 }
             }
             Self::TransientFailure(st) => {
-                if let Some(ah) = &st.abort_handle {
+                if let Some(ah) = &st.task_handle {
                     ah.abort();
                 }
             }
@@ -189,8 +188,8 @@ pub(crate) struct InternalSubchannel {
 struct InnerSubchannel {
     state: InternalSubchannelState,
     watchers: Vec<Arc<SubchannelStateWatcher>>, // TODO(easwars): Revisit the choice for this data structure.
-    backoff_task: Option<JoinHandle<()>>,
-    disconnect_task: Option<JoinHandle<()>>,
+    backoff_task: Option<Box<dyn TaskHandle>>,
+    disconnect_task: Option<Box<dyn TaskHandle>>,
 }
 
 #[async_trait]
@@ -417,7 +416,7 @@ impl InternalSubchannel {
             let mut inner = self.inner.lock().unwrap();
             inner.state = InternalSubchannelState::TransientFailure(
                 InternalSubchannelTransientFailureState {
-                    abort_handle: None,
+                    task_handle: None,
                     error: err.clone(),
                 },
             );
@@ -431,14 +430,14 @@ impl InternalSubchannel {
 
         let backoff_interval = self.backoff.backoff_until();
         let state_machine_tx = self.state_machine_event_sender.clone();
-        let backoff_task = tokio::task::spawn(async move {
+        let backoff_task = self.runtime.spawn(Box::pin(async move {
             tokio::time::sleep_until(backoff_interval).await;
             let _ = state_machine_tx.send(SubchannelStateMachineEvent::BackoffExpired);
-        });
+        }));
         let mut inner = self.inner.lock().unwrap();
         inner.state =
             InternalSubchannelState::TransientFailure(InternalSubchannelTransientFailureState {
-                abort_handle: Some(backoff_task.abort_handle()),
+                task_handle: Some(backoff_task),
                 error: err.clone(),
             });
     }

grpc/src/client/transport/mod.rs

@@ -2,6 +2,10 @@ use crate::{rt::Runtime, service::Service};
 use std::{sync::Arc, time::Duration};
 
 mod registry;
+
+// Using tower/buffer enables tokio's rt feature even though it's possible to
+// create Buffers with a user provided executor.
+#[cfg(feature = "_runtime-tokio")]
 mod tonic;
 
 use ::tonic::async_trait;

grpc/src/rt/mod.rs

@@ -23,9 +23,10 @@
  */
 
 use ::tokio::io::{AsyncRead, AsyncWrite};
-use std::{future::Future, net::SocketAddr, pin::Pin, time::Duration};
+use std::{future::Future, net::SocketAddr, pin::Pin, sync::Arc, time::Duration};
 
 pub(crate) mod hyper_wrapper;
+#[cfg(feature = "_runtime-tokio")]
 pub(crate) mod tokio;
 
 /// An abstraction over an asynchronous runtime.
@@ -91,3 +92,45 @@ pub(crate) struct TcpOptions {
 }
 
 pub(crate) trait TcpStream: AsyncRead + AsyncWrite + Send + Unpin {}
+
+/// A fake runtime to satisfy the compiler when no runtime is enabled. This will
+///
+/// # Panics
+///
+/// Panics if any of its functions are called.
+#[derive(Default)]
+pub(crate) struct NoOpRuntime {}
+
+impl Runtime for NoOpRuntime {
+    fn spawn(
+        &self,
+        task: Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
+    ) -> Box<dyn TaskHandle> {
+        unimplemented!()
+    }
+
+    fn get_dns_resolver(&self, opts: ResolverOptions) -> Result<Box<dyn DnsResolver>, String> {
+        unimplemented!()
+    }
+
+    fn sleep(&self, duration: std::time::Duration) -> Pin<Box<dyn Sleep>> {
+        unimplemented!()
+    }
+
+    fn tcp_stream(
+        &self,
+        target: SocketAddr,
+        opts: TcpOptions,
+    ) -> Pin<Box<dyn Future<Output = Result<Box<dyn TcpStream>, String>> + Send>> {
+        unimplemented!()
+    }
+}
+
+pub(crate) fn default_runtime() -> Arc<dyn Runtime> {
+    #[cfg(feature = "_runtime-tokio")]
+    {
+        return Arc::new(tokio::TokioRuntime {});
+    }
+    #[allow(unreachable_code)]
+    Arc::new(NoOpRuntime::default())
+}