Use tokio's oneshot everywhere (#30516)

Switch our oneshots from `futures` to `tokio` so they're coop budget aware -- if a Tokio task has exhausted its coop budget, `oneshot::Receivers` will preempt the task if awaited. This does require that our futures are all executed within the Tokio runtime, however.

GitOrigin-RevId: da8f8017fd0f07f8e50332682c43ce514ba8c3d4

Author: Sujay Jayakar

Cargo.lock

@@ -193,3 +193,11 @@ codegen-units = 16
 [profile.dev.package.num-bigint-dig]
 opt-level = 3
 codegen-units = 16
+
+[profile.dev.package.tokio]
+opt-level = 3
+codegen-units = 16
+
+[profile.dev.package.tokio-stream]
+opt-level = 3
+codegen-units = 16

Cargo.toml

@@ -52,7 +52,6 @@ use common::{
     },
 };
 use float_next_after::NextAfter;
-use futures::channel::oneshot;
 use http::{
     Method,
     StatusCode,
@@ -74,6 +73,7 @@ use serde_json::{
     json,
     Value as JsonValue,
 };
+use tokio::sync::oneshot;
 use url::Url;
 use usage_tracking::{
     AggregatedFunctionUsageStats,

crates/application/src/function_log.rs

@@ -135,10 +135,7 @@ use function_log::{
     FunctionExecutionPart,
 };
 use function_runner::FunctionRunner;
-use futures::{
-    channel::oneshot,
-    stream::BoxStream,
-};
+use futures::stream::BoxStream;
 use headers::{
     ContentLength,
     ContentType,
@@ -283,6 +280,7 @@ use table_summary_worker::{
     TableSummaryClient,
     TableSummaryWorker,
 };
+use tokio::sync::oneshot;
 use usage_tracking::{
     FunctionUsageStats,
     FunctionUsageTracker,

crates/application/src/lib.rs

@@ -22,12 +22,12 @@ use database::{
     TableSummaryWriter,
 };
 use futures::{
-    channel::oneshot,
     pin_mut,
     select_biased,
     FutureExt,
 };
 use parking_lot::Mutex;
+use tokio::sync::oneshot;
 
 use crate::metrics::log_worker_starting;
 

crates/application/src/table_summary_worker.rs

@@ -1,10 +1,7 @@
 use std::sync::Arc;
 
 use futures::{
-    channel::{
-        mpsc,
-        oneshot,
-    },
+    channel::mpsc,
     future::{
         self,
         BoxFuture,
@@ -17,6 +14,7 @@ use futures::{
     StreamExt,
 };
 use parking_lot::Mutex;
+use tokio::sync::oneshot;
 
 use crate::{
     codel_queue::{

crates/common/src/bounded_thread_pool.rs

@@ -20,7 +20,6 @@ use std::{
 use anyhow::Context;
 use async_trait::async_trait;
 use futures::{
-    channel::oneshot,
     future::{
         BoxFuture,
         FusedFuture,
@@ -52,6 +51,7 @@ use proptest::prelude::*;
 use rand::RngCore;
 use serde::Serialize;
 use thiserror::Error;
+use tokio::sync::oneshot;
 use uuid::Uuid;
 use value::heap_size::HeapSize;
 
@@ -203,6 +203,7 @@ pub trait Runtime: Clone + Sync + Send + 'static {
     /// Spawn a future on a reserved OS thread. This is only really necessary
     /// for libraries like `V8` that care about being called from a
     /// particular thread.
+    #[must_use = "Threads are canceled when their `SpawnHandle` is dropped."]
     fn spawn_thread<Fut: Future<Output = ()>, F: FnOnce() -> Fut + Send + 'static>(
         &self,
         f: F,

crates/common/src/runtime/mod.rs

@@ -18,7 +18,6 @@ use std::{
 };
 
 use futures::{
-    channel::oneshot,
     future::{
         self,
         BoxFuture,
@@ -34,10 +33,13 @@ use rand::{
     SeedableRng,
 };
 use rand_chacha::ChaCha12Rng;
-use tokio::runtime::{
-    Builder,
-    RngSeed,
-    UnhandledPanic,
+use tokio::{
+    runtime::{
+        Builder,
+        RngSeed,
+        UnhandledPanic,
+    },
+    sync::oneshot,
 };
 
 use super::{

crates/common/src/runtime/testing/mod.rs

@@ -1,18 +1,18 @@
 pub mod split_rw_lock;
 pub mod state_channel;
 
-// It's safe to use these `tokio` sync primitives in our runtime-generic code
-// since they don't internally depend on the `tokio` runtime. Feel free to add
-// more if you need them, but generally prefer using `futures`-based primitives
-// if sufficient.
+use futures::future;
+use tokio::sync::oneshot;
 pub use tokio::sync::{
     broadcast,
-    // This channel is useful over `futures::channel::mpsc` since it doesn't require `&mut self` on
-    // `try_send`. The `futures` implementation conforms to their `Sink` trait which unnecessarily
-    // requires mutability.
     mpsc,
     watch,
     Mutex,
     MutexGuard,
     Notify,
 };
+
+/// Wait until a sender's corresponding receiver has been closed.
+pub async fn oneshot_receiver_closed<T>(sender: &mut oneshot::Sender<T>) {
+    future::poll_fn(|cx| sender.poll_closed(cx)).await
+}

crates/common/src/sync/mod.rs

@@ -21,13 +21,13 @@ use convex::{
     Value,
 };
 use futures::{
-    channel::oneshot,
     pin_mut,
     select_biased,
     FutureExt,
     StreamExt,
 };
 use maplit::btreemap;
+use tokio::sync::oneshot;
 
 const SETUP_MSG: &str = r"
 Please run this Convex Chat client from an initialized Convex project.