阻塞线程池

let id = task::Id::next();
let fut =
    blocking_task::<F, BlockingTask<F>>(BlockingTask::new(func), spawn_meta, id.as_u64());

let (task, handle) = task::unowned(
    fut,
    BlockingSchedule::new(rt),
    id,
    task::SpawnLocation::capture(),
);

let spawned = self.spawn_task(Task::new(task, is_mandatory), rt);
(handle, spawned)

fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<R> {
    let me = &mut *self;
    let func = me
        .func
        .take()
        .expect("[internal exception] blocking task ran twice.");

    crate::task::coop::stop();

    Poll::Ready(func())
}

let mut locked = self.mutex.lock();

if locked.thread_mgmt_state.shutdown {
    drop(locked);

    // Shutdown the task: it's fine to shutdown this task
    // (even if mandatory) because it was scheduled after the
    // shutdown of the runtime began.
    task.task.shutdown();
    return Err(SpawnError::ShuttingDown);
}

locked.queue.push_back(task);
metrics.inc_queue_depth();

if metrics.num_idle_threads() == 0 {
    on_no_idle(&mut locked.thread_mgmt_state)?;
} else {
    // Notify an idle worker thread. The notification counter
    // is used to count the needed amount of notifications
    // exactly. Thread libraries may generate spurious
    // wakeups, this counter is used to keep us in a
    // consistent state.
    metrics.dec_num_idle_threads();
    locked.num_notify += 1;
    self.condvar.notify_one();
}

'main: loop {
    // BUSY
    while let Some(task) = locked.queue.pop_front() {
        metrics.dec_queue_depth();
        drop(locked);
        task.run();

        locked = self.mutex.lock();
    }

    // IDLE
    metrics.inc_num_idle_threads();
    // ...

    while !locked.thread_mgmt_state.shutdown {
        let lock_result = self.condvar.wait_timeout(locked, keep_alive).unwrap();

        locked = lock_result.0;
        let timeout_result = lock_result.1;

        if locked.num_notify != 0 {
            // We have received a legitimate wakeup,
            // acknowledge it by decrementing the counter
            // and transition to the BUSY state.
            locked.num_notify -= 1;
            // ...
            break;
        }

        // Even if the condvar "timed out", if the pool is
        // entering the shutdown phase, we want to perform
        // the cleanup logic.
        if !locked.thread_mgmt_state.shutdown && timeout_result.timed_out() {
            join_on_thread = locked.thread_mgmt_state.worker_timed_out(worker_thread_id);

            break 'main;
        }

        // Spurious wakeup detected, go back to sleep.
    }

    if locked.thread_mgmt_state.shutdown {
        // Drain the queue
        while let Some(task) = locked.queue.pop_front() {
            metrics.dec_queue_depth();
            drop(locked);

            task.shutdown_or_run_if_mandatory();

            locked = self.mutex.lock();
        }

        break;
    }
}

use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

#[tokio::main]
async fn main() {
  let ids = Arc::new(Mutex::new(HashSet::new()));
  let start = Instant::now();
  let mut hs = Vec::new();
  for _ in 0..600 {
    let ids = ids.clone();
    hs.push(tokio::task::spawn_blocking(move || {
      ids.lock().unwrap().insert(std::thread::current().id());
      std::thread::sleep(Duration::from_millis(200));
    }));
  }
  for h in hs {
    h.await.unwrap();
  }
  println!(
    "600 blocking tasks ran on {} threads in {:?}",
    ids.lock().unwrap().len(),
    start.elapsed()
  );
}

600 blocking tasks ran on 512 threads in 408.186625ms

match (
    crate::runtime::context::current_enter_context(),
    maybe_cx.is_some(),
) {
    (context::EnterRuntime::Entered { .. }, true) => {
        // We are on a thread pool runtime thread, so we just need to
        // set up blocking.
        had_entered = true;
    }
    (
        context::EnterRuntime::Entered {
            allow_block_in_place,
        },
        false,
    ) => {
        // We are on an executor, but _not_ on the thread pool.  That is
        // _only_ okay if we are in a thread pool runtime's block_on
        // method:
        if allow_block_in_place {
            had_entered = true;
            return Ok(());
        } else {
            // This probably means we are on the current_thread runtime or in a
            // LocalSet, where it is _not_ okay to block.
            return Err(
                "can call blocking only when running on the multi-threaded runtime",
            );
        }
    }
    (context::EnterRuntime::NotEntered, true) => {
        // This is a nested call to block_in_place (we already exited).
        // All the necessary setup has already been done.
        return Ok(());
    }
    (context::EnterRuntime::NotEntered, false) => {
        // We are outside of the tokio runtime, so blocking is fine.
        // We can also skip all of the thread pool blocking setup steps.
        return Ok(());
    }
}

// Since deferred tasks don't stay on `core`, make sure to wake them
// before blocking.
cx.defer.wake();

// Get the worker core. If none is set, then blocking is fine!
let mut core = match cx.core.borrow_mut().take() {
    Some(core) => core,
    None => return Ok(()),
};

// If we heavily call `spawn_blocking`, there might be no available thread to
// run this core. Except for the task in the lifo_slot, all tasks can be
// stolen, so we move the task out of the lifo_slot to the run_queue.
if let Some(task) = core.lifo_slot.take() {
    core.run_queue
        .push_back_or_overflow(task, &*cx.worker.handle, &mut core.stats);
}

// We are taking the core from the context and sending it to another
// thread.
take_core = true;

// ...

// In order to block, the core must be sent to another thread for
// execution.
//
// First, move the core back into the worker's shared core slot.
cx.worker.core.set(core);

// Next, clone the worker handle and send it to a new thread for
// processing.
//
// Once the blocking task is done executing, we will attempt to
// steal the core back.
let worker = cx.worker.clone();
runtime::spawn_blocking(move || run(worker));

impl Drop for Reset {
    fn drop(&mut self) {
        with_current(|maybe_cx| {
            if let Some(cx) = maybe_cx {
                if self.take_core {
                    let core = cx.worker.core.take();

                    // ...

                    let mut cx_core = cx.core.borrow_mut();
                    assert!(cx_core.is_none());
                    *cx_core = core;
                }

                // Reset the task budget as we are re-entering the
                // runtime.
                coop::set(self.budget);
            }
        });
    }
}

use std::time::{Duration, Instant};

fn main() {
  for use_block_in_place in [false, true] {
    let rt = tokio::runtime::Builder::new_multi_thread()
      .worker_threads(1)
      .enable_all()
      .build()
      .unwrap();
    let start = Instant::now();
    rt.block_on(async move {
      let a = tokio::spawn(async move {
        println!("  A starts on {:?}", std::thread::current().id());
        if use_block_in_place {
          tokio::task::block_in_place(|| std::thread::sleep(Duration::from_millis(300)));
        } else {
          std::thread::sleep(Duration::from_millis(300));
        }
        println!("  A ends   on {:?} at {:?}", std::thread::current().id(), start.elapsed());
      });
      tokio::time::sleep(Duration::from_millis(50)).await;
      let b = tokio::spawn(async move {
        println!("  B runs   on {:?} at {:?}", std::thread::current().id(), start.elapsed());
      });
      b.await.unwrap();
      a.await.unwrap();
    });
    println!("block_in_place = {use_block_in_place} done");
  }
}

  A starts on ThreadId(2)
  A ends   on ThreadId(2) at 305.087833ms
  B runs   on ThreadId(2) at 305.278458ms
block_in_place = false done
  A starts on ThreadId(3)
  B runs   on ThreadId(4) at 51.76075ms
  A ends   on ThreadId(3) at 301.761542ms
block_in_place = true done

#[tokio::main(flavor = "current_thread")]
async fn main() {
  tokio::task::block_in_place(|| {});
}

thread 'main' (15533453) panicked at src/bin/bip_ct.rs:3:3:
can call blocking only when running on the multi-threaded runtime

pub(crate) fn shutdown(&mut self, timeout: Option<Duration>) {
    // The function can be called multiple times. First, by explicitly
    // calling `shutdown` then by the drop handler calling `shutdown`. This
    // prevents shutting down twice.
    let inner = &self.spawner.inner;
    let (last_exited_thread, workers) = match inner.inner_impl.begin_shutdown(&inner.metrics) {
        Some(x) => x,
        None => return,
    };

    if self.shutdown_rx.wait(timeout) {
        let _ = last_exited_thread.map(thread::JoinHandle::join);

        // ...

        for (_id, handle) in workers {
            let _ = handle.join();
        }
    }
}

Cannot drop a runtime in a context where blocking is not allowed. This happens when a runtime is dropped from within an asynchronous context.

use std::time::Duration;

fn main() {
  let rt = tokio::runtime::Builder::new_current_thread()
    .max_blocking_threads(1)
    .enable_all()
    .build()
    .unwrap();
  rt.block_on(async {
    let a = tokio::task::spawn_blocking(|| {
      std::thread::sleep(Duration::from_millis(100));
      println!("A closure ran");
    });
    // Give A time to start on the only blocking thread.
    tokio::time::sleep(Duration::from_millis(20)).await;
    let b = tokio::task::spawn_blocking(|| println!("B closure ran"));
    a.abort();
    b.abort();
    println!("A: {:?}", a.await);
    println!("B: {:?}", b.await);
  });
}

A closure ran
A: Ok(())
B: Err(JoinError::Cancelled(Id(2)))