summaryrefslogtreecommitdiffstats
path: root/third_party/rust/tokio-macros/src
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 01:47:29 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 01:47:29 +0000
commit0ebf5bdf043a27fd3dfb7f92e0cb63d88954c44d (patch)
treea31f07c9bcca9d56ce61e9a1ffd30ef350d513aa /third_party/rust/tokio-macros/src
parentInitial commit. (diff)
downloadfirefox-esr-0ebf5bdf043a27fd3dfb7f92e0cb63d88954c44d.tar.xz
firefox-esr-0ebf5bdf043a27fd3dfb7f92e0cb63d88954c44d.zip
Adding upstream version 115.8.0esr.upstream/115.8.0esr
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'third_party/rust/tokio-macros/src')
-rw-r--r--third_party/rust/tokio-macros/src/entry.rs494
-rw-r--r--third_party/rust/tokio-macros/src/lib.rs483
-rw-r--r--third_party/rust/tokio-macros/src/select.rs110
3 files changed, 1087 insertions, 0 deletions
diff --git a/third_party/rust/tokio-macros/src/entry.rs b/third_party/rust/tokio-macros/src/entry.rs
new file mode 100644
index 0000000000..6460e70afa
--- /dev/null
+++ b/third_party/rust/tokio-macros/src/entry.rs
@@ -0,0 +1,494 @@
+use proc_macro::TokenStream;
+use proc_macro2::{Ident, Span};
+use quote::{quote, quote_spanned, ToTokens};
+use syn::parse::Parser;
+
+// syn::AttributeArgs does not implement syn::Parse
+type AttributeArgs = syn::punctuated::Punctuated<syn::NestedMeta, syn::Token![,]>;
+
+#[derive(Clone, Copy, PartialEq)]
+enum RuntimeFlavor {
+ CurrentThread,
+ Threaded,
+}
+
+impl RuntimeFlavor {
+ fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
+ match s {
+ "current_thread" => Ok(RuntimeFlavor::CurrentThread),
+ "multi_thread" => Ok(RuntimeFlavor::Threaded),
+ "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
+ "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
+ "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
+ _ => Err(format!("No such runtime flavor `{}`. The runtime flavors are `current_thread` and `multi_thread`.", s)),
+ }
+ }
+}
+
+struct FinalConfig {
+ flavor: RuntimeFlavor,
+ worker_threads: Option<usize>,
+ start_paused: Option<bool>,
+ crate_name: Option<String>,
+}
+
+/// Config used in case of the attribute not being able to build a valid config
+const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
+ flavor: RuntimeFlavor::CurrentThread,
+ worker_threads: None,
+ start_paused: None,
+ crate_name: None,
+};
+
+struct Configuration {
+ rt_multi_thread_available: bool,
+ default_flavor: RuntimeFlavor,
+ flavor: Option<RuntimeFlavor>,
+ worker_threads: Option<(usize, Span)>,
+ start_paused: Option<(bool, Span)>,
+ is_test: bool,
+ crate_name: Option<String>,
+}
+
+impl Configuration {
+ fn new(is_test: bool, rt_multi_thread: bool) -> Self {
+ Configuration {
+ rt_multi_thread_available: rt_multi_thread,
+ default_flavor: match is_test {
+ true => RuntimeFlavor::CurrentThread,
+ false => RuntimeFlavor::Threaded,
+ },
+ flavor: None,
+ worker_threads: None,
+ start_paused: None,
+ is_test,
+ crate_name: None,
+ }
+ }
+
+ fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> {
+ if self.flavor.is_some() {
+ return Err(syn::Error::new(span, "`flavor` set multiple times."));
+ }
+
+ let runtime_str = parse_string(runtime, span, "flavor")?;
+ let runtime =
+ RuntimeFlavor::from_str(&runtime_str).map_err(|err| syn::Error::new(span, err))?;
+ self.flavor = Some(runtime);
+ Ok(())
+ }
+
+ fn set_worker_threads(
+ &mut self,
+ worker_threads: syn::Lit,
+ span: Span,
+ ) -> Result<(), syn::Error> {
+ if self.worker_threads.is_some() {
+ return Err(syn::Error::new(
+ span,
+ "`worker_threads` set multiple times.",
+ ));
+ }
+
+ let worker_threads = parse_int(worker_threads, span, "worker_threads")?;
+ if worker_threads == 0 {
+ return Err(syn::Error::new(span, "`worker_threads` may not be 0."));
+ }
+ self.worker_threads = Some((worker_threads, span));
+ Ok(())
+ }
+
+ fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
+ if self.start_paused.is_some() {
+ return Err(syn::Error::new(span, "`start_paused` set multiple times."));
+ }
+
+ let start_paused = parse_bool(start_paused, span, "start_paused")?;
+ self.start_paused = Some((start_paused, span));
+ Ok(())
+ }
+
+ fn set_crate_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> {
+ if self.crate_name.is_some() {
+ return Err(syn::Error::new(span, "`crate` set multiple times."));
+ }
+ let name_ident = parse_ident(name, span, "crate")?;
+ self.crate_name = Some(name_ident.to_string());
+ Ok(())
+ }
+
+ fn macro_name(&self) -> &'static str {
+ if self.is_test {
+ "tokio::test"
+ } else {
+ "tokio::main"
+ }
+ }
+
+ fn build(&self) -> Result<FinalConfig, syn::Error> {
+ let flavor = self.flavor.unwrap_or(self.default_flavor);
+ use RuntimeFlavor::*;
+
+ let worker_threads = match (flavor, self.worker_threads) {
+ (CurrentThread, Some((_, worker_threads_span))) => {
+ let msg = format!(
+ "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
+ self.macro_name(),
+ );
+ return Err(syn::Error::new(worker_threads_span, msg));
+ }
+ (CurrentThread, None) => None,
+ (Threaded, worker_threads) if self.rt_multi_thread_available => {
+ worker_threads.map(|(val, _span)| val)
+ }
+ (Threaded, _) => {
+ let msg = if self.flavor.is_none() {
+ "The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled."
+ } else {
+ "The runtime flavor `multi_thread` requires the `rt-multi-thread` feature."
+ };
+ return Err(syn::Error::new(Span::call_site(), msg));
+ }
+ };
+
+ let start_paused = match (flavor, self.start_paused) {
+ (Threaded, Some((_, start_paused_span))) => {
+ let msg = format!(
+ "The `start_paused` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
+ self.macro_name(),
+ );
+ return Err(syn::Error::new(start_paused_span, msg));
+ }
+ (CurrentThread, Some((start_paused, _))) => Some(start_paused),
+ (_, None) => None,
+ };
+
+ Ok(FinalConfig {
+ crate_name: self.crate_name.clone(),
+ flavor,
+ worker_threads,
+ start_paused,
+ })
+ }
+}
+
+fn parse_int(int: syn::Lit, span: Span, field: &str) -> Result<usize, syn::Error> {
+ match int {
+ syn::Lit::Int(lit) => match lit.base10_parse::<usize>() {
+ Ok(value) => Ok(value),
+ Err(e) => Err(syn::Error::new(
+ span,
+ format!("Failed to parse value of `{}` as integer: {}", field, e),
+ )),
+ },
+ _ => Err(syn::Error::new(
+ span,
+ format!("Failed to parse value of `{}` as integer.", field),
+ )),
+ }
+}
+
+fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
+ match int {
+ syn::Lit::Str(s) => Ok(s.value()),
+ syn::Lit::Verbatim(s) => Ok(s.to_string()),
+ _ => Err(syn::Error::new(
+ span,
+ format!("Failed to parse value of `{}` as string.", field),
+ )),
+ }
+}
+
+fn parse_ident(lit: syn::Lit, span: Span, field: &str) -> Result<Ident, syn::Error> {
+ match lit {
+ syn::Lit::Str(s) => {
+ let err = syn::Error::new(
+ span,
+ format!(
+ "Failed to parse value of `{}` as ident: \"{}\"",
+ field,
+ s.value()
+ ),
+ );
+ let path = s.parse::<syn::Path>().map_err(|_| err.clone())?;
+ path.get_ident().cloned().ok_or(err)
+ }
+ _ => Err(syn::Error::new(
+ span,
+ format!("Failed to parse value of `{}` as ident.", field),
+ )),
+ }
+}
+
+fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
+ match bool {
+ syn::Lit::Bool(b) => Ok(b.value),
+ _ => Err(syn::Error::new(
+ span,
+ format!("Failed to parse value of `{}` as bool.", field),
+ )),
+ }
+}
+
+fn build_config(
+ input: syn::ItemFn,
+ args: AttributeArgs,
+ is_test: bool,
+ rt_multi_thread: bool,
+) -> Result<FinalConfig, syn::Error> {
+ if input.sig.asyncness.is_none() {
+ let msg = "the `async` keyword is missing from the function declaration";
+ return Err(syn::Error::new_spanned(input.sig.fn_token, msg));
+ }
+
+ let mut config = Configuration::new(is_test, rt_multi_thread);
+ let macro_name = config.macro_name();
+
+ for arg in args {
+ match arg {
+ syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => {
+ let ident = namevalue
+ .path
+ .get_ident()
+ .ok_or_else(|| {
+ syn::Error::new_spanned(&namevalue, "Must have specified ident")
+ })?
+ .to_string()
+ .to_lowercase();
+ match ident.as_str() {
+ "worker_threads" => {
+ config.set_worker_threads(
+ namevalue.lit.clone(),
+ syn::spanned::Spanned::span(&namevalue.lit),
+ )?;
+ }
+ "flavor" => {
+ config.set_flavor(
+ namevalue.lit.clone(),
+ syn::spanned::Spanned::span(&namevalue.lit),
+ )?;
+ }
+ "start_paused" => {
+ config.set_start_paused(
+ namevalue.lit.clone(),
+ syn::spanned::Spanned::span(&namevalue.lit),
+ )?;
+ }
+ "core_threads" => {
+ let msg = "Attribute `core_threads` is renamed to `worker_threads`";
+ return Err(syn::Error::new_spanned(namevalue, msg));
+ }
+ "crate" => {
+ config.set_crate_name(
+ namevalue.lit.clone(),
+ syn::spanned::Spanned::span(&namevalue.lit),
+ )?;
+ }
+ name => {
+ let msg = format!(
+ "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`",
+ name,
+ );
+ return Err(syn::Error::new_spanned(namevalue, msg));
+ }
+ }
+ }
+ syn::NestedMeta::Meta(syn::Meta::Path(path)) => {
+ let name = path
+ .get_ident()
+ .ok_or_else(|| syn::Error::new_spanned(&path, "Must have specified ident"))?
+ .to_string()
+ .to_lowercase();
+ let msg = match name.as_str() {
+ "threaded_scheduler" | "multi_thread" => {
+ format!(
+ "Set the runtime flavor with #[{}(flavor = \"multi_thread\")].",
+ macro_name
+ )
+ }
+ "basic_scheduler" | "current_thread" | "single_threaded" => {
+ format!(
+ "Set the runtime flavor with #[{}(flavor = \"current_thread\")].",
+ macro_name
+ )
+ }
+ "flavor" | "worker_threads" | "start_paused" => {
+ format!("The `{}` attribute requires an argument.", name)
+ }
+ name => {
+ format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`", name)
+ }
+ };
+ return Err(syn::Error::new_spanned(path, msg));
+ }
+ other => {
+ return Err(syn::Error::new_spanned(
+ other,
+ "Unknown attribute inside the macro",
+ ));
+ }
+ }
+ }
+
+ config.build()
+}
+
+fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> TokenStream {
+ input.sig.asyncness = None;
+
+ // If type mismatch occurs, the current rustc points to the last statement.
+ let (last_stmt_start_span, last_stmt_end_span) = {
+ let mut last_stmt = input
+ .block
+ .stmts
+ .last()
+ .map(ToTokens::into_token_stream)
+ .unwrap_or_default()
+ .into_iter();
+ // `Span` on stable Rust has a limitation that only points to the first
+ // token, not the whole tokens. We can work around this limitation by
+ // using the first/last span of the tokens like
+ // `syn::Error::new_spanned` does.
+ let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
+ let end = last_stmt.last().map_or(start, |t| t.span());
+ (start, end)
+ };
+
+ let crate_name = config.crate_name.as_deref().unwrap_or("tokio");
+
+ let crate_ident = Ident::new(crate_name, last_stmt_start_span);
+
+ let mut rt = match config.flavor {
+ RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
+ #crate_ident::runtime::Builder::new_current_thread()
+ },
+ RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
+ #crate_ident::runtime::Builder::new_multi_thread()
+ },
+ };
+ if let Some(v) = config.worker_threads {
+ rt = quote! { #rt.worker_threads(#v) };
+ }
+ if let Some(v) = config.start_paused {
+ rt = quote! { #rt.start_paused(#v) };
+ }
+
+ let header = if is_test {
+ quote! {
+ #[::core::prelude::v1::test]
+ }
+ } else {
+ quote! {}
+ };
+
+ let body = &input.block;
+ let brace_token = input.block.brace_token;
+ let body_ident = quote! { body };
+ let block_expr = quote_spanned! {last_stmt_end_span=>
+ #[allow(clippy::expect_used, clippy::diverging_sub_expression)]
+ {
+ return #rt
+ .enable_all()
+ .build()
+ .expect("Failed building the Runtime")
+ .block_on(#body_ident);
+ }
+ };
+
+ // For test functions pin the body to the stack and use `Pin<&mut dyn
+ // Future>` to reduce the amount of `Runtime::block_on` (and related
+ // functions) copies we generate during compilation due to the generic
+ // parameter `F` (the future to block on). This could have an impact on
+ // performance, but because it's only for testing it's unlikely to be very
+ // large.
+ //
+ // We don't do this for the main function as it should only be used once so
+ // there will be no benefit.
+ let body = if is_test {
+ let output_type = match &input.sig.output {
+ // For functions with no return value syn doesn't print anything,
+ // but that doesn't work as `Output` for our boxed `Future`, so
+ // default to `()` (the same type as the function output).
+ syn::ReturnType::Default => quote! { () },
+ syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
+ };
+ quote! {
+ let body = async #body;
+ #crate_ident::pin!(body);
+ let body: ::std::pin::Pin<&mut dyn ::std::future::Future<Output = #output_type>> = body;
+ }
+ } else {
+ quote! {
+ let body = async #body;
+ }
+ };
+
+ input.block = syn::parse2(quote! {
+ {
+ #body
+ #block_expr
+ }
+ })
+ .expect("Parsing failure");
+ input.block.brace_token = brace_token;
+
+ let result = quote! {
+ #header
+ #input
+ };
+
+ result.into()
+}
+
+fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
+ tokens.extend(TokenStream::from(error.into_compile_error()));
+ tokens
+}
+
+#[cfg(not(test))] // Work around for rust-lang/rust#62127
+pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
+ // If any of the steps for this macro fail, we still want to expand to an item that is as close
+ // to the expected output as possible. This helps out IDEs such that completions and other
+ // related features keep working.
+ let input: syn::ItemFn = match syn::parse(item.clone()) {
+ Ok(it) => it,
+ Err(e) => return token_stream_with_error(item, e),
+ };
+
+ let config = if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
+ let msg = "the main function cannot accept arguments";
+ Err(syn::Error::new_spanned(&input.sig.ident, msg))
+ } else {
+ AttributeArgs::parse_terminated
+ .parse(args)
+ .and_then(|args| build_config(input.clone(), args, false, rt_multi_thread))
+ };
+
+ match config {
+ Ok(config) => parse_knobs(input, false, config),
+ Err(e) => token_stream_with_error(parse_knobs(input, false, DEFAULT_ERROR_CONFIG), e),
+ }
+}
+
+pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
+ // If any of the steps for this macro fail, we still want to expand to an item that is as close
+ // to the expected output as possible. This helps out IDEs such that completions and other
+ // related features keep working.
+ let input: syn::ItemFn = match syn::parse(item.clone()) {
+ Ok(it) => it,
+ Err(e) => return token_stream_with_error(item, e),
+ };
+ let config = if let Some(attr) = input.attrs.iter().find(|attr| attr.path.is_ident("test")) {
+ let msg = "second test attribute is supplied";
+ Err(syn::Error::new_spanned(attr, msg))
+ } else {
+ AttributeArgs::parse_terminated
+ .parse(args)
+ .and_then(|args| build_config(input.clone(), args, true, rt_multi_thread))
+ };
+
+ match config {
+ Ok(config) => parse_knobs(input, true, config),
+ Err(e) => token_stream_with_error(parse_knobs(input, true, DEFAULT_ERROR_CONFIG), e),
+ }
+}
diff --git a/third_party/rust/tokio-macros/src/lib.rs b/third_party/rust/tokio-macros/src/lib.rs
new file mode 100644
index 0000000000..34041afa4c
--- /dev/null
+++ b/third_party/rust/tokio-macros/src/lib.rs
@@ -0,0 +1,483 @@
+#![allow(clippy::needless_doctest_main)]
+#![warn(
+ missing_debug_implementations,
+ missing_docs,
+ rust_2018_idioms,
+ unreachable_pub
+)]
+#![doc(test(
+ no_crate_inject,
+ attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables))
+))]
+
+//! Macros for use with Tokio
+
+// This `extern` is required for older `rustc` versions but newer `rustc`
+// versions warn about the unused `extern crate`.
+#[allow(unused_extern_crates)]
+extern crate proc_macro;
+
+mod entry;
+mod select;
+
+use proc_macro::TokenStream;
+
+/// Marks async function to be executed by the selected runtime. This macro
+/// helps set up a `Runtime` without requiring the user to use
+/// [Runtime](../tokio/runtime/struct.Runtime.html) or
+/// [Builder](../tokio/runtime/struct.Builder.html) directly.
+///
+/// Note: This macro is designed to be simplistic and targets applications that
+/// do not require a complex setup. If the provided functionality is not
+/// sufficient, you may be interested in using
+/// [Builder](../tokio/runtime/struct.Builder.html), which provides a more
+/// powerful interface.
+///
+/// Note: This macro can be used on any function and not just the `main`
+/// function. Using it on a non-main function makes the function behave as if it
+/// was synchronous by starting a new runtime each time it is called. If the
+/// function is called often, it is preferable to create the runtime using the
+/// runtime builder so the runtime can be reused across calls.
+///
+/// # Multi-threaded runtime
+///
+/// To use the multi-threaded runtime, the macro can be configured using
+///
+/// ```
+/// #[tokio::main(flavor = "multi_thread", worker_threads = 10)]
+/// # async fn main() {}
+/// ```
+///
+/// The `worker_threads` option configures the number of worker threads, and
+/// defaults to the number of cpus on the system. This is the default flavor.
+///
+/// Note: The multi-threaded runtime requires the `rt-multi-thread` feature
+/// flag.
+///
+/// # Current thread runtime
+///
+/// To use the single-threaded runtime known as the `current_thread` runtime,
+/// the macro can be configured using
+///
+/// ```
+/// #[tokio::main(flavor = "current_thread")]
+/// # async fn main() {}
+/// ```
+///
+/// ## Function arguments:
+///
+/// Arguments are allowed for any functions aside from `main` which is special
+///
+/// ## Usage
+///
+/// ### Using the multi-thread runtime
+///
+/// ```rust
+/// #[tokio::main]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// fn main() {
+/// tokio::runtime::Builder::new_multi_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+///
+/// ### Using current thread runtime
+///
+/// The basic scheduler is single-threaded.
+///
+/// ```rust
+/// #[tokio::main(flavor = "current_thread")]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// fn main() {
+/// tokio::runtime::Builder::new_current_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+///
+/// ### Set number of worker threads
+///
+/// ```rust
+/// #[tokio::main(worker_threads = 2)]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// fn main() {
+/// tokio::runtime::Builder::new_multi_thread()
+/// .worker_threads(2)
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+///
+/// ### Configure the runtime to start with time paused
+///
+/// ```rust
+/// #[tokio::main(flavor = "current_thread", start_paused = true)]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// fn main() {
+/// tokio::runtime::Builder::new_current_thread()
+/// .enable_all()
+/// .start_paused(true)
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+///
+/// Note that `start_paused` requires the `test-util` feature to be enabled.
+///
+/// ### Rename package
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::main(crate = "tokio1")]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// fn main() {
+/// tokio1::runtime::Builder::new_multi_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+#[proc_macro_attribute]
+#[cfg(not(test))] // Work around for rust-lang/rust#62127
+pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
+ entry::main(args, item, true)
+}
+
+/// Marks async function to be executed by selected runtime. This macro helps set up a `Runtime`
+/// without requiring the user to use [Runtime](../tokio/runtime/struct.Runtime.html) or
+/// [Builder](../tokio/runtime/struct.builder.html) directly.
+///
+/// ## Function arguments:
+///
+/// Arguments are allowed for any functions aside from `main` which is special
+///
+/// ## Usage
+///
+/// ### Using default
+///
+/// ```rust
+/// #[tokio::main(flavor = "current_thread")]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// fn main() {
+/// tokio::runtime::Builder::new_current_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+///
+/// ### Rename package
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::main(crate = "tokio1")]
+/// async fn main() {
+/// println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// fn main() {
+/// tokio1::runtime::Builder::new_multi_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// println!("Hello world");
+/// })
+/// }
+/// ```
+#[proc_macro_attribute]
+#[cfg(not(test))] // Work around for rust-lang/rust#62127
+pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
+ entry::main(args, item, false)
+}
+
+/// Marks async function to be executed by runtime, suitable to test environment.
+/// This macro helps set up a `Runtime` without requiring the user to use
+/// [Runtime](../tokio/runtime/struct.Runtime.html) or
+/// [Builder](../tokio/runtime/struct.Builder.html) directly.
+///
+/// Note: This macro is designed to be simplistic and targets applications that
+/// do not require a complex setup. If the provided functionality is not
+/// sufficient, you may be interested in using
+/// [Builder](../tokio/runtime/struct.Builder.html), which provides a more
+/// powerful interface.
+///
+/// # Multi-threaded runtime
+///
+/// To use the multi-threaded runtime, the macro can be configured using
+///
+/// ```no_run
+/// #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+///
+/// The `worker_threads` option configures the number of worker threads, and
+/// defaults to the number of cpus on the system. This is the default
+/// flavor.
+///
+/// Note: The multi-threaded runtime requires the `rt-multi-thread` feature
+/// flag.
+///
+/// # Current thread runtime
+///
+/// The default test runtime is single-threaded. Each test gets a
+/// separate current-thread runtime.
+///
+/// ```no_run
+/// #[tokio::test]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+///
+/// ## Usage
+///
+/// ### Using the multi-thread runtime
+///
+/// ```no_run
+/// #[tokio::test(flavor = "multi_thread")]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+/// tokio::runtime::Builder::new_multi_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// assert!(true);
+/// })
+/// }
+/// ```
+///
+/// ### Using current thread runtime
+///
+/// ```no_run
+/// #[tokio::test]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+/// tokio::runtime::Builder::new_current_thread()
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// assert!(true);
+/// })
+/// }
+/// ```
+///
+/// ### Set number of worker threads
+///
+/// ```no_run
+/// #[tokio::test(flavor ="multi_thread", worker_threads = 2)]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+/// tokio::runtime::Builder::new_multi_thread()
+/// .worker_threads(2)
+/// .enable_all()
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// assert!(true);
+/// })
+/// }
+/// ```
+///
+/// ### Configure the runtime to start with time paused
+///
+/// ```no_run
+/// #[tokio::test(start_paused = true)]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+/// tokio::runtime::Builder::new_current_thread()
+/// .enable_all()
+/// .start_paused(true)
+/// .build()
+/// .unwrap()
+/// .block_on(async {
+/// assert!(true);
+/// })
+/// }
+/// ```
+///
+/// Note that `start_paused` requires the `test-util` feature to be enabled.
+///
+/// ### Rename package
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::test(crate = "tokio1")]
+/// async fn my_test() {
+/// println!("Hello world");
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
+ entry::test(args, item, true)
+}
+
+/// Marks async function to be executed by runtime, suitable to test environment
+///
+/// ## Usage
+///
+/// ```no_run
+/// #[tokio::test]
+/// async fn my_test() {
+/// assert!(true);
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn test_rt(args: TokenStream, item: TokenStream) -> TokenStream {
+ entry::test(args, item, false)
+}
+
+/// Always fails with the error message below.
+/// ```text
+/// The #[tokio::main] macro requires rt or rt-multi-thread.
+/// ```
+#[proc_macro_attribute]
+pub fn main_fail(_args: TokenStream, _item: TokenStream) -> TokenStream {
+ syn::Error::new(
+ proc_macro2::Span::call_site(),
+ "The #[tokio::main] macro requires rt or rt-multi-thread.",
+ )
+ .to_compile_error()
+ .into()
+}
+
+/// Always fails with the error message below.
+/// ```text
+/// The #[tokio::test] macro requires rt or rt-multi-thread.
+/// ```
+#[proc_macro_attribute]
+pub fn test_fail(_args: TokenStream, _item: TokenStream) -> TokenStream {
+ syn::Error::new(
+ proc_macro2::Span::call_site(),
+ "The #[tokio::test] macro requires rt or rt-multi-thread.",
+ )
+ .to_compile_error()
+ .into()
+}
+
+/// Implementation detail of the `select!` macro. This macro is **not** intended
+/// to be used as part of the public API and is permitted to change.
+#[proc_macro]
+#[doc(hidden)]
+pub fn select_priv_declare_output_enum(input: TokenStream) -> TokenStream {
+ select::declare_output_enum(input)
+}
+
+/// Implementation detail of the `select!` macro. This macro is **not** intended
+/// to be used as part of the public API and is permitted to change.
+#[proc_macro]
+#[doc(hidden)]
+pub fn select_priv_clean_pattern(input: TokenStream) -> TokenStream {
+ select::clean_pattern_macro(input)
+}
diff --git a/third_party/rust/tokio-macros/src/select.rs b/third_party/rust/tokio-macros/src/select.rs
new file mode 100644
index 0000000000..23e280a105
--- /dev/null
+++ b/third_party/rust/tokio-macros/src/select.rs
@@ -0,0 +1,110 @@
+use proc_macro::{TokenStream, TokenTree};
+use proc_macro2::Span;
+use quote::quote;
+use syn::Ident;
+
+pub(crate) fn declare_output_enum(input: TokenStream) -> TokenStream {
+ // passed in is: `(_ _ _)` with one `_` per branch
+ let branches = match input.into_iter().next() {
+ Some(TokenTree::Group(group)) => group.stream().into_iter().count(),
+ _ => panic!("unexpected macro input"),
+ };
+
+ let variants = (0..branches)
+ .map(|num| Ident::new(&format!("_{}", num), Span::call_site()))
+ .collect::<Vec<_>>();
+
+ // Use a bitfield to track which futures completed
+ let mask = Ident::new(
+ if branches <= 8 {
+ "u8"
+ } else if branches <= 16 {
+ "u16"
+ } else if branches <= 32 {
+ "u32"
+ } else if branches <= 64 {
+ "u64"
+ } else {
+ panic!("up to 64 branches supported");
+ },
+ Span::call_site(),
+ );
+
+ TokenStream::from(quote! {
+ pub(super) enum Out<#( #variants ),*> {
+ #( #variants(#variants), )*
+ // Include a `Disabled` variant signifying that all select branches
+ // failed to resolve.
+ Disabled,
+ }
+
+ pub(super) type Mask = #mask;
+ })
+}
+
+pub(crate) fn clean_pattern_macro(input: TokenStream) -> TokenStream {
+ // If this isn't a pattern, we return the token stream as-is. The select!
+ // macro is using it in a location requiring a pattern, so an error will be
+ // emitted there.
+ let mut input: syn::Pat = match syn::parse(input.clone()) {
+ Ok(it) => it,
+ Err(_) => return input,
+ };
+
+ clean_pattern(&mut input);
+ quote::ToTokens::into_token_stream(input).into()
+}
+
+// Removes any occurrences of ref or mut in the provided pattern.
+fn clean_pattern(pat: &mut syn::Pat) {
+ match pat {
+ syn::Pat::Box(_box) => {}
+ syn::Pat::Lit(_literal) => {}
+ syn::Pat::Macro(_macro) => {}
+ syn::Pat::Path(_path) => {}
+ syn::Pat::Range(_range) => {}
+ syn::Pat::Rest(_rest) => {}
+ syn::Pat::Verbatim(_tokens) => {}
+ syn::Pat::Wild(_underscore) => {}
+ syn::Pat::Ident(ident) => {
+ ident.by_ref = None;
+ ident.mutability = None;
+ if let Some((_at, pat)) = &mut ident.subpat {
+ clean_pattern(&mut *pat);
+ }
+ }
+ syn::Pat::Or(or) => {
+ for case in or.cases.iter_mut() {
+ clean_pattern(case);
+ }
+ }
+ syn::Pat::Slice(slice) => {
+ for elem in slice.elems.iter_mut() {
+ clean_pattern(elem);
+ }
+ }
+ syn::Pat::Struct(struct_pat) => {
+ for field in struct_pat.fields.iter_mut() {
+ clean_pattern(&mut field.pat);
+ }
+ }
+ syn::Pat::Tuple(tuple) => {
+ for elem in tuple.elems.iter_mut() {
+ clean_pattern(elem);
+ }
+ }
+ syn::Pat::TupleStruct(tuple) => {
+ for elem in tuple.pat.elems.iter_mut() {
+ clean_pattern(elem);
+ }
+ }
+ syn::Pat::Reference(reference) => {
+ reference.mutability = None;
+ clean_pattern(&mut *reference.pat);
+ }
+ syn::Pat::Type(type_pat) => {
+ clean_pattern(&mut *type_pat.pat);
+ }
+ _ => {}
+ }
+}