Skip to main content

pingora_runtime/
lib.rs

1// Copyright 2026 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Pingora tokio runtime.
16//!
17//! Tokio runtime comes in two flavors: a single-threaded runtime
18//! and a multi-threaded one which provides work stealing.
19//! Benchmark shows that, compared to the single-threaded runtime, the multi-threaded one
20//! has some overhead due to its more sophisticated work steal scheduling.
21//!
22//! This crate provides a third flavor: a multi-threaded runtime without work stealing.
23//! This flavor is as efficient as the single-threaded runtime while allows the async
24//! program to use multiple cores.
25
26use once_cell::sync::{Lazy, OnceCell};
27use rand::Rng;
28use serde::{Deserialize, Serialize};
29#[cfg(feature = "dial9")]
30use std::path::PathBuf;
31use std::sync::Arc;
32use std::thread::JoinHandle;
33use std::time::Duration;
34use thread_local::ThreadLocal;
35use tokio::runtime::{Builder, Handle};
36use tokio::sync::oneshot::{channel, Sender};
37
38/// Default maximum size of a dial9 trace segment file.
39#[cfg(feature = "dial9")]
40pub const DEFAULT_DIAL9_MAX_FILE_SIZE: u64 = 100 * 1024 * 1024;
41/// Default maximum bytes retained locally by dial9.
42#[cfg(feature = "dial9")]
43pub const DEFAULT_DIAL9_MAX_TOTAL_SIZE: u64 = 512 * 1024 * 1024;
44
45/// Configuration options for the blocking thread pool used by the runtime.
46///
47/// These options control the behavior of the blocking thread pool that handles
48/// [`tokio::task::spawn_blocking`] tasks.
49#[derive(Debug, Clone, Default)]
50pub struct BlockingPoolOpts {
51    /// The maximum number of threads in the blocking thread pool.
52    ///
53    /// When not set, the tokio default (512) is used.
54    pub max_threads: Option<usize>,
55    /// The duration that idle blocking threads are kept alive before being shut down.
56    ///
57    /// When not set, the tokio default (10 seconds) is used.
58    pub thread_keep_alive: Option<Duration>,
59}
60
61/// Configuration options for runtime metrics collection.
62#[derive(Debug, Clone, Default)]
63pub struct RuntimeMetricsOpts {
64    /// Enable Tokio's poll-time histogram on the runtime.
65    ///
66    /// This must be configured before the runtime is built. Enabling it adds
67    /// two timestamp reads to every task poll.
68    pub poll_time_histogram: bool,
69    /// Histogram bucket scale for Tokio's poll-time histogram.
70    pub poll_time_histogram_scale: Option<RuntimeMetricsPollTimeHistogramScale>,
71    /// Width of the first histogram bucket.
72    pub poll_time_histogram_resolution: Option<Duration>,
73    /// Number of histogram buckets. Memory usage scales with runtimes × workers × buckets.
74    pub poll_time_histogram_buckets: Option<usize>,
75}
76
77/// Configuration options for a Tokio runtime.
78#[derive(Debug, Clone, Default)]
79pub struct RuntimeOpts {
80    /// Options for runtime metrics collection.
81    pub metrics: RuntimeMetricsOpts,
82    /// Enable Tokio's experimental alternative timer.
83    ///
84    /// This requires building with `--cfg tokio_unstable` and only applies to
85    /// Tokio's multi-threaded runtime.
86    pub enable_alt_timer: bool,
87    /// Options for dial9 Tokio telemetry.
88    #[cfg(feature = "dial9")]
89    pub dial9: Option<Dial9RuntimeOpts>,
90}
91
92/// Configuration options for dial9 Tokio telemetry.
93#[cfg(feature = "dial9")]
94#[derive(Debug, Clone)]
95pub struct Dial9RuntimeOpts {
96    /// Trace output path after server configuration defaults are applied.
97    pub trace_path: PathBuf,
98    /// Rotate trace segments after this many bytes.
99    pub max_file_size: u64,
100    /// Maximum bytes retained on local disk.
101    pub max_total_size: u64,
102    /// Wall-clock trace rotation period.
103    pub rotation_period: Option<Duration>,
104    /// Enable dial9 task spawn/terminate tracking.
105    pub task_tracking: bool,
106    /// How often the background worker checks for sealed trace segments.
107    pub worker_poll_interval: Option<Duration>,
108    /// Upload sealed trace segments to S3-compatible storage.
109    #[cfg(feature = "dial9-worker-s3")]
110    pub s3_upload: Option<Dial9S3UploadOpts>,
111}
112
113#[cfg(feature = "dial9")]
114impl Dial9RuntimeOpts {
115    /// Create dial9 runtime options using Pingora's dial9 defaults.
116    pub fn new(trace_path: impl Into<PathBuf>) -> Self {
117        Self {
118            trace_path: trace_path.into(),
119            max_file_size: DEFAULT_DIAL9_MAX_FILE_SIZE,
120            max_total_size: DEFAULT_DIAL9_MAX_TOTAL_SIZE,
121            rotation_period: None,
122            task_tracking: true,
123            worker_poll_interval: None,
124            #[cfg(feature = "dial9-worker-s3")]
125            s3_upload: None,
126        }
127    }
128
129    /// Set the maximum size of each trace segment file.
130    pub fn with_max_file_size(mut self, max_file_size: u64) -> Self {
131        self.max_file_size = max_file_size;
132        self
133    }
134
135    /// Set the maximum bytes retained on local disk.
136    pub fn with_max_total_size(mut self, max_total_size: u64) -> Self {
137        self.max_total_size = max_total_size;
138        self
139    }
140
141    /// Set the wall-clock trace rotation period.
142    pub fn with_rotation_period(mut self, rotation_period: Duration) -> Self {
143        self.rotation_period = Some(rotation_period);
144        self
145    }
146
147    /// Enable or disable dial9 task spawn/terminate tracking.
148    pub fn with_task_tracking(mut self, task_tracking: bool) -> Self {
149        self.task_tracking = task_tracking;
150        self
151    }
152
153    /// Set how often the background worker checks for sealed trace segments.
154    pub fn with_worker_poll_interval(mut self, worker_poll_interval: Duration) -> Self {
155        self.worker_poll_interval = Some(worker_poll_interval);
156        self
157    }
158
159    /// Set S3-compatible upload options for sealed trace segments.
160    #[cfg(feature = "dial9-worker-s3")]
161    pub fn with_s3_upload(mut self, s3_upload: Dial9S3UploadOpts) -> Self {
162        self.s3_upload = Some(s3_upload);
163        self
164    }
165}
166
167/// Configuration options for dial9 S3-compatible trace uploads.
168#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))]
169#[derive(Debug, Clone)]
170pub struct Dial9S3UploadOpts {
171    /// S3 bucket that receives sealed trace segments.
172    pub bucket: String,
173    /// Service name included in uploaded object keys.
174    pub service_name: String,
175    /// Optional key prefix.
176    pub prefix: Option<String>,
177    /// Optional region override.
178    pub region: Option<String>,
179    /// Optional instance identifier included in uploaded object keys.
180    pub instance_path: Option<String>,
181    /// Optional pre-built S3 client for custom credentials or endpoints.
182    pub client: Option<aws_sdk_s3::Client>,
183}
184
185#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))]
186impl Dial9S3UploadOpts {
187    /// Create S3-compatible upload options.
188    pub fn new(bucket: impl Into<String>, service_name: impl Into<String>) -> Self {
189        Self {
190            bucket: bucket.into(),
191            service_name: service_name.into(),
192            prefix: None,
193            region: None,
194            instance_path: None,
195            client: None,
196        }
197    }
198
199    /// Set the object key prefix.
200    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
201        self.prefix = Some(prefix.into());
202        self
203    }
204
205    /// Set the AWS region override.
206    pub fn with_region(mut self, region: impl Into<String>) -> Self {
207        self.region = Some(region.into());
208        self
209    }
210
211    /// Set the instance identifier included in uploaded object keys.
212    pub fn with_instance_path(mut self, instance_path: impl Into<String>) -> Self {
213        self.instance_path = Some(instance_path.into());
214        self
215    }
216
217    /// Set a pre-built S3 client for custom credentials or endpoints.
218    pub fn with_client(mut self, client: aws_sdk_s3::Client) -> Self {
219        self.client = Some(client);
220        self
221    }
222}
223
224/// Bucket scale for Tokio's poll-time histogram.
225#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
226#[serde(rename_all = "snake_case")]
227pub enum RuntimeMetricsPollTimeHistogramScale {
228    /// Equal-width buckets.
229    Linear,
230    /// Buckets double in width at each step.
231    Log,
232}
233
234/// Pingora async multi-threaded runtime
235///
236/// The `Steal` flavor is effectively tokio multi-threaded runtime.
237///
238/// The `NoSteal` flavor is backed by multiple tokio single-threaded runtime.
239pub enum Runtime {
240    Steal {
241        runtime: tokio::runtime::Runtime,
242        #[cfg(feature = "dial9")]
243        dial9_guard: Option<dial9_tokio_telemetry::telemetry::TelemetryGuard>,
244    },
245    NoSteal(NoStealRuntime),
246}
247
248/// Apply [`BlockingPoolOpts`] to a tokio [`Builder`].
249fn apply_blocking_opts(builder: &mut Builder, opts: &BlockingPoolOpts) {
250    if let Some(max) = opts.max_threads {
251        builder.max_blocking_threads(max);
252    }
253    if let Some(ttl) = opts.thread_keep_alive {
254        builder.thread_keep_alive(ttl);
255    }
256}
257
258/// Apply [`RuntimeMetricsOpts`] to a tokio [`Builder`].
259// The replacement `metrics_poll_time_histogram_configuration` API is not used
260// here so this crate can continue to compile against older Tokio 1.x versions
261// selected by downstream applications while still honoring these knobs in
262// tokio-unstable builds.
263#[allow(deprecated)]
264fn apply_metrics_opts(builder: &mut Builder, opts: &RuntimeMetricsOpts) {
265    #[cfg(tokio_unstable)]
266    if opts.poll_time_histogram {
267        builder.enable_metrics_poll_time_histogram();
268
269        if let Some(scale) = opts.poll_time_histogram_scale {
270            builder.metrics_poll_count_histogram_scale(match scale {
271                RuntimeMetricsPollTimeHistogramScale::Linear => {
272                    tokio::runtime::HistogramScale::Linear
273                }
274                RuntimeMetricsPollTimeHistogramScale::Log => tokio::runtime::HistogramScale::Log,
275            });
276        }
277        if let Some(resolution) = opts
278            .poll_time_histogram_resolution
279            .filter(|resolution| !resolution.is_zero())
280        {
281            builder.metrics_poll_count_histogram_resolution(resolution);
282        }
283        if let Some(buckets) = opts
284            .poll_time_histogram_buckets
285            .filter(|buckets| *buckets > 0)
286        {
287            builder.metrics_poll_count_histogram_buckets(buckets);
288        }
289    }
290
291    #[cfg(not(tokio_unstable))]
292    let _ = (builder, opts);
293}
294
295/// Apply timer options from [`RuntimeOpts`] to a tokio [`Builder`].
296fn apply_timer_opts(builder: &mut Builder, opts: &RuntimeOpts) {
297    #[cfg(tokio_unstable)]
298    if opts.enable_alt_timer {
299        builder.enable_alt_timer();
300    }
301
302    #[cfg(not(tokio_unstable))]
303    let _ = (builder, opts);
304}
305
306#[cfg(feature = "dial9")]
307fn build_dial9_runtime(
308    builder: Builder,
309    runtime_name: &str,
310    opts: &Dial9RuntimeOpts,
311) -> std::io::Result<(
312    tokio::runtime::Runtime,
313    dial9_tokio_telemetry::telemetry::TelemetryGuard,
314)> {
315    use dial9_tokio_telemetry::telemetry::{RotatingWriter, TracedRuntime};
316    use std::io::{Error, ErrorKind};
317
318    if opts.max_file_size == 0 {
319        return Err(Error::new(
320            ErrorKind::InvalidInput,
321            "dial9 max_file_size must be greater than zero",
322        ));
323    }
324    if opts.max_total_size == 0 {
325        return Err(Error::new(
326            ErrorKind::InvalidInput,
327            "dial9 max_total_size must be greater than zero",
328        ));
329    }
330    if opts.max_file_size > opts.max_total_size {
331        return Err(Error::new(
332            ErrorKind::InvalidInput,
333            "dial9 max_file_size must be less than or equal to max_total_size",
334        ));
335    }
336    if opts.worker_poll_interval == Some(Duration::ZERO) {
337        return Err(Error::new(
338            ErrorKind::InvalidInput,
339            "dial9 worker_poll_interval must be greater than zero",
340        ));
341    }
342
343    if let Some(parent) = opts.trace_path.parent() {
344        std::fs::create_dir_all(parent)?;
345    }
346
347    let writer = RotatingWriter::builder()
348        .base_path(opts.trace_path.clone())
349        .max_file_size(opts.max_file_size)
350        .max_total_size(opts.max_total_size)
351        .maybe_rotation_period(opts.rotation_period)
352        .build()?;
353
354    let mut traced = TracedRuntime::builder()
355        .with_trace_path(opts.trace_path.clone())
356        .with_runtime_name(runtime_name)
357        .with_task_tracking(opts.task_tracking);
358    if let Some(worker_poll_interval) = opts.worker_poll_interval {
359        traced = traced.with_worker_poll_interval(worker_poll_interval);
360    }
361
362    #[cfg(feature = "dial9-worker-s3")]
363    if let Some(s3_upload) = &opts.s3_upload {
364        if s3_upload.bucket.trim().is_empty() {
365            return Err(std::io::Error::new(
366                std::io::ErrorKind::InvalidInput,
367                "dial9 s3 bucket must not be empty",
368            ));
369        }
370        if s3_upload.service_name.trim().is_empty() {
371            return Err(std::io::Error::new(
372                std::io::ErrorKind::InvalidInput,
373                "dial9 s3 service_name must not be empty",
374            ));
375        }
376        let s3_config = dial9_tokio_telemetry::background_task::s3::S3Config::builder()
377            .bucket(s3_upload.bucket.clone())
378            .service_name(s3_upload.service_name.clone())
379            .maybe_prefix(s3_upload.prefix.clone())
380            .maybe_region(s3_upload.region.clone())
381            .maybe_instance_path(s3_upload.instance_path.clone());
382        let traced = traced.with_s3_uploader(s3_config.build());
383        if let Some(client) = s3_upload.client.clone() {
384            return traced
385                .with_s3_client(client)
386                .build_and_start(builder, writer);
387        }
388        return traced.build_and_start(builder, writer);
389    }
390
391    traced.build_and_start(builder, writer)
392}
393
394/// Builder for constructing a [`Runtime`].
395///
396/// # Example
397///
398/// ```
399/// use pingora_runtime::{RuntimeBuilder, BlockingPoolOpts};
400/// use std::time::Duration;
401///
402/// let rt = RuntimeBuilder::new(4, "my-service")
403///     .blocking_pool_opts(BlockingPoolOpts {
404///         max_threads: Some(64),
405///         thread_keep_alive: Some(Duration::from_secs(30)),
406///     })
407///     .build();
408/// ```
409pub struct RuntimeBuilder {
410    threads: usize,
411    name: String,
412    work_steal: bool,
413    blocking_pool_opts: BlockingPoolOpts,
414    runtime_opts: RuntimeOpts,
415}
416
417impl RuntimeBuilder {
418    /// Create a new builder with the given number of worker threads and runtime name.
419    ///
420    /// Work stealing is enabled by default.
421    pub fn new(threads: usize, name: &str) -> Self {
422        Self {
423            threads,
424            name: name.to_string(),
425            work_steal: true,
426            blocking_pool_opts: BlockingPoolOpts::default(),
427            runtime_opts: RuntimeOpts::default(),
428        }
429    }
430
431    /// Set whether work stealing is enabled.
432    ///
433    /// When `true` (the default), a tokio multi-thread runtime is used.
434    /// When `false`, a pool of single-threaded tokio runtimes is used instead.
435    pub fn work_steal(mut self, enabled: bool) -> Self {
436        self.work_steal = enabled;
437        self
438    }
439
440    /// Set the [`BlockingPoolOpts`] for the runtime's blocking thread pool.
441    pub fn blocking_pool_opts(mut self, opts: BlockingPoolOpts) -> Self {
442        self.blocking_pool_opts = opts;
443        self
444    }
445
446    /// Set the [`RuntimeMetricsOpts`] for the runtime.
447    pub fn metrics_opts(mut self, opts: RuntimeMetricsOpts) -> Self {
448        self.runtime_opts.metrics = opts;
449        self
450    }
451
452    /// Set the [`RuntimeOpts`] for the runtime.
453    pub fn runtime_opts(mut self, opts: RuntimeOpts) -> Self {
454        self.runtime_opts = opts;
455        self
456    }
457
458    /// Set whether Tokio's experimental alternative timer is enabled.
459    ///
460    /// This requires building with `--cfg tokio_unstable` and only applies to
461    /// work-stealing runtimes.
462    pub fn enable_alt_timer(mut self, enabled: bool) -> Self {
463        self.runtime_opts.enable_alt_timer = enabled;
464        self
465    }
466
467    fn build_work_stealing_tokio_builder(&self) -> Builder {
468        let mut builder = Builder::new_multi_thread();
469        builder
470            .enable_all()
471            .worker_threads(self.threads)
472            .thread_name(&self.name);
473        apply_blocking_opts(&mut builder, &self.blocking_pool_opts);
474        apply_metrics_opts(&mut builder, &self.runtime_opts.metrics);
475        apply_timer_opts(&mut builder, &self.runtime_opts);
476        builder
477    }
478
479    /// Build the [`Runtime`].
480    pub fn build(self) -> Runtime {
481        if self.work_steal {
482            let mut builder = self.build_work_stealing_tokio_builder();
483            #[cfg(feature = "dial9")]
484            let dial9_guard = if let Some(dial9_opts) = &self.runtime_opts.dial9 {
485                let runtime_name = self.name.clone();
486                match build_dial9_runtime(builder, &runtime_name, dial9_opts) {
487                    Ok((runtime, guard)) => {
488                        return Runtime::Steal {
489                            runtime,
490                            dial9_guard: Some(guard),
491                        };
492                    }
493                    Err(e) => {
494                        log::warn!(
495                            "failed to initialize dial9 runtime telemetry for {runtime_name}: {e}"
496                        );
497                        builder = self.build_work_stealing_tokio_builder();
498                        None
499                    }
500                }
501            } else {
502                None
503            };
504            let runtime = builder
505                .build()
506                .expect("failed to build work-stealing Tokio runtime");
507            Runtime::Steal {
508                runtime,
509                #[cfg(feature = "dial9")]
510                dial9_guard,
511            }
512        } else {
513            #[cfg(feature = "dial9")]
514            if self.runtime_opts.dial9.is_some() {
515                log::warn!("dial9 runtime telemetry is ignored when work stealing is disabled");
516            }
517            Runtime::NoSteal(NoStealRuntime::new(
518                self.threads,
519                &self.name,
520                self.blocking_pool_opts,
521                self.runtime_opts,
522            ))
523        }
524    }
525}
526
527impl Runtime {
528    /// Create a `Steal` flavor runtime. This just a regular tokio runtime
529    pub fn new_steal(threads: usize, name: &str) -> Self {
530        RuntimeBuilder::new(threads, name).build()
531    }
532
533    /// Create a `NoSteal` flavor runtime. This is backed by multiple tokio current-thread runtime
534    pub fn new_no_steal(threads: usize, name: &str) -> Self {
535        RuntimeBuilder::new(threads, name).work_steal(false).build()
536    }
537
538    /// Return the &[Handle] of the [Runtime].
539    /// For `Steal` flavor, it will just return the &[Handle].
540    /// For `NoSteal` flavor, it will return the &[Handle] of a random thread in its pool.
541    /// So if we want tasks to spawn on all the threads, call this function to get a fresh [Handle]
542    /// for each async task.
543    pub fn get_handle(&self) -> &Handle {
544        match self {
545            Self::Steal { runtime, .. } => runtime.handle(),
546            Self::NoSteal(r) => r.get_runtime(),
547        }
548    }
549
550    /// Call tokio's `shutdown_timeout` of all the runtimes. This function is blocking until
551    /// all runtimes exit.
552    pub fn shutdown_timeout(self, timeout: Duration) {
553        match self {
554            Self::Steal {
555                runtime,
556                #[cfg(feature = "dial9")]
557                dial9_guard,
558            } => {
559                #[cfg(feature = "dial9")]
560                drop(dial9_guard);
561                runtime.shutdown_timeout(timeout);
562            }
563            Self::NoSteal(r) => r.shutdown_timeout(timeout),
564        }
565    }
566}
567
568// only NoStealRuntime set the pools in thread threads
569static CURRENT_HANDLE: Lazy<ThreadLocal<Pools>> = Lazy::new(ThreadLocal::new);
570
571/// Return the [Handle] of current runtime.
572/// If the current thread is under a `Steal` runtime, the current [Handle] is returned.
573/// If the current thread is under a `NoSteal` runtime, the [Handle] of a random thread
574/// under this runtime is returned. This function will panic if called outside any runtime.
575pub fn current_handle() -> Handle {
576    if let Some(pools) = CURRENT_HANDLE.get() {
577        // safety: the CURRENT_HANDLE is set when the pool is being initialized in init_pools()
578        let pools = pools.get().unwrap();
579        let mut rng = rand::thread_rng();
580        let index = rng.gen_range(0..pools.len());
581        pools[index].clone()
582    } else {
583        // not NoStealRuntime, just check the current tokio runtime
584        Handle::current()
585    }
586}
587
588type Control = (Sender<Duration>, JoinHandle<()>);
589type Pools = Arc<OnceCell<Box<[Handle]>>>;
590
591/// Multi-threaded runtime backed by a pool of single threaded tokio runtime
592pub struct NoStealRuntime {
593    threads: usize,
594    name: String,
595    blocking_opts: BlockingPoolOpts,
596    runtime_opts: RuntimeOpts,
597    // Lazily init the runtimes so that they are created after pingora
598    // daemonize itself. Otherwise the runtime threads are lost.
599    pools: Pools,
600    controls: OnceCell<Vec<Control>>,
601}
602
603impl NoStealRuntime {
604    /// Create a new [`NoStealRuntime`] with blocking pool options. Panic if `threads` is 0.
605    pub fn new(
606        threads: usize,
607        name: &str,
608        blocking_opts: BlockingPoolOpts,
609        runtime_opts: RuntimeOpts,
610    ) -> Self {
611        assert!(threads != 0);
612        NoStealRuntime {
613            threads,
614            name: name.to_string(),
615            blocking_opts,
616            runtime_opts,
617            pools: Arc::new(OnceCell::new()),
618            controls: OnceCell::new(),
619        }
620    }
621
622    fn init_pools(&self) -> (Box<[Handle]>, Vec<Control>) {
623        let mut pools = Vec::with_capacity(self.threads);
624        let mut controls = Vec::with_capacity(self.threads);
625        for _ in 0..self.threads {
626            let mut builder = Builder::new_current_thread();
627            builder.enable_all();
628            apply_blocking_opts(&mut builder, &self.blocking_opts);
629            apply_metrics_opts(&mut builder, &self.runtime_opts.metrics);
630            let rt = builder
631                .build()
632                .expect("failed to build no-steal Tokio runtime worker");
633            let handler = rt.handle().clone();
634            let (tx, rx) = channel::<Duration>();
635            let pools_ref = self.pools.clone();
636            let join = std::thread::Builder::new()
637                .name(self.name.clone())
638                .spawn(move || {
639                    CURRENT_HANDLE.get_or(|| pools_ref);
640                    if let Ok(timeout) = rt.block_on(rx) {
641                        rt.shutdown_timeout(timeout);
642                    } // else Err(_): tx is dropped, just exit
643                })
644                .unwrap();
645            pools.push(handler);
646            controls.push((tx, join));
647        }
648
649        (pools.into_boxed_slice(), controls)
650    }
651
652    /// Return the &[Handle] of a random thread of this runtime
653    pub fn get_runtime(&self) -> &Handle {
654        let mut rng = rand::thread_rng();
655
656        let index = rng.gen_range(0..self.threads);
657        self.get_runtime_at(index)
658    }
659
660    /// Return the number of threads of this runtime
661    pub fn threads(&self) -> usize {
662        self.threads
663    }
664
665    fn get_pools(&self) -> &[Handle] {
666        if let Some(p) = self.pools.get() {
667            p
668        } else {
669            // TODO: use a mutex to avoid creating a lot threads only to drop them
670            let (pools, controls) = self.init_pools();
671            // there could be another thread racing with this one to init the pools
672            match self.pools.try_insert(pools) {
673                Ok(p) => {
674                    // unwrap to make sure that this is the one that init both pools and controls
675                    self.controls.set(controls).unwrap();
676                    p
677                }
678                // another thread already set it, just return it
679                Err((p, _my_pools)) => p,
680            }
681        }
682    }
683
684    /// Return the &[Handle] of a given thread of this runtime
685    pub fn get_runtime_at(&self, index: usize) -> &Handle {
686        let pools = self.get_pools();
687        &pools[index]
688    }
689
690    /// Call tokio's `shutdown_timeout` of all the runtimes. This function is blocking until
691    /// all runtimes exit.
692    pub fn shutdown_timeout(mut self, timeout: Duration) {
693        if let Some(controls) = self.controls.take() {
694            let (txs, joins): (Vec<Sender<_>>, Vec<JoinHandle<()>>) = controls.into_iter().unzip();
695            for tx in txs {
696                let _ = tx.send(timeout); // Err() when rx is dropped
697            }
698            for join in joins {
699                let _ = join.join(); // ignore thread error
700            }
701        } // else, the controls and the runtimes are not even init yet, just return;
702    }
703
704    // TODO: runtime metrics
705}
706
707#[test]
708fn test_steal_runtime() {
709    use tokio::time::{sleep, Duration};
710    let threads = 2;
711    let rt = Runtime::new_steal(threads, "test");
712    let handle = rt.get_handle();
713    let ret = handle.block_on(async {
714        sleep(Duration::from_secs(1)).await;
715        let handle = current_handle();
716        let join = handle.spawn(async {
717            sleep(Duration::from_secs(1)).await;
718        });
719        join.await.unwrap();
720        1
721    });
722
723    #[cfg(target_os = "linux")]
724    assert_eq!(handle.metrics().num_workers(), threads);
725    assert_eq!(ret, 1);
726}
727
728#[test]
729fn test_no_steal_runtime() {
730    use tokio::time::{sleep, Duration};
731
732    let rt = Runtime::new_no_steal(2, "test");
733    let handle = rt.get_handle();
734    let ret = handle.block_on(async {
735        sleep(Duration::from_secs(1)).await;
736        let handle = current_handle();
737        let join = handle.spawn(async {
738            sleep(Duration::from_secs(1)).await;
739        });
740        join.await.unwrap();
741        1
742    });
743
744    assert_eq!(ret, 1);
745}
746
747#[test]
748fn test_no_steal_shutdown() {
749    use tokio::time::{sleep, Duration};
750
751    let rt = Runtime::new_no_steal(2, "test");
752    let handle = rt.get_handle();
753    let ret = handle.block_on(async {
754        sleep(Duration::from_secs(1)).await;
755        let handle = current_handle();
756        let join = handle.spawn(async {
757            sleep(Duration::from_secs(1)).await;
758        });
759        join.await.unwrap();
760        1
761    });
762    assert_eq!(ret, 1);
763
764    rt.shutdown_timeout(Duration::from_secs(1));
765}
766
767#[cfg(feature = "dial9")]
768#[test]
769fn test_dial9_zero_worker_poll_interval_is_rejected() {
770    let mut opts = Dial9RuntimeOpts::new("trace");
771    opts.worker_poll_interval = Some(Duration::ZERO);
772    let err = match build_dial9_runtime(Builder::new_multi_thread(), "test", &opts) {
773        Ok(_) => panic!("zero worker poll interval should be rejected"),
774        Err(err) => err,
775    };
776
777    assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
778    assert_eq!(
779        err.to_string(),
780        "dial9 worker_poll_interval must be greater than zero"
781    );
782}