remove IdxError and fix tests

tokio-util/src/task/mod.rs

@@ -3,7 +3,7 @@
 #[cfg(tokio_unstable)]
 mod join_map;
 mod spawn_pinned;
-pub use spawn_pinned::{LocalPoolHandle, WorkerIdxError};
+pub use spawn_pinned::LocalPoolHandle;
 
 #[cfg(tokio_unstable)]
 #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]

tokio-util/src/task/spawn_pinned.rs

@@ -1,5 +1,4 @@
 use futures_util::future::{AbortHandle, Abortable};
-use std::error::Error;
 use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::future::Future;
@@ -10,27 +9,6 @@ use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
 use tokio::sync::oneshot;
 use tokio::task::{spawn_local, JoinHandle, LocalSet};
 
-/// Error Type for out-of-bounds indexing error in [`LocalPoolHandle::spawn_pinned_by_idx`].
-///
-/// [`LocalPoolHandle::spawn_pinned_by_idx`]: LocalPoolHandle::spawn_pinned_by_idx
-#[derive(Debug)]
-pub struct WorkerIdxError {
- idx: usize,
- num_workers: usize,
-}
-
-impl fmt::Display for WorkerIdxError {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- write!(
- f,
- "Index {} out of bounds, only {} workers in pool",
- self.idx, self.num_workers
- )
- }
-}
-
-impl Error for WorkerIdxError {}
-
 /// A cloneable handle to a local pool, used for spawning `!Send` tasks.
 ///
 /// Internally the local pool uses a [`tokio::task::LocalSet`] for each worker thread
@@ -156,27 +134,15 @@ impl LocalPoolHandle {
  /// is given by `num_threads() - 1`
  ///
  /// Returns a `WorkerIdxError` if the provided index is out of bounds.
- pub fn spawn_pinned_by_idx<F, Fut>(
- &self,
- create_task: F,
- idx: usize,
- ) -> Result<JoinHandle<Fut::Output>, WorkerIdxError>
+ pub fn spawn_pinned_by_idx<F, Fut>(&self, create_task: F, idx: usize) -> JoinHandle<Fut::Output>
  where
  F: FnOnce() -> Fut,
  F: Send + 'static,
  Fut: Future + 'static,
  Fut::Output: Send + 'static,
  {
- if idx >= self.pool.workers.len() {
- return Err(WorkerIdxError {
- idx,
- num_workers: self.pool.workers.len(),
- });
- }
-
- Ok(self
- .pool
- .spawn_pinned(create_task, WorkerChoice::ByIdx(idx)))
+ self.pool
+ .spawn_pinned(create_task, WorkerChoice::ByIdx(idx))
  }
 
  /// Spawn a task on every worker thread in the pool and pin it so that it
@@ -212,7 +178,7 @@ impl LocalPoolHandle {
  Fut::Output: Send + 'static,
  {
  (0..self.pool.workers.len())
- .map(|idx| self.spawn_pinned_by_idx(create_task.clone(), idx).unwrap())
+ .map(|idx| self.spawn_pinned_by_idx(create_task.clone(), idx))
  .collect::<Vec<_>>()
  }
 }

tokio-util/tests/spawn_pinned.rs

@@ -196,7 +196,7 @@ async fn tasks_are_balanced() {
 #[tokio::test]
 async fn spawn_by_idx() {
  let pool = task::LocalPoolHandle::new(3);
- let barrier = Arc::new(Barrier::new(3));
+ let barrier = Arc::new(Barrier::new(4));
  let barrier1 = barrier.clone();
  let barrier2 = barrier.clone();
  let barrier3 = barrier.clone();
@@ -229,8 +229,8 @@ async fn spawn_by_idx() {
  assert_eq!(loads[1], 1);
  assert_eq!(loads[2], 0);
 
- let thread_id1 = handle1.unwrap().await.unwrap();
- let thread_id2 = handle2.unwrap().await.unwrap();
+ let thread_id1 = handle1.await.unwrap();
+ let thread_id2 = handle2.await.unwrap();
 
  assert_ne!(thread_id1, thread_id2);
 }
@@ -242,7 +242,7 @@ async fn spawn_on_all_workers() {
  let barrier = Arc::new(Barrier::new(NUM_WORKERS + 1));
  let barrier_clone = barrier.clone();
 
- let _ = pool.spawn_pinned_on_all_workers(|| async move {
+ let handles = pool.spawn_pinned_on_all_workers(|| async move {
  barrier_clone.wait().await;
 
  "test"
@@ -253,4 +253,8 @@ async fn spawn_on_all_workers() {
  assert_eq!(loads[0], 1);
  assert_eq!(loads[1], 1);
  assert_eq!(loads[2], 1);
+
+ let _ = handles
+ .into_iter()
+ .map(|handle| async { handle.await.unwrap() });
 }