lance_core/utils/
tokio.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::LazyLock;
5use std::time::Duration;
6
7use crate::Result;
8
9use futures::{Future, FutureExt};
10use tokio::runtime::{Builder, Runtime};
11use tracing::Span;
12
13pub fn get_num_compute_intensive_cpus() -> usize {
14    if let Ok(user_specified) = std::env::var("LANCE_CPU_THREADS") {
15        return user_specified.parse().unwrap();
16    }
17
18    let cpus = num_cpus::get();
19
20    if cpus <= *IO_CORE_RESERVATION {
21        // If the user is not setting a custom value for LANCE_IO_CORE_RESERVATION then we don't emit
22        // a warning because they're just on a small machine and there isn't much they can do about it.
23        if cpus > 2 {
24            log::warn!(
25                "Number of CPUs is less than or equal to the number of IO core reservations. \
26                This is not a supported configuration. using 1 CPU for compute intensive tasks."
27            );
28        }
29        return 1;
30    }
31
32    num_cpus::get() - *IO_CORE_RESERVATION
33}
34
35pub static IO_CORE_RESERVATION: LazyLock<usize> = LazyLock::new(|| {
36    std::env::var("LANCE_IO_CORE_RESERVATION")
37        .unwrap_or("2".to_string())
38        .parse()
39        .unwrap()
40});
41
42pub static CPU_RUNTIME: LazyLock<Runtime> = LazyLock::new(|| {
43    Builder::new_multi_thread()
44        .thread_name("lance-cpu")
45        .max_blocking_threads(get_num_compute_intensive_cpus())
46        .worker_threads(1)
47        // keep the thread alive "forever"
48        .thread_keep_alive(Duration::from_secs(u64::MAX))
49        .build()
50        .unwrap()
51});
52
53/// Spawn a CPU intensive task
54///
55/// This task will be put onto a thread pool dedicated for CPU-intensive work
56/// This keeps the tokio thread pool free so that we can always be ready to service
57/// cheap I/O & control requests.
58///
59/// This can also be used to convert a big chunk of synchronous work into a future
60/// so that it can be run in parallel with something like StreamExt::buffered()
61pub fn spawn_cpu<F: FnOnce() -> Result<R> + Send + 'static, R: Send + 'static>(
62    func: F,
63) -> impl Future<Output = Result<R>> {
64    let (send, recv) = tokio::sync::oneshot::channel();
65    // Propagate the current span into the task
66    let span = Span::current();
67    CPU_RUNTIME.spawn_blocking(move || {
68        let _span_guard = span.enter();
69        let result = func();
70        let _ = send.send(result);
71    });
72    recv.map(|res| res.unwrap())
73}
lance_core/utils/tokio.rs

lance_core/utils/
tokio.rs
