OwnedTasks 与关闭

pub(crate) struct OwnedTasks<S: 'static> {
    list: ShardedList<Task<S>>,
    pub(crate) id: NonZeroU64,
    closed: AtomicBool,
}

pub(crate) struct ShardedList<L: ShardedListItem> {
    lists: Box<[Mutex<LinkedList<L>>]>,
    added: MetricAtomicU64,
    count: MetricAtomicUsize,
    shard_mask: usize,
}

unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
        // SAFETY: The caller guarantees that `target` points at a valid task.
        let task_id = unsafe { Header::get_id(target) };
        task_id.0.get() as usize
    }
}

fn gen_sharded_list_size(num_cores: usize) -> usize {
    const MAX_SHARDED_LIST_SIZE: usize = 1 << 16;
    usize::min(MAX_SHARDED_LIST_SIZE, num_cores.next_power_of_two() * 4)
}

unsafe fn bind_inner(&self, task: Task<S>, notified: Notified<S>) -> Option<Notified<S>>
where
    S: Schedule,
{
    unsafe {
        // safety: We just created the task, so we have exclusive access
        // to the field.
        task.header().set_owner_id(self.id);
    }

    let shard = self.list.lock_shard(&task);
    // Check the closed flag in the lock for ensuring all that tasks
    // will shut down after the OwnedTasks has been closed.
    if self.closed.load(Ordering::Acquire) {
        drop(shard);
        task.shutdown();
        return None;
    }
    shard.push(task);
    Some(notified)
}

pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> {
    // If the task's owner ID is `None` then it is not part of any list and
    // doesn't need removing.
    let task_id = task.header().get_owner_id()?;

    assert_eq!(task_id, self.id);

    // safety: We just checked that the provided task is not in some other
    // linked list.
    unsafe { self.list.remove(task.header_ptr()) }
}

pub(super) fn close(&self) {
    if self.shared.inject.close() {
        self.notify_all();
    }
}

fn pre_shutdown(&mut self, worker: &Worker) {
    // Start from a random inner list
    let start = self
        .rand
        .fastrand_n(worker.handle.shared.owned.get_shard_size() as u32);
    // Signal to all tasks to shut down.
    worker
        .handle
        .shared
        .owned
        .close_and_shutdown_all(start as usize);
    // ...
}

pub(super) fn shutdown(self) {
    if !self.state().transition_to_shutdown() {
        // The task is concurrently running. No further work needed.
        self.drop_reference();
        return;
    }

    // By transitioning the lifecycle to `Running`, we have permission to
    // drop the future.
    cancel_task(self.core());
    self.complete();
}

fn shutdown_core(&self, core: Box<Core>) {
    let mut cores = self.shared.shutdown_cores.lock();
    cores.push(core);

    if cores.len() != self.shared.remotes.len() {
        return;
    }

    debug_assert!(self.shared.owned.is_empty());

    for mut core in cores.drain(..) {
        core.shutdown(self);
    }

    // Drain the injection queue
    //
    // We already shut down every task, so we can simply drop the tasks.
    while let Some(task) = self.next_remote_task() {
        drop(task);
    }
}

struct Guard(u32);
impl Drop for Guard {
  fn drop(&mut self) {
    println!("guard {} dropped on {:?}", self.0, std::thread::current().name());
  }
}

fn main() {
  let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(2).build().unwrap();
  for i in 0..3 {
    rt.spawn(async move {
      let _g = Guard(i);
      std::future::pending::<()>().await;
    });
  }
  rt.block_on(tokio::task::yield_now());
  println!("alive before drop: {}", rt.metrics().num_alive_tasks());
  drop(rt);
  println!("runtime dropped");
}

alive before drop: 3
guard 0 dropped on Some("tokio-rt-worker")
guard 2 dropped on Some("tokio-rt-worker")
guard 1 dropped on Some("tokio-rt-worker")
runtime dropped

struct Guard;
impl Drop for Guard {
  fn drop(&mut self) {
    println!("guard dropped on {:?}", std::thread::current().name());
  }
}

guard dropped on Some("main")
spawned
Err(JoinError::Cancelled(Id(3)))