1use 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#[cfg(feature = "dial9")]
40pub const DEFAULT_DIAL9_MAX_FILE_SIZE: u64 = 100 * 1024 * 1024;
41#[cfg(feature = "dial9")]
43pub const DEFAULT_DIAL9_MAX_TOTAL_SIZE: u64 = 512 * 1024 * 1024;
44
45#[derive(Debug, Clone, Default)]
50pub struct BlockingPoolOpts {
51 pub max_threads: Option<usize>,
55 pub thread_keep_alive: Option<Duration>,
59}
60
61#[derive(Debug, Clone, Default)]
63pub struct RuntimeMetricsOpts {
64 pub poll_time_histogram: bool,
69 pub poll_time_histogram_scale: Option<RuntimeMetricsPollTimeHistogramScale>,
71 pub poll_time_histogram_resolution: Option<Duration>,
73 pub poll_time_histogram_buckets: Option<usize>,
75}
76
77#[derive(Debug, Clone, Default)]
79pub struct RuntimeOpts {
80 pub metrics: RuntimeMetricsOpts,
82 pub enable_alt_timer: bool,
87 #[cfg(feature = "dial9")]
89 pub dial9: Option<Dial9RuntimeOpts>,
90}
91
92#[cfg(feature = "dial9")]
94#[derive(Debug, Clone)]
95pub struct Dial9RuntimeOpts {
96 pub trace_path: PathBuf,
98 pub max_file_size: u64,
100 pub max_total_size: u64,
102 pub rotation_period: Option<Duration>,
104 pub task_tracking: bool,
106 pub worker_poll_interval: Option<Duration>,
108 #[cfg(feature = "dial9-worker-s3")]
110 pub s3_upload: Option<Dial9S3UploadOpts>,
111}
112
113#[cfg(feature = "dial9")]
114impl Dial9RuntimeOpts {
115 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 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 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 pub fn with_rotation_period(mut self, rotation_period: Duration) -> Self {
143 self.rotation_period = Some(rotation_period);
144 self
145 }
146
147 pub fn with_task_tracking(mut self, task_tracking: bool) -> Self {
149 self.task_tracking = task_tracking;
150 self
151 }
152
153 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 #[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#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))]
169#[derive(Debug, Clone)]
170pub struct Dial9S3UploadOpts {
171 pub bucket: String,
173 pub service_name: String,
175 pub prefix: Option<String>,
177 pub region: Option<String>,
179 pub instance_path: Option<String>,
181 pub client: Option<aws_sdk_s3::Client>,
183}
184
185#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))]
186impl Dial9S3UploadOpts {
187 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 pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
201 self.prefix = Some(prefix.into());
202 self
203 }
204
205 pub fn with_region(mut self, region: impl Into<String>) -> Self {
207 self.region = Some(region.into());
208 self
209 }
210
211 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 pub fn with_client(mut self, client: aws_sdk_s3::Client) -> Self {
219 self.client = Some(client);
220 self
221 }
222}
223
224#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
226#[serde(rename_all = "snake_case")]
227pub enum RuntimeMetricsPollTimeHistogramScale {
228 Linear,
230 Log,
232}
233
234pub 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
248fn 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#[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
295fn 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
394pub 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 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 pub fn work_steal(mut self, enabled: bool) -> Self {
436 self.work_steal = enabled;
437 self
438 }
439
440 pub fn blocking_pool_opts(mut self, opts: BlockingPoolOpts) -> Self {
442 self.blocking_pool_opts = opts;
443 self
444 }
445
446 pub fn metrics_opts(mut self, opts: RuntimeMetricsOpts) -> Self {
448 self.runtime_opts.metrics = opts;
449 self
450 }
451
452 pub fn runtime_opts(mut self, opts: RuntimeOpts) -> Self {
454 self.runtime_opts = opts;
455 self
456 }
457
458 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 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 pub fn new_steal(threads: usize, name: &str) -> Self {
530 RuntimeBuilder::new(threads, name).build()
531 }
532
533 pub fn new_no_steal(threads: usize, name: &str) -> Self {
535 RuntimeBuilder::new(threads, name).work_steal(false).build()
536 }
537
538 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 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
568static CURRENT_HANDLE: Lazy<ThreadLocal<Pools>> = Lazy::new(ThreadLocal::new);
570
571pub fn current_handle() -> Handle {
576 if let Some(pools) = CURRENT_HANDLE.get() {
577 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 Handle::current()
585 }
586}
587
588type Control = (Sender<Duration>, JoinHandle<()>);
589type Pools = Arc<OnceCell<Box<[Handle]>>>;
590
591pub struct NoStealRuntime {
593 threads: usize,
594 name: String,
595 blocking_opts: BlockingPoolOpts,
596 runtime_opts: RuntimeOpts,
597 pools: Pools,
600 controls: OnceCell<Vec<Control>>,
601}
602
603impl NoStealRuntime {
604 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 } })
644 .unwrap();
645 pools.push(handler);
646 controls.push((tx, join));
647 }
648
649 (pools.into_boxed_slice(), controls)
650 }
651
652 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 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 let (pools, controls) = self.init_pools();
671 match self.pools.try_insert(pools) {
673 Ok(p) => {
674 self.controls.set(controls).unwrap();
676 p
677 }
678 Err((p, _my_pools)) => p,
680 }
681 }
682 }
683
684 pub fn get_runtime_at(&self, index: usize) -> &Handle {
686 let pools = self.get_pools();
687 &pools[index]
688 }
689
690 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); }
698 for join in joins {
699 let _ = join.join(); }
701 } }
703
704 }
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}