Skip to content

Commit

Permalink
add thread_name_fn method to runtime::Builder, close #1907.
Browse files Browse the repository at this point in the history
Co-Authored-By: Eliza Weisman <eliza@buoyant.io>
  • Loading branch information
biluohc and hawkw committed Apr 10, 2020
1 parent 3137c6f commit 81794ef
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 8 deletions.
5 changes: 3 additions & 2 deletions tokio/src/runtime/blocking/pool.rs
Expand Up @@ -5,6 +5,7 @@ use crate::loom::thread;
use crate::runtime::blocking::schedule::NoopSchedule;
use crate::runtime::blocking::shutdown;
use crate::runtime::blocking::task::BlockingTask;
use crate::runtime::builder::ThreadNameFn;
use crate::runtime::task::{self, JoinHandle};
use crate::runtime::{Builder, Callback, Handle};

Expand All @@ -30,7 +31,7 @@ struct Inner {
condvar: Condvar,

/// Spawned threads use this name
thread_name: String,
thread_name: ThreadNameFn,

/// Spawned thread stack size
stack_size: Option<usize>,
Expand Down Expand Up @@ -194,7 +195,7 @@ impl Spawner {
}

fn spawn_thread(&self, shutdown_tx: shutdown::Sender, rt: &Handle) {
let mut builder = thread::Builder::new().name(self.inner.thread_name.clone());
let mut builder = thread::Builder::new().name((self.inner.thread_name)());

if let Some(stack_size) = self.inner.stack_size {
builder = builder.stack_size(stack_size);
Expand Down
45 changes: 39 additions & 6 deletions tokio/src/runtime/builder.rs
Expand Up @@ -3,7 +3,6 @@ use crate::runtime::shell::Shell;
use crate::runtime::{blocking, io, time, Callback, Runtime, Spawner};

use std::fmt;
#[cfg(not(loom))]
use std::sync::Arc;

/// Builds Tokio Runtime with custom configuration values.
Expand Down Expand Up @@ -55,8 +54,8 @@ pub struct Builder {
/// Cap on thread usage.
max_threads: usize,

/// Name used for threads spawned by the runtime.
pub(super) thread_name: String,
/// Name fn used for threads spawned by the runtime.
pub(super) thread_name: ThreadNameFn,

/// Stack size used for threads spawned by the runtime.
pub(super) thread_stack_size: Option<usize>,
Expand All @@ -68,6 +67,8 @@ pub struct Builder {
pub(super) before_stop: Option<Callback>,
}

pub(crate) type ThreadNameFn = Arc<dyn Fn() -> String + Send + Sync + 'static>;

#[derive(Debug, Clone, Copy)]
enum Kind {
Shell,
Expand Down Expand Up @@ -99,7 +100,7 @@ impl Builder {
max_threads: 512,

// Default thread name
thread_name: "tokio-runtime-worker".into(),
thread_name: Arc::new(|| "tokio-runtime-worker".into()),

// Do not set a stack size by default
thread_stack_size: None,
Expand Down Expand Up @@ -210,7 +211,36 @@ impl Builder {
/// # }
/// ```
pub fn thread_name(&mut self, val: impl Into<String>) -> &mut Self {
self.thread_name = val.into();
let val = val.into();
self.thread_name = Arc::new(move || val.clone());
self
}

/// Sets a function used to generate the name of threads spawned by the `Runtime`'s thread pool.
///
/// The default name fn is `|| "tokio-runtime-worker".into()`.
///
/// # Examples
///
/// ```
/// # use tokio::runtime;
/// # use std::sync::atomic::{AtomicUsize, Ordering};
///
/// # pub fn main() {
/// let rt = runtime::Builder::new()
/// .thread_name_fn(|| {
/// static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
/// let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
/// format!("my-pool-{}", id)
/// })
/// .build();
/// # }
/// ```
pub fn thread_name_fn<F>(&mut self, f: F) -> &mut Self
where
F: Fn() -> String + Send + Sync + 'static,
{
self.thread_name = Arc::new(f);
self
}

Expand Down Expand Up @@ -510,7 +540,10 @@ impl fmt::Debug for Builder {
.field("kind", &self.kind)
.field("core_threads", &self.core_threads)
.field("max_threads", &self.max_threads)
.field("thread_name", &self.thread_name)
.field(
"thread_name",
&"<dyn Fn() -> String + Send + Sync + 'static>",
)
.field("thread_stack_size", &self.thread_stack_size)
.field("after_start", &self.after_start.as_ref().map(|_| "..."))
.field("before_stop", &self.after_start.as_ref().map(|_| "..."))
Expand Down

0 comments on commit 81794ef

Please sign in to comment.