tokio-1.43.1/.cargo_vcs_info.json0000644000000001430000000000100122450ustar { "git": { "sha1": "a7b658c35bd40f6811e557aeb97cbb361b612c56" }, "path_in_vcs": "tokio" }tokio-1.43.1/CHANGELOG.md000064400000000000000000004320131046102023000126530ustar 00000000000000# 1.43.1 (April 2nd, 2025) This release fixes a soundness issue in the broadcast channel. The channel accepts values that are `Send` but `!Sync`. Previously, the channel called `clone()` on these values without synchronizing. This release fixes the channel by synchronizing calls to `.clone()` (Thanks Austin Bonander for finding and reporting the issue). ### Fixed - sync: synchronize `clone()` call in broadcast channel ([#7232]) [#7232]: https://github.com/tokio-rs/tokio/pull/7232 # 1.43.0 (Jan 8th, 2025) ### Added - net: add `UdpSocket::peek` methods ([#7068]) - net: add support for Haiku OS ([#7042]) - process: add `Command::into_std()` ([#7014]) - signal: add `SignalKind::info` on illumos ([#6995]) - signal: add support for realtime signals on illumos ([#7029]) ### Fixed - io: don't call `set_len` before initializing vector in `Blocking` ([#7054]) - macros: suppress `clippy::needless_return` in `#[tokio::main]` ([#6874]) - runtime: fix thread parking on WebAssembly ([#7041]) ### Changes - chore: use unsync loads for `unsync_load` ([#7073]) - io: use `Buf::put_bytes` in `Repeat` read impl ([#7055]) - task: drop the join waker of a task eagerly ([#6986]) ### Changes to unstable APIs - metrics: improve flexibility of H2Histogram Configuration ([#6963]) - taskdump: add accessor methods for backtrace ([#6975]) ### Documented - io: clarify `ReadBuf::uninit` allows initialized buffers as well ([#7053]) - net: fix ambiguity in `TcpStream::try_write_vectored` docs ([#7067]) - runtime: fix `LocalRuntime` doc links ([#7074]) - sync: extend documentation for `watch::Receiver::wait_for` ([#7038]) - sync: fix typos in `OnceCell` docs ([#7047]) [#6874]: https://github.com/tokio-rs/tokio/pull/6874 [#6963]: https://github.com/tokio-rs/tokio/pull/6963 [#6975]: https://github.com/tokio-rs/tokio/pull/6975 [#6986]: https://github.com/tokio-rs/tokio/pull/6986 [#6995]: https://github.com/tokio-rs/tokio/pull/6995 [#7014]: https://github.com/tokio-rs/tokio/pull/7014 [#7029]: https://github.com/tokio-rs/tokio/pull/7029 [#7038]: https://github.com/tokio-rs/tokio/pull/7038 [#7041]: https://github.com/tokio-rs/tokio/pull/7041 [#7042]: https://github.com/tokio-rs/tokio/pull/7042 [#7047]: https://github.com/tokio-rs/tokio/pull/7047 [#7053]: https://github.com/tokio-rs/tokio/pull/7053 [#7054]: https://github.com/tokio-rs/tokio/pull/7054 [#7055]: https://github.com/tokio-rs/tokio/pull/7055 [#7067]: https://github.com/tokio-rs/tokio/pull/7067 [#7068]: https://github.com/tokio-rs/tokio/pull/7068 [#7073]: https://github.com/tokio-rs/tokio/pull/7073 [#7074]: https://github.com/tokio-rs/tokio/pull/7074 # 1.42.0 (Dec 3rd, 2024) ### Added - io: add `AsyncFd::{try_io, try_io_mut}` ([#6967]) ### Fixed - io: avoid `ptr->ref->ptr` roundtrip in RegistrationSet ([#6929]) - runtime: do not defer `yield_now` inside `block_in_place` ([#6999]) ### Changes - io: simplify io readiness logic ([#6966]) ### Documented - net: fix docs for `tokio::net::unix::{pid_t, gid_t, uid_t}` ([#6791]) - time: fix a typo in `Instant` docs ([#6982]) [#6791]: https://github.com/tokio-rs/tokio/pull/6791 [#6929]: https://github.com/tokio-rs/tokio/pull/6929 [#6966]: https://github.com/tokio-rs/tokio/pull/6966 [#6967]: https://github.com/tokio-rs/tokio/pull/6967 [#6982]: https://github.com/tokio-rs/tokio/pull/6982 [#6999]: https://github.com/tokio-rs/tokio/pull/6999 # 1.41.1 (Nov 7th, 2024) ### Fixed - metrics: fix bug with wrong number of buckets for the histogram ([#6957]) - net: display `net` requirement for `net::UdpSocket` in docs ([#6938]) - net: fix typo in `TcpStream` internal comment ([#6944]) [#6957]: https://github.com/tokio-rs/tokio/pull/6957 [#6938]: https://github.com/tokio-rs/tokio/pull/6938 [#6944]: https://github.com/tokio-rs/tokio/pull/6944 # 1.41.0 (Oct 22nd, 2024) ### Added - metrics: stabilize `global_queue_depth` ([#6854], [#6918]) - net: add conversions for unix `SocketAddr` ([#6868]) - sync: add `watch::Sender::sender_count` ([#6836]) - sync: add `mpsc::Receiver::blocking_recv_many` ([#6867]) - task: stabilize `Id` apis ([#6793], [#6891]) ### Added (unstable) - metrics: add H2 Histogram option to improve histogram granularity ([#6897]) - metrics: rename some histogram apis ([#6924]) - runtime: add `LocalRuntime` ([#6808]) ### Changed - runtime: box futures larger than 16k on release mode ([#6826]) - sync: add `#[must_use]` to `Notified` ([#6828]) - sync: make `watch` cooperative ([#6846]) - sync: make `broadcast::Receiver` cooperative ([#6870]) - task: add task size to tracing instrumentation ([#6881]) - wasm: enable `cfg_fs` for `wasi` target ([#6822]) ### Fixed - net: fix regression of abstract socket path in unix socket ([#6838]) ### Documented - io: recommend `OwnedFd` with `AsyncFd` ([#6821]) - io: document cancel safety of `AsyncFd` methods ([#6890]) - macros: render more comprehensible documentation for `join` and `try_join` ([#6814], [#6841]) - net: fix swapped examples for `TcpSocket::set_nodelay` and `TcpSocket::nodelay` ([#6840]) - sync: document runtime compatibility ([#6833]) [#6793]: https://github.com/tokio-rs/tokio/pull/6793 [#6808]: https://github.com/tokio-rs/tokio/pull/6808 [#6810]: https://github.com/tokio-rs/tokio/pull/6810 [#6814]: https://github.com/tokio-rs/tokio/pull/6814 [#6821]: https://github.com/tokio-rs/tokio/pull/6821 [#6822]: https://github.com/tokio-rs/tokio/pull/6822 [#6826]: https://github.com/tokio-rs/tokio/pull/6826 [#6828]: https://github.com/tokio-rs/tokio/pull/6828 [#6833]: https://github.com/tokio-rs/tokio/pull/6833 [#6836]: https://github.com/tokio-rs/tokio/pull/6836 [#6838]: https://github.com/tokio-rs/tokio/pull/6838 [#6840]: https://github.com/tokio-rs/tokio/pull/6840 [#6841]: https://github.com/tokio-rs/tokio/pull/6841 [#6846]: https://github.com/tokio-rs/tokio/pull/6846 [#6854]: https://github.com/tokio-rs/tokio/pull/6854 [#6867]: https://github.com/tokio-rs/tokio/pull/6867 [#6868]: https://github.com/tokio-rs/tokio/pull/6868 [#6870]: https://github.com/tokio-rs/tokio/pull/6870 [#6881]: https://github.com/tokio-rs/tokio/pull/6881 [#6890]: https://github.com/tokio-rs/tokio/pull/6890 [#6891]: https://github.com/tokio-rs/tokio/pull/6891 [#6897]: https://github.com/tokio-rs/tokio/pull/6897 [#6918]: https://github.com/tokio-rs/tokio/pull/6918 [#6924]: https://github.com/tokio-rs/tokio/pull/6924 # 1.40.0 (August 30th, 2024) ### Added - io: add `util::SimplexStream` ([#6589]) - process: stabilize `Command::process_group` ([#6731]) - sync: add `{TrySendError,SendTimeoutError}::into_inner` ([#6755]) - task: add `JoinSet::join_all` ([#6784]) ### Added (unstable) - runtime: add `Builder::{on_task_spawn, on_task_terminate}` ([#6742]) ### Changed - io: use vectored io for `write_all_buf` when possible ([#6724]) - runtime: prevent niche-optimization to avoid triggering miri ([#6744]) - sync: mark mpsc types as `UnwindSafe` ([#6783]) - sync,time: make `Sleep` and `BatchSemaphore` instrumentation explicit roots ([#6727]) - task: use `NonZeroU64` for `task::Id` ([#6733]) - task: include panic message when printing `JoinError` ([#6753]) - task: add `#[must_use]` to `JoinHandle::abort_handle` ([#6762]) - time: eliminate timer wheel allocations ([#6779]) ### Documented - docs: clarify that `[build]` section doesn't go in Cargo.toml ([#6728]) - io: clarify zero remaining capacity case ([#6790]) - macros: improve documentation for `select!` ([#6774]) - sync: document mpsc channel allocation behavior ([#6773]) [#6589]: https://github.com/tokio-rs/tokio/pull/6589 [#6724]: https://github.com/tokio-rs/tokio/pull/6724 [#6727]: https://github.com/tokio-rs/tokio/pull/6727 [#6728]: https://github.com/tokio-rs/tokio/pull/6728 [#6731]: https://github.com/tokio-rs/tokio/pull/6731 [#6733]: https://github.com/tokio-rs/tokio/pull/6733 [#6742]: https://github.com/tokio-rs/tokio/pull/6742 [#6744]: https://github.com/tokio-rs/tokio/pull/6744 [#6753]: https://github.com/tokio-rs/tokio/pull/6753 [#6755]: https://github.com/tokio-rs/tokio/pull/6755 [#6762]: https://github.com/tokio-rs/tokio/pull/6762 [#6773]: https://github.com/tokio-rs/tokio/pull/6773 [#6774]: https://github.com/tokio-rs/tokio/pull/6774 [#6779]: https://github.com/tokio-rs/tokio/pull/6779 [#6783]: https://github.com/tokio-rs/tokio/pull/6783 [#6784]: https://github.com/tokio-rs/tokio/pull/6784 [#6790]: https://github.com/tokio-rs/tokio/pull/6790 # 1.39.3 (August 17th, 2024) This release fixes a regression where the unix socket api stopped accepting the abstract socket namespace. ([#6772]) [#6772]: https://github.com/tokio-rs/tokio/pull/6772 # 1.39.2 (July 27th, 2024) This release fixes a regression where the `select!` macro stopped accepting expressions that make use of temporary lifetime extension. ([#6722]) [#6722]: https://github.com/tokio-rs/tokio/pull/6722 # 1.39.1 (July 23rd, 2024) This release reverts "time: avoid traversing entries in the time wheel twice" because it contains a bug. ([#6715]) [#6715]: https://github.com/tokio-rs/tokio/pull/6715 # 1.39.0 (July 23rd, 2024) Yanked. Please use 1.39.1 instead. - This release bumps the MSRV to 1.70. ([#6645]) - This release upgrades to mio v1. ([#6635]) - This release upgrades to windows-sys v0.52 ([#6154]) ### Added - io: implement `AsyncSeek` for `Empty` ([#6663]) - metrics: stabilize `num_alive_tasks` ([#6619], [#6667]) - process: add `Command::as_std_mut` ([#6608]) - sync: add `watch::Sender::same_channel` ([#6637]) - sync: add `{Receiver,UnboundedReceiver}::{sender_strong_count,sender_weak_count}` ([#6661]) - sync: implement `Default` for `watch::Sender` ([#6626]) - task: implement `Clone` for `AbortHandle` ([#6621]) - task: stabilize `consume_budget` ([#6622]) ### Changed - io: improve panic message of `ReadBuf::put_slice()` ([#6629]) - io: read during write in `copy_bidirectional` and `copy` ([#6532]) - runtime: replace `num_cpus` with `available_parallelism` ([#6709]) - task: avoid stack overflow when passing large future to `block_on` ([#6692]) - time: avoid traversing entries in the time wheel twice ([#6584]) - time: support `IntoFuture` with `timeout` ([#6666]) - macros: support `IntoFuture` with `join!` and `select!` ([#6710]) ### Fixed - docs: fix docsrs builds with the fs feature enabled ([#6585]) - io: only use short-read optimization on known-to-be-compatible platforms ([#6668]) - time: fix overflow panic when using large durations with `Interval` ([#6612]) ### Added (unstable) - macros: allow `unhandled_panic` behavior for `#[tokio::main]` and `#[tokio::test]` ([#6593]) - metrics: add `spawned_tasks_count` ([#6114]) - metrics: add `worker_park_unpark_count` ([#6696]) - metrics: add worker thread id ([#6695]) ### Documented - io: update `tokio::io::stdout` documentation ([#6674]) - macros: typo fix in `join.rs` and `try_join.rs` ([#6641]) - runtime: fix typo in `unhandled_panic` ([#6660]) - task: document behavior of `JoinSet::try_join_next` when all tasks are running ([#6671]) [#6114]: https://github.com/tokio-rs/tokio/pull/6114 [#6154]: https://github.com/tokio-rs/tokio/pull/6154 [#6532]: https://github.com/tokio-rs/tokio/pull/6532 [#6584]: https://github.com/tokio-rs/tokio/pull/6584 [#6585]: https://github.com/tokio-rs/tokio/pull/6585 [#6593]: https://github.com/tokio-rs/tokio/pull/6593 [#6608]: https://github.com/tokio-rs/tokio/pull/6608 [#6612]: https://github.com/tokio-rs/tokio/pull/6612 [#6619]: https://github.com/tokio-rs/tokio/pull/6619 [#6621]: https://github.com/tokio-rs/tokio/pull/6621 [#6622]: https://github.com/tokio-rs/tokio/pull/6622 [#6626]: https://github.com/tokio-rs/tokio/pull/6626 [#6629]: https://github.com/tokio-rs/tokio/pull/6629 [#6635]: https://github.com/tokio-rs/tokio/pull/6635 [#6637]: https://github.com/tokio-rs/tokio/pull/6637 [#6641]: https://github.com/tokio-rs/tokio/pull/6641 [#6645]: https://github.com/tokio-rs/tokio/pull/6645 [#6660]: https://github.com/tokio-rs/tokio/pull/6660 [#6661]: https://github.com/tokio-rs/tokio/pull/6661 [#6663]: https://github.com/tokio-rs/tokio/pull/6663 [#6666]: https://github.com/tokio-rs/tokio/pull/6666 [#6667]: https://github.com/tokio-rs/tokio/pull/6667 [#6668]: https://github.com/tokio-rs/tokio/pull/6668 [#6671]: https://github.com/tokio-rs/tokio/pull/6671 [#6674]: https://github.com/tokio-rs/tokio/pull/6674 [#6692]: https://github.com/tokio-rs/tokio/pull/6692 [#6695]: https://github.com/tokio-rs/tokio/pull/6695 [#6696]: https://github.com/tokio-rs/tokio/pull/6696 [#6709]: https://github.com/tokio-rs/tokio/pull/6709 [#6710]: https://github.com/tokio-rs/tokio/pull/6710 # 1.38.2 (April 2nd, 2025) This release fixes a soundness issue in the broadcast channel. The channel accepts values that are `Send` but `!Sync`. Previously, the channel called `clone()` on these values without synchronizing. This release fixes the channel by synchronizing calls to `.clone()` (Thanks Austin Bonander for finding and reporting the issue). ### Fixed - sync: synchronize `clone()` call in broadcast channel ([#7232]) [#7232]: https://github.com/tokio-rs/tokio/pull/7232 # 1.38.1 (July 16th, 2024) This release fixes the bug identified as ([#6682]), which caused timers not to fire when they should. ### Fixed - time: update `wake_up` while holding all the locks of sharded time wheels ([#6683]) [#6682]: https://github.com/tokio-rs/tokio/pull/6682 [#6683]: https://github.com/tokio-rs/tokio/pull/6683 # 1.38.0 (May 30th, 2024) This release marks the beginning of stabilization for runtime metrics. It stabilizes `RuntimeMetrics::worker_count`. Future releases will continue to stabilize more metrics. ### Added - fs: add `File::create_new` ([#6573]) - io: add `copy_bidirectional_with_sizes` ([#6500]) - io: implement `AsyncBufRead` for `Join` ([#6449]) - net: add Apple visionOS support ([#6465]) - net: implement `Clone` for `NamedPipeInfo` ([#6586]) - net: support QNX OS ([#6421]) - sync: add `Notify::notify_last` ([#6520]) - sync: add `mpsc::Receiver::{capacity,max_capacity}` ([#6511]) - sync: add `split` method to the semaphore permit ([#6472], [#6478]) - task: add `tokio::task::join_set::Builder::spawn_blocking` ([#6578]) - wasm: support rt-multi-thread with wasm32-wasi-preview1-threads ([#6510]) ### Changed - macros: make `#[tokio::test]` append `#[test]` at the end of the attribute list ([#6497]) - metrics: fix `blocking_threads` count ([#6551]) - metrics: stabilize `RuntimeMetrics::worker_count` ([#6556]) - runtime: move task out of the `lifo_slot` in `block_in_place` ([#6596]) - runtime: panic if `global_queue_interval` is zero ([#6445]) - sync: always drop message in destructor for oneshot receiver ([#6558]) - sync: instrument `Semaphore` for task dumps ([#6499]) - sync: use FIFO ordering when waking batches of wakers ([#6521]) - task: make `LocalKey::get` work with Clone types ([#6433]) - tests: update nix and mio-aio dev-dependencies ([#6552]) - time: clean up implementation ([#6517]) - time: lazily init timers on first poll ([#6512]) - time: remove the `true_when` field in `TimerShared` ([#6563]) - time: use sharding for timer implementation ([#6534]) ### Fixed - taskdump: allow building taskdump docs on non-unix machines ([#6564]) - time: check for overflow in `Interval::poll_tick` ([#6487]) - sync: fix incorrect `is_empty` on mpsc block boundaries ([#6603]) ### Documented - fs: rewrite file system docs ([#6467]) - io: fix `stdin` documentation ([#6581]) - io: fix obsolete reference in `ReadHalf::unsplit()` documentation ([#6498]) - macros: render more comprehensible documentation for `select!` ([#6468]) - net: add missing types to module docs ([#6482]) - net: fix misleading `NamedPipeServer` example ([#6590]) - sync: add examples for `SemaphorePermit`, `OwnedSemaphorePermit` ([#6477]) - sync: document that `Barrier::wait` is not cancel safe ([#6494]) - sync: explain relation between `watch::Sender::{subscribe,closed}` ([#6490]) - task: clarify that you can't abort `spawn_blocking` tasks ([#6571]) - task: fix a typo in doc of `LocalSet::run_until` ([#6599]) - time: fix test-util requirement for pause and resume in docs ([#6503]) [#6421]: https://github.com/tokio-rs/tokio/pull/6421 [#6433]: https://github.com/tokio-rs/tokio/pull/6433 [#6445]: https://github.com/tokio-rs/tokio/pull/6445 [#6449]: https://github.com/tokio-rs/tokio/pull/6449 [#6465]: https://github.com/tokio-rs/tokio/pull/6465 [#6467]: https://github.com/tokio-rs/tokio/pull/6467 [#6468]: https://github.com/tokio-rs/tokio/pull/6468 [#6472]: https://github.com/tokio-rs/tokio/pull/6472 [#6477]: https://github.com/tokio-rs/tokio/pull/6477 [#6478]: https://github.com/tokio-rs/tokio/pull/6478 [#6482]: https://github.com/tokio-rs/tokio/pull/6482 [#6487]: https://github.com/tokio-rs/tokio/pull/6487 [#6490]: https://github.com/tokio-rs/tokio/pull/6490 [#6494]: https://github.com/tokio-rs/tokio/pull/6494 [#6497]: https://github.com/tokio-rs/tokio/pull/6497 [#6498]: https://github.com/tokio-rs/tokio/pull/6498 [#6499]: https://github.com/tokio-rs/tokio/pull/6499 [#6500]: https://github.com/tokio-rs/tokio/pull/6500 [#6503]: https://github.com/tokio-rs/tokio/pull/6503 [#6510]: https://github.com/tokio-rs/tokio/pull/6510 [#6511]: https://github.com/tokio-rs/tokio/pull/6511 [#6512]: https://github.com/tokio-rs/tokio/pull/6512 [#6517]: https://github.com/tokio-rs/tokio/pull/6517 [#6520]: https://github.com/tokio-rs/tokio/pull/6520 [#6521]: https://github.com/tokio-rs/tokio/pull/6521 [#6534]: https://github.com/tokio-rs/tokio/pull/6534 [#6551]: https://github.com/tokio-rs/tokio/pull/6551 [#6552]: https://github.com/tokio-rs/tokio/pull/6552 [#6556]: https://github.com/tokio-rs/tokio/pull/6556 [#6558]: https://github.com/tokio-rs/tokio/pull/6558 [#6563]: https://github.com/tokio-rs/tokio/pull/6563 [#6564]: https://github.com/tokio-rs/tokio/pull/6564 [#6571]: https://github.com/tokio-rs/tokio/pull/6571 [#6573]: https://github.com/tokio-rs/tokio/pull/6573 [#6578]: https://github.com/tokio-rs/tokio/pull/6578 [#6581]: https://github.com/tokio-rs/tokio/pull/6581 [#6586]: https://github.com/tokio-rs/tokio/pull/6586 [#6590]: https://github.com/tokio-rs/tokio/pull/6590 [#6596]: https://github.com/tokio-rs/tokio/pull/6596 [#6599]: https://github.com/tokio-rs/tokio/pull/6599 [#6603]: https://github.com/tokio-rs/tokio/pull/6603 # 1.37.0 (March 28th, 2024) ### Added - fs: add `set_max_buf_size` to `tokio::fs::File` ([#6411]) - io: add `try_new` and `try_with_interest` to `AsyncFd` ([#6345]) - sync: add `forget_permits` method to semaphore ([#6331]) - sync: add `is_closed`, `is_empty`, and `len` to mpsc receivers ([#6348]) - sync: add a `rwlock()` method to owned `RwLock` guards ([#6418]) - sync: expose strong and weak counts of mpsc sender handles ([#6405]) - sync: implement `Clone` for `watch::Sender` ([#6388]) - task: add `TaskLocalFuture::take_value` ([#6340]) - task: implement `FromIterator` for `JoinSet` ([#6300]) ### Changed - io: make `io::split` use a mutex instead of a spinlock ([#6403]) ### Fixed - docs: fix docsrs build without net feature ([#6360]) - macros: allow select with only else branch ([#6339]) - runtime: fix leaking registration entries when os registration fails ([#6329]) ### Documented - io: document cancel safety of `AsyncBufReadExt::fill_buf` ([#6431]) - io: document cancel safety of `AsyncReadExt`'s primitive read functions ([#6337]) - runtime: add doc link from `Runtime` to `#[tokio::main]` ([#6366]) - runtime: make the `enter` example deterministic ([#6351]) - sync: add Semaphore example for limiting the number of outgoing requests ([#6419]) - sync: fix missing period in broadcast docs ([#6377]) - sync: mark `mpsc::Sender::downgrade` with `#[must_use]` ([#6326]) - sync: reorder `const_new` before `new_with` ([#6392]) - sync: update watch channel docs ([#6395]) - task: fix documentation links ([#6336]) ### Changed (unstable) - runtime: include task `Id` in taskdumps ([#6328]) - runtime: panic if `unhandled_panic` is enabled when not supported ([#6410]) [#6300]: https://github.com/tokio-rs/tokio/pull/6300 [#6326]: https://github.com/tokio-rs/tokio/pull/6326 [#6328]: https://github.com/tokio-rs/tokio/pull/6328 [#6329]: https://github.com/tokio-rs/tokio/pull/6329 [#6331]: https://github.com/tokio-rs/tokio/pull/6331 [#6336]: https://github.com/tokio-rs/tokio/pull/6336 [#6337]: https://github.com/tokio-rs/tokio/pull/6337 [#6339]: https://github.com/tokio-rs/tokio/pull/6339 [#6340]: https://github.com/tokio-rs/tokio/pull/6340 [#6345]: https://github.com/tokio-rs/tokio/pull/6345 [#6348]: https://github.com/tokio-rs/tokio/pull/6348 [#6351]: https://github.com/tokio-rs/tokio/pull/6351 [#6360]: https://github.com/tokio-rs/tokio/pull/6360 [#6366]: https://github.com/tokio-rs/tokio/pull/6366 [#6377]: https://github.com/tokio-rs/tokio/pull/6377 [#6388]: https://github.com/tokio-rs/tokio/pull/6388 [#6392]: https://github.com/tokio-rs/tokio/pull/6392 [#6395]: https://github.com/tokio-rs/tokio/pull/6395 [#6403]: https://github.com/tokio-rs/tokio/pull/6403 [#6405]: https://github.com/tokio-rs/tokio/pull/6405 [#6410]: https://github.com/tokio-rs/tokio/pull/6410 [#6411]: https://github.com/tokio-rs/tokio/pull/6411 [#6418]: https://github.com/tokio-rs/tokio/pull/6418 [#6419]: https://github.com/tokio-rs/tokio/pull/6419 [#6431]: https://github.com/tokio-rs/tokio/pull/6431 # 1.36.0 (February 2nd, 2024) ### Added - io: add `tokio::io::Join` ([#6220]) - io: implement `AsyncWrite` for `Empty` ([#6235]) - net: add support for anonymous unix pipes ([#6127]) - net: add `UnixSocket` ([#6290]) - net: expose keepalive option on `TcpSocket` ([#6311]) - sync: add `{Receiver,UnboundedReceiver}::poll_recv_many` ([#6236]) - sync: add `Sender::{try_,}reserve_many` ([#6205]) - sync: add `watch::Receiver::mark_unchanged` ([#6252]) - task: add `JoinSet::try_join_next` ([#6280]) ### Changed - io: make `copy` cooperative ([#6265]) - io: make `repeat` and `sink` cooperative ([#6254]) - io: simplify check for empty slice ([#6293]) - process: use pidfd on Linux when available ([#6152]) - sync: use AtomicBool in broadcast channel future ([#6298]) ### Documented - io: clarify `clear_ready` docs ([#6304]) - net: document that `*Fd` traits on `TcpSocket` are unix-only ([#6294]) - sync: document FIFO behavior of `tokio::sync::Mutex` ([#6279]) - chore: typographic improvements ([#6262]) - runtime: remove obsolete comment ([#6303]) - task: fix typo ([#6261]) [#6220]: https://github.com/tokio-rs/tokio/pull/6220 [#6235]: https://github.com/tokio-rs/tokio/pull/6235 [#6127]: https://github.com/tokio-rs/tokio/pull/6127 [#6290]: https://github.com/tokio-rs/tokio/pull/6290 [#6311]: https://github.com/tokio-rs/tokio/pull/6311 [#6236]: https://github.com/tokio-rs/tokio/pull/6236 [#6205]: https://github.com/tokio-rs/tokio/pull/6205 [#6252]: https://github.com/tokio-rs/tokio/pull/6252 [#6280]: https://github.com/tokio-rs/tokio/pull/6280 [#6265]: https://github.com/tokio-rs/tokio/pull/6265 [#6254]: https://github.com/tokio-rs/tokio/pull/6254 [#6293]: https://github.com/tokio-rs/tokio/pull/6293 [#6238]: https://github.com/tokio-rs/tokio/pull/6238 [#6152]: https://github.com/tokio-rs/tokio/pull/6152 [#6298]: https://github.com/tokio-rs/tokio/pull/6298 [#6262]: https://github.com/tokio-rs/tokio/pull/6262 [#6303]: https://github.com/tokio-rs/tokio/pull/6303 [#6261]: https://github.com/tokio-rs/tokio/pull/6261 [#6304]: https://github.com/tokio-rs/tokio/pull/6304 [#6294]: https://github.com/tokio-rs/tokio/pull/6294 [#6279]: https://github.com/tokio-rs/tokio/pull/6279 # 1.35.1 (December 19, 2023) This is a forward part of a change that was backported to 1.25.3. ### Fixed - io: add budgeting to `tokio::runtime::io::registration::async_io` ([#6221]) [#6221]: https://github.com/tokio-rs/tokio/pull/6221 # 1.35.0 (December 8th, 2023) ### Added - net: add Apple watchOS support ([#6176]) ### Changed - io: drop the `Sized` requirements from `AsyncReadExt.read_buf` ([#6169]) - runtime: make `Runtime` unwind safe ([#6189]) - runtime: reduce the lock contention in task spawn ([#6001]) - tokio: update nix dependency to 0.27.1 ([#6190]) ### Fixed - chore: make `--cfg docsrs` work without net feature ([#6166]) - chore: use relaxed load for `unsync_load` on miri ([#6179]) - runtime: handle missing context on wake ([#6148]) - taskdump: fix taskdump cargo config example ([#6150]) - taskdump: skip notified tasks during taskdumps ([#6194]) - tracing: avoid creating resource spans with current parent, use a None parent instead ([#6107]) - tracing: make task span explicit root ([#6158]) ### Documented - io: flush in `AsyncWriteExt` examples ([#6149]) - runtime: document fairness guarantees and current behavior ([#6145]) - task: document cancel safety of `LocalSet::run_until` ([#6147]) [#6001]: https://github.com/tokio-rs/tokio/pull/6001 [#6107]: https://github.com/tokio-rs/tokio/pull/6107 [#6144]: https://github.com/tokio-rs/tokio/pull/6144 [#6145]: https://github.com/tokio-rs/tokio/pull/6145 [#6147]: https://github.com/tokio-rs/tokio/pull/6147 [#6148]: https://github.com/tokio-rs/tokio/pull/6148 [#6149]: https://github.com/tokio-rs/tokio/pull/6149 [#6150]: https://github.com/tokio-rs/tokio/pull/6150 [#6158]: https://github.com/tokio-rs/tokio/pull/6158 [#6166]: https://github.com/tokio-rs/tokio/pull/6166 [#6169]: https://github.com/tokio-rs/tokio/pull/6169 [#6176]: https://github.com/tokio-rs/tokio/pull/6176 [#6179]: https://github.com/tokio-rs/tokio/pull/6179 [#6189]: https://github.com/tokio-rs/tokio/pull/6189 [#6190]: https://github.com/tokio-rs/tokio/pull/6190 [#6194]: https://github.com/tokio-rs/tokio/pull/6194 # 1.34.0 (November 19th, 2023) ### Fixed - io: allow `clear_readiness` after io driver shutdown ([#6067]) - io: fix integer overflow in `take` ([#6080]) - io: fix I/O resource hang ([#6134]) - sync: fix `broadcast::channel` link ([#6100]) ### Changed - macros: use `::core` qualified imports instead of `::std` inside `tokio::test` macro ([#5973]) ### Added - fs: update cfg attr in `fs::read_dir` to include `aix` ([#6075]) - sync: add `mpsc::Receiver::recv_many` ([#6010]) - tokio: added vita target support ([#6094]) [#5973]: https://github.com/tokio-rs/tokio/pull/5973 [#6067]: https://github.com/tokio-rs/tokio/pull/6067 [#6080]: https://github.com/tokio-rs/tokio/pull/6080 [#6134]: https://github.com/tokio-rs/tokio/pull/6134 [#6100]: https://github.com/tokio-rs/tokio/pull/6100 [#6075]: https://github.com/tokio-rs/tokio/pull/6075 [#6010]: https://github.com/tokio-rs/tokio/pull/6010 [#6094]: https://github.com/tokio-rs/tokio/pull/6094 # 1.33.0 (October 9, 2023) ### Fixed - io: mark `Interest::add` with `#[must_use]` ([#6037]) - runtime: fix cache line size for RISC-V ([#5994]) - sync: prevent lock poisoning in `watch::Receiver::wait_for` ([#6021]) - task: fix `spawn_local` source location ([#5984]) ### Changed - sync: use Acquire/Release orderings instead of SeqCst in `watch` ([#6018]) ### Added - fs: add vectored writes to `tokio::fs::File` ([#5958]) - io: add `Interest::remove` method ([#5906]) - io: add vectored writes to `DuplexStream` ([#5985]) - net: add Apple tvOS support ([#6045]) - sync: add `?Sized` bound to `{MutexGuard,OwnedMutexGuard}::map` ([#5997]) - sync: add `watch::Receiver::mark_unseen` ([#5962], [#6014], [#6017]) - sync: add `watch::Sender::new` ([#5998]) - sync: add const fn `OnceCell::from_value` ([#5903]) ### Removed - remove unused `stats` feature ([#5952]) ### Documented - add missing backticks in code examples ([#5938], [#6056]) - fix typos ([#5988], [#6030]) - process: document that `Child::wait` is cancel safe ([#5977]) - sync: add examples for `Semaphore` ([#5939], [#5956], [#5978], [#6031], [#6032], [#6050]) - sync: document that `broadcast` capacity is a lower bound ([#6042]) - sync: document that `const_new` is not instrumented ([#6002]) - sync: improve cancel-safety documentation for `mpsc::Sender::send` ([#5947]) - sync: improve docs for `watch` channel ([#5954]) - taskdump: render taskdump documentation on docs.rs ([#5972]) ### Unstable - taskdump: fix potential deadlock ([#6036]) [#5903]: https://github.com/tokio-rs/tokio/pull/5903 [#5906]: https://github.com/tokio-rs/tokio/pull/5906 [#5938]: https://github.com/tokio-rs/tokio/pull/5938 [#5939]: https://github.com/tokio-rs/tokio/pull/5939 [#5947]: https://github.com/tokio-rs/tokio/pull/5947 [#5952]: https://github.com/tokio-rs/tokio/pull/5952 [#5954]: https://github.com/tokio-rs/tokio/pull/5954 [#5956]: https://github.com/tokio-rs/tokio/pull/5956 [#5958]: https://github.com/tokio-rs/tokio/pull/5958 [#5960]: https://github.com/tokio-rs/tokio/pull/5960 [#5962]: https://github.com/tokio-rs/tokio/pull/5962 [#5971]: https://github.com/tokio-rs/tokio/pull/5971 [#5972]: https://github.com/tokio-rs/tokio/pull/5972 [#5977]: https://github.com/tokio-rs/tokio/pull/5977 [#5978]: https://github.com/tokio-rs/tokio/pull/5978 [#5984]: https://github.com/tokio-rs/tokio/pull/5984 [#5985]: https://github.com/tokio-rs/tokio/pull/5985 [#5988]: https://github.com/tokio-rs/tokio/pull/5988 [#5994]: https://github.com/tokio-rs/tokio/pull/5994 [#5997]: https://github.com/tokio-rs/tokio/pull/5997 [#5998]: https://github.com/tokio-rs/tokio/pull/5998 [#6002]: https://github.com/tokio-rs/tokio/pull/6002 [#6014]: https://github.com/tokio-rs/tokio/pull/6014 [#6017]: https://github.com/tokio-rs/tokio/pull/6017 [#6018]: https://github.com/tokio-rs/tokio/pull/6018 [#6021]: https://github.com/tokio-rs/tokio/pull/6021 [#6030]: https://github.com/tokio-rs/tokio/pull/6030 [#6031]: https://github.com/tokio-rs/tokio/pull/6031 [#6032]: https://github.com/tokio-rs/tokio/pull/6032 [#6036]: https://github.com/tokio-rs/tokio/pull/6036 [#6037]: https://github.com/tokio-rs/tokio/pull/6037 [#6042]: https://github.com/tokio-rs/tokio/pull/6042 [#6045]: https://github.com/tokio-rs/tokio/pull/6045 [#6050]: https://github.com/tokio-rs/tokio/pull/6050 [#6056]: https://github.com/tokio-rs/tokio/pull/6056 [#6058]: https://github.com/tokio-rs/tokio/pull/6058 # 1.32.1 (December 19, 2023) This is a forward part of a change that was backported to 1.25.3. ### Fixed - io: add budgeting to `tokio::runtime::io::registration::async_io` ([#6221]) [#6221]: https://github.com/tokio-rs/tokio/pull/6221 # 1.32.0 (August 16, 2023) ### Fixed - sync: fix potential quadratic behavior in `broadcast::Receiver` ([#5925]) ### Added - process: stabilize `Command::raw_arg` ([#5930]) - io: enable awaiting error readiness ([#5781]) ### Unstable - rt(alt): improve scalability of alt runtime as the number of cores grows ([#5935]) [#5925]: https://github.com/tokio-rs/tokio/pull/5925 [#5930]: https://github.com/tokio-rs/tokio/pull/5930 [#5781]: https://github.com/tokio-rs/tokio/pull/5781 [#5935]: https://github.com/tokio-rs/tokio/pull/5935 # 1.31.0 (August 10, 2023) ### Fixed * io: delegate `WriteHalf::poll_write_vectored` ([#5914]) ### Unstable * rt(alt): fix memory leak in unstable next-gen scheduler prototype ([#5911]) * rt: expose mean task poll time metric ([#5927]) [#5914]: https://github.com/tokio-rs/tokio/pull/5914 [#5911]: https://github.com/tokio-rs/tokio/pull/5911 [#5927]: https://github.com/tokio-rs/tokio/pull/5927 # 1.30.0 (August 9, 2023) This release bumps the MSRV of Tokio to 1.63. ([#5887]) ### Changed - tokio: reduce LLVM code generation ([#5859]) - io: support `--cfg mio_unsupported_force_poll_poll` flag ([#5881]) - sync: make `const_new` methods always available ([#5885]) - sync: avoid false sharing in mpsc channel ([#5829]) - rt: pop at least one task from inject queue ([#5908]) ### Added - sync: add `broadcast::Sender::new` ([#5824]) - net: implement `UCred` for espidf ([#5868]) - fs: add `File::options()` ([#5869]) - time: implement extra reset variants for `Interval` ([#5878]) - process: add `{ChildStd*}::into_owned_{fd, handle}` ([#5899]) ### Removed - tokio: removed unused `tokio_*` cfgs ([#5890]) - remove build script to speed up compilation ([#5887]) ### Documented - sync: mention lagging in docs for `broadcast::send` ([#5820]) - runtime: expand on sharing runtime docs ([#5858]) - io: use vec in example for `AsyncReadExt::read_exact` ([#5863]) - time: mark `Sleep` as `!Unpin` in docs ([#5916]) - process: fix `raw_arg` not showing up in docs ([#5865]) ### Unstable - rt: add runtime ID ([#5864]) - rt: initial implementation of new threaded runtime ([#5823]) [#5820]: https://github.com/tokio-rs/tokio/pull/5820 [#5823]: https://github.com/tokio-rs/tokio/pull/5823 [#5824]: https://github.com/tokio-rs/tokio/pull/5824 [#5829]: https://github.com/tokio-rs/tokio/pull/5829 [#5858]: https://github.com/tokio-rs/tokio/pull/5858 [#5859]: https://github.com/tokio-rs/tokio/pull/5859 [#5863]: https://github.com/tokio-rs/tokio/pull/5863 [#5864]: https://github.com/tokio-rs/tokio/pull/5864 [#5865]: https://github.com/tokio-rs/tokio/pull/5865 [#5868]: https://github.com/tokio-rs/tokio/pull/5868 [#5869]: https://github.com/tokio-rs/tokio/pull/5869 [#5878]: https://github.com/tokio-rs/tokio/pull/5878 [#5881]: https://github.com/tokio-rs/tokio/pull/5881 [#5885]: https://github.com/tokio-rs/tokio/pull/5885 [#5887]: https://github.com/tokio-rs/tokio/pull/5887 [#5890]: https://github.com/tokio-rs/tokio/pull/5890 [#5899]: https://github.com/tokio-rs/tokio/pull/5899 [#5908]: https://github.com/tokio-rs/tokio/pull/5908 [#5916]: https://github.com/tokio-rs/tokio/pull/5916 # 1.29.1 (June 29, 2023) ### Fixed - rt: fix nesting two `block_in_place` with a `block_on` between ([#5837]) [#5837]: https://github.com/tokio-rs/tokio/pull/5837 # 1.29.0 (June 27, 2023) Technically a breaking change, the `Send` implementation is removed from `runtime::EnterGuard`. This change fixes a bug and should not impact most users. ### Breaking - rt: `EnterGuard` should not be `Send` ([#5766]) ### Fixed - fs: reduce blocking ops in `fs::read_dir` ([#5653]) - rt: fix possible starvation ([#5686], [#5712]) - rt: fix stacked borrows issue in `JoinSet` ([#5693]) - rt: panic if `EnterGuard` dropped incorrect order ([#5772]) - time: do not overflow to signal value ([#5710]) - fs: wait for in-flight ops before cloning `File` ([#5803]) ### Changed - rt: reduce time to poll tasks scheduled from outside the runtime ([#5705], [#5720]) ### Added - net: add uds doc alias for unix sockets ([#5659]) - rt: add metric for number of tasks ([#5628]) - sync: implement more traits for channel errors ([#5666]) - net: add nodelay methods on TcpSocket ([#5672]) - sync: add `broadcast::Receiver::blocking_recv` ([#5690]) - process: add `raw_arg` method to `Command` ([#5704]) - io: support PRIORITY epoll events ([#5566]) - task: add `JoinSet::poll_join_next` ([#5721]) - net: add support for Redox OS ([#5790]) ### Unstable - rt: add the ability to dump task backtraces ([#5608], [#5676], [#5708], [#5717]) - rt: instrument task poll times with a histogram ([#5685]) [#5766]: https://github.com/tokio-rs/tokio/pull/5766 [#5653]: https://github.com/tokio-rs/tokio/pull/5653 [#5686]: https://github.com/tokio-rs/tokio/pull/5686 [#5712]: https://github.com/tokio-rs/tokio/pull/5712 [#5693]: https://github.com/tokio-rs/tokio/pull/5693 [#5772]: https://github.com/tokio-rs/tokio/pull/5772 [#5710]: https://github.com/tokio-rs/tokio/pull/5710 [#5803]: https://github.com/tokio-rs/tokio/pull/5803 [#5705]: https://github.com/tokio-rs/tokio/pull/5705 [#5720]: https://github.com/tokio-rs/tokio/pull/5720 [#5659]: https://github.com/tokio-rs/tokio/pull/5659 [#5628]: https://github.com/tokio-rs/tokio/pull/5628 [#5666]: https://github.com/tokio-rs/tokio/pull/5666 [#5672]: https://github.com/tokio-rs/tokio/pull/5672 [#5690]: https://github.com/tokio-rs/tokio/pull/5690 [#5704]: https://github.com/tokio-rs/tokio/pull/5704 [#5566]: https://github.com/tokio-rs/tokio/pull/5566 [#5721]: https://github.com/tokio-rs/tokio/pull/5721 [#5790]: https://github.com/tokio-rs/tokio/pull/5790 [#5608]: https://github.com/tokio-rs/tokio/pull/5608 [#5676]: https://github.com/tokio-rs/tokio/pull/5676 [#5708]: https://github.com/tokio-rs/tokio/pull/5708 [#5717]: https://github.com/tokio-rs/tokio/pull/5717 [#5685]: https://github.com/tokio-rs/tokio/pull/5685 # 1.28.2 (May 28, 2023) Forward ports 1.18.6 changes. ### Fixed - deps: disable default features for mio ([#5728]) [#5728]: https://github.com/tokio-rs/tokio/pull/5728 # 1.28.1 (May 10th, 2023) This release fixes a mistake in the build script that makes `AsFd` implementations unavailable on Rust 1.63. ([#5677]) [#5677]: https://github.com/tokio-rs/tokio/pull/5677 # 1.28.0 (April 25th, 2023) ### Added - io: add `AsyncFd::async_io` ([#5542]) - io: impl BufMut for ReadBuf ([#5590]) - net: add `recv_buf` for `UdpSocket` and `UnixDatagram` ([#5583]) - sync: add `OwnedSemaphorePermit::semaphore` ([#5618]) - sync: add `same_channel` to broadcast channel ([#5607]) - sync: add `watch::Receiver::wait_for` ([#5611]) - task: add `JoinSet::spawn_blocking` and `JoinSet::spawn_blocking_on` ([#5612]) ### Changed - deps: update windows-sys to 0.48 ([#5591]) - io: make `read_to_end` not grow unnecessarily ([#5610]) - macros: make entrypoints more efficient ([#5621]) - sync: improve Debug impl for `RwLock` ([#5647]) - sync: reduce contention in `Notify` ([#5503]) ### Fixed - net: support `get_peer_cred` on AIX ([#5065]) - sync: avoid deadlocks in `broadcast` with custom wakers ([#5578]) ### Documented - sync: fix typo in `Semaphore::MAX_PERMITS` ([#5645]) - sync: fix typo in `tokio::sync::watch::Sender` docs ([#5587]) [#5065]: https://github.com/tokio-rs/tokio/pull/5065 [#5503]: https://github.com/tokio-rs/tokio/pull/5503 [#5542]: https://github.com/tokio-rs/tokio/pull/5542 [#5578]: https://github.com/tokio-rs/tokio/pull/5578 [#5583]: https://github.com/tokio-rs/tokio/pull/5583 [#5587]: https://github.com/tokio-rs/tokio/pull/5587 [#5590]: https://github.com/tokio-rs/tokio/pull/5590 [#5591]: https://github.com/tokio-rs/tokio/pull/5591 [#5607]: https://github.com/tokio-rs/tokio/pull/5607 [#5610]: https://github.com/tokio-rs/tokio/pull/5610 [#5611]: https://github.com/tokio-rs/tokio/pull/5611 [#5612]: https://github.com/tokio-rs/tokio/pull/5612 [#5618]: https://github.com/tokio-rs/tokio/pull/5618 [#5621]: https://github.com/tokio-rs/tokio/pull/5621 [#5645]: https://github.com/tokio-rs/tokio/pull/5645 [#5647]: https://github.com/tokio-rs/tokio/pull/5647 # 1.27.0 (March 27th, 2023) This release bumps the MSRV of Tokio to 1.56. ([#5559]) ### Added - io: add `async_io` helper method to sockets ([#5512]) - io: add implementations of `AsFd`/`AsHandle`/`AsSocket` ([#5514], [#5540]) - net: add `UdpSocket::peek_sender()` ([#5520]) - sync: add `RwLockWriteGuard::{downgrade_map, try_downgrade_map}` ([#5527]) - task: add `JoinHandle::abort_handle` ([#5543]) ### Changed - io: use `memchr` from `libc` ([#5558]) - macros: accept path as crate rename in `#[tokio::main]` ([#5557]) - macros: update to syn 2.0.0 ([#5572]) - time: don't register for a wakeup when `Interval` returns `Ready` ([#5553]) ### Fixed - fs: fuse std iterator in `ReadDir` ([#5555]) - tracing: fix `spawn_blocking` location fields ([#5573]) - time: clean up redundant check in `Wheel::poll()` ([#5574]) ### Documented - macros: define cancellation safety ([#5525]) - io: add details to docs of `tokio::io::copy[_buf]` ([#5575]) - io: refer to `ReaderStream` and `StreamReader` in module docs ([#5576]) [#5512]: https://github.com/tokio-rs/tokio/pull/5512 [#5514]: https://github.com/tokio-rs/tokio/pull/5514 [#5520]: https://github.com/tokio-rs/tokio/pull/5520 [#5525]: https://github.com/tokio-rs/tokio/pull/5525 [#5527]: https://github.com/tokio-rs/tokio/pull/5527 [#5540]: https://github.com/tokio-rs/tokio/pull/5540 [#5543]: https://github.com/tokio-rs/tokio/pull/5543 [#5553]: https://github.com/tokio-rs/tokio/pull/5553 [#5555]: https://github.com/tokio-rs/tokio/pull/5555 [#5557]: https://github.com/tokio-rs/tokio/pull/5557 [#5558]: https://github.com/tokio-rs/tokio/pull/5558 [#5559]: https://github.com/tokio-rs/tokio/pull/5559 [#5572]: https://github.com/tokio-rs/tokio/pull/5572 [#5573]: https://github.com/tokio-rs/tokio/pull/5573 [#5574]: https://github.com/tokio-rs/tokio/pull/5574 [#5575]: https://github.com/tokio-rs/tokio/pull/5575 [#5576]: https://github.com/tokio-rs/tokio/pull/5576 # 1.26.0 (March 1st, 2023) ### Fixed - macros: fix empty `join!` and `try_join!` ([#5504]) - sync: don't leak tracing spans in mutex guards ([#5469]) - sync: drop wakers after unlocking the mutex in Notify ([#5471]) - sync: drop wakers outside lock in semaphore ([#5475]) ### Added - fs: add `fs::try_exists` ([#4299]) - net: add types for named unix pipes ([#5351]) - sync: add `MappedOwnedMutexGuard` ([#5474]) ### Changed - chore: update windows-sys to 0.45 ([#5386]) - net: use Message Read Mode for named pipes ([#5350]) - sync: mark lock guards with `#[clippy::has_significant_drop]` ([#5422]) - sync: reduce contention in watch channel ([#5464]) - time: remove cache padding in timer entries ([#5468]) - time: Improve `Instant::now()` perf with test-util ([#5513]) ### Internal Changes - io: use `poll_fn` in `copy_bidirectional` ([#5486]) - net: refactor named pipe builders to not use bitfields ([#5477]) - rt: remove Arc from Clock ([#5434]) - sync: make `notify_waiters` calls atomic ([#5458]) - time: don't store deadline twice in sleep entries ([#5410]) ### Unstable - metrics: add a new metric for budget exhaustion yields ([#5517]) ### Documented - io: improve AsyncFd example ([#5481]) - runtime: document the nature of the main future ([#5494]) - runtime: remove extra period in docs ([#5511]) - signal: updated Documentation for Signals ([#5459]) - sync: add doc aliases for `blocking_*` methods ([#5448]) - sync: fix docs for Send/Sync bounds in broadcast ([#5480]) - sync: document drop behavior for channels ([#5497]) - task: clarify what happens to spawned work during runtime shutdown ([#5394]) - task: clarify `process::Command` docs ([#5413]) - task: fix wording with 'unsend' ([#5452]) - time: document immediate completion guarantee for timeouts ([#5509]) - tokio: document supported platforms ([#5483]) [#4299]: https://github.com/tokio-rs/tokio/pull/4299 [#5350]: https://github.com/tokio-rs/tokio/pull/5350 [#5351]: https://github.com/tokio-rs/tokio/pull/5351 [#5386]: https://github.com/tokio-rs/tokio/pull/5386 [#5394]: https://github.com/tokio-rs/tokio/pull/5394 [#5410]: https://github.com/tokio-rs/tokio/pull/5410 [#5413]: https://github.com/tokio-rs/tokio/pull/5413 [#5422]: https://github.com/tokio-rs/tokio/pull/5422 [#5434]: https://github.com/tokio-rs/tokio/pull/5434 [#5448]: https://github.com/tokio-rs/tokio/pull/5448 [#5452]: https://github.com/tokio-rs/tokio/pull/5452 [#5458]: https://github.com/tokio-rs/tokio/pull/5458 [#5459]: https://github.com/tokio-rs/tokio/pull/5459 [#5464]: https://github.com/tokio-rs/tokio/pull/5464 [#5468]: https://github.com/tokio-rs/tokio/pull/5468 [#5469]: https://github.com/tokio-rs/tokio/pull/5469 [#5471]: https://github.com/tokio-rs/tokio/pull/5471 [#5474]: https://github.com/tokio-rs/tokio/pull/5474 [#5475]: https://github.com/tokio-rs/tokio/pull/5475 [#5477]: https://github.com/tokio-rs/tokio/pull/5477 [#5480]: https://github.com/tokio-rs/tokio/pull/5480 [#5481]: https://github.com/tokio-rs/tokio/pull/5481 [#5483]: https://github.com/tokio-rs/tokio/pull/5483 [#5486]: https://github.com/tokio-rs/tokio/pull/5486 [#5494]: https://github.com/tokio-rs/tokio/pull/5494 [#5497]: https://github.com/tokio-rs/tokio/pull/5497 [#5504]: https://github.com/tokio-rs/tokio/pull/5504 [#5509]: https://github.com/tokio-rs/tokio/pull/5509 [#5511]: https://github.com/tokio-rs/tokio/pull/5511 [#5513]: https://github.com/tokio-rs/tokio/pull/5513 [#5517]: https://github.com/tokio-rs/tokio/pull/5517 # 1.25.3 (December 17th, 2023) ### Fixed - io: add budgeting to `tokio::runtime::io::registration::async_io` ([#6221]) [#6221]: https://github.com/tokio-rs/tokio/pull/6221 # 1.25.2 (September 22, 2023) Forward ports 1.20.6 changes. ### Changed - io: use `memchr` from `libc` ([#5960]) [#5960]: https://github.com/tokio-rs/tokio/pull/5960 # 1.25.1 (May 28, 2023) Forward ports 1.18.6 changes. ### Fixed - deps: disable default features for mio ([#5728]) [#5728]: https://github.com/tokio-rs/tokio/pull/5728 # 1.25.0 (January 28, 2023) ### Fixed - rt: fix runtime metrics reporting ([#5330]) ### Added - sync: add `broadcast::Sender::len` ([#5343]) ### Changed - fs: increase maximum read buffer size to 2MiB ([#5397]) [#5330]: https://github.com/tokio-rs/tokio/pull/5330 [#5343]: https://github.com/tokio-rs/tokio/pull/5343 [#5397]: https://github.com/tokio-rs/tokio/pull/5397 # 1.24.2 (January 17, 2023) Forward ports 1.18.5 changes. ### Fixed - io: fix unsoundness in `ReadHalf::unsplit` ([#5375]) [#5375]: https://github.com/tokio-rs/tokio/pull/5375 # 1.24.1 (January 6, 2022) This release fixes a compilation failure on targets without `AtomicU64` when using rustc older than 1.63. ([#5356]) [#5356]: https://github.com/tokio-rs/tokio/pull/5356 # 1.24.0 (January 5, 2022) ### Fixed - rt: improve native `AtomicU64` support detection ([#5284]) ### Added - rt: add configuration option for max number of I/O events polled from the OS per tick ([#5186]) - rt: add an environment variable for configuring the default number of worker threads per runtime instance ([#4250]) ### Changed - sync: reduce MPSC channel stack usage ([#5294]) - io: reduce lock contention in I/O operations ([#5300]) - fs: speed up `read_dir()` by chunking operations ([#5309]) - rt: use internal `ThreadId` implementation ([#5329]) - test: don't auto-advance time when a `spawn_blocking` task is running ([#5115]) [#5186]: https://github.com/tokio-rs/tokio/pull/5186 [#5294]: https://github.com/tokio-rs/tokio/pull/5294 [#5284]: https://github.com/tokio-rs/tokio/pull/5284 [#4250]: https://github.com/tokio-rs/tokio/pull/4250 [#5300]: https://github.com/tokio-rs/tokio/pull/5300 [#5329]: https://github.com/tokio-rs/tokio/pull/5329 [#5115]: https://github.com/tokio-rs/tokio/pull/5115 [#5309]: https://github.com/tokio-rs/tokio/pull/5309 # 1.23.1 (January 4, 2022) This release forward ports changes from 1.18.4. ### Fixed - net: fix Windows named pipe server builder to maintain option when toggling pipe mode ([#5336]). [#5336]: https://github.com/tokio-rs/tokio/pull/5336 # 1.23.0 (December 5, 2022) ### Fixed - net: fix Windows named pipe connect ([#5208]) - io: support vectored writes for `ChildStdin` ([#5216]) - io: fix `async fn ready()` false positive for OS-specific events ([#5231]) ### Changed - runtime: `yield_now` defers task until after driver poll ([#5223]) - runtime: reduce amount of codegen needed per spawned task ([#5213]) - windows: replace `winapi` dependency with `windows-sys` ([#5204]) [#5208]: https://github.com/tokio-rs/tokio/pull/5208 [#5216]: https://github.com/tokio-rs/tokio/pull/5216 [#5213]: https://github.com/tokio-rs/tokio/pull/5213 [#5204]: https://github.com/tokio-rs/tokio/pull/5204 [#5223]: https://github.com/tokio-rs/tokio/pull/5223 [#5231]: https://github.com/tokio-rs/tokio/pull/5231 # 1.22.0 (November 17, 2022) ### Added - runtime: add `Handle::runtime_flavor` ([#5138]) - sync: add `Mutex::blocking_lock_owned` ([#5130]) - sync: add `Semaphore::MAX_PERMITS` ([#5144]) - sync: add `merge()` to semaphore permits ([#4948]) - sync: add `mpsc::WeakUnboundedSender` ([#5189]) ### Added (unstable) - process: add `Command::process_group` ([#5114]) - runtime: export metrics about the blocking thread pool ([#5161]) - task: add `task::id()` and `task::try_id()` ([#5171]) ### Fixed - macros: don't take ownership of futures in macros ([#5087]) - runtime: fix Stacked Borrows violation in `LocalOwnedTasks` ([#5099]) - runtime: mitigate ABA with 32-bit queue indices when possible ([#5042]) - task: wake local tasks to the local queue when woken by the same thread ([#5095]) - time: panic in release mode when `mark_pending` called illegally ([#5093]) - runtime: fix typo in expect message ([#5169]) - runtime: fix `unsync_load` on atomic types ([#5175]) - task: elaborate safety comments in task deallocation ([#5172]) - runtime: fix `LocalSet` drop in thread local ([#5179]) - net: remove libc type leakage in a public API ([#5191]) - runtime: update the alignment of `CachePadded` ([#5106]) ### Changed - io: make `tokio::io::copy` continue filling the buffer when writer stalls ([#5066]) - runtime: remove `coop::budget` from `LocalSet::run_until` ([#5155]) - sync: make `Notify` panic safe ([#5154]) ### Documented - io: fix doc for `write_i8` to use signed integers ([#5040]) - net: fix doc typos for TCP and UDP `set_tos` methods ([#5073]) - net: fix function name in `UdpSocket::recv` documentation ([#5150]) - sync: typo in `TryLockError` for `RwLock::try_write` ([#5160]) - task: document that spawned tasks execute immediately ([#5117]) - time: document return type of `timeout` ([#5118]) - time: document that `timeout` checks only before poll ([#5126]) - sync: specify return type of `oneshot::Receiver` in docs ([#5198]) ### Internal changes - runtime: use const `Mutex::new` for globals ([#5061]) - runtime: remove `Option` around `mio::Events` in io driver ([#5078]) - runtime: remove a conditional compilation clause ([#5104]) - runtime: remove a reference to internal time handle ([#5107]) - runtime: misc time driver cleanup ([#5120]) - runtime: move signal driver to runtime module ([#5121]) - runtime: signal driver now uses I/O driver directly ([#5125]) - runtime: start decoupling I/O driver and I/O handle ([#5127]) - runtime: switch `io::handle` refs with scheduler:Handle ([#5128]) - runtime: remove Arc from I/O driver ([#5134]) - runtime: use signal driver handle via `scheduler::Handle` ([#5135]) - runtime: move internal clock fns out of context ([#5139]) - runtime: remove `runtime::context` module ([#5140]) - runtime: keep driver cfgs in `driver.rs` ([#5141]) - runtime: add `runtime::context` to unify thread-locals ([#5143]) - runtime: rename some confusing internal variables/fns ([#5151]) - runtime: move `coop` mod into `runtime` ([#5152]) - runtime: move budget state to context thread-local ([#5157]) - runtime: move park logic into runtime module ([#5158]) - runtime: move `Runtime` into its own file ([#5159]) - runtime: unify entering a runtime with `Handle::enter` ([#5163]) - runtime: remove handle reference from each scheduler ([#5166]) - runtime: move `enter` into `context` ([#5167]) - runtime: combine context and entered thread-locals ([#5168]) - runtime: fix accidental unsetting of current handle ([#5178]) - runtime: move `CoreStage` methods to `Core` ([#5182]) - sync: name mpsc semaphore types ([#5146]) [#4948]: https://github.com/tokio-rs/tokio/pull/4948 [#5040]: https://github.com/tokio-rs/tokio/pull/5040 [#5042]: https://github.com/tokio-rs/tokio/pull/5042 [#5061]: https://github.com/tokio-rs/tokio/pull/5061 [#5066]: https://github.com/tokio-rs/tokio/pull/5066 [#5073]: https://github.com/tokio-rs/tokio/pull/5073 [#5078]: https://github.com/tokio-rs/tokio/pull/5078 [#5087]: https://github.com/tokio-rs/tokio/pull/5087 [#5093]: https://github.com/tokio-rs/tokio/pull/5093 [#5095]: https://github.com/tokio-rs/tokio/pull/5095 [#5099]: https://github.com/tokio-rs/tokio/pull/5099 [#5104]: https://github.com/tokio-rs/tokio/pull/5104 [#5106]: https://github.com/tokio-rs/tokio/pull/5106 [#5107]: https://github.com/tokio-rs/tokio/pull/5107 [#5114]: https://github.com/tokio-rs/tokio/pull/5114 [#5117]: https://github.com/tokio-rs/tokio/pull/5117 [#5118]: https://github.com/tokio-rs/tokio/pull/5118 [#5120]: https://github.com/tokio-rs/tokio/pull/5120 [#5121]: https://github.com/tokio-rs/tokio/pull/5121 [#5125]: https://github.com/tokio-rs/tokio/pull/5125 [#5126]: https://github.com/tokio-rs/tokio/pull/5126 [#5127]: https://github.com/tokio-rs/tokio/pull/5127 [#5128]: https://github.com/tokio-rs/tokio/pull/5128 [#5130]: https://github.com/tokio-rs/tokio/pull/5130 [#5134]: https://github.com/tokio-rs/tokio/pull/5134 [#5135]: https://github.com/tokio-rs/tokio/pull/5135 [#5138]: https://github.com/tokio-rs/tokio/pull/5138 [#5138]: https://github.com/tokio-rs/tokio/pull/5138 [#5139]: https://github.com/tokio-rs/tokio/pull/5139 [#5140]: https://github.com/tokio-rs/tokio/pull/5140 [#5141]: https://github.com/tokio-rs/tokio/pull/5141 [#5143]: https://github.com/tokio-rs/tokio/pull/5143 [#5144]: https://github.com/tokio-rs/tokio/pull/5144 [#5144]: https://github.com/tokio-rs/tokio/pull/5144 [#5146]: https://github.com/tokio-rs/tokio/pull/5146 [#5150]: https://github.com/tokio-rs/tokio/pull/5150 [#5151]: https://github.com/tokio-rs/tokio/pull/5151 [#5152]: https://github.com/tokio-rs/tokio/pull/5152 [#5154]: https://github.com/tokio-rs/tokio/pull/5154 [#5155]: https://github.com/tokio-rs/tokio/pull/5155 [#5157]: https://github.com/tokio-rs/tokio/pull/5157 [#5158]: https://github.com/tokio-rs/tokio/pull/5158 [#5159]: https://github.com/tokio-rs/tokio/pull/5159 [#5160]: https://github.com/tokio-rs/tokio/pull/5160 [#5161]: https://github.com/tokio-rs/tokio/pull/5161 [#5163]: https://github.com/tokio-rs/tokio/pull/5163 [#5166]: https://github.com/tokio-rs/tokio/pull/5166 [#5167]: https://github.com/tokio-rs/tokio/pull/5167 [#5168]: https://github.com/tokio-rs/tokio/pull/5168 [#5169]: https://github.com/tokio-rs/tokio/pull/5169 [#5171]: https://github.com/tokio-rs/tokio/pull/5171 [#5172]: https://github.com/tokio-rs/tokio/pull/5172 [#5175]: https://github.com/tokio-rs/tokio/pull/5175 [#5178]: https://github.com/tokio-rs/tokio/pull/5178 [#5179]: https://github.com/tokio-rs/tokio/pull/5179 [#5182]: https://github.com/tokio-rs/tokio/pull/5182 [#5189]: https://github.com/tokio-rs/tokio/pull/5189 [#5191]: https://github.com/tokio-rs/tokio/pull/5191 [#5198]: https://github.com/tokio-rs/tokio/pull/5198 # 1.21.2 (September 27, 2022) This release removes the dependency on the `once_cell` crate to restore the MSRV of 1.21.x, which is the latest minor version at the time of release. ([#5048]) [#5048]: https://github.com/tokio-rs/tokio/pull/5048 # 1.21.1 (September 13, 2022) ### Fixed - net: fix dependency resolution for socket2 ([#5000]) - task: ignore failure to set TLS in `LocalSet` Drop ([#4976]) [#4976]: https://github.com/tokio-rs/tokio/pull/4976 [#5000]: https://github.com/tokio-rs/tokio/pull/5000 # 1.21.0 (September 2, 2022) This release is the first release of Tokio to intentionally support WASM. The `sync,macros,io-util,rt,time` features are stabilized on WASM. Additionally the wasm32-wasi target is given unstable support for the `net` feature. ### Added - net: add `device` and `bind_device` methods to TCP/UDP sockets ([#4882]) - net: add `tos` and `set_tos` methods to TCP and UDP sockets ([#4877]) - net: add security flags to named pipe `ServerOptions` ([#4845]) - signal: add more windows signal handlers ([#4924]) - sync: add `mpsc::Sender::max_capacity` method ([#4904]) - sync: implement Weak version of `mpsc::Sender` ([#4595]) - task: add `LocalSet::enter` ([#4765]) - task: stabilize `JoinSet` and `AbortHandle` ([#4920]) - tokio: add `track_caller` to public APIs ([#4805], [#4848], [#4852]) - wasm: initial support for `wasm32-wasi` target ([#4716]) ### Fixed - miri: improve miri compatibility by avoiding temporary references in `linked_list::Link` impls ([#4841]) - signal: don't register write interest on signal pipe ([#4898]) - sync: add `#[must_use]` to lock guards ([#4886]) - sync: fix hang when calling `recv` on closed and reopened broadcast channel ([#4867]) - task: propagate attributes on task-locals ([#4837]) ### Changed - fs: change panic to error in `File::start_seek` ([#4897]) - io: reduce syscalls in `poll_read` ([#4840]) - process: use blocking threadpool for child stdio I/O ([#4824]) - signal: make `SignalKind` methods const ([#4956]) ### Internal changes - rt: extract `basic_scheduler::Config` ([#4935]) - rt: move I/O driver into `runtime` module ([#4942]) - rt: rename internal scheduler types ([#4945]) ### Documented - chore: fix typos and grammar ([#4858], [#4894], [#4928]) - io: fix typo in `AsyncSeekExt::rewind` docs ([#4893]) - net: add documentation to `try_read()` for zero-length buffers ([#4937]) - runtime: remove incorrect panic section for `Builder::worker_threads` ([#4849]) - sync: doc of `watch::Sender::send` improved ([#4959]) - task: add cancel safety docs to `JoinHandle` ([#4901]) - task: expand on cancellation of `spawn_blocking` ([#4811]) - time: clarify that the first tick of `Interval::tick` happens immediately ([#4951]) ### Unstable - rt: add unstable option to disable the LIFO slot ([#4936]) - task: fix incorrect signature in `Builder::spawn_on` ([#4953]) - task: make `task::Builder::spawn*` methods fallible ([#4823]) [#4595]: https://github.com/tokio-rs/tokio/pull/4595 [#4716]: https://github.com/tokio-rs/tokio/pull/4716 [#4765]: https://github.com/tokio-rs/tokio/pull/4765 [#4805]: https://github.com/tokio-rs/tokio/pull/4805 [#4811]: https://github.com/tokio-rs/tokio/pull/4811 [#4823]: https://github.com/tokio-rs/tokio/pull/4823 [#4824]: https://github.com/tokio-rs/tokio/pull/4824 [#4837]: https://github.com/tokio-rs/tokio/pull/4837 [#4840]: https://github.com/tokio-rs/tokio/pull/4840 [#4841]: https://github.com/tokio-rs/tokio/pull/4841 [#4845]: https://github.com/tokio-rs/tokio/pull/4845 [#4848]: https://github.com/tokio-rs/tokio/pull/4848 [#4849]: https://github.com/tokio-rs/tokio/pull/4849 [#4852]: https://github.com/tokio-rs/tokio/pull/4852 [#4858]: https://github.com/tokio-rs/tokio/pull/4858 [#4867]: https://github.com/tokio-rs/tokio/pull/4867 [#4877]: https://github.com/tokio-rs/tokio/pull/4877 [#4882]: https://github.com/tokio-rs/tokio/pull/4882 [#4886]: https://github.com/tokio-rs/tokio/pull/4886 [#4893]: https://github.com/tokio-rs/tokio/pull/4893 [#4894]: https://github.com/tokio-rs/tokio/pull/4894 [#4897]: https://github.com/tokio-rs/tokio/pull/4897 [#4898]: https://github.com/tokio-rs/tokio/pull/4898 [#4901]: https://github.com/tokio-rs/tokio/pull/4901 [#4904]: https://github.com/tokio-rs/tokio/pull/4904 [#4920]: https://github.com/tokio-rs/tokio/pull/4920 [#4924]: https://github.com/tokio-rs/tokio/pull/4924 [#4928]: https://github.com/tokio-rs/tokio/pull/4928 [#4935]: https://github.com/tokio-rs/tokio/pull/4935 [#4936]: https://github.com/tokio-rs/tokio/pull/4936 [#4937]: https://github.com/tokio-rs/tokio/pull/4937 [#4942]: https://github.com/tokio-rs/tokio/pull/4942 [#4945]: https://github.com/tokio-rs/tokio/pull/4945 [#4951]: https://github.com/tokio-rs/tokio/pull/4951 [#4953]: https://github.com/tokio-rs/tokio/pull/4953 [#4956]: https://github.com/tokio-rs/tokio/pull/4956 [#4959]: https://github.com/tokio-rs/tokio/pull/4959 # 1.20.6 (September 22, 2023) This is a backport of a change from 1.27.0. ### Changed - io: use `memchr` from `libc` ([#5960]) [#5960]: https://github.com/tokio-rs/tokio/pull/5960 # 1.20.5 (May 28, 2023) Forward ports 1.18.6 changes. ### Fixed - deps: disable default features for mio ([#5728]) [#5728]: https://github.com/tokio-rs/tokio/pull/5728 # 1.20.4 (January 17, 2023) Forward ports 1.18.5 changes. ### Fixed - io: fix unsoundness in `ReadHalf::unsplit` ([#5375]) [#5375]: https://github.com/tokio-rs/tokio/pull/5375 # 1.20.3 (January 3, 2022) This release forward ports changes from 1.18.4. ### Fixed - net: fix Windows named pipe server builder to maintain option when toggling pipe mode ([#5336]). [#5336]: https://github.com/tokio-rs/tokio/pull/5336 # 1.20.2 (September 27, 2022) This release removes the dependency on the `once_cell` crate to restore the MSRV of the 1.20.x LTS release. ([#5048]) [#5048]: https://github.com/tokio-rs/tokio/pull/5048 # 1.20.1 (July 25, 2022) ### Fixed - chore: fix version detection in build script ([#4860]) [#4860]: https://github.com/tokio-rs/tokio/pull/4860 # 1.20.0 (July 12, 2022) ### Added - tokio: add `track_caller` to public APIs ([#4772], [#4791], [#4793], [#4806], [#4808]) - sync: Add `has_changed` method to `watch::Ref` ([#4758]) ### Changed - time: remove `src/time/driver/wheel/stack.rs` ([#4766]) - rt: clean up arguments passed to basic scheduler ([#4767]) - net: be more specific about winapi features ([#4764]) - tokio: use const initialized thread locals where possible ([#4677]) - task: various small improvements to LocalKey ([#4795]) ### Documented - fs: warn about performance pitfall ([#4762]) - chore: fix spelling ([#4769]) - sync: document spurious failures in oneshot ([#4777]) - sync: add warning for watch in non-Send futures ([#4741]) - chore: fix typo ([#4798]) ### Unstable - joinset: rename `join_one` to `join_next` ([#4755]) - rt: unhandled panic config for current thread rt ([#4770]) [#4677]: https://github.com/tokio-rs/tokio/pull/4677 [#4741]: https://github.com/tokio-rs/tokio/pull/4741 [#4755]: https://github.com/tokio-rs/tokio/pull/4755 [#4758]: https://github.com/tokio-rs/tokio/pull/4758 [#4762]: https://github.com/tokio-rs/tokio/pull/4762 [#4764]: https://github.com/tokio-rs/tokio/pull/4764 [#4766]: https://github.com/tokio-rs/tokio/pull/4766 [#4767]: https://github.com/tokio-rs/tokio/pull/4767 [#4769]: https://github.com/tokio-rs/tokio/pull/4769 [#4770]: https://github.com/tokio-rs/tokio/pull/4770 [#4772]: https://github.com/tokio-rs/tokio/pull/4772 [#4777]: https://github.com/tokio-rs/tokio/pull/4777 [#4791]: https://github.com/tokio-rs/tokio/pull/4791 [#4793]: https://github.com/tokio-rs/tokio/pull/4793 [#4795]: https://github.com/tokio-rs/tokio/pull/4795 [#4798]: https://github.com/tokio-rs/tokio/pull/4798 [#4806]: https://github.com/tokio-rs/tokio/pull/4806 [#4808]: https://github.com/tokio-rs/tokio/pull/4808 # 1.19.2 (June 6, 2022) This release fixes another bug in `Notified::enable`. ([#4751]) [#4751]: https://github.com/tokio-rs/tokio/pull/4751 # 1.19.1 (June 5, 2022) This release fixes a bug in `Notified::enable`. ([#4747]) [#4747]: https://github.com/tokio-rs/tokio/pull/4747 # 1.19.0 (June 3, 2022) ### Added - runtime: add `is_finished` method for `JoinHandle` and `AbortHandle` ([#4709]) - runtime: make global queue and event polling intervals configurable ([#4671]) - sync: add `Notified::enable` ([#4705]) - sync: add `watch::Sender::send_if_modified` ([#4591]) - sync: add resubscribe method to broadcast::Receiver ([#4607]) - net: add `take_error` to `TcpSocket` and `TcpStream` ([#4739]) ### Changed - io: refactor out usage of Weak in the io handle ([#4656]) ### Fixed - macros: avoid starvation in `join!` and `try_join!` ([#4624]) ### Documented - runtime: clarify semantics of tasks outliving `block_on` ([#4729]) - time: fix example for `MissedTickBehavior::Burst` ([#4713]) ### Unstable - metrics: correctly update atomics in `IoDriverMetrics` ([#4725]) - metrics: fix compilation with unstable, process, and rt, but without net ([#4682]) - task: add `#[track_caller]` to `JoinSet`/`JoinMap` ([#4697]) - task: add `Builder::{spawn_on, spawn_local_on, spawn_blocking_on}` ([#4683]) - task: add `consume_budget` for cooperative scheduling ([#4498]) - task: add `join_set::Builder` for configuring `JoinSet` tasks ([#4687]) - task: update return value of `JoinSet::join_one` ([#4726]) [#4498]: https://github.com/tokio-rs/tokio/pull/4498 [#4591]: https://github.com/tokio-rs/tokio/pull/4591 [#4607]: https://github.com/tokio-rs/tokio/pull/4607 [#4624]: https://github.com/tokio-rs/tokio/pull/4624 [#4656]: https://github.com/tokio-rs/tokio/pull/4656 [#4671]: https://github.com/tokio-rs/tokio/pull/4671 [#4682]: https://github.com/tokio-rs/tokio/pull/4682 [#4683]: https://github.com/tokio-rs/tokio/pull/4683 [#4687]: https://github.com/tokio-rs/tokio/pull/4687 [#4697]: https://github.com/tokio-rs/tokio/pull/4697 [#4705]: https://github.com/tokio-rs/tokio/pull/4705 [#4709]: https://github.com/tokio-rs/tokio/pull/4709 [#4713]: https://github.com/tokio-rs/tokio/pull/4713 [#4725]: https://github.com/tokio-rs/tokio/pull/4725 [#4726]: https://github.com/tokio-rs/tokio/pull/4726 [#4729]: https://github.com/tokio-rs/tokio/pull/4729 [#4739]: https://github.com/tokio-rs/tokio/pull/4739 # 1.18.6 (May 28, 2023) ### Fixed - deps: disable default features for mio ([#5728]) [#5728]: https://github.com/tokio-rs/tokio/pull/5728 # 1.18.5 (January 17, 2023) ### Fixed - io: fix unsoundness in `ReadHalf::unsplit` ([#5375]) [#5375]: https://github.com/tokio-rs/tokio/pull/5375 # 1.18.4 (January 3, 2022) ### Fixed - net: fix Windows named pipe server builder to maintain option when toggling pipe mode ([#5336]). [#5336]: https://github.com/tokio-rs/tokio/pull/5336 # 1.18.3 (September 27, 2022) This release removes the dependency on the `once_cell` crate to restore the MSRV of the 1.18.x LTS release. ([#5048]) [#5048]: https://github.com/tokio-rs/tokio/pull/5048 # 1.18.2 (May 5, 2022) Add missing features for the `winapi` dependency. ([#4663]) [#4663]: https://github.com/tokio-rs/tokio/pull/4663 # 1.18.1 (May 2, 2022) The 1.18.0 release broke the build for targets without 64-bit atomics when building with `tokio_unstable`. This release fixes that. ([#4649]) [#4649]: https://github.com/tokio-rs/tokio/pull/4649 # 1.18.0 (April 27, 2022) This release adds a number of new APIs in `tokio::net`, `tokio::signal`, and `tokio::sync`. In addition, it adds new unstable APIs to `tokio::task` (`Id`s for uniquely identifying a task, and `AbortHandle` for remotely cancelling a task), as well as a number of bugfixes. ### Fixed - blocking: add missing `#[track_caller]` for `spawn_blocking` ([#4616]) - macros: fix `select` macro to process 64 branches ([#4519]) - net: fix `try_io` methods not calling Mio's `try_io` internally ([#4582]) - runtime: recover when OS fails to spawn a new thread ([#4485]) ### Added - net: add `UdpSocket::peer_addr` ([#4611]) - net: add `try_read_buf` method for named pipes ([#4626]) - signal: add `SignalKind` `Hash`/`Eq` impls and `c_int` conversion ([#4540]) - signal: add support for signals up to `SIGRTMAX` ([#4555]) - sync: add `watch::Sender::send_modify` method ([#4310]) - sync: add `broadcast::Receiver::len` method ([#4542]) - sync: add `watch::Receiver::same_channel` method ([#4581]) - sync: implement `Clone` for `RecvError` types ([#4560]) ### Changed - update `mio` to 0.8.1 ([#4582]) - macros: rename `tokio::select!`'s internal `util` module ([#4543]) - runtime: use `Vec::with_capacity` when building runtime ([#4553]) ### Documented - improve docs for `tokio_unstable` ([#4524]) - runtime: include more documentation for thread_pool/worker ([#4511]) - runtime: update `Handle::current`'s docs to mention `EnterGuard` ([#4567]) - time: clarify platform specific timer resolution ([#4474]) - signal: document that `Signal::recv` is cancel-safe ([#4634]) - sync: `UnboundedReceiver` close docs ([#4548]) ### Unstable The following changes only apply when building with `--cfg tokio_unstable`: - task: add `task::Id` type ([#4630]) - task: add `AbortHandle` type for cancelling tasks in a `JoinSet` ([#4530], [#4640]) - task: fix missing `doc(cfg(...))` attributes for `JoinSet` ([#4531]) - task: fix broken link in `AbortHandle` RustDoc ([#4545]) - metrics: add initial IO driver metrics ([#4507]) [#4616]: https://github.com/tokio-rs/tokio/pull/4616 [#4519]: https://github.com/tokio-rs/tokio/pull/4519 [#4582]: https://github.com/tokio-rs/tokio/pull/4582 [#4485]: https://github.com/tokio-rs/tokio/pull/4485 [#4613]: https://github.com/tokio-rs/tokio/pull/4613 [#4611]: https://github.com/tokio-rs/tokio/pull/4611 [#4626]: https://github.com/tokio-rs/tokio/pull/4626 [#4540]: https://github.com/tokio-rs/tokio/pull/4540 [#4555]: https://github.com/tokio-rs/tokio/pull/4555 [#4310]: https://github.com/tokio-rs/tokio/pull/4310 [#4542]: https://github.com/tokio-rs/tokio/pull/4542 [#4581]: https://github.com/tokio-rs/tokio/pull/4581 [#4560]: https://github.com/tokio-rs/tokio/pull/4560 [#4631]: https://github.com/tokio-rs/tokio/pull/4631 [#4582]: https://github.com/tokio-rs/tokio/pull/4582 [#4543]: https://github.com/tokio-rs/tokio/pull/4543 [#4553]: https://github.com/tokio-rs/tokio/pull/4553 [#4524]: https://github.com/tokio-rs/tokio/pull/4524 [#4511]: https://github.com/tokio-rs/tokio/pull/4511 [#4567]: https://github.com/tokio-rs/tokio/pull/4567 [#4474]: https://github.com/tokio-rs/tokio/pull/4474 [#4634]: https://github.com/tokio-rs/tokio/pull/4634 [#4548]: https://github.com/tokio-rs/tokio/pull/4548 [#4630]: https://github.com/tokio-rs/tokio/pull/4630 [#4530]: https://github.com/tokio-rs/tokio/pull/4530 [#4640]: https://github.com/tokio-rs/tokio/pull/4640 [#4531]: https://github.com/tokio-rs/tokio/pull/4531 [#4545]: https://github.com/tokio-rs/tokio/pull/4545 [#4507]: https://github.com/tokio-rs/tokio/pull/4507 # 1.17.0 (February 16, 2022) This release updates the minimum supported Rust version (MSRV) to 1.49, the `mio` dependency to v0.8, and the (optional) `parking_lot` dependency to v0.12. Additionally, it contains several bug fixes, as well as internal refactoring and performance improvements. ### Fixed - time: prevent panicking in `sleep` with large durations ([#4495]) - time: eliminate potential panics in `Instant` arithmetic on platforms where `Instant::now` is not monotonic ([#4461]) - io: fix `DuplexStream` not participating in cooperative yielding ([#4478]) - rt: fix potential double panic when dropping a `JoinHandle` ([#4430]) ### Changed - update minimum supported Rust version to 1.49 ([#4457]) - update `parking_lot` dependency to v0.12.0 ([#4459]) - update `mio` dependency to v0.8 ([#4449]) - rt: remove an unnecessary lock in the blocking pool ([#4436]) - rt: remove an unnecessary enum in the basic scheduler ([#4462]) - time: use bit manipulation instead of modulo to improve performance ([#4480]) - net: use `std::future::Ready` instead of our own `Ready` future ([#4271]) - replace deprecated `atomic::spin_loop_hint` with `hint::spin_loop` ([#4491]) - fix miri failures in intrusive linked lists ([#4397]) ### Documented - io: add an example for `tokio::process::ChildStdin` ([#4479]) ### Unstable The following changes only apply when building with `--cfg tokio_unstable`: - task: fix missing location information in `tracing` spans generated by `spawn_local` ([#4483]) - task: add `JoinSet` for managing sets of tasks ([#4335]) - metrics: fix compilation error on MIPS ([#4475]) - metrics: fix compilation error on arm32v7 ([#4453]) [#4495]: https://github.com/tokio-rs/tokio/pull/4495 [#4461]: https://github.com/tokio-rs/tokio/pull/4461 [#4478]: https://github.com/tokio-rs/tokio/pull/4478 [#4430]: https://github.com/tokio-rs/tokio/pull/4430 [#4457]: https://github.com/tokio-rs/tokio/pull/4457 [#4459]: https://github.com/tokio-rs/tokio/pull/4459 [#4449]: https://github.com/tokio-rs/tokio/pull/4449 [#4462]: https://github.com/tokio-rs/tokio/pull/4462 [#4436]: https://github.com/tokio-rs/tokio/pull/4436 [#4480]: https://github.com/tokio-rs/tokio/pull/4480 [#4271]: https://github.com/tokio-rs/tokio/pull/4271 [#4491]: https://github.com/tokio-rs/tokio/pull/4491 [#4397]: https://github.com/tokio-rs/tokio/pull/4397 [#4479]: https://github.com/tokio-rs/tokio/pull/4479 [#4483]: https://github.com/tokio-rs/tokio/pull/4483 [#4335]: https://github.com/tokio-rs/tokio/pull/4335 [#4475]: https://github.com/tokio-rs/tokio/pull/4475 [#4453]: https://github.com/tokio-rs/tokio/pull/4453 # 1.16.1 (January 28, 2022) This release fixes a bug in [#4428] with the change [#4437]. [#4428]: https://github.com/tokio-rs/tokio/pull/4428 [#4437]: https://github.com/tokio-rs/tokio/pull/4437 # 1.16.0 (January 27, 2022) Fixes a soundness bug in `io::Take` ([#4428]). The unsoundness is exposed when leaking memory in the given `AsyncRead` implementation and then overwriting the supplied buffer: ```rust impl AsyncRead for Buggy { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_> ) -> Poll> { let new_buf = vec![0; 5].leak(); *buf = ReadBuf::new(new_buf); buf.put_slice(b"hello"); Poll::Ready(Ok(())) } } ``` Also, this release includes improvements to the multi-threaded scheduler that can increase throughput by up to 20% in some cases ([#4383]). ### Fixed - io: **soundness** don't expose uninitialized memory when using `io::Take` in edge case ([#4428]) - fs: ensure `File::write` results in a `write` syscall when the runtime shuts down ([#4316]) - process: drop pipe after child exits in `wait_with_output` ([#4315]) - rt: improve error message when spawning a thread fails ([#4398]) - rt: reduce false-positive thread wakups in the multi-threaded scheduler ([#4383]) - sync: don't inherit `Send` from `parking_lot::*Guard` ([#4359]) ### Added - net: `TcpSocket::linger()` and `set_linger()` ([#4324]) - net: impl `UnwindSafe` for socket types ([#4384]) - rt: impl `UnwindSafe` for `JoinHandle` ([#4418]) - sync: `watch::Receiver::has_changed()` ([#4342]) - sync: `oneshot::Receiver::blocking_recv()` ([#4334]) - sync: `RwLock` blocking operations ([#4425]) ### Unstable The following changes only apply when building with `--cfg tokio_unstable` - rt: **breaking change** overhaul runtime metrics API ([#4373]) [#4428]: https://github.com/tokio-rs/tokio/pull/4428 [#4316]: https://github.com/tokio-rs/tokio/pull/4316 [#4315]: https://github.com/tokio-rs/tokio/pull/4315 [#4398]: https://github.com/tokio-rs/tokio/pull/4398 [#4383]: https://github.com/tokio-rs/tokio/pull/4383 [#4359]: https://github.com/tokio-rs/tokio/pull/4359 [#4324]: https://github.com/tokio-rs/tokio/pull/4324 [#4384]: https://github.com/tokio-rs/tokio/pull/4384 [#4418]: https://github.com/tokio-rs/tokio/pull/4418 [#4342]: https://github.com/tokio-rs/tokio/pull/4342 [#4334]: https://github.com/tokio-rs/tokio/pull/4334 [#4425]: https://github.com/tokio-rs/tokio/pull/4425 [#4373]: https://github.com/tokio-rs/tokio/pull/4373 # 1.15.0 (December 15, 2021) ### Fixed - io: add cooperative yielding support to `io::empty()` ([#4300]) - time: make timeout robust against budget-depleting tasks ([#4314]) ### Changed - update minimum supported Rust version to 1.46. ### Added - time: add `Interval::reset()` ([#4248]) - io: add explicit lifetimes to `AsyncFdReadyGuard` ([#4267]) - process: add `Command::as_std()` ([#4295]) ### Added (unstable) - tracing: instrument `tokio::sync` types ([#4302]) [#4302]: https://github.com/tokio-rs/tokio/pull/4302 [#4300]: https://github.com/tokio-rs/tokio/pull/4300 [#4295]: https://github.com/tokio-rs/tokio/pull/4295 [#4267]: https://github.com/tokio-rs/tokio/pull/4267 [#4248]: https://github.com/tokio-rs/tokio/pull/4248 [#4314]: https://github.com/tokio-rs/tokio/pull/4314 # 1.14.0 (November 15, 2021) ### Fixed - macros: fix compiler errors when using `mut` patterns in `select!` ([#4211]) - sync: fix a data race between `oneshot::Sender::send` and awaiting a `oneshot::Receiver` when the oneshot has been closed ([#4226]) - sync: make `AtomicWaker` panic safe ([#3689]) - runtime: fix basic scheduler dropping tasks outside a runtime context ([#4213]) ### Added - stats: add `RuntimeStats::busy_duration_total` ([#4179], [#4223]) ### Changed - io: updated `copy` buffer size to match `std::io::copy` ([#4209]) ### Documented - io: rename buffer to file in doc-test ([#4230]) - sync: fix Notify example ([#4212]) [#4211]: https://github.com/tokio-rs/tokio/pull/4211 [#4226]: https://github.com/tokio-rs/tokio/pull/4226 [#3689]: https://github.com/tokio-rs/tokio/pull/3689 [#4213]: https://github.com/tokio-rs/tokio/pull/4213 [#4179]: https://github.com/tokio-rs/tokio/pull/4179 [#4223]: https://github.com/tokio-rs/tokio/pull/4223 [#4209]: https://github.com/tokio-rs/tokio/pull/4209 [#4230]: https://github.com/tokio-rs/tokio/pull/4230 [#4212]: https://github.com/tokio-rs/tokio/pull/4212 # 1.13.1 (November 15, 2021) ### Fixed - sync: fix a data race between `oneshot::Sender::send` and awaiting a `oneshot::Receiver` when the oneshot has been closed ([#4226]) [#4226]: https://github.com/tokio-rs/tokio/pull/4226 # 1.13.0 (October 29, 2021) ### Fixed - sync: fix `Notify` to clone the waker before locking its waiter list ([#4129]) - tokio: add riscv32 to non atomic64 architectures ([#4185]) ### Added - net: add `poll_{recv,send}_ready` methods to `udp` and `uds_datagram` ([#4131]) - net: add `try_*`, `readable`, `writable`, `ready`, and `peer_addr` methods to split halves ([#4120]) - sync: add `blocking_lock` to `Mutex` ([#4130]) - sync: add `watch::Sender::send_replace` ([#3962], [#4195]) - sync: expand `Debug` for `Mutex` impl to unsized `T` ([#4134]) - tracing: instrument time::Sleep ([#4072]) - tracing: use structured location fields for spawned tasks ([#4128]) ### Changed - io: add assert in `copy_bidirectional` that `poll_write` is sensible ([#4125]) - macros: use qualified syntax when polling in `select!` ([#4192]) - runtime: handle `block_on` wakeups better ([#4157]) - task: allocate callback on heap immediately in debug mode ([#4203]) - tokio: assert platform-minimum requirements at build time ([#3797]) ### Documented - docs: conversion of doc comments to indicative mood ([#4174]) - docs: add returning on the first error example for `try_join!` ([#4133]) - docs: fixing broken links in `tokio/src/lib.rs` ([#4132]) - signal: add example with background listener ([#4171]) - sync: add more oneshot examples ([#4153]) - time: document `Interval::tick` cancel safety ([#4152]) [#3797]: https://github.com/tokio-rs/tokio/pull/3797 [#3962]: https://github.com/tokio-rs/tokio/pull/3962 [#4072]: https://github.com/tokio-rs/tokio/pull/4072 [#4120]: https://github.com/tokio-rs/tokio/pull/4120 [#4125]: https://github.com/tokio-rs/tokio/pull/4125 [#4128]: https://github.com/tokio-rs/tokio/pull/4128 [#4129]: https://github.com/tokio-rs/tokio/pull/4129 [#4130]: https://github.com/tokio-rs/tokio/pull/4130 [#4131]: https://github.com/tokio-rs/tokio/pull/4131 [#4132]: https://github.com/tokio-rs/tokio/pull/4132 [#4133]: https://github.com/tokio-rs/tokio/pull/4133 [#4134]: https://github.com/tokio-rs/tokio/pull/4134 [#4152]: https://github.com/tokio-rs/tokio/pull/4152 [#4153]: https://github.com/tokio-rs/tokio/pull/4153 [#4157]: https://github.com/tokio-rs/tokio/pull/4157 [#4171]: https://github.com/tokio-rs/tokio/pull/4171 [#4174]: https://github.com/tokio-rs/tokio/pull/4174 [#4185]: https://github.com/tokio-rs/tokio/pull/4185 [#4192]: https://github.com/tokio-rs/tokio/pull/4192 [#4195]: https://github.com/tokio-rs/tokio/pull/4195 [#4203]: https://github.com/tokio-rs/tokio/pull/4203 # 1.12.0 (September 21, 2021) ### Fixed - mpsc: ensure `try_reserve` error is consistent with `try_send` ([#4119]) - mpsc: use `spin_loop_hint` instead of `yield_now` ([#4115]) - sync: make `SendError` field public ([#4097]) ### Added - io: add POSIX AIO on FreeBSD ([#4054]) - io: add convenience method `AsyncSeekExt::rewind` ([#4107]) - runtime: add tracing span for `block_on` futures ([#4094]) - runtime: callback when a worker parks and unparks ([#4070]) - sync: implement `try_recv` for mpsc channels ([#4113]) ### Documented - docs: clarify CPU-bound tasks on Tokio ([#4105]) - mpsc: document spurious failures on `poll_recv` ([#4117]) - mpsc: document that `PollSender` impls `Sink` ([#4110]) - task: document non-guarantees of `yield_now` ([#4091]) - time: document paused time details better ([#4061], [#4103]) [#4027]: https://github.com/tokio-rs/tokio/pull/4027 [#4054]: https://github.com/tokio-rs/tokio/pull/4054 [#4061]: https://github.com/tokio-rs/tokio/pull/4061 [#4070]: https://github.com/tokio-rs/tokio/pull/4070 [#4091]: https://github.com/tokio-rs/tokio/pull/4091 [#4094]: https://github.com/tokio-rs/tokio/pull/4094 [#4097]: https://github.com/tokio-rs/tokio/pull/4097 [#4103]: https://github.com/tokio-rs/tokio/pull/4103 [#4105]: https://github.com/tokio-rs/tokio/pull/4105 [#4107]: https://github.com/tokio-rs/tokio/pull/4107 [#4110]: https://github.com/tokio-rs/tokio/pull/4110 [#4113]: https://github.com/tokio-rs/tokio/pull/4113 [#4115]: https://github.com/tokio-rs/tokio/pull/4115 [#4117]: https://github.com/tokio-rs/tokio/pull/4117 [#4119]: https://github.com/tokio-rs/tokio/pull/4119 # 1.11.0 (August 31, 2021) ### Fixed - time: don't panic when Instant is not monotonic ([#4044]) - io: fix panic in `fill_buf` by not calling `poll_fill_buf` twice ([#4084]) ### Added - watch: add `watch::Sender::subscribe` ([#3800]) - process: add `from_std` to `ChildStd*` ([#4045]) - stats: initial work on runtime stats ([#4043]) ### Changed - tracing: change span naming to new console convention ([#4042]) - io: speed-up waking by using uninitialized array ([#4055], [#4071], [#4075]) ### Documented - time: make Sleep examples easier to find ([#4040]) [#3800]: https://github.com/tokio-rs/tokio/pull/3800 [#4040]: https://github.com/tokio-rs/tokio/pull/4040 [#4042]: https://github.com/tokio-rs/tokio/pull/4042 [#4043]: https://github.com/tokio-rs/tokio/pull/4043 [#4044]: https://github.com/tokio-rs/tokio/pull/4044 [#4045]: https://github.com/tokio-rs/tokio/pull/4045 [#4055]: https://github.com/tokio-rs/tokio/pull/4055 [#4071]: https://github.com/tokio-rs/tokio/pull/4071 [#4075]: https://github.com/tokio-rs/tokio/pull/4075 [#4084]: https://github.com/tokio-rs/tokio/pull/4084 # 1.10.1 (August 24, 2021) ### Fixed - runtime: fix leak in UnownedTask ([#4063]) [#4063]: https://github.com/tokio-rs/tokio/pull/4063 # 1.10.0 (August 12, 2021) ### Added - io: add `(read|write)_f(32|64)[_le]` methods ([#4022]) - io: add `fill_buf` and `consume` to `AsyncBufReadExt` ([#3991]) - process: add `Child::raw_handle()` on windows ([#3998]) ### Fixed - doc: fix non-doc builds with `--cfg docsrs` ([#4020]) - io: flush eagerly in `io::copy` ([#4001]) - runtime: a debug assert was sometimes triggered during shutdown ([#4005]) - sync: use `spin_loop_hint` instead of `yield_now` in mpsc ([#4037]) - tokio: the test-util feature depends on rt, sync, and time ([#4036]) ### Changes - runtime: reorganize parts of the runtime ([#3979], [#4005]) - signal: make windows docs for signal module show up on unix builds ([#3770]) - task: quickly send task to heap on debug mode ([#4009]) ### Documented - io: document cancellation safety of `AsyncBufReadExt` ([#3997]) - sync: document when `watch::send` fails ([#4021]) [#3770]: https://github.com/tokio-rs/tokio/pull/3770 [#3979]: https://github.com/tokio-rs/tokio/pull/3979 [#3991]: https://github.com/tokio-rs/tokio/pull/3991 [#3997]: https://github.com/tokio-rs/tokio/pull/3997 [#3998]: https://github.com/tokio-rs/tokio/pull/3998 [#4001]: https://github.com/tokio-rs/tokio/pull/4001 [#4005]: https://github.com/tokio-rs/tokio/pull/4005 [#4009]: https://github.com/tokio-rs/tokio/pull/4009 [#4020]: https://github.com/tokio-rs/tokio/pull/4020 [#4021]: https://github.com/tokio-rs/tokio/pull/4021 [#4022]: https://github.com/tokio-rs/tokio/pull/4022 [#4036]: https://github.com/tokio-rs/tokio/pull/4036 [#4037]: https://github.com/tokio-rs/tokio/pull/4037 # 1.9.0 (July 22, 2021) ### Added - net: allow customized I/O operations for `TcpStream` ([#3888]) - sync: add getter for the mutex from a guard ([#3928]) - task: expose nameable future for `TaskLocal::scope` ([#3273]) ### Fixed - Fix leak if output of future panics on drop ([#3967]) - Fix leak in `LocalSet` ([#3978]) ### Changes - runtime: reorganize parts of the runtime ([#3909], [#3939], [#3950], [#3955], [#3980]) - sync: clean up `OnceCell` ([#3945]) - task: remove mutex in `JoinError` ([#3959]) [#3273]: https://github.com/tokio-rs/tokio/pull/3273 [#3888]: https://github.com/tokio-rs/tokio/pull/3888 [#3909]: https://github.com/tokio-rs/tokio/pull/3909 [#3928]: https://github.com/tokio-rs/tokio/pull/3928 [#3934]: https://github.com/tokio-rs/tokio/pull/3934 [#3939]: https://github.com/tokio-rs/tokio/pull/3939 [#3945]: https://github.com/tokio-rs/tokio/pull/3945 [#3950]: https://github.com/tokio-rs/tokio/pull/3950 [#3955]: https://github.com/tokio-rs/tokio/pull/3955 [#3959]: https://github.com/tokio-rs/tokio/pull/3959 [#3967]: https://github.com/tokio-rs/tokio/pull/3967 [#3978]: https://github.com/tokio-rs/tokio/pull/3978 [#3980]: https://github.com/tokio-rs/tokio/pull/3980 # 1.8.3 (July 26, 2021) This release backports two fixes from 1.9.0 ### Fixed - Fix leak if output of future panics on drop ([#3967]) - Fix leak in `LocalSet` ([#3978]) [#3967]: https://github.com/tokio-rs/tokio/pull/3967 [#3978]: https://github.com/tokio-rs/tokio/pull/3978 # 1.8.2 (July 19, 2021) Fixes a missed edge case from 1.8.1. ### Fixed - runtime: drop canceled future on next poll ([#3965]) [#3965]: https://github.com/tokio-rs/tokio/pull/3965 # 1.8.1 (July 6, 2021) Forward ports 1.5.1 fixes. ### Fixed - runtime: remotely abort tasks on `JoinHandle::abort` ([#3934]) [#3934]: https://github.com/tokio-rs/tokio/pull/3934 # 1.8.0 (July 2, 2021) ### Added - io: add `get_{ref,mut}` methods to `AsyncFdReadyGuard` and `AsyncFdReadyMutGuard` ([#3807]) - io: efficient implementation of vectored writes for `BufWriter` ([#3163]) - net: add ready/try methods to `NamedPipe{Client,Server}` ([#3866], [#3899]) - sync: add `watch::Receiver::borrow_and_update` ([#3813]) - sync: implement `From` for `OnceCell` ([#3877]) - time: allow users to specify Interval behavior when delayed ([#3721]) ### Added (unstable) - rt: add `tokio::task::Builder` ([#3881]) ### Fixed - net: handle HUP event with `UnixStream` ([#3898]) ### Documented - doc: document cancellation safety ([#3900]) - time: add wait alias to sleep ([#3897]) - time: document auto-advancing behavior of runtime ([#3763]) [#3163]: https://github.com/tokio-rs/tokio/pull/3163 [#3721]: https://github.com/tokio-rs/tokio/pull/3721 [#3763]: https://github.com/tokio-rs/tokio/pull/3763 [#3807]: https://github.com/tokio-rs/tokio/pull/3807 [#3813]: https://github.com/tokio-rs/tokio/pull/3813 [#3866]: https://github.com/tokio-rs/tokio/pull/3866 [#3877]: https://github.com/tokio-rs/tokio/pull/3877 [#3881]: https://github.com/tokio-rs/tokio/pull/3881 [#3897]: https://github.com/tokio-rs/tokio/pull/3897 [#3898]: https://github.com/tokio-rs/tokio/pull/3898 [#3899]: https://github.com/tokio-rs/tokio/pull/3899 [#3900]: https://github.com/tokio-rs/tokio/pull/3900 # 1.7.2 (July 6, 2021) Forward ports 1.5.1 fixes. ### Fixed - runtime: remotely abort tasks on `JoinHandle::abort` ([#3934]) [#3934]: https://github.com/tokio-rs/tokio/pull/3934 # 1.7.1 (June 18, 2021) ### Fixed - runtime: fix early task shutdown during runtime shutdown ([#3870]) [#3870]: https://github.com/tokio-rs/tokio/pull/3870 # 1.7.0 (June 15, 2021) ### Added - net: add named pipes on windows ([#3760]) - net: add `TcpSocket` from `std::net::TcpStream` conversion ([#3838]) - sync: add `receiver_count` to `watch::Sender` ([#3729]) - sync: export `sync::notify::Notified` future publicly ([#3840]) - tracing: instrument task wakers ([#3836]) ### Fixed - macros: suppress `clippy::default_numeric_fallback` lint in generated code ([#3831]) - runtime: immediately drop new tasks when runtime is shut down ([#3752]) - sync: deprecate unused `mpsc::RecvError` type ([#3833]) ### Documented - io: clarify EOF condition for `AsyncReadExt::read_buf` ([#3850]) - io: clarify limits on return values of `AsyncWrite::poll_write` ([#3820]) - sync: add examples to Semaphore ([#3808]) [#3729]: https://github.com/tokio-rs/tokio/pull/3729 [#3752]: https://github.com/tokio-rs/tokio/pull/3752 [#3760]: https://github.com/tokio-rs/tokio/pull/3760 [#3808]: https://github.com/tokio-rs/tokio/pull/3808 [#3820]: https://github.com/tokio-rs/tokio/pull/3820 [#3831]: https://github.com/tokio-rs/tokio/pull/3831 [#3833]: https://github.com/tokio-rs/tokio/pull/3833 [#3836]: https://github.com/tokio-rs/tokio/pull/3836 [#3838]: https://github.com/tokio-rs/tokio/pull/3838 [#3840]: https://github.com/tokio-rs/tokio/pull/3840 [#3850]: https://github.com/tokio-rs/tokio/pull/3850 # 1.6.3 (July 6, 2021) Forward ports 1.5.1 fixes. ### Fixed - runtime: remotely abort tasks on `JoinHandle::abort` ([#3934]) [#3934]: https://github.com/tokio-rs/tokio/pull/3934 # 1.6.2 (June 14, 2021) ### Fixes - test: sub-ms `time:advance` regression introduced in 1.6 ([#3852]) [#3852]: https://github.com/tokio-rs/tokio/pull/3852 # 1.6.1 (May 28, 2021) This release reverts [#3518] because it doesn't work on some kernels due to a kernel bug. ([#3803]) [#3518]: https://github.com/tokio-rs/tokio/issues/3518 [#3803]: https://github.com/tokio-rs/tokio/issues/3803 # 1.6.0 (May 14, 2021) ### Added - fs: try doing a non-blocking read before punting to the threadpool ([#3518]) - io: add `write_all_buf` to `AsyncWriteExt` ([#3737]) - io: implement `AsyncSeek` for `BufReader`, `BufWriter`, and `BufStream` ([#3491]) - net: support non-blocking vectored I/O ([#3761]) - sync: add `mpsc::Sender::{reserve_owned, try_reserve_owned}` ([#3704]) - sync: add a `MutexGuard::map` method that returns a `MappedMutexGuard` ([#2472]) - time: add getter for Interval's period ([#3705]) ### Fixed - io: wake pending writers on `DuplexStream` close ([#3756]) - process: avoid redundant effort to reap orphan processes ([#3743]) - signal: use `std::os::raw::c_int` instead of `libc::c_int` on public API ([#3774]) - sync: preserve permit state in `notify_waiters` ([#3660]) - task: update `JoinHandle` panic message ([#3727]) - time: prevent `time::advance` from going too far ([#3712]) ### Documented - net: hide `net::unix::datagram` module from docs ([#3775]) - process: updated example ([#3748]) - sync: `Barrier` doc should use task, not thread ([#3780]) - task: update documentation on `block_in_place` ([#3753]) [#2472]: https://github.com/tokio-rs/tokio/pull/2472 [#3491]: https://github.com/tokio-rs/tokio/pull/3491 [#3518]: https://github.com/tokio-rs/tokio/pull/3518 [#3660]: https://github.com/tokio-rs/tokio/pull/3660 [#3704]: https://github.com/tokio-rs/tokio/pull/3704 [#3705]: https://github.com/tokio-rs/tokio/pull/3705 [#3712]: https://github.com/tokio-rs/tokio/pull/3712 [#3727]: https://github.com/tokio-rs/tokio/pull/3727 [#3737]: https://github.com/tokio-rs/tokio/pull/3737 [#3743]: https://github.com/tokio-rs/tokio/pull/3743 [#3748]: https://github.com/tokio-rs/tokio/pull/3748 [#3753]: https://github.com/tokio-rs/tokio/pull/3753 [#3756]: https://github.com/tokio-rs/tokio/pull/3756 [#3761]: https://github.com/tokio-rs/tokio/pull/3761 [#3774]: https://github.com/tokio-rs/tokio/pull/3774 [#3775]: https://github.com/tokio-rs/tokio/pull/3775 [#3780]: https://github.com/tokio-rs/tokio/pull/3780 # 1.5.1 (July 6, 2021) ### Fixed - runtime: remotely abort tasks on `JoinHandle::abort` ([#3934]) [#3934]: https://github.com/tokio-rs/tokio/pull/3934 # 1.5.0 (April 12, 2021) ### Added - io: add `AsyncSeekExt::stream_position` ([#3650]) - io: add `AsyncWriteExt::write_vectored` ([#3678]) - io: add a `copy_bidirectional` utility ([#3572]) - net: implement `IntoRawFd` for `TcpSocket` ([#3684]) - sync: add `OnceCell` ([#3591]) - sync: add `OwnedRwLockReadGuard` and `OwnedRwLockWriteGuard` ([#3340]) - sync: add `Semaphore::is_closed` ([#3673]) - sync: add `mpsc::Sender::capacity` ([#3690]) - sync: allow configuring `RwLock` max reads ([#3644]) - task: add `sync_scope` for `LocalKey` ([#3612]) ### Fixed - chore: try to avoid `noalias` attributes on intrusive linked list ([#3654]) - rt: fix panic in `JoinHandle::abort()` when called from other threads ([#3672]) - sync: don't panic in `oneshot::try_recv` ([#3674]) - sync: fix notifications getting dropped on receiver drop ([#3652]) - sync: fix `Semaphore` permit overflow calculation ([#3644]) ### Documented - io: clarify requirements of `AsyncFd` ([#3635]) - runtime: fix unclear docs for `{Handle,Runtime}::block_on` ([#3628]) - sync: document that `Semaphore` is fair ([#3693]) - sync: improve doc on blocking mutex ([#3645]) [#3340]: https://github.com/tokio-rs/tokio/pull/3340 [#3572]: https://github.com/tokio-rs/tokio/pull/3572 [#3591]: https://github.com/tokio-rs/tokio/pull/3591 [#3612]: https://github.com/tokio-rs/tokio/pull/3612 [#3628]: https://github.com/tokio-rs/tokio/pull/3628 [#3635]: https://github.com/tokio-rs/tokio/pull/3635 [#3644]: https://github.com/tokio-rs/tokio/pull/3644 [#3645]: https://github.com/tokio-rs/tokio/pull/3645 [#3650]: https://github.com/tokio-rs/tokio/pull/3650 [#3652]: https://github.com/tokio-rs/tokio/pull/3652 [#3654]: https://github.com/tokio-rs/tokio/pull/3654 [#3672]: https://github.com/tokio-rs/tokio/pull/3672 [#3673]: https://github.com/tokio-rs/tokio/pull/3673 [#3674]: https://github.com/tokio-rs/tokio/pull/3674 [#3678]: https://github.com/tokio-rs/tokio/pull/3678 [#3684]: https://github.com/tokio-rs/tokio/pull/3684 [#3690]: https://github.com/tokio-rs/tokio/pull/3690 [#3693]: https://github.com/tokio-rs/tokio/pull/3693 # 1.4.0 (March 20, 2021) ### Added - macros: introduce biased argument for `select!` ([#3603]) - runtime: add `Handle::block_on` ([#3569]) ### Fixed - runtime: avoid unnecessary polling of `block_on` future ([#3582]) - runtime: fix memory leak/growth when creating many runtimes ([#3564]) - runtime: mark `EnterGuard` with `must_use` ([#3609]) ### Documented - chore: mention fix for building docs in contributing guide ([#3618]) - doc: add link to `PollSender` ([#3613]) - doc: alias sleep to delay ([#3604]) - sync: improve `Mutex` FIFO explanation ([#3615]) - timer: fix double newline in module docs ([#3617]) [#3564]: https://github.com/tokio-rs/tokio/pull/3564 [#3613]: https://github.com/tokio-rs/tokio/pull/3613 [#3618]: https://github.com/tokio-rs/tokio/pull/3618 [#3617]: https://github.com/tokio-rs/tokio/pull/3617 [#3582]: https://github.com/tokio-rs/tokio/pull/3582 [#3615]: https://github.com/tokio-rs/tokio/pull/3615 [#3603]: https://github.com/tokio-rs/tokio/pull/3603 [#3609]: https://github.com/tokio-rs/tokio/pull/3609 [#3604]: https://github.com/tokio-rs/tokio/pull/3604 [#3569]: https://github.com/tokio-rs/tokio/pull/3569 # 1.3.0 (March 9, 2021) ### Added - coop: expose an `unconstrained()` opt-out ([#3547]) - net: add `into_std` for net types without it ([#3509]) - sync: add `same_channel` method to `mpsc::Sender` ([#3532]) - sync: add `{try_,}acquire_many_owned` to `Semaphore` ([#3535]) - sync: add back `RwLockWriteGuard::map` and `RwLockWriteGuard::try_map` ([#3348]) ### Fixed - sync: allow `oneshot::Receiver::close` after successful `try_recv` ([#3552]) - time: do not panic on `timeout(Duration::MAX)` ([#3551]) ### Documented - doc: doc aliases for pre-1.0 function names ([#3523]) - io: fix typos ([#3541]) - io: note the EOF behavior of `read_until` ([#3536]) - io: update `AsyncRead::poll_read` doc ([#3557]) - net: update `UdpSocket` splitting doc ([#3517]) - runtime: add link to `LocalSet` on `new_current_thread` ([#3508]) - runtime: update documentation of thread limits ([#3527]) - sync: do not recommend `join_all` for `Barrier` ([#3514]) - sync: documentation for `oneshot` ([#3592]) - sync: rename `notify` to `notify_one` ([#3526]) - time: fix typo in `Sleep` doc ([#3515]) - time: sync `interval.rs` and `time/mod.rs` docs ([#3533]) [#3348]: https://github.com/tokio-rs/tokio/pull/3348 [#3508]: https://github.com/tokio-rs/tokio/pull/3508 [#3509]: https://github.com/tokio-rs/tokio/pull/3509 [#3514]: https://github.com/tokio-rs/tokio/pull/3514 [#3515]: https://github.com/tokio-rs/tokio/pull/3515 [#3517]: https://github.com/tokio-rs/tokio/pull/3517 [#3523]: https://github.com/tokio-rs/tokio/pull/3523 [#3526]: https://github.com/tokio-rs/tokio/pull/3526 [#3527]: https://github.com/tokio-rs/tokio/pull/3527 [#3532]: https://github.com/tokio-rs/tokio/pull/3532 [#3533]: https://github.com/tokio-rs/tokio/pull/3533 [#3535]: https://github.com/tokio-rs/tokio/pull/3535 [#3536]: https://github.com/tokio-rs/tokio/pull/3536 [#3541]: https://github.com/tokio-rs/tokio/pull/3541 [#3547]: https://github.com/tokio-rs/tokio/pull/3547 [#3551]: https://github.com/tokio-rs/tokio/pull/3551 [#3552]: https://github.com/tokio-rs/tokio/pull/3552 [#3557]: https://github.com/tokio-rs/tokio/pull/3557 [#3592]: https://github.com/tokio-rs/tokio/pull/3592 # 1.2.0 (February 5, 2021) ### Added - signal: make `Signal::poll_recv` method public ([#3383]) ### Fixed - time: make `test-util` paused time fully deterministic ([#3492]) ### Documented - sync: link to new broadcast and watch wrappers ([#3504]) [#3383]: https://github.com/tokio-rs/tokio/pull/3383 [#3492]: https://github.com/tokio-rs/tokio/pull/3492 [#3504]: https://github.com/tokio-rs/tokio/pull/3504 # 1.1.1 (January 29, 2021) Forward ports 1.0.3 fix. ### Fixed - io: memory leak during shutdown ([#3477]). # 1.1.0 (January 22, 2021) ### Added - net: add `try_read_buf` and `try_recv_buf` ([#3351]) - mpsc: Add `Sender::try_reserve` function ([#3418]) - sync: add `RwLock` `try_read` and `try_write` methods ([#3400]) - io: add `ReadBuf::inner_mut` ([#3443]) ### Changed - macros: improve `select!` error message ([#3352]) - io: keep track of initialized bytes in `read_to_end` ([#3426]) - runtime: consolidate errors for context missing ([#3441]) ### Fixed - task: wake `LocalSet` on `spawn_local` ([#3369]) - sync: fix panic in broadcast::Receiver drop ([#3434]) ### Documented - stream: link to new `Stream` wrappers in `tokio-stream` ([#3343]) - docs: mention that `test-util` feature is not enabled with full ([#3397]) - process: add documentation to process::Child fields ([#3437]) - io: clarify `AsyncFd` docs about changes of the inner fd ([#3430]) - net: update datagram docs on splitting ([#3448]) - time: document that `Sleep` is not `Unpin` ([#3457]) - sync: add link to `PollSemaphore` ([#3456]) - task: add `LocalSet` example ([#3438]) - sync: improve bounded `mpsc` documentation ([#3458]) [#3343]: https://github.com/tokio-rs/tokio/pull/3343 [#3351]: https://github.com/tokio-rs/tokio/pull/3351 [#3352]: https://github.com/tokio-rs/tokio/pull/3352 [#3369]: https://github.com/tokio-rs/tokio/pull/3369 [#3397]: https://github.com/tokio-rs/tokio/pull/3397 [#3400]: https://github.com/tokio-rs/tokio/pull/3400 [#3418]: https://github.com/tokio-rs/tokio/pull/3418 [#3426]: https://github.com/tokio-rs/tokio/pull/3426 [#3430]: https://github.com/tokio-rs/tokio/pull/3430 [#3434]: https://github.com/tokio-rs/tokio/pull/3434 [#3437]: https://github.com/tokio-rs/tokio/pull/3437 [#3438]: https://github.com/tokio-rs/tokio/pull/3438 [#3441]: https://github.com/tokio-rs/tokio/pull/3441 [#3443]: https://github.com/tokio-rs/tokio/pull/3443 [#3448]: https://github.com/tokio-rs/tokio/pull/3448 [#3456]: https://github.com/tokio-rs/tokio/pull/3456 [#3457]: https://github.com/tokio-rs/tokio/pull/3457 [#3458]: https://github.com/tokio-rs/tokio/pull/3458 # 1.0.3 (January 28, 2021) ### Fixed - io: memory leak during shutdown ([#3477]). [#3477]: https://github.com/tokio-rs/tokio/pull/3477 # 1.0.2 (January 14, 2021) ### Fixed - io: soundness in `read_to_end` ([#3428]). [#3428]: https://github.com/tokio-rs/tokio/pull/3428 # 1.0.1 (December 25, 2020) This release fixes a soundness hole caused by the combination of `RwLockWriteGuard::map` and `RwLockWriteGuard::downgrade` by removing the `map` function. This is a breaking change, but breaking changes are allowed under our semver policy when they are required to fix a soundness hole. (See [this RFC][semver] for more.) Note that we have chosen not to do a deprecation cycle or similar because Tokio 1.0.0 was released two days ago, and therefore the impact should be minimal. Due to the soundness hole, we have also yanked Tokio version 1.0.0. ### Removed - sync: remove `RwLockWriteGuard::map` and `RwLockWriteGuard::try_map` ([#3345]) ### Fixed - docs: remove stream feature from docs ([#3335]) [semver]: https://github.com/rust-lang/rfcs/blob/master/text/1122-language-semver.md#soundness-changes [#3335]: https://github.com/tokio-rs/tokio/pull/3335 [#3345]: https://github.com/tokio-rs/tokio/pull/3345 # 1.0.0 (December 23, 2020) Commit to the API and long-term support. ### Fixed - sync: spurious wakeup in `watch` ([#3234]). ### Changed - io: rename `AsyncFd::with_io()` to `try_io()` ([#3306]) - fs: avoid OS specific `*Ext` traits in favor of conditionally defining the fn ([#3264]). - fs: `Sleep` is `!Unpin` ([#3278]). - net: pass `SocketAddr` by value ([#3125]). - net: `TcpStream::poll_peek` takes `ReadBuf` ([#3259]). - rt: rename `runtime::Builder::max_threads()` to `max_blocking_threads()` ([#3287]). - time: require `current_thread` runtime when calling `time::pause()` ([#3289]). ### Removed - remove `tokio::prelude` ([#3299]). - io: remove `AsyncFd::with_poll()` ([#3306]). - net: remove `{Tcp,Unix}Stream::shutdown()` in favor of `AsyncWrite::shutdown()` ([#3298]). - stream: move all stream utilities to `tokio-stream` until `Stream` is added to `std` ([#3277]). - sync: mpsc `try_recv()` due to unexpected behavior ([#3263]). - tracing: make unstable as `tracing-core` is not 1.0 yet ([#3266]). ### Added - fs: `poll_*` fns to `DirEntry` ([#3308]). - io: `poll_*` fns to `io::Lines`, `io::Split` ([#3308]). - io: `_mut` method variants to `AsyncFd` ([#3304]). - net: `poll_*` fns to `UnixDatagram` ([#3223]). - net: `UnixStream` readiness and non-blocking ops ([#3246]). - sync: `UnboundedReceiver::blocking_recv()` ([#3262]). - sync: `watch::Sender::borrow()` ([#3269]). - sync: `Semaphore::close()` ([#3065]). - sync: `poll_recv` fns to `mpsc::Receiver`, `mpsc::UnboundedReceiver` ([#3308]). - time: `poll_tick` fn to `time::Interval` ([#3316]). [#3065]: https://github.com/tokio-rs/tokio/pull/3065 [#3125]: https://github.com/tokio-rs/tokio/pull/3125 [#3223]: https://github.com/tokio-rs/tokio/pull/3223 [#3234]: https://github.com/tokio-rs/tokio/pull/3234 [#3246]: https://github.com/tokio-rs/tokio/pull/3246 [#3259]: https://github.com/tokio-rs/tokio/pull/3259 [#3262]: https://github.com/tokio-rs/tokio/pull/3262 [#3263]: https://github.com/tokio-rs/tokio/pull/3263 [#3264]: https://github.com/tokio-rs/tokio/pull/3264 [#3266]: https://github.com/tokio-rs/tokio/pull/3266 [#3269]: https://github.com/tokio-rs/tokio/pull/3269 [#3277]: https://github.com/tokio-rs/tokio/pull/3277 [#3278]: https://github.com/tokio-rs/tokio/pull/3278 [#3287]: https://github.com/tokio-rs/tokio/pull/3287 [#3289]: https://github.com/tokio-rs/tokio/pull/3289 [#3298]: https://github.com/tokio-rs/tokio/pull/3298 [#3299]: https://github.com/tokio-rs/tokio/pull/3299 [#3304]: https://github.com/tokio-rs/tokio/pull/3304 [#3306]: https://github.com/tokio-rs/tokio/pull/3306 [#3308]: https://github.com/tokio-rs/tokio/pull/3308 [#3316]: https://github.com/tokio-rs/tokio/pull/3316 # 0.3.6 (December 14, 2020) ### Fixed - rt: fix deadlock in shutdown ([#3228]) - rt: fix panic in task abort when off rt ([#3159]) - sync: make `add_permits` panic with usize::MAX >> 3 permits ([#3188]) - time: Fix race condition in timer drop ([#3229]) - watch: fix spurious wakeup ([#3244]) ### Added - example: add back udp-codec example ([#3205]) - net: add `TcpStream::into_std` ([#3189]) [#3159]: https://github.com/tokio-rs/tokio/pull/3159 [#3188]: https://github.com/tokio-rs/tokio/pull/3188 [#3189]: https://github.com/tokio-rs/tokio/pull/3189 [#3205]: https://github.com/tokio-rs/tokio/pull/3205 [#3228]: https://github.com/tokio-rs/tokio/pull/3228 [#3229]: https://github.com/tokio-rs/tokio/pull/3229 [#3244]: https://github.com/tokio-rs/tokio/pull/3244 # 0.3.5 (November 30, 2020) ### Fixed - rt: fix `shutdown_timeout(0)` ([#3196]). - time: fixed race condition with small sleeps ([#3069]). ### Added - io: `AsyncFd::with_interest()` ([#3167]). - signal: `CtrlC` stream on windows ([#3186]). [#3069]: https://github.com/tokio-rs/tokio/pull/3069 [#3167]: https://github.com/tokio-rs/tokio/pull/3167 [#3186]: https://github.com/tokio-rs/tokio/pull/3186 [#3196]: https://github.com/tokio-rs/tokio/pull/3196 # 0.3.4 (November 18, 2020) ### Fixed - stream: `StreamMap` `Default` impl bound ([#3093]). - io: `AsyncFd::into_inner()` should deregister the FD ([#3104]). ### Changed - meta: `parking_lot` feature enabled with `full` ([#3119]). ### Added - io: `AsyncWrite` vectored writes ([#3149]). - net: TCP/UDP readiness and non-blocking ops ([#3130], [#2743], [#3138]). - net: TCP socket option (linger, send/recv buf size) ([#3145], [#3143]). - net: PID field in `UCred` with solaris/illumos ([#3085]). - rt: `runtime::Handle` allows spawning onto a runtime ([#3079]). - sync: `Notify::notify_waiters()` ([#3098]). - sync: `acquire_many()`, `try_acquire_many()` to `Semaphore` ([#3067]). [#2743]: https://github.com/tokio-rs/tokio/pull/2743 [#3067]: https://github.com/tokio-rs/tokio/pull/3067 [#3079]: https://github.com/tokio-rs/tokio/pull/3079 [#3085]: https://github.com/tokio-rs/tokio/pull/3085 [#3093]: https://github.com/tokio-rs/tokio/pull/3093 [#3098]: https://github.com/tokio-rs/tokio/pull/3098 [#3104]: https://github.com/tokio-rs/tokio/pull/3104 [#3119]: https://github.com/tokio-rs/tokio/pull/3119 [#3130]: https://github.com/tokio-rs/tokio/pull/3130 [#3138]: https://github.com/tokio-rs/tokio/pull/3138 [#3143]: https://github.com/tokio-rs/tokio/pull/3143 [#3145]: https://github.com/tokio-rs/tokio/pull/3145 [#3149]: https://github.com/tokio-rs/tokio/pull/3149 # 0.3.3 (November 2, 2020) Fixes a soundness hole by adding a missing `Send` bound to `Runtime::spawn_blocking()`. ### Fixed - rt: include missing `Send`, fixing soundness hole ([#3089]). - tracing: avoid huge trace span names ([#3074]). ### Added - net: `TcpSocket::reuseport()`, `TcpSocket::set_reuseport()` ([#3083]). - net: `TcpSocket::reuseaddr()` ([#3093]). - net: `TcpSocket::local_addr()` ([#3093]). - net: add pid to `UCred` ([#2633]). [#2633]: https://github.com/tokio-rs/tokio/pull/2633 [#3074]: https://github.com/tokio-rs/tokio/pull/3074 [#3083]: https://github.com/tokio-rs/tokio/pull/3083 [#3089]: https://github.com/tokio-rs/tokio/pull/3089 [#3093]: https://github.com/tokio-rs/tokio/pull/3093 # 0.3.2 (October 27, 2020) Adds `AsyncFd` as a replacement for v0.2's `PollEvented`. ### Fixed - io: fix a potential deadlock when shutting down the I/O driver ([#2903]). - sync: `RwLockWriteGuard::downgrade()` bug ([#2957]). ### Added - io: `AsyncFd` for receiving readiness events on raw FDs ([#2903]). - net: `poll_*` function on `UdpSocket` ([#2981]). - net: `UdpSocket::take_error()` ([#3051]). - sync: `oneshot::Sender::poll_closed()` ([#3032]). [#2903]: https://github.com/tokio-rs/tokio/pull/2903 [#2957]: https://github.com/tokio-rs/tokio/pull/2957 [#2981]: https://github.com/tokio-rs/tokio/pull/2981 [#3032]: https://github.com/tokio-rs/tokio/pull/3032 [#3051]: https://github.com/tokio-rs/tokio/pull/3051 # 0.3.1 (October 21, 2020) This release fixes an use-after-free in the IO driver. Additionally, the `read_buf` and `write_buf` methods have been added back to the IO traits, as the bytes crate is now on track to reach version 1.0 together with Tokio. ### Fixed - net: fix use-after-free ([#3019]). - fs: ensure buffered data is written on shutdown ([#3009]). ### Added - io: `copy_buf()` ([#2884]). - io: `AsyncReadExt::read_buf()`, `AsyncReadExt::write_buf()` for working with `Buf`/`BufMut` ([#3003]). - rt: `Runtime::spawn_blocking()` ([#2980]). - sync: `watch::Sender::is_closed()` ([#2991]). [#2884]: https://github.com/tokio-rs/tokio/pull/2884 [#2980]: https://github.com/tokio-rs/tokio/pull/2980 [#2991]: https://github.com/tokio-rs/tokio/pull/2991 [#3003]: https://github.com/tokio-rs/tokio/pull/3003 [#3009]: https://github.com/tokio-rs/tokio/pull/3009 [#3019]: https://github.com/tokio-rs/tokio/pull/3019 # 0.3.0 (October 15, 2020) This represents a 1.0 beta release. APIs are polished and future-proofed. APIs not included for 1.0 stabilization have been removed. Biggest changes are: - I/O driver internal rewrite. The windows implementation includes significant changes. - Runtime API is polished, especially with how it interacts with feature flag combinations. - Feature flags are simplified - `rt-core` and `rt-util` are combined to `rt` - `rt-threaded` is renamed to `rt-multi-thread` to match builder API - `tcp`, `udp`, `uds`, `dns` are combined to `net`. - `parking_lot` is included with `full` ### Changes - meta: Minimum supported Rust version is now 1.45. - io: `AsyncRead` trait now takes `ReadBuf` in order to safely handle reading into uninitialized memory ([#2758]). - io: Internal I/O driver storage is now able to compact ([#2757]). - rt: `Runtime::block_on` now takes `&self` ([#2782]). - sync: `watch` reworked to decouple receiving a change notification from receiving the value ([#2814], [#2806]). - sync: `Notify::notify` is renamed to `notify_one` ([#2822]). - process: `Child::kill` is now an `async fn` that cleans zombies ([#2823]). - sync: use `const fn` constructors as possible ([#2833], [#2790]) - signal: reduce cross-thread notification ([#2835]). - net: tcp,udp,uds types support operations with `&self` ([#2828], [#2919], [#2934]). - sync: blocking `mpsc` channel supports `send` with `&self` ([#2861]). - time: rename `delay_for` and `delay_until` to `sleep` and `sleep_until` ([#2826]). - io: upgrade to `mio` 0.7 ([#2893]). - io: `AsyncSeek` trait is tweaked ([#2885]). - fs: `File` operations take `&self` ([#2930]). - rt: runtime API, and `#[tokio::main]` macro polish ([#2876]) - rt: `Runtime::enter` uses an RAII guard instead of a closure ([#2954]). - net: the `from_std` function on all sockets no longer sets socket into non-blocking mode ([#2893]) ### Added - sync: `map` function to lock guards ([#2445]). - sync: `blocking_recv` and `blocking_send` fns to `mpsc` for use outside of Tokio ([#2685]). - rt: `Builder::thread_name_fn` for configuring thread names ([#1921]). - fs: impl `FromRawFd` and `FromRawHandle` for `File` ([#2792]). - process: `Child::wait` and `Child::try_wait` ([#2796]). - rt: support configuring thread keep-alive duration ([#2809]). - rt: `task::JoinHandle::abort` forcibly cancels a spawned task ([#2474]). - sync: `RwLock` write guard to read guard downgrading ([#2733]). - net: add `poll_*` functions that take `&self` to all net types ([#2845]) - sync: `get_mut()` for `Mutex`, `RwLock` ([#2856]). - sync: `mpsc::Sender::closed()` waits for `Receiver` half to close ([#2840]). - sync: `mpsc::Sender::is_closed()` returns true if `Receiver` half is closed ([#2726]). - stream: `iter` and `iter_mut` to `StreamMap` ([#2890]). - net: implement `AsRawSocket` on windows ([#2911]). - net: `TcpSocket` creates a socket without binding or listening ([#2920]). ### Removed - io: vectored ops are removed from `AsyncRead`, `AsyncWrite` traits ([#2882]). - io: `mio` is removed from the public API. `PollEvented` and` Registration` are removed ([#2893]). - io: remove `bytes` from public API. `Buf` and `BufMut` implementation are removed ([#2908]). - time: `DelayQueue` is moved to `tokio-util` ([#2897]). ### Fixed - io: `stdout` and `stderr` buffering on windows ([#2734]). [#1921]: https://github.com/tokio-rs/tokio/pull/1921 [#2445]: https://github.com/tokio-rs/tokio/pull/2445 [#2474]: https://github.com/tokio-rs/tokio/pull/2474 [#2685]: https://github.com/tokio-rs/tokio/pull/2685 [#2726]: https://github.com/tokio-rs/tokio/pull/2726 [#2733]: https://github.com/tokio-rs/tokio/pull/2733 [#2734]: https://github.com/tokio-rs/tokio/pull/2734 [#2757]: https://github.com/tokio-rs/tokio/pull/2757 [#2758]: https://github.com/tokio-rs/tokio/pull/2758 [#2782]: https://github.com/tokio-rs/tokio/pull/2782 [#2790]: https://github.com/tokio-rs/tokio/pull/2790 [#2792]: https://github.com/tokio-rs/tokio/pull/2792 [#2796]: https://github.com/tokio-rs/tokio/pull/2796 [#2806]: https://github.com/tokio-rs/tokio/pull/2806 [#2809]: https://github.com/tokio-rs/tokio/pull/2809 [#2814]: https://github.com/tokio-rs/tokio/pull/2814 [#2822]: https://github.com/tokio-rs/tokio/pull/2822 [#2823]: https://github.com/tokio-rs/tokio/pull/2823 [#2826]: https://github.com/tokio-rs/tokio/pull/2826 [#2828]: https://github.com/tokio-rs/tokio/pull/2828 [#2833]: https://github.com/tokio-rs/tokio/pull/2833 [#2835]: https://github.com/tokio-rs/tokio/pull/2835 [#2840]: https://github.com/tokio-rs/tokio/pull/2840 [#2845]: https://github.com/tokio-rs/tokio/pull/2845 [#2856]: https://github.com/tokio-rs/tokio/pull/2856 [#2861]: https://github.com/tokio-rs/tokio/pull/2861 [#2876]: https://github.com/tokio-rs/tokio/pull/2876 [#2882]: https://github.com/tokio-rs/tokio/pull/2882 [#2885]: https://github.com/tokio-rs/tokio/pull/2885 [#2890]: https://github.com/tokio-rs/tokio/pull/2890 [#2893]: https://github.com/tokio-rs/tokio/pull/2893 [#2897]: https://github.com/tokio-rs/tokio/pull/2897 [#2908]: https://github.com/tokio-rs/tokio/pull/2908 [#2911]: https://github.com/tokio-rs/tokio/pull/2911 [#2919]: https://github.com/tokio-rs/tokio/pull/2919 [#2920]: https://github.com/tokio-rs/tokio/pull/2920 [#2930]: https://github.com/tokio-rs/tokio/pull/2930 [#2934]: https://github.com/tokio-rs/tokio/pull/2934 [#2954]: https://github.com/tokio-rs/tokio/pull/2954 # 0.2.22 (July 21, 2020) ### Fixes - docs: misc improvements ([#2572], [#2658], [#2663], [#2656], [#2647], [#2630], [#2487], [#2621], [#2624], [#2600], [#2623], [#2622], [#2577], [#2569], [#2589], [#2575], [#2540], [#2564], [#2567], [#2520], [#2521], [#2493]) - rt: allow calls to `block_on` inside calls to `block_in_place` that are themselves inside `block_on` ([#2645]) - net: fix non-portable behavior when dropping `TcpStream` `OwnedWriteHalf` ([#2597]) - io: improve stack usage by allocating large buffers on directly on the heap ([#2634]) - io: fix unsound pin projection in `AsyncReadExt::read_buf` and `AsyncWriteExt::write_buf` ([#2612]) - io: fix unnecessary zeroing for `AsyncRead` implementors ([#2525]) - io: Fix `BufReader` not correctly forwarding `poll_write_buf` ([#2654]) - io: fix panic in `AsyncReadExt::read_line` ([#2541]) ### Changes - coop: returning `Poll::Pending` no longer decrements the task budget ([#2549]) ### Added - io: little-endian variants of `AsyncReadExt` and `AsyncWriteExt` methods ([#1915]) - task: add [`tracing`] instrumentation to spawned tasks ([#2655]) - sync: allow unsized types in `Mutex` and `RwLock` (via `default` constructors) ([#2615]) - net: add `ToSocketAddrs` implementation for `&[SocketAddr]` ([#2604]) - fs: add `OpenOptionsExt` for `OpenOptions` ([#2515]) - fs: add `DirBuilder` ([#2524]) [`tracing`]: https://crates.io/crates/tracing [#1915]: https://github.com/tokio-rs/tokio/pull/1915 [#2487]: https://github.com/tokio-rs/tokio/pull/2487 [#2493]: https://github.com/tokio-rs/tokio/pull/2493 [#2515]: https://github.com/tokio-rs/tokio/pull/2515 [#2520]: https://github.com/tokio-rs/tokio/pull/2520 [#2521]: https://github.com/tokio-rs/tokio/pull/2521 [#2524]: https://github.com/tokio-rs/tokio/pull/2524 [#2525]: https://github.com/tokio-rs/tokio/pull/2525 [#2540]: https://github.com/tokio-rs/tokio/pull/2540 [#2541]: https://github.com/tokio-rs/tokio/pull/2541 [#2549]: https://github.com/tokio-rs/tokio/pull/2549 [#2564]: https://github.com/tokio-rs/tokio/pull/2564 [#2567]: https://github.com/tokio-rs/tokio/pull/2567 [#2569]: https://github.com/tokio-rs/tokio/pull/2569 [#2572]: https://github.com/tokio-rs/tokio/pull/2572 [#2575]: https://github.com/tokio-rs/tokio/pull/2575 [#2577]: https://github.com/tokio-rs/tokio/pull/2577 [#2589]: https://github.com/tokio-rs/tokio/pull/2589 [#2597]: https://github.com/tokio-rs/tokio/pull/2597 [#2600]: https://github.com/tokio-rs/tokio/pull/2600 [#2604]: https://github.com/tokio-rs/tokio/pull/2604 [#2612]: https://github.com/tokio-rs/tokio/pull/2612 [#2615]: https://github.com/tokio-rs/tokio/pull/2615 [#2621]: https://github.com/tokio-rs/tokio/pull/2621 [#2622]: https://github.com/tokio-rs/tokio/pull/2622 [#2623]: https://github.com/tokio-rs/tokio/pull/2623 [#2624]: https://github.com/tokio-rs/tokio/pull/2624 [#2630]: https://github.com/tokio-rs/tokio/pull/2630 [#2634]: https://github.com/tokio-rs/tokio/pull/2634 [#2645]: https://github.com/tokio-rs/tokio/pull/2645 [#2647]: https://github.com/tokio-rs/tokio/pull/2647 [#2654]: https://github.com/tokio-rs/tokio/pull/2654 [#2655]: https://github.com/tokio-rs/tokio/pull/2655 [#2656]: https://github.com/tokio-rs/tokio/pull/2656 [#2658]: https://github.com/tokio-rs/tokio/pull/2658 [#2663]: https://github.com/tokio-rs/tokio/pull/2663 # 0.2.21 (May 13, 2020) ### Fixes - macros: disambiguate built-in `#[test]` attribute in macro expansion ([#2503]) - rt: `LocalSet` and task budgeting ([#2462]). - rt: task budgeting with `block_in_place` ([#2502]). - sync: release `broadcast` channel memory without sending a value ([#2509]). - time: notify when resetting a `Delay` to a time in the past ([#2290]) ### Added - io: `get_mut`, `get_ref`, and `into_inner` to `Lines` ([#2450]). - io: `mio::Ready` argument to `PollEvented` ([#2419]). - os: illumos support ([#2486]). - rt: `Handle::spawn_blocking` ([#2501]). - sync: `OwnedMutexGuard` for `Arc>` ([#2455]). [#2290]: https://github.com/tokio-rs/tokio/pull/2290 [#2419]: https://github.com/tokio-rs/tokio/pull/2419 [#2450]: https://github.com/tokio-rs/tokio/pull/2450 [#2455]: https://github.com/tokio-rs/tokio/pull/2455 [#2462]: https://github.com/tokio-rs/tokio/pull/2462 [#2486]: https://github.com/tokio-rs/tokio/pull/2486 [#2501]: https://github.com/tokio-rs/tokio/pull/2501 [#2502]: https://github.com/tokio-rs/tokio/pull/2502 [#2503]: https://github.com/tokio-rs/tokio/pull/2503 [#2509]: https://github.com/tokio-rs/tokio/pull/2509 # 0.2.20 (April 28, 2020) ### Fixes - sync: `broadcast` closing the channel no longer requires capacity ([#2448]). - rt: regression when configuring runtime with `max_threads` less than number of CPUs ([#2457]). [#2448]: https://github.com/tokio-rs/tokio/pull/2448 [#2457]: https://github.com/tokio-rs/tokio/pull/2457 # 0.2.19 (April 24, 2020) ### Fixes - docs: misc improvements ([#2400], [#2405], [#2414], [#2420], [#2423], [#2426], [#2427], [#2434], [#2436], [#2440]). - rt: support `block_in_place` in more contexts ([#2409], [#2410]). - stream: no panic in `merge()` and `chain()` when using `size_hint()` ([#2430]). - task: include visibility modifier when defining a task-local ([#2416]). ### Added - rt: `runtime::Handle::block_on` ([#2437]). - sync: owned `Semaphore` permit ([#2421]). - tcp: owned split ([#2270]). [#2270]: https://github.com/tokio-rs/tokio/pull/2270 [#2400]: https://github.com/tokio-rs/tokio/pull/2400 [#2405]: https://github.com/tokio-rs/tokio/pull/2405 [#2409]: https://github.com/tokio-rs/tokio/pull/2409 [#2410]: https://github.com/tokio-rs/tokio/pull/2410 [#2414]: https://github.com/tokio-rs/tokio/pull/2414 [#2416]: https://github.com/tokio-rs/tokio/pull/2416 [#2420]: https://github.com/tokio-rs/tokio/pull/2420 [#2421]: https://github.com/tokio-rs/tokio/pull/2421 [#2423]: https://github.com/tokio-rs/tokio/pull/2423 [#2426]: https://github.com/tokio-rs/tokio/pull/2426 [#2427]: https://github.com/tokio-rs/tokio/pull/2427 [#2430]: https://github.com/tokio-rs/tokio/pull/2430 [#2434]: https://github.com/tokio-rs/tokio/pull/2434 [#2436]: https://github.com/tokio-rs/tokio/pull/2436 [#2437]: https://github.com/tokio-rs/tokio/pull/2437 [#2440]: https://github.com/tokio-rs/tokio/pull/2440 # 0.2.18 (April 12, 2020) ### Fixes - task: `LocalSet` was incorrectly marked as `Send` ([#2398]) - io: correctly report `WriteZero` failure in `write_int` ([#2334]) [#2334]: https://github.com/tokio-rs/tokio/pull/2334 [#2398]: https://github.com/tokio-rs/tokio/pull/2398 # 0.2.17 (April 9, 2020) ### Fixes - rt: bug in work-stealing queue ([#2387]) ### Changes - rt: threadpool uses logical CPU count instead of physical by default ([#2391]) [#2387]: https://github.com/tokio-rs/tokio/pull/2387 [#2391]: https://github.com/tokio-rs/tokio/pull/2391 # 0.2.16 (April 3, 2020) ### Fixes - sync: fix a regression where `Mutex`, `Semaphore`, and `RwLock` futures no longer implement `Sync` ([#2375]) - fs: fix `fs::copy` not copying file permissions ([#2354]) ### Added - time: added `deadline` method to `delay_queue::Expired` ([#2300]) - io: added `StreamReader` ([#2052]) [#2052]: https://github.com/tokio-rs/tokio/pull/2052 [#2300]: https://github.com/tokio-rs/tokio/pull/2300 [#2354]: https://github.com/tokio-rs/tokio/pull/2354 [#2375]: https://github.com/tokio-rs/tokio/pull/2375 # 0.2.15 (April 2, 2020) ### Fixes - rt: fix queue regression ([#2362]). ### Added - sync: Add disarm to `mpsc::Sender` ([#2358]). [#2358]: https://github.com/tokio-rs/tokio/pull/2358 [#2362]: https://github.com/tokio-rs/tokio/pull/2362 # 0.2.14 (April 1, 2020) ### Fixes - rt: concurrency bug in scheduler ([#2273]). - rt: concurrency bug with shell runtime ([#2333]). - test-util: correct pause/resume of time ([#2253]). - time: `DelayQueue` correct wakeup after `insert` ([#2285]). ### Added - io: impl `RawFd`, `AsRawHandle` for std io types ([#2335]). - rt: automatic cooperative task yielding ([#2160], [#2343], [#2349]). - sync: `RwLock::into_inner` ([#2321]). ### Changed - sync: semaphore, mutex internals rewritten to avoid allocations ([#2325]). [#2160]: https://github.com/tokio-rs/tokio/pull/2160 [#2253]: https://github.com/tokio-rs/tokio/pull/2253 [#2273]: https://github.com/tokio-rs/tokio/pull/2273 [#2285]: https://github.com/tokio-rs/tokio/pull/2285 [#2321]: https://github.com/tokio-rs/tokio/pull/2321 [#2325]: https://github.com/tokio-rs/tokio/pull/2325 [#2333]: https://github.com/tokio-rs/tokio/pull/2333 [#2335]: https://github.com/tokio-rs/tokio/pull/2335 [#2343]: https://github.com/tokio-rs/tokio/pull/2343 [#2349]: https://github.com/tokio-rs/tokio/pull/2349 # 0.2.13 (February 28, 2020) ### Fixes - macros: unresolved import in `pin!` ([#2281]). [#2281]: https://github.com/tokio-rs/tokio/pull/2281 # 0.2.12 (February 27, 2020) ### Fixes - net: `UnixStream::poll_shutdown` should call `shutdown(Write)` ([#2245]). - process: Wake up read and write on `EPOLLERR` ([#2218]). - rt: potential deadlock when using `block_in_place` and shutting down the runtime ([#2119]). - rt: only detect number of CPUs if `core_threads` not specified ([#2238]). - sync: reduce `watch::Receiver` struct size ([#2191]). - time: succeed when setting delay of `$MAX-1` ([#2184]). - time: avoid having to poll `DelayQueue` after inserting new delay ([#2217]). ### Added - macros: `pin!` variant that assigns to identifier and pins ([#2274]). - net: impl `Stream` for `Listener` types ([#2275]). - rt: `Runtime::shutdown_timeout` waits for runtime to shutdown for specified duration ([#2186]). - stream: `StreamMap` merges streams and can insert / remove streams at runtime ([#2185]). - stream: `StreamExt::skip()` skips a fixed number of items ([#2204]). - stream: `StreamExt::skip_while()` skips items based on a predicate ([#2205]). - sync: `Notify` provides basic `async` / `await` task notification ([#2210]). - sync: `Mutex::into_inner` retrieves guarded data ([#2250]). - sync: `mpsc::Sender::send_timeout` sends, waiting for up to specified duration for channel capacity ([#2227]). - time: impl `Ord` and `Hash` for `Instant` ([#2239]). [#2119]: https://github.com/tokio-rs/tokio/pull/2119 [#2184]: https://github.com/tokio-rs/tokio/pull/2184 [#2185]: https://github.com/tokio-rs/tokio/pull/2185 [#2186]: https://github.com/tokio-rs/tokio/pull/2186 [#2191]: https://github.com/tokio-rs/tokio/pull/2191 [#2204]: https://github.com/tokio-rs/tokio/pull/2204 [#2205]: https://github.com/tokio-rs/tokio/pull/2205 [#2210]: https://github.com/tokio-rs/tokio/pull/2210 [#2217]: https://github.com/tokio-rs/tokio/pull/2217 [#2218]: https://github.com/tokio-rs/tokio/pull/2218 [#2227]: https://github.com/tokio-rs/tokio/pull/2227 [#2238]: https://github.com/tokio-rs/tokio/pull/2238 [#2239]: https://github.com/tokio-rs/tokio/pull/2239 [#2245]: https://github.com/tokio-rs/tokio/pull/2245 [#2250]: https://github.com/tokio-rs/tokio/pull/2250 [#2274]: https://github.com/tokio-rs/tokio/pull/2274 [#2275]: https://github.com/tokio-rs/tokio/pull/2275 # 0.2.11 (January 27, 2020) ### Fixes - docs: misc fixes and tweaks ([#2155], [#2103], [#2027], [#2167], [#2175]). - macros: handle generics in `#[tokio::main]` method ([#2177]). - sync: `broadcast` potential lost notifications ([#2135]). - rt: improve "no runtime" panic messages ([#2145]). ### Added - optional support for using `parking_lot` internally ([#2164]). - fs: `fs::copy`, an async version of `std::fs::copy` ([#2079]). - macros: `select!` waits for the first branch to complete ([#2152]). - macros: `join!` waits for all branches to complete ([#2158]). - macros: `try_join!` waits for all branches to complete or the first error ([#2169]). - macros: `pin!` pins a value to the stack ([#2163]). - net: `ReadHalf::poll()` and `ReadHalf::poll_peak` ([#2151]) - stream: `StreamExt::timeout()` sets a per-item max duration ([#2149]). - stream: `StreamExt::fold()` applies a function, producing a single value. ([#2122]). - sync: impl `Eq`, `PartialEq` for `oneshot::RecvError` ([#2168]). - task: methods for inspecting the `JoinError` cause ([#2051]). [#2027]: https://github.com/tokio-rs/tokio/pull/2027 [#2051]: https://github.com/tokio-rs/tokio/pull/2051 [#2079]: https://github.com/tokio-rs/tokio/pull/2079 [#2103]: https://github.com/tokio-rs/tokio/pull/2103 [#2122]: https://github.com/tokio-rs/tokio/pull/2122 [#2135]: https://github.com/tokio-rs/tokio/pull/2135 [#2145]: https://github.com/tokio-rs/tokio/pull/2145 [#2149]: https://github.com/tokio-rs/tokio/pull/2149 [#2151]: https://github.com/tokio-rs/tokio/pull/2151 [#2152]: https://github.com/tokio-rs/tokio/pull/2152 [#2155]: https://github.com/tokio-rs/tokio/pull/2155 [#2158]: https://github.com/tokio-rs/tokio/pull/2158 [#2163]: https://github.com/tokio-rs/tokio/pull/2163 [#2164]: https://github.com/tokio-rs/tokio/pull/2164 [#2167]: https://github.com/tokio-rs/tokio/pull/2167 [#2168]: https://github.com/tokio-rs/tokio/pull/2168 [#2169]: https://github.com/tokio-rs/tokio/pull/2169 [#2175]: https://github.com/tokio-rs/tokio/pull/2175 [#2177]: https://github.com/tokio-rs/tokio/pull/2177 # 0.2.10 (January 21, 2020) ### Fixes - `#[tokio::main]` when `rt-core` feature flag is not enabled ([#2139]). - remove `AsyncBufRead` from `BufStream` impl block ([#2108]). - potential undefined behavior when implementing `AsyncRead` incorrectly ([#2030]). ### Added - `BufStream::with_capacity` ([#2125]). - impl `From` and `Default` for `RwLock` ([#2089]). - `io::ReadHalf::is_pair_of` checks if provided `WriteHalf` is for the same underlying object ([#1762], [#2144]). - `runtime::Handle::try_current()` returns a handle to the current runtime ([#2118]). - `stream::empty()` returns an immediately ready empty stream ([#2092]). - `stream::once(val)` returns a stream that yields a single value: `val` ([#2094]). - `stream::pending()` returns a stream that never becomes ready ([#2092]). - `StreamExt::chain()` sequences a second stream after the first completes ([#2093]). - `StreamExt::collect()` transform a stream into a collection ([#2109]). - `StreamExt::fuse` ends the stream after the first `None` ([#2085]). - `StreamExt::merge` combines two streams, yielding values as they become ready ([#2091]). - Task-local storage ([#2126]). [#1762]: https://github.com/tokio-rs/tokio/pull/1762 [#2030]: https://github.com/tokio-rs/tokio/pull/2030 [#2085]: https://github.com/tokio-rs/tokio/pull/2085 [#2089]: https://github.com/tokio-rs/tokio/pull/2089 [#2091]: https://github.com/tokio-rs/tokio/pull/2091 [#2092]: https://github.com/tokio-rs/tokio/pull/2092 [#2093]: https://github.com/tokio-rs/tokio/pull/2093 [#2094]: https://github.com/tokio-rs/tokio/pull/2094 [#2108]: https://github.com/tokio-rs/tokio/pull/2108 [#2109]: https://github.com/tokio-rs/tokio/pull/2109 [#2118]: https://github.com/tokio-rs/tokio/pull/2118 [#2125]: https://github.com/tokio-rs/tokio/pull/2125 [#2126]: https://github.com/tokio-rs/tokio/pull/2126 [#2139]: https://github.com/tokio-rs/tokio/pull/2139 [#2144]: https://github.com/tokio-rs/tokio/pull/2144 # 0.2.9 (January 9, 2020) ### Fixes - `AsyncSeek` impl for `File` ([#1986]). - rt: shutdown deadlock in `threaded_scheduler` ([#2074], [#2082]). - rt: memory ordering when dropping `JoinHandle` ([#2044]). - docs: misc API documentation fixes and improvements. [#1986]: https://github.com/tokio-rs/tokio/pull/1986 [#2044]: https://github.com/tokio-rs/tokio/pull/2044 [#2074]: https://github.com/tokio-rs/tokio/pull/2074 [#2082]: https://github.com/tokio-rs/tokio/pull/2082 # 0.2.8 (January 7, 2020) ### Fixes - depend on new version of `tokio-macros`. # 0.2.7 (January 7, 2020) ### Fixes - potential deadlock when dropping `basic_scheduler` Runtime. - calling `spawn_blocking` from within a `spawn_blocking` ([#2006]). - storing a `Runtime` instance in a thread-local ([#2011]). - miscellaneous documentation fixes. - rt: fix `Waker::will_wake` to return true when tasks match ([#2045]). - test-util: `time::advance` runs pending tasks before changing the time ([#2059]). ### Added - `net::lookup_host` maps a `T: ToSocketAddrs` to a stream of `SocketAddrs` ([#1870]). - `process::Child` fields are made public to match `std` ([#2014]). - impl `Stream` for `sync::broadcast::Receiver` ([#2012]). - `sync::RwLock` provides an asynchronous read-write lock ([#1699]). - `runtime::Handle::current` returns the handle for the current runtime ([#2040]). - `StreamExt::filter` filters stream values according to a predicate ([#2001]). - `StreamExt::filter_map` simultaneously filter and map stream values ([#2001]). - `StreamExt::try_next` convenience for streams of `Result` ([#2005]). - `StreamExt::take` limits a stream to a specified number of values ([#2025]). - `StreamExt::take_while` limits a stream based on a predicate ([#2029]). - `StreamExt::all` tests if every element of the stream matches a predicate ([#2035]). - `StreamExt::any` tests if any element of the stream matches a predicate ([#2034]). - `task::LocalSet.await` runs spawned tasks until the set is idle ([#1971]). - `time::DelayQueue::len` returns the number entries in the queue ([#1755]). - expose runtime options from the `#[tokio::main]` and `#[tokio::test]` ([#2022]). [#1699]: https://github.com/tokio-rs/tokio/pull/1699 [#1755]: https://github.com/tokio-rs/tokio/pull/1755 [#1870]: https://github.com/tokio-rs/tokio/pull/1870 [#1971]: https://github.com/tokio-rs/tokio/pull/1971 [#2001]: https://github.com/tokio-rs/tokio/pull/2001 [#2005]: https://github.com/tokio-rs/tokio/pull/2005 [#2006]: https://github.com/tokio-rs/tokio/pull/2006 [#2011]: https://github.com/tokio-rs/tokio/pull/2011 [#2012]: https://github.com/tokio-rs/tokio/pull/2012 [#2014]: https://github.com/tokio-rs/tokio/pull/2014 [#2022]: https://github.com/tokio-rs/tokio/pull/2022 [#2025]: https://github.com/tokio-rs/tokio/pull/2025 [#2029]: https://github.com/tokio-rs/tokio/pull/2029 [#2034]: https://github.com/tokio-rs/tokio/pull/2034 [#2035]: https://github.com/tokio-rs/tokio/pull/2035 [#2040]: https://github.com/tokio-rs/tokio/pull/2040 [#2045]: https://github.com/tokio-rs/tokio/pull/2045 [#2059]: https://github.com/tokio-rs/tokio/pull/2059 # 0.2.6 (December 19, 2019) ### Fixes - `fs::File::seek` API regression ([#1991]). [#1991]: https://github.com/tokio-rs/tokio/pull/1991 # 0.2.5 (December 18, 2019) ### Added - `io::AsyncSeek` trait ([#1924]). - `Mutex::try_lock` ([#1939]) - `mpsc::Receiver::try_recv` and `mpsc::UnboundedReceiver::try_recv` ([#1939]). - `writev` support for `TcpStream` ([#1956]). - `time::throttle` for throttling streams ([#1949]). - implement `Stream` for `time::DelayQueue` ([#1975]). - `sync::broadcast` provides a fan-out channel ([#1943]). - `sync::Semaphore` provides an async semaphore ([#1973]). - `stream::StreamExt` provides stream utilities ([#1962]). ### Fixes - deadlock risk while shutting down the runtime ([#1972]). - panic while shutting down the runtime ([#1978]). - `sync::MutexGuard` debug output ([#1961]). - misc doc improvements ([#1933], [#1934], [#1940], [#1942]). ### Changes - runtime threads are configured with `runtime::Builder::core_threads` and `runtime::Builder::max_threads`. `runtime::Builder::num_threads` is deprecated ([#1977]). [#1924]: https://github.com/tokio-rs/tokio/pull/1924 [#1933]: https://github.com/tokio-rs/tokio/pull/1933 [#1934]: https://github.com/tokio-rs/tokio/pull/1934 [#1939]: https://github.com/tokio-rs/tokio/pull/1939 [#1940]: https://github.com/tokio-rs/tokio/pull/1940 [#1942]: https://github.com/tokio-rs/tokio/pull/1942 [#1943]: https://github.com/tokio-rs/tokio/pull/1943 [#1949]: https://github.com/tokio-rs/tokio/pull/1949 [#1956]: https://github.com/tokio-rs/tokio/pull/1956 [#1961]: https://github.com/tokio-rs/tokio/pull/1961 [#1962]: https://github.com/tokio-rs/tokio/pull/1962 [#1972]: https://github.com/tokio-rs/tokio/pull/1972 [#1973]: https://github.com/tokio-rs/tokio/pull/1973 [#1975]: https://github.com/tokio-rs/tokio/pull/1975 [#1977]: https://github.com/tokio-rs/tokio/pull/1977 [#1978]: https://github.com/tokio-rs/tokio/pull/1978 # 0.2.4 (December 6, 2019) ### Fixes - `sync::Mutex` deadlock when `lock()` future is dropped early ([#1898]). [#1898]: https://github.com/tokio-rs/tokio/pull/1898 # 0.2.3 (December 6, 2019) ### Added - read / write integers using `AsyncReadExt` and `AsyncWriteExt` ([#1863]). - `read_buf` / `write_buf` for reading / writing `Buf` / `BufMut` ([#1881]). - `TcpStream::poll_peek` - pollable API for performing TCP peek ([#1864]). - `sync::oneshot::error::TryRecvError` provides variants to detect the error kind ([#1874]). - `LocalSet::block_on` accepts `!'static` task ([#1882]). - `task::JoinError` is now `Sync` ([#1888]). - impl conversions between `tokio::time::Instant` and `std::time::Instant` ([#1904]). ### Fixes - calling `spawn_blocking` after runtime shutdown ([#1875]). - `LocalSet` drop infinite loop ([#1892]). - `LocalSet` hang under load ([#1905]). - improved documentation ([#1865], [#1866], [#1868], [#1874], [#1876], [#1911]). [#1863]: https://github.com/tokio-rs/tokio/pull/1863 [#1864]: https://github.com/tokio-rs/tokio/pull/1864 [#1865]: https://github.com/tokio-rs/tokio/pull/1865 [#1866]: https://github.com/tokio-rs/tokio/pull/1866 [#1868]: https://github.com/tokio-rs/tokio/pull/1868 [#1874]: https://github.com/tokio-rs/tokio/pull/1874 [#1875]: https://github.com/tokio-rs/tokio/pull/1875 [#1876]: https://github.com/tokio-rs/tokio/pull/1876 [#1881]: https://github.com/tokio-rs/tokio/pull/1881 [#1882]: https://github.com/tokio-rs/tokio/pull/1882 [#1888]: https://github.com/tokio-rs/tokio/pull/1888 [#1892]: https://github.com/tokio-rs/tokio/pull/1892 [#1904]: https://github.com/tokio-rs/tokio/pull/1904 [#1905]: https://github.com/tokio-rs/tokio/pull/1905 [#1911]: https://github.com/tokio-rs/tokio/pull/1911 # 0.2.2 (November 29, 2019) ### Fixes - scheduling with `basic_scheduler` ([#1861]). - update `spawn` panic message to specify that a task scheduler is required ([#1839]). - API docs example for `runtime::Builder` to include a task scheduler ([#1841]). - general documentation ([#1834]). - building on illumos/solaris ([#1772]). - panic when dropping `LocalSet` ([#1843]). - API docs mention the required Cargo features for `Builder::{basic, threaded}_scheduler` ([#1858]). ### Added - impl `Stream` for `signal::unix::Signal` ([#1849]). - API docs for platform specific behavior of `signal::ctrl_c` and `signal::unix::Signal` ([#1854]). - API docs for `signal::unix::Signal::{recv, poll_recv}` and `signal::windows::CtrlBreak::{recv, poll_recv}` ([#1854]). - `File::into_std` and `File::try_into_std` methods ([#1856]). [#1772]: https://github.com/tokio-rs/tokio/pull/1772 [#1834]: https://github.com/tokio-rs/tokio/pull/1834 [#1839]: https://github.com/tokio-rs/tokio/pull/1839 [#1841]: https://github.com/tokio-rs/tokio/pull/1841 [#1843]: https://github.com/tokio-rs/tokio/pull/1843 [#1849]: https://github.com/tokio-rs/tokio/pull/1849 [#1854]: https://github.com/tokio-rs/tokio/pull/1854 [#1856]: https://github.com/tokio-rs/tokio/pull/1856 [#1858]: https://github.com/tokio-rs/tokio/pull/1858 [#1861]: https://github.com/tokio-rs/tokio/pull/1861 # 0.2.1 (November 26, 2019) ### Fixes - API docs for `TcpListener::incoming`, `UnixListener::incoming` ([#1831]). ### Added - `tokio::task::LocalSet` provides a strategy for spawning `!Send` tasks ([#1733]). - export `tokio::time::Elapsed` ([#1826]). - impl `AsRawFd`, `AsRawHandle` for `tokio::fs::File` ([#1827]). [#1733]: https://github.com/tokio-rs/tokio/pull/1733 [#1826]: https://github.com/tokio-rs/tokio/pull/1826 [#1827]: https://github.com/tokio-rs/tokio/pull/1827 [#1831]: https://github.com/tokio-rs/tokio/pull/1831 # 0.2.0 (November 26, 2019) A major breaking change. Most implementation and APIs have changed one way or another. This changelog entry contains a highlight ### Changed - APIs are updated to use `async / await`. - most `tokio-*` crates are collapsed into this crate. - Scheduler is rewritten. - `tokio::spawn` returns a `JoinHandle`. - A single I/O / timer is used per runtime. - I/O driver uses a concurrent slab for allocating state. - components are made available via feature flag. - Use `bytes` 0.5 - `tokio::codec` is moved to `tokio-util`. ### Removed - Standalone `timer` and `net` drivers are removed, use `Runtime` instead - `current_thread` runtime is removed, use `tokio::runtime::Runtime` with `basic_scheduler` instead. # 0.1.21 (May 30, 2019) ### Changed - Bump `tokio-trace-core` version to 0.2 ([#1111]). [#1111]: https://github.com/tokio-rs/tokio/pull/1111 # 0.1.20 (May 14, 2019) ### Added - `tokio::runtime::Builder::panic_handler` allows configuring handling panics on the runtime ([#1055]). [#1055]: https://github.com/tokio-rs/tokio/pull/1055 # 0.1.19 (April 22, 2019) ### Added - Re-export `tokio::sync::Mutex` primitive ([#964]). [#964]: https://github.com/tokio-rs/tokio/pull/964 # 0.1.18 (March 22, 2019) ### Added - `TypedExecutor` re-export and implementations ([#993]). [#993]: https://github.com/tokio-rs/tokio/pull/993 # 0.1.17 (March 13, 2019) ### Added - Propagate trace subscriber in the runtime ([#966]). [#966]: https://github.com/tokio-rs/tokio/pull/966 # 0.1.16 (March 1, 2019) ### Fixed - async-await: track latest nightly changes ([#940]). ### Added - `sync::Watch`, a single value broadcast channel ([#922]). - Async equivalent of read / write file helpers being added to `std` ([#896]). [#896]: https://github.com/tokio-rs/tokio/pull/896 [#922]: https://github.com/tokio-rs/tokio/pull/922 [#940]: https://github.com/tokio-rs/tokio/pull/940 # 0.1.15 (January 24, 2019) ### Added - Re-export tokio-sync APIs ([#839]). - Stream enumerate combinator ([#832]). [#832]: https://github.com/tokio-rs/tokio/pull/832 [#839]: https://github.com/tokio-rs/tokio/pull/839 # 0.1.14 (January 6, 2019) - Use feature flags to break up the crate, allowing users to pick & choose components ([#808]). - Export `UnixDatagram` and `UnixDatagramFramed` ([#772]). [#772]: https://github.com/tokio-rs/tokio/pull/772 [#808]: https://github.com/tokio-rs/tokio/pull/808 # 0.1.13 (November 21, 2018) - Fix `Runtime::reactor()` when no tasks are spawned ([#721]). - `runtime::Builder` no longer uses deprecated methods ([#749]). - Provide `after_start` and `before_stop` configuration settings for `Runtime` ([#756]). - Implement throttle stream combinator ([#736]). [#721]: https://github.com/tokio-rs/tokio/pull/721 [#736]: https://github.com/tokio-rs/tokio/pull/736 [#749]: https://github.com/tokio-rs/tokio/pull/749 [#756]: https://github.com/tokio-rs/tokio/pull/756 # 0.1.12 (October 23, 2018) - runtime: expose `keep_alive` on runtime builder ([#676]). - runtime: create a reactor per worker thread ([#660]). - codec: fix panic in `LengthDelimitedCodec` ([#682]). - io: re-export `tokio_io::io::read` function ([#689]). - runtime: check for executor re-entry in more places ([#708]). [#660]: https://github.com/tokio-rs/tokio/pull/660 [#676]: https://github.com/tokio-rs/tokio/pull/676 [#682]: https://github.com/tokio-rs/tokio/pull/682 [#689]: https://github.com/tokio-rs/tokio/pull/689 [#708]: https://github.com/tokio-rs/tokio/pull/708 # 0.1.11 (September 28, 2018) - Fix `tokio-async-await` dependency ([#675]). [#675]: https://github.com/tokio-rs/tokio/pull/675 # 0.1.10 (September 27, 2018) - Fix minimal versions # 0.1.9 (September 27, 2018) - Experimental async/await improvements ([#661]). - Re-export `TaskExecutor` from `tokio-current-thread` ([#652]). - Improve `Runtime` builder API ([#645]). - `tokio::run` panics when called from the context of an executor ([#646]). - Introduce `StreamExt` with a `timeout` helper ([#573]). - Move `length_delimited` into `tokio` ([#575]). - Re-organize `tokio::net` module ([#548]). - Re-export `tokio-current-thread::spawn` in current_thread runtime ([#579]). [#548]: https://github.com/tokio-rs/tokio/pull/548 [#573]: https://github.com/tokio-rs/tokio/pull/573 [#575]: https://github.com/tokio-rs/tokio/pull/575 [#579]: https://github.com/tokio-rs/tokio/pull/579 [#645]: https://github.com/tokio-rs/tokio/pull/645 [#646]: https://github.com/tokio-rs/tokio/pull/646 [#652]: https://github.com/tokio-rs/tokio/pull/652 [#661]: https://github.com/tokio-rs/tokio/pull/661 # 0.1.8 (August 23, 2018) - Extract tokio::executor::current_thread to a sub crate ([#370]) - Add `Runtime::block_on` ([#398]) - Add `runtime::current_thread::block_on_all` ([#477]) - Misc documentation improvements ([#450]) - Implement `std::error::Error` for error types ([#501]) [#370]: https://github.com/tokio-rs/tokio/pull/370 [#398]: https://github.com/tokio-rs/tokio/pull/398 [#450]: https://github.com/tokio-rs/tokio/pull/450 [#477]: https://github.com/tokio-rs/tokio/pull/477 [#501]: https://github.com/tokio-rs/tokio/pull/501 # 0.1.7 (June 6, 2018) - Add `Runtime::block_on` for concurrent runtime ([#391]). - Provide handle to `current_thread::Runtime` that allows spawning tasks from other threads ([#340]). - Provide `clock::now()`, a configurable source of time ([#381]). [#340]: https://github.com/tokio-rs/tokio/pull/340 [#381]: https://github.com/tokio-rs/tokio/pull/381 [#391]: https://github.com/tokio-rs/tokio/pull/391 # 0.1.6 (May 2, 2018) - Add asynchronous filesystem APIs ([#323]). - Add "current thread" runtime variant ([#308]). - `CurrentThread`: Expose inner `Park` instance. - Improve fairness of `CurrentThread` executor ([#313]). [#308]: https://github.com/tokio-rs/tokio/pull/308 [#313]: https://github.com/tokio-rs/tokio/pull/313 [#323]: https://github.com/tokio-rs/tokio/pull/323 # 0.1.5 (March 30, 2018) - Provide timer API ([#266]) [#266]: https://github.com/tokio-rs/tokio/pull/266 # 0.1.4 (March 22, 2018) - Fix build on FreeBSD ([#218]) - Shutdown the Runtime when the handle is dropped ([#214]) - Set Runtime thread name prefix for worker threads ([#232]) - Add builder for Runtime ([#234]) - Extract TCP and UDP types into separate crates ([#224]) - Optionally support futures 0.2. [#214]: https://github.com/tokio-rs/tokio/pull/214 [#218]: https://github.com/tokio-rs/tokio/pull/218 [#224]: https://github.com/tokio-rs/tokio/pull/224 [#232]: https://github.com/tokio-rs/tokio/pull/232 [#234]: https://github.com/tokio-rs/tokio/pull/234 # 0.1.3 (March 09, 2018) - Fix `CurrentThread::turn` to block on idle ([#212]). [#212]: https://github.com/tokio-rs/tokio/pull/212 # 0.1.2 (March 09, 2018) - Introduce Tokio Runtime ([#141]) - Provide `CurrentThread` for more flexible usage of current thread executor ([#141]). - Add Lio for platforms that support it ([#142]). - I/O resources now lazily bind to the reactor ([#160]). - Extract Reactor to dedicated crate ([#169]) - Add facade to sub crates and add prelude ([#166]). - Switch TCP/UDP fns to poll\_ -> Poll<...> style ([#175]) [#141]: https://github.com/tokio-rs/tokio/pull/141 [#142]: https://github.com/tokio-rs/tokio/pull/142 [#160]: https://github.com/tokio-rs/tokio/pull/160 [#166]: https://github.com/tokio-rs/tokio/pull/166 [#169]: https://github.com/tokio-rs/tokio/pull/169 [#175]: https://github.com/tokio-rs/tokio/pull/175 # 0.1.1 (February 09, 2018) - Doc fixes # 0.1.0 (February 07, 2018) - Initial crate released based on [RFC](https://github.com/tokio-rs/tokio-rfcs/pull/3). tokio-1.43.1/Cargo.lock0000644000001160410000000000100102250ustar # This file is automatically @generated by Cargo. # It is not intended for manual editing. version = 3 [[package]] name = "addr2line" version = "0.24.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfbe277e56a376000877090da837660b4427aad530e3028d44e0bffe4f89a1c1" dependencies = [ "gimli", ] [[package]] name = "adler2" version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" [[package]] name = "aho-corasick" version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" dependencies = [ "memchr", ] [[package]] name = "async-stream" version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" dependencies = [ "async-stream-impl", "futures-core", "pin-project-lite", ] [[package]] name = "async-stream-impl" version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "autocfg" version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "backtrace" version = "0.3.74" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a" dependencies = [ "addr2line", "cfg-if", "libc", "miniz_oxide", "object", "rustc-demangle", "windows-targets 0.52.6", ] [[package]] name = "bit-set" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" dependencies = [ "bit-vec", ] [[package]] name = "bit-vec" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" [[package]] name = "bitflags" version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd" [[package]] name = "bumpalo" version = "3.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1628fb46dfa0b37568d12e5edd512553eccf6a22a78e8bde00bb4aed84d5bdbf" [[package]] name = "bytes" version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" [[package]] name = "cc" version = "1.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fcb57c740ae1daf453ae85f16e37396f672b039e00d9d866e07ddb24e328e3a" dependencies = [ "shlex", ] [[package]] name = "cfg-if" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "cfg_aliases" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "difflib" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6184e33543162437515c2e2b48714794e37845ec9851711914eec9d308f6ebe8" [[package]] name = "downcast" version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1435fa1053d8b2fbbe9be7e97eca7f33d37b28409959813daefc1446a14247f1" [[package]] name = "either" version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" [[package]] name = "errno" version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" dependencies = [ "libc", "windows-sys 0.59.0", ] [[package]] name = "fastrand" version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "float-cmp" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "98de4bbd547a563b716d8dfa9aad1cb19bfab00f4fa09a6a4ed21dbcf44ce9c4" dependencies = [ "num-traits", ] [[package]] name = "fnv" version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" [[package]] name = "fragile" version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28dd6caf6059519a65843af8fe2a3ae298b14b80179855aeb4adc2c1934ee619" [[package]] name = "futures" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" dependencies = [ "futures-channel", "futures-core", "futures-executor", "futures-io", "futures-sink", "futures-task", "futures-util", ] [[package]] name = "futures-channel" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" dependencies = [ "futures-core", "futures-sink", ] [[package]] name = "futures-core" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" [[package]] name = "futures-executor" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" dependencies = [ "futures-core", "futures-task", "futures-util", ] [[package]] name = "futures-io" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" [[package]] name = "futures-macro" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "futures-sink" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" [[package]] name = "futures-task" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" [[package]] name = "futures-util" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ "futures-channel", "futures-core", "futures-io", "futures-macro", "futures-sink", "futures-task", "memchr", "pin-project-lite", "pin-utils", "slab", ] [[package]] name = "generator" version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cc6bd114ceda131d3b1d665eba35788690ad37f5916457286b32ab6fd3c438dd" dependencies = [ "cfg-if", "libc", "log", "rustversion", "windows", ] [[package]] name = "getrandom" version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", "libc", "wasi 0.11.0+wasi-snapshot-preview1", ] [[package]] name = "getrandom" version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73fea8450eea4bac3940448fb7ae50d91f034f941199fcd9d909a5a07aa455f0" dependencies = [ "cfg-if", "libc", "r-efi", "wasi 0.14.2+wasi-0.2.4", ] [[package]] name = "gimli" version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" [[package]] name = "itertools" version = "0.10.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" dependencies = [ "either", ] [[package]] name = "itoa" version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "js-sys" version = "0.3.77" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" dependencies = [ "once_cell", "wasm-bindgen", ] [[package]] name = "lazy_static" version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" version = "0.2.171" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" [[package]] name = "linux-raw-sys" version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe7db12097d22ec582439daf8618b8fdd1a7bef6270e9af3b1ebcd30893cf413" [[package]] name = "lock_api" version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" dependencies = [ "autocfg", "scopeguard", ] [[package]] name = "log" version = "0.4.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" [[package]] name = "loom" version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" dependencies = [ "cfg-if", "generator", "pin-utils", "scoped-tls", "serde", "serde_json", "tracing", "tracing-subscriber", ] [[package]] name = "matchers" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" dependencies = [ "regex-automata 0.1.10", ] [[package]] name = "memchr" version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] name = "memoffset" version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" dependencies = [ "autocfg", ] [[package]] name = "minicov" version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f27fe9f1cc3c22e1687f9446c2083c4c5fc7f0bcf1c7a86bdbded14985895b4b" dependencies = [ "cc", "walkdir", ] [[package]] name = "miniz_oxide" version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e3e04debbb59698c15bacbb6d93584a8c0ca9cc3213cb423d31f760d8843ce5" dependencies = [ "adler2", ] [[package]] name = "mio" version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4a650543ca06a924e8b371db273b2756685faae30f8487da1b56505a8f78b0c" dependencies = [ "libc", "log", "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys 0.48.0", ] [[package]] name = "mio" version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd" dependencies = [ "libc", "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys 0.52.0", ] [[package]] name = "mio-aio" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94970d564990c83f7431efebf40f81b4b0d36b4e362fae869081d952530937ff" dependencies = [ "mio 0.8.11", "nix", "pin-utils", ] [[package]] name = "mockall" version = "0.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4c84490118f2ee2d74570d114f3d0493cbf02790df303d2707606c3e14e07c96" dependencies = [ "cfg-if", "downcast", "fragile", "lazy_static", "mockall_derive", "predicates", "predicates-tree", ] [[package]] name = "mockall_derive" version = "0.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22ce75669015c4f47b289fd4d4f56e894e4c96003ffdf3ac51313126f94c6cbb" dependencies = [ "cfg-if", "proc-macro2", "quote", "syn 1.0.109", ] [[package]] name = "nix" version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" dependencies = [ "bitflags", "cfg-if", "cfg_aliases", "libc", "memoffset", "pin-utils", ] [[package]] name = "normalize-line-endings" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61807f77802ff30975e01f4f071c8ba10c022052f98b3294119f3e615d13e5be" [[package]] name = "nu-ansi-term" version = "0.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" dependencies = [ "overload", "winapi", ] [[package]] name = "num-traits" version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", ] [[package]] name = "object" version = "0.36.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" dependencies = [ "memchr", ] [[package]] name = "once_cell" version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "overload" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" [[package]] name = "parking_lot" version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" dependencies = [ "lock_api", "parking_lot_core", ] [[package]] name = "parking_lot_core" version = "0.9.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", "windows-targets 0.52.6", ] [[package]] name = "pin-project-lite" version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" [[package]] name = "pin-utils" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "ppv-lite86" version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" dependencies = [ "zerocopy", ] [[package]] name = "predicates" version = "2.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59230a63c37f3e18569bdb90e4a89cbf5bf8b06fea0b84e65ea10cc4df47addd" dependencies = [ "difflib", "float-cmp", "itertools", "normalize-line-endings", "predicates-core", "regex", ] [[package]] name = "predicates-core" version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "727e462b119fe9c93fd0eb1429a5f7647394014cf3c04ab2c0350eeb09095ffa" [[package]] name = "predicates-tree" version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72dd2d6d381dfb73a193c7fca536518d7caee39fc8503f74e7dc0be0531b425c" dependencies = [ "predicates-core", "termtree", ] [[package]] name = "proc-macro2" version = "1.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a31971752e70b8b2686d7e46ec17fb38dad4051d94024c88df49b667caea9c84" dependencies = [ "unicode-ident", ] [[package]] name = "proptest" version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "14cae93065090804185d3b75f0bf93b8eeda30c7a9b4a33d3bdb3988d6229e50" dependencies = [ "bit-set", "bit-vec", "bitflags", "lazy_static", "num-traits", "rand", "rand_chacha", "rand_xorshift", "regex-syntax 0.8.5", "rusty-fork", "tempfile", "unarray", ] [[package]] name = "quick-error" version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" [[package]] name = "quote" version = "1.0.40" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" dependencies = [ "proc-macro2", ] [[package]] name = "r-efi" version = "5.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5" [[package]] name = "rand" version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", "rand_chacha", "rand_core", ] [[package]] name = "rand_chacha" version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", "rand_core", ] [[package]] name = "rand_core" version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ "getrandom 0.2.15", ] [[package]] name = "rand_xorshift" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f" dependencies = [ "rand_core", ] [[package]] name = "redox_syscall" version = "0.5.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b8c0c260b63a8219631167be35e6a988e9554dbd323f8bd08439c8ed1302bd1" dependencies = [ "bitflags", ] [[package]] name = "regex" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" dependencies = [ "aho-corasick", "memchr", "regex-automata 0.4.9", "regex-syntax 0.8.5", ] [[package]] name = "regex-automata" version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" dependencies = [ "regex-syntax 0.6.29", ] [[package]] name = "regex-automata" version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" dependencies = [ "aho-corasick", "memchr", "regex-syntax 0.8.5", ] [[package]] name = "regex-syntax" version = "0.6.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" [[package]] name = "regex-syntax" version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" [[package]] name = "rustc-demangle" version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" [[package]] name = "rustix" version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e56a18552996ac8d29ecc3b190b4fdbb2d91ca4ec396de7bbffaf43f3d637e96" dependencies = [ "bitflags", "errno", "libc", "linux-raw-sys", "windows-sys 0.59.0", ] [[package]] name = "rustversion" version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eded382c5f5f786b989652c49544c4877d9f015cc22e145a5ea8ea66c2921cd2" [[package]] name = "rusty-fork" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb3dcc6e454c328bb824492db107ab7c0ae8fcffe4ad210136ef014458c1bc4f" dependencies = [ "fnv", "quick-error", "tempfile", "wait-timeout", ] [[package]] name = "ryu" version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" [[package]] name = "same-file" version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" dependencies = [ "winapi-util", ] [[package]] name = "scoped-tls" version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "serde" version = "1.0.219" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" version = "1.0.219" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "serde_json" version = "1.0.140" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373" dependencies = [ "itoa", "memchr", "ryu", "serde", ] [[package]] name = "sharded-slab" version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" dependencies = [ "lazy_static", ] [[package]] name = "shlex" version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] name = "signal-hook-registry" version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" dependencies = [ "libc", ] [[package]] name = "slab" version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" dependencies = [ "autocfg", ] [[package]] name = "smallvec" version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" [[package]] name = "socket2" version = "0.5.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4f5fd57c80058a56cf5c777ab8a126398ece8e442983605d280a44ce79d0edef" dependencies = [ "libc", "windows-sys 0.52.0", ] [[package]] name = "syn" version = "1.0.109" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] [[package]] name = "syn" version = "2.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] [[package]] name = "tempfile" version = "3.19.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7437ac7763b9b123ccf33c338a5cc1bac6f69b45a136c19bdd8a65e3916435bf" dependencies = [ "fastrand", "getrandom 0.3.2", "once_cell", "rustix", "windows-sys 0.59.0", ] [[package]] name = "termtree" version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "thread_local" version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" dependencies = [ "cfg-if", "once_cell", ] [[package]] name = "tokio" version = "1.43.1" dependencies = [ "async-stream", "backtrace", "bytes", "futures", "libc", "loom", "mio 1.0.3", "mio-aio", "mockall", "nix", "parking_lot", "pin-project-lite", "proptest", "rand", "signal-hook-registry", "socket2", "tempfile", "tokio-macros", "tokio-stream", "tokio-test", "tracing", "tracing-mock", "wasm-bindgen-test", "windows-sys 0.52.0", ] [[package]] name = "tokio" version = "1.44.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f382da615b842244d4b8738c82ed1275e6c5dd90c459a30941cd07080b06c91a" dependencies = [ "backtrace", "pin-project-lite", ] [[package]] name = "tokio-macros" version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "tokio-stream" version = "0.1.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eca58d7bba4a75707817a2c44174253f9236b2d5fbd055602e9d5c07c139a047" dependencies = [ "futures-core", "pin-project-lite", "tokio 1.44.1", ] [[package]] name = "tokio-test" version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2468baabc3311435b55dd935f702f42cd1b8abb7e754fb7dfb16bd36aa88f9f7" dependencies = [ "async-stream", "bytes", "futures-core", "tokio 1.44.1", "tokio-stream", ] [[package]] name = "tracing" version = "0.1.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" dependencies = [ "pin-project-lite", "tracing-attributes", "tracing-core", ] [[package]] name = "tracing-attributes" version = "0.1.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "tracing-core" version = "0.1.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c" dependencies = [ "once_cell", "valuable", ] [[package]] name = "tracing-log" version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" dependencies = [ "log", "once_cell", "tracing-core", ] [[package]] name = "tracing-mock" version = "0.1.0-beta.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ff59989fc8854bc58d3daeadbccf5d7fb0b722af043cf839c7785f1ff0daf0e" dependencies = [ "tracing", "tracing-core", ] [[package]] name = "tracing-subscriber" version = "0.3.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008" dependencies = [ "matchers", "nu-ansi-term", "once_cell", "regex", "sharded-slab", "smallvec", "thread_local", "tracing", "tracing-core", "tracing-log", ] [[package]] name = "unarray" version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" [[package]] name = "unicode-ident" version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" [[package]] name = "valuable" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" [[package]] name = "wait-timeout" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" dependencies = [ "libc", ] [[package]] name = "walkdir" version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" dependencies = [ "same-file", "winapi-util", ] [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wasi" version = "0.14.2+wasi-0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" dependencies = [ "wit-bindgen-rt", ] [[package]] name = "wasm-bindgen" version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" dependencies = [ "cfg-if", "once_cell", "wasm-bindgen-macro", ] [[package]] name = "wasm-bindgen-backend" version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" dependencies = [ "bumpalo", "log", "proc-macro2", "quote", "syn 2.0.100", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-futures" version = "0.4.50" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "555d470ec0bc3bb57890405e5d4322cc9ea83cebb085523ced7be4144dac1e61" dependencies = [ "cfg-if", "js-sys", "once_cell", "wasm-bindgen", "web-sys", ] [[package]] name = "wasm-bindgen-macro" version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" dependencies = [ "quote", "wasm-bindgen-macro-support", ] [[package]] name = "wasm-bindgen-macro-support" version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", "wasm-bindgen-backend", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" dependencies = [ "unicode-ident", ] [[package]] name = "wasm-bindgen-test" version = "0.3.50" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66c8d5e33ca3b6d9fa3b4676d774c5778031d27a578c2b007f905acf816152c3" dependencies = [ "js-sys", "minicov", "wasm-bindgen", "wasm-bindgen-futures", "wasm-bindgen-test-macro", ] [[package]] name = "wasm-bindgen-test-macro" version = "0.3.50" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "17d5042cc5fa009658f9a7333ef24291b1291a25b6382dd68862a7f3b969f69b" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "web-sys" version = "0.3.77" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33b6dd2ef9186f1f2072e409e99cd22a975331a6b3591b12c764e0e55c60d5d2" dependencies = [ "js-sys", "wasm-bindgen", ] [[package]] name = "winapi" version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" dependencies = [ "winapi-i686-pc-windows-gnu", "winapi-x86_64-pc-windows-gnu", ] [[package]] name = "winapi-i686-pc-windows-gnu" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ "windows-sys 0.59.0", ] [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] name = "windows" version = "0.58.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dd04d41d93c4992d421894c18c8b43496aa748dd4c081bac0dc93eb0489272b6" dependencies = [ "windows-core", "windows-targets 0.52.6", ] [[package]] name = "windows-core" version = "0.58.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ba6d44ec8c2591c134257ce647b7ea6b20335bf6379a27dac5f1641fcf59f99" dependencies = [ "windows-implement", "windows-interface", "windows-result", "windows-strings", "windows-targets 0.52.6", ] [[package]] name = "windows-implement" version = "0.58.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2bbd5b46c938e506ecbce286b6628a02171d56153ba733b6c741fc627ec9579b" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "windows-interface" version = "0.58.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "053c4c462dc91d3b1504c6fe5a726dd15e216ba718e84a0e46a88fbe5ded3515" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] [[package]] name = "windows-result" version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d1043d8214f791817bab27572aaa8af63732e11bf84aa21a45a78d6c317ae0e" dependencies = [ "windows-targets 0.52.6", ] [[package]] name = "windows-strings" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cd9b125c486025df0eabcb585e62173c6c9eddcec5d117d3b6e8c30e2ee4d10" dependencies = [ "windows-result", "windows-targets 0.52.6", ] [[package]] name = "windows-sys" version = "0.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" dependencies = [ "windows-targets 0.48.5", ] [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ "windows-targets 0.52.6", ] [[package]] name = "windows-sys" version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ "windows-targets 0.52.6", ] [[package]] name = "windows-targets" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" dependencies = [ "windows_aarch64_gnullvm 0.48.5", "windows_aarch64_msvc 0.48.5", "windows_i686_gnu 0.48.5", "windows_i686_msvc 0.48.5", "windows_x86_64_gnu 0.48.5", "windows_x86_64_gnullvm 0.48.5", "windows_x86_64_msvc 0.48.5", ] [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ "windows_aarch64_gnullvm 0.52.6", "windows_aarch64_msvc 0.52.6", "windows_i686_gnu 0.52.6", "windows_i686_gnullvm", "windows_i686_msvc 0.52.6", "windows_x86_64_gnu 0.52.6", "windows_x86_64_gnullvm 0.52.6", "windows_x86_64_msvc 0.52.6", ] [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" [[package]] name = "windows_aarch64_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" [[package]] name = "windows_i686_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" [[package]] name = "windows_i686_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" [[package]] name = "windows_x86_64_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" [[package]] name = "windows_x86_64_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "wit-bindgen-rt" version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" dependencies = [ "bitflags", ] [[package]] name = "zerocopy" version = "0.8.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2586fea28e186957ef732a5f8b3be2da217d65c5969d4b1e17f973ebbe876879" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" version = "0.8.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a996a8f63c5c4448cd959ac1bab0aaa3306ccfd060472f85943ee0750f0169be" dependencies = [ "proc-macro2", "quote", "syn 2.0.100", ] tokio-1.43.1/Cargo.toml0000644000000345110000000000100102510ustar # THIS FILE IS AUTOMATICALLY GENERATED BY CARGO # # When uploading crates to the registry Cargo will automatically # "normalize" Cargo.toml files for maximal compatibility # with all versions of Cargo and also rewrite `path` dependencies # to registry (e.g., crates.io) dependencies. # # If you are reading this file be aware that the original Cargo.toml # will likely look very different (and much more reasonable). # See Cargo.toml.orig for the original contents. [package] edition = "2021" rust-version = "1.70" name = "tokio" version = "1.43.1" authors = ["Tokio Contributors "] build = false autolib = false autobins = false autoexamples = false autotests = false autobenches = false description = """ An event-driven, non-blocking I/O platform for writing asynchronous I/O backed applications. """ homepage = "https://tokio.rs" readme = "README.md" keywords = [ "io", "async", "non-blocking", "futures", ] categories = [ "asynchronous", "network-programming", ] license = "MIT" repository = "https://github.com/tokio-rs/tokio" [package.metadata.cargo_check_external_types] allowed_external_types = [ "bytes::buf::buf_impl::Buf", "bytes::buf::buf_mut::BufMut", "tokio_macros::*", ] [package.metadata.docs.rs] all-features = true rustc-args = [ "--cfg", "tokio_unstable", "--cfg", "tokio_taskdump", ] rustdoc-args = [ "--cfg", "docsrs", "--cfg", "tokio_unstable", "--cfg", "tokio_taskdump", ] [package.metadata.playground] features = [ "full", "test-util", ] [features] default = [] fs = [] full = [ "fs", "io-util", "io-std", "macros", "net", "parking_lot", "process", "rt", "rt-multi-thread", "signal", "sync", "time", ] io-std = [] io-util = ["bytes"] macros = ["tokio-macros"] net = [ "libc", "mio/os-poll", "mio/os-ext", "mio/net", "socket2", "windows-sys/Win32_Foundation", "windows-sys/Win32_Security", "windows-sys/Win32_Storage_FileSystem", "windows-sys/Win32_System_Pipes", "windows-sys/Win32_System_SystemServices", ] process = [ "bytes", "libc", "mio/os-poll", "mio/os-ext", "mio/net", "signal-hook-registry", "windows-sys/Win32_Foundation", "windows-sys/Win32_System_Threading", "windows-sys/Win32_System_WindowsProgramming", ] rt = [] rt-multi-thread = ["rt"] signal = [ "libc", "mio/os-poll", "mio/net", "mio/os-ext", "signal-hook-registry", "windows-sys/Win32_Foundation", "windows-sys/Win32_System_Console", ] sync = [] test-util = [ "rt", "sync", "time", ] time = [] [lib] name = "tokio" path = "src/lib.rs" [[test]] name = "_require_full" path = "tests/_require_full.rs" [[test]] name = "async_send_sync" path = "tests/async_send_sync.rs" [[test]] name = "buffered" path = "tests/buffered.rs" [[test]] name = "coop_budget" path = "tests/coop_budget.rs" [[test]] name = "dump" path = "tests/dump.rs" [[test]] name = "duplex_stream" path = "tests/duplex_stream.rs" [[test]] name = "fs" path = "tests/fs.rs" [[test]] name = "fs_canonicalize_dir" path = "tests/fs_canonicalize_dir.rs" [[test]] name = "fs_copy" path = "tests/fs_copy.rs" [[test]] name = "fs_dir" path = "tests/fs_dir.rs" [[test]] name = "fs_file" path = "tests/fs_file.rs" [[test]] name = "fs_link" path = "tests/fs_link.rs" [[test]] name = "fs_open_options" path = "tests/fs_open_options.rs" [[test]] name = "fs_open_options_windows" path = "tests/fs_open_options_windows.rs" [[test]] name = "fs_remove_dir_all" path = "tests/fs_remove_dir_all.rs" [[test]] name = "fs_remove_file" path = "tests/fs_remove_file.rs" [[test]] name = "fs_rename" path = "tests/fs_rename.rs" [[test]] name = "fs_symlink_dir_windows" path = "tests/fs_symlink_dir_windows.rs" [[test]] name = "fs_symlink_file_windows" path = "tests/fs_symlink_file_windows.rs" [[test]] name = "fs_try_exists" path = "tests/fs_try_exists.rs" [[test]] name = "io_async_fd" path = "tests/io_async_fd.rs" [[test]] name = "io_async_read" path = "tests/io_async_read.rs" [[test]] name = "io_buf_reader" path = "tests/io_buf_reader.rs" [[test]] name = "io_buf_writer" path = "tests/io_buf_writer.rs" [[test]] name = "io_chain" path = "tests/io_chain.rs" [[test]] name = "io_copy" path = "tests/io_copy.rs" [[test]] name = "io_copy_bidirectional" path = "tests/io_copy_bidirectional.rs" [[test]] name = "io_driver" path = "tests/io_driver.rs" [[test]] name = "io_driver_drop" path = "tests/io_driver_drop.rs" [[test]] name = "io_fill_buf" path = "tests/io_fill_buf.rs" [[test]] name = "io_join" path = "tests/io_join.rs" [[test]] name = "io_lines" path = "tests/io_lines.rs" [[test]] name = "io_mem_stream" path = "tests/io_mem_stream.rs" [[test]] name = "io_panic" path = "tests/io_panic.rs" [[test]] name = "io_poll_aio" path = "tests/io_poll_aio.rs" [[test]] name = "io_read" path = "tests/io_read.rs" [[test]] name = "io_read_buf" path = "tests/io_read_buf.rs" [[test]] name = "io_read_exact" path = "tests/io_read_exact.rs" [[test]] name = "io_read_line" path = "tests/io_read_line.rs" [[test]] name = "io_read_to_end" path = "tests/io_read_to_end.rs" [[test]] name = "io_read_to_string" path = "tests/io_read_to_string.rs" [[test]] name = "io_read_until" path = "tests/io_read_until.rs" [[test]] name = "io_repeat" path = "tests/io_repeat.rs" [[test]] name = "io_sink" path = "tests/io_sink.rs" [[test]] name = "io_split" path = "tests/io_split.rs" [[test]] name = "io_take" path = "tests/io_take.rs" [[test]] name = "io_util_empty" path = "tests/io_util_empty.rs" [[test]] name = "io_write" path = "tests/io_write.rs" [[test]] name = "io_write_all" path = "tests/io_write_all.rs" [[test]] name = "io_write_all_buf" path = "tests/io_write_all_buf.rs" [[test]] name = "io_write_buf" path = "tests/io_write_buf.rs" [[test]] name = "io_write_int" path = "tests/io_write_int.rs" [[test]] name = "join_handle_panic" path = "tests/join_handle_panic.rs" [[test]] name = "macros_join" path = "tests/macros_join.rs" [[test]] name = "macros_pin" path = "tests/macros_pin.rs" [[test]] name = "macros_rename_test" path = "tests/macros_rename_test.rs" [[test]] name = "macros_select" path = "tests/macros_select.rs" [[test]] name = "macros_test" path = "tests/macros_test.rs" [[test]] name = "macros_try_join" path = "tests/macros_try_join.rs" [[test]] name = "net_bind_resource" path = "tests/net_bind_resource.rs" [[test]] name = "net_lookup_host" path = "tests/net_lookup_host.rs" [[test]] name = "net_named_pipe" path = "tests/net_named_pipe.rs" [[test]] name = "net_panic" path = "tests/net_panic.rs" [[test]] name = "net_unix_pipe" path = "tests/net_unix_pipe.rs" [[test]] name = "no_rt" path = "tests/no_rt.rs" [[test]] name = "process_arg0" path = "tests/process_arg0.rs" [[test]] name = "process_change_of_runtime" path = "tests/process_change_of_runtime.rs" [[test]] name = "process_issue_2174" path = "tests/process_issue_2174.rs" [[test]] name = "process_issue_42" path = "tests/process_issue_42.rs" [[test]] name = "process_kill_on_drop" path = "tests/process_kill_on_drop.rs" [[test]] name = "process_raw_handle" path = "tests/process_raw_handle.rs" [[test]] name = "process_smoke" path = "tests/process_smoke.rs" [[test]] name = "rt_basic" path = "tests/rt_basic.rs" [[test]] name = "rt_common" path = "tests/rt_common.rs" [[test]] name = "rt_handle" path = "tests/rt_handle.rs" [[test]] name = "rt_handle_block_on" path = "tests/rt_handle_block_on.rs" [[test]] name = "rt_local" path = "tests/rt_local.rs" [[test]] name = "rt_metrics" path = "tests/rt_metrics.rs" [[test]] name = "rt_panic" path = "tests/rt_panic.rs" [[test]] name = "rt_threaded" path = "tests/rt_threaded.rs" [[test]] name = "rt_threaded_alt" path = "tests/rt_threaded_alt.rs" [[test]] name = "rt_time_start_paused" path = "tests/rt_time_start_paused.rs" [[test]] name = "rt_unstable_metrics" path = "tests/rt_unstable_metrics.rs" [[test]] name = "signal_ctrl_c" path = "tests/signal_ctrl_c.rs" [[test]] name = "signal_drop_recv" path = "tests/signal_drop_recv.rs" [[test]] name = "signal_drop_rt" path = "tests/signal_drop_rt.rs" [[test]] name = "signal_drop_signal" path = "tests/signal_drop_signal.rs" [[test]] name = "signal_info" path = "tests/signal_info.rs" [[test]] name = "signal_multi_rt" path = "tests/signal_multi_rt.rs" [[test]] name = "signal_no_rt" path = "tests/signal_no_rt.rs" [[test]] name = "signal_notify_both" path = "tests/signal_notify_both.rs" [[test]] name = "signal_panic" path = "tests/signal_panic.rs" [[test]] name = "signal_realtime" path = "tests/signal_realtime.rs" [[test]] name = "signal_twice" path = "tests/signal_twice.rs" [[test]] name = "signal_usr1" path = "tests/signal_usr1.rs" [[test]] name = "sync_barrier" path = "tests/sync_barrier.rs" [[test]] name = "sync_broadcast" path = "tests/sync_broadcast.rs" [[test]] name = "sync_errors" path = "tests/sync_errors.rs" [[test]] name = "sync_mpsc" path = "tests/sync_mpsc.rs" [[test]] name = "sync_mpsc_weak" path = "tests/sync_mpsc_weak.rs" [[test]] name = "sync_mutex" path = "tests/sync_mutex.rs" [[test]] name = "sync_mutex_owned" path = "tests/sync_mutex_owned.rs" [[test]] name = "sync_notify" path = "tests/sync_notify.rs" [[test]] name = "sync_once_cell" path = "tests/sync_once_cell.rs" [[test]] name = "sync_oneshot" path = "tests/sync_oneshot.rs" [[test]] name = "sync_panic" path = "tests/sync_panic.rs" [[test]] name = "sync_rwlock" path = "tests/sync_rwlock.rs" [[test]] name = "sync_semaphore" path = "tests/sync_semaphore.rs" [[test]] name = "sync_semaphore_owned" path = "tests/sync_semaphore_owned.rs" [[test]] name = "sync_watch" path = "tests/sync_watch.rs" [[test]] name = "task_abort" path = "tests/task_abort.rs" [[test]] name = "task_blocking" path = "tests/task_blocking.rs" [[test]] name = "task_builder" path = "tests/task_builder.rs" [[test]] name = "task_hooks" path = "tests/task_hooks.rs" [[test]] name = "task_id" path = "tests/task_id.rs" [[test]] name = "task_join_set" path = "tests/task_join_set.rs" [[test]] name = "task_local" path = "tests/task_local.rs" [[test]] name = "task_local_set" path = "tests/task_local_set.rs" [[test]] name = "task_panic" path = "tests/task_panic.rs" [[test]] name = "task_yield_now" path = "tests/task_yield_now.rs" [[test]] name = "tcp_accept" path = "tests/tcp_accept.rs" [[test]] name = "tcp_connect" path = "tests/tcp_connect.rs" [[test]] name = "tcp_echo" path = "tests/tcp_echo.rs" [[test]] name = "tcp_into_split" path = "tests/tcp_into_split.rs" [[test]] name = "tcp_into_std" path = "tests/tcp_into_std.rs" [[test]] name = "tcp_peek" path = "tests/tcp_peek.rs" [[test]] name = "tcp_shutdown" path = "tests/tcp_shutdown.rs" [[test]] name = "tcp_socket" path = "tests/tcp_socket.rs" [[test]] name = "tcp_split" path = "tests/tcp_split.rs" [[test]] name = "tcp_stream" path = "tests/tcp_stream.rs" [[test]] name = "test_clock" path = "tests/test_clock.rs" [[test]] name = "time_interval" path = "tests/time_interval.rs" [[test]] name = "time_panic" path = "tests/time_panic.rs" [[test]] name = "time_pause" path = "tests/time_pause.rs" [[test]] name = "time_rt" path = "tests/time_rt.rs" [[test]] name = "time_sleep" path = "tests/time_sleep.rs" [[test]] name = "time_timeout" path = "tests/time_timeout.rs" [[test]] name = "tracing_sync" path = "tests/tracing_sync.rs" [[test]] name = "tracing_task" path = "tests/tracing_task.rs" [[test]] name = "tracing_time" path = "tests/tracing_time.rs" [[test]] name = "udp" path = "tests/udp.rs" [[test]] name = "uds_cred" path = "tests/uds_cred.rs" [[test]] name = "uds_datagram" path = "tests/uds_datagram.rs" [[test]] name = "uds_socket" path = "tests/uds_socket.rs" [[test]] name = "uds_split" path = "tests/uds_split.rs" [[test]] name = "uds_stream" path = "tests/uds_stream.rs" [[test]] name = "unwindsafe" path = "tests/unwindsafe.rs" [dependencies.bytes] version = "1.1.0" optional = true [dependencies.mio] version = "1.0.1" optional = true default-features = false [dependencies.parking_lot] version = "0.12.0" optional = true [dependencies.pin-project-lite] version = "0.2.11" [dependencies.tokio-macros] version = "~2.5.0" optional = true [dev-dependencies.async-stream] version = "0.3" [dev-dependencies.futures] version = "0.3.0" features = ["async-await"] [dev-dependencies.mockall] version = "0.11.1" [dev-dependencies.tokio-stream] version = "0.1" [dev-dependencies.tokio-test] version = "0.4.0" [target.'cfg(all(target_family = "wasm", not(target_os = "wasi")))'.dev-dependencies.wasm-bindgen-test] version = "0.3.0" [target.'cfg(all(tokio_unstable, target_has_atomic = "64"))'.dev-dependencies.tracing-mock] version = "= 0.1.0-beta.1" [target."cfg(loom)".dev-dependencies.loom] version = "0.7" features = [ "futures", "checkpoint", ] [target.'cfg(not(all(target_family = "wasm", target_os = "unknown")))'.dev-dependencies.rand] version = "0.8.0" [target.'cfg(not(target_family = "wasm"))'.dependencies.socket2] version = "0.5.5" features = ["all"] optional = true [target.'cfg(not(target_family = "wasm"))'.dev-dependencies.proptest] version = "1" [target.'cfg(not(target_family = "wasm"))'.dev-dependencies.socket2] version = "0.5.5" [target.'cfg(not(target_family = "wasm"))'.dev-dependencies.tempfile] version = "3.1.0" [target.'cfg(target_os = "freebsd")'.dev-dependencies.mio-aio] version = "0.9.0" features = ["tokio"] [target."cfg(tokio_taskdump)".dependencies.backtrace] version = "0.3.58" [target."cfg(tokio_unstable)".dependencies.tracing] version = "0.1.29" features = ["std"] optional = true default-features = false [target."cfg(unix)".dependencies.libc] version = "0.2.168" optional = true [target."cfg(unix)".dependencies.signal-hook-registry] version = "1.1.1" optional = true [target."cfg(unix)".dev-dependencies.libc] version = "0.2.168" [target."cfg(unix)".dev-dependencies.nix] version = "0.29.0" features = [ "aio", "fs", "socket", ] default-features = false [target."cfg(windows)".dependencies.windows-sys] version = "0.52" optional = true [target."cfg(windows)".dev-dependencies.windows-sys] version = "0.52" features = [ "Win32_Foundation", "Win32_Security_Authorization", ] [lints.rust.unexpected_cfgs] level = "warn" priority = 0 check-cfg = [ "cfg(fuzzing)", "cfg(loom)", "cfg(mio_unsupported_force_poll_poll)", "cfg(tokio_allow_from_blocking_fd)", "cfg(tokio_internal_mt_counters)", "cfg(tokio_no_parking_lot)", "cfg(tokio_no_tuning_tests)", "cfg(tokio_taskdump)", "cfg(tokio_unstable)", ] tokio-1.43.1/Cargo.toml.orig000064400000000000000000000114641046102023000137340ustar 00000000000000[package] name = "tokio" # When releasing to crates.io: # - Remove path dependencies # - Update doc url # - README.md # - Update CHANGELOG.md. # - Create "v1.x.y" git tag. version = "1.43.1" edition = "2021" rust-version = "1.70" authors = ["Tokio Contributors "] license = "MIT" readme = "README.md" repository = "https://github.com/tokio-rs/tokio" homepage = "https://tokio.rs" description = """ An event-driven, non-blocking I/O platform for writing asynchronous I/O backed applications. """ categories = ["asynchronous", "network-programming"] keywords = ["io", "async", "non-blocking", "futures"] [features] # Include nothing by default default = [] # enable everything full = [ "fs", "io-util", "io-std", "macros", "net", "parking_lot", "process", "rt", "rt-multi-thread", "signal", "sync", "time", ] fs = [] io-util = ["bytes"] # stdin, stdout, stderr io-std = [] macros = ["tokio-macros"] net = [ "libc", "mio/os-poll", "mio/os-ext", "mio/net", "socket2", "windows-sys/Win32_Foundation", "windows-sys/Win32_Security", "windows-sys/Win32_Storage_FileSystem", "windows-sys/Win32_System_Pipes", "windows-sys/Win32_System_SystemServices", ] process = [ "bytes", "libc", "mio/os-poll", "mio/os-ext", "mio/net", "signal-hook-registry", "windows-sys/Win32_Foundation", "windows-sys/Win32_System_Threading", "windows-sys/Win32_System_WindowsProgramming", ] # Includes basic task execution capabilities rt = [] rt-multi-thread = ["rt"] signal = [ "libc", "mio/os-poll", "mio/net", "mio/os-ext", "signal-hook-registry", "windows-sys/Win32_Foundation", "windows-sys/Win32_System_Console", ] sync = [] test-util = ["rt", "sync", "time"] time = [] [dependencies] tokio-macros = { version = "~2.5.0", path = "../tokio-macros", optional = true } pin-project-lite = "0.2.11" # Everything else is optional... bytes = { version = "1.1.0", optional = true } mio = { version = "1.0.1", optional = true, default-features = false } parking_lot = { version = "0.12.0", optional = true } [target.'cfg(not(target_family = "wasm"))'.dependencies] socket2 = { version = "0.5.5", optional = true, features = ["all"] } # Currently unstable. The API exposed by these features may be broken at any time. # Requires `--cfg tokio_unstable` to enable. [target.'cfg(tokio_unstable)'.dependencies] tracing = { version = "0.1.29", default-features = false, features = ["std"], optional = true } # Not in full # Currently unstable. The API exposed by these features may be broken at any time. # Requires `--cfg tokio_unstable` to enable. [target.'cfg(tokio_taskdump)'.dependencies] backtrace = { version = "0.3.58" } [target.'cfg(unix)'.dependencies] libc = { version = "0.2.168", optional = true } signal-hook-registry = { version = "1.1.1", optional = true } [target.'cfg(unix)'.dev-dependencies] libc = { version = "0.2.168" } nix = { version = "0.29.0", default-features = false, features = ["aio", "fs", "socket"] } [target.'cfg(windows)'.dependencies.windows-sys] version = "0.52" optional = true [target.'cfg(windows)'.dev-dependencies.windows-sys] version = "0.52" features = [ "Win32_Foundation", "Win32_Security_Authorization", ] [dev-dependencies] tokio-test = { version = "0.4.0", path = "../tokio-test" } tokio-stream = { version = "0.1", path = "../tokio-stream" } futures = { version = "0.3.0", features = ["async-await"] } mockall = "0.11.1" async-stream = "0.3" [target.'cfg(not(target_family = "wasm"))'.dev-dependencies] socket2 = "0.5.5" tempfile = "3.1.0" proptest = "1" [target.'cfg(not(all(target_family = "wasm", target_os = "unknown")))'.dev-dependencies] rand = "0.8.0" [target.'cfg(all(target_family = "wasm", not(target_os = "wasi")))'.dev-dependencies] wasm-bindgen-test = "0.3.0" [target.'cfg(target_os = "freebsd")'.dev-dependencies] mio-aio = { version = "0.9.0", features = ["tokio"] } [target.'cfg(loom)'.dev-dependencies] loom = { version = "0.7", features = ["futures", "checkpoint"] } [target.'cfg(all(tokio_unstable, target_has_atomic = "64"))'.dev-dependencies] tracing-mock = "= 0.1.0-beta.1" [package.metadata.docs.rs] all-features = true # enable unstable features in the documentation rustdoc-args = ["--cfg", "docsrs", "--cfg", "tokio_unstable", "--cfg", "tokio_taskdump"] # it's necessary to _also_ pass `--cfg tokio_unstable` and `--cfg tokio_taskdump` # to rustc, or else dependencies will not be enabled, and the docs build will fail. rustc-args = ["--cfg", "tokio_unstable", "--cfg", "tokio_taskdump"] [package.metadata.playground] features = ["full", "test-util"] [package.metadata.cargo_check_external_types] # The following are types that are allowed to be exposed in Tokio's public API. # The standard library is allowed by default. allowed_external_types = [ "bytes::buf::buf_impl::Buf", "bytes::buf::buf_mut::BufMut", "tokio_macros::*", ] [lints] workspace = true tokio-1.43.1/LICENSE000064400000000000000000000020561046102023000120470ustar 00000000000000MIT License Copyright (c) Tokio Contributors Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. tokio-1.43.1/README.md000064400000000000000000000216731046102023000123270ustar 00000000000000# Tokio A runtime for writing reliable, asynchronous, and slim applications with the Rust programming language. It is: * **Fast**: Tokio's zero-cost abstractions give you bare-metal performance. * **Reliable**: Tokio leverages Rust's ownership, type system, and concurrency model to reduce bugs and ensure thread safety. * **Scalable**: Tokio has a minimal footprint, and handles backpressure and cancellation naturally. [![Crates.io][crates-badge]][crates-url] [![MIT licensed][mit-badge]][mit-url] [![Build Status][actions-badge]][actions-url] [![Discord chat][discord-badge]][discord-url] [crates-badge]: https://img.shields.io/crates/v/tokio.svg [crates-url]: https://crates.io/crates/tokio [mit-badge]: https://img.shields.io/badge/license-MIT-blue.svg [mit-url]: https://github.com/tokio-rs/tokio/blob/master/LICENSE [actions-badge]: https://github.com/tokio-rs/tokio/workflows/CI/badge.svg [actions-url]: https://github.com/tokio-rs/tokio/actions?query=workflow%3ACI+branch%3Amaster [discord-badge]: https://img.shields.io/discord/500028886025895936.svg?logo=discord&style=flat-square [discord-url]: https://discord.gg/tokio [Website](https://tokio.rs) | [Guides](https://tokio.rs/tokio/tutorial) | [API Docs](https://docs.rs/tokio/latest/tokio) | [Chat](https://discord.gg/tokio) ## Overview Tokio is an event-driven, non-blocking I/O platform for writing asynchronous applications with the Rust programming language. At a high level, it provides a few major components: * A multithreaded, work-stealing based task [scheduler]. * A reactor backed by the operating system's event queue (epoll, kqueue, IOCP, etc...). * Asynchronous [TCP and UDP][net] sockets. These components provide the runtime components necessary for building an asynchronous application. [net]: https://docs.rs/tokio/latest/tokio/net/index.html [scheduler]: https://docs.rs/tokio/latest/tokio/runtime/index.html ## Example A basic TCP echo server with Tokio. Make sure you activated the full features of the tokio crate on Cargo.toml: ```toml [dependencies] tokio = { version = "1.43.1", features = ["full"] } ``` Then, on your main.rs: ```rust,no_run use tokio::net::TcpListener; use tokio::io::{AsyncReadExt, AsyncWriteExt}; #[tokio::main] async fn main() -> Result<(), Box> { let listener = TcpListener::bind("127.0.0.1:8080").await?; loop { let (mut socket, _) = listener.accept().await?; tokio::spawn(async move { let mut buf = [0; 1024]; // In a loop, read data from the socket and write the data back. loop { let n = match socket.read(&mut buf).await { // socket closed Ok(0) => return, Ok(n) => n, Err(e) => { eprintln!("failed to read from socket; err = {:?}", e); return; } }; // Write the data back if let Err(e) = socket.write_all(&buf[0..n]).await { eprintln!("failed to write to socket; err = {:?}", e); return; } } }); } } ``` More examples can be found [here][examples]. For a larger "real world" example, see the [mini-redis] repository. [examples]: https://github.com/tokio-rs/tokio/tree/master/examples [mini-redis]: https://github.com/tokio-rs/mini-redis/ To see a list of the available features flags that can be enabled, check our [docs][feature-flag-docs]. ## Getting Help First, see if the answer to your question can be found in the [Guides] or the [API documentation]. If the answer is not there, there is an active community in the [Tokio Discord server][chat]. We would be happy to try to answer your question. You can also ask your question on [the discussions page][discussions]. [Guides]: https://tokio.rs/tokio/tutorial [API documentation]: https://docs.rs/tokio/latest/tokio [chat]: https://discord.gg/tokio [discussions]: https://github.com/tokio-rs/tokio/discussions [feature-flag-docs]: https://docs.rs/tokio/#feature-flags ## Contributing :balloon: Thanks for your help improving the project! We are so happy to have you! We have a [contributing guide][guide] to help you get involved in the Tokio project. [guide]: https://github.com/tokio-rs/tokio/blob/master/CONTRIBUTING.md ## Related Projects In addition to the crates in this repository, the Tokio project also maintains several other libraries, including: * [`axum`]: A web application framework that focuses on ergonomics and modularity. * [`hyper`]: A fast and correct HTTP/1.1 and HTTP/2 implementation for Rust. * [`tonic`]: A gRPC over HTTP/2 implementation focused on high performance, interoperability, and flexibility. * [`warp`]: A super-easy, composable, web server framework for warp speeds. * [`tower`]: A library of modular and reusable components for building robust networking clients and servers. * [`tracing`] (formerly `tokio-trace`): A framework for application-level tracing and async-aware diagnostics. * [`mio`]: A low-level, cross-platform abstraction over OS I/O APIs that powers `tokio`. * [`bytes`]: Utilities for working with bytes, including efficient byte buffers. * [`loom`]: A testing tool for concurrent Rust code. [`axum`]: https://github.com/tokio-rs/axum [`warp`]: https://github.com/seanmonstar/warp [`hyper`]: https://github.com/hyperium/hyper [`tonic`]: https://github.com/hyperium/tonic [`tower`]: https://github.com/tower-rs/tower [`loom`]: https://github.com/tokio-rs/loom [`tracing`]: https://github.com/tokio-rs/tracing [`mio`]: https://github.com/tokio-rs/mio [`bytes`]: https://github.com/tokio-rs/bytes ## Changelog The Tokio repository contains multiple crates. Each crate has its own changelog. * `tokio` - [view changelog](https://github.com/tokio-rs/tokio/blob/master/tokio/CHANGELOG.md) * `tokio-util` - [view changelog](https://github.com/tokio-rs/tokio/blob/master/tokio-util/CHANGELOG.md) * `tokio-stream` - [view changelog](https://github.com/tokio-rs/tokio/blob/master/tokio-stream/CHANGELOG.md) * `tokio-macros` - [view changelog](https://github.com/tokio-rs/tokio/blob/master/tokio-macros/CHANGELOG.md) * `tokio-test` - [view changelog](https://github.com/tokio-rs/tokio/blob/master/tokio-test/CHANGELOG.md) ## Supported Rust Versions Tokio will keep a rolling MSRV (minimum supported rust version) policy of **at least** 6 months. When increasing the MSRV, the new Rust version must have been released at least six months ago. The current MSRV is 1.70. Note that the MSRV is not increased automatically, and only as part of a minor release. The MSRV history for past minor releases can be found below: * 1.39 to now - Rust 1.70 * 1.30 to 1.38 - Rust 1.63 * 1.27 to 1.29 - Rust 1.56 * 1.17 to 1.26 - Rust 1.49 * 1.15 to 1.16 - Rust 1.46 * 1.0 to 1.14 - Rust 1.45 Note that although we try to avoid the situation where a dependency transitively increases the MSRV of Tokio, we do not guarantee that this does not happen. However, every minor release will have some set of versions of dependencies that works with the MSRV of that minor release. ## Release schedule Tokio doesn't follow a fixed release schedule, but we typically make one to two new minor releases each month. We make patch releases for bugfixes as necessary. ## Bug patching policy For the purposes of making patch releases with bugfixes, we have designated certain minor releases as LTS (long term support) releases. Whenever a bug warrants a patch release with a fix for the bug, it will be backported and released as a new patch release for each LTS minor version. Our current LTS releases are: * `1.36.x` - LTS release until March 2025. (MSRV 1.63) * `1.38.x` - LTS release until July 2025. (MSRV 1.63) Each LTS release will continue to receive backported fixes for at least a year. If you wish to use a fixed minor release in your project, we recommend that you use an LTS release. To use a fixed minor version, you can specify the version with a tilde. For example, to specify that you wish to use the newest `1.32.x` patch release, you can use the following dependency specification: ```text tokio = { version = "~1.32", features = [...] } ``` ### Previous LTS releases * `1.8.x` - LTS release until February 2022. * `1.14.x` - LTS release until June 2022. * `1.18.x` - LTS release until June 2023. * `1.20.x` - LTS release until September 2023. * `1.25.x` - LTS release until March 2024. * `1.32.x` - LTS release until September 2024. ## License This project is licensed under the [MIT license]. [MIT license]: https://github.com/tokio-rs/tokio/blob/master/LICENSE ### Contribution Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in Tokio by you, shall be licensed as MIT, without any additional terms or conditions. tokio-1.43.1/docs/reactor-refactor.md000064400000000000000000000226471046102023000155660ustar 00000000000000# Refactor I/O driver Describes changes to the I/O driver for the Tokio 0.3 release. ## Goals * Support `async fn` on I/O types with `&self`. * Refine the `Registration` API. ### Non-goals * Implement `AsyncRead` / `AsyncWrite` for `&TcpStream` or other reference type. ## Overview Currently, I/O types require `&mut self` for `async` functions. The reason for this is the task's waker is stored in the I/O resource's internal state (`ScheduledIo`) instead of in the future returned by the `async` function. Because of this limitation, I/O types limit the number of wakers to one per direction (a direction is either read-related events or write-related events). Moving the waker from the internal I/O resource's state to the operation's future enables multiple wakers to be registered per operation. The "intrusive wake list" strategy used by `Notify` applies to this case, though there are some concerns unique to the I/O driver. ## Reworking the `Registration` type While `Registration` is made private (per #2728), it remains in Tokio as an implementation detail backing I/O resources such as `TcpStream`. The API of `Registration` is updated to support waiting for an arbitrary interest set with `&self`. This supports concurrent waiters with a different readiness interest. ```rust struct Registration { ... } // TODO: naming struct ReadyEvent { tick: u32, ready: mio::Ready, } impl Registration { /// `interest` must be a super set of **all** interest sets specified in /// the other methods. This is the interest set passed to `mio`. pub fn new(io: &T, interest: mio::Ready) -> io::Result where T: mio::Evented; /// Awaits for any readiness event included in `interest`. Returns a /// `ReadyEvent` representing the received readiness event. async fn readiness(&self, interest: mio::Ready) -> io::Result; /// Clears resource level readiness represented by the specified `ReadyEvent` async fn clear_readiness(&self, ready_event: ReadyEvent); ``` A new registration is created for a `T: mio::Evented` and a `interest`. This creates a `ScheduledIo` entry with the I/O driver and registers the resource with `mio`. Because Tokio uses **edge-triggered** notifications, the I/O driver only receives readiness from the OS once the ready state **changes**. The I/O driver must track each resource's known readiness state. This helps prevent syscalls when the process knows the syscall should return with `EWOULDBLOCK`. A call to `readiness()` checks if the currently known resource readiness overlaps with `interest`. If it does, then the `readiness()` immediately returns. If it does not, then the task waits until the I/O driver receives a readiness event. The pseudocode to perform a TCP read is as follows. ```rust async fn read(&self, buf: &mut [u8]) -> io::Result { loop { // Await readiness let event = self.readiness(interest).await?; match self.mio_socket.read(buf) { Ok(v) => return Ok(v), Err(ref e) if e.kind() == WouldBlock => { self.clear_readiness(event); } Err(e) => return Err(e), } } } ``` ## Reworking the `ScheduledIo` type The `ScheduledIo` type is switched to use an intrusive waker linked list. Each entry in the linked list includes the `interest` set passed to `readiness()`. ```rust #[derive(Debug)] pub(crate) struct ScheduledIo { /// Resource's known state packed with other state that must be /// atomically updated. readiness: AtomicUsize, /// Tracks tasks waiting on the resource waiters: Mutex, } #[derive(Debug)] struct Waiters { // List of intrusive waiters. list: LinkedList, /// Waiter used by `AsyncRead` implementations. reader: Option, /// Waiter used by `AsyncWrite` implementations. writer: Option, } // This struct is contained by the **future** returned by `readiness()`. #[derive(Debug)] struct Waiter { /// Intrusive linked-list pointers pointers: linked_list::Pointers, /// Waker for task waiting on I/O resource waiter: Option, /// Readiness events being waited on. This is /// the value passed to `readiness()` interest: mio::Ready, /// Should not be `Unpin`. _p: PhantomPinned, } ``` When an I/O event is received from `mio`, the associated resources' readiness is updated and the waiter list is iterated. All waiters with `interest` that overlap the received readiness event are notified. Any waiter with an `interest` that does not overlap the readiness event remains in the list. ## Cancel interest on drop The future returned by `readiness()` uses an intrusive linked list to store the waker with `ScheduledIo`. Because `readiness()` can be called concurrently, many wakers may be stored simultaneously in the list. If the `readiness()` future is dropped early, it is essential that the waker is removed from the list. This prevents leaking memory. ## Race condition Consider how many tasks may concurrently attempt I/O operations. This, combined with how Tokio uses edge-triggered events, can result in a race condition. Let's revisit the TCP read function: ```rust async fn read(&self, buf: &mut [u8]) -> io::Result { loop { // Await readiness let event = self.readiness(interest).await?; match self.mio_socket.read(buf) { Ok(v) => return Ok(v), Err(ref e) if e.kind() == WouldBlock => { self.clear_readiness(event); } Err(e) => return Err(e), } } } ``` If care is not taken, if between `mio_socket.read(buf)` returning and `clear_readiness(event)` is called, a readiness event arrives, the `read()` function could deadlock. This happens because the readiness event is received, `clear_readiness()` unsets the readiness event, and on the next iteration, `readiness().await` will block forever as a new readiness event is not received. The current I/O driver handles this condition by always registering the task's waker before performing the operation. This is not ideal as it will result in unnecessary task notification. Instead, we will use a strategy to prevent clearing readiness if an "unseen" readiness event has been received. The I/O driver will maintain a "tick" value. Every time the `mio` `poll()` function is called, the tick is incremented. Each readiness event has an associated tick. When the I/O driver sets the resource's readiness, the driver's tick is packed into the atomic `usize`. The `ScheduledIo` readiness `AtomicUsize` is structured as: ``` | shutdown | generation | driver tick | readiness | |----------+------------+--------------+-----------| | 1 bit | 7 bits + 8 bits + 16 bits | ``` The `shutdown` and `generation` components exist today. The `readiness()` function returns a `ReadyEvent` value. This value includes the `tick` component read with the resource's readiness value. When `clear_readiness()` is called, the `ReadyEvent` is provided. Readiness is only cleared if the current `tick` matches the `tick` included in the `ReadyEvent`. If the tick values do not match, the call to `readiness()` on the next iteration will not block and the new `tick` is included in the new `ReadyToken.` TODO ## Implementing `AsyncRead` / `AsyncWrite` The `AsyncRead` and `AsyncWrite` traits use a "poll" based API. This means that it is not possible to use an intrusive linked list to track the waker. Additionally, there is no future associated with the operation which means it is not possible to cancel interest in the readiness events. To implement `AsyncRead` and `AsyncWrite`, `ScheduledIo` includes dedicated waker values for the read direction and the write direction. These values are used to store the waker. Specific `interest` is not tracked for `AsyncRead` and `AsyncWrite` implementations. It is assumed that only events of interest are: * Read ready * Read closed * Write ready * Write closed Note that "read closed" and "write closed" are only available with Mio 0.7. With Mio 0.6, things were a bit messy. It is only possible to implement `AsyncRead` and `AsyncWrite` for resource types themselves and not for `&Resource`. Implementing the traits for `&Resource` would permit concurrent operations to the resource. Because only a single waker is stored per direction, any concurrent usage would result in deadlocks. An alternate implementation would call for a `Vec` but this would result in memory leaks. ## Enabling reads and writes for `&TcpStream` Instead of implementing `AsyncRead` and `AsyncWrite` for `&TcpStream`, a new function is added to `TcpStream`. ```rust impl TcpStream { /// Naming TBD fn by_ref(&self) -> TcpStreamRef<'_>; } struct TcpStreamRef<'a> { stream: &'a TcpStream, // `Waiter` is the node in the intrusive waiter linked-list read_waiter: Waiter, write_waiter: Waiter, } ``` Now, `AsyncRead` and `AsyncWrite` can be implemented on `TcpStreamRef<'a>`. When the `TcpStreamRef` is dropped, all associated waker resources are cleaned up. ### Removing all the `split()` functions With `TcpStream::by_ref()`, `TcpStream::split()` is no longer needed. Instead, it is possible to do something as follows. ```rust let rd = my_stream.by_ref(); let wr = my_stream.by_ref(); select! { // use `rd` and `wr` in separate branches. } ``` It is also possible to store a `TcpStream` in an `Arc`. ```rust let arc_stream = Arc::new(my_tcp_stream); let n = arc_stream.by_ref().read(buf).await?; ``` tokio-1.43.1/src/blocking.rs000064400000000000000000000030511046102023000137630ustar 00000000000000cfg_rt! { pub(crate) use crate::runtime::spawn_blocking; cfg_fs! { #[allow(unused_imports)] pub(crate) use crate::runtime::spawn_mandatory_blocking; } pub(crate) use crate::task::JoinHandle; } cfg_not_rt! { use std::fmt; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; pub(crate) fn spawn_blocking(_f: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { assert_send_sync::>>(); panic!("requires the `rt` Tokio feature flag") } cfg_fs! { pub(crate) fn spawn_mandatory_blocking(_f: F) -> Option> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { panic!("requires the `rt` Tokio feature flag") } } pub(crate) struct JoinHandle { _p: std::marker::PhantomData, } unsafe impl Send for JoinHandle {} unsafe impl Sync for JoinHandle {} impl Future for JoinHandle { type Output = Result; fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { unreachable!() } } impl fmt::Debug for JoinHandle where T: fmt::Debug, { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("JoinHandle").finish() } } fn assert_send_sync() { } } tokio-1.43.1/src/doc/mod.rs000064400000000000000000000024271046102023000135250ustar 00000000000000//! Types which are documented locally in the Tokio crate, but does not actually //! live here. //! //! **Note** this module is only visible on docs.rs, you cannot use it directly //! in your own code. /// The name of a type which is not defined here. /// /// This is typically used as an alias for another type, like so: /// /// ```rust,ignore /// /// See [some::other::location](https://example.com). /// type DEFINED_ELSEWHERE = crate::doc::NotDefinedHere; /// ``` /// /// This type is uninhabitable like the [`never` type] to ensure that no one /// will ever accidentally use it. /// /// [`never` type]: https://doc.rust-lang.org/std/primitive.never.html #[derive(Debug)] pub enum NotDefinedHere {} #[cfg(feature = "net")] impl mio::event::Source for NotDefinedHere { fn register( &mut self, registry: &mio::Registry, token: mio::Token, interests: mio::Interest, ) -> std::io::Result<()> { Ok(()) } fn reregister( &mut self, registry: &mio::Registry, token: mio::Token, interests: mio::Interest, ) -> std::io::Result<()> { Ok(()) } fn deregister(&mut self, registry: &mio::Registry) -> std::io::Result<()> { Ok(()) } } #[cfg(any(feature = "net", feature = "fs"))] pub mod os; tokio-1.43.1/src/doc/os.rs000064400000000000000000000077111046102023000133700ustar 00000000000000//! See [`std::os`](https://doc.rust-lang.org/std/os/index.html). /// Platform-specific extensions to `std` for Windows. /// /// See [`std::os::windows`](https://doc.rust-lang.org/std/os/windows/index.html). pub mod windows { /// Windows-specific extensions to general I/O primitives. /// /// See [`std::os::windows::io`](https://doc.rust-lang.org/std/os/windows/io/index.html). pub mod io { /// See [`std::os::windows::io::RawHandle`](https://doc.rust-lang.org/std/os/windows/io/type.RawHandle.html) pub type RawHandle = crate::doc::NotDefinedHere; /// See [`std::os::windows::io::OwnedHandle`](https://doc.rust-lang.org/std/os/windows/io/struct.OwnedHandle.html) pub type OwnedHandle = crate::doc::NotDefinedHere; /// See [`std::os::windows::io::AsRawHandle`](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawHandle.html) pub trait AsRawHandle { /// See [`std::os::windows::io::AsRawHandle::as_raw_handle`](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawHandle.html#tymethod.as_raw_handle) fn as_raw_handle(&self) -> RawHandle; } /// See [`std::os::windows::io::FromRawHandle`](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawHandle.html) pub trait FromRawHandle { /// See [`std::os::windows::io::FromRawHandle::from_raw_handle`](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawHandle.html#tymethod.from_raw_handle) unsafe fn from_raw_handle(handle: RawHandle) -> Self; } /// See [`std::os::windows::io::RawSocket`](https://doc.rust-lang.org/std/os/windows/io/type.RawSocket.html) pub type RawSocket = crate::doc::NotDefinedHere; /// See [`std::os::windows::io::AsRawSocket`](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html) pub trait AsRawSocket { /// See [`std::os::windows::io::AsRawSocket::as_raw_socket`](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html#tymethod.as_raw_socket) fn as_raw_socket(&self) -> RawSocket; } /// See [`std::os::windows::io::FromRawSocket`](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawSocket.html) pub trait FromRawSocket { /// See [`std::os::windows::io::FromRawSocket::from_raw_socket`](https://doc.rust-lang.org/std/os/windows/io/trait.FromRawSocket.html#tymethod.from_raw_socket) unsafe fn from_raw_socket(sock: RawSocket) -> Self; } /// See [`std::os::windows::io::IntoRawSocket`](https://doc.rust-lang.org/std/os/windows/io/trait.IntoRawSocket.html) pub trait IntoRawSocket { /// See [`std::os::windows::io::IntoRawSocket::into_raw_socket`](https://doc.rust-lang.org/std/os/windows/io/trait.IntoRawSocket.html#tymethod.into_raw_socket) fn into_raw_socket(self) -> RawSocket; } /// See [`std::os::windows::io::BorrowedHandle`](https://doc.rust-lang.org/std/os/windows/io/struct.BorrowedHandle.html) pub type BorrowedHandle<'handle> = crate::doc::NotDefinedHere; /// See [`std::os::windows::io::AsHandle`](https://doc.rust-lang.org/std/os/windows/io/trait.AsHandle.html) pub trait AsHandle { /// See [`std::os::windows::io::AsHandle::as_handle`](https://doc.rust-lang.org/std/os/windows/io/trait.AsHandle.html#tymethod.as_handle) fn as_handle(&self) -> BorrowedHandle<'_>; } /// See [`std::os::windows::io::BorrowedSocket`](https://doc.rust-lang.org/std/os/windows/io/struct.BorrowedSocket.html) pub type BorrowedSocket<'socket> = crate::doc::NotDefinedHere; /// See [`std::os::windows::io::AsSocket`](https://doc.rust-lang.org/std/os/windows/io/trait.AsSocket.html) pub trait AsSocket { /// See [`std::os::windows::io::AsSocket::as_socket`](https://doc.rust-lang.org/std/os/windows/io/trait.AsSocket.html#tymethod.as_socket) fn as_socket(&self) -> BorrowedSocket<'_>; } } } tokio-1.43.1/src/fs/canonicalize.rs000064400000000000000000000032421046102023000152440ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::{Path, PathBuf}; /// Returns the canonical, absolute form of a path with all intermediate /// components normalized and symbolic links resolved. /// /// This is an async version of [`std::fs::canonicalize`]. /// /// # Platform-specific behavior /// /// This function currently corresponds to the `realpath` function on Unix /// and the `CreateFile` and `GetFinalPathNameByHandle` functions on Windows. /// Note that, this [may change in the future][changes]. /// /// On Windows, this converts the path to use [extended length path][path] /// syntax, which allows your program to use longer path names, but means you /// can only join backslash-delimited paths to it, and it may be incompatible /// with other applications (if passed to the application on the command-line, /// or written to a file another application may read). /// /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior /// [path]: https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx#maxpath /// /// # Errors /// /// This function will return an error in the following situations, but is not /// limited to just these cases: /// /// * `path` does not exist. /// * A non-final component in path is not a directory. /// /// # Examples /// /// ```no_run /// use tokio::fs; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let path = fs::canonicalize("../a/../foo.txt").await?; /// Ok(()) /// } /// ``` pub async fn canonicalize(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(move || std::fs::canonicalize(path)).await } tokio-1.43.1/src/fs/copy.rs000064400000000000000000000013051046102023000135550ustar 00000000000000use crate::fs::asyncify; use std::path::Path; /// Copies the contents of one file to another. This function will also copy the permission bits /// of the original file to the destination file. /// This function will overwrite the contents of to. /// /// This is the async equivalent of [`std::fs::copy`]. /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// fs::copy("foo.txt", "bar.txt").await?; /// # Ok(()) /// # } /// ``` pub async fn copy(from: impl AsRef, to: impl AsRef) -> Result { let from = from.as_ref().to_owned(); let to = to.as_ref().to_owned(); asyncify(|| std::fs::copy(from, to)).await } tokio-1.43.1/src/fs/create_dir.rs000064400000000000000000000027611046102023000147130ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Creates a new, empty directory at the provided path. /// /// This is an async version of [`std::fs::create_dir`]. /// /// # Platform-specific behavior /// /// This function currently corresponds to the `mkdir` function on Unix /// and the `CreateDirectory` function on Windows. /// Note that, this [may change in the future][changes]. /// /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior /// /// **NOTE**: If a parent of the given path doesn't exist, this function will /// return an error. To create a directory and all its missing parents at the /// same time, use the [`create_dir_all`] function. /// /// # Errors /// /// This function will return an error in the following situations, but is not /// limited to just these cases: /// /// * User lacks permissions to create directory at `path`. /// * A parent of the given path doesn't exist. (To create a directory and all /// its missing parents at the same time, use the [`create_dir_all`] /// function.) /// * `path` already exists. /// /// [`create_dir_all`]: super::create_dir_all() /// /// # Examples /// /// ```no_run /// use tokio::fs; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// fs::create_dir("/some/dir").await?; /// Ok(()) /// } /// ``` pub async fn create_dir(path: impl AsRef) -> io::Result<()> { let path = path.as_ref().to_owned(); asyncify(move || std::fs::create_dir(path)).await } tokio-1.43.1/src/fs/create_dir_all.rs000064400000000000000000000032771046102023000155460ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Recursively creates a directory and all of its parent components if they /// are missing. /// /// This is an async version of [`std::fs::create_dir_all`]. /// /// # Platform-specific behavior /// /// This function currently corresponds to the `mkdir` function on Unix /// and the `CreateDirectory` function on Windows. /// Note that, this [may change in the future][changes]. /// /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior /// /// # Errors /// /// This function will return an error in the following situations, but is not /// limited to just these cases: /// /// * If any directory in the path specified by `path` does not already exist /// and it could not be created otherwise. The specific error conditions for /// when a directory is being created (after it is determined to not exist) are /// outlined by [`fs::create_dir`]. /// /// Notable exception is made for situations where any of the directories /// specified in the `path` could not be created as it was being created concurrently. /// Such cases are considered to be successful. That is, calling `create_dir_all` /// concurrently from multiple threads or processes is guaranteed not to fail /// due to a race condition with itself. /// /// [`fs::create_dir`]: std::fs::create_dir /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { /// fs::create_dir_all("/some/dir").await?; /// Ok(()) /// } /// ``` pub async fn create_dir_all(path: impl AsRef) -> io::Result<()> { let path = path.as_ref().to_owned(); asyncify(move || std::fs::create_dir_all(path)).await } tokio-1.43.1/src/fs/dir_builder.rs000064400000000000000000000067261046102023000151030ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// A builder for creating directories in various manners. /// /// This is a specialized version of [`std::fs::DirBuilder`] for usage on /// the Tokio runtime. #[derive(Debug, Default)] pub struct DirBuilder { /// Indicates whether to create parent directories if they are missing. recursive: bool, /// Sets the Unix mode for newly created directories. #[cfg(unix)] pub(super) mode: Option, } impl DirBuilder { /// Creates a new set of options with default mode/security settings for all /// platforms and also non-recursive. /// /// This is an async version of [`std::fs::DirBuilder::new`]. /// /// # Examples /// /// ```no_run /// use tokio::fs::DirBuilder; /// /// let builder = DirBuilder::new(); /// ``` pub fn new() -> Self { DirBuilder::default() } /// Indicates whether to create directories recursively (including all parent directories). /// Parents that do not exist are created with the same security and permissions settings. /// /// This option defaults to `false`. /// /// This is an async version of [`std::fs::DirBuilder::recursive`]. /// /// # Examples /// /// ```no_run /// use tokio::fs::DirBuilder; /// /// let mut builder = DirBuilder::new(); /// builder.recursive(true); /// ``` pub fn recursive(&mut self, recursive: bool) -> &mut Self { self.recursive = recursive; self } /// Creates the specified directory with the configured options. /// /// It is considered an error if the directory already exists unless /// recursive mode is enabled. /// /// This is an async version of [`std::fs::DirBuilder::create`]. /// /// # Errors /// /// An error will be returned under the following circumstances: /// /// * Path already points to an existing file. /// * Path already points to an existing directory and the mode is /// non-recursive. /// * The calling process doesn't have permissions to create the directory /// or its missing parents. /// * Other I/O error occurred. /// /// # Examples /// /// ```no_run /// use tokio::fs::DirBuilder; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// DirBuilder::new() /// .recursive(true) /// .create("/tmp/foo/bar/baz") /// .await?; /// /// Ok(()) /// } /// ``` pub async fn create(&self, path: impl AsRef) -> io::Result<()> { let path = path.as_ref().to_owned(); let mut builder = std::fs::DirBuilder::new(); builder.recursive(self.recursive); #[cfg(unix)] { if let Some(mode) = self.mode { std::os::unix::fs::DirBuilderExt::mode(&mut builder, mode); } } asyncify(move || builder.create(path)).await } } feature! { #![unix] impl DirBuilder { /// Sets the mode to create new directories with. /// /// This option defaults to 0o777. /// /// # Examples /// /// /// ```no_run /// use tokio::fs::DirBuilder; /// /// let mut builder = DirBuilder::new(); /// builder.mode(0o775); /// ``` pub fn mode(&mut self, mode: u32) -> &mut Self { self.mode = Some(mode); self } } } tokio-1.43.1/src/fs/file/tests.rs000064400000000000000000000560071046102023000146750ustar 00000000000000use super::*; use crate::{ fs::mocks::*, io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}, }; use mockall::{predicate::eq, Sequence}; use tokio_test::{assert_pending, assert_ready_err, assert_ready_ok, task}; const HELLO: &[u8] = b"hello world..."; const FOO: &[u8] = b"foo bar baz..."; #[test] fn open_read() { let mut file = MockFile::default(); file.expect_inner_read().once().returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); let mut file = File::from_std(file); let mut buf = [0; 1024]; let mut t = task::spawn(file.read(&mut buf)); assert_eq!(0, pool::len()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); let n = assert_ready_ok!(t.poll()); assert_eq!(n, HELLO.len()); assert_eq!(&buf[..n], HELLO); } #[test] fn read_twice_before_dispatch() { let mut file = MockFile::default(); file.expect_inner_read().once().returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); let mut file = File::from_std(file); let mut buf = [0; 1024]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); assert_pending!(t.poll()); assert_eq!(pool::len(), 1); pool::run_one(); assert!(t.is_woken()); let n = assert_ready_ok!(t.poll()); assert_eq!(&buf[..n], HELLO); } #[test] fn read_with_smaller_buf() { let mut file = MockFile::default(); file.expect_inner_read().once().returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); let mut file = File::from_std(file); { let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); } pool::run_one(); { let mut buf = [0; 4]; let mut t = task::spawn(file.read(&mut buf)); let n = assert_ready_ok!(t.poll()); assert_eq!(n, 4); assert_eq!(&buf[..], &HELLO[..n]); } // Calling again immediately succeeds with the rest of the buffer let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); let n = assert_ready_ok!(t.poll()); assert_eq!(n, 10); assert_eq!(&buf[..n], &HELLO[4..]); assert_eq!(0, pool::len()); } #[test] fn read_with_bigger_buf() { let mut seq = Sequence::new(); let mut file = MockFile::default(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..4].copy_from_slice(&HELLO[..4]); Ok(4) }); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len() - 4].copy_from_slice(&HELLO[4..]); Ok(HELLO.len() - 4) }); let mut file = File::from_std(file); { let mut buf = [0; 4]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); } pool::run_one(); { let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); let n = assert_ready_ok!(t.poll()); assert_eq!(n, 4); assert_eq!(&buf[..n], &HELLO[..n]); } // Calling again immediately succeeds with the rest of the buffer let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); let n = assert_ready_ok!(t.poll()); assert_eq!(n, 10); assert_eq!(&buf[..n], &HELLO[4..]); assert_eq!(0, pool::len()); } #[test] fn read_err_then_read_success() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|_| Err(io::ErrorKind::Other.into())); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); let mut file = File::from_std(file); { let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); assert_ready_err!(t.poll()); } { let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let n = assert_ready_ok!(t.poll()); assert_eq!(n, HELLO.len()); assert_eq!(&buf[..n], HELLO); } } #[test] fn open_write() { let mut file = MockFile::default(); file.expect_inner_write() .once() .with(eq(HELLO)) .returning(|buf| Ok(buf.len())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_eq!(0, pool::len()); assert_ready_ok!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(!t.is_woken()); let mut t = task::spawn(file.flush()); assert_ready_ok!(t.poll()); } #[test] fn flush_while_idle() { let file = MockFile::default(); let mut file = File::from_std(file); let mut t = task::spawn(file.flush()); assert_ready_ok!(t.poll()); } #[test] #[cfg_attr(miri, ignore)] // takes a really long time with miri fn read_with_buffer_larger_than_max() { // Chunks let chunk_a = crate::io::blocking::DEFAULT_MAX_BUF_SIZE; let chunk_b = chunk_a * 2; let chunk_c = chunk_a * 3; let chunk_d = chunk_a * 4; assert_eq!(chunk_d / 1024 / 1024, 8); let mut data = vec![]; for i in 0..(chunk_d - 1) { data.push((i % 151) as u8); } let data = Arc::new(data); let d0 = data.clone(); let d1 = data.clone(); let d2 = data.clone(); let d3 = data.clone(); let mut seq = Sequence::new(); let mut file = MockFile::default(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(move |buf| { buf[0..chunk_a].copy_from_slice(&d0[0..chunk_a]); Ok(chunk_a) }); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(move |buf| { buf[..chunk_a].copy_from_slice(&d1[chunk_a..chunk_b]); Ok(chunk_b - chunk_a) }); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(move |buf| { buf[..chunk_a].copy_from_slice(&d2[chunk_b..chunk_c]); Ok(chunk_c - chunk_b) }); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(move |buf| { buf[..chunk_a - 1].copy_from_slice(&d3[chunk_c..]); Ok(chunk_a - 1) }); let mut file = File::from_std(file); let mut actual = vec![0; chunk_d]; let mut pos = 0; while pos < data.len() { let mut t = task::spawn(file.read(&mut actual[pos..])); assert_pending!(t.poll()); pool::run_one(); assert!(t.is_woken()); let n = assert_ready_ok!(t.poll()); assert!(n <= chunk_a); pos += n; } assert_eq!(&data[..], &actual[..data.len()]); } #[test] #[cfg_attr(miri, ignore)] // takes a really long time with miri fn write_with_buffer_larger_than_max() { // Chunks let chunk_a = crate::io::blocking::DEFAULT_MAX_BUF_SIZE; let chunk_b = chunk_a * 2; let chunk_c = chunk_a * 3; let chunk_d = chunk_a * 4; assert_eq!(chunk_d / 1024 / 1024, 8); let mut data = vec![]; for i in 0..(chunk_d - 1) { data.push((i % 151) as u8); } let data = Arc::new(data); let d0 = data.clone(); let d1 = data.clone(); let d2 = data.clone(); let d3 = data.clone(); let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .withf(move |buf| buf == &d0[0..chunk_a]) .returning(|buf| Ok(buf.len())); file.expect_inner_write() .once() .in_sequence(&mut seq) .withf(move |buf| buf == &d1[chunk_a..chunk_b]) .returning(|buf| Ok(buf.len())); file.expect_inner_write() .once() .in_sequence(&mut seq) .withf(move |buf| buf == &d2[chunk_b..chunk_c]) .returning(|buf| Ok(buf.len())); file.expect_inner_write() .once() .in_sequence(&mut seq) .withf(move |buf| buf == &d3[chunk_c..chunk_d - 1]) .returning(|buf| Ok(buf.len())); let mut file = File::from_std(file); let mut rem = &data[..]; let mut first = true; while !rem.is_empty() { let mut task = task::spawn(file.write(rem)); if !first { assert_pending!(task.poll()); pool::run_one(); assert!(task.is_woken()); } first = false; let n = assert_ready_ok!(task.poll()); rem = &rem[n..]; } pool::run_one(); } #[test] fn write_twice_before_dispatch() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(HELLO)) .returning(|buf| Ok(buf.len())); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(FOO)) .returning(|buf| Ok(buf.len())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.write(FOO)); assert_pending!(t.poll()); assert_eq!(pool::len(), 1); pool::run_one(); assert!(t.is_woken()); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.flush()); assert_pending!(t.poll()); assert_eq!(pool::len(), 1); pool::run_one(); assert!(t.is_woken()); assert_ready_ok!(t.poll()); } #[test] fn incomplete_read_followed_by_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); file.expect_inner_seek() .once() .with(eq(SeekFrom::Current(-(HELLO.len() as i64)))) .in_sequence(&mut seq) .returning(|_| Ok(0)); file.expect_inner_write() .once() .with(eq(FOO)) .returning(|_| Ok(FOO.len())); let mut file = File::from_std(file); let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let mut t = task::spawn(file.write(FOO)); assert_ready_ok!(t.poll()); assert_eq!(pool::len(), 1); pool::run_one(); let mut t = task::spawn(file.flush()); assert_ready_ok!(t.poll()); } #[test] fn incomplete_partial_read_followed_by_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); file.expect_inner_seek() .once() .in_sequence(&mut seq) .with(eq(SeekFrom::Current(-10))) .returning(|_| Ok(0)); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(FOO)) .returning(|_| Ok(FOO.len())); let mut file = File::from_std(file); let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let mut buf = [0; 4]; let mut t = task::spawn(file.read(&mut buf)); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.write(FOO)); assert_ready_ok!(t.poll()); assert_eq!(pool::len(), 1); pool::run_one(); let mut t = task::spawn(file.flush()); assert_ready_ok!(t.poll()); } #[test] fn incomplete_read_followed_by_flush() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); file.expect_inner_seek() .once() .in_sequence(&mut seq) .with(eq(SeekFrom::Current(-(HELLO.len() as i64)))) .returning(|_| Ok(0)); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(FOO)) .returning(|_| Ok(FOO.len())); let mut file = File::from_std(file); let mut buf = [0; 32]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let mut t = task::spawn(file.flush()); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.write(FOO)); assert_ready_ok!(t.poll()); pool::run_one(); } #[test] fn incomplete_flush_followed_by_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(HELLO)) .returning(|_| Ok(HELLO.len())); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(FOO)) .returning(|_| Ok(FOO.len())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); let n = assert_ready_ok!(t.poll()); assert_eq!(n, HELLO.len()); let mut t = task::spawn(file.flush()); assert_pending!(t.poll()); // TODO: Move under write pool::run_one(); let mut t = task::spawn(file.write(FOO)); assert_ready_ok!(t.poll()); pool::run_one(); let mut t = task::spawn(file.flush()); assert_ready_ok!(t.poll()); } #[test] fn read_err() { let mut file = MockFile::default(); file.expect_inner_read() .once() .returning(|_| Err(io::ErrorKind::Other.into())); let mut file = File::from_std(file); let mut buf = [0; 1024]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); assert!(t.is_woken()); assert_ready_err!(t.poll()); } #[test] fn write_write_err() { let mut file = MockFile::default(); file.expect_inner_write() .once() .returning(|_| Err(io::ErrorKind::Other.into())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); pool::run_one(); let mut t = task::spawn(file.write(FOO)); assert_ready_err!(t.poll()); } #[test] fn write_read_write_err() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .returning(|_| Err(io::ErrorKind::Other.into())); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); pool::run_one(); let mut buf = [0; 1024]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let mut t = task::spawn(file.write(FOO)); assert_ready_err!(t.poll()); } #[test] fn write_read_flush_err() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .returning(|_| Err(io::ErrorKind::Other.into())); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); pool::run_one(); let mut buf = [0; 1024]; let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let mut t = task::spawn(file.flush()); assert_ready_err!(t.poll()); } #[test] fn write_seek_write_err() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .returning(|_| Err(io::ErrorKind::Other.into())); file.expect_inner_seek() .once() .with(eq(SeekFrom::Start(0))) .in_sequence(&mut seq) .returning(|_| Ok(0)); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); pool::run_one(); { let mut t = task::spawn(file.seek(SeekFrom::Start(0))); assert_pending!(t.poll()); } pool::run_one(); let mut t = task::spawn(file.write(FOO)); assert_ready_err!(t.poll()); } #[test] fn write_seek_flush_err() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .returning(|_| Err(io::ErrorKind::Other.into())); file.expect_inner_seek() .once() .with(eq(SeekFrom::Start(0))) .in_sequence(&mut seq) .returning(|_| Ok(0)); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); pool::run_one(); { let mut t = task::spawn(file.seek(SeekFrom::Start(0))); assert_pending!(t.poll()); } pool::run_one(); let mut t = task::spawn(file.flush()); assert_ready_err!(t.poll()); } #[test] fn sync_all_ordered_after_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(HELLO)) .returning(|_| Ok(HELLO.len())); file.expect_sync_all().once().returning(|| Ok(())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.sync_all()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_ready_ok!(t.poll()); } #[test] fn sync_all_err_ordered_after_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(HELLO)) .returning(|_| Ok(HELLO.len())); file.expect_sync_all() .once() .returning(|| Err(io::ErrorKind::Other.into())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.sync_all()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_ready_err!(t.poll()); } #[test] fn sync_data_ordered_after_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(HELLO)) .returning(|_| Ok(HELLO.len())); file.expect_sync_data().once().returning(|| Ok(())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.sync_data()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_ready_ok!(t.poll()); } #[test] fn sync_data_err_ordered_after_write() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .with(eq(HELLO)) .returning(|_| Ok(HELLO.len())); file.expect_sync_data() .once() .returning(|| Err(io::ErrorKind::Other.into())); let mut file = File::from_std(file); let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); let mut t = task::spawn(file.sync_data()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_pending!(t.poll()); assert_eq!(1, pool::len()); pool::run_one(); assert!(t.is_woken()); assert_ready_err!(t.poll()); } #[test] fn open_set_len_ok() { let mut file = MockFile::default(); file.expect_set_len().with(eq(123)).returning(|_| Ok(())); let file = File::from_std(file); let mut t = task::spawn(file.set_len(123)); assert_pending!(t.poll()); pool::run_one(); assert!(t.is_woken()); assert_ready_ok!(t.poll()); } #[test] fn open_set_len_err() { let mut file = MockFile::default(); file.expect_set_len() .with(eq(123)) .returning(|_| Err(io::ErrorKind::Other.into())); let file = File::from_std(file); let mut t = task::spawn(file.set_len(123)); assert_pending!(t.poll()); pool::run_one(); assert!(t.is_woken()); assert_ready_err!(t.poll()); } #[test] fn partial_read_set_len_ok() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..HELLO.len()].copy_from_slice(HELLO); Ok(HELLO.len()) }); file.expect_inner_seek() .once() .with(eq(SeekFrom::Current(-(HELLO.len() as i64)))) .in_sequence(&mut seq) .returning(|_| Ok(0)); file.expect_set_len() .once() .in_sequence(&mut seq) .with(eq(123)) .returning(|_| Ok(())); file.expect_inner_read() .once() .in_sequence(&mut seq) .returning(|buf| { buf[0..FOO.len()].copy_from_slice(FOO); Ok(FOO.len()) }); let mut buf = [0; 32]; let mut file = File::from_std(file); { let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); } pool::run_one(); { let mut t = task::spawn(file.set_len(123)); assert_pending!(t.poll()); pool::run_one(); assert_ready_ok!(t.poll()); } let mut t = task::spawn(file.read(&mut buf)); assert_pending!(t.poll()); pool::run_one(); let n = assert_ready_ok!(t.poll()); assert_eq!(n, FOO.len()); assert_eq!(&buf[..n], FOO); } #[test] fn busy_file_seek_error() { let mut file = MockFile::default(); let mut seq = Sequence::new(); file.expect_inner_write() .once() .in_sequence(&mut seq) .returning(|_| Err(io::ErrorKind::Other.into())); let mut file = crate::io::BufReader::new(File::from_std(file)); { let mut t = task::spawn(file.write(HELLO)); assert_ready_ok!(t.poll()); } pool::run_one(); let mut t = task::spawn(file.seek(SeekFrom::Start(0))); assert_ready_err!(t.poll()); } tokio-1.43.1/src/fs/file.rs000064400000000000000000000763021046102023000135330ustar 00000000000000//! Types for working with [`File`]. //! //! [`File`]: File use crate::fs::{asyncify, OpenOptions}; use crate::io::blocking::{Buf, DEFAULT_MAX_BUF_SIZE}; use crate::io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use crate::sync::Mutex; use std::cmp; use std::fmt; use std::fs::{Metadata, Permissions}; use std::future::Future; use std::io::{self, Seek, SeekFrom}; use std::path::Path; use std::pin::Pin; use std::sync::Arc; use std::task::{ready, Context, Poll}; #[cfg(test)] use super::mocks::JoinHandle; #[cfg(test)] use super::mocks::MockFile as StdFile; #[cfg(test)] use super::mocks::{spawn_blocking, spawn_mandatory_blocking}; #[cfg(not(test))] use crate::blocking::JoinHandle; #[cfg(not(test))] use crate::blocking::{spawn_blocking, spawn_mandatory_blocking}; #[cfg(not(test))] use std::fs::File as StdFile; /// A reference to an open file on the filesystem. /// /// This is a specialized version of [`std::fs::File`] for usage from the /// Tokio runtime. /// /// An instance of a `File` can be read and/or written depending on what options /// it was opened with. Files also implement [`AsyncSeek`] to alter the logical /// cursor that the file contains internally. /// /// A file will not be closed immediately when it goes out of scope if there /// are any IO operations that have not yet completed. To ensure that a file is /// closed immediately when it is dropped, you should call [`flush`] before /// dropping it. Note that this does not ensure that the file has been fully /// written to disk; the operating system might keep the changes around in an /// in-memory buffer. See the [`sync_all`] method for telling the OS to write /// the data to disk. /// /// Reading and writing to a `File` is usually done using the convenience /// methods found on the [`AsyncReadExt`] and [`AsyncWriteExt`] traits. /// /// [`AsyncSeek`]: trait@crate::io::AsyncSeek /// [`flush`]: fn@crate::io::AsyncWriteExt::flush /// [`sync_all`]: fn@crate::fs::File::sync_all /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt /// /// # Examples /// /// Create a new file and asynchronously write bytes to it: /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; // for write_all() /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// file.write_all(b"hello, world!").await?; /// # Ok(()) /// # } /// ``` /// /// Read the contents of a file into a buffer: /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncReadExt; // for read_to_end() /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::open("foo.txt").await?; /// /// let mut contents = vec![]; /// file.read_to_end(&mut contents).await?; /// /// println!("len = {}", contents.len()); /// # Ok(()) /// # } /// ``` pub struct File { std: Arc, inner: Mutex, max_buf_size: usize, } struct Inner { state: State, /// Errors from writes/flushes are returned in write/flush calls. If a write /// error is observed while performing a read, it is saved until the next /// write / flush call. last_write_err: Option, pos: u64, } #[derive(Debug)] enum State { Idle(Option), Busy(JoinHandle<(Operation, Buf)>), } #[derive(Debug)] enum Operation { Read(io::Result), Write(io::Result<()>), Seek(io::Result), } impl File { /// Attempts to open a file in read-only mode. /// /// See [`OpenOptions`] for more details. /// /// # Errors /// /// This function will return an error if called from outside of the Tokio /// runtime or if path does not already exist. Other errors may also be /// returned according to `OpenOptions::open`. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncReadExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::open("foo.txt").await?; /// /// let mut contents = vec![]; /// file.read_to_end(&mut contents).await?; /// /// println!("len = {}", contents.len()); /// # Ok(()) /// # } /// ``` /// /// The [`read_to_end`] method is defined on the [`AsyncReadExt`] trait. /// /// [`read_to_end`]: fn@crate::io::AsyncReadExt::read_to_end /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt pub async fn open(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); let std = asyncify(|| StdFile::open(path)).await?; Ok(File::from_std(std)) } /// Opens a file in write-only mode. /// /// This function will create a file if it does not exist, and will truncate /// it if it does. /// /// See [`OpenOptions`] for more details. /// /// # Errors /// /// Results in an error if called from outside of the Tokio runtime or if /// the underlying [`create`] call results in an error. /// /// [`create`]: std::fs::File::create /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// file.write_all(b"hello, world!").await?; /// # Ok(()) /// # } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt pub async fn create(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); let std_file = asyncify(move || StdFile::create(path)).await?; Ok(File::from_std(std_file)) } /// Opens a file in read-write mode. /// /// This function will create a file if it does not exist, or return an error /// if it does. This way, if the call succeeds, the file returned is guaranteed /// to be new. /// /// This option is useful because it is atomic. Otherwise between checking /// whether a file exists and creating a new one, the file may have been /// created by another process (a TOCTOU race condition / attack). /// /// This can also be written using `File::options().read(true).write(true).create_new(true).open(...)`. /// /// See [`OpenOptions`] for more details. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::create_new("foo.txt").await?; /// file.write_all(b"hello, world!").await?; /// # Ok(()) /// # } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt pub async fn create_new>(path: P) -> std::io::Result { Self::options() .read(true) .write(true) .create_new(true) .open(path) .await } /// Returns a new [`OpenOptions`] object. /// /// This function returns a new `OpenOptions` object that you can use to /// open or create a file with specific options if `open()` or `create()` /// are not appropriate. /// /// It is equivalent to `OpenOptions::new()`, but allows you to write more /// readable code. Instead of /// `OpenOptions::new().append(true).open("example.log")`, /// you can write `File::options().append(true).open("example.log")`. This /// also avoids the need to import `OpenOptions`. /// /// See the [`OpenOptions::new`] function for more details. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut f = File::options().append(true).open("example.log").await?; /// f.write_all(b"new line\n").await?; /// # Ok(()) /// # } /// ``` #[must_use] pub fn options() -> OpenOptions { OpenOptions::new() } /// Converts a [`std::fs::File`] to a [`tokio::fs::File`](File). /// /// # Examples /// /// ```no_run /// // This line could block. It is not recommended to do this on the Tokio /// // runtime. /// let std_file = std::fs::File::open("foo.txt").unwrap(); /// let file = tokio::fs::File::from_std(std_file); /// ``` pub fn from_std(std: StdFile) -> File { File { std: Arc::new(std), inner: Mutex::new(Inner { state: State::Idle(Some(Buf::with_capacity(0))), last_write_err: None, pos: 0, }), max_buf_size: DEFAULT_MAX_BUF_SIZE, } } /// Attempts to sync all OS-internal metadata to disk. /// /// This function will attempt to ensure that all in-core data reaches the /// filesystem before returning. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// file.write_all(b"hello, world!").await?; /// file.sync_all().await?; /// # Ok(()) /// # } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt pub async fn sync_all(&self) -> io::Result<()> { let mut inner = self.inner.lock().await; inner.complete_inflight().await; let std = self.std.clone(); asyncify(move || std.sync_all()).await } /// This function is similar to `sync_all`, except that it may not /// synchronize file metadata to the filesystem. /// /// This is intended for use cases that must synchronize content, but don't /// need the metadata on disk. The goal of this method is to reduce disk /// operations. /// /// Note that some platforms may simply implement this in terms of `sync_all`. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// file.write_all(b"hello, world!").await?; /// file.sync_data().await?; /// # Ok(()) /// # } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt pub async fn sync_data(&self) -> io::Result<()> { let mut inner = self.inner.lock().await; inner.complete_inflight().await; let std = self.std.clone(); asyncify(move || std.sync_data()).await } /// Truncates or extends the underlying file, updating the size of this file to become size. /// /// If the size is less than the current file's size, then the file will be /// shrunk. If it is greater than the current file's size, then the file /// will be extended to size and have all of the intermediate data filled in /// with 0s. /// /// # Errors /// /// This function will return an error if the file is not opened for /// writing. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// file.write_all(b"hello, world!").await?; /// file.set_len(10).await?; /// # Ok(()) /// # } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt pub async fn set_len(&self, size: u64) -> io::Result<()> { let mut inner = self.inner.lock().await; inner.complete_inflight().await; let mut buf = match inner.state { State::Idle(ref mut buf_cell) => buf_cell.take().unwrap(), _ => unreachable!(), }; let seek = if !buf.is_empty() { Some(SeekFrom::Current(buf.discard_read())) } else { None }; let std = self.std.clone(); inner.state = State::Busy(spawn_blocking(move || { let res = if let Some(seek) = seek { (&*std).seek(seek).and_then(|_| std.set_len(size)) } else { std.set_len(size) } .map(|()| 0); // the value is discarded later // Return the result as a seek (Operation::Seek(res), buf) })); let (op, buf) = match inner.state { State::Idle(_) => unreachable!(), State::Busy(ref mut rx) => rx.await?, }; inner.state = State::Idle(Some(buf)); match op { Operation::Seek(res) => res.map(|pos| { inner.pos = pos; }), _ => unreachable!(), } } /// Queries metadata about the underlying file. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// /// # async fn dox() -> std::io::Result<()> { /// let file = File::open("foo.txt").await?; /// let metadata = file.metadata().await?; /// /// println!("{:?}", metadata); /// # Ok(()) /// # } /// ``` pub async fn metadata(&self) -> io::Result { let std = self.std.clone(); asyncify(move || std.metadata()).await } /// Creates a new `File` instance that shares the same underlying file handle /// as the existing `File` instance. Reads, writes, and seeks will affect both /// File instances simultaneously. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// /// # async fn dox() -> std::io::Result<()> { /// let file = File::open("foo.txt").await?; /// let file_clone = file.try_clone().await?; /// # Ok(()) /// # } /// ``` pub async fn try_clone(&self) -> io::Result { self.inner.lock().await.complete_inflight().await; let std = self.std.clone(); let std_file = asyncify(move || std.try_clone()).await?; Ok(File::from_std(std_file)) } /// Destructures `File` into a [`std::fs::File`]. This function is /// async to allow any in-flight operations to complete. /// /// Use `File::try_into_std` to attempt conversion immediately. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// /// # async fn dox() -> std::io::Result<()> { /// let tokio_file = File::open("foo.txt").await?; /// let std_file = tokio_file.into_std().await; /// # Ok(()) /// # } /// ``` pub async fn into_std(mut self) -> StdFile { self.inner.get_mut().complete_inflight().await; Arc::try_unwrap(self.std).expect("Arc::try_unwrap failed") } /// Tries to immediately destructure `File` into a [`std::fs::File`]. /// /// # Errors /// /// This function will return an error containing the file if some /// operation is in-flight. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// /// # async fn dox() -> std::io::Result<()> { /// let tokio_file = File::open("foo.txt").await?; /// let std_file = tokio_file.try_into_std().unwrap(); /// # Ok(()) /// # } /// ``` pub fn try_into_std(mut self) -> Result { match Arc::try_unwrap(self.std) { Ok(file) => Ok(file), Err(std_file_arc) => { self.std = std_file_arc; Err(self) } } } /// Changes the permissions on the underlying file. /// /// # Platform-specific behavior /// /// This function currently corresponds to the `fchmod` function on Unix and /// the `SetFileInformationByHandle` function on Windows. Note that, this /// [may change in the future][changes]. /// /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior /// /// # Errors /// /// This function will return an error if the user lacks permission change /// attributes on the underlying file. It may also return an error in other /// os-specific unspecified cases. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// /// # async fn dox() -> std::io::Result<()> { /// let file = File::open("foo.txt").await?; /// let mut perms = file.metadata().await?.permissions(); /// perms.set_readonly(true); /// file.set_permissions(perms).await?; /// # Ok(()) /// # } /// ``` pub async fn set_permissions(&self, perm: Permissions) -> io::Result<()> { let std = self.std.clone(); asyncify(move || std.set_permissions(perm)).await } /// Set the maximum buffer size for the underlying [`AsyncRead`] / [`AsyncWrite`] operation. /// /// Although Tokio uses a sensible default value for this buffer size, this function would be /// useful for changing that default depending on the situation. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::AsyncWriteExt; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::open("foo.txt").await?; /// /// // Set maximum buffer size to 8 MiB /// file.set_max_buf_size(8 * 1024 * 1024); /// /// let mut buf = vec![1; 1024 * 1024 * 1024]; /// /// // Write the 1 GiB buffer in chunks up to 8 MiB each. /// file.write_all(&mut buf).await?; /// # Ok(()) /// # } /// ``` pub fn set_max_buf_size(&mut self, max_buf_size: usize) { self.max_buf_size = max_buf_size; } } impl AsyncRead for File { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, dst: &mut ReadBuf<'_>, ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); loop { match inner.state { State::Idle(ref mut buf_cell) => { let mut buf = buf_cell.take().unwrap(); if !buf.is_empty() { buf.copy_to(dst); *buf_cell = Some(buf); return Poll::Ready(Ok(())); } let std = me.std.clone(); let max_buf_size = cmp::min(dst.remaining(), me.max_buf_size); inner.state = State::Busy(spawn_blocking(move || { // SAFETY: the `Read` implementation of `std` does not // read from the buffer it is borrowing and correctly // reports the length of the data written into the buffer. let res = unsafe { buf.read_from(&mut &*std, max_buf_size) }; (Operation::Read(res), buf) })); } State::Busy(ref mut rx) => { let (op, mut buf) = ready!(Pin::new(rx).poll(cx))?; match op { Operation::Read(Ok(_)) => { buf.copy_to(dst); inner.state = State::Idle(Some(buf)); return Poll::Ready(Ok(())); } Operation::Read(Err(e)) => { assert!(buf.is_empty()); inner.state = State::Idle(Some(buf)); return Poll::Ready(Err(e)); } Operation::Write(Ok(())) => { assert!(buf.is_empty()); inner.state = State::Idle(Some(buf)); continue; } Operation::Write(Err(e)) => { assert!(inner.last_write_err.is_none()); inner.last_write_err = Some(e.kind()); inner.state = State::Idle(Some(buf)); } Operation::Seek(result) => { assert!(buf.is_empty()); inner.state = State::Idle(Some(buf)); if let Ok(pos) = result { inner.pos = pos; } continue; } } } } } } } impl AsyncSeek for File { fn start_seek(self: Pin<&mut Self>, mut pos: SeekFrom) -> io::Result<()> { let me = self.get_mut(); let inner = me.inner.get_mut(); match inner.state { State::Busy(_) => Err(io::Error::new( io::ErrorKind::Other, "other file operation is pending, call poll_complete before start_seek", )), State::Idle(ref mut buf_cell) => { let mut buf = buf_cell.take().unwrap(); // Factor in any unread data from the buf if !buf.is_empty() { let n = buf.discard_read(); if let SeekFrom::Current(ref mut offset) = pos { *offset += n; } } let std = me.std.clone(); inner.state = State::Busy(spawn_blocking(move || { let res = (&*std).seek(pos); (Operation::Seek(res), buf) })); Ok(()) } } } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let inner = self.inner.get_mut(); loop { match inner.state { State::Idle(_) => return Poll::Ready(Ok(inner.pos)), State::Busy(ref mut rx) => { let (op, buf) = ready!(Pin::new(rx).poll(cx))?; inner.state = State::Idle(Some(buf)); match op { Operation::Read(_) => {} Operation::Write(Err(e)) => { assert!(inner.last_write_err.is_none()); inner.last_write_err = Some(e.kind()); } Operation::Write(_) => {} Operation::Seek(res) => { if let Ok(pos) = res { inner.pos = pos; } return Poll::Ready(res); } } } } } } } impl AsyncWrite for File { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, src: &[u8], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); if let Some(e) = inner.last_write_err.take() { return Poll::Ready(Err(e.into())); } loop { match inner.state { State::Idle(ref mut buf_cell) => { let mut buf = buf_cell.take().unwrap(); let seek = if !buf.is_empty() { Some(SeekFrom::Current(buf.discard_read())) } else { None }; let n = buf.copy_from(src, me.max_buf_size); let std = me.std.clone(); let blocking_task_join_handle = spawn_mandatory_blocking(move || { let res = if let Some(seek) = seek { (&*std).seek(seek).and_then(|_| buf.write_to(&mut &*std)) } else { buf.write_to(&mut &*std) }; (Operation::Write(res), buf) }) .ok_or_else(|| { io::Error::new(io::ErrorKind::Other, "background task failed") })?; inner.state = State::Busy(blocking_task_join_handle); return Poll::Ready(Ok(n)); } State::Busy(ref mut rx) => { let (op, buf) = ready!(Pin::new(rx).poll(cx))?; inner.state = State::Idle(Some(buf)); match op { Operation::Read(_) => { // We don't care about the result here. The fact // that the cursor has advanced will be reflected in // the next iteration of the loop continue; } Operation::Write(res) => { // If the previous write was successful, continue. // Otherwise, error. res?; continue; } Operation::Seek(_) => { // Ignore the seek continue; } } } } } } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); if let Some(e) = inner.last_write_err.take() { return Poll::Ready(Err(e.into())); } loop { match inner.state { State::Idle(ref mut buf_cell) => { let mut buf = buf_cell.take().unwrap(); let seek = if !buf.is_empty() { Some(SeekFrom::Current(buf.discard_read())) } else { None }; let n = buf.copy_from_bufs(bufs, me.max_buf_size); let std = me.std.clone(); let blocking_task_join_handle = spawn_mandatory_blocking(move || { let res = if let Some(seek) = seek { (&*std).seek(seek).and_then(|_| buf.write_to(&mut &*std)) } else { buf.write_to(&mut &*std) }; (Operation::Write(res), buf) }) .ok_or_else(|| { io::Error::new(io::ErrorKind::Other, "background task failed") })?; inner.state = State::Busy(blocking_task_join_handle); return Poll::Ready(Ok(n)); } State::Busy(ref mut rx) => { let (op, buf) = ready!(Pin::new(rx).poll(cx))?; inner.state = State::Idle(Some(buf)); match op { Operation::Read(_) => { // We don't care about the result here. The fact // that the cursor has advanced will be reflected in // the next iteration of the loop continue; } Operation::Write(res) => { // If the previous write was successful, continue. // Otherwise, error. res?; continue; } Operation::Seek(_) => { // Ignore the seek continue; } } } } } } fn is_write_vectored(&self) -> bool { true } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let inner = self.inner.get_mut(); inner.poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); self.poll_flush(cx) } } impl From for File { fn from(std: StdFile) -> Self { Self::from_std(std) } } impl fmt::Debug for File { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("tokio::fs::File") .field("std", &self.std) .finish() } } #[cfg(unix)] impl std::os::unix::io::AsRawFd for File { fn as_raw_fd(&self) -> std::os::unix::io::RawFd { self.std.as_raw_fd() } } #[cfg(unix)] impl std::os::unix::io::AsFd for File { fn as_fd(&self) -> std::os::unix::io::BorrowedFd<'_> { unsafe { std::os::unix::io::BorrowedFd::borrow_raw(std::os::unix::io::AsRawFd::as_raw_fd(self)) } } } #[cfg(unix)] impl std::os::unix::io::FromRawFd for File { unsafe fn from_raw_fd(fd: std::os::unix::io::RawFd) -> Self { StdFile::from_raw_fd(fd).into() } } cfg_windows! { use crate::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle, AsHandle, BorrowedHandle}; impl AsRawHandle for File { fn as_raw_handle(&self) -> RawHandle { self.std.as_raw_handle() } } impl AsHandle for File { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw( AsRawHandle::as_raw_handle(self), ) } } } impl FromRawHandle for File { unsafe fn from_raw_handle(handle: RawHandle) -> Self { StdFile::from_raw_handle(handle).into() } } } impl Inner { async fn complete_inflight(&mut self) { use std::future::poll_fn; poll_fn(|cx| self.poll_complete_inflight(cx)).await; } fn poll_complete_inflight(&mut self, cx: &mut Context<'_>) -> Poll<()> { ready!(crate::trace::trace_leaf(cx)); match self.poll_flush(cx) { Poll::Ready(Err(e)) => { self.last_write_err = Some(e.kind()); Poll::Ready(()) } Poll::Ready(Ok(())) => Poll::Ready(()), Poll::Pending => Poll::Pending, } } fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll> { if let Some(e) = self.last_write_err.take() { return Poll::Ready(Err(e.into())); } let (op, buf) = match self.state { State::Idle(_) => return Poll::Ready(Ok(())), State::Busy(ref mut rx) => ready!(Pin::new(rx).poll(cx))?, }; // The buffer is not used here self.state = State::Idle(Some(buf)); match op { Operation::Read(_) => Poll::Ready(Ok(())), Operation::Write(res) => Poll::Ready(res), Operation::Seek(_) => Poll::Ready(Ok(())), } } } #[cfg(test)] mod tests; tokio-1.43.1/src/fs/hard_link.rs000064400000000000000000000024041046102023000145370ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Creates a new hard link on the filesystem. /// /// This is an async version of [`std::fs::hard_link`]. /// /// The `dst` path will be a link pointing to the `src` path. Note that systems /// often require these two paths to both be located on the same filesystem. /// /// # Platform-specific behavior /// /// This function currently corresponds to the `link` function on Unix /// and the `CreateHardLink` function on Windows. /// Note that, this [may change in the future][changes]. /// /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior /// /// # Errors /// /// This function will return an error in the following situations, but is not /// limited to just these cases: /// /// * The `src` path is not a file or doesn't exist. /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { /// fs::hard_link("a.txt", "b.txt").await?; // Hard link a.txt to b.txt /// Ok(()) /// } /// ``` pub async fn hard_link(src: impl AsRef, dst: impl AsRef) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); asyncify(move || std::fs::hard_link(src, dst)).await } tokio-1.43.1/src/fs/metadata.rs000064400000000000000000000024311046102023000143640ustar 00000000000000use crate::fs::asyncify; use std::fs::Metadata; use std::io; use std::path::Path; /// Given a path, queries the file system to get information about a file, /// directory, etc. /// /// This is an async version of [`std::fs::metadata`]. /// /// This function will traverse symbolic links to query information about the /// destination file. /// /// # Platform-specific behavior /// /// This function currently corresponds to the `stat` function on Unix and the /// `GetFileAttributesEx` function on Windows. Note that, this [may change in /// the future][changes]. /// /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior /// /// # Errors /// /// This function will return an error in the following situations, but is not /// limited to just these cases: /// /// * The user lacks permissions to perform `metadata` call on `path`. /// * `path` does not exist. /// /// # Examples /// /// ```rust,no_run /// use tokio::fs; /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { /// let attr = fs::metadata("/some/file/path.txt").await?; /// // inspect attr ... /// Ok(()) /// } /// ``` pub async fn metadata(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(|| std::fs::metadata(path)).await } tokio-1.43.1/src/fs/mocks.rs000064400000000000000000000102521046102023000137200ustar 00000000000000//! Mock version of std::fs::File; use mockall::mock; use crate::sync::oneshot; use std::{ cell::RefCell, collections::VecDeque, fs::{Metadata, Permissions}, future::Future, io::{self, Read, Seek, SeekFrom, Write}, path::PathBuf, pin::Pin, task::{Context, Poll}, }; mock! { #[derive(Debug)] pub File { pub fn create(pb: PathBuf) -> io::Result; // These inner_ methods exist because std::fs::File has two // implementations for each of these methods: one on "&mut self" and // one on "&&self". Defining both of those in terms of an inner_ method // allows us to specify the expectation the same way, regardless of // which method is used. pub fn inner_flush(&self) -> io::Result<()>; pub fn inner_read(&self, dst: &mut [u8]) -> io::Result; pub fn inner_seek(&self, pos: SeekFrom) -> io::Result; pub fn inner_write(&self, src: &[u8]) -> io::Result; pub fn metadata(&self) -> io::Result; pub fn open(pb: PathBuf) -> io::Result; pub fn set_len(&self, size: u64) -> io::Result<()>; pub fn set_permissions(&self, _perm: Permissions) -> io::Result<()>; pub fn set_max_buf_size(&self, max_buf_size: usize); pub fn sync_all(&self) -> io::Result<()>; pub fn sync_data(&self) -> io::Result<()>; pub fn try_clone(&self) -> io::Result; } #[cfg(windows)] impl std::os::windows::io::AsRawHandle for File { fn as_raw_handle(&self) -> std::os::windows::io::RawHandle; } #[cfg(windows)] impl std::os::windows::io::FromRawHandle for File { unsafe fn from_raw_handle(h: std::os::windows::io::RawHandle) -> Self; } #[cfg(unix)] impl std::os::unix::io::AsRawFd for File { fn as_raw_fd(&self) -> std::os::unix::io::RawFd; } #[cfg(unix)] impl std::os::unix::io::FromRawFd for File { unsafe fn from_raw_fd(h: std::os::unix::io::RawFd) -> Self; } } impl Read for MockFile { fn read(&mut self, dst: &mut [u8]) -> io::Result { self.inner_read(dst) } } impl Read for &'_ MockFile { fn read(&mut self, dst: &mut [u8]) -> io::Result { self.inner_read(dst) } } impl Seek for &'_ MockFile { fn seek(&mut self, pos: SeekFrom) -> io::Result { self.inner_seek(pos) } } impl Write for &'_ MockFile { fn write(&mut self, src: &[u8]) -> io::Result { self.inner_write(src) } fn flush(&mut self) -> io::Result<()> { self.inner_flush() } } tokio_thread_local! { static QUEUE: RefCell>> = RefCell::new(VecDeque::new()) } #[derive(Debug)] pub(super) struct JoinHandle { rx: oneshot::Receiver, } pub(super) fn spawn_blocking(f: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let (tx, rx) = oneshot::channel(); let task = Box::new(move || { let _ = tx.send(f()); }); QUEUE.with(|cell| cell.borrow_mut().push_back(task)); JoinHandle { rx } } pub(super) fn spawn_mandatory_blocking(f: F) -> Option> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let (tx, rx) = oneshot::channel(); let task = Box::new(move || { let _ = tx.send(f()); }); QUEUE.with(|cell| cell.borrow_mut().push_back(task)); Some(JoinHandle { rx }) } impl Future for JoinHandle { type Output = Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { use std::task::Poll; match Pin::new(&mut self.rx).poll(cx) { Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)), Poll::Ready(Err(e)) => panic!("error = {e:?}"), Poll::Pending => Poll::Pending, } } } pub(super) mod pool { use super::*; pub(in super::super) fn len() -> usize { QUEUE.with(|cell| cell.borrow().len()) } pub(in super::super) fn run_one() { let task = QUEUE .with(|cell| cell.borrow_mut().pop_front()) .expect("expected task to run, but none ready"); task(); } } tokio-1.43.1/src/fs/mod.rs000064400000000000000000000220021046102023000133570ustar 00000000000000#![cfg(not(loom))] //! Asynchronous file utilities. //! //! This module contains utility methods for working with the file system //! asynchronously. This includes reading/writing to files, and working with //! directories. //! //! Be aware that most operating systems do not provide asynchronous file system //! APIs. Because of that, Tokio will use ordinary blocking file operations //! behind the scenes. This is done using the [`spawn_blocking`] threadpool to //! run them in the background. //! //! The `tokio::fs` module should only be used for ordinary files. Trying to use //! it with e.g., a named pipe on Linux can result in surprising behavior, //! such as hangs during runtime shutdown. For special files, you should use a //! dedicated type such as [`tokio::net::unix::pipe`] or [`AsyncFd`] instead. //! //! Currently, Tokio will always use [`spawn_blocking`] on all platforms, but it //! may be changed to use asynchronous file system APIs such as io_uring in the //! future. //! //! # Usage //! //! The easiest way to use this module is to use the utility functions that //! operate on entire files: //! //! * [`tokio::fs::read`](fn@crate::fs::read) //! * [`tokio::fs::read_to_string`](fn@crate::fs::read_to_string) //! * [`tokio::fs::write`](fn@crate::fs::write) //! //! The two `read` functions reads the entire file and returns its contents. //! The `write` function takes the contents of the file and writes those //! contents to the file. It overwrites the existing file, if any. //! //! For example, to read the file: //! //! ``` //! # async fn dox() -> std::io::Result<()> { //! let contents = tokio::fs::read_to_string("my_file.txt").await?; //! //! println!("File has {} lines.", contents.lines().count()); //! # Ok(()) //! # } //! ``` //! //! To overwrite the file: //! //! ``` //! # async fn dox() -> std::io::Result<()> { //! let contents = "First line.\nSecond line.\nThird line.\n"; //! //! tokio::fs::write("my_file.txt", contents.as_bytes()).await?; //! # Ok(()) //! # } //! ``` //! //! ## Using `File` //! //! The main type for interacting with files is [`File`]. It can be used to read //! from and write to a given file. This is done using the [`AsyncRead`] and //! [`AsyncWrite`] traits. This type is generally used when you want to do //! something more complex than just reading or writing the entire contents in //! one go. //! //! **Note:** It is important to use [`flush`] when writing to a Tokio //! [`File`]. This is because calls to `write` will return before the write has //! finished, and [`flush`] will wait for the write to finish. (The write will //! happen even if you don't flush; it will just happen later.) This is //! different from [`std::fs::File`], and is due to the fact that `File` uses //! `spawn_blocking` behind the scenes. //! //! For example, to count the number of lines in a file without loading the //! entire file into memory: //! //! ```no_run //! use tokio::fs::File; //! use tokio::io::AsyncReadExt; //! //! # async fn dox() -> std::io::Result<()> { //! let mut file = File::open("my_file.txt").await?; //! //! let mut chunk = vec![0; 4096]; //! let mut number_of_lines = 0; //! loop { //! let len = file.read(&mut chunk).await?; //! if len == 0 { //! // Length of zero means end of file. //! break; //! } //! for &b in &chunk[..len] { //! if b == b'\n' { //! number_of_lines += 1; //! } //! } //! } //! //! println!("File has {} lines.", number_of_lines); //! # Ok(()) //! # } //! ``` //! //! For example, to write a file line-by-line: //! //! ```no_run //! use tokio::fs::File; //! use tokio::io::AsyncWriteExt; //! //! # async fn dox() -> std::io::Result<()> { //! let mut file = File::create("my_file.txt").await?; //! //! file.write_all(b"First line.\n").await?; //! file.write_all(b"Second line.\n").await?; //! file.write_all(b"Third line.\n").await?; //! //! // Remember to call `flush` after writing! //! file.flush().await?; //! # Ok(()) //! # } //! ``` //! //! ## Tuning your file IO //! //! Tokio's file uses [`spawn_blocking`] behind the scenes, and this has serious //! performance consequences. To get good performance with file IO on Tokio, it //! is recommended to batch your operations into as few `spawn_blocking` calls //! as possible. //! //! One example of this difference can be seen by comparing the two reading //! examples above. The first example uses [`tokio::fs::read`], which reads the //! entire file in a single `spawn_blocking` call, and then returns it. The //! second example will read the file in chunks using many `spawn_blocking` //! calls. This means that the second example will most likely be more expensive //! for large files. (Of course, using chunks may be necessary for very large //! files that don't fit in memory.) //! //! The following examples will show some strategies for this: //! //! When creating a file, write the data to a `String` or `Vec` and then //! write the entire file in a single `spawn_blocking` call with //! `tokio::fs::write`. //! //! ```no_run //! # async fn dox() -> std::io::Result<()> { //! let mut contents = String::new(); //! //! contents.push_str("First line.\n"); //! contents.push_str("Second line.\n"); //! contents.push_str("Third line.\n"); //! //! tokio::fs::write("my_file.txt", contents.as_bytes()).await?; //! # Ok(()) //! # } //! ``` //! //! Use [`BufReader`] and [`BufWriter`] to buffer many small reads or writes //! into a few large ones. This example will most likely only perform one //! `spawn_blocking` call. //! //! ```no_run //! use tokio::fs::File; //! use tokio::io::{AsyncWriteExt, BufWriter}; //! //! # async fn dox() -> std::io::Result<()> { //! let mut file = BufWriter::new(File::create("my_file.txt").await?); //! //! file.write_all(b"First line.\n").await?; //! file.write_all(b"Second line.\n").await?; //! file.write_all(b"Third line.\n").await?; //! //! // Due to the BufWriter, the actual write and spawn_blocking //! // call happens when you flush. //! file.flush().await?; //! # Ok(()) //! # } //! ``` //! //! Manually use [`std::fs`] inside [`spawn_blocking`]. //! //! ```no_run //! use std::fs::File; //! use std::io::{self, Write}; //! use tokio::task::spawn_blocking; //! //! # async fn dox() -> std::io::Result<()> { //! spawn_blocking(move || { //! let mut file = File::create("my_file.txt")?; //! //! file.write_all(b"First line.\n")?; //! file.write_all(b"Second line.\n")?; //! file.write_all(b"Third line.\n")?; //! //! // Unlike Tokio's file, the std::fs file does //! // not need flush. //! //! io::Result::Ok(()) //! }).await.unwrap()?; //! # Ok(()) //! # } //! ``` //! //! It's also good to be aware of [`File::set_max_buf_size`], which controls the //! maximum amount of bytes that Tokio's [`File`] will read or write in a single //! [`spawn_blocking`] call. The default is two megabytes, but this is subject //! to change. //! //! [`spawn_blocking`]: fn@crate::task::spawn_blocking //! [`AsyncRead`]: trait@crate::io::AsyncRead //! [`AsyncWrite`]: trait@crate::io::AsyncWrite //! [`BufReader`]: struct@crate::io::BufReader //! [`BufWriter`]: struct@crate::io::BufWriter //! [`tokio::net::unix::pipe`]: crate::net::unix::pipe //! [`AsyncFd`]: crate::io::unix::AsyncFd //! [`flush`]: crate::io::AsyncWriteExt::flush //! [`tokio::fs::read`]: fn@crate::fs::read mod canonicalize; pub use self::canonicalize::canonicalize; mod create_dir; pub use self::create_dir::create_dir; mod create_dir_all; pub use self::create_dir_all::create_dir_all; mod dir_builder; pub use self::dir_builder::DirBuilder; mod file; pub use self::file::File; mod hard_link; pub use self::hard_link::hard_link; mod metadata; pub use self::metadata::metadata; mod open_options; pub use self::open_options::OpenOptions; mod read; pub use self::read::read; mod read_dir; pub use self::read_dir::{read_dir, DirEntry, ReadDir}; mod read_link; pub use self::read_link::read_link; mod read_to_string; pub use self::read_to_string::read_to_string; mod remove_dir; pub use self::remove_dir::remove_dir; mod remove_dir_all; pub use self::remove_dir_all::remove_dir_all; mod remove_file; pub use self::remove_file::remove_file; mod rename; pub use self::rename::rename; mod set_permissions; pub use self::set_permissions::set_permissions; mod symlink_metadata; pub use self::symlink_metadata::symlink_metadata; mod write; pub use self::write::write; mod copy; pub use self::copy::copy; mod try_exists; pub use self::try_exists::try_exists; #[cfg(test)] mod mocks; feature! { #![unix] mod symlink; pub use self::symlink::symlink; } cfg_windows! { mod symlink_dir; pub use self::symlink_dir::symlink_dir; mod symlink_file; pub use self::symlink_file::symlink_file; } use std::io; #[cfg(not(test))] use crate::blocking::spawn_blocking; #[cfg(test)] use mocks::spawn_blocking; pub(crate) async fn asyncify(f: F) -> io::Result where F: FnOnce() -> io::Result + Send + 'static, T: Send + 'static, { match spawn_blocking(f).await { Ok(res) => res, Err(_) => Err(io::Error::new( io::ErrorKind::Other, "background task failed", )), } } tokio-1.43.1/src/fs/open_options/mock_open_options.rs000064400000000000000000000025701046102023000210510ustar 00000000000000#![allow(unreachable_pub)] //! Mock version of `std::fs::OpenOptions`; use mockall::mock; use crate::fs::mocks::MockFile; #[cfg(unix)] use std::os::unix::fs::OpenOptionsExt; #[cfg(windows)] use std::os::windows::fs::OpenOptionsExt; use std::{io, path::Path}; mock! { #[derive(Debug)] pub OpenOptions { pub fn append(&mut self, append: bool) -> &mut Self; pub fn create(&mut self, create: bool) -> &mut Self; pub fn create_new(&mut self, create_new: bool) -> &mut Self; pub fn open + 'static>(&self, path: P) -> io::Result; pub fn read(&mut self, read: bool) -> &mut Self; pub fn truncate(&mut self, truncate: bool) -> &mut Self; pub fn write(&mut self, write: bool) -> &mut Self; } impl Clone for OpenOptions { fn clone(&self) -> Self; } #[cfg(unix)] impl OpenOptionsExt for OpenOptions { fn custom_flags(&mut self, flags: i32) -> &mut Self; fn mode(&mut self, mode: u32) -> &mut Self; } #[cfg(windows)] impl OpenOptionsExt for OpenOptions { fn access_mode(&mut self, access: u32) -> &mut Self; fn share_mode(&mut self, val: u32) -> &mut Self; fn custom_flags(&mut self, flags: u32) -> &mut Self; fn attributes(&mut self, val: u32) -> &mut Self; fn security_qos_flags(&mut self, flags: u32) -> &mut Self; } } tokio-1.43.1/src/fs/open_options.rs000064400000000000000000000547611046102023000153350ustar 00000000000000use crate::fs::{asyncify, File}; use std::io; use std::path::Path; #[cfg(test)] mod mock_open_options; #[cfg(test)] use mock_open_options::MockOpenOptions as StdOpenOptions; #[cfg(not(test))] use std::fs::OpenOptions as StdOpenOptions; #[cfg(unix)] use std::os::unix::fs::OpenOptionsExt; #[cfg(windows)] use std::os::windows::fs::OpenOptionsExt; /// Options and flags which can be used to configure how a file is opened. /// /// This builder exposes the ability to configure how a [`File`] is opened and /// what operations are permitted on the open file. The [`File::open`] and /// [`File::create`] methods are aliases for commonly used options using this /// builder. /// /// Generally speaking, when using `OpenOptions`, you'll first call [`new`], /// then chain calls to methods to set each option, then call [`open`], passing /// the path of the file you're trying to open. This will give you a /// [`io::Result`] with a [`File`] inside that you can further operate /// on. /// /// This is a specialized version of [`std::fs::OpenOptions`] for usage from /// the Tokio runtime. /// /// `From` is implemented for more advanced configuration /// than the methods provided here. /// /// [`new`]: OpenOptions::new /// [`open`]: OpenOptions::open /// [`File`]: File /// [`File::open`]: File::open /// [`File::create`]: File::create /// /// # Examples /// /// Opening a file to read: /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .read(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` /// /// Opening a file for both reading and writing, as well as creating it if it /// doesn't exist: /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .read(true) /// .write(true) /// .create(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` #[derive(Clone, Debug)] pub struct OpenOptions(StdOpenOptions); impl OpenOptions { /// Creates a blank new set of options ready for configuration. /// /// All options are initially set to `false`. /// /// This is an async version of [`std::fs::OpenOptions::new`][std] /// /// [std]: std::fs::OpenOptions::new /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// /// let mut options = OpenOptions::new(); /// let future = options.read(true).open("foo.txt"); /// ``` pub fn new() -> OpenOptions { OpenOptions(StdOpenOptions::new()) } /// Sets the option for read access. /// /// This option, when true, will indicate that the file should be /// `read`-able if opened. /// /// This is an async version of [`std::fs::OpenOptions::read`][std] /// /// [std]: std::fs::OpenOptions::read /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .read(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` pub fn read(&mut self, read: bool) -> &mut OpenOptions { self.0.read(read); self } /// Sets the option for write access. /// /// This option, when true, will indicate that the file should be /// `write`-able if opened. /// /// This is an async version of [`std::fs::OpenOptions::write`][std] /// /// [std]: std::fs::OpenOptions::write /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .write(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` pub fn write(&mut self, write: bool) -> &mut OpenOptions { self.0.write(write); self } /// Sets the option for the append mode. /// /// This option, when true, means that writes will append to a file instead /// of overwriting previous contents. Note that setting /// `.write(true).append(true)` has the same effect as setting only /// `.append(true)`. /// /// For most filesystems, the operating system guarantees that all writes are /// atomic: no writes get mangled because another process writes at the same /// time. /// /// One maybe obvious note when using append-mode: make sure that all data /// that belongs together is written to the file in one operation. This /// can be done by concatenating strings before passing them to [`write()`], /// or using a buffered writer (with a buffer of adequate size), /// and calling [`flush()`] when the message is complete. /// /// If a file is opened with both read and append access, beware that after /// opening, and after every write, the position for reading may be set at the /// end of the file. So, before writing, save the current position (using /// [`seek`]`(`[`SeekFrom`]`::`[`Current`]`(0))`), and restore it before the next read. /// /// This is an async version of [`std::fs::OpenOptions::append`][std] /// /// [std]: std::fs::OpenOptions::append /// /// ## Note /// /// This function doesn't create the file if it doesn't exist. Use the [`create`] /// method to do so. /// /// [`write()`]: crate::io::AsyncWriteExt::write /// [`flush()`]: crate::io::AsyncWriteExt::flush /// [`seek`]: crate::io::AsyncSeekExt::seek /// [`SeekFrom`]: std::io::SeekFrom /// [`Current`]: std::io::SeekFrom::Current /// [`create`]: OpenOptions::create /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .append(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` pub fn append(&mut self, append: bool) -> &mut OpenOptions { self.0.append(append); self } /// Sets the option for truncating a previous file. /// /// If a file is successfully opened with this option set it will truncate /// the file to 0 length if it already exists. /// /// The file must be opened with write access for truncate to work. /// /// This is an async version of [`std::fs::OpenOptions::truncate`][std] /// /// [std]: std::fs::OpenOptions::truncate /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .write(true) /// .truncate(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` pub fn truncate(&mut self, truncate: bool) -> &mut OpenOptions { self.0.truncate(truncate); self } /// Sets the option for creating a new file. /// /// This option indicates whether a new file will be created if the file /// does not yet already exist. /// /// In order for the file to be created, [`write`] or [`append`] access must /// be used. /// /// This is an async version of [`std::fs::OpenOptions::create`][std] /// /// [std]: std::fs::OpenOptions::create /// [`write`]: OpenOptions::write /// [`append`]: OpenOptions::append /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .write(true) /// .create(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` pub fn create(&mut self, create: bool) -> &mut OpenOptions { self.0.create(create); self } /// Sets the option to always create a new file. /// /// This option indicates whether a new file will be created. No file is /// allowed to exist at the target location, also no (dangling) symlink. /// /// This option is useful because it is atomic. Otherwise between checking /// whether a file exists and creating a new one, the file may have been /// created by another process (a TOCTOU race condition / attack). /// /// If `.create_new(true)` is set, [`.create()`] and [`.truncate()`] are /// ignored. /// /// The file must be opened with write or append access in order to create a /// new file. /// /// This is an async version of [`std::fs::OpenOptions::create_new`][std] /// /// [std]: std::fs::OpenOptions::create_new /// [`.create()`]: OpenOptions::create /// [`.truncate()`]: OpenOptions::truncate /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new() /// .write(true) /// .create_new(true) /// .open("foo.txt") /// .await?; /// /// Ok(()) /// } /// ``` pub fn create_new(&mut self, create_new: bool) -> &mut OpenOptions { self.0.create_new(create_new); self } /// Opens a file at `path` with the options specified by `self`. /// /// This is an async version of [`std::fs::OpenOptions::open`][std] /// /// [std]: std::fs::OpenOptions::open /// /// # Errors /// /// This function will return an error under a number of different /// circumstances. Some of these error conditions are listed here, together /// with their [`ErrorKind`]. The mapping to [`ErrorKind`]s is not part of /// the compatibility contract of the function, especially the `Other` kind /// might change to more specific kinds in the future. /// /// * [`NotFound`]: The specified file does not exist and neither `create` /// or `create_new` is set. /// * [`NotFound`]: One of the directory components of the file path does /// not exist. /// * [`PermissionDenied`]: The user lacks permission to get the specified /// access rights for the file. /// * [`PermissionDenied`]: The user lacks permission to open one of the /// directory components of the specified path. /// * [`AlreadyExists`]: `create_new` was specified and the file already /// exists. /// * [`InvalidInput`]: Invalid combinations of open options (truncate /// without write access, no access mode set, etc.). /// * [`Other`]: One of the directory components of the specified file path /// was not, in fact, a directory. /// * [`Other`]: Filesystem-level errors: full disk, write permission /// requested on a read-only file system, exceeded disk quota, too many /// open files, too long filename, too many symbolic links in the /// specified path (Unix-like systems only), etc. /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let file = OpenOptions::new().open("foo.txt").await?; /// Ok(()) /// } /// ``` /// /// [`ErrorKind`]: std::io::ErrorKind /// [`AlreadyExists`]: std::io::ErrorKind::AlreadyExists /// [`InvalidInput`]: std::io::ErrorKind::InvalidInput /// [`NotFound`]: std::io::ErrorKind::NotFound /// [`Other`]: std::io::ErrorKind::Other /// [`PermissionDenied`]: std::io::ErrorKind::PermissionDenied pub async fn open(&self, path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); let opts = self.0.clone(); let std = asyncify(move || opts.open(path)).await?; Ok(File::from_std(std)) } /// Returns a mutable reference to the underlying `std::fs::OpenOptions` #[cfg(any(windows, unix))] pub(super) fn as_inner_mut(&mut self) -> &mut StdOpenOptions { &mut self.0 } } feature! { #![unix] impl OpenOptions { /// Sets the mode bits that a new file will be created with. /// /// If a new file is created as part of an `OpenOptions::open` call then this /// specified `mode` will be used as the permission bits for the new file. /// If no `mode` is set, the default of `0o666` will be used. /// The operating system masks out bits with the system's `umask`, to produce /// the final permissions. /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut options = OpenOptions::new(); /// options.mode(0o644); // Give read/write for owner and read for others. /// let file = options.open("foo.txt").await?; /// /// Ok(()) /// } /// ``` pub fn mode(&mut self, mode: u32) -> &mut OpenOptions { self.as_inner_mut().mode(mode); self } /// Passes custom flags to the `flags` argument of `open`. /// /// The bits that define the access mode are masked out with `O_ACCMODE`, to /// ensure they do not interfere with the access mode set by Rusts options. /// /// Custom flags can only set flags, not remove flags set by Rusts options. /// This options overwrites any previously set custom flags. /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut options = OpenOptions::new(); /// options.write(true); /// if cfg!(unix) { /// options.custom_flags(libc::O_NOFOLLOW); /// } /// let file = options.open("foo.txt").await?; /// /// Ok(()) /// } /// ``` pub fn custom_flags(&mut self, flags: i32) -> &mut OpenOptions { self.as_inner_mut().custom_flags(flags); self } } } cfg_windows! { impl OpenOptions { /// Overrides the `dwDesiredAccess` argument to the call to [`CreateFile`] /// with the specified value. /// /// This will override the `read`, `write`, and `append` flags on the /// `OpenOptions` structure. This method provides fine-grained control over /// the permissions to read, write and append data, attributes (like hidden /// and system), and extended attributes. /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] /// # async fn main() -> std::io::Result<()> { /// // Open without read and write permission, for example if you only need /// // to call `stat` on the file /// let file = OpenOptions::new().access_mode(0).open("foo.txt").await?; /// # Ok(()) /// # } /// ``` /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea pub fn access_mode(&mut self, access: u32) -> &mut OpenOptions { self.as_inner_mut().access_mode(access); self } /// Overrides the `dwShareMode` argument to the call to [`CreateFile`] with /// the specified value. /// /// By default `share_mode` is set to /// `FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE`. This allows /// other processes to read, write, and delete/rename the same file /// while it is open. Removing any of the flags will prevent other /// processes from performing the corresponding operation until the file /// handle is closed. /// /// # Examples /// /// ```no_run /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] /// # async fn main() -> std::io::Result<()> { /// // Do not allow others to read or modify this file while we have it open /// // for writing. /// let file = OpenOptions::new() /// .write(true) /// .share_mode(0) /// .open("foo.txt").await?; /// # Ok(()) /// # } /// ``` /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea pub fn share_mode(&mut self, share: u32) -> &mut OpenOptions { self.as_inner_mut().share_mode(share); self } /// Sets extra flags for the `dwFileFlags` argument to the call to /// [`CreateFile2`] to the specified value (or combines it with /// `attributes` and `security_qos_flags` to set the `dwFlagsAndAttributes` /// for [`CreateFile`]). /// /// Custom flags can only set flags, not remove flags set by Rust's options. /// This option overwrites any previously set custom flags. /// /// # Examples /// /// ```no_run /// use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_DELETE_ON_CLOSE; /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] /// # async fn main() -> std::io::Result<()> { /// let file = OpenOptions::new() /// .create(true) /// .write(true) /// .custom_flags(FILE_FLAG_DELETE_ON_CLOSE) /// .open("foo.txt").await?; /// # Ok(()) /// # } /// ``` /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea /// [`CreateFile2`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfile2 pub fn custom_flags(&mut self, flags: u32) -> &mut OpenOptions { self.as_inner_mut().custom_flags(flags); self } /// Sets the `dwFileAttributes` argument to the call to [`CreateFile2`] to /// the specified value (or combines it with `custom_flags` and /// `security_qos_flags` to set the `dwFlagsAndAttributes` for /// [`CreateFile`]). /// /// If a _new_ file is created because it does not yet exist and /// `.create(true)` or `.create_new(true)` are specified, the new file is /// given the attributes declared with `.attributes()`. /// /// If an _existing_ file is opened with `.create(true).truncate(true)`, its /// existing attributes are preserved and combined with the ones declared /// with `.attributes()`. /// /// In all other cases the attributes get ignored. /// /// # Examples /// /// ```no_run /// use windows_sys::Win32::Storage::FileSystem::FILE_ATTRIBUTE_HIDDEN; /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] /// # async fn main() -> std::io::Result<()> { /// let file = OpenOptions::new() /// .write(true) /// .create(true) /// .attributes(FILE_ATTRIBUTE_HIDDEN) /// .open("foo.txt").await?; /// # Ok(()) /// # } /// ``` /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea /// [`CreateFile2`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfile2 pub fn attributes(&mut self, attributes: u32) -> &mut OpenOptions { self.as_inner_mut().attributes(attributes); self } /// Sets the `dwSecurityQosFlags` argument to the call to [`CreateFile2`] to /// the specified value (or combines it with `custom_flags` and `attributes` /// to set the `dwFlagsAndAttributes` for [`CreateFile`]). /// /// By default `security_qos_flags` is not set. It should be specified when /// opening a named pipe, to control to which degree a server process can /// act on behalf of a client process (security impersonation level). /// /// When `security_qos_flags` is not set, a malicious program can gain the /// elevated privileges of a privileged Rust process when it allows opening /// user-specified paths, by tricking it into opening a named pipe. So /// arguably `security_qos_flags` should also be set when opening arbitrary /// paths. However the bits can then conflict with other flags, specifically /// `FILE_FLAG_OPEN_NO_RECALL`. /// /// For information about possible values, see [Impersonation Levels] on the /// Windows Dev Center site. The `SECURITY_SQOS_PRESENT` flag is set /// automatically when using this method. /// /// # Examples /// /// ```no_run /// use windows_sys::Win32::Storage::FileSystem::SECURITY_IDENTIFICATION; /// use tokio::fs::OpenOptions; /// /// # #[tokio::main] /// # async fn main() -> std::io::Result<()> { /// let file = OpenOptions::new() /// .write(true) /// .create(true) /// /// // Sets the flag value to `SecurityIdentification`. /// .security_qos_flags(SECURITY_IDENTIFICATION) /// /// .open(r"\\.\pipe\MyPipe").await?; /// # Ok(()) /// # } /// ``` /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea /// [`CreateFile2`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfile2 /// [Impersonation Levels]: /// https://docs.microsoft.com/en-us/windows/win32/api/winnt/ne-winnt-security_impersonation_level pub fn security_qos_flags(&mut self, flags: u32) -> &mut OpenOptions { self.as_inner_mut().security_qos_flags(flags); self } } } impl From for OpenOptions { fn from(options: StdOpenOptions) -> OpenOptions { OpenOptions(options) } } impl Default for OpenOptions { fn default() -> Self { Self::new() } } tokio-1.43.1/src/fs/read.rs000064400000000000000000000032211046102023000135150ustar 00000000000000use crate::fs::asyncify; use std::{io, path::Path}; /// Reads the entire contents of a file into a bytes vector. /// /// This is an async version of [`std::fs::read`]. /// /// This is a convenience function for using [`File::open`] and [`read_to_end`] /// with fewer imports and without an intermediate variable. It pre-allocates a /// buffer based on the file size when available, so it is generally faster than /// reading into a vector created with `Vec::new()`. /// /// This operation is implemented by running the equivalent blocking operation /// on a separate thread pool using [`spawn_blocking`]. /// /// [`File::open`]: super::File::open /// [`read_to_end`]: crate::io::AsyncReadExt::read_to_end /// [`spawn_blocking`]: crate::task::spawn_blocking /// /// # Errors /// /// This function will return an error if `path` does not already exist. /// Other errors may also be returned according to [`OpenOptions::open`]. /// /// [`OpenOptions::open`]: super::OpenOptions::open /// /// It will also return an error if it encounters while reading an error /// of a kind other than [`ErrorKind::Interrupted`]. /// /// [`ErrorKind::Interrupted`]: std::io::ErrorKind::Interrupted /// /// # Examples /// /// ```no_run /// use tokio::fs; /// use std::net::SocketAddr; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let contents = fs::read("address.txt").await?; /// let foo: SocketAddr = String::from_utf8_lossy(&contents).parse()?; /// Ok(()) /// } /// ``` pub async fn read(path: impl AsRef) -> io::Result> { let path = path.as_ref().to_owned(); asyncify(move || std::fs::read(path)).await } tokio-1.43.1/src/fs/read_dir.rs000064400000000000000000000254661046102023000143720ustar 00000000000000use crate::fs::asyncify; use std::collections::VecDeque; use std::ffi::OsString; use std::fs::{FileType, Metadata}; use std::future::Future; use std::io; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::Arc; use std::task::{ready, Context, Poll}; #[cfg(test)] use super::mocks::spawn_blocking; #[cfg(test)] use super::mocks::JoinHandle; #[cfg(not(test))] use crate::blocking::spawn_blocking; #[cfg(not(test))] use crate::blocking::JoinHandle; const CHUNK_SIZE: usize = 32; /// Returns a stream over the entries within a directory. /// /// This is an async version of [`std::fs::read_dir`]. /// /// This operation is implemented by running the equivalent blocking /// operation on a separate thread pool using [`spawn_blocking`]. /// /// [`spawn_blocking`]: crate::task::spawn_blocking pub async fn read_dir(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(|| -> io::Result { let mut std = std::fs::read_dir(path)?; let mut buf = VecDeque::with_capacity(CHUNK_SIZE); let remain = ReadDir::next_chunk(&mut buf, &mut std); Ok(ReadDir(State::Idle(Some((buf, std, remain))))) }) .await } /// Reads the entries in a directory. /// /// This struct is returned from the [`read_dir`] function of this module and /// will yield instances of [`DirEntry`]. Through a [`DirEntry`] information /// like the entry's path and possibly other metadata can be learned. /// /// A `ReadDir` can be turned into a `Stream` with [`ReadDirStream`]. /// /// [`ReadDirStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.ReadDirStream.html /// /// # Errors /// /// This stream will return an [`Err`] if there's some sort of intermittent /// IO error during iteration. /// /// [`read_dir`]: read_dir /// [`DirEntry`]: DirEntry /// [`Err`]: std::result::Result::Err #[derive(Debug)] #[must_use = "streams do nothing unless polled"] pub struct ReadDir(State); #[derive(Debug)] enum State { Idle(Option<(VecDeque>, std::fs::ReadDir, bool)>), Pending(JoinHandle<(VecDeque>, std::fs::ReadDir, bool)>), } impl ReadDir { /// Returns the next entry in the directory stream. /// /// # Cancel safety /// /// This method is cancellation safe. pub async fn next_entry(&mut self) -> io::Result> { use std::future::poll_fn; poll_fn(|cx| self.poll_next_entry(cx)).await } /// Polls for the next directory entry in the stream. /// /// This method returns: /// /// * `Poll::Pending` if the next directory entry is not yet available. /// * `Poll::Ready(Ok(Some(entry)))` if the next directory entry is available. /// * `Poll::Ready(Ok(None))` if there are no more directory entries in this /// stream. /// * `Poll::Ready(Err(err))` if an IO error occurred while reading the next /// directory entry. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when the next directory entry /// becomes available on the underlying IO resource. /// /// Note that on multiple calls to `poll_next_entry`, only the `Waker` from /// the `Context` passed to the most recent call is scheduled to receive a /// wakeup. pub fn poll_next_entry(&mut self, cx: &mut Context<'_>) -> Poll>> { loop { match self.0 { State::Idle(ref mut data) => { let (buf, _, ref remain) = data.as_mut().unwrap(); if let Some(ent) = buf.pop_front() { return Poll::Ready(ent.map(Some)); } else if !remain { return Poll::Ready(Ok(None)); } let (mut buf, mut std, _) = data.take().unwrap(); self.0 = State::Pending(spawn_blocking(move || { let remain = ReadDir::next_chunk(&mut buf, &mut std); (buf, std, remain) })); } State::Pending(ref mut rx) => { self.0 = State::Idle(Some(ready!(Pin::new(rx).poll(cx))?)); } } } } fn next_chunk(buf: &mut VecDeque>, std: &mut std::fs::ReadDir) -> bool { for _ in 0..CHUNK_SIZE { let ret = match std.next() { Some(ret) => ret, None => return false, }; let success = ret.is_ok(); buf.push_back(ret.map(|std| DirEntry { #[cfg(not(any( target_os = "solaris", target_os = "illumos", target_os = "haiku", target_os = "vxworks", target_os = "aix", target_os = "nto", target_os = "vita", )))] file_type: std.file_type().ok(), std: Arc::new(std), })); if !success { break; } } true } } feature! { #![unix] use std::os::unix::fs::DirEntryExt; impl DirEntry { /// Returns the underlying `d_ino` field in the contained `dirent` /// structure. /// /// # Examples /// /// ``` /// use tokio::fs; /// /// # #[tokio::main] /// # async fn main() -> std::io::Result<()> { /// let mut entries = fs::read_dir(".").await?; /// while let Some(entry) = entries.next_entry().await? { /// // Here, `entry` is a `DirEntry`. /// println!("{:?}: {}", entry.file_name(), entry.ino()); /// } /// # Ok(()) /// # } /// ``` pub fn ino(&self) -> u64 { self.as_inner().ino() } } } /// Entries returned by the [`ReadDir`] stream. /// /// [`ReadDir`]: struct@ReadDir /// /// This is a specialized version of [`std::fs::DirEntry`] for usage from the /// Tokio runtime. /// /// An instance of `DirEntry` represents an entry inside of a directory on the /// filesystem. Each entry can be inspected via methods to learn about the full /// path or possibly other metadata through per-platform extension traits. #[derive(Debug)] pub struct DirEntry { #[cfg(not(any( target_os = "solaris", target_os = "illumos", target_os = "haiku", target_os = "vxworks", target_os = "aix", target_os = "nto", target_os = "vita", )))] file_type: Option, std: Arc, } impl DirEntry { /// Returns the full path to the file that this entry represents. /// /// The full path is created by joining the original path to `read_dir` /// with the filename of this entry. /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// let mut entries = fs::read_dir(".").await?; /// /// while let Some(entry) = entries.next_entry().await? { /// println!("{:?}", entry.path()); /// } /// # Ok(()) /// # } /// ``` /// /// This prints output like: /// /// ```text /// "./whatever.txt" /// "./foo.html" /// "./hello_world.rs" /// ``` /// /// The exact text, of course, depends on what files you have in `.`. pub fn path(&self) -> PathBuf { self.std.path() } /// Returns the bare file name of this directory entry without any other /// leading path component. /// /// # Examples /// /// ``` /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// let mut entries = fs::read_dir(".").await?; /// /// while let Some(entry) = entries.next_entry().await? { /// println!("{:?}", entry.file_name()); /// } /// # Ok(()) /// # } /// ``` pub fn file_name(&self) -> OsString { self.std.file_name() } /// Returns the metadata for the file that this entry points at. /// /// This function will not traverse symlinks if this entry points at a /// symlink. /// /// # Platform-specific behavior /// /// On Windows this function is cheap to call (no extra system calls /// needed), but on Unix platforms this function is the equivalent of /// calling `symlink_metadata` on the path. /// /// # Examples /// /// ``` /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// let mut entries = fs::read_dir(".").await?; /// /// while let Some(entry) = entries.next_entry().await? { /// if let Ok(metadata) = entry.metadata().await { /// // Now let's show our entry's permissions! /// println!("{:?}: {:?}", entry.path(), metadata.permissions()); /// } else { /// println!("Couldn't get file type for {:?}", entry.path()); /// } /// } /// # Ok(()) /// # } /// ``` pub async fn metadata(&self) -> io::Result { let std = self.std.clone(); asyncify(move || std.metadata()).await } /// Returns the file type for the file that this entry points at. /// /// This function will not traverse symlinks if this entry points at a /// symlink. /// /// # Platform-specific behavior /// /// On Windows and most Unix platforms this function is free (no extra /// system calls needed), but some Unix platforms may require the equivalent /// call to `symlink_metadata` to learn about the target file type. /// /// # Examples /// /// ``` /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// let mut entries = fs::read_dir(".").await?; /// /// while let Some(entry) = entries.next_entry().await? { /// if let Ok(file_type) = entry.file_type().await { /// // Now let's show our entry's file type! /// println!("{:?}: {:?}", entry.path(), file_type); /// } else { /// println!("Couldn't get file type for {:?}", entry.path()); /// } /// } /// # Ok(()) /// # } /// ``` pub async fn file_type(&self) -> io::Result { #[cfg(not(any( target_os = "solaris", target_os = "illumos", target_os = "haiku", target_os = "vxworks", target_os = "aix", target_os = "nto", target_os = "vita", )))] if let Some(file_type) = self.file_type { return Ok(file_type); } let std = self.std.clone(); asyncify(move || std.file_type()).await } /// Returns a reference to the underlying `std::fs::DirEntry`. #[cfg(unix)] pub(super) fn as_inner(&self) -> &std::fs::DirEntry { &self.std } } tokio-1.43.1/src/fs/read_link.rs000064400000000000000000000005631046102023000145400ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::{Path, PathBuf}; /// Reads a symbolic link, returning the file that the link points to. /// /// This is an async version of [`std::fs::read_link`]. pub async fn read_link(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(move || std::fs::read_link(path)).await } tokio-1.43.1/src/fs/read_to_string.rs000064400000000000000000000016241046102023000156120ustar 00000000000000use crate::fs::asyncify; use std::{io, path::Path}; /// Creates a future which will open a file for reading and read the entire /// contents into a string and return said string. /// /// This is the async equivalent of [`std::fs::read_to_string`][std]. /// /// This operation is implemented by running the equivalent blocking operation /// on a separate thread pool using [`spawn_blocking`]. /// /// [`spawn_blocking`]: crate::task::spawn_blocking /// [std]: fn@std::fs::read_to_string /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// let contents = fs::read_to_string("foo.txt").await?; /// println!("foo.txt contains {} bytes", contents.len()); /// # Ok(()) /// # } /// ``` pub async fn read_to_string(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(move || std::fs::read_to_string(path)).await } tokio-1.43.1/src/fs/remove_dir.rs000064400000000000000000000005111046102023000147340ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Removes an existing, empty directory. /// /// This is an async version of [`std::fs::remove_dir`]. pub async fn remove_dir(path: impl AsRef) -> io::Result<()> { let path = path.as_ref().to_owned(); asyncify(move || std::fs::remove_dir(path)).await } tokio-1.43.1/src/fs/remove_dir_all.rs000064400000000000000000000006571046102023000155770ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Removes a directory at this path, after removing all its contents. Use carefully! /// /// This is an async version of [`std::fs::remove_dir_all`][std] /// /// [std]: fn@std::fs::remove_dir_all pub async fn remove_dir_all(path: impl AsRef) -> io::Result<()> { let path = path.as_ref().to_owned(); asyncify(move || std::fs::remove_dir_all(path)).await } tokio-1.43.1/src/fs/remove_file.rs000064400000000000000000000007701046102023000151040ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Removes a file from the filesystem. /// /// Note that there is no guarantee that the file is immediately deleted (e.g. /// depending on platform, other open file descriptors may prevent immediate /// removal). /// /// This is an async version of [`std::fs::remove_file`]. pub async fn remove_file(path: impl AsRef) -> io::Result<()> { let path = path.as_ref().to_owned(); asyncify(move || std::fs::remove_file(path)).await } tokio-1.43.1/src/fs/rename.rs000064400000000000000000000010041046102023000140460ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Renames a file or directory to a new name, replacing the original file if /// `to` already exists. /// /// This will not work if the new name is on a different mount point. /// /// This is an async version of [`std::fs::rename`]. pub async fn rename(from: impl AsRef, to: impl AsRef) -> io::Result<()> { let from = from.as_ref().to_owned(); let to = to.as_ref().to_owned(); asyncify(move || std::fs::rename(from, to)).await } tokio-1.43.1/src/fs/set_permissions.rs000064400000000000000000000007071046102023000160360ustar 00000000000000use crate::fs::asyncify; use std::fs::Permissions; use std::io; use std::path::Path; /// Changes the permissions found on a file or a directory. /// /// This is an async version of [`std::fs::set_permissions`][std] /// /// [std]: fn@std::fs::set_permissions pub async fn set_permissions(path: impl AsRef, perm: Permissions) -> io::Result<()> { let path = path.as_ref().to_owned(); asyncify(|| std::fs::set_permissions(path, perm)).await } tokio-1.43.1/src/fs/symlink.rs000064400000000000000000000007501046102023000142740ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Creates a new symbolic link on the filesystem. /// /// The `dst` path will be a symbolic link pointing to the `src` path. /// /// This is an async version of [`std::os::unix::fs::symlink`]. pub async fn symlink(src: impl AsRef, dst: impl AsRef) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); asyncify(move || std::os::unix::fs::symlink(src, dst)).await } tokio-1.43.1/src/fs/symlink_dir.rs000064400000000000000000000011371046102023000151320ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Creates a new directory symlink on the filesystem. /// /// The `dst` path will be a directory symbolic link pointing to the `src` /// path. /// /// This is an async version of [`std::os::windows::fs::symlink_dir`][std] /// /// [std]: https://doc.rust-lang.org/std/os/windows/fs/fn.symlink_dir.html pub async fn symlink_dir(src: impl AsRef, dst: impl AsRef) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); asyncify(move || std::os::windows::fs::symlink_dir(src, dst)).await } tokio-1.43.1/src/fs/symlink_file.rs000064400000000000000000000011371046102023000152730ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Creates a new file symbolic link on the filesystem. /// /// The `dst` path will be a file symbolic link pointing to the `src` /// path. /// /// This is an async version of [`std::os::windows::fs::symlink_file`][std] /// /// [std]: https://doc.rust-lang.org/std/os/windows/fs/fn.symlink_file.html pub async fn symlink_file(src: impl AsRef, dst: impl AsRef) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); asyncify(move || std::os::windows::fs::symlink_file(src, dst)).await } tokio-1.43.1/src/fs/symlink_metadata.rs000064400000000000000000000006521046102023000161350ustar 00000000000000use crate::fs::asyncify; use std::fs::Metadata; use std::io; use std::path::Path; /// Queries the file system metadata for a path. /// /// This is an async version of [`std::fs::symlink_metadata`][std] /// /// [std]: fn@std::fs::symlink_metadata pub async fn symlink_metadata(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(|| std::fs::symlink_metadata(path)).await } tokio-1.43.1/src/fs/try_exists.rs000064400000000000000000000013551046102023000150250ustar 00000000000000use crate::fs::asyncify; use std::io; use std::path::Path; /// Returns `Ok(true)` if the path points at an existing entity. /// /// This function will traverse symbolic links to query information about the /// destination file. In case of broken symbolic links this will return `Ok(false)`. /// /// This is the async equivalent of [`std::path::Path::try_exists`][std]. /// /// [std]: fn@std::path::Path::try_exists /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// fs::try_exists("foo.txt").await?; /// # Ok(()) /// # } /// ``` pub async fn try_exists(path: impl AsRef) -> io::Result { let path = path.as_ref().to_owned(); asyncify(move || path.try_exists()).await } tokio-1.43.1/src/fs/write.rs000064400000000000000000000015611046102023000137410ustar 00000000000000use crate::fs::asyncify; use std::{io, path::Path}; /// Creates a future that will open a file for writing and write the entire /// contents of `contents` to it. /// /// This is the async equivalent of [`std::fs::write`][std]. /// /// This operation is implemented by running the equivalent blocking operation /// on a separate thread pool using [`spawn_blocking`]. /// /// [`spawn_blocking`]: crate::task::spawn_blocking /// [std]: fn@std::fs::write /// /// # Examples /// /// ```no_run /// use tokio::fs; /// /// # async fn dox() -> std::io::Result<()> { /// fs::write("foo.txt", b"Hello world!").await?; /// # Ok(()) /// # } /// ``` pub async fn write(path: impl AsRef, contents: impl AsRef<[u8]>) -> io::Result<()> { let path = path.as_ref().to_owned(); let contents = contents.as_ref().to_owned(); asyncify(move || std::fs::write(path, contents)).await } tokio-1.43.1/src/future/block_on.rs000064400000000000000000000012651046102023000153000ustar 00000000000000use std::future::Future; cfg_rt! { #[track_caller] pub(crate) fn block_on(f: F) -> F::Output { let mut e = crate::runtime::context::try_enter_blocking_region().expect( "Cannot block the current thread from within a runtime. This \ happens because a function attempted to block the current \ thread while the thread is being used to drive asynchronous \ tasks." ); e.block_on(f).unwrap() } } cfg_not_rt! { #[track_caller] pub(crate) fn block_on(f: F) -> F::Output { let mut park = crate::runtime::park::CachedParkThread::new(); park.block_on(f).unwrap() } } tokio-1.43.1/src/future/maybe_done.rs000064400000000000000000000073431046102023000156170ustar 00000000000000//! Definition of the [`MaybeDone`] combinator. use pin_project_lite::pin_project; use std::future::{Future, IntoFuture}; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// A future that may have completed. #[derive(Debug)] #[project = MaybeDoneProj] #[project_replace = MaybeDoneProjReplace] #[repr(C)] // https://github.com/rust-lang/miri/issues/3780 pub enum MaybeDone { /// A not-yet-completed future. Future { #[pin] future: Fut }, /// The output of the completed future. Done { output: Fut::Output }, /// The empty variant after the result of a [`MaybeDone`] has been /// taken using the [`take_output`](MaybeDone::take_output) method. Gone, } } /// Wraps a future into a `MaybeDone`. pub fn maybe_done(future: F) -> MaybeDone { MaybeDone::Future { future: future.into_future(), } } impl MaybeDone { /// Returns an [`Option`] containing a mutable reference to the output of the future. /// The output of this method will be [`Some`] if and only if the inner /// future has been completed and [`take_output`](MaybeDone::take_output) /// has not yet been called. pub fn output_mut(self: Pin<&mut Self>) -> Option<&mut Fut::Output> { match self.project() { MaybeDoneProj::Done { output } => Some(output), _ => None, } } /// Attempts to take the output of a `MaybeDone` without driving it /// towards completion. #[inline] pub fn take_output(self: Pin<&mut Self>) -> Option { match *self { MaybeDone::Done { .. } => {} MaybeDone::Future { .. } | MaybeDone::Gone => return None, }; if let MaybeDoneProjReplace::Done { output } = self.project_replace(MaybeDone::Gone) { Some(output) } else { unreachable!() } } } impl Future for MaybeDone { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let output = match self.as_mut().project() { MaybeDoneProj::Future { future } => ready!(future.poll(cx)), MaybeDoneProj::Done { .. } => return Poll::Ready(()), MaybeDoneProj::Gone => panic!("MaybeDone polled after value taken"), }; self.set(MaybeDone::Done { output }); Poll::Ready(()) } } // Test for https://github.com/tokio-rs/tokio/issues/6729 #[cfg(test)] mod miri_tests { use super::maybe_done; use std::{ future::Future, pin::Pin, sync::Arc, task::{Context, Poll, Wake}, }; struct ThingAdder<'a> { thing: &'a mut String, } impl Future for ThingAdder<'_> { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { unsafe { *self.get_unchecked_mut().thing += ", world"; } Poll::Pending } } #[test] fn maybe_done_miri() { let mut thing = "hello".to_owned(); // The async block is necessary to trigger the miri failure. #[allow(clippy::redundant_async_block)] let fut = async move { ThingAdder { thing: &mut thing }.await }; let mut fut = maybe_done(fut); let mut fut = unsafe { Pin::new_unchecked(&mut fut) }; let waker = Arc::new(DummyWaker).into(); let mut ctx = Context::from_waker(&waker); assert_eq!(fut.as_mut().poll(&mut ctx), Poll::Pending); assert_eq!(fut.as_mut().poll(&mut ctx), Poll::Pending); } struct DummyWaker; impl Wake for DummyWaker { fn wake(self: Arc) {} } } tokio-1.43.1/src/future/mod.rs000064400000000000000000000010051046102023000142610ustar 00000000000000#![cfg_attr(not(feature = "macros"), allow(unreachable_pub))] //! Asynchronous values. #[cfg(any(feature = "macros", feature = "process"))] pub(crate) mod maybe_done; cfg_process! { mod try_join; pub(crate) use try_join::try_join3; } cfg_sync! { mod block_on; pub(crate) use block_on::block_on; } cfg_trace! { mod trace; #[allow(unused_imports)] pub(crate) use trace::InstrumentedFuture as Future; } cfg_not_trace! { cfg_rt! { pub(crate) use std::future::Future; } } tokio-1.43.1/src/future/trace.rs000064400000000000000000000004151046102023000146040ustar 00000000000000use std::future::Future; pub(crate) trait InstrumentedFuture: Future { fn id(&self) -> Option; } impl InstrumentedFuture for tracing::instrument::Instrumented { fn id(&self) -> Option { self.span().id() } } tokio-1.43.1/src/future/try_join.rs000064400000000000000000000044211046102023000153440ustar 00000000000000use crate::future::maybe_done::{maybe_done, MaybeDone}; use pin_project_lite::pin_project; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; pub(crate) fn try_join3( future1: F1, future2: F2, future3: F3, ) -> TryJoin3 where F1: Future>, F2: Future>, F3: Future>, { TryJoin3 { future1: maybe_done(future1), future2: maybe_done(future2), future3: maybe_done(future3), } } pin_project! { pub(crate) struct TryJoin3 where F1: Future, F2: Future, F3: Future, { #[pin] future1: MaybeDone, #[pin] future2: MaybeDone, #[pin] future3: MaybeDone, } } impl Future for TryJoin3 where F1: Future>, F2: Future>, F3: Future>, { type Output = Result<(T1, T2, T3), E>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut all_done = true; let mut me = self.project(); if me.future1.as_mut().poll(cx).is_pending() { all_done = false; } else if me.future1.as_mut().output_mut().unwrap().is_err() { return Poll::Ready(Err(me.future1.take_output().unwrap().err().unwrap())); } if me.future2.as_mut().poll(cx).is_pending() { all_done = false; } else if me.future2.as_mut().output_mut().unwrap().is_err() { return Poll::Ready(Err(me.future2.take_output().unwrap().err().unwrap())); } if me.future3.as_mut().poll(cx).is_pending() { all_done = false; } else if me.future3.as_mut().output_mut().unwrap().is_err() { return Poll::Ready(Err(me.future3.take_output().unwrap().err().unwrap())); } if all_done { Poll::Ready(Ok(( me.future1.take_output().unwrap().ok().unwrap(), me.future2.take_output().unwrap().ok().unwrap(), me.future3.take_output().unwrap().ok().unwrap(), ))) } else { Poll::Pending } } } tokio-1.43.1/src/fuzz.rs000064400000000000000000000000731046102023000131720ustar 00000000000000pub use crate::util::linked_list::tests::fuzz_linked_list; tokio-1.43.1/src/io/async_buf_read.rs000064400000000000000000000103711046102023000155510ustar 00000000000000use crate::io::AsyncRead; use std::io; use std::ops::DerefMut; use std::pin::Pin; use std::task::{Context, Poll}; /// Reads bytes asynchronously. /// /// This trait is analogous to [`std::io::BufRead`], but integrates with /// the asynchronous task system. In particular, the [`poll_fill_buf`] method, /// unlike [`BufRead::fill_buf`], will automatically queue the current task for wakeup /// and return if data is not yet available, rather than blocking the calling /// thread. /// /// Utilities for working with `AsyncBufRead` values are provided by /// [`AsyncBufReadExt`]. /// /// [`std::io::BufRead`]: std::io::BufRead /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf /// [`BufRead::fill_buf`]: std::io::BufRead::fill_buf /// [`AsyncBufReadExt`]: crate::io::AsyncBufReadExt pub trait AsyncBufRead: AsyncRead { /// Attempts to return the contents of the internal buffer, filling it with more data /// from the inner reader if it is empty. /// /// On success, returns `Poll::Ready(Ok(buf))`. /// /// If no data is available for reading, the method returns /// `Poll::Pending` and arranges for the current task (via /// `cx.waker().wake_by_ref()`) to receive a notification when the object becomes /// readable or is closed. /// /// This function is a lower-level call. It needs to be paired with the /// [`consume`] method to function properly. When calling this /// method, none of the contents will be "read" in the sense that later /// calling [`poll_read`] may return the same contents. As such, [`consume`] must /// be called with the number of bytes that are consumed from this buffer to /// ensure that the bytes are never returned twice. /// /// An empty buffer returned indicates that the stream has reached EOF. /// /// [`poll_read`]: AsyncRead::poll_read /// [`consume`]: AsyncBufRead::consume fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll>; /// Tells this buffer that `amt` bytes have been consumed from the buffer, /// so they should no longer be returned in calls to [`poll_read`]. /// /// This function is a lower-level call. It needs to be paired with the /// [`poll_fill_buf`] method to function properly. This function does /// not perform any I/O, it simply informs this object that some amount of /// its buffer, returned from [`poll_fill_buf`], has been consumed and should /// no longer be returned. As such, this function may do odd things if /// [`poll_fill_buf`] isn't called before calling it. /// /// The `amt` must be `<=` the number of bytes in the buffer returned by /// [`poll_fill_buf`]. /// /// [`poll_read`]: AsyncRead::poll_read /// [`poll_fill_buf`]: AsyncBufRead::poll_fill_buf fn consume(self: Pin<&mut Self>, amt: usize); } macro_rules! deref_async_buf_read { () => { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut **self.get_mut()).poll_fill_buf(cx) } fn consume(mut self: Pin<&mut Self>, amt: usize) { Pin::new(&mut **self).consume(amt) } }; } impl AsyncBufRead for Box { deref_async_buf_read!(); } impl AsyncBufRead for &mut T { deref_async_buf_read!(); } impl

AsyncBufRead for Pin

where P: DerefMut + Unpin, P::Target: AsyncBufRead, { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_mut().as_mut().poll_fill_buf(cx) } fn consume(self: Pin<&mut Self>, amt: usize) { self.get_mut().as_mut().consume(amt); } } impl AsyncBufRead for &[u8] { fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(*self)) } fn consume(mut self: Pin<&mut Self>, amt: usize) { *self = &self[amt..]; } } impl + Unpin> AsyncBufRead for io::Cursor { fn poll_fill_buf(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(io::BufRead::fill_buf(self.get_mut())) } fn consume(self: Pin<&mut Self>, amt: usize) { io::BufRead::consume(self.get_mut(), amt); } } tokio-1.43.1/src/io/async_fd.rs000064400000000000000000001631171046102023000144020ustar 00000000000000use crate::io::{Interest, Ready}; use crate::runtime::io::{ReadyEvent, Registration}; use crate::runtime::scheduler; use mio::unix::SourceFd; use std::error::Error; use std::fmt; use std::io; use std::os::unix::io::{AsRawFd, RawFd}; use std::task::{ready, Context, Poll}; /// Associates an IO object backed by a Unix file descriptor with the tokio /// reactor, allowing for readiness to be polled. The file descriptor must be of /// a type that can be used with the OS polling facilities (ie, `poll`, `epoll`, /// `kqueue`, etc), such as a network socket or pipe, and the file descriptor /// must have the nonblocking mode set to true. /// /// Creating an [`AsyncFd`] registers the file descriptor with the current tokio /// Reactor, allowing you to directly await the file descriptor being readable /// or writable. Once registered, the file descriptor remains registered until /// the [`AsyncFd`] is dropped. /// /// The [`AsyncFd`] takes ownership of an arbitrary object to represent the IO /// object. It is intended that the inner object will handle closing the file /// descriptor when it is dropped, avoiding resource leaks and ensuring that the /// [`AsyncFd`] can clean up the registration before closing the file descriptor. /// The [`AsyncFd::into_inner`] function can be used to extract the inner object /// to retake control from the tokio IO reactor. The [`OwnedFd`] type is often /// used as the inner object, as it is the simplest type that closes the fd on /// drop. /// /// The inner object is required to implement [`AsRawFd`]. This file descriptor /// must not change while [`AsyncFd`] owns the inner object, i.e. the /// [`AsRawFd::as_raw_fd`] method on the inner type must always return the same /// file descriptor when called multiple times. Failure to uphold this results /// in unspecified behavior in the IO driver, which may include breaking /// notifications for other sockets/etc. /// /// Polling for readiness is done by calling the async functions [`readable`] /// and [`writable`]. These functions complete when the associated readiness /// condition is observed. Any number of tasks can query the same `AsyncFd` in /// parallel, on the same or different conditions. /// /// On some platforms, the readiness detecting mechanism relies on /// edge-triggered notifications. This means that the OS will only notify Tokio /// when the file descriptor transitions from not-ready to ready. For this to /// work you should first try to read or write and only poll for readiness /// if that fails with an error of [`std::io::ErrorKind::WouldBlock`]. /// /// Tokio internally tracks when it has received a ready notification, and when /// readiness checking functions like [`readable`] and [`writable`] are called, /// if the readiness flag is set, these async functions will complete /// immediately. This however does mean that it is critical to ensure that this /// ready flag is cleared when (and only when) the file descriptor ceases to be /// ready. The [`AsyncFdReadyGuard`] returned from readiness checking functions /// serves this function; after calling a readiness-checking async function, /// you must use this [`AsyncFdReadyGuard`] to signal to tokio whether the file /// descriptor is no longer in a ready state. /// /// ## Use with to a poll-based API /// /// In some cases it may be desirable to use `AsyncFd` from APIs similar to /// [`TcpStream::poll_read_ready`]. The [`AsyncFd::poll_read_ready`] and /// [`AsyncFd::poll_write_ready`] functions are provided for this purpose. /// Because these functions don't create a future to hold their state, they have /// the limitation that only one task can wait on each direction (read or write) /// at a time. /// /// # Examples /// /// This example shows how to turn [`std::net::TcpStream`] asynchronous using /// `AsyncFd`. It implements the read/write operations both as an `async fn` /// and using the IO traits [`AsyncRead`] and [`AsyncWrite`]. /// /// ```no_run /// use std::io::{self, Read, Write}; /// use std::net::TcpStream; /// use std::pin::Pin; /// use std::task::{ready, Context, Poll}; /// use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; /// use tokio::io::unix::AsyncFd; /// /// pub struct AsyncTcpStream { /// inner: AsyncFd, /// } /// /// impl AsyncTcpStream { /// pub fn new(tcp: TcpStream) -> io::Result { /// tcp.set_nonblocking(true)?; /// Ok(Self { /// inner: AsyncFd::new(tcp)?, /// }) /// } /// /// pub async fn read(&self, out: &mut [u8]) -> io::Result { /// loop { /// let mut guard = self.inner.readable().await?; /// /// match guard.try_io(|inner| inner.get_ref().read(out)) { /// Ok(result) => return result, /// Err(_would_block) => continue, /// } /// } /// } /// /// pub async fn write(&self, buf: &[u8]) -> io::Result { /// loop { /// let mut guard = self.inner.writable().await?; /// /// match guard.try_io(|inner| inner.get_ref().write(buf)) { /// Ok(result) => return result, /// Err(_would_block) => continue, /// } /// } /// } /// } /// /// impl AsyncRead for AsyncTcpStream { /// fn poll_read( /// self: Pin<&mut Self>, /// cx: &mut Context<'_>, /// buf: &mut ReadBuf<'_> /// ) -> Poll> { /// loop { /// let mut guard = ready!(self.inner.poll_read_ready(cx))?; /// /// let unfilled = buf.initialize_unfilled(); /// match guard.try_io(|inner| inner.get_ref().read(unfilled)) { /// Ok(Ok(len)) => { /// buf.advance(len); /// return Poll::Ready(Ok(())); /// }, /// Ok(Err(err)) => return Poll::Ready(Err(err)), /// Err(_would_block) => continue, /// } /// } /// } /// } /// /// impl AsyncWrite for AsyncTcpStream { /// fn poll_write( /// self: Pin<&mut Self>, /// cx: &mut Context<'_>, /// buf: &[u8] /// ) -> Poll> { /// loop { /// let mut guard = ready!(self.inner.poll_write_ready(cx))?; /// /// match guard.try_io(|inner| inner.get_ref().write(buf)) { /// Ok(result) => return Poll::Ready(result), /// Err(_would_block) => continue, /// } /// } /// } /// /// fn poll_flush( /// self: Pin<&mut Self>, /// cx: &mut Context<'_>, /// ) -> Poll> { /// // tcp flush is a no-op /// Poll::Ready(Ok(())) /// } /// /// fn poll_shutdown( /// self: Pin<&mut Self>, /// cx: &mut Context<'_>, /// ) -> Poll> { /// self.inner.get_ref().shutdown(std::net::Shutdown::Write)?; /// Poll::Ready(Ok(())) /// } /// } /// ``` /// /// [`readable`]: method@Self::readable /// [`writable`]: method@Self::writable /// [`AsyncFdReadyGuard`]: struct@self::AsyncFdReadyGuard /// [`TcpStream::poll_read_ready`]: struct@crate::net::TcpStream /// [`AsyncRead`]: trait@crate::io::AsyncRead /// [`AsyncWrite`]: trait@crate::io::AsyncWrite /// [`OwnedFd`]: struct@std::os::fd::OwnedFd pub struct AsyncFd { registration: Registration, // The inner value is always present. the Option is required for `drop` and `into_inner`. // In all other methods `unwrap` is valid, and will never panic. inner: Option, } /// Represents an IO-ready event detected on a particular file descriptor that /// has not yet been acknowledged. This is a `must_use` structure to help ensure /// that you do not forget to explicitly clear (or not clear) the event. /// /// This type exposes an immutable reference to the underlying IO object. #[must_use = "You must explicitly choose whether to clear the readiness state by calling a method on ReadyGuard"] pub struct AsyncFdReadyGuard<'a, T: AsRawFd> { async_fd: &'a AsyncFd, event: Option, } /// Represents an IO-ready event detected on a particular file descriptor that /// has not yet been acknowledged. This is a `must_use` structure to help ensure /// that you do not forget to explicitly clear (or not clear) the event. /// /// This type exposes a mutable reference to the underlying IO object. #[must_use = "You must explicitly choose whether to clear the readiness state by calling a method on ReadyGuard"] pub struct AsyncFdReadyMutGuard<'a, T: AsRawFd> { async_fd: &'a mut AsyncFd, event: Option, } impl AsyncFd { /// Creates an [`AsyncFd`] backed by (and taking ownership of) an object /// implementing [`AsRawFd`]. The backing file descriptor is cached at the /// time of creation. /// /// Only configures the [`Interest::READABLE`] and [`Interest::WRITABLE`] interests. For more /// control, use [`AsyncFd::with_interest`]. /// /// This method must be called in the context of a tokio runtime. /// /// # Panics /// /// This function panics if there is no current reactor set, or if the `rt` /// feature flag is not enabled. #[inline] #[track_caller] pub fn new(inner: T) -> io::Result where T: AsRawFd, { Self::with_interest(inner, Interest::READABLE | Interest::WRITABLE) } /// Creates an [`AsyncFd`] backed by (and taking ownership of) an object /// implementing [`AsRawFd`], with a specific [`Interest`]. The backing /// file descriptor is cached at the time of creation. /// /// # Panics /// /// This function panics if there is no current reactor set, or if the `rt` /// feature flag is not enabled. #[inline] #[track_caller] pub fn with_interest(inner: T, interest: Interest) -> io::Result where T: AsRawFd, { Self::new_with_handle_and_interest(inner, scheduler::Handle::current(), interest) } #[track_caller] pub(crate) fn new_with_handle_and_interest( inner: T, handle: scheduler::Handle, interest: Interest, ) -> io::Result { Self::try_new_with_handle_and_interest(inner, handle, interest).map_err(Into::into) } /// Creates an [`AsyncFd`] backed by (and taking ownership of) an object /// implementing [`AsRawFd`]. The backing file descriptor is cached at the /// time of creation. /// /// Only configures the [`Interest::READABLE`] and [`Interest::WRITABLE`] interests. For more /// control, use [`AsyncFd::try_with_interest`]. /// /// This method must be called in the context of a tokio runtime. /// /// In the case of failure, it returns [`AsyncFdTryNewError`] that contains the original object /// passed to this function. /// /// # Panics /// /// This function panics if there is no current reactor set, or if the `rt` /// feature flag is not enabled. #[inline] #[track_caller] pub fn try_new(inner: T) -> Result> where T: AsRawFd, { Self::try_with_interest(inner, Interest::READABLE | Interest::WRITABLE) } /// Creates an [`AsyncFd`] backed by (and taking ownership of) an object /// implementing [`AsRawFd`], with a specific [`Interest`]. The backing /// file descriptor is cached at the time of creation. /// /// In the case of failure, it returns [`AsyncFdTryNewError`] that contains the original object /// passed to this function. /// /// # Panics /// /// This function panics if there is no current reactor set, or if the `rt` /// feature flag is not enabled. #[inline] #[track_caller] pub fn try_with_interest(inner: T, interest: Interest) -> Result> where T: AsRawFd, { Self::try_new_with_handle_and_interest(inner, scheduler::Handle::current(), interest) } #[track_caller] pub(crate) fn try_new_with_handle_and_interest( inner: T, handle: scheduler::Handle, interest: Interest, ) -> Result> { let fd = inner.as_raw_fd(); match Registration::new_with_interest_and_handle(&mut SourceFd(&fd), interest, handle) { Ok(registration) => Ok(AsyncFd { registration, inner: Some(inner), }), Err(cause) => Err(AsyncFdTryNewError { inner, cause }), } } /// Returns a shared reference to the backing object of this [`AsyncFd`]. #[inline] pub fn get_ref(&self) -> &T { self.inner.as_ref().unwrap() } /// Returns a mutable reference to the backing object of this [`AsyncFd`]. #[inline] pub fn get_mut(&mut self) -> &mut T { self.inner.as_mut().unwrap() } fn take_inner(&mut self) -> Option { let inner = self.inner.take()?; let fd = inner.as_raw_fd(); let _ = self.registration.deregister(&mut SourceFd(&fd)); Some(inner) } /// Deregisters this file descriptor and returns ownership of the backing /// object. pub fn into_inner(mut self) -> T { self.take_inner().unwrap() } /// Polls for read readiness. /// /// If the file descriptor is not currently ready for reading, this method /// will store a clone of the [`Waker`] from the provided [`Context`]. When the /// file descriptor becomes ready for reading, [`Waker::wake`] will be called. /// /// Note that on multiple calls to [`poll_read_ready`] or /// [`poll_read_ready_mut`], only the `Waker` from the `Context` passed to the /// most recent call is scheduled to receive a wakeup. (However, /// [`poll_write_ready`] retains a second, independent waker). /// /// This method is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// This method takes `&self`, so it is possible to call this method /// concurrently with other methods on this struct. This method only /// provides shared access to the inner IO resource when handling the /// [`AsyncFdReadyGuard`]. /// /// [`poll_read_ready`]: method@Self::poll_read_ready /// [`poll_read_ready_mut`]: method@Self::poll_read_ready_mut /// [`poll_write_ready`]: method@Self::poll_write_ready /// [`readable`]: method@Self::readable /// [`Context`]: struct@std::task::Context /// [`Waker`]: struct@std::task::Waker /// [`Waker::wake`]: method@std::task::Waker::wake pub fn poll_read_ready<'a>( &'a self, cx: &mut Context<'_>, ) -> Poll>> { let event = ready!(self.registration.poll_read_ready(cx))?; Poll::Ready(Ok(AsyncFdReadyGuard { async_fd: self, event: Some(event), })) } /// Polls for read readiness. /// /// If the file descriptor is not currently ready for reading, this method /// will store a clone of the [`Waker`] from the provided [`Context`]. When the /// file descriptor becomes ready for reading, [`Waker::wake`] will be called. /// /// Note that on multiple calls to [`poll_read_ready`] or /// [`poll_read_ready_mut`], only the `Waker` from the `Context` passed to the /// most recent call is scheduled to receive a wakeup. (However, /// [`poll_write_ready`] retains a second, independent waker). /// /// This method is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// This method takes `&mut self`, so it is possible to access the inner IO /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. /// /// [`poll_read_ready`]: method@Self::poll_read_ready /// [`poll_read_ready_mut`]: method@Self::poll_read_ready_mut /// [`poll_write_ready`]: method@Self::poll_write_ready /// [`readable`]: method@Self::readable /// [`Context`]: struct@std::task::Context /// [`Waker`]: struct@std::task::Waker /// [`Waker::wake`]: method@std::task::Waker::wake pub fn poll_read_ready_mut<'a>( &'a mut self, cx: &mut Context<'_>, ) -> Poll>> { let event = ready!(self.registration.poll_read_ready(cx))?; Poll::Ready(Ok(AsyncFdReadyMutGuard { async_fd: self, event: Some(event), })) } /// Polls for write readiness. /// /// If the file descriptor is not currently ready for writing, this method /// will store a clone of the [`Waker`] from the provided [`Context`]. When the /// file descriptor becomes ready for writing, [`Waker::wake`] will be called. /// /// Note that on multiple calls to [`poll_write_ready`] or /// [`poll_write_ready_mut`], only the `Waker` from the `Context` passed to the /// most recent call is scheduled to receive a wakeup. (However, /// [`poll_read_ready`] retains a second, independent waker). /// /// This method is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// This method takes `&self`, so it is possible to call this method /// concurrently with other methods on this struct. This method only /// provides shared access to the inner IO resource when handling the /// [`AsyncFdReadyGuard`]. /// /// [`poll_read_ready`]: method@Self::poll_read_ready /// [`poll_write_ready`]: method@Self::poll_write_ready /// [`poll_write_ready_mut`]: method@Self::poll_write_ready_mut /// [`writable`]: method@Self::readable /// [`Context`]: struct@std::task::Context /// [`Waker`]: struct@std::task::Waker /// [`Waker::wake`]: method@std::task::Waker::wake pub fn poll_write_ready<'a>( &'a self, cx: &mut Context<'_>, ) -> Poll>> { let event = ready!(self.registration.poll_write_ready(cx))?; Poll::Ready(Ok(AsyncFdReadyGuard { async_fd: self, event: Some(event), })) } /// Polls for write readiness. /// /// If the file descriptor is not currently ready for writing, this method /// will store a clone of the [`Waker`] from the provided [`Context`]. When the /// file descriptor becomes ready for writing, [`Waker::wake`] will be called. /// /// Note that on multiple calls to [`poll_write_ready`] or /// [`poll_write_ready_mut`], only the `Waker` from the `Context` passed to the /// most recent call is scheduled to receive a wakeup. (However, /// [`poll_read_ready`] retains a second, independent waker). /// /// This method is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// This method takes `&mut self`, so it is possible to access the inner IO /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. /// /// [`poll_read_ready`]: method@Self::poll_read_ready /// [`poll_write_ready`]: method@Self::poll_write_ready /// [`poll_write_ready_mut`]: method@Self::poll_write_ready_mut /// [`writable`]: method@Self::readable /// [`Context`]: struct@std::task::Context /// [`Waker`]: struct@std::task::Waker /// [`Waker::wake`]: method@std::task::Waker::wake pub fn poll_write_ready_mut<'a>( &'a mut self, cx: &mut Context<'_>, ) -> Poll>> { let event = ready!(self.registration.poll_write_ready(cx))?; Poll::Ready(Ok(AsyncFdReadyMutGuard { async_fd: self, event: Some(event), })) } /// Waits for any of the requested ready states, returning a /// [`AsyncFdReadyGuard`] that must be dropped to resume /// polling for the requested ready states. /// /// The function may complete without the file descriptor being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// When an IO operation does return `io::ErrorKind::WouldBlock`, the readiness must be cleared. /// When a combined interest is used, it is important to clear only the readiness /// that is actually observed to block. For instance when the combined /// interest `Interest::READABLE | Interest::WRITABLE` is used, and a read blocks, only /// read readiness should be cleared using the [`AsyncFdReadyGuard::clear_ready_matching`] method: /// `guard.clear_ready_matching(Ready::READABLE)`. /// Also clearing the write readiness in this case would be incorrect. The [`AsyncFdReadyGuard::clear_ready`] /// method clears all readiness flags. /// /// This method takes `&self`, so it is possible to call this method /// concurrently with other methods on this struct. This method only /// provides shared access to the inner IO resource when handling the /// [`AsyncFdReadyGuard`]. /// /// # Examples /// /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without /// splitting. /// /// ```no_run /// use std::error::Error; /// use std::io; /// use std::io::{Read, Write}; /// use std::net::TcpStream; /// use tokio::io::unix::AsyncFd; /// use tokio::io::{Interest, Ready}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080")?; /// stream.set_nonblocking(true)?; /// let stream = AsyncFd::new(stream)?; /// /// loop { /// let mut guard = stream /// .ready(Interest::READABLE | Interest::WRITABLE) /// .await?; /// /// if guard.ready().is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.get_ref().read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a read has blocked, but a write might still succeed. /// // clear only the read readiness. /// guard.clear_ready_matching(Ready::READABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// if guard.ready().is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.get_ref().write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a write has blocked, but a read might still succeed. /// // clear only the write readiness. /// guard.clear_ready_matching(Ready::WRITABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result> { let event = self.registration.readiness(interest).await?; Ok(AsyncFdReadyGuard { async_fd: self, event: Some(event), }) } /// Waits for any of the requested ready states, returning a /// [`AsyncFdReadyMutGuard`] that must be dropped to resume /// polling for the requested ready states. /// /// The function may complete without the file descriptor being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// When an IO operation does return `io::ErrorKind::WouldBlock`, the readiness must be cleared. /// When a combined interest is used, it is important to clear only the readiness /// that is actually observed to block. For instance when the combined /// interest `Interest::READABLE | Interest::WRITABLE` is used, and a read blocks, only /// read readiness should be cleared using the [`AsyncFdReadyMutGuard::clear_ready_matching`] method: /// `guard.clear_ready_matching(Ready::READABLE)`. /// Also clearing the write readiness in this case would be incorrect. /// The [`AsyncFdReadyMutGuard::clear_ready`] method clears all readiness flags. /// /// This method takes `&mut self`, so it is possible to access the inner IO /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. /// /// # Examples /// /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without /// splitting. /// /// ```no_run /// use std::error::Error; /// use std::io; /// use std::io::{Read, Write}; /// use std::net::TcpStream; /// use tokio::io::unix::AsyncFd; /// use tokio::io::{Interest, Ready}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080")?; /// stream.set_nonblocking(true)?; /// let mut stream = AsyncFd::new(stream)?; /// /// loop { /// let mut guard = stream /// .ready_mut(Interest::READABLE | Interest::WRITABLE) /// .await?; /// /// if guard.ready().is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match guard.get_inner_mut().read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a read has blocked, but a write might still succeed. /// // clear only the read readiness. /// guard.clear_ready_matching(Ready::READABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// if guard.ready().is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match guard.get_inner_mut().write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a write has blocked, but a read might still succeed. /// // clear only the write readiness. /// guard.clear_ready_matching(Ready::WRITABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub async fn ready_mut( &mut self, interest: Interest, ) -> io::Result> { let event = self.registration.readiness(interest).await?; Ok(AsyncFdReadyMutGuard { async_fd: self, event: Some(event), }) } /// Waits for the file descriptor to become readable, returning a /// [`AsyncFdReadyGuard`] that must be dropped to resume read-readiness /// polling. /// /// This method takes `&self`, so it is possible to call this method /// concurrently with other methods on this struct. This method only /// provides shared access to the inner IO resource when handling the /// [`AsyncFdReadyGuard`]. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn readable<'a>(&'a self) -> io::Result> { self.ready(Interest::READABLE).await } /// Waits for the file descriptor to become readable, returning a /// [`AsyncFdReadyMutGuard`] that must be dropped to resume read-readiness /// polling. /// /// This method takes `&mut self`, so it is possible to access the inner IO /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn readable_mut<'a>(&'a mut self) -> io::Result> { self.ready_mut(Interest::READABLE).await } /// Waits for the file descriptor to become writable, returning a /// [`AsyncFdReadyGuard`] that must be dropped to resume write-readiness /// polling. /// /// This method takes `&self`, so it is possible to call this method /// concurrently with other methods on this struct. This method only /// provides shared access to the inner IO resource when handling the /// [`AsyncFdReadyGuard`]. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn writable<'a>(&'a self) -> io::Result> { self.ready(Interest::WRITABLE).await } /// Waits for the file descriptor to become writable, returning a /// [`AsyncFdReadyMutGuard`] that must be dropped to resume write-readiness /// polling. /// /// This method takes `&mut self`, so it is possible to access the inner IO /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn writable_mut<'a>(&'a mut self) -> io::Result> { self.ready_mut(Interest::WRITABLE).await } /// Reads or writes from the file descriptor using a user-provided IO operation. /// /// The `async_io` method is a convenience utility that waits for the file /// descriptor to become ready, and then executes the provided IO operation. /// Since file descriptors may be marked ready spuriously, the closure will /// be called repeatedly until it returns something other than a /// [`WouldBlock`] error. This is done using the following loop: /// /// ```no_run /// # use std::io::{self, Result}; /// # struct Dox { inner: T } /// # impl Dox { /// # async fn writable(&self) -> Result<&Self> { /// # Ok(self) /// # } /// # fn try_io(&self, _: impl FnMut(&T) -> Result) -> Result> { /// # panic!() /// # } /// async fn async_io(&self, mut f: impl FnMut(&T) -> io::Result) -> io::Result { /// loop { /// // or `readable` if called with the read interest. /// let guard = self.writable().await?; /// /// match guard.try_io(&mut f) { /// Ok(result) => return result, /// Err(_would_block) => continue, /// } /// } /// } /// # } /// ``` /// /// The closure should only return a [`WouldBlock`] error if it has performed /// an IO operation on the file descriptor that failed due to the file descriptor not being /// ready. Returning a [`WouldBlock`] error in any other situation will /// incorrectly clear the readiness flag, which can cause the file descriptor to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio [`AsyncFd`] type, as this will mess with the /// readiness flag and can cause the file descriptor to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// # Examples /// /// This example sends some bytes on the inner [`std::net::UdpSocket`]. The `async_io` /// method waits for readiness, and retries if the send operation does block. This example /// is equivalent to the one given for [`try_io`]. /// /// ```no_run /// use tokio::io::{Interest, unix::AsyncFd}; /// /// use std::io; /// use std::net::UdpSocket; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("0.0.0.0:8080")?; /// socket.set_nonblocking(true)?; /// let async_fd = AsyncFd::new(socket)?; /// /// let written = async_fd /// .async_io(Interest::WRITABLE, |inner| inner.send(&[1, 2])) /// .await?; /// /// println!("wrote {written} bytes"); /// /// Ok(()) /// } /// ``` /// /// [`try_io`]: AsyncFdReadyGuard::try_io /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock pub async fn async_io( &self, interest: Interest, mut f: impl FnMut(&T) -> io::Result, ) -> io::Result { self.registration .async_io(interest, || f(self.get_ref())) .await } /// Reads or writes from the file descriptor using a user-provided IO operation. /// /// The behavior is the same as [`async_io`], except that the closure can mutate the inner /// value of the [`AsyncFd`]. /// /// [`async_io`]: AsyncFd::async_io pub async fn async_io_mut( &mut self, interest: Interest, mut f: impl FnMut(&mut T) -> io::Result, ) -> io::Result { self.registration .async_io(interest, || f(self.inner.as_mut().unwrap())) .await } /// Tries to read or write from the file descriptor using a user-provided IO operation. /// /// If the file descriptor is ready, the provided closure is called. The closure /// should attempt to perform IO operation on the file descriptor by manually /// calling the appropriate syscall. If the operation fails because the /// file descriptor is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the file descriptor is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the file descriptor that failed due to the file descriptor not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the file descriptor to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `AsyncFd` type, as this will mess with the /// readiness flag and can cause the file descriptor to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub fn try_io( &self, interest: Interest, f: impl FnOnce(&T) -> io::Result, ) -> io::Result { self.registration .try_io(interest, || f(self.inner.as_ref().unwrap())) } /// Tries to read or write from the file descriptor using a user-provided IO operation. /// /// The behavior is the same as [`try_io`], except that the closure can mutate the inner /// value of the [`AsyncFd`]. /// /// [`try_io`]: AsyncFd::try_io pub fn try_io_mut( &mut self, interest: Interest, f: impl FnOnce(&mut T) -> io::Result, ) -> io::Result { self.registration .try_io(interest, || f(self.inner.as_mut().unwrap())) } } impl AsRawFd for AsyncFd { fn as_raw_fd(&self) -> RawFd { self.inner.as_ref().unwrap().as_raw_fd() } } impl std::os::unix::io::AsFd for AsyncFd { fn as_fd(&self) -> std::os::unix::io::BorrowedFd<'_> { unsafe { std::os::unix::io::BorrowedFd::borrow_raw(self.as_raw_fd()) } } } impl std::fmt::Debug for AsyncFd { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AsyncFd") .field("inner", &self.inner) .finish() } } impl Drop for AsyncFd { fn drop(&mut self) { let _ = self.take_inner(); } } impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { /// Indicates to tokio that the file descriptor is no longer ready. All /// internal readiness flags will be cleared, and tokio will wait for the /// next edge-triggered readiness notification from the OS. /// /// This function is commonly used with guards returned by [`AsyncFd::readable`] and /// [`AsyncFd::writable`]. /// /// It is critical that this function not be called unless your code /// _actually observes_ that the file descriptor is _not_ ready. Do not call /// it simply because, for example, a read succeeded; it should be called /// when a read is observed to block. /// /// This method only clears readiness events that happened before the creation of this guard. /// In other words, if the IO resource becomes ready between the creation of the guard and /// this call to `clear_ready`, then the readiness is not actually cleared. pub fn clear_ready(&mut self) { if let Some(event) = self.event.take() { self.async_fd.registration.clear_readiness(event); } } /// Indicates to tokio that the file descriptor no longer has a specific readiness. /// The internal readiness flag will be cleared, and tokio will wait for the /// next edge-triggered readiness notification from the OS. /// /// This function is useful in combination with the [`AsyncFd::ready`] method when a /// combined interest like `Interest::READABLE | Interest::WRITABLE` is used. /// /// It is critical that this function not be called unless your code /// _actually observes_ that the file descriptor is _not_ ready for the provided `Ready`. /// Do not call it simply because, for example, a read succeeded; it should be called /// when a read is observed to block. Only clear the specific readiness that is observed to /// block. For example when a read blocks when using a combined interest, /// only clear `Ready::READABLE`. /// /// This method only clears readiness events that happened before the creation of this guard. /// In other words, if the IO resource becomes ready between the creation of the guard and /// this call to `clear_ready`, then the readiness is not actually cleared. /// /// # Examples /// /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without /// splitting. /// /// ```no_run /// use std::error::Error; /// use std::io; /// use std::io::{Read, Write}; /// use std::net::TcpStream; /// use tokio::io::unix::AsyncFd; /// use tokio::io::{Interest, Ready}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080")?; /// stream.set_nonblocking(true)?; /// let stream = AsyncFd::new(stream)?; /// /// loop { /// let mut guard = stream /// .ready(Interest::READABLE | Interest::WRITABLE) /// .await?; /// /// if guard.ready().is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.get_ref().read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a read has blocked, but a write might still succeed. /// // clear only the read readiness. /// guard.clear_ready_matching(Ready::READABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// if guard.ready().is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.get_ref().write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a write has blocked, but a read might still succeed. /// // clear only the write readiness. /// guard.clear_ready_matching(Ready::WRITABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub fn clear_ready_matching(&mut self, ready: Ready) { if let Some(mut event) = self.event.take() { self.async_fd .registration .clear_readiness(event.with_ready(ready)); // the event is no longer ready for the readiness that was just cleared event.ready = event.ready - ready; if !event.ready.is_empty() { self.event = Some(event); } } } /// This method should be invoked when you intentionally want to keep the /// ready flag asserted. /// /// While this function is itself a no-op, it satisfies the `#[must_use]` /// constraint on the [`AsyncFdReadyGuard`] type. pub fn retain_ready(&mut self) { // no-op } /// Get the [`Ready`] value associated with this guard. /// /// This method will return the empty readiness state if /// [`AsyncFdReadyGuard::clear_ready`] has been called on /// the guard. /// /// [`Ready`]: crate::io::Ready pub fn ready(&self) -> Ready { match &self.event { Some(event) => event.ready, None => Ready::EMPTY, } } /// Performs the provided IO operation. /// /// If `f` returns a [`WouldBlock`] error, the readiness state associated /// with this file descriptor is cleared, and the method returns /// `Err(TryIoError::WouldBlock)`. You will typically need to poll the /// `AsyncFd` again when this happens. /// /// This method helps ensure that the readiness state of the underlying file /// descriptor remains in sync with the tokio-side readiness state, by /// clearing the tokio-side state only when a [`WouldBlock`] condition /// occurs. It is the responsibility of the caller to ensure that `f` /// returns [`WouldBlock`] only if the file descriptor that originated this /// `AsyncFdReadyGuard` no longer expresses the readiness state that was queried to /// create this `AsyncFdReadyGuard`. /// /// # Examples /// /// This example sends some bytes to the inner [`std::net::UdpSocket`]. Waiting /// for write-readiness and retrying when the send operation does block are explicit. /// This example can be written more succinctly using [`AsyncFd::async_io`]. /// /// ```no_run /// use tokio::io::unix::AsyncFd; /// /// use std::io; /// use std::net::UdpSocket; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("0.0.0.0:8080")?; /// socket.set_nonblocking(true)?; /// let async_fd = AsyncFd::new(socket)?; /// /// let written = loop { /// let mut guard = async_fd.writable().await?; /// match guard.try_io(|inner| inner.get_ref().send(&[1, 2])) { /// Ok(result) => { /// break result?; /// } /// Err(_would_block) => { /// // try_io already cleared the file descriptor's readiness state /// continue; /// } /// } /// }; /// /// println!("wrote {written} bytes"); /// /// Ok(()) /// } /// ``` /// /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock // Alias for old name in 0.x #[cfg_attr(docsrs, doc(alias = "with_io"))] pub fn try_io( &mut self, f: impl FnOnce(&'a AsyncFd) -> io::Result, ) -> Result, TryIoError> { let result = f(self.async_fd); match result { Err(err) if err.kind() == io::ErrorKind::WouldBlock => { self.clear_ready(); Err(TryIoError(())) } result => Ok(result), } } /// Returns a shared reference to the inner [`AsyncFd`]. pub fn get_ref(&self) -> &'a AsyncFd { self.async_fd } /// Returns a shared reference to the backing object of the inner [`AsyncFd`]. pub fn get_inner(&self) -> &'a Inner { self.get_ref().get_ref() } } impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { /// Indicates to tokio that the file descriptor is no longer ready. All /// internal readiness flags will be cleared, and tokio will wait for the /// next edge-triggered readiness notification from the OS. /// /// This function is commonly used with guards returned by [`AsyncFd::readable_mut`] and /// [`AsyncFd::writable_mut`]. /// /// It is critical that this function not be called unless your code /// _actually observes_ that the file descriptor is _not_ ready. Do not call /// it simply because, for example, a read succeeded; it should be called /// when a read is observed to block. /// /// This method only clears readiness events that happened before the creation of this guard. /// In other words, if the IO resource becomes ready between the creation of the guard and /// this call to `clear_ready`, then the readiness is not actually cleared. pub fn clear_ready(&mut self) { if let Some(event) = self.event.take() { self.async_fd.registration.clear_readiness(event); } } /// Indicates to tokio that the file descriptor no longer has a specific readiness. /// The internal readiness flag will be cleared, and tokio will wait for the /// next edge-triggered readiness notification from the OS. /// /// This function is useful in combination with the [`AsyncFd::ready_mut`] method when a /// combined interest like `Interest::READABLE | Interest::WRITABLE` is used. /// /// It is critical that this function not be called unless your code /// _actually observes_ that the file descriptor is _not_ ready for the provided `Ready`. /// Do not call it simply because, for example, a read succeeded; it should be called /// when a read is observed to block. Only clear the specific readiness that is observed to /// block. For example when a read blocks when using a combined interest, /// only clear `Ready::READABLE`. /// /// This method only clears readiness events that happened before the creation of this guard. /// In other words, if the IO resource becomes ready between the creation of the guard and /// this call to `clear_ready`, then the readiness is not actually cleared. /// /// # Examples /// /// Concurrently read and write to a [`std::net::TcpStream`] on the same task without /// splitting. /// /// ```no_run /// use std::error::Error; /// use std::io; /// use std::io::{Read, Write}; /// use std::net::TcpStream; /// use tokio::io::unix::AsyncFd; /// use tokio::io::{Interest, Ready}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080")?; /// stream.set_nonblocking(true)?; /// let mut stream = AsyncFd::new(stream)?; /// /// loop { /// let mut guard = stream /// .ready_mut(Interest::READABLE | Interest::WRITABLE) /// .await?; /// /// if guard.ready().is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match guard.get_inner_mut().read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a read has blocked, but a write might still succeed. /// // clear only the read readiness. /// guard.clear_ready_matching(Ready::READABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// if guard.ready().is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match guard.get_inner_mut().write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // a write has blocked, but a read might still succeed. /// // clear only the write readiness. /// guard.clear_ready_matching(Ready::WRITABLE); /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub fn clear_ready_matching(&mut self, ready: Ready) { if let Some(mut event) = self.event.take() { self.async_fd .registration .clear_readiness(event.with_ready(ready)); // the event is no longer ready for the readiness that was just cleared event.ready = event.ready - ready; if !event.ready.is_empty() { self.event = Some(event); } } } /// This method should be invoked when you intentionally want to keep the /// ready flag asserted. /// /// While this function is itself a no-op, it satisfies the `#[must_use]` /// constraint on the [`AsyncFdReadyGuard`] type. pub fn retain_ready(&mut self) { // no-op } /// Get the [`Ready`] value associated with this guard. /// /// This method will return the empty readiness state if /// [`AsyncFdReadyGuard::clear_ready`] has been called on /// the guard. /// /// [`Ready`]: super::Ready pub fn ready(&self) -> Ready { match &self.event { Some(event) => event.ready, None => Ready::EMPTY, } } /// Performs the provided IO operation. /// /// If `f` returns a [`WouldBlock`] error, the readiness state associated /// with this file descriptor is cleared, and the method returns /// `Err(TryIoError::WouldBlock)`. You will typically need to poll the /// `AsyncFd` again when this happens. /// /// This method helps ensure that the readiness state of the underlying file /// descriptor remains in sync with the tokio-side readiness state, by /// clearing the tokio-side state only when a [`WouldBlock`] condition /// occurs. It is the responsibility of the caller to ensure that `f` /// returns [`WouldBlock`] only if the file descriptor that originated this /// `AsyncFdReadyGuard` no longer expresses the readiness state that was queried to /// create this `AsyncFdReadyGuard`. /// /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock pub fn try_io( &mut self, f: impl FnOnce(&mut AsyncFd) -> io::Result, ) -> Result, TryIoError> { let result = f(self.async_fd); match result { Err(err) if err.kind() == io::ErrorKind::WouldBlock => { self.clear_ready(); Err(TryIoError(())) } result => Ok(result), } } /// Returns a shared reference to the inner [`AsyncFd`]. pub fn get_ref(&self) -> &AsyncFd { self.async_fd } /// Returns a mutable reference to the inner [`AsyncFd`]. pub fn get_mut(&mut self) -> &mut AsyncFd { self.async_fd } /// Returns a shared reference to the backing object of the inner [`AsyncFd`]. pub fn get_inner(&self) -> &Inner { self.get_ref().get_ref() } /// Returns a mutable reference to the backing object of the inner [`AsyncFd`]. pub fn get_inner_mut(&mut self) -> &mut Inner { self.get_mut().get_mut() } } impl<'a, T: std::fmt::Debug + AsRawFd> std::fmt::Debug for AsyncFdReadyGuard<'a, T> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ReadyGuard") .field("async_fd", &self.async_fd) .finish() } } impl<'a, T: std::fmt::Debug + AsRawFd> std::fmt::Debug for AsyncFdReadyMutGuard<'a, T> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("MutReadyGuard") .field("async_fd", &self.async_fd) .finish() } } /// The error type returned by [`try_io`]. /// /// This error indicates that the IO resource returned a [`WouldBlock`] error. /// /// [`WouldBlock`]: std::io::ErrorKind::WouldBlock /// [`try_io`]: method@AsyncFdReadyGuard::try_io #[derive(Debug)] pub struct TryIoError(()); /// Error returned by [`try_new`] or [`try_with_interest`]. /// /// [`try_new`]: AsyncFd::try_new /// [`try_with_interest`]: AsyncFd::try_with_interest pub struct AsyncFdTryNewError { inner: T, cause: io::Error, } impl AsyncFdTryNewError { /// Returns the original object passed to [`try_new`] or [`try_with_interest`] /// alongside the error that caused these functions to fail. /// /// [`try_new`]: AsyncFd::try_new /// [`try_with_interest`]: AsyncFd::try_with_interest pub fn into_parts(self) -> (T, io::Error) { (self.inner, self.cause) } } impl fmt::Display for AsyncFdTryNewError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&self.cause, f) } } impl fmt::Debug for AsyncFdTryNewError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&self.cause, f) } } impl Error for AsyncFdTryNewError { fn source(&self) -> Option<&(dyn Error + 'static)> { Some(&self.cause) } } impl From> for io::Error { fn from(value: AsyncFdTryNewError) -> Self { value.cause } } tokio-1.43.1/src/io/async_read.rs000064400000000000000000000105041046102023000147130ustar 00000000000000use super::ReadBuf; use std::io; use std::ops::DerefMut; use std::pin::Pin; use std::task::{Context, Poll}; /// Reads bytes from a source. /// /// This trait is analogous to the [`std::io::Read`] trait, but integrates with /// the asynchronous task system. In particular, the [`poll_read`] method, /// unlike [`Read::read`], will automatically queue the current task for wakeup /// and return if data is not yet available, rather than blocking the calling /// thread. /// /// Specifically, this means that the `poll_read` function will return one of /// the following: /// /// * `Poll::Ready(Ok(()))` means that data was immediately read and placed into /// the output buffer. The amount of data read can be determined by the /// increase in the length of the slice returned by `ReadBuf::filled`. If the /// difference is 0, either EOF has been reached, or the output buffer had zero /// capacity (i.e. `buf.remaining()` == 0). /// /// * `Poll::Pending` means that no data was read into the buffer /// provided. The I/O object is not currently readable but may become readable /// in the future. Most importantly, **the current future's task is scheduled /// to get unparked when the object is readable**. This means that like /// `Future::poll` you'll receive a notification when the I/O object is /// readable again. /// /// * `Poll::Ready(Err(e))` for other errors are standard I/O errors coming from the /// underlying object. /// /// This trait importantly means that the `read` method only works in the /// context of a future's task. The object may panic if used outside of a task. /// /// Utilities for working with `AsyncRead` values are provided by /// [`AsyncReadExt`]. /// /// [`poll_read`]: AsyncRead::poll_read /// [`std::io::Read`]: std::io::Read /// [`Read::read`]: std::io::Read::read /// [`AsyncReadExt`]: crate::io::AsyncReadExt pub trait AsyncRead { /// Attempts to read from the `AsyncRead` into `buf`. /// /// On success, returns `Poll::Ready(Ok(()))` and places data in the /// unfilled portion of `buf`. If no data was read (`buf.filled().len()` is /// unchanged), it implies that EOF has been reached. /// /// If no data is available for reading, the method returns `Poll::Pending` /// and arranges for the current task (via `cx.waker()`) to receive a /// notification when the object becomes readable or is closed. fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll>; } macro_rules! deref_async_read { () => { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut **self).poll_read(cx, buf) } }; } impl AsyncRead for Box { deref_async_read!(); } impl AsyncRead for &mut T { deref_async_read!(); } impl

AsyncRead for Pin

where P: DerefMut + Unpin, P::Target: AsyncRead, { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.get_mut().as_mut().poll_read(cx, buf) } } impl AsyncRead for &[u8] { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let amt = std::cmp::min(self.len(), buf.remaining()); let (a, b) = self.split_at(amt); buf.put_slice(a); *self = b; Poll::Ready(Ok(())) } } impl + Unpin> AsyncRead for io::Cursor { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let pos = self.position(); let slice: &[u8] = (*self).get_ref().as_ref(); // The position could technically be out of bounds, so don't panic... if pos > slice.len() as u64 { return Poll::Ready(Ok(())); } let start = pos as usize; let amt = std::cmp::min(slice.len() - start, buf.remaining()); // Add won't overflow because of pos check above. let end = start + amt; buf.put_slice(&slice[start..end]); self.set_position(end as u64); Poll::Ready(Ok(())) } } tokio-1.43.1/src/io/async_seek.rs000064400000000000000000000060431046102023000147320ustar 00000000000000use std::io::{self, SeekFrom}; use std::ops::DerefMut; use std::pin::Pin; use std::task::{Context, Poll}; /// Seek bytes asynchronously. /// /// This trait is analogous to the [`std::io::Seek`] trait, but integrates /// with the asynchronous task system. In particular, the `start_seek` /// method, unlike [`Seek::seek`], will not block the calling thread. /// /// Utilities for working with `AsyncSeek` values are provided by /// [`AsyncSeekExt`]. /// /// [`std::io::Seek`]: std::io::Seek /// [`Seek::seek`]: std::io::Seek::seek() /// [`AsyncSeekExt`]: crate::io::AsyncSeekExt pub trait AsyncSeek { /// Attempts to seek to an offset, in bytes, in a stream. /// /// A seek beyond the end of a stream is allowed, but behavior is defined /// by the implementation. /// /// If this function returns successfully, then the job has been submitted. /// To find out when it completes, call `poll_complete`. /// /// # Errors /// /// This function can return [`io::ErrorKind::Other`] in case there is /// another seek in progress. To avoid this, it is advisable that any call /// to `start_seek` is preceded by a call to `poll_complete` to ensure all /// pending seeks have completed. fn start_seek(self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()>; /// Waits for a seek operation to complete. /// /// If the seek operation completed successfully, /// this method returns the new position from the start of the stream. /// That position can be used later with [`SeekFrom::Start`]. Repeatedly /// calling this function without calling `start_seek` might return the /// same result. /// /// # Errors /// /// Seeking to a negative offset is considered an error. fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll>; } macro_rules! deref_async_seek { () => { fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { Pin::new(&mut **self).start_seek(pos) } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut **self).poll_complete(cx) } }; } impl AsyncSeek for Box { deref_async_seek!(); } impl AsyncSeek for &mut T { deref_async_seek!(); } impl

AsyncSeek for Pin

where P: DerefMut + Unpin, P::Target: AsyncSeek, { fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { self.get_mut().as_mut().start_seek(pos) } fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_mut().as_mut().poll_complete(cx) } } impl + Unpin> AsyncSeek for io::Cursor { fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { io::Seek::seek(&mut *self, pos).map(drop) } fn poll_complete(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(self.get_mut().position())) } } tokio-1.43.1/src/io/async_write.rs000064400000000000000000000342561046102023000151440ustar 00000000000000use std::io::{self, IoSlice}; use std::ops::DerefMut; use std::pin::Pin; use std::task::{Context, Poll}; /// Writes bytes asynchronously. /// /// The trait inherits from [`std::io::Write`] and indicates that an I/O object is /// **nonblocking**. All non-blocking I/O objects must return an error when /// bytes cannot be written instead of blocking the current thread. /// /// Specifically, this means that the [`poll_write`] function will return one of /// the following: /// /// * `Poll::Ready(Ok(n))` means that `n` bytes of data was immediately /// written. /// /// * `Poll::Pending` means that no data was written from the buffer /// provided. The I/O object is not currently writable but may become writable /// in the future. Most importantly, **the current future's task is scheduled /// to get unparked when the object is writable**. This means that like /// `Future::poll` you'll receive a notification when the I/O object is /// writable again. /// /// * `Poll::Ready(Err(e))` for other errors are standard I/O errors coming from the /// underlying object. /// /// This trait importantly means that the [`write`][stdwrite] method only works in /// the context of a future's task. The object may panic if used outside of a task. /// /// Note that this trait also represents that the [`Write::flush`][stdflush] method /// works very similarly to the `write` method, notably that `Ok(())` means that the /// writer has successfully been flushed, a "would block" error means that the /// current task is ready to receive a notification when flushing can make more /// progress, and otherwise normal errors can happen as well. /// /// Utilities for working with `AsyncWrite` values are provided by /// [`AsyncWriteExt`]. /// /// [`std::io::Write`]: std::io::Write /// [`poll_write`]: AsyncWrite::poll_write() /// [stdwrite]: std::io::Write::write() /// [stdflush]: std::io::Write::flush() /// [`AsyncWriteExt`]: crate::io::AsyncWriteExt pub trait AsyncWrite { /// Attempt to write bytes from `buf` into the object. /// /// On success, returns `Poll::Ready(Ok(num_bytes_written))`. If successful, /// then it must be guaranteed that `n <= buf.len()`. A return value of `0` /// typically means that the underlying object is no longer able to accept /// bytes and will likely not be able to in the future as well, or that the /// buffer provided is empty. /// /// If the object is not ready for writing, the method returns /// `Poll::Pending` and arranges for the current task (via /// `cx.waker()`) to receive a notification when the object becomes /// writable or is closed. fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll>; /// Attempts to flush the object, ensuring that any buffered data reach /// their destination. /// /// On success, returns `Poll::Ready(Ok(()))`. /// /// If flushing cannot immediately complete, this method returns /// `Poll::Pending` and arranges for the current task (via /// `cx.waker()`) to receive a notification when the object can make /// progress towards flushing. fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll>; /// Initiates or attempts to shut down this writer, returning success when /// the I/O connection has completely shut down. /// /// This method is intended to be used for asynchronous shutdown of I/O /// connections. For example this is suitable for implementing shutdown of a /// TLS connection or calling `TcpStream::shutdown` on a proxied connection. /// Protocols sometimes need to flush out final pieces of data or otherwise /// perform a graceful shutdown handshake, reading/writing more data as /// appropriate. This method is the hook for such protocols to implement the /// graceful shutdown logic. /// /// This `shutdown` method is required by implementers of the /// `AsyncWrite` trait. Wrappers typically just want to proxy this call /// through to the wrapped type, and base types will typically implement /// shutdown logic here or just return `Ok(().into())`. Note that if you're /// wrapping an underlying `AsyncWrite` a call to `shutdown` implies that /// transitively the entire stream has been shut down. After your wrapper's /// shutdown logic has been executed you should shut down the underlying /// stream. /// /// Invocation of a `shutdown` implies an invocation of `flush`. Once this /// method returns `Ready` it implies that a flush successfully happened /// before the shutdown happened. That is, callers don't need to call /// `flush` before calling `shutdown`. They can rely that by calling /// `shutdown` any pending buffered data will be written out. /// /// # Return value /// /// This function returns a `Poll>` classified as such: /// /// * `Poll::Ready(Ok(()))` - indicates that the connection was /// successfully shut down and is now safe to deallocate/drop/close /// resources associated with it. This method means that the current task /// will no longer receive any notifications due to this method and the /// I/O object itself is likely no longer usable. /// /// * `Poll::Pending` - indicates that shutdown is initiated but could /// not complete just yet. This may mean that more I/O needs to happen to /// continue this shutdown operation. The current task is scheduled to /// receive a notification when it's otherwise ready to continue the /// shutdown operation. When woken up this method should be called again. /// /// * `Poll::Ready(Err(e))` - indicates a fatal error has happened with shutdown, /// indicating that the shutdown operation did not complete successfully. /// This typically means that the I/O object is no longer usable. /// /// # Errors /// /// This function can return normal I/O errors through `Err`, described /// above. Additionally this method may also render the underlying /// `Write::write` method no longer usable (e.g. will return errors in the /// future). It's recommended that once `shutdown` is called the /// `write` method is no longer called. /// /// # Panics /// /// This function will panic if not called within the context of a future's /// task. fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll>; /// Like [`poll_write`], except that it writes from a slice of buffers. /// /// Data is copied from each buffer in order, with the final buffer /// read from possibly being only partially consumed. This method must /// behave as a call to [`write`] with the buffers concatenated would. /// /// The default implementation calls [`poll_write`] with either the first nonempty /// buffer provided, or an empty one if none exists. /// /// On success, returns `Poll::Ready(Ok(num_bytes_written))`. /// /// If the object is not ready for writing, the method returns /// `Poll::Pending` and arranges for the current task (via /// `cx.waker()`) to receive a notification when the object becomes /// writable or is closed. /// /// # Note /// /// This should be implemented as a single "atomic" write action. If any /// data has been partially written, it is wrong to return an error or /// pending. /// /// [`poll_write`]: AsyncWrite::poll_write fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { let buf = bufs .iter() .find(|b| !b.is_empty()) .map_or(&[][..], |b| &**b); self.poll_write(cx, buf) } /// Determines if this writer has an efficient [`poll_write_vectored`] /// implementation. /// /// If a writer does not override the default [`poll_write_vectored`] /// implementation, code using it may want to avoid the method all together /// and coalesce writes into a single buffer for higher performance. /// /// The default implementation returns `false`. /// /// [`poll_write_vectored`]: AsyncWrite::poll_write_vectored fn is_write_vectored(&self) -> bool { false } } macro_rules! deref_async_write { () => { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut **self).poll_write(cx, buf) } fn poll_write_vectored( mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Pin::new(&mut **self).poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { (**self).is_write_vectored() } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut **self).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut **self).poll_shutdown(cx) } }; } impl AsyncWrite for Box { deref_async_write!(); } impl AsyncWrite for &mut T { deref_async_write!(); } impl

AsyncWrite for Pin

where P: DerefMut + Unpin, P::Target: AsyncWrite, { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.get_mut().as_mut().poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { self.get_mut().as_mut().poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { (**self).is_write_vectored() } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_mut().as_mut().poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_mut().as_mut().poll_shutdown(cx) } } impl AsyncWrite for Vec { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.get_mut().extend_from_slice(buf); Poll::Ready(Ok(buf.len())) } fn poll_write_vectored( mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } } impl AsyncWrite for io::Cursor<&mut [u8]> { fn poll_write( mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8], ) -> Poll> { Poll::Ready(io::Write::write(&mut *self, buf)) } fn poll_write_vectored( mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(io::Write::flush(&mut *self)) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } impl AsyncWrite for io::Cursor<&mut Vec> { fn poll_write( mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8], ) -> Poll> { Poll::Ready(io::Write::write(&mut *self, buf)) } fn poll_write_vectored( mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(io::Write::flush(&mut *self)) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } impl AsyncWrite for io::Cursor> { fn poll_write( mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8], ) -> Poll> { Poll::Ready(io::Write::write(&mut *self, buf)) } fn poll_write_vectored( mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(io::Write::flush(&mut *self)) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } impl AsyncWrite for io::Cursor> { fn poll_write( mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8], ) -> Poll> { Poll::Ready(io::Write::write(&mut *self, buf)) } fn poll_write_vectored( mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Poll::Ready(io::Write::write_vectored(&mut *self, bufs)) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(io::Write::flush(&mut *self)) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } tokio-1.43.1/src/io/blocking.rs000064400000000000000000000215161046102023000144000ustar 00000000000000use crate::io::sys; use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; use std::cmp; use std::future::Future; use std::io; use std::io::prelude::*; use std::mem::MaybeUninit; use std::pin::Pin; use std::task::{ready, Context, Poll}; /// `T` should not implement _both_ Read and Write. #[derive(Debug)] pub(crate) struct Blocking { inner: Option, state: State, /// `true` if the lower IO layer needs flushing. need_flush: bool, } #[derive(Debug)] pub(crate) struct Buf { buf: Vec, pos: usize, } pub(crate) const DEFAULT_MAX_BUF_SIZE: usize = 2 * 1024 * 1024; #[derive(Debug)] enum State { Idle(Option), Busy(sys::Blocking<(io::Result, Buf, T)>), } cfg_io_blocking! { impl Blocking { /// # Safety /// /// The `Read` implementation of `inner` must never read from the buffer /// it is borrowing and must correctly report the length of the data /// written into the buffer. #[cfg_attr(feature = "fs", allow(dead_code))] pub(crate) unsafe fn new(inner: T) -> Blocking { Blocking { inner: Some(inner), state: State::Idle(Some(Buf::with_capacity(0))), need_flush: false, } } } } impl AsyncRead for Blocking where T: Read + Unpin + Send + 'static, { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, dst: &mut ReadBuf<'_>, ) -> Poll> { loop { match self.state { State::Idle(ref mut buf_cell) => { let mut buf = buf_cell.take().unwrap(); if !buf.is_empty() { buf.copy_to(dst); *buf_cell = Some(buf); return Poll::Ready(Ok(())); } let mut inner = self.inner.take().unwrap(); let max_buf_size = cmp::min(dst.remaining(), DEFAULT_MAX_BUF_SIZE); self.state = State::Busy(sys::run(move || { // SAFETY: the requirements are satisfied by `Blocking::new`. let res = unsafe { buf.read_from(&mut inner, max_buf_size) }; (res, buf, inner) })); } State::Busy(ref mut rx) => { let (res, mut buf, inner) = ready!(Pin::new(rx).poll(cx))?; self.inner = Some(inner); match res { Ok(_) => { buf.copy_to(dst); self.state = State::Idle(Some(buf)); return Poll::Ready(Ok(())); } Err(e) => { assert!(buf.is_empty()); self.state = State::Idle(Some(buf)); return Poll::Ready(Err(e)); } } } } } } } impl AsyncWrite for Blocking where T: Write + Unpin + Send + 'static, { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, src: &[u8], ) -> Poll> { loop { match self.state { State::Idle(ref mut buf_cell) => { let mut buf = buf_cell.take().unwrap(); assert!(buf.is_empty()); let n = buf.copy_from(src, DEFAULT_MAX_BUF_SIZE); let mut inner = self.inner.take().unwrap(); self.state = State::Busy(sys::run(move || { let n = buf.len(); let res = buf.write_to(&mut inner).map(|()| n); (res, buf, inner) })); self.need_flush = true; return Poll::Ready(Ok(n)); } State::Busy(ref mut rx) => { let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?; self.state = State::Idle(Some(buf)); self.inner = Some(inner); // If error, return res?; } } } } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { let need_flush = self.need_flush; match self.state { // The buffer is not used here State::Idle(ref mut buf_cell) => { if need_flush { let buf = buf_cell.take().unwrap(); let mut inner = self.inner.take().unwrap(); self.state = State::Busy(sys::run(move || { let res = inner.flush().map(|()| 0); (res, buf, inner) })); self.need_flush = false; } else { return Poll::Ready(Ok(())); } } State::Busy(ref mut rx) => { let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?; self.state = State::Idle(Some(buf)); self.inner = Some(inner); // If error, return res?; } } } } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } } /// Repeats operations that are interrupted. macro_rules! uninterruptibly { ($e:expr) => {{ loop { match $e { Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {} res => break res, } } }}; } impl Buf { pub(crate) fn with_capacity(n: usize) -> Buf { Buf { buf: Vec::with_capacity(n), pos: 0, } } pub(crate) fn is_empty(&self) -> bool { self.len() == 0 } pub(crate) fn len(&self) -> usize { self.buf.len() - self.pos } pub(crate) fn copy_to(&mut self, dst: &mut ReadBuf<'_>) -> usize { let n = cmp::min(self.len(), dst.remaining()); dst.put_slice(&self.bytes()[..n]); self.pos += n; if self.pos == self.buf.len() { self.buf.truncate(0); self.pos = 0; } n } pub(crate) fn copy_from(&mut self, src: &[u8], max_buf_size: usize) -> usize { assert!(self.is_empty()); let n = cmp::min(src.len(), max_buf_size); self.buf.extend_from_slice(&src[..n]); n } pub(crate) fn bytes(&self) -> &[u8] { &self.buf[self.pos..] } /// # Safety /// /// `rd` must not read from the buffer `read` is borrowing and must correctly /// report the length of the data written into the buffer. pub(crate) unsafe fn read_from( &mut self, rd: &mut T, max_buf_size: usize, ) -> io::Result { assert!(self.is_empty()); self.buf.reserve(max_buf_size); let buf = &mut self.buf.spare_capacity_mut()[..max_buf_size]; // SAFETY: The memory may be uninitialized, but `rd.read` will only write to the buffer. let buf = unsafe { &mut *(buf as *mut [MaybeUninit] as *mut [u8]) }; let res = uninterruptibly!(rd.read(buf)); if let Ok(n) = res { // SAFETY: the caller promises that `rd.read` initializes // a section of `buf` and correctly reports that length. // The `self.is_empty()` assertion verifies that `n` // equals the length of the `buf` capacity that was written // to (and that `buf` isn't being shrunk). unsafe { self.buf.set_len(n) } } else { self.buf.clear(); } assert_eq!(self.pos, 0); res } pub(crate) fn write_to(&mut self, wr: &mut T) -> io::Result<()> { assert_eq!(self.pos, 0); // `write_all` already ignores interrupts let res = wr.write_all(&self.buf); self.buf.clear(); res } } cfg_fs! { impl Buf { pub(crate) fn discard_read(&mut self) -> i64 { let ret = -(self.bytes().len() as i64); self.pos = 0; self.buf.truncate(0); ret } pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>], max_buf_size: usize) -> usize { assert!(self.is_empty()); let mut rem = max_buf_size; for buf in bufs { if rem == 0 { break } let len = buf.len().min(rem); self.buf.extend_from_slice(&buf[..len]); rem -= len; } max_buf_size - rem } } } tokio-1.43.1/src/io/bsd/poll_aio.rs000064400000000000000000000163461046102023000151630ustar 00000000000000//! Use POSIX AIO futures with Tokio. use crate::io::interest::Interest; use crate::runtime::io::{ReadyEvent, Registration}; use crate::runtime::scheduler; use mio::event::Source; use mio::Registry; use mio::Token; use std::fmt; use std::io; use std::ops::{Deref, DerefMut}; use std::os::unix::io::AsRawFd; use std::os::unix::prelude::RawFd; use std::task::{ready, Context, Poll}; /// Like [`mio::event::Source`], but for POSIX AIO only. /// /// Tokio's consumer must pass an implementor of this trait to create a /// [`Aio`] object. pub trait AioSource { /// Registers this AIO event source with Tokio's reactor. fn register(&mut self, kq: RawFd, token: usize); /// Deregisters this AIO event source with Tokio's reactor. fn deregister(&mut self); } /// Wraps the user's AioSource in order to implement mio::event::Source, which /// is what the rest of the crate wants. struct MioSource(T); impl Source for MioSource { fn register( &mut self, registry: &Registry, token: Token, interests: mio::Interest, ) -> io::Result<()> { assert!(interests.is_aio() || interests.is_lio()); self.0.register(registry.as_raw_fd(), usize::from(token)); Ok(()) } fn deregister(&mut self, _registry: &Registry) -> io::Result<()> { self.0.deregister(); Ok(()) } fn reregister( &mut self, registry: &Registry, token: Token, interests: mio::Interest, ) -> io::Result<()> { assert!(interests.is_aio() || interests.is_lio()); self.0.register(registry.as_raw_fd(), usize::from(token)); Ok(()) } } /// Associates a POSIX AIO control block with the reactor that drives it. /// /// `Aio`'s wrapped type must implement [`AioSource`] to be driven /// by the reactor. /// /// The wrapped source may be accessed through the `Aio` via the `Deref` and /// `DerefMut` traits. /// /// ## Clearing readiness /// /// If [`Aio::poll_ready`] returns ready, but the consumer determines that the /// Source is not completely ready and must return to the Pending state, /// [`Aio::clear_ready`] may be used. This can be useful with /// [`lio_listio`], which may generate a kevent when only a portion of the /// operations have completed. /// /// ## Platforms /// /// Only FreeBSD implements POSIX AIO with kqueue notification, so /// `Aio` is only available for that operating system. /// /// [`lio_listio`]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/lio_listio.html // Note: Unlike every other kqueue event source, POSIX AIO registers events not // via kevent(2) but when the aiocb is submitted to the kernel via aio_read, // aio_write, etc. It needs the kqueue's file descriptor to do that. So // AsyncFd can't be used for POSIX AIO. // // Note that Aio doesn't implement Drop. There's no need. Unlike other // kqueue sources, simply dropping the object effectively deregisters it. pub struct Aio { io: MioSource, registration: Registration, } // ===== impl Aio ===== impl Aio { /// Creates a new `Aio` suitable for use with POSIX AIO functions. /// /// It will be associated with the default reactor. The runtime is usually /// set implicitly when this function is called from a future driven by a /// Tokio runtime, otherwise runtime can be set explicitly with /// [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn new_for_aio(io: E) -> io::Result { Self::new_with_interest(io, Interest::AIO) } /// Creates a new `Aio` suitable for use with [`lio_listio`]. /// /// It will be associated with the default reactor. The runtime is usually /// set implicitly when this function is called from a future driven by a /// Tokio runtime, otherwise runtime can be set explicitly with /// [`Runtime::enter`](crate::runtime::Runtime::enter) function. /// /// [`lio_listio`]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/lio_listio.html pub fn new_for_lio(io: E) -> io::Result { Self::new_with_interest(io, Interest::LIO) } fn new_with_interest(io: E, interest: Interest) -> io::Result { let mut io = MioSource(io); let handle = scheduler::Handle::current(); let registration = Registration::new_with_interest_and_handle(&mut io, interest, handle)?; Ok(Self { io, registration }) } /// Indicates to Tokio that the source is no longer ready. The internal /// readiness flag will be cleared, and tokio will wait for the next /// edge-triggered readiness notification from the OS. /// /// It is critical that this method not be called unless your code /// _actually observes_ that the source is _not_ ready. The OS must /// deliver a subsequent notification, or this source will block /// forever. It is equally critical that you `do` call this method if you /// resubmit the same structure to the kernel and poll it again. /// /// This method is not very useful with AIO readiness, since each `aiocb` /// structure is typically only used once. It's main use with /// [`lio_listio`], which will sometimes send notification when only a /// portion of its elements are complete. In that case, the caller must /// call `clear_ready` before resubmitting it. /// /// [`lio_listio`]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/lio_listio.html pub fn clear_ready(&self, ev: AioEvent) { self.registration.clear_readiness(ev.0) } /// Destroy the [`Aio`] and return its inner source. pub fn into_inner(self) -> E { self.io.0 } /// Polls for readiness. Either AIO or LIO counts. /// /// This method returns: /// * `Poll::Pending` if the underlying operation is not complete, whether /// or not it completed successfully. This will be true if the OS is /// still processing it, or if it has not yet been submitted to the OS. /// * `Poll::Ready(Ok(_))` if the underlying operation is complete. /// * `Poll::Ready(Err(_))` if the reactor has been shutdown. This does /// _not_ indicate that the underlying operation encountered an error. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided `Context` /// is scheduled to receive a wakeup when the underlying operation /// completes. Note that on multiple calls to `poll_ready`, only the `Waker` from the /// `Context` passed to the most recent call is scheduled to receive a wakeup. pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll> { let ev = ready!(self.registration.poll_read_ready(cx))?; Poll::Ready(Ok(AioEvent(ev))) } } impl Deref for Aio { type Target = E; fn deref(&self) -> &E { &self.io.0 } } impl DerefMut for Aio { fn deref_mut(&mut self) -> &mut E { &mut self.io.0 } } impl fmt::Debug for Aio { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Aio").field("io", &self.io.0).finish() } } /// Opaque data returned by [`Aio::poll_ready`]. /// /// It can be fed back to [`Aio::clear_ready`]. #[derive(Debug)] pub struct AioEvent(ReadyEvent); tokio-1.43.1/src/io/interest.rs000064400000000000000000000237701046102023000144510ustar 00000000000000#![cfg_attr(not(feature = "net"), allow(dead_code, unreachable_pub))] use crate::io::ready::Ready; use std::fmt; use std::ops; // These must be unique. // same as mio const READABLE: usize = 0b0001; const WRITABLE: usize = 0b0010; // The following are not available on all platforms. #[cfg(target_os = "freebsd")] const AIO: usize = 0b0100; #[cfg(target_os = "freebsd")] const LIO: usize = 0b1000; #[cfg(any(target_os = "linux", target_os = "android"))] const PRIORITY: usize = 0b0001_0000; // error is available on all platforms, but behavior is platform-specific // mio does not have this interest const ERROR: usize = 0b0010_0000; /// Readiness event interest. /// /// Specifies the readiness events the caller is interested in when awaiting on /// I/O resource readiness states. #[cfg_attr(docsrs, doc(cfg(feature = "net")))] #[derive(Clone, Copy, Eq, PartialEq)] pub struct Interest(usize); impl Interest { // The non-FreeBSD definitions in this block are active only when // building documentation. cfg_aio! { /// Interest for POSIX AIO. #[cfg(target_os = "freebsd")] pub const AIO: Interest = Interest(AIO); /// Interest for POSIX AIO. #[cfg(not(target_os = "freebsd"))] pub const AIO: Interest = Interest(READABLE); /// Interest for POSIX AIO `lio_listio` events. #[cfg(target_os = "freebsd")] pub const LIO: Interest = Interest(LIO); /// Interest for POSIX AIO `lio_listio` events. #[cfg(not(target_os = "freebsd"))] pub const LIO: Interest = Interest(READABLE); } /// Interest in all readable events. /// /// Readable interest includes read-closed events. pub const READABLE: Interest = Interest(READABLE); /// Interest in all writable events. /// /// Writable interest includes write-closed events. pub const WRITABLE: Interest = Interest(WRITABLE); /// Interest in error events. /// /// Passes error interest to the underlying OS selector. /// Behavior is platform-specific, read your platform's documentation. pub const ERROR: Interest = Interest(ERROR); /// Returns a `Interest` set representing priority completion interests. #[cfg(any(target_os = "linux", target_os = "android"))] #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] pub const PRIORITY: Interest = Interest(PRIORITY); /// Returns true if the value includes readable interest. /// /// # Examples /// /// ``` /// use tokio::io::Interest; /// /// assert!(Interest::READABLE.is_readable()); /// assert!(!Interest::WRITABLE.is_readable()); /// /// let both = Interest::READABLE | Interest::WRITABLE; /// assert!(both.is_readable()); /// ``` pub const fn is_readable(self) -> bool { self.0 & READABLE != 0 } /// Returns true if the value includes writable interest. /// /// # Examples /// /// ``` /// use tokio::io::Interest; /// /// assert!(!Interest::READABLE.is_writable()); /// assert!(Interest::WRITABLE.is_writable()); /// /// let both = Interest::READABLE | Interest::WRITABLE; /// assert!(both.is_writable()); /// ``` pub const fn is_writable(self) -> bool { self.0 & WRITABLE != 0 } /// Returns true if the value includes error interest. /// /// # Examples /// /// ``` /// use tokio::io::Interest; /// /// assert!(Interest::ERROR.is_error()); /// assert!(!Interest::WRITABLE.is_error()); /// /// let combined = Interest::READABLE | Interest::ERROR; /// assert!(combined.is_error()); /// ``` pub const fn is_error(self) -> bool { self.0 & ERROR != 0 } #[cfg(target_os = "freebsd")] const fn is_aio(self) -> bool { self.0 & AIO != 0 } #[cfg(target_os = "freebsd")] const fn is_lio(self) -> bool { self.0 & LIO != 0 } /// Returns true if the value includes priority interest. /// /// # Examples /// /// ``` /// use tokio::io::Interest; /// /// assert!(!Interest::READABLE.is_priority()); /// assert!(Interest::PRIORITY.is_priority()); /// /// let both = Interest::READABLE | Interest::PRIORITY; /// assert!(both.is_priority()); /// ``` #[cfg(any(target_os = "linux", target_os = "android"))] #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] pub const fn is_priority(self) -> bool { self.0 & PRIORITY != 0 } /// Add together two `Interest` values. /// /// This function works from a `const` context. /// /// # Examples /// /// ``` /// use tokio::io::Interest; /// /// const BOTH: Interest = Interest::READABLE.add(Interest::WRITABLE); /// /// assert!(BOTH.is_readable()); /// assert!(BOTH.is_writable()); #[must_use = "this returns the result of the operation, without modifying the original"] pub const fn add(self, other: Interest) -> Interest { Self(self.0 | other.0) } /// Remove `Interest` from `self`. /// /// Interests present in `other` but *not* in `self` are ignored. /// /// Returns `None` if the set would be empty after removing `Interest`. /// /// # Examples /// /// ``` /// use tokio::io::Interest; /// /// const RW_INTEREST: Interest = Interest::READABLE.add(Interest::WRITABLE); /// /// let w_interest = RW_INTEREST.remove(Interest::READABLE).unwrap(); /// assert!(!w_interest.is_readable()); /// assert!(w_interest.is_writable()); /// /// // Removing all interests from the set returns `None`. /// assert_eq!(w_interest.remove(Interest::WRITABLE), None); /// /// // Remove all interests at once. /// assert_eq!(RW_INTEREST.remove(RW_INTEREST), None); /// ``` #[must_use = "this returns the result of the operation, without modifying the original"] pub fn remove(self, other: Interest) -> Option { let value = self.0 & !other.0; if value != 0 { Some(Self(value)) } else { None } } // This function must be crate-private to avoid exposing a `mio` dependency. pub(crate) fn to_mio(self) -> mio::Interest { fn mio_add(wrapped: &mut Option, add: mio::Interest) { match wrapped { Some(inner) => *inner |= add, None => *wrapped = Some(add), } } // mio does not allow and empty interest, so use None for empty let mut mio = None; if self.is_readable() { mio_add(&mut mio, mio::Interest::READABLE); } if self.is_writable() { mio_add(&mut mio, mio::Interest::WRITABLE); } #[cfg(any(target_os = "linux", target_os = "android"))] if self.is_priority() { mio_add(&mut mio, mio::Interest::PRIORITY); } #[cfg(target_os = "freebsd")] if self.is_aio() { mio_add(&mut mio, mio::Interest::AIO); } #[cfg(target_os = "freebsd")] if self.is_lio() { mio_add(&mut mio, mio::Interest::LIO); } if self.is_error() { // There is no error interest in mio, because error events are always reported. // But mio interests cannot be empty and an interest is needed just for the registration. // // read readiness is filtered out in `Interest::mask` or `Ready::from_interest` if // the read interest was not specified by the user. mio_add(&mut mio, mio::Interest::READABLE); } // the default `mio::Interest::READABLE` should never be used in practice. Either // // - at least one tokio interest with a mio counterpart was used // - only the error tokio interest was specified // // in both cases, `mio` is Some already mio.unwrap_or(mio::Interest::READABLE) } pub(crate) fn mask(self) -> Ready { match self { Interest::READABLE => Ready::READABLE | Ready::READ_CLOSED, Interest::WRITABLE => Ready::WRITABLE | Ready::WRITE_CLOSED, #[cfg(any(target_os = "linux", target_os = "android"))] Interest::PRIORITY => Ready::PRIORITY | Ready::READ_CLOSED, Interest::ERROR => Ready::ERROR, _ => Ready::EMPTY, } } } impl ops::BitOr for Interest { type Output = Self; #[inline] fn bitor(self, other: Self) -> Self { self.add(other) } } impl ops::BitOrAssign for Interest { #[inline] fn bitor_assign(&mut self, other: Self) { *self = *self | other; } } impl fmt::Debug for Interest { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { let mut separator = false; if self.is_readable() { if separator { write!(fmt, " | ")?; } write!(fmt, "READABLE")?; separator = true; } if self.is_writable() { if separator { write!(fmt, " | ")?; } write!(fmt, "WRITABLE")?; separator = true; } #[cfg(any(target_os = "linux", target_os = "android"))] if self.is_priority() { if separator { write!(fmt, " | ")?; } write!(fmt, "PRIORITY")?; separator = true; } #[cfg(target_os = "freebsd")] if self.is_aio() { if separator { write!(fmt, " | ")?; } write!(fmt, "AIO")?; separator = true; } #[cfg(target_os = "freebsd")] if self.is_lio() { if separator { write!(fmt, " | ")?; } write!(fmt, "LIO")?; separator = true; } if self.is_error() { if separator { write!(fmt, " | ")?; } write!(fmt, "ERROR")?; separator = true; } let _ = separator; Ok(()) } } tokio-1.43.1/src/io/join.rs000064400000000000000000000062211046102023000135430ustar 00000000000000//! Join two values implementing `AsyncRead` and `AsyncWrite` into a single one. use crate::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; /// Join two values implementing `AsyncRead` and `AsyncWrite` into a /// single handle. pub fn join(reader: R, writer: W) -> Join where R: AsyncRead, W: AsyncWrite, { Join { reader, writer } } pin_project_lite::pin_project! { /// Joins two values implementing `AsyncRead` and `AsyncWrite` into a /// single handle. #[derive(Debug)] pub struct Join { #[pin] reader: R, #[pin] writer: W, } } impl Join where R: AsyncRead, W: AsyncWrite, { /// Splits this `Join` back into its `AsyncRead` and `AsyncWrite` /// components. pub fn into_inner(self) -> (R, W) { (self.reader, self.writer) } /// Returns a reference to the inner reader. pub fn reader(&self) -> &R { &self.reader } /// Returns a reference to the inner writer. pub fn writer(&self) -> &W { &self.writer } /// Returns a mutable reference to the inner reader. pub fn reader_mut(&mut self) -> &mut R { &mut self.reader } /// Returns a mutable reference to the inner writer. pub fn writer_mut(&mut self) -> &mut W { &mut self.writer } /// Returns a pinned mutable reference to the inner reader. pub fn reader_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { self.project().reader } /// Returns a pinned mutable reference to the inner writer. pub fn writer_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> { self.project().writer } } impl AsyncRead for Join where R: AsyncRead, { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.project().reader.poll_read(cx, buf) } } impl AsyncWrite for Join where W: AsyncWrite, { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.project().writer.poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().writer.poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().writer.poll_shutdown(cx) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.project().writer.poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { self.writer.is_write_vectored() } } impl AsyncBufRead for Join where R: AsyncBufRead, { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().reader.poll_fill_buf(cx) } fn consume(self: Pin<&mut Self>, amt: usize) { self.project().reader.consume(amt) } } tokio-1.43.1/src/io/mod.rs000064400000000000000000000241531046102023000133670ustar 00000000000000//! Traits, helpers, and type definitions for asynchronous I/O functionality. //! //! This module is the asynchronous version of `std::io`. Primarily, it //! defines two traits, [`AsyncRead`] and [`AsyncWrite`], which are asynchronous //! versions of the [`Read`] and [`Write`] traits in the standard library. //! //! # `AsyncRead` and `AsyncWrite` //! //! Like the standard library's [`Read`] and [`Write`] traits, [`AsyncRead`] and //! [`AsyncWrite`] provide the most general interface for reading and writing //! input and output. Unlike the standard library's traits, however, they are //! _asynchronous_ — meaning that reading from or writing to a `tokio::io` //! type will _yield_ to the Tokio scheduler when IO is not ready, rather than //! blocking. This allows other tasks to run while waiting on IO. //! //! Another difference is that `AsyncRead` and `AsyncWrite` only contain //! core methods needed to provide asynchronous reading and writing //! functionality. Instead, utility methods are defined in the [`AsyncReadExt`] //! and [`AsyncWriteExt`] extension traits. These traits are automatically //! implemented for all values that implement `AsyncRead` and `AsyncWrite` //! respectively. //! //! End users will rarely interact directly with `AsyncRead` and //! `AsyncWrite`. Instead, they will use the async functions defined in the //! extension traits. Library authors are expected to implement `AsyncRead` //! and `AsyncWrite` in order to provide types that behave like byte streams. //! //! Even with these differences, Tokio's `AsyncRead` and `AsyncWrite` traits //! can be used in almost exactly the same manner as the standard library's //! `Read` and `Write`. Most types in the standard library that implement `Read` //! and `Write` have asynchronous equivalents in `tokio` that implement //! `AsyncRead` and `AsyncWrite`, such as [`File`] and [`TcpStream`]. //! //! For example, the standard library documentation introduces `Read` by //! [demonstrating][std_example] reading some bytes from a [`std::fs::File`]. We //! can do the same with [`tokio::fs::File`][`File`]: //! //! ```no_run //! use tokio::io::{self, AsyncReadExt}; //! use tokio::fs::File; //! //! #[tokio::main] //! async fn main() -> io::Result<()> { //! let mut f = File::open("foo.txt").await?; //! let mut buffer = [0; 10]; //! //! // read up to 10 bytes //! let n = f.read(&mut buffer).await?; //! //! println!("The bytes: {:?}", &buffer[..n]); //! Ok(()) //! } //! ``` //! //! [`File`]: crate::fs::File //! [`TcpStream`]: crate::net::TcpStream //! [`std::fs::File`]: std::fs::File //! [std_example]: std::io#read-and-write //! //! ## Buffered Readers and Writers //! //! Byte-based interfaces are unwieldy and can be inefficient, as we'd need to be //! making near-constant calls to the operating system. To help with this, //! `std::io` comes with [support for _buffered_ readers and writers][stdbuf], //! and therefore, `tokio::io` does as well. //! //! Tokio provides an async version of the [`std::io::BufRead`] trait, //! [`AsyncBufRead`]; and async [`BufReader`] and [`BufWriter`] structs, which //! wrap readers and writers. These wrappers use a buffer, reducing the number //! of calls and providing nicer methods for accessing exactly what you want. //! //! For example, [`BufReader`] works with the [`AsyncBufRead`] trait to add //! extra methods to any async reader: //! //! ```no_run //! use tokio::io::{self, BufReader, AsyncBufReadExt}; //! use tokio::fs::File; //! //! #[tokio::main] //! async fn main() -> io::Result<()> { //! let f = File::open("foo.txt").await?; //! let mut reader = BufReader::new(f); //! let mut buffer = String::new(); //! //! // read a line into buffer //! reader.read_line(&mut buffer).await?; //! //! println!("{}", buffer); //! Ok(()) //! } //! ``` //! //! [`BufWriter`] doesn't add any new ways of writing; it just buffers every call //! to [`write`](crate::io::AsyncWriteExt::write). However, you **must** flush //! [`BufWriter`] to ensure that any buffered data is written. //! //! ```no_run //! use tokio::io::{self, BufWriter, AsyncWriteExt}; //! use tokio::fs::File; //! //! #[tokio::main] //! async fn main() -> io::Result<()> { //! let f = File::create("foo.txt").await?; //! { //! let mut writer = BufWriter::new(f); //! //! // Write a byte to the buffer. //! writer.write(&[42u8]).await?; //! //! // Flush the buffer before it goes out of scope. //! writer.flush().await?; //! //! } // Unless flushed or shut down, the contents of the buffer is discarded on drop. //! //! Ok(()) //! } //! ``` //! //! [stdbuf]: std::io#bufreader-and-bufwriter //! [`std::io::BufRead`]: std::io::BufRead //! [`AsyncBufRead`]: crate::io::AsyncBufRead //! [`BufReader`]: crate::io::BufReader //! [`BufWriter`]: crate::io::BufWriter //! //! ## Implementing `AsyncRead` and `AsyncWrite` //! //! Because they are traits, we can implement [`AsyncRead`] and [`AsyncWrite`] for //! our own types, as well. Note that these traits must only be implemented for //! non-blocking I/O types that integrate with the futures type system. In //! other words, these types must never block the thread, and instead the //! current task is notified when the I/O resource is ready. //! //! ## Conversion to and from Stream/Sink //! //! It is often convenient to encapsulate the reading and writing of bytes in a //! [`Stream`] or [`Sink`] of data. //! //! Tokio provides simple wrappers for converting [`AsyncRead`] to [`Stream`] //! and vice-versa in the [tokio-util] crate, see [`ReaderStream`] and //! [`StreamReader`]. //! //! There are also utility traits that abstract the asynchronous buffering //! necessary to write your own adaptors for encoding and decoding bytes to/from //! your structured data, allowing to transform something that implements //! [`AsyncRead`]/[`AsyncWrite`] into a [`Stream`]/[`Sink`], see [`Decoder`] and //! [`Encoder`] in the [tokio-util::codec] module. //! //! [tokio-util]: https://docs.rs/tokio-util //! [tokio-util::codec]: https://docs.rs/tokio-util/latest/tokio_util/codec/index.html //! //! # Standard input and output //! //! Tokio provides asynchronous APIs to standard [input], [output], and [error]. //! These APIs are very similar to the ones provided by `std`, but they also //! implement [`AsyncRead`] and [`AsyncWrite`]. //! //! Note that the standard input / output APIs **must** be used from the //! context of the Tokio runtime, as they require Tokio-specific features to //! function. Calling these functions outside of a Tokio runtime will panic. //! //! [input]: fn@stdin //! [output]: fn@stdout //! [error]: fn@stderr //! //! # `std` re-exports //! //! Additionally, [`Error`], [`ErrorKind`], [`Result`], and [`SeekFrom`] are //! re-exported from `std::io` for ease of use. //! //! [`AsyncRead`]: trait@AsyncRead //! [`AsyncWrite`]: trait@AsyncWrite //! [`AsyncReadExt`]: trait@AsyncReadExt //! [`AsyncWriteExt`]: trait@AsyncWriteExt //! ["codec"]: https://docs.rs/tokio-util/latest/tokio_util/codec/index.html //! [`Encoder`]: https://docs.rs/tokio-util/latest/tokio_util/codec/trait.Encoder.html //! [`Decoder`]: https://docs.rs/tokio-util/latest/tokio_util/codec/trait.Decoder.html //! [`ReaderStream`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.ReaderStream.html //! [`StreamReader`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.StreamReader.html //! [`Error`]: struct@Error //! [`ErrorKind`]: enum@ErrorKind //! [`Result`]: type@Result //! [`Read`]: std::io::Read //! [`SeekFrom`]: enum@SeekFrom //! [`Sink`]: https://docs.rs/futures/0.3/futures/sink/trait.Sink.html //! [`Stream`]: https://docs.rs/futures/0.3/futures/stream/trait.Stream.html //! [`Write`]: std::io::Write #![cfg_attr( not(all(feature = "rt", feature = "net")), allow(dead_code, unused_imports) )] cfg_io_blocking! { pub(crate) mod blocking; } mod async_buf_read; pub use self::async_buf_read::AsyncBufRead; mod async_read; pub use self::async_read::AsyncRead; mod async_seek; pub use self::async_seek::AsyncSeek; mod async_write; pub use self::async_write::AsyncWrite; mod read_buf; pub use self::read_buf::ReadBuf; // Re-export some types from `std::io` so that users don't have to deal // with conflicts when `use`ing `tokio::io` and `std::io`. #[doc(no_inline)] pub use std::io::{Error, ErrorKind, Result, SeekFrom}; cfg_io_driver_impl! { pub(crate) mod interest; pub(crate) mod ready; cfg_net! { pub use interest::Interest; pub use ready::Ready; } #[cfg_attr(target_os = "wasi", allow(unused_imports))] mod poll_evented; #[cfg(not(loom))] #[cfg_attr(target_os = "wasi", allow(unused_imports))] pub(crate) use poll_evented::PollEvented; } // The bsd module can't be build on Windows, so we completely ignore it, even // when building documentation. #[cfg(unix)] cfg_aio! { /// BSD-specific I/O types. pub mod bsd { mod poll_aio; pub use poll_aio::{Aio, AioEvent, AioSource}; } } cfg_net_unix! { mod async_fd; pub mod unix { //! Asynchronous IO structures specific to Unix-like operating systems. pub use super::async_fd::{AsyncFd, AsyncFdTryNewError, AsyncFdReadyGuard, AsyncFdReadyMutGuard, TryIoError}; } } cfg_io_std! { mod stdio_common; mod stderr; pub use stderr::{stderr, Stderr}; mod stdin; pub use stdin::{stdin, Stdin}; mod stdout; pub use stdout::{stdout, Stdout}; } cfg_io_util! { mod split; pub use split::{split, ReadHalf, WriteHalf}; mod join; pub use join::{join, Join}; pub(crate) mod seek; pub(crate) mod util; pub use util::{ copy, copy_bidirectional, copy_bidirectional_with_sizes, copy_buf, duplex, empty, repeat, sink, simplex, AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt, BufReader, BufStream, BufWriter, DuplexStream, Empty, Lines, Repeat, Sink, Split, Take, SimplexStream, }; } cfg_not_io_util! { cfg_process! { pub(crate) mod util; } } cfg_io_blocking! { /// Types in this module can be mocked out in tests. mod sys { // TODO: don't rename pub(crate) use crate::blocking::spawn_blocking as run; pub(crate) use crate::blocking::JoinHandle as Blocking; } } tokio-1.43.1/src/io/poll_evented.rs000064400000000000000000000301661046102023000152710ustar 00000000000000use crate::io::interest::Interest; use crate::runtime::io::Registration; use crate::runtime::scheduler; use mio::event::Source; use std::fmt; use std::io; use std::ops::Deref; use std::panic::{RefUnwindSafe, UnwindSafe}; use std::task::ready; cfg_io_driver! { /// Associates an I/O resource that implements the [`std::io::Read`] and/or /// [`std::io::Write`] traits with the reactor that drives it. /// /// `PollEvented` uses [`Registration`] internally to take a type that /// implements [`mio::event::Source`] as well as [`std::io::Read`] and/or /// [`std::io::Write`] and associate it with a reactor that will drive it. /// /// Once the [`mio::event::Source`] type is wrapped by `PollEvented`, it can be /// used from within the future's execution model. As such, the /// `PollEvented` type provides [`AsyncRead`] and [`AsyncWrite`] /// implementations using the underlying I/O resource as well as readiness /// events provided by the reactor. /// /// **Note**: While `PollEvented` is `Sync` (if the underlying I/O type is /// `Sync`), the caller must ensure that there are at most two tasks that /// use a `PollEvented` instance concurrently. One for reading and one for /// writing. While violating this requirement is "safe" from a Rust memory /// model point of view, it will result in unexpected behavior in the form /// of lost notifications and tasks hanging. /// /// ## Readiness events /// /// Besides just providing [`AsyncRead`] and [`AsyncWrite`] implementations, /// this type also supports access to the underlying readiness event stream. /// While similar in function to what [`Registration`] provides, the /// semantics are a bit different. /// /// Two functions are provided to access the readiness events: /// [`poll_read_ready`] and [`poll_write_ready`]. These functions return the /// current readiness state of the `PollEvented` instance. If /// [`poll_read_ready`] indicates read readiness, immediately calling /// [`poll_read_ready`] again will also indicate read readiness. /// /// When the operation is attempted and is unable to succeed due to the I/O /// resource not being ready, the caller must call [`clear_readiness`]. /// This clears the readiness state until a new readiness event is received. /// /// This allows the caller to implement additional functions. For example, /// [`TcpListener`] implements `poll_accept` by using [`poll_read_ready`] and /// [`clear_readiness`]. /// /// ## Platform-specific events /// /// `PollEvented` also allows receiving platform-specific `mio::Ready` events. /// These events are included as part of the read readiness event stream. The /// write readiness event stream is only for `Ready::writable()` events. /// /// [`AsyncRead`]: crate::io::AsyncRead /// [`AsyncWrite`]: crate::io::AsyncWrite /// [`TcpListener`]: crate::net::TcpListener /// [`clear_readiness`]: Registration::clear_readiness /// [`poll_read_ready`]: Registration::poll_read_ready /// [`poll_write_ready`]: Registration::poll_write_ready pub(crate) struct PollEvented { io: Option, registration: Registration, } } // ===== impl PollEvented ===== impl PollEvented { /// Creates a new `PollEvented` associated with the default reactor. /// /// The returned `PollEvented` has readable and writable interests. For more control, use /// [`Self::new_with_interest`]. /// /// # Panics /// /// This function panics if thread-local runtime is not set. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. #[track_caller] #[cfg_attr(feature = "signal", allow(unused))] pub(crate) fn new(io: E) -> io::Result { PollEvented::new_with_interest(io, Interest::READABLE | Interest::WRITABLE) } /// Creates a new `PollEvented` associated with the default reactor, for /// specific `Interest` state. `new_with_interest` should be used over `new` /// when you need control over the readiness state, such as when a file /// descriptor only allows reads. This does not add `hup` or `error` so if /// you are interested in those states, you will need to add them to the /// readiness state passed to this function. /// /// # Panics /// /// This function panics if thread-local runtime is not set. /// /// The runtime is usually set implicitly when this function is called from /// a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) /// function. #[track_caller] #[cfg_attr(feature = "signal", allow(unused))] pub(crate) fn new_with_interest(io: E, interest: Interest) -> io::Result { Self::new_with_interest_and_handle(io, interest, scheduler::Handle::current()) } #[track_caller] pub(crate) fn new_with_interest_and_handle( mut io: E, interest: Interest, handle: scheduler::Handle, ) -> io::Result { let registration = Registration::new_with_interest_and_handle(&mut io, interest, handle)?; Ok(Self { io: Some(io), registration, }) } /// Returns a reference to the registration. #[cfg(feature = "net")] pub(crate) fn registration(&self) -> &Registration { &self.registration } /// Deregisters the inner io from the registration and returns a Result containing the inner io. #[cfg(any(feature = "net", feature = "process"))] pub(crate) fn into_inner(mut self) -> io::Result { let mut inner = self.io.take().unwrap(); // As io shouldn't ever be None, just unwrap here. self.registration.deregister(&mut inner)?; Ok(inner) } #[cfg(all(feature = "process", target_os = "linux"))] pub(crate) fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.registration .poll_read_ready(cx) .map_err(io::Error::from) .map_ok(|_| ()) } /// Re-register under new runtime with `interest`. #[cfg(all(feature = "process", target_os = "linux"))] pub(crate) fn reregister(&mut self, interest: Interest) -> io::Result<()> { let io = self.io.as_mut().unwrap(); // As io shouldn't ever be None, just unwrap here. let _ = self.registration.deregister(io); self.registration = Registration::new_with_interest_and_handle(io, interest, scheduler::Handle::current())?; Ok(()) } } feature! { #![any(feature = "net", all(unix, feature = "process"))] use crate::io::ReadBuf; use std::task::{Context, Poll}; impl PollEvented { // Safety: The caller must ensure that `E` can read into uninitialized memory pub(crate) unsafe fn poll_read<'a>( &'a self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> where &'a E: io::Read + 'a, { use std::io::Read; loop { let evt = ready!(self.registration.poll_read_ready(cx))?; let b = &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]); // used only when the cfgs below apply #[allow(unused_variables)] let len = b.len(); match self.io.as_ref().unwrap().read(b) { Ok(n) => { // When mio is using the epoll or kqueue selector, reading a partially full // buffer is sufficient to show that the socket buffer has been drained. // // This optimization does not work for level-triggered selectors such as // windows or when poll is used. // // Read more: // https://github.com/tokio-rs/tokio/issues/5866 #[cfg(all( not(mio_unsupported_force_poll_poll), any( // epoll target_os = "android", target_os = "illumos", target_os = "linux", target_os = "redox", // kqueue target_os = "dragonfly", target_os = "freebsd", target_os = "ios", target_os = "macos", target_os = "netbsd", target_os = "openbsd", target_os = "tvos", target_os = "visionos", target_os = "watchos", ) ))] if 0 < n && n < len { self.registration.clear_readiness(evt); } // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the // buffer. buf.assume_init(n); buf.advance(n); return Poll::Ready(Ok(())); }, Err(e) if e.kind() == io::ErrorKind::WouldBlock => { self.registration.clear_readiness(evt); } Err(e) => return Poll::Ready(Err(e)), } } } pub(crate) fn poll_write<'a>(&'a self, cx: &mut Context<'_>, buf: &[u8]) -> Poll> where &'a E: io::Write + 'a, { use std::io::Write; loop { let evt = ready!(self.registration.poll_write_ready(cx))?; match self.io.as_ref().unwrap().write(buf) { Ok(n) => { // if we write only part of our buffer, this is sufficient on unix to show // that the socket buffer is full. Unfortunately this assumption // fails for level-triggered selectors (like on Windows or poll even for // UNIX): https://github.com/tokio-rs/tokio/issues/5866 if n > 0 && (!cfg!(windows) && !cfg!(mio_unsupported_force_poll_poll) && n < buf.len()) { self.registration.clear_readiness(evt); } return Poll::Ready(Ok(n)); }, Err(e) if e.kind() == io::ErrorKind::WouldBlock => { self.registration.clear_readiness(evt); } Err(e) => return Poll::Ready(Err(e)), } } } #[cfg(any(feature = "net", feature = "process"))] pub(crate) fn poll_write_vectored<'a>( &'a self, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> where &'a E: io::Write + 'a, { use std::io::Write; self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write_vectored(bufs)) } } } impl UnwindSafe for PollEvented {} impl RefUnwindSafe for PollEvented {} impl Deref for PollEvented { type Target = E; fn deref(&self) -> &E { self.io.as_ref().unwrap() } } impl fmt::Debug for PollEvented { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("PollEvented").field("io", &self.io).finish() } } impl Drop for PollEvented { fn drop(&mut self) { if let Some(mut io) = self.io.take() { // Ignore errors let _ = self.registration.deregister(&mut io); } } } tokio-1.43.1/src/io/read_buf.rs000064400000000000000000000266751046102023000143720ustar 00000000000000use std::fmt; use std::mem::MaybeUninit; /// A wrapper around a byte buffer that is incrementally filled and initialized. /// /// This type is a sort of "double cursor". It tracks three regions in the /// buffer: a region at the beginning of the buffer that has been logically /// filled with data, a region that has been initialized at some point but not /// yet logically filled, and a region at the end that may be uninitialized. /// The filled region is guaranteed to be a subset of the initialized region. /// /// In summary, the contents of the buffer can be visualized as: /// /// ```not_rust /// [ capacity ] /// [ filled | unfilled ] /// [ initialized | uninitialized ] /// ``` /// /// It is undefined behavior to de-initialize any bytes from the uninitialized /// region, since it is merely unknown whether this region is uninitialized or /// not, and if part of it turns out to be initialized, it must stay initialized. pub struct ReadBuf<'a> { buf: &'a mut [MaybeUninit], filled: usize, initialized: usize, } impl<'a> ReadBuf<'a> { /// Creates a new `ReadBuf` from a fully initialized buffer. #[inline] pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> { let initialized = buf.len(); let buf = unsafe { slice_to_uninit_mut(buf) }; ReadBuf { buf, filled: 0, initialized, } } /// Creates a new `ReadBuf` from a buffer that may be uninitialized. /// /// The internal cursor will mark the entire buffer as uninitialized. If /// the buffer is known to be partially initialized, then use `assume_init` /// to move the internal cursor. #[inline] pub fn uninit(buf: &'a mut [MaybeUninit]) -> ReadBuf<'a> { ReadBuf { buf, filled: 0, initialized: 0, } } /// Returns the total capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { self.buf.len() } /// Returns a shared reference to the filled portion of the buffer. #[inline] pub fn filled(&self) -> &[u8] { let slice = &self.buf[..self.filled]; // safety: filled describes how far into the buffer that the // user has filled with bytes, so it's been initialized. unsafe { slice_assume_init(slice) } } /// Returns a mutable reference to the filled portion of the buffer. #[inline] pub fn filled_mut(&mut self) -> &mut [u8] { let slice = &mut self.buf[..self.filled]; // safety: filled describes how far into the buffer that the // user has filled with bytes, so it's been initialized. unsafe { slice_assume_init_mut(slice) } } /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`. #[inline] pub fn take(&mut self, n: usize) -> ReadBuf<'_> { let max = std::cmp::min(self.remaining(), n); // Safety: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`. unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) } } /// Returns a shared reference to the initialized portion of the buffer. /// /// This includes the filled portion. #[inline] pub fn initialized(&self) -> &[u8] { let slice = &self.buf[..self.initialized]; // safety: initialized describes how far into the buffer that the // user has at some point initialized with bytes. unsafe { slice_assume_init(slice) } } /// Returns a mutable reference to the initialized portion of the buffer. /// /// This includes the filled portion. #[inline] pub fn initialized_mut(&mut self) -> &mut [u8] { let slice = &mut self.buf[..self.initialized]; // safety: initialized describes how far into the buffer that the // user has at some point initialized with bytes. unsafe { slice_assume_init_mut(slice) } } /// Returns a mutable reference to the entire buffer, without ensuring that it has been fully /// initialized. /// /// The elements between 0 and `self.filled().len()` are filled, and those between 0 and /// `self.initialized().len()` are initialized (and so can be converted to a `&mut [u8]`). /// /// The caller of this method must ensure that these invariants are upheld. For example, if the /// caller initializes some of the uninitialized section of the buffer, it must call /// [`assume_init`](Self::assume_init) with the number of bytes initialized. /// /// # Safety /// /// The caller must not de-initialize portions of the buffer that have already been initialized. /// This includes any bytes in the region marked as uninitialized by `ReadBuf`. #[inline] pub unsafe fn inner_mut(&mut self) -> &mut [MaybeUninit] { self.buf } /// Returns a mutable reference to the unfilled part of the buffer without ensuring that it has been fully /// initialized. /// /// # Safety /// /// The caller must not de-initialize portions of the buffer that have already been initialized. /// This includes any bytes in the region marked as uninitialized by `ReadBuf`. #[inline] pub unsafe fn unfilled_mut(&mut self) -> &mut [MaybeUninit] { &mut self.buf[self.filled..] } /// Returns a mutable reference to the unfilled part of the buffer, ensuring it is fully initialized. /// /// Since `ReadBuf` tracks the region of the buffer that has been initialized, this is effectively "free" after /// the first use. #[inline] pub fn initialize_unfilled(&mut self) -> &mut [u8] { self.initialize_unfilled_to(self.remaining()) } /// Returns a mutable reference to the first `n` bytes of the unfilled part of the buffer, ensuring it is /// fully initialized. /// /// # Panics /// /// Panics if `self.remaining()` is less than `n`. #[inline] #[track_caller] pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] { assert!(self.remaining() >= n, "n overflows remaining"); // This can't overflow, otherwise the assert above would have failed. let end = self.filled + n; if self.initialized < end { unsafe { self.buf[self.initialized..end] .as_mut_ptr() .write_bytes(0, end - self.initialized); } self.initialized = end; } let slice = &mut self.buf[self.filled..end]; // safety: just above, we checked that the end of the buf has // been initialized to some value. unsafe { slice_assume_init_mut(slice) } } /// Returns the number of bytes at the end of the slice that have not yet been filled. #[inline] pub fn remaining(&self) -> usize { self.capacity() - self.filled } /// Clears the buffer, resetting the filled region to empty. /// /// The number of initialized bytes is not changed, and the contents of the buffer are not modified. #[inline] pub fn clear(&mut self) { self.filled = 0; } /// Advances the size of the filled region of the buffer. /// /// The number of initialized bytes is not changed. /// /// # Panics /// /// Panics if the filled region of the buffer would become larger than the initialized region. #[inline] #[track_caller] pub fn advance(&mut self, n: usize) { let new = self.filled.checked_add(n).expect("filled overflow"); self.set_filled(new); } /// Sets the size of the filled region of the buffer. /// /// The number of initialized bytes is not changed. /// /// Note that this can be used to *shrink* the filled region of the buffer in addition to growing it (for /// example, by a `AsyncRead` implementation that compresses data in-place). /// /// # Panics /// /// Panics if the filled region of the buffer would become larger than the initialized region. #[inline] #[track_caller] pub fn set_filled(&mut self, n: usize) { assert!( n <= self.initialized, "filled must not become larger than initialized" ); self.filled = n; } /// Asserts that the first `n` unfilled bytes of the buffer are initialized. /// /// `ReadBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer /// bytes than are already known to be initialized. /// /// # Safety /// /// The caller must ensure that `n` unfilled bytes of the buffer have already been initialized. #[inline] pub unsafe fn assume_init(&mut self, n: usize) { let new = self.filled + n; if new > self.initialized { self.initialized = new; } } /// Appends data to the buffer, advancing the written position and possibly also the initialized position. /// /// # Panics /// /// Panics if `self.remaining()` is less than `buf.len()`. #[inline] #[track_caller] pub fn put_slice(&mut self, buf: &[u8]) { assert!( self.remaining() >= buf.len(), "buf.len() must fit in remaining(); buf.len() = {}, remaining() = {}", buf.len(), self.remaining() ); let amt = buf.len(); // Cannot overflow, asserted above let end = self.filled + amt; // Safety: the length is asserted above unsafe { self.buf[self.filled..end] .as_mut_ptr() .cast::() .copy_from_nonoverlapping(buf.as_ptr(), amt); } if self.initialized < end { self.initialized = end; } self.filled = end; } } #[cfg(feature = "io-util")] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] unsafe impl<'a> bytes::BufMut for ReadBuf<'a> { fn remaining_mut(&self) -> usize { self.remaining() } // SAFETY: The caller guarantees that at least `cnt` unfilled bytes have been initialized. unsafe fn advance_mut(&mut self, cnt: usize) { self.assume_init(cnt); self.advance(cnt); } fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice { // SAFETY: No region of `unfilled` will be deinitialized because it is // exposed as an `UninitSlice`, whose API guarantees that the memory is // never deinitialized. let unfilled = unsafe { self.unfilled_mut() }; let len = unfilled.len(); let ptr = unfilled.as_mut_ptr() as *mut u8; // SAFETY: The pointer is valid for `len` bytes because it comes from a // slice of that length. unsafe { bytes::buf::UninitSlice::from_raw_parts_mut(ptr, len) } } } impl fmt::Debug for ReadBuf<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ReadBuf") .field("filled", &self.filled) .field("initialized", &self.initialized) .field("capacity", &self.capacity()) .finish() } } unsafe fn slice_to_uninit_mut(slice: &mut [u8]) -> &mut [MaybeUninit] { &mut *(slice as *mut [u8] as *mut [MaybeUninit]) } // TODO: This could use `MaybeUninit::slice_assume_init` when it is stable. unsafe fn slice_assume_init(slice: &[MaybeUninit]) -> &[u8] { &*(slice as *const [MaybeUninit] as *const [u8]) } // TODO: This could use `MaybeUninit::slice_assume_init_mut` when it is stable. unsafe fn slice_assume_init_mut(slice: &mut [MaybeUninit]) -> &mut [u8] { &mut *(slice as *mut [MaybeUninit] as *mut [u8]) } tokio-1.43.1/src/io/ready.rs000064400000000000000000000217601046102023000137150ustar 00000000000000#![cfg_attr(not(feature = "net"), allow(unreachable_pub))] use crate::io::interest::Interest; use std::fmt; use std::ops; const READABLE: usize = 0b0_01; const WRITABLE: usize = 0b0_10; const READ_CLOSED: usize = 0b0_0100; const WRITE_CLOSED: usize = 0b0_1000; #[cfg(any(target_os = "linux", target_os = "android"))] const PRIORITY: usize = 0b1_0000; const ERROR: usize = 0b10_0000; /// Describes the readiness state of an I/O resources. /// /// `Ready` tracks which operation an I/O resource is ready to perform. #[cfg_attr(docsrs, doc(cfg(feature = "net")))] #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct Ready(usize); impl Ready { /// Returns the empty `Ready` set. pub const EMPTY: Ready = Ready(0); /// Returns a `Ready` representing readable readiness. pub const READABLE: Ready = Ready(READABLE); /// Returns a `Ready` representing writable readiness. pub const WRITABLE: Ready = Ready(WRITABLE); /// Returns a `Ready` representing read closed readiness. pub const READ_CLOSED: Ready = Ready(READ_CLOSED); /// Returns a `Ready` representing write closed readiness. pub const WRITE_CLOSED: Ready = Ready(WRITE_CLOSED); /// Returns a `Ready` representing priority readiness. #[cfg(any(target_os = "linux", target_os = "android"))] #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] pub const PRIORITY: Ready = Ready(PRIORITY); /// Returns a `Ready` representing error readiness. pub const ERROR: Ready = Ready(ERROR); /// Returns a `Ready` representing readiness for all operations. #[cfg(any(target_os = "linux", target_os = "android"))] pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED | ERROR | PRIORITY); /// Returns a `Ready` representing readiness for all operations. #[cfg(not(any(target_os = "linux", target_os = "android")))] pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED | ERROR); // Must remain crate-private to avoid adding a public dependency on Mio. pub(crate) fn from_mio(event: &mio::event::Event) -> Ready { let mut ready = Ready::EMPTY; #[cfg(all(target_os = "freebsd", feature = "net"))] { if event.is_aio() { ready |= Ready::READABLE; } if event.is_lio() { ready |= Ready::READABLE; } } if event.is_readable() { ready |= Ready::READABLE; } if event.is_writable() { ready |= Ready::WRITABLE; } if event.is_read_closed() { ready |= Ready::READ_CLOSED; } if event.is_write_closed() { ready |= Ready::WRITE_CLOSED; } if event.is_error() { ready |= Ready::ERROR; } #[cfg(any(target_os = "linux", target_os = "android"))] { if event.is_priority() { ready |= Ready::PRIORITY; } } ready } /// Returns true if `Ready` is the empty set. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(Ready::EMPTY.is_empty()); /// assert!(!Ready::READABLE.is_empty()); /// ``` pub fn is_empty(self) -> bool { self == Ready::EMPTY } /// Returns `true` if the value includes `readable`. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(!Ready::EMPTY.is_readable()); /// assert!(Ready::READABLE.is_readable()); /// assert!(Ready::READ_CLOSED.is_readable()); /// assert!(!Ready::WRITABLE.is_readable()); /// ``` pub fn is_readable(self) -> bool { self.contains(Ready::READABLE) || self.is_read_closed() } /// Returns `true` if the value includes writable `readiness`. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(!Ready::EMPTY.is_writable()); /// assert!(!Ready::READABLE.is_writable()); /// assert!(Ready::WRITABLE.is_writable()); /// assert!(Ready::WRITE_CLOSED.is_writable()); /// ``` pub fn is_writable(self) -> bool { self.contains(Ready::WRITABLE) || self.is_write_closed() } /// Returns `true` if the value includes read-closed `readiness`. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(!Ready::EMPTY.is_read_closed()); /// assert!(!Ready::READABLE.is_read_closed()); /// assert!(Ready::READ_CLOSED.is_read_closed()); /// ``` pub fn is_read_closed(self) -> bool { self.contains(Ready::READ_CLOSED) } /// Returns `true` if the value includes write-closed `readiness`. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(!Ready::EMPTY.is_write_closed()); /// assert!(!Ready::WRITABLE.is_write_closed()); /// assert!(Ready::WRITE_CLOSED.is_write_closed()); /// ``` pub fn is_write_closed(self) -> bool { self.contains(Ready::WRITE_CLOSED) } /// Returns `true` if the value includes priority `readiness`. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(!Ready::EMPTY.is_priority()); /// assert!(!Ready::WRITABLE.is_priority()); /// assert!(Ready::PRIORITY.is_priority()); /// ``` #[cfg(any(target_os = "linux", target_os = "android"))] #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))] pub fn is_priority(self) -> bool { self.contains(Ready::PRIORITY) } /// Returns `true` if the value includes error `readiness`. /// /// # Examples /// /// ``` /// use tokio::io::Ready; /// /// assert!(!Ready::EMPTY.is_error()); /// assert!(!Ready::WRITABLE.is_error()); /// assert!(Ready::ERROR.is_error()); /// ``` pub fn is_error(self) -> bool { self.contains(Ready::ERROR) } /// Returns true if `self` is a superset of `other`. /// /// `other` may represent more than one readiness operations, in which case /// the function only returns true if `self` contains all readiness /// specified in `other`. pub(crate) fn contains>(self, other: T) -> bool { let other = other.into(); (self & other) == other } /// Creates a `Ready` instance using the given `usize` representation. /// /// The `usize` representation must have been obtained from a call to /// `Readiness::as_usize`. /// /// This function is mainly provided to allow the caller to get a /// readiness value from an `AtomicUsize`. pub(crate) fn from_usize(val: usize) -> Ready { Ready(val & Ready::ALL.as_usize()) } /// Returns a `usize` representation of the `Ready` value. /// /// This function is mainly provided to allow the caller to store a /// readiness value in an `AtomicUsize`. pub(crate) fn as_usize(self) -> usize { self.0 } pub(crate) fn from_interest(interest: Interest) -> Ready { let mut ready = Ready::EMPTY; if interest.is_readable() { ready |= Ready::READABLE; ready |= Ready::READ_CLOSED; } if interest.is_writable() { ready |= Ready::WRITABLE; ready |= Ready::WRITE_CLOSED; } #[cfg(any(target_os = "linux", target_os = "android"))] if interest.is_priority() { ready |= Ready::PRIORITY; ready |= Ready::READ_CLOSED; } if interest.is_error() { ready |= Ready::ERROR; } ready } pub(crate) fn intersection(self, interest: Interest) -> Ready { Ready(self.0 & Ready::from_interest(interest).0) } pub(crate) fn satisfies(self, interest: Interest) -> bool { self.0 & Ready::from_interest(interest).0 != 0 } } impl ops::BitOr for Ready { type Output = Ready; #[inline] fn bitor(self, other: Ready) -> Ready { Ready(self.0 | other.0) } } impl ops::BitOrAssign for Ready { #[inline] fn bitor_assign(&mut self, other: Ready) { self.0 |= other.0; } } impl ops::BitAnd for Ready { type Output = Ready; #[inline] fn bitand(self, other: Ready) -> Ready { Ready(self.0 & other.0) } } impl ops::Sub for Ready { type Output = Ready; #[inline] fn sub(self, other: Ready) -> Ready { Ready(self.0 & !other.0) } } impl fmt::Debug for Ready { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { let mut fmt = fmt.debug_struct("Ready"); fmt.field("is_readable", &self.is_readable()) .field("is_writable", &self.is_writable()) .field("is_read_closed", &self.is_read_closed()) .field("is_write_closed", &self.is_write_closed()) .field("is_error", &self.is_error()); #[cfg(any(target_os = "linux", target_os = "android"))] fmt.field("is_priority", &self.is_priority()); fmt.finish() } } tokio-1.43.1/src/io/seek.rs000064400000000000000000000030631046102023000135340ustar 00000000000000use crate::io::AsyncSeek; use pin_project_lite::pin_project; use std::future::Future; use std::io::{self, SeekFrom}; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// Future for the [`seek`](crate::io::AsyncSeekExt::seek) method. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Seek<'a, S: ?Sized> { seek: &'a mut S, pos: Option, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn seek(seek: &mut S, pos: SeekFrom) -> Seek<'_, S> where S: AsyncSeek + ?Sized + Unpin, { Seek { seek, pos: Some(pos), _pin: PhantomPinned, } } impl Future for Seek<'_, S> where S: AsyncSeek + ?Sized + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); match me.pos { Some(pos) => { // ensure no seek in progress ready!(Pin::new(&mut *me.seek).poll_complete(cx))?; match Pin::new(&mut *me.seek).start_seek(*pos) { Ok(()) => { *me.pos = None; Pin::new(&mut *me.seek).poll_complete(cx) } Err(e) => Poll::Ready(Err(e)), } } None => Pin::new(&mut *me.seek).poll_complete(cx), } } } tokio-1.43.1/src/io/split.rs000064400000000000000000000111251046102023000137360ustar 00000000000000//! Split a single value implementing `AsyncRead + AsyncWrite` into separate //! `AsyncRead` and `AsyncWrite` handles. //! //! To restore this read/write object from its `split::ReadHalf` and //! `split::WriteHalf` use `unsplit`. use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; use std::fmt; use std::io; use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; use std::task::{Context, Poll}; cfg_io_util! { /// The readable half of a value returned from [`split`](split()). pub struct ReadHalf { inner: Arc>, } /// The writable half of a value returned from [`split`](split()). pub struct WriteHalf { inner: Arc>, } /// Splits a single value implementing `AsyncRead + AsyncWrite` into separate /// `AsyncRead` and `AsyncWrite` handles. /// /// To restore this read/write object from its `ReadHalf` and /// `WriteHalf` use [`unsplit`](ReadHalf::unsplit()). pub fn split(stream: T) -> (ReadHalf, WriteHalf) where T: AsyncRead + AsyncWrite, { let is_write_vectored = stream.is_write_vectored(); let inner = Arc::new(Inner { stream: Mutex::new(stream), is_write_vectored, }); let rd = ReadHalf { inner: inner.clone(), }; let wr = WriteHalf { inner }; (rd, wr) } } struct Inner { stream: Mutex, is_write_vectored: bool, } impl Inner { fn with_lock(&self, f: impl FnOnce(Pin<&mut T>) -> R) -> R { let mut guard = self.stream.lock().unwrap(); // safety: we do not move the stream. let stream = unsafe { Pin::new_unchecked(&mut *guard) }; f(stream) } } impl ReadHalf { /// Checks if this `ReadHalf` and some `WriteHalf` were split from the same /// stream. pub fn is_pair_of(&self, other: &WriteHalf) -> bool { other.is_pair_of(self) } /// Reunites with a previously split `WriteHalf`. /// /// # Panics /// /// If this `ReadHalf` and the given `WriteHalf` do not originate from the /// same `split` operation this method will panic. /// This can be checked ahead of time by calling [`is_pair_of()`](Self::is_pair_of). #[track_caller] pub fn unsplit(self, wr: WriteHalf) -> T where T: Unpin, { if self.is_pair_of(&wr) { drop(wr); let inner = Arc::try_unwrap(self.inner) .ok() .expect("`Arc::try_unwrap` failed"); inner.stream.into_inner().unwrap() } else { panic!("Unrelated `split::Write` passed to `split::Read::unsplit`.") } } } impl WriteHalf { /// Checks if this `WriteHalf` and some `ReadHalf` were split from the same /// stream. pub fn is_pair_of(&self, other: &ReadHalf) -> bool { Arc::ptr_eq(&self.inner, &other.inner) } } impl AsyncRead for ReadHalf { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.inner.with_lock(|stream| stream.poll_read(cx, buf)) } } impl AsyncWrite for WriteHalf { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.inner.with_lock(|stream| stream.poll_write(cx, buf)) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.inner.with_lock(|stream| stream.poll_flush(cx)) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.inner.with_lock(|stream| stream.poll_shutdown(cx)) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.inner .with_lock(|stream| stream.poll_write_vectored(cx, bufs)) } fn is_write_vectored(&self) -> bool { self.inner.is_write_vectored } } unsafe impl Send for ReadHalf {} unsafe impl Send for WriteHalf {} unsafe impl Sync for ReadHalf {} unsafe impl Sync for WriteHalf {} impl fmt::Debug for ReadHalf { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("split::ReadHalf").finish() } } impl fmt::Debug for WriteHalf { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("split::WriteHalf").finish() } } tokio-1.43.1/src/io/stderr.rs000064400000000000000000000077141046102023000141170ustar 00000000000000use crate::io::blocking::Blocking; use crate::io::stdio_common::SplitByUtf8BoundaryIfWindows; use crate::io::AsyncWrite; use std::io; use std::pin::Pin; use std::task::Context; use std::task::Poll; cfg_io_std! { /// A handle to the standard error stream of a process. /// /// Concurrent writes to stderr must be executed with care: Only individual /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular /// you should be aware that writes using [`write_all`] are not guaranteed /// to occur as a single write, so multiple threads writing data with /// [`write_all`] may result in interleaved output. /// /// Created by the [`stderr`] function. /// /// [`stderr`]: stderr() /// [`AsyncWrite`]: AsyncWrite /// [`write_all`]: crate::io::AsyncWriteExt::write_all() /// /// # Examples /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut stderr = io::stdout(); /// stderr.write_all(b"Print some error here.").await?; /// Ok(()) /// } /// ``` #[derive(Debug)] pub struct Stderr { std: SplitByUtf8BoundaryIfWindows>, } /// Constructs a new handle to the standard error of the current process. /// /// The returned handle allows writing to standard error from the within the /// Tokio runtime. /// /// Concurrent writes to stderr must be executed with care: Only individual /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular /// you should be aware that writes using [`write_all`] are not guaranteed /// to occur as a single write, so multiple threads writing data with /// [`write_all`] may result in interleaved output. /// /// [`AsyncWrite`]: AsyncWrite /// [`write_all`]: crate::io::AsyncWriteExt::write_all() /// /// # Examples /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut stderr = io::stderr(); /// stderr.write_all(b"Print some error here.").await?; /// Ok(()) /// } /// ``` pub fn stderr() -> Stderr { let std = io::stderr(); // SAFETY: The `Read` implementation of `std` does not read from the // buffer it is borrowing and correctly reports the length of the data // written into the buffer. let blocking = unsafe { Blocking::new(std) }; Stderr { std: SplitByUtf8BoundaryIfWindows::new(blocking), } } } #[cfg(unix)] mod sys { use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, RawFd}; use super::Stderr; impl AsRawFd for Stderr { fn as_raw_fd(&self) -> RawFd { std::io::stderr().as_raw_fd() } } impl AsFd for Stderr { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } cfg_windows! { use crate::os::windows::io::{AsHandle, BorrowedHandle, AsRawHandle, RawHandle}; impl AsRawHandle for Stderr { fn as_raw_handle(&self) -> RawHandle { std::io::stderr().as_raw_handle() } } impl AsHandle for Stderr { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } } } } impl AsyncWrite for Stderr { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut self.std).poll_write(cx, buf) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.std).poll_flush(cx) } fn poll_shutdown( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { Pin::new(&mut self.std).poll_shutdown(cx) } } tokio-1.43.1/src/io/stdin.rs000064400000000000000000000060211046102023000137230ustar 00000000000000use crate::io::blocking::Blocking; use crate::io::{AsyncRead, ReadBuf}; use std::io; use std::pin::Pin; use std::task::Context; use std::task::Poll; cfg_io_std! { /// A handle to the standard input stream of a process. /// /// The handle implements the [`AsyncRead`] trait, but beware that concurrent /// reads of `Stdin` must be executed with care. /// /// This handle is best used for non-interactive uses, such as when a file /// is piped into the application. For technical reasons, `stdin` is /// implemented by using an ordinary blocking read on a separate thread, and /// it is impossible to cancel that read. This can make shutdown of the /// runtime hang until the user presses enter. /// /// For interactive uses, it is recommended to spawn a thread dedicated to /// user input and use blocking IO directly in that thread. /// /// Created by the [`stdin`] function. /// /// [`stdin`]: fn@stdin /// [`AsyncRead`]: trait@AsyncRead #[derive(Debug)] pub struct Stdin { std: Blocking, } /// Constructs a new handle to the standard input of the current process. /// /// This handle is best used for non-interactive uses, such as when a file /// is piped into the application. For technical reasons, `stdin` is /// implemented by using an ordinary blocking read on a separate thread, and /// it is impossible to cancel that read. This can make shutdown of the /// runtime hang until the user presses enter. /// /// For interactive uses, it is recommended to spawn a thread dedicated to /// user input and use blocking IO directly in that thread. pub fn stdin() -> Stdin { let std = io::stdin(); // SAFETY: The `Read` implementation of `std` does not read from the // buffer it is borrowing and correctly reports the length of the data // written into the buffer. let std = unsafe { Blocking::new(std) }; Stdin { std, } } } #[cfg(unix)] mod sys { use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, RawFd}; use super::Stdin; impl AsRawFd for Stdin { fn as_raw_fd(&self) -> RawFd { std::io::stdin().as_raw_fd() } } impl AsFd for Stdin { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } cfg_windows! { use crate::os::windows::io::{AsHandle, BorrowedHandle, AsRawHandle, RawHandle}; impl AsRawHandle for Stdin { fn as_raw_handle(&self) -> RawHandle { std::io::stdin().as_raw_handle() } } impl AsHandle for Stdin { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } } } } impl AsyncRead for Stdin { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut self.std).poll_read(cx, buf) } } tokio-1.43.1/src/io/stdio_common.rs000064400000000000000000000170231046102023000153000ustar 00000000000000//! Contains utilities for stdout and stderr. use crate::io::AsyncWrite; use std::pin::Pin; use std::task::{Context, Poll}; /// # Windows /// [`AsyncWrite`] adapter that finds last char boundary in given buffer and does not write the rest, /// if buffer contents seems to be `utf8`. Otherwise it only trims buffer down to `DEFAULT_MAX_BUF_SIZE`. /// That's why, wrapped writer will always receive well-formed utf-8 bytes. /// # Other platforms /// Passes data to `inner` as is. #[derive(Debug)] pub(crate) struct SplitByUtf8BoundaryIfWindows { inner: W, } impl SplitByUtf8BoundaryIfWindows { pub(crate) fn new(inner: W) -> Self { Self { inner } } } // this constant is defined by Unicode standard. const MAX_BYTES_PER_CHAR: usize = 4; // Subject for tweaking here const MAGIC_CONST: usize = 8; impl crate::io::AsyncWrite for SplitByUtf8BoundaryIfWindows where W: AsyncWrite + Unpin, { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, mut buf: &[u8], ) -> Poll> { // just a closure to avoid repetitive code let mut call_inner = move |buf| Pin::new(&mut self.inner).poll_write(cx, buf); // 1. Only windows stdio can suffer from non-utf8. // We also check for `test` so that we can write some tests // for further code. Since `AsyncWrite` can always shrink // buffer at its discretion, excessive (i.e. in tests) shrinking // does not break correctness. // 2. If buffer is small, it will not be shrunk. // That's why, it's "textness" will not change, so we don't have // to fixup it. if cfg!(not(any(target_os = "windows", test))) || buf.len() <= crate::io::blocking::DEFAULT_MAX_BUF_SIZE { return call_inner(buf); } buf = &buf[..crate::io::blocking::DEFAULT_MAX_BUF_SIZE]; // Now there are two possibilities. // If caller gave is binary buffer, we **should not** shrink it // anymore, because excessive shrinking hits performance. // If caller gave as binary buffer, we **must** additionally // shrink it to strip incomplete char at the end of buffer. // that's why check we will perform now is allowed to have // false-positive. // Now let's look at the first MAX_BYTES_PER_CHAR * MAGIC_CONST bytes. // if they are (possibly incomplete) utf8, then we can be quite sure // that input buffer was utf8. let have_to_fix_up = match std::str::from_utf8(&buf[..MAX_BYTES_PER_CHAR * MAGIC_CONST]) { Ok(_) => true, Err(err) => { let incomplete_bytes = MAX_BYTES_PER_CHAR * MAGIC_CONST - err.valid_up_to(); incomplete_bytes < MAX_BYTES_PER_CHAR } }; if have_to_fix_up { // We must pop several bytes at the end which form incomplete // character. To achieve it, we exploit UTF8 encoding: // for any code point, all bytes except first start with 0b10 prefix. // see https://en.wikipedia.org/wiki/UTF-8#Encoding for details let trailing_incomplete_char_size = buf .iter() .rev() .take(MAX_BYTES_PER_CHAR) .position(|byte| *byte < 0b1000_0000 || *byte >= 0b1100_0000) .unwrap_or(0) + 1; buf = &buf[..buf.len() - trailing_incomplete_char_size]; } call_inner(buf) } fn poll_flush( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { Pin::new(&mut self.inner).poll_flush(cx) } fn poll_shutdown( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { Pin::new(&mut self.inner).poll_shutdown(cx) } } #[cfg(test)] #[cfg(not(loom))] mod tests { use crate::io::blocking::DEFAULT_MAX_BUF_SIZE; use crate::io::AsyncWriteExt; use std::io; use std::pin::Pin; use std::task::Context; use std::task::Poll; struct TextMockWriter; impl crate::io::AsyncWrite for TextMockWriter { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { assert!(buf.len() <= DEFAULT_MAX_BUF_SIZE); assert!(std::str::from_utf8(buf).is_ok()); Poll::Ready(Ok(buf.len())) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown( self: Pin<&mut Self>, _cx: &mut Context<'_>, ) -> Poll> { Poll::Ready(Ok(())) } } struct LoggingMockWriter { write_history: Vec, } impl LoggingMockWriter { fn new() -> Self { LoggingMockWriter { write_history: Vec::new(), } } } impl crate::io::AsyncWrite for LoggingMockWriter { fn poll_write( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { assert!(buf.len() <= DEFAULT_MAX_BUF_SIZE); self.write_history.push(buf.len()); Poll::Ready(Ok(buf.len())) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown( self: Pin<&mut Self>, _cx: &mut Context<'_>, ) -> Poll> { Poll::Ready(Ok(())) } } #[test] #[cfg_attr(miri, ignore)] fn test_splitter() { let data = str::repeat("â–ˆ", DEFAULT_MAX_BUF_SIZE); let mut wr = super::SplitByUtf8BoundaryIfWindows::new(TextMockWriter); let fut = async move { wr.write_all(data.as_bytes()).await.unwrap(); }; crate::runtime::Builder::new_current_thread() .build() .unwrap() .block_on(fut); } #[test] #[cfg_attr(miri, ignore)] fn test_pseudo_text() { // In this test we write a piece of binary data, whose beginning is // text though. We then validate that even in this corner case buffer // was not shrunk too much. let checked_count = super::MAGIC_CONST * super::MAX_BYTES_PER_CHAR; let mut data: Vec = str::repeat("a", checked_count).into(); data.extend(std::iter::repeat(0b1010_1010).take(DEFAULT_MAX_BUF_SIZE - checked_count + 1)); let mut writer = LoggingMockWriter::new(); let mut splitter = super::SplitByUtf8BoundaryIfWindows::new(&mut writer); crate::runtime::Builder::new_current_thread() .build() .unwrap() .block_on(async { splitter.write_all(&data).await.unwrap(); }); // Check that at most two writes were performed assert!(writer.write_history.len() <= 2); // Check that all has been written assert_eq!( writer.write_history.iter().copied().sum::(), data.len() ); // Check that at most MAX_BYTES_PER_CHAR + 1 (i.e. 5) bytes were shrunk // from the buffer: one because it was outside of DEFAULT_MAX_BUF_SIZE boundary, and // up to one "utf8 code point". assert!(data.len() - writer.write_history[0] <= super::MAX_BYTES_PER_CHAR + 1); } } tokio-1.43.1/src/io/stdout.rs000064400000000000000000000134511046102023000141310ustar 00000000000000use crate::io::blocking::Blocking; use crate::io::stdio_common::SplitByUtf8BoundaryIfWindows; use crate::io::AsyncWrite; use std::io; use std::pin::Pin; use std::task::Context; use std::task::Poll; cfg_io_std! { /// A handle to the standard output stream of a process. /// /// Concurrent writes to stdout must be executed with care: Only individual /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular /// you should be aware that writes using [`write_all`] are not guaranteed /// to occur as a single write, so multiple threads writing data with /// [`write_all`] may result in interleaved output. /// /// Created by the [`stdout`] function. /// /// [`stdout`]: stdout() /// [`AsyncWrite`]: AsyncWrite /// [`write_all`]: crate::io::AsyncWriteExt::write_all() /// /// # Examples /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut stdout = io::stdout(); /// stdout.write_all(b"Hello world!").await?; /// Ok(()) /// } /// ``` /// /// The following is an example of using `stdio` with loop. /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() { /// let messages = vec!["hello", " world\n"]; /// /// // When you use `stdio` in a loop, it is recommended to create /// // a single `stdio` instance outside the loop and call a write /// // operation against that instance on each loop. /// // /// // Repeatedly creating `stdout` instances inside the loop and /// // writing to that handle could result in mangled output since /// // each write operation is handled by a different blocking thread. /// let mut stdout = io::stdout(); /// /// for message in &messages { /// stdout.write_all(message.as_bytes()).await.unwrap(); /// stdout.flush().await.unwrap(); /// } /// } /// ``` #[derive(Debug)] pub struct Stdout { std: SplitByUtf8BoundaryIfWindows>, } /// Constructs a new handle to the standard output of the current process. /// /// The returned handle allows writing to standard out from the within the /// Tokio runtime. /// /// Concurrent writes to stdout must be executed with care: Only individual /// writes to this [`AsyncWrite`] are guaranteed to be intact. In particular /// you should be aware that writes using [`write_all`] are not guaranteed /// to occur as a single write, so multiple threads writing data with /// [`write_all`] may result in interleaved output. /// /// [`AsyncWrite`]: AsyncWrite /// [`write_all`]: crate::io::AsyncWriteExt::write_all() /// /// # Examples /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut stdout = io::stdout(); /// stdout.write_all(b"Hello world!").await?; /// Ok(()) /// } /// ``` /// /// The following is an example of using `stdio` with loop. /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() { /// let messages = vec!["hello", " world\n"]; /// /// // When you use `stdio` in a loop, it is recommended to create /// // a single `stdio` instance outside the loop and call a write /// // operation against that instance on each loop. /// // /// // Repeatedly creating `stdout` instances inside the loop and /// // writing to that handle could result in mangled output since /// // each write operation is handled by a different blocking thread. /// let mut stdout = io::stdout(); /// /// for message in &messages { /// stdout.write_all(message.as_bytes()).await.unwrap(); /// stdout.flush().await.unwrap(); /// } /// } /// ``` pub fn stdout() -> Stdout { let std = io::stdout(); // SAFETY: The `Read` implementation of `std` does not read from the // buffer it is borrowing and correctly reports the length of the data // written into the buffer. let blocking = unsafe { Blocking::new(std) }; Stdout { std: SplitByUtf8BoundaryIfWindows::new(blocking), } } } #[cfg(unix)] mod sys { use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, RawFd}; use super::Stdout; impl AsRawFd for Stdout { fn as_raw_fd(&self) -> RawFd { std::io::stdout().as_raw_fd() } } impl AsFd for Stdout { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } cfg_windows! { use crate::os::windows::io::{AsHandle, BorrowedHandle, AsRawHandle, RawHandle}; impl AsRawHandle for Stdout { fn as_raw_handle(&self) -> RawHandle { std::io::stdout().as_raw_handle() } } impl AsHandle for Stdout { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } } } } impl AsyncWrite for Stdout { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut self.std).poll_write(cx, buf) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.std).poll_flush(cx) } fn poll_shutdown( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { Pin::new(&mut self.std).poll_shutdown(cx) } } tokio-1.43.1/src/io/util/async_buf_read_ext.rs000064400000000000000000000335141046102023000174120ustar 00000000000000use crate::io::util::fill_buf::{fill_buf, FillBuf}; use crate::io::util::lines::{lines, Lines}; use crate::io::util::read_line::{read_line, ReadLine}; use crate::io::util::read_until::{read_until, ReadUntil}; use crate::io::util::split::{split, Split}; use crate::io::AsyncBufRead; cfg_io_util! { /// An extension trait which adds utility methods to [`AsyncBufRead`] types. /// /// [`AsyncBufRead`]: crate::io::AsyncBufRead pub trait AsyncBufReadExt: AsyncBufRead { /// Reads all bytes into `buf` until the delimiter `byte` or EOF is reached. /// /// Equivalent to: /// /// ```ignore /// async fn read_until(&mut self, byte: u8, buf: &mut Vec) -> io::Result; /// ``` /// /// This function will read bytes from the underlying stream until the /// delimiter or EOF is found. Once found, all bytes up to, and including, /// the delimiter (if found) will be appended to `buf`. /// /// If successful, this function will return the total number of bytes read. /// /// If this function returns `Ok(0)`, the stream has reached EOF. /// /// # Errors /// /// This function will ignore all instances of [`ErrorKind::Interrupted`] and /// will otherwise return any errors returned by [`fill_buf`]. /// /// If an I/O error is encountered then all bytes read so far will be /// present in `buf` and its length will have been adjusted appropriately. /// /// [`fill_buf`]: AsyncBufRead::poll_fill_buf /// [`ErrorKind::Interrupted`]: std::io::ErrorKind::Interrupted /// /// # Cancel safety /// /// If the method is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then some data may have been partially read. Any /// partially read bytes are appended to `buf`, and the method can be /// called again to continue reading until `byte`. /// /// This method returns the total number of bytes read. If you cancel /// the call to `read_until` and then call it again to continue reading, /// the counter is reset. /// /// # Examples /// /// [`std::io::Cursor`][`Cursor`] is a type that implements `BufRead`. In /// this example, we use [`Cursor`] to read all the bytes in a byte slice /// in hyphen delimited segments: /// /// [`Cursor`]: std::io::Cursor /// /// ``` /// use tokio::io::AsyncBufReadExt; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() { /// let mut cursor = Cursor::new(b"lorem-ipsum"); /// let mut buf = vec![]; /// /// // cursor is at 'l' /// let num_bytes = cursor.read_until(b'-', &mut buf) /// .await /// .expect("reading from cursor won't fail"); /// /// assert_eq!(num_bytes, 6); /// assert_eq!(buf, b"lorem-"); /// buf.clear(); /// /// // cursor is at 'i' /// let num_bytes = cursor.read_until(b'-', &mut buf) /// .await /// .expect("reading from cursor won't fail"); /// /// assert_eq!(num_bytes, 5); /// assert_eq!(buf, b"ipsum"); /// buf.clear(); /// /// // cursor is at EOF /// let num_bytes = cursor.read_until(b'-', &mut buf) /// .await /// .expect("reading from cursor won't fail"); /// assert_eq!(num_bytes, 0); /// assert_eq!(buf, b""); /// } /// ``` fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut Vec) -> ReadUntil<'a, Self> where Self: Unpin, { read_until(self, byte, buf) } /// Reads all bytes until a newline (the 0xA byte) is reached, and append /// them to the provided buffer. /// /// Equivalent to: /// /// ```ignore /// async fn read_line(&mut self, buf: &mut String) -> io::Result; /// ``` /// /// This function will read bytes from the underlying stream until the /// newline delimiter (the 0xA byte) or EOF is found. Once found, all bytes /// up to, and including, the delimiter (if found) will be appended to /// `buf`. /// /// If successful, this function will return the total number of bytes read. /// /// If this function returns `Ok(0)`, the stream has reached EOF. /// /// # Errors /// /// This function has the same error semantics as [`read_until`] and will /// also return an error if the read bytes are not valid UTF-8. If an I/O /// error is encountered then `buf` may contain some bytes already read in /// the event that all data read so far was valid UTF-8. /// /// [`read_until`]: AsyncBufReadExt::read_until /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may have been partially /// read, and this data is lost. There are no guarantees regarding the /// contents of `buf` when the call is cancelled. The current /// implementation replaces `buf` with the empty string, but this may /// change in the future. /// /// This function does not behave like [`read_until`] because of the /// requirement that a string contains only valid utf-8. If you need a /// cancellation safe `read_line`, there are three options: /// /// * Call [`read_until`] with a newline character and manually perform the utf-8 check. /// * The stream returned by [`lines`] has a cancellation safe /// [`next_line`] method. /// * Use [`tokio_util::codec::LinesCodec`][LinesCodec]. /// /// [LinesCodec]: https://docs.rs/tokio-util/latest/tokio_util/codec/struct.LinesCodec.html /// [`read_until`]: Self::read_until /// [`lines`]: Self::lines /// [`next_line`]: crate::io::Lines::next_line /// /// # Examples /// /// [`std::io::Cursor`][`Cursor`] is a type that implements /// `AsyncBufRead`. In this example, we use [`Cursor`] to read all the /// lines in a byte slice: /// /// [`Cursor`]: std::io::Cursor /// /// ``` /// use tokio::io::AsyncBufReadExt; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() { /// let mut cursor = Cursor::new(b"foo\nbar"); /// let mut buf = String::new(); /// /// // cursor is at 'f' /// let num_bytes = cursor.read_line(&mut buf) /// .await /// .expect("reading from cursor won't fail"); /// /// assert_eq!(num_bytes, 4); /// assert_eq!(buf, "foo\n"); /// buf.clear(); /// /// // cursor is at 'b' /// let num_bytes = cursor.read_line(&mut buf) /// .await /// .expect("reading from cursor won't fail"); /// /// assert_eq!(num_bytes, 3); /// assert_eq!(buf, "bar"); /// buf.clear(); /// /// // cursor is at EOF /// let num_bytes = cursor.read_line(&mut buf) /// .await /// .expect("reading from cursor won't fail"); /// /// assert_eq!(num_bytes, 0); /// assert_eq!(buf, ""); /// } /// ``` fn read_line<'a>(&'a mut self, buf: &'a mut String) -> ReadLine<'a, Self> where Self: Unpin, { read_line(self, buf) } /// Returns a stream of the contents of this reader split on the byte /// `byte`. /// /// This method is the asynchronous equivalent to /// [`BufRead::split`](std::io::BufRead::split). /// /// The stream returned from this function will yield instances of /// [`io::Result`]`<`[`Option`]`<`[`Vec`]`>>`. Each vector returned will *not* have /// the delimiter byte at the end. /// /// [`io::Result`]: std::io::Result /// [`Option`]: core::option::Option /// [`Vec`]: std::vec::Vec /// /// # Errors /// /// Each item of the stream has the same error semantics as /// [`AsyncBufReadExt::read_until`](AsyncBufReadExt::read_until). /// /// # Examples /// /// ``` /// # use tokio::io::AsyncBufRead; /// use tokio::io::AsyncBufReadExt; /// /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> { /// let mut segments = my_buf_read.split(b'f'); /// /// while let Some(segment) = segments.next_segment().await? { /// println!("length = {}", segment.len()) /// } /// # Ok(()) /// # } /// ``` fn split(self, byte: u8) -> Split where Self: Sized + Unpin, { split(self, byte) } /// Returns the contents of the internal buffer, filling it with more /// data from the inner reader if it is empty. /// /// This function is a lower-level call. It needs to be paired with the /// [`consume`] method to function properly. When calling this method, /// none of the contents will be "read" in the sense that later calling /// `read` may return the same contents. As such, [`consume`] must be /// called with the number of bytes that are consumed from this buffer /// to ensure that the bytes are never returned twice. /// /// An empty buffer returned indicates that the stream has reached EOF. /// /// Equivalent to: /// /// ```ignore /// async fn fill_buf(&mut self) -> io::Result<&[u8]>; /// ``` /// /// # Errors /// /// This function will return an I/O error if the underlying reader was /// read, but returned an error. /// /// # Cancel safety /// /// This method is cancel safe. If you use it as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no data was read. /// /// [`consume`]: crate::io::AsyncBufReadExt::consume fn fill_buf(&mut self) -> FillBuf<'_, Self> where Self: Unpin, { fill_buf(self) } /// Tells this buffer that `amt` bytes have been consumed from the /// buffer, so they should no longer be returned in calls to [`read`]. /// /// This function is a lower-level call. It needs to be paired with the /// [`fill_buf`] method to function properly. This function does not /// perform any I/O, it simply informs this object that some amount of /// its buffer, returned from [`fill_buf`], has been consumed and should /// no longer be returned. As such, this function may do odd things if /// [`fill_buf`] isn't called before calling it. /// /// The `amt` must be less than the number of bytes in the buffer /// returned by [`fill_buf`]. /// /// [`read`]: crate::io::AsyncReadExt::read /// [`fill_buf`]: crate::io::AsyncBufReadExt::fill_buf fn consume(&mut self, amt: usize) where Self: Unpin, { std::pin::Pin::new(self).consume(amt); } /// Returns a stream over the lines of this reader. /// This method is the async equivalent to [`BufRead::lines`](std::io::BufRead::lines). /// /// The stream returned from this function will yield instances of /// [`io::Result`]`<`[`Option`]`<`[`String`]`>>`. Each string returned will *not* have a newline /// byte (the 0xA byte) or `CRLF` (0xD, 0xA bytes) at the end. /// /// [`io::Result`]: std::io::Result /// [`Option`]: core::option::Option /// [`String`]: String /// /// # Errors /// /// Each line of the stream has the same error semantics as [`AsyncBufReadExt::read_line`]. /// /// # Examples /// /// [`std::io::Cursor`][`Cursor`] is a type that implements `BufRead`. In /// this example, we use [`Cursor`] to iterate over all the lines in a byte /// slice. /// /// [`Cursor`]: std::io::Cursor /// /// ``` /// use tokio::io::AsyncBufReadExt; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() { /// let cursor = Cursor::new(b"lorem\nipsum\r\ndolor"); /// /// let mut lines = cursor.lines(); /// /// assert_eq!(lines.next_line().await.unwrap(), Some(String::from("lorem"))); /// assert_eq!(lines.next_line().await.unwrap(), Some(String::from("ipsum"))); /// assert_eq!(lines.next_line().await.unwrap(), Some(String::from("dolor"))); /// assert_eq!(lines.next_line().await.unwrap(), None); /// } /// ``` /// /// [`AsyncBufReadExt::read_line`]: AsyncBufReadExt::read_line fn lines(self) -> Lines where Self: Sized, { lines(self) } } } impl AsyncBufReadExt for R {} tokio-1.43.1/src/io/util/async_read_ext.rs000064400000000000000000001467441046102023000165700ustar 00000000000000use crate::io::util::chain::{chain, Chain}; use crate::io::util::read::{read, Read}; use crate::io::util::read_buf::{read_buf, ReadBuf}; use crate::io::util::read_exact::{read_exact, ReadExact}; use crate::io::util::read_int::{ReadF32, ReadF32Le, ReadF64, ReadF64Le}; use crate::io::util::read_int::{ ReadI128, ReadI128Le, ReadI16, ReadI16Le, ReadI32, ReadI32Le, ReadI64, ReadI64Le, ReadI8, }; use crate::io::util::read_int::{ ReadU128, ReadU128Le, ReadU16, ReadU16Le, ReadU32, ReadU32Le, ReadU64, ReadU64Le, ReadU8, }; use crate::io::util::read_to_end::{read_to_end, ReadToEnd}; use crate::io::util::read_to_string::{read_to_string, ReadToString}; use crate::io::util::take::{take, Take}; use crate::io::AsyncRead; use bytes::BufMut; cfg_io_util! { /// Defines numeric reader macro_rules! read_impl { ( $( $(#[$outer:meta])* fn $name:ident(&mut self) -> $($fut:ident)*; )* ) => { $( $(#[$outer])* fn $name(&mut self) -> $($fut)*<&mut Self> where Self: Unpin { $($fut)*::new(self) } )* } } /// Reads bytes from a source. /// /// Implemented as an extension trait, adding utility methods to all /// [`AsyncRead`] types. Callers will tend to import this trait instead of /// [`AsyncRead`]. /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::{self, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; /// let mut buffer = [0; 10]; /// /// // The `read` method is defined by this trait. /// let n = f.read(&mut buffer[..]).await?; /// /// Ok(()) /// } /// ``` /// /// See [module][crate::io] documentation for more details. /// /// [`AsyncRead`]: AsyncRead pub trait AsyncReadExt: AsyncRead { /// Creates a new `AsyncRead` instance that chains this stream with /// `next`. /// /// The returned `AsyncRead` instance will first read all bytes from this object /// until EOF is encountered. Afterwards the output is equivalent to the /// output of `next`. /// /// # Examples /// /// [`File`][crate::fs::File]s implement `AsyncRead`: /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::{self, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let f1 = File::open("foo.txt").await?; /// let f2 = File::open("bar.txt").await?; /// /// let mut handle = f1.chain(f2); /// let mut buffer = String::new(); /// /// // read the value into a String. We could use any AsyncRead /// // method here, this is just one example. /// handle.read_to_string(&mut buffer).await?; /// Ok(()) /// } /// ``` fn chain(self, next: R) -> Chain where Self: Sized, R: AsyncRead, { chain(self, next) } /// Pulls some bytes from this source into the specified buffer, /// returning how many bytes were read. /// /// Equivalent to: /// /// ```ignore /// async fn read(&mut self, buf: &mut [u8]) -> io::Result; /// ``` /// /// This method does not provide any guarantees about whether it /// completes immediately or asynchronously. /// /// # Return /// /// If the return value of this method is `Ok(n)`, then it must be /// guaranteed that `0 <= n <= buf.len()`. A nonzero `n` value indicates /// that the buffer `buf` has been filled in with `n` bytes of data from /// this source. If `n` is `0`, then it can indicate one of two /// scenarios: /// /// 1. This reader has reached its "end of file" and will likely no longer /// be able to produce bytes. Note that this does not mean that the /// reader will *always* no longer be able to produce bytes. /// 2. The buffer specified was 0 bytes in length. /// /// No guarantees are provided about the contents of `buf` when this /// function is called, implementations cannot rely on any property of the /// contents of `buf` being true. It is recommended that *implementations* /// only write data to `buf` instead of reading its contents. /// /// Correspondingly, however, *callers* of this method may not assume /// any guarantees about how the implementation uses `buf`. It is /// possible that the code that's supposed to write to the buffer might /// also read from it. It is your responsibility to make sure that `buf` /// is initialized before calling `read`. /// /// # Errors /// /// If this function encounters any form of I/O or other error, an error /// variant will be returned. If an error is returned then it must be /// guaranteed that no bytes were read. /// /// # Cancel safety /// /// This method is cancel safe. If you use it as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no data was read. /// /// # Examples /// /// [`File`][crate::fs::File]s implement `Read`: /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::{self, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; /// let mut buffer = [0; 10]; /// /// // read up to 10 bytes /// let n = f.read(&mut buffer[..]).await?; /// /// println!("The bytes: {:?}", &buffer[..n]); /// Ok(()) /// } /// ``` fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Read<'a, Self> where Self: Unpin, { read(self, buf) } /// Pulls some bytes from this source into the specified buffer, /// advancing the buffer's internal cursor. /// /// Equivalent to: /// /// ```ignore /// async fn read_buf(&mut self, buf: &mut B) -> io::Result; /// ``` /// /// Usually, only a single `read` syscall is issued, even if there is /// more space in the supplied buffer. /// /// This method does not provide any guarantees about whether it /// completes immediately or asynchronously. /// /// # Return /// /// A nonzero `n` value indicates that the buffer `buf` has been filled /// in with `n` bytes of data from this source. If `n` is `0`, then it /// can indicate one of two scenarios: /// /// 1. This reader has reached its "end of file" and will likely no longer /// be able to produce bytes. Note that this does not mean that the /// reader will *always* no longer be able to produce bytes. /// 2. The buffer specified had a remaining capacity of zero. /// /// # Errors /// /// If this function encounters any form of I/O or other error, an error /// variant will be returned. If an error is returned then it must be /// guaranteed that no bytes were read. /// /// # Cancel safety /// /// This method is cancel safe. If you use it as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no data was read. /// /// # Examples /// /// [`File`] implements `Read` and [`BytesMut`] implements [`BufMut`]: /// /// [`File`]: crate::fs::File /// [`BytesMut`]: bytes::BytesMut /// [`BufMut`]: bytes::BufMut /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::{self, AsyncReadExt}; /// /// use bytes::BytesMut; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; /// let mut buffer = BytesMut::with_capacity(10); /// /// assert!(buffer.is_empty()); /// assert!(buffer.capacity() >= 10); /// /// // note that the return value is not needed to access the data /// // that was read as `buffer`'s internal cursor is updated. /// // /// // this might read more than 10 bytes if the capacity of `buffer` /// // is larger than 10. /// f.read_buf(&mut buffer).await?; /// /// println!("The bytes: {:?}", &buffer[..]); /// Ok(()) /// } /// ``` fn read_buf<'a, B>(&'a mut self, buf: &'a mut B) -> ReadBuf<'a, Self, B> where Self: Unpin, B: BufMut + ?Sized, { read_buf(self, buf) } /// Reads the exact number of bytes required to fill `buf`. /// /// Equivalent to: /// /// ```ignore /// async fn read_exact(&mut self, buf: &mut [u8]) -> io::Result; /// ``` /// /// This function reads as many bytes as necessary to completely fill /// the specified buffer `buf`. /// /// # Errors /// /// If the operation encounters an "end of file" before completely /// filling the buffer, it returns an error of the kind /// [`ErrorKind::UnexpectedEof`]. The contents of `buf` are unspecified /// in this case. /// /// If any other read error is encountered then the operation /// immediately returns. The contents of `buf` are unspecified in this /// case. /// /// If this operation returns an error, it is unspecified how many bytes /// it has read, but it will never read more than would be necessary to /// completely fill the buffer. /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may already have been /// read into `buf`. /// /// # Examples /// /// [`File`][crate::fs::File]s implement `Read`: /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::{self, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; /// let len = 10; /// let mut buffer = vec![0; len]; /// /// // read exactly 10 bytes /// f.read_exact(&mut buffer).await?; /// Ok(()) /// } /// ``` /// /// [`ErrorKind::UnexpectedEof`]: std::io::ErrorKind::UnexpectedEof fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadExact<'a, Self> where Self: Unpin, { read_exact(self, buf) } read_impl! { /// Reads an unsigned 8 bit integer from the underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u8(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is cancel safe. If this method is used as an event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no data were read. /// /// # Examples /// /// Read unsigned 8 bit integers from an `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![2, 5]); /// /// assert_eq!(2, reader.read_u8().await?); /// assert_eq!(5, reader.read_u8().await?); /// /// Ok(()) /// } /// ``` fn read_u8(&mut self) -> ReadU8; /// Reads a signed 8 bit integer from the underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i8(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is cancel safe. If this method is used as an event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no data were read. /// /// # Examples /// /// Read unsigned 8 bit integers from an `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x02, 0xfb]); /// /// assert_eq!(2, reader.read_i8().await?); /// assert_eq!(-5, reader.read_i8().await?); /// /// Ok(()) /// } /// ``` fn read_i8(&mut self) -> ReadI8; /// Reads an unsigned 16-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u16(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 16 bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![2, 5, 3, 0]); /// /// assert_eq!(517, reader.read_u16().await?); /// assert_eq!(768, reader.read_u16().await?); /// Ok(()) /// } /// ``` fn read_u16(&mut self) -> ReadU16; /// Reads a signed 16-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i16(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 16 bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x00, 0xc1, 0xff, 0x7c]); /// /// assert_eq!(193, reader.read_i16().await?); /// assert_eq!(-132, reader.read_i16().await?); /// Ok(()) /// } /// ``` fn read_i16(&mut self) -> ReadI16; /// Reads an unsigned 32-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u32(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 32-bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x00, 0x00, 0x01, 0x0b]); /// /// assert_eq!(267, reader.read_u32().await?); /// Ok(()) /// } /// ``` fn read_u32(&mut self) -> ReadU32; /// Reads a signed 32-bit integer in big-endian order from the /// underlying reader. /// /// /// Equivalent to: /// /// ```ignore /// async fn read_i32(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 32-bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0xff, 0xff, 0x7a, 0x33]); /// /// assert_eq!(-34253, reader.read_i32().await?); /// Ok(()) /// } /// ``` fn read_i32(&mut self) -> ReadI32; /// Reads an unsigned 64-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u64(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 64-bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 /// ]); /// /// assert_eq!(918733457491587, reader.read_u64().await?); /// Ok(()) /// } /// ``` fn read_u64(&mut self) -> ReadU64; /// Reads an signed 64-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i64(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 64-bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x80, 0, 0, 0, 0, 0, 0, 0]); /// /// assert_eq!(i64::MIN, reader.read_i64().await?); /// Ok(()) /// } /// ``` fn read_i64(&mut self) -> ReadI64; /// Reads an unsigned 128-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u128(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 128-bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83, /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 /// ]); /// /// assert_eq!(16947640962301618749969007319746179, reader.read_u128().await?); /// Ok(()) /// } /// ``` fn read_u128(&mut self) -> ReadU128; /// Reads an signed 128-bit integer in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i128(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 128-bit big-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0x80, 0, 0, 0, 0, 0, 0, 0, /// 0, 0, 0, 0, 0, 0, 0, 0 /// ]); /// /// assert_eq!(i128::MIN, reader.read_i128().await?); /// Ok(()) /// } /// ``` fn read_i128(&mut self) -> ReadI128; /// Reads an 32-bit floating point type in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_f32(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read 32-bit floating point type from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0xff, 0x7f, 0xff, 0xff]); /// /// assert_eq!(f32::MIN, reader.read_f32().await?); /// Ok(()) /// } /// ``` fn read_f32(&mut self) -> ReadF32; /// Reads an 64-bit floating point type in big-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_f64(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read 64-bit floating point type from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0xff, 0xef, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff /// ]); /// /// assert_eq!(f64::MIN, reader.read_f64().await?); /// Ok(()) /// } /// ``` fn read_f64(&mut self) -> ReadF64; /// Reads an unsigned 16-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u16_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 16 bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![2, 5, 3, 0]); /// /// assert_eq!(1282, reader.read_u16_le().await?); /// assert_eq!(3, reader.read_u16_le().await?); /// Ok(()) /// } /// ``` fn read_u16_le(&mut self) -> ReadU16Le; /// Reads a signed 16-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i16_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 16 bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x00, 0xc1, 0xff, 0x7c]); /// /// assert_eq!(-16128, reader.read_i16_le().await?); /// assert_eq!(31999, reader.read_i16_le().await?); /// Ok(()) /// } /// ``` fn read_i16_le(&mut self) -> ReadI16Le; /// Reads an unsigned 32-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u32_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 32-bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x00, 0x00, 0x01, 0x0b]); /// /// assert_eq!(184614912, reader.read_u32_le().await?); /// Ok(()) /// } /// ``` fn read_u32_le(&mut self) -> ReadU32Le; /// Reads a signed 32-bit integer in little-endian order from the /// underlying reader. /// /// /// Equivalent to: /// /// ```ignore /// async fn read_i32_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 32-bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0xff, 0xff, 0x7a, 0x33]); /// /// assert_eq!(863698943, reader.read_i32_le().await?); /// Ok(()) /// } /// ``` fn read_i32_le(&mut self) -> ReadI32Le; /// Reads an unsigned 64-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u64_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 64-bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 /// ]); /// /// assert_eq!(9477368352180732672, reader.read_u64_le().await?); /// Ok(()) /// } /// ``` fn read_u64_le(&mut self) -> ReadU64Le; /// Reads an signed 64-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i64_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 64-bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0x80, 0, 0, 0, 0, 0, 0, 0]); /// /// assert_eq!(128, reader.read_i64_le().await?); /// Ok(()) /// } /// ``` fn read_i64_le(&mut self) -> ReadI64Le; /// Reads an unsigned 128-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_u128_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read unsigned 128-bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83, /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 /// ]); /// /// assert_eq!(174826588484952389081207917399662330624, reader.read_u128_le().await?); /// Ok(()) /// } /// ``` fn read_u128_le(&mut self) -> ReadU128Le; /// Reads an signed 128-bit integer in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_i128_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read signed 128-bit little-endian integers from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0x80, 0, 0, 0, 0, 0, 0, 0, /// 0, 0, 0, 0, 0, 0, 0, 0 /// ]); /// /// assert_eq!(128, reader.read_i128_le().await?); /// Ok(()) /// } /// ``` fn read_i128_le(&mut self) -> ReadI128Le; /// Reads an 32-bit floating point type in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_f32_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read 32-bit floating point type from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![0xff, 0xff, 0x7f, 0xff]); /// /// assert_eq!(f32::MIN, reader.read_f32_le().await?); /// Ok(()) /// } /// ``` fn read_f32_le(&mut self) -> ReadF32Le; /// Reads an 64-bit floating point type in little-endian order from the /// underlying reader. /// /// Equivalent to: /// /// ```ignore /// async fn read_f64_le(&mut self) -> io::Result; /// ``` /// /// It is recommended to use a buffered reader to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncReadExt::read_exact`]. /// /// [`AsyncReadExt::read_exact`]: AsyncReadExt::read_exact /// /// # Cancel safety /// /// This method is not cancellation safe. If the method is used as the /// event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then some data may be lost. /// /// # Examples /// /// Read 64-bit floating point type from a `AsyncRead`: /// /// ```rust /// use tokio::io::{self, AsyncReadExt}; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut reader = Cursor::new(vec![ /// 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xef, 0xff /// ]); /// /// assert_eq!(f64::MIN, reader.read_f64_le().await?); /// Ok(()) /// } /// ``` fn read_f64_le(&mut self) -> ReadF64Le; } /// Reads all bytes until EOF in this source, placing them into `buf`. /// /// Equivalent to: /// /// ```ignore /// async fn read_to_end(&mut self, buf: &mut Vec) -> io::Result; /// ``` /// /// All bytes read from this source will be appended to the specified /// buffer `buf`. This function will continuously call [`read()`] to /// append more data to `buf` until [`read()`] returns `Ok(0)`. /// /// If successful, the total number of bytes read is returned. /// /// [`read()`]: AsyncReadExt::read /// /// # Errors /// /// If a read error is encountered then the `read_to_end` operation /// immediately completes. Any bytes which have already been read will /// be appended to `buf`. /// /// # Examples /// /// [`File`][crate::fs::File]s implement `Read`: /// /// ```no_run /// use tokio::io::{self, AsyncReadExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; /// let mut buffer = Vec::new(); /// /// // read the whole file /// f.read_to_end(&mut buffer).await?; /// Ok(()) /// } /// ``` /// /// (See also the [`tokio::fs::read`] convenience function for reading from a /// file.) /// /// [`tokio::fs::read`]: fn@crate::fs::read fn read_to_end<'a>(&'a mut self, buf: &'a mut Vec) -> ReadToEnd<'a, Self> where Self: Unpin, { read_to_end(self, buf) } /// Reads all bytes until EOF in this source, appending them to `buf`. /// /// Equivalent to: /// /// ```ignore /// async fn read_to_string(&mut self, buf: &mut String) -> io::Result; /// ``` /// /// If successful, the number of bytes which were read and appended to /// `buf` is returned. /// /// # Errors /// /// If the data in this stream is *not* valid UTF-8 then an error is /// returned and `buf` is unchanged. /// /// See [`read_to_end`][AsyncReadExt::read_to_end] for other error semantics. /// /// # Examples /// /// [`File`][crate::fs::File]s implement `Read`: /// /// ```no_run /// use tokio::io::{self, AsyncReadExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; /// let mut buffer = String::new(); /// /// f.read_to_string(&mut buffer).await?; /// Ok(()) /// } /// ``` /// /// (See also the [`crate::fs::read_to_string`] convenience function for /// reading from a file.) /// /// [`crate::fs::read_to_string`]: fn@crate::fs::read_to_string fn read_to_string<'a>(&'a mut self, dst: &'a mut String) -> ReadToString<'a, Self> where Self: Unpin, { read_to_string(self, dst) } /// Creates an adaptor which reads at most `limit` bytes from it. /// /// This function returns a new instance of `AsyncRead` which will read /// at most `limit` bytes, after which it will always return EOF /// (`Ok(0)`). Any read errors will not count towards the number of /// bytes read and future calls to [`read()`] may succeed. /// /// [`read()`]: fn@crate::io::AsyncReadExt::read /// /// [read]: AsyncReadExt::read /// /// # Examples /// /// [`File`][crate::fs::File]s implement `Read`: /// /// ```no_run /// use tokio::io::{self, AsyncReadExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let f = File::open("foo.txt").await?; /// let mut buffer = [0; 5]; /// /// // read at most five bytes /// let mut handle = f.take(5); /// /// handle.read(&mut buffer).await?; /// Ok(()) /// } /// ``` fn take(self, limit: u64) -> Take where Self: Sized, { take(self, limit) } } } impl AsyncReadExt for R {} tokio-1.43.1/src/io/util/async_seek_ext.rs000064400000000000000000000053501046102023000165670ustar 00000000000000use crate::io::seek::{seek, Seek}; use crate::io::AsyncSeek; use std::io::SeekFrom; cfg_io_util! { /// An extension trait that adds utility methods to [`AsyncSeek`] types. /// /// # Examples /// /// ``` /// use std::io::{self, Cursor, SeekFrom}; /// use tokio::io::{AsyncSeekExt, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut cursor = Cursor::new(b"abcdefg"); /// /// // the `seek` method is defined by this trait /// cursor.seek(SeekFrom::Start(3)).await?; /// /// let mut buf = [0; 1]; /// let n = cursor.read(&mut buf).await?; /// assert_eq!(n, 1); /// assert_eq!(buf, [b'd']); /// /// Ok(()) /// } /// ``` /// /// See [module][crate::io] documentation for more details. /// /// [`AsyncSeek`]: AsyncSeek pub trait AsyncSeekExt: AsyncSeek { /// Creates a future which will seek an IO object, and then yield the /// new position in the object and the object itself. /// /// Equivalent to: /// /// ```ignore /// async fn seek(&mut self, pos: SeekFrom) -> io::Result; /// ``` /// /// In the case of an error the buffer and the object will be discarded, with /// the error yielded. /// /// # Examples /// /// ```no_run /// use tokio::fs::File; /// use tokio::io::{AsyncSeekExt, AsyncReadExt}; /// /// use std::io::SeekFrom; /// /// # async fn dox() -> std::io::Result<()> { /// let mut file = File::open("foo.txt").await?; /// file.seek(SeekFrom::Start(6)).await?; /// /// let mut contents = vec![0u8; 10]; /// file.read_exact(&mut contents).await?; /// # Ok(()) /// # } /// ``` fn seek(&mut self, pos: SeekFrom) -> Seek<'_, Self> where Self: Unpin, { seek(self, pos) } /// Creates a future which will rewind to the beginning of the stream. /// /// This is convenience method, equivalent to `self.seek(SeekFrom::Start(0))`. fn rewind(&mut self) -> Seek<'_, Self> where Self: Unpin, { self.seek(SeekFrom::Start(0)) } /// Creates a future which will return the current seek position from the /// start of the stream. /// /// This is equivalent to `self.seek(SeekFrom::Current(0))`. fn stream_position(&mut self) -> Seek<'_, Self> where Self: Unpin, { self.seek(SeekFrom::Current(0)) } } } impl AsyncSeekExt for S {} tokio-1.43.1/src/io/util/async_write_ext.rs000064400000000000000000001317771046102023000170070ustar 00000000000000use crate::io::util::flush::{flush, Flush}; use crate::io::util::shutdown::{shutdown, Shutdown}; use crate::io::util::write::{write, Write}; use crate::io::util::write_all::{write_all, WriteAll}; use crate::io::util::write_all_buf::{write_all_buf, WriteAllBuf}; use crate::io::util::write_buf::{write_buf, WriteBuf}; use crate::io::util::write_int::{WriteF32, WriteF32Le, WriteF64, WriteF64Le}; use crate::io::util::write_int::{ WriteI128, WriteI128Le, WriteI16, WriteI16Le, WriteI32, WriteI32Le, WriteI64, WriteI64Le, WriteI8, }; use crate::io::util::write_int::{ WriteU128, WriteU128Le, WriteU16, WriteU16Le, WriteU32, WriteU32Le, WriteU64, WriteU64Le, WriteU8, }; use crate::io::util::write_vectored::{write_vectored, WriteVectored}; use crate::io::AsyncWrite; use std::io::IoSlice; use bytes::Buf; cfg_io_util! { /// Defines numeric writer. macro_rules! write_impl { ( $( $(#[$outer:meta])* fn $name:ident(&mut self, n: $ty:ty) -> $($fut:ident)*; )* ) => { $( $(#[$outer])* fn $name(&mut self, n: $ty) -> $($fut)*<&mut Self> where Self: Unpin { $($fut)*::new(self, n) } )* } } /// Writes bytes to a sink. /// /// Implemented as an extension trait, adding utility methods to all /// [`AsyncWrite`] types. Callers will tend to import this trait instead of /// [`AsyncWrite`]. /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let data = b"some bytes"; /// /// let mut pos = 0; /// let mut buffer = File::create("foo.txt").await?; /// /// while pos < data.len() { /// let bytes_written = buffer.write(&data[pos..]).await?; /// pos += bytes_written; /// } /// /// Ok(()) /// } /// ``` /// /// See [module][crate::io] documentation for more details. /// /// [`AsyncWrite`]: AsyncWrite pub trait AsyncWriteExt: AsyncWrite { /// Writes a buffer into this writer, returning how many bytes were /// written. /// /// Equivalent to: /// /// ```ignore /// async fn write(&mut self, buf: &[u8]) -> io::Result; /// ``` /// /// This function will attempt to write the entire contents of `buf`, but /// the entire write may not succeed, or the write may also generate an /// error. A call to `write` represents *at most one* attempt to write to /// any wrapped object. /// /// # Return /// /// If the return value is `Ok(n)` then it must be guaranteed that `n <= /// buf.len()`. A return value of `0` typically means that the /// underlying object is no longer able to accept bytes and will likely /// not be able to in the future as well, or that the buffer provided is /// empty. /// /// # Errors /// /// Each call to `write` may generate an I/O error indicating that the /// operation could not be completed. If an error is returned then no bytes /// in the buffer were written to this writer. /// /// It is **not** considered an error if the entire buffer could not be /// written to this writer. /// /// # Cancel safety /// /// This method is cancellation safe in the sense that if it is used as /// the event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then it is guaranteed that no data was /// written to this `AsyncWrite`. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// /// // Writes some prefix of the byte string, not necessarily all of it. /// file.write(b"some bytes").await?; /// file.flush().await?; /// Ok(()) /// } /// ``` fn write<'a>(&'a mut self, src: &'a [u8]) -> Write<'a, Self> where Self: Unpin, { write(self, src) } /// Like [`write`], except that it writes from a slice of buffers. /// /// Equivalent to: /// /// ```ignore /// async fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result; /// ``` /// /// See [`AsyncWrite::poll_write_vectored`] for more details. /// /// # Cancel safety /// /// This method is cancellation safe in the sense that if it is used as /// the event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then it is guaranteed that no data was /// written to this `AsyncWrite`. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; /// use tokio::fs::File; /// use std::io::IoSlice; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// /// let bufs: &[_] = &[ /// IoSlice::new(b"hello"), /// IoSlice::new(b" "), /// IoSlice::new(b"world"), /// ]; /// /// file.write_vectored(&bufs).await?; /// file.flush().await?; /// /// Ok(()) /// } /// ``` /// /// [`write`]: AsyncWriteExt::write fn write_vectored<'a, 'b>(&'a mut self, bufs: &'a [IoSlice<'b>]) -> WriteVectored<'a, 'b, Self> where Self: Unpin, { write_vectored(self, bufs) } /// Writes a buffer into this writer, advancing the buffer's internal /// cursor. /// /// Equivalent to: /// /// ```ignore /// async fn write_buf(&mut self, buf: &mut B) -> io::Result; /// ``` /// /// This function will attempt to write the entire contents of `buf`, but /// the entire write may not succeed, or the write may also generate an /// error. After the operation completes, the buffer's /// internal cursor is advanced by the number of bytes written. A /// subsequent call to `write_buf` using the **same** `buf` value will /// resume from the point that the first call to `write_buf` completed. /// A call to `write_buf` represents *at most one* attempt to write to any /// wrapped object. /// /// # Return /// /// If the return value is `Ok(n)` then it must be guaranteed that `n <= /// buf.len()`. A return value of `0` typically means that the /// underlying object is no longer able to accept bytes and will likely /// not be able to in the future as well, or that the buffer provided is /// empty. /// /// # Errors /// /// Each call to `write` may generate an I/O error indicating that the /// operation could not be completed. If an error is returned then no bytes /// in the buffer were written to this writer. /// /// It is **not** considered an error if the entire buffer could not be /// written to this writer. /// /// # Cancel safety /// /// This method is cancellation safe in the sense that if it is used as /// the event in a [`tokio::select!`](crate::select) statement and some /// other branch completes first, then it is guaranteed that no data was /// written to this `AsyncWrite`. /// /// # Examples /// /// [`File`] implements [`AsyncWrite`] and [`Cursor`]`<&[u8]>` implements [`Buf`]: /// /// [`File`]: crate::fs::File /// [`Buf`]: bytes::Buf /// [`Cursor`]: std::io::Cursor /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; /// use tokio::fs::File; /// /// use bytes::Buf; /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// let mut buffer = Cursor::new(b"data to write"); /// /// // Loop until the entire contents of the buffer are written to /// // the file. /// while buffer.has_remaining() { /// // Writes some prefix of the byte string, not necessarily /// // all of it. /// file.write_buf(&mut buffer).await?; /// } /// file.flush().await?; /// /// Ok(()) /// } /// ``` fn write_buf<'a, B>(&'a mut self, src: &'a mut B) -> WriteBuf<'a, Self, B> where Self: Sized + Unpin, B: Buf, { write_buf(self, src) } /// Attempts to write an entire buffer into this writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_all_buf(&mut self, buf: impl Buf) -> Result<(), io::Error> { /// while buf.has_remaining() { /// self.write_buf(&mut buf).await?; /// } /// Ok(()) /// } /// ``` /// /// This method will continuously call [`write`] until /// [`buf.has_remaining()`](bytes::Buf::has_remaining) returns false. This method will not /// return until the entire buffer has been successfully written or an error occurs. The /// first error generated will be returned. /// /// The buffer is advanced after each chunk is successfully written. After failure, /// `src.chunk()` will return the chunk that failed to write. /// /// # Cancel safety /// /// If `write_all_buf` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then the data in the provided buffer may have been /// partially written. However, it is guaranteed that the provided /// buffer has been [advanced] by the amount of bytes that have been /// partially written. /// /// # Examples /// /// [`File`] implements [`AsyncWrite`] and [`Cursor`]`<&[u8]>` implements [`Buf`]: /// /// [`File`]: crate::fs::File /// [`Buf`]: bytes::Buf /// [`Cursor`]: std::io::Cursor /// [advanced]: bytes::Buf::advance /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; /// use tokio::fs::File; /// /// use std::io::Cursor; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// let mut buffer = Cursor::new(b"data to write"); /// /// file.write_all_buf(&mut buffer).await?; /// file.flush().await?; /// Ok(()) /// } /// ``` /// /// [`write`]: AsyncWriteExt::write fn write_all_buf<'a, B>(&'a mut self, src: &'a mut B) -> WriteAllBuf<'a, Self, B> where Self: Sized + Unpin, B: Buf, { write_all_buf(self, src) } /// Attempts to write an entire buffer into this writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_all(&mut self, buf: &[u8]) -> io::Result<()>; /// ``` /// /// This method will continuously call [`write`] until there is no more data /// to be written. This method will not return until the entire buffer /// has been successfully written or such an error occurs. The first /// error generated from this method will be returned. /// /// # Cancel safety /// /// This method is not cancellation safe. If it is used as the event /// in a [`tokio::select!`](crate::select) statement and some other /// branch completes first, then the provided buffer may have been /// partially written, but future calls to `write_all` will start over /// from the beginning of the buffer. /// /// # Errors /// /// This function will return the first error that [`write`] returns. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut file = File::create("foo.txt").await?; /// /// file.write_all(b"some bytes").await?; /// file.flush().await?; /// Ok(()) /// } /// ``` /// /// [`write`]: AsyncWriteExt::write fn write_all<'a>(&'a mut self, src: &'a [u8]) -> WriteAll<'a, Self> where Self: Unpin, { write_all(self, src) } write_impl! { /// Writes an unsigned 8-bit integer to the underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u8(&mut self, n: u8) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 8 bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u8(2).await?; /// writer.write_u8(5).await?; /// /// assert_eq!(writer, b"\x02\x05"); /// Ok(()) /// } /// ``` fn write_u8(&mut self, n: u8) -> WriteU8; /// Writes a signed 8-bit integer to the underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i8(&mut self, n: i8) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 8 bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i8(-2).await?; /// writer.write_i8(126).await?; /// /// assert_eq!(writer, b"\xFE\x7E"); /// Ok(()) /// } /// ``` fn write_i8(&mut self, n: i8) -> WriteI8; /// Writes an unsigned 16-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u16(&mut self, n: u16) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 16-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u16(517).await?; /// writer.write_u16(768).await?; /// /// assert_eq!(writer, b"\x02\x05\x03\x00"); /// Ok(()) /// } /// ``` fn write_u16(&mut self, n: u16) -> WriteU16; /// Writes a signed 16-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i16(&mut self, n: i16) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 16-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i16(193).await?; /// writer.write_i16(-132).await?; /// /// assert_eq!(writer, b"\x00\xc1\xff\x7c"); /// Ok(()) /// } /// ``` fn write_i16(&mut self, n: i16) -> WriteI16; /// Writes an unsigned 32-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u32(&mut self, n: u32) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 32-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u32(267).await?; /// writer.write_u32(1205419366).await?; /// /// assert_eq!(writer, b"\x00\x00\x01\x0b\x47\xd9\x3d\x66"); /// Ok(()) /// } /// ``` fn write_u32(&mut self, n: u32) -> WriteU32; /// Writes a signed 32-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i32(&mut self, n: i32) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 32-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i32(267).await?; /// writer.write_i32(1205419366).await?; /// /// assert_eq!(writer, b"\x00\x00\x01\x0b\x47\xd9\x3d\x66"); /// Ok(()) /// } /// ``` fn write_i32(&mut self, n: i32) -> WriteI32; /// Writes an unsigned 64-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u64(&mut self, n: u64) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 64-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u64(918733457491587).await?; /// writer.write_u64(143).await?; /// /// assert_eq!(writer, b"\x00\x03\x43\x95\x4d\x60\x86\x83\x00\x00\x00\x00\x00\x00\x00\x8f"); /// Ok(()) /// } /// ``` fn write_u64(&mut self, n: u64) -> WriteU64; /// Writes an signed 64-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i64(&mut self, n: i64) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 64-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i64(i64::MIN).await?; /// writer.write_i64(i64::MAX).await?; /// /// assert_eq!(writer, b"\x80\x00\x00\x00\x00\x00\x00\x00\x7f\xff\xff\xff\xff\xff\xff\xff"); /// Ok(()) /// } /// ``` fn write_i64(&mut self, n: i64) -> WriteI64; /// Writes an unsigned 128-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u128(&mut self, n: u128) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 128-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u128(16947640962301618749969007319746179).await?; /// /// assert_eq!(writer, vec![ /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83, /// 0x00, 0x03, 0x43, 0x95, 0x4d, 0x60, 0x86, 0x83 /// ]); /// Ok(()) /// } /// ``` fn write_u128(&mut self, n: u128) -> WriteU128; /// Writes an signed 128-bit integer in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i128(&mut self, n: i128) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 128-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i128(i128::MIN).await?; /// /// assert_eq!(writer, vec![ /// 0x80, 0, 0, 0, 0, 0, 0, 0, /// 0, 0, 0, 0, 0, 0, 0, 0 /// ]); /// Ok(()) /// } /// ``` fn write_i128(&mut self, n: i128) -> WriteI128; /// Writes an 32-bit floating point type in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_f32(&mut self, n: f32) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write 32-bit floating point type to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_f32(f32::MIN).await?; /// /// assert_eq!(writer, vec![0xff, 0x7f, 0xff, 0xff]); /// Ok(()) /// } /// ``` fn write_f32(&mut self, n: f32) -> WriteF32; /// Writes an 64-bit floating point type in big-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_f64(&mut self, n: f64) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write 64-bit floating point type to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_f64(f64::MIN).await?; /// /// assert_eq!(writer, vec![ /// 0xff, 0xef, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff /// ]); /// Ok(()) /// } /// ``` fn write_f64(&mut self, n: f64) -> WriteF64; /// Writes an unsigned 16-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u16_le(&mut self, n: u16) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 16-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u16_le(517).await?; /// writer.write_u16_le(768).await?; /// /// assert_eq!(writer, b"\x05\x02\x00\x03"); /// Ok(()) /// } /// ``` fn write_u16_le(&mut self, n: u16) -> WriteU16Le; /// Writes a signed 16-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i16_le(&mut self, n: i16) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 16-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i16_le(193).await?; /// writer.write_i16_le(-132).await?; /// /// assert_eq!(writer, b"\xc1\x00\x7c\xff"); /// Ok(()) /// } /// ``` fn write_i16_le(&mut self, n: i16) -> WriteI16Le; /// Writes an unsigned 32-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u32_le(&mut self, n: u32) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 32-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u32_le(267).await?; /// writer.write_u32_le(1205419366).await?; /// /// assert_eq!(writer, b"\x0b\x01\x00\x00\x66\x3d\xd9\x47"); /// Ok(()) /// } /// ``` fn write_u32_le(&mut self, n: u32) -> WriteU32Le; /// Writes a signed 32-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i32_le(&mut self, n: i32) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 32-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i32_le(267).await?; /// writer.write_i32_le(1205419366).await?; /// /// assert_eq!(writer, b"\x0b\x01\x00\x00\x66\x3d\xd9\x47"); /// Ok(()) /// } /// ``` fn write_i32_le(&mut self, n: i32) -> WriteI32Le; /// Writes an unsigned 64-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u64_le(&mut self, n: u64) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 64-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u64_le(918733457491587).await?; /// writer.write_u64_le(143).await?; /// /// assert_eq!(writer, b"\x83\x86\x60\x4d\x95\x43\x03\x00\x8f\x00\x00\x00\x00\x00\x00\x00"); /// Ok(()) /// } /// ``` fn write_u64_le(&mut self, n: u64) -> WriteU64Le; /// Writes an signed 64-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i64_le(&mut self, n: i64) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 64-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i64_le(i64::MIN).await?; /// writer.write_i64_le(i64::MAX).await?; /// /// assert_eq!(writer, b"\x00\x00\x00\x00\x00\x00\x00\x80\xff\xff\xff\xff\xff\xff\xff\x7f"); /// Ok(()) /// } /// ``` fn write_i64_le(&mut self, n: i64) -> WriteI64Le; /// Writes an unsigned 128-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_u128_le(&mut self, n: u128) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write unsigned 128-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_u128_le(16947640962301618749969007319746179).await?; /// /// assert_eq!(writer, vec![ /// 0x83, 0x86, 0x60, 0x4d, 0x95, 0x43, 0x03, 0x00, /// 0x83, 0x86, 0x60, 0x4d, 0x95, 0x43, 0x03, 0x00, /// ]); /// Ok(()) /// } /// ``` fn write_u128_le(&mut self, n: u128) -> WriteU128Le; /// Writes an signed 128-bit integer in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_i128_le(&mut self, n: i128) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write signed 128-bit integers to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_i128_le(i128::MIN).await?; /// /// assert_eq!(writer, vec![ /// 0, 0, 0, 0, 0, 0, 0, /// 0, 0, 0, 0, 0, 0, 0, 0, 0x80 /// ]); /// Ok(()) /// } /// ``` fn write_i128_le(&mut self, n: i128) -> WriteI128Le; /// Writes an 32-bit floating point type in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_f32_le(&mut self, n: f32) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write 32-bit floating point type to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_f32_le(f32::MIN).await?; /// /// assert_eq!(writer, vec![0xff, 0xff, 0x7f, 0xff]); /// Ok(()) /// } /// ``` fn write_f32_le(&mut self, n: f32) -> WriteF32Le; /// Writes an 64-bit floating point type in little-endian order to the /// underlying writer. /// /// Equivalent to: /// /// ```ignore /// async fn write_f64_le(&mut self, n: f64) -> io::Result<()>; /// ``` /// /// It is recommended to use a buffered writer to avoid excessive /// syscalls. /// /// # Errors /// /// This method returns the same errors as [`AsyncWriteExt::write_all`]. /// /// [`AsyncWriteExt::write_all`]: AsyncWriteExt::write_all /// /// # Examples /// /// Write 64-bit floating point type to a `AsyncWrite`: /// /// ```rust /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut writer = Vec::new(); /// /// writer.write_f64_le(f64::MIN).await?; /// /// assert_eq!(writer, vec![ /// 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xef, 0xff /// ]); /// Ok(()) /// } /// ``` fn write_f64_le(&mut self, n: f64) -> WriteF64Le; } /// Flushes this output stream, ensuring that all intermediately buffered /// contents reach their destination. /// /// Equivalent to: /// /// ```ignore /// async fn flush(&mut self) -> io::Result<()>; /// ``` /// /// # Errors /// /// It is considered an error if not all bytes could be written due to /// I/O errors or EOF being reached. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, BufWriter, AsyncWriteExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let f = File::create("foo.txt").await?; /// let mut buffer = BufWriter::new(f); /// /// buffer.write_all(b"some bytes").await?; /// buffer.flush().await?; /// Ok(()) /// } /// ``` fn flush(&mut self) -> Flush<'_, Self> where Self: Unpin, { flush(self) } /// Shuts down the output stream, ensuring that the value can be dropped /// cleanly. /// /// Equivalent to: /// /// ```ignore /// async fn shutdown(&mut self) -> io::Result<()>; /// ``` /// /// Similar to [`flush`], all intermediately buffered is written to the /// underlying stream. Once the operation completes, the caller should /// no longer attempt to write to the stream. For example, the /// `TcpStream` implementation will issue a `shutdown(Write)` sys call. /// /// [`flush`]: fn@crate::io::AsyncWriteExt::flush /// /// # Examples /// /// ```no_run /// use tokio::io::{self, BufWriter, AsyncWriteExt}; /// use tokio::fs::File; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let f = File::create("foo.txt").await?; /// let mut buffer = BufWriter::new(f); /// /// buffer.write_all(b"some bytes").await?; /// buffer.shutdown().await?; /// Ok(()) /// } /// ``` fn shutdown(&mut self) -> Shutdown<'_, Self> where Self: Unpin, { shutdown(self) } } } impl AsyncWriteExt for W {} tokio-1.43.1/src/io/util/buf_reader.rs000064400000000000000000000266011046102023000156630ustar 00000000000000use crate::io::util::DEFAULT_BUF_SIZE; use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use pin_project_lite::pin_project; use std::io::{self, IoSlice, SeekFrom}; use std::pin::Pin; use std::task::{ready, Context, Poll}; use std::{cmp, fmt, mem}; pin_project! { /// The `BufReader` struct adds buffering to any reader. /// /// It can be excessively inefficient to work directly with a [`AsyncRead`] /// instance. A `BufReader` performs large, infrequent reads on the underlying /// [`AsyncRead`] and maintains an in-memory buffer of the results. /// /// `BufReader` can improve the speed of programs that make *small* and /// *repeated* read calls to the same file or network socket. It does not /// help when reading very large amounts at once, or reading just one or a few /// times. It also provides no advantage when reading from a source that is /// already in memory, like a `Vec`. /// /// When the `BufReader` is dropped, the contents of its buffer will be /// discarded. Creating multiple instances of a `BufReader` on the same /// stream can cause data loss. #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct BufReader { #[pin] pub(super) inner: R, pub(super) buf: Box<[u8]>, pub(super) pos: usize, pub(super) cap: usize, pub(super) seek_state: SeekState, } } impl BufReader { /// Creates a new `BufReader` with a default buffer capacity. The default is currently 8 KB, /// but may change in the future. pub fn new(inner: R) -> Self { Self::with_capacity(DEFAULT_BUF_SIZE, inner) } /// Creates a new `BufReader` with the specified buffer capacity. pub fn with_capacity(capacity: usize, inner: R) -> Self { let buffer = vec![0; capacity]; Self { inner, buf: buffer.into_boxed_slice(), pos: 0, cap: 0, seek_state: SeekState::Init, } } /// Gets a reference to the underlying reader. /// /// It is inadvisable to directly read from the underlying reader. pub fn get_ref(&self) -> &R { &self.inner } /// Gets a mutable reference to the underlying reader. /// /// It is inadvisable to directly read from the underlying reader. pub fn get_mut(&mut self) -> &mut R { &mut self.inner } /// Gets a pinned mutable reference to the underlying reader. /// /// It is inadvisable to directly read from the underlying reader. pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { self.project().inner } /// Consumes this `BufReader`, returning the underlying reader. /// /// Note that any leftover data in the internal buffer is lost. pub fn into_inner(self) -> R { self.inner } /// Returns a reference to the internally buffered data. /// /// Unlike `fill_buf`, this will not attempt to fill the buffer if it is empty. pub fn buffer(&self) -> &[u8] { &self.buf[self.pos..self.cap] } /// Invalidates all data in the internal buffer. #[inline] fn discard_buffer(self: Pin<&mut Self>) { let me = self.project(); *me.pos = 0; *me.cap = 0; } } impl AsyncRead for BufReader { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { // If we don't have any buffered data and we're doing a massive read // (larger than our internal buffer), bypass our internal buffer // entirely. if self.pos == self.cap && buf.remaining() >= self.buf.len() { let res = ready!(self.as_mut().get_pin_mut().poll_read(cx, buf)); self.discard_buffer(); return Poll::Ready(res); } let rem = ready!(self.as_mut().poll_fill_buf(cx))?; let amt = std::cmp::min(rem.len(), buf.remaining()); buf.put_slice(&rem[..amt]); self.consume(amt); Poll::Ready(Ok(())) } } impl AsyncBufRead for BufReader { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); // If we've reached the end of our internal buffer then we need to fetch // some more data from the underlying reader. // Branch using `>=` instead of the more correct `==` // to tell the compiler that the pos..cap slice is always valid. if *me.pos >= *me.cap { debug_assert!(*me.pos == *me.cap); let mut buf = ReadBuf::new(me.buf); ready!(me.inner.poll_read(cx, &mut buf))?; *me.cap = buf.filled().len(); *me.pos = 0; } Poll::Ready(Ok(&me.buf[*me.pos..*me.cap])) } fn consume(self: Pin<&mut Self>, amt: usize) { let me = self.project(); *me.pos = cmp::min(*me.pos + amt, *me.cap); } } #[derive(Debug, Clone, Copy)] pub(super) enum SeekState { /// `start_seek` has not been called. Init, /// `start_seek` has been called, but `poll_complete` has not yet been called. Start(SeekFrom), /// Waiting for completion of the first `poll_complete` in the `n.checked_sub(remainder).is_none()` branch. PendingOverflowed(i64), /// Waiting for completion of `poll_complete`. Pending, } /// Seeks to an offset, in bytes, in the underlying reader. /// /// The position used for seeking with `SeekFrom::Current(_)` is the /// position the underlying reader would be at if the `BufReader` had no /// internal buffer. /// /// Seeking always discards the internal buffer, even if the seek position /// would otherwise fall within it. This guarantees that calling /// `.into_inner()` immediately after a seek yields the underlying reader /// at the same position. /// /// See [`AsyncSeek`] for more details. /// /// Note: In the edge case where you're seeking with `SeekFrom::Current(n)` /// where `n` minus the internal buffer length overflows an `i64`, two /// seeks will be performed instead of one. If the second seek returns /// `Err`, the underlying reader will be left at the same position it would /// have if you called `seek` with `SeekFrom::Current(0)`. impl AsyncSeek for BufReader { fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { // We needs to call seek operation multiple times. // And we should always call both start_seek and poll_complete, // as start_seek alone cannot guarantee that the operation will be completed. // poll_complete receives a Context and returns a Poll, so it cannot be called // inside start_seek. *self.project().seek_state = SeekState::Start(pos); Ok(()) } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let res = match mem::replace(self.as_mut().project().seek_state, SeekState::Init) { SeekState::Init => { // 1.x AsyncSeek recommends calling poll_complete before start_seek. // We don't have to guarantee that the value returned by // poll_complete called without start_seek is correct, // so we'll return 0. return Poll::Ready(Ok(0)); } SeekState::Start(SeekFrom::Current(n)) => { let remainder = (self.cap - self.pos) as i64; // it should be safe to assume that remainder fits within an i64 as the alternative // means we managed to allocate 8 exbibytes and that's absurd. // But it's not out of the realm of possibility for some weird underlying reader to // support seeking by i64::MIN so we need to handle underflow when subtracting // remainder. if let Some(offset) = n.checked_sub(remainder) { self.as_mut() .get_pin_mut() .start_seek(SeekFrom::Current(offset))?; } else { // seek backwards by our remainder, and then by the offset self.as_mut() .get_pin_mut() .start_seek(SeekFrom::Current(-remainder))?; if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() { *self.as_mut().project().seek_state = SeekState::PendingOverflowed(n); return Poll::Pending; } // https://github.com/rust-lang/rust/pull/61157#issuecomment-495932676 self.as_mut().discard_buffer(); self.as_mut() .get_pin_mut() .start_seek(SeekFrom::Current(n))?; } self.as_mut().get_pin_mut().poll_complete(cx)? } SeekState::PendingOverflowed(n) => { if self.as_mut().get_pin_mut().poll_complete(cx)?.is_pending() { *self.as_mut().project().seek_state = SeekState::PendingOverflowed(n); return Poll::Pending; } // https://github.com/rust-lang/rust/pull/61157#issuecomment-495932676 self.as_mut().discard_buffer(); self.as_mut() .get_pin_mut() .start_seek(SeekFrom::Current(n))?; self.as_mut().get_pin_mut().poll_complete(cx)? } SeekState::Start(pos) => { // Seeking with Start/End doesn't care about our buffer length. self.as_mut().get_pin_mut().start_seek(pos)?; self.as_mut().get_pin_mut().poll_complete(cx)? } SeekState::Pending => self.as_mut().get_pin_mut().poll_complete(cx)?, }; match res { Poll::Ready(res) => { self.discard_buffer(); Poll::Ready(Ok(res)) } Poll::Pending => { *self.as_mut().project().seek_state = SeekState::Pending; Poll::Pending } } } } impl AsyncWrite for BufReader { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.get_pin_mut().poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { self.get_pin_mut().poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { self.get_ref().is_write_vectored() } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_pin_mut().poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_pin_mut().poll_shutdown(cx) } } impl fmt::Debug for BufReader { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BufReader") .field("reader", &self.inner) .field( "buffer", &format_args!("{}/{}", self.cap - self.pos, self.buf.len()), ) .finish() } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/buf_stream.rs000064400000000000000000000150611046102023000157120ustar 00000000000000use crate::io::util::{BufReader, BufWriter}; use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use pin_project_lite::pin_project; use std::io::{self, IoSlice, SeekFrom}; use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { /// Wraps a type that is [`AsyncWrite`] and [`AsyncRead`], and buffers its input and output. /// /// It can be excessively inefficient to work directly with something that implements [`AsyncWrite`] /// and [`AsyncRead`]. For example, every `write`, however small, has to traverse the syscall /// interface, and similarly, every read has to do the same. The [`BufWriter`] and [`BufReader`] /// types aid with these problems respectively, but do so in only one direction. `BufStream` wraps /// one in the other so that both directions are buffered. See their documentation for details. #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct BufStream { #[pin] inner: BufReader>, } } impl BufStream { /// Wraps a type in both [`BufWriter`] and [`BufReader`]. /// /// See the documentation for those types and [`BufStream`] for details. pub fn new(stream: RW) -> BufStream { BufStream { inner: BufReader::new(BufWriter::new(stream)), } } /// Creates a `BufStream` with the specified [`BufReader`] capacity and [`BufWriter`] /// capacity. /// /// See the documentation for those types and [`BufStream`] for details. pub fn with_capacity( reader_capacity: usize, writer_capacity: usize, stream: RW, ) -> BufStream { BufStream { inner: BufReader::with_capacity( reader_capacity, BufWriter::with_capacity(writer_capacity, stream), ), } } /// Gets a reference to the underlying I/O object. /// /// It is inadvisable to directly read from the underlying I/O object. pub fn get_ref(&self) -> &RW { self.inner.get_ref().get_ref() } /// Gets a mutable reference to the underlying I/O object. /// /// It is inadvisable to directly read from the underlying I/O object. pub fn get_mut(&mut self) -> &mut RW { self.inner.get_mut().get_mut() } /// Gets a pinned mutable reference to the underlying I/O object. /// /// It is inadvisable to directly read from the underlying I/O object. pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut RW> { self.project().inner.get_pin_mut().get_pin_mut() } /// Consumes this `BufStream`, returning the underlying I/O object. /// /// Note that any leftover data in the internal buffer is lost. pub fn into_inner(self) -> RW { self.inner.into_inner().into_inner() } } impl From>> for BufStream { fn from(b: BufReader>) -> Self { BufStream { inner: b } } } impl From>> for BufStream { fn from(b: BufWriter>) -> Self { // we need to "invert" the reader and writer let BufWriter { inner: BufReader { inner, buf: rbuf, pos, cap, seek_state: rseek_state, }, buf: wbuf, written, seek_state: wseek_state, } = b; BufStream { inner: BufReader { inner: BufWriter { inner, buf: wbuf, written, seek_state: wseek_state, }, buf: rbuf, pos, cap, seek_state: rseek_state, }, } } } impl AsyncWrite for BufStream { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.project().inner.poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { self.project().inner.poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { self.inner.is_write_vectored() } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.poll_shutdown(cx) } } impl AsyncRead for BufStream { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.project().inner.poll_read(cx, buf) } } /// Seek to an offset, in bytes, in the underlying stream. /// /// The position used for seeking with `SeekFrom::Current(_)` is the /// position the underlying stream would be at if the `BufStream` had no /// internal buffer. /// /// Seeking always discards the internal buffer, even if the seek position /// would otherwise fall within it. This guarantees that calling /// `.into_inner()` immediately after a seek yields the underlying reader /// at the same position. /// /// See [`AsyncSeek`] for more details. /// /// Note: In the edge case where you're seeking with `SeekFrom::Current(n)` /// where `n` minus the internal buffer length overflows an `i64`, two /// seeks will be performed instead of one. If the second seek returns /// `Err`, the underlying reader will be left at the same position it would /// have if you called `seek` with `SeekFrom::Current(0)`. impl AsyncSeek for BufStream { fn start_seek(self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> { self.project().inner.start_seek(position) } fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.poll_complete(cx) } } impl AsyncBufRead for BufStream { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.poll_fill_buf(cx) } fn consume(self: Pin<&mut Self>, amt: usize) { self.project().inner.consume(amt); } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/buf_writer.rs000064400000000000000000000244041046102023000157340ustar 00000000000000use crate::io::util::DEFAULT_BUF_SIZE; use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use pin_project_lite::pin_project; use std::fmt; use std::io::{self, IoSlice, SeekFrom, Write}; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// Wraps a writer and buffers its output. /// /// It can be excessively inefficient to work directly with something that /// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and /// writes it to an underlying writer in large, infrequent batches. /// /// `BufWriter` can improve the speed of programs that make *small* and /// *repeated* write calls to the same file or network socket. It does not /// help when writing very large amounts at once, or writing just one or a few /// times. It also provides no advantage when writing to a destination that is /// in memory, like a `Vec`. /// /// When the `BufWriter` is dropped, the contents of its buffer will be /// discarded. Creating multiple instances of a `BufWriter` on the same /// stream can cause data loss. If you need to write out the contents of its /// buffer, you must manually call flush before the writer is dropped. /// /// [`AsyncWrite`]: AsyncWrite /// [`flush`]: super::AsyncWriteExt::flush /// #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct BufWriter { #[pin] pub(super) inner: W, pub(super) buf: Vec, pub(super) written: usize, pub(super) seek_state: SeekState, } } impl BufWriter { /// Creates a new `BufWriter` with a default buffer capacity. The default is currently 8 KB, /// but may change in the future. pub fn new(inner: W) -> Self { Self::with_capacity(DEFAULT_BUF_SIZE, inner) } /// Creates a new `BufWriter` with the specified buffer capacity. pub fn with_capacity(cap: usize, inner: W) -> Self { Self { inner, buf: Vec::with_capacity(cap), written: 0, seek_state: SeekState::Init, } } fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut me = self.project(); let len = me.buf.len(); let mut ret = Ok(()); while *me.written < len { match ready!(me.inner.as_mut().poll_write(cx, &me.buf[*me.written..])) { Ok(0) => { ret = Err(io::Error::new( io::ErrorKind::WriteZero, "failed to write the buffered data", )); break; } Ok(n) => *me.written += n, Err(e) => { ret = Err(e); break; } } } if *me.written > 0 { me.buf.drain(..*me.written); } *me.written = 0; Poll::Ready(ret) } /// Gets a reference to the underlying writer. pub fn get_ref(&self) -> &W { &self.inner } /// Gets a mutable reference to the underlying writer. /// /// It is inadvisable to directly write to the underlying writer. pub fn get_mut(&mut self) -> &mut W { &mut self.inner } /// Gets a pinned mutable reference to the underlying writer. /// /// It is inadvisable to directly write to the underlying writer. pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> { self.project().inner } /// Consumes this `BufWriter`, returning the underlying writer. /// /// Note that any leftover data in the internal buffer is lost. pub fn into_inner(self) -> W { self.inner } /// Returns a reference to the internally buffered data. pub fn buffer(&self) -> &[u8] { &self.buf } } impl AsyncWrite for BufWriter { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { if self.buf.len() + buf.len() > self.buf.capacity() { ready!(self.as_mut().flush_buf(cx))?; } let me = self.project(); if buf.len() >= me.buf.capacity() { me.inner.poll_write(cx, buf) } else { Poll::Ready(me.buf.write(buf)) } } fn poll_write_vectored( mut self: Pin<&mut Self>, cx: &mut Context<'_>, mut bufs: &[IoSlice<'_>], ) -> Poll> { if self.inner.is_write_vectored() { let total_len = bufs .iter() .fold(0usize, |acc, b| acc.saturating_add(b.len())); if total_len > self.buf.capacity() - self.buf.len() { ready!(self.as_mut().flush_buf(cx))?; } let me = self.as_mut().project(); if total_len >= me.buf.capacity() { // It's more efficient to pass the slices directly to the // underlying writer than to buffer them. // The case when the total_len calculation saturates at // usize::MAX is also handled here. me.inner.poll_write_vectored(cx, bufs) } else { bufs.iter().for_each(|b| me.buf.extend_from_slice(b)); Poll::Ready(Ok(total_len)) } } else { // Remove empty buffers at the beginning of bufs. while bufs.first().map(|buf| buf.len()) == Some(0) { bufs = &bufs[1..]; } if bufs.is_empty() { return Poll::Ready(Ok(0)); } // Flush if the first buffer doesn't fit. let first_len = bufs[0].len(); if first_len > self.buf.capacity() - self.buf.len() { ready!(self.as_mut().flush_buf(cx))?; debug_assert!(self.buf.is_empty()); } let me = self.as_mut().project(); if first_len >= me.buf.capacity() { // The slice is at least as large as the buffering capacity, // so it's better to write it directly, bypassing the buffer. debug_assert!(me.buf.is_empty()); return me.inner.poll_write(cx, &bufs[0]); } else { me.buf.extend_from_slice(&bufs[0]); bufs = &bufs[1..]; } let mut total_written = first_len; debug_assert!(total_written != 0); // Append the buffers that fit in the internal buffer. for buf in bufs { if buf.len() > me.buf.capacity() - me.buf.len() { break; } else { me.buf.extend_from_slice(buf); total_written += buf.len(); } } Poll::Ready(Ok(total_written)) } } fn is_write_vectored(&self) -> bool { true } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(self.as_mut().flush_buf(cx))?; self.get_pin_mut().poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(self.as_mut().flush_buf(cx))?; self.get_pin_mut().poll_shutdown(cx) } } #[derive(Debug, Clone, Copy)] pub(super) enum SeekState { /// `start_seek` has not been called. Init, /// `start_seek` has been called, but `poll_complete` has not yet been called. Start(SeekFrom), /// Waiting for completion of `poll_complete`. Pending, } /// Seek to the offset, in bytes, in the underlying writer. /// /// Seeking always writes out the internal buffer before seeking. impl AsyncSeek for BufWriter { fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { // We need to flush the internal buffer before seeking. // It receives a `Context` and returns a `Poll`, so it cannot be called // inside `start_seek`. *self.project().seek_state = SeekState::Start(pos); Ok(()) } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let pos = match self.seek_state { SeekState::Init => { return self.project().inner.poll_complete(cx); } SeekState::Start(pos) => Some(pos), SeekState::Pending => None, }; // Flush the internal buffer before seeking. ready!(self.as_mut().flush_buf(cx))?; let mut me = self.project(); if let Some(pos) = pos { // Ensure previous seeks have finished before starting a new one ready!(me.inner.as_mut().poll_complete(cx))?; if let Err(e) = me.inner.as_mut().start_seek(pos) { *me.seek_state = SeekState::Init; return Poll::Ready(Err(e)); } } match me.inner.poll_complete(cx) { Poll::Ready(res) => { *me.seek_state = SeekState::Init; Poll::Ready(res) } Poll::Pending => { *me.seek_state = SeekState::Pending; Poll::Pending } } } } impl AsyncRead for BufWriter { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.get_pin_mut().poll_read(cx, buf) } } impl AsyncBufRead for BufWriter { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.get_pin_mut().poll_fill_buf(cx) } fn consume(self: Pin<&mut Self>, amt: usize) { self.get_pin_mut().consume(amt); } } impl fmt::Debug for BufWriter { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BufWriter") .field("writer", &self.inner) .field( "buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity()), ) .field("written", &self.written) .finish() } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/chain.rs000064400000000000000000000067251046102023000146540ustar 00000000000000use crate::io::{AsyncBufRead, AsyncRead, ReadBuf}; use pin_project_lite::pin_project; use std::fmt; use std::io; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// Stream for the [`chain`](super::AsyncReadExt::chain) method. #[must_use = "streams do nothing unless polled"] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct Chain { #[pin] first: T, #[pin] second: U, done_first: bool, } } pub(super) fn chain(first: T, second: U) -> Chain where T: AsyncRead, U: AsyncRead, { Chain { first, second, done_first: false, } } impl Chain where T: AsyncRead, U: AsyncRead, { /// Gets references to the underlying readers in this `Chain`. pub fn get_ref(&self) -> (&T, &U) { (&self.first, &self.second) } /// Gets mutable references to the underlying readers in this `Chain`. /// /// Care should be taken to avoid modifying the internal I/O state of the /// underlying readers as doing so may corrupt the internal state of this /// `Chain`. pub fn get_mut(&mut self) -> (&mut T, &mut U) { (&mut self.first, &mut self.second) } /// Gets pinned mutable references to the underlying readers in this `Chain`. /// /// Care should be taken to avoid modifying the internal I/O state of the /// underlying readers as doing so may corrupt the internal state of this /// `Chain`. pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) { let me = self.project(); (me.first, me.second) } /// Consumes the `Chain`, returning the wrapped readers. pub fn into_inner(self) -> (T, U) { (self.first, self.second) } } impl fmt::Debug for Chain where T: fmt::Debug, U: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Chain") .field("t", &self.first) .field("u", &self.second) .finish() } } impl AsyncRead for Chain where T: AsyncRead, U: AsyncRead, { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let me = self.project(); if !*me.done_first { let rem = buf.remaining(); ready!(me.first.poll_read(cx, buf))?; if buf.remaining() == rem { *me.done_first = true; } else { return Poll::Ready(Ok(())); } } me.second.poll_read(cx, buf) } } impl AsyncBufRead for Chain where T: AsyncBufRead, U: AsyncBufRead, { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); if !*me.done_first { match ready!(me.first.poll_fill_buf(cx)?) { [] => { *me.done_first = true; } buf => return Poll::Ready(Ok(buf)), } } me.second.poll_fill_buf(cx) } fn consume(self: Pin<&mut Self>, amt: usize) { let me = self.project(); if !*me.done_first { me.first.consume(amt) } else { me.second.consume(amt) } } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/copy.rs000064400000000000000000000233501046102023000145350ustar 00000000000000use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; use std::future::Future; use std::io; use std::pin::Pin; use std::task::{ready, Context, Poll}; #[derive(Debug)] pub(super) struct CopyBuffer { read_done: bool, need_flush: bool, pos: usize, cap: usize, amt: u64, buf: Box<[u8]>, } impl CopyBuffer { pub(super) fn new(buf_size: usize) -> Self { Self { read_done: false, need_flush: false, pos: 0, cap: 0, amt: 0, buf: vec![0; buf_size].into_boxed_slice(), } } fn poll_fill_buf( &mut self, cx: &mut Context<'_>, reader: Pin<&mut R>, ) -> Poll> where R: AsyncRead + ?Sized, { let me = &mut *self; let mut buf = ReadBuf::new(&mut me.buf); buf.set_filled(me.cap); let res = reader.poll_read(cx, &mut buf); if let Poll::Ready(Ok(())) = res { let filled_len = buf.filled().len(); me.read_done = me.cap == filled_len; me.cap = filled_len; } res } fn poll_write_buf( &mut self, cx: &mut Context<'_>, mut reader: Pin<&mut R>, mut writer: Pin<&mut W>, ) -> Poll> where R: AsyncRead + ?Sized, W: AsyncWrite + ?Sized, { let me = &mut *self; match writer.as_mut().poll_write(cx, &me.buf[me.pos..me.cap]) { Poll::Pending => { // Top up the buffer towards full if we can read a bit more // data - this should improve the chances of a large write if !me.read_done && me.cap < me.buf.len() { ready!(me.poll_fill_buf(cx, reader.as_mut()))?; } Poll::Pending } res => res, } } pub(super) fn poll_copy( &mut self, cx: &mut Context<'_>, mut reader: Pin<&mut R>, mut writer: Pin<&mut W>, ) -> Poll> where R: AsyncRead + ?Sized, W: AsyncWrite + ?Sized, { ready!(crate::trace::trace_leaf(cx)); #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); loop { // If there is some space left in our buffer, then we try to read some // data to continue, thus maximizing the chances of a large write. if self.cap < self.buf.len() && !self.read_done { match self.poll_fill_buf(cx, reader.as_mut()) { Poll::Ready(Ok(())) => { #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] coop.made_progress(); } Poll::Ready(Err(err)) => { #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] coop.made_progress(); return Poll::Ready(Err(err)); } Poll::Pending => { // Ignore pending reads when our buffer is not empty, because // we can try to write data immediately. if self.pos == self.cap { // Try flushing when the reader has no progress to avoid deadlock // when the reader depends on buffered writer. if self.need_flush { ready!(writer.as_mut().poll_flush(cx))?; #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] coop.made_progress(); self.need_flush = false; } return Poll::Pending; } } } } // If our buffer has some data, let's write it out! while self.pos < self.cap { let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?; #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] coop.made_progress(); if i == 0 { return Poll::Ready(Err(io::Error::new( io::ErrorKind::WriteZero, "write zero byte into writer", ))); } else { self.pos += i; self.amt += i as u64; self.need_flush = true; } } // If pos larger than cap, this loop will never stop. // In particular, user's wrong poll_write implementation returning // incorrect written length may lead to thread blocking. debug_assert!( self.pos <= self.cap, "writer returned length larger than input slice" ); // All data has been written, the buffer can be considered empty again self.pos = 0; self.cap = 0; // If we've written all the data and we've seen EOF, flush out the // data and finish the transfer. if self.read_done { ready!(writer.as_mut().poll_flush(cx))?; #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] coop.made_progress(); return Poll::Ready(Ok(self.amt)); } } } } /// A future that asynchronously copies the entire contents of a reader into a /// writer. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] struct Copy<'a, R: ?Sized, W: ?Sized> { reader: &'a mut R, writer: &'a mut W, buf: CopyBuffer, } cfg_io_util! { /// Asynchronously copies the entire contents of a reader into a writer. /// /// This function returns a future that will continuously read data from /// `reader` and then write it into `writer` in a streaming fashion until /// `reader` returns EOF or fails. /// /// On success, the total number of bytes that were copied from `reader` to /// `writer` is returned. /// /// This is an asynchronous version of [`std::io::copy`][std]. /// /// A heap-allocated copy buffer with 8 KB is created to take data from the /// reader to the writer, check [`copy_buf`] if you want an alternative for /// [`AsyncBufRead`]. You can use `copy_buf` with [`BufReader`] to change the /// buffer capacity. /// /// [std]: std::io::copy /// [`copy_buf`]: crate::io::copy_buf /// [`AsyncBufRead`]: crate::io::AsyncBufRead /// [`BufReader`]: crate::io::BufReader /// /// # Errors /// /// The returned future will return an error immediately if any call to /// `poll_read` or `poll_write` returns an error. /// /// # Examples /// /// ``` /// use tokio::io; /// /// # async fn dox() -> std::io::Result<()> { /// let mut reader: &[u8] = b"hello"; /// let mut writer: Vec = vec![]; /// /// io::copy(&mut reader, &mut writer).await?; /// /// assert_eq!(&b"hello"[..], &writer[..]); /// # Ok(()) /// # } /// ``` pub async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result where R: AsyncRead + Unpin + ?Sized, W: AsyncWrite + Unpin + ?Sized, { Copy { reader, writer, buf: CopyBuffer::new(super::DEFAULT_BUF_SIZE) }.await } } impl Future for Copy<'_, R, W> where R: AsyncRead + Unpin + ?Sized, W: AsyncWrite + Unpin + ?Sized, { type Output = io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = &mut *self; me.buf .poll_copy(cx, Pin::new(&mut *me.reader), Pin::new(&mut *me.writer)) } } tokio-1.43.1/src/io/util/copy_bidirectional.rs000064400000000000000000000105121046102023000174210ustar 00000000000000use super::copy::CopyBuffer; use crate::io::{AsyncRead, AsyncWrite}; use std::future::poll_fn; use std::io; use std::pin::Pin; use std::task::{ready, Context, Poll}; enum TransferState { Running(CopyBuffer), ShuttingDown(u64), Done(u64), } fn transfer_one_direction( cx: &mut Context<'_>, state: &mut TransferState, r: &mut A, w: &mut B, ) -> Poll> where A: AsyncRead + AsyncWrite + Unpin + ?Sized, B: AsyncRead + AsyncWrite + Unpin + ?Sized, { let mut r = Pin::new(r); let mut w = Pin::new(w); loop { match state { TransferState::Running(buf) => { let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?; *state = TransferState::ShuttingDown(count); } TransferState::ShuttingDown(count) => { ready!(w.as_mut().poll_shutdown(cx))?; *state = TransferState::Done(*count); } TransferState::Done(count) => return Poll::Ready(Ok(*count)), } } } /// Copies data in both directions between `a` and `b`. /// /// This function returns a future that will read from both streams, /// writing any data read to the opposing stream. /// This happens in both directions concurrently. /// /// If an EOF is observed on one stream, [`shutdown()`] will be invoked on /// the other, and reading from that stream will stop. Copying of data in /// the other direction will continue. /// /// The future will complete successfully once both directions of communication has been shut down. /// A direction is shut down when the reader reports EOF, /// at which point [`shutdown()`] is called on the corresponding writer. When finished, /// it will return a tuple of the number of bytes copied from a to b /// and the number of bytes copied from b to a, in that order. /// /// It uses two 8 KB buffers for transferring bytes between `a` and `b` by default. /// To set your own buffers sizes use [`copy_bidirectional_with_sizes()`]. /// /// [`shutdown()`]: crate::io::AsyncWriteExt::shutdown /// /// # Errors /// /// The future will immediately return an error if any IO operation on `a` /// or `b` returns an error. Some data read from either stream may be lost (not /// written to the other stream) in this case. /// /// # Return value /// /// Returns a tuple of bytes copied `a` to `b` and bytes copied `b` to `a`. #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub async fn copy_bidirectional(a: &mut A, b: &mut B) -> io::Result<(u64, u64)> where A: AsyncRead + AsyncWrite + Unpin + ?Sized, B: AsyncRead + AsyncWrite + Unpin + ?Sized, { copy_bidirectional_impl( a, b, CopyBuffer::new(super::DEFAULT_BUF_SIZE), CopyBuffer::new(super::DEFAULT_BUF_SIZE), ) .await } /// Copies data in both directions between `a` and `b` using buffers of the specified size. /// /// This method is the same as the [`copy_bidirectional()`], except that it allows you to set the /// size of the internal buffers used when copying data. #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub async fn copy_bidirectional_with_sizes( a: &mut A, b: &mut B, a_to_b_buf_size: usize, b_to_a_buf_size: usize, ) -> io::Result<(u64, u64)> where A: AsyncRead + AsyncWrite + Unpin + ?Sized, B: AsyncRead + AsyncWrite + Unpin + ?Sized, { copy_bidirectional_impl( a, b, CopyBuffer::new(a_to_b_buf_size), CopyBuffer::new(b_to_a_buf_size), ) .await } async fn copy_bidirectional_impl( a: &mut A, b: &mut B, a_to_b_buffer: CopyBuffer, b_to_a_buffer: CopyBuffer, ) -> io::Result<(u64, u64)> where A: AsyncRead + AsyncWrite + Unpin + ?Sized, B: AsyncRead + AsyncWrite + Unpin + ?Sized, { let mut a_to_b = TransferState::Running(a_to_b_buffer); let mut b_to_a = TransferState::Running(b_to_a_buffer); poll_fn(|cx| { let a_to_b = transfer_one_direction(cx, &mut a_to_b, a, b)?; let b_to_a = transfer_one_direction(cx, &mut b_to_a, b, a)?; // It is not a problem if ready! returns early because transfer_one_direction for the // other direction will keep returning TransferState::Done(count) in future calls to poll let a_to_b = ready!(a_to_b); let b_to_a = ready!(b_to_a); Poll::Ready(Ok((a_to_b, b_to_a))) }) .await } tokio-1.43.1/src/io/util/copy_buf.rs000064400000000000000000000063051046102023000153720ustar 00000000000000use crate::io::{AsyncBufRead, AsyncWrite}; use std::future::Future; use std::io; use std::pin::Pin; use std::task::{ready, Context, Poll}; cfg_io_util! { /// A future that asynchronously copies the entire contents of a reader into a /// writer. /// /// This struct is generally created by calling [`copy_buf`][copy_buf]. Please /// see the documentation of `copy_buf()` for more details. /// /// [copy_buf]: copy_buf() #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] struct CopyBuf<'a, R: ?Sized, W: ?Sized> { reader: &'a mut R, writer: &'a mut W, amt: u64, } /// Asynchronously copies the entire contents of a reader into a writer. /// /// This function returns a future that will continuously read data from /// `reader` and then write it into `writer` in a streaming fashion until /// `reader` returns EOF or fails. /// /// On success, the total number of bytes that were copied from `reader` to /// `writer` is returned. /// /// This is a [`tokio::io::copy`] alternative for [`AsyncBufRead`] readers /// with no extra buffer allocation, since [`AsyncBufRead`] allow access /// to the reader's inner buffer. /// /// [`tokio::io::copy`]: crate::io::copy /// [`AsyncBufRead`]: crate::io::AsyncBufRead /// /// # Errors /// /// The returned future will finish with an error will return an error /// immediately if any call to `poll_fill_buf` or `poll_write` returns an /// error. /// /// # Examples /// /// ``` /// use tokio::io; /// /// # async fn dox() -> std::io::Result<()> { /// let mut reader: &[u8] = b"hello"; /// let mut writer: Vec = vec![]; /// /// io::copy_buf(&mut reader, &mut writer).await?; /// /// assert_eq!(b"hello", &writer[..]); /// # Ok(()) /// # } /// ``` pub async fn copy_buf<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result where R: AsyncBufRead + Unpin + ?Sized, W: AsyncWrite + Unpin + ?Sized, { CopyBuf { reader, writer, amt: 0, }.await } } impl Future for CopyBuf<'_, R, W> where R: AsyncBufRead + Unpin + ?Sized, W: AsyncWrite + Unpin + ?Sized, { type Output = io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { loop { let me = &mut *self; let buffer = ready!(Pin::new(&mut *me.reader).poll_fill_buf(cx))?; if buffer.is_empty() { ready!(Pin::new(&mut self.writer).poll_flush(cx))?; return Poll::Ready(Ok(self.amt)); } let i = ready!(Pin::new(&mut *me.writer).poll_write(cx, buffer))?; if i == 0 { return Poll::Ready(Err(std::io::ErrorKind::WriteZero.into())); } self.amt += i as u64; Pin::new(&mut *self.reader).consume(i); } } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { use std::marker::PhantomPinned; crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/empty.rs000064400000000000000000000107161046102023000147230ustar 00000000000000use crate::io::util::poll_proceed_and_make_progress; use crate::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use std::fmt; use std::io::{self, SeekFrom}; use std::pin::Pin; use std::task::{ready, Context, Poll}; cfg_io_util! { /// `Empty` ignores any data written via [`AsyncWrite`], and will always be empty /// (returning zero bytes) when read via [`AsyncRead`]. /// /// This struct is generally created by calling [`empty`]. Please see /// the documentation of [`empty()`][`empty`] for more details. /// /// This is an asynchronous version of [`std::io::empty`][std]. /// /// [`empty`]: fn@empty /// [std]: std::io::empty pub struct Empty { _p: (), } /// Creates a value that is always at EOF for reads, and ignores all data written. /// /// All writes on the returned instance will return `Poll::Ready(Ok(buf.len()))` /// and the contents of the buffer will not be inspected. /// /// All reads from the returned instance will return `Poll::Ready(Ok(0))`. /// /// This is an asynchronous version of [`std::io::empty`][std]. /// /// [std]: std::io::empty /// /// # Examples /// /// A slightly sad example of not reading anything into a buffer: /// /// ``` /// use tokio::io::{self, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() { /// let mut buffer = String::new(); /// io::empty().read_to_string(&mut buffer).await.unwrap(); /// assert!(buffer.is_empty()); /// } /// ``` /// /// A convoluted way of getting the length of a buffer: /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() { /// let buffer = vec![1, 2, 3, 5, 8]; /// let num_bytes = io::empty().write(&buffer).await.unwrap(); /// assert_eq!(num_bytes, 5); /// } /// ``` pub fn empty() -> Empty { Empty { _p: () } } } impl AsyncRead for Empty { #[inline] fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, _: &mut ReadBuf<'_>, ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } } impl AsyncBufRead for Empty { #[inline] fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(&[])) } #[inline] fn consume(self: Pin<&mut Self>, _: usize) {} } impl AsyncWrite for Empty { #[inline] fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(buf.len())) } #[inline] fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } #[inline] fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } #[inline] fn is_write_vectored(&self) -> bool { true } #[inline] fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); let num_bytes = bufs.iter().map(|b| b.len()).sum(); Poll::Ready(Ok(num_bytes)) } } impl AsyncSeek for Empty { #[inline] fn start_seek(self: Pin<&mut Self>, _position: SeekFrom) -> io::Result<()> { Ok(()) } #[inline] fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(0)) } } impl fmt::Debug for Empty { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.pad("Empty { .. }") } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::(); } } tokio-1.43.1/src/io/util/fill_buf.rs000064400000000000000000000037241046102023000153500ustar 00000000000000use crate::io::AsyncBufRead; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { /// Future for the [`fill_buf`](crate::io::AsyncBufReadExt::fill_buf) method. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct FillBuf<'a, R: ?Sized> { reader: Option<&'a mut R>, #[pin] _pin: PhantomPinned, } } pub(crate) fn fill_buf(reader: &mut R) -> FillBuf<'_, R> where R: AsyncBufRead + ?Sized + Unpin, { FillBuf { reader: Some(reader), _pin: PhantomPinned, } } impl<'a, R: AsyncBufRead + ?Sized + Unpin> Future for FillBuf<'a, R> { type Output = io::Result<&'a [u8]>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); let reader = me.reader.take().expect("Polled after completion."); match Pin::new(&mut *reader).poll_fill_buf(cx) { Poll::Ready(Ok(slice)) => unsafe { // Safety: This is necessary only due to a limitation in the // borrow checker. Once Rust starts using the polonius borrow // checker, this can be simplified. // // The safety of this transmute relies on the fact that the // value of `reader` is `None` when we return in this branch. // Otherwise the caller could poll us again after // completion, and access the mutable reference while the // returned immutable reference still exists. let slice = std::mem::transmute::<&[u8], &'a [u8]>(slice); Poll::Ready(Ok(slice)) }, Poll::Ready(Err(err)) => Poll::Ready(Err(err)), Poll::Pending => { *me.reader = Some(reader); Poll::Pending } } } } tokio-1.43.1/src/io/util/flush.rs000064400000000000000000000022121046102023000146760ustar 00000000000000use crate::io::AsyncWrite; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { /// A future used to fully flush an I/O object. /// /// Created by the [`AsyncWriteExt::flush`][flush] function. /// /// [flush]: crate::io::AsyncWriteExt::flush #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Flush<'a, A: ?Sized> { a: &'a mut A, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } /// Creates a future which will entirely flush an I/O object. pub(super) fn flush(a: &mut A) -> Flush<'_, A> where A: AsyncWrite + Unpin + ?Sized, { Flush { a, _pin: PhantomPinned, } } impl Future for Flush<'_, A> where A: AsyncWrite + Unpin + ?Sized, { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); Pin::new(&mut *me.a).poll_flush(cx) } } tokio-1.43.1/src/io/util/lines.rs000064400000000000000000000075761046102023000147110ustar 00000000000000use crate::io::util::read_line::read_line_internal; use crate::io::AsyncBufRead; use pin_project_lite::pin_project; use std::io; use std::mem; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// Reads lines from an [`AsyncBufRead`]. /// /// A `Lines` can be turned into a `Stream` with [`LinesStream`]. /// /// This type is usually created using the [`lines`] method. /// /// [`AsyncBufRead`]: crate::io::AsyncBufRead /// [`LinesStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.LinesStream.html /// [`lines`]: crate::io::AsyncBufReadExt::lines #[derive(Debug)] #[must_use = "streams do nothing unless polled"] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct Lines { #[pin] reader: R, buf: String, bytes: Vec, read: usize, } } pub(crate) fn lines(reader: R) -> Lines where R: AsyncBufRead, { Lines { reader, buf: String::new(), bytes: Vec::new(), read: 0, } } impl Lines where R: AsyncBufRead + Unpin, { /// Returns the next line in the stream. /// /// # Cancel safety /// /// This method is cancellation safe. /// /// # Examples /// /// ``` /// # use tokio::io::AsyncBufRead; /// use tokio::io::AsyncBufReadExt; /// /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> { /// let mut lines = my_buf_read.lines(); /// /// while let Some(line) = lines.next_line().await? { /// println!("length = {}", line.len()) /// } /// # Ok(()) /// # } /// ``` pub async fn next_line(&mut self) -> io::Result> { use std::future::poll_fn; poll_fn(|cx| Pin::new(&mut *self).poll_next_line(cx)).await } /// Obtains a mutable reference to the underlying reader. pub fn get_mut(&mut self) -> &mut R { &mut self.reader } /// Obtains a reference to the underlying reader. pub fn get_ref(&mut self) -> &R { &self.reader } /// Unwraps this `Lines`, returning the underlying reader. /// /// Note that any leftover data in the internal buffer is lost. /// Therefore, a following read from the underlying reader may lead to data loss. pub fn into_inner(self) -> R { self.reader } } impl Lines where R: AsyncBufRead, { /// Polls for the next line in the stream. /// /// This method returns: /// /// * `Poll::Pending` if the next line is not yet available. /// * `Poll::Ready(Ok(Some(line)))` if the next line is available. /// * `Poll::Ready(Ok(None))` if there are no more lines in this stream. /// * `Poll::Ready(Err(err))` if an IO error occurred while reading the next line. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when more bytes become /// available on the underlying IO resource. Note that on multiple calls to /// `poll_next_line`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. pub fn poll_next_line( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll>> { let me = self.project(); let n = ready!(read_line_internal(me.reader, cx, me.buf, me.bytes, me.read))?; debug_assert_eq!(*me.read, 0); if n == 0 && me.buf.is_empty() { return Poll::Ready(Ok(None)); } if me.buf.ends_with('\n') { me.buf.pop(); if me.buf.ends_with('\r') { me.buf.pop(); } } Poll::Ready(Ok(Some(mem::take(me.buf)))) } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/mem.rs000064400000000000000000000315451046102023000143460ustar 00000000000000//! In-process memory IO types. use crate::io::{split, AsyncRead, AsyncWrite, ReadBuf, ReadHalf, WriteHalf}; use crate::loom::sync::Mutex; use bytes::{Buf, BytesMut}; use std::{ pin::Pin, sync::Arc, task::{self, ready, Poll, Waker}, }; /// A bidirectional pipe to read and write bytes in memory. /// /// A pair of `DuplexStream`s are created together, and they act as a "channel" /// that can be used as in-memory IO types. Writing to one of the pairs will /// allow that data to be read from the other, and vice versa. /// /// # Closing a `DuplexStream` /// /// If one end of the `DuplexStream` channel is dropped, any pending reads on /// the other side will continue to read data until the buffer is drained, then /// they will signal EOF by returning 0 bytes. Any writes to the other side, /// including pending ones (that are waiting for free space in the buffer) will /// return `Err(BrokenPipe)` immediately. /// /// # Example /// /// ``` /// # async fn ex() -> std::io::Result<()> { /// # use tokio::io::{AsyncReadExt, AsyncWriteExt}; /// let (mut client, mut server) = tokio::io::duplex(64); /// /// client.write_all(b"ping").await?; /// /// let mut buf = [0u8; 4]; /// server.read_exact(&mut buf).await?; /// assert_eq!(&buf, b"ping"); /// /// server.write_all(b"pong").await?; /// /// client.read_exact(&mut buf).await?; /// assert_eq!(&buf, b"pong"); /// # Ok(()) /// # } /// ``` #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct DuplexStream { read: Arc>, write: Arc>, } /// A unidirectional pipe to read and write bytes in memory. /// /// It can be constructed by [`simplex`] function which will create a pair of /// reader and writer or by calling [`SimplexStream::new_unsplit`] that will /// create a handle for both reading and writing. /// /// # Example /// /// ``` /// # async fn ex() -> std::io::Result<()> { /// # use tokio::io::{AsyncReadExt, AsyncWriteExt}; /// let (mut receiver, mut sender) = tokio::io::simplex(64); /// /// sender.write_all(b"ping").await?; /// /// let mut buf = [0u8; 4]; /// receiver.read_exact(&mut buf).await?; /// assert_eq!(&buf, b"ping"); /// # Ok(()) /// # } /// ``` #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct SimplexStream { /// The buffer storing the bytes written, also read from. /// /// Using a `BytesMut` because it has efficient `Buf` and `BufMut` /// functionality already. Additionally, it can try to copy data in the /// same buffer if there read index has advanced far enough. buffer: BytesMut, /// Determines if the write side has been closed. is_closed: bool, /// The maximum amount of bytes that can be written before returning /// `Poll::Pending`. max_buf_size: usize, /// If the `read` side has been polled and is pending, this is the waker /// for that parked task. read_waker: Option, /// If the `write` side has filled the `max_buf_size` and returned /// `Poll::Pending`, this is the waker for that parked task. write_waker: Option, } // ===== impl DuplexStream ===== /// Create a new pair of `DuplexStream`s that act like a pair of connected sockets. /// /// The `max_buf_size` argument is the maximum amount of bytes that can be /// written to a side before the write returns `Poll::Pending`. #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub fn duplex(max_buf_size: usize) -> (DuplexStream, DuplexStream) { let one = Arc::new(Mutex::new(SimplexStream::new_unsplit(max_buf_size))); let two = Arc::new(Mutex::new(SimplexStream::new_unsplit(max_buf_size))); ( DuplexStream { read: one.clone(), write: two.clone(), }, DuplexStream { read: two, write: one, }, ) } impl AsyncRead for DuplexStream { // Previous rustc required this `self` to be `mut`, even though newer // versions recognize it isn't needed to call `lock()`. So for // compatibility, we include the `mut` and `allow` the lint. // // See https://github.com/rust-lang/rust/issues/73592 #[allow(unused_mut)] fn poll_read( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut *self.read.lock()).poll_read(cx, buf) } } impl AsyncWrite for DuplexStream { #[allow(unused_mut)] fn poll_write( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut *self.write.lock()).poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut task::Context<'_>, bufs: &[std::io::IoSlice<'_>], ) -> Poll> { Pin::new(&mut *self.write.lock()).poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { true } #[allow(unused_mut)] fn poll_flush( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, ) -> Poll> { Pin::new(&mut *self.write.lock()).poll_flush(cx) } #[allow(unused_mut)] fn poll_shutdown( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, ) -> Poll> { Pin::new(&mut *self.write.lock()).poll_shutdown(cx) } } impl Drop for DuplexStream { fn drop(&mut self) { // notify the other side of the closure self.write.lock().close_write(); self.read.lock().close_read(); } } // ===== impl SimplexStream ===== /// Creates unidirectional buffer that acts like in memory pipe. /// /// The `max_buf_size` argument is the maximum amount of bytes that can be /// written to a buffer before the it returns `Poll::Pending`. /// /// # Unify reader and writer /// /// The reader and writer half can be unified into a single structure /// of `SimplexStream` that supports both reading and writing or /// the `SimplexStream` can be already created as unified structure /// using [`SimplexStream::new_unsplit()`]. /// /// ``` /// # async fn ex() -> std::io::Result<()> { /// # use tokio::io::{AsyncReadExt, AsyncWriteExt}; /// let (writer, reader) = tokio::io::simplex(64); /// let mut simplex_stream = writer.unsplit(reader); /// simplex_stream.write_all(b"hello").await?; /// /// let mut buf = [0u8; 5]; /// simplex_stream.read_exact(&mut buf).await?; /// assert_eq!(&buf, b"hello"); /// # Ok(()) /// # } /// ``` #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub fn simplex(max_buf_size: usize) -> (ReadHalf, WriteHalf) { split(SimplexStream::new_unsplit(max_buf_size)) } impl SimplexStream { /// Creates unidirectional buffer that acts like in memory pipe. To create split /// version with separate reader and writer you can use [`simplex`] function. /// /// The `max_buf_size` argument is the maximum amount of bytes that can be /// written to a buffer before the it returns `Poll::Pending`. #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub fn new_unsplit(max_buf_size: usize) -> SimplexStream { SimplexStream { buffer: BytesMut::new(), is_closed: false, max_buf_size, read_waker: None, write_waker: None, } } fn close_write(&mut self) { self.is_closed = true; // needs to notify any readers that no more data will come if let Some(waker) = self.read_waker.take() { waker.wake(); } } fn close_read(&mut self) { self.is_closed = true; // needs to notify any writers that they have to abort if let Some(waker) = self.write_waker.take() { waker.wake(); } } fn poll_read_internal( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { if self.buffer.has_remaining() { let max = self.buffer.remaining().min(buf.remaining()); buf.put_slice(&self.buffer[..max]); self.buffer.advance(max); if max > 0 { // The passed `buf` might have been empty, don't wake up if // no bytes have been moved. if let Some(waker) = self.write_waker.take() { waker.wake(); } } Poll::Ready(Ok(())) } else if self.is_closed { Poll::Ready(Ok(())) } else { self.read_waker = Some(cx.waker().clone()); Poll::Pending } } fn poll_write_internal( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8], ) -> Poll> { if self.is_closed { return Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into())); } let avail = self.max_buf_size - self.buffer.len(); if avail == 0 { self.write_waker = Some(cx.waker().clone()); return Poll::Pending; } let len = buf.len().min(avail); self.buffer.extend_from_slice(&buf[..len]); if let Some(waker) = self.read_waker.take() { waker.wake(); } Poll::Ready(Ok(len)) } fn poll_write_vectored_internal( mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, bufs: &[std::io::IoSlice<'_>], ) -> Poll> { if self.is_closed { return Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into())); } let avail = self.max_buf_size - self.buffer.len(); if avail == 0 { self.write_waker = Some(cx.waker().clone()); return Poll::Pending; } let mut rem = avail; for buf in bufs { if rem == 0 { break; } let len = buf.len().min(rem); self.buffer.extend_from_slice(&buf[..len]); rem -= len; } if let Some(waker) = self.read_waker.take() { waker.wake(); } Poll::Ready(Ok(avail - rem)) } } impl AsyncRead for SimplexStream { cfg_coop! { fn poll_read( self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = self.poll_read_internal(cx, buf); if ret.is_ready() { coop.made_progress(); } ret } } cfg_not_coop! { fn poll_read( self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); self.poll_read_internal(cx, buf) } } } impl AsyncWrite for SimplexStream { cfg_coop! { fn poll_write( self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = self.poll_write_internal(cx, buf); if ret.is_ready() { coop.made_progress(); } ret } } cfg_not_coop! { fn poll_write( self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); self.poll_write_internal(cx, buf) } } cfg_coop! { fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut task::Context<'_>, bufs: &[std::io::IoSlice<'_>], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = self.poll_write_vectored_internal(cx, bufs); if ret.is_ready() { coop.made_progress(); } ret } } cfg_not_coop! { fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut task::Context<'_>, bufs: &[std::io::IoSlice<'_>], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); self.poll_write_vectored_internal(cx, bufs) } } fn is_write_vectored(&self) -> bool { true } fn poll_flush(self: Pin<&mut Self>, _: &mut task::Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown( mut self: Pin<&mut Self>, _: &mut task::Context<'_>, ) -> Poll> { self.close_write(); Poll::Ready(Ok(())) } } tokio-1.43.1/src/io/util/mod.rs000064400000000000000000000045121046102023000143410ustar 00000000000000#![allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 cfg_io_util! { mod async_buf_read_ext; pub use async_buf_read_ext::AsyncBufReadExt; mod async_read_ext; pub use async_read_ext::AsyncReadExt; mod async_seek_ext; pub use async_seek_ext::AsyncSeekExt; mod async_write_ext; pub use async_write_ext::AsyncWriteExt; mod buf_reader; pub use buf_reader::BufReader; mod buf_stream; pub use buf_stream::BufStream; mod buf_writer; pub use buf_writer::BufWriter; mod chain; mod copy; pub use copy::copy; mod copy_bidirectional; pub use copy_bidirectional::{copy_bidirectional, copy_bidirectional_with_sizes}; mod copy_buf; pub use copy_buf::copy_buf; mod empty; pub use empty::{empty, Empty}; mod flush; mod lines; pub use lines::Lines; mod mem; pub use mem::{duplex, simplex, DuplexStream, SimplexStream}; mod read; mod read_buf; mod read_exact; mod read_int; mod read_line; mod fill_buf; mod read_to_end; mod vec_with_initialized; cfg_process! { pub(crate) use read_to_end::read_to_end; } mod read_to_string; mod read_until; mod repeat; pub use repeat::{repeat, Repeat}; mod shutdown; mod sink; pub use sink::{sink, Sink}; mod split; pub use split::Split; mod take; pub use take::Take; mod write; mod write_vectored; mod write_all; mod write_buf; mod write_all_buf; mod write_int; // used by `BufReader` and `BufWriter` // https://github.com/rust-lang/rust/blob/master/library/std/src/sys_common/io.rs#L1 const DEFAULT_BUF_SIZE: usize = 8 * 1024; cfg_coop! { fn poll_proceed_and_make_progress(cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { let coop = std::task::ready!(crate::runtime::coop::poll_proceed(cx)); coop.made_progress(); std::task::Poll::Ready(()) } } cfg_not_coop! { fn poll_proceed_and_make_progress(_: &mut std::task::Context<'_>) -> std::task::Poll<()> { std::task::Poll::Ready(()) } } } cfg_not_io_util! { cfg_process! { mod vec_with_initialized; mod read_to_end; // Used by process pub(crate) use read_to_end::read_to_end; } } tokio-1.43.1/src/io/util/read.rs000064400000000000000000000030401046102023000144700ustar 00000000000000use crate::io::{AsyncRead, ReadBuf}; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::marker::Unpin; use std::pin::Pin; use std::task::{ready, Context, Poll}; /// Tries to read some bytes directly into the given `buf` in asynchronous /// manner, returning a future type. /// /// The returned future will resolve to both the I/O stream and the buffer /// as well as the number of bytes read once the read operation is completed. pub(crate) fn read<'a, R>(reader: &'a mut R, buf: &'a mut [u8]) -> Read<'a, R> where R: AsyncRead + Unpin + ?Sized, { Read { reader, buf, _pin: PhantomPinned, } } pin_project! { /// A future which can be used to easily read available number of bytes to fill /// a buffer. /// /// Created by the [`read`] function. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Read<'a, R: ?Sized> { reader: &'a mut R, buf: &'a mut [u8], // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } impl Future for Read<'_, R> where R: AsyncRead + Unpin + ?Sized, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); let mut buf = ReadBuf::new(me.buf); ready!(Pin::new(me.reader).poll_read(cx, &mut buf))?; Poll::Ready(Ok(buf.filled().len())) } } tokio-1.43.1/src/io/util/read_buf.rs000064400000000000000000000035771046102023000153430ustar 00000000000000use crate::io::AsyncRead; use bytes::BufMut; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{ready, Context, Poll}; pub(crate) fn read_buf<'a, R, B>(reader: &'a mut R, buf: &'a mut B) -> ReadBuf<'a, R, B> where R: AsyncRead + Unpin + ?Sized, B: BufMut + ?Sized, { ReadBuf { reader, buf, _pin: PhantomPinned, } } pin_project! { /// Future returned by [`read_buf`](crate::io::AsyncReadExt::read_buf). #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadBuf<'a, R: ?Sized, B: ?Sized> { reader: &'a mut R, buf: &'a mut B, #[pin] _pin: PhantomPinned, } } impl Future for ReadBuf<'_, R, B> where R: AsyncRead + Unpin + ?Sized, B: BufMut + ?Sized, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { use crate::io::ReadBuf; use std::mem::MaybeUninit; let me = self.project(); if !me.buf.has_remaining_mut() { return Poll::Ready(Ok(0)); } let n = { let dst = me.buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [MaybeUninit]) }; let mut buf = ReadBuf::uninit(dst); let ptr = buf.filled().as_ptr(); ready!(Pin::new(me.reader).poll_read(cx, &mut buf)?); // Ensure the pointer does not change from under us assert_eq!(ptr, buf.filled().as_ptr()); buf.filled().len() }; // Safety: This is guaranteed to be the number of initialized (and read) // bytes due to the invariants provided by `ReadBuf::filled`. unsafe { me.buf.advance_mut(n); } Poll::Ready(Ok(n)) } } tokio-1.43.1/src/io/util/read_exact.rs000064400000000000000000000037011046102023000156600ustar 00000000000000use crate::io::{AsyncRead, ReadBuf}; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::marker::Unpin; use std::pin::Pin; use std::task::{ready, Context, Poll}; /// A future which can be used to easily read exactly enough bytes to fill /// a buffer. /// /// Created by the [`AsyncReadExt::read_exact`][read_exact]. /// [`read_exact`]: [`crate::io::AsyncReadExt::read_exact`] pub(crate) fn read_exact<'a, A>(reader: &'a mut A, buf: &'a mut [u8]) -> ReadExact<'a, A> where A: AsyncRead + Unpin + ?Sized, { ReadExact { reader, buf: ReadBuf::new(buf), _pin: PhantomPinned, } } pin_project! { /// Creates a future which will read exactly enough bytes to fill `buf`, /// returning an error if EOF is hit sooner. /// /// On success the number of bytes is returned #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadExact<'a, A: ?Sized> { reader: &'a mut A, buf: ReadBuf<'a>, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } fn eof() -> io::Error { io::Error::new(io::ErrorKind::UnexpectedEof, "early eof") } impl Future for ReadExact<'_, A> where A: AsyncRead + Unpin + ?Sized, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); loop { // if our buffer is empty, then we need to read some data to continue. let rem = me.buf.remaining(); if rem != 0 { ready!(Pin::new(&mut *me.reader).poll_read(cx, me.buf))?; if me.buf.remaining() == rem { return Err(eof()).into(); } } else { return Poll::Ready(Ok(me.buf.capacity())); } } } } tokio-1.43.1/src/io/util/read_int.rs000064400000000000000000000112171046102023000153470ustar 00000000000000use crate::io::{AsyncRead, ReadBuf}; use bytes::Buf; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::io::ErrorKind::UnexpectedEof; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; macro_rules! reader { ($name:ident, $ty:ty, $reader:ident) => { reader!($name, $ty, $reader, std::mem::size_of::<$ty>()); }; ($name:ident, $ty:ty, $reader:ident, $bytes:expr) => { pin_project! { #[doc(hidden)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct $name { #[pin] src: R, buf: [u8; $bytes], read: u8, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } impl $name { pub(crate) fn new(src: R) -> Self { $name { src, buf: [0; $bytes], read: 0, _pin: PhantomPinned, } } } impl Future for $name where R: AsyncRead, { type Output = io::Result<$ty>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut me = self.project(); if *me.read == $bytes as u8 { return Poll::Ready(Ok(Buf::$reader(&mut &me.buf[..]))); } while *me.read < $bytes as u8 { let mut buf = ReadBuf::new(&mut me.buf[*me.read as usize..]); *me.read += match me.src.as_mut().poll_read(cx, &mut buf) { Poll::Pending => return Poll::Pending, Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), Poll::Ready(Ok(())) => { let n = buf.filled().len(); if n == 0 { return Poll::Ready(Err(UnexpectedEof.into())); } n as u8 } }; } let num = Buf::$reader(&mut &me.buf[..]); Poll::Ready(Ok(num)) } } }; } macro_rules! reader8 { ($name:ident, $ty:ty) => { pin_project! { /// Future returned from `read_u8` #[doc(hidden)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct $name { #[pin] reader: R, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } impl $name { pub(crate) fn new(reader: R) -> $name { $name { reader, _pin: PhantomPinned, } } } impl Future for $name where R: AsyncRead, { type Output = io::Result<$ty>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); let mut buf = [0; 1]; let mut buf = ReadBuf::new(&mut buf); match me.reader.poll_read(cx, &mut buf) { Poll::Pending => Poll::Pending, Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())), Poll::Ready(Ok(())) => { if buf.filled().len() == 0 { return Poll::Ready(Err(UnexpectedEof.into())); } Poll::Ready(Ok(buf.filled()[0] as $ty)) } } } } }; } reader8!(ReadU8, u8); reader8!(ReadI8, i8); reader!(ReadU16, u16, get_u16); reader!(ReadU32, u32, get_u32); reader!(ReadU64, u64, get_u64); reader!(ReadU128, u128, get_u128); reader!(ReadI16, i16, get_i16); reader!(ReadI32, i32, get_i32); reader!(ReadI64, i64, get_i64); reader!(ReadI128, i128, get_i128); reader!(ReadF32, f32, get_f32); reader!(ReadF64, f64, get_f64); reader!(ReadU16Le, u16, get_u16_le); reader!(ReadU32Le, u32, get_u32_le); reader!(ReadU64Le, u64, get_u64_le); reader!(ReadU128Le, u128, get_u128_le); reader!(ReadI16Le, i16, get_i16_le); reader!(ReadI32Le, i32, get_i32_le); reader!(ReadI64Le, i64, get_i64_le); reader!(ReadI128Le, i128, get_i128_le); reader!(ReadF32Le, f32, get_f32_le); reader!(ReadF64Le, f64, get_f64_le); tokio-1.43.1/src/io/util/read_line.rs000064400000000000000000000076631046102023000155160ustar 00000000000000use crate::io::util::read_until::read_until_internal; use crate::io::AsyncBufRead; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::mem; use std::pin::Pin; use std::string::FromUtf8Error; use std::task::{ready, Context, Poll}; pin_project! { /// Future for the [`read_line`](crate::io::AsyncBufReadExt::read_line) method. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadLine<'a, R: ?Sized> { reader: &'a mut R, // This is the buffer we were provided. It will be replaced with an empty string // while reading to postpone utf-8 handling until after reading. output: &'a mut String, // The actual allocation of the string is moved into this vector instead. buf: Vec, // The number of bytes appended to buf. This can be less than buf.len() if // the buffer was not empty when the operation was started. read: usize, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn read_line<'a, R>(reader: &'a mut R, string: &'a mut String) -> ReadLine<'a, R> where R: AsyncBufRead + ?Sized + Unpin, { ReadLine { reader, buf: mem::take(string).into_bytes(), output: string, read: 0, _pin: PhantomPinned, } } fn put_back_original_data(output: &mut String, mut vector: Vec, num_bytes_read: usize) { let original_len = vector.len() - num_bytes_read; vector.truncate(original_len); *output = String::from_utf8(vector).expect("The original data must be valid utf-8."); } /// This handles the various failure cases and puts the string back into `output`. /// /// The `truncate_on_io_error` `bool` is necessary because `read_to_string` and `read_line` /// disagree on what should happen when an IO error occurs. pub(super) fn finish_string_read( io_res: io::Result, utf8_res: Result, read: usize, output: &mut String, truncate_on_io_error: bool, ) -> Poll> { match (io_res, utf8_res) { (Ok(num_bytes), Ok(string)) => { debug_assert_eq!(read, 0); *output = string; Poll::Ready(Ok(num_bytes)) } (Err(io_err), Ok(string)) => { *output = string; if truncate_on_io_error { let original_len = output.len() - read; output.truncate(original_len); } Poll::Ready(Err(io_err)) } (Ok(num_bytes), Err(utf8_err)) => { debug_assert_eq!(read, 0); put_back_original_data(output, utf8_err.into_bytes(), num_bytes); Poll::Ready(Err(io::Error::new( io::ErrorKind::InvalidData, "stream did not contain valid UTF-8", ))) } (Err(io_err), Err(utf8_err)) => { put_back_original_data(output, utf8_err.into_bytes(), read); Poll::Ready(Err(io_err)) } } } pub(super) fn read_line_internal( reader: Pin<&mut R>, cx: &mut Context<'_>, output: &mut String, buf: &mut Vec, read: &mut usize, ) -> Poll> { let io_res = ready!(read_until_internal(reader, cx, b'\n', buf, read)); let utf8_res = String::from_utf8(mem::take(buf)); // At this point both buf and output are empty. The allocation is in utf8_res. debug_assert!(buf.is_empty()); debug_assert!(output.is_empty()); finish_string_read(io_res, utf8_res, *read, output, false) } impl Future for ReadLine<'_, R> { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); read_line_internal(Pin::new(*me.reader), cx, me.output, me.buf, me.read) } } tokio-1.43.1/src/io/util/read_to_end.rs000064400000000000000000000115601046102023000160260ustar 00000000000000use crate::io::util::vec_with_initialized::{into_read_buf_parts, VecU8, VecWithInitialized}; use crate::io::{AsyncRead, ReadBuf}; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::mem::{self, MaybeUninit}; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadToEnd<'a, R: ?Sized> { reader: &'a mut R, buf: VecWithInitialized<&'a mut Vec>, // The number of bytes appended to buf. This can be less than buf.len() if // the buffer was not empty when the operation was started. read: usize, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn read_to_end<'a, R>(reader: &'a mut R, buffer: &'a mut Vec) -> ReadToEnd<'a, R> where R: AsyncRead + Unpin + ?Sized, { ReadToEnd { reader, buf: VecWithInitialized::new(buffer), read: 0, _pin: PhantomPinned, } } pub(super) fn read_to_end_internal( buf: &mut VecWithInitialized, mut reader: Pin<&mut R>, num_read: &mut usize, cx: &mut Context<'_>, ) -> Poll> { loop { let ret = ready!(poll_read_to_end(buf, reader.as_mut(), cx)); match ret { Err(err) => return Poll::Ready(Err(err)), Ok(0) => return Poll::Ready(Ok(mem::replace(num_read, 0))), Ok(num) => { *num_read += num; } } } } /// Tries to read from the provided [`AsyncRead`]. /// /// The length of the buffer is increased by the number of bytes read. fn poll_read_to_end( buf: &mut VecWithInitialized, read: Pin<&mut R>, cx: &mut Context<'_>, ) -> Poll> { // This uses an adaptive system to extend the vector when it fills. We want to // avoid paying to allocate and zero a huge chunk of memory if the reader only // has 4 bytes while still making large reads if the reader does have a ton // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every // time is 4,500 times (!) slower than this if the reader has a very small // amount of data to return. When the vector is full with its starting // capacity, we first try to read into a small buffer to see if we reached // an EOF. This only happens when the starting capacity is >= NUM_BYTES, since // we allocate at least NUM_BYTES each time. This avoids the unnecessary // allocation that we attempt before reading into the vector. const NUM_BYTES: usize = 32; let try_small_read = buf.try_small_read_first(NUM_BYTES); // Get a ReadBuf into the vector. let mut read_buf; let poll_result; let n = if try_small_read { // Read some bytes using a small read. let mut small_buf: [MaybeUninit; NUM_BYTES] = [MaybeUninit::uninit(); NUM_BYTES]; let mut small_read_buf = ReadBuf::uninit(&mut small_buf); poll_result = read.poll_read(cx, &mut small_read_buf); let to_write = small_read_buf.filled(); // Ensure we have enough space to fill our vector with what we read. read_buf = buf.get_read_buf(); if to_write.len() > read_buf.remaining() { buf.reserve(NUM_BYTES); read_buf = buf.get_read_buf(); } read_buf.put_slice(to_write); to_write.len() } else { // Ensure we have enough space for reading. buf.reserve(NUM_BYTES); read_buf = buf.get_read_buf(); // Read data directly into vector. let filled_before = read_buf.filled().len(); poll_result = read.poll_read(cx, &mut read_buf); // Compute the number of bytes read. read_buf.filled().len() - filled_before }; // Update the length of the vector using the result of poll_read. let read_buf_parts = into_read_buf_parts(read_buf); buf.apply_read_buf(read_buf_parts); match poll_result { Poll::Pending => { // In this case, nothing should have been read. However we still // update the vector in case the poll_read call initialized parts of // the vector's unused capacity. debug_assert_eq!(n, 0); Poll::Pending } Poll::Ready(Err(err)) => { debug_assert_eq!(n, 0); Poll::Ready(Err(err)) } Poll::Ready(Ok(())) => Poll::Ready(Ok(n)), } } impl Future for ReadToEnd<'_, A> where A: AsyncRead + ?Sized + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); read_to_end_internal(me.buf, Pin::new(*me.reader), me.read, cx) } } tokio-1.43.1/src/io/util/read_to_string.rs000064400000000000000000000047211046102023000165670ustar 00000000000000use crate::io::util::read_line::finish_string_read; use crate::io::util::read_to_end::read_to_end_internal; use crate::io::util::vec_with_initialized::VecWithInitialized; use crate::io::AsyncRead; use pin_project_lite::pin_project; use std::future::Future; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{ready, Context, Poll}; use std::{io, mem}; pin_project! { /// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadToString<'a, R: ?Sized> { reader: &'a mut R, // This is the buffer we were provided. It will be replaced with an empty string // while reading to postpone utf-8 handling until after reading. output: &'a mut String, // The actual allocation of the string is moved into this vector instead. buf: VecWithInitialized>, // The number of bytes appended to buf. This can be less than buf.len() if // the buffer was not empty when the operation was started. read: usize, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn read_to_string<'a, R>( reader: &'a mut R, string: &'a mut String, ) -> ReadToString<'a, R> where R: AsyncRead + ?Sized + Unpin, { let buf = mem::take(string).into_bytes(); ReadToString { reader, buf: VecWithInitialized::new(buf), output: string, read: 0, _pin: PhantomPinned, } } fn read_to_string_internal( reader: Pin<&mut R>, output: &mut String, buf: &mut VecWithInitialized>, read: &mut usize, cx: &mut Context<'_>, ) -> Poll> { let io_res = ready!(read_to_end_internal(buf, reader, read, cx)); let utf8_res = String::from_utf8(buf.take()); // At this point both buf and output are empty. The allocation is in utf8_res. debug_assert!(buf.is_empty()); debug_assert!(output.is_empty()); finish_string_read(io_res, utf8_res, *read, output, true) } impl Future for ReadToString<'_, A> where A: AsyncRead + ?Sized + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); read_to_string_internal(Pin::new(*me.reader), me.output, me.buf, me.read, cx) } } tokio-1.43.1/src/io/util/read_until.rs000064400000000000000000000043761046102023000157200ustar 00000000000000use crate::io::AsyncBufRead; use crate::util::memchr; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::mem; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// Future for the [`read_until`](crate::io::AsyncBufReadExt::read_until) method. /// The delimiter is included in the resulting vector. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadUntil<'a, R: ?Sized> { reader: &'a mut R, delimiter: u8, buf: &'a mut Vec, // The number of bytes appended to buf. This can be less than buf.len() if // the buffer was not empty when the operation was started. read: usize, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn read_until<'a, R>( reader: &'a mut R, delimiter: u8, buf: &'a mut Vec, ) -> ReadUntil<'a, R> where R: AsyncBufRead + ?Sized + Unpin, { ReadUntil { reader, delimiter, buf, read: 0, _pin: PhantomPinned, } } pub(super) fn read_until_internal( mut reader: Pin<&mut R>, cx: &mut Context<'_>, delimiter: u8, buf: &mut Vec, read: &mut usize, ) -> Poll> { loop { let (done, used) = { let available = ready!(reader.as_mut().poll_fill_buf(cx))?; if let Some(i) = memchr::memchr(delimiter, available) { buf.extend_from_slice(&available[..=i]); (true, i + 1) } else { buf.extend_from_slice(available); (false, available.len()) } }; reader.as_mut().consume(used); *read += used; if done || used == 0 { return Poll::Ready(Ok(mem::replace(read, 0))); } } } impl Future for ReadUntil<'_, R> { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); read_until_internal(Pin::new(*me.reader), cx, *me.delimiter, me.buf, me.read) } } tokio-1.43.1/src/io/util/repeat.rs000064400000000000000000000035251046102023000150450ustar 00000000000000use bytes::BufMut; use crate::io::util::poll_proceed_and_make_progress; use crate::io::{AsyncRead, ReadBuf}; use std::io; use std::pin::Pin; use std::task::{ready, Context, Poll}; cfg_io_util! { /// An async reader which yields one byte over and over and over and over and /// over and... /// /// This struct is generally created by calling [`repeat`][repeat]. Please /// see the documentation of `repeat()` for more details. /// /// This is an asynchronous version of [`std::io::Repeat`][std]. /// /// [repeat]: fn@repeat /// [std]: std::io::Repeat #[derive(Debug)] pub struct Repeat { byte: u8, } /// Creates an instance of an async reader that infinitely repeats one byte. /// /// All reads from this reader will succeed by filling the specified buffer with /// the given byte. /// /// This is an asynchronous version of [`std::io::repeat`][std]. /// /// [std]: std::io::repeat /// /// # Examples /// /// ``` /// use tokio::io::{self, AsyncReadExt}; /// /// #[tokio::main] /// async fn main() { /// let mut buffer = [0; 3]; /// io::repeat(0b101).read_exact(&mut buffer).await.unwrap(); /// assert_eq!(buffer, [0b101, 0b101, 0b101]); /// } /// ``` pub fn repeat(byte: u8) -> Repeat { Repeat { byte } } } impl AsyncRead for Repeat { #[inline] fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); buf.put_bytes(self.byte, buf.remaining()); Poll::Ready(Ok(())) } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::(); } } tokio-1.43.1/src/io/util/shutdown.rs000064400000000000000000000022251046102023000154340ustar 00000000000000use crate::io::AsyncWrite; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { /// A future used to shutdown an I/O object. /// /// Created by the [`AsyncWriteExt::shutdown`][shutdown] function. /// [shutdown]: [`crate::io::AsyncWriteExt::shutdown`] #[must_use = "futures do nothing unless you `.await` or poll them"] #[derive(Debug)] pub struct Shutdown<'a, A: ?Sized> { a: &'a mut A, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } /// Creates a future which will shutdown an I/O object. pub(super) fn shutdown(a: &mut A) -> Shutdown<'_, A> where A: AsyncWrite + Unpin + ?Sized, { Shutdown { a, _pin: PhantomPinned, } } impl Future for Shutdown<'_, A> where A: AsyncWrite + Unpin + ?Sized, { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); Pin::new(me.a).poll_shutdown(cx) } } tokio-1.43.1/src/io/util/sink.rs000064400000000000000000000046741046102023000145370ustar 00000000000000use crate::io::util::poll_proceed_and_make_progress; use crate::io::AsyncWrite; use std::fmt; use std::io; use std::pin::Pin; use std::task::{ready, Context, Poll}; cfg_io_util! { /// An async writer which will move data into the void. /// /// This struct is generally created by calling [`sink`][sink]. Please /// see the documentation of `sink()` for more details. /// /// This is an asynchronous version of [`std::io::Sink`][std]. /// /// [sink]: sink() /// [std]: std::io::Sink pub struct Sink { _p: (), } /// Creates an instance of an async writer which will successfully consume all /// data. /// /// All calls to [`poll_write`] on the returned instance will return /// `Poll::Ready(Ok(buf.len()))` and the contents of the buffer will not be /// inspected. /// /// This is an asynchronous version of [`std::io::sink`][std]. /// /// [`poll_write`]: crate::io::AsyncWrite::poll_write() /// [std]: std::io::sink /// /// # Examples /// /// ``` /// use tokio::io::{self, AsyncWriteExt}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let buffer = vec![1, 2, 3, 5, 8]; /// let num_bytes = io::sink().write(&buffer).await?; /// assert_eq!(num_bytes, 5); /// Ok(()) /// } /// ``` pub fn sink() -> Sink { Sink { _p: () } } } impl AsyncWrite for Sink { #[inline] fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(buf.len())) } #[inline] fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } #[inline] fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } } impl fmt::Debug for Sink { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.pad("Sink { .. }") } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::(); } } tokio-1.43.1/src/io/util/split.rs000064400000000000000000000063431046102023000147210ustar 00000000000000use crate::io::util::read_until::read_until_internal; use crate::io::AsyncBufRead; use pin_project_lite::pin_project; use std::io; use std::mem; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// Splitter for the [`split`](crate::io::AsyncBufReadExt::split) method. /// /// A `Split` can be turned into a `Stream` with [`SplitStream`]. /// /// [`SplitStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.SplitStream.html #[derive(Debug)] #[must_use = "streams do nothing unless polled"] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct Split { #[pin] reader: R, buf: Vec, delim: u8, read: usize, } } pub(crate) fn split(reader: R, delim: u8) -> Split where R: AsyncBufRead, { Split { reader, buf: Vec::new(), delim, read: 0, } } impl Split where R: AsyncBufRead + Unpin, { /// Returns the next segment in the stream. /// /// # Examples /// /// ``` /// # use tokio::io::AsyncBufRead; /// use tokio::io::AsyncBufReadExt; /// /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> { /// let mut segments = my_buf_read.split(b'f'); /// /// while let Some(segment) = segments.next_segment().await? { /// println!("length = {}", segment.len()) /// } /// # Ok(()) /// # } /// ``` pub async fn next_segment(&mut self) -> io::Result>> { use std::future::poll_fn; poll_fn(|cx| Pin::new(&mut *self).poll_next_segment(cx)).await } } impl Split where R: AsyncBufRead, { /// Polls for the next segment in the stream. /// /// This method returns: /// /// * `Poll::Pending` if the next segment is not yet available. /// * `Poll::Ready(Ok(Some(segment)))` if the next segment is available. /// * `Poll::Ready(Ok(None))` if there are no more segments in this stream. /// * `Poll::Ready(Err(err))` if an IO error occurred while reading the /// next segment. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when more bytes become /// available on the underlying IO resource. /// /// Note that on multiple calls to `poll_next_segment`, only the `Waker` /// from the `Context` passed to the most recent call is scheduled to /// receive a wakeup. pub fn poll_next_segment( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll>>> { let me = self.project(); let n = ready!(read_until_internal( me.reader, cx, *me.delim, me.buf, me.read, ))?; // read_until_internal resets me.read to zero once it finds the delimiter debug_assert_eq!(*me.read, 0); if n == 0 && me.buf.is_empty() { return Poll::Ready(Ok(None)); } if me.buf.last() == Some(me.delim) { me.buf.pop(); } Poll::Ready(Ok(Some(mem::take(me.buf)))) } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/take.rs000064400000000000000000000076501046102023000145140ustar 00000000000000use crate::io::{AsyncBufRead, AsyncRead, ReadBuf}; use pin_project_lite::pin_project; use std::convert::TryFrom; use std::pin::Pin; use std::task::{ready, Context, Poll}; use std::{cmp, io}; pin_project! { /// Stream for the [`take`](super::AsyncReadExt::take) method. #[derive(Debug)] #[must_use = "streams do nothing unless you `.await` or poll them"] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct Take { #[pin] inner: R, // Add '_' to avoid conflicts with `limit` method. limit_: u64, } } pub(super) fn take(inner: R, limit: u64) -> Take { Take { inner, limit_: limit, } } impl Take { /// Returns the remaining number of bytes that can be /// read before this instance will return EOF. /// /// # Note /// /// This instance may reach `EOF` after reading fewer bytes than indicated by /// this method if the underlying [`AsyncRead`] instance reaches EOF. pub fn limit(&self) -> u64 { self.limit_ } /// Sets the number of bytes that can be read before this instance will /// return EOF. This is the same as constructing a new `Take` instance, so /// the amount of bytes read and the previous limit value don't matter when /// calling this method. pub fn set_limit(&mut self, limit: u64) { self.limit_ = limit; } /// Gets a reference to the underlying reader. pub fn get_ref(&self) -> &R { &self.inner } /// Gets a mutable reference to the underlying reader. /// /// Care should be taken to avoid modifying the internal I/O state of the /// underlying reader as doing so may corrupt the internal limit of this /// `Take`. pub fn get_mut(&mut self) -> &mut R { &mut self.inner } /// Gets a pinned mutable reference to the underlying reader. /// /// Care should be taken to avoid modifying the internal I/O state of the /// underlying reader as doing so may corrupt the internal limit of this /// `Take`. pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { self.project().inner } /// Consumes the `Take`, returning the wrapped reader. pub fn into_inner(self) -> R { self.inner } } impl AsyncRead for Take { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { if self.limit_ == 0 { return Poll::Ready(Ok(())); } let me = self.project(); let mut b = buf.take(usize::try_from(*me.limit_).unwrap_or(usize::MAX)); let buf_ptr = b.filled().as_ptr(); ready!(me.inner.poll_read(cx, &mut b))?; assert_eq!(b.filled().as_ptr(), buf_ptr); let n = b.filled().len(); // We need to update the original ReadBuf unsafe { buf.assume_init(n); } buf.advance(n); *me.limit_ -= n as u64; Poll::Ready(Ok(())) } } impl AsyncBufRead for Take { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); // Don't call into inner reader at all at EOF because it may still block if *me.limit_ == 0 { return Poll::Ready(Ok(&[])); } let buf = ready!(me.inner.poll_fill_buf(cx)?); let cap = cmp::min(buf.len() as u64, *me.limit_) as usize; Poll::Ready(Ok(&buf[..cap])) } fn consume(self: Pin<&mut Self>, amt: usize) { let me = self.project(); // Don't let callers reset the limit by passing an overlarge value let amt = cmp::min(amt as u64, *me.limit_) as usize; *me.limit_ -= amt as u64; me.inner.consume(amt); } } #[cfg(test)] mod tests { use super::*; #[test] fn assert_unpin() { crate::is_unpin::>(); } } tokio-1.43.1/src/io/util/vec_with_initialized.rs000064400000000000000000000115631046102023000177630ustar 00000000000000use crate::io::ReadBuf; use std::mem::MaybeUninit; /// Something that looks like a `Vec`. /// /// # Safety /// /// The implementor must guarantee that the vector returned by the /// `as_mut` and `as_mut` methods do not change from one call to /// another. pub(crate) unsafe trait VecU8: AsRef> + AsMut> {} unsafe impl VecU8 for Vec {} unsafe impl VecU8 for &mut Vec {} /// This struct wraps a `Vec` or `&mut Vec`, combining it with a /// `num_initialized`, which keeps track of the number of initialized bytes /// in the unused capacity. /// /// The purpose of this struct is to remember how many bytes were initialized /// through a `ReadBuf` from call to call. /// /// This struct has the safety invariant that the first `num_initialized` of the /// vector's allocation must be initialized at any time. #[derive(Debug)] pub(crate) struct VecWithInitialized { vec: V, // The number of initialized bytes in the vector. // Always between `vec.len()` and `vec.capacity()`. num_initialized: usize, starting_capacity: usize, } impl VecWithInitialized> { #[cfg(feature = "io-util")] pub(crate) fn take(&mut self) -> Vec { self.num_initialized = 0; std::mem::take(&mut self.vec) } } impl VecWithInitialized where V: VecU8, { pub(crate) fn new(mut vec: V) -> Self { // SAFETY: The safety invariants of vector guarantee that the bytes up // to its length are initialized. Self { num_initialized: vec.as_mut().len(), starting_capacity: vec.as_ref().capacity(), vec, } } pub(crate) fn reserve(&mut self, num_bytes: usize) { let vec = self.vec.as_mut(); if vec.capacity() - vec.len() >= num_bytes { return; } // SAFETY: Setting num_initialized to `vec.len()` is correct as // `reserve` does not change the length of the vector. self.num_initialized = vec.len(); vec.reserve(num_bytes); } #[cfg(feature = "io-util")] pub(crate) fn is_empty(&self) -> bool { self.vec.as_ref().is_empty() } pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> { let num_initialized = self.num_initialized; // SAFETY: Creating the slice is safe because of the safety invariants // on Vec. The safety invariants of `ReadBuf` will further guarantee // that no bytes in the slice are de-initialized. let vec = self.vec.as_mut(); let len = vec.len(); let cap = vec.capacity(); let ptr = vec.as_mut_ptr().cast::>(); let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit>(ptr, cap) }; // SAFETY: This is safe because the safety invariants of // VecWithInitialized say that the first num_initialized bytes must be // initialized. let mut read_buf = ReadBuf::uninit(slice); unsafe { read_buf.assume_init(num_initialized); } read_buf.set_filled(len); read_buf } pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) { let vec = self.vec.as_mut(); assert_eq!(vec.as_ptr(), parts.ptr); // SAFETY: // The ReadBufParts really does point inside `self.vec` due to the above // check, and the safety invariants of `ReadBuf` guarantee that the // first `parts.initialized` bytes of `self.vec` really have been // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is // at most `parts.initialized`, so the first `parts.len` bytes are also // initialized. // // Note that this relies on the fact that `V` is either `Vec` or // `&mut Vec`, so the vector returned by `self.vec.as_mut()` cannot // change from call to call. unsafe { self.num_initialized = parts.initialized; vec.set_len(parts.len); } } // Returns a boolean telling the caller to try reading into a small local buffer first if true. // Doing so would avoid overallocating when vec is filled to capacity and we reached EOF. pub(crate) fn try_small_read_first(&self, num_bytes: usize) -> bool { let vec = self.vec.as_ref(); vec.capacity() - vec.len() < num_bytes && self.starting_capacity == vec.capacity() && self.starting_capacity >= num_bytes } } pub(crate) struct ReadBufParts { // Pointer is only used to check that the ReadBuf actually came from the // right VecWithInitialized. ptr: *const u8, len: usize, initialized: usize, } // This is needed to release the borrow on `VecWithInitialized`. pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts { ReadBufParts { ptr: rb.filled().as_ptr(), len: rb.filled().len(), initialized: rb.initialized().len(), } } tokio-1.43.1/src/io/util/write.rs000064400000000000000000000022511046102023000147120ustar 00000000000000use crate::io::AsyncWrite; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { /// A future to write some of the buffer to an `AsyncWrite`. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Write<'a, W: ?Sized> { writer: &'a mut W, buf: &'a [u8], // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } /// Tries to write some bytes from the given `buf` to the writer in an /// asynchronous manner, returning a future. pub(crate) fn write<'a, W>(writer: &'a mut W, buf: &'a [u8]) -> Write<'a, W> where W: AsyncWrite + Unpin + ?Sized, { Write { writer, buf, _pin: PhantomPinned, } } impl Future for Write<'_, W> where W: AsyncWrite + Unpin + ?Sized, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); Pin::new(&mut *me.writer).poll_write(cx, me.buf) } } tokio-1.43.1/src/io/util/write_all.rs000064400000000000000000000025441046102023000155470ustar 00000000000000use crate::io::AsyncWrite; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::mem; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct WriteAll<'a, W: ?Sized> { writer: &'a mut W, buf: &'a [u8], // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn write_all<'a, W>(writer: &'a mut W, buf: &'a [u8]) -> WriteAll<'a, W> where W: AsyncWrite + Unpin + ?Sized, { WriteAll { writer, buf, _pin: PhantomPinned, } } impl Future for WriteAll<'_, W> where W: AsyncWrite + Unpin + ?Sized, { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); while !me.buf.is_empty() { let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf))?; { let (_, rest) = mem::take(&mut *me.buf).split_at(n); *me.buf = rest; } if n == 0 { return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); } } Poll::Ready(Ok(())) } } tokio-1.43.1/src/io/util/write_all_buf.rs000064400000000000000000000034171046102023000164030ustar 00000000000000use crate::io::AsyncWrite; use bytes::Buf; use pin_project_lite::pin_project; use std::future::Future; use std::io::{self, IoSlice}; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// A future to write some of the buffer to an `AsyncWrite`. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct WriteAllBuf<'a, W, B> { writer: &'a mut W, buf: &'a mut B, #[pin] _pin: PhantomPinned, } } /// Tries to write some bytes from the given `buf` to the writer in an /// asynchronous manner, returning a future. pub(crate) fn write_all_buf<'a, W, B>(writer: &'a mut W, buf: &'a mut B) -> WriteAllBuf<'a, W, B> where W: AsyncWrite + Unpin, B: Buf, { WriteAllBuf { writer, buf, _pin: PhantomPinned, } } impl Future for WriteAllBuf<'_, W, B> where W: AsyncWrite + Unpin, B: Buf, { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { const MAX_VECTOR_ELEMENTS: usize = 64; let me = self.project(); while me.buf.has_remaining() { let n = if me.writer.is_write_vectored() { let mut slices = [IoSlice::new(&[]); MAX_VECTOR_ELEMENTS]; let cnt = me.buf.chunks_vectored(&mut slices); ready!(Pin::new(&mut *me.writer).poll_write_vectored(cx, &slices[..cnt]))? } else { ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf.chunk())?) }; me.buf.advance(n); if n == 0 { return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); } } Poll::Ready(Ok(())) } } tokio-1.43.1/src/io/util/write_buf.rs000064400000000000000000000024601046102023000155500ustar 00000000000000use crate::io::AsyncWrite; use bytes::Buf; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{ready, Context, Poll}; pin_project! { /// A future to write some of the buffer to an `AsyncWrite`. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct WriteBuf<'a, W, B> { writer: &'a mut W, buf: &'a mut B, #[pin] _pin: PhantomPinned, } } /// Tries to write some bytes from the given `buf` to the writer in an /// asynchronous manner, returning a future. pub(crate) fn write_buf<'a, W, B>(writer: &'a mut W, buf: &'a mut B) -> WriteBuf<'a, W, B> where W: AsyncWrite + Unpin, B: Buf, { WriteBuf { writer, buf, _pin: PhantomPinned, } } impl Future for WriteBuf<'_, W, B> where W: AsyncWrite + Unpin, B: Buf, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); if !me.buf.has_remaining() { return Poll::Ready(Ok(0)); } let n = ready!(Pin::new(me.writer).poll_write(cx, me.buf.chunk()))?; me.buf.advance(n); Poll::Ready(Ok(n)) } } tokio-1.43.1/src/io/util/write_int.rs000064400000000000000000000107541046102023000155730ustar 00000000000000use crate::io::AsyncWrite; use bytes::BufMut; use pin_project_lite::pin_project; use std::future::Future; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; macro_rules! writer { ($name:ident, $ty:ty, $writer:ident) => { writer!($name, $ty, $writer, std::mem::size_of::<$ty>()); }; ($name:ident, $ty:ty, $writer:ident, $bytes:expr) => { pin_project! { #[doc(hidden)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct $name { #[pin] dst: W, buf: [u8; $bytes], written: u8, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } impl $name { pub(crate) fn new(w: W, value: $ty) -> Self { let mut writer = Self { buf: [0; $bytes], written: 0, dst: w, _pin: PhantomPinned, }; BufMut::$writer(&mut &mut writer.buf[..], value); writer } } impl Future for $name where W: AsyncWrite, { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut me = self.project(); if *me.written == $bytes as u8 { return Poll::Ready(Ok(())); } while *me.written < $bytes as u8 { *me.written += match me .dst .as_mut() .poll_write(cx, &me.buf[*me.written as usize..]) { Poll::Pending => return Poll::Pending, Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), Poll::Ready(Ok(0)) => { return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); } Poll::Ready(Ok(n)) => n as u8, }; } Poll::Ready(Ok(())) } } }; } macro_rules! writer8 { ($name:ident, $ty:ty) => { pin_project! { #[doc(hidden)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct $name { #[pin] dst: W, byte: $ty, // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } impl $name { pub(crate) fn new(dst: W, byte: $ty) -> Self { Self { dst, byte, _pin: PhantomPinned, } } } impl Future for $name where W: AsyncWrite, { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); let buf = [*me.byte as u8]; match me.dst.poll_write(cx, &buf[..]) { Poll::Pending => Poll::Pending, Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())), Poll::Ready(Ok(0)) => Poll::Ready(Err(io::ErrorKind::WriteZero.into())), Poll::Ready(Ok(1)) => Poll::Ready(Ok(())), Poll::Ready(Ok(_)) => unreachable!(), } } } }; } writer8!(WriteU8, u8); writer8!(WriteI8, i8); writer!(WriteU16, u16, put_u16); writer!(WriteU32, u32, put_u32); writer!(WriteU64, u64, put_u64); writer!(WriteU128, u128, put_u128); writer!(WriteI16, i16, put_i16); writer!(WriteI32, i32, put_i32); writer!(WriteI64, i64, put_i64); writer!(WriteI128, i128, put_i128); writer!(WriteF32, f32, put_f32); writer!(WriteF64, f64, put_f64); writer!(WriteU16Le, u16, put_u16_le); writer!(WriteU32Le, u32, put_u32_le); writer!(WriteU64Le, u64, put_u64_le); writer!(WriteU128Le, u128, put_u128_le); writer!(WriteI16Le, i16, put_i16_le); writer!(WriteI32Le, i32, put_i32_le); writer!(WriteI64Le, i64, put_i64_le); writer!(WriteI128Le, i128, put_i128_le); writer!(WriteF32Le, f32, put_f32_le); writer!(WriteF64Le, f64, put_f64_le); tokio-1.43.1/src/io/util/write_vectored.rs000064400000000000000000000022471046102023000166120ustar 00000000000000use crate::io::AsyncWrite; use pin_project_lite::pin_project; use std::io; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; use std::{future::Future, io::IoSlice}; pin_project! { /// A future to write a slice of buffers to an `AsyncWrite`. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct WriteVectored<'a, 'b, W: ?Sized> { writer: &'a mut W, bufs: &'a [IoSlice<'b>], // Make this future `!Unpin` for compatibility with async trait methods. #[pin] _pin: PhantomPinned, } } pub(crate) fn write_vectored<'a, 'b, W>( writer: &'a mut W, bufs: &'a [IoSlice<'b>], ) -> WriteVectored<'a, 'b, W> where W: AsyncWrite + Unpin + ?Sized, { WriteVectored { writer, bufs, _pin: PhantomPinned, } } impl Future for WriteVectored<'_, '_, W> where W: AsyncWrite + Unpin + ?Sized, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let me = self.project(); Pin::new(&mut *me.writer).poll_write_vectored(cx, me.bufs) } } tokio-1.43.1/src/lib.rs000064400000000000000000000577011046102023000127540ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![allow( clippy::cognitive_complexity, clippy::large_enum_variant, clippy::module_inception, clippy::needless_doctest_main )] #![warn( missing_debug_implementations, missing_docs, rust_2018_idioms, unreachable_pub )] #![deny(unused_must_use)] #![doc(test( no_crate_inject, attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables)) ))] #![cfg_attr(docsrs, feature(doc_cfg))] #![cfg_attr(docsrs, allow(unused_attributes))] #![cfg_attr(loom, allow(dead_code, unreachable_pub))] #![cfg_attr(windows, allow(rustdoc::broken_intra_doc_links))] //! A runtime for writing reliable network applications without compromising speed. //! //! Tokio is an event-driven, non-blocking I/O platform for writing asynchronous //! applications with the Rust programming language. At a high level, it //! provides a few major components: //! //! * Tools for [working with asynchronous tasks][tasks], including //! [synchronization primitives and channels][sync] and [timeouts, sleeps, and //! intervals][time]. //! * APIs for [performing asynchronous I/O][io], including [TCP and UDP][net] sockets, //! [filesystem][fs] operations, and [process] and [signal] management. //! * A [runtime] for executing asynchronous code, including a task scheduler, //! an I/O driver backed by the operating system's event queue (`epoll`, `kqueue`, //! `IOCP`, etc...), and a high performance timer. //! //! Guide level documentation is found on the [website]. //! //! [tasks]: #working-with-tasks //! [sync]: crate::sync //! [time]: crate::time //! [io]: #asynchronous-io //! [net]: crate::net //! [fs]: crate::fs //! [process]: crate::process //! [signal]: crate::signal //! [fs]: crate::fs //! [runtime]: crate::runtime //! [website]: https://tokio.rs/tokio/tutorial //! //! # A Tour of Tokio //! //! Tokio consists of a number of modules that provide a range of functionality //! essential for implementing asynchronous applications in Rust. In this //! section, we will take a brief tour of Tokio, summarizing the major APIs and //! their uses. //! //! The easiest way to get started is to enable all features. Do this by //! enabling the `full` feature flag: //! //! ```toml //! tokio = { version = "1", features = ["full"] } //! ``` //! //! ### Authoring applications //! //! Tokio is great for writing applications and most users in this case shouldn't //! worry too much about what features they should pick. If you're unsure, we suggest //! going with `full` to ensure that you don't run into any road blocks while you're //! building your application. //! //! #### Example //! //! This example shows the quickest way to get started with Tokio. //! //! ```toml //! tokio = { version = "1", features = ["full"] } //! ``` //! //! ### Authoring libraries //! //! As a library author your goal should be to provide the lightest weight crate //! that is based on Tokio. To achieve this you should ensure that you only enable //! the features you need. This allows users to pick up your crate without having //! to enable unnecessary features. //! //! #### Example //! //! This example shows how you may want to import features for a library that just //! needs to `tokio::spawn` and use a `TcpStream`. //! //! ```toml //! tokio = { version = "1", features = ["rt", "net"] } //! ``` //! //! ## Working With Tasks //! //! Asynchronous programs in Rust are based around lightweight, non-blocking //! units of execution called [_tasks_][tasks]. The [`tokio::task`] module provides //! important tools for working with tasks: //! //! * The [`spawn`] function and [`JoinHandle`] type, for scheduling a new task //! on the Tokio runtime and awaiting the output of a spawned task, respectively, //! * Functions for [running blocking operations][blocking] in an asynchronous //! task context. //! //! The [`tokio::task`] module is present only when the "rt" feature flag //! is enabled. //! //! [tasks]: task/index.html#what-are-tasks //! [`tokio::task`]: crate::task //! [`spawn`]: crate::task::spawn() //! [`JoinHandle`]: crate::task::JoinHandle //! [blocking]: task/index.html#blocking-and-yielding //! //! The [`tokio::sync`] module contains synchronization primitives to use when //! needing to communicate or share data. These include: //! //! * channels ([`oneshot`], [`mpsc`], [`watch`], and [`broadcast`]), for sending values //! between tasks, //! * a non-blocking [`Mutex`], for controlling access to a shared, mutable //! value, //! * an asynchronous [`Barrier`] type, for multiple tasks to synchronize before //! beginning a computation. //! //! The `tokio::sync` module is present only when the "sync" feature flag is //! enabled. //! //! [`tokio::sync`]: crate::sync //! [`Mutex`]: crate::sync::Mutex //! [`Barrier`]: crate::sync::Barrier //! [`oneshot`]: crate::sync::oneshot //! [`mpsc`]: crate::sync::mpsc //! [`watch`]: crate::sync::watch //! [`broadcast`]: crate::sync::broadcast //! //! The [`tokio::time`] module provides utilities for tracking time and //! scheduling work. This includes functions for setting [timeouts][timeout] for //! tasks, [sleeping][sleep] work to run in the future, or [repeating an operation at an //! interval][interval]. //! //! In order to use `tokio::time`, the "time" feature flag must be enabled. //! //! [`tokio::time`]: crate::time //! [sleep]: crate::time::sleep() //! [interval]: crate::time::interval() //! [timeout]: crate::time::timeout() //! //! Finally, Tokio provides a _runtime_ for executing asynchronous tasks. Most //! applications can use the [`#[tokio::main]`][main] macro to run their code on the //! Tokio runtime. However, this macro provides only basic configuration options. As //! an alternative, the [`tokio::runtime`] module provides more powerful APIs for configuring //! and managing runtimes. You should use that module if the `#[tokio::main]` macro doesn't //! provide the functionality you need. //! //! Using the runtime requires the "rt" or "rt-multi-thread" feature flags, to //! enable the current-thread [single-threaded scheduler][rt] and the [multi-thread //! scheduler][rt-multi-thread], respectively. See the [`runtime` module //! documentation][rt-features] for details. In addition, the "macros" feature //! flag enables the `#[tokio::main]` and `#[tokio::test]` attributes. //! //! [main]: attr.main.html //! [`tokio::runtime`]: crate::runtime //! [`Builder`]: crate::runtime::Builder //! [`Runtime`]: crate::runtime::Runtime //! [rt]: runtime/index.html#current-thread-scheduler //! [rt-multi-thread]: runtime/index.html#multi-thread-scheduler //! [rt-features]: runtime/index.html#runtime-scheduler //! //! ## CPU-bound tasks and blocking code //! //! Tokio is able to concurrently run many tasks on a few threads by repeatedly //! swapping the currently running task on each thread. However, this kind of //! swapping can only happen at `.await` points, so code that spends a long time //! without reaching an `.await` will prevent other tasks from running. To //! combat this, Tokio provides two kinds of threads: Core threads and blocking threads. //! //! The core threads are where all asynchronous code runs, and Tokio will by default //! spawn one for each CPU core. You can use the environment variable `TOKIO_WORKER_THREADS` //! to override the default value. //! //! The blocking threads are spawned on demand, can be used to run blocking code //! that would otherwise block other tasks from running and are kept alive when //! not used for a certain amount of time which can be configured with [`thread_keep_alive`]. //! Since it is not possible for Tokio to swap out blocking tasks, like it //! can do with asynchronous code, the upper limit on the number of blocking //! threads is very large. These limits can be configured on the [`Builder`]. //! //! To spawn a blocking task, you should use the [`spawn_blocking`] function. //! //! [`Builder`]: crate::runtime::Builder //! [`spawn_blocking`]: crate::task::spawn_blocking() //! [`thread_keep_alive`]: crate::runtime::Builder::thread_keep_alive() //! //! ``` //! #[tokio::main] //! async fn main() { //! // This is running on a core thread. //! //! let blocking_task = tokio::task::spawn_blocking(|| { //! // This is running on a blocking thread. //! // Blocking here is ok. //! }); //! //! // We can wait for the blocking task like this: //! // If the blocking task panics, the unwrap below will propagate the //! // panic. //! blocking_task.await.unwrap(); //! } //! ``` //! //! If your code is CPU-bound and you wish to limit the number of threads used //! to run it, you should use a separate thread pool dedicated to CPU bound tasks. //! For example, you could consider using the [rayon] library for CPU-bound //! tasks. It is also possible to create an extra Tokio runtime dedicated to //! CPU-bound tasks, but if you do this, you should be careful that the extra //! runtime runs _only_ CPU-bound tasks, as IO-bound tasks on that runtime //! will behave poorly. //! //! Hint: If using rayon, you can use a [`oneshot`] channel to send the result back //! to Tokio when the rayon task finishes. //! //! [rayon]: https://docs.rs/rayon //! [`oneshot`]: crate::sync::oneshot //! //! ## Asynchronous IO //! //! As well as scheduling and running tasks, Tokio provides everything you need //! to perform input and output asynchronously. //! //! The [`tokio::io`] module provides Tokio's asynchronous core I/O primitives, //! the [`AsyncRead`], [`AsyncWrite`], and [`AsyncBufRead`] traits. In addition, //! when the "io-util" feature flag is enabled, it also provides combinators and //! functions for working with these traits, forming as an asynchronous //! counterpart to [`std::io`]. //! //! Tokio also includes APIs for performing various kinds of I/O and interacting //! with the operating system asynchronously. These include: //! //! * [`tokio::net`], which contains non-blocking versions of [TCP], [UDP], and //! [Unix Domain Sockets][UDS] (enabled by the "net" feature flag), //! * [`tokio::fs`], similar to [`std::fs`] but for performing filesystem I/O //! asynchronously (enabled by the "fs" feature flag), //! * [`tokio::signal`], for asynchronously handling Unix and Windows OS signals //! (enabled by the "signal" feature flag), //! * [`tokio::process`], for spawning and managing child processes (enabled by //! the "process" feature flag). //! //! [`tokio::io`]: crate::io //! [`AsyncRead`]: crate::io::AsyncRead //! [`AsyncWrite`]: crate::io::AsyncWrite //! [`AsyncBufRead`]: crate::io::AsyncBufRead //! [`std::io`]: std::io //! [`tokio::net`]: crate::net //! [TCP]: crate::net::tcp //! [UDP]: crate::net::UdpSocket //! [UDS]: crate::net::unix //! [`tokio::fs`]: crate::fs //! [`std::fs`]: std::fs //! [`tokio::signal`]: crate::signal //! [`tokio::process`]: crate::process //! //! # Examples //! //! A simple TCP echo server: //! //! ```no_run //! use tokio::net::TcpListener; //! use tokio::io::{AsyncReadExt, AsyncWriteExt}; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! let listener = TcpListener::bind("127.0.0.1:8080").await?; //! //! loop { //! let (mut socket, _) = listener.accept().await?; //! //! tokio::spawn(async move { //! let mut buf = [0; 1024]; //! //! // In a loop, read data from the socket and write the data back. //! loop { //! let n = match socket.read(&mut buf).await { //! // socket closed //! Ok(0) => return, //! Ok(n) => n, //! Err(e) => { //! eprintln!("failed to read from socket; err = {:?}", e); //! return; //! } //! }; //! //! // Write the data back //! if let Err(e) = socket.write_all(&buf[0..n]).await { //! eprintln!("failed to write to socket; err = {:?}", e); //! return; //! } //! } //! }); //! } //! } //! ``` //! //! ## Feature flags //! //! Tokio uses a set of [feature flags] to reduce the amount of compiled code. It //! is possible to just enable certain features over others. By default, Tokio //! does not enable any features but allows one to enable a subset for their use //! case. Below is a list of the available feature flags. You may also notice //! above each function, struct and trait there is listed one or more feature flags //! that are required for that item to be used. If you are new to Tokio it is //! recommended that you use the `full` feature flag which will enable all public APIs. //! Beware though that this will pull in many extra dependencies that you may not //! need. //! //! - `full`: Enables all features listed below except `test-util` and `tracing`. //! - `rt`: Enables `tokio::spawn`, the current-thread scheduler, //! and non-scheduler utilities. //! - `rt-multi-thread`: Enables the heavier, multi-threaded, work-stealing scheduler. //! - `io-util`: Enables the IO based `Ext` traits. //! - `io-std`: Enable `Stdout`, `Stdin` and `Stderr` types. //! - `net`: Enables `tokio::net` types such as `TcpStream`, `UnixStream` and //! `UdpSocket`, as well as (on Unix-like systems) `AsyncFd` and (on //! FreeBSD) `PollAio`. //! - `time`: Enables `tokio::time` types and allows the schedulers to enable //! the built in timer. //! - `process`: Enables `tokio::process` types. //! - `macros`: Enables `#[tokio::main]` and `#[tokio::test]` macros. //! - `sync`: Enables all `tokio::sync` types. //! - `signal`: Enables all `tokio::signal` types. //! - `fs`: Enables `tokio::fs` types. //! - `test-util`: Enables testing based infrastructure for the Tokio runtime. //! - `parking_lot`: As a potential optimization, use the `_parking_lot_` crate's //! synchronization primitives internally. Also, this //! dependency is necessary to construct some of our primitives //! in a `const` context. `MSRV` may increase according to the //! `_parking_lot_` release in use. //! //! _Note: `AsyncRead` and `AsyncWrite` traits do not require any features and are //! always available._ //! //! ### Unstable features //! //! Some feature flags are only available when specifying the `tokio_unstable` flag: //! //! - `tracing`: Enables tracing events. //! //! Likewise, some parts of the API are only available with the same flag: //! //! - [`task::Builder`] //! - Some methods on [`task::JoinSet`] //! - [`runtime::RuntimeMetrics`] //! - [`runtime::Builder::on_task_spawn`] //! - [`runtime::Builder::on_task_terminate`] //! - [`runtime::Builder::unhandled_panic`] //! - [`runtime::TaskMeta`] //! //! This flag enables **unstable** features. The public API of these features //! may break in 1.x releases. To enable these features, the `--cfg //! tokio_unstable` argument must be passed to `rustc` when compiling. This //! serves to explicitly opt-in to features which may break semver conventions, //! since Cargo [does not yet directly support such opt-ins][unstable features]. //! //! You can specify it in your project's `.cargo/config.toml` file: //! //! ```toml //! [build] //! rustflags = ["--cfg", "tokio_unstable"] //! ``` //! //!

//! The [build] section does not go in a //! Cargo.toml file. Instead it must be placed in the Cargo config //! file .cargo/config.toml. //!
//! //! Alternatively, you can specify it with an environment variable: //! //! ```sh //! ## Many *nix shells: //! export RUSTFLAGS="--cfg tokio_unstable" //! cargo build //! ``` //! //! ```powershell //! ## Windows PowerShell: //! $Env:RUSTFLAGS="--cfg tokio_unstable" //! cargo build //! ``` //! //! [unstable features]: https://internals.rust-lang.org/t/feature-request-unstable-opt-in-non-transitive-crate-features/16193#why-not-a-crate-feature-2 //! [feature flags]: https://doc.rust-lang.org/cargo/reference/manifest.html#the-features-section //! //! ## Supported platforms //! //! Tokio currently guarantees support for the following platforms: //! //! * Linux //! * Windows //! * Android (API level 21) //! * macOS //! * iOS //! * FreeBSD //! //! Tokio will continue to support these platforms in the future. However, //! future releases may change requirements such as the minimum required libc //! version on Linux, the API level on Android, or the supported FreeBSD //! release. //! //! Beyond the above platforms, Tokio is intended to work on all platforms //! supported by the mio crate. You can find a longer list [in mio's //! documentation][mio-supported]. However, these additional platforms may //! become unsupported in the future. //! //! Note that Wine is considered to be a different platform from Windows. See //! mio's documentation for more information on Wine support. //! //! [mio-supported]: https://crates.io/crates/mio#platforms //! //! ### `WASM` support //! //! Tokio has some limited support for the `WASM` platform. Without the //! `tokio_unstable` flag, the following features are supported: //! //! * `sync` //! * `macros` //! * `io-util` //! * `rt` //! * `time` //! //! Enabling any other feature (including `full`) will cause a compilation //! failure. //! //! The `time` module will only work on `WASM` platforms that have support for //! timers (e.g. wasm32-wasi). The timing functions will panic if used on a `WASM` //! platform that does not support timers. //! //! Note also that if the runtime becomes indefinitely idle, it will panic //! immediately instead of blocking forever. On platforms that don't support //! time, this means that the runtime can never be idle in any way. //! //! ### Unstable `WASM` support //! //! Tokio also has unstable support for some additional `WASM` features. This //! requires the use of the `tokio_unstable` flag. //! //! Using this flag enables the use of `tokio::net` on the wasm32-wasi target. //! However, not all methods are available on the networking types as `WASI` //! currently does not support the creation of new sockets from within `WASM`. //! Because of this, sockets must currently be created via the `FromRawFd` //! trait. // Test that pointer width is compatible. This asserts that e.g. usize is at // least 32 bits, which a lot of components in Tokio currently assumes. // // TODO: improve once we have MSRV access to const eval to make more flexible. #[cfg(not(any(target_pointer_width = "32", target_pointer_width = "64")))] compile_error! { "Tokio requires the platform pointer width to be at least 32 bits" } #[cfg(all( not(tokio_unstable), target_family = "wasm", any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt-multi-thread", feature = "signal" ) ))] compile_error!("Only features sync,macros,io-util,rt,time are supported on wasm."); #[cfg(all(not(tokio_unstable), tokio_taskdump))] compile_error!("The `tokio_taskdump` feature requires `--cfg tokio_unstable`."); #[cfg(all( tokio_taskdump, not(doc), not(all( target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") )) ))] compile_error!( "The `tokio_taskdump` feature is only currently supported on \ linux, on `aarch64`, `x86` and `x86_64`." ); // Includes re-exports used by macros. // // This module is not intended to be part of the public API. In general, any // `doc(hidden)` code is not part of Tokio's public and stable API. #[macro_use] #[doc(hidden)] pub mod macros; cfg_fs! { pub mod fs; } mod future; pub mod io; pub mod net; mod loom; cfg_process! { pub mod process; } #[cfg(any( feature = "fs", feature = "io-std", feature = "net", all(windows, feature = "process"), ))] mod blocking; cfg_rt! { pub mod runtime; } cfg_not_rt! { pub(crate) mod runtime; } cfg_signal! { pub mod signal; } cfg_signal_internal! { #[cfg(not(feature = "signal"))] #[allow(dead_code)] #[allow(unreachable_pub)] pub(crate) mod signal; } cfg_sync! { pub mod sync; } cfg_not_sync! { mod sync; } pub mod task; cfg_rt! { pub use task::spawn; } cfg_time! { pub mod time; } mod trace { use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; cfg_taskdump! { pub(crate) use crate::runtime::task::trace::trace_leaf; } cfg_not_taskdump! { #[inline(always)] #[allow(dead_code)] pub(crate) fn trace_leaf(_: &mut std::task::Context<'_>) -> std::task::Poll<()> { std::task::Poll::Ready(()) } } #[cfg_attr(not(feature = "sync"), allow(dead_code))] pub(crate) fn async_trace_leaf() -> impl Future { struct Trace; impl Future for Trace { type Output = (); #[inline(always)] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { trace_leaf(cx) } } Trace } } mod util; /// Due to the `Stream` trait's inclusion in `std` landing later than Tokio's 1.0 /// release, most of the Tokio stream utilities have been moved into the [`tokio-stream`] /// crate. /// /// # Why was `Stream` not included in Tokio 1.0? /// /// Originally, we had planned to ship Tokio 1.0 with a stable `Stream` type /// but unfortunately the [RFC] had not been merged in time for `Stream` to /// reach `std` on a stable compiler in time for the 1.0 release of Tokio. For /// this reason, the team has decided to move all `Stream` based utilities to /// the [`tokio-stream`] crate. While this is not ideal, once `Stream` has made /// it into the standard library and the `MSRV` period has passed, we will implement /// stream for our different types. /// /// While this may seem unfortunate, not all is lost as you can get much of the /// `Stream` support with `async/await` and `while let` loops. It is also possible /// to create a `impl Stream` from `async fn` using the [`async-stream`] crate. /// /// [`tokio-stream`]: https://docs.rs/tokio-stream /// [`async-stream`]: https://docs.rs/async-stream /// [RFC]: https://github.com/rust-lang/rfcs/pull/2996 /// /// # Example /// /// Convert a [`sync::mpsc::Receiver`] to an `impl Stream`. /// /// ```rust,no_run /// use tokio::sync::mpsc; /// /// let (tx, mut rx) = mpsc::channel::(16); /// /// let stream = async_stream::stream! { /// while let Some(item) = rx.recv().await { /// yield item; /// } /// }; /// ``` pub mod stream {} // local re-exports of platform specific things, allowing for decent // documentation to be shimmed in on docs.rs #[cfg(all(docsrs, unix))] pub mod doc; #[cfg(any(feature = "net", feature = "fs"))] #[cfg(all(docsrs, unix))] #[allow(unused)] pub(crate) use self::doc::os; #[cfg(not(all(docsrs, unix)))] #[allow(unused)] pub(crate) use std::os; cfg_macros! { /// 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. #[doc(hidden)] pub use tokio_macros::select_priv_declare_output_enum; /// 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. #[doc(hidden)] pub use tokio_macros::select_priv_clean_pattern; cfg_rt! { #[cfg(feature = "rt-multi-thread")] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] #[doc(inline)] pub use tokio_macros::main; #[cfg(feature = "rt-multi-thread")] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] #[doc(inline)] pub use tokio_macros::test; cfg_not_rt_multi_thread! { #[doc(inline)] pub use tokio_macros::main_rt as main; #[doc(inline)] pub use tokio_macros::test_rt as test; } } // Always fail if rt is not enabled. cfg_not_rt! { #[doc(inline)] pub use tokio_macros::main_fail as main; #[doc(inline)] pub use tokio_macros::test_fail as test; } } // TODO: rm #[cfg(feature = "io-util")] #[cfg(test)] fn is_unpin() {} /// fuzz test (`fuzz_linked_list`) #[cfg(fuzzing)] pub mod fuzz; tokio-1.43.1/src/loom/mocked.rs000064400000000000000000000037751046102023000144200ustar 00000000000000pub(crate) use loom::*; pub(crate) mod sync { pub(crate) use loom::sync::{MutexGuard, RwLockReadGuard, RwLockWriteGuard}; #[derive(Debug)] pub(crate) struct Mutex(loom::sync::Mutex); #[allow(dead_code)] impl Mutex { #[inline] pub(crate) fn new(t: T) -> Mutex { Mutex(loom::sync::Mutex::new(t)) } #[inline] #[track_caller] pub(crate) fn lock(&self) -> MutexGuard<'_, T> { self.0.lock().unwrap() } #[inline] pub(crate) fn try_lock(&self) -> Option> { self.0.try_lock().ok() } #[inline] pub(crate) fn get_mut(&mut self) -> &mut T { self.0.get_mut().unwrap() } } #[derive(Debug)] pub(crate) struct RwLock(loom::sync::RwLock); #[allow(dead_code)] impl RwLock { #[inline] pub(crate) fn new(t: T) -> Self { Self(loom::sync::RwLock::new(t)) } #[inline] pub(crate) fn read(&self) -> RwLockReadGuard<'_, T> { self.0.read().unwrap() } #[inline] pub(crate) fn try_read(&self) -> Option> { self.0.try_read().ok() } #[inline] pub(crate) fn write(&self) -> RwLockWriteGuard<'_, T> { self.0.write().unwrap() } #[inline] pub(crate) fn try_write(&self) -> Option> { self.0.try_write().ok() } } pub(crate) use loom::sync::*; pub(crate) mod atomic { pub(crate) use loom::sync::atomic::*; // TODO: implement a loom version pub(crate) type StaticAtomicU64 = std::sync::atomic::AtomicU64; } } pub(crate) mod rand { pub(crate) fn seed() -> u64 { 1 } } pub(crate) mod sys { pub(crate) fn num_cpus() -> usize { 2 } } pub(crate) mod thread { pub use loom::lazy_static::AccessError; pub use loom::thread::*; } tokio-1.43.1/src/loom/mod.rs000064400000000000000000000004751046102023000137270ustar 00000000000000//! This module abstracts over `loom` and `std::sync` depending on whether we //! are running tests or not. #![allow(unused)] #[cfg(not(all(test, loom)))] mod std; #[cfg(not(all(test, loom)))] pub(crate) use self::std::*; #[cfg(all(test, loom))] mod mocked; #[cfg(all(test, loom))] pub(crate) use self::mocked::*; tokio-1.43.1/src/loom/std/atomic_u16.rs000064400000000000000000000024411046102023000157040ustar 00000000000000use std::cell::UnsafeCell; use std::fmt; use std::ops::Deref; use std::panic; /// `AtomicU16` providing an additional `unsync_load` function. pub(crate) struct AtomicU16 { inner: UnsafeCell, } unsafe impl Send for AtomicU16 {} unsafe impl Sync for AtomicU16 {} impl panic::RefUnwindSafe for AtomicU16 {} impl panic::UnwindSafe for AtomicU16 {} impl AtomicU16 { pub(crate) const fn new(val: u16) -> AtomicU16 { let inner = UnsafeCell::new(std::sync::atomic::AtomicU16::new(val)); AtomicU16 { inner } } /// Performs an unsynchronized load. /// /// # Safety /// /// All mutations must have happened before the unsynchronized load. /// Additionally, there must be no concurrent mutations. pub(crate) unsafe fn unsync_load(&self) -> u16 { core::ptr::read(self.inner.get() as *const u16) } } impl Deref for AtomicU16 { type Target = std::sync::atomic::AtomicU16; fn deref(&self) -> &Self::Target { // safety: it is always safe to access `&self` fns on the inner value as // we never perform unsafe mutations. unsafe { &*self.inner.get() } } } impl fmt::Debug for AtomicU16 { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.deref().fmt(fmt) } } tokio-1.43.1/src/loom/std/atomic_u32.rs000064400000000000000000000024411046102023000157020ustar 00000000000000use std::cell::UnsafeCell; use std::fmt; use std::ops::Deref; use std::panic; /// `AtomicU32` providing an additional `unsync_load` function. pub(crate) struct AtomicU32 { inner: UnsafeCell, } unsafe impl Send for AtomicU32 {} unsafe impl Sync for AtomicU32 {} impl panic::RefUnwindSafe for AtomicU32 {} impl panic::UnwindSafe for AtomicU32 {} impl AtomicU32 { pub(crate) const fn new(val: u32) -> AtomicU32 { let inner = UnsafeCell::new(std::sync::atomic::AtomicU32::new(val)); AtomicU32 { inner } } /// Performs an unsynchronized load. /// /// # Safety /// /// All mutations must have happened before the unsynchronized load. /// Additionally, there must be no concurrent mutations. pub(crate) unsafe fn unsync_load(&self) -> u32 { core::ptr::read(self.inner.get() as *const u32) } } impl Deref for AtomicU32 { type Target = std::sync::atomic::AtomicU32; fn deref(&self) -> &Self::Target { // safety: it is always safe to access `&self` fns on the inner value as // we never perform unsafe mutations. unsafe { &*self.inner.get() } } } impl fmt::Debug for AtomicU32 { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.deref().fmt(fmt) } } tokio-1.43.1/src/loom/std/atomic_u64.rs000064400000000000000000000012101046102023000157000ustar 00000000000000//! Implementation of an atomic `u64` cell. On 64 bit platforms, this is a //! re-export of `AtomicU64`. On 32 bit platforms, this is implemented using a //! `Mutex`. // `AtomicU64` can only be used on targets with `target_has_atomic` is 64 or greater. // Once `cfg_target_has_atomic` feature is stable, we can replace it with // `#[cfg(target_has_atomic = "64")]`. // Refs: https://github.com/rust-lang/rust/tree/master/src/librustc_target cfg_has_atomic_u64! { #[path = "atomic_u64_native.rs"] mod imp; } cfg_not_has_atomic_u64! { #[path = "atomic_u64_as_mutex.rs"] mod imp; } pub(crate) use imp::{AtomicU64, StaticAtomicU64}; tokio-1.43.1/src/loom/std/atomic_u64_as_mutex.rs000064400000000000000000000031761046102023000176220ustar 00000000000000use crate::loom::sync::Mutex; use std::sync::atomic::Ordering; cfg_has_const_mutex_new! { #[path = "atomic_u64_static_const_new.rs"] mod static_macro; } cfg_not_has_const_mutex_new! { #[path = "atomic_u64_static_once_cell.rs"] mod static_macro; } pub(crate) use static_macro::StaticAtomicU64; #[derive(Debug)] pub(crate) struct AtomicU64 { inner: Mutex, } impl AtomicU64 { pub(crate) fn load(&self, _: Ordering) -> u64 { *self.inner.lock() } pub(crate) fn store(&self, val: u64, _: Ordering) { *self.inner.lock() = val; } pub(crate) fn fetch_add(&self, val: u64, _: Ordering) -> u64 { let mut lock = self.inner.lock(); let prev = *lock; *lock = prev + val; prev } pub(crate) fn fetch_or(&self, val: u64, _: Ordering) -> u64 { let mut lock = self.inner.lock(); let prev = *lock; *lock = prev | val; prev } pub(crate) fn compare_exchange( &self, current: u64, new: u64, _success: Ordering, _failure: Ordering, ) -> Result { let mut lock = self.inner.lock(); if *lock == current { *lock = new; Ok(current) } else { Err(*lock) } } pub(crate) fn compare_exchange_weak( &self, current: u64, new: u64, success: Ordering, failure: Ordering, ) -> Result { self.compare_exchange(current, new, success, failure) } } impl Default for AtomicU64 { fn default() -> AtomicU64 { AtomicU64::new(u64::default()) } } tokio-1.43.1/src/loom/std/atomic_u64_native.rs000064400000000000000000000002221046102023000172500ustar 00000000000000pub(crate) use std::sync::atomic::{AtomicU64, Ordering}; /// Alias `AtomicU64` to `StaticAtomicU64` pub(crate) type StaticAtomicU64 = AtomicU64; tokio-1.43.1/src/loom/std/atomic_u64_static_const_new.rs000064400000000000000000000003571046102023000213410ustar 00000000000000use super::AtomicU64; use crate::loom::sync::Mutex; pub(crate) type StaticAtomicU64 = AtomicU64; impl AtomicU64 { pub(crate) const fn new(val: u64) -> Self { Self { inner: Mutex::const_new(val), } } } tokio-1.43.1/src/loom/std/atomic_u64_static_once_cell.rs000064400000000000000000000023341046102023000212620ustar 00000000000000use super::AtomicU64; use crate::loom::sync::{atomic::Ordering, Mutex}; use crate::util::once_cell::OnceCell; pub(crate) struct StaticAtomicU64 { init: u64, cell: OnceCell>, } impl AtomicU64 { pub(crate) fn new(val: u64) -> Self { Self { inner: Mutex::new(val), } } } impl StaticAtomicU64 { pub(crate) const fn new(val: u64) -> StaticAtomicU64 { StaticAtomicU64 { init: val, cell: OnceCell::new(), } } pub(crate) fn load(&self, order: Ordering) -> u64 { *self.inner().lock() } pub(crate) fn fetch_add(&self, val: u64, order: Ordering) -> u64 { let mut lock = self.inner().lock(); let prev = *lock; *lock = prev + val; prev } pub(crate) fn compare_exchange_weak( &self, current: u64, new: u64, _success: Ordering, _failure: Ordering, ) -> Result { let mut lock = self.inner().lock(); if *lock == current { *lock = new; Ok(current) } else { Err(*lock) } } fn inner(&self) -> &Mutex { self.cell.get(|| Mutex::new(self.init)) } } tokio-1.43.1/src/loom/std/atomic_usize.rs000064400000000000000000000032451046102023000164330ustar 00000000000000use std::cell::UnsafeCell; use std::fmt; use std::ops; use std::panic; /// `AtomicUsize` providing an additional `unsync_load` function. pub(crate) struct AtomicUsize { inner: UnsafeCell, } unsafe impl Send for AtomicUsize {} unsafe impl Sync for AtomicUsize {} impl panic::RefUnwindSafe for AtomicUsize {} impl panic::UnwindSafe for AtomicUsize {} impl AtomicUsize { pub(crate) const fn new(val: usize) -> AtomicUsize { let inner = UnsafeCell::new(std::sync::atomic::AtomicUsize::new(val)); AtomicUsize { inner } } /// Performs an unsynchronized load. /// /// # Safety /// /// All mutations must have happened before the unsynchronized load. /// Additionally, there must be no concurrent mutations. pub(crate) unsafe fn unsync_load(&self) -> usize { core::ptr::read(self.inner.get() as *const usize) } pub(crate) fn with_mut(&mut self, f: impl FnOnce(&mut usize) -> R) -> R { // safety: we have mutable access f(unsafe { (*self.inner.get()).get_mut() }) } } impl ops::Deref for AtomicUsize { type Target = std::sync::atomic::AtomicUsize; fn deref(&self) -> &Self::Target { // safety: it is always safe to access `&self` fns on the inner value as // we never perform unsafe mutations. unsafe { &*self.inner.get() } } } impl ops::DerefMut for AtomicUsize { fn deref_mut(&mut self) -> &mut Self::Target { // safety: we hold `&mut self` unsafe { &mut *self.inner.get() } } } impl fmt::Debug for AtomicUsize { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { (**self).fmt(fmt) } } tokio-1.43.1/src/loom/std/barrier.rs000064400000000000000000000151101046102023000153600ustar 00000000000000//! A `Barrier` that provides `wait_timeout`. //! //! This implementation mirrors that of the Rust standard library. use crate::loom::sync::{Condvar, Mutex}; use std::fmt; use std::time::{Duration, Instant}; /// A barrier enables multiple threads to synchronize the beginning /// of some computation. /// /// # Examples /// /// ``` /// use std::sync::{Arc, Barrier}; /// use std::thread; /// /// let mut handles = Vec::with_capacity(10); /// let barrier = Arc::new(Barrier::new(10)); /// for _ in 0..10 { /// let c = Arc::clone(&barrier); /// // The same messages will be printed together. /// // You will NOT see any interleaving. /// handles.push(thread::spawn(move|| { /// println!("before wait"); /// c.wait(); /// println!("after wait"); /// })); /// } /// // Wait for other threads to finish. /// for handle in handles { /// handle.join().unwrap(); /// } /// ``` pub(crate) struct Barrier { lock: Mutex, cvar: Condvar, num_threads: usize, } // The inner state of a double barrier struct BarrierState { count: usize, generation_id: usize, } /// A `BarrierWaitResult` is returned by [`Barrier::wait()`] when all threads /// in the [`Barrier`] have rendezvoused. /// /// # Examples /// /// ``` /// use std::sync::Barrier; /// /// let barrier = Barrier::new(1); /// let barrier_wait_result = barrier.wait(); /// ``` pub(crate) struct BarrierWaitResult(bool); impl fmt::Debug for Barrier { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Barrier").finish_non_exhaustive() } } impl Barrier { /// Creates a new barrier that can block a given number of threads. /// /// A barrier will block `n`-1 threads which call [`wait()`] and then wake /// up all threads at once when the `n`th thread calls [`wait()`]. /// /// [`wait()`]: Barrier::wait /// /// # Examples /// /// ``` /// use std::sync::Barrier; /// /// let barrier = Barrier::new(10); /// ``` #[must_use] pub(crate) fn new(n: usize) -> Barrier { Barrier { lock: Mutex::new(BarrierState { count: 0, generation_id: 0, }), cvar: Condvar::new(), num_threads: n, } } /// Blocks the current thread until all threads have rendezvoused here. /// /// Barriers are re-usable after all threads have rendezvoused once, and can /// be used continuously. /// /// A single (arbitrary) thread will receive a [`BarrierWaitResult`] that /// returns `true` from [`BarrierWaitResult::is_leader()`] when returning /// from this function, and all other threads will receive a result that /// will return `false` from [`BarrierWaitResult::is_leader()`]. /// /// # Examples /// /// ``` /// use std::sync::{Arc, Barrier}; /// use std::thread; /// /// let mut handles = Vec::with_capacity(10); /// let barrier = Arc::new(Barrier::new(10)); /// for _ in 0..10 { /// let c = Arc::clone(&barrier); /// // The same messages will be printed together. /// // You will NOT see any interleaving. /// handles.push(thread::spawn(move|| { /// println!("before wait"); /// c.wait(); /// println!("after wait"); /// })); /// } /// // Wait for other threads to finish. /// for handle in handles { /// handle.join().unwrap(); /// } /// ``` pub(crate) fn wait(&self) -> BarrierWaitResult { let mut lock = self.lock.lock(); let local_gen = lock.generation_id; lock.count += 1; if lock.count < self.num_threads { // We need a while loop to guard against spurious wakeups. // https://en.wikipedia.org/wiki/Spurious_wakeup while local_gen == lock.generation_id { lock = self.cvar.wait(lock).unwrap(); } BarrierWaitResult(false) } else { lock.count = 0; lock.generation_id = lock.generation_id.wrapping_add(1); self.cvar.notify_all(); BarrierWaitResult(true) } } /// Blocks the current thread until all threads have rendezvoused here for /// at most `timeout` duration. pub(crate) fn wait_timeout(&self, timeout: Duration) -> Option { // This implementation mirrors `wait`, but with each blocking operation // replaced by a timeout-amenable alternative. let deadline = Instant::now() + timeout; // Acquire `self.lock` with at most `timeout` duration. let mut lock = loop { if let Some(guard) = self.lock.try_lock() { break guard; } else if Instant::now() > deadline { return None; } else { std::thread::yield_now(); } }; // Shrink the `timeout` to account for the time taken to acquire `lock`. let timeout = deadline.saturating_duration_since(Instant::now()); let local_gen = lock.generation_id; lock.count += 1; if lock.count < self.num_threads { // We need a while loop to guard against spurious wakeups. // https://en.wikipedia.org/wiki/Spurious_wakeup while local_gen == lock.generation_id { let (guard, timeout_result) = self.cvar.wait_timeout(lock, timeout).unwrap(); lock = guard; if timeout_result.timed_out() { return None; } } Some(BarrierWaitResult(false)) } else { lock.count = 0; lock.generation_id = lock.generation_id.wrapping_add(1); self.cvar.notify_all(); Some(BarrierWaitResult(true)) } } } impl fmt::Debug for BarrierWaitResult { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BarrierWaitResult") .field("is_leader", &self.is_leader()) .finish() } } impl BarrierWaitResult { /// Returns `true` if this thread is the "leader thread" for the call to /// [`Barrier::wait()`]. /// /// Only one thread will have `true` returned from their result, all other /// threads will have `false` returned. /// /// # Examples /// /// ``` /// use std::sync::Barrier; /// /// let barrier = Barrier::new(1); /// let barrier_wait_result = barrier.wait(); /// println!("{:?}", barrier_wait_result.is_leader()); /// ``` #[must_use] pub(crate) fn is_leader(&self) -> bool { self.0 } } tokio-1.43.1/src/loom/std/mod.rs000064400000000000000000000072721046102023000145230ustar 00000000000000#![cfg_attr(any(not(feature = "full"), loom), allow(unused_imports, dead_code))] mod atomic_u16; mod atomic_u32; mod atomic_u64; mod atomic_usize; mod barrier; mod mutex; #[cfg(all(feature = "parking_lot", not(miri)))] mod parking_lot; mod rwlock; mod unsafe_cell; pub(crate) mod cell { pub(crate) use super::unsafe_cell::UnsafeCell; } #[cfg(any( feature = "net", feature = "process", feature = "signal", feature = "sync", ))] pub(crate) mod future { pub(crate) use crate::sync::AtomicWaker; } pub(crate) mod hint { pub(crate) use std::hint::spin_loop; } pub(crate) mod rand { use std::collections::hash_map::RandomState; use std::hash::{BuildHasher, Hash, Hasher}; use std::sync::atomic::AtomicU32; use std::sync::atomic::Ordering::Relaxed; static COUNTER: AtomicU32 = AtomicU32::new(1); pub(crate) fn seed() -> u64 { let rand_state = RandomState::new(); let mut hasher = rand_state.build_hasher(); // Hash some unique-ish data to generate some new state COUNTER.fetch_add(1, Relaxed).hash(&mut hasher); // Get the seed hasher.finish() } } pub(crate) mod sync { pub(crate) use std::sync::{Arc, Weak}; // Below, make sure all the feature-influenced types are exported for // internal use. Note however that some are not _currently_ named by // consuming code. #[cfg(all(feature = "parking_lot", not(miri)))] #[allow(unused_imports)] pub(crate) use crate::loom::std::parking_lot::{ Condvar, Mutex, MutexGuard, RwLock, RwLockReadGuard, WaitTimeoutResult, }; #[cfg(not(all(feature = "parking_lot", not(miri))))] #[allow(unused_imports)] pub(crate) use std::sync::{Condvar, MutexGuard, RwLockReadGuard, WaitTimeoutResult}; #[cfg(not(all(feature = "parking_lot", not(miri))))] pub(crate) use crate::loom::std::mutex::Mutex; #[cfg(not(all(feature = "parking_lot", not(miri))))] pub(crate) use crate::loom::std::rwlock::RwLock; pub(crate) mod atomic { pub(crate) use crate::loom::std::atomic_u16::AtomicU16; pub(crate) use crate::loom::std::atomic_u32::AtomicU32; pub(crate) use crate::loom::std::atomic_u64::{AtomicU64, StaticAtomicU64}; pub(crate) use crate::loom::std::atomic_usize::AtomicUsize; pub(crate) use std::sync::atomic::{fence, AtomicBool, AtomicPtr, AtomicU8, Ordering}; } pub(crate) use super::barrier::Barrier; } pub(crate) mod sys { #[cfg(feature = "rt-multi-thread")] pub(crate) fn num_cpus() -> usize { use std::num::NonZeroUsize; const ENV_WORKER_THREADS: &str = "TOKIO_WORKER_THREADS"; match std::env::var(ENV_WORKER_THREADS) { Ok(s) => { let n = s.parse().unwrap_or_else(|e| { panic!("\"{ENV_WORKER_THREADS}\" must be usize, error: {e}, value: {s}") }); assert!(n > 0, "\"{ENV_WORKER_THREADS}\" cannot be set to 0"); n } Err(std::env::VarError::NotPresent) => { std::thread::available_parallelism().map_or(1, NonZeroUsize::get) } Err(std::env::VarError::NotUnicode(e)) => { panic!("\"{ENV_WORKER_THREADS}\" must be valid unicode, error: {e:?}") } } } #[cfg(not(feature = "rt-multi-thread"))] pub(crate) fn num_cpus() -> usize { 1 } } pub(crate) mod thread { #[inline] pub(crate) fn yield_now() { std::hint::spin_loop(); } #[allow(unused_imports)] pub(crate) use std::thread::{ current, panicking, park, park_timeout, sleep, spawn, AccessError, Builder, JoinHandle, LocalKey, Result, Thread, ThreadId, }; } tokio-1.43.1/src/loom/std/mutex.rs000064400000000000000000000021441046102023000150770ustar 00000000000000use std::sync::{self, MutexGuard, TryLockError}; /// Adapter for `std::Mutex` that removes the poisoning aspects /// from its api. #[derive(Debug)] pub(crate) struct Mutex(sync::Mutex); #[allow(dead_code)] impl Mutex { #[inline] pub(crate) fn new(t: T) -> Mutex { Mutex(sync::Mutex::new(t)) } #[inline] pub(crate) const fn const_new(t: T) -> Mutex { Mutex(sync::Mutex::new(t)) } #[inline] pub(crate) fn lock(&self) -> MutexGuard<'_, T> { match self.0.lock() { Ok(guard) => guard, Err(p_err) => p_err.into_inner(), } } #[inline] pub(crate) fn try_lock(&self) -> Option> { match self.0.try_lock() { Ok(guard) => Some(guard), Err(TryLockError::Poisoned(p_err)) => Some(p_err.into_inner()), Err(TryLockError::WouldBlock) => None, } } #[inline] pub(crate) fn get_mut(&mut self) -> &mut T { match self.0.get_mut() { Ok(val) => val, Err(p_err) => p_err.into_inner(), } } } tokio-1.43.1/src/loom/std/parking_lot.rs000064400000000000000000000116361046102023000162540ustar 00000000000000//! A minimal adaption of the `parking_lot` synchronization primitives to the //! equivalent `std::sync` types. //! //! This can be extended to additional types/methods as required. use std::fmt; use std::marker::PhantomData; use std::ops::{Deref, DerefMut}; use std::sync::LockResult; use std::time::Duration; // All types in this file are marked with PhantomData to ensure that // parking_lot's send_guard feature does not leak through and affect when Tokio // types are Send. // // See for more info. // Types that do not need wrapping pub(crate) use parking_lot::WaitTimeoutResult; #[derive(Debug)] pub(crate) struct Mutex(PhantomData>, parking_lot::Mutex); #[derive(Debug)] pub(crate) struct RwLock(PhantomData>, parking_lot::RwLock); #[derive(Debug)] pub(crate) struct Condvar(PhantomData, parking_lot::Condvar); #[derive(Debug)] pub(crate) struct MutexGuard<'a, T: ?Sized>( PhantomData>, parking_lot::MutexGuard<'a, T>, ); #[derive(Debug)] pub(crate) struct RwLockReadGuard<'a, T: ?Sized>( PhantomData>, parking_lot::RwLockReadGuard<'a, T>, ); #[derive(Debug)] pub(crate) struct RwLockWriteGuard<'a, T: ?Sized>( PhantomData>, parking_lot::RwLockWriteGuard<'a, T>, ); impl Mutex { #[inline] pub(crate) fn new(t: T) -> Mutex { Mutex(PhantomData, parking_lot::Mutex::new(t)) } #[inline] #[cfg(not(all(loom, test)))] pub(crate) const fn const_new(t: T) -> Mutex { Mutex(PhantomData, parking_lot::const_mutex(t)) } #[inline] pub(crate) fn lock(&self) -> MutexGuard<'_, T> { MutexGuard(PhantomData, self.1.lock()) } #[inline] pub(crate) fn try_lock(&self) -> Option> { self.1 .try_lock() .map(|guard| MutexGuard(PhantomData, guard)) } #[inline] pub(crate) fn get_mut(&mut self) -> &mut T { self.1.get_mut() } // Note: Additional methods `is_poisoned` and `into_inner`, can be // provided here as needed. } impl<'a, T: ?Sized> Deref for MutexGuard<'a, T> { type Target = T; fn deref(&self) -> &T { self.1.deref() } } impl<'a, T: ?Sized> DerefMut for MutexGuard<'a, T> { fn deref_mut(&mut self) -> &mut T { self.1.deref_mut() } } impl RwLock { pub(crate) fn new(t: T) -> RwLock { RwLock(PhantomData, parking_lot::RwLock::new(t)) } pub(crate) fn read(&self) -> RwLockReadGuard<'_, T> { RwLockReadGuard(PhantomData, self.1.read()) } pub(crate) fn try_read(&self) -> Option> { Some(RwLockReadGuard(PhantomData, self.1.read())) } pub(crate) fn write(&self) -> RwLockWriteGuard<'_, T> { RwLockWriteGuard(PhantomData, self.1.write()) } pub(crate) fn try_write(&self) -> Option> { Some(RwLockWriteGuard(PhantomData, self.1.write())) } } impl<'a, T: ?Sized> Deref for RwLockReadGuard<'a, T> { type Target = T; fn deref(&self) -> &T { self.1.deref() } } impl<'a, T: ?Sized> Deref for RwLockWriteGuard<'a, T> { type Target = T; fn deref(&self) -> &T { self.1.deref() } } impl<'a, T: ?Sized> DerefMut for RwLockWriteGuard<'a, T> { fn deref_mut(&mut self) -> &mut T { self.1.deref_mut() } } impl Condvar { #[inline] pub(crate) fn new() -> Condvar { Condvar(PhantomData, parking_lot::Condvar::new()) } #[inline] pub(crate) fn notify_one(&self) { self.1.notify_one(); } #[inline] pub(crate) fn notify_all(&self) { self.1.notify_all(); } #[inline] pub(crate) fn wait<'a, T>( &self, mut guard: MutexGuard<'a, T>, ) -> LockResult> { self.1.wait(&mut guard.1); Ok(guard) } #[inline] pub(crate) fn wait_timeout<'a, T>( &self, mut guard: MutexGuard<'a, T>, timeout: Duration, ) -> LockResult<(MutexGuard<'a, T>, WaitTimeoutResult)> { let wtr = self.1.wait_for(&mut guard.1, timeout); Ok((guard, wtr)) } // Note: Additional methods `wait_timeout_ms`, `wait_timeout_until`, // `wait_until` can be provided here as needed. } impl<'a, T: ?Sized + fmt::Display> fmt::Display for MutexGuard<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&self.1, f) } } impl<'a, T: ?Sized + fmt::Display> fmt::Display for RwLockReadGuard<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&self.1, f) } } impl<'a, T: ?Sized + fmt::Display> fmt::Display for RwLockWriteGuard<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&self.1, f) } } tokio-1.43.1/src/loom/std/rwlock.rs000064400000000000000000000025311046102023000152360ustar 00000000000000use std::sync::{self, RwLockReadGuard, RwLockWriteGuard, TryLockError}; /// Adapter for `std::sync::RwLock` that removes the poisoning aspects /// from its api. #[derive(Debug)] pub(crate) struct RwLock(sync::RwLock); #[allow(dead_code)] impl RwLock { #[inline] pub(crate) fn new(t: T) -> Self { Self(sync::RwLock::new(t)) } #[inline] pub(crate) fn read(&self) -> RwLockReadGuard<'_, T> { match self.0.read() { Ok(guard) => guard, Err(p_err) => p_err.into_inner(), } } #[inline] pub(crate) fn try_read(&self) -> Option> { match self.0.try_read() { Ok(guard) => Some(guard), Err(TryLockError::Poisoned(p_err)) => Some(p_err.into_inner()), Err(TryLockError::WouldBlock) => None, } } #[inline] pub(crate) fn write(&self) -> RwLockWriteGuard<'_, T> { match self.0.write() { Ok(guard) => guard, Err(p_err) => p_err.into_inner(), } } #[inline] pub(crate) fn try_write(&self) -> Option> { match self.0.try_write() { Ok(guard) => Some(guard), Err(TryLockError::Poisoned(p_err)) => Some(p_err.into_inner()), Err(TryLockError::WouldBlock) => None, } } } tokio-1.43.1/src/loom/std/unsafe_cell.rs000064400000000000000000000007241046102023000162170ustar 00000000000000#[derive(Debug)] pub(crate) struct UnsafeCell(std::cell::UnsafeCell); impl UnsafeCell { pub(crate) const fn new(data: T) -> UnsafeCell { UnsafeCell(std::cell::UnsafeCell::new(data)) } #[inline(always)] pub(crate) fn with(&self, f: impl FnOnce(*const T) -> R) -> R { f(self.0.get()) } #[inline(always)] pub(crate) fn with_mut(&self, f: impl FnOnce(*mut T) -> R) -> R { f(self.0.get()) } } tokio-1.43.1/src/macros/addr_of.rs000064400000000000000000000014241046102023000150570ustar 00000000000000//! This module defines a macro that lets you go from a raw pointer to a struct //! to a raw pointer to a field of the struct. macro_rules! generate_addr_of_methods { ( impl<$($gen:ident)*> $struct_name:ty {$( $(#[$attrs:meta])* $vis:vis unsafe fn $fn_name:ident(self: NonNull) -> NonNull<$field_type:ty> { &self$(.$field_name:tt)+ } )*} ) => { impl<$($gen)*> $struct_name {$( $(#[$attrs])* $vis unsafe fn $fn_name(me: ::core::ptr::NonNull) -> ::core::ptr::NonNull<$field_type> { let me = me.as_ptr(); let field = ::std::ptr::addr_of_mut!((*me) $(.$field_name)+ ); ::core::ptr::NonNull::new_unchecked(field) } )*} }; } tokio-1.43.1/src/macros/cfg.rs000064400000000000000000000314061046102023000142230ustar 00000000000000#![allow(unused_macros)] macro_rules! feature { ( #![$meta:meta] $($item:item)* ) => { $( #[cfg($meta)] #[cfg_attr(docsrs, doc(cfg($meta)))] $item )* } } /// Enables Windows-specific code. /// Use this macro instead of `cfg(windows)` to generate docs properly. macro_rules! cfg_windows { ($($item:item)*) => { $( #[cfg(any(all(doc, docsrs), windows))] #[cfg_attr(docsrs, doc(cfg(windows)))] $item )* } } /// Enables Unix-specific code. /// Use this macro instead of `cfg(unix)` to generate docs properly. macro_rules! cfg_unix { ($($item:item)*) => { $( #[cfg(any(all(doc, docsrs), unix))] #[cfg_attr(docsrs, doc(cfg(unix)))] $item )* } } /// Enables unstable Windows-specific code. /// Use this macro instead of `cfg(windows)` to generate docs properly. macro_rules! cfg_unstable_windows { ($($item:item)*) => { $( #[cfg(all(any(all(doc, docsrs), windows), tokio_unstable))] #[cfg_attr(docsrs, doc(cfg(all(windows, tokio_unstable))))] $item )* } } /// Enables `enter::block_on`. macro_rules! cfg_block_on { ($($item:item)*) => { $( #[cfg(any( feature = "fs", feature = "net", feature = "io-std", feature = "rt", ))] $item )* } } /// Enables internal `AtomicWaker` impl. macro_rules! cfg_atomic_waker_impl { ($($item:item)*) => { $( #[cfg(any( feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "time", ))] #[cfg(not(loom))] $item )* } } macro_rules! cfg_aio { ($($item:item)*) => { $( #[cfg(all(any(docsrs, target_os = "freebsd"), feature = "net"))] #[cfg_attr(docsrs, doc(cfg(all(target_os = "freebsd", feature = "net"))) )] $item )* } } macro_rules! cfg_fs { ($($item:item)*) => { $( #[cfg(feature = "fs")] #[cfg_attr(docsrs, doc(cfg(feature = "fs")))] $item )* } } macro_rules! cfg_io_blocking { ($($item:item)*) => { $( #[cfg(any( feature = "io-std", feature = "fs", all(windows, feature = "process"), ))] $item )* } } macro_rules! cfg_io_driver { ($($item:item)*) => { $( #[cfg(any( feature = "net", all(unix, feature = "process"), all(unix, feature = "signal"), ))] #[cfg_attr(docsrs, doc(cfg(any( feature = "net", all(unix, feature = "process"), all(unix, feature = "signal"), ))))] $item )* } } macro_rules! cfg_io_driver_impl { ( $( $item:item )* ) => { $( #[cfg(any( feature = "net", all(unix, feature = "process"), all(unix, feature = "signal"), ))] $item )* } } macro_rules! cfg_not_io_driver { ($($item:item)*) => { $( #[cfg(not(any( feature = "net", all(unix, feature = "process"), all(unix, feature = "signal"), )))] $item )* } } macro_rules! cfg_io_readiness { ($($item:item)*) => { $( #[cfg(feature = "net")] $item )* } } macro_rules! cfg_io_std { ($($item:item)*) => { $( #[cfg(feature = "io-std")] #[cfg_attr(docsrs, doc(cfg(feature = "io-std")))] $item )* } } macro_rules! cfg_io_util { ($($item:item)*) => { $( #[cfg(feature = "io-util")] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] $item )* } } macro_rules! cfg_not_io_util { ($($item:item)*) => { $( #[cfg(not(feature = "io-util"))] $item )* } } macro_rules! cfg_loom { ($($item:item)*) => { $( #[cfg(loom)] $item )* } } macro_rules! cfg_not_loom { ($($item:item)*) => { $( #[cfg(not(loom))] $item )* } } macro_rules! cfg_macros { ($($item:item)*) => { $( #[cfg(feature = "macros")] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] $item )* } } macro_rules! cfg_unstable_metrics { ($($item:item)*) => { $( #[cfg(tokio_unstable)] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] $item )* } } /// Some metrics require 64-bit atomics. macro_rules! cfg_64bit_metrics { ($($item:item)*) => { $( #[cfg(target_has_atomic = "64")] #[cfg_attr(docsrs, doc(cfg(target_has_atomic = "64")))] $item )* } } macro_rules! cfg_no_64bit_metrics { ($($item:item)*) => { $( #[cfg(not(target_has_atomic = "64"))] $item )* } } macro_rules! cfg_not_unstable_metrics { ($($item:item)*) => { $( #[cfg(not(tokio_unstable))] $item )* } } macro_rules! cfg_not_rt_and_metrics_and_net { ($($item:item)*) => { $( #[cfg(not(all(feature = "net", feature = "rt", tokio_unstable)))]$item )* } } macro_rules! cfg_net_or_process { ($($item:item)*) => { $( #[cfg(any(feature = "net", feature = "process"))] #[cfg_attr(docsrs, doc(cfg(any(feature = "net", feature = "process"))))] $item )* } } macro_rules! cfg_net { ($($item:item)*) => { $( #[cfg(feature = "net")] #[cfg_attr(docsrs, doc(cfg(feature = "net")))] $item )* } } macro_rules! cfg_net_unix { ($($item:item)*) => { $( #[cfg(all(unix, feature = "net"))] #[cfg_attr(docsrs, doc(cfg(all(unix, feature = "net"))))] $item )* } } macro_rules! cfg_net_windows { ($($item:item)*) => { $( #[cfg(all(any(all(doc, docsrs), windows), feature = "net"))] #[cfg_attr(docsrs, doc(cfg(all(windows, feature = "net"))))] $item )* } } macro_rules! cfg_process { ($($item:item)*) => { $( #[cfg(feature = "process")] #[cfg_attr(docsrs, doc(cfg(feature = "process")))] #[cfg(not(loom))] #[cfg(not(target_os = "wasi"))] $item )* } } macro_rules! cfg_process_driver { ($($item:item)*) => { #[cfg(unix)] #[cfg(not(loom))] cfg_process! { $($item)* } } } macro_rules! cfg_not_process_driver { ($($item:item)*) => { $( #[cfg(not(all(unix, not(loom), feature = "process")))] $item )* } } macro_rules! cfg_signal { ($($item:item)*) => { $( #[cfg(feature = "signal")] #[cfg_attr(docsrs, doc(cfg(feature = "signal")))] #[cfg(not(loom))] #[cfg(not(target_os = "wasi"))] $item )* } } macro_rules! cfg_signal_internal { ($($item:item)*) => { $( #[cfg(any(feature = "signal", all(unix, feature = "process")))] #[cfg(not(loom))] $item )* } } macro_rules! cfg_signal_internal_and_unix { ($($item:item)*) => { #[cfg(unix)] cfg_signal_internal! { $($item)* } } } macro_rules! cfg_not_signal_internal { ($($item:item)*) => { $( #[cfg(any(loom, not(unix), not(any(feature = "signal", all(unix, feature = "process")))))] $item )* } } macro_rules! cfg_sync { ($($item:item)*) => { $( #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(cfg(feature = "sync")))] $item )* } } macro_rules! cfg_not_sync { ($($item:item)*) => { $( #[cfg(not(feature = "sync"))] $item )* } } macro_rules! cfg_rt { ($($item:item)*) => { $( #[cfg(feature = "rt")] #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] $item )* } } macro_rules! cfg_not_rt { ($($item:item)*) => { $( #[cfg(not(feature = "rt"))] $item )* } } macro_rules! cfg_rt_multi_thread { ($($item:item)*) => { $( #[cfg(feature = "rt-multi-thread")] #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] $item )* } } macro_rules! cfg_not_rt_multi_thread { ($($item:item)*) => { $( #[cfg(not(feature = "rt-multi-thread"))] $item )* } } macro_rules! cfg_taskdump { ($($item:item)*) => { $( #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any( target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64" ) ))] $item )* }; } macro_rules! cfg_not_taskdump { ($($item:item)*) => { $( #[cfg(not(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any( target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64" ) )))] $item )* }; } macro_rules! cfg_test_util { ($($item:item)*) => { $( #[cfg(feature = "test-util")] #[cfg_attr(docsrs, doc(cfg(feature = "test-util")))] $item )* } } macro_rules! cfg_not_test_util { ($($item:item)*) => { $( #[cfg(not(feature = "test-util"))] $item )* } } macro_rules! cfg_time { ($($item:item)*) => { $( #[cfg(feature = "time")] #[cfg_attr(docsrs, doc(cfg(feature = "time")))] $item )* } } macro_rules! cfg_not_time { ($($item:item)*) => { $( #[cfg(not(feature = "time"))] $item )* } } macro_rules! cfg_trace { ($($item:item)*) => { $( #[cfg(all(tokio_unstable, feature = "tracing"))] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] $item )* }; } macro_rules! cfg_unstable { ($($item:item)*) => { $( #[cfg(tokio_unstable)] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] $item )* }; } macro_rules! cfg_not_trace { ($($item:item)*) => { $( #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] $item )* } } macro_rules! cfg_coop { ($($item:item)*) => { $( #[cfg(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", ))] $item )* } } macro_rules! cfg_not_coop { ($($item:item)*) => { $( #[cfg(not(any( feature = "fs", feature = "io-std", feature = "net", feature = "process", feature = "rt", feature = "signal", feature = "sync", feature = "time", )))] $item )* } } macro_rules! cfg_has_atomic_u64 { ($($item:item)*) => { $( #[cfg(target_has_atomic = "64")] $item )* } } macro_rules! cfg_not_has_atomic_u64 { ($($item:item)*) => { $( #[cfg(not(target_has_atomic = "64"))] $item )* } } macro_rules! cfg_has_const_mutex_new { ($($item:item)*) => { $( #[cfg(not(all(loom, test)))] $item )* } } macro_rules! cfg_not_has_const_mutex_new { ($($item:item)*) => { $( #[cfg(all(loom, test))] $item )* } } macro_rules! cfg_not_wasi { ($($item:item)*) => { $( #[cfg(not(target_os = "wasi"))] $item )* } } macro_rules! cfg_is_wasm_not_wasi { ($($item:item)*) => { $( #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] $item )* } } tokio-1.43.1/src/macros/join.rs000064400000000000000000000142011046102023000144150ustar 00000000000000macro_rules! doc { ($join:item) => { /// Waits on multiple concurrent branches, returning when **all** branches /// complete. /// /// The `join!` macro must be used inside of async functions, closures, and /// blocks. /// /// The `join!` macro takes a list of async expressions and evaluates them /// concurrently on the same task. Each async expression evaluates to a future /// and the futures from each expression are multiplexed on the current task. /// /// When working with async expressions returning `Result`, `join!` will wait /// for **all** branches complete regardless if any complete with `Err`. Use /// [`try_join!`] to return early when `Err` is encountered. /// /// [`try_join!`]: crate::try_join /// /// # Notes /// /// The supplied futures are stored inline and do not require allocating a /// `Vec`. /// /// ### Runtime characteristics /// /// By running all async expressions on the current task, the expressions are /// able to run **concurrently** but not in **parallel**. This means all /// expressions are run on the same thread and if one branch blocks the thread, /// all other expressions will be unable to continue. If parallelism is /// required, spawn each async expression using [`tokio::spawn`] and pass the /// join handle to `join!`. /// /// [`tokio::spawn`]: crate::spawn /// /// # Examples /// /// Basic join with two branches /// /// ``` /// async fn do_stuff_async() { /// // async work /// } /// /// async fn more_async_work() { /// // more here /// } /// /// #[tokio::main] /// async fn main() { /// let (first, second) = tokio::join!( /// do_stuff_async(), /// more_async_work()); /// /// // do something with the values /// } /// ``` #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] $join }; } #[cfg(doc)] doc! {macro_rules! join { ($($future:expr),*) => { unimplemented!() } }} #[cfg(not(doc))] doc! {macro_rules! join { (@ { // One `_` for each branch in the `join!` macro. This is not used once // normalization is complete. ( $($count:tt)* ) // The expression `0+1+1+ ... +1` equal to the number of branches. ( $($total:tt)* ) // Normalized join! branches $( ( $($skip:tt)* ) $e:expr, )* }) => {{ use $crate::macros::support::{maybe_done, poll_fn, Future, Pin}; use $crate::macros::support::Poll::{Ready, Pending}; // Safety: nothing must be moved out of `futures`. This is to satisfy // the requirement of `Pin::new_unchecked` called below. // // We can't use the `pin!` macro for this because `futures` is a tuple // and the standard library provides no way to pin-project to the fields // of a tuple. let mut futures = ( $( maybe_done($e), )* ); // This assignment makes sure that the `poll_fn` closure only has a // reference to the futures, instead of taking ownership of them. This // mitigates the issue described in // let mut futures = &mut futures; // Each time the future created by poll_fn is polled, a different future will be polled first // to ensure every future passed to join! gets a chance to make progress even if // one of the futures consumes the whole budget. // // This is number of futures that will be skipped in the first loop // iteration the next time. let mut skip_next_time: u32 = 0; poll_fn(move |cx| { const COUNT: u32 = $($total)*; let mut is_pending = false; let mut to_run = COUNT; // The number of futures that will be skipped in the first loop iteration. let mut skip = skip_next_time; skip_next_time = if skip + 1 == COUNT { 0 } else { skip + 1 }; // This loop runs twice and the first `skip` futures // are not polled in the first iteration. loop { $( if skip == 0 { if to_run == 0 { // Every future has been polled break; } to_run -= 1; // Extract the future for this branch from the tuple. let ( $($skip,)* fut, .. ) = &mut *futures; // Safety: future is stored on the stack above // and never moved. let mut fut = unsafe { Pin::new_unchecked(fut) }; // Try polling if fut.poll(cx).is_pending() { is_pending = true; } } else { // Future skipped, one less future to skip in the next iteration skip -= 1; } )* } if is_pending { Pending } else { Ready(($({ // Extract the future for this branch from the tuple. let ( $($skip,)* fut, .. ) = &mut futures; // Safety: future is stored on the stack above // and never moved. let mut fut = unsafe { Pin::new_unchecked(fut) }; fut.take_output().expect("expected completed future") },)*)) } }).await }}; // ===== Normalize ===== (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { $crate::join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) }; // ===== Entry point ===== ( $($e:expr),+ $(,)?) => { $crate::join!(@{ () (0) } $($e,)*) }; () => { async {}.await } }} tokio-1.43.1/src/macros/loom.rs000064400000000000000000000001621046102023000144250ustar 00000000000000macro_rules! if_loom { ($($t:tt)*) => {{ #[cfg(loom)] { $($t)* } }} } tokio-1.43.1/src/macros/mod.rs000064400000000000000000000006641046102023000142450ustar 00000000000000#![cfg_attr(not(feature = "full"), allow(unused_macros))] #[macro_use] mod cfg; #[macro_use] mod loom; #[macro_use] mod pin; #[macro_use] mod thread_local; #[macro_use] mod addr_of; cfg_trace! { #[macro_use] mod trace; } cfg_macros! { #[macro_use] mod select; #[macro_use] mod join; #[macro_use] mod try_join; } // Includes re-exports needed to implement macros #[doc(hidden)] pub mod support; tokio-1.43.1/src/macros/pin.rs000064400000000000000000000065221046102023000142530ustar 00000000000000/// Pins a value on the stack. /// /// Calls to `async fn` return anonymous [`Future`] values that are `!Unpin`. /// These values must be pinned before they can be polled. Calling `.await` will /// handle this, but consumes the future. If it is required to call `.await` on /// a `&mut _` reference, the caller is responsible for pinning the future. /// /// Pinning may be done by allocating with [`Box::pin`] or by using the stack /// with the `pin!` macro. /// /// The following will **fail to compile**: /// /// ```compile_fail /// async fn my_async_fn() { /// // async logic here /// } /// /// #[tokio::main] /// async fn main() { /// let mut future = my_async_fn(); /// (&mut future).await; /// } /// ``` /// /// To make this work requires pinning: /// /// ``` /// use tokio::pin; /// /// async fn my_async_fn() { /// // async logic here /// } /// /// #[tokio::main] /// async fn main() { /// let future = my_async_fn(); /// pin!(future); /// /// (&mut future).await; /// } /// ``` /// /// Pinning is useful when using `select!` and stream operators that require `T: /// Stream + Unpin`. /// /// [`Future`]: trait@std::future::Future /// [`Box::pin`]: std::boxed::Box::pin /// /// # Usage /// /// The `pin!` macro takes **identifiers** as arguments. It does **not** work /// with expressions. /// /// The following does not compile as an expression is passed to `pin!`. /// /// ```compile_fail /// async fn my_async_fn() { /// // async logic here /// } /// /// #[tokio::main] /// async fn main() { /// let mut future = pin!(my_async_fn()); /// (&mut future).await; /// } /// ``` /// /// # Examples /// /// Using with select: /// /// ``` /// use tokio::{pin, select}; /// use tokio_stream::{self as stream, StreamExt}; /// /// async fn my_async_fn() { /// // async logic here /// } /// /// #[tokio::main] /// async fn main() { /// let mut stream = stream::iter(vec![1, 2, 3, 4]); /// /// let future = my_async_fn(); /// pin!(future); /// /// loop { /// select! { /// _ = &mut future => { /// // Stop looping `future` will be polled after completion /// break; /// } /// Some(val) = stream.next() => { /// println!("got value = {}", val); /// } /// } /// } /// } /// ``` /// /// Because assigning to a variable followed by pinning is common, there is also /// a variant of the macro that supports doing both in one go. /// /// ``` /// use tokio::{pin, select}; /// /// async fn my_async_fn() { /// // async logic here /// } /// /// #[tokio::main] /// async fn main() { /// pin! { /// let future1 = my_async_fn(); /// let future2 = my_async_fn(); /// } /// /// select! { /// _ = &mut future1 => {} /// _ = &mut future2 => {} /// } /// } /// ``` #[macro_export] macro_rules! pin { ($($x:ident),*) => { $( // Move the value to ensure that it is owned let mut $x = $x; // Shadow the original binding so that it can't be directly accessed // ever again. #[allow(unused_mut)] let mut $x = unsafe { $crate::macros::support::Pin::new_unchecked(&mut $x) }; )* }; ($( let $x:ident = $init:expr; )*) => { $( let $x = $init; $crate::pin!($x); )* }; } tokio-1.43.1/src/macros/select.rs000064400000000000000000001451541046102023000147510ustar 00000000000000macro_rules! doc { ($select:item) => { /// Waits on multiple concurrent branches, returning when the **first** branch /// completes, cancelling the remaining branches. /// /// The `select!` macro must be used inside of async functions, closures, and /// blocks. /// /// The `select!` macro accepts one or more branches with the following pattern: /// /// ```text /// = (, if )? => , /// ``` /// /// Additionally, the `select!` macro may include a single, optional `else` /// branch, which evaluates if none of the other branches match their patterns: /// /// ```text /// else => /// ``` /// /// The macro aggregates all `` expressions and runs them /// concurrently on the **current** task. Once the **first** expression /// completes with a value that matches its ``, the `select!` macro /// returns the result of evaluating the completed branch's `` /// expression. /// /// Additionally, each branch may include an optional `if` precondition. If the /// precondition returns `false`, then the branch is disabled. The provided /// `` is still evaluated but the resulting future is never /// polled. This capability is useful when using `select!` within a loop. /// /// The complete lifecycle of a `select!` expression is as follows: /// /// 1. Evaluate all provided `` expressions. If the precondition /// returns `false`, disable the branch for the remainder of the current call /// to `select!`. Re-entering `select!` due to a loop clears the "disabled" /// state. /// 2. Aggregate the ``s from each branch, including the /// disabled ones. If the branch is disabled, `` is still /// evaluated, but the resulting future is not polled. /// 3. If **all** branches are disabled: go to step 6. /// 4. Concurrently await on the results for all remaining ``s. /// 5. Once an `` returns a value, attempt to apply the value to the /// provided ``. If the pattern matches, evaluate the `` and return. /// If the pattern **does not** match, disable the current branch for the remainder of /// the current call to `select!`. Continue from step 3. /// 6. Evaluate the `else` expression. If no else expression is provided, panic. /// /// # Runtime characteristics /// /// By running all async expressions on the current task, the expressions are /// able to run **concurrently** but not in **parallel**. This means all /// expressions are run on the same thread and if one branch blocks the thread, /// all other expressions will be unable to continue. If parallelism is /// required, spawn each async expression using [`tokio::spawn`] and pass the /// join handle to `select!`. /// /// [`tokio::spawn`]: crate::spawn /// /// # Fairness /// /// By default, `select!` randomly picks a branch to check first. This provides /// some level of fairness when calling `select!` in a loop with branches that /// are always ready. /// /// This behavior can be overridden by adding `biased;` to the beginning of the /// macro usage. See the examples for details. This will cause `select` to poll /// the futures in the order they appear from top to bottom. There are a few /// reasons you may want this: /// /// - The random number generation of `tokio::select!` has a non-zero CPU cost /// - Your futures may interact in a way where known polling order is significant /// /// But there is an important caveat to this mode. It becomes your responsibility /// to ensure that the polling order of your futures is fair. If for example you /// are selecting between a stream and a shutdown future, and the stream has a /// huge volume of messages and zero or nearly zero time between them, you should /// place the shutdown future earlier in the `select!` list to ensure that it is /// always polled, and will not be ignored due to the stream being constantly /// ready. /// /// # Panics /// /// The `select!` macro panics if all branches are disabled **and** there is no /// provided `else` branch. A branch is disabled when the provided `if` /// precondition returns `false` **or** when the pattern does not match the /// result of ``. /// /// # Cancellation safety /// /// When using `select!` in a loop to receive messages from multiple sources, /// you should make sure that the receive call is cancellation safe to avoid /// losing messages. This section goes through various common methods and /// describes whether they are cancel safe. The lists in this section are not /// exhaustive. /// /// The following methods are cancellation safe: /// /// * [`tokio::sync::mpsc::Receiver::recv`](crate::sync::mpsc::Receiver::recv) /// * [`tokio::sync::mpsc::UnboundedReceiver::recv`](crate::sync::mpsc::UnboundedReceiver::recv) /// * [`tokio::sync::broadcast::Receiver::recv`](crate::sync::broadcast::Receiver::recv) /// * [`tokio::sync::watch::Receiver::changed`](crate::sync::watch::Receiver::changed) /// * [`tokio::net::TcpListener::accept`](crate::net::TcpListener::accept) /// * [`tokio::net::UnixListener::accept`](crate::net::UnixListener::accept) /// * [`tokio::signal::unix::Signal::recv`](crate::signal::unix::Signal::recv) /// * [`tokio::io::AsyncReadExt::read`](crate::io::AsyncReadExt::read) on any `AsyncRead` /// * [`tokio::io::AsyncReadExt::read_buf`](crate::io::AsyncReadExt::read_buf) on any `AsyncRead` /// * [`tokio::io::AsyncWriteExt::write`](crate::io::AsyncWriteExt::write) on any `AsyncWrite` /// * [`tokio::io::AsyncWriteExt::write_buf`](crate::io::AsyncWriteExt::write_buf) on any `AsyncWrite` /// * [`tokio_stream::StreamExt::next`](https://docs.rs/tokio-stream/0.1/tokio_stream/trait.StreamExt.html#method.next) on any `Stream` /// * [`futures::stream::StreamExt::next`](https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.next) on any `Stream` /// /// The following methods are not cancellation safe and can lead to loss of data: /// /// * [`tokio::io::AsyncReadExt::read_exact`](crate::io::AsyncReadExt::read_exact) /// * [`tokio::io::AsyncReadExt::read_to_end`](crate::io::AsyncReadExt::read_to_end) /// * [`tokio::io::AsyncReadExt::read_to_string`](crate::io::AsyncReadExt::read_to_string) /// * [`tokio::io::AsyncWriteExt::write_all`](crate::io::AsyncWriteExt::write_all) /// /// The following methods are not cancellation safe because they use a queue for /// fairness and cancellation makes you lose your place in the queue: /// /// * [`tokio::sync::Mutex::lock`](crate::sync::Mutex::lock) /// * [`tokio::sync::RwLock::read`](crate::sync::RwLock::read) /// * [`tokio::sync::RwLock::write`](crate::sync::RwLock::write) /// * [`tokio::sync::Semaphore::acquire`](crate::sync::Semaphore::acquire) /// * [`tokio::sync::Notify::notified`](crate::sync::Notify::notified) /// /// To determine whether your own methods are cancellation safe, look for the /// location of uses of `.await`. This is because when an asynchronous method is /// cancelled, that always happens at an `.await`. If your function behaves /// correctly even if it is restarted while waiting at an `.await`, then it is /// cancellation safe. /// /// Cancellation safety can be defined in the following way: If you have a /// future that has not yet completed, then it must be a no-op to drop that /// future and recreate it. This definition is motivated by the situation where /// a `select!` is used in a loop. Without this guarantee, you would lose your /// progress when another branch completes and you restart the `select!` by /// going around the loop. /// /// Be aware that cancelling something that is not cancellation safe is not /// necessarily wrong. For example, if you are cancelling a task because the /// application is shutting down, then you probably don't care that partially /// read data is lost. /// /// # Examples /// /// Basic select with two branches. /// /// ``` /// async fn do_stuff_async() { /// // async work /// } /// /// async fn more_async_work() { /// // more here /// } /// /// #[tokio::main] /// async fn main() { /// tokio::select! { /// _ = do_stuff_async() => { /// println!("do_stuff_async() completed first") /// } /// _ = more_async_work() => { /// println!("more_async_work() completed first") /// } /// }; /// } /// ``` /// /// Basic stream selecting. /// /// ``` /// use tokio_stream::{self as stream, StreamExt}; /// /// #[tokio::main] /// async fn main() { /// let mut stream1 = stream::iter(vec![1, 2, 3]); /// let mut stream2 = stream::iter(vec![4, 5, 6]); /// /// let next = tokio::select! { /// v = stream1.next() => v.unwrap(), /// v = stream2.next() => v.unwrap(), /// }; /// /// assert!(next == 1 || next == 4); /// } /// ``` /// /// Collect the contents of two streams. In this example, we rely on pattern /// matching and the fact that `stream::iter` is "fused", i.e. once the stream /// is complete, all calls to `next()` return `None`. /// /// ``` /// use tokio_stream::{self as stream, StreamExt}; /// /// #[tokio::main] /// async fn main() { /// let mut stream1 = stream::iter(vec![1, 2, 3]); /// let mut stream2 = stream::iter(vec![4, 5, 6]); /// /// let mut values = vec![]; /// /// loop { /// tokio::select! { /// Some(v) = stream1.next() => values.push(v), /// Some(v) = stream2.next() => values.push(v), /// else => break, /// } /// } /// /// values.sort(); /// assert_eq!(&[1, 2, 3, 4, 5, 6], &values[..]); /// } /// ``` /// /// Using the same future in multiple `select!` expressions can be done by passing /// a reference to the future. Doing so requires the future to be [`Unpin`]. A /// future can be made [`Unpin`] by either using [`Box::pin`] or stack pinning. /// /// [`Unpin`]: std::marker::Unpin /// [`Box::pin`]: std::boxed::Box::pin /// /// Here, a stream is consumed for at most 1 second. /// /// ``` /// use tokio_stream::{self as stream, StreamExt}; /// use tokio::time::{self, Duration}; /// /// #[tokio::main] /// async fn main() { /// let mut stream = stream::iter(vec![1, 2, 3]); /// let sleep = time::sleep(Duration::from_secs(1)); /// tokio::pin!(sleep); /// /// loop { /// tokio::select! { /// maybe_v = stream.next() => { /// if let Some(v) = maybe_v { /// println!("got = {}", v); /// } else { /// break; /// } /// } /// _ = &mut sleep => { /// println!("timeout"); /// break; /// } /// } /// } /// } /// ``` /// /// Joining two values using `select!`. /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx1, mut rx1) = oneshot::channel(); /// let (tx2, mut rx2) = oneshot::channel(); /// /// tokio::spawn(async move { /// tx1.send("first").unwrap(); /// }); /// /// tokio::spawn(async move { /// tx2.send("second").unwrap(); /// }); /// /// let mut a = None; /// let mut b = None; /// /// while a.is_none() || b.is_none() { /// tokio::select! { /// v1 = (&mut rx1), if a.is_none() => a = Some(v1.unwrap()), /// v2 = (&mut rx2), if b.is_none() => b = Some(v2.unwrap()), /// } /// } /// /// let res = (a.unwrap(), b.unwrap()); /// /// assert_eq!(res.0, "first"); /// assert_eq!(res.1, "second"); /// } /// ``` /// /// Using the `biased;` mode to control polling order. /// /// ``` /// #[tokio::main] /// async fn main() { /// let mut count = 0u8; /// /// loop { /// tokio::select! { /// // If you run this example without `biased;`, the polling order is /// // pseudo-random, and the assertions on the value of count will /// // (probably) fail. /// biased; /// /// _ = async {}, if count < 1 => { /// count += 1; /// assert_eq!(count, 1); /// } /// _ = async {}, if count < 2 => { /// count += 1; /// assert_eq!(count, 2); /// } /// _ = async {}, if count < 3 => { /// count += 1; /// assert_eq!(count, 3); /// } /// _ = async {}, if count < 4 => { /// count += 1; /// assert_eq!(count, 4); /// } /// /// else => { /// break; /// } /// }; /// } /// } /// ``` /// /// ## Avoid racy `if` preconditions /// /// Given that `if` preconditions are used to disable `select!` branches, some /// caution must be used to avoid missing values. /// /// For example, here is **incorrect** usage of `sleep` with `if`. The objective /// is to repeatedly run an asynchronous task for up to 50 milliseconds. /// However, there is a potential for the `sleep` completion to be missed. /// /// ```no_run,should_panic /// use tokio::time::{self, Duration}; /// /// async fn some_async_work() { /// // do work /// } /// /// #[tokio::main] /// async fn main() { /// let sleep = time::sleep(Duration::from_millis(50)); /// tokio::pin!(sleep); /// /// while !sleep.is_elapsed() { /// tokio::select! { /// _ = &mut sleep, if !sleep.is_elapsed() => { /// println!("operation timed out"); /// } /// _ = some_async_work() => { /// println!("operation completed"); /// } /// } /// } /// /// panic!("This example shows how not to do it!"); /// } /// ``` /// /// In the above example, `sleep.is_elapsed()` may return `true` even if /// `sleep.poll()` never returned `Ready`. This opens up a potential race /// condition where `sleep` expires between the `while !sleep.is_elapsed()` /// check and the call to `select!` resulting in the `some_async_work()` call to /// run uninterrupted despite the sleep having elapsed. /// /// One way to write the above example without the race would be: /// /// ``` /// use tokio::time::{self, Duration}; /// /// async fn some_async_work() { /// # time::sleep(Duration::from_millis(10)).await; /// // do work /// } /// /// #[tokio::main] /// async fn main() { /// let sleep = time::sleep(Duration::from_millis(50)); /// tokio::pin!(sleep); /// /// loop { /// tokio::select! { /// _ = &mut sleep => { /// println!("operation timed out"); /// break; /// } /// _ = some_async_work() => { /// println!("operation completed"); /// } /// } /// } /// } /// ``` #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] $select }; } #[cfg(doc)] doc! {macro_rules! select { { $( biased; )? $( $bind:pat = $fut:expr $(, if $cond:expr)? => $handler:expr, )* $( else => $els:expr $(,)? )? } => { unimplemented!() }; }} #[cfg(not(doc))] doc! {macro_rules! select { // Uses a declarative macro to do **most** of the work. While it is possible // to implement fully with a declarative macro, a procedural macro is used // to enable improved error messages. // // The macro is structured as a tt-muncher. All branches are processed and // normalized. Once the input is normalized, it is passed to the top-most // rule. When entering the macro, `@{ }` is inserted at the front. This is // used to collect the normalized input. // // The macro only recurses once per branch. This allows using `select!` // without requiring the user to increase the recursion limit. // All input is normalized, now transform. (@ { // The index of the future to poll first (in bias mode), or the RNG // expression to use to pick a future to poll first. start=$start:expr; // One `_` for each branch in the `select!` macro. Passing this to // `count!` converts $skip to an integer. ( $($count:tt)* ) // Normalized select branches. `( $skip )` is a set of `_` characters. // There is one `_` for each select branch **before** this one. Given // that all input futures are stored in a tuple, $skip is useful for // generating a pattern to reference the future for the current branch. // $skip is also used as an argument to `count!`, returning the index of // the current select branch. $( ( $($skip:tt)* ) $bind:pat = $fut:expr, if $c:expr => $handle:expr, )+ // Fallback expression used when all select branches have been disabled. ; $else:expr }) => {{ // Enter a context where stable "function-like" proc macros can be used. // // This module is defined within a scope and should not leak out of this // macro. #[doc(hidden)] mod __tokio_select_util { // Generate an enum with one variant per select branch $crate::select_priv_declare_output_enum!( ( $($count)* ) ); } // `tokio::macros::support` is a public, but doc(hidden) module // including a re-export of all types needed by this macro. use $crate::macros::support::Future; use $crate::macros::support::Pin; use $crate::macros::support::Poll::{Ready, Pending}; const BRANCHES: u32 = $crate::count!( $($count)* ); let mut disabled: __tokio_select_util::Mask = Default::default(); // First, invoke all the pre-conditions. For any that return true, // set the appropriate bit in `disabled`. $( if !$c { let mask: __tokio_select_util::Mask = 1 << $crate::count!( $($skip)* ); disabled |= mask; } )* // Create a scope to separate polling from handling the output. This // adds borrow checker flexibility when using the macro. let mut output = { // Store each future directly first (that is, without wrapping the future in a call to // `IntoFuture::into_future`). This allows the `$fut` expression to make use of // temporary lifetime extension. // // https://doc.rust-lang.org/1.58.1/reference/destructors.html#temporary-lifetime-extension let futures_init = ($( $fut, )+); // Safety: Nothing must be moved out of `futures`. This is to // satisfy the requirement of `Pin::new_unchecked` called below. // // We can't use the `pin!` macro for this because `futures` is a // tuple and the standard library provides no way to pin-project to // the fields of a tuple. let mut futures = ($( $crate::macros::support::IntoFuture::into_future( $crate::count_field!( futures_init.$($skip)* ) ),)+); // This assignment makes sure that the `poll_fn` closure only has a // reference to the futures, instead of taking ownership of them. // This mitigates the issue described in // let mut futures = &mut futures; $crate::macros::support::poll_fn(|cx| { // Track if any branch returns pending. If no branch completes // **or** returns pending, this implies that all branches are // disabled. let mut is_pending = false; // Choose a starting index to begin polling the futures at. In // practice, this will either be a pseudo-randomly generated // number by default, or the constant 0 if `biased;` is // supplied. let start = $start; for i in 0..BRANCHES { let branch; #[allow(clippy::modulo_one)] { branch = (start + i) % BRANCHES; } match branch { $( #[allow(unreachable_code)] $crate::count!( $($skip)* ) => { // First, if the future has previously been // disabled, do not poll it again. This is done // by checking the associated bit in the // `disabled` bit field. let mask = 1 << branch; if disabled & mask == mask { // The future has been disabled. continue; } // Extract the future for this branch from the // tuple let ( $($skip,)* fut, .. ) = &mut *futures; // Safety: future is stored on the stack above // and never moved. let mut fut = unsafe { Pin::new_unchecked(fut) }; // Try polling it let out = match Future::poll(fut, cx) { Ready(out) => out, Pending => { // Track that at least one future is // still pending and continue polling. is_pending = true; continue; } }; // Disable the future from future polling. disabled |= mask; // The future returned a value, check if matches // the specified pattern. #[allow(unused_variables)] #[allow(unused_mut)] match &out { $crate::select_priv_clean_pattern!($bind) => {} _ => continue, } // The select is complete, return the value return Ready($crate::select_variant!(__tokio_select_util::Out, ($($skip)*))(out)); } )* _ => unreachable!("reaching this means there probably is an off by one bug"), } } if is_pending { Pending } else { // All branches have been disabled. Ready(__tokio_select_util::Out::Disabled) } }).await }; match output { $( $crate::select_variant!(__tokio_select_util::Out, ($($skip)*) ($bind)) => $handle, )* __tokio_select_util::Out::Disabled => $else, _ => unreachable!("failed to match bind"), } }}; // ==== Normalize ===== // These rules match a single `select!` branch and normalize it for // processing by the first rule. (@ { start=$start:expr; $($t:tt)* } ) => { // No `else` branch $crate::select!(@{ start=$start; $($t)*; panic!("all branches are disabled and there is no else branch") }) }; (@ { start=$start:expr; $($t:tt)* } else => $else:expr $(,)?) => { $crate::select!(@{ start=$start; $($t)*; $else }) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:block, $($r:tt)* ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if $c => $h, } $($r)*) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr => $h:block, $($r:tt)* ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if true => $h, } $($r)*) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:block $($r:tt)* ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if $c => $h, } $($r)*) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr => $h:block $($r:tt)* ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if true => $h, } $($r)*) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:expr ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if $c => $h, }) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr => $h:expr ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if true => $h, }) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:expr, $($r:tt)* ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if $c => $h, } $($r)*) }; (@ { start=$start:expr; ( $($s:tt)* ) $($t:tt)* } $p:pat = $f:expr => $h:expr, $($r:tt)* ) => { $crate::select!(@{ start=$start; ($($s)* _) $($t)* ($($s)*) $p = $f, if true => $h, } $($r)*) }; // ===== Entry point ===== ($(biased;)? else => $else:expr $(,)? ) => {{ $else }}; (biased; $p:pat = $($t:tt)* ) => { $crate::select!(@{ start=0; () } $p = $($t)*) }; ( $p:pat = $($t:tt)* ) => { // Randomly generate a starting point. This makes `select!` a bit more // fair and avoids always polling the first future. $crate::select!(@{ start={ $crate::macros::support::thread_rng_n(BRANCHES) }; () } $p = $($t)*) }; () => { compile_error!("select! requires at least one branch.") }; }} // And here... we manually list out matches for up to 64 branches... I'm not // happy about it either, but this is how we manage to use a declarative macro! #[macro_export] #[doc(hidden)] macro_rules! count { () => { 0 }; (_) => { 1 }; (_ _) => { 2 }; (_ _ _) => { 3 }; (_ _ _ _) => { 4 }; (_ _ _ _ _) => { 5 }; (_ _ _ _ _ _) => { 6 }; (_ _ _ _ _ _ _) => { 7 }; (_ _ _ _ _ _ _ _) => { 8 }; (_ _ _ _ _ _ _ _ _) => { 9 }; (_ _ _ _ _ _ _ _ _ _) => { 10 }; (_ _ _ _ _ _ _ _ _ _ _) => { 11 }; (_ _ _ _ _ _ _ _ _ _ _ _) => { 12 }; (_ _ _ _ _ _ _ _ _ _ _ _ _) => { 13 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 14 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 15 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 16 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 17 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 18 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 19 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 20 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 21 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 22 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 23 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 24 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 25 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 26 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 27 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 28 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 29 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 30 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 31 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 32 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 33 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 34 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 35 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 36 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 37 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 38 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 39 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 40 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 41 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 42 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 43 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 44 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 45 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 46 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 47 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 48 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 49 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 50 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 51 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 52 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 53 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 54 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 55 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 56 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 57 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 58 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 59 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 60 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 61 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 62 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 63 }; (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { 64 }; } #[macro_export] #[doc(hidden)] macro_rules! count_field { ($var:ident. ) => { $var.0 }; ($var:ident. _) => { $var.1 }; ($var:ident. _ _) => { $var.2 }; ($var:ident. _ _ _) => { $var.3 }; ($var:ident. _ _ _ _) => { $var.4 }; ($var:ident. _ _ _ _ _) => { $var.5 }; ($var:ident. _ _ _ _ _ _) => { $var.6 }; ($var:ident. _ _ _ _ _ _ _) => { $var.7 }; ($var:ident. _ _ _ _ _ _ _ _) => { $var.8 }; ($var:ident. _ _ _ _ _ _ _ _ _) => { $var.9 }; ($var:ident. _ _ _ _ _ _ _ _ _ _) => { $var.10 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _) => { $var.11 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _) => { $var.12 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.13 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.14 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.15 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.16 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.17 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.18 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.19 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.20 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.21 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.22 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.23 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.24 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.25 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.26 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.27 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.28 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.29 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.30 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.31 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.32 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.33 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.34 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.35 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.36 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.37 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.38 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.39 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.40 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.41 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.42 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.43 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.44 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.45 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.46 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.47 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.48 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.49 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.50 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.51 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.52 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.53 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.54 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.55 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.56 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.57 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.58 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.59 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.60 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.61 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.62 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.63 }; ($var:ident. _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) => { $var.64 }; } #[macro_export] #[doc(hidden)] macro_rules! select_variant { ($($p:ident)::*, () $($t:tt)*) => { $($p)::*::_0 $($t)* }; ($($p:ident)::*, (_) $($t:tt)*) => { $($p)::*::_1 $($t)* }; ($($p:ident)::*, (_ _) $($t:tt)*) => { $($p)::*::_2 $($t)* }; ($($p:ident)::*, (_ _ _) $($t:tt)*) => { $($p)::*::_3 $($t)* }; ($($p:ident)::*, (_ _ _ _) $($t:tt)*) => { $($p)::*::_4 $($t)* }; ($($p:ident)::*, (_ _ _ _ _) $($t:tt)*) => { $($p)::*::_5 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_6 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_7 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_8 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_9 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_10 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_11 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_12 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_13 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_14 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_15 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_16 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_17 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_18 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_19 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_20 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_21 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_22 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_23 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_24 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_25 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_26 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_27 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_28 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_29 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_30 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_31 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_32 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_33 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_34 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_35 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_36 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_37 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_38 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_39 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_40 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_41 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_42 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_43 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_44 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_45 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_46 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_47 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_48 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_49 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_50 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_51 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_52 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_53 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_54 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_55 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_56 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_57 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_58 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_59 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_60 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_61 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_62 $($t)* }; ($($p:ident)::*, (_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _) $($t:tt)*) => { $($p)::*::_63 $($t)* }; } tokio-1.43.1/src/macros/support.rs000064400000000000000000000004661046102023000152020ustar 00000000000000cfg_macros! { pub use crate::future::maybe_done::maybe_done; pub use std::future::poll_fn; #[doc(hidden)] pub fn thread_rng_n(n: u32) -> u32 { crate::runtime::context::thread_rng_n(n) } } pub use std::future::{Future, IntoFuture}; pub use std::pin::Pin; pub use std::task::Poll; tokio-1.43.1/src/macros/thread_local.rs000064400000000000000000000007221046102023000161020ustar 00000000000000#[cfg(all(loom, test))] macro_rules! tokio_thread_local { ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty = const { $expr:expr } $(;)?) => { loom::thread_local! { $(#[$attrs])* $vis static $name: $ty = $expr; } }; ($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } } } #[cfg(not(all(loom, test)))] macro_rules! tokio_thread_local { ($($tts:tt)+) => { ::std::thread_local!{ $($tts)+ } } } tokio-1.43.1/src/macros/trace.rs000064400000000000000000000013261046102023000145600ustar 00000000000000cfg_trace! { macro_rules! trace_op { ($name:expr, $readiness:literal) => { tracing::trace!( target: "runtime::resource::poll_op", op_name = $name, is_ready = $readiness ); } } macro_rules! trace_poll_op { ($name:expr, $poll:expr $(,)*) => { match $poll { std::task::Poll::Ready(t) => { trace_op!($name, true); std::task::Poll::Ready(t) } std::task::Poll::Pending => { trace_op!($name, false); return std::task::Poll::Pending; } } }; } } tokio-1.43.1/src/macros/try_join.rs000064400000000000000000000202201046102023000153110ustar 00000000000000macro_rules! doc { ($try_join:item) => { /// Waits on multiple concurrent branches, returning when **all** branches /// complete with `Ok(_)` or on the first `Err(_)`. /// /// The `try_join!` macro must be used inside of async functions, closures, and /// blocks. /// /// Similar to [`join!`], the `try_join!` macro takes a list of async /// expressions and evaluates them concurrently on the same task. Each async /// expression evaluates to a future and the futures from each expression are /// multiplexed on the current task. The `try_join!` macro returns when **all** /// branches return with `Ok` or when the **first** branch returns with `Err`. /// /// [`join!`]: macro@join /// /// # Notes /// /// The supplied futures are stored inline and do not require allocating a /// `Vec`. /// /// ### Runtime characteristics /// /// By running all async expressions on the current task, the expressions are /// able to run **concurrently** but not in **parallel**. This means all /// expressions are run on the same thread and if one branch blocks the thread, /// all other expressions will be unable to continue. If parallelism is /// required, spawn each async expression using [`tokio::spawn`] and pass the /// join handle to `try_join!`. /// /// [`tokio::spawn`]: crate::spawn /// /// # Examples /// /// Basic `try_join` with two branches. /// /// ``` /// async fn do_stuff_async() -> Result<(), &'static str> { /// // async work /// # Ok(()) /// } /// /// async fn more_async_work() -> Result<(), &'static str> { /// // more here /// # Ok(()) /// } /// /// #[tokio::main] /// async fn main() { /// let res = tokio::try_join!( /// do_stuff_async(), /// more_async_work()); /// /// match res { /// Ok((first, second)) => { /// // do something with the values /// } /// Err(err) => { /// println!("processing failed; error = {}", err); /// } /// } /// } /// ``` /// /// Using `try_join!` with spawned tasks. /// /// ``` /// use tokio::task::JoinHandle; /// /// async fn do_stuff_async() -> Result<(), &'static str> { /// // async work /// # Err("failed") /// } /// /// async fn more_async_work() -> Result<(), &'static str> { /// // more here /// # Ok(()) /// } /// /// async fn flatten(handle: JoinHandle>) -> Result { /// match handle.await { /// Ok(Ok(result)) => Ok(result), /// Ok(Err(err)) => Err(err), /// Err(err) => Err("handling failed"), /// } /// } /// /// #[tokio::main] /// async fn main() { /// let handle1 = tokio::spawn(do_stuff_async()); /// let handle2 = tokio::spawn(more_async_work()); /// match tokio::try_join!(flatten(handle1), flatten(handle2)) { /// Ok(val) => { /// // do something with the values /// } /// Err(err) => { /// println!("Failed with {}.", err); /// # assert_eq!(err, "failed"); /// } /// } /// } /// ``` #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "macros")))] $try_join }; } #[cfg(doc)] doc! {macro_rules! try_join { ($($future:expr),*) => { unimplemented!() } }} #[cfg(not(doc))] doc! {macro_rules! try_join { (@ { // One `_` for each branch in the `try_join!` macro. This is not used once // normalization is complete. ( $($count:tt)* ) // The expression `0+1+1+ ... +1` equal to the number of branches. ( $($total:tt)* ) // Normalized try_join! branches $( ( $($skip:tt)* ) $e:expr, )* }) => {{ use $crate::macros::support::{maybe_done, poll_fn, Future, Pin}; use $crate::macros::support::Poll::{Ready, Pending}; // Safety: nothing must be moved out of `futures`. This is to satisfy // the requirement of `Pin::new_unchecked` called below. // // We can't use the `pin!` macro for this because `futures` is a tuple // and the standard library provides no way to pin-project to the fields // of a tuple. let mut futures = ( $( maybe_done($e), )* ); // This assignment makes sure that the `poll_fn` closure only has a // reference to the futures, instead of taking ownership of them. This // mitigates the issue described in // let mut futures = &mut futures; // Each time the future created by poll_fn is polled, a different future will be polled first // to ensure every future passed to join! gets a chance to make progress even if // one of the futures consumes the whole budget. // // This is number of futures that will be skipped in the first loop // iteration the next time. let mut skip_next_time: u32 = 0; poll_fn(move |cx| { const COUNT: u32 = $($total)*; let mut is_pending = false; let mut to_run = COUNT; // The number of futures that will be skipped in the first loop iteration let mut skip = skip_next_time; skip_next_time = if skip + 1 == COUNT { 0 } else { skip + 1 }; // This loop runs twice and the first `skip` futures // are not polled in the first iteration. loop { $( if skip == 0 { if to_run == 0 { // Every future has been polled break; } to_run -= 1; // Extract the future for this branch from the tuple. let ( $($skip,)* fut, .. ) = &mut *futures; // Safety: future is stored on the stack above // and never moved. let mut fut = unsafe { Pin::new_unchecked(fut) }; // Try polling if fut.as_mut().poll(cx).is_pending() { is_pending = true; } else if fut.as_mut().output_mut().expect("expected completed future").is_err() { return Ready(Err(fut.take_output().expect("expected completed future").err().unwrap())) } } else { // Future skipped, one less future to skip in the next iteration skip -= 1; } )* } if is_pending { Pending } else { Ready(Ok(($({ // Extract the future for this branch from the tuple. let ( $($skip,)* fut, .. ) = &mut futures; // Safety: future is stored on the stack above // and never moved. let mut fut = unsafe { Pin::new_unchecked(fut) }; fut .take_output() .expect("expected completed future") .ok() .expect("expected Ok(_)") },)*))) } }).await }}; // ===== Normalize ===== (@ { ( $($s:tt)* ) ( $($n:tt)* ) $($t:tt)* } $e:expr, $($r:tt)* ) => { $crate::try_join!(@{ ($($s)* _) ($($n)* + 1) $($t)* ($($s)*) $e, } $($r)*) }; // ===== Entry point ===== ( $($e:expr),+ $(,)?) => { $crate::try_join!(@{ () (0) } $($e,)*) }; () => { async { Ok(()) }.await } }} tokio-1.43.1/src/net/addr.rs000064400000000000000000000236011046102023000136760ustar 00000000000000use std::future; use std::io; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; /// Converts or resolves without blocking to one or more `SocketAddr` values. /// /// # DNS /// /// Implementations of `ToSocketAddrs` for string types require a DNS lookup. /// /// # Calling /// /// Currently, this trait is only used as an argument to Tokio functions that /// need to reference a target socket address. To perform a `SocketAddr` /// conversion directly, use [`lookup_host()`](super::lookup_host()). /// /// This trait is sealed and is intended to be opaque. The details of the trait /// will change. Stabilization is pending enhancements to the Rust language. pub trait ToSocketAddrs: sealed::ToSocketAddrsPriv {} type ReadyFuture = future::Ready>; cfg_net! { pub(crate) fn to_socket_addrs(arg: T) -> T::Future where T: ToSocketAddrs, { arg.to_socket_addrs(sealed::Internal) } } // ===== impl &impl ToSocketAddrs ===== impl ToSocketAddrs for &T {} impl sealed::ToSocketAddrsPriv for &T where T: sealed::ToSocketAddrsPriv + ?Sized, { type Iter = T::Iter; type Future = T::Future; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { (**self).to_socket_addrs(sealed::Internal) } } // ===== impl SocketAddr ===== impl ToSocketAddrs for SocketAddr {} impl sealed::ToSocketAddrsPriv for SocketAddr { type Iter = std::option::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { let iter = Some(*self).into_iter(); future::ready(Ok(iter)) } } // ===== impl SocketAddrV4 ===== impl ToSocketAddrs for SocketAddrV4 {} impl sealed::ToSocketAddrsPriv for SocketAddrV4 { type Iter = std::option::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { SocketAddr::V4(*self).to_socket_addrs(sealed::Internal) } } // ===== impl SocketAddrV6 ===== impl ToSocketAddrs for SocketAddrV6 {} impl sealed::ToSocketAddrsPriv for SocketAddrV6 { type Iter = std::option::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { SocketAddr::V6(*self).to_socket_addrs(sealed::Internal) } } // ===== impl (IpAddr, u16) ===== impl ToSocketAddrs for (IpAddr, u16) {} impl sealed::ToSocketAddrsPriv for (IpAddr, u16) { type Iter = std::option::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { let iter = Some(SocketAddr::from(*self)).into_iter(); future::ready(Ok(iter)) } } // ===== impl (Ipv4Addr, u16) ===== impl ToSocketAddrs for (Ipv4Addr, u16) {} impl sealed::ToSocketAddrsPriv for (Ipv4Addr, u16) { type Iter = std::option::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { let (ip, port) = *self; SocketAddrV4::new(ip, port).to_socket_addrs(sealed::Internal) } } // ===== impl (Ipv6Addr, u16) ===== impl ToSocketAddrs for (Ipv6Addr, u16) {} impl sealed::ToSocketAddrsPriv for (Ipv6Addr, u16) { type Iter = std::option::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { let (ip, port) = *self; SocketAddrV6::new(ip, port, 0, 0).to_socket_addrs(sealed::Internal) } } // ===== impl &[SocketAddr] ===== impl ToSocketAddrs for &[SocketAddr] {} impl sealed::ToSocketAddrsPriv for &[SocketAddr] { type Iter = std::vec::IntoIter; type Future = ReadyFuture; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { #[inline] fn slice_to_vec(addrs: &[SocketAddr]) -> Vec { addrs.to_vec() } // This uses a helper method because clippy doesn't like the `to_vec()` // call here (it will allocate, whereas `self.iter().copied()` would // not), but it's actually necessary in order to ensure that the // returned iterator is valid for the `'static` lifetime, which the // borrowed `slice::Iter` iterator would not be. // // Note that we can't actually add an `allow` attribute for // `clippy::unnecessary_to_owned` here, as Tokio's CI runs clippy lints // on Rust 1.52 to avoid breaking LTS releases of Tokio. Users of newer // Rust versions who see this lint should just ignore it. let iter = slice_to_vec(self).into_iter(); future::ready(Ok(iter)) } } cfg_net! { // ===== impl str ===== impl ToSocketAddrs for str {} impl sealed::ToSocketAddrsPriv for str { type Iter = sealed::OneOrMore; type Future = sealed::MaybeReady; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { use crate::blocking::spawn_blocking; use sealed::MaybeReady; // First check if the input parses as a socket address let res: Result = self.parse(); if let Ok(addr) = res { return MaybeReady(sealed::State::Ready(Some(addr))); } // Run DNS lookup on the blocking pool let s = self.to_owned(); MaybeReady(sealed::State::Blocking(spawn_blocking(move || { std::net::ToSocketAddrs::to_socket_addrs(&s) }))) } } // ===== impl (&str, u16) ===== impl ToSocketAddrs for (&str, u16) {} impl sealed::ToSocketAddrsPriv for (&str, u16) { type Iter = sealed::OneOrMore; type Future = sealed::MaybeReady; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { use crate::blocking::spawn_blocking; use sealed::MaybeReady; let (host, port) = *self; // try to parse the host as a regular IP address first if let Ok(addr) = host.parse::() { let addr = SocketAddrV4::new(addr, port); let addr = SocketAddr::V4(addr); return MaybeReady(sealed::State::Ready(Some(addr))); } if let Ok(addr) = host.parse::() { let addr = SocketAddrV6::new(addr, port, 0, 0); let addr = SocketAddr::V6(addr); return MaybeReady(sealed::State::Ready(Some(addr))); } let host = host.to_owned(); MaybeReady(sealed::State::Blocking(spawn_blocking(move || { std::net::ToSocketAddrs::to_socket_addrs(&(&host[..], port)) }))) } } // ===== impl (String, u16) ===== impl ToSocketAddrs for (String, u16) {} impl sealed::ToSocketAddrsPriv for (String, u16) { type Iter = sealed::OneOrMore; type Future = sealed::MaybeReady; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { (self.0.as_str(), self.1).to_socket_addrs(sealed::Internal) } } // ===== impl String ===== impl ToSocketAddrs for String {} impl sealed::ToSocketAddrsPriv for String { type Iter = ::Iter; type Future = ::Future; fn to_socket_addrs(&self, _: sealed::Internal) -> Self::Future { self[..].to_socket_addrs(sealed::Internal) } } } pub(crate) mod sealed { //! The contents of this trait are intended to remain private and __not__ //! part of the `ToSocketAddrs` public API. The details will change over //! time. use std::future::Future; use std::io; use std::net::SocketAddr; #[doc(hidden)] pub trait ToSocketAddrsPriv { type Iter: Iterator + Send + 'static; type Future: Future> + Send + 'static; fn to_socket_addrs(&self, internal: Internal) -> Self::Future; } #[allow(missing_debug_implementations)] pub struct Internal; cfg_net! { use crate::blocking::JoinHandle; use std::option; use std::pin::Pin; use std::task::{ready,Context, Poll}; use std::vec; #[doc(hidden)] #[derive(Debug)] pub struct MaybeReady(pub(super) State); #[derive(Debug)] pub(super) enum State { Ready(Option), Blocking(JoinHandle>>), } #[doc(hidden)] #[derive(Debug)] pub enum OneOrMore { One(option::IntoIter), More(vec::IntoIter), } impl Future for MaybeReady { type Output = io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match self.0 { State::Ready(ref mut i) => { let iter = OneOrMore::One(i.take().into_iter()); Poll::Ready(Ok(iter)) } State::Blocking(ref mut rx) => { let res = ready!(Pin::new(rx).poll(cx))?.map(OneOrMore::More); Poll::Ready(res) } } } } impl Iterator for OneOrMore { type Item = SocketAddr; fn next(&mut self) -> Option { match self { OneOrMore::One(i) => i.next(), OneOrMore::More(i) => i.next(), } } fn size_hint(&self) -> (usize, Option) { match self { OneOrMore::One(i) => i.size_hint(), OneOrMore::More(i) => i.size_hint(), } } } } } tokio-1.43.1/src/net/lookup_host.rs000064400000000000000000000017651046102023000153410ustar 00000000000000cfg_net! { use crate::net::addr::{self, ToSocketAddrs}; use std::io; use std::net::SocketAddr; /// Performs a DNS resolution. /// /// The returned iterator may not actually yield any values depending on the /// outcome of any resolution performed. /// /// This API is not intended to cover all DNS use cases. Anything beyond the /// basic use case should be done with a specialized library. /// /// # Examples /// /// To resolve a DNS entry: /// /// ```no_run /// use tokio::net; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// for addr in net::lookup_host("localhost:3000").await? { /// println!("socket address is {}", addr); /// } /// /// Ok(()) /// } /// ``` pub async fn lookup_host(host: T) -> io::Result> where T: ToSocketAddrs { addr::to_socket_addrs(host).await } } tokio-1.43.1/src/net/mod.rs000064400000000000000000000036351046102023000135500ustar 00000000000000#![cfg(not(loom))] //! TCP/UDP/Unix bindings for `tokio`. //! //! This module contains the TCP/UDP/Unix networking types, similar to the standard //! library, which can be used to implement networking protocols. //! //! # Organization //! //! * [`TcpListener`] and [`TcpStream`] provide functionality for communication over TCP //! * [`UdpSocket`] provides functionality for communication over UDP //! * [`UnixListener`] and [`UnixStream`] provide functionality for communication over a //! Unix Domain Stream Socket **(available on Unix only)** //! * [`UnixDatagram`] provides functionality for communication //! over Unix Domain Datagram Socket **(available on Unix only)** //! * [`tokio::net::unix::pipe`] for FIFO pipes **(available on Unix only)** //! * [`tokio::net::windows::named_pipe`] for Named Pipes **(available on Windows only)** //! //! For IO resources not available in `tokio::net`, you can use [`AsyncFd`]. //! //! [`TcpListener`]: TcpListener //! [`TcpStream`]: TcpStream //! [`UdpSocket`]: UdpSocket //! [`UnixListener`]: UnixListener //! [`UnixStream`]: UnixStream //! [`UnixDatagram`]: UnixDatagram //! [`tokio::net::unix::pipe`]: unix::pipe //! [`tokio::net::windows::named_pipe`]: windows::named_pipe //! [`AsyncFd`]: crate::io::unix::AsyncFd mod addr; cfg_not_wasi! { #[cfg(feature = "net")] pub(crate) use addr::to_socket_addrs; } pub use addr::ToSocketAddrs; cfg_net! { mod lookup_host; pub use lookup_host::lookup_host; pub mod tcp; pub use tcp::listener::TcpListener; pub use tcp::stream::TcpStream; cfg_not_wasi! { pub use tcp::socket::TcpSocket; mod udp; #[doc(inline)] pub use udp::UdpSocket; } } cfg_net_unix! { pub mod unix; pub use unix::datagram::socket::UnixDatagram; pub use unix::listener::UnixListener; pub use unix::stream::UnixStream; pub use unix::socket::UnixSocket; } cfg_net_windows! { pub mod windows; } tokio-1.43.1/src/net/tcp/listener.rs000064400000000000000000000341351046102023000154030ustar 00000000000000use crate::io::{Interest, PollEvented}; use crate::net::tcp::TcpStream; cfg_not_wasi! { use crate::net::{to_socket_addrs, ToSocketAddrs}; } use std::fmt; use std::io; use std::net::{self, SocketAddr}; use std::task::{ready, Context, Poll}; cfg_net! { /// A TCP socket server, listening for connections. /// /// You can accept a new connection by using the [`accept`](`TcpListener::accept`) /// method. /// /// A `TcpListener` can be turned into a `Stream` with [`TcpListenerStream`]. /// /// [`TcpListenerStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.TcpListenerStream.html /// /// # Errors /// /// Note that accepting a connection can lead to various errors and not all /// of them are necessarily fatal ‒ for example having too many open file /// descriptors or the other side closing the connection while it waits in /// an accept queue. These would terminate the stream if not handled in any /// way. /// /// # Examples /// /// Using `accept`: /// ```no_run /// use tokio::net::TcpListener; /// /// use std::io; /// /// async fn process_socket(socket: T) { /// # drop(socket); /// // do work with socket here /// } /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:8080").await?; /// /// loop { /// let (socket, _) = listener.accept().await?; /// process_socket(socket).await; /// } /// } /// ``` pub struct TcpListener { io: PollEvented, } } impl TcpListener { cfg_not_wasi! { /// Creates a new `TcpListener`, which will be bound to the specified address. /// /// The returned listener is ready for accepting connections. /// /// Binding with a port number of 0 will request that the OS assigns a port /// to this listener. The port allocated can be queried via the `local_addr` /// method. /// /// The address type can be any implementor of the [`ToSocketAddrs`] trait. /// If `addr` yields multiple addresses, bind will be attempted with each of /// the addresses until one succeeds and returns the listener. If none of /// the addresses succeed in creating a listener, the error returned from /// the last attempt (the last address) is returned. /// /// This function sets the `SO_REUSEADDR` option on the socket. /// /// To configure the socket before binding, you can use the [`TcpSocket`] /// type. /// /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs /// [`TcpSocket`]: struct@crate::net::TcpSocket /// /// # Examples /// /// ```no_run /// # if cfg!(miri) { return } // No `socket` in miri. /// use tokio::net::TcpListener; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:2345").await?; /// /// // use the listener /// /// # let _ = listener; /// Ok(()) /// } /// ``` pub async fn bind(addr: A) -> io::Result { let addrs = to_socket_addrs(addr).await?; let mut last_err = None; for addr in addrs { match TcpListener::bind_addr(addr) { Ok(listener) => return Ok(listener), Err(e) => last_err = Some(e), } } Err(last_err.unwrap_or_else(|| { io::Error::new( io::ErrorKind::InvalidInput, "could not resolve to any address", ) })) } fn bind_addr(addr: SocketAddr) -> io::Result { let listener = mio::net::TcpListener::bind(addr)?; TcpListener::new(listener) } } /// Accepts a new incoming connection from this listener. /// /// This function will yield once a new TCP connection is established. When /// established, the corresponding [`TcpStream`] and the remote peer's /// address will be returned. /// /// # Cancel safety /// /// This method is cancel safe. If the method is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no new connections were /// accepted by this method. /// /// [`TcpStream`]: struct@crate::net::TcpStream /// /// # Examples /// /// ```no_run /// use tokio::net::TcpListener; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:8080").await?; /// /// match listener.accept().await { /// Ok((_socket, addr)) => println!("new client: {:?}", addr), /// Err(e) => println!("couldn't get client: {:?}", e), /// } /// /// Ok(()) /// } /// ``` pub async fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> { let (mio, addr) = self .io .registration() .async_io(Interest::READABLE, || self.io.accept()) .await?; let stream = TcpStream::new(mio)?; Ok((stream, addr)) } /// Polls to accept a new incoming connection to this listener. /// /// If there is no connection to accept, `Poll::Pending` is returned and the /// current task will be notified by a waker. Note that on multiple calls /// to `poll_accept`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll> { loop { let ev = ready!(self.io.registration().poll_read_ready(cx))?; match self.io.accept() { Ok((io, addr)) => { let io = TcpStream::new(io)?; return Poll::Ready(Ok((io, addr))); } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.io.registration().clear_readiness(ev); } Err(e) => return Poll::Ready(Err(e)), } } } /// Creates new `TcpListener` from a `std::net::TcpListener`. /// /// This function is intended to be used to wrap a TCP listener from the /// standard library in the Tokio equivalent. /// /// This API is typically paired with the `socket2` crate and the `Socket` /// type to build up and customize a listener before it's shipped off to the /// backing event loop. This allows configuration of options like /// `SO_REUSEPORT`, binding to multiple addresses, etc. /// /// # Notes /// /// The caller is responsible for ensuring that the listener is in /// non-blocking mode. Otherwise all I/O operations on the listener /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::net::TcpListener::set_nonblocking /// /// # Examples /// /// ```rust,no_run /// use std::error::Error; /// use tokio::net::TcpListener; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let std_listener = std::net::TcpListener::bind("127.0.0.1:0")?; /// std_listener.set_nonblocking(true)?; /// let listener = TcpListener::from_std(std_listener)?; /// Ok(()) /// } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. #[track_caller] pub fn from_std(listener: net::TcpListener) -> io::Result { let io = mio::net::TcpListener::from_std(listener); let io = PollEvented::new(io)?; Ok(TcpListener { io }) } /// Turns a [`tokio::net::TcpListener`] into a [`std::net::TcpListener`]. /// /// The returned [`std::net::TcpListener`] will have nonblocking mode set as /// `true`. Use [`set_nonblocking`] to change the blocking mode if needed. /// /// # Examples /// /// ```rust,no_run /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let tokio_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; /// let std_listener = tokio_listener.into_std()?; /// std_listener.set_nonblocking(false)?; /// Ok(()) /// } /// ``` /// /// [`tokio::net::TcpListener`]: TcpListener /// [`std::net::TcpListener`]: std::net::TcpListener /// [`set_nonblocking`]: fn@std::net::TcpListener::set_nonblocking pub fn into_std(self) -> io::Result { #[cfg(unix)] { use std::os::unix::io::{FromRawFd, IntoRawFd}; self.io .into_inner() .map(IntoRawFd::into_raw_fd) .map(|raw_fd| unsafe { std::net::TcpListener::from_raw_fd(raw_fd) }) } #[cfg(windows)] { use std::os::windows::io::{FromRawSocket, IntoRawSocket}; self.io .into_inner() .map(|io| io.into_raw_socket()) .map(|raw_socket| unsafe { std::net::TcpListener::from_raw_socket(raw_socket) }) } #[cfg(target_os = "wasi")] { use std::os::wasi::io::{FromRawFd, IntoRawFd}; self.io .into_inner() .map(|io| io.into_raw_fd()) .map(|raw_fd| unsafe { std::net::TcpListener::from_raw_fd(raw_fd) }) } } cfg_not_wasi! { pub(crate) fn new(listener: mio::net::TcpListener) -> io::Result { let io = PollEvented::new(listener)?; Ok(TcpListener { io }) } } /// Returns the local address that this listener is bound to. /// /// This can be useful, for example, when binding to port 0 to figure out /// which port was actually bound. /// /// # Examples /// /// ```rust,no_run /// use tokio::net::TcpListener; /// /// use std::io; /// use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:8080").await?; /// /// assert_eq!(listener.local_addr()?, /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 8080))); /// /// Ok(()) /// } /// ``` pub fn local_addr(&self) -> io::Result { self.io.local_addr() } /// Gets the value of the `IP_TTL` option for this socket. /// /// For more information about this option, see [`set_ttl`]. /// /// [`set_ttl`]: method@Self::set_ttl /// /// # Examples /// /// ```no_run /// use tokio::net::TcpListener; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:0").await?; /// /// listener.set_ttl(100).expect("could not set TTL"); /// assert_eq!(listener.ttl()?, 100); /// /// Ok(()) /// } /// ``` pub fn ttl(&self) -> io::Result { self.io.ttl() } /// Sets the value for the `IP_TTL` option on this socket. /// /// This value sets the time-to-live field that is used in every packet sent /// from this socket. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpListener; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:0").await?; /// /// listener.set_ttl(100).expect("could not set TTL"); /// /// Ok(()) /// } /// ``` pub fn set_ttl(&self, ttl: u32) -> io::Result<()> { self.io.set_ttl(ttl) } } impl TryFrom for TcpListener { type Error = io::Error; /// Consumes stream, returning the tokio I/O object. /// /// This is equivalent to /// [`TcpListener::from_std(stream)`](TcpListener::from_std). fn try_from(stream: net::TcpListener) -> Result { Self::from_std(stream) } } impl fmt::Debug for TcpListener { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.io.fmt(f) } } #[cfg(unix)] mod sys { use super::TcpListener; use std::os::unix::prelude::*; impl AsRawFd for TcpListener { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for TcpListener { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } cfg_unstable! { #[cfg(target_os = "wasi")] mod sys { use super::TcpListener; use std::os::wasi::prelude::*; impl AsRawFd for TcpListener { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for TcpListener { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } } cfg_windows! { use crate::os::windows::io::{AsRawSocket, RawSocket, AsSocket, BorrowedSocket}; impl AsRawSocket for TcpListener { fn as_raw_socket(&self) -> RawSocket { self.io.as_raw_socket() } } impl AsSocket for TcpListener { fn as_socket(&self) -> BorrowedSocket<'_> { unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } } } } tokio-1.43.1/src/net/tcp/mod.rs000064400000000000000000000004411046102023000143260ustar 00000000000000//! TCP utility types. pub(crate) mod listener; cfg_not_wasi! { pub(crate) mod socket; } mod split; pub use split::{ReadHalf, WriteHalf}; mod split_owned; pub use split_owned::{OwnedReadHalf, OwnedWriteHalf, ReuniteError}; pub(crate) mod stream; pub(crate) use stream::TcpStream; tokio-1.43.1/src/net/tcp/socket.rs000064400000000000000000000670621046102023000150530ustar 00000000000000use crate::net::{TcpListener, TcpStream}; use std::fmt; use std::io; use std::net::SocketAddr; #[cfg(unix)] use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::time::Duration; cfg_windows! { use crate::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket, AsSocket, BorrowedSocket}; } cfg_net! { /// A TCP socket that has not yet been converted to a `TcpStream` or /// `TcpListener`. /// /// `TcpSocket` wraps an operating system socket and enables the caller to /// configure the socket before establishing a TCP connection or accepting /// inbound connections. The caller is able to set socket option and explicitly /// bind the socket with a socket address. /// /// The underlying socket is closed when the `TcpSocket` value is dropped. /// /// `TcpSocket` should only be used directly if the default configuration used /// by `TcpStream::connect` and `TcpListener::bind` does not meet the required /// use case. /// /// Calling `TcpStream::connect("127.0.0.1:8080")` is equivalent to: /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// let stream = socket.connect(addr).await?; /// # drop(stream); /// /// Ok(()) /// } /// ``` /// /// Calling `TcpListener::bind("127.0.0.1:8080")` is equivalent to: /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// // On platforms with Berkeley-derived sockets, this allows to quickly /// // rebind a socket, without needing to wait for the OS to clean up the /// // previous one. /// // /// // On Windows, this allows rebinding sockets which are actively in use, /// // which allows "socket hijacking", so we explicitly don't set it here. /// // https://docs.microsoft.com/en-us/windows/win32/winsock/using-so-reuseaddr-and-so-exclusiveaddruse /// socket.set_reuseaddr(true)?; /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// # drop(listener); /// /// Ok(()) /// } /// ``` /// /// Setting socket options not explicitly provided by `TcpSocket` may be done by /// accessing the `RawFd`/`RawSocket` using [`AsRawFd`]/[`AsRawSocket`] and /// setting the option with a crate like [`socket2`]. /// /// [`RawFd`]: https://doc.rust-lang.org/std/os/unix/io/type.RawFd.html /// [`RawSocket`]: https://doc.rust-lang.org/std/os/windows/io/type.RawSocket.html /// [`AsRawFd`]: https://doc.rust-lang.org/std/os/unix/io/trait.AsRawFd.html /// [`AsRawSocket`]: https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html /// [`socket2`]: https://docs.rs/socket2/ #[cfg_attr(docsrs, doc(alias = "connect_std"))] pub struct TcpSocket { inner: socket2::Socket, } } impl TcpSocket { /// Creates a new socket configured for IPv4. /// /// Calls `socket(2)` with `AF_INET` and `SOCK_STREAM`. /// /// # Returns /// /// On success, the newly created `TcpSocket` is returned. If an error is /// encountered, it is returned instead. /// /// # Examples /// /// Create a new IPv4 socket and start listening. /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// let socket = TcpSocket::new_v4()?; /// socket.bind(addr)?; /// /// let listener = socket.listen(128)?; /// # drop(listener); /// Ok(()) /// } /// ``` pub fn new_v4() -> io::Result { TcpSocket::new(socket2::Domain::IPV4) } /// Creates a new socket configured for IPv6. /// /// Calls `socket(2)` with `AF_INET6` and `SOCK_STREAM`. /// /// # Returns /// /// On success, the newly created `TcpSocket` is returned. If an error is /// encountered, it is returned instead. /// /// # Examples /// /// Create a new IPv6 socket and start listening. /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "[::1]:8080".parse().unwrap(); /// let socket = TcpSocket::new_v6()?; /// socket.bind(addr)?; /// /// let listener = socket.listen(128)?; /// # drop(listener); /// Ok(()) /// } /// ``` pub fn new_v6() -> io::Result { TcpSocket::new(socket2::Domain::IPV6) } fn new(domain: socket2::Domain) -> io::Result { let ty = socket2::Type::STREAM; #[cfg(any( target_os = "android", target_os = "dragonfly", target_os = "freebsd", target_os = "fuchsia", target_os = "illumos", target_os = "linux", target_os = "netbsd", target_os = "openbsd" ))] let ty = ty.nonblocking(); let inner = socket2::Socket::new(domain, ty, Some(socket2::Protocol::TCP))?; #[cfg(not(any( target_os = "android", target_os = "dragonfly", target_os = "freebsd", target_os = "fuchsia", target_os = "illumos", target_os = "linux", target_os = "netbsd", target_os = "openbsd" )))] inner.set_nonblocking(true)?; Ok(TcpSocket { inner }) } /// Sets value for the `SO_KEEPALIVE` option on this socket. pub fn set_keepalive(&self, keepalive: bool) -> io::Result<()> { self.inner.set_keepalive(keepalive) } /// Gets the value of the `SO_KEEPALIVE` option on this socket. pub fn keepalive(&self) -> io::Result { self.inner.keepalive() } /// Allows the socket to bind to an in-use address. /// /// Behavior is platform specific. Refer to the target platform's /// documentation for more details. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.set_reuseaddr(true)?; /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// # drop(listener); /// /// Ok(()) /// } /// ``` pub fn set_reuseaddr(&self, reuseaddr: bool) -> io::Result<()> { self.inner.set_reuse_address(reuseaddr) } /// Retrieves the value set for `SO_REUSEADDR` on this socket. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.set_reuseaddr(true)?; /// assert!(socket.reuseaddr().unwrap()); /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// Ok(()) /// } /// ``` pub fn reuseaddr(&self) -> io::Result { self.inner.reuse_address() } /// Allows the socket to bind to an in-use port. Only available for unix systems /// (excluding Solaris & Illumos). /// /// Behavior is platform specific. Refer to the target platform's /// documentation for more details. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.set_reuseport(true)?; /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// Ok(()) /// } /// ``` #[cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))] #[cfg_attr( docsrs, doc(cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))) )] pub fn set_reuseport(&self, reuseport: bool) -> io::Result<()> { self.inner.set_reuse_port(reuseport) } /// Allows the socket to bind to an in-use port. Only available for unix systems /// (excluding Solaris & Illumos). /// /// Behavior is platform specific. Refer to the target platform's /// documentation for more details. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.set_reuseport(true)?; /// assert!(socket.reuseport().unwrap()); /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// Ok(()) /// } /// ``` #[cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))] #[cfg_attr( docsrs, doc(cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))) )] pub fn reuseport(&self) -> io::Result { self.inner.reuse_port() } /// Sets the size of the TCP send buffer on this socket. /// /// On most operating systems, this sets the `SO_SNDBUF` socket option. pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> { self.inner.set_send_buffer_size(size as usize) } /// Returns the size of the TCP send buffer for this socket. /// /// On most operating systems, this is the value of the `SO_SNDBUF` socket /// option. /// /// Note that if [`set_send_buffer_size`] has been called on this socket /// previously, the value returned by this function may not be the same as /// the argument provided to `set_send_buffer_size`. This is for the /// following reasons: /// /// * Most operating systems have minimum and maximum allowed sizes for the /// send buffer, and will clamp the provided value if it is below the /// minimum or above the maximum. The minimum and maximum buffer sizes are /// OS-dependent. /// * Linux will double the buffer size to account for internal bookkeeping /// data, and returns the doubled value from `getsockopt(2)`. As per `man /// 7 socket`: /// > Sets or gets the maximum socket send buffer in bytes. The /// > kernel doubles this value (to allow space for bookkeeping /// > overhead) when it is set using `setsockopt(2)`, and this doubled /// > value is returned by `getsockopt(2)`. /// /// [`set_send_buffer_size`]: #method.set_send_buffer_size pub fn send_buffer_size(&self) -> io::Result { self.inner.send_buffer_size().map(|n| n as u32) } /// Sets the size of the TCP receive buffer on this socket. /// /// On most operating systems, this sets the `SO_RCVBUF` socket option. pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> { self.inner.set_recv_buffer_size(size as usize) } /// Returns the size of the TCP receive buffer for this socket. /// /// On most operating systems, this is the value of the `SO_RCVBUF` socket /// option. /// /// Note that if [`set_recv_buffer_size`] has been called on this socket /// previously, the value returned by this function may not be the same as /// the argument provided to `set_send_buffer_size`. This is for the /// following reasons: /// /// * Most operating systems have minimum and maximum allowed sizes for the /// receive buffer, and will clamp the provided value if it is below the /// minimum or above the maximum. The minimum and maximum buffer sizes are /// OS-dependent. /// * Linux will double the buffer size to account for internal bookkeeping /// data, and returns the doubled value from `getsockopt(2)`. As per `man /// 7 socket`: /// > Sets or gets the maximum socket send buffer in bytes. The /// > kernel doubles this value (to allow space for bookkeeping /// > overhead) when it is set using `setsockopt(2)`, and this doubled /// > value is returned by `getsockopt(2)`. /// /// [`set_recv_buffer_size`]: #method.set_recv_buffer_size pub fn recv_buffer_size(&self) -> io::Result { self.inner.recv_buffer_size().map(|n| n as u32) } /// Sets the linger duration of this socket by setting the `SO_LINGER` option. /// /// This option controls the action taken when a stream has unsent messages and the stream is /// closed. If `SO_LINGER` is set, the system shall block the process until it can transmit the /// data or until the time expires. /// /// If `SO_LINGER` is not specified, and the socket is closed, the system handles the call in a /// way that allows the process to continue as quickly as possible. pub fn set_linger(&self, dur: Option) -> io::Result<()> { self.inner.set_linger(dur) } /// Reads the linger duration for this socket by getting the `SO_LINGER` /// option. /// /// For more information about this option, see [`set_linger`]. /// /// [`set_linger`]: TcpSocket::set_linger pub fn linger(&self) -> io::Result> { self.inner.linger() } /// Sets the value of the `TCP_NODELAY` option on this socket. /// /// If set, this option disables the Nagle algorithm. This means that segments are always /// sent as soon as possible, even if there is only a small amount of data. When not set, /// data is buffered until there is a sufficient amount to send out, thereby avoiding /// the frequent sending of small packets. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// # async fn dox() -> Result<(), Box> { /// let socket = TcpSocket::new_v4()?; /// /// socket.set_nodelay(true)?; /// # Ok(()) /// # } /// ``` pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> { self.inner.set_nodelay(nodelay) } /// Gets the value of the `TCP_NODELAY` option on this socket. /// /// For more information about this option, see [`set_nodelay`]. /// /// [`set_nodelay`]: TcpSocket::set_nodelay /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// # async fn dox() -> Result<(), Box> { /// let socket = TcpSocket::new_v4()?; /// /// println!("{:?}", socket.nodelay()?); /// # Ok(()) /// # } /// ``` pub fn nodelay(&self) -> io::Result { self.inner.nodelay() } /// Gets the value of the `IP_TOS` option for this socket. /// /// For more information about this option, see [`set_tos`]. /// /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) /// /// [`set_tos`]: Self::set_tos // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1464 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))] #[cfg_attr( docsrs, doc(cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))) )] pub fn tos(&self) -> io::Result { self.inner.tos() } /// Sets the value for the `IP_TOS` option on this socket. /// /// This value sets the type-of-service field that is used in every packet /// sent from this socket. /// /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1446 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))] #[cfg_attr( docsrs, doc(cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))) )] pub fn set_tos(&self, tos: u32) -> io::Result<()> { self.inner.set_tos(tos) } /// Gets the value for the `SO_BINDTODEVICE` option on this socket /// /// This value gets the socket binded device's interface name. #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))] #[cfg_attr( docsrs, doc(cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))) )] pub fn device(&self) -> io::Result>> { self.inner.device() } /// Sets the value for the `SO_BINDTODEVICE` option on this socket /// /// If a socket is bound to an interface, only packets received from that /// particular interface are processed by the socket. Note that this only /// works for some socket types, particularly `AF_INET` sockets. /// /// If `interface` is `None` or an empty string it removes the binding. #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[cfg_attr( docsrs, doc(cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))) )] pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> { self.inner.bind_device(interface) } /// Gets the local address of this socket. /// /// Will fail on windows if called before `bind`. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.bind(addr)?; /// assert_eq!(socket.local_addr().unwrap().to_string(), "127.0.0.1:8080"); /// let listener = socket.listen(1024)?; /// Ok(()) /// } /// ``` pub fn local_addr(&self) -> io::Result { self.inner.local_addr().and_then(convert_address) } /// Returns the value of the `SO_ERROR` option. pub fn take_error(&self) -> io::Result> { self.inner.take_error() } /// Binds the socket to the given address. /// /// This calls the `bind(2)` operating-system function. Behavior is /// platform specific. Refer to the target platform's documentation for more /// details. /// /// # Examples /// /// Bind a socket before listening. /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// # drop(listener); /// /// Ok(()) /// } /// ``` pub fn bind(&self, addr: SocketAddr) -> io::Result<()> { self.inner.bind(&addr.into()) } /// Establishes a TCP connection with a peer at the specified socket address. /// /// The `TcpSocket` is consumed. Once the connection is established, a /// connected [`TcpStream`] is returned. If the connection fails, the /// encountered error is returned. /// /// [`TcpStream`]: TcpStream /// /// This calls the `connect(2)` operating-system function. Behavior is /// platform specific. Refer to the target platform's documentation for more /// details. /// /// # Examples /// /// Connecting to a peer. /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// let stream = socket.connect(addr).await?; /// # drop(stream); /// /// Ok(()) /// } /// ``` pub async fn connect(self, addr: SocketAddr) -> io::Result { if let Err(err) = self.inner.connect(&addr.into()) { #[cfg(unix)] if err.raw_os_error() != Some(libc::EINPROGRESS) { return Err(err); } #[cfg(windows)] if err.kind() != io::ErrorKind::WouldBlock { return Err(err); } } #[cfg(unix)] let mio = { use std::os::unix::io::{FromRawFd, IntoRawFd}; let raw_fd = self.inner.into_raw_fd(); unsafe { mio::net::TcpStream::from_raw_fd(raw_fd) } }; #[cfg(windows)] let mio = { use std::os::windows::io::{FromRawSocket, IntoRawSocket}; let raw_socket = self.inner.into_raw_socket(); unsafe { mio::net::TcpStream::from_raw_socket(raw_socket) } }; TcpStream::connect_mio(mio).await } /// Converts the socket into a `TcpListener`. /// /// `backlog` defines the maximum number of pending connections are queued /// by the operating system at any given time. Connection are removed from /// the queue with [`TcpListener::accept`]. When the queue is full, the /// operating-system will start rejecting connections. /// /// [`TcpListener::accept`]: TcpListener::accept /// /// This calls the `listen(2)` operating-system function, marking the socket /// as a passive socket. Behavior is platform specific. Refer to the target /// platform's documentation for more details. /// /// # Examples /// /// Create a `TcpListener`. /// /// ```no_run /// use tokio::net::TcpSocket; /// /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let addr = "127.0.0.1:8080".parse().unwrap(); /// /// let socket = TcpSocket::new_v4()?; /// socket.bind(addr)?; /// /// let listener = socket.listen(1024)?; /// # drop(listener); /// /// Ok(()) /// } /// ``` pub fn listen(self, backlog: u32) -> io::Result { self.inner.listen(backlog as i32)?; #[cfg(unix)] let mio = { use std::os::unix::io::{FromRawFd, IntoRawFd}; let raw_fd = self.inner.into_raw_fd(); unsafe { mio::net::TcpListener::from_raw_fd(raw_fd) } }; #[cfg(windows)] let mio = { use std::os::windows::io::{FromRawSocket, IntoRawSocket}; let raw_socket = self.inner.into_raw_socket(); unsafe { mio::net::TcpListener::from_raw_socket(raw_socket) } }; TcpListener::new(mio) } /// Converts a [`std::net::TcpStream`] into a `TcpSocket`. The provided /// socket must not have been connected prior to calling this function. This /// function is typically used together with crates such as [`socket2`] to /// configure socket options that are not available on `TcpSocket`. /// /// [`std::net::TcpStream`]: struct@std::net::TcpStream /// [`socket2`]: https://docs.rs/socket2/ /// /// # Notes /// /// The caller is responsible for ensuring that the socket is in /// non-blocking mode. Otherwise all I/O operations on the socket /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::net::TcpStream::set_nonblocking /// /// # Examples /// /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// use tokio::net::TcpSocket; /// use socket2::{Domain, Socket, Type}; /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { /// let socket2_socket = Socket::new(Domain::IPV4, Type::STREAM, None)?; /// socket2_socket.set_nonblocking(true)?; /// /// let socket = TcpSocket::from_std_stream(socket2_socket.into()); /// /// Ok(()) /// } /// ``` pub fn from_std_stream(std_stream: std::net::TcpStream) -> TcpSocket { #[cfg(unix)] { use std::os::unix::io::{FromRawFd, IntoRawFd}; let raw_fd = std_stream.into_raw_fd(); unsafe { TcpSocket::from_raw_fd(raw_fd) } } #[cfg(windows)] { use std::os::windows::io::{FromRawSocket, IntoRawSocket}; let raw_socket = std_stream.into_raw_socket(); unsafe { TcpSocket::from_raw_socket(raw_socket) } } } } fn convert_address(address: socket2::SockAddr) -> io::Result { match address.as_socket() { Some(address) => Ok(address), None => Err(io::Error::new( io::ErrorKind::InvalidInput, "invalid address family (not IPv4 or IPv6)", )), } } impl fmt::Debug for TcpSocket { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.inner.fmt(fmt) } } // These trait implementations can't be build on Windows, so we completely // ignore them, even when building documentation. #[cfg(unix)] cfg_unix! { impl AsRawFd for TcpSocket { fn as_raw_fd(&self) -> RawFd { self.inner.as_raw_fd() } } impl AsFd for TcpSocket { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } impl FromRawFd for TcpSocket { /// Converts a `RawFd` to a `TcpSocket`. /// /// # Notes /// /// The caller is responsible for ensuring that the socket is in /// non-blocking mode. unsafe fn from_raw_fd(fd: RawFd) -> TcpSocket { let inner = socket2::Socket::from_raw_fd(fd); TcpSocket { inner } } } impl IntoRawFd for TcpSocket { fn into_raw_fd(self) -> RawFd { self.inner.into_raw_fd() } } } cfg_windows! { impl IntoRawSocket for TcpSocket { fn into_raw_socket(self) -> RawSocket { self.inner.into_raw_socket() } } impl AsRawSocket for TcpSocket { fn as_raw_socket(&self) -> RawSocket { self.inner.as_raw_socket() } } impl AsSocket for TcpSocket { fn as_socket(&self) -> BorrowedSocket<'_> { unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } } } impl FromRawSocket for TcpSocket { /// Converts a `RawSocket` to a `TcpStream`. /// /// # Notes /// /// The caller is responsible for ensuring that the socket is in /// non-blocking mode. unsafe fn from_raw_socket(socket: RawSocket) -> TcpSocket { let inner = socket2::Socket::from_raw_socket(socket); TcpSocket { inner } } } } tokio-1.43.1/src/net/tcp/split.rs000064400000000000000000000357061046102023000147160ustar 00000000000000//! `TcpStream` split support. //! //! A `TcpStream` can be split into a `ReadHalf` and a //! `WriteHalf` with the `TcpStream::split` method. `ReadHalf` //! implements `AsyncRead` while `WriteHalf` implements `AsyncWrite`. //! //! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized //! split has no associated overhead and enforces all invariants at the type //! level. use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::TcpStream; use std::future::poll_fn; use std::io; use std::net::{Shutdown, SocketAddr}; use std::pin::Pin; use std::task::{Context, Poll}; cfg_io_util! { use bytes::BufMut; } /// Borrowed read half of a [`TcpStream`], created by [`split`]. /// /// Reading from a `ReadHalf` is usually done using the convenience methods found on the /// [`AsyncReadExt`] trait. /// /// [`TcpStream`]: TcpStream /// [`split`]: TcpStream::split() /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt #[derive(Debug)] pub struct ReadHalf<'a>(&'a TcpStream); /// Borrowed write half of a [`TcpStream`], created by [`split`]. /// /// Note that in the [`AsyncWrite`] implementation of this type, [`poll_shutdown`] will /// shut down the TCP stream in the write direction. /// /// Writing to an `WriteHalf` is usually done using the convenience methods found /// on the [`AsyncWriteExt`] trait. /// /// [`TcpStream`]: TcpStream /// [`split`]: TcpStream::split() /// [`AsyncWrite`]: trait@crate::io::AsyncWrite /// [`poll_shutdown`]: fn@crate::io::AsyncWrite::poll_shutdown /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt #[derive(Debug)] pub struct WriteHalf<'a>(&'a TcpStream); pub(crate) fn split(stream: &mut TcpStream) -> (ReadHalf<'_>, WriteHalf<'_>) { (ReadHalf(&*stream), WriteHalf(&*stream)) } impl ReadHalf<'_> { /// Attempts to receive data on the socket, without removing that data from /// the queue, registering the current task for wakeup if data is not yet /// available. /// /// Note that on multiple calls to `poll_peek` or `poll_read`, only the /// `Waker` from the `Context` passed to the most recent call is scheduled /// to receive a wakeup. /// /// See the [`TcpStream::poll_peek`] level documentation for more details. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, ReadBuf}; /// use tokio::net::TcpStream; /// /// use std::future::poll_fn; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut stream = TcpStream::connect("127.0.0.1:8000").await?; /// let (mut read_half, _) = stream.split(); /// let mut buf = [0; 10]; /// let mut buf = ReadBuf::new(&mut buf); /// /// poll_fn(|cx| { /// read_half.poll_peek(cx, &mut buf) /// }).await?; /// /// Ok(()) /// } /// ``` /// /// [`TcpStream::poll_peek`]: TcpStream::poll_peek pub fn poll_peek( &mut self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.0.poll_peek(cx, buf) } /// Receives data on the socket from the remote address to which it is /// connected, without removing that data from the queue. On success, /// returns the number of bytes peeked. /// /// See the [`TcpStream::peek`] level documentation for more details. /// /// [`TcpStream::peek`]: TcpStream::peek /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use tokio::io::AsyncReadExt; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; /// let (mut read_half, _) = stream.split(); /// /// let mut b1 = [0; 10]; /// let mut b2 = [0; 10]; /// /// // Peek at the data /// let n = read_half.peek(&mut b1).await?; /// /// // Read the data /// assert_eq!(n, read_half.read(&mut b2[..n]).await?); /// assert_eq!(&b1[..n], &b2[..n]); /// /// Ok(()) /// } /// ``` /// /// The [`read`] method is defined on the [`AsyncReadExt`] trait. /// /// [`read`]: fn@crate::io::AsyncReadExt::read /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result { let mut buf = ReadBuf::new(buf); poll_fn(|cx| self.poll_peek(cx, &mut buf)).await } /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_read()`]. It can be used instead /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] /// and [`Ready::READ_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`TcpStream::ready`]. /// /// [`try_read()`]: Self::try_read /// [`readable()`]: Self::readable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.0.ready(interest).await } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// This function is also equivalent to [`TcpStream::ready`]. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn readable(&self) -> io::Result<()> { self.0.readable().await } /// Tries to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The stream's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.0.try_read(buf) } /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: Self::try_read() /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.0.try_read_vectored(bufs) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.0.try_read_buf(buf) } } /// Returns the remote address that this stream is connected to. pub fn peer_addr(&self) -> io::Result { self.0.peer_addr() } /// Returns the local address that this stream is bound to. pub fn local_addr(&self) -> io::Result { self.0.local_addr() } } impl WriteHalf<'_> { /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_write()`]. It can be used instead /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] /// and [`Ready::WRITE_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`TcpStream::ready`]. /// /// [`try_write()`]: Self::try_write /// [`writable()`]: Self::writable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.0.ready(interest).await } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn writable(&self) -> io::Result<()> { self.0.writable().await } /// Tries to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write(&self, buf: &[u8]) -> io::Result { self.0.try_write(buf) } /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: Self::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result { self.0.try_write_vectored(bufs) } /// Returns the remote address that this stream is connected to. pub fn peer_addr(&self) -> io::Result { self.0.peer_addr() } /// Returns the local address that this stream is bound to. pub fn local_addr(&self) -> io::Result { self.0.local_addr() } } impl AsyncRead for ReadHalf<'_> { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.0.poll_read_priv(cx, buf) } } impl AsyncWrite for WriteHalf<'_> { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.0.poll_write_priv(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.0.poll_write_vectored_priv(cx, bufs) } fn is_write_vectored(&self) -> bool { self.0.is_write_vectored() } #[inline] fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { // tcp flush is a no-op Poll::Ready(Ok(())) } // `poll_shutdown` on a write half shutdowns the stream in the "write" direction. fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { self.0.shutdown_std(Shutdown::Write).into() } } impl AsRef for ReadHalf<'_> { fn as_ref(&self) -> &TcpStream { self.0 } } impl AsRef for WriteHalf<'_> { fn as_ref(&self) -> &TcpStream { self.0 } } tokio-1.43.1/src/net/tcp/split_owned.rs000064400000000000000000000432301046102023000161010ustar 00000000000000//! `TcpStream` owned split support. //! //! A `TcpStream` can be split into an `OwnedReadHalf` and a `OwnedWriteHalf` //! with the `TcpStream::into_split` method. `OwnedReadHalf` implements //! `AsyncRead` while `OwnedWriteHalf` implements `AsyncWrite`. //! //! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized //! split has no associated overhead and enforces all invariants at the type //! level. use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::TcpStream; use std::error::Error; use std::future::poll_fn; use std::net::{Shutdown, SocketAddr}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use std::{fmt, io}; cfg_io_util! { use bytes::BufMut; } /// Owned read half of a [`TcpStream`], created by [`into_split`]. /// /// Reading from an `OwnedReadHalf` is usually done using the convenience methods found /// on the [`AsyncReadExt`] trait. /// /// [`TcpStream`]: TcpStream /// [`into_split`]: TcpStream::into_split() /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt #[derive(Debug)] pub struct OwnedReadHalf { inner: Arc, } /// Owned write half of a [`TcpStream`], created by [`into_split`]. /// /// Note that in the [`AsyncWrite`] implementation of this type, [`poll_shutdown`] will /// shut down the TCP stream in the write direction. Dropping the write half /// will also shut down the write half of the TCP stream. /// /// Writing to an `OwnedWriteHalf` is usually done using the convenience methods found /// on the [`AsyncWriteExt`] trait. /// /// [`TcpStream`]: TcpStream /// [`into_split`]: TcpStream::into_split() /// [`AsyncWrite`]: trait@crate::io::AsyncWrite /// [`poll_shutdown`]: fn@crate::io::AsyncWrite::poll_shutdown /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt #[derive(Debug)] pub struct OwnedWriteHalf { inner: Arc, shutdown_on_drop: bool, } pub(crate) fn split_owned(stream: TcpStream) -> (OwnedReadHalf, OwnedWriteHalf) { let arc = Arc::new(stream); let read = OwnedReadHalf { inner: Arc::clone(&arc), }; let write = OwnedWriteHalf { inner: arc, shutdown_on_drop: true, }; (read, write) } pub(crate) fn reunite( read: OwnedReadHalf, write: OwnedWriteHalf, ) -> Result { if Arc::ptr_eq(&read.inner, &write.inner) { write.forget(); // This unwrap cannot fail as the api does not allow creating more than two Arcs, // and we just dropped the other half. Ok(Arc::try_unwrap(read.inner).expect("TcpStream: try_unwrap failed in reunite")) } else { Err(ReuniteError(read, write)) } } /// Error indicating that two halves were not from the same socket, and thus could /// not be reunited. #[derive(Debug)] pub struct ReuniteError(pub OwnedReadHalf, pub OwnedWriteHalf); impl fmt::Display for ReuniteError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, "tried to reunite halves that are not from the same socket" ) } } impl Error for ReuniteError {} impl OwnedReadHalf { /// Attempts to put the two halves of a `TcpStream` back together and /// recover the original socket. Succeeds only if the two halves /// originated from the same call to [`into_split`]. /// /// [`into_split`]: TcpStream::into_split() pub fn reunite(self, other: OwnedWriteHalf) -> Result { reunite(self, other) } /// Attempt to receive data on the socket, without removing that data from /// the queue, registering the current task for wakeup if data is not yet /// available. /// /// Note that on multiple calls to `poll_peek` or `poll_read`, only the /// `Waker` from the `Context` passed to the most recent call is scheduled /// to receive a wakeup. /// /// See the [`TcpStream::poll_peek`] level documentation for more details. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, ReadBuf}; /// use tokio::net::TcpStream; /// /// use std::future::poll_fn; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let stream = TcpStream::connect("127.0.0.1:8000").await?; /// let (mut read_half, _) = stream.into_split(); /// let mut buf = [0; 10]; /// let mut buf = ReadBuf::new(&mut buf); /// /// poll_fn(|cx| { /// read_half.poll_peek(cx, &mut buf) /// }).await?; /// /// Ok(()) /// } /// ``` /// /// [`TcpStream::poll_peek`]: TcpStream::poll_peek pub fn poll_peek( &mut self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.inner.poll_peek(cx, buf) } /// Receives data on the socket from the remote address to which it is /// connected, without removing that data from the queue. On success, /// returns the number of bytes peeked. /// /// See the [`TcpStream::peek`] level documentation for more details. /// /// [`TcpStream::peek`]: TcpStream::peek /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use tokio::io::AsyncReadExt; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// let (mut read_half, _) = stream.into_split(); /// /// let mut b1 = [0; 10]; /// let mut b2 = [0; 10]; /// /// // Peek at the data /// let n = read_half.peek(&mut b1).await?; /// /// // Read the data /// assert_eq!(n, read_half.read(&mut b2[..n]).await?); /// assert_eq!(&b1[..n], &b2[..n]); /// /// Ok(()) /// } /// ``` /// /// The [`read`] method is defined on the [`AsyncReadExt`] trait. /// /// [`read`]: fn@crate::io::AsyncReadExt::read /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result { let mut buf = ReadBuf::new(buf); poll_fn(|cx| self.poll_peek(cx, &mut buf)).await } /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_read()`]. It can be used instead /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] /// and [`Ready::READ_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`TcpStream::ready`]. /// /// [`try_read()`]: Self::try_read /// [`readable()`]: Self::readable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.inner.ready(interest).await } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// This function is also equivalent to [`TcpStream::ready`]. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn readable(&self) -> io::Result<()> { self.inner.readable().await } /// Tries to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The stream's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.inner.try_read(buf) } /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: Self::try_read() /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.inner.try_read_vectored(bufs) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.inner.try_read_buf(buf) } } /// Returns the remote address that this stream is connected to. pub fn peer_addr(&self) -> io::Result { self.inner.peer_addr() } /// Returns the local address that this stream is bound to. pub fn local_addr(&self) -> io::Result { self.inner.local_addr() } } impl AsyncRead for OwnedReadHalf { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.inner.poll_read_priv(cx, buf) } } impl OwnedWriteHalf { /// Attempts to put the two halves of a `TcpStream` back together and /// recover the original socket. Succeeds only if the two halves /// originated from the same call to [`into_split`]. /// /// [`into_split`]: TcpStream::into_split() pub fn reunite(self, other: OwnedReadHalf) -> Result { reunite(other, self) } /// Destroys the write half, but don't close the write half of the stream /// until the read half is dropped. If the read half has already been /// dropped, this closes the stream. pub fn forget(mut self) { self.shutdown_on_drop = false; drop(self); } /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_write()`]. It can be used instead /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] /// and [`Ready::WRITE_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`TcpStream::ready`]. /// /// [`try_write()`]: Self::try_write /// [`writable()`]: Self::writable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.inner.ready(interest).await } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn writable(&self) -> io::Result<()> { self.inner.writable().await } /// Tries to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write(&self, buf: &[u8]) -> io::Result { self.inner.try_write(buf) } /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: Self::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result { self.inner.try_write_vectored(bufs) } /// Returns the remote address that this stream is connected to. pub fn peer_addr(&self) -> io::Result { self.inner.peer_addr() } /// Returns the local address that this stream is bound to. pub fn local_addr(&self) -> io::Result { self.inner.local_addr() } } impl Drop for OwnedWriteHalf { fn drop(&mut self) { if self.shutdown_on_drop { let _ = self.inner.shutdown_std(Shutdown::Write); } } } impl AsyncWrite for OwnedWriteHalf { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.inner.poll_write_priv(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.inner.poll_write_vectored_priv(cx, bufs) } fn is_write_vectored(&self) -> bool { self.inner.is_write_vectored() } #[inline] fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { // tcp flush is a no-op Poll::Ready(Ok(())) } // `poll_shutdown` on a write half shutdowns the stream in the "write" direction. fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { let res = self.inner.shutdown_std(Shutdown::Write); if res.is_ok() { Pin::into_inner(self).shutdown_on_drop = false; } res.into() } } impl AsRef for OwnedReadHalf { fn as_ref(&self) -> &TcpStream { &self.inner } } impl AsRef for OwnedWriteHalf { fn as_ref(&self) -> &TcpStream { &self.inner } } tokio-1.43.1/src/net/tcp/stream.rs000064400000000000000000001417321046102023000150530ustar 00000000000000cfg_not_wasi! { use crate::net::{to_socket_addrs, ToSocketAddrs}; use std::future::poll_fn; use std::time::Duration; } use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; use crate::net::tcp::split::{split, ReadHalf, WriteHalf}; use crate::net::tcp::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf}; use std::fmt; use std::io; use std::net::{Shutdown, SocketAddr}; use std::pin::Pin; use std::task::{ready, Context, Poll}; cfg_io_util! { use bytes::BufMut; } cfg_net! { /// A TCP stream between a local and a remote socket. /// /// A TCP stream can either be created by connecting to an endpoint, via the /// [`connect`] method, or by [accepting] a connection from a [listener]. A /// TCP stream can also be created via the [`TcpSocket`] type. /// /// Reading and writing to a `TcpStream` is usually done using the /// convenience methods found on the [`AsyncReadExt`] and [`AsyncWriteExt`] /// traits. /// /// [`connect`]: method@TcpStream::connect /// [accepting]: method@crate::net::TcpListener::accept /// [listener]: struct@crate::net::TcpListener /// [`TcpSocket`]: struct@crate::net::TcpSocket /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use tokio::io::AsyncWriteExt; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// // Write some data. /// stream.write_all(b"hello world!").await?; /// /// Ok(()) /// } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt /// /// To shut down the stream in the write direction, you can call the /// [`shutdown()`] method. This will cause the other peer to receive a read of /// length 0, indicating that no more data will be sent. This only closes /// the stream in one direction. /// /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown pub struct TcpStream { io: PollEvented, } } impl TcpStream { cfg_not_wasi! { /// Opens a TCP connection to a remote host. /// /// `addr` is an address of the remote host. Anything which implements the /// [`ToSocketAddrs`] trait can be supplied as the address. If `addr` /// yields multiple addresses, connect will be attempted with each of the /// addresses until a connection is successful. If none of the addresses /// result in a successful connection, the error returned from the last /// connection attempt (the last address) is returned. /// /// To configure the socket before connecting, you can use the [`TcpSocket`] /// type. /// /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs /// [`TcpSocket`]: struct@crate::net::TcpSocket /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use tokio::io::AsyncWriteExt; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// // Write some data. /// stream.write_all(b"hello world!").await?; /// /// Ok(()) /// } /// ``` /// /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. /// /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt pub async fn connect(addr: A) -> io::Result { let addrs = to_socket_addrs(addr).await?; let mut last_err = None; for addr in addrs { match TcpStream::connect_addr(addr).await { Ok(stream) => return Ok(stream), Err(e) => last_err = Some(e), } } Err(last_err.unwrap_or_else(|| { io::Error::new( io::ErrorKind::InvalidInput, "could not resolve to any address", ) })) } /// Establishes a connection to the specified `addr`. async fn connect_addr(addr: SocketAddr) -> io::Result { let sys = mio::net::TcpStream::connect(addr)?; TcpStream::connect_mio(sys).await } pub(crate) async fn connect_mio(sys: mio::net::TcpStream) -> io::Result { let stream = TcpStream::new(sys)?; // Once we've connected, wait for the stream to be writable as // that's when the actual connection has been initiated. Once we're // writable we check for `take_socket_error` to see if the connect // actually hit an error or not. // // If all that succeeded then we ship everything on up. poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?; if let Some(e) = stream.io.take_error()? { return Err(e); } Ok(stream) } } pub(crate) fn new(connected: mio::net::TcpStream) -> io::Result { let io = PollEvented::new(connected)?; Ok(TcpStream { io }) } /// Creates new `TcpStream` from a `std::net::TcpStream`. /// /// This function is intended to be used to wrap a TCP stream from the /// standard library in the Tokio equivalent. /// /// # Notes /// /// The caller is responsible for ensuring that the stream is in /// non-blocking mode. Otherwise all I/O operations on the stream /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::net::TcpStream::set_nonblocking /// /// # Examples /// /// ```rust,no_run /// use std::error::Error; /// use tokio::net::TcpStream; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let std_stream = std::net::TcpStream::connect("127.0.0.1:34254")?; /// std_stream.set_nonblocking(true)?; /// let stream = TcpStream::from_std(std_stream)?; /// Ok(()) /// } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. #[track_caller] pub fn from_std(stream: std::net::TcpStream) -> io::Result { let io = mio::net::TcpStream::from_std(stream); let io = PollEvented::new(io)?; Ok(TcpStream { io }) } /// Turns a [`tokio::net::TcpStream`] into a [`std::net::TcpStream`]. /// /// The returned [`std::net::TcpStream`] will have nonblocking mode set as `true`. /// Use [`set_nonblocking`] to change the blocking mode if needed. /// /// # Examples /// /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// use std::error::Error; /// use std::io::Read; /// use tokio::net::TcpListener; /// # use tokio::net::TcpStream; /// # use tokio::io::AsyncWriteExt; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let mut data = [0u8; 12]; /// # if false { /// let listener = TcpListener::bind("127.0.0.1:34254").await?; /// # } /// # let listener = TcpListener::bind("127.0.0.1:0").await?; /// # let addr = listener.local_addr().unwrap(); /// # let handle = tokio::spawn(async move { /// # let mut stream: TcpStream = TcpStream::connect(addr).await.unwrap(); /// # stream.write_all(b"Hello world!").await.unwrap(); /// # }); /// let (tokio_tcp_stream, _) = listener.accept().await?; /// let mut std_tcp_stream = tokio_tcp_stream.into_std()?; /// # handle.await.expect("The task being joined has panicked"); /// std_tcp_stream.set_nonblocking(false)?; /// std_tcp_stream.read_exact(&mut data)?; /// # assert_eq!(b"Hello world!", &data); /// Ok(()) /// } /// ``` /// [`tokio::net::TcpStream`]: TcpStream /// [`std::net::TcpStream`]: std::net::TcpStream /// [`set_nonblocking`]: fn@std::net::TcpStream::set_nonblocking pub fn into_std(self) -> io::Result { #[cfg(unix)] { use std::os::unix::io::{FromRawFd, IntoRawFd}; self.io .into_inner() .map(IntoRawFd::into_raw_fd) .map(|raw_fd| unsafe { std::net::TcpStream::from_raw_fd(raw_fd) }) } #[cfg(windows)] { use std::os::windows::io::{FromRawSocket, IntoRawSocket}; self.io .into_inner() .map(|io| io.into_raw_socket()) .map(|raw_socket| unsafe { std::net::TcpStream::from_raw_socket(raw_socket) }) } #[cfg(target_os = "wasi")] { use std::os::wasi::io::{FromRawFd, IntoRawFd}; self.io .into_inner() .map(|io| io.into_raw_fd()) .map(|raw_fd| unsafe { std::net::TcpStream::from_raw_fd(raw_fd) }) } } /// Returns the local address that this stream is bound to. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// println!("{:?}", stream.local_addr()?); /// # Ok(()) /// # } /// ``` pub fn local_addr(&self) -> io::Result { self.io.local_addr() } /// Returns the value of the `SO_ERROR` option. pub fn take_error(&self) -> io::Result> { self.io.take_error() } /// Returns the remote address that this stream is connected to. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// println!("{:?}", stream.peer_addr()?); /// # Ok(()) /// # } /// ``` pub fn peer_addr(&self) -> io::Result { self.io.peer_addr() } /// Attempts to receive data on the socket, without removing that data from /// the queue, registering the current task for wakeup if data is not yet /// available. /// /// Note that on multiple calls to `poll_peek`, `poll_read` or /// `poll_read_ready`, only the `Waker` from the `Context` passed to the /// most recent call is scheduled to receive a wakeup. (However, /// `poll_write` retains a second, independent waker.) /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if data is not yet available. /// * `Poll::Ready(Ok(n))` if data is available. `n` is the number of bytes peeked. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// # Examples /// /// ```no_run /// use tokio::io::{self, ReadBuf}; /// use tokio::net::TcpStream; /// /// use std::future::poll_fn; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let stream = TcpStream::connect("127.0.0.1:8000").await?; /// let mut buf = [0; 10]; /// let mut buf = ReadBuf::new(&mut buf); /// /// poll_fn(|cx| { /// stream.poll_peek(cx, &mut buf) /// }).await?; /// /// Ok(()) /// } /// ``` pub fn poll_peek( &self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { loop { let ev = ready!(self.io.registration().poll_read_ready(cx))?; let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; match self.io.peek(b) { Ok(ret) => { unsafe { buf.assume_init(ret) }; buf.advance(ret); return Poll::Ready(Ok(ret)); } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.io.registration().clear_readiness(ev); } Err(e) => return Poll::Ready(Err(e)), } } } /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same socket on a single /// task without splitting the socket. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// Concurrently read and write to the stream on the same task without /// splitting. /// /// ```no_run /// use tokio::io::Interest; /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// loop { /// let ready = stream.ready(Interest::READABLE | Interest::WRITABLE).await?; /// /// if ready.is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// /// } /// /// if ready.is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// let mut msg = vec![0; 1024]; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read(&mut msg) { /// Ok(n) => { /// msg.truncate(n); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// println!("GOT = {:?}", msg); /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read readiness. /// /// If the tcp stream is not currently ready for reading, this method will /// store a clone of the `Waker` from the provided `Context`. When the tcp /// stream becomes ready for reading, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_read_ready`, `poll_read` or /// `poll_peek`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. (However, /// `poll_write_ready` retains a second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the tcp stream is not ready for reading. /// * `Poll::Ready(Ok(()))` if the tcp stream is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`readable`]: method@Self::readable pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Tries to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: TcpStream::readable() /// [`ready()`]: TcpStream::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The stream's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf = [0; 4096]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read(&self, buf: &mut [u8]) -> io::Result { use std::io::Read; self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read(buf)) } /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: TcpStream::try_read() /// [`readable()`]: TcpStream::readable() /// [`ready()`]: TcpStream::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io::{self, IoSliceMut}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf_a = [0; 512]; /// let mut buf_b = [0; 1024]; /// let mut bufs = [ /// IoSliceMut::new(&mut buf_a), /// IoSliceMut::new(&mut buf_b), /// ]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read_vectored(&mut bufs) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { use std::io::Read; self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: TcpStream::readable() /// [`ready()`]: TcpStream::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// let mut buf = Vec::with_capacity(4096); /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read_buf(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { use std::io::Read; let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the // buffer. let n = (&*self.io).read(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }) } } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be writable /// stream.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write readiness. /// /// If the tcp stream is not currently ready for writing, this method will /// store a clone of the `Waker` from the provided `Context`. When the tcp /// stream becomes ready for writing, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_read_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the tcp stream is not ready for writing. /// * `Poll::Ready(Ok(()))` if the tcp stream is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`writable`]: method@Self::writable pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Try to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be writable /// stream.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write(&self, buf: &[u8]) -> io::Result { use std::io::Write; self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possibly being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: TcpStream::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; /// /// loop { /// // Wait for the socket to be writable /// stream.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write_vectored(&bufs) { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result { use std::io::Write; self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(bufs)) } /// Tries to read or write from the socket using a user-provided IO operation. /// /// If the socket is ready, the provided closure is called. The closure /// should attempt to perform IO operation on the socket by manually /// calling the appropriate syscall. If the operation fails because the /// socket is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the socket is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `TcpStream` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: TcpStream::readable() /// [`writable()`]: TcpStream::writable() /// [`ready()`]: TcpStream::ready() pub fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { self.io .registration() .try_io(interest, || self.io.try_io(f)) } /// Reads or writes from the socket using a user-provided IO operation. /// /// The readiness of the socket is awaited and when the socket is ready, /// the provided closure is called. The closure should attempt to perform /// IO operation on the socket by manually calling the appropriate syscall. /// If the operation fails because the socket is not actually ready, /// then the closure should return a `WouldBlock` error. In such case the /// readiness flag is cleared and the socket readiness is awaited again. /// This loop is repeated until the closure returns an `Ok` or an error /// other than `WouldBlock`. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `TcpStream` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub async fn async_io( &self, interest: Interest, mut f: impl FnMut() -> io::Result, ) -> io::Result { self.io .registration() .async_io(interest, || self.io.try_io(&mut f)) .await } /// Receives data on the socket from the remote address to which it is /// connected, without removing that data from the queue. On success, /// returns the number of bytes peeked. /// /// Successive calls return the same data. This is accomplished by passing /// `MSG_PEEK` as a flag to the underlying `recv` system call. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// use tokio::io::AsyncReadExt; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// let mut b1 = [0; 10]; /// let mut b2 = [0; 10]; /// /// // Peek at the data /// let n = stream.peek(&mut b1).await?; /// /// // Read the data /// assert_eq!(n, stream.read(&mut b2[..n]).await?); /// assert_eq!(&b1[..n], &b2[..n]); /// /// Ok(()) /// } /// ``` /// /// The [`read`] method is defined on the [`AsyncReadExt`] trait. /// /// [`read`]: fn@crate::io::AsyncReadExt::read /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt pub async fn peek(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .async_io(Interest::READABLE, || self.io.peek(buf)) .await } /// Shuts down the read, write, or both halves of this connection. /// /// This function will cause all pending and future I/O on the specified /// portions to return immediately with an appropriate value (see the /// documentation of `Shutdown`). pub(super) fn shutdown_std(&self, how: Shutdown) -> io::Result<()> { self.io.shutdown(how) } /// Gets the value of the `TCP_NODELAY` option on this socket. /// /// For more information about this option, see [`set_nodelay`]. /// /// [`set_nodelay`]: TcpStream::set_nodelay /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// println!("{:?}", stream.nodelay()?); /// # Ok(()) /// # } /// ``` pub fn nodelay(&self) -> io::Result { self.io.nodelay() } /// Sets the value of the `TCP_NODELAY` option on this socket. /// /// If set, this option disables the Nagle algorithm. This means that /// segments are always sent as soon as possible, even if there is only a /// small amount of data. When not set, data is buffered until there is a /// sufficient amount to send out, thereby avoiding the frequent sending of /// small packets. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// stream.set_nodelay(true)?; /// # Ok(()) /// # } /// ``` pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> { self.io.set_nodelay(nodelay) } cfg_not_wasi! { /// Reads the linger duration for this socket by getting the `SO_LINGER` /// option. /// /// For more information about this option, see [`set_linger`]. /// /// [`set_linger`]: TcpStream::set_linger /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// println!("{:?}", stream.linger()?); /// # Ok(()) /// # } /// ``` pub fn linger(&self) -> io::Result> { socket2::SockRef::from(self).linger() } /// Sets the linger duration of this socket by setting the `SO_LINGER` option. /// /// This option controls the action taken when a stream has unsent messages and the stream is /// closed. If `SO_LINGER` is set, the system shall block the process until it can transmit the /// data or until the time expires. /// /// If `SO_LINGER` is not specified, and the stream is closed, the system handles the call in a /// way that allows the process to continue as quickly as possible. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// stream.set_linger(None)?; /// # Ok(()) /// # } /// ``` pub fn set_linger(&self, dur: Option) -> io::Result<()> { socket2::SockRef::from(self).set_linger(dur) } } /// Gets the value of the `IP_TTL` option for this socket. /// /// For more information about this option, see [`set_ttl`]. /// /// [`set_ttl`]: TcpStream::set_ttl /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// println!("{:?}", stream.ttl()?); /// # Ok(()) /// # } /// ``` pub fn ttl(&self) -> io::Result { self.io.ttl() } /// Sets the value for the `IP_TTL` option on this socket. /// /// This value sets the time-to-live field that is used in every packet sent /// from this socket. /// /// # Examples /// /// ```no_run /// use tokio::net::TcpStream; /// /// # async fn dox() -> Result<(), Box> { /// let stream = TcpStream::connect("127.0.0.1:8080").await?; /// /// stream.set_ttl(123)?; /// # Ok(()) /// # } /// ``` pub fn set_ttl(&self, ttl: u32) -> io::Result<()> { self.io.set_ttl(ttl) } // These lifetime markers also appear in the generated documentation, and make // it more clear that this is a *borrowed* split. #[allow(clippy::needless_lifetimes)] /// Splits a `TcpStream` into a read half and a write half, which can be used /// to read and write the stream concurrently. /// /// This method is more efficient than [`into_split`], but the halves cannot be /// moved into independently spawned tasks. /// /// [`into_split`]: TcpStream::into_split() pub fn split<'a>(&'a mut self) -> (ReadHalf<'a>, WriteHalf<'a>) { split(self) } /// Splits a `TcpStream` into a read half and a write half, which can be used /// to read and write the stream concurrently. /// /// Unlike [`split`], the owned halves can be moved to separate tasks, however /// this comes at the cost of a heap allocation. /// /// **Note:** Dropping the write half will shut down the write half of the TCP /// stream. This is equivalent to calling [`shutdown()`] on the `TcpStream`. /// /// [`split`]: TcpStream::split() /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown pub fn into_split(self) -> (OwnedReadHalf, OwnedWriteHalf) { split_owned(self) } // == Poll IO functions that takes `&self` == // // To read or write without mutable access to the `TcpStream`, combine the // `poll_read_ready` or `poll_write_ready` methods with the `try_read` or // `try_write` methods. pub(crate) fn poll_read_priv( &self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { // Safety: `TcpStream::read` correctly handles reads into uninitialized memory unsafe { self.io.poll_read(cx, buf) } } pub(super) fn poll_write_priv( &self, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.io.poll_write(cx, buf) } pub(super) fn poll_write_vectored_priv( &self, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.io.poll_write_vectored(cx, bufs) } } impl TryFrom for TcpStream { type Error = io::Error; /// Consumes stream, returning the tokio I/O object. /// /// This is equivalent to /// [`TcpStream::from_std(stream)`](TcpStream::from_std). fn try_from(stream: std::net::TcpStream) -> Result { Self::from_std(stream) } } // ===== impl Read / Write ===== impl AsyncRead for TcpStream { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.poll_read_priv(cx, buf) } } impl AsyncWrite for TcpStream { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.poll_write_priv(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.poll_write_vectored_priv(cx, bufs) } fn is_write_vectored(&self) -> bool { true } #[inline] fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { // tcp flush is a no-op Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { self.shutdown_std(std::net::Shutdown::Write)?; Poll::Ready(Ok(())) } } impl fmt::Debug for TcpStream { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.io.fmt(f) } } #[cfg(unix)] mod sys { use super::TcpStream; use std::os::unix::prelude::*; impl AsRawFd for TcpStream { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for TcpStream { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } cfg_windows! { use crate::os::windows::io::{AsRawSocket, RawSocket, AsSocket, BorrowedSocket}; impl AsRawSocket for TcpStream { fn as_raw_socket(&self) -> RawSocket { self.io.as_raw_socket() } } impl AsSocket for TcpStream { fn as_socket(&self) -> BorrowedSocket<'_> { unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } } } } #[cfg(all(tokio_unstable, target_os = "wasi"))] mod sys { use super::TcpStream; use std::os::wasi::prelude::*; impl AsRawFd for TcpStream { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for TcpStream { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } tokio-1.43.1/src/net/udp.rs000064400000000000000000002441441046102023000135630ustar 00000000000000use crate::io::{Interest, PollEvented, ReadBuf, Ready}; use crate::net::{to_socket_addrs, ToSocketAddrs}; use std::fmt; use std::io; use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::task::{ready, Context, Poll}; cfg_io_util! { use bytes::BufMut; } cfg_net! { /// A UDP socket. /// /// UDP is "connectionless", unlike TCP. Meaning, regardless of what address you've bound to, a `UdpSocket` /// is free to communicate with many different remotes. In tokio there are basically two main ways to use `UdpSocket`: /// /// * one to many: [`bind`](`UdpSocket::bind`) and use [`send_to`](`UdpSocket::send_to`) /// and [`recv_from`](`UdpSocket::recv_from`) to communicate with many different addresses /// * one to one: [`connect`](`UdpSocket::connect`) and associate with a single address, using [`send`](`UdpSocket::send`) /// and [`recv`](`UdpSocket::recv`) to communicate only with that remote address /// /// This type does not provide a `split` method, because this functionality /// can be achieved by instead wrapping the socket in an [`Arc`]. Note that /// you do not need a `Mutex` to share the `UdpSocket` — an `Arc` /// is enough. This is because all of the methods take `&self` instead of /// `&mut self`. Once you have wrapped it in an `Arc`, you can call /// `.clone()` on the `Arc` to get multiple shared handles to the /// same socket. An example of such usage can be found further down. /// /// [`Arc`]: std::sync::Arc /// /// # Streams /// /// If you need to listen over UDP and produce a [`Stream`], you can look /// at [`UdpFramed`]. /// /// [`UdpFramed`]: https://docs.rs/tokio-util/latest/tokio_util/udp/struct.UdpFramed.html /// [`Stream`]: https://docs.rs/futures/0.3/futures/stream/trait.Stream.html /// /// # Example: one to many (bind) /// /// Using `bind` we can create a simple echo server that sends and recv's with many different clients: /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let sock = UdpSocket::bind("0.0.0.0:8080").await?; /// let mut buf = [0; 1024]; /// loop { /// let (len, addr) = sock.recv_from(&mut buf).await?; /// println!("{:?} bytes received from {:?}", len, addr); /// /// let len = sock.send_to(&buf[..len], addr).await?; /// println!("{:?} bytes sent", len); /// } /// } /// ``` /// /// # Example: one to one (connect) /// /// Or using `connect` we can echo with a single remote address using `send` and `recv`: /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let sock = UdpSocket::bind("0.0.0.0:8080").await?; /// /// let remote_addr = "127.0.0.1:59611"; /// sock.connect(remote_addr).await?; /// let mut buf = [0; 1024]; /// loop { /// let len = sock.recv(&mut buf).await?; /// println!("{:?} bytes received from {:?}", len, remote_addr); /// /// let len = sock.send(&buf[..len]).await?; /// println!("{:?} bytes sent", len); /// } /// } /// ``` /// /// # Example: Splitting with `Arc` /// /// Because `send_to` and `recv_from` take `&self`. It's perfectly alright /// to use an `Arc` and share the references to multiple tasks. /// Here is a similar "echo" example that supports concurrent /// sending/receiving: /// /// ```no_run /// use tokio::{net::UdpSocket, sync::mpsc}; /// use std::{io, net::SocketAddr, sync::Arc}; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let sock = UdpSocket::bind("0.0.0.0:8080".parse::().unwrap()).await?; /// let r = Arc::new(sock); /// let s = r.clone(); /// let (tx, mut rx) = mpsc::channel::<(Vec, SocketAddr)>(1_000); /// /// tokio::spawn(async move { /// while let Some((bytes, addr)) = rx.recv().await { /// let len = s.send_to(&bytes, &addr).await.unwrap(); /// println!("{:?} bytes sent", len); /// } /// }); /// /// let mut buf = [0; 1024]; /// loop { /// let (len, addr) = r.recv_from(&mut buf).await?; /// println!("{:?} bytes received from {:?}", len, addr); /// tx.send((buf[..len].to_vec(), addr)).await.unwrap(); /// } /// } /// ``` /// pub struct UdpSocket { io: PollEvented, } } impl UdpSocket { /// This function will create a new UDP socket and attempt to bind it to /// the `addr` provided. /// /// Binding with a port number of 0 will request that the OS assigns a port /// to this listener. The port allocated can be queried via the `local_addr` /// method. /// /// # Example /// /// ```no_run /// # if cfg!(miri) { return } // No `socket` in miri. /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let sock = UdpSocket::bind("0.0.0.0:8080").await?; /// // use `sock` /// # let _ = sock; /// Ok(()) /// } /// ``` pub async fn bind(addr: A) -> io::Result { let addrs = to_socket_addrs(addr).await?; let mut last_err = None; for addr in addrs { match UdpSocket::bind_addr(addr) { Ok(socket) => return Ok(socket), Err(e) => last_err = Some(e), } } Err(last_err.unwrap_or_else(|| { io::Error::new( io::ErrorKind::InvalidInput, "could not resolve to any address", ) })) } fn bind_addr(addr: SocketAddr) -> io::Result { let sys = mio::net::UdpSocket::bind(addr)?; UdpSocket::new(sys) } #[track_caller] fn new(socket: mio::net::UdpSocket) -> io::Result { let io = PollEvented::new(socket)?; Ok(UdpSocket { io }) } /// Creates new `UdpSocket` from a previously bound `std::net::UdpSocket`. /// /// This function is intended to be used to wrap a UDP socket from the /// standard library in the Tokio equivalent. /// /// This can be used in conjunction with `socket2`'s `Socket` interface to /// configure a socket before it's handed off, such as setting options like /// `reuse_address` or binding to multiple addresses. /// /// # Notes /// /// The caller is responsible for ensuring that the socket is in /// non-blocking mode. Otherwise all I/O operations on the socket /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::net::UdpSocket::set_nonblocking /// /// # Panics /// /// This function panics if thread-local runtime is not set. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. /// /// # Example /// /// ```no_run /// use tokio::net::UdpSocket; /// # use std::{io, net::SocketAddr}; /// /// # #[tokio::main] /// # async fn main() -> io::Result<()> { /// let addr = "0.0.0.0:8080".parse::().unwrap(); /// let std_sock = std::net::UdpSocket::bind(addr)?; /// std_sock.set_nonblocking(true)?; /// let sock = UdpSocket::from_std(std_sock)?; /// // use `sock` /// # Ok(()) /// # } /// ``` #[track_caller] pub fn from_std(socket: net::UdpSocket) -> io::Result { let io = mio::net::UdpSocket::from_std(socket); UdpSocket::new(io) } /// Turns a [`tokio::net::UdpSocket`] into a [`std::net::UdpSocket`]. /// /// The returned [`std::net::UdpSocket`] will have nonblocking mode set as /// `true`. Use [`set_nonblocking`] to change the blocking mode if needed. /// /// # Examples /// /// ```rust,no_run /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let tokio_socket = tokio::net::UdpSocket::bind("127.0.0.1:0").await?; /// let std_socket = tokio_socket.into_std()?; /// std_socket.set_nonblocking(false)?; /// Ok(()) /// } /// ``` /// /// [`tokio::net::UdpSocket`]: UdpSocket /// [`std::net::UdpSocket`]: std::net::UdpSocket /// [`set_nonblocking`]: fn@std::net::UdpSocket::set_nonblocking pub fn into_std(self) -> io::Result { #[cfg(unix)] { use std::os::unix::io::{FromRawFd, IntoRawFd}; self.io .into_inner() .map(IntoRawFd::into_raw_fd) .map(|raw_fd| unsafe { std::net::UdpSocket::from_raw_fd(raw_fd) }) } #[cfg(windows)] { use std::os::windows::io::{FromRawSocket, IntoRawSocket}; self.io .into_inner() .map(|io| io.into_raw_socket()) .map(|raw_socket| unsafe { std::net::UdpSocket::from_raw_socket(raw_socket) }) } } fn as_socket(&self) -> socket2::SockRef<'_> { socket2::SockRef::from(self) } /// Returns the local address that this socket is bound to. /// /// # Example /// /// ```no_run /// use tokio::net::UdpSocket; /// # use std::{io, net::SocketAddr}; /// /// # #[tokio::main] /// # async fn main() -> io::Result<()> { /// let addr = "0.0.0.0:8080".parse::().unwrap(); /// let sock = UdpSocket::bind(addr).await?; /// // the address the socket is bound to /// let local_addr = sock.local_addr()?; /// # Ok(()) /// # } /// ``` pub fn local_addr(&self) -> io::Result { self.io.local_addr() } /// Returns the socket address of the remote peer this socket was connected to. /// /// # Example /// /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// use tokio::net::UdpSocket; /// /// # use std::{io, net::SocketAddr}; /// # #[tokio::main] /// # async fn main() -> io::Result<()> { /// let addr = "0.0.0.0:8080".parse::().unwrap(); /// let peer = "127.0.0.1:11100".parse::().unwrap(); /// let sock = UdpSocket::bind(addr).await?; /// sock.connect(peer).await?; /// assert_eq!(peer, sock.peer_addr()?); /// # Ok(()) /// # } /// ``` pub fn peer_addr(&self) -> io::Result { self.io.peer_addr() } /// Connects the UDP socket setting the default destination for send() and /// limiting packets that are read via `recv` from the address specified in /// `addr`. /// /// # Example /// /// ```no_run /// use tokio::net::UdpSocket; /// # use std::{io, net::SocketAddr}; /// /// # #[tokio::main] /// # async fn main() -> io::Result<()> { /// let sock = UdpSocket::bind("0.0.0.0:8080".parse::().unwrap()).await?; /// /// let remote_addr = "127.0.0.1:59600".parse::().unwrap(); /// sock.connect(remote_addr).await?; /// let mut buf = [0u8; 32]; /// // recv from remote_addr /// let len = sock.recv(&mut buf).await?; /// // send to remote_addr /// let _len = sock.send(&buf[..len]).await?; /// # Ok(()) /// # } /// ``` pub async fn connect(&self, addr: A) -> io::Result<()> { let addrs = to_socket_addrs(addr).await?; let mut last_err = None; for addr in addrs { match self.io.connect(addr) { Ok(()) => return Ok(()), Err(e) => last_err = Some(e), } } Err(last_err.unwrap_or_else(|| { io::Error::new( io::ErrorKind::InvalidInput, "could not resolve to any address", ) })) } /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_recv()` or `try_send()`. It /// can be used to concurrently `recv` / `send` to the same socket on a single /// task without splitting the socket. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// Concurrently receive from and send to the socket on the same task /// without splitting. /// /// ```no_run /// use tokio::io::{self, Interest}; /// use tokio::net::UdpSocket; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// loop { /// let ready = socket.ready(Interest::READABLE | Interest::WRITABLE).await?; /// /// if ready.is_readable() { /// // The buffer is **not** included in the async task and will only exist /// // on the stack. /// let mut data = [0; 1024]; /// match socket.try_recv(&mut data[..]) { /// Ok(n) => { /// println!("received {:?}", &data[..n]); /// } /// // False-positive, continue /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {} /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// if ready.is_writable() { /// // Write some data /// match socket.try_send(b"hello world") { /// Ok(n) => { /// println!("sent {} bytes", n); /// } /// // False-positive, continue /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {} /// Err(e) => { /// return Err(e); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is /// usually paired with `try_send()` or `try_send_to()`. /// /// The function may complete without the socket being writable. This is a /// false-positive and attempting a `try_send()` will return with /// `io::ErrorKind::WouldBlock`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Bind socket /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// loop { /// // Wait for the socket to be writable /// socket.writable().await?; /// /// // Try to send data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_send(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write/send readiness. /// /// If the udp stream is not currently ready for sending, this method will /// store a clone of the `Waker` from the provided `Context`. When the udp /// stream becomes ready for sending, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_send_ready` or `poll_send`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_recv_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the udp stream is not ready for writing. /// * `Poll::Ready(Ok(()))` if the udp stream is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`writable`]: method@Self::writable pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Sends data on the socket to the remote address that the socket is /// connected to. /// /// The [`connect`] method will connect this socket to a remote address. /// This method will fail if the socket is not connected. /// /// [`connect`]: method@Self::connect /// /// # Return /// /// On success, the number of bytes sent is returned, otherwise, the /// encountered error is returned. /// /// # Cancel safety /// /// This method is cancel safe. If `send` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that the message was not sent. /// /// # Examples /// /// ```no_run /// use tokio::io; /// use tokio::net::UdpSocket; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Bind socket /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// // Send a message /// socket.send(b"hello world").await?; /// /// Ok(()) /// } /// ``` pub async fn send(&self, buf: &[u8]) -> io::Result { self.io .registration() .async_io(Interest::WRITABLE, || self.io.send(buf)) .await } /// Attempts to send data on the socket to the remote address to which it /// was previously `connect`ed. /// /// The [`connect`] method will connect this socket to a remote address. /// This method will fail if the socket is not connected. /// /// Note that on multiple calls to a `poll_*` method in the send direction, /// only the `Waker` from the `Context` passed to the most recent call will /// be scheduled to receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not available to write /// * `Poll::Ready(Ok(n))` `n` is the number of bytes sent /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`connect`]: method@Self::connect pub fn poll_send(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { self.io .registration() .poll_write_io(cx, || self.io.send(buf)) } /// Tries to send data on the socket to the remote address to which it is /// connected. /// /// When the socket buffer is full, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `writable()`. /// /// # Returns /// /// If successful, `Ok(n)` is returned, where `n` is the number of bytes /// sent. If the socket is not ready to send data, /// `Err(ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Bind a UDP socket /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// // Connect to a peer /// socket.connect("127.0.0.1:8081").await?; /// /// loop { /// // Wait for the socket to be writable /// socket.writable().await?; /// /// // Try to send data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_send(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_send(&self, buf: &[u8]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || self.io.send(buf)) } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_recv()`. /// /// The function may complete without the socket being readable. This is a /// false-positive and attempting a `try_recv()` will return with /// `io::ErrorKind::WouldBlock`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// // The buffer is **not** included in the async task and will /// // only exist on the stack. /// let mut buf = [0; 1024]; /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv(&mut buf) { /// Ok(n) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read/receive readiness. /// /// If the udp stream is not currently ready for receiving, this method will /// store a clone of the `Waker` from the provided `Context`. When the udp /// socket becomes ready for reading, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_recv_ready`, `poll_recv` or /// `poll_peek`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. (However, /// `poll_send_ready` retains a second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the udp stream is not ready for reading. /// * `Poll::Ready(Ok(()))` if the udp stream is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`readable`]: method@Self::readable pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Receives a single datagram message on the socket from the remote address /// to which it is connected. On success, returns the number of bytes read. /// /// The function must be called with valid byte array `buf` of sufficient /// size to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// The [`connect`] method will connect this socket to a remote address. /// This method will fail if the socket is not connected. /// /// # Cancel safety /// /// This method is cancel safe. If `recv` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// socket. /// /// [`connect`]: method@Self::connect /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Bind socket /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// let mut buf = vec![0; 10]; /// let n = socket.recv(&mut buf).await?; /// /// println!("received {} bytes {:?}", n, &buf[..n]); /// /// Ok(()) /// } /// ``` pub async fn recv(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .async_io(Interest::READABLE, || self.io.recv(buf)) .await } /// Attempts to receive a single datagram message on the socket from the remote /// address to which it is `connect`ed. /// /// The [`connect`] method will connect this socket to a remote address. This method /// resolves to an error if the socket is not connected. /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to read /// * `Poll::Ready(Ok(()))` reads data `ReadBuf` if the socket is ready /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`connect`]: method@Self::connect pub fn poll_recv(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { #[allow(clippy::blocks_in_conditions)] let n = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; self.io.recv(b) }))?; // Safety: We trust `recv` to have filled up `n` bytes in the buffer. unsafe { buf.assume_init(n); } buf.advance(n); Poll::Ready(Ok(())) } /// Tries to receive a single datagram message on the socket from the remote /// address to which it is connected. On success, returns the number of /// bytes read. /// /// This method must be called with valid byte array `buf` of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// // The buffer is **not** included in the async task and will /// // only exist on the stack. /// let mut buf = [0; 1024]; /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv(&mut buf) { /// Ok(n) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || self.io.recv(buf)) } cfg_io_util! { /// Tries to receive data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// This method must be called with valid byte array `buf` of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// This method can be used even if `buf` is uninitialized. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// let mut buf = Vec::with_capacity(1024); /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv_buf(&mut buf) { /// Ok(n) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; let n = (*self.io).recv(dst)?; // Safety: We trust `UdpSocket::recv` to have filled up `n` bytes in the // buffer. unsafe { buf.advance_mut(n); } Ok(n) }) } /// Receives a single datagram message on the socket from the remote address /// to which it is connected, advancing the buffer's internal cursor, /// returning how many bytes were read. /// /// This method must be called with valid byte array `buf` of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// This method can be used even if `buf` is uninitialized. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// let mut buf = Vec::with_capacity(512); /// let len = socket.recv_buf(&mut buf).await?; /// /// println!("received {} bytes {:?}", len, &buf[..len]); /// /// Ok(()) /// } /// ``` pub async fn recv_buf(&self, buf: &mut B) -> io::Result { self.io.registration().async_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; let n = (*self.io).recv(dst)?; // Safety: We trust `UdpSocket::recv` to have filled up `n` bytes in the // buffer. unsafe { buf.advance_mut(n); } Ok(n) }).await } /// Tries to receive a single datagram message on the socket. On success, /// returns the number of bytes read and the origin. /// /// This method must be called with valid byte array `buf` of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// This method can be used even if `buf` is uninitialized. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// # Notes /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// let mut buf = Vec::with_capacity(1024); /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv_buf_from(&mut buf) { /// Ok((n, _addr)) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv_buf_from(&self, buf: &mut B) -> io::Result<(usize, SocketAddr)> { self.io.registration().try_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; let (n, addr) = (*self.io).recv_from(dst)?; // Safety: We trust `UdpSocket::recv_from` to have filled up `n` bytes in the // buffer. unsafe { buf.advance_mut(n); } Ok((n, addr)) }) } /// Receives a single datagram message on the socket, advancing the /// buffer's internal cursor, returning how many bytes were read and the origin. /// /// This method must be called with valid byte array `buf` of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// This method can be used even if `buf` is uninitialized. /// /// # Notes /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// socket.connect("127.0.0.1:8081").await?; /// /// let mut buf = Vec::with_capacity(512); /// let (len, addr) = socket.recv_buf_from(&mut buf).await?; /// /// println!("received {:?} bytes from {:?}", len, addr); /// /// Ok(()) /// } /// ``` pub async fn recv_buf_from(&self, buf: &mut B) -> io::Result<(usize, SocketAddr)> { self.io.registration().async_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; let (n, addr) = (*self.io).recv_from(dst)?; // Safety: We trust `UdpSocket::recv_from` to have filled up `n` bytes in the // buffer. unsafe { buf.advance_mut(n); } Ok((n,addr)) }).await } } /// Sends data on the socket to the given address. On success, returns the /// number of bytes written. /// /// Address type can be any implementor of [`ToSocketAddrs`] trait. See its /// documentation for concrete examples. /// /// It is possible for `addr` to yield multiple addresses, but `send_to` /// will only send data to the first address yielded by `addr`. /// /// This will return an error when the IP version of the local socket does /// not match that returned from [`ToSocketAddrs`]. /// /// [`ToSocketAddrs`]: crate::net::ToSocketAddrs /// /// # Cancel safety /// /// This method is cancel safe. If `send_to` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that the message was not sent. /// /// # Example /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// let len = socket.send_to(b"hello world", "127.0.0.1:8081").await?; /// /// println!("Sent {} bytes", len); /// /// Ok(()) /// } /// ``` pub async fn send_to(&self, buf: &[u8], target: A) -> io::Result { let mut addrs = to_socket_addrs(target).await?; match addrs.next() { Some(target) => self.send_to_addr(buf, target).await, None => Err(io::Error::new( io::ErrorKind::InvalidInput, "no addresses to send data to", )), } } /// Attempts to send data on the socket to a given address. /// /// Note that on multiple calls to a `poll_*` method in the send direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to write /// * `Poll::Ready(Ok(n))` `n` is the number of bytes sent. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. pub fn poll_send_to( &self, cx: &mut Context<'_>, buf: &[u8], target: SocketAddr, ) -> Poll> { self.io .registration() .poll_write_io(cx, || self.io.send_to(buf, target)) } /// Tries to send data on the socket to the given address, but if the send is /// blocked this will return right away. /// /// This function is usually paired with `writable()`. /// /// # Returns /// /// If successful, returns the number of bytes sent /// /// Users should ensure that when the remote cannot receive, the /// [`ErrorKind::WouldBlock`] is properly handled. An error can also occur /// if the IP version of the socket does not match that of `target`. /// /// [`ErrorKind::WouldBlock`]: std::io::ErrorKind::WouldBlock /// /// # Example /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// let dst = "127.0.0.1:8081".parse()?; /// /// loop { /// socket.writable().await?; /// /// match socket.try_send_to(&b"hello world"[..], dst) { /// Ok(sent) => { /// println!("sent {} bytes", sent); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// // Writable false positive. /// continue; /// } /// Err(e) => return Err(e.into()), /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_send_to(&self, buf: &[u8], target: SocketAddr) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || self.io.send_to(buf, target)) } async fn send_to_addr(&self, buf: &[u8], target: SocketAddr) -> io::Result { self.io .registration() .async_io(Interest::WRITABLE, || self.io.send_to(buf, target)) .await } /// Receives a single datagram message on the socket. On success, returns /// the number of bytes read and the origin. /// /// The function must be called with valid byte array `buf` of sufficient /// size to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// # Cancel safety /// /// This method is cancel safe. If `recv_from` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// socket. /// /// # Example /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// let mut buf = vec![0u8; 32]; /// let (len, addr) = socket.recv_from(&mut buf).await?; /// /// println!("received {:?} bytes from {:?}", len, addr); /// /// Ok(()) /// } /// ``` /// /// # Notes /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { self.io .registration() .async_io(Interest::READABLE, || self.io.recv_from(buf)) .await } /// Attempts to receive a single datagram on the socket. /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to read /// * `Poll::Ready(Ok(addr))` reads data from `addr` into `ReadBuf` if the socket is ready /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// # Notes /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn poll_recv_from( &self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { #[allow(clippy::blocks_in_conditions)] let (n, addr) = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; self.io.recv_from(b) }))?; // Safety: We trust `recv` to have filled up `n` bytes in the buffer. unsafe { buf.assume_init(n); } buf.advance(n); Poll::Ready(Ok(addr)) } /// Tries to receive a single datagram message on the socket. On success, /// returns the number of bytes read and the origin. /// /// This method must be called with valid byte array `buf` of sufficient size /// to hold the message bytes. If a message is too long to fit in the /// supplied buffer, excess bytes may be discarded. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// # Notes /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// // The buffer is **not** included in the async task and will /// // only exist on the stack. /// let mut buf = [0; 1024]; /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv_from(&mut buf) { /// Ok((n, _addr)) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { self.io .registration() .try_io(Interest::READABLE, || self.io.recv_from(buf)) } /// Tries to read or write from the socket using a user-provided IO operation. /// /// If the socket is ready, the provided closure is called. The closure /// should attempt to perform IO operation on the socket by manually /// calling the appropriate syscall. If the operation fails because the /// socket is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the socket is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `UdpSocket` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: UdpSocket::readable() /// [`writable()`]: UdpSocket::writable() /// [`ready()`]: UdpSocket::ready() pub fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { self.io .registration() .try_io(interest, || self.io.try_io(f)) } /// Reads or writes from the socket using a user-provided IO operation. /// /// The readiness of the socket is awaited and when the socket is ready, /// the provided closure is called. The closure should attempt to perform /// IO operation on the socket by manually calling the appropriate syscall. /// If the operation fails because the socket is not actually ready, /// then the closure should return a `WouldBlock` error. In such case the /// readiness flag is cleared and the socket readiness is awaited again. /// This loop is repeated until the closure returns an `Ok` or an error /// other than `WouldBlock`. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `UdpSocket` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub async fn async_io( &self, interest: Interest, mut f: impl FnMut() -> io::Result, ) -> io::Result { self.io .registration() .async_io(interest, || self.io.try_io(&mut f)) .await } /// Receives a single datagram from the connected address without removing it from the queue. /// On success, returns the number of bytes read from whence the data came. /// /// # Notes /// /// On Windows, if the data is larger than the buffer specified, the buffer /// is filled with the first part of the data, and `peek_from` returns the error /// `WSAEMSGSIZE(10040)`. The excess data is lost. /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// /// MacOS will return an error if you pass a zero-sized buffer. /// /// If you're merely interested in learning the sender of the data at the head of the queue, /// try [`peek_sender`]. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// let mut buf = vec![0u8; 32]; /// let len = socket.peek(&mut buf).await?; /// /// println!("peeked {:?} bytes", len); /// /// Ok(()) /// } /// ``` /// /// [`peek_sender`]: method@Self::peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub async fn peek(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .async_io(Interest::READABLE, || self.io.peek(buf)) .await } /// Receives data from the connected address, without removing it from the input queue. /// On success, returns the sending address of the datagram. /// /// # Notes /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup /// /// On Windows, if the data is larger than the buffer specified, the buffer /// is filled with the first part of the data, and peek returns the error /// `WSAEMSGSIZE(10040)`. The excess data is lost. /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// /// MacOS will return an error if you pass a zero-sized buffer. /// /// If you're merely interested in learning the sender of the data at the head of the queue, /// try [`poll_peek_sender`]. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to read /// * `Poll::Ready(Ok(()))` reads data into `ReadBuf` if the socket is ready /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`poll_peek_sender`]: method@Self::poll_peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn poll_peek(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { #[allow(clippy::blocks_in_conditions)] let n = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; self.io.peek(b) }))?; // Safety: We trust `recv` to have filled up `n` bytes in the buffer. unsafe { buf.assume_init(n); } buf.advance(n); Poll::Ready(Ok(())) } /// Tries to receive data on the connected address without removing it from the input queue. /// On success, returns the number of bytes read. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// # Notes /// /// On Windows, if the data is larger than the buffer specified, the buffer /// is filled with the first part of the data, and peek returns the error /// `WSAEMSGSIZE(10040)`. The excess data is lost. /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// /// MacOS will return an error if you pass a zero-sized buffer. /// /// If you're merely interested in learning the sender of the data at the head of the queue, /// try [`try_peek_sender`]. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [`try_peek_sender`]: method@Self::try_peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn try_peek(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || self.io.peek(buf)) } /// Receives data from the socket, without removing it from the input queue. /// On success, returns the number of bytes read and the address from whence /// the data came. /// /// # Notes /// /// On Windows, if the data is larger than the buffer specified, the buffer /// is filled with the first part of the data, and `peek_from` returns the error /// `WSAEMSGSIZE(10040)`. The excess data is lost. /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// /// MacOS will return an error if you pass a zero-sized buffer. /// /// If you're merely interested in learning the sender of the data at the head of the queue, /// try [`peek_sender`]. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; /// /// let mut buf = vec![0u8; 32]; /// let (len, addr) = socket.peek_from(&mut buf).await?; /// /// println!("peeked {:?} bytes from {:?}", len, addr); /// /// Ok(()) /// } /// ``` /// /// [`peek_sender`]: method@Self::peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub async fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { self.io .registration() .async_io(Interest::READABLE, || self.io.peek_from(buf)) .await } /// Receives data from the socket, without removing it from the input queue. /// On success, returns the sending address of the datagram. /// /// # Notes /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup /// /// On Windows, if the data is larger than the buffer specified, the buffer /// is filled with the first part of the data, and peek returns the error /// `WSAEMSGSIZE(10040)`. The excess data is lost. /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// /// MacOS will return an error if you pass a zero-sized buffer. /// /// If you're merely interested in learning the sender of the data at the head of the queue, /// try [`poll_peek_sender`]. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to read /// * `Poll::Ready(Ok(addr))` reads data from `addr` into `ReadBuf` if the socket is ready /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`poll_peek_sender`]: method@Self::poll_peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn poll_peek_from( &self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { #[allow(clippy::blocks_in_conditions)] let (n, addr) = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; self.io.peek_from(b) }))?; // Safety: We trust `recv` to have filled up `n` bytes in the buffer. unsafe { buf.assume_init(n); } buf.advance(n); Poll::Ready(Ok(addr)) } /// Tries to receive data on the socket without removing it from the input queue. /// On success, returns the number of bytes read and the sending address of the /// datagram. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// # Notes /// /// On Windows, if the data is larger than the buffer specified, the buffer /// is filled with the first part of the data, and peek returns the error /// `WSAEMSGSIZE(10040)`. The excess data is lost. /// Make sure to always use a sufficiently large buffer to hold the /// maximum UDP packet size, which can be up to 65536 bytes in size. /// /// MacOS will return an error if you pass a zero-sized buffer. /// /// If you're merely interested in learning the sender of the data at the head of the queue, /// try [`try_peek_sender`]. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [`try_peek_sender`]: method@Self::try_peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn try_peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { self.io .registration() .try_io(Interest::READABLE, || self.io.peek_from(buf)) } /// Retrieve the sender of the data at the head of the input queue, waiting if empty. /// /// This is equivalent to calling [`peek_from`] with a zero-sized buffer, /// but suppresses the `WSAEMSGSIZE` error on Windows and the "invalid argument" error on macOS. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [`peek_from`]: method@Self::peek_from /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub async fn peek_sender(&self) -> io::Result { self.io .registration() .async_io(Interest::READABLE, || self.peek_sender_inner()) .await } /// Retrieve the sender of the data at the head of the input queue, /// scheduling a wakeup if empty. /// /// This is equivalent to calling [`poll_peek_from`] with a zero-sized buffer, /// but suppresses the `WSAEMSGSIZE` error on Windows and the "invalid argument" error on macOS. /// /// # Notes /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [`poll_peek_from`]: method@Self::poll_peek_from /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn poll_peek_sender(&self, cx: &mut Context<'_>) -> Poll> { self.io .registration() .poll_read_io(cx, || self.peek_sender_inner()) } /// Try to retrieve the sender of the data at the head of the input queue. /// /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is /// returned. This function is usually paired with `readable()`. /// /// Note that the socket address **cannot** be implicitly trusted, because it is relatively /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. /// Because UDP is stateless and does not validate the origin of a packet, /// the attacker does not need to be able to intercept traffic in order to interfere. /// It is important to be aware of this when designing your application-level protocol. /// /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection pub fn try_peek_sender(&self) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || self.peek_sender_inner()) } #[inline] fn peek_sender_inner(&self) -> io::Result { self.io.try_io(|| { self.as_socket() .peek_sender()? // May be `None` if the platform doesn't populate the sender for some reason. // In testing, that only occurred on macOS if you pass a zero-sized buffer, // but the implementation of `Socket::peek_sender()` covers that. .as_socket() .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "sender not available")) }) } /// Gets the value of the `SO_BROADCAST` option for this socket. /// /// For more information about this option, see [`set_broadcast`]. /// /// [`set_broadcast`]: method@Self::set_broadcast pub fn broadcast(&self) -> io::Result { self.io.broadcast() } /// Sets the value of the `SO_BROADCAST` option for this socket. /// /// When enabled, this socket is allowed to send packets to a broadcast /// address. pub fn set_broadcast(&self, on: bool) -> io::Result<()> { self.io.set_broadcast(on) } /// Gets the value of the `IP_MULTICAST_LOOP` option for this socket. /// /// For more information about this option, see [`set_multicast_loop_v4`]. /// /// [`set_multicast_loop_v4`]: method@Self::set_multicast_loop_v4 pub fn multicast_loop_v4(&self) -> io::Result { self.io.multicast_loop_v4() } /// Sets the value of the `IP_MULTICAST_LOOP` option for this socket. /// /// If enabled, multicast packets will be looped back to the local socket. /// /// # Note /// /// This may not have any affect on IPv6 sockets. pub fn set_multicast_loop_v4(&self, on: bool) -> io::Result<()> { self.io.set_multicast_loop_v4(on) } /// Gets the value of the `IP_MULTICAST_TTL` option for this socket. /// /// For more information about this option, see [`set_multicast_ttl_v4`]. /// /// [`set_multicast_ttl_v4`]: method@Self::set_multicast_ttl_v4 pub fn multicast_ttl_v4(&self) -> io::Result { self.io.multicast_ttl_v4() } /// Sets the value of the `IP_MULTICAST_TTL` option for this socket. /// /// Indicates the time-to-live value of outgoing multicast packets for /// this socket. The default value is 1 which means that multicast packets /// don't leave the local network unless explicitly requested. /// /// # Note /// /// This may not have any affect on IPv6 sockets. pub fn set_multicast_ttl_v4(&self, ttl: u32) -> io::Result<()> { self.io.set_multicast_ttl_v4(ttl) } /// Gets the value of the `IPV6_MULTICAST_LOOP` option for this socket. /// /// For more information about this option, see [`set_multicast_loop_v6`]. /// /// [`set_multicast_loop_v6`]: method@Self::set_multicast_loop_v6 pub fn multicast_loop_v6(&self) -> io::Result { self.io.multicast_loop_v6() } /// Sets the value of the `IPV6_MULTICAST_LOOP` option for this socket. /// /// Controls whether this socket sees the multicast packets it sends itself. /// /// # Note /// /// This may not have any affect on IPv4 sockets. pub fn set_multicast_loop_v6(&self, on: bool) -> io::Result<()> { self.io.set_multicast_loop_v6(on) } /// Gets the value of the `IP_TTL` option for this socket. /// /// For more information about this option, see [`set_ttl`]. /// /// [`set_ttl`]: method@Self::set_ttl /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// # use std::io; /// /// # async fn dox() -> io::Result<()> { /// let sock = UdpSocket::bind("127.0.0.1:8080").await?; /// /// println!("{:?}", sock.ttl()?); /// # Ok(()) /// # } /// ``` pub fn ttl(&self) -> io::Result { self.io.ttl() } /// Sets the value for the `IP_TTL` option on this socket. /// /// This value sets the time-to-live field that is used in every packet sent /// from this socket. /// /// # Examples /// /// ```no_run /// use tokio::net::UdpSocket; /// # use std::io; /// /// # async fn dox() -> io::Result<()> { /// let sock = UdpSocket::bind("127.0.0.1:8080").await?; /// sock.set_ttl(60)?; /// /// # Ok(()) /// # } /// ``` pub fn set_ttl(&self, ttl: u32) -> io::Result<()> { self.io.set_ttl(ttl) } /// Gets the value of the `IP_TOS` option for this socket. /// /// For more information about this option, see [`set_tos`]. /// /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) /// /// [`set_tos`]: Self::set_tos // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1464 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))] #[cfg_attr( docsrs, doc(cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))) )] pub fn tos(&self) -> io::Result { self.as_socket().tos() } /// Sets the value for the `IP_TOS` option on this socket. /// /// This value sets the type-of-service field that is used in every packet /// sent from this socket. /// /// **NOTE:** On Windows, `IP_TOS` is only supported on [Windows 8+ or /// Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options) // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1446 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))] #[cfg_attr( docsrs, doc(cfg(not(any( target_os = "fuchsia", target_os = "redox", target_os = "solaris", target_os = "illumos", target_os = "haiku" )))) )] pub fn set_tos(&self, tos: u32) -> io::Result<()> { self.as_socket().set_tos(tos) } /// Gets the value for the `SO_BINDTODEVICE` option on this socket /// /// This value gets the socket-bound device's interface name. #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))] #[cfg_attr( docsrs, doc(cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))) )] pub fn device(&self) -> io::Result>> { self.as_socket().device() } /// Sets the value for the `SO_BINDTODEVICE` option on this socket /// /// If a socket is bound to an interface, only packets received from that /// particular interface are processed by the socket. Note that this only /// works for some socket types, particularly `AF_INET` sockets. /// /// If `interface` is `None` or an empty string it removes the binding. #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] #[cfg_attr( docsrs, doc(cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))) )] pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> { self.as_socket().bind_device(interface) } /// Executes an operation of the `IP_ADD_MEMBERSHIP` type. /// /// This function specifies a new multicast group for this socket to join. /// The address must be a valid multicast address, and `interface` is the /// address of the local interface with which the system should join the /// multicast group. If it's equal to `INADDR_ANY` then an appropriate /// interface is chosen by the system. pub fn join_multicast_v4(&self, multiaddr: Ipv4Addr, interface: Ipv4Addr) -> io::Result<()> { self.io.join_multicast_v4(&multiaddr, &interface) } /// Executes an operation of the `IPV6_ADD_MEMBERSHIP` type. /// /// This function specifies a new multicast group for this socket to join. /// The address must be a valid multicast address, and `interface` is the /// index of the interface to join/leave (or 0 to indicate any interface). pub fn join_multicast_v6(&self, multiaddr: &Ipv6Addr, interface: u32) -> io::Result<()> { self.io.join_multicast_v6(multiaddr, interface) } /// Executes an operation of the `IP_DROP_MEMBERSHIP` type. /// /// For more information about this option, see [`join_multicast_v4`]. /// /// [`join_multicast_v4`]: method@Self::join_multicast_v4 pub fn leave_multicast_v4(&self, multiaddr: Ipv4Addr, interface: Ipv4Addr) -> io::Result<()> { self.io.leave_multicast_v4(&multiaddr, &interface) } /// Executes an operation of the `IPV6_DROP_MEMBERSHIP` type. /// /// For more information about this option, see [`join_multicast_v6`]. /// /// [`join_multicast_v6`]: method@Self::join_multicast_v6 pub fn leave_multicast_v6(&self, multiaddr: &Ipv6Addr, interface: u32) -> io::Result<()> { self.io.leave_multicast_v6(multiaddr, interface) } /// Returns the value of the `SO_ERROR` option. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// use tokio::net::UdpSocket; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Create a socket /// let socket = UdpSocket::bind("0.0.0.0:8080").await?; /// /// if let Ok(Some(err)) = socket.take_error() { /// println!("Got error: {:?}", err); /// } /// /// Ok(()) /// } /// ``` pub fn take_error(&self) -> io::Result> { self.io.take_error() } } impl TryFrom for UdpSocket { type Error = io::Error; /// Consumes stream, returning the tokio I/O object. /// /// This is equivalent to /// [`UdpSocket::from_std(stream)`](UdpSocket::from_std). fn try_from(stream: std::net::UdpSocket) -> Result { Self::from_std(stream) } } impl fmt::Debug for UdpSocket { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.io.fmt(f) } } #[cfg(unix)] mod sys { use super::UdpSocket; use std::os::unix::prelude::*; impl AsRawFd for UdpSocket { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for UdpSocket { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } } cfg_windows! { use crate::os::windows::io::{AsRawSocket, RawSocket}; use crate::os::windows::io::{AsSocket, BorrowedSocket}; impl AsRawSocket for UdpSocket { fn as_raw_socket(&self) -> RawSocket { self.io.as_raw_socket() } } impl AsSocket for UdpSocket { fn as_socket(&self) -> BorrowedSocket<'_> { unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } } } } tokio-1.43.1/src/net/unix/datagram/mod.rs000064400000000000000000000000611046102023000163010ustar 00000000000000//! Unix datagram types. pub(crate) mod socket; tokio-1.43.1/src/net/unix/datagram/socket.rs000064400000000000000000001565461046102023000170360ustar 00000000000000use crate::io::{Interest, PollEvented, ReadBuf, Ready}; use crate::net::unix::SocketAddr; use std::fmt; use std::io; use std::net::Shutdown; use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; use std::task::{ready, Context, Poll}; cfg_io_util! { use bytes::BufMut; } cfg_net_unix! { /// An I/O object representing a Unix datagram socket. /// /// A socket can be either named (associated with a filesystem path) or /// unnamed. /// /// This type does not provide a `split` method, because this functionality /// can be achieved by wrapping the socket in an [`Arc`]. Note that you do /// not need a `Mutex` to share the `UnixDatagram` — an `Arc` /// is enough. This is because all of the methods take `&self` instead of /// `&mut self`. /// /// **Note:** named sockets are persisted even after the object is dropped /// and the program has exited, and cannot be reconnected. It is advised /// that you either check for and unlink the existing socket if it exists, /// or use a temporary file that is guaranteed to not already exist. /// /// [`Arc`]: std::sync::Arc /// /// # Examples /// Using named sockets, associated with a filesystem path: /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind each socket to a filesystem path /// let tx_path = tmp.path().join("tx"); /// let tx = UnixDatagram::bind(&tx_path)?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// let bytes = b"hello world"; /// tx.send_to(bytes, &rx_path).await?; /// /// let mut buf = vec![0u8; 24]; /// let (size, addr) = rx.recv_from(&mut buf).await?; /// /// let dgram = &buf[..size]; /// assert_eq!(dgram, bytes); /// assert_eq!(addr.as_pathname().unwrap(), &tx_path); /// /// # Ok(()) /// # } /// ``` /// /// Using unnamed sockets, created as a pair /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create the pair of sockets /// let (sock1, sock2) = UnixDatagram::pair()?; /// /// // Since the sockets are paired, the paired send/recv /// // functions can be used /// let bytes = b"hello world"; /// sock1.send(bytes).await?; /// /// let mut buff = vec![0u8; 24]; /// let size = sock2.recv(&mut buff).await?; /// /// let dgram = &buff[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixDatagram { io: PollEvented, } } impl UnixDatagram { pub(crate) fn from_mio(sys: mio::net::UnixDatagram) -> io::Result { let datagram = UnixDatagram::new(sys)?; if let Some(e) = datagram.io.take_error()? { return Err(e); } Ok(datagram) } /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_recv()` or `try_send()`. It /// can be used to concurrently `recv` / `send` to the same socket on a single /// task without splitting the socket. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// Concurrently receive from and send to the socket on the same task /// without splitting. /// /// ```no_run /// use tokio::io::Interest; /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// socket.connect(&server_path)?; /// /// loop { /// let ready = socket.ready(Interest::READABLE | Interest::WRITABLE).await?; /// /// if ready.is_readable() { /// let mut data = [0; 1024]; /// match socket.try_recv(&mut data[..]) { /// Ok(n) => { /// println!("received {:?}", &data[..n]); /// } /// // False-positive, continue /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {} /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// if ready.is_writable() { /// // Write some data /// match socket.try_send(b"hello world") { /// Ok(n) => { /// println!("sent {} bytes", n); /// } /// // False-positive, continue /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {} /// Err(e) => { /// return Err(e); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is /// usually paired with `try_send()` or `try_send_to()`. /// /// The function may complete without the socket being writable. This is a /// false-positive and attempting a `try_send()` will return with /// `io::ErrorKind::WouldBlock`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// socket.connect(&server_path)?; /// /// loop { /// // Wait for the socket to be writable /// socket.writable().await?; /// /// // Try to send data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_send(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write/send readiness. /// /// If the socket is not currently ready for sending, this method will /// store a clone of the `Waker` from the provided `Context`. When the socket /// becomes ready for sending, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_send_ready` or `poll_send`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_recv_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready for writing. /// * `Poll::Ready(Ok(()))` if the socket is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`writable`]: method@Self::writable pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_recv()`. /// /// The function may complete without the socket being readable. This is a /// false-positive and attempting a `try_recv()` will return with /// `io::ErrorKind::WouldBlock`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// socket.connect(&server_path)?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// // The buffer is **not** included in the async task and will /// // only exist on the stack. /// let mut buf = [0; 1024]; /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv(&mut buf) { /// Ok(n) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read/receive readiness. /// /// If the socket is not currently ready for receiving, this method will /// store a clone of the `Waker` from the provided `Context`. When the /// socket becomes ready for reading, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_recv_ready`, `poll_recv` or /// `poll_peek`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. (However, /// `poll_send_ready` retains a second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready for reading. /// * `Poll::Ready(Ok(()))` if the socket is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`readable`]: method@Self::readable pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Creates a new `UnixDatagram` bound to the specified path. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind the socket to a filesystem path /// let socket_path = tmp.path().join("socket"); /// let socket = UnixDatagram::bind(&socket_path)?; /// /// # Ok(()) /// # } /// ``` pub fn bind

(path: P) -> io::Result where P: AsRef, { let socket = mio::net::UnixDatagram::bind(path)?; UnixDatagram::new(socket) } /// Creates an unnamed pair of connected sockets. /// /// This function will create a pair of interconnected Unix sockets for /// communicating back and forth between one another. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create the pair of sockets /// let (sock1, sock2) = UnixDatagram::pair()?; /// /// // Since the sockets are paired, the paired send/recv /// // functions can be used /// let bytes = b"hail eris"; /// sock1.send(bytes).await?; /// /// let mut buff = vec![0u8; 24]; /// let size = sock2.recv(&mut buff).await?; /// /// let dgram = &buff[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` pub fn pair() -> io::Result<(UnixDatagram, UnixDatagram)> { let (a, b) = mio::net::UnixDatagram::pair()?; let a = UnixDatagram::new(a)?; let b = UnixDatagram::new(b)?; Ok((a, b)) } /// Creates new [`UnixDatagram`] from a [`std::os::unix::net::UnixDatagram`]. /// /// This function is intended to be used to wrap a `UnixDatagram` from the /// standard library in the Tokio equivalent. /// /// # Notes /// /// The caller is responsible for ensuring that the socket is in /// non-blocking mode. Otherwise all I/O operations on the socket /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::os::unix::net::UnixDatagram::set_nonblocking /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a Tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use std::os::unix::net::UnixDatagram as StdUDS; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind the socket to a filesystem path /// let socket_path = tmp.path().join("socket"); /// let std_socket = StdUDS::bind(&socket_path)?; /// std_socket.set_nonblocking(true)?; /// let tokio_socket = UnixDatagram::from_std(std_socket)?; /// /// # Ok(()) /// # } /// ``` #[track_caller] pub fn from_std(datagram: net::UnixDatagram) -> io::Result { let socket = mio::net::UnixDatagram::from_std(datagram); let io = PollEvented::new(socket)?; Ok(UnixDatagram { io }) } /// Turns a [`tokio::net::UnixDatagram`] into a [`std::os::unix::net::UnixDatagram`]. /// /// The returned [`std::os::unix::net::UnixDatagram`] will have nonblocking /// mode set as `true`. Use [`set_nonblocking`] to change the blocking mode /// if needed. /// /// # Examples /// /// ```rust,no_run /// # use std::error::Error; /// # async fn dox() -> Result<(), Box> { /// let tokio_socket = tokio::net::UnixDatagram::bind("/path/to/the/socket")?; /// let std_socket = tokio_socket.into_std()?; /// std_socket.set_nonblocking(false)?; /// # Ok(()) /// # } /// ``` /// /// [`tokio::net::UnixDatagram`]: UnixDatagram /// [`std::os::unix::net::UnixDatagram`]: std::os::unix::net::UnixDatagram /// [`set_nonblocking`]: fn@std::os::unix::net::UnixDatagram::set_nonblocking pub fn into_std(self) -> io::Result { self.io .into_inner() .map(IntoRawFd::into_raw_fd) .map(|raw_fd| unsafe { std::os::unix::net::UnixDatagram::from_raw_fd(raw_fd) }) } fn new(socket: mio::net::UnixDatagram) -> io::Result { let io = PollEvented::new(socket)?; Ok(UnixDatagram { io }) } /// Creates a new `UnixDatagram` which is not bound to any address. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // Create an unbound socket /// let tx = UnixDatagram::unbound()?; /// /// // Create another, bound socket /// let tmp = tempdir()?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// // Send to the bound socket /// let bytes = b"hello world"; /// tx.send_to(bytes, &rx_path).await?; /// /// let mut buf = vec![0u8; 24]; /// let (size, addr) = rx.recv_from(&mut buf).await?; /// /// let dgram = &buf[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` pub fn unbound() -> io::Result { let socket = mio::net::UnixDatagram::unbound()?; UnixDatagram::new(socket) } /// Connects the socket to the specified address. /// /// The `send` method may be used to send data to the specified address. /// `recv` and `recv_from` will only receive data from that address. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // Create an unbound socket /// let tx = UnixDatagram::unbound()?; /// /// // Create another, bound socket /// let tmp = tempdir()?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// // Connect to the bound socket /// tx.connect(&rx_path)?; /// /// // Send to the bound socket /// let bytes = b"hello world"; /// tx.send(bytes).await?; /// /// let mut buf = vec![0u8; 24]; /// let (size, addr) = rx.recv_from(&mut buf).await?; /// /// let dgram = &buf[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` pub fn connect>(&self, path: P) -> io::Result<()> { self.io.connect(path) } /// Sends data on the socket to the socket's peer. /// /// # Cancel safety /// /// This method is cancel safe. If `send` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that the message was not sent. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create the pair of sockets /// let (sock1, sock2) = UnixDatagram::pair()?; /// /// // Since the sockets are paired, the paired send/recv /// // functions can be used /// let bytes = b"hello world"; /// sock1.send(bytes).await?; /// /// let mut buff = vec![0u8; 24]; /// let size = sock2.recv(&mut buff).await?; /// /// let dgram = &buff[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` pub async fn send(&self, buf: &[u8]) -> io::Result { self.io .registration() .async_io(Interest::WRITABLE, || self.io.send(buf)) .await } /// Tries to send a datagram to the peer without waiting. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// socket.connect(&server_path)?; /// /// loop { /// // Wait for the socket to be writable /// socket.writable().await?; /// /// // Try to send data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_send(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_send(&self, buf: &[u8]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || self.io.send(buf)) } /// Tries to send a datagram to the peer without waiting. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// /// loop { /// // Wait for the socket to be writable /// socket.writable().await?; /// /// // Try to send data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_send_to(b"hello world", &server_path) { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_send_to

(&self, buf: &[u8], target: P) -> io::Result where P: AsRef, { self.io .registration() .try_io(Interest::WRITABLE, || self.io.send_to(buf, target)) } /// Receives data from the socket. /// /// # Cancel safety /// /// This method is cancel safe. If `recv` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// socket. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create the pair of sockets /// let (sock1, sock2) = UnixDatagram::pair()?; /// /// // Since the sockets are paired, the paired send/recv /// // functions can be used /// let bytes = b"hello world"; /// sock1.send(bytes).await?; /// /// let mut buff = vec![0u8; 24]; /// let size = sock2.recv(&mut buff).await?; /// /// let dgram = &buff[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` pub async fn recv(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .async_io(Interest::READABLE, || self.io.recv(buf)) .await } /// Tries to receive a datagram from the peer without waiting. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// socket.connect(&server_path)?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// // The buffer is **not** included in the async task and will /// // only exist on the stack. /// let mut buf = [0; 1024]; /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv(&mut buf) { /// Ok(n) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || self.io.recv(buf)) } cfg_io_util! { /// Tries to receive data from the socket without waiting. /// /// This method can be used even if `buf` is uninitialized. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// let mut buf = Vec::with_capacity(1024); /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv_buf_from(&mut buf) { /// Ok((n, _addr)) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv_buf_from(&self, buf: &mut B) -> io::Result<(usize, SocketAddr)> { let (n, addr) = self.io.registration().try_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `UnixDatagram::recv_from` to have filled up `n` bytes in the // buffer. let (n, addr) = (*self.io).recv_from(dst)?; unsafe { buf.advance_mut(n); } Ok((n, addr)) })?; Ok((n, SocketAddr(addr))) } /// Receives from the socket, advances the /// buffer's internal cursor and returns how many bytes were read and the origin. /// /// This method can be used even if `buf` is uninitialized. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind each socket to a filesystem path /// let tx_path = tmp.path().join("tx"); /// let tx = UnixDatagram::bind(&tx_path)?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// let bytes = b"hello world"; /// tx.send_to(bytes, &rx_path).await?; /// /// let mut buf = Vec::with_capacity(24); /// let (size, addr) = rx.recv_buf_from(&mut buf).await?; /// /// let dgram = &buf[..size]; /// assert_eq!(dgram, bytes); /// assert_eq!(addr.as_pathname().unwrap(), &tx_path); /// /// # Ok(()) /// # } /// ``` pub async fn recv_buf_from(&self, buf: &mut B) -> io::Result<(usize, SocketAddr)> { self.io.registration().async_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `UnixDatagram::recv_from` to have filled up `n` bytes in the // buffer. let (n, addr) = (*self.io).recv_from(dst)?; unsafe { buf.advance_mut(n); } Ok((n,SocketAddr(addr))) }).await } /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// This method can be used even if `buf` is uninitialized. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// socket.connect(&server_path)?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// let mut buf = Vec::with_capacity(1024); /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv_buf(&mut buf) { /// Ok(n) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `UnixDatagram::recv` to have filled up `n` bytes in the // buffer. let n = (*self.io).recv(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }) } /// Receives data from the socket from the address to which it is connected, /// advancing the buffer's internal cursor, returning how many bytes were read. /// /// This method can be used even if `buf` is uninitialized. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create the pair of sockets /// let (sock1, sock2) = UnixDatagram::pair()?; /// /// // Since the sockets are paired, the paired send/recv /// // functions can be used /// let bytes = b"hello world"; /// sock1.send(bytes).await?; /// /// let mut buff = Vec::with_capacity(24); /// let size = sock2.recv_buf(&mut buff).await?; /// /// let dgram = &buff[..size]; /// assert_eq!(dgram, bytes); /// /// # Ok(()) /// # } /// ``` pub async fn recv_buf(&self, buf: &mut B) -> io::Result { self.io.registration().async_io(Interest::READABLE, || { let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `UnixDatagram::recv_from` to have filled up `n` bytes in the // buffer. let n = (*self.io).recv(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }).await } } /// Sends data on the socket to the specified address. /// /// # Cancel safety /// /// This method is cancel safe. If `send_to` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that the message was not sent. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind each socket to a filesystem path /// let tx_path = tmp.path().join("tx"); /// let tx = UnixDatagram::bind(&tx_path)?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// let bytes = b"hello world"; /// tx.send_to(bytes, &rx_path).await?; /// /// let mut buf = vec![0u8; 24]; /// let (size, addr) = rx.recv_from(&mut buf).await?; /// /// let dgram = &buf[..size]; /// assert_eq!(dgram, bytes); /// assert_eq!(addr.as_pathname().unwrap(), &tx_path); /// /// # Ok(()) /// # } /// ``` pub async fn send_to

(&self, buf: &[u8], target: P) -> io::Result where P: AsRef, { self.io .registration() .async_io(Interest::WRITABLE, || self.io.send_to(buf, target.as_ref())) .await } /// Receives data from the socket. /// /// # Cancel safety /// /// This method is cancel safe. If `recv_from` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// socket. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind each socket to a filesystem path /// let tx_path = tmp.path().join("tx"); /// let tx = UnixDatagram::bind(&tx_path)?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// let bytes = b"hello world"; /// tx.send_to(bytes, &rx_path).await?; /// /// let mut buf = vec![0u8; 24]; /// let (size, addr) = rx.recv_from(&mut buf).await?; /// /// let dgram = &buf[..size]; /// assert_eq!(dgram, bytes); /// assert_eq!(addr.as_pathname().unwrap(), &tx_path); /// /// # Ok(()) /// # } /// ``` pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { let (n, addr) = self .io .registration() .async_io(Interest::READABLE, || self.io.recv_from(buf)) .await?; Ok((n, SocketAddr(addr))) } /// Attempts to receive a single datagram on the specified address. /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to read /// * `Poll::Ready(Ok(addr))` reads data from `addr` into `ReadBuf` if the socket is ready /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. pub fn poll_recv_from( &self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { #[allow(clippy::blocks_in_conditions)] let (n, addr) = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; self.io.recv_from(b) }))?; // Safety: We trust `recv` to have filled up `n` bytes in the buffer. unsafe { buf.assume_init(n); } buf.advance(n); Poll::Ready(Ok(SocketAddr(addr))) } /// Attempts to send data to the specified address. /// /// Note that on multiple calls to a `poll_*` method in the send direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to write /// * `Poll::Ready(Ok(n))` `n` is the number of bytes sent. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. pub fn poll_send_to

( &self, cx: &mut Context<'_>, buf: &[u8], target: P, ) -> Poll> where P: AsRef, { self.io .registration() .poll_write_io(cx, || self.io.send_to(buf, target.as_ref())) } /// Attempts to send data on the socket to the remote address to which it /// was previously `connect`ed. /// /// The [`connect`] method will connect this socket to a remote address. /// This method will fail if the socket is not connected. /// /// Note that on multiple calls to a `poll_*` method in the send direction, /// only the `Waker` from the `Context` passed to the most recent call will /// be scheduled to receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not available to write /// * `Poll::Ready(Ok(n))` `n` is the number of bytes sent /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`connect`]: method@Self::connect pub fn poll_send(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { self.io .registration() .poll_write_io(cx, || self.io.send(buf)) } /// Attempts to receive a single datagram message on the socket from the remote /// address to which it is `connect`ed. /// /// The [`connect`] method will connect this socket to a remote address. This method /// resolves to an error if the socket is not connected. /// /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the /// `Waker` from the `Context` passed to the most recent call will be scheduled to /// receive a wakeup. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the socket is not ready to read /// * `Poll::Ready(Ok(()))` reads data `ReadBuf` if the socket is ready /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`connect`]: method@Self::connect pub fn poll_recv(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { #[allow(clippy::blocks_in_conditions)] let n = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) }; self.io.recv(b) }))?; // Safety: We trust `recv` to have filled up `n` bytes in the buffer. unsafe { buf.assume_init(n); } buf.advance(n); Poll::Ready(Ok(())) } /// Tries to receive data from the socket without waiting. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixDatagram; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let client_path = dir.path().join("client.sock"); /// let server_path = dir.path().join("server.sock"); /// let socket = UnixDatagram::bind(&client_path)?; /// /// loop { /// // Wait for the socket to be readable /// socket.readable().await?; /// /// // The buffer is **not** included in the async task and will /// // only exist on the stack. /// let mut buf = [0; 1024]; /// /// // Try to recv data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match socket.try_recv_from(&mut buf) { /// Ok((n, _addr)) => { /// println!("GOT {:?}", &buf[..n]); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { let (n, addr) = self .io .registration() .try_io(Interest::READABLE, || self.io.recv_from(buf))?; Ok((n, SocketAddr(addr))) } /// Tries to read or write from the socket using a user-provided IO operation. /// /// If the socket is ready, the provided closure is called. The closure /// should attempt to perform IO operation on the socket by manually /// calling the appropriate syscall. If the operation fails because the /// socket is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the socket is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `UnixDatagram` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: UnixDatagram::readable() /// [`writable()`]: UnixDatagram::writable() /// [`ready()`]: UnixDatagram::ready() pub fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { self.io .registration() .try_io(interest, || self.io.try_io(f)) } /// Reads or writes from the socket using a user-provided IO operation. /// /// The readiness of the socket is awaited and when the socket is ready, /// the provided closure is called. The closure should attempt to perform /// IO operation on the socket by manually calling the appropriate syscall. /// If the operation fails because the socket is not actually ready, /// then the closure should return a `WouldBlock` error. In such case the /// readiness flag is cleared and the socket readiness is awaited again. /// This loop is repeated until the closure returns an `Ok` or an error /// other than `WouldBlock`. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `UnixDatagram` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub async fn async_io( &self, interest: Interest, mut f: impl FnMut() -> io::Result, ) -> io::Result { self.io .registration() .async_io(interest, || self.io.try_io(&mut f)) .await } /// Returns the local address that this socket is bound to. /// /// # Examples /// For a socket bound to a local path /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // We use a temporary directory so that the socket /// // files left by the bound sockets will get cleaned up. /// let tmp = tempdir()?; /// /// // Bind socket to a filesystem path /// let socket_path = tmp.path().join("socket"); /// let socket = UnixDatagram::bind(&socket_path)?; /// /// assert_eq!(socket.local_addr()?.as_pathname().unwrap(), &socket_path); /// /// # Ok(()) /// # } /// ``` /// /// For an unbound socket /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create an unbound socket /// let socket = UnixDatagram::unbound()?; /// /// assert!(socket.local_addr()?.is_unnamed()); /// /// # Ok(()) /// # } /// ``` pub fn local_addr(&self) -> io::Result { self.io.local_addr().map(SocketAddr) } /// Returns the address of this socket's peer. /// /// The `connect` method will connect the socket to a peer. /// /// # Examples /// For a peer with a local path /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use tempfile::tempdir; /// /// // Create an unbound socket /// let tx = UnixDatagram::unbound()?; /// /// // Create another, bound socket /// let tmp = tempdir()?; /// let rx_path = tmp.path().join("rx"); /// let rx = UnixDatagram::bind(&rx_path)?; /// /// // Connect to the bound socket /// tx.connect(&rx_path)?; /// /// assert_eq!(tx.peer_addr()?.as_pathname().unwrap(), &rx_path); /// /// # Ok(()) /// # } /// ``` /// /// For an unbound peer /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create the pair of sockets /// let (sock1, sock2) = UnixDatagram::pair()?; /// /// assert!(sock1.peer_addr()?.is_unnamed()); /// /// # Ok(()) /// # } /// ``` pub fn peer_addr(&self) -> io::Result { self.io.peer_addr().map(SocketAddr) } /// Returns the value of the `SO_ERROR` option. /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// /// // Create an unbound socket /// let socket = UnixDatagram::unbound()?; /// /// if let Ok(Some(err)) = socket.take_error() { /// println!("Got error: {:?}", err); /// } /// /// # Ok(()) /// # } /// ``` pub fn take_error(&self) -> io::Result> { self.io.take_error() } /// Shuts down the read, write, or both halves of this connection. /// /// This function will cause all pending and future I/O calls on the /// specified portions to immediately return with an appropriate value /// (see the documentation of `Shutdown`). /// /// # Examples /// ``` /// # if cfg!(miri) { return } // No SOCK_DGRAM for `socketpair` in miri. /// # use std::error::Error; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// use tokio::net::UnixDatagram; /// use std::net::Shutdown; /// /// // Create an unbound socket /// let (socket, other) = UnixDatagram::pair()?; /// /// socket.shutdown(Shutdown::Both)?; /// /// // NOTE: the following commented out code does NOT work as expected. /// // Due to an underlying issue, the recv call will block indefinitely. /// // See: https://github.com/tokio-rs/tokio/issues/1679 /// //let mut buff = vec![0u8; 24]; /// //let size = socket.recv(&mut buff).await?; /// //assert_eq!(size, 0); /// /// let send_result = socket.send(b"hello world").await; /// assert!(send_result.is_err()); /// /// # Ok(()) /// # } /// ``` pub fn shutdown(&self, how: Shutdown) -> io::Result<()> { self.io.shutdown(how) } } impl TryFrom for UnixDatagram { type Error = io::Error; /// Consumes stream, returning the Tokio I/O object. /// /// This is equivalent to /// [`UnixDatagram::from_std(stream)`](UnixDatagram::from_std). fn try_from(stream: std::os::unix::net::UnixDatagram) -> Result { Self::from_std(stream) } } impl fmt::Debug for UnixDatagram { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.io.fmt(f) } } impl AsRawFd for UnixDatagram { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for UnixDatagram { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } tokio-1.43.1/src/net/unix/listener.rs000064400000000000000000000210451046102023000155740ustar 00000000000000use crate::io::{Interest, PollEvented}; use crate::net::unix::{SocketAddr, UnixStream}; use std::fmt; use std::io; #[cfg(target_os = "android")] use std::os::android::net::SocketAddrExt; #[cfg(target_os = "linux")] use std::os::linux::net::SocketAddrExt; #[cfg(any(target_os = "linux", target_os = "android"))] use std::os::unix::ffi::OsStrExt; use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net::{self, SocketAddr as StdSocketAddr}; use std::path::Path; use std::task::{ready, Context, Poll}; cfg_net_unix! { /// A Unix socket which can accept connections from other Unix sockets. /// /// You can accept a new connection by using the [`accept`](`UnixListener::accept`) method. /// /// A `UnixListener` can be turned into a `Stream` with [`UnixListenerStream`]. /// /// [`UnixListenerStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.UnixListenerStream.html /// /// # Errors /// /// Note that accepting a connection can lead to various errors and not all /// of them are necessarily fatal ‒ for example having too many open file /// descriptors or the other side closing the connection while it waits in /// an accept queue. These would terminate the stream if not handled in any /// way. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixListener; /// /// #[tokio::main] /// async fn main() { /// let listener = UnixListener::bind("/path/to/the/socket").unwrap(); /// loop { /// match listener.accept().await { /// Ok((stream, _addr)) => { /// println!("new client!"); /// } /// Err(e) => { /* connection failed */ } /// } /// } /// } /// ``` #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixListener { io: PollEvented, } } impl UnixListener { pub(crate) fn new(listener: mio::net::UnixListener) -> io::Result { let io = PollEvented::new(listener)?; Ok(UnixListener { io }) } /// Creates a new `UnixListener` bound to the specified path. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. #[track_caller] pub fn bind

(path: P) -> io::Result where P: AsRef, { // For now, we handle abstract socket paths on linux here. #[cfg(any(target_os = "linux", target_os = "android"))] let addr = { let os_str_bytes = path.as_ref().as_os_str().as_bytes(); if os_str_bytes.starts_with(b"\0") { StdSocketAddr::from_abstract_name(&os_str_bytes[1..])? } else { StdSocketAddr::from_pathname(path)? } }; #[cfg(not(any(target_os = "linux", target_os = "android")))] let addr = StdSocketAddr::from_pathname(path)?; let listener = mio::net::UnixListener::bind_addr(&addr)?; let io = PollEvented::new(listener)?; Ok(UnixListener { io }) } /// Creates new [`UnixListener`] from a [`std::os::unix::net::UnixListener`]. /// /// This function is intended to be used to wrap a `UnixListener` from the /// standard library in the Tokio equivalent. /// /// # Notes /// /// The caller is responsible for ensuring that the listener is in /// non-blocking mode. Otherwise all I/O operations on the listener /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::os::unix::net::UnixListener::set_nonblocking /// /// # Examples /// /// ```no_run /// use tokio::net::UnixListener; /// use std::os::unix::net::UnixListener as StdUnixListener; /// # use std::error::Error; /// /// # async fn dox() -> Result<(), Box> { /// let std_listener = StdUnixListener::bind("/path/to/the/socket")?; /// std_listener.set_nonblocking(true)?; /// let listener = UnixListener::from_std(std_listener)?; /// # Ok(()) /// # } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. #[track_caller] pub fn from_std(listener: net::UnixListener) -> io::Result { let listener = mio::net::UnixListener::from_std(listener); let io = PollEvented::new(listener)?; Ok(UnixListener { io }) } /// Turns a [`tokio::net::UnixListener`] into a [`std::os::unix::net::UnixListener`]. /// /// The returned [`std::os::unix::net::UnixListener`] will have nonblocking mode /// set as `true`. Use [`set_nonblocking`] to change the blocking mode if needed. /// /// # Examples /// /// ```rust,no_run /// # use std::error::Error; /// # async fn dox() -> Result<(), Box> { /// let tokio_listener = tokio::net::UnixListener::bind("/path/to/the/socket")?; /// let std_listener = tokio_listener.into_std()?; /// std_listener.set_nonblocking(false)?; /// # Ok(()) /// # } /// ``` /// /// [`tokio::net::UnixListener`]: UnixListener /// [`std::os::unix::net::UnixListener`]: std::os::unix::net::UnixListener /// [`set_nonblocking`]: fn@std::os::unix::net::UnixListener::set_nonblocking pub fn into_std(self) -> io::Result { self.io .into_inner() .map(IntoRawFd::into_raw_fd) .map(|raw_fd| unsafe { net::UnixListener::from_raw_fd(raw_fd) }) } /// Returns the local socket address of this listener. pub fn local_addr(&self) -> io::Result { self.io.local_addr().map(SocketAddr) } /// Returns the value of the `SO_ERROR` option. pub fn take_error(&self) -> io::Result> { self.io.take_error() } /// Accepts a new incoming connection to this listener. /// /// # Cancel safety /// /// This method is cancel safe. If the method is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no new connections were /// accepted by this method. pub async fn accept(&self) -> io::Result<(UnixStream, SocketAddr)> { let (mio, addr) = self .io .registration() .async_io(Interest::READABLE, || self.io.accept()) .await?; let addr = SocketAddr(addr); let stream = UnixStream::new(mio)?; Ok((stream, addr)) } /// Polls to accept a new incoming connection to this listener. /// /// If there is no connection to accept, `Poll::Pending` is returned and the /// current task will be notified by a waker. Note that on multiple calls /// to `poll_accept`, only the `Waker` from the `Context` passed to the most /// recent call is scheduled to receive a wakeup. pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll> { let (sock, addr) = ready!(self.io.registration().poll_read_io(cx, || self.io.accept()))?; let addr = SocketAddr(addr); let sock = UnixStream::new(sock)?; Poll::Ready(Ok((sock, addr))) } } impl TryFrom for UnixListener { type Error = io::Error; /// Consumes stream, returning the tokio I/O object. /// /// This is equivalent to /// [`UnixListener::from_std(stream)`](UnixListener::from_std). fn try_from(stream: std::os::unix::net::UnixListener) -> io::Result { Self::from_std(stream) } } impl fmt::Debug for UnixListener { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.io.fmt(f) } } impl AsRawFd for UnixListener { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for UnixListener { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } tokio-1.43.1/src/net/unix/mod.rs000064400000000000000000000016111046102023000145230ustar 00000000000000//! Unix specific network types. // This module does not currently provide any public API, but it was // unintentionally defined as a public module. Hide it from the documentation // instead of changing it to a private module to avoid breakage. #[doc(hidden)] pub mod datagram; pub(crate) mod listener; pub(crate) mod socket; mod split; pub use split::{ReadHalf, WriteHalf}; mod split_owned; pub use split_owned::{OwnedReadHalf, OwnedWriteHalf, ReuniteError}; mod socketaddr; pub use socketaddr::SocketAddr; pub(crate) mod stream; pub(crate) use stream::UnixStream; mod ucred; pub use ucred::UCred; pub mod pipe; /// A type representing user ID. #[allow(non_camel_case_types)] pub type uid_t = u32; /// A type representing group ID. #[allow(non_camel_case_types)] pub type gid_t = u32; /// A type representing process and process group IDs. #[allow(non_camel_case_types)] pub type pid_t = i32; tokio-1.43.1/src/net/unix/pipe.rs000064400000000000000000001450231046102023000147070ustar 00000000000000//! Unix pipe types. use crate::io::interest::Interest; use crate::io::{AsyncRead, AsyncWrite, PollEvented, ReadBuf, Ready}; use mio::unix::pipe as mio_pipe; use std::fs::File; use std::io::{self, Read, Write}; use std::os::unix::fs::OpenOptionsExt; use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd}; use std::path::Path; use std::pin::Pin; use std::task::{Context, Poll}; cfg_io_util! { use bytes::BufMut; } /// Creates a new anonymous Unix pipe. /// /// This function will open a new pipe and associate both pipe ends with the default /// event loop. /// /// If you need to create a pipe for communication with a spawned process, you can /// use [`Stdio::piped()`] instead. /// /// [`Stdio::piped()`]: std::process::Stdio::piped /// /// # Errors /// /// If creating a pipe fails, this function will return with the related OS error. /// /// # Examples /// /// Create a pipe and pass the writing end to a spawned process. /// /// ```no_run /// use tokio::net::unix::pipe; /// use tokio::process::Command; /// # use tokio::io::AsyncReadExt; /// # use std::error::Error; /// /// # async fn dox() -> Result<(), Box> { /// let (tx, mut rx) = pipe::pipe()?; /// let mut buffer = String::new(); /// /// let status = Command::new("echo") /// .arg("Hello, world!") /// .stdout(tx.into_blocking_fd()?) /// .status(); /// rx.read_to_string(&mut buffer).await?; /// /// assert!(status.await?.success()); /// assert_eq!(buffer, "Hello, world!\n"); /// # Ok(()) /// # } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn pipe() -> io::Result<(Sender, Receiver)> { let (tx, rx) = mio_pipe::new()?; Ok((Sender::from_mio(tx)?, Receiver::from_mio(rx)?)) } /// Options and flags which can be used to configure how a FIFO file is opened. /// /// This builder allows configuring how to create a pipe end from a FIFO file. /// Generally speaking, when using `OpenOptions`, you'll first call [`new`], /// then chain calls to methods to set each option, then call either /// [`open_receiver`] or [`open_sender`], passing the path of the FIFO file you /// are trying to open. This will give you a [`io::Result`] with a pipe end /// inside that you can further operate on. /// /// [`new`]: OpenOptions::new /// [`open_receiver`]: OpenOptions::open_receiver /// [`open_sender`]: OpenOptions::open_sender /// /// # Examples /// /// Opening a pair of pipe ends from a FIFO file: /// /// ```no_run /// use tokio::net::unix::pipe; /// # use std::error::Error; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// let rx = pipe::OpenOptions::new().open_receiver(FIFO_NAME)?; /// let tx = pipe::OpenOptions::new().open_sender(FIFO_NAME)?; /// # Ok(()) /// # } /// ``` /// /// Opening a [`Sender`] on Linux when you are sure the file is a FIFO: /// /// ```ignore /// use tokio::net::unix::pipe; /// use nix::{unistd::mkfifo, sys::stat::Mode}; /// # use std::error::Error; /// /// // Our program has exclusive access to this path. /// const FIFO_NAME: &str = "path/to/a/new/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// mkfifo(FIFO_NAME, Mode::S_IRWXU)?; /// let tx = pipe::OpenOptions::new() /// .read_write(true) /// .unchecked(true) /// .open_sender(FIFO_NAME)?; /// # Ok(()) /// # } /// ``` #[derive(Clone, Debug)] pub struct OpenOptions { #[cfg(target_os = "linux")] read_write: bool, unchecked: bool, } impl OpenOptions { /// Creates a blank new set of options ready for configuration. /// /// All options are initially set to `false`. pub fn new() -> OpenOptions { OpenOptions { #[cfg(target_os = "linux")] read_write: false, unchecked: false, } } /// Sets the option for read-write access. /// /// This option, when true, will indicate that a FIFO file will be opened /// in read-write access mode. This operation is not defined by the POSIX /// standard and is only guaranteed to work on Linux. /// /// # Examples /// /// Opening a [`Sender`] even if there are no open reading ends: /// /// ```ignore /// use tokio::net::unix::pipe; /// /// let tx = pipe::OpenOptions::new() /// .read_write(true) /// .open_sender("path/to/a/fifo"); /// ``` /// /// Opening a resilient [`Receiver`] i.e. a reading pipe end which will not /// fail with [`UnexpectedEof`] during reading if all writing ends of the /// pipe close the FIFO file. /// /// [`UnexpectedEof`]: std::io::ErrorKind::UnexpectedEof /// /// ```ignore /// use tokio::net::unix::pipe; /// /// let tx = pipe::OpenOptions::new() /// .read_write(true) /// .open_receiver("path/to/a/fifo"); /// ``` #[cfg(target_os = "linux")] #[cfg_attr(docsrs, doc(cfg(target_os = "linux")))] pub fn read_write(&mut self, value: bool) -> &mut Self { self.read_write = value; self } /// Sets the option to skip the check for FIFO file type. /// /// By default, [`open_receiver`] and [`open_sender`] functions will check /// if the opened file is a FIFO file. Set this option to `true` if you are /// sure the file is a FIFO file. /// /// [`open_receiver`]: OpenOptions::open_receiver /// [`open_sender`]: OpenOptions::open_sender /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use nix::{unistd::mkfifo, sys::stat::Mode}; /// # use std::error::Error; /// /// // Our program has exclusive access to this path. /// const FIFO_NAME: &str = "path/to/a/new/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// mkfifo(FIFO_NAME, Mode::S_IRWXU)?; /// let rx = pipe::OpenOptions::new() /// .unchecked(true) /// .open_receiver(FIFO_NAME)?; /// # Ok(()) /// # } /// ``` pub fn unchecked(&mut self, value: bool) -> &mut Self { self.unchecked = value; self } /// Creates a [`Receiver`] from a FIFO file with the options specified by `self`. /// /// This function will open the FIFO file at the specified path, possibly /// check if it is a pipe, and associate the pipe with the default event /// loop for reading. /// /// # Errors /// /// If the file type check fails, this function will fail with `io::ErrorKind::InvalidInput`. /// This function may also fail with other standard OS errors. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn open_receiver>(&self, path: P) -> io::Result { let file = self.open(path.as_ref(), PipeEnd::Receiver)?; Receiver::from_file_unchecked(file) } /// Creates a [`Sender`] from a FIFO file with the options specified by `self`. /// /// This function will open the FIFO file at the specified path, possibly /// check if it is a pipe, and associate the pipe with the default event /// loop for writing. /// /// # Errors /// /// If the file type check fails, this function will fail with `io::ErrorKind::InvalidInput`. /// If the file is not opened in read-write access mode and the file is not /// currently open for reading, this function will fail with `ENXIO`. /// This function may also fail with other standard OS errors. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn open_sender>(&self, path: P) -> io::Result { let file = self.open(path.as_ref(), PipeEnd::Sender)?; Sender::from_file_unchecked(file) } fn open(&self, path: &Path, pipe_end: PipeEnd) -> io::Result { let mut options = std::fs::OpenOptions::new(); options .read(pipe_end == PipeEnd::Receiver) .write(pipe_end == PipeEnd::Sender) .custom_flags(libc::O_NONBLOCK); #[cfg(target_os = "linux")] if self.read_write { options.read(true).write(true); } let file = options.open(path)?; if !self.unchecked && !is_pipe(file.as_fd())? { return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); } Ok(file) } } impl Default for OpenOptions { fn default() -> OpenOptions { OpenOptions::new() } } #[derive(Clone, Copy, PartialEq, Eq, Debug)] enum PipeEnd { Sender, Receiver, } /// Writing end of a Unix pipe. /// /// It can be constructed from a FIFO file with [`OpenOptions::open_sender`]. /// /// Opening a named pipe for writing involves a few steps. /// Call to [`OpenOptions::open_sender`] might fail with an error indicating /// different things: /// /// * [`io::ErrorKind::NotFound`] - There is no file at the specified path. /// * [`io::ErrorKind::InvalidInput`] - The file exists, but it is not a FIFO. /// * [`ENXIO`] - The file is a FIFO, but no process has it open for reading. /// Sleep for a while and try again. /// * Other OS errors not specific to opening FIFO files. /// /// Opening a `Sender` from a FIFO file should look like this: /// /// ```no_run /// use tokio::net::unix::pipe; /// use tokio::time::{self, Duration}; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// // Wait for a reader to open the file. /// let tx = loop { /// match pipe::OpenOptions::new().open_sender(FIFO_NAME) { /// Ok(tx) => break tx, /// Err(e) if e.raw_os_error() == Some(libc::ENXIO) => {}, /// Err(e) => return Err(e.into()), /// } /// /// time::sleep(Duration::from_millis(50)).await; /// }; /// # Ok(()) /// # } /// ``` /// /// On Linux, it is possible to create a `Sender` without waiting in a sleeping /// loop. This is done by opening a named pipe in read-write access mode with /// `OpenOptions::read_write`. This way, a `Sender` can at the same time hold /// both a writing end and a reading end, and the latter allows to open a FIFO /// without [`ENXIO`] error since the pipe is open for reading as well. /// /// `Sender` cannot be used to read from a pipe, so in practice the read access /// is only used when a FIFO is opened. However, using a `Sender` in read-write /// mode **may lead to lost data**, because written data will be dropped by the /// system as soon as all pipe ends are closed. To avoid lost data you have to /// make sure that a reading end has been opened before dropping a `Sender`. /// /// Note that using read-write access mode with FIFO files is not defined by /// the POSIX standard and it is only guaranteed to work on Linux. /// /// ```ignore /// use tokio::io::AsyncWriteExt; /// use tokio::net::unix::pipe; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// let mut tx = pipe::OpenOptions::new() /// .read_write(true) /// .open_sender(FIFO_NAME)?; /// /// // Asynchronously write to the pipe before a reader. /// tx.write_all(b"hello world").await?; /// # Ok(()) /// # } /// ``` /// /// [`ENXIO`]: https://docs.rs/libc/latest/libc/constant.ENXIO.html #[derive(Debug)] pub struct Sender { io: PollEvented, } impl Sender { fn from_mio(mio_tx: mio_pipe::Sender) -> io::Result { let io = PollEvented::new_with_interest(mio_tx, Interest::WRITABLE)?; Ok(Sender { io }) } /// Creates a new `Sender` from a [`File`]. /// /// This function is intended to construct a pipe from a [`File`] representing /// a special FIFO file. It will check if the file is a pipe and has write access, /// set it in non-blocking mode and perform the conversion. /// /// # Errors /// /// Fails with `io::ErrorKind::InvalidInput` if the file is not a pipe or it /// does not have write access. Also fails with any standard OS error if it occurs. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_file(file: File) -> io::Result { Sender::from_owned_fd(file.into()) } /// Creates a new `Sender` from an [`OwnedFd`]. /// /// This function is intended to construct a pipe from an [`OwnedFd`] representing /// an anonymous pipe or a special FIFO file. It will check if the file descriptor /// is a pipe and has write access, set it in non-blocking mode and perform the /// conversion. /// /// # Errors /// /// Fails with `io::ErrorKind::InvalidInput` if the file descriptor is not a pipe /// or it does not have write access. Also fails with any standard OS error if it /// occurs. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_owned_fd(owned_fd: OwnedFd) -> io::Result { if !is_pipe(owned_fd.as_fd())? { return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); } let flags = get_file_flags(owned_fd.as_fd())?; if has_write_access(flags) { set_nonblocking(owned_fd.as_fd(), flags)?; Sender::from_owned_fd_unchecked(owned_fd) } else { Err(io::Error::new( io::ErrorKind::InvalidInput, "not in O_WRONLY or O_RDWR access mode", )) } } /// Creates a new `Sender` from a [`File`] without checking pipe properties. /// /// This function is intended to construct a pipe from a File representing /// a special FIFO file. The conversion assumes nothing about the underlying /// file; it is left up to the user to make sure it is opened with write access, /// represents a pipe and is set in non-blocking mode. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::fs::OpenOptions; /// use std::os::unix::fs::{FileTypeExt, OpenOptionsExt}; /// # use std::error::Error; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// let file = OpenOptions::new() /// .write(true) /// .custom_flags(libc::O_NONBLOCK) /// .open(FIFO_NAME)?; /// if file.metadata()?.file_type().is_fifo() { /// let tx = pipe::Sender::from_file_unchecked(file)?; /// /* use the Sender */ /// } /// # Ok(()) /// # } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_file_unchecked(file: File) -> io::Result { Sender::from_owned_fd_unchecked(file.into()) } /// Creates a new `Sender` from an [`OwnedFd`] without checking pipe properties. /// /// This function is intended to construct a pipe from an [`OwnedFd`] representing /// an anonymous pipe or a special FIFO file. The conversion assumes nothing about /// the underlying pipe; it is left up to the user to make sure that the file /// descriptor represents the writing end of a pipe and the pipe is set in /// non-blocking mode. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_owned_fd_unchecked(owned_fd: OwnedFd) -> io::Result { // Safety: OwnedFd represents a valid, open file descriptor. let mio_tx = unsafe { mio_pipe::Sender::from_raw_fd(owned_fd.into_raw_fd()) }; Sender::from_mio(mio_tx) } /// Waits for any of the requested ready states. /// /// This function can be used instead of [`writable()`] to check the returned /// ready set for [`Ready::WRITABLE`] and [`Ready::WRITE_CLOSED`] events. /// /// The function may complete without the pipe being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// [`writable()`]: Self::writable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the pipe to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with [`try_write()`]. /// /// [`try_write()`]: Self::try_write /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a writing end of a fifo /// let tx = pipe::OpenOptions::new().open_sender("path/to/a/fifo")?; /// /// loop { /// // Wait for the pipe to be writable /// tx.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match tx.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write readiness. /// /// If the pipe is not currently ready for writing, this method will /// store a clone of the `Waker` from the provided `Context`. When the pipe /// becomes ready for writing, `Waker::wake` will be called on the waker. /// /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// [`writable`]: Self::writable /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the pipe is not ready for writing. /// * `Poll::Ready(Ok(()))` if the pipe is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Tries to write a buffer to the pipe, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. If the length of `buf` is not /// greater than `PIPE_BUF` (an OS constant, 4096 under Linux), then the /// write is guaranteed to be atomic, i.e. either the entire content of /// `buf` will be written or this method will fail with `WouldBlock`. There /// is no such guarantee if `buf` is larger than `PIPE_BUF`. /// /// This function is usually paired with [`writable`]. /// /// [`writable`]: Self::writable /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the pipe is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a writing end of a fifo /// let tx = pipe::OpenOptions::new().open_sender("path/to/a/fifo")?; /// /// loop { /// // Wait for the pipe to be writable /// tx.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match tx.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write(&self, buf: &[u8]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } /// Tries to write several buffers to the pipe, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// If the total length of buffers is not greater than `PIPE_BUF` (an OS /// constant, 4096 under Linux), then the write is guaranteed to be atomic, /// i.e. either the entire contents of buffers will be written or this /// method will fail with `WouldBlock`. There is no such guarantee if the /// total length of buffers is greater than `PIPE_BUF`. /// /// This function is usually paired with [`writable`]. /// /// [`try_write()`]: Self::try_write() /// [`writable`]: Self::writable /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the pipe is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a writing end of a fifo /// let tx = pipe::OpenOptions::new().open_sender("path/to/a/fifo")?; /// /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; /// /// loop { /// // Wait for the pipe to be writable /// tx.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match tx.try_write_vectored(&bufs) { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } /// Converts the pipe into an [`OwnedFd`] in blocking mode. /// /// This function will deregister this pipe end from the event loop, set /// it in blocking mode and perform the conversion. pub fn into_blocking_fd(self) -> io::Result { let fd = self.into_nonblocking_fd()?; set_blocking(&fd)?; Ok(fd) } /// Converts the pipe into an [`OwnedFd`] in nonblocking mode. /// /// This function will deregister this pipe end from the event loop and /// perform the conversion. The returned file descriptor will be in nonblocking /// mode. pub fn into_nonblocking_fd(self) -> io::Result { let mio_pipe = self.io.into_inner()?; // Safety: the pipe is now deregistered from the event loop // and we are the only owner of this pipe end. let owned_fd = unsafe { OwnedFd::from_raw_fd(mio_pipe.into_raw_fd()) }; Ok(owned_fd) } } impl AsyncWrite for Sender { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.io.poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.io.poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } } impl AsRawFd for Sender { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for Sender { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } /// Reading end of a Unix pipe. /// /// It can be constructed from a FIFO file with [`OpenOptions::open_receiver`]. /// /// # Examples /// /// Receiving messages from a named pipe in a loop: /// /// ```no_run /// use tokio::net::unix::pipe; /// use tokio::io::{self, AsyncReadExt}; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// let mut rx = pipe::OpenOptions::new().open_receiver(FIFO_NAME)?; /// loop { /// let mut msg = vec![0; 256]; /// match rx.read_exact(&mut msg).await { /// Ok(_) => { /// /* handle the message */ /// } /// Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => { /// // Writing end has been closed, we should reopen the pipe. /// rx = pipe::OpenOptions::new().open_receiver(FIFO_NAME)?; /// } /// Err(e) => return Err(e.into()), /// } /// } /// # } /// ``` /// /// On Linux, you can use a `Receiver` in read-write access mode to implement /// resilient reading from a named pipe. Unlike `Receiver` opened in read-only /// mode, read from a pipe in read-write mode will not fail with `UnexpectedEof` /// when the writing end is closed. This way, a `Receiver` can asynchronously /// wait for the next writer to open the pipe. /// /// You should not use functions waiting for EOF such as [`read_to_end`] with /// a `Receiver` in read-write access mode, since it **may wait forever**. /// `Receiver` in this mode also holds an open writing end, which prevents /// receiving EOF. /// /// To set the read-write access mode you can use `OpenOptions::read_write`. /// Note that using read-write access mode with FIFO files is not defined by /// the POSIX standard and it is only guaranteed to work on Linux. /// /// ```ignore /// use tokio::net::unix::pipe; /// use tokio::io::AsyncReadExt; /// # use std::error::Error; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// let mut rx = pipe::OpenOptions::new() /// .read_write(true) /// .open_receiver(FIFO_NAME)?; /// loop { /// let mut msg = vec![0; 256]; /// rx.read_exact(&mut msg).await?; /// /* handle the message */ /// } /// # } /// ``` /// /// [`read_to_end`]: crate::io::AsyncReadExt::read_to_end #[derive(Debug)] pub struct Receiver { io: PollEvented, } impl Receiver { fn from_mio(mio_rx: mio_pipe::Receiver) -> io::Result { let io = PollEvented::new_with_interest(mio_rx, Interest::READABLE)?; Ok(Receiver { io }) } /// Creates a new `Receiver` from a [`File`]. /// /// This function is intended to construct a pipe from a [`File`] representing /// a special FIFO file. It will check if the file is a pipe and has read access, /// set it in non-blocking mode and perform the conversion. /// /// # Errors /// /// Fails with `io::ErrorKind::InvalidInput` if the file is not a pipe or it /// does not have read access. Also fails with any standard OS error if it occurs. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_file(file: File) -> io::Result { Receiver::from_owned_fd(file.into()) } /// Creates a new `Receiver` from an [`OwnedFd`]. /// /// This function is intended to construct a pipe from an [`OwnedFd`] representing /// an anonymous pipe or a special FIFO file. It will check if the file descriptor /// is a pipe and has read access, set it in non-blocking mode and perform the /// conversion. /// /// # Errors /// /// Fails with `io::ErrorKind::InvalidInput` if the file descriptor is not a pipe /// or it does not have read access. Also fails with any standard OS error if it /// occurs. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_owned_fd(owned_fd: OwnedFd) -> io::Result { if !is_pipe(owned_fd.as_fd())? { return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); } let flags = get_file_flags(owned_fd.as_fd())?; if has_read_access(flags) { set_nonblocking(owned_fd.as_fd(), flags)?; Receiver::from_owned_fd_unchecked(owned_fd) } else { Err(io::Error::new( io::ErrorKind::InvalidInput, "not in O_RDONLY or O_RDWR access mode", )) } } /// Creates a new `Receiver` from a [`File`] without checking pipe properties. /// /// This function is intended to construct a pipe from a File representing /// a special FIFO file. The conversion assumes nothing about the underlying /// file; it is left up to the user to make sure it is opened with read access, /// represents a pipe and is set in non-blocking mode. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::fs::OpenOptions; /// use std::os::unix::fs::{FileTypeExt, OpenOptionsExt}; /// # use std::error::Error; /// /// const FIFO_NAME: &str = "path/to/a/fifo"; /// /// # async fn dox() -> Result<(), Box> { /// let file = OpenOptions::new() /// .read(true) /// .custom_flags(libc::O_NONBLOCK) /// .open(FIFO_NAME)?; /// if file.metadata()?.file_type().is_fifo() { /// let rx = pipe::Receiver::from_file_unchecked(file)?; /// /* use the Receiver */ /// } /// # Ok(()) /// # } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_file_unchecked(file: File) -> io::Result { Receiver::from_owned_fd_unchecked(file.into()) } /// Creates a new `Receiver` from an [`OwnedFd`] without checking pipe properties. /// /// This function is intended to construct a pipe from an [`OwnedFd`] representing /// an anonymous pipe or a special FIFO file. The conversion assumes nothing about /// the underlying pipe; it is left up to the user to make sure that the file /// descriptor represents the reading end of a pipe and the pipe is set in /// non-blocking mode. /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. pub fn from_owned_fd_unchecked(owned_fd: OwnedFd) -> io::Result { // Safety: OwnedFd represents a valid, open file descriptor. let mio_rx = unsafe { mio_pipe::Receiver::from_raw_fd(owned_fd.into_raw_fd()) }; Receiver::from_mio(mio_rx) } /// Waits for any of the requested ready states. /// /// This function can be used instead of [`readable()`] to check the returned /// ready set for [`Ready::READABLE`] and [`Ready::READ_CLOSED`] events. /// /// The function may complete without the pipe being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// [`readable()`]: Self::readable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the pipe to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with [`try_read()`]. /// /// [`try_read()`]: Self::try_read() /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a reading end of a fifo /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; /// /// let mut msg = vec![0; 1024]; /// /// loop { /// // Wait for the pipe to be readable /// rx.readable().await?; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match rx.try_read(&mut msg) { /// Ok(n) => { /// msg.truncate(n); /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// println!("GOT = {:?}", msg); /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read readiness. /// /// If the pipe is not currently ready for reading, this method will /// store a clone of the `Waker` from the provided `Context`. When the pipe /// becomes ready for reading, `Waker::wake` will be called on the waker. /// /// Note that on multiple calls to `poll_read_ready` or `poll_read`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// [`readable`]: Self::readable /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the pipe is not ready for reading. /// * `Poll::Ready(Ok(()))` if the pipe is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Tries to read data from the pipe into the provided buffer, returning how /// many bytes were read. /// /// Reads any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually [`readable()`] is used with this function. /// /// [`readable()`]: Self::readable() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The pipe's writing end is closed and will no longer write data. /// 2. The specified buffer was 0 bytes in length. /// /// If the pipe is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a reading end of a fifo /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; /// /// let mut msg = vec![0; 1024]; /// /// loop { /// // Wait for the pipe to be readable /// rx.readable().await?; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match rx.try_read(&mut msg) { /// Ok(n) => { /// msg.truncate(n); /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// println!("GOT = {:?}", msg); /// Ok(()) /// } /// ``` pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read(buf)) } /// Tries to read data from the pipe into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Reads any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] is used with this function. /// /// [`try_read()`]: Self::try_read() /// [`readable()`]: Self::readable() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the pipe's writing end is /// closed and will no longer write data. If the pipe is not ready to read /// data `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a reading end of a fifo /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; /// /// loop { /// // Wait for the pipe to be readable /// rx.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf_a = [0; 512]; /// let mut buf_b = [0; 1024]; /// let mut bufs = [ /// io::IoSliceMut::new(&mut buf_a), /// io::IoSliceMut::new(&mut buf_b), /// ]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match rx.try_read_vectored(&mut bufs) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } cfg_io_util! { /// Tries to read data from the pipe into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Reads any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable /// [`ready()`]: Self::ready /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the pipe's writing end is /// closed and will no longer write data. If the pipe is not ready to read /// data `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::unix::pipe; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// // Open a reading end of a fifo /// let rx = pipe::OpenOptions::new().open_receiver("path/to/a/fifo")?; /// /// loop { /// // Wait for the pipe to be readable /// rx.readable().await?; /// /// let mut buf = Vec::with_capacity(4096); /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match rx.try_read_buf(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { use std::io::Read; let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: `mio_pipe::Receiver` uses a `std::fs::File` underneath, // which correctly handles reads into uninitialized memory. let n = (&*self.io).read(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }) } } /// Converts the pipe into an [`OwnedFd`] in blocking mode. /// /// This function will deregister this pipe end from the event loop, set /// it in blocking mode and perform the conversion. pub fn into_blocking_fd(self) -> io::Result { let fd = self.into_nonblocking_fd()?; set_blocking(&fd)?; Ok(fd) } /// Converts the pipe into an [`OwnedFd`] in nonblocking mode. /// /// This function will deregister this pipe end from the event loop and /// perform the conversion. Returned file descriptor will be in nonblocking /// mode. pub fn into_nonblocking_fd(self) -> io::Result { let mio_pipe = self.io.into_inner()?; // Safety: the pipe is now deregistered from the event loop // and we are the only owner of this pipe end. let owned_fd = unsafe { OwnedFd::from_raw_fd(mio_pipe.into_raw_fd()) }; Ok(owned_fd) } } impl AsyncRead for Receiver { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { // Safety: `mio_pipe::Receiver` uses a `std::fs::File` underneath, // which correctly handles reads into uninitialized memory. unsafe { self.io.poll_read(cx, buf) } } } impl AsRawFd for Receiver { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for Receiver { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } /// Checks if the file descriptor is a pipe or a FIFO. fn is_pipe(fd: BorrowedFd<'_>) -> io::Result { // Safety: `libc::stat` is C-like struct used for syscalls and all-zero // byte pattern forms a valid value. let mut stat: libc::stat = unsafe { std::mem::zeroed() }; // Safety: it's safe to call `fstat` with a valid, open file descriptor // and a valid pointer to a `stat` struct. let r = unsafe { libc::fstat(fd.as_raw_fd(), &mut stat) }; if r == -1 { Err(io::Error::last_os_error()) } else { Ok((stat.st_mode as libc::mode_t & libc::S_IFMT) == libc::S_IFIFO) } } /// Gets file descriptor's flags by fcntl. fn get_file_flags(fd: BorrowedFd<'_>) -> io::Result { // Safety: it's safe to use `fcntl` to read flags of a valid, open file descriptor. let flags = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFL) }; if flags < 0 { Err(io::Error::last_os_error()) } else { Ok(flags) } } /// Checks for `O_RDONLY` or `O_RDWR` access mode. fn has_read_access(flags: libc::c_int) -> bool { let mode = flags & libc::O_ACCMODE; mode == libc::O_RDONLY || mode == libc::O_RDWR } /// Checks for `O_WRONLY` or `O_RDWR` access mode. fn has_write_access(flags: libc::c_int) -> bool { let mode = flags & libc::O_ACCMODE; mode == libc::O_WRONLY || mode == libc::O_RDWR } /// Sets file descriptor's flags with `O_NONBLOCK` by fcntl. fn set_nonblocking(fd: BorrowedFd<'_>, current_flags: libc::c_int) -> io::Result<()> { let flags = current_flags | libc::O_NONBLOCK; if flags != current_flags { // Safety: it's safe to use `fcntl` to set the `O_NONBLOCK` flag of a valid, // open file descriptor. let ret = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFL, flags) }; if ret < 0 { return Err(io::Error::last_os_error()); } } Ok(()) } /// Removes `O_NONBLOCK` from fd's flags. fn set_blocking(fd: &T) -> io::Result<()> { // Safety: it's safe to use `fcntl` to read flags of a valid, open file descriptor. let previous = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFL) }; if previous == -1 { return Err(io::Error::last_os_error()); } let new = previous & !libc::O_NONBLOCK; // Safety: it's safe to use `fcntl` to unset the `O_NONBLOCK` flag of a valid, // open file descriptor. let r = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFL, new) }; if r == -1 { Err(io::Error::last_os_error()) } else { Ok(()) } } tokio-1.43.1/src/net/unix/socket.rs000064400000000000000000000214441046102023000152420ustar 00000000000000use std::io; use std::path::Path; use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use crate::net::{UnixDatagram, UnixListener, UnixStream}; cfg_net_unix! { /// A Unix socket that has not yet been converted to a [`UnixStream`], [`UnixDatagram`], or /// [`UnixListener`]. /// /// `UnixSocket` wraps an operating system socket and enables the caller to /// configure the socket before establishing a connection or accepting /// inbound connections. The caller is able to set socket option and explicitly /// bind the socket with a socket address. /// /// The underlying socket is closed when the `UnixSocket` value is dropped. /// /// `UnixSocket` should only be used directly if the default configuration used /// by [`UnixStream::connect`], [`UnixDatagram::bind`], and [`UnixListener::bind`] /// does not meet the required use case. /// /// Calling `UnixStream::connect(path)` effectively performs the same function as: /// /// ```no_run /// use tokio::net::UnixSocket; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let path = dir.path().join("bind_path"); /// let socket = UnixSocket::new_stream()?; /// /// let stream = socket.connect(path).await?; /// /// Ok(()) /// } /// ``` /// /// Calling `UnixDatagram::bind(path)` effectively performs the same function as: /// /// ```no_run /// use tokio::net::UnixSocket; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let path = dir.path().join("bind_path"); /// let socket = UnixSocket::new_datagram()?; /// socket.bind(path)?; /// /// let datagram = socket.datagram()?; /// /// Ok(()) /// } /// ``` /// /// Calling `UnixListener::bind(path)` effectively performs the same function as: /// /// ```no_run /// use tokio::net::UnixSocket; /// use std::error::Error; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let path = dir.path().join("bind_path"); /// let socket = UnixSocket::new_stream()?; /// socket.bind(path)?; /// /// let listener = socket.listen(1024)?; /// /// Ok(()) /// } /// ``` /// /// Setting socket options not explicitly provided by `UnixSocket` may be done by /// accessing the [`RawFd`]/[`RawSocket`] using [`AsRawFd`]/[`AsRawSocket`] and /// setting the option with a crate like [`socket2`]. /// /// [`RawFd`]: std::os::fd::RawFd /// [`RawSocket`]: https://doc.rust-lang.org/std/os/windows/io/type.RawSocket.html /// [`AsRawFd`]: std::os::fd::AsRawFd /// [`AsRawSocket`]: https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html /// [`socket2`]: https://docs.rs/socket2/ #[derive(Debug)] pub struct UnixSocket { inner: socket2::Socket, } } impl UnixSocket { fn ty(&self) -> socket2::Type { self.inner.r#type().unwrap() } /// Creates a new Unix datagram socket. /// /// Calls `socket(2)` with `AF_UNIX` and `SOCK_DGRAM`. /// /// # Returns /// /// On success, the newly created [`UnixSocket`] is returned. If an error is /// encountered, it is returned instead. pub fn new_datagram() -> io::Result { UnixSocket::new(socket2::Type::DGRAM) } /// Creates a new Unix stream socket. /// /// Calls `socket(2)` with `AF_UNIX` and `SOCK_STREAM`. /// /// # Returns /// /// On success, the newly created [`UnixSocket`] is returned. If an error is /// encountered, it is returned instead. pub fn new_stream() -> io::Result { UnixSocket::new(socket2::Type::STREAM) } fn new(ty: socket2::Type) -> io::Result { #[cfg(any( target_os = "android", target_os = "dragonfly", target_os = "freebsd", target_os = "fuchsia", target_os = "illumos", target_os = "linux", target_os = "netbsd", target_os = "openbsd" ))] let ty = ty.nonblocking(); let inner = socket2::Socket::new(socket2::Domain::UNIX, ty, None)?; #[cfg(not(any( target_os = "android", target_os = "dragonfly", target_os = "freebsd", target_os = "fuchsia", target_os = "illumos", target_os = "linux", target_os = "netbsd", target_os = "openbsd" )))] inner.set_nonblocking(true)?; Ok(UnixSocket { inner }) } /// Binds the socket to the given address. /// /// This calls the `bind(2)` operating-system function. pub fn bind(&self, path: impl AsRef) -> io::Result<()> { let addr = socket2::SockAddr::unix(path)?; self.inner.bind(&addr) } /// Converts the socket into a `UnixListener`. /// /// `backlog` defines the maximum number of pending connections are queued /// by the operating system at any given time. Connection are removed from /// the queue with [`UnixListener::accept`]. When the queue is full, the /// operating-system will start rejecting connections. /// /// Calling this function on a socket created by [`new_datagram`] will return an error. /// /// This calls the `listen(2)` operating-system function, marking the socket /// as a passive socket. /// /// [`new_datagram`]: `UnixSocket::new_datagram` pub fn listen(self, backlog: u32) -> io::Result { if self.ty() == socket2::Type::DGRAM { return Err(io::Error::new( io::ErrorKind::Other, "listen cannot be called on a datagram socket", )); } self.inner.listen(backlog as i32)?; let mio = { use std::os::unix::io::{FromRawFd, IntoRawFd}; let raw_fd = self.inner.into_raw_fd(); unsafe { mio::net::UnixListener::from_raw_fd(raw_fd) } }; UnixListener::new(mio) } /// Establishes a Unix connection with a peer at the specified socket address. /// /// The `UnixSocket` is consumed. Once the connection is established, a /// connected [`UnixStream`] is returned. If the connection fails, the /// encountered error is returned. /// /// Calling this function on a socket created by [`new_datagram`] will return an error. /// /// This calls the `connect(2)` operating-system function. /// /// [`new_datagram`]: `UnixSocket::new_datagram` pub async fn connect(self, path: impl AsRef) -> io::Result { if self.ty() == socket2::Type::DGRAM { return Err(io::Error::new( io::ErrorKind::Other, "connect cannot be called on a datagram socket", )); } let addr = socket2::SockAddr::unix(path)?; if let Err(err) = self.inner.connect(&addr) { if err.raw_os_error() != Some(libc::EINPROGRESS) { return Err(err); } } let mio = { use std::os::unix::io::{FromRawFd, IntoRawFd}; let raw_fd = self.inner.into_raw_fd(); unsafe { mio::net::UnixStream::from_raw_fd(raw_fd) } }; UnixStream::connect_mio(mio).await } /// Converts the socket into a [`UnixDatagram`]. /// /// Calling this function on a socket created by [`new_stream`] will return an error. /// /// [`new_stream`]: `UnixSocket::new_stream` pub fn datagram(self) -> io::Result { if self.ty() == socket2::Type::STREAM { return Err(io::Error::new( io::ErrorKind::Other, "datagram cannot be called on a stream socket", )); } let mio = { use std::os::unix::io::{FromRawFd, IntoRawFd}; let raw_fd = self.inner.into_raw_fd(); unsafe { mio::net::UnixDatagram::from_raw_fd(raw_fd) } }; UnixDatagram::from_mio(mio) } } impl AsRawFd for UnixSocket { fn as_raw_fd(&self) -> RawFd { self.inner.as_raw_fd() } } impl AsFd for UnixSocket { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } impl FromRawFd for UnixSocket { unsafe fn from_raw_fd(fd: RawFd) -> UnixSocket { let inner = socket2::Socket::from_raw_fd(fd); UnixSocket { inner } } } impl IntoRawFd for UnixSocket { fn into_raw_fd(self) -> RawFd { self.inner.into_raw_fd() } } tokio-1.43.1/src/net/unix/socketaddr.rs000064400000000000000000000024141046102023000160710ustar 00000000000000use std::fmt; use std::path::Path; /// An address associated with a Tokio Unix socket. /// /// This type is a thin wrapper around [`std::os::unix::net::SocketAddr`]. You /// can convert to and from the standard library `SocketAddr` type using the /// [`From`] trait. pub struct SocketAddr(pub(super) std::os::unix::net::SocketAddr); impl SocketAddr { /// Returns `true` if the address is unnamed. /// /// Documentation reflected in [`SocketAddr`] /// /// [`SocketAddr`]: std::os::unix::net::SocketAddr pub fn is_unnamed(&self) -> bool { self.0.is_unnamed() } /// Returns the contents of this address if it is a `pathname` address. /// /// Documentation reflected in [`SocketAddr`] /// /// [`SocketAddr`]: std::os::unix::net::SocketAddr pub fn as_pathname(&self) -> Option<&Path> { self.0.as_pathname() } } impl fmt::Debug for SocketAddr { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.0.fmt(fmt) } } impl From for SocketAddr { fn from(value: std::os::unix::net::SocketAddr) -> Self { SocketAddr(value) } } impl From for std::os::unix::net::SocketAddr { fn from(value: SocketAddr) -> Self { value.0 } } tokio-1.43.1/src/net/unix/split.rs000064400000000000000000000300201046102023000150730ustar 00000000000000//! `UnixStream` split support. //! //! A `UnixStream` can be split into a read half and a write half with //! `UnixStream::split`. The read half implements `AsyncRead` while the write //! half implements `AsyncWrite`. //! //! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized //! split has no associated overhead and enforces all invariants at the type //! level. use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::UnixStream; use crate::net::unix::SocketAddr; use std::io; use std::net::Shutdown; use std::pin::Pin; use std::task::{Context, Poll}; cfg_io_util! { use bytes::BufMut; } /// Borrowed read half of a [`UnixStream`], created by [`split`]. /// /// Reading from a `ReadHalf` is usually done using the convenience methods found on the /// [`AsyncReadExt`] trait. /// /// [`UnixStream`]: UnixStream /// [`split`]: UnixStream::split() /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt #[derive(Debug)] pub struct ReadHalf<'a>(&'a UnixStream); /// Borrowed write half of a [`UnixStream`], created by [`split`]. /// /// Note that in the [`AsyncWrite`] implementation of this type, [`poll_shutdown`] will /// shut down the [`UnixStream`] stream in the write direction. /// /// Writing to an `WriteHalf` is usually done using the convenience methods found /// on the [`AsyncWriteExt`] trait. /// /// [`UnixStream`]: UnixStream /// [`split`]: UnixStream::split() /// [`AsyncWrite`]: trait@crate::io::AsyncWrite /// [`poll_shutdown`]: fn@crate::io::AsyncWrite::poll_shutdown /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt #[derive(Debug)] pub struct WriteHalf<'a>(&'a UnixStream); pub(crate) fn split(stream: &mut UnixStream) -> (ReadHalf<'_>, WriteHalf<'_>) { (ReadHalf(stream), WriteHalf(stream)) } impl ReadHalf<'_> { /// Wait for any of the requested ready states. /// /// This function is usually paired with [`try_read()`]. It can be used instead /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] /// and [`Ready::READ_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`UnixStream::ready`]. /// /// [`try_read()`]: Self::try_read /// [`readable()`]: Self::readable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.0.ready(interest).await } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn readable(&self) -> io::Result<()> { self.0.readable().await } /// Tries to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The stream's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.0.try_read(buf) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.0.try_read_buf(buf) } } /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: Self::try_read() /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.0.try_read_vectored(bufs) } /// Returns the socket address of the remote half of this connection. pub fn peer_addr(&self) -> io::Result { self.0.peer_addr() } /// Returns the socket address of the local half of this connection. pub fn local_addr(&self) -> io::Result { self.0.local_addr() } } impl WriteHalf<'_> { /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_write()`]. It can be used instead /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] /// and [`Ready::WRITE_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`UnixStream::ready`]. /// /// [`try_write()`]: Self::try_write /// [`writable()`]: Self::writable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.0.ready(interest).await } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn writable(&self) -> io::Result<()> { self.0.writable().await } /// Tries to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write(&self, buf: &[u8]) -> io::Result { self.0.try_write(buf) } /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: Self::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result { self.0.try_write_vectored(buf) } /// Returns the socket address of the remote half of this connection. pub fn peer_addr(&self) -> io::Result { self.0.peer_addr() } /// Returns the socket address of the local half of this connection. pub fn local_addr(&self) -> io::Result { self.0.local_addr() } } impl AsyncRead for ReadHalf<'_> { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.0.poll_read_priv(cx, buf) } } impl AsyncWrite for WriteHalf<'_> { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.0.poll_write_priv(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.0.poll_write_vectored_priv(cx, bufs) } fn is_write_vectored(&self) -> bool { self.0.is_write_vectored() } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { self.0.shutdown_std(Shutdown::Write).into() } } impl AsRef for ReadHalf<'_> { fn as_ref(&self) -> &UnixStream { self.0 } } impl AsRef for WriteHalf<'_> { fn as_ref(&self) -> &UnixStream { self.0 } } tokio-1.43.1/src/net/unix/split_owned.rs000064400000000000000000000357371046102023000163130ustar 00000000000000//! `UnixStream` owned split support. //! //! A `UnixStream` can be split into an `OwnedReadHalf` and a `OwnedWriteHalf` //! with the `UnixStream::into_split` method. `OwnedReadHalf` implements //! `AsyncRead` while `OwnedWriteHalf` implements `AsyncWrite`. //! //! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized //! split has no associated overhead and enforces all invariants at the type //! level. use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; use crate::net::UnixStream; use crate::net::unix::SocketAddr; use std::error::Error; use std::net::Shutdown; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use std::{fmt, io}; cfg_io_util! { use bytes::BufMut; } /// Owned read half of a [`UnixStream`], created by [`into_split`]. /// /// Reading from an `OwnedReadHalf` is usually done using the convenience methods found /// on the [`AsyncReadExt`] trait. /// /// [`UnixStream`]: crate::net::UnixStream /// [`into_split`]: crate::net::UnixStream::into_split() /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt #[derive(Debug)] pub struct OwnedReadHalf { inner: Arc, } /// Owned write half of a [`UnixStream`], created by [`into_split`]. /// /// Note that in the [`AsyncWrite`] implementation of this type, /// [`poll_shutdown`] will shut down the stream in the write direction. /// Dropping the write half will also shut down the write half of the stream. /// /// Writing to an `OwnedWriteHalf` is usually done using the convenience methods /// found on the [`AsyncWriteExt`] trait. /// /// [`UnixStream`]: crate::net::UnixStream /// [`into_split`]: crate::net::UnixStream::into_split() /// [`AsyncWrite`]: trait@crate::io::AsyncWrite /// [`poll_shutdown`]: fn@crate::io::AsyncWrite::poll_shutdown /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt #[derive(Debug)] pub struct OwnedWriteHalf { inner: Arc, shutdown_on_drop: bool, } pub(crate) fn split_owned(stream: UnixStream) -> (OwnedReadHalf, OwnedWriteHalf) { let arc = Arc::new(stream); let read = OwnedReadHalf { inner: Arc::clone(&arc), }; let write = OwnedWriteHalf { inner: arc, shutdown_on_drop: true, }; (read, write) } pub(crate) fn reunite( read: OwnedReadHalf, write: OwnedWriteHalf, ) -> Result { if Arc::ptr_eq(&read.inner, &write.inner) { write.forget(); // This unwrap cannot fail as the api does not allow creating more than two Arcs, // and we just dropped the other half. Ok(Arc::try_unwrap(read.inner).expect("UnixStream: try_unwrap failed in reunite")) } else { Err(ReuniteError(read, write)) } } /// Error indicating that two halves were not from the same socket, and thus could /// not be reunited. #[derive(Debug)] pub struct ReuniteError(pub OwnedReadHalf, pub OwnedWriteHalf); impl fmt::Display for ReuniteError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, "tried to reunite halves that are not from the same socket" ) } } impl Error for ReuniteError {} impl OwnedReadHalf { /// Attempts to put the two halves of a `UnixStream` back together and /// recover the original socket. Succeeds only if the two halves /// originated from the same call to [`into_split`]. /// /// [`into_split`]: crate::net::UnixStream::into_split() pub fn reunite(self, other: OwnedWriteHalf) -> Result { reunite(self, other) } /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_read()`]. It can be used instead /// of [`readable()`] to check the returned ready set for [`Ready::READABLE`] /// and [`Ready::READ_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`UnixStream::ready`]. /// /// [`try_read()`]: Self::try_read /// [`readable()`]: Self::readable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.inner.ready(interest).await } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn readable(&self) -> io::Result<()> { self.inner.readable().await } /// Tries to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The stream's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.inner.try_read(buf) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.inner.try_read_buf(buf) } } /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: Self::try_read() /// [`readable()`]: Self::readable() /// [`ready()`]: Self::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.inner.try_read_vectored(bufs) } /// Returns the socket address of the remote half of this connection. pub fn peer_addr(&self) -> io::Result { self.inner.peer_addr() } /// Returns the socket address of the local half of this connection. pub fn local_addr(&self) -> io::Result { self.inner.local_addr() } } impl AsyncRead for OwnedReadHalf { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.inner.poll_read_priv(cx, buf) } } impl OwnedWriteHalf { /// Attempts to put the two halves of a `UnixStream` back together and /// recover the original socket. Succeeds only if the two halves /// originated from the same call to [`into_split`]. /// /// [`into_split`]: crate::net::UnixStream::into_split() pub fn reunite(self, other: OwnedReadHalf) -> Result { reunite(other, self) } /// Destroys the write half, but don't close the write half of the stream /// until the read half is dropped. If the read half has already been /// dropped, this closes the stream. pub fn forget(mut self) { self.shutdown_on_drop = false; drop(self); } /// Waits for any of the requested ready states. /// /// This function is usually paired with [`try_write()`]. It can be used instead /// of [`writable()`] to check the returned ready set for [`Ready::WRITABLE`] /// and [`Ready::WRITE_CLOSED`] events. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// This function is equivalent to [`UnixStream::ready`]. /// /// [`try_write()`]: Self::try_write /// [`writable()`]: Self::writable /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn ready(&self, interest: Interest) -> io::Result { self.inner.ready(interest).await } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. pub async fn writable(&self) -> io::Result<()> { self.inner.writable().await } /// Tries to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write(&self, buf: &[u8]) -> io::Result { self.inner.try_write(buf) } /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: Self::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result { self.inner.try_write_vectored(buf) } /// Returns the socket address of the remote half of this connection. pub fn peer_addr(&self) -> io::Result { self.inner.peer_addr() } /// Returns the socket address of the local half of this connection. pub fn local_addr(&self) -> io::Result { self.inner.local_addr() } } impl Drop for OwnedWriteHalf { fn drop(&mut self) { if self.shutdown_on_drop { let _ = self.inner.shutdown_std(Shutdown::Write); } } } impl AsyncWrite for OwnedWriteHalf { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.inner.poll_write_priv(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.inner.poll_write_vectored_priv(cx, bufs) } fn is_write_vectored(&self) -> bool { self.inner.is_write_vectored() } #[inline] fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { // flush is a no-op Poll::Ready(Ok(())) } // `poll_shutdown` on a write half shutdowns the stream in the "write" direction. fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { let res = self.inner.shutdown_std(Shutdown::Write); if res.is_ok() { Pin::into_inner(self).shutdown_on_drop = false; } res.into() } } impl AsRef for OwnedReadHalf { fn as_ref(&self) -> &UnixStream { &self.inner } } impl AsRef for OwnedWriteHalf { fn as_ref(&self) -> &UnixStream { &self.inner } } tokio-1.43.1/src/net/unix/stream.rs000064400000000000000000001160651046102023000152510ustar 00000000000000use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; use crate::net::unix::split::{split, ReadHalf, WriteHalf}; use crate::net::unix::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf}; use crate::net::unix::ucred::{self, UCred}; use crate::net::unix::SocketAddr; use std::fmt; use std::future::poll_fn; use std::io::{self, Read, Write}; use std::net::Shutdown; #[cfg(target_os = "android")] use std::os::android::net::SocketAddrExt; #[cfg(target_os = "linux")] use std::os::linux::net::SocketAddrExt; #[cfg(any(target_os = "linux", target_os = "android"))] use std::os::unix::ffi::OsStrExt; use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net::{self, SocketAddr as StdSocketAddr}; use std::path::Path; use std::pin::Pin; use std::task::{Context, Poll}; cfg_io_util! { use bytes::BufMut; } cfg_net_unix! { /// A structure representing a connected Unix socket. /// /// This socket can be connected directly with [`UnixStream::connect`] or accepted /// from a listener with [`UnixListener::accept`]. Additionally, a pair of /// anonymous Unix sockets can be created with `UnixStream::pair`. /// /// To shut down the stream in the write direction, you can call the /// [`shutdown()`] method. This will cause the other peer to receive a read of /// length 0, indicating that no more data will be sent. This only closes /// the stream in one direction. /// /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown /// [`UnixListener::accept`]: crate::net::UnixListener::accept #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixStream { io: PollEvented, } } impl UnixStream { pub(crate) async fn connect_mio(sys: mio::net::UnixStream) -> io::Result { let stream = UnixStream::new(sys)?; // Once we've connected, wait for the stream to be writable as // that's when the actual connection has been initiated. Once we're // writable we check for `take_socket_error` to see if the connect // actually hit an error or not. // // If all that succeeded then we ship everything on up. poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?; if let Some(e) = stream.io.take_error()? { return Err(e); } Ok(stream) } /// Connects to the socket named by `path`. /// /// This function will create a new Unix socket and connect to the path /// specified, associating the returned stream with the default event loop's /// handle. pub async fn connect

(path: P) -> io::Result where P: AsRef, { // On linux, abstract socket paths need to be considered. #[cfg(any(target_os = "linux", target_os = "android"))] let addr = { let os_str_bytes = path.as_ref().as_os_str().as_bytes(); if os_str_bytes.starts_with(b"\0") { StdSocketAddr::from_abstract_name(&os_str_bytes[1..])? } else { StdSocketAddr::from_pathname(path)? } }; #[cfg(not(any(target_os = "linux", target_os = "android")))] let addr = StdSocketAddr::from_pathname(path)?; let stream = mio::net::UnixStream::connect_addr(&addr)?; let stream = UnixStream::new(stream)?; poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?; if let Some(e) = stream.io.take_error()? { return Err(e); } Ok(stream) } /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same socket on a single /// task without splitting the socket. /// /// The function may complete without the socket being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read or write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// Concurrently read and write to the stream on the same task without /// splitting. /// /// ```no_run /// use tokio::io::Interest; /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// loop { /// let ready = stream.ready(Interest::READABLE | Interest::WRITABLE).await?; /// /// if ready.is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// /// } /// /// if ready.is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the socket to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to read that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// let mut msg = vec![0; 1024]; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read(&mut msg) { /// Ok(n) => { /// msg.truncate(n); /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// println!("GOT = {:?}", msg); /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read readiness. /// /// If the unix stream is not currently ready for reading, this method will /// store a clone of the `Waker` from the provided `Context`. When the unix /// stream becomes ready for reading, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_read_ready` or `poll_read`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_write_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the unix stream is not ready for reading. /// * `Poll::Ready(Ok(()))` if the unix stream is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`readable`]: method@Self::readable pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Try to read data from the stream into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: UnixStream::readable() /// [`ready()`]: UnixStream::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The stream's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the stream is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf = [0; 4096]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read(buf)) } /// Tries to read data from the stream into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: UnixStream::try_read() /// [`readable()`]: UnixStream::readable() /// [`ready()`]: UnixStream::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io::{self, IoSliceMut}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf_a = [0; 512]; /// let mut buf_b = [0; 1024]; /// let mut bufs = [ /// IoSliceMut::new(&mut buf_a), /// IoSliceMut::new(&mut buf_b), /// ]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read_vectored(&mut bufs) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the socket but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: UnixStream::readable() /// [`ready()`]: UnixStream::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// loop { /// // Wait for the socket to be readable /// stream.readable().await?; /// /// let mut buf = Vec::with_capacity(4096); /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_read_buf(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { use std::io::Read; let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `UnixStream::read` to have filled up `n` bytes in the // buffer. let n = (&*self.io).read(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }) } } /// Waits for the socket to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Cancel safety /// /// This method is cancel safe. Once a readiness event occurs, the method /// will continue to return immediately until the readiness event is /// consumed by an attempt to write that fails with `WouldBlock` or /// `Poll::Pending`. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// loop { /// // Wait for the socket to be writable /// stream.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write readiness. /// /// If the unix stream is not currently ready for writing, this method will /// store a clone of the `Waker` from the provided `Context`. When the unix /// stream becomes ready for writing, `Waker::wake` will be called on the /// waker. /// /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_read_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the unix stream is not ready for writing. /// * `Poll::Ready(Ok(()))` if the unix stream is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`writable`]: method@Self::writable pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Tries to write a buffer to the stream, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// loop { /// // Wait for the socket to be writable /// stream.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write(&self, buf: &[u8]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } /// Tries to write several buffers to the stream, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: UnixStream::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the stream is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::error::Error; /// use std::io; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // Connect to a peer /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; /// /// loop { /// // Wait for the socket to be writable /// stream.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match stream.try_write_vectored(&bufs) { /// Ok(n) => { /// break; /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } /// Tries to read or write from the socket using a user-provided IO operation. /// /// If the socket is ready, the provided closure is called. The closure /// should attempt to perform IO operation on the socket by manually /// calling the appropriate syscall. If the operation fails because the /// socket is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the socket is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `UnixStream` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: UnixStream::readable() /// [`writable()`]: UnixStream::writable() /// [`ready()`]: UnixStream::ready() pub fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { self.io .registration() .try_io(interest, || self.io.try_io(f)) } /// Reads or writes from the socket using a user-provided IO operation. /// /// The readiness of the socket is awaited and when the socket is ready, /// the provided closure is called. The closure should attempt to perform /// IO operation on the socket by manually calling the appropriate syscall. /// If the operation fails because the socket is not actually ready, /// then the closure should return a `WouldBlock` error. In such case the /// readiness flag is cleared and the socket readiness is awaited again. /// This loop is repeated until the closure returns an `Ok` or an error /// other than `WouldBlock`. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the socket that failed due to the socket not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the socket to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `UnixStream` type, as this will mess with the /// readiness flag and can cause the socket to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub async fn async_io( &self, interest: Interest, mut f: impl FnMut() -> io::Result, ) -> io::Result { self.io .registration() .async_io(interest, || self.io.try_io(&mut f)) .await } /// Creates new [`UnixStream`] from a [`std::os::unix::net::UnixStream`]. /// /// This function is intended to be used to wrap a `UnixStream` from the /// standard library in the Tokio equivalent. /// /// # Notes /// /// The caller is responsible for ensuring that the stream is in /// non-blocking mode. Otherwise all I/O operations on the stream /// will block the thread, which will cause unexpected behavior. /// Non-blocking mode can be set using [`set_nonblocking`]. /// /// [`set_nonblocking`]: std::os::unix::net::UnixStream::set_nonblocking /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// use std::os::unix::net::UnixStream as StdUnixStream; /// # use std::error::Error; /// /// # async fn dox() -> Result<(), Box> { /// let std_stream = StdUnixStream::connect("/path/to/the/socket")?; /// std_stream.set_nonblocking(true)?; /// let stream = UnixStream::from_std(std_stream)?; /// # Ok(()) /// # } /// ``` /// /// # Panics /// /// This function panics if it is not called from within a runtime with /// IO enabled. /// /// The runtime is usually set implicitly when this function is called /// from a future driven by a tokio runtime, otherwise runtime can be set /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. #[track_caller] pub fn from_std(stream: net::UnixStream) -> io::Result { let stream = mio::net::UnixStream::from_std(stream); let io = PollEvented::new(stream)?; Ok(UnixStream { io }) } /// Turns a [`tokio::net::UnixStream`] into a [`std::os::unix::net::UnixStream`]. /// /// The returned [`std::os::unix::net::UnixStream`] will have nonblocking /// mode set as `true`. Use [`set_nonblocking`] to change the blocking /// mode if needed. /// /// # Examples /// /// ``` /// # if cfg!(miri) { return } // No `socket` in miri. /// use std::error::Error; /// use std::io::Read; /// use tokio::net::UnixListener; /// # use tokio::net::UnixStream; /// # use tokio::io::AsyncWriteExt; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// /// let mut data = [0u8; 12]; /// let listener = UnixListener::bind(&bind_path)?; /// # let handle = tokio::spawn(async { /// # let mut stream = UnixStream::connect(bind_path).await.unwrap(); /// # stream.write(b"Hello world!").await.unwrap(); /// # }); /// let (tokio_unix_stream, _) = listener.accept().await?; /// let mut std_unix_stream = tokio_unix_stream.into_std()?; /// # handle.await.expect("The task being joined has panicked"); /// std_unix_stream.set_nonblocking(false)?; /// std_unix_stream.read_exact(&mut data)?; /// # assert_eq!(b"Hello world!", &data); /// Ok(()) /// } /// ``` /// [`tokio::net::UnixStream`]: UnixStream /// [`std::os::unix::net::UnixStream`]: std::os::unix::net::UnixStream /// [`set_nonblocking`]: fn@std::os::unix::net::UnixStream::set_nonblocking pub fn into_std(self) -> io::Result { self.io .into_inner() .map(IntoRawFd::into_raw_fd) .map(|raw_fd| unsafe { std::os::unix::net::UnixStream::from_raw_fd(raw_fd) }) } /// Creates an unnamed pair of connected sockets. /// /// This function will create a pair of interconnected Unix sockets for /// communicating back and forth between one another. Each socket will /// be associated with the default event loop's handle. pub fn pair() -> io::Result<(UnixStream, UnixStream)> { let (a, b) = mio::net::UnixStream::pair()?; let a = UnixStream::new(a)?; let b = UnixStream::new(b)?; Ok((a, b)) } pub(crate) fn new(stream: mio::net::UnixStream) -> io::Result { let io = PollEvented::new(stream)?; Ok(UnixStream { io }) } /// Returns the socket address of the local half of this connection. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// /// # async fn dox() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// println!("{:?}", stream.local_addr()?); /// # Ok(()) /// # } /// ``` pub fn local_addr(&self) -> io::Result { self.io.local_addr().map(SocketAddr) } /// Returns the socket address of the remote half of this connection. /// /// # Examples /// /// ```no_run /// use tokio::net::UnixStream; /// /// # async fn dox() -> Result<(), Box> { /// let dir = tempfile::tempdir().unwrap(); /// let bind_path = dir.path().join("bind_path"); /// let stream = UnixStream::connect(bind_path).await?; /// /// println!("{:?}", stream.peer_addr()?); /// # Ok(()) /// # } /// ``` pub fn peer_addr(&self) -> io::Result { self.io.peer_addr().map(SocketAddr) } /// Returns effective credentials of the process which called `connect` or `pair`. pub fn peer_cred(&self) -> io::Result { ucred::get_peer_cred(self) } /// Returns the value of the `SO_ERROR` option. pub fn take_error(&self) -> io::Result> { self.io.take_error() } /// Shuts down the read, write, or both halves of this connection. /// /// This function will cause all pending and future I/O calls on the /// specified portions to immediately return with an appropriate value /// (see the documentation of `Shutdown`). pub(super) fn shutdown_std(&self, how: Shutdown) -> io::Result<()> { self.io.shutdown(how) } // These lifetime markers also appear in the generated documentation, and make // it more clear that this is a *borrowed* split. #[allow(clippy::needless_lifetimes)] /// Splits a `UnixStream` into a read half and a write half, which can be used /// to read and write the stream concurrently. /// /// This method is more efficient than [`into_split`], but the halves cannot be /// moved into independently spawned tasks. /// /// [`into_split`]: Self::into_split() pub fn split<'a>(&'a mut self) -> (ReadHalf<'a>, WriteHalf<'a>) { split(self) } /// Splits a `UnixStream` into a read half and a write half, which can be used /// to read and write the stream concurrently. /// /// Unlike [`split`], the owned halves can be moved to separate tasks, however /// this comes at the cost of a heap allocation. /// /// **Note:** Dropping the write half will shut down the write half of the /// stream. This is equivalent to calling [`shutdown()`] on the `UnixStream`. /// /// [`split`]: Self::split() /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown pub fn into_split(self) -> (OwnedReadHalf, OwnedWriteHalf) { split_owned(self) } } impl TryFrom for UnixStream { type Error = io::Error; /// Consumes stream, returning the tokio I/O object. /// /// This is equivalent to /// [`UnixStream::from_std(stream)`](UnixStream::from_std). fn try_from(stream: net::UnixStream) -> io::Result { Self::from_std(stream) } } impl AsyncRead for UnixStream { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.poll_read_priv(cx, buf) } } impl AsyncWrite for UnixStream { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.poll_write_priv(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.poll_write_vectored_priv(cx, bufs) } fn is_write_vectored(&self) -> bool { true } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { self.shutdown_std(std::net::Shutdown::Write)?; Poll::Ready(Ok(())) } } impl UnixStream { // == Poll IO functions that takes `&self` == // // To read or write without mutable access to the `UnixStream`, combine the // `poll_read_ready` or `poll_write_ready` methods with the `try_read` or // `try_write` methods. pub(crate) fn poll_read_priv( &self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { // Safety: `UnixStream::read` correctly handles reads into uninitialized memory unsafe { self.io.poll_read(cx, buf) } } pub(crate) fn poll_write_priv( &self, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.io.poll_write(cx, buf) } pub(super) fn poll_write_vectored_priv( &self, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.io.poll_write_vectored(cx, bufs) } } impl fmt::Debug for UnixStream { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.io.fmt(f) } } impl AsRawFd for UnixStream { fn as_raw_fd(&self) -> RawFd { self.io.as_raw_fd() } } impl AsFd for UnixStream { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } tokio-1.43.1/src/net/unix/ucred.rs000064400000000000000000000223111046102023000150460ustar 00000000000000use crate::net::unix; /// Credentials of a process. #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] pub struct UCred { /// PID (process ID) of the process. pid: Option, /// UID (user ID) of the process. uid: unix::uid_t, /// GID (group ID) of the process. gid: unix::gid_t, } impl UCred { /// Gets UID (user ID) of the process. pub fn uid(&self) -> unix::uid_t { self.uid } /// Gets GID (group ID) of the process. pub fn gid(&self) -> unix::gid_t { self.gid } /// Gets PID (process ID) of the process. /// /// This is only implemented under Linux, Android, iOS, macOS, Solaris and /// Illumos. On other platforms this will always return `None`. pub fn pid(&self) -> Option { self.pid } } #[cfg(any( target_os = "linux", target_os = "redox", target_os = "android", target_os = "openbsd", target_os = "haiku" ))] pub(crate) use self::impl_linux::get_peer_cred; #[cfg(any(target_os = "netbsd", target_os = "nto"))] pub(crate) use self::impl_netbsd::get_peer_cred; #[cfg(any(target_os = "dragonfly", target_os = "freebsd"))] pub(crate) use self::impl_bsd::get_peer_cred; #[cfg(any( target_os = "macos", target_os = "ios", target_os = "tvos", target_os = "watchos", target_os = "visionos" ))] pub(crate) use self::impl_macos::get_peer_cred; #[cfg(any(target_os = "solaris", target_os = "illumos"))] pub(crate) use self::impl_solaris::get_peer_cred; #[cfg(target_os = "aix")] pub(crate) use self::impl_aix::get_peer_cred; #[cfg(any(target_os = "espidf", target_os = "vita"))] pub(crate) use self::impl_noproc::get_peer_cred; #[cfg(any( target_os = "linux", target_os = "redox", target_os = "android", target_os = "openbsd", target_os = "haiku" ))] pub(crate) mod impl_linux { use crate::net::unix::{self, UnixStream}; use libc::{c_void, getsockopt, socklen_t, SOL_SOCKET, SO_PEERCRED}; use std::{io, mem}; #[cfg(target_os = "openbsd")] use libc::sockpeercred as ucred; #[cfg(any( target_os = "linux", target_os = "redox", target_os = "android", target_os = "haiku" ))] use libc::ucred; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result { use std::os::unix::io::AsRawFd; unsafe { let raw_fd = sock.as_raw_fd(); let mut ucred = ucred { pid: 0, uid: 0, gid: 0, }; let ucred_size = mem::size_of::(); // These paranoid checks should be optimized-out assert!(mem::size_of::() <= mem::size_of::()); assert!(ucred_size <= u32::MAX as usize); let mut ucred_size = ucred_size as socklen_t; let ret = getsockopt( raw_fd, SOL_SOCKET, SO_PEERCRED, &mut ucred as *mut ucred as *mut c_void, &mut ucred_size, ); if ret == 0 && ucred_size as usize == mem::size_of::() { Ok(super::UCred { uid: ucred.uid as unix::uid_t, gid: ucred.gid as unix::gid_t, pid: Some(ucred.pid as unix::pid_t), }) } else { Err(io::Error::last_os_error()) } } } } #[cfg(any(target_os = "netbsd", target_os = "nto"))] pub(crate) mod impl_netbsd { use crate::net::unix::{self, UnixStream}; use libc::{c_void, getsockopt, socklen_t, unpcbid, LOCAL_PEEREID, SOL_SOCKET}; use std::io; use std::mem::size_of; use std::os::unix::io::AsRawFd; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result { unsafe { let raw_fd = sock.as_raw_fd(); let mut unpcbid = unpcbid { unp_pid: 0, unp_euid: 0, unp_egid: 0, }; let unpcbid_size = size_of::(); let mut unpcbid_size = unpcbid_size as socklen_t; let ret = getsockopt( raw_fd, SOL_SOCKET, LOCAL_PEEREID, &mut unpcbid as *mut unpcbid as *mut c_void, &mut unpcbid_size, ); if ret == 0 && unpcbid_size as usize == size_of::() { Ok(super::UCred { uid: unpcbid.unp_euid as unix::uid_t, gid: unpcbid.unp_egid as unix::gid_t, pid: Some(unpcbid.unp_pid as unix::pid_t), }) } else { Err(io::Error::last_os_error()) } } } } #[cfg(any(target_os = "dragonfly", target_os = "freebsd"))] pub(crate) mod impl_bsd { use crate::net::unix::{self, UnixStream}; use libc::getpeereid; use std::io; use std::mem::MaybeUninit; use std::os::unix::io::AsRawFd; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result { unsafe { let raw_fd = sock.as_raw_fd(); let mut uid = MaybeUninit::uninit(); let mut gid = MaybeUninit::uninit(); let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr()); if ret == 0 { Ok(super::UCred { uid: uid.assume_init() as unix::uid_t, gid: gid.assume_init() as unix::gid_t, pid: None, }) } else { Err(io::Error::last_os_error()) } } } } #[cfg(any( target_os = "macos", target_os = "ios", target_os = "tvos", target_os = "watchos", target_os = "visionos" ))] pub(crate) mod impl_macos { use crate::net::unix::{self, UnixStream}; use libc::{c_void, getpeereid, getsockopt, pid_t, LOCAL_PEEREPID, SOL_LOCAL}; use std::io; use std::mem::size_of; use std::mem::MaybeUninit; use std::os::unix::io::AsRawFd; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result { unsafe { let raw_fd = sock.as_raw_fd(); let mut uid = MaybeUninit::uninit(); let mut gid = MaybeUninit::uninit(); let mut pid: MaybeUninit = MaybeUninit::uninit(); let mut pid_size: MaybeUninit = MaybeUninit::new(size_of::() as u32); if getsockopt( raw_fd, SOL_LOCAL, LOCAL_PEEREPID, pid.as_mut_ptr() as *mut c_void, pid_size.as_mut_ptr(), ) != 0 { return Err(io::Error::last_os_error()); } assert!(pid_size.assume_init() == (size_of::() as u32)); let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr()); if ret == 0 { Ok(super::UCred { uid: uid.assume_init() as unix::uid_t, gid: gid.assume_init() as unix::gid_t, pid: Some(pid.assume_init() as unix::pid_t), }) } else { Err(io::Error::last_os_error()) } } } } #[cfg(any(target_os = "solaris", target_os = "illumos"))] pub(crate) mod impl_solaris { use crate::net::unix::{self, UnixStream}; use std::io; use std::os::unix::io::AsRawFd; use std::ptr; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result { unsafe { let raw_fd = sock.as_raw_fd(); let mut cred = ptr::null_mut(); let ret = libc::getpeerucred(raw_fd, &mut cred); if ret == 0 { let uid = libc::ucred_geteuid(cred); let gid = libc::ucred_getegid(cred); let pid = libc::ucred_getpid(cred); libc::ucred_free(cred); Ok(super::UCred { uid: uid as unix::uid_t, gid: gid as unix::gid_t, pid: Some(pid as unix::pid_t), }) } else { Err(io::Error::last_os_error()) } } } } #[cfg(target_os = "aix")] pub(crate) mod impl_aix { use crate::net::unix::UnixStream; use std::io; use std::os::unix::io::AsRawFd; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result { unsafe { let raw_fd = sock.as_raw_fd(); let mut uid = std::mem::MaybeUninit::uninit(); let mut gid = std::mem::MaybeUninit::uninit(); let ret = libc::getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr()); if ret == 0 { Ok(super::UCred { uid: uid.assume_init(), gid: gid.assume_init(), pid: None, }) } else { Err(io::Error::last_os_error()) } } } } #[cfg(any(target_os = "espidf", target_os = "vita"))] pub(crate) mod impl_noproc { use crate::net::unix::UnixStream; use std::io; pub(crate) fn get_peer_cred(_sock: &UnixStream) -> io::Result { Ok(super::UCred { uid: 0, gid: 0, pid: None, }) } } tokio-1.43.1/src/net/windows/mod.rs000064400000000000000000000000711046102023000152310ustar 00000000000000//! Windows specific network types. pub mod named_pipe; tokio-1.43.1/src/net/windows/named_pipe.rs000064400000000000000000003015221046102023000165600ustar 00000000000000//! Tokio support for [Windows named pipes]. //! //! [Windows named pipes]: https://docs.microsoft.com/en-us/windows/win32/ipc/named-pipes use std::ffi::c_void; use std::ffi::OsStr; use std::io::{self, Read, Write}; use std::pin::Pin; use std::ptr; use std::task::{Context, Poll}; use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; use crate::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, RawHandle}; cfg_io_util! { use bytes::BufMut; } // Hide imports which are not used when generating documentation. #[cfg(windows)] mod doc { pub(super) use crate::os::windows::ffi::OsStrExt; pub(super) mod windows_sys { pub(crate) use windows_sys::{ Win32::Foundation::*, Win32::Storage::FileSystem::*, Win32::System::Pipes::*, Win32::System::SystemServices::*, }; } pub(super) use mio::windows as mio_windows; } // NB: none of these shows up in public API, so don't document them. #[cfg(not(windows))] mod doc { pub(super) mod mio_windows { pub type NamedPipe = crate::doc::NotDefinedHere; } } use self::doc::*; /// A [Windows named pipe] server. /// /// Accepting client connections involves creating a server with /// [`ServerOptions::create`] and waiting for clients to connect using /// [`NamedPipeServer::connect`]. /// /// To avoid having clients sporadically fail with /// [`std::io::ErrorKind::NotFound`] when they connect to a server, we must /// ensure that at least one server instance is available at all times. This /// means that the typical listen loop for a server is a bit involved, because /// we have to ensure that we never drop a server accidentally while a client /// might connect. /// /// So a correctly implemented server looks like this: /// /// ```no_run /// use std::io; /// use tokio::net::windows::named_pipe::ServerOptions; /// /// const PIPE_NAME: &str = r"\\.\pipe\named-pipe-idiomatic-server"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// // The first server needs to be constructed early so that clients can /// // be correctly connected. Otherwise calling .wait will cause the client to /// // error. /// // /// // Here we also make use of `first_pipe_instance`, which will ensure that /// // there are no other servers up and running already. /// let mut server = ServerOptions::new() /// .first_pipe_instance(true) /// .create(PIPE_NAME)?; /// /// // Spawn the server loop. /// let server = tokio::spawn(async move { /// loop { /// // Wait for a client to connect. /// server.connect().await?; /// let connected_client = server; /// /// // Construct the next server to be connected before sending the one /// // we already have of onto a task. This ensures that the server /// // isn't closed (after it's done in the task) before a new one is /// // available. Otherwise the client might error with /// // `io::ErrorKind::NotFound`. /// server = ServerOptions::new().create(PIPE_NAME)?; /// /// let client = tokio::spawn(async move { /// /* use the connected client */ /// # Ok::<_, std::io::Error>(()) /// }); /// # if true { break } // needed for type inference to work /// } /// /// Ok::<_, io::Error>(()) /// }); /// /// /* do something else not server related here */ /// # Ok(()) } /// ``` /// /// [Windows named pipe]: https://docs.microsoft.com/en-us/windows/win32/ipc/named-pipes #[derive(Debug)] pub struct NamedPipeServer { io: PollEvented, } impl NamedPipeServer { /// Constructs a new named pipe server from the specified raw handle. /// /// This function will consume ownership of the handle given, passing /// responsibility for closing the handle to the returned object. /// /// This function is also unsafe as the primitives currently returned have /// the contract that they are the sole owner of the file descriptor they /// are wrapping. Usage of this function could accidentally allow violating /// this contract which can cause memory unsafety in code that relies on it /// being true. /// /// # Errors /// /// This errors if called outside of a [Tokio Runtime], or in a runtime that /// has not [enabled I/O], or if any OS-specific I/O errors occur. /// /// [Tokio Runtime]: crate::runtime::Runtime /// [enabled I/O]: crate::runtime::Builder::enable_io pub unsafe fn from_raw_handle(handle: RawHandle) -> io::Result { let named_pipe = mio_windows::NamedPipe::from_raw_handle(handle); Ok(Self { io: PollEvented::new(named_pipe)?, }) } /// Retrieves information about the named pipe the server is associated /// with. /// /// ```no_run /// use tokio::net::windows::named_pipe::{PipeEnd, PipeMode, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-info"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let server = ServerOptions::new() /// .pipe_mode(PipeMode::Message) /// .max_instances(5) /// .create(PIPE_NAME)?; /// /// let server_info = server.info()?; /// /// assert_eq!(server_info.end, PipeEnd::Server); /// assert_eq!(server_info.mode, PipeMode::Message); /// assert_eq!(server_info.max_instances, 5); /// # Ok(()) } /// ``` pub fn info(&self) -> io::Result { // Safety: we're ensuring the lifetime of the named pipe. unsafe { named_pipe_info(self.io.as_raw_handle()) } } /// Enables a named pipe server process to wait for a client process to /// connect to an instance of a named pipe. A client process connects by /// creating a named pipe with the same name. /// /// This corresponds to the [`ConnectNamedPipe`] system call. /// /// # Cancel safety /// /// This method is cancellation safe in the sense that if it is used as the /// event in a [`select!`](crate::select) statement and some other branch /// completes first, then no connection events have been lost. /// /// [`ConnectNamedPipe`]: https://docs.microsoft.com/en-us/windows/win32/api/namedpipeapi/nf-namedpipeapi-connectnamedpipe /// /// # Example /// /// ```no_run /// use tokio::net::windows::named_pipe::ServerOptions; /// /// const PIPE_NAME: &str = r"\\.\pipe\mynamedpipe"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let pipe = ServerOptions::new().create(PIPE_NAME)?; /// /// // Wait for a client to connect. /// pipe.connect().await?; /// /// // Use the connected client... /// # Ok(()) } /// ``` pub async fn connect(&self) -> io::Result<()> { match self.io.connect() { Err(e) if e.kind() == io::ErrorKind::WouldBlock => { self.io .registration() .async_io(Interest::WRITABLE, || self.io.connect()) .await } x => x, } } /// Disconnects the server end of a named pipe instance from a client /// process. /// /// ``` /// use tokio::io::AsyncWriteExt; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// use windows_sys::Win32::Foundation::ERROR_PIPE_NOT_CONNECTED; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-disconnect"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let server = ServerOptions::new() /// .create(PIPE_NAME)?; /// /// let mut client = ClientOptions::new() /// .open(PIPE_NAME)?; /// /// // Wait for a client to become connected. /// server.connect().await?; /// /// // Forcibly disconnect the client. /// server.disconnect()?; /// /// // Write fails with an OS-specific error after client has been /// // disconnected. /// let e = client.write(b"ping").await.unwrap_err(); /// assert_eq!(e.raw_os_error(), Some(ERROR_PIPE_NOT_CONNECTED as i32)); /// # Ok(()) } /// ``` pub fn disconnect(&self) -> io::Result<()> { self.io.disconnect() } /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same pipe on a single /// task without splitting the pipe. /// /// The function may complete without the pipe being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// # Examples /// /// Concurrently read and write to the pipe on the same task without /// splitting. /// /// ```no_run /// use tokio::io::Interest; /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-ready"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// loop { /// let ready = server.ready(Interest::READABLE | Interest::WRITABLE).await?; /// /// if ready.is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// if ready.is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the pipe to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-readable"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// let mut msg = vec![0; 1024]; /// /// loop { /// // Wait for the pipe to be readable /// server.readable().await?; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_read(&mut msg) { /// Ok(n) => { /// msg.truncate(n); /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// println!("GOT = {:?}", msg); /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read readiness. /// /// If the pipe is not currently ready for reading, this method will /// store a clone of the `Waker` from the provided `Context`. When the pipe /// becomes ready for reading, `Waker::wake` will be called on the waker. /// /// Note that on multiple calls to `poll_read_ready` or `poll_read`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_write_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the pipe is not ready for reading. /// * `Poll::Ready(Ok(()))` if the pipe is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`readable`]: method@Self::readable pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Tries to read data from the pipe into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: NamedPipeServer::readable() /// [`ready()`]: NamedPipeServer::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The pipe's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the pipe is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-try-read"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be readable /// server.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf = [0; 4096]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_read(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read(buf)) } /// Tries to read data from the pipe into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: NamedPipeServer::try_read() /// [`readable()`]: NamedPipeServer::readable() /// [`ready()`]: NamedPipeServer::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the pipe's read half is closed /// and will no longer yield data. If the pipe is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io::{self, IoSliceMut}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-try-read-vectored"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be readable /// server.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf_a = [0; 512]; /// let mut buf_b = [0; 1024]; /// let mut bufs = [ /// IoSliceMut::new(&mut buf_a), /// IoSliceMut::new(&mut buf_b), /// ]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_read_vectored(&mut bufs) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: NamedPipeServer::readable() /// [`ready()`]: NamedPipeServer::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-readable"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new().create(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be readable /// server.readable().await?; /// /// let mut buf = Vec::with_capacity(4096); /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_read_buf(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { use std::io::Read; let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `NamedPipeServer::read` to have filled up `n` bytes in the // buffer. let n = (&*self.io).read(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }) } } /// Waits for the pipe to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-writable"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be writable /// server.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write readiness. /// /// If the pipe is not currently ready for writing, this method will /// store a clone of the `Waker` from the provided `Context`. When the pipe /// becomes ready for writing, `Waker::wake` will be called on the waker. /// /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_read_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the pipe is not ready for writing. /// * `Poll::Ready(Ok(()))` if the pipe is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`writable`]: method@Self::writable pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Tries to write a buffer to the pipe, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the pipe is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-try-write"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be writable /// server.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write(&self, buf: &[u8]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } /// Tries to write several buffers to the pipe, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: NamedPipeServer::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the pipe is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-server-try-write-vectored"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let server = named_pipe::ServerOptions::new() /// .create(PIPE_NAME)?; /// /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; /// /// loop { /// // Wait for the pipe to be writable /// server.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match server.try_write_vectored(&bufs) { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } /// Tries to read or write from the pipe using a user-provided IO operation. /// /// If the pipe is ready, the provided closure is called. The closure /// should attempt to perform IO operation from the pipe by manually /// calling the appropriate syscall. If the operation fails because the /// pipe is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the pipe is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the pipe that failed due to the pipe not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the pipe to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the /// methods defined on the Tokio `NamedPipeServer` type, as this will mess with /// the readiness flag and can cause the pipe to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: NamedPipeServer::readable() /// [`writable()`]: NamedPipeServer::writable() /// [`ready()`]: NamedPipeServer::ready() pub fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { self.io.registration().try_io(interest, f) } /// Reads or writes from the pipe using a user-provided IO operation. /// /// The readiness of the pipe is awaited and when the pipe is ready, /// the provided closure is called. The closure should attempt to perform /// IO operation on the pipe by manually calling the appropriate syscall. /// If the operation fails because the pipe is not actually ready, /// then the closure should return a `WouldBlock` error. In such case the /// readiness flag is cleared and the pipe readiness is awaited again. /// This loop is repeated until the closure returns an `Ok` or an error /// other than `WouldBlock`. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the pipe that failed due to the pipe not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the pipe to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `NamedPipeServer` type, as this will mess with the /// readiness flag and can cause the pipe to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub async fn async_io( &self, interest: Interest, f: impl FnMut() -> io::Result, ) -> io::Result { self.io.registration().async_io(interest, f).await } } impl AsyncRead for NamedPipeServer { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { unsafe { self.io.poll_read(cx, buf) } } } impl AsyncWrite for NamedPipeServer { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.io.poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.io.poll_write_vectored(cx, bufs) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } impl AsRawHandle for NamedPipeServer { fn as_raw_handle(&self) -> RawHandle { self.io.as_raw_handle() } } impl AsHandle for NamedPipeServer { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } } } /// A [Windows named pipe] client. /// /// Constructed using [`ClientOptions::open`]. /// /// Connecting a client correctly involves a few steps. When connecting through /// [`ClientOptions::open`], it might error indicating one of two things: /// /// * [`std::io::ErrorKind::NotFound`] - There is no server available. /// * [`ERROR_PIPE_BUSY`] - There is a server available, but it is busy. Sleep /// for a while and try again. /// /// So a correctly implemented client looks like this: /// /// ```no_run /// use std::time::Duration; /// use tokio::net::windows::named_pipe::ClientOptions; /// use tokio::time; /// use windows_sys::Win32::Foundation::ERROR_PIPE_BUSY; /// /// const PIPE_NAME: &str = r"\\.\pipe\named-pipe-idiomatic-client"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let client = loop { /// match ClientOptions::new().open(PIPE_NAME) { /// Ok(client) => break client, /// Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY as i32) => (), /// Err(e) => return Err(e), /// } /// /// time::sleep(Duration::from_millis(50)).await; /// }; /// /// /* use the connected client */ /// # Ok(()) } /// ``` /// /// [`ERROR_PIPE_BUSY`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_PIPE_BUSY.html /// [Windows named pipe]: https://docs.microsoft.com/en-us/windows/win32/ipc/named-pipes #[derive(Debug)] pub struct NamedPipeClient { io: PollEvented, } impl NamedPipeClient { /// Constructs a new named pipe client from the specified raw handle. /// /// This function will consume ownership of the handle given, passing /// responsibility for closing the handle to the returned object. /// /// This function is also unsafe as the primitives currently returned have /// the contract that they are the sole owner of the file descriptor they /// are wrapping. Usage of this function could accidentally allow violating /// this contract which can cause memory unsafety in code that relies on it /// being true. /// /// # Errors /// /// This errors if called outside of a [Tokio Runtime], or in a runtime that /// has not [enabled I/O], or if any OS-specific I/O errors occur. /// /// [Tokio Runtime]: crate::runtime::Runtime /// [enabled I/O]: crate::runtime::Builder::enable_io pub unsafe fn from_raw_handle(handle: RawHandle) -> io::Result { let named_pipe = mio_windows::NamedPipe::from_raw_handle(handle); Ok(Self { io: PollEvented::new(named_pipe)?, }) } /// Retrieves information about the named pipe the client is associated /// with. /// /// ```no_run /// use tokio::net::windows::named_pipe::{ClientOptions, PipeEnd, PipeMode}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-info"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let client = ClientOptions::new() /// .open(PIPE_NAME)?; /// /// let client_info = client.info()?; /// /// assert_eq!(client_info.end, PipeEnd::Client); /// assert_eq!(client_info.mode, PipeMode::Message); /// assert_eq!(client_info.max_instances, 5); /// # Ok(()) } /// ``` pub fn info(&self) -> io::Result { // Safety: we're ensuring the lifetime of the named pipe. unsafe { named_pipe_info(self.io.as_raw_handle()) } } /// Waits for any of the requested ready states. /// /// This function is usually paired with `try_read()` or `try_write()`. It /// can be used to concurrently read / write to the same pipe on a single /// task without splitting the pipe. /// /// The function may complete without the pipe being ready. This is a /// false-positive and attempting an operation will return with /// `io::ErrorKind::WouldBlock`. The function can also return with an empty /// [`Ready`] set, so you should always check the returned value and possibly /// wait again if the requested states are not set. /// /// # Examples /// /// Concurrently read and write to the pipe on the same task without /// splitting. /// /// ```no_run /// use tokio::io::Interest; /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-ready"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// loop { /// let ready = client.ready(Interest::READABLE | Interest::WRITABLE).await?; /// /// if ready.is_readable() { /// let mut data = vec![0; 1024]; /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_read(&mut data) { /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// if ready.is_writable() { /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_write(b"hello world") { /// Ok(n) => { /// println!("write {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// } /// } /// ``` pub async fn ready(&self, interest: Interest) -> io::Result { let event = self.io.registration().readiness(interest).await?; Ok(event.ready) } /// Waits for the pipe to become readable. /// /// This function is equivalent to `ready(Interest::READABLE)` and is usually /// paired with `try_read()`. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-readable"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// let mut msg = vec![0; 1024]; /// /// loop { /// // Wait for the pipe to be readable /// client.readable().await?; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_read(&mut msg) { /// Ok(n) => { /// msg.truncate(n); /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// println!("GOT = {:?}", msg); /// Ok(()) /// } /// ``` pub async fn readable(&self) -> io::Result<()> { self.ready(Interest::READABLE).await?; Ok(()) } /// Polls for read readiness. /// /// If the pipe is not currently ready for reading, this method will /// store a clone of the `Waker` from the provided `Context`. When the pipe /// becomes ready for reading, `Waker::wake` will be called on the waker. /// /// Note that on multiple calls to `poll_read_ready` or `poll_read`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_write_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`readable`] is not feasible. Where possible, using [`readable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the pipe is not ready for reading. /// * `Poll::Ready(Ok(()))` if the pipe is ready for reading. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`readable`]: method@Self::readable pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_read_ready(cx).map_ok(|_| ()) } /// Tries to read data from the pipe into the provided buffer, returning how /// many bytes were read. /// /// Receives any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: NamedPipeClient::readable() /// [`ready()`]: NamedPipeClient::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. If `n` is `0`, then it can indicate one of two scenarios: /// /// 1. The pipe's read half is closed and will no longer yield data. /// 2. The specified buffer was 0 bytes in length. /// /// If the pipe is not ready to read data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-try-read"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be readable /// client.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf = [0; 4096]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_read(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read(&self, buf: &mut [u8]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read(buf)) } /// Tries to read data from the pipe into the provided buffers, returning /// how many bytes were read. /// /// Data is copied to fill each buffer in order, with the final buffer /// written to possibly being only partially filled. This method behaves /// equivalently to a single call to [`try_read()`] with concatenated /// buffers. /// /// Receives any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_vectored()` is non-blocking, the buffer does not have to be /// stored by the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`try_read()`]: NamedPipeClient::try_read() /// [`readable()`]: NamedPipeClient::readable() /// [`ready()`]: NamedPipeClient::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the pipe's read half is closed /// and will no longer yield data. If the pipe is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io::{self, IoSliceMut}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-try-read-vectored"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be readable /// client.readable().await?; /// /// // Creating the buffer **after** the `await` prevents it from /// // being stored in the async task. /// let mut buf_a = [0; 512]; /// let mut buf_b = [0; 1024]; /// let mut bufs = [ /// IoSliceMut::new(&mut buf_a), /// IoSliceMut::new(&mut buf_b), /// ]; /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_read_vectored(&mut bufs) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result { self.io .registration() .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } cfg_io_util! { /// Tries to read data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// /// Receives any pending data from the pipe but does not wait for new data /// to arrive. On success, returns the number of bytes read. Because /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by /// the async task and can exist entirely on the stack. /// /// Usually, [`readable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: NamedPipeClient::readable() /// [`ready()`]: NamedPipeClient::ready() /// /// # Return /// /// If data is successfully read, `Ok(n)` is returned, where `n` is the /// number of bytes read. `Ok(0)` indicates the stream's read half is closed /// and will no longer yield data. If the stream is not ready to read data /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-readable"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be readable /// client.readable().await?; /// /// let mut buf = Vec::with_capacity(4096); /// /// // Try to read data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_read_buf(&mut buf) { /// Ok(0) => break, /// Ok(n) => { /// println!("read {} bytes", n); /// } /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_read_buf(&self, buf: &mut B) -> io::Result { self.io.registration().try_io(Interest::READABLE, || { use std::io::Read; let dst = buf.chunk_mut(); let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit] as *mut [u8]) }; // Safety: We trust `NamedPipeClient::read` to have filled up `n` bytes in the // buffer. let n = (&*self.io).read(dst)?; unsafe { buf.advance_mut(n); } Ok(n) }) } } /// Waits for the pipe to become writable. /// /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually /// paired with `try_write()`. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-writable"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be writable /// client.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub async fn writable(&self) -> io::Result<()> { self.ready(Interest::WRITABLE).await?; Ok(()) } /// Polls for write readiness. /// /// If the pipe is not currently ready for writing, this method will /// store a clone of the `Waker` from the provided `Context`. When the pipe /// becomes ready for writing, `Waker::wake` will be called on the waker. /// /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only /// the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. (However, `poll_read_ready` retains a /// second, independent waker.) /// /// This function is intended for cases where creating and pinning a future /// via [`writable`] is not feasible. Where possible, using [`writable`] is /// preferred, as this supports polling from multiple tasks at once. /// /// # Return value /// /// The function returns: /// /// * `Poll::Pending` if the pipe is not ready for writing. /// * `Poll::Ready(Ok(()))` if the pipe is ready for writing. /// * `Poll::Ready(Err(e))` if an error is encountered. /// /// # Errors /// /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`writable`]: method@Self::writable pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { self.io.registration().poll_write_ready(cx).map_ok(|_| ()) } /// Tries to write a buffer to the pipe, returning how many bytes were /// written. /// /// The function will attempt to write the entire contents of `buf`, but /// only part of the buffer may be written. /// /// This function is usually paired with `writable()`. /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the pipe is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-try-write"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// loop { /// // Wait for the pipe to be writable /// client.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_write(b"hello world") { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write(&self, buf: &[u8]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) } /// Tries to write several buffers to the pipe, returning how many bytes /// were written. /// /// Data is written from each buffer in order, with the final buffer read /// from possible being only partially consumed. This method behaves /// equivalently to a single call to [`try_write()`] with concatenated /// buffers. /// /// This function is usually paired with `writable()`. /// /// [`try_write()`]: NamedPipeClient::try_write() /// /// # Return /// /// If data is successfully written, `Ok(n)` is returned, where `n` is the /// number of bytes written. If the pipe is not ready to write data, /// `Err(io::ErrorKind::WouldBlock)` is returned. /// /// # Examples /// /// ```no_run /// use tokio::net::windows::named_pipe; /// use std::error::Error; /// use std::io; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-try-write-vectored"; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let client = named_pipe::ClientOptions::new().open(PIPE_NAME)?; /// /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; /// /// loop { /// // Wait for the pipe to be writable /// client.writable().await?; /// /// // Try to write data, this may still fail with `WouldBlock` /// // if the readiness event is a false positive. /// match client.try_write_vectored(&bufs) { /// Ok(n) => { /// break; /// } /// Err(e) if e.kind() == io::ErrorKind::WouldBlock => { /// continue; /// } /// Err(e) => { /// return Err(e.into()); /// } /// } /// } /// /// Ok(()) /// } /// ``` pub fn try_write_vectored(&self, buf: &[io::IoSlice<'_>]) -> io::Result { self.io .registration() .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } /// Tries to read or write from the pipe using a user-provided IO operation. /// /// If the pipe is ready, the provided closure is called. The closure /// should attempt to perform IO operation from the pipe by manually /// calling the appropriate syscall. If the operation fails because the /// pipe is not actually ready, then the closure should return a /// `WouldBlock` error and the readiness flag is cleared. The return value /// of the closure is then returned by `try_io`. /// /// If the pipe is not ready, then the closure is not called /// and a `WouldBlock` error is returned. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the pipe that failed due to the pipe not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the pipe to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `NamedPipeClient` type, as this will mess with the /// readiness flag and can cause the pipe to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. /// /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. /// /// [`readable()`]: NamedPipeClient::readable() /// [`writable()`]: NamedPipeClient::writable() /// [`ready()`]: NamedPipeClient::ready() pub fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { self.io.registration().try_io(interest, f) } /// Reads or writes from the pipe using a user-provided IO operation. /// /// The readiness of the pipe is awaited and when the pipe is ready, /// the provided closure is called. The closure should attempt to perform /// IO operation on the pipe by manually calling the appropriate syscall. /// If the operation fails because the pipe is not actually ready, /// then the closure should return a `WouldBlock` error. In such case the /// readiness flag is cleared and the pipe readiness is awaited again. /// This loop is repeated until the closure returns an `Ok` or an error /// other than `WouldBlock`. /// /// The closure should only return a `WouldBlock` error if it has performed /// an IO operation on the pipe that failed due to the pipe not being /// ready. Returning a `WouldBlock` error in any other situation will /// incorrectly clear the readiness flag, which can cause the pipe to /// behave incorrectly. /// /// The closure should not perform the IO operation using any of the methods /// defined on the Tokio `NamedPipeClient` type, as this will mess with the /// readiness flag and can cause the pipe to behave incorrectly. /// /// This method is not intended to be used with combined interests. /// The closure should perform only one type of IO operation, so it should not /// require more than one ready state. This method may panic or sleep forever /// if it is called with a combined interest. pub async fn async_io( &self, interest: Interest, f: impl FnMut() -> io::Result, ) -> io::Result { self.io.registration().async_io(interest, f).await } } impl AsyncRead for NamedPipeClient { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { unsafe { self.io.poll_read(cx, buf) } } } impl AsyncWrite for NamedPipeClient { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.io.poll_write(cx, buf) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.io.poll_write_vectored(cx, bufs) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } impl AsRawHandle for NamedPipeClient { fn as_raw_handle(&self) -> RawHandle { self.io.as_raw_handle() } } impl AsHandle for NamedPipeClient { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } } } /// A builder structure for construct a named pipe with named pipe-specific /// options. This is required to use for named pipe servers who wants to modify /// pipe-related options. /// /// See [`ServerOptions::create`]. #[derive(Debug, Clone)] pub struct ServerOptions { // dwOpenMode access_inbound: bool, access_outbound: bool, first_pipe_instance: bool, write_dac: bool, write_owner: bool, access_system_security: bool, // dwPipeMode pipe_mode: PipeMode, reject_remote_clients: bool, // other options max_instances: u32, out_buffer_size: u32, in_buffer_size: u32, default_timeout: u32, } impl ServerOptions { /// Creates a new named pipe builder with the default settings. /// /// ``` /// use tokio::net::windows::named_pipe::ServerOptions; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-new"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let server = ServerOptions::new().create(PIPE_NAME)?; /// # Ok(()) } /// ``` pub fn new() -> ServerOptions { ServerOptions { access_inbound: true, access_outbound: true, first_pipe_instance: false, write_dac: false, write_owner: false, access_system_security: false, pipe_mode: PipeMode::Byte, reject_remote_clients: true, max_instances: windows_sys::PIPE_UNLIMITED_INSTANCES, out_buffer_size: 65536, in_buffer_size: 65536, default_timeout: 0, } } /// The pipe mode. /// /// The default pipe mode is [`PipeMode::Byte`]. See [`PipeMode`] for /// documentation of what each mode means. /// /// This corresponds to specifying `PIPE_TYPE_` and `PIPE_READMODE_` in [`dwPipeMode`]. /// /// [`dwPipeMode`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn pipe_mode(&mut self, pipe_mode: PipeMode) -> &mut Self { self.pipe_mode = pipe_mode; self } /// The flow of data in the pipe goes from client to server only. /// /// This corresponds to setting [`PIPE_ACCESS_INBOUND`]. /// /// [`PIPE_ACCESS_INBOUND`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea#pipe_access_inbound /// /// # Errors /// /// Server side prevents connecting by denying inbound access, client errors /// with [`std::io::ErrorKind::PermissionDenied`] when attempting to create /// the connection. /// /// ``` /// use std::io; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-access-inbound-err1"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let _server = ServerOptions::new() /// .access_inbound(false) /// .create(PIPE_NAME)?; /// /// let e = ClientOptions::new() /// .open(PIPE_NAME) /// .unwrap_err(); /// /// assert_eq!(e.kind(), io::ErrorKind::PermissionDenied); /// # Ok(()) } /// ``` /// /// Disabling writing allows a client to connect, but errors with /// [`std::io::ErrorKind::PermissionDenied`] if a write is attempted. /// /// ``` /// use std::io; /// use tokio::io::AsyncWriteExt; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-access-inbound-err2"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let server = ServerOptions::new() /// .access_inbound(false) /// .create(PIPE_NAME)?; /// /// let mut client = ClientOptions::new() /// .write(false) /// .open(PIPE_NAME)?; /// /// server.connect().await?; /// /// let e = client.write(b"ping").await.unwrap_err(); /// assert_eq!(e.kind(), io::ErrorKind::PermissionDenied); /// # Ok(()) } /// ``` /// /// # Examples /// /// A unidirectional named pipe that only supports server-to-client /// communication. /// /// ``` /// use std::io; /// use tokio::io::{AsyncReadExt, AsyncWriteExt}; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-access-inbound"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let mut server = ServerOptions::new() /// .access_inbound(false) /// .create(PIPE_NAME)?; /// /// let mut client = ClientOptions::new() /// .write(false) /// .open(PIPE_NAME)?; /// /// server.connect().await?; /// /// let write = server.write_all(b"ping"); /// /// let mut buf = [0u8; 4]; /// let read = client.read_exact(&mut buf); /// /// let ((), read) = tokio::try_join!(write, read)?; /// /// assert_eq!(read, 4); /// assert_eq!(&buf[..], b"ping"); /// # Ok(()) } /// ``` pub fn access_inbound(&mut self, allowed: bool) -> &mut Self { self.access_inbound = allowed; self } /// The flow of data in the pipe goes from server to client only. /// /// This corresponds to setting [`PIPE_ACCESS_OUTBOUND`]. /// /// [`PIPE_ACCESS_OUTBOUND`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea#pipe_access_outbound /// /// # Errors /// /// Server side prevents connecting by denying outbound access, client /// errors with [`std::io::ErrorKind::PermissionDenied`] when attempting to /// create the connection. /// /// ``` /// use std::io; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-access-outbound-err1"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let server = ServerOptions::new() /// .access_outbound(false) /// .create(PIPE_NAME)?; /// /// let e = ClientOptions::new() /// .open(PIPE_NAME) /// .unwrap_err(); /// /// assert_eq!(e.kind(), io::ErrorKind::PermissionDenied); /// # Ok(()) } /// ``` /// /// Disabling reading allows a client to connect, but attempting to read /// will error with [`std::io::ErrorKind::PermissionDenied`]. /// /// ``` /// use std::io; /// use tokio::io::AsyncReadExt; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-access-outbound-err2"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let server = ServerOptions::new() /// .access_outbound(false) /// .create(PIPE_NAME)?; /// /// let mut client = ClientOptions::new() /// .read(false) /// .open(PIPE_NAME)?; /// /// server.connect().await?; /// /// let mut buf = [0u8; 4]; /// let e = client.read(&mut buf).await.unwrap_err(); /// assert_eq!(e.kind(), io::ErrorKind::PermissionDenied); /// # Ok(()) } /// ``` /// /// # Examples /// /// A unidirectional named pipe that only supports client-to-server /// communication. /// /// ``` /// use tokio::io::{AsyncReadExt, AsyncWriteExt}; /// use tokio::net::windows::named_pipe::{ClientOptions, ServerOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-access-outbound"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let mut server = ServerOptions::new() /// .access_outbound(false) /// .create(PIPE_NAME)?; /// /// let mut client = ClientOptions::new() /// .read(false) /// .open(PIPE_NAME)?; /// /// server.connect().await?; /// /// let write = client.write_all(b"ping"); /// /// let mut buf = [0u8; 4]; /// let read = server.read_exact(&mut buf); /// /// let ((), read) = tokio::try_join!(write, read)?; /// /// println!("done reading and writing"); /// /// assert_eq!(read, 4); /// assert_eq!(&buf[..], b"ping"); /// # Ok(()) } /// ``` pub fn access_outbound(&mut self, allowed: bool) -> &mut Self { self.access_outbound = allowed; self } /// If you attempt to create multiple instances of a pipe with this flag /// set, creation of the first server instance succeeds, but creation of any /// subsequent instances will fail with /// [`std::io::ErrorKind::PermissionDenied`]. /// /// This option is intended to be used with servers that want to ensure that /// they are the only process listening for clients on a given named pipe. /// This is accomplished by enabling it for the first server instance /// created in a process. /// /// This corresponds to setting [`FILE_FLAG_FIRST_PIPE_INSTANCE`]. /// /// # Errors /// /// If this option is set and more than one instance of the server for a /// given named pipe exists, calling [`create`] will fail with /// [`std::io::ErrorKind::PermissionDenied`]. /// /// ``` /// use std::io; /// use tokio::net::windows::named_pipe::ServerOptions; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-first-instance-error"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let server1 = ServerOptions::new() /// .first_pipe_instance(true) /// .create(PIPE_NAME)?; /// /// // Second server errs, since it's not the first instance. /// let e = ServerOptions::new() /// .first_pipe_instance(true) /// .create(PIPE_NAME) /// .unwrap_err(); /// /// assert_eq!(e.kind(), io::ErrorKind::PermissionDenied); /// # Ok(()) } /// ``` /// /// # Examples /// /// ``` /// use std::io; /// use tokio::net::windows::named_pipe::ServerOptions; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-first-instance"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let mut builder = ServerOptions::new(); /// builder.first_pipe_instance(true); /// /// let server = builder.create(PIPE_NAME)?; /// let e = builder.create(PIPE_NAME).unwrap_err(); /// assert_eq!(e.kind(), io::ErrorKind::PermissionDenied); /// drop(server); /// /// // OK: since, we've closed the other instance. /// let _server2 = builder.create(PIPE_NAME)?; /// # Ok(()) } /// ``` /// /// [`create`]: ServerOptions::create /// [`FILE_FLAG_FIRST_PIPE_INSTANCE`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea#pipe_first_pipe_instance pub fn first_pipe_instance(&mut self, first: bool) -> &mut Self { self.first_pipe_instance = first; self } /// Requests permission to modify the pipe's discretionary access control list. /// /// This corresponds to setting [`WRITE_DAC`] in dwOpenMode. /// /// # Examples /// /// ``` /// use std::{io, os::windows::prelude::AsRawHandle, ptr}; /// /// use tokio::net::windows::named_pipe::ServerOptions; /// use windows_sys::{ /// Win32::Foundation::ERROR_SUCCESS, /// Win32::Security::DACL_SECURITY_INFORMATION, /// Win32::Security::Authorization::{SetSecurityInfo, SE_KERNEL_OBJECT}, /// }; /// /// const PIPE_NAME: &str = r"\\.\pipe\write_dac_pipe"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let mut pipe_template = ServerOptions::new(); /// pipe_template.write_dac(true); /// let pipe = pipe_template.create(PIPE_NAME)?; /// /// unsafe { /// assert_eq!( /// ERROR_SUCCESS, /// SetSecurityInfo( /// pipe.as_raw_handle() as _, /// SE_KERNEL_OBJECT, /// DACL_SECURITY_INFORMATION, /// ptr::null_mut(), /// ptr::null_mut(), /// ptr::null_mut(), /// ptr::null_mut(), /// ) /// ); /// } /// /// # Ok(()) } /// ``` /// /// ``` /// use std::{io, os::windows::prelude::AsRawHandle, ptr}; /// /// use tokio::net::windows::named_pipe::ServerOptions; /// use windows_sys::{ /// Win32::Foundation::ERROR_ACCESS_DENIED, /// Win32::Security::DACL_SECURITY_INFORMATION, /// Win32::Security::Authorization::{SetSecurityInfo, SE_KERNEL_OBJECT}, /// }; /// /// const PIPE_NAME: &str = r"\\.\pipe\write_dac_pipe_fail"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let mut pipe_template = ServerOptions::new(); /// pipe_template.write_dac(false); /// let pipe = pipe_template.create(PIPE_NAME)?; /// /// unsafe { /// assert_eq!( /// ERROR_ACCESS_DENIED, /// SetSecurityInfo( /// pipe.as_raw_handle() as _, /// SE_KERNEL_OBJECT, /// DACL_SECURITY_INFORMATION, /// ptr::null_mut(), /// ptr::null_mut(), /// ptr::null_mut(), /// ptr::null_mut(), /// ) /// ); /// } /// /// # Ok(()) } /// ``` /// /// [`WRITE_DAC`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn write_dac(&mut self, requested: bool) -> &mut Self { self.write_dac = requested; self } /// Requests permission to modify the pipe's owner. /// /// This corresponds to setting [`WRITE_OWNER`] in dwOpenMode. /// /// [`WRITE_OWNER`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn write_owner(&mut self, requested: bool) -> &mut Self { self.write_owner = requested; self } /// Requests permission to modify the pipe's system access control list. /// /// This corresponds to setting [`ACCESS_SYSTEM_SECURITY`] in dwOpenMode. /// /// [`ACCESS_SYSTEM_SECURITY`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn access_system_security(&mut self, requested: bool) -> &mut Self { self.access_system_security = requested; self } /// Indicates whether this server can accept remote clients or not. Remote /// clients are disabled by default. /// /// This corresponds to setting [`PIPE_REJECT_REMOTE_CLIENTS`]. /// /// [`PIPE_REJECT_REMOTE_CLIENTS`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea#pipe_reject_remote_clients pub fn reject_remote_clients(&mut self, reject: bool) -> &mut Self { self.reject_remote_clients = reject; self } /// The maximum number of instances that can be created for this pipe. The /// first instance of the pipe can specify this value; the same number must /// be specified for other instances of the pipe. Acceptable values are in /// the range 1 through 254. The default value is unlimited. /// /// This corresponds to specifying [`nMaxInstances`]. /// /// [`nMaxInstances`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea /// /// # Errors /// /// The same numbers of `max_instances` have to be used by all servers. Any /// additional servers trying to be built which uses a mismatching value /// might error. /// /// ``` /// use std::io; /// use tokio::net::windows::named_pipe::{ServerOptions, ClientOptions}; /// use windows_sys::Win32::Foundation::ERROR_PIPE_BUSY; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-max-instances"; /// /// # #[tokio::main] async fn main() -> io::Result<()> { /// let mut server = ServerOptions::new(); /// server.max_instances(2); /// /// let s1 = server.create(PIPE_NAME)?; /// let c1 = ClientOptions::new().open(PIPE_NAME); /// /// let s2 = server.create(PIPE_NAME)?; /// let c2 = ClientOptions::new().open(PIPE_NAME); /// /// // Too many servers! /// let e = server.create(PIPE_NAME).unwrap_err(); /// assert_eq!(e.raw_os_error(), Some(ERROR_PIPE_BUSY as i32)); /// /// // Still too many servers even if we specify a higher value! /// let e = server.max_instances(100).create(PIPE_NAME).unwrap_err(); /// assert_eq!(e.raw_os_error(), Some(ERROR_PIPE_BUSY as i32)); /// # Ok(()) } /// ``` /// /// # Panics /// /// This function will panic if more than 254 instances are specified. If /// you do not wish to set an instance limit, leave it unspecified. /// /// ```should_panic /// use tokio::net::windows::named_pipe::ServerOptions; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let builder = ServerOptions::new().max_instances(255); /// # Ok(()) } /// ``` #[track_caller] pub fn max_instances(&mut self, instances: usize) -> &mut Self { assert!(instances < 255, "cannot specify more than 254 instances"); self.max_instances = instances as u32; self } /// The number of bytes to reserve for the output buffer. /// /// This corresponds to specifying [`nOutBufferSize`]. /// /// [`nOutBufferSize`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn out_buffer_size(&mut self, buffer: u32) -> &mut Self { self.out_buffer_size = buffer; self } /// The number of bytes to reserve for the input buffer. /// /// This corresponds to specifying [`nInBufferSize`]. /// /// [`nInBufferSize`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea pub fn in_buffer_size(&mut self, buffer: u32) -> &mut Self { self.in_buffer_size = buffer; self } /// Creates the named pipe identified by `addr` for use as a server. /// /// This uses the [`CreateNamedPipe`] function. /// /// [`CreateNamedPipe`]: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-createnamedpipea /// /// # Errors /// /// This errors if called outside of a [Tokio Runtime], or in a runtime that /// has not [enabled I/O], or if any OS-specific I/O errors occur. /// /// [Tokio Runtime]: crate::runtime::Runtime /// [enabled I/O]: crate::runtime::Builder::enable_io /// /// # Examples /// /// ``` /// use tokio::net::windows::named_pipe::ServerOptions; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-create"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let server = ServerOptions::new().create(PIPE_NAME)?; /// # Ok(()) } /// ``` pub fn create(&self, addr: impl AsRef) -> io::Result { // Safety: We're calling create_with_security_attributes_raw w/ a null // pointer which disables it. unsafe { self.create_with_security_attributes_raw(addr, ptr::null_mut()) } } /// Creates the named pipe identified by `addr` for use as a server. /// /// This is the same as [`create`] except that it supports providing the raw /// pointer to a structure of [`SECURITY_ATTRIBUTES`] which will be passed /// as the `lpSecurityAttributes` argument to [`CreateFile`]. /// /// # Errors /// /// This errors if called outside of a [Tokio Runtime], or in a runtime that /// has not [enabled I/O], or if any OS-specific I/O errors occur. /// /// [Tokio Runtime]: crate::runtime::Runtime /// [enabled I/O]: crate::runtime::Builder::enable_io /// /// # Safety /// /// The `attrs` argument must either be null or point at a valid instance of /// the [`SECURITY_ATTRIBUTES`] structure. If the argument is null, the /// behavior is identical to calling the [`create`] method. /// /// [`create`]: ServerOptions::create /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew /// [`SECURITY_ATTRIBUTES`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Security/struct.SECURITY_ATTRIBUTES.html pub unsafe fn create_with_security_attributes_raw( &self, addr: impl AsRef, attrs: *mut c_void, ) -> io::Result { let addr = encode_addr(addr); let pipe_mode = { let mut mode = if matches!(self.pipe_mode, PipeMode::Message) { windows_sys::PIPE_TYPE_MESSAGE | windows_sys::PIPE_READMODE_MESSAGE } else { windows_sys::PIPE_TYPE_BYTE | windows_sys::PIPE_READMODE_BYTE }; if self.reject_remote_clients { mode |= windows_sys::PIPE_REJECT_REMOTE_CLIENTS; } else { mode |= windows_sys::PIPE_ACCEPT_REMOTE_CLIENTS; } mode }; let open_mode = { let mut mode = windows_sys::FILE_FLAG_OVERLAPPED; if self.access_inbound { mode |= windows_sys::PIPE_ACCESS_INBOUND; } if self.access_outbound { mode |= windows_sys::PIPE_ACCESS_OUTBOUND; } if self.first_pipe_instance { mode |= windows_sys::FILE_FLAG_FIRST_PIPE_INSTANCE; } if self.write_dac { mode |= windows_sys::WRITE_DAC; } if self.write_owner { mode |= windows_sys::WRITE_OWNER; } if self.access_system_security { mode |= windows_sys::ACCESS_SYSTEM_SECURITY; } mode }; let h = windows_sys::CreateNamedPipeW( addr.as_ptr(), open_mode, pipe_mode, self.max_instances, self.out_buffer_size, self.in_buffer_size, self.default_timeout, attrs as *mut _, ); if h == windows_sys::INVALID_HANDLE_VALUE { return Err(io::Error::last_os_error()); } NamedPipeServer::from_raw_handle(h as _) } } /// A builder suitable for building and interacting with named pipes from the /// client side. /// /// See [`ClientOptions::open`]. #[derive(Debug, Clone)] pub struct ClientOptions { generic_read: bool, generic_write: bool, security_qos_flags: u32, pipe_mode: PipeMode, } impl ClientOptions { /// Creates a new named pipe builder with the default settings. /// /// ``` /// use tokio::net::windows::named_pipe::{ServerOptions, ClientOptions}; /// /// const PIPE_NAME: &str = r"\\.\pipe\tokio-named-pipe-client-new"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// // Server must be created in order for the client creation to succeed. /// let server = ServerOptions::new().create(PIPE_NAME)?; /// let client = ClientOptions::new().open(PIPE_NAME)?; /// # Ok(()) } /// ``` pub fn new() -> Self { Self { generic_read: true, generic_write: true, security_qos_flags: windows_sys::SECURITY_IDENTIFICATION | windows_sys::SECURITY_SQOS_PRESENT, pipe_mode: PipeMode::Byte, } } /// If the client supports reading data. This is enabled by default. /// /// This corresponds to setting [`GENERIC_READ`] in the call to [`CreateFile`]. /// /// [`GENERIC_READ`]: https://docs.microsoft.com/en-us/windows/win32/secauthz/generic-access-rights /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew pub fn read(&mut self, allowed: bool) -> &mut Self { self.generic_read = allowed; self } /// If the created pipe supports writing data. This is enabled by default. /// /// This corresponds to setting [`GENERIC_WRITE`] in the call to [`CreateFile`]. /// /// [`GENERIC_WRITE`]: https://docs.microsoft.com/en-us/windows/win32/secauthz/generic-access-rights /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew pub fn write(&mut self, allowed: bool) -> &mut Self { self.generic_write = allowed; self } /// Sets qos flags which are combined with other flags and attributes in the /// call to [`CreateFile`]. /// /// By default `security_qos_flags` is set to [`SECURITY_IDENTIFICATION`], /// calling this function would override that value completely with the /// argument specified. /// /// When `security_qos_flags` is not set, a malicious program can gain the /// elevated privileges of a privileged Rust process when it allows opening /// user-specified paths, by tricking it into opening a named pipe. So /// arguably `security_qos_flags` should also be set when opening arbitrary /// paths. However the bits can then conflict with other flags, specifically /// `FILE_FLAG_OPEN_NO_RECALL`. /// /// For information about possible values, see [Impersonation Levels] on the /// Windows Dev Center site. The `SECURITY_SQOS_PRESENT` flag is set /// automatically when using this method. /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea /// [`SECURITY_IDENTIFICATION`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Storage/FileSystem/constant.SECURITY_IDENTIFICATION.html /// [Impersonation Levels]: https://docs.microsoft.com/en-us/windows/win32/api/winnt/ne-winnt-security_impersonation_level pub fn security_qos_flags(&mut self, flags: u32) -> &mut Self { // See: https://github.com/rust-lang/rust/pull/58216 self.security_qos_flags = flags | windows_sys::SECURITY_SQOS_PRESENT; self } /// The pipe mode. /// /// The default pipe mode is [`PipeMode::Byte`]. See [`PipeMode`] for /// documentation of what each mode means. pub fn pipe_mode(&mut self, pipe_mode: PipeMode) -> &mut Self { self.pipe_mode = pipe_mode; self } /// Opens the named pipe identified by `addr`. /// /// This opens the client using [`CreateFile`] with the /// `dwCreationDisposition` option set to `OPEN_EXISTING`. /// /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilea /// /// # Errors /// /// This errors if called outside of a [Tokio Runtime], or in a runtime that /// has not [enabled I/O], or if any OS-specific I/O errors occur. /// /// There are a few errors you need to take into account when creating a /// named pipe on the client side: /// /// * [`std::io::ErrorKind::NotFound`] - This indicates that the named pipe /// does not exist. Presumably the server is not up. /// * [`ERROR_PIPE_BUSY`] - This error is raised when the named pipe exists, /// but the server is not currently waiting for a connection. Please see the /// examples for how to check for this error. /// /// [`ERROR_PIPE_BUSY`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_PIPE_BUSY.html /// [enabled I/O]: crate::runtime::Builder::enable_io /// [Tokio Runtime]: crate::runtime::Runtime /// /// A connect loop that waits until a pipe becomes available looks like /// this: /// /// ```no_run /// use std::time::Duration; /// use tokio::net::windows::named_pipe::ClientOptions; /// use tokio::time; /// use windows_sys::Win32::Foundation::ERROR_PIPE_BUSY; /// /// const PIPE_NAME: &str = r"\\.\pipe\mynamedpipe"; /// /// # #[tokio::main] async fn main() -> std::io::Result<()> { /// let client = loop { /// match ClientOptions::new().open(PIPE_NAME) { /// Ok(client) => break client, /// Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY as i32) => (), /// Err(e) => return Err(e), /// } /// /// time::sleep(Duration::from_millis(50)).await; /// }; /// /// // use the connected client. /// # Ok(()) } /// ``` pub fn open(&self, addr: impl AsRef) -> io::Result { // Safety: We're calling open_with_security_attributes_raw w/ a null // pointer which disables it. unsafe { self.open_with_security_attributes_raw(addr, ptr::null_mut()) } } /// Opens the named pipe identified by `addr`. /// /// This is the same as [`open`] except that it supports providing the raw /// pointer to a structure of [`SECURITY_ATTRIBUTES`] which will be passed /// as the `lpSecurityAttributes` argument to [`CreateFile`]. /// /// # Safety /// /// The `attrs` argument must either be null or point at a valid instance of /// the [`SECURITY_ATTRIBUTES`] structure. If the argument is null, the /// behavior is identical to calling the [`open`] method. /// /// [`open`]: ClientOptions::open /// [`CreateFile`]: https://docs.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew /// [`SECURITY_ATTRIBUTES`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Security/struct.SECURITY_ATTRIBUTES.html pub unsafe fn open_with_security_attributes_raw( &self, addr: impl AsRef, attrs: *mut c_void, ) -> io::Result { let addr = encode_addr(addr); let desired_access = { let mut access = 0; if self.generic_read { access |= windows_sys::GENERIC_READ; } if self.generic_write { access |= windows_sys::GENERIC_WRITE; } access }; // NB: We could use a platform specialized `OpenOptions` here, but since // we have access to windows_sys it ultimately doesn't hurt to use // `CreateFile` explicitly since it allows the use of our already // well-structured wide `addr` to pass into CreateFileW. let h = windows_sys::CreateFileW( addr.as_ptr(), desired_access, 0, attrs as *mut _, windows_sys::OPEN_EXISTING, self.get_flags(), 0, ); if h == windows_sys::INVALID_HANDLE_VALUE { return Err(io::Error::last_os_error()); } if matches!(self.pipe_mode, PipeMode::Message) { let mode = windows_sys::PIPE_READMODE_MESSAGE; let result = windows_sys::SetNamedPipeHandleState(h, &mode, ptr::null_mut(), ptr::null_mut()); if result == 0 { return Err(io::Error::last_os_error()); } } NamedPipeClient::from_raw_handle(h as _) } fn get_flags(&self) -> u32 { self.security_qos_flags | windows_sys::FILE_FLAG_OVERLAPPED } } /// The pipe mode of a named pipe. /// /// Set through [`ServerOptions::pipe_mode`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[non_exhaustive] pub enum PipeMode { /// Data is written to the pipe as a stream of bytes. The pipe does not /// distinguish bytes written during different write operations. /// /// Corresponds to [`PIPE_TYPE_BYTE`]. /// /// [`PIPE_TYPE_BYTE`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_TYPE_BYTE.html Byte, /// Data is written to the pipe as a stream of messages. The pipe treats the /// bytes written during each write operation as a message unit. Any reading /// on a named pipe returns [`ERROR_MORE_DATA`] when a message is not read /// completely. /// /// Corresponds to [`PIPE_TYPE_MESSAGE`]. /// /// [`ERROR_MORE_DATA`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/Foundation/constant.ERROR_MORE_DATA.html /// [`PIPE_TYPE_MESSAGE`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_TYPE_MESSAGE.html Message, } /// Indicates the end of a named pipe. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[non_exhaustive] pub enum PipeEnd { /// The named pipe refers to the client end of a named pipe instance. /// /// Corresponds to [`PIPE_CLIENT_END`]. /// /// [`PIPE_CLIENT_END`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_CLIENT_END.html Client, /// The named pipe refers to the server end of a named pipe instance. /// /// Corresponds to [`PIPE_SERVER_END`]. /// /// [`PIPE_SERVER_END`]: https://docs.rs/windows-sys/latest/windows_sys/Win32/System/Pipes/constant.PIPE_SERVER_END.html Server, } /// Information about a named pipe. /// /// Constructed through [`NamedPipeServer::info`] or [`NamedPipeClient::info`]. #[derive(Debug, Clone)] #[non_exhaustive] pub struct PipeInfo { /// Indicates the mode of a named pipe. pub mode: PipeMode, /// Indicates the end of a named pipe. pub end: PipeEnd, /// The maximum number of instances that can be created for this pipe. pub max_instances: u32, /// The number of bytes to reserve for the output buffer. pub out_buffer_size: u32, /// The number of bytes to reserve for the input buffer. pub in_buffer_size: u32, } /// Encodes an address so that it is a null-terminated wide string. fn encode_addr(addr: impl AsRef) -> Box<[u16]> { let len = addr.as_ref().encode_wide().count(); let mut vec = Vec::with_capacity(len + 1); vec.extend(addr.as_ref().encode_wide()); vec.push(0); vec.into_boxed_slice() } /// Internal function to get the info out of a raw named pipe. unsafe fn named_pipe_info(handle: RawHandle) -> io::Result { let mut flags = 0; let mut out_buffer_size = 0; let mut in_buffer_size = 0; let mut max_instances = 0; let result = windows_sys::GetNamedPipeInfo( handle as _, &mut flags, &mut out_buffer_size, &mut in_buffer_size, &mut max_instances, ); if result == 0 { return Err(io::Error::last_os_error()); } let mut end = PipeEnd::Client; let mut mode = PipeMode::Byte; if flags & windows_sys::PIPE_SERVER_END != 0 { end = PipeEnd::Server; } if flags & windows_sys::PIPE_TYPE_MESSAGE != 0 { mode = PipeMode::Message; } Ok(PipeInfo { end, mode, out_buffer_size, in_buffer_size, max_instances, }) } tokio-1.43.1/src/process/kill.rs000064400000000000000000000004241046102023000146050ustar 00000000000000use std::io; /// An interface for killing a running process. pub(crate) trait Kill { /// Forcefully kills the process. fn kill(&mut self) -> io::Result<()>; } impl Kill for &mut T { fn kill(&mut self) -> io::Result<()> { (**self).kill() } } tokio-1.43.1/src/process/mod.rs000064400000000000000000001536431046102023000144450ustar 00000000000000//! An implementation of asynchronous process management for Tokio. //! //! This module provides a [`Command`] struct that imitates the interface of the //! [`std::process::Command`] type in the standard library, but provides asynchronous versions of //! functions that create processes. These functions (`spawn`, `status`, `output` and their //! variants) return "future aware" types that interoperate with Tokio. The asynchronous process //! support is provided through signal handling on Unix and system APIs on Windows. //! //! [`std::process::Command`]: std::process::Command //! //! # Examples //! //! Here's an example program which will spawn `echo hello world` and then wait //! for it complete. //! //! ```no_run //! use tokio::process::Command; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! // The usage is similar as with the standard library's `Command` type //! let mut child = Command::new("echo") //! .arg("hello") //! .arg("world") //! .spawn() //! .expect("failed to spawn"); //! //! // Await until the command completes //! let status = child.wait().await?; //! println!("the command exited with: {}", status); //! Ok(()) //! } //! ``` //! //! Next, let's take a look at an example where we not only spawn `echo hello //! world` but we also capture its output. //! //! ```no_run //! use tokio::process::Command; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! // Like above, but use `output` which returns a future instead of //! // immediately returning the `Child`. //! let output = Command::new("echo").arg("hello").arg("world") //! .output(); //! //! let output = output.await?; //! //! assert!(output.status.success()); //! assert_eq!(output.stdout, b"hello world\n"); //! Ok(()) //! } //! ``` //! //! We can also read input line by line. //! //! ```no_run //! use tokio::io::{BufReader, AsyncBufReadExt}; //! use tokio::process::Command; //! //! use std::process::Stdio; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! let mut cmd = Command::new("cat"); //! //! // Specify that we want the command's standard output piped back to us. //! // By default, standard input/output/error will be inherited from the //! // current process (for example, this means that standard input will //! // come from the keyboard and standard output/error will go directly to //! // the terminal if this process is invoked from the command line). //! cmd.stdout(Stdio::piped()); //! //! let mut child = cmd.spawn() //! .expect("failed to spawn command"); //! //! let stdout = child.stdout.take() //! .expect("child did not have a handle to stdout"); //! //! let mut reader = BufReader::new(stdout).lines(); //! //! // Ensure the child process is spawned in the runtime so it can //! // make progress on its own while we await for any output. //! tokio::spawn(async move { //! let status = child.wait().await //! .expect("child process encountered an error"); //! //! println!("child status was: {}", status); //! }); //! //! while let Some(line) = reader.next_line().await? { //! println!("Line: {}", line); //! } //! //! Ok(()) //! } //! ``` //! //! Here is another example using `sort` writing into the child process //! standard input, capturing the output of the sorted text. //! //! ```no_run //! use tokio::io::AsyncWriteExt; //! use tokio::process::Command; //! //! use std::process::Stdio; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! let mut cmd = Command::new("sort"); //! //! // Specifying that we want pipe both the output and the input. //! // Similarly to capturing the output, by configuring the pipe //! // to stdin it can now be used as an asynchronous writer. //! cmd.stdout(Stdio::piped()); //! cmd.stdin(Stdio::piped()); //! //! let mut child = cmd.spawn().expect("failed to spawn command"); //! //! // These are the animals we want to sort //! let animals: &[&str] = &["dog", "bird", "frog", "cat", "fish"]; //! //! let mut stdin = child //! .stdin //! .take() //! .expect("child did not have a handle to stdin"); //! //! // Write our animals to the child process //! // Note that the behavior of `sort` is to buffer _all input_ before writing any output. //! // In the general sense, it is recommended to write to the child in a separate task as //! // awaiting its exit (or output) to avoid deadlocks (for example, the child tries to write //! // some output but gets stuck waiting on the parent to read from it, meanwhile the parent //! // is stuck waiting to write its input completely before reading the output). //! stdin //! .write(animals.join("\n").as_bytes()) //! .await //! .expect("could not write to stdin"); //! //! // We drop the handle here which signals EOF to the child process. //! // This tells the child process that it there is no more data on the pipe. //! drop(stdin); //! //! let op = child.wait_with_output().await?; //! //! // Results should come back in sorted order //! assert_eq!(op.stdout, "bird\ncat\ndog\nfish\nfrog\n".as_bytes()); //! //! Ok(()) //! } //! ``` //! //! With some coordination, we can also pipe the output of one command into //! another. //! //! ```no_run //! use tokio::join; //! use tokio::process::Command; //! use std::process::Stdio; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! let mut echo = Command::new("echo") //! .arg("hello world!") //! .stdout(Stdio::piped()) //! .spawn() //! .expect("failed to spawn echo"); //! //! let tr_stdin: Stdio = echo //! .stdout //! .take() //! .unwrap() //! .try_into() //! .expect("failed to convert to Stdio"); //! //! let tr = Command::new("tr") //! .arg("a-z") //! .arg("A-Z") //! .stdin(tr_stdin) //! .stdout(Stdio::piped()) //! .spawn() //! .expect("failed to spawn tr"); //! //! let (echo_result, tr_output) = join!(echo.wait(), tr.wait_with_output()); //! //! assert!(echo_result.unwrap().success()); //! //! let tr_output = tr_output.expect("failed to await tr"); //! assert!(tr_output.status.success()); //! //! assert_eq!(tr_output.stdout, b"HELLO WORLD!\n"); //! //! Ok(()) //! } //! ``` //! //! # Caveats //! //! ## Dropping/Cancellation //! //! Similar to the behavior to the standard library, and unlike the futures //! paradigm of dropping-implies-cancellation, a spawned process will, by //! default, continue to execute even after the `Child` handle has been dropped. //! //! The [`Command::kill_on_drop`] method can be used to modify this behavior //! and kill the child process if the `Child` wrapper is dropped before it //! has exited. //! //! ## Unix Processes //! //! On Unix platforms processes must be "reaped" by their parent process after //! they have exited in order to release all OS resources. A child process which //! has exited, but has not yet been reaped by its parent is considered a "zombie" //! process. Such processes continue to count against limits imposed by the system, //! and having too many zombie processes present can prevent additional processes //! from being spawned. //! //! The tokio runtime will, on a best-effort basis, attempt to reap and clean up //! any process which it has spawned. No additional guarantees are made with regard to //! how quickly or how often this procedure will take place. //! //! It is recommended to avoid dropping a [`Child`] process handle before it has been //! fully `await`ed if stricter cleanup guarantees are required. //! //! [`Command`]: crate::process::Command //! [`Command::kill_on_drop`]: crate::process::Command::kill_on_drop //! [`Child`]: crate::process::Child #[path = "unix/mod.rs"] #[cfg(unix)] mod imp; #[cfg(unix)] pub(crate) mod unix { pub(crate) use super::imp::*; } #[path = "windows.rs"] #[cfg(windows)] mod imp; mod kill; use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; use crate::process::kill::Kill; use std::ffi::OsStr; use std::future::Future; use std::io; use std::path::Path; use std::pin::Pin; use std::process::{Command as StdCommand, ExitStatus, Output, Stdio}; use std::task::{ready, Context, Poll}; #[cfg(unix)] use std::os::unix::process::CommandExt; #[cfg(windows)] use std::os::windows::process::CommandExt; cfg_windows! { use crate::os::windows::io::{AsRawHandle, RawHandle}; } /// This structure mimics the API of [`std::process::Command`] found in the standard library, but /// replaces functions that create a process with an asynchronous variant. The main provided /// asynchronous functions are [spawn](Command::spawn), [status](Command::status), and /// [output](Command::output). /// /// `Command` uses asynchronous versions of some `std` types (for example [`Child`]). /// /// [`std::process::Command`]: std::process::Command /// [`Child`]: struct@Child #[derive(Debug)] pub struct Command { std: StdCommand, kill_on_drop: bool, } pub(crate) struct SpawnedChild { child: imp::Child, stdin: Option, stdout: Option, stderr: Option, } impl Command { /// Constructs a new `Command` for launching the program at /// path `program`, with the following default configuration: /// /// * No arguments to the program /// * Inherit the current process's environment /// * Inherit the current process's working directory /// * Inherit stdin/stdout/stderr for `spawn` or `status`, but create pipes for `output` /// /// Builder methods are provided to change these defaults and /// otherwise configure the process. /// /// If `program` is not an absolute path, the `PATH` will be searched in /// an OS-defined way. /// /// The search path to be used may be controlled by setting the /// `PATH` environment variable on the Command, /// but this has some implementation limitations on Windows /// (see issue [rust-lang/rust#37519]). /// /// # Examples /// /// Basic usage: /// /// ```no_run /// use tokio::process::Command; /// let mut command = Command::new("sh"); /// # let _ = command.output(); // assert borrow checker /// ``` /// /// [rust-lang/rust#37519]: https://github.com/rust-lang/rust/issues/37519 pub fn new>(program: S) -> Command { Self::from(StdCommand::new(program)) } /// Cheaply convert to a `&std::process::Command` for places where the type from the standard /// library is expected. pub fn as_std(&self) -> &StdCommand { &self.std } /// Cheaply convert to a `&mut std::process::Command` for places where the type from the /// standard library is expected. pub fn as_std_mut(&mut self) -> &mut StdCommand { &mut self.std } /// Cheaply convert into a `std::process::Command`. /// /// Note that Tokio specific options will be lost. Currently, this only applies to [`kill_on_drop`]. /// /// [`kill_on_drop`]: Command::kill_on_drop pub fn into_std(self) -> StdCommand { self.std } /// Adds an argument to pass to the program. /// /// Only one argument can be passed per use. So instead of: /// /// ```no_run /// let mut command = tokio::process::Command::new("sh"); /// command.arg("-C /path/to/repo"); /// /// # let _ = command.output(); // assert borrow checker /// ``` /// /// usage would be: /// /// ```no_run /// let mut command = tokio::process::Command::new("sh"); /// command.arg("-C"); /// command.arg("/path/to/repo"); /// /// # let _ = command.output(); // assert borrow checker /// ``` /// /// To pass multiple arguments see [`args`]. /// /// [`args`]: method@Self::args /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .arg("-l") /// .arg("-a") /// .output().await.unwrap(); /// # } /// /// ``` pub fn arg>(&mut self, arg: S) -> &mut Command { self.std.arg(arg); self } /// Adds multiple arguments to pass to the program. /// /// To pass a single argument see [`arg`]. /// /// [`arg`]: method@Self::arg /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .args(&["-l", "-a"]) /// .output().await.unwrap(); /// # } /// ``` pub fn args(&mut self, args: I) -> &mut Command where I: IntoIterator, S: AsRef, { self.std.args(args); self } cfg_windows! { /// Append literal text to the command line without any quoting or escaping. /// /// This is useful for passing arguments to `cmd.exe /c`, which doesn't follow /// `CommandLineToArgvW` escaping rules. pub fn raw_arg>(&mut self, text_to_append_as_is: S) -> &mut Command { self.std.raw_arg(text_to_append_as_is); self } } /// Inserts or updates an environment variable mapping. /// /// Note that environment variable names are case-insensitive (but case-preserving) on Windows, /// and case-sensitive on all other platforms. /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .env("PATH", "/bin") /// .output().await.unwrap(); /// # } /// ``` pub fn env(&mut self, key: K, val: V) -> &mut Command where K: AsRef, V: AsRef, { self.std.env(key, val); self } /// Adds or updates multiple environment variable mappings. /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::{Stdio}; /// use std::env; /// use std::collections::HashMap; /// /// let filtered_env : HashMap = /// env::vars().filter(|&(ref k, _)| /// k == "TERM" || k == "TZ" || k == "LANG" || k == "PATH" /// ).collect(); /// /// let output = Command::new("printenv") /// .stdin(Stdio::null()) /// .stdout(Stdio::inherit()) /// .env_clear() /// .envs(&filtered_env) /// .output().await.unwrap(); /// # } /// ``` pub fn envs(&mut self, vars: I) -> &mut Command where I: IntoIterator, K: AsRef, V: AsRef, { self.std.envs(vars); self } /// Removes an environment variable mapping. /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .env_remove("PATH") /// .output().await.unwrap(); /// # } /// ``` pub fn env_remove>(&mut self, key: K) -> &mut Command { self.std.env_remove(key); self } /// Clears the entire environment map for the child process. /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .env_clear() /// .output().await.unwrap(); /// # } /// ``` pub fn env_clear(&mut self) -> &mut Command { self.std.env_clear(); self } /// Sets the working directory for the child process. /// /// # Platform-specific behavior /// /// If the program path is relative (e.g., `"./script.sh"`), it's ambiguous /// whether it should be interpreted relative to the parent's working /// directory or relative to `current_dir`. The behavior in this case is /// platform specific and unstable, and it's recommended to use /// [`canonicalize`] to get an absolute program path instead. /// /// [`canonicalize`]: crate::fs::canonicalize() /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .current_dir("/bin") /// .output().await.unwrap(); /// # } /// ``` pub fn current_dir>(&mut self, dir: P) -> &mut Command { self.std.current_dir(dir); self } /// Sets configuration for the child process's standard input (stdin) handle. /// /// Defaults to [`inherit`]. /// /// [`inherit`]: std::process::Stdio::inherit /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use std::process::{Stdio}; /// use tokio::process::Command; /// /// let output = Command::new("ls") /// .stdin(Stdio::null()) /// .output().await.unwrap(); /// # } /// ``` pub fn stdin>(&mut self, cfg: T) -> &mut Command { self.std.stdin(cfg); self } /// Sets configuration for the child process's standard output (stdout) handle. /// /// Defaults to [`inherit`] when used with `spawn` or `status`, and /// defaults to [`piped`] when used with `output`. /// /// [`inherit`]: std::process::Stdio::inherit /// [`piped`]: std::process::Stdio::piped /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::Stdio; /// /// let output = Command::new("ls") /// .stdout(Stdio::null()) /// .output().await.unwrap(); /// # } /// ``` pub fn stdout>(&mut self, cfg: T) -> &mut Command { self.std.stdout(cfg); self } /// Sets configuration for the child process's standard error (stderr) handle. /// /// Defaults to [`inherit`] when used with `spawn` or `status`, and /// defaults to [`piped`] when used with `output`. /// /// [`inherit`]: std::process::Stdio::inherit /// [`piped`]: std::process::Stdio::piped /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::{Stdio}; /// /// let output = Command::new("ls") /// .stderr(Stdio::null()) /// .output().await.unwrap(); /// # } /// ``` pub fn stderr>(&mut self, cfg: T) -> &mut Command { self.std.stderr(cfg); self } /// Controls whether a `kill` operation should be invoked on a spawned child /// process when its corresponding `Child` handle is dropped. /// /// By default, this value is assumed to be `false`, meaning the next spawned /// process will not be killed on drop, similar to the behavior of the standard /// library. /// /// # Caveats /// /// On Unix platforms processes must be "reaped" by their parent process after /// they have exited in order to release all OS resources. A child process which /// has exited, but has not yet been reaped by its parent is considered a "zombie" /// process. Such processes continue to count against limits imposed by the system, /// and having too many zombie processes present can prevent additional processes /// from being spawned. /// /// Although issuing a `kill` signal to the child process is a synchronous /// operation, the resulting zombie process cannot be `.await`ed inside of the /// destructor to avoid blocking other tasks. The tokio runtime will, on a /// best-effort basis, attempt to reap and clean up such processes in the /// background, but no additional guarantees are made with regard to /// how quickly or how often this procedure will take place. /// /// If stronger guarantees are required, it is recommended to avoid dropping /// a [`Child`] handle where possible, and instead utilize `child.wait().await` /// or `child.kill().await` where possible. pub fn kill_on_drop(&mut self, kill_on_drop: bool) -> &mut Command { self.kill_on_drop = kill_on_drop; self } cfg_windows! { /// Sets the [process creation flags][1] to be passed to `CreateProcess`. /// /// These will always be ORed with `CREATE_UNICODE_ENVIRONMENT`. /// /// [1]: https://msdn.microsoft.com/en-us/library/windows/desktop/ms684863(v=vs.85).aspx pub fn creation_flags(&mut self, flags: u32) -> &mut Command { self.std.creation_flags(flags); self } } /// Sets the child process's user ID. This translates to a /// `setuid` call in the child process. Failure in the `setuid` /// call will cause the spawn to fail. #[cfg(unix)] #[cfg_attr(docsrs, doc(cfg(unix)))] pub fn uid(&mut self, id: u32) -> &mut Command { #[cfg(target_os = "nto")] let id = id as i32; self.std.uid(id); self } /// Similar to `uid` but sets the group ID of the child process. This has /// the same semantics as the `uid` field. #[cfg(unix)] #[cfg_attr(docsrs, doc(cfg(unix)))] pub fn gid(&mut self, id: u32) -> &mut Command { #[cfg(target_os = "nto")] let id = id as i32; self.std.gid(id); self } /// Sets executable argument. /// /// Set the first process argument, `argv[0]`, to something other than the /// default executable path. #[cfg(unix)] #[cfg_attr(docsrs, doc(cfg(unix)))] pub fn arg0(&mut self, arg: S) -> &mut Command where S: AsRef, { self.std.arg0(arg); self } /// Schedules a closure to be run just before the `exec` function is /// invoked. /// /// The closure is allowed to return an I/O error whose OS error code will /// be communicated back to the parent and returned as an error from when /// the spawn was requested. /// /// Multiple closures can be registered and they will be called in order of /// their registration. If a closure returns `Err` then no further closures /// will be called and the spawn operation will immediately return with a /// failure. /// /// # Safety /// /// This closure will be run in the context of the child process after a /// `fork`. This primarily means that any modifications made to memory on /// behalf of this closure will **not** be visible to the parent process. /// This is often a very constrained environment where normal operations /// like `malloc` or acquiring a mutex are not guaranteed to work (due to /// other threads perhaps still running when the `fork` was run). /// /// This also means that all resources such as file descriptors and /// memory-mapped regions got duplicated. It is your responsibility to make /// sure that the closure does not violate library invariants by making /// invalid use of these duplicates. /// /// When this closure is run, aspects such as the stdio file descriptors and /// working directory have successfully been changed, so output to these /// locations may not appear where intended. #[cfg(unix)] #[cfg_attr(docsrs, doc(cfg(unix)))] pub unsafe fn pre_exec(&mut self, f: F) -> &mut Command where F: FnMut() -> io::Result<()> + Send + Sync + 'static, { self.std.pre_exec(f); self } /// Sets the process group ID (PGID) of the child process. Equivalent to a /// `setpgid` call in the child process, but may be more efficient. /// /// Process groups determine which processes receive signals. /// /// # Examples /// /// Pressing Ctrl-C in a terminal will send `SIGINT` to all processes /// in the current foreground process group. By spawning the `sleep` /// subprocess in a new process group, it will not receive `SIGINT` /// from the terminal. /// /// The parent process could install a [signal handler] and manage the /// process on its own terms. /// /// A process group ID of 0 will use the process ID as the PGID. /// /// ```no_run /// # async fn test() { // allow using await /// use tokio::process::Command; /// /// let output = Command::new("sleep") /// .arg("10") /// .process_group(0) /// .output() /// .await /// .unwrap(); /// # } /// ``` /// /// [signal handler]: crate::signal #[cfg(unix)] #[cfg_attr(docsrs, doc(cfg(unix)))] pub fn process_group(&mut self, pgroup: i32) -> &mut Command { self.std.process_group(pgroup); self } /// Executes the command as a child process, returning a handle to it. /// /// By default, stdin, stdout and stderr are inherited from the parent. /// /// This method will spawn the child process synchronously and return a /// handle to a future-aware child process. The `Child` returned implements /// `Future` itself to acquire the `ExitStatus` of the child, and otherwise /// the `Child` has methods to acquire handles to the stdin, stdout, and /// stderr streams. /// /// All I/O this child does will be associated with the current default /// event loop. /// /// # Examples /// /// Basic usage: /// /// ```no_run /// # if cfg!(miri) { return } // No `pidfd_spawnp` in miri. /// use tokio::process::Command; /// /// async fn run_ls() -> std::process::ExitStatus { /// Command::new("ls") /// .spawn() /// .expect("ls command failed to start") /// .wait() /// .await /// .expect("ls command failed to run") /// } /// ``` /// /// # Caveats /// /// ## Dropping/Cancellation /// /// Similar to the behavior to the standard library, and unlike the futures /// paradigm of dropping-implies-cancellation, a spawned process will, by /// default, continue to execute even after the `Child` handle has been dropped. /// /// The [`Command::kill_on_drop`] method can be used to modify this behavior /// and kill the child process if the `Child` wrapper is dropped before it /// has exited. /// /// ## Unix Processes /// /// On Unix platforms processes must be "reaped" by their parent process after /// they have exited in order to release all OS resources. A child process which /// has exited, but has not yet been reaped by its parent is considered a "zombie" /// process. Such processes continue to count against limits imposed by the system, /// and having too many zombie processes present can prevent additional processes /// from being spawned. /// /// The tokio runtime will, on a best-effort basis, attempt to reap and clean up /// any process which it has spawned. No additional guarantees are made with regard to /// how quickly or how often this procedure will take place. /// /// It is recommended to avoid dropping a [`Child`] process handle before it has been /// fully `await`ed if stricter cleanup guarantees are required. /// /// [`Command`]: crate::process::Command /// [`Command::kill_on_drop`]: crate::process::Command::kill_on_drop /// [`Child`]: crate::process::Child /// /// # Errors /// /// On Unix platforms this method will fail with `std::io::ErrorKind::WouldBlock` /// if the system process limit is reached (which includes other applications /// running on the system). pub fn spawn(&mut self) -> io::Result { imp::spawn_child(&mut self.std).map(|spawned_child| Child { child: FusedChild::Child(ChildDropGuard { inner: spawned_child.child, kill_on_drop: self.kill_on_drop, }), stdin: spawned_child.stdin.map(|inner| ChildStdin { inner }), stdout: spawned_child.stdout.map(|inner| ChildStdout { inner }), stderr: spawned_child.stderr.map(|inner| ChildStderr { inner }), }) } /// Executes the command as a child process, waiting for it to finish and /// collecting its exit status. /// /// By default, stdin, stdout and stderr are inherited from the parent. /// If any input/output handles are set to a pipe then they will be immediately /// closed after the child is spawned. /// /// All I/O this child does will be associated with the current default /// event loop. /// /// The destructor of the future returned by this function will kill /// the child if [`kill_on_drop`] is set to true. /// /// [`kill_on_drop`]: fn@Self::kill_on_drop /// /// # Errors /// /// This future will return an error if the child process cannot be spawned /// or if there is an error while awaiting its status. /// /// On Unix platforms this method will fail with `std::io::ErrorKind::WouldBlock` /// if the system process limit is reached (which includes other applications /// running on the system). /// /// # Examples /// /// Basic usage: /// /// ```no_run /// use tokio::process::Command; /// /// async fn run_ls() -> std::process::ExitStatus { /// Command::new("ls") /// .status() /// .await /// .expect("ls command failed to run") /// } /// ``` pub fn status(&mut self) -> impl Future> { let child = self.spawn(); async { let mut child = child?; // Ensure we close any stdio handles so we can't deadlock // waiting on the child which may be waiting to read/write // to a pipe we're holding. child.stdin.take(); child.stdout.take(); child.stderr.take(); child.wait().await } } /// Executes the command as a child process, waiting for it to finish and /// collecting all of its output. /// /// > **Note**: this method, unlike the standard library, will /// > unconditionally configure the stdout/stderr handles to be pipes, even /// > if they have been previously configured. If this is not desired then /// > the `spawn` method should be used in combination with the /// > `wait_with_output` method on child. /// /// This method will return a future representing the collection of the /// child process's stdout/stderr. It will resolve to /// the `Output` type in the standard library, containing `stdout` and /// `stderr` as `Vec` along with an `ExitStatus` representing how the /// process exited. /// /// All I/O this child does will be associated with the current default /// event loop. /// /// The destructor of the future returned by this function will kill /// the child if [`kill_on_drop`] is set to true. /// /// [`kill_on_drop`]: fn@Self::kill_on_drop /// /// # Errors /// /// This future will return an error if the child process cannot be spawned /// or if there is an error while awaiting its status. /// /// On Unix platforms this method will fail with `std::io::ErrorKind::WouldBlock` /// if the system process limit is reached (which includes other applications /// running on the system). /// # Examples /// /// Basic usage: /// /// ```no_run /// use tokio::process::Command; /// /// async fn run_ls() { /// let output: std::process::Output = Command::new("ls") /// .output() /// .await /// .expect("ls command failed to run"); /// println!("stderr of ls: {:?}", output.stderr); /// } /// ``` pub fn output(&mut self) -> impl Future> { self.std.stdout(Stdio::piped()); self.std.stderr(Stdio::piped()); let child = self.spawn(); async { child?.wait_with_output().await } } } impl From for Command { fn from(std: StdCommand) -> Command { Command { std, kill_on_drop: false, } } } /// A drop guard which can ensure the child process is killed on drop if specified. #[derive(Debug)] struct ChildDropGuard { inner: T, kill_on_drop: bool, } impl Kill for ChildDropGuard { fn kill(&mut self) -> io::Result<()> { let ret = self.inner.kill(); if ret.is_ok() { self.kill_on_drop = false; } ret } } impl Drop for ChildDropGuard { fn drop(&mut self) { if self.kill_on_drop { drop(self.kill()); } } } impl Future for ChildDropGuard where F: Future> + Kill + Unpin, { type Output = Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = Pin::new(&mut self.inner).poll(cx); if let Poll::Ready(Ok(_)) = ret { // Avoid the overhead of trying to kill a reaped process self.kill_on_drop = false; } if ret.is_ready() { coop.made_progress(); } ret } } /// Keeps track of the exit status of a child process without worrying about /// polling the underlying futures even after they have completed. #[derive(Debug)] enum FusedChild { Child(ChildDropGuard), Done(ExitStatus), } /// Representation of a child process spawned onto an event loop. /// /// # Caveats /// Similar to the behavior to the standard library, and unlike the futures /// paradigm of dropping-implies-cancellation, a spawned process will, by /// default, continue to execute even after the `Child` handle has been dropped. /// /// The `Command::kill_on_drop` method can be used to modify this behavior /// and kill the child process if the `Child` wrapper is dropped before it /// has exited. #[derive(Debug)] pub struct Child { child: FusedChild, /// The handle for writing to the child's standard input (stdin), if it has /// been captured. To avoid partially moving the `child` and thus blocking /// yourself from calling functions on `child` while using `stdin`, you might /// find it helpful to do: /// /// ```no_run /// # let mut child = tokio::process::Command::new("echo").spawn().unwrap(); /// let stdin = child.stdin.take().unwrap(); /// ``` pub stdin: Option, /// The handle for reading from the child's standard output (stdout), if it /// has been captured. You might find it helpful to do /// /// ```no_run /// # let mut child = tokio::process::Command::new("echo").spawn().unwrap(); /// let stdout = child.stdout.take().unwrap(); /// ``` /// /// to avoid partially moving the `child` and thus blocking yourself from calling /// functions on `child` while using `stdout`. pub stdout: Option, /// The handle for reading from the child's standard error (stderr), if it /// has been captured. You might find it helpful to do /// /// ```no_run /// # let mut child = tokio::process::Command::new("echo").spawn().unwrap(); /// let stderr = child.stderr.take().unwrap(); /// ``` /// /// to avoid partially moving the `child` and thus blocking yourself from calling /// functions on `child` while using `stderr`. pub stderr: Option, } impl Child { /// Returns the OS-assigned process identifier associated with this child /// while it is still running. /// /// Once the child has been polled to completion this will return `None`. /// This is done to avoid confusion on platforms like Unix where the OS /// identifier could be reused once the process has completed. pub fn id(&self) -> Option { match &self.child { FusedChild::Child(child) => Some(child.inner.id()), FusedChild::Done(_) => None, } } cfg_windows! { /// Extracts the raw handle of the process associated with this child while /// it is still running. Returns `None` if the child has exited. pub fn raw_handle(&self) -> Option { match &self.child { FusedChild::Child(c) => Some(c.inner.as_raw_handle()), FusedChild::Done(_) => None, } } } /// Attempts to force the child to exit, but does not wait for the request /// to take effect. /// /// On Unix platforms, this is the equivalent to sending a `SIGKILL`. Note /// that on Unix platforms it is possible for a zombie process to remain /// after a kill is sent; to avoid this, the caller should ensure that either /// `child.wait().await` or `child.try_wait()` is invoked successfully. pub fn start_kill(&mut self) -> io::Result<()> { match &mut self.child { FusedChild::Child(child) => child.kill(), FusedChild::Done(_) => Err(io::Error::new( io::ErrorKind::InvalidInput, "invalid argument: can't kill an exited process", )), } } /// Forces the child to exit. /// /// This is equivalent to sending a `SIGKILL` on unix platforms. /// /// If the child has to be killed remotely, it is possible to do it using /// a combination of the select! macro and a `oneshot` channel. In the following /// example, the child will run until completion unless a message is sent on /// the `oneshot` channel. If that happens, the child is killed immediately /// using the `.kill()` method. /// /// ```no_run /// use tokio::process::Command; /// use tokio::sync::oneshot::channel; /// /// #[tokio::main] /// async fn main() { /// let (send, recv) = channel::<()>(); /// let mut child = Command::new("sleep").arg("1").spawn().unwrap(); /// tokio::spawn(async move { send.send(()) }); /// tokio::select! { /// _ = child.wait() => {} /// _ = recv => child.kill().await.expect("kill failed"), /// } /// } /// ``` pub async fn kill(&mut self) -> io::Result<()> { self.start_kill()?; self.wait().await?; Ok(()) } /// Waits for the child to exit completely, returning the status that it /// exited with. This function will continue to have the same return value /// after it has been called at least once. /// /// The stdin handle to the child process, if any, will be closed /// before waiting. This helps avoid deadlock: it ensures that the /// child does not block waiting for input from the parent, while /// the parent waits for the child to exit. /// /// If the caller wishes to explicitly control when the child's stdin /// handle is closed, they may `.take()` it before calling `.wait()`: /// /// # Cancel safety /// /// This function is cancel safe. /// /// ``` /// # if cfg!(miri) { return } // No `pidfd_spawnp` in miri. /// # #[cfg(not(unix))]fn main(){} /// # #[cfg(unix)] /// use tokio::io::AsyncWriteExt; /// # #[cfg(unix)] /// use tokio::process::Command; /// # #[cfg(unix)] /// use std::process::Stdio; /// /// # #[cfg(unix)] /// #[tokio::main] /// async fn main() { /// let mut child = Command::new("cat") /// .stdin(Stdio::piped()) /// .spawn() /// .unwrap(); /// /// let mut stdin = child.stdin.take().unwrap(); /// tokio::spawn(async move { /// // do something with stdin here... /// stdin.write_all(b"hello world\n").await.unwrap(); /// /// // then drop when finished /// drop(stdin); /// }); /// /// // wait for the process to complete /// let _ = child.wait().await; /// } /// ``` pub async fn wait(&mut self) -> io::Result { // Ensure stdin is closed so the child isn't stuck waiting on // input while the parent is waiting for it to exit. drop(self.stdin.take()); match &mut self.child { FusedChild::Done(exit) => Ok(*exit), FusedChild::Child(child) => { let ret = child.await; if let Ok(exit) = ret { self.child = FusedChild::Done(exit); } ret } } } /// Attempts to collect the exit status of the child if it has already /// exited. /// /// This function will not block the calling thread and will only /// check to see if the child process has exited or not. If the child has /// exited then on Unix the process ID is reaped. This function is /// guaranteed to repeatedly return a successful exit status so long as the /// child has already exited. /// /// If the child has exited, then `Ok(Some(status))` is returned. If the /// exit status is not available at this time then `Ok(None)` is returned. /// If an error occurs, then that error is returned. /// /// Note that unlike `wait`, this function will not attempt to drop stdin, /// nor will it wake the current task if the child exits. pub fn try_wait(&mut self) -> io::Result> { match &mut self.child { FusedChild::Done(exit) => Ok(Some(*exit)), FusedChild::Child(guard) => { let ret = guard.inner.try_wait(); if let Ok(Some(exit)) = ret { // Avoid the overhead of trying to kill a reaped process guard.kill_on_drop = false; self.child = FusedChild::Done(exit); } ret } } } /// Returns a future that will resolve to an `Output`, containing the exit /// status, stdout, and stderr of the child process. /// /// The returned future will simultaneously waits for the child to exit and /// collect all remaining output on the stdout/stderr handles, returning an /// `Output` instance. /// /// The stdin handle to the child process, if any, will be closed before /// waiting. This helps avoid deadlock: it ensures that the child does not /// block waiting for input from the parent, while the parent waits for the /// child to exit. /// /// By default, stdin, stdout and stderr are inherited from the parent. In /// order to capture the output into this `Output` it is necessary to create /// new pipes between parent and child. Use `stdout(Stdio::piped())` or /// `stderr(Stdio::piped())`, respectively, when creating a `Command`. pub async fn wait_with_output(mut self) -> io::Result { use crate::future::try_join3; async fn read_to_end(io: &mut Option) -> io::Result> { let mut vec = Vec::new(); if let Some(io) = io.as_mut() { crate::io::util::read_to_end(io, &mut vec).await?; } Ok(vec) } let mut stdout_pipe = self.stdout.take(); let mut stderr_pipe = self.stderr.take(); let stdout_fut = read_to_end(&mut stdout_pipe); let stderr_fut = read_to_end(&mut stderr_pipe); let (status, stdout, stderr) = try_join3(self.wait(), stdout_fut, stderr_fut).await?; // Drop happens after `try_join` due to drop(stdout_pipe); drop(stderr_pipe); Ok(Output { status, stdout, stderr, }) } } /// The standard input stream for spawned children. /// /// This type implements the `AsyncWrite` trait to pass data to the stdin handle of /// handle of a child process asynchronously. #[derive(Debug)] pub struct ChildStdin { inner: imp::ChildStdio, } /// The standard output stream for spawned children. /// /// This type implements the `AsyncRead` trait to read data from the stdout /// handle of a child process asynchronously. #[derive(Debug)] pub struct ChildStdout { inner: imp::ChildStdio, } /// The standard error stream for spawned children. /// /// This type implements the `AsyncRead` trait to read data from the stderr /// handle of a child process asynchronously. #[derive(Debug)] pub struct ChildStderr { inner: imp::ChildStdio, } impl ChildStdin { /// Creates an asynchronous `ChildStdin` from a synchronous one. /// /// # Errors /// /// This method may fail if an error is encountered when setting the pipe to /// non-blocking mode, or when registering the pipe with the runtime's IO /// driver. pub fn from_std(inner: std::process::ChildStdin) -> io::Result { Ok(Self { inner: imp::stdio(inner)?, }) } } impl ChildStdout { /// Creates an asynchronous `ChildStdout` from a synchronous one. /// /// # Errors /// /// This method may fail if an error is encountered when setting the pipe to /// non-blocking mode, or when registering the pipe with the runtime's IO /// driver. pub fn from_std(inner: std::process::ChildStdout) -> io::Result { Ok(Self { inner: imp::stdio(inner)?, }) } } impl ChildStderr { /// Creates an asynchronous `ChildStderr` from a synchronous one. /// /// # Errors /// /// This method may fail if an error is encountered when setting the pipe to /// non-blocking mode, or when registering the pipe with the runtime's IO /// driver. pub fn from_std(inner: std::process::ChildStderr) -> io::Result { Ok(Self { inner: imp::stdio(inner)?, }) } } impl AsyncWrite for ChildStdin { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut self.inner).poll_write(cx, buf) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_shutdown(cx) } fn poll_write_vectored( mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { Pin::new(&mut self.inner).poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { self.inner.is_write_vectored() } } impl AsyncRead for ChildStdout { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut self.inner).poll_read(cx, buf) } } impl AsyncRead for ChildStderr { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut self.inner).poll_read(cx, buf) } } impl TryInto for ChildStdin { type Error = io::Error; fn try_into(self) -> Result { imp::convert_to_stdio(self.inner) } } impl TryInto for ChildStdout { type Error = io::Error; fn try_into(self) -> Result { imp::convert_to_stdio(self.inner) } } impl TryInto for ChildStderr { type Error = io::Error; fn try_into(self) -> Result { imp::convert_to_stdio(self.inner) } } #[cfg(unix)] #[cfg_attr(docsrs, doc(cfg(unix)))] mod sys { use std::{ io, os::unix::io::{AsFd, AsRawFd, BorrowedFd, OwnedFd, RawFd}, }; use super::{ChildStderr, ChildStdin, ChildStdout}; macro_rules! impl_traits { ($type:ty) => { impl $type { /// Convert into [`OwnedFd`]. pub fn into_owned_fd(self) -> io::Result { self.inner.into_owned_fd() } } impl AsRawFd for $type { fn as_raw_fd(&self) -> RawFd { self.inner.as_raw_fd() } } impl AsFd for $type { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } }; } impl_traits!(ChildStdin); impl_traits!(ChildStdout); impl_traits!(ChildStderr); } #[cfg(any(windows, docsrs))] #[cfg_attr(docsrs, doc(cfg(windows)))] mod windows { use super::*; use crate::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, OwnedHandle, RawHandle}; #[cfg(not(docsrs))] macro_rules! impl_traits { ($type:ty) => { impl $type { /// Convert into [`OwnedHandle`]. pub fn into_owned_handle(self) -> io::Result { self.inner.into_owned_handle() } } impl AsRawHandle for $type { fn as_raw_handle(&self) -> RawHandle { self.inner.as_raw_handle() } } impl AsHandle for $type { fn as_handle(&self) -> BorrowedHandle<'_> { unsafe { BorrowedHandle::borrow_raw(self.as_raw_handle()) } } } }; } #[cfg(docsrs)] macro_rules! impl_traits { ($type:ty) => { impl $type { /// Convert into [`OwnedHandle`]. pub fn into_owned_handle(self) -> io::Result { todo!("For doc generation only") } } impl AsRawHandle for $type { fn as_raw_handle(&self) -> RawHandle { todo!("For doc generation only") } } impl AsHandle for $type { fn as_handle(&self) -> BorrowedHandle<'_> { todo!("For doc generation only") } } }; } impl_traits!(ChildStdin); impl_traits!(ChildStdout); impl_traits!(ChildStderr); } #[cfg(all(test, not(loom)))] mod test { use super::kill::Kill; use super::ChildDropGuard; use futures::future::FutureExt; use std::future::Future; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; struct Mock { num_kills: usize, num_polls: usize, poll_result: Poll>, } impl Mock { fn new() -> Self { Self::with_result(Poll::Pending) } fn with_result(result: Poll>) -> Self { Self { num_kills: 0, num_polls: 0, poll_result: result, } } } impl Kill for Mock { fn kill(&mut self) -> io::Result<()> { self.num_kills += 1; Ok(()) } } impl Future for Mock { type Output = Result<(), ()>; fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { let inner = Pin::get_mut(self); inner.num_polls += 1; inner.poll_result } } #[test] fn kills_on_drop_if_specified() { let mut mock = Mock::new(); { let guard = ChildDropGuard { inner: &mut mock, kill_on_drop: true, }; drop(guard); } assert_eq!(1, mock.num_kills); assert_eq!(0, mock.num_polls); } #[test] fn no_kill_on_drop_by_default() { let mut mock = Mock::new(); { let guard = ChildDropGuard { inner: &mut mock, kill_on_drop: false, }; drop(guard); } assert_eq!(0, mock.num_kills); assert_eq!(0, mock.num_polls); } #[test] fn no_kill_if_already_killed() { let mut mock = Mock::new(); { let mut guard = ChildDropGuard { inner: &mut mock, kill_on_drop: true, }; let _ = guard.kill(); drop(guard); } assert_eq!(1, mock.num_kills); assert_eq!(0, mock.num_polls); } #[test] fn no_kill_if_reaped() { let mut mock_pending = Mock::with_result(Poll::Pending); let mut mock_reaped = Mock::with_result(Poll::Ready(Ok(()))); let mut mock_err = Mock::with_result(Poll::Ready(Err(()))); let waker = futures::task::noop_waker(); let mut context = Context::from_waker(&waker); { let mut guard = ChildDropGuard { inner: &mut mock_pending, kill_on_drop: true, }; let _ = guard.poll_unpin(&mut context); let mut guard = ChildDropGuard { inner: &mut mock_reaped, kill_on_drop: true, }; let _ = guard.poll_unpin(&mut context); let mut guard = ChildDropGuard { inner: &mut mock_err, kill_on_drop: true, }; let _ = guard.poll_unpin(&mut context); } assert_eq!(1, mock_pending.num_kills); assert_eq!(1, mock_pending.num_polls); assert_eq!(0, mock_reaped.num_kills); assert_eq!(1, mock_reaped.num_polls); assert_eq!(1, mock_err.num_kills); assert_eq!(1, mock_err.num_polls); } } tokio-1.43.1/src/process/unix/mod.rs000064400000000000000000000243661046102023000154270ustar 00000000000000//! Unix handling of child processes. //! //! Right now the only "fancy" thing about this is how we implement the //! `Future` implementation on `Child` to get the exit status. Unix offers //! no way to register a child with epoll, and the only real way to get a //! notification when a process exits is the SIGCHLD signal. //! //! Signal handling in general is *super* hairy and complicated, and it's even //! more complicated here with the fact that signals are coalesced, so we may //! not get a SIGCHLD-per-child. //! //! Our best approximation here is to check *all spawned processes* for all //! SIGCHLD signals received. To do that we create a `Signal`, implemented in //! the `tokio-net` crate, which is a stream over signals being received. //! //! Later when we poll the process's exit status we simply check to see if a //! SIGCHLD has happened since we last checked, and while that returns "yes" we //! keep trying. //! //! Note that this means that this isn't really scalable, but then again //! processes in general aren't scalable (e.g. millions) so it shouldn't be that //! bad in theory... pub(crate) mod orphan; use orphan::{OrphanQueue, OrphanQueueImpl, Wait}; mod reap; use reap::Reaper; #[cfg(all(target_os = "linux", feature = "rt"))] mod pidfd_reaper; use crate::io::{AsyncRead, AsyncWrite, PollEvented, ReadBuf}; use crate::process::kill::Kill; use crate::process::SpawnedChild; use crate::runtime::signal::Handle as SignalHandle; use crate::signal::unix::{signal, Signal, SignalKind}; use mio::event::Source; use mio::unix::SourceFd; use std::fmt; use std::fs::File; use std::future::Future; use std::io; use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd}; use std::pin::Pin; use std::process::{Child as StdChild, ExitStatus, Stdio}; use std::task::Context; use std::task::Poll; impl Wait for StdChild { fn id(&self) -> u32 { self.id() } fn try_wait(&mut self) -> io::Result> { self.try_wait() } } impl Kill for StdChild { fn kill(&mut self) -> io::Result<()> { self.kill() } } cfg_not_has_const_mutex_new! { fn get_orphan_queue() -> &'static OrphanQueueImpl { use crate::util::once_cell::OnceCell; static ORPHAN_QUEUE: OnceCell> = OnceCell::new(); ORPHAN_QUEUE.get(OrphanQueueImpl::new) } } cfg_has_const_mutex_new! { fn get_orphan_queue() -> &'static OrphanQueueImpl { static ORPHAN_QUEUE: OrphanQueueImpl = OrphanQueueImpl::new(); &ORPHAN_QUEUE } } pub(crate) struct GlobalOrphanQueue; impl fmt::Debug for GlobalOrphanQueue { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { get_orphan_queue().fmt(fmt) } } impl GlobalOrphanQueue { pub(crate) fn reap_orphans(handle: &SignalHandle) { get_orphan_queue().reap_orphans(handle); } } impl OrphanQueue for GlobalOrphanQueue { fn push_orphan(&self, orphan: StdChild) { get_orphan_queue().push_orphan(orphan); } } #[must_use = "futures do nothing unless polled"] pub(crate) enum Child { SignalReaper(Reaper), #[cfg(all(target_os = "linux", feature = "rt"))] PidfdReaper(pidfd_reaper::PidfdReaper), } impl fmt::Debug for Child { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Child").field("pid", &self.id()).finish() } } pub(crate) fn spawn_child(cmd: &mut std::process::Command) -> io::Result { let mut child = cmd.spawn()?; let stdin = child.stdin.take().map(stdio).transpose()?; let stdout = child.stdout.take().map(stdio).transpose()?; let stderr = child.stderr.take().map(stdio).transpose()?; #[cfg(all(target_os = "linux", feature = "rt"))] match pidfd_reaper::PidfdReaper::new(child, GlobalOrphanQueue) { Ok(pidfd_reaper) => { return Ok(SpawnedChild { child: Child::PidfdReaper(pidfd_reaper), stdin, stdout, stderr, }) } Err((Some(err), _child)) => return Err(err), Err((None, child_returned)) => child = child_returned, } let signal = signal(SignalKind::child())?; Ok(SpawnedChild { child: Child::SignalReaper(Reaper::new(child, GlobalOrphanQueue, signal)), stdin, stdout, stderr, }) } impl Child { pub(crate) fn id(&self) -> u32 { match self { Self::SignalReaper(signal_reaper) => signal_reaper.id(), #[cfg(all(target_os = "linux", feature = "rt"))] Self::PidfdReaper(pidfd_reaper) => pidfd_reaper.id(), } } fn std_child(&mut self) -> &mut StdChild { match self { Self::SignalReaper(signal_reaper) => signal_reaper.inner_mut(), #[cfg(all(target_os = "linux", feature = "rt"))] Self::PidfdReaper(pidfd_reaper) => pidfd_reaper.inner_mut(), } } pub(crate) fn try_wait(&mut self) -> io::Result> { self.std_child().try_wait() } } impl Kill for Child { fn kill(&mut self) -> io::Result<()> { self.std_child().kill() } } impl Future for Child { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match Pin::into_inner(self) { Self::SignalReaper(signal_reaper) => Pin::new(signal_reaper).poll(cx), #[cfg(all(target_os = "linux", feature = "rt"))] Self::PidfdReaper(pidfd_reaper) => Pin::new(pidfd_reaper).poll(cx), } } } #[derive(Debug)] pub(crate) struct Pipe { // Actually a pipe is not a File. However, we are reusing `File` to get // close on drop. This is a similar trick as `mio`. fd: File, } impl From for Pipe { fn from(fd: T) -> Self { let fd = unsafe { File::from_raw_fd(fd.into_raw_fd()) }; Self { fd } } } impl<'a> io::Read for &'a Pipe { fn read(&mut self, bytes: &mut [u8]) -> io::Result { (&self.fd).read(bytes) } } impl<'a> io::Write for &'a Pipe { fn write(&mut self, bytes: &[u8]) -> io::Result { (&self.fd).write(bytes) } fn flush(&mut self) -> io::Result<()> { (&self.fd).flush() } fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result { (&self.fd).write_vectored(bufs) } } impl AsRawFd for Pipe { fn as_raw_fd(&self) -> RawFd { self.fd.as_raw_fd() } } impl AsFd for Pipe { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } fn convert_to_blocking_file(io: ChildStdio) -> io::Result { let mut fd = io.inner.into_inner()?.fd; // Ensure that the fd to be inherited is set to *blocking* mode, as this // is the default that virtually all programs expect to have. Those // programs that know how to work with nonblocking stdio will know how to // change it to nonblocking mode. set_nonblocking(&mut fd, false)?; Ok(fd) } pub(crate) fn convert_to_stdio(io: ChildStdio) -> io::Result { convert_to_blocking_file(io).map(Stdio::from) } impl Source for Pipe { fn register( &mut self, registry: &mio::Registry, token: mio::Token, interest: mio::Interest, ) -> io::Result<()> { SourceFd(&self.as_raw_fd()).register(registry, token, interest) } fn reregister( &mut self, registry: &mio::Registry, token: mio::Token, interest: mio::Interest, ) -> io::Result<()> { SourceFd(&self.as_raw_fd()).reregister(registry, token, interest) } fn deregister(&mut self, registry: &mio::Registry) -> io::Result<()> { SourceFd(&self.as_raw_fd()).deregister(registry) } } pub(crate) struct ChildStdio { inner: PollEvented, } impl ChildStdio { pub(super) fn into_owned_fd(self) -> io::Result { convert_to_blocking_file(self).map(OwnedFd::from) } } impl fmt::Debug for ChildStdio { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.inner.fmt(fmt) } } impl AsRawFd for ChildStdio { fn as_raw_fd(&self) -> RawFd { self.inner.as_raw_fd() } } impl AsFd for ChildStdio { fn as_fd(&self) -> BorrowedFd<'_> { unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } } } impl AsyncWrite for ChildStdio { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.inner.poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { self.inner.poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { true } } impl AsyncRead for ChildStdio { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { // Safety: pipes support reading into uninitialized memory unsafe { self.inner.poll_read(cx, buf) } } } fn set_nonblocking(fd: &mut T, nonblocking: bool) -> io::Result<()> { unsafe { let fd = fd.as_raw_fd(); let previous = libc::fcntl(fd, libc::F_GETFL); if previous == -1 { return Err(io::Error::last_os_error()); } let new = if nonblocking { previous | libc::O_NONBLOCK } else { previous & !libc::O_NONBLOCK }; let r = libc::fcntl(fd, libc::F_SETFL, new); if r == -1 { return Err(io::Error::last_os_error()); } } Ok(()) } pub(super) fn stdio(io: T) -> io::Result where T: IntoRawFd, { // Set the fd to nonblocking before we pass it to the event loop let mut pipe = Pipe::from(io); set_nonblocking(&mut pipe, true)?; PollEvented::new(pipe).map(|inner| ChildStdio { inner }) } tokio-1.43.1/src/process/unix/orphan.rs000064400000000000000000000240361046102023000161310ustar 00000000000000use crate::loom::sync::{Mutex, MutexGuard}; use crate::runtime::signal::Handle as SignalHandle; use crate::signal::unix::{signal_with_handle, SignalKind}; use crate::sync::watch; use std::io; use std::process::ExitStatus; /// An interface for waiting on a process to exit. pub(crate) trait Wait { /// Get the identifier for this process or diagnostics. #[allow(dead_code)] fn id(&self) -> u32; /// Try waiting for a process to exit in a non-blocking manner. fn try_wait(&mut self) -> io::Result>; } impl Wait for &mut T { fn id(&self) -> u32 { (**self).id() } fn try_wait(&mut self) -> io::Result> { (**self).try_wait() } } /// An interface for queueing up an orphaned process so that it can be reaped. pub(crate) trait OrphanQueue { /// Adds an orphan to the queue. fn push_orphan(&self, orphan: T); } impl> OrphanQueue for &O { fn push_orphan(&self, orphan: T) { (**self).push_orphan(orphan); } } /// An implementation of `OrphanQueue`. #[derive(Debug)] pub(crate) struct OrphanQueueImpl { sigchild: Mutex>>, queue: Mutex>, } impl OrphanQueueImpl { cfg_not_has_const_mutex_new! { pub(crate) fn new() -> Self { Self { sigchild: Mutex::new(None), queue: Mutex::new(Vec::new()), } } } cfg_has_const_mutex_new! { pub(crate) const fn new() -> Self { Self { sigchild: Mutex::const_new(None), queue: Mutex::const_new(Vec::new()), } } } #[cfg(test)] fn len(&self) -> usize { self.queue.lock().len() } pub(crate) fn push_orphan(&self, orphan: T) where T: Wait, { self.queue.lock().push(orphan); } /// Attempts to reap every process in the queue, ignoring any errors and /// enqueueing any orphans which have not yet exited. pub(crate) fn reap_orphans(&self, handle: &SignalHandle) where T: Wait, { // If someone else is holding the lock, they will be responsible for draining // the queue as necessary, so we can safely bail if that happens if let Some(mut sigchild_guard) = self.sigchild.try_lock() { match &mut *sigchild_guard { Some(sigchild) => { if sigchild.try_has_changed().and_then(Result::ok).is_some() { drain_orphan_queue(self.queue.lock()); } } None => { let queue = self.queue.lock(); // Be lazy and only initialize the SIGCHLD listener if there // are any orphaned processes in the queue. if !queue.is_empty() { // An errors shouldn't really happen here, but if it does it // means that the signal driver isn't running, in // which case there isn't anything we can // register/initialize here, so we can try again later if let Ok(sigchild) = signal_with_handle(SignalKind::child(), handle) { *sigchild_guard = Some(sigchild); drain_orphan_queue(queue); } } } } } } } fn drain_orphan_queue(mut queue: MutexGuard<'_, Vec>) where T: Wait, { for i in (0..queue.len()).rev() { match queue[i].try_wait() { Ok(None) => {} Ok(Some(_)) | Err(_) => { // The stdlib handles interruption errors (EINTR) when polling a child process. // All other errors represent invalid inputs or pids that have already been // reaped, so we can drop the orphan in case an error is raised. queue.swap_remove(i); } } } drop(queue); } #[cfg(all(test, not(loom)))] pub(crate) mod test { use super::*; use crate::runtime::io::Driver as IoDriver; use crate::runtime::signal::{Driver as SignalDriver, Handle as SignalHandle}; use crate::sync::watch; use std::cell::{Cell, RefCell}; use std::io; use std::os::unix::process::ExitStatusExt; use std::process::ExitStatus; use std::rc::Rc; pub(crate) struct MockQueue { pub(crate) all_enqueued: RefCell>, } impl MockQueue { pub(crate) fn new() -> Self { Self { all_enqueued: RefCell::new(Vec::new()), } } } impl OrphanQueue for MockQueue { fn push_orphan(&self, orphan: W) { self.all_enqueued.borrow_mut().push(orphan); } } struct MockWait { total_waits: Rc>, num_wait_until_status: usize, return_err: bool, } impl MockWait { fn new(num_wait_until_status: usize) -> Self { Self { total_waits: Rc::new(Cell::new(0)), num_wait_until_status, return_err: false, } } fn with_err() -> Self { Self { total_waits: Rc::new(Cell::new(0)), num_wait_until_status: 0, return_err: true, } } } impl Wait for MockWait { fn id(&self) -> u32 { 42 } fn try_wait(&mut self) -> io::Result> { let waits = self.total_waits.get(); let ret = if self.num_wait_until_status == waits { if self.return_err { Ok(Some(ExitStatus::from_raw(0))) } else { Err(io::Error::new(io::ErrorKind::Other, "mock err")) } } else { Ok(None) }; self.total_waits.set(waits + 1); ret } } #[test] fn drain_attempts_a_single_reap_of_all_queued_orphans() { let first_orphan = MockWait::new(0); let second_orphan = MockWait::new(1); let third_orphan = MockWait::new(2); let fourth_orphan = MockWait::with_err(); let first_waits = first_orphan.total_waits.clone(); let second_waits = second_orphan.total_waits.clone(); let third_waits = third_orphan.total_waits.clone(); let fourth_waits = fourth_orphan.total_waits.clone(); let orphanage = OrphanQueueImpl::new(); orphanage.push_orphan(first_orphan); orphanage.push_orphan(third_orphan); orphanage.push_orphan(second_orphan); orphanage.push_orphan(fourth_orphan); assert_eq!(orphanage.len(), 4); drain_orphan_queue(orphanage.queue.lock()); assert_eq!(orphanage.len(), 2); assert_eq!(first_waits.get(), 1); assert_eq!(second_waits.get(), 1); assert_eq!(third_waits.get(), 1); assert_eq!(fourth_waits.get(), 1); drain_orphan_queue(orphanage.queue.lock()); assert_eq!(orphanage.len(), 1); assert_eq!(first_waits.get(), 1); assert_eq!(second_waits.get(), 2); assert_eq!(third_waits.get(), 2); assert_eq!(fourth_waits.get(), 1); drain_orphan_queue(orphanage.queue.lock()); assert_eq!(orphanage.len(), 0); assert_eq!(first_waits.get(), 1); assert_eq!(second_waits.get(), 2); assert_eq!(third_waits.get(), 3); assert_eq!(fourth_waits.get(), 1); // Safe to reap when empty drain_orphan_queue(orphanage.queue.lock()); } #[test] fn no_reap_if_no_signal_received() { let (tx, rx) = watch::channel(()); let handle = SignalHandle::default(); let orphanage = OrphanQueueImpl::new(); *orphanage.sigchild.lock() = Some(rx); let orphan = MockWait::new(2); let waits = orphan.total_waits.clone(); orphanage.push_orphan(orphan); orphanage.reap_orphans(&handle); assert_eq!(waits.get(), 0); orphanage.reap_orphans(&handle); assert_eq!(waits.get(), 0); tx.send(()).unwrap(); orphanage.reap_orphans(&handle); assert_eq!(waits.get(), 1); } #[test] fn no_reap_if_signal_lock_held() { let handle = SignalHandle::default(); let orphanage = OrphanQueueImpl::new(); let signal_guard = orphanage.sigchild.lock(); let orphan = MockWait::new(2); let waits = orphan.total_waits.clone(); orphanage.push_orphan(orphan); orphanage.reap_orphans(&handle); assert_eq!(waits.get(), 0); drop(signal_guard); } #[cfg_attr(miri, ignore)] // Miri does not support epoll. #[test] fn does_not_register_signal_if_queue_empty() { let (io_driver, io_handle) = IoDriver::new(1024).unwrap(); let signal_driver = SignalDriver::new(io_driver, &io_handle).unwrap(); let handle = signal_driver.handle(); let orphanage = OrphanQueueImpl::new(); assert!(orphanage.sigchild.lock().is_none()); // Sanity // No register when queue empty orphanage.reap_orphans(&handle); assert!(orphanage.sigchild.lock().is_none()); let orphan = MockWait::new(2); let waits = orphan.total_waits.clone(); orphanage.push_orphan(orphan); orphanage.reap_orphans(&handle); assert!(orphanage.sigchild.lock().is_some()); assert_eq!(waits.get(), 1); // Eager reap when registering listener } #[test] fn does_nothing_if_signal_could_not_be_registered() { let handle = SignalHandle::default(); let orphanage = OrphanQueueImpl::new(); assert!(orphanage.sigchild.lock().is_none()); let orphan = MockWait::new(2); let waits = orphan.total_waits.clone(); orphanage.push_orphan(orphan); // Signal handler has "gone away", nothing to register or reap orphanage.reap_orphans(&handle); assert!(orphanage.sigchild.lock().is_none()); assert_eq!(waits.get(), 0); } } tokio-1.43.1/src/process/unix/pidfd_reaper.rs000064400000000000000000000204051046102023000172620ustar 00000000000000use crate::{ io::{interest::Interest, PollEvented}, process::{ imp::{orphan::Wait, OrphanQueue}, kill::Kill, }, util::error::RUNTIME_SHUTTING_DOWN_ERROR, }; use libc::{syscall, SYS_pidfd_open, ENOSYS, PIDFD_NONBLOCK}; use mio::{event::Source, unix::SourceFd}; use std::{ fs::File, future::Future, io, marker::Unpin, ops::Deref, os::unix::io::{AsRawFd, FromRawFd, RawFd}, pin::Pin, process::ExitStatus, sync::atomic::{AtomicBool, Ordering::Relaxed}, task::{ready, Context, Poll}, }; #[derive(Debug)] struct Pidfd { fd: File, } impl Pidfd { fn open(pid: u32) -> Option { // Store false (0) to reduce executable size static NO_PIDFD_SUPPORT: AtomicBool = AtomicBool::new(false); if NO_PIDFD_SUPPORT.load(Relaxed) { return None; } // Safety: The following function calls invovkes syscall pidfd_open, // which takes two parameter: pidfd_open(fd: c_int, flag: c_int) let fd = unsafe { syscall(SYS_pidfd_open, pid, PIDFD_NONBLOCK) }; if fd == -1 { let errno = io::Error::last_os_error().raw_os_error().unwrap(); if errno == ENOSYS { NO_PIDFD_SUPPORT.store(true, Relaxed) } None } else { // Safety: pidfd_open returns -1 on error or a valid fd with ownership. Some(Pidfd { fd: unsafe { File::from_raw_fd(fd as i32) }, }) } } } impl AsRawFd for Pidfd { fn as_raw_fd(&self) -> RawFd { self.fd.as_raw_fd() } } impl Source for Pidfd { fn register( &mut self, registry: &mio::Registry, token: mio::Token, interest: mio::Interest, ) -> io::Result<()> { SourceFd(&self.as_raw_fd()).register(registry, token, interest) } fn reregister( &mut self, registry: &mio::Registry, token: mio::Token, interest: mio::Interest, ) -> io::Result<()> { SourceFd(&self.as_raw_fd()).reregister(registry, token, interest) } fn deregister(&mut self, registry: &mio::Registry) -> io::Result<()> { SourceFd(&self.as_raw_fd()).deregister(registry) } } #[derive(Debug)] struct PidfdReaperInner where W: Unpin, { inner: W, pidfd: PollEvented, } #[allow(deprecated)] fn is_rt_shutdown_err(err: &io::Error) -> bool { if let Some(inner) = err.get_ref() { // Using `Error::description()` is more efficient than `format!("{inner}")`, // so we use it here even if it is deprecated. err.kind() == io::ErrorKind::Other && inner.source().is_none() && inner.description() == RUNTIME_SHUTTING_DOWN_ERROR } else { false } } impl Future for PidfdReaperInner where W: Wait + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = Pin::into_inner(self); match ready!(this.pidfd.poll_read_ready(cx)) { Err(err) if is_rt_shutdown_err(&err) => { this.pidfd.reregister(Interest::READABLE)?; ready!(this.pidfd.poll_read_ready(cx))? } res => res?, } Poll::Ready(Ok(this .inner .try_wait()? .expect("pidfd is ready to read, the process should have exited"))) } } #[derive(Debug)] pub(crate) struct PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { inner: Option>, orphan_queue: Q, } impl Deref for PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { type Target = W; fn deref(&self) -> &Self::Target { &self.inner.as_ref().expect("inner has gone away").inner } } impl PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { pub(crate) fn new(inner: W, orphan_queue: Q) -> Result, W)> { if let Some(pidfd) = Pidfd::open(inner.id()) { match PollEvented::new_with_interest(pidfd, Interest::READABLE) { Ok(pidfd) => Ok(Self { inner: Some(PidfdReaperInner { pidfd, inner }), orphan_queue, }), Err(io_error) => Err((Some(io_error), inner)), } } else { Err((None, inner)) } } pub(crate) fn inner_mut(&mut self) -> &mut W { &mut self.inner.as_mut().expect("inner has gone away").inner } } impl Future for PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new( Pin::into_inner(self) .inner .as_mut() .expect("inner has gone away"), ) .poll(cx) } } impl Kill for PidfdReaper where W: Wait + Unpin + Kill, Q: OrphanQueue + Unpin, { fn kill(&mut self) -> io::Result<()> { self.inner_mut().kill() } } impl Drop for PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { fn drop(&mut self) { let mut orphan = self.inner.take().expect("inner has gone away").inner; if let Ok(Some(_)) = orphan.try_wait() { return; } self.orphan_queue.push_orphan(orphan); } } #[cfg(all(test, not(loom), not(miri)))] mod test { use super::*; use crate::{ process::unix::orphan::test::MockQueue, runtime::{Builder as RuntimeBuilder, Runtime}, }; use std::process::{Command, Output}; fn create_runtime() -> Runtime { RuntimeBuilder::new_current_thread() .enable_io() .build() .unwrap() } fn run_test(fut: impl Future) { create_runtime().block_on(fut) } fn is_pidfd_available() -> bool { let Output { stdout, status, .. } = Command::new("uname").arg("-r").output().unwrap(); assert!(status.success()); let stdout = String::from_utf8_lossy(&stdout); let mut kernel_version_iter = match stdout.split_once('-') { Some((version, _)) => version, _ => &stdout, } .split('.'); let major: u32 = kernel_version_iter.next().unwrap().parse().unwrap(); let minor: u32 = kernel_version_iter.next().unwrap().parse().unwrap(); major >= 6 || (major == 5 && minor >= 10) } #[test] fn test_pidfd_reaper_poll() { if !is_pidfd_available() { eprintln!("pidfd is not available on this linux kernel, skip this test"); return; } let queue = MockQueue::new(); run_test(async { let child = Command::new("true").spawn().unwrap(); let pidfd_reaper = PidfdReaper::new(child, &queue).unwrap(); let exit_status = pidfd_reaper.await.unwrap(); assert!(exit_status.success()); }); assert!(queue.all_enqueued.borrow().is_empty()); } #[test] fn test_pidfd_reaper_kill() { if !is_pidfd_available() { eprintln!("pidfd is not available on this linux kernel, skip this test"); return; } let queue = MockQueue::new(); run_test(async { let child = Command::new("sleep").arg("1800").spawn().unwrap(); let mut pidfd_reaper = PidfdReaper::new(child, &queue).unwrap(); pidfd_reaper.kill().unwrap(); let exit_status = pidfd_reaper.await.unwrap(); assert!(!exit_status.success()); }); assert!(queue.all_enqueued.borrow().is_empty()); } #[test] fn test_pidfd_reaper_drop() { if !is_pidfd_available() { eprintln!("pidfd is not available on this linux kernel, skip this test"); return; } let queue = MockQueue::new(); let mut child = Command::new("sleep").arg("1800").spawn().unwrap(); run_test(async { let _pidfd_reaper = PidfdReaper::new(&mut child, &queue).unwrap(); }); assert_eq!(queue.all_enqueued.borrow().len(), 1); child.kill().unwrap(); child.wait().unwrap(); } } tokio-1.43.1/src/process/unix/reap.rs000064400000000000000000000205231046102023000155660ustar 00000000000000use crate::process::imp::orphan::{OrphanQueue, Wait}; use crate::process::kill::Kill; use crate::signal::unix::InternalStream; use std::future::Future; use std::io; use std::ops::Deref; use std::pin::Pin; use std::process::ExitStatus; use std::task::Context; use std::task::Poll; /// Orchestrates between registering interest for receiving signals when a /// child process has exited, and attempting to poll for process completion. #[derive(Debug)] pub(crate) struct Reaper where W: Wait, Q: OrphanQueue, { inner: Option, orphan_queue: Q, signal: S, } impl Deref for Reaper where W: Wait, Q: OrphanQueue, { type Target = W; fn deref(&self) -> &Self::Target { self.inner() } } impl Reaper where W: Wait, Q: OrphanQueue, { pub(crate) fn new(inner: W, orphan_queue: Q, signal: S) -> Self { Self { inner: Some(inner), orphan_queue, signal, } } fn inner(&self) -> &W { self.inner.as_ref().expect("inner has gone away") } pub(crate) fn inner_mut(&mut self) -> &mut W { self.inner.as_mut().expect("inner has gone away") } } impl Future for Reaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, S: InternalStream + Unpin, { type Output = io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { loop { // If the child hasn't exited yet, then it's our responsibility to // ensure the current task gets notified when it might be able to // make progress. We can use the delivery of a SIGCHLD signal as a // sign that we can potentially make progress. // // However, we will register for a notification on the next signal // BEFORE we poll the child. Otherwise it is possible that the child // can exit and the signal can arrive after we last polled the child, // but before we've registered for a notification on the next signal // (this can cause a deadlock if there are no more spawned children // which can generate a different signal for us). A side effect of // pre-registering for signal notifications is that when the child // exits, we will have already registered for an additional // notification we don't need to consume. If another signal arrives, // this future's task will be notified/woken up again. Since the // futures model allows for spurious wake ups this extra wakeup // should not cause significant issues with parent futures. let registered_interest = self.signal.poll_recv(cx).is_pending(); if let Some(status) = self.inner_mut().try_wait()? { return Poll::Ready(Ok(status)); } // If our attempt to poll for the next signal was not ready, then // we've arranged for our task to get notified and we can bail out. if registered_interest { return Poll::Pending; } else { // Otherwise, if the signal stream delivered a signal to us, we // won't get notified at the next signal, so we'll loop and try // again. continue; } } } } impl Kill for Reaper where W: Kill + Wait, Q: OrphanQueue, { fn kill(&mut self) -> io::Result<()> { self.inner_mut().kill() } } impl Drop for Reaper where W: Wait, Q: OrphanQueue, { fn drop(&mut self) { if let Ok(Some(_)) = self.inner_mut().try_wait() { return; } let orphan = self.inner.take().unwrap(); self.orphan_queue.push_orphan(orphan); } } #[cfg(all(test, not(loom)))] mod test { use super::*; use crate::process::unix::orphan::test::MockQueue; use futures::future::FutureExt; use std::os::unix::process::ExitStatusExt; use std::process::ExitStatus; use std::task::Context; use std::task::Poll; #[derive(Debug)] struct MockWait { total_kills: usize, total_waits: usize, num_wait_until_status: usize, status: ExitStatus, } impl MockWait { fn new(status: ExitStatus, num_wait_until_status: usize) -> Self { Self { total_kills: 0, total_waits: 0, num_wait_until_status, status, } } } impl Wait for MockWait { fn id(&self) -> u32 { 0 } fn try_wait(&mut self) -> io::Result> { let ret = if self.num_wait_until_status == self.total_waits { Some(self.status) } else { None }; self.total_waits += 1; Ok(ret) } } impl Kill for MockWait { fn kill(&mut self) -> io::Result<()> { self.total_kills += 1; Ok(()) } } struct MockStream { total_polls: usize, values: Vec>, } impl MockStream { fn new(values: Vec>) -> Self { Self { total_polls: 0, values, } } } impl InternalStream for MockStream { fn poll_recv(&mut self, _cx: &mut Context<'_>) -> Poll> { self.total_polls += 1; match self.values.remove(0) { Some(()) => Poll::Ready(Some(())), None => Poll::Pending, } } } #[test] fn reaper() { let exit = ExitStatus::from_raw(0); let mock = MockWait::new(exit, 3); let mut grim = Reaper::new( mock, MockQueue::new(), MockStream::new(vec![None, Some(()), None, None, None]), ); let waker = futures::task::noop_waker(); let mut context = Context::from_waker(&waker); // Not yet exited, interest registered assert!(grim.poll_unpin(&mut context).is_pending()); assert_eq!(1, grim.signal.total_polls); assert_eq!(1, grim.total_waits); assert!(grim.orphan_queue.all_enqueued.borrow().is_empty()); // Not yet exited, couldn't register interest the first time // but managed to register interest the second time around assert!(grim.poll_unpin(&mut context).is_pending()); assert_eq!(3, grim.signal.total_polls); assert_eq!(3, grim.total_waits); assert!(grim.orphan_queue.all_enqueued.borrow().is_empty()); // Exited if let Poll::Ready(r) = grim.poll_unpin(&mut context) { assert!(r.is_ok()); let exit_code = r.unwrap(); assert_eq!(exit_code, exit); } else { unreachable!(); } assert_eq!(4, grim.signal.total_polls); assert_eq!(4, grim.total_waits); assert!(grim.orphan_queue.all_enqueued.borrow().is_empty()); } #[test] fn kill() { let exit = ExitStatus::from_raw(0); let mut grim = Reaper::new( MockWait::new(exit, 0), MockQueue::new(), MockStream::new(vec![None]), ); grim.kill().unwrap(); assert_eq!(1, grim.total_kills); assert!(grim.orphan_queue.all_enqueued.borrow().is_empty()); } #[test] fn drop_reaps_if_possible() { let exit = ExitStatus::from_raw(0); let mut mock = MockWait::new(exit, 0); { let queue = MockQueue::new(); let grim = Reaper::new(&mut mock, &queue, MockStream::new(vec![])); drop(grim); assert!(queue.all_enqueued.borrow().is_empty()); } assert_eq!(1, mock.total_waits); assert_eq!(0, mock.total_kills); } #[test] fn drop_enqueues_orphan_if_wait_fails() { let exit = ExitStatus::from_raw(0); let mut mock = MockWait::new(exit, 2); { let queue = MockQueue::<&mut MockWait>::new(); let grim = Reaper::new(&mut mock, &queue, MockStream::new(vec![])); drop(grim); assert_eq!(1, queue.all_enqueued.borrow().len()); } assert_eq!(1, mock.total_waits); assert_eq!(0, mock.total_kills); } } tokio-1.43.1/src/process/windows.rs000064400000000000000000000202771046102023000153540ustar 00000000000000//! Windows asynchronous process handling. //! //! Like with Unix we don't actually have a way of registering a process with an //! IOCP object. As a result we similarly need another mechanism for getting a //! signal when a process has exited. For now this is implemented with the //! `RegisterWaitForSingleObject` function in the kernel32.dll. //! //! This strategy is the same that libuv takes and essentially just queues up a //! wait for the process in a kernel32-specific thread pool. Once the object is //! notified (e.g. the process exits) then we have a callback that basically //! just completes a `Oneshot`. //! //! The `poll_exit` implementation will attempt to wait for the process in a //! nonblocking fashion, but failing that it'll fire off a //! `RegisterWaitForSingleObject` and then wait on the other end of the oneshot //! from then on out. use crate::io::{blocking::Blocking, AsyncRead, AsyncWrite, ReadBuf}; use crate::process::kill::Kill; use crate::process::SpawnedChild; use crate::sync::oneshot; use std::fmt; use std::fs::File as StdFile; use std::future::Future; use std::io; use std::os::windows::prelude::{AsRawHandle, IntoRawHandle, OwnedHandle, RawHandle}; use std::pin::Pin; use std::process::Stdio; use std::process::{Child as StdChild, Command as StdCommand, ExitStatus}; use std::sync::Arc; use std::task::{Context, Poll}; use windows_sys::{ Win32::Foundation::{ DuplicateHandle, BOOLEAN, DUPLICATE_SAME_ACCESS, HANDLE, INVALID_HANDLE_VALUE, }, Win32::System::Threading::{ GetCurrentProcess, RegisterWaitForSingleObject, UnregisterWaitEx, INFINITE, WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE, }, }; #[must_use = "futures do nothing unless polled"] pub(crate) struct Child { child: StdChild, waiting: Option, } impl fmt::Debug for Child { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Child") .field("pid", &self.id()) .field("child", &self.child) .field("waiting", &"..") .finish() } } struct Waiting { rx: oneshot::Receiver<()>, wait_object: HANDLE, tx: *mut Option>, } unsafe impl Sync for Waiting {} unsafe impl Send for Waiting {} pub(crate) fn spawn_child(cmd: &mut StdCommand) -> io::Result { let mut child = cmd.spawn()?; let stdin = child.stdin.take().map(stdio).transpose()?; let stdout = child.stdout.take().map(stdio).transpose()?; let stderr = child.stderr.take().map(stdio).transpose()?; Ok(SpawnedChild { child: Child { child, waiting: None, }, stdin, stdout, stderr, }) } impl Child { pub(crate) fn id(&self) -> u32 { self.child.id() } pub(crate) fn try_wait(&mut self) -> io::Result> { self.child.try_wait() } } impl Kill for Child { fn kill(&mut self) -> io::Result<()> { self.child.kill() } } impl Future for Child { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let inner = Pin::get_mut(self); loop { if let Some(ref mut w) = inner.waiting { match Pin::new(&mut w.rx).poll(cx) { Poll::Ready(Ok(())) => {} Poll::Ready(Err(_)) => panic!("should not be canceled"), Poll::Pending => return Poll::Pending, } let status = inner.try_wait()?.expect("not ready yet"); return Poll::Ready(Ok(status)); } if let Some(e) = inner.try_wait()? { return Poll::Ready(Ok(e)); } let (tx, rx) = oneshot::channel(); let ptr = Box::into_raw(Box::new(Some(tx))); let mut wait_object = 0; let rc = unsafe { RegisterWaitForSingleObject( &mut wait_object, inner.child.as_raw_handle() as _, Some(callback), ptr as *mut _, INFINITE, WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE, ) }; if rc == 0 { let err = io::Error::last_os_error(); drop(unsafe { Box::from_raw(ptr) }); return Poll::Ready(Err(err)); } inner.waiting = Some(Waiting { rx, wait_object, tx: ptr, }); } } } impl AsRawHandle for Child { fn as_raw_handle(&self) -> RawHandle { self.child.as_raw_handle() } } impl Drop for Waiting { fn drop(&mut self) { unsafe { let rc = UnregisterWaitEx(self.wait_object, INVALID_HANDLE_VALUE); if rc == 0 { panic!("failed to unregister: {}", io::Error::last_os_error()); } drop(Box::from_raw(self.tx)); } } } unsafe extern "system" fn callback(ptr: *mut std::ffi::c_void, _timer_fired: BOOLEAN) { let complete = &mut *(ptr as *mut Option>); let _ = complete.take().unwrap().send(()); } #[derive(Debug)] struct ArcFile(Arc); impl io::Read for ArcFile { fn read(&mut self, bytes: &mut [u8]) -> io::Result { (&*self.0).read(bytes) } } impl io::Write for ArcFile { fn write(&mut self, bytes: &[u8]) -> io::Result { (&*self.0).write(bytes) } fn flush(&mut self) -> io::Result<()> { (&*self.0).flush() } } #[derive(Debug)] pub(crate) struct ChildStdio { // Used for accessing the raw handle, even if the io version is busy raw: Arc, // For doing I/O operations asynchronously io: Blocking, } impl ChildStdio { pub(super) fn into_owned_handle(self) -> io::Result { convert_to_file(self).map(OwnedHandle::from) } } impl AsRawHandle for ChildStdio { fn as_raw_handle(&self) -> RawHandle { self.raw.as_raw_handle() } } impl AsyncRead for ChildStdio { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut self.io).poll_read(cx, buf) } } impl AsyncWrite for ChildStdio { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut self.io).poll_write(cx, buf) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.io).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.io).poll_shutdown(cx) } } pub(super) fn stdio(io: T) -> io::Result where T: IntoRawHandle, { use std::os::windows::prelude::FromRawHandle; let raw = Arc::new(unsafe { StdFile::from_raw_handle(io.into_raw_handle()) }); let io = ArcFile(raw.clone()); // SAFETY: the `Read` implementation of `io` does not // read from the buffer it is borrowing and correctly // reports the length of the data written into the buffer. let io = unsafe { Blocking::new(io) }; Ok(ChildStdio { raw, io }) } fn convert_to_file(child_stdio: ChildStdio) -> io::Result { let ChildStdio { raw, io } = child_stdio; drop(io); // Try to drop the Arc count here Arc::try_unwrap(raw).or_else(|raw| duplicate_handle(&*raw)) } pub(crate) fn convert_to_stdio(child_stdio: ChildStdio) -> io::Result { convert_to_file(child_stdio).map(Stdio::from) } fn duplicate_handle(io: &T) -> io::Result { use std::os::windows::prelude::FromRawHandle; unsafe { let mut dup_handle = INVALID_HANDLE_VALUE; let cur_proc = GetCurrentProcess(); let status = DuplicateHandle( cur_proc, io.as_raw_handle() as _, cur_proc, &mut dup_handle, 0, 0, DUPLICATE_SAME_ACCESS, ); if status == 0 { return Err(io::Error::last_os_error()); } Ok(StdFile::from_raw_handle(dup_handle as _)) } } tokio-1.43.1/src/runtime/blocking/mod.rs000064400000000000000000000012411046102023000162240ustar 00000000000000//! Abstracts out the APIs necessary to `Runtime` for integrating the blocking //! pool. When the `blocking` feature flag is **not** enabled, these APIs are //! shells. This isolates the complexity of dealing with conditional //! compilation. mod pool; pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner}; cfg_fs! { pub(crate) use pool::spawn_mandatory_blocking; } cfg_trace! { pub(crate) use pool::Mandatory; } mod schedule; mod shutdown; mod task; pub(crate) use task::BlockingTask; use crate::runtime::Builder; pub(crate) fn create_blocking_pool(builder: &Builder, thread_cap: usize) -> BlockingPool { BlockingPool::new(builder, thread_cap) } tokio-1.43.1/src/runtime/blocking/pool.rs000064400000000000000000000452031046102023000164240ustar 00000000000000//! Thread pool for blocking operations use crate::loom::sync::{Arc, Condvar, Mutex}; use crate::loom::thread; use crate::runtime::blocking::schedule::BlockingSchedule; use crate::runtime::blocking::{shutdown, BlockingTask}; use crate::runtime::builder::ThreadNameFn; use crate::runtime::task::{self, JoinHandle}; use crate::runtime::{Builder, Callback, Handle, BOX_FUTURE_THRESHOLD}; use crate::util::metric_atomics::MetricAtomicUsize; use crate::util::trace::{blocking_task, SpawnMeta}; use std::collections::{HashMap, VecDeque}; use std::fmt; use std::io; use std::sync::atomic::Ordering; use std::time::Duration; pub(crate) struct BlockingPool { spawner: Spawner, shutdown_rx: shutdown::Receiver, } #[derive(Clone)] pub(crate) struct Spawner { inner: Arc, } #[derive(Default)] pub(crate) struct SpawnerMetrics { num_threads: MetricAtomicUsize, num_idle_threads: MetricAtomicUsize, queue_depth: MetricAtomicUsize, } impl SpawnerMetrics { fn num_threads(&self) -> usize { self.num_threads.load(Ordering::Relaxed) } fn num_idle_threads(&self) -> usize { self.num_idle_threads.load(Ordering::Relaxed) } cfg_unstable_metrics! { fn queue_depth(&self) -> usize { self.queue_depth.load(Ordering::Relaxed) } } fn inc_num_threads(&self) { self.num_threads.increment(); } fn dec_num_threads(&self) { self.num_threads.decrement(); } fn inc_num_idle_threads(&self) { self.num_idle_threads.increment(); } fn dec_num_idle_threads(&self) -> usize { self.num_idle_threads.decrement() } fn inc_queue_depth(&self) { self.queue_depth.increment(); } fn dec_queue_depth(&self) { self.queue_depth.decrement(); } } struct Inner { /// State shared between worker threads. shared: Mutex, /// Pool threads wait on this. condvar: Condvar, /// Spawned threads use this name. thread_name: ThreadNameFn, /// Spawned thread stack size. stack_size: Option, /// Call after a thread starts. after_start: Option, /// Call before a thread stops. before_stop: Option, // Maximum number of threads. thread_cap: usize, // Customizable wait timeout. keep_alive: Duration, // Metrics about the pool. metrics: SpawnerMetrics, } struct Shared { queue: VecDeque, num_notify: u32, shutdown: bool, shutdown_tx: Option, /// Prior to shutdown, we clean up `JoinHandles` by having each timed-out /// thread join on the previous timed-out thread. This is not strictly /// necessary but helps avoid Valgrind false positives, see /// /// for more information. last_exiting_thread: Option>, /// This holds the `JoinHandles` for all running threads; on shutdown, the thread /// calling shutdown handles joining on these. worker_threads: HashMap>, /// This is a counter used to iterate `worker_threads` in a consistent order (for loom's /// benefit). worker_thread_index: usize, } pub(crate) struct Task { task: task::UnownedTask, mandatory: Mandatory, } #[derive(PartialEq, Eq)] pub(crate) enum Mandatory { #[cfg_attr(not(fs), allow(dead_code))] Mandatory, NonMandatory, } pub(crate) enum SpawnError { /// Pool is shutting down and the task was not scheduled ShuttingDown, /// There are no worker threads available to take the task /// and the OS failed to spawn a new one NoThreads(io::Error), } impl From for io::Error { fn from(e: SpawnError) -> Self { match e { SpawnError::ShuttingDown => { io::Error::new(io::ErrorKind::Other, "blocking pool shutting down") } SpawnError::NoThreads(e) => e, } } } impl Task { pub(crate) fn new(task: task::UnownedTask, mandatory: Mandatory) -> Task { Task { task, mandatory } } fn run(self) { self.task.run(); } fn shutdown_or_run_if_mandatory(self) { match self.mandatory { Mandatory::NonMandatory => self.task.shutdown(), Mandatory::Mandatory => self.task.run(), } } } const KEEP_ALIVE: Duration = Duration::from_secs(10); /// Runs the provided function on an executor dedicated to blocking operations. /// Tasks will be scheduled as non-mandatory, meaning they may not get executed /// in case of runtime shutdown. #[track_caller] #[cfg_attr(target_os = "wasi", allow(dead_code))] pub(crate) fn spawn_blocking(func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let rt = Handle::current(); rt.spawn_blocking(func) } cfg_fs! { #[cfg_attr(any( all(loom, not(test)), // the function is covered by loom tests test ), allow(dead_code))] /// Runs the provided function on an executor dedicated to blocking /// operations. Tasks will be scheduled as mandatory, meaning they are /// guaranteed to run unless a shutdown is already taking place. In case a /// shutdown is already taking place, `None` will be returned. pub(crate) fn spawn_mandatory_blocking(func: F) -> Option> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let rt = Handle::current(); rt.inner.blocking_spawner().spawn_mandatory_blocking(&rt, func) } } // ===== impl BlockingPool ===== impl BlockingPool { pub(crate) fn new(builder: &Builder, thread_cap: usize) -> BlockingPool { let (shutdown_tx, shutdown_rx) = shutdown::channel(); let keep_alive = builder.keep_alive.unwrap_or(KEEP_ALIVE); BlockingPool { spawner: Spawner { inner: Arc::new(Inner { shared: Mutex::new(Shared { queue: VecDeque::new(), num_notify: 0, shutdown: false, shutdown_tx: Some(shutdown_tx), last_exiting_thread: None, worker_threads: HashMap::new(), worker_thread_index: 0, }), condvar: Condvar::new(), thread_name: builder.thread_name.clone(), stack_size: builder.thread_stack_size, after_start: builder.after_start.clone(), before_stop: builder.before_stop.clone(), thread_cap, keep_alive, metrics: SpawnerMetrics::default(), }), }, shutdown_rx, } } pub(crate) fn spawner(&self) -> &Spawner { &self.spawner } pub(crate) fn shutdown(&mut self, timeout: Option) { let mut shared = self.spawner.inner.shared.lock(); // The function can be called multiple times. First, by explicitly // calling `shutdown` then by the drop handler calling `shutdown`. This // prevents shutting down twice. if shared.shutdown { return; } shared.shutdown = true; shared.shutdown_tx = None; self.spawner.inner.condvar.notify_all(); let last_exited_thread = std::mem::take(&mut shared.last_exiting_thread); let workers = std::mem::take(&mut shared.worker_threads); drop(shared); if self.shutdown_rx.wait(timeout) { let _ = last_exited_thread.map(thread::JoinHandle::join); // Loom requires that execution be deterministic, so sort by thread ID before joining. // (HashMaps use a randomly-seeded hash function, so the order is nondeterministic) #[cfg(loom)] let workers: Vec<(usize, thread::JoinHandle<()>)> = { let mut workers: Vec<_> = workers.into_iter().collect(); workers.sort_by_key(|(id, _)| *id); workers }; for (_id, handle) in workers { let _ = handle.join(); } } } } impl Drop for BlockingPool { fn drop(&mut self) { self.shutdown(None); } } impl fmt::Debug for BlockingPool { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("BlockingPool").finish() } } // ===== impl Spawner ===== impl Spawner { #[track_caller] pub(crate) fn spawn_blocking(&self, rt: &Handle, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let fn_size = std::mem::size_of::(); let (join_handle, spawn_result) = if fn_size > BOX_FUTURE_THRESHOLD { self.spawn_blocking_inner( Box::new(func), Mandatory::NonMandatory, SpawnMeta::new_unnamed(fn_size), rt, ) } else { self.spawn_blocking_inner( func, Mandatory::NonMandatory, SpawnMeta::new_unnamed(fn_size), rt, ) }; match spawn_result { Ok(()) => join_handle, // Compat: do not panic here, return the join_handle even though it will never resolve Err(SpawnError::ShuttingDown) => join_handle, Err(SpawnError::NoThreads(e)) => { panic!("OS can't spawn worker thread: {e}") } } } cfg_fs! { #[track_caller] #[cfg_attr(any( all(loom, not(test)), // the function is covered by loom tests test ), allow(dead_code))] pub(crate) fn spawn_mandatory_blocking(&self, rt: &Handle, func: F) -> Option> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let fn_size = std::mem::size_of::(); let (join_handle, spawn_result) = if fn_size > BOX_FUTURE_THRESHOLD { self.spawn_blocking_inner( Box::new(func), Mandatory::Mandatory, SpawnMeta::new_unnamed(fn_size), rt, ) } else { self.spawn_blocking_inner( func, Mandatory::Mandatory, SpawnMeta::new_unnamed(fn_size), rt, ) }; if spawn_result.is_ok() { Some(join_handle) } else { None } } } #[track_caller] pub(crate) fn spawn_blocking_inner( &self, func: F, is_mandatory: Mandatory, spawn_meta: SpawnMeta<'_>, rt: &Handle, ) -> (JoinHandle, Result<(), SpawnError>) where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let id = task::Id::next(); let fut = blocking_task::>(BlockingTask::new(func), spawn_meta, id.as_u64()); let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id); let spawned = self.spawn_task(Task::new(task, is_mandatory), rt); (handle, spawned) } fn spawn_task(&self, task: Task, rt: &Handle) -> Result<(), SpawnError> { let mut shared = self.inner.shared.lock(); if shared.shutdown { // Shutdown the task: it's fine to shutdown this task (even if // mandatory) because it was scheduled after the shutdown of the // runtime began. task.task.shutdown(); // no need to even push this task; it would never get picked up return Err(SpawnError::ShuttingDown); } shared.queue.push_back(task); self.inner.metrics.inc_queue_depth(); if self.inner.metrics.num_idle_threads() == 0 { // No threads are able to process the task. if self.inner.metrics.num_threads() == self.inner.thread_cap { // At max number of threads } else { assert!(shared.shutdown_tx.is_some()); let shutdown_tx = shared.shutdown_tx.clone(); if let Some(shutdown_tx) = shutdown_tx { let id = shared.worker_thread_index; match self.spawn_thread(shutdown_tx, rt, id) { Ok(handle) => { self.inner.metrics.inc_num_threads(); shared.worker_thread_index += 1; shared.worker_threads.insert(id, handle); } Err(ref e) if is_temporary_os_thread_error(e) && self.inner.metrics.num_threads() > 0 => { // OS temporarily failed to spawn a new thread. // The task will be picked up eventually by a currently // busy thread. } Err(e) => { // The OS refused to spawn the thread and there is no thread // to pick up the task that has just been pushed to the queue. return Err(SpawnError::NoThreads(e)); } } } } } else { // Notify an idle worker thread. The notification counter // is used to count the needed amount of notifications // exactly. Thread libraries may generate spurious // wakeups, this counter is used to keep us in a // consistent state. self.inner.metrics.dec_num_idle_threads(); shared.num_notify += 1; self.inner.condvar.notify_one(); } Ok(()) } fn spawn_thread( &self, shutdown_tx: shutdown::Sender, rt: &Handle, id: usize, ) -> io::Result> { 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); } let rt = rt.clone(); builder.spawn(move || { // Only the reference should be moved into the closure let _enter = rt.enter(); rt.inner.blocking_spawner().inner.run(id); drop(shutdown_tx); }) } } cfg_unstable_metrics! { impl Spawner { pub(crate) fn num_threads(&self) -> usize { self.inner.metrics.num_threads() } pub(crate) fn num_idle_threads(&self) -> usize { self.inner.metrics.num_idle_threads() } pub(crate) fn queue_depth(&self) -> usize { self.inner.metrics.queue_depth() } } } // Tells whether the error when spawning a thread is temporary. #[inline] fn is_temporary_os_thread_error(error: &io::Error) -> bool { matches!(error.kind(), io::ErrorKind::WouldBlock) } impl Inner { fn run(&self, worker_thread_id: usize) { if let Some(f) = &self.after_start { f(); } let mut shared = self.shared.lock(); let mut join_on_thread = None; 'main: loop { // BUSY while let Some(task) = shared.queue.pop_front() { self.metrics.dec_queue_depth(); drop(shared); task.run(); shared = self.shared.lock(); } // IDLE self.metrics.inc_num_idle_threads(); while !shared.shutdown { let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap(); shared = lock_result.0; let timeout_result = lock_result.1; if shared.num_notify != 0 { // We have received a legitimate wakeup, // acknowledge it by decrementing the counter // and transition to the BUSY state. shared.num_notify -= 1; break; } // Even if the condvar "timed out", if the pool is entering the // shutdown phase, we want to perform the cleanup logic. if !shared.shutdown && timeout_result.timed_out() { // We'll join the prior timed-out thread's JoinHandle after dropping the lock. // This isn't done when shutting down, because the thread calling shutdown will // handle joining everything. let my_handle = shared.worker_threads.remove(&worker_thread_id); join_on_thread = std::mem::replace(&mut shared.last_exiting_thread, my_handle); break 'main; } // Spurious wakeup detected, go back to sleep. } if shared.shutdown { // Drain the queue while let Some(task) = shared.queue.pop_front() { self.metrics.dec_queue_depth(); drop(shared); task.shutdown_or_run_if_mandatory(); shared = self.shared.lock(); } // Work was produced, and we "took" it (by decrementing num_notify). // This means that num_idle was decremented once for our wakeup. // But, since we are exiting, we need to "undo" that, as we'll stay idle. self.metrics.inc_num_idle_threads(); // NOTE: Technically we should also do num_notify++ and notify again, // but since we're shutting down anyway, that won't be necessary. break; } } // Thread exit self.metrics.dec_num_threads(); // num_idle should now be tracked exactly, panic // with a descriptive message if it is not the // case. let prev_idle = self.metrics.dec_num_idle_threads(); assert!( prev_idle >= self.metrics.num_idle_threads(), "num_idle_threads underflowed on thread exit" ); if shared.shutdown && self.metrics.num_threads() == 0 { self.condvar.notify_one(); } drop(shared); if let Some(f) = &self.before_stop { f(); } if let Some(handle) = join_on_thread { let _ = handle.join(); } } } impl fmt::Debug for Spawner { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("blocking::Spawner").finish() } } tokio-1.43.1/src/runtime/blocking/schedule.rs000064400000000000000000000046261046102023000172530ustar 00000000000000#[cfg(feature = "test-util")] use crate::runtime::scheduler; use crate::runtime::task::{self, Task, TaskHarnessScheduleHooks}; use crate::runtime::Handle; /// `task::Schedule` implementation that does nothing (except some bookkeeping /// in test-util builds). This is unique to the blocking scheduler as tasks /// scheduled are not really futures but blocking operations. /// /// We avoid storing the task by forgetting it in `bind` and re-materializing it /// in `release`. pub(crate) struct BlockingSchedule { #[cfg(feature = "test-util")] handle: Handle, hooks: TaskHarnessScheduleHooks, } impl BlockingSchedule { #[cfg_attr(not(feature = "test-util"), allow(unused_variables))] pub(crate) fn new(handle: &Handle) -> Self { #[cfg(feature = "test-util")] { match &handle.inner { scheduler::Handle::CurrentThread(handle) => { handle.driver.clock.inhibit_auto_advance(); } #[cfg(feature = "rt-multi-thread")] scheduler::Handle::MultiThread(_) => {} #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] scheduler::Handle::MultiThreadAlt(_) => {} } } BlockingSchedule { #[cfg(feature = "test-util")] handle: handle.clone(), hooks: TaskHarnessScheduleHooks { task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(), }, } } } impl task::Schedule for BlockingSchedule { fn release(&self, _task: &Task) -> Option> { #[cfg(feature = "test-util")] { match &self.handle.inner { scheduler::Handle::CurrentThread(handle) => { handle.driver.clock.allow_auto_advance(); handle.driver.unpark(); } #[cfg(feature = "rt-multi-thread")] scheduler::Handle::MultiThread(_) => {} #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] scheduler::Handle::MultiThreadAlt(_) => {} } } None } fn schedule(&self, _task: task::Notified) { unreachable!(); } fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: self.hooks.task_terminate_callback.clone(), } } } tokio-1.43.1/src/runtime/blocking/shutdown.rs000064400000000000000000000041711046102023000173250ustar 00000000000000//! A shutdown channel. //! //! Each worker holds the `Sender` half. When all the `Sender` halves are //! dropped, the `Receiver` receives a notification. use crate::loom::sync::Arc; use crate::sync::oneshot; use std::time::Duration; #[derive(Debug, Clone)] pub(super) struct Sender { _tx: Arc>, } #[derive(Debug)] pub(super) struct Receiver { rx: oneshot::Receiver<()>, } pub(super) fn channel() -> (Sender, Receiver) { let (tx, rx) = oneshot::channel(); let tx = Sender { _tx: Arc::new(tx) }; let rx = Receiver { rx }; (tx, rx) } impl Receiver { /// Blocks the current thread until all `Sender` handles drop. /// /// If `timeout` is `Some`, the thread is blocked for **at most** `timeout` /// duration. If `timeout` is `None`, then the thread is blocked until the /// shutdown signal is received. /// /// If the timeout has elapsed, it returns `false`, otherwise it returns `true`. pub(crate) fn wait(&mut self, timeout: Option) -> bool { use crate::runtime::context::try_enter_blocking_region; if timeout == Some(Duration::from_nanos(0)) { return false; } let mut e = match try_enter_blocking_region() { Some(enter) => enter, _ => { if std::thread::panicking() { // Don't panic in a panic return false; } else { panic!( "Cannot drop a runtime in a context where blocking is not allowed. \ This happens when a runtime is dropped from within an asynchronous context." ); } } }; // The oneshot completes with an Err // // If blocking fails to wait, this indicates a problem parking the // current thread (usually, shutting down a runtime stored in a // thread-local). if let Some(timeout) = timeout { e.block_on_timeout(&mut self.rx, timeout).is_ok() } else { let _ = e.block_on(&mut self.rx); true } } } tokio-1.43.1/src/runtime/blocking/task.rs000064400000000000000000000025531046102023000164160ustar 00000000000000use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; /// Converts a function to a future that completes on poll. pub(crate) struct BlockingTask { func: Option, } impl BlockingTask { /// Initializes a new blocking task from the given function. pub(crate) fn new(func: T) -> BlockingTask { BlockingTask { func: Some(func) } } } // The closure `F` is never pinned impl Unpin for BlockingTask {} impl Future for BlockingTask where T: FnOnce() -> R + Send + 'static, R: Send + 'static, { type Output = R; fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { let me = &mut *self; let func = me .func .take() .expect("[internal exception] blocking task ran twice."); // This is a little subtle: // For convenience, we'd like _every_ call tokio ever makes to Task::poll() to be budgeted // using coop. However, the way things are currently modeled, even running a blocking task // currently goes through Task::poll(), and so is subject to budgeting. That isn't really // what we want; a blocking task may itself want to run tasks (it might be a Worker!), so // we want it to start without any budgeting. crate::runtime::coop::stop(); Poll::Ready(func()) } } tokio-1.43.1/src/runtime/builder.rs000064400000000000000000001650231046102023000153140ustar 00000000000000#![cfg_attr(loom, allow(unused_imports))] use crate::runtime::handle::Handle; use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime, TaskCallback}; #[cfg(tokio_unstable)] use crate::runtime::{metrics::HistogramConfiguration, LocalOptions, LocalRuntime, TaskMeta}; use crate::util::rand::{RngSeed, RngSeedGenerator}; use crate::runtime::blocking::BlockingPool; use crate::runtime::scheduler::CurrentThread; use std::fmt; use std::io; use std::thread::ThreadId; use std::time::Duration; /// Builds Tokio Runtime with custom configuration values. /// /// Methods can be chained in order to set the configuration values. The /// Runtime is constructed by calling [`build`]. /// /// New instances of `Builder` are obtained via [`Builder::new_multi_thread`] /// or [`Builder::new_current_thread`]. /// /// See function level documentation for details on the various configuration /// settings. /// /// [`build`]: method@Self::build /// [`Builder::new_multi_thread`]: method@Self::new_multi_thread /// [`Builder::new_current_thread`]: method@Self::new_current_thread /// /// # Examples /// /// ``` /// use tokio::runtime::Builder; /// /// fn main() { /// // build runtime /// let runtime = Builder::new_multi_thread() /// .worker_threads(4) /// .thread_name("my-custom-name") /// .thread_stack_size(3 * 1024 * 1024) /// .build() /// .unwrap(); /// /// // use runtime ... /// } /// ``` pub struct Builder { /// Runtime type kind: Kind, /// Whether or not to enable the I/O driver enable_io: bool, nevents: usize, /// Whether or not to enable the time driver enable_time: bool, /// Whether or not the clock should start paused. start_paused: bool, /// The number of worker threads, used by Runtime. /// /// Only used when not using the current-thread executor. worker_threads: Option, /// Cap on thread usage. max_blocking_threads: usize, /// 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, /// Callback to run after each thread starts. pub(super) after_start: Option, /// To run before each worker thread stops pub(super) before_stop: Option, /// To run before each worker thread is parked. pub(super) before_park: Option, /// To run after each thread is unparked. pub(super) after_unpark: Option, /// To run before each task is spawned. pub(super) before_spawn: Option, /// To run after each task is terminated. pub(super) after_termination: Option, /// Customizable keep alive timeout for `BlockingPool` pub(super) keep_alive: Option, /// How many ticks before pulling a task from the global/remote queue? /// /// When `None`, the value is unspecified and behavior details are left to /// the scheduler. Each scheduler flavor could choose to either pick its own /// default value or use some other strategy to decide when to poll from the /// global queue. For example, the multi-threaded scheduler uses a /// self-tuning strategy based on mean task poll times. pub(super) global_queue_interval: Option, /// How many ticks before yielding to the driver for timer and I/O events? pub(super) event_interval: u32, pub(super) local_queue_capacity: usize, /// When true, the multi-threade scheduler LIFO slot should not be used. /// /// This option should only be exposed as unstable. pub(super) disable_lifo_slot: bool, /// Specify a random number generator seed to provide deterministic results pub(super) seed_generator: RngSeedGenerator, /// When true, enables task poll count histogram instrumentation. pub(super) metrics_poll_count_histogram_enable: bool, /// Configures the task poll count histogram pub(super) metrics_poll_count_histogram: HistogramBuilder, #[cfg(tokio_unstable)] pub(super) unhandled_panic: UnhandledPanic, } cfg_unstable! { /// How the runtime should respond to unhandled panics. /// /// Instances of `UnhandledPanic` are passed to `Builder::unhandled_panic` /// to configure the runtime behavior when a spawned task panics. /// /// See [`Builder::unhandled_panic`] for more details. #[derive(Debug, Clone)] #[non_exhaustive] pub enum UnhandledPanic { /// The runtime should ignore panics on spawned tasks. /// /// The panic is forwarded to the task's [`JoinHandle`] and all spawned /// tasks continue running normally. /// /// This is the default behavior. /// /// # Examples /// /// ``` /// use tokio::runtime::{self, UnhandledPanic}; /// /// # pub fn main() { /// let rt = runtime::Builder::new_current_thread() /// .unhandled_panic(UnhandledPanic::Ignore) /// .build() /// .unwrap(); /// /// let task1 = rt.spawn(async { panic!("boom"); }); /// let task2 = rt.spawn(async { /// // This task completes normally /// "done" /// }); /// /// rt.block_on(async { /// // The panic on the first task is forwarded to the `JoinHandle` /// assert!(task1.await.is_err()); /// /// // The second task completes normally /// assert!(task2.await.is_ok()); /// }) /// # } /// ``` /// /// [`JoinHandle`]: struct@crate::task::JoinHandle Ignore, /// The runtime should immediately shutdown if a spawned task panics. /// /// The runtime will immediately shutdown even if the panicked task's /// [`JoinHandle`] is still available. All further spawned tasks will be /// immediately dropped and call to [`Runtime::block_on`] will panic. /// /// # Examples /// /// ```should_panic /// use tokio::runtime::{self, UnhandledPanic}; /// /// # pub fn main() { /// let rt = runtime::Builder::new_current_thread() /// .unhandled_panic(UnhandledPanic::ShutdownRuntime) /// .build() /// .unwrap(); /// /// rt.spawn(async { panic!("boom"); }); /// rt.spawn(async { /// // This task never completes. /// }); /// /// rt.block_on(async { /// // Do some work /// # loop { tokio::task::yield_now().await; } /// }) /// # } /// ``` /// /// [`JoinHandle`]: struct@crate::task::JoinHandle ShutdownRuntime, } } pub(crate) type ThreadNameFn = std::sync::Arc String + Send + Sync + 'static>; #[derive(Clone, Copy)] pub(crate) enum Kind { CurrentThread, #[cfg(feature = "rt-multi-thread")] MultiThread, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] MultiThreadAlt, } impl Builder { /// Returns a new builder with the current thread scheduler selected. /// /// Configuration methods can be chained on the return value. /// /// To spawn non-`Send` tasks on the resulting runtime, combine it with a /// [`LocalSet`]. /// /// [`LocalSet`]: crate::task::LocalSet pub fn new_current_thread() -> Builder { #[cfg(loom)] const EVENT_INTERVAL: u32 = 4; // The number `61` is fairly arbitrary. I believe this value was copied from golang. #[cfg(not(loom))] const EVENT_INTERVAL: u32 = 61; Builder::new(Kind::CurrentThread, EVENT_INTERVAL) } /// Returns a new builder with the multi thread scheduler selected. /// /// Configuration methods can be chained on the return value. #[cfg(feature = "rt-multi-thread")] #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] pub fn new_multi_thread() -> Builder { // The number `61` is fairly arbitrary. I believe this value was copied from golang. Builder::new(Kind::MultiThread, 61) } cfg_unstable! { /// Returns a new builder with the alternate multi thread scheduler /// selected. /// /// The alternate multi threaded scheduler is an in-progress /// candidate to replace the existing multi threaded scheduler. It /// currently does not scale as well to 16+ processors. /// /// This runtime flavor is currently **not considered production /// ready**. /// /// Configuration methods can be chained on the return value. #[cfg(feature = "rt-multi-thread")] #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] pub fn new_multi_thread_alt() -> Builder { // The number `61` is fairly arbitrary. I believe this value was copied from golang. Builder::new(Kind::MultiThreadAlt, 61) } } /// Returns a new runtime builder initialized with default configuration /// values. /// /// Configuration methods can be chained on the return value. pub(crate) fn new(kind: Kind, event_interval: u32) -> Builder { Builder { kind, // I/O defaults to "off" enable_io: false, nevents: 1024, // Time defaults to "off" enable_time: false, // The clock starts not-paused start_paused: false, // Read from environment variable first in multi-threaded mode. // Default to lazy auto-detection (one thread per CPU core) worker_threads: None, max_blocking_threads: 512, // Default thread name thread_name: std::sync::Arc::new(|| "tokio-runtime-worker".into()), // Do not set a stack size by default thread_stack_size: None, // No worker thread callbacks after_start: None, before_stop: None, before_park: None, after_unpark: None, before_spawn: None, after_termination: None, keep_alive: None, // Defaults for these values depend on the scheduler kind, so we get them // as parameters. global_queue_interval: None, event_interval, #[cfg(not(loom))] local_queue_capacity: 256, #[cfg(loom)] local_queue_capacity: 4, seed_generator: RngSeedGenerator::new(RngSeed::new()), #[cfg(tokio_unstable)] unhandled_panic: UnhandledPanic::Ignore, metrics_poll_count_histogram_enable: false, metrics_poll_count_histogram: HistogramBuilder::default(), disable_lifo_slot: false, } } /// Enables both I/O and time drivers. /// /// Doing this is a shorthand for calling `enable_io` and `enable_time` /// individually. If additional components are added to Tokio in the future, /// `enable_all` will include these future components. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_all() /// .build() /// .unwrap(); /// ``` pub fn enable_all(&mut self) -> &mut Self { #[cfg(any( feature = "net", all(unix, feature = "process"), all(unix, feature = "signal") ))] self.enable_io(); #[cfg(feature = "time")] self.enable_time(); self } /// Sets the number of worker threads the `Runtime` will use. /// /// This can be any number above 0 though it is advised to keep this value /// on the smaller side. /// /// This will override the value read from environment variable `TOKIO_WORKER_THREADS`. /// /// # Default /// /// The default value is the number of cores available to the system. /// /// When using the `current_thread` runtime this method has no effect. /// /// # Examples /// /// ## Multi threaded runtime with 4 threads /// /// ``` /// use tokio::runtime; /// /// // This will spawn a work-stealing runtime with 4 worker threads. /// let rt = runtime::Builder::new_multi_thread() /// .worker_threads(4) /// .build() /// .unwrap(); /// /// rt.spawn(async move {}); /// ``` /// /// ## Current thread runtime (will only run on the current thread via `Runtime::block_on`) /// /// ``` /// use tokio::runtime; /// /// // Create a runtime that _must_ be driven from a call /// // to `Runtime::block_on`. /// let rt = runtime::Builder::new_current_thread() /// .build() /// .unwrap(); /// /// // This will run the runtime and future on the current thread /// rt.block_on(async move {}); /// ``` /// /// # Panics /// /// This will panic if `val` is not larger than `0`. #[track_caller] pub fn worker_threads(&mut self, val: usize) -> &mut Self { assert!(val > 0, "Worker threads cannot be set to 0"); self.worker_threads = Some(val); self } /// Specifies the limit for additional threads spawned by the Runtime. /// /// These threads are used for blocking operations like tasks spawned /// through [`spawn_blocking`], this includes but is not limited to: /// - [`fs`] operations /// - dns resolution through [`ToSocketAddrs`] /// - writing to [`Stdout`] or [`Stderr`] /// - reading from [`Stdin`] /// /// Unlike the [`worker_threads`], they are not always active and will exit /// if left idle for too long. You can change this timeout duration with [`thread_keep_alive`]. /// /// It's recommended to not set this limit too low in order to avoid hanging on operations /// requiring [`spawn_blocking`]. /// /// The default value is 512. /// /// # Panics /// /// This will panic if `val` is not larger than `0`. /// /// # Upgrading from 0.x /// /// In old versions `max_threads` limited both blocking and worker threads, but the /// current `max_blocking_threads` does not include async worker threads in the count. /// /// [`spawn_blocking`]: fn@crate::task::spawn_blocking /// [`fs`]: mod@crate::fs /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs /// [`Stdout`]: struct@crate::io::Stdout /// [`Stdin`]: struct@crate::io::Stdin /// [`Stderr`]: struct@crate::io::Stderr /// [`worker_threads`]: Self::worker_threads /// [`thread_keep_alive`]: Self::thread_keep_alive #[track_caller] #[cfg_attr(docsrs, doc(alias = "max_threads"))] pub fn max_blocking_threads(&mut self, val: usize) -> &mut Self { assert!(val > 0, "Max blocking threads cannot be set to 0"); self.max_blocking_threads = val; self } /// Sets name of threads spawned by the `Runtime`'s thread pool. /// /// The default name is "tokio-runtime-worker". /// /// # Examples /// /// ``` /// # use tokio::runtime; /// /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .thread_name("my-pool") /// .build(); /// # } /// ``` pub fn thread_name(&mut self, val: impl Into) -> &mut Self { let val = val.into(); self.thread_name = std::sync::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_multi_thread() /// .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(&mut self, f: F) -> &mut Self where F: Fn() -> String + Send + Sync + 'static, { self.thread_name = std::sync::Arc::new(f); self } /// Sets the stack size (in bytes) for worker threads. /// /// The actual stack size may be greater than this value if the platform /// specifies minimal stack size. /// /// The default stack size for spawned threads is 2 MiB, though this /// particular stack size is subject to change in the future. /// /// # Examples /// /// ``` /// # use tokio::runtime; /// /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .thread_stack_size(32 * 1024) /// .build(); /// # } /// ``` pub fn thread_stack_size(&mut self, val: usize) -> &mut Self { self.thread_stack_size = Some(val); self } /// Executes function `f` after each thread is started but before it starts /// doing work. /// /// This is intended for bookkeeping and monitoring use cases. /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let runtime = runtime::Builder::new_multi_thread() /// .on_thread_start(|| { /// println!("thread started"); /// }) /// .build(); /// # } /// ``` #[cfg(not(loom))] pub fn on_thread_start(&mut self, f: F) -> &mut Self where F: Fn() + Send + Sync + 'static, { self.after_start = Some(std::sync::Arc::new(f)); self } /// Executes function `f` before each thread stops. /// /// This is intended for bookkeeping and monitoring use cases. /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let runtime = runtime::Builder::new_multi_thread() /// .on_thread_stop(|| { /// println!("thread stopping"); /// }) /// .build(); /// # } /// ``` #[cfg(not(loom))] pub fn on_thread_stop(&mut self, f: F) -> &mut Self where F: Fn() + Send + Sync + 'static, { self.before_stop = Some(std::sync::Arc::new(f)); self } /// Executes function `f` just before a thread is parked (goes idle). /// `f` is called within the Tokio context, so functions like [`tokio::spawn`](crate::spawn) /// can be called, and may result in this thread being unparked immediately. /// /// This can be used to start work only when the executor is idle, or for bookkeeping /// and monitoring purposes. /// /// Note: There can only be one park callback for a runtime; calling this function /// more than once replaces the last callback defined, rather than adding to it. /// /// # Examples /// /// ## Multithreaded executor /// ``` /// # use std::sync::Arc; /// # use std::sync::atomic::{AtomicBool, Ordering}; /// # use tokio::runtime; /// # use tokio::sync::Barrier; /// # pub fn main() { /// let once = AtomicBool::new(true); /// let barrier = Arc::new(Barrier::new(2)); /// /// let runtime = runtime::Builder::new_multi_thread() /// .worker_threads(1) /// .on_thread_park({ /// let barrier = barrier.clone(); /// move || { /// let barrier = barrier.clone(); /// if once.swap(false, Ordering::Relaxed) { /// tokio::spawn(async move { barrier.wait().await; }); /// } /// } /// }) /// .build() /// .unwrap(); /// /// runtime.block_on(async { /// barrier.wait().await; /// }) /// # } /// ``` /// ## Current thread executor /// ``` /// # use std::sync::Arc; /// # use std::sync::atomic::{AtomicBool, Ordering}; /// # use tokio::runtime; /// # use tokio::sync::Barrier; /// # pub fn main() { /// let once = AtomicBool::new(true); /// let barrier = Arc::new(Barrier::new(2)); /// /// let runtime = runtime::Builder::new_current_thread() /// .on_thread_park({ /// let barrier = barrier.clone(); /// move || { /// let barrier = barrier.clone(); /// if once.swap(false, Ordering::Relaxed) { /// tokio::spawn(async move { barrier.wait().await; }); /// } /// } /// }) /// .build() /// .unwrap(); /// /// runtime.block_on(async { /// barrier.wait().await; /// }) /// # } /// ``` #[cfg(not(loom))] pub fn on_thread_park(&mut self, f: F) -> &mut Self where F: Fn() + Send + Sync + 'static, { self.before_park = Some(std::sync::Arc::new(f)); self } /// Executes function `f` just after a thread unparks (starts executing tasks). /// /// This is intended for bookkeeping and monitoring use cases; note that work /// in this callback will increase latencies when the application has allowed one or /// more runtime threads to go idle. /// /// Note: There can only be one unpark callback for a runtime; calling this function /// more than once replaces the last callback defined, rather than adding to it. /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let runtime = runtime::Builder::new_multi_thread() /// .on_thread_unpark(|| { /// println!("thread unparking"); /// }) /// .build(); /// /// runtime.unwrap().block_on(async { /// tokio::task::yield_now().await; /// println!("Hello from Tokio!"); /// }) /// # } /// ``` #[cfg(not(loom))] pub fn on_thread_unpark(&mut self, f: F) -> &mut Self where F: Fn() + Send + Sync + 'static, { self.after_unpark = Some(std::sync::Arc::new(f)); self } /// Executes function `f` just before a task is spawned. /// /// `f` is called within the Tokio context, so functions like /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being /// invoked immediately. /// /// This can be used for bookkeeping or monitoring purposes. /// /// Note: There can only be one spawn callback for a runtime; calling this function more /// than once replaces the last callback defined, rather than adding to it. /// /// This *does not* support [`LocalSet`](crate::task::LocalSet) at this time. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let runtime = runtime::Builder::new_current_thread() /// .on_task_spawn(|_| { /// println!("spawning task"); /// }) /// .build() /// .unwrap(); /// /// runtime.block_on(async { /// tokio::task::spawn(std::future::ready(())); /// /// for _ in 0..64 { /// tokio::task::yield_now().await; /// } /// }) /// # } /// ``` #[cfg(all(not(loom), tokio_unstable))] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn on_task_spawn(&mut self, f: F) -> &mut Self where F: Fn(&TaskMeta<'_>) + Send + Sync + 'static, { self.before_spawn = Some(std::sync::Arc::new(f)); self } /// Executes function `f` just after a task is terminated. /// /// `f` is called within the Tokio context, so functions like /// [`tokio::spawn`](crate::spawn) can be called. /// /// This can be used for bookkeeping or monitoring purposes. /// /// Note: There can only be one task termination callback for a runtime; calling this /// function more than once replaces the last callback defined, rather than adding to it. /// /// This *does not* support [`LocalSet`](crate::task::LocalSet) at this time. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let runtime = runtime::Builder::new_current_thread() /// .on_task_terminate(|_| { /// println!("killing task"); /// }) /// .build() /// .unwrap(); /// /// runtime.block_on(async { /// tokio::task::spawn(std::future::ready(())); /// /// for _ in 0..64 { /// tokio::task::yield_now().await; /// } /// }) /// # } /// ``` #[cfg(all(not(loom), tokio_unstable))] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn on_task_terminate(&mut self, f: F) -> &mut Self where F: Fn(&TaskMeta<'_>) + Send + Sync + 'static, { self.after_termination = Some(std::sync::Arc::new(f)); self } /// Creates the configured `Runtime`. /// /// The returned `Runtime` instance is ready to spawn tasks. /// /// # Examples /// /// ``` /// use tokio::runtime::Builder; /// /// let rt = Builder::new_multi_thread().build().unwrap(); /// /// rt.block_on(async { /// println!("Hello from the Tokio runtime"); /// }); /// ``` pub fn build(&mut self) -> io::Result { match &self.kind { Kind::CurrentThread => self.build_current_thread_runtime(), #[cfg(feature = "rt-multi-thread")] Kind::MultiThread => self.build_threaded_runtime(), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Kind::MultiThreadAlt => self.build_alt_threaded_runtime(), } } /// Creates the configured `LocalRuntime`. /// /// The returned `LocalRuntime` instance is ready to spawn tasks. /// /// # Panics /// This will panic if `current_thread` is not the selected runtime flavor. /// All other runtime flavors are unsupported by [`LocalRuntime`]. /// /// [`LocalRuntime`]: [crate::runtime::LocalRuntime] /// /// # Examples /// /// ``` /// use tokio::runtime::Builder; /// /// let rt = Builder::new_current_thread().build_local(&mut Default::default()).unwrap(); /// /// rt.block_on(async { /// println!("Hello from the Tokio runtime"); /// }); /// ``` #[allow(unused_variables, unreachable_patterns)] #[cfg(tokio_unstable)] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn build_local(&mut self, options: &LocalOptions) -> io::Result { match &self.kind { Kind::CurrentThread => self.build_current_thread_local_runtime(), _ => panic!("Only current_thread is supported when building a local runtime"), } } fn get_cfg(&self, workers: usize) -> driver::Cfg { driver::Cfg { enable_pause_time: match self.kind { Kind::CurrentThread => true, #[cfg(feature = "rt-multi-thread")] Kind::MultiThread => false, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Kind::MultiThreadAlt => false, }, enable_io: self.enable_io, enable_time: self.enable_time, start_paused: self.start_paused, nevents: self.nevents, workers, } } /// Sets a custom timeout for a thread in the blocking pool. /// /// By default, the timeout for a thread is set to 10 seconds. This can /// be overridden using `.thread_keep_alive()`. /// /// # Example /// /// ``` /// # use tokio::runtime; /// # use std::time::Duration; /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .thread_keep_alive(Duration::from_millis(100)) /// .build(); /// # } /// ``` pub fn thread_keep_alive(&mut self, duration: Duration) -> &mut Self { self.keep_alive = Some(duration); self } /// Sets the number of scheduler ticks after which the scheduler will poll the global /// task queue. /// /// A scheduler "tick" roughly corresponds to one `poll` invocation on a task. /// /// By default the global queue interval is 31 for the current-thread scheduler. Please see /// [the module documentation] for the default behavior of the multi-thread scheduler. /// /// Schedulers have a local queue of already-claimed tasks, and a global queue of incoming /// tasks. Setting the interval to a smaller value increases the fairness of the scheduler, /// at the cost of more synchronization overhead. That can be beneficial for prioritizing /// getting started on new work, especially if tasks frequently yield rather than complete /// or await on further I/O. Conversely, a higher value prioritizes existing work, and /// is a good choice when most tasks quickly complete polling. /// /// [the module documentation]: crate::runtime#multi-threaded-runtime-behavior-at-the-time-of-writing /// /// # Panics /// /// This function will panic if 0 is passed as an argument. /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .global_queue_interval(31) /// .build(); /// # } /// ``` #[track_caller] pub fn global_queue_interval(&mut self, val: u32) -> &mut Self { assert!(val > 0, "global_queue_interval must be greater than 0"); self.global_queue_interval = Some(val); self } /// Sets the number of scheduler ticks after which the scheduler will poll for /// external events (timers, I/O, and so on). /// /// A scheduler "tick" roughly corresponds to one `poll` invocation on a task. /// /// By default, the event interval is `61` for all scheduler types. /// /// Setting the event interval determines the effective "priority" of delivering /// these external events (which may wake up additional tasks), compared to /// executing tasks that are currently ready to run. A smaller value is useful /// when tasks frequently spend a long time in polling, or frequently yield, /// which can result in overly long delays picking up I/O events. Conversely, /// picking up new events requires extra synchronization and syscall overhead, /// so if tasks generally complete their polling quickly, a higher event interval /// will minimize that overhead while still keeping the scheduler responsive to /// events. /// /// # Examples /// /// ``` /// # use tokio::runtime; /// # pub fn main() { /// let rt = runtime::Builder::new_multi_thread() /// .event_interval(31) /// .build(); /// # } /// ``` pub fn event_interval(&mut self, val: u32) -> &mut Self { self.event_interval = val; self } cfg_unstable! { /// Configure how the runtime responds to an unhandled panic on a /// spawned task. /// /// By default, an unhandled panic (i.e. a panic not caught by /// [`std::panic::catch_unwind`]) has no impact on the runtime's /// execution. The panic's error value is forwarded to the task's /// [`JoinHandle`] and all other spawned tasks continue running. /// /// The `unhandled_panic` option enables configuring this behavior. /// /// * `UnhandledPanic::Ignore` is the default behavior. Panics on /// spawned tasks have no impact on the runtime's execution. /// * `UnhandledPanic::ShutdownRuntime` will force the runtime to /// shutdown immediately when a spawned task panics even if that /// task's `JoinHandle` has not been dropped. All other spawned tasks /// will immediately terminate and further calls to /// [`Runtime::block_on`] will panic. /// /// # Panics /// This method panics if called with [`UnhandledPanic::ShutdownRuntime`] /// on a runtime other than the current thread runtime. /// /// # Unstable /// /// This option is currently unstable and its implementation is /// incomplete. The API may change or be removed in the future. See /// issue [tokio-rs/tokio#4516] for more details. /// /// # Examples /// /// The following demonstrates a runtime configured to shutdown on /// panic. The first spawned task panics and results in the runtime /// shutting down. The second spawned task never has a chance to /// execute. The call to `block_on` will panic due to the runtime being /// forcibly shutdown. /// /// ```should_panic /// use tokio::runtime::{self, UnhandledPanic}; /// /// # pub fn main() { /// let rt = runtime::Builder::new_current_thread() /// .unhandled_panic(UnhandledPanic::ShutdownRuntime) /// .build() /// .unwrap(); /// /// rt.spawn(async { panic!("boom"); }); /// rt.spawn(async { /// // This task never completes. /// }); /// /// rt.block_on(async { /// // Do some work /// # loop { tokio::task::yield_now().await; } /// }) /// # } /// ``` /// /// [`JoinHandle`]: struct@crate::task::JoinHandle /// [tokio-rs/tokio#4516]: https://github.com/tokio-rs/tokio/issues/4516 pub fn unhandled_panic(&mut self, behavior: UnhandledPanic) -> &mut Self { if !matches!(self.kind, Kind::CurrentThread) && matches!(behavior, UnhandledPanic::ShutdownRuntime) { panic!("UnhandledPanic::ShutdownRuntime is only supported in current thread runtime"); } self.unhandled_panic = behavior; self } /// Disables the LIFO task scheduler heuristic. /// /// The multi-threaded scheduler includes a heuristic for optimizing /// message-passing patterns. This heuristic results in the **last** /// scheduled task being polled first. /// /// To implement this heuristic, each worker thread has a slot which /// holds the task that should be polled next. However, this slot cannot /// be stolen by other worker threads, which can result in lower total /// throughput when tasks tend to have longer poll times. /// /// This configuration option will disable this heuristic resulting in /// all scheduled tasks being pushed into the worker-local queue, which /// is stealable. /// /// Consider trying this option when the task "scheduled" time is high /// but the runtime is underutilized. Use [tokio-rs/tokio-metrics] to /// collect this data. /// /// # Unstable /// /// This configuration option is considered a workaround for the LIFO /// slot not being stealable. When the slot becomes stealable, we will /// revisit whether or not this option is necessary. See /// issue [tokio-rs/tokio#4941]. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_multi_thread() /// .disable_lifo_slot() /// .build() /// .unwrap(); /// ``` /// /// [tokio-rs/tokio-metrics]: https://github.com/tokio-rs/tokio-metrics /// [tokio-rs/tokio#4941]: https://github.com/tokio-rs/tokio/issues/4941 pub fn disable_lifo_slot(&mut self) -> &mut Self { self.disable_lifo_slot = true; self } /// Specifies the random number generation seed to use within all /// threads associated with the runtime being built. /// /// This option is intended to make certain parts of the runtime /// deterministic (e.g. the [`tokio::select!`] macro). In the case of /// [`tokio::select!`] it will ensure that the order that branches are /// polled is deterministic. /// /// In addition to the code specifying `rng_seed` and interacting with /// the runtime, the internals of Tokio and the Rust compiler may affect /// the sequences of random numbers. In order to ensure repeatable /// results, the version of Tokio, the versions of all other /// dependencies that interact with Tokio, and the Rust compiler version /// should also all remain constant. /// /// # Examples /// /// ``` /// # use tokio::runtime::{self, RngSeed}; /// # pub fn main() { /// let seed = RngSeed::from_bytes(b"place your seed here"); /// let rt = runtime::Builder::new_current_thread() /// .rng_seed(seed) /// .build(); /// # } /// ``` /// /// [`tokio::select!`]: crate::select pub fn rng_seed(&mut self, seed: RngSeed) -> &mut Self { self.seed_generator = RngSeedGenerator::new(seed); self } } cfg_unstable_metrics! { /// Enables tracking the distribution of task poll times. /// /// Task poll times are not instrumented by default as doing so requires /// calling [`Instant::now()`] twice per task poll, which could add /// measurable overhead. Use the [`Handle::metrics()`] to access the /// metrics data. /// /// The histogram uses fixed bucket sizes. In other words, the histogram /// buckets are not dynamic based on input values. Use the /// `metrics_poll_time_histogram` builder methods to configure the /// histogram details. /// /// By default, a linear histogram with 10 buckets each 100 microseconds wide will be used. /// This has an extremely low memory footprint, but may not provide enough granularity. For /// better granularity with low memory usage, use [`metrics_poll_time_histogram_configuration()`] /// to select [`LogHistogram`] instead. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .build() /// .unwrap(); /// # // Test default values here /// # fn us(n: u64) -> std::time::Duration { std::time::Duration::from_micros(n) } /// # let m = rt.handle().metrics(); /// # assert_eq!(m.poll_time_histogram_num_buckets(), 10); /// # assert_eq!(m.poll_time_histogram_bucket_range(0), us(0)..us(100)); /// # assert_eq!(m.poll_time_histogram_bucket_range(1), us(100)..us(200)); /// ``` /// /// [`Handle::metrics()`]: crate::runtime::Handle::metrics /// [`Instant::now()`]: std::time::Instant::now /// [`LogHistogram`]: crate::runtime::LogHistogram /// [`metrics_poll_time_histogram_configuration()`]: Builder::metrics_poll_time_histogram_configuration pub fn enable_metrics_poll_time_histogram(&mut self) -> &mut Self { self.metrics_poll_count_histogram_enable = true; self } /// Deprecated. Use [`enable_metrics_poll_time_histogram()`] instead. /// /// [`enable_metrics_poll_time_histogram()`]: Builder::enable_metrics_poll_time_histogram #[deprecated(note = "`poll_count_histogram` related methods have been renamed `poll_time_histogram` to better reflect their functionality.")] #[doc(hidden)] pub fn enable_metrics_poll_count_histogram(&mut self) -> &mut Self { self.enable_metrics_poll_time_histogram() } /// Sets the histogram scale for tracking the distribution of task poll /// times. /// /// Tracking the distribution of task poll times can be done using a /// linear or log scale. When using linear scale, each histogram bucket /// will represent the same range of poll times. When using log scale, /// each histogram bucket will cover a range twice as big as the /// previous bucket. /// /// **Default:** linear scale. /// /// # Examples /// /// ``` /// use tokio::runtime::{self, HistogramScale}; /// /// # #[allow(deprecated)] /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_count_histogram_scale(HistogramScale::Log) /// .build() /// .unwrap(); /// ``` #[deprecated(note = "use `metrics_poll_time_histogram_configuration`")] pub fn metrics_poll_count_histogram_scale(&mut self, histogram_scale: crate::runtime::HistogramScale) -> &mut Self { self.metrics_poll_count_histogram.legacy_mut(|b|b.scale = histogram_scale); self } /// Configure the histogram for tracking poll times /// /// By default, a linear histogram with 10 buckets each 100 microseconds wide will be used. /// This has an extremely low memory footprint, but may not provide enough granularity. For /// better granularity with low memory usage, use [`LogHistogram`] instead. /// /// # Examples /// Configure a [`LogHistogram`] with [default configuration]: /// ``` /// use tokio::runtime; /// use tokio::runtime::{HistogramConfiguration, LogHistogram}; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_time_histogram_configuration( /// HistogramConfiguration::log(LogHistogram::default()) /// ) /// .build() /// .unwrap(); /// ``` /// /// Configure a linear histogram with 100 buckets, each 10μs wide /// ``` /// use tokio::runtime; /// use std::time::Duration; /// use tokio::runtime::HistogramConfiguration; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_time_histogram_configuration( /// HistogramConfiguration::linear(Duration::from_micros(10), 100) /// ) /// .build() /// .unwrap(); /// ``` /// /// Configure a [`LogHistogram`] with the following settings: /// - Measure times from 100ns to 120s /// - Max error of 0.1 /// - No more than 1024 buckets /// ``` /// use std::time::Duration; /// use tokio::runtime; /// use tokio::runtime::{HistogramConfiguration, LogHistogram}; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_time_histogram_configuration( /// HistogramConfiguration::log(LogHistogram::builder() /// .max_value(Duration::from_secs(120)) /// .min_value(Duration::from_nanos(100)) /// .max_error(0.1) /// .max_buckets(1024) /// .expect("configuration uses 488 buckets") /// ) /// ) /// .build() /// .unwrap(); /// ``` /// /// When migrating from the legacy histogram ([`HistogramScale::Log`]) and wanting /// to match the previous behavior, use `precision_exact(0)`. This creates a histogram /// where each bucket is twice the size of the previous bucket. /// ```rust /// use std::time::Duration; /// use tokio::runtime::{HistogramConfiguration, LogHistogram}; /// let rt = tokio::runtime::Builder::new_current_thread() /// .enable_all() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_time_histogram_configuration(HistogramConfiguration::log( /// LogHistogram::builder() /// .min_value(Duration::from_micros(20)) /// .max_value(Duration::from_millis(4)) /// // Set `precision_exact` to `0` to match `HistogramScale::Log` /// .precision_exact(0) /// .max_buckets(10) /// .unwrap(), /// )) /// .build() /// .unwrap(); /// ``` /// /// [`LogHistogram`]: crate::runtime::LogHistogram /// [default configuration]: crate::runtime::LogHistogramBuilder /// [`HistogramScale::Log`]: crate::runtime::HistogramScale::Log pub fn metrics_poll_time_histogram_configuration(&mut self, configuration: HistogramConfiguration) -> &mut Self { self.metrics_poll_count_histogram.histogram_type = configuration.inner; self } /// Sets the histogram resolution for tracking the distribution of task /// poll times. /// /// The resolution is the histogram's first bucket's range. When using a /// linear histogram scale, each bucket will cover the same range. When /// using a log scale, each bucket will cover a range twice as big as /// the previous bucket. In the log case, the resolution represents the /// smallest bucket range. /// /// Note that, when using log scale, the resolution is rounded up to the /// nearest power of 2 in nanoseconds. /// /// **Default:** 100 microseconds. /// /// # Examples /// /// ``` /// use tokio::runtime; /// use std::time::Duration; /// /// # #[allow(deprecated)] /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_count_histogram_resolution(Duration::from_micros(100)) /// .build() /// .unwrap(); /// ``` #[deprecated(note = "use `metrics_poll_time_histogram_configuration`")] pub fn metrics_poll_count_histogram_resolution(&mut self, resolution: Duration) -> &mut Self { assert!(resolution > Duration::from_secs(0)); // Sanity check the argument and also make the cast below safe. assert!(resolution <= Duration::from_secs(1)); let resolution = resolution.as_nanos() as u64; self.metrics_poll_count_histogram.legacy_mut(|b|b.resolution = resolution); self } /// Sets the number of buckets for the histogram tracking the /// distribution of task poll times. /// /// The last bucket tracks all greater values that fall out of other /// ranges. So, configuring the histogram using a linear scale, /// resolution of 50ms, and 10 buckets, the 10th bucket will track task /// polls that take more than 450ms to complete. /// /// **Default:** 10 /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// # #[allow(deprecated)] /// let rt = runtime::Builder::new_multi_thread() /// .enable_metrics_poll_time_histogram() /// .metrics_poll_count_histogram_buckets(15) /// .build() /// .unwrap(); /// ``` #[deprecated(note = "use `metrics_poll_time_histogram_configuration`")] pub fn metrics_poll_count_histogram_buckets(&mut self, buckets: usize) -> &mut Self { self.metrics_poll_count_histogram.legacy_mut(|b|b.num_buckets = buckets); self } } cfg_loom! { pub(crate) fn local_queue_capacity(&mut self, value: usize) -> &mut Self { assert!(value.is_power_of_two()); self.local_queue_capacity = value; self } } fn build_current_thread_runtime(&mut self) -> io::Result { use crate::runtime::runtime::Scheduler; let (scheduler, handle, blocking_pool) = self.build_current_thread_runtime_components(None)?; Ok(Runtime::from_parts( Scheduler::CurrentThread(scheduler), handle, blocking_pool, )) } #[cfg(tokio_unstable)] fn build_current_thread_local_runtime(&mut self) -> io::Result { use crate::runtime::local_runtime::LocalRuntimeScheduler; let tid = std::thread::current().id(); let (scheduler, handle, blocking_pool) = self.build_current_thread_runtime_components(Some(tid))?; Ok(LocalRuntime::from_parts( LocalRuntimeScheduler::CurrentThread(scheduler), handle, blocking_pool, )) } fn build_current_thread_runtime_components( &mut self, local_tid: Option, ) -> io::Result<(CurrentThread, Handle, BlockingPool)> { use crate::runtime::scheduler; use crate::runtime::Config; let (driver, driver_handle) = driver::Driver::new(self.get_cfg(1))?; // Blocking pool let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads); let blocking_spawner = blocking_pool.spawner().clone(); // Generate a rng seed for this runtime. let seed_generator_1 = self.seed_generator.next_generator(); let seed_generator_2 = self.seed_generator.next_generator(); // And now put a single-threaded scheduler on top of the timer. When // there are no futures ready to do something, it'll let the timer or // the reactor to generate some new stimuli for the futures to continue // in their life. let (scheduler, handle) = CurrentThread::new( driver, driver_handle, blocking_spawner, seed_generator_2, Config { before_park: self.before_park.clone(), after_unpark: self.after_unpark.clone(), before_spawn: self.before_spawn.clone(), after_termination: self.after_termination.clone(), global_queue_interval: self.global_queue_interval, event_interval: self.event_interval, local_queue_capacity: self.local_queue_capacity, #[cfg(tokio_unstable)] unhandled_panic: self.unhandled_panic.clone(), disable_lifo_slot: self.disable_lifo_slot, seed_generator: seed_generator_1, metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), }, local_tid, ); let handle = Handle { inner: scheduler::Handle::CurrentThread(handle), }; Ok((scheduler, handle, blocking_pool)) } fn metrics_poll_count_histogram_builder(&self) -> Option { if self.metrics_poll_count_histogram_enable { Some(self.metrics_poll_count_histogram.clone()) } else { None } } } cfg_io_driver! { impl Builder { /// Enables the I/O driver. /// /// Doing this enables using net, process, signal, and some I/O types on /// the runtime. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_io() /// .build() /// .unwrap(); /// ``` pub fn enable_io(&mut self) -> &mut Self { self.enable_io = true; self } /// Enables the I/O driver and configures the max number of events to be /// processed per tick. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_current_thread() /// .enable_io() /// .max_io_events_per_tick(1024) /// .build() /// .unwrap(); /// ``` pub fn max_io_events_per_tick(&mut self, capacity: usize) -> &mut Self { self.nevents = capacity; self } } } cfg_time! { impl Builder { /// Enables the time driver. /// /// Doing this enables using `tokio::time` on the runtime. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_multi_thread() /// .enable_time() /// .build() /// .unwrap(); /// ``` pub fn enable_time(&mut self) -> &mut Self { self.enable_time = true; self } } } cfg_test_util! { impl Builder { /// Controls if the runtime's clock starts paused or advancing. /// /// Pausing time requires the current-thread runtime; construction of /// the runtime will panic otherwise. /// /// # Examples /// /// ``` /// use tokio::runtime; /// /// let rt = runtime::Builder::new_current_thread() /// .enable_time() /// .start_paused(true) /// .build() /// .unwrap(); /// ``` pub fn start_paused(&mut self, start_paused: bool) -> &mut Self { self.start_paused = start_paused; self } } } cfg_rt_multi_thread! { impl Builder { fn build_threaded_runtime(&mut self) -> io::Result { use crate::loom::sys::num_cpus; use crate::runtime::{Config, runtime::Scheduler}; use crate::runtime::scheduler::{self, MultiThread}; let core_threads = self.worker_threads.unwrap_or_else(num_cpus); let (driver, driver_handle) = driver::Driver::new(self.get_cfg(core_threads))?; // Create the blocking pool let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads + core_threads); let blocking_spawner = blocking_pool.spawner().clone(); // Generate a rng seed for this runtime. let seed_generator_1 = self.seed_generator.next_generator(); let seed_generator_2 = self.seed_generator.next_generator(); let (scheduler, handle, launch) = MultiThread::new( core_threads, driver, driver_handle, blocking_spawner, seed_generator_2, Config { before_park: self.before_park.clone(), after_unpark: self.after_unpark.clone(), before_spawn: self.before_spawn.clone(), after_termination: self.after_termination.clone(), global_queue_interval: self.global_queue_interval, event_interval: self.event_interval, local_queue_capacity: self.local_queue_capacity, #[cfg(tokio_unstable)] unhandled_panic: self.unhandled_panic.clone(), disable_lifo_slot: self.disable_lifo_slot, seed_generator: seed_generator_1, metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), }, ); let handle = Handle { inner: scheduler::Handle::MultiThread(handle) }; // Spawn the thread pool workers let _enter = handle.enter(); launch.launch(); Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool)) } cfg_unstable! { fn build_alt_threaded_runtime(&mut self) -> io::Result { use crate::loom::sys::num_cpus; use crate::runtime::{Config, runtime::Scheduler}; use crate::runtime::scheduler::MultiThreadAlt; let core_threads = self.worker_threads.unwrap_or_else(num_cpus); let (driver, driver_handle) = driver::Driver::new(self.get_cfg(core_threads))?; // Create the blocking pool let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads + core_threads); let blocking_spawner = blocking_pool.spawner().clone(); // Generate a rng seed for this runtime. let seed_generator_1 = self.seed_generator.next_generator(); let seed_generator_2 = self.seed_generator.next_generator(); let (scheduler, handle) = MultiThreadAlt::new( core_threads, driver, driver_handle, blocking_spawner, seed_generator_2, Config { before_park: self.before_park.clone(), after_unpark: self.after_unpark.clone(), before_spawn: self.before_spawn.clone(), after_termination: self.after_termination.clone(), global_queue_interval: self.global_queue_interval, event_interval: self.event_interval, local_queue_capacity: self.local_queue_capacity, #[cfg(tokio_unstable)] unhandled_panic: self.unhandled_panic.clone(), disable_lifo_slot: self.disable_lifo_slot, seed_generator: seed_generator_1, metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), }, ); Ok(Runtime::from_parts(Scheduler::MultiThreadAlt(scheduler), handle, blocking_pool)) } } } } impl fmt::Debug for Builder { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Builder") .field("worker_threads", &self.worker_threads) .field("max_blocking_threads", &self.max_blocking_threads) .field( "thread_name", &" 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.before_stop.as_ref().map(|_| "...")) .field("before_park", &self.before_park.as_ref().map(|_| "...")) .field("after_unpark", &self.after_unpark.as_ref().map(|_| "...")) .finish() } } tokio-1.43.1/src/runtime/config.rs000064400000000000000000000034461046102023000151330ustar 00000000000000#![cfg_attr( any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"), allow(dead_code) )] use crate::runtime::{Callback, TaskCallback}; use crate::util::RngSeedGenerator; pub(crate) struct Config { /// How many ticks before pulling a task from the global/remote queue? pub(crate) global_queue_interval: Option, /// How many ticks before yielding to the driver for timer and I/O events? pub(crate) event_interval: u32, /// How big to make each worker's local queue pub(crate) local_queue_capacity: usize, /// Callback for a worker parking itself pub(crate) before_park: Option, /// Callback for a worker unparking itself pub(crate) after_unpark: Option, /// To run before each task is spawned. pub(crate) before_spawn: Option, /// To run after each task is terminated. pub(crate) after_termination: Option, /// The multi-threaded scheduler includes a per-worker LIFO slot used to /// store the last scheduled task. This can improve certain usage patterns, /// especially message passing between tasks. However, this LIFO slot is not /// currently stealable. /// /// Eventually, the LIFO slot **will** become stealable, however as a /// stop-gap, this unstable option lets users disable the LIFO task. pub(crate) disable_lifo_slot: bool, /// Random number generator seed to configure runtimes to act in a /// deterministic way. pub(crate) seed_generator: RngSeedGenerator, /// How to build poll time histograms pub(crate) metrics_poll_count_histogram: Option, #[cfg(tokio_unstable)] /// How to respond to unhandled task panics. pub(crate) unhandled_panic: crate::runtime::UnhandledPanic, } tokio-1.43.1/src/runtime/context/blocking.rs000064400000000000000000000070201046102023000171320ustar 00000000000000use super::{EnterRuntime, CONTEXT}; use crate::loom::thread::AccessError; use crate::util::markers::NotSendOrSync; use std::marker::PhantomData; use std::time::Duration; /// Guard tracking that a caller has entered a blocking region. #[must_use] pub(crate) struct BlockingRegionGuard { _p: PhantomData, } pub(crate) struct DisallowBlockInPlaceGuard(bool); pub(crate) fn try_enter_blocking_region() -> Option { CONTEXT .try_with(|c| { if c.runtime.get().is_entered() { None } else { Some(BlockingRegionGuard::new()) } // If accessing the thread-local fails, the thread is terminating // and thread-locals are being destroyed. Because we don't know if // we are currently in a runtime or not, we default to being // permissive. }) .unwrap_or_else(|_| Some(BlockingRegionGuard::new())) } /// Disallows blocking in the current runtime context until the guard is dropped. pub(crate) fn disallow_block_in_place() -> DisallowBlockInPlaceGuard { let reset = CONTEXT.with(|c| { if let EnterRuntime::Entered { allow_block_in_place: true, } = c.runtime.get() { c.runtime.set(EnterRuntime::Entered { allow_block_in_place: false, }); true } else { false } }); DisallowBlockInPlaceGuard(reset) } impl BlockingRegionGuard { pub(super) fn new() -> BlockingRegionGuard { BlockingRegionGuard { _p: PhantomData } } /// Blocks the thread on the specified future, returning the value with /// which that future completes. pub(crate) fn block_on(&mut self, f: F) -> Result where F: std::future::Future, { use crate::runtime::park::CachedParkThread; let mut park = CachedParkThread::new(); park.block_on(f) } /// Blocks the thread on the specified future for **at most** `timeout` /// /// If the future completes before `timeout`, the result is returned. If /// `timeout` elapses, then `Err` is returned. pub(crate) fn block_on_timeout(&mut self, f: F, timeout: Duration) -> Result where F: std::future::Future, { use crate::runtime::park::CachedParkThread; use std::task::Context; use std::task::Poll::Ready; use std::time::Instant; let mut park = CachedParkThread::new(); let waker = park.waker().map_err(|_| ())?; let mut cx = Context::from_waker(&waker); pin!(f); let when = Instant::now() + timeout; loop { if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) { return Ok(v); } let now = Instant::now(); if now >= when { return Err(()); } park.park_timeout(when - now); } } } impl Drop for DisallowBlockInPlaceGuard { fn drop(&mut self) { if self.0 { // XXX: Do we want some kind of assertion here, or is "best effort" okay? CONTEXT.with(|c| { if let EnterRuntime::Entered { allow_block_in_place: false, } = c.runtime.get() { c.runtime.set(EnterRuntime::Entered { allow_block_in_place: true, }); } }); } } } tokio-1.43.1/src/runtime/context/current.rs000064400000000000000000000053201046102023000170250ustar 00000000000000use super::{Context, CONTEXT}; use crate::runtime::{scheduler, TryCurrentError}; use crate::util::markers::SyncNotSend; use std::cell::{Cell, RefCell}; use std::marker::PhantomData; #[derive(Debug)] #[must_use] pub(crate) struct SetCurrentGuard { // The previous handle prev: Option, // The depth for this guard depth: usize, // Don't let the type move across threads. _p: PhantomData, } pub(super) struct HandleCell { /// Current handle handle: RefCell>, /// Tracks the number of nested calls to `try_set_current`. depth: Cell, } /// Sets this [`Handle`] as the current active [`Handle`]. /// /// [`Handle`]: crate::runtime::scheduler::Handle pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option { CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok() } pub(crate) fn with_current(f: F) -> Result where F: FnOnce(&scheduler::Handle) -> R, { match CONTEXT.try_with(|ctx| ctx.current.handle.borrow().as_ref().map(f)) { Ok(Some(ret)) => Ok(ret), Ok(None) => Err(TryCurrentError::new_no_context()), Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()), } } impl Context { pub(super) fn set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard { let old_handle = self.current.handle.borrow_mut().replace(handle.clone()); let depth = self.current.depth.get(); assert!(depth != usize::MAX, "reached max `enter` depth"); let depth = depth + 1; self.current.depth.set(depth); SetCurrentGuard { prev: old_handle, depth, _p: PhantomData, } } } impl HandleCell { pub(super) const fn new() -> HandleCell { HandleCell { handle: RefCell::new(None), depth: Cell::new(0), } } } impl Drop for SetCurrentGuard { fn drop(&mut self) { CONTEXT.with(|ctx| { let depth = ctx.current.depth.get(); if depth != self.depth { if !std::thread::panicking() { panic!( "`EnterGuard` values dropped out of order. Guards returned by \ `tokio::runtime::Handle::enter()` must be dropped in the reverse \ order as they were acquired." ); } else { // Just return... this will leave handles in a wonky state though... return; } } *ctx.current.handle.borrow_mut() = self.prev.take(); ctx.current.depth.set(depth - 1); }); } } tokio-1.43.1/src/runtime/context/runtime.rs000064400000000000000000000055331046102023000170340ustar 00000000000000use super::{BlockingRegionGuard, SetCurrentGuard, CONTEXT}; use crate::runtime::scheduler; use crate::util::rand::{FastRand, RngSeed}; use std::fmt; #[derive(Debug, Clone, Copy)] #[must_use] pub(crate) enum EnterRuntime { /// Currently in a runtime context. #[cfg_attr(not(feature = "rt"), allow(dead_code))] Entered { allow_block_in_place: bool }, /// Not in a runtime context **or** a blocking region. NotEntered, } /// Guard tracking that a caller has entered a runtime context. #[must_use] pub(crate) struct EnterRuntimeGuard { /// Tracks that the current thread has entered a blocking function call. pub(crate) blocking: BlockingRegionGuard, #[allow(dead_code)] // Only tracking the guard. pub(crate) handle: SetCurrentGuard, // Tracks the previous random number generator seed old_seed: RngSeed, } /// Marks the current thread as being within the dynamic extent of an /// executor. #[track_caller] pub(crate) fn enter_runtime(handle: &scheduler::Handle, allow_block_in_place: bool, f: F) -> R where F: FnOnce(&mut BlockingRegionGuard) -> R, { let maybe_guard = CONTEXT.with(|c| { if c.runtime.get().is_entered() { None } else { // Set the entered flag c.runtime.set(EnterRuntime::Entered { allow_block_in_place, }); // Generate a new seed let rng_seed = handle.seed_generator().next_seed(); // Swap the RNG seed let mut rng = c.rng.get().unwrap_or_else(FastRand::new); let old_seed = rng.replace_seed(rng_seed); c.rng.set(Some(rng)); Some(EnterRuntimeGuard { blocking: BlockingRegionGuard::new(), handle: c.set_current(handle), old_seed, }) } }); if let Some(mut guard) = maybe_guard { return f(&mut guard.blocking); } panic!( "Cannot start a runtime from within a runtime. This happens \ because a function (like `block_on`) attempted to block the \ current thread while the thread is being used to drive \ asynchronous tasks." ); } impl fmt::Debug for EnterRuntimeGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Enter").finish() } } impl Drop for EnterRuntimeGuard { fn drop(&mut self) { CONTEXT.with(|c| { assert!(c.runtime.get().is_entered()); c.runtime.set(EnterRuntime::NotEntered); // Replace the previous RNG seed let mut rng = c.rng.get().unwrap_or_else(FastRand::new); rng.replace_seed(self.old_seed.clone()); c.rng.set(Some(rng)); }); } } impl EnterRuntime { pub(crate) fn is_entered(self) -> bool { matches!(self, EnterRuntime::Entered { .. }) } } tokio-1.43.1/src/runtime/context/runtime_mt.rs000064400000000000000000000017501046102023000175310ustar 00000000000000use super::{EnterRuntime, CONTEXT}; /// Returns true if in a runtime context. pub(crate) fn current_enter_context() -> EnterRuntime { CONTEXT.with(|c| c.runtime.get()) } /// Forces the current "entered" state to be cleared while the closure /// is executed. pub(crate) fn exit_runtime R, R>(f: F) -> R { // Reset in case the closure panics struct Reset(EnterRuntime); impl Drop for Reset { fn drop(&mut self) { CONTEXT.with(|c| { assert!( !c.runtime.get().is_entered(), "closure claimed permanent executor" ); c.runtime.set(self.0); }); } } let was = CONTEXT.with(|c| { let e = c.runtime.get(); assert!(e.is_entered(), "asked to exit when not entered"); c.runtime.set(EnterRuntime::NotEntered); e }); let _reset = Reset(was); // dropping _reset after f() will reset ENTERED f() } tokio-1.43.1/src/runtime/context/scoped.rs000064400000000000000000000022451046102023000166230ustar 00000000000000use std::cell::Cell; use std::ptr; /// Scoped thread-local storage pub(super) struct Scoped { pub(super) inner: Cell<*const T>, } impl Scoped { pub(super) const fn new() -> Scoped { Scoped { inner: Cell::new(ptr::null()), } } /// Inserts a value into the scoped cell for the duration of the closure pub(super) fn set(&self, t: &T, f: F) -> R where F: FnOnce() -> R, { struct Reset<'a, T> { cell: &'a Cell<*const T>, prev: *const T, } impl Drop for Reset<'_, T> { fn drop(&mut self) { self.cell.set(self.prev); } } let prev = self.inner.get(); self.inner.set(t as *const _); let _reset = Reset { cell: &self.inner, prev, }; f() } /// Gets the value out of the scoped cell; pub(super) fn with(&self, f: F) -> R where F: FnOnce(Option<&T>) -> R, { let val = self.inner.get(); if val.is_null() { f(None) } else { unsafe { f(Some(&*val)) } } } } tokio-1.43.1/src/runtime/context.rs000064400000000000000000000137671046102023000153610ustar 00000000000000use crate::loom::thread::AccessError; use crate::runtime::coop; use std::cell::Cell; #[cfg(any(feature = "rt", feature = "macros", feature = "time"))] use crate::util::rand::FastRand; cfg_rt! { mod blocking; pub(crate) use blocking::{disallow_block_in_place, try_enter_blocking_region, BlockingRegionGuard}; mod current; pub(crate) use current::{with_current, try_set_current, SetCurrentGuard}; mod runtime; pub(crate) use runtime::{EnterRuntime, enter_runtime}; mod scoped; use scoped::Scoped; use crate::runtime::{scheduler, task::Id}; use std::task::Waker; cfg_taskdump! { use crate::runtime::task::trace; } } cfg_rt_multi_thread! { mod runtime_mt; pub(crate) use runtime_mt::{current_enter_context, exit_runtime}; } struct Context { /// Uniquely identifies the current thread #[cfg(feature = "rt")] thread_id: Cell>, /// Handle to the runtime scheduler running on the current thread. #[cfg(feature = "rt")] current: current::HandleCell, /// Handle to the scheduler's internal "context" #[cfg(feature = "rt")] scheduler: Scoped, #[cfg(feature = "rt")] current_task_id: Cell>, /// Tracks if the current thread is currently driving a runtime. /// Note, that if this is set to "entered", the current scheduler /// handle may not reference the runtime currently executing. This /// is because other runtime handles may be set to current from /// within a runtime. #[cfg(feature = "rt")] runtime: Cell, #[cfg(any(feature = "rt", feature = "macros", feature = "time"))] rng: Cell>, /// Tracks the amount of "work" a task may still do before yielding back to /// the scheduler budget: Cell, #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] trace: trace::Context, } tokio_thread_local! { static CONTEXT: Context = const { Context { #[cfg(feature = "rt")] thread_id: Cell::new(None), // Tracks the current runtime handle to use when spawning, // accessing drivers, etc... #[cfg(feature = "rt")] current: current::HandleCell::new(), // Tracks the current scheduler internal context #[cfg(feature = "rt")] scheduler: Scoped::new(), #[cfg(feature = "rt")] current_task_id: Cell::new(None), // Tracks if the current thread is currently driving a runtime. // Note, that if this is set to "entered", the current scheduler // handle may not reference the runtime currently executing. This // is because other runtime handles may be set to current from // within a runtime. #[cfg(feature = "rt")] runtime: Cell::new(EnterRuntime::NotEntered), #[cfg(any(feature = "rt", feature = "macros", feature = "time"))] rng: Cell::new(None), budget: Cell::new(coop::Budget::unconstrained()), #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any( target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64" ) ))] trace: trace::Context::new(), } } } #[cfg(any( feature = "time", feature = "macros", all(feature = "sync", feature = "rt") ))] pub(crate) fn thread_rng_n(n: u32) -> u32 { CONTEXT.with(|ctx| { let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new); let ret = rng.fastrand_n(n); ctx.rng.set(Some(rng)); ret }) } pub(super) fn budget(f: impl FnOnce(&Cell) -> R) -> Result { CONTEXT.try_with(|ctx| f(&ctx.budget)) } cfg_rt! { use crate::runtime::ThreadId; pub(crate) fn thread_id() -> Result { CONTEXT.try_with(|ctx| { match ctx.thread_id.get() { Some(id) => id, None => { let id = ThreadId::next(); ctx.thread_id.set(Some(id)); id } } }) } pub(crate) fn set_current_task_id(id: Option) -> Option { CONTEXT.try_with(|ctx| ctx.current_task_id.replace(id)).unwrap_or(None) } pub(crate) fn current_task_id() -> Option { CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None) } #[track_caller] pub(crate) fn defer(waker: &Waker) { with_scheduler(|maybe_scheduler| { if let Some(scheduler) = maybe_scheduler { scheduler.defer(waker); } else { // Called from outside of the runtime, immediately wake the // task. waker.wake_by_ref(); } }); } pub(super) fn set_scheduler(v: &scheduler::Context, f: impl FnOnce() -> R) -> R { CONTEXT.with(|c| c.scheduler.set(v, f)) } #[track_caller] pub(super) fn with_scheduler(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R { let mut f = Some(f); CONTEXT.try_with(|c| { let f = f.take().unwrap(); if matches!(c.runtime.get(), EnterRuntime::Entered { .. }) { c.scheduler.with(f) } else { f(None) } }) .unwrap_or_else(|_| (f.take().unwrap())(None)) } cfg_taskdump! { /// SAFETY: Callers of this function must ensure that trace frames always /// form a valid linked list. pub(crate) unsafe fn with_trace(f: impl FnOnce(&trace::Context) -> R) -> Option { CONTEXT.try_with(|c| f(&c.trace)).ok() } } } tokio-1.43.1/src/runtime/coop.rs000064400000000000000000000274511046102023000146300ustar 00000000000000#![cfg_attr(not(feature = "full"), allow(dead_code))] //! Yield points for improved cooperative scheduling. //! //! Documentation for this can be found in the [`tokio::task`] module. //! //! [`tokio::task`]: crate::task. // ```ignore // # use tokio_stream::{Stream, StreamExt}; // async fn drop_all(mut input: I) { // while let Some(_) = input.next().await { // tokio::coop::proceed().await; // } // } // ``` // // The `proceed` future will coordinate with the executor to make sure that // every so often control is yielded back to the executor so it can run other // tasks. // // # Placing yield points // // Voluntary yield points should be placed _after_ at least some work has been // done. If they are not, a future sufficiently deep in the task hierarchy may // end up _never_ getting to run because of the number of yield points that // inevitably appear before it is reached. In general, you will want yield // points to only appear in "leaf" futures -- those that do not themselves poll // other futures. By doing this, you avoid double-counting each iteration of // the outer future against the cooperating budget. use crate::runtime::context; /// Opaque type tracking the amount of "work" a task may still do before /// yielding back to the scheduler. #[derive(Debug, Copy, Clone)] pub(crate) struct Budget(Option); pub(crate) struct BudgetDecrement { success: bool, hit_zero: bool, } impl Budget { /// Budget assigned to a task on each poll. /// /// The value itself is chosen somewhat arbitrarily. It needs to be high /// enough to amortize wakeup and scheduling costs, but low enough that we /// do not starve other tasks for too long. The value also needs to be high /// enough that particularly deep tasks are able to do at least some useful /// work at all. /// /// Note that as more yield points are added in the ecosystem, this value /// will probably also have to be raised. const fn initial() -> Budget { Budget(Some(128)) } /// Returns an unconstrained budget. Operations will not be limited. pub(super) const fn unconstrained() -> Budget { Budget(None) } fn has_remaining(self) -> bool { self.0.map_or(true, |budget| budget > 0) } } /// Runs the given closure with a cooperative task budget. When the function /// returns, the budget is reset to the value prior to calling the function. #[inline(always)] pub(crate) fn budget(f: impl FnOnce() -> R) -> R { with_budget(Budget::initial(), f) } /// Runs the given closure with an unconstrained task budget. When the function returns, the budget /// is reset to the value prior to calling the function. #[inline(always)] pub(crate) fn with_unconstrained(f: impl FnOnce() -> R) -> R { with_budget(Budget::unconstrained(), f) } #[inline(always)] fn with_budget(budget: Budget, f: impl FnOnce() -> R) -> R { struct ResetGuard { prev: Budget, } impl Drop for ResetGuard { fn drop(&mut self) { let _ = context::budget(|cell| { cell.set(self.prev); }); } } #[allow(unused_variables)] let maybe_guard = context::budget(|cell| { let prev = cell.get(); cell.set(budget); ResetGuard { prev } }); // The function is called regardless even if the budget is not successfully // set due to the thread-local being destroyed. f() } #[inline(always)] pub(crate) fn has_budget_remaining() -> bool { // If the current budget cannot be accessed due to the thread-local being // shutdown, then we assume there is budget remaining. context::budget(|cell| cell.get().has_remaining()).unwrap_or(true) } cfg_rt_multi_thread! { /// Sets the current task's budget. pub(crate) fn set(budget: Budget) { let _ = context::budget(|cell| cell.set(budget)); } } cfg_rt! { /// Forcibly removes the budgeting constraints early. /// /// Returns the remaining budget pub(crate) fn stop() -> Budget { context::budget(|cell| { let prev = cell.get(); cell.set(Budget::unconstrained()); prev }).unwrap_or(Budget::unconstrained()) } } cfg_coop! { use pin_project_lite::pin_project; use std::cell::Cell; use std::future::Future; use std::pin::Pin; use std::task::{ready, Context, Poll}; #[must_use] pub(crate) struct RestoreOnPending(Cell); impl RestoreOnPending { pub(crate) fn made_progress(&self) { self.0.set(Budget::unconstrained()); } } impl Drop for RestoreOnPending { fn drop(&mut self) { // Don't reset if budget was unconstrained or if we made progress. // They are both represented as the remembered budget being unconstrained. let budget = self.0.get(); if !budget.is_unconstrained() { let _ = context::budget(|cell| { cell.set(budget); }); } } } /// Returns `Poll::Pending` if the current task has exceeded its budget and should yield. /// /// When you call this method, the current budget is decremented. However, to ensure that /// progress is made every time a task is polled, the budget is automatically restored to its /// former value if the returned `RestoreOnPending` is dropped. It is the caller's /// responsibility to call `RestoreOnPending::made_progress` if it made progress, to ensure /// that the budget empties appropriately. /// /// Note that `RestoreOnPending` restores the budget **as it was before `poll_proceed`**. /// Therefore, if the budget is _further_ adjusted between when `poll_proceed` returns and /// `RestRestoreOnPending` is dropped, those adjustments are erased unless the caller indicates /// that progress was made. #[inline] pub(crate) fn poll_proceed(cx: &mut Context<'_>) -> Poll { context::budget(|cell| { let mut budget = cell.get(); let decrement = budget.decrement(); if decrement.success { let restore = RestoreOnPending(Cell::new(cell.get())); cell.set(budget); // avoid double counting if decrement.hit_zero { inc_budget_forced_yield_count(); } Poll::Ready(restore) } else { cx.waker().wake_by_ref(); Poll::Pending } }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained())))) } cfg_rt! { cfg_unstable_metrics! { #[inline(always)] fn inc_budget_forced_yield_count() { let _ = context::with_current(|handle| { handle.scheduler_metrics().inc_budget_forced_yield_count(); }); } } cfg_not_unstable_metrics! { #[inline(always)] fn inc_budget_forced_yield_count() {} } } cfg_not_rt! { #[inline(always)] fn inc_budget_forced_yield_count() {} } impl Budget { /// Decrements the budget. Returns `true` if successful. Decrementing fails /// when there is not enough remaining budget. fn decrement(&mut self) -> BudgetDecrement { if let Some(num) = &mut self.0 { if *num > 0 { *num -= 1; let hit_zero = *num == 0; BudgetDecrement { success: true, hit_zero } } else { BudgetDecrement { success: false, hit_zero: false } } } else { BudgetDecrement { success: true, hit_zero: false } } } fn is_unconstrained(self) -> bool { self.0.is_none() } } pin_project! { /// Future wrapper to ensure cooperative scheduling. /// /// When being polled `poll_proceed` is called before the inner future is polled to check /// if the inner future has exceeded its budget. If the inner future resolves, this will /// automatically call `RestoreOnPending::made_progress` before resolving this future with /// the result of the inner one. If polling the inner future is pending, polling this future /// type will also return a `Poll::Pending`. #[must_use = "futures do nothing unless polled"] pub(crate) struct Coop { #[pin] pub(crate) fut: F, } } impl Future for Coop { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let coop = ready!(poll_proceed(cx)); let me = self.project(); if let Poll::Ready(ret) = me.fut.poll(cx) { coop.made_progress(); Poll::Ready(ret) } else { Poll::Pending } } } /// Run a future with a budget constraint for cooperative scheduling. /// If the future exceeds its budget while being polled, control is yielded back to the /// runtime. #[inline] pub(crate) fn cooperative(fut: F) -> Coop { Coop { fut } } } #[cfg(all(test, not(loom)))] mod test { use super::*; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; fn get() -> Budget { context::budget(|cell| cell.get()).unwrap_or(Budget::unconstrained()) } #[test] fn budgeting() { use std::future::poll_fn; use tokio_test::*; assert!(get().0.is_none()); let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx))); assert!(get().0.is_none()); drop(coop); assert!(get().0.is_none()); budget(|| { assert_eq!(get().0, Budget::initial().0); let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx))); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1); drop(coop); // we didn't make progress assert_eq!(get().0, Budget::initial().0); let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx))); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1); coop.made_progress(); drop(coop); // we _did_ make progress assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1); let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx))); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 2); coop.made_progress(); drop(coop); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 2); budget(|| { assert_eq!(get().0, Budget::initial().0); let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx))); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1); coop.made_progress(); drop(coop); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1); }); assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 2); }); assert!(get().0.is_none()); budget(|| { let n = get().0.unwrap(); for _ in 0..n { let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx))); coop.made_progress(); } let mut task = task::spawn(poll_fn(|cx| { let coop = std::task::ready!(poll_proceed(cx)); coop.made_progress(); Poll::Ready(()) })); assert_pending!(task.poll()); }); } } tokio-1.43.1/src/runtime/driver.rs000064400000000000000000000245751046102023000151670ustar 00000000000000//! Abstracts out the entire chain of runtime sub-drivers into common types. // Eventually, this file will see significant refactoring / cleanup. For now, we // don't need to worry much about dead code with certain feature permutations. #![cfg_attr( any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"), allow(dead_code) )] use crate::runtime::park::{ParkThread, UnparkThread}; use std::io; use std::time::Duration; #[derive(Debug)] pub(crate) struct Driver { inner: TimeDriver, } #[derive(Debug)] pub(crate) struct Handle { /// IO driver handle pub(crate) io: IoHandle, /// Signal driver handle #[cfg_attr(any(not(unix), loom), allow(dead_code))] pub(crate) signal: SignalHandle, /// Time driver handle pub(crate) time: TimeHandle, /// Source of `Instant::now()` #[cfg_attr(not(all(feature = "time", feature = "test-util")), allow(dead_code))] pub(crate) clock: Clock, } pub(crate) struct Cfg { pub(crate) enable_io: bool, pub(crate) enable_time: bool, pub(crate) enable_pause_time: bool, pub(crate) start_paused: bool, pub(crate) nevents: usize, pub(crate) workers: usize, } impl Driver { pub(crate) fn new(cfg: Cfg) -> io::Result<(Self, Handle)> { let (io_stack, io_handle, signal_handle) = create_io_stack(cfg.enable_io, cfg.nevents)?; let clock = create_clock(cfg.enable_pause_time, cfg.start_paused); let (time_driver, time_handle) = create_time_driver(cfg.enable_time, io_stack, &clock, cfg.workers); Ok(( Self { inner: time_driver }, Handle { io: io_handle, signal: signal_handle, time: time_handle, clock, }, )) } pub(crate) fn is_enabled(&self) -> bool { self.inner.is_enabled() } pub(crate) fn park(&mut self, handle: &Handle) { self.inner.park(handle); } pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) { self.inner.park_timeout(handle, duration); } pub(crate) fn shutdown(&mut self, handle: &Handle) { self.inner.shutdown(handle); } } impl Handle { pub(crate) fn unpark(&self) { #[cfg(feature = "time")] if let Some(handle) = &self.time { handle.unpark(); } self.io.unpark(); } cfg_io_driver! { #[track_caller] pub(crate) fn io(&self) -> &crate::runtime::io::Handle { self.io .as_ref() .expect("A Tokio 1.x context was found, but IO is disabled. Call `enable_io` on the runtime builder to enable IO.") } } cfg_signal_internal_and_unix! { #[track_caller] pub(crate) fn signal(&self) -> &crate::runtime::signal::Handle { self.signal .as_ref() .expect("there is no signal driver running, must be called from the context of Tokio runtime") } } cfg_time! { /// Returns a reference to the time driver handle. /// /// Panics if no time driver is present. #[track_caller] pub(crate) fn time(&self) -> &crate::runtime::time::Handle { self.time .as_ref() .expect("A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.") } pub(crate) fn clock(&self) -> &Clock { &self.clock } } } // ===== io driver ===== cfg_io_driver! { pub(crate) type IoDriver = crate::runtime::io::Driver; #[derive(Debug)] pub(crate) enum IoStack { Enabled(ProcessDriver), Disabled(ParkThread), } #[derive(Debug)] pub(crate) enum IoHandle { Enabled(crate::runtime::io::Handle), Disabled(UnparkThread), } fn create_io_stack(enabled: bool, nevents: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> { #[cfg(loom)] assert!(!enabled); let ret = if enabled { let (io_driver, io_handle) = crate::runtime::io::Driver::new(nevents)?; let (signal_driver, signal_handle) = create_signal_driver(io_driver, &io_handle)?; let process_driver = create_process_driver(signal_driver); (IoStack::Enabled(process_driver), IoHandle::Enabled(io_handle), signal_handle) } else { let park_thread = ParkThread::new(); let unpark_thread = park_thread.unpark(); (IoStack::Disabled(park_thread), IoHandle::Disabled(unpark_thread), Default::default()) }; Ok(ret) } impl IoStack { pub(crate) fn is_enabled(&self) -> bool { match self { IoStack::Enabled(..) => true, IoStack::Disabled(..) => false, } } pub(crate) fn park(&mut self, handle: &Handle) { match self { IoStack::Enabled(v) => v.park(handle), IoStack::Disabled(v) => v.park(), } } pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) { match self { IoStack::Enabled(v) => v.park_timeout(handle, duration), IoStack::Disabled(v) => v.park_timeout(duration), } } pub(crate) fn shutdown(&mut self, handle: &Handle) { match self { IoStack::Enabled(v) => v.shutdown(handle), IoStack::Disabled(v) => v.shutdown(), } } } impl IoHandle { pub(crate) fn unpark(&self) { match self { IoHandle::Enabled(handle) => handle.unpark(), IoHandle::Disabled(handle) => handle.unpark(), } } pub(crate) fn as_ref(&self) -> Option<&crate::runtime::io::Handle> { match self { IoHandle::Enabled(v) => Some(v), IoHandle::Disabled(..) => None, } } } } cfg_not_io_driver! { pub(crate) type IoHandle = UnparkThread; #[derive(Debug)] pub(crate) struct IoStack(ParkThread); fn create_io_stack(_enabled: bool, _nevents: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> { let park_thread = ParkThread::new(); let unpark_thread = park_thread.unpark(); Ok((IoStack(park_thread), unpark_thread, Default::default())) } impl IoStack { pub(crate) fn park(&mut self, _handle: &Handle) { self.0.park(); } pub(crate) fn park_timeout(&mut self, _handle: &Handle, duration: Duration) { self.0.park_timeout(duration); } pub(crate) fn shutdown(&mut self, _handle: &Handle) { self.0.shutdown(); } /// This is not a "real" driver, so it is not considered enabled. pub(crate) fn is_enabled(&self) -> bool { false } } } // ===== signal driver ===== cfg_signal_internal_and_unix! { type SignalDriver = crate::runtime::signal::Driver; pub(crate) type SignalHandle = Option; fn create_signal_driver(io_driver: IoDriver, io_handle: &crate::runtime::io::Handle) -> io::Result<(SignalDriver, SignalHandle)> { let driver = crate::runtime::signal::Driver::new(io_driver, io_handle)?; let handle = driver.handle(); Ok((driver, Some(handle))) } } cfg_not_signal_internal! { pub(crate) type SignalHandle = (); cfg_io_driver! { type SignalDriver = IoDriver; fn create_signal_driver(io_driver: IoDriver, _io_handle: &crate::runtime::io::Handle) -> io::Result<(SignalDriver, SignalHandle)> { Ok((io_driver, ())) } } } // ===== process driver ===== cfg_process_driver! { type ProcessDriver = crate::runtime::process::Driver; fn create_process_driver(signal_driver: SignalDriver) -> ProcessDriver { ProcessDriver::new(signal_driver) } } cfg_not_process_driver! { cfg_io_driver! { type ProcessDriver = SignalDriver; fn create_process_driver(signal_driver: SignalDriver) -> ProcessDriver { signal_driver } } } // ===== time driver ===== cfg_time! { #[derive(Debug)] pub(crate) enum TimeDriver { Enabled { driver: crate::runtime::time::Driver, }, Disabled(IoStack), } pub(crate) type Clock = crate::time::Clock; pub(crate) type TimeHandle = Option; fn create_clock(enable_pausing: bool, start_paused: bool) -> Clock { crate::time::Clock::new(enable_pausing, start_paused) } fn create_time_driver( enable: bool, io_stack: IoStack, clock: &Clock, workers: usize, ) -> (TimeDriver, TimeHandle) { if enable { let (driver, handle) = crate::runtime::time::Driver::new(io_stack, clock, workers as u32); (TimeDriver::Enabled { driver }, Some(handle)) } else { (TimeDriver::Disabled(io_stack), None) } } impl TimeDriver { pub(crate) fn is_enabled(&self) -> bool { match self { TimeDriver::Enabled { .. } => true, TimeDriver::Disabled(inner) => inner.is_enabled(), } } pub(crate) fn park(&mut self, handle: &Handle) { match self { TimeDriver::Enabled { driver, .. } => driver.park(handle), TimeDriver::Disabled(v) => v.park(handle), } } pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) { match self { TimeDriver::Enabled { driver } => driver.park_timeout(handle, duration), TimeDriver::Disabled(v) => v.park_timeout(handle, duration), } } pub(crate) fn shutdown(&mut self, handle: &Handle) { match self { TimeDriver::Enabled { driver } => driver.shutdown(handle), TimeDriver::Disabled(v) => v.shutdown(handle), } } } } cfg_not_time! { type TimeDriver = IoStack; pub(crate) type Clock = (); pub(crate) type TimeHandle = (); fn create_clock(_enable_pausing: bool, _start_paused: bool) -> Clock { () } fn create_time_driver( _enable: bool, io_stack: IoStack, _clock: &Clock, _workers: usize, ) -> (TimeDriver, TimeHandle) { (io_stack, ()) } } tokio-1.43.1/src/runtime/dump.rs000064400000000000000000000201021046102023000146170ustar 00000000000000//! Snapshots of runtime state. //! //! See [`Handle::dump`][crate::runtime::Handle::dump]. use crate::task::Id; use std::{fmt, path::Path}; /// A snapshot of a runtime's state. /// /// See [`Handle::dump`][crate::runtime::Handle::dump]. #[derive(Debug)] pub struct Dump { tasks: Tasks, } /// Snapshots of tasks. /// /// See [`Handle::dump`][crate::runtime::Handle::dump]. #[derive(Debug)] pub struct Tasks { tasks: Vec, } /// A snapshot of a task. /// /// See [`Handle::dump`][crate::runtime::Handle::dump]. #[derive(Debug)] pub struct Task { id: Id, trace: Trace, } /// Represents an address that should not be dereferenced. /// /// This type exists to get the auto traits correct, the public API /// uses raw pointers to make life easier for users. #[derive(Copy, Clone, Debug)] struct Address(*mut std::ffi::c_void); // Safe since Address should not be dereferenced unsafe impl Send for Address {} unsafe impl Sync for Address {} /// A backtrace symbol. /// /// This struct provides accessors for backtrace symbols, similar to [`backtrace::BacktraceSymbol`]. #[derive(Clone, Debug)] pub struct BacktraceSymbol { name: Option>, name_demangled: Option>, addr: Option

, filename: Option, lineno: Option, colno: Option, } impl BacktraceSymbol { pub(crate) fn from_backtrace_symbol(sym: &backtrace::BacktraceSymbol) -> Self { let name = sym.name(); Self { name: name.as_ref().map(|name| name.as_bytes().into()), name_demangled: name.map(|name| format!("{}", name).into()), addr: sym.addr().map(Address), filename: sym.filename().map(From::from), lineno: sym.lineno(), colno: sym.colno(), } } /// Return the raw name of the symbol. pub fn name_raw(&self) -> Option<&[u8]> { self.name.as_deref() } /// Return the demangled name of the symbol. pub fn name_demangled(&self) -> Option<&str> { self.name_demangled.as_deref() } /// Returns the starting address of this symbol. pub fn addr(&self) -> Option<*mut std::ffi::c_void> { self.addr.map(|addr| addr.0) } /// Returns the file name where this function was defined. If debuginfo /// is missing, this is likely to return None. pub fn filename(&self) -> Option<&Path> { self.filename.as_deref() } /// Returns the line number for where this symbol is currently executing. /// /// If debuginfo is missing, this is likely to return `None`. pub fn lineno(&self) -> Option { self.lineno } /// Returns the column number for where this symbol is currently executing. /// /// If debuginfo is missing, this is likely to return `None`. pub fn colno(&self) -> Option { self.colno } } /// A backtrace frame. /// /// This struct represents one stack frame in a captured backtrace, similar to [`backtrace::BacktraceFrame`]. #[derive(Clone, Debug)] pub struct BacktraceFrame { ip: Address, symbol_address: Address, symbols: Box<[BacktraceSymbol]>, } impl BacktraceFrame { pub(crate) fn from_resolved_backtrace_frame(frame: &backtrace::BacktraceFrame) -> Self { Self { ip: Address(frame.ip()), symbol_address: Address(frame.symbol_address()), symbols: frame .symbols() .iter() .map(BacktraceSymbol::from_backtrace_symbol) .collect(), } } /// Return the instruction pointer of this frame. /// /// See the ABI docs for your platform for the exact meaning. pub fn ip(&self) -> *mut std::ffi::c_void { self.ip.0 } /// Returns the starting symbol address of the frame of this function. pub fn symbol_address(&self) -> *mut std::ffi::c_void { self.symbol_address.0 } /// Return an iterator over the symbols of this backtrace frame. /// /// Due to inlining, it is possible for there to be multiple [`BacktraceSymbol`] items relating /// to a single frame. The first symbol listed is the "innermost function", /// whereas the last symbol is the outermost (last caller). pub fn symbols(&self) -> impl Iterator { self.symbols.iter() } } /// A captured backtrace. /// /// This struct provides access to each backtrace frame, similar to [`backtrace::Backtrace`]. #[derive(Clone, Debug)] pub struct Backtrace { frames: Box<[BacktraceFrame]>, } impl Backtrace { /// Return the frames in this backtrace, innermost (in a task dump, /// likely to be a leaf future's poll function) first. pub fn frames(&self) -> impl Iterator { self.frames.iter() } } /// An execution trace of a task's last poll. /// ///
/// /// Resolving a backtrace, either via the [`Display`][std::fmt::Display] impl or via /// [`resolve_backtraces`][Trace::resolve_backtraces], parses debuginfo, which is /// possibly a CPU-expensive operation that can take a platform-specific but /// long time to run - often over 100 milliseconds, especially if the current /// process's binary is big. In some cases, the platform might internally cache some of the /// debuginfo, so successive calls to `resolve_backtraces` might be faster than /// the first call, but all guarantees are platform-dependent. /// /// To avoid blocking the runtime, it is recommended /// that you resolve backtraces inside of a [`spawn_blocking()`][crate::task::spawn_blocking] /// and to have some concurrency-limiting mechanism to avoid unexpected performance impact. ///
/// /// See [`Handle::dump`][crate::runtime::Handle::dump]. #[derive(Debug)] pub struct Trace { inner: super::task::trace::Trace, } impl Trace { /// Resolve and return a list of backtraces that are involved in polls in this trace. /// /// The exact backtraces included here are unstable and might change in the future, /// but you can expect one [`Backtrace`] for every call to /// [`poll`] to a bottom-level Tokio future - so if something like [`join!`] is /// used, there will be a backtrace for each future in the join. /// /// [`poll`]: std::future::Future::poll /// [`join!`]: macro@join pub fn resolve_backtraces(&self) -> Vec { self.inner .backtraces() .iter() .map(|backtrace| { let mut backtrace = backtrace::Backtrace::from(backtrace.clone()); backtrace.resolve(); Backtrace { frames: backtrace .frames() .iter() .map(BacktraceFrame::from_resolved_backtrace_frame) .collect(), } }) .collect() } } impl Dump { pub(crate) fn new(tasks: Vec) -> Self { Self { tasks: Tasks { tasks }, } } /// Tasks in this snapshot. pub fn tasks(&self) -> &Tasks { &self.tasks } } impl Tasks { /// Iterate over tasks. pub fn iter(&self) -> impl Iterator { self.tasks.iter() } } impl Task { pub(crate) fn new(id: Id, trace: super::task::trace::Trace) -> Self { Self { id, trace: Trace { inner: trace }, } } /// Returns a [task ID] that uniquely identifies this task relative to other /// tasks spawned at the time of the dump. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [task ID]: crate::task::Id /// [unstable]: crate#unstable-features #[cfg(tokio_unstable)] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn id(&self) -> Id { self.id } /// A trace of this task's state. pub fn trace(&self) -> &Trace { &self.trace } } impl fmt::Display for Trace { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.inner.fmt(f) } } tokio-1.43.1/src/runtime/handle.rs000064400000000000000000000601511046102023000151150ustar 00000000000000#[cfg(tokio_unstable)] use crate::runtime; use crate::runtime::{context, scheduler, RuntimeFlavor, RuntimeMetrics}; /// Handle to the runtime. /// /// The handle is internally reference-counted and can be freely cloned. A handle can be /// obtained using the [`Runtime::handle`] method. /// /// [`Runtime::handle`]: crate::runtime::Runtime::handle() #[derive(Debug, Clone)] // When the `rt` feature is *not* enabled, this type is still defined, but not // included in the public API. pub struct Handle { pub(crate) inner: scheduler::Handle, } use crate::runtime::task::JoinHandle; use crate::runtime::BOX_FUTURE_THRESHOLD; use crate::util::error::{CONTEXT_MISSING_ERROR, THREAD_LOCAL_DESTROYED_ERROR}; use crate::util::trace::SpawnMeta; use std::future::Future; use std::marker::PhantomData; use std::{error, fmt, mem}; /// Runtime context guard. /// /// Returned by [`Runtime::enter`] and [`Handle::enter`], the context guard exits /// the runtime context on drop. /// /// [`Runtime::enter`]: fn@crate::runtime::Runtime::enter #[derive(Debug)] #[must_use = "Creating and dropping a guard does nothing"] pub struct EnterGuard<'a> { _guard: context::SetCurrentGuard, _handle_lifetime: PhantomData<&'a Handle>, } impl Handle { /// Enters the runtime context. This allows you to construct types that must /// have an executor available on creation such as [`Sleep`] or /// [`TcpStream`]. It will also allow you to call methods such as /// [`tokio::spawn`] and [`Handle::current`] without panicking. /// /// # Panics /// /// When calling `Handle::enter` multiple times, the returned guards /// **must** be dropped in the reverse order that they were acquired. /// Failure to do so will result in a panic and possible memory leaks. /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// let rt = Runtime::new().unwrap(); /// /// let _guard = rt.enter(); /// tokio::spawn(async { /// println!("Hello world!"); /// }); /// ``` /// /// Do **not** do the following, this shows a scenario that will result in a /// panic and possible memory leak. /// /// ```should_panic /// use tokio::runtime::Runtime; /// /// let rt1 = Runtime::new().unwrap(); /// let rt2 = Runtime::new().unwrap(); /// /// let enter1 = rt1.enter(); /// let enter2 = rt2.enter(); /// /// drop(enter1); /// drop(enter2); /// ``` /// /// [`Sleep`]: struct@crate::time::Sleep /// [`TcpStream`]: struct@crate::net::TcpStream /// [`tokio::spawn`]: fn@crate::spawn pub fn enter(&self) -> EnterGuard<'_> { EnterGuard { _guard: match context::try_set_current(&self.inner) { Some(guard) => guard, None => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), }, _handle_lifetime: PhantomData, } } /// Returns a `Handle` view over the currently running `Runtime`. /// /// # Panics /// /// This will panic if called outside the context of a Tokio runtime. That means that you must /// call this on one of the threads **being run by the runtime**, or from a thread with an active /// `EnterGuard`. Calling this from within a thread created by `std::thread::spawn` (for example) /// will cause a panic unless that thread has an active `EnterGuard`. /// /// # Examples /// /// This can be used to obtain the handle of the surrounding runtime from an async /// block or function running on that runtime. /// /// ``` /// # use std::thread; /// # use tokio::runtime::Runtime; /// # fn dox() { /// # let rt = Runtime::new().unwrap(); /// # rt.spawn(async { /// use tokio::runtime::Handle; /// /// // Inside an async block or function. /// let handle = Handle::current(); /// handle.spawn(async { /// println!("now running in the existing Runtime"); /// }); /// /// # let handle = /// thread::spawn(move || { /// // Notice that the handle is created outside of this thread and then moved in /// handle.spawn(async { /* ... */ }); /// // This next line would cause a panic because we haven't entered the runtime /// // and created an EnterGuard /// // let handle2 = Handle::current(); // panic /// // So we create a guard here with Handle::enter(); /// let _guard = handle.enter(); /// // Now we can call Handle::current(); /// let handle2 = Handle::current(); /// }); /// # handle.join().unwrap(); /// # }); /// # } /// ``` #[track_caller] pub fn current() -> Self { Handle { inner: scheduler::Handle::current(), } } /// Returns a Handle view over the currently running Runtime /// /// Returns an error if no Runtime has been started /// /// Contrary to `current`, this never panics pub fn try_current() -> Result { context::with_current(|inner| Handle { inner: inner.clone(), }) } /// Spawns a future onto the Tokio runtime. /// /// This spawns the given future onto the runtime's executor, usually a /// thread pool. The thread pool is then responsible for polling the future /// until it completes. /// /// The provided future will start running in the background immediately /// when `spawn` is called, even if you don't await the returned /// `JoinHandle`. /// /// See [module level][mod] documentation for more details. /// /// [mod]: index.html /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// # fn dox() { /// // Create the runtime /// let rt = Runtime::new().unwrap(); /// // Get a handle from this runtime /// let handle = rt.handle(); /// /// // Spawn a future onto the runtime using the handle /// handle.spawn(async { /// println!("now running on a worker thread"); /// }); /// # } /// ``` #[track_caller] pub fn spawn(&self, future: F) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, { let fut_size = mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { self.spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { self.spawn_named(future, SpawnMeta::new_unnamed(fut_size)) } } /// Runs the provided function on an executor dedicated to blocking /// operations. /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// # fn dox() { /// // Create the runtime /// let rt = Runtime::new().unwrap(); /// // Get a handle from this runtime /// let handle = rt.handle(); /// /// // Spawn a blocking function onto the runtime using the handle /// handle.spawn_blocking(|| { /// println!("now running on a worker thread"); /// }); /// # } #[track_caller] pub fn spawn_blocking(&self, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { self.inner.blocking_spawner().spawn_blocking(self, func) } /// Runs a future to completion on this `Handle`'s associated `Runtime`. /// /// This runs the given future on the current thread, blocking until it is /// complete, and yielding its resolved result. Any tasks or timers which /// the future spawns internally will be executed on the runtime. /// /// When this is used on a `current_thread` runtime, only the /// [`Runtime::block_on`] method can drive the IO and timer drivers, but the /// `Handle::block_on` method cannot drive them. This means that, when using /// this method on a `current_thread` runtime, anything that relies on IO or /// timers will not work unless there is another thread currently calling /// [`Runtime::block_on`] on the same runtime. /// /// # If the runtime has been shut down /// /// If the `Handle`'s associated `Runtime` has been shut down (through /// [`Runtime::shutdown_background`], [`Runtime::shutdown_timeout`], or by /// dropping it) and `Handle::block_on` is used it might return an error or /// panic. Specifically IO resources will return an error and timers will /// panic. Runtime independent futures will run as normal. /// /// # Panics /// /// This function panics if the provided future panics, if called within an /// asynchronous execution context, or if a timer future is executed on a runtime that has been /// shut down. /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// // Create the runtime /// let rt = Runtime::new().unwrap(); /// /// // Get a handle from this runtime /// let handle = rt.handle(); /// /// // Execute the future, blocking the current thread until completion /// handle.block_on(async { /// println!("hello"); /// }); /// ``` /// /// Or using `Handle::current`: /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main () { /// let handle = Handle::current(); /// std::thread::spawn(move || { /// // Using Handle::block_on to run async code in the new thread. /// handle.block_on(async { /// println!("hello"); /// }); /// }); /// } /// ``` /// /// [`JoinError`]: struct@crate::task::JoinError /// [`JoinHandle`]: struct@crate::task::JoinHandle /// [`Runtime::block_on`]: fn@crate::runtime::Runtime::block_on /// [`Runtime::shutdown_background`]: fn@crate::runtime::Runtime::shutdown_background /// [`Runtime::shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout /// [`spawn_blocking`]: crate::task::spawn_blocking /// [`tokio::fs`]: crate::fs /// [`tokio::net`]: crate::net /// [`tokio::time`]: crate::time #[track_caller] pub fn block_on(&self, future: F) -> F::Output { let fut_size = mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { self.block_on_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { self.block_on_inner(future, SpawnMeta::new_unnamed(fut_size)) } } #[track_caller] fn block_on_inner(&self, future: F, _meta: SpawnMeta<'_>) -> F::Output { #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] let future = super::task::trace::Trace::root(future); #[cfg(all(tokio_unstable, feature = "tracing"))] let future = crate::util::trace::task(future, "block_on", _meta, super::task::Id::next().as_u64()); // Enter the runtime context. This sets the current driver handles and // prevents blocking an existing runtime. context::enter_runtime(&self.inner, true, |blocking| { blocking.block_on(future).expect("failed to park thread") }) } #[track_caller] pub(crate) fn spawn_named(&self, future: F, _meta: SpawnMeta<'_>) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, { let id = crate::runtime::task::Id::next(); #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] let future = super::task::trace::Trace::root(future); #[cfg(all(tokio_unstable, feature = "tracing"))] let future = crate::util::trace::task(future, "task", _meta, id.as_u64()); self.inner.spawn(future, id) } #[track_caller] #[allow(dead_code)] pub(crate) unsafe fn spawn_local_named( &self, future: F, _meta: SpawnMeta<'_>, ) -> JoinHandle where F: Future + 'static, F::Output: 'static, { let id = crate::runtime::task::Id::next(); #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] let future = super::task::trace::Trace::root(future); #[cfg(all(tokio_unstable, feature = "tracing"))] let future = crate::util::trace::task(future, "task", _meta, id.as_u64()); self.inner.spawn_local(future, id) } /// Returns the flavor of the current `Runtime`. /// /// # Examples /// /// ``` /// use tokio::runtime::{Handle, RuntimeFlavor}; /// /// #[tokio::main(flavor = "current_thread")] /// async fn main() { /// assert_eq!(RuntimeFlavor::CurrentThread, Handle::current().runtime_flavor()); /// } /// ``` /// /// ``` /// use tokio::runtime::{Handle, RuntimeFlavor}; /// /// #[tokio::main(flavor = "multi_thread", worker_threads = 4)] /// async fn main() { /// assert_eq!(RuntimeFlavor::MultiThread, Handle::current().runtime_flavor()); /// } /// ``` pub fn runtime_flavor(&self) -> RuntimeFlavor { match self.inner { scheduler::Handle::CurrentThread(_) => RuntimeFlavor::CurrentThread, #[cfg(feature = "rt-multi-thread")] scheduler::Handle::MultiThread(_) => RuntimeFlavor::MultiThread, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] scheduler::Handle::MultiThreadAlt(_) => RuntimeFlavor::MultiThreadAlt, } } cfg_unstable! { /// Returns the [`Id`] of the current `Runtime`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main(flavor = "current_thread")] /// async fn main() { /// println!("Current runtime id: {}", Handle::current().id()); /// } /// ``` /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features /// [`Id`]: struct@crate::runtime::Id pub fn id(&self) -> runtime::Id { let owned_id = match &self.inner { scheduler::Handle::CurrentThread(handle) => handle.owned_id(), #[cfg(feature = "rt-multi-thread")] scheduler::Handle::MultiThread(handle) => handle.owned_id(), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] scheduler::Handle::MultiThreadAlt(handle) => handle.owned_id(), }; owned_id.into() } } /// Returns a view that lets you get information about how the runtime /// is performing. pub fn metrics(&self) -> RuntimeMetrics { RuntimeMetrics::new(self.clone()) } } cfg_taskdump! { impl Handle { /// Captures a snapshot of the runtime's state. /// /// This functionality is experimental, and comes with a number of /// requirements and limitations. /// /// # Examples /// /// This can be used to get call traces of each task in the runtime. /// Calls to `Handle::dump` should usually be enclosed in a /// [timeout][crate::time::timeout], so that dumping does not escalate a /// single blocked runtime thread into an entirely blocked runtime. /// /// ``` /// # use tokio::runtime::Runtime; /// # fn dox() { /// # let rt = Runtime::new().unwrap(); /// # rt.spawn(async { /// use tokio::runtime::Handle; /// use tokio::time::{timeout, Duration}; /// /// // Inside an async block or function. /// let handle = Handle::current(); /// if let Ok(dump) = timeout(Duration::from_secs(2), handle.dump()).await { /// for (i, task) in dump.tasks().iter().enumerate() { /// let trace = task.trace(); /// println!("TASK {i}:"); /// println!("{trace}\n"); /// } /// } /// # }); /// # } /// ``` /// /// This produces highly detailed traces of tasks; e.g.: /// /// ```plain /// TASK 0: /// ╼ dump::main::{{closure}}::a::{{closure}} at /tokio/examples/dump.rs:18:20 /// └╼ dump::main::{{closure}}::b::{{closure}} at /tokio/examples/dump.rs:23:20 /// └╼ dump::main::{{closure}}::c::{{closure}} at /tokio/examples/dump.rs:28:24 /// └╼ tokio::sync::barrier::Barrier::wait::{{closure}} at /tokio/tokio/src/sync/barrier.rs:129:10 /// └╼ as core::future::future::Future>::poll at /tokio/tokio/src/util/trace.rs:77:46 /// └╼ tokio::sync::barrier::Barrier::wait_internal::{{closure}} at /tokio/tokio/src/sync/barrier.rs:183:36 /// └╼ tokio::sync::watch::Receiver::changed::{{closure}} at /tokio/tokio/src/sync/watch.rs:604:55 /// └╼ tokio::sync::watch::changed_impl::{{closure}} at /tokio/tokio/src/sync/watch.rs:755:18 /// └╼ ::poll at /tokio/tokio/src/sync/notify.rs:1103:9 /// └╼ tokio::sync::notify::Notified::poll_notified at /tokio/tokio/src/sync/notify.rs:996:32 /// ``` /// /// # Requirements /// /// ## Debug Info Must Be Available /// /// To produce task traces, the application must **not** be compiled /// with `split debuginfo`. On Linux, including `debuginfo` within the /// application binary is the (correct) default. You can further ensure /// this behavior with the following directive in your `Cargo.toml`: /// /// ```toml /// [profile.*] /// split-debuginfo = "off" /// ``` /// /// ## Unstable Features /// /// This functionality is **unstable**, and requires both the /// `tokio_unstable` and `tokio_taskdump` `cfg` flags to be set. /// /// You can do this by setting the `RUSTFLAGS` environment variable /// before invoking `cargo`; e.g.: /// ```bash /// RUSTFLAGS="--cfg tokio_unstable --cfg tokio_taskdump" cargo run --example dump /// ``` /// /// Or by [configuring][cargo-config] `rustflags` in /// `.cargo/config.toml`: /// ```text /// [build] /// rustflags = ["--cfg", "tokio_unstable", "--cfg", "tokio_taskdump"] /// ``` /// /// [cargo-config]: /// https://doc.rust-lang.org/cargo/reference/config.html /// /// ## Platform Requirements /// /// Task dumps are supported on Linux atop `aarch64`, `x86` and `x86_64`. /// /// ## Current Thread Runtime Requirements /// /// On the `current_thread` runtime, task dumps may only be requested /// from *within* the context of the runtime being dumped. Do not, for /// example, await `Handle::dump()` on a different runtime. /// /// # Limitations /// /// ## Performance /// /// Although enabling the `tokio_taskdump` feature imposes virtually no /// additional runtime overhead, actually calling `Handle::dump` is /// expensive. The runtime must synchronize and pause its workers, then /// re-poll every task in a special tracing mode. Avoid requesting dumps /// often. /// /// ## Local Executors /// /// Tasks managed by local executors (e.g., `FuturesUnordered` and /// [`LocalSet`][crate::task::LocalSet]) may not appear in task dumps. /// /// ## Non-Termination When Workers Are Blocked /// /// The future produced by `Handle::dump` may never produce `Ready` if /// another runtime worker is blocked for more than 250ms. This may /// occur if a dump is requested during shutdown, or if another runtime /// worker is infinite looping or synchronously deadlocked. For these /// reasons, task dumping should usually be paired with an explicit /// [timeout][crate::time::timeout]. pub async fn dump(&self) -> crate::runtime::Dump { match &self.inner { scheduler::Handle::CurrentThread(handle) => handle.dump(), #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] scheduler::Handle::MultiThread(handle) => { // perform the trace in a separate thread so that the // trace itself does not appear in the taskdump. let handle = handle.clone(); spawn_thread(async { let handle = handle; handle.dump().await }).await }, #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] scheduler::Handle::MultiThreadAlt(_) => panic!("task dump not implemented for this runtime flavor"), } } /// Produces `true` if the current task is being traced for a dump; /// otherwise false. This function is only public for integration /// testing purposes. Do not rely on it. #[doc(hidden)] pub fn is_tracing() -> bool { super::task::trace::Context::is_tracing() } } cfg_rt_multi_thread! { /// Spawn a new thread and asynchronously await on its result. async fn spawn_thread(f: F) -> ::Output where F: Future + Send + 'static, ::Output: Send + 'static { let (tx, rx) = crate::sync::oneshot::channel(); crate::loom::thread::spawn(|| { let rt = crate::runtime::Builder::new_current_thread().build().unwrap(); rt.block_on(async { let _ = tx.send(f.await); }); }); rx.await.unwrap() } } } /// Error returned by `try_current` when no Runtime has been started #[derive(Debug)] pub struct TryCurrentError { kind: TryCurrentErrorKind, } impl TryCurrentError { pub(crate) fn new_no_context() -> Self { Self { kind: TryCurrentErrorKind::NoContext, } } pub(crate) fn new_thread_local_destroyed() -> Self { Self { kind: TryCurrentErrorKind::ThreadLocalDestroyed, } } /// Returns true if the call failed because there is currently no runtime in /// the Tokio context. pub fn is_missing_context(&self) -> bool { matches!(self.kind, TryCurrentErrorKind::NoContext) } /// Returns true if the call failed because the Tokio context thread-local /// had been destroyed. This can usually only happen if in the destructor of /// other thread-locals. pub fn is_thread_local_destroyed(&self) -> bool { matches!(self.kind, TryCurrentErrorKind::ThreadLocalDestroyed) } } enum TryCurrentErrorKind { NoContext, ThreadLocalDestroyed, } impl fmt::Debug for TryCurrentErrorKind { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { TryCurrentErrorKind::NoContext => f.write_str("NoContext"), TryCurrentErrorKind::ThreadLocalDestroyed => f.write_str("ThreadLocalDestroyed"), } } } impl fmt::Display for TryCurrentError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { use TryCurrentErrorKind as E; match self.kind { E::NoContext => f.write_str(CONTEXT_MISSING_ERROR), E::ThreadLocalDestroyed => f.write_str(THREAD_LOCAL_DESTROYED_ERROR), } } } impl error::Error for TryCurrentError {} tokio-1.43.1/src/runtime/id.rs000064400000000000000000000027021046102023000142540ustar 00000000000000use std::fmt; use std::num::{NonZeroU32, NonZeroU64}; /// An opaque ID that uniquely identifies a runtime relative to all other currently /// running runtimes. /// /// # Notes /// /// - Runtime IDs are unique relative to other *currently running* runtimes. /// When a runtime completes, the same ID may be used for another runtime. /// - Runtime IDs are *not* sequential, and do not indicate the order in which /// runtimes are started or any other data. /// - The runtime ID of the currently running task can be obtained from the /// Handle. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main(flavor = "multi_thread", worker_threads = 4)] /// async fn main() { /// println!("Current runtime id: {}", Handle::current().id()); /// } /// ``` /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)] pub struct Id(NonZeroU64); impl From for Id { fn from(value: NonZeroU64) -> Self { Id(value) } } impl From for Id { fn from(value: NonZeroU32) -> Self { Id(value.into()) } } impl fmt::Display for Id { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.0.fmt(f) } } tokio-1.43.1/src/runtime/io/driver/signal.rs000064400000000000000000000007471046102023000170460ustar 00000000000000use super::{Driver, Handle, TOKEN_SIGNAL}; use std::io; impl Handle { pub(crate) fn register_signal_receiver( &self, receiver: &mut mio::net::UnixStream, ) -> io::Result<()> { self.registry .register(receiver, TOKEN_SIGNAL, mio::Interest::READABLE)?; Ok(()) } } impl Driver { pub(crate) fn consume_signal_ready(&mut self) -> bool { let ret = self.signal_ready; self.signal_ready = false; ret } } tokio-1.43.1/src/runtime/io/driver.rs000064400000000000000000000204771046102023000155730ustar 00000000000000// Signal handling cfg_signal_internal_and_unix! { mod signal; } use crate::io::interest::Interest; use crate::io::ready::Ready; use crate::loom::sync::Mutex; use crate::runtime::driver; use crate::runtime::io::registration_set; use crate::runtime::io::{IoDriverMetrics, RegistrationSet, ScheduledIo}; use mio::event::Source; use std::fmt; use std::io; use std::sync::Arc; use std::time::Duration; /// I/O driver, backed by Mio. pub(crate) struct Driver { /// True when an event with the signal token is received signal_ready: bool, /// Reuse the `mio::Events` value across calls to poll. events: mio::Events, /// The system event queue. poll: mio::Poll, } /// A reference to an I/O driver. pub(crate) struct Handle { /// Registers I/O resources. registry: mio::Registry, /// Tracks all registrations registrations: RegistrationSet, /// State that should be synchronized synced: Mutex, /// Used to wake up the reactor from a call to `turn`. /// Not supported on `Wasi` due to lack of threading support. #[cfg(not(target_os = "wasi"))] waker: mio::Waker, pub(crate) metrics: IoDriverMetrics, } #[derive(Debug)] pub(crate) struct ReadyEvent { pub(super) tick: u8, pub(crate) ready: Ready, pub(super) is_shutdown: bool, } cfg_net_unix!( impl ReadyEvent { pub(crate) fn with_ready(&self, ready: Ready) -> Self { Self { ready, tick: self.tick, is_shutdown: self.is_shutdown, } } } ); #[derive(Debug, Eq, PartialEq, Clone, Copy)] pub(super) enum Direction { Read, Write, } pub(super) enum Tick { Set, Clear(u8), } const TOKEN_WAKEUP: mio::Token = mio::Token(0); const TOKEN_SIGNAL: mio::Token = mio::Token(1); fn _assert_kinds() { fn _assert() {} _assert::(); } // ===== impl Driver ===== impl Driver { /// Creates a new event loop, returning any error that happened during the /// creation. pub(crate) fn new(nevents: usize) -> io::Result<(Driver, Handle)> { let poll = mio::Poll::new()?; #[cfg(not(target_os = "wasi"))] let waker = mio::Waker::new(poll.registry(), TOKEN_WAKEUP)?; let registry = poll.registry().try_clone()?; let driver = Driver { signal_ready: false, events: mio::Events::with_capacity(nevents), poll, }; let (registrations, synced) = RegistrationSet::new(); let handle = Handle { registry, registrations, synced: Mutex::new(synced), #[cfg(not(target_os = "wasi"))] waker, metrics: IoDriverMetrics::default(), }; Ok((driver, handle)) } pub(crate) fn park(&mut self, rt_handle: &driver::Handle) { let handle = rt_handle.io(); self.turn(handle, None); } pub(crate) fn park_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) { let handle = rt_handle.io(); self.turn(handle, Some(duration)); } pub(crate) fn shutdown(&mut self, rt_handle: &driver::Handle) { let handle = rt_handle.io(); let ios = handle.registrations.shutdown(&mut handle.synced.lock()); // `shutdown()` must be called without holding the lock. for io in ios { io.shutdown(); } } fn turn(&mut self, handle: &Handle, max_wait: Option) { debug_assert!(!handle.registrations.is_shutdown(&handle.synced.lock())); handle.release_pending_registrations(); let events = &mut self.events; // Block waiting for an event to happen, peeling out how many events // happened. match self.poll.poll(events, max_wait) { Ok(()) => {} Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {} #[cfg(target_os = "wasi")] Err(e) if e.kind() == io::ErrorKind::InvalidInput => { // In case of wasm32_wasi this error happens, when trying to poll without subscriptions // just return from the park, as there would be nothing, which wakes us up. } Err(e) => panic!("unexpected error when polling the I/O driver: {e:?}"), } // Process all the events that came in, dispatching appropriately let mut ready_count = 0; for event in events.iter() { let token = event.token(); if token == TOKEN_WAKEUP { // Nothing to do, the event is used to unblock the I/O driver } else if token == TOKEN_SIGNAL { self.signal_ready = true; } else { let ready = Ready::from_mio(event); let ptr = super::EXPOSE_IO.from_exposed_addr(token.0); // Safety: we ensure that the pointers used as tokens are not freed // until they are both deregistered from mio **and** we know the I/O // driver is not concurrently polling. The I/O driver holds ownership of // an `Arc` so we can safely cast this to a ref. let io: &ScheduledIo = unsafe { &*ptr }; io.set_readiness(Tick::Set, |curr| curr | ready); io.wake(ready); ready_count += 1; } } handle.metrics.incr_ready_count_by(ready_count); } } impl fmt::Debug for Driver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "Driver") } } impl Handle { /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise /// makes the next call to `turn` return immediately. /// /// This method is intended to be used in situations where a notification /// needs to otherwise be sent to the main reactor. If the reactor is /// currently blocked inside of `turn` then it will wake up and soon return /// after this method has been called. If the reactor is not currently /// blocked in `turn`, then the next call to `turn` will not block and /// return immediately. pub(crate) fn unpark(&self) { #[cfg(not(target_os = "wasi"))] self.waker.wake().expect("failed to wake I/O driver"); } /// Registers an I/O resource with the reactor for a given `mio::Ready` state. /// /// The registration token is returned. pub(super) fn add_source( &self, source: &mut impl mio::event::Source, interest: Interest, ) -> io::Result> { let scheduled_io = self.registrations.allocate(&mut self.synced.lock())?; let token = scheduled_io.token(); // we should remove the `scheduled_io` from the `registrations` set if registering // the `source` with the OS fails. Otherwise it will leak the `scheduled_io`. if let Err(e) = self.registry.register(source, token, interest.to_mio()) { // safety: `scheduled_io` is part of the `registrations` set. unsafe { self.registrations .remove(&mut self.synced.lock(), &scheduled_io) }; return Err(e); } // TODO: move this logic to `RegistrationSet` and use a `CountedLinkedList` self.metrics.incr_fd_count(); Ok(scheduled_io) } /// Deregisters an I/O resource from the reactor. pub(super) fn deregister_source( &self, registration: &Arc, source: &mut impl Source, ) -> io::Result<()> { // Deregister the source with the OS poller **first** self.registry.deregister(source)?; if self .registrations .deregister(&mut self.synced.lock(), registration) { self.unpark(); } self.metrics.dec_fd_count(); Ok(()) } fn release_pending_registrations(&self) { if self.registrations.needs_release() { self.registrations.release(&mut self.synced.lock()); } } } impl fmt::Debug for Handle { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "Handle") } } impl Direction { pub(super) fn mask(self) -> Ready { match self { Direction::Read => Ready::READABLE | Ready::READ_CLOSED, Direction::Write => Ready::WRITABLE | Ready::WRITE_CLOSED, } } } tokio-1.43.1/src/runtime/io/metrics.rs000064400000000000000000000012271046102023000157360ustar 00000000000000//! This file contains mocks of the metrics types used in the I/O driver. //! //! The reason these mocks don't live in `src/runtime/mock.rs` is because //! these need to be available in the case when `net` is enabled but //! `rt` is not. cfg_not_rt_and_metrics_and_net! { #[derive(Default)] pub(crate) struct IoDriverMetrics {} impl IoDriverMetrics { pub(crate) fn incr_fd_count(&self) {} pub(crate) fn dec_fd_count(&self) {} pub(crate) fn incr_ready_count_by(&self, _amt: u64) {} } } cfg_net! { cfg_rt! { cfg_unstable_metrics! { pub(crate) use crate::runtime::IoDriverMetrics; } } } tokio-1.43.1/src/runtime/io/mod.rs000064400000000000000000000007741046102023000150550ustar 00000000000000#![cfg_attr(not(all(feature = "rt", feature = "net")), allow(dead_code))] mod driver; use driver::{Direction, Tick}; pub(crate) use driver::{Driver, Handle, ReadyEvent}; mod registration; pub(crate) use registration::Registration; mod registration_set; use registration_set::RegistrationSet; mod scheduled_io; use scheduled_io::ScheduledIo; mod metrics; use metrics::IoDriverMetrics; use crate::util::ptr_expose::PtrExposeDomain; static EXPOSE_IO: PtrExposeDomain = PtrExposeDomain::new(); tokio-1.43.1/src/runtime/io/registration.rs000064400000000000000000000215441046102023000170060ustar 00000000000000#![cfg_attr(not(feature = "net"), allow(dead_code))] use crate::io::interest::Interest; use crate::runtime::io::{Direction, Handle, ReadyEvent, ScheduledIo}; use crate::runtime::scheduler; use mio::event::Source; use std::io; use std::sync::Arc; use std::task::{ready, Context, Poll}; cfg_io_driver! { /// Associates an I/O resource with the reactor instance that drives it. /// /// A registration represents an I/O resource registered with a Reactor such /// that it will receive task notifications on readiness. This is the lowest /// level API for integrating with a reactor. /// /// The association between an I/O resource is made by calling /// [`new_with_interest_and_handle`]. /// Once the association is established, it remains established until the /// registration instance is dropped. /// /// A registration instance represents two separate readiness streams. One /// for the read readiness and one for write readiness. These streams are /// independent and can be consumed from separate tasks. /// /// **Note**: while `Registration` is `Sync`, the caller must ensure that /// there are at most two tasks that use a registration instance /// concurrently. One task for [`poll_read_ready`] and one task for /// [`poll_write_ready`]. While violating this requirement is "safe" from a /// Rust memory safety point of view, it will result in unexpected behavior /// in the form of lost notifications and tasks hanging. /// /// ## Platform-specific events /// /// `Registration` also allows receiving platform-specific `mio::Ready` /// events. These events are included as part of the read readiness event /// stream. The write readiness event stream is only for `Ready::writable()` /// events. /// /// [`new_with_interest_and_handle`]: method@Self::new_with_interest_and_handle /// [`poll_read_ready`]: method@Self::poll_read_ready` /// [`poll_write_ready`]: method@Self::poll_write_ready` #[derive(Debug)] pub(crate) struct Registration { /// Handle to the associated runtime. /// /// TODO: this can probably be moved into `ScheduledIo`. handle: scheduler::Handle, /// Reference to state stored by the driver. shared: Arc, } } unsafe impl Send for Registration {} unsafe impl Sync for Registration {} // ===== impl Registration ===== impl Registration { /// Registers the I/O resource with the reactor for the provided handle, for /// a specific `Interest`. This does not add `hup` or `error` so if you are /// interested in those states, you will need to add them to the readiness /// state passed to this function. /// /// # Return /// /// - `Ok` if the registration happened successfully /// - `Err` if an error was encountered during registration #[track_caller] pub(crate) fn new_with_interest_and_handle( io: &mut impl Source, interest: Interest, handle: scheduler::Handle, ) -> io::Result { let shared = handle.driver().io().add_source(io, interest)?; Ok(Registration { handle, shared }) } /// Deregisters the I/O resource from the reactor it is associated with. /// /// This function must be called before the I/O resource associated with the /// registration is dropped. /// /// Note that deregistering does not guarantee that the I/O resource can be /// registered with a different reactor. Some I/O resource types can only be /// associated with a single reactor instance for their lifetime. /// /// # Return /// /// If the deregistration was successful, `Ok` is returned. Any calls to /// `Reactor::turn` that happen after a successful call to `deregister` will /// no longer result in notifications getting sent for this registration. /// /// `Err` is returned if an error is encountered. pub(crate) fn deregister(&mut self, io: &mut impl Source) -> io::Result<()> { self.handle().deregister_source(&self.shared, io) } pub(crate) fn clear_readiness(&self, event: ReadyEvent) { self.shared.clear_readiness(event); } // Uses the poll path, requiring the caller to ensure mutual exclusion for // correctness. Only the last task to call this function is notified. pub(crate) fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { self.poll_ready(cx, Direction::Read) } // Uses the poll path, requiring the caller to ensure mutual exclusion for // correctness. Only the last task to call this function is notified. pub(crate) fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { self.poll_ready(cx, Direction::Write) } // Uses the poll path, requiring the caller to ensure mutual exclusion for // correctness. Only the last task to call this function is notified. #[cfg(not(target_os = "wasi"))] pub(crate) fn poll_read_io( &self, cx: &mut Context<'_>, f: impl FnMut() -> io::Result, ) -> Poll> { self.poll_io(cx, Direction::Read, f) } // Uses the poll path, requiring the caller to ensure mutual exclusion for // correctness. Only the last task to call this function is notified. pub(crate) fn poll_write_io( &self, cx: &mut Context<'_>, f: impl FnMut() -> io::Result, ) -> Poll> { self.poll_io(cx, Direction::Write, f) } /// Polls for events on the I/O resource's `direction` readiness stream. /// /// If called with a task context, notify the task when a new event is /// received. fn poll_ready( &self, cx: &mut Context<'_>, direction: Direction, ) -> Poll> { ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ev = ready!(self.shared.poll_readiness(cx, direction)); if ev.is_shutdown { return Poll::Ready(Err(gone())); } coop.made_progress(); Poll::Ready(Ok(ev)) } fn poll_io( &self, cx: &mut Context<'_>, direction: Direction, mut f: impl FnMut() -> io::Result, ) -> Poll> { loop { let ev = ready!(self.poll_ready(cx, direction))?; match f() { Ok(ret) => { return Poll::Ready(Ok(ret)); } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.clear_readiness(ev); } Err(e) => return Poll::Ready(Err(e)), } } } pub(crate) fn try_io( &self, interest: Interest, f: impl FnOnce() -> io::Result, ) -> io::Result { let ev = self.shared.ready_event(interest); // Don't attempt the operation if the resource is not ready. if ev.ready.is_empty() { return Err(io::ErrorKind::WouldBlock.into()); } match f() { Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.clear_readiness(ev); Err(io::ErrorKind::WouldBlock.into()) } res => res, } } pub(crate) async fn readiness(&self, interest: Interest) -> io::Result { let ev = self.shared.readiness(interest).await; if ev.is_shutdown { return Err(gone()); } Ok(ev) } pub(crate) async fn async_io( &self, interest: Interest, mut f: impl FnMut() -> io::Result, ) -> io::Result { loop { let event = self.readiness(interest).await?; let coop = std::future::poll_fn(crate::runtime::coop::poll_proceed).await; match f() { Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.clear_readiness(event); } x => { coop.made_progress(); return x; } } } } fn handle(&self) -> &Handle { self.handle.driver().io() } } impl Drop for Registration { fn drop(&mut self) { // It is possible for a cycle to be created between wakers stored in // `ScheduledIo` instances and `Arc`. To break this // cycle, wakers are cleared. This is an imperfect solution as it is // possible to store a `Registration` in a waker. In this case, the // cycle would remain. // // See tokio-rs/tokio#3481 for more details. self.shared.clear_wakers(); } } fn gone() -> io::Error { io::Error::new( io::ErrorKind::Other, crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR, ) } tokio-1.43.1/src/runtime/io/registration_set.rs000064400000000000000000000113021046102023000176500ustar 00000000000000use crate::loom::sync::atomic::AtomicUsize; use crate::runtime::io::ScheduledIo; use crate::util::linked_list::{self, LinkedList}; use std::io; use std::ptr::NonNull; use std::sync::atomic::Ordering::{Acquire, Release}; use std::sync::Arc; pub(super) struct RegistrationSet { num_pending_release: AtomicUsize, } pub(super) struct Synced { // True when the I/O driver shutdown. At this point, no more registrations // should be added to the set. is_shutdown: bool, // List of all registrations tracked by the set registrations: LinkedList, ScheduledIo>, // Registrations that are pending drop. When a `Registration` is dropped, it // stores its `ScheduledIo` in this list. The I/O driver is responsible for // dropping it. This ensures the `ScheduledIo` is not freed while it can // still be included in an I/O event. pending_release: Vec>, } impl RegistrationSet { pub(super) fn new() -> (RegistrationSet, Synced) { let set = RegistrationSet { num_pending_release: AtomicUsize::new(0), }; let synced = Synced { is_shutdown: false, registrations: LinkedList::new(), pending_release: Vec::with_capacity(16), }; (set, synced) } pub(super) fn is_shutdown(&self, synced: &Synced) -> bool { synced.is_shutdown } /// Returns `true` if there are registrations that need to be released pub(super) fn needs_release(&self) -> bool { self.num_pending_release.load(Acquire) != 0 } pub(super) fn allocate(&self, synced: &mut Synced) -> io::Result> { if synced.is_shutdown { return Err(io::Error::new( io::ErrorKind::Other, crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR, )); } let ret = Arc::new(ScheduledIo::default()); // Push a ref into the list of all resources. synced.registrations.push_front(ret.clone()); Ok(ret) } // Returns `true` if the caller should unblock the I/O driver to purge // registrations pending release. pub(super) fn deregister(&self, synced: &mut Synced, registration: &Arc) -> bool { // Kind of arbitrary, but buffering 16 `ScheduledIo`s doesn't seem like much const NOTIFY_AFTER: usize = 16; synced.pending_release.push(registration.clone()); let len = synced.pending_release.len(); self.num_pending_release.store(len, Release); len == NOTIFY_AFTER } pub(super) fn shutdown(&self, synced: &mut Synced) -> Vec> { if synced.is_shutdown { return vec![]; } synced.is_shutdown = true; synced.pending_release.clear(); // Building a vec of all outstanding I/O handles could be expensive, but // this is the shutdown operation. In theory, shutdowns should be // "clean" with no outstanding I/O resources. Even if it is slow, we // aren't optimizing for shutdown. let mut ret = vec![]; while let Some(io) = synced.registrations.pop_back() { ret.push(io); } ret } pub(super) fn release(&self, synced: &mut Synced) { let pending = std::mem::take(&mut synced.pending_release); for io in pending { // safety: the registration is part of our list unsafe { self.remove(synced, &io) } } self.num_pending_release.store(0, Release); } // This function is marked as unsafe, because the caller must make sure that // `io` is part of the registration set. pub(super) unsafe fn remove(&self, synced: &mut Synced, io: &Arc) { // SAFETY: Pointers into an Arc are never null. let io = unsafe { NonNull::new_unchecked(Arc::as_ptr(io).cast_mut()) }; super::EXPOSE_IO.unexpose_provenance(io.as_ptr()); let _ = synced.registrations.remove(io); } } // Safety: `Arc` pins the inner data unsafe impl linked_list::Link for Arc { type Handle = Arc; type Target = ScheduledIo; fn as_raw(handle: &Self::Handle) -> NonNull { // safety: Arc::as_ptr never returns null unsafe { NonNull::new_unchecked(Arc::as_ptr(handle) as *mut _) } } unsafe fn from_raw(ptr: NonNull) -> Arc { // safety: the linked list currently owns a ref count unsafe { Arc::from_raw(ptr.as_ptr() as *const _) } } unsafe fn pointers( target: NonNull, ) -> NonNull> { NonNull::new_unchecked(target.as_ref().linked_list_pointers.get()) } } tokio-1.43.1/src/runtime/io/scheduled_io.rs000064400000000000000000000464671046102023000167360ustar 00000000000000use crate::io::interest::Interest; use crate::io::ready::Ready; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Mutex; use crate::runtime::io::{Direction, ReadyEvent, Tick}; use crate::util::bit; use crate::util::linked_list::{self, LinkedList}; use crate::util::WakeList; use std::cell::UnsafeCell; use std::future::Future; use std::marker::PhantomPinned; use std::pin::Pin; use std::ptr::NonNull; use std::sync::atomic::Ordering::{AcqRel, Acquire}; use std::task::{Context, Poll, Waker}; /// Stored in the I/O driver resource slab. #[derive(Debug)] // # This struct should be cache padded to avoid false sharing. The cache padding rules are copied // from crossbeam-utils/src/cache_padded.rs // // Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache // lines at a time, so we have to align to 128 bytes rather than 64. // // Sources: // - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf // - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107 // // ARM's big.LITTLE architecture has asymmetric cores and "big" cores have 128-byte cache line size. // // Sources: // - https://www.mono-project.com/news/2016/09/12/arm64-icache/ // // powerpc64 has 128-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_ppc64x.go#L9 #[cfg_attr( any( target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", ), repr(align(128)) )] // arm, mips, mips64, sparc, and hexagon have 32-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_arm.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mipsle.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips64x.go#L9 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L17 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/hexagon/include/asm/cache.h#L12 #[cfg_attr( any( target_arch = "arm", target_arch = "mips", target_arch = "mips64", target_arch = "sparc", target_arch = "hexagon", ), repr(align(32)) )] // m68k has 16-byte cache line size. // // Sources: // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/m68k/include/asm/cache.h#L9 #[cfg_attr(target_arch = "m68k", repr(align(16)))] // s390x has 256-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_s390x.go#L7 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/s390/include/asm/cache.h#L13 #[cfg_attr(target_arch = "s390x", repr(align(256)))] // x86, riscv, wasm, and sparc64 have 64-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/dda2991c2ea0c5914714469c4defc2562a907230/src/internal/cpu/cpu_x86.go#L9 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_wasm.go#L7 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L19 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/riscv/include/asm/cache.h#L10 // // All others are assumed to have 64-byte cache line size. #[cfg_attr( not(any( target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", target_arch = "arm", target_arch = "mips", target_arch = "mips64", target_arch = "sparc", target_arch = "hexagon", target_arch = "m68k", target_arch = "s390x", )), repr(align(64)) )] pub(crate) struct ScheduledIo { pub(super) linked_list_pointers: UnsafeCell>, /// Packs the resource's readiness and I/O driver latest tick. readiness: AtomicUsize, waiters: Mutex, } type WaitList = LinkedList::Target>; #[derive(Debug, Default)] struct Waiters { /// List of all current waiters. list: WaitList, /// Waker used for `AsyncRead`. reader: Option, /// Waker used for `AsyncWrite`. writer: Option, } #[derive(Debug)] struct Waiter { pointers: linked_list::Pointers, /// The waker for this task. waker: Option, /// The interest this waiter is waiting on. interest: Interest, is_ready: bool, /// Should never be `!Unpin`. _p: PhantomPinned, } generate_addr_of_methods! { impl<> Waiter { unsafe fn addr_of_pointers(self: NonNull) -> NonNull> { &self.pointers } } } /// Future returned by `readiness()`. struct Readiness<'a> { scheduled_io: &'a ScheduledIo, state: State, /// Entry in the waiter `LinkedList`. waiter: UnsafeCell, } enum State { Init, Waiting, Done, } // The `ScheduledIo::readiness` (`AtomicUsize`) is packed full of goodness. // // | shutdown | driver tick | readiness | // |----------+-------------+-----------| // | 1 bit | 15 bits + 16 bits | const READINESS: bit::Pack = bit::Pack::least_significant(16); const TICK: bit::Pack = READINESS.then(15); const SHUTDOWN: bit::Pack = TICK.then(1); // ===== impl ScheduledIo ===== impl Default for ScheduledIo { fn default() -> ScheduledIo { ScheduledIo { linked_list_pointers: UnsafeCell::new(linked_list::Pointers::new()), readiness: AtomicUsize::new(0), waiters: Mutex::new(Waiters::default()), } } } impl ScheduledIo { pub(crate) fn token(&self) -> mio::Token { mio::Token(super::EXPOSE_IO.expose_provenance(self)) } /// Invoked when the IO driver is shut down; forces this `ScheduledIo` into a /// permanently shutdown state. pub(super) fn shutdown(&self) { let mask = SHUTDOWN.pack(1, 0); self.readiness.fetch_or(mask, AcqRel); self.wake(Ready::ALL); } /// Sets the readiness on this `ScheduledIo` by invoking the given closure on /// the current value, returning the previous readiness value. /// /// # Arguments /// - `tick`: whether setting the tick or trying to clear readiness for a /// specific tick. /// - `f`: a closure returning a new readiness value given the previous /// readiness. pub(super) fn set_readiness(&self, tick_op: Tick, f: impl Fn(Ready) -> Ready) { let _ = self.readiness.fetch_update(AcqRel, Acquire, |curr| { // If the io driver is shut down, then you are only allowed to clear readiness. debug_assert!(SHUTDOWN.unpack(curr) == 0 || matches!(tick_op, Tick::Clear(_))); const MAX_TICK: usize = TICK.max_value() + 1; let tick = TICK.unpack(curr); let new_tick = match tick_op { // Trying to clear readiness with an old event! Tick::Clear(t) if tick as u8 != t => return None, Tick::Clear(t) => t as usize, Tick::Set => tick.wrapping_add(1) % MAX_TICK, }; let ready = Ready::from_usize(READINESS.unpack(curr)); Some(TICK.pack(new_tick, f(ready).as_usize())) }); } /// Notifies all pending waiters that have registered interest in `ready`. /// /// There may be many waiters to notify. Waking the pending task **must** be /// done from outside of the lock otherwise there is a potential for a /// deadlock. /// /// A stack array of wakers is created and filled with wakers to notify, the /// lock is released, and the wakers are notified. Because there may be more /// than 32 wakers to notify, if the stack array fills up, the lock is /// released, the array is cleared, and the iteration continues. pub(super) fn wake(&self, ready: Ready) { let mut wakers = WakeList::new(); let mut waiters = self.waiters.lock(); // check for AsyncRead slot if ready.is_readable() { if let Some(waker) = waiters.reader.take() { wakers.push(waker); } } // check for AsyncWrite slot if ready.is_writable() { if let Some(waker) = waiters.writer.take() { wakers.push(waker); } } 'outer: loop { let mut iter = waiters.list.drain_filter(|w| ready.satisfies(w.interest)); while wakers.can_push() { match iter.next() { Some(waiter) => { let waiter = unsafe { &mut *waiter.as_ptr() }; if let Some(waker) = waiter.waker.take() { waiter.is_ready = true; wakers.push(waker); } } None => { break 'outer; } } } drop(waiters); wakers.wake_all(); // Acquire the lock again. waiters = self.waiters.lock(); } // Release the lock before notifying drop(waiters); wakers.wake_all(); } pub(super) fn ready_event(&self, interest: Interest) -> ReadyEvent { let curr = self.readiness.load(Acquire); ReadyEvent { tick: TICK.unpack(curr) as u8, ready: interest.mask() & Ready::from_usize(READINESS.unpack(curr)), is_shutdown: SHUTDOWN.unpack(curr) != 0, } } /// Polls for readiness events in a given direction. /// /// These are to support `AsyncRead` and `AsyncWrite` polling methods, /// which cannot use the `async fn` version. This uses reserved reader /// and writer slots. pub(super) fn poll_readiness( &self, cx: &mut Context<'_>, direction: Direction, ) -> Poll { let curr = self.readiness.load(Acquire); let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr)); let is_shutdown = SHUTDOWN.unpack(curr) != 0; if ready.is_empty() && !is_shutdown { // Update the task info let mut waiters = self.waiters.lock(); let waker = match direction { Direction::Read => &mut waiters.reader, Direction::Write => &mut waiters.writer, }; // Avoid cloning the waker if one is already stored that matches the // current task. match waker { Some(waker) => waker.clone_from(cx.waker()), None => *waker = Some(cx.waker().clone()), } // Try again, in case the readiness was changed while we were // taking the waiters lock let curr = self.readiness.load(Acquire); let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr)); let is_shutdown = SHUTDOWN.unpack(curr) != 0; if is_shutdown { Poll::Ready(ReadyEvent { tick: TICK.unpack(curr) as u8, ready: direction.mask(), is_shutdown, }) } else if ready.is_empty() { Poll::Pending } else { Poll::Ready(ReadyEvent { tick: TICK.unpack(curr) as u8, ready, is_shutdown, }) } } else { Poll::Ready(ReadyEvent { tick: TICK.unpack(curr) as u8, ready, is_shutdown, }) } } pub(crate) fn clear_readiness(&self, event: ReadyEvent) { // This consumes the current readiness state **except** for closed // states. Closed states are excluded because they are final states. let mask_no_closed = event.ready - Ready::READ_CLOSED - Ready::WRITE_CLOSED; self.set_readiness(Tick::Clear(event.tick), |curr| curr - mask_no_closed); } pub(crate) fn clear_wakers(&self) { let mut waiters = self.waiters.lock(); waiters.reader.take(); waiters.writer.take(); } } impl Drop for ScheduledIo { fn drop(&mut self) { self.wake(Ready::ALL); } } unsafe impl Send for ScheduledIo {} unsafe impl Sync for ScheduledIo {} impl ScheduledIo { /// An async version of `poll_readiness` which uses a linked list of wakers. pub(crate) async fn readiness(&self, interest: Interest) -> ReadyEvent { self.readiness_fut(interest).await } // This is in a separate function so that the borrow checker doesn't think // we are borrowing the `UnsafeCell` possibly over await boundaries. // // Go figure. fn readiness_fut(&self, interest: Interest) -> Readiness<'_> { Readiness { scheduled_io: self, state: State::Init, waiter: UnsafeCell::new(Waiter { pointers: linked_list::Pointers::new(), waker: None, is_ready: false, interest, _p: PhantomPinned, }), } } } unsafe impl linked_list::Link for Waiter { type Handle = NonNull; type Target = Waiter; fn as_raw(handle: &NonNull) -> NonNull { *handle } unsafe fn from_raw(ptr: NonNull) -> NonNull { ptr } unsafe fn pointers(target: NonNull) -> NonNull> { Waiter::addr_of_pointers(target) } } // ===== impl Readiness ===== impl Future for Readiness<'_> { type Output = ReadyEvent; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { use std::sync::atomic::Ordering::SeqCst; let (scheduled_io, state, waiter) = unsafe { let me = self.get_unchecked_mut(); (&me.scheduled_io, &mut me.state, &me.waiter) }; loop { match *state { State::Init => { // Optimistically check existing readiness let curr = scheduled_io.readiness.load(SeqCst); let is_shutdown = SHUTDOWN.unpack(curr) != 0; // Safety: `waiter.interest` never changes let interest = unsafe { (*waiter.get()).interest }; let ready = Ready::from_usize(READINESS.unpack(curr)).intersection(interest); if !ready.is_empty() || is_shutdown { // Currently ready! let tick = TICK.unpack(curr) as u8; *state = State::Done; return Poll::Ready(ReadyEvent { tick, ready, is_shutdown, }); } // Wasn't ready, take the lock (and check again while locked). let mut waiters = scheduled_io.waiters.lock(); let curr = scheduled_io.readiness.load(SeqCst); let mut ready = Ready::from_usize(READINESS.unpack(curr)); let is_shutdown = SHUTDOWN.unpack(curr) != 0; if is_shutdown { ready = Ready::ALL; } let ready = ready.intersection(interest); if !ready.is_empty() || is_shutdown { // Currently ready! let tick = TICK.unpack(curr) as u8; *state = State::Done; return Poll::Ready(ReadyEvent { tick, ready, is_shutdown, }); } // Not ready even after locked, insert into list... // Safety: called while locked unsafe { (*waiter.get()).waker = Some(cx.waker().clone()); } // Insert the waiter into the linked list // // safety: pointers from `UnsafeCell` are never null. waiters .list .push_front(unsafe { NonNull::new_unchecked(waiter.get()) }); *state = State::Waiting; } State::Waiting => { // Currently in the "Waiting" state, implying the caller has // a waiter stored in the waiter list (guarded by // `notify.waiters`). In order to access the waker fields, // we must hold the lock. let waiters = scheduled_io.waiters.lock(); // Safety: called while locked let w = unsafe { &mut *waiter.get() }; if w.is_ready { // Our waker has been notified. *state = State::Done; } else { // Update the waker, if necessary. w.waker.as_mut().unwrap().clone_from(cx.waker()); return Poll::Pending; } // Explicit drop of the lock to indicate the scope that the // lock is held. Because holding the lock is required to // ensure safe access to fields not held within the lock, it // is helpful to visualize the scope of the critical // section. drop(waiters); } State::Done => { // Safety: State::Done means it is no longer shared let w = unsafe { &mut *waiter.get() }; let curr = scheduled_io.readiness.load(Acquire); let is_shutdown = SHUTDOWN.unpack(curr) != 0; // The returned tick might be newer than the event // which notified our waker. This is ok because the future // still didn't return `Poll::Ready`. let tick = TICK.unpack(curr) as u8; // The readiness state could have been cleared in the meantime, // but we allow the returned ready set to be empty. let ready = Ready::from_usize(READINESS.unpack(curr)).intersection(w.interest); return Poll::Ready(ReadyEvent { tick, ready, is_shutdown, }); } } } } } impl Drop for Readiness<'_> { fn drop(&mut self) { let mut waiters = self.scheduled_io.waiters.lock(); // Safety: `waiter` is only ever stored in `waiters` unsafe { waiters .list .remove(NonNull::new_unchecked(self.waiter.get())) }; } } unsafe impl Send for Readiness<'_> {} unsafe impl Sync for Readiness<'_> {} tokio-1.43.1/src/runtime/local_runtime/mod.rs000064400000000000000000000002111046102023000172650ustar 00000000000000mod runtime; mod options; pub use options::LocalOptions; pub use runtime::LocalRuntime; pub(super) use runtime::LocalRuntimeScheduler; tokio-1.43.1/src/runtime/local_runtime/options.rs000064400000000000000000000005311046102023000202060ustar 00000000000000use std::marker::PhantomData; /// `LocalRuntime`-only config options /// /// Currently, there are no such options, but in the future, things like `!Send + !Sync` hooks may /// be added. #[derive(Default, Debug)] #[non_exhaustive] pub struct LocalOptions { /// Marker used to make this !Send and !Sync. _phantom: PhantomData<*mut u8>, } tokio-1.43.1/src/runtime/local_runtime/runtime.rs000064400000000000000000000311021046102023000201740ustar 00000000000000#![allow(irrefutable_let_patterns)] use crate::runtime::blocking::BlockingPool; use crate::runtime::scheduler::CurrentThread; use crate::runtime::{context, Builder, EnterGuard, Handle, BOX_FUTURE_THRESHOLD}; use crate::task::JoinHandle; use crate::util::trace::SpawnMeta; use std::future::Future; use std::marker::PhantomData; use std::mem; use std::time::Duration; /// A local Tokio runtime. /// /// This runtime is capable of driving tasks which are not `Send + Sync` without the use of a /// `LocalSet`, and thus supports `spawn_local` without the need for a `LocalSet` context. /// /// This runtime cannot be moved between threads or driven from different threads. /// /// This runtime is incompatible with `LocalSet`. You should not attempt to drive a `LocalSet` within a /// `LocalRuntime`. /// /// Currently, this runtime supports one flavor, which is internally identical to `current_thread`, /// save for the aforementioned differences related to `spawn_local`. /// /// For more general information on how to use runtimes, see the [module] docs. /// /// [runtime]: crate::runtime::Runtime /// [module]: crate::runtime #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub struct LocalRuntime { /// Task scheduler scheduler: LocalRuntimeScheduler, /// Handle to runtime, also contains driver handles handle: Handle, /// Blocking pool handle, used to signal shutdown blocking_pool: BlockingPool, /// Marker used to make this !Send and !Sync. _phantom: PhantomData<*mut u8>, } /// The runtime scheduler is always a `current_thread` scheduler right now. #[derive(Debug)] pub(crate) enum LocalRuntimeScheduler { /// Execute all tasks on the current-thread. CurrentThread(CurrentThread), } impl LocalRuntime { pub(crate) fn from_parts( scheduler: LocalRuntimeScheduler, handle: Handle, blocking_pool: BlockingPool, ) -> LocalRuntime { LocalRuntime { scheduler, handle, blocking_pool, _phantom: Default::default(), } } /// Creates a new local runtime instance with default configuration values. /// /// This results in the scheduler, I/O driver, and time driver being /// initialized. /// /// When a more complex configuration is necessary, the [runtime builder] may be used. /// /// See [module level][mod] documentation for more details. /// /// # Examples /// /// Creating a new `LocalRuntime` with default configuration values. /// /// ``` /// use tokio::runtime::LocalRuntime; /// /// let rt = LocalRuntime::new() /// .unwrap(); /// /// // Use the runtime... /// ``` /// /// [mod]: crate::runtime /// [runtime builder]: crate::runtime::Builder pub fn new() -> std::io::Result { Builder::new_current_thread() .enable_all() .build_local(&Default::default()) } /// Returns a handle to the runtime's spawner. /// /// The returned handle can be used to spawn tasks that run on this runtime, and can /// be cloned to allow moving the `Handle` to other threads. /// /// As the handle can be sent to other threads, it can only be used to spawn tasks that are `Send`. /// /// Calling [`Handle::block_on`] on a handle to a `LocalRuntime` is error-prone. /// Refer to the documentation of [`Handle::block_on`] for more. /// /// # Examples /// /// ``` /// use tokio::runtime::LocalRuntime; /// /// let rt = LocalRuntime::new() /// .unwrap(); /// /// let handle = rt.handle(); /// /// // Use the handle... /// ``` pub fn handle(&self) -> &Handle { &self.handle } /// Spawns a task on the runtime. /// /// This is analogous to the [`spawn`] method on the standard [`Runtime`], but works even if the task is not thread-safe. /// /// [`spawn`]: crate::runtime::Runtime::spawn /// [`Runtime`]: crate::runtime::Runtime /// /// # Examples /// /// ``` /// use tokio::runtime::LocalRuntime; /// /// # fn dox() { /// // Create the runtime /// let rt = LocalRuntime::new().unwrap(); /// /// // Spawn a future onto the runtime /// rt.spawn_local(async { /// println!("now running on a worker thread"); /// }); /// # } /// ``` #[track_caller] pub fn spawn_local(&self, future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, { let fut_size = std::mem::size_of::(); let meta = SpawnMeta::new_unnamed(fut_size); // safety: spawn_local can only be called from `LocalRuntime`, which this is unsafe { if std::mem::size_of::() > BOX_FUTURE_THRESHOLD { self.handle.spawn_local_named(Box::pin(future), meta) } else { self.handle.spawn_local_named(future, meta) } } } /// Runs the provided function on a thread from a dedicated blocking thread pool. /// /// This function _will_ be run on another thread. /// /// See the [documentation in the non-local runtime][Runtime] for more /// information. /// /// [Runtime]: crate::runtime::Runtime::spawn_blocking /// /// # Examples /// /// ``` /// use tokio::runtime::LocalRuntime; /// /// # fn dox() { /// // Create the runtime /// let rt = LocalRuntime::new().unwrap(); /// /// // Spawn a blocking function onto the runtime /// rt.spawn_blocking(|| { /// println!("now running on a worker thread"); /// }); /// # } /// ``` #[track_caller] pub fn spawn_blocking(&self, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { self.handle.spawn_blocking(func) } /// Runs a future to completion on the Tokio runtime. This is the /// runtime's entry point. /// /// See the documentation for [the equivalent method on Runtime][Runtime] /// for more information. /// /// [Runtime]: crate::runtime::Runtime::block_on /// /// # Examples /// /// ```no_run /// use tokio::runtime::LocalRuntime; /// /// // Create the runtime /// let rt = LocalRuntime::new().unwrap(); /// /// // Execute the future, blocking the current thread until completion /// rt.block_on(async { /// println!("hello"); /// }); /// ``` #[track_caller] pub fn block_on(&self, future: F) -> F::Output { let fut_size = mem::size_of::(); let meta = SpawnMeta::new_unnamed(fut_size); if std::mem::size_of::() > BOX_FUTURE_THRESHOLD { self.block_on_inner(Box::pin(future), meta) } else { self.block_on_inner(future, meta) } } #[track_caller] fn block_on_inner(&self, future: F, _meta: SpawnMeta<'_>) -> F::Output { #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] let future = crate::runtime::task::trace::Trace::root(future); #[cfg(all(tokio_unstable, feature = "tracing"))] let future = crate::util::trace::task( future, "block_on", _meta, crate::runtime::task::Id::next().as_u64(), ); let _enter = self.enter(); if let LocalRuntimeScheduler::CurrentThread(exec) = &self.scheduler { exec.block_on(&self.handle.inner, future) } else { unreachable!("LocalRuntime only supports current_thread") } } /// Enters the runtime context. /// /// This allows you to construct types that must have an executor /// available on creation such as [`Sleep`] or [`TcpStream`]. It will /// also allow you to call methods such as [`tokio::spawn`]. /// /// If this is a handle to a [`LocalRuntime`], and this function is being invoked from the same /// thread that the runtime was created on, you will also be able to call /// [`tokio::task::spawn_local`]. /// /// [`Sleep`]: struct@crate::time::Sleep /// [`TcpStream`]: struct@crate::net::TcpStream /// [`tokio::spawn`]: fn@crate::spawn /// [`LocalRuntime`]: struct@crate::runtime::LocalRuntime /// [`tokio::task::spawn_local`]: fn@crate::task::spawn_local /// /// # Example /// /// ``` /// use tokio::runtime::LocalRuntime; /// use tokio::task::JoinHandle; /// /// fn function_that_spawns(msg: String) -> JoinHandle<()> { /// // Had we not used `rt.enter` below, this would panic. /// tokio::spawn(async move { /// println!("{}", msg); /// }) /// } /// /// fn main() { /// let rt = LocalRuntime::new().unwrap(); /// /// let s = "Hello World!".to_string(); /// /// // By entering the context, we tie `tokio::spawn` to this executor. /// let _guard = rt.enter(); /// let handle = function_that_spawns(s); /// /// // Wait for the task before we end the test. /// rt.block_on(handle).unwrap(); /// } /// ``` pub fn enter(&self) -> EnterGuard<'_> { self.handle.enter() } /// Shuts down the runtime, waiting for at most `duration` for all spawned /// work to stop. /// /// Note that `spawn_blocking` tasks, and only `spawn_blocking` tasks, can get left behind if /// the timeout expires. /// /// See the [struct level documentation](LocalRuntime#shutdown) for more details. /// /// # Examples /// /// ``` /// use tokio::runtime::LocalRuntime; /// use tokio::task; /// /// use std::thread; /// use std::time::Duration; /// /// fn main() { /// let runtime = LocalRuntime::new().unwrap(); /// /// runtime.block_on(async move { /// task::spawn_blocking(move || { /// thread::sleep(Duration::from_secs(10_000)); /// }); /// }); /// /// runtime.shutdown_timeout(Duration::from_millis(100)); /// } /// ``` pub fn shutdown_timeout(mut self, duration: Duration) { // Wakeup and shutdown all the worker threads self.handle.inner.shutdown(); self.blocking_pool.shutdown(Some(duration)); } /// Shuts down the runtime, without waiting for any spawned work to stop. /// /// This can be useful if you want to drop a runtime from within another runtime. /// Normally, dropping a runtime will block indefinitely for spawned blocking tasks /// to complete, which would normally not be permitted within an asynchronous context. /// By calling `shutdown_background()`, you can drop the runtime from such a context. /// /// Note however, that because we do not wait for any blocking tasks to complete, this /// may result in a resource leak (in that any blocking tasks are still running until they /// return. No other tasks will leak. /// /// See the [struct level documentation](LocalRuntime#shutdown) for more details. /// /// This function is equivalent to calling `shutdown_timeout(Duration::from_nanos(0))`. /// /// ``` /// use tokio::runtime::LocalRuntime; /// /// fn main() { /// let runtime = LocalRuntime::new().unwrap(); /// /// runtime.block_on(async move { /// let inner_runtime = LocalRuntime::new().unwrap(); /// // ... /// inner_runtime.shutdown_background(); /// }); /// } /// ``` pub fn shutdown_background(self) { self.shutdown_timeout(Duration::from_nanos(0)); } /// Returns a view that lets you get information about how the runtime /// is performing. pub fn metrics(&self) -> crate::runtime::RuntimeMetrics { self.handle.metrics() } } #[allow(clippy::single_match)] // there are comments in the error branch, so we don't want if-let impl Drop for LocalRuntime { fn drop(&mut self) { if let LocalRuntimeScheduler::CurrentThread(current_thread) = &mut self.scheduler { // This ensures that tasks spawned on the current-thread // runtime are dropped inside the runtime's context. let _guard = context::try_set_current(&self.handle.inner); current_thread.shutdown(&self.handle.inner); } else { unreachable!("LocalRuntime only supports current-thread") } } } impl std::panic::UnwindSafe for LocalRuntime {} impl std::panic::RefUnwindSafe for LocalRuntime {} tokio-1.43.1/src/runtime/metrics/batch.rs000064400000000000000000000123761046102023000164170ustar 00000000000000use crate::runtime::metrics::{HistogramBatch, WorkerMetrics}; use std::sync::atomic::Ordering::Relaxed; use std::time::{Duration, Instant}; pub(crate) struct MetricsBatch { /// Number of times the worker parked. park_count: u64, /// Number of times the worker parked and unparked. park_unpark_count: u64, /// Number of times the worker woke w/o doing work. noop_count: u64, /// Number of tasks stolen. steal_count: u64, /// Number of times tasks where stolen. steal_operations: u64, /// Number of tasks that were polled by the worker. poll_count: u64, /// Number of tasks polled when the worker entered park. This is used to /// track the noop count. poll_count_on_last_park: u64, /// Number of tasks that were scheduled locally on this worker. local_schedule_count: u64, /// Number of tasks moved to the global queue to make space in the local /// queue overflow_count: u64, /// The total busy duration in nanoseconds. busy_duration_total: u64, /// Instant at which work last resumed (continued after park). processing_scheduled_tasks_started_at: Instant, /// If `Some`, tracks poll times in nanoseconds poll_timer: Option, } struct PollTimer { /// Histogram of poll counts within each band. poll_counts: HistogramBatch, /// Instant when the most recent task started polling. poll_started_at: Instant, } impl MetricsBatch { pub(crate) fn new(worker_metrics: &WorkerMetrics) -> MetricsBatch { let now = Instant::now(); MetricsBatch { park_count: 0, park_unpark_count: 0, noop_count: 0, steal_count: 0, steal_operations: 0, poll_count: 0, poll_count_on_last_park: 0, local_schedule_count: 0, overflow_count: 0, busy_duration_total: 0, processing_scheduled_tasks_started_at: now, poll_timer: worker_metrics .poll_count_histogram .as_ref() .map(|worker_poll_counts| PollTimer { poll_counts: HistogramBatch::from_histogram(worker_poll_counts), poll_started_at: now, }), } } pub(crate) fn submit(&mut self, worker: &WorkerMetrics, mean_poll_time: u64) { worker.mean_poll_time.store(mean_poll_time, Relaxed); worker.park_count.store(self.park_count, Relaxed); worker .park_unpark_count .store(self.park_unpark_count, Relaxed); worker.noop_count.store(self.noop_count, Relaxed); worker.steal_count.store(self.steal_count, Relaxed); worker .steal_operations .store(self.steal_operations, Relaxed); worker.poll_count.store(self.poll_count, Relaxed); worker .busy_duration_total .store(self.busy_duration_total, Relaxed); worker .local_schedule_count .store(self.local_schedule_count, Relaxed); worker.overflow_count.store(self.overflow_count, Relaxed); if let Some(poll_timer) = &self.poll_timer { let dst = worker.poll_count_histogram.as_ref().unwrap(); poll_timer.poll_counts.submit(dst); } } /// The worker is about to park. pub(crate) fn about_to_park(&mut self) { self.park_count += 1; self.park_unpark_count += 1; if self.poll_count_on_last_park == self.poll_count { self.noop_count += 1; } else { self.poll_count_on_last_park = self.poll_count; } } /// The worker was unparked. pub(crate) fn unparked(&mut self) { self.park_unpark_count += 1; } /// Start processing a batch of tasks pub(crate) fn start_processing_scheduled_tasks(&mut self) { self.processing_scheduled_tasks_started_at = Instant::now(); } /// Stop processing a batch of tasks pub(crate) fn end_processing_scheduled_tasks(&mut self) { let busy_duration = self.processing_scheduled_tasks_started_at.elapsed(); self.busy_duration_total += duration_as_u64(busy_duration); } /// Start polling an individual task pub(crate) fn start_poll(&mut self) { self.poll_count += 1; if let Some(poll_timer) = &mut self.poll_timer { poll_timer.poll_started_at = Instant::now(); } } /// Stop polling an individual task pub(crate) fn end_poll(&mut self) { if let Some(poll_timer) = &mut self.poll_timer { let elapsed = duration_as_u64(poll_timer.poll_started_at.elapsed()); poll_timer.poll_counts.measure(elapsed, 1); } } pub(crate) fn inc_local_schedule_count(&mut self) { self.local_schedule_count += 1; } } cfg_rt_multi_thread! { impl MetricsBatch { pub(crate) fn incr_steal_count(&mut self, by: u16) { self.steal_count += by as u64; } pub(crate) fn incr_steal_operations(&mut self) { self.steal_operations += 1; } pub(crate) fn incr_overflow_count(&mut self) { self.overflow_count += 1; } } } pub(crate) fn duration_as_u64(dur: Duration) -> u64 { u64::try_from(dur.as_nanos()).unwrap_or(u64::MAX) } tokio-1.43.1/src/runtime/metrics/histogram/h2_histogram.rs000064400000000000000000000500341046102023000217120ustar 00000000000000use crate::runtime::metrics::batch::duration_as_u64; use std::cmp; use std::error::Error; use std::fmt::{Display, Formatter}; use std::time::Duration; const DEFAULT_MIN_VALUE: Duration = Duration::from_nanos(100); const DEFAULT_MAX_VALUE: Duration = Duration::from_secs(60); /// Default precision is 2^-2 = 25% max error const DEFAULT_PRECISION: u32 = 2; const MAX_PRECISION: u32 = 10; /// Log Histogram /// /// This implements an [H2 Histogram](https://iop.systems/blog/h2-histogram/), a histogram similar /// to HdrHistogram, but with better performance. It guarantees an error bound of `2^-p`. /// /// Unlike a traditional H2 histogram this has two small changes: /// 1. The 0th bucket runs for `0..min_value`. This allows truncating a large number of buckets that /// would cover extremely short timescales that customers usually don't care about. /// 2. The final bucket runs all the way to `u64::MAX` — traditional H2 histograms would truncate /// or reject these values. /// /// For information about the default configuration, see [`LogHistogramBuilder`]. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct LogHistogram { /// Number of buckets in the histogram pub(crate) num_buckets: usize, /// Precision of histogram. Error is bounded to 2^-p. pub(crate) p: u32, /// All buckets `idx < bucket_offset` are grouped into bucket 0. /// /// This increases the smallest measurable value of the histogram. pub(crate) bucket_offset: usize, } impl Default for LogHistogram { fn default() -> Self { LogHistogramBuilder::default().build() } } impl LogHistogram { /// Create a Histogram configuration directly from values for `n` and `p`. /// /// # Panics /// - If `bucket_offset` is greater than the specified number of buckets, `(n - p + 1) * 2^p` fn from_n_p(n: u32, p: u32, bucket_offset: usize) -> Self { assert!(n >= p, "{n} (n) must be at least as large as {p} (p)"); let num_buckets = ((n - p + 1) * 1 << p) as usize - bucket_offset; Self { num_buckets, p, bucket_offset, } } fn truncate_to_max_value(&self, max_value: u64) -> LogHistogram { let mut hist = self.clone(); while hist.max_value() >= max_value { hist.num_buckets -= 1; } hist.num_buckets += 1; hist } /// Creates a builder for [`LogHistogram`] pub fn builder() -> LogHistogramBuilder { LogHistogramBuilder::default() } /// The maximum value that can be stored before truncation in this histogram pub fn max_value(&self) -> u64 { self.bucket_range(self.num_buckets - 2).end } pub(crate) fn value_to_bucket(&self, value: u64) -> usize { let index = bucket_index(value, self.p); let offset_bucket = if index < self.bucket_offset as u64 { 0 } else { index - self.bucket_offset as u64 }; let max = self.num_buckets - 1; offset_bucket.min(max as u64) as usize } pub(crate) fn bucket_range(&self, bucket: usize) -> std::ops::Range { let LogHistogram { p, bucket_offset, num_buckets, } = self; let input_bucket = bucket; let bucket = bucket + bucket_offset; let range_start_0th_bucket = match input_bucket { 0 => Some(0_u64), _ => None, }; let range_end_last_bucket = match input_bucket { n if n == num_buckets - 1 => Some(u64::MAX), _ => None, }; if bucket < 1 << p { // The first set of buckets are all size 1 let bucket = bucket as u64; range_start_0th_bucket.unwrap_or(bucket)..range_end_last_bucket.unwrap_or(bucket + 1) } else { // Determine which range of buckets we're in, then determine which bucket in the range it is let bucket = bucket as u64; let p = *p as u64; let w = (bucket >> p) - 1; let base_bucket = (w + 1) * (1_u64 << p); let offset = bucket - base_bucket; let s = 1_u64 << (w + p); let start = s + (offset << w); let end = s + ((offset + 1) << w); range_start_0th_bucket.unwrap_or(start)..range_end_last_bucket.unwrap_or(end) } } } /// Configuration for a [`LogHistogram`] /// /// The log-scaled histogram implements an H2 histogram where the first bucket covers /// the range from 0 to [`LogHistogramBuilder::min_value`] and the final bucket covers /// [`LogHistogramBuilder::max_value`] to infinity. The precision is bounded to the specified /// [`LogHistogramBuilder::max_error`]. Specifically, the precision is the next smallest value /// of `2^-p` such that it is smaller than the requested max error. You can also select `p` directly /// with [`LogHistogramBuilder::precision_exact`]. /// /// Depending on the selected parameters, the number of buckets required is variable. To ensure /// that the histogram size is acceptable, callers may call [`LogHistogramBuilder::max_buckets`]. /// If the resulting histogram would require more buckets, then the method will return an error. /// /// ## Default values /// The default configuration provides the following settings: /// 1. `min_value`: 100ns /// 2. `max_value`: 68 seconds. The final bucket covers all values >68 seconds /// 3. `precision`: max error of 25% /// /// This uses 237 64-bit buckets. #[derive(Default, Debug, Copy, Clone)] pub struct LogHistogramBuilder { max_value: Option, min_value: Option, precision: Option, } impl From for LogHistogram { fn from(value: LogHistogramBuilder) -> Self { value.build() } } impl LogHistogramBuilder { /// Set the precision for this histogram /// /// This function determines the smallest value of `p` that would satisfy the requested precision /// such that `2^-p` is less than `precision`. To set `p` directly, use /// [`LogHistogramBuilder::precision_exact`]. /// /// Precision controls the size of the "bucket groups" (consecutive buckets with identical /// ranges). When `p` is 0, each bucket will be twice the size of the previous bucket. To match /// the behavior of the legacy log histogram implementation, use `builder.precision_exact(0)`. /// /// The default value is 25% (2^-2) /// /// The highest supported precision is `0.0977%` `(2^-10)`. Provided values /// less than this will be truncated. /// /// # Panics /// - `max_error` < 0 /// - `max_error` > 1 pub fn max_error(mut self, max_error: f64) -> Self { assert!(max_error > 0.0, "max_error must be greater than 0"); assert!(max_error < 1.0, "max_error must be less than 1"); let mut p = 2; while 2_f64.powf(-1.0 * p as f64) > max_error && p <= MAX_PRECISION { p += 1; } self.precision = Some(p); self } /// Sets the precision of this histogram directly. /// /// The precision (meaning: the ratio `n/bucket_range(n)` for some given `n`) will be `2^-p`. /// /// Precision controls the number consecutive buckets with identically sized ranges. /// When `p` is 0, each bucket will be twice the size of the previous bucket (bucket groups are /// only a single bucket wide). /// /// To match the behavior of the legacy implementation ([`HistogramScale::Log`]), use `builder.precision_exact(0)`. /// /// # Panics /// - `p` > 10 /// /// [`HistogramScale::Log`]: [crate::runtime::HistogramScale] pub fn precision_exact(mut self, p: u32) -> Self { assert!(p <= MAX_PRECISION, "precision must be <= {MAX_PRECISION}"); self.precision = Some(p); self } /// Sets the minimum duration that can be accurately stored by this histogram. /// /// This sets the resolution. The first bucket will be no larger than /// the provided duration. Setting this value will reduce the number of required buckets, /// sometimes quite significantly. pub fn min_value(mut self, duration: Duration) -> Self { self.min_value = Some(duration); self } /// Sets the maximum value that can by this histogram without truncation /// /// Values greater than this fall in the final bucket that stretches to `u64::MAX`. /// /// # Panics /// The provided value is 0 pub fn max_value(mut self, duration: Duration) -> Self { if duration.is_zero() { panic!("max value must be greater than 0"); } self.max_value = Some(duration); self } /// Builds the log histogram, enforcing the max buckets requirement pub fn max_buckets( &mut self, max_buckets: usize, ) -> Result { let histogram = self.build(); if histogram.num_buckets > max_buckets { return Err(InvalidHistogramConfiguration::TooManyBuckets { required_bucket_count: histogram.num_buckets, }); } Ok(histogram) } /// Builds the log histogram pub fn build(&self) -> LogHistogram { let requested_max_value = duration_as_u64(self.max_value.unwrap_or(DEFAULT_MAX_VALUE)); let max_value = requested_max_value.next_power_of_two(); let min_value = duration_as_u64(self.min_value.unwrap_or(DEFAULT_MIN_VALUE)); let p = self.precision.unwrap_or(DEFAULT_PRECISION); // determine the bucket offset by finding the bucket for the minimum value. We need to lower // this by one to ensure we are at least as granular as requested. let bucket_offset = cmp::max(bucket_index(min_value, p), 1) - 1; // n must be at least as large as p let n = max_value.ilog2().max(p) + 1; LogHistogram::from_n_p(n, p, bucket_offset as usize) .truncate_to_max_value(requested_max_value) } } /// Error constructing a histogram #[derive(Debug)] pub enum InvalidHistogramConfiguration { /// This histogram required more than the specified number of buckets TooManyBuckets { /// The number of buckets that would have been required required_bucket_count: usize, }, } impl Display for InvalidHistogramConfiguration { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { InvalidHistogramConfiguration::TooManyBuckets { required_bucket_count } => write!(f, "The configuration for this histogram would have required {required_bucket_count} buckets") } } } impl Error for InvalidHistogramConfiguration {} /// Compute the index for a given value + p combination /// /// This function does NOT enforce that the value is within the number of expected buckets. fn bucket_index(value: u64, p: u32) -> u64 { // Algorithm described here: https://iop.systems/blog/h2-histogram/ // find the highest non-zero digit if value == 0 { return 0; } let h = 63 - value.leading_zeros(); if h <= p { value } else { let w = h - p; ((w + 1) * (1_u32 << p)) as u64 + ((value - (1_u64 << h)) >> w) } } #[cfg(test)] mod test { use super::InvalidHistogramConfiguration; use crate::runtime::metrics::histogram::h2_histogram::LogHistogram; use crate::runtime::metrics::histogram::HistogramType; #[cfg(not(target_family = "wasm"))] mod proptests { use super::*; use crate::runtime::metrics::batch::duration_as_u64; use crate::runtime::metrics::histogram::h2_histogram::MAX_PRECISION; use proptest::prelude::*; use std::time::Duration; fn valid_log_histogram_strategy() -> impl Strategy { (1..=50u32, 0..=MAX_PRECISION, 0..100usize).prop_map(|(n, p, bucket_offset)| { let p = p.min(n); let base = LogHistogram::from_n_p(n, p, 0); LogHistogram::from_n_p(n, p, bucket_offset.min(base.num_buckets - 1)) }) } fn log_histogram_settings() -> impl Strategy { ( duration_as_u64(Duration::from_nanos(1))..duration_as_u64(Duration::from_secs(20)), duration_as_u64(Duration::from_secs(1))..duration_as_u64(Duration::from_secs(1000)), 0..MAX_PRECISION, ) } // test against a wide assortment of different histogram configurations to ensure invariants hold proptest! { #[test] fn log_histogram_settings_maintain_invariants((min_value, max_value, p) in log_histogram_settings()) { if max_value < min_value { return Ok(()) } let (min_value, max_value) = (Duration::from_nanos(min_value), Duration::from_nanos(max_value)); let histogram = LogHistogram::builder().min_value(min_value).max_value(max_value).precision_exact(p).build(); let first_bucket_end = Duration::from_nanos(histogram.bucket_range(0).end); let last_bucket_start = Duration::from_nanos(histogram.bucket_range(histogram.num_buckets - 1).start); let second_last_bucket_start = Duration::from_nanos(histogram.bucket_range(histogram.num_buckets - 2).start); prop_assert!( first_bucket_end <= min_value, "first bucket {first_bucket_end:?} must be less than {min_value:?}" ); prop_assert!( last_bucket_start > max_value, "last bucket start ({last_bucket_start:?} must be at least as big as `max_value` ({max_value:?})" ); // We should have the exact right number of buckets. The second to last bucket should be strictly less than max value. prop_assert!( second_last_bucket_start < max_value, "second last bucket end ({second_last_bucket_start:?} must be at least as big as `max_value` ({max_value:?})" ); } #[test] fn proptest_log_histogram_invariants(histogram in valid_log_histogram_strategy()) { // 1. Assert that the first bucket always starts at 0 let first_range = histogram.bucket_range(0); prop_assert_eq!(first_range.start, 0, "First bucket doesn't start at 0"); // Check that bucket ranges are disjoint and contiguous let mut prev_end = 0; let mut prev_size = 0; for bucket in 0..histogram.num_buckets { let range = histogram.bucket_range(bucket); prop_assert_eq!(range.start, prev_end, "Bucket ranges are not contiguous"); prop_assert!(range.start < range.end, "Bucket range is empty or reversed"); let size = range.end - range.start; // 2. Assert that the sizes of the buckets are always powers of 2 if bucket > 0 && bucket < histogram.num_buckets - 1 { prop_assert!(size.is_power_of_two(), "Bucket size is not a power of 2"); } if bucket > 1 { // Assert that the sizes of the buckets are monotonically increasing // (after the first bucket, which may be smaller than the 0 bucket) prop_assert!(size >= prev_size, "Bucket sizes are not monotonically increasing: This size {size} (previous: {prev_size}). Bucket: {bucket}"); } // 4. Assert that the size of the buckets is always within the error bound of 2^-p if bucket > 0 && bucket < histogram.num_buckets - 1 { let p = histogram.p as f64; let error_bound = 2.0_f64.powf(-p); // the most it could be wrong is by the length of the range / 2 let relative_error = ((size as f64 - 1.0) / 2.0) / range.start as f64; prop_assert!( relative_error <= error_bound, "Bucket size error exceeds bound: {:?} > {:?} ({range:?})", relative_error, error_bound ); } prev_end = range.end; prev_size = size; } prop_assert_eq!(prev_end, u64::MAX, "Last bucket should end at u64::MAX"); // Check bijection between value_to_bucket and bucket_range for bucket in 0..histogram.num_buckets { let range = histogram.bucket_range(bucket); for value in [range.start, range.end - 1] { prop_assert_eq!( histogram.value_to_bucket(value), bucket, "value_to_bucket is not consistent with bucket_range" ); } } } } } #[test] fn bucket_ranges_are_correct() { let p = 2; let config = HistogramType::H2(LogHistogram { num_buckets: 1024, p, bucket_offset: 0, }); // check precise buckets. There are 2^(p+1) precise buckets for i in 0..2_usize.pow(p + 1) { assert_eq!( config.value_to_bucket(i as u64), i, "{} should be in bucket {}", i, i ); } let mut value = 2_usize.pow(p + 1); let current_bucket = value; while value < current_bucket * 2 { assert_eq!( config.value_to_bucket(value as u64), current_bucket + ((value - current_bucket) / 2), "bucket for {value}" ); value += 1; } } // test buckets against known values #[test] fn bucket_computation_spot_check() { let p = 9; let config = HistogramType::H2(LogHistogram { num_buckets: 4096, p, bucket_offset: 0, }); struct T { v: u64, bucket: usize, } let tests = [ T { v: 1, bucket: 1 }, T { v: 1023, bucket: 1023, }, T { v: 1024, bucket: 1024, }, T { v: 2048, bucket: 1536, }, T { v: 2052, bucket: 1537, }, ]; for test in tests { assert_eq!(config.value_to_bucket(test.v), test.bucket); } } #[test] fn last_bucket_goes_to_infinity() { let conf = HistogramType::H2(LogHistogram::from_n_p(16, 3, 10)); assert_eq!(conf.bucket_range(conf.num_buckets() - 1).end, u64::MAX); } #[test] fn bucket_offset() { // skip the first 10 buckets let conf = HistogramType::H2(LogHistogram::from_n_p(16, 3, 10)); for i in 0..10 { assert_eq!(conf.value_to_bucket(i), 0); } // There are 16 1-element buckets. We skipped 10 of them. The first 2 element bucket starts // at 16 assert_eq!(conf.value_to_bucket(10), 0); assert_eq!(conf.value_to_bucket(16), 6); assert_eq!(conf.value_to_bucket(17), 6); assert_eq!(conf.bucket_range(6), 16..18); } #[test] fn max_buckets_enforcement() { let error = LogHistogram::builder() .max_error(0.001) .max_buckets(5) .expect_err("this produces way more than 5 buckets"); let num_buckets = match error { InvalidHistogramConfiguration::TooManyBuckets { required_bucket_count, } => required_bucket_count, }; assert_eq!(num_buckets, 27291); } #[test] fn default_configuration_size() { let conf = LogHistogram::builder().build(); assert_eq!(conf.num_buckets, 119); } } tokio-1.43.1/src/runtime/metrics/histogram.rs000064400000000000000000000417711046102023000173340ustar 00000000000000mod h2_histogram; pub use h2_histogram::{InvalidHistogramConfiguration, LogHistogram, LogHistogramBuilder}; use crate::util::metric_atomics::MetricAtomicU64; use std::sync::atomic::Ordering::Relaxed; use crate::runtime::metrics::batch::duration_as_u64; use std::cmp; use std::ops::Range; use std::time::Duration; #[derive(Debug)] pub(crate) struct Histogram { /// The histogram buckets buckets: Box<[MetricAtomicU64]>, /// The type of the histogram /// /// This handles `fn(bucket) -> Range` and `fn(value) -> bucket` histogram_type: HistogramType, } #[derive(Debug, Clone)] pub(crate) struct HistogramBuilder { pub(crate) histogram_type: HistogramType, pub(crate) legacy: Option, } #[derive(Debug, Clone)] pub(crate) struct LegacyBuilder { pub(crate) resolution: u64, pub(crate) scale: HistogramScale, pub(crate) num_buckets: usize, } impl Default for LegacyBuilder { fn default() -> Self { Self { resolution: 100_000, num_buckets: 10, scale: HistogramScale::Linear, } } } #[derive(Debug)] pub(crate) struct HistogramBatch { buckets: Box<[u64]>, configuration: HistogramType, } cfg_unstable! { /// Whether the histogram used to aggregate a metric uses a linear or /// logarithmic scale. #[derive(Debug, Copy, Clone, Eq, PartialEq)] #[non_exhaustive] pub enum HistogramScale { /// Linear bucket scale Linear, /// Logarithmic bucket scale Log, } /// Configuration for the poll count histogram #[derive(Debug, Clone)] pub struct HistogramConfiguration { pub(crate) inner: HistogramType } impl HistogramConfiguration { /// Create a linear bucketed histogram /// /// # Arguments /// /// * `bucket_width`: The width of each bucket /// * `num_buckets`: The number of buckets pub fn linear(bucket_width: Duration, num_buckets: usize) -> Self { Self { inner: HistogramType::Linear(LinearHistogram { num_buckets, bucket_width: duration_as_u64(bucket_width), }), } } /// Creates a log-scaled bucketed histogram /// /// See [`LogHistogramBuilder`] for information about configuration & defaults pub fn log(configuration: impl Into) -> Self { Self { inner: HistogramType::H2(configuration.into()), } } } } #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub(crate) enum HistogramType { /// Linear histogram with fixed width buckets Linear(LinearHistogram), /// Old log histogram where each bucket doubles in size LogLegacy(LegacyLogHistogram), /// Log histogram implementation based on H2 Histograms H2(LogHistogram), } impl HistogramType { pub(crate) fn num_buckets(&self) -> usize { match self { HistogramType::Linear(linear) => linear.num_buckets, HistogramType::LogLegacy(log) => log.num_buckets, HistogramType::H2(h2) => h2.num_buckets, } } fn value_to_bucket(&self, value: u64) -> usize { match self { HistogramType::Linear(LinearHistogram { num_buckets, bucket_width, }) => { let max = num_buckets - 1; cmp::min(value / *bucket_width, max as u64) as usize } HistogramType::LogLegacy(LegacyLogHistogram { num_buckets, first_bucket_width, }) => { let max = num_buckets - 1; if value < *first_bucket_width { 0 } else { let significant_digits = 64 - value.leading_zeros(); let bucket_digits = 64 - (first_bucket_width - 1).leading_zeros(); cmp::min(significant_digits as usize - bucket_digits as usize, max) } } HistogramType::H2(log_histogram) => log_histogram.value_to_bucket(value), } } fn bucket_range(&self, bucket: usize) -> Range { match self { HistogramType::Linear(LinearHistogram { num_buckets, bucket_width, }) => Range { start: bucket_width * bucket as u64, end: if bucket == num_buckets - 1 { u64::MAX } else { bucket_width * (bucket as u64 + 1) }, }, HistogramType::LogLegacy(LegacyLogHistogram { num_buckets, first_bucket_width, }) => Range { start: if bucket == 0 { 0 } else { first_bucket_width << (bucket - 1) }, end: if bucket == num_buckets - 1 { u64::MAX } else { first_bucket_width << bucket }, }, HistogramType::H2(log) => log.bucket_range(bucket), } } } #[derive(Debug, Copy, Clone, Eq, PartialEq)] pub(crate) struct LinearHistogram { num_buckets: usize, bucket_width: u64, } #[derive(Debug, Copy, Clone, Eq, PartialEq)] pub(crate) struct LegacyLogHistogram { num_buckets: usize, first_bucket_width: u64, } impl Histogram { pub(crate) fn num_buckets(&self) -> usize { self.buckets.len() } cfg_64bit_metrics! { pub(crate) fn get(&self, bucket: usize) -> u64 { self.buckets[bucket].load(Relaxed) } } pub(crate) fn bucket_range(&self, bucket: usize) -> Range { self.histogram_type.bucket_range(bucket) } } impl HistogramBatch { pub(crate) fn from_histogram(histogram: &Histogram) -> HistogramBatch { let buckets = vec![0; histogram.buckets.len()].into_boxed_slice(); HistogramBatch { buckets, configuration: histogram.histogram_type, } } pub(crate) fn measure(&mut self, value: u64, count: u64) { self.buckets[self.value_to_bucket(value)] += count; } pub(crate) fn submit(&self, histogram: &Histogram) { debug_assert_eq!(self.configuration, histogram.histogram_type); debug_assert_eq!(self.buckets.len(), histogram.buckets.len()); for i in 0..self.buckets.len() { histogram.buckets[i].store(self.buckets[i], Relaxed); } } fn value_to_bucket(&self, value: u64) -> usize { self.configuration.value_to_bucket(value) } } impl HistogramBuilder { pub(crate) fn new() -> HistogramBuilder { HistogramBuilder { histogram_type: HistogramType::Linear(LinearHistogram { num_buckets: 10, bucket_width: 100_000, }), legacy: None, } } pub(crate) fn legacy_mut(&mut self, f: impl Fn(&mut LegacyBuilder)) { let legacy = self.legacy.get_or_insert_with(|| LegacyBuilder::default()); f(legacy); } pub(crate) fn build(&self) -> Histogram { let histogram_type = match &self.legacy { Some(legacy) => { assert!(legacy.resolution > 0); match legacy.scale { HistogramScale::Linear => HistogramType::Linear(LinearHistogram { num_buckets: legacy.num_buckets, bucket_width: legacy.resolution, }), HistogramScale::Log => HistogramType::LogLegacy(LegacyLogHistogram { num_buckets: legacy.num_buckets, first_bucket_width: legacy.resolution.next_power_of_two(), }), } } None => self.histogram_type, }; let num_buckets = histogram_type.num_buckets(); Histogram { buckets: (0..num_buckets) .map(|_| MetricAtomicU64::new(0)) .collect::>() .into_boxed_slice(), histogram_type, } } } impl Default for HistogramBuilder { fn default() -> HistogramBuilder { HistogramBuilder::new() } } #[cfg(all(test, target_has_atomic = "64"))] mod test { use super::*; macro_rules! assert_bucket_eq { ($h:expr, $bucket:expr, $val:expr) => {{ assert_eq!($h.buckets[$bucket], $val); }}; } fn linear(resolution: u64, num_buckets: usize) -> Histogram { HistogramBuilder { histogram_type: HistogramType::Linear(LinearHistogram { bucket_width: resolution, num_buckets, }), legacy: None, } .build() } #[test] fn test_legacy_builder() { let mut builder = HistogramBuilder::new(); builder.legacy_mut(|b| b.num_buckets = 20); assert_eq!(builder.build().num_buckets(), 20); } #[test] fn log_scale_resolution_1() { let h = HistogramBuilder { histogram_type: HistogramType::LogLegacy(LegacyLogHistogram { first_bucket_width: 1, num_buckets: 10, }), legacy: None, } .build(); assert_eq!(h.bucket_range(0), 0..1); assert_eq!(h.bucket_range(1), 1..2); assert_eq!(h.bucket_range(2), 2..4); assert_eq!(h.bucket_range(3), 4..8); assert_eq!(h.bucket_range(9), 256..u64::MAX); let mut b = HistogramBatch::from_histogram(&h); b.measure(0, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 0); b.measure(1, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 0); b.measure(2, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 1); b.measure(3, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 2); b.measure(4, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 2); assert_bucket_eq!(b, 3, 1); b.measure(100, 1); assert_bucket_eq!(b, 7, 1); b.measure(128, 1); assert_bucket_eq!(b, 8, 1); b.measure(4096, 1); assert_bucket_eq!(b, 9, 1); b.measure(u64::MAX, 1); assert_bucket_eq!(b, 9, 2); } #[test] fn log_scale_resolution_2() { let h = HistogramBuilder { histogram_type: HistogramType::LogLegacy(LegacyLogHistogram { num_buckets: 10, first_bucket_width: 2, }), legacy: None, } .build(); assert_eq!(h.bucket_range(0), 0..2); assert_eq!(h.bucket_range(1), 2..4); assert_eq!(h.bucket_range(2), 4..8); assert_eq!(h.bucket_range(3), 8..16); assert_eq!(h.bucket_range(9), 512..u64::MAX); let mut b = HistogramBatch::from_histogram(&h); b.measure(0, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 0); b.measure(1, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 0); b.measure(2, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 0); b.measure(3, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 0); b.measure(4, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 1); b.measure(5, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 2); b.measure(6, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 3); b.measure(7, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 4); b.measure(8, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 4); assert_bucket_eq!(b, 3, 1); b.measure(100, 1); assert_bucket_eq!(b, 6, 1); b.measure(128, 1); assert_bucket_eq!(b, 7, 1); b.measure(4096, 1); assert_bucket_eq!(b, 9, 1); for bucket in h.buckets.iter() { assert_eq!(bucket.load(Relaxed), 0); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } } #[test] fn linear_scale_resolution_1() { let h = linear(1, 10); assert_eq!(h.bucket_range(0), 0..1); assert_eq!(h.bucket_range(1), 1..2); assert_eq!(h.bucket_range(2), 2..3); assert_eq!(h.bucket_range(3), 3..4); assert_eq!(h.bucket_range(9), 9..u64::MAX); let mut b = HistogramBatch::from_histogram(&h); b.measure(0, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 0); b.measure(1, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 0); b.measure(2, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 1); assert_bucket_eq!(b, 3, 0); b.measure(3, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 1); assert_bucket_eq!(b, 3, 1); b.measure(5, 1); assert_bucket_eq!(b, 5, 1); b.measure(4096, 1); assert_bucket_eq!(b, 9, 1); for bucket in h.buckets.iter() { assert_eq!(bucket.load(Relaxed), 0); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } } #[test] fn linear_scale_resolution_100() { let h = linear(100, 10); assert_eq!(h.bucket_range(0), 0..100); assert_eq!(h.bucket_range(1), 100..200); assert_eq!(h.bucket_range(2), 200..300); assert_eq!(h.bucket_range(3), 300..400); assert_eq!(h.bucket_range(9), 900..u64::MAX); let mut b = HistogramBatch::from_histogram(&h); b.measure(0, 1); assert_bucket_eq!(b, 0, 1); assert_bucket_eq!(b, 1, 0); b.measure(50, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 0); b.measure(100, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 1); assert_bucket_eq!(b, 2, 0); b.measure(101, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 0); b.measure(200, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 1); b.measure(299, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 2); b.measure(222, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 3); b.measure(300, 1); assert_bucket_eq!(b, 0, 2); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 3); assert_bucket_eq!(b, 3, 1); b.measure(888, 1); assert_bucket_eq!(b, 8, 1); b.measure(4096, 1); assert_bucket_eq!(b, 9, 1); for bucket in h.buckets.iter() { assert_eq!(bucket.load(Relaxed), 0); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } } #[test] fn inc_by_more_than_one() { let h = linear(100, 10); let mut b = HistogramBatch::from_histogram(&h); b.measure(0, 3); assert_bucket_eq!(b, 0, 3); assert_bucket_eq!(b, 1, 0); b.measure(50, 5); assert_bucket_eq!(b, 0, 8); assert_bucket_eq!(b, 1, 0); b.measure(100, 2); assert_bucket_eq!(b, 0, 8); assert_bucket_eq!(b, 1, 2); assert_bucket_eq!(b, 2, 0); b.measure(101, 19); assert_bucket_eq!(b, 0, 8); assert_bucket_eq!(b, 1, 21); assert_bucket_eq!(b, 2, 0); for bucket in h.buckets.iter() { assert_eq!(bucket.load(Relaxed), 0); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } b.submit(&h); for i in 0..h.buckets.len() { assert_eq!(h.buckets[i].load(Relaxed), b.buckets[i]); } } } tokio-1.43.1/src/runtime/metrics/io.rs000064400000000000000000000012571046102023000157410ustar 00000000000000#![cfg_attr(not(feature = "net"), allow(dead_code))] use crate::util::metric_atomics::MetricAtomicU64; use std::sync::atomic::Ordering::Relaxed; #[derive(Default)] pub(crate) struct IoDriverMetrics { pub(super) fd_registered_count: MetricAtomicU64, pub(super) fd_deregistered_count: MetricAtomicU64, pub(super) ready_count: MetricAtomicU64, } impl IoDriverMetrics { pub(crate) fn incr_fd_count(&self) { self.fd_registered_count.add(1, Relaxed); } pub(crate) fn dec_fd_count(&self) { self.fd_deregistered_count.add(1, Relaxed); } pub(crate) fn incr_ready_count_by(&self, amt: u64) { self.ready_count.add(amt, Relaxed); } } tokio-1.43.1/src/runtime/metrics/mock.rs000064400000000000000000000031631046102023000162610ustar 00000000000000//! This file contains mocks of the types in src/runtime/metrics use std::thread::ThreadId; pub(crate) struct SchedulerMetrics {} pub(crate) struct WorkerMetrics {} pub(crate) struct MetricsBatch {} #[derive(Clone, Default)] pub(crate) struct HistogramBuilder {} impl SchedulerMetrics { pub(crate) fn new() -> Self { Self {} } /// Increment the number of tasks scheduled externally pub(crate) fn inc_remote_schedule_count(&self) {} } impl WorkerMetrics { pub(crate) fn new() -> Self { Self {} } pub(crate) fn from_config(config: &crate::runtime::Config) -> Self { // Prevent the dead-code warning from being triggered let _ = &config.metrics_poll_count_histogram; Self::new() } pub(crate) fn set_queue_depth(&self, _len: usize) {} pub(crate) fn set_thread_id(&self, _thread_id: ThreadId) {} } impl MetricsBatch { pub(crate) fn new(_: &WorkerMetrics) -> Self { Self {} } pub(crate) fn submit(&mut self, _to: &WorkerMetrics, _mean_poll_time: u64) {} pub(crate) fn about_to_park(&mut self) {} pub(crate) fn unparked(&mut self) {} pub(crate) fn inc_local_schedule_count(&mut self) {} pub(crate) fn start_processing_scheduled_tasks(&mut self) {} pub(crate) fn end_processing_scheduled_tasks(&mut self) {} pub(crate) fn start_poll(&mut self) {} pub(crate) fn end_poll(&mut self) {} } cfg_rt_multi_thread! { impl MetricsBatch { pub(crate) fn incr_steal_count(&mut self, _by: u16) {} pub(crate) fn incr_steal_operations(&mut self) {} pub(crate) fn incr_overflow_count(&mut self) {} } } tokio-1.43.1/src/runtime/metrics/mod.rs000064400000000000000000000021631046102023000161060ustar 00000000000000//! This module contains information need to view information about how the //! runtime is performing. //! //! **Note**: This is an [unstable API][unstable]. The public API of types in //! this module may break in 1.x releases. See [the documentation on unstable //! features][unstable] for details. //! //! [unstable]: crate#unstable-features #![allow(clippy::module_inception)] mod runtime; pub use runtime::RuntimeMetrics; cfg_unstable_metrics! { mod batch; pub(crate) use batch::MetricsBatch; mod histogram; pub(crate) use histogram::{Histogram, HistogramBatch, HistogramBuilder}; #[allow(unreachable_pub)] // rust-lang/rust#57411 pub use histogram::{HistogramScale, HistogramConfiguration, LogHistogram, LogHistogramBuilder, InvalidHistogramConfiguration}; mod scheduler; pub(crate) use scheduler::SchedulerMetrics; mod worker; pub(crate) use worker::WorkerMetrics; cfg_net! { mod io; pub(crate) use io::IoDriverMetrics; } } cfg_not_unstable_metrics! { mod mock; pub(crate) use mock::{SchedulerMetrics, WorkerMetrics, MetricsBatch, HistogramBuilder}; } tokio-1.43.1/src/runtime/metrics/runtime.rs000064400000000000000000001263251046102023000170210ustar 00000000000000use crate::runtime::Handle; cfg_unstable_metrics! { use std::ops::Range; use std::thread::ThreadId; cfg_64bit_metrics! { use std::sync::atomic::Ordering::Relaxed; } use std::time::Duration; } /// Handle to the runtime's metrics. /// /// This handle is internally reference-counted and can be freely cloned. A /// `RuntimeMetrics` handle is obtained using the [`Runtime::metrics`] method. /// /// [`Runtime::metrics`]: crate::runtime::Runtime::metrics() #[derive(Clone, Debug)] pub struct RuntimeMetrics { handle: Handle, } impl RuntimeMetrics { pub(crate) fn new(handle: Handle) -> RuntimeMetrics { RuntimeMetrics { handle } } /// Returns the number of worker threads used by the runtime. /// /// The number of workers is set by configuring `worker_threads` on /// `runtime::Builder`. When using the `current_thread` runtime, the return /// value is always `1`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.num_workers(); /// println!("Runtime is using {} workers", n); /// } /// ``` pub fn num_workers(&self) -> usize { self.handle.inner.num_workers() } /// Returns the current number of alive tasks in the runtime. /// /// This counter increases when a task is spawned and decreases when a /// task exits. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.num_alive_tasks(); /// println!("Runtime has {} alive tasks", n); /// } /// ``` pub fn num_alive_tasks(&self) -> usize { self.handle.inner.num_alive_tasks() } /// Returns the number of tasks currently scheduled in the runtime's /// global queue. /// /// Tasks that are spawned or notified from a non-runtime thread are /// scheduled using the runtime's global queue. This metric returns the /// **current** number of tasks pending in the global queue. As such, the /// returned value may increase or decrease as new tasks are scheduled and /// processed. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.global_queue_depth(); /// println!("{} tasks currently pending in the runtime's global queue", n); /// } /// ``` pub fn global_queue_depth(&self) -> usize { self.handle.inner.injection_queue_depth() } cfg_unstable_metrics! { /// Returns the number of additional threads spawned by the runtime. /// /// The number of workers is set by configuring `max_blocking_threads` on /// `runtime::Builder`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let _ = tokio::task::spawn_blocking(move || { /// // Stand-in for compute-heavy work or using synchronous APIs /// 1 + 1 /// }).await; /// let metrics = Handle::current().metrics(); /// /// let n = metrics.num_blocking_threads(); /// println!("Runtime has created {} threads", n); /// } /// ``` pub fn num_blocking_threads(&self) -> usize { self.handle.inner.num_blocking_threads() } #[deprecated = "Renamed to num_alive_tasks"] /// Renamed to [`RuntimeMetrics::num_alive_tasks`] pub fn active_tasks_count(&self) -> usize { self.num_alive_tasks() } /// Returns the number of idle threads, which have spawned by the runtime /// for `spawn_blocking` calls. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let _ = tokio::task::spawn_blocking(move || { /// // Stand-in for compute-heavy work or using synchronous APIs /// 1 + 1 /// }).await; /// let metrics = Handle::current().metrics(); /// /// let n = metrics.num_idle_blocking_threads(); /// println!("Runtime has {} idle blocking thread pool threads", n); /// } /// ``` pub fn num_idle_blocking_threads(&self) -> usize { self.handle.inner.num_idle_blocking_threads() } /// Returns the thread id of the given worker thread. /// /// The returned value is `None` if the worker thread has not yet finished /// starting up. /// /// If additional information about the thread, such as its native id, are /// required, those can be collected in [`on_thread_start`] and correlated /// using the thread id. /// /// [`on_thread_start`]: crate::runtime::Builder::on_thread_start /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let id = metrics.worker_thread_id(0); /// println!("worker 0 has id {:?}", id); /// } /// ``` pub fn worker_thread_id(&self, worker: usize) -> Option { self.handle .inner .worker_metrics(worker) .thread_id() } cfg_64bit_metrics! { /// Returns the number of tasks spawned in this runtime since it was created. /// /// This count starts at zero when the runtime is created and increases by one each time a task is spawned. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.spawned_tasks_count(); /// println!("Runtime has had {} tasks spawned", n); /// } /// ``` pub fn spawned_tasks_count(&self) -> u64 { self.handle.inner.spawned_tasks_count() } /// Returns the number of tasks scheduled from **outside** of the runtime. /// /// The remote schedule count starts at zero when the runtime is created and /// increases by one each time a task is woken from **outside** of the /// runtime. This usually means that a task is spawned or notified from a /// non-runtime thread and must be queued using the Runtime's injection /// queue, which tends to be slower. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.remote_schedule_count(); /// println!("{} tasks were scheduled from outside the runtime", n); /// } /// ``` pub fn remote_schedule_count(&self) -> u64 { self.handle .inner .scheduler_metrics() .remote_schedule_count .load(Relaxed) } /// Returns the number of times that tasks have been forced to yield back to the scheduler /// after exhausting their task budgets. /// /// This count starts at zero when the runtime is created and increases by one each time a task yields due to exhausting its budget. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. pub fn budget_forced_yield_count(&self) -> u64 { self.handle .inner .scheduler_metrics() .budget_forced_yield_count .load(Relaxed) } /// Returns the total number of times the given worker thread has parked. /// /// The worker park count starts at zero when the runtime is created and /// increases by one each time the worker parks the thread waiting for new /// inbound events to process. This usually means the worker has processed /// all pending work and is currently idle. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_park_count(0); /// println!("worker 0 parked {} times", n); /// } /// ``` pub fn worker_park_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .park_count .load(Relaxed) } /// Returns the total number of times the given worker thread has parked /// and unparked. /// /// The worker park/unpark count starts at zero when the runtime is created /// and increases by one each time the worker parks the thread waiting for /// new inbound events to process. This usually means the worker has processed /// all pending work and is currently idle. When new work becomes available, /// the worker is unparked and the park/unpark count is again increased by one. /// /// An odd count means that the worker is currently parked. /// An even count means that the worker is currently active. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// let n = metrics.worker_park_unpark_count(0); /// /// println!("worker 0 parked and unparked {} times", n); /// /// if n % 2 == 0 { /// println!("worker 0 is active"); /// } else { /// println!("worker 0 is parked"); /// } /// } /// ``` pub fn worker_park_unpark_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .park_unpark_count .load(Relaxed) } /// Returns the number of times the given worker thread unparked but /// performed no work before parking again. /// /// The worker no-op count starts at zero when the runtime is created and /// increases by one each time the worker unparks the thread but finds no /// new work and goes back to sleep. This indicates a false-positive wake up. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_noop_count(0); /// println!("worker 0 had {} no-op unparks", n); /// } /// ``` pub fn worker_noop_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .noop_count .load(Relaxed) } /// Returns the number of tasks the given worker thread stole from /// another worker thread. /// /// This metric only applies to the **multi-threaded** runtime and will /// always return `0` when using the current thread runtime. /// /// The worker steal count starts at zero when the runtime is created and /// increases by `N` each time the worker has processed its scheduled queue /// and successfully steals `N` more pending tasks from another worker. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_steal_count(0); /// println!("worker 0 has stolen {} tasks", n); /// } /// ``` pub fn worker_steal_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .steal_count .load(Relaxed) } /// Returns the number of times the given worker thread stole tasks from /// another worker thread. /// /// This metric only applies to the **multi-threaded** runtime and will /// always return `0` when using the current thread runtime. /// /// The worker steal count starts at zero when the runtime is created and /// increases by one each time the worker has processed its scheduled queue /// and successfully steals more pending tasks from another worker. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_steal_operations(0); /// println!("worker 0 has stolen tasks {} times", n); /// } /// ``` pub fn worker_steal_operations(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .steal_operations .load(Relaxed) } /// Returns the number of tasks the given worker thread has polled. /// /// The worker poll count starts at zero when the runtime is created and /// increases by one each time the worker polls a scheduled task. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_poll_count(0); /// println!("worker 0 has polled {} tasks", n); /// } /// ``` pub fn worker_poll_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .poll_count .load(Relaxed) } /// Returns the amount of time the given worker thread has been busy. /// /// The worker busy duration starts at zero when the runtime is created and /// increases whenever the worker is spending time processing work. Using /// this value can indicate the load of the given worker. If a lot of time /// is spent busy, then the worker is under load and will check for inbound /// events less often. /// /// The timer is monotonically increasing. It is never decremented or reset /// to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_total_busy_duration(0); /// println!("worker 0 was busy for a total of {:?}", n); /// } /// ``` pub fn worker_total_busy_duration(&self, worker: usize) -> Duration { let nanos = self .handle .inner .worker_metrics(worker) .busy_duration_total .load(Relaxed); Duration::from_nanos(nanos) } /// Returns the number of tasks scheduled from **within** the runtime on the /// given worker's local queue. /// /// The local schedule count starts at zero when the runtime is created and /// increases by one each time a task is woken from **inside** of the /// runtime on the given worker. This usually means that a task is spawned /// or notified from within a runtime thread and will be queued on the /// worker-local queue. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_local_schedule_count(0); /// println!("{} tasks were scheduled on the worker's local queue", n); /// } /// ``` pub fn worker_local_schedule_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .local_schedule_count .load(Relaxed) } /// Returns the number of times the given worker thread saturated its local /// queue. /// /// This metric only applies to the **multi-threaded** scheduler. /// /// The worker overflow count starts at zero when the runtime is created and /// increases by one each time the worker attempts to schedule a task /// locally, but its local queue is full. When this happens, half of the /// local queue is moved to the injection queue. /// /// The counter is monotonically increasing. It is never decremented or /// reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_overflow_count(0); /// println!("worker 0 has overflowed its queue {} times", n); /// } /// ``` pub fn worker_overflow_count(&self, worker: usize) -> u64 { self.handle .inner .worker_metrics(worker) .overflow_count .load(Relaxed) } } /// Renamed to [`RuntimeMetrics::global_queue_depth`] #[deprecated = "Renamed to global_queue_depth"] #[doc(hidden)] pub fn injection_queue_depth(&self) -> usize { self.handle.inner.injection_queue_depth() } /// Returns the number of tasks currently scheduled in the given worker's /// local queue. /// /// Tasks that are spawned or notified from within a runtime thread are /// scheduled using that worker's local queue. This metric returns the /// **current** number of tasks pending in the worker's local queue. As /// such, the returned value may increase or decrease as new tasks are /// scheduled and processed. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_local_queue_depth(0); /// println!("{} tasks currently pending in worker 0's local queue", n); /// } /// ``` pub fn worker_local_queue_depth(&self, worker: usize) -> usize { self.handle.inner.worker_local_queue_depth(worker) } /// Returns `true` if the runtime is tracking the distribution of task poll /// times. /// /// Task poll times are not instrumented by default as doing so requires /// calling [`Instant::now()`] twice per task poll. The feature is enabled /// by calling [`enable_metrics_poll_time_histogram()`] when building the /// runtime. /// /// # Examples /// /// ``` /// use tokio::runtime::{self, Handle}; /// /// fn main() { /// runtime::Builder::new_current_thread() /// .enable_metrics_poll_time_histogram() /// .build() /// .unwrap() /// .block_on(async { /// let metrics = Handle::current().metrics(); /// let enabled = metrics.poll_time_histogram_enabled(); /// /// println!("Tracking task poll time distribution: {:?}", enabled); /// }); /// } /// ``` /// /// [`enable_metrics_poll_time_histogram()`]: crate::runtime::Builder::enable_metrics_poll_time_histogram /// [`Instant::now()`]: std::time::Instant::now pub fn poll_time_histogram_enabled(&self) -> bool { self.handle .inner .worker_metrics(0) .poll_count_histogram .is_some() } #[deprecated(note = "Renamed to `poll_time_histogram_enabled`")] #[doc(hidden)] pub fn poll_count_histogram_enabled(&self) -> bool { self.poll_time_histogram_enabled() } /// Returns the number of histogram buckets tracking the distribution of /// task poll times. /// /// This value is configured by calling /// [`metrics_poll_time_histogram_configuration()`] when building the runtime. /// /// # Examples /// /// ``` /// use tokio::runtime::{self, Handle}; /// /// fn main() { /// runtime::Builder::new_current_thread() /// .enable_metrics_poll_time_histogram() /// .build() /// .unwrap() /// .block_on(async { /// let metrics = Handle::current().metrics(); /// let buckets = metrics.poll_time_histogram_num_buckets(); /// /// println!("Histogram buckets: {:?}", buckets); /// }); /// } /// ``` /// /// [`metrics_poll_time_histogram_configuration()`]: /// crate::runtime::Builder::metrics_poll_time_histogram_configuration pub fn poll_time_histogram_num_buckets(&self) -> usize { self.handle .inner .worker_metrics(0) .poll_count_histogram .as_ref() .map(|histogram| histogram.num_buckets()) .unwrap_or_default() } /// Deprecated. Use [`poll_time_histogram_num_buckets()`] instead. /// /// [`poll_time_histogram_num_buckets()`]: Self::poll_time_histogram_num_buckets #[doc(hidden)] #[deprecated(note = "renamed to `poll_time_histogram_num_buckets`.")] pub fn poll_count_histogram_num_buckets(&self) -> usize { self.poll_time_histogram_num_buckets() } /// Returns the range of task poll times tracked by the given bucket. /// /// This value is configured by calling /// [`metrics_poll_time_histogram_configuration()`] when building the runtime. /// /// # Panics /// /// The method panics if `bucket` represents an invalid bucket index, i.e. /// is greater than or equal to `poll_time_histogram_num_buckets()`. /// /// # Examples /// /// ``` /// use tokio::runtime::{self, Handle}; /// /// fn main() { /// runtime::Builder::new_current_thread() /// .enable_metrics_poll_time_histogram() /// .build() /// .unwrap() /// .block_on(async { /// let metrics = Handle::current().metrics(); /// let buckets = metrics.poll_time_histogram_num_buckets(); /// /// for i in 0..buckets { /// let range = metrics.poll_time_histogram_bucket_range(i); /// println!("Histogram bucket {} range: {:?}", i, range); /// } /// }); /// } /// ``` /// /// [`metrics_poll_time_histogram_configuration()`]: /// crate::runtime::Builder::metrics_poll_time_histogram_configuration #[track_caller] pub fn poll_time_histogram_bucket_range(&self, bucket: usize) -> Range { self.handle .inner .worker_metrics(0) .poll_count_histogram .as_ref() .map(|histogram| { let range = histogram.bucket_range(bucket); std::ops::Range { start: Duration::from_nanos(range.start), end: Duration::from_nanos(range.end), } }) .unwrap_or_default() } /// Deprecated. Use [`poll_time_histogram_bucket_range()`] instead. /// /// [`poll_time_histogram_bucket_range()`]: Self::poll_time_histogram_bucket_range #[track_caller] #[doc(hidden)] #[deprecated(note = "renamed to `poll_time_histogram_bucket_range`")] pub fn poll_count_histogram_bucket_range(&self, bucket: usize) -> Range { self.poll_time_histogram_bucket_range(bucket) } cfg_64bit_metrics! { /// Returns the number of times the given worker polled tasks with a poll /// duration within the given bucket's range. /// /// Each worker maintains its own histogram and the counts for each bucket /// starts at zero when the runtime is created. Each time the worker polls a /// task, it tracks the duration the task poll time took and increments the /// associated bucket by 1. /// /// Each bucket is a monotonically increasing counter. It is never /// decremented or reset to zero. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// `bucket` is the index of the bucket being queried. The bucket is scoped /// to the worker. The range represented by the bucket can be queried by /// calling [`poll_time_histogram_bucket_range()`]. Each worker maintains /// identical bucket ranges. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()` or if `bucket` represents an /// invalid bucket. /// /// # Examples /// /// ``` /// use tokio::runtime::{self, Handle}; /// /// fn main() { /// runtime::Builder::new_current_thread() /// .enable_metrics_poll_time_histogram() /// .build() /// .unwrap() /// .block_on(async { /// let metrics = Handle::current().metrics(); /// let buckets = metrics.poll_time_histogram_num_buckets(); /// /// for worker in 0..metrics.num_workers() { /// for i in 0..buckets { /// let count = metrics.poll_time_histogram_bucket_count(worker, i); /// println!("Poll count {}", count); /// } /// } /// }); /// } /// ``` /// /// [`poll_time_histogram_bucket_range()`]: crate::runtime::RuntimeMetrics::poll_time_histogram_bucket_range #[track_caller] pub fn poll_time_histogram_bucket_count(&self, worker: usize, bucket: usize) -> u64 { self.handle .inner .worker_metrics(worker) .poll_count_histogram .as_ref() .map(|histogram| histogram.get(bucket)) .unwrap_or_default() } #[doc(hidden)] #[deprecated(note = "use `poll_time_histogram_bucket_count` instead")] pub fn poll_count_histogram_bucket_count(&self, worker: usize, bucket: usize) -> u64 { self.poll_time_histogram_bucket_count(worker, bucket) } /// Returns the mean duration of task polls, in nanoseconds. /// /// This is an exponentially weighted moving average. Currently, this metric /// is only provided by the multi-threaded runtime. /// /// # Arguments /// /// `worker` is the index of the worker being queried. The given value must /// be between 0 and `num_workers()`. The index uniquely identifies a single /// worker and will continue to identify the worker throughout the lifetime /// of the runtime instance. /// /// # Panics /// /// The method panics when `worker` represents an invalid worker, i.e. is /// greater than or equal to `num_workers()`. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.worker_mean_poll_time(0); /// println!("worker 0 has a mean poll time of {:?}", n); /// } /// ``` #[track_caller] pub fn worker_mean_poll_time(&self, worker: usize) -> Duration { let nanos = self .handle .inner .worker_metrics(worker) .mean_poll_time .load(Relaxed); Duration::from_nanos(nanos) } } /// Returns the number of tasks currently scheduled in the blocking /// thread pool, spawned using `spawn_blocking`. /// /// This metric returns the **current** number of tasks pending in /// blocking thread pool. As such, the returned value may increase /// or decrease as new tasks are scheduled and processed. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.blocking_queue_depth(); /// println!("{} tasks currently pending in the blocking thread pool", n); /// } /// ``` pub fn blocking_queue_depth(&self) -> usize { self.handle.inner.blocking_queue_depth() } cfg_net! { cfg_64bit_metrics! { /// Returns the number of file descriptors that have been registered with the /// runtime's I/O driver. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let registered_fds = metrics.io_driver_fd_registered_count(); /// println!("{} fds have been registered with the runtime's I/O driver.", registered_fds); /// /// let deregistered_fds = metrics.io_driver_fd_deregistered_count(); /// /// let current_fd_count = registered_fds - deregistered_fds; /// println!("{} fds are currently registered by the runtime's I/O driver.", current_fd_count); /// } /// ``` pub fn io_driver_fd_registered_count(&self) -> u64 { self.with_io_driver_metrics(|m| { m.fd_registered_count.load(Relaxed) }) } /// Returns the number of file descriptors that have been deregistered by the /// runtime's I/O driver. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.io_driver_fd_deregistered_count(); /// println!("{} fds have been deregistered by the runtime's I/O driver.", n); /// } /// ``` pub fn io_driver_fd_deregistered_count(&self) -> u64 { self.with_io_driver_metrics(|m| { m.fd_deregistered_count.load(Relaxed) }) } /// Returns the number of ready events processed by the runtime's /// I/O driver. /// /// # Examples /// /// ``` /// use tokio::runtime::Handle; /// /// #[tokio::main] /// async fn main() { /// let metrics = Handle::current().metrics(); /// /// let n = metrics.io_driver_ready_count(); /// println!("{} ready events processed by the runtime's I/O driver.", n); /// } /// ``` pub fn io_driver_ready_count(&self) -> u64 { self.with_io_driver_metrics(|m| m.ready_count.load(Relaxed)) } fn with_io_driver_metrics(&self, f: F) -> u64 where F: Fn(&super::IoDriverMetrics) -> u64, { // TODO: Investigate if this should return 0, most of our metrics always increase // thus this breaks that guarantee. self.handle .inner .driver() .io .as_ref() .map(|h| f(&h.metrics)) .unwrap_or(0) } } } } } tokio-1.43.1/src/runtime/metrics/scheduler.rs000064400000000000000000000023121046102023000173010ustar 00000000000000use crate::loom::sync::atomic::Ordering::Relaxed; use crate::util::metric_atomics::MetricAtomicU64; /// Retrieves metrics from the Tokio runtime. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features #[derive(Debug)] pub(crate) struct SchedulerMetrics { /// Number of tasks that are scheduled from outside the runtime. pub(super) remote_schedule_count: MetricAtomicU64, pub(super) budget_forced_yield_count: MetricAtomicU64, } impl SchedulerMetrics { pub(crate) fn new() -> SchedulerMetrics { SchedulerMetrics { remote_schedule_count: MetricAtomicU64::new(0), budget_forced_yield_count: MetricAtomicU64::new(0), } } /// Increment the number of tasks scheduled externally pub(crate) fn inc_remote_schedule_count(&self) { self.remote_schedule_count.add(1, Relaxed); } /// Increment the number of tasks forced to yield due to budget exhaustion pub(crate) fn inc_budget_forced_yield_count(&self) { self.budget_forced_yield_count.add(1, Relaxed); } } tokio-1.43.1/src/runtime/metrics/worker.rs000064400000000000000000000054701046102023000166440ustar 00000000000000use crate::runtime::metrics::Histogram; use crate::runtime::Config; use crate::util::metric_atomics::{MetricAtomicU64, MetricAtomicUsize}; use std::sync::atomic::Ordering::Relaxed; use std::sync::Mutex; use std::thread::ThreadId; /// Retrieve runtime worker metrics. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features #[derive(Debug, Default)] #[repr(align(128))] pub(crate) struct WorkerMetrics { /// Number of times the worker parked. pub(crate) park_count: MetricAtomicU64, /// Number of times the worker parked and unparked. pub(crate) park_unpark_count: MetricAtomicU64, /// Number of times the worker woke then parked again without doing work. pub(crate) noop_count: MetricAtomicU64, /// Number of tasks the worker stole. pub(crate) steal_count: MetricAtomicU64, /// Number of times the worker stole pub(crate) steal_operations: MetricAtomicU64, /// Number of tasks the worker polled. pub(crate) poll_count: MetricAtomicU64, /// EWMA task poll time, in nanoseconds. pub(crate) mean_poll_time: MetricAtomicU64, /// Amount of time the worker spent doing work vs. parking. pub(crate) busy_duration_total: MetricAtomicU64, /// Number of tasks scheduled for execution on the worker's local queue. pub(crate) local_schedule_count: MetricAtomicU64, /// Number of tasks moved from the local queue to the global queue to free space. pub(crate) overflow_count: MetricAtomicU64, /// Number of tasks currently in the local queue. Used only by the /// current-thread scheduler. pub(crate) queue_depth: MetricAtomicUsize, /// If `Some`, tracks the number of polls by duration range. pub(super) poll_count_histogram: Option, /// Thread id of worker thread. thread_id: Mutex>, } impl WorkerMetrics { pub(crate) fn from_config(config: &Config) -> WorkerMetrics { let mut worker_metrics = WorkerMetrics::new(); worker_metrics.poll_count_histogram = config .metrics_poll_count_histogram .as_ref() .map(|histogram_builder| histogram_builder.build()); worker_metrics } pub(crate) fn new() -> WorkerMetrics { WorkerMetrics::default() } pub(crate) fn queue_depth(&self) -> usize { self.queue_depth.load(Relaxed) } pub(crate) fn set_queue_depth(&self, len: usize) { self.queue_depth.store(len, Relaxed); } pub(crate) fn thread_id(&self) -> Option { *self.thread_id.lock().unwrap() } pub(crate) fn set_thread_id(&self, thread_id: ThreadId) { *self.thread_id.lock().unwrap() = Some(thread_id); } } tokio-1.43.1/src/runtime/mod.rs000064400000000000000000000406211046102023000144410ustar 00000000000000//! The Tokio runtime. //! //! Unlike other Rust programs, asynchronous applications require runtime //! support. In particular, the following runtime services are necessary: //! //! * An **I/O event loop**, called the driver, which drives I/O resources and //! dispatches I/O events to tasks that depend on them. //! * A **scheduler** to execute [tasks] that use these I/O resources. //! * A **timer** for scheduling work to run after a set period of time. //! //! Tokio's [`Runtime`] bundles all of these services as a single type, allowing //! them to be started, shut down, and configured together. However, often it is //! not required to configure a [`Runtime`] manually, and a user may just use the //! [`tokio::main`] attribute macro, which creates a [`Runtime`] under the hood. //! //! # Usage //! //! When no fine tuning is required, the [`tokio::main`] attribute macro can be //! used. //! //! ```no_run //! use tokio::net::TcpListener; //! use tokio::io::{AsyncReadExt, AsyncWriteExt}; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! let listener = TcpListener::bind("127.0.0.1:8080").await?; //! //! loop { //! let (mut socket, _) = listener.accept().await?; //! //! tokio::spawn(async move { //! let mut buf = [0; 1024]; //! //! // In a loop, read data from the socket and write the data back. //! loop { //! let n = match socket.read(&mut buf).await { //! // socket closed //! Ok(0) => return, //! Ok(n) => n, //! Err(e) => { //! println!("failed to read from socket; err = {:?}", e); //! return; //! } //! }; //! //! // Write the data back //! if let Err(e) = socket.write_all(&buf[0..n]).await { //! println!("failed to write to socket; err = {:?}", e); //! return; //! } //! } //! }); //! } //! } //! ``` //! //! From within the context of the runtime, additional tasks are spawned using //! the [`tokio::spawn`] function. Futures spawned using this function will be //! executed on the same thread pool used by the [`Runtime`]. //! //! A [`Runtime`] instance can also be used directly. //! //! ```no_run //! use tokio::net::TcpListener; //! use tokio::io::{AsyncReadExt, AsyncWriteExt}; //! use tokio::runtime::Runtime; //! //! fn main() -> Result<(), Box> { //! // Create the runtime //! let rt = Runtime::new()?; //! //! // Spawn the root task //! rt.block_on(async { //! let listener = TcpListener::bind("127.0.0.1:8080").await?; //! //! loop { //! let (mut socket, _) = listener.accept().await?; //! //! tokio::spawn(async move { //! let mut buf = [0; 1024]; //! //! // In a loop, read data from the socket and write the data back. //! loop { //! let n = match socket.read(&mut buf).await { //! // socket closed //! Ok(0) => return, //! Ok(n) => n, //! Err(e) => { //! println!("failed to read from socket; err = {:?}", e); //! return; //! } //! }; //! //! // Write the data back //! if let Err(e) = socket.write_all(&buf[0..n]).await { //! println!("failed to write to socket; err = {:?}", e); //! return; //! } //! } //! }); //! } //! }) //! } //! ``` //! //! ## Runtime Configurations //! //! Tokio provides multiple task scheduling strategies, suitable for different //! applications. The [runtime builder] or `#[tokio::main]` attribute may be //! used to select which scheduler to use. //! //! #### Multi-Thread Scheduler //! //! The multi-thread scheduler executes futures on a _thread pool_, using a //! work-stealing strategy. By default, it will start a worker thread for each //! CPU core available on the system. This tends to be the ideal configuration //! for most applications. The multi-thread scheduler requires the `rt-multi-thread` //! feature flag, and is selected by default: //! ``` //! use tokio::runtime; //! //! # fn main() -> Result<(), Box> { //! let threaded_rt = runtime::Runtime::new()?; //! # Ok(()) } //! ``` //! //! Most applications should use the multi-thread scheduler, except in some //! niche use-cases, such as when running only a single thread is required. //! //! #### Current-Thread Scheduler //! //! The current-thread scheduler provides a _single-threaded_ future executor. //! All tasks will be created and executed on the current thread. This requires //! the `rt` feature flag. //! ``` //! use tokio::runtime; //! //! # fn main() -> Result<(), Box> { //! let rt = runtime::Builder::new_current_thread() //! .build()?; //! # Ok(()) } //! ``` //! //! #### Resource drivers //! //! When configuring a runtime by hand, no resource drivers are enabled by //! default. In this case, attempting to use networking types or time types will //! fail. In order to enable these types, the resource drivers must be enabled. //! This is done with [`Builder::enable_io`] and [`Builder::enable_time`]. As a //! shorthand, [`Builder::enable_all`] enables both resource drivers. //! //! ## Lifetime of spawned threads //! //! The runtime may spawn threads depending on its configuration and usage. The //! multi-thread scheduler spawns threads to schedule tasks and for `spawn_blocking` //! calls. //! //! While the `Runtime` is active, threads may shut down after periods of being //! idle. Once `Runtime` is dropped, all runtime threads have usually been //! terminated, but in the presence of unstoppable spawned work are not //! guaranteed to have been terminated. See the //! [struct level documentation](Runtime#shutdown) for more details. //! //! [tasks]: crate::task //! [`Runtime`]: Runtime //! [`tokio::spawn`]: crate::spawn //! [`tokio::main`]: ../attr.main.html //! [runtime builder]: crate::runtime::Builder //! [`Runtime::new`]: crate::runtime::Runtime::new //! [`Builder::threaded_scheduler`]: crate::runtime::Builder::threaded_scheduler //! [`Builder::enable_io`]: crate::runtime::Builder::enable_io //! [`Builder::enable_time`]: crate::runtime::Builder::enable_time //! [`Builder::enable_all`]: crate::runtime::Builder::enable_all //! //! # Detailed runtime behavior //! //! This section gives more details into how the Tokio runtime will schedule //! tasks for execution. //! //! At its most basic level, a runtime has a collection of tasks that need to be //! scheduled. It will repeatedly remove a task from that collection and //! schedule it (by calling [`poll`]). When the collection is empty, the thread //! will go to sleep until a task is added to the collection. //! //! However, the above is not sufficient to guarantee a well-behaved runtime. //! For example, the runtime might have a single task that is always ready to be //! scheduled, and schedule that task every time. This is a problem because it //! starves other tasks by not scheduling them. To solve this, Tokio provides //! the following fairness guarantee: //! //! > If the total number of tasks does not grow without bound, and no task is //! > [blocking the thread], then it is guaranteed that tasks are scheduled //! > fairly. //! //! Or, more formally: //! //! > Under the following two assumptions: //! > //! > * There is some number `MAX_TASKS` such that the total number of tasks on //! > the runtime at any specific point in time never exceeds `MAX_TASKS`. //! > * There is some number `MAX_SCHEDULE` such that calling [`poll`] on any //! > task spawned on the runtime returns within `MAX_SCHEDULE` time units. //! > //! > Then, there is some number `MAX_DELAY` such that when a task is woken, it //! > will be scheduled by the runtime within `MAX_DELAY` time units. //! //! (Here, `MAX_TASKS` and `MAX_SCHEDULE` can be any number and the user of //! the runtime may choose them. The `MAX_DELAY` number is controlled by the //! runtime, and depends on the value of `MAX_TASKS` and `MAX_SCHEDULE`.) //! //! Other than the above fairness guarantee, there is no guarantee about the //! order in which tasks are scheduled. There is also no guarantee that the //! runtime is equally fair to all tasks. For example, if the runtime has two //! tasks A and B that are both ready, then the runtime may schedule A five //! times before it schedules B. This is the case even if A yields using //! [`yield_now`]. All that is guaranteed is that it will schedule B eventually. //! //! Normally, tasks are scheduled only if they have been woken by calling //! [`wake`] on their waker. However, this is not guaranteed, and Tokio may //! schedule tasks that have not been woken under some circumstances. This is //! called a spurious wakeup. //! //! ## IO and timers //! //! Beyond just scheduling tasks, the runtime must also manage IO resources and //! timers. It does this by periodically checking whether there are any IO //! resources or timers that are ready, and waking the relevant task so that //! it will be scheduled. //! //! These checks are performed periodically between scheduling tasks. Under the //! same assumptions as the previous fairness guarantee, Tokio guarantees that //! it will wake tasks with an IO or timer event within some maximum number of //! time units. //! //! ## Current thread runtime (behavior at the time of writing) //! //! This section describes how the [current thread runtime] behaves today. This //! behavior may change in future versions of Tokio. //! //! The current thread runtime maintains two FIFO queues of tasks that are ready //! to be scheduled: the global queue and the local queue. The runtime will prefer //! to choose the next task to schedule from the local queue, and will only pick a //! task from the global queue if the local queue is empty, or if it has picked //! a task from the local queue 31 times in a row. The number 31 can be //! changed using the [`global_queue_interval`] setting. //! //! The runtime will check for new IO or timer events whenever there are no //! tasks ready to be scheduled, or when it has scheduled 61 tasks in a row. The //! number 61 may be changed using the [`event_interval`] setting. //! //! When a task is woken from within a task running on the runtime, then the //! woken task is added directly to the local queue. Otherwise, the task is //! added to the global queue. The current thread runtime does not use [the lifo //! slot optimization]. //! //! ## Multi threaded runtime (behavior at the time of writing) //! //! This section describes how the [multi thread runtime] behaves today. This //! behavior may change in future versions of Tokio. //! //! A multi thread runtime has a fixed number of worker threads, which are all //! created on startup. The multi thread runtime maintains one global queue, and //! a local queue for each worker thread. The local queue of a worker thread can //! fit at most 256 tasks. If more than 256 tasks are added to the local queue, //! then half of them are moved to the global queue to make space. //! //! The runtime will prefer to choose the next task to schedule from the local //! queue, and will only pick a task from the global queue if the local queue is //! empty, or if it has picked a task from the local queue //! [`global_queue_interval`] times in a row. If the value of //! [`global_queue_interval`] is not explicitly set using the runtime builder, //! then the runtime will dynamically compute it using a heuristic that targets //! 10ms intervals between each check of the global queue (based on the //! [`worker_mean_poll_time`] metric). //! //! If both the local queue and global queue is empty, then the worker thread //! will attempt to steal tasks from the local queue of another worker thread. //! Stealing is done by moving half of the tasks in one local queue to another //! local queue. //! //! The runtime will check for new IO or timer events whenever there are no //! tasks ready to be scheduled, or when it has scheduled 61 tasks in a row. The //! number 61 may be changed using the [`event_interval`] setting. //! //! The multi thread runtime uses [the lifo slot optimization]: Whenever a task //! wakes up another task, the other task is added to the worker thread's lifo //! slot instead of being added to a queue. If there was already a task in the //! lifo slot when this happened, then the lifo slot is replaced, and the task //! that used to be in the lifo slot is placed in the thread's local queue. //! When the runtime finishes scheduling a task, it will schedule the task in //! the lifo slot immediately, if any. When the lifo slot is used, the [coop //! budget] is not reset. Furthermore, if a worker thread uses the lifo slot //! three times in a row, it is temporarily disabled until the worker thread has //! scheduled a task that didn't come from the lifo slot. The lifo slot can be //! disabled using the [`disable_lifo_slot`] setting. The lifo slot is separate //! from the local queue, so other worker threads cannot steal the task in the //! lifo slot. //! //! When a task is woken from a thread that is not a worker thread, then the //! task is placed in the global queue. //! //! [`poll`]: std::future::Future::poll //! [`wake`]: std::task::Waker::wake //! [`yield_now`]: crate::task::yield_now //! [blocking the thread]: https://ryhl.io/blog/async-what-is-blocking/ //! [current thread runtime]: crate::runtime::Builder::new_current_thread //! [multi thread runtime]: crate::runtime::Builder::new_multi_thread //! [`global_queue_interval`]: crate::runtime::Builder::global_queue_interval //! [`event_interval`]: crate::runtime::Builder::event_interval //! [`disable_lifo_slot`]: crate::runtime::Builder::disable_lifo_slot //! [the lifo slot optimization]: crate::runtime::Builder::disable_lifo_slot //! [coop budget]: crate::task#cooperative-scheduling //! [`worker_mean_poll_time`]: crate::runtime::RuntimeMetrics::worker_mean_poll_time // At the top due to macros #[cfg(test)] #[cfg(not(target_family = "wasm"))] #[macro_use] mod tests; pub(crate) mod context; pub(crate) mod coop; pub(crate) mod park; mod driver; pub(crate) mod scheduler; cfg_io_driver_impl! { pub(crate) mod io; } cfg_process_driver! { mod process; } cfg_time! { pub(crate) mod time; } cfg_signal_internal_and_unix! { pub(crate) mod signal; } cfg_rt! { pub(crate) mod task; mod config; use config::Config; mod blocking; #[cfg_attr(target_os = "wasi", allow(unused_imports))] pub(crate) use blocking::spawn_blocking; cfg_trace! { pub(crate) use blocking::Mandatory; } cfg_fs! { pub(crate) use blocking::spawn_mandatory_blocking; } mod builder; pub use self::builder::Builder; cfg_unstable! { mod id; #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] pub use id::Id; pub use self::builder::UnhandledPanic; pub use crate::util::rand::RngSeed; mod local_runtime; pub use local_runtime::{LocalRuntime, LocalOptions}; } cfg_taskdump! { pub mod dump; pub use dump::Dump; } mod task_hooks; pub(crate) use task_hooks::{TaskHooks, TaskCallback}; cfg_unstable! { pub use task_hooks::TaskMeta; } #[cfg(not(tokio_unstable))] pub(crate) use task_hooks::TaskMeta; mod handle; pub use handle::{EnterGuard, Handle, TryCurrentError}; mod runtime; pub use runtime::{Runtime, RuntimeFlavor}; /// Boundary value to prevent stack overflow caused by a large-sized /// Future being placed in the stack. pub(crate) const BOX_FUTURE_THRESHOLD: usize = if cfg!(debug_assertions) { 2048 } else { 16384 }; mod thread_id; pub(crate) use thread_id::ThreadId; pub(crate) mod metrics; pub use metrics::RuntimeMetrics; cfg_unstable_metrics! { pub use metrics::{HistogramScale, HistogramConfiguration, LogHistogram, LogHistogramBuilder, InvalidHistogramConfiguration} ; cfg_net! { pub(crate) use metrics::IoDriverMetrics; } } pub(crate) use metrics::{MetricsBatch, SchedulerMetrics, WorkerMetrics, HistogramBuilder}; /// After thread starts / before thread stops type Callback = std::sync::Arc; } tokio-1.43.1/src/runtime/park.rs000064400000000000000000000236461046102023000146270ustar 00000000000000#![cfg_attr(not(feature = "full"), allow(dead_code))] use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Arc, Condvar, Mutex}; use std::sync::atomic::Ordering::SeqCst; use std::time::Duration; #[derive(Debug)] pub(crate) struct ParkThread { inner: Arc, } /// Unblocks a thread that was blocked by `ParkThread`. #[derive(Clone, Debug)] pub(crate) struct UnparkThread { inner: Arc, } #[derive(Debug)] struct Inner { state: AtomicUsize, mutex: Mutex<()>, condvar: Condvar, } const EMPTY: usize = 0; const PARKED: usize = 1; const NOTIFIED: usize = 2; tokio_thread_local! { static CURRENT_PARKER: ParkThread = ParkThread::new(); } // Bit of a hack, but it is only for loom #[cfg(loom)] tokio_thread_local! { pub(crate) static CURRENT_THREAD_PARK_COUNT: AtomicUsize = AtomicUsize::new(0); } // ==== impl ParkThread ==== impl ParkThread { pub(crate) fn new() -> Self { Self { inner: Arc::new(Inner { state: AtomicUsize::new(EMPTY), mutex: Mutex::new(()), condvar: Condvar::new(), }), } } pub(crate) fn unpark(&self) -> UnparkThread { let inner = self.inner.clone(); UnparkThread { inner } } pub(crate) fn park(&mut self) { #[cfg(loom)] CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst)); self.inner.park(); } pub(crate) fn park_timeout(&mut self, duration: Duration) { #[cfg(loom)] CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst)); self.inner.park_timeout(duration); } pub(crate) fn shutdown(&mut self) { self.inner.shutdown(); } } // ==== impl Inner ==== impl Inner { fn park(&self) { // If we were previously notified then we consume this notification and // return quickly. if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { return; } // Otherwise we need to coordinate going to sleep let mut m = self.mutex.lock(); match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read here, even though we know it will be `NOTIFIED`. // This is because `unpark` may have been called again since we read // `NOTIFIED` in the `compare_exchange` above. We must perform an // acquire operation that synchronizes with that `unpark` to observe // any writes it made before the call to unpark. To do that we must // read from the write it made to `state`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } Err(actual) => panic!("inconsistent park state; actual = {actual}"), } loop { m = self.condvar.wait(m).unwrap(); if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { // got a notification return; } // spurious wakeup, go back to sleep } } /// Parks the current thread for at most `dur`. fn park_timeout(&self, dur: Duration) { // Like `park` above we have a fast path for an already-notified thread, // and afterwards we start coordinating for a sleep. Return quickly. if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { return; } if dur == Duration::from_millis(0) { return; } let m = self.mutex.lock(); match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read again here, see `park`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } Err(actual) => panic!("inconsistent park_timeout state; actual = {actual}"), } #[cfg(not(all(target_family = "wasm", not(target_feature = "atomics"))))] // Wait with a timeout, and if we spuriously wake up or otherwise wake up // from a notification, we just want to unconditionally set the state back to // empty, either consuming a notification or un-flagging ourselves as // parked. let (_m, _result) = self.condvar.wait_timeout(m, dur).unwrap(); #[cfg(all(target_family = "wasm", not(target_feature = "atomics")))] // Wasm without atomics doesn't have threads, so just sleep. { let _m = m; std::thread::sleep(dur); } match self.state.swap(EMPTY, SeqCst) { NOTIFIED => {} // got a notification, hurray! PARKED => {} // no notification, alas n => panic!("inconsistent park_timeout state: {n}"), } } fn unpark(&self) { // To ensure the unparked thread will observe any writes we made before // this call, we must perform a release operation that `park` can // synchronize with. To do that we must write `NOTIFIED` even if `state` // is already `NOTIFIED`. That is why this must be a swap rather than a // compare-and-swap that returns if it reads `NOTIFIED` on failure. match self.state.swap(NOTIFIED, SeqCst) { EMPTY => return, // no one was waiting NOTIFIED => return, // already unparked PARKED => {} // gotta go wake someone up _ => panic!("inconsistent state in unpark"), } // There is a period between when the parked thread sets `state` to // `PARKED` (or last checked `state` in the case of a spurious wake // up) and when it actually waits on `cvar`. If we were to notify // during this period it would be ignored and then when the parked // thread went to sleep it would never wake up. Fortunately, it has // `lock` locked at this stage so we can acquire `lock` to wait until // it is ready to receive the notification. // // Releasing `lock` before the call to `notify_one` means that when the // parked thread wakes it doesn't get woken only to have to wait for us // to release `lock`. drop(self.mutex.lock()); self.condvar.notify_one(); } fn shutdown(&self) { self.condvar.notify_all(); } } impl Default for ParkThread { fn default() -> Self { Self::new() } } // ===== impl UnparkThread ===== impl UnparkThread { pub(crate) fn unpark(&self) { self.inner.unpark(); } } use crate::loom::thread::AccessError; use std::future::Future; use std::marker::PhantomData; use std::rc::Rc; use std::task::{RawWaker, RawWakerVTable, Waker}; /// Blocks the current thread using a condition variable. #[derive(Debug)] pub(crate) struct CachedParkThread { _anchor: PhantomData>, } impl CachedParkThread { /// Creates a new `ParkThread` handle for the current thread. /// /// This type cannot be moved to other threads, so it should be created on /// the thread that the caller intends to park. pub(crate) fn new() -> CachedParkThread { CachedParkThread { _anchor: PhantomData, } } pub(crate) fn waker(&self) -> Result { self.unpark().map(UnparkThread::into_waker) } fn unpark(&self) -> Result { self.with_current(ParkThread::unpark) } pub(crate) fn park(&mut self) { self.with_current(|park_thread| park_thread.inner.park()) .unwrap(); } pub(crate) fn park_timeout(&mut self, duration: Duration) { self.with_current(|park_thread| park_thread.inner.park_timeout(duration)) .unwrap(); } /// Gets a reference to the `ParkThread` handle for this thread. fn with_current(&self, f: F) -> Result where F: FnOnce(&ParkThread) -> R, { CURRENT_PARKER.try_with(|inner| f(inner)) } pub(crate) fn block_on(&mut self, f: F) -> Result { use std::task::Context; use std::task::Poll::Ready; let waker = self.waker()?; let mut cx = Context::from_waker(&waker); pin!(f); loop { if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) { return Ok(v); } self.park(); } } } impl UnparkThread { pub(crate) fn into_waker(self) -> Waker { unsafe { let raw = unparker_to_raw_waker(self.inner); Waker::from_raw(raw) } } } impl Inner { #[allow(clippy::wrong_self_convention)] fn into_raw(this: Arc) -> *const () { Arc::into_raw(this) as *const () } unsafe fn from_raw(ptr: *const ()) -> Arc { Arc::from_raw(ptr as *const Inner) } } unsafe fn unparker_to_raw_waker(unparker: Arc) -> RawWaker { RawWaker::new( Inner::into_raw(unparker), &RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker), ) } unsafe fn clone(raw: *const ()) -> RawWaker { Arc::increment_strong_count(raw as *const Inner); unparker_to_raw_waker(Inner::from_raw(raw)) } unsafe fn drop_waker(raw: *const ()) { drop(Inner::from_raw(raw)); } unsafe fn wake(raw: *const ()) { let unparker = Inner::from_raw(raw); unparker.unpark(); } unsafe fn wake_by_ref(raw: *const ()) { let raw = raw as *const Inner; (*raw).unpark(); } #[cfg(loom)] pub(crate) fn current_thread_park_count() -> usize { CURRENT_THREAD_PARK_COUNT.with(|count| count.load(SeqCst)) } tokio-1.43.1/src/runtime/process.rs000064400000000000000000000022671046102023000153440ustar 00000000000000#![cfg_attr(not(feature = "rt"), allow(dead_code))] //! Process driver. use crate::process::unix::GlobalOrphanQueue; use crate::runtime::driver; use crate::runtime::signal::{Driver as SignalDriver, Handle as SignalHandle}; use std::time::Duration; /// Responsible for cleaning up orphaned child processes on Unix platforms. #[derive(Debug)] pub(crate) struct Driver { park: SignalDriver, signal_handle: SignalHandle, } // ===== impl Driver ===== impl Driver { /// Creates a new signal `Driver` instance that delegates wakeups to `park`. pub(crate) fn new(park: SignalDriver) -> Self { let signal_handle = park.handle(); Self { park, signal_handle, } } pub(crate) fn park(&mut self, handle: &driver::Handle) { self.park.park(handle); GlobalOrphanQueue::reap_orphans(&self.signal_handle); } pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { self.park.park_timeout(handle, duration); GlobalOrphanQueue::reap_orphans(&self.signal_handle); } pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { self.park.shutdown(handle); } } tokio-1.43.1/src/runtime/runtime.rs000064400000000000000000000433461046102023000153540ustar 00000000000000use super::BOX_FUTURE_THRESHOLD; use crate::runtime::blocking::BlockingPool; use crate::runtime::scheduler::CurrentThread; use crate::runtime::{context, EnterGuard, Handle}; use crate::task::JoinHandle; use crate::util::trace::SpawnMeta; use std::future::Future; use std::mem; use std::time::Duration; cfg_rt_multi_thread! { use crate::runtime::Builder; use crate::runtime::scheduler::MultiThread; cfg_unstable! { use crate::runtime::scheduler::MultiThreadAlt; } } /// The Tokio runtime. /// /// The runtime provides an I/O driver, task scheduler, [timer], and /// blocking pool, necessary for running asynchronous tasks. /// /// Instances of `Runtime` can be created using [`new`], or [`Builder`]. /// However, most users will use the [`#[tokio::main]`][main] annotation on /// their entry point instead. /// /// See [module level][mod] documentation for more details. /// /// # Shutdown /// /// Shutting down the runtime is done by dropping the value, or calling /// [`shutdown_background`] or [`shutdown_timeout`]. /// /// Tasks spawned through [`Runtime::spawn`] keep running until they yield. /// Then they are dropped. They are not *guaranteed* to run to completion, but /// *might* do so if they do not yield until completion. /// /// Blocking functions spawned through [`Runtime::spawn_blocking`] keep running /// until they return. /// /// The thread initiating the shutdown blocks until all spawned work has been /// stopped. This can take an indefinite amount of time. The `Drop` /// implementation waits forever for this. /// /// The [`shutdown_background`] and [`shutdown_timeout`] methods can be used if /// waiting forever is undesired. When the timeout is reached, spawned work that /// did not stop in time and threads running it are leaked. The work continues /// to run until one of the stopping conditions is fulfilled, but the thread /// initiating the shutdown is unblocked. /// /// Once the runtime has been dropped, any outstanding I/O resources bound to /// it will no longer function. Calling any method on them will result in an /// error. /// /// # Sharing /// /// There are several ways to establish shared access to a Tokio runtime: /// /// * Using an [Arc]\. /// * Using a [`Handle`]. /// * Entering the runtime context. /// /// Using an [Arc]\ or [`Handle`] allows you to do various /// things with the runtime such as spawning new tasks or entering the runtime /// context. Both types can be cloned to create a new handle that allows access /// to the same runtime. By passing clones into different tasks or threads, you /// will be able to access the runtime from those tasks or threads. /// /// The difference between [Arc]\ and [`Handle`] is that /// an [Arc]\ will prevent the runtime from shutting down, /// whereas a [`Handle`] does not prevent that. This is because shutdown of the /// runtime happens when the destructor of the `Runtime` object runs. /// /// Calls to [`shutdown_background`] and [`shutdown_timeout`] require exclusive /// ownership of the `Runtime` type. When using an [Arc]\, /// this can be achieved via [`Arc::try_unwrap`] when only one strong count /// reference is left over. /// /// The runtime context is entered using the [`Runtime::enter`] or /// [`Handle::enter`] methods, which use a thread-local variable to store the /// current runtime. Whenever you are inside the runtime context, methods such /// as [`tokio::spawn`] will use the runtime whose context you are inside. /// /// [timer]: crate::time /// [mod]: index.html /// [`new`]: method@Self::new /// [`Builder`]: struct@Builder /// [`Handle`]: struct@Handle /// [main]: macro@crate::main /// [`tokio::spawn`]: crate::spawn /// [`Arc::try_unwrap`]: std::sync::Arc::try_unwrap /// [Arc]: std::sync::Arc /// [`shutdown_background`]: method@Runtime::shutdown_background /// [`shutdown_timeout`]: method@Runtime::shutdown_timeout #[derive(Debug)] pub struct Runtime { /// Task scheduler scheduler: Scheduler, /// Handle to runtime, also contains driver handles handle: Handle, /// Blocking pool handle, used to signal shutdown blocking_pool: BlockingPool, } /// The flavor of a `Runtime`. /// /// This is the return type for [`Handle::runtime_flavor`](crate::runtime::Handle::runtime_flavor()). #[derive(Debug, PartialEq, Eq)] #[non_exhaustive] pub enum RuntimeFlavor { /// The flavor that executes all tasks on the current thread. CurrentThread, /// The flavor that executes tasks across multiple threads. MultiThread, /// The flavor that executes tasks across multiple threads. #[cfg(tokio_unstable)] #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] MultiThreadAlt, } /// The runtime scheduler is either a multi-thread or a current-thread executor. #[derive(Debug)] pub(super) enum Scheduler { /// Execute all tasks on the current-thread. CurrentThread(CurrentThread), /// Execute tasks across multiple threads. #[cfg(feature = "rt-multi-thread")] MultiThread(MultiThread), /// Execute tasks across multiple threads. #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] MultiThreadAlt(MultiThreadAlt), } impl Runtime { pub(super) fn from_parts( scheduler: Scheduler, handle: Handle, blocking_pool: BlockingPool, ) -> Runtime { Runtime { scheduler, handle, blocking_pool, } } /// Creates a new runtime instance with default configuration values. /// /// This results in the multi threaded scheduler, I/O driver, and time driver being /// initialized. /// /// Most applications will not need to call this function directly. Instead, /// they will use the [`#[tokio::main]` attribute][main]. When a more complex /// configuration is necessary, the [runtime builder] may be used. /// /// See [module level][mod] documentation for more details. /// /// # Examples /// /// Creating a new `Runtime` with default configuration values. /// /// ``` /// use tokio::runtime::Runtime; /// /// let rt = Runtime::new() /// .unwrap(); /// /// // Use the runtime... /// ``` /// /// [mod]: index.html /// [main]: ../attr.main.html /// [threaded scheduler]: index.html#threaded-scheduler /// [runtime builder]: crate::runtime::Builder #[cfg(feature = "rt-multi-thread")] #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] pub fn new() -> std::io::Result { Builder::new_multi_thread().enable_all().build() } /// Returns a handle to the runtime's spawner. /// /// The returned handle can be used to spawn tasks that run on this runtime, and can /// be cloned to allow moving the `Handle` to other threads. /// /// Calling [`Handle::block_on`] on a handle to a `current_thread` runtime is error-prone. /// Refer to the documentation of [`Handle::block_on`] for more. /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// let rt = Runtime::new() /// .unwrap(); /// /// let handle = rt.handle(); /// /// // Use the handle... /// ``` pub fn handle(&self) -> &Handle { &self.handle } /// Spawns a future onto the Tokio runtime. /// /// This spawns the given future onto the runtime's executor, usually a /// thread pool. The thread pool is then responsible for polling the future /// until it completes. /// /// The provided future will start running in the background immediately /// when `spawn` is called, even if you don't await the returned /// `JoinHandle`. /// /// See [module level][mod] documentation for more details. /// /// [mod]: index.html /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// # fn dox() { /// // Create the runtime /// let rt = Runtime::new().unwrap(); /// /// // Spawn a future onto the runtime /// rt.spawn(async { /// println!("now running on a worker thread"); /// }); /// # } /// ``` #[track_caller] pub fn spawn(&self, future: F) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, { let fut_size = mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { self.handle .spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { self.handle .spawn_named(future, SpawnMeta::new_unnamed(fut_size)) } } /// Runs the provided function on an executor dedicated to blocking operations. /// /// # Examples /// /// ``` /// use tokio::runtime::Runtime; /// /// # fn dox() { /// // Create the runtime /// let rt = Runtime::new().unwrap(); /// /// // Spawn a blocking function onto the runtime /// rt.spawn_blocking(|| { /// println!("now running on a worker thread"); /// }); /// # } /// ``` #[track_caller] pub fn spawn_blocking(&self, func: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { self.handle.spawn_blocking(func) } /// Runs a future to completion on the Tokio runtime. This is the /// runtime's entry point. /// /// This runs the given future on the current thread, blocking until it is /// complete, and yielding its resolved result. Any tasks or timers /// which the future spawns internally will be executed on the runtime. /// /// # Non-worker future /// /// Note that the future required by this function does not run as a /// worker. The expectation is that other tasks are spawned by the future here. /// Awaiting on other futures from the future provided here will not /// perform as fast as those spawned as workers. /// /// # Multi thread scheduler /// /// When the multi thread scheduler is used this will allow futures /// to run within the io driver and timer context of the overall runtime. /// /// Any spawned tasks will continue running after `block_on` returns. /// /// # Current thread scheduler /// /// When the current thread scheduler is enabled `block_on` /// can be called concurrently from multiple threads. The first call /// will take ownership of the io and timer drivers. This means /// other threads which do not own the drivers will hook into that one. /// When the first `block_on` completes, other threads will be able to /// "steal" the driver to allow continued execution of their futures. /// /// Any spawned tasks will be suspended after `block_on` returns. Calling /// `block_on` again will resume previously spawned tasks. /// /// # Panics /// /// This function panics if the provided future panics, or if called within an /// asynchronous execution context. /// /// # Examples /// /// ```no_run /// use tokio::runtime::Runtime; /// /// // Create the runtime /// let rt = Runtime::new().unwrap(); /// /// // Execute the future, blocking the current thread until completion /// rt.block_on(async { /// println!("hello"); /// }); /// ``` /// /// [handle]: fn@Handle::block_on #[track_caller] pub fn block_on(&self, future: F) -> F::Output { let fut_size = mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { self.block_on_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { self.block_on_inner(future, SpawnMeta::new_unnamed(fut_size)) } } #[track_caller] fn block_on_inner(&self, future: F, _meta: SpawnMeta<'_>) -> F::Output { #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] let future = super::task::trace::Trace::root(future); #[cfg(all(tokio_unstable, feature = "tracing"))] let future = crate::util::trace::task( future, "block_on", _meta, crate::runtime::task::Id::next().as_u64(), ); let _enter = self.enter(); match &self.scheduler { Scheduler::CurrentThread(exec) => exec.block_on(&self.handle.inner, future), #[cfg(feature = "rt-multi-thread")] Scheduler::MultiThread(exec) => exec.block_on(&self.handle.inner, future), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Scheduler::MultiThreadAlt(exec) => exec.block_on(&self.handle.inner, future), } } /// Enters the runtime context. /// /// This allows you to construct types that must have an executor /// available on creation such as [`Sleep`] or [`TcpStream`]. It will /// also allow you to call methods such as [`tokio::spawn`]. /// /// [`Sleep`]: struct@crate::time::Sleep /// [`TcpStream`]: struct@crate::net::TcpStream /// [`tokio::spawn`]: fn@crate::spawn /// /// # Example /// /// ``` /// use tokio::runtime::Runtime; /// use tokio::task::JoinHandle; /// /// fn function_that_spawns(msg: String) -> JoinHandle<()> { /// // Had we not used `rt.enter` below, this would panic. /// tokio::spawn(async move { /// println!("{}", msg); /// }) /// } /// /// fn main() { /// let rt = Runtime::new().unwrap(); /// /// let s = "Hello World!".to_string(); /// /// // By entering the context, we tie `tokio::spawn` to this executor. /// let _guard = rt.enter(); /// let handle = function_that_spawns(s); /// /// // Wait for the task before we end the test. /// rt.block_on(handle).unwrap(); /// } /// ``` pub fn enter(&self) -> EnterGuard<'_> { self.handle.enter() } /// Shuts down the runtime, waiting for at most `duration` for all spawned /// work to stop. /// /// See the [struct level documentation](Runtime#shutdown) for more details. /// /// # Examples /// /// ``` /// # if cfg!(miri) { return } // Miri reports error when main thread terminated without waiting all remaining threads. /// use tokio::runtime::Runtime; /// use tokio::task; /// /// use std::thread; /// use std::time::Duration; /// /// fn main() { /// let runtime = Runtime::new().unwrap(); /// /// runtime.block_on(async move { /// task::spawn_blocking(move || { /// thread::sleep(Duration::from_secs(10_000)); /// }); /// }); /// /// runtime.shutdown_timeout(Duration::from_millis(100)); /// } /// ``` pub fn shutdown_timeout(mut self, duration: Duration) { // Wakeup and shutdown all the worker threads self.handle.inner.shutdown(); self.blocking_pool.shutdown(Some(duration)); } /// Shuts down the runtime, without waiting for any spawned work to stop. /// /// This can be useful if you want to drop a runtime from within another runtime. /// Normally, dropping a runtime will block indefinitely for spawned blocking tasks /// to complete, which would normally not be permitted within an asynchronous context. /// By calling `shutdown_background()`, you can drop the runtime from such a context. /// /// Note however, that because we do not wait for any blocking tasks to complete, this /// may result in a resource leak (in that any blocking tasks are still running until they /// return. /// /// See the [struct level documentation](Runtime#shutdown) for more details. /// /// This function is equivalent to calling `shutdown_timeout(Duration::from_nanos(0))`. /// /// ``` /// use tokio::runtime::Runtime; /// /// fn main() { /// let runtime = Runtime::new().unwrap(); /// /// runtime.block_on(async move { /// let inner_runtime = Runtime::new().unwrap(); /// // ... /// inner_runtime.shutdown_background(); /// }); /// } /// ``` pub fn shutdown_background(self) { self.shutdown_timeout(Duration::from_nanos(0)); } /// Returns a view that lets you get information about how the runtime /// is performing. pub fn metrics(&self) -> crate::runtime::RuntimeMetrics { self.handle.metrics() } } #[allow(clippy::single_match)] // there are comments in the error branch, so we don't want if-let impl Drop for Runtime { fn drop(&mut self) { match &mut self.scheduler { Scheduler::CurrentThread(current_thread) => { // This ensures that tasks spawned on the current-thread // runtime are dropped inside the runtime's context. let _guard = context::try_set_current(&self.handle.inner); current_thread.shutdown(&self.handle.inner); } #[cfg(feature = "rt-multi-thread")] Scheduler::MultiThread(multi_thread) => { // The threaded scheduler drops its tasks on its worker threads, which is // already in the runtime's context. multi_thread.shutdown(&self.handle.inner); } #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Scheduler::MultiThreadAlt(multi_thread) => { // The threaded scheduler drops its tasks on its worker threads, which is // already in the runtime's context. multi_thread.shutdown(&self.handle.inner); } } } } impl std::panic::UnwindSafe for Runtime {} impl std::panic::RefUnwindSafe for Runtime {} tokio-1.43.1/src/runtime/scheduler/block_in_place.rs000064400000000000000000000007641046102023000205700ustar 00000000000000use crate::runtime::scheduler; #[track_caller] pub(crate) fn block_in_place(f: F) -> R where F: FnOnce() -> R, { #[cfg(tokio_unstable)] { use crate::runtime::{Handle, RuntimeFlavor::MultiThreadAlt}; match Handle::try_current().map(|h| h.runtime_flavor()) { Ok(MultiThreadAlt) => { return scheduler::multi_thread_alt::block_in_place(f); } _ => {} } } scheduler::multi_thread::block_in_place(f) } tokio-1.43.1/src/runtime/scheduler/current_thread/mod.rs000064400000000000000000000622741046102023000214400ustar 00000000000000use crate::loom::sync::atomic::AtomicBool; use crate::loom::sync::Arc; use crate::runtime::driver::{self, Driver}; use crate::runtime::scheduler::{self, Defer, Inject}; use crate::runtime::task::{ self, JoinHandle, OwnedTasks, Schedule, Task, TaskHarnessScheduleHooks, }; use crate::runtime::{ blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics, }; use crate::sync::notify::Notify; use crate::util::atomic_cell::AtomicCell; use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef}; use std::cell::RefCell; use std::collections::VecDeque; use std::future::{poll_fn, Future}; use std::sync::atomic::Ordering::{AcqRel, Release}; use std::task::Poll::{Pending, Ready}; use std::task::Waker; use std::thread::ThreadId; use std::time::Duration; use std::{fmt, thread}; /// Executes tasks on the current thread pub(crate) struct CurrentThread { /// Core scheduler data is acquired by a thread entering `block_on`. core: AtomicCell, /// Notifier for waking up other threads to steal the /// driver. notify: Notify, } /// Handle to the current thread scheduler pub(crate) struct Handle { /// Scheduler state shared across threads shared: Shared, /// Resource driver handles pub(crate) driver: driver::Handle, /// Blocking pool spawner pub(crate) blocking_spawner: blocking::Spawner, /// Current random number generator seed pub(crate) seed_generator: RngSeedGenerator, /// User-supplied hooks to invoke for things pub(crate) task_hooks: TaskHooks, /// If this is a `LocalRuntime`, flags the owning thread ID. pub(crate) local_tid: Option, } /// Data required for executing the scheduler. The struct is passed around to /// a function that will perform the scheduling work and acts as a capability token. struct Core { /// Scheduler run queue tasks: VecDeque, /// Current tick tick: u32, /// Runtime driver /// /// The driver is removed before starting to park the thread driver: Option, /// Metrics batch metrics: MetricsBatch, /// How often to check the global queue global_queue_interval: u32, /// True if a task panicked without being handled and the runtime is /// configured to shutdown on unhandled panic. unhandled_panic: bool, } /// Scheduler state shared between threads. struct Shared { /// Remote run queue inject: Inject>, /// Collection of all active tasks spawned onto this executor. owned: OwnedTasks>, /// Indicates whether the blocked on thread was woken. woken: AtomicBool, /// Scheduler configuration options config: Config, /// Keeps track of various runtime metrics. scheduler_metrics: SchedulerMetrics, /// This scheduler only has one worker. worker_metrics: WorkerMetrics, } /// Thread-local context. /// /// pub(crate) to store in `runtime::context`. pub(crate) struct Context { /// Scheduler handle handle: Arc, /// Scheduler core, enabling the holder of `Context` to execute the /// scheduler. core: RefCell>>, /// Deferred tasks, usually ones that called `task::yield_now()`. pub(crate) defer: Defer, } type Notified = task::Notified>; /// Initial queue capacity. const INITIAL_CAPACITY: usize = 64; /// Used if none is specified. This is a temporary constant and will be removed /// as we unify tuning logic between the multi-thread and current-thread /// schedulers. const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = 31; impl CurrentThread { pub(crate) fn new( driver: Driver, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, local_tid: Option, ) -> (CurrentThread, Arc) { let worker_metrics = WorkerMetrics::from_config(&config); worker_metrics.set_thread_id(thread::current().id()); // Get the configured global queue interval, or use the default. let global_queue_interval = config .global_queue_interval .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL); let handle = Arc::new(Handle { task_hooks: TaskHooks { task_spawn_callback: config.before_spawn.clone(), task_terminate_callback: config.after_termination.clone(), }, shared: Shared { inject: Inject::new(), owned: OwnedTasks::new(1), woken: AtomicBool::new(false), config, scheduler_metrics: SchedulerMetrics::new(), worker_metrics, }, driver: driver_handle, blocking_spawner, seed_generator, local_tid, }); let core = AtomicCell::new(Some(Box::new(Core { tasks: VecDeque::with_capacity(INITIAL_CAPACITY), tick: 0, driver: Some(driver), metrics: MetricsBatch::new(&handle.shared.worker_metrics), global_queue_interval, unhandled_panic: false, }))); let scheduler = CurrentThread { core, notify: Notify::new(), }; (scheduler, handle) } #[track_caller] pub(crate) fn block_on(&self, handle: &scheduler::Handle, future: F) -> F::Output { pin!(future); crate::runtime::context::enter_runtime(handle, false, |blocking| { let handle = handle.as_current_thread(); // Attempt to steal the scheduler core and block_on the future if we can // there, otherwise, lets select on a notification that the core is // available or the future is complete. loop { if let Some(core) = self.take_core(handle) { handle .shared .worker_metrics .set_thread_id(thread::current().id()); return core.block_on(future); } else { let notified = self.notify.notified(); pin!(notified); if let Some(out) = blocking .block_on(poll_fn(|cx| { if notified.as_mut().poll(cx).is_ready() { return Ready(None); } if let Ready(out) = future.as_mut().poll(cx) { return Ready(Some(out)); } Pending })) .expect("Failed to `Enter::block_on`") { return out; } } } }) } fn take_core(&self, handle: &Arc) -> Option> { let core = self.core.take()?; Some(CoreGuard { context: scheduler::Context::CurrentThread(Context { handle: handle.clone(), core: RefCell::new(Some(core)), defer: Defer::new(), }), scheduler: self, }) } pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) { let handle = handle.as_current_thread(); // Avoid a double panic if we are currently panicking and // the lock may be poisoned. let core = match self.take_core(handle) { Some(core) => core, None if std::thread::panicking() => return, None => panic!("Oh no! We never placed the Core back, this is a bug!"), }; // Check that the thread-local is not being destroyed let tls_available = context::with_current(|_| ()).is_ok(); if tls_available { core.enter(|core, _context| { let core = shutdown2(core, handle); (core, ()) }); } else { // Shutdown without setting the context. `tokio::spawn` calls will // fail, but those will fail either way because the thread-local is // not available anymore. let context = core.context.expect_current_thread(); let core = context.core.borrow_mut().take().unwrap(); let core = shutdown2(core, handle); *context.core.borrow_mut() = Some(core); } } } fn shutdown2(mut core: Box, handle: &Handle) -> Box { // Drain the OwnedTasks collection. This call also closes the // collection, ensuring that no tasks are ever pushed after this // call returns. handle.shared.owned.close_and_shutdown_all(0); // Drain local queue // We already shut down every task, so we just need to drop the task. while let Some(task) = core.next_local_task(handle) { drop(task); } // Close the injection queue handle.shared.inject.close(); // Drain remote queue while let Some(task) = handle.shared.inject.pop() { drop(task); } assert!(handle.shared.owned.is_empty()); // Submit metrics core.submit_metrics(handle); // Shutdown the resource drivers if let Some(driver) = core.driver.as_mut() { driver.shutdown(&handle.driver); } core } impl fmt::Debug for CurrentThread { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("CurrentThread").finish() } } // ===== impl Core ===== impl Core { /// Get and increment the current tick fn tick(&mut self) { self.tick = self.tick.wrapping_add(1); } fn next_task(&mut self, handle: &Handle) -> Option { if self.tick % self.global_queue_interval == 0 { handle .next_remote_task() .or_else(|| self.next_local_task(handle)) } else { self.next_local_task(handle) .or_else(|| handle.next_remote_task()) } } fn next_local_task(&mut self, handle: &Handle) -> Option { let ret = self.tasks.pop_front(); handle .shared .worker_metrics .set_queue_depth(self.tasks.len()); ret } fn push_task(&mut self, handle: &Handle, task: Notified) { self.tasks.push_back(task); self.metrics.inc_local_schedule_count(); handle .shared .worker_metrics .set_queue_depth(self.tasks.len()); } fn submit_metrics(&mut self, handle: &Handle) { self.metrics.submit(&handle.shared.worker_metrics, 0); } } #[cfg(tokio_taskdump)] fn wake_deferred_tasks_and_free(context: &Context) { let wakers = context.defer.take_deferred(); for waker in wakers { waker.wake(); } } // ===== impl Context ===== impl Context { /// Execute the closure with the given scheduler core stored in the /// thread-local context. fn run_task(&self, mut core: Box, f: impl FnOnce() -> R) -> (Box, R) { core.metrics.start_poll(); let mut ret = self.enter(core, || crate::runtime::coop::budget(f)); ret.0.metrics.end_poll(); ret } /// Blocks the current thread until an event is received by the driver, /// including I/O events, timer events, ... fn park(&self, mut core: Box, handle: &Handle) -> Box { let mut driver = core.driver.take().expect("driver missing"); if let Some(f) = &handle.shared.config.before_park { let (c, ()) = self.enter(core, || f()); core = c; } // This check will fail if `before_park` spawns a task for us to run // instead of parking the thread if core.tasks.is_empty() { // Park until the thread is signaled core.metrics.about_to_park(); core.submit_metrics(handle); let (c, ()) = self.enter(core, || { driver.park(&handle.driver); self.defer.wake(); }); core = c; core.metrics.unparked(); core.submit_metrics(handle); } if let Some(f) = &handle.shared.config.after_unpark { let (c, ()) = self.enter(core, || f()); core = c; } core.driver = Some(driver); core } /// Checks the driver for new events without blocking the thread. fn park_yield(&self, mut core: Box, handle: &Handle) -> Box { let mut driver = core.driver.take().expect("driver missing"); core.submit_metrics(handle); let (mut core, ()) = self.enter(core, || { driver.park_timeout(&handle.driver, Duration::from_millis(0)); self.defer.wake(); }); core.driver = Some(driver); core } fn enter(&self, core: Box, f: impl FnOnce() -> R) -> (Box, R) { // Store the scheduler core in the thread-local context // // A drop-guard is employed at a higher level. *self.core.borrow_mut() = Some(core); // Execute the closure while tracking the execution budget let ret = f(); // Take the scheduler core back let core = self.core.borrow_mut().take().expect("core missing"); (core, ret) } pub(crate) fn defer(&self, waker: &Waker) { self.defer.defer(waker); } } // ===== impl Handle ===== impl Handle { /// Spawns a future onto the `CurrentThread` scheduler pub(crate) fn spawn( me: &Arc, future: F, id: crate::runtime::task::Id, ) -> JoinHandle where F: crate::future::Future + Send + 'static, F::Output: Send + 'static, { let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); me.task_hooks.spawn(&TaskMeta { id, _phantom: Default::default(), }); if let Some(notified) = notified { me.schedule(notified); } handle } /// Spawn a task which isn't safe to send across thread boundaries onto the runtime. /// /// # Safety /// This should only be used when this is a `LocalRuntime` or in another case where the runtime /// provably cannot be driven from or moved to different threads from the one on which the task /// is spawned. pub(crate) unsafe fn spawn_local( me: &Arc, future: F, id: crate::runtime::task::Id, ) -> JoinHandle where F: crate::future::Future + 'static, F::Output: 'static, { let (handle, notified) = me.shared.owned.bind_local(future, me.clone(), id); me.task_hooks.spawn(&TaskMeta { id, _phantom: Default::default(), }); if let Some(notified) = notified { me.schedule(notified); } handle } /// Capture a snapshot of this runtime's state. #[cfg(all( tokio_unstable, tokio_taskdump, target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] pub(crate) fn dump(&self) -> crate::runtime::Dump { use crate::runtime::dump; use task::trace::trace_current_thread; let mut traces = vec![]; // todo: how to make this work outside of a runtime context? context::with_scheduler(|maybe_context| { // drain the local queue let context = if let Some(context) = maybe_context { context.expect_current_thread() } else { return; }; let mut maybe_core = context.core.borrow_mut(); let core = if let Some(core) = maybe_core.as_mut() { core } else { return; }; let local = &mut core.tasks; if self.shared.inject.is_closed() { return; } traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject) .into_iter() .map(|(id, trace)| dump::Task::new(id, trace)) .collect(); // Avoid double borrow panic drop(maybe_core); // Taking a taskdump could wakes every task, but we probably don't want // the `yield_now` vector to be that large under normal circumstances. // Therefore, we free its allocation. wake_deferred_tasks_and_free(context); }); dump::Dump::new(traces) } fn next_remote_task(&self) -> Option { self.shared.inject.pop() } fn waker_ref(me: &Arc) -> WakerRef<'_> { // Set woken to true when enter block_on, ensure outer future // be polled for the first time when enter loop me.shared.woken.store(true, Release); waker_ref(me) } // reset woken to false and return original value pub(crate) fn reset_woken(&self) -> bool { self.shared.woken.swap(false, AcqRel) } pub(crate) fn num_alive_tasks(&self) -> usize { self.shared.owned.num_alive_tasks() } pub(crate) fn injection_queue_depth(&self) -> usize { self.shared.inject.len() } } cfg_unstable_metrics! { impl Handle { pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { &self.shared.scheduler_metrics } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { assert_eq!(0, worker); &self.shared.worker_metrics } pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { self.worker_metrics(worker).queue_depth() } pub(crate) fn num_blocking_threads(&self) -> usize { self.blocking_spawner.num_threads() } pub(crate) fn num_idle_blocking_threads(&self) -> usize { self.blocking_spawner.num_idle_threads() } pub(crate) fn blocking_queue_depth(&self) -> usize { self.blocking_spawner.queue_depth() } cfg_64bit_metrics! { pub(crate) fn spawned_tasks_count(&self) -> u64 { self.shared.owned.spawned_tasks_count() } } } } cfg_unstable! { use std::num::NonZeroU64; impl Handle { pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } } } impl fmt::Debug for Handle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("current_thread::Handle { ... }").finish() } } // ===== impl Shared ===== impl Schedule for Arc { fn release(&self, task: &Task) -> Option> { self.shared.owned.remove(task) } fn schedule(&self, task: task::Notified) { use scheduler::Context::CurrentThread; context::with_scheduler(|maybe_cx| match maybe_cx { Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => { let mut core = cx.core.borrow_mut(); // If `None`, the runtime is shutting down, so there is no need // to schedule the task. if let Some(core) = core.as_mut() { core.push_task(self, task); } } _ => { // Track that a task was scheduled from **outside** of the runtime. self.shared.scheduler_metrics.inc_remote_schedule_count(); // Schedule the task self.shared.inject.push(task); self.driver.unpark(); } }); } fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: self.task_hooks.task_terminate_callback.clone(), } } cfg_unstable! { fn unhandled_panic(&self) { use crate::runtime::UnhandledPanic; match self.shared.config.unhandled_panic { UnhandledPanic::Ignore => { // Do nothing } UnhandledPanic::ShutdownRuntime => { use scheduler::Context::CurrentThread; // This hook is only called from within the runtime, so // `context::with_scheduler` should match with `&self`, i.e. // there is no opportunity for a nested scheduler to be // called. context::with_scheduler(|maybe_cx| match maybe_cx { Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => { let mut core = cx.core.borrow_mut(); // If `None`, the runtime is shutting down, so there is no need to signal shutdown if let Some(core) = core.as_mut() { core.unhandled_panic = true; self.shared.owned.close_and_shutdown_all(0); } } _ => unreachable!("runtime core not set in CURRENT thread-local"), }) } } } } } impl Wake for Handle { fn wake(arc_self: Arc) { Wake::wake_by_ref(&arc_self); } /// Wake by reference fn wake_by_ref(arc_self: &Arc) { arc_self.shared.woken.store(true, Release); arc_self.driver.unpark(); } } // ===== CoreGuard ===== /// Used to ensure we always place the `Core` value back into its slot in /// `CurrentThread`, even if the future panics. struct CoreGuard<'a> { context: scheduler::Context, scheduler: &'a CurrentThread, } impl CoreGuard<'_> { #[track_caller] fn block_on(self, future: F) -> F::Output { let ret = self.enter(|mut core, context| { let waker = Handle::waker_ref(&context.handle); let mut cx = std::task::Context::from_waker(&waker); pin!(future); core.metrics.start_processing_scheduled_tasks(); 'outer: loop { let handle = &context.handle; if handle.reset_woken() { let (c, res) = context.enter(core, || { crate::runtime::coop::budget(|| future.as_mut().poll(&mut cx)) }); core = c; if let Ready(v) = res { return (core, Some(v)); } } for _ in 0..handle.shared.config.event_interval { // Make sure we didn't hit an unhandled_panic if core.unhandled_panic { return (core, None); } core.tick(); let entry = core.next_task(handle); let task = match entry { Some(entry) => entry, None => { core.metrics.end_processing_scheduled_tasks(); core = if !context.defer.is_empty() { context.park_yield(core, handle) } else { context.park(core, handle) }; core.metrics.start_processing_scheduled_tasks(); // Try polling the `block_on` future next continue 'outer; } }; let task = context.handle.shared.owned.assert_owner(task); let (c, ()) = context.run_task(core, || { task.run(); }); core = c; } core.metrics.end_processing_scheduled_tasks(); // Yield to the driver, this drives the timer and pulls any // pending I/O events. core = context.park_yield(core, handle); core.metrics.start_processing_scheduled_tasks(); } }); match ret { Some(ret) => ret, None => { // `block_on` panicked. panic!("a spawned task panicked and the runtime is configured to shut down on unhandled panic"); } } } /// Enters the scheduler context. This sets the queue and other necessary /// scheduler state in the thread-local. fn enter(self, f: F) -> R where F: FnOnce(Box, &Context) -> (Box, R), { let context = self.context.expect_current_thread(); // Remove `core` from `context` to pass into the closure. let core = context.core.borrow_mut().take().expect("core missing"); // Call the closure and place `core` back let (core, ret) = context::set_scheduler(&self.context, || f(core, context)); *context.core.borrow_mut() = Some(core); ret } } impl Drop for CoreGuard<'_> { fn drop(&mut self) { let context = self.context.expect_current_thread(); if let Some(core) = context.core.borrow_mut().take() { // Replace old scheduler back into the state to allow // other threads to pick it up and drive it. self.scheduler.core.set(core); // Wake up other possible threads that could steal the driver. self.scheduler.notify.notify_one(); } } } tokio-1.43.1/src/runtime/scheduler/defer.rs000064400000000000000000000017711046102023000167300ustar 00000000000000use std::cell::RefCell; use std::task::Waker; pub(crate) struct Defer { deferred: RefCell>, } impl Defer { pub(crate) fn new() -> Defer { Defer { deferred: RefCell::default(), } } pub(crate) fn defer(&self, waker: &Waker) { let mut deferred = self.deferred.borrow_mut(); // If the same task adds itself a bunch of times, then only add it once. if let Some(last) = deferred.last() { if last.will_wake(waker) { return; } } deferred.push(waker.clone()); } pub(crate) fn is_empty(&self) -> bool { self.deferred.borrow().is_empty() } pub(crate) fn wake(&self) { while let Some(waker) = self.deferred.borrow_mut().pop() { waker.wake(); } } #[cfg(tokio_taskdump)] pub(crate) fn take_deferred(&self) -> Vec { let mut deferred = self.deferred.borrow_mut(); std::mem::take(&mut *deferred) } } tokio-1.43.1/src/runtime/scheduler/inject/metrics.rs000064400000000000000000000001731046102023000205600ustar 00000000000000use super::Inject; impl Inject { pub(crate) fn len(&self) -> usize { self.shared.len() } } tokio-1.43.1/src/runtime/scheduler/inject/pop.rs000064400000000000000000000020541046102023000177100ustar 00000000000000use super::Synced; use crate::runtime::task; use std::marker::PhantomData; pub(crate) struct Pop<'a, T: 'static> { len: usize, synced: &'a mut Synced, _p: PhantomData, } impl<'a, T: 'static> Pop<'a, T> { pub(super) fn new(len: usize, synced: &'a mut Synced) -> Pop<'a, T> { Pop { len, synced, _p: PhantomData, } } } impl<'a, T: 'static> Iterator for Pop<'a, T> { type Item = task::Notified; fn next(&mut self) -> Option { if self.len == 0 { return None; } let ret = self.synced.pop(); // Should be `Some` when `len > 0` debug_assert!(ret.is_some()); self.len -= 1; ret } fn size_hint(&self) -> (usize, Option) { (self.len, Some(self.len)) } } impl<'a, T: 'static> ExactSizeIterator for Pop<'a, T> { fn len(&self) -> usize { self.len } } impl<'a, T: 'static> Drop for Pop<'a, T> { fn drop(&mut self) { for _ in self.by_ref() {} } } tokio-1.43.1/src/runtime/scheduler/inject/rt_multi_thread.rs000064400000000000000000000057351046102023000223110ustar 00000000000000use super::{Shared, Synced}; use crate::runtime::scheduler::Lock; use crate::runtime::task; use std::sync::atomic::Ordering::Release; impl<'a> Lock for &'a mut Synced { type Handle = &'a mut Synced; fn lock(self) -> Self::Handle { self } } impl AsMut for Synced { fn as_mut(&mut self) -> &mut Synced { self } } impl Shared { /// Pushes several values into the queue. /// /// # Safety /// /// Must be called with the same `Synced` instance returned by `Inject::new` #[inline] pub(crate) unsafe fn push_batch(&self, shared: L, mut iter: I) where L: Lock, I: Iterator>, { let first = match iter.next() { Some(first) => first.into_raw(), None => return, }; // Link up all the tasks. let mut prev = first; let mut counter = 1; // We are going to be called with an `std::iter::Chain`, and that // iterator overrides `for_each` to something that is easier for the // compiler to optimize than a loop. iter.for_each(|next| { let next = next.into_raw(); // safety: Holding the Notified for a task guarantees exclusive // access to the `queue_next` field. unsafe { prev.set_queue_next(Some(next)) }; prev = next; counter += 1; }); // Now that the tasks are linked together, insert them into the // linked list. self.push_batch_inner(shared, first, prev, counter); } /// Inserts several tasks that have been linked together into the queue. /// /// The provided head and tail may be be the same task. In this case, a /// single task is inserted. #[inline] unsafe fn push_batch_inner( &self, shared: L, batch_head: task::RawTask, batch_tail: task::RawTask, num: usize, ) where L: Lock, { debug_assert!(unsafe { batch_tail.get_queue_next().is_none() }); let mut synced = shared.lock(); if synced.as_mut().is_closed { drop(synced); let mut curr = Some(batch_head); while let Some(task) = curr { curr = task.get_queue_next(); let _ = unsafe { task::Notified::::from_raw(task) }; } return; } let synced = synced.as_mut(); if let Some(tail) = synced.tail { unsafe { tail.set_queue_next(Some(batch_head)); } } else { synced.head = Some(batch_head); } synced.tail = Some(batch_tail); // Increment the count. // // safety: All updates to the len atomic are guarded by the mutex. As // such, a non-atomic load followed by a store is safe. let len = self.len.unsync_load(); self.len.store(len + num, Release); } } tokio-1.43.1/src/runtime/scheduler/inject/shared.rs000064400000000000000000000064301046102023000203620ustar 00000000000000use super::{Pop, Synced}; use crate::loom::sync::atomic::AtomicUsize; use crate::runtime::task; use std::marker::PhantomData; use std::sync::atomic::Ordering::{Acquire, Release}; pub(crate) struct Shared { /// Number of pending tasks in the queue. This helps prevent unnecessary /// locking in the hot path. pub(super) len: AtomicUsize, _p: PhantomData, } unsafe impl Send for Shared {} unsafe impl Sync for Shared {} impl Shared { pub(crate) fn new() -> (Shared, Synced) { let inject = Shared { len: AtomicUsize::new(0), _p: PhantomData, }; let synced = Synced { is_closed: false, head: None, tail: None, }; (inject, synced) } pub(crate) fn is_empty(&self) -> bool { self.len() == 0 } // Kind of annoying to have to include the cfg here #[cfg(any(tokio_taskdump, feature = "rt-multi-thread"))] pub(crate) fn is_closed(&self, synced: &Synced) -> bool { synced.is_closed } /// Closes the injection queue, returns `true` if the queue is open when the /// transition is made. pub(crate) fn close(&self, synced: &mut Synced) -> bool { if synced.is_closed { return false; } synced.is_closed = true; true } pub(crate) fn len(&self) -> usize { self.len.load(Acquire) } /// Pushes a value into the queue. /// /// This does nothing if the queue is closed. /// /// # Safety /// /// Must be called with the same `Synced` instance returned by `Inject::new` pub(crate) unsafe fn push(&self, synced: &mut Synced, task: task::Notified) { if synced.is_closed { return; } // safety: only mutated with the lock held let len = self.len.unsync_load(); let task = task.into_raw(); // The next pointer should already be null debug_assert!(unsafe { task.get_queue_next().is_none() }); if let Some(tail) = synced.tail { // safety: Holding the Notified for a task guarantees exclusive // access to the `queue_next` field. unsafe { tail.set_queue_next(Some(task)) }; } else { synced.head = Some(task); } synced.tail = Some(task); self.len.store(len + 1, Release); } /// Pop a value from the queue. /// /// # Safety /// /// Must be called with the same `Synced` instance returned by `Inject::new` pub(crate) unsafe fn pop(&self, synced: &mut Synced) -> Option> { self.pop_n(synced, 1).next() } /// Pop `n` values from the queue /// /// # Safety /// /// Must be called with the same `Synced` instance returned by `Inject::new` pub(crate) unsafe fn pop_n<'a>(&'a self, synced: &'a mut Synced, n: usize) -> Pop<'a, T> { use std::cmp; debug_assert!(n > 0); // safety: All updates to the len atomic are guarded by the mutex. As // such, a non-atomic load followed by a store is safe. let len = self.len.unsync_load(); let n = cmp::min(n, len); // Decrement the count. self.len.store(len - n, Release); Pop::new(n, synced) } } tokio-1.43.1/src/runtime/scheduler/inject/synced.rs000064400000000000000000000017061046102023000204020ustar 00000000000000#![cfg_attr( any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"), allow(dead_code) )] use crate::runtime::task; pub(crate) struct Synced { /// True if the queue is closed. pub(super) is_closed: bool, /// Linked-list head. pub(super) head: Option, /// Linked-list tail. pub(super) tail: Option, } unsafe impl Send for Synced {} unsafe impl Sync for Synced {} impl Synced { pub(super) fn pop(&mut self) -> Option> { let task = self.head?; self.head = unsafe { task.get_queue_next() }; if self.head.is_none() { self.tail = None; } unsafe { task.set_queue_next(None) }; // safety: a `Notified` is pushed into the queue and now it is popped! Some(unsafe { task::Notified::from_raw(task) }) } pub(crate) fn is_empty(&self) -> bool { self.head.is_none() } } tokio-1.43.1/src/runtime/scheduler/inject.rs000064400000000000000000000034111046102023000171100ustar 00000000000000//! Inject queue used to send wakeups to a work-stealing scheduler use crate::loom::sync::Mutex; use crate::runtime::task; mod pop; pub(crate) use pop::Pop; mod shared; pub(crate) use shared::Shared; mod synced; pub(crate) use synced::Synced; cfg_rt_multi_thread! { mod rt_multi_thread; } mod metrics; /// Growable, MPMC queue used to inject new tasks into the scheduler and as an /// overflow queue when the local, fixed-size, array queue overflows. pub(crate) struct Inject { shared: Shared, synced: Mutex, } impl Inject { pub(crate) fn new() -> Inject { let (shared, synced) = Shared::new(); Inject { shared, synced: Mutex::new(synced), } } // Kind of annoying to have to include the cfg here #[cfg(tokio_taskdump)] pub(crate) fn is_closed(&self) -> bool { let synced = self.synced.lock(); self.shared.is_closed(&synced) } /// Closes the injection queue, returns `true` if the queue is open when the /// transition is made. pub(crate) fn close(&self) -> bool { let mut synced = self.synced.lock(); self.shared.close(&mut synced) } /// Pushes a value into the queue. /// /// This does nothing if the queue is closed. pub(crate) fn push(&self, task: task::Notified) { let mut synced = self.synced.lock(); // safety: passing correct `Synced` unsafe { self.shared.push(&mut synced, task) } } pub(crate) fn pop(&self) -> Option> { if self.shared.is_empty() { return None; } let mut synced = self.synced.lock(); // safety: passing correct `Synced` unsafe { self.shared.pop(&mut synced) } } } tokio-1.43.1/src/runtime/scheduler/lock.rs000064400000000000000000000002061046102023000165630ustar 00000000000000/// A lock (mutex) yielding generic data. pub(crate) trait Lock { type Handle: AsMut; fn lock(self) -> Self::Handle; } tokio-1.43.1/src/runtime/scheduler/mod.rs000064400000000000000000000246711046102023000164260ustar 00000000000000cfg_rt! { pub(crate) mod current_thread; pub(crate) use current_thread::CurrentThread; mod defer; use defer::Defer; pub(crate) mod inject; pub(crate) use inject::Inject; use crate::runtime::TaskHooks; } cfg_rt_multi_thread! { mod block_in_place; pub(crate) use block_in_place::block_in_place; mod lock; use lock::Lock; pub(crate) mod multi_thread; pub(crate) use multi_thread::MultiThread; cfg_unstable! { pub(crate) mod multi_thread_alt; pub(crate) use multi_thread_alt::MultiThread as MultiThreadAlt; } } use crate::runtime::driver; #[derive(Debug, Clone)] pub(crate) enum Handle { #[cfg(feature = "rt")] CurrentThread(Arc), #[cfg(feature = "rt-multi-thread")] MultiThread(Arc), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] MultiThreadAlt(Arc), // TODO: This is to avoid triggering "dead code" warnings many other places // in the codebase. Remove this during a later cleanup #[cfg(not(feature = "rt"))] #[allow(dead_code)] Disabled, } #[cfg(feature = "rt")] pub(super) enum Context { CurrentThread(current_thread::Context), #[cfg(feature = "rt-multi-thread")] MultiThread(multi_thread::Context), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] MultiThreadAlt(multi_thread_alt::Context), } impl Handle { #[cfg_attr(not(feature = "full"), allow(dead_code))] pub(crate) fn driver(&self) -> &driver::Handle { match *self { #[cfg(feature = "rt")] Handle::CurrentThread(ref h) => &h.driver, #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(ref h) => &h.driver, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(ref h) => &h.driver, #[cfg(not(feature = "rt"))] Handle::Disabled => unreachable!(), } } } cfg_rt! { use crate::future::Future; use crate::loom::sync::Arc; use crate::runtime::{blocking, task::Id}; use crate::runtime::context; use crate::task::JoinHandle; use crate::util::RngSeedGenerator; use std::task::Waker; macro_rules! match_flavor { ($self:expr, $ty:ident($h:ident) => $e:expr) => { match $self { $ty::CurrentThread($h) => $e, #[cfg(feature = "rt-multi-thread")] $ty::MultiThread($h) => $e, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] $ty::MultiThreadAlt($h) => $e, } } } impl Handle { #[track_caller] pub(crate) fn current() -> Handle { match context::with_current(Clone::clone) { Ok(handle) => handle, Err(e) => panic!("{}", e), } } pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner { match_flavor!(self, Handle(h) => &h.blocking_spawner) } pub(crate) fn is_local(&self) -> bool { match self { Handle::CurrentThread(h) => h.local_tid.is_some(), #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(_) => false, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(_) => false, } } /// Returns true if this is a local runtime and the runtime is owned by the current thread. pub(crate) fn can_spawn_local_on_local_runtime(&self) -> bool { match self { Handle::CurrentThread(h) => h.local_tid.map(|x| std::thread::current().id() == x).unwrap_or(false), #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(_) => false, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(_) => false, } } pub(crate) fn spawn(&self, future: F, id: Id) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, { match self { Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id), #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(h) => multi_thread_alt::Handle::spawn(h, future, id), } } /// Spawn a local task /// /// # Safety /// This should only be called in `LocalRuntime` if the runtime has been verified to be owned /// by the current thread. #[allow(irrefutable_let_patterns)] pub(crate) unsafe fn spawn_local(&self, future: F, id: Id) -> JoinHandle where F: Future + 'static, F::Output: 'static, { if let Handle::CurrentThread(h) = self { current_thread::Handle::spawn_local(h, future, id) } else { panic!("Only current_thread and LocalSet have spawn_local internals implemented") } } pub(crate) fn shutdown(&self) { match *self { Handle::CurrentThread(_) => {}, #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(ref h) => h.shutdown(), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(ref h) => h.shutdown(), } } pub(crate) fn seed_generator(&self) -> &RngSeedGenerator { match_flavor!(self, Handle(h) => &h.seed_generator) } pub(crate) fn as_current_thread(&self) -> &Arc { match self { Handle::CurrentThread(handle) => handle, #[cfg(feature = "rt-multi-thread")] _ => panic!("not a CurrentThread handle"), } } pub(crate) fn hooks(&self) -> &TaskHooks { match self { Handle::CurrentThread(h) => &h.task_hooks, #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(h) => &h.task_hooks, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(h) => &h.task_hooks, } } cfg_rt_multi_thread! { cfg_unstable! { pub(crate) fn expect_multi_thread_alt(&self) -> &Arc { match self { Handle::MultiThreadAlt(handle) => handle, _ => panic!("not a `MultiThreadAlt` handle"), } } } } } impl Handle { pub(crate) fn num_workers(&self) -> usize { match self { Handle::CurrentThread(_) => 1, #[cfg(feature = "rt-multi-thread")] Handle::MultiThread(handle) => handle.num_workers(), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Handle::MultiThreadAlt(handle) => handle.num_workers(), } } pub(crate) fn num_alive_tasks(&self) -> usize { match_flavor!(self, Handle(handle) => handle.num_alive_tasks()) } pub(crate) fn injection_queue_depth(&self) -> usize { match_flavor!(self, Handle(handle) => handle.injection_queue_depth()) } } cfg_unstable_metrics! { use crate::runtime::{SchedulerMetrics, WorkerMetrics}; impl Handle { cfg_64bit_metrics! { pub(crate) fn spawned_tasks_count(&self) -> u64 { match_flavor!(self, Handle(handle) => handle.spawned_tasks_count()) } } pub(crate) fn num_blocking_threads(&self) -> usize { match_flavor!(self, Handle(handle) => handle.num_blocking_threads()) } pub(crate) fn num_idle_blocking_threads(&self) -> usize { match_flavor!(self, Handle(handle) => handle.num_idle_blocking_threads()) } pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { match_flavor!(self, Handle(handle) => handle.scheduler_metrics()) } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { match_flavor!(self, Handle(handle) => handle.worker_metrics(worker)) } pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { match_flavor!(self, Handle(handle) => handle.worker_local_queue_depth(worker)) } pub(crate) fn blocking_queue_depth(&self) -> usize { match_flavor!(self, Handle(handle) => handle.blocking_queue_depth()) } } } impl Context { #[track_caller] pub(crate) fn expect_current_thread(&self) -> ¤t_thread::Context { match self { Context::CurrentThread(context) => context, #[cfg(feature = "rt-multi-thread")] _ => panic!("expected `CurrentThread::Context`") } } pub(crate) fn defer(&self, waker: &Waker) { match_flavor!(self, Context(context) => context.defer(waker)); } cfg_rt_multi_thread! { #[track_caller] pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context { match self { Context::MultiThread(context) => context, _ => panic!("expected `MultiThread::Context`") } } cfg_unstable! { #[track_caller] pub(crate) fn expect_multi_thread_alt(&self) -> &multi_thread_alt::Context { match self { Context::MultiThreadAlt(context) => context, _ => panic!("expected `MultiThreadAlt::Context`") } } } } } } cfg_not_rt! { #[cfg(any( feature = "net", all(unix, feature = "process"), all(unix, feature = "signal"), feature = "time", ))] impl Handle { #[track_caller] pub(crate) fn current() -> Handle { panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR) } } } tokio-1.43.1/src/runtime/scheduler/multi_thread/counters.rs000064400000000000000000000037331046102023000221660ustar 00000000000000#[cfg(tokio_internal_mt_counters)] mod imp { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::Relaxed; static NUM_MAINTENANCE: AtomicUsize = AtomicUsize::new(0); static NUM_NOTIFY_LOCAL: AtomicUsize = AtomicUsize::new(0); static NUM_UNPARKS_LOCAL: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_SCHEDULES: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_CAPPED: AtomicUsize = AtomicUsize::new(0); impl Drop for super::Counters { fn drop(&mut self) { let notifies_local = NUM_NOTIFY_LOCAL.load(Relaxed); let unparks_local = NUM_UNPARKS_LOCAL.load(Relaxed); let maintenance = NUM_MAINTENANCE.load(Relaxed); let lifo_scheds = NUM_LIFO_SCHEDULES.load(Relaxed); let lifo_capped = NUM_LIFO_CAPPED.load(Relaxed); println!("---"); println!("notifies (local): {}", notifies_local); println!(" unparks (local): {}", unparks_local); println!(" maintenance: {}", maintenance); println!(" LIFO schedules: {}", lifo_scheds); println!(" LIFO capped: {}", lifo_capped); } } pub(crate) fn inc_num_inc_notify_local() { NUM_NOTIFY_LOCAL.fetch_add(1, Relaxed); } pub(crate) fn inc_num_unparks_local() { NUM_UNPARKS_LOCAL.fetch_add(1, Relaxed); } pub(crate) fn inc_num_maintenance() { NUM_MAINTENANCE.fetch_add(1, Relaxed); } pub(crate) fn inc_lifo_schedules() { NUM_LIFO_SCHEDULES.fetch_add(1, Relaxed); } pub(crate) fn inc_lifo_capped() { NUM_LIFO_CAPPED.fetch_add(1, Relaxed); } } #[cfg(not(tokio_internal_mt_counters))] mod imp { pub(crate) fn inc_num_inc_notify_local() {} pub(crate) fn inc_num_unparks_local() {} pub(crate) fn inc_num_maintenance() {} pub(crate) fn inc_lifo_schedules() {} pub(crate) fn inc_lifo_capped() {} } #[derive(Debug)] pub(crate) struct Counters; pub(super) use imp::*; tokio-1.43.1/src/runtime/scheduler/multi_thread/handle/metrics.rs000064400000000000000000000030161046102023000232170ustar 00000000000000use super::Handle; cfg_unstable_metrics! { use crate::runtime::{SchedulerMetrics, WorkerMetrics}; } impl Handle { pub(crate) fn num_workers(&self) -> usize { self.shared.worker_metrics.len() } pub(crate) fn num_alive_tasks(&self) -> usize { self.shared.owned.num_alive_tasks() } pub(crate) fn injection_queue_depth(&self) -> usize { self.shared.injection_queue_depth() } cfg_unstable_metrics! { cfg_64bit_metrics! { pub(crate) fn spawned_tasks_count(&self) -> u64 { self.shared.owned.spawned_tasks_count() } } pub(crate) fn num_blocking_threads(&self) -> usize { // workers are currently spawned using spawn_blocking self.blocking_spawner .num_threads() .saturating_sub(self.num_workers()) } pub(crate) fn num_idle_blocking_threads(&self) -> usize { self.blocking_spawner.num_idle_threads() } pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { &self.shared.scheduler_metrics } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { &self.shared.worker_metrics[worker] } pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { self.shared.worker_local_queue_depth(worker) } pub(crate) fn blocking_queue_depth(&self) -> usize { self.blocking_spawner.queue_depth() } } } tokio-1.43.1/src/runtime/scheduler/multi_thread/handle/taskdump.rs000064400000000000000000000012121046102023000233750ustar 00000000000000use super::Handle; use crate::runtime::Dump; impl Handle { pub(crate) async fn dump(&self) -> Dump { let trace_status = &self.shared.trace_status; // If a dump is in progress, block. trace_status.start_trace_request(&self).await; let result = loop { if let Some(result) = trace_status.take_result() { break result; } else { self.notify_all(); trace_status.result_ready.notified().await; } }; // Allow other queued dumps to proceed. trace_status.end_trace_request(&self).await; result } } tokio-1.43.1/src/runtime/scheduler/multi_thread/handle.rs000064400000000000000000000036371046102023000215620ustar 00000000000000use crate::future::Future; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread::worker; use crate::runtime::{ blocking, driver, task::{self, JoinHandle}, TaskHooks, TaskMeta, }; use crate::util::RngSeedGenerator; use std::fmt; mod metrics; cfg_taskdump! { mod taskdump; } /// Handle to the multi thread scheduler pub(crate) struct Handle { /// Task spawner pub(super) shared: worker::Shared, /// Resource driver handles pub(crate) driver: driver::Handle, /// Blocking pool spawner pub(crate) blocking_spawner: blocking::Spawner, /// Current random number generator seed pub(crate) seed_generator: RngSeedGenerator, /// User-supplied hooks to invoke for things pub(crate) task_hooks: TaskHooks, } impl Handle { /// Spawns a future onto the thread pool pub(crate) fn spawn(me: &Arc, future: F, id: task::Id) -> JoinHandle where F: crate::future::Future + Send + 'static, F::Output: Send + 'static, { Self::bind_new_task(me, future, id) } pub(crate) fn shutdown(&self) { self.close(); } pub(super) fn bind_new_task(me: &Arc, future: T, id: task::Id) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); me.task_hooks.spawn(&TaskMeta { id, _phantom: Default::default(), }); me.schedule_option_task_without_yield(notified); handle } } cfg_unstable! { use std::num::NonZeroU64; impl Handle { pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } } } impl fmt::Debug for Handle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("multi_thread::Handle { ... }").finish() } } tokio-1.43.1/src/runtime/scheduler/multi_thread/idle.rs000064400000000000000000000160171046102023000212400ustar 00000000000000//! Coordinates idling workers use crate::loom::sync::atomic::AtomicUsize; use crate::runtime::scheduler::multi_thread::Shared; use std::fmt; use std::sync::atomic::Ordering::{self, SeqCst}; pub(super) struct Idle { /// Tracks both the number of searching workers and the number of unparked /// workers. /// /// Used as a fast-path to avoid acquiring the lock when needed. state: AtomicUsize, /// Total number of workers. num_workers: usize, } /// Data synchronized by the scheduler mutex pub(super) struct Synced { /// Sleeping workers sleepers: Vec, } const UNPARK_SHIFT: usize = 16; const UNPARK_MASK: usize = !SEARCH_MASK; const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1; #[derive(Copy, Clone)] struct State(usize); impl Idle { pub(super) fn new(num_workers: usize) -> (Idle, Synced) { let init = State::new(num_workers); let idle = Idle { state: AtomicUsize::new(init.into()), num_workers, }; let synced = Synced { sleepers: Vec::with_capacity(num_workers), }; (idle, synced) } /// If there are no workers actively searching, returns the index of a /// worker currently sleeping. pub(super) fn worker_to_notify(&self, shared: &Shared) -> Option { // If at least one worker is spinning, work being notified will // eventually be found. A searching thread will find **some** work and // notify another worker, eventually leading to our work being found. // // For this to happen, this load must happen before the thread // transitioning `num_searching` to zero. Acquire / Release does not // provide sufficient guarantees, so this load is done with `SeqCst` and // will pair with the `fetch_sub(1)` when transitioning out of // searching. if !self.notify_should_wakeup() { return None; } // Acquire the lock let mut lock = shared.synced.lock(); // Check again, now that the lock is acquired if !self.notify_should_wakeup() { return None; } // A worker should be woken up, atomically increment the number of // searching workers as well as the number of unparked workers. State::unpark_one(&self.state, 1); // Get the worker to unpark let ret = lock.idle.sleepers.pop(); debug_assert!(ret.is_some()); ret } /// Returns `true` if the worker needs to do a final check for submitted /// work. pub(super) fn transition_worker_to_parked( &self, shared: &Shared, worker: usize, is_searching: bool, ) -> bool { // Acquire the lock let mut lock = shared.synced.lock(); // Decrement the number of unparked threads let ret = State::dec_num_unparked(&self.state, is_searching); // Track the sleeping worker lock.idle.sleepers.push(worker); ret } pub(super) fn transition_worker_to_searching(&self) -> bool { let state = State::load(&self.state, SeqCst); if 2 * state.num_searching() >= self.num_workers { return false; } // It is possible for this routine to allow more than 50% of the workers // to search. That is OK. Limiting searchers is only an optimization to // prevent too much contention. State::inc_num_searching(&self.state, SeqCst); true } /// A lightweight transition from searching -> running. /// /// Returns `true` if this is the final searching worker. The caller /// **must** notify a new worker. pub(super) fn transition_worker_from_searching(&self) -> bool { State::dec_num_searching(&self.state) } /// Unpark a specific worker. This happens if tasks are submitted from /// within the worker's park routine. /// /// Returns `true` if the worker was parked before calling the method. pub(super) fn unpark_worker_by_id(&self, shared: &Shared, worker_id: usize) -> bool { let mut lock = shared.synced.lock(); let sleepers = &mut lock.idle.sleepers; for index in 0..sleepers.len() { if sleepers[index] == worker_id { sleepers.swap_remove(index); // Update the state accordingly while the lock is held. State::unpark_one(&self.state, 0); return true; } } false } /// Returns `true` if `worker_id` is contained in the sleep set. pub(super) fn is_parked(&self, shared: &Shared, worker_id: usize) -> bool { let lock = shared.synced.lock(); lock.idle.sleepers.contains(&worker_id) } fn notify_should_wakeup(&self) -> bool { let state = State(self.state.fetch_add(0, SeqCst)); state.num_searching() == 0 && state.num_unparked() < self.num_workers } } impl State { fn new(num_workers: usize) -> State { // All workers start in the unparked state let ret = State(num_workers << UNPARK_SHIFT); debug_assert_eq!(num_workers, ret.num_unparked()); debug_assert_eq!(0, ret.num_searching()); ret } fn load(cell: &AtomicUsize, ordering: Ordering) -> State { State(cell.load(ordering)) } fn unpark_one(cell: &AtomicUsize, num_searching: usize) { cell.fetch_add(num_searching | (1 << UNPARK_SHIFT), SeqCst); } fn inc_num_searching(cell: &AtomicUsize, ordering: Ordering) { cell.fetch_add(1, ordering); } /// Returns `true` if this is the final searching worker fn dec_num_searching(cell: &AtomicUsize) -> bool { let state = State(cell.fetch_sub(1, SeqCst)); state.num_searching() == 1 } /// Track a sleeping worker /// /// Returns `true` if this is the final searching worker. fn dec_num_unparked(cell: &AtomicUsize, is_searching: bool) -> bool { let mut dec = 1 << UNPARK_SHIFT; if is_searching { dec += 1; } let prev = State(cell.fetch_sub(dec, SeqCst)); is_searching && prev.num_searching() == 1 } /// Number of workers currently searching fn num_searching(self) -> usize { self.0 & SEARCH_MASK } /// Number of workers currently unparked fn num_unparked(self) -> usize { (self.0 & UNPARK_MASK) >> UNPARK_SHIFT } } impl From for State { fn from(src: usize) -> State { State(src) } } impl From for usize { fn from(src: State) -> usize { src.0 } } impl fmt::Debug for State { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("worker::State") .field("num_unparked", &self.num_unparked()) .field("num_searching", &self.num_searching()) .finish() } } #[test] fn test_state() { assert_eq!(0, UNPARK_MASK & SEARCH_MASK); assert_eq!(0, !(UNPARK_MASK | SEARCH_MASK)); let state = State::new(10); assert_eq!(10, state.num_unparked()); assert_eq!(0, state.num_searching()); } tokio-1.43.1/src/runtime/scheduler/multi_thread/mod.rs000064400000000000000000000044711046102023000211030ustar 00000000000000//! Multi-threaded runtime mod counters; use counters::Counters; mod handle; pub(crate) use handle::Handle; mod overflow; pub(crate) use overflow::Overflow; mod idle; use self::idle::Idle; mod stats; pub(crate) use stats::Stats; mod park; pub(crate) use park::{Parker, Unparker}; pub(crate) mod queue; mod worker; pub(crate) use worker::{Context, Launch, Shared}; cfg_taskdump! { mod trace; use trace::TraceStatus; pub(crate) use worker::Synced; } cfg_not_taskdump! { mod trace_mock; use trace_mock::TraceStatus; } pub(crate) use worker::block_in_place; use crate::loom::sync::Arc; use crate::runtime::{ blocking, driver::{self, Driver}, scheduler, Config, }; use crate::util::RngSeedGenerator; use std::fmt; use std::future::Future; /// Work-stealing based thread pool for executing futures. pub(crate) struct MultiThread; // ===== impl MultiThread ===== impl MultiThread { pub(crate) fn new( size: usize, driver: Driver, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, ) -> (MultiThread, Arc, Launch) { let parker = Parker::new(driver); let (handle, launch) = worker::create( size, parker, driver_handle, blocking_spawner, seed_generator, config, ); (MultiThread, handle, launch) } /// Blocks the current thread waiting for the future to complete. /// /// The future will execute on the current thread, but all spawned tasks /// will be executed on the thread pool. pub(crate) fn block_on(&self, handle: &scheduler::Handle, future: F) -> F::Output where F: Future, { crate::runtime::context::enter_runtime(handle, true, |blocking| { blocking.block_on(future).expect("failed to park thread") }) } pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) { match handle { scheduler::Handle::MultiThread(handle) => handle.shutdown(), _ => panic!("expected MultiThread scheduler"), } } } impl fmt::Debug for MultiThread { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("MultiThread").finish() } } tokio-1.43.1/src/runtime/scheduler/multi_thread/overflow.rs000064400000000000000000000010671046102023000221650ustar 00000000000000use crate::runtime::task; #[cfg(test)] use std::cell::RefCell; pub(crate) trait Overflow { fn push(&self, task: task::Notified); fn push_batch(&self, iter: I) where I: Iterator>; } #[cfg(test)] impl Overflow for RefCell>> { fn push(&self, task: task::Notified) { self.borrow_mut().push(task); } fn push_batch(&self, iter: I) where I: Iterator>, { self.borrow_mut().extend(iter); } } tokio-1.43.1/src/runtime/scheduler/multi_thread/park.rs000064400000000000000000000175201046102023000212600ustar 00000000000000//! Parks the runtime. //! //! A combination of the various resource driver park handles. use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Arc, Condvar, Mutex}; use crate::runtime::driver::{self, Driver}; use crate::util::TryLock; use std::sync::atomic::Ordering::SeqCst; use std::time::Duration; #[cfg(loom)] use crate::runtime::park::CURRENT_THREAD_PARK_COUNT; pub(crate) struct Parker { inner: Arc, } pub(crate) struct Unparker { inner: Arc, } struct Inner { /// Avoids entering the park if possible state: AtomicUsize, /// Used to coordinate access to the driver / `condvar` mutex: Mutex<()>, /// `Condvar` to block on if the driver is unavailable. condvar: Condvar, /// Resource (I/O, time, ...) driver shared: Arc, } const EMPTY: usize = 0; const PARKED_CONDVAR: usize = 1; const PARKED_DRIVER: usize = 2; const NOTIFIED: usize = 3; /// Shared across multiple Parker handles struct Shared { /// Shared driver. Only one thread at a time can use this driver: TryLock, } impl Parker { pub(crate) fn new(driver: Driver) -> Parker { Parker { inner: Arc::new(Inner { state: AtomicUsize::new(EMPTY), mutex: Mutex::new(()), condvar: Condvar::new(), shared: Arc::new(Shared { driver: TryLock::new(driver), }), }), } } pub(crate) fn unpark(&self) -> Unparker { Unparker { inner: self.inner.clone(), } } pub(crate) fn park(&mut self, handle: &driver::Handle) { self.inner.park(handle); } pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { // Only parking with zero is supported... assert_eq!(duration, Duration::from_millis(0)); if let Some(mut driver) = self.inner.shared.driver.try_lock() { driver.park_timeout(handle, duration); } else { // https://github.com/tokio-rs/tokio/issues/6536 // Hacky, but it's just for loom tests. The counter gets incremented during // `park_timeout`, but we still have to increment the counter if we can't acquire the // lock. #[cfg(loom)] CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst)); } } pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { self.inner.shutdown(handle); } } impl Clone for Parker { fn clone(&self) -> Parker { Parker { inner: Arc::new(Inner { state: AtomicUsize::new(EMPTY), mutex: Mutex::new(()), condvar: Condvar::new(), shared: self.inner.shared.clone(), }), } } } impl Unparker { pub(crate) fn unpark(&self, driver: &driver::Handle) { self.inner.unpark(driver); } } impl Inner { /// Parks the current thread for at most `dur`. fn park(&self, handle: &driver::Handle) { // If we were previously notified then we consume this notification and // return quickly. if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { return; } if let Some(mut driver) = self.shared.driver.try_lock() { self.park_driver(&mut driver, handle); } else { self.park_condvar(); } } fn park_condvar(&self) { // Otherwise we need to coordinate going to sleep let mut m = self.mutex.lock(); match self .state .compare_exchange(EMPTY, PARKED_CONDVAR, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read here, even though we know it will be `NOTIFIED`. // This is because `unpark` may have been called again since we read // `NOTIFIED` in the `compare_exchange` above. We must perform an // acquire operation that synchronizes with that `unpark` to observe // any writes it made before the call to unpark. To do that we must // read from the write it made to `state`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } Err(actual) => panic!("inconsistent park state; actual = {actual}"), } loop { m = self.condvar.wait(m).unwrap(); if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { // got a notification return; } // spurious wakeup, go back to sleep } } fn park_driver(&self, driver: &mut Driver, handle: &driver::Handle) { match self .state .compare_exchange(EMPTY, PARKED_DRIVER, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read here, even though we know it will be `NOTIFIED`. // This is because `unpark` may have been called again since we read // `NOTIFIED` in the `compare_exchange` above. We must perform an // acquire operation that synchronizes with that `unpark` to observe // any writes it made before the call to unpark. To do that we must // read from the write it made to `state`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } Err(actual) => panic!("inconsistent park state; actual = {actual}"), } driver.park(handle); match self.state.swap(EMPTY, SeqCst) { NOTIFIED => {} // got a notification, hurray! PARKED_DRIVER => {} // no notification, alas n => panic!("inconsistent park_timeout state: {n}"), } } fn unpark(&self, driver: &driver::Handle) { // To ensure the unparked thread will observe any writes we made before // this call, we must perform a release operation that `park` can // synchronize with. To do that we must write `NOTIFIED` even if `state` // is already `NOTIFIED`. That is why this must be a swap rather than a // compare-and-swap that returns if it reads `NOTIFIED` on failure. match self.state.swap(NOTIFIED, SeqCst) { EMPTY => {} // no one was waiting NOTIFIED => {} // already unparked PARKED_CONDVAR => self.unpark_condvar(), PARKED_DRIVER => driver.unpark(), actual => panic!("inconsistent state in unpark; actual = {actual}"), } } fn unpark_condvar(&self) { // There is a period between when the parked thread sets `state` to // `PARKED` (or last checked `state` in the case of a spurious wake // up) and when it actually waits on `cvar`. If we were to notify // during this period it would be ignored and then when the parked // thread went to sleep it would never wake up. Fortunately, it has // `lock` locked at this stage so we can acquire `lock` to wait until // it is ready to receive the notification. // // Releasing `lock` before the call to `notify_one` means that when the // parked thread wakes it doesn't get woken only to have to wait for us // to release `lock`. drop(self.mutex.lock()); self.condvar.notify_one(); } fn shutdown(&self, handle: &driver::Handle) { if let Some(mut driver) = self.shared.driver.try_lock() { driver.shutdown(handle); } self.condvar.notify_all(); } } tokio-1.43.1/src/runtime/scheduler/multi_thread/queue.rs000064400000000000000000000473351046102023000214560ustar 00000000000000//! Run-queue structures to support a work-stealing scheduler use crate::loom::cell::UnsafeCell; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread::{Overflow, Stats}; use crate::runtime::task; use std::mem::{self, MaybeUninit}; use std::ptr; use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; // Use wider integers when possible to increase ABA resilience. // // See issue #5041: . cfg_has_atomic_u64! { type UnsignedShort = u32; type UnsignedLong = u64; type AtomicUnsignedShort = crate::loom::sync::atomic::AtomicU32; type AtomicUnsignedLong = crate::loom::sync::atomic::AtomicU64; } cfg_not_has_atomic_u64! { type UnsignedShort = u16; type UnsignedLong = u32; type AtomicUnsignedShort = crate::loom::sync::atomic::AtomicU16; type AtomicUnsignedLong = crate::loom::sync::atomic::AtomicU32; } /// Producer handle. May only be used from a single thread. pub(crate) struct Local { inner: Arc>, } /// Consumer handle. May be used from many threads. pub(crate) struct Steal(Arc>); pub(crate) struct Inner { /// Concurrently updated by many threads. /// /// Contains two `UnsignedShort` values. The `LSB` byte is the "real" head of /// the queue. The `UnsignedShort` in the `MSB` is set by a stealer in process /// of stealing values. It represents the first value being stolen in the /// batch. The `UnsignedShort` indices are intentionally wider than strictly /// required for buffer indexing in order to provide ABA mitigation and make /// it possible to distinguish between full and empty buffers. /// /// When both `UnsignedShort` values are the same, there is no active /// stealer. /// /// Tracking an in-progress stealer prevents a wrapping scenario. head: AtomicUnsignedLong, /// Only updated by producer thread but read by many threads. tail: AtomicUnsignedShort, /// Elements buffer: Box<[UnsafeCell>>; LOCAL_QUEUE_CAPACITY]>, } unsafe impl Send for Inner {} unsafe impl Sync for Inner {} #[cfg(not(loom))] const LOCAL_QUEUE_CAPACITY: usize = 256; // Shrink the size of the local queue when using loom. This shouldn't impact // logic, but allows loom to test more edge cases in a reasonable a mount of // time. #[cfg(loom)] const LOCAL_QUEUE_CAPACITY: usize = 4; const MASK: usize = LOCAL_QUEUE_CAPACITY - 1; // Constructing the fixed size array directly is very awkward. The only way to // do it is to repeat `UnsafeCell::new(MaybeUninit::uninit())` 256 times, as // the contents are not Copy. The trick with defining a const doesn't work for // generic types. fn make_fixed_size(buffer: Box<[T]>) -> Box<[T; LOCAL_QUEUE_CAPACITY]> { assert_eq!(buffer.len(), LOCAL_QUEUE_CAPACITY); // safety: We check that the length is correct. unsafe { Box::from_raw(Box::into_raw(buffer).cast()) } } /// Create a new local run-queue pub(crate) fn local() -> (Steal, Local) { let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY); for _ in 0..LOCAL_QUEUE_CAPACITY { buffer.push(UnsafeCell::new(MaybeUninit::uninit())); } let inner = Arc::new(Inner { head: AtomicUnsignedLong::new(0), tail: AtomicUnsignedShort::new(0), buffer: make_fixed_size(buffer.into_boxed_slice()), }); let local = Local { inner: inner.clone(), }; let remote = Steal(inner); (remote, local) } impl Local { /// Returns the number of entries in the queue pub(crate) fn len(&self) -> usize { self.inner.len() as usize } /// How many tasks can be pushed into the queue pub(crate) fn remaining_slots(&self) -> usize { self.inner.remaining_slots() } pub(crate) fn max_capacity(&self) -> usize { LOCAL_QUEUE_CAPACITY } /// Returns false if there are any entries in the queue /// /// Separate to `is_stealable` so that refactors of `is_stealable` to "protect" /// some tasks from stealing won't affect this pub(crate) fn has_tasks(&self) -> bool { !self.inner.is_empty() } /// Pushes a batch of tasks to the back of the queue. All tasks must fit in /// the local queue. /// /// # Panics /// /// The method panics if there is not enough capacity to fit in the queue. pub(crate) fn push_back(&mut self, tasks: impl ExactSizeIterator>) { let len = tasks.len(); assert!(len <= LOCAL_QUEUE_CAPACITY); if len == 0 { // Nothing to do return; } let head = self.inner.head.load(Acquire); let (steal, _) = unpack(head); // safety: this is the **only** thread that updates this cell. let mut tail = unsafe { self.inner.tail.unsync_load() }; if tail.wrapping_sub(steal) <= (LOCAL_QUEUE_CAPACITY - len) as UnsignedShort { // Yes, this if condition is structured a bit weird (first block // does nothing, second returns an error). It is this way to match // `push_back_or_overflow`. } else { panic!() } for task in tasks { let idx = tail as usize & MASK; self.inner.buffer[idx].with_mut(|ptr| { // Write the task to the slot // // Safety: There is only one producer and the above `if` // condition ensures we don't touch a cell if there is a // value, thus no consumer. unsafe { ptr::write((*ptr).as_mut_ptr(), task); } }); tail = tail.wrapping_add(1); } self.inner.tail.store(tail, Release); } /// Pushes a task to the back of the local queue, if there is not enough /// capacity in the queue, this triggers the overflow operation. /// /// When the queue overflows, half of the current contents of the queue is /// moved to the given Injection queue. This frees up capacity for more /// tasks to be pushed into the local queue. pub(crate) fn push_back_or_overflow>( &mut self, mut task: task::Notified, overflow: &O, stats: &mut Stats, ) { let tail = loop { let head = self.inner.head.load(Acquire); let (steal, real) = unpack(head); // safety: this is the **only** thread that updates this cell. let tail = unsafe { self.inner.tail.unsync_load() }; if tail.wrapping_sub(steal) < LOCAL_QUEUE_CAPACITY as UnsignedShort { // There is capacity for the task break tail; } else if steal != real { // Concurrently stealing, this will free up capacity, so only // push the task onto the inject queue overflow.push(task); return; } else { // Push the current task and half of the queue into the // inject queue. match self.push_overflow(task, real, tail, overflow, stats) { Ok(_) => return, // Lost the race, try again Err(v) => { task = v; } } } }; self.push_back_finish(task, tail); } // Second half of `push_back` fn push_back_finish(&self, task: task::Notified, tail: UnsignedShort) { // Map the position to a slot index. let idx = tail as usize & MASK; self.inner.buffer[idx].with_mut(|ptr| { // Write the task to the slot // // Safety: There is only one producer and the above `if` // condition ensures we don't touch a cell if there is a // value, thus no consumer. unsafe { ptr::write((*ptr).as_mut_ptr(), task); } }); // Make the task available. Synchronizes with a load in // `steal_into2`. self.inner.tail.store(tail.wrapping_add(1), Release); } /// Moves a batch of tasks into the inject queue. /// /// This will temporarily make some of the tasks unavailable to stealers. /// Once `push_overflow` is done, a notification is sent out, so if other /// workers "missed" some of the tasks during a steal, they will get /// another opportunity. #[inline(never)] fn push_overflow>( &mut self, task: task::Notified, head: UnsignedShort, tail: UnsignedShort, overflow: &O, stats: &mut Stats, ) -> Result<(), task::Notified> { /// How many elements are we taking from the local queue. /// /// This is one less than the number of tasks pushed to the inject /// queue as we are also inserting the `task` argument. const NUM_TASKS_TAKEN: UnsignedShort = (LOCAL_QUEUE_CAPACITY / 2) as UnsignedShort; assert_eq!( tail.wrapping_sub(head) as usize, LOCAL_QUEUE_CAPACITY, "queue is not full; tail = {tail}; head = {head}" ); let prev = pack(head, head); // Claim a bunch of tasks // // We are claiming the tasks **before** reading them out of the buffer. // This is safe because only the **current** thread is able to push new // tasks. // // There isn't really any need for memory ordering... Relaxed would // work. This is because all tasks are pushed into the queue from the // current thread (or memory has been acquired if the local queue handle // moved). if self .inner .head .compare_exchange( prev, pack( head.wrapping_add(NUM_TASKS_TAKEN), head.wrapping_add(NUM_TASKS_TAKEN), ), Release, Relaxed, ) .is_err() { // We failed to claim the tasks, losing the race. Return out of // this function and try the full `push` routine again. The queue // may not be full anymore. return Err(task); } /// An iterator that takes elements out of the run queue. struct BatchTaskIter<'a, T: 'static> { buffer: &'a [UnsafeCell>>; LOCAL_QUEUE_CAPACITY], head: UnsignedLong, i: UnsignedLong, } impl<'a, T: 'static> Iterator for BatchTaskIter<'a, T> { type Item = task::Notified; #[inline] fn next(&mut self) -> Option> { if self.i == UnsignedLong::from(NUM_TASKS_TAKEN) { None } else { let i_idx = self.i.wrapping_add(self.head) as usize & MASK; let slot = &self.buffer[i_idx]; // safety: Our CAS from before has assumed exclusive ownership // of the task pointers in this range. let task = slot.with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); self.i += 1; Some(task) } } } // safety: The CAS above ensures that no consumer will look at these // values again, and we are the only producer. let batch_iter = BatchTaskIter { buffer: &self.inner.buffer, head: head as UnsignedLong, i: 0, }; overflow.push_batch(batch_iter.chain(std::iter::once(task))); // Add 1 to factor in the task currently being scheduled. stats.incr_overflow_count(); Ok(()) } /// Pops a task from the local queue. pub(crate) fn pop(&mut self) -> Option> { let mut head = self.inner.head.load(Acquire); let idx = loop { let (steal, real) = unpack(head); // safety: this is the **only** thread that updates this cell. let tail = unsafe { self.inner.tail.unsync_load() }; if real == tail { // queue is empty return None; } let next_real = real.wrapping_add(1); // If `steal == real` there are no concurrent stealers. Both `steal` // and `real` are updated. let next = if steal == real { pack(next_real, next_real) } else { assert_ne!(steal, next_real); pack(steal, next_real) }; // Attempt to claim a task. let res = self .inner .head .compare_exchange(head, next, AcqRel, Acquire); match res { Ok(_) => break real as usize & MASK, Err(actual) => head = actual, } }; Some(self.inner.buffer[idx].with(|ptr| unsafe { ptr::read(ptr).assume_init() })) } } impl Steal { pub(crate) fn is_empty(&self) -> bool { self.0.is_empty() } /// Steals half the tasks from self and place them into `dst`. pub(crate) fn steal_into( &self, dst: &mut Local, dst_stats: &mut Stats, ) -> Option> { // Safety: the caller is the only thread that mutates `dst.tail` and // holds a mutable reference. let dst_tail = unsafe { dst.inner.tail.unsync_load() }; // To the caller, `dst` may **look** empty but still have values // contained in the buffer. If another thread is concurrently stealing // from `dst` there may not be enough capacity to steal. let (steal, _) = unpack(dst.inner.head.load(Acquire)); if dst_tail.wrapping_sub(steal) > LOCAL_QUEUE_CAPACITY as UnsignedShort / 2 { // we *could* try to steal less here, but for simplicity, we're just // going to abort. return None; } // Steal the tasks into `dst`'s buffer. This does not yet expose the // tasks in `dst`. let mut n = self.steal_into2(dst, dst_tail); if n == 0 { // No tasks were stolen return None; } dst_stats.incr_steal_count(n as u16); dst_stats.incr_steal_operations(); // We are returning a task here n -= 1; let ret_pos = dst_tail.wrapping_add(n); let ret_idx = ret_pos as usize & MASK; // safety: the value was written as part of `steal_into2` and not // exposed to stealers, so no other thread can access it. let ret = dst.inner.buffer[ret_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); if n == 0 { // The `dst` queue is empty, but a single task was stolen return Some(ret); } // Make the stolen items available to consumers dst.inner.tail.store(dst_tail.wrapping_add(n), Release); Some(ret) } // Steal tasks from `self`, placing them into `dst`. Returns the number of // tasks that were stolen. fn steal_into2(&self, dst: &mut Local, dst_tail: UnsignedShort) -> UnsignedShort { let mut prev_packed = self.0.head.load(Acquire); let mut next_packed; let n = loop { let (src_head_steal, src_head_real) = unpack(prev_packed); let src_tail = self.0.tail.load(Acquire); // If these two do not match, another thread is concurrently // stealing from the queue. if src_head_steal != src_head_real { return 0; } // Number of available tasks to steal let n = src_tail.wrapping_sub(src_head_real); let n = n - n / 2; if n == 0 { // No tasks available to steal return 0; } // Update the real head index to acquire the tasks. let steal_to = src_head_real.wrapping_add(n); assert_ne!(src_head_steal, steal_to); next_packed = pack(src_head_steal, steal_to); // Claim all those tasks. This is done by incrementing the "real" // head but not the steal. By doing this, no other thread is able to // steal from this queue until the current thread completes. let res = self .0 .head .compare_exchange(prev_packed, next_packed, AcqRel, Acquire); match res { Ok(_) => break n, Err(actual) => prev_packed = actual, } }; assert!( n <= LOCAL_QUEUE_CAPACITY as UnsignedShort / 2, "actual = {n}" ); let (first, _) = unpack(next_packed); // Take all the tasks for i in 0..n { // Compute the positions let src_pos = first.wrapping_add(i); let dst_pos = dst_tail.wrapping_add(i); // Map to slots let src_idx = src_pos as usize & MASK; let dst_idx = dst_pos as usize & MASK; // Read the task // // safety: We acquired the task with the atomic exchange above. let task = self.0.buffer[src_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); // Write the task to the new slot // // safety: `dst` queue is empty and we are the only producer to // this queue. dst.inner.buffer[dst_idx] .with_mut(|ptr| unsafe { ptr::write((*ptr).as_mut_ptr(), task) }); } let mut prev_packed = next_packed; // Update `src_head_steal` to match `src_head_real` signalling that the // stealing routine is complete. loop { let head = unpack(prev_packed).1; next_packed = pack(head, head); let res = self .0 .head .compare_exchange(prev_packed, next_packed, AcqRel, Acquire); match res { Ok(_) => return n, Err(actual) => { let (actual_steal, actual_real) = unpack(actual); assert_ne!(actual_steal, actual_real); prev_packed = actual; } } } } } cfg_unstable_metrics! { impl Steal { pub(crate) fn len(&self) -> usize { self.0.len() as _ } } } impl Clone for Steal { fn clone(&self) -> Steal { Steal(self.0.clone()) } } impl Drop for Local { fn drop(&mut self) { if !std::thread::panicking() { assert!(self.pop().is_none(), "queue not empty"); } } } impl Inner { fn remaining_slots(&self) -> usize { let (steal, _) = unpack(self.head.load(Acquire)); let tail = self.tail.load(Acquire); LOCAL_QUEUE_CAPACITY - (tail.wrapping_sub(steal) as usize) } fn len(&self) -> UnsignedShort { let (_, head) = unpack(self.head.load(Acquire)); let tail = self.tail.load(Acquire); tail.wrapping_sub(head) } fn is_empty(&self) -> bool { self.len() == 0 } } /// Split the head value into the real head and the index a stealer is working /// on. fn unpack(n: UnsignedLong) -> (UnsignedShort, UnsignedShort) { let real = n & UnsignedShort::MAX as UnsignedLong; let steal = n >> (mem::size_of::() * 8); (steal as UnsignedShort, real as UnsignedShort) } /// Join the two head values fn pack(steal: UnsignedShort, real: UnsignedShort) -> UnsignedLong { (real as UnsignedLong) | ((steal as UnsignedLong) << (mem::size_of::() * 8)) } #[test] fn test_local_queue_capacity() { assert!(LOCAL_QUEUE_CAPACITY - 1 <= u8::MAX as usize); } tokio-1.43.1/src/runtime/scheduler/multi_thread/stats.rs000064400000000000000000000113641046102023000214610ustar 00000000000000use crate::runtime::{Config, MetricsBatch, WorkerMetrics}; use std::time::{Duration, Instant}; /// Per-worker statistics. This is used for both tuning the scheduler and /// reporting runtime-level metrics/stats. pub(crate) struct Stats { /// The metrics batch used to report runtime-level metrics/stats to the /// user. batch: MetricsBatch, /// Instant at which work last resumed (continued after park). /// /// This duplicates the value stored in `MetricsBatch`. We will unify /// `Stats` and `MetricsBatch` when we stabilize metrics. processing_scheduled_tasks_started_at: Instant, /// Number of tasks polled in the batch of scheduled tasks tasks_polled_in_batch: usize, /// Exponentially-weighted moving average of time spent polling scheduled a /// task. /// /// Tracked in nanoseconds, stored as a `f64` since that is what we use with /// the EWMA calculations task_poll_time_ewma: f64, } /// How to weigh each individual poll time, value is plucked from thin air. const TASK_POLL_TIME_EWMA_ALPHA: f64 = 0.1; /// Ideally, we wouldn't go above this, value is plucked from thin air. const TARGET_GLOBAL_QUEUE_INTERVAL: f64 = Duration::from_micros(200).as_nanos() as f64; /// Max value for the global queue interval. This is 2x the previous default const MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 127; /// This is the previous default const TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 61; impl Stats { pub(crate) fn new(worker_metrics: &WorkerMetrics) -> Stats { // Seed the value with what we hope to see. let task_poll_time_ewma = TARGET_GLOBAL_QUEUE_INTERVAL / TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL as f64; Stats { batch: MetricsBatch::new(worker_metrics), processing_scheduled_tasks_started_at: Instant::now(), tasks_polled_in_batch: 0, task_poll_time_ewma, } } pub(crate) fn tuned_global_queue_interval(&self, config: &Config) -> u32 { // If an interval is explicitly set, don't tune. if let Some(configured) = config.global_queue_interval { return configured; } // As of Rust 1.45, casts from f64 -> u32 are saturating, which is fine here. let tasks_per_interval = (TARGET_GLOBAL_QUEUE_INTERVAL / self.task_poll_time_ewma) as u32; // If we are using self-tuning, we don't want to return less than 2 as that would result in the // global queue always getting checked first. tasks_per_interval.clamp(2, MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL) } pub(crate) fn submit(&mut self, to: &WorkerMetrics) { self.batch.submit(to, self.task_poll_time_ewma as u64); } pub(crate) fn about_to_park(&mut self) { self.batch.about_to_park(); } pub(crate) fn unparked(&mut self) { self.batch.unparked(); } pub(crate) fn inc_local_schedule_count(&mut self) { self.batch.inc_local_schedule_count(); } pub(crate) fn start_processing_scheduled_tasks(&mut self) { self.batch.start_processing_scheduled_tasks(); self.processing_scheduled_tasks_started_at = Instant::now(); self.tasks_polled_in_batch = 0; } pub(crate) fn end_processing_scheduled_tasks(&mut self) { self.batch.end_processing_scheduled_tasks(); // Update the EWMA task poll time if self.tasks_polled_in_batch > 0 { let now = Instant::now(); // If we "overflow" this conversion, we have bigger problems than // slightly off stats. let elapsed = (now - self.processing_scheduled_tasks_started_at).as_nanos() as f64; let num_polls = self.tasks_polled_in_batch as f64; // Calculate the mean poll duration for a single task in the batch let mean_poll_duration = elapsed / num_polls; // Compute the alpha weighted by the number of tasks polled this batch. let weighted_alpha = 1.0 - (1.0 - TASK_POLL_TIME_EWMA_ALPHA).powf(num_polls); // Now compute the new weighted average task poll time. self.task_poll_time_ewma = weighted_alpha * mean_poll_duration + (1.0 - weighted_alpha) * self.task_poll_time_ewma; } } pub(crate) fn start_poll(&mut self) { self.batch.start_poll(); self.tasks_polled_in_batch += 1; } pub(crate) fn end_poll(&mut self) { self.batch.end_poll(); } pub(crate) fn incr_steal_count(&mut self, by: u16) { self.batch.incr_steal_count(by); } pub(crate) fn incr_steal_operations(&mut self) { self.batch.incr_steal_operations(); } pub(crate) fn incr_overflow_count(&mut self) { self.batch.incr_overflow_count(); } } tokio-1.43.1/src/runtime/scheduler/multi_thread/trace.rs000064400000000000000000000034661046102023000214250ustar 00000000000000use crate::loom::sync::atomic::{AtomicBool, Ordering}; use crate::loom::sync::{Barrier, Mutex}; use crate::runtime::dump::Dump; use crate::runtime::scheduler::multi_thread::Handle; use crate::sync::notify::Notify; /// Tracing status of the worker. pub(super) struct TraceStatus { pub(super) trace_requested: AtomicBool, pub(super) trace_start: Barrier, pub(super) trace_end: Barrier, pub(super) result_ready: Notify, pub(super) trace_result: Mutex>, } impl TraceStatus { pub(super) fn new(remotes_len: usize) -> Self { Self { trace_requested: AtomicBool::new(false), trace_start: Barrier::new(remotes_len), trace_end: Barrier::new(remotes_len), result_ready: Notify::new(), trace_result: Mutex::new(None), } } pub(super) fn trace_requested(&self) -> bool { self.trace_requested.load(Ordering::Relaxed) } pub(super) async fn start_trace_request(&self, handle: &Handle) { while self .trace_requested .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) .is_err() { handle.notify_all(); crate::task::yield_now().await; } } pub(super) fn stash_result(&self, dump: Dump) { let _ = self.trace_result.lock().insert(dump); self.result_ready.notify_one(); } pub(super) fn take_result(&self) -> Option { self.trace_result.lock().take() } pub(super) async fn end_trace_request(&self, handle: &Handle) { while self .trace_requested .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed) .is_err() { handle.notify_all(); crate::task::yield_now().await; } } } tokio-1.43.1/src/runtime/scheduler/multi_thread/trace_mock.rs000064400000000000000000000002771046102023000224330ustar 00000000000000pub(super) struct TraceStatus {} impl TraceStatus { pub(super) fn new(_: usize) -> Self { Self {} } pub(super) fn trace_requested(&self) -> bool { false } } tokio-1.43.1/src/runtime/scheduler/multi_thread/worker/metrics.rs000064400000000000000000000004701046102023000232760ustar 00000000000000use super::Shared; impl Shared { pub(crate) fn injection_queue_depth(&self) -> usize { self.inject.len() } } cfg_unstable_metrics! { impl Shared { pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { self.remotes[worker].steal.len() } } } tokio-1.43.1/src/runtime/scheduler/multi_thread/worker/taskdump.rs000064400000000000000000000045211046102023000234610ustar 00000000000000use super::{Core, Handle, Shared}; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread::Stats; use crate::runtime::task::trace::trace_multi_thread; use crate::runtime::{dump, WorkerMetrics}; use std::time::Duration; impl Handle { pub(super) fn trace_core(&self, mut core: Box) -> Box { core.is_traced = false; if core.is_shutdown { return core; } // wait for other workers, or timeout without tracing let timeout = Duration::from_millis(250); // a _very_ generous timeout let barrier = if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) { barrier } else { // don't attempt to trace return core; }; if !barrier.is_leader() { // wait for leader to finish tracing self.shared.trace_status.trace_end.wait(); return core; } // trace let owned = &self.shared.owned; let mut local = self.shared.steal_all(); let synced = &self.shared.synced; let injection = &self.shared.inject; // safety: `trace_multi_thread` is invoked with the same `synced` that `injection` // was created with. let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) } .into_iter() .map(|(id, trace)| dump::Task::new(id, trace)) .collect(); let result = dump::Dump::new(traces); // stash the result self.shared.trace_status.stash_result(result); // allow other workers to proceed self.shared.trace_status.trace_end.wait(); core } } impl Shared { /// Steal all tasks from remotes into a single local queue. pub(super) fn steal_all(&self) -> super::queue::Local> { let (_steal, mut local) = super::queue::local(); let worker_metrics = WorkerMetrics::new(); let mut stats = Stats::new(&worker_metrics); for remote in self.remotes.iter() { let steal = &remote.steal; while !steal.is_empty() { if let Some(task) = steal.steal_into(&mut local, &mut stats) { local.push_back([task].into_iter()); } } } local } } tokio-1.43.1/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs000064400000000000000000000002031046102023000244630ustar 00000000000000use super::{Core, Handle}; impl Handle { pub(super) fn trace_core(&self, core: Box) -> Box { core } } tokio-1.43.1/src/runtime/scheduler/multi_thread/worker.rs000064400000000000000000001224411046102023000216330ustar 00000000000000//! A scheduler is initialized with a fixed number of workers. Each worker is //! driven by a thread. Each worker has a "core" which contains data such as the //! run queue and other state. When `block_in_place` is called, the worker's //! "core" is handed off to a new thread allowing the scheduler to continue to //! make progress while the originating thread blocks. //! //! # Shutdown //! //! Shutting down the runtime involves the following steps: //! //! 1. The Shared::close method is called. This closes the inject queue and //! `OwnedTasks` instance and wakes up all worker threads. //! //! 2. Each worker thread observes the close signal next time it runs //! Core::maintenance by checking whether the inject queue is closed. //! The `Core::is_shutdown` flag is set to true. //! //! 3. The worker thread calls `pre_shutdown` in parallel. Here, the worker //! will keep removing tasks from `OwnedTasks` until it is empty. No new //! tasks can be pushed to the `OwnedTasks` during or after this step as it //! was closed in step 1. //! //! 5. The workers call Shared::shutdown to enter the single-threaded phase of //! shutdown. These calls will push their core to `Shared::shutdown_cores`, //! and the last thread to push its core will finish the shutdown procedure. //! //! 6. The local run queue of each core is emptied, then the inject queue is //! emptied. //! //! At this point, shutdown has completed. It is not possible for any of the //! collections to contain any tasks at this point, as each collection was //! closed first, then emptied afterwards. //! //! ## Spawns during shutdown //! //! When spawning tasks during shutdown, there are two cases: //! //! * The spawner observes the `OwnedTasks` being open, and the inject queue is //! closed. //! * The spawner observes the `OwnedTasks` being closed and doesn't check the //! inject queue. //! //! The first case can only happen if the `OwnedTasks::bind` call happens before //! or during step 1 of shutdown. In this case, the runtime will clean up the //! task in step 3 of shutdown. //! //! In the latter case, the task was not spawned and the task is immediately //! cancelled by the spawner. //! //! The correctness of shutdown requires both the inject queue and `OwnedTasks` //! collection to have a closed bit. With a close bit on only the inject queue, //! spawning could run in to a situation where a task is successfully bound long //! after the runtime has shut down. With a close bit on only the `OwnedTasks`, //! the first spawning situation could result in the notification being pushed //! to the inject queue after step 6 of shutdown, which would leave a task in //! the inject queue indefinitely. This would be a ref-count cycle and a memory //! leak. use crate::loom::sync::{Arc, Mutex}; use crate::runtime; use crate::runtime::scheduler::multi_thread::{ idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker, }; use crate::runtime::scheduler::{inject, Defer, Lock}; use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks}; use crate::runtime::{ blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics, }; use crate::runtime::{context, TaskHooks}; use crate::util::atomic_cell::AtomicCell; use crate::util::rand::{FastRand, RngSeedGenerator}; use std::cell::RefCell; use std::task::Waker; use std::thread; use std::time::Duration; mod metrics; cfg_taskdump! { mod taskdump; } cfg_not_taskdump! { mod taskdump_mock; } /// A scheduler worker pub(super) struct Worker { /// Reference to scheduler's handle handle: Arc, /// Index holding this worker's remote state index: usize, /// Used to hand-off a worker's core to another thread. core: AtomicCell, } /// Core data struct Core { /// Used to schedule bookkeeping tasks every so often. tick: u32, /// When a task is scheduled from a worker, it is stored in this slot. The /// worker will check this slot for a task **before** checking the run /// queue. This effectively results in the **last** scheduled task to be run /// next (LIFO). This is an optimization for improving locality which /// benefits message passing patterns and helps to reduce latency. lifo_slot: Option, /// When `true`, locally scheduled tasks go to the LIFO slot. When `false`, /// they go to the back of the `run_queue`. lifo_enabled: bool, /// The worker-local run queue. run_queue: queue::Local>, /// True if the worker is currently searching for more work. Searching /// involves attempting to steal from other workers. is_searching: bool, /// True if the scheduler is being shutdown is_shutdown: bool, /// True if the scheduler is being traced is_traced: bool, /// Parker /// /// Stored in an `Option` as the parker is added / removed to make the /// borrow checker happy. park: Option, /// Per-worker runtime stats stats: Stats, /// How often to check the global queue global_queue_interval: u32, /// Fast random number generator. rand: FastRand, } /// State shared across all workers pub(crate) struct Shared { /// Per-worker remote state. All other workers have access to this and is /// how they communicate between each other. remotes: Box<[Remote]>, /// Global task queue used for: /// 1. Submit work to the scheduler while **not** currently on a worker thread. /// 2. Submit work to the scheduler when a worker run queue is saturated pub(super) inject: inject::Shared>, /// Coordinates idle workers idle: Idle, /// Collection of all active tasks spawned onto this executor. pub(crate) owned: OwnedTasks>, /// Data synchronized by the scheduler mutex pub(super) synced: Mutex, /// Cores that have observed the shutdown signal /// /// The core is **not** placed back in the worker to avoid it from being /// stolen by a thread that was spawned as part of `block_in_place`. #[allow(clippy::vec_box)] // we're moving an already-boxed value shutdown_cores: Mutex>>, /// The number of cores that have observed the trace signal. pub(super) trace_status: TraceStatus, /// Scheduler configuration options config: Config, /// Collects metrics from the runtime. pub(super) scheduler_metrics: SchedulerMetrics, pub(super) worker_metrics: Box<[WorkerMetrics]>, /// Only held to trigger some code on drop. This is used to get internal /// runtime metrics that can be useful when doing performance /// investigations. This does nothing (empty struct, no drop impl) unless /// the `tokio_internal_mt_counters` `cfg` flag is set. _counters: Counters, } /// Data synchronized by the scheduler mutex pub(crate) struct Synced { /// Synchronized state for `Idle`. pub(super) idle: idle::Synced, /// Synchronized state for `Inject`. pub(crate) inject: inject::Synced, } /// Used to communicate with a worker from other threads. struct Remote { /// Steals tasks from this worker. pub(super) steal: queue::Steal>, /// Unparks the associated worker thread unpark: Unparker, } /// Thread-local context pub(crate) struct Context { /// Worker worker: Arc, /// Core data core: RefCell>>, /// Tasks to wake after resource drivers are polled. This is mostly to /// handle yielded tasks. pub(crate) defer: Defer, } /// Starts the workers pub(crate) struct Launch(Vec>); /// Running a task may consume the core. If the core is still available when /// running the task completes, it is returned. Otherwise, the worker will need /// to stop processing. type RunResult = Result, ()>; /// A task handle type Task = task::Task>; /// A notified task handle type Notified = task::Notified>; /// Value picked out of thin-air. Running the LIFO slot a handful of times /// seems sufficient to benefit from locality. More than 3 times probably is /// overweighing. The value can be tuned in the future with data that shows /// improvements. const MAX_LIFO_POLLS_PER_TICK: usize = 3; pub(super) fn create( size: usize, park: Parker, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, ) -> (Arc, Launch) { let mut cores = Vec::with_capacity(size); let mut remotes = Vec::with_capacity(size); let mut worker_metrics = Vec::with_capacity(size); // Create the local queues for _ in 0..size { let (steal, run_queue) = queue::local(); let park = park.clone(); let unpark = park.unpark(); let metrics = WorkerMetrics::from_config(&config); let stats = Stats::new(&metrics); cores.push(Box::new(Core { tick: 0, lifo_slot: None, lifo_enabled: !config.disable_lifo_slot, run_queue, is_searching: false, is_shutdown: false, is_traced: false, park: Some(park), global_queue_interval: stats.tuned_global_queue_interval(&config), stats, rand: FastRand::from_seed(config.seed_generator.next_seed()), })); remotes.push(Remote { steal, unpark }); worker_metrics.push(metrics); } let (idle, idle_synced) = Idle::new(size); let (inject, inject_synced) = inject::Shared::new(); let remotes_len = remotes.len(); let handle = Arc::new(Handle { task_hooks: TaskHooks { task_spawn_callback: config.before_spawn.clone(), task_terminate_callback: config.after_termination.clone(), }, shared: Shared { remotes: remotes.into_boxed_slice(), inject, idle, owned: OwnedTasks::new(size), synced: Mutex::new(Synced { idle: idle_synced, inject: inject_synced, }), shutdown_cores: Mutex::new(vec![]), trace_status: TraceStatus::new(remotes_len), config, scheduler_metrics: SchedulerMetrics::new(), worker_metrics: worker_metrics.into_boxed_slice(), _counters: Counters, }, driver: driver_handle, blocking_spawner, seed_generator, }); let mut launch = Launch(vec![]); for (index, core) in cores.drain(..).enumerate() { launch.0.push(Arc::new(Worker { handle: handle.clone(), index, core: AtomicCell::new(Some(core)), })); } (handle, launch) } #[track_caller] pub(crate) fn block_in_place(f: F) -> R where F: FnOnce() -> R, { // Try to steal the worker core back struct Reset { take_core: bool, budget: coop::Budget, } impl Drop for Reset { fn drop(&mut self) { with_current(|maybe_cx| { if let Some(cx) = maybe_cx { if self.take_core { let core = cx.worker.core.take(); if core.is_some() { cx.worker.handle.shared.worker_metrics[cx.worker.index] .set_thread_id(thread::current().id()); } let mut cx_core = cx.core.borrow_mut(); assert!(cx_core.is_none()); *cx_core = core; } // Reset the task budget as we are re-entering the // runtime. coop::set(self.budget); } }); } } let mut had_entered = false; let mut take_core = false; let setup_result = with_current(|maybe_cx| { match ( crate::runtime::context::current_enter_context(), maybe_cx.is_some(), ) { (context::EnterRuntime::Entered { .. }, true) => { // We are on a thread pool runtime thread, so we just need to // set up blocking. had_entered = true; } ( context::EnterRuntime::Entered { allow_block_in_place, }, false, ) => { // We are on an executor, but _not_ on the thread pool. That is // _only_ okay if we are in a thread pool runtime's block_on // method: if allow_block_in_place { had_entered = true; return Ok(()); } else { // This probably means we are on the current_thread runtime or in a // LocalSet, where it is _not_ okay to block. return Err( "can call blocking only when running on the multi-threaded runtime", ); } } (context::EnterRuntime::NotEntered, true) => { // This is a nested call to block_in_place (we already exited). // All the necessary setup has already been done. return Ok(()); } (context::EnterRuntime::NotEntered, false) => { // We are outside of the tokio runtime, so blocking is fine. // We can also skip all of the thread pool blocking setup steps. return Ok(()); } } let cx = maybe_cx.expect("no .is_some() == false cases above should lead here"); // Get the worker core. If none is set, then blocking is fine! let mut core = match cx.core.borrow_mut().take() { Some(core) => core, None => return Ok(()), }; // If we heavily call `spawn_blocking`, there might be no available thread to // run this core. Except for the task in the lifo_slot, all tasks can be // stolen, so we move the task out of the lifo_slot to the run_queue. if let Some(task) = core.lifo_slot.take() { core.run_queue .push_back_or_overflow(task, &*cx.worker.handle, &mut core.stats); } // We are taking the core from the context and sending it to another // thread. take_core = true; // The parker should be set here assert!(core.park.is_some()); // In order to block, the core must be sent to another thread for // execution. // // First, move the core back into the worker's shared core slot. cx.worker.core.set(core); // Next, clone the worker handle and send it to a new thread for // processing. // // Once the blocking task is done executing, we will attempt to // steal the core back. let worker = cx.worker.clone(); runtime::spawn_blocking(move || run(worker)); Ok(()) }); if let Err(panic_message) = setup_result { panic!("{}", panic_message); } if had_entered { // Unset the current task's budget. Blocking sections are not // constrained by task budgets. let _reset = Reset { take_core, budget: coop::stop(), }; crate::runtime::context::exit_runtime(f) } else { f() } } impl Launch { pub(crate) fn launch(mut self) { for worker in self.0.drain(..) { runtime::spawn_blocking(move || run(worker)); } } } fn run(worker: Arc) { #[allow(dead_code)] struct AbortOnPanic; impl Drop for AbortOnPanic { fn drop(&mut self) { if std::thread::panicking() { eprintln!("worker thread panicking; aborting process"); std::process::abort(); } } } // Catching panics on worker threads in tests is quite tricky. Instead, when // debug assertions are enabled, we just abort the process. #[cfg(debug_assertions)] let _abort_on_panic = AbortOnPanic; // Acquire a core. If this fails, then another thread is running this // worker and there is nothing further to do. let core = match worker.core.take() { Some(core) => core, None => return, }; worker.handle.shared.worker_metrics[worker.index].set_thread_id(thread::current().id()); let handle = scheduler::Handle::MultiThread(worker.handle.clone()); crate::runtime::context::enter_runtime(&handle, true, |_| { // Set the worker context. let cx = scheduler::Context::MultiThread(Context { worker, core: RefCell::new(None), defer: Defer::new(), }); context::set_scheduler(&cx, || { let cx = cx.expect_multi_thread(); // This should always be an error. It only returns a `Result` to support // using `?` to short circuit. assert!(cx.run(core).is_err()); // Check if there are any deferred tasks to notify. This can happen when // the worker core is lost due to `block_in_place()` being called from // within the task. cx.defer.wake(); }); }); } impl Context { fn run(&self, mut core: Box) -> RunResult { // Reset `lifo_enabled` here in case the core was previously stolen from // a task that had the LIFO slot disabled. self.reset_lifo_enabled(&mut core); // Start as "processing" tasks as polling tasks from the local queue // will be one of the first things we do. core.stats.start_processing_scheduled_tasks(); while !core.is_shutdown { self.assert_lifo_enabled_is_correct(&core); if core.is_traced { core = self.worker.handle.trace_core(core); } // Increment the tick core.tick(); // Run maintenance, if needed core = self.maintenance(core); // First, check work available to the current worker. if let Some(task) = core.next_task(&self.worker) { core = self.run_task(task, core)?; continue; } // We consumed all work in the queues and will start searching for work. core.stats.end_processing_scheduled_tasks(); // There is no more **local** work to process, try to steal work // from other workers. if let Some(task) = core.steal_work(&self.worker) { // Found work, switch back to processing core.stats.start_processing_scheduled_tasks(); core = self.run_task(task, core)?; } else { // Wait for work core = if !self.defer.is_empty() { self.park_timeout(core, Some(Duration::from_millis(0))) } else { self.park(core) }; core.stats.start_processing_scheduled_tasks(); } } core.pre_shutdown(&self.worker); // Signal shutdown self.worker.handle.shutdown_core(core); Err(()) } fn run_task(&self, task: Notified, mut core: Box) -> RunResult { let task = self.worker.handle.shared.owned.assert_owner(task); // Make sure the worker is not in the **searching** state. This enables // another idle worker to try to steal work. core.transition_from_searching(&self.worker); self.assert_lifo_enabled_is_correct(&core); // Measure the poll start time. Note that we may end up polling other // tasks under this measurement. In this case, the tasks came from the // LIFO slot and are considered part of the current task for scheduling // purposes. These tasks inherent the "parent"'s limits. core.stats.start_poll(); // Make the core available to the runtime context *self.core.borrow_mut() = Some(core); // Run the task coop::budget(|| { task.run(); let mut lifo_polls = 0; // As long as there is budget remaining and a task exists in the // `lifo_slot`, then keep running. loop { // Check if we still have the core. If not, the core was stolen // by another worker. let mut core = match self.core.borrow_mut().take() { Some(core) => core, None => { // In this case, we cannot call `reset_lifo_enabled()` // because the core was stolen. The stealer will handle // that at the top of `Context::run` return Err(()); } }; // Check for a task in the LIFO slot let task = match core.lifo_slot.take() { Some(task) => task, None => { self.reset_lifo_enabled(&mut core); core.stats.end_poll(); return Ok(core); } }; if !coop::has_budget_remaining() { core.stats.end_poll(); // Not enough budget left to run the LIFO task, push it to // the back of the queue and return. core.run_queue.push_back_or_overflow( task, &*self.worker.handle, &mut core.stats, ); // If we hit this point, the LIFO slot should be enabled. // There is no need to reset it. debug_assert!(core.lifo_enabled); return Ok(core); } // Track that we are about to run a task from the LIFO slot. lifo_polls += 1; super::counters::inc_lifo_schedules(); // Disable the LIFO slot if we reach our limit // // In ping-ping style workloads where task A notifies task B, // which notifies task A again, continuously prioritizing the // LIFO slot can cause starvation as these two tasks will // repeatedly schedule the other. To mitigate this, we limit the // number of times the LIFO slot is prioritized. if lifo_polls >= MAX_LIFO_POLLS_PER_TICK { core.lifo_enabled = false; super::counters::inc_lifo_capped(); } // Run the LIFO task, then loop *self.core.borrow_mut() = Some(core); let task = self.worker.handle.shared.owned.assert_owner(task); task.run(); } }) } fn reset_lifo_enabled(&self, core: &mut Core) { core.lifo_enabled = !self.worker.handle.shared.config.disable_lifo_slot; } fn assert_lifo_enabled_is_correct(&self, core: &Core) { debug_assert_eq!( core.lifo_enabled, !self.worker.handle.shared.config.disable_lifo_slot ); } fn maintenance(&self, mut core: Box) -> Box { if core.tick % self.worker.handle.shared.config.event_interval == 0 { super::counters::inc_num_maintenance(); core.stats.end_processing_scheduled_tasks(); // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... // to run without actually putting the thread to sleep. core = self.park_timeout(core, Some(Duration::from_millis(0))); // Run regularly scheduled maintenance core.maintenance(&self.worker); core.stats.start_processing_scheduled_tasks(); } core } /// Parks the worker thread while waiting for tasks to execute. /// /// This function checks if indeed there's no more work left to be done before parking. /// Also important to notice that, before parking, the worker thread will try to take /// ownership of the Driver (IO/Time) and dispatch any events that might have fired. /// Whenever a worker thread executes the Driver loop, all waken tasks are scheduled /// in its own local queue until the queue saturates (ntasks > `LOCAL_QUEUE_CAPACITY`). /// When the local queue is saturated, the overflow tasks are added to the injection queue /// from where other workers can pick them up. /// Also, we rely on the workstealing algorithm to spread the tasks amongst workers /// after all the IOs get dispatched fn park(&self, mut core: Box) -> Box { if let Some(f) = &self.worker.handle.shared.config.before_park { f(); } if core.transition_to_parked(&self.worker) { while !core.is_shutdown && !core.is_traced { core.stats.about_to_park(); core.stats .submit(&self.worker.handle.shared.worker_metrics[self.worker.index]); core = self.park_timeout(core, None); core.stats.unparked(); // Run regularly scheduled maintenance core.maintenance(&self.worker); if core.transition_from_parked(&self.worker) { break; } } } if let Some(f) = &self.worker.handle.shared.config.after_unpark { f(); } core } fn park_timeout(&self, mut core: Box, duration: Option) -> Box { self.assert_lifo_enabled_is_correct(&core); // Take the parker out of core let mut park = core.park.take().expect("park missing"); // Store `core` in context *self.core.borrow_mut() = Some(core); // Park thread if let Some(timeout) = duration { park.park_timeout(&self.worker.handle.driver, timeout); } else { park.park(&self.worker.handle.driver); } self.defer.wake(); // Remove `core` from context core = self.core.borrow_mut().take().expect("core missing"); // Place `park` back in `core` core.park = Some(park); if core.should_notify_others() { self.worker.handle.notify_parked_local(); } core } pub(crate) fn defer(&self, waker: &Waker) { self.defer.defer(waker); } #[allow(dead_code)] pub(crate) fn get_worker_index(&self) -> usize { self.worker.index } } impl Core { /// Increment the tick fn tick(&mut self) { self.tick = self.tick.wrapping_add(1); } /// Return the next notified task available to this worker. fn next_task(&mut self, worker: &Worker) -> Option { if self.tick % self.global_queue_interval == 0 { // Update the global queue interval, if needed self.tune_global_queue_interval(worker); worker .handle .next_remote_task() .or_else(|| self.next_local_task()) } else { let maybe_task = self.next_local_task(); if maybe_task.is_some() { return maybe_task; } if worker.inject().is_empty() { return None; } // Other threads can only **remove** tasks from the current worker's // `run_queue`. So, we can be confident that by the time we call // `run_queue.push_back` below, there will be *at least* `cap` // available slots in the queue. let cap = usize::min( self.run_queue.remaining_slots(), self.run_queue.max_capacity() / 2, ); // The worker is currently idle, pull a batch of work from the // injection queue. We don't want to pull *all* the work so other // workers can also get some. let n = usize::min( worker.inject().len() / worker.handle.shared.remotes.len() + 1, cap, ); // Take at least one task since the first task is returned directly // and not pushed onto the local queue. let n = usize::max(1, n); let mut synced = worker.handle.shared.synced.lock(); // safety: passing in the correct `inject::Synced`. let mut tasks = unsafe { worker.inject().pop_n(&mut synced.inject, n) }; // Pop the first task to return immediately let ret = tasks.next(); // Push the rest of the on the run queue self.run_queue.push_back(tasks); ret } } fn next_local_task(&mut self) -> Option { self.lifo_slot.take().or_else(|| self.run_queue.pop()) } /// Function responsible for stealing tasks from another worker /// /// Note: Only if less than half the workers are searching for tasks to steal /// a new worker will actually try to steal. The idea is to make sure not all /// workers will be trying to steal at the same time. fn steal_work(&mut self, worker: &Worker) -> Option { if !self.transition_to_searching(worker) { return None; } let num = worker.handle.shared.remotes.len(); // Start from a random worker let start = self.rand.fastrand_n(num as u32) as usize; for i in 0..num { let i = (start + i) % num; // Don't steal from ourself! We know we don't have work. if i == worker.index { continue; } let target = &worker.handle.shared.remotes[i]; if let Some(task) = target .steal .steal_into(&mut self.run_queue, &mut self.stats) { return Some(task); } } // Fallback on checking the global queue worker.handle.next_remote_task() } fn transition_to_searching(&mut self, worker: &Worker) -> bool { if !self.is_searching { self.is_searching = worker.handle.shared.idle.transition_worker_to_searching(); } self.is_searching } fn transition_from_searching(&mut self, worker: &Worker) { if !self.is_searching { return; } self.is_searching = false; worker.handle.transition_worker_from_searching(); } fn has_tasks(&self) -> bool { self.lifo_slot.is_some() || self.run_queue.has_tasks() } fn should_notify_others(&self) -> bool { // If there are tasks available to steal, but this worker is not // looking for tasks to steal, notify another worker. if self.is_searching { return false; } self.lifo_slot.is_some() as usize + self.run_queue.len() > 1 } /// Prepares the worker state for parking. /// /// Returns true if the transition happened, false if there is work to do first. fn transition_to_parked(&mut self, worker: &Worker) -> bool { // Workers should not park if they have work to do if self.has_tasks() || self.is_traced { return false; } // When the final worker transitions **out** of searching to parked, it // must check all the queues one last time in case work materialized // between the last work scan and transitioning out of searching. let is_last_searcher = worker.handle.shared.idle.transition_worker_to_parked( &worker.handle.shared, worker.index, self.is_searching, ); // The worker is no longer searching. Setting this is the local cache // only. self.is_searching = false; if is_last_searcher { worker.handle.notify_if_work_pending(); } true } /// Returns `true` if the transition happened. fn transition_from_parked(&mut self, worker: &Worker) -> bool { // If a task is in the lifo slot/run queue, then we must unpark regardless of // being notified if self.has_tasks() { // When a worker wakes, it should only transition to the "searching" // state when the wake originates from another worker *or* a new task // is pushed. We do *not* want the worker to transition to "searching" // when it wakes when the I/O driver receives new events. self.is_searching = !worker .handle .shared .idle .unpark_worker_by_id(&worker.handle.shared, worker.index); return true; } if worker .handle .shared .idle .is_parked(&worker.handle.shared, worker.index) { return false; } // When unparked, the worker is in the searching state. self.is_searching = true; true } /// Runs maintenance work such as checking the pool's state. fn maintenance(&mut self, worker: &Worker) { self.stats .submit(&worker.handle.shared.worker_metrics[worker.index]); if !self.is_shutdown { // Check if the scheduler has been shutdown let synced = worker.handle.shared.synced.lock(); self.is_shutdown = worker.inject().is_closed(&synced.inject); } if !self.is_traced { // Check if the worker should be tracing. self.is_traced = worker.handle.shared.trace_status.trace_requested(); } } /// Signals all tasks to shut down, and waits for them to complete. Must run /// before we enter the single-threaded phase of shutdown processing. fn pre_shutdown(&mut self, worker: &Worker) { // Start from a random inner list let start = self .rand .fastrand_n(worker.handle.shared.owned.get_shard_size() as u32); // Signal to all tasks to shut down. worker .handle .shared .owned .close_and_shutdown_all(start as usize); self.stats .submit(&worker.handle.shared.worker_metrics[worker.index]); } /// Shuts down the core. fn shutdown(&mut self, handle: &Handle) { // Take the core let mut park = self.park.take().expect("park missing"); // Drain the queue while self.next_local_task().is_some() {} park.shutdown(&handle.driver); } fn tune_global_queue_interval(&mut self, worker: &Worker) { let next = self .stats .tuned_global_queue_interval(&worker.handle.shared.config); // Smooth out jitter if u32::abs_diff(self.global_queue_interval, next) > 2 { self.global_queue_interval = next; } } } impl Worker { /// Returns a reference to the scheduler's injection queue. fn inject(&self) -> &inject::Shared> { &self.handle.shared.inject } } // TODO: Move `Handle` impls into handle.rs impl task::Schedule for Arc { fn release(&self, task: &Task) -> Option { self.shared.owned.remove(task) } fn schedule(&self, task: Notified) { self.schedule_task(task, false); } fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: self.task_hooks.task_terminate_callback.clone(), } } fn yield_now(&self, task: Notified) { self.schedule_task(task, true); } } impl Handle { pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { with_current(|maybe_cx| { if let Some(cx) = maybe_cx { // Make sure the task is part of the **current** scheduler. if self.ptr_eq(&cx.worker.handle) { // And the current thread still holds a core if let Some(core) = cx.core.borrow_mut().as_mut() { self.schedule_local(core, task, is_yield); return; } } } // Otherwise, use the inject queue. self.push_remote_task(task); self.notify_parked_remote(); }); } pub(super) fn schedule_option_task_without_yield(&self, task: Option) { if let Some(task) = task { self.schedule_task(task, false); } } fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { core.stats.inc_local_schedule_count(); // Spawning from the worker thread. If scheduling a "yield" then the // task must always be pushed to the back of the queue, enabling other // tasks to be executed. If **not** a yield, then there is more // flexibility and the task may go to the front of the queue. let should_notify = if is_yield || !core.lifo_enabled { core.run_queue .push_back_or_overflow(task, self, &mut core.stats); true } else { // Push to the LIFO slot let prev = core.lifo_slot.take(); let ret = prev.is_some(); if let Some(prev) = prev { core.run_queue .push_back_or_overflow(prev, self, &mut core.stats); } core.lifo_slot = Some(task); ret }; // Only notify if not currently parked. If `park` is `None`, then the // scheduling is from a resource driver. As notifications often come in // batches, the notification is delayed until the park is complete. if should_notify && core.park.is_some() { self.notify_parked_local(); } } fn next_remote_task(&self) -> Option { if self.shared.inject.is_empty() { return None; } let mut synced = self.shared.synced.lock(); // safety: passing in correct `idle::Synced` unsafe { self.shared.inject.pop(&mut synced.inject) } } fn push_remote_task(&self, task: Notified) { self.shared.scheduler_metrics.inc_remote_schedule_count(); let mut synced = self.shared.synced.lock(); // safety: passing in correct `idle::Synced` unsafe { self.shared.inject.push(&mut synced.inject, task); } } pub(super) fn close(&self) { if self .shared .inject .close(&mut self.shared.synced.lock().inject) { self.notify_all(); } } fn notify_parked_local(&self) { super::counters::inc_num_inc_notify_local(); if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) { super::counters::inc_num_unparks_local(); self.shared.remotes[index].unpark.unpark(&self.driver); } } fn notify_parked_remote(&self) { if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) { self.shared.remotes[index].unpark.unpark(&self.driver); } } pub(super) fn notify_all(&self) { for remote in &self.shared.remotes[..] { remote.unpark.unpark(&self.driver); } } fn notify_if_work_pending(&self) { for remote in &self.shared.remotes[..] { if !remote.steal.is_empty() { self.notify_parked_local(); return; } } if !self.shared.inject.is_empty() { self.notify_parked_local(); } } fn transition_worker_from_searching(&self) { if self.shared.idle.transition_worker_from_searching() { // We are the final searching worker. Because work was found, we // need to notify another worker. self.notify_parked_local(); } } /// Signals that a worker has observed the shutdown signal and has replaced /// its core back into its handle. /// /// If all workers have reached this point, the final cleanup is performed. fn shutdown_core(&self, core: Box) { let mut cores = self.shared.shutdown_cores.lock(); cores.push(core); if cores.len() != self.shared.remotes.len() { return; } debug_assert!(self.shared.owned.is_empty()); for mut core in cores.drain(..) { core.shutdown(self); } // Drain the injection queue // // We already shut down every task, so we can simply drop the tasks. while let Some(task) = self.next_remote_task() { drop(task); } } fn ptr_eq(&self, other: &Handle) -> bool { std::ptr::eq(self, other) } } impl Overflow> for Handle { fn push(&self, task: task::Notified>) { self.push_remote_task(task); } fn push_batch(&self, iter: I) where I: Iterator>>, { unsafe { self.shared.inject.push_batch(self, iter); } } } pub(crate) struct InjectGuard<'a> { lock: crate::loom::sync::MutexGuard<'a, Synced>, } impl<'a> AsMut for InjectGuard<'a> { fn as_mut(&mut self) -> &mut inject::Synced { &mut self.lock.inject } } impl<'a> Lock for &'a Handle { type Handle = InjectGuard<'a>; fn lock(self) -> Self::Handle { InjectGuard { lock: self.shared.synced.lock(), } } } #[track_caller] fn with_current(f: impl FnOnce(Option<&Context>) -> R) -> R { use scheduler::Context::MultiThread; context::with_scheduler(|ctx| match ctx { Some(MultiThread(ctx)) => f(Some(ctx)), _ => f(None), }) } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/counters.rs000064400000000000000000000142131046102023000230210ustar 00000000000000#[cfg(tokio_internal_mt_counters)] mod imp { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::Relaxed; static NUM_MAINTENANCE: AtomicUsize = AtomicUsize::new(0); static NUM_NOTIFY_LOCAL: AtomicUsize = AtomicUsize::new(0); static NUM_NOTIFY_REMOTE: AtomicUsize = AtomicUsize::new(0); static NUM_UNPARKS_LOCAL: AtomicUsize = AtomicUsize::new(0); static NUM_UNPARKS_REMOTE: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_SCHEDULES: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_CAPPED: AtomicUsize = AtomicUsize::new(0); static NUM_STEALS: AtomicUsize = AtomicUsize::new(0); static NUM_OVERFLOW: AtomicUsize = AtomicUsize::new(0); static NUM_PARK: AtomicUsize = AtomicUsize::new(0); static NUM_POLLS: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_POLLS: AtomicUsize = AtomicUsize::new(0); static NUM_REMOTE_BATCH: AtomicUsize = AtomicUsize::new(0); static NUM_GLOBAL_QUEUE_INTERVAL: AtomicUsize = AtomicUsize::new(0); static NUM_NO_AVAIL_CORE: AtomicUsize = AtomicUsize::new(0); static NUM_RELAY_SEARCH: AtomicUsize = AtomicUsize::new(0); static NUM_SPIN_STALL: AtomicUsize = AtomicUsize::new(0); static NUM_NO_LOCAL_WORK: AtomicUsize = AtomicUsize::new(0); impl Drop for super::Counters { fn drop(&mut self) { let notifies_local = NUM_NOTIFY_LOCAL.load(Relaxed); let notifies_remote = NUM_NOTIFY_REMOTE.load(Relaxed); let unparks_local = NUM_UNPARKS_LOCAL.load(Relaxed); let unparks_remote = NUM_UNPARKS_REMOTE.load(Relaxed); let maintenance = NUM_MAINTENANCE.load(Relaxed); let lifo_scheds = NUM_LIFO_SCHEDULES.load(Relaxed); let lifo_capped = NUM_LIFO_CAPPED.load(Relaxed); let num_steals = NUM_STEALS.load(Relaxed); let num_overflow = NUM_OVERFLOW.load(Relaxed); let num_park = NUM_PARK.load(Relaxed); let num_polls = NUM_POLLS.load(Relaxed); let num_lifo_polls = NUM_LIFO_POLLS.load(Relaxed); let num_remote_batch = NUM_REMOTE_BATCH.load(Relaxed); let num_global_queue_interval = NUM_GLOBAL_QUEUE_INTERVAL.load(Relaxed); let num_no_avail_core = NUM_NO_AVAIL_CORE.load(Relaxed); let num_relay_search = NUM_RELAY_SEARCH.load(Relaxed); let num_spin_stall = NUM_SPIN_STALL.load(Relaxed); let num_no_local_work = NUM_NO_LOCAL_WORK.load(Relaxed); println!("---"); println!("notifies (remote): {}", notifies_remote); println!(" notifies (local): {}", notifies_local); println!(" unparks (local): {}", unparks_local); println!(" unparks (remote): {}", unparks_remote); println!(" notify, no core: {}", num_no_avail_core); println!(" maintenance: {}", maintenance); println!(" LIFO schedules: {}", lifo_scheds); println!(" LIFO capped: {}", lifo_capped); println!(" steals: {}", num_steals); println!(" queue overflows: {}", num_overflow); println!(" parks: {}", num_park); println!(" polls: {}", num_polls); println!(" polls (LIFO): {}", num_lifo_polls); println!("remote task batch: {}", num_remote_batch); println!("global Q interval: {}", num_global_queue_interval); println!(" relay search: {}", num_relay_search); println!(" spin stall: {}", num_spin_stall); println!(" no local work: {}", num_no_local_work); } } pub(crate) fn inc_num_inc_notify_local() { NUM_NOTIFY_LOCAL.fetch_add(1, Relaxed); } pub(crate) fn inc_num_notify_remote() { NUM_NOTIFY_REMOTE.fetch_add(1, Relaxed); } pub(crate) fn inc_num_unparks_local() { NUM_UNPARKS_LOCAL.fetch_add(1, Relaxed); } pub(crate) fn inc_num_unparks_remote() { NUM_UNPARKS_REMOTE.fetch_add(1, Relaxed); } pub(crate) fn inc_num_maintenance() { NUM_MAINTENANCE.fetch_add(1, Relaxed); } pub(crate) fn inc_lifo_schedules() { NUM_LIFO_SCHEDULES.fetch_add(1, Relaxed); } pub(crate) fn inc_lifo_capped() { NUM_LIFO_CAPPED.fetch_add(1, Relaxed); } pub(crate) fn inc_num_steals() { NUM_STEALS.fetch_add(1, Relaxed); } pub(crate) fn inc_num_overflows() { NUM_OVERFLOW.fetch_add(1, Relaxed); } pub(crate) fn inc_num_parks() { NUM_PARK.fetch_add(1, Relaxed); } pub(crate) fn inc_num_polls() { NUM_POLLS.fetch_add(1, Relaxed); } pub(crate) fn inc_num_lifo_polls() { NUM_LIFO_POLLS.fetch_add(1, Relaxed); } pub(crate) fn inc_num_remote_batch() { NUM_REMOTE_BATCH.fetch_add(1, Relaxed); } pub(crate) fn inc_global_queue_interval() { NUM_GLOBAL_QUEUE_INTERVAL.fetch_add(1, Relaxed); } pub(crate) fn inc_notify_no_core() { NUM_NO_AVAIL_CORE.fetch_add(1, Relaxed); } pub(crate) fn inc_num_relay_search() { NUM_RELAY_SEARCH.fetch_add(1, Relaxed); } pub(crate) fn inc_num_spin_stall() { NUM_SPIN_STALL.fetch_add(1, Relaxed); } pub(crate) fn inc_num_no_local_work() { NUM_NO_LOCAL_WORK.fetch_add(1, Relaxed); } } #[cfg(not(tokio_internal_mt_counters))] mod imp { pub(crate) fn inc_num_inc_notify_local() {} pub(crate) fn inc_num_notify_remote() {} pub(crate) fn inc_num_unparks_local() {} pub(crate) fn inc_num_unparks_remote() {} pub(crate) fn inc_num_maintenance() {} pub(crate) fn inc_lifo_schedules() {} pub(crate) fn inc_lifo_capped() {} pub(crate) fn inc_num_steals() {} pub(crate) fn inc_num_overflows() {} pub(crate) fn inc_num_parks() {} pub(crate) fn inc_num_polls() {} pub(crate) fn inc_num_lifo_polls() {} pub(crate) fn inc_num_remote_batch() {} pub(crate) fn inc_global_queue_interval() {} pub(crate) fn inc_notify_no_core() {} pub(crate) fn inc_num_relay_search() {} pub(crate) fn inc_num_spin_stall() {} pub(crate) fn inc_num_no_local_work() {} } #[derive(Debug)] pub(crate) struct Counters; pub(super) use imp::*; tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/handle/metrics.rs000064400000000000000000000025461046102023000240660ustar 00000000000000use super::Handle; use crate::runtime::{SchedulerMetrics, WorkerMetrics}; impl Handle { pub(crate) fn num_workers(&self) -> usize { self.shared.worker_metrics.len() } pub(crate) fn num_blocking_threads(&self) -> usize { // workers are currently spawned using spawn_blocking self.blocking_spawner .num_threads() .saturating_sub(self.num_workers()) } pub(crate) fn num_idle_blocking_threads(&self) -> usize { self.blocking_spawner.num_idle_threads() } pub(crate) fn num_alive_tasks(&self) -> usize { self.shared.owned.num_alive_tasks() } cfg_64bit_metrics! { pub(crate) fn spawned_tasks_count(&self) -> u64 { self.shared.owned.spawned_tasks_count() } } pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { &self.shared.scheduler_metrics } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { &self.shared.worker_metrics[worker] } pub(crate) fn injection_queue_depth(&self) -> usize { self.shared.injection_queue_depth() } pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { self.shared.worker_local_queue_depth(worker) } pub(crate) fn blocking_queue_depth(&self) -> usize { self.blocking_spawner.queue_depth() } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/handle/taskdump.rs000064400000000000000000000012121046102023000242350ustar 00000000000000use super::Handle; use crate::runtime::Dump; impl Handle { pub(crate) async fn dump(&self) -> Dump { let trace_status = &self.shared.trace_status; // If a dump is in progress, block. trace_status.start_trace_request(&self).await; let result = loop { if let Some(result) = trace_status.take_result() { break result; } else { self.notify_all(); trace_status.result_ready.notified().await; } }; // Allow other queued dumps to proceed. trace_status.end_trace_request(&self).await; result } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/handle.rs000064400000000000000000000040321046102023000224100ustar 00000000000000use crate::future::Future; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread_alt::worker; use crate::runtime::{ blocking, driver, task::{self, JoinHandle}, TaskHooks, TaskMeta, }; use crate::util::RngSeedGenerator; use std::fmt; cfg_unstable_metrics! { mod metrics; } /// Handle to the multi thread scheduler pub(crate) struct Handle { /// Task spawner pub(super) shared: worker::Shared, /// Resource driver handles pub(crate) driver: driver::Handle, /// Blocking pool spawner pub(crate) blocking_spawner: blocking::Spawner, /// Current random number generator seed pub(crate) seed_generator: RngSeedGenerator, /// User-supplied hooks to invoke for things pub(crate) task_hooks: TaskHooks, } impl Handle { /// Spawns a future onto the thread pool pub(crate) fn spawn(me: &Arc, future: F, id: task::Id) -> JoinHandle where F: crate::future::Future + Send + 'static, F::Output: Send + 'static, { Self::bind_new_task(me, future, id) } pub(crate) fn shutdown(&self) { self.shared.close(self); self.driver.unpark(); } pub(super) fn bind_new_task(me: &Arc, future: T, id: task::Id) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); me.task_hooks.spawn(&TaskMeta { #[cfg(tokio_unstable)] id, _phantom: Default::default(), }); if let Some(notified) = notified { me.shared.schedule_task(notified, false); } handle } } cfg_unstable! { use std::num::NonZeroU64; impl Handle { pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } } } impl fmt::Debug for Handle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("multi_thread::Handle { ... }").finish() } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/idle.rs000064400000000000000000000316361046102023000221040ustar 00000000000000//! Coordinates idling workers #![allow(dead_code)] use crate::loom::sync::atomic::{AtomicBool, AtomicUsize}; use crate::loom::sync::MutexGuard; use crate::runtime::scheduler::multi_thread_alt::{worker, Core, Handle, Shared}; use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; pub(super) struct Idle { /// Number of searching cores num_searching: AtomicUsize, /// Number of idle cores num_idle: AtomicUsize, /// Map of idle cores idle_map: IdleMap, /// Used to catch false-negatives when waking workers needs_searching: AtomicBool, /// Total number of cores num_cores: usize, } pub(super) struct IdleMap { chunks: Vec, } pub(super) struct Snapshot { chunks: Vec, } /// Data synchronized by the scheduler mutex pub(super) struct Synced { /// Worker IDs that are currently sleeping sleepers: Vec, /// Cores available for workers available_cores: Vec>, } impl Idle { pub(super) fn new(cores: Vec>, num_workers: usize) -> (Idle, Synced) { let idle = Idle { num_searching: AtomicUsize::new(0), num_idle: AtomicUsize::new(cores.len()), idle_map: IdleMap::new(&cores), needs_searching: AtomicBool::new(false), num_cores: cores.len(), }; let synced = Synced { sleepers: Vec::with_capacity(num_workers), available_cores: cores, }; (idle, synced) } pub(super) fn needs_searching(&self) -> bool { self.needs_searching.load(Acquire) } pub(super) fn num_idle(&self, synced: &Synced) -> usize { #[cfg(not(loom))] debug_assert_eq!(synced.available_cores.len(), self.num_idle.load(Acquire)); synced.available_cores.len() } pub(super) fn num_searching(&self) -> usize { self.num_searching.load(Acquire) } pub(super) fn snapshot(&self, snapshot: &mut Snapshot) { snapshot.update(&self.idle_map) } /// Try to acquire an available core pub(super) fn try_acquire_available_core(&self, synced: &mut Synced) -> Option> { let ret = synced.available_cores.pop(); if let Some(core) = &ret { // Decrement the number of idle cores let num_idle = self.num_idle.load(Acquire) - 1; debug_assert_eq!(num_idle, synced.available_cores.len()); self.num_idle.store(num_idle, Release); self.idle_map.unset(core.index); debug_assert!(self.idle_map.matches(&synced.available_cores)); } ret } /// We need at least one searching worker pub(super) fn notify_local(&self, shared: &Shared) { if self.num_searching.load(Acquire) != 0 { // There already is a searching worker. Note, that this could be a // false positive. However, because this method is called **from** a // worker, we know that there is at least one worker currently // awake, so the scheduler won't deadlock. return; } if self.num_idle.load(Acquire) == 0 { self.needs_searching.store(true, Release); return; } // There aren't any searching workers. Try to initialize one if self .num_searching .compare_exchange(0, 1, AcqRel, Acquire) .is_err() { // Failing the compare_exchange means another thread concurrently // launched a searching worker. return; } super::counters::inc_num_unparks_local(); // Acquire the lock let synced = shared.synced.lock(); self.notify_synced(synced, shared); } /// Notifies a single worker pub(super) fn notify_remote(&self, synced: MutexGuard<'_, worker::Synced>, shared: &Shared) { if synced.idle.sleepers.is_empty() { self.needs_searching.store(true, Release); return; } // We need to establish a stronger barrier than with `notify_local` self.num_searching.fetch_add(1, AcqRel); self.notify_synced(synced, shared); } /// Notify a worker while synced fn notify_synced(&self, mut synced: MutexGuard<'_, worker::Synced>, shared: &Shared) { // Find a sleeping worker if let Some(worker) = synced.idle.sleepers.pop() { // Find an available core if let Some(mut core) = self.try_acquire_available_core(&mut synced.idle) { debug_assert!(!core.is_searching); core.is_searching = true; // Assign the core to the worker synced.assigned_cores[worker] = Some(core); // Drop the lock before notifying the condvar. drop(synced); super::counters::inc_num_unparks_remote(); // Notify the worker shared.condvars[worker].notify_one(); return; } else { synced.idle.sleepers.push(worker); } } super::counters::inc_notify_no_core(); // Set the `needs_searching` flag, this happens *while* the lock is held. self.needs_searching.store(true, Release); self.num_searching.fetch_sub(1, Release); // Explicit mutex guard drop to show that holding the guard to this // point is significant. `needs_searching` and `num_searching` must be // updated in the critical section. drop(synced); } pub(super) fn notify_mult( &self, synced: &mut worker::Synced, workers: &mut Vec, num: usize, ) { debug_assert!(workers.is_empty()); for _ in 0..num { if let Some(worker) = synced.idle.sleepers.pop() { // TODO: can this be switched to use next_available_core? if let Some(core) = synced.idle.available_cores.pop() { debug_assert!(!core.is_searching); self.idle_map.unset(core.index); synced.assigned_cores[worker] = Some(core); workers.push(worker); continue; } else { synced.idle.sleepers.push(worker); } } break; } if !workers.is_empty() { debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); let num_idle = synced.idle.available_cores.len(); self.num_idle.store(num_idle, Release); } else { #[cfg(not(loom))] debug_assert_eq!( synced.idle.available_cores.len(), self.num_idle.load(Acquire) ); self.needs_searching.store(true, Release); } } pub(super) fn shutdown(&self, synced: &mut worker::Synced, shared: &Shared) { // Wake every sleeping worker and assign a core to it. There may not be // enough sleeping workers for all cores, but other workers will // eventually find the cores and shut them down. while !synced.idle.sleepers.is_empty() && !synced.idle.available_cores.is_empty() { let worker = synced.idle.sleepers.pop().unwrap(); let core = self.try_acquire_available_core(&mut synced.idle).unwrap(); synced.assigned_cores[worker] = Some(core); shared.condvars[worker].notify_one(); } debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); // Wake up any other workers while let Some(index) = synced.idle.sleepers.pop() { shared.condvars[index].notify_one(); } } pub(super) fn shutdown_unassigned_cores(&self, handle: &Handle, shared: &Shared) { // If there are any remaining cores, shut them down here. // // This code is a bit convoluted to avoid lock-reentry. while let Some(core) = { let mut synced = shared.synced.lock(); self.try_acquire_available_core(&mut synced.idle) } { shared.shutdown_core(handle, core); } } /// The worker releases the given core, making it available to other workers /// that are waiting. pub(super) fn release_core(&self, synced: &mut worker::Synced, core: Box) { // The core should not be searching at this point debug_assert!(!core.is_searching); // Check that there are no pending tasks in the global queue debug_assert!(synced.inject.is_empty()); let num_idle = synced.idle.available_cores.len(); #[cfg(not(loom))] debug_assert_eq!(num_idle, self.num_idle.load(Acquire)); self.idle_map.set(core.index); // Store the core in the list of available cores synced.idle.available_cores.push(core); debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); // Update `num_idle` self.num_idle.store(num_idle + 1, Release); } pub(super) fn transition_worker_to_parked(&self, synced: &mut worker::Synced, index: usize) { // Store the worker index in the list of sleepers synced.idle.sleepers.push(index); // The worker's assigned core slot should be empty debug_assert!(synced.assigned_cores[index].is_none()); } pub(super) fn try_transition_worker_to_searching(&self, core: &mut Core) { debug_assert!(!core.is_searching); let num_searching = self.num_searching.load(Acquire); let num_idle = self.num_idle.load(Acquire); if 2 * num_searching >= self.num_cores - num_idle { return; } self.transition_worker_to_searching(core); } /// Needs to happen while synchronized in order to avoid races pub(super) fn transition_worker_to_searching_if_needed( &self, _synced: &mut Synced, core: &mut Core, ) -> bool { if self.needs_searching.load(Acquire) { // Needs to be called while holding the lock self.transition_worker_to_searching(core); true } else { false } } pub(super) fn transition_worker_to_searching(&self, core: &mut Core) { core.is_searching = true; self.num_searching.fetch_add(1, AcqRel); self.needs_searching.store(false, Release); } /// A lightweight transition from searching -> running. /// /// Returns `true` if this is the final searching worker. The caller /// **must** notify a new worker. pub(super) fn transition_worker_from_searching(&self) -> bool { let prev = self.num_searching.fetch_sub(1, AcqRel); debug_assert!(prev > 0); prev == 1 } } const BITS: usize = usize::BITS as usize; const BIT_MASK: usize = (usize::BITS - 1) as usize; impl IdleMap { fn new(cores: &[Box]) -> IdleMap { let ret = IdleMap::new_n(num_chunks(cores.len())); ret.set_all(cores); ret } fn new_n(n: usize) -> IdleMap { let chunks = (0..n).map(|_| AtomicUsize::new(0)).collect(); IdleMap { chunks } } fn set(&self, index: usize) { let (chunk, mask) = index_to_mask(index); let prev = self.chunks[chunk].load(Acquire); let next = prev | mask; self.chunks[chunk].store(next, Release); } fn set_all(&self, cores: &[Box]) { for core in cores { self.set(core.index); } } fn unset(&self, index: usize) { let (chunk, mask) = index_to_mask(index); let prev = self.chunks[chunk].load(Acquire); let next = prev & !mask; self.chunks[chunk].store(next, Release); } fn matches(&self, idle_cores: &[Box]) -> bool { let expect = IdleMap::new_n(self.chunks.len()); expect.set_all(idle_cores); for (i, chunk) in expect.chunks.iter().enumerate() { if chunk.load(Acquire) != self.chunks[i].load(Acquire) { return false; } } true } } impl Snapshot { pub(crate) fn new(idle: &Idle) -> Snapshot { let chunks = vec![0; idle.idle_map.chunks.len()]; let mut ret = Snapshot { chunks }; ret.update(&idle.idle_map); ret } fn update(&mut self, idle_map: &IdleMap) { for i in 0..self.chunks.len() { self.chunks[i] = idle_map.chunks[i].load(Acquire); } } pub(super) fn is_idle(&self, index: usize) -> bool { let (chunk, mask) = index_to_mask(index); debug_assert!( chunk < self.chunks.len(), "index={}; chunks={}", index, self.chunks.len() ); self.chunks[chunk] & mask == mask } } fn num_chunks(max_cores: usize) -> usize { (max_cores / BITS) + 1 } fn index_to_mask(index: usize) -> (usize, usize) { let mask = 1 << (index & BIT_MASK); let chunk = index / BITS; (chunk, mask) } fn num_active_workers(synced: &Synced) -> usize { synced.available_cores.capacity() - synced.available_cores.len() } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/mod.rs000064400000000000000000000041311046102023000217340ustar 00000000000000//! Multi-threaded runtime mod counters; use counters::Counters; mod handle; pub(crate) use handle::Handle; mod overflow; pub(crate) use overflow::Overflow; mod idle; use self::idle::Idle; mod stats; pub(crate) use stats::Stats; pub(crate) mod queue; mod worker; use worker::Core; pub(crate) use worker::{Context, Shared}; // TODO: implement task dump mod trace_mock; use trace_mock::TraceStatus; pub(crate) use worker::block_in_place; use crate::runtime::{ self, blocking, driver::{self, Driver}, scheduler, Config, }; use crate::util::RngSeedGenerator; use std::fmt; use std::future::Future; /// Work-stealing based thread pool for executing futures. pub(crate) struct MultiThread; // ===== impl MultiThread ===== impl MultiThread { pub(crate) fn new( size: usize, driver: Driver, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, ) -> (MultiThread, runtime::Handle) { let handle = worker::create( size, driver, driver_handle, blocking_spawner, seed_generator, config, ); (MultiThread, handle) } /// Blocks the current thread waiting for the future to complete. /// /// The future will execute on the current thread, but all spawned tasks /// will be executed on the thread pool. pub(crate) fn block_on(&self, handle: &scheduler::Handle, future: F) -> F::Output where F: Future, { crate::runtime::context::enter_runtime(handle, true, |blocking| { blocking.block_on(future).expect("failed to park thread") }) } pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) { match handle { scheduler::Handle::MultiThreadAlt(handle) => handle.shutdown(), _ => panic!("expected MultiThread scheduler"), } } } impl fmt::Debug for MultiThread { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("MultiThread").finish() } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/overflow.rs000064400000000000000000000010671046102023000230250ustar 00000000000000use crate::runtime::task; #[cfg(test)] use std::cell::RefCell; pub(crate) trait Overflow { fn push(&self, task: task::Notified); fn push_batch(&self, iter: I) where I: Iterator>; } #[cfg(test)] impl Overflow for RefCell>> { fn push(&self, task: task::Notified) { self.borrow_mut().push(task); } fn push_batch(&self, iter: I) where I: Iterator>, { self.borrow_mut().extend(iter); } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/park.rs000064400000000000000000000166111046102023000221200ustar 00000000000000//! Parks the runtime. //! //! A combination of the various resource driver park handles. use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Arc, Condvar, Mutex}; use crate::runtime::driver::{self, Driver}; use crate::util::TryLock; use std::sync::atomic::Ordering::SeqCst; use std::time::Duration; pub(crate) struct Parker { inner: Arc, } pub(crate) struct Unparker { inner: Arc, } struct Inner { /// Avoids entering the park if possible state: AtomicUsize, /// Used to coordinate access to the driver / condvar mutex: Mutex<()>, /// Condvar to block on if the driver is unavailable. condvar: Condvar, /// Resource (I/O, time, ...) driver shared: Arc, } const EMPTY: usize = 0; const PARKED_CONDVAR: usize = 1; const PARKED_DRIVER: usize = 2; const NOTIFIED: usize = 3; /// Shared across multiple Parker handles struct Shared { /// Shared driver. Only one thread at a time can use this driver: TryLock, } impl Parker { pub(crate) fn new(driver: Driver) -> Parker { Parker { inner: Arc::new(Inner { state: AtomicUsize::new(EMPTY), mutex: Mutex::new(()), condvar: Condvar::new(), shared: Arc::new(Shared { driver: TryLock::new(driver), }), }), } } pub(crate) fn unpark(&self) -> Unparker { Unparker { inner: self.inner.clone(), } } pub(crate) fn park(&mut self, handle: &driver::Handle) { self.inner.park(handle); } pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { // Only parking with zero is supported... assert_eq!(duration, Duration::from_millis(0)); if let Some(mut driver) = self.inner.shared.driver.try_lock() { driver.park_timeout(handle, duration) } } pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { self.inner.shutdown(handle); } } impl Clone for Parker { fn clone(&self) -> Parker { Parker { inner: Arc::new(Inner { state: AtomicUsize::new(EMPTY), mutex: Mutex::new(()), condvar: Condvar::new(), shared: self.inner.shared.clone(), }), } } } impl Unparker { pub(crate) fn unpark(&self, driver: &driver::Handle) { self.inner.unpark(driver); } } impl Inner { /// Parks the current thread for at most `dur`. fn park(&self, handle: &driver::Handle) { // If we were previously notified then we consume this notification and // return quickly. if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { return; } if let Some(mut driver) = self.shared.driver.try_lock() { self.park_driver(&mut driver, handle); } else { self.park_condvar(); } } fn park_condvar(&self) { // Otherwise we need to coordinate going to sleep let mut m = self.mutex.lock(); match self .state .compare_exchange(EMPTY, PARKED_CONDVAR, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read here, even though we know it will be `NOTIFIED`. // This is because `unpark` may have been called again since we read // `NOTIFIED` in the `compare_exchange` above. We must perform an // acquire operation that synchronizes with that `unpark` to observe // any writes it made before the call to unpark. To do that we must // read from the write it made to `state`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } Err(actual) => panic!("inconsistent park state; actual = {}", actual), } loop { m = self.condvar.wait(m).unwrap(); if self .state .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst) .is_ok() { // got a notification return; } // spurious wakeup, go back to sleep } } fn park_driver(&self, driver: &mut Driver, handle: &driver::Handle) { match self .state .compare_exchange(EMPTY, PARKED_DRIVER, SeqCst, SeqCst) { Ok(_) => {} Err(NOTIFIED) => { // We must read here, even though we know it will be `NOTIFIED`. // This is because `unpark` may have been called again since we read // `NOTIFIED` in the `compare_exchange` above. We must perform an // acquire operation that synchronizes with that `unpark` to observe // any writes it made before the call to unpark. To do that we must // read from the write it made to `state`. let old = self.state.swap(EMPTY, SeqCst); debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly"); return; } Err(actual) => panic!("inconsistent park state; actual = {}", actual), } driver.park(handle); match self.state.swap(EMPTY, SeqCst) { NOTIFIED => {} // got a notification, hurray! PARKED_DRIVER => {} // no notification, alas n => panic!("inconsistent park_timeout state: {}", n), } } fn unpark(&self, driver: &driver::Handle) { // To ensure the unparked thread will observe any writes we made before // this call, we must perform a release operation that `park` can // synchronize with. To do that we must write `NOTIFIED` even if `state` // is already `NOTIFIED`. That is why this must be a swap rather than a // compare-and-swap that returns if it reads `NOTIFIED` on failure. match self.state.swap(NOTIFIED, SeqCst) { EMPTY => {} // no one was waiting NOTIFIED => {} // already unparked PARKED_CONDVAR => self.unpark_condvar(), PARKED_DRIVER => driver.unpark(), actual => panic!("inconsistent state in unpark; actual = {}", actual), } } fn unpark_condvar(&self) { // There is a period between when the parked thread sets `state` to // `PARKED` (or last checked `state` in the case of a spurious wake // up) and when it actually waits on `cvar`. If we were to notify // during this period it would be ignored and then when the parked // thread went to sleep it would never wake up. Fortunately, it has // `lock` locked at this stage so we can acquire `lock` to wait until // it is ready to receive the notification. // // Releasing `lock` before the call to `notify_one` means that when the // parked thread wakes it doesn't get woken only to have to wait for us // to release `lock`. drop(self.mutex.lock()); self.condvar.notify_one() } fn shutdown(&self, handle: &driver::Handle) { if let Some(mut driver) = self.shared.driver.try_lock() { driver.shutdown(handle); } self.condvar.notify_all(); } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/queue.rs000064400000000000000000000464151046102023000223140ustar 00000000000000//! Run-queue structures to support a work-stealing scheduler use crate::loom::cell::UnsafeCell; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread_alt::{Overflow, Stats}; use crate::runtime::task; use std::mem::{self, MaybeUninit}; use std::ptr; use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; // Use wider integers when possible to increase ABA resilience. // // See issue #5041: . cfg_has_atomic_u64! { type UnsignedShort = u32; type UnsignedLong = u64; type AtomicUnsignedShort = crate::loom::sync::atomic::AtomicU32; type AtomicUnsignedLong = crate::loom::sync::atomic::AtomicU64; } cfg_not_has_atomic_u64! { type UnsignedShort = u16; type UnsignedLong = u32; type AtomicUnsignedShort = crate::loom::sync::atomic::AtomicU16; type AtomicUnsignedLong = crate::loom::sync::atomic::AtomicU32; } /// Producer handle. May only be used from a single thread. pub(crate) struct Local { inner: Arc>, } /// Consumer handle. May be used from many threads. pub(crate) struct Steal(Arc>); #[repr(align(128))] pub(crate) struct Inner { /// Concurrently updated by many threads. /// /// Contains two `UnsignedShort` values. The `LSB` byte is the "real" head of /// the queue. The `UnsignedShort` in the `MSB` is set by a stealer in process /// of stealing values. It represents the first value being stolen in the /// batch. The `UnsignedShort` indices are intentionally wider than strictly /// required for buffer indexing in order to provide ABA mitigation and make /// it possible to distinguish between full and empty buffers. /// /// When both `UnsignedShort` values are the same, there is no active /// stealer. /// /// Tracking an in-progress stealer prevents a wrapping scenario. head: AtomicUnsignedLong, /// Only updated by producer thread but read by many threads. tail: AtomicUnsignedShort, /// Elements buffer: Box<[UnsafeCell>>]>, mask: usize, } unsafe impl Send for Inner {} unsafe impl Sync for Inner {} /// Create a new local run-queue pub(crate) fn local(capacity: usize) -> (Steal, Local) { assert!(capacity <= 4096); assert!(capacity >= 1); let mut buffer = Vec::with_capacity(capacity); for _ in 0..capacity { buffer.push(UnsafeCell::new(MaybeUninit::uninit())); } let inner = Arc::new(Inner { head: AtomicUnsignedLong::new(0), tail: AtomicUnsignedShort::new(0), buffer: buffer.into_boxed_slice(), mask: capacity - 1, }); let local = Local { inner: inner.clone(), }; let remote = Steal(inner); (remote, local) } impl Local { /// How many tasks can be pushed into the queue pub(crate) fn remaining_slots(&self) -> usize { self.inner.remaining_slots() } pub(crate) fn max_capacity(&self) -> usize { self.inner.buffer.len() } /// Returns `true` if there are no entries in the queue pub(crate) fn is_empty(&self) -> bool { self.inner.is_empty() } pub(crate) fn can_steal(&self) -> bool { self.remaining_slots() >= self.max_capacity() - self.max_capacity() / 2 } /// Pushes a batch of tasks to the back of the queue. All tasks must fit in /// the local queue. /// /// # Panics /// /// The method panics if there is not enough capacity to fit in the queue. pub(crate) fn push_back(&mut self, tasks: impl ExactSizeIterator>) { let len = tasks.len(); assert!(len <= self.inner.buffer.len()); if len == 0 { // Nothing to do return; } let head = self.inner.head.load(Acquire); let (steal, real) = unpack(head); // safety: this is the **only** thread that updates this cell. let mut tail = unsafe { self.inner.tail.unsync_load() }; if tail.wrapping_sub(steal) <= (self.inner.buffer.len() - len) as UnsignedShort { // Yes, this if condition is structured a bit weird (first block // does nothing, second returns an error). It is this way to match // `push_back_or_overflow`. } else { panic!( "not enough capacity; len={}; tail={}; steal={}; real={}", len, tail, steal, real ); } for task in tasks { let idx = tail as usize & self.inner.mask; self.inner.buffer[idx].with_mut(|ptr| { // Write the task to the slot // // Safety: There is only one producer and the above `if` // condition ensures we don't touch a cell if there is a // value, thus no consumer. unsafe { ptr::write((*ptr).as_mut_ptr(), task); } }); tail = tail.wrapping_add(1); } self.inner.tail.store(tail, Release); } /// Pushes a task to the back of the local queue, if there is not enough /// capacity in the queue, this triggers the overflow operation. /// /// When the queue overflows, half of the current contents of the queue is /// moved to the given Injection queue. This frees up capacity for more /// tasks to be pushed into the local queue. pub(crate) fn push_back_or_overflow>( &mut self, mut task: task::Notified, overflow: &O, stats: &mut Stats, ) { let tail = loop { let head = self.inner.head.load(Acquire); let (steal, real) = unpack(head); // safety: this is the **only** thread that updates this cell. let tail = unsafe { self.inner.tail.unsync_load() }; if tail.wrapping_sub(steal) < self.inner.buffer.len() as UnsignedShort { // There is capacity for the task break tail; } else if steal != real { super::counters::inc_num_overflows(); // Concurrently stealing, this will free up capacity, so only // push the task onto the inject queue overflow.push(task); return; } else { super::counters::inc_num_overflows(); // Push the current task and half of the queue into the // inject queue. match self.push_overflow(task, real, tail, overflow, stats) { Ok(_) => return, // Lost the race, try again Err(v) => { task = v; } } } }; self.push_back_finish(task, tail); } // Second half of `push_back` fn push_back_finish(&self, task: task::Notified, tail: UnsignedShort) { // Map the position to a slot index. let idx = tail as usize & self.inner.mask; self.inner.buffer[idx].with_mut(|ptr| { // Write the task to the slot // // Safety: There is only one producer and the above `if` // condition ensures we don't touch a cell if there is a // value, thus no consumer. unsafe { ptr::write((*ptr).as_mut_ptr(), task); } }); // Make the task available. Synchronizes with a load in // `steal_into2`. self.inner.tail.store(tail.wrapping_add(1), Release); } /// Moves a batch of tasks into the inject queue. /// /// This will temporarily make some of the tasks unavailable to stealers. /// Once `push_overflow` is done, a notification is sent out, so if other /// workers "missed" some of the tasks during a steal, they will get /// another opportunity. #[inline(never)] fn push_overflow>( &mut self, task: task::Notified, head: UnsignedShort, tail: UnsignedShort, overflow: &O, stats: &mut Stats, ) -> Result<(), task::Notified> { // How many elements are we taking from the local queue. // // This is one less than the number of tasks pushed to the inject // queue as we are also inserting the `task` argument. let num_tasks_taken: UnsignedShort = (self.inner.buffer.len() / 2) as UnsignedShort; assert_eq!( tail.wrapping_sub(head) as usize, self.inner.buffer.len(), "queue is not full; tail = {}; head = {}", tail, head ); let prev = pack(head, head); // Claim a bunch of tasks // // We are claiming the tasks **before** reading them out of the buffer. // This is safe because only the **current** thread is able to push new // tasks. // // There isn't really any need for memory ordering... Relaxed would // work. This is because all tasks are pushed into the queue from the // current thread (or memory has been acquired if the local queue handle // moved). if self .inner .head .compare_exchange( prev, pack( head.wrapping_add(num_tasks_taken), head.wrapping_add(num_tasks_taken), ), Release, Relaxed, ) .is_err() { // We failed to claim the tasks, losing the race. Return out of // this function and try the full `push` routine again. The queue // may not be full anymore. return Err(task); } /// An iterator that takes elements out of the run queue. struct BatchTaskIter<'a, T: 'static> { buffer: &'a [UnsafeCell>>], mask: usize, head: UnsignedLong, i: UnsignedLong, num: UnsignedShort, } impl<'a, T: 'static> Iterator for BatchTaskIter<'a, T> { type Item = task::Notified; #[inline] fn next(&mut self) -> Option> { if self.i == UnsignedLong::from(self.num) { None } else { let i_idx = self.i.wrapping_add(self.head) as usize & self.mask; let slot = &self.buffer[i_idx]; // safety: Our CAS from before has assumed exclusive ownership // of the task pointers in this range. let task = slot.with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); self.i += 1; Some(task) } } } // safety: The CAS above ensures that no consumer will look at these // values again, and we are the only producer. let batch_iter = BatchTaskIter { buffer: &self.inner.buffer, mask: self.inner.mask, head: head as UnsignedLong, i: 0, num: num_tasks_taken, }; overflow.push_batch(batch_iter.chain(std::iter::once(task))); // Add 1 to factor in the task currently being scheduled. stats.incr_overflow_count(); Ok(()) } /// Pops a task from the local queue. pub(crate) fn pop(&mut self) -> Option> { let mut head = self.inner.head.load(Acquire); let idx = loop { let (steal, real) = unpack(head); // safety: this is the **only** thread that updates this cell. let tail = unsafe { self.inner.tail.unsync_load() }; if real == tail { // queue is empty return None; } let next_real = real.wrapping_add(1); // If `steal == real` there are no concurrent stealers. Both `steal` // and `real` are updated. let next = if steal == real { pack(next_real, next_real) } else { assert_ne!(steal, next_real); pack(steal, next_real) }; // Attempt to claim a task. let res = self .inner .head .compare_exchange(head, next, AcqRel, Acquire); match res { Ok(_) => break real as usize & self.inner.mask, Err(actual) => head = actual, } }; Some(self.inner.buffer[idx].with(|ptr| unsafe { ptr::read(ptr).assume_init() })) } } impl Steal { pub(crate) fn is_empty(&self) -> bool { self.0.is_empty() } /// Steals half the tasks from self and place them into `dst`. pub(crate) fn steal_into( &self, dst: &mut Local, dst_stats: &mut Stats, ) -> Option> { // Safety: the caller is the only thread that mutates `dst.tail` and // holds a mutable reference. let dst_tail = unsafe { dst.inner.tail.unsync_load() }; // To the caller, `dst` may **look** empty but still have values // contained in the buffer. If another thread is concurrently stealing // from `dst` there may not be enough capacity to steal. let (steal, _) = unpack(dst.inner.head.load(Acquire)); if dst_tail.wrapping_sub(steal) > self.0.buffer.len() as UnsignedShort / 2 { // we *could* try to steal less here, but for simplicity, we're just // going to abort. return None; } // Steal the tasks into `dst`'s buffer. This does not yet expose the // tasks in `dst`. let mut n = self.steal_into2(dst, dst_tail); if n == 0 { // No tasks were stolen return None; } super::counters::inc_num_steals(); dst_stats.incr_steal_count(n as u16); dst_stats.incr_steal_operations(); // We are returning a task here n -= 1; let ret_pos = dst_tail.wrapping_add(n); let ret_idx = ret_pos as usize & dst.inner.mask; // safety: the value was written as part of `steal_into2` and not // exposed to stealers, so no other thread can access it. let ret = dst.inner.buffer[ret_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); if n == 0 { // The `dst` queue is empty, but a single task was stolen return Some(ret); } // Make the stolen items available to consumers dst.inner.tail.store(dst_tail.wrapping_add(n), Release); Some(ret) } // Steal tasks from `self`, placing them into `dst`. Returns the number of // tasks that were stolen. fn steal_into2(&self, dst: &mut Local, dst_tail: UnsignedShort) -> UnsignedShort { let mut prev_packed = self.0.head.load(Acquire); let mut next_packed; let n = loop { let (src_head_steal, src_head_real) = unpack(prev_packed); let src_tail = self.0.tail.load(Acquire); // If these two do not match, another thread is concurrently // stealing from the queue. if src_head_steal != src_head_real { return 0; } // Number of available tasks to steal let n = src_tail.wrapping_sub(src_head_real); let n = n - n / 2; if n == 0 { // No tasks available to steal return 0; } // Update the real head index to acquire the tasks. let steal_to = src_head_real.wrapping_add(n); assert_ne!(src_head_steal, steal_to); next_packed = pack(src_head_steal, steal_to); // Claim all those tasks. This is done by incrementing the "real" // head but not the steal. By doing this, no other thread is able to // steal from this queue until the current thread completes. let res = self .0 .head .compare_exchange(prev_packed, next_packed, AcqRel, Acquire); match res { Ok(_) => break n, Err(actual) => prev_packed = actual, } }; debug_assert!( n <= (self.0.buffer.len() - self.0.buffer.len() / 2) as UnsignedShort, "actual = {}", n ); let (first, _) = unpack(next_packed); // Take all the tasks for i in 0..n { // Compute the positions let src_pos = first.wrapping_add(i); let dst_pos = dst_tail.wrapping_add(i); // Map to slots let src_idx = src_pos as usize & self.0.mask; let dst_idx = dst_pos as usize & self.0.mask; // Read the task // // safety: We acquired the task with the atomic exchange above. let task = self.0.buffer[src_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); // Write the task to the new slot // // safety: `dst` queue is empty and we are the only producer to // this queue. dst.inner.buffer[dst_idx] .with_mut(|ptr| unsafe { ptr::write((*ptr).as_mut_ptr(), task) }); } let mut prev_packed = next_packed; // Update `src_head_steal` to match `src_head_real` signalling that the // stealing routine is complete. loop { let head = unpack(prev_packed).1; next_packed = pack(head, head); let res = self .0 .head .compare_exchange(prev_packed, next_packed, AcqRel, Acquire); match res { Ok(_) => return n, Err(actual) => { let (actual_steal, actual_real) = unpack(actual); assert_ne!(actual_steal, actual_real); prev_packed = actual; } } } } } cfg_unstable_metrics! { impl Steal { pub(crate) fn len(&self) -> usize { self.0.len() as _ } } } impl Clone for Steal { fn clone(&self) -> Steal { Steal(self.0.clone()) } } impl Drop for Local { fn drop(&mut self) { if !std::thread::panicking() { assert!(self.pop().is_none(), "queue not empty"); } } } impl Inner { fn remaining_slots(&self) -> usize { let (steal, _) = unpack(self.head.load(Acquire)); let tail = self.tail.load(Acquire); self.buffer.len() - (tail.wrapping_sub(steal) as usize) } fn len(&self) -> UnsignedShort { let (_, head) = unpack(self.head.load(Acquire)); let tail = self.tail.load(Acquire); tail.wrapping_sub(head) } fn is_empty(&self) -> bool { self.len() == 0 } } /// Split the head value into the real head and the index a stealer is working /// on. fn unpack(n: UnsignedLong) -> (UnsignedShort, UnsignedShort) { let real = n & UnsignedShort::MAX as UnsignedLong; let steal = n >> (mem::size_of::() * 8); (steal as UnsignedShort, real as UnsignedShort) } /// Join the two head values fn pack(steal: UnsignedShort, real: UnsignedShort) -> UnsignedLong { (real as UnsignedLong) | ((steal as UnsignedLong) << (mem::size_of::() * 8)) } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/stats.rs000064400000000000000000000133101046102023000223120ustar 00000000000000use crate::runtime::{Config, MetricsBatch, WorkerMetrics}; use std::cmp; use std::time::{Duration, Instant}; /// Per-worker statistics. This is used for both tuning the scheduler and /// reporting runtime-level metrics/stats. pub(crate) struct Stats { /// The metrics batch used to report runtime-level metrics/stats to the /// user. batch: MetricsBatch, /// Exponentially-weighted moving average of time spent polling scheduled a /// task. /// /// Tracked in nanoseconds, stored as a `f64` since that is what we use with /// the EWMA calculations task_poll_time_ewma: f64, } /// Transient state pub(crate) struct Ephemeral { /// Instant at which work last resumed (continued after park). /// /// This duplicates the value stored in `MetricsBatch`. We will unify /// `Stats` and `MetricsBatch` when we stabilize metrics. processing_scheduled_tasks_started_at: Instant, /// Number of tasks polled in the batch of scheduled tasks tasks_polled_in_batch: usize, /// Used to ensure calls to start / stop batch are paired #[cfg(debug_assertions)] batch_started: bool, } impl Ephemeral { pub(crate) fn new() -> Ephemeral { Ephemeral { processing_scheduled_tasks_started_at: Instant::now(), tasks_polled_in_batch: 0, #[cfg(debug_assertions)] batch_started: false, } } } /// How to weigh each individual poll time, value is plucked from thin air. const TASK_POLL_TIME_EWMA_ALPHA: f64 = 0.1; /// Ideally, we wouldn't go above this, value is plucked from thin air. const TARGET_GLOBAL_QUEUE_INTERVAL: f64 = Duration::from_micros(200).as_nanos() as f64; /// Max value for the global queue interval. This is 2x the previous default const MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 127; /// This is the previous default const TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 61; impl Stats { pub(crate) const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL; pub(crate) fn new(worker_metrics: &WorkerMetrics) -> Stats { // Seed the value with what we hope to see. let task_poll_time_ewma = TARGET_GLOBAL_QUEUE_INTERVAL / TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL as f64; Stats { batch: MetricsBatch::new(worker_metrics), task_poll_time_ewma, } } pub(crate) fn tuned_global_queue_interval(&self, config: &Config) -> u32 { // If an interval is explicitly set, don't tune. if let Some(configured) = config.global_queue_interval { return configured; } // As of Rust 1.45, casts from f64 -> u32 are saturating, which is fine here. let tasks_per_interval = (TARGET_GLOBAL_QUEUE_INTERVAL / self.task_poll_time_ewma) as u32; cmp::max( // If we are using self-tuning, we don't want to return less than 2 as that would result in the // global queue always getting checked first. 2, cmp::min( MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL, tasks_per_interval, ), ) } pub(crate) fn submit(&mut self, to: &WorkerMetrics) { self.batch.submit(to, self.task_poll_time_ewma as u64); } pub(crate) fn about_to_park(&mut self) { self.batch.about_to_park(); } pub(crate) fn unparked(&mut self) { self.batch.unparked(); } pub(crate) fn inc_local_schedule_count(&mut self) { self.batch.inc_local_schedule_count(); } pub(crate) fn start_processing_scheduled_tasks(&mut self, ephemeral: &mut Ephemeral) { self.batch.start_processing_scheduled_tasks(); #[cfg(debug_assertions)] { debug_assert!(!ephemeral.batch_started); ephemeral.batch_started = true; } ephemeral.processing_scheduled_tasks_started_at = Instant::now(); ephemeral.tasks_polled_in_batch = 0; } pub(crate) fn end_processing_scheduled_tasks(&mut self, ephemeral: &mut Ephemeral) { self.batch.end_processing_scheduled_tasks(); #[cfg(debug_assertions)] { debug_assert!(ephemeral.batch_started); ephemeral.batch_started = false; } // Update the EWMA task poll time if ephemeral.tasks_polled_in_batch > 0 { let now = Instant::now(); // If we "overflow" this conversion, we have bigger problems than // slightly off stats. let elapsed = (now - ephemeral.processing_scheduled_tasks_started_at).as_nanos() as f64; let num_polls = ephemeral.tasks_polled_in_batch as f64; // Calculate the mean poll duration for a single task in the batch let mean_poll_duration = elapsed / num_polls; // Compute the alpha weighted by the number of tasks polled this batch. let weighted_alpha = 1.0 - (1.0 - TASK_POLL_TIME_EWMA_ALPHA).powf(num_polls); // Now compute the new weighted average task poll time. self.task_poll_time_ewma = weighted_alpha * mean_poll_duration + (1.0 - weighted_alpha) * self.task_poll_time_ewma; } } pub(crate) fn start_poll(&mut self, ephemeral: &mut Ephemeral) { self.batch.start_poll(); ephemeral.tasks_polled_in_batch += 1; } pub(crate) fn end_poll(&mut self) { self.batch.end_poll(); } pub(crate) fn incr_steal_count(&mut self, by: u16) { self.batch.incr_steal_count(by); } pub(crate) fn incr_steal_operations(&mut self) { self.batch.incr_steal_operations(); } pub(crate) fn incr_overflow_count(&mut self) { self.batch.incr_overflow_count(); } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/trace.rs000064400000000000000000000034721046102023000222620ustar 00000000000000use crate::loom::sync::atomic::{AtomicBool, Ordering}; use crate::loom::sync::{Barrier, Mutex}; use crate::runtime::dump::Dump; use crate::runtime::scheduler::multi_thread_alt::Handle; use crate::sync::notify::Notify; /// Tracing status of the worker. pub(super) struct TraceStatus { pub(super) trace_requested: AtomicBool, pub(super) trace_start: Barrier, pub(super) trace_end: Barrier, pub(super) result_ready: Notify, pub(super) trace_result: Mutex>, } impl TraceStatus { pub(super) fn new(remotes_len: usize) -> Self { Self { trace_requested: AtomicBool::new(false), trace_start: Barrier::new(remotes_len), trace_end: Barrier::new(remotes_len), result_ready: Notify::new(), trace_result: Mutex::new(None), } } pub(super) fn trace_requested(&self) -> bool { self.trace_requested.load(Ordering::Relaxed) } pub(super) async fn start_trace_request(&self, handle: &Handle) { while self .trace_requested .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) .is_err() { handle.notify_all(); crate::task::yield_now().await; } } pub(super) fn stash_result(&self, dump: Dump) { let _ = self.trace_result.lock().insert(dump); self.result_ready.notify_one(); } pub(super) fn take_result(&self) -> Option { self.trace_result.lock().take() } pub(super) async fn end_trace_request(&self, handle: &Handle) { while self .trace_requested .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed) .is_err() { handle.notify_all(); crate::task::yield_now().await; } } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/trace_mock.rs000064400000000000000000000002771046102023000232730ustar 00000000000000pub(super) struct TraceStatus {} impl TraceStatus { pub(super) fn new(_: usize) -> Self { Self {} } pub(super) fn trace_requested(&self) -> bool { false } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/worker/metrics.rs000064400000000000000000000003721046102023000241370ustar 00000000000000use super::Shared; impl Shared { pub(crate) fn injection_queue_depth(&self) -> usize { self.inject.len() } pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { self.remotes[worker].steal.len() } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/worker/taskdump.rs000064400000000000000000000044741046102023000243300ustar 00000000000000use super::{Core, Handle, Shared}; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread_alt::Stats; use crate::runtime::task::trace::trace_multi_thread; use crate::runtime::{dump, WorkerMetrics}; use std::time::Duration; impl Handle { pub(super) fn trace_core(&self, mut core: Box) -> Box { core.is_traced = false; if core.is_shutdown { return core; } // wait for other workers, or timeout without tracing let timeout = Duration::from_millis(250); // a _very_ generous timeout let barrier = if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) { barrier } else { // don't attempt to trace return core; }; if !barrier.is_leader() { // wait for leader to finish tracing self.shared.trace_status.trace_end.wait(); return core; } // trace let owned = &self.shared.owned; let mut local = self.shared.steal_all(); let synced = &self.shared.synced; let injection = &self.shared.inject; // safety: `trace_multi_thread` is invoked with the same `synced` that `injection` // was created with. let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) } .into_iter() .map(dump::Task::new) .collect(); let result = dump::Dump::new(traces); // stash the result self.shared.trace_status.stash_result(result); // allow other workers to proceed self.shared.trace_status.trace_end.wait(); core } } impl Shared { /// Steal all tasks from remotes into a single local queue. pub(super) fn steal_all(&self) -> super::queue::Local> { let (_steal, mut local) = super::queue::local(); let worker_metrics = WorkerMetrics::new(); let mut stats = Stats::new(&worker_metrics); for remote in self.remotes.iter() { let steal = &remote.steal; while !steal.is_empty() { if let Some(task) = steal.steal_into(&mut local, &mut stats) { local.push_back([task].into_iter()); } } } local } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/worker/taskdump_mock.rs000064400000000000000000000002031046102023000253230ustar 00000000000000use super::{Core, Handle}; impl Handle { pub(super) fn trace_core(&self, core: Box) -> Box { core } } tokio-1.43.1/src/runtime/scheduler/multi_thread_alt/worker.rs000064400000000000000000001500601046102023000224710ustar 00000000000000//! A scheduler is initialized with a fixed number of workers. Each worker is //! driven by a thread. Each worker has a "core" which contains data such as the //! run queue and other state. When `block_in_place` is called, the worker's //! "core" is handed off to a new thread allowing the scheduler to continue to //! make progress while the originating thread blocks. //! //! # Shutdown //! //! Shutting down the runtime involves the following steps: //! //! 1. The Shared::close method is called. This closes the inject queue and //! `OwnedTasks` instance and wakes up all worker threads. //! //! 2. Each worker thread observes the close signal next time it runs //! Core::maintenance by checking whether the inject queue is closed. //! The `Core::is_shutdown` flag is set to true. //! //! 3. The worker thread calls `pre_shutdown` in parallel. Here, the worker //! will keep removing tasks from `OwnedTasks` until it is empty. No new //! tasks can be pushed to the `OwnedTasks` during or after this step as it //! was closed in step 1. //! //! 5. The workers call Shared::shutdown to enter the single-threaded phase of //! shutdown. These calls will push their core to `Shared::shutdown_cores`, //! and the last thread to push its core will finish the shutdown procedure. //! //! 6. The local run queue of each core is emptied, then the inject queue is //! emptied. //! //! At this point, shutdown has completed. It is not possible for any of the //! collections to contain any tasks at this point, as each collection was //! closed first, then emptied afterwards. //! //! ## Spawns during shutdown //! //! When spawning tasks during shutdown, there are two cases: //! //! * The spawner observes the `OwnedTasks` being open, and the inject queue is //! closed. //! * The spawner observes the `OwnedTasks` being closed and doesn't check the //! inject queue. //! //! The first case can only happen if the `OwnedTasks::bind` call happens before //! or during step 1 of shutdown. In this case, the runtime will clean up the //! task in step 3 of shutdown. //! //! In the latter case, the task was not spawned and the task is immediately //! cancelled by the spawner. //! //! The correctness of shutdown requires both the inject queue and `OwnedTasks` //! collection to have a closed bit. With a close bit on only the inject queue, //! spawning could run in to a situation where a task is successfully bound long //! after the runtime has shut down. With a close bit on only the `OwnedTasks`, //! the first spawning situation could result in the notification being pushed //! to the inject queue after step 6 of shutdown, which would leave a task in //! the inject queue indefinitely. This would be a ref-count cycle and a memory //! leak. use crate::loom::sync::{Arc, Condvar, Mutex, MutexGuard}; use crate::runtime; use crate::runtime::driver::Driver; use crate::runtime::scheduler::multi_thread_alt::{ idle, queue, stats, Counters, Handle, Idle, Overflow, Stats, TraceStatus, }; use crate::runtime::scheduler::{self, inject, Lock}; use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks}; use crate::runtime::{blocking, coop, driver, task, Config, SchedulerMetrics, WorkerMetrics}; use crate::runtime::{context, TaskHooks}; use crate::util::atomic_cell::AtomicCell; use crate::util::rand::{FastRand, RngSeedGenerator}; use std::cell::{Cell, RefCell}; use std::task::Waker; use std::time::Duration; use std::{cmp, thread}; cfg_unstable_metrics! { mod metrics; } mod taskdump_mock; /// A scheduler worker /// /// Data is stack-allocated and never migrates threads pub(super) struct Worker { /// Used to schedule bookkeeping tasks every so often. tick: u32, /// True if the scheduler is being shutdown pub(super) is_shutdown: bool, /// True if the scheduler is being traced is_traced: bool, /// Counter used to track when to poll from the local queue vs. the /// injection queue num_seq_local_queue_polls: u32, /// How often to check the global queue global_queue_interval: u32, /// Used to collect a list of workers to notify workers_to_notify: Vec, /// Snapshot of idle core list. This helps speedup stealing idle_snapshot: idle::Snapshot, stats: stats::Ephemeral, } /// Core data /// /// Data is heap-allocated and migrates threads. #[repr(align(128))] pub(super) struct Core { /// Index holding this core's remote/shared state. pub(super) index: usize, lifo_slot: Option, /// The worker-local run queue. run_queue: queue::Local>, /// True if the worker is currently searching for more work. Searching /// involves attempting to steal from other workers. pub(super) is_searching: bool, /// Per-worker runtime stats stats: Stats, /// Fast random number generator. rand: FastRand, } /// State shared across all workers pub(crate) struct Shared { /// Per-core remote state. remotes: Box<[Remote]>, /// Global task queue used for: /// 1. Submit work to the scheduler while **not** currently on a worker thread. /// 2. Submit work to the scheduler when a worker run queue is saturated pub(super) inject: inject::Shared>, /// Coordinates idle workers idle: Idle, /// Collection of all active tasks spawned onto this executor. pub(super) owned: OwnedTasks>, /// Data synchronized by the scheduler mutex pub(super) synced: Mutex, /// Power's Tokio's I/O, timers, etc... the responsibility of polling the /// driver is shared across workers. driver: AtomicCell, /// Condition variables used to unblock worker threads. Each worker thread /// has its own `condvar` it waits on. pub(super) condvars: Vec, /// The number of cores that have observed the trace signal. pub(super) trace_status: TraceStatus, /// Scheduler configuration options config: Config, /// Collects metrics from the runtime. pub(super) scheduler_metrics: SchedulerMetrics, pub(super) worker_metrics: Box<[WorkerMetrics]>, /// Only held to trigger some code on drop. This is used to get internal /// runtime metrics that can be useful when doing performance /// investigations. This does nothing (empty struct, no drop impl) unless /// the `tokio_internal_mt_counters` `cfg` flag is set. _counters: Counters, } /// Data synchronized by the scheduler mutex pub(crate) struct Synced { /// When worker is notified, it is assigned a core. The core is placed here /// until the worker wakes up to take it. pub(super) assigned_cores: Vec>>, /// Cores that have observed the shutdown signal /// /// The core is **not** placed back in the worker to avoid it from being /// stolen by a thread that was spawned as part of `block_in_place`. shutdown_cores: Vec>, /// The driver goes here when shutting down shutdown_driver: Option>, /// Synchronized state for `Idle`. pub(super) idle: idle::Synced, /// Synchronized state for `Inject`. pub(crate) inject: inject::Synced, } /// Used to communicate with a worker from other threads. struct Remote { /// When a task is scheduled from a worker, it is stored in this slot. The /// worker will check this slot for a task **before** checking the run /// queue. This effectively results in the **last** scheduled task to be run /// next (LIFO). This is an optimization for improving locality which /// benefits message passing patterns and helps to reduce latency. // lifo_slot: Lifo, /// Steals tasks from this worker. pub(super) steal: queue::Steal>, } /// Thread-local context pub(crate) struct Context { // Current scheduler's handle handle: Arc, /// Worker index index: usize, /// True when the LIFO slot is enabled lifo_enabled: Cell, /// Core data core: RefCell>>, /// Used to pass cores to other threads when `block_in_place` is called handoff_core: Arc>, /// Tasks to wake after resource drivers are polled. This is mostly to /// handle yielded tasks. pub(crate) defer: RefCell>, } /// Running a task may consume the core. If the core is still available when /// running the task completes, it is returned. Otherwise, the worker will need /// to stop processing. type RunResult = Result, ()>; type NextTaskResult = Result<(Option, Box), ()>; /// A task handle type Task = task::Task>; /// A notified task handle type Notified = task::Notified>; /// Value picked out of thin-air. Running the LIFO slot a handful of times /// seems sufficient to benefit from locality. More than 3 times probably is /// overweighing. The value can be tuned in the future with data that shows /// improvements. const MAX_LIFO_POLLS_PER_TICK: usize = 3; pub(super) fn create( num_cores: usize, driver: Driver, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, ) -> runtime::Handle { let mut num_workers = num_cores; // If the driver is enabled, we need an extra thread to handle polling the // driver when all cores are busy. if driver.is_enabled() { num_workers += 1; } let mut cores = Vec::with_capacity(num_cores); let mut remotes = Vec::with_capacity(num_cores); // Worker metrics are actually core based let mut worker_metrics = Vec::with_capacity(num_cores); // Create the local queues for i in 0..num_cores { let (steal, run_queue) = queue::local(config.local_queue_capacity); let metrics = WorkerMetrics::from_config(&config); let stats = Stats::new(&metrics); cores.push(Box::new(Core { index: i, lifo_slot: None, run_queue, is_searching: false, stats, rand: FastRand::from_seed(config.seed_generator.next_seed()), })); remotes.push(Remote { steal, // lifo_slot: Lifo::new(), }); worker_metrics.push(metrics); } // Allocate num-cores + 1 workers, so one worker can handle the I/O driver, // if needed. let (idle, idle_synced) = Idle::new(cores, num_workers); let (inject, inject_synced) = inject::Shared::new(); let handle = Arc::new(Handle { task_hooks: TaskHooks { task_spawn_callback: config.before_spawn.clone(), task_terminate_callback: config.after_termination.clone(), }, shared: Shared { remotes: remotes.into_boxed_slice(), inject, idle, owned: OwnedTasks::new(num_cores), synced: Mutex::new(Synced { assigned_cores: (0..num_workers).map(|_| None).collect(), shutdown_cores: Vec::with_capacity(num_cores), shutdown_driver: None, idle: idle_synced, inject: inject_synced, }), driver: AtomicCell::new(Some(Box::new(driver))), condvars: (0..num_workers).map(|_| Condvar::new()).collect(), trace_status: TraceStatus::new(num_cores), config, scheduler_metrics: SchedulerMetrics::new(), worker_metrics: worker_metrics.into_boxed_slice(), _counters: Counters, }, driver: driver_handle, blocking_spawner, seed_generator, }); let rt_handle = runtime::Handle { inner: scheduler::Handle::MultiThreadAlt(handle), }; // Eagerly start worker threads for index in 0..num_workers { let handle = rt_handle.inner.expect_multi_thread_alt(); let h2 = handle.clone(); let handoff_core = Arc::new(AtomicCell::new(None)); handle .blocking_spawner .spawn_blocking(&rt_handle, move || run(index, h2, handoff_core, false)); } rt_handle } #[track_caller] pub(crate) fn block_in_place(f: F) -> R where F: FnOnce() -> R, { // Try to steal the worker core back struct Reset(coop::Budget); impl Drop for Reset { fn drop(&mut self) { with_current(|maybe_cx| { if let Some(cx) = maybe_cx { let core = cx.handoff_core.take(); let mut cx_core = cx.core.borrow_mut(); assert!(cx_core.is_none()); *cx_core = core; // Reset the task budget as we are re-entering the // runtime. coop::set(self.0); } }); } } let mut had_entered = false; let setup_result = with_current(|maybe_cx| { match ( crate::runtime::context::current_enter_context(), maybe_cx.is_some(), ) { (context::EnterRuntime::Entered { .. }, true) => { // We are on a thread pool runtime thread, so we just need to // set up blocking. had_entered = true; } ( context::EnterRuntime::Entered { allow_block_in_place, }, false, ) => { // We are on an executor, but _not_ on the thread pool. That is // _only_ okay if we are in a thread pool runtime's block_on // method: if allow_block_in_place { had_entered = true; return Ok(()); } else { // This probably means we are on the current_thread runtime or in a // LocalSet, where it is _not_ okay to block. return Err( "can call blocking only when running on the multi-threaded runtime", ); } } (context::EnterRuntime::NotEntered, true) => { // This is a nested call to block_in_place (we already exited). // All the necessary setup has already been done. return Ok(()); } (context::EnterRuntime::NotEntered, false) => { // We are outside of the tokio runtime, so blocking is fine. // We can also skip all of the thread pool blocking setup steps. return Ok(()); } } let cx = maybe_cx.expect("no .is_some() == false cases above should lead here"); // Get the worker core. If none is set, then blocking is fine! let core = match cx.core.borrow_mut().take() { Some(core) => core, None => return Ok(()), }; // In order to block, the core must be sent to another thread for // execution. // // First, move the core back into the worker's shared core slot. cx.handoff_core.set(core); // Next, clone the worker handle and send it to a new thread for // processing. // // Once the blocking task is done executing, we will attempt to // steal the core back. let index = cx.index; let handle = cx.handle.clone(); let handoff_core = cx.handoff_core.clone(); runtime::spawn_blocking(move || run(index, handle, handoff_core, true)); Ok(()) }); if let Err(panic_message) = setup_result { panic!("{}", panic_message); } if had_entered { // Unset the current task's budget. Blocking sections are not // constrained by task budgets. let _reset = Reset(coop::stop()); crate::runtime::context::exit_runtime(f) } else { f() } } fn run( index: usize, handle: Arc, handoff_core: Arc>, blocking_in_place: bool, ) { struct AbortOnPanic; impl Drop for AbortOnPanic { fn drop(&mut self) { if std::thread::panicking() { eprintln!("worker thread panicking; aborting process"); std::process::abort(); } } } // Catching panics on worker threads in tests is quite tricky. Instead, when // debug assertions are enabled, we just abort the process. #[cfg(debug_assertions)] let _abort_on_panic = AbortOnPanic; let num_workers = handle.shared.condvars.len(); let mut worker = Worker { tick: 0, num_seq_local_queue_polls: 0, global_queue_interval: Stats::DEFAULT_GLOBAL_QUEUE_INTERVAL, is_shutdown: false, is_traced: false, workers_to_notify: Vec::with_capacity(num_workers - 1), idle_snapshot: idle::Snapshot::new(&handle.shared.idle), stats: stats::Ephemeral::new(), }; let sched_handle = scheduler::Handle::MultiThreadAlt(handle.clone()); crate::runtime::context::enter_runtime(&sched_handle, true, |_| { // Set the worker context. let cx = scheduler::Context::MultiThreadAlt(Context { index, lifo_enabled: Cell::new(!handle.shared.config.disable_lifo_slot), handle, core: RefCell::new(None), handoff_core, defer: RefCell::new(Vec::with_capacity(64)), }); context::set_scheduler(&cx, || { let cx = cx.expect_multi_thread_alt(); // Run the worker let res = worker.run(&cx, blocking_in_place); // `err` here signifies the core was lost, this is an expected end // state for a worker. debug_assert!(res.is_err()); // Check if there are any deferred tasks to notify. This can happen when // the worker core is lost due to `block_in_place()` being called from // within the task. if !cx.defer.borrow().is_empty() { worker.schedule_deferred_without_core(&cx, &mut cx.shared().synced.lock()); } }); }); } macro_rules! try_task { ($e:expr) => {{ let (task, core) = $e?; if task.is_some() { return Ok((task, core)); } core }}; } macro_rules! try_task_new_batch { ($w:expr, $e:expr) => {{ let (task, mut core) = $e?; if task.is_some() { core.stats.start_processing_scheduled_tasks(&mut $w.stats); return Ok((task, core)); } core }}; } impl Worker { fn run(&mut self, cx: &Context, blocking_in_place: bool) -> RunResult { let (maybe_task, mut core) = { if blocking_in_place { if let Some(core) = cx.handoff_core.take() { (None, core) } else { // Just shutdown return Err(()); } } else { let mut synced = cx.shared().synced.lock(); // First try to acquire an available core if let Some(core) = self.try_acquire_available_core(cx, &mut synced) { // Try to poll a task from the global queue let maybe_task = cx.shared().next_remote_task_synced(&mut synced); (maybe_task, core) } else { // block the thread to wait for a core to be assigned to us self.wait_for_core(cx, synced)? } } }; cx.shared().worker_metrics[core.index].set_thread_id(thread::current().id()); core.stats.start_processing_scheduled_tasks(&mut self.stats); if let Some(task) = maybe_task { core = self.run_task(cx, core, task)?; } while !self.is_shutdown { let (maybe_task, c) = self.next_task(cx, core)?; core = c; if let Some(task) = maybe_task { core = self.run_task(cx, core, task)?; } else { // The only reason to get `None` from `next_task` is we have // entered the shutdown phase. assert!(self.is_shutdown); break; } } cx.shared().shutdown_core(&cx.handle, core); // It is possible that tasks wake others during drop, so we need to // clear the defer list. self.shutdown_clear_defer(cx); Err(()) } // Try to acquire an available core, but do not block the thread fn try_acquire_available_core( &mut self, cx: &Context, synced: &mut Synced, ) -> Option> { if let Some(mut core) = cx .shared() .idle .try_acquire_available_core(&mut synced.idle) { self.reset_acquired_core(cx, synced, &mut core); Some(core) } else { None } } // Block the current thread, waiting for an available core fn wait_for_core( &mut self, cx: &Context, mut synced: MutexGuard<'_, Synced>, ) -> NextTaskResult { if cx.shared().idle.needs_searching() { if let Some(mut core) = self.try_acquire_available_core(cx, &mut synced) { cx.shared().idle.transition_worker_to_searching(&mut core); return Ok((None, core)); } } cx.shared() .idle .transition_worker_to_parked(&mut synced, cx.index); // Wait until a core is available, then exit the loop. let mut core = loop { if let Some(core) = synced.assigned_cores[cx.index].take() { break core; } // If shutting down, abort if cx.shared().inject.is_closed(&synced.inject) { self.shutdown_clear_defer(cx); return Err(()); } synced = cx.shared().condvars[cx.index].wait(synced).unwrap(); }; self.reset_acquired_core(cx, &mut synced, &mut core); if self.is_shutdown { // Currently shutting down, don't do any more work return Ok((None, core)); } let n = cmp::max(core.run_queue.remaining_slots() / 2, 1); let maybe_task = self.next_remote_task_batch_synced(cx, &mut synced, &mut core, n); core.stats.unparked(); self.flush_metrics(cx, &mut core); Ok((maybe_task, core)) } /// Ensure core's state is set correctly for the worker to start using. fn reset_acquired_core(&mut self, cx: &Context, synced: &mut Synced, core: &mut Core) { self.global_queue_interval = core.stats.tuned_global_queue_interval(&cx.shared().config); // Reset `lifo_enabled` here in case the core was previously stolen from // a task that had the LIFO slot disabled. self.reset_lifo_enabled(cx); // At this point, the local queue should be empty #[cfg(not(loom))] debug_assert!(core.run_queue.is_empty()); // Update shutdown state while locked self.update_global_flags(cx, synced); } /// Finds the next task to run, this could be from a queue or stealing. If /// none are available, the thread sleeps and tries again. fn next_task(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { self.assert_lifo_enabled_is_correct(cx); if self.is_traced { core = cx.handle.trace_core(core); } // Increment the tick self.tick = self.tick.wrapping_add(1); // Runs maintenance every so often. When maintenance is run, the // driver is checked, which may result in a task being found. core = try_task!(self.maybe_maintenance(&cx, core)); // Check the LIFO slot, local run queue, and the injection queue for // a notified task. core = try_task!(self.next_notified_task(cx, core)); // We consumed all work in the queues and will start searching for work. core.stats.end_processing_scheduled_tasks(&mut self.stats); super::counters::inc_num_no_local_work(); if !cx.defer.borrow().is_empty() { // We are deferring tasks, so poll the resource driver and schedule // the deferred tasks. try_task_new_batch!(self, self.park_yield(cx, core)); panic!("what happened to the deferred tasks? 🤔"); } while !self.is_shutdown { // Search for more work, this involves trying to poll the resource // driver, steal from other workers, and check the global queue // again. core = try_task_new_batch!(self, self.search_for_work(cx, core)); debug_assert!(cx.defer.borrow().is_empty()); core = try_task_new_batch!(self, self.park(cx, core)); } // Shutting down, drop any deferred tasks self.shutdown_clear_defer(cx); Ok((None, core)) } fn next_notified_task(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { self.num_seq_local_queue_polls += 1; if self.num_seq_local_queue_polls % self.global_queue_interval == 0 { super::counters::inc_global_queue_interval(); self.num_seq_local_queue_polls = 0; // Update the global queue interval, if needed self.tune_global_queue_interval(cx, &mut core); if let Some(task) = self.next_remote_task(cx) { return Ok((Some(task), core)); } } if let Some(task) = core.next_local_task() { return Ok((Some(task), core)); } self.next_remote_task_batch(cx, core) } fn next_remote_task(&self, cx: &Context) -> Option { if cx.shared().inject.is_empty() { return None; } let mut synced = cx.shared().synced.lock(); cx.shared().next_remote_task_synced(&mut synced) } fn next_remote_task_batch(&self, cx: &Context, mut core: Box) -> NextTaskResult { if cx.shared().inject.is_empty() { return Ok((None, core)); } // Other threads can only **remove** tasks from the current worker's // `run_queue`. So, we can be confident that by the time we call // `run_queue.push_back` below, there will be *at least* `cap` // available slots in the queue. let cap = usize::min( core.run_queue.remaining_slots(), usize::max(core.run_queue.max_capacity() / 2, 1), ); let mut synced = cx.shared().synced.lock(); let maybe_task = self.next_remote_task_batch_synced(cx, &mut synced, &mut core, cap); Ok((maybe_task, core)) } fn next_remote_task_batch_synced( &self, cx: &Context, synced: &mut Synced, core: &mut Core, max: usize, ) -> Option { super::counters::inc_num_remote_batch(); // The worker is currently idle, pull a batch of work from the // injection queue. We don't want to pull *all* the work so other // workers can also get some. let n = if core.is_searching { cx.shared().inject.len() / cx.shared().idle.num_searching() + 1 } else { cx.shared().inject.len() / cx.shared().remotes.len() + 1 }; let n = usize::min(n, max) + 1; // safety: passing in the correct `inject::Synced`. let mut tasks = unsafe { cx.shared().inject.pop_n(&mut synced.inject, n) }; // Pop the first task to return immediately let ret = tasks.next(); // Push the rest of the on the run queue core.run_queue.push_back(tasks); ret } /// Function responsible for stealing tasks from another worker /// /// Note: Only if less than half the workers are searching for tasks to steal /// a new worker will actually try to steal. The idea is to make sure not all /// workers will be trying to steal at the same time. fn search_for_work(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { #[cfg(not(loom))] const ROUNDS: usize = 4; #[cfg(loom)] const ROUNDS: usize = 1; debug_assert!(core.lifo_slot.is_none()); #[cfg(not(loom))] debug_assert!(core.run_queue.is_empty()); if !core.run_queue.can_steal() { return Ok((None, core)); } if !self.transition_to_searching(cx, &mut core) { return Ok((None, core)); } // core = try_task!(self, self.poll_driver(cx, core)); // Get a snapshot of which workers are idle cx.shared().idle.snapshot(&mut self.idle_snapshot); let num = cx.shared().remotes.len(); for i in 0..ROUNDS { // Start from a random worker let start = core.rand.fastrand_n(num as u32) as usize; if let Some(task) = self.steal_one_round(cx, &mut core, start) { return Ok((Some(task), core)); } core = try_task!(self.next_remote_task_batch(cx, core)); if i > 0 { super::counters::inc_num_spin_stall(); std::thread::sleep(std::time::Duration::from_micros(i as u64)); } } Ok((None, core)) } fn steal_one_round(&self, cx: &Context, core: &mut Core, start: usize) -> Option { let num = cx.shared().remotes.len(); for i in 0..num { let i = (start + i) % num; // Don't steal from ourself! We know we don't have work. if i == core.index { continue; } // If the core is currently idle, then there is nothing to steal. if self.idle_snapshot.is_idle(i) { continue; } let target = &cx.shared().remotes[i]; if let Some(task) = target .steal .steal_into(&mut core.run_queue, &mut core.stats) { return Some(task); } } None } fn run_task(&mut self, cx: &Context, mut core: Box, task: Notified) -> RunResult { let task = cx.shared().owned.assert_owner(task); // Make sure the worker is not in the **searching** state. This enables // another idle worker to try to steal work. if self.transition_from_searching(cx, &mut core) { super::counters::inc_num_relay_search(); cx.shared().notify_parked_local(); } self.assert_lifo_enabled_is_correct(cx); // Measure the poll start time. Note that we may end up polling other // tasks under this measurement. In this case, the tasks came from the // LIFO slot and are considered part of the current task for scheduling // purposes. These tasks inherent the "parent"'s limits. core.stats.start_poll(&mut self.stats); // Make the core available to the runtime context *cx.core.borrow_mut() = Some(core); // Run the task coop::budget(|| { super::counters::inc_num_polls(); task.run(); let mut lifo_polls = 0; // As long as there is budget remaining and a task exists in the // `lifo_slot`, then keep running. loop { // Check if we still have the core. If not, the core was stolen // by another worker. let mut core = match cx.core.borrow_mut().take() { Some(core) => core, None => { // In this case, we cannot call `reset_lifo_enabled()` // because the core was stolen. The stealer will handle // that at the top of `Context::run` return Err(()); } }; // Check for a task in the LIFO slot let task = match core.next_lifo_task() { Some(task) => task, None => { self.reset_lifo_enabled(cx); core.stats.end_poll(); return Ok(core); } }; if !coop::has_budget_remaining() { core.stats.end_poll(); // Not enough budget left to run the LIFO task, push it to // the back of the queue and return. core.run_queue .push_back_or_overflow(task, cx.shared(), &mut core.stats); // If we hit this point, the LIFO slot should be enabled. // There is no need to reset it. debug_assert!(cx.lifo_enabled.get()); return Ok(core); } // Track that we are about to run a task from the LIFO slot. lifo_polls += 1; super::counters::inc_lifo_schedules(); // Disable the LIFO slot if we reach our limit // // In ping-ping style workloads where task A notifies task B, // which notifies task A again, continuously prioritizing the // LIFO slot can cause starvation as these two tasks will // repeatedly schedule the other. To mitigate this, we limit the // number of times the LIFO slot is prioritized. if lifo_polls >= MAX_LIFO_POLLS_PER_TICK { cx.lifo_enabled.set(false); super::counters::inc_lifo_capped(); } // Run the LIFO task, then loop *cx.core.borrow_mut() = Some(core); let task = cx.shared().owned.assert_owner(task); super::counters::inc_num_lifo_polls(); task.run(); } }) } fn schedule_deferred_with_core<'a>( &mut self, cx: &'a Context, mut core: Box, synced: impl FnOnce() -> MutexGuard<'a, Synced>, ) -> NextTaskResult { let mut defer = cx.defer.borrow_mut(); // Grab a task to run next let task = defer.pop(); if task.is_none() { return Ok((None, core)); } if !defer.is_empty() { let mut synced = synced(); // Number of tasks we want to try to spread across idle workers let num_fanout = cmp::min(defer.len(), cx.shared().idle.num_idle(&synced.idle)); // Cap the number of threads woken up at one time. This is to limit // the number of no-op wakes and reduce mutext contention. // // This number was picked after some basic benchmarks, but it can // probably be tuned using the mean poll time value (slower task // polls can leverage more woken workers). let num_fanout = cmp::min(2, num_fanout); if num_fanout > 0 { cx.shared() .push_remote_task_batch_synced(&mut synced, defer.drain(..num_fanout)); cx.shared() .idle .notify_mult(&mut synced, &mut self.workers_to_notify, num_fanout); } // Do not run the task while holding the lock... drop(synced); } // Notify any workers for worker in self.workers_to_notify.drain(..) { cx.shared().condvars[worker].notify_one() } if !defer.is_empty() { // Push the rest of the tasks on the local queue for task in defer.drain(..) { core.run_queue .push_back_or_overflow(task, cx.shared(), &mut core.stats); } cx.shared().notify_parked_local(); } Ok((task, core)) } fn schedule_deferred_without_core<'a>(&mut self, cx: &Context, synced: &mut Synced) { let mut defer = cx.defer.borrow_mut(); let num = defer.len(); if num > 0 { // Push all tasks to the injection queue cx.shared() .push_remote_task_batch_synced(synced, defer.drain(..)); debug_assert!(self.workers_to_notify.is_empty()); // Notify workers cx.shared() .idle .notify_mult(synced, &mut self.workers_to_notify, num); // Notify any workers for worker in self.workers_to_notify.drain(..) { cx.shared().condvars[worker].notify_one() } } } fn maybe_maintenance(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { if self.tick % cx.shared().config.event_interval == 0 { super::counters::inc_num_maintenance(); core.stats.end_processing_scheduled_tasks(&mut self.stats); // Run regularly scheduled maintenance core = try_task_new_batch!(self, self.park_yield(cx, core)); core.stats.start_processing_scheduled_tasks(&mut self.stats); } Ok((None, core)) } fn flush_metrics(&self, cx: &Context, core: &mut Core) { core.stats.submit(&cx.shared().worker_metrics[core.index]); } fn update_global_flags(&mut self, cx: &Context, synced: &mut Synced) { if !self.is_shutdown { self.is_shutdown = cx.shared().inject.is_closed(&synced.inject); } if !self.is_traced { self.is_traced = cx.shared().trace_status.trace_requested(); } } fn park_yield(&mut self, cx: &Context, core: Box) -> NextTaskResult { // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... // to run without actually putting the thread to sleep. if let Some(mut driver) = cx.shared().driver.take() { driver.park_timeout(&cx.handle.driver, Duration::from_millis(0)); cx.shared().driver.set(driver); } // If there are more I/O events, schedule them. let (maybe_task, mut core) = self.schedule_deferred_with_core(cx, core, || cx.shared().synced.lock())?; self.flush_metrics(cx, &mut core); self.update_global_flags(cx, &mut cx.shared().synced.lock()); Ok((maybe_task, core)) } /* fn poll_driver(&mut self, cx: &Context, core: Box) -> NextTaskResult { // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... // to run without actually putting the thread to sleep. if let Some(mut driver) = cx.shared().driver.take() { driver.park_timeout(&cx.handle.driver, Duration::from_millis(0)); cx.shared().driver.set(driver); // If there are more I/O events, schedule them. self.schedule_deferred_with_core(cx, core, || cx.shared().synced.lock()) } else { Ok((None, core)) } } */ fn park(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { if let Some(f) = &cx.shared().config.before_park { f(); } if self.can_transition_to_parked(&mut core) { debug_assert!(!self.is_shutdown); debug_assert!(!self.is_traced); core = try_task!(self.do_park(cx, core)); } if let Some(f) = &cx.shared().config.after_unpark { f(); } Ok((None, core)) } fn do_park(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { let was_searching = core.is_searching; // Acquire the lock let mut synced = cx.shared().synced.lock(); // The local queue should be empty at this point #[cfg(not(loom))] debug_assert!(core.run_queue.is_empty()); // Try one last time to get tasks let n = cmp::max(core.run_queue.remaining_slots() / 2, 1); if let Some(task) = self.next_remote_task_batch_synced(cx, &mut synced, &mut core, n) { return Ok((Some(task), core)); } if !was_searching { if cx .shared() .idle .transition_worker_to_searching_if_needed(&mut synced.idle, &mut core) { // Skip parking, go back to searching return Ok((None, core)); } } super::counters::inc_num_parks(); core.stats.about_to_park(); // Flush metrics to the runtime metrics aggregator self.flush_metrics(cx, &mut core); // If the runtime is shutdown, skip parking self.update_global_flags(cx, &mut synced); if self.is_shutdown { return Ok((None, core)); } // Release the core core.is_searching = false; cx.shared().idle.release_core(&mut synced, core); drop(synced); if was_searching { if cx.shared().idle.transition_worker_from_searching() { // cx.shared().idle.snapshot(&mut self.idle_snapshot); // We were the last searching worker, we need to do one last check for i in 0..cx.shared().remotes.len() { if !cx.shared().remotes[i].steal.is_empty() { let mut synced = cx.shared().synced.lock(); // Try to get a core if let Some(mut core) = self.try_acquire_available_core(cx, &mut synced) { cx.shared().idle.transition_worker_to_searching(&mut core); return Ok((None, core)); } else { // Fall back to the park routine break; } } } } } if let Some(mut driver) = cx.shared().take_driver() { // Wait for driver events driver.park(&cx.handle.driver); synced = cx.shared().synced.lock(); if cx.shared().inject.is_closed(&mut synced.inject) { synced.shutdown_driver = Some(driver); self.shutdown_clear_defer(cx); cx.shared().shutdown_finalize(&cx.handle, &mut synced); return Err(()); } // Put the driver back cx.shared().driver.set(driver); // Try to acquire an available core to schedule I/O events if let Some(core) = self.try_acquire_available_core(cx, &mut synced) { // This may result in a task being run self.schedule_deferred_with_core(cx, core, move || synced) } else { // Schedule any deferred tasks self.schedule_deferred_without_core(cx, &mut synced); // Wait for a core. self.wait_for_core(cx, synced) } } else { synced = cx.shared().synced.lock(); // Wait for a core to be assigned to us self.wait_for_core(cx, synced) } } fn transition_to_searching(&self, cx: &Context, core: &mut Core) -> bool { if !core.is_searching { cx.shared().idle.try_transition_worker_to_searching(core); } core.is_searching } /// Returns `true` if another worker must be notified fn transition_from_searching(&self, cx: &Context, core: &mut Core) -> bool { if !core.is_searching { return false; } core.is_searching = false; cx.shared().idle.transition_worker_from_searching() } fn can_transition_to_parked(&self, core: &mut Core) -> bool { !self.has_tasks(core) && !self.is_shutdown && !self.is_traced } fn has_tasks(&self, core: &Core) -> bool { core.lifo_slot.is_some() || !core.run_queue.is_empty() } fn reset_lifo_enabled(&self, cx: &Context) { cx.lifo_enabled .set(!cx.handle.shared.config.disable_lifo_slot); } fn assert_lifo_enabled_is_correct(&self, cx: &Context) { debug_assert_eq!( cx.lifo_enabled.get(), !cx.handle.shared.config.disable_lifo_slot ); } fn tune_global_queue_interval(&mut self, cx: &Context, core: &mut Core) { let next = core.stats.tuned_global_queue_interval(&cx.shared().config); // Smooth out jitter if u32::abs_diff(self.global_queue_interval, next) > 2 { self.global_queue_interval = next; } } fn shutdown_clear_defer(&self, cx: &Context) { let mut defer = cx.defer.borrow_mut(); for task in defer.drain(..) { drop(task); } } } impl Context { pub(crate) fn defer(&self, waker: &Waker) { // TODO: refactor defer across all runtimes waker.wake_by_ref(); } fn shared(&self) -> &Shared { &self.handle.shared } #[cfg_attr(not(feature = "time"), allow(dead_code))] pub(crate) fn get_worker_index(&self) -> usize { self.index } } impl Core { fn next_local_task(&mut self) -> Option { self.next_lifo_task().or_else(|| self.run_queue.pop()) } fn next_lifo_task(&mut self) -> Option { self.lifo_slot.take() } } impl Shared { fn next_remote_task_synced(&self, synced: &mut Synced) -> Option { // safety: we only have access to a valid `Synced` in this file. unsafe { self.inject.pop(&mut synced.inject) } } pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { use std::ptr; with_current(|maybe_cx| { if let Some(cx) = maybe_cx { // Make sure the task is part of the **current** scheduler. if ptr::eq(self, &cx.handle.shared) { // And the current thread still holds a core if let Some(core) = cx.core.borrow_mut().as_mut() { if is_yield { cx.defer.borrow_mut().push(task); } else { self.schedule_local(cx, core, task); } } else { // This can happen if either the core was stolen // (`block_in_place`) or the notification happens from // the driver. cx.defer.borrow_mut().push(task); } return; } } // Otherwise, use the inject queue. self.schedule_remote(task); }) } fn schedule_local(&self, cx: &Context, core: &mut Core, task: Notified) { core.stats.inc_local_schedule_count(); if cx.lifo_enabled.get() { // Push to the LIFO slot let prev = std::mem::replace(&mut core.lifo_slot, Some(task)); // let prev = cx.shared().remotes[core.index].lifo_slot.swap_local(task); if let Some(prev) = prev { core.run_queue .push_back_or_overflow(prev, self, &mut core.stats); } else { return; } } else { core.run_queue .push_back_or_overflow(task, self, &mut core.stats); } self.notify_parked_local(); } fn notify_parked_local(&self) { super::counters::inc_num_inc_notify_local(); self.idle.notify_local(self); } fn schedule_remote(&self, task: Notified) { super::counters::inc_num_notify_remote(); self.scheduler_metrics.inc_remote_schedule_count(); let mut synced = self.synced.lock(); // Push the task in the self.push_remote_task(&mut synced, task); // Notify a worker. The mutex is passed in and will be released as part // of the method call. self.idle.notify_remote(synced, self); } pub(super) fn close(&self, handle: &Handle) { { let mut synced = self.synced.lock(); if let Some(driver) = self.driver.take() { synced.shutdown_driver = Some(driver); } if !self.inject.close(&mut synced.inject) { return; } // Set the shutdown flag on all available cores self.idle.shutdown(&mut synced, self); } // Any unassigned cores need to be shutdown, but we have to first drop // the lock self.idle.shutdown_unassigned_cores(handle, self); } fn push_remote_task(&self, synced: &mut Synced, task: Notified) { // safety: passing in correct `idle::Synced` unsafe { self.inject.push(&mut synced.inject, task); } } fn push_remote_task_batch(&self, iter: I) where I: Iterator>>, { unsafe { self.inject.push_batch(self, iter); } } fn push_remote_task_batch_synced(&self, synced: &mut Synced, iter: I) where I: Iterator>>, { unsafe { self.inject.push_batch(&mut synced.inject, iter); } } fn take_driver(&self) -> Option> { if !self.driver_enabled() { return None; } self.driver.take() } fn driver_enabled(&self) -> bool { self.condvars.len() > self.remotes.len() } pub(super) fn shutdown_core(&self, handle: &Handle, mut core: Box) { // Start from a random inner list let start = core.rand.fastrand_n(self.owned.get_shard_size() as u32); self.owned.close_and_shutdown_all(start as usize); core.stats.submit(&self.worker_metrics[core.index]); let mut synced = self.synced.lock(); synced.shutdown_cores.push(core); self.shutdown_finalize(handle, &mut synced); } pub(super) fn shutdown_finalize(&self, handle: &Handle, synced: &mut Synced) { // Wait for all cores if synced.shutdown_cores.len() != self.remotes.len() { return; } let driver = synced.shutdown_driver.take(); if self.driver_enabled() && driver.is_none() { return; } debug_assert!(self.owned.is_empty()); for mut core in synced.shutdown_cores.drain(..) { // Drain tasks from the local queue while core.next_local_task().is_some() {} } // Shutdown the driver if let Some(mut driver) = driver { driver.shutdown(&handle.driver); } // Drain the injection queue // // We already shut down every task, so we can simply drop the tasks. We // cannot call `next_remote_task()` because we already hold the lock. // // safety: passing in correct `idle::Synced` while let Some(task) = self.next_remote_task_synced(synced) { drop(task); } } } impl Overflow> for Shared { fn push(&self, task: task::Notified>) { self.push_remote_task(&mut self.synced.lock(), task); } fn push_batch(&self, iter: I) where I: Iterator>>, { self.push_remote_task_batch(iter) } } impl<'a> Lock for &'a Shared { type Handle = SyncedGuard<'a>; fn lock(self) -> Self::Handle { SyncedGuard { lock: self.synced.lock(), } } } impl<'a> Lock for &'a Shared { type Handle = SyncedGuard<'a>; fn lock(self) -> Self::Handle { SyncedGuard { lock: self.synced.lock(), } } } impl task::Schedule for Arc { fn release(&self, task: &Task) -> Option { self.shared.owned.remove(task) } fn schedule(&self, task: Notified) { self.shared.schedule_task(task, false); } fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: self.task_hooks.task_terminate_callback.clone(), } } fn yield_now(&self, task: Notified) { self.shared.schedule_task(task, true); } } impl AsMut for Synced { fn as_mut(&mut self) -> &mut Synced { self } } pub(crate) struct SyncedGuard<'a> { lock: crate::loom::sync::MutexGuard<'a, Synced>, } impl<'a> AsMut for SyncedGuard<'a> { fn as_mut(&mut self) -> &mut inject::Synced { &mut self.lock.inject } } impl<'a> AsMut for SyncedGuard<'a> { fn as_mut(&mut self) -> &mut Synced { &mut self.lock } } #[track_caller] fn with_current(f: impl FnOnce(Option<&Context>) -> R) -> R { use scheduler::Context::MultiThreadAlt; context::with_scheduler(|ctx| match ctx { Some(MultiThreadAlt(ctx)) => f(Some(ctx)), _ => f(None), }) } tokio-1.43.1/src/runtime/signal/mod.rs000064400000000000000000000116721046102023000157220ustar 00000000000000#![cfg_attr(not(feature = "rt"), allow(dead_code))] //! Signal driver use crate::runtime::{driver, io}; use crate::signal::registry::globals; use mio::net::UnixStream; use std::io::{self as std_io, Read}; use std::sync::{Arc, Weak}; use std::time::Duration; /// Responsible for registering wakeups when an OS signal is received, and /// subsequently dispatching notifications to any signal listeners as appropriate. /// /// Note: this driver relies on having an enabled IO driver in order to listen to /// pipe write wakeups. #[derive(Debug)] pub(crate) struct Driver { /// Thread parker. The `Driver` park implementation delegates to this. io: io::Driver, /// A pipe for receiving wake events from the signal handler receiver: UnixStream, /// Shared state. The driver keeps a strong ref and the handle keeps a weak /// ref. The weak ref is used to check if the driver is still active before /// trying to register a signal handler. inner: Arc<()>, } #[derive(Debug, Default)] pub(crate) struct Handle { /// Paired w/ the `Arc` above and is used to check if the driver is still /// around before attempting to register a signal handler. inner: Weak<()>, } // ===== impl Driver ===== impl Driver { /// Creates a new signal `Driver` instance that delegates wakeups to `park`. pub(crate) fn new(io: io::Driver, io_handle: &io::Handle) -> std_io::Result { use std::mem::ManuallyDrop; use std::os::unix::io::{AsRawFd, FromRawFd}; // NB: We give each driver a "fresh" receiver file descriptor to avoid // the issues described in alexcrichton/tokio-process#42. // // In the past we would reuse the actual receiver file descriptor and // swallow any errors around double registration of the same descriptor. // I'm not sure if the second (failed) registration simply doesn't end // up receiving wake up notifications, or there could be some race // condition when consuming readiness events, but having distinct // descriptors appears to mitigate this. // // Unfortunately we cannot just use a single global UnixStream instance // either, since we can't assume they will always be registered with the // exact same reactor. // // Mio 0.7 removed `try_clone()` as an API due to unexpected behavior // with registering dups with the same reactor. In this case, duping is // safe as each dup is registered with separate reactors **and** we // only expect at least one dup to receive the notification. // Manually drop as we don't actually own this instance of UnixStream. let receiver_fd = globals().receiver.as_raw_fd(); // safety: there is nothing unsafe about this, but the `from_raw_fd` fn is marked as unsafe. let original = ManuallyDrop::new(unsafe { std::os::unix::net::UnixStream::from_raw_fd(receiver_fd) }); let mut receiver = UnixStream::from_std(original.try_clone()?); io_handle.register_signal_receiver(&mut receiver)?; Ok(Self { io, receiver, inner: Arc::new(()), }) } /// Returns a handle to this event loop which can be sent across threads /// and can be used as a proxy to the event loop itself. pub(crate) fn handle(&self) -> Handle { Handle { inner: Arc::downgrade(&self.inner), } } pub(crate) fn park(&mut self, handle: &driver::Handle) { self.io.park(handle); self.process(); } pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { self.io.park_timeout(handle, duration); self.process(); } pub(crate) fn shutdown(&mut self, handle: &driver::Handle) { self.io.shutdown(handle); } fn process(&mut self) { // If the signal pipe has not received a readiness event, then there is // nothing else to do. if !self.io.consume_signal_ready() { return; } // Drain the pipe completely so we can receive a new readiness event // if another signal has come in. let mut buf = [0; 128]; #[allow(clippy::unused_io_amount)] loop { match self.receiver.read(&mut buf) { Ok(0) => panic!("EOF on self-pipe"), Ok(_) => continue, // Keep reading Err(e) if e.kind() == std_io::ErrorKind::WouldBlock => break, Err(e) => panic!("Bad read on self-pipe: {e}"), } } // Broadcast any signals which were received globals().broadcast(); } } // ===== impl Handle ===== impl Handle { pub(crate) fn check_inner(&self) -> std_io::Result<()> { if self.inner.strong_count() > 0 { Ok(()) } else { Err(std_io::Error::new( std_io::ErrorKind::Other, "signal driver gone", )) } } } tokio-1.43.1/src/runtime/task/abort.rs000064400000000000000000000075301046102023000157350ustar 00000000000000use crate::runtime::task::{Header, RawTask}; use std::fmt; use std::panic::{RefUnwindSafe, UnwindSafe}; /// An owned permission to abort a spawned task, without awaiting its completion. /// /// Unlike a [`JoinHandle`], an `AbortHandle` does *not* represent the /// permission to await the task's completion, only to terminate it. /// /// The task may be aborted by calling the [`AbortHandle::abort`] method. /// Dropping an `AbortHandle` releases the permission to terminate the task /// --- it does *not* abort the task. /// /// Be aware that tasks spawned using [`spawn_blocking`] cannot be aborted /// because they are not async. If you call `abort` on a `spawn_blocking` task, /// then this *will not have any effect*, and the task will continue running /// normally. The exception is if the task has not started running yet; in that /// case, calling `abort` may prevent the task from starting. /// /// [`JoinHandle`]: crate::task::JoinHandle /// [`spawn_blocking`]: crate::task::spawn_blocking #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub struct AbortHandle { raw: RawTask, } impl AbortHandle { pub(super) fn new(raw: RawTask) -> Self { Self { raw } } /// Abort the task associated with the handle. /// /// Awaiting a cancelled task might complete as usual if the task was /// already completed at the time it was cancelled, but most likely it /// will fail with a [cancelled] `JoinError`. /// /// If the task was already cancelled, such as by [`JoinHandle::abort`], /// this method will do nothing. /// /// Be aware that tasks spawned using [`spawn_blocking`] cannot be aborted /// because they are not async. If you call `abort` on a `spawn_blocking` /// task, then this *will not have any effect*, and the task will continue /// running normally. The exception is if the task has not started running /// yet; in that case, calling `abort` may prevent the task from starting. /// /// See also [the module level docs] for more information on cancellation. /// /// [cancelled]: method@super::error::JoinError::is_cancelled /// [`JoinHandle::abort`]: method@super::JoinHandle::abort /// [the module level docs]: crate::task#cancellation /// [`spawn_blocking`]: crate::task::spawn_blocking pub fn abort(&self) { self.raw.remote_abort(); } /// Checks if the task associated with this `AbortHandle` has finished. /// /// Please note that this method can return `false` even if `abort` has been /// called on the task. This is because the cancellation process may take /// some time, and this method does not return `true` until it has /// completed. pub fn is_finished(&self) -> bool { let state = self.raw.state().load(); state.is_complete() } /// Returns a [task ID] that uniquely identifies this task relative to other /// currently spawned tasks. /// /// [task ID]: crate::task::Id pub fn id(&self) -> super::Id { // Safety: The header pointer is valid. unsafe { Header::get_id(self.raw.header_ptr()) } } } unsafe impl Send for AbortHandle {} unsafe impl Sync for AbortHandle {} impl UnwindSafe for AbortHandle {} impl RefUnwindSafe for AbortHandle {} impl fmt::Debug for AbortHandle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { // Safety: The header pointer is valid. let id_ptr = unsafe { Header::get_id_ptr(self.raw.header_ptr()) }; let id = unsafe { id_ptr.as_ref() }; fmt.debug_struct("AbortHandle").field("id", id).finish() } } impl Drop for AbortHandle { fn drop(&mut self) { self.raw.drop_abort_handle(); } } impl Clone for AbortHandle { /// Returns a cloned `AbortHandle` that can be used to remotely abort this task. fn clone(&self) -> Self { self.raw.ref_inc(); Self::new(self.raw) } } tokio-1.43.1/src/runtime/task/core.rs000064400000000000000000000415141046102023000155560ustar 00000000000000//! Core task module. //! //! # Safety //! //! The functions in this module are private to the `task` module. All of them //! should be considered `unsafe` to use, but are not marked as such since it //! would be too noisy. //! //! Make sure to consult the relevant safety section of each function before //! use. use crate::future::Future; use crate::loom::cell::UnsafeCell; use crate::runtime::context; use crate::runtime::task::raw::{self, Vtable}; use crate::runtime::task::state::State; use crate::runtime::task::{Id, Schedule, TaskHarnessScheduleHooks}; use crate::util::linked_list; use std::num::NonZeroU64; use std::pin::Pin; use std::ptr::NonNull; use std::task::{Context, Poll, Waker}; /// The task cell. Contains the components of the task. /// /// It is critical for `Header` to be the first field as the task structure will /// be referenced by both *mut Cell and *mut Header. /// /// Any changes to the layout of this struct _must_ also be reflected in the /// `const` fns in raw.rs. /// // # This struct should be cache padded to avoid false sharing. The cache padding rules are copied // from crossbeam-utils/src/cache_padded.rs // // Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache // lines at a time, so we have to align to 128 bytes rather than 64. // // Sources: // - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf // - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107 // // ARM's big.LITTLE architecture has asymmetric cores and "big" cores have 128-byte cache line size. // // Sources: // - https://www.mono-project.com/news/2016/09/12/arm64-icache/ // // powerpc64 has 128-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_ppc64x.go#L9 #[cfg_attr( any( target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", ), repr(align(128)) )] // arm, mips, mips64, sparc, and hexagon have 32-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_arm.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mipsle.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips64x.go#L9 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L17 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/hexagon/include/asm/cache.h#L12 #[cfg_attr( any( target_arch = "arm", target_arch = "mips", target_arch = "mips64", target_arch = "sparc", target_arch = "hexagon", ), repr(align(32)) )] // m68k has 16-byte cache line size. // // Sources: // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/m68k/include/asm/cache.h#L9 #[cfg_attr(target_arch = "m68k", repr(align(16)))] // s390x has 256-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_s390x.go#L7 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/s390/include/asm/cache.h#L13 #[cfg_attr(target_arch = "s390x", repr(align(256)))] // x86, riscv, wasm, and sparc64 have 64-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/dda2991c2ea0c5914714469c4defc2562a907230/src/internal/cpu/cpu_x86.go#L9 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_wasm.go#L7 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L19 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/riscv/include/asm/cache.h#L10 // // All others are assumed to have 64-byte cache line size. #[cfg_attr( not(any( target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", target_arch = "arm", target_arch = "mips", target_arch = "mips64", target_arch = "sparc", target_arch = "hexagon", target_arch = "m68k", target_arch = "s390x", )), repr(align(64)) )] #[repr(C)] pub(super) struct Cell { /// Hot task state data pub(super) header: Header, /// Either the future or output, depending on the execution stage. pub(super) core: Core, /// Cold data pub(super) trailer: Trailer, } pub(super) struct CoreStage { stage: UnsafeCell>, } /// The core of the task. /// /// Holds the future or output, depending on the stage of execution. /// /// Any changes to the layout of this struct _must_ also be reflected in the /// `const` fns in raw.rs. #[repr(C)] pub(super) struct Core { /// Scheduler used to drive this future. pub(super) scheduler: S, /// The task's ID, used for populating `JoinError`s. pub(super) task_id: Id, /// Either the future or the output. pub(super) stage: CoreStage, } /// Crate public as this is also needed by the pool. #[repr(C)] pub(crate) struct Header { /// Task state. pub(super) state: State, /// Pointer to next task, used with the injection queue. pub(super) queue_next: UnsafeCell>>, /// Table of function pointers for executing actions on the task. pub(super) vtable: &'static Vtable, /// This integer contains the id of the `OwnedTasks` or `LocalOwnedTasks` /// that this task is stored in. If the task is not in any list, should be /// the id of the list that it was previously in, or `None` if it has never /// been in any list. /// /// Once a task has been bound to a list, it can never be bound to another /// list, even if removed from the first list. /// /// The id is not unset when removed from a list because we want to be able /// to read the id without synchronization, even if it is concurrently being /// removed from the list. pub(super) owner_id: UnsafeCell>, /// The tracing ID for this instrumented task. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) tracing_id: Option, } unsafe impl Send for Header {} unsafe impl Sync for Header {} /// Cold data is stored after the future. Data is considered cold if it is only /// used during creation or shutdown of the task. pub(super) struct Trailer { /// Pointers for the linked list in the `OwnedTasks` that owns this task. pub(super) owned: linked_list::Pointers
, /// Consumer task waiting on completion of this task. pub(super) waker: UnsafeCell>, /// Optional hooks needed in the harness. pub(super) hooks: TaskHarnessScheduleHooks, } generate_addr_of_methods! { impl<> Trailer { pub(super) unsafe fn addr_of_owned(self: NonNull) -> NonNull> { &self.owned } } } /// Either the future or the output. #[repr(C)] // https://github.com/rust-lang/miri/issues/3780 pub(super) enum Stage { Running(T), Finished(super::Result), Consumed, } impl Cell { /// Allocates a new task cell, containing the header, trailer, and core /// structures. pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box> { // Separated into a non-generic function to reduce LLVM codegen fn new_header( state: State, vtable: &'static Vtable, #[cfg(all(tokio_unstable, feature = "tracing"))] tracing_id: Option, ) -> Header { Header { state, queue_next: UnsafeCell::new(None), vtable, owner_id: UnsafeCell::new(None), #[cfg(all(tokio_unstable, feature = "tracing"))] tracing_id, } } #[cfg(all(tokio_unstable, feature = "tracing"))] let tracing_id = future.id(); let vtable = raw::vtable::(); let result = Box::new(Cell { trailer: Trailer::new(scheduler.hooks()), header: new_header( state, vtable, #[cfg(all(tokio_unstable, feature = "tracing"))] tracing_id, ), core: Core { scheduler, stage: CoreStage { stage: UnsafeCell::new(Stage::Running(future)), }, task_id, }, }); #[cfg(debug_assertions)] { // Using a separate function for this code avoids instantiating it separately for every `T`. unsafe fn check(header: &Header, trailer: &Trailer, scheduler: &S, task_id: &Id) { let trailer_addr = trailer as *const Trailer as usize; let trailer_ptr = unsafe { Header::get_trailer(NonNull::from(header)) }; assert_eq!(trailer_addr, trailer_ptr.as_ptr() as usize); let scheduler_addr = scheduler as *const S as usize; let scheduler_ptr = unsafe { Header::get_scheduler::(NonNull::from(header)) }; assert_eq!(scheduler_addr, scheduler_ptr.as_ptr() as usize); let id_addr = task_id as *const Id as usize; let id_ptr = unsafe { Header::get_id_ptr(NonNull::from(header)) }; assert_eq!(id_addr, id_ptr.as_ptr() as usize); } unsafe { check( &result.header, &result.trailer, &result.core.scheduler, &result.core.task_id, ); } } result } } impl CoreStage { pub(super) fn with_mut(&self, f: impl FnOnce(*mut Stage) -> R) -> R { self.stage.with_mut(f) } } /// Set and clear the task id in the context when the future is executed or /// dropped, or when the output produced by the future is dropped. pub(crate) struct TaskIdGuard { parent_task_id: Option, } impl TaskIdGuard { fn enter(id: Id) -> Self { TaskIdGuard { parent_task_id: context::set_current_task_id(Some(id)), } } } impl Drop for TaskIdGuard { fn drop(&mut self) { context::set_current_task_id(self.parent_task_id); } } impl Core { /// Polls the future. /// /// # Safety /// /// The caller must ensure it is safe to mutate the `state` field. This /// requires ensuring mutual exclusion between any concurrent thread that /// might modify the future or output field. /// /// The mutual exclusion is implemented by `Harness` and the `Lifecycle` /// component of the task state. /// /// `self` must also be pinned. This is handled by storing the task on the /// heap. pub(super) fn poll(&self, mut cx: Context<'_>) -> Poll { let res = { self.stage.stage.with_mut(|ptr| { // Safety: The caller ensures mutual exclusion to the field. let future = match unsafe { &mut *ptr } { Stage::Running(future) => future, _ => unreachable!("unexpected stage"), }; // Safety: The caller ensures the future is pinned. let future = unsafe { Pin::new_unchecked(future) }; let _guard = TaskIdGuard::enter(self.task_id); future.poll(&mut cx) }) }; if res.is_ready() { self.drop_future_or_output(); } res } /// Drops the future. /// /// # Safety /// /// The caller must ensure it is safe to mutate the `stage` field. pub(super) fn drop_future_or_output(&self) { // Safety: the caller ensures mutual exclusion to the field. unsafe { self.set_stage(Stage::Consumed); } } /// Stores the task output. /// /// # Safety /// /// The caller must ensure it is safe to mutate the `stage` field. pub(super) fn store_output(&self, output: super::Result) { // Safety: the caller ensures mutual exclusion to the field. unsafe { self.set_stage(Stage::Finished(output)); } } /// Takes the task output. /// /// # Safety /// /// The caller must ensure it is safe to mutate the `stage` field. pub(super) fn take_output(&self) -> super::Result { use std::mem; self.stage.stage.with_mut(|ptr| { // Safety:: the caller ensures mutual exclusion to the field. match mem::replace(unsafe { &mut *ptr }, Stage::Consumed) { Stage::Finished(output) => output, _ => panic!("JoinHandle polled after completion"), } }) } unsafe fn set_stage(&self, stage: Stage) { let _guard = TaskIdGuard::enter(self.task_id); self.stage.stage.with_mut(|ptr| *ptr = stage); } } impl Header { pub(super) unsafe fn set_next(&self, next: Option>) { self.queue_next.with_mut(|ptr| *ptr = next); } // safety: The caller must guarantee exclusive access to this field, and // must ensure that the id is either `None` or the id of the OwnedTasks // containing this task. pub(super) unsafe fn set_owner_id(&self, owner: NonZeroU64) { self.owner_id.with_mut(|ptr| *ptr = Some(owner)); } pub(super) fn get_owner_id(&self) -> Option { // safety: If there are concurrent writes, then that write has violated // the safety requirements on `set_owner_id`. unsafe { self.owner_id.with(|ptr| *ptr) } } /// Gets a pointer to the `Trailer` of the task containing this `Header`. /// /// # Safety /// /// The provided raw pointer must point at the header of a task. pub(super) unsafe fn get_trailer(me: NonNull
) -> NonNull { let offset = me.as_ref().vtable.trailer_offset; let trailer = me.as_ptr().cast::().add(offset).cast::(); NonNull::new_unchecked(trailer) } /// Gets a pointer to the scheduler of the task containing this `Header`. /// /// # Safety /// /// The provided raw pointer must point at the header of a task. /// /// The generic type S must be set to the correct scheduler type for this /// task. pub(super) unsafe fn get_scheduler(me: NonNull
) -> NonNull { let offset = me.as_ref().vtable.scheduler_offset; let scheduler = me.as_ptr().cast::().add(offset).cast::(); NonNull::new_unchecked(scheduler) } /// Gets a pointer to the id of the task containing this `Header`. /// /// # Safety /// /// The provided raw pointer must point at the header of a task. pub(super) unsafe fn get_id_ptr(me: NonNull
) -> NonNull { let offset = me.as_ref().vtable.id_offset; let id = me.as_ptr().cast::().add(offset).cast::(); NonNull::new_unchecked(id) } /// Gets the id of the task containing this `Header`. /// /// # Safety /// /// The provided raw pointer must point at the header of a task. pub(super) unsafe fn get_id(me: NonNull
) -> Id { let ptr = Header::get_id_ptr(me).as_ptr(); *ptr } /// Gets the tracing id of the task containing this `Header`. /// /// # Safety /// /// The provided raw pointer must point at the header of a task. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) unsafe fn get_tracing_id(me: &NonNull
) -> Option<&tracing::Id> { me.as_ref().tracing_id.as_ref() } } impl Trailer { fn new(hooks: TaskHarnessScheduleHooks) -> Self { Trailer { waker: UnsafeCell::new(None), owned: linked_list::Pointers::new(), hooks, } } pub(super) unsafe fn set_waker(&self, waker: Option) { self.waker.with_mut(|ptr| { *ptr = waker; }); } pub(super) unsafe fn will_wake(&self, waker: &Waker) -> bool { self.waker .with(|ptr| (*ptr).as_ref().unwrap().will_wake(waker)) } pub(super) fn wake_join(&self) { self.waker.with(|ptr| match unsafe { &*ptr } { Some(waker) => waker.wake_by_ref(), None => panic!("waker missing"), }); } } #[test] #[cfg(not(loom))] fn header_lte_cache_line() { assert!(std::mem::size_of::
() <= 8 * std::mem::size_of::<*const ()>()); } tokio-1.43.1/src/runtime/task/error.rs000064400000000000000000000125611046102023000157570ustar 00000000000000use std::any::Any; use std::fmt; use std::io; use super::Id; use crate::util::SyncWrapper; cfg_rt! { /// Task failed to execute to completion. pub struct JoinError { repr: Repr, id: Id, } } enum Repr { Cancelled, Panic(SyncWrapper>), } impl JoinError { pub(crate) fn cancelled(id: Id) -> JoinError { JoinError { repr: Repr::Cancelled, id, } } pub(crate) fn panic(id: Id, err: Box) -> JoinError { JoinError { repr: Repr::Panic(SyncWrapper::new(err)), id, } } /// Returns true if the error was caused by the task being cancelled. /// /// See [the module level docs] for more information on cancellation. /// /// [the module level docs]: crate::task#cancellation pub fn is_cancelled(&self) -> bool { matches!(&self.repr, Repr::Cancelled) } /// Returns true if the error was caused by the task panicking. /// /// # Examples /// /// ``` /// use std::panic; /// /// #[tokio::main] /// async fn main() { /// let err = tokio::spawn(async { /// panic!("boom"); /// }).await.unwrap_err(); /// /// assert!(err.is_panic()); /// } /// ``` pub fn is_panic(&self) -> bool { matches!(&self.repr, Repr::Panic(_)) } /// Consumes the join error, returning the object with which the task panicked. /// /// # Panics /// /// `into_panic()` panics if the `Error` does not represent the underlying /// task terminating with a panic. Use `is_panic` to check the error reason /// or `try_into_panic` for a variant that does not panic. /// /// # Examples /// /// ```should_panic /// use std::panic; /// /// #[tokio::main] /// async fn main() { /// let err = tokio::spawn(async { /// panic!("boom"); /// }).await.unwrap_err(); /// /// if err.is_panic() { /// // Resume the panic on the main task /// panic::resume_unwind(err.into_panic()); /// } /// } /// ``` #[track_caller] pub fn into_panic(self) -> Box { self.try_into_panic() .expect("`JoinError` reason is not a panic.") } /// Consumes the join error, returning the object with which the task /// panicked if the task terminated due to a panic. Otherwise, `self` is /// returned. /// /// # Examples /// /// ```should_panic /// use std::panic; /// /// #[tokio::main] /// async fn main() { /// let err = tokio::spawn(async { /// panic!("boom"); /// }).await.unwrap_err(); /// /// if let Ok(reason) = err.try_into_panic() { /// // Resume the panic on the main task /// panic::resume_unwind(reason); /// } /// } /// ``` pub fn try_into_panic(self) -> Result, JoinError> { match self.repr { Repr::Panic(p) => Ok(p.into_inner()), _ => Err(self), } } /// Returns a [task ID] that identifies the task which errored relative to /// other currently spawned tasks. /// /// [task ID]: crate::task::Id pub fn id(&self) -> Id { self.id } } impl fmt::Display for JoinError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match &self.repr { Repr::Cancelled => write!(fmt, "task {} was cancelled", self.id), Repr::Panic(p) => match panic_payload_as_str(p) { Some(panic_str) => { write!( fmt, "task {} panicked with message {:?}", self.id, panic_str ) } None => { write!(fmt, "task {} panicked", self.id) } }, } } } impl fmt::Debug for JoinError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match &self.repr { Repr::Cancelled => write!(fmt, "JoinError::Cancelled({:?})", self.id), Repr::Panic(p) => match panic_payload_as_str(p) { Some(panic_str) => { write!(fmt, "JoinError::Panic({:?}, {:?}, ...)", self.id, panic_str) } None => write!(fmt, "JoinError::Panic({:?}, ...)", self.id), }, } } } impl std::error::Error for JoinError {} impl From for io::Error { fn from(src: JoinError) -> io::Error { io::Error::new( io::ErrorKind::Other, match src.repr { Repr::Cancelled => "task was cancelled", Repr::Panic(_) => "task panicked", }, ) } } fn panic_payload_as_str(payload: &SyncWrapper>) -> Option<&str> { // Panic payloads are almost always `String` (if invoked with formatting arguments) // or `&'static str` (if invoked with a string literal). // // Non-string panic payloads have niche use-cases, // so we don't really need to worry about those. if let Some(s) = payload.downcast_ref_sync::() { return Some(s); } if let Some(s) = payload.downcast_ref_sync::<&'static str>() { return Some(s); } None } tokio-1.43.1/src/runtime/task/harness.rs000064400000000000000000000515111046102023000162670ustar 00000000000000use crate::future::Future; use crate::runtime::task::core::{Cell, Core, Header, Trailer}; use crate::runtime::task::state::{Snapshot, State}; use crate::runtime::task::waker::waker_ref; use crate::runtime::task::{Id, JoinError, Notified, RawTask, Schedule, Task}; use crate::runtime::TaskMeta; use std::any::Any; use std::mem; use std::mem::ManuallyDrop; use std::panic; use std::ptr::NonNull; use std::task::{Context, Poll, Waker}; /// Typed raw task handle. pub(super) struct Harness { cell: NonNull>, } impl Harness where T: Future, S: 'static, { pub(super) unsafe fn from_raw(ptr: NonNull
) -> Harness { Harness { cell: ptr.cast::>(), } } fn header_ptr(&self) -> NonNull
{ self.cell.cast() } fn header(&self) -> &Header { unsafe { &*self.header_ptr().as_ptr() } } fn state(&self) -> &State { &self.header().state } fn trailer(&self) -> &Trailer { unsafe { &self.cell.as_ref().trailer } } fn core(&self) -> &Core { unsafe { &self.cell.as_ref().core } } } /// Task operations that can be implemented without being generic over the /// scheduler or task. Only one version of these methods should exist in the /// final binary. impl RawTask { pub(super) fn drop_reference(self) { if self.state().ref_dec() { self.dealloc(); } } /// This call consumes a ref-count and notifies the task. This will create a /// new Notified and submit it if necessary. /// /// The caller does not need to hold a ref-count besides the one that was /// passed to this call. pub(super) fn wake_by_val(&self) { use super::state::TransitionToNotifiedByVal; match self.state().transition_to_notified_by_val() { TransitionToNotifiedByVal::Submit => { // The caller has given us a ref-count, and the transition has // created a new ref-count, so we now hold two. We turn the new // ref-count Notified and pass it to the call to `schedule`. // // The old ref-count is retained for now to ensure that the task // is not dropped during the call to `schedule` if the call // drops the task it was given. self.schedule(); // Now that we have completed the call to schedule, we can // release our ref-count. self.drop_reference(); } TransitionToNotifiedByVal::Dealloc => { self.dealloc(); } TransitionToNotifiedByVal::DoNothing => {} } } /// This call notifies the task. It will not consume any ref-counts, but the /// caller should hold a ref-count. This will create a new Notified and /// submit it if necessary. pub(super) fn wake_by_ref(&self) { use super::state::TransitionToNotifiedByRef; match self.state().transition_to_notified_by_ref() { TransitionToNotifiedByRef::Submit => { // The transition above incremented the ref-count for a new task // and the caller also holds a ref-count. The caller's ref-count // ensures that the task is not destroyed even if the new task // is dropped before `schedule` returns. self.schedule(); } TransitionToNotifiedByRef::DoNothing => {} } } /// Remotely aborts the task. /// /// The caller should hold a ref-count, but we do not consume it. /// /// This is similar to `shutdown` except that it asks the runtime to perform /// the shutdown. This is necessary to avoid the shutdown happening in the /// wrong thread for non-Send tasks. pub(super) fn remote_abort(&self) { if self.state().transition_to_notified_and_cancel() { // The transition has created a new ref-count, which we turn into // a Notified and pass to the task. // // Since the caller holds a ref-count, the task cannot be destroyed // before the call to `schedule` returns even if the call drops the // `Notified` internally. self.schedule(); } } /// Try to set the waker notified when the task is complete. Returns true if /// the task has already completed. If this call returns false, then the /// waker will not be notified. pub(super) fn try_set_join_waker(&self, waker: &Waker) -> bool { can_read_output(self.header(), self.trailer(), waker) } } impl Harness where T: Future, S: Schedule, { pub(super) fn drop_reference(self) { if self.state().ref_dec() { self.dealloc(); } } /// Polls the inner future. A ref-count is consumed. /// /// All necessary state checks and transitions are performed. /// Panics raised while polling the future are handled. pub(super) fn poll(self) { // We pass our ref-count to `poll_inner`. match self.poll_inner() { PollFuture::Notified => { // The `poll_inner` call has given us two ref-counts back. // We give one of them to a new task and call `yield_now`. self.core() .scheduler .yield_now(Notified(self.get_new_task())); // The remaining ref-count is now dropped. We kept the extra // ref-count until now to ensure that even if the `yield_now` // call drops the provided task, the task isn't deallocated // before after `yield_now` returns. self.drop_reference(); } PollFuture::Complete => { self.complete(); } PollFuture::Dealloc => { self.dealloc(); } PollFuture::Done => (), } } /// Polls the task and cancel it if necessary. This takes ownership of a /// ref-count. /// /// If the return value is Notified, the caller is given ownership of two /// ref-counts. /// /// If the return value is Complete, the caller is given ownership of a /// single ref-count, which should be passed on to `complete`. /// /// If the return value is `Dealloc`, then this call consumed the last /// ref-count and the caller should call `dealloc`. /// /// Otherwise the ref-count is consumed and the caller should not access /// `self` again. fn poll_inner(&self) -> PollFuture { use super::state::{TransitionToIdle, TransitionToRunning}; match self.state().transition_to_running() { TransitionToRunning::Success => { // Separated to reduce LLVM codegen fn transition_result_to_poll_future(result: TransitionToIdle) -> PollFuture { match result { TransitionToIdle::Ok => PollFuture::Done, TransitionToIdle::OkNotified => PollFuture::Notified, TransitionToIdle::OkDealloc => PollFuture::Dealloc, TransitionToIdle::Cancelled => PollFuture::Complete, } } let header_ptr = self.header_ptr(); let waker_ref = waker_ref::(&header_ptr); let cx = Context::from_waker(&waker_ref); let res = poll_future(self.core(), cx); if res == Poll::Ready(()) { // The future completed. Move on to complete the task. return PollFuture::Complete; } let transition_res = self.state().transition_to_idle(); if let TransitionToIdle::Cancelled = transition_res { // The transition to idle failed because the task was // cancelled during the poll. cancel_task(self.core()); } transition_result_to_poll_future(transition_res) } TransitionToRunning::Cancelled => { cancel_task(self.core()); PollFuture::Complete } TransitionToRunning::Failed => PollFuture::Done, TransitionToRunning::Dealloc => PollFuture::Dealloc, } } /// Forcibly shuts down the task. /// /// Attempt to transition to `Running` in order to forcibly shutdown the /// task. If the task is currently running or in a state of completion, then /// there is nothing further to do. When the task completes running, it will /// notice the `CANCELLED` bit and finalize the task. pub(super) fn shutdown(self) { if !self.state().transition_to_shutdown() { // The task is concurrently running. No further work needed. self.drop_reference(); return; } // By transitioning the lifecycle to `Running`, we have permission to // drop the future. cancel_task(self.core()); self.complete(); } pub(super) fn dealloc(self) { // Observe that we expect to have mutable access to these objects // because we are going to drop them. This only matters when running // under loom. self.trailer().waker.with_mut(|_| ()); self.core().stage.with_mut(|_| ()); // Safety: The caller of this method just transitioned our ref-count to // zero, so it is our responsibility to release the allocation. // // We don't hold any references into the allocation at this point, but // it is possible for another thread to still hold a `&State` into the // allocation if that other thread has decremented its last ref-count, // but has not yet returned from the relevant method on `State`. // // However, the `State` type consists of just an `AtomicUsize`, and an // `AtomicUsize` wraps the entirety of its contents in an `UnsafeCell`. // As explained in the documentation for `UnsafeCell`, such references // are allowed to be dangling after their last use, even if the // reference has not yet gone out of scope. unsafe { drop(Box::from_raw(self.cell.as_ptr())); } } // ===== join handle ===== /// Read the task output into `dst`. pub(super) fn try_read_output(self, dst: &mut Poll>, waker: &Waker) { if can_read_output(self.header(), self.trailer(), waker) { *dst = Poll::Ready(self.core().take_output()); } } pub(super) fn drop_join_handle_slow(self) { // Try to unset `JOIN_INTEREST` and `JOIN_WAKER`. This must be done as a first step in // case the task concurrently completed. let transition = self.state().transition_to_join_handle_dropped(); if transition.drop_output { // It is our responsibility to drop the output. This is critical as // the task output may not be `Send` and as such must remain with // the scheduler or `JoinHandle`. i.e. if the output remains in the // task structure until the task is deallocated, it may be dropped // by a Waker on any arbitrary thread. // // Panics are delivered to the user via the `JoinHandle`. Given that // they are dropping the `JoinHandle`, we assume they are not // interested in the panic and swallow it. let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { self.core().drop_future_or_output(); })); } if transition.drop_waker { // If the JOIN_WAKER flag is unset at this point, the task is either // already terminal or not complete so the `JoinHandle` is responsible // for dropping the waker. // Safety: // If the JOIN_WAKER bit is not set the join handle has exclusive // access to the waker as per rule 2 in task/mod.rs. // This can only be the case at this point in two scenarios: // 1. The task completed and the runtime unset `JOIN_WAKER` flag // after accessing the waker during task completion. So the // `JoinHandle` is the only one to access the join waker here. // 2. The task is not completed so the `JoinHandle` was able to unset // `JOIN_WAKER` bit itself to get mutable access to the waker. // The runtime will not access the waker when this flag is unset. unsafe { self.trailer().set_waker(None) }; } // Drop the `JoinHandle` reference, possibly deallocating the task self.drop_reference(); } // ====== internal ====== /// Completes the task. This method assumes that the state is RUNNING. fn complete(self) { // The future has completed and its output has been written to the task // stage. We transition from running to complete. let snapshot = self.state().transition_to_complete(); // We catch panics here in case dropping the future or waking the // JoinHandle panics. let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { if !snapshot.is_join_interested() { // The `JoinHandle` is not interested in the output of // this task. It is our responsibility to drop the // output. The join waker was already dropped by the // `JoinHandle` before. self.core().drop_future_or_output(); } else if snapshot.is_join_waker_set() { // Notify the waker. Reading the waker field is safe per rule 4 // in task/mod.rs, since the JOIN_WAKER bit is set and the call // to transition_to_complete() above set the COMPLETE bit. self.trailer().wake_join(); // Inform the `JoinHandle` that we are done waking the waker by // unsetting the `JOIN_WAKER` bit. If the `JoinHandle` has // already been dropped and `JOIN_INTEREST` is unset, then we must // drop the waker ourselves. if !self .state() .unset_waker_after_complete() .is_join_interested() { // SAFETY: We have COMPLETE=1 and JOIN_INTEREST=0, so // we have exclusive access to the waker. unsafe { self.trailer().set_waker(None) }; } } })); // We catch panics here in case invoking a hook panics. // // We call this in a separate block so that it runs after the task appears to have // completed and will still run if the destructor panics. if let Some(f) = self.trailer().hooks.task_terminate_callback.as_ref() { let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { f(&TaskMeta { id: self.core().task_id, _phantom: Default::default(), }) })); } // The task has completed execution and will no longer be scheduled. let num_release = self.release(); if self.state().transition_to_terminal(num_release) { self.dealloc(); } } /// Releases the task from the scheduler. Returns the number of ref-counts /// that should be decremented. fn release(&self) -> usize { // We don't actually increment the ref-count here, but the new task is // never destroyed, so that's ok. let me = ManuallyDrop::new(self.get_new_task()); if let Some(task) = self.core().scheduler.release(&me) { mem::forget(task); 2 } else { 1 } } /// Creates a new task that holds its own ref-count. /// /// # Safety /// /// Any use of `self` after this call must ensure that a ref-count to the /// task holds the task alive until after the use of `self`. Passing the /// returned Task to any method on `self` is unsound if dropping the Task /// could drop `self` before the call on `self` returned. fn get_new_task(&self) -> Task { // safety: The header is at the beginning of the cell, so this cast is // safe. unsafe { Task::from_raw(self.cell.cast()) } } } fn can_read_output(header: &Header, trailer: &Trailer, waker: &Waker) -> bool { // Load a snapshot of the current task state let snapshot = header.state.load(); debug_assert!(snapshot.is_join_interested()); if !snapshot.is_complete() { // If the task is not complete, try storing the provided waker in the // task's waker field. let res = if snapshot.is_join_waker_set() { // If JOIN_WAKER is set, then JoinHandle has previously stored a // waker in the waker field per step (iii) of rule 5 in task/mod.rs. // Optimization: if the stored waker and the provided waker wake the // same task, then return without touching the waker field. (Reading // the waker field below is safe per rule 3 in task/mod.rs.) if unsafe { trailer.will_wake(waker) } { return false; } // Otherwise swap the stored waker with the provided waker by // following the rule 5 in task/mod.rs. header .state .unset_waker() .and_then(|snapshot| set_join_waker(header, trailer, waker.clone(), snapshot)) } else { // If JOIN_WAKER is unset, then JoinHandle has mutable access to the // waker field per rule 2 in task/mod.rs; therefore, skip step (i) // of rule 5 and try to store the provided waker in the waker field. set_join_waker(header, trailer, waker.clone(), snapshot) }; match res { Ok(_) => return false, Err(snapshot) => { assert!(snapshot.is_complete()); } } } true } fn set_join_waker( header: &Header, trailer: &Trailer, waker: Waker, snapshot: Snapshot, ) -> Result { assert!(snapshot.is_join_interested()); assert!(!snapshot.is_join_waker_set()); // Safety: Only the `JoinHandle` may set the `waker` field. When // `JOIN_INTEREST` is **not** set, nothing else will touch the field. unsafe { trailer.set_waker(Some(waker)); } // Update the `JoinWaker` state accordingly let res = header.state.set_join_waker(); // If the state could not be updated, then clear the join waker if res.is_err() { unsafe { trailer.set_waker(None); } } res } enum PollFuture { Complete, Notified, Done, Dealloc, } /// Cancels the task and store the appropriate error in the stage field. fn cancel_task(core: &Core) { // Drop the future from a panic guard. let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { core.drop_future_or_output(); })); core.store_output(Err(panic_result_to_join_error(core.task_id, res))); } fn panic_result_to_join_error( task_id: Id, res: Result<(), Box>, ) -> JoinError { match res { Ok(()) => JoinError::cancelled(task_id), Err(panic) => JoinError::panic(task_id, panic), } } /// Polls the future. If the future completes, the output is written to the /// stage field. fn poll_future(core: &Core, cx: Context<'_>) -> Poll<()> { // Poll the future. let output = panic::catch_unwind(panic::AssertUnwindSafe(|| { struct Guard<'a, T: Future, S: Schedule> { core: &'a Core, } impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> { fn drop(&mut self) { // If the future panics on poll, we drop it inside the panic // guard. self.core.drop_future_or_output(); } } let guard = Guard { core }; let res = guard.core.poll(cx); mem::forget(guard); res })); // Prepare output for being placed in the core stage. let output = match output { Ok(Poll::Pending) => return Poll::Pending, Ok(Poll::Ready(output)) => Ok(output), Err(panic) => Err(panic_to_error(&core.scheduler, core.task_id, panic)), }; // Catch and ignore panics if the future panics on drop. let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { core.store_output(output); })); if res.is_err() { core.scheduler.unhandled_panic(); } Poll::Ready(()) } #[cold] fn panic_to_error( scheduler: &S, task_id: Id, panic: Box, ) -> JoinError { scheduler.unhandled_panic(); JoinError::panic(task_id, panic) } tokio-1.43.1/src/runtime/task/id.rs000064400000000000000000000051201046102023000152130ustar 00000000000000use crate::runtime::context; use std::{fmt, num::NonZeroU64}; /// An opaque ID that uniquely identifies a task relative to all other currently /// running tasks. /// /// # Notes /// /// - Task IDs are unique relative to other *currently running* tasks. When a /// task completes, the same ID may be used for another task. /// - Task IDs are *not* sequential, and do not indicate the order in which /// tasks are spawned, what runtime a task is spawned on, or any other data. /// - The task ID of the currently running task can be obtained from inside the /// task via the [`task::try_id()`](crate::task::try_id()) and /// [`task::id()`](crate::task::id()) functions and from outside the task via /// the [`JoinHandle::id()`](crate::task::JoinHandle::id()) function. #[cfg_attr(docsrs, doc(cfg(all(feature = "rt"))))] #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)] pub struct Id(pub(crate) NonZeroU64); /// Returns the [`Id`] of the currently running task. /// /// # Panics /// /// This function panics if called from outside a task. Please note that calls /// to `block_on` do not have task IDs, so the method will panic if called from /// within a call to `block_on`. For a version of this function that doesn't /// panic, see [`task::try_id()`](crate::runtime::task::try_id()). /// /// [task ID]: crate::task::Id #[track_caller] pub fn id() -> Id { context::current_task_id().expect("Can't get a task id when not inside a task") } /// Returns the [`Id`] of the currently running task, or `None` if called outside /// of a task. /// /// This function is similar to [`task::id()`](crate::runtime::task::id()), except /// that it returns `None` rather than panicking if called outside of a task /// context. /// /// [task ID]: crate::task::Id #[track_caller] pub fn try_id() -> Option { context::current_task_id() } impl fmt::Display for Id { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.0.fmt(f) } } impl Id { pub(crate) fn next() -> Self { use crate::loom::sync::atomic::Ordering::Relaxed; use crate::loom::sync::atomic::StaticAtomicU64; #[cfg(all(test, loom))] crate::loom::lazy_static! { static ref NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1); } #[cfg(not(all(test, loom)))] static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1); loop { let id = NEXT_ID.fetch_add(1, Relaxed); if let Some(id) = NonZeroU64::new(id) { return Self(id); } } } pub(crate) fn as_u64(&self) -> u64 { self.0.get() } } tokio-1.43.1/src/runtime/task/join.rs000064400000000000000000000272671046102023000155760ustar 00000000000000use crate::runtime::task::{Header, RawTask}; use std::fmt; use std::future::Future; use std::marker::PhantomData; use std::panic::{RefUnwindSafe, UnwindSafe}; use std::pin::Pin; use std::task::{ready, Context, Poll, Waker}; cfg_rt! { /// An owned permission to join on a task (await its termination). /// /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] /// for a Tokio task rather than a thread. Note that the background task /// associated with this `JoinHandle` started running immediately when you /// called spawn, even if you have not yet awaited the `JoinHandle`. /// /// A `JoinHandle` *detaches* the associated task when it is dropped, which /// means that there is no longer any handle to the task, and no way to `join` /// on it. /// /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`] /// functions. /// /// # Cancel safety /// /// The `&mut JoinHandle` type is cancel safe. If it is used as the event /// in a `tokio::select!` statement and some other branch completes first, /// then it is guaranteed that the output of the task is not lost. /// /// If a `JoinHandle` is dropped, then the task continues running in the /// background and its return value is lost. /// /// # Examples /// /// Creation from [`task::spawn`]: /// /// ``` /// use tokio::task; /// /// # async fn doc() { /// let join_handle: task::JoinHandle<_> = task::spawn(async { /// // some work here /// }); /// # } /// ``` /// /// Creation from [`task::spawn_blocking`]: /// /// ``` /// use tokio::task; /// /// # async fn doc() { /// let join_handle: task::JoinHandle<_> = task::spawn_blocking(|| { /// // some blocking work here /// }); /// # } /// ``` /// /// The generic parameter `T` in `JoinHandle` is the return type of the spawned task. /// If the return value is an `i32`, the join handle has type `JoinHandle`: /// /// ``` /// use tokio::task; /// /// # async fn doc() { /// let join_handle: task::JoinHandle = task::spawn(async { /// 5 + 3 /// }); /// # } /// /// ``` /// /// If the task does not have a return value, the join handle has type `JoinHandle<()>`: /// /// ``` /// use tokio::task; /// /// # async fn doc() { /// let join_handle: task::JoinHandle<()> = task::spawn(async { /// println!("I return nothing."); /// }); /// # } /// ``` /// /// Note that `handle.await` doesn't give you the return type directly. It is wrapped in a /// `Result` because panics in the spawned task are caught by Tokio. The `?` operator has /// to be double chained to extract the returned value: /// /// ``` /// use tokio::task; /// use std::io; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let join_handle: task::JoinHandle> = tokio::spawn(async { /// Ok(5 + 3) /// }); /// /// let result = join_handle.await??; /// assert_eq!(result, 8); /// Ok(()) /// } /// ``` /// /// If the task panics, the error is a [`JoinError`] that contains the panic: /// /// ``` /// use tokio::task; /// use std::io; /// use std::panic; /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let join_handle: task::JoinHandle> = tokio::spawn(async { /// panic!("boom"); /// }); /// /// let err = join_handle.await.unwrap_err(); /// assert!(err.is_panic()); /// Ok(()) /// } /// /// ``` /// Child being detached and outliving its parent: /// /// ```no_run /// use tokio::task; /// use tokio::time; /// use std::time::Duration; /// /// # #[tokio::main] async fn main() { /// let original_task = task::spawn(async { /// let _detached_task = task::spawn(async { /// // Here we sleep to make sure that the first task returns before. /// time::sleep(Duration::from_millis(10)).await; /// // This will be called, even though the JoinHandle is dropped. /// println!("♫ Still alive ♫"); /// }); /// }); /// /// original_task.await.expect("The task being joined has panicked"); /// println!("Original task is joined."); /// /// // We make sure that the new task has time to run, before the main /// // task returns. /// /// time::sleep(Duration::from_millis(1000)).await; /// # } /// ``` /// /// [`task::spawn`]: crate::task::spawn() /// [`task::spawn_blocking`]: crate::task::spawn_blocking /// [`std::thread::JoinHandle`]: std::thread::JoinHandle /// [`JoinError`]: crate::task::JoinError pub struct JoinHandle { raw: RawTask, _p: PhantomData, } } unsafe impl Send for JoinHandle {} unsafe impl Sync for JoinHandle {} impl UnwindSafe for JoinHandle {} impl RefUnwindSafe for JoinHandle {} impl JoinHandle { pub(super) fn new(raw: RawTask) -> JoinHandle { JoinHandle { raw, _p: PhantomData, } } /// Abort the task associated with the handle. /// /// Awaiting a cancelled task might complete as usual if the task was /// already completed at the time it was cancelled, but most likely it /// will fail with a [cancelled] `JoinError`. /// /// Be aware that tasks spawned using [`spawn_blocking`] cannot be aborted /// because they are not async. If you call `abort` on a `spawn_blocking` /// task, then this *will not have any effect*, and the task will continue /// running normally. The exception is if the task has not started running /// yet; in that case, calling `abort` may prevent the task from starting. /// /// See also [the module level docs] for more information on cancellation. /// /// ```rust /// use tokio::time; /// /// # #[tokio::main(flavor = "current_thread", start_paused = true)] /// # async fn main() { /// let mut handles = Vec::new(); /// /// handles.push(tokio::spawn(async { /// time::sleep(time::Duration::from_secs(10)).await; /// true /// })); /// /// handles.push(tokio::spawn(async { /// time::sleep(time::Duration::from_secs(10)).await; /// false /// })); /// /// for handle in &handles { /// handle.abort(); /// } /// /// for handle in handles { /// assert!(handle.await.unwrap_err().is_cancelled()); /// } /// # } /// ``` /// /// [cancelled]: method@super::error::JoinError::is_cancelled /// [the module level docs]: crate::task#cancellation /// [`spawn_blocking`]: crate::task::spawn_blocking pub fn abort(&self) { self.raw.remote_abort(); } /// Checks if the task associated with this `JoinHandle` has finished. /// /// Please note that this method can return `false` even if [`abort`] has been /// called on the task. This is because the cancellation process may take /// some time, and this method does not return `true` until it has /// completed. /// /// ```rust /// use tokio::time; /// /// # #[tokio::main(flavor = "current_thread", start_paused = true)] /// # async fn main() { /// let handle1 = tokio::spawn(async { /// // do some stuff here /// }); /// let handle2 = tokio::spawn(async { /// // do some other stuff here /// time::sleep(time::Duration::from_secs(10)).await; /// }); /// // Wait for the task to finish /// handle2.abort(); /// time::sleep(time::Duration::from_secs(1)).await; /// assert!(handle1.is_finished()); /// assert!(handle2.is_finished()); /// # } /// ``` /// [`abort`]: method@JoinHandle::abort pub fn is_finished(&self) -> bool { let state = self.raw.header().state.load(); state.is_complete() } /// Set the waker that is notified when the task completes. pub(crate) fn set_join_waker(&mut self, waker: &Waker) { if self.raw.try_set_join_waker(waker) { // In this case the task has already completed. We wake the waker immediately. waker.wake_by_ref(); } } /// Returns a new `AbortHandle` that can be used to remotely abort this task. /// /// Awaiting a task cancelled by the `AbortHandle` might complete as usual if the task was /// already completed at the time it was cancelled, but most likely it /// will fail with a [cancelled] `JoinError`. /// /// ```rust /// use tokio::{time, task}; /// /// # #[tokio::main(flavor = "current_thread", start_paused = true)] /// # async fn main() { /// let mut handles = Vec::new(); /// /// handles.push(tokio::spawn(async { /// time::sleep(time::Duration::from_secs(10)).await; /// true /// })); /// /// handles.push(tokio::spawn(async { /// time::sleep(time::Duration::from_secs(10)).await; /// false /// })); /// /// let abort_handles: Vec = handles.iter().map(|h| h.abort_handle()).collect(); /// /// for handle in abort_handles { /// handle.abort(); /// } /// /// for handle in handles { /// assert!(handle.await.unwrap_err().is_cancelled()); /// } /// # } /// ``` /// [cancelled]: method@super::error::JoinError::is_cancelled #[must_use = "abort handles do nothing unless `.abort` is called"] pub fn abort_handle(&self) -> super::AbortHandle { self.raw.ref_inc(); super::AbortHandle::new(self.raw) } /// Returns a [task ID] that uniquely identifies this task relative to other /// currently spawned tasks. /// /// [task ID]: crate::task::Id pub fn id(&self) -> super::Id { // Safety: The header pointer is valid. unsafe { Header::get_id(self.raw.header_ptr()) } } } impl Unpin for JoinHandle {} impl Future for JoinHandle { type Output = super::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { ready!(crate::trace::trace_leaf(cx)); let mut ret = Poll::Pending; // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); // Try to read the task output. If the task is not yet complete, the // waker is stored and is notified once the task does complete. // // The function must go via the vtable, which requires erasing generic // types. To do this, the function "return" is placed on the stack // **before** calling the function and is passed into the function using // `*mut ()`. // // Safety: // // The type of `T` must match the task's output type. unsafe { self.raw .try_read_output(&mut ret as *mut _ as *mut (), cx.waker()); } if ret.is_ready() { coop.made_progress(); } ret } } impl Drop for JoinHandle { fn drop(&mut self) { if self.raw.state().drop_join_handle_fast().is_ok() { return; } self.raw.drop_join_handle_slow(); } } impl fmt::Debug for JoinHandle where T: fmt::Debug, { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { // Safety: The header pointer is valid. let id_ptr = unsafe { Header::get_id_ptr(self.raw.header_ptr()) }; let id = unsafe { id_ptr.as_ref() }; fmt.debug_struct("JoinHandle").field("id", id).finish() } } tokio-1.43.1/src/runtime/task/list.rs000064400000000000000000000264631046102023000156070ustar 00000000000000//! This module has containers for storing the tasks spawned on a scheduler. The //! `OwnedTasks` container is thread-safe but can only store tasks that //! implement Send. The `LocalOwnedTasks` container is not thread safe, but can //! store non-Send tasks. //! //! The collections can be closed to prevent adding new tasks during shutdown of //! the scheduler with the collection. use crate::future::Future; use crate::loom::cell::UnsafeCell; use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task}; use crate::util::linked_list::{Link, LinkedList}; use crate::util::sharded_list; use crate::loom::sync::atomic::{AtomicBool, Ordering}; use std::marker::PhantomData; use std::num::NonZeroU64; // The id from the module below is used to verify whether a given task is stored // in this OwnedTasks, or some other task. The counter starts at one so we can // use `None` for tasks not owned by any list. // // The safety checks in this file can technically be violated if the counter is // overflown, but the checks are not supposed to ever fail unless there is a // bug in Tokio, so we accept that certain bugs would not be caught if the two // mixed up runtimes happen to have the same id. cfg_has_atomic_u64! { use std::sync::atomic::AtomicU64; static NEXT_OWNED_TASKS_ID: AtomicU64 = AtomicU64::new(1); fn get_next_id() -> NonZeroU64 { loop { let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed); if let Some(id) = NonZeroU64::new(id) { return id; } } } } cfg_not_has_atomic_u64! { use std::sync::atomic::AtomicU32; static NEXT_OWNED_TASKS_ID: AtomicU32 = AtomicU32::new(1); fn get_next_id() -> NonZeroU64 { loop { let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed); if let Some(id) = NonZeroU64::new(u64::from(id)) { return id; } } } } pub(crate) struct OwnedTasks { list: List, pub(crate) id: NonZeroU64, closed: AtomicBool, } type List = sharded_list::ShardedList, as Link>::Target>; pub(crate) struct LocalOwnedTasks { inner: UnsafeCell>, pub(crate) id: NonZeroU64, _not_send_or_sync: PhantomData<*const ()>, } struct OwnedTasksInner { list: LinkedList, as Link>::Target>, closed: bool, } impl OwnedTasks { pub(crate) fn new(num_cores: usize) -> Self { let shard_size = Self::gen_shared_list_size(num_cores); Self { list: List::new(shard_size), closed: AtomicBool::new(false), id: get_next_id(), } } /// Binds the provided task to this `OwnedTasks` instance. This fails if the /// `OwnedTasks` has been closed. pub(crate) fn bind( &self, task: T, scheduler: S, id: super::Id, ) -> (JoinHandle, Option>) where S: Schedule, T: Future + Send + 'static, T::Output: Send + 'static, { let (task, notified, join) = super::new_task(task, scheduler, id); let notified = unsafe { self.bind_inner(task, notified) }; (join, notified) } /// Bind a task that isn't safe to transfer across thread boundaries. /// /// # Safety /// Only use this in `LocalRuntime` where the task cannot move pub(crate) unsafe fn bind_local( &self, task: T, scheduler: S, id: super::Id, ) -> (JoinHandle, Option>) where S: Schedule, T: Future + 'static, T::Output: 'static, { let (task, notified, join) = super::new_task(task, scheduler, id); let notified = unsafe { self.bind_inner(task, notified) }; (join, notified) } /// The part of `bind` that's the same for every type of future. unsafe fn bind_inner(&self, task: Task, notified: Notified) -> Option> where S: Schedule, { unsafe { // safety: We just created the task, so we have exclusive access // to the field. task.header().set_owner_id(self.id); } let shard = self.list.lock_shard(&task); // Check the closed flag in the lock for ensuring all that tasks // will shut down after the OwnedTasks has been closed. if self.closed.load(Ordering::Acquire) { drop(shard); task.shutdown(); return None; } shard.push(task); Some(notified) } /// Asserts that the given task is owned by this `OwnedTasks` and convert it to /// a `LocalNotified`, giving the thread permission to poll this task. #[inline] pub(crate) fn assert_owner(&self, task: Notified) -> LocalNotified { debug_assert_eq!(task.header().get_owner_id(), Some(self.id)); // safety: All tasks bound to this OwnedTasks are Send, so it is safe // to poll it on this thread no matter what thread we are on. LocalNotified { task: task.0, _not_send: PhantomData, } } /// Shuts down all tasks in the collection. This call also closes the /// collection, preventing new items from being added. /// /// The parameter start determines which shard this method will start at. /// Using different values for each worker thread reduces contention. pub(crate) fn close_and_shutdown_all(&self, start: usize) where S: Schedule, { self.closed.store(true, Ordering::Release); for i in start..self.get_shard_size() + start { loop { let task = self.list.pop_back(i); match task { Some(task) => { task.shutdown(); } None => break, } } } } #[inline] pub(crate) fn get_shard_size(&self) -> usize { self.list.shard_size() } pub(crate) fn num_alive_tasks(&self) -> usize { self.list.len() } cfg_64bit_metrics! { pub(crate) fn spawned_tasks_count(&self) -> u64 { self.list.added() } } pub(crate) fn remove(&self, task: &Task) -> Option> { // If the task's owner ID is `None` then it is not part of any list and // doesn't need removing. let task_id = task.header().get_owner_id()?; assert_eq!(task_id, self.id); // safety: We just checked that the provided task is not in some other // linked list. unsafe { self.list.remove(task.header_ptr()) } } pub(crate) fn is_empty(&self) -> bool { self.list.is_empty() } /// Generates the size of the sharded list based on the number of worker threads. /// /// The sharded lock design can effectively alleviate /// lock contention performance problems caused by high concurrency. /// /// However, as the number of shards increases, the memory continuity between /// nodes in the intrusive linked list will diminish. Furthermore, /// the construction time of the sharded list will also increase with a higher number of shards. /// /// Due to the above reasons, we set a maximum value for the shared list size, /// denoted as `MAX_SHARED_LIST_SIZE`. fn gen_shared_list_size(num_cores: usize) -> usize { const MAX_SHARED_LIST_SIZE: usize = 1 << 16; usize::min(MAX_SHARED_LIST_SIZE, num_cores.next_power_of_two() * 4) } } cfg_taskdump! { impl OwnedTasks { /// Locks the tasks, and calls `f` on an iterator over them. pub(crate) fn for_each(&self, f: F) where F: FnMut(&Task), { self.list.for_each(f); } } } impl LocalOwnedTasks { pub(crate) fn new() -> Self { Self { inner: UnsafeCell::new(OwnedTasksInner { list: LinkedList::new(), closed: false, }), id: get_next_id(), _not_send_or_sync: PhantomData, } } pub(crate) fn bind( &self, task: T, scheduler: S, id: super::Id, ) -> (JoinHandle, Option>) where S: Schedule, T: Future + 'static, T::Output: 'static, { let (task, notified, join) = super::new_task(task, scheduler, id); unsafe { // safety: We just created the task, so we have exclusive access // to the field. task.header().set_owner_id(self.id); } if self.is_closed() { drop(notified); task.shutdown(); (join, None) } else { self.with_inner(|inner| { inner.list.push_front(task); }); (join, Some(notified)) } } /// Shuts down all tasks in the collection. This call also closes the /// collection, preventing new items from being added. pub(crate) fn close_and_shutdown_all(&self) where S: Schedule, { self.with_inner(|inner| inner.closed = true); while let Some(task) = self.with_inner(|inner| inner.list.pop_back()) { task.shutdown(); } } pub(crate) fn remove(&self, task: &Task) -> Option> { // If the task's owner ID is `None` then it is not part of any list and // doesn't need removing. let task_id = task.header().get_owner_id()?; assert_eq!(task_id, self.id); self.with_inner(|inner| // safety: We just checked that the provided task is not in some // other linked list. unsafe { inner.list.remove(task.header_ptr()) }) } /// Asserts that the given task is owned by this `LocalOwnedTasks` and convert /// it to a `LocalNotified`, giving the thread permission to poll this task. #[inline] pub(crate) fn assert_owner(&self, task: Notified) -> LocalNotified { assert_eq!(task.header().get_owner_id(), Some(self.id)); // safety: The task was bound to this LocalOwnedTasks, and the // LocalOwnedTasks is not Send or Sync, so we are on the right thread // for polling this task. LocalNotified { task: task.0, _not_send: PhantomData, } } #[inline] fn with_inner(&self, f: F) -> T where F: FnOnce(&mut OwnedTasksInner) -> T, { // safety: This type is not Sync, so concurrent calls of this method // can't happen. Furthermore, all uses of this method in this file make // sure that they don't call `with_inner` recursively. self.inner.with_mut(|ptr| unsafe { f(&mut *ptr) }) } pub(crate) fn is_closed(&self) -> bool { self.with_inner(|inner| inner.closed) } pub(crate) fn is_empty(&self) -> bool { self.with_inner(|inner| inner.list.is_empty()) } } #[cfg(test)] mod tests { use super::*; // This test may run in parallel with other tests, so we only test that ids // come in increasing order. #[test] fn test_id_not_broken() { let mut last_id = get_next_id(); for _ in 0..1000 { let next_id = get_next_id(); assert!(last_id < next_id); last_id = next_id; } } } tokio-1.43.1/src/runtime/task/mod.rs000064400000000000000000000445611046102023000154120ustar 00000000000000//! The task module. //! //! The task module contains the code that manages spawned tasks and provides a //! safe API for the rest of the runtime to use. Each task in a runtime is //! stored in an `OwnedTasks` or `LocalOwnedTasks` object. //! //! # Task reference types //! //! A task is usually referenced by multiple handles, and there are several //! types of handles. //! //! * `OwnedTask` - tasks stored in an `OwnedTasks` or `LocalOwnedTasks` are of this //! reference type. //! //! * `JoinHandle` - each task has a `JoinHandle` that allows access to the output //! of the task. //! //! * `Waker` - every waker for a task has this reference type. There can be any //! number of waker references. //! //! * `Notified` - tracks whether the task is notified. //! //! * `Unowned` - this task reference type is used for tasks not stored in any //! runtime. Mainly used for blocking tasks, but also in tests. //! //! The task uses a reference count to keep track of how many active references //! exist. The `Unowned` reference type takes up two ref-counts. All other //! reference types take up a single ref-count. //! //! Besides the waker type, each task has at most one of each reference type. //! //! # State //! //! The task stores its state in an atomic `usize` with various bitfields for the //! necessary information. The state has the following bitfields: //! //! * `RUNNING` - Tracks whether the task is currently being polled or cancelled. //! This bit functions as a lock around the task. //! //! * `COMPLETE` - Is one once the future has fully completed and has been //! dropped. Never unset once set. Never set together with RUNNING. //! //! * `NOTIFIED` - Tracks whether a Notified object currently exists. //! //! * `CANCELLED` - Is set to one for tasks that should be cancelled as soon as //! possible. May take any value for completed tasks. //! //! * `JOIN_INTEREST` - Is set to one if there exists a `JoinHandle`. //! //! * `JOIN_WAKER` - Acts as an access control bit for the join handle waker. The //! protocol for its usage is described below. //! //! The rest of the bits are used for the ref-count. //! //! # Fields in the task //! //! The task has various fields. This section describes how and when it is safe //! to access a field. //! //! * The state field is accessed with atomic instructions. //! //! * The `OwnedTask` reference has exclusive access to the `owned` field. //! //! * The Notified reference has exclusive access to the `queue_next` field. //! //! * The `owner_id` field can be set as part of construction of the task, but //! is otherwise immutable and anyone can access the field immutably without //! synchronization. //! //! * If COMPLETE is one, then the `JoinHandle` has exclusive access to the //! stage field. If COMPLETE is zero, then the RUNNING bitfield functions as //! a lock for the stage field, and it can be accessed only by the thread //! that set RUNNING to one. //! //! * The waker field may be concurrently accessed by different threads: in one //! thread the runtime may complete a task and *read* the waker field to //! invoke the waker, and in another thread the task's `JoinHandle` may be //! polled, and if the task hasn't yet completed, the `JoinHandle` may *write* //! a waker to the waker field. The `JOIN_WAKER` bit ensures safe access by //! multiple threads to the waker field using the following rules: //! //! 1. `JOIN_WAKER` is initialized to zero. //! //! 2. If `JOIN_WAKER` is zero, then the `JoinHandle` has exclusive (mutable) //! access to the waker field. //! //! 3. If `JOIN_WAKER` is one, then the `JoinHandle` has shared (read-only) //! access to the waker field. //! //! 4. If `JOIN_WAKER` is one and COMPLETE is one, then the runtime has shared //! (read-only) access to the waker field. //! //! 5. If the `JoinHandle` needs to write to the waker field, then the //! `JoinHandle` needs to (i) successfully set `JOIN_WAKER` to zero if it is //! not already zero to gain exclusive access to the waker field per rule //! 2, (ii) write a waker, and (iii) successfully set `JOIN_WAKER` to one. //! If the `JoinHandle` unsets `JOIN_WAKER` in the process of being dropped //! to clear the waker field, only steps (i) and (ii) are relevant. //! //! 6. The `JoinHandle` can change `JOIN_WAKER` only if COMPLETE is zero (i.e. //! the task hasn't yet completed). The runtime can change `JOIN_WAKER` only //! if COMPLETE is one. //! //! 7. If `JOIN_INTEREST` is zero and COMPLETE is one, then the runtime has //! exclusive (mutable) access to the waker field. This might happen if the //! `JoinHandle` gets dropped right after the task completes and the runtime //! sets the `COMPLETE` bit. In this case the runtime needs the mutable access //! to the waker field to drop it. //! //! Rule 6 implies that the steps (i) or (iii) of rule 5 may fail due to a //! race. If step (i) fails, then the attempt to write a waker is aborted. If //! step (iii) fails because COMPLETE is set to one by another thread after //! step (i), then the waker field is cleared. Once COMPLETE is one (i.e. //! task has completed), the `JoinHandle` will not modify `JOIN_WAKER`. After the //! runtime sets COMPLETE to one, it invokes the waker if there is one so in this //! case when a task completes the `JOIN_WAKER` bit implicates to the runtime //! whether it should invoke the waker or not. After the runtime is done with //! using the waker during task completion, it unsets the `JOIN_WAKER` bit to give //! the `JoinHandle` exclusive access again so that it is able to drop the waker //! at a later point. //! //! All other fields are immutable and can be accessed immutably without //! synchronization by anyone. //! //! # Safety //! //! This section goes through various situations and explains why the API is //! safe in that situation. //! //! ## Polling or dropping the future //! //! Any mutable access to the future happens after obtaining a lock by modifying //! the RUNNING field, so exclusive access is ensured. //! //! When the task completes, exclusive access to the output is transferred to //! the `JoinHandle`. If the `JoinHandle` is already dropped when the transition to //! complete happens, the thread performing that transition retains exclusive //! access to the output and should immediately drop it. //! //! ## Non-Send futures //! //! If a future is not Send, then it is bound to a `LocalOwnedTasks`. The future //! will only ever be polled or dropped given a `LocalNotified` or inside a call //! to `LocalOwnedTasks::shutdown_all`. In either case, it is guaranteed that the //! future is on the right thread. //! //! If the task is never removed from the `LocalOwnedTasks`, then it is leaked, so //! there is no risk that the task is dropped on some other thread when the last //! ref-count drops. //! //! ## Non-Send output //! //! When a task completes, the output is placed in the stage of the task. Then, //! a transition that sets COMPLETE to true is performed, and the value of //! `JOIN_INTEREST` when this transition happens is read. //! //! If `JOIN_INTEREST` is zero when the transition to COMPLETE happens, then the //! output is immediately dropped. //! //! If `JOIN_INTEREST` is one when the transition to COMPLETE happens, then the //! `JoinHandle` is responsible for cleaning up the output. If the output is not //! Send, then this happens: //! //! 1. The output is created on the thread that the future was polled on. Since //! only non-Send futures can have non-Send output, the future was polled on //! the thread that the future was spawned from. //! 2. Since `JoinHandle` is not Send if Output is not Send, the //! `JoinHandle` is also on the thread that the future was spawned from. //! 3. Thus, the `JoinHandle` will not move the output across threads when it //! takes or drops the output. //! //! ## Recursive poll/shutdown //! //! Calling poll from inside a shutdown call or vice-versa is not prevented by //! the API exposed by the task module, so this has to be safe. In either case, //! the lock in the RUNNING bitfield makes the inner call return immediately. If //! the inner call is a `shutdown` call, then the CANCELLED bit is set, and the //! poll call will notice it when the poll finishes, and the task is cancelled //! at that point. // Some task infrastructure is here to support `JoinSet`, which is currently // unstable. This should be removed once `JoinSet` is stabilized. #![cfg_attr(not(tokio_unstable), allow(dead_code))] mod core; use self::core::Cell; use self::core::Header; mod error; pub use self::error::JoinError; mod harness; use self::harness::Harness; mod id; #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, unused_imports))] pub use id::{id, try_id, Id}; #[cfg(feature = "rt")] mod abort; mod join; #[cfg(feature = "rt")] pub use self::abort::AbortHandle; pub use self::join::JoinHandle; mod list; pub(crate) use self::list::{LocalOwnedTasks, OwnedTasks}; mod raw; pub(crate) use self::raw::RawTask; mod state; use self::state::State; mod waker; cfg_taskdump! { pub(crate) mod trace; } use crate::future::Future; use crate::util::linked_list; use crate::util::sharded_list; use crate::runtime::TaskCallback; use std::marker::PhantomData; use std::ptr::NonNull; use std::{fmt, mem}; /// An owned handle to the task, tracked by ref count. #[repr(transparent)] pub(crate) struct Task { raw: RawTask, _p: PhantomData, } unsafe impl Send for Task {} unsafe impl Sync for Task {} /// A task was notified. #[repr(transparent)] pub(crate) struct Notified(Task); // safety: This type cannot be used to touch the task without first verifying // that the value is on a thread where it is safe to poll the task. unsafe impl Send for Notified {} unsafe impl Sync for Notified {} /// A non-Send variant of Notified with the invariant that it is on a thread /// where it is safe to poll it. #[repr(transparent)] pub(crate) struct LocalNotified { task: Task, _not_send: PhantomData<*const ()>, } /// A task that is not owned by any `OwnedTasks`. Used for blocking tasks. /// This type holds two ref-counts. pub(crate) struct UnownedTask { raw: RawTask, _p: PhantomData, } // safety: This type can only be created given a Send task. unsafe impl Send for UnownedTask {} unsafe impl Sync for UnownedTask {} /// Task result sent back. pub(crate) type Result = std::result::Result; /// Hooks for scheduling tasks which are needed in the task harness. #[derive(Clone)] pub(crate) struct TaskHarnessScheduleHooks { pub(crate) task_terminate_callback: Option, } pub(crate) trait Schedule: Sync + Sized + 'static { /// The task has completed work and is ready to be released. The scheduler /// should release it immediately and return it. The task module will batch /// the ref-dec with setting other options. /// /// If the scheduler has already released the task, then None is returned. fn release(&self, task: &Task) -> Option>; /// Schedule the task fn schedule(&self, task: Notified); fn hooks(&self) -> TaskHarnessScheduleHooks; /// Schedule the task to run in the near future, yielding the thread to /// other tasks. fn yield_now(&self, task: Notified) { self.schedule(task); } /// Polling the task resulted in a panic. Should the runtime shutdown? fn unhandled_panic(&self) { // By default, do nothing. This maintains the 1.0 behavior. } } cfg_rt! { /// This is the constructor for a new task. Three references to the task are /// created. The first task reference is usually put into an `OwnedTasks` /// immediately. The Notified is sent to the scheduler as an ordinary /// notification. fn new_task( task: T, scheduler: S, id: Id, ) -> (Task, Notified, JoinHandle) where S: Schedule, T: Future + 'static, T::Output: 'static, { let raw = RawTask::new::(task, scheduler, id); let task = Task { raw, _p: PhantomData, }; let notified = Notified(Task { raw, _p: PhantomData, }); let join = JoinHandle::new(raw); (task, notified, join) } /// Creates a new task with an associated join handle. This method is used /// only when the task is not going to be stored in an `OwnedTasks` list. /// /// Currently only blocking tasks use this method. pub(crate) fn unowned(task: T, scheduler: S, id: Id) -> (UnownedTask, JoinHandle) where S: Schedule, T: Send + Future + 'static, T::Output: Send + 'static, { let (task, notified, join) = new_task(task, scheduler, id); // This transfers the ref-count of task and notified into an UnownedTask. // This is valid because an UnownedTask holds two ref-counts. let unowned = UnownedTask { raw: task.raw, _p: PhantomData, }; std::mem::forget(task); std::mem::forget(notified); (unowned, join) } } impl Task { unsafe fn new(raw: RawTask) -> Task { Task { raw, _p: PhantomData, } } unsafe fn from_raw(ptr: NonNull
) -> Task { Task::new(RawTask::from_raw(ptr)) } #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] pub(super) fn as_raw(&self) -> RawTask { self.raw } fn header(&self) -> &Header { self.raw.header() } fn header_ptr(&self) -> NonNull
{ self.raw.header_ptr() } cfg_taskdump! { /// Notify the task for task dumping. /// /// Returns `None` if the task has already been notified. pub(super) fn notify_for_tracing(&self) -> Option> { if self.as_raw().state().transition_to_notified_for_tracing() { // SAFETY: `transition_to_notified_for_tracing` increments the // refcount. Some(unsafe { Notified(Task::new(self.raw)) }) } else { None } } /// Returns a [task ID] that uniquely identifies this task relative to other /// currently spawned tasks. /// /// [task ID]: crate::task::Id #[cfg(tokio_unstable)] pub(crate) fn id(&self) -> crate::task::Id { // Safety: The header pointer is valid. unsafe { Header::get_id(self.raw.header_ptr()) } } } } impl Notified { fn header(&self) -> &Header { self.0.header() } } impl Notified { pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified { Notified(Task::new(ptr)) } } impl Notified { pub(crate) fn into_raw(self) -> RawTask { let raw = self.0.raw; mem::forget(self); raw } } impl Task { /// Preemptively cancels the task as part of the shutdown process. pub(crate) fn shutdown(self) { let raw = self.raw; mem::forget(self); raw.shutdown(); } } impl LocalNotified { /// Runs the task. pub(crate) fn run(self) { let raw = self.task.raw; mem::forget(self); raw.poll(); } } impl UnownedTask { // Used in test of the inject queue. #[cfg(test)] #[cfg_attr(target_family = "wasm", allow(dead_code))] pub(super) fn into_notified(self) -> Notified { Notified(self.into_task()) } fn into_task(self) -> Task { // Convert into a task. let task = Task { raw: self.raw, _p: PhantomData, }; mem::forget(self); // Drop a ref-count since an UnownedTask holds two. task.header().state.ref_dec(); task } pub(crate) fn run(self) { let raw = self.raw; mem::forget(self); // Transfer one ref-count to a Task object. let task = Task:: { raw, _p: PhantomData, }; // Use the other ref-count to poll the task. raw.poll(); // Decrement our extra ref-count drop(task); } pub(crate) fn shutdown(self) { self.into_task().shutdown(); } } impl Drop for Task { fn drop(&mut self) { // Decrement the ref count if self.header().state.ref_dec() { // Deallocate if this is the final ref count self.raw.dealloc(); } } } impl Drop for UnownedTask { fn drop(&mut self) { // Decrement the ref count if self.raw.header().state.ref_dec_twice() { // Deallocate if this is the final ref count self.raw.dealloc(); } } } impl fmt::Debug for Task { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "Task({:p})", self.header()) } } impl fmt::Debug for Notified { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "task::Notified({:p})", self.0.header()) } } /// # Safety /// /// Tasks are pinned. unsafe impl linked_list::Link for Task { type Handle = Task; type Target = Header; fn as_raw(handle: &Task) -> NonNull
{ handle.raw.header_ptr() } unsafe fn from_raw(ptr: NonNull
) -> Task { Task::from_raw(ptr) } unsafe fn pointers(target: NonNull
) -> NonNull> { self::core::Trailer::addr_of_owned(Header::get_trailer(target)) } } /// # Safety /// /// The id of a task is never changed after creation of the task, so the return value of /// `get_shard_id` will not change. (The cast may throw away the upper 32 bits of the task id, but /// the shard id still won't change from call to call.) unsafe impl sharded_list::ShardedListItem for Task { unsafe fn get_shard_id(target: NonNull) -> usize { // SAFETY: The caller guarantees that `target` points at a valid task. let task_id = unsafe { Header::get_id(target) }; task_id.0.get() as usize } } tokio-1.43.1/src/runtime/task/raw.rs000064400000000000000000000224161046102023000154170ustar 00000000000000use crate::future::Future; use crate::runtime::task::core::{Core, Trailer}; use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State}; use std::ptr::NonNull; use std::task::{Poll, Waker}; /// Raw task handle #[derive(Clone)] pub(crate) struct RawTask { ptr: NonNull
, } pub(super) struct Vtable { /// Polls the future. pub(super) poll: unsafe fn(NonNull
), /// Schedules the task for execution on the runtime. pub(super) schedule: unsafe fn(NonNull
), /// Deallocates the memory. pub(super) dealloc: unsafe fn(NonNull
), /// Reads the task output, if complete. pub(super) try_read_output: unsafe fn(NonNull
, *mut (), &Waker), /// The join handle has been dropped. pub(super) drop_join_handle_slow: unsafe fn(NonNull
), /// An abort handle has been dropped. pub(super) drop_abort_handle: unsafe fn(NonNull
), /// Scheduler is being shutdown. pub(super) shutdown: unsafe fn(NonNull
), /// The number of bytes that the `trailer` field is offset from the header. pub(super) trailer_offset: usize, /// The number of bytes that the `scheduler` field is offset from the header. pub(super) scheduler_offset: usize, /// The number of bytes that the `id` field is offset from the header. pub(super) id_offset: usize, } /// Get the vtable for the requested `T` and `S` generics. pub(super) fn vtable() -> &'static Vtable { &Vtable { poll: poll::, schedule: schedule::, dealloc: dealloc::, try_read_output: try_read_output::, drop_join_handle_slow: drop_join_handle_slow::, drop_abort_handle: drop_abort_handle::, shutdown: shutdown::, trailer_offset: OffsetHelper::::TRAILER_OFFSET, scheduler_offset: OffsetHelper::::SCHEDULER_OFFSET, id_offset: OffsetHelper::::ID_OFFSET, } } /// Calling `get_trailer_offset` directly in vtable doesn't work because it /// prevents the vtable from being promoted to a static reference. /// /// See this thread for more info: /// struct OffsetHelper(T, S); impl OffsetHelper { // Pass `size_of`/`align_of` as arguments rather than calling them directly // inside `get_trailer_offset` because trait bounds on generic parameters // of const fn are unstable on our MSRV. const TRAILER_OFFSET: usize = get_trailer_offset( std::mem::size_of::
(), std::mem::size_of::>(), std::mem::align_of::>(), std::mem::align_of::(), ); // The `scheduler` is the first field of `Core`, so it has the same // offset as `Core`. const SCHEDULER_OFFSET: usize = get_core_offset( std::mem::size_of::
(), std::mem::align_of::>(), ); const ID_OFFSET: usize = get_id_offset( std::mem::size_of::
(), std::mem::align_of::>(), std::mem::size_of::(), std::mem::align_of::(), ); } /// Compute the offset of the `Trailer` field in `Cell` using the /// `#[repr(C)]` algorithm. /// /// Pseudo-code for the `#[repr(C)]` algorithm can be found here: /// const fn get_trailer_offset( header_size: usize, core_size: usize, core_align: usize, trailer_align: usize, ) -> usize { let mut offset = header_size; let core_misalign = offset % core_align; if core_misalign > 0 { offset += core_align - core_misalign; } offset += core_size; let trailer_misalign = offset % trailer_align; if trailer_misalign > 0 { offset += trailer_align - trailer_misalign; } offset } /// Compute the offset of the `Core` field in `Cell` using the /// `#[repr(C)]` algorithm. /// /// Pseudo-code for the `#[repr(C)]` algorithm can be found here: /// const fn get_core_offset(header_size: usize, core_align: usize) -> usize { let mut offset = header_size; let core_misalign = offset % core_align; if core_misalign > 0 { offset += core_align - core_misalign; } offset } /// Compute the offset of the `Id` field in `Cell` using the /// `#[repr(C)]` algorithm. /// /// Pseudo-code for the `#[repr(C)]` algorithm can be found here: /// const fn get_id_offset( header_size: usize, core_align: usize, scheduler_size: usize, id_align: usize, ) -> usize { let mut offset = get_core_offset(header_size, core_align); offset += scheduler_size; let id_misalign = offset % id_align; if id_misalign > 0 { offset += id_align - id_misalign; } offset } impl RawTask { pub(super) fn new(task: T, scheduler: S, id: Id) -> RawTask where T: Future, S: Schedule, { let ptr = Box::into_raw(Cell::<_, S>::new(task, scheduler, State::new(), id)); let ptr = unsafe { NonNull::new_unchecked(ptr.cast()) }; RawTask { ptr } } pub(super) unsafe fn from_raw(ptr: NonNull
) -> RawTask { RawTask { ptr } } pub(super) fn header_ptr(&self) -> NonNull
{ self.ptr } pub(super) fn trailer_ptr(&self) -> NonNull { unsafe { Header::get_trailer(self.ptr) } } /// Returns a reference to the task's header. pub(super) fn header(&self) -> &Header { unsafe { self.ptr.as_ref() } } /// Returns a reference to the task's trailer. pub(super) fn trailer(&self) -> &Trailer { unsafe { &*self.trailer_ptr().as_ptr() } } /// Returns a reference to the task's state. pub(super) fn state(&self) -> &State { &self.header().state } /// Safety: mutual exclusion is required to call this function. pub(crate) fn poll(self) { let vtable = self.header().vtable; unsafe { (vtable.poll)(self.ptr) } } pub(super) fn schedule(self) { let vtable = self.header().vtable; unsafe { (vtable.schedule)(self.ptr) } } pub(super) fn dealloc(self) { let vtable = self.header().vtable; unsafe { (vtable.dealloc)(self.ptr); } } /// Safety: `dst` must be a `*mut Poll>` where `T` /// is the future stored by the task. pub(super) unsafe fn try_read_output(self, dst: *mut (), waker: &Waker) { let vtable = self.header().vtable; (vtable.try_read_output)(self.ptr, dst, waker); } pub(super) fn drop_join_handle_slow(self) { let vtable = self.header().vtable; unsafe { (vtable.drop_join_handle_slow)(self.ptr) } } pub(super) fn drop_abort_handle(self) { let vtable = self.header().vtable; unsafe { (vtable.drop_abort_handle)(self.ptr) } } pub(super) fn shutdown(self) { let vtable = self.header().vtable; unsafe { (vtable.shutdown)(self.ptr) } } /// Increment the task's reference count. /// /// Currently, this is used only when creating an `AbortHandle`. pub(super) fn ref_inc(self) { self.header().state.ref_inc(); } /// Get the queue-next pointer /// /// This is for usage by the injection queue /// /// Safety: make sure only one queue uses this and access is synchronized. pub(crate) unsafe fn get_queue_next(self) -> Option { self.header() .queue_next .with(|ptr| *ptr) .map(|p| RawTask::from_raw(p)) } /// Sets the queue-next pointer /// /// This is for usage by the injection queue /// /// Safety: make sure only one queue uses this and access is synchronized. pub(crate) unsafe fn set_queue_next(self, val: Option) { self.header().set_next(val.map(|task| task.ptr)); } } impl Copy for RawTask {} unsafe fn poll(ptr: NonNull
) { let harness = Harness::::from_raw(ptr); harness.poll(); } unsafe fn schedule(ptr: NonNull
) { use crate::runtime::task::{Notified, Task}; let scheduler = Header::get_scheduler::(ptr); scheduler .as_ref() .schedule(Notified(Task::from_raw(ptr.cast()))); } unsafe fn dealloc(ptr: NonNull
) { let harness = Harness::::from_raw(ptr); harness.dealloc(); } unsafe fn try_read_output( ptr: NonNull
, dst: *mut (), waker: &Waker, ) { let out = &mut *(dst as *mut Poll>); let harness = Harness::::from_raw(ptr); harness.try_read_output(out, waker); } unsafe fn drop_join_handle_slow(ptr: NonNull
) { let harness = Harness::::from_raw(ptr); harness.drop_join_handle_slow(); } unsafe fn drop_abort_handle(ptr: NonNull
) { let harness = Harness::::from_raw(ptr); harness.drop_reference(); } unsafe fn shutdown(ptr: NonNull
) { let harness = Harness::::from_raw(ptr); harness.shutdown(); } tokio-1.43.1/src/runtime/task/state.rs000064400000000000000000000522031046102023000157430ustar 00000000000000use crate::loom::sync::atomic::AtomicUsize; use std::fmt; use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; pub(super) struct State { val: AtomicUsize, } /// Current state value. #[derive(Copy, Clone)] pub(super) struct Snapshot(usize); type UpdateResult = Result; /// The task is currently being run. const RUNNING: usize = 0b0001; /// The task is complete. /// /// Once this bit is set, it is never unset. const COMPLETE: usize = 0b0010; /// Extracts the task's lifecycle value from the state. const LIFECYCLE_MASK: usize = 0b11; /// Flag tracking if the task has been pushed into a run queue. const NOTIFIED: usize = 0b100; /// The join handle is still around. const JOIN_INTEREST: usize = 0b1_000; /// A join handle waker has been set. const JOIN_WAKER: usize = 0b10_000; /// The task has been forcibly cancelled. const CANCELLED: usize = 0b100_000; /// All bits. const STATE_MASK: usize = LIFECYCLE_MASK | NOTIFIED | JOIN_INTEREST | JOIN_WAKER | CANCELLED; /// Bits used by the ref count portion of the state. const REF_COUNT_MASK: usize = !STATE_MASK; /// Number of positions to shift the ref count. const REF_COUNT_SHIFT: usize = REF_COUNT_MASK.count_zeros() as usize; /// One ref count. const REF_ONE: usize = 1 << REF_COUNT_SHIFT; /// State a task is initialized with. /// /// A task is initialized with three references: /// /// * A reference that will be stored in an `OwnedTasks` or `LocalOwnedTasks`. /// * A reference that will be sent to the scheduler as an ordinary notification. /// * A reference for the `JoinHandle`. /// /// As the task starts with a `JoinHandle`, `JOIN_INTEREST` is set. /// As the task starts with a `Notified`, `NOTIFIED` is set. const INITIAL_STATE: usize = (REF_ONE * 3) | JOIN_INTEREST | NOTIFIED; #[must_use] pub(super) enum TransitionToRunning { Success, Cancelled, Failed, Dealloc, } #[must_use] pub(super) enum TransitionToIdle { Ok, OkNotified, OkDealloc, Cancelled, } #[must_use] pub(super) enum TransitionToNotifiedByVal { DoNothing, Submit, Dealloc, } #[must_use] pub(crate) enum TransitionToNotifiedByRef { DoNothing, Submit, } #[must_use] pub(super) struct TransitionToJoinHandleDrop { pub(super) drop_waker: bool, pub(super) drop_output: bool, } /// All transitions are performed via RMW operations. This establishes an /// unambiguous modification order. impl State { /// Returns a task's initial state. pub(super) fn new() -> State { // The raw task returned by this method has a ref-count of three. See // the comment on INITIAL_STATE for more. State { val: AtomicUsize::new(INITIAL_STATE), } } /// Loads the current state, establishes `Acquire` ordering. pub(super) fn load(&self) -> Snapshot { Snapshot(self.val.load(Acquire)) } /// Attempts to transition the lifecycle to `Running`. This sets the /// notified bit to false so notifications during the poll can be detected. pub(super) fn transition_to_running(&self) -> TransitionToRunning { self.fetch_update_action(|mut next| { let action; assert!(next.is_notified()); if !next.is_idle() { // This happens if the task is either currently running or if it // has already completed, e.g. if it was cancelled during // shutdown. Consume the ref-count and return. next.ref_dec(); if next.ref_count() == 0 { action = TransitionToRunning::Dealloc; } else { action = TransitionToRunning::Failed; } } else { // We are able to lock the RUNNING bit. next.set_running(); next.unset_notified(); if next.is_cancelled() { action = TransitionToRunning::Cancelled; } else { action = TransitionToRunning::Success; } } (action, Some(next)) }) } /// Transitions the task from `Running` -> `Idle`. /// /// The transition to `Idle` fails if the task has been flagged to be /// cancelled. pub(super) fn transition_to_idle(&self) -> TransitionToIdle { self.fetch_update_action(|curr| { assert!(curr.is_running()); if curr.is_cancelled() { return (TransitionToIdle::Cancelled, None); } let mut next = curr; let action; next.unset_running(); if !next.is_notified() { // Polling the future consumes the ref-count of the Notified. next.ref_dec(); if next.ref_count() == 0 { action = TransitionToIdle::OkDealloc; } else { action = TransitionToIdle::Ok; } } else { // The caller will schedule a new notification, so we create a // new ref-count for the notification. Our own ref-count is kept // for now, and the caller will drop it shortly. next.ref_inc(); action = TransitionToIdle::OkNotified; } (action, Some(next)) }) } /// Transitions the task from `Running` -> `Complete`. pub(super) fn transition_to_complete(&self) -> Snapshot { const DELTA: usize = RUNNING | COMPLETE; let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel)); assert!(prev.is_running()); assert!(!prev.is_complete()); Snapshot(prev.0 ^ DELTA) } /// Transitions from `Complete` -> `Terminal`, decrementing the reference /// count the specified number of times. /// /// Returns true if the task should be deallocated. pub(super) fn transition_to_terminal(&self, count: usize) -> bool { let prev = Snapshot(self.val.fetch_sub(count * REF_ONE, AcqRel)); assert!( prev.ref_count() >= count, "current: {}, sub: {}", prev.ref_count(), count ); prev.ref_count() == count } /// Transitions the state to `NOTIFIED`. /// /// If no task needs to be submitted, a ref-count is consumed. /// /// If a task needs to be submitted, the ref-count is incremented for the /// new Notified. pub(super) fn transition_to_notified_by_val(&self) -> TransitionToNotifiedByVal { self.fetch_update_action(|mut snapshot| { let action; if snapshot.is_running() { // If the task is running, we mark it as notified, but we should // not submit anything as the thread currently running the // future is responsible for that. snapshot.set_notified(); snapshot.ref_dec(); // The thread that set the running bit also holds a ref-count. assert!(snapshot.ref_count() > 0); action = TransitionToNotifiedByVal::DoNothing; } else if snapshot.is_complete() || snapshot.is_notified() { // We do not need to submit any notifications, but we have to // decrement the ref-count. snapshot.ref_dec(); if snapshot.ref_count() == 0 { action = TransitionToNotifiedByVal::Dealloc; } else { action = TransitionToNotifiedByVal::DoNothing; } } else { // We create a new notified that we can submit. The caller // retains ownership of the ref-count they passed in. snapshot.set_notified(); snapshot.ref_inc(); action = TransitionToNotifiedByVal::Submit; } (action, Some(snapshot)) }) } /// Transitions the state to `NOTIFIED`. pub(super) fn transition_to_notified_by_ref(&self) -> TransitionToNotifiedByRef { self.fetch_update_action(|mut snapshot| { if snapshot.is_complete() || snapshot.is_notified() { // There is nothing to do in this case. (TransitionToNotifiedByRef::DoNothing, None) } else if snapshot.is_running() { // If the task is running, we mark it as notified, but we should // not submit as the thread currently running the future is // responsible for that. snapshot.set_notified(); (TransitionToNotifiedByRef::DoNothing, Some(snapshot)) } else { // The task is idle and not notified. We should submit a // notification. snapshot.set_notified(); snapshot.ref_inc(); (TransitionToNotifiedByRef::Submit, Some(snapshot)) } }) } /// Transitions the state to `NOTIFIED`, unconditionally increasing the ref /// count. /// /// Returns `true` if the notified bit was transitioned from `0` to `1`; /// otherwise `false.` #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] pub(super) fn transition_to_notified_for_tracing(&self) -> bool { self.fetch_update_action(|mut snapshot| { if snapshot.is_notified() { (false, None) } else { snapshot.set_notified(); snapshot.ref_inc(); (true, Some(snapshot)) } }) } /// Sets the cancelled bit and transitions the state to `NOTIFIED` if idle. /// /// Returns `true` if the task needs to be submitted to the pool for /// execution. pub(super) fn transition_to_notified_and_cancel(&self) -> bool { self.fetch_update_action(|mut snapshot| { if snapshot.is_cancelled() || snapshot.is_complete() { // Aborts to completed or cancelled tasks are no-ops. (false, None) } else if snapshot.is_running() { // If the task is running, we mark it as cancelled. The thread // running the task will notice the cancelled bit when it // stops polling and it will kill the task. // // The set_notified() call is not strictly necessary but it will // in some cases let a wake_by_ref call return without having // to perform a compare_exchange. snapshot.set_notified(); snapshot.set_cancelled(); (false, Some(snapshot)) } else { // The task is idle. We set the cancelled and notified bits and // submit a notification if the notified bit was not already // set. snapshot.set_cancelled(); if !snapshot.is_notified() { snapshot.set_notified(); snapshot.ref_inc(); (true, Some(snapshot)) } else { (false, Some(snapshot)) } } }) } /// Sets the `CANCELLED` bit and attempts to transition to `Running`. /// /// Returns `true` if the transition to `Running` succeeded. pub(super) fn transition_to_shutdown(&self) -> bool { let mut prev = Snapshot(0); let _ = self.fetch_update(|mut snapshot| { prev = snapshot; if snapshot.is_idle() { snapshot.set_running(); } // If the task was not idle, the thread currently running the task // will notice the cancelled bit and cancel it once the poll // completes. snapshot.set_cancelled(); Some(snapshot) }); prev.is_idle() } /// Optimistically tries to swap the state assuming the join handle is /// __immediately__ dropped on spawn. pub(super) fn drop_join_handle_fast(&self) -> Result<(), ()> { use std::sync::atomic::Ordering::Relaxed; // Relaxed is acceptable as if this function is called and succeeds, // then nothing has been done w/ the join handle. // // The moment the join handle is used (polled), the `JOIN_WAKER` flag is // set, at which point the CAS will fail. // // Given this, there is no risk if this operation is reordered. self.val .compare_exchange_weak( INITIAL_STATE, (INITIAL_STATE - REF_ONE) & !JOIN_INTEREST, Release, Relaxed, ) .map(|_| ()) .map_err(|_| ()) } /// Unsets the `JOIN_INTEREST` flag. If `COMPLETE` is not set, the `JOIN_WAKER` /// flag is also unset. /// The returned `TransitionToJoinHandleDrop` indicates whether the `JoinHandle` should drop /// the output of the future or the join waker after the transition. pub(super) fn transition_to_join_handle_dropped(&self) -> TransitionToJoinHandleDrop { self.fetch_update_action(|mut snapshot| { assert!(snapshot.is_join_interested()); let mut transition = TransitionToJoinHandleDrop { drop_waker: false, drop_output: false, }; snapshot.unset_join_interested(); if !snapshot.is_complete() { // If `COMPLETE` is unset we also unset `JOIN_WAKER` to give the // `JoinHandle` exclusive access to the waker following rule 6 in task/mod.rs. // The `JoinHandle` will drop the waker if it has exclusive access // to drop it. snapshot.unset_join_waker(); } else { // If `COMPLETE` is set the task is completed so the `JoinHandle` is responsible // for dropping the output. transition.drop_output = true; } if !snapshot.is_join_waker_set() { // If the `JOIN_WAKER` bit is unset and the `JOIN_HANDLE` has exclusive access to // the join waker and should drop it following this transition. // This might happen in two situations: // 1. The task is not completed and we just unset the `JOIN_WAKer` above in this // function. // 2. The task is completed. In that case the `JOIN_WAKER` bit was already unset // by the runtime during completion. transition.drop_waker = true; } (transition, Some(snapshot)) }) } /// Sets the `JOIN_WAKER` bit. /// /// Returns `Ok` if the bit is set, `Err` otherwise. This operation fails if /// the task has completed. pub(super) fn set_join_waker(&self) -> UpdateResult { self.fetch_update(|curr| { assert!(curr.is_join_interested()); assert!(!curr.is_join_waker_set()); if curr.is_complete() { return None; } let mut next = curr; next.set_join_waker(); Some(next) }) } /// Unsets the `JOIN_WAKER` bit. /// /// Returns `Ok` has been unset, `Err` otherwise. This operation fails if /// the task has completed. pub(super) fn unset_waker(&self) -> UpdateResult { self.fetch_update(|curr| { assert!(curr.is_join_interested()); if curr.is_complete() { return None; } // If the task is completed, this bit may have been unset by // `unset_waker_after_complete`. assert!(curr.is_join_waker_set()); let mut next = curr; next.unset_join_waker(); Some(next) }) } /// Unsets the `JOIN_WAKER` bit unconditionally after task completion. /// /// This operation requires the task to be completed. pub(super) fn unset_waker_after_complete(&self) -> Snapshot { let prev = Snapshot(self.val.fetch_and(!JOIN_WAKER, AcqRel)); assert!(prev.is_complete()); assert!(prev.is_join_waker_set()); Snapshot(prev.0 & !JOIN_WAKER) } pub(super) fn ref_inc(&self) { use std::process; use std::sync::atomic::Ordering::Relaxed; // Using a relaxed ordering is alright here, as knowledge of the // original reference prevents other threads from erroneously deleting // the object. // // As explained in the [Boost documentation][1], Increasing the // reference counter can always be done with memory_order_relaxed: New // references to an object can only be formed from an existing // reference, and passing an existing reference from one thread to // another must already provide any required synchronization. // // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) let prev = self.val.fetch_add(REF_ONE, Relaxed); // If the reference count overflowed, abort. if prev > isize::MAX as usize { process::abort(); } } /// Returns `true` if the task should be released. pub(super) fn ref_dec(&self) -> bool { let prev = Snapshot(self.val.fetch_sub(REF_ONE, AcqRel)); assert!(prev.ref_count() >= 1); prev.ref_count() == 1 } /// Returns `true` if the task should be released. pub(super) fn ref_dec_twice(&self) -> bool { let prev = Snapshot(self.val.fetch_sub(2 * REF_ONE, AcqRel)); assert!(prev.ref_count() >= 2); prev.ref_count() == 2 } fn fetch_update_action(&self, mut f: F) -> T where F: FnMut(Snapshot) -> (T, Option), { let mut curr = self.load(); loop { let (output, next) = f(curr); let next = match next { Some(next) => next, None => return output, }; let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); match res { Ok(_) => return output, Err(actual) => curr = Snapshot(actual), } } } fn fetch_update(&self, mut f: F) -> Result where F: FnMut(Snapshot) -> Option, { let mut curr = self.load(); loop { let next = match f(curr) { Some(next) => next, None => return Err(curr), }; let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); match res { Ok(_) => return Ok(next), Err(actual) => curr = Snapshot(actual), } } } } // ===== impl Snapshot ===== impl Snapshot { /// Returns `true` if the task is in an idle state. pub(super) fn is_idle(self) -> bool { self.0 & (RUNNING | COMPLETE) == 0 } /// Returns `true` if the task has been flagged as notified. pub(super) fn is_notified(self) -> bool { self.0 & NOTIFIED == NOTIFIED } fn unset_notified(&mut self) { self.0 &= !NOTIFIED; } fn set_notified(&mut self) { self.0 |= NOTIFIED; } pub(super) fn is_running(self) -> bool { self.0 & RUNNING == RUNNING } fn set_running(&mut self) { self.0 |= RUNNING; } fn unset_running(&mut self) { self.0 &= !RUNNING; } pub(super) fn is_cancelled(self) -> bool { self.0 & CANCELLED == CANCELLED } fn set_cancelled(&mut self) { self.0 |= CANCELLED; } /// Returns `true` if the task's future has completed execution. pub(super) fn is_complete(self) -> bool { self.0 & COMPLETE == COMPLETE } pub(super) fn is_join_interested(self) -> bool { self.0 & JOIN_INTEREST == JOIN_INTEREST } fn unset_join_interested(&mut self) { self.0 &= !JOIN_INTEREST; } pub(super) fn is_join_waker_set(self) -> bool { self.0 & JOIN_WAKER == JOIN_WAKER } fn set_join_waker(&mut self) { self.0 |= JOIN_WAKER; } fn unset_join_waker(&mut self) { self.0 &= !JOIN_WAKER; } pub(super) fn ref_count(self) -> usize { (self.0 & REF_COUNT_MASK) >> REF_COUNT_SHIFT } fn ref_inc(&mut self) { assert!(self.0 <= isize::MAX as usize); self.0 += REF_ONE; } pub(super) fn ref_dec(&mut self) { assert!(self.ref_count() > 0); self.0 -= REF_ONE; } } impl fmt::Debug for State { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { let snapshot = self.load(); snapshot.fmt(fmt) } } impl fmt::Debug for Snapshot { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Snapshot") .field("is_running", &self.is_running()) .field("is_complete", &self.is_complete()) .field("is_notified", &self.is_notified()) .field("is_cancelled", &self.is_cancelled()) .field("is_join_interested", &self.is_join_interested()) .field("is_join_waker_set", &self.is_join_waker_set()) .field("ref_count", &self.ref_count()) .finish() } } tokio-1.43.1/src/runtime/task/trace/mod.rs000064400000000000000000000262411046102023000165030ustar 00000000000000use crate::loom::sync::Arc; use crate::runtime::context; use crate::runtime::scheduler::{self, current_thread, Inject}; use crate::task::Id; use backtrace::BacktraceFrame; use std::cell::Cell; use std::collections::VecDeque; use std::ffi::c_void; use std::fmt; use std::future::Future; use std::pin::Pin; use std::ptr::{self, NonNull}; use std::task::{self, Poll}; mod symbol; mod tree; use symbol::Symbol; use tree::Tree; use super::{Notified, OwnedTasks, Schedule}; type Backtrace = Vec; type SymbolTrace = Vec; /// The ambient backtracing context. pub(crate) struct Context { /// The address of [`Trace::root`] establishes an upper unwinding bound on /// the backtraces in `Trace`. active_frame: Cell>>, /// The place to stash backtraces. collector: Cell>, } /// A [`Frame`] in an intrusive, doubly-linked tree of [`Frame`]s. struct Frame { /// The location associated with this frame. inner_addr: *const c_void, /// The parent frame, if any. parent: Option>, } /// An tree execution trace. /// /// Traces are captured with [`Trace::capture`], rooted with [`Trace::root`] /// and leaved with [`trace_leaf`]. #[derive(Clone, Debug)] pub(crate) struct Trace { // The linear backtraces that comprise this trace. These linear traces can // be re-knitted into a tree. backtraces: Vec, } pin_project_lite::pin_project! { #[derive(Debug, Clone)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub(crate) struct Root { #[pin] future: T, } } const FAIL_NO_THREAD_LOCAL: &str = "The Tokio thread-local has been destroyed \ as part of shutting down the current \ thread, so collecting a taskdump is not \ possible."; impl Context { pub(crate) const fn new() -> Self { Context { active_frame: Cell::new(None), collector: Cell::new(None), } } /// SAFETY: Callers of this function must ensure that trace frames always /// form a valid linked list. unsafe fn try_with_current(f: F) -> Option where F: FnOnce(&Self) -> R, { crate::runtime::context::with_trace(f) } unsafe fn with_current_frame(f: F) -> R where F: FnOnce(&Cell>>) -> R, { Self::try_with_current(|context| f(&context.active_frame)).expect(FAIL_NO_THREAD_LOCAL) } fn with_current_collector(f: F) -> R where F: FnOnce(&Cell>) -> R, { // SAFETY: This call can only access the collector field, so it cannot // break the trace frame linked list. unsafe { Self::try_with_current(|context| f(&context.collector)).expect(FAIL_NO_THREAD_LOCAL) } } /// Produces `true` if the current task is being traced; otherwise false. pub(crate) fn is_tracing() -> bool { Self::with_current_collector(|maybe_collector| { let collector = maybe_collector.take(); let result = collector.is_some(); maybe_collector.set(collector); result }) } } impl Trace { /// Invokes `f`, returning both its result and the collection of backtraces /// captured at each sub-invocation of [`trace_leaf`]. #[inline(never)] pub(crate) fn capture(f: F) -> (R, Trace) where F: FnOnce() -> R, { let collector = Trace { backtraces: vec![] }; let previous = Context::with_current_collector(|current| current.replace(Some(collector))); let result = f(); let collector = Context::with_current_collector(|current| current.replace(previous)).unwrap(); (result, collector) } /// The root of a trace. #[inline(never)] pub(crate) fn root(future: F) -> Root { Root { future } } pub(crate) fn backtraces(&self) -> &[Backtrace] { &self.backtraces } } /// If this is a sub-invocation of [`Trace::capture`], capture a backtrace. /// /// The captured backtrace will be returned by [`Trace::capture`]. /// /// Invoking this function does nothing when it is not a sub-invocation /// [`Trace::capture`]. // This function is marked `#[inline(never)]` to ensure that it gets a distinct `Frame` in the // backtrace, below which frames should not be included in the backtrace (since they reflect the // internal implementation details of this crate). #[inline(never)] pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> { // Safety: We don't manipulate the current context's active frame. let did_trace = unsafe { Context::try_with_current(|context_cell| { if let Some(mut collector) = context_cell.collector.take() { let mut frames = vec![]; let mut above_leaf = false; if let Some(active_frame) = context_cell.active_frame.get() { let active_frame = active_frame.as_ref(); backtrace::trace(|frame| { let below_root = !ptr::eq(frame.symbol_address(), active_frame.inner_addr); // only capture frames above `Trace::leaf` and below // `Trace::root`. if above_leaf && below_root { frames.push(frame.to_owned().into()); } if ptr::eq(frame.symbol_address(), trace_leaf as *const _) { above_leaf = true; } // only continue unwinding if we're below `Trace::root` below_root }); } collector.backtraces.push(frames); context_cell.collector.set(Some(collector)); true } else { false } }) .unwrap_or(false) }; if did_trace { // Use the same logic that `yield_now` uses to send out wakeups after // the task yields. context::with_scheduler(|scheduler| { if let Some(scheduler) = scheduler { match scheduler { scheduler::Context::CurrentThread(s) => s.defer.defer(cx.waker()), #[cfg(feature = "rt-multi-thread")] scheduler::Context::MultiThread(s) => s.defer.defer(cx.waker()), #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] scheduler::Context::MultiThreadAlt(_) => unimplemented!(), } } }); Poll::Pending } else { Poll::Ready(()) } } impl fmt::Display for Trace { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { Tree::from_trace(self.clone()).fmt(f) } } fn defer R, R>(f: F) -> impl Drop { use std::mem::ManuallyDrop; struct Defer R, R>(ManuallyDrop); impl R, R> Drop for Defer { #[inline(always)] fn drop(&mut self) { unsafe { ManuallyDrop::take(&mut self.0)(); } } } Defer(ManuallyDrop::new(f)) } impl Future for Root { type Output = T::Output; #[inline(never)] fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll { // SAFETY: The context's current frame is restored to its original state // before `frame` is dropped. unsafe { let mut frame = Frame { inner_addr: Self::poll as *const c_void, parent: None, }; Context::with_current_frame(|current| { frame.parent = current.take(); current.set(Some(NonNull::from(&frame))); }); let _restore = defer(|| { Context::with_current_frame(|current| { current.set(frame.parent); }); }); let this = self.project(); this.future.poll(cx) } } } /// Trace and poll all tasks of the `current_thread` runtime. pub(in crate::runtime) fn trace_current_thread( owned: &OwnedTasks>, local: &mut VecDeque>>, injection: &Inject>, ) -> Vec<(Id, Trace)> { // clear the local and injection queues let mut dequeued = Vec::new(); while let Some(task) = local.pop_back() { dequeued.push(task); } while let Some(task) = injection.pop() { dequeued.push(task); } // precondition: We have drained the tasks from the injection queue. trace_owned(owned, dequeued) } cfg_rt_multi_thread! { use crate::loom::sync::Mutex; use crate::runtime::scheduler::multi_thread; use crate::runtime::scheduler::multi_thread::Synced; use crate::runtime::scheduler::inject::Shared; /// Trace and poll all tasks of the `current_thread` runtime. /// /// ## Safety /// /// Must be called with the same `synced` that `injection` was created with. pub(in crate::runtime) unsafe fn trace_multi_thread( owned: &OwnedTasks>, local: &mut multi_thread::queue::Local>, synced: &Mutex, injection: &Shared>, ) -> Vec<(Id, Trace)> { let mut dequeued = Vec::new(); // clear the local queue while let Some(notified) = local.pop() { dequeued.push(notified); } // clear the injection queue let mut synced = synced.lock(); while let Some(notified) = injection.pop(&mut synced.inject) { dequeued.push(notified); } drop(synced); // precondition: we have drained the tasks from the local and injection // queues. trace_owned(owned, dequeued) } } /// Trace the `OwnedTasks`. /// /// # Preconditions /// /// This helper presumes exclusive access to each task. The tasks must not exist /// in any other queue. fn trace_owned(owned: &OwnedTasks, dequeued: Vec>) -> Vec<(Id, Trace)> { let mut tasks = dequeued; // Notify and trace all un-notified tasks. The dequeued tasks are already // notified and so do not need to be re-notified. owned.for_each(|task| { // Notify the task (and thus make it poll-able) and stash it. This fails // if the task is already notified. In these cases, we skip tracing the // task. if let Some(notified) = task.notify_for_tracing() { tasks.push(notified); } // We do not poll tasks here, since we hold a lock on `owned` and the // task may complete and need to remove itself from `owned`. Polling // such a task here would result in a deadlock. }); tasks .into_iter() .map(|task| { let local_notified = owned.assert_owner(task); let id = local_notified.task.id(); let ((), trace) = Trace::capture(|| local_notified.run()); (id, trace) }) .collect() } tokio-1.43.1/src/runtime/task/trace/symbol.rs000064400000000000000000000056711046102023000172350ustar 00000000000000use backtrace::BacktraceSymbol; use std::fmt; use std::hash::{Hash, Hasher}; use std::ptr; /// A symbol in a backtrace. /// /// This wrapper type serves two purposes. The first is that it provides a /// representation of a symbol that can be inserted into hashmaps and hashsets; /// the [`backtrace`] crate does not define [`Hash`], [`PartialEq`], or [`Eq`] /// on [`BacktraceSymbol`], and recommends that users define their own wrapper /// which implements these traits. /// /// Second, this wrapper includes a `parent_hash` field that uniquely /// identifies this symbol's position in its trace. Otherwise, e.g., our code /// would not be able to distinguish between recursive calls of a function at /// different depths. #[derive(Clone)] pub(super) struct Symbol { pub(super) symbol: BacktraceSymbol, pub(super) parent_hash: u64, } impl Hash for Symbol { fn hash(&self, state: &mut H) where H: Hasher, { if let Some(name) = self.symbol.name() { name.as_bytes().hash(state); } if let Some(addr) = self.symbol.addr() { ptr::hash(addr, state); } self.symbol.filename().hash(state); self.symbol.lineno().hash(state); self.symbol.colno().hash(state); self.parent_hash.hash(state); } } impl PartialEq for Symbol { fn eq(&self, other: &Self) -> bool { (self.parent_hash == other.parent_hash) && match (self.symbol.name(), other.symbol.name()) { (None, None) => true, (Some(lhs_name), Some(rhs_name)) => lhs_name.as_bytes() == rhs_name.as_bytes(), _ => false, } && match (self.symbol.addr(), other.symbol.addr()) { (None, None) => true, (Some(lhs_addr), Some(rhs_addr)) => ptr::eq(lhs_addr, rhs_addr), _ => false, } && (self.symbol.filename() == other.symbol.filename()) && (self.symbol.lineno() == other.symbol.lineno()) && (self.symbol.colno() == other.symbol.colno()) } } impl Eq for Symbol {} impl fmt::Display for Symbol { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { if let Some(name) = self.symbol.name() { let name = name.to_string(); let name = if let Some((name, _)) = name.rsplit_once("::") { name } else { &name }; fmt::Display::fmt(&name, f)?; } if let Some(filename) = self.symbol.filename() { f.write_str(" at ")?; filename.to_string_lossy().fmt(f)?; if let Some(lineno) = self.symbol.lineno() { f.write_str(":")?; fmt::Display::fmt(&lineno, f)?; if let Some(colno) = self.symbol.colno() { f.write_str(":")?; fmt::Display::fmt(&colno, f)?; } } } Ok(()) } } tokio-1.43.1/src/runtime/task/trace/tree.rs000064400000000000000000000073701046102023000166650ustar 00000000000000use std::collections::{hash_map::DefaultHasher, HashMap, HashSet}; use std::fmt; use std::hash::{Hash, Hasher}; use super::{Backtrace, Symbol, SymbolTrace, Trace}; /// An adjacency list representation of an execution tree. /// /// This tree provides a convenient intermediate representation for formatting /// [`Trace`] as a tree. pub(super) struct Tree { /// The roots of the trees. /// /// There should only be one root, but the code is robust to multiple roots. roots: HashSet, /// The adjacency list of symbols in the execution tree(s). edges: HashMap>, } impl Tree { /// Constructs a [`Tree`] from [`Trace`] pub(super) fn from_trace(trace: Trace) -> Self { let mut roots: HashSet = HashSet::default(); let mut edges: HashMap> = HashMap::default(); for trace in trace.backtraces { let trace = to_symboltrace(trace); if let Some(first) = trace.first() { roots.insert(first.to_owned()); } let mut trace = trace.into_iter().peekable(); while let Some(frame) = trace.next() { let subframes = edges.entry(frame).or_default(); if let Some(subframe) = trace.peek() { subframes.insert(subframe.clone()); } } } Tree { roots, edges } } /// Produces the sub-symbols of a given symbol. fn consequences(&self, frame: &Symbol) -> Option> { Some(self.edges.get(frame)?.iter()) } /// Format this [`Tree`] as a textual tree. fn display( &self, f: &mut W, root: &Symbol, is_last: bool, prefix: &str, ) -> fmt::Result { let root_fmt = format!("{}", root); let current; let next; if is_last { current = format!("{prefix}└╼\u{a0}{root_fmt}"); next = format!("{}\u{a0}\u{a0}\u{a0}", prefix); } else { current = format!("{prefix}├╼\u{a0}{root_fmt}"); next = format!("{}│\u{a0}\u{a0}", prefix); } write!(f, "{}", { let mut current = current.chars(); current.next().unwrap(); current.next().unwrap(); ¤t.as_str() })?; if let Some(consequences) = self.consequences(root) { let len = consequences.len(); for (i, consequence) in consequences.enumerate() { let is_last = i == len - 1; writeln!(f)?; self.display(f, consequence, is_last, &next)?; } } Ok(()) } } impl fmt::Display for Tree { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { for root in &self.roots { self.display(f, root, true, " ")?; } Ok(()) } } /// Resolve a sequence of [`backtrace::BacktraceFrame`]s into a sequence of /// [`Symbol`]s. fn to_symboltrace(backtrace: Backtrace) -> SymbolTrace { // Resolve the backtrace frames to symbols. let backtrace: Backtrace = { let mut backtrace = backtrace::Backtrace::from(backtrace); backtrace.resolve(); backtrace.into() }; // Accumulate the symbols in descending order into `symboltrace`. let mut symboltrace: SymbolTrace = vec![]; let mut state = DefaultHasher::new(); for frame in backtrace.into_iter().rev() { for symbol in frame.symbols().iter().rev() { let symbol = Symbol { symbol: symbol.clone(), parent_hash: state.finish(), }; symbol.hash(&mut state); symboltrace.push(symbol); } } symboltrace } tokio-1.43.1/src/runtime/task/waker.rs000064400000000000000000000055161046102023000157410ustar 00000000000000use crate::runtime::task::{Header, RawTask, Schedule}; use std::marker::PhantomData; use std::mem::ManuallyDrop; use std::ops; use std::ptr::NonNull; use std::task::{RawWaker, RawWakerVTable, Waker}; pub(super) struct WakerRef<'a, S: 'static> { waker: ManuallyDrop, _p: PhantomData<(&'a Header, S)>, } /// Returns a `WakerRef` which avoids having to preemptively increase the /// refcount if there is no need to do so. pub(super) fn waker_ref(header: &NonNull
) -> WakerRef<'_, S> where S: Schedule, { // `Waker::will_wake` uses the VTABLE pointer as part of the check. This // means that `will_wake` will always return false when using the current // task's waker. (discussion at rust-lang/rust#66281). // // To fix this, we use a single vtable. Since we pass in a reference at this // point and not an *owned* waker, we must ensure that `drop` is never // called on this waker instance. This is done by wrapping it with // `ManuallyDrop` and then never calling drop. let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker(*header))) }; WakerRef { waker, _p: PhantomData, } } impl ops::Deref for WakerRef<'_, S> { type Target = Waker; fn deref(&self) -> &Waker { &self.waker } } cfg_trace! { macro_rules! trace { ($header:expr, $op:expr) => { if let Some(id) = Header::get_tracing_id(&$header) { tracing::trace!( target: "tokio::task::waker", op = $op, task.id = id.into_u64(), ); } } } } cfg_not_trace! { macro_rules! trace { ($header:expr, $op:expr) => { // noop let _ = &$header; } } } unsafe fn clone_waker(ptr: *const ()) -> RawWaker { let header = NonNull::new_unchecked(ptr as *mut Header); trace!(header, "waker.clone"); header.as_ref().state.ref_inc(); raw_waker(header) } unsafe fn drop_waker(ptr: *const ()) { let ptr = NonNull::new_unchecked(ptr as *mut Header); trace!(ptr, "waker.drop"); let raw = RawTask::from_raw(ptr); raw.drop_reference(); } unsafe fn wake_by_val(ptr: *const ()) { let ptr = NonNull::new_unchecked(ptr as *mut Header); trace!(ptr, "waker.wake"); let raw = RawTask::from_raw(ptr); raw.wake_by_val(); } // Wake without consuming the waker unsafe fn wake_by_ref(ptr: *const ()) { let ptr = NonNull::new_unchecked(ptr as *mut Header); trace!(ptr, "waker.wake_by_ref"); let raw = RawTask::from_raw(ptr); raw.wake_by_ref(); } static WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker); fn raw_waker(header: NonNull
) -> RawWaker { let ptr = header.as_ptr() as *const (); RawWaker::new(ptr, &WAKER_VTABLE) } tokio-1.43.1/src/runtime/task_hooks.rs000064400000000000000000000023321046102023000160240ustar 00000000000000use std::marker::PhantomData; impl TaskHooks { pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) { if let Some(f) = self.task_spawn_callback.as_ref() { f(meta) } } } #[derive(Clone)] pub(crate) struct TaskHooks { pub(crate) task_spawn_callback: Option, pub(crate) task_terminate_callback: Option, } /// Task metadata supplied to user-provided hooks for task events. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features #[allow(missing_debug_implementations)] #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] pub struct TaskMeta<'a> { /// The opaque ID of the task. pub(crate) id: super::task::Id, pub(crate) _phantom: PhantomData<&'a ()>, } impl<'a> TaskMeta<'a> { /// Return the opaque ID of the task. #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))] pub fn id(&self) -> super::task::Id { self.id } } /// Runs on specific task-related events pub(crate) type TaskCallback = std::sync::Arc) + Send + Sync>; tokio-1.43.1/src/runtime/tests/inject.rs000064400000000000000000000024571046102023000163050ustar 00000000000000use crate::runtime::scheduler::inject; #[test] fn push_and_pop() { const N: usize = 2; let (inject, mut synced) = inject::Shared::new(); for i in 0..N { assert_eq!(inject.len(), i); let (task, _) = super::unowned(async {}); unsafe { inject.push(&mut synced, task) }; } for i in 0..N { assert_eq!(inject.len(), N - i); assert!(unsafe { inject.pop(&mut synced) }.is_some()); } println!("--------------"); assert!(unsafe { inject.pop(&mut synced) }.is_none()); } #[test] fn push_batch_and_pop() { let (inject, mut inject_synced) = inject::Shared::new(); unsafe { inject.push_batch( &mut inject_synced, (0..10).map(|_| super::unowned(async {}).0), ); assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count()); assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count()); assert_eq!(0, inject.pop_n(&mut inject_synced, 5).count()); } } #[test] fn pop_n_drains_on_drop() { let (inject, mut inject_synced) = inject::Shared::new(); unsafe { inject.push_batch( &mut inject_synced, (0..10).map(|_| super::unowned(async {}).0), ); let _ = inject.pop_n(&mut inject_synced, 10); assert_eq!(inject.len(), 0); } } tokio-1.43.1/src/runtime/tests/loom_blocking.rs000064400000000000000000000054511046102023000176440ustar 00000000000000use crate::runtime::{self, Runtime}; use std::sync::Arc; #[test] fn blocking_shutdown() { loom::model(|| { let v = Arc::new(()); let rt = mk_runtime(1); { let _enter = rt.enter(); for _ in 0..2 { let v = v.clone(); crate::task::spawn_blocking(move || { assert!(1 < Arc::strong_count(&v)); }); } } drop(rt); assert_eq!(1, Arc::strong_count(&v)); }); } #[test] fn spawn_mandatory_blocking_should_always_run() { use crate::runtime::tests::loom_oneshot; loom::model(|| { let rt = runtime::Builder::new_current_thread().build().unwrap(); let (tx, rx) = loom_oneshot::channel(); let _enter = rt.enter(); runtime::spawn_blocking(|| {}); runtime::spawn_mandatory_blocking(move || { let _ = tx.send(()); }) .unwrap(); drop(rt); // This call will deadlock if `spawn_mandatory_blocking` doesn't run. let () = rx.recv(); }); } #[test] fn spawn_mandatory_blocking_should_run_even_when_shutting_down_from_other_thread() { use crate::runtime::tests::loom_oneshot; loom::model(|| { let rt = runtime::Builder::new_current_thread().build().unwrap(); let handle = rt.handle().clone(); // Drop the runtime in a different thread { loom::thread::spawn(move || { drop(rt); }); } let _enter = handle.enter(); let (tx, rx) = loom_oneshot::channel(); let handle = runtime::spawn_mandatory_blocking(move || { let _ = tx.send(()); }); // handle.is_some() means that `spawn_mandatory_blocking` // promised us to run the blocking task if handle.is_some() { // This call will deadlock if `spawn_mandatory_blocking` doesn't run. let () = rx.recv(); } }); } #[test] fn spawn_blocking_when_paused() { use std::time::Duration; loom::model(|| { let rt = crate::runtime::Builder::new_current_thread() .enable_time() .start_paused(true) .build() .unwrap(); let handle = rt.handle(); let _enter = handle.enter(); let a = crate::task::spawn_blocking(|| {}); let b = crate::task::spawn_blocking(|| {}); rt.block_on(crate::time::timeout(Duration::from_millis(1), async move { a.await.expect("blocking task should finish"); b.await.expect("blocking task should finish"); })) .expect("timeout should not trigger"); }); } fn mk_runtime(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) .build() .unwrap() } tokio-1.43.1/src/runtime/tests/loom_current_thread/yield_now.rs000064400000000000000000000015211046102023000230500ustar 00000000000000use crate::runtime::park; use crate::runtime::{self, Runtime}; #[test] fn yield_calls_park_before_scheduling_again() { // Don't need to check all permutations let mut loom = loom::model::Builder::default(); loom.max_permutations = Some(1); loom.check(|| { let rt = mk_runtime(); let jh = rt.spawn(async { let tid = loom::thread::current().id(); let park_count = park::current_thread_park_count(); crate::task::yield_now().await; if tid == loom::thread::current().id() { let new_park_count = park::current_thread_park_count(); assert_eq!(park_count + 1, new_park_count); } }); rt.block_on(jh).unwrap(); }); } fn mk_runtime() -> Runtime { runtime::Builder::new_current_thread().build().unwrap() } tokio-1.43.1/src/runtime/tests/loom_current_thread.rs000064400000000000000000000142571046102023000210710ustar 00000000000000mod yield_now; use crate::loom::sync::atomic::{AtomicUsize, Ordering}; use crate::loom::sync::Arc; use crate::loom::thread; use crate::runtime::{Builder, Runtime}; use crate::sync::oneshot::{self, Receiver}; use crate::task; use std::future::Future; use std::pin::Pin; use std::sync::atomic::Ordering::{Acquire, Release}; use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; fn assert_at_most_num_polls(rt: Arc, at_most_polls: usize) { let (tx, rx) = oneshot::channel(); let num_polls = Arc::new(AtomicUsize::new(0)); rt.spawn(async move { for _ in 0..12 { task::yield_now().await; } tx.send(()).unwrap(); }); rt.block_on(async { BlockedFuture { rx, num_polls: num_polls.clone(), } .await; }); let polls = num_polls.load(Acquire); assert!(polls <= at_most_polls); } #[test] fn block_on_num_polls() { loom::model(|| { // we expect at most 4 number of polls because there are three points at // which we poll the future and an opportunity for a false-positive.. At // any of these points it can be ready: // // - when we fail to steal the parker and we block on a notification // that it is available. // // - when we steal the parker and we schedule the future // // - when the future is woken up and we have ran the max number of tasks // for the current tick or there are no more tasks to run. // // - a thread is notified that the parker is available but a third // thread acquires it before the notified thread can. // let at_most = 4; let rt1 = Arc::new(Builder::new_current_thread().build().unwrap()); let rt2 = rt1.clone(); let rt3 = rt1.clone(); let th1 = thread::spawn(move || assert_at_most_num_polls(rt1, at_most)); let th2 = thread::spawn(move || assert_at_most_num_polls(rt2, at_most)); let th3 = thread::spawn(move || assert_at_most_num_polls(rt3, at_most)); th1.join().unwrap(); th2.join().unwrap(); th3.join().unwrap(); }); } #[test] fn assert_no_unnecessary_polls() { loom::model(|| { // // After we poll outer future, woken should reset to false let rt = Builder::new_current_thread().build().unwrap(); let (tx, rx) = oneshot::channel(); let pending_cnt = Arc::new(AtomicUsize::new(0)); rt.spawn(async move { for _ in 0..24 { task::yield_now().await; } tx.send(()).unwrap(); }); let pending_cnt_clone = pending_cnt.clone(); rt.block_on(async move { // use task::yield_now() to ensure woken set to true // ResetFuture will be polled at most once // Here comes two cases // 1. recv no message from channel, ResetFuture will be polled // but get Pending and we record ResetFuture.pending_cnt ++. // Then when message arrive, ResetFuture returns Ready. So we // expect ResetFuture.pending_cnt = 1 // 2. recv message from channel, ResetFuture returns Ready immediately. // We expect ResetFuture.pending_cnt = 0 task::yield_now().await; ResetFuture { rx, pending_cnt: pending_cnt_clone, } .await; }); let pending_cnt = pending_cnt.load(Acquire); assert!(pending_cnt <= 1); }); } #[test] fn drop_jh_during_schedule() { unsafe fn waker_clone(ptr: *const ()) -> RawWaker { let atomic = unsafe { &*(ptr as *const AtomicUsize) }; atomic.fetch_add(1, Ordering::Relaxed); RawWaker::new(ptr, &VTABLE) } unsafe fn waker_drop(ptr: *const ()) { let atomic = unsafe { &*(ptr as *const AtomicUsize) }; atomic.fetch_sub(1, Ordering::Relaxed); } unsafe fn waker_nop(_ptr: *const ()) {} static VTABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, waker_drop, waker_nop, waker_drop); loom::model(|| { let rt = Builder::new_current_thread().build().unwrap(); let mut jh = rt.spawn(async {}); // Using AbortHandle to increment task refcount. This ensures that the waker is not // destroyed due to the refcount hitting zero. let task_refcnt = jh.abort_handle(); let waker_refcnt = AtomicUsize::new(1); { // Set up the join waker. use std::future::Future; use std::pin::Pin; // SAFETY: Before `waker_refcnt` goes out of scope, this test asserts that the refcnt // has dropped to zero. let join_waker = unsafe { Waker::from_raw(RawWaker::new( (&waker_refcnt) as *const AtomicUsize as *const (), &VTABLE, )) }; assert!(Pin::new(&mut jh) .poll(&mut Context::from_waker(&join_waker)) .is_pending()); } assert_eq!(waker_refcnt.load(Ordering::Relaxed), 1); let bg_thread = loom::thread::spawn(move || drop(jh)); rt.block_on(crate::task::yield_now()); bg_thread.join().unwrap(); assert_eq!(waker_refcnt.load(Ordering::Relaxed), 0); drop(task_refcnt); }); } struct BlockedFuture { rx: Receiver<()>, num_polls: Arc, } impl Future for BlockedFuture { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { self.num_polls.fetch_add(1, Release); match Pin::new(&mut self.rx).poll(cx) { Poll::Pending => Poll::Pending, _ => Poll::Ready(()), } } } struct ResetFuture { rx: Receiver<()>, pending_cnt: Arc, } impl Future for ResetFuture { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match Pin::new(&mut self.rx).poll(cx) { Poll::Pending => { self.pending_cnt.fetch_add(1, Release); Poll::Pending } _ => Poll::Ready(()), } } } tokio-1.43.1/src/runtime/tests/loom_join_set.rs000064400000000000000000000047531046102023000176720ustar 00000000000000use crate::runtime::Builder; use crate::task::JoinSet; #[test] fn test_join_set() { loom::model(|| { let rt = Builder::new_multi_thread() .worker_threads(1) .build() .unwrap(); let mut set = JoinSet::new(); rt.block_on(async { assert_eq!(set.len(), 0); set.spawn(async { () }); assert_eq!(set.len(), 1); set.spawn(async { () }); assert_eq!(set.len(), 2); let () = set.join_next().await.unwrap().unwrap(); assert_eq!(set.len(), 1); set.spawn(async { () }); assert_eq!(set.len(), 2); let () = set.join_next().await.unwrap().unwrap(); assert_eq!(set.len(), 1); let () = set.join_next().await.unwrap().unwrap(); assert_eq!(set.len(), 0); set.spawn(async { () }); assert_eq!(set.len(), 1); }); drop(set); drop(rt); }); } #[test] fn abort_all_during_completion() { use std::sync::{ atomic::{AtomicBool, Ordering::SeqCst}, Arc, }; // These booleans assert that at least one execution had the task complete first, and that at // least one execution had the task be cancelled before it completed. let complete_happened = Arc::new(AtomicBool::new(false)); let cancel_happened = Arc::new(AtomicBool::new(false)); { let complete_happened = complete_happened.clone(); let cancel_happened = cancel_happened.clone(); loom::model(move || { let rt = Builder::new_multi_thread() .worker_threads(1) .build() .unwrap(); let mut set = JoinSet::new(); rt.block_on(async { set.spawn(async { () }); set.abort_all(); match set.join_next().await { Some(Ok(())) => complete_happened.store(true, SeqCst), Some(Err(err)) if err.is_cancelled() => cancel_happened.store(true, SeqCst), Some(Err(err)) => panic!("fail: {}", err), None => { unreachable!("Aborting the task does not remove it from the JoinSet.") } } assert!(matches!(set.join_next().await, None)); }); drop(set); drop(rt); }); } assert!(complete_happened.load(SeqCst)); assert!(cancel_happened.load(SeqCst)); } tokio-1.43.1/src/runtime/tests/loom_local.rs000064400000000000000000000026061046102023000171450ustar 00000000000000use crate::runtime::tests::loom_oneshot as oneshot; use crate::runtime::Builder; use crate::task::LocalSet; use std::task::Poll; /// Waking a runtime will attempt to push a task into a queue of notifications /// in the runtime, however the tasks in such a queue usually have a reference /// to the runtime itself. This means that if they are not properly removed at /// runtime shutdown, this will cause a memory leak. /// /// This test verifies that waking something during shutdown of a `LocalSet` does /// not result in tasks lingering in the queue once shutdown is complete. This /// is verified using loom's leak finder. #[test] fn wake_during_shutdown() { loom::model(|| { let rt = Builder::new_current_thread().build().unwrap(); let ls = LocalSet::new(); let (send, recv) = oneshot::channel(); ls.spawn_local(async move { let mut send = Some(send); let () = std::future::poll_fn(|cx| { if let Some(send) = send.take() { send.send(cx.waker().clone()); } Poll::Pending }) .await; }); let handle = loom::thread::spawn(move || { let waker = recv.recv(); waker.wake(); }); ls.block_on(&rt, crate::task::yield_now()); drop(ls); handle.join().unwrap(); drop(rt); }); } tokio-1.43.1/src/runtime/tests/loom_multi_thread/queue.rs000064400000000000000000000114121046102023000216530ustar 00000000000000use crate::runtime::scheduler::multi_thread::{queue, Stats}; use crate::runtime::tests::{unowned, NoopSchedule}; use loom::thread; use std::cell::RefCell; fn new_stats() -> Stats { Stats::new(&crate::runtime::WorkerMetrics::new()) } #[test] fn basic() { loom::model(|| { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..3 { if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } while local.pop().is_some() { n += 1; } } n }); let mut n = 0; for _ in 0..2 { for _ in 0..2 { let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } if local.pop().is_some() { n += 1; } // Push another task let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); while local.pop().is_some() { n += 1; } } n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); assert_eq!(6, n); }); } #[test] fn steal_overflow() { loom::model(|| { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } while local.pop().is_some() { n += 1; } n }); let mut n = 0; // push a task, pop a task let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); if local.pop().is_some() { n += 1; } for _ in 0..6 { let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } n += th.join().unwrap(); while local.pop().is_some() { n += 1; } n += inject.borrow_mut().drain(..).count(); assert_eq!(7, n); }); } #[test] fn multi_stealer() { const NUM_TASKS: usize = 5; fn steal_tasks(steal: queue::Steal) -> usize { let mut stats = new_stats(); let (_, mut local) = queue::local(); if steal.steal_into(&mut local, &mut stats).is_none() { return 0; } let mut n = 1; while local.pop().is_some() { n += 1; } n } loom::model(|| { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); // Push work for _ in 0..NUM_TASKS { let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } let th1 = { let steal = steal.clone(); thread::spawn(move || steal_tasks(steal)) }; let th2 = thread::spawn(move || steal_tasks(steal)); let mut n = 0; while local.pop().is_some() { n += 1; } n += inject.borrow_mut().drain(..).count(); n += th1.join().unwrap(); n += th2.join().unwrap(); assert_eq!(n, NUM_TASKS); }); } #[test] fn chained_steal() { loom::model(|| { let mut stats = new_stats(); let (s1, mut l1) = queue::local(); let (s2, mut l2) = queue::local(); let inject = RefCell::new(vec![]); // Load up some tasks for _ in 0..4 { let (task, _) = unowned(async {}); l1.push_back_or_overflow(task, &inject, &mut stats); let (task, _) = unowned(async {}); l2.push_back_or_overflow(task, &inject, &mut stats); } // Spawn a task to steal from **our** queue let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); s1.steal_into(&mut local, &mut stats); while local.pop().is_some() {} }); // Drain our tasks, then attempt to steal while l1.pop().is_some() {} s2.steal_into(&mut l1, &mut stats); th.join().unwrap(); while l1.pop().is_some() {} while l2.pop().is_some() {} }); } tokio-1.43.1/src/runtime/tests/loom_multi_thread/shutdown.rs000064400000000000000000000014301046102023000224010ustar 00000000000000use crate::runtime::{Builder, Handle}; #[test] fn join_handle_cancel_on_shutdown() { let mut builder = loom::model::Builder::new(); builder.preemption_bound = Some(2); builder.check(|| { use futures::future::FutureExt; let rt = Builder::new_multi_thread() .worker_threads(2) .build() .unwrap(); let handle = rt.block_on(async move { Handle::current() }); let jh1 = handle.spawn(futures::future::pending::<()>()); drop(rt); let jh2 = handle.spawn(futures::future::pending::<()>()); let err1 = jh1.now_or_never().unwrap().unwrap_err(); let err2 = jh2.now_or_never().unwrap().unwrap_err(); assert!(err1.is_cancelled()); assert!(err2.is_cancelled()); }); } tokio-1.43.1/src/runtime/tests/loom_multi_thread/yield_now.rs000064400000000000000000000020001046102023000225110ustar 00000000000000use crate::runtime::park; use crate::runtime::tests::loom_oneshot as oneshot; use crate::runtime::{self, Runtime}; #[test] fn yield_calls_park_before_scheduling_again() { // Don't need to check all permutations let mut loom = loom::model::Builder::default(); loom.max_permutations = Some(1); loom.check(|| { let rt = mk_runtime(2); let (tx, rx) = oneshot::channel::<()>(); rt.spawn(async { let tid = loom::thread::current().id(); let park_count = park::current_thread_park_count(); crate::task::yield_now().await; if tid == loom::thread::current().id() { let new_park_count = park::current_thread_park_count(); assert_eq!(park_count + 1, new_park_count); } tx.send(()); }); rx.recv(); }); } fn mk_runtime(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) .build() .unwrap() } tokio-1.43.1/src/runtime/tests/loom_multi_thread.rs000064400000000000000000000261301046102023000205320ustar 00000000000000mod queue; mod shutdown; mod yield_now; /// Full runtime loom tests. These are heavy tests and take significant time to /// run on CI. /// /// Use `LOOM_MAX_PREEMPTIONS=1` to do a "quick" run as a smoke test. /// /// In order to speed up the C use crate::runtime::tests::loom_oneshot as oneshot; use crate::runtime::{self, Runtime}; use crate::{spawn, task}; use tokio_test::assert_ok; use loom::sync::atomic::{AtomicBool, AtomicUsize}; use loom::sync::Arc; use pin_project_lite::pin_project; use std::future::{poll_fn, Future}; use std::pin::Pin; use std::sync::atomic::Ordering::{Relaxed, SeqCst}; use std::task::{ready, Context, Poll}; mod atomic_take { use loom::sync::atomic::AtomicBool; use std::mem::MaybeUninit; use std::sync::atomic::Ordering::SeqCst; pub(super) struct AtomicTake { inner: MaybeUninit, taken: AtomicBool, } impl AtomicTake { pub(super) fn new(value: T) -> Self { Self { inner: MaybeUninit::new(value), taken: AtomicBool::new(false), } } pub(super) fn take(&self) -> Option { // safety: Only one thread will see the boolean change from false // to true, so that thread is able to take the value. match self.taken.fetch_or(true, SeqCst) { false => unsafe { Some(std::ptr::read(self.inner.as_ptr())) }, true => None, } } } impl Drop for AtomicTake { fn drop(&mut self) { drop(self.take()); } } } #[derive(Clone)] struct AtomicOneshot { value: std::sync::Arc>>, } impl AtomicOneshot { fn new(sender: oneshot::Sender) -> Self { Self { value: std::sync::Arc::new(atomic_take::AtomicTake::new(sender)), } } fn assert_send(&self, value: T) { self.value.take().unwrap().send(value); } } /// Tests are divided into groups to make the runs faster on CI. mod group_a { use super::*; #[test] fn racy_shutdown() { loom::model(|| { let pool = mk_pool(1); // here's the case we want to exercise: // // a worker that still has tasks in its local queue gets sent to the blocking pool (due to // block_in_place). the blocking pool is shut down, so drops the worker. the worker's // shutdown method never gets run. // // we do this by spawning two tasks on one worker, the first of which does block_in_place, // and then immediately drop the pool. pool.spawn(track(async { crate::task::block_in_place(|| {}); })); pool.spawn(track(async {})); drop(pool); }); } #[test] fn pool_multi_spawn() { loom::model(|| { let pool = mk_pool(2); let c1 = Arc::new(AtomicUsize::new(0)); let (tx, rx) = oneshot::channel(); let tx1 = AtomicOneshot::new(tx); // Spawn a task let c2 = c1.clone(); let tx2 = tx1.clone(); pool.spawn(track(async move { spawn(track(async move { if 1 == c1.fetch_add(1, Relaxed) { tx1.assert_send(()); } })); })); // Spawn a second task pool.spawn(track(async move { spawn(track(async move { if 1 == c2.fetch_add(1, Relaxed) { tx2.assert_send(()); } })); })); rx.recv(); }); } fn only_blocking_inner(first_pending: bool) { loom::model(move || { let pool = mk_pool(1); let (block_tx, block_rx) = oneshot::channel(); pool.spawn(track(async move { crate::task::block_in_place(move || { block_tx.send(()); }); if first_pending { task::yield_now().await } })); block_rx.recv(); drop(pool); }); } #[test] fn only_blocking_without_pending() { only_blocking_inner(false) } #[test] fn only_blocking_with_pending() { only_blocking_inner(true) } } mod group_b { use super::*; fn blocking_and_regular_inner(first_pending: bool) { const NUM: usize = 3; loom::model(move || { let pool = mk_pool(1); let cnt = Arc::new(AtomicUsize::new(0)); let (block_tx, block_rx) = oneshot::channel(); let (done_tx, done_rx) = oneshot::channel(); let done_tx = AtomicOneshot::new(done_tx); pool.spawn(track(async move { crate::task::block_in_place(move || { block_tx.send(()); }); if first_pending { task::yield_now().await } })); for _ in 0..NUM { let cnt = cnt.clone(); let done_tx = done_tx.clone(); pool.spawn(track(async move { if NUM == cnt.fetch_add(1, Relaxed) + 1 { done_tx.assert_send(()); } })); } done_rx.recv(); block_rx.recv(); drop(pool); }); } #[test] fn blocking_and_regular() { blocking_and_regular_inner(false); } #[test] fn blocking_and_regular_with_pending() { blocking_and_regular_inner(true); } #[test] fn join_output() { loom::model(|| { let rt = mk_pool(1); rt.block_on(async { let t = crate::spawn(track(async { "hello" })); let out = assert_ok!(t.await); assert_eq!("hello", out.into_inner()); }); }); } #[test] fn poll_drop_handle_then_drop() { loom::model(|| { let rt = mk_pool(1); rt.block_on(async move { let mut t = crate::spawn(track(async { "hello" })); poll_fn(|cx| { let _ = Pin::new(&mut t).poll(cx); Poll::Ready(()) }) .await; }); }) } #[test] fn complete_block_on_under_load() { loom::model(|| { let pool = mk_pool(1); pool.block_on(async { // Trigger a re-schedule crate::spawn(track(async { for _ in 0..2 { task::yield_now().await; } })); gated2(true).await }); }); } #[test] fn shutdown_with_notification() { use crate::sync::oneshot; loom::model(|| { let rt = mk_pool(2); let (done_tx, done_rx) = oneshot::channel::<()>(); rt.spawn(track(async move { let (tx, rx) = oneshot::channel::<()>(); crate::spawn(async move { crate::task::spawn_blocking(move || { let _ = tx.send(()); }); let _ = done_rx.await; }); let _ = rx.await; let _ = done_tx.send(()); })); }); } } mod group_c { use super::*; #[test] fn pool_shutdown() { loom::model(|| { let pool = mk_pool(2); pool.spawn(track(async move { gated2(true).await; })); pool.spawn(track(async move { gated2(false).await; })); drop(pool); }); } } mod group_d { use super::*; #[test] fn pool_multi_notify() { loom::model(|| { let pool = mk_pool(2); let c1 = Arc::new(AtomicUsize::new(0)); let (done_tx, done_rx) = oneshot::channel(); let done_tx1 = AtomicOneshot::new(done_tx); let done_tx2 = done_tx1.clone(); // Spawn a task let c2 = c1.clone(); pool.spawn(track(async move { multi_gated().await; if 1 == c1.fetch_add(1, Relaxed) { done_tx1.assert_send(()); } })); // Spawn a second task pool.spawn(track(async move { multi_gated().await; if 1 == c2.fetch_add(1, Relaxed) { done_tx2.assert_send(()); } })); done_rx.recv(); }); } } fn mk_pool(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) // Set the intervals to avoid tuning logic .event_interval(2) .build() .unwrap() } fn gated2(thread: bool) -> impl Future { use loom::thread; use std::sync::Arc; let gate = Arc::new(AtomicBool::new(false)); let mut fired = false; poll_fn(move |cx| { if !fired { let gate = gate.clone(); let waker = cx.waker().clone(); if thread { thread::spawn(move || { gate.store(true, SeqCst); waker.wake_by_ref(); }); } else { spawn(track(async move { gate.store(true, SeqCst); waker.wake_by_ref(); })); } fired = true; return Poll::Pending; } if gate.load(SeqCst) { Poll::Ready("hello world") } else { Poll::Pending } }) } async fn multi_gated() { struct Gate { waker: loom::future::AtomicWaker, count: AtomicUsize, } let gate = Arc::new(Gate { waker: loom::future::AtomicWaker::new(), count: AtomicUsize::new(0), }); { let gate = gate.clone(); spawn(track(async move { for i in 1..3 { gate.count.store(i, SeqCst); gate.waker.wake(); } })); } poll_fn(move |cx| { gate.waker.register_by_ref(cx.waker()); if gate.count.load(SeqCst) < 2 { Poll::Pending } else { Poll::Ready(()) } }) .await; } fn track(f: T) -> Track { Track { inner: f, arc: Arc::new(()), } } pin_project! { struct Track { #[pin] inner: T, // Arc is used to hook into loom's leak tracking. arc: Arc<()>, } } impl Track { fn into_inner(self) -> T { self.inner } } impl Future for Track { type Output = Track; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); Poll::Ready(Track { inner: ready!(me.inner.poll(cx)), arc: me.arc.clone(), }) } } tokio-1.43.1/src/runtime/tests/loom_multi_thread_alt/queue.rs000064400000000000000000000114121046102023000225130ustar 00000000000000use crate::runtime::scheduler::multi_thread::{queue, Stats}; use crate::runtime::tests::{unowned, NoopSchedule}; use loom::thread; use std::cell::RefCell; fn new_stats() -> Stats { Stats::new(&crate::runtime::WorkerMetrics::new()) } #[test] fn basic() { loom::model(|| { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..3 { if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } while local.pop().is_some() { n += 1; } } n }); let mut n = 0; for _ in 0..2 { for _ in 0..2 { let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } if local.pop().is_some() { n += 1; } // Push another task let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); while local.pop().is_some() { n += 1; } } n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); assert_eq!(6, n); }); } #[test] fn steal_overflow() { loom::model(|| { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } while local.pop().is_some() { n += 1; } n }); let mut n = 0; // push a task, pop a task let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); if local.pop().is_some() { n += 1; } for _ in 0..6 { let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } n += th.join().unwrap(); while local.pop().is_some() { n += 1; } n += inject.borrow_mut().drain(..).count(); assert_eq!(7, n); }); } #[test] fn multi_stealer() { const NUM_TASKS: usize = 5; fn steal_tasks(steal: queue::Steal) -> usize { let mut stats = new_stats(); let (_, mut local) = queue::local(); if steal.steal_into(&mut local, &mut stats).is_none() { return 0; } let mut n = 1; while local.pop().is_some() { n += 1; } n } loom::model(|| { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); // Push work for _ in 0..NUM_TASKS { let (task, _) = unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } let th1 = { let steal = steal.clone(); thread::spawn(move || steal_tasks(steal)) }; let th2 = thread::spawn(move || steal_tasks(steal)); let mut n = 0; while local.pop().is_some() { n += 1; } n += inject.borrow_mut().drain(..).count(); n += th1.join().unwrap(); n += th2.join().unwrap(); assert_eq!(n, NUM_TASKS); }); } #[test] fn chained_steal() { loom::model(|| { let mut stats = new_stats(); let (s1, mut l1) = queue::local(); let (s2, mut l2) = queue::local(); let inject = RefCell::new(vec![]); // Load up some tasks for _ in 0..4 { let (task, _) = unowned(async {}); l1.push_back_or_overflow(task, &inject, &mut stats); let (task, _) = unowned(async {}); l2.push_back_or_overflow(task, &inject, &mut stats); } // Spawn a task to steal from **our** queue let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); s1.steal_into(&mut local, &mut stats); while local.pop().is_some() {} }); // Drain our tasks, then attempt to steal while l1.pop().is_some() {} s2.steal_into(&mut l1, &mut stats); th.join().unwrap(); while l1.pop().is_some() {} while l2.pop().is_some() {} }); } tokio-1.43.1/src/runtime/tests/loom_multi_thread_alt/shutdown.rs000064400000000000000000000014301046102023000232410ustar 00000000000000use crate::runtime::{Builder, Handle}; #[test] fn join_handle_cancel_on_shutdown() { let mut builder = loom::model::Builder::new(); builder.preemption_bound = Some(2); builder.check(|| { use futures::future::FutureExt; let rt = Builder::new_multi_thread() .worker_threads(2) .build() .unwrap(); let handle = rt.block_on(async move { Handle::current() }); let jh1 = handle.spawn(futures::future::pending::<()>()); drop(rt); let jh2 = handle.spawn(futures::future::pending::<()>()); let err1 = jh1.now_or_never().unwrap().unwrap_err(); let err2 = jh2.now_or_never().unwrap().unwrap_err(); assert!(err1.is_cancelled()); assert!(err2.is_cancelled()); }); } tokio-1.43.1/src/runtime/tests/loom_multi_thread_alt/yield_now.rs000064400000000000000000000020121046102023000233540ustar 00000000000000use crate::runtime::park; use crate::runtime::tests::loom_oneshot as oneshot; use crate::runtime::{self, Runtime}; #[test] #[ignore] fn yield_calls_park_before_scheduling_again() { // Don't need to check all permutations let mut loom = loom::model::Builder::default(); loom.max_permutations = Some(1); loom.check(|| { let rt = mk_runtime(2); let (tx, rx) = oneshot::channel::<()>(); rt.spawn(async { let tid = loom::thread::current().id(); let park_count = park::current_thread_park_count(); crate::task::yield_now().await; if tid == loom::thread::current().id() { let new_park_count = park::current_thread_park_count(); assert_eq!(park_count + 1, new_park_count); } tx.send(()); }); rx.recv(); }); } fn mk_runtime(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) .build() .unwrap() } tokio-1.43.1/src/runtime/tests/loom_multi_thread_alt.rs000064400000000000000000000351341046102023000213760ustar 00000000000000#![cfg(tokio_unstable)] mod queue; mod shutdown; mod yield_now; /// Full runtime loom tests. These are heavy tests and take significant time to /// run on CI. /// /// Use `LOOM_MAX_PREEMPTIONS=1` to do a "quick" run as a smoke test. /// /// In order to speed up the C use crate::runtime::tests::loom_oneshot as oneshot; use crate::runtime::{self, Runtime}; use crate::{spawn, task}; use tokio_test::assert_ok; use loom::sync::atomic::{AtomicBool, AtomicUsize}; use loom::sync::Arc; use pin_project_lite::pin_project; use std::future::{poll_fn, Future}; use std::pin::Pin; use std::sync::atomic::Ordering::{Relaxed, SeqCst}; use std::task::{ready, Context, Poll}; mod atomic_take { use loom::sync::atomic::AtomicBool; use std::mem::MaybeUninit; use std::sync::atomic::Ordering::SeqCst; pub(super) struct AtomicTake { inner: MaybeUninit, taken: AtomicBool, } impl AtomicTake { pub(super) fn new(value: T) -> Self { Self { inner: MaybeUninit::new(value), taken: AtomicBool::new(false), } } pub(super) fn take(&self) -> Option { // safety: Only one thread will see the boolean change from false // to true, so that thread is able to take the value. match self.taken.fetch_or(true, SeqCst) { false => unsafe { Some(std::ptr::read(self.inner.as_ptr())) }, true => None, } } } impl Drop for AtomicTake { fn drop(&mut self) { drop(self.take()); } } } #[derive(Clone)] struct AtomicOneshot { value: std::sync::Arc>>, } impl AtomicOneshot { fn new(sender: oneshot::Sender) -> Self { Self { value: std::sync::Arc::new(atomic_take::AtomicTake::new(sender)), } } fn assert_send(&self, value: T) { self.value.take().unwrap().send(value); } } /// Tests are divided into groups to make the runs faster on CI. mod group_a { use super::*; #[test] fn racy_shutdown() { loom::model(|| { let pool = mk_pool(1); // here's the case we want to exercise: // // a worker that still has tasks in its local queue gets sent to the blocking pool (due to // block_in_place). the blocking pool is shut down, so drops the worker. the worker's // shutdown method never gets run. // // we do this by spawning two tasks on one worker, the first of which does block_in_place, // and then immediately drop the pool. pool.spawn(track(async { crate::task::block_in_place(|| {}); })); pool.spawn(track(async {})); drop(pool); }); } #[test] fn pool_multi_spawn() { loom::model(|| { let pool = mk_pool(2); let c1 = Arc::new(AtomicUsize::new(0)); let (tx, rx) = oneshot::channel(); let tx1 = AtomicOneshot::new(tx); // Spawn a task let c2 = c1.clone(); let tx2 = tx1.clone(); pool.spawn(track(async move { spawn(track(async move { if 1 == c1.fetch_add(1, Relaxed) { tx1.assert_send(()); } })); })); // Spawn a second task pool.spawn(track(async move { spawn(track(async move { if 1 == c2.fetch_add(1, Relaxed) { tx2.assert_send(()); } })); })); rx.recv(); }); } fn only_blocking_inner(first_pending: bool) { loom::model(move || { let pool = mk_pool(1); let (block_tx, block_rx) = oneshot::channel(); pool.spawn(track(async move { crate::task::block_in_place(move || { block_tx.send(()); }); if first_pending { task::yield_now().await } })); block_rx.recv(); drop(pool); }); } #[test] fn only_blocking_without_pending() { only_blocking_inner(false) } #[test] fn only_blocking_with_pending() { only_blocking_inner(true) } } mod group_b { use super::*; fn blocking_and_regular_inner(first_pending: bool) { const NUM: usize = 3; loom::model(move || { let pool = mk_pool(1); let cnt = Arc::new(AtomicUsize::new(0)); let (block_tx, block_rx) = oneshot::channel(); let (done_tx, done_rx) = oneshot::channel(); let done_tx = AtomicOneshot::new(done_tx); pool.spawn(track(async move { crate::task::block_in_place(move || { block_tx.send(()); }); if first_pending { task::yield_now().await } })); for _ in 0..NUM { let cnt = cnt.clone(); let done_tx = done_tx.clone(); pool.spawn(track(async move { if NUM == cnt.fetch_add(1, Relaxed) + 1 { done_tx.assert_send(()); } })); } done_rx.recv(); block_rx.recv(); drop(pool); }); } #[test] #[ignore] // TODO: uncomment fn blocking_and_regular_without_pending() { blocking_and_regular_inner(false); } #[test] fn blocking_and_regular_with_pending() { blocking_and_regular_inner(true); } #[test] fn join_output() { loom::model(|| { let rt = mk_pool(1); rt.block_on(async { let t = crate::spawn(track(async { "hello" })); let out = assert_ok!(t.await); assert_eq!("hello", out.into_inner()); }); }); } #[test] fn poll_drop_handle_then_drop() { loom::model(|| { let rt = mk_pool(1); rt.block_on(async move { let mut t = crate::spawn(track(async { "hello" })); poll_fn(|cx| { let _ = Pin::new(&mut t).poll(cx); Poll::Ready(()) }) .await; }); }) } #[test] fn complete_block_on_under_load() { loom::model(|| { let pool = mk_pool(1); pool.block_on(async { // Trigger a re-schedule crate::spawn(track(async { for _ in 0..2 { task::yield_now().await; } })); gated2(true).await }); }); } #[test] fn shutdown_with_notification() { use crate::sync::oneshot; loom::model(|| { let rt = mk_pool(2); let (done_tx, done_rx) = oneshot::channel::<()>(); rt.spawn(track(async move { let (tx, rx) = oneshot::channel::<()>(); crate::spawn(async move { crate::task::spawn_blocking(move || { let _ = tx.send(()); }); let _ = done_rx.await; }); let _ = rx.await; let _ = done_tx.send(()); })); }); } } mod group_c { use super::*; #[test] fn pool_shutdown() { loom::model(|| { let pool = mk_pool(2); pool.spawn(track(async move { gated2(true).await; })); pool.spawn(track(async move { gated2(false).await; })); drop(pool); }); } #[test] fn fill_local_queue() { const NUM_SPAWNS: usize = 3; loom::model(|| { // using std versions here as it is just to control shutdown. let cnt = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let (tx, rx) = oneshot::channel(); let tx = AtomicOneshot::new(tx); let pool = runtime::Builder::new_multi_thread_alt() .worker_threads(2) // Set the intervals to avoid tuning logic .global_queue_interval(61) .local_queue_capacity(1) .build() .unwrap(); for _ in 0..NUM_SPAWNS { let cnt = cnt.clone(); let tx = tx.clone(); pool.spawn(track(async move { if NUM_SPAWNS == 1 + cnt.fetch_add(1, Relaxed) { tx.assert_send(()); } })); } rx.recv(); }); } // This tests a very specific case that happened when a worker has no more // available work to process because a peer is in the process of stealing // (but does not finish stealing), and the worker happens to find more work // from the injection queue *right* before parking. #[test] fn pool_concurrent_park_with_steal_with_inject() { const DEPTH: usize = 4; let mut model = loom::model::Builder::new(); model.expect_explicit_explore = true; model.preemption_bound = Some(3); model.check(|| { let pool = runtime::Builder::new_multi_thread_alt() .worker_threads(2) // Set the intervals to avoid tuning logic .global_queue_interval(61) .local_queue_capacity(DEPTH) .build() .unwrap(); // Use std types to avoid adding backtracking. type Flag = std::sync::Arc; let flag: Flag = Default::default(); let flag1 = flag.clone(); let (tx1, rx1) = oneshot::channel(); async fn task(expect: isize, flag: Flag) { if expect == flag.load(Relaxed) { flag.store(expect + 1, Relaxed); } else { flag.store(-1, Relaxed); loom::skip_branch(); } } pool.spawn(track(async move { let flag = flag1; // First 2 spawned task should be stolen crate::spawn(task(1, flag.clone())); crate::spawn(task(2, flag.clone())); crate::spawn(async move { task(0, flag.clone()).await; tx1.send(()); }); // One to fill the LIFO slot crate::spawn(async move {}); loom::explore(); })); rx1.recv(); if 1 == flag.load(Relaxed) { loom::stop_exploring(); let (tx3, rx3) = oneshot::channel(); pool.spawn(async move { loom::skip_branch(); tx3.send(()); }); pool.spawn(async {}); pool.spawn(async {}); loom::explore(); rx3.recv(); } else { loom::skip_branch(); } }); } } mod group_d { use super::*; #[test] fn pool_multi_notify() { loom::model(|| { let pool = mk_pool(2); let c1 = Arc::new(AtomicUsize::new(0)); let (done_tx, done_rx) = oneshot::channel(); let done_tx1 = AtomicOneshot::new(done_tx); let done_tx2 = done_tx1.clone(); // Spawn a task let c2 = c1.clone(); pool.spawn(track(async move { multi_gated().await; if 1 == c1.fetch_add(1, Relaxed) { done_tx1.assert_send(()); } })); // Spawn a second task pool.spawn(track(async move { multi_gated().await; if 1 == c2.fetch_add(1, Relaxed) { done_tx2.assert_send(()); } })); done_rx.recv(); }); } } fn mk_pool(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread_alt() .worker_threads(num_threads) // Set the intervals to avoid tuning logic .global_queue_interval(61) .build() .unwrap() } fn gated2(thread: bool) -> impl Future { use loom::thread; use std::sync::Arc; let gate = Arc::new(AtomicBool::new(false)); let mut fired = false; poll_fn(move |cx| { if !fired { let gate = gate.clone(); let waker = cx.waker().clone(); if thread { thread::spawn(move || { gate.store(true, SeqCst); waker.wake_by_ref(); }); } else { spawn(track(async move { gate.store(true, SeqCst); waker.wake_by_ref(); })); } fired = true; return Poll::Pending; } if gate.load(SeqCst) { Poll::Ready("hello world") } else { Poll::Pending } }) } async fn multi_gated() { struct Gate { waker: loom::future::AtomicWaker, count: AtomicUsize, } let gate = Arc::new(Gate { waker: loom::future::AtomicWaker::new(), count: AtomicUsize::new(0), }); { let gate = gate.clone(); spawn(track(async move { for i in 1..3 { gate.count.store(i, SeqCst); gate.waker.wake(); } })); } poll_fn(move |cx| { gate.waker.register_by_ref(cx.waker()); if gate.count.load(SeqCst) < 2 { Poll::Pending } else { Poll::Ready(()) } }) .await; } fn track(f: T) -> Track { Track { inner: f, arc: Arc::new(()), } } pin_project! { struct Track { #[pin] inner: T, // Arc is used to hook into loom's leak tracking. arc: Arc<()>, } } impl Track { fn into_inner(self) -> T { self.inner } } impl Future for Track { type Output = Track; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.project(); Poll::Ready(Track { inner: ready!(me.inner.poll(cx)), arc: me.arc.clone(), }) } } tokio-1.43.1/src/runtime/tests/loom_oneshot.rs000064400000000000000000000016161046102023000175320ustar 00000000000000use crate::loom::sync::{Arc, Mutex}; use loom::sync::Notify; pub(crate) fn channel() -> (Sender, Receiver) { let inner = Arc::new(Inner { notify: Notify::new(), value: Mutex::new(None), }); let tx = Sender { inner: inner.clone(), }; let rx = Receiver { inner }; (tx, rx) } pub(crate) struct Sender { inner: Arc>, } pub(crate) struct Receiver { inner: Arc>, } struct Inner { notify: Notify, value: Mutex>, } impl Sender { pub(crate) fn send(self, value: T) { *self.inner.value.lock() = Some(value); self.inner.notify.notify(); } } impl Receiver { pub(crate) fn recv(self) -> T { loop { if let Some(v) = self.inner.value.lock().take() { return v; } self.inner.notify.wait(); } } } tokio-1.43.1/src/runtime/tests/mod.rs000064400000000000000000000044671046102023000156130ustar 00000000000000// Enable dead_code / unreachable_pub here. It has been disabled in lib.rs for // other code when running loom tests. #![cfg_attr(loom, warn(dead_code, unreachable_pub))] use self::noop_scheduler::NoopSchedule; use self::unowned_wrapper::unowned; mod noop_scheduler { use crate::runtime::task::{self, Task, TaskHarnessScheduleHooks}; /// `task::Schedule` implementation that does nothing, for testing. pub(crate) struct NoopSchedule; impl task::Schedule for NoopSchedule { fn release(&self, _task: &Task) -> Option> { None } fn schedule(&self, _task: task::Notified) { unreachable!(); } fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: None, } } } } mod unowned_wrapper { use crate::runtime::task::{Id, JoinHandle, Notified}; use crate::runtime::tests::NoopSchedule; #[cfg(all(tokio_unstable, feature = "tracing"))] pub(crate) fn unowned(task: T) -> (Notified, JoinHandle) where T: std::future::Future + Send + 'static, T::Output: Send + 'static, { use tracing::Instrument; let span = tracing::trace_span!("test_span"); let task = task.instrument(span); let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next()); (task.into_notified(), handle) } #[cfg(not(all(tokio_unstable, feature = "tracing")))] pub(crate) fn unowned(task: T) -> (Notified, JoinHandle) where T: std::future::Future + Send + 'static, T::Output: Send + 'static, { let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next()); (task.into_notified(), handle) } } cfg_loom! { mod loom_blocking; mod loom_current_thread; mod loom_join_set; mod loom_local; mod loom_multi_thread; mod loom_multi_thread_alt; mod loom_oneshot; // Make sure debug assertions are enabled #[cfg(not(debug_assertions))] compile_error!("these tests require debug assertions to be enabled"); } cfg_not_loom! { mod inject; mod queue; #[cfg(not(miri))] mod task_combinations; #[cfg(miri)] mod task; } tokio-1.43.1/src/runtime/tests/queue.rs000064400000000000000000000143421046102023000161510ustar 00000000000000use crate::runtime::scheduler::multi_thread::{queue, Stats}; use std::cell::RefCell; use std::thread; use std::time::Duration; #[allow(unused)] macro_rules! assert_metrics { ($stats:ident, $field:ident == $v:expr) => { #[cfg(target_has_atomic = "64")] { use crate::runtime::WorkerMetrics; use std::sync::atomic::Ordering::Relaxed; let worker = WorkerMetrics::new(); $stats.submit(&worker); let expect = $v; let actual = worker.$field.load(Relaxed); assert!(actual == expect, "expect = {}; actual = {}", expect, actual) } }; } fn new_stats() -> Stats { use crate::runtime::WorkerMetrics; Stats::new(&WorkerMetrics::new()) } #[test] fn fits_256_one_at_a_time() { let (_, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); for _ in 0..256 { let (task, _) = super::unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } cfg_unstable_metrics! { assert_metrics!(stats, overflow_count == 0); } assert!(inject.borrow_mut().pop().is_none()); while local.pop().is_some() {} } #[test] fn fits_256_all_at_once() { let (_, mut local) = queue::local(); let mut tasks = (0..256) .map(|_| super::unowned(async {}).0) .collect::>(); local.push_back(tasks.drain(..)); let mut i = 0; while local.pop().is_some() { i += 1; } assert_eq!(i, 256); } #[test] fn fits_256_all_in_chunks() { let (_, mut local) = queue::local(); let mut tasks = (0..256) .map(|_| super::unowned(async {}).0) .collect::>(); local.push_back(tasks.drain(..10)); local.push_back(tasks.drain(..100)); local.push_back(tasks.drain(..46)); local.push_back(tasks.drain(..100)); let mut i = 0; while local.pop().is_some() { i += 1; } assert_eq!(i, 256); } #[test] fn overflow() { let (_, mut local) = queue::local(); let inject = RefCell::new(vec![]); let mut stats = new_stats(); for _ in 0..257 { let (task, _) = super::unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } cfg_unstable_metrics! { assert_metrics!(stats, overflow_count == 1); } let mut n = 0; n += inject.borrow_mut().drain(..).count(); while local.pop().is_some() { n += 1; } assert_eq!(n, 257); } #[test] fn steal_batch() { let mut stats = new_stats(); let (steal1, mut local1) = queue::local(); let (_, mut local2) = queue::local(); let inject = RefCell::new(vec![]); for _ in 0..4 { let (task, _) = super::unowned(async {}); local1.push_back_or_overflow(task, &inject, &mut stats); } assert!(steal1.steal_into(&mut local2, &mut stats).is_some()); cfg_unstable_metrics! { assert_metrics!(stats, steal_count == 2); } for _ in 0..1 { assert!(local2.pop().is_some()); } assert!(local2.pop().is_none()); for _ in 0..2 { assert!(local1.pop().is_some()); } assert!(local1.pop().is_none()); } const fn normal_or_miri(normal: usize, miri: usize) -> usize { if cfg!(miri) { miri } else { normal } } #[test] fn stress1() { const NUM_ITER: usize = 5; const NUM_STEAL: usize = normal_or_miri(1_000, 10); const NUM_LOCAL: usize = normal_or_miri(1_000, 10); const NUM_PUSH: usize = normal_or_miri(500, 10); const NUM_POP: usize = normal_or_miri(250, 10); let mut stats = new_stats(); for _ in 0..NUM_ITER { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..NUM_STEAL { if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } while local.pop().is_some() { n += 1; } thread::yield_now(); } cfg_unstable_metrics! { assert_metrics!(stats, steal_count == n as _); } n }); let mut n = 0; for _ in 0..NUM_LOCAL { for _ in 0..NUM_PUSH { let (task, _) = super::unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); } for _ in 0..NUM_POP { if local.pop().is_some() { n += 1; } else { break; } } } n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); assert_eq!(n, NUM_LOCAL * NUM_PUSH); } } #[test] fn stress2() { const NUM_ITER: usize = 1; const NUM_TASKS: usize = normal_or_miri(1_000_000, 50); const NUM_STEAL: usize = normal_or_miri(1_000, 10); let mut stats = new_stats(); for _ in 0..NUM_ITER { let (steal, mut local) = queue::local(); let inject = RefCell::new(vec![]); let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; for _ in 0..NUM_STEAL { if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } while local.pop().is_some() { n += 1; } thread::sleep(Duration::from_micros(10)); } n }); let mut num_pop = 0; for i in 0..NUM_TASKS { let (task, _) = super::unowned(async {}); local.push_back_or_overflow(task, &inject, &mut stats); if i % 128 == 0 && local.pop().is_some() { num_pop += 1; } num_pop += inject.borrow_mut().drain(..).count(); } num_pop += th.join().unwrap(); while local.pop().is_some() { num_pop += 1; } num_pop += inject.borrow_mut().drain(..).count(); assert_eq!(num_pop, NUM_TASKS); } } tokio-1.43.1/src/runtime/tests/task.rs000064400000000000000000000235751046102023000157770ustar 00000000000000use crate::runtime::task::{ self, unowned, Id, JoinHandle, OwnedTasks, Schedule, Task, TaskHarnessScheduleHooks, }; use crate::runtime::tests::NoopSchedule; use std::collections::VecDeque; use std::future::Future; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; struct AssertDropHandle { is_dropped: Arc, } impl AssertDropHandle { #[track_caller] fn assert_dropped(&self) { assert!(self.is_dropped.load(Ordering::SeqCst)); } #[track_caller] fn assert_not_dropped(&self) { assert!(!self.is_dropped.load(Ordering::SeqCst)); } } struct AssertDrop { is_dropped: Arc, } impl AssertDrop { fn new() -> (Self, AssertDropHandle) { let shared = Arc::new(AtomicBool::new(false)); ( AssertDrop { is_dropped: shared.clone(), }, AssertDropHandle { is_dropped: shared.clone(), }, ) } } impl Drop for AssertDrop { fn drop(&mut self) { self.is_dropped.store(true, Ordering::SeqCst); } } // A Notified does not shut down on drop, but it is dropped once the ref-count // hits zero. #[test] fn create_drop1() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); drop(notified); handle.assert_not_dropped(); drop(join); handle.assert_dropped(); } #[test] fn create_drop2() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); drop(join); handle.assert_not_dropped(); drop(notified); handle.assert_dropped(); } #[test] fn drop_abort_handle1() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); let abort = join.abort_handle(); drop(join); handle.assert_not_dropped(); drop(notified); handle.assert_not_dropped(); drop(abort); handle.assert_dropped(); } #[test] fn drop_abort_handle2() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); let abort = join.abort_handle(); drop(notified); handle.assert_not_dropped(); drop(abort); handle.assert_not_dropped(); drop(join); handle.assert_dropped(); } #[test] fn drop_abort_handle_clone() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); let abort = join.abort_handle(); let abort_clone = abort.clone(); drop(join); handle.assert_not_dropped(); drop(notified); handle.assert_not_dropped(); drop(abort); handle.assert_not_dropped(); drop(abort_clone); handle.assert_dropped(); } // Shutting down through Notified works #[test] fn create_shutdown1() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); drop(join); handle.assert_not_dropped(); notified.shutdown(); handle.assert_dropped(); } #[test] fn create_shutdown2() { let (ad, handle) = AssertDrop::new(); let (notified, join) = unowned( async { drop(ad); unreachable!() }, NoopSchedule, Id::next(), ); handle.assert_not_dropped(); notified.shutdown(); handle.assert_dropped(); drop(join); } #[test] fn unowned_poll() { let (task, _) = unowned(async {}, NoopSchedule, Id::next()); task.run(); } #[test] fn schedule() { with(|rt| { rt.spawn(async { crate::task::yield_now().await; }); assert_eq!(2, rt.tick()); rt.shutdown(); }) } #[test] fn shutdown() { with(|rt| { rt.spawn(async { loop { crate::task::yield_now().await; } }); rt.tick_max(1); rt.shutdown(); }) } #[test] fn shutdown_immediately() { with(|rt| { rt.spawn(async { loop { crate::task::yield_now().await; } }); rt.shutdown(); }) } // Test for https://github.com/tokio-rs/tokio/issues/6729 #[test] fn spawn_niche_in_task() { use std::future::poll_fn; use std::task::{Context, Poll, Waker}; with(|rt| { let state = Arc::new(Mutex::new(State::new())); let mut subscriber = Subscriber::new(Arc::clone(&state), 1); rt.spawn(async move { subscriber.wait().await; subscriber.wait().await; }); rt.spawn(async move { state.lock().unwrap().set_version(2); state.lock().unwrap().set_version(0); }); rt.tick_max(10); assert!(rt.is_empty()); rt.shutdown(); }); pub(crate) struct Subscriber { state: Arc>, observed_version: u64, waker_key: Option, } impl Subscriber { pub(crate) fn new(state: Arc>, version: u64) -> Self { Self { state, observed_version: version, waker_key: None, } } pub(crate) async fn wait(&mut self) { poll_fn(|cx| { self.state .lock() .unwrap() .poll_update(&mut self.observed_version, &mut self.waker_key, cx) .map(|_| ()) }) .await; } } struct State { version: u64, wakers: Vec, } impl State { pub(crate) fn new() -> Self { Self { version: 1, wakers: Vec::new(), } } pub(crate) fn poll_update( &mut self, observed_version: &mut u64, waker_key: &mut Option, cx: &Context<'_>, ) -> Poll> { if self.version == 0 { *waker_key = None; Poll::Ready(None) } else if *observed_version < self.version { *waker_key = None; *observed_version = self.version; Poll::Ready(Some(())) } else { self.wakers.push(cx.waker().clone()); *waker_key = Some(self.wakers.len()); Poll::Pending } } pub(crate) fn set_version(&mut self, version: u64) { self.version = version; for waker in self.wakers.drain(..) { waker.wake(); } } } } #[test] fn spawn_during_shutdown() { static DID_SPAWN: AtomicBool = AtomicBool::new(false); struct SpawnOnDrop(Runtime); impl Drop for SpawnOnDrop { fn drop(&mut self) { DID_SPAWN.store(true, Ordering::SeqCst); self.0.spawn(async {}); } } with(|rt| { let rt2 = rt.clone(); rt.spawn(async move { let _spawn_on_drop = SpawnOnDrop(rt2); loop { crate::task::yield_now().await; } }); rt.tick_max(1); rt.shutdown(); }); assert!(DID_SPAWN.load(Ordering::SeqCst)); } fn with(f: impl FnOnce(Runtime)) { struct Reset; impl Drop for Reset { fn drop(&mut self) { let _rt = CURRENT.try_lock().unwrap().take(); } } let _reset = Reset; let rt = Runtime(Arc::new(Inner { owned: OwnedTasks::new(16), core: Mutex::new(Core { queue: VecDeque::new(), }), })); *CURRENT.try_lock().unwrap() = Some(rt.clone()); f(rt) } #[derive(Clone)] struct Runtime(Arc); struct Inner { core: Mutex, owned: OwnedTasks, } struct Core { queue: VecDeque>, } static CURRENT: Mutex> = Mutex::new(None); impl Runtime { fn spawn(&self, future: T) -> JoinHandle where T: 'static + Send + Future, T::Output: 'static + Send, { let (handle, notified) = self.0.owned.bind(future, self.clone(), Id::next()); if let Some(notified) = notified { self.schedule(notified); } handle } fn tick(&self) -> usize { self.tick_max(usize::MAX) } fn tick_max(&self, max: usize) -> usize { let mut n = 0; while !self.is_empty() && n < max { let task = self.next_task(); n += 1; let task = self.0.owned.assert_owner(task); task.run(); } n } fn is_empty(&self) -> bool { self.0.core.try_lock().unwrap().queue.is_empty() } fn next_task(&self) -> task::Notified { self.0.core.try_lock().unwrap().queue.pop_front().unwrap() } fn shutdown(&self) { let mut core = self.0.core.try_lock().unwrap(); self.0.owned.close_and_shutdown_all(0); while let Some(task) = core.queue.pop_back() { drop(task); } drop(core); assert!(self.0.owned.is_empty()); } } impl Schedule for Runtime { fn release(&self, task: &Task) -> Option> { self.0.owned.remove(task) } fn schedule(&self, task: task::Notified) { self.0.core.try_lock().unwrap().queue.push_back(task); } fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: None, } } } tokio-1.43.1/src/runtime/tests/task_combinations.rs000064400000000000000000000365501046102023000205410ustar 00000000000000use std::fmt; use std::future::Future; use std::panic; use std::pin::Pin; use std::task::{Context, Poll}; use crate::runtime::task::AbortHandle; use crate::runtime::Builder; use crate::sync::oneshot; use crate::task::JoinHandle; use futures::future::FutureExt; // Enums for each option in the combinations being tested #[derive(Copy, Clone, Debug, PartialEq)] enum CombiRuntime { CurrentThread, Multi1, Multi2, } #[derive(Copy, Clone, Debug, PartialEq)] enum CombiLocalSet { Yes, No, } #[derive(Copy, Clone, Debug, PartialEq)] enum CombiTask { PanicOnRun, PanicOnDrop, PanicOnRunAndDrop, NoPanic, } #[derive(Copy, Clone, Debug, PartialEq)] enum CombiOutput { PanicOnDrop, NoPanic, } #[derive(Copy, Clone, Debug, PartialEq)] enum CombiJoinInterest { Polled, NotPolled, } #[allow(clippy::enum_variant_names)] // we aren't using glob imports #[derive(Copy, Clone, Debug, PartialEq)] enum CombiJoinHandle { DropImmediately = 1, DropFirstPoll = 2, DropAfterNoConsume = 3, DropAfterConsume = 4, } #[derive(Copy, Clone, Debug, PartialEq)] enum CombiAbort { NotAborted = 0, AbortedImmediately = 1, AbortedFirstPoll = 2, AbortedAfterFinish = 3, AbortedAfterConsumeOutput = 4, } #[derive(Copy, Clone, Debug, PartialEq)] enum CombiAbortSource { JoinHandle, AbortHandle, } #[test] #[cfg_attr(panic = "abort", ignore)] fn test_combinations() { let mut rt = &[ CombiRuntime::CurrentThread, CombiRuntime::Multi1, CombiRuntime::Multi2, ][..]; if cfg!(miri) { rt = &[CombiRuntime::CurrentThread]; } let ls = [CombiLocalSet::Yes, CombiLocalSet::No]; let task = [ CombiTask::NoPanic, CombiTask::PanicOnRun, CombiTask::PanicOnDrop, CombiTask::PanicOnRunAndDrop, ]; let output = [CombiOutput::NoPanic, CombiOutput::PanicOnDrop]; let ji = [CombiJoinInterest::Polled, CombiJoinInterest::NotPolled]; let jh = [ CombiJoinHandle::DropImmediately, CombiJoinHandle::DropFirstPoll, CombiJoinHandle::DropAfterNoConsume, CombiJoinHandle::DropAfterConsume, ]; let abort = [ CombiAbort::NotAborted, CombiAbort::AbortedImmediately, CombiAbort::AbortedFirstPoll, CombiAbort::AbortedAfterFinish, CombiAbort::AbortedAfterConsumeOutput, ]; let ah = [ None, Some(CombiJoinHandle::DropImmediately), Some(CombiJoinHandle::DropFirstPoll), Some(CombiJoinHandle::DropAfterNoConsume), Some(CombiJoinHandle::DropAfterConsume), ]; for rt in rt.iter().copied() { for ls in ls.iter().copied() { for task in task.iter().copied() { for output in output.iter().copied() { for ji in ji.iter().copied() { for jh in jh.iter().copied() { for abort in abort.iter().copied() { // abort via join handle --- abort handles // may be dropped at any point for ah in ah.iter().copied() { test_combination( rt, ls, task, output, ji, jh, ah, abort, CombiAbortSource::JoinHandle, ); } // if aborting via AbortHandle, it will // never be dropped. test_combination( rt, ls, task, output, ji, jh, None, abort, CombiAbortSource::AbortHandle, ); } } } } } } } } fn is_debug(_: &T) {} #[allow(clippy::too_many_arguments)] fn test_combination( rt: CombiRuntime, ls: CombiLocalSet, task: CombiTask, output: CombiOutput, ji: CombiJoinInterest, jh: CombiJoinHandle, ah: Option, abort: CombiAbort, abort_src: CombiAbortSource, ) { match (abort_src, ah) { (CombiAbortSource::JoinHandle, _) if (jh as usize) < (abort as usize) => { // join handle dropped prior to abort return; } (CombiAbortSource::AbortHandle, Some(_)) => { // abort handle dropped, we can't abort through the // abort handle return; } _ => {} } if (task == CombiTask::PanicOnDrop) && (output == CombiOutput::PanicOnDrop) { // this causes double panic return; } if (task == CombiTask::PanicOnRunAndDrop) && (abort != CombiAbort::AbortedImmediately) { // this causes double panic return; } is_debug(&rt); is_debug(&ls); is_debug(&task); is_debug(&output); is_debug(&ji); is_debug(&jh); is_debug(&ah); is_debug(&abort); is_debug(&abort_src); // A runtime optionally with a LocalSet struct Rt { rt: crate::runtime::Runtime, ls: Option, } impl Rt { fn new(rt: CombiRuntime, ls: CombiLocalSet) -> Self { let rt = match rt { CombiRuntime::CurrentThread => Builder::new_current_thread().build().unwrap(), CombiRuntime::Multi1 => Builder::new_multi_thread() .worker_threads(1) .build() .unwrap(), CombiRuntime::Multi2 => Builder::new_multi_thread() .worker_threads(2) .build() .unwrap(), }; let ls = match ls { CombiLocalSet::Yes => Some(crate::task::LocalSet::new()), CombiLocalSet::No => None, }; Self { rt, ls } } fn block_on(&self, task: T) -> T::Output where T: Future, { match &self.ls { Some(ls) => ls.block_on(&self.rt, task), None => self.rt.block_on(task), } } fn spawn(&self, task: T) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { match &self.ls { Some(ls) => ls.spawn_local(task), None => self.rt.spawn(task), } } } // The type used for the output of the future struct Output { panic_on_drop: bool, on_drop: Option>, } impl Output { fn disarm(&mut self) { self.panic_on_drop = false; } } impl Drop for Output { fn drop(&mut self) { let _ = self.on_drop.take().unwrap().send(()); if self.panic_on_drop { panic!("Panicking in Output"); } } } // A wrapper around the future that is spawned struct FutWrapper { inner: F, on_drop: Option>, panic_on_drop: bool, } impl Future for FutWrapper { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { unsafe { let me = Pin::into_inner_unchecked(self); let inner = Pin::new_unchecked(&mut me.inner); inner.poll(cx) } } } impl Drop for FutWrapper { fn drop(&mut self) { let _: Result<(), ()> = self.on_drop.take().unwrap().send(()); if self.panic_on_drop { panic!("Panicking in FutWrapper"); } } } // The channels passed to the task struct Signals { on_first_poll: Option>, wait_complete: Option>, on_output_drop: Option>, } // The task we will spawn async fn my_task(mut signal: Signals, task: CombiTask, out: CombiOutput) -> Output { // Signal that we have been polled once let _ = signal.on_first_poll.take().unwrap().send(()); // Wait for a signal, then complete the future let _ = signal.wait_complete.take().unwrap().await; // If the task gets past wait_complete without yielding, then aborts // may not be caught without this yield_now. crate::task::yield_now().await; if task == CombiTask::PanicOnRun || task == CombiTask::PanicOnRunAndDrop { panic!("Panicking in my_task on {:?}", std::thread::current().id()); } Output { panic_on_drop: out == CombiOutput::PanicOnDrop, on_drop: signal.on_output_drop.take(), } } let rt = Rt::new(rt, ls); let (on_first_poll, wait_first_poll) = oneshot::channel(); let (on_complete, wait_complete) = oneshot::channel(); let (on_future_drop, wait_future_drop) = oneshot::channel(); let (on_output_drop, wait_output_drop) = oneshot::channel(); let signal = Signals { on_first_poll: Some(on_first_poll), wait_complete: Some(wait_complete), on_output_drop: Some(on_output_drop), }; // === Spawn task === let mut handle = Some(rt.spawn(FutWrapper { inner: my_task(signal, task, output), on_drop: Some(on_future_drop), panic_on_drop: task == CombiTask::PanicOnDrop || task == CombiTask::PanicOnRunAndDrop, })); // Keep track of whether the task has been killed with an abort let mut aborted = false; // If we want to poll the JoinHandle, do it now if ji == CombiJoinInterest::Polled { assert!( handle.as_mut().unwrap().now_or_never().is_none(), "Polling handle succeeded" ); } // If we are either aborting the task via an abort handle, or dropping via // an abort handle, do that now. let mut abort_handle = if ah.is_some() || abort_src == CombiAbortSource::AbortHandle { handle.as_ref().map(JoinHandle::abort_handle) } else { None }; let do_abort = |abort_handle: &mut Option, join_handle: Option<&mut JoinHandle<_>>| { match abort_src { CombiAbortSource::AbortHandle => abort_handle.take().unwrap().abort(), CombiAbortSource::JoinHandle => join_handle.unwrap().abort(), } }; if abort == CombiAbort::AbortedImmediately { do_abort(&mut abort_handle, handle.as_mut()); aborted = true; } if jh == CombiJoinHandle::DropImmediately { drop(handle.take().unwrap()); } // === Wait for first poll === let got_polled = rt.block_on(wait_first_poll).is_ok(); if !got_polled { // it's possible that we are aborted but still got polled assert!( aborted, "Task completed without ever being polled but was not aborted." ); } if abort == CombiAbort::AbortedFirstPoll { do_abort(&mut abort_handle, handle.as_mut()); aborted = true; } if jh == CombiJoinHandle::DropFirstPoll { drop(handle.take().unwrap()); } if ah == Some(CombiJoinHandle::DropFirstPoll) { drop(abort_handle.take().unwrap()); } // Signal the future that it can return now let _ = on_complete.send(()); // === Wait for future to be dropped === assert!( rt.block_on(wait_future_drop).is_ok(), "The future should always be dropped." ); if abort == CombiAbort::AbortedAfterFinish { // Don't set aborted to true here as the task already finished do_abort(&mut abort_handle, handle.as_mut()); } if jh == CombiJoinHandle::DropAfterNoConsume { if ah == Some(CombiJoinHandle::DropAfterNoConsume) { drop(handle.take().unwrap()); // The runtime will usually have dropped every ref-count at this point, // in which case dropping the AbortHandle drops the output. // // (But it might race and still hold a ref-count) let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { drop(abort_handle.take().unwrap()); })); if panic.is_err() { assert!( (output == CombiOutput::PanicOnDrop) && (!matches!(task, CombiTask::PanicOnRun | CombiTask::PanicOnRunAndDrop)) && !aborted, "Dropping AbortHandle shouldn't panic here" ); } } else { // The runtime will usually have dropped every ref-count at this point, // in which case dropping the JoinHandle drops the output. // // (But it might race and still hold a ref-count) let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { drop(handle.take().unwrap()); })); if panic.is_err() { assert!( (output == CombiOutput::PanicOnDrop) && (!matches!(task, CombiTask::PanicOnRun | CombiTask::PanicOnRunAndDrop)) && !aborted, "Dropping JoinHandle shouldn't panic here" ); } } } // Check whether we drop after consuming the output if jh == CombiJoinHandle::DropAfterConsume { // Using as_mut() to not immediately drop the handle let result = rt.block_on(handle.as_mut().unwrap()); match result { Ok(mut output) => { // Don't panic here. output.disarm(); assert!(!aborted, "Task was aborted but returned output"); } Err(err) if err.is_cancelled() => assert!(aborted, "Cancelled output but not aborted"), Err(err) if err.is_panic() => { assert!( (task == CombiTask::PanicOnRun) || (task == CombiTask::PanicOnDrop) || (task == CombiTask::PanicOnRunAndDrop) || (output == CombiOutput::PanicOnDrop), "Panic but nothing should panic" ); } _ => unreachable!(), } let mut handle = handle.take().unwrap(); if abort == CombiAbort::AbortedAfterConsumeOutput { do_abort(&mut abort_handle, Some(&mut handle)); } drop(handle); if ah == Some(CombiJoinHandle::DropAfterConsume) { drop(abort_handle.take()); } } // The output should have been dropped now. Check whether the output // object was created at all. let output_created = rt.block_on(wait_output_drop).is_ok(); assert_eq!( output_created, (!matches!(task, CombiTask::PanicOnRun | CombiTask::PanicOnRunAndDrop)) && !aborted, "Creation of output object" ); } tokio-1.43.1/src/runtime/thread_id.rs000064400000000000000000000015261046102023000156060ustar 00000000000000use std::num::NonZeroU64; #[derive(Eq, PartialEq, Clone, Copy, Hash, Debug)] pub(crate) struct ThreadId(NonZeroU64); impl ThreadId { pub(crate) fn next() -> Self { use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64}; static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(0); let mut last = NEXT_ID.load(Relaxed); loop { let id = match last.checked_add(1) { Some(id) => id, None => exhausted(), }; match NEXT_ID.compare_exchange_weak(last, id, Relaxed, Relaxed) { Ok(_) => return ThreadId(NonZeroU64::new(id).unwrap()), Err(id) => last = id, } } } } #[cold] #[allow(dead_code)] fn exhausted() -> ! { panic!("failed to generate unique thread ID: bitspace exhausted") } tokio-1.43.1/src/runtime/time/entry.rs000064400000000000000000000616201046102023000157630ustar 00000000000000//! Timer state structures. //! //! This module contains the heart of the intrusive timer implementation, and as //! such the structures inside are full of tricky concurrency and unsafe code. //! //! # Ground rules //! //! The heart of the timer implementation here is the [`TimerShared`] structure, //! shared between the [`TimerEntry`] and the driver. Generally, we permit access //! to [`TimerShared`] ONLY via either 1) a mutable reference to [`TimerEntry`] or //! 2) a held driver lock. //! //! It follows from this that any changes made while holding BOTH 1 and 2 will //! be reliably visible, regardless of ordering. This is because of the `acq/rel` //! fences on the driver lock ensuring ordering with 2, and rust mutable //! reference rules for 1 (a mutable reference to an object can't be passed //! between threads without an `acq/rel` barrier, and same-thread we have local //! happens-before ordering). //! //! # State field //! //! Each timer has a state field associated with it. This field contains either //! the current scheduled time, or a special flag value indicating its state. //! This state can either indicate that the timer is on the 'pending' queue (and //! thus will be fired with an `Ok(())` result soon) or that it has already been //! fired/deregistered. //! //! This single state field allows for code that is firing the timer to //! synchronize with any racing `reset` calls reliably. //! //! # Cached vs true timeouts //! //! To allow for the use case of a timeout that is periodically reset before //! expiration to be as lightweight as possible, we support optimistically //! lock-free timer resets, in the case where a timer is rescheduled to a later //! point than it was originally scheduled for. //! //! This is accomplished by lazily rescheduling timers. That is, we update the //! state field with the true expiration of the timer from the holder of //! the [`TimerEntry`]. When the driver services timers (ie, whenever it's //! walking lists of timers), it checks this "true when" value, and reschedules //! based on it. //! //! We do, however, also need to track what the expiration time was when we //! originally registered the timer; this is used to locate the right linked //! list when the timer is being cancelled. This is referred to as the "cached //! when" internally. //! //! There is of course a race condition between timer reset and timer //! expiration. If the driver fails to observe the updated expiration time, it //! could trigger expiration of the timer too early. However, because //! [`mark_pending`][mark_pending] performs a compare-and-swap, it will identify this race and //! refuse to mark the timer as pending. //! //! [mark_pending]: TimerHandle::mark_pending use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicU64; use crate::loom::sync::atomic::Ordering; use crate::runtime::context; use crate::runtime::scheduler; use crate::sync::AtomicWaker; use crate::time::Instant; use crate::util::linked_list; use std::cell::UnsafeCell as StdUnsafeCell; use std::task::{Context, Poll, Waker}; use std::{marker::PhantomPinned, pin::Pin, ptr::NonNull}; type TimerResult = Result<(), crate::time::error::Error>; const STATE_DEREGISTERED: u64 = u64::MAX; const STATE_PENDING_FIRE: u64 = STATE_DEREGISTERED - 1; const STATE_MIN_VALUE: u64 = STATE_PENDING_FIRE; /// The largest safe integer to use for ticks. /// /// This value should be updated if any other signal values are added above. pub(super) const MAX_SAFE_MILLIS_DURATION: u64 = STATE_MIN_VALUE - 1; /// This structure holds the current shared state of the timer - its scheduled /// time (if registered), or otherwise the result of the timer completing, as /// well as the registered waker. /// /// Generally, the `StateCell` is only permitted to be accessed from two contexts: /// Either a thread holding the corresponding `&mut TimerEntry`, or a thread /// holding the timer driver lock. The write actions on the `StateCell` amount to /// passing "ownership" of the `StateCell` between these contexts; moving a timer /// from the `TimerEntry` to the driver requires _both_ holding the `&mut /// TimerEntry` and the driver lock, while moving it back (firing the timer) /// requires only the driver lock. pub(super) struct StateCell { /// Holds either the scheduled expiration time for this timer, or (if the /// timer has been fired and is unregistered), `u64::MAX`. state: AtomicU64, /// If the timer is fired (an Acquire order read on state shows /// `u64::MAX`), holds the result that should be returned from /// polling the timer. Otherwise, the contents are unspecified and reading /// without holding the driver lock is undefined behavior. result: UnsafeCell, /// The currently-registered waker waker: AtomicWaker, } impl Default for StateCell { fn default() -> Self { Self::new() } } impl std::fmt::Debug for StateCell { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "StateCell({:?})", self.read_state()) } } impl StateCell { fn new() -> Self { Self { state: AtomicU64::new(STATE_DEREGISTERED), result: UnsafeCell::new(Ok(())), waker: AtomicWaker::new(), } } fn is_pending(&self) -> bool { self.state.load(Ordering::Relaxed) == STATE_PENDING_FIRE } /// Returns the current expiration time, or None if not currently scheduled. fn when(&self) -> Option { let cur_state = self.state.load(Ordering::Relaxed); if cur_state == STATE_DEREGISTERED { None } else { Some(cur_state) } } /// If the timer is completed, returns the result of the timer. Otherwise, /// returns None and registers the waker. fn poll(&self, waker: &Waker) -> Poll { // We must register first. This ensures that either `fire` will // observe the new waker, or we will observe a racing fire to have set // the state, or both. self.waker.register_by_ref(waker); self.read_state() } fn read_state(&self) -> Poll { let cur_state = self.state.load(Ordering::Acquire); if cur_state == STATE_DEREGISTERED { // SAFETY: The driver has fired this timer; this involves writing // the result, and then writing (with release ordering) the state // field. Poll::Ready(unsafe { self.result.with(|p| *p) }) } else { Poll::Pending } } /// Marks this timer as being moved to the pending list, if its scheduled /// time is not after `not_after`. /// /// If the timer is scheduled for a time after `not_after`, returns an Err /// containing the current scheduled time. /// /// SAFETY: Must hold the driver lock. unsafe fn mark_pending(&self, not_after: u64) -> Result<(), u64> { // Quick initial debug check to see if the timer is already fired. Since // firing the timer can only happen with the driver lock held, we know // we shouldn't be able to "miss" a transition to a fired state, even // with relaxed ordering. let mut cur_state = self.state.load(Ordering::Relaxed); loop { // improve the error message for things like // https://github.com/tokio-rs/tokio/issues/3675 assert!( cur_state < STATE_MIN_VALUE, "mark_pending called when the timer entry is in an invalid state" ); if cur_state > not_after { break Err(cur_state); } match self.state.compare_exchange_weak( cur_state, STATE_PENDING_FIRE, Ordering::AcqRel, Ordering::Acquire, ) { Ok(_) => break Ok(()), Err(actual_state) => cur_state = actual_state, } } } /// Fires the timer, setting the result to the provided result. /// /// Returns: /// * `Some(waker)` - if fired and a waker needs to be invoked once the /// driver lock is released /// * `None` - if fired and a waker does not need to be invoked, or if /// already fired /// /// SAFETY: The driver lock must be held. unsafe fn fire(&self, result: TimerResult) -> Option { // Quick initial check to see if the timer is already fired. Since // firing the timer can only happen with the driver lock held, we know // we shouldn't be able to "miss" a transition to a fired state, even // with relaxed ordering. let cur_state = self.state.load(Ordering::Relaxed); if cur_state == STATE_DEREGISTERED { return None; } // SAFETY: We assume the driver lock is held and the timer is not // fired, so only the driver is accessing this field. // // We perform a release-ordered store to state below, to ensure this // write is visible before the state update is visible. unsafe { self.result.with_mut(|p| *p = result) }; self.state.store(STATE_DEREGISTERED, Ordering::Release); self.waker.take_waker() } /// Marks the timer as registered (poll will return None) and sets the /// expiration time. /// /// While this function is memory-safe, it should only be called from a /// context holding both `&mut TimerEntry` and the driver lock. fn set_expiration(&self, timestamp: u64) { debug_assert!(timestamp < STATE_MIN_VALUE); // We can use relaxed ordering because we hold the driver lock and will // fence when we release the lock. self.state.store(timestamp, Ordering::Relaxed); } /// Attempts to adjust the timer to a new timestamp. /// /// If the timer has already been fired, is pending firing, or the new /// timestamp is earlier than the old timestamp, (or occasionally /// spuriously) returns Err without changing the timer's state. In this /// case, the timer must be deregistered and re-registered. fn extend_expiration(&self, new_timestamp: u64) -> Result<(), ()> { let mut prior = self.state.load(Ordering::Relaxed); loop { if new_timestamp < prior || prior >= STATE_MIN_VALUE { return Err(()); } match self.state.compare_exchange_weak( prior, new_timestamp, Ordering::AcqRel, Ordering::Acquire, ) { Ok(_) => return Ok(()), Err(true_prior) => prior = true_prior, } } } /// Returns true if the state of this timer indicates that the timer might /// be registered with the driver. This check is performed with relaxed /// ordering, but is conservative - if it returns false, the timer is /// definitely _not_ registered. pub(super) fn might_be_registered(&self) -> bool { self.state.load(Ordering::Relaxed) != u64::MAX } } /// A timer entry. /// /// This is the handle to a timer that is controlled by the requester of the /// timer. As this participates in intrusive data structures, it must be pinned /// before polling. #[derive(Debug)] pub(crate) struct TimerEntry { /// Arc reference to the runtime handle. We can only free the driver after /// deregistering everything from their respective timer wheels. driver: scheduler::Handle, /// Shared inner structure; this is part of an intrusive linked list, and /// therefore other references can exist to it while mutable references to /// Entry exist. /// /// This is manipulated only under the inner mutex. TODO: Can we use loom /// cells for this? inner: StdUnsafeCell>, /// Deadline for the timer. This is used to register on the first /// poll, as we can't register prior to being pinned. deadline: Instant, /// Whether the deadline has been registered. registered: bool, /// Ensure the type is !Unpin _m: std::marker::PhantomPinned, } unsafe impl Send for TimerEntry {} unsafe impl Sync for TimerEntry {} /// An `TimerHandle` is the (non-enforced) "unique" pointer from the driver to the /// timer entry. Generally, at most one `TimerHandle` exists for a timer at a time /// (enforced by the timer state machine). /// /// SAFETY: An `TimerHandle` is essentially a raw pointer, and the usual caveats /// of pointer safety apply. In particular, `TimerHandle` does not itself enforce /// that the timer does still exist; however, normally an `TimerHandle` is created /// immediately before registering the timer, and is consumed when firing the /// timer, to help minimize mistakes. Still, because `TimerHandle` cannot enforce /// memory safety, all operations are unsafe. #[derive(Debug)] pub(crate) struct TimerHandle { inner: NonNull, } pub(super) type EntryList = crate::util::linked_list::LinkedList; /// The shared state structure of a timer. This structure is shared between the /// frontend (`Entry`) and driver backend. /// /// Note that this structure is located inside the `TimerEntry` structure. pub(crate) struct TimerShared { /// The shard id. We should never change it. shard_id: u32, /// A link within the doubly-linked list of timers on a particular level and /// slot. Valid only if state is equal to Registered. /// /// Only accessed under the entry lock. pointers: linked_list::Pointers, /// The expiration time for which this entry is currently registered. /// Generally owned by the driver, but is accessed by the entry when not /// registered. cached_when: AtomicU64, /// Current state. This records whether the timer entry is currently under /// the ownership of the driver, and if not, its current state (not /// complete, fired, error, etc). state: StateCell, _p: PhantomPinned, } unsafe impl Send for TimerShared {} unsafe impl Sync for TimerShared {} impl std::fmt::Debug for TimerShared { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("TimerShared") .field("cached_when", &self.cached_when.load(Ordering::Relaxed)) .field("state", &self.state) .finish() } } generate_addr_of_methods! { impl<> TimerShared { unsafe fn addr_of_pointers(self: NonNull) -> NonNull> { &self.pointers } } } impl TimerShared { pub(super) fn new(shard_id: u32) -> Self { Self { shard_id, cached_when: AtomicU64::new(0), pointers: linked_list::Pointers::new(), state: StateCell::default(), _p: PhantomPinned, } } /// Gets the cached time-of-expiration value. pub(super) fn cached_when(&self) -> u64 { // Cached-when is only accessed under the driver lock, so we can use relaxed self.cached_when.load(Ordering::Relaxed) } /// Gets the true time-of-expiration value, and copies it into the cached /// time-of-expiration value. /// /// SAFETY: Must be called with the driver lock held, and when this entry is /// not in any timer wheel lists. pub(super) unsafe fn sync_when(&self) -> u64 { let true_when = self.true_when(); self.cached_when.store(true_when, Ordering::Relaxed); true_when } /// Sets the cached time-of-expiration value. /// /// SAFETY: Must be called with the driver lock held, and when this entry is /// not in any timer wheel lists. unsafe fn set_cached_when(&self, when: u64) { self.cached_when.store(when, Ordering::Relaxed); } /// Returns the true time-of-expiration value, with relaxed memory ordering. pub(super) fn true_when(&self) -> u64 { self.state.when().expect("Timer already fired") } /// Sets the true time-of-expiration value, even if it is less than the /// current expiration or the timer is deregistered. /// /// SAFETY: Must only be called with the driver lock held and the entry not /// in the timer wheel. pub(super) unsafe fn set_expiration(&self, t: u64) { self.state.set_expiration(t); self.cached_when.store(t, Ordering::Relaxed); } /// Sets the true time-of-expiration only if it is after the current. pub(super) fn extend_expiration(&self, t: u64) -> Result<(), ()> { self.state.extend_expiration(t) } /// Returns a `TimerHandle` for this timer. pub(super) fn handle(&self) -> TimerHandle { TimerHandle { inner: NonNull::from(self), } } /// Returns true if the state of this timer indicates that the timer might /// be registered with the driver. This check is performed with relaxed /// ordering, but is conservative - if it returns false, the timer is /// definitely _not_ registered. pub(super) fn might_be_registered(&self) -> bool { self.state.might_be_registered() } /// Gets the shard id. pub(super) fn shard_id(&self) -> u32 { self.shard_id } } unsafe impl linked_list::Link for TimerShared { type Handle = TimerHandle; type Target = TimerShared; fn as_raw(handle: &Self::Handle) -> NonNull { handle.inner } unsafe fn from_raw(ptr: NonNull) -> Self::Handle { TimerHandle { inner: ptr } } unsafe fn pointers( target: NonNull, ) -> NonNull> { TimerShared::addr_of_pointers(target) } } // ===== impl Entry ===== impl TimerEntry { #[track_caller] pub(crate) fn new(handle: scheduler::Handle, deadline: Instant) -> Self { // Panic if the time driver is not enabled let _ = handle.driver().time(); Self { driver: handle, inner: StdUnsafeCell::new(None), deadline, registered: false, _m: std::marker::PhantomPinned, } } fn is_inner_init(&self) -> bool { unsafe { &*self.inner.get() }.is_some() } // This lazy initialization is for performance purposes. fn inner(&self) -> &TimerShared { let inner = unsafe { &*self.inner.get() }; if inner.is_none() { let shard_size = self.driver.driver().time().inner.get_shard_size(); let shard_id = generate_shard_id(shard_size); unsafe { *self.inner.get() = Some(TimerShared::new(shard_id)); } } return inner.as_ref().unwrap(); } pub(crate) fn deadline(&self) -> Instant { self.deadline } pub(crate) fn is_elapsed(&self) -> bool { self.is_inner_init() && !self.inner().state.might_be_registered() && self.registered } /// Cancels and deregisters the timer. This operation is irreversible. pub(crate) fn cancel(self: Pin<&mut Self>) { // Avoid calling the `clear_entry` method, because it has not been initialized yet. if !self.is_inner_init() { return; } // We need to perform an acq/rel fence with the driver thread, and the // simplest way to do so is to grab the driver lock. // // Why is this necessary? We're about to release this timer's memory for // some other non-timer use. However, we've been doing a bunch of // relaxed (or even non-atomic) writes from the driver thread, and we'll // be doing more from _this thread_ (as this memory is interpreted as // something else). // // It is critical to ensure that, from the point of view of the driver, // those future non-timer writes happen-after the timer is fully fired, // and from the purpose of this thread, the driver's writes all // happen-before we drop the timer. This in turn requires us to perform // an acquire-release barrier in _both_ directions between the driver // and dropping thread. // // The lock acquisition in clear_entry serves this purpose. All of the // driver manipulations happen with the lock held, so we can just take // the lock and be sure that this drop happens-after everything the // driver did so far and happens-before everything the driver does in // the future. While we have the lock held, we also go ahead and // deregister the entry if necessary. unsafe { self.driver().clear_entry(NonNull::from(self.inner())) }; } pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant, reregister: bool) { let this = unsafe { self.as_mut().get_unchecked_mut() }; this.deadline = new_time; this.registered = reregister; let tick = self.driver().time_source().deadline_to_tick(new_time); if self.inner().extend_expiration(tick).is_ok() { return; } if reregister { unsafe { self.driver() .reregister(&self.driver.driver().io, tick, self.inner().into()); } } } pub(crate) fn poll_elapsed( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { assert!( !self.driver().is_shutdown(), "{}", crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR ); if !self.registered { let deadline = self.deadline; self.as_mut().reset(deadline, true); } self.inner().state.poll(cx.waker()) } pub(crate) fn driver(&self) -> &super::Handle { self.driver.driver().time() } #[cfg(all(tokio_unstable, feature = "tracing"))] pub(crate) fn clock(&self) -> &super::Clock { self.driver.driver().clock() } } impl TimerHandle { pub(super) unsafe fn cached_when(&self) -> u64 { unsafe { self.inner.as_ref().cached_when() } } pub(super) unsafe fn sync_when(&self) -> u64 { unsafe { self.inner.as_ref().sync_when() } } pub(super) unsafe fn is_pending(&self) -> bool { unsafe { self.inner.as_ref().state.is_pending() } } /// Forcibly sets the true and cached expiration times to the given tick. /// /// SAFETY: The caller must ensure that the handle remains valid, the driver /// lock is held, and that the timer is not in any wheel linked lists. pub(super) unsafe fn set_expiration(&self, tick: u64) { self.inner.as_ref().set_expiration(tick); } /// Attempts to mark this entry as pending. If the expiration time is after /// `not_after`, however, returns an Err with the current expiration time. /// /// If an `Err` is returned, the `cached_when` value will be updated to this /// new expiration time. /// /// SAFETY: The caller must ensure that the handle remains valid, the driver /// lock is held, and that the timer is not in any wheel linked lists. /// After returning Ok, the entry must be added to the pending list. pub(super) unsafe fn mark_pending(&self, not_after: u64) -> Result<(), u64> { match self.inner.as_ref().state.mark_pending(not_after) { Ok(()) => { // mark this as being on the pending queue in cached_when self.inner.as_ref().set_cached_when(u64::MAX); Ok(()) } Err(tick) => { self.inner.as_ref().set_cached_when(tick); Err(tick) } } } /// Attempts to transition to a terminal state. If the state is already a /// terminal state, does nothing. /// /// Because the entry might be dropped after the state is moved to a /// terminal state, this function consumes the handle to ensure we don't /// access the entry afterwards. /// /// Returns the last-registered waker, if any. /// /// SAFETY: The driver lock must be held while invoking this function, and /// the entry must not be in any wheel linked lists. pub(super) unsafe fn fire(self, completed_state: TimerResult) -> Option { self.inner.as_ref().state.fire(completed_state) } } impl Drop for TimerEntry { fn drop(&mut self) { unsafe { Pin::new_unchecked(self) }.as_mut().cancel(); } } // Generates a shard id. If current thread is a worker thread, we use its worker index as a shard id. // Otherwise, we use a random number generator to obtain the shard id. cfg_rt! { fn generate_shard_id(shard_size: u32) -> u32 { let id = context::with_scheduler(|ctx| match ctx { Some(scheduler::Context::CurrentThread(_ctx)) => 0, #[cfg(feature = "rt-multi-thread")] Some(scheduler::Context::MultiThread(ctx)) => ctx.get_worker_index() as u32, #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] Some(scheduler::Context::MultiThreadAlt(ctx)) => ctx.get_worker_index() as u32, None => context::thread_rng_n(shard_size), }); id % shard_size } } cfg_not_rt! { fn generate_shard_id(shard_size: u32) -> u32 { context::thread_rng_n(shard_size) } } tokio-1.43.1/src/runtime/time/handle.rs000064400000000000000000000037141046102023000160550ustar 00000000000000use crate::runtime::time::TimeSource; use std::fmt; /// Handle to time driver instance. pub(crate) struct Handle { pub(super) time_source: TimeSource, pub(super) inner: super::Inner, } impl Handle { /// Returns the time source associated with this handle. pub(crate) fn time_source(&self) -> &TimeSource { &self.time_source } /// Checks whether the driver has been shutdown. pub(super) fn is_shutdown(&self) -> bool { self.inner.is_shutdown() } /// Track that the driver is being unparked pub(crate) fn unpark(&self) { #[cfg(feature = "test-util")] self.inner .did_wake .store(true, std::sync::atomic::Ordering::SeqCst); } } cfg_not_rt! { impl Handle { /// Tries to get a handle to the current timer. /// /// # Panics /// /// This function panics if there is no current timer set. /// /// It can be triggered when [`Builder::enable_time`] or /// [`Builder::enable_all`] are not included in the builder. /// /// It can also panic whenever a timer is created outside of a /// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, /// since the function is executed outside of the runtime. /// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. /// And this is because wrapping the function on an async makes it lazy, /// and so gets executed inside the runtime successfully without /// panicking. /// /// [`Builder::enable_time`]: crate::runtime::Builder::enable_time /// [`Builder::enable_all`]: crate::runtime::Builder::enable_all #[track_caller] pub(crate) fn current() -> Self { panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR) } } } impl fmt::Debug for Handle { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "Handle") } } tokio-1.43.1/src/runtime/time/mod.rs000064400000000000000000000415541046102023000154050ustar 00000000000000// Currently, rust warns when an unsafe fn contains an unsafe {} block. However, // in the future, this will change to the reverse. For now, suppress this // warning and generally stick with being explicit about unsafety. #![allow(unused_unsafe)] #![cfg_attr(not(feature = "rt"), allow(dead_code))] //! Time driver. mod entry; pub(crate) use entry::TimerEntry; use entry::{EntryList, TimerHandle, TimerShared, MAX_SAFE_MILLIS_DURATION}; mod handle; pub(crate) use self::handle::Handle; use self::wheel::Wheel; mod source; pub(crate) use source::TimeSource; mod wheel; use crate::loom::sync::atomic::{AtomicBool, Ordering}; use crate::loom::sync::{Mutex, RwLock}; use crate::runtime::driver::{self, IoHandle, IoStack}; use crate::time::error::Error; use crate::time::{Clock, Duration}; use crate::util::WakeList; use crate::loom::sync::atomic::AtomicU64; use std::fmt; use std::{num::NonZeroU64, ptr::NonNull}; struct AtomicOptionNonZeroU64(AtomicU64); // A helper type to store the `next_wake`. impl AtomicOptionNonZeroU64 { fn new(val: Option) -> Self { Self(AtomicU64::new(val.map_or(0, NonZeroU64::get))) } fn store(&self, val: Option) { self.0 .store(val.map_or(0, NonZeroU64::get), Ordering::Relaxed); } fn load(&self) -> Option { NonZeroU64::new(self.0.load(Ordering::Relaxed)) } } /// Time implementation that drives [`Sleep`][sleep], [`Interval`][interval], and [`Timeout`][timeout]. /// /// A `Driver` instance tracks the state necessary for managing time and /// notifying the [`Sleep`][sleep] instances once their deadlines are reached. /// /// It is expected that a single instance manages many individual [`Sleep`][sleep] /// instances. The `Driver` implementation is thread-safe and, as such, is able /// to handle callers from across threads. /// /// After creating the `Driver` instance, the caller must repeatedly call `park` /// or `park_timeout`. The time driver will perform no work unless `park` or /// `park_timeout` is called repeatedly. /// /// The driver has a resolution of one millisecond. Any unit of time that falls /// between milliseconds are rounded up to the next millisecond. /// /// When an instance is dropped, any outstanding [`Sleep`][sleep] instance that has not /// elapsed will be notified with an error. At this point, calling `poll` on the /// [`Sleep`][sleep] instance will result in panic. /// /// # Implementation /// /// The time driver is based on the [paper by Varghese and Lauck][paper]. /// /// A hashed timing wheel is a vector of slots, where each slot handles a time /// slice. As time progresses, the timer walks over the slot for the current /// instant, and processes each entry for that slot. When the timer reaches the /// end of the wheel, it starts again at the beginning. /// /// The implementation maintains six wheels arranged in a set of levels. As the /// levels go up, the slots of the associated wheel represent larger intervals /// of time. At each level, the wheel has 64 slots. Each slot covers a range of /// time equal to the wheel at the lower level. At level zero, each slot /// represents one millisecond of time. /// /// The wheels are: /// /// * Level 0: 64 x 1 millisecond slots. /// * Level 1: 64 x 64 millisecond slots. /// * Level 2: 64 x ~4 second slots. /// * Level 3: 64 x ~4 minute slots. /// * Level 4: 64 x ~4 hour slots. /// * Level 5: 64 x ~12 day slots. /// /// When the timer processes entries at level zero, it will notify all the /// `Sleep` instances as their deadlines have been reached. For all higher /// levels, all entries will be redistributed across the wheel at the next level /// down. Eventually, as time progresses, entries with [`Sleep`][sleep] instances will /// either be canceled (dropped) or their associated entries will reach level /// zero and be notified. /// /// [paper]: http://www.cs.columbia.edu/~nahum/w6998/papers/ton97-timing-wheels.pdf /// [sleep]: crate::time::Sleep /// [timeout]: crate::time::Timeout /// [interval]: crate::time::Interval #[derive(Debug)] pub(crate) struct Driver { /// Parker to delegate to. park: IoStack, } /// Timer state shared between `Driver`, `Handle`, and `Registration`. struct Inner { /// The earliest time at which we promise to wake up without unparking. next_wake: AtomicOptionNonZeroU64, /// Sharded Timer wheels. wheels: RwLock, /// Number of entries in the sharded timer wheels. wheels_len: u32, /// True if the driver is being shutdown. pub(super) is_shutdown: AtomicBool, // When `true`, a call to `park_timeout` should immediately return and time // should not advance. One reason for this to be `true` is if the task // passed to `Runtime::block_on` called `task::yield_now()`. // // While it may look racy, it only has any effect when the clock is paused // and pausing the clock is restricted to a single-threaded runtime. #[cfg(feature = "test-util")] did_wake: AtomicBool, } /// Wrapper around the sharded timer wheels. struct ShardedWheel(Box<[Mutex]>); // ===== impl Driver ===== impl Driver { /// Creates a new `Driver` instance that uses `park` to block the current /// thread and `time_source` to get the current time and convert to ticks. /// /// Specifying the source of time is useful when testing. pub(crate) fn new(park: IoStack, clock: &Clock, shards: u32) -> (Driver, Handle) { assert!(shards > 0); let time_source = TimeSource::new(clock); let wheels: Vec<_> = (0..shards) .map(|_| Mutex::new(wheel::Wheel::new())) .collect(); let handle = Handle { time_source, inner: Inner { next_wake: AtomicOptionNonZeroU64::new(None), wheels: RwLock::new(ShardedWheel(wheels.into_boxed_slice())), wheels_len: shards, is_shutdown: AtomicBool::new(false), #[cfg(feature = "test-util")] did_wake: AtomicBool::new(false), }, }; let driver = Driver { park }; (driver, handle) } pub(crate) fn park(&mut self, handle: &driver::Handle) { self.park_internal(handle, None); } pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) { self.park_internal(handle, Some(duration)); } pub(crate) fn shutdown(&mut self, rt_handle: &driver::Handle) { let handle = rt_handle.time(); if handle.is_shutdown() { return; } handle.inner.is_shutdown.store(true, Ordering::SeqCst); // Advance time forward to the end of time. handle.process_at_time(0, u64::MAX); self.park.shutdown(rt_handle); } fn park_internal(&mut self, rt_handle: &driver::Handle, limit: Option) { let handle = rt_handle.time(); assert!(!handle.is_shutdown()); // Finds out the min expiration time to park. let expiration_time = { let mut wheels_lock = rt_handle.time().inner.wheels.write(); let expiration_time = wheels_lock .0 .iter_mut() .filter_map(|wheel| wheel.get_mut().next_expiration_time()) .min(); rt_handle .time() .inner .next_wake .store(next_wake_time(expiration_time)); expiration_time }; match expiration_time { Some(when) => { let now = handle.time_source.now(rt_handle.clock()); // Note that we effectively round up to 1ms here - this avoids // very short-duration microsecond-resolution sleeps that the OS // might treat as zero-length. let mut duration = handle .time_source .tick_to_duration(when.saturating_sub(now)); if duration > Duration::from_millis(0) { if let Some(limit) = limit { duration = std::cmp::min(limit, duration); } self.park_thread_timeout(rt_handle, duration); } else { self.park.park_timeout(rt_handle, Duration::from_secs(0)); } } None => { if let Some(duration) = limit { self.park_thread_timeout(rt_handle, duration); } else { self.park.park(rt_handle); } } } // Process pending timers after waking up handle.process(rt_handle.clock()); } cfg_test_util! { fn park_thread_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) { let handle = rt_handle.time(); let clock = rt_handle.clock(); if clock.can_auto_advance() { self.park.park_timeout(rt_handle, Duration::from_secs(0)); // If the time driver was woken, then the park completed // before the "duration" elapsed (usually caused by a // yield in `Runtime::block_on`). In this case, we don't // advance the clock. if !handle.did_wake() { // Simulate advancing time if let Err(msg) = clock.advance(duration) { panic!("{}", msg); } } } else { self.park.park_timeout(rt_handle, duration); } } } cfg_not_test_util! { fn park_thread_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) { self.park.park_timeout(rt_handle, duration); } } } // Helper function to turn expiration_time into next_wake_time. // Since the `park_timeout` will round up to 1ms for avoiding very // short-duration microsecond-resolution sleeps, we do the same here. // The conversion is as follows // None => None // Some(0) => Some(1) // Some(i) => Some(i) fn next_wake_time(expiration_time: Option) -> Option { expiration_time.and_then(|v| { if v == 0 { NonZeroU64::new(1) } else { NonZeroU64::new(v) } }) } impl Handle { /// Runs timer related logic, and returns the next wakeup time pub(self) fn process(&self, clock: &Clock) { let now = self.time_source().now(clock); // For fairness, randomly select one to start. let shards = self.inner.get_shard_size(); let start = crate::runtime::context::thread_rng_n(shards); self.process_at_time(start, now); } pub(self) fn process_at_time(&self, start: u32, now: u64) { let shards = self.inner.get_shard_size(); let expiration_time = (start..shards + start) .filter_map(|i| self.process_at_sharded_time(i, now)) .min(); self.inner.next_wake.store(next_wake_time(expiration_time)); } // Returns the next wakeup time of this shard. pub(self) fn process_at_sharded_time(&self, id: u32, mut now: u64) -> Option { let mut waker_list = WakeList::new(); let mut wheels_lock = self.inner.wheels.read(); let mut lock = wheels_lock.lock_sharded_wheel(id); if now < lock.elapsed() { // Time went backwards! This normally shouldn't happen as the Rust language // guarantees that an Instant is monotonic, but can happen when running // Linux in a VM on a Windows host due to std incorrectly trusting the // hardware clock to be monotonic. // // See for more information. now = lock.elapsed(); } while let Some(entry) = lock.poll(now) { debug_assert!(unsafe { entry.is_pending() }); // SAFETY: We hold the driver lock, and just removed the entry from any linked lists. if let Some(waker) = unsafe { entry.fire(Ok(())) } { waker_list.push(waker); if !waker_list.can_push() { // Wake a batch of wakers. To avoid deadlock, we must do this with the lock temporarily dropped. drop(lock); drop(wheels_lock); waker_list.wake_all(); wheels_lock = self.inner.wheels.read(); lock = wheels_lock.lock_sharded_wheel(id); } } } let next_wake_up = lock.poll_at(); drop(lock); drop(wheels_lock); waker_list.wake_all(); next_wake_up } /// Removes a registered timer from the driver. /// /// The timer will be moved to the cancelled state. Wakers will _not_ be /// invoked. If the timer is already completed, this function is a no-op. /// /// This function always acquires the driver lock, even if the entry does /// not appear to be registered. /// /// SAFETY: The timer must not be registered with some other driver, and /// `add_entry` must not be called concurrently. pub(self) unsafe fn clear_entry(&self, entry: NonNull) { unsafe { let wheels_lock = self.inner.wheels.read(); let mut lock = wheels_lock.lock_sharded_wheel(entry.as_ref().shard_id()); if entry.as_ref().might_be_registered() { lock.remove(entry); } entry.as_ref().handle().fire(Ok(())); } } /// Removes and re-adds an entry to the driver. /// /// SAFETY: The timer must be either unregistered, or registered with this /// driver. No other threads are allowed to concurrently manipulate the /// timer at all (the current thread should hold an exclusive reference to /// the `TimerEntry`) pub(self) unsafe fn reregister( &self, unpark: &IoHandle, new_tick: u64, entry: NonNull, ) { let waker = unsafe { let wheels_lock = self.inner.wheels.read(); let mut lock = wheels_lock.lock_sharded_wheel(entry.as_ref().shard_id()); // We may have raced with a firing/deregistration, so check before // deregistering. if unsafe { entry.as_ref().might_be_registered() } { lock.remove(entry); } // Now that we have exclusive control of this entry, mint a handle to reinsert it. let entry = entry.as_ref().handle(); if self.is_shutdown() { unsafe { entry.fire(Err(crate::time::error::Error::shutdown())) } } else { entry.set_expiration(new_tick); // Note: We don't have to worry about racing with some other resetting // thread, because add_entry and reregister require exclusive control of // the timer entry. match unsafe { lock.insert(entry) } { Ok(when) => { if self .inner .next_wake .load() .map(|next_wake| when < next_wake.get()) .unwrap_or(true) { unpark.unpark(); } None } Err((entry, crate::time::error::InsertError::Elapsed)) => unsafe { entry.fire(Ok(())) }, } } // Must release lock before invoking waker to avoid the risk of deadlock. }; // The timer was fired synchronously as a result of the reregistration. // Wake the waker; this is needed because we might reset _after_ a poll, // and otherwise the task won't be awoken to poll again. if let Some(waker) = waker { waker.wake(); } } cfg_test_util! { fn did_wake(&self) -> bool { self.inner.did_wake.swap(false, Ordering::SeqCst) } } } // ===== impl Inner ===== impl Inner { // Check whether the driver has been shutdown pub(super) fn is_shutdown(&self) -> bool { self.is_shutdown.load(Ordering::SeqCst) } // Gets the number of shards. fn get_shard_size(&self) -> u32 { self.wheels_len } } impl fmt::Debug for Inner { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Inner").finish() } } // ===== impl ShardedWheel ===== impl ShardedWheel { /// Locks the driver's sharded wheel structure. pub(super) fn lock_sharded_wheel( &self, shard_id: u32, ) -> crate::loom::sync::MutexGuard<'_, Wheel> { let index = shard_id % (self.0.len() as u32); // Safety: This modulo operation ensures that the index is not out of bounds. unsafe { self.0.get_unchecked(index as usize) }.lock() } } #[cfg(test)] mod tests; tokio-1.43.1/src/runtime/time/source.rs000064400000000000000000000022521046102023000161160ustar 00000000000000use super::MAX_SAFE_MILLIS_DURATION; use crate::time::{Clock, Duration, Instant}; /// A structure which handles conversion from Instants to `u64` timestamps. #[derive(Debug)] pub(crate) struct TimeSource { start_time: Instant, } impl TimeSource { pub(crate) fn new(clock: &Clock) -> Self { Self { start_time: clock.now(), } } pub(crate) fn deadline_to_tick(&self, t: Instant) -> u64 { // Round up to the end of a ms self.instant_to_tick(t + Duration::from_nanos(999_999)) } pub(crate) fn instant_to_tick(&self, t: Instant) -> u64 { // round up let dur: Duration = t.saturating_duration_since(self.start_time); let ms = dur .as_millis() .try_into() .unwrap_or(MAX_SAFE_MILLIS_DURATION); ms.min(MAX_SAFE_MILLIS_DURATION) } pub(crate) fn tick_to_duration(&self, t: u64) -> Duration { Duration::from_millis(t) } pub(crate) fn now(&self, clock: &Clock) -> u64 { self.instant_to_tick(clock.now()) } #[cfg(test)] #[allow(dead_code)] pub(super) fn start_time(&self) -> Instant { self.start_time } } tokio-1.43.1/src/runtime/time/tests/mod.rs000064400000000000000000000162761046102023000165520ustar 00000000000000#![cfg(not(target_os = "wasi"))] use std::{task::Context, time::Duration}; #[cfg(not(loom))] use futures::task::noop_waker_ref; use crate::loom::sync::atomic::{AtomicBool, Ordering}; use crate::loom::sync::Arc; use crate::loom::thread; use super::TimerEntry; fn block_on(f: impl std::future::Future) -> T { #[cfg(loom)] return loom::future::block_on(f); #[cfg(not(loom))] { let rt = crate::runtime::Builder::new_current_thread() .build() .unwrap(); rt.block_on(f) } } fn model(f: impl Fn() + Send + Sync + 'static) { #[cfg(loom)] loom::model(f); #[cfg(not(loom))] f(); } fn rt(start_paused: bool) -> crate::runtime::Runtime { crate::runtime::Builder::new_current_thread() .enable_time() .start_paused(start_paused) .build() .unwrap() } #[test] fn single_timer() { model(|| { let rt = rt(false); let handle = rt.handle(); let handle_ = handle.clone(); let jh = thread::spawn(move || { let entry = TimerEntry::new( handle_.inner.clone(), handle_.inner.driver().clock().now() + Duration::from_secs(1), ); pin!(entry); block_on(std::future::poll_fn(|cx| entry.as_mut().poll_elapsed(cx))).unwrap(); }); thread::yield_now(); let time = handle.inner.driver().time(); let clock = handle.inner.driver().clock(); // This may or may not return Some (depending on how it races with the // thread). If it does return None, however, the timer should complete // synchronously. time.process_at_time(0, time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) } #[test] fn drop_timer() { model(|| { let rt = rt(false); let handle = rt.handle(); let handle_ = handle.clone(); let jh = thread::spawn(move || { let entry = TimerEntry::new( handle_.inner.clone(), handle_.inner.driver().clock().now() + Duration::from_secs(1), ); pin!(entry); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); }); thread::yield_now(); let time = handle.inner.driver().time(); let clock = handle.inner.driver().clock(); // advance 2s in the future. time.process_at_time(0, time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) } #[test] fn change_waker() { model(|| { let rt = rt(false); let handle = rt.handle(); let handle_ = handle.clone(); let jh = thread::spawn(move || { let entry = TimerEntry::new( handle_.inner.clone(), handle_.inner.driver().clock().now() + Duration::from_secs(1), ); pin!(entry); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); block_on(std::future::poll_fn(|cx| entry.as_mut().poll_elapsed(cx))).unwrap(); }); thread::yield_now(); let time = handle.inner.driver().time(); let clock = handle.inner.driver().clock(); // advance 2s time.process_at_time(0, time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) } #[test] fn reset_future() { model(|| { let finished_early = Arc::new(AtomicBool::new(false)); let rt = rt(false); let handle = rt.handle(); let handle_ = handle.clone(); let finished_early_ = finished_early.clone(); let start = handle.inner.driver().clock().now(); let jh = thread::spawn(move || { let entry = TimerEntry::new(handle_.inner.clone(), start + Duration::from_secs(1)); pin!(entry); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); entry.as_mut().reset(start + Duration::from_secs(2), true); // shouldn't complete before 2s block_on(std::future::poll_fn(|cx| entry.as_mut().poll_elapsed(cx))).unwrap(); finished_early_.store(true, Ordering::Relaxed); }); thread::yield_now(); let handle = handle.inner.driver().time(); // This may or may not return a wakeup time. handle.process_at_time( 0, handle .time_source() .instant_to_tick(start + Duration::from_millis(1500)), ); assert!(!finished_early.load(Ordering::Relaxed)); handle.process_at_time( 0, handle .time_source() .instant_to_tick(start + Duration::from_millis(2500)), ); jh.join().unwrap(); assert!(finished_early.load(Ordering::Relaxed)); }) } #[cfg(not(loom))] fn normal_or_miri(normal: T, miri: T) -> T { if cfg!(miri) { miri } else { normal } } #[test] #[cfg(not(loom))] fn poll_process_levels() { let rt = rt(true); let handle = rt.handle(); let mut entries = vec![]; for i in 0..normal_or_miri(1024, 64) { let mut entry = Box::pin(TimerEntry::new( handle.inner.clone(), handle.inner.driver().clock().now() + Duration::from_millis(i), )); let _ = entry .as_mut() .poll_elapsed(&mut Context::from_waker(noop_waker_ref())); entries.push(entry); } for t in 1..normal_or_miri(1024, 64) { handle.inner.driver().time().process_at_time(0, t as u64); for (deadline, future) in entries.iter_mut().enumerate() { let mut context = Context::from_waker(noop_waker_ref()); if deadline <= t { assert!(future.as_mut().poll_elapsed(&mut context).is_ready()); } else { assert!(future.as_mut().poll_elapsed(&mut context).is_pending()); } } } } #[test] #[cfg(not(loom))] fn poll_process_levels_targeted() { let mut context = Context::from_waker(noop_waker_ref()); let rt = rt(true); let handle = rt.handle(); let e1 = TimerEntry::new( handle.inner.clone(), handle.inner.driver().clock().now() + Duration::from_millis(193), ); pin!(e1); let handle = handle.inner.driver().time(); handle.process_at_time(0, 62); assert!(e1.as_mut().poll_elapsed(&mut context).is_pending()); handle.process_at_time(0, 192); handle.process_at_time(0, 192); } #[test] #[cfg(not(loom))] fn instant_to_tick_max() { use crate::runtime::time::entry::MAX_SAFE_MILLIS_DURATION; let rt = rt(true); let handle = rt.handle().inner.driver().time(); let start_time = handle.time_source.start_time(); let long_future = start_time + std::time::Duration::from_millis(MAX_SAFE_MILLIS_DURATION + 1); assert!(handle.time_source.instant_to_tick(long_future) <= MAX_SAFE_MILLIS_DURATION); } tokio-1.43.1/src/runtime/time/wheel/level.rs000064400000000000000000000137521046102023000170400ustar 00000000000000use crate::runtime::time::{EntryList, TimerHandle, TimerShared}; use std::{array, fmt, ptr::NonNull}; /// Wheel for a single level in the timer. This wheel contains 64 slots. pub(crate) struct Level { level: usize, /// Bit field tracking which slots currently contain entries. /// /// Using a bit field to track slots that contain entries allows avoiding a /// scan to find entries. This field is updated when entries are added or /// removed from a slot. /// /// The least-significant bit represents slot zero. occupied: u64, /// Slots. We access these via the EntryInner `current_list` as well, so this needs to be an `UnsafeCell`. slot: [EntryList; LEVEL_MULT], } /// Indicates when a slot must be processed next. #[derive(Debug)] pub(crate) struct Expiration { /// The level containing the slot. pub(crate) level: usize, /// The slot index. pub(crate) slot: usize, /// The instant at which the slot needs to be processed. pub(crate) deadline: u64, } /// Level multiplier. /// /// Being a power of 2 is very important. const LEVEL_MULT: usize = 64; impl Level { pub(crate) fn new(level: usize) -> Level { Level { level, occupied: 0, slot: array::from_fn(|_| EntryList::default()), } } /// Finds the slot that needs to be processed next and returns the slot and /// `Instant` at which this slot must be processed. pub(crate) fn next_expiration(&self, now: u64) -> Option { // Use the `occupied` bit field to get the index of the next slot that // needs to be processed. let slot = self.next_occupied_slot(now)?; // From the slot index, calculate the `Instant` at which it needs to be // processed. This value *must* be in the future with respect to `now`. let level_range = level_range(self.level); let slot_range = slot_range(self.level); // Compute the start date of the current level by masking the low bits // of `now` (`level_range` is a power of 2). let level_start = now & !(level_range - 1); let mut deadline = level_start + slot as u64 * slot_range; if deadline <= now { // A timer is in a slot "prior" to the current time. This can occur // because we do not have an infinite hierarchy of timer levels, and // eventually a timer scheduled for a very distant time might end up // being placed in a slot that is beyond the end of all of the // arrays. // // To deal with this, we first limit timers to being scheduled no // more than MAX_DURATION ticks in the future; that is, they're at // most one rotation of the top level away. Then, we force timers // that logically would go into the top+1 level, to instead go into // the top level's slots. // // What this means is that the top level's slots act as a // pseudo-ring buffer, and we rotate around them indefinitely. If we // compute a deadline before now, and it's the top level, it // therefore means we're actually looking at a slot in the future. debug_assert_eq!(self.level, super::NUM_LEVELS - 1); deadline += level_range; } debug_assert!( deadline >= now, "deadline={:016X}; now={:016X}; level={}; lr={:016X}, sr={:016X}, slot={}; occupied={:b}", deadline, now, self.level, level_range, slot_range, slot, self.occupied ); Some(Expiration { level: self.level, slot, deadline, }) } fn next_occupied_slot(&self, now: u64) -> Option { if self.occupied == 0 { return None; } // Get the slot for now using Maths let now_slot = (now / slot_range(self.level)) as usize; let occupied = self.occupied.rotate_right(now_slot as u32); let zeros = occupied.trailing_zeros() as usize; let slot = (zeros + now_slot) % LEVEL_MULT; Some(slot) } pub(crate) unsafe fn add_entry(&mut self, item: TimerHandle) { let slot = slot_for(item.cached_when(), self.level); self.slot[slot].push_front(item); self.occupied |= occupied_bit(slot); } pub(crate) unsafe fn remove_entry(&mut self, item: NonNull) { let slot = slot_for(unsafe { item.as_ref().cached_when() }, self.level); unsafe { self.slot[slot].remove(item) }; if self.slot[slot].is_empty() { // The bit is currently set debug_assert!(self.occupied & occupied_bit(slot) != 0); // Unset the bit self.occupied ^= occupied_bit(slot); } } pub(crate) fn take_slot(&mut self, slot: usize) -> EntryList { self.occupied &= !occupied_bit(slot); std::mem::take(&mut self.slot[slot]) } } impl fmt::Debug for Level { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Level") .field("occupied", &self.occupied) .finish() } } fn occupied_bit(slot: usize) -> u64 { 1 << slot } fn slot_range(level: usize) -> u64 { LEVEL_MULT.pow(level as u32) as u64 } fn level_range(level: usize) -> u64 { LEVEL_MULT as u64 * slot_range(level) } /// Converts a duration (milliseconds) and a level to a slot position. fn slot_for(duration: u64, level: usize) -> usize { ((duration >> (level * 6)) % LEVEL_MULT as u64) as usize } #[cfg(all(test, not(loom)))] mod test { use super::*; #[test] fn test_slot_for() { for pos in 0..64 { assert_eq!(pos as usize, slot_for(pos, 0)); } for level in 1..5 { for pos in level..64 { let a = pos * 64_usize.pow(level as u32); assert_eq!(pos, slot_for(a as u64, level)); } } } } tokio-1.43.1/src/runtime/time/wheel/mod.rs000064400000000000000000000247761046102023000165200ustar 00000000000000use crate::runtime::time::{TimerHandle, TimerShared}; use crate::time::error::InsertError; mod level; pub(crate) use self::level::Expiration; use self::level::Level; use std::{array, ptr::NonNull}; use super::EntryList; /// Timing wheel implementation. /// /// This type provides the hashed timing wheel implementation that backs `Timer` /// and `DelayQueue`. /// /// The structure is generic over `T: Stack`. This allows handling timeout data /// being stored on the heap or in a slab. In order to support the latter case, /// the slab must be passed into each function allowing the implementation to /// lookup timer entries. /// /// See `Timer` documentation for some implementation notes. #[derive(Debug)] pub(crate) struct Wheel { /// The number of milliseconds elapsed since the wheel started. elapsed: u64, /// Timer wheel. /// /// Levels: /// /// * 1 ms slots / 64 ms range /// * 64 ms slots / ~ 4 sec range /// * ~ 4 sec slots / ~ 4 min range /// * ~ 4 min slots / ~ 4 hr range /// * ~ 4 hr slots / ~ 12 day range /// * ~ 12 day slots / ~ 2 yr range levels: Box<[Level; NUM_LEVELS]>, /// Entries queued for firing pending: EntryList, } /// Number of levels. Each level has 64 slots. By using 6 levels with 64 slots /// each, the timer is able to track time up to 2 years into the future with a /// precision of 1 millisecond. const NUM_LEVELS: usize = 6; /// The maximum duration of a `Sleep`. pub(super) const MAX_DURATION: u64 = (1 << (6 * NUM_LEVELS)) - 1; impl Wheel { /// Creates a new timing wheel. pub(crate) fn new() -> Wheel { Wheel { elapsed: 0, levels: Box::new(array::from_fn(Level::new)), pending: EntryList::new(), } } /// Returns the number of milliseconds that have elapsed since the timing /// wheel's creation. pub(crate) fn elapsed(&self) -> u64 { self.elapsed } /// Inserts an entry into the timing wheel. /// /// # Arguments /// /// * `item`: The item to insert into the wheel. /// /// # Return /// /// Returns `Ok` when the item is successfully inserted, `Err` otherwise. /// /// `Err(Elapsed)` indicates that `when` represents an instant that has /// already passed. In this case, the caller should fire the timeout /// immediately. /// /// `Err(Invalid)` indicates an invalid `when` argument as been supplied. /// /// # Safety /// /// This function registers item into an intrusive linked list. The caller /// must ensure that `item` is pinned and will not be dropped without first /// being deregistered. pub(crate) unsafe fn insert( &mut self, item: TimerHandle, ) -> Result { let when = item.sync_when(); if when <= self.elapsed { return Err((item, InsertError::Elapsed)); } // Get the level at which the entry should be stored let level = self.level_for(when); unsafe { self.levels[level].add_entry(item); } debug_assert!({ self.levels[level] .next_expiration(self.elapsed) .map(|e| e.deadline >= self.elapsed) .unwrap_or(true) }); Ok(when) } /// Removes `item` from the timing wheel. pub(crate) unsafe fn remove(&mut self, item: NonNull) { unsafe { let when = item.as_ref().cached_when(); if when == u64::MAX { self.pending.remove(item); } else { debug_assert!( self.elapsed <= when, "elapsed={}; when={}", self.elapsed, when ); let level = self.level_for(when); self.levels[level].remove_entry(item); } } } /// Instant at which to poll. pub(crate) fn poll_at(&self) -> Option { self.next_expiration().map(|expiration| expiration.deadline) } /// Advances the timer up to the instant represented by `now`. pub(crate) fn poll(&mut self, now: u64) -> Option { loop { if let Some(handle) = self.pending.pop_back() { return Some(handle); } match self.next_expiration() { Some(ref expiration) if expiration.deadline <= now => { self.process_expiration(expiration); self.set_elapsed(expiration.deadline); } _ => { // in this case the poll did not indicate an expiration // _and_ we were not able to find a next expiration in // the current list of timers. advance to the poll's // current time and do nothing else. self.set_elapsed(now); break; } } } self.pending.pop_back() } /// Returns the instant at which the next timeout expires. fn next_expiration(&self) -> Option { if !self.pending.is_empty() { // Expire immediately as we have things pending firing return Some(Expiration { level: 0, slot: 0, deadline: self.elapsed, }); } // Check all levels for (level_num, level) in self.levels.iter().enumerate() { if let Some(expiration) = level.next_expiration(self.elapsed) { // There cannot be any expirations at a higher level that happen // before this one. debug_assert!(self.no_expirations_before(level_num + 1, expiration.deadline)); return Some(expiration); } } None } /// Returns the tick at which this timer wheel next needs to perform some /// processing, or None if there are no timers registered. pub(super) fn next_expiration_time(&self) -> Option { self.next_expiration().map(|ex| ex.deadline) } /// Used for debug assertions fn no_expirations_before(&self, start_level: usize, before: u64) -> bool { let mut res = true; for level in &self.levels[start_level..] { if let Some(e2) = level.next_expiration(self.elapsed) { if e2.deadline < before { res = false; } } } res } /// iteratively find entries that are between the wheel's current /// time and the expiration time. for each in that population either /// queue it for notification (in the case of the last level) or tier /// it down to the next level (in all other cases). pub(crate) fn process_expiration(&mut self, expiration: &Expiration) { // Note that we need to take _all_ of the entries off the list before // processing any of them. This is important because it's possible that // those entries might need to be reinserted into the same slot. // // This happens only on the highest level, when an entry is inserted // more than MAX_DURATION into the future. When this happens, we wrap // around, and process some entries a multiple of MAX_DURATION before // they actually need to be dropped down a level. We then reinsert them // back into the same position; we must make sure we don't then process // those entries again or we'll end up in an infinite loop. let mut entries = self.take_entries(expiration); while let Some(item) = entries.pop_back() { if expiration.level == 0 { debug_assert_eq!(unsafe { item.cached_when() }, expiration.deadline); } // Try to expire the entry; this is cheap (doesn't synchronize) if // the timer is not expired, and updates cached_when. match unsafe { item.mark_pending(expiration.deadline) } { Ok(()) => { // Item was expired self.pending.push_front(item); } Err(expiration_tick) => { let level = level_for(expiration.deadline, expiration_tick); unsafe { self.levels[level].add_entry(item); } } } } } fn set_elapsed(&mut self, when: u64) { assert!( self.elapsed <= when, "elapsed={:?}; when={:?}", self.elapsed, when ); if when > self.elapsed { self.elapsed = when; } } /// Obtains the list of entries that need processing for the given expiration. fn take_entries(&mut self, expiration: &Expiration) -> EntryList { self.levels[expiration.level].take_slot(expiration.slot) } fn level_for(&self, when: u64) -> usize { level_for(self.elapsed, when) } } fn level_for(elapsed: u64, when: u64) -> usize { const SLOT_MASK: u64 = (1 << 6) - 1; // Mask in the trailing bits ignored by the level calculation in order to cap // the possible leading zeros let mut masked = elapsed ^ when | SLOT_MASK; if masked >= MAX_DURATION { // Fudge the timer into the top level masked = MAX_DURATION - 1; } let leading_zeros = masked.leading_zeros() as usize; let significant = 63 - leading_zeros; significant / NUM_LEVELS } #[cfg(all(test, not(loom)))] mod test { use super::*; #[test] fn test_level_for() { for pos in 0..64 { assert_eq!(0, level_for(0, pos), "level_for({pos}) -- binary = {pos:b}"); } for level in 1..5 { for pos in level..64 { let a = pos * 64_usize.pow(level as u32); assert_eq!( level, level_for(0, a as u64), "level_for({a}) -- binary = {a:b}" ); if pos > level { let a = a - 1; assert_eq!( level, level_for(0, a as u64), "level_for({a}) -- binary = {a:b}" ); } if pos < 64 { let a = a + 1; assert_eq!( level, level_for(0, a as u64), "level_for({a}) -- binary = {a:b}" ); } } } } } tokio-1.43.1/src/signal/ctrl_c.rs000064400000000000000000000040441046102023000147210ustar 00000000000000#[cfg(unix)] use super::unix::{self as os_impl}; #[cfg(windows)] use super::windows::{self as os_impl}; use std::io; /// Completes when a "ctrl-c" notification is sent to the process. /// /// While signals are handled very differently between Unix and Windows, both /// platforms support receiving a signal on "ctrl-c". This function provides a /// portable API for receiving this notification. /// /// Once the returned future is polled, a listener is registered. The future /// will complete on the first received `ctrl-c` **after** the initial call to /// either `Future::poll` or `.await`. /// /// # Caveats /// /// On Unix platforms, the first time that a `Signal` instance is registered for a /// particular signal kind, an OS signal-handler is installed which replaces the /// default platform behavior when that signal is received, **for the duration of /// the entire process**. /// /// For example, Unix systems will terminate a process by default when it /// receives a signal generated by `"CTRL+C"` on the terminal. But, when a /// `ctrl_c` stream is created to listen for this signal, the time it arrives, /// it will be translated to a stream event, and the process will continue to /// execute. **Even if this `Signal` instance is dropped, subsequent `SIGINT` /// deliveries will end up captured by Tokio, and the default platform behavior /// will NOT be reset**. /// /// Thus, applications should take care to ensure the expected signal behavior /// occurs as expected after listening for specific signals. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal; /// /// #[tokio::main] /// async fn main() { /// println!("waiting for ctrl-c"); /// /// signal::ctrl_c().await.expect("failed to listen for event"); /// /// println!("received ctrl-c event"); /// } /// ``` /// /// Listen in the background: /// /// ```rust,no_run /// tokio::spawn(async move { /// tokio::signal::ctrl_c().await.unwrap(); /// // Your handler here /// }); /// ``` pub async fn ctrl_c() -> io::Result<()> { os_impl::ctrl_c()?.recv().await; Ok(()) } tokio-1.43.1/src/signal/mod.rs000064400000000000000000000046351046102023000142400ustar 00000000000000//! Asynchronous signal handling for Tokio. //! //! Note that signal handling is in general a very tricky topic and should be //! used with great care. This crate attempts to implement 'best practice' for //! signal handling, but it should be evaluated for your own applications' needs //! to see if it's suitable. //! //! There are some fundamental limitations of this crate documented on the OS //! specific structures, as well. //! //! # Examples //! //! Print on "ctrl-c" notification. //! //! ```rust,no_run //! use tokio::signal; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! signal::ctrl_c().await?; //! println!("ctrl-c received!"); //! Ok(()) //! } //! ``` //! //! Wait for `SIGHUP` on Unix //! //! ```rust,no_run //! # #[cfg(unix)] { //! use tokio::signal::unix::{signal, SignalKind}; //! //! #[tokio::main] //! async fn main() -> Result<(), Box> { //! // An infinite stream of hangup signals. //! let mut stream = signal(SignalKind::hangup())?; //! //! // Print whenever a HUP signal is received //! loop { //! stream.recv().await; //! println!("got signal HUP"); //! } //! } //! # } //! ``` use crate::sync::watch::Receiver; use std::task::{Context, Poll}; #[cfg(feature = "signal")] mod ctrl_c; #[cfg(feature = "signal")] pub use ctrl_c::ctrl_c; pub(crate) mod registry; mod os { #[cfg(unix)] pub(crate) use super::unix::{OsExtraData, OsStorage}; #[cfg(windows)] pub(crate) use super::windows::{OsExtraData, OsStorage}; } pub mod unix; pub mod windows; mod reusable_box; use self::reusable_box::ReusableBoxFuture; #[derive(Debug)] struct RxFuture { inner: ReusableBoxFuture>, } async fn make_future(mut rx: Receiver<()>) -> Receiver<()> { rx.changed().await.expect("signal sender went away"); rx } impl RxFuture { fn new(rx: Receiver<()>) -> Self { Self { inner: ReusableBoxFuture::new(make_future(rx)), } } async fn recv(&mut self) -> Option<()> { use std::future::poll_fn; poll_fn(|cx| self.poll_recv(cx)).await } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { match self.inner.poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(rx) => { self.inner.set(make_future(rx)); Poll::Ready(Some(())) } } } } tokio-1.43.1/src/signal/registry.rs000064400000000000000000000167741046102023000153400ustar 00000000000000use crate::signal::os::{OsExtraData, OsStorage}; use crate::sync::watch; use crate::util::once_cell::OnceCell; use std::ops; use std::sync::atomic::{AtomicBool, Ordering}; pub(crate) type EventId = usize; /// State for a specific event, whether a notification is pending delivery, /// and what listeners are registered. #[derive(Debug)] pub(crate) struct EventInfo { pending: AtomicBool, tx: watch::Sender<()>, } impl Default for EventInfo { fn default() -> Self { let (tx, _rx) = watch::channel(()); Self { pending: AtomicBool::new(false), tx, } } } /// An interface for retrieving the `EventInfo` for a particular `eventId`. pub(crate) trait Storage { /// Gets the `EventInfo` for `id` if it exists. fn event_info(&self, id: EventId) -> Option<&EventInfo>; /// Invokes `f` once for each defined `EventInfo` in this storage. fn for_each<'a, F>(&'a self, f: F) where F: FnMut(&'a EventInfo); } impl Storage for Vec { fn event_info(&self, id: EventId) -> Option<&EventInfo> { self.get(id) } fn for_each<'a, F>(&'a self, f: F) where F: FnMut(&'a EventInfo), { self.iter().for_each(f); } } /// An interface for initializing a type. Useful for situations where we cannot /// inject a configured instance in the constructor of another type. pub(crate) trait Init { fn init() -> Self; } /// Manages and distributes event notifications to any registered listeners. /// /// Generic over the underlying storage to allow for domain specific /// optimizations (e.g. `eventIds` may or may not be contiguous). #[derive(Debug)] pub(crate) struct Registry { storage: S, } impl Registry { fn new(storage: S) -> Self { Self { storage } } } impl Registry { /// Registers a new listener for `event_id`. fn register_listener(&self, event_id: EventId) -> watch::Receiver<()> { self.storage .event_info(event_id) .unwrap_or_else(|| panic!("invalid event_id: {event_id}")) .tx .subscribe() } /// Marks `event_id` as having been delivered, without broadcasting it to /// any listeners. fn record_event(&self, event_id: EventId) { if let Some(event_info) = self.storage.event_info(event_id) { event_info.pending.store(true, Ordering::SeqCst); } } /// Broadcasts all previously recorded events to their respective listeners. /// /// Returns `true` if an event was delivered to at least one listener. fn broadcast(&self) -> bool { let mut did_notify = false; self.storage.for_each(|event_info| { // Any signal of this kind arrived since we checked last? if !event_info.pending.swap(false, Ordering::SeqCst) { return; } // Ignore errors if there are no listeners if event_info.tx.send(()).is_ok() { did_notify = true; } }); did_notify } } pub(crate) struct Globals { extra: OsExtraData, registry: Registry, } impl ops::Deref for Globals { type Target = OsExtraData; fn deref(&self) -> &Self::Target { &self.extra } } impl Globals { /// Registers a new listener for `event_id`. pub(crate) fn register_listener(&self, event_id: EventId) -> watch::Receiver<()> { self.registry.register_listener(event_id) } /// Marks `event_id` as having been delivered, without broadcasting it to /// any listeners. pub(crate) fn record_event(&self, event_id: EventId) { self.registry.record_event(event_id); } /// Broadcasts all previously recorded events to their respective listeners. /// /// Returns `true` if an event was delivered to at least one listener. pub(crate) fn broadcast(&self) -> bool { self.registry.broadcast() } #[cfg(unix)] pub(crate) fn storage(&self) -> &OsStorage { &self.registry.storage } } fn globals_init() -> Globals where OsExtraData: 'static + Send + Sync + Init, OsStorage: 'static + Send + Sync + Init, { Globals { extra: OsExtraData::init(), registry: Registry::new(OsStorage::init()), } } pub(crate) fn globals() -> &'static Globals where OsExtraData: 'static + Send + Sync + Init, OsStorage: 'static + Send + Sync + Init, { static GLOBALS: OnceCell = OnceCell::new(); GLOBALS.get(globals_init) } #[cfg(all(test, not(loom)))] mod tests { use super::*; use crate::runtime::{self, Runtime}; use crate::sync::{oneshot, watch}; use futures::future; #[test] fn smoke() { let rt = rt(); rt.block_on(async move { let registry = Registry::new(vec![ EventInfo::default(), EventInfo::default(), EventInfo::default(), ]); let first = registry.register_listener(0); let second = registry.register_listener(1); let third = registry.register_listener(2); let (fire, wait) = oneshot::channel(); crate::spawn(async { wait.await.expect("wait failed"); // Record some events which should get coalesced registry.record_event(0); registry.record_event(0); registry.record_event(1); registry.record_event(1); registry.broadcast(); // Yield so the previous broadcast can get received // // This yields many times since the block_on task is only polled every 61 // ticks. for _ in 0..100 { crate::task::yield_now().await; } // Send subsequent signal registry.record_event(0); registry.broadcast(); drop(registry); }); let _ = fire.send(()); let all = future::join3(collect(first), collect(second), collect(third)); let (first_results, second_results, third_results) = all.await; assert_eq!(2, first_results.len()); assert_eq!(1, second_results.len()); assert_eq!(0, third_results.len()); }); } #[test] #[should_panic = "invalid event_id: 1"] fn register_panics_on_invalid_input() { let registry = Registry::new(vec![EventInfo::default()]); registry.register_listener(1); } #[test] fn record_invalid_event_does_nothing() { let registry = Registry::new(vec![EventInfo::default()]); registry.record_event(1302); } #[test] fn broadcast_returns_if_at_least_one_event_fired() { let registry = Registry::new(vec![EventInfo::default(), EventInfo::default()]); registry.record_event(0); assert!(!registry.broadcast()); let first = registry.register_listener(0); let second = registry.register_listener(1); registry.record_event(0); assert!(registry.broadcast()); drop(first); registry.record_event(0); assert!(!registry.broadcast()); drop(second); } fn rt() -> Runtime { runtime::Builder::new_current_thread() .enable_time() .build() .unwrap() } async fn collect(mut rx: watch::Receiver<()>) -> Vec<()> { let mut ret = vec![]; while let Ok(v) = rx.changed().await { ret.push(v); } ret } } tokio-1.43.1/src/signal/reusable_box.rs000064400000000000000000000157161046102023000161350ustar 00000000000000use std::alloc::Layout; use std::future::Future; use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::ptr::{self, NonNull}; use std::task::{Context, Poll}; use std::{fmt, panic}; /// A reusable `Pin + Send>>`. /// /// This type lets you replace the future stored in the box without /// reallocating when the size and alignment permits this. pub(crate) struct ReusableBoxFuture { boxed: NonNull + Send>, } impl ReusableBoxFuture { /// Create a new `ReusableBoxFuture` containing the provided future. pub(crate) fn new(future: F) -> Self where F: Future + Send + 'static, { let boxed: Box + Send> = Box::new(future); let boxed = Box::into_raw(boxed); // SAFETY: Box::into_raw does not return null pointers. let boxed = unsafe { NonNull::new_unchecked(boxed) }; Self { boxed } } /// Replaces the future currently stored in this box. /// /// This reallocates if and only if the layout of the provided future is /// different from the layout of the currently stored future. pub(crate) fn set(&mut self, future: F) where F: Future + Send + 'static, { if let Err(future) = self.try_set(future) { *self = Self::new(future); } } /// Replaces the future currently stored in this box. /// /// This function never reallocates, but returns an error if the provided /// future has a different size or alignment from the currently stored /// future. pub(crate) fn try_set(&mut self, future: F) -> Result<(), F> where F: Future + Send + 'static, { // SAFETY: The pointer is not dangling. let self_layout = { let dyn_future: &(dyn Future + Send) = unsafe { self.boxed.as_ref() }; Layout::for_value(dyn_future) }; if Layout::new::() == self_layout { // SAFETY: We just checked that the layout of F is correct. unsafe { self.set_same_layout(future); } Ok(()) } else { Err(future) } } /// Sets the current future. /// /// # Safety /// /// This function requires that the layout of the provided future is the /// same as `self.layout`. unsafe fn set_same_layout(&mut self, future: F) where F: Future + Send + 'static, { // Drop the existing future, catching any panics. let result = panic::catch_unwind(AssertUnwindSafe(|| { ptr::drop_in_place(self.boxed.as_ptr()); })); // Overwrite the future behind the pointer. This is safe because the // allocation was allocated with the same size and alignment as the type F. let self_ptr: *mut F = self.boxed.as_ptr() as *mut F; ptr::write(self_ptr, future); // Update the vtable of self.boxed. The pointer is not null because we // just got it from self.boxed, which is not null. self.boxed = NonNull::new_unchecked(self_ptr); // If the old future's destructor panicked, resume unwinding. match result { Ok(()) => {} Err(payload) => { panic::resume_unwind(payload); } } } /// Gets a pinned reference to the underlying future. pub(crate) fn get_pin(&mut self) -> Pin<&mut (dyn Future + Send)> { // SAFETY: The user of this box cannot move the box, and we do not move it // either. unsafe { Pin::new_unchecked(self.boxed.as_mut()) } } /// Polls the future stored inside this box. pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll { self.get_pin().poll(cx) } } impl Future for ReusableBoxFuture { type Output = T; /// Polls the future stored inside this box. fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::into_inner(self).get_pin().poll(cx) } } // The future stored inside ReusableBoxFuture must be Send. unsafe impl Send for ReusableBoxFuture {} // The only method called on self.boxed is poll, which takes &mut self, so this // struct being Sync does not permit any invalid access to the Future, even if // the future is not Sync. unsafe impl Sync for ReusableBoxFuture {} // Just like a Pin> is always Unpin, so is this type. impl Unpin for ReusableBoxFuture {} impl Drop for ReusableBoxFuture { fn drop(&mut self) { unsafe { drop(Box::from_raw(self.boxed.as_ptr())); } } } impl fmt::Debug for ReusableBoxFuture { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ReusableBoxFuture").finish() } } #[cfg(test)] mod test { use super::ReusableBoxFuture; use futures::future::FutureExt; use std::alloc::Layout; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; #[test] fn test_different_futures() { let fut = async move { 10 }; // Not zero sized! assert_eq!(Layout::for_value(&fut).size(), 1); let mut b = ReusableBoxFuture::new(fut); assert_eq!(b.get_pin().now_or_never(), Some(10)); b.try_set(async move { 20 }) .unwrap_or_else(|_| panic!("incorrect size")); assert_eq!(b.get_pin().now_or_never(), Some(20)); b.try_set(async move { 30 }) .unwrap_or_else(|_| panic!("incorrect size")); assert_eq!(b.get_pin().now_or_never(), Some(30)); } #[test] fn test_different_sizes() { let fut1 = async move { 10 }; let val = [0u32; 1000]; let fut2 = async move { val[0] }; let fut3 = ZeroSizedFuture {}; assert_eq!(Layout::for_value(&fut1).size(), 1); assert_eq!(Layout::for_value(&fut2).size(), 4004); assert_eq!(Layout::for_value(&fut3).size(), 0); let mut b = ReusableBoxFuture::new(fut1); assert_eq!(b.get_pin().now_or_never(), Some(10)); b.set(fut2); assert_eq!(b.get_pin().now_or_never(), Some(0)); b.set(fut3); assert_eq!(b.get_pin().now_or_never(), Some(5)); } struct ZeroSizedFuture {} impl Future for ZeroSizedFuture { type Output = u32; fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { Poll::Ready(5) } } #[test] fn test_zero_sized() { let fut = ZeroSizedFuture {}; // Zero sized! assert_eq!(Layout::for_value(&fut).size(), 0); let mut b = ReusableBoxFuture::new(fut); assert_eq!(b.get_pin().now_or_never(), Some(5)); assert_eq!(b.get_pin().now_or_never(), Some(5)); b.try_set(ZeroSizedFuture {}) .unwrap_or_else(|_| panic!("incorrect size")); assert_eq!(b.get_pin().now_or_never(), Some(5)); assert_eq!(b.get_pin().now_or_never(), Some(5)); } } tokio-1.43.1/src/signal/unix.rs000064400000000000000000000425471046102023000144500ustar 00000000000000//! Unix-specific types for signal handling. //! //! This module is only defined on Unix platforms and contains the primary //! `Signal` type for receiving notifications of signals. #![cfg(unix)] #![cfg_attr(docsrs, doc(cfg(all(unix, feature = "signal"))))] use crate::runtime::scheduler; use crate::runtime::signal::Handle; use crate::signal::registry::{globals, EventId, EventInfo, Globals, Init, Storage}; use crate::signal::RxFuture; use crate::sync::watch; use mio::net::UnixStream; use std::io::{self, Error, ErrorKind, Write}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Once; use std::task::{Context, Poll}; pub(crate) type OsStorage = Box<[SignalInfo]>; impl Init for OsStorage { fn init() -> Self { // There are reliable signals ranging from 1 to 33 available on every Unix platform. #[cfg(not(any(target_os = "linux", target_os = "illumos")))] let possible = 0..=33; // On Linux and illumos, there are additional real-time signals // available. (This is also likely true on Solaris, but this should be // verified before being enabled.) #[cfg(any(target_os = "linux", target_os = "illumos"))] let possible = 0..=libc::SIGRTMAX(); possible.map(|_| SignalInfo::default()).collect() } } impl Storage for OsStorage { fn event_info(&self, id: EventId) -> Option<&EventInfo> { self.get(id).map(|si| &si.event_info) } fn for_each<'a, F>(&'a self, f: F) where F: FnMut(&'a EventInfo), { self.iter().map(|si| &si.event_info).for_each(f); } } #[derive(Debug)] pub(crate) struct OsExtraData { sender: UnixStream, pub(crate) receiver: UnixStream, } impl Init for OsExtraData { fn init() -> Self { let (receiver, sender) = UnixStream::pair().expect("failed to create UnixStream"); Self { sender, receiver } } } /// Represents the specific kind of signal to listen for. #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] pub struct SignalKind(libc::c_int); impl SignalKind { /// Allows for listening to any valid OS signal. /// /// For example, this can be used for listening for platform-specific /// signals. /// ```rust,no_run /// # use tokio::signal::unix::SignalKind; /// # let signum = -1; /// // let signum = libc::OS_SPECIFIC_SIGNAL; /// let kind = SignalKind::from_raw(signum); /// ``` // Use `std::os::raw::c_int` on public API to prevent leaking a non-stable // type alias from libc. // `libc::c_int` and `std::os::raw::c_int` are currently the same type, and are // unlikely to change to other types, but technically libc can change this // in the future minor version. // See https://github.com/tokio-rs/tokio/issues/3767 for more. pub const fn from_raw(signum: std::os::raw::c_int) -> Self { Self(signum as libc::c_int) } /// Get the signal's numeric value. /// /// ```rust /// # use tokio::signal::unix::SignalKind; /// let kind = SignalKind::interrupt(); /// assert_eq!(kind.as_raw_value(), libc::SIGINT); /// ``` pub const fn as_raw_value(&self) -> std::os::raw::c_int { self.0 } /// Represents the `SIGALRM` signal. /// /// On Unix systems this signal is sent when a real-time timer has expired. /// By default, the process is terminated by this signal. pub const fn alarm() -> Self { Self(libc::SIGALRM) } /// Represents the `SIGCHLD` signal. /// /// On Unix systems this signal is sent when the status of a child process /// has changed. By default, this signal is ignored. pub const fn child() -> Self { Self(libc::SIGCHLD) } /// Represents the `SIGHUP` signal. /// /// On Unix systems this signal is sent when the terminal is disconnected. /// By default, the process is terminated by this signal. pub const fn hangup() -> Self { Self(libc::SIGHUP) } /// Represents the `SIGINFO` signal. /// /// On Unix systems this signal is sent to request a status update from the /// process. By default, this signal is ignored. #[cfg(any( target_os = "dragonfly", target_os = "freebsd", target_os = "macos", target_os = "netbsd", target_os = "openbsd", target_os = "illumos" ))] pub const fn info() -> Self { Self(libc::SIGINFO) } /// Represents the `SIGINT` signal. /// /// On Unix systems this signal is sent to interrupt a program. /// By default, the process is terminated by this signal. pub const fn interrupt() -> Self { Self(libc::SIGINT) } #[cfg(target_os = "haiku")] /// Represents the `SIGPOLL` signal. /// /// On POSIX systems this signal is sent when I/O operations are possible /// on some file descriptor. By default, this signal is ignored. pub const fn io() -> Self { Self(libc::SIGPOLL) } #[cfg(not(target_os = "haiku"))] /// Represents the `SIGIO` signal. /// /// On Unix systems this signal is sent when I/O operations are possible /// on some file descriptor. By default, this signal is ignored. pub const fn io() -> Self { Self(libc::SIGIO) } /// Represents the `SIGPIPE` signal. /// /// On Unix systems this signal is sent when the process attempts to write /// to a pipe which has no reader. By default, the process is terminated by /// this signal. pub const fn pipe() -> Self { Self(libc::SIGPIPE) } /// Represents the `SIGQUIT` signal. /// /// On Unix systems this signal is sent to issue a shutdown of the /// process, after which the OS will dump the process core. /// By default, the process is terminated by this signal. pub const fn quit() -> Self { Self(libc::SIGQUIT) } /// Represents the `SIGTERM` signal. /// /// On Unix systems this signal is sent to issue a shutdown of the /// process. By default, the process is terminated by this signal. pub const fn terminate() -> Self { Self(libc::SIGTERM) } /// Represents the `SIGUSR1` signal. /// /// On Unix systems this is a user defined signal. /// By default, the process is terminated by this signal. pub const fn user_defined1() -> Self { Self(libc::SIGUSR1) } /// Represents the `SIGUSR2` signal. /// /// On Unix systems this is a user defined signal. /// By default, the process is terminated by this signal. pub const fn user_defined2() -> Self { Self(libc::SIGUSR2) } /// Represents the `SIGWINCH` signal. /// /// On Unix systems this signal is sent when the terminal window is resized. /// By default, this signal is ignored. pub const fn window_change() -> Self { Self(libc::SIGWINCH) } } impl From for SignalKind { fn from(signum: std::os::raw::c_int) -> Self { Self::from_raw(signum as libc::c_int) } } impl From for std::os::raw::c_int { fn from(kind: SignalKind) -> Self { kind.as_raw_value() } } pub(crate) struct SignalInfo { event_info: EventInfo, init: Once, initialized: AtomicBool, } impl Default for SignalInfo { fn default() -> SignalInfo { SignalInfo { event_info: EventInfo::default(), init: Once::new(), initialized: AtomicBool::new(false), } } } /// Our global signal handler for all signals registered by this module. /// /// The purpose of this signal handler is to primarily: /// /// 1. Flag that our specific signal was received (e.g. store an atomic flag) /// 2. Wake up the driver by writing a byte to a pipe /// /// Those two operations should both be async-signal safe. fn action(globals: &'static Globals, signal: libc::c_int) { globals.record_event(signal as EventId); // Send a wakeup, ignore any errors (anything reasonably possible is // full pipe and then it will wake up anyway). let mut sender = &globals.sender; drop(sender.write(&[1])); } /// Enables this module to receive signal notifications for the `signal` /// provided. /// /// This will register the signal handler if it hasn't already been registered, /// returning any error along the way if that fails. fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { let signal = signal.0; if signal < 0 || signal_hook_registry::FORBIDDEN.contains(&signal) { return Err(Error::new( ErrorKind::Other, format!("Refusing to register signal {signal}"), )); } // Check that we have a signal driver running handle.check_inner()?; let globals = globals(); let siginfo = match globals.storage().get(signal as EventId) { Some(slot) => slot, None => return Err(io::Error::new(io::ErrorKind::Other, "signal too large")), }; let mut registered = Ok(()); siginfo.init.call_once(|| { registered = unsafe { signal_hook_registry::register(signal, move || action(globals, signal)).map(|_| ()) }; if registered.is_ok() { siginfo.initialized.store(true, Ordering::Relaxed); } }); registered?; // If the call_once failed, it won't be retried on the next attempt to register the signal. In // such case it is not run, registered is still `Ok(())`, initialized is still `false`. if siginfo.initialized.load(Ordering::Relaxed) { Ok(()) } else { Err(Error::new( ErrorKind::Other, "Failed to register signal handler", )) } } /// An listener for receiving a particular type of OS signal. /// /// The listener can be turned into a `Stream` using [`SignalStream`]. /// /// [`SignalStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.SignalStream.html /// /// In general signal handling on Unix is a pretty tricky topic, and this /// structure is no exception! There are some important limitations to keep in /// mind when using `Signal` streams: /// /// * Signals handling in Unix already necessitates coalescing signals /// together sometimes. This `Signal` stream is also no exception here in /// that it will also coalesce signals. That is, even if the signal handler /// for this process runs multiple times, the `Signal` stream may only return /// one signal notification. Specifically, before `poll` is called, all /// signal notifications are coalesced into one item returned from `poll`. /// Once `poll` has been called, however, a further signal is guaranteed to /// be yielded as an item. /// /// Put another way, any element pulled off the returned listener corresponds to /// *at least one* signal, but possibly more. /// /// * Signal handling in general is relatively inefficient. Although some /// improvements are possible in this crate, it's recommended to not plan on /// having millions of signal channels open. /// /// If you've got any questions about this feel free to open an issue on the /// repo! New approaches to alleviate some of these limitations are always /// appreciated! /// /// # Caveats /// /// The first time that a `Signal` instance is registered for a particular /// signal kind, an OS signal-handler is installed which replaces the default /// platform behavior when that signal is received, **for the duration of the /// entire process**. /// /// For example, Unix systems will terminate a process by default when it /// receives `SIGINT`. But, when a `Signal` instance is created to listen for /// this signal, the next `SIGINT` that arrives will be translated to a stream /// event, and the process will continue to execute. **Even if this `Signal` /// instance is dropped, subsequent `SIGINT` deliveries will end up captured by /// Tokio, and the default platform behavior will NOT be reset**. /// /// Thus, applications should take care to ensure the expected signal behavior /// occurs as expected after listening for specific signals. /// /// # Examples /// /// Wait for `SIGHUP` /// /// ```rust,no_run /// use tokio::signal::unix::{signal, SignalKind}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // An infinite stream of hangup signals. /// let mut sig = signal(SignalKind::hangup())?; /// /// // Print whenever a HUP signal is received /// loop { /// sig.recv().await; /// println!("got signal HUP"); /// } /// } /// ``` #[must_use = "streams do nothing unless polled"] #[derive(Debug)] pub struct Signal { inner: RxFuture, } /// Creates a new listener which will receive notifications when the current /// process receives the specified signal `kind`. /// /// This function will create a new stream which binds to the default reactor. /// The `Signal` stream is an infinite stream which will receive /// notifications whenever a signal is received. More documentation can be /// found on `Signal` itself, but to reiterate: /// /// * Signals may be coalesced beyond what the kernel already does. /// * Once a signal handler is registered with the process the underlying /// libc signal handler is never unregistered. /// /// A `Signal` stream can be created for a particular signal number /// multiple times. When a signal is received then all the associated /// channels will receive the signal notification. /// /// # Errors /// /// * If the lower-level C functions fail for some reason. /// * If the previous initialization of this specific signal failed. /// * If the signal is one of /// [`signal_hook::FORBIDDEN`](fn@signal_hook_registry::register#panics) /// /// # Panics /// /// This function panics if there is no current reactor set, or if the `rt` /// feature flag is not enabled. #[track_caller] pub fn signal(kind: SignalKind) -> io::Result { let handle = scheduler::Handle::current(); let rx = signal_with_handle(kind, handle.driver().signal())?; Ok(Signal { inner: RxFuture::new(rx), }) } pub(crate) fn signal_with_handle( kind: SignalKind, handle: &Handle, ) -> io::Result> { // Turn the signal delivery on once we are ready for it signal_enable(kind, handle)?; Ok(globals().register_listener(kind.0 as EventId)) } impl Signal { /// Receives the next signal notification event. /// /// `None` is returned if no more events can be received by this stream. /// /// # Cancel safety /// /// This method is cancel safe. If you use it as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no signal is lost. /// /// # Examples /// /// Wait for `SIGHUP` /// /// ```rust,no_run /// use tokio::signal::unix::{signal, SignalKind}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // An infinite stream of hangup signals. /// let mut stream = signal(SignalKind::hangup())?; /// /// // Print whenever a HUP signal is received /// loop { /// stream.recv().await; /// println!("got signal HUP"); /// } /// } /// ``` pub async fn recv(&mut self) -> Option<()> { self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// /// This method returns: /// /// * `Poll::Pending` if no signals are available but the channel is not /// closed. /// * `Poll::Ready(Some(()))` if a signal is available. /// * `Poll::Ready(None)` if the channel has been closed and all signals /// sent before it was closed have been received. /// /// # Examples /// /// Polling from a manually implemented future /// /// ```rust,no_run /// use std::pin::Pin; /// use std::future::Future; /// use std::task::{Context, Poll}; /// use tokio::signal::unix::Signal; /// /// struct MyFuture { /// signal: Signal, /// } /// /// impl Future for MyFuture { /// type Output = Option<()>; /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// println!("polling MyFuture"); /// self.signal.poll_recv(cx) /// } /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_recv(cx) } } // Work around for abstracting streams internally #[cfg(feature = "process")] pub(crate) trait InternalStream { fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll>; } #[cfg(feature = "process")] impl InternalStream for Signal { fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.poll_recv(cx) } } pub(crate) fn ctrl_c() -> io::Result { signal(SignalKind::interrupt()) } #[cfg(all(test, not(loom)))] mod tests { use super::*; #[test] fn signal_enable_error_on_invalid_input() { signal_enable(SignalKind::from_raw(-1), &Handle::default()).unwrap_err(); } #[test] fn signal_enable_error_on_forbidden_input() { signal_enable( SignalKind::from_raw(signal_hook_registry::FORBIDDEN[0]), &Handle::default(), ) .unwrap_err(); } #[test] fn from_c_int() { assert_eq!(SignalKind::from(2), SignalKind::interrupt()); } #[test] fn into_c_int() { let value: std::os::raw::c_int = SignalKind::interrupt().into(); assert_eq!(value, libc::SIGINT as _); } } tokio-1.43.1/src/signal/windows/stub.rs000064400000000000000000000010011046102023000161100ustar 00000000000000//! Stub implementations for the platform API so that rustdoc can build linkable //! documentation on non-windows platforms. use crate::signal::RxFuture; use std::io; pub(super) fn ctrl_break() -> io::Result { panic!() } pub(super) fn ctrl_close() -> io::Result { panic!() } pub(super) fn ctrl_c() -> io::Result { panic!() } pub(super) fn ctrl_logoff() -> io::Result { panic!() } pub(super) fn ctrl_shutdown() -> io::Result { panic!() } tokio-1.43.1/src/signal/windows/sys.rs000064400000000000000000000141471046102023000157700ustar 00000000000000use std::io; use std::sync::Once; use crate::signal::registry::{globals, EventId, EventInfo, Init, Storage}; use crate::signal::RxFuture; use windows_sys::Win32::Foundation::BOOL; use windows_sys::Win32::System::Console as console; pub(super) fn ctrl_break() -> io::Result { new(console::CTRL_BREAK_EVENT) } pub(super) fn ctrl_close() -> io::Result { new(console::CTRL_CLOSE_EVENT) } pub(super) fn ctrl_c() -> io::Result { new(console::CTRL_C_EVENT) } pub(super) fn ctrl_logoff() -> io::Result { new(console::CTRL_LOGOFF_EVENT) } pub(super) fn ctrl_shutdown() -> io::Result { new(console::CTRL_SHUTDOWN_EVENT) } fn new(signum: u32) -> io::Result { global_init()?; let rx = globals().register_listener(signum as EventId); Ok(RxFuture::new(rx)) } #[derive(Debug)] pub(crate) struct OsStorage { ctrl_break: EventInfo, ctrl_close: EventInfo, ctrl_c: EventInfo, ctrl_logoff: EventInfo, ctrl_shutdown: EventInfo, } impl Init for OsStorage { fn init() -> Self { Self { ctrl_break: Default::default(), ctrl_close: Default::default(), ctrl_c: Default::default(), ctrl_logoff: Default::default(), ctrl_shutdown: Default::default(), } } } impl Storage for OsStorage { fn event_info(&self, id: EventId) -> Option<&EventInfo> { match u32::try_from(id) { Ok(console::CTRL_BREAK_EVENT) => Some(&self.ctrl_break), Ok(console::CTRL_CLOSE_EVENT) => Some(&self.ctrl_close), Ok(console::CTRL_C_EVENT) => Some(&self.ctrl_c), Ok(console::CTRL_LOGOFF_EVENT) => Some(&self.ctrl_logoff), Ok(console::CTRL_SHUTDOWN_EVENT) => Some(&self.ctrl_shutdown), _ => None, } } fn for_each<'a, F>(&'a self, mut f: F) where F: FnMut(&'a EventInfo), { f(&self.ctrl_break); f(&self.ctrl_close); f(&self.ctrl_c); f(&self.ctrl_logoff); f(&self.ctrl_shutdown); } } #[derive(Debug)] pub(crate) struct OsExtraData {} impl Init for OsExtraData { fn init() -> Self { Self {} } } fn global_init() -> io::Result<()> { static INIT: Once = Once::new(); let mut init = None; INIT.call_once(|| unsafe { let rc = console::SetConsoleCtrlHandler(Some(handler), 1); let ret = if rc == 0 { Err(io::Error::last_os_error()) } else { Ok(()) }; init = Some(ret); }); init.unwrap_or_else(|| Ok(())) } unsafe extern "system" fn handler(ty: u32) -> BOOL { let globals = globals(); globals.record_event(ty as EventId); // According to https://docs.microsoft.com/en-us/windows/console/handlerroutine // the handler routine is always invoked in a new thread, thus we don't // have the same restrictions as in Unix signal handlers, meaning we can // go ahead and perform the broadcast here. if globals.broadcast() { 1 } else { // No one is listening for this notification any more // let the OS fire the next (possibly the default) handler. 0 } } #[cfg(all(test, not(loom)))] mod tests { use super::*; use crate::runtime::Runtime; use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task}; #[test] fn ctrl_c() { let rt = rt(); let _enter = rt.enter(); let mut ctrl_c = task::spawn(crate::signal::ctrl_c()); assert_pending!(ctrl_c.poll()); // Windows doesn't have a good programmatic way of sending events // like sending signals on Unix, so we'll stub out the actual OS // integration and test that our handling works. unsafe { super::handler(console::CTRL_C_EVENT); } assert_ready_ok!(ctrl_c.poll()); } #[test] fn ctrl_break() { let rt = rt(); rt.block_on(async { let mut ctrl_break = assert_ok!(crate::signal::windows::ctrl_break()); // Windows doesn't have a good programmatic way of sending events // like sending signals on Unix, so we'll stub out the actual OS // integration and test that our handling works. unsafe { super::handler(console::CTRL_BREAK_EVENT); } ctrl_break.recv().await.unwrap(); }); } #[test] fn ctrl_close() { let rt = rt(); rt.block_on(async { let mut ctrl_close = assert_ok!(crate::signal::windows::ctrl_close()); // Windows doesn't have a good programmatic way of sending events // like sending signals on Unix, so we'll stub out the actual OS // integration and test that our handling works. unsafe { super::handler(console::CTRL_CLOSE_EVENT); } ctrl_close.recv().await.unwrap(); }); } #[test] fn ctrl_shutdown() { let rt = rt(); rt.block_on(async { let mut ctrl_shutdown = assert_ok!(crate::signal::windows::ctrl_shutdown()); // Windows doesn't have a good programmatic way of sending events // like sending signals on Unix, so we'll stub out the actual OS // integration and test that our handling works. unsafe { super::handler(console::CTRL_SHUTDOWN_EVENT); } ctrl_shutdown.recv().await.unwrap(); }); } #[test] fn ctrl_logoff() { let rt = rt(); rt.block_on(async { let mut ctrl_logoff = assert_ok!(crate::signal::windows::ctrl_logoff()); // Windows doesn't have a good programmatic way of sending events // like sending signals on Unix, so we'll stub out the actual OS // integration and test that our handling works. unsafe { super::handler(console::CTRL_LOGOFF_EVENT); } ctrl_logoff.recv().await.unwrap(); }); } fn rt() -> Runtime { crate::runtime::Builder::new_current_thread() .build() .unwrap() } } tokio-1.43.1/src/signal/windows.rs000064400000000000000000000366711046102023000151600ustar 00000000000000//! Windows-specific types for signal handling. //! //! This module is only defined on Windows and allows receiving "ctrl-c", //! "ctrl-break", "ctrl-logoff", "ctrl-shutdown", and "ctrl-close" //! notifications. These events are listened for via the `SetConsoleCtrlHandler` //! function which receives the corresponding `windows_sys` event type. #![cfg(any(windows, docsrs))] #![cfg_attr(docsrs, doc(cfg(all(windows, feature = "signal"))))] use crate::signal::RxFuture; use std::io; use std::task::{Context, Poll}; #[cfg(windows)] #[path = "windows/sys.rs"] mod imp; #[cfg(windows)] pub(crate) use self::imp::{OsExtraData, OsStorage}; // For building documentation on Unix machines when the `docsrs` flag is set. #[cfg(not(windows))] #[path = "windows/stub.rs"] mod imp; /// Creates a new listener which receives "ctrl-c" notifications sent to the /// process. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_c; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-C events. /// let mut signal = ctrl_c()?; /// /// // Print whenever a CTRL-C event is received. /// for countdown in (0..3).rev() { /// signal.recv().await; /// println!("got CTRL-C. {} more to exit", countdown); /// } /// /// Ok(()) /// } /// ``` pub fn ctrl_c() -> io::Result { Ok(CtrlC { inner: self::imp::ctrl_c()?, }) } /// Represents a listener which receives "ctrl-c" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// /// This event can be turned into a `Stream` using [`CtrlCStream`]. /// /// [`CtrlCStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.CtrlCStream.html /// /// A notification to this process notifies *all* receivers for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, /// then the listener may only receive one item about the two notifications. #[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlC { inner: RxFuture, } impl CtrlC { /// Receives the next signal notification event. /// /// `None` is returned if no more events can be received by the listener. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_c; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let mut signal = ctrl_c()?; /// /// // Print whenever a CTRL-C event is received. /// for countdown in (0..3).rev() { /// signal.recv().await; /// println!("got CTRL-C. {} more to exit", countdown); /// } /// /// Ok(()) /// } /// ``` pub async fn recv(&mut self) -> Option<()> { self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// /// `None` is returned if no more events can be received. /// /// # Examples /// /// Polling from a manually implemented future /// /// ```rust,no_run /// use std::pin::Pin; /// use std::future::Future; /// use std::task::{Context, Poll}; /// use tokio::signal::windows::CtrlC; /// /// struct MyFuture { /// ctrl_c: CtrlC, /// } /// /// impl Future for MyFuture { /// type Output = Option<()>; /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// println!("polling MyFuture"); /// self.ctrl_c.poll_recv(cx) /// } /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_recv(cx) } } /// Represents a listener which receives "ctrl-break" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// /// This listener can be turned into a `Stream` using [`CtrlBreakStream`]. /// /// [`CtrlBreakStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.CtrlBreakStream.html /// /// A notification to this process notifies *all* receivers for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, /// then the listener may only receive one item about the two notifications. #[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlBreak { inner: RxFuture, } impl CtrlBreak { /// Receives the next signal notification event. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_break; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-BREAK events. /// let mut signal = ctrl_break()?; /// /// // Print whenever a CTRL-BREAK event is received. /// loop { /// signal.recv().await; /// println!("got signal CTRL-BREAK"); /// } /// } /// ``` pub async fn recv(&mut self) -> Option<()> { self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// Polling from a manually implemented future /// /// ```rust,no_run /// use std::pin::Pin; /// use std::future::Future; /// use std::task::{Context, Poll}; /// use tokio::signal::windows::CtrlBreak; /// /// struct MyFuture { /// ctrl_break: CtrlBreak, /// } /// /// impl Future for MyFuture { /// type Output = Option<()>; /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// println!("polling MyFuture"); /// self.ctrl_break.poll_recv(cx) /// } /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_recv(cx) } } /// Creates a new listener which receives "ctrl-break" notifications sent to the /// process. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_break; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-BREAK events. /// let mut signal = ctrl_break()?; /// /// // Print whenever a CTRL-BREAK event is received. /// loop { /// signal.recv().await; /// println!("got signal CTRL-BREAK"); /// } /// } /// ``` pub fn ctrl_break() -> io::Result { Ok(CtrlBreak { inner: self::imp::ctrl_break()?, }) } /// Creates a new listener which receives "ctrl-close" notifications sent to the /// process. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_close; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-CLOSE events. /// let mut signal = ctrl_close()?; /// /// // Print whenever a CTRL-CLOSE event is received. /// for countdown in (0..3).rev() { /// signal.recv().await; /// println!("got CTRL-CLOSE. {} more to exit", countdown); /// } /// /// Ok(()) /// } /// ``` pub fn ctrl_close() -> io::Result { Ok(CtrlClose { inner: self::imp::ctrl_close()?, }) } /// Represents a listener which receives "ctrl-close" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// /// A notification to this process notifies *all* listeners listening for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, /// then the listener may only receive one item about the two notifications. #[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlClose { inner: RxFuture, } impl CtrlClose { /// Receives the next signal notification event. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_close; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-CLOSE events. /// let mut signal = ctrl_close()?; /// /// // Print whenever a CTRL-CLOSE event is received. /// signal.recv().await; /// println!("got CTRL-CLOSE. Cleaning up before exiting"); /// /// Ok(()) /// } /// ``` pub async fn recv(&mut self) -> Option<()> { self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// Polling from a manually implemented future /// /// ```rust,no_run /// use std::pin::Pin; /// use std::future::Future; /// use std::task::{Context, Poll}; /// use tokio::signal::windows::CtrlClose; /// /// struct MyFuture { /// ctrl_close: CtrlClose, /// } /// /// impl Future for MyFuture { /// type Output = Option<()>; /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// println!("polling MyFuture"); /// self.ctrl_close.poll_recv(cx) /// } /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_recv(cx) } } /// Creates a new listener which receives "ctrl-shutdown" notifications sent to the /// process. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_shutdown; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-SHUTDOWN events. /// let mut signal = ctrl_shutdown()?; /// /// signal.recv().await; /// println!("got CTRL-SHUTDOWN. Cleaning up before exiting"); /// /// Ok(()) /// } /// ``` pub fn ctrl_shutdown() -> io::Result { Ok(CtrlShutdown { inner: self::imp::ctrl_shutdown()?, }) } /// Represents a listener which receives "ctrl-shutdown" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// /// A notification to this process notifies *all* listeners listening for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, /// then the listener may only receive one item about the two notifications. #[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlShutdown { inner: RxFuture, } impl CtrlShutdown { /// Receives the next signal notification event. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_shutdown; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-SHUTDOWN events. /// let mut signal = ctrl_shutdown()?; /// /// // Print whenever a CTRL-SHUTDOWN event is received. /// signal.recv().await; /// println!("got CTRL-SHUTDOWN. Cleaning up before exiting"); /// /// Ok(()) /// } /// ``` pub async fn recv(&mut self) -> Option<()> { self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// Polling from a manually implemented future /// /// ```rust,no_run /// use std::pin::Pin; /// use std::future::Future; /// use std::task::{Context, Poll}; /// use tokio::signal::windows::CtrlShutdown; /// /// struct MyFuture { /// ctrl_shutdown: CtrlShutdown, /// } /// /// impl Future for MyFuture { /// type Output = Option<()>; /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// println!("polling MyFuture"); /// self.ctrl_shutdown.poll_recv(cx) /// } /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_recv(cx) } } /// Creates a new listener which receives "ctrl-logoff" notifications sent to the /// process. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_logoff; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // A listener of CTRL-LOGOFF events. /// let mut signal = ctrl_logoff()?; /// /// signal.recv().await; /// println!("got CTRL-LOGOFF. Cleaning up before exiting"); /// /// Ok(()) /// } /// ``` pub fn ctrl_logoff() -> io::Result { Ok(CtrlLogoff { inner: self::imp::ctrl_logoff()?, }) } /// Represents a listener which receives "ctrl-logoff" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// /// A notification to this process notifies *all* listeners listening for /// this event. Moreover, the notifications **are coalesced** if they aren't processed /// quickly enough. This means that if two notifications are received back-to-back, /// then the listener may only receive one item about the two notifications. #[must_use = "listeners do nothing unless polled"] #[derive(Debug)] pub struct CtrlLogoff { inner: RxFuture, } impl CtrlLogoff { /// Receives the next signal notification event. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// ```rust,no_run /// use tokio::signal::windows::ctrl_logoff; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// // An listener of CTRL-LOGOFF events. /// let mut signal = ctrl_logoff()?; /// /// // Print whenever a CTRL-LOGOFF event is received. /// signal.recv().await; /// println!("got CTRL-LOGOFF. Cleaning up before exiting"); /// /// Ok(()) /// } /// ``` pub async fn recv(&mut self) -> Option<()> { self.inner.recv().await } /// Polls to receive the next signal notification event, outside of an /// `async` context. /// /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// /// Polling from a manually implemented future /// /// ```rust,no_run /// use std::pin::Pin; /// use std::future::Future; /// use std::task::{Context, Poll}; /// use tokio::signal::windows::CtrlLogoff; /// /// struct MyFuture { /// ctrl_logoff: CtrlLogoff, /// } /// /// impl Future for MyFuture { /// type Output = Option<()>; /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// println!("polling MyFuture"); /// self.ctrl_logoff.poll_recv(cx) /// } /// } /// ``` pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_recv(cx) } } tokio-1.43.1/src/sync/barrier.rs000064400000000000000000000161101046102023000145750ustar 00000000000000use crate::loom::sync::Mutex; use crate::sync::watch; #[cfg(all(tokio_unstable, feature = "tracing"))] use crate::util::trace; /// A barrier enables multiple tasks to synchronize the beginning of some computation. /// /// ``` /// # #[tokio::main] /// # async fn main() { /// use tokio::sync::Barrier; /// use std::sync::Arc; /// /// let mut handles = Vec::with_capacity(10); /// let barrier = Arc::new(Barrier::new(10)); /// for _ in 0..10 { /// let c = barrier.clone(); /// // The same messages will be printed together. /// // You will NOT see any interleaving. /// handles.push(tokio::spawn(async move { /// println!("before wait"); /// let wait_result = c.wait().await; /// println!("after wait"); /// wait_result /// })); /// } /// /// // Will not resolve until all "after wait" messages have been printed /// let mut num_leaders = 0; /// for handle in handles { /// let wait_result = handle.await.unwrap(); /// if wait_result.is_leader() { /// num_leaders += 1; /// } /// } /// /// // Exactly one barrier will resolve as the "leader" /// assert_eq!(num_leaders, 1); /// # } /// ``` #[derive(Debug)] pub struct Barrier { state: Mutex, wait: watch::Receiver, n: usize, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, } #[derive(Debug)] struct BarrierState { waker: watch::Sender, arrived: usize, generation: usize, } impl Barrier { /// Creates a new barrier that can block a given number of tasks. /// /// A barrier will block `n`-1 tasks which call [`Barrier::wait`] and then wake up all /// tasks at once when the `n`th task calls `wait`. #[track_caller] pub fn new(mut n: usize) -> Barrier { let (waker, wait) = crate::sync::watch::channel(0); if n == 0 { // if n is 0, it's not clear what behavior the user wants. // in std::sync::Barrier, an n of 0 exhibits the same behavior as n == 1, where every // .wait() immediately unblocks, so we adopt that here as well. n = 1; } #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = std::panic::Location::caller(); let resource_span = tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Barrier", kind = "Sync", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", size = n, ); tracing::trace!( target: "runtime::resource::state_update", arrived = 0, ) }); resource_span }; Barrier { state: Mutex::new(BarrierState { waker, arrived: 0, generation: 1, }), n, wait, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Does not resolve until all tasks have rendezvoused here. /// /// Barriers are re-usable after all tasks have rendezvoused once, and can /// be used continuously. /// /// A single (arbitrary) future will receive a [`BarrierWaitResult`] that returns `true` from /// [`BarrierWaitResult::is_leader`] when returning from this function, and all other tasks /// will receive a result that will return `false` from `is_leader`. /// /// # Cancel safety /// /// This method is not cancel safe. pub async fn wait(&self) -> BarrierWaitResult { #[cfg(all(tokio_unstable, feature = "tracing"))] return trace::async_op( || self.wait_internal(), self.resource_span.clone(), "Barrier::wait", "poll", false, ) .await; #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] return self.wait_internal().await; } async fn wait_internal(&self) -> BarrierWaitResult { crate::trace::async_trace_leaf().await; // NOTE: we are taking a _synchronous_ lock here. // It is okay to do so because the critical section is fast and never yields, so it cannot // deadlock even if another future is concurrently holding the lock. // It is _desirable_ to do so as synchronous Mutexes are, at least in theory, faster than // the asynchronous counter-parts, so we should use them where possible [citation needed]. // NOTE: the extra scope here is so that the compiler doesn't think `state` is held across // a yield point, and thus marks the returned future as !Send. let generation = { let mut state = self.state.lock(); let generation = state.generation; state.arrived += 1; #[cfg(all(tokio_unstable, feature = "tracing"))] tracing::trace!( target: "runtime::resource::state_update", arrived = 1, arrived.op = "add", ); #[cfg(all(tokio_unstable, feature = "tracing"))] tracing::trace!( target: "runtime::resource::async_op::state_update", arrived = true, ); if state.arrived == self.n { #[cfg(all(tokio_unstable, feature = "tracing"))] tracing::trace!( target: "runtime::resource::async_op::state_update", is_leader = true, ); // we are the leader for this generation // wake everyone, increment the generation, and return state .waker .send(state.generation) .expect("there is at least one receiver"); state.arrived = 0; state.generation += 1; return BarrierWaitResult(true); } generation }; // we're going to have to wait for the last of the generation to arrive let mut wait = self.wait.clone(); loop { let _ = wait.changed().await; // note that the first time through the loop, this _will_ yield a generation // immediately, since we cloned a receiver that has never seen any values. if *wait.borrow() >= generation { break; } } BarrierWaitResult(false) } } /// A `BarrierWaitResult` is returned by `wait` when all tasks in the `Barrier` have rendezvoused. #[derive(Debug, Clone)] pub struct BarrierWaitResult(bool); impl BarrierWaitResult { /// Returns `true` if this task from wait is the "leader task". /// /// Only one task will have `true` returned from their result, all other tasks will have /// `false` returned. pub fn is_leader(&self) -> bool { self.0 } } tokio-1.43.1/src/sync/batch_semaphore.rs000064400000000000000000000653061046102023000163060ustar 00000000000000#![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))] //! # Implementation Details. //! //! The semaphore is implemented using an intrusive linked list of waiters. An //! atomic counter tracks the number of available permits. If the semaphore does //! not contain the required number of permits, the task attempting to acquire //! permits places its waker at the end of a queue. When new permits are made //! available (such as by releasing an initial acquisition), they are assigned //! to the task at the front of the queue, waking that task if its requested //! number of permits is met. //! //! Because waiters are enqueued at the back of the linked list and dequeued //! from the front, the semaphore is fair. Tasks trying to acquire large numbers //! of permits at a time will always be woken eventually, even if many other //! tasks are acquiring smaller numbers of permits. This means that in a //! use-case like tokio's read-write lock, writers will not be starved by //! readers. use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Mutex, MutexGuard}; use crate::util::linked_list::{self, LinkedList}; #[cfg(all(tokio_unstable, feature = "tracing"))] use crate::util::trace; use crate::util::WakeList; use std::future::Future; use std::marker::PhantomPinned; use std::pin::Pin; use std::ptr::NonNull; use std::sync::atomic::Ordering::*; use std::task::{ready, Context, Poll, Waker}; use std::{cmp, fmt}; /// An asynchronous counting semaphore which permits waiting on multiple permits at once. pub(crate) struct Semaphore { waiters: Mutex, /// The current number of available permits in the semaphore. permits: AtomicUsize, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, } struct Waitlist { queue: LinkedList::Target>, closed: bool, } /// Error returned from the [`Semaphore::try_acquire`] function. /// /// [`Semaphore::try_acquire`]: crate::sync::Semaphore::try_acquire #[derive(Debug, PartialEq, Eq)] pub enum TryAcquireError { /// The semaphore has been [closed] and cannot issue new permits. /// /// [closed]: crate::sync::Semaphore::close Closed, /// The semaphore has no available permits. NoPermits, } /// Error returned from the [`Semaphore::acquire`] function. /// /// An `acquire` operation can only fail if the semaphore has been /// [closed]. /// /// [closed]: crate::sync::Semaphore::close /// [`Semaphore::acquire`]: crate::sync::Semaphore::acquire #[derive(Debug)] pub struct AcquireError(()); pub(crate) struct Acquire<'a> { node: Waiter, semaphore: &'a Semaphore, num_permits: usize, queued: bool, } /// An entry in the wait queue. struct Waiter { /// The current state of the waiter. /// /// This is either the number of remaining permits required by /// the waiter, or a flag indicating that the waiter is not yet queued. state: AtomicUsize, /// The waker to notify the task awaiting permits. /// /// # Safety /// /// This may only be accessed while the wait queue is locked. waker: UnsafeCell>, /// Intrusive linked-list pointers. /// /// # Safety /// /// This may only be accessed while the wait queue is locked. /// /// TODO: Ideally, we would be able to use loom to enforce that /// this isn't accessed concurrently. However, it is difficult to /// use a `UnsafeCell` here, since the `Link` trait requires _returning_ /// references to `Pointers`, and `UnsafeCell` requires that checked access /// take place inside a closure. We should consider changing `Pointers` to /// use `UnsafeCell` internally. pointers: linked_list::Pointers, #[cfg(all(tokio_unstable, feature = "tracing"))] ctx: trace::AsyncOpTracingCtx, /// Should not be `Unpin`. _p: PhantomPinned, } generate_addr_of_methods! { impl<> Waiter { unsafe fn addr_of_pointers(self: NonNull) -> NonNull> { &self.pointers } } } impl Semaphore { /// The maximum number of permits which a semaphore can hold. /// /// Note that this reserves three bits of flags in the permit counter, but /// we only actually use one of them. However, the previous semaphore /// implementation used three bits, so we will continue to reserve them to /// avoid a breaking change if additional flags need to be added in the /// future. pub(crate) const MAX_PERMITS: usize = usize::MAX >> 3; const CLOSED: usize = 1; // The least-significant bit in the number of permits is reserved to use // as a flag indicating that the semaphore has been closed. Consequently // PERMIT_SHIFT is used to leave that bit for that purpose. const PERMIT_SHIFT: usize = 1; /// Creates a new semaphore with the initial number of permits /// /// Maximum number of permits on 32-bit platforms is `1<<29`. pub(crate) fn new(permits: usize) -> Self { assert!( permits <= Self::MAX_PERMITS, "a semaphore may not have more than MAX_PERMITS permits ({})", Self::MAX_PERMITS ); #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let resource_span = tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Semaphore", kind = "Sync", is_internal = true ); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", permits = permits, permits.op = "override", ) }); resource_span }; Self { permits: AtomicUsize::new(permits << Self::PERMIT_SHIFT), waiters: Mutex::new(Waitlist { queue: LinkedList::new(), closed: false, }), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Creates a new semaphore with the initial number of permits. /// /// Maximum number of permits on 32-bit platforms is `1<<29`. #[cfg(not(all(loom, test)))] pub(crate) const fn const_new(permits: usize) -> Self { assert!(permits <= Self::MAX_PERMITS); Self { permits: AtomicUsize::new(permits << Self::PERMIT_SHIFT), waiters: Mutex::const_new(Waitlist { queue: LinkedList::new(), closed: false, }), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Creates a new closed semaphore with 0 permits. pub(crate) fn new_closed() -> Self { Self { permits: AtomicUsize::new(Self::CLOSED), waiters: Mutex::new(Waitlist { queue: LinkedList::new(), closed: true, }), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Creates a new closed semaphore with 0 permits. #[cfg(not(all(loom, test)))] pub(crate) const fn const_new_closed() -> Self { Self { permits: AtomicUsize::new(Self::CLOSED), waiters: Mutex::const_new(Waitlist { queue: LinkedList::new(), closed: true, }), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Returns the current number of available permits. pub(crate) fn available_permits(&self) -> usize { self.permits.load(Acquire) >> Self::PERMIT_SHIFT } /// Adds `added` new permits to the semaphore. /// /// The maximum number of permits is `usize::MAX >> 3`, and this function will panic if the limit is exceeded. pub(crate) fn release(&self, added: usize) { if added == 0 { return; } // Assign permits to the wait queue self.add_permits_locked(added, self.waiters.lock()); } /// Closes the semaphore. This prevents the semaphore from issuing new /// permits and notifies all pending waiters. pub(crate) fn close(&self) { let mut waiters = self.waiters.lock(); // If the semaphore's permits counter has enough permits for an // unqueued waiter to acquire all the permits it needs immediately, // it won't touch the wait list. Therefore, we have to set a bit on // the permit counter as well. However, we must do this while // holding the lock --- otherwise, if we set the bit and then wait // to acquire the lock we'll enter an inconsistent state where the // permit counter is closed, but the wait list is not. self.permits.fetch_or(Self::CLOSED, Release); waiters.closed = true; while let Some(mut waiter) = waiters.queue.pop_back() { let waker = unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) }; if let Some(waker) = waker { waker.wake(); } } } /// Returns true if the semaphore is closed. pub(crate) fn is_closed(&self) -> bool { self.permits.load(Acquire) & Self::CLOSED == Self::CLOSED } pub(crate) fn try_acquire(&self, num_permits: usize) -> Result<(), TryAcquireError> { assert!( num_permits <= Self::MAX_PERMITS, "a semaphore may not have more than MAX_PERMITS permits ({})", Self::MAX_PERMITS ); let num_permits = num_permits << Self::PERMIT_SHIFT; let mut curr = self.permits.load(Acquire); loop { // Has the semaphore closed? if curr & Self::CLOSED == Self::CLOSED { return Err(TryAcquireError::Closed); } // Are there enough permits remaining? if curr < num_permits { return Err(TryAcquireError::NoPermits); } let next = curr - num_permits; match self.permits.compare_exchange(curr, next, AcqRel, Acquire) { Ok(_) => { // TODO: Instrument once issue has been solved return Ok(()); } Err(actual) => curr = actual, } } } pub(crate) fn acquire(&self, num_permits: usize) -> Acquire<'_> { Acquire::new(self, num_permits) } /// Release `rem` permits to the semaphore's wait list, starting from the /// end of the queue. /// /// If `rem` exceeds the number of permits needed by the wait list, the /// remainder are assigned back to the semaphore. fn add_permits_locked(&self, mut rem: usize, waiters: MutexGuard<'_, Waitlist>) { let mut wakers = WakeList::new(); let mut lock = Some(waiters); let mut is_empty = false; while rem > 0 { let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); 'inner: while wakers.can_push() { // Was the waiter assigned enough permits to wake it? match waiters.queue.last() { Some(waiter) => { if !waiter.assign_permits(&mut rem) { break 'inner; } } None => { is_empty = true; // If we assigned permits to all the waiters in the queue, and there are // still permits left over, assign them back to the semaphore. break 'inner; } }; let mut waiter = waiters.queue.pop_back().unwrap(); if let Some(waker) = unsafe { waiter.as_mut().waker.with_mut(|waker| (*waker).take()) } { wakers.push(waker); } } if rem > 0 && is_empty { let permits = rem; assert!( permits <= Self::MAX_PERMITS, "cannot add more than MAX_PERMITS permits ({})", Self::MAX_PERMITS ); let prev = self.permits.fetch_add(rem << Self::PERMIT_SHIFT, Release); let prev = prev >> Self::PERMIT_SHIFT; assert!( prev + permits <= Self::MAX_PERMITS, "number of added permits ({}) would overflow MAX_PERMITS ({})", rem, Self::MAX_PERMITS ); // add remaining permits back #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", permits = rem, permits.op = "add", ) }); rem = 0; } drop(waiters); // release the lock wakers.wake_all(); } assert_eq!(rem, 0); } /// Decrease a semaphore's permits by a maximum of `n`. /// /// If there are insufficient permits and it's not possible to reduce by `n`, /// return the number of permits that were actually reduced. pub(crate) fn forget_permits(&self, n: usize) -> usize { if n == 0 { return 0; } let mut curr_bits = self.permits.load(Acquire); loop { let curr = curr_bits >> Self::PERMIT_SHIFT; let new = curr.saturating_sub(n); match self.permits.compare_exchange_weak( curr_bits, new << Self::PERMIT_SHIFT, AcqRel, Acquire, ) { Ok(_) => return std::cmp::min(curr, n), Err(actual) => curr_bits = actual, }; } } fn poll_acquire( &self, cx: &mut Context<'_>, num_permits: usize, node: Pin<&mut Waiter>, queued: bool, ) -> Poll> { let mut acquired = 0; let needed = if queued { node.state.load(Acquire) << Self::PERMIT_SHIFT } else { num_permits << Self::PERMIT_SHIFT }; let mut lock = None; // First, try to take the requested number of permits from the // semaphore. let mut curr = self.permits.load(Acquire); let mut waiters = loop { // Has the semaphore closed? if curr & Self::CLOSED > 0 { return Poll::Ready(Err(AcquireError::closed())); } let mut remaining = 0; let total = curr .checked_add(acquired) .expect("number of permits must not overflow"); let (next, acq) = if total >= needed { let next = curr - (needed - acquired); (next, needed >> Self::PERMIT_SHIFT) } else { remaining = (needed - acquired) - curr; (0, curr >> Self::PERMIT_SHIFT) }; if remaining > 0 && lock.is_none() { // No permits were immediately available, so this permit will // (probably) need to wait. We'll need to acquire a lock on the // wait queue before continuing. We need to do this _before_ the // CAS that sets the new value of the semaphore's `permits` // counter. Otherwise, if we subtract the permits and then // acquire the lock, we might miss additional permits being // added while waiting for the lock. lock = Some(self.waiters.lock()); } match self.permits.compare_exchange(curr, next, AcqRel, Acquire) { Ok(_) => { acquired += acq; if remaining == 0 { if !queued { #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", permits = acquired, permits.op = "sub", ); tracing::trace!( target: "runtime::resource::async_op::state_update", permits_obtained = acquired, permits.op = "add", ) }); return Poll::Ready(Ok(())); } else if lock.is_none() { break self.waiters.lock(); } } break lock.expect("lock must be acquired before waiting"); } Err(actual) => curr = actual, } }; if waiters.closed { return Poll::Ready(Err(AcquireError::closed())); } #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", permits = acquired, permits.op = "sub", ) }); if node.assign_permits(&mut acquired) { self.add_permits_locked(acquired, waiters); return Poll::Ready(Ok(())); } assert_eq!(acquired, 0); let mut old_waker = None; // Otherwise, register the waker & enqueue the node. node.waker.with_mut(|waker| { // Safety: the wait list is locked, so we may modify the waker. let waker = unsafe { &mut *waker }; // Do we need to register the new waker? if waker .as_ref() .map_or(true, |waker| !waker.will_wake(cx.waker())) { old_waker = std::mem::replace(waker, Some(cx.waker().clone())); } }); // If the waiter is not already in the wait queue, enqueue it. if !queued { let node = unsafe { let node = Pin::into_inner_unchecked(node) as *mut _; NonNull::new_unchecked(node) }; waiters.queue.push_front(node); } drop(waiters); drop(old_waker); Poll::Pending } } impl fmt::Debug for Semaphore { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Semaphore") .field("permits", &self.available_permits()) .finish() } } impl Waiter { fn new( num_permits: usize, #[cfg(all(tokio_unstable, feature = "tracing"))] ctx: trace::AsyncOpTracingCtx, ) -> Self { Waiter { waker: UnsafeCell::new(None), state: AtomicUsize::new(num_permits), pointers: linked_list::Pointers::new(), #[cfg(all(tokio_unstable, feature = "tracing"))] ctx, _p: PhantomPinned, } } /// Assign permits to the waiter. /// /// Returns `true` if the waiter should be removed from the queue fn assign_permits(&self, n: &mut usize) -> bool { let mut curr = self.state.load(Acquire); loop { let assign = cmp::min(curr, *n); let next = curr - assign; match self.state.compare_exchange(curr, next, AcqRel, Acquire) { Ok(_) => { *n -= assign; #[cfg(all(tokio_unstable, feature = "tracing"))] self.ctx.async_op_span.in_scope(|| { tracing::trace!( target: "runtime::resource::async_op::state_update", permits_obtained = assign, permits.op = "add", ); }); return next == 0; } Err(actual) => curr = actual, } } } } impl Future for Acquire<'_> { type Output = Result<(), AcquireError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { ready!(crate::trace::trace_leaf(cx)); #[cfg(all(tokio_unstable, feature = "tracing"))] let _resource_span = self.node.ctx.resource_span.clone().entered(); #[cfg(all(tokio_unstable, feature = "tracing"))] let _async_op_span = self.node.ctx.async_op_span.clone().entered(); #[cfg(all(tokio_unstable, feature = "tracing"))] let _async_op_poll_span = self.node.ctx.async_op_poll_span.clone().entered(); let (node, semaphore, needed, queued) = self.project(); // First, ensure the current task has enough budget to proceed. #[cfg(all(tokio_unstable, feature = "tracing"))] let coop = ready!(trace_poll_op!( "poll_acquire", crate::runtime::coop::poll_proceed(cx), )); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let result = match semaphore.poll_acquire(cx, needed, node, *queued) { Poll::Pending => { *queued = true; Poll::Pending } Poll::Ready(r) => { coop.made_progress(); r?; *queued = false; Poll::Ready(Ok(())) } }; #[cfg(all(tokio_unstable, feature = "tracing"))] return trace_poll_op!("poll_acquire", result); #[cfg(not(all(tokio_unstable, feature = "tracing")))] return result; } } impl<'a> Acquire<'a> { fn new(semaphore: &'a Semaphore, num_permits: usize) -> Self { #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] return Self { node: Waiter::new(num_permits), semaphore, num_permits, queued: false, }; #[cfg(all(tokio_unstable, feature = "tracing"))] return semaphore.resource_span.in_scope(|| { let async_op_span = tracing::trace_span!("runtime.resource.async_op", source = "Acquire::new"); let async_op_poll_span = async_op_span.in_scope(|| { tracing::trace!( target: "runtime::resource::async_op::state_update", permits_requested = num_permits, permits.op = "override", ); tracing::trace!( target: "runtime::resource::async_op::state_update", permits_obtained = 0usize, permits.op = "override", ); tracing::trace_span!("runtime.resource.async_op.poll") }); let ctx = trace::AsyncOpTracingCtx { async_op_span, async_op_poll_span, resource_span: semaphore.resource_span.clone(), }; Self { node: Waiter::new(num_permits, ctx), semaphore, num_permits, queued: false, } }); } fn project(self: Pin<&mut Self>) -> (Pin<&mut Waiter>, &Semaphore, usize, &mut bool) { fn is_unpin() {} unsafe { // Safety: all fields other than `node` are `Unpin` is_unpin::<&Semaphore>(); is_unpin::<&mut bool>(); is_unpin::(); let this = self.get_unchecked_mut(); ( Pin::new_unchecked(&mut this.node), this.semaphore, this.num_permits, &mut this.queued, ) } } } impl Drop for Acquire<'_> { fn drop(&mut self) { // If the future is completed, there is no node in the wait list, so we // can skip acquiring the lock. if !self.queued { return; } // This is where we ensure safety. The future is being dropped, // which means we must ensure that the waiter entry is no longer stored // in the linked list. let mut waiters = self.semaphore.waiters.lock(); // remove the entry from the list let node = NonNull::from(&mut self.node); // Safety: we have locked the wait list. unsafe { waiters.queue.remove(node) }; let acquired_permits = self.num_permits - self.node.state.load(Acquire); if acquired_permits > 0 { self.semaphore.add_permits_locked(acquired_permits, waiters); } } } // Safety: the `Acquire` future is not `Sync` automatically because it contains // a `Waiter`, which, in turn, contains an `UnsafeCell`. However, the // `UnsafeCell` is only accessed when the future is borrowed mutably (either in // `poll` or in `drop`). Therefore, it is safe (although not particularly // _useful_) for the future to be borrowed immutably across threads. unsafe impl Sync for Acquire<'_> {} // ===== impl AcquireError ==== impl AcquireError { fn closed() -> AcquireError { AcquireError(()) } } impl fmt::Display for AcquireError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "semaphore closed") } } impl std::error::Error for AcquireError {} // ===== impl TryAcquireError ===== impl TryAcquireError { /// Returns `true` if the error was caused by a closed semaphore. #[allow(dead_code)] // may be used later! pub(crate) fn is_closed(&self) -> bool { matches!(self, TryAcquireError::Closed) } /// Returns `true` if the error was caused by calling `try_acquire` on a /// semaphore with no available permits. #[allow(dead_code)] // may be used later! pub(crate) fn is_no_permits(&self) -> bool { matches!(self, TryAcquireError::NoPermits) } } impl fmt::Display for TryAcquireError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match self { TryAcquireError::Closed => write!(fmt, "semaphore closed"), TryAcquireError::NoPermits => write!(fmt, "no permits available"), } } } impl std::error::Error for TryAcquireError {} /// # Safety /// /// `Waiter` is forced to be !Unpin. unsafe impl linked_list::Link for Waiter { type Handle = NonNull; type Target = Waiter; fn as_raw(handle: &Self::Handle) -> NonNull { *handle } unsafe fn from_raw(ptr: NonNull) -> NonNull { ptr } unsafe fn pointers(target: NonNull) -> NonNull> { Waiter::addr_of_pointers(target) } } tokio-1.43.1/src/sync/broadcast.rs000064400000000000000000001414401046102023000151160ustar 00000000000000//! A multi-producer, multi-consumer broadcast queue. Each sent value is seen by //! all consumers. //! //! A [`Sender`] is used to broadcast values to **all** connected [`Receiver`] //! values. [`Sender`] handles are clone-able, allowing concurrent send and //! receive actions. [`Sender`] and [`Receiver`] are both `Send` and `Sync` as //! long as `T` is `Send`. //! //! When a value is sent, **all** [`Receiver`] handles are notified and will //! receive the value. The value is stored once inside the channel and cloned on //! demand for each receiver. Once all receivers have received a clone of the //! value, the value is released from the channel. //! //! A channel is created by calling [`channel`], specifying the maximum number //! of messages the channel can retain at any given time. //! //! New [`Receiver`] handles are created by calling [`Sender::subscribe`]. The //! returned [`Receiver`] will receive values sent **after** the call to //! `subscribe`. //! //! This channel is also suitable for the single-producer multi-consumer //! use-case, where a single sender broadcasts values to many receivers. //! //! ## Lagging //! //! As sent messages must be retained until **all** [`Receiver`] handles receive //! a clone, broadcast channels are susceptible to the "slow receiver" problem. //! In this case, all but one receiver are able to receive values at the rate //! they are sent. Because one receiver is stalled, the channel starts to fill //! up. //! //! This broadcast channel implementation handles this case by setting a hard //! upper bound on the number of values the channel may retain at any given //! time. This upper bound is passed to the [`channel`] function as an argument. //! //! If a value is sent when the channel is at capacity, the oldest value //! currently held by the channel is released. This frees up space for the new //! value. Any receiver that has not yet seen the released value will return //! [`RecvError::Lagged`] the next time [`recv`] is called. //! //! Once [`RecvError::Lagged`] is returned, the lagging receiver's position is //! updated to the oldest value contained by the channel. The next call to //! [`recv`] will return this value. //! //! This behavior enables a receiver to detect when it has lagged so far behind //! that data has been dropped. The caller may decide how to respond to this: //! either by aborting its task or by tolerating lost messages and resuming //! consumption of the channel. //! //! ## Closing //! //! When **all** [`Sender`] handles have been dropped, no new values may be //! sent. At this point, the channel is "closed". Once a receiver has received //! all values retained by the channel, the next call to [`recv`] will return //! with [`RecvError::Closed`]. //! //! When a [`Receiver`] handle is dropped, any messages not read by the receiver //! will be marked as read. If this receiver was the only one not to have read //! that message, the message will be dropped at this point. //! //! [`Sender`]: crate::sync::broadcast::Sender //! [`Sender::subscribe`]: crate::sync::broadcast::Sender::subscribe //! [`Receiver`]: crate::sync::broadcast::Receiver //! [`channel`]: crate::sync::broadcast::channel //! [`RecvError::Lagged`]: crate::sync::broadcast::error::RecvError::Lagged //! [`RecvError::Closed`]: crate::sync::broadcast::error::RecvError::Closed //! [`recv`]: crate::sync::broadcast::Receiver::recv //! //! # Examples //! //! Basic usage //! //! ``` //! use tokio::sync::broadcast; //! //! #[tokio::main] //! async fn main() { //! let (tx, mut rx1) = broadcast::channel(16); //! let mut rx2 = tx.subscribe(); //! //! tokio::spawn(async move { //! assert_eq!(rx1.recv().await.unwrap(), 10); //! assert_eq!(rx1.recv().await.unwrap(), 20); //! }); //! //! tokio::spawn(async move { //! assert_eq!(rx2.recv().await.unwrap(), 10); //! assert_eq!(rx2.recv().await.unwrap(), 20); //! }); //! //! tx.send(10).unwrap(); //! tx.send(20).unwrap(); //! } //! ``` //! //! Handling lag //! //! ``` //! use tokio::sync::broadcast; //! //! #[tokio::main] //! async fn main() { //! let (tx, mut rx) = broadcast::channel(2); //! //! tx.send(10).unwrap(); //! tx.send(20).unwrap(); //! tx.send(30).unwrap(); //! //! // The receiver lagged behind //! assert!(rx.recv().await.is_err()); //! //! // At this point, we can abort or continue with lost messages //! //! assert_eq!(20, rx.recv().await.unwrap()); //! assert_eq!(30, rx.recv().await.unwrap()); //! } //! ``` use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::{AtomicBool, AtomicUsize}; use crate::loom::sync::{Arc, Mutex, MutexGuard}; use crate::runtime::coop::cooperative; use crate::util::linked_list::{self, GuardedLinkedList, LinkedList}; use crate::util::WakeList; use std::fmt; use std::future::Future; use std::marker::PhantomPinned; use std::pin::Pin; use std::ptr::NonNull; use std::sync::atomic::Ordering::{Acquire, Relaxed, Release, SeqCst}; use std::task::{ready, Context, Poll, Waker}; /// Sending-half of the [`broadcast`] channel. /// /// May be used from many threads. Messages can be sent with /// [`send`][Sender::send]. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// tokio::spawn(async move { /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// }); /// /// tokio::spawn(async move { /// assert_eq!(rx2.recv().await.unwrap(), 10); /// assert_eq!(rx2.recv().await.unwrap(), 20); /// }); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// } /// ``` /// /// [`broadcast`]: crate::sync::broadcast pub struct Sender { shared: Arc>, } /// Receiving-half of the [`broadcast`] channel. /// /// Must not be used concurrently. Messages may be retrieved using /// [`recv`][Receiver::recv]. /// /// To turn this receiver into a `Stream`, you can use the [`BroadcastStream`] /// wrapper. /// /// [`BroadcastStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.BroadcastStream.html /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// tokio::spawn(async move { /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// }); /// /// tokio::spawn(async move { /// assert_eq!(rx2.recv().await.unwrap(), 10); /// assert_eq!(rx2.recv().await.unwrap(), 20); /// }); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// } /// ``` /// /// [`broadcast`]: crate::sync::broadcast pub struct Receiver { /// State shared with all receivers and senders. shared: Arc>, /// Next position to read from next: u64, } pub mod error { //! Broadcast error types use std::fmt; /// Error returned by the [`send`] function on a [`Sender`]. /// /// A **send** operation can only fail if there are no active receivers, /// implying that the message could never be received. The error contains the /// message being sent as a payload so it can be recovered. /// /// [`send`]: crate::sync::broadcast::Sender::send /// [`Sender`]: crate::sync::broadcast::Sender #[derive(Debug)] pub struct SendError(pub T); impl fmt::Display for SendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "channel closed") } } impl std::error::Error for SendError {} /// An error returned from the [`recv`] function on a [`Receiver`]. /// /// [`recv`]: crate::sync::broadcast::Receiver::recv /// [`Receiver`]: crate::sync::broadcast::Receiver #[derive(Debug, PartialEq, Eq, Clone)] pub enum RecvError { /// There are no more active senders implying no further messages will ever /// be sent. Closed, /// The receiver lagged too far behind. Attempting to receive again will /// return the oldest message still retained by the channel. /// /// Includes the number of skipped messages. Lagged(u64), } impl fmt::Display for RecvError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { RecvError::Closed => write!(f, "channel closed"), RecvError::Lagged(amt) => write!(f, "channel lagged by {amt}"), } } } impl std::error::Error for RecvError {} /// An error returned from the [`try_recv`] function on a [`Receiver`]. /// /// [`try_recv`]: crate::sync::broadcast::Receiver::try_recv /// [`Receiver`]: crate::sync::broadcast::Receiver #[derive(Debug, PartialEq, Eq, Clone)] pub enum TryRecvError { /// The channel is currently empty. There are still active /// [`Sender`] handles, so data may yet become available. /// /// [`Sender`]: crate::sync::broadcast::Sender Empty, /// There are no more active senders implying no further messages will ever /// be sent. Closed, /// The receiver lagged too far behind and has been forcibly disconnected. /// Attempting to receive again will return the oldest message still /// retained by the channel. /// /// Includes the number of skipped messages. Lagged(u64), } impl fmt::Display for TryRecvError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { TryRecvError::Empty => write!(f, "channel empty"), TryRecvError::Closed => write!(f, "channel closed"), TryRecvError::Lagged(amt) => write!(f, "channel lagged by {amt}"), } } } impl std::error::Error for TryRecvError {} } use self::error::{RecvError, SendError, TryRecvError}; /// Data shared between senders and receivers. struct Shared { /// slots in the channel. buffer: Box<[Mutex>]>, /// Mask a position -> index. mask: usize, /// Tail of the queue. Includes the rx wait list. tail: Mutex, /// Number of outstanding Sender handles. num_tx: AtomicUsize, } /// Next position to write a value. struct Tail { /// Next position to write to. pos: u64, /// Number of active receivers. rx_cnt: usize, /// True if the channel is closed. closed: bool, /// Receivers waiting for a value. waiters: LinkedList::Target>, } /// Slot in the buffer. struct Slot { /// Remaining number of receivers that are expected to see this value. /// /// When this goes to zero, the value is released. /// /// An atomic is used as it is mutated concurrently with the slot read lock /// acquired. rem: AtomicUsize, /// Uniquely identifies the `send` stored in the slot. pos: u64, /// The value being broadcast. /// /// The value is set by `send` when the write lock is held. When a reader /// drops, `rem` is decremented. When it hits zero, the value is dropped. val: Option, } /// An entry in the wait queue. struct Waiter { /// True if queued. queued: AtomicBool, /// Task waiting on the broadcast channel. waker: Option, /// Intrusive linked-list pointers. pointers: linked_list::Pointers, /// Should not be `Unpin`. _p: PhantomPinned, } impl Waiter { fn new() -> Self { Self { queued: AtomicBool::new(false), waker: None, pointers: linked_list::Pointers::new(), _p: PhantomPinned, } } } generate_addr_of_methods! { impl<> Waiter { unsafe fn addr_of_pointers(self: NonNull) -> NonNull> { &self.pointers } } } struct RecvGuard<'a, T> { slot: MutexGuard<'a, Slot>, } /// Receive a value future. struct Recv<'a, T> { /// Receiver being waited on. receiver: &'a mut Receiver, /// Entry in the waiter `LinkedList`. waiter: WaiterCell, } // The wrapper around `UnsafeCell` isolates the unsafe impl `Send` and `Sync` // from `Recv`. struct WaiterCell(UnsafeCell); unsafe impl Send for WaiterCell {} unsafe impl Sync for WaiterCell {} /// Max number of receivers. Reserve space to lock. const MAX_RECEIVERS: usize = usize::MAX >> 2; /// Create a bounded, multi-producer, multi-consumer channel where each sent /// value is broadcasted to all active receivers. /// /// **Note:** The actual capacity may be greater than the provided `capacity`. /// /// All data sent on [`Sender`] will become available on every active /// [`Receiver`] in the same order as it was sent. /// /// The `Sender` can be cloned to `send` to the same channel from multiple /// points in the process or it can be used concurrently from an `Arc`. New /// `Receiver` handles are created by calling [`Sender::subscribe`]. /// /// If all [`Receiver`] handles are dropped, the `send` method will return a /// [`SendError`]. Similarly, if all [`Sender`] handles are dropped, the [`recv`] /// method will return a [`RecvError`]. /// /// [`Sender`]: crate::sync::broadcast::Sender /// [`Sender::subscribe`]: crate::sync::broadcast::Sender::subscribe /// [`Receiver`]: crate::sync::broadcast::Receiver /// [`recv`]: crate::sync::broadcast::Receiver::recv /// [`SendError`]: crate::sync::broadcast::error::SendError /// [`RecvError`]: crate::sync::broadcast::error::RecvError /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// tokio::spawn(async move { /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// }); /// /// tokio::spawn(async move { /// assert_eq!(rx2.recv().await.unwrap(), 10); /// assert_eq!(rx2.recv().await.unwrap(), 20); /// }); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// } /// ``` /// /// # Panics /// /// This will panic if `capacity` is equal to `0` or larger /// than `usize::MAX / 2`. #[track_caller] pub fn channel(capacity: usize) -> (Sender, Receiver) { // SAFETY: In the line below we are creating one extra receiver, so there will be 1 in total. let tx = unsafe { Sender::new_with_receiver_count(1, capacity) }; let rx = Receiver { shared: tx.shared.clone(), next: 0, }; (tx, rx) } impl Sender { /// Creates the sending-half of the [`broadcast`] channel. /// /// See the documentation of [`broadcast::channel`] for more information on this method. /// /// [`broadcast`]: crate::sync::broadcast /// [`broadcast::channel`]: crate::sync::broadcast::channel #[track_caller] pub fn new(capacity: usize) -> Self { // SAFETY: We don't create extra receivers, so there are 0. unsafe { Self::new_with_receiver_count(0, capacity) } } /// Creates the sending-half of the [`broadcast`](self) channel, and provide the receiver /// count. /// /// See the documentation of [`broadcast::channel`](self::channel) for more errors when /// calling this function. /// /// # Safety: /// /// The caller must ensure that the amount of receivers for this Sender is correct before /// the channel functionalities are used, the count is zero by default, as this function /// does not create any receivers by itself. #[track_caller] unsafe fn new_with_receiver_count(receiver_count: usize, mut capacity: usize) -> Self { assert!(capacity > 0, "broadcast channel capacity cannot be zero"); assert!( capacity <= usize::MAX >> 1, "broadcast channel capacity exceeded `usize::MAX / 2`" ); // Round to a power of two capacity = capacity.next_power_of_two(); let mut buffer = Vec::with_capacity(capacity); for i in 0..capacity { buffer.push(Mutex::new(Slot { rem: AtomicUsize::new(0), pos: (i as u64).wrapping_sub(capacity as u64), val: None, })); } let shared = Arc::new(Shared { buffer: buffer.into_boxed_slice(), mask: capacity - 1, tail: Mutex::new(Tail { pos: 0, rx_cnt: receiver_count, closed: false, waiters: LinkedList::new(), }), num_tx: AtomicUsize::new(1), }); Sender { shared } } /// Attempts to send a value to all active [`Receiver`] handles, returning /// it back if it could not be sent. /// /// A successful send occurs when there is at least one active [`Receiver`] /// handle. An unsuccessful send would be one where all associated /// [`Receiver`] handles have already been dropped. /// /// # Return /// /// On success, the number of subscribed [`Receiver`] handles is returned. /// This does not mean that this number of receivers will see the message as /// a receiver may drop or lag ([see lagging](self#lagging)) before receiving /// the message. /// /// # Note /// /// A return value of `Ok` **does not** mean that the sent value will be /// observed by all or any of the active [`Receiver`] handles. [`Receiver`] /// handles may be dropped before receiving the sent message. /// /// A return value of `Err` **does not** mean that future calls to `send` /// will fail. New [`Receiver`] handles may be created by calling /// [`subscribe`]. /// /// [`Receiver`]: crate::sync::broadcast::Receiver /// [`subscribe`]: crate::sync::broadcast::Sender::subscribe /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// tokio::spawn(async move { /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// }); /// /// tokio::spawn(async move { /// assert_eq!(rx2.recv().await.unwrap(), 10); /// assert_eq!(rx2.recv().await.unwrap(), 20); /// }); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// } /// ``` pub fn send(&self, value: T) -> Result> { let mut tail = self.shared.tail.lock(); if tail.rx_cnt == 0 { return Err(SendError(value)); } // Position to write into let pos = tail.pos; let rem = tail.rx_cnt; let idx = (pos & self.shared.mask as u64) as usize; // Update the tail position tail.pos = tail.pos.wrapping_add(1); // Get the slot let mut slot = self.shared.buffer[idx].lock(); // Track the position slot.pos = pos; // Set remaining receivers slot.rem.with_mut(|v| *v = rem); // Write the value slot.val = Some(value); // Release the slot lock before notifying the receivers. drop(slot); // Notify and release the mutex. This must happen after the slot lock is // released, otherwise the writer lock bit could be cleared while another // thread is in the critical section. self.shared.notify_rx(tail); Ok(rem) } /// Creates a new [`Receiver`] handle that will receive values sent **after** /// this call to `subscribe`. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = broadcast::channel(16); /// /// // Will not be seen /// tx.send(10).unwrap(); /// /// let mut rx = tx.subscribe(); /// /// tx.send(20).unwrap(); /// /// let value = rx.recv().await.unwrap(); /// assert_eq!(20, value); /// } /// ``` pub fn subscribe(&self) -> Receiver { let shared = self.shared.clone(); new_receiver(shared) } /// Returns the number of queued values. /// /// A value is queued until it has either been seen by all receivers that were alive at the time /// it was sent, or has been evicted from the queue by subsequent sends that exceeded the /// queue's capacity. /// /// # Note /// /// In contrast to [`Receiver::len`], this method only reports queued values and not values that /// have been evicted from the queue before being seen by all receivers. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// tx.send(30).unwrap(); /// /// assert_eq!(tx.len(), 3); /// /// rx1.recv().await.unwrap(); /// /// // The len is still 3 since rx2 hasn't seen the first value yet. /// assert_eq!(tx.len(), 3); /// /// rx2.recv().await.unwrap(); /// /// assert_eq!(tx.len(), 2); /// } /// ``` pub fn len(&self) -> usize { let tail = self.shared.tail.lock(); let base_idx = (tail.pos & self.shared.mask as u64) as usize; let mut low = 0; let mut high = self.shared.buffer.len(); while low < high { let mid = low + (high - low) / 2; let idx = base_idx.wrapping_add(mid) & self.shared.mask; if self.shared.buffer[idx].lock().rem.load(SeqCst) == 0 { low = mid + 1; } else { high = mid; } } self.shared.buffer.len() - low } /// Returns true if there are no queued values. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// assert!(tx.is_empty()); /// /// tx.send(10).unwrap(); /// /// assert!(!tx.is_empty()); /// /// rx1.recv().await.unwrap(); /// /// // The queue is still not empty since rx2 hasn't seen the value. /// assert!(!tx.is_empty()); /// /// rx2.recv().await.unwrap(); /// /// assert!(tx.is_empty()); /// } /// ``` pub fn is_empty(&self) -> bool { let tail = self.shared.tail.lock(); let idx = (tail.pos.wrapping_sub(1) & self.shared.mask as u64) as usize; self.shared.buffer[idx].lock().rem.load(SeqCst) == 0 } /// Returns the number of active receivers. /// /// An active receiver is a [`Receiver`] handle returned from [`channel`] or /// [`subscribe`]. These are the handles that will receive values sent on /// this [`Sender`]. /// /// # Note /// /// It is not guaranteed that a sent message will reach this number of /// receivers. Active receivers may never call [`recv`] again before /// dropping. /// /// [`recv`]: crate::sync::broadcast::Receiver::recv /// [`Receiver`]: crate::sync::broadcast::Receiver /// [`Sender`]: crate::sync::broadcast::Sender /// [`subscribe`]: crate::sync::broadcast::Sender::subscribe /// [`channel`]: crate::sync::broadcast::channel /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx1) = broadcast::channel(16); /// /// assert_eq!(1, tx.receiver_count()); /// /// let mut _rx2 = tx.subscribe(); /// /// assert_eq!(2, tx.receiver_count()); /// /// tx.send(10).unwrap(); /// } /// ``` pub fn receiver_count(&self) -> usize { let tail = self.shared.tail.lock(); tail.rx_cnt } /// Returns `true` if senders belong to the same channel. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = broadcast::channel::<()>(16); /// let tx2 = tx.clone(); /// /// assert!(tx.same_channel(&tx2)); /// /// let (tx3, _rx3) = broadcast::channel::<()>(16); /// /// assert!(!tx3.same_channel(&tx2)); /// } /// ``` pub fn same_channel(&self, other: &Self) -> bool { Arc::ptr_eq(&self.shared, &other.shared) } fn close_channel(&self) { let mut tail = self.shared.tail.lock(); tail.closed = true; self.shared.notify_rx(tail); } } /// Create a new `Receiver` which reads starting from the tail. fn new_receiver(shared: Arc>) -> Receiver { let mut tail = shared.tail.lock(); assert!(tail.rx_cnt != MAX_RECEIVERS, "max receivers"); tail.rx_cnt = tail.rx_cnt.checked_add(1).expect("overflow"); let next = tail.pos; drop(tail); Receiver { shared, next } } /// List used in `Shared::notify_rx`. It wraps a guarded linked list /// and gates the access to it on the `Shared.tail` mutex. It also empties /// the list on drop. struct WaitersList<'a, T> { list: GuardedLinkedList::Target>, is_empty: bool, shared: &'a Shared, } impl<'a, T> Drop for WaitersList<'a, T> { fn drop(&mut self) { // If the list is not empty, we unlink all waiters from it. // We do not wake the waiters to avoid double panics. if !self.is_empty { let _lock_guard = self.shared.tail.lock(); while self.list.pop_back().is_some() {} } } } impl<'a, T> WaitersList<'a, T> { fn new( unguarded_list: LinkedList::Target>, guard: Pin<&'a Waiter>, shared: &'a Shared, ) -> Self { let guard_ptr = NonNull::from(guard.get_ref()); let list = unguarded_list.into_guarded(guard_ptr); WaitersList { list, is_empty: false, shared, } } /// Removes the last element from the guarded list. Modifying this list /// requires an exclusive access to the main list in `Notify`. fn pop_back_locked(&mut self, _tail: &mut Tail) -> Option> { let result = self.list.pop_back(); if result.is_none() { // Save information about emptiness to avoid waiting for lock // in the destructor. self.is_empty = true; } result } } impl Shared { fn notify_rx<'a, 'b: 'a>(&'b self, mut tail: MutexGuard<'a, Tail>) { // It is critical for `GuardedLinkedList` safety that the guard node is // pinned in memory and is not dropped until the guarded list is dropped. let guard = Waiter::new(); pin!(guard); // We move all waiters to a secondary list. It uses a `GuardedLinkedList` // underneath to allow every waiter to safely remove itself from it. // // * This list will be still guarded by the `waiters` lock. // `NotifyWaitersList` wrapper makes sure we hold the lock to modify it. // * This wrapper will empty the list on drop. It is critical for safety // that we will not leave any list entry with a pointer to the local // guard node after this function returns / panics. let mut list = WaitersList::new(std::mem::take(&mut tail.waiters), guard.as_ref(), self); let mut wakers = WakeList::new(); 'outer: loop { while wakers.can_push() { match list.pop_back_locked(&mut tail) { Some(waiter) => { unsafe { // Safety: accessing `waker` is safe because // the tail lock is held. if let Some(waker) = (*waiter.as_ptr()).waker.take() { wakers.push(waker); } // Safety: `queued` is atomic. let queued = &(*waiter.as_ptr()).queued; // `Relaxed` suffices because the tail lock is held. assert!(queued.load(Relaxed)); // `Release` is needed to synchronize with `Recv::drop`. // It is critical to set this variable **after** waker // is extracted, otherwise we may data race with `Recv::drop`. queued.store(false, Release); } } None => { break 'outer; } } } // Release the lock before waking. drop(tail); // Before we acquire the lock again all sorts of things can happen: // some waiters may remove themselves from the list and new waiters // may be added. This is fine since at worst we will unnecessarily // wake up waiters which will then queue themselves again. wakers.wake_all(); // Acquire the lock again. tail = self.tail.lock(); } // Release the lock before waking. drop(tail); wakers.wake_all(); } } impl Clone for Sender { fn clone(&self) -> Sender { let shared = self.shared.clone(); shared.num_tx.fetch_add(1, SeqCst); Sender { shared } } } impl Drop for Sender { fn drop(&mut self) { if 1 == self.shared.num_tx.fetch_sub(1, SeqCst) { self.close_channel(); } } } impl Receiver { /// Returns the number of messages that were sent into the channel and that /// this [`Receiver`] has yet to receive. /// /// If the returned value from `len` is larger than the next largest power of 2 /// of the capacity of the channel any call to [`recv`] will return an /// `Err(RecvError::Lagged)` and any call to [`try_recv`] will return an /// `Err(TryRecvError::Lagged)`, e.g. if the capacity of the channel is 10, /// [`recv`] will start to return `Err(RecvError::Lagged)` once `len` returns /// values larger than 16. /// /// [`Receiver`]: crate::sync::broadcast::Receiver /// [`recv`]: crate::sync::broadcast::Receiver::recv /// [`try_recv`]: crate::sync::broadcast::Receiver::try_recv /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// /// assert_eq!(rx1.len(), 2); /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.len(), 1); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// assert_eq!(rx1.len(), 0); /// } /// ``` pub fn len(&self) -> usize { let next_send_pos = self.shared.tail.lock().pos; (next_send_pos - self.next) as usize } /// Returns true if there aren't any messages in the channel that the [`Receiver`] /// has yet to receive. /// /// [`Receiver]: create::sync::broadcast::Receiver /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// /// assert!(rx1.is_empty()); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// /// assert!(!rx1.is_empty()); /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// assert!(rx1.is_empty()); /// } /// ``` pub fn is_empty(&self) -> bool { self.len() == 0 } /// Returns `true` if receivers belong to the same channel. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = broadcast::channel::<()>(16); /// let rx2 = tx.subscribe(); /// /// assert!(rx.same_channel(&rx2)); /// /// let (_tx3, rx3) = broadcast::channel::<()>(16); /// /// assert!(!rx3.same_channel(&rx2)); /// } /// ``` pub fn same_channel(&self, other: &Self) -> bool { Arc::ptr_eq(&self.shared, &other.shared) } /// Locks the next value if there is one. fn recv_ref( &mut self, waiter: Option<(&UnsafeCell, &Waker)>, ) -> Result, TryRecvError> { let idx = (self.next & self.shared.mask as u64) as usize; // The slot holding the next value to read let mut slot = self.shared.buffer[idx].lock(); if slot.pos != self.next { // Release the `slot` lock before attempting to acquire the `tail` // lock. This is required because `send2` acquires the tail lock // first followed by the slot lock. Acquiring the locks in reverse // order here would result in a potential deadlock: `recv_ref` // acquires the `slot` lock and attempts to acquire the `tail` lock // while `send2` acquired the `tail` lock and attempts to acquire // the slot lock. drop(slot); let mut old_waker = None; let mut tail = self.shared.tail.lock(); // Acquire slot lock again slot = self.shared.buffer[idx].lock(); // Make sure the position did not change. This could happen in the // unlikely event that the buffer is wrapped between dropping the // read lock and acquiring the tail lock. if slot.pos != self.next { let next_pos = slot.pos.wrapping_add(self.shared.buffer.len() as u64); if next_pos == self.next { // At this point the channel is empty for *this* receiver. If // it's been closed, then that's what we return, otherwise we // set a waker and return empty. if tail.closed { return Err(TryRecvError::Closed); } // Store the waker if let Some((waiter, waker)) = waiter { // Safety: called while locked. unsafe { // Only queue if not already queued waiter.with_mut(|ptr| { // If there is no waker **or** if the currently // stored waker references a **different** task, // track the tasks' waker to be notified on // receipt of a new value. match (*ptr).waker { Some(ref w) if w.will_wake(waker) => {} _ => { old_waker = std::mem::replace( &mut (*ptr).waker, Some(waker.clone()), ); } } // If the waiter is not already queued, enqueue it. // `Relaxed` order suffices: we have synchronized with // all writers through the tail lock that we hold. if !(*ptr).queued.load(Relaxed) { // `Relaxed` order suffices: all the readers will // synchronize with this write through the tail lock. (*ptr).queued.store(true, Relaxed); tail.waiters.push_front(NonNull::new_unchecked(&mut *ptr)); } }); } } // Drop the old waker after releasing the locks. drop(slot); drop(tail); drop(old_waker); return Err(TryRecvError::Empty); } // At this point, the receiver has lagged behind the sender by // more than the channel capacity. The receiver will attempt to // catch up by skipping dropped messages and setting the // internal cursor to the **oldest** message stored by the // channel. let next = tail.pos.wrapping_sub(self.shared.buffer.len() as u64); let missed = next.wrapping_sub(self.next); drop(tail); // The receiver is slow but no values have been missed if missed == 0 { self.next = self.next.wrapping_add(1); return Ok(RecvGuard { slot }); } self.next = next; return Err(TryRecvError::Lagged(missed)); } } self.next = self.next.wrapping_add(1); Ok(RecvGuard { slot }) } } impl Receiver { /// Re-subscribes to the channel starting from the current tail element. /// /// This [`Receiver`] handle will receive a clone of all values sent /// **after** it has resubscribed. This will not include elements that are /// in the queue of the current receiver. Consider the following example. /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = broadcast::channel(2); /// /// tx.send(1).unwrap(); /// let mut rx2 = rx.resubscribe(); /// tx.send(2).unwrap(); /// /// assert_eq!(rx2.recv().await.unwrap(), 2); /// assert_eq!(rx.recv().await.unwrap(), 1); /// } /// ``` pub fn resubscribe(&self) -> Self { let shared = self.shared.clone(); new_receiver(shared) } /// Receives the next value for this receiver. /// /// Each [`Receiver`] handle will receive a clone of all values sent /// **after** it has subscribed. /// /// `Err(RecvError::Closed)` is returned when all `Sender` halves have /// dropped, indicating that no further values can be sent on the channel. /// /// If the [`Receiver`] handle falls behind, once the channel is full, newly /// sent values will overwrite old values. At this point, a call to [`recv`] /// will return with `Err(RecvError::Lagged)` and the [`Receiver`]'s /// internal cursor is updated to point to the oldest value still held by /// the channel. A subsequent call to [`recv`] will return this value /// **unless** it has been since overwritten. /// /// # Cancel safety /// /// This method is cancel safe. If `recv` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// channel. /// /// [`Receiver`]: crate::sync::broadcast::Receiver /// [`recv`]: crate::sync::broadcast::Receiver::recv /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx1) = broadcast::channel(16); /// let mut rx2 = tx.subscribe(); /// /// tokio::spawn(async move { /// assert_eq!(rx1.recv().await.unwrap(), 10); /// assert_eq!(rx1.recv().await.unwrap(), 20); /// }); /// /// tokio::spawn(async move { /// assert_eq!(rx2.recv().await.unwrap(), 10); /// assert_eq!(rx2.recv().await.unwrap(), 20); /// }); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// } /// ``` /// /// Handling lag /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = broadcast::channel(2); /// /// tx.send(10).unwrap(); /// tx.send(20).unwrap(); /// tx.send(30).unwrap(); /// /// // The receiver lagged behind /// assert!(rx.recv().await.is_err()); /// /// // At this point, we can abort or continue with lost messages /// /// assert_eq!(20, rx.recv().await.unwrap()); /// assert_eq!(30, rx.recv().await.unwrap()); /// } /// ``` pub async fn recv(&mut self) -> Result { cooperative(Recv::new(self)).await } /// Attempts to return a pending value on this receiver without awaiting. /// /// This is useful for a flavor of "optimistic check" before deciding to /// await on a receiver. /// /// Compared with [`recv`], this function has three failure cases instead of two /// (one for closed, one for an empty buffer, one for a lagging receiver). /// /// `Err(TryRecvError::Closed)` is returned when all `Sender` halves have /// dropped, indicating that no further values can be sent on the channel. /// /// If the [`Receiver`] handle falls behind, once the channel is full, newly /// sent values will overwrite old values. At this point, a call to [`recv`] /// will return with `Err(TryRecvError::Lagged)` and the [`Receiver`]'s /// internal cursor is updated to point to the oldest value still held by /// the channel. A subsequent call to [`try_recv`] will return this value /// **unless** it has been since overwritten. If there are no values to /// receive, `Err(TryRecvError::Empty)` is returned. /// /// [`recv`]: crate::sync::broadcast::Receiver::recv /// [`try_recv`]: crate::sync::broadcast::Receiver::try_recv /// [`Receiver`]: crate::sync::broadcast::Receiver /// /// # Examples /// /// ``` /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = broadcast::channel(16); /// /// assert!(rx.try_recv().is_err()); /// /// tx.send(10).unwrap(); /// /// let value = rx.try_recv().unwrap(); /// assert_eq!(10, value); /// } /// ``` pub fn try_recv(&mut self) -> Result { let guard = self.recv_ref(None)?; guard.clone_value().ok_or(TryRecvError::Closed) } /// Blocking receive to call outside of asynchronous contexts. /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// ``` /// use std::thread; /// use tokio::sync::broadcast; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = broadcast::channel(16); /// /// let sync_code = thread::spawn(move || { /// assert_eq!(rx.blocking_recv(), Ok(10)); /// }); /// /// let _ = tx.send(10); /// sync_code.join().unwrap(); /// } /// ``` pub fn blocking_recv(&mut self) -> Result { crate::future::block_on(self.recv()) } } impl Drop for Receiver { fn drop(&mut self) { let mut tail = self.shared.tail.lock(); tail.rx_cnt -= 1; let until = tail.pos; drop(tail); while self.next < until { match self.recv_ref(None) { Ok(_) => {} // The channel is closed Err(TryRecvError::Closed) => break, // Ignore lagging, we will catch up Err(TryRecvError::Lagged(..)) => {} // Can't be empty Err(TryRecvError::Empty) => panic!("unexpected empty broadcast channel"), } } } } impl<'a, T> Recv<'a, T> { fn new(receiver: &'a mut Receiver) -> Recv<'a, T> { Recv { receiver, waiter: WaiterCell(UnsafeCell::new(Waiter { queued: AtomicBool::new(false), waker: None, pointers: linked_list::Pointers::new(), _p: PhantomPinned, })), } } /// A custom `project` implementation is used in place of `pin-project-lite` /// as a custom drop implementation is needed. fn project(self: Pin<&mut Self>) -> (&mut Receiver, &UnsafeCell) { unsafe { // Safety: Receiver is Unpin is_unpin::<&mut Receiver>(); let me = self.get_unchecked_mut(); (me.receiver, &me.waiter.0) } } } impl<'a, T> Future for Recv<'a, T> where T: Clone, { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); let (receiver, waiter) = self.project(); let guard = match receiver.recv_ref(Some((waiter, cx.waker()))) { Ok(value) => value, Err(TryRecvError::Empty) => return Poll::Pending, Err(TryRecvError::Lagged(n)) => return Poll::Ready(Err(RecvError::Lagged(n))), Err(TryRecvError::Closed) => return Poll::Ready(Err(RecvError::Closed)), }; Poll::Ready(guard.clone_value().ok_or(RecvError::Closed)) } } impl<'a, T> Drop for Recv<'a, T> { fn drop(&mut self) { // Safety: `waiter.queued` is atomic. // Acquire ordering is required to synchronize with // `Shared::notify_rx` before we drop the object. let queued = self .waiter .0 .with(|ptr| unsafe { (*ptr).queued.load(Acquire) }); // If the waiter is queued, we need to unlink it from the waiters list. // If not, no further synchronization is required, since the waiter // is not in the list and, as such, is not shared with any other threads. if queued { // Acquire the tail lock. This is required for safety before accessing // the waiter node. let mut tail = self.receiver.shared.tail.lock(); // Safety: tail lock is held. // `Relaxed` order suffices because we hold the tail lock. let queued = self .waiter .0 .with_mut(|ptr| unsafe { (*ptr).queued.load(Relaxed) }); if queued { // Remove the node // // safety: tail lock is held and the wait node is verified to be in // the list. unsafe { self.waiter.0.with_mut(|ptr| { tail.waiters.remove((&mut *ptr).into()); }); } } } } } /// # Safety /// /// `Waiter` is forced to be !Unpin. unsafe impl linked_list::Link for Waiter { type Handle = NonNull; type Target = Waiter; fn as_raw(handle: &NonNull) -> NonNull { *handle } unsafe fn from_raw(ptr: NonNull) -> NonNull { ptr } unsafe fn pointers(target: NonNull) -> NonNull> { Waiter::addr_of_pointers(target) } } impl fmt::Debug for Sender { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "broadcast::Sender") } } impl fmt::Debug for Receiver { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "broadcast::Receiver") } } impl<'a, T> RecvGuard<'a, T> { fn clone_value(&self) -> Option where T: Clone, { self.slot.val.clone() } } impl<'a, T> Drop for RecvGuard<'a, T> { fn drop(&mut self) { // Decrement the remaining counter if 1 == self.slot.rem.fetch_sub(1, SeqCst) { self.slot.val = None; } } } fn is_unpin() {} #[cfg(not(loom))] #[cfg(test)] mod tests { use super::*; #[test] fn receiver_count_on_sender_constructor() { let sender = Sender::::new(16); assert_eq!(sender.receiver_count(), 0); let rx_1 = sender.subscribe(); assert_eq!(sender.receiver_count(), 1); let rx_2 = rx_1.resubscribe(); assert_eq!(sender.receiver_count(), 2); let rx_3 = sender.subscribe(); assert_eq!(sender.receiver_count(), 3); drop(rx_3); drop(rx_1); assert_eq!(sender.receiver_count(), 1); drop(rx_2); assert_eq!(sender.receiver_count(), 0); } #[cfg(not(loom))] #[test] fn receiver_count_on_channel_constructor() { let (sender, rx) = channel::(16); assert_eq!(sender.receiver_count(), 1); let _rx_2 = rx.resubscribe(); assert_eq!(sender.receiver_count(), 2); } } tokio-1.43.1/src/sync/mod.rs000064400000000000000000000424651046102023000137420ustar 00000000000000#![cfg_attr(loom, allow(dead_code, unreachable_pub, unused_imports))] //! Synchronization primitives for use in asynchronous contexts. //! //! Tokio programs tend to be organized as a set of [tasks] where each task //! operates independently and may be executed on separate physical threads. The //! synchronization primitives provided in this module permit these independent //! tasks to communicate together. //! //! [tasks]: crate::task //! //! # Message passing //! //! The most common form of synchronization in a Tokio program is message //! passing. Two tasks operate independently and send messages to each other to //! synchronize. Doing so has the advantage of avoiding shared state. //! //! Message passing is implemented using channels. A channel supports sending a //! message from one producer task to one or more consumer tasks. There are a //! few flavors of channels provided by Tokio. Each channel flavor supports //! different message passing patterns. When a channel supports multiple //! producers, many separate tasks may **send** messages. When a channel //! supports multiple consumers, many different separate tasks may **receive** //! messages. //! //! Tokio provides many different channel flavors as different message passing //! patterns are best handled with different implementations. //! //! ## `oneshot` channel //! //! The [`oneshot` channel][oneshot] supports sending a **single** value from a //! single producer to a single consumer. This channel is usually used to send //! the result of a computation to a waiter. //! //! **Example:** using a [`oneshot` channel][oneshot] to receive the result of a //! computation. //! //! ``` //! use tokio::sync::oneshot; //! //! async fn some_computation() -> String { //! "represents the result of the computation".to_string() //! } //! //! #[tokio::main] //! async fn main() { //! let (tx, rx) = oneshot::channel(); //! //! tokio::spawn(async move { //! let res = some_computation().await; //! tx.send(res).unwrap(); //! }); //! //! // Do other work while the computation is happening in the background //! //! // Wait for the computation result //! let res = rx.await.unwrap(); //! } //! ``` //! //! Note, if the task produces a computation result as its final //! action before terminating, the [`JoinHandle`] can be used to //! receive that value instead of allocating resources for the //! `oneshot` channel. Awaiting on [`JoinHandle`] returns `Result`. If //! the task panics, the `Joinhandle` yields `Err` with the panic //! cause. //! //! **Example:** //! //! ``` //! async fn some_computation() -> String { //! "the result of the computation".to_string() //! } //! //! #[tokio::main] //! async fn main() { //! let join_handle = tokio::spawn(async move { //! some_computation().await //! }); //! //! // Do other work while the computation is happening in the background //! //! // Wait for the computation result //! let res = join_handle.await.unwrap(); //! } //! ``` //! //! [`JoinHandle`]: crate::task::JoinHandle //! //! ## `mpsc` channel //! //! The [`mpsc` channel][mpsc] supports sending **many** values from **many** //! producers to a single consumer. This channel is often used to send work to a //! task or to receive the result of many computations. //! //! This is also the channel you should use if you want to send many messages //! from a single producer to a single consumer. There is no dedicated spsc //! channel. //! //! **Example:** using an mpsc to incrementally stream the results of a series //! of computations. //! //! ``` //! use tokio::sync::mpsc; //! //! async fn some_computation(input: u32) -> String { //! format!("the result of computation {}", input) //! } //! //! #[tokio::main] //! async fn main() { //! let (tx, mut rx) = mpsc::channel(100); //! //! tokio::spawn(async move { //! for i in 0..10 { //! let res = some_computation(i).await; //! tx.send(res).await.unwrap(); //! } //! }); //! //! while let Some(res) = rx.recv().await { //! println!("got = {}", res); //! } //! } //! ``` //! //! The argument to `mpsc::channel` is the channel capacity. This is the maximum //! number of values that can be stored in the channel pending receipt at any //! given time. Properly setting this value is key in implementing robust //! programs as the channel capacity plays a critical part in handling back //! pressure. //! //! A common concurrency pattern for resource management is to spawn a task //! dedicated to managing that resource and using message passing between other //! tasks to interact with the resource. The resource may be anything that may //! not be concurrently used. Some examples include a socket and program state. //! For example, if multiple tasks need to send data over a single socket, spawn //! a task to manage the socket and use a channel to synchronize. //! //! **Example:** sending data from many tasks over a single socket using message //! passing. //! //! ```no_run //! use tokio::io::{self, AsyncWriteExt}; //! use tokio::net::TcpStream; //! use tokio::sync::mpsc; //! //! #[tokio::main] //! async fn main() -> io::Result<()> { //! let mut socket = TcpStream::connect("www.example.com:1234").await?; //! let (tx, mut rx) = mpsc::channel(100); //! //! for _ in 0..10 { //! // Each task needs its own `tx` handle. This is done by cloning the //! // original handle. //! let tx = tx.clone(); //! //! tokio::spawn(async move { //! tx.send(&b"data to write"[..]).await.unwrap(); //! }); //! } //! //! // The `rx` half of the channel returns `None` once **all** `tx` clones //! // drop. To ensure `None` is returned, drop the handle owned by the //! // current task. If this `tx` handle is not dropped, there will always //! // be a single outstanding `tx` handle. //! drop(tx); //! //! while let Some(res) = rx.recv().await { //! socket.write_all(res).await?; //! } //! //! Ok(()) //! } //! ``` //! //! The [`mpsc`] and [`oneshot`] channels can be combined to provide a request / //! response type synchronization pattern with a shared resource. A task is //! spawned to synchronize a resource and waits on commands received on a //! [`mpsc`] channel. Each command includes a [`oneshot`] `Sender` on which the //! result of the command is sent. //! //! **Example:** use a task to synchronize a `u64` counter. Each task sends an //! "fetch and increment" command. The counter value **before** the increment is //! sent over the provided `oneshot` channel. //! //! ``` //! use tokio::sync::{oneshot, mpsc}; //! use Command::Increment; //! //! enum Command { //! Increment, //! // Other commands can be added here //! } //! //! #[tokio::main] //! async fn main() { //! let (cmd_tx, mut cmd_rx) = mpsc::channel::<(Command, oneshot::Sender)>(100); //! //! // Spawn a task to manage the counter //! tokio::spawn(async move { //! let mut counter: u64 = 0; //! //! while let Some((cmd, response)) = cmd_rx.recv().await { //! match cmd { //! Increment => { //! let prev = counter; //! counter += 1; //! response.send(prev).unwrap(); //! } //! } //! } //! }); //! //! let mut join_handles = vec![]; //! //! // Spawn tasks that will send the increment command. //! for _ in 0..10 { //! let cmd_tx = cmd_tx.clone(); //! //! join_handles.push(tokio::spawn(async move { //! let (resp_tx, resp_rx) = oneshot::channel(); //! //! cmd_tx.send((Increment, resp_tx)).await.ok().unwrap(); //! let res = resp_rx.await.unwrap(); //! //! println!("previous value = {}", res); //! })); //! } //! //! // Wait for all tasks to complete //! for join_handle in join_handles.drain(..) { //! join_handle.await.unwrap(); //! } //! } //! ``` //! //! ## `broadcast` channel //! //! The [`broadcast` channel] supports sending **many** values from //! **many** producers to **many** consumers. Each consumer will receive //! **each** value. This channel can be used to implement "fan out" style //! patterns common with pub / sub or "chat" systems. //! //! This channel tends to be used less often than `oneshot` and `mpsc` but still //! has its use cases. //! //! This is also the channel you should use if you want to broadcast values from //! a single producer to many consumers. There is no dedicated spmc broadcast //! channel. //! //! Basic usage //! //! ``` //! use tokio::sync::broadcast; //! //! #[tokio::main] //! async fn main() { //! let (tx, mut rx1) = broadcast::channel(16); //! let mut rx2 = tx.subscribe(); //! //! tokio::spawn(async move { //! assert_eq!(rx1.recv().await.unwrap(), 10); //! assert_eq!(rx1.recv().await.unwrap(), 20); //! }); //! //! tokio::spawn(async move { //! assert_eq!(rx2.recv().await.unwrap(), 10); //! assert_eq!(rx2.recv().await.unwrap(), 20); //! }); //! //! tx.send(10).unwrap(); //! tx.send(20).unwrap(); //! } //! ``` //! //! [`broadcast` channel]: crate::sync::broadcast //! //! ## `watch` channel //! //! The [`watch` channel] supports sending **many** values from a **many** //! producer to **many** consumers. However, only the **most recent** value is //! stored in the channel. Consumers are notified when a new value is sent, but //! there is no guarantee that consumers will see **all** values. //! //! The [`watch` channel] is similar to a [`broadcast` channel] with capacity 1. //! //! Use cases for the [`watch` channel] include broadcasting configuration //! changes or signalling program state changes, such as transitioning to //! shutdown. //! //! **Example:** use a [`watch` channel] to notify tasks of configuration //! changes. In this example, a configuration file is checked periodically. When //! the file changes, the configuration changes are signalled to consumers. //! //! ``` //! use tokio::sync::watch; //! use tokio::time::{self, Duration, Instant}; //! //! use std::io; //! //! #[derive(Debug, Clone, Eq, PartialEq)] //! struct Config { //! timeout: Duration, //! } //! //! impl Config { //! async fn load_from_file() -> io::Result { //! // file loading and deserialization logic here //! # Ok(Config { timeout: Duration::from_secs(1) }) //! } //! } //! //! async fn my_async_operation() { //! // Do something here //! } //! //! #[tokio::main] //! async fn main() { //! // Load initial configuration value //! let mut config = Config::load_from_file().await.unwrap(); //! //! // Create the watch channel, initialized with the loaded configuration //! let (tx, rx) = watch::channel(config.clone()); //! //! // Spawn a task to monitor the file. //! tokio::spawn(async move { //! loop { //! // Wait 10 seconds between checks //! time::sleep(Duration::from_secs(10)).await; //! //! // Load the configuration file //! let new_config = Config::load_from_file().await.unwrap(); //! //! // If the configuration changed, send the new config value //! // on the watch channel. //! if new_config != config { //! tx.send(new_config.clone()).unwrap(); //! config = new_config; //! } //! } //! }); //! //! let mut handles = vec![]; //! //! // Spawn tasks that runs the async operation for at most `timeout`. If //! // the timeout elapses, restart the operation. //! // //! // The task simultaneously watches the `Config` for changes. When the //! // timeout duration changes, the timeout is updated without restarting //! // the in-flight operation. //! for _ in 0..5 { //! // Clone a config watch handle for use in this task //! let mut rx = rx.clone(); //! //! let handle = tokio::spawn(async move { //! // Start the initial operation and pin the future to the stack. //! // Pinning to the stack is required to resume the operation //! // across multiple calls to `select!` //! let op = my_async_operation(); //! tokio::pin!(op); //! //! // Get the initial config value //! let mut conf = rx.borrow().clone(); //! //! let mut op_start = Instant::now(); //! let sleep = time::sleep_until(op_start + conf.timeout); //! tokio::pin!(sleep); //! //! loop { //! tokio::select! { //! _ = &mut sleep => { //! // The operation elapsed. Restart it //! op.set(my_async_operation()); //! //! // Track the new start time //! op_start = Instant::now(); //! //! // Restart the timeout //! sleep.set(time::sleep_until(op_start + conf.timeout)); //! } //! _ = rx.changed() => { //! conf = rx.borrow_and_update().clone(); //! //! // The configuration has been updated. Update the //! // `sleep` using the new `timeout` value. //! sleep.as_mut().reset(op_start + conf.timeout); //! } //! _ = &mut op => { //! // The operation completed! //! return //! } //! } //! } //! }); //! //! handles.push(handle); //! } //! //! for handle in handles.drain(..) { //! handle.await.unwrap(); //! } //! } //! ``` //! //! [`watch` channel]: mod@crate::sync::watch //! [`broadcast` channel]: mod@crate::sync::broadcast //! //! # State synchronization //! //! The remaining synchronization primitives focus on synchronizing state. //! These are asynchronous equivalents to versions provided by `std`. They //! operate in a similar way as their `std` counterparts but will wait //! asynchronously instead of blocking the thread. //! //! * [`Barrier`] Ensures multiple tasks will wait for each other to reach a //! point in the program, before continuing execution all together. //! //! * [`Mutex`] Mutual Exclusion mechanism, which ensures that at most one //! thread at a time is able to access some data. //! //! * [`Notify`] Basic task notification. `Notify` supports notifying a //! receiving task without sending data. In this case, the task wakes up and //! resumes processing. //! //! * [`RwLock`] Provides a mutual exclusion mechanism which allows multiple //! readers at the same time, while allowing only one writer at a time. In //! some cases, this can be more efficient than a mutex. //! //! * [`Semaphore`] Limits the amount of concurrency. A semaphore holds a //! number of permits, which tasks may request in order to enter a critical //! section. Semaphores are useful for implementing limiting or bounding of //! any kind. //! //! # Runtime compatibility //! //! All synchronization primitives provided in this module are runtime agnostic. //! You can freely move them between different instances of the Tokio runtime //! or even use them from non-Tokio runtimes. //! //! When used in a Tokio runtime, the synchronization primitives participate in //! [cooperative scheduling](crate::task#cooperative-scheduling) to avoid //! starvation. This feature does not apply when used from non-Tokio runtimes. //! //! As an exception, methods ending in `_timeout` are not runtime agnostic //! because they require access to the Tokio timer. See the documentation of //! each `*_timeout` method for more information on its use. cfg_sync! { /// Named future types. pub mod futures { pub use super::notify::Notified; } mod barrier; pub use barrier::{Barrier, BarrierWaitResult}; pub mod broadcast; pub mod mpsc; mod mutex; pub use mutex::{Mutex, MutexGuard, TryLockError, OwnedMutexGuard, MappedMutexGuard, OwnedMappedMutexGuard}; pub(crate) mod notify; pub use notify::Notify; pub mod oneshot; pub(crate) mod batch_semaphore; pub use batch_semaphore::{AcquireError, TryAcquireError}; mod semaphore; pub use semaphore::{Semaphore, SemaphorePermit, OwnedSemaphorePermit}; mod rwlock; pub use rwlock::RwLock; pub use rwlock::owned_read_guard::OwnedRwLockReadGuard; pub use rwlock::owned_write_guard::OwnedRwLockWriteGuard; pub use rwlock::owned_write_guard_mapped::OwnedRwLockMappedWriteGuard; pub use rwlock::read_guard::RwLockReadGuard; pub use rwlock::write_guard::RwLockWriteGuard; pub use rwlock::write_guard_mapped::RwLockMappedWriteGuard; mod task; pub(crate) use task::AtomicWaker; mod once_cell; pub use self::once_cell::{OnceCell, SetError}; pub mod watch; } cfg_not_sync! { cfg_fs! { pub(crate) mod batch_semaphore; mod mutex; pub(crate) use mutex::Mutex; } #[cfg(any(feature = "rt", feature = "signal", all(unix, feature = "process")))] pub(crate) mod notify; #[cfg(any(feature = "rt", all(windows, feature = "process")))] pub(crate) mod oneshot; cfg_atomic_waker_impl! { mod task; pub(crate) use task::AtomicWaker; } #[cfg(any(feature = "signal", all(unix, feature = "process")))] pub(crate) mod watch; } /// Unit tests #[cfg(test)] mod tests; tokio-1.43.1/src/sync/mpsc/block.rs000064400000000000000000000356141046102023000152150ustar 00000000000000use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize}; use std::alloc::Layout; use std::mem::MaybeUninit; use std::ops; use std::ptr::{self, NonNull}; use std::sync::atomic::Ordering::{self, AcqRel, Acquire, Release}; /// A block in a linked list. /// /// Each block in the list can hold up to `BLOCK_CAP` messages. pub(crate) struct Block { /// The header fields. header: BlockHeader, /// Array containing values pushed into the block. Values are stored in a /// continuous array in order to improve cache line behavior when reading. /// The values must be manually dropped. values: Values, } /// Extra fields for a `Block`. struct BlockHeader { /// The start index of this block. /// /// Slots in this block have indices in `start_index .. start_index + BLOCK_CAP`. start_index: usize, /// The next block in the linked list. next: AtomicPtr>, /// Bitfield tracking slots that are ready to have their values consumed. ready_slots: AtomicUsize, /// The observed `tail_position` value *after* the block has been passed by /// `block_tail`. observed_tail_position: UnsafeCell, } pub(crate) enum Read { Value(T), Closed, } #[repr(transparent)] struct Values([UnsafeCell>; BLOCK_CAP]); use super::BLOCK_CAP; /// Masks an index to get the block identifier. const BLOCK_MASK: usize = !(BLOCK_CAP - 1); /// Masks an index to get the value offset in a block. const SLOT_MASK: usize = BLOCK_CAP - 1; /// Flag tracking that a block has gone through the sender's release routine. /// /// When this is set, the receiver may consider freeing the block. const RELEASED: usize = 1 << BLOCK_CAP; /// Flag tracking all senders dropped. /// /// When this flag is set, the send half of the channel has closed. const TX_CLOSED: usize = RELEASED << 1; /// Mask covering all bits used to track slot readiness. const READY_MASK: usize = RELEASED - 1; /// Returns the index of the first slot in the block referenced by `slot_index`. #[inline(always)] pub(crate) fn start_index(slot_index: usize) -> usize { BLOCK_MASK & slot_index } /// Returns the offset into the block referenced by `slot_index`. #[inline(always)] pub(crate) fn offset(slot_index: usize) -> usize { SLOT_MASK & slot_index } generate_addr_of_methods! { impl Block { unsafe fn addr_of_header(self: NonNull) -> NonNull> { &self.header } unsafe fn addr_of_values(self: NonNull) -> NonNull> { &self.values } } } impl Block { pub(crate) fn new(start_index: usize) -> Box> { unsafe { // Allocate the block on the heap. // SAFETY: The size of the Block is non-zero, since it is at least the size of the header. let block = std::alloc::alloc(Layout::new::>()) as *mut Block; let block = match NonNull::new(block) { Some(block) => block, None => std::alloc::handle_alloc_error(Layout::new::>()), }; // Write the header to the block. Block::addr_of_header(block).as_ptr().write(BlockHeader { // The absolute index in the channel of the first slot in the block. start_index, // Pointer to the next block in the linked list. next: AtomicPtr::new(ptr::null_mut()), ready_slots: AtomicUsize::new(0), observed_tail_position: UnsafeCell::new(0), }); // Initialize the values array. Values::initialize(Block::addr_of_values(block)); // Convert the pointer to a `Box`. // Safety: The raw pointer was allocated using the global allocator, and with // the layout for a `Block`, so it's valid to convert it to box. Box::from_raw(block.as_ptr()) } } /// Returns `true` if the block matches the given index. pub(crate) fn is_at_index(&self, index: usize) -> bool { debug_assert!(offset(index) == 0); self.header.start_index == index } /// Returns the number of blocks between `self` and the block at the /// specified index. /// /// `start_index` must represent a block *after* `self`. pub(crate) fn distance(&self, other_index: usize) -> usize { debug_assert!(offset(other_index) == 0); other_index.wrapping_sub(self.header.start_index) / BLOCK_CAP } /// Reads the value at the given offset. /// /// Returns `None` if the slot is empty. /// /// # Safety /// /// To maintain safety, the caller must ensure: /// /// * No concurrent access to the slot. pub(crate) unsafe fn read(&self, slot_index: usize) -> Option> { let offset = offset(slot_index); let ready_bits = self.header.ready_slots.load(Acquire); if !is_ready(ready_bits, offset) { if is_tx_closed(ready_bits) { return Some(Read::Closed); } return None; } // Get the value let value = self.values[offset].with(|ptr| ptr::read(ptr)); Some(Read::Value(value.assume_init())) } /// Returns true if *this* block has a value in the given slot. /// /// Always returns false when given an index from a different block. pub(crate) fn has_value(&self, slot_index: usize) -> bool { if slot_index < self.header.start_index { return false; } if slot_index >= self.header.start_index + super::BLOCK_CAP { return false; } let offset = offset(slot_index); let ready_bits = self.header.ready_slots.load(Acquire); is_ready(ready_bits, offset) } /// Writes a value to the block at the given offset. /// /// # Safety /// /// To maintain safety, the caller must ensure: /// /// * The slot is empty. /// * No concurrent access to the slot. pub(crate) unsafe fn write(&self, slot_index: usize, value: T) { // Get the offset into the block let slot_offset = offset(slot_index); self.values[slot_offset].with_mut(|ptr| { ptr::write(ptr, MaybeUninit::new(value)); }); // Release the value. After this point, the slot ref may no longer // be used. It is possible for the receiver to free the memory at // any point. self.set_ready(slot_offset); } /// Signal to the receiver that the sender half of the list is closed. pub(crate) unsafe fn tx_close(&self) { self.header.ready_slots.fetch_or(TX_CLOSED, Release); } pub(crate) unsafe fn is_closed(&self) -> bool { let ready_bits = self.header.ready_slots.load(Acquire); is_tx_closed(ready_bits) } /// Resets the block to a blank state. This enables reusing blocks in the /// channel. /// /// # Safety /// /// To maintain safety, the caller must ensure: /// /// * All slots are empty. /// * The caller holds a unique pointer to the block. pub(crate) unsafe fn reclaim(&mut self) { self.header.start_index = 0; self.header.next = AtomicPtr::new(ptr::null_mut()); self.header.ready_slots = AtomicUsize::new(0); } /// Releases the block to the rx half for freeing. /// /// This function is called by the tx half once it can be guaranteed that no /// more senders will attempt to access the block. /// /// # Safety /// /// To maintain safety, the caller must ensure: /// /// * The block will no longer be accessed by any sender. pub(crate) unsafe fn tx_release(&self, tail_position: usize) { // Track the observed tail_position. Any sender targeting a greater // tail_position is guaranteed to not access this block. self.header .observed_tail_position .with_mut(|ptr| *ptr = tail_position); // Set the released bit, signalling to the receiver that it is safe to // free the block's memory as soon as all slots **prior** to // `observed_tail_position` have been filled. self.header.ready_slots.fetch_or(RELEASED, Release); } /// Mark a slot as ready fn set_ready(&self, slot: usize) { let mask = 1 << slot; self.header.ready_slots.fetch_or(mask, Release); } /// Returns `true` when all slots have their `ready` bits set. /// /// This indicates that the block is in its final state and will no longer /// be mutated. pub(crate) fn is_final(&self) -> bool { self.header.ready_slots.load(Acquire) & READY_MASK == READY_MASK } /// Returns the `observed_tail_position` value, if set pub(crate) fn observed_tail_position(&self) -> Option { if 0 == RELEASED & self.header.ready_slots.load(Acquire) { None } else { Some( self.header .observed_tail_position .with(|ptr| unsafe { *ptr }), ) } } /// Loads the next block pub(crate) fn load_next(&self, ordering: Ordering) -> Option>> { let ret = NonNull::new(self.header.next.load(ordering)); debug_assert!(unsafe { ret.map_or(true, |block| { block.as_ref().header.start_index == self.header.start_index.wrapping_add(BLOCK_CAP) }) }); ret } /// Pushes `block` as the next block in the link. /// /// Returns Ok if successful, otherwise, a pointer to the next block in /// the list is returned. /// /// This requires that the next pointer is null. /// /// # Ordering /// /// This performs a compare-and-swap on `next` using `AcqRel` ordering. /// /// # Safety /// /// To maintain safety, the caller must ensure: /// /// * `block` is not freed until it has been removed from the list. pub(crate) unsafe fn try_push( &self, block: &mut NonNull>, success: Ordering, failure: Ordering, ) -> Result<(), NonNull>> { block.as_mut().header.start_index = self.header.start_index.wrapping_add(BLOCK_CAP); let next_ptr = self .header .next .compare_exchange(ptr::null_mut(), block.as_ptr(), success, failure) .unwrap_or_else(|x| x); match NonNull::new(next_ptr) { Some(next_ptr) => Err(next_ptr), None => Ok(()), } } /// Grows the `Block` linked list by allocating and appending a new block. /// /// The next block in the linked list is returned. This may or may not be /// the one allocated by the function call. /// /// # Implementation /// /// It is assumed that `self.next` is null. A new block is allocated with /// `start_index` set to be the next block. A compare-and-swap is performed /// with `AcqRel` memory ordering. If the compare-and-swap is successful, the /// newly allocated block is released to other threads walking the block /// linked list. If the compare-and-swap fails, the current thread acquires /// the next block in the linked list, allowing the current thread to access /// the slots. pub(crate) fn grow(&self) -> NonNull> { // Create the new block. It is assumed that the block will become the // next one after `&self`. If this turns out to not be the case, // `start_index` is updated accordingly. let new_block = Block::new(self.header.start_index + BLOCK_CAP); let mut new_block = unsafe { NonNull::new_unchecked(Box::into_raw(new_block)) }; // Attempt to store the block. The first compare-and-swap attempt is // "unrolled" due to minor differences in logic // // `AcqRel` is used as the ordering **only** when attempting the // compare-and-swap on self.next. // // If the compare-and-swap fails, then the actual value of the cell is // returned from this function and accessed by the caller. Given this, // the memory must be acquired. // // `Release` ensures that the newly allocated block is available to // other threads acquiring the next pointer. let next = NonNull::new( self.header .next .compare_exchange(ptr::null_mut(), new_block.as_ptr(), AcqRel, Acquire) .unwrap_or_else(|x| x), ); let next = match next { Some(next) => next, None => { // The compare-and-swap succeeded and the newly allocated block // is successfully pushed. return new_block; } }; // There already is a next block in the linked list. The newly allocated // block could be dropped and the discovered next block returned; // however, that would be wasteful. Instead, the linked list is walked // by repeatedly attempting to compare-and-swap the pointer into the // `next` register until the compare-and-swap succeed. // // Care is taken to update new_block's start_index field as appropriate. let mut curr = next; // TODO: Should this iteration be capped? loop { let actual = unsafe { curr.as_ref().try_push(&mut new_block, AcqRel, Acquire) }; curr = match actual { Ok(()) => { return next; } Err(curr) => curr, }; crate::loom::thread::yield_now(); } } } /// Returns `true` if the specified slot has a value ready to be consumed. fn is_ready(bits: usize, slot: usize) -> bool { let mask = 1 << slot; mask == mask & bits } /// Returns `true` if the closed flag has been set. fn is_tx_closed(bits: usize) -> bool { TX_CLOSED == bits & TX_CLOSED } impl Values { /// Initialize a `Values` struct from a pointer. /// /// # Safety /// /// The raw pointer must be valid for writing a `Values`. unsafe fn initialize(_value: NonNull>) { // When fuzzing, `UnsafeCell` needs to be initialized. if_loom! { let p = _value.as_ptr() as *mut UnsafeCell>; for i in 0..BLOCK_CAP { p.add(i) .write(UnsafeCell::new(MaybeUninit::uninit())); } } } } impl ops::Index for Values { type Output = UnsafeCell>; fn index(&self, index: usize) -> &Self::Output { self.0.index(index) } } #[cfg(all(test, not(loom)))] #[test] fn assert_no_stack_overflow() { // https://github.com/tokio-rs/tokio/issues/5293 struct Foo { _a: [u8; 2_000_000], } assert_eq!( Layout::new::>>(), Layout::new::>() ); let _block = Block::::new(0); } tokio-1.43.1/src/sync/mpsc/bounded.rs000064400000000000000000001765411046102023000155500ustar 00000000000000use crate::loom::sync::Arc; use crate::sync::batch_semaphore::{self as semaphore, TryAcquireError}; use crate::sync::mpsc::chan; use crate::sync::mpsc::error::{SendError, TryRecvError, TrySendError}; cfg_time! { use crate::sync::mpsc::error::SendTimeoutError; use crate::time::Duration; } use std::fmt; use std::task::{Context, Poll}; /// Sends values to the associated `Receiver`. /// /// Instances are created by the [`channel`] function. /// /// To convert the `Sender` into a `Sink` or use it in a poll function, you can /// use the [`PollSender`] utility. /// /// [`PollSender`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSender.html pub struct Sender { chan: chan::Tx, } /// A sender that does not prevent the channel from being closed. /// /// If all [`Sender`] instances of a channel were dropped and only `WeakSender` /// instances remain, the channel is closed. /// /// In order to send messages, the `WeakSender` needs to be upgraded using /// [`WeakSender::upgrade`], which returns `Option`. It returns `None` /// if all `Sender`s have been dropped, and otherwise it returns a `Sender`. /// /// [`Sender`]: Sender /// [`WeakSender::upgrade`]: WeakSender::upgrade /// /// # Examples /// /// ``` /// use tokio::sync::mpsc::channel; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = channel::(15); /// let tx_weak = tx.downgrade(); /// /// // Upgrading will succeed because `tx` still exists. /// assert!(tx_weak.upgrade().is_some()); /// /// // If we drop `tx`, then it will fail. /// drop(tx); /// assert!(tx_weak.clone().upgrade().is_none()); /// } /// ``` pub struct WeakSender { chan: Arc>, } /// Permits to send one value into the channel. /// /// `Permit` values are returned by [`Sender::reserve()`] and [`Sender::try_reserve()`] /// and are used to guarantee channel capacity before generating a message to send. /// /// [`Sender::reserve()`]: Sender::reserve /// [`Sender::try_reserve()`]: Sender::try_reserve pub struct Permit<'a, T> { chan: &'a chan::Tx, } /// An [`Iterator`] of [`Permit`] that can be used to hold `n` slots in the channel. /// /// `PermitIterator` values are returned by [`Sender::reserve_many()`] and [`Sender::try_reserve_many()`] /// and are used to guarantee channel capacity before generating `n` messages to send. /// /// [`Sender::reserve_many()`]: Sender::reserve_many /// [`Sender::try_reserve_many()`]: Sender::try_reserve_many pub struct PermitIterator<'a, T> { chan: &'a chan::Tx, n: usize, } /// Owned permit to send one value into the channel. /// /// This is identical to the [`Permit`] type, except that it moves the sender /// rather than borrowing it. /// /// `OwnedPermit` values are returned by [`Sender::reserve_owned()`] and /// [`Sender::try_reserve_owned()`] and are used to guarantee channel capacity /// before generating a message to send. /// /// [`Permit`]: Permit /// [`Sender::reserve_owned()`]: Sender::reserve_owned /// [`Sender::try_reserve_owned()`]: Sender::try_reserve_owned pub struct OwnedPermit { chan: Option>, } /// Receives values from the associated `Sender`. /// /// Instances are created by the [`channel`] function. /// /// This receiver can be turned into a `Stream` using [`ReceiverStream`]. /// /// [`ReceiverStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.ReceiverStream.html pub struct Receiver { /// The channel receiver. chan: chan::Rx, } /// Creates a bounded mpsc channel for communicating between asynchronous tasks /// with backpressure. /// /// The channel will buffer up to the provided number of messages. Once the /// buffer is full, attempts to send new messages will wait until a message is /// received from the channel. The provided buffer capacity must be at least 1. /// /// All data sent on `Sender` will become available on `Receiver` in the same /// order as it was sent. /// /// The `Sender` can be cloned to `send` to the same channel from multiple code /// locations. Only one `Receiver` is supported. /// /// If the `Receiver` is disconnected while trying to `send`, the `send` method /// will return a `SendError`. Similarly, if `Sender` is disconnected while /// trying to `recv`, the `recv` method will return `None`. /// /// # Panics /// /// Panics if the buffer capacity is 0. /// /// # Examples /// /// ```rust /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tokio::spawn(async move { /// for i in 0..10 { /// if let Err(_) = tx.send(i).await { /// println!("receiver dropped"); /// return; /// } /// } /// }); /// /// while let Some(i) = rx.recv().await { /// println!("got = {}", i); /// } /// } /// ``` #[track_caller] pub fn channel(buffer: usize) -> (Sender, Receiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); let semaphore = Semaphore { semaphore: semaphore::Semaphore::new(buffer), bound: buffer, }; let (tx, rx) = chan::channel(semaphore); let tx = Sender::new(tx); let rx = Receiver::new(rx); (tx, rx) } /// Channel semaphore is a tuple of the semaphore implementation and a `usize` /// representing the channel bound. #[derive(Debug)] pub(crate) struct Semaphore { pub(crate) semaphore: semaphore::Semaphore, pub(crate) bound: usize, } impl Receiver { pub(crate) fn new(chan: chan::Rx) -> Receiver { Receiver { chan } } /// Receives the next value for this receiver. /// /// This method returns `None` if the channel has been closed and there are /// no remaining messages in the channel's buffer. This indicates that no /// further values can ever be received from this `Receiver`. The channel is /// closed when all senders have been dropped, or when [`close`] is called. /// /// If there are no messages in the channel's buffer, but the channel has /// not yet been closed, this method will sleep until a message is sent or /// the channel is closed. Note that if [`close`] is called, but there are /// still outstanding [`Permits`] from before it was closed, the channel is /// not considered closed by `recv` until the permits are released. /// /// # Cancel safety /// /// This method is cancel safe. If `recv` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// channel. /// /// [`close`]: Self::close /// [`Permits`]: struct@crate::sync::mpsc::Permit /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tokio::spawn(async move { /// tx.send("hello").await.unwrap(); /// }); /// /// assert_eq!(Some("hello"), rx.recv().await); /// assert_eq!(None, rx.recv().await); /// } /// ``` /// /// Values are buffered: /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tx.send("hello").await.unwrap(); /// tx.send("world").await.unwrap(); /// /// assert_eq!(Some("hello"), rx.recv().await); /// assert_eq!(Some("world"), rx.recv().await); /// } /// ``` pub async fn recv(&mut self) -> Option { use std::future::poll_fn; poll_fn(|cx| self.chan.recv(cx)).await } /// Receives the next values for this receiver and extends `buffer`. /// /// This method extends `buffer` by no more than a fixed number of values /// as specified by `limit`. If `limit` is zero, the function immediately /// returns `0`. The return value is the number of values added to `buffer`. /// /// For `limit > 0`, if there are no messages in the channel's queue, but /// the channel has not yet been closed, this method will sleep until a /// message is sent or the channel is closed. Note that if [`close`] is /// called, but there are still outstanding [`Permits`] from before it was /// closed, the channel is not considered closed by `recv_many` until the /// permits are released. /// /// For non-zero values of `limit`, this method will never return `0` unless /// the channel has been closed and there are no remaining messages in the /// channel's queue. This indicates that no further values can ever be /// received from this `Receiver`. The channel is closed when all senders /// have been dropped, or when [`close`] is called. /// /// The capacity of `buffer` is increased as needed. /// /// # Cancel safety /// /// This method is cancel safe. If `recv_many` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// channel. /// /// [`close`]: Self::close /// [`Permits`]: struct@crate::sync::mpsc::Permit /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let mut buffer: Vec<&str> = Vec::with_capacity(2); /// let limit = 2; /// let (tx, mut rx) = mpsc::channel(100); /// let tx2 = tx.clone(); /// tx2.send("first").await.unwrap(); /// tx2.send("second").await.unwrap(); /// tx2.send("third").await.unwrap(); /// /// // Call `recv_many` to receive up to `limit` (2) values. /// assert_eq!(2, rx.recv_many(&mut buffer, limit).await); /// assert_eq!(vec!["first", "second"], buffer); /// /// // If the buffer is full, the next call to `recv_many` /// // reserves additional capacity. /// assert_eq!(1, rx.recv_many(&mut buffer, 1).await); /// /// tokio::spawn(async move { /// tx.send("fourth").await.unwrap(); /// }); /// /// // 'tx' is dropped, but `recv_many` /// // is guaranteed not to return 0 as the channel /// // is not yet closed. /// assert_eq!(1, rx.recv_many(&mut buffer, 1).await); /// assert_eq!(vec!["first", "second", "third", "fourth"], buffer); /// /// // Once the last sender is dropped, the channel is /// // closed and `recv_many` returns 0, capacity unchanged. /// drop(tx2); /// assert_eq!(0, rx.recv_many(&mut buffer, limit).await); /// assert_eq!(vec!["first", "second", "third", "fourth"], buffer); /// } /// ``` pub async fn recv_many(&mut self, buffer: &mut Vec, limit: usize) -> usize { use std::future::poll_fn; poll_fn(|cx| self.chan.recv_many(cx, buffer, limit)).await } /// Tries to receive the next value for this receiver. /// /// This method returns the [`Empty`] error if the channel is currently /// empty, but there are still outstanding [senders] or [permits]. /// /// This method returns the [`Disconnected`] error if the channel is /// currently empty, and there are no outstanding [senders] or [permits]. /// /// Unlike the [`poll_recv`] method, this method will never return an /// [`Empty`] error spuriously. /// /// [`Empty`]: crate::sync::mpsc::error::TryRecvError::Empty /// [`Disconnected`]: crate::sync::mpsc::error::TryRecvError::Disconnected /// [`poll_recv`]: Self::poll_recv /// [senders]: crate::sync::mpsc::Sender /// [permits]: crate::sync::mpsc::Permit /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// use tokio::sync::mpsc::error::TryRecvError; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tx.send("hello").await.unwrap(); /// /// assert_eq!(Ok("hello"), rx.try_recv()); /// assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); /// /// tx.send("hello").await.unwrap(); /// // Drop the last sender, closing the channel. /// drop(tx); /// /// assert_eq!(Ok("hello"), rx.try_recv()); /// assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); /// } /// ``` pub fn try_recv(&mut self) -> Result { self.chan.try_recv() } /// Blocking receive to call outside of asynchronous contexts. /// /// This method returns `None` if the channel has been closed and there are /// no remaining messages in the channel's buffer. This indicates that no /// further values can ever be received from this `Receiver`. The channel is /// closed when all senders have been dropped, or when [`close`] is called. /// /// If there are no messages in the channel's buffer, but the channel has /// not yet been closed, this method will block until a message is sent or /// the channel is closed. /// /// This method is intended for use cases where you are sending from /// asynchronous code to synchronous code, and will work even if the sender /// is not using [`blocking_send`] to send the message. /// /// Note that if [`close`] is called, but there are still outstanding /// [`Permits`] from before it was closed, the channel is not considered /// closed by `blocking_recv` until the permits are released. /// /// [`close`]: Self::close /// [`Permits`]: struct@crate::sync::mpsc::Permit /// [`blocking_send`]: fn@crate::sync::mpsc::Sender::blocking_send /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// /// ``` /// use std::thread; /// use tokio::runtime::Runtime; /// use tokio::sync::mpsc; /// /// fn main() { /// let (tx, mut rx) = mpsc::channel::(10); /// /// let sync_code = thread::spawn(move || { /// assert_eq!(Some(10), rx.blocking_recv()); /// }); /// /// Runtime::new() /// .unwrap() /// .block_on(async move { /// let _ = tx.send(10).await; /// }); /// sync_code.join().unwrap() /// } /// ``` #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(&mut self) -> Option { crate::future::block_on(self.recv()) } /// Variant of [`Self::recv_many`] for blocking contexts. /// /// The same conditions as in [`Self::blocking_recv`] apply. #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "recv_many_blocking"))] pub fn blocking_recv_many(&mut self, buffer: &mut Vec, limit: usize) -> usize { crate::future::block_on(self.recv_many(buffer, limit)) } /// Closes the receiving half of a channel without dropping it. /// /// This prevents any further messages from being sent on the channel while /// still enabling the receiver to drain messages that are buffered. Any /// outstanding [`Permit`] values will still be able to send messages. /// /// To guarantee that no messages are dropped, after calling `close()`, /// `recv()` must be called until `None` is returned. If there are /// outstanding [`Permit`] or [`OwnedPermit`] values, the `recv` method will /// not return `None` until those are released. /// /// [`Permit`]: Permit /// [`OwnedPermit`]: OwnedPermit /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(20); /// /// tokio::spawn(async move { /// let mut i = 0; /// while let Ok(permit) = tx.reserve().await { /// permit.send(i); /// i += 1; /// } /// }); /// /// rx.close(); /// /// while let Some(msg) = rx.recv().await { /// println!("got {}", msg); /// } /// /// // Channel closed and no messages are lost. /// } /// ``` pub fn close(&mut self) { self.chan.close(); } /// Checks if a channel is closed. /// /// This method returns `true` if the channel has been closed. The channel is closed /// when all [`Sender`] have been dropped, or when [`Receiver::close`] is called. /// /// [`Sender`]: crate::sync::mpsc::Sender /// [`Receiver::close`]: crate::sync::mpsc::Receiver::close /// /// # Examples /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (_tx, mut rx) = mpsc::channel::<()>(10); /// assert!(!rx.is_closed()); /// /// rx.close(); /// /// assert!(rx.is_closed()); /// } /// ``` pub fn is_closed(&self) -> bool { self.chan.is_closed() } /// Checks if a channel is empty. /// /// This method returns `true` if the channel has no messages. /// /// # Examples /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::channel(10); /// assert!(rx.is_empty()); /// /// tx.send(0).await.unwrap(); /// assert!(!rx.is_empty()); /// } /// /// ``` pub fn is_empty(&self) -> bool { self.chan.is_empty() } /// Returns the number of messages in the channel. /// /// # Examples /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::channel(10); /// assert_eq!(0, rx.len()); /// /// tx.send(0).await.unwrap(); /// assert_eq!(1, rx.len()); /// } /// ``` pub fn len(&self) -> usize { self.chan.len() } /// Returns the current capacity of the channel. /// /// The capacity goes down when the sender sends a value by calling [`Sender::send`] or by reserving /// capacity with [`Sender::reserve`]. The capacity goes up when values are received. /// This is distinct from [`max_capacity`], which always returns buffer capacity initially /// specified when calling [`channel`]. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel::<()>(5); /// /// assert_eq!(rx.capacity(), 5); /// /// // Making a reservation drops the capacity by one. /// let permit = tx.reserve().await.unwrap(); /// assert_eq!(rx.capacity(), 4); /// assert_eq!(rx.len(), 0); /// /// // Sending and receiving a value increases the capacity by one. /// permit.send(()); /// assert_eq!(rx.len(), 1); /// rx.recv().await.unwrap(); /// assert_eq!(rx.capacity(), 5); /// /// // Directly sending a message drops the capacity by one. /// tx.send(()).await.unwrap(); /// assert_eq!(rx.capacity(), 4); /// assert_eq!(rx.len(), 1); /// /// // Receiving the message increases the capacity by one. /// rx.recv().await.unwrap(); /// assert_eq!(rx.capacity(), 5); /// assert_eq!(rx.len(), 0); /// } /// ``` /// [`capacity`]: Receiver::capacity /// [`max_capacity`]: Receiver::max_capacity pub fn capacity(&self) -> usize { self.chan.semaphore().semaphore.available_permits() } /// Returns the maximum buffer capacity of the channel. /// /// The maximum capacity is the buffer capacity initially specified when calling /// [`channel`]. This is distinct from [`capacity`], which returns the *current* /// available buffer capacity: as messages are sent and received, the value /// returned by [`capacity`] will go up or down, whereas the value /// returned by [`max_capacity`] will remain constant. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::channel::<()>(5); /// /// // both max capacity and capacity are the same at first /// assert_eq!(rx.max_capacity(), 5); /// assert_eq!(rx.capacity(), 5); /// /// // Making a reservation doesn't change the max capacity. /// let permit = tx.reserve().await.unwrap(); /// assert_eq!(rx.max_capacity(), 5); /// // but drops the capacity by one /// assert_eq!(rx.capacity(), 4); /// } /// ``` /// [`capacity`]: Receiver::capacity /// [`max_capacity`]: Receiver::max_capacity pub fn max_capacity(&self) -> usize { self.chan.semaphore().bound } /// Polls to receive the next message on this channel. /// /// This method returns: /// /// * `Poll::Pending` if no messages are available but the channel is not /// closed, or if a spurious failure happens. /// * `Poll::Ready(Some(message))` if a message is available. /// * `Poll::Ready(None)` if the channel has been closed and all messages /// sent before it was closed have been received. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when a message is sent on any /// receiver, or when the channel is closed. Note that on multiple calls to /// `poll_recv` or `poll_recv_many`, only the `Waker` from the `Context` /// passed to the most recent call is scheduled to receive a wakeup. /// /// If this method returns `Poll::Pending` due to a spurious failure, then /// the `Waker` will be notified when the situation causing the spurious /// failure has been resolved. Note that receiving such a wakeup does not /// guarantee that the next call will succeed — it could fail with another /// spurious failure. pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.chan.recv(cx) } /// Polls to receive multiple messages on this channel, extending the provided buffer. /// /// This method returns: /// * `Poll::Pending` if no messages are available but the channel is not closed, or if a /// spurious failure happens. /// * `Poll::Ready(count)` where `count` is the number of messages successfully received and /// stored in `buffer`. This can be less than, or equal to, `limit`. /// * `Poll::Ready(0)` if `limit` is set to zero or when the channel is closed. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when a message is sent on any /// receiver, or when the channel is closed. Note that on multiple calls to /// `poll_recv` or `poll_recv_many`, only the `Waker` from the `Context` /// passed to the most recent call is scheduled to receive a wakeup. /// /// Note that this method does not guarantee that exactly `limit` messages /// are received. Rather, if at least one message is available, it returns /// as many messages as it can up to the given limit. This method returns /// zero only if the channel is closed (or if `limit` is zero). /// /// # Examples /// /// ``` /// use std::task::{Context, Poll}; /// use std::pin::Pin; /// use tokio::sync::mpsc; /// use futures::Future; /// /// struct MyReceiverFuture<'a> { /// receiver: mpsc::Receiver, /// buffer: &'a mut Vec, /// limit: usize, /// } /// /// impl<'a> Future for MyReceiverFuture<'a> { /// type Output = usize; // Number of messages received /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// let MyReceiverFuture { receiver, buffer, limit } = &mut *self; /// /// // Now `receiver` and `buffer` are mutable references, and `limit` is copied /// match receiver.poll_recv_many(cx, *buffer, *limit) { /// Poll::Pending => Poll::Pending, /// Poll::Ready(count) => Poll::Ready(count), /// } /// } /// } /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::channel(32); /// let mut buffer = Vec::new(); /// /// let my_receiver_future = MyReceiverFuture { /// receiver: rx, /// buffer: &mut buffer, /// limit: 3, /// }; /// /// for i in 0..10 { /// tx.send(i).await.unwrap(); /// } /// /// let count = my_receiver_future.await; /// assert_eq!(count, 3); /// assert_eq!(buffer, vec![0,1,2]) /// } /// ``` pub fn poll_recv_many( &mut self, cx: &mut Context<'_>, buffer: &mut Vec, limit: usize, ) -> Poll { self.chan.recv_many(cx, buffer, limit) } /// Returns the number of [`Sender`] handles. pub fn sender_strong_count(&self) -> usize { self.chan.sender_strong_count() } /// Returns the number of [`WeakSender`] handles. pub fn sender_weak_count(&self) -> usize { self.chan.sender_weak_count() } } impl fmt::Debug for Receiver { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Receiver") .field("chan", &self.chan) .finish() } } impl Unpin for Receiver {} impl Sender { pub(crate) fn new(chan: chan::Tx) -> Sender { Sender { chan } } /// Sends a value, waiting until there is capacity. /// /// A successful send occurs when it is determined that the other end of the /// channel has not hung up already. An unsuccessful send would be one where /// the corresponding receiver has already been closed. Note that a return /// value of `Err` means that the data will never be received, but a return /// value of `Ok` does not mean that the data will be received. It is /// possible for the corresponding receiver to hang up immediately after /// this function returns `Ok`. /// /// # Errors /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`Receiver`] handle dropping, the function returns /// an error. The error includes the value passed to `send`. /// /// [`close`]: Receiver::close /// [`Receiver`]: Receiver /// /// # Cancel safety /// /// If `send` is used as the event in a [`tokio::select!`](crate::select) /// statement and some other branch completes first, then it is guaranteed /// that the message was not sent. **However, in that case, the message /// is dropped and will be lost.** /// /// To avoid losing messages, use [`reserve`](Self::reserve) to reserve /// capacity, then use the returned [`Permit`] to send the message. /// /// This channel uses a queue to ensure that calls to `send` and `reserve` /// complete in the order they were requested. Cancelling a call to /// `send` makes you lose your place in the queue. /// /// # Examples /// /// In the following example, each call to `send` will block until the /// previously sent value was received. /// /// ```rust /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// tokio::spawn(async move { /// for i in 0..10 { /// if let Err(_) = tx.send(i).await { /// println!("receiver dropped"); /// return; /// } /// } /// }); /// /// while let Some(i) = rx.recv().await { /// println!("got = {}", i); /// } /// } /// ``` pub async fn send(&self, value: T) -> Result<(), SendError> { match self.reserve().await { Ok(permit) => { permit.send(value); Ok(()) } Err(_) => Err(SendError(value)), } } /// Completes when the receiver has dropped. /// /// This allows the producers to get notified when interest in the produced /// values is canceled and immediately stop doing work. /// /// # Cancel safety /// /// This method is cancel safe. Once the channel is closed, it stays closed /// forever and all future calls to `closed` will return immediately. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx1, rx) = mpsc::channel::<()>(1); /// let tx2 = tx1.clone(); /// let tx3 = tx1.clone(); /// let tx4 = tx1.clone(); /// let tx5 = tx1.clone(); /// tokio::spawn(async move { /// drop(rx); /// }); /// /// futures::join!( /// tx1.closed(), /// tx2.closed(), /// tx3.closed(), /// tx4.closed(), /// tx5.closed() /// ); /// println!("Receiver dropped"); /// } /// ``` pub async fn closed(&self) { self.chan.closed().await; } /// Attempts to immediately send a message on this `Sender` /// /// This method differs from [`send`] by returning immediately if the channel's /// buffer is full or no receiver is waiting to acquire some data. Compared /// with [`send`], this function has two failure cases instead of one (one for /// disconnection, one for a full buffer). /// /// # Errors /// /// If the channel capacity has been reached, i.e., the channel has `n` /// buffered values where `n` is the argument passed to [`channel`], then an /// error is returned. /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`Receiver`] handle dropping, the function returns /// an error. The error includes the value passed to `send`. /// /// [`send`]: Sender::send /// [`channel`]: channel /// [`close`]: Receiver::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// // Create a channel with buffer size 1 /// let (tx1, mut rx) = mpsc::channel(1); /// let tx2 = tx1.clone(); /// /// tokio::spawn(async move { /// tx1.send(1).await.unwrap(); /// tx1.send(2).await.unwrap(); /// // task waits until the receiver receives a value. /// }); /// /// tokio::spawn(async move { /// // This will return an error and send /// // no message if the buffer is full /// let _ = tx2.try_send(3); /// }); /// /// let mut msg; /// msg = rx.recv().await.unwrap(); /// println!("message {} received", msg); /// /// msg = rx.recv().await.unwrap(); /// println!("message {} received", msg); /// /// // Third message may have never been sent /// match rx.recv().await { /// Some(msg) => println!("message {} received", msg), /// None => println!("the third message was never sent"), /// } /// } /// ``` pub fn try_send(&self, message: T) -> Result<(), TrySendError> { match self.chan.semaphore().semaphore.try_acquire(1) { Ok(()) => {} Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(message)), Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(message)), } // Send the message self.chan.send(message); Ok(()) } /// Sends a value, waiting until there is capacity, but only for a limited time. /// /// Shares the same success and error conditions as [`send`], adding one more /// condition for an unsuccessful send, which is when the provided timeout has /// elapsed, and there is no capacity available. /// /// [`send`]: Sender::send /// /// # Errors /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`Receiver`] having been dropped, /// the function returns an error. The error includes the value passed to `send`. /// /// [`close`]: Receiver::close /// [`Receiver`]: Receiver /// /// # Panics /// /// This function panics if it is called outside the context of a Tokio /// runtime [with time enabled](crate::runtime::Builder::enable_time). /// /// # Examples /// /// In the following example, each call to `send_timeout` will block until the /// previously sent value was received, unless the timeout has elapsed. /// /// ```rust /// use tokio::sync::mpsc; /// use tokio::time::{sleep, Duration}; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// tokio::spawn(async move { /// for i in 0..10 { /// if let Err(e) = tx.send_timeout(i, Duration::from_millis(100)).await { /// println!("send error: #{:?}", e); /// return; /// } /// } /// }); /// /// while let Some(i) = rx.recv().await { /// println!("got = {}", i); /// sleep(Duration::from_millis(200)).await; /// } /// } /// ``` #[cfg(feature = "time")] #[cfg_attr(docsrs, doc(cfg(feature = "time")))] pub async fn send_timeout( &self, value: T, timeout: Duration, ) -> Result<(), SendTimeoutError> { let permit = match crate::time::timeout(timeout, self.reserve()).await { Err(_) => { return Err(SendTimeoutError::Timeout(value)); } Ok(Err(_)) => { return Err(SendTimeoutError::Closed(value)); } Ok(Ok(permit)) => permit, }; permit.send(value); Ok(()) } /// Blocking send to call outside of asynchronous contexts. /// /// This method is intended for use cases where you are sending from /// synchronous code to asynchronous code, and will work even if the /// receiver is not using [`blocking_recv`] to receive the message. /// /// [`blocking_recv`]: fn@crate::sync::mpsc::Receiver::blocking_recv /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// /// ``` /// use std::thread; /// use tokio::runtime::Runtime; /// use tokio::sync::mpsc; /// /// fn main() { /// let (tx, mut rx) = mpsc::channel::(1); /// /// let sync_code = thread::spawn(move || { /// tx.blocking_send(10).unwrap(); /// }); /// /// Runtime::new().unwrap().block_on(async move { /// assert_eq!(Some(10), rx.recv().await); /// }); /// sync_code.join().unwrap() /// } /// ``` #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "send_blocking"))] pub fn blocking_send(&self, value: T) -> Result<(), SendError> { crate::future::block_on(self.send(value)) } /// Checks if the channel has been closed. This happens when the /// [`Receiver`] is dropped, or when the [`Receiver::close`] method is /// called. /// /// [`Receiver`]: crate::sync::mpsc::Receiver /// [`Receiver::close`]: crate::sync::mpsc::Receiver::close /// /// ``` /// let (tx, rx) = tokio::sync::mpsc::channel::<()>(42); /// assert!(!tx.is_closed()); /// /// let tx2 = tx.clone(); /// assert!(!tx2.is_closed()); /// /// drop(rx); /// assert!(tx.is_closed()); /// assert!(tx2.is_closed()); /// ``` pub fn is_closed(&self) -> bool { self.chan.is_closed() } /// Waits for channel capacity. Once capacity to send one message is /// available, it is reserved for the caller. /// /// If the channel is full, the function waits for the number of unreceived /// messages to become less than the channel capacity. Capacity to send one /// message is reserved for the caller. A [`Permit`] is returned to track /// the reserved capacity. The [`send`] function on [`Permit`] consumes the /// reserved capacity. /// /// Dropping [`Permit`] without sending a message releases the capacity back /// to the channel. /// /// [`Permit`]: Permit /// [`send`]: Permit::send /// /// # Cancel safety /// /// This channel uses a queue to ensure that calls to `send` and `reserve` /// complete in the order they were requested. Cancelling a call to /// `reserve` makes you lose your place in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.reserve().await.unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Sending on the permit succeeds /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// } /// ``` pub async fn reserve(&self) -> Result, SendError<()>> { self.reserve_inner(1).await?; Ok(Permit { chan: &self.chan }) } /// Waits for channel capacity. Once capacity to send `n` messages is /// available, it is reserved for the caller. /// /// If the channel is full or if there are fewer than `n` permits available, the function waits /// for the number of unreceived messages to become `n` less than the channel capacity. /// Capacity to send `n` message is then reserved for the caller. /// /// A [`PermitIterator`] is returned to track the reserved capacity. /// You can call this [`Iterator`] until it is exhausted to /// get a [`Permit`] and then call [`Permit::send`]. This function is similar to /// [`try_reserve_many`] except it awaits for the slots to become available. /// /// If the channel is closed, the function returns a [`SendError`]. /// /// Dropping [`PermitIterator`] without consuming it entirely releases the remaining /// permits back to the channel. /// /// [`PermitIterator`]: PermitIterator /// [`Permit`]: Permit /// [`send`]: Permit::send /// [`try_reserve_many`]: Sender::try_reserve_many /// /// # Cancel safety /// /// This channel uses a queue to ensure that calls to `send` and `reserve_many` /// complete in the order they were requested. Cancelling a call to /// `reserve_many` makes you lose your place in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(2); /// /// // Reserve capacity /// let mut permit = tx.reserve_many(2).await.unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Sending with the permit iterator succeeds /// permit.next().unwrap().send(456); /// permit.next().unwrap().send(457); /// /// // The iterator should now be exhausted /// assert!(permit.next().is_none()); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// assert_eq!(rx.recv().await.unwrap(), 457); /// } /// ``` pub async fn reserve_many(&self, n: usize) -> Result, SendError<()>> { self.reserve_inner(n).await?; Ok(PermitIterator { chan: &self.chan, n, }) } /// Waits for channel capacity, moving the `Sender` and returning an owned /// permit. Once capacity to send one message is available, it is reserved /// for the caller. /// /// This moves the sender _by value_, and returns an owned permit that can /// be used to send a message into the channel. Unlike [`Sender::reserve`], /// this method may be used in cases where the permit must be valid for the /// `'static` lifetime. `Sender`s may be cloned cheaply (`Sender::clone` is /// essentially a reference count increment, comparable to [`Arc::clone`]), /// so when multiple [`OwnedPermit`]s are needed or the `Sender` cannot be /// moved, it can be cloned prior to calling `reserve_owned`. /// /// If the channel is full, the function waits for the number of unreceived /// messages to become less than the channel capacity. Capacity to send one /// message is reserved for the caller. An [`OwnedPermit`] is returned to /// track the reserved capacity. The [`send`] function on [`OwnedPermit`] /// consumes the reserved capacity. /// /// Dropping the [`OwnedPermit`] without sending a message releases the /// capacity back to the channel. /// /// # Cancel safety /// /// This channel uses a queue to ensure that calls to `send` and `reserve` /// complete in the order they were requested. Cancelling a call to /// `reserve_owned` makes you lose your place in the queue. /// /// # Examples /// Sending a message using an [`OwnedPermit`]: /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity, moving the sender. /// let permit = tx.reserve_owned().await.unwrap(); /// /// // Send a message, consuming the permit and returning /// // the moved sender. /// let tx = permit.send(123); /// /// // The value sent on the permit is received. /// assert_eq!(rx.recv().await.unwrap(), 123); /// /// // The sender can now be used again. /// tx.send(456).await.unwrap(); /// } /// ``` /// /// When multiple [`OwnedPermit`]s are needed, or the sender cannot be moved /// by value, it can be inexpensively cloned before calling `reserve_owned`: /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Clone the sender and reserve capacity. /// let permit = tx.clone().reserve_owned().await.unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Sending on the permit succeeds. /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// } /// ``` /// /// [`Sender::reserve`]: Sender::reserve /// [`OwnedPermit`]: OwnedPermit /// [`send`]: OwnedPermit::send /// [`Arc::clone`]: std::sync::Arc::clone pub async fn reserve_owned(self) -> Result, SendError<()>> { self.reserve_inner(1).await?; Ok(OwnedPermit { chan: Some(self.chan), }) } async fn reserve_inner(&self, n: usize) -> Result<(), SendError<()>> { crate::trace::async_trace_leaf().await; if n > self.max_capacity() { return Err(SendError(())); } match self.chan.semaphore().semaphore.acquire(n).await { Ok(()) => Ok(()), Err(_) => Err(SendError(())), } } /// Tries to acquire a slot in the channel without waiting for the slot to become /// available. /// /// If the channel is full this function will return [`TrySendError`], otherwise /// if there is a slot available it will return a [`Permit`] that will then allow you /// to [`send`] on the channel with a guaranteed slot. This function is similar to /// [`reserve`] except it does not await for the slot to become available. /// /// Dropping [`Permit`] without sending a message releases the capacity back /// to the channel. /// /// [`Permit`]: Permit /// [`send`]: Permit::send /// [`reserve`]: Sender::reserve /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.try_reserve().unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Trying to reserve an additional slot on the `tx` will /// // fail because there is no capacity. /// assert!(tx.try_reserve().is_err()); /// /// // Sending on the permit succeeds /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// /// } /// ``` pub fn try_reserve(&self) -> Result, TrySendError<()>> { match self.chan.semaphore().semaphore.try_acquire(1) { Ok(()) => {} Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(())), Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(())), } Ok(Permit { chan: &self.chan }) } /// Tries to acquire `n` slots in the channel without waiting for the slot to become /// available. /// /// A [`PermitIterator`] is returned to track the reserved capacity. /// You can call this [`Iterator`] until it is exhausted to /// get a [`Permit`] and then call [`Permit::send`]. This function is similar to /// [`reserve_many`] except it does not await for the slots to become available. /// /// If there are fewer than `n` permits available on the channel, then /// this function will return a [`TrySendError::Full`]. If the channel is closed /// this function will return a [`TrySendError::Closed`]. /// /// Dropping [`PermitIterator`] without consuming it entirely releases the remaining /// permits back to the channel. /// /// [`PermitIterator`]: PermitIterator /// [`send`]: Permit::send /// [`reserve_many`]: Sender::reserve_many /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(2); /// /// // Reserve capacity /// let mut permit = tx.try_reserve_many(2).unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Trying to reserve an additional slot on the `tx` will /// // fail because there is no capacity. /// assert!(tx.try_reserve().is_err()); /// /// // Sending with the permit iterator succeeds /// permit.next().unwrap().send(456); /// permit.next().unwrap().send(457); /// /// // The iterator should now be exhausted /// assert!(permit.next().is_none()); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// assert_eq!(rx.recv().await.unwrap(), 457); /// /// // Trying to call try_reserve_many with 0 will return an empty iterator /// let mut permit = tx.try_reserve_many(0).unwrap(); /// assert!(permit.next().is_none()); /// /// // Trying to call try_reserve_many with a number greater than the channel /// // capacity will return an error /// let permit = tx.try_reserve_many(3); /// assert!(permit.is_err()); /// /// // Trying to call try_reserve_many on a closed channel will return an error /// drop(rx); /// let permit = tx.try_reserve_many(1); /// assert!(permit.is_err()); /// /// let permit = tx.try_reserve_many(0); /// assert!(permit.is_err()); /// } /// ``` pub fn try_reserve_many(&self, n: usize) -> Result, TrySendError<()>> { if n > self.max_capacity() { return Err(TrySendError::Full(())); } match self.chan.semaphore().semaphore.try_acquire(n) { Ok(()) => {} Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(())), Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(())), } Ok(PermitIterator { chan: &self.chan, n, }) } /// Tries to acquire a slot in the channel without waiting for the slot to become /// available, returning an owned permit. /// /// This moves the sender _by value_, and returns an owned permit that can /// be used to send a message into the channel. Unlike [`Sender::try_reserve`], /// this method may be used in cases where the permit must be valid for the /// `'static` lifetime. `Sender`s may be cloned cheaply (`Sender::clone` is /// essentially a reference count increment, comparable to [`Arc::clone`]), /// so when multiple [`OwnedPermit`]s are needed or the `Sender` cannot be /// moved, it can be cloned prior to calling `try_reserve_owned`. /// /// If the channel is full this function will return a [`TrySendError`]. /// Since the sender is taken by value, the `TrySendError` returned in this /// case contains the sender, so that it may be used again. Otherwise, if /// there is a slot available, this method will return an [`OwnedPermit`] /// that can then be used to [`send`] on the channel with a guaranteed slot. /// This function is similar to [`reserve_owned`] except it does not await /// for the slot to become available. /// /// Dropping the [`OwnedPermit`] without sending a message releases the capacity back /// to the channel. /// /// [`OwnedPermit`]: OwnedPermit /// [`send`]: OwnedPermit::send /// [`reserve_owned`]: Sender::reserve_owned /// [`Arc::clone`]: std::sync::Arc::clone /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.clone().try_reserve_owned().unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Trying to reserve an additional slot on the `tx` will /// // fail because there is no capacity. /// assert!(tx.try_reserve().is_err()); /// /// // Sending on the permit succeeds /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// /// } /// ``` pub fn try_reserve_owned(self) -> Result, TrySendError> { match self.chan.semaphore().semaphore.try_acquire(1) { Ok(()) => {} Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(self)), Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(self)), } Ok(OwnedPermit { chan: Some(self.chan), }) } /// Returns `true` if senders belong to the same channel. /// /// # Examples /// /// ``` /// let (tx, rx) = tokio::sync::mpsc::channel::<()>(1); /// let tx2 = tx.clone(); /// assert!(tx.same_channel(&tx2)); /// /// let (tx3, rx3) = tokio::sync::mpsc::channel::<()>(1); /// assert!(!tx3.same_channel(&tx2)); /// ``` pub fn same_channel(&self, other: &Self) -> bool { self.chan.same_channel(&other.chan) } /// Returns the current capacity of the channel. /// /// The capacity goes down when sending a value by calling [`send`] or by reserving capacity /// with [`reserve`]. The capacity goes up when values are received by the [`Receiver`]. /// This is distinct from [`max_capacity`], which always returns buffer capacity initially /// specified when calling [`channel`] /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel::<()>(5); /// /// assert_eq!(tx.capacity(), 5); /// /// // Making a reservation drops the capacity by one. /// let permit = tx.reserve().await.unwrap(); /// assert_eq!(tx.capacity(), 4); /// /// // Sending and receiving a value increases the capacity by one. /// permit.send(()); /// rx.recv().await.unwrap(); /// assert_eq!(tx.capacity(), 5); /// } /// ``` /// /// [`send`]: Sender::send /// [`reserve`]: Sender::reserve /// [`channel`]: channel /// [`max_capacity`]: Sender::max_capacity pub fn capacity(&self) -> usize { self.chan.semaphore().semaphore.available_permits() } /// Converts the `Sender` to a [`WeakSender`] that does not count /// towards RAII semantics, i.e. if all `Sender` instances of the /// channel were dropped and only `WeakSender` instances remain, /// the channel is closed. #[must_use = "Downgrade creates a WeakSender without destroying the original non-weak sender."] pub fn downgrade(&self) -> WeakSender { WeakSender { chan: self.chan.downgrade(), } } /// Returns the maximum buffer capacity of the channel. /// /// The maximum capacity is the buffer capacity initially specified when calling /// [`channel`]. This is distinct from [`capacity`], which returns the *current* /// available buffer capacity: as messages are sent and received, the /// value returned by [`capacity`] will go up or down, whereas the value /// returned by [`max_capacity`] will remain constant. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = mpsc::channel::<()>(5); /// /// // both max capacity and capacity are the same at first /// assert_eq!(tx.max_capacity(), 5); /// assert_eq!(tx.capacity(), 5); /// /// // Making a reservation doesn't change the max capacity. /// let permit = tx.reserve().await.unwrap(); /// assert_eq!(tx.max_capacity(), 5); /// // but drops the capacity by one /// assert_eq!(tx.capacity(), 4); /// } /// ``` /// /// [`channel`]: channel /// [`max_capacity`]: Sender::max_capacity /// [`capacity`]: Sender::capacity pub fn max_capacity(&self) -> usize { self.chan.semaphore().bound } /// Returns the number of [`Sender`] handles. pub fn strong_count(&self) -> usize { self.chan.strong_count() } /// Returns the number of [`WeakSender`] handles. pub fn weak_count(&self) -> usize { self.chan.weak_count() } } impl Clone for Sender { fn clone(&self) -> Self { Sender { chan: self.chan.clone(), } } } impl fmt::Debug for Sender { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Sender") .field("chan", &self.chan) .finish() } } impl Clone for WeakSender { fn clone(&self) -> Self { self.chan.increment_weak_count(); WeakSender { chan: self.chan.clone(), } } } impl Drop for WeakSender { fn drop(&mut self) { self.chan.decrement_weak_count(); } } impl WeakSender { /// Tries to convert a `WeakSender` into a [`Sender`]. This will return `Some` /// if there are other `Sender` instances alive and the channel wasn't /// previously dropped, otherwise `None` is returned. pub fn upgrade(&self) -> Option> { chan::Tx::upgrade(self.chan.clone()).map(Sender::new) } /// Returns the number of [`Sender`] handles. pub fn strong_count(&self) -> usize { self.chan.strong_count() } /// Returns the number of [`WeakSender`] handles. pub fn weak_count(&self) -> usize { self.chan.weak_count() } } impl fmt::Debug for WeakSender { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("WeakSender").finish() } } // ===== impl Permit ===== impl Permit<'_, T> { /// Sends a value using the reserved capacity. /// /// Capacity for the message has already been reserved. The message is sent /// to the receiver and the permit is consumed. The operation will succeed /// even if the receiver half has been closed. See [`Receiver::close`] for /// more details on performing a clean shutdown. /// /// [`Receiver::close`]: Receiver::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.reserve().await.unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Send a message on the permit /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// } /// ``` pub fn send(self, value: T) { use std::mem; self.chan.send(value); // Avoid the drop logic mem::forget(self); } } impl Drop for Permit<'_, T> { fn drop(&mut self) { use chan::Semaphore; let semaphore = self.chan.semaphore(); // Add the permit back to the semaphore semaphore.add_permit(); // If this is the last sender for this channel, wake the receiver so // that it can be notified that the channel is closed. if semaphore.is_closed() && semaphore.is_idle() { self.chan.wake_rx(); } } } impl fmt::Debug for Permit<'_, T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Permit") .field("chan", &self.chan) .finish() } } // ===== impl PermitIterator ===== impl<'a, T> Iterator for PermitIterator<'a, T> { type Item = Permit<'a, T>; fn next(&mut self) -> Option { if self.n == 0 { return None; } self.n -= 1; Some(Permit { chan: self.chan }) } fn size_hint(&self) -> (usize, Option) { let n = self.n; (n, Some(n)) } } impl ExactSizeIterator for PermitIterator<'_, T> {} impl std::iter::FusedIterator for PermitIterator<'_, T> {} impl Drop for PermitIterator<'_, T> { fn drop(&mut self) { use chan::Semaphore; if self.n == 0 { return; } let semaphore = self.chan.semaphore(); // Add the remaining permits back to the semaphore semaphore.add_permits(self.n); // If this is the last sender for this channel, wake the receiver so // that it can be notified that the channel is closed. if semaphore.is_closed() && semaphore.is_idle() { self.chan.wake_rx(); } } } impl fmt::Debug for PermitIterator<'_, T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("PermitIterator") .field("chan", &self.chan) .field("capacity", &self.n) .finish() } } // ===== impl Permit ===== impl OwnedPermit { /// Sends a value using the reserved capacity. /// /// Capacity for the message has already been reserved. The message is sent /// to the receiver and the permit is consumed. The operation will succeed /// even if the receiver half has been closed. See [`Receiver::close`] for /// more details on performing a clean shutdown. /// /// Unlike [`Permit::send`], this method returns the [`Sender`] from which /// the `OwnedPermit` was reserved. /// /// [`Receiver::close`]: Receiver::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.reserve_owned().await.unwrap(); /// /// // Send a message on the permit, returning the sender. /// let tx = permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// /// // We may now reuse `tx` to send another message. /// tx.send(789).await.unwrap(); /// } /// ``` pub fn send(mut self, value: T) -> Sender { let chan = self.chan.take().unwrap_or_else(|| { unreachable!("OwnedPermit channel is only taken when the permit is moved") }); chan.send(value); Sender { chan } } /// Releases the reserved capacity *without* sending a message, returning the /// [`Sender`]. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::channel(1); /// /// // Clone the sender and reserve capacity /// let permit = tx.clone().reserve_owned().await.unwrap(); /// /// // Trying to send on the original `tx` will fail, since the `permit` /// // has reserved all the available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Release the permit without sending a message, returning the clone /// // of the sender. /// let tx2 = permit.release(); /// /// // We may now reuse `tx` to send another message. /// tx.send(789).await.unwrap(); /// # drop(rx); drop(tx2); /// } /// ``` /// /// [`Sender`]: Sender pub fn release(mut self) -> Sender { use chan::Semaphore; let chan = self.chan.take().unwrap_or_else(|| { unreachable!("OwnedPermit channel is only taken when the permit is moved") }); // Add the permit back to the semaphore chan.semaphore().add_permit(); Sender { chan } } } impl Drop for OwnedPermit { fn drop(&mut self) { use chan::Semaphore; // Are we still holding onto the sender? if let Some(chan) = self.chan.take() { let semaphore = chan.semaphore(); // Add the permit back to the semaphore semaphore.add_permit(); // If this `OwnedPermit` is holding the last sender for this // channel, wake the receiver so that it can be notified that the // channel is closed. if semaphore.is_closed() && semaphore.is_idle() { chan.wake_rx(); } } // Otherwise, do nothing. } } impl fmt::Debug for OwnedPermit { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("OwnedPermit") .field("chan", &self.chan) .finish() } } tokio-1.43.1/src/sync/mpsc/chan.rs000064400000000000000000000434041046102023000150300ustar 00000000000000use crate::loom::cell::UnsafeCell; use crate::loom::future::AtomicWaker; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Arc; use crate::runtime::park::CachedParkThread; use crate::sync::mpsc::error::TryRecvError; use crate::sync::mpsc::{bounded, list, unbounded}; use crate::sync::notify::Notify; use crate::util::cacheline::CachePadded; use std::fmt; use std::panic; use std::process; use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; use std::task::Poll::{Pending, Ready}; use std::task::{ready, Context, Poll}; /// Channel sender. pub(crate) struct Tx { inner: Arc>, } impl fmt::Debug for Tx { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Tx").field("inner", &self.inner).finish() } } /// Channel receiver. pub(crate) struct Rx { inner: Arc>, } impl fmt::Debug for Rx { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Rx").field("inner", &self.inner).finish() } } pub(crate) trait Semaphore { fn is_idle(&self) -> bool; fn add_permit(&self); fn add_permits(&self, n: usize); fn close(&self); fn is_closed(&self) -> bool; } pub(super) struct Chan { /// Handle to the push half of the lock-free list. tx: CachePadded>, /// Receiver waker. Notified when a value is pushed into the channel. rx_waker: CachePadded, /// Notifies all tasks listening for the receiver being dropped. notify_rx_closed: Notify, /// Coordinates access to channel's capacity. semaphore: S, /// Tracks the number of outstanding sender handles. /// /// When this drops to zero, the send half of the channel is closed. tx_count: AtomicUsize, /// Tracks the number of outstanding weak sender handles. tx_weak_count: AtomicUsize, /// Only accessed by `Rx` handle. rx_fields: UnsafeCell>, } impl fmt::Debug for Chan where S: fmt::Debug, { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Chan") .field("tx", &*self.tx) .field("semaphore", &self.semaphore) .field("rx_waker", &*self.rx_waker) .field("tx_count", &self.tx_count) .field("rx_fields", &"...") .finish() } } /// Fields only accessed by `Rx` handle. struct RxFields { /// Channel receiver. This field is only accessed by the `Receiver` type. list: list::Rx, /// `true` if `Rx::close` is called. rx_closed: bool, } impl fmt::Debug for RxFields { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("RxFields") .field("list", &self.list) .field("rx_closed", &self.rx_closed) .finish() } } unsafe impl Send for Chan {} unsafe impl Sync for Chan {} impl panic::RefUnwindSafe for Chan {} impl panic::UnwindSafe for Chan {} pub(crate) fn channel(semaphore: S) -> (Tx, Rx) { let (tx, rx) = list::channel(); let chan = Arc::new(Chan { notify_rx_closed: Notify::new(), tx: CachePadded::new(tx), semaphore, rx_waker: CachePadded::new(AtomicWaker::new()), tx_count: AtomicUsize::new(1), tx_weak_count: AtomicUsize::new(0), rx_fields: UnsafeCell::new(RxFields { list: rx, rx_closed: false, }), }); (Tx::new(chan.clone()), Rx::new(chan)) } // ===== impl Tx ===== impl Tx { fn new(chan: Arc>) -> Tx { Tx { inner: chan } } pub(super) fn strong_count(&self) -> usize { self.inner.tx_count.load(Acquire) } pub(super) fn weak_count(&self) -> usize { self.inner.tx_weak_count.load(Relaxed) } pub(super) fn downgrade(&self) -> Arc> { self.inner.increment_weak_count(); self.inner.clone() } // Returns the upgraded channel or None if the upgrade failed. pub(super) fn upgrade(chan: Arc>) -> Option { let mut tx_count = chan.tx_count.load(Acquire); loop { if tx_count == 0 { // channel is closed return None; } match chan .tx_count .compare_exchange_weak(tx_count, tx_count + 1, AcqRel, Acquire) { Ok(_) => return Some(Tx { inner: chan }), Err(prev_count) => tx_count = prev_count, } } } pub(super) fn semaphore(&self) -> &S { &self.inner.semaphore } /// Send a message and notify the receiver. pub(crate) fn send(&self, value: T) { self.inner.send(value); } /// Wake the receive half pub(crate) fn wake_rx(&self) { self.inner.rx_waker.wake(); } /// Returns `true` if senders belong to the same channel. pub(crate) fn same_channel(&self, other: &Self) -> bool { Arc::ptr_eq(&self.inner, &other.inner) } } impl Tx { pub(crate) fn is_closed(&self) -> bool { self.inner.semaphore.is_closed() } pub(crate) async fn closed(&self) { // In order to avoid a race condition, we first request a notification, // **then** check whether the semaphore is closed. If the semaphore is // closed the notification request is dropped. let notified = self.inner.notify_rx_closed.notified(); if self.inner.semaphore.is_closed() { return; } notified.await; } } impl Clone for Tx { fn clone(&self) -> Tx { // Using a Relaxed ordering here is sufficient as the caller holds a // strong ref to `self`, preventing a concurrent decrement to zero. self.inner.tx_count.fetch_add(1, Relaxed); Tx { inner: self.inner.clone(), } } } impl Drop for Tx { fn drop(&mut self) { if self.inner.tx_count.fetch_sub(1, AcqRel) != 1 { return; } // Close the list, which sends a `Close` message self.inner.tx.close(); // Notify the receiver self.wake_rx(); } } // ===== impl Rx ===== impl Rx { fn new(chan: Arc>) -> Rx { Rx { inner: chan } } pub(crate) fn close(&mut self) { self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; if rx_fields.rx_closed { return; } rx_fields.rx_closed = true; }); self.inner.semaphore.close(); self.inner.notify_rx_closed.notify_waiters(); } pub(crate) fn is_closed(&self) -> bool { // There two internal states that can represent a closed channel // // 1. When `close` is called. // In this case, the inner semaphore will be closed. // // 2. When all senders are dropped. // In this case, the semaphore remains unclosed, and the `index` in the list won't // reach the tail position. It is necessary to check the list if the last block is // `closed`. self.inner.semaphore.is_closed() || self.inner.tx_count.load(Acquire) == 0 } pub(crate) fn is_empty(&self) -> bool { self.inner.rx_fields.with(|rx_fields_ptr| { let rx_fields = unsafe { &*rx_fields_ptr }; rx_fields.list.is_empty(&self.inner.tx) }) } pub(crate) fn len(&self) -> usize { self.inner.rx_fields.with(|rx_fields_ptr| { let rx_fields = unsafe { &*rx_fields_ptr }; rx_fields.list.len(&self.inner.tx) }) } /// Receive the next value pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll> { use super::block::Read; ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; macro_rules! try_recv { () => { match rx_fields.list.pop(&self.inner.tx) { Some(Read::Value(value)) => { self.inner.semaphore.add_permit(); coop.made_progress(); return Ready(Some(value)); } Some(Read::Closed) => { // TODO: This check may not be required as it most // likely can only return `true` at this point. A // channel is closed when all tx handles are // dropped. Dropping a tx handle releases memory, // which ensures that if dropping the tx handle is // visible, then all messages sent are also visible. assert!(self.inner.semaphore.is_idle()); coop.made_progress(); return Ready(None); } None => {} // fall through } }; } try_recv!(); self.inner.rx_waker.register_by_ref(cx.waker()); // It is possible that a value was pushed between attempting to read // and registering the task, so we have to check the channel a // second time here. try_recv!(); if rx_fields.rx_closed && self.inner.semaphore.is_idle() { coop.made_progress(); Ready(None) } else { Pending } }) } /// Receives up to `limit` values into `buffer` /// /// For `limit > 0`, receives up to limit values into `buffer`. /// For `limit == 0`, immediately returns Ready(0). pub(crate) fn recv_many( &mut self, cx: &mut Context<'_>, buffer: &mut Vec, limit: usize, ) -> Poll { use super::block::Read; ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); if limit == 0 { coop.made_progress(); return Ready(0usize); } let mut remaining = limit; let initial_length = buffer.len(); self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; macro_rules! try_recv { () => { while remaining > 0 { match rx_fields.list.pop(&self.inner.tx) { Some(Read::Value(value)) => { remaining -= 1; buffer.push(value); } Some(Read::Closed) => { let number_added = buffer.len() - initial_length; if number_added > 0 { self.inner.semaphore.add_permits(number_added); } // TODO: This check may not be required as it most // likely can only return `true` at this point. A // channel is closed when all tx handles are // dropped. Dropping a tx handle releases memory, // which ensures that if dropping the tx handle is // visible, then all messages sent are also visible. assert!(self.inner.semaphore.is_idle()); coop.made_progress(); return Ready(number_added); } None => { break; // fall through } } } let number_added = buffer.len() - initial_length; if number_added > 0 { self.inner.semaphore.add_permits(number_added); coop.made_progress(); return Ready(number_added); } }; } try_recv!(); self.inner.rx_waker.register_by_ref(cx.waker()); // It is possible that a value was pushed between attempting to read // and registering the task, so we have to check the channel a // second time here. try_recv!(); if rx_fields.rx_closed && self.inner.semaphore.is_idle() { assert!(buffer.is_empty()); coop.made_progress(); Ready(0usize) } else { Pending } }) } /// Try to receive the next value. pub(crate) fn try_recv(&mut self) -> Result { use super::list::TryPopResult; self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; macro_rules! try_recv { () => { match rx_fields.list.try_pop(&self.inner.tx) { TryPopResult::Ok(value) => { self.inner.semaphore.add_permit(); return Ok(value); } TryPopResult::Closed => return Err(TryRecvError::Disconnected), TryPopResult::Empty => return Err(TryRecvError::Empty), TryPopResult::Busy => {} // fall through } }; } try_recv!(); // If a previous `poll_recv` call has set a waker, we wake it here. // This allows us to put our own CachedParkThread waker in the // AtomicWaker slot instead. // // This is not a spurious wakeup to `poll_recv` since we just got a // Busy from `try_pop`, which only happens if there are messages in // the queue. self.inner.rx_waker.wake(); // Park the thread until the problematic send has completed. let mut park = CachedParkThread::new(); let waker = park.waker().unwrap(); loop { self.inner.rx_waker.register_by_ref(&waker); // It is possible that the problematic send has now completed, // so we have to check for messages again. try_recv!(); park.park(); } }) } pub(super) fn semaphore(&self) -> &S { &self.inner.semaphore } pub(super) fn sender_strong_count(&self) -> usize { self.inner.tx_count.load(Acquire) } pub(super) fn sender_weak_count(&self) -> usize { self.inner.tx_weak_count.load(Relaxed) } } impl Drop for Rx { fn drop(&mut self) { use super::block::Read::Value; self.close(); self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; while let Some(Value(_)) = rx_fields.list.pop(&self.inner.tx) { self.inner.semaphore.add_permit(); } }); } } // ===== impl Chan ===== impl Chan { fn send(&self, value: T) { // Push the value self.tx.push(value); // Notify the rx task self.rx_waker.wake(); } pub(super) fn decrement_weak_count(&self) { self.tx_weak_count.fetch_sub(1, Relaxed); } pub(super) fn increment_weak_count(&self) { self.tx_weak_count.fetch_add(1, Relaxed); } pub(super) fn strong_count(&self) -> usize { self.tx_count.load(Acquire) } pub(super) fn weak_count(&self) -> usize { self.tx_weak_count.load(Relaxed) } } impl Drop for Chan { fn drop(&mut self) { use super::block::Read::Value; // Safety: the only owner of the rx fields is Chan, and being // inside its own Drop means we're the last ones to touch it. self.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; while let Some(Value(_)) = rx_fields.list.pop(&self.tx) {} unsafe { rx_fields.list.free_blocks() }; }); } } // ===== impl Semaphore for (::Semaphore, capacity) ===== impl Semaphore for bounded::Semaphore { fn add_permit(&self) { self.semaphore.release(1); } fn add_permits(&self, n: usize) { self.semaphore.release(n) } fn is_idle(&self) -> bool { self.semaphore.available_permits() == self.bound } fn close(&self) { self.semaphore.close(); } fn is_closed(&self) -> bool { self.semaphore.is_closed() } } // ===== impl Semaphore for AtomicUsize ===== impl Semaphore for unbounded::Semaphore { fn add_permit(&self) { let prev = self.0.fetch_sub(2, Release); if prev >> 1 == 0 { // Something went wrong process::abort(); } } fn add_permits(&self, n: usize) { let prev = self.0.fetch_sub(n << 1, Release); if (prev >> 1) < n { // Something went wrong process::abort(); } } fn is_idle(&self) -> bool { self.0.load(Acquire) >> 1 == 0 } fn close(&self) { self.0.fetch_or(1, Release); } fn is_closed(&self) -> bool { self.0.load(Acquire) & 1 == 1 } } tokio-1.43.1/src/sync/mpsc/error.rs000064400000000000000000000113071046102023000152450ustar 00000000000000//! Channel error types. use std::error::Error; use std::fmt; /// Error returned by the `Sender`. #[derive(PartialEq, Eq, Clone, Copy)] pub struct SendError(pub T); impl fmt::Debug for SendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("SendError").finish_non_exhaustive() } } impl fmt::Display for SendError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } impl Error for SendError {} // ===== TrySendError ===== /// This enumeration is the list of the possible error outcomes for the /// [`try_send`](super::Sender::try_send) method. #[derive(PartialEq, Eq, Clone, Copy)] pub enum TrySendError { /// The data could not be sent on the channel because the channel is /// currently full and sending would require blocking. Full(T), /// The receive half of the channel was explicitly closed or has been /// dropped. Closed(T), } impl TrySendError { /// Consume the `TrySendError`, returning the unsent value. pub fn into_inner(self) -> T { match self { TrySendError::Full(val) => val, TrySendError::Closed(val) => val, } } } impl fmt::Debug for TrySendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match *self { TrySendError::Full(..) => "Full(..)".fmt(f), TrySendError::Closed(..) => "Closed(..)".fmt(f), } } } impl fmt::Display for TrySendError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!( fmt, "{}", match self { TrySendError::Full(..) => "no available capacity", TrySendError::Closed(..) => "channel closed", } ) } } impl Error for TrySendError {} impl From> for TrySendError { fn from(src: SendError) -> TrySendError { TrySendError::Closed(src.0) } } // ===== TryRecvError ===== /// Error returned by `try_recv`. #[derive(PartialEq, Eq, Clone, Copy, Debug)] pub enum TryRecvError { /// This **channel** is currently empty, but the **Sender**(s) have not yet /// disconnected, so data may yet become available. Empty, /// The **channel**'s sending half has become disconnected, and there will /// never be any more data received on it. Disconnected, } impl fmt::Display for TryRecvError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match *self { TryRecvError::Empty => "receiving on an empty channel".fmt(fmt), TryRecvError::Disconnected => "receiving on a closed channel".fmt(fmt), } } } impl Error for TryRecvError {} // ===== RecvError ===== /// Error returned by `Receiver`. #[derive(Debug, Clone)] #[doc(hidden)] #[deprecated(note = "This type is unused because recv returns an Option.")] pub struct RecvError(()); #[allow(deprecated)] impl fmt::Display for RecvError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } #[allow(deprecated)] impl Error for RecvError {} cfg_time! { // ===== SendTimeoutError ===== #[derive(PartialEq, Eq, Clone, Copy)] /// Error returned by [`Sender::send_timeout`](super::Sender::send_timeout)]. pub enum SendTimeoutError { /// The data could not be sent on the channel because the channel is /// full, and the timeout to send has elapsed. Timeout(T), /// The receive half of the channel was explicitly closed or has been /// dropped. Closed(T), } impl SendTimeoutError { /// Consume the `SendTimeoutError`, returning the unsent value. pub fn into_inner(self) -> T { match self { SendTimeoutError::Timeout(val) => val, SendTimeoutError::Closed(val) => val, } } } impl fmt::Debug for SendTimeoutError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match *self { SendTimeoutError::Timeout(..) => "Timeout(..)".fmt(f), SendTimeoutError::Closed(..) => "Closed(..)".fmt(f), } } } impl fmt::Display for SendTimeoutError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!( fmt, "{}", match self { SendTimeoutError::Timeout(..) => "timed out waiting on send operation", SendTimeoutError::Closed(..) => "channel closed", } ) } } impl Error for SendTimeoutError {} } tokio-1.43.1/src/sync/mpsc/list.rs000064400000000000000000000314561046102023000150760ustar 00000000000000//! A concurrent, lock-free, FIFO list. use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize}; use crate::loom::thread; use crate::sync::mpsc::block::{self, Block}; use std::fmt; use std::ptr::NonNull; use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; /// List queue transmit handle. pub(crate) struct Tx { /// Tail in the `Block` mpmc list. block_tail: AtomicPtr>, /// Position to push the next message. This references a block and offset /// into the block. tail_position: AtomicUsize, } /// List queue receive handle pub(crate) struct Rx { /// Pointer to the block being processed. head: NonNull>, /// Next slot index to process. index: usize, /// Pointer to the next block pending release. free_head: NonNull>, } /// Return value of `Rx::try_pop`. pub(crate) enum TryPopResult { /// Successfully popped a value. Ok(T), /// The channel is empty. Empty, /// The channel is empty and closed. Closed, /// The channel is not empty, but the first value is being written. Busy, } pub(crate) fn channel() -> (Tx, Rx) { // Create the initial block shared between the tx and rx halves. let initial_block = Block::new(0); let initial_block_ptr = Box::into_raw(initial_block); let tx = Tx { block_tail: AtomicPtr::new(initial_block_ptr), tail_position: AtomicUsize::new(0), }; let head = NonNull::new(initial_block_ptr).unwrap(); let rx = Rx { head, index: 0, free_head: head, }; (tx, rx) } impl Tx { /// Pushes a value into the list. pub(crate) fn push(&self, value: T) { // First, claim a slot for the value. `Acquire` is used here to // synchronize with the `fetch_add` in `reclaim_blocks`. let slot_index = self.tail_position.fetch_add(1, Acquire); // Load the current block and write the value let block = self.find_block(slot_index); unsafe { // Write the value to the block block.as_ref().write(slot_index, value); } } /// Closes the send half of the list. /// /// Similar process as pushing a value, but instead of writing the value & /// setting the ready flag, the `TX_CLOSED` flag is set on the block. pub(crate) fn close(&self) { // First, claim a slot for the value. This is the last slot that will be // claimed. let slot_index = self.tail_position.fetch_add(1, Acquire); let block = self.find_block(slot_index); unsafe { block.as_ref().tx_close() } } fn find_block(&self, slot_index: usize) -> NonNull> { // The start index of the block that contains `index`. let start_index = block::start_index(slot_index); // The index offset into the block let offset = block::offset(slot_index); // Load the current head of the block let mut block_ptr = self.block_tail.load(Acquire); let block = unsafe { &*block_ptr }; // Calculate the distance between the tail ptr and the target block let distance = block.distance(start_index); // Decide if this call to `find_block` should attempt to update the // `block_tail` pointer. // // Updating `block_tail` is not always performed in order to reduce // contention. // // When set, as the routine walks the linked list, it attempts to update // `block_tail`. If the update cannot be performed, `try_updating_tail` // is unset. let mut try_updating_tail = distance > offset; // Walk the linked list of blocks until the block with `start_index` is // found. loop { let block = unsafe { &(*block_ptr) }; if block.is_at_index(start_index) { return unsafe { NonNull::new_unchecked(block_ptr) }; } let next_block = block .load_next(Acquire) // There is no allocated next block, grow the linked list. .unwrap_or_else(|| block.grow()); // If the block is **not** final, then the tail pointer cannot be // advanced any more. try_updating_tail &= block.is_final(); if try_updating_tail { // Advancing `block_tail` must happen when walking the linked // list. `block_tail` may not advance passed any blocks that are // not "final". At the point a block is finalized, it is unknown // if there are any prior blocks that are unfinalized, which // makes it impossible to advance `block_tail`. // // While walking the linked list, `block_tail` can be advanced // as long as finalized blocks are traversed. // // Release ordering is used to ensure that any subsequent reads // are able to see the memory pointed to by `block_tail`. // // Acquire is not needed as any "actual" value is not accessed. // At this point, the linked list is walked to acquire blocks. if self .block_tail .compare_exchange(block_ptr, next_block.as_ptr(), Release, Relaxed) .is_ok() { // Synchronize with any senders let tail_position = self.tail_position.fetch_add(0, Release); unsafe { block.tx_release(tail_position); } } else { // A concurrent sender is also working on advancing // `block_tail` and this thread is falling behind. // // Stop trying to advance the tail pointer try_updating_tail = false; } } block_ptr = next_block.as_ptr(); thread::yield_now(); } } pub(crate) unsafe fn reclaim_block(&self, mut block: NonNull>) { // The block has been removed from the linked list and ownership // is reclaimed. // // Before dropping the block, see if it can be reused by // inserting it back at the end of the linked list. // // First, reset the data block.as_mut().reclaim(); let mut reused = false; // Attempt to insert the block at the end // // Walk at most three times // let curr_ptr = self.block_tail.load(Acquire); // The pointer can never be null debug_assert!(!curr_ptr.is_null()); let mut curr = NonNull::new_unchecked(curr_ptr); // TODO: Unify this logic with Block::grow for _ in 0..3 { match curr.as_ref().try_push(&mut block, AcqRel, Acquire) { Ok(()) => { reused = true; break; } Err(next) => { curr = next; } } } if !reused { let _ = Box::from_raw(block.as_ptr()); } } pub(crate) fn is_closed(&self) -> bool { let tail = self.block_tail.load(Acquire); unsafe { let tail_block = &*tail; tail_block.is_closed() } } } impl fmt::Debug for Tx { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Tx") .field("block_tail", &self.block_tail.load(Relaxed)) .field("tail_position", &self.tail_position.load(Relaxed)) .finish() } } impl Rx { pub(crate) fn is_empty(&self, tx: &Tx) -> bool { let block = unsafe { self.head.as_ref() }; if block.has_value(self.index) { return false; } // It is possible that a block has no value "now" but the list is still not empty. // To be sure, it is necessary to check the length of the list. self.len(tx) == 0 } pub(crate) fn len(&self, tx: &Tx) -> usize { // When all the senders are dropped, there will be a last block in the tail position, // but it will be closed let tail_position = tx.tail_position.load(Acquire); tail_position - self.index - (tx.is_closed() as usize) } /// Pops the next value off the queue. pub(crate) fn pop(&mut self, tx: &Tx) -> Option> { // Advance `head`, if needed if !self.try_advancing_head() { return None; } self.reclaim_blocks(tx); unsafe { let block = self.head.as_ref(); let ret = block.read(self.index); if let Some(block::Read::Value(..)) = ret { self.index = self.index.wrapping_add(1); } ret } } /// Pops the next value off the queue, detecting whether the block /// is busy or empty on failure. /// /// This function exists because `Rx::pop` can return `None` even if the /// channel's queue contains a message that has been completely written. /// This can happen if the fully delivered message is behind another message /// that is in the middle of being written to the block, since the channel /// can't return the messages out of order. pub(crate) fn try_pop(&mut self, tx: &Tx) -> TryPopResult { let tail_position = tx.tail_position.load(Acquire); let result = self.pop(tx); match result { Some(block::Read::Value(t)) => TryPopResult::Ok(t), Some(block::Read::Closed) => TryPopResult::Closed, None if tail_position == self.index => TryPopResult::Empty, None => TryPopResult::Busy, } } /// Tries advancing the block pointer to the block referenced by `self.index`. /// /// Returns `true` if successful, `false` if there is no next block to load. fn try_advancing_head(&mut self) -> bool { let block_index = block::start_index(self.index); loop { let next_block = { let block = unsafe { self.head.as_ref() }; if block.is_at_index(block_index) { return true; } block.load_next(Acquire) }; let next_block = match next_block { Some(next_block) => next_block, None => { return false; } }; self.head = next_block; thread::yield_now(); } } fn reclaim_blocks(&mut self, tx: &Tx) { while self.free_head != self.head { unsafe { // Get a handle to the block that will be freed and update // `free_head` to point to the next block. let block = self.free_head; let observed_tail_position = block.as_ref().observed_tail_position(); let required_index = match observed_tail_position { Some(i) => i, None => return, }; if required_index > self.index { return; } // We may read the next pointer with `Relaxed` ordering as it is // guaranteed that the `reclaim_blocks` routine trails the `recv` // routine. Any memory accessed by `reclaim_blocks` has already // been acquired by `recv`. let next_block = block.as_ref().load_next(Relaxed); // Update the free list head self.free_head = next_block.unwrap(); // Push the emptied block onto the back of the queue, making it // available to senders. tx.reclaim_block(block); } thread::yield_now(); } } /// Effectively `Drop` all the blocks. Should only be called once, when /// the list is dropping. pub(super) unsafe fn free_blocks(&mut self) { debug_assert_ne!(self.free_head, NonNull::dangling()); let mut cur = Some(self.free_head); #[cfg(debug_assertions)] { // to trigger the debug assert above so as to catch that we // don't call `free_blocks` more than once. self.free_head = NonNull::dangling(); self.head = NonNull::dangling(); } while let Some(block) = cur { cur = block.as_ref().load_next(Relaxed); drop(Box::from_raw(block.as_ptr())); } } } impl fmt::Debug for Rx { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Rx") .field("head", &self.head) .field("index", &self.index) .field("free_head", &self.free_head) .finish() } } tokio-1.43.1/src/sync/mpsc/mod.rs000064400000000000000000000150751046102023000147010ustar 00000000000000#![cfg_attr(not(feature = "sync"), allow(dead_code, unreachable_pub))] //! A multi-producer, single-consumer queue for sending values between //! asynchronous tasks. //! //! This module provides two variants of the channel: bounded and unbounded. The //! bounded variant has a limit on the number of messages that the channel can //! store, and if this limit is reached, trying to send another message will //! wait until a message is received from the channel. An unbounded channel has //! an infinite capacity, so the `send` method will always complete immediately. //! This makes the [`UnboundedSender`] usable from both synchronous and //! asynchronous code. //! //! Similar to the `mpsc` channels provided by `std`, the channel constructor //! functions provide separate send and receive handles, [`Sender`] and //! [`Receiver`] for the bounded channel, [`UnboundedSender`] and //! [`UnboundedReceiver`] for the unbounded channel. If there is no message to read, //! the current task will be notified when a new value is sent. [`Sender`] and //! [`UnboundedSender`] allow sending values into the channel. If the bounded //! channel is at capacity, the send is rejected and the task will be notified //! when additional capacity is available. In other words, the channel provides //! backpressure. //! //! This channel is also suitable for the single-producer single-consumer //! use-case. (Unless you only need to send one message, in which case you //! should use the [oneshot] channel.) //! //! # Disconnection //! //! When all [`Sender`] handles have been dropped, it is no longer //! possible to send values into the channel. This is considered the termination //! event of the stream. As such, `Receiver::poll` returns `Ok(Ready(None))`. //! //! If the [`Receiver`] handle is dropped, then messages can no longer //! be read out of the channel. In this case, all further attempts to send will //! result in an error. Additionally, all unread messages will be drained from the //! channel and dropped. //! //! # Clean Shutdown //! //! When the [`Receiver`] is dropped, it is possible for unprocessed messages to //! remain in the channel. Instead, it is usually desirable to perform a "clean" //! shutdown. To do this, the receiver first calls `close`, which will prevent //! any further messages to be sent into the channel. Then, the receiver //! consumes the channel to completion, at which point the receiver can be //! dropped. //! //! # Communicating between sync and async code //! //! When you want to communicate between synchronous and asynchronous code, there //! are two situations to consider: //! //! **Bounded channel**: If you need a bounded channel, you should use a bounded //! Tokio `mpsc` channel for both directions of communication. Instead of calling //! the async [`send`][bounded-send] or [`recv`][bounded-recv] methods, in //! synchronous code you will need to use the [`blocking_send`][blocking-send] or //! [`blocking_recv`][blocking-recv] methods. //! //! **Unbounded channel**: You should use the kind of channel that matches where //! the receiver is. So for sending a message _from async to sync_, you should //! use [the standard library unbounded channel][std-unbounded] or //! [crossbeam][crossbeam-unbounded]. Similarly, for sending a message _from sync //! to async_, you should use an unbounded Tokio `mpsc` channel. //! //! Please be aware that the above remarks were written with the `mpsc` channel //! in mind, but they can also be generalized to other kinds of channels. In //! general, any channel method that isn't marked async can be called anywhere, //! including outside of the runtime. For example, sending a message on a //! [oneshot] channel from outside the runtime is perfectly fine. //! //! # Multiple runtimes //! //! The `mpsc` channel is runtime agnostic. You can freely move it between //! different instances of the Tokio runtime or even use it from non-Tokio //! runtimes. //! //! When used in a Tokio runtime, it participates in //! [cooperative scheduling](crate::task#cooperative-scheduling) to avoid //! starvation. This feature does not apply when used from non-Tokio runtimes. //! //! As an exception, methods ending in `_timeout` are not runtime agnostic //! because they require access to the Tokio timer. See the documentation of //! each `*_timeout` method for more information on its use. //! //! # Allocation behavior //! //!
The implementation details described in this section may change in future //! Tokio releases.
//! //! The mpsc channel stores elements in blocks. Blocks are organized in a linked list. Sending //! pushes new elements onto the block at the front of the list, and receiving pops them off the //! one at the back. A block can hold 32 messages on a 64-bit target and 16 messages on a 32-bit //! target. This number is independent of channel and message size. Each block also stores 4 //! pointer-sized values for bookkeeping (so on a 64-bit machine, each message has 1 byte of //! overhead). //! //! When all values in a block have been received, it becomes empty. It will then be freed, unless //! the channel's first block (where newly-sent elements are being stored) has no next block. In //! that case, the empty block is reused as the next block. //! //! [`Sender`]: crate::sync::mpsc::Sender //! [`Receiver`]: crate::sync::mpsc::Receiver //! [bounded-send]: crate::sync::mpsc::Sender::send() //! [bounded-recv]: crate::sync::mpsc::Receiver::recv() //! [blocking-send]: crate::sync::mpsc::Sender::blocking_send() //! [blocking-recv]: crate::sync::mpsc::Receiver::blocking_recv() //! [`UnboundedSender`]: crate::sync::mpsc::UnboundedSender //! [`UnboundedReceiver`]: crate::sync::mpsc::UnboundedReceiver //! [oneshot]: crate::sync::oneshot //! [`Handle::block_on`]: crate::runtime::Handle::block_on() //! [std-unbounded]: std::sync::mpsc::channel //! [crossbeam-unbounded]: https://docs.rs/crossbeam/*/crossbeam/channel/fn.unbounded.html //! [`send_timeout`]: crate::sync::mpsc::Sender::send_timeout pub(super) mod block; mod bounded; pub use self::bounded::{ channel, OwnedPermit, Permit, PermitIterator, Receiver, Sender, WeakSender, }; mod chan; pub(super) mod list; mod unbounded; pub use self::unbounded::{ unbounded_channel, UnboundedReceiver, UnboundedSender, WeakUnboundedSender, }; pub mod error; /// The number of values a block can contain. /// /// This value must be a power of 2. It also must be smaller than the number of /// bits in `usize`. #[cfg(all(target_pointer_width = "64", not(loom)))] const BLOCK_CAP: usize = 32; #[cfg(all(not(target_pointer_width = "64"), not(loom)))] const BLOCK_CAP: usize = 16; #[cfg(loom)] const BLOCK_CAP: usize = 2; tokio-1.43.1/src/sync/mpsc/unbounded.rs000064400000000000000000000574751046102023000161170ustar 00000000000000use crate::loom::sync::{atomic::AtomicUsize, Arc}; use crate::sync::mpsc::chan; use crate::sync::mpsc::error::{SendError, TryRecvError}; use std::fmt; use std::task::{Context, Poll}; /// Send values to the associated `UnboundedReceiver`. /// /// Instances are created by the [`unbounded_channel`] function. pub struct UnboundedSender { chan: chan::Tx, } /// An unbounded sender that does not prevent the channel from being closed. /// /// If all [`UnboundedSender`] instances of a channel were dropped and only /// `WeakUnboundedSender` instances remain, the channel is closed. /// /// In order to send messages, the `WeakUnboundedSender` needs to be upgraded using /// [`WeakUnboundedSender::upgrade`], which returns `Option`. It returns `None` /// if all `UnboundedSender`s have been dropped, and otherwise it returns an `UnboundedSender`. /// /// [`UnboundedSender`]: UnboundedSender /// [`WeakUnboundedSender::upgrade`]: WeakUnboundedSender::upgrade /// /// # Examples /// /// ``` /// use tokio::sync::mpsc::unbounded_channel; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = unbounded_channel::(); /// let tx_weak = tx.downgrade(); /// /// // Upgrading will succeed because `tx` still exists. /// assert!(tx_weak.upgrade().is_some()); /// /// // If we drop `tx`, then it will fail. /// drop(tx); /// assert!(tx_weak.clone().upgrade().is_none()); /// } /// ``` pub struct WeakUnboundedSender { chan: Arc>, } impl Clone for UnboundedSender { fn clone(&self) -> Self { UnboundedSender { chan: self.chan.clone(), } } } impl fmt::Debug for UnboundedSender { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("UnboundedSender") .field("chan", &self.chan) .finish() } } /// Receive values from the associated `UnboundedSender`. /// /// Instances are created by the [`unbounded_channel`] function. /// /// This receiver can be turned into a `Stream` using [`UnboundedReceiverStream`]. /// /// [`UnboundedReceiverStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.UnboundedReceiverStream.html pub struct UnboundedReceiver { /// The channel receiver chan: chan::Rx, } impl fmt::Debug for UnboundedReceiver { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("UnboundedReceiver") .field("chan", &self.chan) .finish() } } /// Creates an unbounded mpsc channel for communicating between asynchronous /// tasks without backpressure. /// /// A `send` on this channel will always succeed as long as the receive half has /// not been closed. If the receiver falls behind, messages will be arbitrarily /// buffered. /// /// **Note** that the amount of available system memory is an implicit bound to /// the channel. Using an `unbounded` channel has the ability of causing the /// process to run out of memory. In this case, the process will be aborted. pub fn unbounded_channel() -> (UnboundedSender, UnboundedReceiver) { let (tx, rx) = chan::channel(Semaphore(AtomicUsize::new(0))); let tx = UnboundedSender::new(tx); let rx = UnboundedReceiver::new(rx); (tx, rx) } /// No capacity #[derive(Debug)] pub(crate) struct Semaphore(pub(crate) AtomicUsize); impl UnboundedReceiver { pub(crate) fn new(chan: chan::Rx) -> UnboundedReceiver { UnboundedReceiver { chan } } /// Receives the next value for this receiver. /// /// This method returns `None` if the channel has been closed and there are /// no remaining messages in the channel's buffer. This indicates that no /// further values can ever be received from this `Receiver`. The channel is /// closed when all senders have been dropped, or when [`close`] is called. /// /// If there are no messages in the channel's buffer, but the channel has /// not yet been closed, this method will sleep until a message is sent or /// the channel is closed. /// /// # Cancel safety /// /// This method is cancel safe. If `recv` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// channel. /// /// [`close`]: Self::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::unbounded_channel(); /// /// tokio::spawn(async move { /// tx.send("hello").unwrap(); /// }); /// /// assert_eq!(Some("hello"), rx.recv().await); /// assert_eq!(None, rx.recv().await); /// } /// ``` /// /// Values are buffered: /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::unbounded_channel(); /// /// tx.send("hello").unwrap(); /// tx.send("world").unwrap(); /// /// assert_eq!(Some("hello"), rx.recv().await); /// assert_eq!(Some("world"), rx.recv().await); /// } /// ``` pub async fn recv(&mut self) -> Option { use std::future::poll_fn; poll_fn(|cx| self.poll_recv(cx)).await } /// Receives the next values for this receiver and extends `buffer`. /// /// This method extends `buffer` by no more than a fixed number of values /// as specified by `limit`. If `limit` is zero, the function returns /// immediately with `0`. The return value is the number of values added to /// `buffer`. /// /// For `limit > 0`, if there are no messages in the channel's queue, /// but the channel has not yet been closed, this method will sleep /// until a message is sent or the channel is closed. /// /// For non-zero values of `limit`, this method will never return `0` unless /// the channel has been closed and there are no remaining messages in the /// channel's queue. This indicates that no further values can ever be /// received from this `Receiver`. The channel is closed when all senders /// have been dropped, or when [`close`] is called. /// /// The capacity of `buffer` is increased as needed. /// /// # Cancel safety /// /// This method is cancel safe. If `recv_many` is used as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, it is guaranteed that no messages were received on this /// channel. /// /// [`close`]: Self::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let mut buffer: Vec<&str> = Vec::with_capacity(2); /// let limit = 2; /// let (tx, mut rx) = mpsc::unbounded_channel(); /// let tx2 = tx.clone(); /// tx2.send("first").unwrap(); /// tx2.send("second").unwrap(); /// tx2.send("third").unwrap(); /// /// // Call `recv_many` to receive up to `limit` (2) values. /// assert_eq!(2, rx.recv_many(&mut buffer, limit).await); /// assert_eq!(vec!["first", "second"], buffer); /// /// // If the buffer is full, the next call to `recv_many` /// // reserves additional capacity. /// assert_eq!(1, rx.recv_many(&mut buffer, limit).await); /// /// tokio::spawn(async move { /// tx.send("fourth").unwrap(); /// }); /// /// // 'tx' is dropped, but `recv_many` /// // is guaranteed not to return 0 as the channel /// // is not yet closed. /// assert_eq!(1, rx.recv_many(&mut buffer, limit).await); /// assert_eq!(vec!["first", "second", "third", "fourth"], buffer); /// /// // Once the last sender is dropped, the channel is /// // closed and `recv_many` returns 0, capacity unchanged. /// drop(tx2); /// assert_eq!(0, rx.recv_many(&mut buffer, limit).await); /// assert_eq!(vec!["first", "second", "third", "fourth"], buffer); /// } /// ``` pub async fn recv_many(&mut self, buffer: &mut Vec, limit: usize) -> usize { use std::future::poll_fn; poll_fn(|cx| self.chan.recv_many(cx, buffer, limit)).await } /// Tries to receive the next value for this receiver. /// /// This method returns the [`Empty`] error if the channel is currently /// empty, but there are still outstanding [senders] or [permits]. /// /// This method returns the [`Disconnected`] error if the channel is /// currently empty, and there are no outstanding [senders] or [permits]. /// /// Unlike the [`poll_recv`] method, this method will never return an /// [`Empty`] error spuriously. /// /// [`Empty`]: crate::sync::mpsc::error::TryRecvError::Empty /// [`Disconnected`]: crate::sync::mpsc::error::TryRecvError::Disconnected /// [`poll_recv`]: Self::poll_recv /// [senders]: crate::sync::mpsc::Sender /// [permits]: crate::sync::mpsc::Permit /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// use tokio::sync::mpsc::error::TryRecvError; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::unbounded_channel(); /// /// tx.send("hello").unwrap(); /// /// assert_eq!(Ok("hello"), rx.try_recv()); /// assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); /// /// tx.send("hello").unwrap(); /// // Drop the last sender, closing the channel. /// drop(tx); /// /// assert_eq!(Ok("hello"), rx.try_recv()); /// assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); /// } /// ``` pub fn try_recv(&mut self) -> Result { self.chan.try_recv() } /// Blocking receive to call outside of asynchronous contexts. /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// /// ``` /// use std::thread; /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::unbounded_channel::(); /// /// let sync_code = thread::spawn(move || { /// assert_eq!(Some(10), rx.blocking_recv()); /// }); /// /// let _ = tx.send(10); /// sync_code.join().unwrap(); /// } /// ``` #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(&mut self) -> Option { crate::future::block_on(self.recv()) } /// Variant of [`Self::recv_many`] for blocking contexts. /// /// The same conditions as in [`Self::blocking_recv`] apply. #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "recv_many_blocking"))] pub fn blocking_recv_many(&mut self, buffer: &mut Vec, limit: usize) -> usize { crate::future::block_on(self.recv_many(buffer, limit)) } /// Closes the receiving half of a channel, without dropping it. /// /// This prevents any further messages from being sent on the channel while /// still enabling the receiver to drain messages that are buffered. /// /// To guarantee that no messages are dropped, after calling `close()`, /// `recv()` must be called until `None` is returned. pub fn close(&mut self) { self.chan.close(); } /// Checks if a channel is closed. /// /// This method returns `true` if the channel has been closed. The channel is closed /// when all [`UnboundedSender`] have been dropped, or when [`UnboundedReceiver::close`] is called. /// /// [`UnboundedSender`]: crate::sync::mpsc::UnboundedSender /// [`UnboundedReceiver::close`]: crate::sync::mpsc::UnboundedReceiver::close /// /// # Examples /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (_tx, mut rx) = mpsc::unbounded_channel::<()>(); /// assert!(!rx.is_closed()); /// /// rx.close(); /// /// assert!(rx.is_closed()); /// } /// ``` pub fn is_closed(&self) -> bool { self.chan.is_closed() } /// Checks if a channel is empty. /// /// This method returns `true` if the channel has no messages. /// /// # Examples /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::unbounded_channel(); /// assert!(rx.is_empty()); /// /// tx.send(0).unwrap(); /// assert!(!rx.is_empty()); /// } /// /// ``` pub fn is_empty(&self) -> bool { self.chan.is_empty() } /// Returns the number of messages in the channel. /// /// # Examples /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::unbounded_channel(); /// assert_eq!(0, rx.len()); /// /// tx.send(0).unwrap(); /// assert_eq!(1, rx.len()); /// } /// ``` pub fn len(&self) -> usize { self.chan.len() } /// Polls to receive the next message on this channel. /// /// This method returns: /// /// * `Poll::Pending` if no messages are available but the channel is not /// closed, or if a spurious failure happens. /// * `Poll::Ready(Some(message))` if a message is available. /// * `Poll::Ready(None)` if the channel has been closed and all messages /// sent before it was closed have been received. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when a message is sent on any /// receiver, or when the channel is closed. Note that on multiple calls to /// `poll_recv` or `poll_recv_many`, only the `Waker` from the `Context` /// passed to the most recent call is scheduled to receive a wakeup. /// /// If this method returns `Poll::Pending` due to a spurious failure, then /// the `Waker` will be notified when the situation causing the spurious /// failure has been resolved. Note that receiving such a wakeup does not /// guarantee that the next call will succeed — it could fail with another /// spurious failure. pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.chan.recv(cx) } /// Polls to receive multiple messages on this channel, extending the provided buffer. /// /// This method returns: /// * `Poll::Pending` if no messages are available but the channel is not closed, or if a /// spurious failure happens. /// * `Poll::Ready(count)` where `count` is the number of messages successfully received and /// stored in `buffer`. This can be less than, or equal to, `limit`. /// * `Poll::Ready(0)` if `limit` is set to zero or when the channel is closed. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided /// `Context` is scheduled to receive a wakeup when a message is sent on any /// receiver, or when the channel is closed. Note that on multiple calls to /// `poll_recv` or `poll_recv_many`, only the `Waker` from the `Context` /// passed to the most recent call is scheduled to receive a wakeup. /// /// Note that this method does not guarantee that exactly `limit` messages /// are received. Rather, if at least one message is available, it returns /// as many messages as it can up to the given limit. This method returns /// zero only if the channel is closed (or if `limit` is zero). /// /// # Examples /// /// ``` /// use std::task::{Context, Poll}; /// use std::pin::Pin; /// use tokio::sync::mpsc; /// use futures::Future; /// /// struct MyReceiverFuture<'a> { /// receiver: mpsc::UnboundedReceiver, /// buffer: &'a mut Vec, /// limit: usize, /// } /// /// impl<'a> Future for MyReceiverFuture<'a> { /// type Output = usize; // Number of messages received /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { /// let MyReceiverFuture { receiver, buffer, limit } = &mut *self; /// /// // Now `receiver` and `buffer` are mutable references, and `limit` is copied /// match receiver.poll_recv_many(cx, *buffer, *limit) { /// Poll::Pending => Poll::Pending, /// Poll::Ready(count) => Poll::Ready(count), /// } /// } /// } /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = mpsc::unbounded_channel::(); /// let mut buffer = Vec::new(); /// /// let my_receiver_future = MyReceiverFuture { /// receiver: rx, /// buffer: &mut buffer, /// limit: 3, /// }; /// /// for i in 0..10 { /// tx.send(i).expect("Unable to send integer"); /// } /// /// let count = my_receiver_future.await; /// assert_eq!(count, 3); /// assert_eq!(buffer, vec![0,1,2]) /// } /// ``` pub fn poll_recv_many( &mut self, cx: &mut Context<'_>, buffer: &mut Vec, limit: usize, ) -> Poll { self.chan.recv_many(cx, buffer, limit) } /// Returns the number of [`UnboundedSender`] handles. pub fn sender_strong_count(&self) -> usize { self.chan.sender_strong_count() } /// Returns the number of [`WeakUnboundedSender`] handles. pub fn sender_weak_count(&self) -> usize { self.chan.sender_weak_count() } } impl UnboundedSender { pub(crate) fn new(chan: chan::Tx) -> UnboundedSender { UnboundedSender { chan } } /// Attempts to send a message on this `UnboundedSender` without blocking. /// /// This method is not marked async because sending a message to an unbounded channel /// never requires any form of waiting. Because of this, the `send` method can be /// used in both synchronous and asynchronous code without problems. /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`UnboundedReceiver`] having been dropped, this /// function returns an error. The error includes the value passed to `send`. /// /// [`close`]: UnboundedReceiver::close /// [`UnboundedReceiver`]: UnboundedReceiver pub fn send(&self, message: T) -> Result<(), SendError> { if !self.inc_num_messages() { return Err(SendError(message)); } self.chan.send(message); Ok(()) } fn inc_num_messages(&self) -> bool { use std::process; use std::sync::atomic::Ordering::{AcqRel, Acquire}; let mut curr = self.chan.semaphore().0.load(Acquire); loop { if curr & 1 == 1 { return false; } if curr == usize::MAX ^ 1 { // Overflowed the ref count. There is no safe way to recover, so // abort the process. In practice, this should never happen. process::abort() } match self .chan .semaphore() .0 .compare_exchange(curr, curr + 2, AcqRel, Acquire) { Ok(_) => return true, Err(actual) => { curr = actual; } } } } /// Completes when the receiver has dropped. /// /// This allows the producers to get notified when interest in the produced /// values is canceled and immediately stop doing work. /// /// # Cancel safety /// /// This method is cancel safe. Once the channel is closed, it stays closed /// forever and all future calls to `closed` will return immediately. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx1, rx) = mpsc::unbounded_channel::<()>(); /// let tx2 = tx1.clone(); /// let tx3 = tx1.clone(); /// let tx4 = tx1.clone(); /// let tx5 = tx1.clone(); /// tokio::spawn(async move { /// drop(rx); /// }); /// /// futures::join!( /// tx1.closed(), /// tx2.closed(), /// tx3.closed(), /// tx4.closed(), /// tx5.closed() /// ); //// println!("Receiver dropped"); /// } /// ``` pub async fn closed(&self) { self.chan.closed().await; } /// Checks if the channel has been closed. This happens when the /// [`UnboundedReceiver`] is dropped, or when the /// [`UnboundedReceiver::close`] method is called. /// /// [`UnboundedReceiver`]: crate::sync::mpsc::UnboundedReceiver /// [`UnboundedReceiver::close`]: crate::sync::mpsc::UnboundedReceiver::close /// /// ``` /// let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<()>(); /// assert!(!tx.is_closed()); /// /// let tx2 = tx.clone(); /// assert!(!tx2.is_closed()); /// /// drop(rx); /// assert!(tx.is_closed()); /// assert!(tx2.is_closed()); /// ``` pub fn is_closed(&self) -> bool { self.chan.is_closed() } /// Returns `true` if senders belong to the same channel. /// /// # Examples /// /// ``` /// let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<()>(); /// let tx2 = tx.clone(); /// assert!(tx.same_channel(&tx2)); /// /// let (tx3, rx3) = tokio::sync::mpsc::unbounded_channel::<()>(); /// assert!(!tx3.same_channel(&tx2)); /// ``` pub fn same_channel(&self, other: &Self) -> bool { self.chan.same_channel(&other.chan) } /// Converts the `UnboundedSender` to a [`WeakUnboundedSender`] that does not count /// towards RAII semantics, i.e. if all `UnboundedSender` instances of the /// channel were dropped and only `WeakUnboundedSender` instances remain, /// the channel is closed. #[must_use = "Downgrade creates a WeakSender without destroying the original non-weak sender."] pub fn downgrade(&self) -> WeakUnboundedSender { WeakUnboundedSender { chan: self.chan.downgrade(), } } /// Returns the number of [`UnboundedSender`] handles. pub fn strong_count(&self) -> usize { self.chan.strong_count() } /// Returns the number of [`WeakUnboundedSender`] handles. pub fn weak_count(&self) -> usize { self.chan.weak_count() } } impl Clone for WeakUnboundedSender { fn clone(&self) -> Self { self.chan.increment_weak_count(); WeakUnboundedSender { chan: self.chan.clone(), } } } impl Drop for WeakUnboundedSender { fn drop(&mut self) { self.chan.decrement_weak_count(); } } impl WeakUnboundedSender { /// Tries to convert a `WeakUnboundedSender` into an [`UnboundedSender`]. /// This will return `Some` if there are other `Sender` instances alive and /// the channel wasn't previously dropped, otherwise `None` is returned. pub fn upgrade(&self) -> Option> { chan::Tx::upgrade(self.chan.clone()).map(UnboundedSender::new) } /// Returns the number of [`UnboundedSender`] handles. pub fn strong_count(&self) -> usize { self.chan.strong_count() } /// Returns the number of [`WeakUnboundedSender`] handles. pub fn weak_count(&self) -> usize { self.chan.weak_count() } } impl fmt::Debug for WeakUnboundedSender { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("WeakUnboundedSender").finish() } } tokio-1.43.1/src/sync/mutex.rs000064400000000000000000001320661046102023000143220ustar 00000000000000#![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))] use crate::sync::batch_semaphore as semaphore; #[cfg(all(tokio_unstable, feature = "tracing"))] use crate::util::trace; use std::cell::UnsafeCell; use std::error::Error; use std::marker::PhantomData; use std::ops::{Deref, DerefMut}; use std::sync::Arc; use std::{fmt, mem, ptr}; /// An asynchronous `Mutex`-like type. /// /// This type acts similarly to [`std::sync::Mutex`], with two major /// differences: [`lock`] is an async method so does not block, and the lock /// guard is designed to be held across `.await` points. /// /// Tokio's Mutex operates on a guaranteed FIFO basis. /// This means that the order in which tasks call the [`lock`] method is /// the exact order in which they will acquire the lock. /// /// # Which kind of mutex should you use? /// /// Contrary to popular belief, it is ok and often preferred to use the ordinary /// [`Mutex`][std] from the standard library in asynchronous code. /// /// The feature that the async mutex offers over the blocking mutex is the /// ability to keep it locked across an `.await` point. This makes the async /// mutex more expensive than the blocking mutex, so the blocking mutex should /// be preferred in the cases where it can be used. The primary use case for the /// async mutex is to provide shared mutable access to IO resources such as a /// database connection. If the value behind the mutex is just data, it's /// usually appropriate to use a blocking mutex such as the one in the standard /// library or [`parking_lot`]. /// /// Note that, although the compiler will not prevent the std `Mutex` from holding /// its guard across `.await` points in situations where the task is not movable /// between threads, this virtually never leads to correct concurrent code in /// practice as it can easily lead to deadlocks. /// /// A common pattern is to wrap the `Arc>` in a struct that provides /// non-async methods for performing operations on the data within, and only /// lock the mutex inside these methods. The [mini-redis] example provides an /// illustration of this pattern. /// /// Additionally, when you _do_ want shared access to an IO resource, it is /// often better to spawn a task to manage the IO resource, and to use message /// passing to communicate with that task. /// /// [std]: std::sync::Mutex /// [`parking_lot`]: https://docs.rs/parking_lot /// [mini-redis]: https://github.com/tokio-rs/mini-redis/blob/master/src/db.rs /// /// # Examples: /// /// ```rust,no_run /// use tokio::sync::Mutex; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let data1 = Arc::new(Mutex::new(0)); /// let data2 = Arc::clone(&data1); /// /// tokio::spawn(async move { /// let mut lock = data2.lock().await; /// *lock += 1; /// }); /// /// let mut lock = data1.lock().await; /// *lock += 1; /// } /// ``` /// /// /// ```rust,no_run /// use tokio::sync::Mutex; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let count = Arc::new(Mutex::new(0)); /// /// for i in 0..5 { /// let my_count = Arc::clone(&count); /// tokio::spawn(async move { /// for j in 0..10 { /// let mut lock = my_count.lock().await; /// *lock += 1; /// println!("{} {} {}", i, j, lock); /// } /// }); /// } /// /// loop { /// if *count.lock().await >= 50 { /// break; /// } /// } /// println!("Count hit 50."); /// } /// ``` /// There are a few things of note here to pay attention to in this example. /// 1. The mutex is wrapped in an [`Arc`] to allow it to be shared across /// threads. /// 2. Each spawned task obtains a lock and releases it on every iteration. /// 3. Mutation of the data protected by the Mutex is done by de-referencing /// the obtained lock as seen on lines 13 and 20. /// /// Tokio's Mutex works in a simple FIFO (first in, first out) style where all /// calls to [`lock`] complete in the order they were performed. In that way the /// Mutex is "fair" and predictable in how it distributes the locks to inner /// data. Locks are released and reacquired after every iteration, so basically, /// each thread goes to the back of the line after it increments the value once. /// Note that there's some unpredictability to the timing between when the /// threads are started, but once they are going they alternate predictably. /// Finally, since there is only a single valid lock at any given time, there is /// no possibility of a race condition when mutating the inner value. /// /// Note that in contrast to [`std::sync::Mutex`], this implementation does not /// poison the mutex when a thread holding the [`MutexGuard`] panics. In such a /// case, the mutex will be unlocked. If the panic is caught, this might leave /// the data protected by the mutex in an inconsistent state. /// /// [`Mutex`]: struct@Mutex /// [`MutexGuard`]: struct@MutexGuard /// [`Arc`]: struct@std::sync::Arc /// [`std::sync::Mutex`]: struct@std::sync::Mutex /// [`Send`]: trait@std::marker::Send /// [`lock`]: method@Mutex::lock pub struct Mutex { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, s: semaphore::Semaphore, c: UnsafeCell, } /// A handle to a held `Mutex`. The guard can be held across any `.await` point /// as it is [`Send`]. /// /// As long as you have this guard, you have exclusive access to the underlying /// `T`. The guard internally borrows the `Mutex`, so the mutex will not be /// dropped while a guard exists. /// /// The lock is automatically released whenever the guard is dropped, at which /// point `lock` will succeed yet again. #[clippy::has_significant_drop] #[must_use = "if unused the Mutex will immediately unlock"] pub struct MutexGuard<'a, T: ?Sized> { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, lock: &'a Mutex, } /// An owned handle to a held `Mutex`. /// /// This guard is only available from a `Mutex` that is wrapped in an [`Arc`]. It /// is identical to `MutexGuard`, except that rather than borrowing the `Mutex`, /// it clones the `Arc`, incrementing the reference count. This means that /// unlike `MutexGuard`, it will have the `'static` lifetime. /// /// As long as you have this guard, you have exclusive access to the underlying /// `T`. The guard internally keeps a reference-counted pointer to the original /// `Mutex`, so even if the lock goes away, the guard remains valid. /// /// The lock is automatically released whenever the guard is dropped, at which /// point `lock` will succeed yet again. /// /// [`Arc`]: std::sync::Arc #[clippy::has_significant_drop] pub struct OwnedMutexGuard { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, lock: Arc>, } /// A handle to a held `Mutex` that has had a function applied to it via [`MutexGuard::map`]. /// /// This can be used to hold a subfield of the protected data. /// /// [`MutexGuard::map`]: method@MutexGuard::map #[clippy::has_significant_drop] #[must_use = "if unused the Mutex will immediately unlock"] pub struct MappedMutexGuard<'a, T: ?Sized> { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, s: &'a semaphore::Semaphore, data: *mut T, // Needed to tell the borrow checker that we are holding a `&mut T` marker: PhantomData<&'a mut T>, } /// A owned handle to a held `Mutex` that has had a function applied to it via /// [`OwnedMutexGuard::map`]. /// /// This can be used to hold a subfield of the protected data. /// /// [`OwnedMutexGuard::map`]: method@OwnedMutexGuard::map #[clippy::has_significant_drop] #[must_use = "if unused the Mutex will immediately unlock"] pub struct OwnedMappedMutexGuard { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, data: *mut U, lock: Arc>, } /// A helper type used when taking apart a `MutexGuard` without running its /// Drop implementation. #[allow(dead_code)] // Unused fields are still used in Drop. struct MutexGuardInner<'a, T: ?Sized> { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, lock: &'a Mutex, } /// A helper type used when taking apart a `OwnedMutexGuard` without running /// its Drop implementation. struct OwnedMutexGuardInner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, lock: Arc>, } /// A helper type used when taking apart a `MappedMutexGuard` without running /// its Drop implementation. #[allow(dead_code)] // Unused fields are still used in Drop. struct MappedMutexGuardInner<'a, T: ?Sized> { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, s: &'a semaphore::Semaphore, data: *mut T, } /// A helper type used when taking apart a `OwnedMappedMutexGuard` without running /// its Drop implementation. #[allow(dead_code)] // Unused fields are still used in Drop. struct OwnedMappedMutexGuardInner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, data: *mut U, lock: Arc>, } // As long as T: Send, it's fine to send and share Mutex between threads. // If T was not Send, sending and sharing a Mutex would be bad, since you can // access T through Mutex. unsafe impl Send for Mutex where T: ?Sized + Send {} unsafe impl Sync for Mutex where T: ?Sized + Send {} unsafe impl Sync for MutexGuard<'_, T> where T: ?Sized + Send + Sync {} unsafe impl Sync for OwnedMutexGuard where T: ?Sized + Send + Sync {} unsafe impl<'a, T> Sync for MappedMutexGuard<'a, T> where T: ?Sized + Sync + 'a {} unsafe impl<'a, T> Send for MappedMutexGuard<'a, T> where T: ?Sized + Send + 'a {} unsafe impl Sync for OwnedMappedMutexGuard where T: ?Sized + Send + Sync, U: ?Sized + Send + Sync, { } unsafe impl Send for OwnedMappedMutexGuard where T: ?Sized + Send, U: ?Sized + Send, { } /// Error returned from the [`Mutex::try_lock`], [`RwLock::try_read`] and /// [`RwLock::try_write`] functions. /// /// `Mutex::try_lock` operation will only fail if the mutex is already locked. /// /// `RwLock::try_read` operation will only fail if the lock is currently held /// by an exclusive writer. /// /// `RwLock::try_write` operation will only fail if the lock is currently held /// by any reader or by an exclusive writer. /// /// [`Mutex::try_lock`]: Mutex::try_lock /// [`RwLock::try_read`]: fn@super::RwLock::try_read /// [`RwLock::try_write`]: fn@super::RwLock::try_write #[derive(Debug)] pub struct TryLockError(pub(super) ()); impl fmt::Display for TryLockError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "operation would block") } } impl Error for TryLockError {} #[test] #[cfg(not(loom))] fn bounds() { fn check_send() {} fn check_unpin() {} // This has to take a value, since the async fn's return type is unnameable. fn check_send_sync_val(_t: T) {} fn check_send_sync() {} fn check_static() {} fn check_static_val(_t: T) {} check_send::>(); check_send::>(); check_unpin::>(); check_send_sync::>(); check_static::>(); let mutex = Mutex::new(1); check_send_sync_val(mutex.lock()); let arc_mutex = Arc::new(Mutex::new(1)); check_send_sync_val(arc_mutex.clone().lock_owned()); check_static_val(arc_mutex.lock_owned()); } impl Mutex { /// Creates a new lock in an unlocked state ready for use. /// /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// /// let lock = Mutex::new(5); /// ``` #[track_caller] pub fn new(t: T) -> Self where T: Sized, { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = std::panic::Location::caller(); tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Mutex", kind = "Sync", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ) }; #[cfg(all(tokio_unstable, feature = "tracing"))] let s = resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = false, ); semaphore::Semaphore::new(1) }); #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let s = semaphore::Semaphore::new(1); Self { c: UnsafeCell::new(t), s, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Creates a new lock in an unlocked state ready for use. /// /// When using the `tracing` [unstable feature], a `Mutex` created with /// `const_new` will not be instrumented. As such, it will not be visible /// in [`tokio-console`]. Instead, [`Mutex::new`] should be used to create /// an instrumented object if that is needed. /// /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// /// static LOCK: Mutex = Mutex::const_new(5); /// ``` /// /// [`tokio-console`]: https://github.com/tokio-rs/console /// [unstable feature]: crate#unstable-features #[cfg(not(all(loom, test)))] pub const fn const_new(t: T) -> Self where T: Sized, { Self { c: UnsafeCell::new(t), s: semaphore::Semaphore::const_new(1), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Locks this mutex, causing the current task to yield until the lock has /// been acquired. When the lock has been acquired, function returns a /// [`MutexGuard`]. /// /// If the mutex is available to be acquired immediately, then this call /// will typically not yield to the runtime. However, this is not guaranteed /// under all circumstances. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute locks in the order they /// were requested. Cancelling a call to `lock` makes you lose your place in /// the queue. /// /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// /// #[tokio::main] /// async fn main() { /// let mutex = Mutex::new(1); /// /// let mut n = mutex.lock().await; /// *n = 2; /// } /// ``` pub async fn lock(&self) -> MutexGuard<'_, T> { let acquire_fut = async { self.acquire().await; MutexGuard { lock: self, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), } }; #[cfg(all(tokio_unstable, feature = "tracing"))] let acquire_fut = trace::async_op( move || acquire_fut, self.resource_span.clone(), "Mutex::lock", "poll", false, ); #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing let guard = acquire_fut.await; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = true, ); }); guard } /// Blockingly locks this `Mutex`. When the lock has been acquired, function returns a /// [`MutexGuard`]. /// /// This method is intended for use cases where you /// need to use this mutex in asynchronous code as well as in synchronous code. /// /// # Panics /// /// This function panics if called within an asynchronous execution context. /// /// - If you find yourself in an asynchronous execution context and needing /// to call some (synchronous) function which performs one of these /// `blocking_` operations, then consider wrapping that call inside /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] /// (or [`block_in_place()`][crate::task::block_in_place]). /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Mutex; /// /// #[tokio::main] /// async fn main() { /// let mutex = Arc::new(Mutex::new(1)); /// let lock = mutex.lock().await; /// /// let mutex1 = Arc::clone(&mutex); /// let blocking_task = tokio::task::spawn_blocking(move || { /// // This shall block until the `lock` is released. /// let mut n = mutex1.blocking_lock(); /// *n = 2; /// }); /// /// assert_eq!(*lock, 1); /// // Release the lock. /// drop(lock); /// /// // Await the completion of the blocking task. /// blocking_task.await.unwrap(); /// /// // Assert uncontended. /// let n = mutex.try_lock().unwrap(); /// assert_eq!(*n, 2); /// } /// /// ``` #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "lock_blocking"))] pub fn blocking_lock(&self) -> MutexGuard<'_, T> { crate::future::block_on(self.lock()) } /// Blockingly locks this `Mutex`. When the lock has been acquired, function returns an /// [`OwnedMutexGuard`]. /// /// This method is identical to [`Mutex::blocking_lock`], except that the returned /// guard references the `Mutex` with an [`Arc`] rather than by borrowing /// it. Therefore, the `Mutex` must be wrapped in an `Arc` to call this /// method, and the guard will live for the `'static` lifetime, as it keeps /// the `Mutex` alive by holding an `Arc`. /// /// # Panics /// /// This function panics if called within an asynchronous execution context. /// /// - If you find yourself in an asynchronous execution context and needing /// to call some (synchronous) function which performs one of these /// `blocking_` operations, then consider wrapping that call inside /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] /// (or [`block_in_place()`][crate::task::block_in_place]). /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Mutex; /// /// #[tokio::main] /// async fn main() { /// let mutex = Arc::new(Mutex::new(1)); /// let lock = mutex.lock().await; /// /// let mutex1 = Arc::clone(&mutex); /// let blocking_task = tokio::task::spawn_blocking(move || { /// // This shall block until the `lock` is released. /// let mut n = mutex1.blocking_lock_owned(); /// *n = 2; /// }); /// /// assert_eq!(*lock, 1); /// // Release the lock. /// drop(lock); /// /// // Await the completion of the blocking task. /// blocking_task.await.unwrap(); /// /// // Assert uncontended. /// let n = mutex.try_lock().unwrap(); /// assert_eq!(*n, 2); /// } /// /// ``` #[track_caller] #[cfg(feature = "sync")] pub fn blocking_lock_owned(self: Arc) -> OwnedMutexGuard { crate::future::block_on(self.lock_owned()) } /// Locks this mutex, causing the current task to yield until the lock has /// been acquired. When the lock has been acquired, this returns an /// [`OwnedMutexGuard`]. /// /// If the mutex is available to be acquired immediately, then this call /// will typically not yield to the runtime. However, this is not guaranteed /// under all circumstances. /// /// This method is identical to [`Mutex::lock`], except that the returned /// guard references the `Mutex` with an [`Arc`] rather than by borrowing /// it. Therefore, the `Mutex` must be wrapped in an `Arc` to call this /// method, and the guard will live for the `'static` lifetime, as it keeps /// the `Mutex` alive by holding an `Arc`. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute locks in the order they /// were requested. Cancelling a call to `lock_owned` makes you lose your /// place in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let mutex = Arc::new(Mutex::new(1)); /// /// let mut n = mutex.clone().lock_owned().await; /// *n = 2; /// } /// ``` /// /// [`Arc`]: std::sync::Arc pub async fn lock_owned(self: Arc) -> OwnedMutexGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = self.resource_span.clone(); let acquire_fut = async { self.acquire().await; OwnedMutexGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), lock: self, } }; #[cfg(all(tokio_unstable, feature = "tracing"))] let acquire_fut = trace::async_op( move || acquire_fut, resource_span, "Mutex::lock_owned", "poll", false, ); #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing let guard = acquire_fut.await; #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = true, ); }); guard } async fn acquire(&self) { crate::trace::async_trace_leaf().await; self.s.acquire(1).await.unwrap_or_else(|_| { // The semaphore was closed. but, we never explicitly close it, and // we own it exclusively, which means that this can never happen. unreachable!() }); } /// Attempts to acquire the lock, and returns [`TryLockError`] if the /// lock is currently held somewhere else. /// /// [`TryLockError`]: TryLockError /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// # async fn dox() -> Result<(), tokio::sync::TryLockError> { /// /// let mutex = Mutex::new(1); /// /// let n = mutex.try_lock()?; /// assert_eq!(*n, 1); /// # Ok(()) /// # } /// ``` pub fn try_lock(&self) -> Result, TryLockError> { match self.s.try_acquire(1) { Ok(()) => { let guard = MutexGuard { lock: self, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), }; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = true, ); }); Ok(guard) } Err(_) => Err(TryLockError(())), } } /// Returns a mutable reference to the underlying data. /// /// Since this call borrows the `Mutex` mutably, no actual locking needs to /// take place -- the mutable borrow statically guarantees no locks exist. /// /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// /// fn main() { /// let mut mutex = Mutex::new(1); /// /// let n = mutex.get_mut(); /// *n = 2; /// } /// ``` pub fn get_mut(&mut self) -> &mut T { unsafe { // Safety: This is https://github.com/rust-lang/rust/pull/76936 &mut *self.c.get() } } /// Attempts to acquire the lock, and returns [`TryLockError`] if the lock /// is currently held somewhere else. /// /// This method is identical to [`Mutex::try_lock`], except that the /// returned guard references the `Mutex` with an [`Arc`] rather than by /// borrowing it. Therefore, the `Mutex` must be wrapped in an `Arc` to call /// this method, and the guard will live for the `'static` lifetime, as it /// keeps the `Mutex` alive by holding an `Arc`. /// /// [`TryLockError`]: TryLockError /// [`Arc`]: std::sync::Arc /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// use std::sync::Arc; /// # async fn dox() -> Result<(), tokio::sync::TryLockError> { /// /// let mutex = Arc::new(Mutex::new(1)); /// /// let n = mutex.clone().try_lock_owned()?; /// assert_eq!(*n, 1); /// # Ok(()) /// # } pub fn try_lock_owned(self: Arc) -> Result, TryLockError> { match self.s.try_acquire(1) { Ok(()) => { let guard = OwnedMutexGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), lock: self, }; #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = true, ); }); Ok(guard) } Err(_) => Err(TryLockError(())), } } /// Consumes the mutex, returning the underlying data. /// # Examples /// /// ``` /// use tokio::sync::Mutex; /// /// #[tokio::main] /// async fn main() { /// let mutex = Mutex::new(1); /// /// let n = mutex.into_inner(); /// assert_eq!(n, 1); /// } /// ``` pub fn into_inner(self) -> T where T: Sized, { self.c.into_inner() } } impl From for Mutex { fn from(s: T) -> Self { Self::new(s) } } impl Default for Mutex where T: Default, { fn default() -> Self { Self::new(T::default()) } } impl std::fmt::Debug for Mutex where T: std::fmt::Debug, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let mut d = f.debug_struct("Mutex"); match self.try_lock() { Ok(inner) => d.field("data", &&*inner), Err(_) => d.field("data", &format_args!("")), }; d.finish() } } // === impl MutexGuard === impl<'a, T: ?Sized> MutexGuard<'a, T> { fn skip_drop(self) -> MutexGuardInner<'a, T> { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the `resource_span` and then forgets the // original. In the end, we have not duplicated or forgotten any values. MutexGuardInner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: unsafe { std::ptr::read(&me.resource_span) }, lock: me.lock, } } /// Makes a new [`MappedMutexGuard`] for a component of the locked data. /// /// This operation cannot fail as the [`MutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `MutexGuard::map(...)`. A method /// would interfere with methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use tokio::sync::{Mutex, MutexGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let foo = Mutex::new(Foo(1)); /// /// { /// let mut mapped = MutexGuard::map(foo.lock().await, |f| &mut f.0); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *foo.lock().await); /// # } /// ``` /// /// [`MutexGuard`]: struct@MutexGuard /// [`MappedMutexGuard`]: struct@MappedMutexGuard #[inline] pub fn map(mut this: Self, f: F) -> MappedMutexGuard<'a, U> where U: ?Sized, F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; let inner = this.skip_drop(); MappedMutexGuard { s: &inner.lock.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, } } /// Attempts to make a new [`MappedMutexGuard`] for a component of the locked data. The /// original guard is returned if the closure returns `None`. /// /// This operation cannot fail as the [`MutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `MutexGuard::try_map(...)`. A /// method would interfere with methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use tokio::sync::{Mutex, MutexGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let foo = Mutex::new(Foo(1)); /// /// { /// let mut mapped = MutexGuard::try_map(foo.lock().await, |f| Some(&mut f.0)) /// .expect("should not fail"); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *foo.lock().await); /// # } /// ``` /// /// [`MutexGuard`]: struct@MutexGuard /// [`MappedMutexGuard`]: struct@MappedMutexGuard #[inline] pub fn try_map(mut this: Self, f: F) -> Result, Self> where U: ?Sized, F: FnOnce(&mut T) -> Option<&mut U>, { let data = match f(&mut *this) { Some(data) => data as *mut U, None => return Err(this), }; let inner = this.skip_drop(); Ok(MappedMutexGuard { s: &inner.lock.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, }) } /// Returns a reference to the original `Mutex`. /// /// ``` /// use tokio::sync::{Mutex, MutexGuard}; /// /// async fn unlock_and_relock<'l>(guard: MutexGuard<'l, u32>) -> MutexGuard<'l, u32> { /// println!("1. contains: {:?}", *guard); /// let mutex = MutexGuard::mutex(&guard); /// drop(guard); /// let guard = mutex.lock().await; /// println!("2. contains: {:?}", *guard); /// guard /// } /// # /// # #[tokio::main] /// # async fn main() { /// # let mutex = Mutex::new(0u32); /// # let guard = mutex.lock().await; /// # let _guard = unlock_and_relock(guard).await; /// # } /// ``` #[inline] pub fn mutex(this: &Self) -> &'a Mutex { this.lock } } impl Drop for MutexGuard<'_, T> { fn drop(&mut self) { self.lock.s.release(1); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = false, ); }); } } impl Deref for MutexGuard<'_, T> { type Target = T; fn deref(&self) -> &Self::Target { unsafe { &*self.lock.c.get() } } } impl DerefMut for MutexGuard<'_, T> { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { &mut *self.lock.c.get() } } } impl fmt::Debug for MutexGuard<'_, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for MutexGuard<'_, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } // === impl OwnedMutexGuard === impl OwnedMutexGuard { fn skip_drop(self) -> OwnedMutexGuardInner { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. unsafe { OwnedMutexGuardInner { lock: ptr::read(&me.lock), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: ptr::read(&me.resource_span), } } } /// Makes a new [`OwnedMappedMutexGuard`] for a component of the locked data. /// /// This operation cannot fail as the [`OwnedMutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `OwnedMutexGuard::map(...)`. A method /// would interfere with methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use tokio::sync::{Mutex, OwnedMutexGuard}; /// use std::sync::Arc; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let foo = Arc::new(Mutex::new(Foo(1))); /// /// { /// let mut mapped = OwnedMutexGuard::map(foo.clone().lock_owned().await, |f| &mut f.0); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *foo.lock().await); /// # } /// ``` /// /// [`OwnedMutexGuard`]: struct@OwnedMutexGuard /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard #[inline] pub fn map(mut this: Self, f: F) -> OwnedMappedMutexGuard where U: ?Sized, F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; let inner = this.skip_drop(); OwnedMappedMutexGuard { data, lock: inner.lock, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, } } /// Attempts to make a new [`OwnedMappedMutexGuard`] for a component of the locked data. The /// original guard is returned if the closure returns `None`. /// /// This operation cannot fail as the [`OwnedMutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `OwnedMutexGuard::try_map(...)`. A /// method would interfere with methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use tokio::sync::{Mutex, OwnedMutexGuard}; /// use std::sync::Arc; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let foo = Arc::new(Mutex::new(Foo(1))); /// /// { /// let mut mapped = OwnedMutexGuard::try_map(foo.clone().lock_owned().await, |f| Some(&mut f.0)) /// .expect("should not fail"); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *foo.lock().await); /// # } /// ``` /// /// [`OwnedMutexGuard`]: struct@OwnedMutexGuard /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard #[inline] pub fn try_map(mut this: Self, f: F) -> Result, Self> where U: ?Sized, F: FnOnce(&mut T) -> Option<&mut U>, { let data = match f(&mut *this) { Some(data) => data as *mut U, None => return Err(this), }; let inner = this.skip_drop(); Ok(OwnedMappedMutexGuard { data, lock: inner.lock, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, }) } /// Returns a reference to the original `Arc`. /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{Mutex, OwnedMutexGuard}; /// /// async fn unlock_and_relock(guard: OwnedMutexGuard) -> OwnedMutexGuard { /// println!("1. contains: {:?}", *guard); /// let mutex: Arc> = OwnedMutexGuard::mutex(&guard).clone(); /// drop(guard); /// let guard = mutex.lock_owned().await; /// println!("2. contains: {:?}", *guard); /// guard /// } /// # /// # #[tokio::main] /// # async fn main() { /// # let mutex = Arc::new(Mutex::new(0u32)); /// # let guard = mutex.lock_owned().await; /// # unlock_and_relock(guard).await; /// # } /// ``` #[inline] pub fn mutex(this: &Self) -> &Arc> { &this.lock } } impl Drop for OwnedMutexGuard { fn drop(&mut self) { self.lock.s.release(1); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = false, ); }); } } impl Deref for OwnedMutexGuard { type Target = T; fn deref(&self) -> &Self::Target { unsafe { &*self.lock.c.get() } } } impl DerefMut for OwnedMutexGuard { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { &mut *self.lock.c.get() } } } impl fmt::Debug for OwnedMutexGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for OwnedMutexGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } // === impl MappedMutexGuard === impl<'a, T: ?Sized> MappedMutexGuard<'a, T> { fn skip_drop(self) -> MappedMutexGuardInner<'a, T> { let me = mem::ManuallyDrop::new(self); MappedMutexGuardInner { s: me.s, data: me.data, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: unsafe { std::ptr::read(&me.resource_span) }, } } /// Makes a new [`MappedMutexGuard`] for a component of the locked data. /// /// This operation cannot fail as the [`MappedMutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `MappedMutexGuard::map(...)`. A /// method would interfere with methods of the same name on the contents of the locked data. /// /// [`MappedMutexGuard`]: struct@MappedMutexGuard #[inline] pub fn map(mut this: Self, f: F) -> MappedMutexGuard<'a, U> where F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; let inner = this.skip_drop(); MappedMutexGuard { s: inner.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, } } /// Attempts to make a new [`MappedMutexGuard`] for a component of the locked data. The /// original guard is returned if the closure returns `None`. /// /// This operation cannot fail as the [`MappedMutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `MappedMutexGuard::try_map(...)`. A /// method would interfere with methods of the same name on the contents of the locked data. /// /// [`MappedMutexGuard`]: struct@MappedMutexGuard #[inline] pub fn try_map(mut this: Self, f: F) -> Result, Self> where F: FnOnce(&mut T) -> Option<&mut U>, { let data = match f(&mut *this) { Some(data) => data as *mut U, None => return Err(this), }; let inner = this.skip_drop(); Ok(MappedMutexGuard { s: inner.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, }) } } impl<'a, T: ?Sized> Drop for MappedMutexGuard<'a, T> { fn drop(&mut self) { self.s.release(1); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = false, ); }); } } impl<'a, T: ?Sized> Deref for MappedMutexGuard<'a, T> { type Target = T; fn deref(&self) -> &Self::Target { unsafe { &*self.data } } } impl<'a, T: ?Sized> DerefMut for MappedMutexGuard<'a, T> { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { &mut *self.data } } } impl<'a, T: ?Sized + fmt::Debug> fmt::Debug for MappedMutexGuard<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl<'a, T: ?Sized + fmt::Display> fmt::Display for MappedMutexGuard<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } // === impl OwnedMappedMutexGuard === impl OwnedMappedMutexGuard { fn skip_drop(self) -> OwnedMappedMutexGuardInner { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. unsafe { OwnedMappedMutexGuardInner { data: me.data, lock: ptr::read(&me.lock), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: ptr::read(&me.resource_span), } } } /// Makes a new [`OwnedMappedMutexGuard`] for a component of the locked data. /// /// This operation cannot fail as the [`OwnedMappedMutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `OwnedMappedMutexGuard::map(...)`. A method /// would interfere with methods of the same name on the contents of the locked data. /// /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard #[inline] pub fn map(mut this: Self, f: F) -> OwnedMappedMutexGuard where F: FnOnce(&mut U) -> &mut S, { let data = f(&mut *this) as *mut S; let inner = this.skip_drop(); OwnedMappedMutexGuard { data, lock: inner.lock, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, } } /// Attempts to make a new [`OwnedMappedMutexGuard`] for a component of the locked data. The /// original guard is returned if the closure returns `None`. /// /// This operation cannot fail as the [`OwnedMutexGuard`] passed in already locked the mutex. /// /// This is an associated function that needs to be used as `OwnedMutexGuard::try_map(...)`. A /// method would interfere with methods of the same name on the contents of the locked data. /// /// [`OwnedMutexGuard`]: struct@OwnedMutexGuard /// [`OwnedMappedMutexGuard`]: struct@OwnedMappedMutexGuard #[inline] pub fn try_map(mut this: Self, f: F) -> Result, Self> where F: FnOnce(&mut U) -> Option<&mut S>, { let data = match f(&mut *this) { Some(data) => data as *mut S, None => return Err(this), }; let inner = this.skip_drop(); Ok(OwnedMappedMutexGuard { data, lock: inner.lock, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: inner.resource_span, }) } } impl Drop for OwnedMappedMutexGuard { fn drop(&mut self) { self.lock.s.release(1); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = false, ); }); } } impl Deref for OwnedMappedMutexGuard { type Target = U; fn deref(&self) -> &Self::Target { unsafe { &*self.data } } } impl DerefMut for OwnedMappedMutexGuard { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { &mut *self.data } } } impl fmt::Debug for OwnedMappedMutexGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for OwnedMappedMutexGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } tokio-1.43.1/src/sync/notify.rs000064400000000000000000001270431046102023000144670ustar 00000000000000// Allow `unreachable_pub` warnings when sync is not enabled // due to the usage of `Notify` within the `rt` feature set. // When this module is compiled with `sync` enabled we will warn on // this lint. When `rt` is enabled we use `pub(crate)` which // triggers this warning but it is safe to ignore in this case. #![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))] use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Mutex; use crate::util::linked_list::{self, GuardedLinkedList, LinkedList}; use crate::util::WakeList; use std::future::Future; use std::marker::PhantomPinned; use std::panic::{RefUnwindSafe, UnwindSafe}; use std::pin::Pin; use std::ptr::NonNull; use std::sync::atomic::Ordering::{self, Acquire, Relaxed, Release, SeqCst}; use std::task::{Context, Poll, Waker}; type WaitList = LinkedList::Target>; type GuardedWaitList = GuardedLinkedList::Target>; /// Notifies a single task to wake up. /// /// `Notify` provides a basic mechanism to notify a single task of an event. /// `Notify` itself does not carry any data. Instead, it is to be used to signal /// another task to perform an operation. /// /// A `Notify` can be thought of as a [`Semaphore`] starting with 0 permits. The /// [`notified().await`] method waits for a permit to become available, and /// [`notify_one()`] sets a permit **if there currently are no available /// permits**. /// /// The synchronization details of `Notify` are similar to /// [`thread::park`][park] and [`Thread::unpark`][unpark] from std. A [`Notify`] /// value contains a single permit. [`notified().await`] waits for the permit to /// be made available, consumes the permit, and resumes. [`notify_one()`] sets /// the permit, waking a pending task if there is one. /// /// If `notify_one()` is called **before** `notified().await`, then the next /// call to `notified().await` will complete immediately, consuming the permit. /// Any subsequent calls to `notified().await` will wait for a new permit. /// /// If `notify_one()` is called **multiple** times before `notified().await`, /// only a **single** permit is stored. The next call to `notified().await` will /// complete immediately, but the one after will wait for a new permit. /// /// # Examples /// /// Basic usage. /// /// ``` /// use tokio::sync::Notify; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let notify = Arc::new(Notify::new()); /// let notify2 = notify.clone(); /// /// let handle = tokio::spawn(async move { /// notify2.notified().await; /// println!("received notification"); /// }); /// /// println!("sending notification"); /// notify.notify_one(); /// /// // Wait for task to receive notification. /// handle.await.unwrap(); /// } /// ``` /// /// Unbound multi-producer single-consumer (mpsc) channel. /// /// No wakeups can be lost when using this channel because the call to /// `notify_one()` will store a permit in the `Notify`, which the following call /// to `notified()` will consume. /// /// ``` /// use tokio::sync::Notify; /// /// use std::collections::VecDeque; /// use std::sync::Mutex; /// /// struct Channel { /// values: Mutex>, /// notify: Notify, /// } /// /// impl Channel { /// pub fn send(&self, value: T) { /// self.values.lock().unwrap() /// .push_back(value); /// /// // Notify the consumer a value is available /// self.notify.notify_one(); /// } /// /// // This is a single-consumer channel, so several concurrent calls to /// // `recv` are not allowed. /// pub async fn recv(&self) -> T { /// loop { /// // Drain values /// if let Some(value) = self.values.lock().unwrap().pop_front() { /// return value; /// } /// /// // Wait for values to be available /// self.notify.notified().await; /// } /// } /// } /// ``` /// /// Unbound multi-producer multi-consumer (mpmc) channel. /// /// The call to [`enable`] is important because otherwise if you have two /// calls to `recv` and two calls to `send` in parallel, the following could /// happen: /// /// 1. Both calls to `try_recv` return `None`. /// 2. Both new elements are added to the vector. /// 3. The `notify_one` method is called twice, adding only a single /// permit to the `Notify`. /// 4. Both calls to `recv` reach the `Notified` future. One of them /// consumes the permit, and the other sleeps forever. /// /// By adding the `Notified` futures to the list by calling `enable` before /// `try_recv`, the `notify_one` calls in step three would remove the /// futures from the list and mark them notified instead of adding a permit /// to the `Notify`. This ensures that both futures are woken. /// /// Notice that this failure can only happen if there are two concurrent calls /// to `recv`. This is why the mpsc example above does not require a call to /// `enable`. /// /// ``` /// use tokio::sync::Notify; /// /// use std::collections::VecDeque; /// use std::sync::Mutex; /// /// struct Channel { /// messages: Mutex>, /// notify_on_sent: Notify, /// } /// /// impl Channel { /// pub fn send(&self, msg: T) { /// let mut locked_queue = self.messages.lock().unwrap(); /// locked_queue.push_back(msg); /// drop(locked_queue); /// /// // Send a notification to one of the calls currently /// // waiting in a call to `recv`. /// self.notify_on_sent.notify_one(); /// } /// /// pub fn try_recv(&self) -> Option { /// let mut locked_queue = self.messages.lock().unwrap(); /// locked_queue.pop_front() /// } /// /// pub async fn recv(&self) -> T { /// let future = self.notify_on_sent.notified(); /// tokio::pin!(future); /// /// loop { /// // Make sure that no wakeup is lost if we get /// // `None` from `try_recv`. /// future.as_mut().enable(); /// /// if let Some(msg) = self.try_recv() { /// return msg; /// } /// /// // Wait for a call to `notify_one`. /// // /// // This uses `.as_mut()` to avoid consuming the future, /// // which lets us call `Pin::set` below. /// future.as_mut().await; /// /// // Reset the future in case another call to /// // `try_recv` got the message before us. /// future.set(self.notify_on_sent.notified()); /// } /// } /// } /// ``` /// /// [park]: std::thread::park /// [unpark]: std::thread::Thread::unpark /// [`notified().await`]: Notify::notified() /// [`notify_one()`]: Notify::notify_one() /// [`enable`]: Notified::enable() /// [`Semaphore`]: crate::sync::Semaphore #[derive(Debug)] pub struct Notify { // `state` uses 2 bits to store one of `EMPTY`, // `WAITING` or `NOTIFIED`. The rest of the bits // are used to store the number of times `notify_waiters` // was called. // // Throughout the code there are two assumptions: // - state can be transitioned *from* `WAITING` only if // `waiters` lock is held // - number of times `notify_waiters` was called can // be modified only if `waiters` lock is held state: AtomicUsize, waiters: Mutex, } #[derive(Debug)] struct Waiter { /// Intrusive linked-list pointers. pointers: linked_list::Pointers, /// Waiting task's waker. Depending on the value of `notification`, /// this field is either protected by the `waiters` lock in /// `Notify`, or it is exclusively owned by the enclosing `Waiter`. waker: UnsafeCell>, /// Notification for this waiter. Uses 2 bits to store if and how was /// notified, 1 bit for storing if it was woken up using FIFO or LIFO, and /// the rest of it is unused. /// * if it's `None`, then `waker` is protected by the `waiters` lock. /// * if it's `Some`, then `waker` is exclusively owned by the /// enclosing `Waiter` and can be accessed without locking. notification: AtomicNotification, /// Should not be `Unpin`. _p: PhantomPinned, } impl Waiter { fn new() -> Waiter { Waiter { pointers: linked_list::Pointers::new(), waker: UnsafeCell::new(None), notification: AtomicNotification::none(), _p: PhantomPinned, } } } generate_addr_of_methods! { impl<> Waiter { unsafe fn addr_of_pointers(self: NonNull) -> NonNull> { &self.pointers } } } // No notification. const NOTIFICATION_NONE: usize = 0b000; // Notification type used by `notify_one`. const NOTIFICATION_ONE: usize = 0b001; // Notification type used by `notify_last`. const NOTIFICATION_LAST: usize = 0b101; // Notification type used by `notify_waiters`. const NOTIFICATION_ALL: usize = 0b010; /// Notification for a `Waiter`. /// This struct is equivalent to `Option`, but uses /// `AtomicUsize` inside for atomic operations. #[derive(Debug)] struct AtomicNotification(AtomicUsize); impl AtomicNotification { fn none() -> Self { AtomicNotification(AtomicUsize::new(NOTIFICATION_NONE)) } /// Store-release a notification. /// This method should be called exactly once. fn store_release(&self, notification: Notification) { let data: usize = match notification { Notification::All => NOTIFICATION_ALL, Notification::One(NotifyOneStrategy::Fifo) => NOTIFICATION_ONE, Notification::One(NotifyOneStrategy::Lifo) => NOTIFICATION_LAST, }; self.0.store(data, Release); } fn load(&self, ordering: Ordering) -> Option { let data = self.0.load(ordering); match data { NOTIFICATION_NONE => None, NOTIFICATION_ONE => Some(Notification::One(NotifyOneStrategy::Fifo)), NOTIFICATION_LAST => Some(Notification::One(NotifyOneStrategy::Lifo)), NOTIFICATION_ALL => Some(Notification::All), _ => unreachable!(), } } /// Clears the notification. /// This method is used by a `Notified` future to consume the /// notification. It uses relaxed ordering and should be only /// used once the atomic notification is no longer shared. fn clear(&self) { self.0.store(NOTIFICATION_NONE, Relaxed); } } #[derive(Debug, PartialEq, Eq)] #[repr(usize)] enum NotifyOneStrategy { Fifo, Lifo, } #[derive(Debug, PartialEq, Eq)] #[repr(usize)] enum Notification { One(NotifyOneStrategy), All, } /// List used in `Notify::notify_waiters`. It wraps a guarded linked list /// and gates the access to it on `notify.waiters` mutex. It also empties /// the list on drop. struct NotifyWaitersList<'a> { list: GuardedWaitList, is_empty: bool, notify: &'a Notify, } impl<'a> NotifyWaitersList<'a> { fn new( unguarded_list: WaitList, guard: Pin<&'a Waiter>, notify: &'a Notify, ) -> NotifyWaitersList<'a> { let guard_ptr = NonNull::from(guard.get_ref()); let list = unguarded_list.into_guarded(guard_ptr); NotifyWaitersList { list, is_empty: false, notify, } } /// Removes the last element from the guarded list. Modifying this list /// requires an exclusive access to the main list in `Notify`. fn pop_back_locked(&mut self, _waiters: &mut WaitList) -> Option> { let result = self.list.pop_back(); if result.is_none() { // Save information about emptiness to avoid waiting for lock // in the destructor. self.is_empty = true; } result } } impl Drop for NotifyWaitersList<'_> { fn drop(&mut self) { // If the list is not empty, we unlink all waiters from it. // We do not wake the waiters to avoid double panics. if !self.is_empty { let _lock_guard = self.notify.waiters.lock(); while let Some(waiter) = self.list.pop_back() { // Safety: we never make mutable references to waiters. let waiter = unsafe { waiter.as_ref() }; waiter.notification.store_release(Notification::All); } } } } /// Future returned from [`Notify::notified()`]. /// /// This future is fused, so once it has completed, any future calls to poll /// will immediately return `Poll::Ready`. #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Notified<'a> { /// The `Notify` being received on. notify: &'a Notify, /// The current state of the receiving process. state: State, /// Number of calls to `notify_waiters` at the time of creation. notify_waiters_calls: usize, /// Entry in the waiter `LinkedList`. waiter: Waiter, } unsafe impl<'a> Send for Notified<'a> {} unsafe impl<'a> Sync for Notified<'a> {} #[derive(Debug)] enum State { Init, Waiting, Done, } const NOTIFY_WAITERS_SHIFT: usize = 2; const STATE_MASK: usize = (1 << NOTIFY_WAITERS_SHIFT) - 1; const NOTIFY_WAITERS_CALLS_MASK: usize = !STATE_MASK; /// Initial "idle" state. const EMPTY: usize = 0; /// One or more threads are currently waiting to be notified. const WAITING: usize = 1; /// Pending notification. const NOTIFIED: usize = 2; fn set_state(data: usize, state: usize) -> usize { (data & NOTIFY_WAITERS_CALLS_MASK) | (state & STATE_MASK) } fn get_state(data: usize) -> usize { data & STATE_MASK } fn get_num_notify_waiters_calls(data: usize) -> usize { (data & NOTIFY_WAITERS_CALLS_MASK) >> NOTIFY_WAITERS_SHIFT } fn inc_num_notify_waiters_calls(data: usize) -> usize { data + (1 << NOTIFY_WAITERS_SHIFT) } fn atomic_inc_num_notify_waiters_calls(data: &AtomicUsize) { data.fetch_add(1 << NOTIFY_WAITERS_SHIFT, SeqCst); } impl Notify { /// Create a new `Notify`, initialized without a permit. /// /// # Examples /// /// ``` /// use tokio::sync::Notify; /// /// let notify = Notify::new(); /// ``` pub fn new() -> Notify { Notify { state: AtomicUsize::new(0), waiters: Mutex::new(LinkedList::new()), } } /// Create a new `Notify`, initialized without a permit. /// /// When using the `tracing` [unstable feature], a `Notify` created with /// `const_new` will not be instrumented. As such, it will not be visible /// in [`tokio-console`]. Instead, [`Notify::new`] should be used to create /// an instrumented object if that is needed. /// /// # Examples /// /// ``` /// use tokio::sync::Notify; /// /// static NOTIFY: Notify = Notify::const_new(); /// ``` /// /// [`tokio-console`]: https://github.com/tokio-rs/console /// [unstable feature]: crate#unstable-features #[cfg(not(all(loom, test)))] pub const fn const_new() -> Notify { Notify { state: AtomicUsize::new(0), waiters: Mutex::const_new(LinkedList::new()), } } /// Wait for a notification. /// /// Equivalent to: /// /// ```ignore /// async fn notified(&self); /// ``` /// /// Each `Notify` value holds a single permit. If a permit is available from /// an earlier call to [`notify_one()`], then `notified().await` will complete /// immediately, consuming that permit. Otherwise, `notified().await` waits /// for a permit to be made available by the next call to `notify_one()`. /// /// The `Notified` future is not guaranteed to receive wakeups from calls to /// `notify_one()` if it has not yet been polled. See the documentation for /// [`Notified::enable()`] for more details. /// /// The `Notified` future is guaranteed to receive wakeups from /// `notify_waiters()` as soon as it has been created, even if it has not /// yet been polled. /// /// [`notify_one()`]: Notify::notify_one /// [`Notified::enable()`]: Notified::enable /// /// # Cancel safety /// /// This method uses a queue to fairly distribute notifications in the order /// they were requested. Cancelling a call to `notified` makes you lose your /// place in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::Notify; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let notify = Arc::new(Notify::new()); /// let notify2 = notify.clone(); /// /// tokio::spawn(async move { /// notify2.notified().await; /// println!("received notification"); /// }); /// /// println!("sending notification"); /// notify.notify_one(); /// } /// ``` pub fn notified(&self) -> Notified<'_> { // we load the number of times notify_waiters // was called and store that in the future. let state = self.state.load(SeqCst); Notified { notify: self, state: State::Init, notify_waiters_calls: get_num_notify_waiters_calls(state), waiter: Waiter::new(), } } /// Notifies the first waiting task. /// /// If a task is currently waiting, that task is notified. Otherwise, a /// permit is stored in this `Notify` value and the **next** call to /// [`notified().await`] will complete immediately consuming the permit made /// available by this call to `notify_one()`. /// /// At most one permit may be stored by `Notify`. Many sequential calls to /// `notify_one` will result in a single permit being stored. The next call to /// `notified().await` will complete immediately, but the one after that /// will wait. /// /// [`notified().await`]: Notify::notified() /// /// # Examples /// /// ``` /// use tokio::sync::Notify; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let notify = Arc::new(Notify::new()); /// let notify2 = notify.clone(); /// /// tokio::spawn(async move { /// notify2.notified().await; /// println!("received notification"); /// }); /// /// println!("sending notification"); /// notify.notify_one(); /// } /// ``` // Alias for old name in 0.x #[cfg_attr(docsrs, doc(alias = "notify"))] pub fn notify_one(&self) { self.notify_with_strategy(NotifyOneStrategy::Fifo); } /// Notifies the last waiting task. /// /// This function behaves similar to `notify_one`. The only difference is that it wakes /// the most recently added waiter instead of the oldest waiter. /// /// Check the [`notify_one()`] documentation for more info and /// examples. /// /// [`notify_one()`]: Notify::notify_one pub fn notify_last(&self) { self.notify_with_strategy(NotifyOneStrategy::Lifo); } fn notify_with_strategy(&self, strategy: NotifyOneStrategy) { // Load the current state let mut curr = self.state.load(SeqCst); // If the state is `EMPTY`, transition to `NOTIFIED` and return. while let EMPTY | NOTIFIED = get_state(curr) { // The compare-exchange from `NOTIFIED` -> `NOTIFIED` is intended. A // happens-before synchronization must happen between this atomic // operation and a task calling `notified().await`. let new = set_state(curr, NOTIFIED); let res = self.state.compare_exchange(curr, new, SeqCst, SeqCst); match res { // No waiters, no further work to do Ok(_) => return, Err(actual) => { curr = actual; } } } // There are waiters, the lock must be acquired to notify. let mut waiters = self.waiters.lock(); // The state must be reloaded while the lock is held. The state may only // transition out of WAITING while the lock is held. curr = self.state.load(SeqCst); if let Some(waker) = notify_locked(&mut waiters, &self.state, curr, strategy) { drop(waiters); waker.wake(); } } /// Notifies all waiting tasks. /// /// If a task is currently waiting, that task is notified. Unlike with /// `notify_one()`, no permit is stored to be used by the next call to /// `notified().await`. The purpose of this method is to notify all /// already registered waiters. Registering for notification is done by /// acquiring an instance of the `Notified` future via calling `notified()`. /// /// # Examples /// /// ``` /// use tokio::sync::Notify; /// use std::sync::Arc; /// /// #[tokio::main] /// async fn main() { /// let notify = Arc::new(Notify::new()); /// let notify2 = notify.clone(); /// /// let notified1 = notify.notified(); /// let notified2 = notify.notified(); /// /// let handle = tokio::spawn(async move { /// println!("sending notifications"); /// notify2.notify_waiters(); /// }); /// /// notified1.await; /// notified2.await; /// println!("received notifications"); /// } /// ``` pub fn notify_waiters(&self) { let mut waiters = self.waiters.lock(); // The state must be loaded while the lock is held. The state may only // transition out of WAITING while the lock is held. let curr = self.state.load(SeqCst); if matches!(get_state(curr), EMPTY | NOTIFIED) { // There are no waiting tasks. All we need to do is increment the // number of times this method was called. atomic_inc_num_notify_waiters_calls(&self.state); return; } // Increment the number of times this method was called // and transition to empty. let new_state = set_state(inc_num_notify_waiters_calls(curr), EMPTY); self.state.store(new_state, SeqCst); // It is critical for `GuardedLinkedList` safety that the guard node is // pinned in memory and is not dropped until the guarded list is dropped. let guard = Waiter::new(); pin!(guard); // We move all waiters to a secondary list. It uses a `GuardedLinkedList` // underneath to allow every waiter to safely remove itself from it. // // * This list will be still guarded by the `waiters` lock. // `NotifyWaitersList` wrapper makes sure we hold the lock to modify it. // * This wrapper will empty the list on drop. It is critical for safety // that we will not leave any list entry with a pointer to the local // guard node after this function returns / panics. let mut list = NotifyWaitersList::new(std::mem::take(&mut *waiters), guard.as_ref(), self); let mut wakers = WakeList::new(); 'outer: loop { while wakers.can_push() { match list.pop_back_locked(&mut waiters) { Some(waiter) => { // Safety: we never make mutable references to waiters. let waiter = unsafe { waiter.as_ref() }; // Safety: we hold the lock, so we can access the waker. if let Some(waker) = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) } { wakers.push(waker); } // This waiter is unlinked and will not be shared ever again, release it. waiter.notification.store_release(Notification::All); } None => { break 'outer; } } } // Release the lock before notifying. drop(waiters); // One of the wakers may panic, but the remaining waiters will still // be unlinked from the list in `NotifyWaitersList` destructor. wakers.wake_all(); // Acquire the lock again. waiters = self.waiters.lock(); } // Release the lock before notifying drop(waiters); wakers.wake_all(); } } impl Default for Notify { fn default() -> Notify { Notify::new() } } impl UnwindSafe for Notify {} impl RefUnwindSafe for Notify {} fn notify_locked( waiters: &mut WaitList, state: &AtomicUsize, curr: usize, strategy: NotifyOneStrategy, ) -> Option { match get_state(curr) { EMPTY | NOTIFIED => { let res = state.compare_exchange(curr, set_state(curr, NOTIFIED), SeqCst, SeqCst); match res { Ok(_) => None, Err(actual) => { let actual_state = get_state(actual); assert!(actual_state == EMPTY || actual_state == NOTIFIED); state.store(set_state(actual, NOTIFIED), SeqCst); None } } } WAITING => { // At this point, it is guaranteed that the state will not // concurrently change as holding the lock is required to // transition **out** of `WAITING`. // // Get a pending waiter using one of the available dequeue strategies. let waiter = match strategy { NotifyOneStrategy::Fifo => waiters.pop_back().unwrap(), NotifyOneStrategy::Lifo => waiters.pop_front().unwrap(), }; // Safety: we never make mutable references to waiters. let waiter = unsafe { waiter.as_ref() }; // Safety: we hold the lock, so we can access the waker. let waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; // This waiter is unlinked and will not be shared ever again, release it. waiter .notification .store_release(Notification::One(strategy)); if waiters.is_empty() { // As this the **final** waiter in the list, the state // must be transitioned to `EMPTY`. As transitioning // **from** `WAITING` requires the lock to be held, a // `store` is sufficient. state.store(set_state(curr, EMPTY), SeqCst); } waker } _ => unreachable!(), } } // ===== impl Notified ===== impl Notified<'_> { /// Adds this future to the list of futures that are ready to receive /// wakeups from calls to [`notify_one`]. /// /// Polling the future also adds it to the list, so this method should only /// be used if you want to add the future to the list before the first call /// to `poll`. (In fact, this method is equivalent to calling `poll` except /// that no `Waker` is registered.) /// /// This has no effect on notifications sent using [`notify_waiters`], which /// are received as long as they happen after the creation of the `Notified` /// regardless of whether `enable` or `poll` has been called. /// /// This method returns true if the `Notified` is ready. This happens in the /// following situations: /// /// 1. The `notify_waiters` method was called between the creation of the /// `Notified` and the call to this method. /// 2. This is the first call to `enable` or `poll` on this future, and the /// `Notify` was holding a permit from a previous call to `notify_one`. /// The call consumes the permit in that case. /// 3. The future has previously been enabled or polled, and it has since /// then been marked ready by either consuming a permit from the /// `Notify`, or by a call to `notify_one` or `notify_waiters` that /// removed it from the list of futures ready to receive wakeups. /// /// If this method returns true, any future calls to poll on the same future /// will immediately return `Poll::Ready`. /// /// # Examples /// /// Unbound multi-producer multi-consumer (mpmc) channel. /// /// The call to `enable` is important because otherwise if you have two /// calls to `recv` and two calls to `send` in parallel, the following could /// happen: /// /// 1. Both calls to `try_recv` return `None`. /// 2. Both new elements are added to the vector. /// 3. The `notify_one` method is called twice, adding only a single /// permit to the `Notify`. /// 4. Both calls to `recv` reach the `Notified` future. One of them /// consumes the permit, and the other sleeps forever. /// /// By adding the `Notified` futures to the list by calling `enable` before /// `try_recv`, the `notify_one` calls in step three would remove the /// futures from the list and mark them notified instead of adding a permit /// to the `Notify`. This ensures that both futures are woken. /// /// ``` /// use tokio::sync::Notify; /// /// use std::collections::VecDeque; /// use std::sync::Mutex; /// /// struct Channel { /// messages: Mutex>, /// notify_on_sent: Notify, /// } /// /// impl Channel { /// pub fn send(&self, msg: T) { /// let mut locked_queue = self.messages.lock().unwrap(); /// locked_queue.push_back(msg); /// drop(locked_queue); /// /// // Send a notification to one of the calls currently /// // waiting in a call to `recv`. /// self.notify_on_sent.notify_one(); /// } /// /// pub fn try_recv(&self) -> Option { /// let mut locked_queue = self.messages.lock().unwrap(); /// locked_queue.pop_front() /// } /// /// pub async fn recv(&self) -> T { /// let future = self.notify_on_sent.notified(); /// tokio::pin!(future); /// /// loop { /// // Make sure that no wakeup is lost if we get /// // `None` from `try_recv`. /// future.as_mut().enable(); /// /// if let Some(msg) = self.try_recv() { /// return msg; /// } /// /// // Wait for a call to `notify_one`. /// // /// // This uses `.as_mut()` to avoid consuming the future, /// // which lets us call `Pin::set` below. /// future.as_mut().await; /// /// // Reset the future in case another call to /// // `try_recv` got the message before us. /// future.set(self.notify_on_sent.notified()); /// } /// } /// } /// ``` /// /// [`notify_one`]: Notify::notify_one() /// [`notify_waiters`]: Notify::notify_waiters() pub fn enable(self: Pin<&mut Self>) -> bool { self.poll_notified(None).is_ready() } /// A custom `project` implementation is used in place of `pin-project-lite` /// as a custom drop implementation is needed. fn project(self: Pin<&mut Self>) -> (&Notify, &mut State, &usize, &Waiter) { unsafe { // Safety: `notify`, `state` and `notify_waiters_calls` are `Unpin`. is_unpin::<&Notify>(); is_unpin::(); is_unpin::(); let me = self.get_unchecked_mut(); ( me.notify, &mut me.state, &me.notify_waiters_calls, &me.waiter, ) } } fn poll_notified(self: Pin<&mut Self>, waker: Option<&Waker>) -> Poll<()> { let (notify, state, notify_waiters_calls, waiter) = self.project(); 'outer_loop: loop { match *state { State::Init => { let curr = notify.state.load(SeqCst); // Optimistically try acquiring a pending notification let res = notify.state.compare_exchange( set_state(curr, NOTIFIED), set_state(curr, EMPTY), SeqCst, SeqCst, ); if res.is_ok() { // Acquired the notification *state = State::Done; continue 'outer_loop; } // Clone the waker before locking, a waker clone can be // triggering arbitrary code. let waker = waker.cloned(); // Acquire the lock and attempt to transition to the waiting // state. let mut waiters = notify.waiters.lock(); // Reload the state with the lock held let mut curr = notify.state.load(SeqCst); // if notify_waiters has been called after the future // was created, then we are done if get_num_notify_waiters_calls(curr) != *notify_waiters_calls { *state = State::Done; continue 'outer_loop; } // Transition the state to WAITING. loop { match get_state(curr) { EMPTY => { // Transition to WAITING let res = notify.state.compare_exchange( set_state(curr, EMPTY), set_state(curr, WAITING), SeqCst, SeqCst, ); if let Err(actual) = res { assert_eq!(get_state(actual), NOTIFIED); curr = actual; } else { break; } } WAITING => break, NOTIFIED => { // Try consuming the notification let res = notify.state.compare_exchange( set_state(curr, NOTIFIED), set_state(curr, EMPTY), SeqCst, SeqCst, ); match res { Ok(_) => { // Acquired the notification *state = State::Done; continue 'outer_loop; } Err(actual) => { assert_eq!(get_state(actual), EMPTY); curr = actual; } } } _ => unreachable!(), } } let mut old_waker = None; if waker.is_some() { // Safety: called while locked. // // The use of `old_waiter` here is not necessary, as the field is always // None when we reach this line. unsafe { old_waker = waiter.waker.with_mut(|v| std::mem::replace(&mut *v, waker)); } } // Insert the waiter into the linked list waiters.push_front(NonNull::from(waiter)); *state = State::Waiting; drop(waiters); drop(old_waker); return Poll::Pending; } State::Waiting => { #[cfg(tokio_taskdump)] if let Some(waker) = waker { let mut ctx = Context::from_waker(waker); std::task::ready!(crate::trace::trace_leaf(&mut ctx)); } if waiter.notification.load(Acquire).is_some() { // Safety: waiter is already unlinked and will not be shared again, // so we have an exclusive access to `waker`. drop(unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }); waiter.notification.clear(); *state = State::Done; return Poll::Ready(()); } // Our waiter was not notified, implying it is still stored in a waiter // list (guarded by `notify.waiters`). In order to access the waker // fields, we must acquire the lock. let mut old_waker = None; let mut waiters = notify.waiters.lock(); // We hold the lock and notifications are set only with the lock held, // so this can be relaxed, because the happens-before relationship is // established through the mutex. if waiter.notification.load(Relaxed).is_some() { // Safety: waiter is already unlinked and will not be shared again, // so we have an exclusive access to `waker`. old_waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; waiter.notification.clear(); // Drop the old waker after releasing the lock. drop(waiters); drop(old_waker); *state = State::Done; return Poll::Ready(()); } // Load the state with the lock held. let curr = notify.state.load(SeqCst); if get_num_notify_waiters_calls(curr) != *notify_waiters_calls { // Before we add a waiter to the list we check if these numbers are // different while holding the lock. If these numbers are different now, // it means that there is a call to `notify_waiters` in progress and this // waiter must be contained by a guarded list used in `notify_waiters`. // We can treat the waiter as notified and remove it from the list, as // it would have been notified in the `notify_waiters` call anyways. // Safety: we hold the lock, so we can modify the waker. old_waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; // Safety: we hold the lock, so we have an exclusive access to the list. // The list is used in `notify_waiters`, so it must be guarded. unsafe { waiters.remove(NonNull::from(waiter)) }; *state = State::Done; } else { // Safety: we hold the lock, so we can modify the waker. unsafe { waiter.waker.with_mut(|v| { if let Some(waker) = waker { let should_update = match &*v { Some(current_waker) => !current_waker.will_wake(waker), None => true, }; if should_update { old_waker = std::mem::replace(&mut *v, Some(waker.clone())); } } }); } // Drop the old waker after releasing the lock. drop(waiters); drop(old_waker); return Poll::Pending; } // Explicit drop of the lock to indicate the scope that the // lock is held. Because holding the lock is required to // ensure safe access to fields not held within the lock, it // is helpful to visualize the scope of the critical // section. drop(waiters); // Drop the old waker after releasing the lock. drop(old_waker); } State::Done => { #[cfg(tokio_taskdump)] if let Some(waker) = waker { let mut ctx = Context::from_waker(waker); std::task::ready!(crate::trace::trace_leaf(&mut ctx)); } return Poll::Ready(()); } } } } } impl Future for Notified<'_> { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { self.poll_notified(Some(cx.waker())) } } impl Drop for Notified<'_> { fn drop(&mut self) { // Safety: The type only transitions to a "Waiting" state when pinned. let (notify, state, _, waiter) = unsafe { Pin::new_unchecked(self).project() }; // This is where we ensure safety. The `Notified` value is being // dropped, which means we must ensure that the waiter entry is no // longer stored in the linked list. if matches!(*state, State::Waiting) { let mut waiters = notify.waiters.lock(); let mut notify_state = notify.state.load(SeqCst); // We hold the lock, so this field is not concurrently accessed by // `notify_*` functions and we can use the relaxed ordering. let notification = waiter.notification.load(Relaxed); // remove the entry from the list (if not already removed) // // Safety: we hold the lock, so we have an exclusive access to every list the // waiter may be contained in. If the node is not contained in the `waiters` // list, then it is contained by a guarded list used by `notify_waiters`. unsafe { waiters.remove(NonNull::from(waiter)) }; if waiters.is_empty() && get_state(notify_state) == WAITING { notify_state = set_state(notify_state, EMPTY); notify.state.store(notify_state, SeqCst); } // See if the node was notified but not received. In this case, if // the notification was triggered via `notify_one`, it must be sent // to the next waiter. if let Some(Notification::One(strategy)) = notification { if let Some(waker) = notify_locked(&mut waiters, ¬ify.state, notify_state, strategy) { drop(waiters); waker.wake(); } } } } } /// # Safety /// /// `Waiter` is forced to be !Unpin. unsafe impl linked_list::Link for Waiter { type Handle = NonNull; type Target = Waiter; fn as_raw(handle: &NonNull) -> NonNull { *handle } unsafe fn from_raw(ptr: NonNull) -> NonNull { ptr } unsafe fn pointers(target: NonNull) -> NonNull> { Waiter::addr_of_pointers(target) } } fn is_unpin() {} tokio-1.43.1/src/sync/once_cell.rs000064400000000000000000000422131046102023000150750ustar 00000000000000use super::{Semaphore, SemaphorePermit, TryAcquireError}; use crate::loom::cell::UnsafeCell; use std::error::Error; use std::fmt; use std::future::Future; use std::mem::MaybeUninit; use std::ops::Drop; use std::ptr; use std::sync::atomic::{AtomicBool, Ordering}; // This file contains an implementation of an OnceCell. The principle // behind the safety of the cell is that any thread with an `&OnceCell` may // access the `value` field according the following rules: // // 1. When `value_set` is false, the `value` field may be modified by the // thread holding the permit on the semaphore. // 2. When `value_set` is true, the `value` field may be accessed immutably by // any thread. // // It is an invariant that if the semaphore is closed, then `value_set` is true. // The reverse does not necessarily hold — but if not, the semaphore may not // have any available permits. // // A thread with a `&mut OnceCell` may modify the value in any way it wants as // long as the invariants are upheld. /// A thread-safe cell that can be written to only once. /// /// A `OnceCell` is typically used for global variables that need to be /// initialized once on first use, but need no further changes. The `OnceCell` /// in Tokio allows the initialization procedure to be asynchronous. /// /// # Examples /// /// ``` /// use tokio::sync::OnceCell; /// /// async fn some_computation() -> u32 { /// 1 + 1 /// } /// /// static ONCE: OnceCell = OnceCell::const_new(); /// /// #[tokio::main] /// async fn main() { /// let result = ONCE.get_or_init(some_computation).await; /// assert_eq!(*result, 2); /// } /// ``` /// /// It is often useful to write a wrapper method for accessing the value. /// /// ``` /// use tokio::sync::OnceCell; /// /// static ONCE: OnceCell = OnceCell::const_new(); /// /// async fn get_global_integer() -> &'static u32 { /// ONCE.get_or_init(|| async { /// 1 + 1 /// }).await /// } /// /// #[tokio::main] /// async fn main() { /// let result = get_global_integer().await; /// assert_eq!(*result, 2); /// } /// ``` pub struct OnceCell { value_set: AtomicBool, value: UnsafeCell>, semaphore: Semaphore, } impl Default for OnceCell { fn default() -> OnceCell { OnceCell::new() } } impl fmt::Debug for OnceCell { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("OnceCell") .field("value", &self.get()) .finish() } } impl Clone for OnceCell { fn clone(&self) -> OnceCell { OnceCell::new_with(self.get().cloned()) } } impl PartialEq for OnceCell { fn eq(&self, other: &OnceCell) -> bool { self.get() == other.get() } } impl Eq for OnceCell {} impl Drop for OnceCell { fn drop(&mut self) { if self.initialized_mut() { unsafe { self.value .with_mut(|ptr| ptr::drop_in_place((*ptr).as_mut_ptr())); }; } } } impl From for OnceCell { fn from(value: T) -> Self { OnceCell { value_set: AtomicBool::new(true), value: UnsafeCell::new(MaybeUninit::new(value)), semaphore: Semaphore::new_closed(), } } } impl OnceCell { /// Creates a new empty `OnceCell` instance. pub fn new() -> Self { OnceCell { value_set: AtomicBool::new(false), value: UnsafeCell::new(MaybeUninit::uninit()), semaphore: Semaphore::new(1), } } /// Creates a new empty `OnceCell` instance. /// /// Equivalent to `OnceCell::new`, except that it can be used in static /// variables. /// /// When using the `tracing` [unstable feature], a `OnceCell` created with /// `const_new` will not be instrumented. As such, it will not be visible /// in [`tokio-console`]. Instead, [`OnceCell::new`] should be used to /// create an instrumented object if that is needed. /// /// # Example /// /// ``` /// use tokio::sync::OnceCell; /// /// static ONCE: OnceCell = OnceCell::const_new(); /// /// async fn get_global_integer() -> &'static u32 { /// ONCE.get_or_init(|| async { /// 1 + 1 /// }).await /// } /// /// #[tokio::main] /// async fn main() { /// let result = get_global_integer().await; /// assert_eq!(*result, 2); /// } /// ``` /// /// [`tokio-console`]: https://github.com/tokio-rs/console /// [unstable feature]: crate#unstable-features #[cfg(not(all(loom, test)))] pub const fn const_new() -> Self { OnceCell { value_set: AtomicBool::new(false), value: UnsafeCell::new(MaybeUninit::uninit()), semaphore: Semaphore::const_new(1), } } /// Creates a new `OnceCell` that contains the provided value, if any. /// /// If the `Option` is `None`, this is equivalent to `OnceCell::new`. /// /// [`OnceCell::new`]: crate::sync::OnceCell::new // Once https://github.com/rust-lang/rust/issues/73255 lands // and tokio MSRV is bumped to the rustc version with it stabilised, // we can make this function available in const context, // by creating `Semaphore::const_new_closed`. pub fn new_with(value: Option) -> Self { if let Some(v) = value { OnceCell::from(v) } else { OnceCell::new() } } /// Creates a new `OnceCell` that contains the provided value. /// /// # Example /// /// When using the `tracing` [unstable feature], a `OnceCell` created with /// `const_new_with` will not be instrumented. As such, it will not be /// visible in [`tokio-console`]. Instead, [`OnceCell::new_with`] should be /// used to create an instrumented object if that is needed. /// /// ``` /// use tokio::sync::OnceCell; /// /// static ONCE: OnceCell = OnceCell::const_new_with(1); /// /// async fn get_global_integer() -> &'static u32 { /// ONCE.get_or_init(|| async { /// 1 + 1 /// }).await /// } /// /// #[tokio::main] /// async fn main() { /// let result = get_global_integer().await; /// assert_eq!(*result, 1); /// } /// ``` /// /// [`tokio-console`]: https://github.com/tokio-rs/console /// [unstable feature]: crate#unstable-features #[cfg(not(all(loom, test)))] pub const fn const_new_with(value: T) -> Self { OnceCell { value_set: AtomicBool::new(true), value: UnsafeCell::new(MaybeUninit::new(value)), semaphore: Semaphore::const_new_closed(), } } /// Returns `true` if the `OnceCell` currently contains a value, and `false` /// otherwise. pub fn initialized(&self) -> bool { // Using acquire ordering so any threads that read a true from this // atomic is able to read the value. self.value_set.load(Ordering::Acquire) } /// Returns `true` if the `OnceCell` currently contains a value, and `false` /// otherwise. fn initialized_mut(&mut self) -> bool { *self.value_set.get_mut() } // SAFETY: The OnceCell must not be empty. unsafe fn get_unchecked(&self) -> &T { &*self.value.with(|ptr| (*ptr).as_ptr()) } // SAFETY: The OnceCell must not be empty. unsafe fn get_unchecked_mut(&mut self) -> &mut T { &mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr()) } fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T { // SAFETY: We are holding the only permit on the semaphore. unsafe { self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); } // Using release ordering so any threads that read a true from this // atomic is able to read the value we just stored. self.value_set.store(true, Ordering::Release); self.semaphore.close(); permit.forget(); // SAFETY: We just initialized the cell. unsafe { self.get_unchecked() } } /// Returns a reference to the value currently stored in the `OnceCell`, or /// `None` if the `OnceCell` is empty. pub fn get(&self) -> Option<&T> { if self.initialized() { Some(unsafe { self.get_unchecked() }) } else { None } } /// Returns a mutable reference to the value currently stored in the /// `OnceCell`, or `None` if the `OnceCell` is empty. /// /// Since this call borrows the `OnceCell` mutably, it is safe to mutate the /// value inside the `OnceCell` — the mutable borrow statically guarantees /// no other references exist. pub fn get_mut(&mut self) -> Option<&mut T> { if self.initialized_mut() { Some(unsafe { self.get_unchecked_mut() }) } else { None } } /// Sets the value of the `OnceCell` to the given value if the `OnceCell` is /// empty. /// /// If the `OnceCell` already has a value, this call will fail with an /// [`SetError::AlreadyInitializedError`]. /// /// If the `OnceCell` is empty, but some other task is currently trying to /// set the value, this call will fail with [`SetError::InitializingError`]. /// /// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError /// [`SetError::InitializingError`]: crate::sync::SetError::InitializingError pub fn set(&self, value: T) -> Result<(), SetError> { if self.initialized() { return Err(SetError::AlreadyInitializedError(value)); } // Another task might be initializing the cell, in which case // `try_acquire` will return an error. If we succeed to acquire the // permit, then we can set the value. match self.semaphore.try_acquire() { Ok(permit) => { debug_assert!(!self.initialized()); self.set_value(value, permit); Ok(()) } Err(TryAcquireError::NoPermits) => { // Some other task is holding the permit. That task is // currently trying to initialize the value. Err(SetError::InitializingError(value)) } Err(TryAcquireError::Closed) => { // The semaphore was closed. Some other task has initialized // the value. Err(SetError::AlreadyInitializedError(value)) } } } /// Gets the value currently in the `OnceCell`, or initialize it with the /// given asynchronous operation. /// /// If some other task is currently working on initializing the `OnceCell`, /// this call will wait for that other task to finish, then return the value /// that the other task produced. /// /// If the provided operation is cancelled or panics, the initialization /// attempt is cancelled. If there are other tasks waiting for the value to /// be initialized, one of them will start another attempt at initializing /// the value. /// /// This will deadlock if `f` tries to initialize the cell recursively. pub async fn get_or_init(&self, f: F) -> &T where F: FnOnce() -> Fut, Fut: Future, { crate::trace::async_trace_leaf().await; if self.initialized() { // SAFETY: The OnceCell has been fully initialized. unsafe { self.get_unchecked() } } else { // Here we try to acquire the semaphore permit. Holding the permit // will allow us to set the value of the OnceCell, and prevents // other tasks from initializing the OnceCell while we are holding // it. match self.semaphore.acquire().await { Ok(permit) => { debug_assert!(!self.initialized()); // If `f()` panics or `select!` is called, this // `get_or_init` call is aborted and the semaphore permit is // dropped. let value = f().await; self.set_value(value, permit) } Err(_) => { debug_assert!(self.initialized()); // SAFETY: The semaphore has been closed. This only happens // when the OnceCell is fully initialized. unsafe { self.get_unchecked() } } } } } /// Gets the value currently in the `OnceCell`, or initialize it with the /// given asynchronous operation. /// /// If some other task is currently working on initializing the `OnceCell`, /// this call will wait for that other task to finish, then return the value /// that the other task produced. /// /// If the provided operation returns an error, is cancelled or panics, the /// initialization attempt is cancelled. If there are other tasks waiting /// for the value to be initialized, one of them will start another attempt /// at initializing the value. /// /// This will deadlock if `f` tries to initialize the cell recursively. pub async fn get_or_try_init(&self, f: F) -> Result<&T, E> where F: FnOnce() -> Fut, Fut: Future>, { crate::trace::async_trace_leaf().await; if self.initialized() { // SAFETY: The OnceCell has been fully initialized. unsafe { Ok(self.get_unchecked()) } } else { // Here we try to acquire the semaphore permit. Holding the permit // will allow us to set the value of the OnceCell, and prevents // other tasks from initializing the OnceCell while we are holding // it. match self.semaphore.acquire().await { Ok(permit) => { debug_assert!(!self.initialized()); // If `f()` panics or `select!` is called, this // `get_or_try_init` call is aborted and the semaphore // permit is dropped. let value = f().await; match value { Ok(value) => Ok(self.set_value(value, permit)), Err(e) => Err(e), } } Err(_) => { debug_assert!(self.initialized()); // SAFETY: The semaphore has been closed. This only happens // when the OnceCell is fully initialized. unsafe { Ok(self.get_unchecked()) } } } } } /// Takes the value from the cell, destroying the cell in the process. /// Returns `None` if the cell is empty. pub fn into_inner(mut self) -> Option { if self.initialized_mut() { // Set to uninitialized for the destructor of `OnceCell` to work properly *self.value_set.get_mut() = false; Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) }) } else { None } } /// Takes ownership of the current value, leaving the cell empty. Returns /// `None` if the cell is empty. pub fn take(&mut self) -> Option { std::mem::take(self).into_inner() } } // Since `get` gives us access to immutable references of the OnceCell, OnceCell // can only be Sync if T is Sync, otherwise OnceCell would allow sharing // references of !Sync values across threads. We need T to be Send in order for // OnceCell to by Sync because we can use `set` on `&OnceCell` to send values // (of type T) across threads. unsafe impl Sync for OnceCell {} // Access to OnceCell's value is guarded by the semaphore permit // and atomic operations on `value_set`, so as long as T itself is Send // it's safe to send it to another thread unsafe impl Send for OnceCell {} /// Errors that can be returned from [`OnceCell::set`]. /// /// [`OnceCell::set`]: crate::sync::OnceCell::set #[derive(Debug, PartialEq, Eq)] pub enum SetError { /// The cell was already initialized when [`OnceCell::set`] was called. /// /// [`OnceCell::set`]: crate::sync::OnceCell::set AlreadyInitializedError(T), /// The cell is currently being initialized. InitializingError(T), } impl fmt::Display for SetError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { SetError::AlreadyInitializedError(_) => write!(f, "AlreadyInitializedError"), SetError::InitializingError(_) => write!(f, "InitializingError"), } } } impl Error for SetError {} impl SetError { /// Whether `SetError` is `SetError::AlreadyInitializedError`. pub fn is_already_init_err(&self) -> bool { match self { SetError::AlreadyInitializedError(_) => true, SetError::InitializingError(_) => false, } } /// Whether `SetError` is `SetError::InitializingError` pub fn is_initializing_err(&self) -> bool { match self { SetError::AlreadyInitializedError(_) => false, SetError::InitializingError(_) => true, } } } tokio-1.43.1/src/sync/oneshot.rs000064400000000000000000001243171046102023000146370ustar 00000000000000#![cfg_attr(not(feature = "sync"), allow(dead_code, unreachable_pub))] //! A one-shot channel is used for sending a single message between //! asynchronous tasks. The [`channel`] function is used to create a //! [`Sender`] and [`Receiver`] handle pair that form the channel. //! //! The `Sender` handle is used by the producer to send the value. //! The `Receiver` handle is used by the consumer to receive the value. //! //! Each handle can be used on separate tasks. //! //! Since the `send` method is not async, it can be used anywhere. This includes //! sending between two runtimes, and using it from non-async code. //! //! If the [`Receiver`] is closed before receiving a message which has already //! been sent, the message will remain in the channel until the receiver is //! dropped, at which point the message will be dropped immediately. //! //! # Examples //! //! ``` //! use tokio::sync::oneshot; //! //! #[tokio::main] //! async fn main() { //! let (tx, rx) = oneshot::channel(); //! //! tokio::spawn(async move { //! if let Err(_) = tx.send(3) { //! println!("the receiver dropped"); //! } //! }); //! //! match rx.await { //! Ok(v) => println!("got = {:?}", v), //! Err(_) => println!("the sender dropped"), //! } //! } //! ``` //! //! If the sender is dropped without sending, the receiver will fail with //! [`error::RecvError`]: //! //! ``` //! use tokio::sync::oneshot; //! //! #[tokio::main] //! async fn main() { //! let (tx, rx) = oneshot::channel::(); //! //! tokio::spawn(async move { //! drop(tx); //! }); //! //! match rx.await { //! Ok(_) => panic!("This doesn't happen"), //! Err(_) => println!("the sender dropped"), //! } //! } //! ``` //! //! To use a `oneshot` channel in a `tokio::select!` loop, add `&mut` in front of //! the channel. //! //! ``` //! use tokio::sync::oneshot; //! use tokio::time::{interval, sleep, Duration}; //! //! #[tokio::main] //! # async fn _doc() {} //! # #[tokio::main(flavor = "current_thread", start_paused = true)] //! async fn main() { //! let (send, mut recv) = oneshot::channel(); //! let mut interval = interval(Duration::from_millis(100)); //! //! # let handle = //! tokio::spawn(async move { //! sleep(Duration::from_secs(1)).await; //! send.send("shut down").unwrap(); //! }); //! //! loop { //! tokio::select! { //! _ = interval.tick() => println!("Another 100ms"), //! msg = &mut recv => { //! println!("Got message: {}", msg.unwrap()); //! break; //! } //! } //! } //! # handle.await.unwrap(); //! } //! ``` //! //! To use a `Sender` from a destructor, put it in an [`Option`] and call //! [`Option::take`]. //! //! ``` //! use tokio::sync::oneshot; //! //! struct SendOnDrop { //! sender: Option>, //! } //! impl Drop for SendOnDrop { //! fn drop(&mut self) { //! if let Some(sender) = self.sender.take() { //! // Using `let _ =` to ignore send errors. //! let _ = sender.send("I got dropped!"); //! } //! } //! } //! //! #[tokio::main] //! # async fn _doc() {} //! # #[tokio::main(flavor = "current_thread")] //! async fn main() { //! let (send, recv) = oneshot::channel(); //! //! let send_on_drop = SendOnDrop { sender: Some(send) }; //! drop(send_on_drop); //! //! assert_eq!(recv.await, Ok("I got dropped!")); //! } //! ``` use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Arc; #[cfg(all(tokio_unstable, feature = "tracing"))] use crate::util::trace; use std::fmt; use std::future::Future; use std::mem::MaybeUninit; use std::pin::Pin; use std::sync::atomic::Ordering::{self, AcqRel, Acquire}; use std::task::Poll::{Pending, Ready}; use std::task::{ready, Context, Poll, Waker}; /// Sends a value to the associated [`Receiver`]. /// /// A pair of both a [`Sender`] and a [`Receiver`] are created by the /// [`channel`](fn@channel) function. /// /// # Examples /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel(); /// /// tokio::spawn(async move { /// if let Err(_) = tx.send(3) { /// println!("the receiver dropped"); /// } /// }); /// /// match rx.await { /// Ok(v) => println!("got = {:?}", v), /// Err(_) => println!("the sender dropped"), /// } /// } /// ``` /// /// If the sender is dropped without sending, the receiver will fail with /// [`error::RecvError`]: /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel::(); /// /// tokio::spawn(async move { /// drop(tx); /// }); /// /// match rx.await { /// Ok(_) => panic!("This doesn't happen"), /// Err(_) => println!("the sender dropped"), /// } /// } /// ``` /// /// To use a `Sender` from a destructor, put it in an [`Option`] and call /// [`Option::take`]. /// /// ``` /// use tokio::sync::oneshot; /// /// struct SendOnDrop { /// sender: Option>, /// } /// impl Drop for SendOnDrop { /// fn drop(&mut self) { /// if let Some(sender) = self.sender.take() { /// // Using `let _ =` to ignore send errors. /// let _ = sender.send("I got dropped!"); /// } /// } /// } /// /// #[tokio::main] /// # async fn _doc() {} /// # #[tokio::main(flavor = "current_thread")] /// async fn main() { /// let (send, recv) = oneshot::channel(); /// /// let send_on_drop = SendOnDrop { sender: Some(send) }; /// drop(send_on_drop); /// /// assert_eq!(recv.await, Ok("I got dropped!")); /// } /// ``` /// /// [`Option`]: std::option::Option /// [`Option::take`]: std::option::Option::take #[derive(Debug)] pub struct Sender { inner: Option>>, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, } /// Receives a value from the associated [`Sender`]. /// /// A pair of both a [`Sender`] and a [`Receiver`] are created by the /// [`channel`](fn@channel) function. /// /// This channel has no `recv` method because the receiver itself implements the /// [`Future`] trait. To receive a `Result`, `.await` the `Receiver` object directly. /// /// The `poll` method on the `Future` trait is allowed to spuriously return /// `Poll::Pending` even if the message has been sent. If such a spurious /// failure happens, then the caller will be woken when the spurious failure has /// been resolved so that the caller can attempt to receive the message again. /// Note that receiving such a wakeup does not guarantee that the next call will /// succeed — it could fail with another spurious failure. (A spurious failure /// does not mean that the message is lost. It is just delayed.) /// /// [`Future`]: trait@std::future::Future /// /// # Examples /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel(); /// /// tokio::spawn(async move { /// if let Err(_) = tx.send(3) { /// println!("the receiver dropped"); /// } /// }); /// /// match rx.await { /// Ok(v) => println!("got = {:?}", v), /// Err(_) => println!("the sender dropped"), /// } /// } /// ``` /// /// If the sender is dropped without sending, the receiver will fail with /// [`error::RecvError`]: /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel::(); /// /// tokio::spawn(async move { /// drop(tx); /// }); /// /// match rx.await { /// Ok(_) => panic!("This doesn't happen"), /// Err(_) => println!("the sender dropped"), /// } /// } /// ``` /// /// To use a `Receiver` in a `tokio::select!` loop, add `&mut` in front of the /// channel. /// /// ``` /// use tokio::sync::oneshot; /// use tokio::time::{interval, sleep, Duration}; /// /// #[tokio::main] /// # async fn _doc() {} /// # #[tokio::main(flavor = "current_thread", start_paused = true)] /// async fn main() { /// let (send, mut recv) = oneshot::channel(); /// let mut interval = interval(Duration::from_millis(100)); /// /// # let handle = /// tokio::spawn(async move { /// sleep(Duration::from_secs(1)).await; /// send.send("shut down").unwrap(); /// }); /// /// loop { /// tokio::select! { /// _ = interval.tick() => println!("Another 100ms"), /// msg = &mut recv => { /// println!("Got message: {}", msg.unwrap()); /// break; /// } /// } /// } /// # handle.await.unwrap(); /// } /// ``` #[derive(Debug)] pub struct Receiver { inner: Option>>, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, #[cfg(all(tokio_unstable, feature = "tracing"))] async_op_span: tracing::Span, #[cfg(all(tokio_unstable, feature = "tracing"))] async_op_poll_span: tracing::Span, } pub mod error { //! `Oneshot` error types. use std::fmt; /// Error returned by the `Future` implementation for `Receiver`. /// /// This error is returned by the receiver when the sender is dropped without sending. #[derive(Debug, Eq, PartialEq, Clone)] pub struct RecvError(pub(super) ()); /// Error returned by the `try_recv` function on `Receiver`. #[derive(Debug, Eq, PartialEq, Clone)] pub enum TryRecvError { /// The send half of the channel has not yet sent a value. Empty, /// The send half of the channel was dropped without sending a value. Closed, } // ===== impl RecvError ===== impl fmt::Display for RecvError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } impl std::error::Error for RecvError {} // ===== impl TryRecvError ===== impl fmt::Display for TryRecvError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { match self { TryRecvError::Empty => write!(fmt, "channel empty"), TryRecvError::Closed => write!(fmt, "channel closed"), } } } impl std::error::Error for TryRecvError {} } use self::error::*; struct Inner { /// Manages the state of the inner cell. state: AtomicUsize, /// The value. This is set by `Sender` and read by `Receiver`. The state of /// the cell is tracked by `state`. value: UnsafeCell>, /// The task to notify when the receiver drops without consuming the value. /// /// ## Safety /// /// The `TX_TASK_SET` bit in the `state` field is set if this field is /// initialized. If that bit is unset, this field may be uninitialized. tx_task: Task, /// The task to notify when the value is sent. /// /// ## Safety /// /// The `RX_TASK_SET` bit in the `state` field is set if this field is /// initialized. If that bit is unset, this field may be uninitialized. rx_task: Task, } struct Task(UnsafeCell>); impl Task { unsafe fn will_wake(&self, cx: &mut Context<'_>) -> bool { self.with_task(|w| w.will_wake(cx.waker())) } unsafe fn with_task(&self, f: F) -> R where F: FnOnce(&Waker) -> R, { self.0.with(|ptr| { let waker: *const Waker = (*ptr).as_ptr(); f(&*waker) }) } unsafe fn drop_task(&self) { self.0.with_mut(|ptr| { let ptr: *mut Waker = (*ptr).as_mut_ptr(); ptr.drop_in_place(); }); } unsafe fn set_task(&self, cx: &mut Context<'_>) { self.0.with_mut(|ptr| { let ptr: *mut Waker = (*ptr).as_mut_ptr(); ptr.write(cx.waker().clone()); }); } } #[derive(Clone, Copy)] struct State(usize); /// Creates a new one-shot channel for sending single values across asynchronous /// tasks. /// /// The function returns separate "send" and "receive" handles. The `Sender` /// handle is used by the producer to send the value. The `Receiver` handle is /// used by the consumer to receive the value. /// /// Each handle can be used on separate tasks. /// /// # Examples /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel(); /// /// tokio::spawn(async move { /// if let Err(_) = tx.send(3) { /// println!("the receiver dropped"); /// } /// }); /// /// match rx.await { /// Ok(v) => println!("got = {:?}", v), /// Err(_) => println!("the sender dropped"), /// } /// } /// ``` #[track_caller] pub fn channel() -> (Sender, Receiver) { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = std::panic::Location::caller(); let resource_span = tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Sender|Receiver", kind = "Sync", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", tx_dropped = false, tx_dropped.op = "override", ) }); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", rx_dropped = false, rx_dropped.op = "override", ) }); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", value_sent = false, value_sent.op = "override", ) }); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", value_received = false, value_received.op = "override", ) }); resource_span }; let inner = Arc::new(Inner { state: AtomicUsize::new(State::new().as_usize()), value: UnsafeCell::new(None), tx_task: Task(UnsafeCell::new(MaybeUninit::uninit())), rx_task: Task(UnsafeCell::new(MaybeUninit::uninit())), }); let tx = Sender { inner: Some(inner.clone()), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: resource_span.clone(), }; #[cfg(all(tokio_unstable, feature = "tracing"))] let async_op_span = resource_span .in_scope(|| tracing::trace_span!("runtime.resource.async_op", source = "Receiver::await")); #[cfg(all(tokio_unstable, feature = "tracing"))] let async_op_poll_span = async_op_span.in_scope(|| tracing::trace_span!("runtime.resource.async_op.poll")); let rx = Receiver { inner: Some(inner), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, #[cfg(all(tokio_unstable, feature = "tracing"))] async_op_span, #[cfg(all(tokio_unstable, feature = "tracing"))] async_op_poll_span, }; (tx, rx) } impl Sender { /// Attempts to send a value on this channel, returning it back if it could /// not be sent. /// /// This method consumes `self` as only one value may ever be sent on a `oneshot` /// channel. It is not marked async because sending a message to an `oneshot` /// channel never requires any form of waiting. Because of this, the `send` /// method can be used in both synchronous and asynchronous code without /// problems. /// /// A successful send occurs when it is determined that the other end of the /// channel has not hung up already. An unsuccessful send would be one where /// the corresponding receiver has already been deallocated. Note that a /// return value of `Err` means that the data will never be received, but /// a return value of `Ok` does *not* mean that the data will be received. /// It is possible for the corresponding receiver to hang up immediately /// after this function returns `Ok`. /// /// # Examples /// /// Send a value to another task /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel(); /// /// tokio::spawn(async move { /// if let Err(_) = tx.send(3) { /// println!("the receiver dropped"); /// } /// }); /// /// match rx.await { /// Ok(v) => println!("got = {:?}", v), /// Err(_) => println!("the sender dropped"), /// } /// } /// ``` pub fn send(mut self, t: T) -> Result<(), T> { let inner = self.inner.take().unwrap(); inner.value.with_mut(|ptr| unsafe { // SAFETY: The receiver will not access the `UnsafeCell` unless the // channel has been marked as "complete" (the `VALUE_SENT` state bit // is set). // That bit is only set by the sender later on in this method, and // calling this method consumes `self`. Therefore, if it was possible to // call this method, we know that the `VALUE_SENT` bit is unset, and // the receiver is not currently accessing the `UnsafeCell`. *ptr = Some(t); }); if !inner.complete() { unsafe { // SAFETY: The receiver will not access the `UnsafeCell` unless // the channel has been marked as "complete". Calling // `complete()` will return true if this bit is set, and false // if it is not set. Thus, if `complete()` returned false, it is // safe for us to access the value, because we know that the // receiver will not. return Err(inner.consume_value().unwrap()); } } #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", value_sent = true, value_sent.op = "override", ) }); Ok(()) } /// Waits for the associated [`Receiver`] handle to close. /// /// A [`Receiver`] is closed by either calling [`close`] explicitly or the /// [`Receiver`] value is dropped. /// /// This function is useful when paired with `select!` to abort a /// computation when the receiver is no longer interested in the result. /// /// # Return /// /// Returns a `Future` which must be awaited on. /// /// [`Receiver`]: Receiver /// [`close`]: Receiver::close /// /// # Examples /// /// Basic usage /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (mut tx, rx) = oneshot::channel::<()>(); /// /// tokio::spawn(async move { /// drop(rx); /// }); /// /// tx.closed().await; /// println!("the receiver dropped"); /// } /// ``` /// /// Paired with select /// /// ``` /// use tokio::sync::oneshot; /// use tokio::time::{self, Duration}; /// /// async fn compute() -> String { /// // Complex computation returning a `String` /// # "hello".to_string() /// } /// /// #[tokio::main] /// async fn main() { /// let (mut tx, rx) = oneshot::channel(); /// /// tokio::spawn(async move { /// tokio::select! { /// _ = tx.closed() => { /// // The receiver dropped, no need to do any further work /// } /// value = compute() => { /// // The send can fail if the channel was closed at the exact same /// // time as when compute() finished, so just ignore the failure. /// let _ = tx.send(value); /// } /// } /// }); /// /// // Wait for up to 10 seconds /// let _ = time::timeout(Duration::from_secs(10), rx).await; /// } /// ``` pub async fn closed(&mut self) { use std::future::poll_fn; #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = self.resource_span.clone(); #[cfg(all(tokio_unstable, feature = "tracing"))] let closed = trace::async_op( || poll_fn(|cx| self.poll_closed(cx)), resource_span, "Sender::closed", "poll_closed", false, ); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let closed = poll_fn(|cx| self.poll_closed(cx)); closed.await; } /// Returns `true` if the associated [`Receiver`] handle has been dropped. /// /// A [`Receiver`] is closed by either calling [`close`] explicitly or the /// [`Receiver`] value is dropped. /// /// If `true` is returned, a call to `send` will always result in an error. /// /// [`Receiver`]: Receiver /// [`close`]: Receiver::close /// /// # Examples /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel(); /// /// assert!(!tx.is_closed()); /// /// drop(rx); /// /// assert!(tx.is_closed()); /// assert!(tx.send("never received").is_err()); /// } /// ``` pub fn is_closed(&self) -> bool { let inner = self.inner.as_ref().unwrap(); let state = State::load(&inner.state, Acquire); state.is_closed() } /// Checks whether the `oneshot` channel has been closed, and if not, schedules the /// `Waker` in the provided `Context` to receive a notification when the channel is /// closed. /// /// A [`Receiver`] is closed by either calling [`close`] explicitly, or when the /// [`Receiver`] value is dropped. /// /// Note that on multiple calls to poll, only the `Waker` from the `Context` passed /// to the most recent call will be scheduled to receive a wakeup. /// /// [`Receiver`]: struct@crate::sync::oneshot::Receiver /// [`close`]: fn@crate::sync::oneshot::Receiver::close /// /// # Return value /// /// This function returns: /// /// * `Poll::Pending` if the channel is still open. /// * `Poll::Ready(())` if the channel is closed. /// /// # Examples /// /// ``` /// use tokio::sync::oneshot; /// /// use std::future::poll_fn; /// /// #[tokio::main] /// async fn main() { /// let (mut tx, mut rx) = oneshot::channel::<()>(); /// /// tokio::spawn(async move { /// rx.close(); /// }); /// /// poll_fn(|cx| tx.poll_closed(cx)).await; /// /// println!("the receiver dropped"); /// } /// ``` pub fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()> { ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let inner = self.inner.as_ref().unwrap(); let mut state = State::load(&inner.state, Acquire); if state.is_closed() { coop.made_progress(); return Ready(()); } if state.is_tx_task_set() { let will_notify = unsafe { inner.tx_task.will_wake(cx) }; if !will_notify { state = State::unset_tx_task(&inner.state); if state.is_closed() { // Set the flag again so that the waker is released in drop State::set_tx_task(&inner.state); coop.made_progress(); return Ready(()); } else { unsafe { inner.tx_task.drop_task() }; } } } if !state.is_tx_task_set() { // Attempt to set the task unsafe { inner.tx_task.set_task(cx); } // Update the state state = State::set_tx_task(&inner.state); if state.is_closed() { coop.made_progress(); return Ready(()); } } Pending } } impl Drop for Sender { fn drop(&mut self) { if let Some(inner) = self.inner.as_ref() { inner.complete(); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", tx_dropped = true, tx_dropped.op = "override", ) }); } } } impl Receiver { /// Prevents the associated [`Sender`] handle from sending a value. /// /// Any `send` operation which happens after calling `close` is guaranteed /// to fail. After calling `close`, [`try_recv`] should be called to /// receive a value if one was sent **before** the call to `close` /// completed. /// /// This function is useful to perform a graceful shutdown and ensure that a /// value will not be sent into the channel and never received. /// /// `close` is no-op if a message is already received or the channel /// is already closed. /// /// [`Sender`]: Sender /// [`try_recv`]: Receiver::try_recv /// /// # Examples /// /// Prevent a value from being sent /// /// ``` /// use tokio::sync::oneshot; /// use tokio::sync::oneshot::error::TryRecvError; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = oneshot::channel(); /// /// assert!(!tx.is_closed()); /// /// rx.close(); /// /// assert!(tx.is_closed()); /// assert!(tx.send("never received").is_err()); /// /// match rx.try_recv() { /// Err(TryRecvError::Closed) => {} /// _ => unreachable!(), /// } /// } /// ``` /// /// Receive a value sent **before** calling `close` /// /// ``` /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = oneshot::channel(); /// /// assert!(tx.send("will receive").is_ok()); /// /// rx.close(); /// /// let msg = rx.try_recv().unwrap(); /// assert_eq!(msg, "will receive"); /// } /// ``` pub fn close(&mut self) { if let Some(inner) = self.inner.as_ref() { inner.close(); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", rx_dropped = true, rx_dropped.op = "override", ) }); } } /// Attempts to receive a value. /// /// If a pending value exists in the channel, it is returned. If no value /// has been sent, the current task **will not** be registered for /// future notification. /// /// This function is useful to call from outside the context of an /// asynchronous task. /// /// Note that unlike the `poll` method, the `try_recv` method cannot fail /// spuriously. Any send or close event that happens before this call to /// `try_recv` will be correctly returned to the caller. /// /// # Return /// /// - `Ok(T)` if a value is pending in the channel. /// - `Err(TryRecvError::Empty)` if no value has been sent yet. /// - `Err(TryRecvError::Closed)` if the sender has dropped without sending /// a value, or if the message has already been received. /// /// # Examples /// /// `try_recv` before a value is sent, then after. /// /// ``` /// use tokio::sync::oneshot; /// use tokio::sync::oneshot::error::TryRecvError; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = oneshot::channel(); /// /// match rx.try_recv() { /// // The channel is currently empty /// Err(TryRecvError::Empty) => {} /// _ => unreachable!(), /// } /// /// // Send a value /// tx.send("hello").unwrap(); /// /// match rx.try_recv() { /// Ok(value) => assert_eq!(value, "hello"), /// _ => unreachable!(), /// } /// } /// ``` /// /// `try_recv` when the sender dropped before sending a value /// /// ``` /// use tokio::sync::oneshot; /// use tokio::sync::oneshot::error::TryRecvError; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = oneshot::channel::<()>(); /// /// drop(tx); /// /// match rx.try_recv() { /// // The channel will never receive a value. /// Err(TryRecvError::Closed) => {} /// _ => unreachable!(), /// } /// } /// ``` pub fn try_recv(&mut self) -> Result { let result = if let Some(inner) = self.inner.as_ref() { let state = State::load(&inner.state, Acquire); if state.is_complete() { // SAFETY: If `state.is_complete()` returns true, then the // `VALUE_SENT` bit has been set and the sender side of the // channel will no longer attempt to access the inner // `UnsafeCell`. Therefore, it is now safe for us to access the // cell. match unsafe { inner.consume_value() } { Some(value) => { #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", value_received = true, value_received.op = "override", ) }); Ok(value) } None => Err(TryRecvError::Closed), } } else if state.is_closed() { Err(TryRecvError::Closed) } else { // Not ready, this does not clear `inner` return Err(TryRecvError::Empty); } } else { Err(TryRecvError::Closed) }; self.inner = None; result } /// Blocking receive to call outside of asynchronous contexts. /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// /// ``` /// use std::thread; /// use tokio::sync::oneshot; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = oneshot::channel::(); /// /// let sync_code = thread::spawn(move || { /// assert_eq!(Ok(10), rx.blocking_recv()); /// }); /// /// let _ = tx.send(10); /// sync_code.join().unwrap(); /// } /// ``` #[track_caller] #[cfg(feature = "sync")] #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(self) -> Result { crate::future::block_on(self) } } impl Drop for Receiver { fn drop(&mut self) { if let Some(inner) = self.inner.as_ref() { let state = inner.close(); if state.is_complete() { // SAFETY: we have ensured that the `VALUE_SENT` bit has been set, // so only the receiver can access the value. drop(unsafe { inner.consume_value() }); } #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", rx_dropped = true, rx_dropped.op = "override", ) }); } } } impl Future for Receiver { type Output = Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { // If `inner` is `None`, then `poll()` has already completed. #[cfg(all(tokio_unstable, feature = "tracing"))] let _res_span = self.resource_span.clone().entered(); #[cfg(all(tokio_unstable, feature = "tracing"))] let _ao_span = self.async_op_span.clone().entered(); #[cfg(all(tokio_unstable, feature = "tracing"))] let _ao_poll_span = self.async_op_poll_span.clone().entered(); let ret = if let Some(inner) = self.as_ref().get_ref().inner.as_ref() { #[cfg(all(tokio_unstable, feature = "tracing"))] let res = ready!(trace_poll_op!("poll_recv", inner.poll_recv(cx)))?; #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let res = ready!(inner.poll_recv(cx))?; res } else { panic!("called after complete"); }; self.inner = None; Ready(Ok(ret)) } } impl Inner { fn complete(&self) -> bool { let prev = State::set_complete(&self.state); if prev.is_closed() { return false; } if prev.is_rx_task_set() { // TODO: Consume waker? unsafe { self.rx_task.with_task(Waker::wake_by_ref); } } true } fn poll_recv(&self, cx: &mut Context<'_>) -> Poll> { ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); // Load the state let mut state = State::load(&self.state, Acquire); if state.is_complete() { coop.made_progress(); match unsafe { self.consume_value() } { Some(value) => Ready(Ok(value)), None => Ready(Err(RecvError(()))), } } else if state.is_closed() { coop.made_progress(); Ready(Err(RecvError(()))) } else { if state.is_rx_task_set() { let will_notify = unsafe { self.rx_task.will_wake(cx) }; // Check if the task is still the same if !will_notify { // Unset the task state = State::unset_rx_task(&self.state); if state.is_complete() { // Set the flag again so that the waker is released in drop State::set_rx_task(&self.state); coop.made_progress(); // SAFETY: If `state.is_complete()` returns true, then the // `VALUE_SENT` bit has been set and the sender side of the // channel will no longer attempt to access the inner // `UnsafeCell`. Therefore, it is now safe for us to access the // cell. return match unsafe { self.consume_value() } { Some(value) => Ready(Ok(value)), None => Ready(Err(RecvError(()))), }; } else { unsafe { self.rx_task.drop_task() }; } } } if !state.is_rx_task_set() { // Attempt to set the task unsafe { self.rx_task.set_task(cx); } // Update the state state = State::set_rx_task(&self.state); if state.is_complete() { coop.made_progress(); match unsafe { self.consume_value() } { Some(value) => Ready(Ok(value)), None => Ready(Err(RecvError(()))), } } else { Pending } } else { Pending } } } /// Called by `Receiver` to indicate that the value will never be received. fn close(&self) -> State { let prev = State::set_closed(&self.state); if prev.is_tx_task_set() && !prev.is_complete() { unsafe { self.tx_task.with_task(Waker::wake_by_ref); } } prev } /// Consumes the value. This function does not check `state`. /// /// # Safety /// /// Calling this method concurrently on multiple threads will result in a /// data race. The `VALUE_SENT` state bit is used to ensure that only the /// sender *or* the receiver will call this method at a given point in time. /// If `VALUE_SENT` is not set, then only the sender may call this method; /// if it is set, then only the receiver may call this method. unsafe fn consume_value(&self) -> Option { self.value.with_mut(|ptr| (*ptr).take()) } } unsafe impl Send for Inner {} unsafe impl Sync for Inner {} fn mut_load(this: &mut AtomicUsize) -> usize { this.with_mut(|v| *v) } impl Drop for Inner { fn drop(&mut self) { let state = State(mut_load(&mut self.state)); if state.is_rx_task_set() { unsafe { self.rx_task.drop_task(); } } if state.is_tx_task_set() { unsafe { self.tx_task.drop_task(); } } // SAFETY: we have `&mut self`, and therefore we have // exclusive access to the value. unsafe { // Note: the assertion holds because if the value has been sent by sender, // we must ensure that the value must have been consumed by the receiver before // dropping the `Inner`. debug_assert!(self.consume_value().is_none()); } } } impl fmt::Debug for Inner { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { use std::sync::atomic::Ordering::Relaxed; fmt.debug_struct("Inner") .field("state", &State::load(&self.state, Relaxed)) .finish() } } /// Indicates that a waker for the receiving task has been set. /// /// # Safety /// /// If this bit is not set, the `rx_task` field may be uninitialized. const RX_TASK_SET: usize = 0b00001; /// Indicates that a value has been stored in the channel's inner `UnsafeCell`. /// /// # Safety /// /// This bit controls which side of the channel is permitted to access the /// `UnsafeCell`. If it is set, the `UnsafeCell` may ONLY be accessed by the /// receiver. If this bit is NOT set, the `UnsafeCell` may ONLY be accessed by /// the sender. const VALUE_SENT: usize = 0b00010; const CLOSED: usize = 0b00100; /// Indicates that a waker for the sending task has been set. /// /// # Safety /// /// If this bit is not set, the `tx_task` field may be uninitialized. const TX_TASK_SET: usize = 0b01000; impl State { fn new() -> State { State(0) } fn is_complete(self) -> bool { self.0 & VALUE_SENT == VALUE_SENT } fn set_complete(cell: &AtomicUsize) -> State { // This method is a compare-and-swap loop rather than a fetch-or like // other `set_$WHATEVER` methods on `State`. This is because we must // check if the state has been closed before setting the `VALUE_SENT` // bit. // // We don't want to set both the `VALUE_SENT` bit if the `CLOSED` // bit is already set, because `VALUE_SENT` will tell the receiver that // it's okay to access the inner `UnsafeCell`. Immediately after calling // `set_complete`, if the channel was closed, the sender will _also_ // access the `UnsafeCell` to take the value back out, so if a // `poll_recv` or `try_recv` call is occurring concurrently, both // threads may try to access the `UnsafeCell` if we were to set the // `VALUE_SENT` bit on a closed channel. let mut state = cell.load(Ordering::Relaxed); loop { if State(state).is_closed() { break; } // TODO: This could be `Release`, followed by an `Acquire` fence *if* // the `RX_TASK_SET` flag is set. However, `loom` does not support // fences yet. match cell.compare_exchange_weak( state, state | VALUE_SENT, Ordering::AcqRel, Ordering::Acquire, ) { Ok(_) => break, Err(actual) => state = actual, } } State(state) } fn is_rx_task_set(self) -> bool { self.0 & RX_TASK_SET == RX_TASK_SET } fn set_rx_task(cell: &AtomicUsize) -> State { let val = cell.fetch_or(RX_TASK_SET, AcqRel); State(val | RX_TASK_SET) } fn unset_rx_task(cell: &AtomicUsize) -> State { let val = cell.fetch_and(!RX_TASK_SET, AcqRel); State(val & !RX_TASK_SET) } fn is_closed(self) -> bool { self.0 & CLOSED == CLOSED } fn set_closed(cell: &AtomicUsize) -> State { // Acquire because we want all later writes (attempting to poll) to be // ordered after this. let val = cell.fetch_or(CLOSED, Acquire); State(val) } fn set_tx_task(cell: &AtomicUsize) -> State { let val = cell.fetch_or(TX_TASK_SET, AcqRel); State(val | TX_TASK_SET) } fn unset_tx_task(cell: &AtomicUsize) -> State { let val = cell.fetch_and(!TX_TASK_SET, AcqRel); State(val & !TX_TASK_SET) } fn is_tx_task_set(self) -> bool { self.0 & TX_TASK_SET == TX_TASK_SET } fn as_usize(self) -> usize { self.0 } fn load(cell: &AtomicUsize, order: Ordering) -> State { let val = cell.load(order); State(val) } } impl fmt::Debug for State { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("State") .field("is_complete", &self.is_complete()) .field("is_closed", &self.is_closed()) .field("is_rx_task_set", &self.is_rx_task_set()) .field("is_tx_task_set", &self.is_tx_task_set()) .finish() } } tokio-1.43.1/src/sync/rwlock/owned_read_guard.rs000064400000000000000000000142571046102023000177530ustar 00000000000000use crate::sync::rwlock::RwLock; use std::marker::PhantomData; use std::sync::Arc; use std::{fmt, mem, ops, ptr}; /// Owned RAII structure used to release the shared read access of a lock when /// dropped. /// /// This structure is created by the [`read_owned`] method on /// [`RwLock`]. /// /// [`read_owned`]: method@crate::sync::RwLock::read_owned /// [`RwLock`]: struct@crate::sync::RwLock #[clippy::has_significant_drop] pub struct OwnedRwLockReadGuard { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) resource_span: tracing::Span, pub(super) lock: Arc>, pub(super) data: *const U, pub(super) _p: PhantomData, } #[allow(dead_code)] // Unused fields are still used in Drop. struct Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, lock: Arc>, data: *const U, } impl OwnedRwLockReadGuard { fn skip_drop(self) -> Inner { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. unsafe { Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: ptr::read(&me.resource_span), lock: ptr::read(&me.lock), data: me.data, } } } /// Makes a new `OwnedRwLockReadGuard` for a component of the locked data. /// This operation cannot fail as the `OwnedRwLockReadGuard` passed in /// already locked the data. /// /// This is an associated function that needs to be /// used as `OwnedRwLockReadGuard::map(...)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockReadGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// let guard = lock.read_owned().await; /// let guard = OwnedRwLockReadGuard::map(guard, |f| &f.0); /// /// assert_eq!(1, *guard); /// # } /// ``` #[inline] pub fn map(this: Self, f: F) -> OwnedRwLockReadGuard where F: FnOnce(&U) -> &V, { let data = f(&*this) as *const V; let this = this.skip_drop(); OwnedRwLockReadGuard { lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, } } /// Attempts to make a new [`OwnedRwLockReadGuard`] for a component of the /// locked data. The original guard is returned if the closure returns /// `None`. /// /// This operation cannot fail as the `OwnedRwLockReadGuard` passed in /// already locked the data. /// /// This is an associated function that needs to be used as /// `OwnedRwLockReadGuard::try_map(..)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockReadGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// let guard = lock.read_owned().await; /// let guard = OwnedRwLockReadGuard::try_map(guard, |f| Some(&f.0)).expect("should not fail"); /// /// assert_eq!(1, *guard); /// # } /// ``` #[inline] pub fn try_map(this: Self, f: F) -> Result, Self> where F: FnOnce(&U) -> Option<&V>, { let data = match f(&*this) { Some(data) => data as *const V, None => return Err(this), }; let this = this.skip_drop(); Ok(OwnedRwLockReadGuard { lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }) } /// Returns a reference to the original `Arc`. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockReadGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// let guard = lock.clone().read_owned().await; /// assert!(Arc::ptr_eq(&lock, OwnedRwLockReadGuard::rwlock(&guard))); /// /// let guard = OwnedRwLockReadGuard::map(guard, |f| &f.0); /// assert!(Arc::ptr_eq(&lock, OwnedRwLockReadGuard::rwlock(&guard))); /// # } /// ``` pub fn rwlock(this: &Self) -> &Arc> { &this.lock } } impl ops::Deref for OwnedRwLockReadGuard { type Target = U; fn deref(&self) -> &U { unsafe { &*self.data } } } impl fmt::Debug for OwnedRwLockReadGuard where U: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for OwnedRwLockReadGuard where U: fmt::Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } impl Drop for OwnedRwLockReadGuard { fn drop(&mut self) { self.lock.s.release(1); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "sub", ) }); } } tokio-1.43.1/src/sync/rwlock/owned_write_guard.rs000064400000000000000000000347401046102023000201710ustar 00000000000000use crate::sync::rwlock::owned_read_guard::OwnedRwLockReadGuard; use crate::sync::rwlock::owned_write_guard_mapped::OwnedRwLockMappedWriteGuard; use crate::sync::rwlock::RwLock; use std::marker::PhantomData; use std::sync::Arc; use std::{fmt, mem, ops, ptr}; /// Owned RAII structure used to release the exclusive write access of a lock when /// dropped. /// /// This structure is created by the [`write_owned`] method /// on [`RwLock`]. /// /// [`write_owned`]: method@crate::sync::RwLock::write_owned /// [`RwLock`]: struct@crate::sync::RwLock #[clippy::has_significant_drop] pub struct OwnedRwLockWriteGuard { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, pub(super) lock: Arc>, pub(super) data: *mut T, pub(super) _p: PhantomData, } #[allow(dead_code)] // Unused fields are still used in Drop. struct Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, permits_acquired: u32, lock: Arc>, data: *const T, } impl OwnedRwLockWriteGuard { fn skip_drop(self) -> Inner { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. unsafe { Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: ptr::read(&me.resource_span), permits_acquired: me.permits_acquired, lock: ptr::read(&me.lock), data: me.data, } } } /// Makes a new [`OwnedRwLockMappedWriteGuard`] for a component of the locked /// data. /// /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in /// already locked the data. /// /// This is an associated function that needs to be used as /// `OwnedRwLockWriteGuard::map(..)`. A method would interfere with methods /// of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// { /// let lock = Arc::clone(&lock); /// let mut mapped = OwnedRwLockWriteGuard::map(lock.write_owned().await, |f| &mut f.0); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn map(mut this: Self, f: F) -> OwnedRwLockMappedWriteGuard where F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; let this = this.skip_drop(); OwnedRwLockMappedWriteGuard { permits_acquired: this.permits_acquired, lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, } } /// Makes a new [`OwnedRwLockReadGuard`] for a component of the locked data. /// /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be used as /// `OwnedRwLockWriteGuard::downgrade_map(..)`. A method would interfere with methods of /// the same name on the contents of the locked data. /// /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a /// `&mut T` would result in unsoundness, as you could use interior mutability. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// let guard = Arc::clone(&lock).write_owned().await; /// let mapped = OwnedRwLockWriteGuard::downgrade_map(guard, |f| &f.0); /// let foo = lock.read_owned().await; /// assert_eq!(foo.0, *mapped); /// # } /// ``` #[inline] pub fn downgrade_map(this: Self, f: F) -> OwnedRwLockReadGuard where F: FnOnce(&T) -> &U, { let data = f(&*this) as *const U; let this = this.skip_drop(); let guard = OwnedRwLockReadGuard { lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }; // Release all but one of the permits held by the write guard let to_release = (this.permits_acquired - 1) as usize; guard.lock.s.release(to_release); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); guard } /// Attempts to make a new [`OwnedRwLockMappedWriteGuard`] for a component /// of the locked data. The original guard is returned if the closure /// returns `None`. /// /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in /// already locked the data. /// /// This is an associated function that needs to be /// used as `OwnedRwLockWriteGuard::try_map(...)`. A method would interfere /// with methods of the same name on the contents of the locked data. /// /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// { /// let guard = Arc::clone(&lock).write_owned().await; /// let mut guard = OwnedRwLockWriteGuard::try_map(guard, |f| Some(&mut f.0)).expect("should not fail"); /// *guard = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn try_map( mut this: Self, f: F, ) -> Result, Self> where F: FnOnce(&mut T) -> Option<&mut U>, { let data = match f(&mut *this) { Some(data) => data as *mut U, None => return Err(this), }; let this = this.skip_drop(); Ok(OwnedRwLockMappedWriteGuard { permits_acquired: this.permits_acquired, lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }) } /// Attempts to make a new [`OwnedRwLockReadGuard`] for a component of /// the locked data. The original guard is returned if the closure returns /// `None`. /// /// This operation cannot fail as the `OwnedRwLockWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be /// used as `OwnedRwLockWriteGuard::try_downgrade_map(...)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a /// `&mut T` would result in unsoundness, as you could use interior mutability. /// /// If this function returns `Err(...)`, the lock is never unlocked nor downgraded. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// let guard = Arc::clone(&lock).write_owned().await; /// let guard = OwnedRwLockWriteGuard::try_downgrade_map(guard, |f| Some(&f.0)).expect("should not fail"); /// let foo = lock.read_owned().await; /// assert_eq!(foo.0, *guard); /// # } /// ``` #[inline] pub fn try_downgrade_map( this: Self, f: F, ) -> Result, Self> where F: FnOnce(&T) -> Option<&U>, { let data = match f(&*this) { Some(data) => data as *const U, None => return Err(this), }; let this = this.skip_drop(); let guard = OwnedRwLockReadGuard { lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }; // Release all but one of the permits held by the write guard let to_release = (this.permits_acquired - 1) as usize; guard.lock.s.release(to_release); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); Ok(guard) } /// Converts this `OwnedRwLockWriteGuard` into an /// `OwnedRwLockMappedWriteGuard`. This method can be used to store a /// non-mapped guard in a struct field that expects a mapped guard. /// /// This is equivalent to calling `OwnedRwLockWriteGuard::map(guard, |me| me)`. #[inline] pub fn into_mapped(this: Self) -> OwnedRwLockMappedWriteGuard { Self::map(this, |me| me) } /// Atomically downgrades a write lock into a read lock without allowing /// any writers to take exclusive access of the lock in the meantime. /// /// **Note:** This won't *necessarily* allow any additional readers to acquire /// locks, since [`RwLock`] is fair and it is possible that a writer is next /// in line. /// /// Returns an RAII guard which will drop this read access of the `RwLock` /// when dropped. /// /// # Examples /// /// ``` /// # use tokio::sync::RwLock; /// # use std::sync::Arc; /// # /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// /// let n = lock.clone().write_owned().await; /// /// let cloned_lock = lock.clone(); /// let handle = tokio::spawn(async move { /// *cloned_lock.write_owned().await = 2; /// }); /// /// let n = n.downgrade(); /// assert_eq!(*n, 1, "downgrade is atomic"); /// /// drop(n); /// handle.await.unwrap(); /// assert_eq!(*lock.read().await, 2, "second writer obtained write lock"); /// # } /// ``` pub fn downgrade(self) -> OwnedRwLockReadGuard { let this = self.skip_drop(); let guard = OwnedRwLockReadGuard { lock: this.lock, data: this.data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }; // Release all but one of the permits held by the write guard let to_release = (this.permits_acquired - 1) as usize; guard.lock.s.release(to_release); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); guard } /// Returns a reference to the original `Arc`. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// /// let guard = lock.clone().write_owned().await; /// assert!(Arc::ptr_eq(&lock, OwnedRwLockWriteGuard::rwlock(&guard))); /// # } /// ``` pub fn rwlock(this: &Self) -> &Arc> { &this.lock } } impl ops::Deref for OwnedRwLockWriteGuard { type Target = T; fn deref(&self) -> &T { unsafe { &*self.data } } } impl ops::DerefMut for OwnedRwLockWriteGuard { fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.data } } } impl fmt::Debug for OwnedRwLockWriteGuard where T: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for OwnedRwLockWriteGuard where T: fmt::Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } impl Drop for OwnedRwLockWriteGuard { fn drop(&mut self) { self.lock.s.release(self.permits_acquired as usize); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); } } tokio-1.43.1/src/sync/rwlock/owned_write_guard_mapped.rs000064400000000000000000000160161046102023000215130ustar 00000000000000use crate::sync::rwlock::RwLock; use std::marker::PhantomData; use std::sync::Arc; use std::{fmt, mem, ops, ptr}; /// Owned RAII structure used to release the exclusive write access of a lock when /// dropped. /// /// This structure is created by [mapping] an [`OwnedRwLockWriteGuard`]. It is a /// separate type from `OwnedRwLockWriteGuard` to disallow downgrading a mapped /// guard, since doing so can cause undefined behavior. /// /// [mapping]: method@crate::sync::OwnedRwLockWriteGuard::map /// [`OwnedRwLockWriteGuard`]: struct@crate::sync::OwnedRwLockWriteGuard #[clippy::has_significant_drop] pub struct OwnedRwLockMappedWriteGuard { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, pub(super) lock: Arc>, pub(super) data: *mut U, pub(super) _p: PhantomData, } #[allow(dead_code)] // Unused fields are still used in Drop. struct Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, permits_acquired: u32, lock: Arc>, data: *const U, } impl OwnedRwLockMappedWriteGuard { fn skip_drop(self) -> Inner { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. unsafe { Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: ptr::read(&me.resource_span), permits_acquired: me.permits_acquired, lock: ptr::read(&me.lock), data: me.data, } } } /// Makes a new `OwnedRwLockMappedWriteGuard` for a component of the locked /// data. /// /// This operation cannot fail as the `OwnedRwLockMappedWriteGuard` passed /// in already locked the data. /// /// This is an associated function that needs to be used as /// `OwnedRwLockWriteGuard::map(..)`. A method would interfere with methods /// of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// { /// let lock = Arc::clone(&lock); /// let mut mapped = OwnedRwLockWriteGuard::map(lock.write_owned().await, |f| &mut f.0); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn map(mut this: Self, f: F) -> OwnedRwLockMappedWriteGuard where F: FnOnce(&mut U) -> &mut V, { let data = f(&mut *this) as *mut V; let this = this.skip_drop(); OwnedRwLockMappedWriteGuard { permits_acquired: this.permits_acquired, lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, } } /// Attempts to make a new `OwnedRwLockMappedWriteGuard` for a component /// of the locked data. The original guard is returned if the closure /// returns `None`. /// /// This operation cannot fail as the `OwnedRwLockMappedWriteGuard` passed /// in already locked the data. /// /// This is an associated function that needs to be /// used as `OwnedRwLockMappedWriteGuard::try_map(...)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(Foo(1))); /// /// { /// let guard = Arc::clone(&lock).write_owned().await; /// let mut guard = OwnedRwLockWriteGuard::try_map(guard, |f| Some(&mut f.0)).expect("should not fail"); /// *guard = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn try_map( mut this: Self, f: F, ) -> Result, Self> where F: FnOnce(&mut U) -> Option<&mut V>, { let data = match f(&mut *this) { Some(data) => data as *mut V, None => return Err(this), }; let this = this.skip_drop(); Ok(OwnedRwLockMappedWriteGuard { permits_acquired: this.permits_acquired, lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }) } /// Returns a reference to the original `Arc`. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{ /// RwLock, /// OwnedRwLockWriteGuard, /// OwnedRwLockMappedWriteGuard, /// }; /// /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// /// let guard = lock.clone().write_owned().await; /// let guard = OwnedRwLockWriteGuard::map(guard, |x| x); /// assert!(Arc::ptr_eq(&lock, OwnedRwLockMappedWriteGuard::rwlock(&guard))); /// # } /// ``` pub fn rwlock(this: &Self) -> &Arc> { &this.lock } } impl ops::Deref for OwnedRwLockMappedWriteGuard { type Target = U; fn deref(&self) -> &U { unsafe { &*self.data } } } impl ops::DerefMut for OwnedRwLockMappedWriteGuard { fn deref_mut(&mut self) -> &mut U { unsafe { &mut *self.data } } } impl fmt::Debug for OwnedRwLockMappedWriteGuard where U: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for OwnedRwLockMappedWriteGuard where U: fmt::Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } impl Drop for OwnedRwLockMappedWriteGuard { fn drop(&mut self) { self.lock.s.release(self.permits_acquired as usize); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); } } tokio-1.43.1/src/sync/rwlock/read_guard.rs000064400000000000000000000134171046102023000165540ustar 00000000000000use crate::sync::batch_semaphore::Semaphore; use std::marker::PhantomData; use std::{fmt, mem, ops}; /// RAII structure used to release the shared read access of a lock when /// dropped. /// /// This structure is created by the [`read`] method on /// [`RwLock`]. /// /// [`read`]: method@crate::sync::RwLock::read /// [`RwLock`]: struct@crate::sync::RwLock #[clippy::has_significant_drop] #[must_use = "if unused the RwLock will immediately unlock"] pub struct RwLockReadGuard<'a, T: ?Sized> { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) resource_span: tracing::Span, pub(super) s: &'a Semaphore, pub(super) data: *const T, pub(super) marker: PhantomData<&'a T>, } #[allow(dead_code)] // Unused fields are still used in Drop. struct Inner<'a, T: ?Sized> { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, s: &'a Semaphore, data: *const T, } impl<'a, T: ?Sized> RwLockReadGuard<'a, T> { fn skip_drop(self) -> Inner<'a, T> { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: unsafe { std::ptr::read(&me.resource_span) }, s: me.s, data: me.data, } } /// Makes a new `RwLockReadGuard` for a component of the locked data. /// /// This operation cannot fail as the `RwLockReadGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be /// used as `RwLockReadGuard::map(...)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// This is an asynchronous version of [`RwLockReadGuard::map`] from the /// [`parking_lot` crate]. /// /// [`RwLockReadGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockReadGuard.html#method.map /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockReadGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// let guard = lock.read().await; /// let guard = RwLockReadGuard::map(guard, |f| &f.0); /// /// assert_eq!(1, *guard); /// # } /// ``` #[inline] pub fn map(this: Self, f: F) -> RwLockReadGuard<'a, U> where F: FnOnce(&T) -> &U, { let data = f(&*this) as *const U; let this = this.skip_drop(); RwLockReadGuard { s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, } } /// Attempts to make a new [`RwLockReadGuard`] for a component of the /// locked data. The original guard is returned if the closure returns /// `None`. /// /// This operation cannot fail as the `RwLockReadGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be used as /// `RwLockReadGuard::try_map(..)`. A method would interfere with methods of the /// same name on the contents of the locked data. /// /// This is an asynchronous version of [`RwLockReadGuard::try_map`] from the /// [`parking_lot` crate]. /// /// [`RwLockReadGuard::try_map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockReadGuard.html#method.try_map /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockReadGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// let guard = lock.read().await; /// let guard = RwLockReadGuard::try_map(guard, |f| Some(&f.0)).expect("should not fail"); /// /// assert_eq!(1, *guard); /// # } /// ``` #[inline] pub fn try_map(this: Self, f: F) -> Result, Self> where F: FnOnce(&T) -> Option<&U>, { let data = match f(&*this) { Some(data) => data as *const U, None => return Err(this), }; let this = this.skip_drop(); Ok(RwLockReadGuard { s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }) } } impl ops::Deref for RwLockReadGuard<'_, T> { type Target = T; fn deref(&self) -> &T { unsafe { &*self.data } } } impl<'a, T: ?Sized> fmt::Debug for RwLockReadGuard<'a, T> where T: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl<'a, T: ?Sized> fmt::Display for RwLockReadGuard<'a, T> where T: fmt::Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } impl<'a, T: ?Sized> Drop for RwLockReadGuard<'a, T> { fn drop(&mut self) { self.s.release(1); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "sub", ) }); } } tokio-1.43.1/src/sync/rwlock/write_guard.rs000064400000000000000000000363431046102023000167760ustar 00000000000000use crate::sync::batch_semaphore::Semaphore; use crate::sync::rwlock::read_guard::RwLockReadGuard; use crate::sync::rwlock::write_guard_mapped::RwLockMappedWriteGuard; use std::marker::PhantomData; use std::{fmt, mem, ops}; /// RAII structure used to release the exclusive write access of a lock when /// dropped. /// /// This structure is created by the [`write`] method /// on [`RwLock`]. /// /// [`write`]: method@crate::sync::RwLock::write /// [`RwLock`]: struct@crate::sync::RwLock #[clippy::has_significant_drop] #[must_use = "if unused the RwLock will immediately unlock"] pub struct RwLockWriteGuard<'a, T: ?Sized> { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, pub(super) s: &'a Semaphore, pub(super) data: *mut T, pub(super) marker: PhantomData<&'a mut T>, } #[allow(dead_code)] // Unused fields are still used in Drop. struct Inner<'a, T: ?Sized> { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, permits_acquired: u32, s: &'a Semaphore, data: *mut T, } impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { fn skip_drop(self) -> Inner<'a, T> { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: unsafe { std::ptr::read(&me.resource_span) }, permits_acquired: me.permits_acquired, s: me.s, data: me.data, } } /// Makes a new [`RwLockMappedWriteGuard`] for a component of the locked data. /// /// This operation cannot fail as the `RwLockWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be used as /// `RwLockWriteGuard::map(..)`. A method would interfere with methods of /// the same name on the contents of the locked data. /// /// This is an asynchronous version of [`RwLockWriteGuard::map`] from the /// [`parking_lot` crate]. /// /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// { /// let mut mapped = RwLockWriteGuard::map(lock.write().await, |f| &mut f.0); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn map(mut this: Self, f: F) -> RwLockMappedWriteGuard<'a, U> where F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; let this = this.skip_drop(); RwLockMappedWriteGuard { permits_acquired: this.permits_acquired, s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, } } /// Makes a new [`RwLockReadGuard`] for a component of the locked data. /// /// This operation cannot fail as the `RwLockWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be used as /// `RwLockWriteGuard::downgrade_map(..)`. A method would interfere with methods of /// the same name on the contents of the locked data. /// /// This is equivalent to a combination of asynchronous [`RwLockWriteGuard::map`] and [`RwLockWriteGuard::downgrade`] /// from the [`parking_lot` crate]. /// /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a /// `&mut T` would result in unsoundness, as you could use interior mutability. /// /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map /// [`RwLockWriteGuard::downgrade`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.downgrade /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// let mapped = RwLockWriteGuard::downgrade_map(lock.write().await, |f| &f.0); /// let foo = lock.read().await; /// assert_eq!(foo.0, *mapped); /// # } /// ``` #[inline] pub fn downgrade_map(this: Self, f: F) -> RwLockReadGuard<'a, U> where F: FnOnce(&T) -> &U, { let data = f(&*this) as *const U; let this = this.skip_drop(); let guard = RwLockReadGuard { s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }; // Release all but one of the permits held by the write guard let to_release = (this.permits_acquired - 1) as usize; this.s.release(to_release); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); guard } /// Attempts to make a new [`RwLockMappedWriteGuard`] for a component of /// the locked data. The original guard is returned if the closure returns /// `None`. /// /// This operation cannot fail as the `RwLockWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be /// used as `RwLockWriteGuard::try_map(...)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// This is an asynchronous version of [`RwLockWriteGuard::try_map`] from /// the [`parking_lot` crate]. /// /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard /// [`RwLockWriteGuard::try_map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.try_map /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// { /// let guard = lock.write().await; /// let mut guard = RwLockWriteGuard::try_map(guard, |f| Some(&mut f.0)).expect("should not fail"); /// *guard = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn try_map( mut this: Self, f: F, ) -> Result, Self> where F: FnOnce(&mut T) -> Option<&mut U>, { let data = match f(&mut *this) { Some(data) => data as *mut U, None => return Err(this), }; let this = this.skip_drop(); Ok(RwLockMappedWriteGuard { permits_acquired: this.permits_acquired, s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }) } /// Attempts to make a new [`RwLockReadGuard`] for a component of /// the locked data. The original guard is returned if the closure returns /// `None`. /// /// This operation cannot fail as the `RwLockWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be /// used as `RwLockWriteGuard::try_downgrade_map(...)`. A method would interfere with /// methods of the same name on the contents of the locked data. /// /// This is equivalent to a combination of asynchronous [`RwLockWriteGuard::try_map`] and [`RwLockWriteGuard::downgrade`] /// from the [`parking_lot` crate]. /// /// Inside of `f`, you retain exclusive access to the data, despite only being given a `&T`. Handing out a /// `&mut T` would result in unsoundness, as you could use interior mutability. /// /// If this function returns `Err(...)`, the lock is never unlocked nor downgraded. /// /// [`RwLockMappedWriteGuard`]: struct@crate::sync::RwLockMappedWriteGuard /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map /// [`RwLockWriteGuard::downgrade`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.downgrade /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// let guard = RwLockWriteGuard::try_downgrade_map(lock.write().await, |f| Some(&f.0)).expect("should not fail"); /// let foo = lock.read().await; /// assert_eq!(foo.0, *guard); /// # } /// ``` #[inline] pub fn try_downgrade_map(this: Self, f: F) -> Result, Self> where F: FnOnce(&T) -> Option<&U>, { let data = match f(&*this) { Some(data) => data as *const U, None => return Err(this), }; let this = this.skip_drop(); let guard = RwLockReadGuard { s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }; // Release all but one of the permits held by the write guard let to_release = (this.permits_acquired - 1) as usize; this.s.release(to_release); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); Ok(guard) } /// Converts this `RwLockWriteGuard` into an `RwLockMappedWriteGuard`. This /// method can be used to store a non-mapped guard in a struct field that /// expects a mapped guard. /// /// This is equivalent to calling `RwLockWriteGuard::map(guard, |me| me)`. #[inline] pub fn into_mapped(this: Self) -> RwLockMappedWriteGuard<'a, T> { RwLockWriteGuard::map(this, |me| me) } /// Atomically downgrades a write lock into a read lock without allowing /// any writers to take exclusive access of the lock in the meantime. /// /// **Note:** This won't *necessarily* allow any additional readers to acquire /// locks, since [`RwLock`] is fair and it is possible that a writer is next /// in line. /// /// Returns an RAII guard which will drop this read access of the `RwLock` /// when dropped. /// /// # Examples /// /// ``` /// # use tokio::sync::RwLock; /// # use std::sync::Arc; /// # /// # #[tokio::main] /// # async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// /// let n = lock.write().await; /// /// let cloned_lock = lock.clone(); /// let handle = tokio::spawn(async move { /// *cloned_lock.write().await = 2; /// }); /// /// let n = n.downgrade(); /// assert_eq!(*n, 1, "downgrade is atomic"); /// /// drop(n); /// handle.await.unwrap(); /// assert_eq!(*lock.read().await, 2, "second writer obtained write lock"); /// # } /// ``` /// /// [`RwLock`]: struct@crate::sync::RwLock pub fn downgrade(self) -> RwLockReadGuard<'a, T> { let this = self.skip_drop(); let guard = RwLockReadGuard { s: this.s, data: this.data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }; // Release all but one of the permits held by the write guard let to_release = (this.permits_acquired - 1) as usize; this.s.release(to_release); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); guard } } impl ops::Deref for RwLockWriteGuard<'_, T> { type Target = T; fn deref(&self) -> &T { unsafe { &*self.data } } } impl ops::DerefMut for RwLockWriteGuard<'_, T> { fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.data } } } impl<'a, T: ?Sized> fmt::Debug for RwLockWriteGuard<'a, T> where T: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl<'a, T: ?Sized> fmt::Display for RwLockWriteGuard<'a, T> where T: fmt::Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } impl<'a, T: ?Sized> Drop for RwLockWriteGuard<'a, T> { fn drop(&mut self) { self.s.release(self.permits_acquired as usize); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); } } tokio-1.43.1/src/sync/rwlock/write_guard_mapped.rs000064400000000000000000000153601046102023000203200ustar 00000000000000use crate::sync::batch_semaphore::Semaphore; use std::marker::PhantomData; use std::{fmt, mem, ops}; /// RAII structure used to release the exclusive write access of a lock when /// dropped. /// /// This structure is created by [mapping] an [`RwLockWriteGuard`]. It is a /// separate type from `RwLockWriteGuard` to disallow downgrading a mapped /// guard, since doing so can cause undefined behavior. /// /// [mapping]: method@crate::sync::RwLockWriteGuard::map /// [`RwLockWriteGuard`]: struct@crate::sync::RwLockWriteGuard #[clippy::has_significant_drop] pub struct RwLockMappedWriteGuard<'a, T: ?Sized> { // When changing the fields in this struct, make sure to update the // `skip_drop` method. #[cfg(all(tokio_unstable, feature = "tracing"))] pub(super) resource_span: tracing::Span, pub(super) permits_acquired: u32, pub(super) s: &'a Semaphore, pub(super) data: *mut T, pub(super) marker: PhantomData<&'a mut T>, } #[allow(dead_code)] // Unused fields are still used in Drop. struct Inner<'a, T: ?Sized> { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, permits_acquired: u32, s: &'a Semaphore, data: *mut T, } impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> { fn skip_drop(self) -> Inner<'a, T> { let me = mem::ManuallyDrop::new(self); // SAFETY: This duplicates the values in every field of the guard, then // forgets the originals, so in the end no value is duplicated. Inner { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: unsafe { std::ptr::read(&me.resource_span) }, permits_acquired: me.permits_acquired, s: me.s, data: me.data, } } /// Makes a new `RwLockMappedWriteGuard` for a component of the locked data. /// /// This operation cannot fail as the `RwLockMappedWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be used as /// `RwLockMappedWriteGuard::map(..)`. A method would interfere with methods /// of the same name on the contents of the locked data. /// /// This is an asynchronous version of [`RwLockWriteGuard::map`] from the /// [`parking_lot` crate]. /// /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// { /// let mut mapped = RwLockWriteGuard::map(lock.write().await, |f| &mut f.0); /// *mapped = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn map(mut this: Self, f: F) -> RwLockMappedWriteGuard<'a, U> where F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; let this = this.skip_drop(); RwLockMappedWriteGuard { permits_acquired: this.permits_acquired, s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, } } /// Attempts to make a new [`RwLockMappedWriteGuard`] for a component of /// the locked data. The original guard is returned if the closure returns /// `None`. /// /// This operation cannot fail as the `RwLockMappedWriteGuard` passed in already /// locked the data. /// /// This is an associated function that needs to be /// used as `RwLockMappedWriteGuard::try_map(...)`. A method would interfere /// with methods of the same name on the contents of the locked data. /// /// This is an asynchronous version of [`RwLockWriteGuard::try_map`] from /// the [`parking_lot` crate]. /// /// [`RwLockWriteGuard::try_map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.try_map /// [`parking_lot` crate]: https://crates.io/crates/parking_lot /// /// # Examples /// /// ``` /// use tokio::sync::{RwLock, RwLockWriteGuard}; /// /// #[derive(Debug, Clone, Copy, PartialEq, Eq)] /// struct Foo(u32); /// /// # #[tokio::main] /// # async fn main() { /// let lock = RwLock::new(Foo(1)); /// /// { /// let guard = lock.write().await; /// let mut guard = RwLockWriteGuard::try_map(guard, |f| Some(&mut f.0)).expect("should not fail"); /// *guard = 2; /// } /// /// assert_eq!(Foo(2), *lock.read().await); /// # } /// ``` #[inline] pub fn try_map( mut this: Self, f: F, ) -> Result, Self> where F: FnOnce(&mut T) -> Option<&mut U>, { let data = match f(&mut *this) { Some(data) => data as *mut U, None => return Err(this), }; let this = this.skip_drop(); Ok(RwLockMappedWriteGuard { permits_acquired: this.permits_acquired, s: this.s, data, marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: this.resource_span, }) } // Note: No `downgrade`, `downgrade_map` nor `try_downgrade_map` because they would be unsound, as we're already // potentially been mapped with internal mutability. } impl ops::Deref for RwLockMappedWriteGuard<'_, T> { type Target = T; fn deref(&self) -> &T { unsafe { &*self.data } } } impl ops::DerefMut for RwLockMappedWriteGuard<'_, T> { fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.data } } } impl<'a, T: ?Sized> fmt::Debug for RwLockMappedWriteGuard<'a, T> where T: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl<'a, T: ?Sized> fmt::Display for RwLockMappedWriteGuard<'a, T> where T: fmt::Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt(&**self, f) } } impl<'a, T: ?Sized> Drop for RwLockMappedWriteGuard<'a, T> { fn drop(&mut self) { self.s.release(self.permits_acquired as usize); #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, write_locked.op = "override", ) }); } } tokio-1.43.1/src/sync/rwlock.rs000064400000000000000000001115431046102023000144560ustar 00000000000000use crate::sync::batch_semaphore::{Semaphore, TryAcquireError}; use crate::sync::mutex::TryLockError; #[cfg(all(tokio_unstable, feature = "tracing"))] use crate::util::trace; use std::cell::UnsafeCell; use std::marker; use std::marker::PhantomData; use std::sync::Arc; pub(crate) mod owned_read_guard; pub(crate) mod owned_write_guard; pub(crate) mod owned_write_guard_mapped; pub(crate) mod read_guard; pub(crate) mod write_guard; pub(crate) mod write_guard_mapped; pub(crate) use owned_read_guard::OwnedRwLockReadGuard; pub(crate) use owned_write_guard::OwnedRwLockWriteGuard; pub(crate) use owned_write_guard_mapped::OwnedRwLockMappedWriteGuard; pub(crate) use read_guard::RwLockReadGuard; pub(crate) use write_guard::RwLockWriteGuard; pub(crate) use write_guard_mapped::RwLockMappedWriteGuard; #[cfg(not(loom))] const MAX_READS: u32 = u32::MAX >> 3; #[cfg(loom)] const MAX_READS: u32 = 10; /// An asynchronous reader-writer lock. /// /// This type of lock allows a number of readers or at most one writer at any /// point in time. The write portion of this lock typically allows modification /// of the underlying data (exclusive access) and the read portion of this lock /// typically allows for read-only access (shared access). /// /// In comparison, a [`Mutex`] does not distinguish between readers or writers /// that acquire the lock, therefore causing any tasks waiting for the lock to /// become available to yield. An `RwLock` will allow any number of readers to /// acquire the lock as long as a writer is not holding the lock. /// /// The priority policy of Tokio's read-write lock is _fair_ (or /// [_write-preferring_]), in order to ensure that readers cannot starve /// writers. Fairness is ensured using a first-in, first-out queue for the tasks /// awaiting the lock; if a task that wishes to acquire the write lock is at the /// head of the queue, read locks will not be given out until the write lock has /// been released. This is in contrast to the Rust standard library's /// `std::sync::RwLock`, where the priority policy is dependent on the /// operating system's implementation. /// /// The type parameter `T` represents the data that this lock protects. It is /// required that `T` satisfies [`Send`] to be shared across threads. The RAII guards /// returned from the locking methods implement [`Deref`](trait@std::ops::Deref) /// (and [`DerefMut`](trait@std::ops::DerefMut) /// for the `write` methods) to allow access to the content of the lock. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = RwLock::new(5); /// /// // many reader locks can be held at once /// { /// let r1 = lock.read().await; /// let r2 = lock.read().await; /// assert_eq!(*r1, 5); /// assert_eq!(*r2, 5); /// } // read locks are dropped at this point /// /// // only one write lock may be held, however /// { /// let mut w = lock.write().await; /// *w += 1; /// assert_eq!(*w, 6); /// } // write lock is dropped here /// } /// ``` /// /// [`Mutex`]: struct@super::Mutex /// [`RwLock`]: struct@RwLock /// [`RwLockReadGuard`]: struct@RwLockReadGuard /// [`RwLockWriteGuard`]: struct@RwLockWriteGuard /// [`Send`]: trait@std::marker::Send /// [_write-preferring_]: https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock#Priority_policies pub struct RwLock { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, // maximum number of concurrent readers mr: u32, //semaphore to coordinate read and write access to T s: Semaphore, //inner data T c: UnsafeCell, } #[test] #[cfg(not(loom))] fn bounds() { fn check_send() {} fn check_sync() {} fn check_unpin() {} // This has to take a value, since the async fn's return type is unnameable. fn check_send_sync_val(_t: T) {} check_send::>(); check_sync::>(); check_unpin::>(); check_send::>(); check_sync::>(); check_unpin::>(); check_send::>(); check_sync::>(); check_unpin::>(); check_send::>(); check_sync::>(); check_unpin::>(); check_send::>(); check_sync::>(); check_unpin::>(); check_send::>(); check_sync::>(); check_unpin::>(); check_send::>(); check_sync::>(); check_unpin::>(); let rwlock = Arc::new(RwLock::new(0)); check_send_sync_val(rwlock.read()); check_send_sync_val(Arc::clone(&rwlock).read_owned()); check_send_sync_val(rwlock.write()); check_send_sync_val(Arc::clone(&rwlock).write_owned()); } // As long as T: Send + Sync, it's fine to send and share RwLock between threads. // If T were not Send, sending and sharing a RwLock would be bad, since you can access T through // RwLock. unsafe impl Send for RwLock where T: ?Sized + Send {} unsafe impl Sync for RwLock where T: ?Sized + Send + Sync {} // NB: These impls need to be explicit since we're storing a raw pointer. // Safety: Stores a raw pointer to `T`, so if `T` is `Sync`, the lock guard over // `T` is `Send`. unsafe impl Send for RwLockReadGuard<'_, T> where T: ?Sized + Sync {} unsafe impl Sync for RwLockReadGuard<'_, T> where T: ?Sized + Send + Sync {} // T is required to be `Send` because an OwnedRwLockReadGuard can be used to drop the value held in // the RwLock, unlike RwLockReadGuard. unsafe impl Send for OwnedRwLockReadGuard where T: ?Sized + Send + Sync, U: ?Sized + Sync, { } unsafe impl Sync for OwnedRwLockReadGuard where T: ?Sized + Send + Sync, U: ?Sized + Send + Sync, { } unsafe impl Sync for RwLockWriteGuard<'_, T> where T: ?Sized + Send + Sync {} unsafe impl Sync for OwnedRwLockWriteGuard where T: ?Sized + Send + Sync {} unsafe impl Sync for RwLockMappedWriteGuard<'_, T> where T: ?Sized + Send + Sync {} unsafe impl Sync for OwnedRwLockMappedWriteGuard where T: ?Sized + Send + Sync, U: ?Sized + Send + Sync, { } // Safety: Stores a raw pointer to `T`, so if `T` is `Sync`, the lock guard over // `T` is `Send` - but since this is also provides mutable access, we need to // make sure that `T` is `Send` since its value can be sent across thread // boundaries. unsafe impl Send for RwLockWriteGuard<'_, T> where T: ?Sized + Send + Sync {} unsafe impl Send for OwnedRwLockWriteGuard where T: ?Sized + Send + Sync {} unsafe impl Send for RwLockMappedWriteGuard<'_, T> where T: ?Sized + Send + Sync {} unsafe impl Send for OwnedRwLockMappedWriteGuard where T: ?Sized + Send + Sync, U: ?Sized + Send + Sync, { } impl RwLock { /// Creates a new instance of an `RwLock` which is unlocked. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// let lock = RwLock::new(5); /// ``` #[track_caller] pub fn new(value: T) -> RwLock where T: Sized, { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = std::panic::Location::caller(); let resource_span = tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "RwLock", kind = "Sync", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", max_readers = MAX_READS, ); tracing::trace!( target: "runtime::resource::state_update", write_locked = false, ); tracing::trace!( target: "runtime::resource::state_update", current_readers = 0, ); }); resource_span }; #[cfg(all(tokio_unstable, feature = "tracing"))] let s = resource_span.in_scope(|| Semaphore::new(MAX_READS as usize)); #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let s = Semaphore::new(MAX_READS as usize); RwLock { mr: MAX_READS, c: UnsafeCell::new(value), s, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Creates a new instance of an `RwLock` which is unlocked /// and allows a maximum of `max_reads` concurrent readers. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// let lock = RwLock::with_max_readers(5, 1024); /// ``` /// /// # Panics /// /// Panics if `max_reads` is more than `u32::MAX >> 3`. #[track_caller] pub fn with_max_readers(value: T, max_reads: u32) -> RwLock where T: Sized, { assert!( max_reads <= MAX_READS, "a RwLock may not be created with more than {MAX_READS} readers" ); #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = std::panic::Location::caller(); let resource_span = tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "RwLock", kind = "Sync", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ); resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", max_readers = max_reads, ); tracing::trace!( target: "runtime::resource::state_update", write_locked = false, ); tracing::trace!( target: "runtime::resource::state_update", current_readers = 0, ); }); resource_span }; #[cfg(all(tokio_unstable, feature = "tracing"))] let s = resource_span.in_scope(|| Semaphore::new(max_reads as usize)); #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let s = Semaphore::new(max_reads as usize); RwLock { mr: max_reads, c: UnsafeCell::new(value), s, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Creates a new instance of an `RwLock` which is unlocked. /// /// When using the `tracing` [unstable feature], a `RwLock` created with /// `const_new` will not be instrumented. As such, it will not be visible /// in [`tokio-console`]. Instead, [`RwLock::new`] should be used to create /// an instrumented object if that is needed. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// static LOCK: RwLock = RwLock::const_new(5); /// ``` /// /// [`tokio-console`]: https://github.com/tokio-rs/console /// [unstable feature]: crate#unstable-features #[cfg(not(all(loom, test)))] pub const fn const_new(value: T) -> RwLock where T: Sized, { RwLock { mr: MAX_READS, c: UnsafeCell::new(value), s: Semaphore::const_new(MAX_READS as usize), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Creates a new instance of an `RwLock` which is unlocked /// and allows a maximum of `max_reads` concurrent readers. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// static LOCK: RwLock = RwLock::const_with_max_readers(5, 1024); /// ``` #[cfg(not(all(loom, test)))] pub const fn const_with_max_readers(value: T, max_reads: u32) -> RwLock where T: Sized, { assert!(max_reads <= MAX_READS); RwLock { mr: max_reads, c: UnsafeCell::new(value), s: Semaphore::const_new(max_reads as usize), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Locks this `RwLock` with shared read access, causing the current task /// to yield until the lock has been acquired. /// /// The calling task will yield until there are no writers which hold the /// lock. There may be other readers inside the lock when the task resumes. /// /// Note that under the priority policy of [`RwLock`], read locks are not /// granted until prior write locks, to prevent starvation. Therefore /// deadlock may occur if a read lock is held by the current task, a write /// lock attempt is made, and then a subsequent read lock attempt is made /// by the current task. /// /// Returns an RAII guard which will drop this read access of the `RwLock` /// when dropped. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute locks in the order they /// were requested. Cancelling a call to `read` makes you lose your place in /// the queue. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// let c_lock = lock.clone(); /// /// let n = lock.read().await; /// assert_eq!(*n, 1); /// /// tokio::spawn(async move { /// // While main has an active read lock, we acquire one too. /// let r = c_lock.read().await; /// assert_eq!(*r, 1); /// }).await.expect("The spawned task has panicked"); /// /// // Drop the guard after the spawned task finishes. /// drop(n); /// } /// ``` pub async fn read(&self) -> RwLockReadGuard<'_, T> { let acquire_fut = async { self.s.acquire(1).await.unwrap_or_else(|_| { // The semaphore was closed. but, we never explicitly close it, and we have a // handle to it through the Arc, which means that this can never happen. unreachable!() }); RwLockReadGuard { s: &self.s, data: self.c.get(), marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), } }; #[cfg(all(tokio_unstable, feature = "tracing"))] let acquire_fut = trace::async_op( move || acquire_fut, self.resource_span.clone(), "RwLock::read", "poll", false, ); #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing let guard = acquire_fut.await; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); guard } /// Blockingly locks this `RwLock` with shared read access. /// /// This method is intended for use cases where you /// need to use this rwlock in asynchronous code as well as in synchronous code. /// /// Returns an RAII guard which will drop the read access of this `RwLock` when dropped. /// /// # Panics /// /// This function panics if called within an asynchronous execution context. /// /// - If you find yourself in an asynchronous execution context and needing /// to call some (synchronous) function which performs one of these /// `blocking_` operations, then consider wrapping that call inside /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] /// (or [`block_in_place()`][crate::task::block_in_place]). /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let rwlock = Arc::new(RwLock::new(1)); /// let mut write_lock = rwlock.write().await; /// /// let blocking_task = tokio::task::spawn_blocking({ /// let rwlock = Arc::clone(&rwlock); /// move || { /// // This shall block until the `write_lock` is released. /// let read_lock = rwlock.blocking_read(); /// assert_eq!(*read_lock, 0); /// } /// }); /// /// *write_lock -= 1; /// drop(write_lock); // release the lock. /// /// // Await the completion of the blocking task. /// blocking_task.await.unwrap(); /// /// // Assert uncontended. /// assert!(rwlock.try_write().is_ok()); /// } /// ``` #[track_caller] #[cfg(feature = "sync")] pub fn blocking_read(&self) -> RwLockReadGuard<'_, T> { crate::future::block_on(self.read()) } /// Locks this `RwLock` with shared read access, causing the current task /// to yield until the lock has been acquired. /// /// The calling task will yield until there are no writers which hold the /// lock. There may be other readers inside the lock when the task resumes. /// /// This method is identical to [`RwLock::read`], except that the returned /// guard references the `RwLock` with an [`Arc`] rather than by borrowing /// it. Therefore, the `RwLock` must be wrapped in an `Arc` to call this /// method, and the guard will live for the `'static` lifetime, as it keeps /// the `RwLock` alive by holding an `Arc`. /// /// Note that under the priority policy of [`RwLock`], read locks are not /// granted until prior write locks, to prevent starvation. Therefore /// deadlock may occur if a read lock is held by the current task, a write /// lock attempt is made, and then a subsequent read lock attempt is made /// by the current task. /// /// Returns an RAII guard which will drop this read access of the `RwLock` /// when dropped. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute locks in the order they /// were requested. Cancelling a call to `read_owned` makes you lose your /// place in the queue. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// let c_lock = lock.clone(); /// /// let n = lock.read_owned().await; /// assert_eq!(*n, 1); /// /// tokio::spawn(async move { /// // While main has an active read lock, we acquire one too. /// let r = c_lock.read_owned().await; /// assert_eq!(*r, 1); /// }).await.expect("The spawned task has panicked"); /// /// // Drop the guard after the spawned task finishes. /// drop(n); ///} /// ``` pub async fn read_owned(self: Arc) -> OwnedRwLockReadGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = self.resource_span.clone(); let acquire_fut = async { self.s.acquire(1).await.unwrap_or_else(|_| { // The semaphore was closed. but, we never explicitly close it, and we have a // handle to it through the Arc, which means that this can never happen. unreachable!() }); OwnedRwLockReadGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), data: self.c.get(), lock: self, _p: PhantomData, } }; #[cfg(all(tokio_unstable, feature = "tracing"))] let acquire_fut = trace::async_op( move || acquire_fut, resource_span, "RwLock::read_owned", "poll", false, ); #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing let guard = acquire_fut.await; #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); guard } /// Attempts to acquire this `RwLock` with shared read access. /// /// If the access couldn't be acquired immediately, returns [`TryLockError`]. /// Otherwise, an RAII guard is returned which will release read access /// when dropped. /// /// [`TryLockError`]: TryLockError /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// let c_lock = lock.clone(); /// /// let v = lock.try_read().unwrap(); /// assert_eq!(*v, 1); /// /// tokio::spawn(async move { /// // While main has an active read lock, we acquire one too. /// let n = c_lock.read().await; /// assert_eq!(*n, 1); /// }).await.expect("The spawned task has panicked"); /// /// // Drop the guard when spawned task finishes. /// drop(v); /// } /// ``` pub fn try_read(&self) -> Result, TryLockError> { match self.s.try_acquire(1) { Ok(permit) => permit, Err(TryAcquireError::NoPermits) => return Err(TryLockError(())), Err(TryAcquireError::Closed) => unreachable!(), } let guard = RwLockReadGuard { s: &self.s, data: self.c.get(), marker: marker::PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), }; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); Ok(guard) } /// Attempts to acquire this `RwLock` with shared read access. /// /// If the access couldn't be acquired immediately, returns [`TryLockError`]. /// Otherwise, an RAII guard is returned which will release read access /// when dropped. /// /// This method is identical to [`RwLock::try_read`], except that the /// returned guard references the `RwLock` with an [`Arc`] rather than by /// borrowing it. Therefore, the `RwLock` must be wrapped in an `Arc` to /// call this method, and the guard will live for the `'static` lifetime, /// as it keeps the `RwLock` alive by holding an `Arc`. /// /// [`TryLockError`]: TryLockError /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// let c_lock = lock.clone(); /// /// let v = lock.try_read_owned().unwrap(); /// assert_eq!(*v, 1); /// /// tokio::spawn(async move { /// // While main has an active read lock, we acquire one too. /// let n = c_lock.read_owned().await; /// assert_eq!(*n, 1); /// }).await.expect("The spawned task has panicked"); /// /// // Drop the guard when spawned task finishes. /// drop(v); /// } /// ``` pub fn try_read_owned(self: Arc) -> Result, TryLockError> { match self.s.try_acquire(1) { Ok(permit) => permit, Err(TryAcquireError::NoPermits) => return Err(TryLockError(())), Err(TryAcquireError::Closed) => unreachable!(), } let guard = OwnedRwLockReadGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), data: self.c.get(), lock: self, _p: PhantomData, }; #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, current_readers.op = "add", ) }); Ok(guard) } /// Locks this `RwLock` with exclusive write access, causing the current /// task to yield until the lock has been acquired. /// /// The calling task will yield while other writers or readers currently /// have access to the lock. /// /// Returns an RAII guard which will drop the write access of this `RwLock` /// when dropped. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute locks in the order they /// were requested. Cancelling a call to `write` makes you lose your place /// in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = RwLock::new(1); /// /// let mut n = lock.write().await; /// *n = 2; ///} /// ``` pub async fn write(&self) -> RwLockWriteGuard<'_, T> { let acquire_fut = async { self.s.acquire(self.mr as usize).await.unwrap_or_else(|_| { // The semaphore was closed. but, we never explicitly close it, and we have a // handle to it through the Arc, which means that this can never happen. unreachable!() }); RwLockWriteGuard { permits_acquired: self.mr, s: &self.s, data: self.c.get(), marker: marker::PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), } }; #[cfg(all(tokio_unstable, feature = "tracing"))] let acquire_fut = trace::async_op( move || acquire_fut, self.resource_span.clone(), "RwLock::write", "poll", false, ); #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing let guard = acquire_fut.await; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = true, write_locked.op = "override", ) }); guard } /// Blockingly locks this `RwLock` with exclusive write access. /// /// This method is intended for use cases where you /// need to use this rwlock in asynchronous code as well as in synchronous code. /// /// Returns an RAII guard which will drop the write access of this `RwLock` when dropped. /// /// # Panics /// /// This function panics if called within an asynchronous execution context. /// /// - If you find yourself in an asynchronous execution context and needing /// to call some (synchronous) function which performs one of these /// `blocking_` operations, then consider wrapping that call inside /// [`spawn_blocking()`][crate::runtime::Handle::spawn_blocking] /// (or [`block_in_place()`][crate::task::block_in_place]). /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::{sync::RwLock}; /// /// #[tokio::main] /// async fn main() { /// let rwlock = Arc::new(RwLock::new(1)); /// let read_lock = rwlock.read().await; /// /// let blocking_task = tokio::task::spawn_blocking({ /// let rwlock = Arc::clone(&rwlock); /// move || { /// // This shall block until the `read_lock` is released. /// let mut write_lock = rwlock.blocking_write(); /// *write_lock = 2; /// } /// }); /// /// assert_eq!(*read_lock, 1); /// // Release the last outstanding read lock. /// drop(read_lock); /// /// // Await the completion of the blocking task. /// blocking_task.await.unwrap(); /// /// // Assert uncontended. /// let read_lock = rwlock.try_read().unwrap(); /// assert_eq!(*read_lock, 2); /// } /// ``` #[track_caller] #[cfg(feature = "sync")] pub fn blocking_write(&self) -> RwLockWriteGuard<'_, T> { crate::future::block_on(self.write()) } /// Locks this `RwLock` with exclusive write access, causing the current /// task to yield until the lock has been acquired. /// /// The calling task will yield while other writers or readers currently /// have access to the lock. /// /// This method is identical to [`RwLock::write`], except that the returned /// guard references the `RwLock` with an [`Arc`] rather than by borrowing /// it. Therefore, the `RwLock` must be wrapped in an `Arc` to call this /// method, and the guard will live for the `'static` lifetime, as it keeps /// the `RwLock` alive by holding an `Arc`. /// /// Returns an RAII guard which will drop the write access of this `RwLock` /// when dropped. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute locks in the order they /// were requested. Cancelling a call to `write_owned` makes you lose your /// place in the queue. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let lock = Arc::new(RwLock::new(1)); /// /// let mut n = lock.write_owned().await; /// *n = 2; ///} /// ``` pub async fn write_owned(self: Arc) -> OwnedRwLockWriteGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = self.resource_span.clone(); let acquire_fut = async { self.s.acquire(self.mr as usize).await.unwrap_or_else(|_| { // The semaphore was closed. but, we never explicitly close it, and we have a // handle to it through the Arc, which means that this can never happen. unreachable!() }); OwnedRwLockWriteGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), permits_acquired: self.mr, data: self.c.get(), lock: self, _p: PhantomData, } }; #[cfg(all(tokio_unstable, feature = "tracing"))] let acquire_fut = trace::async_op( move || acquire_fut, resource_span, "RwLock::write_owned", "poll", false, ); #[allow(clippy::let_and_return)] // this lint triggers when disabling tracing let guard = acquire_fut.await; #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = true, write_locked.op = "override", ) }); guard } /// Attempts to acquire this `RwLock` with exclusive write access. /// /// If the access couldn't be acquired immediately, returns [`TryLockError`]. /// Otherwise, an RAII guard is returned which will release write access /// when dropped. /// /// [`TryLockError`]: TryLockError /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let rw = RwLock::new(1); /// /// let v = rw.read().await; /// assert_eq!(*v, 1); /// /// assert!(rw.try_write().is_err()); /// } /// ``` pub fn try_write(&self) -> Result, TryLockError> { match self.s.try_acquire(self.mr as usize) { Ok(permit) => permit, Err(TryAcquireError::NoPermits) => return Err(TryLockError(())), Err(TryAcquireError::Closed) => unreachable!(), } let guard = RwLockWriteGuard { permits_acquired: self.mr, s: &self.s, data: self.c.get(), marker: marker::PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), }; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = true, write_locked.op = "override", ) }); Ok(guard) } /// Attempts to acquire this `RwLock` with exclusive write access. /// /// If the access couldn't be acquired immediately, returns [`TryLockError`]. /// Otherwise, an RAII guard is returned which will release write access /// when dropped. /// /// This method is identical to [`RwLock::try_write`], except that the /// returned guard references the `RwLock` with an [`Arc`] rather than by /// borrowing it. Therefore, the `RwLock` must be wrapped in an `Arc` to /// call this method, and the guard will live for the `'static` lifetime, /// as it keeps the `RwLock` alive by holding an `Arc`. /// /// [`TryLockError`]: TryLockError /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::RwLock; /// /// #[tokio::main] /// async fn main() { /// let rw = Arc::new(RwLock::new(1)); /// /// let v = Arc::clone(&rw).read_owned().await; /// assert_eq!(*v, 1); /// /// assert!(rw.try_write_owned().is_err()); /// } /// ``` pub fn try_write_owned(self: Arc) -> Result, TryLockError> { match self.s.try_acquire(self.mr as usize) { Ok(permit) => permit, Err(TryAcquireError::NoPermits) => return Err(TryLockError(())), Err(TryAcquireError::Closed) => unreachable!(), } let guard = OwnedRwLockWriteGuard { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: self.resource_span.clone(), permits_acquired: self.mr, data: self.c.get(), lock: self, _p: PhantomData, }; #[cfg(all(tokio_unstable, feature = "tracing"))] guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = true, write_locked.op = "override", ) }); Ok(guard) } /// Returns a mutable reference to the underlying data. /// /// Since this call borrows the `RwLock` mutably, no actual locking needs to /// take place -- the mutable borrow statically guarantees no locks exist. /// /// # Examples /// /// ``` /// use tokio::sync::RwLock; /// /// fn main() { /// let mut lock = RwLock::new(1); /// /// let n = lock.get_mut(); /// *n = 2; /// } /// ``` pub fn get_mut(&mut self) -> &mut T { unsafe { // Safety: This is https://github.com/rust-lang/rust/pull/76936 &mut *self.c.get() } } /// Consumes the lock, returning the underlying data. pub fn into_inner(self) -> T where T: Sized, { self.c.into_inner() } } impl From for RwLock { fn from(s: T) -> Self { Self::new(s) } } impl Default for RwLock where T: Default, { fn default() -> Self { Self::new(T::default()) } } impl std::fmt::Debug for RwLock where T: std::fmt::Debug, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let mut d = f.debug_struct("RwLock"); match self.try_read() { Ok(inner) => d.field("data", &&*inner), Err(_) => d.field("data", &format_args!("")), }; d.finish() } } tokio-1.43.1/src/sync/semaphore.rs000064400000000000000000001174061046102023000151440ustar 00000000000000use super::batch_semaphore as ll; // low level implementation use super::{AcquireError, TryAcquireError}; #[cfg(all(tokio_unstable, feature = "tracing"))] use crate::util::trace; use std::sync::Arc; /// Counting semaphore performing asynchronous permit acquisition. /// /// A semaphore maintains a set of permits. Permits are used to synchronize /// access to a shared resource. A semaphore differs from a mutex in that it /// can allow more than one concurrent caller to access the shared resource at a /// time. /// /// When `acquire` is called and the semaphore has remaining permits, the /// function immediately returns a permit. However, if no remaining permits are /// available, `acquire` (asynchronously) waits until an outstanding permit is /// dropped. At this point, the freed permit is assigned to the caller. /// /// This `Semaphore` is fair, which means that permits are given out in the order /// they were requested. This fairness is also applied when `acquire_many` gets /// involved, so if a call to `acquire_many` at the front of the queue requests /// more permits than currently available, this can prevent a call to `acquire` /// from completing, even if the semaphore has enough permits complete the call /// to `acquire`. /// /// To use the `Semaphore` in a poll function, you can use the [`PollSemaphore`] /// utility. /// /// # Examples /// /// Basic usage: /// /// ``` /// use tokio::sync::{Semaphore, TryAcquireError}; /// /// #[tokio::main] /// async fn main() { /// let semaphore = Semaphore::new(3); /// /// let a_permit = semaphore.acquire().await.unwrap(); /// let two_permits = semaphore.acquire_many(2).await.unwrap(); /// /// assert_eq!(semaphore.available_permits(), 0); /// /// let permit_attempt = semaphore.try_acquire(); /// assert_eq!(permit_attempt.err(), Some(TryAcquireError::NoPermits)); /// } /// ``` /// /// ## Limit the number of simultaneously opened files in your program /// /// Most operating systems have limits on the number of open file /// handles. Even in systems without explicit limits, resource constraints /// implicitly set an upper bound on the number of open files. If your /// program attempts to open a large number of files and exceeds this /// limit, it will result in an error. /// /// This example uses a Semaphore with 100 permits. By acquiring a permit from /// the Semaphore before accessing a file, you ensure that your program opens /// no more than 100 files at a time. When trying to open the 101st /// file, the program will wait until a permit becomes available before /// proceeding to open another file. /// ``` /// use std::io::Result; /// use tokio::fs::File; /// use tokio::sync::Semaphore; /// use tokio::io::AsyncWriteExt; /// /// static PERMITS: Semaphore = Semaphore::const_new(100); /// /// async fn write_to_file(message: &[u8]) -> Result<()> { /// let _permit = PERMITS.acquire().await.unwrap(); /// let mut buffer = File::create("example.txt").await?; /// buffer.write_all(message).await?; /// Ok(()) // Permit goes out of scope here, and is available again for acquisition /// } /// ``` /// /// ## Limit the number of outgoing requests being sent at the same time /// /// In some scenarios, it might be required to limit the number of outgoing /// requests being sent in parallel. This could be due to limits of a consumed /// API or the network resources of the system the application is running on. /// /// This example uses an `Arc` with 10 permits. Each task spawned is /// given a reference to the semaphore by cloning the `Arc`. Before /// a task sends a request, it must acquire a permit from the semaphore by /// calling [`Semaphore::acquire`]. This ensures that at most 10 requests are /// sent in parallel at any given time. After a task has sent a request, it /// drops the permit to allow other tasks to send requests. /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// #[tokio::main] /// async fn main() { /// // Define maximum number of parallel requests. /// let semaphore = Arc::new(Semaphore::new(10)); /// // Spawn many tasks that will send requests. /// let mut jhs = Vec::new(); /// for task_id in 0..100 { /// let semaphore = semaphore.clone(); /// let jh = tokio::spawn(async move { /// // Acquire permit before sending request. /// let _permit = semaphore.acquire().await.unwrap(); /// // Send the request. /// let response = send_request(task_id).await; /// // Drop the permit after the request has been sent. /// drop(_permit); /// // Handle response. /// // ... /// /// response /// }); /// jhs.push(jh); /// } /// // Collect responses from tasks. /// let mut responses = Vec::new(); /// for jh in jhs { /// let response = jh.await.unwrap(); /// responses.push(response); /// } /// // Process responses. /// // ... /// } /// # async fn send_request(task_id: usize) { /// # // Send request. /// # } /// ``` /// /// ## Limit the number of incoming requests being handled at the same time /// /// Similar to limiting the number of simultaneously opened files, network handles /// are a limited resource. Allowing an unbounded amount of requests to be processed /// could result in a denial-of-service, among many other issues. /// /// This example uses an `Arc` instead of a global variable. /// To limit the number of requests that can be processed at the time, /// we acquire a permit for each task before spawning it. Once acquired, /// a new task is spawned; and once finished, the permit is dropped inside /// of the task to allow others to spawn. Permits must be acquired via /// [`Semaphore::acquire_owned`] to be movable across the task boundary. /// (Since our semaphore is not a global variable — if it was, then `acquire` would be enough.) /// /// ```no_run /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// use tokio::net::TcpListener; /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { /// let semaphore = Arc::new(Semaphore::new(3)); /// let listener = TcpListener::bind("127.0.0.1:8080").await?; /// /// loop { /// // Acquire permit before accepting the next socket. /// // /// // We use `acquire_owned` so that we can move `permit` into /// // other tasks. /// let permit = semaphore.clone().acquire_owned().await.unwrap(); /// let (mut socket, _) = listener.accept().await?; /// /// tokio::spawn(async move { /// // Do work using the socket. /// handle_connection(&mut socket).await; /// // Drop socket while the permit is still live. /// drop(socket); /// // Drop the permit, so more tasks can be created. /// drop(permit); /// }); /// } /// } /// # async fn handle_connection(_socket: &mut tokio::net::TcpStream) { /// # // Do work /// # } /// ``` /// /// ## Prevent tests from running in parallel /// /// By default, Rust runs tests in the same file in parallel. However, in some /// cases, running two tests in parallel may lead to problems. For example, this /// can happen when tests use the same database. /// /// Consider the following scenario: /// 1. `test_insert`: Inserts a key-value pair into the database, then retrieves /// the value using the same key to verify the insertion. /// 2. `test_update`: Inserts a key, then updates the key to a new value and /// verifies that the value has been accurately updated. /// 3. `test_others`: A third test that doesn't modify the database state. It /// can run in parallel with the other tests. /// /// In this example, `test_insert` and `test_update` need to run in sequence to /// work, but it doesn't matter which test runs first. We can leverage a /// semaphore with a single permit to address this challenge. /// /// ``` /// # use tokio::sync::Mutex; /// # use std::collections::BTreeMap; /// # struct Database { /// # map: Mutex>, /// # } /// # impl Database { /// # pub const fn setup() -> Database { /// # Database { /// # map: Mutex::const_new(BTreeMap::new()), /// # } /// # } /// # pub async fn insert(&self, key: &str, value: i32) { /// # self.map.lock().await.insert(key.to_string(), value); /// # } /// # pub async fn update(&self, key: &str, value: i32) { /// # self.map.lock().await /// # .entry(key.to_string()) /// # .and_modify(|origin| *origin = value); /// # } /// # pub async fn delete(&self, key: &str) { /// # self.map.lock().await.remove(key); /// # } /// # pub async fn get(&self, key: &str) -> i32 { /// # *self.map.lock().await.get(key).unwrap() /// # } /// # } /// use tokio::sync::Semaphore; /// /// // Initialize a static semaphore with only one permit, which is used to /// // prevent test_insert and test_update from running in parallel. /// static PERMIT: Semaphore = Semaphore::const_new(1); /// /// // Initialize the database that will be used by the subsequent tests. /// static DB: Database = Database::setup(); /// /// #[tokio::test] /// # async fn fake_test_insert() {} /// async fn test_insert() { /// // Acquire permit before proceeding. Since the semaphore has only one permit, /// // the test will wait if the permit is already acquired by other tests. /// let permit = PERMIT.acquire().await.unwrap(); /// /// // Do the actual test stuff with database /// /// // Insert a key-value pair to database /// let (key, value) = ("name", 0); /// DB.insert(key, value).await; /// /// // Verify that the value has been inserted correctly. /// assert_eq!(DB.get(key).await, value); /// /// // Undo the insertion, so the database is empty at the end of the test. /// DB.delete(key).await; /// /// // Drop permit. This allows the other test to start running. /// drop(permit); /// } /// /// #[tokio::test] /// # async fn fake_test_update() {} /// async fn test_update() { /// // Acquire permit before proceeding. Since the semaphore has only one permit, /// // the test will wait if the permit is already acquired by other tests. /// let permit = PERMIT.acquire().await.unwrap(); /// /// // Do the same insert. /// let (key, value) = ("name", 0); /// DB.insert(key, value).await; /// /// // Update the existing value with a new one. /// let new_value = 1; /// DB.update(key, new_value).await; /// /// // Verify that the value has been updated correctly. /// assert_eq!(DB.get(key).await, new_value); /// /// // Undo any modificattion. /// DB.delete(key).await; /// /// // Drop permit. This allows the other test to start running. /// drop(permit); /// } /// /// #[tokio::test] /// # async fn fake_test_others() {} /// async fn test_others() { /// // This test can run in parallel with test_insert and test_update, /// // so it does not use PERMIT. /// } /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() { /// # test_insert().await; /// # test_update().await; /// # test_others().await; /// # } /// ``` /// /// ## Rate limiting using a token bucket /// /// This example showcases the [`add_permits`] and [`SemaphorePermit::forget`] methods. /// /// Many applications and systems have constraints on the rate at which certain /// operations should occur. Exceeding this rate can result in suboptimal /// performance or even errors. /// /// This example implements rate limiting using a [token bucket]. A token bucket is a form of rate /// limiting that doesn't kick in immediately, to allow for short bursts of incoming requests that /// arrive at the same time. /// /// With a token bucket, each incoming request consumes a token, and the tokens are refilled at a /// certain rate that defines the rate limit. When a burst of requests arrives, tokens are /// immediately given out until the bucket is empty. Once the bucket is empty, requests will have to /// wait for new tokens to be added. /// /// Unlike the example that limits how many requests can be handled at the same time, we do not add /// tokens back when we finish handling a request. Instead, tokens are added only by a timer task. /// /// Note that this implementation is suboptimal when the duration is small, because it consumes a /// lot of cpu constantly looping and sleeping. /// /// [token bucket]: https://en.wikipedia.org/wiki/Token_bucket /// [`add_permits`]: crate::sync::Semaphore::add_permits /// [`SemaphorePermit::forget`]: crate::sync::SemaphorePermit::forget /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// use tokio::time::{interval, Duration}; /// /// struct TokenBucket { /// sem: Arc, /// jh: tokio::task::JoinHandle<()>, /// } /// /// impl TokenBucket { /// fn new(duration: Duration, capacity: usize) -> Self { /// let sem = Arc::new(Semaphore::new(capacity)); /// /// // refills the tokens at the end of each interval /// let jh = tokio::spawn({ /// let sem = sem.clone(); /// let mut interval = interval(duration); /// interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); /// /// async move { /// loop { /// interval.tick().await; /// /// if sem.available_permits() < capacity { /// sem.add_permits(1); /// } /// } /// } /// }); /// /// Self { jh, sem } /// } /// /// async fn acquire(&self) { /// // This can return an error if the semaphore is closed, but we /// // never close it, so this error can never happen. /// let permit = self.sem.acquire().await.unwrap(); /// // To avoid releasing the permit back to the semaphore, we use /// // the `SemaphorePermit::forget` method. /// permit.forget(); /// } /// } /// /// impl Drop for TokenBucket { /// fn drop(&mut self) { /// // Kill the background task so it stops taking up resources when we /// // don't need it anymore. /// self.jh.abort(); /// } /// } /// /// #[tokio::main] /// # async fn _hidden() {} /// # #[tokio::main(flavor = "current_thread", start_paused = true)] /// async fn main() { /// let capacity = 5; /// let update_interval = Duration::from_secs_f32(1.0 / capacity as f32); /// let bucket = TokenBucket::new(update_interval, capacity); /// /// for _ in 0..5 { /// bucket.acquire().await; /// /// // do the operation /// } /// } /// ``` /// /// [`PollSemaphore`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSemaphore.html /// [`Semaphore::acquire_owned`]: crate::sync::Semaphore::acquire_owned #[derive(Debug)] pub struct Semaphore { /// The low level semaphore ll_sem: ll::Semaphore, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, } /// A permit from the semaphore. /// /// This type is created by the [`acquire`] method. /// /// [`acquire`]: crate::sync::Semaphore::acquire() #[must_use] #[clippy::has_significant_drop] #[derive(Debug)] pub struct SemaphorePermit<'a> { sem: &'a Semaphore, permits: u32, } /// An owned permit from the semaphore. /// /// This type is created by the [`acquire_owned`] method. /// /// [`acquire_owned`]: crate::sync::Semaphore::acquire_owned() #[must_use] #[clippy::has_significant_drop] #[derive(Debug)] pub struct OwnedSemaphorePermit { sem: Arc, permits: u32, } #[test] #[cfg(not(loom))] fn bounds() { fn check_unpin() {} // This has to take a value, since the async fn's return type is unnameable. fn check_send_sync_val(_t: T) {} fn check_send_sync() {} check_unpin::(); check_unpin::>(); check_send_sync::(); let semaphore = Semaphore::new(0); check_send_sync_val(semaphore.acquire()); } impl Semaphore { /// The maximum number of permits which a semaphore can hold. It is `usize::MAX >> 3`. /// /// Exceeding this limit typically results in a panic. pub const MAX_PERMITS: usize = super::batch_semaphore::Semaphore::MAX_PERMITS; /// Creates a new semaphore with the initial number of permits. /// /// Panics if `permits` exceeds [`Semaphore::MAX_PERMITS`]. #[track_caller] pub fn new(permits: usize) -> Self { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = std::panic::Location::caller(); tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Semaphore", kind = "Sync", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), inherits_child_attrs = true, ) }; #[cfg(all(tokio_unstable, feature = "tracing"))] let ll_sem = resource_span.in_scope(|| ll::Semaphore::new(permits)); #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let ll_sem = ll::Semaphore::new(permits); Self { ll_sem, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Creates a new semaphore with the initial number of permits. /// /// When using the `tracing` [unstable feature], a `Semaphore` created with /// `const_new` will not be instrumented. As such, it will not be visible /// in [`tokio-console`]. Instead, [`Semaphore::new`] should be used to /// create an instrumented object if that is needed. /// /// # Examples /// /// ``` /// use tokio::sync::Semaphore; /// /// static SEM: Semaphore = Semaphore::const_new(10); /// ``` /// /// [`tokio-console`]: https://github.com/tokio-rs/console /// [unstable feature]: crate#unstable-features #[cfg(not(all(loom, test)))] pub const fn const_new(permits: usize) -> Self { Self { ll_sem: ll::Semaphore::const_new(permits), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Creates a new closed semaphore with 0 permits. pub(crate) fn new_closed() -> Self { Self { ll_sem: ll::Semaphore::new_closed(), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Creates a new closed semaphore with 0 permits. #[cfg(not(all(loom, test)))] pub(crate) const fn const_new_closed() -> Self { Self { ll_sem: ll::Semaphore::const_new_closed(), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), } } /// Returns the current number of available permits. pub fn available_permits(&self) -> usize { self.ll_sem.available_permits() } /// Adds `n` new permits to the semaphore. /// /// The maximum number of permits is [`Semaphore::MAX_PERMITS`], and this function will panic if the limit is exceeded. pub fn add_permits(&self, n: usize) { self.ll_sem.release(n); } /// Decrease a semaphore's permits by a maximum of `n`. /// /// If there are insufficient permits and it's not possible to reduce by `n`, /// return the number of permits that were actually reduced. pub fn forget_permits(&self, n: usize) -> usize { self.ll_sem.forget_permits(n) } /// Acquires a permit from the semaphore. /// /// If the semaphore has been closed, this returns an [`AcquireError`]. /// Otherwise, this returns a [`SemaphorePermit`] representing the /// acquired permit. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute permits in the order they /// were requested. Cancelling a call to `acquire` makes you lose your place /// in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::Semaphore; /// /// #[tokio::main] /// async fn main() { /// let semaphore = Semaphore::new(2); /// /// let permit_1 = semaphore.acquire().await.unwrap(); /// assert_eq!(semaphore.available_permits(), 1); /// /// let permit_2 = semaphore.acquire().await.unwrap(); /// assert_eq!(semaphore.available_permits(), 0); /// /// drop(permit_1); /// assert_eq!(semaphore.available_permits(), 1); /// } /// ``` /// /// [`AcquireError`]: crate::sync::AcquireError /// [`SemaphorePermit`]: crate::sync::SemaphorePermit pub async fn acquire(&self) -> Result, AcquireError> { #[cfg(all(tokio_unstable, feature = "tracing"))] let inner = trace::async_op( || self.ll_sem.acquire(1), self.resource_span.clone(), "Semaphore::acquire", "poll", true, ); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let inner = self.ll_sem.acquire(1); inner.await?; Ok(SemaphorePermit { sem: self, permits: 1, }) } /// Acquires `n` permits from the semaphore. /// /// If the semaphore has been closed, this returns an [`AcquireError`]. /// Otherwise, this returns a [`SemaphorePermit`] representing the /// acquired permits. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute permits in the order they /// were requested. Cancelling a call to `acquire_many` makes you lose your /// place in the queue. /// /// # Examples /// /// ``` /// use tokio::sync::Semaphore; /// /// #[tokio::main] /// async fn main() { /// let semaphore = Semaphore::new(5); /// /// let permit = semaphore.acquire_many(3).await.unwrap(); /// assert_eq!(semaphore.available_permits(), 2); /// } /// ``` /// /// [`AcquireError`]: crate::sync::AcquireError /// [`SemaphorePermit`]: crate::sync::SemaphorePermit pub async fn acquire_many(&self, n: u32) -> Result, AcquireError> { #[cfg(all(tokio_unstable, feature = "tracing"))] trace::async_op( || self.ll_sem.acquire(n as usize), self.resource_span.clone(), "Semaphore::acquire_many", "poll", true, ) .await?; #[cfg(not(all(tokio_unstable, feature = "tracing")))] self.ll_sem.acquire(n as usize).await?; Ok(SemaphorePermit { sem: self, permits: n, }) } /// Tries to acquire a permit from the semaphore. /// /// If the semaphore has been closed, this returns a [`TryAcquireError::Closed`] /// and a [`TryAcquireError::NoPermits`] if there are no permits left. Otherwise, /// this returns a [`SemaphorePermit`] representing the acquired permits. /// /// # Examples /// /// ``` /// use tokio::sync::{Semaphore, TryAcquireError}; /// /// # fn main() { /// let semaphore = Semaphore::new(2); /// /// let permit_1 = semaphore.try_acquire().unwrap(); /// assert_eq!(semaphore.available_permits(), 1); /// /// let permit_2 = semaphore.try_acquire().unwrap(); /// assert_eq!(semaphore.available_permits(), 0); /// /// let permit_3 = semaphore.try_acquire(); /// assert_eq!(permit_3.err(), Some(TryAcquireError::NoPermits)); /// # } /// ``` /// /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits /// [`SemaphorePermit`]: crate::sync::SemaphorePermit pub fn try_acquire(&self) -> Result, TryAcquireError> { match self.ll_sem.try_acquire(1) { Ok(()) => Ok(SemaphorePermit { sem: self, permits: 1, }), Err(e) => Err(e), } } /// Tries to acquire `n` permits from the semaphore. /// /// If the semaphore has been closed, this returns a [`TryAcquireError::Closed`] /// and a [`TryAcquireError::NoPermits`] if there are not enough permits left. /// Otherwise, this returns a [`SemaphorePermit`] representing the acquired permits. /// /// # Examples /// /// ``` /// use tokio::sync::{Semaphore, TryAcquireError}; /// /// # fn main() { /// let semaphore = Semaphore::new(4); /// /// let permit_1 = semaphore.try_acquire_many(3).unwrap(); /// assert_eq!(semaphore.available_permits(), 1); /// /// let permit_2 = semaphore.try_acquire_many(2); /// assert_eq!(permit_2.err(), Some(TryAcquireError::NoPermits)); /// # } /// ``` /// /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits /// [`SemaphorePermit`]: crate::sync::SemaphorePermit pub fn try_acquire_many(&self, n: u32) -> Result, TryAcquireError> { match self.ll_sem.try_acquire(n as usize) { Ok(()) => Ok(SemaphorePermit { sem: self, permits: n, }), Err(e) => Err(e), } } /// Acquires a permit from the semaphore. /// /// The semaphore must be wrapped in an [`Arc`] to call this method. /// If the semaphore has been closed, this returns an [`AcquireError`]. /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the /// acquired permit. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute permits in the order they /// were requested. Cancelling a call to `acquire_owned` makes you lose your /// place in the queue. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// #[tokio::main] /// async fn main() { /// let semaphore = Arc::new(Semaphore::new(3)); /// let mut join_handles = Vec::new(); /// /// for _ in 0..5 { /// let permit = semaphore.clone().acquire_owned().await.unwrap(); /// join_handles.push(tokio::spawn(async move { /// // perform task... /// // explicitly own `permit` in the task /// drop(permit); /// })); /// } /// /// for handle in join_handles { /// handle.await.unwrap(); /// } /// } /// ``` /// /// [`Arc`]: std::sync::Arc /// [`AcquireError`]: crate::sync::AcquireError /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit pub async fn acquire_owned(self: Arc) -> Result { #[cfg(all(tokio_unstable, feature = "tracing"))] let inner = trace::async_op( || self.ll_sem.acquire(1), self.resource_span.clone(), "Semaphore::acquire_owned", "poll", true, ); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let inner = self.ll_sem.acquire(1); inner.await?; Ok(OwnedSemaphorePermit { sem: self, permits: 1, }) } /// Acquires `n` permits from the semaphore. /// /// The semaphore must be wrapped in an [`Arc`] to call this method. /// If the semaphore has been closed, this returns an [`AcquireError`]. /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the /// acquired permit. /// /// # Cancel safety /// /// This method uses a queue to fairly distribute permits in the order they /// were requested. Cancelling a call to `acquire_many_owned` makes you lose /// your place in the queue. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// #[tokio::main] /// async fn main() { /// let semaphore = Arc::new(Semaphore::new(10)); /// let mut join_handles = Vec::new(); /// /// for _ in 0..5 { /// let permit = semaphore.clone().acquire_many_owned(2).await.unwrap(); /// join_handles.push(tokio::spawn(async move { /// // perform task... /// // explicitly own `permit` in the task /// drop(permit); /// })); /// } /// /// for handle in join_handles { /// handle.await.unwrap(); /// } /// } /// ``` /// /// [`Arc`]: std::sync::Arc /// [`AcquireError`]: crate::sync::AcquireError /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit pub async fn acquire_many_owned( self: Arc, n: u32, ) -> Result { #[cfg(all(tokio_unstable, feature = "tracing"))] let inner = trace::async_op( || self.ll_sem.acquire(n as usize), self.resource_span.clone(), "Semaphore::acquire_many_owned", "poll", true, ); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let inner = self.ll_sem.acquire(n as usize); inner.await?; Ok(OwnedSemaphorePermit { sem: self, permits: n, }) } /// Tries to acquire a permit from the semaphore. /// /// The semaphore must be wrapped in an [`Arc`] to call this method. If /// the semaphore has been closed, this returns a [`TryAcquireError::Closed`] /// and a [`TryAcquireError::NoPermits`] if there are no permits left. /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the /// acquired permit. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{Semaphore, TryAcquireError}; /// /// # fn main() { /// let semaphore = Arc::new(Semaphore::new(2)); /// /// let permit_1 = Arc::clone(&semaphore).try_acquire_owned().unwrap(); /// assert_eq!(semaphore.available_permits(), 1); /// /// let permit_2 = Arc::clone(&semaphore).try_acquire_owned().unwrap(); /// assert_eq!(semaphore.available_permits(), 0); /// /// let permit_3 = semaphore.try_acquire_owned(); /// assert_eq!(permit_3.err(), Some(TryAcquireError::NoPermits)); /// # } /// ``` /// /// [`Arc`]: std::sync::Arc /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit pub fn try_acquire_owned(self: Arc) -> Result { match self.ll_sem.try_acquire(1) { Ok(()) => Ok(OwnedSemaphorePermit { sem: self, permits: 1, }), Err(e) => Err(e), } } /// Tries to acquire `n` permits from the semaphore. /// /// The semaphore must be wrapped in an [`Arc`] to call this method. If /// the semaphore has been closed, this returns a [`TryAcquireError::Closed`] /// and a [`TryAcquireError::NoPermits`] if there are no permits left. /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the /// acquired permit. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::{Semaphore, TryAcquireError}; /// /// # fn main() { /// let semaphore = Arc::new(Semaphore::new(4)); /// /// let permit_1 = Arc::clone(&semaphore).try_acquire_many_owned(3).unwrap(); /// assert_eq!(semaphore.available_permits(), 1); /// /// let permit_2 = semaphore.try_acquire_many_owned(2); /// assert_eq!(permit_2.err(), Some(TryAcquireError::NoPermits)); /// # } /// ``` /// /// [`Arc`]: std::sync::Arc /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit pub fn try_acquire_many_owned( self: Arc, n: u32, ) -> Result { match self.ll_sem.try_acquire(n as usize) { Ok(()) => Ok(OwnedSemaphorePermit { sem: self, permits: n, }), Err(e) => Err(e), } } /// Closes the semaphore. /// /// This prevents the semaphore from issuing new permits and notifies all pending waiters. /// /// # Examples /// /// ``` /// use tokio::sync::Semaphore; /// use std::sync::Arc; /// use tokio::sync::TryAcquireError; /// /// #[tokio::main] /// async fn main() { /// let semaphore = Arc::new(Semaphore::new(1)); /// let semaphore2 = semaphore.clone(); /// /// tokio::spawn(async move { /// let permit = semaphore.acquire_many(2).await; /// assert!(permit.is_err()); /// println!("waiter received error"); /// }); /// /// println!("closing semaphore"); /// semaphore2.close(); /// /// // Cannot obtain more permits /// assert_eq!(semaphore2.try_acquire().err(), Some(TryAcquireError::Closed)) /// } /// ``` pub fn close(&self) { self.ll_sem.close(); } /// Returns true if the semaphore is closed pub fn is_closed(&self) -> bool { self.ll_sem.is_closed() } } impl<'a> SemaphorePermit<'a> { /// Forgets the permit **without** releasing it back to the semaphore. /// This can be used to reduce the amount of permits available from a /// semaphore. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// let sem = Arc::new(Semaphore::new(10)); /// { /// let permit = sem.try_acquire_many(5).unwrap(); /// assert_eq!(sem.available_permits(), 5); /// permit.forget(); /// } /// /// // Since we forgot the permit, available permits won't go back to its initial value /// // even after the permit is dropped. /// assert_eq!(sem.available_permits(), 5); /// ``` pub fn forget(mut self) { self.permits = 0; } /// Merge two [`SemaphorePermit`] instances together, consuming `other` /// without releasing the permits it holds. /// /// Permits held by both `self` and `other` are released when `self` drops. /// /// # Panics /// /// This function panics if permits from different [`Semaphore`] instances /// are merged. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// let sem = Arc::new(Semaphore::new(10)); /// let mut permit = sem.try_acquire().unwrap(); /// /// for _ in 0..9 { /// let _permit = sem.try_acquire().unwrap(); /// // Merge individual permits into a single one. /// permit.merge(_permit) /// } /// /// assert_eq!(sem.available_permits(), 0); /// /// // Release all permits in a single batch. /// drop(permit); /// /// assert_eq!(sem.available_permits(), 10); /// ``` #[track_caller] pub fn merge(&mut self, mut other: Self) { assert!( std::ptr::eq(self.sem, other.sem), "merging permits from different semaphore instances" ); self.permits += other.permits; other.permits = 0; } /// Splits `n` permits from `self` and returns a new [`SemaphorePermit`] instance that holds `n` permits. /// /// If there are insufficient permits and it's not possible to reduce by `n`, returns `None`. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// let sem = Arc::new(Semaphore::new(3)); /// /// let mut p1 = sem.try_acquire_many(3).unwrap(); /// let p2 = p1.split(1).unwrap(); /// /// assert_eq!(p1.num_permits(), 2); /// assert_eq!(p2.num_permits(), 1); /// ``` pub fn split(&mut self, n: usize) -> Option { let n = u32::try_from(n).ok()?; if n > self.permits { return None; } self.permits -= n; Some(Self { sem: self.sem, permits: n, }) } /// Returns the number of permits held by `self`. pub fn num_permits(&self) -> usize { self.permits as usize } } impl OwnedSemaphorePermit { /// Forgets the permit **without** releasing it back to the semaphore. /// This can be used to reduce the amount of permits available from a /// semaphore. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// let sem = Arc::new(Semaphore::new(10)); /// { /// let permit = sem.clone().try_acquire_many_owned(5).unwrap(); /// assert_eq!(sem.available_permits(), 5); /// permit.forget(); /// } /// /// // Since we forgot the permit, available permits won't go back to its initial value /// // even after the permit is dropped. /// assert_eq!(sem.available_permits(), 5); /// ``` pub fn forget(mut self) { self.permits = 0; } /// Merge two [`OwnedSemaphorePermit`] instances together, consuming `other` /// without releasing the permits it holds. /// /// Permits held by both `self` and `other` are released when `self` drops. /// /// # Panics /// /// This function panics if permits from different [`Semaphore`] instances /// are merged. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// let sem = Arc::new(Semaphore::new(10)); /// let mut permit = sem.clone().try_acquire_owned().unwrap(); /// /// for _ in 0..9 { /// let _permit = sem.clone().try_acquire_owned().unwrap(); /// // Merge individual permits into a single one. /// permit.merge(_permit) /// } /// /// assert_eq!(sem.available_permits(), 0); /// /// // Release all permits in a single batch. /// drop(permit); /// /// assert_eq!(sem.available_permits(), 10); /// ``` #[track_caller] pub fn merge(&mut self, mut other: Self) { assert!( Arc::ptr_eq(&self.sem, &other.sem), "merging permits from different semaphore instances" ); self.permits += other.permits; other.permits = 0; } /// Splits `n` permits from `self` and returns a new [`OwnedSemaphorePermit`] instance that holds `n` permits. /// /// If there are insufficient permits and it's not possible to reduce by `n`, returns `None`. /// /// # Note /// /// It will clone the owned `Arc` to construct the new instance. /// /// # Examples /// /// ``` /// use std::sync::Arc; /// use tokio::sync::Semaphore; /// /// let sem = Arc::new(Semaphore::new(3)); /// /// let mut p1 = sem.try_acquire_many_owned(3).unwrap(); /// let p2 = p1.split(1).unwrap(); /// /// assert_eq!(p1.num_permits(), 2); /// assert_eq!(p2.num_permits(), 1); /// ``` pub fn split(&mut self, n: usize) -> Option { let n = u32::try_from(n).ok()?; if n > self.permits { return None; } self.permits -= n; Some(Self { sem: self.sem.clone(), permits: n, }) } /// Returns the [`Semaphore`] from which this permit was acquired. pub fn semaphore(&self) -> &Arc { &self.sem } /// Returns the number of permits held by `self`. pub fn num_permits(&self) -> usize { self.permits as usize } } impl Drop for SemaphorePermit<'_> { fn drop(&mut self) { self.sem.add_permits(self.permits as usize); } } impl Drop for OwnedSemaphorePermit { fn drop(&mut self) { self.sem.add_permits(self.permits as usize); } } tokio-1.43.1/src/sync/task/atomic_waker.rs000064400000000000000000000364311046102023000165660ustar 00000000000000#![cfg_attr(any(loom, not(feature = "sync")), allow(dead_code, unreachable_pub))] use crate::loom::cell::UnsafeCell; use crate::loom::hint; use crate::loom::sync::atomic::AtomicUsize; use std::fmt; use std::panic::{resume_unwind, AssertUnwindSafe, RefUnwindSafe, UnwindSafe}; use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; use std::task::Waker; /// A synchronization primitive for task waking. /// /// `AtomicWaker` will coordinate concurrent wakes with the consumer /// potentially "waking" the underlying task. This is useful in scenarios /// where a computation completes in another thread and wants to wake the /// consumer, but the consumer is in the process of being migrated to a new /// logical task. /// /// Consumers should call `register` before checking the result of a computation /// and producers should call `wake` after producing the computation (this /// differs from the usual `thread::park` pattern). It is also permitted for /// `wake` to be called **before** `register`. This results in a no-op. /// /// A single `AtomicWaker` may be reused for any number of calls to `register` or /// `wake`. pub(crate) struct AtomicWaker { state: AtomicUsize, waker: UnsafeCell>, } impl RefUnwindSafe for AtomicWaker {} impl UnwindSafe for AtomicWaker {} // `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell // stores a `Waker` value produced by calls to `register` and many threads can // race to take the waker by calling `wake`. // // If a new `Waker` instance is produced by calling `register` before an existing // one is consumed, then the existing one is overwritten. // // While `AtomicWaker` is single-producer, the implementation ensures memory // safety. In the event of concurrent calls to `register`, there will be a // single winner whose waker will get stored in the cell. The losers will not // have their tasks woken. As such, callers should ensure to add synchronization // to calls to `register`. // // The implementation uses a single `AtomicUsize` value to coordinate access to // the `Waker` cell. There are two bits that are operated on independently. These // are represented by `REGISTERING` and `WAKING`. // // The `REGISTERING` bit is set when a producer enters the critical section. The // `WAKING` bit is set when a consumer enters the critical section. Neither // bit being set is represented by `WAITING`. // // A thread obtains an exclusive lock on the waker cell by transitioning the // state from `WAITING` to `REGISTERING` or `WAKING`, depending on the // operation the thread wishes to perform. When this transition is made, it is // guaranteed that no other thread will access the waker cell. // // # Registering // // On a call to `register`, an attempt to transition the state from WAITING to // REGISTERING is made. On success, the caller obtains a lock on the waker cell. // // If the lock is obtained, then the thread sets the waker cell to the waker // provided as an argument. Then it attempts to transition the state back from // `REGISTERING` -> `WAITING`. // // If this transition is successful, then the registering process is complete // and the next call to `wake` will observe the waker. // // If the transition fails, then there was a concurrent call to `wake` that // was unable to access the waker cell (due to the registering thread holding the // lock). To handle this, the registering thread removes the waker it just set // from the cell and calls `wake` on it. This call to wake represents the // attempt to wake by the other thread (that set the `WAKING` bit). The // state is then transitioned from `REGISTERING | WAKING` back to `WAITING`. // This transition must succeed because, at this point, the state cannot be // transitioned by another thread. // // # Waking // // On a call to `wake`, an attempt to transition the state from `WAITING` to // `WAKING` is made. On success, the caller obtains a lock on the waker cell. // // If the lock is obtained, then the thread takes ownership of the current value // in the waker cell, and calls `wake` on it. The state is then transitioned // back to `WAITING`. This transition must succeed as, at this point, the state // cannot be transitioned by another thread. // // If the thread is unable to obtain the lock, the `WAKING` bit is still set. // This is because it has either been set by the current thread but the previous // value included the `REGISTERING` bit **or** a concurrent thread is in the // `WAKING` critical section. Either way, no action must be taken. // // If the current thread is the only concurrent call to `wake` and another // thread is in the `register` critical section, when the other thread **exits** // the `register` critical section, it will observe the `WAKING` bit and // handle the waker itself. // // If another thread is in the `waker` critical section, then it will handle // waking the caller task. // // # A potential race (is safely handled). // // Imagine the following situation: // // * Thread A obtains the `wake` lock and wakes a task. // // * Before thread A releases the `wake` lock, the woken task is scheduled. // // * Thread B attempts to wake the task. In theory this should result in the // task being woken, but it cannot because thread A still holds the wake // lock. // // This case is handled by requiring users of `AtomicWaker` to call `register` // **before** attempting to observe the application state change that resulted // in the task being woken. The wakers also change the application state // before calling wake. // // Because of this, the task will do one of two things. // // 1) Observe the application state change that Thread B is waking on. In // this case, it is OK for Thread B's wake to be lost. // // 2) Call register before attempting to observe the application state. Since // Thread A still holds the `wake` lock, the call to `register` will result // in the task waking itself and get scheduled again. /// Idle state. const WAITING: usize = 0; /// A new waker value is being registered with the `AtomicWaker` cell. const REGISTERING: usize = 0b01; /// The task currently registered with the `AtomicWaker` cell is being woken. const WAKING: usize = 0b10; impl AtomicWaker { /// Create an `AtomicWaker` pub(crate) fn new() -> AtomicWaker { AtomicWaker { state: AtomicUsize::new(WAITING), waker: UnsafeCell::new(None), } } /* /// Registers the current waker to be notified on calls to `wake`. pub(crate) fn register(&self, waker: Waker) { self.do_register(waker); } */ /// Registers the provided waker to be notified on calls to `wake`. /// /// The new waker will take place of any previous wakers that were registered /// by previous calls to `register`. Any calls to `wake` that happen after /// a call to `register` (as defined by the memory ordering rules), will /// wake the `register` caller's task. /// /// It is safe to call `register` with multiple other threads concurrently /// calling `wake`. This will result in the `register` caller's current /// task being woken once. /// /// This function is safe to call concurrently, but this is generally a bad /// idea. Concurrent calls to `register` will attempt to register different /// tasks to be woken. One of the callers will win and have its task set, /// but there is no guarantee as to which caller will succeed. pub(crate) fn register_by_ref(&self, waker: &Waker) { self.do_register(waker); } fn do_register(&self, waker: W) where W: WakerRef, { fn catch_unwind R, R>(f: F) -> std::thread::Result { std::panic::catch_unwind(AssertUnwindSafe(f)) } match self .state .compare_exchange(WAITING, REGISTERING, Acquire, Acquire) .unwrap_or_else(|x| x) { WAITING => { unsafe { // If `into_waker` panics (because it's code outside of // AtomicWaker) we need to prime a guard that is called on // unwind to restore the waker to a WAITING state. Otherwise // any future calls to register will incorrectly be stuck // believing it's being updated by someone else. let new_waker_or_panic = catch_unwind(move || waker.into_waker()); // Set the field to contain the new waker, or if // `into_waker` panicked, leave the old value. let mut maybe_panic = None; let mut old_waker = None; match new_waker_or_panic { Ok(new_waker) => { old_waker = self.waker.with_mut(|t| (*t).take()); self.waker.with_mut(|t| *t = Some(new_waker)); } Err(panic) => maybe_panic = Some(panic), } // Release the lock. If the state transitioned to include // the `WAKING` bit, this means that a wake has been // called concurrently, so we have to remove the waker and // wake it.` // // Start by assuming that the state is `REGISTERING` as this // is what we jut set it to. let res = self .state .compare_exchange(REGISTERING, WAITING, AcqRel, Acquire); match res { Ok(_) => { // We don't want to give the caller the panic if it // was someone else who put in that waker. let _ = catch_unwind(move || { drop(old_waker); }); } Err(actual) => { // This branch can only be reached if a // concurrent thread called `wake`. In this // case, `actual` **must** be `REGISTERING | // WAKING`. debug_assert_eq!(actual, REGISTERING | WAKING); // Take the waker to wake once the atomic operation has // completed. let mut waker = self.waker.with_mut(|t| (*t).take()); // Just swap, because no one could change state // while state == `Registering | `Waking` self.state.swap(WAITING, AcqRel); // If `into_waker` panicked, then the waker in the // waker slot is actually the old waker. if maybe_panic.is_some() { old_waker = waker.take(); } // We don't want to give the caller the panic if it // was someone else who put in that waker. if let Some(old_waker) = old_waker { let _ = catch_unwind(move || { old_waker.wake(); }); } // The atomic swap was complete, now wake the waker // and return. // // If this panics, we end up in a consumed state and // return the panic to the caller. if let Some(waker) = waker { debug_assert!(maybe_panic.is_none()); waker.wake(); } } } if let Some(panic) = maybe_panic { // If `into_waker` panicked, return the panic to the caller. resume_unwind(panic); } } } WAKING => { // Currently in the process of waking the task, i.e., // `wake` is currently being called on the old waker. // So, we call wake on the new waker. // // If this panics, someone else is responsible for restoring the // state of the waker. waker.wake(); // This is equivalent to a spin lock, so use a spin hint. hint::spin_loop(); } state => { // In this case, a concurrent thread is holding the // "registering" lock. This probably indicates a bug in the // caller's code as racing to call `register` doesn't make much // sense. // // We just want to maintain memory safety. It is ok to drop the // call to `register`. debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); } } } /// Wakes the task that last called `register`. /// /// If `register` has not been called yet, then this does nothing. pub(crate) fn wake(&self) { if let Some(waker) = self.take_waker() { // If wake panics, we've consumed the waker which is a legitimate // outcome. waker.wake(); } } /// Attempts to take the `Waker` value out of the `AtomicWaker` with the /// intention that the caller will wake the task later. pub(crate) fn take_waker(&self) -> Option { // AcqRel ordering is used in order to acquire the value of the `waker` // cell as well as to establish a `release` ordering with whatever // memory the `AtomicWaker` is associated with. match self.state.fetch_or(WAKING, AcqRel) { WAITING => { // The waking lock has been acquired. let waker = unsafe { self.waker.with_mut(|t| (*t).take()) }; // Release the lock self.state.fetch_and(!WAKING, Release); waker } state => { // There is a concurrent thread currently updating the // associated waker. // // Nothing more to do as the `WAKING` bit has been set. It // doesn't matter if there are concurrent registering threads or // not. // debug_assert!( state == REGISTERING || state == REGISTERING | WAKING || state == WAKING ); None } } } } impl Default for AtomicWaker { fn default() -> Self { AtomicWaker::new() } } impl fmt::Debug for AtomicWaker { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "AtomicWaker") } } unsafe impl Send for AtomicWaker {} unsafe impl Sync for AtomicWaker {} trait WakerRef { fn wake(self); fn into_waker(self) -> Waker; } impl WakerRef for Waker { fn wake(self) { self.wake(); } fn into_waker(self) -> Waker { self } } impl WakerRef for &Waker { fn wake(self) { self.wake_by_ref(); } fn into_waker(self) -> Waker { self.clone() } } tokio-1.43.1/src/sync/task/mod.rs000064400000000000000000000001611046102023000146670ustar 00000000000000//! Thread-safe task notification primitives. mod atomic_waker; pub(crate) use self::atomic_waker::AtomicWaker; tokio-1.43.1/src/sync/tests/atomic_waker.rs000064400000000000000000000042041046102023000167570ustar 00000000000000use crate::sync::AtomicWaker; use tokio_test::task; use std::task::Waker; #[allow(unused)] trait AssertSend: Send {} #[allow(unused)] trait AssertSync: Sync {} impl AssertSend for AtomicWaker {} impl AssertSync for AtomicWaker {} impl AssertSend for Waker {} impl AssertSync for Waker {} #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[test] fn basic_usage() { let mut waker = task::spawn(AtomicWaker::new()); waker.enter(|cx, waker| waker.register_by_ref(cx.waker())); waker.wake(); assert!(waker.is_woken()); } #[test] fn wake_without_register() { let mut waker = task::spawn(AtomicWaker::new()); waker.wake(); // Registering should not result in a notification waker.enter(|cx, waker| waker.register_by_ref(cx.waker())); assert!(!waker.is_woken()); } #[cfg(panic = "unwind")] #[test] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn atomic_waker_panic_safe() { use std::panic; use std::ptr; use std::task::{RawWaker, RawWakerVTable, Waker}; static PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( |_| panic!("clone"), |_| unimplemented!("wake"), |_| unimplemented!("wake_by_ref"), |_| (), ); static NONPANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( |_| RawWaker::new(ptr::null(), &NONPANICKING_VTABLE), |_| unimplemented!("wake"), |_| unimplemented!("wake_by_ref"), |_| (), ); let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) }; let nonpanicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &NONPANICKING_VTABLE)) }; let atomic_waker = AtomicWaker::new(); let panicking = panic::AssertUnwindSafe(&panicking); let result = panic::catch_unwind(|| { let panic::AssertUnwindSafe(panicking) = panicking; atomic_waker.register_by_ref(panicking); }); assert!(result.is_err()); assert!(atomic_waker.take_waker().is_none()); atomic_waker.register_by_ref(&nonpanicking); assert!(atomic_waker.take_waker().is_some()); } tokio-1.43.1/src/sync/tests/loom_atomic_waker.rs000064400000000000000000000052151046102023000200100ustar 00000000000000use crate::sync::task::AtomicWaker; use loom::future::block_on; use loom::sync::atomic::AtomicUsize; use loom::thread; use std::future::poll_fn; use std::sync::atomic::Ordering::Relaxed; use std::sync::Arc; use std::task::Poll::{Pending, Ready}; struct Chan { num: AtomicUsize, task: AtomicWaker, } #[test] fn basic_notification() { const NUM_NOTIFY: usize = 2; loom::model(|| { let chan = Arc::new(Chan { num: AtomicUsize::new(0), task: AtomicWaker::new(), }); for _ in 0..NUM_NOTIFY { let chan = chan.clone(); thread::spawn(move || { chan.num.fetch_add(1, Relaxed); chan.task.wake(); }); } block_on(poll_fn(move |cx| { chan.task.register_by_ref(cx.waker()); if NUM_NOTIFY == chan.num.load(Relaxed) { return Ready(()); } Pending })); }); } #[test] fn test_panicky_waker() { use std::panic; use std::ptr; use std::task::{RawWaker, RawWakerVTable, Waker}; static PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new(|_| panic!("clone"), |_| (), |_| (), |_| ()); let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) }; // If you're working with this test (and I sure hope you never have to!), // uncomment the following section because there will be a lot of panics // which would otherwise log. // // We can't however leaved it uncommented, because it's global. // panic::set_hook(Box::new(|_| ())); const NUM_NOTIFY: usize = 2; loom::model(move || { let chan = Arc::new(Chan { num: AtomicUsize::new(0), task: AtomicWaker::new(), }); for _ in 0..NUM_NOTIFY { let chan = chan.clone(); thread::spawn(move || { chan.num.fetch_add(1, Relaxed); chan.task.wake(); }); } // Note: this panic should have no effect on the overall state of the // waker and it should proceed as normal. // // A thread above might race to flag a wakeup, and a WAKING state will // be preserved if this expected panic races with that so the below // procedure should be allowed to continue uninterrupted. let _ = panic::catch_unwind(|| chan.task.register_by_ref(&panicking)); block_on(poll_fn(move |cx| { chan.task.register_by_ref(cx.waker()); if NUM_NOTIFY == chan.num.load(Relaxed) { return Ready(()); } Pending })); }); } tokio-1.43.1/src/sync/tests/loom_broadcast.rs000064400000000000000000000124171046102023000173070ustar 00000000000000use crate::sync::broadcast; use crate::sync::broadcast::error::RecvError::{Closed, Lagged}; use loom::future::block_on; use loom::sync::Arc; use loom::thread; use tokio_test::{assert_err, assert_ok}; #[test] fn broadcast_send() { loom::model(|| { let (tx1, mut rx) = broadcast::channel(2); let tx1 = Arc::new(tx1); let tx2 = tx1.clone(); let th1 = thread::spawn(move || { block_on(async { assert_ok!(tx1.send("one")); assert_ok!(tx1.send("two")); assert_ok!(tx1.send("three")); }); }); let th2 = thread::spawn(move || { block_on(async { assert_ok!(tx2.send("eins")); assert_ok!(tx2.send("zwei")); assert_ok!(tx2.send("drei")); }); }); block_on(async { let mut num = 0; loop { match rx.recv().await { Ok(_) => num += 1, Err(Closed) => break, Err(Lagged(n)) => num += n as usize, } } assert_eq!(num, 6); }); assert_ok!(th1.join()); assert_ok!(th2.join()); }); } // An `Arc` is used as the value in order to detect memory leaks. #[test] fn broadcast_two() { loom::model(|| { let (tx, mut rx1) = broadcast::channel::>(16); let mut rx2 = tx.subscribe(); let th1 = thread::spawn(move || { block_on(async { let v = assert_ok!(rx1.recv().await); assert_eq!(*v, "hello"); let v = assert_ok!(rx1.recv().await); assert_eq!(*v, "world"); match assert_err!(rx1.recv().await) { Closed => {} _ => panic!(), } }); }); let th2 = thread::spawn(move || { block_on(async { let v = assert_ok!(rx2.recv().await); assert_eq!(*v, "hello"); let v = assert_ok!(rx2.recv().await); assert_eq!(*v, "world"); match assert_err!(rx2.recv().await) { Closed => {} _ => panic!(), } }); }); assert_ok!(tx.send(Arc::new("hello"))); assert_ok!(tx.send(Arc::new("world"))); drop(tx); assert_ok!(th1.join()); assert_ok!(th2.join()); }); } #[test] fn broadcast_wrap() { loom::model(|| { let (tx, mut rx1) = broadcast::channel(2); let mut rx2 = tx.subscribe(); let th1 = thread::spawn(move || { block_on(async { let mut num = 0; loop { match rx1.recv().await { Ok(_) => num += 1, Err(Closed) => break, Err(Lagged(n)) => num += n as usize, } } assert_eq!(num, 3); }); }); let th2 = thread::spawn(move || { block_on(async { let mut num = 0; loop { match rx2.recv().await { Ok(_) => num += 1, Err(Closed) => break, Err(Lagged(n)) => num += n as usize, } } assert_eq!(num, 3); }); }); assert_ok!(tx.send("one")); assert_ok!(tx.send("two")); assert_ok!(tx.send("three")); drop(tx); assert_ok!(th1.join()); assert_ok!(th2.join()); }); } #[test] fn drop_rx() { loom::model(|| { let (tx, mut rx1) = broadcast::channel(16); let rx2 = tx.subscribe(); let th1 = thread::spawn(move || { block_on(async { let v = assert_ok!(rx1.recv().await); assert_eq!(v, "one"); let v = assert_ok!(rx1.recv().await); assert_eq!(v, "two"); let v = assert_ok!(rx1.recv().await); assert_eq!(v, "three"); match assert_err!(rx1.recv().await) { Closed => {} _ => panic!(), } }); }); let th2 = thread::spawn(move || { drop(rx2); }); assert_ok!(tx.send("one")); assert_ok!(tx.send("two")); assert_ok!(tx.send("three")); drop(tx); assert_ok!(th1.join()); assert_ok!(th2.join()); }); } #[test] fn drop_multiple_rx_with_overflow() { loom::model(move || { // It is essential to have multiple senders and receivers in this test case. let (tx, mut rx) = broadcast::channel(1); let _rx2 = tx.subscribe(); let _ = tx.send(()); let tx2 = tx.clone(); let th1 = thread::spawn(move || { block_on(async { for _ in 0..100 { let _ = tx2.send(()); } }); }); let _ = tx.send(()); let th2 = thread::spawn(move || { block_on(async { while let Ok(_) = rx.recv().await {} }); }); assert_ok!(th1.join()); assert_ok!(th2.join()); }); } tokio-1.43.1/src/sync/tests/loom_list.rs000064400000000000000000000020501046102023000163100ustar 00000000000000use crate::sync::mpsc::list; use loom::thread; use std::sync::Arc; #[test] fn smoke() { use crate::sync::mpsc::block::Read; const NUM_TX: usize = 2; const NUM_MSG: usize = 2; loom::model(|| { let (tx, mut rx) = list::channel(); let tx = Arc::new(tx); for th in 0..NUM_TX { let tx = tx.clone(); thread::spawn(move || { for i in 0..NUM_MSG { tx.push((th, i)); } }); } let mut next = vec![0; NUM_TX]; loop { match rx.pop(&tx) { Some(Read::Value((th, v))) => { assert_eq!(v, next[th]); next[th] += 1; if next.iter().all(|&i| i == NUM_MSG) { break; } } Some(Read::Closed) => { panic!(); } None => { thread::yield_now(); } } } }); } tokio-1.43.1/src/sync/tests/loom_mpsc.rs000064400000000000000000000113051046102023000163020ustar 00000000000000use crate::sync::mpsc; use loom::future::block_on; use loom::sync::Arc; use loom::thread; use std::future::poll_fn; use tokio_test::assert_ok; #[test] fn closing_tx() { loom::model(|| { let (tx, mut rx) = mpsc::channel(16); thread::spawn(move || { tx.try_send(()).unwrap(); drop(tx); }); let v = block_on(rx.recv()); assert!(v.is_some()); let v = block_on(rx.recv()); assert!(v.is_none()); }); } #[test] fn closing_unbounded_tx() { loom::model(|| { let (tx, mut rx) = mpsc::unbounded_channel(); thread::spawn(move || { tx.send(()).unwrap(); drop(tx); }); let v = block_on(rx.recv()); assert!(v.is_some()); let v = block_on(rx.recv()); assert!(v.is_none()); }); } #[test] fn closing_bounded_rx() { loom::model(|| { let (tx1, rx) = mpsc::channel::<()>(16); let tx2 = tx1.clone(); thread::spawn(move || { drop(rx); }); block_on(tx1.closed()); block_on(tx2.closed()); }); } #[test] fn closing_and_sending() { loom::model(|| { let (tx1, mut rx) = mpsc::channel::<()>(16); let tx1 = Arc::new(tx1); let tx2 = tx1.clone(); let th1 = thread::spawn(move || { tx1.try_send(()).unwrap(); }); let th2 = thread::spawn(move || { block_on(tx2.closed()); }); let th3 = thread::spawn(move || { let v = block_on(rx.recv()); assert!(v.is_some()); drop(rx); }); assert_ok!(th1.join()); assert_ok!(th2.join()); assert_ok!(th3.join()); }); } #[test] fn closing_unbounded_rx() { loom::model(|| { let (tx1, rx) = mpsc::unbounded_channel::<()>(); let tx2 = tx1.clone(); thread::spawn(move || { drop(rx); }); block_on(tx1.closed()); block_on(tx2.closed()); }); } #[test] fn dropping_tx() { loom::model(|| { let (tx, mut rx) = mpsc::channel::<()>(16); for _ in 0..2 { let tx = tx.clone(); thread::spawn(move || { drop(tx); }); } drop(tx); let v = block_on(rx.recv()); assert!(v.is_none()); }); } #[test] fn dropping_unbounded_tx() { loom::model(|| { let (tx, mut rx) = mpsc::unbounded_channel::<()>(); for _ in 0..2 { let tx = tx.clone(); thread::spawn(move || { drop(tx); }); } drop(tx); let v = block_on(rx.recv()); assert!(v.is_none()); }); } #[test] fn try_recv() { loom::model(|| { use crate::sync::{mpsc, Semaphore}; use loom::sync::{Arc, Mutex}; const PERMITS: usize = 2; const TASKS: usize = 2; const CYCLES: usize = 1; struct Context { sem: Arc, tx: mpsc::Sender<()>, rx: Mutex>, } fn run(ctx: &Context) { block_on(async { let permit = ctx.sem.acquire().await; assert_ok!(ctx.rx.lock().unwrap().try_recv()); crate::task::yield_now().await; assert_ok!(ctx.tx.clone().try_send(())); drop(permit); }); } let (tx, rx) = mpsc::channel(PERMITS); let sem = Arc::new(Semaphore::new(PERMITS)); let ctx = Arc::new(Context { sem, tx, rx: Mutex::new(rx), }); for _ in 0..PERMITS { assert_ok!(ctx.tx.clone().try_send(())); } let mut ths = Vec::new(); for _ in 0..TASKS { let ctx = ctx.clone(); ths.push(thread::spawn(move || { run(&ctx); })); } run(&ctx); for th in ths { th.join().unwrap(); } }); } #[test] fn len_nonzero_after_send() { loom::model(|| { let (send, recv) = mpsc::channel(10); let send2 = send.clone(); let join = thread::spawn(move || { block_on(send2.send("message2")).unwrap(); }); block_on(send.send("message1")).unwrap(); assert!(recv.len() != 0); join.join().unwrap(); }); } #[test] fn nonempty_after_send() { loom::model(|| { let (send, recv) = mpsc::channel(10); let send2 = send.clone(); let join = thread::spawn(move || { block_on(send2.send("message2")).unwrap(); }); block_on(send.send("message1")).unwrap(); assert!(!recv.is_empty()); join.join().unwrap(); }); } tokio-1.43.1/src/sync/tests/loom_notify.rs000064400000000000000000000230271046102023000166540ustar 00000000000000use crate::sync::Notify; use loom::future::block_on; use loom::sync::Arc; use loom::thread; use tokio_test::{assert_pending, assert_ready}; /// `util::wake_list::NUM_WAKERS` const WAKE_LIST_SIZE: usize = 32; #[test] fn notify_one() { loom::model(|| { let tx = Arc::new(Notify::new()); let rx = tx.clone(); let th = thread::spawn(move || { block_on(async { rx.notified().await; }); }); tx.notify_one(); th.join().unwrap(); }); } #[test] fn notify_waiters() { loom::model(|| { let notify = Arc::new(Notify::new()); let tx = notify.clone(); let notified1 = notify.notified(); let notified2 = notify.notified(); let th = thread::spawn(move || { tx.notify_waiters(); }); block_on(async { notified1.await; notified2.await; }); th.join().unwrap(); }); } #[test] fn notify_waiters_and_one() { loom::model(|| { let notify = Arc::new(Notify::new()); let tx1 = notify.clone(); let tx2 = notify.clone(); let th1 = thread::spawn(move || { tx1.notify_waiters(); }); let th2 = thread::spawn(move || { tx2.notify_one(); }); let th3 = thread::spawn(move || { let notified = notify.notified(); block_on(async { notified.await; }); }); th1.join().unwrap(); th2.join().unwrap(); th3.join().unwrap(); }); } #[test] fn notify_multi() { loom::model(|| { let notify = Arc::new(Notify::new()); let mut ths = vec![]; for _ in 0..2 { let notify = notify.clone(); ths.push(thread::spawn(move || { block_on(async { notify.notified().await; notify.notify_one(); }) })); } notify.notify_one(); for th in ths.drain(..) { th.join().unwrap(); } block_on(async { notify.notified().await; }); }); } #[test] fn notify_drop() { use std::future::{poll_fn, Future}; use std::task::Poll; loom::model(|| { let notify = Arc::new(Notify::new()); let rx1 = notify.clone(); let rx2 = notify.clone(); let th1 = thread::spawn(move || { let mut recv = Box::pin(rx1.notified()); block_on(poll_fn(|cx| { if recv.as_mut().poll(cx).is_ready() { rx1.notify_one(); } Poll::Ready(()) })); }); let th2 = thread::spawn(move || { block_on(async { rx2.notified().await; // Trigger second notification rx2.notify_one(); rx2.notified().await; }); }); notify.notify_one(); th1.join().unwrap(); th2.join().unwrap(); }); } /// Polls two `Notified` futures and checks if poll results are consistent /// with each other. If the first future is notified by a `notify_waiters` /// call, then the second one must be notified as well. #[test] fn notify_waiters_poll_consistency() { fn notify_waiters_poll_consistency_variant(poll_setting: [bool; 2]) { let notify = Arc::new(Notify::new()); let mut notified = [ tokio_test::task::spawn(notify.notified()), tokio_test::task::spawn(notify.notified()), ]; for i in 0..2 { if poll_setting[i] { assert_pending!(notified[i].poll()); } } let tx = notify.clone(); let th = thread::spawn(move || { tx.notify_waiters(); }); let res1 = notified[0].poll(); let res2 = notified[1].poll(); // If res1 is ready, then res2 must also be ready. assert!(res1.is_pending() || res2.is_ready()); th.join().unwrap(); } // We test different scenarios in which pending futures had or had not // been polled before the call to `notify_waiters`. loom::model(|| notify_waiters_poll_consistency_variant([false, false])); loom::model(|| notify_waiters_poll_consistency_variant([true, false])); loom::model(|| notify_waiters_poll_consistency_variant([false, true])); loom::model(|| notify_waiters_poll_consistency_variant([true, true])); } /// Polls two `Notified` futures and checks if poll results are consistent /// with each other. If the first future is notified by a `notify_waiters` /// call, then the second one must be notified as well. /// /// Here we also add other `Notified` futures in between to force the two /// tested futures to end up in different chunks. #[test] fn notify_waiters_poll_consistency_many() { fn notify_waiters_poll_consistency_many_variant(order: [usize; 2]) { let notify = Arc::new(Notify::new()); let mut futs = (0..WAKE_LIST_SIZE + 1) .map(|_| tokio_test::task::spawn(notify.notified())) .collect::>(); assert_pending!(futs[order[0]].poll()); for i in 2..futs.len() { assert_pending!(futs[i].poll()); } assert_pending!(futs[order[1]].poll()); let tx = notify.clone(); let th = thread::spawn(move || { tx.notify_waiters(); }); let res1 = futs[0].poll(); let res2 = futs[1].poll(); // If res1 is ready, then res2 must also be ready. assert!(res1.is_pending() || res2.is_ready()); th.join().unwrap(); } // We test different scenarios in which futures are polled in different order. loom::model(|| notify_waiters_poll_consistency_many_variant([0, 1])); loom::model(|| notify_waiters_poll_consistency_many_variant([1, 0])); } /// Checks if a call to `notify_waiters` is observed as atomic when combined /// with a concurrent call to `notify_one`. #[test] fn notify_waiters_is_atomic() { fn notify_waiters_is_atomic_variant(tested_fut_index: usize) { let notify = Arc::new(Notify::new()); let mut futs = (0..WAKE_LIST_SIZE + 1) .map(|_| tokio_test::task::spawn(notify.notified())) .collect::>(); for fut in &mut futs { assert_pending!(fut.poll()); } let tx = notify.clone(); let th = thread::spawn(move || { tx.notify_waiters(); }); block_on(async { // If awaiting one of the futures completes, then we should be // able to assume that all pending futures are notified. Therefore // a notification from a subsequent `notify_one` call should not // be consumed by an old future. futs.remove(tested_fut_index).await; let mut new_fut = tokio_test::task::spawn(notify.notified()); assert_pending!(new_fut.poll()); notify.notify_one(); // `new_fut` must consume the notification from `notify_one`. assert_ready!(new_fut.poll()); }); th.join().unwrap(); } // We test different scenarios in which the tested future is at the beginning // or at the end of the waiters queue used by `Notify`. loom::model(|| notify_waiters_is_atomic_variant(0)); loom::model(|| notify_waiters_is_atomic_variant(32)); } /// Checks if a single call to `notify_waiters` does not get through two `Notified` /// futures created and awaited sequentially like this: /// ```ignore /// notify.notified().await; /// notify.notified().await; /// ``` #[test] fn notify_waiters_sequential_notified_await() { use crate::sync::oneshot; loom::model(|| { let notify = Arc::new(Notify::new()); let (tx_fst, rx_fst) = oneshot::channel(); let (tx_snd, rx_snd) = oneshot::channel(); let receiver = thread::spawn({ let notify = notify.clone(); move || { block_on(async { // Poll the first `Notified` to put it as the first waiter // in the queue. let mut first_notified = tokio_test::task::spawn(notify.notified()); assert_pending!(first_notified.poll()); // Create additional waiters to force `notify_waiters` to // release the lock at least once. let _task_pile = (0..WAKE_LIST_SIZE + 1) .map(|_| { let mut fut = tokio_test::task::spawn(notify.notified()); assert_pending!(fut.poll()); fut }) .collect::>(); // We are ready for the notify_waiters call. tx_fst.send(()).unwrap(); first_notified.await; // Poll the second `Notified` future to try to insert // it to the waiters queue. let mut second_notified = tokio_test::task::spawn(notify.notified()); assert_pending!(second_notified.poll()); // Wait for the `notify_waiters` to end and check if we // are woken up. rx_snd.await.unwrap(); assert_pending!(second_notified.poll()); }); } }); // Wait for the signal and call `notify_waiters`. block_on(rx_fst).unwrap(); notify.notify_waiters(); tx_snd.send(()).unwrap(); receiver.join().unwrap(); }); } tokio-1.43.1/src/sync/tests/loom_oneshot.rs000064400000000000000000000101461046102023000170210ustar 00000000000000use crate::sync::oneshot; use loom::future::block_on; use loom::thread; use std::future::poll_fn; use std::task::Poll::{Pending, Ready}; #[test] fn smoke() { loom::model(|| { let (tx, rx) = oneshot::channel(); thread::spawn(move || { tx.send(1).unwrap(); }); let value = block_on(rx).unwrap(); assert_eq!(1, value); }); } #[test] fn changing_rx_task() { loom::model(|| { let (tx, mut rx) = oneshot::channel(); thread::spawn(move || { tx.send(1).unwrap(); }); let rx = thread::spawn(move || { let ready = block_on(poll_fn(|cx| match Pin::new(&mut rx).poll(cx) { Ready(Ok(value)) => { assert_eq!(1, value); Ready(true) } Ready(Err(_)) => unimplemented!(), Pending => Ready(false), })); if ready { None } else { Some(rx) } }) .join() .unwrap(); if let Some(rx) = rx { // Previous task parked, use a new task... let value = block_on(rx).unwrap(); assert_eq!(1, value); } }); } #[test] fn try_recv_close() { // reproduces https://github.com/tokio-rs/tokio/issues/4225 loom::model(|| { let (tx, mut rx) = oneshot::channel(); thread::spawn(move || { let _ = tx.send(()); }); rx.close(); let _ = rx.try_recv(); }) } #[test] fn recv_closed() { // reproduces https://github.com/tokio-rs/tokio/issues/4225 loom::model(|| { let (tx, mut rx) = oneshot::channel(); thread::spawn(move || { let _ = tx.send(1); }); rx.close(); let _ = block_on(rx); }); } // TODO: Move this into `oneshot` proper. use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; struct OnClose<'a> { tx: &'a mut oneshot::Sender, } impl<'a> OnClose<'a> { fn new(tx: &'a mut oneshot::Sender) -> Self { OnClose { tx } } } impl Future for OnClose<'_> { type Output = bool; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let fut = self.get_mut().tx.closed(); crate::pin!(fut); Ready(fut.poll(cx).is_ready()) } } #[test] fn changing_tx_task() { loom::model(|| { let (mut tx, rx) = oneshot::channel::(); thread::spawn(move || { drop(rx); }); let tx = thread::spawn(move || { let t1 = block_on(OnClose::new(&mut tx)); if t1 { None } else { Some(tx) } }) .join() .unwrap(); if let Some(mut tx) = tx { // Previous task parked, use a new task... block_on(OnClose::new(&mut tx)); } }); } #[test] fn checking_tx_send_ok_not_drop() { use std::borrow::Borrow; use std::cell::Cell; loom::thread_local! { static IS_RX: Cell = Cell::new(true); } struct Msg; impl Drop for Msg { fn drop(&mut self) { IS_RX.with(|is_rx: &Cell<_>| { // On `tx.send(msg)` returning `Err(msg)`, // we call `std::mem::forget(msg)`, so that // `drop` is not expected to be called in the // tx thread. assert!(is_rx.get()); }); } } let mut builder = loom::model::Builder::new(); builder.preemption_bound = Some(2); builder.check(|| { let (tx, rx) = oneshot::channel(); // tx thread let tx_thread_join_handle = thread::spawn(move || { // Ensure that `Msg::drop` in this thread will see is_rx == false IS_RX.with(|is_rx: &Cell<_>| { is_rx.set(false); }); if let Err(msg) = tx.send(Msg) { std::mem::forget(msg); } }); // main thread is the rx thread drop(rx); tx_thread_join_handle.join().unwrap(); }); } tokio-1.43.1/src/sync/tests/loom_rwlock.rs000064400000000000000000000056211046102023000166450ustar 00000000000000use crate::sync::rwlock::*; use loom::future::block_on; use loom::thread; use std::sync::Arc; #[test] fn concurrent_write() { let b = loom::model::Builder::new(); b.check(|| { let rwlock = Arc::new(RwLock::::new(0)); let rwclone = rwlock.clone(); let t1 = thread::spawn(move || { block_on(async { let mut guard = rwclone.write().await; *guard += 5; }); }); let rwclone = rwlock.clone(); let t2 = thread::spawn(move || { block_on(async { let mut guard = rwclone.write_owned().await; *guard += 5; }); }); t1.join().expect("thread 1 write should not panic"); t2.join().expect("thread 2 write should not panic"); //when all threads have finished the value on the lock should be 10 let guard = block_on(rwlock.read()); assert_eq!(10, *guard); }); } #[test] fn concurrent_read_write() { let b = loom::model::Builder::new(); b.check(|| { let rwlock = Arc::new(RwLock::::new(0)); let rwclone = rwlock.clone(); let t1 = thread::spawn(move || { block_on(async { let mut guard = rwclone.write().await; *guard += 5; }); }); let rwclone = rwlock.clone(); let t2 = thread::spawn(move || { block_on(async { let mut guard = rwclone.write_owned().await; *guard += 5; }); }); let rwclone = rwlock.clone(); let t3 = thread::spawn(move || { block_on(async { let guard = rwclone.read().await; //at this state the value on the lock may either be 0, 5, or 10 assert!(*guard == 0 || *guard == 5 || *guard == 10); }); }); { let guard = block_on(rwlock.clone().read_owned()); //at this state the value on the lock may either be 0, 5, or 10 assert!(*guard == 0 || *guard == 5 || *guard == 10); } t1.join().expect("thread 1 write should not panic"); t2.join().expect("thread 2 write should not panic"); t3.join().expect("thread 3 read should not panic"); let guard = block_on(rwlock.read()); //when all threads have finished the value on the lock should be 10 assert_eq!(10, *guard); }); } #[test] fn downgrade() { loom::model(|| { let lock = Arc::new(RwLock::new(1)); let n = block_on(lock.write()); let cloned_lock = lock.clone(); let handle = thread::spawn(move || { let mut guard = block_on(cloned_lock.write()); *guard = 2; }); let n = n.downgrade(); assert_eq!(*n, 1); drop(n); handle.join().unwrap(); assert_eq!(*block_on(lock.read()), 2); }); } tokio-1.43.1/src/sync/tests/loom_semaphore_batch.rs000064400000000000000000000141461046102023000204720ustar 00000000000000use crate::sync::batch_semaphore::*; use loom::future::block_on; use loom::sync::atomic::AtomicUsize; use loom::thread; use std::future::{poll_fn, Future}; use std::pin::Pin; use std::sync::atomic::Ordering::SeqCst; use std::sync::Arc; use std::task::Poll::Ready; use std::task::{Context, Poll}; #[test] fn basic_usage() { const NUM: usize = 2; struct Shared { semaphore: Semaphore, active: AtomicUsize, } async fn actor(shared: Arc) { shared.semaphore.acquire(1).await.unwrap(); let actual = shared.active.fetch_add(1, SeqCst); assert!(actual <= NUM - 1); let actual = shared.active.fetch_sub(1, SeqCst); assert!(actual <= NUM); shared.semaphore.release(1); } loom::model(|| { let shared = Arc::new(Shared { semaphore: Semaphore::new(NUM), active: AtomicUsize::new(0), }); for _ in 0..NUM { let shared = shared.clone(); thread::spawn(move || { block_on(actor(shared)); }); } block_on(actor(shared)); }); } #[test] fn release() { loom::model(|| { let semaphore = Arc::new(Semaphore::new(1)); { let semaphore = semaphore.clone(); thread::spawn(move || { block_on(semaphore.acquire(1)).unwrap(); semaphore.release(1); }); } block_on(semaphore.acquire(1)).unwrap(); semaphore.release(1); }); } #[test] fn basic_closing() { const NUM: usize = 2; loom::model(|| { let semaphore = Arc::new(Semaphore::new(1)); for _ in 0..NUM { let semaphore = semaphore.clone(); thread::spawn(move || { for _ in 0..2 { block_on(semaphore.acquire(1)).map_err(|_| ())?; semaphore.release(1); } Ok::<(), ()>(()) }); } semaphore.close(); }); } #[test] fn concurrent_close() { const NUM: usize = 3; loom::model(|| { let semaphore = Arc::new(Semaphore::new(1)); for _ in 0..NUM { let semaphore = semaphore.clone(); thread::spawn(move || { block_on(semaphore.acquire(1)).map_err(|_| ())?; semaphore.release(1); semaphore.close(); Ok::<(), ()>(()) }); } }); } #[test] fn concurrent_cancel() { async fn poll_and_cancel(semaphore: Arc) { let mut acquire1 = Some(semaphore.acquire(1)); let mut acquire2 = Some(semaphore.acquire(1)); poll_fn(|cx| { // poll the acquire future once, and then immediately throw // it away. this simulates a situation where a future is // polled and then cancelled, such as by a timeout. if let Some(acquire) = acquire1.take() { pin!(acquire); let _ = acquire.poll(cx); } if let Some(acquire) = acquire2.take() { pin!(acquire); let _ = acquire.poll(cx); } Poll::Ready(()) }) .await } loom::model(|| { let semaphore = Arc::new(Semaphore::new(0)); let t1 = { let semaphore = semaphore.clone(); thread::spawn(move || block_on(poll_and_cancel(semaphore))) }; let t2 = { let semaphore = semaphore.clone(); thread::spawn(move || block_on(poll_and_cancel(semaphore))) }; let t3 = { let semaphore = semaphore.clone(); thread::spawn(move || block_on(poll_and_cancel(semaphore))) }; t1.join().unwrap(); semaphore.release(10); t2.join().unwrap(); t3.join().unwrap(); }); } #[test] fn batch() { let mut b = loom::model::Builder::new(); b.preemption_bound = Some(1); b.check(|| { let semaphore = Arc::new(Semaphore::new(10)); let active = Arc::new(AtomicUsize::new(0)); let mut ths = vec![]; for _ in 0..2 { let semaphore = semaphore.clone(); let active = active.clone(); ths.push(thread::spawn(move || { for n in &[4, 10, 8] { block_on(semaphore.acquire(*n)).unwrap(); active.fetch_add(*n as usize, SeqCst); let num_active = active.load(SeqCst); assert!(num_active <= 10); thread::yield_now(); active.fetch_sub(*n as usize, SeqCst); semaphore.release(*n as usize); } })); } for th in ths.into_iter() { th.join().unwrap(); } assert_eq!(10, semaphore.available_permits()); }); } #[test] fn release_during_acquire() { loom::model(|| { let semaphore = Arc::new(Semaphore::new(10)); semaphore .try_acquire(8) .expect("try_acquire should succeed; semaphore uncontended"); let semaphore2 = semaphore.clone(); let thread = thread::spawn(move || block_on(semaphore2.acquire(4)).unwrap()); semaphore.release(8); thread.join().unwrap(); semaphore.release(4); assert_eq!(10, semaphore.available_permits()); }) } #[test] fn concurrent_permit_updates() { loom::model(move || { let semaphore = Arc::new(Semaphore::new(5)); let t1 = { let semaphore = semaphore.clone(); thread::spawn(move || semaphore.release(3)) }; let t2 = { let semaphore = semaphore.clone(); thread::spawn(move || { semaphore .try_acquire(1) .expect("try_acquire should succeed") }) }; let t3 = { let semaphore = semaphore.clone(); thread::spawn(move || semaphore.forget_permits(2)) }; t1.join().unwrap(); t2.join().unwrap(); t3.join().unwrap(); assert_eq!(semaphore.available_permits(), 5); }) } tokio-1.43.1/src/sync/tests/loom_watch.rs000064400000000000000000000047561046102023000164620ustar 00000000000000use crate::sync::watch; use loom::future::block_on; use loom::thread; use std::sync::Arc; #[test] fn smoke() { loom::model(|| { let (tx, mut rx1) = watch::channel(1); let mut rx2 = rx1.clone(); let mut rx3 = rx1.clone(); let mut rx4 = rx1.clone(); let mut rx5 = rx1.clone(); let th = thread::spawn(move || { tx.send(2).unwrap(); }); block_on(rx1.changed()).unwrap(); assert_eq!(*rx1.borrow(), 2); block_on(rx2.changed()).unwrap(); assert_eq!(*rx2.borrow(), 2); block_on(rx3.changed()).unwrap(); assert_eq!(*rx3.borrow(), 2); block_on(rx4.changed()).unwrap(); assert_eq!(*rx4.borrow(), 2); block_on(rx5.changed()).unwrap(); assert_eq!(*rx5.borrow(), 2); th.join().unwrap(); }) } #[test] fn wait_for_test() { loom::model(move || { let (tx, mut rx) = watch::channel(false); let tx_arc = Arc::new(tx); let tx1 = tx_arc.clone(); let tx2 = tx_arc.clone(); let th1 = thread::spawn(move || { for _ in 0..2 { tx1.send_modify(|_x| {}); } }); let th2 = thread::spawn(move || { tx2.send(true).unwrap(); }); assert_eq!(*block_on(rx.wait_for(|x| *x)).unwrap(), true); th1.join().unwrap(); th2.join().unwrap(); }); } #[test] fn wait_for_returns_correct_value() { loom::model(move || { let (tx, mut rx) = watch::channel(0); let jh = thread::spawn(move || { tx.send(1).unwrap(); tx.send(2).unwrap(); tx.send(3).unwrap(); }); // Stop at the first value we are called at. let mut stopped_at = usize::MAX; let returned = *block_on(rx.wait_for(|x| { stopped_at = *x; true })) .unwrap(); // Check that it returned the same value as the one we returned // `true` for. assert_eq!(stopped_at, returned); jh.join().unwrap(); }); } #[test] fn multiple_sender_drop_concurrently() { loom::model(move || { let (tx1, rx) = watch::channel(0); let tx2 = tx1.clone(); let jh = thread::spawn(move || { drop(tx2); }); assert!(rx.has_changed().is_ok()); drop(tx1); jh.join().unwrap(); // Check if all sender are dropped and closed flag is set. assert!(rx.has_changed().is_err()); }); } tokio-1.43.1/src/sync/tests/mod.rs000064400000000000000000000004531046102023000150730ustar 00000000000000cfg_not_loom! { mod atomic_waker; mod notify; mod semaphore_batch; } cfg_loom! { mod loom_atomic_waker; mod loom_broadcast; mod loom_list; mod loom_mpsc; mod loom_notify; mod loom_oneshot; mod loom_semaphore_batch; mod loom_watch; mod loom_rwlock; } tokio-1.43.1/src/sync/tests/notify.rs000064400000000000000000000057651046102023000156370ustar 00000000000000use crate::sync::Notify; use std::future::Future; use std::sync::Arc; use std::task::{Context, RawWaker, RawWakerVTable, Waker}; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[test] fn notify_clones_waker_before_lock() { const VTABLE: &RawWakerVTable = &RawWakerVTable::new(clone_w, wake, wake_by_ref, drop_w); unsafe fn clone_w(data: *const ()) -> RawWaker { let ptr = data as *const Notify; Arc::::increment_strong_count(ptr); // Or some other arbitrary code that shouldn't be executed while the // Notify wait list is locked. (*ptr).notify_one(); RawWaker::new(data, VTABLE) } unsafe fn drop_w(data: *const ()) { drop(Arc::::from_raw(data as *const Notify)); } unsafe fn wake(_data: *const ()) { unreachable!() } unsafe fn wake_by_ref(_data: *const ()) { unreachable!() } let notify = Arc::new(Notify::new()); let notify2 = notify.clone(); let waker = unsafe { Waker::from_raw(RawWaker::new(Arc::into_raw(notify2) as *const _, VTABLE)) }; let mut cx = Context::from_waker(&waker); let future = notify.notified(); pin!(future); // The result doesn't matter, we're just testing that we don't deadlock. let _ = future.poll(&mut cx); } #[cfg(panic = "unwind")] #[test] fn notify_waiters_handles_panicking_waker() { use futures::task::ArcWake; let notify = Arc::new(Notify::new()); struct PanickingWaker(#[allow(dead_code)] Arc); impl ArcWake for PanickingWaker { fn wake_by_ref(_arc_self: &Arc) { panic!("waker panicked"); } } let bad_fut = notify.notified(); pin!(bad_fut); let waker = futures::task::waker(Arc::new(PanickingWaker(notify.clone()))); let mut cx = Context::from_waker(&waker); let _ = bad_fut.poll(&mut cx); let mut futs = Vec::new(); for _ in 0..32 { let mut fut = tokio_test::task::spawn(notify.notified()); assert!(fut.poll().is_pending()); futs.push(fut); } assert!(std::panic::catch_unwind(|| { notify.notify_waiters(); }) .is_err()); for mut fut in futs { assert!(fut.poll().is_ready()); } } #[test] fn notify_simple() { let notify = Notify::new(); let mut fut1 = tokio_test::task::spawn(notify.notified()); assert!(fut1.poll().is_pending()); let mut fut2 = tokio_test::task::spawn(notify.notified()); assert!(fut2.poll().is_pending()); notify.notify_waiters(); assert!(fut1.poll().is_ready()); assert!(fut2.poll().is_ready()); } #[test] #[cfg(not(target_family = "wasm"))] fn watch_test() { let rt = crate::runtime::Builder::new_current_thread() .build() .unwrap(); rt.block_on(async { let (tx, mut rx) = crate::sync::watch::channel(()); crate::spawn(async move { let _ = tx.send(()); }); let _ = rx.changed().await; }); } tokio-1.43.1/src/sync/tests/semaphore_batch.rs000064400000000000000000000166641046102023000174530ustar 00000000000000use crate::sync::batch_semaphore::Semaphore; use tokio_test::*; const MAX_PERMITS: usize = crate::sync::Semaphore::MAX_PERMITS; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[test] fn poll_acquire_one_available() { let s = Semaphore::new(100); assert_eq!(s.available_permits(), 100); // Polling for a permit succeeds immediately assert_ready_ok!(task::spawn(s.acquire(1)).poll()); assert_eq!(s.available_permits(), 99); } #[test] fn poll_acquire_many_available() { let s = Semaphore::new(100); assert_eq!(s.available_permits(), 100); // Polling for a permit succeeds immediately assert_ready_ok!(task::spawn(s.acquire(5)).poll()); assert_eq!(s.available_permits(), 95); assert_ready_ok!(task::spawn(s.acquire(5)).poll()); assert_eq!(s.available_permits(), 90); } #[test] fn try_acquire_one_available() { let s = Semaphore::new(100); assert_eq!(s.available_permits(), 100); assert_ok!(s.try_acquire(1)); assert_eq!(s.available_permits(), 99); assert_ok!(s.try_acquire(1)); assert_eq!(s.available_permits(), 98); } #[test] fn try_acquire_many_available() { let s = Semaphore::new(100); assert_eq!(s.available_permits(), 100); assert_ok!(s.try_acquire(5)); assert_eq!(s.available_permits(), 95); assert_ok!(s.try_acquire(5)); assert_eq!(s.available_permits(), 90); } #[test] fn poll_acquire_one_unavailable() { let s = Semaphore::new(1); // Acquire the first permit assert_ready_ok!(task::spawn(s.acquire(1)).poll()); assert_eq!(s.available_permits(), 0); let mut acquire_2 = task::spawn(s.acquire(1)); // Try to acquire the second permit assert_pending!(acquire_2.poll()); assert_eq!(s.available_permits(), 0); s.release(1); assert_eq!(s.available_permits(), 0); assert!(acquire_2.is_woken()); assert_ready_ok!(acquire_2.poll()); assert_eq!(s.available_permits(), 0); s.release(1); assert_eq!(s.available_permits(), 1); } #[test] fn poll_acquire_many_unavailable() { let s = Semaphore::new(5); // Acquire the first permit assert_ready_ok!(task::spawn(s.acquire(1)).poll()); assert_eq!(s.available_permits(), 4); // Try to acquire the second permit let mut acquire_2 = task::spawn(s.acquire(5)); assert_pending!(acquire_2.poll()); assert_eq!(s.available_permits(), 0); // Try to acquire the third permit let mut acquire_3 = task::spawn(s.acquire(3)); assert_pending!(acquire_3.poll()); assert_eq!(s.available_permits(), 0); s.release(1); assert_eq!(s.available_permits(), 0); assert!(acquire_2.is_woken()); assert_ready_ok!(acquire_2.poll()); assert!(!acquire_3.is_woken()); assert_eq!(s.available_permits(), 0); s.release(1); assert!(!acquire_3.is_woken()); assert_eq!(s.available_permits(), 0); s.release(2); assert!(acquire_3.is_woken()); assert_ready_ok!(acquire_3.poll()); } #[test] fn try_acquire_one_unavailable() { let s = Semaphore::new(1); // Acquire the first permit assert_ok!(s.try_acquire(1)); assert_eq!(s.available_permits(), 0); assert_err!(s.try_acquire(1)); s.release(1); assert_eq!(s.available_permits(), 1); assert_ok!(s.try_acquire(1)); s.release(1); assert_eq!(s.available_permits(), 1); } #[test] fn try_acquire_many_unavailable() { let s = Semaphore::new(5); // Acquire the first permit assert_ok!(s.try_acquire(1)); assert_eq!(s.available_permits(), 4); assert_err!(s.try_acquire(5)); s.release(1); assert_eq!(s.available_permits(), 5); assert_ok!(s.try_acquire(5)); s.release(1); assert_eq!(s.available_permits(), 1); s.release(1); assert_eq!(s.available_permits(), 2); } #[test] fn poll_acquire_one_zero_permits() { let s = Semaphore::new(0); assert_eq!(s.available_permits(), 0); // Try to acquire the permit let mut acquire = task::spawn(s.acquire(1)); assert_pending!(acquire.poll()); s.release(1); assert!(acquire.is_woken()); assert_ready_ok!(acquire.poll()); } #[test] fn max_permits_doesnt_panic() { Semaphore::new(MAX_PERMITS); } #[test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn validates_max_permits() { Semaphore::new(MAX_PERMITS + 1); } #[test] fn close_semaphore_prevents_acquire() { let s = Semaphore::new(5); s.close(); assert_eq!(5, s.available_permits()); assert_ready_err!(task::spawn(s.acquire(1)).poll()); assert_eq!(5, s.available_permits()); assert_ready_err!(task::spawn(s.acquire(1)).poll()); assert_eq!(5, s.available_permits()); } #[test] fn close_semaphore_notifies_permit1() { let s = Semaphore::new(0); let mut acquire = task::spawn(s.acquire(1)); assert_pending!(acquire.poll()); s.close(); assert!(acquire.is_woken()); assert_ready_err!(acquire.poll()); } #[test] fn close_semaphore_notifies_permit2() { let s = Semaphore::new(2); // Acquire a couple of permits assert_ready_ok!(task::spawn(s.acquire(1)).poll()); assert_ready_ok!(task::spawn(s.acquire(1)).poll()); let mut acquire3 = task::spawn(s.acquire(1)); let mut acquire4 = task::spawn(s.acquire(1)); assert_pending!(acquire3.poll()); assert_pending!(acquire4.poll()); s.close(); assert!(acquire3.is_woken()); assert!(acquire4.is_woken()); assert_ready_err!(acquire3.poll()); assert_ready_err!(acquire4.poll()); assert_eq!(0, s.available_permits()); s.release(1); assert_eq!(1, s.available_permits()); assert_ready_err!(task::spawn(s.acquire(1)).poll()); s.release(1); assert_eq!(2, s.available_permits()); } #[test] fn cancel_acquire_releases_permits() { let s = Semaphore::new(10); s.try_acquire(4).expect("uncontended try_acquire succeeds"); assert_eq!(6, s.available_permits()); let mut acquire = task::spawn(s.acquire(8)); assert_pending!(acquire.poll()); assert_eq!(0, s.available_permits()); drop(acquire); assert_eq!(6, s.available_permits()); assert_ok!(s.try_acquire(6)); } #[test] fn release_permits_at_drop() { use crate::sync::semaphore::*; use futures::task::ArcWake; use std::future::Future; use std::sync::Arc; let sem = Arc::new(Semaphore::new(1)); struct ReleaseOnDrop(#[allow(dead_code)] Option); impl ArcWake for ReleaseOnDrop { fn wake_by_ref(_arc_self: &Arc) {} } let mut fut = Box::pin(async { let _permit = sem.acquire().await.unwrap(); }); // Second iteration shouldn't deadlock. for _ in 0..=1 { let waker = futures::task::waker(Arc::new(ReleaseOnDrop( sem.clone().try_acquire_owned().ok(), ))); let mut cx = std::task::Context::from_waker(&waker); assert!(fut.as_mut().poll(&mut cx).is_pending()); } } #[test] fn forget_permits_basic() { let s = Semaphore::new(10); assert_eq!(s.forget_permits(4), 4); assert_eq!(s.available_permits(), 6); assert_eq!(s.forget_permits(10), 6); assert_eq!(s.available_permits(), 0); } #[test] fn update_permits_many_times() { let s = Semaphore::new(5); let mut acquire = task::spawn(s.acquire(7)); assert_pending!(acquire.poll()); s.release(5); assert_ready_ok!(acquire.poll()); assert_eq!(s.available_permits(), 3); assert_eq!(s.forget_permits(3), 3); assert_eq!(s.available_permits(), 0); } tokio-1.43.1/src/sync/watch.rs000064400000000000000000001431521046102023000142640ustar 00000000000000#![cfg_attr(not(feature = "sync"), allow(dead_code, unreachable_pub))] //! A multi-producer, multi-consumer channel that only retains the *last* sent //! value. //! //! This channel is useful for watching for changes to a value from multiple //! points in the code base, for example, changes to configuration values. //! //! # Usage //! //! [`channel`] returns a [`Sender`] / [`Receiver`] pair. These are the producer //! and consumer halves of the channel. The channel is created with an initial //! value. //! //! Each [`Receiver`] independently tracks the last value *seen* by its caller. //! //! To access the **current** value stored in the channel and mark it as *seen* //! by a given [`Receiver`], use [`Receiver::borrow_and_update()`]. //! //! To access the current value **without** marking it as *seen*, use //! [`Receiver::borrow()`]. (If the value has already been marked *seen*, //! [`Receiver::borrow()`] is equivalent to [`Receiver::borrow_and_update()`].) //! //! For more information on when to use these methods, see //! [here](#borrow_and_update-versus-borrow). //! //! ## Change notifications //! //! The [`Receiver`] half provides an asynchronous [`changed`] method. This //! method is ready when a new, *unseen* value is sent via the [`Sender`] half. //! //! * [`Receiver::changed()`] returns `Ok(())` on receiving a new value, or //! `Err(`[`error::RecvError`]`)` if the [`Sender`] has been dropped. //! * If the current value is *unseen* when calling [`changed`], then //! [`changed`] will return immediately. If the current value is *seen*, then //! it will sleep until either a new message is sent via the [`Sender`] half, //! or the [`Sender`] is dropped. //! * On completion, the [`changed`] method marks the new value as *seen*. //! * At creation, the initial value is considered *seen*. In other words, //! [`Receiver::changed()`] will not return until a subsequent value is sent. //! * New [`Receiver`] instances can be created with [`Sender::subscribe()`]. //! The current value at the time the [`Receiver`] is created is considered //! *seen*. //! //! ## `borrow_and_update` versus `borrow` //! //! If the receiver intends to await notifications from [`changed`] in a loop, //! [`Receiver::borrow_and_update()`] should be preferred over //! [`Receiver::borrow()`]. This avoids a potential race where a new value is //! sent between [`changed`] being ready and the value being read. (If //! [`Receiver::borrow()`] is used, the loop may run twice with the same value.) //! //! If the receiver is only interested in the current value, and does not intend //! to wait for changes, then [`Receiver::borrow()`] can be used. It may be more //! convenient to use [`borrow`](Receiver::borrow) since it's an `&self` //! method---[`borrow_and_update`](Receiver::borrow_and_update) requires `&mut //! self`. //! //! # Examples //! //! The following example prints `hello! world! `. //! //! ``` //! use tokio::sync::watch; //! use tokio::time::{Duration, sleep}; //! //! # async fn dox() -> Result<(), Box> { //! let (tx, mut rx) = watch::channel("hello"); //! //! tokio::spawn(async move { //! // Use the equivalent of a "do-while" loop so the initial value is //! // processed before awaiting the `changed()` future. //! loop { //! println!("{}! ", *rx.borrow_and_update()); //! if rx.changed().await.is_err() { //! break; //! } //! } //! }); //! //! sleep(Duration::from_millis(100)).await; //! tx.send("world")?; //! # Ok(()) //! # } //! ``` //! //! # Closing //! //! [`Sender::is_closed`] and [`Sender::closed`] allow the producer to detect //! when all [`Receiver`] handles have been dropped. This indicates that there //! is no further interest in the values being produced and work can be stopped. //! //! The value in the channel will not be dropped until the sender and all //! receivers have been dropped. //! //! # Thread safety //! //! Both [`Sender`] and [`Receiver`] are thread safe. They can be moved to other //! threads and can be used in a concurrent environment. Clones of [`Receiver`] //! handles may be moved to separate threads and also used concurrently. //! //! [`Sender`]: crate::sync::watch::Sender //! [`Receiver`]: crate::sync::watch::Receiver //! [`changed`]: crate::sync::watch::Receiver::changed //! [`Receiver::changed()`]: crate::sync::watch::Receiver::changed //! [`Receiver::borrow()`]: crate::sync::watch::Receiver::borrow //! [`Receiver::borrow_and_update()`]: //! crate::sync::watch::Receiver::borrow_and_update //! [`channel`]: crate::sync::watch::channel //! [`Sender::is_closed`]: crate::sync::watch::Sender::is_closed //! [`Sender::closed`]: crate::sync::watch::Sender::closed //! [`Sender::subscribe()`]: crate::sync::watch::Sender::subscribe use crate::runtime::coop::cooperative; use crate::sync::notify::Notify; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::atomic::Ordering::{AcqRel, Relaxed}; use crate::loom::sync::{Arc, RwLock, RwLockReadGuard}; use std::fmt; use std::mem; use std::ops; use std::panic; /// Receives values from the associated [`Sender`](struct@Sender). /// /// Instances are created by the [`channel`](fn@channel) function. /// /// To turn this receiver into a `Stream`, you can use the [`WatchStream`] /// wrapper. /// /// [`WatchStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.WatchStream.html #[derive(Debug)] pub struct Receiver { /// Pointer to the shared state shared: Arc>, /// Last observed version version: Version, } /// Sends values to the associated [`Receiver`](struct@Receiver). /// /// Instances are created by the [`channel`](fn@channel) function. #[derive(Debug)] pub struct Sender { shared: Arc>, } impl Clone for Sender { fn clone(&self) -> Self { self.shared.ref_count_tx.fetch_add(1, Relaxed); Self { shared: self.shared.clone(), } } } impl Default for Sender { fn default() -> Self { Self::new(T::default()) } } /// Returns a reference to the inner value. /// /// Outstanding borrows hold a read lock on the inner value. This means that /// long-lived borrows could cause the producer half to block. It is recommended /// to keep the borrow as short-lived as possible. Additionally, if you are /// running in an environment that allows `!Send` futures, you must ensure that /// the returned `Ref` type is never held alive across an `.await` point, /// otherwise, it can lead to a deadlock. /// /// The priority policy of the lock is dependent on the underlying lock /// implementation, and this type does not guarantee that any particular policy /// will be used. In particular, a producer which is waiting to acquire the lock /// in `send` might or might not block concurrent calls to `borrow`, e.g.: /// ///
Potential deadlock example /// /// ```text /// // Task 1 (on thread A) | // Task 2 (on thread B) /// let _ref1 = rx.borrow(); | /// | // will block /// | let _ = tx.send(()); /// // may deadlock | /// let _ref2 = rx.borrow(); | /// ``` ///
#[derive(Debug)] pub struct Ref<'a, T> { inner: RwLockReadGuard<'a, T>, has_changed: bool, } impl<'a, T> Ref<'a, T> { /// Indicates if the borrowed value is considered as _changed_ since the last /// time it has been marked as seen. /// /// Unlike [`Receiver::has_changed()`], this method does not fail if the channel is closed. /// /// When borrowed from the [`Sender`] this function will always return `false`. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = watch::channel("hello"); /// /// tx.send("goodbye").unwrap(); /// // The sender does never consider the value as changed. /// assert!(!tx.borrow().has_changed()); /// /// // Drop the sender immediately, just for testing purposes. /// drop(tx); /// /// // Even if the sender has already been dropped... /// assert!(rx.has_changed().is_err()); /// // ...the modified value is still readable and detected as changed. /// assert_eq!(*rx.borrow(), "goodbye"); /// assert!(rx.borrow().has_changed()); /// /// // Read the changed value and mark it as seen. /// { /// let received = rx.borrow_and_update(); /// assert_eq!(*received, "goodbye"); /// assert!(received.has_changed()); /// // Release the read lock when leaving this scope. /// } /// /// // Now the value has already been marked as seen and could /// // never be modified again (after the sender has been dropped). /// assert!(!rx.borrow().has_changed()); /// } /// ``` pub fn has_changed(&self) -> bool { self.has_changed } } struct Shared { /// The most recent value. value: RwLock, /// The current version. /// /// The lowest bit represents a "closed" state. The rest of the bits /// represent the current version. state: AtomicState, /// Tracks the number of `Receiver` instances. ref_count_rx: AtomicUsize, /// Tracks the number of `Sender` instances. ref_count_tx: AtomicUsize, /// Notifies waiting receivers that the value changed. notify_rx: big_notify::BigNotify, /// Notifies any task listening for `Receiver` dropped events. notify_tx: Notify, } impl fmt::Debug for Shared { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let state = self.state.load(); f.debug_struct("Shared") .field("value", &self.value) .field("version", &state.version()) .field("is_closed", &state.is_closed()) .field("ref_count_rx", &self.ref_count_rx) .finish() } } pub mod error { //! Watch error types. use std::error::Error; use std::fmt; /// Error produced when sending a value fails. #[derive(PartialEq, Eq, Clone, Copy)] pub struct SendError(pub T); // ===== impl SendError ===== impl fmt::Debug for SendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("SendError").finish_non_exhaustive() } } impl fmt::Display for SendError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } impl Error for SendError {} /// Error produced when receiving a change notification. #[derive(Debug, Clone)] pub struct RecvError(pub(super) ()); // ===== impl RecvError ===== impl fmt::Display for RecvError { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } impl Error for RecvError {} } mod big_notify { use super::Notify; use crate::sync::notify::Notified; // To avoid contention on the lock inside the `Notify`, we store multiple // copies of it. Then, we use either circular access or randomness to spread // out threads over different `Notify` objects. // // Some simple benchmarks show that randomness performs slightly better than // circular access (probably due to contention on `next`), so we prefer to // use randomness when Tokio is compiled with a random number generator. // // When the random number generator is not available, we fall back to // circular access. pub(super) struct BigNotify { #[cfg(not(all(not(loom), feature = "sync", any(feature = "rt", feature = "macros"))))] next: std::sync::atomic::AtomicUsize, inner: [Notify; 8], } impl BigNotify { pub(super) fn new() -> Self { Self { #[cfg(not(all( not(loom), feature = "sync", any(feature = "rt", feature = "macros") )))] next: std::sync::atomic::AtomicUsize::new(0), inner: Default::default(), } } pub(super) fn notify_waiters(&self) { for notify in &self.inner { notify.notify_waiters(); } } /// This function implements the case where randomness is not available. #[cfg(not(all(not(loom), feature = "sync", any(feature = "rt", feature = "macros"))))] pub(super) fn notified(&self) -> Notified<'_> { let i = self.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed) % 8; self.inner[i].notified() } /// This function implements the case where randomness is available. #[cfg(all(not(loom), feature = "sync", any(feature = "rt", feature = "macros")))] pub(super) fn notified(&self) -> Notified<'_> { let i = crate::runtime::context::thread_rng_n(8) as usize; self.inner[i].notified() } } } use self::state::{AtomicState, Version}; mod state { use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::atomic::Ordering; const CLOSED_BIT: usize = 1; // Using 2 as the step size preserves the `CLOSED_BIT`. const STEP_SIZE: usize = 2; /// The version part of the state. The lowest bit is always zero. #[derive(Copy, Clone, Debug, Eq, PartialEq)] pub(super) struct Version(usize); /// Snapshot of the state. The first bit is used as the CLOSED bit. /// The remaining bits are used as the version. /// /// The CLOSED bit tracks whether the Sender has been dropped. Dropping all /// receivers does not set it. #[derive(Copy, Clone, Debug)] pub(super) struct StateSnapshot(usize); /// The state stored in an atomic integer. /// /// The `Sender` uses `Release` ordering for storing a new state /// and the `Receiver`s use `Acquire` ordering for loading the /// current state. This ensures that written values are seen by /// the `Receiver`s for a proper handover. #[derive(Debug)] pub(super) struct AtomicState(AtomicUsize); impl Version { /// Decrements the version. pub(super) fn decrement(&mut self) { // Using a wrapping decrement here is required to ensure that the // operation is consistent with `std::sync::atomic::AtomicUsize::fetch_add()` // which wraps on overflow. self.0 = self.0.wrapping_sub(STEP_SIZE); } pub(super) const INITIAL: Self = Version(0); } impl StateSnapshot { /// Extract the version from the state. pub(super) fn version(self) -> Version { Version(self.0 & !CLOSED_BIT) } /// Is the closed bit set? pub(super) fn is_closed(self) -> bool { (self.0 & CLOSED_BIT) == CLOSED_BIT } } impl AtomicState { /// Create a new `AtomicState` that is not closed and which has the /// version set to `Version::INITIAL`. pub(super) fn new() -> Self { AtomicState(AtomicUsize::new(Version::INITIAL.0)) } /// Load the current value of the state. /// /// Only used by the receiver and for debugging purposes. /// /// The receiver side (read-only) uses `Acquire` ordering for a proper handover /// of the shared value with the sender side (single writer). The state is always /// updated after modifying and before releasing the (exclusive) lock on the /// shared value. pub(super) fn load(&self) -> StateSnapshot { StateSnapshot(self.0.load(Ordering::Acquire)) } /// Increment the version counter. pub(super) fn increment_version_while_locked(&self) { // Use `Release` ordering to ensure that the shared value // has been written before updating the version. The shared // value is still protected by an exclusive lock during this // method. self.0.fetch_add(STEP_SIZE, Ordering::Release); } /// Set the closed bit in the state. pub(super) fn set_closed(&self) { self.0.fetch_or(CLOSED_BIT, Ordering::Release); } } } /// Creates a new watch channel, returning the "send" and "receive" handles. /// /// All values sent by [`Sender`] will become visible to the [`Receiver`] handles. /// Only the last value sent is made available to the [`Receiver`] half. All /// intermediate values are dropped. /// /// # Examples /// /// The following example prints `hello! world! `. /// /// ``` /// use tokio::sync::watch; /// use tokio::time::{Duration, sleep}; /// /// # async fn dox() -> Result<(), Box> { /// let (tx, mut rx) = watch::channel("hello"); /// /// tokio::spawn(async move { /// // Use the equivalent of a "do-while" loop so the initial value is /// // processed before awaiting the `changed()` future. /// loop { /// println!("{}! ", *rx.borrow_and_update()); /// if rx.changed().await.is_err() { /// break; /// } /// } /// }); /// /// sleep(Duration::from_millis(100)).await; /// tx.send("world")?; /// # Ok(()) /// # } /// ``` /// /// [`Sender`]: struct@Sender /// [`Receiver`]: struct@Receiver pub fn channel(init: T) -> (Sender, Receiver) { let shared = Arc::new(Shared { value: RwLock::new(init), state: AtomicState::new(), ref_count_rx: AtomicUsize::new(1), ref_count_tx: AtomicUsize::new(1), notify_rx: big_notify::BigNotify::new(), notify_tx: Notify::new(), }); let tx = Sender { shared: shared.clone(), }; let rx = Receiver { shared, version: Version::INITIAL, }; (tx, rx) } impl Receiver { fn from_shared(version: Version, shared: Arc>) -> Self { // No synchronization necessary as this is only used as a counter and // not memory access. shared.ref_count_rx.fetch_add(1, Relaxed); Self { shared, version } } /// Returns a reference to the most recently sent value. /// /// This method does not mark the returned value as seen, so future calls to /// [`changed`] may return immediately even if you have already seen the /// value with a call to `borrow`. /// /// Outstanding borrows hold a read lock on the inner value. This means that /// long-lived borrows could cause the producer half to block. It is recommended /// to keep the borrow as short-lived as possible. Additionally, if you are /// running in an environment that allows `!Send` futures, you must ensure that /// the returned `Ref` type is never held alive across an `.await` point, /// otherwise, it can lead to a deadlock. /// /// The priority policy of the lock is dependent on the underlying lock /// implementation, and this type does not guarantee that any particular policy /// will be used. In particular, a producer which is waiting to acquire the lock /// in `send` might or might not block concurrent calls to `borrow`, e.g.: /// ///
Potential deadlock example /// /// ```text /// // Task 1 (on thread A) | // Task 2 (on thread B) /// let _ref1 = rx.borrow(); | /// | // will block /// | let _ = tx.send(()); /// // may deadlock | /// let _ref2 = rx.borrow(); | /// ``` ///
/// /// For more information on when to use this method versus /// [`borrow_and_update`], see [here](self#borrow_and_update-versus-borrow). /// /// [`changed`]: Receiver::changed /// [`borrow_and_update`]: Receiver::borrow_and_update /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// let (_, rx) = watch::channel("hello"); /// assert_eq!(*rx.borrow(), "hello"); /// ``` pub fn borrow(&self) -> Ref<'_, T> { let inner = self.shared.value.read(); // After obtaining a read-lock no concurrent writes could occur // and the loaded version matches that of the borrowed reference. let new_version = self.shared.state.load().version(); let has_changed = self.version != new_version; Ref { inner, has_changed } } /// Returns a reference to the most recently sent value and marks that value /// as seen. /// /// This method marks the current value as seen. Subsequent calls to [`changed`] /// will not return immediately until the [`Sender`] has modified the shared /// value again. /// /// Outstanding borrows hold a read lock on the inner value. This means that /// long-lived borrows could cause the producer half to block. It is recommended /// to keep the borrow as short-lived as possible. Additionally, if you are /// running in an environment that allows `!Send` futures, you must ensure that /// the returned `Ref` type is never held alive across an `.await` point, /// otherwise, it can lead to a deadlock. /// /// The priority policy of the lock is dependent on the underlying lock /// implementation, and this type does not guarantee that any particular policy /// will be used. In particular, a producer which is waiting to acquire the lock /// in `send` might or might not block concurrent calls to `borrow`, e.g.: /// ///
Potential deadlock example /// /// ```text /// // Task 1 (on thread A) | // Task 2 (on thread B) /// let _ref1 = rx1.borrow_and_update(); | /// | // will block /// | let _ = tx.send(()); /// // may deadlock | /// let _ref2 = rx2.borrow_and_update(); | /// ``` ///
/// /// For more information on when to use this method versus [`borrow`], see /// [here](self#borrow_and_update-versus-borrow). /// /// [`changed`]: Receiver::changed /// [`borrow`]: Receiver::borrow pub fn borrow_and_update(&mut self) -> Ref<'_, T> { let inner = self.shared.value.read(); // After obtaining a read-lock no concurrent writes could occur // and the loaded version matches that of the borrowed reference. let new_version = self.shared.state.load().version(); let has_changed = self.version != new_version; // Mark the shared value as seen by updating the version self.version = new_version; Ref { inner, has_changed } } /// Checks if this channel contains a message that this receiver has not yet /// seen. The new value is not marked as seen. /// /// Although this method is called `has_changed`, it does not check new /// messages for equality, so this call will return true even if the new /// message is equal to the old message. /// /// Returns an error if the channel has been closed. /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = watch::channel("hello"); /// /// tx.send("goodbye").unwrap(); /// /// assert!(rx.has_changed().unwrap()); /// assert_eq!(*rx.borrow_and_update(), "goodbye"); /// /// // The value has been marked as seen /// assert!(!rx.has_changed().unwrap()); /// /// drop(tx); /// // The `tx` handle has been dropped /// assert!(rx.has_changed().is_err()); /// } /// ``` pub fn has_changed(&self) -> Result { // Load the version from the state let state = self.shared.state.load(); if state.is_closed() { // The sender has dropped. return Err(error::RecvError(())); } let new_version = state.version(); Ok(self.version != new_version) } /// Marks the state as changed. /// /// After invoking this method [`has_changed()`](Self::has_changed) /// returns `true` and [`changed()`](Self::changed) returns /// immediately, regardless of whether a new value has been sent. /// /// This is useful for triggering an initial change notification after /// subscribing to synchronize new receivers. pub fn mark_changed(&mut self) { self.version.decrement(); } /// Marks the state as unchanged. /// /// The current value will be considered seen by the receiver. /// /// This is useful if you are not interested in the current value /// visible in the receiver. pub fn mark_unchanged(&mut self) { let current_version = self.shared.state.load().version(); self.version = current_version; } /// Waits for a change notification, then marks the newest value as seen. /// /// If the newest value in the channel has not yet been marked seen when /// this method is called, the method marks that value seen and returns /// immediately. If the newest value has already been marked seen, then the /// method sleeps until a new message is sent by the [`Sender`] connected to /// this `Receiver`, or until the [`Sender`] is dropped. /// /// This method returns an error if and only if the [`Sender`] is dropped. /// /// For more information, see /// [*Change notifications*](self#change-notifications) in the module-level documentation. /// /// # Cancel safety /// /// This method is cancel safe. If you use it as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that no values have been marked /// seen by this call to `changed`. /// /// [`Sender`]: struct@Sender /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = watch::channel("hello"); /// /// tokio::spawn(async move { /// tx.send("goodbye").unwrap(); /// }); /// /// assert!(rx.changed().await.is_ok()); /// assert_eq!(*rx.borrow_and_update(), "goodbye"); /// /// // The `tx` handle has been dropped /// assert!(rx.changed().await.is_err()); /// } /// ``` pub async fn changed(&mut self) -> Result<(), error::RecvError> { cooperative(changed_impl(&self.shared, &mut self.version)).await } /// Waits for a value that satisfies the provided condition. /// /// This method will call the provided closure whenever something is sent on /// the channel. Once the closure returns `true`, this method will return a /// reference to the value that was passed to the closure. /// /// Before `wait_for` starts waiting for changes, it will call the closure /// on the current value. If the closure returns `true` when given the /// current value, then `wait_for` will immediately return a reference to /// the current value. This is the case even if the current value is already /// considered seen. /// /// The watch channel only keeps track of the most recent value, so if /// several messages are sent faster than `wait_for` is able to call the /// closure, then it may skip some updates. Whenever the closure is called, /// it will be called with the most recent value. /// /// When this function returns, the value that was passed to the closure /// when it returned `true` will be considered seen. /// /// If the channel is closed, then `wait_for` will return a [`RecvError`]. /// Once this happens, no more messages can ever be sent on the channel. /// When an error is returned, it is guaranteed that the closure has been /// called on the last value, and that it returned `false` for that value. /// (If the closure returned `true`, then the last value would have been /// returned instead of the error.) /// /// Like the [`borrow`] method, the returned borrow holds a read lock on the /// inner value. This means that long-lived borrows could cause the producer /// half to block. It is recommended to keep the borrow as short-lived as /// possible. See the documentation of `borrow` for more information on /// this. /// /// [`borrow`]: Receiver::borrow /// [`RecvError`]: error::RecvError /// /// # Cancel safety /// /// This method is cancel safe. If you use it as the event in a /// [`tokio::select!`](crate::select) statement and some other branch /// completes first, then it is guaranteed that the last seen value `val` /// (if any) satisfies `f(val) == false`. /// /// # Panics /// /// If and only if the closure `f` panics. In that case, no resource owned /// or shared by this [`Receiver`] will be poisoned. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// use tokio::time::{sleep, Duration}; /// /// #[tokio::main(flavor = "current_thread", start_paused = true)] /// async fn main() { /// let (tx, mut rx) = watch::channel("hello"); /// /// tokio::spawn(async move { /// sleep(Duration::from_secs(1)).await; /// tx.send("goodbye").unwrap(); /// }); /// /// assert!(rx.wait_for(|val| *val == "goodbye").await.is_ok()); /// assert_eq!(*rx.borrow(), "goodbye"); /// } /// ``` pub async fn wait_for( &mut self, f: impl FnMut(&T) -> bool, ) -> Result, error::RecvError> { cooperative(self.wait_for_inner(f)).await } async fn wait_for_inner( &mut self, mut f: impl FnMut(&T) -> bool, ) -> Result, error::RecvError> { let mut closed = false; loop { { let inner = self.shared.value.read(); let new_version = self.shared.state.load().version(); let has_changed = self.version != new_version; self.version = new_version; if !closed || has_changed { let result = panic::catch_unwind(panic::AssertUnwindSafe(|| f(&inner))); match result { Ok(true) => { return Ok(Ref { inner, has_changed }); } Ok(false) => { // Skip the value. } Err(panicked) => { // Drop the read-lock to avoid poisoning it. drop(inner); // Forward the panic to the caller. panic::resume_unwind(panicked); // Unreachable } }; } } if closed { return Err(error::RecvError(())); } // Wait for the value to change. closed = changed_impl(&self.shared, &mut self.version).await.is_err(); } } /// Returns `true` if receivers belong to the same channel. /// /// # Examples /// /// ``` /// let (tx, rx) = tokio::sync::watch::channel(true); /// let rx2 = rx.clone(); /// assert!(rx.same_channel(&rx2)); /// /// let (tx3, rx3) = tokio::sync::watch::channel(true); /// assert!(!rx3.same_channel(&rx2)); /// ``` pub fn same_channel(&self, other: &Self) -> bool { Arc::ptr_eq(&self.shared, &other.shared) } cfg_process_driver! { pub(crate) fn try_has_changed(&mut self) -> Option> { maybe_changed(&self.shared, &mut self.version) } } } fn maybe_changed( shared: &Shared, version: &mut Version, ) -> Option> { // Load the version from the state let state = shared.state.load(); let new_version = state.version(); if *version != new_version { // Observe the new version and return *version = new_version; return Some(Ok(())); } if state.is_closed() { // The sender has been dropped. return Some(Err(error::RecvError(()))); } None } async fn changed_impl( shared: &Shared, version: &mut Version, ) -> Result<(), error::RecvError> { crate::trace::async_trace_leaf().await; loop { // In order to avoid a race condition, we first request a notification, // **then** check the current value's version. If a new version exists, // the notification request is dropped. let notified = shared.notify_rx.notified(); if let Some(ret) = maybe_changed(shared, version) { return ret; } notified.await; // loop around again in case the wake-up was spurious } } impl Clone for Receiver { fn clone(&self) -> Self { let version = self.version; let shared = self.shared.clone(); Self::from_shared(version, shared) } } impl Drop for Receiver { fn drop(&mut self) { // No synchronization necessary as this is only used as a counter and // not memory access. if 1 == self.shared.ref_count_rx.fetch_sub(1, Relaxed) { // This is the last `Receiver` handle, tasks waiting on `Sender::closed()` self.shared.notify_tx.notify_waiters(); } } } impl Sender { /// Creates the sending-half of the [`watch`] channel. /// /// See documentation of [`watch::channel`] for errors when calling this function. /// Beware that attempting to send a value when there are no receivers will /// return an error. /// /// [`watch`]: crate::sync::watch /// [`watch::channel`]: crate::sync::watch /// /// # Examples /// ``` /// let sender = tokio::sync::watch::Sender::new(0u8); /// assert!(sender.send(3).is_err()); /// let _rec = sender.subscribe(); /// assert!(sender.send(4).is_ok()); /// ``` pub fn new(init: T) -> Self { let (tx, _) = channel(init); tx } /// Sends a new value via the channel, notifying all receivers. /// /// This method fails if the channel is closed, which is the case when /// every receiver has been dropped. It is possible to reopen the channel /// using the [`subscribe`] method. However, when `send` fails, the value /// isn't made available for future receivers (but returned with the /// [`SendError`]). /// /// To always make a new value available for future receivers, even if no /// receiver currently exists, one of the other send methods /// ([`send_if_modified`], [`send_modify`], or [`send_replace`]) can be /// used instead. /// /// [`subscribe`]: Sender::subscribe /// [`SendError`]: error::SendError /// [`send_if_modified`]: Sender::send_if_modified /// [`send_modify`]: Sender::send_modify /// [`send_replace`]: Sender::send_replace pub fn send(&self, value: T) -> Result<(), error::SendError> { // This is pretty much only useful as a hint anyway, so synchronization isn't critical. if 0 == self.receiver_count() { return Err(error::SendError(value)); } self.send_replace(value); Ok(()) } /// Modifies the watched value **unconditionally** in-place, /// notifying all receivers. /// /// This can be useful for modifying the watched value, without /// having to allocate a new instance. Additionally, this /// method permits sending values even when there are no receivers. /// /// Prefer to use the more versatile function [`Self::send_if_modified()`] /// if the value is only modified conditionally during the mutable borrow /// to prevent unneeded change notifications for unmodified values. /// /// # Panics /// /// This function panics when the invocation of the `modify` closure panics. /// No receivers are notified when panicking. All changes of the watched /// value applied by the closure before panicking will be visible in /// subsequent calls to `borrow`. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// struct State { /// counter: usize, /// } /// let (state_tx, state_rx) = watch::channel(State { counter: 0 }); /// state_tx.send_modify(|state| state.counter += 1); /// assert_eq!(state_rx.borrow().counter, 1); /// ``` pub fn send_modify(&self, modify: F) where F: FnOnce(&mut T), { self.send_if_modified(|value| { modify(value); true }); } /// Modifies the watched value **conditionally** in-place, /// notifying all receivers only if modified. /// /// This can be useful for modifying the watched value, without /// having to allocate a new instance. Additionally, this /// method permits sending values even when there are no receivers. /// /// The `modify` closure must return `true` if the value has actually /// been modified during the mutable borrow. It should only return `false` /// if the value is guaranteed to be unmodified despite the mutable /// borrow. /// /// Receivers are only notified if the closure returned `true`. If the /// closure has modified the value but returned `false` this results /// in a *silent modification*, i.e. the modified value will be visible /// in subsequent calls to `borrow`, but receivers will not receive /// a change notification. /// /// Returns the result of the closure, i.e. `true` if the value has /// been modified and `false` otherwise. /// /// # Panics /// /// This function panics when the invocation of the `modify` closure panics. /// No receivers are notified when panicking. All changes of the watched /// value applied by the closure before panicking will be visible in /// subsequent calls to `borrow`. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// struct State { /// counter: usize, /// } /// let (state_tx, mut state_rx) = watch::channel(State { counter: 1 }); /// let inc_counter_if_odd = |state: &mut State| { /// if state.counter % 2 == 1 { /// state.counter += 1; /// return true; /// } /// false /// }; /// /// assert_eq!(state_rx.borrow().counter, 1); /// /// assert!(!state_rx.has_changed().unwrap()); /// assert!(state_tx.send_if_modified(inc_counter_if_odd)); /// assert!(state_rx.has_changed().unwrap()); /// assert_eq!(state_rx.borrow_and_update().counter, 2); /// /// assert!(!state_rx.has_changed().unwrap()); /// assert!(!state_tx.send_if_modified(inc_counter_if_odd)); /// assert!(!state_rx.has_changed().unwrap()); /// assert_eq!(state_rx.borrow_and_update().counter, 2); /// ``` pub fn send_if_modified(&self, modify: F) -> bool where F: FnOnce(&mut T) -> bool, { { // Acquire the write lock and update the value. let mut lock = self.shared.value.write(); // Update the value and catch possible panic inside func. let result = panic::catch_unwind(panic::AssertUnwindSafe(|| modify(&mut lock))); match result { Ok(modified) => { if !modified { // Abort, i.e. don't notify receivers if unmodified return false; } // Continue if modified } Err(panicked) => { // Drop the lock to avoid poisoning it. drop(lock); // Forward the panic to the caller. panic::resume_unwind(panicked); // Unreachable } }; self.shared.state.increment_version_while_locked(); // Release the write lock. // // Incrementing the version counter while holding the lock ensures // that receivers are able to figure out the version number of the // value they are currently looking at. drop(lock); } self.shared.notify_rx.notify_waiters(); true } /// Sends a new value via the channel, notifying all receivers and returning /// the previous value in the channel. /// /// This can be useful for reusing the buffers inside a watched value. /// Additionally, this method permits sending values even when there are no /// receivers. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// let (tx, _rx) = watch::channel(1); /// assert_eq!(tx.send_replace(2), 1); /// assert_eq!(tx.send_replace(3), 2); /// ``` pub fn send_replace(&self, mut value: T) -> T { // swap old watched value with the new one self.send_modify(|old| mem::swap(old, &mut value)); value } /// Returns a reference to the most recently sent value /// /// Outstanding borrows hold a read lock on the inner value. This means that /// long-lived borrows could cause the producer half to block. It is recommended /// to keep the borrow as short-lived as possible. Additionally, if you are /// running in an environment that allows `!Send` futures, you must ensure that /// the returned `Ref` type is never held alive across an `.await` point, /// otherwise, it can lead to a deadlock. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// let (tx, _) = watch::channel("hello"); /// assert_eq!(*tx.borrow(), "hello"); /// ``` pub fn borrow(&self) -> Ref<'_, T> { let inner = self.shared.value.read(); // The sender/producer always sees the current version let has_changed = false; Ref { inner, has_changed } } /// Checks if the channel has been closed. This happens when all receivers /// have dropped. /// /// # Examples /// /// ``` /// let (tx, rx) = tokio::sync::watch::channel(()); /// assert!(!tx.is_closed()); /// /// drop(rx); /// assert!(tx.is_closed()); /// ``` pub fn is_closed(&self) -> bool { self.receiver_count() == 0 } /// Completes when all receivers have dropped. /// /// This allows the producer to get notified when interest in the produced /// values is canceled and immediately stop doing work. Once a channel is /// closed, the only way to reopen it is to call [`Sender::subscribe`] to /// get a new receiver. /// /// If the channel becomes closed for a brief amount of time (e.g., the last /// receiver is dropped and then `subscribe` is called), then this call to /// `closed` might return, but it is also possible that it does not "notice" /// that the channel was closed for a brief amount of time. /// /// # Cancel safety /// /// This method is cancel safe. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx) = watch::channel("hello"); /// /// tokio::spawn(async move { /// // use `rx` /// drop(rx); /// }); /// /// // Waits for `rx` to drop /// tx.closed().await; /// println!("the `rx` handles dropped") /// } /// ``` pub async fn closed(&self) { cooperative(async { crate::trace::async_trace_leaf().await; while self.receiver_count() > 0 { let notified = self.shared.notify_tx.notified(); if self.receiver_count() == 0 { return; } notified.await; // The channel could have been reopened in the meantime by calling // `subscribe`, so we loop again. } }) .await; } /// Creates a new [`Receiver`] connected to this `Sender`. /// /// All messages sent before this call to `subscribe` are initially marked /// as seen by the new `Receiver`. /// /// This method can be called even if there are no other receivers. In this /// case, the channel is reopened. /// /// # Examples /// /// The new channel will receive messages sent on this `Sender`. /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = watch::channel(0u64); /// /// tx.send(5).unwrap(); /// /// let rx = tx.subscribe(); /// assert_eq!(5, *rx.borrow()); /// /// tx.send(10).unwrap(); /// assert_eq!(10, *rx.borrow()); /// } /// ``` /// /// The most recent message is considered seen by the channel, so this test /// is guaranteed to pass. /// /// ``` /// use tokio::sync::watch; /// use tokio::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let (tx, _rx) = watch::channel(0u64); /// tx.send(5).unwrap(); /// let mut rx = tx.subscribe(); /// /// tokio::spawn(async move { /// // by spawning and sleeping, the message is sent after `main` /// // hits the call to `changed`. /// # if false { /// tokio::time::sleep(Duration::from_millis(10)).await; /// # } /// tx.send(100).unwrap(); /// }); /// /// rx.changed().await.unwrap(); /// assert_eq!(100, *rx.borrow()); /// } /// ``` pub fn subscribe(&self) -> Receiver { let shared = self.shared.clone(); let version = shared.state.load().version(); // The CLOSED bit in the state tracks only whether the sender is // dropped, so we do not need to unset it if this reopens the channel. Receiver::from_shared(version, shared) } /// Returns the number of receivers that currently exist. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx, rx1) = watch::channel("hello"); /// /// assert_eq!(1, tx.receiver_count()); /// /// let mut _rx2 = rx1.clone(); /// /// assert_eq!(2, tx.receiver_count()); /// } /// ``` pub fn receiver_count(&self) -> usize { self.shared.ref_count_rx.load(Relaxed) } /// Returns the number of senders that currently exist. /// /// # Examples /// /// ``` /// use tokio::sync::watch; /// /// #[tokio::main] /// async fn main() { /// let (tx1, rx) = watch::channel("hello"); /// /// assert_eq!(1, tx1.sender_count()); /// /// let tx2 = tx1.clone(); /// /// assert_eq!(2, tx1.sender_count()); /// assert_eq!(2, tx2.sender_count()); /// } /// ``` pub fn sender_count(&self) -> usize { self.shared.ref_count_tx.load(Relaxed) } /// Returns `true` if senders belong to the same channel. /// /// # Examples /// /// ``` /// let (tx, rx) = tokio::sync::watch::channel(true); /// let tx2 = tx.clone(); /// assert!(tx.same_channel(&tx2)); /// /// let (tx3, rx3) = tokio::sync::watch::channel(true); /// assert!(!tx3.same_channel(&tx2)); /// ``` pub fn same_channel(&self, other: &Self) -> bool { Arc::ptr_eq(&self.shared, &other.shared) } } impl Drop for Sender { fn drop(&mut self) { if self.shared.ref_count_tx.fetch_sub(1, AcqRel) == 1 { self.shared.state.set_closed(); self.shared.notify_rx.notify_waiters(); } } } // ===== impl Ref ===== impl ops::Deref for Ref<'_, T> { type Target = T; fn deref(&self) -> &T { self.inner.deref() } } #[cfg(all(test, loom))] mod tests { use futures::future::FutureExt; use loom::thread; // test for https://github.com/tokio-rs/tokio/issues/3168 #[test] fn watch_spurious_wakeup() { loom::model(|| { let (send, mut recv) = crate::sync::watch::channel(0i32); send.send(1).unwrap(); let send_thread = thread::spawn(move || { send.send(2).unwrap(); send }); recv.changed().now_or_never(); let send = send_thread.join().unwrap(); let recv_thread = thread::spawn(move || { recv.changed().now_or_never(); recv.changed().now_or_never(); recv }); send.send(3).unwrap(); let mut recv = recv_thread.join().unwrap(); let send_thread = thread::spawn(move || { send.send(2).unwrap(); }); recv.changed().now_or_never(); send_thread.join().unwrap(); }); } #[test] fn watch_borrow() { loom::model(|| { let (send, mut recv) = crate::sync::watch::channel(0i32); assert!(send.borrow().eq(&0)); assert!(recv.borrow().eq(&0)); send.send(1).unwrap(); assert!(send.borrow().eq(&1)); let send_thread = thread::spawn(move || { send.send(2).unwrap(); send }); recv.changed().now_or_never(); let send = send_thread.join().unwrap(); let recv_thread = thread::spawn(move || { recv.changed().now_or_never(); recv.changed().now_or_never(); recv }); send.send(3).unwrap(); let recv = recv_thread.join().unwrap(); assert!(recv.borrow().eq(&3)); assert!(send.borrow().eq(&3)); send.send(2).unwrap(); thread::spawn(move || { assert!(recv.borrow().eq(&2)); }); assert!(send.borrow().eq(&2)); }); } } tokio-1.43.1/src/task/blocking.rs000064400000000000000000000214061046102023000147310ustar 00000000000000use crate::task::JoinHandle; cfg_rt_multi_thread! { /// Runs the provided blocking function on the current thread without /// blocking the executor. /// /// In general, issuing a blocking call or performing a lot of compute in a /// future without yielding is problematic, as it may prevent the executor /// from driving other tasks forward. Calling this function informs the /// executor that the currently executing task is about to block the thread, /// so the executor is able to hand off any other tasks it has to a new /// worker thread before that happens. See the [CPU-bound tasks and blocking /// code][blocking] section for more information. /// /// Be aware that although this function avoids starving other independently /// spawned tasks, any other code running concurrently in the same task will /// be suspended during the call to `block_in_place`. This can happen e.g. /// when using the [`join!`] macro. To avoid this issue, use /// [`spawn_blocking`] instead of `block_in_place`. /// /// Note that this function cannot be used within a [`current_thread`] runtime /// because in this case there are no other worker threads to hand off tasks /// to. On the other hand, calling the function outside a runtime is /// allowed. In this case, `block_in_place` just calls the provided closure /// normally. /// /// Code running behind `block_in_place` cannot be cancelled. When you shut /// down the executor, it will wait indefinitely for all blocking operations /// to finish. You can use [`shutdown_timeout`] to stop waiting for them /// after a certain timeout. Be aware that this will still not cancel the /// tasks — they are simply allowed to keep running after the method /// returns. /// /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code /// [`spawn_blocking`]: fn@crate::task::spawn_blocking /// [`join!`]: macro@join /// [`thread::spawn`]: fn@std::thread::spawn /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout /// /// # Examples /// /// ``` /// use tokio::task; /// /// # async fn docs() { /// task::block_in_place(move || { /// // do some compute-heavy work or call synchronous code /// }); /// # } /// ``` /// /// Code running inside `block_in_place` may use `block_on` to reenter the /// async context. /// /// ``` /// use tokio::task; /// use tokio::runtime::Handle; /// /// # async fn docs() { /// task::block_in_place(move || { /// Handle::current().block_on(async move { /// // do something async /// }); /// }); /// # } /// ``` /// /// # Panics /// /// This function panics if called from a [`current_thread`] runtime. /// /// [`current_thread`]: fn@crate::runtime::Builder::new_current_thread #[track_caller] pub fn block_in_place(f: F) -> R where F: FnOnce() -> R, { crate::runtime::scheduler::block_in_place(f) } } cfg_rt! { /// Runs the provided closure on a thread where blocking is acceptable. /// /// In general, issuing a blocking call or performing a lot of compute in a /// future without yielding is problematic, as it may prevent the executor from /// driving other futures forward. This function runs the provided closure on a /// thread dedicated to blocking operations. See the [CPU-bound tasks and /// blocking code][blocking] section for more information. /// /// Tokio will spawn more blocking threads when they are requested through this /// function until the upper limit configured on the [`Builder`] is reached. /// After reaching the upper limit, the tasks are put in a queue. /// The thread limit is very large by default, because `spawn_blocking` is often /// used for various kinds of IO operations that cannot be performed /// asynchronously. When you run CPU-bound code using `spawn_blocking`, you /// should keep this large upper limit in mind. When running many CPU-bound /// computations, a semaphore or some other synchronization primitive should be /// used to limit the number of computation executed in parallel. Specialized /// CPU-bound executors, such as [rayon], may also be a good fit. /// /// This function is intended for non-async operations that eventually finish on /// their own. If you want to spawn an ordinary thread, you should use /// [`thread::spawn`] instead. /// /// Be aware that tasks spawned using `spawn_blocking` cannot be aborted /// because they are not async. If you call [`abort`] on a `spawn_blocking` /// task, then this *will not have any effect*, and the task will continue /// running normally. The exception is if the task has not started running /// yet; in that case, calling `abort` may prevent the task from starting. /// /// When you shut down the executor, it will wait indefinitely for all blocking operations to /// finish. You can use [`shutdown_timeout`] to stop waiting for them after a /// certain timeout. Be aware that this will still not cancel the tasks — they /// are simply allowed to keep running after the method returns. It is possible /// for a blocking task to be cancelled if it has not yet started running, but this /// is not guaranteed. /// /// Note that if you are using the single threaded runtime, this function will /// still spawn additional threads for blocking operations. The current-thread /// scheduler's single thread is only used for asynchronous code. /// /// # Related APIs and patterns for bridging asynchronous and blocking code /// /// In simple cases, it is sufficient to have the closure accept input /// parameters at creation time and return a single value (or struct/tuple, etc.). /// /// For more complex situations in which it is desirable to stream data to or from /// the synchronous context, the [`mpsc channel`] has `blocking_send` and /// `blocking_recv` methods for use in non-async code such as the thread created /// by `spawn_blocking`. /// /// Another option is [`SyncIoBridge`] for cases where the synchronous context /// is operating on byte streams. For example, you might use an asynchronous /// HTTP client such as [hyper] to fetch data, but perform complex parsing /// of the payload body using a library written for synchronous I/O. /// /// Finally, see also [Bridging with sync code][bridgesync] for discussions /// around the opposite case of using Tokio as part of a larger synchronous /// codebase. /// /// [`Builder`]: struct@crate::runtime::Builder /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code /// [rayon]: https://docs.rs/rayon /// [`mpsc channel`]: crate::sync::mpsc /// [`SyncIoBridge`]: https://docs.rs/tokio-util/latest/tokio_util/io/struct.SyncIoBridge.html /// [hyper]: https://docs.rs/hyper /// [`thread::spawn`]: fn@std::thread::spawn /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout /// [bridgesync]: https://tokio.rs/tokio/topics/bridging /// [`AtomicBool`]: struct@std::sync::atomic::AtomicBool /// [`abort`]: crate::task::JoinHandle::abort /// /// # Examples /// /// Pass an input value and receive result of computation: /// /// ``` /// use tokio::task; /// /// # async fn docs() -> Result<(), Box>{ /// // Initial input /// let mut v = "Hello, ".to_string(); /// let res = task::spawn_blocking(move || { /// // Stand-in for compute-heavy work or using synchronous APIs /// v.push_str("world"); /// // Pass ownership of the value back to the asynchronous context /// v /// }).await?; /// /// // `res` is the value returned from the thread /// assert_eq!(res.as_str(), "Hello, world"); /// # Ok(()) /// # } /// ``` /// /// Use a channel: /// /// ``` /// use tokio::task; /// use tokio::sync::mpsc; /// /// # async fn docs() { /// let (tx, mut rx) = mpsc::channel(2); /// let start = 5; /// let worker = task::spawn_blocking(move || { /// for x in 0..10 { /// // Stand in for complex computation /// tx.blocking_send(start + x).unwrap(); /// } /// }); /// /// let mut acc = 0; /// while let Some(v) = rx.recv().await { /// acc += v; /// } /// assert_eq!(acc, 95); /// worker.await.unwrap(); /// # } /// ``` #[track_caller] pub fn spawn_blocking(f: F) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { crate::runtime::spawn_blocking(f) } } tokio-1.43.1/src/task/builder.rs000064400000000000000000000164401046102023000145710ustar 00000000000000#![allow(unreachable_pub)] use crate::{ runtime::{Handle, BOX_FUTURE_THRESHOLD}, task::{JoinHandle, LocalSet}, util::trace::SpawnMeta, }; use std::{future::Future, io, mem}; /// Factory which is used to configure the properties of a new task. /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// Methods can be chained in order to configure it. /// /// Currently, there is only one configuration option: /// /// - [`name`], which specifies an associated name for /// the task /// /// There are three types of task that can be spawned from a Builder: /// - [`spawn_local`] for executing futures on the current thread /// - [`spawn`] for executing [`Send`] futures on the runtime /// - [`spawn_blocking`] for executing blocking code in the /// blocking thread pool. /// /// ## Example /// /// ```no_run /// use tokio::net::{TcpListener, TcpStream}; /// /// use std::io; /// /// async fn process(socket: TcpStream) { /// // ... /// # drop(socket); /// } /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:8080").await?; /// /// loop { /// let (socket, _) = listener.accept().await?; /// /// tokio::task::Builder::new() /// .name("tcp connection handler") /// .spawn(async move { /// // Process each socket concurrently. /// process(socket).await /// })?; /// } /// } /// ``` /// [unstable]: crate#unstable-features /// [`name`]: Builder::name /// [`spawn_local`]: Builder::spawn_local /// [`spawn`]: Builder::spawn /// [`spawn_blocking`]: Builder::spawn_blocking #[derive(Default, Debug)] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] pub struct Builder<'a> { name: Option<&'a str>, } impl<'a> Builder<'a> { /// Creates a new task builder. pub fn new() -> Self { Self::default() } /// Assigns a name to the task which will be spawned. pub fn name(&self, name: &'a str) -> Self { Self { name: Some(name) } } /// Spawns a task with this builder's settings on the current runtime. /// /// # Panics /// /// This method panics if called outside of a Tokio runtime. /// /// See [`task::spawn`](crate::task::spawn()) for /// more details. #[track_caller] pub fn spawn(self, future: Fut) -> io::Result> where Fut: Future + Send + 'static, Fut::Output: Send + 'static, { let fut_size = mem::size_of::(); Ok(if fut_size > BOX_FUTURE_THRESHOLD { super::spawn::spawn_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size)) } else { super::spawn::spawn_inner(future, SpawnMeta::new(self.name, fut_size)) }) } /// Spawn a task with this builder's settings on the provided [runtime /// handle]. /// /// See [`Handle::spawn`] for more details. /// /// [runtime handle]: crate::runtime::Handle /// [`Handle::spawn`]: crate::runtime::Handle::spawn #[track_caller] pub fn spawn_on(self, future: Fut, handle: &Handle) -> io::Result> where Fut: Future + Send + 'static, Fut::Output: Send + 'static, { let fut_size = mem::size_of::(); Ok(if fut_size > BOX_FUTURE_THRESHOLD { handle.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size)) } else { handle.spawn_named(future, SpawnMeta::new(self.name, fut_size)) }) } /// Spawns `!Send` a task on the current [`LocalSet`] with this builder's /// settings. /// /// The spawned future will be run on the same thread that called `spawn_local`. /// This may only be called from the context of a [local task set][`LocalSet`]. /// /// # Panics /// /// This function panics if called outside of a [local task set][`LocalSet`]. /// /// See [`task::spawn_local`] for more details. /// /// [`task::spawn_local`]: crate::task::spawn_local /// [`LocalSet`]: crate::task::LocalSet #[track_caller] pub fn spawn_local(self, future: Fut) -> io::Result> where Fut: Future + 'static, Fut::Output: 'static, { let fut_size = mem::size_of::(); Ok(if fut_size > BOX_FUTURE_THRESHOLD { super::local::spawn_local_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size)) } else { super::local::spawn_local_inner(future, SpawnMeta::new(self.name, fut_size)) }) } /// Spawns `!Send` a task on the provided [`LocalSet`] with this builder's /// settings. /// /// See [`LocalSet::spawn_local`] for more details. /// /// [`LocalSet::spawn_local`]: crate::task::LocalSet::spawn_local /// [`LocalSet`]: crate::task::LocalSet #[track_caller] pub fn spawn_local_on( self, future: Fut, local_set: &LocalSet, ) -> io::Result> where Fut: Future + 'static, Fut::Output: 'static, { let fut_size = mem::size_of::(); Ok(if fut_size > BOX_FUTURE_THRESHOLD { local_set.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size)) } else { local_set.spawn_named(future, SpawnMeta::new(self.name, fut_size)) }) } /// Spawns blocking code on the blocking threadpool. /// /// # Panics /// /// This method panics if called outside of a Tokio runtime. /// /// See [`task::spawn_blocking`](crate::task::spawn_blocking) /// for more details. #[track_caller] pub fn spawn_blocking( self, function: Function, ) -> io::Result> where Function: FnOnce() -> Output + Send + 'static, Output: Send + 'static, { let handle = Handle::current(); self.spawn_blocking_on(function, &handle) } /// Spawns blocking code on the provided [runtime handle]'s blocking threadpool. /// /// See [`Handle::spawn_blocking`] for more details. /// /// [runtime handle]: crate::runtime::Handle /// [`Handle::spawn_blocking`]: crate::runtime::Handle::spawn_blocking #[track_caller] pub fn spawn_blocking_on( self, function: Function, handle: &Handle, ) -> io::Result> where Function: FnOnce() -> Output + Send + 'static, Output: Send + 'static, { use crate::runtime::Mandatory; let fn_size = mem::size_of::(); let (join_handle, spawn_result) = if fn_size > BOX_FUTURE_THRESHOLD { handle.inner.blocking_spawner().spawn_blocking_inner( Box::new(function), Mandatory::NonMandatory, SpawnMeta::new(self.name, fn_size), handle, ) } else { handle.inner.blocking_spawner().spawn_blocking_inner( function, Mandatory::NonMandatory, SpawnMeta::new(self.name, fn_size), handle, ) }; spawn_result?; Ok(join_handle) } } tokio-1.43.1/src/task/consume_budget.rs000064400000000000000000000024061046102023000161430ustar 00000000000000use std::task::{ready, Poll}; /// Consumes a unit of budget and returns the execution back to the Tokio /// runtime *if* the task's coop budget was exhausted. /// /// The task will only yield if its entire coop budget has been exhausted. /// This function can be used in order to insert optional yield points into long /// computations that do not use Tokio resources like sockets or semaphores, /// without redundantly yielding to the runtime each time. /// /// # Examples /// /// Make sure that a function which returns a sum of (potentially lots of) /// iterated values is cooperative. /// /// ``` /// async fn sum_iterator(input: &mut impl std::iter::Iterator) -> i64 { /// let mut sum: i64 = 0; /// while let Some(i) = input.next() { /// sum += i; /// tokio::task::consume_budget().await /// } /// sum /// } /// ``` #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub async fn consume_budget() { let mut status = Poll::Pending; std::future::poll_fn(move |cx| { ready!(crate::trace::trace_leaf(cx)); if status.is_ready() { return status; } status = crate::runtime::coop::poll_proceed(cx).map(|restore| { restore.made_progress(); }); status }) .await } tokio-1.43.1/src/task/join_set.rs000064400000000000000000000647001046102023000147570ustar 00000000000000//! A collection of tasks spawned on a Tokio runtime. //! //! This module provides the [`JoinSet`] type, a collection which stores a set //! of spawned tasks and allows asynchronously awaiting the output of those //! tasks as they complete. See the documentation for the [`JoinSet`] type for //! details. use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use std::{fmt, panic}; use crate::runtime::Handle; use crate::task::Id; use crate::task::{unconstrained, AbortHandle, JoinError, JoinHandle, LocalSet}; use crate::util::IdleNotifiedSet; /// A collection of tasks spawned on a Tokio runtime. /// /// A `JoinSet` can be used to await the completion of some or all of the tasks /// in the set. The set is not ordered, and the tasks will be returned in the /// order they complete. /// /// All of the tasks must have the same return type `T`. /// /// When the `JoinSet` is dropped, all tasks in the `JoinSet` are immediately aborted. /// /// # Examples /// /// Spawn multiple tasks and wait for them. /// /// ``` /// use tokio::task::JoinSet; /// /// #[tokio::main] /// async fn main() { /// let mut set = JoinSet::new(); /// /// for i in 0..10 { /// set.spawn(async move { i }); /// } /// /// let mut seen = [false; 10]; /// while let Some(res) = set.join_next().await { /// let idx = res.unwrap(); /// seen[idx] = true; /// } /// /// for i in 0..10 { /// assert!(seen[i]); /// } /// } /// ``` #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub struct JoinSet { inner: IdleNotifiedSet>, } /// A variant of [`task::Builder`] that spawns tasks on a [`JoinSet`] rather /// than on the current default runtime. /// /// [`task::Builder`]: crate::task::Builder #[cfg(all(tokio_unstable, feature = "tracing"))] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] #[must_use = "builders do nothing unless used to spawn a task"] pub struct Builder<'a, T> { joinset: &'a mut JoinSet, builder: super::Builder<'a>, } impl JoinSet { /// Create a new `JoinSet`. pub fn new() -> Self { Self { inner: IdleNotifiedSet::new(), } } /// Returns the number of tasks currently in the `JoinSet`. pub fn len(&self) -> usize { self.inner.len() } /// Returns whether the `JoinSet` is empty. pub fn is_empty(&self) -> bool { self.inner.is_empty() } } impl JoinSet { /// Returns a [`Builder`] that can be used to configure a task prior to /// spawning it on this `JoinSet`. /// /// # Examples /// /// ``` /// use tokio::task::JoinSet; /// /// #[tokio::main] /// async fn main() -> std::io::Result<()> { /// let mut set = JoinSet::new(); /// /// // Use the builder to configure a task's name before spawning it. /// set.build_task() /// .name("my_task") /// .spawn(async { /* ... */ })?; /// /// Ok(()) /// } /// ``` #[cfg(all(tokio_unstable, feature = "tracing"))] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] pub fn build_task(&mut self) -> Builder<'_, T> { Builder { builder: super::Builder::new(), joinset: self, } } /// Spawn the provided task on the `JoinSet`, returning an [`AbortHandle`] /// that can be used to remotely cancel the task. /// /// The provided future will start running in the background immediately /// when this method is called, even if you don't await anything on this /// `JoinSet`. /// /// # Panics /// /// This method panics if called outside of a Tokio runtime. /// /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn(&mut self, task: F) -> AbortHandle where F: Future, F: Send + 'static, T: Send, { self.insert(crate::spawn(task)) } /// Spawn the provided task on the provided runtime and store it in this /// `JoinSet` returning an [`AbortHandle`] that can be used to remotely /// cancel the task. /// /// The provided future will start running in the background immediately /// when this method is called, even if you don't await anything on this /// `JoinSet`. /// /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_on(&mut self, task: F, handle: &Handle) -> AbortHandle where F: Future, F: Send + 'static, T: Send, { self.insert(handle.spawn(task)) } /// Spawn the provided task on the current [`LocalSet`] and store it in this /// `JoinSet`, returning an [`AbortHandle`] that can be used to remotely /// cancel the task. /// /// The provided future will start running in the background immediately /// when this method is called, even if you don't await anything on this /// `JoinSet`. /// /// # Panics /// /// This method panics if it is called outside of a `LocalSet`. /// /// [`LocalSet`]: crate::task::LocalSet /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_local(&mut self, task: F) -> AbortHandle where F: Future, F: 'static, { self.insert(crate::task::spawn_local(task)) } /// Spawn the provided task on the provided [`LocalSet`] and store it in /// this `JoinSet`, returning an [`AbortHandle`] that can be used to /// remotely cancel the task. /// /// Unlike the [`spawn_local`] method, this method may be used to spawn local /// tasks on a `LocalSet` that is _not_ currently running. The provided /// future will start running whenever the `LocalSet` is next started. /// /// [`LocalSet`]: crate::task::LocalSet /// [`AbortHandle`]: crate::task::AbortHandle /// [`spawn_local`]: Self::spawn_local #[track_caller] pub fn spawn_local_on(&mut self, task: F, local_set: &LocalSet) -> AbortHandle where F: Future, F: 'static, { self.insert(local_set.spawn_local(task)) } /// Spawn the blocking code on the blocking threadpool and store /// it in this `JoinSet`, returning an [`AbortHandle`] that can be /// used to remotely cancel the task. /// /// # Examples /// /// Spawn multiple blocking tasks and wait for them. /// /// ``` /// use tokio::task::JoinSet; /// /// #[tokio::main] /// async fn main() { /// let mut set = JoinSet::new(); /// /// for i in 0..10 { /// set.spawn_blocking(move || { i }); /// } /// /// let mut seen = [false; 10]; /// while let Some(res) = set.join_next().await { /// let idx = res.unwrap(); /// seen[idx] = true; /// } /// /// for i in 0..10 { /// assert!(seen[i]); /// } /// } /// ``` /// /// # Panics /// /// This method panics if called outside of a Tokio runtime. /// /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_blocking(&mut self, f: F) -> AbortHandle where F: FnOnce() -> T, F: Send + 'static, T: Send, { self.insert(crate::runtime::spawn_blocking(f)) } /// Spawn the blocking code on the blocking threadpool of the /// provided runtime and store it in this `JoinSet`, returning an /// [`AbortHandle`] that can be used to remotely cancel the task. /// /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_blocking_on(&mut self, f: F, handle: &Handle) -> AbortHandle where F: FnOnce() -> T, F: Send + 'static, T: Send, { self.insert(handle.spawn_blocking(f)) } fn insert(&mut self, jh: JoinHandle) -> AbortHandle { let abort = jh.abort_handle(); let mut entry = self.inner.insert_idle(jh); // Set the waker that is notified when the task completes. entry.with_value_and_context(|jh, ctx| jh.set_join_waker(ctx.waker())); abort } /// Waits until one of the tasks in the set completes and returns its output. /// /// Returns `None` if the set is empty. /// /// # Cancel Safety /// /// This method is cancel safe. If `join_next` is used as the event in a `tokio::select!` /// statement and some other branch completes first, it is guaranteed that no tasks were /// removed from this `JoinSet`. pub async fn join_next(&mut self) -> Option> { std::future::poll_fn(|cx| self.poll_join_next(cx)).await } /// Waits until one of the tasks in the set completes and returns its /// output, along with the [task ID] of the completed task. /// /// Returns `None` if the set is empty. /// /// When this method returns an error, then the id of the task that failed can be accessed /// using the [`JoinError::id`] method. /// /// # Cancel Safety /// /// This method is cancel safe. If `join_next_with_id` is used as the event in a `tokio::select!` /// statement and some other branch completes first, it is guaranteed that no tasks were /// removed from this `JoinSet`. /// /// [task ID]: crate::task::Id /// [`JoinError::id`]: fn@crate::task::JoinError::id pub async fn join_next_with_id(&mut self) -> Option> { std::future::poll_fn(|cx| self.poll_join_next_with_id(cx)).await } /// Tries to join one of the tasks in the set that has completed and return its output. /// /// Returns `None` if there are no completed tasks, or if the set is empty. pub fn try_join_next(&mut self) -> Option> { // Loop over all notified `JoinHandle`s to find one that's ready, or until none are left. loop { let mut entry = self.inner.try_pop_notified()?; let res = entry.with_value_and_context(|jh, ctx| { // Since this function is not async and cannot be forced to yield, we should // disable budgeting when we want to check for the `JoinHandle` readiness. Pin::new(&mut unconstrained(jh)).poll(ctx) }); if let Poll::Ready(res) = res { let _entry = entry.remove(); return Some(res); } } } /// Tries to join one of the tasks in the set that has completed and return its output, /// along with the [task ID] of the completed task. /// /// Returns `None` if there are no completed tasks, or if the set is empty. /// /// When this method returns an error, then the id of the task that failed can be accessed /// using the [`JoinError::id`] method. /// /// [task ID]: crate::task::Id /// [`JoinError::id`]: fn@crate::task::JoinError::id pub fn try_join_next_with_id(&mut self) -> Option> { // Loop over all notified `JoinHandle`s to find one that's ready, or until none are left. loop { let mut entry = self.inner.try_pop_notified()?; let res = entry.with_value_and_context(|jh, ctx| { // Since this function is not async and cannot be forced to yield, we should // disable budgeting when we want to check for the `JoinHandle` readiness. Pin::new(&mut unconstrained(jh)).poll(ctx) }); if let Poll::Ready(res) = res { let entry = entry.remove(); return Some(res.map(|output| (entry.id(), output))); } } } /// Aborts all tasks and waits for them to finish shutting down. /// /// Calling this method is equivalent to calling [`abort_all`] and then calling [`join_next`] in /// a loop until it returns `None`. /// /// This method ignores any panics in the tasks shutting down. When this call returns, the /// `JoinSet` will be empty. /// /// [`abort_all`]: fn@Self::abort_all /// [`join_next`]: fn@Self::join_next pub async fn shutdown(&mut self) { self.abort_all(); while self.join_next().await.is_some() {} } /// Awaits the completion of all tasks in this `JoinSet`, returning a vector of their results. /// /// The results will be stored in the order they completed not the order they were spawned. /// This is a convenience method that is equivalent to calling [`join_next`] in /// a loop. If any tasks on the `JoinSet` fail with an [`JoinError`], then this call /// to `join_all` will panic and all remaining tasks on the `JoinSet` are /// cancelled. To handle errors in any other way, manually call [`join_next`] /// in a loop. /// /// # Examples /// /// Spawn multiple tasks and `join_all` them. /// /// ``` /// use tokio::task::JoinSet; /// use std::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let mut set = JoinSet::new(); /// /// for i in 0..3 { /// set.spawn(async move { /// tokio::time::sleep(Duration::from_secs(3 - i)).await; /// i /// }); /// } /// /// let output = set.join_all().await; /// assert_eq!(output, vec![2, 1, 0]); /// } /// ``` /// /// Equivalent implementation of `join_all`, using [`join_next`] and loop. /// /// ``` /// use tokio::task::JoinSet; /// use std::panic; /// /// #[tokio::main] /// async fn main() { /// let mut set = JoinSet::new(); /// /// for i in 0..3 { /// set.spawn(async move {i}); /// } /// /// let mut output = Vec::new(); /// while let Some(res) = set.join_next().await{ /// match res { /// Ok(t) => output.push(t), /// Err(err) if err.is_panic() => panic::resume_unwind(err.into_panic()), /// Err(err) => panic!("{err}"), /// } /// } /// assert_eq!(output.len(),3); /// } /// ``` /// [`join_next`]: fn@Self::join_next /// [`JoinError::id`]: fn@crate::task::JoinError::id pub async fn join_all(mut self) -> Vec { let mut output = Vec::with_capacity(self.len()); while let Some(res) = self.join_next().await { match res { Ok(t) => output.push(t), Err(err) if err.is_panic() => panic::resume_unwind(err.into_panic()), Err(err) => panic!("{err}"), } } output } /// Aborts all tasks on this `JoinSet`. /// /// This does not remove the tasks from the `JoinSet`. To wait for the tasks to complete /// cancellation, you should call `join_next` in a loop until the `JoinSet` is empty. pub fn abort_all(&mut self) { self.inner.for_each(|jh| jh.abort()); } /// Removes all tasks from this `JoinSet` without aborting them. /// /// The tasks removed by this call will continue to run in the background even if the `JoinSet` /// is dropped. pub fn detach_all(&mut self) { self.inner.drain(drop); } /// Polls for one of the tasks in the set to complete. /// /// If this returns `Poll::Ready(Some(_))`, then the task that completed is removed from the set. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided `Context` is scheduled /// to receive a wakeup when a task in the `JoinSet` completes. Note that on multiple calls to /// `poll_join_next`, only the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. /// /// # Returns /// /// This function returns: /// /// * `Poll::Pending` if the `JoinSet` is not empty but there is no task whose output is /// available right now. /// * `Poll::Ready(Some(Ok(value)))` if one of the tasks in this `JoinSet` has completed. /// The `value` is the return value of one of the tasks that completed. /// * `Poll::Ready(Some(Err(err)))` if one of the tasks in this `JoinSet` has panicked or been /// aborted. The `err` is the `JoinError` from the panicked/aborted task. /// * `Poll::Ready(None)` if the `JoinSet` is empty. /// /// Note that this method may return `Poll::Pending` even if one of the tasks has completed. /// This can happen if the [coop budget] is reached. /// /// [coop budget]: crate::task#cooperative-scheduling pub fn poll_join_next(&mut self, cx: &mut Context<'_>) -> Poll>> { // The call to `pop_notified` moves the entry to the `idle` list. It is moved back to // the `notified` list if the waker is notified in the `poll` call below. let mut entry = match self.inner.pop_notified(cx.waker()) { Some(entry) => entry, None => { if self.is_empty() { return Poll::Ready(None); } else { // The waker was set by `pop_notified`. return Poll::Pending; } } }; let res = entry.with_value_and_context(|jh, ctx| Pin::new(jh).poll(ctx)); if let Poll::Ready(res) = res { let _entry = entry.remove(); Poll::Ready(Some(res)) } else { // A JoinHandle generally won't emit a wakeup without being ready unless // the coop limit has been reached. We yield to the executor in this // case. cx.waker().wake_by_ref(); Poll::Pending } } /// Polls for one of the tasks in the set to complete. /// /// If this returns `Poll::Ready(Some(_))`, then the task that completed is removed from the set. /// /// When the method returns `Poll::Pending`, the `Waker` in the provided `Context` is scheduled /// to receive a wakeup when a task in the `JoinSet` completes. Note that on multiple calls to /// `poll_join_next`, only the `Waker` from the `Context` passed to the most recent call is /// scheduled to receive a wakeup. /// /// # Returns /// /// This function returns: /// /// * `Poll::Pending` if the `JoinSet` is not empty but there is no task whose output is /// available right now. /// * `Poll::Ready(Some(Ok((id, value))))` if one of the tasks in this `JoinSet` has completed. /// The `value` is the return value of one of the tasks that completed, and /// `id` is the [task ID] of that task. /// * `Poll::Ready(Some(Err(err)))` if one of the tasks in this `JoinSet` has panicked or been /// aborted. The `err` is the `JoinError` from the panicked/aborted task. /// * `Poll::Ready(None)` if the `JoinSet` is empty. /// /// Note that this method may return `Poll::Pending` even if one of the tasks has completed. /// This can happen if the [coop budget] is reached. /// /// [coop budget]: crate::task#cooperative-scheduling /// [task ID]: crate::task::Id pub fn poll_join_next_with_id( &mut self, cx: &mut Context<'_>, ) -> Poll>> { // The call to `pop_notified` moves the entry to the `idle` list. It is moved back to // the `notified` list if the waker is notified in the `poll` call below. let mut entry = match self.inner.pop_notified(cx.waker()) { Some(entry) => entry, None => { if self.is_empty() { return Poll::Ready(None); } else { // The waker was set by `pop_notified`. return Poll::Pending; } } }; let res = entry.with_value_and_context(|jh, ctx| Pin::new(jh).poll(ctx)); if let Poll::Ready(res) = res { let entry = entry.remove(); // If the task succeeded, add the task ID to the output. Otherwise, the // `JoinError` will already have the task's ID. Poll::Ready(Some(res.map(|output| (entry.id(), output)))) } else { // A JoinHandle generally won't emit a wakeup without being ready unless // the coop limit has been reached. We yield to the executor in this // case. cx.waker().wake_by_ref(); Poll::Pending } } } impl Drop for JoinSet { fn drop(&mut self) { self.inner.drain(|join_handle| join_handle.abort()); } } impl fmt::Debug for JoinSet { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("JoinSet").field("len", &self.len()).finish() } } impl Default for JoinSet { fn default() -> Self { Self::new() } } /// Collect an iterator of futures into a [`JoinSet`]. /// /// This is equivalent to calling [`JoinSet::spawn`] on each element of the iterator. /// /// # Examples /// /// The main example from [`JoinSet`]'s documentation can also be written using [`collect`]: /// /// ``` /// use tokio::task::JoinSet; /// /// #[tokio::main] /// async fn main() { /// let mut set: JoinSet<_> = (0..10).map(|i| async move { i }).collect(); /// /// let mut seen = [false; 10]; /// while let Some(res) = set.join_next().await { /// let idx = res.unwrap(); /// seen[idx] = true; /// } /// /// for i in 0..10 { /// assert!(seen[i]); /// } /// } /// ``` /// /// [`collect`]: std::iter::Iterator::collect impl std::iter::FromIterator for JoinSet where F: Future, F: Send + 'static, T: Send + 'static, { fn from_iter>(iter: I) -> Self { let mut set = Self::new(); iter.into_iter().for_each(|task| { set.spawn(task); }); set } } // === impl Builder === #[cfg(all(tokio_unstable, feature = "tracing"))] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] impl<'a, T: 'static> Builder<'a, T> { /// Assigns a name to the task which will be spawned. pub fn name(self, name: &'a str) -> Self { let builder = self.builder.name(name); Self { builder, ..self } } /// Spawn the provided task with this builder's settings and store it in the /// [`JoinSet`], returning an [`AbortHandle`] that can be used to remotely /// cancel the task. /// /// # Returns /// /// An [`AbortHandle`] that can be used to remotely cancel the task. /// /// # Panics /// /// This method panics if called outside of a Tokio runtime. /// /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn(self, future: F) -> std::io::Result where F: Future, F: Send + 'static, T: Send, { Ok(self.joinset.insert(self.builder.spawn(future)?)) } /// Spawn the provided task on the provided [runtime handle] with this /// builder's settings, and store it in the [`JoinSet`]. /// /// # Returns /// /// An [`AbortHandle`] that can be used to remotely cancel the task. /// /// /// [`AbortHandle`]: crate::task::AbortHandle /// [runtime handle]: crate::runtime::Handle #[track_caller] pub fn spawn_on(self, future: F, handle: &Handle) -> std::io::Result where F: Future, F: Send + 'static, T: Send, { Ok(self.joinset.insert(self.builder.spawn_on(future, handle)?)) } /// Spawn the blocking code on the blocking threadpool with this builder's /// settings, and store it in the [`JoinSet`]. /// /// # Returns /// /// An [`AbortHandle`] that can be used to remotely cancel the task. /// /// # Panics /// /// This method panics if called outside of a Tokio runtime. /// /// [`JoinSet`]: crate::task::JoinSet /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_blocking(self, f: F) -> std::io::Result where F: FnOnce() -> T, F: Send + 'static, T: Send, { Ok(self.joinset.insert(self.builder.spawn_blocking(f)?)) } /// Spawn the blocking code on the blocking threadpool of the provided /// runtime handle with this builder's settings, and store it in the /// [`JoinSet`]. /// /// # Returns /// /// An [`AbortHandle`] that can be used to remotely cancel the task. /// /// [`JoinSet`]: crate::task::JoinSet /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_blocking_on(self, f: F, handle: &Handle) -> std::io::Result where F: FnOnce() -> T, F: Send + 'static, T: Send, { Ok(self .joinset .insert(self.builder.spawn_blocking_on(f, handle)?)) } /// Spawn the provided task on the current [`LocalSet`] with this builder's /// settings, and store it in the [`JoinSet`]. /// /// # Returns /// /// An [`AbortHandle`] that can be used to remotely cancel the task. /// /// # Panics /// /// This method panics if it is called outside of a `LocalSet`. /// /// [`LocalSet`]: crate::task::LocalSet /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_local(self, future: F) -> std::io::Result where F: Future, F: 'static, { Ok(self.joinset.insert(self.builder.spawn_local(future)?)) } /// Spawn the provided task on the provided [`LocalSet`] with this builder's /// settings, and store it in the [`JoinSet`]. /// /// # Returns /// /// An [`AbortHandle`] that can be used to remotely cancel the task. /// /// [`LocalSet`]: crate::task::LocalSet /// [`AbortHandle`]: crate::task::AbortHandle #[track_caller] pub fn spawn_local_on(self, future: F, local_set: &LocalSet) -> std::io::Result where F: Future, F: 'static, { Ok(self .joinset .insert(self.builder.spawn_local_on(future, local_set)?)) } } // Manual `Debug` impl so that `Builder` is `Debug` regardless of whether `T` is // `Debug`. #[cfg(all(tokio_unstable, feature = "tracing"))] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] impl<'a, T> fmt::Debug for Builder<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("join_set::Builder") .field("joinset", &self.joinset) .field("builder", &self.builder) .finish() } } tokio-1.43.1/src/task/local.rs000064400000000000000000001313501046102023000142330ustar 00000000000000//! Runs `!Send` futures on the current thread. use crate::loom::cell::UnsafeCell; use crate::loom::sync::{Arc, Mutex}; #[cfg(tokio_unstable)] use crate::runtime; use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task, TaskHarnessScheduleHooks}; use crate::runtime::{context, ThreadId, BOX_FUTURE_THRESHOLD}; use crate::sync::AtomicWaker; use crate::util::trace::SpawnMeta; use crate::util::RcCell; use std::cell::Cell; use std::collections::VecDeque; use std::fmt; use std::future::Future; use std::marker::PhantomData; use std::mem; use std::pin::Pin; use std::rc::Rc; use std::task::Poll; use pin_project_lite::pin_project; cfg_rt! { /// A set of tasks which are executed on the same thread. /// /// In some cases, it is necessary to run one or more futures that do not /// implement [`Send`] and thus are unsafe to send between threads. In these /// cases, a [local task set] may be used to schedule one or more `!Send` /// futures to run together on the same thread. /// /// For example, the following code will not compile: /// /// ```rust,compile_fail /// use std::rc::Rc; /// /// #[tokio::main] /// async fn main() { /// // `Rc` does not implement `Send`, and thus may not be sent between /// // threads safely. /// let nonsend_data = Rc::new("my nonsend data..."); /// /// let nonsend_data = nonsend_data.clone(); /// // Because the `async` block here moves `nonsend_data`, the future is `!Send`. /// // Since `tokio::spawn` requires the spawned future to implement `Send`, this /// // will not compile. /// tokio::spawn(async move { /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// } /// ``` /// /// # Use with `run_until` /// /// To spawn `!Send` futures, we can use a local task set to schedule them /// on the thread calling [`Runtime::block_on`]. When running inside of the /// local task set, we can use [`task::spawn_local`], which can spawn /// `!Send` futures. For example: /// /// ```rust /// use std::rc::Rc; /// use tokio::task; /// /// #[tokio::main] /// async fn main() { /// let nonsend_data = Rc::new("my nonsend data..."); /// /// // Construct a local task set that can run `!Send` futures. /// let local = task::LocalSet::new(); /// /// // Run the local task set. /// local.run_until(async move { /// let nonsend_data = nonsend_data.clone(); /// // `spawn_local` ensures that the future is spawned on the local /// // task set. /// task::spawn_local(async move { /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// }).await; /// } /// ``` /// **Note:** The `run_until` method can only be used in `#[tokio::main]`, /// `#[tokio::test]` or directly inside a call to [`Runtime::block_on`]. It /// cannot be used inside a task spawned with `tokio::spawn`. /// /// ## Awaiting a `LocalSet` /// /// Additionally, a `LocalSet` itself implements `Future`, completing when /// *all* tasks spawned on the `LocalSet` complete. This can be used to run /// several futures on a `LocalSet` and drive the whole set until they /// complete. For example, /// /// ```rust /// use tokio::{task, time}; /// use std::rc::Rc; /// /// #[tokio::main] /// async fn main() { /// let nonsend_data = Rc::new("world"); /// let local = task::LocalSet::new(); /// /// let nonsend_data2 = nonsend_data.clone(); /// local.spawn_local(async move { /// // ... /// println!("hello {}", nonsend_data2) /// }); /// /// local.spawn_local(async move { /// time::sleep(time::Duration::from_millis(100)).await; /// println!("goodbye {}", nonsend_data) /// }); /// /// // ... /// /// local.await; /// } /// ``` /// **Note:** Awaiting a `LocalSet` can only be done inside /// `#[tokio::main]`, `#[tokio::test]` or directly inside a call to /// [`Runtime::block_on`]. It cannot be used inside a task spawned with /// `tokio::spawn`. /// /// ## Use inside `tokio::spawn` /// /// The two methods mentioned above cannot be used inside `tokio::spawn`, so /// to spawn `!Send` futures from inside `tokio::spawn`, we need to do /// something else. The solution is to create the `LocalSet` somewhere else, /// and communicate with it using an [`mpsc`] channel. /// /// The following example puts the `LocalSet` inside a new thread. /// ``` /// use tokio::runtime::Builder; /// use tokio::sync::{mpsc, oneshot}; /// use tokio::task::LocalSet; /// /// // This struct describes the task you want to spawn. Here we include /// // some simple examples. The oneshot channel allows sending a response /// // to the spawner. /// #[derive(Debug)] /// enum Task { /// PrintNumber(u32), /// AddOne(u32, oneshot::Sender), /// } /// /// #[derive(Clone)] /// struct LocalSpawner { /// send: mpsc::UnboundedSender, /// } /// /// impl LocalSpawner { /// pub fn new() -> Self { /// let (send, mut recv) = mpsc::unbounded_channel(); /// /// let rt = Builder::new_current_thread() /// .enable_all() /// .build() /// .unwrap(); /// /// std::thread::spawn(move || { /// let local = LocalSet::new(); /// /// local.spawn_local(async move { /// while let Some(new_task) = recv.recv().await { /// tokio::task::spawn_local(run_task(new_task)); /// } /// // If the while loop returns, then all the LocalSpawner /// // objects have been dropped. /// }); /// /// // This will return once all senders are dropped and all /// // spawned tasks have returned. /// rt.block_on(local); /// }); /// /// Self { /// send, /// } /// } /// /// pub fn spawn(&self, task: Task) { /// self.send.send(task).expect("Thread with LocalSet has shut down."); /// } /// } /// /// // This task may do !Send stuff. We use printing a number as an example, /// // but it could be anything. /// // /// // The Task struct is an enum to support spawning many different kinds /// // of operations. /// async fn run_task(task: Task) { /// match task { /// Task::PrintNumber(n) => { /// println!("{}", n); /// }, /// Task::AddOne(n, response) => { /// // We ignore failures to send the response. /// let _ = response.send(n + 1); /// }, /// } /// } /// /// #[tokio::main] /// async fn main() { /// let spawner = LocalSpawner::new(); /// /// let (send, response) = oneshot::channel(); /// spawner.spawn(Task::AddOne(10, send)); /// let eleven = response.await.unwrap(); /// assert_eq!(eleven, 11); /// } /// ``` /// /// [`Send`]: trait@std::marker::Send /// [local task set]: struct@LocalSet /// [`Runtime::block_on`]: method@crate::runtime::Runtime::block_on /// [`task::spawn_local`]: fn@spawn_local /// [`mpsc`]: mod@crate::sync::mpsc pub struct LocalSet { /// Current scheduler tick. tick: Cell, /// State available from thread-local. context: Rc, /// This type should not be Send. _not_send: PhantomData<*const ()>, } } /// State available from the thread-local. struct Context { /// State shared between threads. shared: Arc, /// True if a task panicked without being handled and the local set is /// configured to shutdown on unhandled panic. unhandled_panic: Cell, } /// `LocalSet` state shared between threads. struct Shared { /// # Safety /// /// This field must *only* be accessed from the thread that owns the /// `LocalSet` (i.e., `Thread::current().id() == owner`). local_state: LocalState, /// Remote run queue sender. queue: Mutex>>>>, /// Wake the `LocalSet` task. waker: AtomicWaker, /// How to respond to unhandled task panics. #[cfg(tokio_unstable)] pub(crate) unhandled_panic: crate::runtime::UnhandledPanic, } /// Tracks the `LocalSet` state that must only be accessed from the thread that /// created the `LocalSet`. struct LocalState { /// The `ThreadId` of the thread that owns the `LocalSet`. owner: ThreadId, /// Local run queue sender and receiver. local_queue: UnsafeCell>>>, /// Collection of all active tasks spawned onto this executor. owned: LocalOwnedTasks>, } pin_project! { #[derive(Debug)] struct RunUntil<'a, F> { local_set: &'a LocalSet, #[pin] future: F, } } tokio_thread_local!(static CURRENT: LocalData = const { LocalData { ctx: RcCell::new(), wake_on_schedule: Cell::new(false), } }); struct LocalData { ctx: RcCell, wake_on_schedule: Cell, } impl LocalData { /// Should be called except when we call `LocalSet::enter`. /// Especially when we poll a `LocalSet`. #[must_use = "dropping this guard will reset the entered state"] fn enter(&self, ctx: Rc) -> LocalDataEnterGuard<'_> { let ctx = self.ctx.replace(Some(ctx)); let wake_on_schedule = self.wake_on_schedule.replace(false); LocalDataEnterGuard { local_data_ref: self, ctx, wake_on_schedule, } } } /// A guard for `LocalData::enter()` struct LocalDataEnterGuard<'a> { local_data_ref: &'a LocalData, ctx: Option>, wake_on_schedule: bool, } impl<'a> Drop for LocalDataEnterGuard<'a> { fn drop(&mut self) { self.local_data_ref.ctx.set(self.ctx.take()); self.local_data_ref .wake_on_schedule .set(self.wake_on_schedule) } } cfg_rt! { /// Spawns a `!Send` future on the current [`LocalSet`] or [`LocalRuntime`]. /// /// The spawned future will run on the same thread that called `spawn_local`. /// /// The provided future will start running in the background immediately /// when `spawn_local` is called, even if you don't await the returned /// `JoinHandle`. /// /// # Panics /// /// This function panics if called outside of a [`LocalSet`]. /// /// Note that if [`tokio::spawn`] is used from within a `LocalSet`, the /// resulting new task will _not_ be inside the `LocalSet`, so you must use /// `spawn_local` if you want to stay within the `LocalSet`. /// /// # Examples /// /// ```rust /// use std::rc::Rc; /// use tokio::task; /// /// #[tokio::main] /// async fn main() { /// let nonsend_data = Rc::new("my nonsend data..."); /// /// let local = task::LocalSet::new(); /// /// // Run the local task set. /// local.run_until(async move { /// let nonsend_data = nonsend_data.clone(); /// task::spawn_local(async move { /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// }).await; /// } /// ``` /// /// [`LocalSet`]: struct@crate::task::LocalSet /// [`LocalRuntime`]: struct@crate::runtime::LocalRuntime /// [`tokio::spawn`]: fn@crate::task::spawn #[track_caller] pub fn spawn_local(future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, { let fut_size = std::mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { spawn_local_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { spawn_local_inner(future, SpawnMeta::new_unnamed(fut_size)) } } #[track_caller] pub(super) fn spawn_local_inner(future: F, meta: SpawnMeta<'_>) -> JoinHandle where F: Future + 'static, F::Output: 'static { use crate::runtime::{context, task}; let mut future = Some(future); let res = context::with_current(|handle| { Some(if handle.is_local() { if !handle.can_spawn_local_on_local_runtime() { return None; } let future = future.take().unwrap(); #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any( target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64" ) ))] let future = task::trace::Trace::root(future); let id = task::Id::next(); let task = crate::util::trace::task(future, "task", meta, id.as_u64()); // safety: we have verified that this is a `LocalRuntime` owned by the current thread unsafe { handle.spawn_local(task, id) } } else { match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) { None => panic!("`spawn_local` called from outside of a `task::LocalSet` or LocalRuntime"), Some(cx) => cx.spawn(future.take().unwrap(), meta) } }) }); match res { Ok(None) => panic!("Local tasks can only be spawned on a LocalRuntime from the thread the runtime was created on"), Ok(Some(join_handle)) => join_handle, Err(_) => match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) { None => panic!("`spawn_local` called from outside of a `task::LocalSet` or LocalRuntime"), Some(cx) => cx.spawn(future.unwrap(), meta) } } } } /// Initial queue capacity. const INITIAL_CAPACITY: usize = 64; /// Max number of tasks to poll per tick. const MAX_TASKS_PER_TICK: usize = 61; /// How often it check the remote queue first. const REMOTE_FIRST_INTERVAL: u8 = 31; /// Context guard for `LocalSet` pub struct LocalEnterGuard { ctx: Option>, /// Distinguishes whether the context was entered or being polled. /// When we enter it, the value `wake_on_schedule` is set. In this case /// `spawn_local` refers the context, whereas it is not being polled now. wake_on_schedule: bool, } impl Drop for LocalEnterGuard { fn drop(&mut self) { CURRENT.with( |LocalData { ctx, wake_on_schedule, }| { ctx.set(self.ctx.take()); wake_on_schedule.set(self.wake_on_schedule); }, ); } } impl fmt::Debug for LocalEnterGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("LocalEnterGuard").finish() } } impl LocalSet { /// Returns a new local task set. pub fn new() -> LocalSet { let owner = context::thread_id().expect("cannot create LocalSet during thread shutdown"); LocalSet { tick: Cell::new(0), context: Rc::new(Context { shared: Arc::new(Shared { local_state: LocalState { owner, owned: LocalOwnedTasks::new(), local_queue: UnsafeCell::new(VecDeque::with_capacity(INITIAL_CAPACITY)), }, queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))), waker: AtomicWaker::new(), #[cfg(tokio_unstable)] unhandled_panic: crate::runtime::UnhandledPanic::Ignore, }), unhandled_panic: Cell::new(false), }), _not_send: PhantomData, } } /// Enters the context of this `LocalSet`. /// /// The [`spawn_local`] method will spawn tasks on the `LocalSet` whose /// context you are inside. /// /// [`spawn_local`]: fn@crate::task::spawn_local pub fn enter(&self) -> LocalEnterGuard { CURRENT.with( |LocalData { ctx, wake_on_schedule, .. }| { let ctx = ctx.replace(Some(self.context.clone())); let wake_on_schedule = wake_on_schedule.replace(true); LocalEnterGuard { ctx, wake_on_schedule, } }, ) } /// Spawns a `!Send` task onto the local task set. /// /// This task is guaranteed to be run on the current thread. /// /// Unlike the free function [`spawn_local`], this method may be used to /// spawn local tasks when the `LocalSet` is _not_ running. The provided /// future will start running once the `LocalSet` is next started, even if /// you don't await the returned `JoinHandle`. /// /// # Examples /// /// ```rust /// use tokio::task; /// /// #[tokio::main] /// async fn main() { /// let local = task::LocalSet::new(); /// /// // Spawn a future on the local set. This future will be run when /// // we call `run_until` to drive the task set. /// local.spawn_local(async { /// // ... /// }); /// /// // Run the local task set. /// local.run_until(async move { /// // ... /// }).await; /// /// // When `run` finishes, we can spawn _more_ futures, which will /// // run in subsequent calls to `run_until`. /// local.spawn_local(async { /// // ... /// }); /// /// local.run_until(async move { /// // ... /// }).await; /// } /// ``` /// [`spawn_local`]: fn@spawn_local #[track_caller] pub fn spawn_local(&self, future: F) -> JoinHandle where F: Future + 'static, F::Output: 'static, { let fut_size = mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { self.spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { self.spawn_named(future, SpawnMeta::new_unnamed(fut_size)) } } /// Runs a future to completion on the provided runtime, driving any local /// futures spawned on this task set on the current thread. /// /// This runs the given future on the runtime, blocking until it is /// complete, and yielding its resolved result. Any tasks or timers which /// the future spawns internally will be executed on the runtime. The future /// may also call [`spawn_local`] to `spawn_local` additional local futures on the /// current thread. /// /// This method should not be called from an asynchronous context. /// /// # Panics /// /// This function panics if the executor is at capacity, if the provided /// future panics, or if called within an asynchronous execution context. /// /// # Notes /// /// Since this function internally calls [`Runtime::block_on`], and drives /// futures in the local task set inside that call to `block_on`, the local /// futures may not use [in-place blocking]. If a blocking call needs to be /// issued from a local task, the [`spawn_blocking`] API may be used instead. /// /// For example, this will panic: /// ```should_panic /// use tokio::runtime::Runtime; /// use tokio::task; /// /// let rt = Runtime::new().unwrap(); /// let local = task::LocalSet::new(); /// local.block_on(&rt, async { /// let join = task::spawn_local(async { /// let blocking_result = task::block_in_place(|| { /// // ... /// }); /// // ... /// }); /// join.await.unwrap(); /// }) /// ``` /// This, however, will not panic: /// ``` /// use tokio::runtime::Runtime; /// use tokio::task; /// /// let rt = Runtime::new().unwrap(); /// let local = task::LocalSet::new(); /// local.block_on(&rt, async { /// let join = task::spawn_local(async { /// let blocking_result = task::spawn_blocking(|| { /// // ... /// }).await; /// // ... /// }); /// join.await.unwrap(); /// }) /// ``` /// /// [`spawn_local`]: fn@spawn_local /// [`Runtime::block_on`]: method@crate::runtime::Runtime::block_on /// [in-place blocking]: fn@crate::task::block_in_place /// [`spawn_blocking`]: fn@crate::task::spawn_blocking #[track_caller] #[cfg(feature = "rt")] #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub fn block_on(&self, rt: &crate::runtime::Runtime, future: F) -> F::Output where F: Future, { rt.block_on(self.run_until(future)) } /// Runs a future to completion on the local set, returning its output. /// /// This returns a future that runs the given future with a local set, /// allowing it to call [`spawn_local`] to spawn additional `!Send` futures. /// Any local futures spawned on the local set will be driven in the /// background until the future passed to `run_until` completes. When the future /// passed to `run_until` finishes, any local futures which have not completed /// will remain on the local set, and will be driven on subsequent calls to /// `run_until` or when [awaiting the local set] itself. /// /// # Cancel safety /// /// This method is cancel safe when `future` is cancel safe. /// /// # Examples /// /// ```rust /// use tokio::task; /// /// #[tokio::main] /// async fn main() { /// task::LocalSet::new().run_until(async { /// task::spawn_local(async move { /// // ... /// }).await.unwrap(); /// // ... /// }).await; /// } /// ``` /// /// [`spawn_local`]: fn@spawn_local /// [awaiting the local set]: #awaiting-a-localset pub async fn run_until(&self, future: F) -> F::Output where F: Future, { let run_until = RunUntil { future, local_set: self, }; run_until.await } #[track_caller] pub(in crate::task) fn spawn_named( &self, future: F, meta: SpawnMeta<'_>, ) -> JoinHandle where F: Future + 'static, F::Output: 'static, { self.spawn_named_inner(future, meta) } #[track_caller] fn spawn_named_inner(&self, future: F, meta: SpawnMeta<'_>) -> JoinHandle where F: Future + 'static, F::Output: 'static, { let handle = self.context.spawn(future, meta); // Because a task was spawned from *outside* the `LocalSet`, wake the // `LocalSet` future to execute the new task, if it hasn't been woken. // // Spawning via the free fn `spawn` does not require this, as it can // only be called from *within* a future executing on the `LocalSet` — // in that case, the `LocalSet` must already be awake. self.context.shared.waker.wake(); handle } /// Ticks the scheduler, returning whether the local future needs to be /// notified again. fn tick(&self) -> bool { for _ in 0..MAX_TASKS_PER_TICK { // Make sure we didn't hit an unhandled panic assert!(!self.context.unhandled_panic.get(), "a spawned task panicked and the LocalSet is configured to shutdown on unhandled panic"); match self.next_task() { // Run the task // // Safety: As spawned tasks are `!Send`, `run_unchecked` must be // used. We are responsible for maintaining the invariant that // `run_unchecked` is only called on threads that spawned the // task initially. Because `LocalSet` itself is `!Send`, and // `spawn_local` spawns into the `LocalSet` on the current // thread, the invariant is maintained. Some(task) => crate::runtime::coop::budget(|| task.run()), // We have fully drained the queue of notified tasks, so the // local future doesn't need to be notified again — it can wait // until something else wakes a task in the local set. None => return false, } } true } fn next_task(&self) -> Option>> { let tick = self.tick.get(); self.tick.set(tick.wrapping_add(1)); let task = if tick % REMOTE_FIRST_INTERVAL == 0 { self.context .shared .queue .lock() .as_mut() .and_then(|queue| queue.pop_front()) .or_else(|| self.pop_local()) } else { self.pop_local().or_else(|| { self.context .shared .queue .lock() .as_mut() .and_then(VecDeque::pop_front) }) }; task.map(|task| unsafe { // Safety: because the `LocalSet` itself is `!Send`, we know we are // on the same thread if we have access to the `LocalSet`, and can // therefore access the local run queue. self.context.shared.local_state.assert_owner(task) }) } fn pop_local(&self) -> Option>> { unsafe { // Safety: because the `LocalSet` itself is `!Send`, we know we are // on the same thread if we have access to the `LocalSet`, and can // therefore access the local run queue. self.context.shared.local_state.task_pop_front() } } fn with(&self, f: impl FnOnce() -> T) -> T { CURRENT.with(|local_data| { let _guard = local_data.enter(self.context.clone()); f() }) } /// This method is like `with`, but it just calls `f` without setting the thread-local if that /// fails. fn with_if_possible(&self, f: impl FnOnce() -> T) -> T { let mut f = Some(f); let res = CURRENT.try_with(|local_data| { let _guard = local_data.enter(self.context.clone()); (f.take().unwrap())() }); match res { Ok(res) => res, Err(_access_error) => (f.take().unwrap())(), } } } cfg_unstable! { impl LocalSet { /// Configure how the `LocalSet` responds to an unhandled panic on a /// spawned task. /// /// By default, an unhandled panic (i.e. a panic not caught by /// [`std::panic::catch_unwind`]) has no impact on the `LocalSet`'s /// execution. The panic is error value is forwarded to the task's /// [`JoinHandle`] and all other spawned tasks continue running. /// /// The `unhandled_panic` option enables configuring this behavior. /// /// * `UnhandledPanic::Ignore` is the default behavior. Panics on /// spawned tasks have no impact on the `LocalSet`'s execution. /// * `UnhandledPanic::ShutdownRuntime` will force the `LocalSet` to /// shutdown immediately when a spawned task panics even if that /// task's `JoinHandle` has not been dropped. All other spawned tasks /// will immediately terminate and further calls to /// [`LocalSet::block_on`] and [`LocalSet::run_until`] will panic. /// /// # Panics /// /// This method panics if called after the `LocalSet` has started /// running. /// /// # Unstable /// /// This option is currently unstable and its implementation is /// incomplete. The API may change or be removed in the future. See /// tokio-rs/tokio#4516 for more details. /// /// # Examples /// /// The following demonstrates a `LocalSet` configured to shutdown on /// panic. The first spawned task panics and results in the `LocalSet` /// shutting down. The second spawned task never has a chance to /// execute. The call to `run_until` will panic due to the runtime being /// forcibly shutdown. /// /// ```should_panic /// use tokio::runtime::UnhandledPanic; /// /// # #[tokio::main] /// # async fn main() { /// tokio::task::LocalSet::new() /// .unhandled_panic(UnhandledPanic::ShutdownRuntime) /// .run_until(async { /// tokio::task::spawn_local(async { panic!("boom"); }); /// tokio::task::spawn_local(async { /// // This task never completes /// }); /// /// // Do some work, but `run_until` will panic before it completes /// # loop { tokio::task::yield_now().await; } /// }) /// .await; /// # } /// ``` /// /// [`JoinHandle`]: struct@crate::task::JoinHandle pub fn unhandled_panic(&mut self, behavior: crate::runtime::UnhandledPanic) -> &mut Self { // TODO: This should be set as a builder Rc::get_mut(&mut self.context) .and_then(|ctx| Arc::get_mut(&mut ctx.shared)) .expect("Unhandled Panic behavior modified after starting LocalSet") .unhandled_panic = behavior; self } /// Returns the [`Id`] of the current `LocalSet` runtime. /// /// # Examples /// /// ```rust /// use tokio::task; /// /// #[tokio::main] /// async fn main() { /// let local_set = task::LocalSet::new(); /// println!("Local set id: {}", local_set.id()); /// } /// ``` /// /// **Note**: This is an [unstable API][unstable]. The public API of this type /// may break in 1.x releases. See [the documentation on unstable /// features][unstable] for details. /// /// [unstable]: crate#unstable-features /// [`Id`]: struct@crate::runtime::Id pub fn id(&self) -> runtime::Id { self.context.shared.local_state.owned.id.into() } } } impl fmt::Debug for LocalSet { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("LocalSet").finish() } } impl Future for LocalSet { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { // Register the waker before starting to work self.context.shared.waker.register_by_ref(cx.waker()); if self.with(|| self.tick()) { // If `tick` returns true, we need to notify the local future again: // there are still tasks remaining in the run queue. cx.waker().wake_by_ref(); Poll::Pending // Safety: called from the thread that owns `LocalSet`. Because // `LocalSet` is `!Send`, this is safe. } else if unsafe { self.context.shared.local_state.owned_is_empty() } { // If the scheduler has no remaining futures, we're done! Poll::Ready(()) } else { // There are still futures in the local set, but we've polled all the // futures in the run queue. Therefore, we can just return Pending // since the remaining futures will be woken from somewhere else. Poll::Pending } } } impl Default for LocalSet { fn default() -> LocalSet { LocalSet::new() } } impl Drop for LocalSet { fn drop(&mut self) { self.with_if_possible(|| { // Shut down all tasks in the LocalOwnedTasks and close it to // prevent new tasks from ever being added. unsafe { // Safety: called from the thread that owns `LocalSet` self.context.shared.local_state.close_and_shutdown_all(); } // We already called shutdown on all tasks above, so there is no // need to call shutdown. // Safety: note that this *intentionally* bypasses the unsafe // `Shared::local_queue()` method. This is in order to avoid the // debug assertion that we are on the thread that owns the // `LocalSet`, because on some systems (e.g. at least some macOS // versions), attempting to get the current thread ID can panic due // to the thread's local data that stores the thread ID being // dropped *before* the `LocalSet`. // // Despite avoiding the assertion here, it is safe for us to access // the local queue in `Drop`, because the `LocalSet` itself is // `!Send`, so we can reasonably guarantee that it will not be // `Drop`ped from another thread. let local_queue = unsafe { // Safety: called from the thread that owns `LocalSet` self.context.shared.local_state.take_local_queue() }; for task in local_queue { drop(task); } // Take the queue from the Shared object to prevent pushing // notifications to it in the future. let queue = self.context.shared.queue.lock().take().unwrap(); for task in queue { drop(task); } // Safety: called from the thread that owns `LocalSet` assert!(unsafe { self.context.shared.local_state.owned_is_empty() }); }); } } // === impl Context === impl Context { #[track_caller] fn spawn(&self, future: F, meta: SpawnMeta<'_>) -> JoinHandle where F: Future + 'static, F::Output: 'static, { let id = crate::runtime::task::Id::next(); let future = crate::util::trace::task(future, "local", meta, id.as_u64()); // Safety: called from the thread that owns the `LocalSet` let (handle, notified) = { self.shared.local_state.assert_called_from_owner_thread(); self.shared .local_state .owned .bind(future, self.shared.clone(), id) }; if let Some(notified) = notified { self.shared.schedule(notified); } handle } } // === impl LocalFuture === impl Future for RunUntil<'_, T> { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { let me = self.project(); me.local_set.with(|| { me.local_set .context .shared .waker .register_by_ref(cx.waker()); let _no_blocking = crate::runtime::context::disallow_block_in_place(); let f = me.future; if let Poll::Ready(output) = f.poll(cx) { return Poll::Ready(output); } if me.local_set.tick() { // If `tick` returns `true`, we need to notify the local future again: // there are still tasks remaining in the run queue. cx.waker().wake_by_ref(); } Poll::Pending }) } } impl Shared { /// Schedule the provided task on the scheduler. fn schedule(&self, task: task::Notified>) { CURRENT.with(|localdata| { match localdata.ctx.get() { // If the current `LocalSet` is being polled, we don't need to wake it. // When we `enter` it, then the value `wake_on_schedule` is set to be true. // In this case it is not being polled, so we need to wake it. Some(cx) if cx.shared.ptr_eq(self) && !localdata.wake_on_schedule.get() => unsafe { // Safety: if the current `LocalSet` context points to this // `LocalSet`, then we are on the thread that owns it. cx.shared.local_state.task_push_back(task); }, // We are on the thread that owns the `LocalSet`, so we can // wake to the local queue. _ if context::thread_id().ok() == Some(self.local_state.owner) => { unsafe { // Safety: we just checked that the thread ID matches // the localset's owner, so this is safe. self.local_state.task_push_back(task); } // We still have to wake the `LocalSet`, because it isn't // currently being polled. self.waker.wake(); } // We are *not* on the thread that owns the `LocalSet`, so we // have to wake to the remote queue. _ => { // First, check whether the queue is still there (if not, the // LocalSet is dropped). Then push to it if so, and if not, // do nothing. let mut lock = self.queue.lock(); if let Some(queue) = lock.as_mut() { queue.push_back(task); drop(lock); self.waker.wake(); } } } }); } fn ptr_eq(&self, other: &Shared) -> bool { std::ptr::eq(self, other) } } // This is safe because (and only because) we *pinky pwomise* to never touch the // local run queue except from the thread that owns the `LocalSet`. unsafe impl Sync for Shared {} impl task::Schedule for Arc { fn release(&self, task: &Task) -> Option> { // Safety, this is always called from the thread that owns `LocalSet` unsafe { self.local_state.task_remove(task) } } fn schedule(&self, task: task::Notified) { Shared::schedule(self, task); } // localset does not currently support task hooks fn hooks(&self) -> TaskHarnessScheduleHooks { TaskHarnessScheduleHooks { task_terminate_callback: None, } } cfg_unstable! { fn unhandled_panic(&self) { use crate::runtime::UnhandledPanic; match self.unhandled_panic { UnhandledPanic::Ignore => { // Do nothing } UnhandledPanic::ShutdownRuntime => { // This hook is only called from within the runtime, so // `CURRENT` should match with `&self`, i.e. there is no // opportunity for a nested scheduler to be called. CURRENT.with(|LocalData { ctx, .. }| match ctx.get() { Some(cx) if Arc::ptr_eq(self, &cx.shared) => { cx.unhandled_panic.set(true); // Safety: this is always called from the thread that owns `LocalSet` unsafe { cx.shared.local_state.close_and_shutdown_all(); } } _ => unreachable!("runtime core not set in CURRENT thread-local"), }) } } } } } impl LocalState { unsafe fn task_pop_front(&self) -> Option>> { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.local_queue.with_mut(|ptr| (*ptr).pop_front()) } unsafe fn task_push_back(&self, task: task::Notified>) { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.local_queue.with_mut(|ptr| (*ptr).push_back(task)); } unsafe fn take_local_queue(&self) -> VecDeque>> { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.local_queue.with_mut(|ptr| std::mem::take(&mut (*ptr))) } unsafe fn task_remove(&self, task: &Task>) -> Option>> { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.owned.remove(task) } /// Returns true if the `LocalSet` does not have any spawned tasks unsafe fn owned_is_empty(&self) -> bool { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.owned.is_empty() } unsafe fn assert_owner( &self, task: task::Notified>, ) -> task::LocalNotified> { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.owned.assert_owner(task) } unsafe fn close_and_shutdown_all(&self) { // The caller ensures it is called from the same thread that owns // the LocalSet. self.assert_called_from_owner_thread(); self.owned.close_and_shutdown_all(); } #[track_caller] fn assert_called_from_owner_thread(&self) { // FreeBSD has some weirdness around thread-local destruction. // TODO: remove this hack when thread id is cleaned up #[cfg(not(any(target_os = "openbsd", target_os = "freebsd")))] debug_assert!( // if we couldn't get the thread ID because we're dropping the local // data, skip the assertion --- the `Drop` impl is not going to be // called from another thread, because `LocalSet` is `!Send` context::thread_id() .map(|id| id == self.owner) .unwrap_or(true), "`LocalSet`'s local run queue must not be accessed by another thread!" ); } } // This is `Send` because it is stored in `Shared`. It is up to the caller to // ensure they are on the same thread that owns the `LocalSet`. unsafe impl Send for LocalState {} #[cfg(all(test, not(loom)))] mod tests { use super::*; // Does a `LocalSet` running on a current-thread runtime...basically work? // // This duplicates a test in `tests/task_local_set.rs`, but because this is // a lib test, it will run under Miri, so this is necessary to catch stacked // borrows violations in the `LocalSet` implementation. #[test] fn local_current_thread_scheduler() { let f = async { LocalSet::new() .run_until(async { spawn_local(async {}).await.unwrap(); }) .await; }; crate::runtime::Builder::new_current_thread() .build() .expect("rt") .block_on(f) } // Tests that when a task on a `LocalSet` is woken by an io driver on the // same thread, the task is woken to the localset's local queue rather than // its remote queue. // // This test has to be defined in the `local.rs` file as a lib test, rather // than in `tests/`, because it makes assertions about the local set's // internal state. #[test] fn wakes_to_local_queue() { use super::*; use crate::sync::Notify; let rt = crate::runtime::Builder::new_current_thread() .build() .expect("rt"); rt.block_on(async { let local = LocalSet::new(); let notify = Arc::new(Notify::new()); let task = local.spawn_local({ let notify = notify.clone(); async move { notify.notified().await; } }); let mut run_until = Box::pin(local.run_until(async move { task.await.unwrap(); })); // poll the run until future once std::future::poll_fn(|cx| { let _ = run_until.as_mut().poll(cx); Poll::Ready(()) }) .await; notify.notify_one(); let task = unsafe { local.context.shared.local_state.task_pop_front() }; // TODO(eliza): it would be nice to be able to assert that this is // the local task. assert!( task.is_some(), "task should have been notified to the LocalSet's local queue" ); }) } } tokio-1.43.1/src/task/mod.rs000064400000000000000000000327051046102023000137240ustar 00000000000000//! Asynchronous green-threads. //! //! ## What are Tasks? //! //! A _task_ is a light weight, non-blocking unit of execution. A task is similar //! to an OS thread, but rather than being managed by the OS scheduler, they are //! managed by the [Tokio runtime][rt]. Another name for this general pattern is //! [green threads]. If you are familiar with [Go's goroutines], [Kotlin's //! coroutines], or [Erlang's processes], you can think of Tokio's tasks as //! something similar. //! //! Key points about tasks include: //! //! * Tasks are **light weight**. Because tasks are scheduled by the Tokio //! runtime rather than the operating system, creating new tasks or switching //! between tasks does not require a context switch and has fairly low //! overhead. Creating, running, and destroying large numbers of tasks is //! quite cheap, especially compared to OS threads. //! //! * Tasks are scheduled **cooperatively**. Most operating systems implement //! _preemptive multitasking_. This is a scheduling technique where the //! operating system allows each thread to run for a period of time, and then //! _preempts_ it, temporarily pausing that thread and switching to another. //! Tasks, on the other hand, implement _cooperative multitasking_. In //! cooperative multitasking, a task is allowed to run until it _yields_, //! indicating to the Tokio runtime's scheduler that it cannot currently //! continue executing. When a task yields, the Tokio runtime switches to //! executing the next task. //! //! * Tasks are **non-blocking**. Typically, when an OS thread performs I/O or //! must synchronize with another thread, it _blocks_, allowing the OS to //! schedule another thread. When a task cannot continue executing, it must //! yield instead, allowing the Tokio runtime to schedule another task. Tasks //! should generally not perform system calls or other operations that could //! block a thread, as this would prevent other tasks running on the same //! thread from executing as well. Instead, this module provides APIs for //! running blocking operations in an asynchronous context. //! //! [rt]: crate::runtime //! [green threads]: https://en.wikipedia.org/wiki/Green_threads //! [Go's goroutines]: https://tour.golang.org/concurrency/1 //! [Kotlin's coroutines]: https://kotlinlang.org/docs/reference/coroutines-overview.html //! [Erlang's processes]: http://erlang.org/doc/getting_started/conc_prog.html#processes //! //! ## Working with Tasks //! //! This module provides the following APIs for working with tasks: //! //! ### Spawning //! //! Perhaps the most important function in this module is [`task::spawn`]. This //! function can be thought of as an async equivalent to the standard library's //! [`thread::spawn`][`std::thread::spawn`]. It takes an `async` block or other //! [future], and creates a new task to run that work concurrently: //! //! ``` //! use tokio::task; //! //! # async fn doc() { //! task::spawn(async { //! // perform some work here... //! }); //! # } //! ``` //! //! Like [`std::thread::spawn`], `task::spawn` returns a [`JoinHandle`] struct. //! A `JoinHandle` is itself a future which may be used to await the output of //! the spawned task. For example: //! //! ``` //! use tokio::task; //! //! # #[tokio::main] async fn main() -> Result<(), Box> { //! let join = task::spawn(async { //! // ... //! "hello world!" //! }); //! //! // ... //! //! // Await the result of the spawned task. //! let result = join.await?; //! assert_eq!(result, "hello world!"); //! # Ok(()) //! # } //! ``` //! //! Again, like `std::thread`'s [`JoinHandle` type][thread_join], if the spawned //! task panics, awaiting its `JoinHandle` will return a [`JoinError`]. For //! example: //! //! ``` //! use tokio::task; //! //! # #[tokio::main] async fn main() { //! let join = task::spawn(async { //! panic!("something bad happened!") //! }); //! //! // The returned result indicates that the task failed. //! assert!(join.await.is_err()); //! # } //! ``` //! //! `spawn`, `JoinHandle`, and `JoinError` are present when the "rt" //! feature flag is enabled. //! //! [`task::spawn`]: crate::task::spawn() //! [future]: std::future::Future //! [`std::thread::spawn`]: std::thread::spawn //! [`JoinHandle`]: crate::task::JoinHandle //! [thread_join]: std::thread::JoinHandle //! [`JoinError`]: crate::task::JoinError //! //! #### Cancellation //! //! Spawned tasks may be cancelled using the [`JoinHandle::abort`] or //! [`AbortHandle::abort`] methods. When one of these methods are called, the //! task is signalled to shut down next time it yields at an `.await` point. If //! the task is already idle, then it will be shut down as soon as possible //! without running again before being shut down. Additionally, shutting down a //! Tokio runtime (e.g. by returning from `#[tokio::main]`) immediately cancels //! all tasks on it. //! //! When tasks are shut down, it will stop running at whichever `.await` it has //! yielded at. All local variables are destroyed by running their destructor. //! Once shutdown has completed, awaiting the [`JoinHandle`] will fail with a //! [cancelled error](crate::task::JoinError::is_cancelled). //! //! Note that aborting a task does not guarantee that it fails with a cancelled //! error, since it may complete normally first. For example, if the task does //! not yield to the runtime at any point between the call to `abort` and the //! end of the task, then the [`JoinHandle`] will instead report that the task //! exited normally. //! //! Be aware that tasks spawned using [`spawn_blocking`] cannot be aborted //! because they are not async. If you call `abort` on a `spawn_blocking` //! task, then this *will not have any effect*, and the task will continue //! running normally. The exception is if the task has not started running //! yet; in that case, calling `abort` may prevent the task from starting. //! //! Be aware that calls to [`JoinHandle::abort`] just schedule the task for //! cancellation, and will return before the cancellation has completed. To wait //! for cancellation to complete, wait for the task to finish by awaiting the //! [`JoinHandle`]. Similarly, the [`JoinHandle::is_finished`] method does not //! return `true` until the cancellation has finished. //! //! Calling [`JoinHandle::abort`] multiple times has the same effect as calling //! it once. //! //! Tokio also provides an [`AbortHandle`], which is like the [`JoinHandle`], //! except that it does not provide a mechanism to wait for the task to finish. //! Each task can only have one [`JoinHandle`], but it can have more than one //! [`AbortHandle`]. //! //! [`JoinHandle::abort`]: crate::task::JoinHandle::abort //! [`AbortHandle::abort`]: crate::task::AbortHandle::abort //! [`AbortHandle`]: crate::task::AbortHandle //! [`JoinHandle::is_finished`]: crate::task::JoinHandle::is_finished //! //! ### Blocking and Yielding //! //! As we discussed above, code running in asynchronous tasks should not perform //! operations that can block. A blocking operation performed in a task running //! on a thread that is also running other tasks would block the entire thread, //! preventing other tasks from running. //! //! Instead, Tokio provides two APIs for running blocking operations in an //! asynchronous context: [`task::spawn_blocking`] and [`task::block_in_place`]. //! //! Be aware that if you call a non-async method from async code, that non-async //! method is still inside the asynchronous context, so you should also avoid //! blocking operations there. This includes destructors of objects destroyed in //! async code. //! //! #### `spawn_blocking` //! //! The `task::spawn_blocking` function is similar to the `task::spawn` function //! discussed in the previous section, but rather than spawning an //! _non-blocking_ future on the Tokio runtime, it instead spawns a //! _blocking_ function on a dedicated thread pool for blocking tasks. For //! example: //! //! ``` //! use tokio::task; //! //! # async fn docs() { //! task::spawn_blocking(|| { //! // do some compute-heavy work or call synchronous code //! }); //! # } //! ``` //! //! Just like `task::spawn`, `task::spawn_blocking` returns a `JoinHandle` //! which we can use to await the result of the blocking operation: //! //! ```rust //! # use tokio::task; //! # async fn docs() -> Result<(), Box>{ //! let join = task::spawn_blocking(|| { //! // do some compute-heavy work or call synchronous code //! "blocking completed" //! }); //! //! let result = join.await?; //! assert_eq!(result, "blocking completed"); //! # Ok(()) //! # } //! ``` //! //! #### `block_in_place` //! //! When using the [multi-threaded runtime][rt-multi-thread], the [`task::block_in_place`] //! function is also available. Like `task::spawn_blocking`, this function //! allows running a blocking operation from an asynchronous context. Unlike //! `spawn_blocking`, however, `block_in_place` works by transitioning the //! _current_ worker thread to a blocking thread, moving other tasks running on //! that thread to another worker thread. This can improve performance by avoiding //! context switches. //! //! For example: //! //! ``` //! use tokio::task; //! //! # async fn docs() { //! let result = task::block_in_place(|| { //! // do some compute-heavy work or call synchronous code //! "blocking completed" //! }); //! //! assert_eq!(result, "blocking completed"); //! # } //! ``` //! //! #### `yield_now` //! //! In addition, this module provides a [`task::yield_now`] async function //! that is analogous to the standard library's [`thread::yield_now`]. Calling //! and `await`ing this function will cause the current task to yield to the //! Tokio runtime's scheduler, allowing other tasks to be //! scheduled. Eventually, the yielding task will be polled again, allowing it //! to execute. For example: //! //! ```rust //! use tokio::task; //! //! # #[tokio::main] async fn main() { //! async { //! task::spawn(async { //! // ... //! println!("spawned task done!") //! }); //! //! // Yield, allowing the newly-spawned task to execute first. //! task::yield_now().await; //! println!("main task done!"); //! } //! # .await; //! # } //! ``` //! //! ### Cooperative scheduling //! //! A single call to [`poll`] on a top-level task may potentially do a lot of //! work before it returns `Poll::Pending`. If a task runs for a long period of //! time without yielding back to the executor, it can starve other tasks //! waiting on that executor to execute them, or drive underlying resources. //! Since Rust does not have a runtime, it is difficult to forcibly preempt a //! long-running task. Instead, this module provides an opt-in mechanism for //! futures to collaborate with the executor to avoid starvation. //! //! Consider a future like this one: //! //! ``` //! # use tokio_stream::{Stream, StreamExt}; //! async fn drop_all(mut input: I) { //! while let Some(_) = input.next().await {} //! } //! ``` //! //! It may look harmless, but consider what happens under heavy load if the //! input stream is _always_ ready. If we spawn `drop_all`, the task will never //! yield, and will starve other tasks and resources on the same executor. //! //! To account for this, Tokio has explicit yield points in a number of library //! functions, which force tasks to return to the executor periodically. //! //! //! #### unconstrained //! //! If necessary, [`task::unconstrained`] lets you opt a future out of Tokio's cooperative //! scheduling. When a future is wrapped with `unconstrained`, it will never be forced to yield to //! Tokio. For example: //! //! ``` //! # #[tokio::main] //! # async fn main() { //! use tokio::{task, sync::mpsc}; //! //! let fut = async { //! let (tx, mut rx) = mpsc::unbounded_channel(); //! //! for i in 0..1000 { //! let _ = tx.send(()); //! // This will always be ready. If coop was in effect, this code would be forced to yield //! // periodically. However, if left unconstrained, then this code will never yield. //! rx.recv().await; //! } //! }; //! //! task::unconstrained(fut).await; //! # } //! ``` //! //! [`task::spawn_blocking`]: crate::task::spawn_blocking //! [`task::block_in_place`]: crate::task::block_in_place //! [rt-multi-thread]: ../runtime/index.html#threaded-scheduler //! [`task::yield_now`]: crate::task::yield_now() //! [`thread::yield_now`]: std::thread::yield_now //! [`task::unconstrained`]: crate::task::unconstrained() //! [`poll`]: method@std::future::Future::poll cfg_rt! { pub use crate::runtime::task::{JoinError, JoinHandle}; mod blocking; pub use blocking::spawn_blocking; mod spawn; pub use spawn::spawn; cfg_rt_multi_thread! { pub use blocking::block_in_place; } mod yield_now; pub use yield_now::yield_now; mod consume_budget; pub use consume_budget::consume_budget; mod local; pub use local::{spawn_local, LocalSet, LocalEnterGuard}; mod task_local; pub use task_local::LocalKey; mod unconstrained; pub use unconstrained::{unconstrained, Unconstrained}; #[doc(inline)] pub use join_set::JoinSet; pub use crate::runtime::task::AbortHandle; // Uses #[cfg(...)] instead of macro since the macro adds docsrs annotations. #[cfg(not(tokio_unstable))] mod join_set; #[cfg(tokio_unstable)] pub mod join_set; pub use crate::runtime::task::{Id, id, try_id}; cfg_trace! { mod builder; pub use builder::Builder; } /// Task-related futures. pub mod futures { pub use super::task_local::TaskLocalFuture; } } tokio-1.43.1/src/task/spawn.rs000064400000000000000000000145541046102023000142770ustar 00000000000000use crate::runtime::BOX_FUTURE_THRESHOLD; use crate::task::JoinHandle; use crate::util::trace::SpawnMeta; use std::future::Future; cfg_rt! { /// Spawns a new asynchronous task, returning a /// [`JoinHandle`](JoinHandle) for it. /// /// The provided future will start running in the background immediately /// when `spawn` is called, even if you don't await the returned /// `JoinHandle`. /// /// Spawning a task enables the task to execute concurrently to other tasks. The /// spawned task may execute on the current thread, or it may be sent to a /// different thread to be executed. The specifics depend on the current /// [`Runtime`](crate::runtime::Runtime) configuration. /// /// It is guaranteed that spawn will not synchronously poll the task being spawned. /// This means that calling spawn while holding a lock does not pose a risk of /// deadlocking with the spawned task. /// /// There is no guarantee that a spawned task will execute to completion. /// When a runtime is shutdown, all outstanding tasks are dropped, /// regardless of the lifecycle of that task. /// /// This function must be called from the context of a Tokio runtime. Tasks running on /// the Tokio runtime are always inside its context, but you can also enter the context /// using the [`Runtime::enter`](crate::runtime::Runtime::enter()) method. /// /// # Examples /// /// In this example, a server is started and `spawn` is used to start a new task /// that processes each received connection. /// /// ```no_run /// use tokio::net::{TcpListener, TcpStream}; /// /// use std::io; /// /// async fn process(socket: TcpStream) { /// // ... /// # drop(socket); /// } /// /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let listener = TcpListener::bind("127.0.0.1:8080").await?; /// /// loop { /// let (socket, _) = listener.accept().await?; /// /// tokio::spawn(async move { /// // Process each socket concurrently. /// process(socket).await /// }); /// } /// } /// ``` /// /// To run multiple tasks in parallel and receive their results, join /// handles can be stored in a vector. /// ``` /// # #[tokio::main(flavor = "current_thread")] async fn main() { /// async fn my_background_op(id: i32) -> String { /// let s = format!("Starting background task {}.", id); /// println!("{}", s); /// s /// } /// /// let ops = vec![1, 2, 3]; /// let mut tasks = Vec::with_capacity(ops.len()); /// for op in ops { /// // This call will make them start running in the background /// // immediately. /// tasks.push(tokio::spawn(my_background_op(op))); /// } /// /// let mut outputs = Vec::with_capacity(tasks.len()); /// for task in tasks { /// outputs.push(task.await.unwrap()); /// } /// println!("{:?}", outputs); /// # } /// ``` /// This example pushes the tasks to `outputs` in the order they were /// started in. If you do not care about the ordering of the outputs, then /// you can also use a [`JoinSet`]. /// /// [`JoinSet`]: struct@crate::task::JoinSet /// /// # Panics /// /// Panics if called from **outside** of the Tokio runtime. /// /// # Using `!Send` values from a task /// /// The task supplied to `spawn` must implement `Send`. However, it is /// possible to **use** `!Send` values from the task as long as they only /// exist between calls to `.await`. /// /// For example, this will work: /// /// ``` /// use tokio::task; /// /// use std::rc::Rc; /// /// fn use_rc(rc: Rc<()>) { /// // Do stuff w/ rc /// # drop(rc); /// } /// /// #[tokio::main] /// async fn main() { /// tokio::spawn(async { /// // Force the `Rc` to stay in a scope with no `.await` /// { /// let rc = Rc::new(()); /// use_rc(rc.clone()); /// } /// /// task::yield_now().await; /// }).await.unwrap(); /// } /// ``` /// /// This will **not** work: /// /// ```compile_fail /// use tokio::task; /// /// use std::rc::Rc; /// /// fn use_rc(rc: Rc<()>) { /// // Do stuff w/ rc /// # drop(rc); /// } /// /// #[tokio::main] /// async fn main() { /// tokio::spawn(async { /// let rc = Rc::new(()); /// /// task::yield_now().await; /// /// use_rc(rc.clone()); /// }).await.unwrap(); /// } /// ``` /// /// Holding on to a `!Send` value across calls to `.await` will result in /// an unfriendly compile error message similar to: /// /// ```text /// `[... some type ...]` cannot be sent between threads safely /// ``` /// /// or: /// /// ```text /// error[E0391]: cycle detected when processing `main` /// ``` #[track_caller] pub fn spawn(future: F) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, { let fut_size = std::mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { spawn_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) } else { spawn_inner(future, SpawnMeta::new_unnamed(fut_size)) } } #[track_caller] pub(super) fn spawn_inner(future: T, meta: SpawnMeta<'_>) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { use crate::runtime::{context, task}; #[cfg(all( tokio_unstable, tokio_taskdump, feature = "rt", target_os = "linux", any( target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64" ) ))] let future = task::trace::Trace::root(future); let id = task::Id::next(); let task = crate::util::trace::task(future, "task", meta, id.as_u64()); match context::with_current(|handle| handle.spawn(task, id)) { Ok(join_handle) => join_handle, Err(e) => panic!("{}", e), } } } tokio-1.43.1/src/task/task_local.rs000064400000000000000000000327541046102023000152650ustar 00000000000000use pin_project_lite::pin_project; use std::cell::RefCell; use std::error::Error; use std::future::Future; use std::marker::PhantomPinned; use std::pin::Pin; use std::task::{Context, Poll}; use std::{fmt, mem, thread}; /// Declares a new task-local key of type [`tokio::task::LocalKey`]. /// /// # Syntax /// /// The macro wraps any number of static declarations and makes them local to the current task. /// Publicity and attributes for each static is preserved. For example: /// /// # Examples /// /// ``` /// # use tokio::task_local; /// task_local! { /// pub static ONE: u32; /// /// #[allow(unused)] /// static TWO: f32; /// } /// # fn main() {} /// ``` /// /// See [`LocalKey` documentation][`tokio::task::LocalKey`] for more /// information. /// /// [`tokio::task::LocalKey`]: struct@crate::task::LocalKey #[macro_export] #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] macro_rules! task_local { // empty (base case for the recursion) () => {}; ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty; $($rest:tt)*) => { $crate::__task_local_inner!($(#[$attr])* $vis $name, $t); $crate::task_local!($($rest)*); }; ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty) => { $crate::__task_local_inner!($(#[$attr])* $vis $name, $t); } } #[doc(hidden)] #[macro_export] macro_rules! __task_local_inner { ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => { $(#[$attr])* $vis static $name: $crate::task::LocalKey<$t> = { std::thread_local! { static __KEY: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; } $crate::task::LocalKey { inner: __KEY } }; }; } /// A key for task-local data. /// /// This type is generated by the [`task_local!`] macro. /// /// Unlike [`std::thread::LocalKey`], `tokio::task::LocalKey` will /// _not_ lazily initialize the value on first access. Instead, the /// value is first initialized when the future containing /// the task-local is first polled by a futures executor, like Tokio. /// /// # Examples /// /// ``` /// # async fn dox() { /// tokio::task_local! { /// static NUMBER: u32; /// } /// /// NUMBER.scope(1, async move { /// assert_eq!(NUMBER.get(), 1); /// }).await; /// /// NUMBER.scope(2, async move { /// assert_eq!(NUMBER.get(), 2); /// /// NUMBER.scope(3, async move { /// assert_eq!(NUMBER.get(), 3); /// }).await; /// }).await; /// # } /// ``` /// /// [`std::thread::LocalKey`]: struct@std::thread::LocalKey /// [`task_local!`]: ../macro.task_local.html #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub struct LocalKey { #[doc(hidden)] pub inner: thread::LocalKey>>, } impl LocalKey { /// Sets a value `T` as the task-local value for the future `F`. /// /// On completion of `scope`, the task-local will be dropped. /// /// ### Panics /// /// If you poll the returned future inside a call to [`with`] or /// [`try_with`] on the same `LocalKey`, then the call to `poll` will panic. /// /// ### Examples /// /// ``` /// # async fn dox() { /// tokio::task_local! { /// static NUMBER: u32; /// } /// /// NUMBER.scope(1, async move { /// println!("task local value: {}", NUMBER.get()); /// }).await; /// # } /// ``` /// /// [`with`]: fn@Self::with /// [`try_with`]: fn@Self::try_with pub fn scope(&'static self, value: T, f: F) -> TaskLocalFuture where F: Future, { TaskLocalFuture { local: self, slot: Some(value), future: Some(f), _pinned: PhantomPinned, } } /// Sets a value `T` as the task-local value for the closure `F`. /// /// On completion of `sync_scope`, the task-local will be dropped. /// /// ### Panics /// /// This method panics if called inside a call to [`with`] or [`try_with`] /// on the same `LocalKey`. /// /// ### Examples /// /// ``` /// # async fn dox() { /// tokio::task_local! { /// static NUMBER: u32; /// } /// /// NUMBER.sync_scope(1, || { /// println!("task local value: {}", NUMBER.get()); /// }); /// # } /// ``` /// /// [`with`]: fn@Self::with /// [`try_with`]: fn@Self::try_with #[track_caller] pub fn sync_scope(&'static self, value: T, f: F) -> R where F: FnOnce() -> R, { let mut value = Some(value); match self.scope_inner(&mut value, f) { Ok(res) => res, Err(err) => err.panic(), } } fn scope_inner(&'static self, slot: &mut Option, f: F) -> Result where F: FnOnce() -> R, { struct Guard<'a, T: 'static> { local: &'static LocalKey, slot: &'a mut Option, } impl<'a, T: 'static> Drop for Guard<'a, T> { fn drop(&mut self) { // This should not panic. // // We know that the RefCell was not borrowed before the call to // `scope_inner`, so the only way for this to panic is if the // closure has created but not destroyed a RefCell guard. // However, we never give user-code access to the guards, so // there's no way for user-code to forget to destroy a guard. // // The call to `with` also should not panic, since the // thread-local wasn't destroyed when we first called // `scope_inner`, and it shouldn't have gotten destroyed since // then. self.local.inner.with(|inner| { let mut ref_mut = inner.borrow_mut(); mem::swap(self.slot, &mut *ref_mut); }); } } self.inner.try_with(|inner| { inner .try_borrow_mut() .map(|mut ref_mut| mem::swap(slot, &mut *ref_mut)) })??; let guard = Guard { local: self, slot }; let res = f(); drop(guard); Ok(res) } /// Accesses the current task-local and runs the provided closure. /// /// # Panics /// /// This function will panic if the task local doesn't have a value set. #[track_caller] pub fn with(&'static self, f: F) -> R where F: FnOnce(&T) -> R, { match self.try_with(f) { Ok(res) => res, Err(_) => panic!("cannot access a task-local storage value without setting it first"), } } /// Accesses the current task-local and runs the provided closure. /// /// If the task-local with the associated key is not present, this /// method will return an `AccessError`. For a panicking variant, /// see `with`. pub fn try_with(&'static self, f: F) -> Result where F: FnOnce(&T) -> R, { // If called after the thread-local storing the task-local is destroyed, // then we are outside of a closure where the task-local is set. // // Therefore, it is correct to return an AccessError if `try_with` // returns an error. let try_with_res = self.inner.try_with(|v| { // This call to `borrow` cannot panic because no user-defined code // runs while a `borrow_mut` call is active. v.borrow().as_ref().map(f) }); match try_with_res { Ok(Some(res)) => Ok(res), Ok(None) | Err(_) => Err(AccessError { _private: () }), } } } impl LocalKey { /// Returns a copy of the task-local value /// if the task-local value implements `Clone`. /// /// # Panics /// /// This function will panic if the task local doesn't have a value set. #[track_caller] pub fn get(&'static self) -> T { self.with(|v| v.clone()) } } impl fmt::Debug for LocalKey { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.pad("LocalKey { .. }") } } pin_project! { /// A future that sets a value `T` of a task local for the future `F` during /// its execution. /// /// The value of the task-local must be `'static` and will be dropped on the /// completion of the future. /// /// Created by the function [`LocalKey::scope`](self::LocalKey::scope). /// /// ### Examples /// /// ``` /// # async fn dox() { /// tokio::task_local! { /// static NUMBER: u32; /// } /// /// NUMBER.scope(1, async move { /// println!("task local value: {}", NUMBER.get()); /// }).await; /// # } /// ``` pub struct TaskLocalFuture where T: 'static, { local: &'static LocalKey, slot: Option, #[pin] future: Option, #[pin] _pinned: PhantomPinned, } impl PinnedDrop for TaskLocalFuture { fn drop(this: Pin<&mut Self>) { let this = this.project(); if mem::needs_drop::() && this.future.is_some() { // Drop the future while the task-local is set, if possible. Otherwise // the future is dropped normally when the `Option` field drops. let mut future = this.future; let _ = this.local.scope_inner(this.slot, || { future.set(None); }); } } } } impl TaskLocalFuture where T: 'static, { /// Returns the value stored in the task local by this `TaskLocalFuture`. /// /// The function returns: /// /// * `Some(T)` if the task local value exists. /// * `None` if the task local value has already been taken. /// /// Note that this function attempts to take the task local value even if /// the future has not yet completed. In that case, the value will no longer /// be available via the task local after the call to `take_value`. /// /// # Examples /// /// ``` /// # async fn dox() { /// tokio::task_local! { /// static KEY: u32; /// } /// /// let fut = KEY.scope(42, async { /// // Do some async work /// }); /// /// let mut pinned = Box::pin(fut); /// /// // Complete the TaskLocalFuture /// let _ = pinned.as_mut().await; /// /// // And here, we can take task local value /// let value = pinned.as_mut().take_value(); /// /// assert_eq!(value, Some(42)); /// # } /// ``` pub fn take_value(self: Pin<&mut Self>) -> Option { let this = self.project(); this.slot.take() } } impl Future for TaskLocalFuture { type Output = F::Output; #[track_caller] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); let mut future_opt = this.future; let res = this .local .scope_inner(this.slot, || match future_opt.as_mut().as_pin_mut() { Some(fut) => { let res = fut.poll(cx); if res.is_ready() { future_opt.set(None); } Some(res) } None => None, }); match res { Ok(Some(res)) => res, Ok(None) => panic!("`TaskLocalFuture` polled after completion"), Err(err) => err.panic(), } } } impl fmt::Debug for TaskLocalFuture where T: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { /// Format the Option without Some. struct TransparentOption<'a, T> { value: &'a Option, } impl<'a, T: fmt::Debug> fmt::Debug for TransparentOption<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self.value.as_ref() { Some(value) => value.fmt(f), // Hitting the None branch should not be possible. None => f.pad(""), } } } f.debug_struct("TaskLocalFuture") .field("value", &TransparentOption { value: &self.slot }) .finish() } } /// An error returned by [`LocalKey::try_with`](method@LocalKey::try_with). #[derive(Clone, Copy, Eq, PartialEq)] pub struct AccessError { _private: (), } impl fmt::Debug for AccessError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("AccessError").finish() } } impl fmt::Display for AccessError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Display::fmt("task-local value not set", f) } } impl Error for AccessError {} enum ScopeInnerErr { BorrowError, AccessError, } impl ScopeInnerErr { #[track_caller] fn panic(&self) -> ! { match self { Self::BorrowError => panic!("cannot enter a task-local scope while the task-local storage is borrowed"), Self::AccessError => panic!("cannot enter a task-local scope during or after destruction of the underlying thread-local"), } } } impl From for ScopeInnerErr { fn from(_: std::cell::BorrowMutError) -> Self { Self::BorrowError } } impl From for ScopeInnerErr { fn from(_: std::thread::AccessError) -> Self { Self::AccessError } } tokio-1.43.1/src/task/unconstrained.rs000064400000000000000000000024641046102023000160200ustar 00000000000000use pin_project_lite::pin_project; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; pin_project! { /// Future for the [`unconstrained`](unconstrained) method. #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] #[must_use = "Unconstrained does nothing unless polled"] pub struct Unconstrained { #[pin] inner: F, } } impl Future for Unconstrained where F: Future, { type Output = ::Output; cfg_coop! { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let inner = self.project().inner; crate::runtime::coop::with_unconstrained(|| inner.poll(cx)) } } cfg_not_coop! { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let inner = self.project().inner; inner.poll(cx) } } } /// Turn off cooperative scheduling for a future. The future will never be forced to yield by /// Tokio. Using this exposes your service to starvation if the unconstrained future never yields /// otherwise. /// /// See also the usage example in the [task module](index.html#unconstrained). #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub fn unconstrained(inner: F) -> Unconstrained { Unconstrained { inner } } tokio-1.43.1/src/task/yield_now.rs000064400000000000000000000044611046102023000151340ustar 00000000000000use crate::runtime::context; use std::future::Future; use std::pin::Pin; use std::task::{ready, Context, Poll}; /// Yields execution back to the Tokio runtime. /// /// A task yields by awaiting on `yield_now()`, and may resume when that future /// completes (with no output.) The current task will be re-added as a pending /// task at the _back_ of the pending queue. Any other pending tasks will be /// scheduled. No other waking is required for the task to continue. /// /// See also the usage example in the [task module](index.html#yield_now). /// /// ## Non-guarantees /// /// This function may not yield all the way up to the executor if there are any /// special combinators above it in the call stack. For example, if a /// [`tokio::select!`] has another branch complete during the same poll as the /// `yield_now()`, then the yield is not propagated all the way up to the /// runtime. /// /// It is generally not guaranteed that the runtime behaves like you expect it /// to when deciding which task to schedule next after a call to `yield_now()`. /// In particular, the runtime may choose to poll the task that just ran /// `yield_now()` again immediately without polling any other tasks first. For /// example, the runtime will not drive the IO driver between every poll of a /// task, and this could result in the runtime polling the current task again /// immediately even if there is another task that could make progress if that /// other task is waiting for a notification from the IO driver. /// /// In general, changes to the order in which the runtime polls tasks is not /// considered a breaking change, and your program should be correct no matter /// which order the runtime polls your tasks in. /// /// [`tokio::select!`]: macro@crate::select #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] pub async fn yield_now() { /// Yield implementation struct YieldNow { yielded: bool, } impl Future for YieldNow { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { ready!(crate::trace::trace_leaf(cx)); if self.yielded { return Poll::Ready(()); } self.yielded = true; context::defer(cx.waker()); Poll::Pending } } YieldNow { yielded: false }.await; } tokio-1.43.1/src/time/clock.rs000064400000000000000000000240121046102023000142240ustar 00000000000000#![cfg_attr(not(feature = "rt"), allow(dead_code))] //! Source of time abstraction. //! //! By default, `std::time::Instant::now()` is used. However, when the //! `test-util` feature flag is enabled, the values returned for `now()` are //! configurable. cfg_not_test_util! { use crate::time::{Instant}; #[derive(Debug, Clone)] pub(crate) struct Clock {} pub(crate) fn now() -> Instant { Instant::from_std(std::time::Instant::now()) } impl Clock { pub(crate) fn new(_enable_pausing: bool, _start_paused: bool) -> Clock { Clock {} } pub(crate) fn now(&self) -> Instant { now() } } } cfg_test_util! { use crate::time::{Duration, Instant}; use crate::loom::sync::Mutex; use crate::loom::sync::atomic::Ordering; use std::sync::atomic::AtomicBool as StdAtomicBool; cfg_rt! { #[track_caller] fn with_clock(f: impl FnOnce(Option<&Clock>) -> Result) -> R { use crate::runtime::Handle; let res = match Handle::try_current() { Ok(handle) => f(Some(handle.inner.driver().clock())), Err(ref e) if e.is_missing_context() => f(None), Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), }; match res { Ok(ret) => ret, Err(msg) => panic!("{}", msg), } } } cfg_not_rt! { #[track_caller] fn with_clock(f: impl FnOnce(Option<&Clock>) -> Result) -> R { match f(None) { Ok(ret) => ret, Err(msg) => panic!("{}", msg), } } } /// A handle to a source of time. #[derive(Debug)] pub(crate) struct Clock { inner: Mutex, } // Used to track if the clock was ever paused. This is an optimization to // avoid touching the mutex if `test-util` was accidentally enabled in // release mode. // // A static is used so we can avoid accessing the thread-local as well. The // `std` AtomicBool is used directly because loom does not support static // atomics. static DID_PAUSE_CLOCK: StdAtomicBool = StdAtomicBool::new(false); #[derive(Debug)] struct Inner { /// True if the ability to pause time is enabled. enable_pausing: bool, /// Instant to use as the clock's base instant. base: std::time::Instant, /// Instant at which the clock was last unfrozen. unfrozen: Option, /// Number of `inhibit_auto_advance` calls still in effect. auto_advance_inhibit_count: usize, } /// Pauses time. /// /// The current value of `Instant::now()` is saved and all subsequent calls /// to `Instant::now()` will return the saved value. The saved value can be /// changed by [`advance`] or by the time auto-advancing once the runtime /// has no work to do. This only affects the `Instant` type in Tokio, and /// the `Instant` in std continues to work as normal. /// /// Pausing time requires the `current_thread` Tokio runtime. This is the /// default runtime used by `#[tokio::test]`. The runtime can be initialized /// with time in a paused state using the `Builder::start_paused` method. /// /// For cases where time is immediately paused, it is better to pause /// the time using the `main` or `test` macro: /// ``` /// #[tokio::main(flavor = "current_thread", start_paused = true)] /// async fn main() { /// println!("Hello world"); /// } /// ``` /// /// # Panics /// /// Panics if time is already frozen or if called from outside of a /// `current_thread` Tokio runtime. /// /// # Auto-advance /// /// If time is paused and the runtime has no work to do, the clock is /// auto-advanced to the next pending timer. This means that [`Sleep`] or /// other timer-backed primitives can cause the runtime to advance the /// current time when awaited. /// /// [`Sleep`]: crate::time::Sleep /// [`advance`]: crate::time::advance #[track_caller] pub fn pause() { with_clock(|maybe_clock| { match maybe_clock { Some(clock) => clock.pause(), None => Err("time cannot be frozen from outside the Tokio runtime"), } }); } /// Resumes time. /// /// Clears the saved `Instant::now()` value. Subsequent calls to /// `Instant::now()` will return the value returned by the system call. /// /// # Panics /// /// Panics if time is not frozen or if called from outside of the Tokio /// runtime. #[track_caller] pub fn resume() { with_clock(|maybe_clock| { let clock = match maybe_clock { Some(clock) => clock, None => return Err("time cannot be frozen from outside the Tokio runtime"), }; let mut inner = clock.inner.lock(); if inner.unfrozen.is_some() { return Err("time is not frozen"); } inner.unfrozen = Some(std::time::Instant::now()); Ok(()) }); } /// Advances time. /// /// Increments the saved `Instant::now()` value by `duration`. Subsequent /// calls to `Instant::now()` will return the result of the increment. /// /// This function will make the current time jump forward by the given /// duration in one jump. This means that all `sleep` calls with a deadline /// before the new time will immediately complete "at the same time", and /// the runtime is free to poll them in any order. Additionally, this /// method will not wait for the `sleep` calls it advanced past to complete. /// If you want to do that, you should instead call [`sleep`] and rely on /// the runtime's auto-advance feature. /// /// Note that calls to `sleep` are not guaranteed to complete the first time /// they are polled after a call to `advance`. For example, this can happen /// if the runtime has not yet touched the timer driver after the call to /// `advance`. However if they don't, the runtime will poll the task again /// shortly. /// /// # Panics /// /// Panics if time is not frozen or if called from outside of the Tokio /// runtime. /// /// # Auto-advance /// /// If the time is paused and there is no work to do, the runtime advances /// time to the next timer. See [`pause`](pause#auto-advance) for more /// details. /// /// [`sleep`]: fn@crate::time::sleep pub async fn advance(duration: Duration) { with_clock(|maybe_clock| { let clock = match maybe_clock { Some(clock) => clock, None => return Err("time cannot be frozen from outside the Tokio runtime"), }; clock.advance(duration) }); crate::task::yield_now().await; } /// Returns the current instant, factoring in frozen time. pub(crate) fn now() -> Instant { if !DID_PAUSE_CLOCK.load(Ordering::Acquire) { return Instant::from_std(std::time::Instant::now()); } with_clock(|maybe_clock| { Ok(if let Some(clock) = maybe_clock { clock.now() } else { Instant::from_std(std::time::Instant::now()) }) }) } impl Clock { /// Returns a new `Clock` instance that uses the current execution context's /// source of time. pub(crate) fn new(enable_pausing: bool, start_paused: bool) -> Clock { let now = std::time::Instant::now(); let clock = Clock { inner: Mutex::new(Inner { enable_pausing, base: now, unfrozen: Some(now), auto_advance_inhibit_count: 0, }), }; if start_paused { if let Err(msg) = clock.pause() { panic!("{}", msg); } } clock } pub(crate) fn pause(&self) -> Result<(), &'static str> { let mut inner = self.inner.lock(); if !inner.enable_pausing { drop(inner); // avoid poisoning the lock return Err("`time::pause()` requires the `current_thread` Tokio runtime. \ This is the default Runtime used by `#[tokio::test]."); } // Track that we paused the clock DID_PAUSE_CLOCK.store(true, Ordering::Release); let elapsed = match inner.unfrozen.as_ref() { Some(v) => v.elapsed(), None => return Err("time is already frozen") }; inner.base += elapsed; inner.unfrozen = None; Ok(()) } /// Temporarily stop auto-advancing the clock (see `tokio::time::pause`). pub(crate) fn inhibit_auto_advance(&self) { let mut inner = self.inner.lock(); inner.auto_advance_inhibit_count += 1; } pub(crate) fn allow_auto_advance(&self) { let mut inner = self.inner.lock(); inner.auto_advance_inhibit_count -= 1; } pub(crate) fn can_auto_advance(&self) -> bool { let inner = self.inner.lock(); inner.unfrozen.is_none() && inner.auto_advance_inhibit_count == 0 } pub(crate) fn advance(&self, duration: Duration) -> Result<(), &'static str> { let mut inner = self.inner.lock(); if inner.unfrozen.is_some() { return Err("time is not frozen"); } inner.base += duration; Ok(()) } pub(crate) fn now(&self) -> Instant { let inner = self.inner.lock(); let mut ret = inner.base; if let Some(unfrozen) = inner.unfrozen { ret += unfrozen.elapsed(); } Instant::from_std(ret) } } } tokio-1.43.1/src/time/error.rs000064400000000000000000000067361046102023000142770ustar 00000000000000//! Time error types. use std::error; use std::fmt; /// Errors encountered by the timer implementation. /// /// Currently, there are two different errors that can occur: /// /// * `shutdown` occurs when a timer operation is attempted, but the timer /// instance has been dropped. In this case, the operation will never be able /// to complete and the `shutdown` error is returned. This is a permanent /// error, i.e., once this error is observed, timer operations will never /// succeed in the future. /// /// * `at_capacity` occurs when a timer operation is attempted, but the timer /// instance is currently handling its maximum number of outstanding sleep instances. /// In this case, the operation is not able to be performed at the current /// moment, and `at_capacity` is returned. This is a transient error, i.e., at /// some point in the future, if the operation is attempted again, it might /// succeed. Callers that observe this error should attempt to [shed load]. One /// way to do this would be dropping the future that issued the timer operation. /// /// [shed load]: https://en.wikipedia.org/wiki/Load_Shedding #[derive(Debug, Copy, Clone)] pub struct Error(Kind); #[derive(Debug, Clone, Copy, Eq, PartialEq)] #[repr(u8)] pub(crate) enum Kind { Shutdown = 1, AtCapacity = 2, Invalid = 3, } impl From for Error { fn from(k: Kind) -> Self { Error(k) } } /// Errors returned by `Timeout`. /// /// This error is returned when a timeout expires before the function was able /// to finish. #[derive(Debug, PartialEq, Eq)] pub struct Elapsed(()); #[derive(Debug)] pub(crate) enum InsertError { Elapsed, } // ===== impl Error ===== impl Error { /// Creates an error representing a shutdown timer. pub fn shutdown() -> Error { Error(Kind::Shutdown) } /// Returns `true` if the error was caused by the timer being shutdown. pub fn is_shutdown(&self) -> bool { matches!(self.0, Kind::Shutdown) } /// Creates an error representing a timer at capacity. pub fn at_capacity() -> Error { Error(Kind::AtCapacity) } /// Returns `true` if the error was caused by the timer being at capacity. pub fn is_at_capacity(&self) -> bool { matches!(self.0, Kind::AtCapacity) } /// Creates an error representing a misconfigured timer. pub fn invalid() -> Error { Error(Kind::Invalid) } /// Returns `true` if the error was caused by the timer being misconfigured. pub fn is_invalid(&self) -> bool { matches!(self.0, Kind::Invalid) } } impl error::Error for Error {} impl fmt::Display for Error { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { let descr = match self.0 { Kind::Shutdown => { "the timer is shutdown, must be called from the context of Tokio runtime" } Kind::AtCapacity => "timer is at capacity and cannot create a new entry", Kind::Invalid => "timer duration exceeds maximum duration", }; write!(fmt, "{descr}") } } // ===== impl Elapsed ===== impl Elapsed { pub(crate) fn new() -> Self { Elapsed(()) } } impl fmt::Display for Elapsed { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { "deadline has elapsed".fmt(fmt) } } impl std::error::Error for Elapsed {} impl From for std::io::Error { fn from(_err: Elapsed) -> std::io::Error { std::io::ErrorKind::TimedOut.into() } } tokio-1.43.1/src/time/instant.rs000064400000000000000000000152731046102023000146220ustar 00000000000000#![allow(clippy::trivially_copy_pass_by_ref)] use std::fmt; use std::ops; use std::time::Duration; /// A measurement of a monotonically nondecreasing clock. /// Opaque and useful only with `Duration`. /// /// Instants are always guaranteed to be no less than any previously measured /// instant when created, and are often useful for tasks such as measuring /// benchmarks or timing how long an operation takes. /// /// Note, however, that instants are not guaranteed to be **steady**. In other /// words, each tick of the underlying clock may not be the same length (e.g. /// some seconds may be longer than others). An instant may jump forwards or /// experience time dilation (slow down or speed up), but it will never go /// backwards. /// /// Instants are opaque types that can only be compared to one another. There is /// no method to get "the number of seconds" from an instant. Instead, it only /// allows measuring the duration between two instants (or comparing two /// instants). /// /// The size of an `Instant` struct may vary depending on the target operating /// system. /// /// # Note /// /// This type wraps the inner `std` variant and is used to align the Tokio /// clock for uses of `now()`. This can be useful for testing where you can /// take advantage of `time::pause()` and `time::advance()`. #[derive(Clone, Copy, Eq, PartialEq, PartialOrd, Ord, Hash)] pub struct Instant { std: std::time::Instant, } impl Instant { /// Returns an instant corresponding to "now". /// /// # Examples /// /// ``` /// use tokio::time::Instant; /// /// let now = Instant::now(); /// ``` pub fn now() -> Instant { variant::now() } /// Create a `tokio::time::Instant` from a `std::time::Instant`. pub fn from_std(std: std::time::Instant) -> Instant { Instant { std } } pub(crate) fn far_future() -> Instant { // Roughly 30 years from now. // API does not provide a way to obtain max `Instant` // or convert specific date in the future to instant. // 1000 years overflows on macOS, 100 years overflows on FreeBSD. Self::now() + Duration::from_secs(86400 * 365 * 30) } /// Convert the value into a `std::time::Instant`. pub fn into_std(self) -> std::time::Instant { self.std } /// Returns the amount of time elapsed from another instant to this one, or /// zero duration if that instant is later than this one. pub fn duration_since(&self, earlier: Instant) -> Duration { self.std.saturating_duration_since(earlier.std) } /// Returns the amount of time elapsed from another instant to this one, or /// None if that instant is later than this one. /// /// # Examples /// /// ``` /// use tokio::time::{Duration, Instant, sleep}; /// /// #[tokio::main] /// async fn main() { /// let now = Instant::now(); /// sleep(Duration::new(1, 0)).await; /// let new_now = Instant::now(); /// println!("{:?}", new_now.checked_duration_since(now)); /// println!("{:?}", now.checked_duration_since(new_now)); // None /// } /// ``` pub fn checked_duration_since(&self, earlier: Instant) -> Option { self.std.checked_duration_since(earlier.std) } /// Returns the amount of time elapsed from another instant to this one, or /// zero duration if that instant is later than this one. /// /// # Examples /// /// ``` /// use tokio::time::{Duration, Instant, sleep}; /// /// #[tokio::main] /// async fn main() { /// let now = Instant::now(); /// sleep(Duration::new(1, 0)).await; /// let new_now = Instant::now(); /// println!("{:?}", new_now.saturating_duration_since(now)); /// println!("{:?}", now.saturating_duration_since(new_now)); // 0ns /// } /// ``` pub fn saturating_duration_since(&self, earlier: Instant) -> Duration { self.std.saturating_duration_since(earlier.std) } /// Returns the amount of time elapsed since this instant was created, /// or zero duration if this instant is in the future. /// /// # Examples /// /// ``` /// use tokio::time::{Duration, Instant, sleep}; /// /// #[tokio::main] /// async fn main() { /// let instant = Instant::now(); /// let three_secs = Duration::from_secs(3); /// sleep(three_secs).await; /// assert!(instant.elapsed() >= three_secs); /// } /// ``` pub fn elapsed(&self) -> Duration { Instant::now().saturating_duration_since(*self) } /// Returns `Some(t)` where `t` is the time `self + duration` if `t` can be /// represented as `Instant` (which means it's inside the bounds of the /// underlying data structure), `None` otherwise. pub fn checked_add(&self, duration: Duration) -> Option { self.std.checked_add(duration).map(Instant::from_std) } /// Returns `Some(t)` where `t` is the time `self - duration` if `t` can be /// represented as `Instant` (which means it's inside the bounds of the /// underlying data structure), `None` otherwise. pub fn checked_sub(&self, duration: Duration) -> Option { self.std.checked_sub(duration).map(Instant::from_std) } } impl From for Instant { fn from(time: std::time::Instant) -> Instant { Instant::from_std(time) } } impl From for std::time::Instant { fn from(time: Instant) -> std::time::Instant { time.into_std() } } impl ops::Add for Instant { type Output = Instant; fn add(self, other: Duration) -> Instant { Instant::from_std(self.std + other) } } impl ops::AddAssign for Instant { fn add_assign(&mut self, rhs: Duration) { *self = *self + rhs; } } impl ops::Sub for Instant { type Output = Duration; fn sub(self, rhs: Instant) -> Duration { self.std.saturating_duration_since(rhs.std) } } impl ops::Sub for Instant { type Output = Instant; fn sub(self, rhs: Duration) -> Instant { Instant::from_std(std::time::Instant::sub(self.std, rhs)) } } impl ops::SubAssign for Instant { fn sub_assign(&mut self, rhs: Duration) { *self = *self - rhs; } } impl fmt::Debug for Instant { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { self.std.fmt(fmt) } } #[cfg(not(feature = "test-util"))] mod variant { use super::Instant; pub(super) fn now() -> Instant { Instant::from_std(std::time::Instant::now()) } } #[cfg(feature = "test-util")] mod variant { use super::Instant; pub(super) fn now() -> Instant { crate::time::clock::now() } } tokio-1.43.1/src/time/interval.rs000064400000000000000000000546241046102023000147710ustar 00000000000000use crate::time::{sleep_until, Duration, Instant, Sleep}; use crate::util::trace; use std::future::{poll_fn, Future}; use std::panic::Location; use std::pin::Pin; use std::task::{ready, Context, Poll}; /// Creates new [`Interval`] that yields with interval of `period`. The first /// tick completes immediately. The default [`MissedTickBehavior`] is /// [`Burst`](MissedTickBehavior::Burst), but this can be configured /// by calling [`set_missed_tick_behavior`](Interval::set_missed_tick_behavior). /// /// An interval will tick indefinitely. At any time, the [`Interval`] value can /// be dropped. This cancels the interval. /// /// This function is equivalent to /// [`interval_at(Instant::now(), period)`](interval_at). /// /// # Panics /// /// This function panics if `period` is zero. /// /// # Examples /// /// ``` /// use tokio::time::{self, Duration}; /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(Duration::from_millis(10)); /// /// interval.tick().await; // ticks immediately /// interval.tick().await; // ticks after 10ms /// interval.tick().await; // ticks after 10ms /// /// // approximately 20ms have elapsed. /// } /// ``` /// /// A simple example using `interval` to execute a task every two seconds. /// /// The difference between `interval` and [`sleep`] is that an [`Interval`] /// measures the time since the last tick, which means that [`.tick().await`] /// may wait for a shorter time than the duration specified for the interval /// if some time has passed between calls to [`.tick().await`]. /// /// If the tick in the example below was replaced with [`sleep`], the task /// would only be executed once every three seconds, and not every two /// seconds. /// /// ``` /// use tokio::time; /// /// async fn task_that_takes_a_second() { /// println!("hello"); /// time::sleep(time::Duration::from_secs(1)).await /// } /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(time::Duration::from_secs(2)); /// for _i in 0..5 { /// interval.tick().await; /// task_that_takes_a_second().await; /// } /// } /// ``` /// /// [`sleep`]: crate::time::sleep() /// [`.tick().await`]: Interval::tick #[track_caller] pub fn interval(period: Duration) -> Interval { assert!(period > Duration::new(0, 0), "`period` must be non-zero."); internal_interval_at(Instant::now(), period, trace::caller_location()) } /// Creates new [`Interval`] that yields with interval of `period` with the /// first tick completing at `start`. The default [`MissedTickBehavior`] is /// [`Burst`](MissedTickBehavior::Burst), but this can be configured /// by calling [`set_missed_tick_behavior`](Interval::set_missed_tick_behavior). /// /// An interval will tick indefinitely. At any time, the [`Interval`] value can /// be dropped. This cancels the interval. /// /// # Panics /// /// This function panics if `period` is zero. /// /// # Examples /// /// ``` /// use tokio::time::{interval_at, Duration, Instant}; /// /// #[tokio::main] /// async fn main() { /// let start = Instant::now() + Duration::from_millis(50); /// let mut interval = interval_at(start, Duration::from_millis(10)); /// /// interval.tick().await; // ticks after 50ms /// interval.tick().await; // ticks after 10ms /// interval.tick().await; // ticks after 10ms /// /// // approximately 70ms have elapsed. /// } /// ``` #[track_caller] pub fn interval_at(start: Instant, period: Duration) -> Interval { assert!(period > Duration::new(0, 0), "`period` must be non-zero."); internal_interval_at(start, period, trace::caller_location()) } #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_variables))] fn internal_interval_at( start: Instant, period: Duration, location: Option<&'static Location<'static>>, ) -> Interval { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = { let location = location.expect("should have location if tracing"); tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Interval", kind = "timer", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ) }; #[cfg(all(tokio_unstable, feature = "tracing"))] let delay = resource_span.in_scope(|| Box::pin(sleep_until(start))); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let delay = Box::pin(sleep_until(start)); Interval { delay, period, missed_tick_behavior: MissedTickBehavior::default(), #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span, } } /// Defines the behavior of an [`Interval`] when it misses a tick. /// /// Sometimes, an [`Interval`]'s tick is missed. For example, consider the /// following: /// /// ``` /// use tokio::time::{self, Duration}; /// # async fn task_that_takes_one_to_three_millis() {} /// /// #[tokio::main] /// async fn main() { /// // ticks every 2 milliseconds /// let mut interval = time::interval(Duration::from_millis(2)); /// for _ in 0..5 { /// interval.tick().await; /// // if this takes more than 2 milliseconds, a tick will be delayed /// task_that_takes_one_to_three_millis().await; /// } /// } /// ``` /// /// Generally, a tick is missed if too much time is spent without calling /// [`Interval::tick()`]. /// /// By default, when a tick is missed, [`Interval`] fires ticks as quickly as it /// can until it is "caught up" in time to where it should be. /// `MissedTickBehavior` can be used to specify a different behavior for /// [`Interval`] to exhibit. Each variant represents a different strategy. /// /// Note that because the executor cannot guarantee exact precision with timers, /// these strategies will only apply when the delay is greater than 5 /// milliseconds. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum MissedTickBehavior { /// Ticks as fast as possible until caught up. /// /// When this strategy is used, [`Interval`] schedules ticks "normally" (the /// same as it would have if the ticks hadn't been delayed), which results /// in it firing ticks as fast as possible until it is caught up in time to /// where it should be. Unlike [`Delay`] and [`Skip`], the ticks yielded /// when `Burst` is used (the [`Instant`]s that [`tick`](Interval::tick) /// yields) aren't different than they would have been if a tick had not /// been missed. Like [`Skip`], and unlike [`Delay`], the ticks may be /// shortened. /// /// This looks something like this: /// ```text /// Expected ticks: | 1 | 2 | 3 | 4 | 5 | 6 | /// Actual ticks: | work -----| delay | work | work | work -| work -----| /// ``` /// /// In code: /// /// ``` /// use tokio::time::{interval, Duration}; /// # async fn task_that_takes_200_millis() {} /// /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() { /// let mut interval = interval(Duration::from_millis(50)); /// /// // First tick resolves immediately after creation /// interval.tick().await; /// /// task_that_takes_200_millis().await; /// // The `Interval` has missed a tick /// /// // Since we have exceeded our timeout, this will resolve immediately /// interval.tick().await; /// /// // Since we are more than 100ms after the start of `interval`, this will /// // also resolve immediately. /// interval.tick().await; /// /// // Also resolves immediately, because it was supposed to resolve at /// // 150ms after the start of `interval` /// interval.tick().await; /// /// // Resolves immediately /// interval.tick().await; /// /// // Since we have gotten to 200ms after the start of `interval`, this /// // will resolve after 50ms /// interval.tick().await; /// # } /// ``` /// /// This is the default behavior when [`Interval`] is created with /// [`interval`] and [`interval_at`]. /// /// [`Delay`]: MissedTickBehavior::Delay /// [`Skip`]: MissedTickBehavior::Skip Burst, /// Tick at multiples of `period` from when [`tick`] was called, rather than /// from `start`. /// /// When this strategy is used and [`Interval`] has missed a tick, instead /// of scheduling ticks to fire at multiples of `period` from `start` (the /// time when the first tick was fired), it schedules all future ticks to /// happen at a regular `period` from the point when [`tick`] was called. /// Unlike [`Burst`] and [`Skip`], ticks are not shortened, and they aren't /// guaranteed to happen at a multiple of `period` from `start` any longer. /// /// This looks something like this: /// ```text /// Expected ticks: | 1 | 2 | 3 | 4 | 5 | 6 | /// Actual ticks: | work -----| delay | work -----| work -----| work -----| /// ``` /// /// In code: /// /// ``` /// use tokio::time::{interval, Duration, MissedTickBehavior}; /// # async fn task_that_takes_more_than_50_millis() {} /// /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() { /// let mut interval = interval(Duration::from_millis(50)); /// interval.set_missed_tick_behavior(MissedTickBehavior::Delay); /// /// task_that_takes_more_than_50_millis().await; /// // The `Interval` has missed a tick /// /// // Since we have exceeded our timeout, this will resolve immediately /// interval.tick().await; /// /// // But this one, rather than also resolving immediately, as might happen /// // with the `Burst` or `Skip` behaviors, will not resolve until /// // 50ms after the call to `tick` up above. That is, in `tick`, when we /// // recognize that we missed a tick, we schedule the next tick to happen /// // 50ms (or whatever the `period` is) from right then, not from when /// // were *supposed* to tick /// interval.tick().await; /// # } /// ``` /// /// [`Burst`]: MissedTickBehavior::Burst /// [`Skip`]: MissedTickBehavior::Skip /// [`tick`]: Interval::tick Delay, /// Skips missed ticks and tick on the next multiple of `period` from /// `start`. /// /// When this strategy is used, [`Interval`] schedules the next tick to fire /// at the next-closest tick that is a multiple of `period` away from /// `start` (the point where [`Interval`] first ticked). Like [`Burst`], all /// ticks remain multiples of `period` away from `start`, but unlike /// [`Burst`], the ticks may not be *one* multiple of `period` away from the /// last tick. Like [`Delay`], the ticks are no longer the same as they /// would have been if ticks had not been missed, but unlike [`Delay`], and /// like [`Burst`], the ticks may be shortened to be less than one `period` /// away from each other. /// /// This looks something like this: /// ```text /// Expected ticks: | 1 | 2 | 3 | 4 | 5 | 6 | /// Actual ticks: | work -----| delay | work ---| work -----| work -----| /// ``` /// /// In code: /// /// ``` /// use tokio::time::{interval, Duration, MissedTickBehavior}; /// # async fn task_that_takes_75_millis() {} /// /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() { /// let mut interval = interval(Duration::from_millis(50)); /// interval.set_missed_tick_behavior(MissedTickBehavior::Skip); /// /// task_that_takes_75_millis().await; /// // The `Interval` has missed a tick /// /// // Since we have exceeded our timeout, this will resolve immediately /// interval.tick().await; /// /// // This one will resolve after 25ms, 100ms after the start of /// // `interval`, which is the closest multiple of `period` from the start /// // of `interval` after the call to `tick` up above. /// interval.tick().await; /// # } /// ``` /// /// [`Burst`]: MissedTickBehavior::Burst /// [`Delay`]: MissedTickBehavior::Delay Skip, } impl MissedTickBehavior { /// If a tick is missed, this method is called to determine when the next tick should happen. fn next_timeout(&self, timeout: Instant, now: Instant, period: Duration) -> Instant { match self { Self::Burst => timeout + period, Self::Delay => now + period, Self::Skip => { now + period - Duration::from_nanos( ((now - timeout).as_nanos() % period.as_nanos()) .try_into() // This operation is practically guaranteed not to // fail, as in order for it to fail, `period` would // have to be longer than `now - timeout`, and both // would have to be longer than 584 years. // // If it did fail, there's not a good way to pass // the error along to the user, so we just panic. .expect( "too much time has elapsed since the interval was supposed to tick", ), ) } } } } impl Default for MissedTickBehavior { /// Returns [`MissedTickBehavior::Burst`]. /// /// For most usecases, the [`Burst`] strategy is what is desired. /// Additionally, to preserve backwards compatibility, the [`Burst`] /// strategy must be the default. For these reasons, /// [`MissedTickBehavior::Burst`] is the default for [`MissedTickBehavior`]. /// See [`Burst`] for more details. /// /// [`Burst`]: MissedTickBehavior::Burst fn default() -> Self { Self::Burst } } /// Interval returned by [`interval`] and [`interval_at`]. /// /// This type allows you to wait on a sequence of instants with a certain /// duration between each instant. Unlike calling [`sleep`] in a loop, this lets /// you count the time spent between the calls to [`sleep`] as well. /// /// An `Interval` can be turned into a `Stream` with [`IntervalStream`]. /// /// [`IntervalStream`]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.IntervalStream.html /// [`sleep`]: crate::time::sleep() #[derive(Debug)] pub struct Interval { /// Future that completes the next time the `Interval` yields a value. delay: Pin>, /// The duration between values yielded by `Interval`. period: Duration, /// The strategy `Interval` should use when a tick is missed. missed_tick_behavior: MissedTickBehavior, #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, } impl Interval { /// Completes when the next instant in the interval has been reached. /// /// # Cancel safety /// /// This method is cancellation safe. If `tick` is used as the branch in a `tokio::select!` and /// another branch completes first, then no tick has been consumed. /// /// # Examples /// /// ``` /// use tokio::time; /// /// use std::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(Duration::from_millis(10)); /// /// interval.tick().await; /// // approximately 0ms have elapsed. The first tick completes immediately. /// interval.tick().await; /// interval.tick().await; /// /// // approximately 20ms have elapsed. /// } /// ``` pub async fn tick(&mut self) -> Instant { #[cfg(all(tokio_unstable, feature = "tracing"))] let resource_span = self.resource_span.clone(); #[cfg(all(tokio_unstable, feature = "tracing"))] let instant = trace::async_op( || poll_fn(|cx| self.poll_tick(cx)), resource_span, "Interval::tick", "poll_tick", false, ); #[cfg(not(all(tokio_unstable, feature = "tracing")))] let instant = poll_fn(|cx| self.poll_tick(cx)); instant.await } /// Polls for the next instant in the interval to be reached. /// /// This method can return the following values: /// /// * `Poll::Pending` if the next instant has not yet been reached. /// * `Poll::Ready(instant)` if the next instant has been reached. /// /// When this method returns `Poll::Pending`, the current task is scheduled /// to receive a wakeup when the instant has elapsed. Note that on multiple /// calls to `poll_tick`, only the [`Waker`](std::task::Waker) from the /// [`Context`] passed to the most recent call is scheduled to receive a /// wakeup. pub fn poll_tick(&mut self, cx: &mut Context<'_>) -> Poll { // Wait for the delay to be done ready!(Pin::new(&mut self.delay).poll(cx)); // Get the time when we were scheduled to tick let timeout = self.delay.deadline(); let now = Instant::now(); // If a tick was not missed, and thus we are being called before the // next tick is due, just schedule the next tick normally, one `period` // after `timeout` // // However, if a tick took excessively long and we are now behind, // schedule the next tick according to how the user specified with // `MissedTickBehavior` let next = if now > timeout + Duration::from_millis(5) { self.missed_tick_behavior .next_timeout(timeout, now, self.period) } else { timeout .checked_add(self.period) .unwrap_or_else(Instant::far_future) }; // When we arrive here, the internal delay returned `Poll::Ready`. // Reset the delay but do not register it. It should be registered with // the next call to [`poll_tick`]. self.delay.as_mut().reset_without_reregister(next); // Return the time when we were scheduled to tick Poll::Ready(timeout) } /// Resets the interval to complete one period after the current time. /// /// This method ignores [`MissedTickBehavior`] strategy. /// /// This is equivalent to calling `reset_at(Instant::now() + period)`. /// /// # Examples /// /// ``` /// use tokio::time; /// /// use std::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(Duration::from_millis(100)); /// /// interval.tick().await; /// /// time::sleep(Duration::from_millis(50)).await; /// interval.reset(); /// /// interval.tick().await; /// interval.tick().await; /// /// // approximately 250ms have elapsed. /// } /// ``` pub fn reset(&mut self) { self.delay.as_mut().reset(Instant::now() + self.period); } /// Resets the interval immediately. /// /// This method ignores [`MissedTickBehavior`] strategy. /// /// This is equivalent to calling `reset_at(Instant::now())`. /// /// # Examples /// /// ``` /// use tokio::time; /// /// use std::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(Duration::from_millis(100)); /// /// interval.tick().await; /// /// time::sleep(Duration::from_millis(50)).await; /// interval.reset_immediately(); /// /// interval.tick().await; /// interval.tick().await; /// /// // approximately 150ms have elapsed. /// } /// ``` pub fn reset_immediately(&mut self) { self.delay.as_mut().reset(Instant::now()); } /// Resets the interval after the specified [`std::time::Duration`]. /// /// This method ignores [`MissedTickBehavior`] strategy. /// /// This is equivalent to calling `reset_at(Instant::now() + after)`. /// /// # Examples /// /// ``` /// use tokio::time; /// /// use std::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(Duration::from_millis(100)); /// interval.tick().await; /// /// time::sleep(Duration::from_millis(50)).await; /// /// let after = Duration::from_millis(20); /// interval.reset_after(after); /// /// interval.tick().await; /// interval.tick().await; /// /// // approximately 170ms have elapsed. /// } /// ``` pub fn reset_after(&mut self, after: Duration) { self.delay.as_mut().reset(Instant::now() + after); } /// Resets the interval to a [`crate::time::Instant`] deadline. /// /// Sets the next tick to expire at the given instant. If the instant is in /// the past, then the [`MissedTickBehavior`] strategy will be used to /// catch up. If the instant is in the future, then the next tick will /// complete at the given instant, even if that means that it will sleep for /// longer than the duration of this [`Interval`]. If the [`Interval`] had /// any missed ticks before calling this method, then those are discarded. /// /// # Examples /// /// ``` /// use tokio::time::{self, Instant}; /// /// use std::time::Duration; /// /// #[tokio::main] /// async fn main() { /// let mut interval = time::interval(Duration::from_millis(100)); /// interval.tick().await; /// /// time::sleep(Duration::from_millis(50)).await; /// /// let deadline = Instant::now() + Duration::from_millis(30); /// interval.reset_at(deadline); /// /// interval.tick().await; /// interval.tick().await; /// /// // approximately 180ms have elapsed. /// } /// ``` pub fn reset_at(&mut self, deadline: Instant) { self.delay.as_mut().reset(deadline); } /// Returns the [`MissedTickBehavior`] strategy currently being used. pub fn missed_tick_behavior(&self) -> MissedTickBehavior { self.missed_tick_behavior } /// Sets the [`MissedTickBehavior`] strategy that should be used. pub fn set_missed_tick_behavior(&mut self, behavior: MissedTickBehavior) { self.missed_tick_behavior = behavior; } /// Returns the period of the interval. pub fn period(&self) -> Duration { self.period } } tokio-1.43.1/src/time/mod.rs000064400000000000000000000056461046102023000137240ustar 00000000000000//! Utilities for tracking time. //! //! This module provides a number of types for executing code after a set period //! of time. //! //! * [`Sleep`] is a future that does no work and completes at a specific [`Instant`] //! in time. //! //! * [`Interval`] is a stream yielding a value at a fixed period. It is //! initialized with a [`Duration`] and repeatedly yields each time the duration //! elapses. //! //! * [`Timeout`]: Wraps a future or stream, setting an upper bound to the amount //! of time it is allowed to execute. If the future or stream does not //! complete in time, then it is canceled and an error is returned. //! //! These types are sufficient for handling a large number of scenarios //! involving time. //! //! These types must be used from within the context of the [`Runtime`](crate::runtime::Runtime). //! //! # Examples //! //! Wait 100ms and print "100 ms have elapsed" //! //! ``` //! use std::time::Duration; //! use tokio::time::sleep; //! //! #[tokio::main] //! async fn main() { //! sleep(Duration::from_millis(100)).await; //! println!("100 ms have elapsed"); //! } //! ``` //! //! Require that an operation takes no more than 1s. //! //! ``` //! use tokio::time::{timeout, Duration}; //! //! async fn long_future() { //! // do work here //! } //! //! # async fn dox() { //! let res = timeout(Duration::from_secs(1), long_future()).await; //! //! if res.is_err() { //! println!("operation timed out"); //! } //! # } //! ``` //! //! A simple example using [`interval`] to execute a task every two seconds. //! //! The difference between [`interval`] and [`sleep`] is that an [`interval`] //! measures the time since the last tick, which means that `.tick().await` may //! wait for a shorter time than the duration specified for the interval //! if some time has passed between calls to `.tick().await`. //! //! If the tick in the example below was replaced with [`sleep`], the task //! would only be executed once every three seconds, and not every two //! seconds. //! //! ``` //! use tokio::time; //! //! async fn task_that_takes_a_second() { //! println!("hello"); //! time::sleep(time::Duration::from_secs(1)).await //! } //! //! #[tokio::main] //! async fn main() { //! let mut interval = time::interval(time::Duration::from_secs(2)); //! for _i in 0..5 { //! interval.tick().await; //! task_that_takes_a_second().await; //! } //! } //! ``` //! //! [`interval`]: crate::time::interval() //! [`sleep`]: sleep() mod clock; pub(crate) use self::clock::Clock; cfg_test_util! { pub use clock::{advance, pause, resume}; } pub mod error; mod instant; pub use self::instant::Instant; mod interval; pub use interval::{interval, interval_at, Interval, MissedTickBehavior}; mod sleep; pub use sleep::{sleep, sleep_until, Sleep}; mod timeout; #[doc(inline)] pub use timeout::{timeout, timeout_at, Timeout}; // Re-export for convenience #[doc(no_inline)] pub use std::time::Duration; tokio-1.43.1/src/time/sleep.rs000064400000000000000000000361071046102023000142510ustar 00000000000000use crate::runtime::time::TimerEntry; use crate::time::{error::Error, Duration, Instant}; use crate::util::trace; use pin_project_lite::pin_project; use std::future::Future; use std::panic::Location; use std::pin::Pin; use std::task::{self, ready, Poll}; /// Waits until `deadline` is reached. /// /// No work is performed while awaiting on the sleep future to complete. `Sleep` /// operates at millisecond granularity and should not be used for tasks that /// require high-resolution timers. /// /// To run something regularly on a schedule, see [`interval`]. /// /// # Cancellation /// /// Canceling a sleep instance is done by dropping the returned future. No additional /// cleanup work is required. /// /// # Examples /// /// Wait 100ms and print "100 ms have elapsed". /// /// ``` /// use tokio::time::{sleep_until, Instant, Duration}; /// /// #[tokio::main] /// async fn main() { /// sleep_until(Instant::now() + Duration::from_millis(100)).await; /// println!("100 ms have elapsed"); /// } /// ``` /// /// See the documentation for the [`Sleep`] type for more examples. /// /// # Panics /// /// This function panics if there is no current timer set. /// /// It can be triggered when [`Builder::enable_time`] or /// [`Builder::enable_all`] are not included in the builder. /// /// It can also panic whenever a timer is created outside of a /// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, /// since the function is executed outside of the runtime. /// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. /// And this is because wrapping the function on an async makes it lazy, /// and so gets executed inside the runtime successfully without /// panicking. /// /// [`Sleep`]: struct@crate::time::Sleep /// [`interval`]: crate::time::interval() /// [`Builder::enable_time`]: crate::runtime::Builder::enable_time /// [`Builder::enable_all`]: crate::runtime::Builder::enable_all // Alias for old name in 0.x #[cfg_attr(docsrs, doc(alias = "delay_until"))] #[track_caller] pub fn sleep_until(deadline: Instant) -> Sleep { Sleep::new_timeout(deadline, trace::caller_location()) } /// Waits until `duration` has elapsed. /// /// Equivalent to `sleep_until(Instant::now() + duration)`. An asynchronous /// analog to `std::thread::sleep`. /// /// No work is performed while awaiting on the sleep future to complete. `Sleep` /// operates at millisecond granularity and should not be used for tasks that /// require high-resolution timers. The implementation is platform specific, /// and some platforms (specifically Windows) will provide timers with a /// larger resolution than 1 ms. /// /// To run something regularly on a schedule, see [`interval`]. /// /// The maximum duration for a sleep is 68719476734 milliseconds (approximately 2.2 years). /// /// # Cancellation /// /// Canceling a sleep instance is done by dropping the returned future. No additional /// cleanup work is required. /// /// # Examples /// /// Wait 100ms and print "100 ms have elapsed". /// /// ``` /// use tokio::time::{sleep, Duration}; /// /// #[tokio::main] /// async fn main() { /// sleep(Duration::from_millis(100)).await; /// println!("100 ms have elapsed"); /// } /// ``` /// /// See the documentation for the [`Sleep`] type for more examples. /// /// # Panics /// /// This function panics if there is no current timer set. /// /// It can be triggered when [`Builder::enable_time`] or /// [`Builder::enable_all`] are not included in the builder. /// /// It can also panic whenever a timer is created outside of a /// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, /// since the function is executed outside of the runtime. /// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. /// And this is because wrapping the function on an async makes it lazy, /// and so gets executed inside the runtime successfully without /// panicking. /// /// [`Sleep`]: struct@crate::time::Sleep /// [`interval`]: crate::time::interval() /// [`Builder::enable_time`]: crate::runtime::Builder::enable_time /// [`Builder::enable_all`]: crate::runtime::Builder::enable_all // Alias for old name in 0.x #[cfg_attr(docsrs, doc(alias = "delay_for"))] #[cfg_attr(docsrs, doc(alias = "wait"))] #[track_caller] pub fn sleep(duration: Duration) -> Sleep { let location = trace::caller_location(); match Instant::now().checked_add(duration) { Some(deadline) => Sleep::new_timeout(deadline, location), None => Sleep::new_timeout(Instant::far_future(), location), } } pin_project! { /// Future returned by [`sleep`](sleep) and [`sleep_until`](sleep_until). /// /// This type does not implement the `Unpin` trait, which means that if you /// use it with [`select!`] or by calling `poll`, you have to pin it first. /// If you use it with `.await`, this does not apply. /// /// # Examples /// /// Wait 100ms and print "100 ms have elapsed". /// /// ``` /// use tokio::time::{sleep, Duration}; /// /// #[tokio::main] /// async fn main() { /// sleep(Duration::from_millis(100)).await; /// println!("100 ms have elapsed"); /// } /// ``` /// /// Use with [`select!`]. Pinning the `Sleep` with [`tokio::pin!`] is /// necessary when the same `Sleep` is selected on multiple times. /// ```no_run /// use tokio::time::{self, Duration, Instant}; /// /// #[tokio::main] /// async fn main() { /// let sleep = time::sleep(Duration::from_millis(10)); /// tokio::pin!(sleep); /// /// loop { /// tokio::select! { /// () = &mut sleep => { /// println!("timer elapsed"); /// sleep.as_mut().reset(Instant::now() + Duration::from_millis(50)); /// }, /// } /// } /// } /// ``` /// Use in a struct with boxing. By pinning the `Sleep` with a `Box`, the /// `HasSleep` struct implements `Unpin`, even though `Sleep` does not. /// ``` /// use std::future::Future; /// use std::pin::Pin; /// use std::task::{Context, Poll}; /// use tokio::time::Sleep; /// /// struct HasSleep { /// sleep: Pin>, /// } /// /// impl Future for HasSleep { /// type Output = (); /// /// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { /// self.sleep.as_mut().poll(cx) /// } /// } /// ``` /// Use in a struct with pin projection. This method avoids the `Box`, but /// the `HasSleep` struct will not be `Unpin` as a consequence. /// ``` /// use std::future::Future; /// use std::pin::Pin; /// use std::task::{Context, Poll}; /// use tokio::time::Sleep; /// use pin_project_lite::pin_project; /// /// pin_project! { /// struct HasSleep { /// #[pin] /// sleep: Sleep, /// } /// } /// /// impl Future for HasSleep { /// type Output = (); /// /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { /// self.project().sleep.poll(cx) /// } /// } /// ``` /// /// [`select!`]: ../macro.select.html /// [`tokio::pin!`]: ../macro.pin.html #[project(!Unpin)] // Alias for old name in 0.2 #[cfg_attr(docsrs, doc(alias = "Delay"))] #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Sleep { inner: Inner, // The link between the `Sleep` instance and the timer that drives it. #[pin] entry: TimerEntry, } } cfg_trace! { #[derive(Debug)] struct Inner { ctx: trace::AsyncOpTracingCtx, } } cfg_not_trace! { #[derive(Debug)] struct Inner { } } impl Sleep { #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_variables))] #[track_caller] pub(crate) fn new_timeout( deadline: Instant, location: Option<&'static Location<'static>>, ) -> Sleep { use crate::runtime::scheduler; let handle = scheduler::Handle::current(); let entry = TimerEntry::new(handle, deadline); #[cfg(all(tokio_unstable, feature = "tracing"))] let inner = { let handle = scheduler::Handle::current(); let clock = handle.driver().clock(); let handle = &handle.driver().time(); let time_source = handle.time_source(); let deadline_tick = time_source.deadline_to_tick(deadline); let duration = deadline_tick.saturating_sub(time_source.now(clock)); let location = location.expect("should have location if tracing"); let resource_span = tracing::trace_span!( parent: None, "runtime.resource", concrete_type = "Sleep", kind = "timer", loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ); let async_op_span = resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", duration = duration, duration.unit = "ms", duration.op = "override", ); tracing::trace_span!("runtime.resource.async_op", source = "Sleep::new_timeout") }); let async_op_poll_span = async_op_span.in_scope(|| tracing::trace_span!("runtime.resource.async_op.poll")); let ctx = trace::AsyncOpTracingCtx { async_op_span, async_op_poll_span, resource_span, }; Inner { ctx } }; #[cfg(not(all(tokio_unstable, feature = "tracing")))] let inner = Inner {}; Sleep { inner, entry } } pub(crate) fn far_future(location: Option<&'static Location<'static>>) -> Sleep { Self::new_timeout(Instant::far_future(), location) } /// Returns the instant at which the future will complete. pub fn deadline(&self) -> Instant { self.entry.deadline() } /// Returns `true` if `Sleep` has elapsed. /// /// A `Sleep` instance is elapsed when the requested duration has elapsed. pub fn is_elapsed(&self) -> bool { self.entry.is_elapsed() } /// Resets the `Sleep` instance to a new deadline. /// /// Calling this function allows changing the instant at which the `Sleep` /// future completes without having to create new associated state. /// /// This function can be called both before and after the future has /// completed. /// /// To call this method, you will usually combine the call with /// [`Pin::as_mut`], which lets you call the method without consuming the /// `Sleep` itself. /// /// # Example /// /// ``` /// use tokio::time::{Duration, Instant}; /// /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() { /// let sleep = tokio::time::sleep(Duration::from_millis(10)); /// tokio::pin!(sleep); /// /// sleep.as_mut().reset(Instant::now() + Duration::from_millis(20)); /// # } /// ``` /// /// See also the top-level examples. /// /// [`Pin::as_mut`]: fn@std::pin::Pin::as_mut pub fn reset(self: Pin<&mut Self>, deadline: Instant) { self.reset_inner(deadline); } /// Resets the `Sleep` instance to a new deadline without reregistering it /// to be woken up. /// /// Calling this function allows changing the instant at which the `Sleep` /// future completes without having to create new associated state and /// without having it registered. This is required in e.g. the /// [`crate::time::Interval`] where we want to reset the internal [Sleep] /// without having it wake up the last task that polled it. pub(crate) fn reset_without_reregister(self: Pin<&mut Self>, deadline: Instant) { let mut me = self.project(); me.entry.as_mut().reset(deadline, false); } fn reset_inner(self: Pin<&mut Self>, deadline: Instant) { let mut me = self.project(); me.entry.as_mut().reset(deadline, true); #[cfg(all(tokio_unstable, feature = "tracing"))] { let _resource_enter = me.inner.ctx.resource_span.enter(); me.inner.ctx.async_op_span = tracing::trace_span!("runtime.resource.async_op", source = "Sleep::reset"); let _async_op_enter = me.inner.ctx.async_op_span.enter(); me.inner.ctx.async_op_poll_span = tracing::trace_span!("runtime.resource.async_op.poll"); let duration = { let clock = me.entry.clock(); let time_source = me.entry.driver().time_source(); let now = time_source.now(clock); let deadline_tick = time_source.deadline_to_tick(deadline); deadline_tick.saturating_sub(now) }; tracing::trace!( target: "runtime::resource::state_update", duration = duration, duration.unit = "ms", duration.op = "override", ); } } fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll> { let me = self.project(); ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget #[cfg(all(tokio_unstable, feature = "tracing"))] let coop = ready!(trace_poll_op!( "poll_elapsed", crate::runtime::coop::poll_proceed(cx), )); #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let result = me.entry.poll_elapsed(cx).map(move |r| { coop.made_progress(); r }); #[cfg(all(tokio_unstable, feature = "tracing"))] return trace_poll_op!("poll_elapsed", result); #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] return result; } } impl Future for Sleep { type Output = (); // `poll_elapsed` can return an error in two cases: // // - AtCapacity: this is a pathological case where far too many // sleep instances have been scheduled. // - Shutdown: No timer has been setup, which is a mis-use error. // // Both cases are extremely rare, and pretty accurately fit into // "logic errors", so we just panic in this case. A user couldn't // really do much better if we passed the error onwards. fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll { #[cfg(all(tokio_unstable, feature = "tracing"))] let _res_span = self.inner.ctx.resource_span.clone().entered(); #[cfg(all(tokio_unstable, feature = "tracing"))] let _ao_span = self.inner.ctx.async_op_span.clone().entered(); #[cfg(all(tokio_unstable, feature = "tracing"))] let _ao_poll_span = self.inner.ctx.async_op_poll_span.clone().entered(); match ready!(self.as_mut().poll_elapsed(cx)) { Ok(()) => Poll::Ready(()), Err(e) => panic!("timer error: {e}"), } } } tokio-1.43.1/src/time/timeout.rs000064400000000000000000000156701046102023000146310ustar 00000000000000//! Allows a future to execute for a maximum amount of time. //! //! See [`Timeout`] documentation for more details. //! //! [`Timeout`]: struct@Timeout use crate::{ runtime::coop, time::{error::Elapsed, sleep_until, Duration, Instant, Sleep}, util::trace, }; use pin_project_lite::pin_project; use std::future::{Future, IntoFuture}; use std::pin::Pin; use std::task::{self, Poll}; /// Requires a `Future` to complete before the specified duration has elapsed. /// /// If the future completes before the duration has elapsed, then the completed /// value is returned. Otherwise, an error is returned and the future is /// canceled. /// /// Note that the timeout is checked before polling the future, so if the future /// does not yield during execution then it is possible for the future to complete /// and exceed the timeout _without_ returning an error. /// /// This function returns a future whose return type is [`Result`]``, where `T` is the /// return type of the provided future. /// /// If the provided future completes immediately, then the future returned from /// this function is guaranteed to complete immediately with an [`Ok`] variant /// no matter the provided duration. /// /// [`Ok`]: std::result::Result::Ok /// [`Result`]: std::result::Result /// [`Elapsed`]: crate::time::error::Elapsed /// /// # Cancellation /// /// Cancelling a timeout is done by dropping the future. No additional cleanup /// or other work is required. /// /// The original future may be obtained by calling [`Timeout::into_inner`]. This /// consumes the `Timeout`. /// /// # Examples /// /// Create a new `Timeout` set to expire in 10 milliseconds. /// /// ```rust /// use tokio::time::timeout; /// use tokio::sync::oneshot; /// /// use std::time::Duration; /// /// # async fn dox() { /// let (tx, rx) = oneshot::channel(); /// # tx.send(()).unwrap(); /// /// // Wrap the future with a `Timeout` set to expire in 10 milliseconds. /// if let Err(_) = timeout(Duration::from_millis(10), rx).await { /// println!("did not receive value within 10 ms"); /// } /// # } /// ``` /// /// # Panics /// /// This function panics if there is no current timer set. /// /// It can be triggered when [`Builder::enable_time`] or /// [`Builder::enable_all`] are not included in the builder. /// /// It can also panic whenever a timer is created outside of a /// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic, /// since the function is executed outside of the runtime. /// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic. /// And this is because wrapping the function on an async makes it lazy, /// and so gets executed inside the runtime successfully without /// panicking. /// /// [`Builder::enable_time`]: crate::runtime::Builder::enable_time /// [`Builder::enable_all`]: crate::runtime::Builder::enable_all #[track_caller] pub fn timeout(duration: Duration, future: F) -> Timeout where F: IntoFuture, { let location = trace::caller_location(); let deadline = Instant::now().checked_add(duration); let delay = match deadline { Some(deadline) => Sleep::new_timeout(deadline, location), None => Sleep::far_future(location), }; Timeout::new_with_delay(future.into_future(), delay) } /// Requires a `Future` to complete before the specified instant in time. /// /// If the future completes before the instant is reached, then the completed /// value is returned. Otherwise, an error is returned. /// /// This function returns a future whose return type is [`Result`]``, where `T` is the /// return type of the provided future. /// /// If the provided future completes immediately, then the future returned from /// this function is guaranteed to complete immediately with an [`Ok`] variant /// no matter the provided deadline. /// /// [`Ok`]: std::result::Result::Ok /// [`Result`]: std::result::Result /// [`Elapsed`]: crate::time::error::Elapsed /// /// # Cancellation /// /// Cancelling a timeout is done by dropping the future. No additional cleanup /// or other work is required. /// /// The original future may be obtained by calling [`Timeout::into_inner`]. This /// consumes the `Timeout`. /// /// # Examples /// /// Create a new `Timeout` set to expire in 10 milliseconds. /// /// ```rust /// use tokio::time::{Instant, timeout_at}; /// use tokio::sync::oneshot; /// /// use std::time::Duration; /// /// # async fn dox() { /// let (tx, rx) = oneshot::channel(); /// # tx.send(()).unwrap(); /// /// // Wrap the future with a `Timeout` set to expire 10 milliseconds into the /// // future. /// if let Err(_) = timeout_at(Instant::now() + Duration::from_millis(10), rx).await { /// println!("did not receive value within 10 ms"); /// } /// # } /// ``` pub fn timeout_at(deadline: Instant, future: F) -> Timeout where F: IntoFuture, { let delay = sleep_until(deadline); Timeout { value: future.into_future(), delay, } } pin_project! { /// Future returned by [`timeout`](timeout) and [`timeout_at`](timeout_at). #[must_use = "futures do nothing unless you `.await` or poll them"] #[derive(Debug)] pub struct Timeout { #[pin] value: T, #[pin] delay: Sleep, } } impl Timeout { pub(crate) fn new_with_delay(value: T, delay: Sleep) -> Timeout { Timeout { value, delay } } /// Gets a reference to the underlying value in this timeout. pub fn get_ref(&self) -> &T { &self.value } /// Gets a mutable reference to the underlying value in this timeout. pub fn get_mut(&mut self) -> &mut T { &mut self.value } /// Consumes this timeout, returning the underlying value. pub fn into_inner(self) -> T { self.value } } impl Future for Timeout where T: Future, { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll { let me = self.project(); let had_budget_before = coop::has_budget_remaining(); // First, try polling the future if let Poll::Ready(v) = me.value.poll(cx) { return Poll::Ready(Ok(v)); } let has_budget_now = coop::has_budget_remaining(); let delay = me.delay; let poll_delay = || -> Poll { match delay.poll(cx) { Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())), Poll::Pending => Poll::Pending, } }; if let (true, false) = (had_budget_before, has_budget_now) { // if it is the underlying future that exhausted the budget, we poll // the `delay` with an unconstrained one. This prevents pathological // cases where the underlying future always exhausts the budget and // we never get a chance to evaluate whether the timeout was hit or // not. coop::with_unconstrained(poll_delay) } else { poll_delay() } } } tokio-1.43.1/src/util/atomic_cell.rs000064400000000000000000000021721046102023000154260ustar 00000000000000use crate::loom::sync::atomic::AtomicPtr; use std::ptr; use std::sync::atomic::Ordering::AcqRel; pub(crate) struct AtomicCell { data: AtomicPtr, } unsafe impl Send for AtomicCell {} unsafe impl Sync for AtomicCell {} impl AtomicCell { pub(crate) fn new(data: Option>) -> AtomicCell { AtomicCell { data: AtomicPtr::new(to_raw(data)), } } pub(crate) fn swap(&self, val: Option>) -> Option> { let old = self.data.swap(to_raw(val), AcqRel); from_raw(old) } pub(crate) fn set(&self, val: Box) { let _ = self.swap(Some(val)); } pub(crate) fn take(&self) -> Option> { self.swap(None) } } fn to_raw(data: Option>) -> *mut T { data.map_or(ptr::null_mut(), Box::into_raw) } fn from_raw(val: *mut T) -> Option> { if val.is_null() { None } else { Some(unsafe { Box::from_raw(val) }) } } impl Drop for AtomicCell { fn drop(&mut self) { // Free any data still held by the cell let _ = self.take(); } } tokio-1.43.1/src/util/bit.rs000064400000000000000000000032551046102023000137340ustar 00000000000000use std::fmt; #[derive(Clone, Copy, PartialEq)] pub(crate) struct Pack { mask: usize, shift: u32, } impl Pack { /// Value is packed in the `width` least-significant bits. pub(crate) const fn least_significant(width: u32) -> Pack { let mask = mask_for(width); Pack { mask, shift: 0 } } /// Value is packed in the `width` more-significant bits. pub(crate) const fn then(&self, width: u32) -> Pack { let shift = usize::BITS - self.mask.leading_zeros(); let mask = mask_for(width) << shift; Pack { mask, shift } } /// Width, in bits, dedicated to storing the value. pub(crate) const fn width(&self) -> u32 { usize::BITS - (self.mask >> self.shift).leading_zeros() } /// Max representable value. pub(crate) const fn max_value(&self) -> usize { (1 << self.width()) - 1 } pub(crate) fn pack(&self, value: usize, base: usize) -> usize { assert!(value <= self.max_value()); (base & !self.mask) | (value << self.shift) } pub(crate) fn unpack(&self, src: usize) -> usize { unpack(src, self.mask, self.shift) } } impl fmt::Debug for Pack { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!( fmt, "Pack {{ mask: {:b}, shift: {} }}", self.mask, self.shift ) } } /// Returns a `usize` with the right-most `n` bits set. pub(crate) const fn mask_for(n: u32) -> usize { let shift = 1usize.wrapping_shl(n - 1); shift | (shift - 1) } /// Unpacks a value using a mask & shift. pub(crate) const fn unpack(src: usize, mask: usize, shift: u32) -> usize { (src & mask) >> shift } tokio-1.43.1/src/util/cacheline.rs000064400000000000000000000062011046102023000150630ustar 00000000000000#![cfg_attr(not(feature = "sync"), allow(dead_code, unreachable_pub))] use std::ops::{Deref, DerefMut}; /// Pads and aligns a value to the length of a cache line. #[derive(Clone, Copy, Default, Hash, PartialEq, Eq)] // Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache // lines at a time, so we have to align to 128 bytes rather than 64. // // Sources: // - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf // - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107 // // ARM's big.LITTLE architecture has asymmetric cores and "big" cores have 128-byte cache line size. // // Sources: // - https://www.mono-project.com/news/2016/09/12/arm64-icache/ // // powerpc64 has 128-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_ppc64x.go#L9 #[cfg_attr( any( target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", ), repr(align(128)) )] // arm, mips and mips64 have 32-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_arm.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mipsle.go#L7 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips64x.go#L9 #[cfg_attr( any(target_arch = "arm", target_arch = "mips", target_arch = "mips64",), repr(align(32)) )] // s390x has 256-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_s390x.go#L7 #[cfg_attr(target_arch = "s390x", repr(align(256)))] // x86, riscv and wasm have 64-byte cache line size. // // Sources: // - https://github.com/golang/go/blob/dda2991c2ea0c5914714469c4defc2562a907230/src/internal/cpu/cpu_x86.go#L9 // - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_wasm.go#L7 // - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/riscv/include/asm/cache.h#L10 // // All others are assumed to have 64-byte cache line size. #[cfg_attr( not(any( target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64", target_arch = "arm", target_arch = "mips", target_arch = "mips64", target_arch = "s390x", )), repr(align(64)) )] pub(crate) struct CachePadded { value: T, } impl CachePadded { /// Pads and aligns a value to the length of a cache line. pub(crate) fn new(value: T) -> CachePadded { CachePadded:: { value } } } impl Deref for CachePadded { type Target = T; fn deref(&self) -> &T { &self.value } } impl DerefMut for CachePadded { fn deref_mut(&mut self) -> &mut T { &mut self.value } } tokio-1.43.1/src/util/error.rs000064400000000000000000000015421046102023000143040ustar 00000000000000// Some combinations of features may not use these constants. #![cfg_attr(not(feature = "full"), allow(dead_code))] /// Error string explaining that the Tokio context hasn't been instantiated. pub(crate) const CONTEXT_MISSING_ERROR: &str = "there is no reactor running, must be called from the context of a Tokio 1.x runtime"; /// Error string explaining that the Tokio context is shutting down and cannot drive timers. pub(crate) const RUNTIME_SHUTTING_DOWN_ERROR: &str = "A Tokio 1.x context was found, but it is being shutdown."; /// Error string explaining that the Tokio context is not available because the /// thread-local storing it has been destroyed. This usually only happens during /// destructors of other thread-locals. pub(crate) const THREAD_LOCAL_DESTROYED_ERROR: &str = "The Tokio context thread-local variable has been destroyed."; tokio-1.43.1/src/util/idle_notified_set.rs000064400000000000000000000423511046102023000166270ustar 00000000000000//! This module defines an `IdleNotifiedSet`, which is a collection of elements. //! Each element is intended to correspond to a task, and the collection will //! keep track of which tasks have had their waker notified, and which have not. //! //! Each entry in the set holds some user-specified value. The value's type is //! specified using the `T` parameter. It will usually be a `JoinHandle` or //! similar. use std::marker::PhantomPinned; use std::mem::ManuallyDrop; use std::ptr::NonNull; use std::task::{Context, Waker}; use crate::loom::cell::UnsafeCell; use crate::loom::sync::{Arc, Mutex}; use crate::util::linked_list::{self, Link}; use crate::util::{waker_ref, Wake}; type LinkedList = linked_list::LinkedList, as linked_list::Link>::Target>; /// This is the main handle to the collection. pub(crate) struct IdleNotifiedSet { lists: Arc>, length: usize, } /// A handle to an entry that is guaranteed to be stored in the idle or notified /// list of its `IdleNotifiedSet`. This value borrows the `IdleNotifiedSet` /// mutably to prevent the entry from being moved to the `Neither` list, which /// only the `IdleNotifiedSet` may do. /// /// The main consequence of being stored in one of the lists is that the `value` /// field has not yet been consumed. /// /// Note: This entry can be moved from the idle to the notified list while this /// object exists by waking its waker. pub(crate) struct EntryInOneOfTheLists<'a, T> { entry: Arc>, set: &'a mut IdleNotifiedSet, } type Lists = Mutex>; /// The linked lists hold strong references to the `ListEntry` items, and the /// `ListEntry` items also hold a strong reference back to the Lists object, but /// the destructor of the `IdleNotifiedSet` will clear the two lists, so once /// that object is destroyed, no ref-cycles will remain. struct ListsInner { notified: LinkedList, idle: LinkedList, /// Whenever an element in the `notified` list is woken, this waker will be /// notified and consumed, if it exists. waker: Option, } /// Which of the two lists in the shared Lists object is this entry stored in? /// /// If the value is `Idle`, then an entry's waker may move it to the notified /// list. Otherwise, only the `IdleNotifiedSet` may move it. /// /// If the value is `Neither`, then it is still possible that the entry is in /// some third external list (this happens in `drain`). #[derive(Copy, Clone, Eq, PartialEq)] enum List { Notified, Idle, Neither, } /// An entry in the list. /// /// # Safety /// /// The `my_list` field must only be accessed while holding the mutex in /// `parent`. It is an invariant that the value of `my_list` corresponds to /// which linked list in the `parent` holds this entry. Once this field takes /// the value `Neither`, then it may never be modified again. /// /// If the value of `my_list` is `Notified` or `Idle`, then the `pointers` field /// must only be accessed while holding the mutex. If the value of `my_list` is /// `Neither`, then the `pointers` field may be accessed by the /// `IdleNotifiedSet` (this happens inside `drain`). /// /// The `value` field is owned by the `IdleNotifiedSet` and may only be accessed /// by the `IdleNotifiedSet`. The operation that sets the value of `my_list` to /// `Neither` assumes ownership of the `value`, and it must either drop it or /// move it out from this entry to prevent it from getting leaked. (Since the /// two linked lists are emptied in the destructor of `IdleNotifiedSet`, the /// value should not be leaked.) struct ListEntry { /// The linked list pointers of the list this entry is in. pointers: linked_list::Pointers>, /// Pointer to the shared `Lists` struct. parent: Arc>, /// The value stored in this entry. value: UnsafeCell>, /// Used to remember which list this entry is in. my_list: UnsafeCell, /// Required by the `linked_list::Pointers` field. _pin: PhantomPinned, } generate_addr_of_methods! { impl ListEntry { unsafe fn addr_of_pointers(self: NonNull) -> NonNull>> { &self.pointers } } } // With mutable access to the `IdleNotifiedSet`, you can get mutable access to // the values. unsafe impl Send for IdleNotifiedSet {} // With the current API we strictly speaking don't even need `T: Sync`, but we // require it anyway to support adding &self APIs that access the values in the // future. unsafe impl Sync for IdleNotifiedSet {} // These impls control when it is safe to create a Waker. Since the waker does // not allow access to the value in any way (including its destructor), it is // not necessary for `T` to be Send or Sync. unsafe impl Send for ListEntry {} unsafe impl Sync for ListEntry {} impl IdleNotifiedSet { /// Create a new `IdleNotifiedSet`. pub(crate) fn new() -> Self { let lists = Mutex::new(ListsInner { notified: LinkedList::new(), idle: LinkedList::new(), waker: None, }); IdleNotifiedSet { lists: Arc::new(lists), length: 0, } } pub(crate) fn len(&self) -> usize { self.length } pub(crate) fn is_empty(&self) -> bool { self.length == 0 } /// Insert the given value into the `idle` list. pub(crate) fn insert_idle(&mut self, value: T) -> EntryInOneOfTheLists<'_, T> { self.length += 1; let entry = Arc::new(ListEntry { parent: self.lists.clone(), value: UnsafeCell::new(ManuallyDrop::new(value)), my_list: UnsafeCell::new(List::Idle), pointers: linked_list::Pointers::new(), _pin: PhantomPinned, }); { let mut lock = self.lists.lock(); lock.idle.push_front(entry.clone()); } // Safety: We just put the entry in the idle list, so it is in one of the lists. EntryInOneOfTheLists { entry, set: self } } /// Pop an entry from the notified list to poll it. The entry is moved to /// the idle list atomically. pub(crate) fn pop_notified(&mut self, waker: &Waker) -> Option> { // We don't decrement the length because this call moves the entry to // the idle list rather than removing it. if self.length == 0 { // Fast path. return None; } let mut lock = self.lists.lock(); let should_update_waker = match lock.waker.as_mut() { Some(cur_waker) => !waker.will_wake(cur_waker), None => true, }; if should_update_waker { lock.waker = Some(waker.clone()); } // Pop the entry, returning None if empty. let entry = lock.notified.pop_back()?; lock.idle.push_front(entry.clone()); // Safety: We are holding the lock. entry.my_list.with_mut(|ptr| unsafe { *ptr = List::Idle; }); drop(lock); // Safety: We just put the entry in the idle list, so it is in one of the lists. Some(EntryInOneOfTheLists { entry, set: self }) } /// Tries to pop an entry from the notified list to poll it. The entry is moved to /// the idle list atomically. pub(crate) fn try_pop_notified(&mut self) -> Option> { // We don't decrement the length because this call moves the entry to // the idle list rather than removing it. if self.length == 0 { // Fast path. return None; } let mut lock = self.lists.lock(); // Pop the entry, returning None if empty. let entry = lock.notified.pop_back()?; lock.idle.push_front(entry.clone()); // Safety: We are holding the lock. entry.my_list.with_mut(|ptr| unsafe { *ptr = List::Idle; }); drop(lock); // Safety: We just put the entry in the idle list, so it is in one of the lists. Some(EntryInOneOfTheLists { entry, set: self }) } /// Call a function on every element in this list. pub(crate) fn for_each(&mut self, mut func: F) { fn get_ptrs(list: &mut LinkedList, ptrs: &mut Vec<*mut T>) { let mut node = list.last(); while let Some(entry) = node { ptrs.push(entry.value.with_mut(|ptr| { let ptr: *mut ManuallyDrop = ptr; let ptr: *mut T = ptr.cast(); ptr })); let prev = entry.pointers.get_prev(); node = prev.map(|prev| unsafe { &*prev.as_ptr() }); } } // Atomically get a raw pointer to the value of every entry. // // Since this only locks the mutex once, it is not possible for a value // to get moved from the idle list to the notified list during the // operation, which would otherwise result in some value being listed // twice. let mut ptrs = Vec::with_capacity(self.len()); { let mut lock = self.lists.lock(); get_ptrs(&mut lock.idle, &mut ptrs); get_ptrs(&mut lock.notified, &mut ptrs); } debug_assert_eq!(ptrs.len(), ptrs.capacity()); for ptr in ptrs { // Safety: When we grabbed the pointers, the entries were in one of // the two lists. This means that their value was valid at the time, // and it must still be valid because we are the IdleNotifiedSet, // and only we can remove an entry from the two lists. (It's // possible that an entry is moved from one list to the other during // this loop, but that is ok.) func(unsafe { &mut *ptr }); } } /// Remove all entries in both lists, applying some function to each element. /// /// The closure is called on all elements even if it panics. Having it panic /// twice is a double-panic, and will abort the application. pub(crate) fn drain(&mut self, func: F) { if self.length == 0 { // Fast path. return; } self.length = 0; // The LinkedList is not cleared on panic, so we use a bomb to clear it. // // This value has the invariant that any entry in its `all_entries` list // has `my_list` set to `Neither` and that the value has not yet been // dropped. struct AllEntries { all_entries: LinkedList, func: F, } impl AllEntries { fn pop_next(&mut self) -> bool { if let Some(entry) = self.all_entries.pop_back() { // Safety: We just took this value from the list, so we can // destroy the value in the entry. entry .value .with_mut(|ptr| unsafe { (self.func)(ManuallyDrop::take(&mut *ptr)) }); true } else { false } } } impl Drop for AllEntries { fn drop(&mut self) { while self.pop_next() {} } } let mut all_entries = AllEntries { all_entries: LinkedList::new(), func, }; // Atomically move all entries to the new linked list in the AllEntries // object. { let mut lock = self.lists.lock(); unsafe { // Safety: We are holding the lock and `all_entries` is a new // LinkedList. move_to_new_list(&mut lock.idle, &mut all_entries.all_entries); move_to_new_list(&mut lock.notified, &mut all_entries.all_entries); } } // Keep destroying entries in the list until it is empty. // // If the closure panics, then the destructor of the `AllEntries` bomb // ensures that we keep running the destructor on the remaining values. // A second panic will abort the program. while all_entries.pop_next() {} } } /// # Safety /// /// The mutex for the entries must be held, and the target list must be such /// that setting `my_list` to `Neither` is ok. unsafe fn move_to_new_list(from: &mut LinkedList, to: &mut LinkedList) { while let Some(entry) = from.pop_back() { entry.my_list.with_mut(|ptr| { *ptr = List::Neither; }); to.push_front(entry); } } impl<'a, T> EntryInOneOfTheLists<'a, T> { /// Remove this entry from the list it is in, returning the value associated /// with the entry. /// /// This consumes the value, since it is no longer guaranteed to be in a /// list. pub(crate) fn remove(self) -> T { self.set.length -= 1; { let mut lock = self.set.lists.lock(); // Safety: We are holding the lock so there is no race, and we will // remove the entry afterwards to uphold invariants. let old_my_list = self.entry.my_list.with_mut(|ptr| unsafe { let old_my_list = *ptr; *ptr = List::Neither; old_my_list }); let list = match old_my_list { List::Idle => &mut lock.idle, List::Notified => &mut lock.notified, // An entry in one of the lists is in one of the lists. List::Neither => unreachable!(), }; unsafe { // Safety: We just checked that the entry is in this particular // list. list.remove(ListEntry::as_raw(&self.entry)).unwrap(); } } // By setting `my_list` to `Neither`, we have taken ownership of the // value. We return it to the caller. // // Safety: We have a mutable reference to the `IdleNotifiedSet` that // owns this entry, so we can use its permission to access the value. self.entry .value .with_mut(|ptr| unsafe { ManuallyDrop::take(&mut *ptr) }) } /// Access the value in this entry together with a context for its waker. pub(crate) fn with_value_and_context(&mut self, func: F) -> U where F: FnOnce(&mut T, &mut Context<'_>) -> U, T: 'static, { let waker = waker_ref(&self.entry); let mut context = Context::from_waker(&waker); // Safety: We have a mutable reference to the `IdleNotifiedSet` that // owns this entry, so we can use its permission to access the value. self.entry .value .with_mut(|ptr| unsafe { func(&mut *ptr, &mut context) }) } } impl Drop for IdleNotifiedSet { fn drop(&mut self) { // Clear both lists. self.drain(drop); #[cfg(debug_assertions)] if !std::thread::panicking() { let lock = self.lists.lock(); assert!(lock.idle.is_empty()); assert!(lock.notified.is_empty()); } } } impl Wake for ListEntry { fn wake_by_ref(me: &Arc) { let mut lock = me.parent.lock(); // Safety: We are holding the lock and we will update the lists to // maintain invariants. let old_my_list = me.my_list.with_mut(|ptr| unsafe { let old_my_list = *ptr; if old_my_list == List::Idle { *ptr = List::Notified; } old_my_list }); if old_my_list == List::Idle { // We move ourself to the notified list. let me = unsafe { // Safety: We just checked that we are in this particular list. lock.idle.remove(ListEntry::as_raw(me)).unwrap() }; lock.notified.push_front(me); if let Some(waker) = lock.waker.take() { drop(lock); waker.wake(); } } } fn wake(me: Arc) { Self::wake_by_ref(&me); } } /// # Safety /// /// `ListEntry` is forced to be !Unpin. unsafe impl linked_list::Link for ListEntry { type Handle = Arc>; type Target = ListEntry; fn as_raw(handle: &Self::Handle) -> NonNull> { let ptr: *const ListEntry = Arc::as_ptr(handle); // Safety: We can't get a null pointer from `Arc::as_ptr`. unsafe { NonNull::new_unchecked(ptr as *mut ListEntry) } } unsafe fn from_raw(ptr: NonNull>) -> Arc> { Arc::from_raw(ptr.as_ptr()) } unsafe fn pointers( target: NonNull>, ) -> NonNull>> { ListEntry::addr_of_pointers(target) } } #[cfg(all(test, not(loom)))] mod tests { use crate::runtime::Builder; use crate::task::JoinSet; // A test that runs under miri. // // https://github.com/tokio-rs/tokio/pull/5693 #[test] fn join_set_test() { let rt = Builder::new_current_thread().build().unwrap(); let mut set = JoinSet::new(); set.spawn_on(futures::future::ready(()), rt.handle()); rt.block_on(set.join_next()).unwrap().unwrap(); } } tokio-1.43.1/src/util/linked_list.rs000064400000000000000000000601301046102023000154520ustar 00000000000000#![cfg_attr(not(feature = "full"), allow(dead_code))] //! An intrusive double linked list of data. //! //! The data structure supports tracking pinned nodes. Most of the data //! structure's APIs are `unsafe` as they require the caller to ensure the //! specified node is actually contained by the list. use core::cell::UnsafeCell; use core::fmt; use core::marker::{PhantomData, PhantomPinned}; use core::mem::ManuallyDrop; use core::ptr::{self, NonNull}; /// An intrusive linked list. /// /// Currently, the list is not emptied on drop. It is the caller's /// responsibility to ensure the list is empty before dropping it. pub(crate) struct LinkedList { /// Linked list head head: Option>, /// Linked list tail tail: Option>, /// Node type marker. _marker: PhantomData<*const L>, } unsafe impl Send for LinkedList where L::Target: Send {} unsafe impl Sync for LinkedList where L::Target: Sync {} /// Defines how a type is tracked within a linked list. /// /// In order to support storing a single type within multiple lists, accessing /// the list pointers is decoupled from the entry type. /// /// # Safety /// /// Implementations must guarantee that `Target` types are pinned in memory. In /// other words, when a node is inserted, the value will not be moved as long as /// it is stored in the list. pub(crate) unsafe trait Link { /// Handle to the list entry. /// /// This is usually a pointer-ish type. type Handle; /// Node type. type Target; /// Convert the handle to a raw pointer without consuming the handle. #[allow(clippy::wrong_self_convention)] fn as_raw(handle: &Self::Handle) -> NonNull; /// Convert the raw pointer to a handle unsafe fn from_raw(ptr: NonNull) -> Self::Handle; /// Return the pointers for a node /// /// # Safety /// /// The resulting pointer should have the same tag in the stacked-borrows /// stack as the argument. In particular, the method may not create an /// intermediate reference in the process of creating the resulting raw /// pointer. unsafe fn pointers(target: NonNull) -> NonNull>; } /// Previous / next pointers. pub(crate) struct Pointers { inner: UnsafeCell>, } /// We do not want the compiler to put the `noalias` attribute on mutable /// references to this type, so the type has been made `!Unpin` with a /// `PhantomPinned` field. /// /// Additionally, we never access the `prev` or `next` fields directly, as any /// such access would implicitly involve the creation of a reference to the /// field, which we want to avoid since the fields are not `!Unpin`, and would /// hence be given the `noalias` attribute if we were to do such an access. As /// an alternative to accessing the fields directly, the `Pointers` type /// provides getters and setters for the two fields, and those are implemented /// using `ptr`-specific methods which avoids the creation of intermediate /// references. /// /// See this link for more information: /// struct PointersInner { /// The previous node in the list. null if there is no previous node. prev: Option>, /// The next node in the list. null if there is no previous node. next: Option>, /// This type is !Unpin due to the heuristic from: /// _pin: PhantomPinned, } unsafe impl Send for Pointers {} unsafe impl Sync for Pointers {} // ===== impl LinkedList ===== impl LinkedList { /// Creates an empty linked list. pub(crate) const fn new() -> LinkedList { LinkedList { head: None, tail: None, _marker: PhantomData, } } } impl LinkedList { /// Adds an element first in the list. pub(crate) fn push_front(&mut self, val: L::Handle) { // The value should not be dropped, it is being inserted into the list let val = ManuallyDrop::new(val); let ptr = L::as_raw(&val); assert_ne!(self.head, Some(ptr)); unsafe { L::pointers(ptr).as_mut().set_next(self.head); L::pointers(ptr).as_mut().set_prev(None); if let Some(head) = self.head { L::pointers(head).as_mut().set_prev(Some(ptr)); } self.head = Some(ptr); if self.tail.is_none() { self.tail = Some(ptr); } } } /// Removes the first element from a list and returns it, or None if it is /// empty. pub(crate) fn pop_front(&mut self) -> Option { unsafe { let head = self.head?; self.head = L::pointers(head).as_ref().get_next(); if let Some(new_head) = L::pointers(head).as_ref().get_next() { L::pointers(new_head).as_mut().set_prev(None); } else { self.tail = None; } L::pointers(head).as_mut().set_prev(None); L::pointers(head).as_mut().set_next(None); Some(L::from_raw(head)) } } /// Removes the last element from a list and returns it, or None if it is /// empty. pub(crate) fn pop_back(&mut self) -> Option { unsafe { let last = self.tail?; self.tail = L::pointers(last).as_ref().get_prev(); if let Some(prev) = L::pointers(last).as_ref().get_prev() { L::pointers(prev).as_mut().set_next(None); } else { self.head = None; } L::pointers(last).as_mut().set_prev(None); L::pointers(last).as_mut().set_next(None); Some(L::from_raw(last)) } } /// Returns whether the linked list does not contain any node pub(crate) fn is_empty(&self) -> bool { if self.head.is_some() { return false; } assert!(self.tail.is_none()); true } /// Removes the specified node from the list /// /// # Safety /// /// The caller **must** ensure that exactly one of the following is true: /// - `node` is currently contained by `self`, /// - `node` is not contained by any list, /// - `node` is currently contained by some other `GuardedLinkedList` **and** /// the caller has an exclusive access to that list. This condition is /// used by the linked list in `sync::Notify`. pub(crate) unsafe fn remove(&mut self, node: NonNull) -> Option { if let Some(prev) = L::pointers(node).as_ref().get_prev() { debug_assert_eq!(L::pointers(prev).as_ref().get_next(), Some(node)); L::pointers(prev) .as_mut() .set_next(L::pointers(node).as_ref().get_next()); } else { if self.head != Some(node) { return None; } self.head = L::pointers(node).as_ref().get_next(); } if let Some(next) = L::pointers(node).as_ref().get_next() { debug_assert_eq!(L::pointers(next).as_ref().get_prev(), Some(node)); L::pointers(next) .as_mut() .set_prev(L::pointers(node).as_ref().get_prev()); } else { // This might be the last item in the list if self.tail != Some(node) { return None; } self.tail = L::pointers(node).as_ref().get_prev(); } L::pointers(node).as_mut().set_next(None); L::pointers(node).as_mut().set_prev(None); Some(L::from_raw(node)) } } impl fmt::Debug for LinkedList { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("LinkedList") .field("head", &self.head) .field("tail", &self.tail) .finish() } } #[cfg(any( feature = "fs", feature = "rt", all(unix, feature = "process"), feature = "signal", feature = "sync", ))] impl LinkedList { pub(crate) fn last(&self) -> Option<&L::Target> { let tail = self.tail.as_ref()?; unsafe { Some(&*tail.as_ptr()) } } } impl Default for LinkedList { fn default() -> Self { Self::new() } } // ===== impl DrainFilter ===== cfg_io_driver_impl! { pub(crate) struct DrainFilter<'a, T: Link, F> { list: &'a mut LinkedList, filter: F, curr: Option>, } impl LinkedList { pub(crate) fn drain_filter(&mut self, filter: F) -> DrainFilter<'_, T, F> where F: FnMut(&T::Target) -> bool, { let curr = self.head; DrainFilter { curr, filter, list: self, } } } impl<'a, T, F> Iterator for DrainFilter<'a, T, F> where T: Link, F: FnMut(&T::Target) -> bool, { type Item = T::Handle; fn next(&mut self) -> Option { while let Some(curr) = self.curr { // safety: the pointer references data contained by the list self.curr = unsafe { T::pointers(curr).as_ref() }.get_next(); // safety: the value is still owned by the linked list. if (self.filter)(unsafe { &mut *curr.as_ptr() }) { return unsafe { self.list.remove(curr) }; } } None } } } cfg_taskdump! { impl LinkedList { pub(crate) fn for_each(&mut self, mut f: F) where F: FnMut(&T::Handle), { let mut next = self.head; while let Some(curr) = next { unsafe { let handle = ManuallyDrop::new(T::from_raw(curr)); f(&handle); next = T::pointers(curr).as_ref().get_next(); } } } } } // ===== impl GuardedLinkedList ===== feature! { #![any( feature = "process", feature = "sync", feature = "rt", feature = "signal", )] /// An intrusive linked list, but instead of keeping pointers to the head /// and tail nodes, it uses a special guard node linked with those nodes. /// It means that the list is circular and every pointer of a node from /// the list is not `None`, including pointers from the guard node. /// /// If a list is empty, then both pointers of the guard node are pointing /// at the guard node itself. pub(crate) struct GuardedLinkedList { /// Pointer to the guard node. guard: NonNull, /// Node type marker. _marker: PhantomData<*const L>, } impl LinkedList { /// Turns a linked list into the guarded version by linking the guard node /// with the head and tail nodes. Like with other nodes, you should guarantee /// that the guard node is pinned in memory. pub(crate) fn into_guarded(self, guard_handle: L::Handle) -> GuardedLinkedList { // `guard_handle` is a NonNull pointer, we don't have to care about dropping it. let guard = L::as_raw(&guard_handle); unsafe { if let Some(head) = self.head { debug_assert!(L::pointers(head).as_ref().get_prev().is_none()); L::pointers(head).as_mut().set_prev(Some(guard)); L::pointers(guard).as_mut().set_next(Some(head)); // The list is not empty, so the tail cannot be `None`. let tail = self.tail.unwrap(); debug_assert!(L::pointers(tail).as_ref().get_next().is_none()); L::pointers(tail).as_mut().set_next(Some(guard)); L::pointers(guard).as_mut().set_prev(Some(tail)); } else { // The list is empty. L::pointers(guard).as_mut().set_prev(Some(guard)); L::pointers(guard).as_mut().set_next(Some(guard)); } } GuardedLinkedList { guard, _marker: PhantomData } } } impl GuardedLinkedList { fn tail(&self) -> Option> { let tail_ptr = unsafe { L::pointers(self.guard).as_ref().get_prev().unwrap() }; // Compare the tail pointer with the address of the guard node itself. // If the guard points at itself, then there are no other nodes and // the list is considered empty. if tail_ptr != self.guard { Some(tail_ptr) } else { None } } /// Removes the last element from a list and returns it, or None if it is /// empty. pub(crate) fn pop_back(&mut self) -> Option { unsafe { let last = self.tail()?; let before_last = L::pointers(last).as_ref().get_prev().unwrap(); L::pointers(self.guard).as_mut().set_prev(Some(before_last)); L::pointers(before_last).as_mut().set_next(Some(self.guard)); L::pointers(last).as_mut().set_prev(None); L::pointers(last).as_mut().set_next(None); Some(L::from_raw(last)) } } } } // ===== impl Pointers ===== impl Pointers { /// Create a new set of empty pointers pub(crate) fn new() -> Pointers { Pointers { inner: UnsafeCell::new(PointersInner { prev: None, next: None, _pin: PhantomPinned, }), } } pub(crate) fn get_prev(&self) -> Option> { // SAFETY: Field is accessed immutably through a reference. unsafe { ptr::addr_of!((*self.inner.get()).prev).read() } } pub(crate) fn get_next(&self) -> Option> { // SAFETY: Field is accessed immutably through a reference. unsafe { ptr::addr_of!((*self.inner.get()).next).read() } } fn set_prev(&mut self, value: Option>) { // SAFETY: Field is accessed mutably through a mutable reference. unsafe { ptr::addr_of_mut!((*self.inner.get()).prev).write(value); } } fn set_next(&mut self, value: Option>) { // SAFETY: Field is accessed mutably through a mutable reference. unsafe { ptr::addr_of_mut!((*self.inner.get()).next).write(value); } } } impl fmt::Debug for Pointers { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let prev = self.get_prev(); let next = self.get_next(); f.debug_struct("Pointers") .field("prev", &prev) .field("next", &next) .finish() } } #[cfg(any(test, fuzzing))] #[cfg(not(loom))] pub(crate) mod tests { use super::*; use std::pin::Pin; #[derive(Debug)] #[repr(C)] struct Entry { pointers: Pointers, val: i32, } unsafe impl<'a> Link for &'a Entry { type Handle = Pin<&'a Entry>; type Target = Entry; fn as_raw(handle: &Pin<&'_ Entry>) -> NonNull { NonNull::from(handle.get_ref()) } unsafe fn from_raw(ptr: NonNull) -> Pin<&'a Entry> { Pin::new_unchecked(&*ptr.as_ptr()) } unsafe fn pointers(target: NonNull) -> NonNull> { target.cast() } } fn entry(val: i32) -> Pin> { Box::pin(Entry { pointers: Pointers::new(), val, }) } fn ptr(r: &Pin>) -> NonNull { r.as_ref().get_ref().into() } fn collect_list(list: &mut LinkedList<&'_ Entry, <&'_ Entry as Link>::Target>) -> Vec { let mut ret = vec![]; while let Some(entry) = list.pop_back() { ret.push(entry.val); } ret } fn push_all<'a>( list: &mut LinkedList<&'a Entry, <&'_ Entry as Link>::Target>, entries: &[Pin<&'a Entry>], ) { for entry in entries.iter() { list.push_front(*entry); } } #[cfg(test)] macro_rules! assert_clean { ($e:ident) => {{ assert!($e.pointers.get_next().is_none()); assert!($e.pointers.get_prev().is_none()); }}; } #[cfg(test)] macro_rules! assert_ptr_eq { ($a:expr, $b:expr) => {{ // Deal with mapping a Pin<&mut T> -> Option> assert_eq!(Some($a.as_ref().get_ref().into()), $b) }}; } #[test] fn const_new() { const _: LinkedList<&Entry, <&Entry as Link>::Target> = LinkedList::new(); } #[test] fn push_and_drain() { let a = entry(5); let b = entry(7); let c = entry(31); let mut list = LinkedList::new(); assert!(list.is_empty()); list.push_front(a.as_ref()); assert!(!list.is_empty()); list.push_front(b.as_ref()); list.push_front(c.as_ref()); let items: Vec = collect_list(&mut list); assert_eq!([5, 7, 31].to_vec(), items); assert!(list.is_empty()); } #[test] fn push_pop_push_pop() { let a = entry(5); let b = entry(7); let mut list = LinkedList::<&Entry, <&Entry as Link>::Target>::new(); list.push_front(a.as_ref()); let entry = list.pop_back().unwrap(); assert_eq!(5, entry.val); assert!(list.is_empty()); list.push_front(b.as_ref()); let entry = list.pop_back().unwrap(); assert_eq!(7, entry.val); assert!(list.is_empty()); assert!(list.pop_back().is_none()); } #[test] fn remove_by_address() { let a = entry(5); let b = entry(7); let c = entry(31); unsafe { // Remove first let mut list = LinkedList::new(); push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); // `a` should be no longer there and can't be removed twice assert!(list.remove(ptr(&a)).is_none()); assert!(!list.is_empty()); assert!(list.remove(ptr(&b)).is_some()); assert_clean!(b); // `b` should be no longer there and can't be removed twice assert!(list.remove(ptr(&b)).is_none()); assert!(!list.is_empty()); assert!(list.remove(ptr(&c)).is_some()); assert_clean!(c); // `b` should be no longer there and can't be removed twice assert!(list.remove(ptr(&c)).is_none()); assert!(list.is_empty()); } unsafe { // Remove middle let mut list = LinkedList::new(); push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); assert_ptr_eq!(b, list.head); assert_ptr_eq!(c, b.pointers.get_next()); assert_ptr_eq!(b, c.pointers.get_prev()); let items = collect_list(&mut list); assert_eq!([31, 7].to_vec(), items); } unsafe { // Remove middle let mut list = LinkedList::new(); push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); assert!(list.remove(ptr(&b)).is_some()); assert_clean!(b); assert_ptr_eq!(c, a.pointers.get_next()); assert_ptr_eq!(a, c.pointers.get_prev()); let items = collect_list(&mut list); assert_eq!([31, 5].to_vec(), items); } unsafe { // Remove last // Remove middle let mut list = LinkedList::new(); push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); assert!(list.remove(ptr(&c)).is_some()); assert_clean!(c); assert!(b.pointers.get_next().is_none()); assert_ptr_eq!(b, list.tail); let items = collect_list(&mut list); assert_eq!([7, 5].to_vec(), items); } unsafe { // Remove first of two let mut list = LinkedList::new(); push_all(&mut list, &[b.as_ref(), a.as_ref()]); assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); // a should be no longer there and can't be removed twice assert!(list.remove(ptr(&a)).is_none()); assert_ptr_eq!(b, list.head); assert_ptr_eq!(b, list.tail); assert!(b.pointers.get_next().is_none()); assert!(b.pointers.get_prev().is_none()); let items = collect_list(&mut list); assert_eq!([7].to_vec(), items); } unsafe { // Remove last of two let mut list = LinkedList::new(); push_all(&mut list, &[b.as_ref(), a.as_ref()]); assert!(list.remove(ptr(&b)).is_some()); assert_clean!(b); assert_ptr_eq!(a, list.head); assert_ptr_eq!(a, list.tail); assert!(a.pointers.get_next().is_none()); assert!(a.pointers.get_prev().is_none()); let items = collect_list(&mut list); assert_eq!([5].to_vec(), items); } unsafe { // Remove last item let mut list = LinkedList::new(); push_all(&mut list, &[a.as_ref()]); assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); assert!(list.head.is_none()); assert!(list.tail.is_none()); let items = collect_list(&mut list); assert!(items.is_empty()); } unsafe { // Remove missing let mut list = LinkedList::<&Entry, <&Entry as Link>::Target>::new(); list.push_front(b.as_ref()); list.push_front(a.as_ref()); assert!(list.remove(ptr(&c)).is_none()); } } /// This is a fuzz test. You run it by entering `cargo fuzz run fuzz_linked_list` in CLI in `/tokio/` module. #[cfg(fuzzing)] pub fn fuzz_linked_list(ops: &[u8]) { enum Op { Push, Pop, Remove(usize), } use std::collections::VecDeque; let ops = ops .iter() .map(|i| match i % 3u8 { 0 => Op::Push, 1 => Op::Pop, 2 => Op::Remove((i / 3u8) as usize), _ => unreachable!(), }) .collect::>(); let mut ll = LinkedList::<&Entry, <&Entry as Link>::Target>::new(); let mut reference = VecDeque::new(); let entries: Vec<_> = (0..ops.len()).map(|i| entry(i as i32)).collect(); for (i, op) in ops.iter().enumerate() { match op { Op::Push => { reference.push_front(i as i32); assert_eq!(entries[i].val, i as i32); ll.push_front(entries[i].as_ref()); } Op::Pop => { if reference.is_empty() { assert!(ll.is_empty()); continue; } let v = reference.pop_back(); assert_eq!(v, ll.pop_back().map(|v| v.val)); } Op::Remove(n) => { if reference.is_empty() { assert!(ll.is_empty()); continue; } let idx = n % reference.len(); let expect = reference.remove(idx).unwrap(); unsafe { let entry = ll.remove(ptr(&entries[expect as usize])).unwrap(); assert_eq!(expect, entry.val); } } } } } } tokio-1.43.1/src/util/markers.rs000064400000000000000000000003701046102023000146150ustar 00000000000000/// Marker for types that are `Sync` but not `Send` #[allow(dead_code)] pub(crate) struct SyncNotSend(#[allow(dead_code)] *mut ()); unsafe impl Sync for SyncNotSend {} cfg_rt! { pub(crate) struct NotSendOrSync(#[allow(dead_code)] *mut ()); } tokio-1.43.1/src/util/memchr.rs000064400000000000000000000043521046102023000144300ustar 00000000000000//! Search for a byte in a byte array using libc. //! //! When nothing pulls in libc, then just use a trivial implementation. Note //! that we only depend on libc on unix. #[cfg(not(all(unix, feature = "libc")))] pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option { haystack.iter().position(|val| needle == *val) } #[cfg(all(unix, feature = "libc"))] pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option { let start = haystack.as_ptr(); // SAFETY: `start` is valid for `haystack.len()` bytes. let ptr = unsafe { libc::memchr(start.cast(), needle as _, haystack.len()) }; if ptr.is_null() { None } else { Some(ptr as usize - start as usize) } } #[cfg(test)] mod tests { use super::memchr; #[test] fn memchr_test() { let haystack = b"123abc456\0\xffabc\n"; assert_eq!(memchr(b'1', haystack), Some(0)); assert_eq!(memchr(b'2', haystack), Some(1)); assert_eq!(memchr(b'3', haystack), Some(2)); assert_eq!(memchr(b'4', haystack), Some(6)); assert_eq!(memchr(b'5', haystack), Some(7)); assert_eq!(memchr(b'6', haystack), Some(8)); assert_eq!(memchr(b'7', haystack), None); assert_eq!(memchr(b'a', haystack), Some(3)); assert_eq!(memchr(b'b', haystack), Some(4)); assert_eq!(memchr(b'c', haystack), Some(5)); assert_eq!(memchr(b'd', haystack), None); assert_eq!(memchr(b'A', haystack), None); assert_eq!(memchr(0, haystack), Some(9)); assert_eq!(memchr(0xff, haystack), Some(10)); assert_eq!(memchr(0xfe, haystack), None); assert_eq!(memchr(1, haystack), None); assert_eq!(memchr(b'\n', haystack), Some(14)); assert_eq!(memchr(b'\r', haystack), None); } #[test] fn memchr_all() { let mut arr = Vec::new(); for b in 0..=255 { arr.push(b); } for b in 0..=255 { assert_eq!(memchr(b, &arr), Some(b as usize)); } arr.reverse(); for b in 0..=255 { assert_eq!(memchr(b, &arr), Some(255 - b as usize)); } } #[test] fn memchr_empty() { for b in 0..=255 { assert_eq!(memchr(b, b""), None); } } } tokio-1.43.1/src/util/metric_atomics.rs000064400000000000000000000046111046102023000161550ustar 00000000000000use std::sync::atomic::{AtomicUsize, Ordering}; cfg_64bit_metrics! { use std::sync::atomic::AtomicU64; } /// `AtomicU64` that is is a no-op on platforms without 64-bit atomics /// /// When used on platforms without 64-bit atomics, writes to this are no-ops. /// The `load` method is only defined when 64-bit atomics are available. #[derive(Debug, Default)] pub(crate) struct MetricAtomicU64 { #[cfg(target_has_atomic = "64")] value: AtomicU64, } // some of these are currently only used behind cfg_unstable #[allow(dead_code)] impl MetricAtomicU64 { // Load is only defined when supported cfg_64bit_metrics! { pub(crate) fn load(&self, ordering: Ordering) -> u64 { self.value.load(ordering) } } cfg_64bit_metrics! { pub(crate) fn store(&self, val: u64, ordering: Ordering) { self.value.store(val, ordering) } pub(crate) fn new(value: u64) -> Self { Self { value: AtomicU64::new(value) } } pub(crate) fn add(&self, value: u64, ordering: Ordering) { self.value.fetch_add(value, ordering); } } cfg_no_64bit_metrics! { pub(crate) fn store(&self, _val: u64, _ordering: Ordering) { } // on platforms without 64-bit atomics, fetch-add returns unit pub(crate) fn add(&self, _value: u64, _ordering: Ordering) { } pub(crate) fn new(_value: u64) -> Self { Self { } } } } #[cfg_attr(not(all(tokio_unstable, feature = "rt")), allow(dead_code))] /// `AtomicUsize` for use in metrics. /// /// This exposes simplified APIs for use in metrics & uses `std::sync` instead of Loom to avoid polluting loom logs with metric information. #[derive(Debug, Default)] pub(crate) struct MetricAtomicUsize { value: AtomicUsize, } #[cfg_attr(not(all(tokio_unstable, feature = "rt")), allow(dead_code))] impl MetricAtomicUsize { pub(crate) fn new(value: usize) -> Self { Self { value: AtomicUsize::new(value), } } pub(crate) fn load(&self, ordering: Ordering) -> usize { self.value.load(ordering) } pub(crate) fn store(&self, val: usize, ordering: Ordering) { self.value.store(val, ordering) } pub(crate) fn increment(&self) -> usize { self.value.fetch_add(1, Ordering::Relaxed) } pub(crate) fn decrement(&self) -> usize { self.value.fetch_sub(1, Ordering::Relaxed) } } tokio-1.43.1/src/util/mod.rs000064400000000000000000000035631046102023000137370ustar 00000000000000cfg_io_driver! { pub(crate) mod bit; } #[cfg(feature = "rt")] pub(crate) mod atomic_cell; pub(crate) mod metric_atomics; #[cfg(any(feature = "rt", feature = "signal", feature = "process"))] pub(crate) mod once_cell; #[cfg(any( // io driver uses `WakeList` directly feature = "net", feature = "process", // `sync` enables `Notify` and `batch_semaphore`, which require `WakeList`. feature = "sync", // `fs` uses `batch_semaphore`, which requires `WakeList`. feature = "fs", // rt and signal use `Notify`, which requires `WakeList`. feature = "rt", feature = "signal", // time driver uses `WakeList` in `Handle::process_at_time`. feature = "time", ))] mod wake_list; #[cfg(any( feature = "net", feature = "process", feature = "sync", feature = "fs", feature = "rt", feature = "signal", feature = "time", ))] pub(crate) use wake_list::WakeList; #[cfg(any( feature = "fs", feature = "net", feature = "process", feature = "rt", feature = "sync", feature = "signal", feature = "time", ))] pub(crate) mod linked_list; cfg_rt! { pub(crate) mod sharded_list; } #[cfg(any(feature = "rt", feature = "macros", feature = "time"))] pub(crate) mod rand; cfg_rt! { mod idle_notified_set; pub(crate) use idle_notified_set::IdleNotifiedSet; pub(crate) use self::rand::RngSeedGenerator; mod wake; pub(crate) use wake::WakerRef; pub(crate) use wake::{waker_ref, Wake}; mod sync_wrapper; pub(crate) use sync_wrapper::SyncWrapper; mod rc_cell; pub(crate) use rc_cell::RcCell; } cfg_rt_multi_thread! { mod try_lock; pub(crate) use try_lock::TryLock; } pub(crate) mod trace; pub(crate) mod error; #[cfg(feature = "io-util")] pub(crate) mod memchr; pub(crate) mod markers; pub(crate) mod cacheline; cfg_io_driver_impl! { pub(crate) mod ptr_expose; } tokio-1.43.1/src/util/once_cell.rs000064400000000000000000000042221046102023000150740ustar 00000000000000#![allow(dead_code)] use std::cell::UnsafeCell; use std::mem::MaybeUninit; use std::sync::Once; pub(crate) struct OnceCell { once: Once, value: UnsafeCell>, } unsafe impl Send for OnceCell {} unsafe impl Sync for OnceCell {} impl OnceCell { pub(crate) const fn new() -> Self { Self { once: Once::new(), value: UnsafeCell::new(MaybeUninit::uninit()), } } /// Get the value inside this cell, initializing it using the provided /// function if necessary. /// /// If the `init` closure panics, then the `OnceCell` is poisoned and all /// future calls to `get` will panic. #[inline] pub(crate) fn get(&self, init: impl FnOnce() -> T) -> &T { if !self.once.is_completed() { self.do_init(init); } // Safety: The `std::sync::Once` guarantees that we can only reach this // line if a `call_once` closure has been run exactly once and without // panicking. Thus, the value is not uninitialized. // // There is also no race because the only `&self` method that modifies // `value` is `do_init`, but if the `call_once` closure is still // running, then no thread has gotten past the `call_once`. unsafe { &*(self.value.get() as *const T) } } #[cold] fn do_init(&self, init: impl FnOnce() -> T) { let value_ptr = self.value.get() as *mut T; self.once.call_once(|| { let set_to = init(); // Safety: The `std::sync::Once` guarantees that this initialization // will run at most once, and that no thread can get past the // `call_once` until it has run exactly once. Thus, we have // exclusive access to `value`. unsafe { std::ptr::write(value_ptr, set_to); } }); } } impl Drop for OnceCell { fn drop(&mut self) { if self.once.is_completed() { let value_ptr = self.value.get() as *mut T; unsafe { std::ptr::drop_in_place(value_ptr); } } } } tokio-1.43.1/src/util/ptr_expose.rs000064400000000000000000000045451046102023000153510ustar 00000000000000//! Utility for helping miri understand our exposed pointers. //! //! During normal execution, this module is equivalent to pointer casts. However, when running //! under miri, pointer casts are replaced with lookups in a hash map. This makes Tokio compatible //! with strict provenance when running under miri (which comes with a performance cost). use std::marker::PhantomData; #[cfg(miri)] use {crate::loom::sync::Mutex, std::collections::BTreeMap}; pub(crate) struct PtrExposeDomain { #[cfg(miri)] map: Mutex>, _phantom: PhantomData, } // SAFETY: Actually using the pointers is unsafe, so it's sound to transfer them across threads. unsafe impl Sync for PtrExposeDomain {} impl PtrExposeDomain { pub(crate) const fn new() -> Self { Self { #[cfg(miri)] map: Mutex::const_new(BTreeMap::new()), _phantom: PhantomData, } } #[inline] pub(crate) fn expose_provenance(&self, ptr: *const T) -> usize { #[cfg(miri)] { // FIXME: Use `pointer:addr` when it is stable. // SAFETY: Equivalent to `pointer::addr` which is safe. let addr: usize = unsafe { std::mem::transmute(ptr) }; self.map.lock().insert(addr, ptr); addr } #[cfg(not(miri))] { ptr as usize } } #[inline] #[allow(clippy::wrong_self_convention)] // mirrors std name pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T { #[cfg(miri)] { let maybe_ptr = self.map.lock().get(&addr).copied(); // SAFETY: Intentionally trigger a miri failure if the provenance we want is not // exposed. unsafe { maybe_ptr.unwrap_unchecked() } } #[cfg(not(miri))] { addr as *const T } } #[inline] pub(crate) fn unexpose_provenance(&self, _ptr: *const T) { #[cfg(miri)] { // SAFETY: Equivalent to `pointer::addr` which is safe. let addr: usize = unsafe { std::mem::transmute(_ptr) }; let maybe_ptr = self.map.lock().remove(&addr); // SAFETY: Intentionally trigger a miri failure if the provenance we want is not // exposed. unsafe { maybe_ptr.unwrap_unchecked() }; } } } tokio-1.43.1/src/util/rand/rt.rs000064400000000000000000000037051046102023000145270ustar 00000000000000use super::{FastRand, RngSeed}; use std::sync::Mutex; /// A deterministic generator for seeds (and other generators). /// /// Given the same initial seed, the generator will output the same sequence of seeds. /// /// Since the seed generator will be kept in a runtime handle, we need to wrap `FastRand` /// in a Mutex to make it thread safe. Different to the `FastRand` that we keep in a /// thread local store, the expectation is that seed generation will not need to happen /// very frequently, so the cost of the mutex should be minimal. #[derive(Debug)] pub(crate) struct RngSeedGenerator { /// Internal state for the seed generator. We keep it in a Mutex so that we can safely /// use it across multiple threads. state: Mutex, } impl RngSeedGenerator { /// Returns a new generator from the provided seed. pub(crate) fn new(seed: RngSeed) -> Self { Self { state: Mutex::new(FastRand::from_seed(seed)), } } /// Returns the next seed in the sequence. pub(crate) fn next_seed(&self) -> RngSeed { let mut rng = self .state .lock() .expect("RNG seed generator is internally corrupt"); let s = rng.fastrand(); let r = rng.fastrand(); RngSeed::from_pair(s, r) } /// Directly creates a generator using the next seed. pub(crate) fn next_generator(&self) -> Self { RngSeedGenerator::new(self.next_seed()) } } impl FastRand { /// Replaces the state of the random number generator with the provided seed, returning /// the seed that represents the previous state of the random number generator. /// /// The random number generator will become equivalent to one created with /// the same seed. pub(crate) fn replace_seed(&mut self, seed: RngSeed) -> RngSeed { let old_seed = RngSeed::from_pair(self.one, self.two); self.one = seed.s; self.two = seed.r; old_seed } } tokio-1.43.1/src/util/rand/rt_unstable.rs000064400000000000000000000007541046102023000164250ustar 00000000000000use super::RngSeed; use std::collections::hash_map::DefaultHasher; use std::hash::Hasher; impl RngSeed { /// Generates a seed from the provided byte slice. /// /// # Example /// /// ``` /// # use tokio::runtime::RngSeed; /// let seed = RngSeed::from_bytes(b"make me a seed"); /// ``` pub fn from_bytes(bytes: &[u8]) -> Self { let mut hasher = DefaultHasher::default(); hasher.write(bytes); Self::from_u64(hasher.finish()) } } tokio-1.43.1/src/util/rand.rs000064400000000000000000000046571046102023000141110ustar 00000000000000cfg_rt! { mod rt; pub(crate) use rt::RngSeedGenerator; cfg_unstable! { mod rt_unstable; } } /// A seed for random number generation. /// /// In order to make certain functions within a runtime deterministic, a seed /// can be specified at the time of creation. #[allow(unreachable_pub)] #[derive(Clone, Debug)] pub struct RngSeed { s: u32, r: u32, } /// Fast random number generate. /// /// Implement `xorshift64+`: 2 32-bit `xorshift` sequences added together. /// Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's /// `Xorshift` paper: /// This generator passes the SmallCrush suite, part of TestU01 framework: /// #[derive(Clone, Copy, Debug)] pub(crate) struct FastRand { one: u32, two: u32, } impl RngSeed { /// Creates a random seed using loom internally. pub(crate) fn new() -> Self { Self::from_u64(crate::loom::rand::seed()) } fn from_u64(seed: u64) -> Self { let one = (seed >> 32) as u32; let mut two = seed as u32; if two == 0 { // This value cannot be zero two = 1; } Self::from_pair(one, two) } fn from_pair(s: u32, r: u32) -> Self { Self { s, r } } } impl FastRand { /// Initialize a new fast random number generator using the default source of entropy. pub(crate) fn new() -> FastRand { FastRand::from_seed(RngSeed::new()) } /// Initializes a new, thread-local, fast random number generator. pub(crate) fn from_seed(seed: RngSeed) -> FastRand { FastRand { one: seed.s, two: seed.r, } } #[cfg(any( feature = "macros", feature = "rt-multi-thread", feature = "time", all(feature = "sync", feature = "rt") ))] pub(crate) fn fastrand_n(&mut self, n: u32) -> u32 { // This is similar to fastrand() % n, but faster. // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/ let mul = (self.fastrand() as u64).wrapping_mul(n as u64); (mul >> 32) as u32 } fn fastrand(&mut self) -> u32 { let mut s1 = self.one; let s0 = self.two; s1 ^= s1 << 17; s1 = s1 ^ s0 ^ s1 >> 7 ^ s0 >> 16; self.one = s0; self.two = s1; s0.wrapping_add(s1) } } tokio-1.43.1/src/util/rc_cell.rs000064400000000000000000000035471046102023000145650ustar 00000000000000use crate::loom::cell::UnsafeCell; use std::rc::Rc; /// This is exactly like `Cell>>`, except that it provides a `get` /// method even though `Rc` is not `Copy`. pub(crate) struct RcCell { inner: UnsafeCell>>, } impl RcCell { #[cfg(not(all(loom, test)))] pub(crate) const fn new() -> Self { Self { inner: UnsafeCell::new(None), } } // The UnsafeCell in loom does not have a const `new` fn. #[cfg(all(loom, test))] pub(crate) fn new() -> Self { Self { inner: UnsafeCell::new(None), } } /// Safety: This method may not be called recursively. #[inline] unsafe fn with_inner(&self, f: F) -> R where F: FnOnce(&mut Option>) -> R, { // safety: This type is not Sync, so concurrent calls of this method // cannot happen. Furthermore, the caller guarantees that the method is // not called recursively. Finally, this is the only place that can // create mutable references to the inner Rc. This ensures that any // mutable references created here are exclusive. self.inner.with_mut(|ptr| f(&mut *ptr)) } pub(crate) fn get(&self) -> Option> { // safety: The `Rc::clone` method will not call any unknown user-code, // so it will not result in a recursive call to `with_inner`. unsafe { self.with_inner(|rc| rc.clone()) } } pub(crate) fn replace(&self, val: Option>) -> Option> { // safety: No destructors or other unknown user-code will run inside the // `with_inner` call, so no recursive call to `with_inner` can happen. unsafe { self.with_inner(|rc| std::mem::replace(rc, val)) } } pub(crate) fn set(&self, val: Option>) { let old = self.replace(val); drop(old); } } tokio-1.43.1/src/util/sharded_list.rs000064400000000000000000000122211046102023000156140ustar 00000000000000use std::ptr::NonNull; use std::sync::atomic::Ordering; use crate::loom::sync::{Mutex, MutexGuard}; use crate::util::metric_atomics::{MetricAtomicU64, MetricAtomicUsize}; use super::linked_list::{Link, LinkedList}; /// An intrusive linked list supporting highly concurrent updates. /// /// It currently relies on `LinkedList`, so it is the caller's /// responsibility to ensure the list is empty before dropping it. /// /// Note: Due to its inner sharded design, the order of nodes cannot be guaranteed. pub(crate) struct ShardedList { lists: Box<[Mutex>]>, added: MetricAtomicU64, count: MetricAtomicUsize, shard_mask: usize, } /// Determines which linked list an item should be stored in. /// /// # Safety /// /// Implementations must guarantee that the id of an item does not change from /// call to call. pub(crate) unsafe trait ShardedListItem: Link { /// # Safety /// The provided pointer must point at a valid list item. unsafe fn get_shard_id(target: NonNull) -> usize; } impl ShardedList { /// Creates a new and empty sharded linked list with the specified size. pub(crate) fn new(sharded_size: usize) -> Self { assert!(sharded_size.is_power_of_two()); let shard_mask = sharded_size - 1; let mut lists = Vec::with_capacity(sharded_size); for _ in 0..sharded_size { lists.push(Mutex::new(LinkedList::::new())) } Self { lists: lists.into_boxed_slice(), added: MetricAtomicU64::new(0), count: MetricAtomicUsize::new(0), shard_mask, } } } /// Used to get the lock of shard. pub(crate) struct ShardGuard<'a, L, T> { lock: MutexGuard<'a, LinkedList>, added: &'a MetricAtomicU64, count: &'a MetricAtomicUsize, id: usize, } impl ShardedList { /// Removes the last element from a list specified by `shard_id` and returns it, or None if it is /// empty. pub(crate) fn pop_back(&self, shard_id: usize) -> Option { let mut lock = self.shard_inner(shard_id); let node = lock.pop_back(); if node.is_some() { self.count.decrement(); } node } /// Removes the specified node from the list. /// /// # Safety /// /// The caller **must** ensure that exactly one of the following is true: /// - `node` is currently contained by `self`, /// - `node` is not contained by any list, /// - `node` is currently contained by some other `GuardedLinkedList`. pub(crate) unsafe fn remove(&self, node: NonNull) -> Option { let id = L::get_shard_id(node); let mut lock = self.shard_inner(id); // SAFETY: Since the shard id cannot change, it's not possible for this node // to be in any other list of the same sharded list. let node = unsafe { lock.remove(node) }; if node.is_some() { self.count.decrement(); } node } /// Gets the lock of `ShardedList`, makes us have the write permission. pub(crate) fn lock_shard(&self, val: &L::Handle) -> ShardGuard<'_, L, L::Target> { let id = unsafe { L::get_shard_id(L::as_raw(val)) }; ShardGuard { lock: self.shard_inner(id), added: &self.added, count: &self.count, id, } } /// Gets the count of elements in this list. pub(crate) fn len(&self) -> usize { self.count.load(Ordering::Relaxed) } cfg_64bit_metrics! { /// Gets the total number of elements added to this list. pub(crate) fn added(&self) -> u64 { self.added.load(Ordering::Relaxed) } } /// Returns whether the linked list does not contain any node. pub(crate) fn is_empty(&self) -> bool { self.len() == 0 } /// Gets the shard size of this `SharedList`. /// /// Used to help us to decide the parameter `shard_id` of the `pop_back` method. pub(crate) fn shard_size(&self) -> usize { self.shard_mask + 1 } #[inline] fn shard_inner(&self, id: usize) -> MutexGuard<'_, LinkedList::Target>> { // Safety: This modulo operation ensures that the index is not out of bounds. unsafe { self.lists.get_unchecked(id & self.shard_mask).lock() } } } impl<'a, L: ShardedListItem> ShardGuard<'a, L, L::Target> { /// Push a value to this shard. pub(crate) fn push(mut self, val: L::Handle) { let id = unsafe { L::get_shard_id(L::as_raw(&val)) }; assert_eq!(id, self.id); self.lock.push_front(val); self.added.add(1, Ordering::Relaxed); self.count.increment(); } } cfg_taskdump! { impl ShardedList { pub(crate) fn for_each(&self, mut f: F) where F: FnMut(&L::Handle), { let mut guards = Vec::with_capacity(self.lists.len()); for list in self.lists.iter() { guards.push(list.lock()); } for g in &mut guards { g.for_each(&mut f); } } } } tokio-1.43.1/src/util/sync_wrapper.rs000064400000000000000000000021571046102023000156720ustar 00000000000000//! This module contains a type that can make `Send + !Sync` types `Sync` by //! disallowing all immutable access to the value. //! //! A similar primitive is provided in the `sync_wrapper` crate. use std::any::Any; pub(crate) struct SyncWrapper { value: T, } // safety: The SyncWrapper being send allows you to send the inner value across // thread boundaries. unsafe impl Send for SyncWrapper {} // safety: An immutable reference to a SyncWrapper is useless, so moving such an // immutable reference across threads is safe. unsafe impl Sync for SyncWrapper {} impl SyncWrapper { pub(crate) fn new(value: T) -> Self { Self { value } } pub(crate) fn into_inner(self) -> T { self.value } } impl SyncWrapper> { /// Attempt to downcast using `Any::downcast_ref()` to a type that is known to be `Sync`. pub(crate) fn downcast_ref_sync(&self) -> Option<&T> { // SAFETY: if the downcast fails, the inner value is not touched, // so no thread-safety violation can occur. self.value.downcast_ref() } } tokio-1.43.1/src/util/trace.rs000064400000000000000000000154041046102023000142530ustar 00000000000000cfg_rt! { use std::marker::PhantomData; #[derive(Copy, Clone)] pub(crate) struct SpawnMeta<'a> { /// The name of the task #[cfg(all(tokio_unstable, feature = "tracing"))] pub(crate) name: Option<&'a str>, /// The original size of the future or function being spawned #[cfg(all(tokio_unstable, feature = "tracing"))] pub(crate) original_size: usize, _pd: PhantomData<&'a ()>, } impl<'a> SpawnMeta<'a> { /// Create new spawn meta with a name and original size (before possible auto-boxing) #[cfg(all(tokio_unstable, feature = "tracing"))] pub(crate) fn new(name: Option<&'a str>, original_size: usize) -> Self { Self { name, original_size, _pd: PhantomData, } } /// Create a new unnamed spawn meta with the original size (before possible auto-boxing) pub(crate) fn new_unnamed(original_size: usize) -> Self { #[cfg(not(all(tokio_unstable, feature = "tracing")))] let _original_size = original_size; Self { #[cfg(all(tokio_unstable, feature = "tracing"))] name: None, #[cfg(all(tokio_unstable, feature = "tracing"))] original_size, _pd: PhantomData, } } } cfg_trace! { use core::{ pin::Pin, task::{Context, Poll}, }; use pin_project_lite::pin_project; use std::mem; use std::future::Future; use tracing::instrument::Instrument; pub(crate) use tracing::instrument::Instrumented; #[inline] #[track_caller] pub(crate) fn task(task: F, kind: &'static str, meta: SpawnMeta<'_>, id: u64) -> Instrumented { #[track_caller] fn get_span(kind: &'static str, spawn_meta: SpawnMeta<'_>, id: u64, task_size: usize) -> tracing::Span { let location = std::panic::Location::caller(); let original_size = if spawn_meta.original_size != task_size { Some(spawn_meta.original_size) } else { None }; tracing::trace_span!( target: "tokio::task", parent: None, "runtime.spawn", %kind, task.name = %spawn_meta.name.unwrap_or_default(), task.id = id, original_size.bytes = original_size, size.bytes = task_size, loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ) } use tracing::instrument::Instrument; let span = get_span(kind, meta, id, mem::size_of::()); task.instrument(span) } #[inline] #[track_caller] pub(crate) fn blocking_task(task: Fut, spawn_meta: SpawnMeta<'_>, id: u64) -> Instrumented { let location = std::panic::Location::caller(); let fn_size = mem::size_of::(); let original_size = if spawn_meta.original_size != fn_size { Some(spawn_meta.original_size) } else { None }; let span = tracing::trace_span!( target: "tokio::task::blocking", "runtime.spawn", kind = %"blocking", task.name = %spawn_meta.name.unwrap_or_default(), task.id = id, "fn" = %std::any::type_name::(), original_size.bytes = original_size, size.bytes = fn_size, loc.file = location.file(), loc.line = location.line(), loc.col = location.column(), ); task.instrument(span) } pub(crate) fn async_op(inner: P, resource_span: tracing::Span, source: &str, poll_op_name: &'static str, inherits_child_attrs: bool) -> InstrumentedAsyncOp where P: FnOnce() -> F { resource_span.in_scope(|| { let async_op_span = tracing::trace_span!("runtime.resource.async_op", source = source, inherits_child_attrs = inherits_child_attrs); let enter = async_op_span.enter(); let async_op_poll_span = tracing::trace_span!("runtime.resource.async_op.poll"); let inner = inner(); drop(enter); let tracing_ctx = AsyncOpTracingCtx { async_op_span, async_op_poll_span, resource_span: resource_span.clone(), }; InstrumentedAsyncOp { inner, tracing_ctx, poll_op_name, } }) } #[derive(Debug, Clone)] pub(crate) struct AsyncOpTracingCtx { pub(crate) async_op_span: tracing::Span, pub(crate) async_op_poll_span: tracing::Span, pub(crate) resource_span: tracing::Span, } pin_project! { #[derive(Debug, Clone)] pub(crate) struct InstrumentedAsyncOp { #[pin] pub(crate) inner: F, pub(crate) tracing_ctx: AsyncOpTracingCtx, pub(crate) poll_op_name: &'static str } } impl Future for InstrumentedAsyncOp { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); let poll_op_name = &*this.poll_op_name; let _res_enter = this.tracing_ctx.resource_span.enter(); let _async_op_enter = this.tracing_ctx.async_op_span.enter(); let _async_op_poll_enter = this.tracing_ctx.async_op_poll_span.enter(); trace_poll_op!(poll_op_name, this.inner.poll(cx)) } } } cfg_not_trace! { #[inline] pub(crate) fn task(task: F, _kind: &'static str, _meta: SpawnMeta<'_>, _id: u64) -> F { // nop task } #[inline] pub(crate) fn blocking_task(task: Fut, _spawn_meta: SpawnMeta<'_>, _id: u64) -> Fut { let _ = PhantomData::<&Fn>; // nop task } } } cfg_time! { #[track_caller] pub(crate) fn caller_location() -> Option<&'static std::panic::Location<'static>> { #[cfg(all(tokio_unstable, feature = "tracing"))] return Some(std::panic::Location::caller()); #[cfg(not(all(tokio_unstable, feature = "tracing")))] None } } tokio-1.43.1/src/util/try_lock.rs000064400000000000000000000032311046102023000147760ustar 00000000000000use crate::loom::sync::atomic::AtomicBool; use std::cell::UnsafeCell; use std::marker::PhantomData; use std::ops::{Deref, DerefMut}; use std::sync::atomic::Ordering::SeqCst; pub(crate) struct TryLock { locked: AtomicBool, data: UnsafeCell, } pub(crate) struct LockGuard<'a, T> { lock: &'a TryLock, _p: PhantomData>, } unsafe impl Send for TryLock {} unsafe impl Sync for TryLock {} unsafe impl Sync for LockGuard<'_, T> {} macro_rules! new { ($data:ident) => { TryLock { locked: AtomicBool::new(false), data: UnsafeCell::new($data), } }; } impl TryLock { #[cfg(not(loom))] /// Create a new `TryLock` pub(crate) const fn new(data: T) -> TryLock { new!(data) } #[cfg(loom)] /// Create a new `TryLock` pub(crate) fn new(data: T) -> TryLock { new!(data) } /// Attempt to acquire lock pub(crate) fn try_lock(&self) -> Option> { if self .locked .compare_exchange(false, true, SeqCst, SeqCst) .is_err() { return None; } Some(LockGuard { lock: self, _p: PhantomData, }) } } impl Deref for LockGuard<'_, T> { type Target = T; fn deref(&self) -> &T { unsafe { &*self.lock.data.get() } } } impl DerefMut for LockGuard<'_, T> { fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.lock.data.get() } } } impl Drop for LockGuard<'_, T> { fn drop(&mut self) { self.lock.locked.store(false, SeqCst); } } tokio-1.43.1/src/util/wake.rs000064400000000000000000000035451046102023000141070ustar 00000000000000use crate::loom::sync::Arc; use std::marker::PhantomData; use std::mem::ManuallyDrop; use std::ops::Deref; use std::task::{RawWaker, RawWakerVTable, Waker}; /// Simplified waking interface based on Arcs. pub(crate) trait Wake: Send + Sync + Sized + 'static { /// Wake by value. fn wake(arc_self: Arc); /// Wake by reference. fn wake_by_ref(arc_self: &Arc); } /// A `Waker` that is only valid for a given lifetime. #[derive(Debug)] pub(crate) struct WakerRef<'a> { waker: ManuallyDrop, _p: PhantomData<&'a ()>, } impl Deref for WakerRef<'_> { type Target = Waker; fn deref(&self) -> &Waker { &self.waker } } /// Creates a reference to a `Waker` from a reference to `Arc`. pub(crate) fn waker_ref(wake: &Arc) -> WakerRef<'_> { let ptr = Arc::as_ptr(wake).cast::<()>(); let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::())) }; WakerRef { waker: ManuallyDrop::new(waker), _p: PhantomData, } } fn waker_vtable() -> &'static RawWakerVTable { &RawWakerVTable::new( clone_arc_raw::, wake_arc_raw::, wake_by_ref_arc_raw::, drop_arc_raw::, ) } unsafe fn clone_arc_raw(data: *const ()) -> RawWaker { Arc::::increment_strong_count(data as *const T); RawWaker::new(data, waker_vtable::()) } unsafe fn wake_arc_raw(data: *const ()) { let arc: Arc = Arc::from_raw(data as *const T); Wake::wake(arc); } // used by `waker_ref` unsafe fn wake_by_ref_arc_raw(data: *const ()) { // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop let arc = ManuallyDrop::new(Arc::::from_raw(data.cast())); Wake::wake_by_ref(&arc); } unsafe fn drop_arc_raw(data: *const ()) { drop(Arc::::from_raw(data.cast())); } tokio-1.43.1/src/util/wake_list.rs000064400000000000000000000051201046102023000151310ustar 00000000000000use core::mem::MaybeUninit; use core::ptr; use std::task::Waker; const NUM_WAKERS: usize = 32; /// A list of wakers to be woken. /// /// # Invariants /// /// The first `curr` elements of `inner` are initialized. pub(crate) struct WakeList { inner: [MaybeUninit; NUM_WAKERS], curr: usize, } impl WakeList { pub(crate) fn new() -> Self { const UNINIT_WAKER: MaybeUninit = MaybeUninit::uninit(); Self { inner: [UNINIT_WAKER; NUM_WAKERS], curr: 0, } } #[inline] pub(crate) fn can_push(&self) -> bool { self.curr < NUM_WAKERS } pub(crate) fn push(&mut self, val: Waker) { debug_assert!(self.can_push()); self.inner[self.curr] = MaybeUninit::new(val); self.curr += 1; } pub(crate) fn wake_all(&mut self) { struct DropGuard { start: *mut Waker, end: *mut Waker, } impl Drop for DropGuard { fn drop(&mut self) { // SAFETY: Both pointers are part of the same object, with `start <= end`. let len = unsafe { self.end.offset_from(self.start) } as usize; let slice = ptr::slice_from_raw_parts_mut(self.start, len); // SAFETY: All elements in `start..len` are initialized, so we can drop them. unsafe { ptr::drop_in_place(slice) }; } } debug_assert!(self.curr <= NUM_WAKERS); let mut guard = { let start = self.inner.as_mut_ptr().cast::(); // SAFETY: The resulting pointer is in bounds or one after the length of the same object. let end = unsafe { start.add(self.curr) }; // Transfer ownership of the wakers in `inner` to `DropGuard`. self.curr = 0; DropGuard { start, end } }; while !ptr::eq(guard.start, guard.end) { // SAFETY: `start` is always initialized if `start != end`. let waker = unsafe { ptr::read(guard.start) }; // SAFETY: The resulting pointer is in bounds or one after the length of the same object. guard.start = unsafe { guard.start.add(1) }; // If this panics, then `guard` will clean up the remaining wakers. waker.wake(); } } } impl Drop for WakeList { fn drop(&mut self) { let slice = ptr::slice_from_raw_parts_mut(self.inner.as_mut_ptr().cast::(), self.curr); // SAFETY: The first `curr` elements are initialized, so we can drop them. unsafe { ptr::drop_in_place(slice) }; } } tokio-1.43.1/tests/_require_full.rs000064400000000000000000000007331046102023000154070ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #[cfg(not(any(feature = "full", target_family = "wasm")))] compile_error!("run main Tokio tests with `--features full`"); // CI sets `--cfg tokio_no_parking_lot` when trying to run tests with // `parking_lot` disabled. This check prevents "silent failure" if `parking_lot` // accidentally gets enabled. #[cfg(all(tokio_no_parking_lot, feature = "parking_lot"))] compile_error!("parking_lot feature enabled when it should not be"); tokio-1.43.1/tests/async_send_sync.rs000064400000000000000000001252711046102023000157410ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![allow(clippy::type_complexity, clippy::diverging_sub_expression)] use std::cell::Cell; use std::future::Future; use std::io::SeekFrom; use std::net::SocketAddr; use std::pin::Pin; use std::rc::Rc; use tokio::net::TcpStream; use tokio::time::{Duration, Instant}; // The names of these structs behaves better when sorted. // Send: Yes, Sync: Yes #[derive(Clone)] #[allow(unused)] struct YY {} // Send: Yes, Sync: No #[derive(Clone)] #[allow(unused)] struct YN { _value: Cell, } // Send: No, Sync: No #[derive(Clone)] #[allow(unused)] struct NN { _value: Rc, } #[allow(dead_code)] type BoxFutureSync = std::pin::Pin + Send + Sync>>; #[allow(dead_code)] type BoxFutureSend = std::pin::Pin + Send>>; #[allow(dead_code)] type BoxFuture = std::pin::Pin>>; #[allow(dead_code)] type BoxAsyncRead = std::pin::Pin>; #[allow(dead_code)] type BoxAsyncSeek = std::pin::Pin>; #[allow(dead_code)] type BoxAsyncWrite = std::pin::Pin>; #[allow(dead_code)] fn require_send(_t: &T) {} #[allow(dead_code)] fn require_sync(_t: &T) {} #[allow(dead_code)] fn require_unpin(_t: &T) {} #[allow(dead_code)] struct Invalid; #[allow(unused)] trait AmbiguousIfSend { fn some_item(&self) {} } impl AmbiguousIfSend<()> for T {} impl AmbiguousIfSend for T {} #[allow(unused)] trait AmbiguousIfSync { fn some_item(&self) {} } impl AmbiguousIfSync<()> for T {} impl AmbiguousIfSync for T {} #[allow(unused)] trait AmbiguousIfUnpin { fn some_item(&self) {} } impl AmbiguousIfUnpin<()> for T {} impl AmbiguousIfUnpin for T {} macro_rules! into_todo { ($typ:ty) => {{ let x: $typ = todo!(); x }}; } macro_rules! async_assert_fn_send { (Send & $(!)?Sync & $(!)?Unpin, $value:expr) => { require_send(&$value); }; (!Send & $(!)?Sync & $(!)?Unpin, $value:expr) => { AmbiguousIfSend::some_item(&$value); }; } macro_rules! async_assert_fn_sync { ($(!)?Send & Sync & $(!)?Unpin, $value:expr) => { require_sync(&$value); }; ($(!)?Send & !Sync & $(!)?Unpin, $value:expr) => { AmbiguousIfSync::some_item(&$value); }; } macro_rules! async_assert_fn_unpin { ($(!)?Send & $(!)?Sync & Unpin, $value:expr) => { require_unpin(&$value); }; ($(!)?Send & $(!)?Sync & !Unpin, $value:expr) => { AmbiguousIfUnpin::some_item(&$value); }; } macro_rules! async_assert_fn { ($($f:ident $(< $($generic:ty),* > )? )::+($($arg:ty),*): $($tok:tt)*) => { #[allow(unreachable_code)] #[allow(unused_variables)] const _: fn() = || { let f = $($f $(::<$($generic),*>)? )::+( $( into_todo!($arg) ),* ); async_assert_fn_send!($($tok)*, f); async_assert_fn_sync!($($tok)*, f); async_assert_fn_unpin!($($tok)*, f); }; }; } macro_rules! assert_value { ($type:ty: $($tok:tt)*) => { #[allow(unreachable_code)] #[allow(unused_variables)] const _: fn() = || { let f: $type = todo!(); async_assert_fn_send!($($tok)*, f); async_assert_fn_sync!($($tok)*, f); async_assert_fn_unpin!($($tok)*, f); }; }; } macro_rules! cfg_not_wasi { ($($item:item)*) => { $( #[cfg(not(target_os = "wasi"))] $item )* } } // Manually re-implementation of `async_assert_fn` for `poll_fn`. The macro // doesn't work for this particular case because constructing the closure // is too complicated. const _: fn() = || { let pinned = std::marker::PhantomPinned; let f = tokio::macros::support::poll_fn(move |_| { // Use `pinned` to take ownership of it. let _ = &pinned; std::task::Poll::Pending::<()> }); require_send(&f); require_sync(&f); AmbiguousIfUnpin::some_item(&f); }; cfg_not_wasi! { mod fs { use super::*; assert_value!(tokio::fs::DirBuilder: Send & Sync & Unpin); assert_value!(tokio::fs::DirEntry: Send & Sync & Unpin); assert_value!(tokio::fs::File: Send & Sync & Unpin); assert_value!(tokio::fs::OpenOptions: Send & Sync & Unpin); assert_value!(tokio::fs::ReadDir: Send & Sync & Unpin); async_assert_fn!(tokio::fs::canonicalize(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::copy(&str, &str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::create_dir(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::create_dir_all(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::hard_link(&str, &str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::metadata(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::read(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::read_dir(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::read_link(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::read_to_string(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::remove_dir(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::remove_dir_all(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::remove_file(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::rename(&str, &str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::set_permissions(&str, std::fs::Permissions): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::symlink_metadata(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::write(&str, Vec): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::ReadDir::next_entry(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::OpenOptions::open(_, &str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::DirBuilder::create(_, &str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::DirEntry::metadata(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::DirEntry::file_type(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::open(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::create(&str): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::sync_all(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::sync_data(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::set_len(_, u64): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::metadata(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::try_clone(_): Send & Sync & !Unpin); async_assert_fn!(tokio::fs::File::into_std(_): Send & Sync & !Unpin); async_assert_fn!( tokio::fs::File::set_permissions(_, std::fs::Permissions): Send & Sync & !Unpin ); } } cfg_not_wasi! { assert_value!(tokio::net::TcpSocket: Send & Sync & Unpin); async_assert_fn!(tokio::net::TcpListener::bind(SocketAddr): Send & Sync & !Unpin); async_assert_fn!(tokio::net::TcpStream::connect(SocketAddr): Send & Sync & !Unpin); } assert_value!(tokio::net::TcpListener: Send & Sync & Unpin); assert_value!(tokio::net::TcpStream: Send & Sync & Unpin); assert_value!(tokio::net::tcp::OwnedReadHalf: Send & Sync & Unpin); assert_value!(tokio::net::tcp::OwnedWriteHalf: Send & Sync & Unpin); assert_value!(tokio::net::tcp::ReadHalf<'_>: Send & Sync & Unpin); assert_value!(tokio::net::tcp::ReuniteError: Send & Sync & Unpin); assert_value!(tokio::net::tcp::WriteHalf<'_>: Send & Sync & Unpin); async_assert_fn!(tokio::net::TcpListener::accept(_): Send & Sync & !Unpin); async_assert_fn!(tokio::net::TcpStream::peek(_, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(tokio::net::TcpStream::readable(_): Send & Sync & !Unpin); async_assert_fn!(tokio::net::TcpStream::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(tokio::net::TcpStream::writable(_): Send & Sync & !Unpin); // Wasi does not support UDP cfg_not_wasi! { mod udp_socket { use super::*; assert_value!(tokio::net::UdpSocket: Send & Sync & Unpin); async_assert_fn!(tokio::net::UdpSocket::bind(SocketAddr): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::connect(_, SocketAddr): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::peek_from(_, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::readable(_): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::recv(_, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::recv_from(_, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::send(_, &[u8]): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::send_to(_, &[u8], SocketAddr): Send & Sync & !Unpin); async_assert_fn!(tokio::net::UdpSocket::writable(_): Send & Sync & !Unpin); } } async_assert_fn!(tokio::net::lookup_host(SocketAddr): Send & Sync & !Unpin); async_assert_fn!(tokio::net::tcp::ReadHalf::peek(_, &mut [u8]): Send & Sync & !Unpin); #[cfg(unix)] mod unix_datagram { use super::*; use tokio::net::*; assert_value!(UnixDatagram: Send & Sync & Unpin); assert_value!(UnixListener: Send & Sync & Unpin); assert_value!(UnixStream: Send & Sync & Unpin); assert_value!(unix::OwnedReadHalf: Send & Sync & Unpin); assert_value!(unix::OwnedWriteHalf: Send & Sync & Unpin); assert_value!(unix::ReadHalf<'_>: Send & Sync & Unpin); assert_value!(unix::ReuniteError: Send & Sync & Unpin); assert_value!(unix::SocketAddr: Send & Sync & Unpin); assert_value!(unix::UCred: Send & Sync & Unpin); assert_value!(unix::WriteHalf<'_>: Send & Sync & Unpin); async_assert_fn!(UnixDatagram::readable(_): Send & Sync & !Unpin); async_assert_fn!(UnixDatagram::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(UnixDatagram::recv(_, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(UnixDatagram::recv_from(_, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(UnixDatagram::send(_, &[u8]): Send & Sync & !Unpin); async_assert_fn!(UnixDatagram::send_to(_, &[u8], &str): Send & Sync & !Unpin); async_assert_fn!(UnixDatagram::writable(_): Send & Sync & !Unpin); async_assert_fn!(UnixListener::accept(_): Send & Sync & !Unpin); async_assert_fn!(UnixStream::connect(&str): Send & Sync & !Unpin); async_assert_fn!(UnixStream::readable(_): Send & Sync & !Unpin); async_assert_fn!(UnixStream::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(UnixStream::writable(_): Send & Sync & !Unpin); } #[cfg(unix)] mod unix_pipe { use super::*; use tokio::net::unix::pipe::*; assert_value!(OpenOptions: Send & Sync & Unpin); assert_value!(Receiver: Send & Sync & Unpin); assert_value!(Sender: Send & Sync & Unpin); async_assert_fn!(Receiver::readable(_): Send & Sync & !Unpin); async_assert_fn!(Receiver::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(Sender::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(Sender::writable(_): Send & Sync & !Unpin); } #[cfg(windows)] mod windows_named_pipe { use super::*; use tokio::net::windows::named_pipe::*; assert_value!(ClientOptions: Send & Sync & Unpin); assert_value!(NamedPipeClient: Send & Sync & Unpin); assert_value!(NamedPipeServer: Send & Sync & Unpin); assert_value!(PipeEnd: Send & Sync & Unpin); assert_value!(PipeInfo: Send & Sync & Unpin); assert_value!(PipeMode: Send & Sync & Unpin); assert_value!(ServerOptions: Send & Sync & Unpin); async_assert_fn!(NamedPipeClient::readable(_): Send & Sync & !Unpin); async_assert_fn!(NamedPipeClient::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(NamedPipeClient::writable(_): Send & Sync & !Unpin); async_assert_fn!(NamedPipeServer::connect(_): Send & Sync & !Unpin); async_assert_fn!(NamedPipeServer::readable(_): Send & Sync & !Unpin); async_assert_fn!(NamedPipeServer::ready(_, tokio::io::Interest): Send & Sync & !Unpin); async_assert_fn!(NamedPipeServer::writable(_): Send & Sync & !Unpin); } cfg_not_wasi! { mod test_process { use super::*; assert_value!(tokio::process::Child: Send & Sync & Unpin); assert_value!(tokio::process::ChildStderr: Send & Sync & Unpin); assert_value!(tokio::process::ChildStdin: Send & Sync & Unpin); assert_value!(tokio::process::ChildStdout: Send & Sync & Unpin); assert_value!(tokio::process::Command: Send & Sync & Unpin); async_assert_fn!(tokio::process::Child::kill(_): Send & Sync & !Unpin); async_assert_fn!(tokio::process::Child::wait(_): Send & Sync & !Unpin); async_assert_fn!(tokio::process::Child::wait_with_output(_): Send & Sync & !Unpin); } async_assert_fn!(tokio::signal::ctrl_c(): Send & Sync & !Unpin); } #[cfg(unix)] mod unix_signal { use super::*; assert_value!(tokio::signal::unix::Signal: Send & Sync & Unpin); assert_value!(tokio::signal::unix::SignalKind: Send & Sync & Unpin); async_assert_fn!(tokio::signal::unix::Signal::recv(_): Send & Sync & !Unpin); } #[cfg(windows)] mod windows_signal { use super::*; assert_value!(tokio::signal::windows::CtrlC: Send & Sync & Unpin); assert_value!(tokio::signal::windows::CtrlBreak: Send & Sync & Unpin); async_assert_fn!(tokio::signal::windows::CtrlC::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::signal::windows::CtrlBreak::recv(_): Send & Sync & !Unpin); } assert_value!(tokio::sync::AcquireError: Send & Sync & Unpin); assert_value!(tokio::sync::Barrier: Send & Sync & Unpin); assert_value!(tokio::sync::BarrierWaitResult: Send & Sync & Unpin); assert_value!(tokio::sync::MappedMutexGuard<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::MappedMutexGuard<'_, YN>: Send & !Sync & Unpin); assert_value!(tokio::sync::MappedMutexGuard<'_, YY>: Send & Sync & Unpin); assert_value!(tokio::sync::Mutex: !Send & !Sync & Unpin); assert_value!(tokio::sync::Mutex: Send & Sync & Unpin); assert_value!(tokio::sync::Mutex: Send & Sync & Unpin); assert_value!(tokio::sync::MutexGuard<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::MutexGuard<'_, YN>: Send & !Sync & Unpin); assert_value!(tokio::sync::MutexGuard<'_, YY>: Send & Sync & Unpin); assert_value!(tokio::sync::Notify: Send & Sync & Unpin); assert_value!(tokio::sync::OnceCell: !Send & !Sync & Unpin); assert_value!(tokio::sync::OnceCell: Send & !Sync & Unpin); assert_value!(tokio::sync::OnceCell: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard: Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMappedMutexGuard: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockMappedWriteGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockMappedWriteGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockMappedWriteGuard: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockReadGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockReadGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockReadGuard: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockWriteGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockWriteGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockWriteGuard: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedSemaphorePermit: Send & Sync & Unpin); assert_value!(tokio::sync::RwLock: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLock: Send & !Sync & Unpin); assert_value!(tokio::sync::RwLock: Send & Sync & Unpin); assert_value!(tokio::sync::RwLockMappedWriteGuard<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLockMappedWriteGuard<'_, YN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLockMappedWriteGuard<'_, YY>: Send & Sync & Unpin); assert_value!(tokio::sync::RwLockReadGuard<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLockReadGuard<'_, YN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLockReadGuard<'_, YY>: Send & Sync & Unpin); assert_value!(tokio::sync::RwLockWriteGuard<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLockWriteGuard<'_, YN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::RwLockWriteGuard<'_, YY>: Send & Sync & Unpin); assert_value!(tokio::sync::Semaphore: Send & Sync & Unpin); assert_value!(tokio::sync::SemaphorePermit<'_>: Send & Sync & Unpin); assert_value!(tokio::sync::TryAcquireError: Send & Sync & Unpin); assert_value!(tokio::sync::TryLockError: Send & Sync & Unpin); assert_value!(tokio::sync::broadcast::Receiver: !Send & !Sync & Unpin); assert_value!(tokio::sync::broadcast::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::broadcast::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::broadcast::Sender: !Send & !Sync & Unpin); assert_value!(tokio::sync::broadcast::Sender: Send & Sync & Unpin); assert_value!(tokio::sync::broadcast::Sender: Send & Sync & Unpin); assert_value!(tokio::sync::futures::Notified<'_>: Send & Sync & !Unpin); assert_value!(tokio::sync::mpsc::OwnedPermit: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::OwnedPermit: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::OwnedPermit: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::Permit<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::Permit<'_, YN>: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::Permit<'_, YY>: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::Receiver: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::Sender: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::Sender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::Sender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::UnboundedReceiver: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::UnboundedReceiver: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::UnboundedReceiver: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::UnboundedSender: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::UnboundedSender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::UnboundedSender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::WeakSender: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::WeakSender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::WeakSender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::WeakUnboundedSender: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::WeakUnboundedSender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::WeakUnboundedSender: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::error::SendError: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::error::SendError: Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::error::SendError: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::error::SendTimeoutError: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::error::SendTimeoutError: Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::error::SendTimeoutError: Send & Sync & Unpin); assert_value!(tokio::sync::mpsc::error::TrySendError: !Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::error::TrySendError: Send & !Sync & Unpin); assert_value!(tokio::sync::mpsc::error::TrySendError: Send & Sync & Unpin); assert_value!(tokio::sync::oneshot::Receiver: !Send & !Sync & Unpin); assert_value!(tokio::sync::oneshot::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::oneshot::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::oneshot::Sender: !Send & !Sync & Unpin); assert_value!(tokio::sync::oneshot::Sender: Send & Sync & Unpin); assert_value!(tokio::sync::oneshot::Sender: Send & Sync & Unpin); assert_value!(tokio::sync::watch::Receiver: !Send & !Sync & Unpin); assert_value!(tokio::sync::watch::Receiver: !Send & !Sync & Unpin); assert_value!(tokio::sync::watch::Receiver: Send & Sync & Unpin); assert_value!(tokio::sync::watch::Ref<'_, NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::watch::Ref<'_, YN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::watch::Ref<'_, YY>: !Send & Sync & Unpin); assert_value!(tokio::sync::watch::Sender: !Send & !Sync & Unpin); assert_value!(tokio::sync::watch::Sender: !Send & !Sync & Unpin); assert_value!(tokio::sync::watch::Sender: Send & Sync & Unpin); assert_value!(tokio::task::JoinError: Send & Sync & Unpin); assert_value!(tokio::task::JoinHandle: !Send & !Sync & Unpin); assert_value!(tokio::task::JoinHandle: Send & Sync & Unpin); assert_value!(tokio::task::JoinHandle: Send & Sync & Unpin); assert_value!(tokio::task::JoinSet: !Send & !Sync & Unpin); assert_value!(tokio::task::JoinSet: Send & Sync & Unpin); assert_value!(tokio::task::JoinSet: Send & Sync & Unpin); assert_value!(tokio::task::LocalSet: !Send & !Sync & Unpin); async_assert_fn!(tokio::sync::Barrier::wait(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Mutex::lock(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::Mutex::lock_owned(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::Mutex::lock(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Mutex::lock_owned(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Mutex::lock(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Mutex::lock_owned(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Notify::notified(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin + Send + Sync>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin + Send>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send + Sync>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin>>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin + Send + Sync>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin + Send>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send + Sync>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin>>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin + Send + Sync>>): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin + Send>>): Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send + Sync>>): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send>>): Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin>>>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::read(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::write(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::read(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::write(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::read(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::write(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Semaphore::acquire(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Semaphore::acquire_many(_, u32): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Semaphore::acquire_many_owned(_, u32): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::Semaphore::acquire_owned(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::broadcast::Receiver::recv(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::broadcast::Receiver::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::broadcast::Receiver::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Receiver::recv(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Receiver::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Receiver::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::closed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::reserve(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::reserve_owned(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::send(_, NN): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::send_timeout(_, NN, Duration): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::reserve(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::reserve_owned(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::send(_, YN): Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::send_timeout(_, YN, Duration): Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::reserve(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::reserve_owned(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::send(_, YY): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::Sender::send_timeout(_, YY, Duration): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::UnboundedReceiver::recv(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::UnboundedReceiver::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::UnboundedReceiver::recv(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::UnboundedSender::closed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::UnboundedSender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::mpsc::UnboundedSender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::oneshot::Sender::closed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::oneshot::Sender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::oneshot::Sender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::watch::Receiver::changed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::watch::Receiver::changed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::watch::Receiver::changed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::watch::Sender::closed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::watch::Sender::closed(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::watch::Sender::closed(_): Send & Sync & !Unpin); async_assert_fn!(tokio::task::JoinSet>::join_next(_): Send & Sync & !Unpin); async_assert_fn!(tokio::task::JoinSet>::shutdown(_): Send & Sync & !Unpin); async_assert_fn!(tokio::task::JoinSet>::join_next(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::JoinSet>::shutdown(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::JoinSet::join_next(_): Send & Sync & !Unpin); async_assert_fn!(tokio::task::JoinSet::shutdown(_): Send & Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey>::scope(_, Cell, BoxFuture<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey>::scope(_, Cell, BoxFutureSend<()>): Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey>::scope(_, Cell, BoxFutureSync<()>): Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey>::scope(_, Rc, BoxFuture<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey>::scope(_, Rc, BoxFutureSend<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey>::scope(_, Rc, BoxFutureSync<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey::scope(_, u32, BoxFuture<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey::scope(_, u32, BoxFutureSend<()>): Send & !Sync & !Unpin); async_assert_fn!(tokio::task::LocalKey::scope(_, u32, BoxFutureSync<()>): Send & Sync & !Unpin); async_assert_fn!(tokio::task::LocalSet::run_until(_, BoxFutureSync<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::task::unconstrained(BoxFuture<()>): !Send & !Sync & Unpin); async_assert_fn!(tokio::task::unconstrained(BoxFutureSend<()>): Send & !Sync & Unpin); async_assert_fn!(tokio::task::unconstrained(BoxFutureSync<()>): Send & Sync & Unpin); assert_value!(tokio::runtime::Builder: Send & Sync & Unpin); assert_value!(tokio::runtime::EnterGuard<'_>: !Send & Sync & Unpin); assert_value!(tokio::runtime::Handle: Send & Sync & Unpin); assert_value!(tokio::runtime::Runtime: Send & Sync & Unpin); assert_value!(tokio::time::Interval: Send & Sync & Unpin); assert_value!(tokio::time::Instant: Send & Sync & Unpin); assert_value!(tokio::time::Sleep: Send & Sync & !Unpin); assert_value!(tokio::time::Timeout>: Send & Sync & !Unpin); assert_value!(tokio::time::Timeout>: Send & !Sync & !Unpin); assert_value!(tokio::time::Timeout>: !Send & !Sync & !Unpin); assert_value!(tokio::time::error::Elapsed: Send & Sync & Unpin); assert_value!(tokio::time::error::Error: Send & Sync & Unpin); async_assert_fn!(tokio::time::advance(Duration): Send & Sync & !Unpin); async_assert_fn!(tokio::time::sleep(Duration): Send & Sync & !Unpin); async_assert_fn!(tokio::time::sleep_until(Instant): Send & Sync & !Unpin); async_assert_fn!(tokio::time::timeout(Duration, BoxFutureSync<()>): Send & Sync & !Unpin); async_assert_fn!(tokio::time::timeout(Duration, BoxFutureSend<()>): Send & !Sync & !Unpin); async_assert_fn!(tokio::time::timeout(Duration, BoxFuture<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::time::timeout_at(Instant, BoxFutureSync<()>): Send & Sync & !Unpin); async_assert_fn!(tokio::time::timeout_at(Instant, BoxFutureSend<()>): Send & !Sync & !Unpin); async_assert_fn!(tokio::time::timeout_at(Instant, BoxFuture<()>): !Send & !Sync & !Unpin); async_assert_fn!(tokio::time::Interval::tick(_): Send & Sync & !Unpin); assert_value!(tokio::io::BufReader: Send & Sync & Unpin); assert_value!(tokio::io::BufStream: Send & Sync & Unpin); assert_value!(tokio::io::BufWriter: Send & Sync & Unpin); assert_value!(tokio::io::DuplexStream: Send & Sync & Unpin); assert_value!(tokio::io::Empty: Send & Sync & Unpin); assert_value!(tokio::io::Interest: Send & Sync & Unpin); assert_value!(tokio::io::Lines: Send & Sync & Unpin); assert_value!(tokio::io::ReadBuf<'_>: Send & Sync & Unpin); assert_value!(tokio::io::ReadHalf: Send & Sync & Unpin); assert_value!(tokio::io::Ready: Send & Sync & Unpin); assert_value!(tokio::io::Repeat: Send & Sync & Unpin); assert_value!(tokio::io::Sink: Send & Sync & Unpin); assert_value!(tokio::io::Split: Send & Sync & Unpin); assert_value!(tokio::io::Stderr: Send & Sync & Unpin); assert_value!(tokio::io::Stdin: Send & Sync & Unpin); assert_value!(tokio::io::Stdout: Send & Sync & Unpin); assert_value!(tokio::io::Take: Send & Sync & Unpin); assert_value!(tokio::io::WriteHalf: Send & Sync & Unpin); async_assert_fn!(tokio::io::copy(&mut TcpStream, &mut TcpStream): Send & Sync & !Unpin); async_assert_fn!( tokio::io::copy_bidirectional(&mut TcpStream, &mut TcpStream): Send & Sync & !Unpin ); async_assert_fn!(tokio::io::copy_buf(&mut tokio::io::BufReader, &mut TcpStream): Send & Sync & !Unpin); async_assert_fn!(tokio::io::empty(): Send & Sync & Unpin); async_assert_fn!(tokio::io::repeat(u8): Send & Sync & Unpin); async_assert_fn!(tokio::io::sink(): Send & Sync & Unpin); async_assert_fn!(tokio::io::split(TcpStream): Send & Sync & Unpin); async_assert_fn!(tokio::io::stderr(): Send & Sync & Unpin); async_assert_fn!(tokio::io::stdin(): Send & Sync & Unpin); async_assert_fn!(tokio::io::stdout(): Send & Sync & Unpin); async_assert_fn!(tokio::io::Split>::next_segment(_): Send & Sync & !Unpin); async_assert_fn!(tokio::io::Lines>::next_line(_): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncBufReadExt::read_until(&mut BoxAsyncRead, u8, &mut Vec): Send & Sync & !Unpin); async_assert_fn!( tokio::io::AsyncBufReadExt::read_line(&mut BoxAsyncRead, &mut String): Send & Sync & !Unpin ); async_assert_fn!(tokio::io::AsyncBufReadExt::fill_buf(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read(&mut BoxAsyncRead, &mut [u8]): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_buf(&mut BoxAsyncRead, &mut Vec): Send & Sync & !Unpin); async_assert_fn!( tokio::io::AsyncReadExt::read_exact(&mut BoxAsyncRead, &mut [u8]): Send & Sync & !Unpin ); async_assert_fn!(tokio::io::AsyncReadExt::read_u8(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i8(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u16(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i16(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u32(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i32(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u64(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i64(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u128(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i128(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_f32(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_f64(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u16_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i16_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u32_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i32_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u64_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i64_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_u128_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_i128_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_f32_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_f64_le(&mut BoxAsyncRead): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncReadExt::read_to_end(&mut BoxAsyncRead, &mut Vec): Send & Sync & !Unpin); async_assert_fn!( tokio::io::AsyncReadExt::read_to_string(&mut BoxAsyncRead, &mut String): Send & Sync & !Unpin ); async_assert_fn!(tokio::io::AsyncSeekExt::seek(&mut BoxAsyncSeek, SeekFrom): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncSeekExt::stream_position(&mut BoxAsyncSeek): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncWriteExt::write(&mut BoxAsyncWrite, &[u8]): Send & Sync & !Unpin); async_assert_fn!( tokio::io::AsyncWriteExt::write_vectored(&mut BoxAsyncWrite, _): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_buf(&mut BoxAsyncWrite, &mut bytes::Bytes): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_all_buf(&mut BoxAsyncWrite, &mut bytes::Bytes): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_all(&mut BoxAsyncWrite, &[u8]): Send & Sync & !Unpin ); async_assert_fn!(tokio::io::AsyncWriteExt::write_u8(&mut BoxAsyncWrite, u8): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncWriteExt::write_i8(&mut BoxAsyncWrite, i8): Send & Sync & !Unpin); async_assert_fn!( tokio::io::AsyncWriteExt::write_u16(&mut BoxAsyncWrite, u16): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i16(&mut BoxAsyncWrite, i16): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u32(&mut BoxAsyncWrite, u32): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i32(&mut BoxAsyncWrite, i32): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u64(&mut BoxAsyncWrite, u64): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i64(&mut BoxAsyncWrite, i64): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u128(&mut BoxAsyncWrite, u128): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i128(&mut BoxAsyncWrite, i128): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_f32(&mut BoxAsyncWrite, f32): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_f64(&mut BoxAsyncWrite, f64): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u16_le(&mut BoxAsyncWrite, u16): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i16_le(&mut BoxAsyncWrite, i16): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u32_le(&mut BoxAsyncWrite, u32): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i32_le(&mut BoxAsyncWrite, i32): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u64_le(&mut BoxAsyncWrite, u64): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i64_le(&mut BoxAsyncWrite, i64): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_u128_le(&mut BoxAsyncWrite, u128): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_i128_le(&mut BoxAsyncWrite, i128): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_f32_le(&mut BoxAsyncWrite, f32): Send & Sync & !Unpin ); async_assert_fn!( tokio::io::AsyncWriteExt::write_f64_le(&mut BoxAsyncWrite, f64): Send & Sync & !Unpin ); async_assert_fn!(tokio::io::AsyncWriteExt::flush(&mut BoxAsyncWrite): Send & Sync & !Unpin); async_assert_fn!(tokio::io::AsyncWriteExt::shutdown(&mut BoxAsyncWrite): Send & Sync & !Unpin); #[cfg(unix)] mod unix_asyncfd { use super::*; use tokio::io::unix::*; #[allow(unused)] struct ImplsFd { _t: T, } impl std::os::unix::io::AsRawFd for ImplsFd { fn as_raw_fd(&self) -> std::os::unix::io::RawFd { unreachable!() } } assert_value!(AsyncFd>: Send & Sync & Unpin); assert_value!(AsyncFd>: Send & !Sync & Unpin); assert_value!(AsyncFd>: !Send & !Sync & Unpin); assert_value!(AsyncFdReadyGuard<'_, ImplsFd>: Send & Sync & Unpin); assert_value!(AsyncFdReadyGuard<'_, ImplsFd>: !Send & !Sync & Unpin); assert_value!(AsyncFdReadyGuard<'_, ImplsFd>: !Send & !Sync & Unpin); assert_value!(AsyncFdReadyMutGuard<'_, ImplsFd>: Send & Sync & Unpin); assert_value!(AsyncFdReadyMutGuard<'_, ImplsFd>: Send & !Sync & Unpin); assert_value!(AsyncFdReadyMutGuard<'_, ImplsFd>: !Send & !Sync & Unpin); assert_value!(TryIoError: Send & Sync & Unpin); async_assert_fn!(AsyncFd>::readable(_): Send & Sync & !Unpin); async_assert_fn!(AsyncFd>::readable_mut(_): Send & Sync & !Unpin); async_assert_fn!(AsyncFd>::writable(_): Send & Sync & !Unpin); async_assert_fn!(AsyncFd>::writable_mut(_): Send & Sync & !Unpin); async_assert_fn!(AsyncFd>::readable(_): !Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::readable_mut(_): Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::writable(_): !Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::writable_mut(_): Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::readable(_): !Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::readable_mut(_): !Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::writable(_): !Send & !Sync & !Unpin); async_assert_fn!(AsyncFd>::writable_mut(_): !Send & !Sync & !Unpin); } tokio-1.43.1/tests/buffered.rs000064400000000000000000000026011046102023000143300ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support bind() use tokio::net::TcpListener; use tokio_test::assert_ok; use std::io::prelude::*; use std::net::TcpStream; use std::thread; #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn echo_server() { const N: usize = 1024; let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let msg = "foo bar baz"; let t = thread::spawn(move || { let mut s = assert_ok!(TcpStream::connect(addr)); let t2 = thread::spawn(move || { let mut s = assert_ok!(TcpStream::connect(addr)); let mut b = vec![0; msg.len() * N]; assert_ok!(s.read_exact(&mut b)); b }); let mut expected = Vec::::new(); for _i in 0..N { expected.extend(msg.as_bytes()); let res = assert_ok!(s.write(msg.as_bytes())); assert_eq!(res, msg.len()); } (expected, t2) }); let (mut a, _) = assert_ok!(srv.accept().await); let (mut b, _) = assert_ok!(srv.accept().await); let n = assert_ok!(tokio::io::copy(&mut a, &mut b).await); let (expected, t2) = t.join().unwrap(); let actual = t2.join().unwrap(); assert!(expected == actual); assert_eq!(n, msg.len() as u64 * 1024); } tokio-1.43.1/tests/coop_budget.rs000064400000000000000000000052561046102023000150510ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", target_os = "linux"))] use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use tokio::net::UdpSocket; /// Ensure that UDP sockets have functional budgeting /// /// # Design /// Two sockets communicate by spamming packets from one to the other. /// /// In Linux, this packet will be slammed through the entire network stack and into the receiver's buffer during the /// send system call because we are using the loopback interface. /// This happens because the softirq chain invoked on send when using the loopback interface covers virtually the /// entirety of the lifecycle of a packet within the kernel network stack. /// /// As a result, neither socket will ever encounter an EWOULDBLOCK, and the only way for these to yield during the loop /// is through budgeting. /// /// A second task runs in the background and increments a counter before yielding, allowing us to know how many times sockets yielded. /// Since we are both sending and receiving, that should happen once per 64 packets, because budgets are of size 128 /// and there are two budget events per packet, a send and a recv. #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn coop_budget_udp_send_recv() { const BUDGET: usize = 128; const N_ITERATIONS: usize = 1024; const PACKET: &[u8] = b"Hello, world"; const PACKET_LEN: usize = 12; assert_eq!( PACKET_LEN, PACKET.len(), "Defect in test, programmer can't do math" ); // bind each socket to a dynamic port, forcing IPv4 addressing on the localhost interface let tx = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let rx = UdpSocket::bind("127.0.0.1:0").await.unwrap(); tx.connect(rx.local_addr().unwrap()).await.unwrap(); rx.connect(tx.local_addr().unwrap()).await.unwrap(); let tracker = Arc::new(AtomicUsize::default()); let tracker_clone = Arc::clone(&tracker); tokio::task::yield_now().await; tokio::spawn(async move { loop { tracker_clone.fetch_add(1, Ordering::SeqCst); tokio::task::yield_now().await; } }); for _ in 0..N_ITERATIONS { tx.send(PACKET).await.unwrap(); let mut tmp = [0; PACKET_LEN]; // ensure that we aren't somehow accumulating other assert_eq!( PACKET_LEN, rx.recv(&mut tmp).await.unwrap(), "Defect in test case, received unexpected result from socket" ); assert_eq!( PACKET, &tmp, "Defect in test case, received unexpected result from socket" ); } assert_eq!(N_ITERATIONS / (BUDGET / 2), tracker.load(Ordering::SeqCst)); } tokio-1.43.1/tests/dump.rs000064400000000000000000000114101046102023000135110ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![cfg(all( tokio_unstable, tokio_taskdump, target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") ))] use std::hint::black_box; use tokio::runtime::{self, Handle}; #[inline(never)] async fn a() { black_box(b()).await } #[inline(never)] async fn b() { black_box(c()).await } #[inline(never)] async fn c() { loop { black_box(tokio::task::yield_now()).await } } #[test] fn current_thread() { let rt = runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); async fn dump() { let handle = Handle::current(); let dump = handle.dump().await; let tasks: Vec<_> = dump.tasks().iter().collect(); assert_eq!(tasks.len(), 3); for task in tasks { let id = task.id(); let trace = task.trace().to_string(); eprintln!("\n\n{id}:\n{trace}\n\n"); assert!(trace.contains("dump::a")); assert!(trace.contains("dump::b")); assert!(trace.contains("dump::c")); assert!(trace.contains("tokio::task::yield_now")); } } rt.block_on(async { tokio::select!( biased; _ = tokio::spawn(a()) => {}, _ = tokio::spawn(a()) => {}, _ = tokio::spawn(a()) => {}, _ = dump() => {}, ); }); } #[test] fn multi_thread() { let rt = runtime::Builder::new_multi_thread() .enable_all() .worker_threads(3) .build() .unwrap(); async fn dump() { let handle = Handle::current(); let dump = handle.dump().await; let tasks: Vec<_> = dump.tasks().iter().collect(); assert_eq!(tasks.len(), 3); for task in tasks { let id = task.id(); let trace = task.trace().to_string(); eprintln!("\n\n{id}:\n{trace}\n\n"); assert!(trace.contains("dump::a")); assert!(trace.contains("dump::b")); assert!(trace.contains("dump::c")); assert!(trace.contains("tokio::task::yield_now")); } } rt.block_on(async { tokio::select!( biased; _ = tokio::spawn(a()) => {}, _ = tokio::spawn(a()) => {}, _ = tokio::spawn(a()) => {}, _ = dump() => {}, ); }); } /// Regression tests for #6035. /// /// These tests ensure that dumping will not deadlock if a future completes /// during a trace. mod future_completes_during_trace { use super::*; use core::future::{poll_fn, Future}; /// A future that completes only during a trace. fn complete_during_trace() -> impl Future + Send { use std::task::Poll; poll_fn(|cx| { if Handle::is_tracing() { Poll::Ready(()) } else { cx.waker().wake_by_ref(); Poll::Pending } }) } #[test] fn current_thread() { let rt = runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); async fn dump() { let handle = Handle::current(); let _dump = handle.dump().await; } rt.block_on(async { let _ = tokio::join!(tokio::spawn(complete_during_trace()), dump()); }); } #[test] fn multi_thread() { let rt = runtime::Builder::new_multi_thread() .enable_all() .build() .unwrap(); async fn dump() { let handle = Handle::current(); let _dump = handle.dump().await; tokio::task::yield_now().await; } rt.block_on(async { let _ = tokio::join!(tokio::spawn(complete_during_trace()), dump()); }); } } /// Regression test for #6051. /// /// This test ensures that tasks notified outside of a worker will not be /// traced, since doing so will un-set their notified bit prior to them being /// run and panic. #[test] fn notified_during_tracing() { let rt = runtime::Builder::new_multi_thread() .enable_all() .worker_threads(3) .build() .unwrap(); let timeout = async { tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; }; let timer = rt.spawn(async { loop { tokio::time::sleep(tokio::time::Duration::from_nanos(1)).await; } }); let dump = async { loop { let handle = Handle::current(); let _dump = handle.dump().await; } }; rt.block_on(async { tokio::select!( biased; _ = timeout => {}, _ = timer => {}, _ = dump => {}, ); }); } tokio-1.43.1/tests/duplex_stream.rs000064400000000000000000000023411046102023000154230ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::io::IoSlice; use tokio::io::{AsyncReadExt, AsyncWriteExt}; const HELLO: &[u8] = b"hello world..."; #[tokio::test] async fn write_vectored() { let (mut client, mut server) = tokio::io::duplex(64); let ret = client .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)]) .await .unwrap(); assert_eq!(ret, HELLO.len() * 2); client.flush().await.unwrap(); drop(client); let mut buf = Vec::with_capacity(HELLO.len() * 2); let bytes_read = server.read_to_end(&mut buf).await.unwrap(); assert_eq!(bytes_read, HELLO.len() * 2); assert_eq!(buf, [HELLO, HELLO].concat()); } #[tokio::test] async fn write_vectored_and_shutdown() { let (mut client, mut server) = tokio::io::duplex(64); let ret = client .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)]) .await .unwrap(); assert_eq!(ret, HELLO.len() * 2); client.shutdown().await.unwrap(); drop(client); let mut buf = Vec::with_capacity(HELLO.len() * 2); let bytes_read = server.read_to_end(&mut buf).await.unwrap(); assert_eq!(bytes_read, HELLO.len() * 2); assert_eq!(buf, [HELLO, HELLO].concat()); } tokio-1.43.1/tests/fs.rs000064400000000000000000000007531046102023000131640ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support file operations use tokio::fs; use tokio_test::assert_ok; #[tokio::test] async fn path_read_write() { let temp = tempdir(); let dir = temp.path(); assert_ok!(fs::write(dir.join("bar"), b"bytes").await); let out = assert_ok!(fs::read(dir.join("bar")).await); assert_eq!(out, b"bytes"); } fn tempdir() -> tempfile::TempDir { tempfile::tempdir().unwrap() } tokio-1.43.1/tests/fs_canonicalize_dir.rs000064400000000000000000000011541046102023000165350ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tokio::fs; #[tokio::test] #[cfg(unix)] async fn canonicalize_root_dir_unix() { assert_eq!(fs::canonicalize("/.").await.unwrap().to_str().unwrap(), "/"); } #[tokio::test] #[cfg(windows)] async fn canonicalize_root_dir_windows() { // 2-step let bindings due to Rust memory semantics let dir_path = fs::canonicalize("C:\\.\\").await.unwrap(); let dir_name = dir_path.to_str().unwrap(); assert!(dir_name.starts_with("\\\\")); assert!(dir_name.ends_with("C:\\")); } tokio-1.43.1/tests/fs_copy.rs000064400000000000000000000023651046102023000142170ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tempfile::tempdir; use tokio::fs; #[tokio::test] #[cfg_attr(miri, ignore)] // No `fchmod` in miri. async fn copy() { let dir = tempdir().unwrap(); let source_path = dir.path().join("foo.txt"); let dest_path = dir.path().join("bar.txt"); fs::write(&source_path, b"Hello File!").await.unwrap(); fs::copy(&source_path, &dest_path).await.unwrap(); let from = fs::read(&source_path).await.unwrap(); let to = fs::read(&dest_path).await.unwrap(); assert_eq!(from, to); } #[tokio::test] #[cfg_attr(miri, ignore)] // No `fchmod` in miri. async fn copy_permissions() { let dir = tempdir().unwrap(); let from_path = dir.path().join("foo.txt"); let to_path = dir.path().join("bar.txt"); let from = tokio::fs::File::create(&from_path).await.unwrap(); let mut from_perms = from.metadata().await.unwrap().permissions(); from_perms.set_readonly(true); from.set_permissions(from_perms.clone()).await.unwrap(); tokio::fs::copy(from_path, &to_path).await.unwrap(); let to_perms = tokio::fs::metadata(to_path).await.unwrap().permissions(); assert_eq!(from_perms, to_perms); } tokio-1.43.1/tests/fs_dir.rs000064400000000000000000000062661046102023000140270ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tokio::fs; use tokio_test::{assert_err, assert_ok}; use std::sync::{Arc, Mutex}; use tempfile::tempdir; #[tokio::test] async fn create_dir() { let base_dir = tempdir().unwrap(); let new_dir = base_dir.path().join("foo"); let new_dir_2 = new_dir.clone(); assert_ok!(fs::create_dir(new_dir).await); assert!(new_dir_2.is_dir()); } #[tokio::test] async fn create_all() { let base_dir = tempdir().unwrap(); let new_dir = base_dir.path().join("foo").join("bar"); let new_dir_2 = new_dir.clone(); assert_ok!(fs::create_dir_all(new_dir).await); assert!(new_dir_2.is_dir()); } #[tokio::test] async fn build_dir() { let base_dir = tempdir().unwrap(); let new_dir = base_dir.path().join("foo").join("bar"); let new_dir_2 = new_dir.clone(); assert_ok!(fs::DirBuilder::new().recursive(true).create(new_dir).await); assert!(new_dir_2.is_dir()); assert_err!( fs::DirBuilder::new() .recursive(false) .create(new_dir_2) .await ); } #[tokio::test] #[cfg(unix)] async fn build_dir_mode_read_only() { let base_dir = tempdir().unwrap(); let new_dir = base_dir.path().join("abc"); assert_ok!( fs::DirBuilder::new() .recursive(true) .mode(0o444) .create(&new_dir) .await ); assert!(fs::metadata(new_dir) .await .expect("metadata result") .permissions() .readonly()); } #[tokio::test] async fn remove() { let base_dir = tempdir().unwrap(); let new_dir = base_dir.path().join("foo"); let new_dir_2 = new_dir.clone(); std::fs::create_dir(new_dir.clone()).unwrap(); assert_ok!(fs::remove_dir(new_dir).await); assert!(!new_dir_2.exists()); } #[tokio::test] async fn read_inherent() { let base_dir = tempdir().unwrap(); let p = base_dir.path(); std::fs::create_dir(p.join("aa")).unwrap(); std::fs::create_dir(p.join("bb")).unwrap(); std::fs::create_dir(p.join("cc")).unwrap(); let files = Arc::new(Mutex::new(Vec::new())); let f = files.clone(); let p = p.to_path_buf(); let mut entries = fs::read_dir(p).await.unwrap(); while let Some(e) = assert_ok!(entries.next_entry().await) { let s = e.file_name().to_str().unwrap().to_string(); f.lock().unwrap().push(s); } let mut files = files.lock().unwrap(); files.sort(); // because the order is not guaranteed assert_eq!( *files, vec!["aa".to_string(), "bb".to_string(), "cc".to_string()] ); } #[tokio::test] async fn read_dir_entry_info() { let temp_dir = tempdir().unwrap(); let file_path = temp_dir.path().join("a.txt"); fs::write(&file_path, b"Hello File!").await.unwrap(); let mut dir = fs::read_dir(temp_dir.path()).await.unwrap(); let first_entry = dir.next_entry().await.unwrap().unwrap(); assert_eq!(first_entry.path(), file_path); assert_eq!(first_entry.file_name(), "a.txt"); assert!(first_entry.metadata().await.unwrap().is_file()); assert!(first_entry.file_type().await.unwrap().is_file()); } tokio-1.43.1/tests/fs_file.rs000064400000000000000000000151021046102023000141550ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use std::io::prelude::*; use std::io::IoSlice; use tempfile::NamedTempFile; use tokio::fs::File; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, SeekFrom}; use tokio_test::task; const HELLO: &[u8] = b"hello world..."; #[tokio::test] async fn basic_read() { let mut tempfile = tempfile(); tempfile.write_all(HELLO).unwrap(); let mut file = File::open(tempfile.path()).await.unwrap(); let mut buf = [0; 1024]; let n = file.read(&mut buf).await.unwrap(); assert_eq!(n, HELLO.len()); assert_eq!(&buf[..n], HELLO); } #[tokio::test] async fn basic_write() { let tempfile = tempfile(); let mut file = File::create(tempfile.path()).await.unwrap(); file.write_all(HELLO).await.unwrap(); file.flush().await.unwrap(); let file = std::fs::read(tempfile.path()).unwrap(); assert_eq!(file, HELLO); } #[tokio::test] async fn basic_write_and_shutdown() { let tempfile = tempfile(); let mut file = File::create(tempfile.path()).await.unwrap(); file.write_all(HELLO).await.unwrap(); file.shutdown().await.unwrap(); let file = std::fs::read(tempfile.path()).unwrap(); assert_eq!(file, HELLO); } #[tokio::test] async fn write_vectored() { let tempfile = tempfile(); let mut file = File::create(tempfile.path()).await.unwrap(); let ret = file .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)]) .await .unwrap(); assert_eq!(ret, HELLO.len() * 2); file.flush().await.unwrap(); let file = std::fs::read(tempfile.path()).unwrap(); assert_eq!(file, [HELLO, HELLO].concat()); } #[tokio::test] async fn write_vectored_and_shutdown() { let tempfile = tempfile(); let mut file = File::create(tempfile.path()).await.unwrap(); let ret = file .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)]) .await .unwrap(); assert_eq!(ret, HELLO.len() * 2); file.shutdown().await.unwrap(); let file = std::fs::read(tempfile.path()).unwrap(); assert_eq!(file, [HELLO, HELLO].concat()); } #[tokio::test] async fn rewind_seek_position() { let tempfile = tempfile(); let mut file = File::create(tempfile.path()).await.unwrap(); file.seek(SeekFrom::Current(10)).await.unwrap(); file.rewind().await.unwrap(); assert_eq!(file.stream_position().await.unwrap(), 0); } #[tokio::test] async fn coop() { let mut tempfile = tempfile(); tempfile.write_all(HELLO).unwrap(); let mut task = task::spawn(async { let mut file = File::open(tempfile.path()).await.unwrap(); let mut buf = [0; 1024]; loop { let _ = file.read(&mut buf).await.unwrap(); file.seek(std::io::SeekFrom::Start(0)).await.unwrap(); } }); for _ in 0..1_000 { if task.poll().is_pending() { return; } } panic!("did not yield"); } #[tokio::test] async fn write_to_clone() { let tempfile = tempfile(); let file = File::create(tempfile.path()).await.unwrap(); let mut clone = file.try_clone().await.unwrap(); clone.write_all(HELLO).await.unwrap(); clone.flush().await.unwrap(); let contents = std::fs::read(tempfile.path()).unwrap(); assert_eq!(contents, HELLO); } #[tokio::test] async fn write_into_std() { let tempfile = tempfile(); let file = File::create(tempfile.path()).await.unwrap(); let mut std_file = file.into_std().await; std_file.write_all(HELLO).unwrap(); let contents = std::fs::read(tempfile.path()).unwrap(); assert_eq!(contents, HELLO); } #[tokio::test] async fn write_into_std_immediate() { let tempfile = tempfile(); let file = File::create(tempfile.path()).await.unwrap(); let mut std_file = file.try_into_std().unwrap(); std_file.write_all(HELLO).unwrap(); let contents = std::fs::read(tempfile.path()).unwrap(); assert_eq!(contents, HELLO); } #[tokio::test] async fn read_file_from_std() { let mut tempfile = tempfile(); tempfile.write_all(HELLO).unwrap(); let std_file = std::fs::File::open(tempfile.path()).unwrap(); let mut file = File::from(std_file); let mut buf = [0; 1024]; let n = file.read(&mut buf).await.unwrap(); assert_eq!(n, HELLO.len()); assert_eq!(&buf[..n], HELLO); } fn tempfile() -> NamedTempFile { NamedTempFile::new().unwrap() } #[tokio::test] async fn set_max_buf_size_read() { let mut tempfile = tempfile(); tempfile.write_all(HELLO).unwrap(); let mut file = File::open(tempfile.path()).await.unwrap(); let mut buf = [0; 1024]; file.set_max_buf_size(1); // A single read operation reads a maximum of 1 byte. assert_eq!(file.read(&mut buf).await.unwrap(), 1); } #[tokio::test] async fn set_max_buf_size_write() { let tempfile = tempfile(); let mut file = File::create(tempfile.path()).await.unwrap(); file.set_max_buf_size(1); // A single write operation writes a maximum of 1 byte. assert_eq!(file.write(HELLO).await.unwrap(), 1); } #[tokio::test] #[cfg_attr(miri, ignore)] #[cfg(unix)] async fn file_debug_fmt() { let tempfile = tempfile(); let file = File::open(tempfile.path()).await.unwrap(); assert_eq!( &format!("{file:?}")[0..33], "tokio::fs::File { std: File { fd:" ); } #[tokio::test] #[cfg(windows)] async fn file_debug_fmt() { let tempfile = tempfile(); let file = File::open(tempfile.path()).await.unwrap(); assert_eq!( &format!("{:?}", file)[0..37], "tokio::fs::File { std: File { handle:" ); } #[tokio::test] #[cfg(unix)] async fn unix_fd_is_valid() { use std::os::unix::io::AsRawFd; let tempfile = tempfile(); let file = File::create(tempfile.path()).await.unwrap(); assert!(file.as_raw_fd() as u64 > 0); } #[tokio::test] #[cfg(unix)] async fn read_file_from_unix_fd() { use std::os::unix::io::{FromRawFd, IntoRawFd}; let mut tempfile = tempfile(); tempfile.write_all(HELLO).unwrap(); let file1 = File::open(tempfile.path()).await.unwrap(); let raw_fd = file1.into_std().await.into_raw_fd(); assert!(raw_fd > 0); let mut file2 = unsafe { File::from_raw_fd(raw_fd) }; let mut buf = [0; 1024]; let n = file2.read(&mut buf).await.unwrap(); assert_eq!(n, HELLO.len()); assert_eq!(&buf[..n], HELLO); } #[tokio::test] #[cfg(windows)] async fn windows_handle() { use std::os::windows::io::AsRawHandle; let tempfile = tempfile(); let file = File::create(tempfile.path()).await.unwrap(); assert!(file.as_raw_handle() as u64 > 0); } tokio-1.43.1/tests/fs_link.rs000064400000000000000000000030121046102023000141700ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tokio::fs; use std::io::Write; use tempfile::tempdir; #[tokio::test] #[cfg_attr(miri, ignore)] // No `linkat` in miri. async fn test_hard_link() { let dir = tempdir().unwrap(); let src = dir.path().join("src.txt"); let dst = dir.path().join("dst.txt"); std::fs::File::create(&src) .unwrap() .write_all(b"hello") .unwrap(); fs::hard_link(&src, &dst).await.unwrap(); std::fs::File::create(&src) .unwrap() .write_all(b"new-data") .unwrap(); let content = fs::read(&dst).await.unwrap(); assert_eq!(content, b"new-data"); // test that this is not a symlink: assert!(fs::read_link(&dst).await.is_err()); } #[cfg(unix)] #[tokio::test] async fn test_symlink() { let dir = tempdir().unwrap(); let src = dir.path().join("src.txt"); let dst = dir.path().join("dst.txt"); std::fs::File::create(&src) .unwrap() .write_all(b"hello") .unwrap(); fs::symlink(&src, &dst).await.unwrap(); std::fs::File::create(&src) .unwrap() .write_all(b"new-data") .unwrap(); let content = fs::read(&dst).await.unwrap(); assert_eq!(content, b"new-data"); let read = fs::read_link(dst.clone()).await.unwrap(); assert!(read == src); let symlink_meta = fs::symlink_metadata(dst.clone()).await.unwrap(); assert!(symlink_meta.file_type().is_symlink()); } tokio-1.43.1/tests/fs_open_options.rs000064400000000000000000000050561046102023000157610ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use std::io::Write; use tempfile::NamedTempFile; use tokio::fs::OpenOptions; use tokio::io::AsyncReadExt; const HELLO: &[u8] = b"hello world..."; #[tokio::test] async fn open_with_open_options_and_read() { let mut tempfile = NamedTempFile::new().unwrap(); tempfile.write_all(HELLO).unwrap(); let mut file = OpenOptions::new().read(true).open(tempfile).await.unwrap(); let mut buf = [0; 1024]; let n = file.read(&mut buf).await.unwrap(); assert_eq!(n, HELLO.len()); assert_eq!(&buf[..n], HELLO); } #[tokio::test] async fn open_options_write() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().write(true)).contains("write: true")); } #[tokio::test] async fn open_options_append() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().append(true)).contains("append: true")); } #[tokio::test] async fn open_options_truncate() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().truncate(true)).contains("truncate: true")); } #[tokio::test] async fn open_options_create() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().create(true)).contains("create: true")); } #[tokio::test] async fn open_options_create_new() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().create_new(true)).contains("create_new: true")); } #[tokio::test] #[cfg(unix)] async fn open_options_mode() { let mode = format!("{:?}", OpenOptions::new().mode(0o644)); // TESTING HACK: use Debug output to check the stored data assert!( mode.contains("mode: 420 ") || mode.contains("mode: 0o000644 "), "mode is: {mode}" ); } #[tokio::test] #[cfg(target_os = "linux")] async fn open_options_custom_flags_linux() { // TESTING HACK: use Debug output to check the stored data assert!( format!("{:?}", OpenOptions::new().custom_flags(libc::O_TRUNC)) .contains("custom_flags: 512,") ); } #[tokio::test] #[cfg(any(target_os = "freebsd", target_os = "macos"))] async fn open_options_custom_flags_bsd_family() { // TESTING HACK: use Debug output to check the stored data assert!( format!("{:?}", OpenOptions::new().custom_flags(libc::O_NOFOLLOW)) .contains("custom_flags: 256,") ); } tokio-1.43.1/tests/fs_open_options_windows.rs000064400000000000000000000027221046102023000175300ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations #![cfg(windows)] use tokio::fs::OpenOptions; use windows_sys::Win32::Storage::FileSystem; #[tokio::test] #[cfg(windows)] async fn open_options_windows_access_mode() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().access_mode(0)).contains("access_mode: Some(0)")); } #[tokio::test] #[cfg(windows)] async fn open_options_windows_share_mode() { // TESTING HACK: use Debug output to check the stored data assert!(format!("{:?}", OpenOptions::new().share_mode(0)).contains("share_mode: 0,")); } #[tokio::test] #[cfg(windows)] async fn open_options_windows_custom_flags() { // TESTING HACK: use Debug output to check the stored data assert!(format!( "{:?}", OpenOptions::new().custom_flags(FileSystem::FILE_FLAG_DELETE_ON_CLOSE) ) .contains("custom_flags: 67108864,")); } #[tokio::test] #[cfg(windows)] async fn open_options_windows_attributes() { assert!(format!( "{:?}", OpenOptions::new().attributes(FileSystem::FILE_ATTRIBUTE_HIDDEN) ) .contains("attributes: 2,")); } #[tokio::test] #[cfg(windows)] async fn open_options_windows_security_qos_flags() { assert!(format!( "{:?}", OpenOptions::new().security_qos_flags(FileSystem::SECURITY_IDENTIFICATION) ) .contains("security_qos_flags: 1114112,")); } tokio-1.43.1/tests/fs_remove_dir_all.rs000064400000000000000000000016471046102023000162320ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tempfile::tempdir; use tokio::fs; #[tokio::test] async fn remove_dir_all() { let temp_dir = tempdir().unwrap(); let test_dir = temp_dir.path().join("test"); fs::create_dir(&test_dir).await.unwrap(); let file_path = test_dir.as_path().join("a.txt"); fs::write(&file_path, b"Hello File!").await.unwrap(); fs::remove_dir_all(test_dir.as_path()).await.unwrap(); // test dir should no longer exist match fs::try_exists(test_dir).await { Ok(exists) => assert!(!exists), Err(_) => println!("ignored try_exists error after remove_dir_all"), }; // contents should no longer exist match fs::try_exists(file_path).await { Ok(exists) => assert!(!exists), Err(_) => println!("ignored try_exists error after remove_dir_all"), }; } tokio-1.43.1/tests/fs_remove_file.rs000064400000000000000000000012241046102023000155320ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tempfile::tempdir; use tokio::fs; #[tokio::test] async fn remove_file() { let temp_dir = tempdir().unwrap(); let file_path = temp_dir.path().join("a.txt"); fs::write(&file_path, b"Hello File!").await.unwrap(); assert!(fs::try_exists(&file_path).await.unwrap()); fs::remove_file(&file_path).await.unwrap(); // should no longer exist match fs::try_exists(file_path).await { Ok(exists) => assert!(!exists), Err(_) => println!("ignored try_exists error after remove_file"), }; } tokio-1.43.1/tests/fs_rename.rs000064400000000000000000000014341046102023000145100ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tempfile::tempdir; use tokio::fs; #[tokio::test] async fn rename_file() { let temp_dir = tempdir().unwrap(); let file_path = temp_dir.path().join("a.txt"); fs::write(&file_path, b"Hello File!").await.unwrap(); assert!(fs::try_exists(&file_path).await.unwrap()); let new_file_path = temp_dir.path().join("b.txt"); fs::rename(&file_path, &new_file_path).await.unwrap(); assert!(fs::try_exists(new_file_path).await.unwrap()); // original file should no longer exist match fs::try_exists(file_path).await { Ok(exists) => assert!(!exists), Err(_) => println!("ignored try_exists error after rename"), }; } tokio-1.43.1/tests/fs_symlink_dir_windows.rs000064400000000000000000000014621046102023000173400ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations #![cfg(windows)] use tempfile::tempdir; use tokio::fs; #[tokio::test] async fn symlink_file_windows() { const FILE_NAME: &str = "abc.txt"; let temp_dir = tempdir().unwrap(); let dir1 = temp_dir.path().join("a"); fs::create_dir(&dir1).await.unwrap(); let file1 = dir1.as_path().join(FILE_NAME); fs::write(&file1, b"Hello File!").await.unwrap(); let dir2 = temp_dir.path().join("b"); fs::symlink_dir(&dir1, &dir2).await.unwrap(); fs::write(&file1, b"new data!").await.unwrap(); let file2 = dir2.as_path().join(FILE_NAME); let from = fs::read(&file1).await.unwrap(); let to = fs::read(&file2).await.unwrap(); assert_eq!(from, to); } tokio-1.43.1/tests/fs_symlink_file_windows.rs000064400000000000000000000012521046102023000174760ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations #![cfg(windows)] use tempfile::tempdir; use tokio::fs; #[tokio::test] async fn symlink_file_windows() { let dir = tempdir().unwrap(); let source_path = dir.path().join("foo.txt"); let dest_path = dir.path().join("bar.txt"); fs::write(&source_path, b"Hello File!").await.unwrap(); fs::symlink_file(&source_path, &dest_path).await.unwrap(); fs::write(&source_path, b"new data!").await.unwrap(); let from = fs::read(&source_path).await.unwrap(); let to = fs::read(&dest_path).await.unwrap(); assert_eq!(from, to); } tokio-1.43.1/tests/fs_try_exists.rs000064400000000000000000000031541046102023000154570ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tempfile::tempdir; use tokio::fs; #[tokio::test] #[cfg_attr(miri, ignore)] // No `chmod` in miri. async fn try_exists() { let dir = tempdir().unwrap(); let existing_path = dir.path().join("foo.txt"); fs::write(&existing_path, b"Hello File!").await.unwrap(); let nonexisting_path = dir.path().join("bar.txt"); assert!(fs::try_exists(existing_path).await.unwrap()); assert!(!fs::try_exists(nonexisting_path).await.unwrap()); // FreeBSD root user always has permission to stat. #[cfg(all(unix, not(target_os = "freebsd")))] { use std::os::unix::prelude::PermissionsExt; let permission_denied_directory_path = dir.path().join("baz"); fs::create_dir(&permission_denied_directory_path) .await .unwrap(); let permission_denied_file_path = permission_denied_directory_path.join("baz.txt"); fs::write(&permission_denied_file_path, b"Hello File!") .await .unwrap(); let mut perms = tokio::fs::metadata(&permission_denied_directory_path) .await .unwrap() .permissions(); perms.set_mode(0o244); fs::set_permissions(&permission_denied_directory_path, perms) .await .unwrap(); let permission_denied_result = fs::try_exists(permission_denied_file_path).await; assert_eq!( permission_denied_result.err().unwrap().kind(), std::io::ErrorKind::PermissionDenied ); } } tokio-1.43.1/tests/io_async_fd.rs000064400000000000000000000647531046102023000150430ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(unix, feature = "full"))] use std::os::unix::io::{AsRawFd, RawFd}; use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, }; use std::time::Duration; use std::{ future::Future, io::{self, ErrorKind, Read, Write}, task::{Context, Waker}, }; use nix::unistd::{read, write}; use futures::poll; use tokio::io::unix::{AsyncFd, AsyncFdReadyGuard}; use tokio::io::Interest; use tokio_test::{assert_err, assert_pending}; struct TestWaker { inner: Arc, waker: Waker, } #[derive(Default)] struct TestWakerInner { awoken: AtomicBool, } impl futures::task::ArcWake for TestWakerInner { fn wake_by_ref(arc_self: &Arc) { arc_self.awoken.store(true, Ordering::SeqCst); } } impl TestWaker { fn new() -> Self { let inner: Arc = Default::default(); Self { inner: inner.clone(), waker: futures::task::waker(inner), } } fn awoken(&self) -> bool { self.inner.awoken.swap(false, Ordering::SeqCst) } fn context(&self) -> Context<'_> { Context::from_waker(&self.waker) } } #[derive(Debug)] struct FileDescriptor { fd: std::os::fd::OwnedFd, } impl AsRawFd for FileDescriptor { fn as_raw_fd(&self) -> RawFd { self.fd.as_raw_fd() } } impl Read for &FileDescriptor { fn read(&mut self, buf: &mut [u8]) -> io::Result { read(self.fd.as_raw_fd(), buf).map_err(io::Error::from) } } impl Read for FileDescriptor { fn read(&mut self, buf: &mut [u8]) -> io::Result { (self as &Self).read(buf) } } impl Write for &FileDescriptor { fn write(&mut self, buf: &[u8]) -> io::Result { write(&self.fd, buf).map_err(io::Error::from) } fn flush(&mut self) -> io::Result<()> { Ok(()) } } impl Write for FileDescriptor { fn write(&mut self, buf: &[u8]) -> io::Result { (self as &Self).write(buf) } fn flush(&mut self) -> io::Result<()> { (self as &Self).flush() } } fn set_nonblocking(fd: RawFd) { use nix::fcntl::{OFlag, F_GETFL, F_SETFL}; let flags = nix::fcntl::fcntl(fd, F_GETFL).expect("fcntl(F_GETFD)"); if flags < 0 { panic!( "bad return value from fcntl(F_GETFL): {} ({:?})", flags, nix::Error::last() ); } let flags = OFlag::from_bits_truncate(flags) | OFlag::O_NONBLOCK; nix::fcntl::fcntl(fd, F_SETFL(flags)).expect("fcntl(F_SETFD)"); } fn socketpair() -> (FileDescriptor, FileDescriptor) { use nix::sys::socket::{self, AddressFamily, SockFlag, SockType}; let (fd_a, fd_b) = socket::socketpair( AddressFamily::Unix, SockType::Stream, None, SockFlag::empty(), ) .expect("socketpair"); let fds = (FileDescriptor { fd: fd_a }, FileDescriptor { fd: fd_b }); set_nonblocking(fds.0.fd.as_raw_fd()); set_nonblocking(fds.1.fd.as_raw_fd()); fds } fn drain(mut fd: &FileDescriptor, mut amt: usize) { let mut buf = [0u8; 512]; while amt > 0 { match fd.read(&mut buf[..]) { Err(e) if e.kind() == ErrorKind::WouldBlock => {} Ok(0) => panic!("unexpected EOF"), Err(e) => panic!("unexpected error: {e:?}"), Ok(x) => amt -= x, } } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn initially_writable() { let (a, b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); let afd_b = AsyncFd::new(b).unwrap(); afd_a.writable().await.unwrap().clear_ready(); afd_b.writable().await.unwrap().clear_ready(); tokio::select! { biased; _ = tokio::time::sleep(Duration::from_millis(10)) => {}, _ = afd_a.readable() => panic!("Unexpected readable state"), _ = afd_b.readable() => panic!("Unexpected readable state"), } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn reset_readable() { let (a, mut b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); let readable = afd_a.readable(); tokio::pin!(readable); tokio::select! { _ = readable.as_mut() => panic!(), _ = tokio::time::sleep(Duration::from_millis(10)) => {} } b.write_all(b"0").unwrap(); let mut guard = readable.await.unwrap(); guard .try_io(|_| afd_a.get_ref().read(&mut [0])) .unwrap() .unwrap(); // `a` is not readable, but the reactor still thinks it is // (because we have not observed a not-ready error yet) afd_a.readable().await.unwrap().retain_ready(); // Explicitly clear the ready state guard.clear_ready(); let readable = afd_a.readable(); tokio::pin!(readable); tokio::select! { _ = readable.as_mut() => panic!(), _ = tokio::time::sleep(Duration::from_millis(10)) => {} } b.write_all(b"0").unwrap(); // We can observe the new readable event afd_a.readable().await.unwrap().clear_ready(); } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn reset_writable() { let (a, b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); let mut guard = afd_a.writable().await.unwrap(); // Write until we get a WouldBlock. This also clears the ready state. let mut bytes = 0; while let Ok(Ok(amt)) = guard.try_io(|_| afd_a.get_ref().write(&[0; 512][..])) { bytes += amt; } // Writable state should be cleared now. let writable = afd_a.writable(); tokio::pin!(writable); tokio::select! { _ = writable.as_mut() => panic!(), _ = tokio::time::sleep(Duration::from_millis(10)) => {} } // Read from the other side; we should become writable now. drain(&b, bytes); let _ = writable.await.unwrap(); } #[derive(Debug)] struct ArcFd(Arc); impl AsRawFd for ArcFd { fn as_raw_fd(&self) -> RawFd { self.0.as_raw_fd() } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn drop_closes() { let (a, mut b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); assert_eq!( ErrorKind::WouldBlock, b.read(&mut [0]).err().unwrap().kind() ); std::mem::drop(afd_a); assert_eq!(0, b.read(&mut [0]).unwrap()); // into_inner does not close the fd let (a, mut b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); let _a: FileDescriptor = afd_a.into_inner(); assert_eq!( ErrorKind::WouldBlock, b.read(&mut [0]).err().unwrap().kind() ); // Drop closure behavior is delegated to the inner object let (a, mut b) = socketpair(); let arc_fd = Arc::new(a); let afd_a = AsyncFd::new(ArcFd(arc_fd.clone())).unwrap(); std::mem::drop(afd_a); assert_eq!( ErrorKind::WouldBlock, b.read(&mut [0]).err().unwrap().kind() ); std::mem::drop(arc_fd); // suppress unnecessary clone clippy warning } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn reregister() { let (a, _b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); let a = afd_a.into_inner(); AsyncFd::new(a).unwrap(); } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn guard_try_io() { let (a, mut b) = socketpair(); b.write_all(b"0").unwrap(); let afd_a = AsyncFd::new(a).unwrap(); let mut guard = afd_a.readable().await.unwrap(); afd_a.get_ref().read_exact(&mut [0]).unwrap(); // Should not clear the readable state let _ = guard.try_io(|_| Ok(())); // Still readable... let _ = afd_a.readable().await.unwrap(); // Should clear the readable state let _ = guard.try_io(|_| io::Result::<()>::Err(ErrorKind::WouldBlock.into())); // Assert not readable let readable = afd_a.readable(); tokio::pin!(readable); tokio::select! { _ = readable.as_mut() => panic!(), _ = tokio::time::sleep(Duration::from_millis(10)) => {} } // Write something down b again and make sure we're reawoken b.write_all(b"0").unwrap(); let _ = readable.await.unwrap(); } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn try_io_readable() { let (a, mut b) = socketpair(); let mut afd_a = AsyncFd::new(a).unwrap(); // Give the runtime some time to update bookkeeping. tokio::task::yield_now().await; { let mut called = false; let _ = afd_a.try_io_mut(Interest::READABLE, |_| { called = true; Ok(()) }); assert!( !called, "closure should not have been called, since socket should not be readable" ); } // Make `a` readable by writing to `b`. // Give the runtime some time to update bookkeeping. b.write_all(&[0]).unwrap(); tokio::task::yield_now().await; { let mut called = false; let _ = afd_a.try_io(Interest::READABLE, |_| { called = true; Ok(()) }); assert!( called, "closure should have been called, since socket should have data available to read" ); } { let mut called = false; let _ = afd_a.try_io(Interest::READABLE, |_| { called = true; io::Result::<()>::Err(ErrorKind::WouldBlock.into()) }); assert!( called, "closure should have been called, since socket should have data available to read" ); } { let mut called = false; let _ = afd_a.try_io(Interest::READABLE, |_| { called = true; Ok(()) }); assert!(!called, "closure should not have been called, since socket readable state should have been cleared"); } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn try_io_writable() { let (a, _b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); // Give the runtime some time to update bookkeeping. tokio::task::yield_now().await; { let mut called = false; let _ = afd_a.try_io(Interest::WRITABLE, |_| { called = true; Ok(()) }); assert!( called, "closure should have been called, since socket should still be marked as writable" ); } { let mut called = false; let _ = afd_a.try_io(Interest::WRITABLE, |_| { called = true; io::Result::<()>::Err(ErrorKind::WouldBlock.into()) }); assert!( called, "closure should have been called, since socket should still be marked as writable" ); } { let mut called = false; let _ = afd_a.try_io(Interest::WRITABLE, |_| { called = true; Ok(()) }); assert!(!called, "closure should not have been called, since socket writable state should have been cleared"); } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn multiple_waiters() { let (a, mut b) = socketpair(); let afd_a = Arc::new(AsyncFd::new(a).unwrap()); let barrier = Arc::new(tokio::sync::Barrier::new(11)); let mut tasks = Vec::new(); for _ in 0..10 { let afd_a = afd_a.clone(); let barrier = barrier.clone(); let f = async move { let notify_barrier = async { barrier.wait().await; futures::future::pending::<()>().await; }; tokio::select! { biased; guard = afd_a.readable() => { tokio::task::yield_now().await; guard.unwrap().clear_ready() }, _ = notify_barrier => unreachable!(), } std::mem::drop(afd_a); }; tasks.push(tokio::spawn(f)); } let mut all_tasks = futures::future::try_join_all(tasks); tokio::select! { r = std::pin::Pin::new(&mut all_tasks) => { r.unwrap(); // propagate panic panic!("Tasks exited unexpectedly") }, _ = barrier.wait() => {} } b.write_all(b"0").unwrap(); all_tasks.await.unwrap(); } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn poll_fns() { let (a, b) = socketpair(); let afd_a = Arc::new(AsyncFd::new(a).unwrap()); let afd_b = Arc::new(AsyncFd::new(b).unwrap()); // Fill up the write side of A let mut bytes = 0; while let Ok(amt) = afd_a.get_ref().write(&[0; 512]) { bytes += amt; } let waker = TestWaker::new(); assert_pending!(afd_a.as_ref().poll_read_ready(&mut waker.context())); let afd_a_2 = afd_a.clone(); let r_barrier = Arc::new(tokio::sync::Barrier::new(2)); let barrier_clone = r_barrier.clone(); let read_fut = tokio::spawn(async move { // Move waker onto this task first assert_pending!(poll!(std::future::poll_fn(|cx| afd_a_2 .as_ref() .poll_read_ready(cx)))); barrier_clone.wait().await; let _ = std::future::poll_fn(|cx| afd_a_2.as_ref().poll_read_ready(cx)).await; }); let afd_a_2 = afd_a.clone(); let w_barrier = Arc::new(tokio::sync::Barrier::new(2)); let barrier_clone = w_barrier.clone(); let mut write_fut = tokio::spawn(async move { // Move waker onto this task first assert_pending!(poll!(std::future::poll_fn(|cx| afd_a_2 .as_ref() .poll_write_ready(cx)))); barrier_clone.wait().await; let _ = std::future::poll_fn(|cx| afd_a_2.as_ref().poll_write_ready(cx)).await; }); r_barrier.wait().await; w_barrier.wait().await; let readable = afd_a.readable(); tokio::pin!(readable); tokio::select! { _ = &mut readable => unreachable!(), _ = tokio::task::yield_now() => {} } // Make A readable. We expect that 'readable' and 'read_fut' will both complete quickly afd_b.get_ref().write_all(b"0").unwrap(); let _ = tokio::join!(readable, read_fut); // Our original waker should _not_ be awoken (poll_read_ready retains only the last context) assert!(!waker.awoken()); // The writable side should not be awoken tokio::select! { _ = &mut write_fut => unreachable!(), _ = tokio::time::sleep(Duration::from_millis(5)) => {} } // Make it writable now drain(afd_b.get_ref(), bytes); // now we should be writable (ie - the waker for poll_write should still be registered after we wake the read side) let _ = write_fut.await; } fn assert_pending>(f: F) -> std::pin::Pin> { let mut pinned = Box::pin(f); assert_pending!(pinned .as_mut() .poll(&mut Context::from_waker(futures::task::noop_waker_ref()))); pinned } fn rt() -> tokio::runtime::Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_wakes_currently_pending() { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; let readable = assert_pending(afd_a.readable()); std::mem::drop(rt); // The future was initialized **before** dropping the rt assert_err!(futures::executor::block_on(readable)); // The future is initialized **after** dropping the rt. assert_err!(futures::executor::block_on(afd_a.readable())); } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_wakes_future_pending() { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; std::mem::drop(rt); assert_err!(futures::executor::block_on(afd_a.readable())); } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_wakes_pending_race() { // TODO: make this a loom test for _ in 0..100 { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; let _ = std::thread::spawn(move || std::mem::drop(rt)); // This may or may not return an error (but will be awoken) let _ = futures::executor::block_on(afd_a.readable()); // However retrying will always return an error assert_err!(futures::executor::block_on(afd_a.readable())); } } async fn poll_readable(fd: &AsyncFd) -> std::io::Result> { std::future::poll_fn(|cx| fd.poll_read_ready(cx)).await } async fn poll_writable(fd: &AsyncFd) -> std::io::Result> { std::future::poll_fn(|cx| fd.poll_write_ready(cx)).await } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_wakes_currently_pending_polls() { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; while afd_a.get_ref().write(&[0; 512]).is_ok() {} // make not writable let readable = assert_pending(poll_readable(&afd_a)); let writable = assert_pending(poll_writable(&afd_a)); std::mem::drop(rt); // Attempting to poll readiness when the rt is dropped is an error assert_err!(futures::executor::block_on(readable)); assert_err!(futures::executor::block_on(writable)); } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_wakes_poll() { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; std::mem::drop(rt); assert_err!(futures::executor::block_on(poll_readable(&afd_a))); assert_err!(futures::executor::block_on(poll_writable(&afd_a))); } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_then_clear_readiness() { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; let mut write_ready = rt.block_on(afd_a.writable()).unwrap(); std::mem::drop(rt); write_ready.clear_ready(); } #[test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. fn driver_shutdown_wakes_poll_race() { // TODO: make this a loom test for _ in 0..100 { let rt = rt(); let (a, _b) = socketpair(); let afd_a = { let _enter = rt.enter(); AsyncFd::new(a).unwrap() }; while afd_a.get_ref().write(&[0; 512]).is_ok() {} // make not writable let _ = std::thread::spawn(move || std::mem::drop(rt)); // The poll variants will always return an error in this case assert_err!(futures::executor::block_on(poll_readable(&afd_a))); assert_err!(futures::executor::block_on(poll_writable(&afd_a))); } } #[tokio::test] #[cfg_attr(miri, ignore)] // No socket in miri. #[cfg(any(target_os = "linux", target_os = "android"))] async fn priority_event_on_oob_data() { use std::net::SocketAddr; use tokio::io::Interest; let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); let listener = std::net::TcpListener::bind(addr).unwrap(); let addr = listener.local_addr().unwrap(); let client = std::net::TcpStream::connect(addr).unwrap(); let client = AsyncFd::with_interest(client, Interest::PRIORITY).unwrap(); let (stream, _) = listener.accept().unwrap(); // Sending out of band data should trigger priority event. send_oob_data(&stream, b"hello").unwrap(); let _ = client.ready(Interest::PRIORITY).await.unwrap(); } #[cfg(any(target_os = "linux", target_os = "android"))] fn send_oob_data(stream: &S, data: &[u8]) -> io::Result { unsafe { let res = libc::send( stream.as_raw_fd(), data.as_ptr().cast(), data.len(), libc::MSG_OOB, ); if res == -1 { Err(io::Error::last_os_error()) } else { Ok(res as usize) } } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn clear_ready_matching_clears_ready() { use tokio::io::{Interest, Ready}; let (a, mut b) = socketpair(); let afd_a = AsyncFd::new(a).unwrap(); b.write_all(b"0").unwrap(); let mut guard = afd_a .ready(Interest::READABLE | Interest::WRITABLE) .await .unwrap(); assert_eq!(guard.ready(), Ready::READABLE | Ready::WRITABLE); guard.clear_ready_matching(Ready::READABLE); assert_eq!(guard.ready(), Ready::WRITABLE); guard.clear_ready_matching(Ready::WRITABLE); assert_eq!(guard.ready(), Ready::EMPTY); } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn clear_ready_matching_clears_ready_mut() { use tokio::io::{Interest, Ready}; let (a, mut b) = socketpair(); let mut afd_a = AsyncFd::new(a).unwrap(); b.write_all(b"0").unwrap(); let mut guard = afd_a .ready_mut(Interest::READABLE | Interest::WRITABLE) .await .unwrap(); assert_eq!(guard.ready(), Ready::READABLE | Ready::WRITABLE); guard.clear_ready_matching(Ready::READABLE); assert_eq!(guard.ready(), Ready::WRITABLE); guard.clear_ready_matching(Ready::WRITABLE); assert_eq!(guard.ready(), Ready::EMPTY); } #[tokio::test] #[cfg_attr(miri, ignore)] // No socket in miri. #[cfg(target_os = "linux")] async fn await_error_readiness_timestamping() { use std::net::{Ipv4Addr, SocketAddr}; use tokio::io::{Interest, Ready}; let address_a = SocketAddr::from((Ipv4Addr::LOCALHOST, 0)); let address_b = SocketAddr::from((Ipv4Addr::LOCALHOST, 0)); let socket = std::net::UdpSocket::bind(address_a).unwrap(); socket.set_nonblocking(true).unwrap(); // configure send timestamps configure_timestamping_socket(&socket).unwrap(); socket.connect(address_b).unwrap(); let fd = AsyncFd::new(socket).unwrap(); tokio::select! { _ = fd.ready(Interest::ERROR) => panic!(), _ = tokio::time::sleep(Duration::from_millis(10)) => {} } let buf = b"hello there"; fd.get_ref().send(buf).unwrap(); // the send timestamp should now be in the error queue let guard = fd.ready(Interest::ERROR).await.unwrap(); assert_eq!(guard.ready(), Ready::ERROR); } #[cfg(target_os = "linux")] fn configure_timestamping_socket(udp_socket: &std::net::UdpSocket) -> std::io::Result { // enable software timestamping, and specifically software send timestamping let options = libc::SOF_TIMESTAMPING_SOFTWARE | libc::SOF_TIMESTAMPING_TX_SOFTWARE; let res = unsafe { libc::setsockopt( udp_socket.as_raw_fd(), libc::SOL_SOCKET, libc::SO_TIMESTAMP, &options as *const _ as *const libc::c_void, std::mem::size_of_val(&options) as libc::socklen_t, ) }; if res == -1 { Err(std::io::Error::last_os_error()) } else { Ok(res) } } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. #[cfg(target_os = "linux")] async fn await_error_readiness_invalid_address() { use std::net::{Ipv4Addr, SocketAddr}; use tokio::io::{Interest, Ready}; let socket_addr = SocketAddr::from((Ipv4Addr::LOCALHOST, 0)); let socket = std::net::UdpSocket::bind(socket_addr).unwrap(); let socket_fd = socket.as_raw_fd(); // Enable IP_RECVERR option to receive error messages // https://man7.org/linux/man-pages/man7/ip.7.html has some extra information let recv_err: libc::c_int = 1; unsafe { let res = libc::setsockopt( socket.as_raw_fd(), libc::SOL_IP, libc::IP_RECVERR, &recv_err as *const _ as *const libc::c_void, std::mem::size_of_val(&recv_err) as libc::socklen_t, ); if res == -1 { panic!("{:?}", std::io::Error::last_os_error()); } } // Spawn a separate thread for sending messages tokio::spawn(async move { // Set the destination address. This address is invalid in this context. the OS will notice // that nobody is listening on port this port. Normally this is ignored (UDP is "fire and forget"), // but because IP_RECVERR is enabled, the error will actually be reported to the sending socket let mut dest_addr = unsafe { std::mem::MaybeUninit::::zeroed().assume_init() }; dest_addr.sin_family = libc::AF_INET as _; // based on https://en.wikipedia.org/wiki/Ephemeral_port, we should pick a port number // below 1024 to guarantee that other tests don't select this port by accident when they // use port 0 to select an ephemeral port. dest_addr.sin_port = 512u16.to_be(); // Destination port // Prepare the message data let message = "Hello, Socket!"; // Prepare the message structure for sendmsg let mut iov = libc::iovec { iov_base: message.as_ptr() as *mut libc::c_void, iov_len: message.len(), }; // Prepare the destination address for the sendmsg call let dest_sockaddr: *const libc::sockaddr = &dest_addr as *const _ as *const libc::sockaddr; let dest_addrlen: libc::socklen_t = std::mem::size_of_val(&dest_addr) as libc::socklen_t; let mut msg: libc::msghdr = unsafe { std::mem::MaybeUninit::zeroed().assume_init() }; msg.msg_name = dest_sockaddr as *mut libc::c_void; msg.msg_namelen = dest_addrlen; msg.msg_iov = &mut iov; msg.msg_iovlen = 1; if unsafe { libc::sendmsg(socket_fd, &msg, 0) } == -1 { panic!("{:?}", std::io::Error::last_os_error()) } }); let fd = AsyncFd::new(socket).unwrap(); let guard = fd.ready(Interest::ERROR).await.unwrap(); assert_eq!(guard.ready(), Ready::ERROR); } #[derive(Debug, PartialEq, Eq)] struct InvalidSource; impl AsRawFd for InvalidSource { fn as_raw_fd(&self) -> RawFd { -1 } } #[tokio::test] async fn try_new() { let original = Arc::new(InvalidSource); let error = AsyncFd::try_new(original.clone()).unwrap_err(); let (returned, _cause) = error.into_parts(); assert!(Arc::ptr_eq(&original, &returned)); } #[tokio::test] async fn try_with_interest() { let original = Arc::new(InvalidSource); let error = AsyncFd::try_with_interest(original.clone(), Interest::READABLE).unwrap_err(); let (returned, _cause) = error.into_parts(); assert!(Arc::ptr_eq(&original, &returned)); } tokio-1.43.1/tests/io_async_read.rs000064400000000000000000000002561046102023000153510ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::AsyncRead; #[test] fn assert_obj_safe() { fn _assert() {} _assert::>(); } tokio-1.43.1/tests/io_buf_reader.rs000064400000000000000000000274371046102023000153510ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] // https://github.com/rust-lang/futures-rs/blob/1803948ff091b4eabf7f3bf39e16bbbdefca5cc8/futures/tests/io_buf_reader.rs use futures::task::{noop_waker_ref, Context, Poll}; use std::cmp; use std::io::{self, Cursor}; use std::pin::Pin; use tokio::io::{ AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, AsyncWriteExt, BufReader, ReadBuf, SeekFrom, }; use tokio_test::task::spawn; use tokio_test::{assert_pending, assert_ready}; macro_rules! run_fill_buf { ($reader:expr) => {{ let mut cx = Context::from_waker(noop_waker_ref()); loop { if let Poll::Ready(x) = Pin::new(&mut $reader).poll_fill_buf(&mut cx) { break x; } } }}; } struct MaybePending<'a> { inner: &'a [u8], ready_read: bool, ready_fill_buf: bool, } impl<'a> MaybePending<'a> { fn new(inner: &'a [u8]) -> Self { Self { inner, ready_read: false, ready_fill_buf: false, } } } impl AsyncRead for MaybePending<'_> { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { if self.ready_read { self.ready_read = false; Pin::new(&mut self.inner).poll_read(cx, buf) } else { self.ready_read = true; cx.waker().wake_by_ref(); Poll::Pending } } } impl AsyncBufRead for MaybePending<'_> { fn poll_fill_buf(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { if self.ready_fill_buf { self.ready_fill_buf = false; if self.inner.is_empty() { return Poll::Ready(Ok(&[])); } let len = cmp::min(2, self.inner.len()); Poll::Ready(Ok(&self.inner[0..len])) } else { self.ready_fill_buf = true; Poll::Pending } } fn consume(mut self: Pin<&mut Self>, amt: usize) { self.inner = &self.inner[amt..]; } } #[tokio::test] async fn test_buffered_reader() { let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4]; let mut reader = BufReader::with_capacity(2, inner); let mut buf = [0, 0, 0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 3); assert_eq!(buf, [5, 6, 7]); assert_eq!(reader.buffer(), []); let mut buf = [0, 0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 2); assert_eq!(buf, [0, 1]); assert_eq!(reader.buffer(), []); let mut buf = [0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 1); assert_eq!(buf, [2]); assert_eq!(reader.buffer(), [3]); let mut buf = [0, 0, 0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 1); assert_eq!(buf, [3, 0, 0]); assert_eq!(reader.buffer(), []); let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 1); assert_eq!(buf, [4, 0, 0]); assert_eq!(reader.buffer(), []); assert_eq!(reader.read(&mut buf).await.unwrap(), 0); } #[tokio::test] async fn test_buffered_reader_seek() { let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4]; let mut reader = BufReader::with_capacity(2, Cursor::new(inner)); assert_eq!(reader.seek(SeekFrom::Start(3)).await.unwrap(), 3); assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]); assert!(reader.seek(SeekFrom::Current(i64::MIN)).await.is_err()); assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]); assert_eq!(reader.seek(SeekFrom::Current(1)).await.unwrap(), 4); assert_eq!(run_fill_buf!(reader).unwrap(), &[1, 2][..]); Pin::new(&mut reader).consume(1); assert_eq!(reader.seek(SeekFrom::Current(-2)).await.unwrap(), 3); } #[tokio::test] async fn test_buffered_reader_seek_underflow() { // gimmick reader that yields its position modulo 256 for each byte struct PositionReader { pos: u64, } impl AsyncRead for PositionReader { fn poll_read( mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let b = buf.initialize_unfilled(); let len = b.len(); for x in b { *x = self.pos as u8; self.pos = self.pos.wrapping_add(1); } buf.advance(len); Poll::Ready(Ok(())) } } impl AsyncSeek for PositionReader { fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { match pos { SeekFrom::Start(n) => { self.pos = n; } SeekFrom::Current(n) => { self.pos = self.pos.wrapping_add(n as u64); } SeekFrom::End(n) => { self.pos = u64::MAX.wrapping_add(n as u64); } } Ok(()) } fn poll_complete(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(self.pos)) } } let mut reader = BufReader::with_capacity(5, PositionReader { pos: 0 }); assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1, 2, 3, 4][..]); assert_eq!(reader.seek(SeekFrom::End(-5)).await.unwrap(), u64::MAX - 5); assert_eq!(run_fill_buf!(reader).unwrap().len(), 5); // the following seek will require two underlying seeks let expected = 9_223_372_036_854_775_802; assert_eq!( reader.seek(SeekFrom::Current(i64::MIN)).await.unwrap(), expected ); assert_eq!(run_fill_buf!(reader).unwrap().len(), 5); // seeking to 0 should empty the buffer. assert_eq!(reader.seek(SeekFrom::Current(0)).await.unwrap(), expected); assert_eq!(reader.get_ref().pos, expected); } #[tokio::test] async fn test_short_reads() { /// A dummy reader intended at testing short-reads propagation. struct ShortReader { lengths: Vec, } impl AsyncRead for ShortReader { fn poll_read( mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { if !self.lengths.is_empty() { buf.advance(self.lengths.remove(0)); } Poll::Ready(Ok(())) } } let inner = ShortReader { lengths: vec![0, 1, 2, 0, 1, 0], }; let mut reader = BufReader::new(inner); let mut buf = [0, 0]; assert_eq!(reader.read(&mut buf).await.unwrap(), 0); assert_eq!(reader.read(&mut buf).await.unwrap(), 1); assert_eq!(reader.read(&mut buf).await.unwrap(), 2); assert_eq!(reader.read(&mut buf).await.unwrap(), 0); assert_eq!(reader.read(&mut buf).await.unwrap(), 1); assert_eq!(reader.read(&mut buf).await.unwrap(), 0); assert_eq!(reader.read(&mut buf).await.unwrap(), 0); } #[tokio::test] async fn maybe_pending() { let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4]; let mut reader = BufReader::with_capacity(2, MaybePending::new(inner)); let mut buf = [0, 0, 0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 3); assert_eq!(buf, [5, 6, 7]); assert_eq!(reader.buffer(), []); let mut buf = [0, 0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 2); assert_eq!(buf, [0, 1]); assert_eq!(reader.buffer(), []); let mut buf = [0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 1); assert_eq!(buf, [2]); assert_eq!(reader.buffer(), [3]); let mut buf = [0, 0, 0]; let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 1); assert_eq!(buf, [3, 0, 0]); assert_eq!(reader.buffer(), []); let nread = reader.read(&mut buf).await.unwrap(); assert_eq!(nread, 1); assert_eq!(buf, [4, 0, 0]); assert_eq!(reader.buffer(), []); assert_eq!(reader.read(&mut buf).await.unwrap(), 0); } #[tokio::test] async fn maybe_pending_buf_read() { let inner = MaybePending::new(&[0, 1, 2, 3, 1, 0]); let mut reader = BufReader::with_capacity(2, inner); let mut v = Vec::new(); reader.read_until(3, &mut v).await.unwrap(); assert_eq!(v, [0, 1, 2, 3]); v.clear(); reader.read_until(1, &mut v).await.unwrap(); assert_eq!(v, [1]); v.clear(); reader.read_until(8, &mut v).await.unwrap(); assert_eq!(v, [0]); v.clear(); reader.read_until(9, &mut v).await.unwrap(); assert_eq!(v, []); } // https://github.com/rust-lang/futures-rs/pull/1573#discussion_r281162309 #[tokio::test] async fn maybe_pending_seek() { struct MaybePendingSeek<'a> { inner: Cursor<&'a [u8]>, ready: bool, seek_res: Option>, } impl<'a> MaybePendingSeek<'a> { fn new(inner: &'a [u8]) -> Self { Self { inner: Cursor::new(inner), ready: true, seek_res: None, } } } impl AsyncRead for MaybePendingSeek<'_> { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut self.inner).poll_read(cx, buf) } } impl AsyncBufRead for MaybePendingSeek<'_> { fn poll_fill_buf( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { let this: *mut Self = &mut *self as *mut _; Pin::new(&mut unsafe { &mut *this }.inner).poll_fill_buf(cx) } fn consume(mut self: Pin<&mut Self>, amt: usize) { Pin::new(&mut self.inner).consume(amt) } } impl AsyncSeek for MaybePendingSeek<'_> { fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { self.seek_res = Some(Pin::new(&mut self.inner).start_seek(pos)); Ok(()) } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.ready { self.ready = false; self.seek_res.take().unwrap_or(Ok(()))?; Pin::new(&mut self.inner).poll_complete(cx) } else { self.ready = true; cx.waker().wake_by_ref(); Poll::Pending } } } let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4]; let mut reader = BufReader::with_capacity(2, MaybePendingSeek::new(inner)); assert_eq!(reader.seek(SeekFrom::Current(3)).await.unwrap(), 3); assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]); assert!(reader.seek(SeekFrom::Current(i64::MIN)).await.is_err()); assert_eq!(run_fill_buf!(reader).unwrap(), &[0, 1][..]); assert_eq!(reader.seek(SeekFrom::Current(1)).await.unwrap(), 4); assert_eq!(run_fill_buf!(reader).unwrap(), &[1, 2][..]); Pin::new(&mut reader).consume(1); assert_eq!(reader.seek(SeekFrom::Current(-2)).await.unwrap(), 3); } // This tests the AsyncBufReadExt::fill_buf wrapper. #[tokio::test] async fn test_fill_buf_wrapper() { let (mut write, read) = tokio::io::duplex(16); let mut read = BufReader::new(read); write.write_all(b"hello world").await.unwrap(); assert_eq!(read.fill_buf().await.unwrap(), b"hello world"); read.consume(b"hello ".len()); assert_eq!(read.fill_buf().await.unwrap(), b"world"); assert_eq!(read.fill_buf().await.unwrap(), b"world"); read.consume(b"world".len()); let mut fill = spawn(read.fill_buf()); assert_pending!(fill.poll()); write.write_all(b"foo bar").await.unwrap(); assert_eq!(assert_ready!(fill.poll()).unwrap(), b"foo bar"); drop(fill); drop(write); assert_eq!(read.fill_buf().await.unwrap(), b"foo bar"); read.consume(b"foo bar".len()); assert_eq!(read.fill_buf().await.unwrap(), b""); } tokio-1.43.1/tests/io_buf_writer.rs000064400000000000000000000375061046102023000154210ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] // https://github.com/rust-lang/futures-rs/blob/1803948ff091b4eabf7f3bf39e16bbbdefca5cc8/futures/tests/io_buf_writer.rs use futures::task::{Context, Poll}; use std::io::{self, Cursor}; use std::pin::Pin; use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, BufWriter, SeekFrom}; use futures::future; use tokio_test::assert_ok; use std::cmp; use std::io::IoSlice; mod support { pub(crate) mod io_vec; } use support::io_vec::IoBufs; struct MaybePending { inner: Vec, ready: bool, } impl MaybePending { fn new(inner: Vec) -> Self { Self { inner, ready: false, } } } impl AsyncWrite for MaybePending { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { if self.ready { self.ready = false; Pin::new(&mut self.inner).poll_write(cx, buf) } else { self.ready = true; cx.waker().wake_by_ref(); Poll::Pending } } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_shutdown(cx) } } async fn write_vectored(writer: &mut W, bufs: &[IoSlice<'_>]) -> io::Result where W: AsyncWrite + Unpin, { let mut writer = Pin::new(writer); future::poll_fn(|cx| writer.as_mut().poll_write_vectored(cx, bufs)).await } #[tokio::test] async fn buf_writer() { let mut writer = BufWriter::with_capacity(2, Vec::new()); assert_eq!(writer.write(&[0, 1]).await.unwrap(), 2); assert_eq!(writer.buffer(), []); assert_eq!(*writer.get_ref(), [0, 1]); assert_eq!(writer.write(&[2]).await.unwrap(), 1); assert_eq!(writer.buffer(), [2]); assert_eq!(*writer.get_ref(), [0, 1]); assert_eq!(writer.write(&[3]).await.unwrap(), 1); assert_eq!(writer.buffer(), [2, 3]); assert_eq!(*writer.get_ref(), [0, 1]); writer.flush().await.unwrap(); assert_eq!(writer.buffer(), []); assert_eq!(*writer.get_ref(), [0, 1, 2, 3]); assert_eq!(writer.write(&[4]).await.unwrap(), 1); assert_eq!(writer.write(&[5]).await.unwrap(), 1); assert_eq!(writer.buffer(), [4, 5]); assert_eq!(*writer.get_ref(), [0, 1, 2, 3]); assert_eq!(writer.write(&[6]).await.unwrap(), 1); assert_eq!(writer.buffer(), [6]); assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5]); assert_eq!(writer.write(&[7, 8]).await.unwrap(), 2); assert_eq!(writer.buffer(), []); assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8]); assert_eq!(writer.write(&[9, 10, 11]).await.unwrap(), 3); assert_eq!(writer.buffer(), []); assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]); writer.flush().await.unwrap(); assert_eq!(writer.buffer(), []); assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]); } #[tokio::test] async fn buf_writer_inner_flushes() { let mut w = BufWriter::with_capacity(3, Vec::new()); assert_eq!(w.write(&[0, 1]).await.unwrap(), 2); assert_eq!(*w.get_ref(), []); w.flush().await.unwrap(); let w = w.into_inner(); assert_eq!(w, [0, 1]); } #[tokio::test] async fn buf_writer_seek() { let mut w = BufWriter::with_capacity(3, Cursor::new(Vec::new())); w.write_all(&[0, 1, 2, 3, 4, 5]).await.unwrap(); w.write_all(&[6, 7]).await.unwrap(); assert_eq!(w.seek(SeekFrom::Current(0)).await.unwrap(), 8); assert_eq!(&w.get_ref().get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..]); assert_eq!(w.seek(SeekFrom::Start(2)).await.unwrap(), 2); w.write_all(&[8, 9]).await.unwrap(); w.flush().await.unwrap(); assert_eq!(&w.into_inner().into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7]); } #[tokio::test] async fn maybe_pending_buf_writer() { let mut writer = BufWriter::with_capacity(2, MaybePending::new(Vec::new())); assert_eq!(writer.write(&[0, 1]).await.unwrap(), 2); assert_eq!(writer.buffer(), []); assert_eq!(&writer.get_ref().inner, &[0, 1]); assert_eq!(writer.write(&[2]).await.unwrap(), 1); assert_eq!(writer.buffer(), [2]); assert_eq!(&writer.get_ref().inner, &[0, 1]); assert_eq!(writer.write(&[3]).await.unwrap(), 1); assert_eq!(writer.buffer(), [2, 3]); assert_eq!(&writer.get_ref().inner, &[0, 1]); writer.flush().await.unwrap(); assert_eq!(writer.buffer(), []); assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3]); assert_eq!(writer.write(&[4]).await.unwrap(), 1); assert_eq!(writer.write(&[5]).await.unwrap(), 1); assert_eq!(writer.buffer(), [4, 5]); assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3]); assert_eq!(writer.write(&[6]).await.unwrap(), 1); assert_eq!(writer.buffer(), [6]); assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5]); assert_eq!(writer.write(&[7, 8]).await.unwrap(), 2); assert_eq!(writer.buffer(), []); assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8]); assert_eq!(writer.write(&[9, 10, 11]).await.unwrap(), 3); assert_eq!(writer.buffer(), []); assert_eq!( writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] ); writer.flush().await.unwrap(); assert_eq!(writer.buffer(), []); assert_eq!( &writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] ); } #[tokio::test] async fn maybe_pending_buf_writer_inner_flushes() { let mut w = BufWriter::with_capacity(3, MaybePending::new(Vec::new())); assert_eq!(w.write(&[0, 1]).await.unwrap(), 2); assert_eq!(&w.get_ref().inner, &[]); w.flush().await.unwrap(); let w = w.into_inner().inner; assert_eq!(w, [0, 1]); } #[tokio::test] async fn maybe_pending_buf_writer_seek() { struct MaybePendingSeek { inner: Cursor>, ready_write: bool, ready_seek: bool, seek_res: Option>, } impl MaybePendingSeek { fn new(inner: Vec) -> Self { Self { inner: Cursor::new(inner), ready_write: false, ready_seek: false, seek_res: None, } } } impl AsyncWrite for MaybePendingSeek { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { if self.ready_write { self.ready_write = false; Pin::new(&mut self.inner).poll_write(cx, buf) } else { self.ready_write = true; cx.waker().wake_by_ref(); Poll::Pending } } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.inner).poll_shutdown(cx) } } impl AsyncSeek for MaybePendingSeek { fn start_seek(mut self: Pin<&mut Self>, pos: SeekFrom) -> io::Result<()> { self.seek_res = Some(Pin::new(&mut self.inner).start_seek(pos)); Ok(()) } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.ready_seek { self.ready_seek = false; self.seek_res.take().unwrap_or(Ok(()))?; Pin::new(&mut self.inner).poll_complete(cx) } else { self.ready_seek = true; cx.waker().wake_by_ref(); Poll::Pending } } } let mut w = BufWriter::with_capacity(3, MaybePendingSeek::new(Vec::new())); w.write_all(&[0, 1, 2, 3, 4, 5]).await.unwrap(); w.write_all(&[6, 7]).await.unwrap(); assert_eq!(w.seek(SeekFrom::Current(0)).await.unwrap(), 8); assert_eq!( &w.get_ref().inner.get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..] ); assert_eq!(w.seek(SeekFrom::Start(2)).await.unwrap(), 2); w.write_all(&[8, 9]).await.unwrap(); w.flush().await.unwrap(); assert_eq!( &w.into_inner().inner.into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7] ); } struct MockWriter { data: Vec, write_len: usize, vectored: bool, } impl MockWriter { fn new(write_len: usize) -> Self { MockWriter { data: Vec::new(), write_len, vectored: false, } } fn vectored(write_len: usize) -> Self { MockWriter { data: Vec::new(), write_len, vectored: true, } } fn write_up_to(&mut self, buf: &[u8], limit: usize) -> usize { let len = cmp::min(buf.len(), limit); self.data.extend_from_slice(&buf[..len]); len } } impl AsyncWrite for MockWriter { fn poll_write( self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8], ) -> Poll> { let this = self.get_mut(); let n = this.write_up_to(buf, this.write_len); Ok(n).into() } fn poll_write_vectored( self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { let this = self.get_mut(); let mut total_written = 0; for buf in bufs { let n = this.write_up_to(buf, this.write_len - total_written); total_written += n; if total_written == this.write_len { break; } } Ok(total_written).into() } fn is_write_vectored(&self) -> bool { self.vectored } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Ok(()).into() } } #[tokio::test] async fn write_vectored_empty_on_non_vectored() { let mut w = BufWriter::new(MockWriter::new(4)); let n = assert_ok!(write_vectored(&mut w, &[]).await); assert_eq!(n, 0); let io_vec = [IoSlice::new(&[]); 3]; let n = assert_ok!(write_vectored(&mut w, &io_vec).await); assert_eq!(n, 0); assert_ok!(w.flush().await); assert!(w.get_ref().data.is_empty()); } #[tokio::test] async fn write_vectored_empty_on_vectored() { let mut w = BufWriter::new(MockWriter::vectored(4)); let n = assert_ok!(write_vectored(&mut w, &[]).await); assert_eq!(n, 0); let io_vec = [IoSlice::new(&[]); 3]; let n = assert_ok!(write_vectored(&mut w, &io_vec).await); assert_eq!(n, 0); assert_ok!(w.flush().await); assert!(w.get_ref().data.is_empty()); } #[tokio::test] async fn write_vectored_basic_on_non_vectored() { let msg = b"foo bar baz"; let bufs = [ IoSlice::new(&msg[0..4]), IoSlice::new(&msg[4..8]), IoSlice::new(&msg[8..]), ]; let mut w = BufWriter::new(MockWriter::new(4)); let n = assert_ok!(write_vectored(&mut w, &bufs).await); assert_eq!(n, msg.len()); assert!(w.buffer() == &msg[..]); assert_ok!(w.flush().await); assert_eq!(w.get_ref().data, msg); } #[tokio::test] async fn write_vectored_basic_on_vectored() { let msg = b"foo bar baz"; let bufs = [ IoSlice::new(&msg[0..4]), IoSlice::new(&msg[4..8]), IoSlice::new(&msg[8..]), ]; let mut w = BufWriter::new(MockWriter::vectored(4)); let n = assert_ok!(write_vectored(&mut w, &bufs).await); assert_eq!(n, msg.len()); assert!(w.buffer() == &msg[..]); assert_ok!(w.flush().await); assert_eq!(w.get_ref().data, msg); } #[tokio::test] async fn write_vectored_large_total_on_non_vectored() { let msg = b"foo bar baz"; let mut bufs = [ IoSlice::new(&msg[0..4]), IoSlice::new(&msg[4..8]), IoSlice::new(&msg[8..]), ]; let io_vec = IoBufs::new(&mut bufs); let mut w = BufWriter::with_capacity(8, MockWriter::new(4)); let n = assert_ok!(write_vectored(&mut w, &io_vec).await); assert_eq!(n, 8); assert!(w.buffer() == &msg[..8]); let io_vec = io_vec.advance(n); let n = assert_ok!(write_vectored(&mut w, &io_vec).await); assert_eq!(n, 3); assert!(w.get_ref().data.as_slice() == &msg[..8]); assert!(w.buffer() == &msg[8..]); } #[tokio::test] async fn write_vectored_large_total_on_vectored() { let msg = b"foo bar baz"; let mut bufs = [ IoSlice::new(&msg[0..4]), IoSlice::new(&msg[4..8]), IoSlice::new(&msg[8..]), ]; let io_vec = IoBufs::new(&mut bufs); let mut w = BufWriter::with_capacity(8, MockWriter::vectored(10)); let n = assert_ok!(write_vectored(&mut w, &io_vec).await); assert_eq!(n, 10); assert!(w.buffer().is_empty()); let io_vec = io_vec.advance(n); let n = assert_ok!(write_vectored(&mut w, &io_vec).await); assert_eq!(n, 1); assert!(w.get_ref().data.as_slice() == &msg[..10]); assert!(w.buffer() == &msg[10..]); } struct VectoredWriteHarness { writer: BufWriter, buf_capacity: usize, } impl VectoredWriteHarness { fn new(buf_capacity: usize) -> Self { VectoredWriteHarness { writer: BufWriter::with_capacity(buf_capacity, MockWriter::new(4)), buf_capacity, } } fn with_vectored_backend(buf_capacity: usize) -> Self { VectoredWriteHarness { writer: BufWriter::with_capacity(buf_capacity, MockWriter::vectored(4)), buf_capacity, } } async fn write_all<'a, 'b>(&mut self, mut io_vec: IoBufs<'a, 'b>) -> usize { let mut total_written = 0; while !io_vec.is_empty() { let n = assert_ok!(write_vectored(&mut self.writer, &io_vec).await); assert!(n != 0); assert!(self.writer.buffer().len() <= self.buf_capacity); total_written += n; io_vec = io_vec.advance(n); } total_written } async fn flush(&mut self) -> &[u8] { assert_ok!(self.writer.flush().await); &self.writer.get_ref().data } } #[tokio::test] async fn write_vectored_odd_on_non_vectored() { let msg = b"foo bar baz"; let mut bufs = [ IoSlice::new(&msg[0..4]), IoSlice::new(&[]), IoSlice::new(&msg[4..9]), IoSlice::new(&msg[9..]), ]; let mut h = VectoredWriteHarness::new(8); let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; assert_eq!(bytes_written, msg.len()); assert_eq!(h.flush().await, msg); } #[tokio::test] async fn write_vectored_odd_on_vectored() { let msg = b"foo bar baz"; let mut bufs = [ IoSlice::new(&msg[0..4]), IoSlice::new(&[]), IoSlice::new(&msg[4..9]), IoSlice::new(&msg[9..]), ]; let mut h = VectoredWriteHarness::with_vectored_backend(8); let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; assert_eq!(bytes_written, msg.len()); assert_eq!(h.flush().await, msg); } #[tokio::test] async fn write_vectored_large_slice_on_non_vectored() { let msg = b"foo bar baz"; let mut bufs = [ IoSlice::new(&[]), IoSlice::new(&msg[..9]), IoSlice::new(&msg[9..]), ]; let mut h = VectoredWriteHarness::new(8); let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; assert_eq!(bytes_written, msg.len()); assert_eq!(h.flush().await, msg); } #[tokio::test] async fn write_vectored_large_slice_on_vectored() { let msg = b"foo bar baz"; let mut bufs = [ IoSlice::new(&[]), IoSlice::new(&msg[..9]), IoSlice::new(&msg[9..]), ]; let mut h = VectoredWriteHarness::with_vectored_backend(8); let bytes_written = h.write_all(IoBufs::new(&mut bufs)).await; assert_eq!(bytes_written, msg.len()); assert_eq!(h.flush().await, msg); } tokio-1.43.1/tests/io_chain.rs000064400000000000000000000005471046102023000143260ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::AsyncReadExt; use tokio_test::assert_ok; #[tokio::test] async fn chain() { let mut buf = Vec::new(); let rd1: &[u8] = b"hello "; let rd2: &[u8] = b"world"; let mut rd = rd1.chain(rd2); assert_ok!(rd.read_to_end(&mut buf).await); assert_eq!(buf, b"hello world"); } tokio-1.43.1/tests/io_copy.rs000064400000000000000000000050531046102023000142130ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use bytes::BytesMut; use tokio::io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio_test::assert_ok; use std::pin::Pin; use std::task::{ready, Context, Poll}; #[tokio::test] async fn copy() { struct Rd(bool); impl AsyncRead for Rd { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { if self.0 { buf.put_slice(b"hello world"); self.0 = false; Poll::Ready(Ok(())) } else { Poll::Ready(Ok(())) } } } let mut rd = Rd(true); let mut wr = Vec::new(); let n = assert_ok!(io::copy(&mut rd, &mut wr).await); assert_eq!(n, 11); assert_eq!(wr, b"hello world"); } #[tokio::test] async fn proxy() { struct BufferedWd { buf: BytesMut, writer: io::DuplexStream, } impl AsyncWrite for BufferedWd { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.get_mut().buf.extend_from_slice(buf); Poll::Ready(Ok(buf.len())) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); while !this.buf.is_empty() { let n = ready!(Pin::new(&mut this.writer).poll_write(cx, &this.buf))?; let _ = this.buf.split_to(n); } Pin::new(&mut this.writer).poll_flush(cx) } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.writer).poll_shutdown(cx) } } let (rd, wd) = io::duplex(1024); let mut rd = rd.take(1024); let mut wd = BufferedWd { buf: BytesMut::new(), writer: wd, }; // write start bytes assert_ok!(wd.write_all(&[0x42; 512]).await); assert_ok!(wd.flush().await); let n = assert_ok!(io::copy(&mut rd, &mut wd).await); assert_eq!(n, 1024); } #[tokio::test] async fn copy_is_cooperative() { tokio::select! { biased; _ = async { loop { let mut reader: &[u8] = b"hello"; let mut writer: Vec = vec![]; let _ = io::copy(&mut reader, &mut writer).await; } } => {}, _ = tokio::task::yield_now() => {} } } tokio-1.43.1/tests/io_copy_bidirectional.rs000064400000000000000000000113241046102023000171010ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support bind() use std::time::Duration; use tokio::io::{self, copy_bidirectional, AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; use tokio::task::JoinHandle; async fn make_socketpair() -> (TcpStream, TcpStream) { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let connector = TcpStream::connect(addr); let acceptor = listener.accept(); let (c1, c2) = tokio::join!(connector, acceptor); (c1.unwrap(), c2.unwrap().0) } async fn block_write(s: &mut TcpStream) -> usize { static BUF: [u8; 2048] = [0; 2048]; let mut copied = 0; loop { tokio::select! { result = s.write(&BUF) => { copied += result.expect("write error") }, _ = tokio::time::sleep(Duration::from_millis(10)) => { break; } } } copied } async fn symmetric(mut cb: F) where F: FnMut(JoinHandle>, TcpStream, TcpStream) -> Fut, Fut: std::future::Future, { // We run the test twice, with streams passed to copy_bidirectional in // different orders, in order to ensure that the two arguments are // interchangeable. let (a, mut a1) = make_socketpair().await; let (b, mut b1) = make_socketpair().await; let handle = tokio::spawn(async move { copy_bidirectional(&mut a1, &mut b1).await }); cb(handle, a, b).await; let (a, mut a1) = make_socketpair().await; let (b, mut b1) = make_socketpair().await; let handle = tokio::spawn(async move { copy_bidirectional(&mut b1, &mut a1).await }); cb(handle, b, a).await; } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn test_basic_transfer() { symmetric(|_handle, mut a, mut b| async move { a.write_all(b"test").await.unwrap(); let mut tmp = [0; 4]; b.read_exact(&mut tmp).await.unwrap(); assert_eq!(&tmp[..], b"test"); }) .await } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn test_transfer_after_close() { symmetric(|handle, mut a, mut b| async move { AsyncWriteExt::shutdown(&mut a).await.unwrap(); b.read_to_end(&mut Vec::new()).await.unwrap(); b.write_all(b"quux").await.unwrap(); let mut tmp = [0; 4]; a.read_exact(&mut tmp).await.unwrap(); assert_eq!(&tmp[..], b"quux"); // Once both are closed, we should have our handle back drop(b); assert_eq!(handle.await.unwrap().unwrap(), (0, 4)); }) .await } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn blocking_one_side_does_not_block_other() { symmetric(|handle, mut a, mut b| async move { block_write(&mut a).await; b.write_all(b"quux").await.unwrap(); let mut tmp = [0; 4]; a.read_exact(&mut tmp).await.unwrap(); assert_eq!(&tmp[..], b"quux"); AsyncWriteExt::shutdown(&mut a).await.unwrap(); let mut buf = Vec::new(); b.read_to_end(&mut buf).await.unwrap(); drop(b); assert_eq!(handle.await.unwrap().unwrap(), (buf.len() as u64, 4)); }) .await } #[tokio::test] async fn immediate_exit_on_write_error() { let payload = b"here, take this"; let error = || io::Error::new(io::ErrorKind::Other, "no thanks!"); let mut a = tokio_test::io::Builder::new() .read(payload) .write_error(error()) .build(); let mut b = tokio_test::io::Builder::new() .read(payload) .write_error(error()) .build(); assert!(copy_bidirectional(&mut a, &mut b).await.is_err()); } #[tokio::test] async fn immediate_exit_on_read_error() { let error = || io::Error::new(io::ErrorKind::Other, "got nothing!"); let mut a = tokio_test::io::Builder::new().read_error(error()).build(); let mut b = tokio_test::io::Builder::new().read_error(error()).build(); assert!(copy_bidirectional(&mut a, &mut b).await.is_err()); } #[tokio::test] async fn copy_bidirectional_is_cooperative() { tokio::select! { biased; _ = async { loop { let payload = b"here, take this"; let mut a = tokio_test::io::Builder::new() .read(payload) .write(payload) .build(); let mut b = tokio_test::io::Builder::new() .read(payload) .write(payload) .build(); let _ = copy_bidirectional(&mut a, &mut b).await; } } => {}, _ = tokio::task::yield_now() => {} } } tokio-1.43.1/tests/io_driver.rs000064400000000000000000000055511046102023000145370ustar 00000000000000#![warn(rust_2018_idioms)] // Wasi does not support panic recovery or threading #![cfg(all(feature = "full", not(target_os = "wasi")))] use tokio::net::TcpListener; use tokio::runtime; use tokio_test::{assert_ok, assert_pending}; use futures::task::{waker_ref, ArcWake}; use std::future::Future; use std::net::TcpStream; use std::pin::Pin; use std::sync::{mpsc, Arc, Mutex}; use std::task::Context; struct Task { future: Mutex>>, } impl ArcWake for Task { fn wake_by_ref(_: &Arc) { // Do nothing... } } impl Task { fn new(future: T) -> Task { Task { future: Mutex::new(Box::pin(future)), } } } #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn test_drop_on_notify() { // When the reactor receives a kernel notification, it notifies the // task that holds the associated socket. If this notification results in // the task being dropped, the socket will also be dropped. // // Previously, there was a deadlock scenario where the reactor, while // notifying, held a lock and the task being dropped attempted to acquire // that same lock in order to clean up state. // // To simulate this case, we create a fake executor that does nothing when // the task is notified. This simulates an executor in the process of // shutting down. Then, when the task handle is dropped, the task itself is // dropped. let rt = runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); let (addr_tx, addr_rx) = mpsc::channel(); // Define a task that just drains the listener let task = Arc::new(Task::new(async move { // Create a listener let listener = assert_ok!(TcpListener::bind("127.0.0.1:0").await); // Send the address let addr = listener.local_addr().unwrap(); addr_tx.send(addr).unwrap(); loop { let _ = listener.accept().await; } })); { let _enter = rt.enter(); let waker = waker_ref(&task); let mut cx = Context::from_waker(&waker); assert_pending!(task.future.lock().unwrap().as_mut().poll(&mut cx)); } // Get the address let addr = addr_rx.recv().unwrap(); drop(task); // Establish a connection to the acceptor let _s = TcpStream::connect(addr).unwrap(); // Force the reactor to turn rt.block_on(async {}); } #[test] #[should_panic( expected = "A Tokio 1.x context was found, but IO is disabled. Call `enable_io` on the runtime builder to enable IO." )] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn panics_when_io_disabled() { let rt = runtime::Builder::new_current_thread().build().unwrap(); rt.block_on(async { let _ = tokio::net::TcpListener::from_std(std::net::TcpListener::bind("127.0.0.1:0").unwrap()); }); } tokio-1.43.1/tests/io_driver_drop.rs000064400000000000000000000024211046102023000155540ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support bind use tokio::net::TcpListener; use tokio::runtime; use tokio_test::{assert_err, assert_pending, assert_ready, task}; #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn tcp_doesnt_block() { let rt = rt(); let listener = { let _enter = rt.enter(); let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); TcpListener::from_std(listener).unwrap() }; drop(rt); let mut task = task::spawn(async move { assert_err!(listener.accept().await); }); assert_ready!(task.poll()); } #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn drop_wakes() { let rt = rt(); let listener = { let _enter = rt.enter(); let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); TcpListener::from_std(listener).unwrap() }; let mut task = task::spawn(async move { assert_err!(listener.accept().await); }); assert_pending!(task.poll()); drop(rt); assert!(task.is_woken()); assert_ready!(task.poll()); } fn rt() -> runtime::Runtime { runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/io_fill_buf.rs000064400000000000000000000015331046102023000150220ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support file operations use tempfile::NamedTempFile; use tokio::fs::File; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio_test::assert_ok; #[tokio::test] async fn fill_buf_file() { let file = NamedTempFile::new().unwrap(); assert_ok!(std::fs::write(file.path(), b"hello")); let file = assert_ok!(File::open(file.path()).await); let mut file = BufReader::new(file); let mut contents = Vec::new(); loop { let consumed = { let buffer = assert_ok!(file.fill_buf().await); if buffer.is_empty() { break; } contents.extend_from_slice(buffer); buffer.len() }; file.consume(consumed); } assert_eq!(contents, b"hello"); } tokio-1.43.1/tests/io_join.rs000064400000000000000000000035271046102023000142040ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{join, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, Join, ReadBuf}; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; struct R; impl AsyncRead for R { fn poll_read( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { buf.put_slice(&[b'z']); Poll::Ready(Ok(())) } } struct W; impl AsyncWrite for W { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8], ) -> Poll> { Poll::Ready(Ok(1)) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_write_vectored( self: Pin<&mut Self>, _cx: &mut Context<'_>, _bufs: &[io::IoSlice<'_>], ) -> Poll> { Poll::Ready(Ok(2)) } fn is_write_vectored(&self) -> bool { true } } #[test] fn is_send_and_sync() { fn assert_bound() {} assert_bound::>(); } #[test] fn method_delegation() { let mut rw = join(R, W); let mut buf = [0; 1]; tokio_test::block_on(async move { assert_eq!(1, rw.read(&mut buf).await.unwrap()); assert_eq!(b'z', buf[0]); assert_eq!(1, rw.write(&[b'x']).await.unwrap()); assert_eq!( 2, rw.write_vectored(&[io::IoSlice::new(&[b'x'])]) .await .unwrap() ); assert!(rw.is_write_vectored()); assert!(rw.flush().await.is_ok()); assert!(rw.shutdown().await.is_ok()); }); } tokio-1.43.1/tests/io_lines.rs000064400000000000000000000010251046102023000143460ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::AsyncBufReadExt; use tokio_test::assert_ok; #[tokio::test] async fn lines_inherent() { let rd: &[u8] = b"hello\r\nworld\n\n"; let mut st = rd.lines(); let b = assert_ok!(st.next_line().await).unwrap(); assert_eq!(b, "hello"); let b = assert_ok!(st.next_line().await).unwrap(); assert_eq!(b, "world"); let b = assert_ok!(st.next_line().await).unwrap(); assert_eq!(b, ""); assert!(assert_ok!(st.next_line().await).is_none()); } tokio-1.43.1/tests/io_mem_stream.rs000064400000000000000000000054431046102023000153750ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt}; #[tokio::test] async fn ping_pong() { let (mut a, mut b) = duplex(32); let mut buf = [0u8; 4]; a.write_all(b"ping").await.unwrap(); b.read_exact(&mut buf).await.unwrap(); assert_eq!(&buf, b"ping"); b.write_all(b"pong").await.unwrap(); a.read_exact(&mut buf).await.unwrap(); assert_eq!(&buf, b"pong"); } #[tokio::test] async fn across_tasks() { let (mut a, mut b) = duplex(32); let t1 = tokio::spawn(async move { a.write_all(b"ping").await.unwrap(); let mut buf = [0u8; 4]; a.read_exact(&mut buf).await.unwrap(); assert_eq!(&buf, b"pong"); }); let t2 = tokio::spawn(async move { let mut buf = [0u8; 4]; b.read_exact(&mut buf).await.unwrap(); assert_eq!(&buf, b"ping"); b.write_all(b"pong").await.unwrap(); }); t1.await.unwrap(); t2.await.unwrap(); } #[tokio::test] async fn disconnect() { let (mut a, mut b) = duplex(32); let t1 = tokio::spawn(async move { a.write_all(b"ping").await.unwrap(); // and dropped }); let t2 = tokio::spawn(async move { let mut buf = [0u8; 32]; let n = b.read(&mut buf).await.unwrap(); assert_eq!(&buf[..n], b"ping"); let n = b.read(&mut buf).await.unwrap(); assert_eq!(n, 0); }); t1.await.unwrap(); t2.await.unwrap(); } #[tokio::test] async fn disconnect_reader() { let (a, mut b) = duplex(2); let t1 = tokio::spawn(async move { // this will block, as not all data fits into duplex b.write_all(b"ping").await.unwrap_err(); }); let t2 = tokio::spawn(async move { // here we drop the reader side, and we expect the writer in the other // task to exit with an error drop(a); }); t2.await.unwrap(); t1.await.unwrap(); } #[tokio::test] async fn max_write_size() { let (mut a, mut b) = duplex(32); let t1 = tokio::spawn(async move { let n = a.write(&[0u8; 64]).await.unwrap(); assert_eq!(n, 32); let n = a.write(&[0u8; 64]).await.unwrap(); assert_eq!(n, 4); }); let mut buf = [0u8; 4]; b.read_exact(&mut buf).await.unwrap(); t1.await.unwrap(); // drop b only after task t1 finishes writing drop(b); } #[tokio::test] async fn duplex_is_cooperative() { let (mut tx, mut rx) = tokio::io::duplex(1024 * 8); tokio::select! { biased; _ = async { loop { let buf = [3u8; 4096]; tx.write_all(&buf).await.unwrap(); let mut buf = [0u8; 4096]; let _ = rx.read(&mut buf).await.unwrap(); } } => {}, _ = tokio::task::yield_now() => {} } } tokio-1.43.1/tests/io_panic.rs000064400000000000000000000130171046102023000143320ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery #![cfg(panic = "unwind")] use std::task::{Context, Poll}; use std::{error::Error, pin::Pin}; use tokio::io::{self, split, AsyncRead, AsyncWrite, ReadBuf}; mod support { pub mod panic; } use support::panic::test_panic; struct RW; impl AsyncRead for RW { fn poll_read( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { buf.put_slice(&[b'z']); Poll::Ready(Ok(())) } } impl AsyncWrite for RW { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8], ) -> Poll> { Poll::Ready(Ok(1)) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } } #[cfg(unix)] mod unix { use std::os::unix::prelude::{AsRawFd, RawFd}; pub struct MockFd; impl AsRawFd for MockFd { fn as_raw_fd(&self) -> RawFd { 0 } } } #[test] fn read_buf_initialize_unfilled_to_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let mut buffer = Vec::::new(); let mut read_buf = ReadBuf::new(&mut buffer); read_buf.initialize_unfilled_to(2); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn read_buf_advance_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let mut buffer = Vec::::new(); let mut read_buf = ReadBuf::new(&mut buffer); read_buf.advance(2); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn read_buf_set_filled_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let mut buffer = Vec::::new(); let mut read_buf = ReadBuf::new(&mut buffer); read_buf.set_filled(2); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn read_buf_put_slice_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let mut buffer = Vec::::new(); let mut read_buf = ReadBuf::new(&mut buffer); let new_slice = [0x40_u8, 0x41_u8]; read_buf.put_slice(&new_slice); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn unsplit_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let (r1, _w1) = split(RW); let (_r2, w2) = split(RW); r1.unsplit(w2); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] fn async_fd_new_panic_caller() -> Result<(), Box> { use tokio::io::unix::AsyncFd; use tokio::runtime::Builder; let panic_location_file = test_panic(|| { // Runtime without `enable_io` so it has no IO driver set. let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async { let fd = unix::MockFd; let _ = AsyncFd::new(fd); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] fn async_fd_with_interest_panic_caller() -> Result<(), Box> { use tokio::io::unix::AsyncFd; use tokio::io::Interest; use tokio::runtime::Builder; let panic_location_file = test_panic(|| { // Runtime without `enable_io` so it has no IO driver set. let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async { let fd = unix::MockFd; let _ = AsyncFd::with_interest(fd, Interest::READABLE); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] fn async_fd_try_new_panic_caller() -> Result<(), Box> { use tokio::io::unix::AsyncFd; use tokio::runtime::Builder; let panic_location_file = test_panic(|| { // Runtime without `enable_io` so it has no IO driver set. let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async { let fd = unix::MockFd; let _ = AsyncFd::try_new(fd); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] fn async_fd_try_with_interest_panic_caller() -> Result<(), Box> { use tokio::io::unix::AsyncFd; use tokio::io::Interest; use tokio::runtime::Builder; let panic_location_file = test_panic(|| { // Runtime without `enable_io` so it has no IO driver set. let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async { let fd = unix::MockFd; let _ = AsyncFd::try_with_interest(fd, Interest::READABLE); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } tokio-1.43.1/tests/io_poll_aio.rs000064400000000000000000000263671046102023000150520ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(target_os = "freebsd", feature = "net"))] use mio_aio::{AioFsyncMode, SourceApi}; use std::{ future::Future, io, mem, os::fd::AsFd, os::unix::io::{AsRawFd, RawFd}, pin::{pin, Pin}, task::{Context, Poll}, }; use tempfile::tempfile; use tokio::io::bsd::{Aio, AioSource}; use tokio_test::assert_pending; mod aio { use super::*; #[derive(Debug)] struct TokioSource<'fd>(mio_aio::Source>); impl<'fd> AioSource for TokioSource<'fd> { fn register(&mut self, kq: RawFd, token: usize) { self.0.register_raw(kq, token) } fn deregister(&mut self) { self.0.deregister_raw() } } /// A very crude implementation of an AIO-based future struct FsyncFut<'fd>(Aio>); impl<'fd> FsyncFut<'fd> { pub fn submit(self: Pin<&mut Self>) -> io::Result<()> { let p = unsafe { self.map_unchecked_mut(|s| &mut s.0 .0) }; match p.submit() { Ok(()) => Ok(()), Err(e) => Err(io::Error::from_raw_os_error(e as i32)), } } } impl<'fd> Future for FsyncFut<'fd> { type Output = io::Result<()>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let poll_result = self.0.poll_ready(cx); match poll_result { Poll::Pending => Poll::Pending, Poll::Ready(Err(e)) => Poll::Ready(Err(e)), Poll::Ready(Ok(_ev)) => { // At this point, we could clear readiness. But there's no // point, since we're about to drop the Aio. let p = unsafe { self.map_unchecked_mut(|s| &mut s.0 .0) }; let result = p.aio_return(); match result { Ok(r) => Poll::Ready(Ok(r)), Err(e) => Poll::Ready(Err(io::Error::from_raw_os_error(e as i32))), } } } } } /// Low-level AIO Source /// /// An example bypassing mio_aio and Nix to demonstrate how the kevent /// registration actually works, under the hood. struct LlSource(Pin>); impl LlSource { fn fsync(mut self: Pin<&mut Self>) { let r = unsafe { let p = self.0.as_mut().get_unchecked_mut(); libc::aio_fsync(libc::O_SYNC, p) }; assert_eq!(0, r); } } impl AioSource for LlSource { fn register(&mut self, kq: RawFd, token: usize) { let mut sev: libc::sigevent = unsafe { mem::MaybeUninit::zeroed().assume_init() }; sev.sigev_notify = libc::SIGEV_KEVENT; sev.sigev_signo = kq; sev.sigev_value = libc::sigval { sival_ptr: token as *mut libc::c_void, }; self.0.aio_sigevent = sev; } fn deregister(&mut self) { unsafe { self.0.aio_sigevent = mem::zeroed(); } } } struct LlFut(Aio); impl LlFut { pub fn fsync(self: Pin<&mut Self>) { let p = unsafe { self.map_unchecked_mut(|s| &mut *(s.0)) }; p.fsync(); } } impl Future for LlFut { type Output = std::io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let poll_result = self.0.poll_ready(cx); match poll_result { Poll::Pending => Poll::Pending, Poll::Ready(Err(e)) => Poll::Ready(Err(e)), Poll::Ready(Ok(ev)) => { // Clearing readiness makes the future non-idempotent; the // caller can't poll it repeatedly after it has already // returned Ready. But that's ok; most futures behave this // way. self.0.clear_ready(ev); let r = unsafe { libc::aio_return(self.0 .0.as_mut().get_unchecked_mut()) }; if r >= 0 { Poll::Ready(Ok(r as usize)) } else { Poll::Ready(Err(io::Error::last_os_error())) } } } } } #[tokio::test] async fn fsync() { let f = tempfile().unwrap(); let fd = f.as_fd(); let mode = AioFsyncMode::O_SYNC; let source = TokioSource(mio_aio::Fsync::fsync(fd, mode, 0)); let poll_aio = Aio::new_for_aio(source).unwrap(); let mut fut = pin!(FsyncFut(poll_aio)); fut.as_mut().submit().unwrap(); fut.await.unwrap(); } #[tokio::test] async fn ll_fsync() { let f = tempfile().unwrap(); let fd = f.as_raw_fd(); let mut aiocb: libc::aiocb = unsafe { mem::MaybeUninit::zeroed().assume_init() }; aiocb.aio_fildes = fd; let source = LlSource(Box::pin(aiocb)); let mut poll_aio = Aio::new_for_aio(source).unwrap(); let r = unsafe { let p = poll_aio.0.as_mut().get_unchecked_mut(); libc::aio_fsync(libc::O_SYNC, p) }; assert_eq!(0, r); let fut = LlFut(poll_aio); fut.await.unwrap(); } /// A suitably crafted future type can reuse an Aio object #[tokio::test] async fn reuse() { let f = tempfile().unwrap(); let fd = f.as_raw_fd(); let mut aiocb: libc::aiocb = unsafe { mem::MaybeUninit::zeroed().assume_init() }; aiocb.aio_fildes = fd; let source = LlSource(Box::pin(aiocb)); let poll_aio = Aio::new_for_aio(source).unwrap(); // Send the operation to the kernel the first time let mut fut = LlFut(poll_aio); { let mut pfut = Pin::new(&mut fut); pfut.as_mut().fsync(); pfut.as_mut().await.unwrap(); } // Check that readiness was cleared let mut ctx = Context::from_waker(futures::task::noop_waker_ref()); assert_pending!(fut.0.poll_ready(&mut ctx)); // and reuse the future and its Aio object { let mut pfut = Pin::new(&mut fut); pfut.as_mut().fsync(); pfut.as_mut().await.unwrap(); } } } mod lio { use super::*; /// Low-level source based on lio_listio /// /// An example demonstrating using AIO with `Interest::Lio`. mio_aio 0.8 /// doesn't include any bindings for lio_listio, so we've got to go /// low-level. struct LioSource<'a> { aiocb: Pin<&'a mut [&'a mut libc::aiocb; 1]>, sev: libc::sigevent, } impl<'a> LioSource<'a> { fn new(aiocb: Pin<&'a mut [&'a mut libc::aiocb; 1]>) -> Self { LioSource { aiocb, sev: unsafe { mem::zeroed() }, } } fn submit(mut self: Pin<&mut Self>) { let p: *const *mut libc::aiocb = unsafe { self.aiocb.as_mut().get_unchecked_mut() } as *const _ as *const *mut _; let r = unsafe { libc::lio_listio(libc::LIO_NOWAIT, p, 1, &mut self.sev) }; assert_eq!(r, 0); } } impl<'a> AioSource for LioSource<'a> { fn register(&mut self, kq: RawFd, token: usize) { let mut sev: libc::sigevent = unsafe { mem::MaybeUninit::zeroed().assume_init() }; sev.sigev_notify = libc::SIGEV_KEVENT; sev.sigev_signo = kq; sev.sigev_value = libc::sigval { sival_ptr: token as *mut libc::c_void, }; self.sev = sev; } fn deregister(&mut self) { unsafe { self.sev = mem::zeroed(); } } } struct LioFut<'a>(Aio>); impl<'a> LioFut<'a> { pub fn submit(self: Pin<&mut Self>) { let p = unsafe { self.map_unchecked_mut(|s| &mut *(s.0)) }; p.submit(); } } impl<'a> Future for LioFut<'a> { type Output = std::io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let poll_result = self.0.poll_ready(cx); match poll_result { Poll::Pending => Poll::Pending, Poll::Ready(Err(e)) => Poll::Ready(Err(e)), Poll::Ready(Ok(ev)) => { // Clearing readiness makes the future non-idempotent; the // caller can't poll it repeatedly after it has already // returned Ready. But that's ok; most futures behave this // way. Clearing readiness is especially useful for // lio_listio, because sometimes some operations will be // ready but not all. self.0.clear_ready(ev); let r = unsafe { let p1 = self.get_unchecked_mut(); let p2: &mut [&mut libc::aiocb; 1] = p1.0.aiocb.as_mut().get_unchecked_mut(); let p3: &mut libc::aiocb = p2[0]; libc::aio_return(p3) }; if r >= 0 { Poll::Ready(Ok(r as usize)) } else { Poll::Ready(Err(io::Error::last_os_error())) } } } } } /// An lio_listio operation with one fsync element #[tokio::test] async fn onewrite() { const WBUF: &[u8] = b"abcdef"; let f = tempfile().unwrap(); let mut aiocb: libc::aiocb = unsafe { mem::zeroed() }; aiocb.aio_fildes = f.as_raw_fd(); aiocb.aio_lio_opcode = libc::LIO_WRITE; aiocb.aio_nbytes = WBUF.len(); aiocb.aio_buf = WBUF.as_ptr() as *mut _; let aiocb = pin!([&mut aiocb]); let source = LioSource::new(aiocb); let poll_aio = Aio::new_for_lio(source).unwrap(); // Send the operation to the kernel let mut fut = pin!(LioFut(poll_aio)); fut.as_mut().submit(); fut.await.unwrap(); } /// A suitably crafted future type can reuse an Aio object #[tokio::test] async fn reuse() { const WBUF: &[u8] = b"abcdef"; let f = tempfile().unwrap(); let mut aiocb: libc::aiocb = unsafe { mem::zeroed() }; aiocb.aio_fildes = f.as_raw_fd(); aiocb.aio_lio_opcode = libc::LIO_WRITE; aiocb.aio_nbytes = WBUF.len(); aiocb.aio_buf = WBUF.as_ptr() as *mut _; let aiocb = pin!([&mut aiocb]); let source = LioSource::new(aiocb); let poll_aio = Aio::new_for_lio(source).unwrap(); // Send the operation to the kernel the first time let mut fut = LioFut(poll_aio); { let mut pfut = Pin::new(&mut fut); pfut.as_mut().submit(); pfut.as_mut().await.unwrap(); } // Check that readiness was cleared let mut ctx = Context::from_waker(futures::task::noop_waker_ref()); assert_pending!(fut.0.poll_ready(&mut ctx)); // and reuse the future and its Aio object { let mut pfut = Pin::new(&mut fut); pfut.as_mut().submit(); pfut.as_mut().await.unwrap(); } } } tokio-1.43.1/tests/io_read.rs000064400000000000000000000032701046102023000141530ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf}; use tokio_test::assert_ok; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; mod support { pub(crate) mod leaked_buffers; } use support::leaked_buffers::LeakedBuffers; #[tokio::test] async fn read() { #[derive(Default)] struct Rd { poll_cnt: usize, } impl AsyncRead for Rd { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { assert_eq!(0, self.poll_cnt); self.poll_cnt += 1; buf.put_slice(b"hello world"); Poll::Ready(Ok(())) } } let mut buf = Box::new([0; 11]); let mut rd = Rd::default(); let n = assert_ok!(rd.read(&mut buf[..]).await); assert_eq!(n, 11); assert_eq!(buf[..], b"hello world"[..]); } struct BadAsyncRead { leaked_buffers: LeakedBuffers, } impl BadAsyncRead { fn new() -> Self { Self { leaked_buffers: LeakedBuffers::new(), } } } impl AsyncRead for BadAsyncRead { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { *buf = ReadBuf::new(unsafe { self.leaked_buffers.create(buf.capacity()) }); buf.advance(buf.capacity()); Poll::Ready(Ok(())) } } #[tokio::test] #[should_panic] async fn read_buf_bad_async_read() { let mut buf = Vec::with_capacity(10); BadAsyncRead::new().read_buf(&mut buf).await.unwrap(); } tokio-1.43.1/tests/io_read_buf.rs000064400000000000000000000047361046102023000150170ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf}; use tokio_test::assert_ok; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; #[tokio::test] async fn read_buf() { struct Rd { cnt: usize, } impl AsyncRead for Rd { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { self.cnt += 1; buf.put_slice(b"hello world"); Poll::Ready(Ok(())) } } let mut buf = vec![]; let mut rd = Rd { cnt: 0 }; let n = assert_ok!(rd.read_buf(&mut buf).await); assert_eq!(1, rd.cnt); assert_eq!(n, 11); assert_eq!(buf[..], b"hello world"[..]); } #[tokio::test] #[cfg(feature = "io-util")] async fn issue_5588() { use bytes::BufMut; // steps to zero let mut buf = [0; 8]; let mut read_buf = ReadBuf::new(&mut buf); assert_eq!(read_buf.remaining_mut(), 8); assert_eq!(read_buf.chunk_mut().len(), 8); unsafe { read_buf.advance_mut(1); } assert_eq!(read_buf.remaining_mut(), 7); assert_eq!(read_buf.chunk_mut().len(), 7); unsafe { read_buf.advance_mut(5); } assert_eq!(read_buf.remaining_mut(), 2); assert_eq!(read_buf.chunk_mut().len(), 2); unsafe { read_buf.advance_mut(2); } assert_eq!(read_buf.remaining_mut(), 0); assert_eq!(read_buf.chunk_mut().len(), 0); // directly to zero let mut buf = [0; 8]; let mut read_buf = ReadBuf::new(&mut buf); assert_eq!(read_buf.remaining_mut(), 8); assert_eq!(read_buf.chunk_mut().len(), 8); unsafe { read_buf.advance_mut(8); } assert_eq!(read_buf.remaining_mut(), 0); assert_eq!(read_buf.chunk_mut().len(), 0); // uninit let mut buf = [std::mem::MaybeUninit::new(1); 8]; let mut uninit = ReadBuf::uninit(&mut buf); assert_eq!(uninit.remaining_mut(), 8); assert_eq!(uninit.chunk_mut().len(), 8); let mut buf = [std::mem::MaybeUninit::uninit(); 8]; let mut uninit = ReadBuf::uninit(&mut buf); unsafe { uninit.advance_mut(4); } assert_eq!(uninit.remaining_mut(), 4); assert_eq!(uninit.chunk_mut().len(), 4); uninit.put_u8(1); assert_eq!(uninit.remaining_mut(), 3); assert_eq!(uninit.chunk_mut().len(), 3); uninit.put_slice(&[1, 2, 3]); assert_eq!(uninit.remaining_mut(), 0); assert_eq!(uninit.chunk_mut().len(), 0); } tokio-1.43.1/tests/io_read_exact.rs000064400000000000000000000005401046102023000153340ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::AsyncReadExt; use tokio_test::assert_ok; #[tokio::test] async fn read_exact() { let mut buf = Box::new([0; 8]); let mut rd: &[u8] = b"hello world"; let n = assert_ok!(rd.read_exact(&mut buf[..]).await); assert_eq!(n, 8); assert_eq!(buf[..], b"hello wo"[..]); } tokio-1.43.1/tests/io_read_line.rs000064400000000000000000000061341046102023000151640ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::io::ErrorKind; use tokio::io::{AsyncBufReadExt, BufReader, Error}; use tokio_test::{assert_ok, io::Builder}; use std::io::Cursor; #[tokio::test] async fn read_line() { let mut buf = String::new(); let mut rd = Cursor::new(b"hello\nworld\n\n"); let n = assert_ok!(rd.read_line(&mut buf).await); assert_eq!(n, 6); assert_eq!(buf, "hello\n"); buf.clear(); let n = assert_ok!(rd.read_line(&mut buf).await); assert_eq!(n, 6); assert_eq!(buf, "world\n"); buf.clear(); let n = assert_ok!(rd.read_line(&mut buf).await); assert_eq!(n, 1); assert_eq!(buf, "\n"); buf.clear(); let n = assert_ok!(rd.read_line(&mut buf).await); assert_eq!(n, 0); assert_eq!(buf, ""); } #[tokio::test] async fn read_line_not_all_ready() { let mock = Builder::new() .read(b"Hello Wor") .read(b"ld\nFizzBuz") .read(b"z\n1\n2") .build(); let mut read = BufReader::new(mock); let mut line = "We say ".to_string(); let bytes = read.read_line(&mut line).await.unwrap(); assert_eq!(bytes, "Hello World\n".len()); assert_eq!(line.as_str(), "We say Hello World\n"); line = "I solve ".to_string(); let bytes = read.read_line(&mut line).await.unwrap(); assert_eq!(bytes, "FizzBuzz\n".len()); assert_eq!(line.as_str(), "I solve FizzBuzz\n"); line.clear(); let bytes = read.read_line(&mut line).await.unwrap(); assert_eq!(bytes, 2); assert_eq!(line.as_str(), "1\n"); line.clear(); let bytes = read.read_line(&mut line).await.unwrap(); assert_eq!(bytes, 1); assert_eq!(line.as_str(), "2"); } #[tokio::test] async fn read_line_invalid_utf8() { let mock = Builder::new().read(b"Hello Wor\xffld.\n").build(); let mut read = BufReader::new(mock); let mut line = "Foo".to_string(); let err = read.read_line(&mut line).await.expect_err("Should fail"); assert_eq!(err.kind(), ErrorKind::InvalidData); assert_eq!(err.to_string(), "stream did not contain valid UTF-8"); assert_eq!(line.as_str(), "Foo"); } #[tokio::test] async fn read_line_fail() { let mock = Builder::new() .read(b"Hello Wor") .read_error(Error::new(ErrorKind::Other, "The world has no end")) .build(); let mut read = BufReader::new(mock); let mut line = "Foo".to_string(); let err = read.read_line(&mut line).await.expect_err("Should fail"); assert_eq!(err.kind(), ErrorKind::Other); assert_eq!(err.to_string(), "The world has no end"); assert_eq!(line.as_str(), "FooHello Wor"); } #[tokio::test] async fn read_line_fail_and_utf8_fail() { let mock = Builder::new() .read(b"Hello Wor") .read(b"\xff\xff\xff") .read_error(Error::new(ErrorKind::Other, "The world has no end")) .build(); let mut read = BufReader::new(mock); let mut line = "Foo".to_string(); let err = read.read_line(&mut line).await.expect_err("Should fail"); assert_eq!(err.kind(), ErrorKind::Other); assert_eq!(err.to_string(), "The world has no end"); assert_eq!(line.as_str(), "Foo"); } tokio-1.43.1/tests/io_read_to_end.rs000064400000000000000000000070151046102023000155040ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::pin::Pin; use std::task::{Context, Poll}; use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf}; use tokio_test::assert_ok; use tokio_test::io::Builder; #[tokio::test] async fn read_to_end() { let mut buf = vec![]; let mut rd: &[u8] = b"hello world"; let n = assert_ok!(rd.read_to_end(&mut buf).await); assert_eq!(n, 11); assert_eq!(buf[..], b"hello world"[..]); } #[derive(Copy, Clone, Debug)] enum State { Initializing, JustFilling, Done, } struct UninitTest { num_init: usize, state: State, } impl AsyncRead for UninitTest { fn poll_read( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let me = Pin::into_inner(self); let real_num_init = buf.initialized().len() - buf.filled().len(); assert_eq!(real_num_init, me.num_init, "{:?}", me.state); match me.state { State::Initializing => { buf.initialize_unfilled_to(me.num_init + 2); buf.advance(1); me.num_init += 1; if me.num_init == 24 { me.state = State::JustFilling; } } State::JustFilling => { buf.advance(1); me.num_init -= 1; if me.num_init == 15 { // The buffer is resized on next call. me.num_init = 0; me.state = State::Done; } } State::Done => { /* .. do nothing .. */ } } Poll::Ready(Ok(())) } } #[tokio::test] async fn read_to_end_uninit() { let mut buf = Vec::with_capacity(64); let mut test = UninitTest { num_init: 0, state: State::Initializing, }; test.read_to_end(&mut buf).await.unwrap(); assert_eq!(buf.len(), 33); } #[tokio::test] #[cfg_attr(miri, ignore)] // too slow with miri async fn read_to_end_doesnt_grow_with_capacity() { let arr: Vec = (0..100).collect(); // We only test from 32 since we allocate at least 32 bytes each time for len in 32..100 { let bytes = &arr[..len]; for split in 0..len { for cap in 0..101 { let mut mock = if split == 0 { Builder::new().read(bytes).build() } else { Builder::new() .read(&bytes[..split]) .read(&bytes[split..]) .build() }; let mut buf = Vec::with_capacity(cap); AsyncReadExt::read_to_end(&mut mock, &mut buf) .await .unwrap(); // It has the right data. assert_eq!(buf.as_slice(), bytes); // Unless cap was smaller than length, then we did not reallocate. if cap >= len { assert_eq!(buf.capacity(), cap); } } } } } #[tokio::test] async fn read_to_end_grows_capacity_if_unfit() { let bytes = b"the_vector_startingcap_will_be_smaller"; let mut mock = Builder::new().read(bytes).build(); let initial_capacity = bytes.len() - 4; let mut buf = Vec::with_capacity(initial_capacity); AsyncReadExt::read_to_end(&mut mock, &mut buf) .await .unwrap(); // *4 since it doubles when it doesn't fit and again when reaching EOF assert_eq!(buf.capacity(), initial_capacity * 4); } tokio-1.43.1/tests/io_read_to_string.rs000064400000000000000000000031241046102023000162410ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::io; use tokio::io::AsyncReadExt; use tokio_test::assert_ok; use tokio_test::io::Builder; #[tokio::test] async fn read_to_string() { let mut buf = String::new(); let mut rd: &[u8] = b"hello world"; let n = assert_ok!(rd.read_to_string(&mut buf).await); assert_eq!(n, 11); assert_eq!(buf[..], "hello world"[..]); } #[tokio::test] async fn to_string_does_not_truncate_on_utf8_error() { let data = vec![0xff, 0xff, 0xff]; let mut s = "abc".to_string(); match AsyncReadExt::read_to_string(&mut data.as_slice(), &mut s).await { Ok(len) => panic!("Should fail: {len} bytes."), Err(err) if err.to_string() == "stream did not contain valid UTF-8" => {} Err(err) => panic!("Fail: {err}."), } assert_eq!(s, "abc"); } #[tokio::test] async fn to_string_does_not_truncate_on_io_error() { let mut mock = Builder::new() .read(b"def") .read_error(io::Error::new(io::ErrorKind::Other, "whoops")) .build(); let mut s = "abc".to_string(); match AsyncReadExt::read_to_string(&mut mock, &mut s).await { Ok(len) => panic!("Should fail: {len} bytes."), Err(err) if err.to_string() == "whoops" => {} Err(err) => panic!("Fail: {err}."), } assert_eq!(s, "abc"); } #[tokio::test] async fn to_string_appends() { let data = b"def".to_vec(); let mut s = "abc".to_string(); let len = AsyncReadExt::read_to_string(&mut data.as_slice(), &mut s) .await .unwrap(); assert_eq!(len, 3); assert_eq!(s, "abcdef"); } tokio-1.43.1/tests/io_read_until.rs000064400000000000000000000040651046102023000153710ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::io::ErrorKind; use tokio::io::{AsyncBufReadExt, BufReader, Error}; use tokio_test::{assert_ok, io::Builder}; #[tokio::test] async fn read_until() { let mut buf = vec![]; let mut rd: &[u8] = b"hello world"; let n = assert_ok!(rd.read_until(b' ', &mut buf).await); assert_eq!(n, 6); assert_eq!(buf, b"hello "); buf.clear(); let n = assert_ok!(rd.read_until(b' ', &mut buf).await); assert_eq!(n, 5); assert_eq!(buf, b"world"); buf.clear(); let n = assert_ok!(rd.read_until(b' ', &mut buf).await); assert_eq!(n, 0); assert_eq!(buf, []); } #[tokio::test] async fn read_until_not_all_ready() { let mock = Builder::new() .read(b"Hello Wor") .read(b"ld#Fizz\xffBuz") .read(b"z#1#2") .build(); let mut read = BufReader::new(mock); let mut chunk = b"We say ".to_vec(); let bytes = read.read_until(b'#', &mut chunk).await.unwrap(); assert_eq!(bytes, b"Hello World#".len()); assert_eq!(chunk, b"We say Hello World#"); chunk = b"I solve ".to_vec(); let bytes = read.read_until(b'#', &mut chunk).await.unwrap(); assert_eq!(bytes, b"Fizz\xffBuzz\n".len()); assert_eq!(chunk, b"I solve Fizz\xffBuzz#"); chunk.clear(); let bytes = read.read_until(b'#', &mut chunk).await.unwrap(); assert_eq!(bytes, 2); assert_eq!(chunk, b"1#"); chunk.clear(); let bytes = read.read_until(b'#', &mut chunk).await.unwrap(); assert_eq!(bytes, 1); assert_eq!(chunk, b"2"); } #[tokio::test] async fn read_until_fail() { let mock = Builder::new() .read(b"Hello \xffWor") .read_error(Error::new(ErrorKind::Other, "The world has no end")) .build(); let mut read = BufReader::new(mock); let mut chunk = b"Foo".to_vec(); let err = read .read_until(b'#', &mut chunk) .await .expect_err("Should fail"); assert_eq!(err.kind(), ErrorKind::Other); assert_eq!(err.to_string(), "The world has no end"); assert_eq!(chunk, b"FooHello \xffWor"); } tokio-1.43.1/tests/io_repeat.rs000064400000000000000000000006671046102023000145270ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(miri)))] use tokio::io::AsyncReadExt; #[tokio::test] async fn repeat_poll_read_is_cooperative() { tokio::select! { biased; _ = async { loop { let mut buf = [0u8; 4096]; tokio::io::repeat(0b101).read_exact(&mut buf).await.unwrap(); } } => {}, _ = tokio::task::yield_now() => {} } } tokio-1.43.1/tests/io_sink.rs000064400000000000000000000017071046102023000142070ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::AsyncWriteExt; #[tokio::test] async fn sink_poll_write_is_cooperative() { tokio::select! { biased; _ = async { loop { let buf = vec![1, 2, 3]; tokio::io::sink().write_all(&buf).await.unwrap(); } } => {}, _ = tokio::task::yield_now() => {} } } #[tokio::test] async fn sink_poll_flush_is_cooperative() { tokio::select! { biased; _ = async { loop { tokio::io::sink().flush().await.unwrap(); } } => {}, _ = tokio::task::yield_now() => {} } } #[tokio::test] async fn sink_poll_shutdown_is_cooperative() { tokio::select! { biased; _ = async { loop { tokio::io::sink().shutdown().await.unwrap(); } } => {}, _ = tokio::task::yield_now() => {} } } tokio-1.43.1/tests/io_split.rs000064400000000000000000000047651046102023000144050ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use tokio::io::{ split, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf, ReadHalf, WriteHalf, }; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; struct RW; impl AsyncRead for RW { fn poll_read( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { buf.put_slice(&[b'z']); Poll::Ready(Ok(())) } } impl AsyncWrite for RW { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8], ) -> Poll> { Poll::Ready(Ok(1)) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_write_vectored( self: Pin<&mut Self>, _cx: &mut Context<'_>, _bufs: &[io::IoSlice<'_>], ) -> Poll> { Poll::Ready(Ok(2)) } fn is_write_vectored(&self) -> bool { true } } #[test] fn is_send_and_sync() { fn assert_bound() {} assert_bound::>(); assert_bound::>(); } #[test] fn split_stream_id() { let (r1, w1) = split(RW); let (r2, w2) = split(RW); assert!(r1.is_pair_of(&w1)); assert!(!r1.is_pair_of(&w2)); assert!(r2.is_pair_of(&w2)); assert!(!r2.is_pair_of(&w1)); } #[test] fn unsplit_ok() { let (r, w) = split(RW); r.unsplit(w); } #[test] #[should_panic] fn unsplit_err1() { let (r, _) = split(RW); let (_, w) = split(RW); r.unsplit(w); } #[test] #[should_panic] fn unsplit_err2() { let (_, w) = split(RW); let (r, _) = split(RW); r.unsplit(w); } #[test] fn method_delegation() { let (mut r, mut w) = split(RW); let mut buf = [0; 1]; tokio_test::block_on(async move { assert_eq!(1, r.read(&mut buf).await.unwrap()); assert_eq!(b'z', buf[0]); assert_eq!(1, w.write(&[b'x']).await.unwrap()); assert_eq!( 2, w.write_vectored(&[io::IoSlice::new(&[b'x'])]) .await .unwrap() ); assert!(w.is_write_vectored()); assert!(w.flush().await.is_ok()); assert!(w.shutdown().await.is_ok()); }); } tokio-1.43.1/tests/io_take.rs000064400000000000000000000032251046102023000141640ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use std::pin::Pin; use std::task::{Context, Poll}; use tokio::io::{self, AsyncRead, AsyncReadExt, ReadBuf}; use tokio_test::assert_ok; mod support { pub(crate) mod leaked_buffers; } use support::leaked_buffers::LeakedBuffers; #[tokio::test] async fn take() { let mut buf = [0; 6]; let rd: &[u8] = b"hello world"; let mut rd = rd.take(4); let n = assert_ok!(rd.read(&mut buf).await); assert_eq!(n, 4); assert_eq!(&buf, &b"hell\0\0"[..]); } #[tokio::test] async fn issue_4435() { let mut buf = [0; 8]; let rd: &[u8] = b"hello world"; let rd = rd.take(4); tokio::pin!(rd); let mut read_buf = ReadBuf::new(&mut buf); read_buf.put_slice(b"AB"); std::future::poll_fn(|cx| rd.as_mut().poll_read(cx, &mut read_buf)) .await .unwrap(); assert_eq!(&buf, &b"ABhell\0\0"[..]); } struct BadReader { leaked_buffers: LeakedBuffers, } impl BadReader { fn new() -> Self { Self { leaked_buffers: LeakedBuffers::new(), } } } impl AsyncRead for BadReader { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, read_buf: &mut ReadBuf<'_>, ) -> Poll> { let mut buf = ReadBuf::new(unsafe { self.leaked_buffers.create(10) }); buf.put_slice(&[123; 10]); *read_buf = buf; Poll::Ready(Ok(())) } } #[tokio::test] #[should_panic] async fn bad_reader_fails() { let mut buf = Vec::with_capacity(10); BadReader::new().take(10).read_buf(&mut buf).await.unwrap(); } tokio-1.43.1/tests/io_util_empty.rs000064400000000000000000000034371046102023000154400ustar 00000000000000#![cfg(feature = "full")] use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt}; #[tokio::test] async fn empty_read_is_cooperative() { tokio::select! { biased; _ = async { loop { let mut buf = [0u8; 4096]; let _ = tokio::io::empty().read(&mut buf).await; } } => {}, _ = tokio::task::yield_now() => {} } } #[tokio::test] async fn empty_buf_reads_are_cooperative() { tokio::select! { biased; _ = async { loop { let mut buf = String::new(); let _ = tokio::io::empty().read_line(&mut buf).await; } } => {}, _ = tokio::task::yield_now() => {} } } #[tokio::test] async fn empty_seek() { use std::io::SeekFrom; let mut empty = tokio::io::empty(); assert!(matches!(empty.seek(SeekFrom::Start(0)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::Start(1)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::Start(u64::MAX)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::End(i64::MIN)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::End(-1)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::End(0)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::End(1)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::End(i64::MAX)).await, Ok(0))); assert!(matches!( empty.seek(SeekFrom::Current(i64::MIN)).await, Ok(0) )); assert!(matches!(empty.seek(SeekFrom::Current(-1)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::Current(0)).await, Ok(0))); assert!(matches!(empty.seek(SeekFrom::Current(1)).await, Ok(0))); assert!(matches!( empty.seek(SeekFrom::Current(i64::MAX)).await, Ok(0) )); } tokio-1.43.1/tests/io_write.rs000064400000000000000000000024761046102023000144010ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio_test::assert_ok; use bytes::BytesMut; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; #[tokio::test] async fn write() { struct Wr { buf: BytesMut, cnt: usize, } impl AsyncWrite for Wr { fn poll_write( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { assert_eq!(self.cnt, 0); self.buf.extend(&buf[0..4]); Ok(4).into() } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } } let mut wr = Wr { buf: BytesMut::with_capacity(64), cnt: 0, }; let n = assert_ok!(wr.write(b"hello world").await); assert_eq!(n, 4); assert_eq!(wr.buf, b"hell"[..]); } #[tokio::test] async fn write_cursor() { use std::io::Cursor; let mut wr = Cursor::new(Vec::new()); let n = assert_ok!(wr.write(b"hello world").await); assert_eq!(n, 11); assert_eq!(wr.get_ref().as_slice(), &b"hello world"[..]); } tokio-1.43.1/tests/io_write_all.rs000064400000000000000000000022401046102023000152160ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio_test::assert_ok; use bytes::BytesMut; use std::cmp; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; #[tokio::test] async fn write_all() { struct Wr { buf: BytesMut, cnt: usize, } impl AsyncWrite for Wr { fn poll_write( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { let n = cmp::min(4, buf.len()); let buf = &buf[0..n]; self.cnt += 1; self.buf.extend(buf); Ok(buf.len()).into() } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } } let mut wr = Wr { buf: BytesMut::with_capacity(64), cnt: 0, }; assert_ok!(wr.write_all(b"hello world").await); assert_eq!(wr.buf, b"hello world"[..]); assert_eq!(wr.cnt, 3); } tokio-1.43.1/tests/io_write_all_buf.rs000064400000000000000000000075541046102023000160670ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio_test::{assert_err, assert_ok}; use bytes::{Buf, Bytes, BytesMut}; use std::cmp; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; #[tokio::test] async fn write_all_buf() { struct Wr { buf: BytesMut, cnt: usize, } impl AsyncWrite for Wr { fn poll_write( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { let n = cmp::min(4, buf.len()); dbg!(buf); let buf = &buf[0..n]; self.cnt += 1; self.buf.extend(buf); Ok(buf.len()).into() } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } } let mut wr = Wr { buf: BytesMut::with_capacity(64), cnt: 0, }; let mut buf = Bytes::from_static(b"hello").chain(Bytes::from_static(b"world")); assert_ok!(wr.write_all_buf(&mut buf).await); assert_eq!(wr.buf, b"helloworld"[..]); // expect 4 writes, [hell],[o],[worl],[d] assert_eq!(wr.cnt, 4); assert!(!buf.has_remaining()); } #[tokio::test] async fn write_buf_err() { /// Error out after writing the first 4 bytes struct Wr { cnt: usize, } impl AsyncWrite for Wr { fn poll_write( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8], ) -> Poll> { self.cnt += 1; if self.cnt == 2 { return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, "whoops"))); } Poll::Ready(Ok(4)) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } } let mut wr = Wr { cnt: 0 }; let mut buf = Bytes::from_static(b"hello").chain(Bytes::from_static(b"world")); assert_err!(wr.write_all_buf(&mut buf).await); assert_eq!( buf.copy_to_bytes(buf.remaining()), Bytes::from_static(b"oworld") ); } #[tokio::test] async fn write_all_buf_vectored() { struct Wr { buf: BytesMut, } impl AsyncWrite for Wr { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8], ) -> Poll> { // When executing `write_all_buf` with this writer, // `poll_write` is not called. panic!("shouldn't be called") } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_write_vectored( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, bufs: &[io::IoSlice<'_>], ) -> Poll> { for buf in bufs { self.buf.extend_from_slice(buf); } let n = self.buf.len(); Ok(n).into() } fn is_write_vectored(&self) -> bool { // Enable vectored write. true } } let mut wr = Wr { buf: BytesMut::with_capacity(64), }; let mut buf = Bytes::from_static(b"hello") .chain(Bytes::from_static(b" ")) .chain(Bytes::from_static(b"world")); wr.write_all_buf(&mut buf).await.unwrap(); assert_eq!(&wr.buf[..], b"hello world"); } tokio-1.43.1/tests/io_write_buf.rs000064400000000000000000000024411046102023000152250ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio_test::assert_ok; use bytes::BytesMut; use std::cmp; use std::io::{self, Cursor}; use std::pin::Pin; use std::task::{Context, Poll}; #[tokio::test] async fn write_all() { struct Wr { buf: BytesMut, cnt: usize, } impl AsyncWrite for Wr { fn poll_write( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { assert_eq!(self.cnt, 0); let n = cmp::min(4, buf.len()); let buf = &buf[0..n]; self.cnt += 1; self.buf.extend(buf); Ok(buf.len()).into() } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } } let mut wr = Wr { buf: BytesMut::with_capacity(64), cnt: 0, }; let mut buf = Cursor::new(&b"hello world"[..]); assert_ok!(wr.write_buf(&mut buf).await); assert_eq!(wr.buf, b"hell"[..]); assert_eq!(wr.cnt, 1); assert_eq!(buf.position(), 4); } tokio-1.43.1/tests/io_write_int.rs000064400000000000000000000016561046102023000152520ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::io::{AsyncWrite, AsyncWriteExt}; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; #[tokio::test] async fn write_int_should_err_if_write_count_0() { struct Wr {} impl AsyncWrite for Wr { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8], ) -> Poll> { Ok(0).into() } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Ok(()).into() } } let mut wr = Wr {}; // should be ok just to test these 2, other cases actually expanded by same macro. assert!(wr.write_i8(0).await.is_err()); assert!(wr.write_i32(12).await.is_err()); } tokio-1.43.1/tests/join_handle_panic.rs000064400000000000000000000011551046102023000161750ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery #![cfg(panic = "unwind")] struct PanicsOnDrop; impl Drop for PanicsOnDrop { fn drop(&mut self) { panic!("I told you so"); } } #[tokio::test] async fn test_panics_do_not_propagate_when_dropping_join_handle() { let join_handle = tokio::spawn(async move { PanicsOnDrop }); // only drop the JoinHandle when the task has completed // (which is difficult to synchronize precisely) tokio::time::sleep(std::time::Duration::from_millis(3)).await; drop(join_handle); } tokio-1.43.1/tests/macros_join.rs000064400000000000000000000105541046102023000150570ustar 00000000000000#![cfg(feature = "macros")] #![allow(clippy::disallowed_names)] use std::sync::Arc; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] #[cfg(target_pointer_width = "64")] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::{oneshot, Semaphore}; use tokio_test::{assert_pending, assert_ready, task}; #[maybe_tokio_test] async fn sync_one_lit_expr_comma() { let foo = tokio::join!(async { 1 },); assert_eq!(foo, (1,)); } #[maybe_tokio_test] async fn sync_one_lit_expr_no_comma() { let foo = tokio::join!(async { 1 }); assert_eq!(foo, (1,)); } #[maybe_tokio_test] async fn sync_two_lit_expr_comma() { let foo = tokio::join!(async { 1 }, async { 2 },); assert_eq!(foo, (1, 2)); } #[maybe_tokio_test] async fn sync_two_lit_expr_no_comma() { let foo = tokio::join!(async { 1 }, async { 2 }); assert_eq!(foo, (1, 2)); } #[maybe_tokio_test] async fn two_await() { let (tx1, rx1) = oneshot::channel::<&str>(); let (tx2, rx2) = oneshot::channel::(); let mut join = task::spawn(async { tokio::join!(async { rx1.await.unwrap() }, async { rx2.await.unwrap() }) }); assert_pending!(join.poll()); tx2.send(123).unwrap(); assert!(join.is_woken()); assert_pending!(join.poll()); tx1.send("hello").unwrap(); assert!(join.is_woken()); let res = assert_ready!(join.poll()); assert_eq!(("hello", 123), res); } #[test] #[cfg(target_pointer_width = "64")] fn join_size() { use futures::future; use std::mem; let fut = async { let ready = future::ready(0i32); tokio::join!(ready) }; assert_eq!(mem::size_of_val(&fut), 32); let fut = async { let ready1 = future::ready(0i32); let ready2 = future::ready(0i32); tokio::join!(ready1, ready2) }; assert_eq!(mem::size_of_val(&fut), 48); } async fn non_cooperative_task(permits: Arc) -> usize { let mut exceeded_budget = 0; for _ in 0..5 { // Another task should run after this task uses its whole budget for _ in 0..128 { let _permit = permits.clone().acquire_owned().await.unwrap(); } exceeded_budget += 1; } exceeded_budget } async fn poor_little_task(permits: Arc) -> usize { let mut how_many_times_i_got_to_run = 0; for _ in 0..5 { let _permit = permits.clone().acquire_owned().await.unwrap(); how_many_times_i_got_to_run += 1; } how_many_times_i_got_to_run } #[tokio::test] async fn join_does_not_allow_tasks_to_starve() { let permits = Arc::new(Semaphore::new(1)); // non_cooperative_task should yield after its budget is exceeded and then poor_little_task should run. let (non_cooperative_result, little_task_result) = tokio::join!( non_cooperative_task(Arc::clone(&permits)), poor_little_task(permits) ); assert_eq!(5, non_cooperative_result); assert_eq!(5, little_task_result); } #[tokio::test] async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { let poll_order = Arc::new(std::sync::Mutex::new(vec![])); let fut = |x, poll_order: Arc>>| async move { for _ in 0..4 { { let mut guard = poll_order.lock().unwrap(); guard.push(x); } tokio::task::yield_now().await; } }; tokio::join!( fut(1, Arc::clone(&poll_order)), fut(2, Arc::clone(&poll_order)), fut(3, Arc::clone(&poll_order)), ); // Each time the future created by join! is polled, it should start // by polling a different future first. assert_eq!( vec![1, 2, 3, 2, 3, 1, 3, 1, 2, 1, 2, 3], *poll_order.lock().unwrap() ); } #[tokio::test] #[allow(clippy::unit_cmp)] async fn empty_join() { assert_eq!(tokio::join!(), ()); } #[tokio::test] async fn join_into_future() { struct NotAFuture; impl std::future::IntoFuture for NotAFuture { type Output = (); type IntoFuture = std::future::Ready<()>; fn into_future(self) -> Self::IntoFuture { std::future::ready(()) } } tokio::join!(NotAFuture); } tokio-1.43.1/tests/macros_pin.rs000064400000000000000000000007061046102023000147040ustar 00000000000000#![cfg(feature = "macros")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; async fn one() {} async fn two() {} #[maybe_tokio_test] async fn multi_pin() { tokio::pin! { let f1 = one(); let f2 = two(); } (&mut f1).await; (&mut f2).await; } tokio-1.43.1/tests/macros_rename_test.rs000064400000000000000000000012401046102023000164160ustar 00000000000000#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threading #[allow(unused_imports)] use std as tokio; use ::tokio as tokio1; mod test { pub use ::tokio; } async fn compute() -> usize { let join = tokio1::spawn(async { 1 }); join.await.unwrap() } #[tokio1::main(crate = "tokio1")] async fn compute_main() -> usize { compute().await } #[test] fn crate_rename_main() { assert_eq!(1, compute_main()); } #[tokio1::test(crate = "tokio1")] async fn crate_rename_test() { assert_eq!(1, compute().await); } #[test::tokio::test(crate = "test::tokio")] async fn crate_path_test() { assert_eq!(1, compute().await); } tokio-1.43.1/tests/macros_select.rs000064400000000000000000000401351046102023000153750ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![cfg(feature = "macros")] #![allow(clippy::disallowed_names)] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::oneshot; use tokio_test::{assert_ok, assert_pending, assert_ready}; use std::future::poll_fn; use std::task::Poll::Ready; #[maybe_tokio_test] async fn sync_one_lit_expr_comma() { let foo = tokio::select! { foo = async { 1 } => foo, }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn no_branch_else_only() { let foo = tokio::select! { else => 1, }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn no_branch_else_only_biased() { let foo = tokio::select! { biased; else => 1, }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn nested_one() { let foo = tokio::select! { foo = async { 1 } => tokio::select! { bar = async { foo } => bar, }, }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn sync_one_lit_expr_no_comma() { let foo = tokio::select! { foo = async { 1 } => foo }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn sync_one_lit_expr_block() { let foo = tokio::select! { foo = async { 1 } => { foo } }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn sync_one_await() { let foo = tokio::select! { foo = one() => foo, }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn sync_one_ident() { let one = one(); let foo = tokio::select! { foo = one => foo, }; assert_eq!(foo, 1); } #[maybe_tokio_test] async fn sync_two() { use std::cell::Cell; let cnt = Cell::new(0); let res = tokio::select! { foo = async { cnt.set(cnt.get() + 1); 1 } => foo, bar = async { cnt.set(cnt.get() + 1); 2 } => bar, }; assert_eq!(1, cnt.get()); assert!(res == 1 || res == 2); } #[maybe_tokio_test] async fn drop_in_fut() { let s = "hello".to_string(); let res = tokio::select! { foo = async { let v = one().await; drop(s); v } => foo }; assert_eq!(res, 1); } #[maybe_tokio_test] #[cfg(feature = "full")] async fn one_ready() { let (tx1, rx1) = oneshot::channel::(); let (_tx2, rx2) = oneshot::channel::(); tx1.send(1).unwrap(); let v = tokio::select! { res = rx1 => { assert_ok!(res) }, _ = rx2 => unreachable!(), }; assert_eq!(1, v); } #[maybe_tokio_test] #[cfg(feature = "full")] async fn select_streams() { use tokio::sync::mpsc; let (tx1, mut rx1) = mpsc::unbounded_channel::(); let (tx2, mut rx2) = mpsc::unbounded_channel::(); tokio::spawn(async move { assert_ok!(tx2.send(1)); tokio::task::yield_now().await; assert_ok!(tx1.send(2)); tokio::task::yield_now().await; assert_ok!(tx2.send(3)); tokio::task::yield_now().await; drop((tx1, tx2)); }); let mut rem = true; let mut msgs = vec![]; while rem { tokio::select! { Some(x) = rx1.recv() => { msgs.push(x); } Some(y) = rx2.recv() => { msgs.push(y); } else => { rem = false; } } } msgs.sort_unstable(); assert_eq!(&msgs[..], &[1, 2, 3]); } #[maybe_tokio_test] async fn move_uncompleted_futures() { let (tx1, mut rx1) = oneshot::channel::(); let (tx2, mut rx2) = oneshot::channel::(); tx1.send(1).unwrap(); tx2.send(2).unwrap(); let ran; tokio::select! { res = &mut rx1 => { assert_eq!(1, assert_ok!(res)); assert_eq!(2, assert_ok!(rx2.await)); ran = true; }, res = &mut rx2 => { assert_eq!(2, assert_ok!(res)); assert_eq!(1, assert_ok!(rx1.await)); ran = true; }, } assert!(ran); } #[maybe_tokio_test] async fn nested() { let res = tokio::select! { x = async { 1 } => { tokio::select! { y = async { 2 } => x + y, } } }; assert_eq!(res, 3); } #[cfg(target_pointer_width = "64")] mod pointer_64_tests { use super::maybe_tokio_test; use futures::future; use std::mem; #[maybe_tokio_test] async fn struct_size_1() { let fut = async { let ready = future::ready(0i32); tokio::select! { _ = ready => {}, } }; assert_eq!(mem::size_of_val(&fut), 32); } #[maybe_tokio_test] async fn struct_size_2() { let fut = async { let ready1 = future::ready(0i32); let ready2 = future::ready(0i32); tokio::select! { _ = ready1 => {}, _ = ready2 => {}, } }; assert_eq!(mem::size_of_val(&fut), 40); } #[maybe_tokio_test] async fn struct_size_3() { let fut = async { let ready1 = future::ready(0i32); let ready2 = future::ready(0i32); let ready3 = future::ready(0i32); tokio::select! { _ = ready1 => {}, _ = ready2 => {}, _ = ready3 => {}, } }; assert_eq!(mem::size_of_val(&fut), 48); } } #[maybe_tokio_test] async fn mutable_borrowing_future_with_same_borrow_in_block() { let mut value = 234; tokio::select! { _ = require_mutable(&mut value) => { }, _ = async_noop() => { value += 5; }, } assert!(value >= 234); } #[maybe_tokio_test] async fn mutable_borrowing_future_with_same_borrow_in_block_and_else() { let mut value = 234; tokio::select! { _ = require_mutable(&mut value) => { }, _ = async_noop() => { value += 5; }, else => { value += 27; }, } assert!(value >= 234); } #[maybe_tokio_test] async fn future_panics_after_poll() { use tokio_test::task; let (tx, rx) = oneshot::channel(); let mut polled = false; let f = poll_fn(|_| { assert!(!polled); polled = true; Ready(None::<()>) }); let mut f = task::spawn(async { tokio::select! { Some(_) = f => unreachable!(), ret = rx => ret.unwrap(), } }); assert_pending!(f.poll()); assert_pending!(f.poll()); assert_ok!(tx.send(1)); let res = assert_ready!(f.poll()); assert_eq!(1, res); } #[maybe_tokio_test] async fn disable_with_if() { use tokio_test::task; let f = poll_fn(|_| panic!()); let (tx, rx) = oneshot::channel(); let mut f = task::spawn(async { tokio::select! { _ = f, if false => unreachable!(), _ = rx => (), } }); assert_pending!(f.poll()); assert_ok!(tx.send(())); assert!(f.is_woken()); assert_ready!(f.poll()); } #[maybe_tokio_test] async fn join_with_select() { use tokio_test::task; let (tx1, mut rx1) = oneshot::channel(); let (tx2, mut rx2) = oneshot::channel(); let mut f = task::spawn(async { let mut a = None; let mut b = None; while a.is_none() || b.is_none() { tokio::select! { v1 = &mut rx1, if a.is_none() => a = Some(assert_ok!(v1)), v2 = &mut rx2, if b.is_none() => b = Some(assert_ok!(v2)) } } (a.unwrap(), b.unwrap()) }); assert_pending!(f.poll()); assert_ok!(tx1.send(123)); assert!(f.is_woken()); assert_pending!(f.poll()); assert_ok!(tx2.send(456)); assert!(f.is_woken()); let (a, b) = assert_ready!(f.poll()); assert_eq!(a, 123); assert_eq!(b, 456); } #[tokio::test] #[cfg(feature = "full")] async fn use_future_in_if_condition() { use tokio::time::{self, Duration}; tokio::select! { _ = time::sleep(Duration::from_millis(10)), if false => { panic!("if condition ignored") } _ = async { 1u32 } => { } } } #[tokio::test] #[cfg(feature = "full")] async fn use_future_in_if_condition_biased() { use tokio::time::{self, Duration}; tokio::select! { biased; _ = time::sleep(Duration::from_millis(10)), if false => { panic!("if condition ignored") } _ = async { 1u32 } => { } } } #[maybe_tokio_test] async fn many_branches() { let num = tokio::select! { x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, x = async { 1 } => x, }; assert_eq!(1, num); } #[maybe_tokio_test] async fn never_branch_no_warnings() { let t = tokio::select! { _ = async_never() => 0, one_async_ready = one() => one_async_ready, }; assert_eq!(t, 1); } async fn one() -> usize { 1 } async fn require_mutable(_: &mut i32) {} async fn async_noop() {} async fn async_never() -> ! { futures::future::pending().await } // From https://github.com/tokio-rs/tokio/issues/2857 #[maybe_tokio_test] async fn mut_on_left_hand_side() { let v = async move { let ok = async { 1 }; tokio::pin!(ok); tokio::select! { mut a = &mut ok => { a += 1; a } } } .await; assert_eq!(v, 2); } #[maybe_tokio_test] async fn biased_one_not_ready() { let (_tx1, rx1) = oneshot::channel::(); let (tx2, rx2) = oneshot::channel::(); let (tx3, rx3) = oneshot::channel::(); tx2.send(2).unwrap(); tx3.send(3).unwrap(); let v = tokio::select! { biased; _ = rx1 => unreachable!(), res = rx2 => { assert_ok!(res) }, _ = rx3 => { panic!("This branch should never be activated because `rx2` should be polled before `rx3` due to `biased;`.") } }; assert_eq!(2, v); } #[maybe_tokio_test] #[cfg(feature = "full")] async fn biased_eventually_ready() { use tokio::task::yield_now; let one = async {}; let two = async { yield_now().await }; let three = async { yield_now().await }; let mut count = 0u8; tokio::pin!(one, two, three); loop { tokio::select! { biased; _ = &mut two, if count < 2 => { count += 1; assert_eq!(count, 2); } _ = &mut three, if count < 3 => { count += 1; assert_eq!(count, 3); } _ = &mut one, if count < 1 => { count += 1; assert_eq!(count, 1); } else => break, } } assert_eq!(count, 3); } // https://github.com/tokio-rs/tokio/issues/3830 // https://github.com/rust-lang/rust-clippy/issues/7304 #[warn(clippy::default_numeric_fallback)] pub async fn default_numeric_fallback() { tokio::select! { _ = async {} => (), else => (), } } // https://github.com/tokio-rs/tokio/issues/4182 #[maybe_tokio_test] async fn mut_ref_patterns() { tokio::select! { Some(mut foo) = async { Some("1".to_string()) } => { assert_eq!(foo, "1"); foo = "2".to_string(); assert_eq!(foo, "2"); }, }; tokio::select! { Some(ref foo) = async { Some("1".to_string()) } => { assert_eq!(*foo, "1"); }, }; tokio::select! { Some(ref mut foo) = async { Some("1".to_string()) } => { assert_eq!(*foo, "1"); *foo = "2".to_string(); assert_eq!(*foo, "2"); }, }; } #[cfg(tokio_unstable)] mod unstable { use tokio::runtime::RngSeed; #[test] fn deterministic_select_current_thread() { let seed = b"bytes used to generate seed"; let rt1 = tokio::runtime::Builder::new_current_thread() .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt1_values = rt1.block_on(async { (select_0_to_9().await, select_0_to_9().await) }); let rt2 = tokio::runtime::Builder::new_current_thread() .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt2_values = rt2.block_on(async { (select_0_to_9().await, select_0_to_9().await) }); assert_eq!(rt1_values, rt2_values); } #[test] #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] fn deterministic_select_multi_thread() { let seed = b"bytes used to generate seed"; let rt1 = tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt1_values = rt1.block_on(async { let _ = tokio::spawn(async { (select_0_to_9().await, select_0_to_9().await) }).await; }); let rt2 = tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt2_values = rt2.block_on(async { let _ = tokio::spawn(async { (select_0_to_9().await, select_0_to_9().await) }).await; }); assert_eq!(rt1_values, rt2_values); } async fn select_0_to_9() -> u32 { tokio::select!( x = async { 0 } => x, x = async { 1 } => x, x = async { 2 } => x, x = async { 3 } => x, x = async { 4 } => x, x = async { 5 } => x, x = async { 6 } => x, x = async { 7 } => x, x = async { 8 } => x, x = async { 9 } => x, ) } } #[tokio::test] async fn select_into_future() { struct NotAFuture; impl std::future::IntoFuture for NotAFuture { type Output = (); type IntoFuture = std::future::Ready<()>; fn into_future(self) -> Self::IntoFuture { std::future::ready(()) } } tokio::select! { () = NotAFuture => {}, } } // regression test for https://github.com/tokio-rs/tokio/issues/6721 #[tokio::test] async fn temporary_lifetime_extension() { tokio::select! { () = &mut std::future::ready(()) => {}, } } tokio-1.43.1/tests/macros_test.rs000064400000000000000000000056341046102023000151020ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threading use tokio::test; #[test] async fn test_macro_can_be_used_via_use() { tokio::spawn(async {}).await.unwrap(); } #[tokio::test] async fn test_macro_is_resilient_to_shadowing() { tokio::spawn(async {}).await.unwrap(); } // https://github.com/tokio-rs/tokio/issues/3403 #[rustfmt::skip] // this `rustfmt::skip` is necessary because unused_braces does not warn if the block contains newline. #[tokio::main] pub async fn unused_braces_main() { println!("hello") } #[rustfmt::skip] // this `rustfmt::skip` is necessary because unused_braces does not warn if the block contains newline. #[tokio::test] async fn unused_braces_test() { assert_eq!(1 + 1, 2) } // https://github.com/tokio-rs/tokio/pull/3766#issuecomment-835508651 #[std::prelude::v1::test] fn trait_method() { trait A { fn f(self); fn g(self); } impl A for () { #[tokio::main] async fn f(self) { self.g() } fn g(self) {} } ().f() } // https://github.com/tokio-rs/tokio/issues/4175 #[tokio::main] pub async fn issue_4175_main_1() -> ! { panic!(); } #[tokio::main] pub async fn issue_4175_main_2() -> std::io::Result<()> { panic!(); } #[allow(unreachable_code)] #[tokio::test] pub async fn issue_4175_test() -> std::io::Result<()> { return Ok(()); panic!(); } // https://github.com/tokio-rs/tokio/issues/4175 #[allow(clippy::let_unit_value)] pub mod clippy_semicolon_if_nothing_returned { #![deny(clippy::semicolon_if_nothing_returned)] #[tokio::main] pub async fn local() { let _x = (); } #[tokio::main] pub async fn item() { fn _f() {} } #[tokio::main] pub async fn semi() { panic!(); } #[tokio::main] pub async fn empty() { // To trigger clippy::semicolon_if_nothing_returned lint, the block needs to contain newline. } } // https://github.com/tokio-rs/tokio/issues/5243 pub mod issue_5243 { macro_rules! mac { (async fn $name:ident() $b:block) => { #[::tokio::test] async fn $name() { $b } }; } mac!( async fn foo() {} ); } #[cfg(tokio_unstable)] pub mod macro_rt_arg_unhandled_panic { use tokio_test::assert_err; #[tokio::test(flavor = "current_thread", unhandled_panic = "shutdown_runtime")] #[should_panic] async fn unhandled_panic_shutdown_runtime() { let _ = tokio::spawn(async { panic!("This panic should shutdown the runtime."); }) .await; } #[tokio::test(flavor = "current_thread", unhandled_panic = "ignore")] async fn unhandled_panic_ignore() { let rt = tokio::spawn(async { panic!("This panic should be forwarded to rt as an error."); }) .await; assert_err!(rt); } } tokio-1.43.1/tests/macros_try_join.rs000064400000000000000000000111161046102023000157500ustar 00000000000000#![cfg(feature = "macros")] #![allow(clippy::disallowed_names)] use std::sync::Arc; use tokio::sync::{oneshot, Semaphore}; use tokio_test::{assert_pending, assert_ready, task}; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; #[maybe_tokio_test] async fn sync_one_lit_expr_comma() { let foo = tokio::try_join!(async { ok(1) },); assert_eq!(foo, Ok((1,))); } #[maybe_tokio_test] async fn sync_one_lit_expr_no_comma() { let foo = tokio::try_join!(async { ok(1) }); assert_eq!(foo, Ok((1,))); } #[maybe_tokio_test] async fn sync_two_lit_expr_comma() { let foo = tokio::try_join!(async { ok(1) }, async { ok(2) },); assert_eq!(foo, Ok((1, 2))); } #[maybe_tokio_test] async fn sync_two_lit_expr_no_comma() { let foo = tokio::try_join!(async { ok(1) }, async { ok(2) }); assert_eq!(foo, Ok((1, 2))); } #[maybe_tokio_test] async fn two_await() { let (tx1, rx1) = oneshot::channel::<&str>(); let (tx2, rx2) = oneshot::channel::(); let mut join = task::spawn(async { tokio::try_join!(rx1, rx2) }); assert_pending!(join.poll()); tx2.send(123).unwrap(); assert!(join.is_woken()); assert_pending!(join.poll()); tx1.send("hello").unwrap(); assert!(join.is_woken()); let res: Result<(&str, u32), _> = assert_ready!(join.poll()); assert_eq!(Ok(("hello", 123)), res); } #[maybe_tokio_test] async fn err_abort_early() { let (tx1, rx1) = oneshot::channel::<&str>(); let (tx2, rx2) = oneshot::channel::(); let (_tx3, rx3) = oneshot::channel::(); let mut join = task::spawn(async { tokio::try_join!(rx1, rx2, rx3) }); assert_pending!(join.poll()); tx2.send(123).unwrap(); assert!(join.is_woken()); assert_pending!(join.poll()); drop(tx1); assert!(join.is_woken()); let res = assert_ready!(join.poll()); assert!(res.is_err()); } #[test] #[cfg(target_pointer_width = "64")] fn join_size() { use futures::future; use std::mem; let fut = async { let ready = future::ready(ok(0i32)); tokio::try_join!(ready) }; assert_eq!(mem::size_of_val(&fut), 32); let fut = async { let ready1 = future::ready(ok(0i32)); let ready2 = future::ready(ok(0i32)); tokio::try_join!(ready1, ready2) }; assert_eq!(mem::size_of_val(&fut), 48); } fn ok(val: T) -> Result { Ok(val) } async fn non_cooperative_task(permits: Arc) -> Result { let mut exceeded_budget = 0; for _ in 0..5 { // Another task should run after this task uses its whole budget for _ in 0..128 { let _permit = permits.clone().acquire_owned().await.unwrap(); } exceeded_budget += 1; } Ok(exceeded_budget) } async fn poor_little_task(permits: Arc) -> Result { let mut how_many_times_i_got_to_run = 0; for _ in 0..5 { let _permit = permits.clone().acquire_owned().await.unwrap(); how_many_times_i_got_to_run += 1; } Ok(how_many_times_i_got_to_run) } #[tokio::test] async fn try_join_does_not_allow_tasks_to_starve() { let permits = Arc::new(Semaphore::new(10)); // non_cooperative_task should yield after its budget is exceeded and then poor_little_task should run. let result = tokio::try_join!( non_cooperative_task(Arc::clone(&permits)), poor_little_task(permits) ); let (non_cooperative_result, little_task_result) = result.unwrap(); assert_eq!(5, non_cooperative_result); assert_eq!(5, little_task_result); } #[tokio::test] async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { let poll_order = Arc::new(std::sync::Mutex::new(vec![])); let fut = |x, poll_order: Arc>>| async move { for _ in 0..4 { { let mut guard = poll_order.lock().unwrap(); guard.push(x); } tokio::task::yield_now().await; } }; tokio::join!( fut(1, Arc::clone(&poll_order)), fut(2, Arc::clone(&poll_order)), fut(3, Arc::clone(&poll_order)), ); // Each time the future created by join! is polled, it should start // by polling a different future first. assert_eq!( vec![1, 2, 3, 2, 3, 1, 3, 1, 2, 1, 2, 3], *poll_order.lock().unwrap() ); } #[tokio::test] async fn empty_try_join() { assert_eq!(tokio::try_join!() as Result<_, ()>, Ok(())); } tokio-1.43.1/tests/net_bind_resource.rs000064400000000000000000000006671046102023000162510ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery or bind use tokio::net::TcpListener; use std::net; #[test] #[should_panic] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn no_runtime_panics_binding_net_tcp_listener() { let listener = net::TcpListener::bind("127.0.0.1:0").expect("failed to bind listener"); let _ = TcpListener::try_from(listener); } tokio-1.43.1/tests/net_lookup_host.rs000064400000000000000000000022151046102023000157630ustar 00000000000000#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support direct socket operations use tokio::net; use tokio_test::assert_ok; use std::io; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; #[tokio::test] async fn lookup_socket_addr() { let addr: SocketAddr = "127.0.0.1:8000".parse().unwrap(); let actual = assert_ok!(net::lookup_host(addr).await).collect::>(); assert_eq!(vec![addr], actual); } #[tokio::test] async fn lookup_str_socket_addr() { let addr: SocketAddr = "127.0.0.1:8000".parse().unwrap(); let actual = assert_ok!(net::lookup_host("127.0.0.1:8000").await).collect::>(); assert_eq!(vec![addr], actual); } #[tokio::test] #[cfg_attr(miri, ignore)] // No `getaddrinfo` in miri. async fn resolve_dns() -> io::Result<()> { let mut hosts = net::lookup_host("localhost:3000").await?; let host = hosts.next().unwrap(); let expected = if host.is_ipv4() { SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 3000) } else { SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 3000) }; assert_eq!(host, expected); Ok(()) } tokio-1.43.1/tests/net_named_pipe.rs000064400000000000000000000315071046102023000155240ustar 00000000000000#![cfg(feature = "full")] #![cfg(windows)] use std::io; use std::time::Duration; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::windows::named_pipe::{ClientOptions, PipeMode, ServerOptions}; use tokio::time; use windows_sys::Win32::Foundation::{ERROR_NO_DATA, ERROR_PIPE_BUSY}; #[tokio::test] async fn test_named_pipe_client_drop() -> io::Result<()> { const PIPE_NAME: &str = r"\\.\pipe\test-named-pipe-client-drop"; let mut server = ServerOptions::new().create(PIPE_NAME)?; let client = ClientOptions::new().open(PIPE_NAME)?; server.connect().await?; drop(client); // instance will be broken because client is gone match server.write_all(b"ping").await { Err(e) if e.raw_os_error() == Some(ERROR_NO_DATA as i32) => (), x => panic!("{:?}", x), } Ok(()) } #[tokio::test] async fn test_named_pipe_single_client() -> io::Result<()> { use tokio::io::{AsyncBufReadExt as _, BufReader}; const PIPE_NAME: &str = r"\\.\pipe\test-named-pipe-single-client"; let server = ServerOptions::new().create(PIPE_NAME)?; let server = tokio::spawn(async move { // Note: we wait for a client to connect. server.connect().await?; let mut server = BufReader::new(server); let mut buf = String::new(); server.read_line(&mut buf).await?; server.write_all(b"pong\n").await?; Ok::<_, io::Error>(buf) }); let client = tokio::spawn(async move { let client = ClientOptions::new().open(PIPE_NAME)?; let mut client = BufReader::new(client); let mut buf = String::new(); client.write_all(b"ping\n").await?; client.read_line(&mut buf).await?; Ok::<_, io::Error>(buf) }); let (server, client) = tokio::try_join!(server, client)?; assert_eq!(server?, "ping\n"); assert_eq!(client?, "pong\n"); Ok(()) } #[tokio::test] async fn test_named_pipe_multi_client() -> io::Result<()> { use tokio::io::{AsyncBufReadExt as _, BufReader}; const PIPE_NAME: &str = r"\\.\pipe\test-named-pipe-multi-client"; const N: usize = 10; // The first server needs to be constructed early so that clients can // be correctly connected. Otherwise calling .wait will cause the client to // error. let mut server = ServerOptions::new().create(PIPE_NAME)?; let server = tokio::spawn(async move { for _ in 0..N { // Wait for client to connect. server.connect().await?; let mut inner = BufReader::new(server); // Construct the next server to be connected before sending the one // we already have of onto a task. This ensures that the server // isn't closed (after it's done in the task) before a new one is // available. Otherwise the client might error with // `io::ErrorKind::NotFound`. server = ServerOptions::new().create(PIPE_NAME)?; tokio::spawn(async move { let mut buf = String::new(); inner.read_line(&mut buf).await?; inner.write_all(b"pong\n").await?; inner.flush().await?; Ok::<_, io::Error>(()) }); } Ok::<_, io::Error>(()) }); let mut clients = Vec::new(); for _ in 0..N { clients.push(tokio::spawn(async move { // This showcases a generic connect loop. // // We immediately try to create a client, if it's not found or the // pipe is busy we use the specialized wait function on the client // builder. let client = loop { match ClientOptions::new().open(PIPE_NAME) { Ok(client) => break client, Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY as i32) => (), Err(e) if e.kind() == io::ErrorKind::NotFound => (), Err(e) => return Err(e), } // Wait for a named pipe to become available. time::sleep(Duration::from_millis(10)).await; }; let mut client = BufReader::new(client); let mut buf = String::new(); client.write_all(b"ping\n").await?; client.flush().await?; client.read_line(&mut buf).await?; Ok::<_, io::Error>(buf) })); } for client in clients { let result = client.await?; assert_eq!(result?, "pong\n"); } server.await??; Ok(()) } #[tokio::test] async fn test_named_pipe_multi_client_ready() -> io::Result<()> { use tokio::io::Interest; const PIPE_NAME: &str = r"\\.\pipe\test-named-pipe-multi-client-ready"; const N: usize = 10; // The first server needs to be constructed early so that clients can // be correctly connected. Otherwise calling .wait will cause the client to // error. let mut server = ServerOptions::new().create(PIPE_NAME)?; let server = tokio::spawn(async move { for _ in 0..N { // Wait for client to connect. server.connect().await?; let inner_server = server; // Construct the next server to be connected before sending the one // we already have of onto a task. This ensures that the server // isn't closed (after it's done in the task) before a new one is // available. Otherwise the client might error with // `io::ErrorKind::NotFound`. server = ServerOptions::new().create(PIPE_NAME)?; tokio::spawn(async move { let server = inner_server; { let mut read_buf = [0u8; 5]; let mut read_buf_cursor = 0; loop { server.readable().await?; let buf = &mut read_buf[read_buf_cursor..]; match server.try_read(buf) { Ok(n) => { read_buf_cursor += n; if read_buf_cursor == read_buf.len() { break; } } Err(e) if e.kind() == io::ErrorKind::WouldBlock => { continue; } Err(e) => { return Err(e); } } } }; { let write_buf = b"pong\n"; let mut write_buf_cursor = 0; loop { server.writable().await?; let buf = &write_buf[write_buf_cursor..]; match server.try_write(buf) { Ok(n) => { write_buf_cursor += n; if write_buf_cursor == write_buf.len() { break; } } Err(e) if e.kind() == io::ErrorKind::WouldBlock => { continue; } Err(e) => { return Err(e); } } } } Ok::<_, io::Error>(()) }); } Ok::<_, io::Error>(()) }); let mut clients = Vec::new(); for _ in 0..N { clients.push(tokio::spawn(async move { // This showcases a generic connect loop. // // We immediately try to create a client, if it's not found or the // pipe is busy we use the specialized wait function on the client // builder. let client = loop { match ClientOptions::new().open(PIPE_NAME) { Ok(client) => break client, Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY as i32) => (), Err(e) if e.kind() == io::ErrorKind::NotFound => (), Err(e) => return Err(e), } // Wait for a named pipe to become available. time::sleep(Duration::from_millis(10)).await; }; let mut read_buf = [0u8; 5]; let mut read_buf_cursor = 0; let write_buf = b"ping\n"; let mut write_buf_cursor = 0; loop { let mut interest = Interest::READABLE; if write_buf_cursor < write_buf.len() { interest |= Interest::WRITABLE; } let ready = client.ready(interest).await?; if ready.is_readable() { let buf = &mut read_buf[read_buf_cursor..]; match client.try_read(buf) { Ok(n) => { read_buf_cursor += n; if read_buf_cursor == read_buf.len() { break; } } Err(e) if e.kind() == io::ErrorKind::WouldBlock => { continue; } Err(e) => { return Err(e); } } } if ready.is_writable() { let buf = &write_buf[write_buf_cursor..]; if buf.is_empty() { continue; } match client.try_write(buf) { Ok(n) => { write_buf_cursor += n; } Err(e) if e.kind() == io::ErrorKind::WouldBlock => { continue; } Err(e) => { return Err(e); } } } } let buf = String::from_utf8_lossy(&read_buf).into_owned(); Ok::<_, io::Error>(buf) })); } for client in clients { let result = client.await?; assert_eq!(result?, "pong\n"); } server.await??; Ok(()) } // This tests that message mode works as expected. #[tokio::test] async fn test_named_pipe_mode_message() -> io::Result<()> { // it's easy to accidentally get a seemingly working test here because byte pipes // often return contents at write boundaries. to make sure we're doing the right thing we // explicitly test that it doesn't work in byte mode. _named_pipe_mode_message(PipeMode::Message).await?; _named_pipe_mode_message(PipeMode::Byte).await } async fn _named_pipe_mode_message(mode: PipeMode) -> io::Result<()> { let pipe_name = format!( r"\\.\pipe\test-named-pipe-mode-message-{}", matches!(mode, PipeMode::Message) ); let mut buf = [0u8; 32]; let mut server = ServerOptions::new() .first_pipe_instance(true) .pipe_mode(mode) .create(&pipe_name)?; let mut client = ClientOptions::new().pipe_mode(mode).open(&pipe_name)?; server.connect().await?; // this needs a few iterations, presumably Windows waits for a few calls before merging buffers for _ in 0..10 { client.write_all(b"hello").await?; server.write_all(b"world").await?; } for _ in 0..10 { let n = server.read(&mut buf).await?; if buf[..n] != b"hello"[..] { assert!(matches!(mode, PipeMode::Byte)); return Ok(()); } let n = client.read(&mut buf).await?; if buf[..n] != b"world"[..] { assert!(matches!(mode, PipeMode::Byte)); return Ok(()); } } // byte mode should have errored before. assert!(matches!(mode, PipeMode::Message)); Ok(()) } // This tests `NamedPipeServer::connect` with various access settings. #[tokio::test] async fn test_named_pipe_access() -> io::Result<()> { const PIPE_NAME: &str = r"\\.\pipe\test-named-pipe-access"; for (inb, outb) in [(true, true), (true, false), (false, true)] { let (tx, rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { let s = ServerOptions::new() .access_inbound(inb) .access_outbound(outb) .create(PIPE_NAME)?; let mut connect_fut = tokio_test::task::spawn(s.connect()); assert!(connect_fut.poll().is_pending()); tx.send(()).unwrap(); connect_fut.await }); // Wait for the server to call connect. rx.await.unwrap(); let _ = ClientOptions::new().read(outb).write(inb).open(PIPE_NAME)?; server.await??; } Ok(()) } tokio-1.43.1/tests/net_panic.rs000064400000000000000000000130261046102023000145110ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] #![cfg(panic = "unwind")] use std::error::Error; use tokio::net::{TcpListener, TcpStream}; use tokio::runtime::{Builder, Runtime}; mod support { pub mod panic; } use support::panic::test_panic; #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn udp_socket_from_std_panic_caller() -> Result<(), Box> { use std::net::SocketAddr; use tokio::net::UdpSocket; let addr = "127.0.0.1:0".parse::().unwrap(); let std_sock = std::net::UdpSocket::bind(addr).unwrap(); std_sock.set_nonblocking(true).unwrap(); let panic_location_file = test_panic(|| { let rt = runtime_without_io(); rt.block_on(async { let _sock = UdpSocket::from_std(std_sock); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn tcp_listener_from_std_panic_caller() -> Result<(), Box> { let std_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); std_listener.set_nonblocking(true).unwrap(); let panic_location_file = test_panic(|| { let rt = runtime_without_io(); rt.block_on(async { let _ = TcpListener::from_std(std_listener); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn tcp_stream_from_std_panic_caller() -> Result<(), Box> { let std_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); let std_stream = std::net::TcpStream::connect(std_listener.local_addr().unwrap()).unwrap(); std_stream.set_nonblocking(true).unwrap(); let panic_location_file = test_panic(|| { let rt = runtime_without_io(); rt.block_on(async { let _ = TcpStream::from_std(std_stream); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn unix_listener_bind_panic_caller() -> Result<(), Box> { use tokio::net::UnixListener; let dir = tempfile::tempdir().unwrap(); let sock_path = dir.path().join("socket"); let panic_location_file = test_panic(|| { let rt = runtime_without_io(); rt.block_on(async { let _ = UnixListener::bind(&sock_path); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn unix_listener_from_std_panic_caller() -> Result<(), Box> { use tokio::net::UnixListener; let dir = tempfile::tempdir().unwrap(); let sock_path = dir.path().join("socket"); let std_listener = std::os::unix::net::UnixListener::bind(sock_path).unwrap(); let panic_location_file = test_panic(|| { let rt = runtime_without_io(); rt.block_on(async { let _ = UnixListener::from_std(std_listener); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn unix_stream_from_std_panic_caller() -> Result<(), Box> { use tokio::net::UnixStream; let dir = tempfile::tempdir().unwrap(); let sock_path = dir.path().join("socket"); let _std_listener = std::os::unix::net::UnixListener::bind(&sock_path).unwrap(); let std_stream = std::os::unix::net::UnixStream::connect(&sock_path).unwrap(); let panic_location_file = test_panic(|| { let rt = runtime_without_io(); rt.block_on(async { let _ = UnixStream::from_std(std_stream); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(unix)] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn unix_datagram_from_std_panic_caller() -> Result<(), Box> { use std::os::unix::net::UnixDatagram as StdUDS; use tokio::net::UnixDatagram; let dir = tempfile::tempdir().unwrap(); let sock_path = dir.path().join("socket"); // Bind the socket to a filesystem path // /let socket_path = tmp.path().join("socket"); let std_socket = StdUDS::bind(sock_path).unwrap(); std_socket.set_nonblocking(true).unwrap(); let panic_location_file = test_panic(move || { let rt = runtime_without_io(); rt.block_on(async { let _ = UnixDatagram::from_std(std_socket); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg(windows)] fn server_options_max_instances_panic_caller() -> Result<(), Box> { use tokio::net::windows::named_pipe::ServerOptions; let panic_location_file = test_panic(move || { let rt = runtime_without_io(); rt.block_on(async { let mut options = ServerOptions::new(); options.max_instances(255); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } // Runtime without `enable_io` so it has no IO driver set. fn runtime_without_io() -> Runtime { Builder::new_current_thread().build().unwrap() } tokio-1.43.1/tests/net_unix_pipe.rs000064400000000000000000000406711046102023000154250ustar 00000000000000#![cfg(feature = "full")] #![cfg(unix)] use tokio::io::{AsyncReadExt, AsyncWriteExt, Interest}; use tokio::net::unix::pipe; use tokio_test::task; use tokio_test::{assert_err, assert_ok, assert_pending, assert_ready_ok}; use std::fs::File; use std::io; use std::os::unix::fs::OpenOptionsExt; use std::os::unix::io::AsRawFd; use std::path::{Path, PathBuf}; /// Helper struct which will clean up temporary files once dropped. struct TempFifo { path: PathBuf, _dir: tempfile::TempDir, } impl TempFifo { fn new(name: &str) -> io::Result { let dir = tempfile::Builder::new() .prefix("tokio-fifo-tests") .tempdir()?; let path = dir.path().join(name); nix::unistd::mkfifo(&path, nix::sys::stat::Mode::S_IRWXU)?; Ok(TempFifo { path, _dir: dir }) } } impl AsRef for TempFifo { fn as_ref(&self) -> &Path { self.path.as_ref() } } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn fifo_simple_send() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; let fifo = TempFifo::new("simple_send")?; // Create a reading task which should wait for data from the pipe. let mut reader = pipe::OpenOptions::new().open_receiver(&fifo)?; let mut read_fut = task::spawn(async move { let mut buf = vec![0; DATA.len()]; reader.read_exact(&mut buf).await?; Ok::<_, io::Error>(buf) }); assert_pending!(read_fut.poll()); let mut writer = pipe::OpenOptions::new().open_sender(&fifo)?; writer.write_all(DATA).await?; // Let the IO driver poll events for the reader. while !read_fut.is_woken() { tokio::task::yield_now().await; } // Reading task should be ready now. let read_data = assert_ready_ok!(read_fut.poll()); assert_eq!(&read_data, DATA); Ok(()) } #[tokio::test] #[cfg(target_os = "linux")] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn fifo_simple_send_sender_first() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; // Create a new fifo file with *no reading ends open*. let fifo = TempFifo::new("simple_send_sender_first")?; // Simple `open_sender` should fail with ENXIO (no such device or address). let err = assert_err!(pipe::OpenOptions::new().open_sender(&fifo)); assert_eq!(err.raw_os_error(), Some(libc::ENXIO)); // `open_sender` in read-write mode should succeed and the pipe should be ready to write. let mut writer = pipe::OpenOptions::new() .read_write(true) .open_sender(&fifo)?; writer.write_all(DATA).await?; // Read the written data and validate. let mut reader = pipe::OpenOptions::new().open_receiver(&fifo)?; let mut read_data = vec![0; DATA.len()]; reader.read_exact(&mut read_data).await?; assert_eq!(&read_data, DATA); Ok(()) } // Opens a FIFO file, write and *close the writer*. async fn write_and_close(path: impl AsRef, msg: &[u8]) -> io::Result<()> { let mut writer = pipe::OpenOptions::new().open_sender(path)?; writer.write_all(msg).await?; drop(writer); // Explicit drop. Ok(()) } /// Checks EOF behavior with single reader and writers sequentially opening /// and closing a FIFO. #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn fifo_multiple_writes() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; let fifo = TempFifo::new("fifo_multiple_writes")?; let mut reader = pipe::OpenOptions::new().open_receiver(&fifo)?; write_and_close(&fifo, DATA).await?; let ev = reader.ready(Interest::READABLE).await?; assert!(ev.is_readable()); let mut read_data = vec![0; DATA.len()]; assert_ok!(reader.read_exact(&mut read_data).await); // Check that reader hits EOF. let err = assert_err!(reader.read_exact(&mut read_data).await); assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); // Write more data and read again. write_and_close(&fifo, DATA).await?; assert_ok!(reader.read_exact(&mut read_data).await); Ok(()) } /// Checks behavior of a resilient reader (Receiver in O_RDWR access mode) /// with writers sequentially opening and closing a FIFO. #[tokio::test] #[cfg(target_os = "linux")] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn fifo_resilient_reader() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; let fifo = TempFifo::new("fifo_resilient_reader")?; // Open reader in read-write access mode. let mut reader = pipe::OpenOptions::new() .read_write(true) .open_receiver(&fifo)?; write_and_close(&fifo, DATA).await?; let ev = reader.ready(Interest::READABLE).await?; let mut read_data = vec![0; DATA.len()]; reader.read_exact(&mut read_data).await?; // Check that reader didn't hit EOF. assert!(!ev.is_read_closed()); // Resilient reader can asynchronously wait for the next writer. let mut second_read_fut = task::spawn(reader.read_exact(&mut read_data)); assert_pending!(second_read_fut.poll()); // Write more data and read again. write_and_close(&fifo, DATA).await?; assert_ok!(second_read_fut.await); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `O_NONBLOCK` for open64 in miri. async fn open_detects_not_a_fifo() -> io::Result<()> { let dir = tempfile::Builder::new() .prefix("tokio-fifo-tests") .tempdir() .unwrap(); let path = dir.path().join("not_a_fifo"); // Create an ordinary file. File::create(&path)?; // Check if Sender detects invalid file type. let err = assert_err!(pipe::OpenOptions::new().open_sender(&path)); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); // Check if Receiver detects invalid file type. let err = assert_err!(pipe::OpenOptions::new().open_sender(&path)); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn from_file() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; let fifo = TempFifo::new("from_file")?; // Construct a Receiver from a File. let file = std::fs::OpenOptions::new() .read(true) .custom_flags(libc::O_NONBLOCK) .open(&fifo)?; let mut reader = pipe::Receiver::from_file(file)?; // Construct a Sender from a File. let file = std::fs::OpenOptions::new() .write(true) .custom_flags(libc::O_NONBLOCK) .open(&fifo)?; let mut writer = pipe::Sender::from_file(file)?; // Write and read some data to test async. let mut read_fut = task::spawn(async move { let mut buf = vec![0; DATA.len()]; reader.read_exact(&mut buf).await?; Ok::<_, io::Error>(buf) }); assert_pending!(read_fut.poll()); writer.write_all(DATA).await?; let read_data = assert_ok!(read_fut.await); assert_eq!(&read_data, DATA); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `fstat` in miri. async fn from_file_detects_not_a_fifo() -> io::Result<()> { let dir = tempfile::Builder::new() .prefix("tokio-fifo-tests") .tempdir() .unwrap(); let path = dir.path().join("not_a_fifo"); // Create an ordinary file. File::create(&path)?; // Check if Sender detects invalid file type. let file = std::fs::OpenOptions::new().write(true).open(&path)?; let err = assert_err!(pipe::Sender::from_file(file)); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); // Check if Receiver detects invalid file type. let file = std::fs::OpenOptions::new().read(true).open(&path)?; let err = assert_err!(pipe::Receiver::from_file(file)); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn from_file_detects_wrong_access_mode() -> io::Result<()> { let fifo = TempFifo::new("wrong_access_mode")?; // Open a read end to open the fifo for writing. let _reader = pipe::OpenOptions::new().open_receiver(&fifo)?; // Check if Receiver detects write-only access mode. let wronly = std::fs::OpenOptions::new() .write(true) .custom_flags(libc::O_NONBLOCK) .open(&fifo)?; let err = assert_err!(pipe::Receiver::from_file(wronly)); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); // Check if Sender detects read-only access mode. let rdonly = std::fs::OpenOptions::new() .read(true) .custom_flags(libc::O_NONBLOCK) .open(&fifo)?; let err = assert_err!(pipe::Sender::from_file(rdonly)); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); Ok(()) } fn is_nonblocking(fd: &T) -> io::Result { let flags = nix::fcntl::fcntl(fd.as_raw_fd(), nix::fcntl::F_GETFL)?; Ok((flags & libc::O_NONBLOCK) != 0) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn from_file_sets_nonblock() -> io::Result<()> { let fifo = TempFifo::new("sets_nonblock")?; // Open read and write ends to let blocking files open. let _reader = pipe::OpenOptions::new().open_receiver(&fifo)?; let _writer = pipe::OpenOptions::new().open_sender(&fifo)?; // Check if Receiver sets the pipe in non-blocking mode. let rdonly = std::fs::OpenOptions::new().read(true).open(&fifo)?; assert!(!is_nonblocking(&rdonly)?); let reader = pipe::Receiver::from_file(rdonly)?; assert!(is_nonblocking(&reader)?); // Check if Sender sets the pipe in non-blocking mode. let wronly = std::fs::OpenOptions::new().write(true).open(&fifo)?; assert!(!is_nonblocking(&wronly)?); let writer = pipe::Sender::from_file(wronly)?; assert!(is_nonblocking(&writer)?); Ok(()) } fn writable_by_poll(writer: &pipe::Sender) -> bool { task::spawn(writer.writable()).poll().is_ready() } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn try_read_write() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; // Create a pipe pair over a fifo file. let fifo = TempFifo::new("try_read_write")?; let reader = pipe::OpenOptions::new().open_receiver(&fifo)?; let writer = pipe::OpenOptions::new().open_sender(&fifo)?; // Fill the pipe buffer with `try_write`. let mut write_data = Vec::new(); while writable_by_poll(&writer) { match writer.try_write(DATA) { Ok(n) => write_data.extend(&DATA[..n]), Err(e) => { assert_eq!(e.kind(), io::ErrorKind::WouldBlock); break; } } } // Drain the pipe buffer with `try_read`. let mut read_data = vec![0; write_data.len()]; let mut i = 0; while i < write_data.len() { reader.readable().await?; match reader.try_read(&mut read_data[i..]) { Ok(n) => i += n, Err(e) => { assert_eq!(e.kind(), io::ErrorKind::WouldBlock); continue; } } } assert_eq!(read_data, write_data); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn try_read_write_vectored() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; // Create a pipe pair over a fifo file. let fifo = TempFifo::new("try_read_write_vectored")?; let reader = pipe::OpenOptions::new().open_receiver(&fifo)?; let writer = pipe::OpenOptions::new().open_sender(&fifo)?; let write_bufs: Vec<_> = DATA.chunks(3).map(io::IoSlice::new).collect(); // Fill the pipe buffer with `try_write_vectored`. let mut write_data = Vec::new(); while writable_by_poll(&writer) { match writer.try_write_vectored(&write_bufs) { Ok(n) => write_data.extend(&DATA[..n]), Err(e) => { assert_eq!(e.kind(), io::ErrorKind::WouldBlock); break; } } } // Drain the pipe buffer with `try_read_vectored`. let mut read_data = vec![0; write_data.len()]; let mut i = 0; while i < write_data.len() { reader.readable().await?; let mut read_bufs: Vec<_> = read_data[i..] .chunks_mut(0x10000) .map(io::IoSliceMut::new) .collect(); match reader.try_read_vectored(&mut read_bufs) { Ok(n) => i += n, Err(e) => { assert_eq!(e.kind(), io::ErrorKind::WouldBlock); continue; } } } assert_eq!(read_data, write_data); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `mkfifo` in miri. async fn try_read_buf() -> std::io::Result<()> { const DATA: &[u8] = b"this is some data to write to the fifo"; // Create a pipe pair over a fifo file. let fifo = TempFifo::new("try_read_write_vectored")?; let reader = pipe::OpenOptions::new().open_receiver(&fifo)?; let writer = pipe::OpenOptions::new().open_sender(&fifo)?; // Fill the pipe buffer with `try_write`. let mut write_data = Vec::new(); while writable_by_poll(&writer) { match writer.try_write(DATA) { Ok(n) => write_data.extend(&DATA[..n]), Err(e) => { assert_eq!(e.kind(), io::ErrorKind::WouldBlock); break; } } } // Drain the pipe buffer with `try_read_buf`. let mut read_data = vec![0; write_data.len()]; let mut i = 0; while i < write_data.len() { reader.readable().await?; match reader.try_read_buf(&mut read_data) { Ok(n) => i += n, Err(e) => { assert_eq!(e.kind(), io::ErrorKind::WouldBlock); continue; } } } assert_eq!(read_data, write_data); Ok(()) } #[tokio::test] async fn anon_pipe_simple_send() -> io::Result<()> { const DATA: &[u8] = b"this is some data to write to the pipe"; let (mut writer, mut reader) = pipe::pipe()?; // Create a reading task which should wait for data from the pipe. let mut read_fut = task::spawn(async move { let mut buf = vec![0; DATA.len()]; reader.read_exact(&mut buf).await?; Ok::<_, io::Error>(buf) }); assert_pending!(read_fut.poll()); writer.write_all(DATA).await?; // Let the IO driver poll events for the reader. while !read_fut.is_woken() { tokio::task::yield_now().await; } // Reading task should be ready now. let read_data = assert_ready_ok!(read_fut.poll()); assert_eq!(&read_data, DATA); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn anon_pipe_spawn_echo() -> std::io::Result<()> { use tokio::process::Command; const DATA: &str = "this is some data to write to the pipe"; let (tx, mut rx) = pipe::pipe()?; let status = Command::new("echo") .arg("-n") .arg(DATA) .stdout(tx.into_blocking_fd()?) .status(); let mut buf = vec![0; DATA.len()]; rx.read_exact(&mut buf).await?; assert_eq!(String::from_utf8(buf).unwrap(), DATA); let exit_code = status.await?; assert!(exit_code.success()); // Check if the pipe is closed. buf = Vec::new(); let total = assert_ok!(rx.try_read(&mut buf)); assert_eq!(total, 0); Ok(()) } #[tokio::test] #[cfg(target_os = "linux")] #[cfg_attr(miri, ignore)] // No `fstat` in miri. async fn anon_pipe_from_owned_fd() -> std::io::Result<()> { use nix::fcntl::OFlag; const DATA: &[u8] = b"this is some data to write to the pipe"; let (rx_fd, tx_fd) = nix::unistd::pipe2(OFlag::O_CLOEXEC | OFlag::O_NONBLOCK)?; let mut rx = pipe::Receiver::from_owned_fd(rx_fd)?; let mut tx = pipe::Sender::from_owned_fd(tx_fd)?; let mut buf = vec![0; DATA.len()]; tx.write_all(DATA).await?; rx.read_exact(&mut buf).await?; assert_eq!(buf, DATA); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn anon_pipe_into_nonblocking_fd() -> std::io::Result<()> { let (tx, rx) = pipe::pipe()?; let tx_fd = tx.into_nonblocking_fd()?; let rx_fd = rx.into_nonblocking_fd()?; assert!(is_nonblocking(&tx_fd)?); assert!(is_nonblocking(&rx_fd)?); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri. async fn anon_pipe_into_blocking_fd() -> std::io::Result<()> { let (tx, rx) = pipe::pipe()?; let tx_fd = tx.into_blocking_fd()?; let rx_fd = rx.into_blocking_fd()?; assert!(!is_nonblocking(&tx_fd)?); assert!(!is_nonblocking(&rx_fd)?); Ok(()) } tokio-1.43.1/tests/no_rt.rs000064400000000000000000000024111046102023000136660ustar 00000000000000#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use tokio::net::TcpStream; use tokio::sync::oneshot; use tokio::time::{timeout, Duration}; use futures::executor::block_on; use std::net::TcpListener; #[test] #[should_panic( expected = "there is no reactor running, must be called from the context of a Tokio 1.x runtime" )] fn timeout_panics_when_no_tokio_context() { block_on(timeout_value()); } #[test] #[should_panic( expected = "there is no reactor running, must be called from the context of a Tokio 1.x runtime" )] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn panics_when_no_reactor() { let srv = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = srv.local_addr().unwrap(); block_on(TcpStream::connect(&addr)).unwrap(); } async fn timeout_value() { let (_tx, rx) = oneshot::channel::<()>(); let dur = Duration::from_millis(10); let _ = timeout(dur, rx).await; } #[test] #[should_panic( expected = "there is no reactor running, must be called from the context of a Tokio 1.x runtime" )] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn io_panics_when_no_tokio_context() { let _ = tokio::net::TcpListener::from_std(std::net::TcpListener::bind("127.0.0.1:0").unwrap()); } tokio-1.43.1/tests/process_arg0.rs000064400000000000000000000005111046102023000151330ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", unix, not(miri)))] use tokio::process::Command; #[tokio::test] async fn arg0() { let mut cmd = Command::new("sh"); cmd.arg0("test_string").arg("-c").arg("echo $0"); let output = cmd.output().await.unwrap(); assert_eq!(output.stdout, b"test_string\n"); } tokio-1.43.1/tests/process_change_of_runtime.rs000064400000000000000000000020371046102023000177630ustar 00000000000000#![cfg(feature = "process")] #![warn(rust_2018_idioms)] // This tests test the behavior of `process::Command::spawn` when it is used // outside runtime, and when `process::Child::wait ` is used in a different // runtime from which `process::Command::spawn` is used. #![cfg(all(unix, not(target_os = "freebsd"), not(miri)))] use std::process::Stdio; use tokio::{process::Command, runtime::Runtime}; #[test] fn process_spawned_and_wait_in_different_runtime() { let mut child = Runtime::new().unwrap().block_on(async { Command::new("true") .stdin(Stdio::piped()) .stdout(Stdio::null()) .spawn() .unwrap() }); Runtime::new().unwrap().block_on(async { let _ = child.wait().await.unwrap(); }); } #[test] #[should_panic( expected = "there is no reactor running, must be called from the context of a Tokio 1.x runtime" )] fn process_spawned_outside_runtime() { let _ = Command::new("true") .stdin(Stdio::piped()) .stdout(Stdio::null()) .spawn(); } tokio-1.43.1/tests/process_issue_2174.rs000064400000000000000000000026071046102023000161170ustar 00000000000000#![cfg(feature = "process")] #![warn(rust_2018_idioms)] // This test reveals a difference in behavior of kqueue on FreeBSD. When the // reader disconnects, there does not seem to be an `EVFILT_WRITE` filter that // is returned. // // It is expected that `EVFILT_WRITE` would be returned with either the // `EV_EOF` or `EV_ERROR` flag set. If either flag is set a write would be // attempted, but that does not seem to occur. #![cfg(all(unix, not(target_os = "freebsd"), not(miri)))] use std::process::Stdio; use std::time::Duration; use tokio::io::AsyncWriteExt; use tokio::process::Command; use tokio::time; use tokio_test::assert_err; #[tokio::test] #[cfg_attr(panic = "abort", ignore)] async fn issue_2174() { let mut child = Command::new("sleep") .arg("2") .stdin(Stdio::piped()) .stdout(Stdio::null()) .spawn() .unwrap(); let mut input = child.stdin.take().unwrap(); // Writes will buffer up to 65_636. This *should* loop at least 8 times // and then register interest. let handle = tokio::spawn(async move { let data = [0u8; 8192]; loop { input.write_all(&data).await.unwrap(); } }); // Sleep enough time so that the child process's stdin's buffer fills. time::sleep(Duration::from_secs(1)).await; // Kill the child process. child.kill().await.unwrap(); assert_err!(handle.await); } tokio-1.43.1/tests/process_issue_42.rs000064400000000000000000000024111046102023000157400ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] use futures::future::join_all; use std::process::Stdio; use tokio::process::Command; use tokio::task; #[tokio::test] async fn issue_42() { // We spawn a many batches of processes which should exit at roughly the // same time (modulo OS scheduling delays), to make sure that consuming // a readiness event for one process doesn't inadvertently starve another. // We then do this many times (in parallel) in an effort to stress test the // implementation to ensure there are no race conditions. // See alexcrichton/tokio-process#42 for background let join_handles = (0..10usize).map(|_| { task::spawn(async { let processes = (0..10usize).map(|i| { let mut child = Command::new("echo") .arg(format!("I am spawned process #{i}")) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::null()) .kill_on_drop(true) .spawn() .unwrap(); async move { child.wait().await } }); join_all(processes).await; }) }); join_all(join_handles).await; } tokio-1.43.1/tests/process_kill_on_drop.rs000064400000000000000000000020621046102023000167600ustar 00000000000000#![cfg(all(unix, feature = "process", not(miri)))] #![warn(rust_2018_idioms)] use std::io::ErrorKind; use std::process::Stdio; use std::time::Duration; use tokio::io::AsyncReadExt; use tokio::process::Command; use tokio::time::sleep; use tokio_test::assert_ok; #[tokio::test] async fn kill_on_drop() { let mut cmd = Command::new("bash"); cmd.args([ "-c", " # Fork another child that won't get killed sh -c 'sleep 1; echo child ran' & disown -a # Await our death sleep 5 echo hello from beyond the grave ", ]); let e = cmd.kill_on_drop(true).stdout(Stdio::piped()).spawn(); if e.is_err() && e.as_ref().unwrap_err().kind() == ErrorKind::NotFound { println!("bash not available; skipping test"); return; } let mut child = e.unwrap(); sleep(Duration::from_secs(2)).await; let mut out = child.stdout.take().unwrap(); drop(child); let mut msg = String::new(); assert_ok!(out.read_to_string(&mut msg).await); assert_eq!("child ran\n", msg); } tokio-1.43.1/tests/process_raw_handle.rs000064400000000000000000000011421046102023000164070ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(windows)] #![cfg(not(miri))] use tokio::process::Command; use windows_sys::Win32::System::Threading::GetProcessId; #[tokio::test] async fn obtain_raw_handle() { let mut cmd = Command::new("cmd"); cmd.kill_on_drop(true); cmd.arg("/c"); cmd.arg("pause"); let child = cmd.spawn().unwrap(); let orig_id = child.id().expect("missing id"); assert!(orig_id > 0); let handle = child.raw_handle().expect("process stopped"); let handled_id = unsafe { GetProcessId(handle as _) }; assert_eq!(handled_id, orig_id); } tokio-1.43.1/tests/process_smoke.rs000064400000000000000000000015541046102023000154300ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi cannot run system commands use tokio::process::Command; use tokio_test::assert_ok; #[tokio::test] async fn simple() { let mut cmd; if cfg!(windows) { cmd = Command::new("cmd"); cmd.arg("/c"); } else { cmd = Command::new("sh"); cmd.arg("-c"); } let mut child = cmd.arg("exit 2").spawn().unwrap(); let id = child.id().expect("missing id"); assert!(id > 0); let status = assert_ok!(child.wait().await); assert_eq!(status.code(), Some(2)); // test that the `.wait()` method is fused just like the stdlib let status = assert_ok!(child.wait().await); assert_eq!(status.code(), Some(2)); // Can't get id after process has exited assert_eq!(child.id(), None); drop(child.kill()); } tokio-1.43.1/tests/rt_basic.rs000064400000000000000000000265661046102023000143540ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(not(miri))] // Possible bug on Miri. use tokio::runtime::Runtime; use tokio::sync::oneshot; use tokio::time::{timeout, Duration}; use tokio_test::{assert_err, assert_ok}; use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::task::{Context, Poll}; use std::thread; mod support { pub(crate) mod mpsc_stream; } macro_rules! cfg_metrics { ($($t:tt)*) => { #[cfg(all(tokio_unstable, target_has_atomic = "64"))] { $( $t )* } } } #[test] fn spawned_task_does_not_progress_without_block_on() { let (tx, mut rx) = oneshot::channel(); let rt = rt(); rt.spawn(async move { assert_ok!(tx.send("hello")); }); thread::sleep(Duration::from_millis(50)); assert_err!(rx.try_recv()); let out = rt.block_on(async { assert_ok!(rx.await) }); assert_eq!(out, "hello"); } #[test] fn no_extra_poll() { use pin_project_lite::pin_project; use std::pin::Pin; use std::sync::{ atomic::{AtomicUsize, Ordering::SeqCst}, Arc, }; use std::task::{Context, Poll}; use tokio_stream::{Stream, StreamExt}; pin_project! { struct TrackPolls { npolls: Arc, #[pin] s: S, } } impl Stream for TrackPolls where S: Stream, { type Item = S::Item; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.project(); this.npolls.fetch_add(1, SeqCst); this.s.poll_next(cx) } } let (tx, rx) = support::mpsc_stream::unbounded_channel_stream::<()>(); let rx = TrackPolls { npolls: Arc::new(AtomicUsize::new(0)), s: rx, }; let npolls = Arc::clone(&rx.npolls); let rt = rt(); // TODO: could probably avoid this, but why not. let mut rx = Box::pin(rx); rt.spawn(async move { while rx.next().await.is_some() {} }); rt.block_on(async { tokio::task::yield_now().await; }); // should have been polled exactly once: the initial poll assert_eq!(npolls.load(SeqCst), 1); tx.send(()).unwrap(); rt.block_on(async { tokio::task::yield_now().await; }); // should have been polled twice more: once to yield Some(), then once to yield Pending assert_eq!(npolls.load(SeqCst), 1 + 2); drop(tx); rt.block_on(async { tokio::task::yield_now().await; }); // should have been polled once more: to yield None assert_eq!(npolls.load(SeqCst), 1 + 2 + 1); } #[test] fn acquire_mutex_in_drop() { use futures::future::pending; use tokio::task; let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); let rt = rt(); rt.spawn(async move { let _ = rx2.await; unreachable!(); }); rt.spawn(async move { let _ = rx1.await; tx2.send(()).unwrap(); unreachable!(); }); // Spawn a task that will never notify rt.spawn(async move { pending::<()>().await; tx1.send(()).unwrap(); }); // Tick the loop rt.block_on(async { task::yield_now().await; }); // Drop the rt drop(rt); } #[test] fn drop_tasks_in_context() { static SUCCESS: AtomicBool = AtomicBool::new(false); struct ContextOnDrop; impl Future for ContextOnDrop { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Pending } } impl Drop for ContextOnDrop { fn drop(&mut self) { if tokio::runtime::Handle::try_current().is_ok() { SUCCESS.store(true, Ordering::SeqCst); } } } let rt = rt(); rt.spawn(ContextOnDrop); drop(rt); assert!(SUCCESS.load(Ordering::SeqCst)); } #[test] #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] #[should_panic(expected = "boom")] fn wake_in_drop_after_panic() { struct WakeOnDrop(Option>); impl Drop for WakeOnDrop { fn drop(&mut self) { let _ = self.0.take().unwrap().send(()); } } let rt = rt(); let (tx1, rx1) = oneshot::channel::<()>(); let (tx2, rx2) = oneshot::channel::<()>(); // Spawn two tasks. We don't know the order in which they are dropped, so we // make both tasks identical. When the first task is dropped, we wake up the // second task. This ensures that we trigger a wakeup on a live task while // handling the "boom" panic, no matter the order in which the tasks are // dropped. rt.spawn(async move { let _wake_on_drop = WakeOnDrop(Some(tx2)); let _ = rx1.await; unreachable!() }); rt.spawn(async move { let _wake_on_drop = WakeOnDrop(Some(tx1)); let _ = rx2.await; unreachable!() }); rt.block_on(async { tokio::task::yield_now().await; panic!("boom"); }); } #[test] fn spawn_two() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); tokio::spawn(async move { tokio::spawn(async move { tx.send("ZOMG").unwrap(); }); }); assert_ok!(rx.await) }); assert_eq!(out, "ZOMG"); cfg_metrics! { let metrics = rt.metrics(); drop(rt); assert_eq!(0, metrics.remote_schedule_count()); let mut local = 0; for i in 0..metrics.num_workers() { local += metrics.worker_local_schedule_count(i); } assert_eq!(2, local); } } #[cfg_attr(target_os = "wasi", ignore = "WASI: std::thread::spawn not supported")] #[test] fn spawn_remote() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); let handle = tokio::spawn(async move { std::thread::spawn(move || { std::thread::sleep(Duration::from_millis(10)); tx.send("ZOMG").unwrap(); }); rx.await.unwrap() }); handle.await.unwrap() }); assert_eq!(out, "ZOMG"); cfg_metrics! { let metrics = rt.metrics(); drop(rt); assert_eq!(1, metrics.remote_schedule_count()); let mut local = 0; for i in 0..metrics.num_workers() { local += metrics.worker_local_schedule_count(i); } assert_eq!(1, local); } } #[test] #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] #[should_panic( expected = "A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers." )] fn timeout_panics_when_no_time_handle() { let rt = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); rt.block_on(async { let (_tx, rx) = oneshot::channel::<()>(); let dur = Duration::from_millis(20); let _ = timeout(dur, rx).await; }); } #[cfg(tokio_unstable)] mod unstable { use tokio::runtime::{Builder, RngSeed, UnhandledPanic}; #[test] #[should_panic( expected = "a spawned task panicked and the runtime is configured to shut down on unhandled panic" )] fn shutdown_on_panic() { let rt = Builder::new_current_thread() .unhandled_panic(UnhandledPanic::ShutdownRuntime) .build() .unwrap(); rt.block_on(async { tokio::spawn(async { panic!("boom"); }); futures::future::pending::<()>().await; }) } #[test] #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] fn spawns_do_nothing() { use std::sync::Arc; let rt = Builder::new_current_thread() .unhandled_panic(UnhandledPanic::ShutdownRuntime) .build() .unwrap(); let rt1 = Arc::new(rt); let rt2 = rt1.clone(); let _ = std::thread::spawn(move || { rt2.block_on(async { tokio::spawn(async { panic!("boom"); }); futures::future::pending::<()>().await; }) }) .join(); let task = rt1.spawn(async {}); let res = futures::executor::block_on(task); assert!(res.is_err()); } #[test] #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] fn shutdown_all_concurrent_block_on() { const N: usize = 2; use std::sync::{mpsc, Arc}; let rt = Builder::new_current_thread() .unhandled_panic(UnhandledPanic::ShutdownRuntime) .build() .unwrap(); let rt = Arc::new(rt); let mut ths = vec![]; let (tx, rx) = mpsc::channel(); for _ in 0..N { let rt = rt.clone(); let tx = tx.clone(); ths.push(std::thread::spawn(move || { rt.block_on(async { tx.send(()).unwrap(); futures::future::pending::<()>().await; }); })); } for _ in 0..N { rx.recv().unwrap(); } rt.spawn(async { panic!("boom"); }); for th in ths { assert!(th.join().is_err()); } } #[test] fn rng_seed() { let seed = b"bytes used to generate seed"; let rt1 = tokio::runtime::Builder::new_current_thread() .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt1_values = rt1.block_on(async { let rand_1 = tokio::macros::support::thread_rng_n(100); let rand_2 = tokio::macros::support::thread_rng_n(100); (rand_1, rand_2) }); let rt2 = tokio::runtime::Builder::new_current_thread() .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt2_values = rt2.block_on(async { let rand_1 = tokio::macros::support::thread_rng_n(100); let rand_2 = tokio::macros::support::thread_rng_n(100); (rand_1, rand_2) }); assert_eq!(rt1_values, rt2_values); } #[test] fn rng_seed_multi_enter() { let seed = b"bytes used to generate seed"; fn two_rand_values() -> (u32, u32) { let rand_1 = tokio::macros::support::thread_rng_n(100); let rand_2 = tokio::macros::support::thread_rng_n(100); (rand_1, rand_2) } let rt1 = tokio::runtime::Builder::new_current_thread() .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt1_values_1 = rt1.block_on(async { two_rand_values() }); let rt1_values_2 = rt1.block_on(async { two_rand_values() }); let rt2 = tokio::runtime::Builder::new_current_thread() .rng_seed(RngSeed::from_bytes(seed)) .build() .unwrap(); let rt2_values_1 = rt2.block_on(async { two_rand_values() }); let rt2_values_2 = rt2.block_on(async { two_rand_values() }); assert_eq!(rt1_values_1, rt2_values_1); assert_eq!(rt1_values_2, rt2_values_2); } } fn rt() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/rt_common.rs000064400000000000000000001173161046102023000145550ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![allow(clippy::needless_range_loop)] #![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(not(miri))] // Tests to run on both current-thread & multi-thread runtime variants. macro_rules! rt_test { ($($t:tt)*) => { mod current_thread_scheduler { $($t)* #[cfg(not(target_os="wasi"))] const NUM_WORKERS: usize = 1; fn rt() -> Arc { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() .into() } } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads mod threaded_scheduler_4_threads { $($t)* const NUM_WORKERS: usize = 4; fn rt() -> Arc { tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build() .unwrap() .into() } } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads mod threaded_scheduler_1_thread { $($t)* const NUM_WORKERS: usize = 1; fn rt() -> Arc { tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .enable_all() .build() .unwrap() .into() } } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[cfg(tokio_unstable)] mod alt_threaded_scheduler_4_threads { $($t)* const NUM_WORKERS: usize = 4; fn rt() -> Arc { tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build() .unwrap() .into() } } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[cfg(tokio_unstable)] mod alt_threaded_scheduler_1_thread { $($t)* const NUM_WORKERS: usize = 1; fn rt() -> Arc { tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .enable_all() .build() .unwrap() .into() } } } } #[test] fn send_sync_bound() { use tokio::runtime::Runtime; fn is_send() {} is_send::(); } rt_test! { #[cfg(not(target_os="wasi"))] use tokio::net::{TcpListener, TcpStream}; #[cfg(not(target_os="wasi"))] use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::runtime::Runtime; use tokio::sync::oneshot; use tokio::{task, time}; #[cfg(not(target_os="wasi"))] use tokio_test::assert_err; use tokio_test::assert_ok; use std::future::{poll_fn, Future}; use std::pin::Pin; #[cfg(not(target_os="wasi"))] use std::sync::mpsc; use std::sync::Arc; use std::task::{Context, Poll}; #[cfg(not(target_os="wasi"))] use std::thread; use std::time::{Duration, Instant}; #[test] fn block_on_sync() { let rt = rt(); let mut win = false; rt.block_on(async { win = true; }); assert!(win); } #[cfg(not(target_os="wasi"))] #[test] fn block_on_async() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); thread::spawn(move || { thread::sleep(Duration::from_millis(50)); tx.send("ZOMG").unwrap(); }); assert_ok!(rx.await) }); assert_eq!(out, "ZOMG"); } #[test] fn spawn_one_bg() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); tokio::spawn(async move { tx.send("ZOMG").unwrap(); }); assert_ok!(rx.await) }); assert_eq!(out, "ZOMG"); } #[test] fn spawn_one_join() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); let handle = tokio::spawn(async move { tx.send("ZOMG").unwrap(); "DONE" }); let msg = assert_ok!(rx.await); let out = assert_ok!(handle.await); assert_eq!(out, "DONE"); msg }); assert_eq!(out, "ZOMG"); } #[test] fn spawn_two() { let rt = rt(); let out = rt.block_on(async { let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); tokio::spawn(async move { assert_ok!(tx1.send("ZOMG")); }); tokio::spawn(async move { let msg = assert_ok!(rx1.await); assert_ok!(tx2.send(msg)); }); assert_ok!(rx2.await) }); assert_eq!(out, "ZOMG"); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn spawn_many_from_block_on() { use tokio::sync::mpsc; const ITER: usize = 200; let rt = rt(); let out = rt.block_on(async { let (done_tx, mut done_rx) = mpsc::unbounded_channel(); let mut txs = (0..ITER) .map(|i| { let (tx, rx) = oneshot::channel(); let done_tx = done_tx.clone(); tokio::spawn(async move { let msg = assert_ok!(rx.await); assert_eq!(i, msg); assert_ok!(done_tx.send(msg)); }); tx }) .collect::>(); drop(done_tx); thread::spawn(move || { for (i, tx) in txs.drain(..).enumerate() { assert_ok!(tx.send(i)); } }); let mut out = vec![]; while let Some(i) = done_rx.recv().await { out.push(i); } out.sort_unstable(); out }); assert_eq!(ITER, out.len()); for i in 0..ITER { assert_eq!(i, out[i]); } } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn spawn_many_from_task() { use tokio::sync::mpsc; const ITER: usize = 500; let rt = rt(); let out = rt.block_on(async { tokio::spawn(async move { let (done_tx, mut done_rx) = mpsc::unbounded_channel(); let mut txs = (0..ITER) .map(|i| { let (tx, rx) = oneshot::channel(); let done_tx = done_tx.clone(); tokio::spawn(async move { let msg = assert_ok!(rx.await); assert_eq!(i, msg); assert_ok!(done_tx.send(msg)); }); tx }) .collect::>(); drop(done_tx); thread::spawn(move || { for (i, tx) in txs.drain(..).enumerate() { assert_ok!(tx.send(i)); } }); let mut out = vec![]; while let Some(i) = done_rx.recv().await { out.push(i); } out.sort_unstable(); out }).await.unwrap() }); assert_eq!(ITER, out.len()); for i in 0..ITER { assert_eq!(i, out[i]); } } #[test] fn spawn_one_from_block_on_called_on_handle() { let rt = rt(); let (tx, rx) = oneshot::channel(); #[allow(clippy::async_yields_async)] let handle = rt.handle().block_on(async { tokio::spawn(async move { tx.send("ZOMG").unwrap(); "DONE" }) }); let out = rt.block_on(async { let msg = assert_ok!(rx.await); let out = assert_ok!(handle.await); assert_eq!(out, "DONE"); msg }); assert_eq!(out, "ZOMG"); } #[test] fn spawn_await_chain() { let rt = rt(); let out = rt.block_on(async { assert_ok!(tokio::spawn(async { assert_ok!(tokio::spawn(async { "hello" }).await) }).await) }); assert_eq!(out, "hello"); } #[test] fn outstanding_tasks_dropped() { let rt = rt(); let cnt = Arc::new(()); rt.block_on(async { let cnt = cnt.clone(); tokio::spawn(poll_fn(move |_| { assert_eq!(2, Arc::strong_count(&cnt)); Poll::<()>::Pending })); }); assert_eq!(2, Arc::strong_count(&cnt)); drop(rt); assert_eq!(1, Arc::strong_count(&cnt)); } #[test] #[should_panic] fn nested_rt() { let rt1 = rt(); let rt2 = rt(); rt1.block_on(async { rt2.block_on(async { "hello" }) }); } #[test] fn create_rt_in_block_on() { let rt1 = rt(); let rt2 = rt1.block_on(async { rt() }); let out = rt2.block_on(async { "ZOMG" }); assert_eq!(out, "ZOMG"); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn complete_block_on_under_load() { let rt = rt(); rt.block_on(async { let (tx, rx) = oneshot::channel(); // Spin hard tokio::spawn(async { loop { yield_once().await; } }); thread::spawn(move || { thread::sleep(Duration::from_millis(50)); assert_ok!(tx.send(())); }); assert_ok!(rx.await); }); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn complete_task_under_load() { let rt = rt(); rt.block_on(async { let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); // Spin hard tokio::spawn(async { loop { yield_once().await; } }); thread::spawn(move || { thread::sleep(Duration::from_millis(50)); assert_ok!(tx1.send(())); }); tokio::spawn(async move { assert_ok!(rx1.await); assert_ok!(tx2.send(())); }); assert_ok!(rx2.await); }); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn spawn_from_other_thread_idle() { let rt = rt(); let handle = rt.clone(); let (tx, rx) = oneshot::channel(); thread::spawn(move || { thread::sleep(Duration::from_millis(50)); handle.spawn(async move { assert_ok!(tx.send(())); }); }); rt.block_on(async move { assert_ok!(rx.await); }); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn spawn_from_other_thread_under_load() { let rt = rt(); let handle = rt.clone(); let (tx, rx) = oneshot::channel(); thread::spawn(move || { handle.spawn(async move { assert_ok!(tx.send(())); }); }); rt.block_on(async move { // Spin hard tokio::spawn(async { loop { yield_once().await; } }); assert_ok!(rx.await); }); } #[test] fn sleep_at_root() { let rt = rt(); let now = Instant::now(); let dur = Duration::from_millis(50); rt.block_on(async move { time::sleep(dur).await; }); assert!(now.elapsed() >= dur); } #[test] fn sleep_in_spawn() { let rt = rt(); let now = Instant::now(); let dur = Duration::from_millis(50); rt.block_on(async move { let (tx, rx) = oneshot::channel(); tokio::spawn(async move { time::sleep(dur).await; assert_ok!(tx.send(())); }); assert_ok!(rx.await); }); assert!(now.elapsed() >= dur); } #[cfg(not(target_os="wasi"))] // Wasi does not support bind #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn block_on_socket() { let rt = rt(); rt.block_on(async move { let (tx, rx) = oneshot::channel(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { let _ = listener.accept().await; tx.send(()).unwrap(); }); TcpStream::connect(&addr).await.unwrap(); rx.await.unwrap(); }); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn spawn_from_blocking() { let rt = rt(); let out = rt.block_on(async move { let inner = assert_ok!(tokio::task::spawn_blocking(|| { tokio::spawn(async move { "hello" }) }).await); assert_ok!(inner.await) }); assert_eq!(out, "hello") } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn spawn_blocking_from_blocking() { let rt = rt(); let out = rt.block_on(async move { let inner = assert_ok!(tokio::task::spawn_blocking(|| { tokio::task::spawn_blocking(|| "hello") }).await); assert_ok!(inner.await) }); assert_eq!(out, "hello") } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn sleep_from_blocking() { let rt = rt(); rt.block_on(async move { assert_ok!(tokio::task::spawn_blocking(|| { let now = std::time::Instant::now(); let dur = Duration::from_millis(1); // use the futures' block_on fn to make sure we aren't setting // any Tokio context futures::executor::block_on(async { tokio::time::sleep(dur).await; }); assert!(now.elapsed() >= dur); }).await); }); } #[cfg(not(target_os="wasi"))] // Wasi does not support bind #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn socket_from_blocking() { let rt = rt(); rt.block_on(async move { let listener = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(listener.local_addr()); let peer = tokio::task::spawn_blocking(move || { // use the futures' block_on fn to make sure we aren't setting // any Tokio context futures::executor::block_on(async { assert_ok!(TcpStream::connect(addr).await); }); }); // Wait for the client to connect let _ = assert_ok!(listener.accept().await); assert_ok!(peer.await); }); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn always_active_parker() { // This test it to show that we will always have // an active parker even if we call block_on concurrently let rt = rt(); let rt2 = rt.clone(); let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); let jh1 = thread::spawn(move || { rt.block_on(async move { rx2.await.unwrap(); time::sleep(Duration::from_millis(5)).await; tx1.send(()).unwrap(); }); }); let jh2 = thread::spawn(move || { rt2.block_on(async move { tx2.send(()).unwrap(); time::sleep(Duration::from_millis(5)).await; rx1.await.unwrap(); time::sleep(Duration::from_millis(5)).await; }); }); jh1.join().unwrap(); jh2.join().unwrap(); } #[test] // IOCP requires setting the "max thread" concurrency value. The sane, // default, is to set this to the number of cores. Threads that poll I/O // become associated with the IOCP handle. Once those threads sleep for any // reason (mutex), they yield their ownership. // // This test hits an edge case on windows where more threads than cores are // created, none of those threads ever yield due to being at capacity, so // IOCP gets "starved". // // For now, this is a very edge case that is probably not a real production // concern. There also isn't a great/obvious solution to take. For now, the // test is disabled. #[cfg(not(windows))] #[cfg_attr(miri, ignore)] // No `socket` in miri. #[cfg(not(target_os="wasi"))] // Wasi does not support bind or threads fn io_driver_called_when_under_load() { let rt = rt(); // Create a lot of constant load. The scheduler will always be busy. for _ in 0..100 { rt.spawn(async { loop { // Don't use Tokio's `yield_now()` to avoid special defer // logic. std::future::poll_fn::<(), _>(|cx| { cx.waker().wake_by_ref(); std::task::Poll::Pending }).await; } }); } // Do some I/O work rt.block_on(async { let listener = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(listener.local_addr()); let srv = tokio::spawn(async move { let (mut stream, _) = assert_ok!(listener.accept().await); assert_ok!(stream.write_all(b"hello world").await); }); let cli = tokio::spawn(async move { let mut stream = assert_ok!(TcpStream::connect(addr).await); let mut dst = vec![0; 11]; assert_ok!(stream.read_exact(&mut dst).await); assert_eq!(dst, b"hello world"); }); assert_ok!(srv.await); assert_ok!(cli.await); }); } /// Tests that yielded tasks are not scheduled until **after** resource /// drivers are polled. /// /// The OS does not guarantee when I/O events are delivered, so there may be /// more yields than anticipated. This makes the test slightly flaky. To /// help avoid flakiness, we run the test 10 times and only fail it after /// 10 failures in a row. /// /// Note that if the test fails by panicking rather than by returning false, /// then we fail it immediately. That kind of failure should not happen /// spuriously. #[test] #[cfg(not(target_os="wasi"))] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn yield_defers_until_park() { for _ in 0..10 { if yield_defers_until_park_inner() { // test passed return; } // Wait a bit and run the test again. std::thread::sleep(std::time::Duration::from_secs(2)); } panic!("yield_defers_until_park is failing consistently"); } /// Implementation of `yield_defers_until_park` test. Returns `true` if the /// test passed. #[cfg(not(target_os="wasi"))] fn yield_defers_until_park_inner() -> bool { use std::sync::atomic::{AtomicBool, Ordering::SeqCst}; use std::sync::Barrier; let rt = rt(); let flag = Arc::new(AtomicBool::new(false)); let barrier = Arc::new(Barrier::new(NUM_WORKERS)); rt.block_on(async { // Make sure other workers cannot steal tasks #[allow(clippy::reversed_empty_ranges)] for _ in 0..(NUM_WORKERS-1) { let flag = flag.clone(); let barrier = barrier.clone(); tokio::spawn(async move { barrier.wait(); while !flag.load(SeqCst) { std::thread::sleep(std::time::Duration::from_millis(1)); } }); } barrier.wait(); let (fail_test, fail_test_recv) = oneshot::channel::<()>(); let flag_clone = flag.clone(); let jh = tokio::spawn(async move { // Create a TCP listener let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::join!( async { // Done in a blocking manner intentionally. let _socket = std::net::TcpStream::connect(addr).unwrap(); // Yield until connected let mut cnt = 0; while !flag_clone.load(SeqCst){ tokio::task::yield_now().await; cnt += 1; if cnt >= 10 { // yielded too many times; report failure and // sleep forever so that the `fail_test` branch // of the `select!` below triggers. let _ = fail_test.send(()); futures::future::pending::<()>().await; break; } } }, async { let _ = listener.accept().await.unwrap(); flag_clone.store(true, SeqCst); } ); }); // Wait until the spawned task completes or fails. If no message is // sent on `fail_test`, then the test succeeds. Otherwise, it fails. let success = fail_test_recv.await.is_err(); if success { // Setting flag to true ensures that the tasks we spawned at // the beginning of the test will exit. // If we don't do this, the test will hang since the runtime waits // for all spawned tasks to finish when dropping. flag.store(true, SeqCst); // Check for panics in spawned task. jh.abort(); jh.await.unwrap(); } success }) } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn client_server_block_on() { let rt = rt(); let (tx, rx) = mpsc::channel(); rt.block_on(async move { client_server(tx).await }); assert_ok!(rx.try_recv()); assert_err!(rx.try_recv()); } #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support threads or panic recovery")] #[cfg(panic = "unwind")] #[test] fn panic_in_task() { let rt = rt(); let (tx, rx) = oneshot::channel(); struct Boom(Option>); impl Future for Boom { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { panic!(); } } impl Drop for Boom { fn drop(&mut self) { assert!(std::thread::panicking()); self.0.take().unwrap().send(()).unwrap(); } } rt.spawn(Boom(Some(tx))); assert_ok!(rt.block_on(rx)); } #[test] #[should_panic] #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] fn panic_in_block_on() { let rt = rt(); rt.block_on(async { panic!() }); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads async fn yield_once() { let mut yielded = false; poll_fn(|cx| { if yielded { Poll::Ready(()) } else { yielded = true; cx.waker().wake_by_ref(); Poll::Pending } }) .await } #[test] fn enter_and_spawn() { let rt = rt(); let handle = { let _enter = rt.enter(); tokio::spawn(async {}) }; assert_ok!(rt.block_on(handle)); } #[test] fn eagerly_drops_futures_on_shutdown() { use std::sync::mpsc; struct Never { drop_tx: mpsc::Sender<()>, } impl Future for Never { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Pending } } impl Drop for Never { fn drop(&mut self) { self.drop_tx.send(()).unwrap(); } } let rt = rt(); let (drop_tx, drop_rx) = mpsc::channel(); let (run_tx, run_rx) = oneshot::channel(); rt.block_on(async move { tokio::spawn(async move { assert_ok!(run_tx.send(())); Never { drop_tx }.await }); assert_ok!(run_rx.await); }); drop(rt); assert_ok!(drop_rx.recv()); } #[test] fn wake_while_rt_is_dropping() { use tokio::sync::Barrier; struct OnDrop(F); impl Drop for OnDrop { fn drop(&mut self) { (self.0)() } } let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); let barrier = Arc::new(Barrier::new(3)); let barrier1 = barrier.clone(); let barrier2 = barrier.clone(); let rt = rt(); rt.spawn(async move { let mut tx2 = Some(tx2); let _d = OnDrop(move || { let _ = tx2.take().unwrap().send(()); }); // Ensure a waker gets stored in oneshot 1. let _ = tokio::join!(rx1, barrier1.wait()); }); rt.spawn(async move { let mut tx1 = Some(tx1); let _d = OnDrop(move || { let _ = tx1.take().unwrap().send(()); }); // Ensure a waker gets stored in oneshot 2. let _ = tokio::join!(rx2, barrier2.wait()); }); // Wait until every oneshot channel has been polled. rt.block_on(barrier.wait()); // Drop the rt. Regardless of which task is dropped first, its destructor will wake the // other task. drop(rt); } #[cfg(not(target_os="wasi"))] // Wasi doesn't support UDP or bind() #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn io_notify_while_shutting_down() { use tokio::net::UdpSocket; use std::sync::Arc; for _ in 1..10 { let runtime = rt(); runtime.block_on(async { let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let addr = socket.local_addr().unwrap(); let send_half = Arc::new(socket); let recv_half = send_half.clone(); tokio::spawn(async move { let mut buf = [0]; loop { recv_half.recv_from(&mut buf).await.unwrap(); std::thread::sleep(Duration::from_millis(2)); } }); tokio::spawn(async move { let buf = [0]; loop { send_half.send_to(&buf, &addr).await.unwrap(); tokio::time::sleep(Duration::from_millis(1)).await; } }); tokio::time::sleep(Duration::from_millis(5)).await; }); } } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn shutdown_timeout() { let (tx, rx) = oneshot::channel(); let runtime = rt(); runtime.block_on(async move { task::spawn_blocking(move || { tx.send(()).unwrap(); thread::sleep(Duration::from_secs(10_000)); }); rx.await.unwrap(); }); Arc::try_unwrap(runtime).unwrap().shutdown_timeout(Duration::from_millis(100)); } #[cfg(not(target_os="wasi"))] // Wasi does not support threads #[test] fn shutdown_timeout_0() { let runtime = rt(); runtime.block_on(async move { task::spawn_blocking(move || { thread::sleep(Duration::from_secs(10_000)); }); }); let now = Instant::now(); Arc::try_unwrap(runtime).unwrap().shutdown_timeout(Duration::from_nanos(0)); assert!(now.elapsed().as_secs() < 1); } #[test] fn shutdown_wakeup_time() { let runtime = rt(); runtime.block_on(async move { tokio::time::sleep(std::time::Duration::from_millis(100)).await; }); Arc::try_unwrap(runtime).unwrap().shutdown_timeout(Duration::from_secs(10_000)); } // This test is currently ignored on Windows because of a // rust-lang issue in thread local storage destructors. // See https://github.com/rust-lang/rust/issues/74875 #[test] #[cfg(not(windows))] #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support threads")] fn runtime_in_thread_local() { use std::cell::RefCell; use std::thread; thread_local!( static R: RefCell> = const { RefCell::new(None) }; ); thread::spawn(|| { R.with(|cell| { let rt = rt(); let rt = Arc::try_unwrap(rt).unwrap(); *cell.borrow_mut() = Some(rt); }); let _rt = rt(); }).join().unwrap(); } #[cfg(not(target_os="wasi"))] // Wasi does not support bind async fn client_server(tx: mpsc::Sender<()>) { let server = assert_ok!(TcpListener::bind("127.0.0.1:0").await); // Get the assigned address let addr = assert_ok!(server.local_addr()); // Spawn the server tokio::spawn(async move { // Accept a socket let (mut socket, _) = server.accept().await.unwrap(); // Write some data socket.write_all(b"hello").await.unwrap(); }); let mut client = TcpStream::connect(&addr).await.unwrap(); let mut buf = vec![]; client.read_to_end(&mut buf).await.unwrap(); assert_eq!(buf, b"hello"); tx.send(()).unwrap(); } #[cfg(not(target_os = "wasi"))] // Wasi does not support bind #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn local_set_block_on_socket() { let rt = rt(); let local = task::LocalSet::new(); local.block_on(&rt, async move { let (tx, rx) = oneshot::channel(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); task::spawn_local(async move { let _ = listener.accept().await; tx.send(()).unwrap(); }); TcpStream::connect(&addr).await.unwrap(); rx.await.unwrap(); }); } #[cfg(not(target_os = "wasi"))] // Wasi does not support bind #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn local_set_client_server_block_on() { let rt = rt(); let (tx, rx) = mpsc::channel(); let local = task::LocalSet::new(); local.block_on(&rt, async move { client_server_local(tx).await }); assert_ok!(rx.try_recv()); assert_err!(rx.try_recv()); } #[cfg(not(target_os = "wasi"))] // Wasi does not support bind async fn client_server_local(tx: mpsc::Sender<()>) { let server = assert_ok!(TcpListener::bind("127.0.0.1:0").await); // Get the assigned address let addr = assert_ok!(server.local_addr()); // Spawn the server task::spawn_local(async move { // Accept a socket let (mut socket, _) = server.accept().await.unwrap(); // Write some data socket.write_all(b"hello").await.unwrap(); }); let mut client = TcpStream::connect(&addr).await.unwrap(); let mut buf = vec![]; client.read_to_end(&mut buf).await.unwrap(); assert_eq!(buf, b"hello"); tx.send(()).unwrap(); } #[test] fn coop() { use std::task::Poll::Ready; use tokio::sync::mpsc; let rt = rt(); rt.block_on(async { let (send, mut recv) = mpsc::unbounded_channel(); // Send a bunch of messages. for _ in 0..1_000 { send.send(()).unwrap(); } poll_fn(|cx| { // At least one response should return pending. for _ in 0..1_000 { if recv.poll_recv(cx).is_pending() { return Ready(()); } } panic!("did not yield"); }).await; }); } #[test] fn coop_unconstrained() { use std::task::Poll::Ready; use tokio::sync::mpsc; let rt = rt(); rt.block_on(async { let (send, mut recv) = mpsc::unbounded_channel(); // Send a bunch of messages. for _ in 0..1_000 { send.send(()).unwrap(); } tokio::task::unconstrained(poll_fn(|cx| { // All the responses should be ready. for _ in 0..1_000 { assert_eq!(recv.poll_recv(cx), Poll::Ready(Some(()))); } Ready(()) })).await; }); } #[cfg(tokio_unstable)] #[test] fn coop_consume_budget() { let rt = rt(); rt.block_on(async { poll_fn(|cx| { let counter = Arc::new(std::sync::Mutex::new(0)); let counter_clone = Arc::clone(&counter); let mut worker = Box::pin(async move { // Consume the budget until a yield happens for _ in 0..1000 { *counter.lock().unwrap() += 1; task::consume_budget().await } }); // Assert that the worker was yielded and it didn't manage // to finish the whole work (assuming the total budget of 128) assert!(Pin::new(&mut worker).poll(cx).is_pending()); assert!(*counter_clone.lock().unwrap() < 1000); std::task::Poll::Ready(()) }).await; }); } // Tests that the "next task" scheduler optimization is not able to starve // other tasks. #[test] fn ping_pong_saturation() { use std::sync::atomic::{Ordering, AtomicBool}; use tokio::sync::mpsc; const NUM: usize = 100; let rt = rt(); let running = Arc::new(AtomicBool::new(true)); rt.block_on(async { let (spawned_tx, mut spawned_rx) = mpsc::unbounded_channel(); let mut tasks = vec![]; // Spawn a bunch of tasks that ping-pong between each other to // saturate the runtime. for _ in 0..NUM { let (tx1, mut rx1) = mpsc::unbounded_channel(); let (tx2, mut rx2) = mpsc::unbounded_channel(); let spawned_tx = spawned_tx.clone(); let running = running.clone(); tasks.push(task::spawn(async move { spawned_tx.send(()).unwrap(); while running.load(Ordering::Relaxed) { tx1.send(()).unwrap(); rx2.recv().await.unwrap(); } // Close the channel and wait for the other task to exit. drop(tx1); assert!(rx2.recv().await.is_none()); })); tasks.push(task::spawn(async move { while rx1.recv().await.is_some() { tx2.send(()).unwrap(); } })); } for _ in 0..NUM { spawned_rx.recv().await.unwrap(); } // spawn another task and wait for it to complete let handle = task::spawn(async { for _ in 0..5 { // Yielding forces it back into the local queue. task::yield_now().await; } }); handle.await.unwrap(); running.store(false, Ordering::Relaxed); for t in tasks { t.await.unwrap(); } }); } #[test] #[cfg(not(target_os="wasi"))] fn shutdown_concurrent_spawn() { const NUM_TASKS: usize = 10_000; for _ in 0..5 { let (tx, rx) = std::sync::mpsc::channel(); let rt = rt(); let mut txs = vec![]; for _ in 0..NUM_TASKS { let (tx, rx) = tokio::sync::oneshot::channel(); txs.push(tx); rt.spawn(async move { rx.await.unwrap(); }); } // Prime the tasks rt.block_on(async { tokio::task::yield_now().await }); let th = std::thread::spawn(move || { tx.send(()).unwrap(); for tx in txs.drain(..) { let _ = tx.send(()); } }); rx.recv().unwrap(); drop(rt); th.join().unwrap(); } } #[test] #[cfg_attr(target_family = "wasm", ignore)] fn wake_by_ref_from_thread_local() { wake_from_thread_local(true); } #[test] #[cfg_attr(target_family = "wasm", ignore)] fn wake_by_val_from_thread_local() { wake_from_thread_local(false); } fn wake_from_thread_local(by_ref: bool) { use std::cell::RefCell; use std::sync::mpsc::{channel, Sender}; use std::task::Waker; struct TLData { by_ref: bool, waker: Option, done: Sender, } impl Drop for TLData { fn drop(&mut self) { if self.by_ref { self.waker.take().unwrap().wake_by_ref(); } else { self.waker.take().unwrap().wake(); } let _ = self.done.send(true); } } std::thread_local! { static TL_DATA: RefCell> = const { RefCell::new(None) }; }; let (send, recv) = channel(); std::thread::spawn(move || { let rt = rt(); rt.block_on(rt.spawn(poll_fn(move |cx| { let waker = cx.waker().clone(); let send = send.clone(); TL_DATA.with(|tl| { tl.replace(Some(TLData { by_ref, waker: Some(waker), done: send, })); }); Poll::Ready(()) }))) .unwrap(); }) .join() .unwrap(); assert!(recv.recv().unwrap()); } } tokio-1.43.1/tests/rt_handle.rs000064400000000000000000000054451046102023000145170ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::sync::Arc; use tokio::runtime::Runtime; use tokio::sync::{mpsc, Barrier}; #[test] #[cfg_attr(panic = "abort", ignore)] fn basic_enter() { let rt1 = rt(); let rt2 = rt(); let enter1 = rt1.enter(); let enter2 = rt2.enter(); drop(enter2); drop(enter1); } #[test] #[should_panic] #[cfg_attr(panic = "abort", ignore)] fn interleave_enter_different_rt() { let rt1 = rt(); let rt2 = rt(); let enter1 = rt1.enter(); let enter2 = rt2.enter(); drop(enter1); drop(enter2); } #[test] #[should_panic] #[cfg_attr(panic = "abort", ignore)] fn interleave_enter_same_rt() { let rt1 = rt(); let _enter1 = rt1.enter(); let enter2 = rt1.enter(); let enter3 = rt1.enter(); drop(enter2); drop(enter3); } #[test] #[cfg(not(target_os = "wasi"))] #[cfg_attr(panic = "abort", ignore)] fn interleave_then_enter() { let _ = std::panic::catch_unwind(|| { let rt1 = rt(); let rt2 = rt(); let enter1 = rt1.enter(); let enter2 = rt2.enter(); drop(enter1); drop(enter2); }); // Can still enter let rt3 = rt(); let _enter = rt3.enter(); } // If the cycle causes a leak, then miri will catch it. #[test] fn drop_tasks_with_reference_cycle() { rt().block_on(async { let (tx, mut rx) = mpsc::channel(1); let barrier = Arc::new(Barrier::new(3)); let barrier_a = barrier.clone(); let barrier_b = barrier.clone(); let a = tokio::spawn(async move { let b = rx.recv().await.unwrap(); // Poll the JoinHandle once. This registers the waker. // The other task cannot have finished at this point due to the barrier below. futures::future::select(b, std::future::ready(())).await; barrier_a.wait().await; }); let b = tokio::spawn(async move { // Poll the JoinHandle once. This registers the waker. // The other task cannot have finished at this point due to the barrier below. futures::future::select(a, std::future::ready(())).await; barrier_b.wait().await; }); tx.send(b).await.unwrap(); barrier.wait().await; }); } #[cfg(tokio_unstable)] mod unstable { use super::*; #[test] fn runtime_id_is_same() { let rt = rt(); let handle1 = rt.handle(); let handle2 = rt.handle(); assert_eq!(handle1.id(), handle2.id()); } #[test] fn runtime_ids_different() { let rt1 = rt(); let rt2 = rt(); assert_ne!(rt1.handle().id(), rt2.handle().id()); } } fn rt() -> Runtime { tokio::runtime::Builder::new_current_thread() .build() .unwrap() } tokio-1.43.1/tests/rt_handle_block_on.rs000064400000000000000000000344161046102023000163650ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(miri)))] // All io tests that deal with shutdown is currently ignored because there are known bugs in with // shutting down the io driver while concurrently registering new resources. See // https://github.com/tokio-rs/tokio/pull/3569#pullrequestreview-612703467 for more details. // // When this has been fixed we want to re-enable these tests. use std::time::Duration; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; #[cfg(not(target_os = "wasi"))] use tokio::{net, time}; #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads macro_rules! multi_threaded_rt_test { ($($t:tt)*) => { mod threaded_scheduler_4_threads_only { use super::*; $($t)* fn rt() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build() .unwrap() } } mod threaded_scheduler_1_thread_only { use super::*; $($t)* fn rt() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .enable_all() .build() .unwrap() } } } } #[cfg(not(target_os = "wasi"))] macro_rules! rt_test { ($($t:tt)*) => { mod current_thread_scheduler { use super::*; $($t)* fn rt() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } } mod threaded_scheduler_4_threads { use super::*; $($t)* fn rt() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build() .unwrap() } } mod threaded_scheduler_1_thread { use super::*; $($t)* fn rt() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .enable_all() .build() .unwrap() } } } } // ==== runtime independent futures ====== #[test] fn basic() { test_with_runtimes(|| { let one = Handle::current().block_on(async { 1 }); assert_eq!(1, one); }); } #[test] fn bounded_mpsc_channel() { test_with_runtimes(|| { let (tx, mut rx) = mpsc::channel(1024); Handle::current().block_on(tx.send(42)).unwrap(); let value = Handle::current().block_on(rx.recv()).unwrap(); assert_eq!(value, 42); }); } #[test] fn unbounded_mpsc_channel() { test_with_runtimes(|| { let (tx, mut rx) = mpsc::unbounded_channel(); let _ = tx.send(42); let value = Handle::current().block_on(rx.recv()).unwrap(); assert_eq!(value, 42); }) } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support file operations or bind rt_test! { use tokio::fs; // ==== spawn blocking futures ====== #[test] fn basic_fs() { let rt = rt(); let _enter = rt.enter(); let contents = Handle::current() .block_on(fs::read_to_string("Cargo.toml")) .unwrap(); assert!(contents.contains("https://tokio.rs")); } #[test] fn fs_shutdown_before_started() { let rt = rt(); let _enter = rt.enter(); rt.shutdown_timeout(Duration::from_secs(1000)); let err: std::io::Error = Handle::current() .block_on(fs::read_to_string("Cargo.toml")) .unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); let inner_err = err.get_ref().expect("no inner error"); assert_eq!(inner_err.to_string(), "background task failed"); } #[test] fn basic_spawn_blocking() { use tokio::task::spawn_blocking; let rt = rt(); let _enter = rt.enter(); let answer = Handle::current() .block_on(spawn_blocking(|| { std::thread::sleep(Duration::from_millis(100)); 42 })) .unwrap(); assert_eq!(answer, 42); } #[test] fn spawn_blocking_after_shutdown_fails() { use tokio::task::spawn_blocking; let rt = rt(); let _enter = rt.enter(); rt.shutdown_timeout(Duration::from_secs(1000)); let join_err = Handle::current() .block_on(spawn_blocking(|| { std::thread::sleep(Duration::from_millis(100)); 42 })) .unwrap_err(); assert!(join_err.is_cancelled()); } #[test] fn spawn_blocking_started_before_shutdown_continues() { use tokio::task::spawn_blocking; let rt = rt(); let _enter = rt.enter(); let handle = spawn_blocking(|| { std::thread::sleep(Duration::from_secs(1)); 42 }); rt.shutdown_timeout(Duration::from_secs(1000)); let answer = Handle::current().block_on(handle).unwrap(); assert_eq!(answer, 42); } // ==== net ====== #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn tcp_listener_bind() { let rt = rt(); let _enter = rt.enter(); Handle::current() .block_on(net::TcpListener::bind("127.0.0.1:0")) .unwrap(); } // All io tests are ignored for now. See above why that is. #[ignore] #[test] fn tcp_listener_connect_after_shutdown() { let rt = rt(); let _enter = rt.enter(); rt.shutdown_timeout(Duration::from_secs(1000)); let err = Handle::current() .block_on(net::TcpListener::bind("127.0.0.1:0")) .unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!( err.get_ref().unwrap().to_string(), "A Tokio 1.x context was found, but it is being shutdown.", ); } // All io tests are ignored for now. See above why that is. #[ignore] #[test] fn tcp_listener_connect_before_shutdown() { let rt = rt(); let _enter = rt.enter(); let bind_future = net::TcpListener::bind("127.0.0.1:0"); rt.shutdown_timeout(Duration::from_secs(1000)); let err = Handle::current().block_on(bind_future).unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!( err.get_ref().unwrap().to_string(), "A Tokio 1.x context was found, but it is being shutdown.", ); } #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn udp_socket_bind() { let rt = rt(); let _enter = rt.enter(); Handle::current() .block_on(net::UdpSocket::bind("127.0.0.1:0")) .unwrap(); } // All io tests are ignored for now. See above why that is. #[ignore] #[test] fn udp_stream_bind_after_shutdown() { let rt = rt(); let _enter = rt.enter(); rt.shutdown_timeout(Duration::from_secs(1000)); let err = Handle::current() .block_on(net::UdpSocket::bind("127.0.0.1:0")) .unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!( err.get_ref().unwrap().to_string(), "A Tokio 1.x context was found, but it is being shutdown.", ); } // All io tests are ignored for now. See above why that is. #[ignore] #[test] fn udp_stream_bind_before_shutdown() { let rt = rt(); let _enter = rt.enter(); let bind_future = net::UdpSocket::bind("127.0.0.1:0"); rt.shutdown_timeout(Duration::from_secs(1000)); let err = Handle::current().block_on(bind_future).unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!( err.get_ref().unwrap().to_string(), "A Tokio 1.x context was found, but it is being shutdown.", ); } // All io tests are ignored for now. See above why that is. #[ignore] #[cfg(unix)] #[test] fn unix_listener_bind_after_shutdown() { let rt = rt(); let _enter = rt.enter(); let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("socket"); rt.shutdown_timeout(Duration::from_secs(1000)); let err = net::UnixListener::bind(path).unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!( err.get_ref().unwrap().to_string(), "A Tokio 1.x context was found, but it is being shutdown.", ); } // All io tests are ignored for now. See above why that is. #[ignore] #[cfg(unix)] #[test] fn unix_listener_shutdown_after_bind() { let rt = rt(); let _enter = rt.enter(); let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("socket"); let listener = net::UnixListener::bind(path).unwrap(); rt.shutdown_timeout(Duration::from_secs(1000)); // this should not timeout but fail immediately since the runtime has been shutdown let err = Handle::current().block_on(listener.accept()).unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!(err.get_ref().unwrap().to_string(), "reactor gone"); } // All io tests are ignored for now. See above why that is. #[ignore] #[cfg(unix)] #[test] fn unix_listener_shutdown_after_accept() { let rt = rt(); let _enter = rt.enter(); let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("socket"); let listener = net::UnixListener::bind(path).unwrap(); let accept_future = listener.accept(); rt.shutdown_timeout(Duration::from_secs(1000)); // this should not timeout but fail immediately since the runtime has been shutdown let err = Handle::current().block_on(accept_future).unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!(err.get_ref().unwrap().to_string(), "reactor gone"); } // ==== nesting ====== #[test] #[should_panic( expected = "Cannot start a runtime from within a runtime. This happens because a function (like `block_on`) attempted to block the current thread while the thread is being used to drive asynchronous tasks." )] fn nesting() { fn some_non_async_function() -> i32 { Handle::current().block_on(time::sleep(Duration::from_millis(10))); 1 } let rt = rt(); rt.block_on(async { some_non_async_function() }); } #[test] fn spawn_after_runtime_dropped() { use futures::future::FutureExt; let rt = rt(); let handle = rt.block_on(async move { Handle::current() }); let jh1 = handle.spawn(futures::future::pending::<()>()); drop(rt); let jh2 = handle.spawn(futures::future::pending::<()>()); let err1 = jh1.now_or_never().unwrap().unwrap_err(); let err2 = jh2.now_or_never().unwrap().unwrap_err(); assert!(err1.is_cancelled()); assert!(err2.is_cancelled()); } } #[cfg(not(target_os = "wasi"))] multi_threaded_rt_test! { #[cfg(unix)] #[cfg_attr(miri, ignore)] // No `socket` in miri. #[test] fn unix_listener_bind() { let rt = rt(); let _enter = rt.enter(); let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("socket"); let listener = net::UnixListener::bind(path).unwrap(); // this should timeout and not fail immediately since the runtime has not been shutdown let _: tokio::time::error::Elapsed = Handle::current() .block_on(tokio::time::timeout( Duration::from_millis(10), listener.accept(), )) .unwrap_err(); } // ==== timers ====== // `Handle::block_on` doesn't work with timer futures on a current thread runtime as there is no // one to drive the timers so they will just hang forever. Therefore they are not tested. #[test] fn sleep() { let rt = rt(); let _enter = rt.enter(); Handle::current().block_on(time::sleep(Duration::from_millis(100))); } #[test] #[should_panic(expected = "A Tokio 1.x context was found, but it is being shutdown.")] fn sleep_before_shutdown_panics() { let rt = rt(); let _enter = rt.enter(); let f = time::sleep(Duration::from_millis(100)); rt.shutdown_timeout(Duration::from_secs(1000)); Handle::current().block_on(f); } #[test] #[should_panic(expected = "A Tokio 1.x context was found, but it is being shutdown.")] fn sleep_after_shutdown_panics() { let rt = rt(); let _enter = rt.enter(); rt.shutdown_timeout(Duration::from_secs(1000)); Handle::current().block_on(time::sleep(Duration::from_millis(100))); } } // ==== utils ====== /// Create a new multi threaded runtime #[cfg(not(target_os = "wasi"))] fn new_multi_thread(n: usize) -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(n) .enable_all() .build() .unwrap() } /// Create a new single threaded runtime fn new_current_thread() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } /// Utility to test things on both kinds of runtimes both before and after shutting it down. fn test_with_runtimes(f: F) where F: Fn(), { { let rt = new_current_thread(); let _enter = rt.enter(); f(); rt.shutdown_timeout(Duration::from_secs(1000)); f(); } #[cfg(not(target_os = "wasi"))] { let rt = new_multi_thread(1); let _enter = rt.enter(); f(); rt.shutdown_timeout(Duration::from_secs(1000)); f(); } #[cfg(not(target_os = "wasi"))] { let rt = new_multi_thread(4); let _enter = rt.enter(); f(); rt.shutdown_timeout(Duration::from_secs(1000)); f(); } } tokio-1.43.1/tests/rt_local.rs000064400000000000000000000037541046102023000143570ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(feature = "full", tokio_unstable))] use tokio::runtime::LocalOptions; use tokio::task::spawn_local; #[test] fn test_spawn_local_in_runtime() { let rt = rt(); let res = rt.block_on(async move { let (tx, rx) = tokio::sync::oneshot::channel(); spawn_local(async { tokio::task::yield_now().await; tx.send(5).unwrap(); }); rx.await.unwrap() }); assert_eq!(res, 5); } #[test] fn test_spawn_from_handle() { let rt = rt(); let (tx, rx) = tokio::sync::oneshot::channel(); rt.handle().spawn(async { tokio::task::yield_now().await; tx.send(5).unwrap(); }); let res = rt.block_on(async move { rx.await.unwrap() }); assert_eq!(res, 5); } #[test] fn test_spawn_local_on_runtime_object() { let rt = rt(); let (tx, rx) = tokio::sync::oneshot::channel(); rt.spawn_local(async { tokio::task::yield_now().await; tx.send(5).unwrap(); }); let res = rt.block_on(async move { rx.await.unwrap() }); assert_eq!(res, 5); } #[test] fn test_spawn_local_from_guard() { let rt = rt(); let (tx, rx) = tokio::sync::oneshot::channel(); let _guard = rt.enter(); spawn_local(async { tokio::task::yield_now().await; tx.send(5).unwrap(); }); let res = rt.block_on(async move { rx.await.unwrap() }); assert_eq!(res, 5); } #[test] #[should_panic] fn test_spawn_local_from_guard_other_thread() { let (tx, rx) = std::sync::mpsc::channel(); std::thread::spawn(move || { let rt = rt(); let handle = rt.handle().clone(); tx.send(handle).unwrap(); }); let handle = rx.recv().unwrap(); let _guard = handle.enter(); spawn_local(async {}); } fn rt() -> tokio::runtime::LocalRuntime { tokio::runtime::Builder::new_current_thread() .enable_all() .build_local(&LocalOptions::default()) .unwrap() } tokio-1.43.1/tests/rt_metrics.rs000064400000000000000000000067601046102023000147330ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), target_has_atomic = "64"))] use std::sync::mpsc; use std::time::Duration; use tokio::runtime::Runtime; #[test] fn num_workers() { let rt = current_thread(); assert_eq!(1, rt.metrics().num_workers()); let rt = threaded(); assert_eq!(2, rt.metrics().num_workers()); } #[test] fn num_alive_tasks() { let rt = current_thread(); let metrics = rt.metrics(); assert_eq!(0, metrics.num_alive_tasks()); rt.block_on(rt.spawn(async move { assert_eq!(1, metrics.num_alive_tasks()); })) .unwrap(); assert_eq!(0, rt.metrics().num_alive_tasks()); let rt = threaded(); let metrics = rt.metrics(); assert_eq!(0, metrics.num_alive_tasks()); rt.block_on(rt.spawn(async move { assert_eq!(1, metrics.num_alive_tasks()); })) .unwrap(); // try for 10 seconds to see if this eventually succeeds. // wake_join() is called before the task is released, so in multithreaded // code, this means we sometimes exit the block_on before the counter decrements. for _ in 0..100 { if rt.metrics().num_alive_tasks() == 0 { break; } std::thread::sleep(std::time::Duration::from_millis(100)); } assert_eq!(0, rt.metrics().num_alive_tasks()); } #[test] fn global_queue_depth_current_thread() { use std::thread; let rt = current_thread(); let handle = rt.handle().clone(); let metrics = rt.metrics(); thread::spawn(move || { handle.spawn(async {}); }) .join() .unwrap(); assert_eq!(1, metrics.global_queue_depth()); } #[test] fn global_queue_depth_multi_thread() { for _ in 0..10 { let rt = threaded(); let metrics = rt.metrics(); if let Ok(_blocking_tasks) = try_block_threaded(&rt) { for i in 0..10 { assert_eq!(i, metrics.global_queue_depth()); rt.spawn(async {}); } return; } } panic!("exhausted every try to block the runtime"); } fn try_block_threaded(rt: &Runtime) -> Result>, mpsc::RecvTimeoutError> { let (tx, rx) = mpsc::channel(); let blocking_tasks = (0..rt.metrics().num_workers()) .map(|_| { let tx = tx.clone(); let (task, barrier) = mpsc::channel(); // Spawn a task per runtime worker to block it. rt.spawn(async move { tx.send(()).ok(); barrier.recv().ok(); }); task }) .collect(); // Make sure the previously spawned tasks are blocking the runtime by // receiving a message from each blocking task. // // If this times out we were unsuccessful in blocking the runtime and hit // a deadlock instead (which might happen and is expected behaviour). for _ in 0..rt.metrics().num_workers() { rx.recv_timeout(Duration::from_secs(1))?; } // Return senders of the mpsc channels used for blocking the runtime as a // surrogate handle for the tasks. Sending a message or dropping the senders // will unblock the runtime. Ok(blocking_tasks) } fn current_thread() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } fn threaded() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .build() .unwrap() } tokio-1.43.1/tests/rt_panic.rs000064400000000000000000000043621046102023000143530ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery #![cfg(panic = "unwind")] use futures::future; use std::error::Error; use tokio::runtime::{Builder, Handle, Runtime}; mod support { pub mod panic; } use support::panic::test_panic; #[test] fn current_handle_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = Handle::current(); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn into_panic_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(move || { let rt = current_thread(); rt.block_on(async { let handle = tokio::spawn(future::pending::<()>()); handle.abort(); let err = handle.await.unwrap_err(); assert!(!&err.is_panic()); let _ = err.into_panic(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn builder_worker_threads_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = Builder::new_multi_thread().worker_threads(0).build(); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn builder_max_blocking_threads_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = Builder::new_multi_thread().max_blocking_threads(0).build(); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn builder_global_queue_interval_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = Builder::new_multi_thread().global_queue_interval(0).build(); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } fn current_thread() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/rt_threaded.rs000064400000000000000000000522731046102023000150450ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] // Too slow on miri. #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::runtime; use tokio::sync::oneshot; use tokio_test::{assert_err, assert_ok}; use std::future::{poll_fn, Future}; use std::pin::Pin; use std::sync::atomic::Ordering::Relaxed; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{mpsc, Arc, Mutex}; use std::task::{Context, Poll, Waker}; macro_rules! cfg_metrics { ($($t:tt)*) => { #[cfg(all(tokio_unstable, target_has_atomic = "64"))] { $( $t )* } } } #[test] fn single_thread() { // No panic when starting a runtime w/ a single thread let _ = runtime::Builder::new_multi_thread() .enable_all() .worker_threads(1) .build() .unwrap(); } #[test] fn many_oneshot_futures() { // used for notifying the main thread const NUM: usize = 1_000; for _ in 0..5 { let (tx, rx) = mpsc::channel(); let rt = rt(); let cnt = Arc::new(AtomicUsize::new(0)); for _ in 0..NUM { let cnt = cnt.clone(); let tx = tx.clone(); rt.spawn(async move { let num = cnt.fetch_add(1, Relaxed) + 1; if num == NUM { tx.send(()).unwrap(); } }); } rx.recv().unwrap(); // Wait for the pool to shutdown drop(rt); } } #[test] fn spawn_two() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); tokio::spawn(async move { tokio::spawn(async move { tx.send("ZOMG").unwrap(); }); }); assert_ok!(rx.await) }); assert_eq!(out, "ZOMG"); cfg_metrics! { let metrics = rt.metrics(); drop(rt); assert_eq!(1, metrics.remote_schedule_count()); let mut local = 0; for i in 0..metrics.num_workers() { local += metrics.worker_local_schedule_count(i); } assert_eq!(1, local); } } #[test] fn many_multishot_futures() { const CHAIN: usize = 200; const CYCLES: usize = 5; const TRACKS: usize = 50; for _ in 0..50 { let rt = rt(); let mut start_txs = Vec::with_capacity(TRACKS); let mut final_rxs = Vec::with_capacity(TRACKS); for _ in 0..TRACKS { let (start_tx, mut chain_rx) = tokio::sync::mpsc::channel(10); for _ in 0..CHAIN { let (next_tx, next_rx) = tokio::sync::mpsc::channel(10); // Forward all the messages rt.spawn(async move { while let Some(v) = chain_rx.recv().await { next_tx.send(v).await.unwrap(); } }); chain_rx = next_rx; } // This final task cycles if needed let (final_tx, final_rx) = tokio::sync::mpsc::channel(10); let cycle_tx = start_tx.clone(); let mut rem = CYCLES; rt.spawn(async move { for _ in 0..CYCLES { let msg = chain_rx.recv().await.unwrap(); rem -= 1; if rem == 0 { final_tx.send(msg).await.unwrap(); } else { cycle_tx.send(msg).await.unwrap(); } } }); start_txs.push(start_tx); final_rxs.push(final_rx); } { rt.block_on(async move { for start_tx in start_txs { start_tx.send("ping").await.unwrap(); } for mut final_rx in final_rxs { final_rx.recv().await.unwrap(); } }); } } } #[test] fn lifo_slot_budget() { async fn my_fn() { spawn_another(); } fn spawn_another() { tokio::spawn(my_fn()); } let rt = runtime::Builder::new_multi_thread() .enable_all() .worker_threads(1) .build() .unwrap(); let (send, recv) = oneshot::channel(); rt.spawn(async move { tokio::spawn(my_fn()); let _ = send.send(()); }); let _ = rt.block_on(recv); } #[test] #[cfg_attr(miri, ignore)] // No `socket` in miri. fn spawn_shutdown() { let rt = rt(); let (tx, rx) = mpsc::channel(); rt.block_on(async { tokio::spawn(client_server(tx.clone())); }); // Use spawner rt.spawn(client_server(tx)); assert_ok!(rx.recv()); assert_ok!(rx.recv()); drop(rt); assert_err!(rx.try_recv()); } async fn client_server(tx: mpsc::Sender<()>) { let server = assert_ok!(TcpListener::bind("127.0.0.1:0").await); // Get the assigned address let addr = assert_ok!(server.local_addr()); // Spawn the server tokio::spawn(async move { // Accept a socket let (mut socket, _) = server.accept().await.unwrap(); // Write some data socket.write_all(b"hello").await.unwrap(); }); let mut client = TcpStream::connect(&addr).await.unwrap(); let mut buf = vec![]; client.read_to_end(&mut buf).await.unwrap(); assert_eq!(buf, b"hello"); tx.send(()).unwrap(); } #[test] fn drop_threadpool_drops_futures() { for _ in 0..1_000 { let num_inc = Arc::new(AtomicUsize::new(0)); let num_dec = Arc::new(AtomicUsize::new(0)); let num_drop = Arc::new(AtomicUsize::new(0)); struct Never(Arc); impl Future for Never { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Pending } } impl Drop for Never { fn drop(&mut self) { self.0.fetch_add(1, Relaxed); } } let a = num_inc.clone(); let b = num_dec.clone(); let rt = runtime::Builder::new_multi_thread() .enable_all() .on_thread_start(move || { a.fetch_add(1, Relaxed); }) .on_thread_stop(move || { b.fetch_add(1, Relaxed); }) .build() .unwrap(); rt.spawn(Never(num_drop.clone())); // Wait for the pool to shutdown drop(rt); // Assert that only a single thread was spawned. let a = num_inc.load(Relaxed); assert!(a >= 1); // Assert that all threads shutdown let b = num_dec.load(Relaxed); assert_eq!(a, b); // Assert that the future was dropped let c = num_drop.load(Relaxed); assert_eq!(c, 1); } } #[test] fn start_stop_callbacks_called() { use std::sync::atomic::{AtomicUsize, Ordering}; let after_start = Arc::new(AtomicUsize::new(0)); let before_stop = Arc::new(AtomicUsize::new(0)); let after_inner = after_start.clone(); let before_inner = before_stop.clone(); let rt = tokio::runtime::Builder::new_multi_thread() .enable_all() .on_thread_start(move || { after_inner.clone().fetch_add(1, Ordering::Relaxed); }) .on_thread_stop(move || { before_inner.clone().fetch_add(1, Ordering::Relaxed); }) .build() .unwrap(); let (tx, rx) = oneshot::channel(); rt.spawn(async move { assert_ok!(tx.send(())); }); assert_ok!(rt.block_on(rx)); drop(rt); assert!(after_start.load(Ordering::Relaxed) > 0); assert!(before_stop.load(Ordering::Relaxed) > 0); } #[test] // too slow on miri #[cfg_attr(miri, ignore)] fn blocking() { // used for notifying the main thread const NUM: usize = 1_000; for _ in 0..10 { let (tx, rx) = mpsc::channel(); let rt = rt(); let cnt = Arc::new(AtomicUsize::new(0)); // there are four workers in the pool // so, if we run 4 blocking tasks, we know that handoff must have happened let block = Arc::new(std::sync::Barrier::new(5)); for _ in 0..4 { let block = block.clone(); rt.spawn(async move { tokio::task::block_in_place(move || { block.wait(); block.wait(); }) }); } block.wait(); for _ in 0..NUM { let cnt = cnt.clone(); let tx = tx.clone(); rt.spawn(async move { let num = cnt.fetch_add(1, Relaxed) + 1; if num == NUM { tx.send(()).unwrap(); } }); } rx.recv().unwrap(); // Wait for the pool to shutdown block.wait(); } } #[test] fn multi_threadpool() { use tokio::sync::oneshot; let rt1 = rt(); let rt2 = rt(); let (tx, rx) = oneshot::channel(); let (done_tx, done_rx) = mpsc::channel(); rt2.spawn(async move { rx.await.unwrap(); done_tx.send(()).unwrap(); }); rt1.spawn(async move { tx.send(()).unwrap(); }); done_rx.recv().unwrap(); } // When `block_in_place` returns, it attempts to reclaim the yielded runtime // worker. In this case, the remainder of the task is on the runtime worker and // must take part in the cooperative task budgeting system. // // The test ensures that, when this happens, attempting to consume from a // channel yields occasionally even if there are values ready to receive. #[test] fn coop_and_block_in_place() { let rt = tokio::runtime::Builder::new_multi_thread() // Setting max threads to 1 prevents another thread from claiming the // runtime worker yielded as part of `block_in_place` and guarantees the // same thread will reclaim the worker at the end of the // `block_in_place` call. .max_blocking_threads(1) .build() .unwrap(); rt.block_on(async move { let (tx, mut rx) = tokio::sync::mpsc::channel(1024); // Fill the channel for _ in 0..1024 { tx.send(()).await.unwrap(); } drop(tx); tokio::spawn(async move { // Block in place without doing anything tokio::task::block_in_place(|| {}); // Receive all the values, this should trigger a `Pending` as the // coop limit will be reached. poll_fn(|cx| { while let Poll::Ready(v) = { tokio::pin! { let fut = rx.recv(); } Pin::new(&mut fut).poll(cx) } { if v.is_none() { panic!("did not yield"); } } Poll::Ready(()) }) .await }) .await .unwrap(); }); } #[test] fn yield_after_block_in_place() { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .build() .unwrap(); rt.block_on(async { tokio::spawn(async move { // Block in place then enter a new runtime tokio::task::block_in_place(|| { let rt = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); rt.block_on(async {}); }); // Yield, then complete tokio::task::yield_now().await; }) .await .unwrap() }); } // Testing this does not panic #[test] fn max_blocking_threads() { let _rt = tokio::runtime::Builder::new_multi_thread() .max_blocking_threads(1) .build() .unwrap(); } #[test] #[should_panic] fn max_blocking_threads_set_to_zero() { let _rt = tokio::runtime::Builder::new_multi_thread() .max_blocking_threads(0) .build() .unwrap(); } /// Regression test for #6445. /// /// After #6445, setting `global_queue_interval` to 1 is now technically valid. /// This test confirms that there is no regression in `multi_thread_runtime` /// when global_queue_interval is set to 1. #[test] fn global_queue_interval_set_to_one() { let rt = tokio::runtime::Builder::new_multi_thread() .global_queue_interval(1) .build() .unwrap(); // Perform a simple work. let cnt = Arc::new(AtomicUsize::new(0)); rt.block_on(async { let mut set = tokio::task::JoinSet::new(); for _ in 0..10 { let cnt = cnt.clone(); set.spawn(async move { cnt.fetch_add(1, Ordering::Relaxed) }); } while let Some(res) = set.join_next().await { res.unwrap(); } }); assert_eq!(cnt.load(Relaxed), 10); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn hang_on_shutdown() { let (sync_tx, sync_rx) = std::sync::mpsc::channel::<()>(); tokio::spawn(async move { tokio::task::block_in_place(|| sync_rx.recv().ok()); }); tokio::spawn(async { tokio::time::sleep(std::time::Duration::from_secs(2)).await; drop(sync_tx); }); tokio::time::sleep(std::time::Duration::from_secs(1)).await; } /// Demonstrates tokio-rs/tokio#3869 #[test] fn wake_during_shutdown() { struct Shared { waker: Option, } struct MyFuture { shared: Arc>, put_waker: bool, } impl MyFuture { fn new() -> (Self, Self) { let shared = Arc::new(Mutex::new(Shared { waker: None })); let f1 = MyFuture { shared: shared.clone(), put_waker: true, }; let f2 = MyFuture { shared, put_waker: false, }; (f1, f2) } } impl Future for MyFuture { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { let me = Pin::into_inner(self); let mut lock = me.shared.lock().unwrap(); if me.put_waker { lock.waker = Some(cx.waker().clone()); } Poll::Pending } } impl Drop for MyFuture { fn drop(&mut self) { let mut lock = self.shared.lock().unwrap(); if !self.put_waker { lock.waker.take().unwrap().wake(); } drop(lock); } } let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .enable_all() .build() .unwrap(); let (f1, f2) = MyFuture::new(); rt.spawn(f1); rt.spawn(f2); rt.block_on(async { tokio::time::sleep(tokio::time::Duration::from_millis(20)).await }); } #[should_panic] #[tokio::test] async fn test_block_in_place1() { tokio::task::block_in_place(|| {}); } #[tokio::test(flavor = "multi_thread")] async fn test_block_in_place2() { tokio::task::block_in_place(|| {}); } #[should_panic] #[tokio::main(flavor = "current_thread")] #[test] async fn test_block_in_place3() { tokio::task::block_in_place(|| {}); } #[tokio::main] #[test] async fn test_block_in_place4() { tokio::task::block_in_place(|| {}); } // Repro for tokio-rs/tokio#5239 #[test] fn test_nested_block_in_place_with_block_on_between() { let rt = runtime::Builder::new_multi_thread() .worker_threads(1) // Needs to be more than 0 .max_blocking_threads(1) .build() .unwrap(); // Triggered by a race condition, so run a few times to make sure it is OK. for _ in 0..100 { let h = rt.handle().clone(); rt.block_on(async move { tokio::spawn(async move { tokio::task::block_in_place(|| { h.block_on(async { tokio::task::block_in_place(|| {}); }); }) }) .await .unwrap() }); } } // Testing the tuning logic is tricky as it is inherently timing based, and more // of a heuristic than an exact behavior. This test checks that the interval // changes over time based on load factors. There are no assertions, completion // is sufficient. If there is a regression, this test will hang. In theory, we // could add limits, but that would be likely to fail on CI. #[test] #[cfg(not(tokio_no_tuning_tests))] fn test_tuning() { use std::sync::atomic::AtomicBool; use std::time::Duration; let rt = runtime::Builder::new_multi_thread() .worker_threads(1) .build() .unwrap(); fn iter(flag: Arc, counter: Arc, stall: bool) { if flag.load(Relaxed) { if stall { std::thread::sleep(Duration::from_micros(5)); } counter.fetch_add(1, Relaxed); tokio::spawn(async move { iter(flag, counter, stall) }); } } let flag = Arc::new(AtomicBool::new(true)); let counter = Arc::new(AtomicUsize::new(61)); let interval = Arc::new(AtomicUsize::new(61)); { let flag = flag.clone(); let counter = counter.clone(); rt.spawn(async move { iter(flag, counter, true) }); } // Now, hammer the injection queue until the interval drops. let mut n = 0; loop { let curr = interval.load(Relaxed); if curr <= 8 { n += 1; } else { n = 0; } // Make sure we get a few good rounds. Jitter in the tuning could result // in one "good" value without being representative of reaching a good // state. if n == 3 { break; } if Arc::strong_count(&interval) < 5_000 { let counter = counter.clone(); let interval = interval.clone(); rt.spawn(async move { let prev = counter.swap(0, Relaxed); interval.store(prev, Relaxed); }); std::thread::yield_now(); } } flag.store(false, Relaxed); let w = Arc::downgrade(&interval); drop(interval); while w.strong_count() > 0 { std::thread::sleep(Duration::from_micros(500)); } // Now, run it again with a faster task let flag = Arc::new(AtomicBool::new(true)); // Set it high, we know it shouldn't ever really be this high let counter = Arc::new(AtomicUsize::new(10_000)); let interval = Arc::new(AtomicUsize::new(10_000)); { let flag = flag.clone(); let counter = counter.clone(); rt.spawn(async move { iter(flag, counter, false) }); } // Now, hammer the injection queue until the interval reaches the expected range. let mut n = 0; loop { let curr = interval.load(Relaxed); if curr <= 1_000 && curr > 32 { n += 1; } else { n = 0; } if n == 3 { break; } if Arc::strong_count(&interval) <= 5_000 { let counter = counter.clone(); let interval = interval.clone(); rt.spawn(async move { let prev = counter.swap(0, Relaxed); interval.store(prev, Relaxed); }); } std::thread::yield_now(); } flag.store(false, Relaxed); } fn rt() -> runtime::Runtime { runtime::Runtime::new().unwrap() } #[cfg(tokio_unstable)] mod unstable { use super::*; #[test] fn test_disable_lifo_slot() { use std::sync::mpsc::{channel, RecvTimeoutError}; let rt = runtime::Builder::new_multi_thread() .disable_lifo_slot() .worker_threads(2) .build() .unwrap(); // Spawn a background thread to poke the runtime periodically. // // This is necessary because we may end up triggering the issue in: // // // Spawning a task will wake up the second worker, which will then steal // the task. However, the steal will fail if the task is in the LIFO // slot, because the LIFO slot cannot be stolen. // // Note that this only happens rarely. Most of the time, this thread is // not necessary. let (kill_bg_thread, recv) = channel::<()>(); let handle = rt.handle().clone(); let bg_thread = std::thread::spawn(move || { let one_sec = std::time::Duration::from_secs(1); while recv.recv_timeout(one_sec) == Err(RecvTimeoutError::Timeout) { handle.spawn(async {}); } }); rt.block_on(async { tokio::spawn(async { // Spawn another task and block the thread until completion. If the LIFO slot // is used then the test doesn't complete. futures::executor::block_on(tokio::spawn(async {})).unwrap(); }) .await .unwrap(); }); drop(kill_bg_thread); bg_thread.join().unwrap(); } #[test] fn runtime_id_is_same() { let rt = rt(); let handle1 = rt.handle(); let handle2 = rt.handle(); assert_eq!(handle1.id(), handle2.id()); } #[test] fn runtime_ids_different() { let rt1 = rt(); let rt2 = rt(); assert_ne!(rt1.handle().id(), rt2.handle().id()); } } tokio-1.43.1/tests/rt_threaded_alt.rs000064400000000000000000000450051046102023000157000ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] #![cfg(tokio_unstable)] use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::runtime; use tokio::sync::oneshot; use tokio_test::{assert_err, assert_ok}; use std::future::{poll_fn, Future}; use std::pin::Pin; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::Relaxed; use std::sync::{mpsc, Arc, Mutex}; use std::task::{Context, Poll, Waker}; macro_rules! cfg_metrics { ($($t:tt)*) => { #[cfg(all(tokio_unstable, target_has_atomic = "64"))] { $( $t )* } } } #[test] fn single_thread() { // No panic when starting a runtime w/ a single thread let _ = runtime::Builder::new_multi_thread_alt() .enable_all() .worker_threads(1) .build() .unwrap(); } #[test] #[ignore] // https://github.com/tokio-rs/tokio/issues/5995 fn many_oneshot_futures() { // used for notifying the main thread const NUM: usize = 1_000; for _ in 0..5 { let (tx, rx) = mpsc::channel(); let rt = rt(); let cnt = Arc::new(AtomicUsize::new(0)); for _ in 0..NUM { let cnt = cnt.clone(); let tx = tx.clone(); rt.spawn(async move { let num = cnt.fetch_add(1, Relaxed) + 1; if num == NUM { tx.send(()).unwrap(); } }); } rx.recv().unwrap(); // Wait for the pool to shutdown drop(rt); } } #[test] fn spawn_two() { let rt = rt(); let out = rt.block_on(async { let (tx, rx) = oneshot::channel(); tokio::spawn(async move { tokio::spawn(async move { tx.send("ZOMG").unwrap(); }); }); assert_ok!(rx.await) }); assert_eq!(out, "ZOMG"); cfg_metrics! { let metrics = rt.metrics(); drop(rt); assert_eq!(1, metrics.remote_schedule_count()); let mut local = 0; for i in 0..metrics.num_workers() { local += metrics.worker_local_schedule_count(i); } assert_eq!(1, local); } } #[test] fn many_multishot_futures() { const CHAIN: usize = 200; const CYCLES: usize = 5; const TRACKS: usize = 50; for _ in 0..50 { let rt = rt(); let mut start_txs = Vec::with_capacity(TRACKS); let mut final_rxs = Vec::with_capacity(TRACKS); for _ in 0..TRACKS { let (start_tx, mut chain_rx) = tokio::sync::mpsc::channel(10); for _ in 0..CHAIN { let (next_tx, next_rx) = tokio::sync::mpsc::channel(10); // Forward all the messages rt.spawn(async move { while let Some(v) = chain_rx.recv().await { next_tx.send(v).await.unwrap(); } }); chain_rx = next_rx; } // This final task cycles if needed let (final_tx, final_rx) = tokio::sync::mpsc::channel(10); let cycle_tx = start_tx.clone(); let mut rem = CYCLES; rt.spawn(async move { for _ in 0..CYCLES { let msg = chain_rx.recv().await.unwrap(); rem -= 1; if rem == 0 { final_tx.send(msg).await.unwrap(); } else { cycle_tx.send(msg).await.unwrap(); } } }); start_txs.push(start_tx); final_rxs.push(final_rx); } { rt.block_on(async move { for start_tx in start_txs { start_tx.send("ping").await.unwrap(); } for mut final_rx in final_rxs { final_rx.recv().await.unwrap(); } }); } } } #[test] fn lifo_slot_budget() { async fn my_fn() { spawn_another(); } fn spawn_another() { tokio::spawn(my_fn()); } let rt = runtime::Builder::new_multi_thread_alt() .enable_all() .worker_threads(1) .build() .unwrap(); let (send, recv) = oneshot::channel(); rt.spawn(async move { tokio::spawn(my_fn()); let _ = send.send(()); }); let _ = rt.block_on(recv); } #[test] fn spawn_shutdown() { let rt = rt(); let (tx, rx) = mpsc::channel(); rt.block_on(async { tokio::spawn(client_server(tx.clone())); }); // Use spawner rt.spawn(client_server(tx)); assert_ok!(rx.recv()); assert_ok!(rx.recv()); drop(rt); assert_err!(rx.try_recv()); } async fn client_server(tx: mpsc::Sender<()>) { let server = assert_ok!(TcpListener::bind("127.0.0.1:0").await); // Get the assigned address let addr = assert_ok!(server.local_addr()); // Spawn the server tokio::spawn(async move { // Accept a socket let (mut socket, _) = server.accept().await.unwrap(); // Write some data socket.write_all(b"hello").await.unwrap(); }); let mut client = TcpStream::connect(&addr).await.unwrap(); let mut buf = vec![]; client.read_to_end(&mut buf).await.unwrap(); assert_eq!(buf, b"hello"); tx.send(()).unwrap(); } #[test] fn drop_threadpool_drops_futures() { for _ in 0..1_000 { let num_inc = Arc::new(AtomicUsize::new(0)); let num_dec = Arc::new(AtomicUsize::new(0)); let num_drop = Arc::new(AtomicUsize::new(0)); struct Never(Arc); impl Future for Never { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Pending } } impl Drop for Never { fn drop(&mut self) { self.0.fetch_add(1, Relaxed); } } let a = num_inc.clone(); let b = num_dec.clone(); let rt = runtime::Builder::new_multi_thread_alt() .enable_all() .on_thread_start(move || { a.fetch_add(1, Relaxed); }) .on_thread_stop(move || { b.fetch_add(1, Relaxed); }) .build() .unwrap(); rt.spawn(Never(num_drop.clone())); // Wait for the pool to shutdown drop(rt); // Assert that only a single thread was spawned. let a = num_inc.load(Relaxed); assert!(a >= 1); // Assert that all threads shutdown let b = num_dec.load(Relaxed); assert_eq!(a, b); // Assert that the future was dropped let c = num_drop.load(Relaxed); assert_eq!(c, 1); } } #[test] fn start_stop_callbacks_called() { use std::sync::atomic::{AtomicUsize, Ordering}; let after_start = Arc::new(AtomicUsize::new(0)); let before_stop = Arc::new(AtomicUsize::new(0)); let after_inner = after_start.clone(); let before_inner = before_stop.clone(); let rt = tokio::runtime::Builder::new_multi_thread_alt() .enable_all() .on_thread_start(move || { after_inner.clone().fetch_add(1, Ordering::Relaxed); }) .on_thread_stop(move || { before_inner.clone().fetch_add(1, Ordering::Relaxed); }) .build() .unwrap(); let (tx, rx) = oneshot::channel(); rt.spawn(async move { assert_ok!(tx.send(())); }); assert_ok!(rt.block_on(rx)); drop(rt); assert!(after_start.load(Ordering::Relaxed) > 0); assert!(before_stop.load(Ordering::Relaxed) > 0); } #[test] fn blocking_task() { // used for notifying the main thread const NUM: usize = 1_000; for _ in 0..10 { let (tx, rx) = mpsc::channel(); let rt = rt(); let cnt = Arc::new(AtomicUsize::new(0)); // there are four workers in the pool // so, if we run 4 blocking tasks, we know that handoff must have happened let block = Arc::new(std::sync::Barrier::new(5)); for _ in 0..4 { let block = block.clone(); rt.spawn(async move { tokio::task::block_in_place(move || { block.wait(); block.wait(); }) }); } block.wait(); for _ in 0..NUM { let cnt = cnt.clone(); let tx = tx.clone(); rt.spawn(async move { let num = cnt.fetch_add(1, Relaxed) + 1; if num == NUM { tx.send(()).unwrap(); } }); } rx.recv().unwrap(); // Wait for the pool to shutdown block.wait(); } } #[test] fn multi_threadpool() { use tokio::sync::oneshot; let rt1 = rt(); let rt2 = rt(); let (tx, rx) = oneshot::channel(); let (done_tx, done_rx) = mpsc::channel(); rt2.spawn(async move { rx.await.unwrap(); done_tx.send(()).unwrap(); }); rt1.spawn(async move { tx.send(()).unwrap(); }); done_rx.recv().unwrap(); } // When `block_in_place` returns, it attempts to reclaim the yielded runtime // worker. In this case, the remainder of the task is on the runtime worker and // must take part in the cooperative task budgeting system. // // The test ensures that, when this happens, attempting to consume from a // channel yields occasionally even if there are values ready to receive. #[test] fn coop_and_block_in_place() { let rt = tokio::runtime::Builder::new_multi_thread_alt() // Setting max threads to 1 prevents another thread from claiming the // runtime worker yielded as part of `block_in_place` and guarantees the // same thread will reclaim the worker at the end of the // `block_in_place` call. .max_blocking_threads(1) .build() .unwrap(); rt.block_on(async move { let (tx, mut rx) = tokio::sync::mpsc::channel(1024); // Fill the channel for _ in 0..1024 { tx.send(()).await.unwrap(); } drop(tx); tokio::spawn(async move { // Block in place without doing anything tokio::task::block_in_place(|| {}); // Receive all the values, this should trigger a `Pending` as the // coop limit will be reached. poll_fn(|cx| { while let Poll::Ready(v) = { tokio::pin! { let fut = rx.recv(); } Pin::new(&mut fut).poll(cx) } { if v.is_none() { panic!("did not yield"); } } Poll::Ready(()) }) .await }) .await .unwrap(); }); } #[test] fn yield_after_block_in_place() { let rt = tokio::runtime::Builder::new_multi_thread_alt() .worker_threads(1) .build() .unwrap(); rt.block_on(async { tokio::spawn(async move { // Block in place then enter a new runtime tokio::task::block_in_place(|| { let rt = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); rt.block_on(async {}); }); // Yield, then complete tokio::task::yield_now().await; }) .await .unwrap() }); } // Testing this does not panic #[test] fn max_blocking_threads() { let _rt = tokio::runtime::Builder::new_multi_thread_alt() .max_blocking_threads(1) .build() .unwrap(); } #[test] #[should_panic] fn max_blocking_threads_set_to_zero() { let _rt = tokio::runtime::Builder::new_multi_thread_alt() .max_blocking_threads(0) .build() .unwrap(); } /// Regression test for #6445. /// /// After #6445, setting `global_queue_interval` to 1 is now technically valid. /// This test confirms that there is no regression in `multi_thread_runtime` /// when global_queue_interval is set to 1. #[test] fn global_queue_interval_set_to_one() { let rt = tokio::runtime::Builder::new_multi_thread_alt() .global_queue_interval(1) .build() .unwrap(); // Perform a simple work. let cnt = Arc::new(AtomicUsize::new(0)); rt.block_on(async { let mut set = tokio::task::JoinSet::new(); for _ in 0..10 { let cnt = cnt.clone(); set.spawn(async move { cnt.fetch_add(1, Relaxed) }); } while let Some(res) = set.join_next().await { res.unwrap(); } }); assert_eq!(cnt.load(Relaxed), 10); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn hang_on_shutdown() { let (sync_tx, sync_rx) = std::sync::mpsc::channel::<()>(); tokio::spawn(async move { tokio::task::block_in_place(|| sync_rx.recv().ok()); }); tokio::spawn(async { tokio::time::sleep(std::time::Duration::from_secs(2)).await; drop(sync_tx); }); tokio::time::sleep(std::time::Duration::from_secs(1)).await; } /// Demonstrates tokio-rs/tokio#3869 #[test] fn wake_during_shutdown() { struct Shared { waker: Option, } struct MyFuture { shared: Arc>, put_waker: bool, } impl MyFuture { fn new() -> (Self, Self) { let shared = Arc::new(Mutex::new(Shared { waker: None })); let f1 = MyFuture { shared: shared.clone(), put_waker: true, }; let f2 = MyFuture { shared, put_waker: false, }; (f1, f2) } } impl Future for MyFuture { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { let me = Pin::into_inner(self); let mut lock = me.shared.lock().unwrap(); if me.put_waker { lock.waker = Some(cx.waker().clone()); } Poll::Pending } } impl Drop for MyFuture { fn drop(&mut self) { let mut lock = self.shared.lock().unwrap(); if !self.put_waker { lock.waker.take().unwrap().wake(); } drop(lock); } } let rt = tokio::runtime::Builder::new_multi_thread_alt() .worker_threads(1) .enable_all() .build() .unwrap(); let (f1, f2) = MyFuture::new(); rt.spawn(f1); rt.spawn(f2); rt.block_on(async { tokio::time::sleep(tokio::time::Duration::from_millis(20)).await }); } #[should_panic] #[tokio::test] async fn test_block_in_place1() { tokio::task::block_in_place(|| {}); } #[tokio::test(flavor = "multi_thread")] async fn test_block_in_place2() { tokio::task::block_in_place(|| {}); } #[should_panic] #[tokio::main(flavor = "current_thread")] #[test] async fn test_block_in_place3() { tokio::task::block_in_place(|| {}); } #[tokio::main] #[test] async fn test_block_in_place4() { tokio::task::block_in_place(|| {}); } // Testing the tuning logic is tricky as it is inherently timing based, and more // of a heuristic than an exact behavior. This test checks that the interval // changes over time based on load factors. There are no assertions, completion // is sufficient. If there is a regression, this test will hang. In theory, we // could add limits, but that would be likely to fail on CI. #[test] #[cfg(not(tokio_no_tuning_tests))] fn test_tuning() { use std::sync::atomic::AtomicBool; use std::time::Duration; let rt = runtime::Builder::new_multi_thread_alt() .worker_threads(1) .build() .unwrap(); fn iter(flag: Arc, counter: Arc, stall: bool) { if flag.load(Relaxed) { if stall { std::thread::sleep(Duration::from_micros(5)); } counter.fetch_add(1, Relaxed); tokio::spawn(async move { iter(flag, counter, stall) }); } } let flag = Arc::new(AtomicBool::new(true)); let counter = Arc::new(AtomicUsize::new(61)); let interval = Arc::new(AtomicUsize::new(61)); { let flag = flag.clone(); let counter = counter.clone(); rt.spawn(async move { iter(flag, counter, true) }); } // Now, hammer the injection queue until the interval drops. let mut n = 0; loop { let curr = interval.load(Relaxed); if curr <= 8 { n += 1; } else { n = 0; } // Make sure we get a few good rounds. Jitter in the tuning could result // in one "good" value without being representative of reaching a good // state. if n == 3 { break; } if Arc::strong_count(&interval) < 5_000 { let counter = counter.clone(); let interval = interval.clone(); rt.spawn(async move { let prev = counter.swap(0, Relaxed); interval.store(prev, Relaxed); }); std::thread::yield_now(); } } flag.store(false, Relaxed); let w = Arc::downgrade(&interval); drop(interval); while w.strong_count() > 0 { std::thread::sleep(Duration::from_micros(500)); } // Now, run it again with a faster task let flag = Arc::new(AtomicBool::new(true)); // Set it high, we know it shouldn't ever really be this high let counter = Arc::new(AtomicUsize::new(10_000)); let interval = Arc::new(AtomicUsize::new(10_000)); { let flag = flag.clone(); let counter = counter.clone(); rt.spawn(async move { iter(flag, counter, false) }); } // Now, hammer the injection queue until the interval reaches the expected range. let mut n = 0; loop { let curr = interval.load(Relaxed); if curr <= 1_000 && curr > 32 { n += 1; } else { n = 0; } if n == 3 { break; } if Arc::strong_count(&interval) <= 5_000 { let counter = counter.clone(); let interval = interval.clone(); rt.spawn(async move { let prev = counter.swap(0, Relaxed); interval.store(prev, Relaxed); }); } std::thread::yield_now(); } flag.store(false, Relaxed); } fn rt() -> runtime::Runtime { runtime::Builder::new_multi_thread_alt() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/rt_time_start_paused.rs000064400000000000000000000005451046102023000167740ustar 00000000000000#![cfg(feature = "full")] use tokio::time::{Duration, Instant}; #[tokio::test(start_paused = true)] async fn test_start_paused() { let now = Instant::now(); // Pause a few times w/ std sleep and ensure `now` stays the same for _ in 0..5 { std::thread::sleep(Duration::from_millis(1)); assert_eq!(now, Instant::now()); } } tokio-1.43.1/tests/rt_unstable_metrics.rs000064400000000000000000000630221046102023000166220ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all( feature = "full", tokio_unstable, not(target_os = "wasi"), target_has_atomic = "64" ))] use std::future::Future; use std::sync::{mpsc, Arc, Mutex}; use std::task::Poll; use std::thread; use tokio::macros::support::poll_fn; use tokio::runtime::{HistogramConfiguration, HistogramScale, LogHistogram, Runtime}; use tokio::task::consume_budget; use tokio::time::{self, Duration}; #[test] fn num_workers() { let rt = current_thread(); assert_eq!(1, rt.metrics().num_workers()); let rt = threaded(); assert_eq!(2, rt.metrics().num_workers()); } #[test] fn num_blocking_threads() { let rt = current_thread(); assert_eq!(0, rt.metrics().num_blocking_threads()); let _ = rt.block_on(rt.spawn_blocking(move || {})); assert_eq!(1, rt.metrics().num_blocking_threads()); let rt = threaded(); assert_eq!(0, rt.metrics().num_blocking_threads()); let _ = rt.block_on(rt.spawn_blocking(move || {})); assert_eq!(1, rt.metrics().num_blocking_threads()); } #[test] fn num_idle_blocking_threads() { let rt = current_thread(); assert_eq!(0, rt.metrics().num_idle_blocking_threads()); let _ = rt.block_on(rt.spawn_blocking(move || {})); rt.block_on(async { time::sleep(Duration::from_millis(5)).await; }); // We need to wait until the blocking thread has become idle. Usually 5ms is // enough for this to happen, but not always. When it isn't enough, sleep // for another second. We don't always wait for a whole second since we want // the test suite to finish quickly. // // Note that the timeout for idle threads to be killed is 10 seconds. if 0 == rt.metrics().num_idle_blocking_threads() { rt.block_on(async { time::sleep(Duration::from_secs(1)).await; }); } assert_eq!(1, rt.metrics().num_idle_blocking_threads()); } #[test] fn blocking_queue_depth() { let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .max_blocking_threads(1) .build() .unwrap(); assert_eq!(0, rt.metrics().blocking_queue_depth()); let ready = Arc::new(Mutex::new(())); let guard = ready.lock().unwrap(); let ready_cloned = ready.clone(); let wait_until_ready = move || { let _unused = ready_cloned.lock().unwrap(); }; let h1 = rt.spawn_blocking(wait_until_ready.clone()); let h2 = rt.spawn_blocking(wait_until_ready); assert!(rt.metrics().blocking_queue_depth() > 0); drop(guard); let _ = rt.block_on(h1); let _ = rt.block_on(h2); assert_eq!(0, rt.metrics().blocking_queue_depth()); } #[test] fn spawned_tasks_count() { let rt = current_thread(); let metrics = rt.metrics(); assert_eq!(0, metrics.spawned_tasks_count()); rt.block_on(rt.spawn(async move { assert_eq!(1, metrics.spawned_tasks_count()); })) .unwrap(); assert_eq!(1, rt.metrics().spawned_tasks_count()); let rt = threaded(); let metrics = rt.metrics(); assert_eq!(0, metrics.spawned_tasks_count()); rt.block_on(rt.spawn(async move { assert_eq!(1, metrics.spawned_tasks_count()); })) .unwrap(); assert_eq!(1, rt.metrics().spawned_tasks_count()); } #[test] fn remote_schedule_count() { use std::thread; let rt = current_thread(); let handle = rt.handle().clone(); let task = thread::spawn(move || { handle.spawn(async { // DO nothing }) }) .join() .unwrap(); rt.block_on(task).unwrap(); assert_eq!(1, rt.metrics().remote_schedule_count()); let rt = threaded(); let handle = rt.handle().clone(); let task = thread::spawn(move || { handle.spawn(async { // DO nothing }) }) .join() .unwrap(); rt.block_on(task).unwrap(); assert_eq!(1, rt.metrics().remote_schedule_count()); } #[test] fn worker_thread_id_current_thread() { let rt = current_thread(); let metrics = rt.metrics(); // Check that runtime is on this thread. rt.block_on(async {}); assert_eq!(Some(thread::current().id()), metrics.worker_thread_id(0)); // Move runtime to another thread. let thread_id = std::thread::scope(|scope| { let join_handle = scope.spawn(|| { rt.block_on(async {}); }); join_handle.thread().id() }); assert_eq!(Some(thread_id), metrics.worker_thread_id(0)); // Move runtime back to this thread. rt.block_on(async {}); assert_eq!(Some(thread::current().id()), metrics.worker_thread_id(0)); } #[test] fn worker_thread_id_threaded() { let rt = threaded(); let metrics = rt.metrics(); rt.block_on(rt.spawn(async move { // Check that we are running on a worker thread and determine // the index of our worker. let thread_id = std::thread::current().id(); let this_worker = (0..2) .position(|w| metrics.worker_thread_id(w) == Some(thread_id)) .expect("task not running on any worker thread"); // Force worker to another thread. let moved_thread_id = tokio::task::block_in_place(|| { assert_eq!(thread_id, std::thread::current().id()); // Wait for worker to move to another thread. for _ in 0..100 { let new_id = metrics.worker_thread_id(this_worker).unwrap(); if thread_id != new_id { return new_id; } std::thread::sleep(Duration::from_millis(100)); } panic!("worker did not move to new thread"); }); // After blocking task worker either stays on new thread or // is moved back to current thread. assert!( metrics.worker_thread_id(this_worker) == Some(moved_thread_id) || metrics.worker_thread_id(this_worker) == Some(thread_id) ); })) .unwrap() } #[test] fn worker_park_count() { let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(async { time::sleep(Duration::from_millis(1)).await; }); drop(rt); assert!(1 <= metrics.worker_park_count(0)); let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { time::sleep(Duration::from_millis(1)).await; }); drop(rt); assert!(1 <= metrics.worker_park_count(0)); assert!(1 <= metrics.worker_park_count(1)); } #[test] fn worker_park_unpark_count() { let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(rt.spawn(async {})).unwrap(); drop(rt); assert!(2 <= metrics.worker_park_unpark_count(0)); let rt = threaded(); let metrics = rt.metrics(); // Wait for workers to be parked after runtime startup. for _ in 0..100 { if 1 <= metrics.worker_park_unpark_count(0) && 1 <= metrics.worker_park_unpark_count(1) { break; } std::thread::sleep(std::time::Duration::from_millis(100)); } assert_eq!(1, metrics.worker_park_unpark_count(0)); assert_eq!(1, metrics.worker_park_unpark_count(1)); // Spawn a task to unpark and then park a worker. rt.block_on(rt.spawn(async {})).unwrap(); for _ in 0..100 { if 3 <= metrics.worker_park_unpark_count(0) || 3 <= metrics.worker_park_unpark_count(1) { break; } std::thread::sleep(std::time::Duration::from_millis(100)); } assert!(3 <= metrics.worker_park_unpark_count(0) || 3 <= metrics.worker_park_unpark_count(1)); // Both threads unpark for runtime shutdown. drop(rt); assert_eq!(0, metrics.worker_park_unpark_count(0) % 2); assert_eq!(0, metrics.worker_park_unpark_count(1) % 2); assert!(4 <= metrics.worker_park_unpark_count(0) || 4 <= metrics.worker_park_unpark_count(1)); } #[test] fn worker_noop_count() { // There isn't really a great way to generate no-op parks as they happen as // false-positive events under concurrency. let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(async { time::sleep(Duration::from_millis(1)).await; }); drop(rt); assert!(0 < metrics.worker_noop_count(0)); let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { time::sleep(Duration::from_millis(1)).await; }); drop(rt); assert!(0 < metrics.worker_noop_count(0)); assert!(0 < metrics.worker_noop_count(1)); } #[test] fn worker_steal_count() { // This metric only applies to the multi-threaded runtime. for _ in 0..10 { let rt = threaded_no_lifo(); let metrics = rt.metrics(); let successfully_spawned_stealable_task = rt.block_on(async { // The call to `try_spawn_stealable_task` may time out, which means // that the sending task couldn't be scheduled due to a deadlock in // the runtime. // This is expected behaviour, we just retry until we succeed or // exhaust all tries, the latter causing this test to fail. try_spawn_stealable_task().await.is_ok() }); drop(rt); if successfully_spawned_stealable_task { let n: u64 = (0..metrics.num_workers()) .map(|i| metrics.worker_steal_count(i)) .sum(); assert_eq!(1, n); return; } } panic!("exhausted every try to schedule the stealable task"); } #[test] fn worker_poll_count_and_time() { const N: u64 = 5; async fn task() { // Sync sleep std::thread::sleep(std::time::Duration::from_micros(10)); } let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { tokio::spawn(task()).await.unwrap(); } }); drop(rt); assert_eq!(N, metrics.worker_poll_count(0)); // Not currently supported for current-thread runtime assert_eq!(Duration::default(), metrics.worker_mean_poll_time(0)); // Does not populate the histogram assert!(!metrics.poll_time_histogram_enabled()); for i in 0..10 { assert_eq!(0, metrics.poll_time_histogram_bucket_count(0, i)); } let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { tokio::spawn(task()).await.unwrap(); } }); drop(rt); // Account for the `block_on` task let n = (0..metrics.num_workers()) .map(|i| metrics.worker_poll_count(i)) .sum(); assert_eq!(N, n); let n: Duration = (0..metrics.num_workers()) .map(|i| metrics.worker_mean_poll_time(i)) .sum(); assert!(n > Duration::default()); // Does not populate the histogram assert!(!metrics.poll_time_histogram_enabled()); for n in 0..metrics.num_workers() { for i in 0..10 { assert_eq!(0, metrics.poll_time_histogram_bucket_count(n, i)); } } } #[test] fn log_histogram() { const N: u64 = 50; let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_time_histogram_configuration(HistogramConfiguration::log( LogHistogram::builder() .max_value(Duration::from_secs(60)) .min_value(Duration::from_nanos(100)) .max_error(0.25), )) .build() .unwrap(); let metrics = rt.metrics(); let num_buckets = rt.metrics().poll_time_histogram_num_buckets(); assert_eq!(num_buckets, 119); rt.block_on(async { for _ in 0..N { tokio::spawn(async {}).await.unwrap(); } }); drop(rt); assert_eq!( metrics.poll_time_histogram_bucket_range(0), Duration::from_nanos(0)..Duration::from_nanos(96) ); assert_eq!( metrics.poll_time_histogram_bucket_range(1), Duration::from_nanos(96)..Duration::from_nanos(96 + 2_u64.pow(4)) ); assert_eq!( metrics.poll_time_histogram_bucket_range(118).end, Duration::from_nanos(u64::MAX) ); let n = (0..metrics.num_workers()) .flat_map(|i| (0..num_buckets).map(move |j| (i, j))) .map(|(worker, bucket)| metrics.poll_time_histogram_bucket_count(worker, bucket)) .sum(); assert_eq!(N, n); } #[test] fn minimal_log_histogram() { let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_time_histogram_configuration(HistogramConfiguration::log( LogHistogram::builder() .max_value(Duration::from_millis(4)) .min_value(Duration::from_micros(20)) .precision_exact(0), )) .build() .unwrap(); let metrics = rt.metrics(); let num_buckets = rt.metrics().poll_time_histogram_num_buckets(); for b in 1..num_buckets - 1 { let range = metrics.poll_time_histogram_bucket_range(b); let size = range.end - range.start; // Assert the buckets continue doubling in size assert_eq!( size, Duration::from_nanos((1 << (b - 1)) * 16384), "incorrect range for {b}" ); } assert_eq!(num_buckets, 10); } #[test] #[allow(deprecated)] fn legacy_log_histogram() { let rt = tokio::runtime::Builder::new_multi_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_count_histogram_scale(HistogramScale::Log) .metrics_poll_count_histogram_resolution(Duration::from_micros(50)) .metrics_poll_count_histogram_buckets(20) .build() .unwrap(); let num_buckets = rt.metrics().poll_time_histogram_num_buckets(); assert_eq!(num_buckets, 20); } #[test] fn log_histogram_default_configuration() { let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_time_histogram_configuration(HistogramConfiguration::log( LogHistogram::default(), )) .build() .unwrap(); let num_buckets = rt.metrics().poll_time_histogram_num_buckets(); assert_eq!(num_buckets, 119); } #[test] fn worker_poll_count_histogram() { const N: u64 = 5; let rts = [ tokio::runtime::Builder::new_current_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_time_histogram_configuration(HistogramConfiguration::linear( Duration::from_millis(50), 3, )) .build() .unwrap(), tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_time_histogram_configuration(HistogramConfiguration::linear( Duration::from_millis(50), 3, )) .build() .unwrap(), ]; for rt in rts { let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { tokio::spawn(async {}).await.unwrap(); } }); drop(rt); let num_workers = metrics.num_workers(); let num_buckets = metrics.poll_time_histogram_num_buckets(); assert!(metrics.poll_time_histogram_enabled()); assert_eq!(num_buckets, 3); let n = (0..num_workers) .flat_map(|i| (0..num_buckets).map(move |j| (i, j))) .map(|(worker, bucket)| metrics.poll_time_histogram_bucket_count(worker, bucket)) .sum(); assert_eq!(N, n); } } #[test] fn worker_poll_count_histogram_range() { let max = Duration::from_nanos(u64::MAX); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_time_histogram_configuration(HistogramConfiguration::linear(us(50), 3)) .build() .unwrap(); let metrics = rt.metrics(); assert_eq!(metrics.poll_time_histogram_bucket_range(0), us(0)..us(50)); assert_eq!(metrics.poll_time_histogram_bucket_range(1), us(50)..us(100)); assert_eq!(metrics.poll_time_histogram_bucket_range(2), us(100)..max); // ensure the old methods work too #[allow(deprecated)] let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .enable_metrics_poll_time_histogram() .metrics_poll_count_histogram_scale(tokio::runtime::HistogramScale::Log) .metrics_poll_count_histogram_buckets(3) .metrics_poll_count_histogram_resolution(us(50)) .build() .unwrap(); let metrics = rt.metrics(); let a = Duration::from_nanos(50000_u64.next_power_of_two()); let b = a * 2; assert_eq!(metrics.poll_time_histogram_bucket_range(0), us(0)..a); assert_eq!(metrics.poll_time_histogram_bucket_range(1), a..b); assert_eq!(metrics.poll_time_histogram_bucket_range(2), b..max); } #[test] fn worker_poll_count_histogram_disabled_without_explicit_enable() { let rts = [ tokio::runtime::Builder::new_current_thread() .enable_all() .metrics_poll_time_histogram_configuration(HistogramConfiguration::linear( Duration::from_millis(50), 3, )) .build() .unwrap(), tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .metrics_poll_time_histogram_configuration(HistogramConfiguration::linear( Duration::from_millis(50), 3, )) .build() .unwrap(), ]; for rt in rts { let metrics = rt.metrics(); assert!(!metrics.poll_time_histogram_enabled()); } } #[test] fn worker_total_busy_duration() { const N: usize = 5; let zero = Duration::from_millis(0); let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { tokio::spawn(async { tokio::task::yield_now().await; }) .await .unwrap(); } }); drop(rt); assert!(zero < metrics.worker_total_busy_duration(0)); let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { tokio::spawn(async { tokio::task::yield_now().await; }) .await .unwrap(); } }); drop(rt); for i in 0..metrics.num_workers() { assert!(zero < metrics.worker_total_busy_duration(i)); } } #[test] fn worker_local_schedule_count() { let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(async { tokio::spawn(async {}).await.unwrap(); }); drop(rt); assert_eq!(1, metrics.worker_local_schedule_count(0)); assert_eq!(0, metrics.remote_schedule_count()); let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { // Move to the runtime tokio::spawn(async { tokio::spawn(async {}).await.unwrap(); }) .await .unwrap(); }); drop(rt); let n: u64 = (0..metrics.num_workers()) .map(|i| metrics.worker_local_schedule_count(i)) .sum(); assert_eq!(2, n); assert_eq!(1, metrics.remote_schedule_count()); } #[test] fn worker_overflow_count() { // Only applies to the threaded worker let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { // Move to the runtime tokio::spawn(async { let (tx1, rx1) = std::sync::mpsc::channel(); let (tx2, rx2) = std::sync::mpsc::channel(); // First, we need to block the other worker until all tasks have // been spawned. // // We spawn from outside the runtime to ensure that the other worker // will pick it up: // tokio::task::spawn_blocking(|| { tokio::spawn(async move { tx1.send(()).unwrap(); rx2.recv().unwrap(); }); }); rx1.recv().unwrap(); // Spawn many tasks for _ in 0..300 { tokio::spawn(async {}); } tx2.send(()).unwrap(); }) .await .unwrap(); }); drop(rt); let n: u64 = (0..metrics.num_workers()) .map(|i| metrics.worker_overflow_count(i)) .sum(); assert_eq!(1, n); } #[test] fn worker_local_queue_depth() { const N: usize = 100; let rt = current_thread(); let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { tokio::spawn(async {}); } assert_eq!(N, metrics.worker_local_queue_depth(0)); }); let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async move { // Move to the runtime tokio::spawn(async move { let (tx1, rx1) = std::sync::mpsc::channel(); let (tx2, rx2) = std::sync::mpsc::channel(); // First, we need to block the other worker until all tasks have // been spawned. tokio::spawn(async move { tx1.send(()).unwrap(); rx2.recv().unwrap(); }); // Bump the next-run spawn tokio::spawn(async {}); rx1.recv().unwrap(); // Spawn some tasks for _ in 0..100 { tokio::spawn(async {}); } let n: usize = (0..metrics.num_workers()) .map(|i| metrics.worker_local_queue_depth(i)) .sum(); assert_eq!(n, N); tx2.send(()).unwrap(); }) .await .unwrap(); }); } #[test] fn budget_exhaustion_yield() { let rt = current_thread(); let metrics = rt.metrics(); assert_eq!(0, metrics.budget_forced_yield_count()); let mut did_yield = false; // block on a task which consumes budget until it yields rt.block_on(poll_fn(|cx| loop { if did_yield { return Poll::Ready(()); } let fut = consume_budget(); tokio::pin!(fut); if fut.poll(cx).is_pending() { did_yield = true; return Poll::Pending; } })); assert_eq!(1, rt.metrics().budget_forced_yield_count()); } #[test] fn budget_exhaustion_yield_with_joins() { let rt = current_thread(); let metrics = rt.metrics(); assert_eq!(0, metrics.budget_forced_yield_count()); let mut did_yield_1 = false; let mut did_yield_2 = false; // block on a task which consumes budget until it yields rt.block_on(async { tokio::join!( poll_fn(|cx| loop { if did_yield_1 { return Poll::Ready(()); } let fut = consume_budget(); tokio::pin!(fut); if fut.poll(cx).is_pending() { did_yield_1 = true; return Poll::Pending; } }), poll_fn(|cx| loop { if did_yield_2 { return Poll::Ready(()); } let fut = consume_budget(); tokio::pin!(fut); if fut.poll(cx).is_pending() { did_yield_2 = true; return Poll::Pending; } }) ) }); assert_eq!(1, rt.metrics().budget_forced_yield_count()); } #[cfg(any(target_os = "linux", target_os = "macos"))] #[test] fn io_driver_fd_count() { let rt = current_thread(); let metrics = rt.metrics(); assert_eq!(metrics.io_driver_fd_registered_count(), 0); let stream = tokio::net::TcpStream::connect("google.com:80"); let stream = rt.block_on(async move { stream.await.unwrap() }); assert_eq!(metrics.io_driver_fd_registered_count(), 1); assert_eq!(metrics.io_driver_fd_deregistered_count(), 0); drop(stream); assert_eq!(metrics.io_driver_fd_deregistered_count(), 1); assert_eq!(metrics.io_driver_fd_registered_count(), 1); } #[cfg(any(target_os = "linux", target_os = "macos"))] #[test] fn io_driver_ready_count() { let rt = current_thread(); let metrics = rt.metrics(); let stream = tokio::net::TcpStream::connect("google.com:80"); let _stream = rt.block_on(async move { stream.await.unwrap() }); assert_eq!(metrics.io_driver_ready_count(), 1); } async fn try_spawn_stealable_task() -> Result<(), mpsc::RecvTimeoutError> { // We use a blocking channel to synchronize the tasks. let (tx, rx) = mpsc::channel(); // Make sure we are in the context of the runtime. tokio::spawn(async move { // Spawn the task that sends to the channel. // // Note that the runtime needs to have the lifo slot disabled to make // this task stealable. tokio::spawn(async move { tx.send(()).unwrap(); }); // Blocking receive on the channel, timing out if the sending task // wasn't scheduled in time. rx.recv_timeout(Duration::from_secs(1)) }) .await .unwrap()?; Ok(()) } fn current_thread() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } fn threaded() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .build() .unwrap() } fn threaded_no_lifo() -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .disable_lifo_slot() .enable_all() .build() .unwrap() } fn us(n: u64) -> Duration { Duration::from_micros(n) } tokio-1.43.1/tests/signal_ctrl_c.rs000064400000000000000000000006231046102023000153530ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` on Miri mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal; use tokio_test::assert_ok; #[tokio::test] async fn ctrl_c() { let ctrl_c = signal::ctrl_c(); tokio::spawn(async { send_signal(libc::SIGINT); }); assert_ok!(ctrl_c.await); } tokio-1.43.1/tests/signal_drop_recv.rs000064400000000000000000000010461046102023000160700ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` in Miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal::unix::{signal, SignalKind}; #[tokio::test] async fn drop_then_get_a_signal() { let kind = SignalKind::user_defined1(); let sig = signal(kind).expect("failed to create first signal"); drop(sig); send_signal(libc::SIGUSR1); let mut sig = signal(kind).expect("failed to create second signal"); let _ = sig.recv().await; } tokio-1.43.1/tests/signal_drop_rt.rs000064400000000000000000000020331046102023000155530ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` in miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::runtime::Runtime; use tokio::signal::unix::{signal, SignalKind}; #[test] fn dropping_loops_does_not_cause_starvation() { let kind = SignalKind::user_defined1(); let first_rt = rt(); let mut first_signal = first_rt.block_on(async { signal(kind).expect("failed to register first signal") }); let second_rt = rt(); let mut second_signal = second_rt.block_on(async { signal(kind).expect("failed to register second signal") }); send_signal(libc::SIGUSR1); first_rt .block_on(first_signal.recv()) .expect("failed to await first signal"); drop(first_rt); drop(first_signal); send_signal(libc::SIGUSR1); second_rt.block_on(second_signal.recv()); } fn rt() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/signal_drop_signal.rs000064400000000000000000000014601046102023000164060ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` in miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal::unix::{signal, SignalKind}; #[tokio::test] async fn dropping_signal_does_not_deregister_any_other_instances() { let kind = SignalKind::user_defined1(); // Signals should not starve based on ordering let first_duplicate_signal = signal(kind).expect("failed to register first duplicate signal"); let mut sig = signal(kind).expect("failed to register signal"); let second_duplicate_signal = signal(kind).expect("failed to register second duplicate signal"); drop(first_duplicate_signal); drop(second_duplicate_signal); send_signal(libc::SIGUSR1); let _ = sig.recv().await; } tokio-1.43.1/tests/signal_info.rs000064400000000000000000000015311046102023000150370ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(any( target_os = "dragonfly", target_os = "freebsd", target_os = "macos", target_os = "netbsd", target_os = "openbsd", target_os = "illumos" ))] #![cfg(not(miri))] // No `sigaction` on Miri mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal; use tokio::signal::unix::SignalKind; use tokio::time::{timeout, Duration}; #[tokio::test] async fn siginfo() { let mut sig = signal::unix::signal(SignalKind::info()).expect("installed signal handler"); tokio::spawn(async { send_signal(libc::SIGINFO); }); // Add a timeout to ensure the test doesn't hang. timeout(Duration::from_secs(5), sig.recv()) .await .expect("received SIGINFO signal in time") .expect("received SIGINFO signal"); } tokio-1.43.1/tests/signal_multi_rt.rs000064400000000000000000000027461046102023000157540ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` on Miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::runtime::Runtime; use tokio::signal::unix::{signal, SignalKind}; use std::sync::mpsc::channel; use std::thread; #[test] fn multi_loop() { // An "ordinary" (non-future) channel let (sender, receiver) = channel(); // Run multiple times, to make sure there are no race conditions for _ in 0..10 { // Run multiple event loops, each one in its own thread let threads: Vec<_> = (0..4) .map(|_| { let sender = sender.clone(); thread::spawn(move || { let rt = rt(); let _ = rt.block_on(async { let mut signal = signal(SignalKind::hangup()).unwrap(); sender.send(()).unwrap(); signal.recv().await }); }) }) .collect(); // Wait for them to declare they're ready for &_ in threads.iter() { receiver.recv().unwrap(); } // Send a signal send_signal(libc::SIGHUP); // Make sure the threads terminated correctly for t in threads { t.join().unwrap(); } } } fn rt() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/signal_no_rt.rs000064400000000000000000000005411046102023000152250ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` on Miri. use tokio::signal::unix::{signal, SignalKind}; #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] #[test] #[should_panic] fn no_runtime_panics_creating_signals() { let _ = signal(SignalKind::hangup()); } tokio-1.43.1/tests/signal_notify_both.rs000064400000000000000000000010441046102023000164270ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` on Miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal::unix::{signal, SignalKind}; #[tokio::test] async fn notify_both() { let kind = SignalKind::user_defined2(); let mut signal1 = signal(kind).expect("failed to create signal1"); let mut signal2 = signal(kind).expect("failed to create signal2"); send_signal(libc::SIGUSR2); signal1.recv().await; signal2.recv().await; } tokio-1.43.1/tests/signal_panic.rs000064400000000000000000000013401046102023000151740ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(panic = "unwind")] #![cfg(not(miri))] // No `sigaction` on Miri. use std::error::Error; use tokio::runtime::Builder; use tokio::signal::unix::{signal, SignalKind}; mod support { pub mod panic; } use support::panic::test_panic; #[test] fn signal_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async { let kind = SignalKind::from_raw(-1); let _ = signal(kind); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } tokio-1.43.1/tests/signal_realtime.rs000064400000000000000000000062261046102023000157140ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(any(target_os = "linux", target_os = "illumos"))] #![cfg(not(miri))] // No `sigaction` in Miri. mod support { pub mod signal; } use libc::c_int; use support::signal::send_signal; use futures::stream::{FuturesUnordered, StreamExt}; use std::collections::HashMap; use tokio::signal::unix::{signal, SignalKind}; use tokio::time::{sleep, Duration}; use tokio_test::assert_ok; #[tokio::test] async fn signal_realtime() { // Attempt to register a real-time signal for everything between SIGRTMIN // and SIGRTMAX. let signals = (libc::SIGRTMIN()..=sigrt_max()) .map(|signum| { let sig = assert_ok!( signal(SignalKind::from_raw(signum)), "failed to create signal for {}", sigrtnum_to_string(signum), ); (signum, sig) }) .collect::>(); eprintln!( "registered {} signals in the range {}..={}", signals.len(), libc::SIGRTMIN(), libc::SIGRTMAX() ); // Now send signals to each of the registered signals. for signum in libc::SIGRTMIN()..=sigrt_max() { send_signal(signum); } let futures = signals .into_iter() .map(|(signum, mut sig)| async move { let res = sig.recv().await; (signum, res) }) .collect::>(); // Ensure that all signals are received in time -- attempt to get whatever // we can. let sleep = std::pin::pin!(sleep(Duration::from_secs(5))); let done = futures.take_until(sleep).collect::>().await; let mut none = Vec::new(); let mut missing = Vec::new(); for signum in libc::SIGRTMIN()..=sigrt_max() { match done.get(&signum) { Some(Some(())) => {} Some(None) => none.push(signum), None => missing.push(signum), } } if none.is_empty() && missing.is_empty() { return; } let mut msg = String::new(); if !none.is_empty() { msg.push_str("no signals received for:\n"); for signum in none { msg.push_str(&format!("- {}\n", sigrtnum_to_string(signum))); } } if !missing.is_empty() { msg.push_str("missing signals for:\n"); for signum in missing { msg.push_str(&format!("- {}\n", sigrtnum_to_string(signum))); } } panic!("{}", msg); } fn sigrt_max() -> c_int { // Generally, you would expect this to be SIGRTMAX. But QEMU only supports // 28 real-time signals even though it might report SIGRTMAX to be higher. // See https://wiki.qemu.org/ChangeLog/9.2#signals. // // The goal of this test is to test that real-time signals are supported in // general, not necessarily that every single signal is supported (which, as // this example suggests, depends on the execution environment). So cap this // at SIGRTMIN+27 (i.e. SIGRTMIN..=SIGRTMIN+27, so 28 signals inclusive). libc::SIGRTMAX().min(libc::SIGRTMIN() + 27) } fn sigrtnum_to_string(signum: i32) -> String { format!("SIGRTMIN+{} (signal {})", signum - libc::SIGRTMIN(), signum) } tokio-1.43.1/tests/signal_twice.rs000064400000000000000000000007451046102023000152250ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` on Miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal::unix::{signal, SignalKind}; #[tokio::test] async fn twice() { let kind = SignalKind::user_defined1(); let mut sig = signal(kind).expect("failed to get signal"); for _ in 0..2 { send_signal(libc::SIGUSR1); assert!(sig.recv().await.is_some()); } } tokio-1.43.1/tests/signal_usr1.rs000064400000000000000000000007401046102023000147770ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `sigaction` in Miri. mod support { pub mod signal; } use support::signal::send_signal; use tokio::signal::unix::{signal, SignalKind}; use tokio_test::assert_ok; #[tokio::test] async fn signal_usr1() { let mut signal = assert_ok!( signal(SignalKind::user_defined1()), "failed to create signal" ); send_signal(libc::SIGUSR1); signal.recv().await; } tokio-1.43.1/tests/support/io_vec.rs000064400000000000000000000023511046102023000155300ustar 00000000000000use std::io::IoSlice; use std::ops::Deref; use std::slice; pub struct IoBufs<'a, 'b>(&'b mut [IoSlice<'a>]); impl<'a, 'b> IoBufs<'a, 'b> { pub fn new(slices: &'b mut [IoSlice<'a>]) -> Self { IoBufs(slices) } pub fn is_empty(&self) -> bool { self.0.is_empty() } pub fn advance(mut self, n: usize) -> IoBufs<'a, 'b> { let mut to_remove = 0; let mut remaining_len = n; for slice in self.0.iter() { if remaining_len < slice.len() { break; } else { remaining_len -= slice.len(); to_remove += 1; } } self.0 = self.0.split_at_mut(to_remove).1; if let Some(slice) = self.0.first_mut() { let tail = &slice[remaining_len..]; // Safety: recasts slice to the original lifetime let tail = unsafe { slice::from_raw_parts(tail.as_ptr(), tail.len()) }; *slice = IoSlice::new(tail); } else if remaining_len != 0 { panic!("advance past the end of the slice vector"); } self } } impl<'a, 'b> Deref for IoBufs<'a, 'b> { type Target = [IoSlice<'a>]; fn deref(&self) -> &[IoSlice<'a>] { self.0 } } tokio-1.43.1/tests/support/leaked_buffers.rs000064400000000000000000000014661046102023000172330ustar 00000000000000/// Can create buffers of arbitrary lifetime. /// Frees created buffers when dropped. /// /// This struct is of course unsafe and the fact that /// it must outlive the created slices has to be ensured by /// the programmer. /// /// Used at certain test scenarios as a safer version of /// Vec::leak, to satisfy the address sanitizer. pub struct LeakedBuffers { leaked_vecs: Vec>, } impl LeakedBuffers { pub fn new() -> Self { Self { leaked_vecs: vec![], } } pub unsafe fn create<'a>(&mut self, size: usize) -> &'a mut [u8] { let new_mem = vec![0u8; size].into_boxed_slice(); self.leaked_vecs.push(new_mem); let new_mem = self.leaked_vecs.last_mut().unwrap(); std::slice::from_raw_parts_mut(new_mem.as_mut_ptr(), new_mem.len()) } } tokio-1.43.1/tests/support/mpsc_stream.rs000064400000000000000000000021051046102023000165760ustar 00000000000000#![allow(dead_code)] use std::pin::Pin; use std::task::{Context, Poll}; use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver, UnboundedSender}; use tokio_stream::Stream; struct UnboundedStream { recv: UnboundedReceiver, } impl Stream for UnboundedStream { type Item = T; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::into_inner(self).recv.poll_recv(cx) } } pub fn unbounded_channel_stream() -> (UnboundedSender, impl Stream) { let (tx, rx) = mpsc::unbounded_channel(); let stream = UnboundedStream { recv: rx }; (tx, stream) } struct BoundedStream { recv: Receiver, } impl Stream for BoundedStream { type Item = T; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::into_inner(self).recv.poll_recv(cx) } } pub fn channel_stream(size: usize) -> (Sender, impl Stream) { let (tx, rx) = mpsc::channel(size); let stream = BoundedStream { recv: rx }; (tx, stream) } tokio-1.43.1/tests/support/panic.rs000064400000000000000000000020571046102023000153610ustar 00000000000000use std::panic; use std::sync::{Arc, Mutex}; pub fn test_panic(func: Func) -> Option { static PANIC_MUTEX: Mutex<()> = Mutex::new(()); { let _guard = PANIC_MUTEX.lock(); let panic_file: Arc>> = Arc::new(Mutex::new(None)); let prev_hook = panic::take_hook(); { let panic_file = panic_file.clone(); panic::set_hook(Box::new(move |panic_info| { let panic_location = panic_info.location().unwrap(); panic_file .lock() .unwrap() .clone_from(&Some(panic_location.file().to_string())); })); } let result = panic::catch_unwind(func); // Return to the previously set panic hook (maybe default) so that we get nice error // messages in the tests. panic::set_hook(prev_hook); if result.is_err() { panic_file.lock().unwrap().clone() } else { None } } } tokio-1.43.1/tests/support/signal.rs000064400000000000000000000005211046102023000155360ustar 00000000000000pub fn send_signal(signal: libc::c_int) { use libc::{getpid, kill}; unsafe { let pid = getpid(); assert_eq!( kill(pid, signal), 0, "kill(pid = {}, {}) failed with error: {}", pid, signal, std::io::Error::last_os_error(), ); } } tokio-1.43.1/tests/sync_barrier.rs000064400000000000000000000042521046102023000152340ustar 00000000000000#![allow(clippy::unnecessary_operation)] #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::Barrier; use tokio_test::task::spawn; use tokio_test::{assert_pending, assert_ready}; struct IsSend(T); #[test] fn barrier_future_is_send() { let b = Barrier::new(0); IsSend(b.wait()); } #[test] fn zero_does_not_block() { let b = Barrier::new(0); { let mut w = spawn(b.wait()); let wr = assert_ready!(w.poll()); assert!(wr.is_leader()); } { let mut w = spawn(b.wait()); let wr = assert_ready!(w.poll()); assert!(wr.is_leader()); } } #[test] fn single() { let b = Barrier::new(1); { let mut w = spawn(b.wait()); let wr = assert_ready!(w.poll()); assert!(wr.is_leader()); } { let mut w = spawn(b.wait()); let wr = assert_ready!(w.poll()); assert!(wr.is_leader()); } { let mut w = spawn(b.wait()); let wr = assert_ready!(w.poll()); assert!(wr.is_leader()); } } #[test] fn tango() { let b = Barrier::new(2); let mut w1 = spawn(b.wait()); assert_pending!(w1.poll()); let mut w2 = spawn(b.wait()); let wr2 = assert_ready!(w2.poll()); let wr1 = assert_ready!(w1.poll()); assert!(wr1.is_leader() || wr2.is_leader()); assert!(!(wr1.is_leader() && wr2.is_leader())); } #[test] fn lots() { let b = Barrier::new(100); for _ in 0..10 { let mut wait = Vec::new(); for _ in 0..99 { let mut w = spawn(b.wait()); assert_pending!(w.poll()); wait.push(w); } for w in &mut wait { assert_pending!(w.poll()); } // pass the barrier let mut w = spawn(b.wait()); let mut found_leader = assert_ready!(w.poll()).is_leader(); for mut w in wait { let wr = assert_ready!(w.poll()); if wr.is_leader() { assert!(!found_leader); found_leader = true; } } assert!(found_leader); } } tokio-1.43.1/tests/sync_broadcast.rs000064400000000000000000000345641046102023000155610ustar 00000000000000#![allow(clippy::cognitive_complexity)] #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::broadcast; use tokio_test::task; use tokio_test::{ assert_err, assert_ok, assert_pending, assert_ready, assert_ready_err, assert_ready_ok, }; use std::sync::Arc; macro_rules! assert_recv { ($e:expr) => { match $e.try_recv() { Ok(value) => value, Err(e) => panic!("expected recv; got = {:?}", e), } }; } macro_rules! assert_empty { ($e:expr) => { match $e.try_recv() { Ok(value) => panic!("expected empty; got = {:?}", value), Err(broadcast::error::TryRecvError::Empty) => {} Err(e) => panic!("expected empty; got = {:?}", e), } }; } macro_rules! assert_lagged { ($e:expr, $n:expr) => { match assert_err!($e) { broadcast::error::TryRecvError::Lagged(n) => { assert_eq!(n, $n); } _ => panic!("did not lag"), } }; } macro_rules! assert_closed { ($e:expr) => { match assert_err!($e) { broadcast::error::TryRecvError::Closed => {} _ => panic!("is not closed"), } }; } #[allow(unused)] trait AssertSend: Send + Sync {} impl AssertSend for broadcast::Sender {} impl AssertSend for broadcast::Receiver {} #[test] fn send_try_recv_bounded() { let (tx, mut rx) = broadcast::channel(16); assert_empty!(rx); let n = assert_ok!(tx.send("hello")); assert_eq!(n, 1); let val = assert_recv!(rx); assert_eq!(val, "hello"); assert_empty!(rx); } #[test] fn send_two_recv() { let (tx, mut rx1) = broadcast::channel(16); let mut rx2 = tx.subscribe(); assert_empty!(rx1); assert_empty!(rx2); let n = assert_ok!(tx.send("hello")); assert_eq!(n, 2); let val = assert_recv!(rx1); assert_eq!(val, "hello"); let val = assert_recv!(rx2); assert_eq!(val, "hello"); assert_empty!(rx1); assert_empty!(rx2); } #[test] fn send_recv_bounded() { let (tx, mut rx) = broadcast::channel(16); let mut recv = task::spawn(rx.recv()); assert_pending!(recv.poll()); assert_ok!(tx.send("hello")); assert!(recv.is_woken()); let val = assert_ready_ok!(recv.poll()); assert_eq!(val, "hello"); } #[test] fn send_two_recv_bounded() { let (tx, mut rx1) = broadcast::channel(16); let mut rx2 = tx.subscribe(); let mut recv1 = task::spawn(rx1.recv()); let mut recv2 = task::spawn(rx2.recv()); assert_pending!(recv1.poll()); assert_pending!(recv2.poll()); assert_ok!(tx.send("hello")); assert!(recv1.is_woken()); assert!(recv2.is_woken()); let val1 = assert_ready_ok!(recv1.poll()); let val2 = assert_ready_ok!(recv2.poll()); assert_eq!(val1, "hello"); assert_eq!(val2, "hello"); drop((recv1, recv2)); let mut recv1 = task::spawn(rx1.recv()); let mut recv2 = task::spawn(rx2.recv()); assert_pending!(recv1.poll()); assert_ok!(tx.send("world")); assert!(recv1.is_woken()); assert!(!recv2.is_woken()); let val1 = assert_ready_ok!(recv1.poll()); let val2 = assert_ready_ok!(recv2.poll()); assert_eq!(val1, "world"); assert_eq!(val2, "world"); } #[test] fn change_tasks() { let (tx, mut rx) = broadcast::channel(1); let mut recv = Box::pin(rx.recv()); let mut task1 = task::spawn(&mut recv); assert_pending!(task1.poll()); let mut task2 = task::spawn(&mut recv); assert_pending!(task2.poll()); tx.send("hello").unwrap(); assert!(task2.is_woken()); } #[test] fn send_slow_rx() { let (tx, mut rx1) = broadcast::channel(16); let mut rx2 = tx.subscribe(); { let mut recv2 = task::spawn(rx2.recv()); { let mut recv1 = task::spawn(rx1.recv()); assert_pending!(recv1.poll()); assert_pending!(recv2.poll()); assert_ok!(tx.send("one")); assert!(recv1.is_woken()); assert!(recv2.is_woken()); assert_ok!(tx.send("two")); let val = assert_ready_ok!(recv1.poll()); assert_eq!(val, "one"); } let val = assert_ready_ok!(task::spawn(rx1.recv()).poll()); assert_eq!(val, "two"); let mut recv1 = task::spawn(rx1.recv()); assert_pending!(recv1.poll()); assert_ok!(tx.send("three")); assert!(recv1.is_woken()); let val = assert_ready_ok!(recv1.poll()); assert_eq!(val, "three"); let val = assert_ready_ok!(recv2.poll()); assert_eq!(val, "one"); } let val = assert_recv!(rx2); assert_eq!(val, "two"); let val = assert_recv!(rx2); assert_eq!(val, "three"); } #[test] fn drop_rx_while_values_remain() { let (tx, mut rx1) = broadcast::channel(16); let mut rx2 = tx.subscribe(); assert_ok!(tx.send("one")); assert_ok!(tx.send("two")); assert_recv!(rx1); assert_recv!(rx2); drop(rx2); drop(rx1); } #[test] fn lagging_rx() { let (tx, mut rx1) = broadcast::channel(2); let mut rx2 = tx.subscribe(); assert_ok!(tx.send("one")); assert_ok!(tx.send("two")); assert_eq!("one", assert_recv!(rx1)); assert_ok!(tx.send("three")); // Lagged too far let x = dbg!(rx2.try_recv()); assert_lagged!(x, 1); // Calling again gets the next value assert_eq!("two", assert_recv!(rx2)); assert_eq!("two", assert_recv!(rx1)); assert_eq!("three", assert_recv!(rx1)); assert_ok!(tx.send("four")); assert_ok!(tx.send("five")); assert_lagged!(rx2.try_recv(), 1); assert_ok!(tx.send("six")); assert_lagged!(rx2.try_recv(), 1); } #[test] fn send_no_rx() { let (tx, _) = broadcast::channel(16); assert_err!(tx.send("hello")); let mut rx = tx.subscribe(); assert_ok!(tx.send("world")); let val = assert_recv!(rx); assert_eq!("world", val); } #[test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn zero_capacity() { broadcast::channel::<()>(0); } #[test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn capacity_too_big() { broadcast::channel::<()>(1 + (usize::MAX >> 1)); } #[test] #[cfg(panic = "unwind")] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn panic_in_clone() { use std::panic::{self, AssertUnwindSafe}; #[derive(Eq, PartialEq, Debug)] struct MyVal(usize); impl Clone for MyVal { fn clone(&self) -> MyVal { assert_ne!(0, self.0); MyVal(self.0) } } let (tx, mut rx) = broadcast::channel(16); assert_ok!(tx.send(MyVal(0))); assert_ok!(tx.send(MyVal(1))); let res = panic::catch_unwind(AssertUnwindSafe(|| { let _ = rx.try_recv(); })); assert_err!(res); let val = assert_recv!(rx); assert_eq!(val, MyVal(1)); } #[test] fn dropping_tx_notifies_rx() { let (tx, mut rx1) = broadcast::channel::<()>(16); let mut rx2 = tx.subscribe(); let tx2 = tx.clone(); let mut recv1 = task::spawn(rx1.recv()); let mut recv2 = task::spawn(rx2.recv()); assert_pending!(recv1.poll()); assert_pending!(recv2.poll()); drop(tx); assert_pending!(recv1.poll()); assert_pending!(recv2.poll()); drop(tx2); assert!(recv1.is_woken()); assert!(recv2.is_woken()); let err = assert_ready_err!(recv1.poll()); assert!(is_closed(err)); let err = assert_ready_err!(recv2.poll()); assert!(is_closed(err)); } #[test] fn unconsumed_messages_are_dropped() { let (tx, rx) = broadcast::channel(16); let msg = Arc::new(()); assert_ok!(tx.send(msg.clone())); assert_eq!(2, Arc::strong_count(&msg)); drop(rx); assert_eq!(1, Arc::strong_count(&msg)); } #[test] fn single_capacity_recvs() { let (tx, mut rx) = broadcast::channel(1); assert_ok!(tx.send(1)); assert_eq!(assert_recv!(rx), 1); assert_empty!(rx); } #[test] fn single_capacity_recvs_after_drop_1() { let (tx, mut rx) = broadcast::channel(1); assert_ok!(tx.send(1)); drop(tx); assert_eq!(assert_recv!(rx), 1); assert_closed!(rx.try_recv()); } #[test] fn single_capacity_recvs_after_drop_2() { let (tx, mut rx) = broadcast::channel(1); assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); drop(tx); assert_lagged!(rx.try_recv(), 1); assert_eq!(assert_recv!(rx), 2); assert_closed!(rx.try_recv()); } #[test] fn dropping_sender_does_not_overwrite() { let (tx, mut rx) = broadcast::channel(2); assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); drop(tx); assert_eq!(assert_recv!(rx), 1); assert_eq!(assert_recv!(rx), 2); assert_closed!(rx.try_recv()); } #[test] fn lagging_receiver_recovers_after_wrap_closed_1() { let (tx, mut rx) = broadcast::channel(2); assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); assert_ok!(tx.send(3)); drop(tx); assert_lagged!(rx.try_recv(), 1); assert_eq!(assert_recv!(rx), 2); assert_eq!(assert_recv!(rx), 3); assert_closed!(rx.try_recv()); } #[test] fn lagging_receiver_recovers_after_wrap_closed_2() { let (tx, mut rx) = broadcast::channel(2); assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); assert_ok!(tx.send(3)); assert_ok!(tx.send(4)); drop(tx); assert_lagged!(rx.try_recv(), 2); assert_eq!(assert_recv!(rx), 3); assert_eq!(assert_recv!(rx), 4); assert_closed!(rx.try_recv()); } #[test] fn lagging_receiver_recovers_after_wrap_open() { let (tx, mut rx) = broadcast::channel(2); assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); assert_ok!(tx.send(3)); assert_lagged!(rx.try_recv(), 1); assert_eq!(assert_recv!(rx), 2); assert_eq!(assert_recv!(rx), 3); assert_empty!(rx); } #[test] fn receiver_len_with_lagged() { let (tx, mut rx) = broadcast::channel(3); tx.send(10).unwrap(); tx.send(20).unwrap(); tx.send(30).unwrap(); tx.send(40).unwrap(); assert_eq!(rx.len(), 4); assert_eq!(assert_recv!(rx), 10); tx.send(50).unwrap(); tx.send(60).unwrap(); assert_eq!(rx.len(), 5); assert_lagged!(rx.try_recv(), 1); } fn is_closed(err: broadcast::error::RecvError) -> bool { matches!(err, broadcast::error::RecvError::Closed) } #[test] fn resubscribe_points_to_tail() { let (tx, mut rx) = broadcast::channel(3); tx.send(1).unwrap(); let mut rx_resub = rx.resubscribe(); // verify we're one behind at the start assert_empty!(rx_resub); assert_eq!(assert_recv!(rx), 1); // verify we do not affect rx tx.send(2).unwrap(); assert_eq!(assert_recv!(rx_resub), 2); tx.send(3).unwrap(); assert_eq!(assert_recv!(rx), 2); assert_eq!(assert_recv!(rx), 3); assert_empty!(rx); assert_eq!(assert_recv!(rx_resub), 3); assert_empty!(rx_resub); } #[test] fn resubscribe_lagged() { let (tx, mut rx) = broadcast::channel(1); tx.send(1).unwrap(); tx.send(2).unwrap(); let mut rx_resub = rx.resubscribe(); assert_lagged!(rx.try_recv(), 1); assert_empty!(rx_resub); assert_eq!(assert_recv!(rx), 2); assert_empty!(rx); assert_empty!(rx_resub); } #[test] fn resubscribe_to_closed_channel() { let (tx, rx) = tokio::sync::broadcast::channel::(2); drop(tx); let mut rx_resub = rx.resubscribe(); assert_closed!(rx_resub.try_recv()); } #[test] fn sender_len() { let (tx, mut rx1) = broadcast::channel(4); let mut rx2 = tx.subscribe(); assert_eq!(tx.len(), 0); assert!(tx.is_empty()); tx.send(1).unwrap(); tx.send(2).unwrap(); tx.send(3).unwrap(); assert_eq!(tx.len(), 3); assert!(!tx.is_empty()); assert_recv!(rx1); assert_recv!(rx1); assert_eq!(tx.len(), 3); assert!(!tx.is_empty()); assert_recv!(rx2); assert_eq!(tx.len(), 2); assert!(!tx.is_empty()); tx.send(4).unwrap(); tx.send(5).unwrap(); tx.send(6).unwrap(); assert_eq!(tx.len(), 4); assert!(!tx.is_empty()); } #[test] #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] fn sender_len_random() { use rand::Rng; let (tx, mut rx1) = broadcast::channel(16); let mut rx2 = tx.subscribe(); for _ in 0..1000 { match rand::thread_rng().gen_range(0..4) { 0 => { let _ = rx1.try_recv(); } 1 => { let _ = rx2.try_recv(); } _ => { tx.send(0).unwrap(); } } let expected_len = usize::min(usize::max(rx1.len(), rx2.len()), 16); assert_eq!(tx.len(), expected_len); } } #[test] fn send_in_waker_drop() { use futures::task::ArcWake; use std::future::Future; use std::task::Context; struct SendOnDrop(broadcast::Sender<()>); impl Drop for SendOnDrop { fn drop(&mut self) { let _ = self.0.send(()); } } impl ArcWake for SendOnDrop { fn wake_by_ref(_arc_self: &Arc) {} } // Test if there is no deadlock when replacing the old waker. let (tx, mut rx) = broadcast::channel(16); let mut fut = Box::pin(async { let _ = rx.recv().await; }); // Store our special waker in the receiving future. let waker = futures::task::waker(Arc::new(SendOnDrop(tx))); let mut cx = Context::from_waker(&waker); assert!(fut.as_mut().poll(&mut cx).is_pending()); drop(waker); // Second poll shouldn't deadlock. let mut cx = Context::from_waker(futures::task::noop_waker_ref()); let _ = fut.as_mut().poll(&mut cx); // Test if there is no deadlock when calling waker.wake(). let (tx, mut rx) = broadcast::channel(16); let mut fut = Box::pin(async { let _ = rx.recv().await; }); // Store our special waker in the receiving future. let waker = futures::task::waker(Arc::new(SendOnDrop(tx.clone()))); let mut cx = Context::from_waker(&waker); assert!(fut.as_mut().poll(&mut cx).is_pending()); drop(waker); // Shouldn't deadlock. let _ = tx.send(()); } #[tokio::test] async fn receiver_recv_is_cooperative() { let (tx, mut rx) = broadcast::channel(8); tokio::select! { biased; _ = async { loop { assert!(tx.send(()).is_ok()); assert!(rx.recv().await.is_ok()); } } => {}, _ = tokio::task::yield_now() => {}, } } tokio-1.43.1/tests/sync_errors.rs000064400000000000000000000011731046102023000151210ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; fn is_error() {} #[test] fn mpsc_error_bound() { use tokio::sync::mpsc::error; is_error::>(); is_error::>(); } #[test] fn oneshot_error_bound() { use tokio::sync::oneshot::error; is_error::(); is_error::(); } #[test] fn watch_error_bound() { use tokio::sync::watch::error; is_error::>(); } tokio-1.43.1/tests/sync_mpsc.rs000064400000000000000000001143361046102023000145550ustar 00000000000000#![allow(clippy::redundant_clone)] #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use std::fmt; use std::panic; use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::mpsc::error::{TryRecvError, TrySendError}; use tokio_test::*; #[cfg(not(target_family = "wasm"))] mod support { pub(crate) mod mpsc_stream; } #[allow(unused)] trait AssertRefUnwindSafe: panic::RefUnwindSafe {} impl AssertRefUnwindSafe for mpsc::OwnedPermit {} impl<'a, T> AssertRefUnwindSafe for mpsc::Permit<'a, T> {} impl<'a, T> AssertRefUnwindSafe for mpsc::PermitIterator<'a, T> {} impl AssertRefUnwindSafe for mpsc::Receiver {} impl AssertRefUnwindSafe for mpsc::Sender {} impl AssertRefUnwindSafe for mpsc::UnboundedReceiver {} impl AssertRefUnwindSafe for mpsc::UnboundedSender {} impl AssertRefUnwindSafe for mpsc::WeakSender {} impl AssertRefUnwindSafe for mpsc::WeakUnboundedSender {} #[allow(unused)] trait AssertUnwindSafe: panic::UnwindSafe {} impl AssertUnwindSafe for mpsc::OwnedPermit {} impl<'a, T> AssertUnwindSafe for mpsc::Permit<'a, T> {} impl<'a, T> AssertUnwindSafe for mpsc::PermitIterator<'a, T> {} impl AssertUnwindSafe for mpsc::Receiver {} impl AssertUnwindSafe for mpsc::Sender {} impl AssertUnwindSafe for mpsc::UnboundedReceiver {} impl AssertUnwindSafe for mpsc::UnboundedSender {} impl AssertUnwindSafe for mpsc::WeakSender {} impl AssertUnwindSafe for mpsc::WeakUnboundedSender {} #[maybe_tokio_test] async fn send_recv_with_buffer() { let (tx, mut rx) = mpsc::channel::(16); // Using poll_ready / try_send // let permit assert_ready_ok!(tx.reserve()); let permit = tx.reserve().await.unwrap(); permit.send(1); // Without poll_ready tx.try_send(2).unwrap(); drop(tx); let val = rx.recv().await; assert_eq!(val, Some(1)); let val = rx.recv().await; assert_eq!(val, Some(2)); let val = rx.recv().await; assert!(val.is_none()); } #[tokio::test] #[cfg(feature = "full")] async fn reserve_disarm() { let (tx, mut rx) = mpsc::channel::(2); let tx1 = tx.clone(); let tx2 = tx.clone(); let tx3 = tx.clone(); let tx4 = tx; // We should be able to `poll_ready` two handles without problem let permit1 = assert_ok!(tx1.reserve().await); let permit2 = assert_ok!(tx2.reserve().await); // But a third should not be ready let mut r3 = tokio_test::task::spawn(tx3.reserve()); assert_pending!(r3.poll()); let mut r4 = tokio_test::task::spawn(tx4.reserve()); assert_pending!(r4.poll()); // Using one of the reserved slots should allow a new handle to become ready permit1.send(1); // We also need to receive for the slot to be free assert!(!r3.is_woken()); rx.recv().await.unwrap(); // Now there's a free slot! assert!(r3.is_woken()); assert!(!r4.is_woken()); // Dropping a permit should also open up a slot drop(permit2); assert!(r4.is_woken()); let mut r1 = tokio_test::task::spawn(tx1.reserve()); assert_pending!(r1.poll()); } #[tokio::test] #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads async fn send_recv_stream_with_buffer() { use tokio_stream::StreamExt; let (tx, rx) = support::mpsc_stream::channel_stream::(16); let mut rx = Box::pin(rx); tokio::spawn(async move { assert_ok!(tx.send(1).await); assert_ok!(tx.send(2).await); }); assert_eq!(Some(1), rx.next().await); assert_eq!(Some(2), rx.next().await); assert_eq!(None, rx.next().await); } #[tokio::test] #[cfg(feature = "full")] async fn async_send_recv_with_buffer() { let (tx, mut rx) = mpsc::channel(16); tokio::spawn(async move { assert_ok!(tx.send(1).await); assert_ok!(tx.send(2).await); }); assert_eq!(Some(1), rx.recv().await); assert_eq!(Some(2), rx.recv().await); assert_eq!(None, rx.recv().await); } #[tokio::test] #[cfg(feature = "full")] async fn async_send_recv_many_with_buffer() { let (tx, mut rx) = mpsc::channel(2); let mut buffer = Vec::::with_capacity(3); // With `limit=0` does not sleep, returns immediately assert_eq!(0, rx.recv_many(&mut buffer, 0).await); let handle = tokio::spawn(async move { assert_ok!(tx.send(1).await); assert_ok!(tx.send(2).await); assert_ok!(tx.send(7).await); assert_ok!(tx.send(0).await); }); let limit = 3; let mut recv_count = 0usize; while recv_count < 4 { recv_count += rx.recv_many(&mut buffer, limit).await; assert_eq!(buffer.len(), recv_count); } assert_eq!(vec![1, 2, 7, 0], buffer); assert_eq!(0, rx.recv_many(&mut buffer, limit).await); handle.await.unwrap(); } #[tokio::test] #[cfg(feature = "full")] async fn start_send_past_cap() { use std::future::Future; let mut t1 = tokio_test::task::spawn(()); let (tx1, mut rx) = mpsc::channel(1); let tx2 = tx1.clone(); assert_ok!(tx1.try_send(())); let mut r1 = Box::pin(tx1.reserve()); t1.enter(|cx, _| assert_pending!(r1.as_mut().poll(cx))); { let mut r2 = tokio_test::task::spawn(tx2.reserve()); assert_pending!(r2.poll()); drop(r1); assert!(rx.recv().await.is_some()); assert!(r2.is_woken()); assert!(!t1.is_woken()); } drop(tx1); drop(tx2); assert!(rx.recv().await.is_none()); } #[test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn buffer_gteq_one() { mpsc::channel::(0); } #[maybe_tokio_test] async fn send_recv_unbounded() { let (tx, mut rx) = mpsc::unbounded_channel::(); // Using `try_send` assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); assert_eq!(rx.recv().await, Some(1)); assert_eq!(rx.recv().await, Some(2)); drop(tx); assert!(rx.recv().await.is_none()); } #[maybe_tokio_test] async fn send_recv_many_unbounded() { let (tx, mut rx) = mpsc::unbounded_channel::(); let mut buffer: Vec = Vec::new(); // With `limit=0` does not sleep, returns immediately rx.recv_many(&mut buffer, 0).await; assert_eq!(0, buffer.len()); assert_ok!(tx.send(7)); assert_ok!(tx.send(13)); assert_ok!(tx.send(100)); assert_ok!(tx.send(1002)); rx.recv_many(&mut buffer, 0).await; assert_eq!(0, buffer.len()); let mut count = 0; while count < 4 { count += rx.recv_many(&mut buffer, 1).await; } assert_eq!(count, 4); assert_eq!(vec![7, 13, 100, 1002], buffer); let final_capacity = buffer.capacity(); assert!(final_capacity > 0); buffer.clear(); assert_ok!(tx.send(5)); assert_ok!(tx.send(6)); assert_ok!(tx.send(7)); assert_ok!(tx.send(2)); // Re-use existing capacity count = rx.recv_many(&mut buffer, 32).await; assert_eq!(final_capacity, buffer.capacity()); assert_eq!(count, 4); assert_eq!(vec![5, 6, 7, 2], buffer); drop(tx); // recv_many will immediately return zero if the channel // is closed and no more messages are waiting assert_eq!(0, rx.recv_many(&mut buffer, 4).await); assert!(rx.recv().await.is_none()); } #[tokio::test] #[cfg(feature = "full")] async fn send_recv_many_bounded_capacity() { let mut buffer: Vec = Vec::with_capacity(9); let limit = buffer.capacity(); let (tx, mut rx) = mpsc::channel(100); let mut expected: Vec = (0..limit) .map(|x: usize| format!("{x}")) .collect::>(); for x in expected.clone() { tx.send(x).await.unwrap() } tx.send("one more".to_string()).await.unwrap(); // Here `recv_many` receives all but the last value; // the initial capacity is adequate, so the buffer does // not increase in side. assert_eq!(buffer.capacity(), rx.recv_many(&mut buffer, limit).await); assert_eq!(expected, buffer); assert_eq!(limit, buffer.capacity()); // Receive up more values: assert_eq!(1, rx.recv_many(&mut buffer, limit).await); assert!(buffer.capacity() > limit); expected.push("one more".to_string()); assert_eq!(expected, buffer); tokio::spawn(async move { tx.send("final".to_string()).await.unwrap(); }); // 'tx' is dropped, but `recv_many` is guaranteed not // to return 0 as the channel has outstanding permits assert_eq!(1, rx.recv_many(&mut buffer, limit).await); expected.push("final".to_string()); assert_eq!(expected, buffer); // The channel is now closed and `recv_many` returns 0. assert_eq!(0, rx.recv_many(&mut buffer, limit).await); assert_eq!(expected, buffer); } #[tokio::test] #[cfg(feature = "full")] async fn send_recv_many_unbounded_capacity() { let mut buffer: Vec = Vec::with_capacity(9); // capacity >= 9 let limit = buffer.capacity(); let (tx, mut rx) = mpsc::unbounded_channel(); let mut expected: Vec = (0..limit) .map(|x: usize| format!("{x}")) .collect::>(); for x in expected.clone() { tx.send(x).unwrap() } tx.send("one more".to_string()).unwrap(); // Here `recv_many` receives all but the last value; // the initial capacity is adequate, so the buffer does // not increase in side. assert_eq!(buffer.capacity(), rx.recv_many(&mut buffer, limit).await); assert_eq!(expected, buffer); assert_eq!(limit, buffer.capacity()); // Receive up more values: assert_eq!(1, rx.recv_many(&mut buffer, limit).await); assert!(buffer.capacity() > limit); expected.push("one more".to_string()); assert_eq!(expected, buffer); tokio::spawn(async move { tx.send("final".to_string()).unwrap(); }); // 'tx' is dropped, but `recv_many` is guaranteed not // to return 0 as the channel has outstanding permits assert_eq!(1, rx.recv_many(&mut buffer, limit).await); expected.push("final".to_string()); assert_eq!(expected, buffer); // The channel is now closed and `recv_many` returns 0. assert_eq!(0, rx.recv_many(&mut buffer, limit).await); assert_eq!(expected, buffer); } #[tokio::test] #[cfg(feature = "full")] async fn async_send_recv_unbounded() { let (tx, mut rx) = mpsc::unbounded_channel(); tokio::spawn(async move { assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); }); assert_eq!(Some(1), rx.recv().await); assert_eq!(Some(2), rx.recv().await); assert_eq!(None, rx.recv().await); } #[tokio::test] #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads async fn send_recv_stream_unbounded() { use tokio_stream::StreamExt; let (tx, rx) = support::mpsc_stream::unbounded_channel_stream::(); let mut rx = Box::pin(rx); tokio::spawn(async move { assert_ok!(tx.send(1)); assert_ok!(tx.send(2)); }); assert_eq!(Some(1), rx.next().await); assert_eq!(Some(2), rx.next().await); assert_eq!(None, rx.next().await); } #[maybe_tokio_test] async fn no_t_bounds_buffer() { struct NoImpls; let (tx, mut rx) = mpsc::channel(100); // sender should be Debug even though T isn't Debug is_debug(&tx); // same with Receiver is_debug(&rx); // and sender should be Clone even though T isn't Clone assert!(tx.clone().try_send(NoImpls).is_ok()); assert!(rx.recv().await.is_some()); } #[maybe_tokio_test] async fn no_t_bounds_unbounded() { struct NoImpls; let (tx, mut rx) = mpsc::unbounded_channel(); // sender should be Debug even though T isn't Debug is_debug(&tx); // same with Receiver is_debug(&rx); // and sender should be Clone even though T isn't Clone assert!(tx.clone().send(NoImpls).is_ok()); assert!(rx.recv().await.is_some()); } #[tokio::test] #[cfg(feature = "full")] async fn send_recv_buffer_limited() { let (tx, mut rx) = mpsc::channel::(1); // Reserve capacity let p1 = assert_ok!(tx.reserve().await); // Send first message p1.send(1); // Not ready let mut p2 = tokio_test::task::spawn(tx.reserve()); assert_pending!(p2.poll()); // Take the value assert!(rx.recv().await.is_some()); // Notified assert!(p2.is_woken()); // Trying to send fails assert_err!(tx.try_send(1337)); // Send second let permit = assert_ready_ok!(p2.poll()); permit.send(2); assert!(rx.recv().await.is_some()); } #[maybe_tokio_test] async fn recv_close_gets_none_idle() { let (tx, mut rx) = mpsc::channel::(10); rx.close(); assert!(rx.recv().await.is_none()); assert_err!(tx.send(1).await); } #[tokio::test] #[cfg(feature = "full")] async fn recv_close_gets_none_reserved() { let (tx1, mut rx) = mpsc::channel::(1); let tx2 = tx1.clone(); let permit1 = assert_ok!(tx1.reserve().await); let mut permit2 = tokio_test::task::spawn(tx2.reserve()); assert_pending!(permit2.poll()); rx.close(); assert!(permit2.is_woken()); assert_ready_err!(permit2.poll()); { let mut recv = tokio_test::task::spawn(rx.recv()); assert_pending!(recv.poll()); permit1.send(123); assert!(recv.is_woken()); let v = assert_ready!(recv.poll()); assert_eq!(v, Some(123)); } assert!(rx.recv().await.is_none()); } #[maybe_tokio_test] async fn tx_close_gets_none() { let (_, mut rx) = mpsc::channel::(10); assert!(rx.recv().await.is_none()); } #[maybe_tokio_test] async fn try_send_fail() { let (tx, mut rx) = mpsc::channel(1); tx.try_send("hello").unwrap(); // This should fail match assert_err!(tx.try_send("fail")) { TrySendError::Full(..) => {} _ => panic!(), } assert_eq!(rx.recv().await, Some("hello")); assert_ok!(tx.try_send("goodbye")); drop(tx); assert_eq!(rx.recv().await, Some("goodbye")); assert!(rx.recv().await.is_none()); } #[maybe_tokio_test] async fn try_send_fail_with_try_recv() { let (tx, mut rx) = mpsc::channel(1); tx.try_send("hello").unwrap(); // This should fail match assert_err!(tx.try_send("fail")) { TrySendError::Full(..) => {} _ => panic!(), } assert_eq!(rx.try_recv(), Ok("hello")); assert_ok!(tx.try_send("goodbye")); drop(tx); assert_eq!(rx.try_recv(), Ok("goodbye")); assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); } #[maybe_tokio_test] async fn reserve_many_above_cap() { const MAX_PERMITS: usize = tokio::sync::Semaphore::MAX_PERMITS; let (tx, _rx) = mpsc::channel::<()>(1); assert_err!(tx.reserve_many(2).await); assert_err!(tx.reserve_many(MAX_PERMITS + 1).await); assert_err!(tx.reserve_many(usize::MAX).await); } #[test] fn try_reserve_many_zero() { let (tx, rx) = mpsc::channel::<()>(1); // Succeeds when not closed. assert!(assert_ok!(tx.try_reserve_many(0)).next().is_none()); // Even when channel is full. tx.try_send(()).unwrap(); assert!(assert_ok!(tx.try_reserve_many(0)).next().is_none()); drop(rx); // Closed error when closed. assert_eq!( assert_err!(tx.try_reserve_many(0)), TrySendError::Closed(()) ); } #[maybe_tokio_test] async fn reserve_many_zero() { let (tx, rx) = mpsc::channel::<()>(1); // Succeeds when not closed. assert!(assert_ok!(tx.reserve_many(0).await).next().is_none()); // Even when channel is full. tx.send(()).await.unwrap(); assert!(assert_ok!(tx.reserve_many(0).await).next().is_none()); drop(rx); // Closed error when closed. assert_err!(tx.reserve_many(0).await); } #[maybe_tokio_test] async fn try_reserve_many_edge_cases() { const MAX_PERMITS: usize = tokio::sync::Semaphore::MAX_PERMITS; let (tx, rx) = mpsc::channel::<()>(1); let mut permit = assert_ok!(tx.try_reserve_many(0)); assert!(permit.next().is_none()); let permit = tx.try_reserve_many(MAX_PERMITS + 1); match assert_err!(permit) { TrySendError::Full(..) => {} _ => panic!(), } let permit = tx.try_reserve_many(usize::MAX); match assert_err!(permit) { TrySendError::Full(..) => {} _ => panic!(), } // Dropping the receiver should close the channel drop(rx); assert_err!(tx.reserve_many(0).await); } #[maybe_tokio_test] async fn try_reserve_fails() { let (tx, mut rx) = mpsc::channel(1); let permit = tx.try_reserve().unwrap(); // This should fail match assert_err!(tx.try_reserve()) { TrySendError::Full(()) => {} _ => panic!(), } permit.send("foo"); assert_eq!(rx.recv().await, Some("foo")); // Dropping permit releases the slot. let permit = tx.try_reserve().unwrap(); drop(permit); let _permit = tx.try_reserve().unwrap(); } #[maybe_tokio_test] async fn reserve_many_and_send() { let (tx, mut rx) = mpsc::channel(100); for i in 0..100 { for permit in assert_ok!(tx.reserve_many(i).await) { permit.send("foo"); assert_eq!(rx.recv().await, Some("foo")); } assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } } #[maybe_tokio_test] async fn try_reserve_many_and_send() { let (tx, mut rx) = mpsc::channel(100); for i in 0..100 { for permit in assert_ok!(tx.try_reserve_many(i)) { permit.send("foo"); assert_eq!(rx.recv().await, Some("foo")); } assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } } #[maybe_tokio_test] async fn reserve_many_on_closed_channel() { let (tx, rx) = mpsc::channel::<()>(100); drop(rx); assert_err!(tx.reserve_many(10).await); } #[maybe_tokio_test] async fn try_reserve_many_on_closed_channel() { let (tx, rx) = mpsc::channel::(100); drop(rx); match assert_err!(tx.try_reserve_many(10)) { TrySendError::Closed(()) => {} _ => panic!(), }; } #[maybe_tokio_test] #[cfg_attr(miri, ignore)] // Too slow on miri. async fn try_reserve_many_full() { // Reserve n capacity and send k messages for n in 1..100 { for k in 0..n { let (tx, mut rx) = mpsc::channel::(n); let permits = assert_ok!(tx.try_reserve_many(n)); assert_eq!(permits.len(), n); assert_eq!(tx.capacity(), 0); match assert_err!(tx.try_reserve_many(1)) { TrySendError::Full(..) => {} _ => panic!(), }; for permit in permits.take(k) { permit.send(0); } // We only used k permits on the n reserved assert_eq!(tx.capacity(), n - k); // We can reserve more permits assert_ok!(tx.try_reserve_many(1)); // But not more than the current capacity match assert_err!(tx.try_reserve_many(n - k + 1)) { TrySendError::Full(..) => {} _ => panic!(), }; for _i in 0..k { assert_eq!(rx.recv().await, Some(0)); } // Now that we've received everything, capacity should be back to n assert_eq!(tx.capacity(), n); } } } #[tokio::test] #[cfg(feature = "full")] async fn drop_permit_releases_permit() { // poll_ready reserves capacity, ensure that the capacity is released if tx // is dropped w/o sending a value. let (tx1, _rx) = mpsc::channel::(1); let tx2 = tx1.clone(); let permit = assert_ok!(tx1.reserve().await); let mut reserve2 = tokio_test::task::spawn(tx2.reserve()); assert_pending!(reserve2.poll()); drop(permit); assert!(reserve2.is_woken()); assert_ready_ok!(reserve2.poll()); } #[maybe_tokio_test] async fn drop_permit_iterator_releases_permits() { // poll_ready reserves capacity, ensure that the capacity is released if tx // is dropped w/o sending a value. for n in 1..100 { let (tx1, _rx) = mpsc::channel::(n); let tx2 = tx1.clone(); let permits = assert_ok!(tx1.reserve_many(n).await); let mut reserve2 = tokio_test::task::spawn(tx2.reserve_many(n)); assert_pending!(reserve2.poll()); drop(permits); assert!(reserve2.is_woken()); let permits = assert_ready_ok!(reserve2.poll()); drop(permits); assert_eq!(tx1.capacity(), n); } } #[maybe_tokio_test] async fn dropping_rx_closes_channel() { let (tx, rx) = mpsc::channel(100); let msg = Arc::new(()); assert_ok!(tx.try_send(msg.clone())); drop(rx); assert_err!(tx.reserve().await); assert_err!(tx.reserve_many(10).await); assert_eq!(1, Arc::strong_count(&msg)); } #[test] fn dropping_rx_closes_channel_for_try() { let (tx, rx) = mpsc::channel(100); let msg = Arc::new(()); tx.try_send(msg.clone()).unwrap(); drop(rx); assert!(matches!( tx.try_send(msg.clone()), Err(TrySendError::Closed(_)) )); assert!(matches!(tx.try_reserve(), Err(TrySendError::Closed(_)))); assert!(matches!( tx.try_reserve_owned(), Err(TrySendError::Closed(_)) )); assert_eq!(1, Arc::strong_count(&msg)); } #[test] fn unconsumed_messages_are_dropped() { let msg = Arc::new(()); let (tx, rx) = mpsc::channel(100); tx.try_send(msg.clone()).unwrap(); assert_eq!(2, Arc::strong_count(&msg)); drop((tx, rx)); assert_eq!(1, Arc::strong_count(&msg)); } #[test] #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads fn blocking_recv() { let (tx, mut rx) = mpsc::channel::(1); let sync_code = std::thread::spawn(move || { assert_eq!(Some(10), rx.blocking_recv()); }); tokio::runtime::Runtime::new() .unwrap() .block_on(async move { let _ = tx.send(10).await; }); sync_code.join().unwrap() } #[tokio::test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding async fn blocking_recv_async() { let (_tx, mut rx) = mpsc::channel::<()>(1); let _ = rx.blocking_recv(); } #[test] #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads fn blocking_send() { let (tx, mut rx) = mpsc::channel::(1); let sync_code = std::thread::spawn(move || { tx.blocking_send(10).unwrap(); }); tokio::runtime::Runtime::new() .unwrap() .block_on(async move { assert_eq!(Some(10), rx.recv().await); }); sync_code.join().unwrap() } #[tokio::test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding async fn blocking_send_async() { let (tx, _rx) = mpsc::channel::<()>(1); let _ = tx.blocking_send(()); } #[tokio::test] #[cfg(feature = "full")] async fn ready_close_cancel_bounded() { let (tx, mut rx) = mpsc::channel::<()>(100); let _tx2 = tx.clone(); let permit = assert_ok!(tx.reserve().await); rx.close(); let mut recv = tokio_test::task::spawn(rx.recv()); assert_pending!(recv.poll()); drop(permit); assert!(recv.is_woken()); let val = assert_ready!(recv.poll()); assert!(val.is_none()); } #[tokio::test] #[cfg(feature = "full")] async fn permit_available_not_acquired_close() { let (tx1, mut rx) = mpsc::channel::<()>(1); let tx2 = tx1.clone(); let permit1 = assert_ok!(tx1.reserve().await); let mut permit2 = tokio_test::task::spawn(tx2.reserve()); assert_pending!(permit2.poll()); rx.close(); drop(permit1); assert!(permit2.is_woken()); drop(permit2); assert!(rx.recv().await.is_none()); } #[test] fn try_recv_bounded() { let (tx, mut rx) = mpsc::channel(5); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); assert!(tx.try_send("hello").is_err()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); assert_eq!(Ok("hello"), rx.try_recv()); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); assert!(tx.try_send("hello").is_err()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); tx.try_send("hello").unwrap(); drop(tx); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Ok("hello"), rx.try_recv()); assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); } #[test] fn try_recv_unbounded() { for num in 0..100 { let (tx, mut rx) = mpsc::unbounded_channel(); for i in 0..num { tx.send(i).unwrap(); } for i in 0..num { assert_eq!(rx.try_recv(), Ok(i)); } assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); drop(tx); assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); } } #[test] fn try_recv_close_while_empty_bounded() { let (tx, mut rx) = mpsc::channel::<()>(5); assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); drop(tx); assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); } #[test] fn try_recv_close_while_empty_unbounded() { let (tx, mut rx) = mpsc::unbounded_channel::<()>(); assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); drop(tx); assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); } #[tokio::test(start_paused = true)] #[cfg(feature = "full")] async fn recv_timeout() { use tokio::sync::mpsc::error::SendTimeoutError::{Closed, Timeout}; use tokio::time::Duration; let (tx, rx) = mpsc::channel(5); assert_eq!(tx.send_timeout(10, Duration::from_secs(1)).await, Ok(())); assert_eq!(tx.send_timeout(20, Duration::from_secs(1)).await, Ok(())); assert_eq!(tx.send_timeout(30, Duration::from_secs(1)).await, Ok(())); assert_eq!(tx.send_timeout(40, Duration::from_secs(1)).await, Ok(())); assert_eq!(tx.send_timeout(50, Duration::from_secs(1)).await, Ok(())); assert_eq!( tx.send_timeout(60, Duration::from_secs(1)).await, Err(Timeout(60)) ); drop(rx); assert_eq!( tx.send_timeout(70, Duration::from_secs(1)).await, Err(Closed(70)) ); } #[test] #[should_panic = "there is no reactor running, must be called from the context of a Tokio 1.x runtime"] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn recv_timeout_panic() { use futures::future::FutureExt; use tokio::time::Duration; let (tx, _rx) = mpsc::channel(5); tx.send_timeout(10, Duration::from_secs(1)).now_or_never(); } // Tests that channel `capacity` changes and `max_capacity` stays the same #[tokio::test] async fn test_tx_capacity() { let (tx, _rx) = mpsc::channel::<()>(10); // both capacities are same before assert_eq!(tx.capacity(), 10); assert_eq!(tx.max_capacity(), 10); let _permit = tx.reserve().await.unwrap(); // after reserve, only capacity should drop by one assert_eq!(tx.capacity(), 9); assert_eq!(tx.max_capacity(), 10); tx.send(()).await.unwrap(); // after send, capacity should drop by one again assert_eq!(tx.capacity(), 8); assert_eq!(tx.max_capacity(), 10); } #[tokio::test] async fn test_rx_is_closed_when_calling_close_with_sender() { // is_closed should return true after calling close but still has a sender let (_tx, mut rx) = mpsc::channel::<()>(10); rx.close(); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_is_closed_when_dropping_all_senders() { // is_closed should return true after dropping all senders let (tx, rx) = mpsc::channel::<()>(10); let another_tx = tx.clone(); let task = tokio::spawn(async move { drop(another_tx); }); drop(tx); let _ = task.await; assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_is_not_closed_when_there_are_senders() { // is_closed should return false when there is a sender let (_tx, rx) = mpsc::channel::<()>(10); assert!(!rx.is_closed()); } #[tokio::test] async fn test_rx_is_not_closed_when_there_are_senders_and_buffer_filled() { // is_closed should return false when there is a sender, even if enough messages have been sent to fill the channel let (tx, rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } assert!(!rx.is_closed()); } #[tokio::test] async fn test_rx_is_closed_when_there_are_no_senders_and_there_are_messages() { // is_closed should return true when there are messages in the buffer, but no senders let (tx, rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } drop(tx); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_is_closed_when_there_are_messages_and_close_is_called() { // is_closed should return true when there are messages in the buffer, and close is called let (tx, mut rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } rx.close(); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_is_not_closed_when_there_are_permits_but_not_senders() { // is_closed should return false when there is a permit (but no senders) let (tx, rx) = mpsc::channel::<()>(10); let _permit = tx.reserve_owned().await.expect("Failed to reserve permit"); assert!(!rx.is_closed()); } #[tokio::test] async fn test_rx_is_empty_when_no_messages_were_sent() { let (_tx, rx) = mpsc::channel::<()>(10); assert!(rx.is_empty()) } #[tokio::test] async fn test_rx_is_not_empty_when_there_are_messages_in_the_buffer() { let (tx, rx) = mpsc::channel::<()>(10); assert!(tx.send(()).await.is_ok()); assert!(!rx.is_empty()) } #[tokio::test] async fn test_rx_is_not_empty_when_the_buffer_is_full() { let (tx, rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } assert!(!rx.is_empty()) } #[tokio::test] async fn test_rx_is_not_empty_when_all_but_one_messages_are_consumed() { let (tx, mut rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } for _ in 0..9 { assert!(rx.recv().await.is_some()); } assert!(!rx.is_empty()) } #[tokio::test] async fn test_rx_is_empty_when_all_messages_are_consumed() { let (tx, mut rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } while rx.try_recv().is_ok() {} assert!(rx.is_empty()) } #[tokio::test] async fn test_rx_is_empty_all_senders_are_dropped_and_messages_consumed() { let (tx, mut rx) = mpsc::channel(10); for i in 0..10 { assert!(tx.send(i).await.is_ok()); } drop(tx); for _ in 0..10 { assert!(rx.recv().await.is_some()); } assert!(rx.is_empty()) } #[tokio::test] async fn test_rx_len_on_empty_channel() { let (_tx, rx) = mpsc::channel::<()>(100); assert_eq!(rx.len(), 0); } #[tokio::test] async fn test_rx_len_on_empty_channel_without_senders() { // when all senders are dropped, a "closed" value is added to the end of the linked list. // here we test that the "closed" value does not change the len of the channel. let (tx, rx) = mpsc::channel::<()>(100); drop(tx); assert_eq!(rx.len(), 0); } #[tokio::test] async fn test_rx_len_on_filled_channel() { let (tx, rx) = mpsc::channel(100); for i in 0..100 { assert!(tx.send(i).await.is_ok()); } assert_eq!(rx.len(), 100); } #[tokio::test] async fn test_rx_len_on_filled_channel_without_senders() { let (tx, rx) = mpsc::channel(100); for i in 0..100 { assert!(tx.send(i).await.is_ok()); } drop(tx); assert_eq!(rx.len(), 100); } #[tokio::test] async fn test_rx_len_when_consuming_all_messages() { let (tx, mut rx) = mpsc::channel(100); for i in 0..100 { assert!(tx.send(i).await.is_ok()); assert_eq!(rx.len(), i + 1); } drop(tx); for i in (0..100).rev() { assert!(rx.recv().await.is_some()); assert_eq!(rx.len(), i); } } #[tokio::test] async fn test_rx_len_when_close_is_called() { let (tx, mut rx) = mpsc::channel(100); tx.send(()).await.unwrap(); rx.close(); assert_eq!(rx.len(), 1); } #[tokio::test] async fn test_rx_len_when_close_is_called_before_dropping_sender() { let (tx, mut rx) = mpsc::channel(100); tx.send(()).await.unwrap(); rx.close(); drop(tx); assert_eq!(rx.len(), 1); } #[tokio::test] async fn test_rx_len_when_close_is_called_after_dropping_sender() { let (tx, mut rx) = mpsc::channel(100); tx.send(()).await.unwrap(); drop(tx); rx.close(); assert_eq!(rx.len(), 1); } #[tokio::test] async fn test_rx_unbounded_is_closed_when_calling_close_with_sender() { // is_closed should return true after calling close but still has a sender let (_tx, mut rx) = mpsc::unbounded_channel::<()>(); rx.close(); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_unbounded_is_closed_when_dropping_all_senders() { // is_closed should return true after dropping all senders let (tx, rx) = mpsc::unbounded_channel::<()>(); let another_tx = tx.clone(); let task = tokio::spawn(async move { drop(another_tx); }); drop(tx); let _ = task.await; assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_unbounded_is_not_closed_when_there_are_senders() { // is_closed should return false when there is a sender let (_tx, rx) = mpsc::unbounded_channel::<()>(); assert!(!rx.is_closed()); } #[tokio::test] async fn test_rx_unbounded_is_closed_when_there_are_no_senders_and_there_are_messages() { // is_closed should return true when there are messages in the buffer, but no senders let (tx, rx) = mpsc::unbounded_channel(); for i in 0..10 { assert!(tx.send(i).is_ok()); } drop(tx); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_unbounded_is_closed_when_there_are_messages_and_close_is_called() { // is_closed should return true when there are messages in the buffer, and close is called let (tx, mut rx) = mpsc::unbounded_channel(); for i in 0..10 { assert!(tx.send(i).is_ok()); } rx.close(); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_unbounded_is_empty_when_no_messages_were_sent() { let (_tx, rx) = mpsc::unbounded_channel::<()>(); assert!(rx.is_empty()) } #[tokio::test] async fn test_rx_unbounded_is_not_empty_when_there_are_messages_in_the_buffer() { let (tx, rx) = mpsc::unbounded_channel(); assert!(tx.send(()).is_ok()); assert!(!rx.is_empty()) } #[tokio::test] async fn test_rx_unbounded_is_not_empty_when_all_but_one_messages_are_consumed() { let (tx, mut rx) = mpsc::unbounded_channel(); for i in 0..10 { assert!(tx.send(i).is_ok()); } for _ in 0..9 { assert!(rx.recv().await.is_some()); } assert!(!rx.is_empty()) } #[tokio::test] async fn test_rx_unbounded_is_empty_when_all_messages_are_consumed() { let (tx, mut rx) = mpsc::unbounded_channel(); for i in 0..10 { assert!(tx.send(i).is_ok()); } while rx.try_recv().is_ok() {} assert!(rx.is_empty()) } #[tokio::test] async fn test_rx_unbounded_is_empty_all_senders_are_dropped_and_messages_consumed() { let (tx, mut rx) = mpsc::unbounded_channel(); for i in 0..10 { assert!(tx.send(i).is_ok()); } drop(tx); for _ in 0..10 { assert!(rx.recv().await.is_some()); } assert!(rx.is_empty()) } #[tokio::test] async fn test_rx_unbounded_len_on_empty_channel() { let (_tx, rx) = mpsc::unbounded_channel::<()>(); assert_eq!(rx.len(), 0); } #[tokio::test] async fn test_rx_unbounded_len_on_empty_channel_without_senders() { // when all senders are dropped, a "closed" value is added to the end of the linked list. // here we test that the "closed" value does not change the len of the channel. let (tx, rx) = mpsc::unbounded_channel::<()>(); drop(tx); assert_eq!(rx.len(), 0); } #[tokio::test] async fn test_rx_unbounded_len_with_multiple_messages() { let (tx, rx) = mpsc::unbounded_channel(); for i in 0..100 { assert!(tx.send(i).is_ok()); } assert_eq!(rx.len(), 100); } #[tokio::test] async fn test_rx_unbounded_len_with_multiple_messages_and_dropped_senders() { let (tx, rx) = mpsc::unbounded_channel(); for i in 0..100 { assert!(tx.send(i).is_ok()); } drop(tx); assert_eq!(rx.len(), 100); } #[tokio::test] async fn test_rx_unbounded_len_when_consuming_all_messages() { let (tx, mut rx) = mpsc::unbounded_channel(); for i in 0..100 { assert!(tx.send(i).is_ok()); assert_eq!(rx.len(), i + 1); } drop(tx); for i in (0..100).rev() { assert!(rx.recv().await.is_some()); assert_eq!(rx.len(), i); } } #[tokio::test] async fn test_rx_unbounded_len_when_close_is_called() { let (tx, mut rx) = mpsc::unbounded_channel(); tx.send(()).unwrap(); rx.close(); assert_eq!(rx.len(), 1); } #[tokio::test] async fn test_rx_unbounded_len_when_close_is_called_before_dropping_sender() { let (tx, mut rx) = mpsc::unbounded_channel(); tx.send(()).unwrap(); rx.close(); drop(tx); assert_eq!(rx.len(), 1); } #[tokio::test] async fn test_rx_unbounded_len_when_close_is_called_after_dropping_sender() { let (tx, mut rx) = mpsc::unbounded_channel(); tx.send(()).unwrap(); drop(tx); rx.close(); assert_eq!(rx.len(), 1); } // Regression test for https://github.com/tokio-rs/tokio/issues/6602 #[tokio::test] async fn test_is_empty_32_msgs() { let (sender, mut receiver) = mpsc::channel(33); for value in 1..257 { sender.send(value).await.unwrap(); receiver.recv().await.unwrap(); assert!(receiver.is_empty(), "{value}. len: {}", receiver.len()); } } fn is_debug(_: &T) {} tokio-1.43.1/tests/sync_mpsc_weak.rs000064400000000000000000000446451046102023000155710ustar 00000000000000#![allow(clippy::redundant_clone)] #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::{Acquire, Release}; use tokio::sync::mpsc::{self, channel, unbounded_channel}; use tokio::sync::oneshot; #[tokio::test] async fn weak_sender() { let (tx, mut rx) = channel(11); let tx_weak = tokio::spawn(async move { let tx_weak = tx.clone().downgrade(); for i in 0..10 { if tx.send(i).await.is_err() { return None; } } let tx2 = tx_weak .upgrade() .expect("expected to be able to upgrade tx_weak"); let _ = tx2.send(20).await; let tx_weak = tx2.downgrade(); Some(tx_weak) }) .await .unwrap(); for i in 0..12 { let recvd = rx.recv().await; match recvd { Some(msg) => { if i == 10 { assert_eq!(msg, 20); } } None => { assert_eq!(i, 11); break; } } } let tx_weak = tx_weak.unwrap(); let upgraded = tx_weak.upgrade(); assert!(upgraded.is_none()); } #[tokio::test] async fn actor_weak_sender() { pub struct MyActor { receiver: mpsc::Receiver, sender: mpsc::WeakSender, next_id: u32, pub received_self_msg: bool, } enum ActorMessage { GetUniqueId { respond_to: oneshot::Sender }, SelfMessage {}, } impl MyActor { fn new( receiver: mpsc::Receiver, sender: mpsc::WeakSender, ) -> Self { MyActor { receiver, sender, next_id: 0, received_self_msg: false, } } fn handle_message(&mut self, msg: ActorMessage) { match msg { ActorMessage::GetUniqueId { respond_to } => { self.next_id += 1; // The `let _ =` ignores any errors when sending. // // This can happen if the `select!` macro is used // to cancel waiting for the response. let _ = respond_to.send(self.next_id); } ActorMessage::SelfMessage { .. } => { self.received_self_msg = true; } } } async fn send_message_to_self(&mut self) { let msg = ActorMessage::SelfMessage {}; let sender = self.sender.clone(); // cannot move self.sender here if let Some(sender) = sender.upgrade() { let _ = sender.send(msg).await; self.sender = sender.downgrade(); } } async fn run(&mut self) { let mut i = 0; while let Some(msg) = self.receiver.recv().await { self.handle_message(msg); if i == 0 { self.send_message_to_self().await; } i += 1 } assert!(self.received_self_msg); } } #[derive(Clone)] pub struct MyActorHandle { sender: mpsc::Sender, } impl MyActorHandle { pub fn new() -> (Self, MyActor) { let (sender, receiver) = mpsc::channel(8); let actor = MyActor::new(receiver, sender.clone().downgrade()); (Self { sender }, actor) } pub async fn get_unique_id(&self) -> u32 { let (send, recv) = oneshot::channel(); let msg = ActorMessage::GetUniqueId { respond_to: send }; // Ignore send errors. If this send fails, so does the // recv.await below. There's no reason to check the // failure twice. let _ = self.sender.send(msg).await; recv.await.expect("Actor task has been killed") } } let (handle, mut actor) = MyActorHandle::new(); let actor_handle = tokio::spawn(async move { actor.run().await }); let _ = tokio::spawn(async move { let _ = handle.get_unique_id().await; drop(handle); }) .await; let _ = actor_handle.await; } static NUM_DROPPED: AtomicUsize = AtomicUsize::new(0); #[derive(Debug)] struct Msg; impl Drop for Msg { fn drop(&mut self) { NUM_DROPPED.fetch_add(1, Release); } } // Tests that no pending messages are put onto the channel after `Rx` was // dropped. // // Note: After the introduction of `WeakSender`, which internally // used `Arc` and doesn't call a drop of the channel after the last strong // `Sender` was dropped while more than one `WeakSender` remains, we want to // ensure that no messages are kept in the channel, which were sent after // the receiver was dropped. #[tokio::test] async fn test_msgs_dropped_on_rx_drop() { let (tx, mut rx) = mpsc::channel(3); tx.send(Msg {}).await.unwrap(); tx.send(Msg {}).await.unwrap(); // This msg will be pending and should be dropped when `rx` is dropped let sent_fut = tx.send(Msg {}); let _ = rx.recv().await.unwrap(); let _ = rx.recv().await.unwrap(); sent_fut.await.unwrap(); drop(rx); assert_eq!(NUM_DROPPED.load(Acquire), 3); // This msg will not be put onto `Tx` list anymore, since `Rx` is closed. assert!(tx.send(Msg {}).await.is_err()); assert_eq!(NUM_DROPPED.load(Acquire), 4); } // Tests that a `WeakSender` is upgradeable when other `Sender`s exist. #[test] fn downgrade_upgrade_sender_success() { let (tx, _rx) = mpsc::channel::(1); let weak_tx = tx.downgrade(); assert!(weak_tx.upgrade().is_some()); } // Tests that a `WeakSender` fails to upgrade when no other `Sender` exists. #[test] fn downgrade_upgrade_sender_failure() { let (tx, _rx) = mpsc::channel::(1); let weak_tx = tx.downgrade(); drop(tx); assert!(weak_tx.upgrade().is_none()); } // Tests that a `WeakSender` cannot be upgraded after a `Sender` was dropped, // which existed at the time of the `downgrade` call. #[test] fn downgrade_drop_upgrade() { let (tx, _rx) = mpsc::channel::(1); // the cloned `Tx` is dropped right away let weak_tx = tx.clone().downgrade(); drop(tx); assert!(weak_tx.upgrade().is_none()); } // Tests that we can upgrade a weak sender with an outstanding permit // but no other strong senders. #[tokio::test] async fn downgrade_get_permit_upgrade_no_senders() { let (tx, _rx) = mpsc::channel::(1); let weak_tx = tx.downgrade(); let _permit = tx.reserve_owned().await.unwrap(); assert!(weak_tx.upgrade().is_some()); } // Tests that you can downgrade and upgrade a sender with an outstanding permit // but no other senders left. #[tokio::test] async fn downgrade_upgrade_get_permit_no_senders() { let (tx, _rx) = mpsc::channel::(1); let tx2 = tx.clone(); let _permit = tx.reserve_owned().await.unwrap(); let weak_tx = tx2.downgrade(); drop(tx2); assert!(weak_tx.upgrade().is_some()); } // Tests that `downgrade` does not change the `tx_count` of the channel. #[test] fn test_tx_count_weak_sender() { let (tx, _rx) = mpsc::channel::(1); let tx_weak = tx.downgrade(); let tx_weak2 = tx.downgrade(); drop(tx); assert!(tx_weak.upgrade().is_none() && tx_weak2.upgrade().is_none()); } #[tokio::test] async fn weak_unbounded_sender() { let (tx, mut rx) = unbounded_channel(); let tx_weak = tokio::spawn(async move { let tx_weak = tx.clone().downgrade(); for i in 0..10 { if tx.send(i).is_err() { return None; } } let tx2 = tx_weak .upgrade() .expect("expected to be able to upgrade tx_weak"); let _ = tx2.send(20); let tx_weak = tx2.downgrade(); Some(tx_weak) }) .await .unwrap(); for i in 0..12 { let recvd = rx.recv().await; match recvd { Some(msg) => { if i == 10 { assert_eq!(msg, 20); } } None => { assert_eq!(i, 11); break; } } } let tx_weak = tx_weak.unwrap(); let upgraded = tx_weak.upgrade(); assert!(upgraded.is_none()); } #[tokio::test] async fn actor_weak_unbounded_sender() { pub struct MyActor { receiver: mpsc::UnboundedReceiver, sender: mpsc::WeakUnboundedSender, next_id: u32, pub received_self_msg: bool, } enum ActorMessage { GetUniqueId { respond_to: oneshot::Sender }, SelfMessage {}, } impl MyActor { fn new( receiver: mpsc::UnboundedReceiver, sender: mpsc::WeakUnboundedSender, ) -> Self { MyActor { receiver, sender, next_id: 0, received_self_msg: false, } } fn handle_message(&mut self, msg: ActorMessage) { match msg { ActorMessage::GetUniqueId { respond_to } => { self.next_id += 1; // The `let _ =` ignores any errors when sending. // // This can happen if the `select!` macro is used // to cancel waiting for the response. let _ = respond_to.send(self.next_id); } ActorMessage::SelfMessage { .. } => { self.received_self_msg = true; } } } async fn send_message_to_self(&mut self) { let msg = ActorMessage::SelfMessage {}; let sender = self.sender.clone(); // cannot move self.sender here if let Some(sender) = sender.upgrade() { let _ = sender.send(msg); self.sender = sender.downgrade(); } } async fn run(&mut self) { let mut i = 0; while let Some(msg) = self.receiver.recv().await { self.handle_message(msg); if i == 0 { self.send_message_to_self().await; } i += 1 } assert!(self.received_self_msg); } } #[derive(Clone)] pub struct MyActorHandle { sender: mpsc::UnboundedSender, } impl MyActorHandle { pub fn new() -> (Self, MyActor) { let (sender, receiver) = mpsc::unbounded_channel(); let actor = MyActor::new(receiver, sender.clone().downgrade()); (Self { sender }, actor) } pub async fn get_unique_id(&self) -> u32 { let (send, recv) = oneshot::channel(); let msg = ActorMessage::GetUniqueId { respond_to: send }; // Ignore send errors. If this send fails, so does the // recv.await below. There's no reason to check the // failure twice. let _ = self.sender.send(msg); recv.await.expect("Actor task has been killed") } } let (handle, mut actor) = MyActorHandle::new(); let actor_handle = tokio::spawn(async move { actor.run().await }); let _ = tokio::spawn(async move { let _ = handle.get_unique_id().await; drop(handle); }) .await; let _ = actor_handle.await; } static NUM_DROPPED_UNBOUNDED: AtomicUsize = AtomicUsize::new(0); #[derive(Debug)] struct MsgUnbounded; impl Drop for MsgUnbounded { fn drop(&mut self) { NUM_DROPPED_UNBOUNDED.fetch_add(1, Release); } } // Tests that no pending messages are put onto the channel after `Rx` was // dropped. // // Note: After the introduction of `UnboundedWeakSender`, which internally // used `Arc` and doesn't call a drop of the channel after the last strong // `UnboundedSender` was dropped while more than one `UnboundedWeakSender` // remains, we want to ensure that no messages are kept in the channel, which // were sent after the receiver was dropped. #[tokio::test] async fn test_msgs_dropped_on_unbounded_rx_drop() { let (tx, mut rx) = mpsc::unbounded_channel(); tx.send(MsgUnbounded {}).unwrap(); tx.send(MsgUnbounded {}).unwrap(); // This msg will be pending and should be dropped when `rx` is dropped let sent = tx.send(MsgUnbounded {}); let _ = rx.recv().await.unwrap(); let _ = rx.recv().await.unwrap(); sent.unwrap(); drop(rx); assert_eq!(NUM_DROPPED_UNBOUNDED.load(Acquire), 3); // This msg will not be put onto `Tx` list anymore, since `Rx` is closed. assert!(tx.send(MsgUnbounded {}).is_err()); assert_eq!(NUM_DROPPED_UNBOUNDED.load(Acquire), 4); } // Tests that an `WeakUnboundedSender` is upgradeable when other // `UnboundedSender`s exist. #[test] fn downgrade_upgrade_unbounded_sender_success() { let (tx, _rx) = mpsc::unbounded_channel::(); let weak_tx = tx.downgrade(); assert!(weak_tx.upgrade().is_some()); } // Tests that a `WeakUnboundedSender` fails to upgrade when no other // `UnboundedSender` exists. #[test] fn downgrade_upgrade_unbounded_sender_failure() { let (tx, _rx) = mpsc::unbounded_channel::(); let weak_tx = tx.downgrade(); drop(tx); assert!(weak_tx.upgrade().is_none()); } // Tests that an `WeakUnboundedSender` cannot be upgraded after an // `UnboundedSender` was dropped, which existed at the time of the `downgrade` call. #[test] fn downgrade_drop_upgrade_unbounded() { let (tx, _rx) = mpsc::unbounded_channel::(); // the cloned `Tx` is dropped right away let weak_tx = tx.clone().downgrade(); drop(tx); assert!(weak_tx.upgrade().is_none()); } // Tests that `downgrade` does not change the `tx_count` of the channel. #[test] fn test_tx_count_weak_unbounded_sender() { let (tx, _rx) = mpsc::unbounded_channel::(); let tx_weak = tx.downgrade(); let tx_weak2 = tx.downgrade(); drop(tx); assert!(tx_weak.upgrade().is_none() && tx_weak2.upgrade().is_none()); } #[tokio::test] async fn test_rx_is_closed_when_dropping_all_senders_except_weak_senders() { // is_closed should return true after dropping all senders except for a weak sender let (tx, rx) = mpsc::channel::<()>(10); let _weak_sender = tx.clone().downgrade(); drop(tx); assert!(rx.is_closed()); } #[tokio::test] async fn test_rx_unbounded_is_closed_when_dropping_all_senders_except_weak_senders() { // is_closed should return true after dropping all senders except for a weak sender let (tx, rx) = mpsc::unbounded_channel::<()>(); let _weak_sender = tx.clone().downgrade(); drop(tx); assert!(rx.is_closed()); } #[tokio::test] async fn sender_strong_count_when_cloned() { let (tx, rx) = mpsc::channel::<()>(1); let tx2 = tx.clone(); assert_eq!(tx.strong_count(), 2); assert_eq!(tx2.strong_count(), 2); assert_eq!(rx.sender_strong_count(), 2); } #[tokio::test] async fn sender_weak_count_when_downgraded() { let (tx, _rx) = mpsc::channel::<()>(1); let weak = tx.downgrade(); assert_eq!(tx.weak_count(), 1); assert_eq!(weak.weak_count(), 1); } #[tokio::test] async fn sender_strong_count_when_dropped() { let (tx, rx) = mpsc::channel::<()>(1); let tx2 = tx.clone(); drop(tx2); assert_eq!(tx.strong_count(), 1); assert_eq!(rx.sender_strong_count(), 1); } #[tokio::test] async fn sender_weak_count_when_dropped() { let (tx, rx) = mpsc::channel::<()>(1); let weak = tx.downgrade(); drop(weak); assert_eq!(tx.weak_count(), 0); assert_eq!(rx.sender_weak_count(), 0); } #[tokio::test] async fn sender_strong_and_weak_conut() { let (tx, rx) = mpsc::channel::<()>(1); let tx2 = tx.clone(); let weak = tx.downgrade(); let weak2 = tx2.downgrade(); assert_eq!(tx.strong_count(), 2); assert_eq!(tx2.strong_count(), 2); assert_eq!(weak.strong_count(), 2); assert_eq!(weak2.strong_count(), 2); assert_eq!(rx.sender_strong_count(), 2); assert_eq!(tx.weak_count(), 2); assert_eq!(tx2.weak_count(), 2); assert_eq!(weak.weak_count(), 2); assert_eq!(weak2.weak_count(), 2); assert_eq!(rx.sender_weak_count(), 2); drop(tx2); drop(weak2); assert_eq!(tx.strong_count(), 1); assert_eq!(weak.strong_count(), 1); assert_eq!(rx.sender_strong_count(), 1); assert_eq!(tx.weak_count(), 1); assert_eq!(weak.weak_count(), 1); assert_eq!(rx.sender_weak_count(), 1); } #[tokio::test] async fn unbounded_sender_strong_count_when_cloned() { let (tx, rx) = mpsc::unbounded_channel::<()>(); let tx2 = tx.clone(); assert_eq!(tx.strong_count(), 2); assert_eq!(tx2.strong_count(), 2); assert_eq!(rx.sender_strong_count(), 2); } #[tokio::test] async fn unbounded_sender_weak_count_when_downgraded() { let (tx, rx) = mpsc::unbounded_channel::<()>(); let weak = tx.downgrade(); assert_eq!(tx.weak_count(), 1); assert_eq!(weak.weak_count(), 1); assert_eq!(rx.sender_weak_count(), 1); } #[tokio::test] async fn unbounded_sender_strong_count_when_dropped() { let (tx, rx) = mpsc::unbounded_channel::<()>(); let tx2 = tx.clone(); drop(tx2); assert_eq!(tx.strong_count(), 1); assert_eq!(rx.sender_strong_count(), 1); } #[tokio::test] async fn unbounded_sender_weak_count_when_dropped() { let (tx, rx) = mpsc::unbounded_channel::<()>(); let weak = tx.downgrade(); drop(weak); assert_eq!(tx.weak_count(), 0); assert_eq!(rx.sender_weak_count(), 0); } #[tokio::test] async fn unbounded_sender_strong_and_weak_conut() { let (tx, rx) = mpsc::unbounded_channel::<()>(); let tx2 = tx.clone(); let weak = tx.downgrade(); let weak2 = tx2.downgrade(); assert_eq!(tx.strong_count(), 2); assert_eq!(tx2.strong_count(), 2); assert_eq!(weak.strong_count(), 2); assert_eq!(weak2.strong_count(), 2); assert_eq!(rx.sender_strong_count(), 2); assert_eq!(tx.weak_count(), 2); assert_eq!(tx2.weak_count(), 2); assert_eq!(weak.weak_count(), 2); assert_eq!(weak2.weak_count(), 2); assert_eq!(rx.sender_weak_count(), 2); drop(tx2); drop(weak2); assert_eq!(tx.strong_count(), 1); assert_eq!(weak.strong_count(), 1); assert_eq!(rx.sender_strong_count(), 1); assert_eq!(tx.weak_count(), 1); assert_eq!(weak.weak_count(), 1); assert_eq!(rx.sender_weak_count(), 1); } tokio-1.43.1/tests/sync_mutex.rs000064400000000000000000000110471046102023000147500ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::Mutex; use tokio_test::task::spawn; use tokio_test::{assert_pending, assert_ready}; use std::sync::Arc; #[test] fn straight_execution() { let l = Mutex::new(100); { let mut t = spawn(l.lock()); let mut g = assert_ready!(t.poll()); assert_eq!(&*g, &100); *g = 99; } { let mut t = spawn(l.lock()); let mut g = assert_ready!(t.poll()); assert_eq!(&*g, &99); *g = 98; } { let mut t = spawn(l.lock()); let g = assert_ready!(t.poll()); assert_eq!(&*g, &98); } } #[test] fn readiness() { let l1 = Arc::new(Mutex::new(100)); let l2 = Arc::clone(&l1); let mut t1 = spawn(l1.lock()); let mut t2 = spawn(l2.lock()); let g = assert_ready!(t1.poll()); // We can't now acquire the lease since it's already held in g assert_pending!(t2.poll()); // But once g unlocks, we can acquire it drop(g); assert!(t2.is_woken()); let _t2 = assert_ready!(t2.poll()); } /* #[test] #[ignore] fn lock() { let mut lock = Mutex::new(false); let mut lock2 = lock.clone(); std::thread::spawn(move || { let l = lock2.lock(); pin_mut!(l); let mut task = MockTask::new(); let mut g = assert_ready!(task.poll(&mut l)); std::thread::sleep(std::time::Duration::from_millis(500)); *g = true; drop(g); }); std::thread::sleep(std::time::Duration::from_millis(50)); let mut task = MockTask::new(); let l = lock.lock(); pin_mut!(l); assert_pending!(task.poll(&mut l)); std::thread::sleep(std::time::Duration::from_millis(500)); assert!(task.is_woken()); let result = assert_ready!(task.poll(&mut l)); assert!(*result); } */ /// Ensure a mutex is unlocked if a future holding the lock /// is aborted prematurely. #[tokio::test] #[cfg(feature = "full")] async fn aborted_future_1() { use std::time::Duration; use tokio::time::{interval, timeout}; let m1: Arc> = Arc::new(Mutex::new(0)); { let m2 = m1.clone(); // Try to lock mutex in a future that is aborted prematurely timeout(Duration::from_millis(1u64), async move { let iv = interval(Duration::from_millis(1000)); tokio::pin!(iv); let _g = m2.lock().await; iv.as_mut().tick().await; iv.as_mut().tick().await; }) .await .unwrap_err(); } // This should succeed as there is no lock left for the mutex. timeout(Duration::from_millis(1u64), async move { let _g = m1.lock().await; }) .await .expect("Mutex is locked"); } /// This test is similar to `aborted_future_1` but this time the /// aborted future is waiting for the lock. #[tokio::test] #[cfg(feature = "full")] async fn aborted_future_2() { use std::time::Duration; use tokio::time::timeout; let m1: Arc> = Arc::new(Mutex::new(0)); { // Lock mutex let _lock = m1.lock().await; { let m2 = m1.clone(); // Try to lock mutex in a future that is aborted prematurely timeout(Duration::from_millis(1u64), async move { let _g = m2.lock().await; }) .await .unwrap_err(); } } // This should succeed as there is no lock left for the mutex. timeout(Duration::from_millis(1u64), async move { let _g = m1.lock().await; }) .await .expect("Mutex is locked"); } #[test] fn try_lock() { let m: Mutex = Mutex::new(0); { let g1 = m.try_lock(); assert!(g1.is_ok()); let g2 = m.try_lock(); assert!(g2.is_err()); } let g3 = m.try_lock(); assert!(g3.is_ok()); } #[maybe_tokio_test] async fn debug_format() { let s = "debug"; let m = Mutex::new(s.to_string()); assert_eq!(format!("{s:?}"), format!("{:?}", m.lock().await)); } #[maybe_tokio_test] async fn mutex_debug() { let s = "data"; let m = Mutex::new(s.to_string()); assert_eq!(format!("{m:?}"), r#"Mutex { data: "data" }"#); let _guard = m.lock().await; assert_eq!(format!("{m:?}"), r#"Mutex { data: }"#) } tokio-1.43.1/tests/sync_mutex_owned.rs000064400000000000000000000072451046102023000161510ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::Mutex; use tokio_test::task::spawn; use tokio_test::{assert_pending, assert_ready}; use std::sync::Arc; #[test] fn straight_execution() { let l = Arc::new(Mutex::new(100)); { let mut t = spawn(l.clone().lock_owned()); let mut g = assert_ready!(t.poll()); assert_eq!(&*g, &100); *g = 99; } { let mut t = spawn(l.clone().lock_owned()); let mut g = assert_ready!(t.poll()); assert_eq!(&*g, &99); *g = 98; } { let mut t = spawn(l.lock_owned()); let g = assert_ready!(t.poll()); assert_eq!(&*g, &98); } } #[test] fn readiness() { let l = Arc::new(Mutex::new(100)); let mut t1 = spawn(l.clone().lock_owned()); let mut t2 = spawn(l.lock_owned()); let g = assert_ready!(t1.poll()); // We can't now acquire the lease since it's already held in g assert_pending!(t2.poll()); // But once g unlocks, we can acquire it drop(g); assert!(t2.is_woken()); assert_ready!(t2.poll()); } /// Ensure a mutex is unlocked if a future holding the lock /// is aborted prematurely. #[tokio::test] #[cfg(feature = "full")] #[cfg_attr(miri, ignore)] async fn aborted_future_1() { use std::time::Duration; use tokio::time::{interval, timeout}; let m1: Arc> = Arc::new(Mutex::new(0)); { let m2 = m1.clone(); // Try to lock mutex in a future that is aborted prematurely timeout(Duration::from_millis(1u64), async move { let iv = interval(Duration::from_millis(1000)); tokio::pin!(iv); m2.lock_owned().await; iv.as_mut().tick().await; iv.as_mut().tick().await; }) .await .unwrap_err(); } // This should succeed as there is no lock left for the mutex. timeout(Duration::from_millis(1u64), async move { m1.lock_owned().await; }) .await .expect("Mutex is locked"); } /// This test is similar to `aborted_future_1` but this time the /// aborted future is waiting for the lock. #[tokio::test] #[cfg(feature = "full")] async fn aborted_future_2() { use std::time::Duration; use tokio::time::timeout; let m1: Arc> = Arc::new(Mutex::new(0)); { // Lock mutex let _lock = m1.clone().lock_owned().await; { let m2 = m1.clone(); // Try to lock mutex in a future that is aborted prematurely timeout(Duration::from_millis(1u64), async move { m2.lock_owned().await; }) .await .unwrap_err(); } } // This should succeed as there is no lock left for the mutex. timeout(Duration::from_millis(1u64), async move { m1.lock_owned().await; }) .await .expect("Mutex is locked"); } #[test] fn try_lock_owned() { let m: Arc> = Arc::new(Mutex::new(0)); { let g1 = m.clone().try_lock_owned(); assert!(g1.is_ok()); let g2 = m.clone().try_lock_owned(); assert!(g2.is_err()); } let g3 = m.try_lock_owned(); assert!(g3.is_ok()); } #[maybe_tokio_test] async fn debug_format() { let s = "debug"; let m = Arc::new(Mutex::new(s.to_string())); assert_eq!(format!("{s:?}"), format!("{:?}", m.lock_owned().await)); } tokio-1.43.1/tests/sync_notify.rs000064400000000000000000000165561046102023000151300ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::Notify; use tokio_test::task::spawn; use tokio_test::*; #[allow(unused)] trait AssertSend: Send + Sync {} impl AssertSend for Notify {} #[test] fn notify_notified_one() { let notify = Notify::new(); let mut notified = spawn(async { notify.notified().await }); notify.notify_one(); assert_ready!(notified.poll()); } #[test] fn notify_multi_notified_one() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); // add two waiters into the queue assert_pending!(notified1.poll()); assert_pending!(notified2.poll()); // should wakeup the first one notify.notify_one(); assert_ready!(notified1.poll()); assert_pending!(notified2.poll()); } #[test] fn notify_multi_notified_last() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); // add two waiters into the queue assert_pending!(notified1.poll()); assert_pending!(notified2.poll()); // should wakeup the last one notify.notify_last(); assert_pending!(notified1.poll()); assert_ready!(notified2.poll()); } #[test] fn notified_one_notify() { let notify = Notify::new(); let mut notified = spawn(async { notify.notified().await }); assert_pending!(notified.poll()); notify.notify_one(); assert!(notified.is_woken()); assert_ready!(notified.poll()); } #[test] fn notified_multi_notify() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); assert_pending!(notified1.poll()); assert_pending!(notified2.poll()); notify.notify_one(); assert!(notified1.is_woken()); assert!(!notified2.is_woken()); assert_ready!(notified1.poll()); assert_pending!(notified2.poll()); } #[test] fn notify_notified_multi() { let notify = Notify::new(); notify.notify_one(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); assert_ready!(notified1.poll()); assert_pending!(notified2.poll()); notify.notify_one(); assert!(notified2.is_woken()); assert_ready!(notified2.poll()); } #[test] fn notified_drop_notified_notify() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); assert_pending!(notified1.poll()); drop(notified1); assert_pending!(notified2.poll()); notify.notify_one(); assert!(notified2.is_woken()); assert_ready!(notified2.poll()); } #[test] fn notified_multi_notify_drop_one() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); assert_pending!(notified1.poll()); assert_pending!(notified2.poll()); notify.notify_one(); assert!(notified1.is_woken()); assert!(!notified2.is_woken()); drop(notified1); assert!(notified2.is_woken()); assert_ready!(notified2.poll()); } #[test] fn notified_multi_notify_one_drop() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); let mut notified3 = spawn(async { notify.notified().await }); // add waiters by order of poll execution assert_pending!(notified1.poll()); assert_pending!(notified2.poll()); assert_pending!(notified3.poll()); // by default fifo notify.notify_one(); drop(notified1); // next waiter should be the one to be to woken up assert_ready!(notified2.poll()); assert_pending!(notified3.poll()); } #[test] fn notified_multi_notify_last_drop() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); let mut notified2 = spawn(async { notify.notified().await }); let mut notified3 = spawn(async { notify.notified().await }); // add waiters by order of poll execution assert_pending!(notified1.poll()); assert_pending!(notified2.poll()); assert_pending!(notified3.poll()); notify.notify_last(); drop(notified3); // latest waiter added should be the one to woken up assert_ready!(notified2.poll()); assert_pending!(notified1.poll()); } #[test] fn notify_in_drop_after_wake() { use futures::task::ArcWake; use std::future::Future; use std::sync::Arc; let notify = Arc::new(Notify::new()); struct NotifyOnDrop(Arc); impl ArcWake for NotifyOnDrop { fn wake_by_ref(_arc_self: &Arc) {} } impl Drop for NotifyOnDrop { fn drop(&mut self) { self.0.notify_waiters(); } } let mut fut = Box::pin(async { notify.notified().await; }); { let waker = futures::task::waker(Arc::new(NotifyOnDrop(notify.clone()))); let mut cx = std::task::Context::from_waker(&waker); assert!(fut.as_mut().poll(&mut cx).is_pending()); } // Now, notifying **should not** deadlock notify.notify_waiters(); } #[test] fn notify_one_after_dropped_all() { let notify = Notify::new(); let mut notified1 = spawn(async { notify.notified().await }); assert_pending!(notified1.poll()); notify.notify_waiters(); notify.notify_one(); drop(notified1); let mut notified2 = spawn(async { notify.notified().await }); assert_ready!(notified2.poll()); } #[test] fn test_notify_one_not_enabled() { let notify = Notify::new(); let mut future = spawn(notify.notified()); notify.notify_one(); assert_ready!(future.poll()); } #[test] fn test_notify_one_after_enable() { let notify = Notify::new(); let mut future = spawn(notify.notified()); future.enter(|_, fut| assert!(!fut.enable())); notify.notify_one(); assert_ready!(future.poll()); future.enter(|_, fut| assert!(fut.enable())); } #[test] fn test_poll_after_enable() { let notify = Notify::new(); let mut future = spawn(notify.notified()); future.enter(|_, fut| assert!(!fut.enable())); assert_pending!(future.poll()); } #[test] fn test_enable_after_poll() { let notify = Notify::new(); let mut future = spawn(notify.notified()); assert_pending!(future.poll()); future.enter(|_, fut| assert!(!fut.enable())); } #[test] fn test_enable_consumes_permit() { let notify = Notify::new(); // Add a permit. notify.notify_one(); let mut future1 = spawn(notify.notified()); future1.enter(|_, fut| assert!(fut.enable())); let mut future2 = spawn(notify.notified()); future2.enter(|_, fut| assert!(!fut.enable())); } #[test] fn test_waker_update() { use futures::task::noop_waker; use std::future::Future; use std::task::Context; let notify = Notify::new(); let mut future = spawn(notify.notified()); let noop = noop_waker(); future.enter(|_, fut| assert_pending!(fut.poll(&mut Context::from_waker(&noop)))); assert_pending!(future.poll()); notify.notify_one(); assert!(future.is_woken()); } tokio-1.43.1/tests/sync_once_cell.rs000064400000000000000000000143501046102023000155310ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::mem; use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; use tokio::runtime; use tokio::sync::OnceCell; use tokio::sync::SetError; use tokio::time; #[test] fn drop_cell() { static NUM_DROPS: AtomicU32 = AtomicU32::new(0); struct Foo {} let fooer = Foo {}; impl Drop for Foo { fn drop(&mut self) { NUM_DROPS.fetch_add(1, Ordering::Release); } } { let once_cell = OnceCell::new(); let prev = once_cell.set(fooer); assert!(prev.is_ok()) } assert!(NUM_DROPS.load(Ordering::Acquire) == 1); } #[test] fn drop_cell_new_with() { static NUM_DROPS: AtomicU32 = AtomicU32::new(0); struct Foo {} let fooer = Foo {}; impl Drop for Foo { fn drop(&mut self) { NUM_DROPS.fetch_add(1, Ordering::Release); } } { let once_cell = OnceCell::new_with(Some(fooer)); assert!(once_cell.initialized()); } assert!(NUM_DROPS.load(Ordering::Acquire) == 1); } #[test] fn drop_into_inner() { static NUM_DROPS: AtomicU32 = AtomicU32::new(0); struct Foo {} let fooer = Foo {}; impl Drop for Foo { fn drop(&mut self) { NUM_DROPS.fetch_add(1, Ordering::Release); } } let once_cell = OnceCell::new(); assert!(once_cell.set(fooer).is_ok()); let fooer = once_cell.into_inner(); let count = NUM_DROPS.load(Ordering::Acquire); assert!(count == 0); drop(fooer); let count = NUM_DROPS.load(Ordering::Acquire); assert!(count == 1); } #[test] fn drop_into_inner_new_with() { static NUM_DROPS: AtomicU32 = AtomicU32::new(0); struct Foo {} let fooer = Foo {}; impl Drop for Foo { fn drop(&mut self) { NUM_DROPS.fetch_add(1, Ordering::Release); } } let once_cell = OnceCell::new_with(Some(fooer)); let fooer = once_cell.into_inner(); let count = NUM_DROPS.load(Ordering::Acquire); assert!(count == 0); mem::drop(fooer); let count = NUM_DROPS.load(Ordering::Acquire); assert!(count == 1); } #[test] fn from() { let cell = OnceCell::from(2); assert_eq!(*cell.get().unwrap(), 2); } async fn func1() -> u32 { 5 } async fn func2() -> u32 { time::sleep(Duration::from_millis(1)).await; 10 } async fn func_err() -> Result { Err(()) } async fn func_ok() -> Result { Ok(10) } async fn func_panic() -> u32 { time::sleep(Duration::from_millis(1)).await; panic!(); } async fn sleep_and_set() -> u32 { // Simulate sleep by pausing time and waiting for another thread to // resume clock when calling `set`, then finding the cell being initialized // by this call time::sleep(Duration::from_millis(2)).await; 5 } async fn advance_time_and_set(cell: &'static OnceCell, v: u32) -> Result<(), SetError> { time::advance(Duration::from_millis(1)).await; cell.set(v) } #[test] fn get_or_init() { let rt = runtime::Builder::new_current_thread() .enable_time() .start_paused(true) .build() .unwrap(); static ONCE: OnceCell = OnceCell::const_new(); rt.block_on(async { let handle1 = rt.spawn(async { ONCE.get_or_init(func1).await }); let handle2 = rt.spawn(async { ONCE.get_or_init(func2).await }); time::advance(Duration::from_millis(1)).await; time::resume(); let result1 = handle1.await.unwrap(); let result2 = handle2.await.unwrap(); assert_eq!(*result1, 5); assert_eq!(*result2, 5); }); } #[test] fn get_or_init_panic() { let rt = runtime::Builder::new_current_thread() .enable_time() .build() .unwrap(); static ONCE: OnceCell = OnceCell::const_new(); rt.block_on(async { time::pause(); let handle1 = rt.spawn(async { ONCE.get_or_init(func1).await }); let handle2 = rt.spawn(async { ONCE.get_or_init(func_panic).await }); time::advance(Duration::from_millis(1)).await; let result1 = handle1.await.unwrap(); let result2 = handle2.await.unwrap(); assert_eq!(*result1, 5); assert_eq!(*result2, 5); }); } #[test] fn set_and_get() { let rt = runtime::Builder::new_current_thread() .enable_time() .build() .unwrap(); static ONCE: OnceCell = OnceCell::const_new(); rt.block_on(async { let _ = rt.spawn(async { ONCE.set(5) }).await; let value = ONCE.get().unwrap(); assert_eq!(*value, 5); }); } #[test] fn get_uninit() { static ONCE: OnceCell = OnceCell::const_new(); let uninit = ONCE.get(); assert!(uninit.is_none()); } #[test] fn set_twice() { static ONCE: OnceCell = OnceCell::const_new(); let first = ONCE.set(5); assert_eq!(first, Ok(())); let second = ONCE.set(6); assert!(second.err().unwrap().is_already_init_err()); } #[test] fn set_while_initializing() { let rt = runtime::Builder::new_current_thread() .enable_time() .build() .unwrap(); static ONCE: OnceCell = OnceCell::const_new(); rt.block_on(async { time::pause(); let handle1 = rt.spawn(async { ONCE.get_or_init(sleep_and_set).await }); let handle2 = rt.spawn(async { advance_time_and_set(&ONCE, 10).await }); time::advance(Duration::from_millis(2)).await; let result1 = handle1.await.unwrap(); let result2 = handle2.await.unwrap(); assert_eq!(*result1, 5); assert!(result2.err().unwrap().is_initializing_err()); }); } #[test] fn get_or_try_init() { let rt = runtime::Builder::new_current_thread() .enable_time() .start_paused(true) .build() .unwrap(); static ONCE: OnceCell = OnceCell::const_new(); rt.block_on(async { let handle1 = rt.spawn(async { ONCE.get_or_try_init(func_err).await }); let handle2 = rt.spawn(async { ONCE.get_or_try_init(func_ok).await }); time::advance(Duration::from_millis(1)).await; time::resume(); let result1 = handle1.await.unwrap(); assert!(result1.is_err()); let result2 = handle2.await.unwrap(); assert_eq!(*result2.unwrap(), 10); }); } tokio-1.43.1/tests/sync_oneshot.rs000064400000000000000000000150441046102023000152660ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::oneshot; use tokio::sync::oneshot::error::TryRecvError; use tokio_test::*; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; #[allow(unused)] trait AssertSend: Send {} impl AssertSend for oneshot::Sender {} impl AssertSend for oneshot::Receiver {} #[allow(unused)] trait SenderExt { fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()>; } impl SenderExt for oneshot::Sender { fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()> { tokio::pin! { let fut = self.closed(); } fut.poll(cx) } } #[test] fn send_recv() { let (tx, rx) = oneshot::channel(); let mut rx = task::spawn(rx); assert_pending!(rx.poll()); assert_ok!(tx.send(1)); assert!(rx.is_woken()); let val = assert_ready_ok!(rx.poll()); assert_eq!(val, 1); } #[maybe_tokio_test] async fn async_send_recv() { let (tx, rx) = oneshot::channel(); assert_ok!(tx.send(1)); assert_eq!(1, assert_ok!(rx.await)); } #[test] fn close_tx() { let (tx, rx) = oneshot::channel::(); let mut rx = task::spawn(rx); assert_pending!(rx.poll()); drop(tx); assert!(rx.is_woken()); assert_ready_err!(rx.poll()); } #[test] fn close_rx() { // First, without checking poll_closed() // let (tx, _) = oneshot::channel(); assert_err!(tx.send(1)); // Second, via poll_closed(); let (tx, rx) = oneshot::channel(); let mut tx = task::spawn(tx); assert_pending!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); drop(rx); assert!(tx.is_woken()); assert!(tx.is_closed()); assert_ready!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); assert_err!(tx.into_inner().send(1)); } #[tokio::test] #[cfg(feature = "full")] async fn async_rx_closed() { let (mut tx, rx) = oneshot::channel::<()>(); tokio::spawn(async move { drop(rx); }); tx.closed().await; } #[test] fn explicit_close_poll() { // First, with message sent let (tx, rx) = oneshot::channel(); let mut rx = task::spawn(rx); assert_ok!(tx.send(1)); rx.close(); let value = assert_ready_ok!(rx.poll()); assert_eq!(value, 1); // Second, without the message sent let (tx, rx) = oneshot::channel::(); let mut tx = task::spawn(tx); let mut rx = task::spawn(rx); assert_pending!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); rx.close(); assert!(tx.is_woken()); assert!(tx.is_closed()); assert_ready!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); assert_err!(tx.into_inner().send(1)); assert_ready_err!(rx.poll()); // Again, but without sending the value this time let (tx, rx) = oneshot::channel::(); let mut tx = task::spawn(tx); let mut rx = task::spawn(rx); assert_pending!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); rx.close(); assert!(tx.is_woken()); assert!(tx.is_closed()); assert_ready!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); assert_ready_err!(rx.poll()); } #[test] fn explicit_close_try_recv() { // First, with message sent let (tx, mut rx) = oneshot::channel(); assert_ok!(tx.send(1)); rx.close(); let val = assert_ok!(rx.try_recv()); assert_eq!(1, val); // Second, without the message sent let (tx, mut rx) = oneshot::channel::(); let mut tx = task::spawn(tx); assert_pending!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); rx.close(); assert!(tx.is_woken()); assert!(tx.is_closed()); assert_ready!(tx.enter(|cx, mut tx| tx.poll_closed(cx))); assert_err!(rx.try_recv()); } #[test] #[should_panic] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn close_try_recv_poll() { let (_tx, rx) = oneshot::channel::(); let mut rx = task::spawn(rx); rx.close(); assert_err!(rx.try_recv()); let _ = rx.poll(); } #[test] fn close_after_recv() { let (tx, mut rx) = oneshot::channel::(); tx.send(17).unwrap(); assert_eq!(17, rx.try_recv().unwrap()); rx.close(); } #[test] fn try_recv_after_completion() { let (tx, mut rx) = oneshot::channel::(); tx.send(17).unwrap(); assert_eq!(17, rx.try_recv().unwrap()); assert_eq!(Err(TryRecvError::Closed), rx.try_recv()); rx.close(); } #[test] fn try_recv_after_completion_await() { let (tx, rx) = oneshot::channel::(); let mut rx = task::spawn(rx); tx.send(17).unwrap(); assert_eq!(Ok(17), assert_ready!(rx.poll())); assert_eq!(Err(TryRecvError::Closed), rx.try_recv()); rx.close(); } #[test] fn drops_tasks() { let (mut tx, mut rx) = oneshot::channel::(); let mut tx_task = task::spawn(()); let mut rx_task = task::spawn(()); assert_pending!(tx_task.enter(|cx, _| tx.poll_closed(cx))); assert_pending!(rx_task.enter(|cx, _| Pin::new(&mut rx).poll(cx))); drop(tx); drop(rx); assert_eq!(1, tx_task.waker_ref_count()); assert_eq!(1, rx_task.waker_ref_count()); } #[test] fn receiver_changes_task() { let (tx, mut rx) = oneshot::channel(); let mut task1 = task::spawn(()); let mut task2 = task::spawn(()); assert_pending!(task1.enter(|cx, _| Pin::new(&mut rx).poll(cx))); assert_eq!(2, task1.waker_ref_count()); assert_eq!(1, task2.waker_ref_count()); assert_pending!(task2.enter(|cx, _| Pin::new(&mut rx).poll(cx))); assert_eq!(1, task1.waker_ref_count()); assert_eq!(2, task2.waker_ref_count()); assert_ok!(tx.send(1)); assert!(!task1.is_woken()); assert!(task2.is_woken()); assert_ready_ok!(task2.enter(|cx, _| Pin::new(&mut rx).poll(cx))); } #[test] fn sender_changes_task() { let (mut tx, rx) = oneshot::channel::(); let mut task1 = task::spawn(()); let mut task2 = task::spawn(()); assert_pending!(task1.enter(|cx, _| tx.poll_closed(cx))); assert_eq!(2, task1.waker_ref_count()); assert_eq!(1, task2.waker_ref_count()); assert_pending!(task2.enter(|cx, _| tx.poll_closed(cx))); assert_eq!(1, task1.waker_ref_count()); assert_eq!(2, task2.waker_ref_count()); drop(rx); assert!(!task1.is_woken()); assert!(task2.is_woken()); assert_ready!(task2.enter(|cx, _| tx.poll_closed(cx))); } tokio-1.43.1/tests/sync_panic.rs000064400000000000000000000141151046102023000146770ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] #![cfg(panic = "unwind")] use std::{error::Error, sync::Arc}; use tokio::{ runtime::{Builder, Runtime}, sync::{broadcast, mpsc, oneshot, Mutex, RwLock, Semaphore}, }; mod support { pub mod panic; } use support::panic::test_panic; #[test] fn broadcast_channel_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let (_, _) = broadcast::channel::(0); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mutex_blocking_lock_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); rt.block_on(async { let mutex = Mutex::new(5_u32); let _g = mutex.blocking_lock(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn oneshot_blocking_recv_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); rt.block_on(async { let (_tx, rx) = oneshot::channel::(); let _ = rx.blocking_recv(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn rwlock_with_max_readers_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = RwLock::::with_max_readers(0, (u32::MAX >> 3) + 1); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn rwlock_blocking_read_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); rt.block_on(async { let lock = RwLock::::new(0); let _ = lock.blocking_read(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn rwlock_blocking_write_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); rt.block_on(async { let lock = RwLock::::new(0); let _ = lock.blocking_write(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mpsc_bounded_channel_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let (_, _) = mpsc::channel::(0); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mpsc_bounded_receiver_blocking_recv_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); let (_tx, mut rx) = mpsc::channel::(1); rt.block_on(async { let _ = rx.blocking_recv(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mpsc_bounded_receiver_blocking_recv_many_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); let (_tx, mut rx) = mpsc::channel::(1); rt.block_on(async { let _ = rx.blocking_recv(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mpsc_bounded_sender_blocking_send_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); let (tx, _rx) = mpsc::channel::(1); rt.block_on(async { let _ = tx.blocking_send(3); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mpsc_unbounded_receiver_blocking_recv_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); let (_tx, mut rx) = mpsc::unbounded_channel::(); rt.block_on(async { let _ = rx.blocking_recv(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn mpsc_unbounded_receiver_blocking_recv_many_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); let (_tx, mut rx) = mpsc::unbounded_channel::(); let mut vec = vec![]; rt.block_on(async { let _ = rx.blocking_recv_many(&mut vec, 1); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn semaphore_merge_unrelated_owned_permits() -> Result<(), Box> { let panic_location_file = test_panic(|| { let sem1 = Arc::new(Semaphore::new(42)); let sem2 = Arc::new(Semaphore::new(42)); let mut p1 = sem1.try_acquire_owned().unwrap(); let p2 = sem2.try_acquire_owned().unwrap(); p1.merge(p2); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn semaphore_merge_unrelated_permits() -> Result<(), Box> { let panic_location_file = test_panic(|| { let sem1 = Semaphore::new(42); let sem2 = Semaphore::new(42); let mut p1 = sem1.try_acquire().unwrap(); let p2 = sem2.try_acquire().unwrap(); p1.merge(p2); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } fn current_thread() -> Runtime { Builder::new_current_thread().enable_all().build().unwrap() } tokio-1.43.1/tests/sync_rwlock.rs000064400000000000000000000223341046102023000151100ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use std::task::Poll; use futures::future::FutureExt; use tokio::sync::{RwLock, RwLockWriteGuard}; use tokio_test::task::spawn; use tokio_test::{assert_pending, assert_ready}; #[test] fn into_inner() { let rwlock = RwLock::new(42); assert_eq!(rwlock.into_inner(), 42); } // multiple reads should be Ready #[test] fn read_shared() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.read()); let _g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.read()); let _g2 = assert_ready!(t2.poll()); } // When there is an active shared owner, exclusive access should not be possible #[test] fn write_shared_pending() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.read()); let _g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.write()); assert_pending!(t2.poll()); } // When there is an active exclusive owner, subsequent exclusive access should not be possible #[test] fn read_exclusive_pending() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.write()); let _g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.read()); assert_pending!(t2.poll()); } // If the max shared access is reached and subsequent shared access is pending // should be made available when one of the shared accesses is dropped #[test] fn exhaust_reading() { let rwlock = RwLock::with_max_readers(100, 1024); let mut reads = Vec::new(); loop { let mut t = spawn(rwlock.read()); match t.poll() { Poll::Ready(guard) => reads.push(guard), Poll::Pending => break, } } let mut t1 = spawn(rwlock.read()); assert_pending!(t1.poll()); let g2 = reads.pop().unwrap(); drop(g2); assert!(t1.is_woken()); let _g1 = assert_ready!(t1.poll()); } // When there is an active exclusive owner, subsequent exclusive access should not be possible #[test] fn write_exclusive_pending() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.write()); let _g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.write()); assert_pending!(t2.poll()); } // When there is an active shared owner, exclusive access should be possible after shared is dropped #[test] fn write_shared_drop() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.read()); let g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.write()); assert_pending!(t2.poll()); drop(g1); assert!(t2.is_woken()); let _g2 = assert_ready!(t2.poll()); } // when there is an active shared owner, and exclusive access is triggered, // subsequent shared access should not be possible as write gathers all the available semaphore permits #[test] fn write_read_shared_pending() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.read()); let _g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.read()); let _g2 = assert_ready!(t2.poll()); let mut t3 = spawn(rwlock.write()); assert_pending!(t3.poll()); let mut t4 = spawn(rwlock.read()); assert_pending!(t4.poll()); } // when there is an active shared owner, and exclusive access is triggered, // reading should be possible after pending exclusive access is dropped #[test] fn write_read_shared_drop_pending() { let rwlock = RwLock::new(100); let mut t1 = spawn(rwlock.read()); let _g1 = assert_ready!(t1.poll()); let mut t2 = spawn(rwlock.write()); assert_pending!(t2.poll()); let mut t3 = spawn(rwlock.read()); assert_pending!(t3.poll()); drop(t2); assert!(t3.is_woken()); let _t3 = assert_ready!(t3.poll()); } // Acquire an RwLock nonexclusively by a single task #[maybe_tokio_test] async fn read_uncontested() { let rwlock = RwLock::new(100); let result = *rwlock.read().await; assert_eq!(result, 100); } // Acquire an uncontested RwLock in exclusive mode #[maybe_tokio_test] async fn write_uncontested() { let rwlock = RwLock::new(100); let mut result = rwlock.write().await; *result += 50; assert_eq!(*result, 150); } // RwLocks should be acquired in the order that their Futures are waited upon. #[maybe_tokio_test] async fn write_order() { let rwlock = RwLock::>::new(vec![]); let fut2 = rwlock.write().map(|mut guard| guard.push(2)); let fut1 = rwlock.write().map(|mut guard| guard.push(1)); fut1.await; fut2.await; let g = rwlock.read().await; assert_eq!(*g, vec![1, 2]); } // A single RwLock is contested by tasks in multiple threads #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads #[cfg_attr(miri, ignore)] // Too slow on miri. #[tokio::test(flavor = "multi_thread", worker_threads = 8)] async fn multithreaded() { use futures::stream::{self, StreamExt}; use std::sync::Arc; use tokio::sync::Barrier; let barrier = Arc::new(Barrier::new(5)); let rwlock = Arc::new(RwLock::::new(0)); let rwclone1 = rwlock.clone(); let rwclone2 = rwlock.clone(); let rwclone3 = rwlock.clone(); let rwclone4 = rwlock.clone(); let b1 = barrier.clone(); tokio::spawn(async move { stream::iter(0..1000) .for_each(move |_| { let rwlock = rwclone1.clone(); async move { let mut guard = rwlock.write().await; *guard += 2; } }) .await; b1.wait().await; }); let b2 = barrier.clone(); tokio::spawn(async move { stream::iter(0..1000) .for_each(move |_| { let rwlock = rwclone2.clone(); async move { let mut guard = rwlock.write().await; *guard += 3; } }) .await; b2.wait().await; }); let b3 = barrier.clone(); tokio::spawn(async move { stream::iter(0..1000) .for_each(move |_| { let rwlock = rwclone3.clone(); async move { let mut guard = rwlock.write().await; *guard += 5; } }) .await; b3.wait().await; }); let b4 = barrier.clone(); tokio::spawn(async move { stream::iter(0..1000) .for_each(move |_| { let rwlock = rwclone4.clone(); async move { let mut guard = rwlock.write().await; *guard += 7; } }) .await; b4.wait().await; }); barrier.wait().await; let g = rwlock.read().await; assert_eq!(*g, 17_000); } #[maybe_tokio_test] async fn try_write() { let lock = RwLock::new(0); let read_guard = lock.read().await; assert!(lock.try_write().is_err()); drop(read_guard); assert!(lock.try_write().is_ok()); } #[test] fn try_read_try_write() { let lock: RwLock = RwLock::new(15); { let rg1 = lock.try_read().unwrap(); assert_eq!(*rg1, 15); assert!(lock.try_write().is_err()); let rg2 = lock.try_read().unwrap(); assert_eq!(*rg2, 15) } { let mut wg = lock.try_write().unwrap(); *wg = 1515; assert!(lock.try_read().is_err()) } assert_eq!(*lock.try_read().unwrap(), 1515); } #[maybe_tokio_test] async fn downgrade_map() { let lock = RwLock::new(0); let write_guard = lock.write().await; let mut read_t = spawn(lock.read()); // We can't create a read when a write exists assert_pending!(read_t.poll()); // During the call to `f`, `read_t` doesn't have access yet. let read_guard1 = RwLockWriteGuard::downgrade_map(write_guard, |v| { assert_pending!(read_t.poll()); v }); // After the downgrade, `read_t` got the lock let read_guard2 = assert_ready!(read_t.poll()); // Ensure they're equal, as we return the original value assert_eq!(&*read_guard1 as *const _, &*read_guard2 as *const _); } #[maybe_tokio_test] async fn try_downgrade_map() { let lock = RwLock::new(0); let write_guard = lock.write().await; let mut read_t = spawn(lock.read()); // We can't create a read when a write exists assert_pending!(read_t.poll()); // During the call to `f`, `read_t` doesn't have access yet. let write_guard = RwLockWriteGuard::try_downgrade_map(write_guard, |_| { assert_pending!(read_t.poll()); None::<&()> }) .expect_err("downgrade didn't fail"); // After `f` returns `None`, `read_t` doesn't have access assert_pending!(read_t.poll()); // After `f` returns `Some`, `read_t` does have access let read_guard1 = RwLockWriteGuard::try_downgrade_map(write_guard, |v| Some(v)) .expect("downgrade didn't succeed"); let read_guard2 = assert_ready!(read_t.poll()); // Ensure they're equal, as we return the original value assert_eq!(&*read_guard1 as *const _, &*read_guard2 as *const _); } tokio-1.43.1/tests/sync_semaphore.rs000064400000000000000000000112011046102023000155610ustar 00000000000000#![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use std::sync::Arc; use tokio::sync::Semaphore; #[test] fn no_permits() { // this should not panic Semaphore::new(0); } #[test] fn try_acquire() { let sem = Semaphore::new(1); { let p1 = sem.try_acquire(); assert!(p1.is_ok()); let p2 = sem.try_acquire(); assert!(p2.is_err()); } let p3 = sem.try_acquire(); assert!(p3.is_ok()); } #[tokio::test] #[cfg(feature = "full")] async fn acquire() { let sem = Arc::new(Semaphore::new(1)); let p1 = sem.try_acquire().unwrap(); let sem_clone = sem.clone(); let j = tokio::spawn(async move { let _p2 = sem_clone.acquire().await; }); drop(p1); j.await.unwrap(); } #[tokio::test] #[cfg(feature = "full")] async fn add_permits() { let sem = Arc::new(Semaphore::new(0)); let sem_clone = sem.clone(); let j = tokio::spawn(async move { let _p2 = sem_clone.acquire().await; }); sem.add_permits(1); j.await.unwrap(); } #[test] fn forget() { let sem = Arc::new(Semaphore::new(1)); { let p = sem.try_acquire().unwrap(); assert_eq!(sem.available_permits(), 0); p.forget(); assert_eq!(sem.available_permits(), 0); } assert_eq!(sem.available_permits(), 0); assert!(sem.try_acquire().is_err()); } #[test] fn merge() { let sem = Arc::new(Semaphore::new(3)); { let mut p1 = sem.try_acquire().unwrap(); assert_eq!(sem.available_permits(), 2); let p2 = sem.try_acquire_many(2).unwrap(); assert_eq!(sem.available_permits(), 0); p1.merge(p2); assert_eq!(sem.available_permits(), 0); } assert_eq!(sem.available_permits(), 3); } #[test] #[cfg(not(target_family = "wasm"))] // No stack unwinding on wasm targets #[should_panic] fn merge_unrelated_permits() { let sem1 = Arc::new(Semaphore::new(3)); let sem2 = Arc::new(Semaphore::new(3)); let mut p1 = sem1.try_acquire().unwrap(); let p2 = sem2.try_acquire().unwrap(); p1.merge(p2); } #[test] fn split() { let sem = Semaphore::new(5); let mut p1 = sem.try_acquire_many(3).unwrap(); assert_eq!(sem.available_permits(), 2); assert_eq!(p1.num_permits(), 3); let mut p2 = p1.split(1).unwrap(); assert_eq!(sem.available_permits(), 2); assert_eq!(p1.num_permits(), 2); assert_eq!(p2.num_permits(), 1); let p3 = p1.split(0).unwrap(); assert_eq!(p3.num_permits(), 0); drop(p1); assert_eq!(sem.available_permits(), 4); let p4 = p2.split(1).unwrap(); assert_eq!(p2.num_permits(), 0); assert_eq!(p4.num_permits(), 1); assert!(p2.split(1).is_none()); drop(p2); assert_eq!(sem.available_permits(), 4); drop(p3); assert_eq!(sem.available_permits(), 4); drop(p4); assert_eq!(sem.available_permits(), 5); } #[tokio::test] #[cfg(feature = "full")] async fn stress_test() { let sem = Arc::new(Semaphore::new(5)); let mut join_handles = Vec::new(); for _ in 0..1000 { let sem_clone = sem.clone(); join_handles.push(tokio::spawn(async move { let _p = sem_clone.acquire().await; })); } for j in join_handles { j.await.unwrap(); } // there should be exactly 5 semaphores available now let _p1 = sem.try_acquire().unwrap(); let _p2 = sem.try_acquire().unwrap(); let _p3 = sem.try_acquire().unwrap(); let _p4 = sem.try_acquire().unwrap(); let _p5 = sem.try_acquire().unwrap(); assert!(sem.try_acquire().is_err()); } #[test] fn add_max_amount_permits() { let s = tokio::sync::Semaphore::new(0); s.add_permits(tokio::sync::Semaphore::MAX_PERMITS); assert_eq!(s.available_permits(), tokio::sync::Semaphore::MAX_PERMITS); } #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding #[test] #[should_panic] fn add_more_than_max_amount_permits1() { let s = tokio::sync::Semaphore::new(1); s.add_permits(tokio::sync::Semaphore::MAX_PERMITS); } #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding #[test] #[should_panic] fn add_more_than_max_amount_permits2() { let s = Semaphore::new(Semaphore::MAX_PERMITS - 1); s.add_permits(1); s.add_permits(1); } #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding #[test] #[should_panic] fn panic_when_exceeds_maxpermits() { let _ = Semaphore::new(Semaphore::MAX_PERMITS + 1); } #[test] fn no_panic_at_maxpermits() { let _ = Semaphore::new(Semaphore::MAX_PERMITS); let s = Semaphore::new(Semaphore::MAX_PERMITS - 1); s.add_permits(1); } tokio-1.43.1/tests/sync_semaphore_owned.rs000064400000000000000000000113571046102023000167710ustar 00000000000000#![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use std::sync::Arc; use tokio::sync::Semaphore; #[test] fn try_acquire() { let sem = Arc::new(Semaphore::new(1)); { let p1 = sem.clone().try_acquire_owned(); assert!(p1.is_ok()); let p2 = sem.clone().try_acquire_owned(); assert!(p2.is_err()); } let p3 = sem.try_acquire_owned(); assert!(p3.is_ok()); } #[test] fn try_acquire_many() { let sem = Arc::new(Semaphore::new(42)); { let p1 = sem.clone().try_acquire_many_owned(42); assert!(p1.is_ok()); let p2 = sem.clone().try_acquire_owned(); assert!(p2.is_err()); } let p3 = sem.clone().try_acquire_many_owned(32); assert!(p3.is_ok()); let p4 = sem.clone().try_acquire_many_owned(10); assert!(p4.is_ok()); assert!(sem.try_acquire_owned().is_err()); } #[tokio::test] #[cfg(feature = "full")] async fn acquire() { let sem = Arc::new(Semaphore::new(1)); let p1 = sem.clone().try_acquire_owned().unwrap(); let sem_clone = sem.clone(); let j = tokio::spawn(async move { let _p2 = sem_clone.acquire_owned().await; }); drop(p1); j.await.unwrap(); } #[tokio::test] #[cfg(feature = "full")] async fn acquire_many() { let semaphore = Arc::new(Semaphore::new(42)); let permit32 = semaphore.clone().try_acquire_many_owned(32).unwrap(); let (sender, receiver) = tokio::sync::oneshot::channel(); let join_handle = tokio::spawn(async move { let _permit10 = semaphore.clone().acquire_many_owned(10).await.unwrap(); sender.send(()).unwrap(); let _permit32 = semaphore.acquire_many_owned(32).await.unwrap(); }); receiver.await.unwrap(); drop(permit32); join_handle.await.unwrap(); } #[tokio::test] #[cfg(feature = "full")] async fn add_permits() { let sem = Arc::new(Semaphore::new(0)); let sem_clone = sem.clone(); let j = tokio::spawn(async move { let _p2 = sem_clone.acquire_owned().await; }); sem.add_permits(1); j.await.unwrap(); } #[test] fn forget() { let sem = Arc::new(Semaphore::new(1)); { let p = sem.clone().try_acquire_owned().unwrap(); assert_eq!(sem.available_permits(), 0); p.forget(); assert_eq!(sem.available_permits(), 0); } assert_eq!(sem.available_permits(), 0); assert!(sem.try_acquire_owned().is_err()); } #[test] fn merge() { let sem = Arc::new(Semaphore::new(3)); { let mut p1 = sem.clone().try_acquire_owned().unwrap(); assert_eq!(sem.available_permits(), 2); let p2 = sem.clone().try_acquire_many_owned(2).unwrap(); assert_eq!(sem.available_permits(), 0); p1.merge(p2); assert_eq!(sem.available_permits(), 0); } assert_eq!(sem.available_permits(), 3); } #[test] #[cfg(not(target_family = "wasm"))] // No stack unwinding on wasm targets #[should_panic] fn merge_unrelated_permits() { let sem1 = Arc::new(Semaphore::new(3)); let sem2 = Arc::new(Semaphore::new(3)); let mut p1 = sem1.try_acquire_owned().unwrap(); let p2 = sem2.try_acquire_owned().unwrap(); p1.merge(p2) } #[test] fn split() { let sem = Arc::new(Semaphore::new(5)); let mut p1 = sem.clone().try_acquire_many_owned(3).unwrap(); assert_eq!(sem.available_permits(), 2); assert_eq!(p1.num_permits(), 3); let mut p2 = p1.split(1).unwrap(); assert_eq!(sem.available_permits(), 2); assert_eq!(p1.num_permits(), 2); assert_eq!(p2.num_permits(), 1); let p3 = p1.split(0).unwrap(); assert_eq!(p3.num_permits(), 0); drop(p1); assert_eq!(sem.available_permits(), 4); let p4 = p2.split(1).unwrap(); assert_eq!(p2.num_permits(), 0); assert_eq!(p4.num_permits(), 1); assert!(p2.split(1).is_none()); drop(p2); assert_eq!(sem.available_permits(), 4); drop(p3); assert_eq!(sem.available_permits(), 4); drop(p4); assert_eq!(sem.available_permits(), 5); } #[tokio::test] #[cfg(feature = "full")] async fn stress_test() { let sem = Arc::new(Semaphore::new(5)); let mut join_handles = Vec::new(); for _ in 0..1000 { let sem_clone = sem.clone(); join_handles.push(tokio::spawn(async move { let _p = sem_clone.acquire_owned().await; })); } for j in join_handles { j.await.unwrap(); } // there should be exactly 5 semaphores available now let _p1 = sem.clone().try_acquire_owned().unwrap(); let _p2 = sem.clone().try_acquire_owned().unwrap(); let _p3 = sem.clone().try_acquire_owned().unwrap(); let _p4 = sem.clone().try_acquire_owned().unwrap(); let _p5 = sem.clone().try_acquire_owned().unwrap(); assert!(sem.try_acquire_owned().is_err()); } tokio-1.43.1/tests/sync_watch.rs000064400000000000000000000232221046102023000147120ustar 00000000000000#![allow(clippy::cognitive_complexity)] #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::watch; use tokio_test::task::spawn; use tokio_test::{ assert_pending, assert_ready, assert_ready_eq, assert_ready_err, assert_ready_ok, }; #[test] fn single_rx_recv() { let (tx, mut rx) = watch::channel("one"); { // Not initially notified let mut t = spawn(rx.changed()); assert_pending!(t.poll()); } assert_eq!(*rx.borrow(), "one"); { let mut t = spawn(rx.changed()); assert_pending!(t.poll()); tx.send("two").unwrap(); assert!(t.is_woken()); assert_ready_ok!(t.poll()); } assert_eq!(*rx.borrow(), "two"); { let mut t = spawn(rx.changed()); assert_pending!(t.poll()); drop(tx); assert!(t.is_woken()); assert_ready_err!(t.poll()); } assert_eq!(*rx.borrow(), "two"); } #[test] fn rx_version_underflow() { let (_tx, mut rx) = watch::channel("one"); // Version starts at 2, validate we do not underflow rx.mark_changed(); rx.mark_changed(); } #[test] fn rx_mark_changed() { let (tx, mut rx) = watch::channel("one"); let mut rx2 = rx.clone(); let mut rx3 = rx.clone(); let mut rx4 = rx.clone(); { rx.mark_changed(); assert!(rx.has_changed().unwrap()); let mut t = spawn(rx.changed()); assert_ready_ok!(t.poll()); } { assert!(!rx2.has_changed().unwrap()); let mut t = spawn(rx2.changed()); assert_pending!(t.poll()); } { rx3.mark_changed(); assert_eq!(*rx3.borrow(), "one"); assert!(rx3.has_changed().unwrap()); assert_eq!(*rx3.borrow_and_update(), "one"); assert!(!rx3.has_changed().unwrap()); let mut t = spawn(rx3.changed()); assert_pending!(t.poll()); } { tx.send("two").unwrap(); assert!(rx4.has_changed().unwrap()); assert_eq!(*rx4.borrow_and_update(), "two"); rx4.mark_changed(); assert!(rx4.has_changed().unwrap()); assert_eq!(*rx4.borrow_and_update(), "two") } assert_eq!(*rx.borrow(), "two"); } #[test] fn rx_mark_unchanged() { let (tx, mut rx) = watch::channel("one"); let mut rx2 = rx.clone(); { assert!(!rx.has_changed().unwrap()); rx.mark_changed(); assert!(rx.has_changed().unwrap()); rx.mark_unchanged(); assert!(!rx.has_changed().unwrap()); let mut t = spawn(rx.changed()); assert_pending!(t.poll()); } { assert!(!rx2.has_changed().unwrap()); tx.send("two").unwrap(); assert!(rx2.has_changed().unwrap()); rx2.mark_unchanged(); assert!(!rx2.has_changed().unwrap()); assert_eq!(*rx2.borrow_and_update(), "two"); } assert_eq!(*rx.borrow(), "two"); } #[test] fn multi_rx() { let (tx, mut rx1) = watch::channel("one"); let mut rx2 = rx1.clone(); { let mut t1 = spawn(rx1.changed()); let mut t2 = spawn(rx2.changed()); assert_pending!(t1.poll()); assert_pending!(t2.poll()); } assert_eq!(*rx1.borrow(), "one"); assert_eq!(*rx2.borrow(), "one"); let mut t2 = spawn(rx2.changed()); { let mut t1 = spawn(rx1.changed()); assert_pending!(t1.poll()); assert_pending!(t2.poll()); tx.send("two").unwrap(); assert!(t1.is_woken()); assert!(t2.is_woken()); assert_ready_ok!(t1.poll()); } assert_eq!(*rx1.borrow(), "two"); { let mut t1 = spawn(rx1.changed()); assert_pending!(t1.poll()); tx.send("three").unwrap(); assert!(t1.is_woken()); assert!(t2.is_woken()); assert_ready_ok!(t1.poll()); assert_ready_ok!(t2.poll()); } assert_eq!(*rx1.borrow(), "three"); drop(t2); assert_eq!(*rx2.borrow(), "three"); { let mut t1 = spawn(rx1.changed()); let mut t2 = spawn(rx2.changed()); assert_pending!(t1.poll()); assert_pending!(t2.poll()); tx.send("four").unwrap(); assert_ready_ok!(t1.poll()); assert_ready_ok!(t2.poll()); } assert_eq!(*rx1.borrow(), "four"); assert_eq!(*rx2.borrow(), "four"); } #[test] fn rx_observes_final_value() { // Initial value let (tx, mut rx) = watch::channel("one"); drop(tx); { let mut t1 = spawn(rx.changed()); assert_ready_err!(t1.poll()); } assert_eq!(*rx.borrow(), "one"); // Sending a value let (tx, mut rx) = watch::channel("one"); tx.send("two").unwrap(); { let mut t1 = spawn(rx.changed()); assert_ready_ok!(t1.poll()); } assert_eq!(*rx.borrow(), "two"); { let mut t1 = spawn(rx.changed()); assert_pending!(t1.poll()); tx.send("three").unwrap(); drop(tx); assert!(t1.is_woken()); assert_ready_ok!(t1.poll()); } assert_eq!(*rx.borrow(), "three"); { let mut t1 = spawn(rx.changed()); assert_ready_err!(t1.poll()); } assert_eq!(*rx.borrow(), "three"); } #[test] fn poll_close() { let (tx, rx) = watch::channel("one"); { let mut t = spawn(tx.closed()); assert_pending!(t.poll()); drop(rx); assert!(t.is_woken()); assert_ready!(t.poll()); } assert!(tx.send("two").is_err()); } #[test] fn borrow_and_update() { let (tx, mut rx) = watch::channel("one"); assert!(!rx.has_changed().unwrap()); tx.send("two").unwrap(); assert!(rx.has_changed().unwrap()); assert_ready!(spawn(rx.changed()).poll()).unwrap(); assert_pending!(spawn(rx.changed()).poll()); assert!(!rx.has_changed().unwrap()); tx.send("three").unwrap(); assert!(rx.has_changed().unwrap()); assert_eq!(*rx.borrow_and_update(), "three"); assert_pending!(spawn(rx.changed()).poll()); assert!(!rx.has_changed().unwrap()); drop(tx); assert_eq!(*rx.borrow_and_update(), "three"); assert_ready!(spawn(rx.changed()).poll()).unwrap_err(); assert!(rx.has_changed().is_err()); } #[test] fn reopened_after_subscribe() { let (tx, rx) = watch::channel("one"); assert!(!tx.is_closed()); drop(rx); assert!(tx.is_closed()); let rx = tx.subscribe(); assert!(!tx.is_closed()); drop(rx); assert!(tx.is_closed()); } #[test] #[cfg(panic = "unwind")] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn send_modify_panic() { let (tx, mut rx) = watch::channel("one"); tx.send_modify(|old| *old = "two"); assert_eq!(*rx.borrow_and_update(), "two"); let mut rx2 = rx.clone(); assert_eq!(*rx2.borrow_and_update(), "two"); let mut task = spawn(rx2.changed()); let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { tx.send_modify(|old| { *old = "panicked"; panic!(); }) })); assert!(result.is_err()); assert_pending!(task.poll()); assert_eq!(*rx.borrow(), "panicked"); tx.send_modify(|old| *old = "three"); assert_ready_ok!(task.poll()); assert_eq!(*rx.borrow_and_update(), "three"); } #[tokio::test] async fn multiple_sender() { let (tx1, mut rx) = watch::channel(0); let tx2 = tx1.clone(); let mut t = spawn(async { rx.changed().await.unwrap(); let v1 = *rx.borrow_and_update(); rx.changed().await.unwrap(); let v2 = *rx.borrow_and_update(); (v1, v2) }); tx1.send(1).unwrap(); assert_pending!(t.poll()); tx2.send(2).unwrap(); assert_ready_eq!(t.poll(), (1, 2)); } #[tokio::test] async fn receiver_is_notified_when_last_sender_is_dropped() { let (tx1, mut rx) = watch::channel(0); let tx2 = tx1.clone(); let mut t = spawn(rx.changed()); assert_pending!(t.poll()); drop(tx1); assert!(!t.is_woken()); drop(tx2); assert!(t.is_woken()); } #[tokio::test] async fn receiver_changed_is_cooperative() { let (tx, mut rx) = watch::channel(()); drop(tx); tokio::select! { biased; _ = async { loop { assert!(rx.changed().await.is_err()); } } => {}, _ = tokio::task::yield_now() => {}, } } #[tokio::test] async fn receiver_changed_is_cooperative_ok() { let (tx, mut rx) = watch::channel(()); tokio::select! { biased; _ = async { loop { assert!(tx.send(()).is_ok()); assert!(rx.changed().await.is_ok()); } } => {}, _ = tokio::task::yield_now() => {}, } } #[tokio::test] async fn receiver_wait_for_is_cooperative() { let (tx, mut rx) = watch::channel(0); drop(tx); tokio::select! { biased; _ = async { loop { assert!(rx.wait_for(|val| *val == 1).await.is_err()); } } => {}, _ = tokio::task::yield_now() => {}, } } #[tokio::test] async fn receiver_wait_for_is_cooperative_ok() { let (tx, mut rx) = watch::channel(0); tokio::select! { biased; _ = async { loop { assert!(tx.send(1).is_ok()); assert!(rx.wait_for(|val| *val == 1).await.is_ok()); } } => {}, _ = tokio::task::yield_now() => {}, } } #[tokio::test] async fn sender_closed_is_cooperative() { let (tx, rx) = watch::channel(()); drop(rx); tokio::select! { _ = async { loop { tx.closed().await; } } => {}, _ = tokio::task::yield_now() => {}, } } tokio-1.43.1/tests/task_abort.rs000064400000000000000000000240661046102023000147100ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery use std::sync::Arc; use std::thread::sleep; use tokio::time::Duration; use tokio::runtime::Builder; #[cfg(panic = "unwind")] struct PanicOnDrop; #[cfg(panic = "unwind")] impl Drop for PanicOnDrop { fn drop(&mut self) { panic!("Well what did you expect would happen..."); } } /// Checks that a suspended task can be aborted without panicking as reported in /// issue #3157: . #[test] fn test_abort_without_panic_3157() { let rt = Builder::new_multi_thread() .enable_time() .worker_threads(1) .build() .unwrap(); rt.block_on(async move { let handle = tokio::spawn(async move { tokio::time::sleep(Duration::new(100, 0)).await }); // wait for task to sleep. tokio::time::sleep(Duration::from_millis(10)).await; handle.abort(); let _ = handle.await; }); } /// Checks that a suspended task can be aborted inside of a current_thread /// executor without panicking as reported in issue #3662: /// . #[test] fn test_abort_without_panic_3662() { use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; struct DropCheck(Arc); impl Drop for DropCheck { fn drop(&mut self) { self.0.store(true, Ordering::SeqCst); } } let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async move { let drop_flag = Arc::new(AtomicBool::new(false)); let drop_check = DropCheck(drop_flag.clone()); let j = tokio::spawn(async move { // NB: just grab the drop check here so that it becomes part of the // task. let _drop_check = drop_check; futures::future::pending::<()>().await; }); let drop_flag2 = drop_flag.clone(); let task = std::thread::spawn(move || { // This runs in a separate thread so it doesn't have immediate // thread-local access to the executor. It does however transition // the underlying task to be completed, which will cause it to be // dropped (but not in this thread). assert!(!drop_flag2.load(Ordering::SeqCst)); j.abort(); j }) .join() .unwrap(); let result = task.await; assert!(drop_flag.load(Ordering::SeqCst)); assert!(result.unwrap_err().is_cancelled()); // Note: We do the following to trigger a deferred task cleanup. // // The relevant piece of code you want to look at is in: // `Inner::block_on` of `scheduler/current_thread.rs`. // // We cause the cleanup to happen by having a poll return Pending once // so that the scheduler can go into the "auxiliary tasks" mode, at // which point the task is removed from the scheduler. let i = tokio::spawn(async move { tokio::task::yield_now().await; }); i.await.unwrap(); }); } /// Checks that a suspended LocalSet task can be aborted from a remote thread /// without panicking and without running the tasks destructor on the wrong thread. /// #[test] fn remote_abort_local_set_3929() { struct DropCheck { created_on: std::thread::ThreadId, not_send: std::marker::PhantomData<*const ()>, } impl DropCheck { fn new() -> Self { Self { created_on: std::thread::current().id(), not_send: std::marker::PhantomData, } } } impl Drop for DropCheck { fn drop(&mut self) { if std::thread::current().id() != self.created_on { panic!("non-Send value dropped in another thread!"); } } } let rt = Builder::new_current_thread().build().unwrap(); let local = tokio::task::LocalSet::new(); let check = DropCheck::new(); let jh = local.spawn_local(async move { futures::future::pending::<()>().await; drop(check); }); let jh2 = std::thread::spawn(move || { sleep(Duration::from_millis(10)); jh.abort(); }); rt.block_on(local); jh2.join().unwrap(); } /// Checks that a suspended task can be aborted even if the `JoinHandle` is immediately dropped. /// issue #3964: . #[test] fn test_abort_wakes_task_3964() { let rt = Builder::new_current_thread().enable_time().build().unwrap(); rt.block_on(async move { let notify_dropped = Arc::new(()); let weak_notify_dropped = Arc::downgrade(¬ify_dropped); let handle = tokio::spawn(async move { // Make sure the Arc is moved into the task let _notify_dropped = notify_dropped; tokio::time::sleep(Duration::new(100, 0)).await }); // wait for task to sleep. tokio::time::sleep(Duration::from_millis(10)).await; handle.abort(); drop(handle); // wait for task to abort. tokio::time::sleep(Duration::from_millis(10)).await; // Check that the Arc has been dropped. assert!(weak_notify_dropped.upgrade().is_none()); }); } /// Checks that aborting a task whose destructor panics does not allow the /// panic to escape the task. #[test] #[cfg(panic = "unwind")] fn test_abort_task_that_panics_on_drop_contained() { let rt = Builder::new_current_thread().enable_time().build().unwrap(); rt.block_on(async move { let handle = tokio::spawn(async move { // Make sure the Arc is moved into the task let _panic_dropped = PanicOnDrop; tokio::time::sleep(Duration::new(100, 0)).await }); // wait for task to sleep. tokio::time::sleep(Duration::from_millis(10)).await; handle.abort(); drop(handle); // wait for task to abort. tokio::time::sleep(Duration::from_millis(10)).await; }); } /// Checks that aborting a task whose destructor panics has the expected result. #[test] #[cfg(panic = "unwind")] fn test_abort_task_that_panics_on_drop_returned() { let rt = Builder::new_current_thread().enable_time().build().unwrap(); rt.block_on(async move { let handle = tokio::spawn(async move { // Make sure the Arc is moved into the task let _panic_dropped = PanicOnDrop; tokio::time::sleep(Duration::new(100, 0)).await }); // wait for task to sleep. tokio::time::sleep(Duration::from_millis(10)).await; handle.abort(); assert!(handle.await.unwrap_err().is_panic()); }); } // It's not clear where these tests belong. This was the place suggested by @Darksonn: // https://github.com/tokio-rs/tokio/pull/6753#issuecomment-2271434176 /// Checks that a `JoinError` with a panic payload prints the expected text. #[test] #[cfg(panic = "unwind")] fn test_join_error_display() { let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async move { // `String` payload let join_err = tokio::spawn(async move { let value = 1234; panic!("Format-args payload: {value}") }) .await .unwrap_err(); // We can't assert the full output because the task ID can change. let join_err_str = join_err.to_string(); assert!( join_err_str.starts_with("task ") && join_err_str.ends_with(" panicked with message \"Format-args payload: 1234\""), "Unexpected join_err_str {join_err_str:?}" ); // `&'static str` payload let join_err = tokio::spawn(async move { panic!("Const payload") }) .await .unwrap_err(); let join_err_str = join_err.to_string(); assert!( join_err_str.starts_with("task ") && join_err_str.ends_with(" panicked with message \"Const payload\""), "Unexpected join_err_str {join_err_str:?}" ); // Non-string payload let join_err = tokio::spawn(async move { std::panic::panic_any(1234i32) }) .await .unwrap_err(); let join_err_str = join_err.to_string(); assert!( join_err_str.starts_with("task ") && join_err_str.ends_with(" panicked"), "Unexpected join_err_str {join_err_str:?}" ); }); } /// Checks that a `JoinError` with a panic payload prints the expected text from `Debug`. #[test] #[cfg(panic = "unwind")] fn test_join_error_debug() { let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async move { // `String` payload let join_err = tokio::spawn(async move { let value = 1234; panic!("Format-args payload: {value}") }) .await .unwrap_err(); // We can't assert the full output because the task ID can change. let join_err_str = format!("{join_err:?}"); assert!( join_err_str.starts_with("JoinError::Panic(Id(") && join_err_str.ends_with("), \"Format-args payload: 1234\", ...)"), "Unexpected join_err_str {join_err_str:?}" ); // `&'static str` payload let join_err = tokio::spawn(async move { panic!("Const payload") }) .await .unwrap_err(); let join_err_str = format!("{join_err:?}"); assert!( join_err_str.starts_with("JoinError::Panic(Id(") && join_err_str.ends_with("), \"Const payload\", ...)"), "Unexpected join_err_str {join_err_str:?}" ); // Non-string payload let join_err = tokio::spawn(async move { std::panic::panic_any(1234i32) }) .await .unwrap_err(); let join_err_str = format!("{join_err:?}"); assert!( join_err_str.starts_with("JoinError::Panic(Id(") && join_err_str.ends_with("), ...)"), "Unexpected join_err_str {join_err_str:?}" ); }); } tokio-1.43.1/tests/task_blocking.rs000064400000000000000000000205141046102023000153630ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support threads use tokio::{runtime, task, time}; use tokio_test::assert_ok; use std::thread; use std::time::Duration; mod support { pub(crate) mod mpsc_stream; } #[tokio::test] async fn basic_blocking() { // Run a few times for _ in 0..100 { let out = assert_ok!( tokio::spawn(async { assert_ok!( task::spawn_blocking(|| { thread::sleep(Duration::from_millis(5)); "hello" }) .await ) }) .await ); assert_eq!(out, "hello"); } } #[tokio::test(flavor = "multi_thread")] async fn block_in_blocking() { // Run a few times for _ in 0..100 { let out = assert_ok!( tokio::spawn(async { assert_ok!( task::spawn_blocking(|| { task::block_in_place(|| { thread::sleep(Duration::from_millis(5)); }); "hello" }) .await ) }) .await ); assert_eq!(out, "hello"); } } #[tokio::test(flavor = "multi_thread")] async fn block_in_block() { // Run a few times for _ in 0..100 { let out = assert_ok!( tokio::spawn(async { task::block_in_place(|| { task::block_in_place(|| { thread::sleep(Duration::from_millis(5)); }); "hello" }) }) .await ); assert_eq!(out, "hello"); } } #[tokio::test(flavor = "current_thread")] #[should_panic] async fn no_block_in_current_thread_scheduler() { task::block_in_place(|| {}); } #[test] fn yes_block_in_threaded_block_on() { let rt = runtime::Runtime::new().unwrap(); rt.block_on(async { task::block_in_place(|| {}); }); } #[test] #[should_panic] fn no_block_in_current_thread_block_on() { let rt = runtime::Builder::new_current_thread().build().unwrap(); rt.block_on(async { task::block_in_place(|| {}); }); } #[test] fn can_enter_current_thread_rt_from_within_block_in_place() { let outer = tokio::runtime::Runtime::new().unwrap(); outer.block_on(async { tokio::task::block_in_place(|| { let inner = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); inner.block_on(async {}) }) }); } #[test] #[cfg(panic = "unwind")] fn useful_panic_message_when_dropping_rt_in_rt() { use std::panic::{catch_unwind, AssertUnwindSafe}; let outer = tokio::runtime::Runtime::new().unwrap(); let result = catch_unwind(AssertUnwindSafe(|| { outer.block_on(async { let _ = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); }); })); assert!(result.is_err()); let err = result.unwrap_err(); let err: &'static str = err.downcast_ref::<&'static str>().unwrap(); assert!( err.contains("Cannot drop a runtime"), "Wrong panic message: {err:?}" ); } #[test] fn can_shutdown_with_zero_timeout_in_runtime() { let outer = tokio::runtime::Runtime::new().unwrap(); outer.block_on(async { let rt = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); rt.shutdown_timeout(Duration::from_nanos(0)); }); } #[test] fn can_shutdown_now_in_runtime() { let outer = tokio::runtime::Runtime::new().unwrap(); outer.block_on(async { let rt = tokio::runtime::Builder::new_current_thread() .build() .unwrap(); rt.shutdown_background(); }); } #[test] fn coop_disabled_in_block_in_place() { let outer = tokio::runtime::Builder::new_multi_thread() .enable_time() .build() .unwrap(); let (tx, rx) = support::mpsc_stream::unbounded_channel_stream(); for i in 0..200 { tx.send(i).unwrap(); } drop(tx); outer.block_on(async move { let jh = tokio::spawn(async move { tokio::task::block_in_place(move || { futures::executor::block_on(async move { use tokio_stream::StreamExt; assert_eq!(rx.fold(0, |n, _| n + 1).await, 200); }) }) }); tokio::time::timeout(Duration::from_secs(1), jh) .await .expect("timed out (probably hanging)") .unwrap() }); } #[test] fn coop_disabled_in_block_in_place_in_block_on() { let (done_tx, done_rx) = std::sync::mpsc::channel(); let done = done_tx.clone(); thread::spawn(move || { let outer = tokio::runtime::Runtime::new().unwrap(); let (tx, rx) = support::mpsc_stream::unbounded_channel_stream(); for i in 0..200 { tx.send(i).unwrap(); } drop(tx); outer.block_on(async move { tokio::task::block_in_place(move || { futures::executor::block_on(async move { use tokio_stream::StreamExt; assert_eq!(rx.fold(0, |n, _| n + 1).await, 200); }) }) }); let _ = done.send(Ok(())); }); thread::spawn(move || { thread::sleep(Duration::from_secs(1)); let _ = done_tx.send(Err("timed out (probably hanging)")); }); done_rx.recv().unwrap().unwrap(); } #[cfg(feature = "test-util")] #[tokio::test(start_paused = true)] async fn blocking_when_paused() { // Do not auto-advance time when we have started a blocking task that has // not yet finished. time::timeout( Duration::from_secs(3), task::spawn_blocking(|| thread::sleep(Duration::from_millis(1))), ) .await .expect("timeout should not trigger") .expect("blocking task should finish"); // Really: Do not auto-advance time, even if the timeout is short and the // blocking task runs for longer than that. It doesn't matter: Tokio time // is paused; system time is not. time::timeout( Duration::from_millis(1), task::spawn_blocking(|| thread::sleep(Duration::from_millis(50))), ) .await .expect("timeout should not trigger") .expect("blocking task should finish"); } #[cfg(feature = "test-util")] #[tokio::test(start_paused = true)] async fn blocking_task_wakes_paused_runtime() { let t0 = std::time::Instant::now(); time::timeout( Duration::from_secs(15), task::spawn_blocking(|| thread::sleep(Duration::from_millis(1))), ) .await .expect("timeout should not trigger") .expect("blocking task should finish"); assert!( t0.elapsed() < Duration::from_secs(10), "completing a spawn_blocking should wake the scheduler if it's parked while time is paused" ); } #[cfg(feature = "test-util")] #[tokio::test(start_paused = true)] async fn unawaited_blocking_task_wakes_paused_runtime() { let t0 = std::time::Instant::now(); // When this task finishes, time should auto-advance, even though the // JoinHandle has not been awaited yet. let a = task::spawn_blocking(|| { thread::sleep(Duration::from_millis(1)); }); crate::time::sleep(Duration::from_secs(15)).await; a.await.expect("blocking task should finish"); assert!( t0.elapsed() < Duration::from_secs(10), "completing a spawn_blocking should wake the scheduler if it's parked while time is paused" ); } #[cfg(panic = "unwind")] #[cfg(feature = "test-util")] #[tokio::test(start_paused = true)] async fn panicking_blocking_task_wakes_paused_runtime() { let t0 = std::time::Instant::now(); let result = time::timeout( Duration::from_secs(15), task::spawn_blocking(|| { thread::sleep(Duration::from_millis(1)); panic!("blocking task panicked"); }), ) .await .expect("timeout should not trigger"); assert!(result.is_err(), "blocking task should have panicked"); assert!( t0.elapsed() < Duration::from_secs(10), "completing a spawn_blocking should wake the scheduler if it's parked while time is paused" ); } tokio-1.43.1/tests/task_builder.rs000064400000000000000000000034221046102023000152200ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![cfg(all(tokio_unstable, feature = "tracing"))] use std::rc::Rc; use tokio::{ task::{Builder, LocalSet}, test, }; #[test] async fn spawn_with_name() { let result = Builder::new() .name("name") .spawn(async { "task executed" }) .unwrap() .await; assert_eq!(result.unwrap(), "task executed"); } #[test] async fn spawn_blocking_with_name() { let result = Builder::new() .name("name") .spawn_blocking(|| "task executed") .unwrap() .await; assert_eq!(result.unwrap(), "task executed"); } #[test] async fn spawn_local_with_name() { let unsend_data = Rc::new("task executed"); let result = LocalSet::new() .run_until(async move { Builder::new() .name("name") .spawn_local(async move { unsend_data }) .unwrap() .await }) .await; assert_eq!(*result.unwrap(), "task executed"); } #[test] async fn spawn_without_name() { let result = Builder::new() .spawn(async { "task executed" }) .unwrap() .await; assert_eq!(result.unwrap(), "task executed"); } #[test] async fn spawn_blocking_without_name() { let result = Builder::new() .spawn_blocking(|| "task executed") .unwrap() .await; assert_eq!(result.unwrap(), "task executed"); } #[test] async fn spawn_local_without_name() { let unsend_data = Rc::new("task executed"); let result = LocalSet::new() .run_until(async move { Builder::new() .spawn_local(async move { unsend_data }) .unwrap() .await }) .await; assert_eq!(*result.unwrap(), "task executed"); } tokio-1.43.1/tests/task_hooks.rs000064400000000000000000000040431046102023000147150ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(feature = "full", tokio_unstable, target_has_atomic = "64"))] use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use tokio::runtime::Builder; const TASKS: usize = 8; const ITERATIONS: usize = 64; /// Assert that the spawn task hook always fires when set. #[test] fn spawn_task_hook_fires() { let count = Arc::new(AtomicUsize::new(0)); let count2 = Arc::clone(&count); let ids = Arc::new(Mutex::new(HashSet::new())); let ids2 = Arc::clone(&ids); let runtime = Builder::new_current_thread() .on_task_spawn(move |data| { ids2.lock().unwrap().insert(data.id()); count2.fetch_add(1, Ordering::SeqCst); }) .build() .unwrap(); for _ in 0..TASKS { runtime.spawn(std::future::pending::<()>()); } let count_realized = count.load(Ordering::SeqCst); assert_eq!( TASKS, count_realized, "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}", count_realized ); let count_ids_realized = ids.lock().unwrap().len(); assert_eq!( TASKS, count_ids_realized, "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}", count_realized ); } /// Assert that the terminate task hook always fires when set. #[test] fn terminate_task_hook_fires() { let count = Arc::new(AtomicUsize::new(0)); let count2 = Arc::clone(&count); let runtime = Builder::new_current_thread() .on_task_terminate(move |_data| { count2.fetch_add(1, Ordering::SeqCst); }) .build() .unwrap(); for _ in 0..TASKS { runtime.spawn(std::future::ready(())); } runtime.block_on(async { // tick the runtime a bunch to close out tasks for _ in 0..ITERATIONS { tokio::task::yield_now().await; } }); assert_eq!(TASKS, count.load(Ordering::SeqCst)); } tokio-1.43.1/tests/task_id.rs000064400000000000000000000174361046102023000142000ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::error::Error; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::runtime::Runtime; use tokio::sync::oneshot; use tokio::task::{self, Id, LocalSet}; mod support { pub mod panic; } use support::panic::test_panic; #[tokio::test(flavor = "current_thread")] async fn task_id_spawn() { tokio::spawn(async { println!("task id: {}", task::id()) }) .await .unwrap(); } #[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "current_thread")] async fn task_id_spawn_blocking() { task::spawn_blocking(|| println!("task id: {}", task::id())) .await .unwrap(); } #[tokio::test(flavor = "current_thread")] async fn task_id_collision_current_thread() { let handle1 = tokio::spawn(async { task::id() }); let handle2 = tokio::spawn(async { task::id() }); let (id1, id2) = tokio::join!(handle1, handle2); assert_ne!(id1.unwrap(), id2.unwrap()); } #[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_collision_multi_thread() { let handle1 = tokio::spawn(async { task::id() }); let handle2 = tokio::spawn(async { task::id() }); let (id1, id2) = tokio::join!(handle1, handle2); assert_ne!(id1.unwrap(), id2.unwrap()); } #[tokio::test(flavor = "current_thread")] async fn task_ids_match_current_thread() { let (tx, rx) = oneshot::channel(); let handle = tokio::spawn(async { let id = rx.await.unwrap(); assert_eq!(id, task::id()); }); tx.send(handle.id()).unwrap(); handle.await.unwrap(); } #[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_ids_match_multi_thread() { let (tx, rx) = oneshot::channel(); let handle = tokio::spawn(async { let id = rx.await.unwrap(); assert_eq!(id, task::id()); }); tx.send(handle.id()).unwrap(); handle.await.unwrap(); } #[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_future_destructor_completion() { struct MyFuture { tx: Option>, } impl Future for MyFuture { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Ready(()) } } impl Drop for MyFuture { fn drop(&mut self) { let _ = self.tx.take().unwrap().send(task::id()); } } let (tx, rx) = oneshot::channel(); let handle = tokio::spawn(MyFuture { tx: Some(tx) }); let id = handle.id(); handle.await.unwrap(); assert_eq!(rx.await.unwrap(), id); } #[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_future_destructor_abort() { struct MyFuture { tx: Option>, } impl Future for MyFuture { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Pending } } impl Drop for MyFuture { fn drop(&mut self) { let _ = self.tx.take().unwrap().send(task::id()); } } let (tx, rx) = oneshot::channel(); let handle = tokio::spawn(MyFuture { tx: Some(tx) }); let id = handle.id(); handle.abort(); assert!(handle.await.unwrap_err().is_cancelled()); assert_eq!(rx.await.unwrap(), id); } #[tokio::test(flavor = "current_thread")] async fn task_id_output_destructor_handle_dropped_before_completion() { struct MyOutput { tx: Option>, } impl Drop for MyOutput { fn drop(&mut self) { let _ = self.tx.take().unwrap().send(task::id()); } } struct MyFuture { tx: Option>, } impl Future for MyFuture { type Output = MyOutput; fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { Poll::Ready(MyOutput { tx: self.tx.take() }) } } let (tx, mut rx) = oneshot::channel(); let handle = tokio::spawn(MyFuture { tx: Some(tx) }); let id = handle.id(); drop(handle); assert!(rx.try_recv().is_err()); assert_eq!(rx.await.unwrap(), id); } #[tokio::test(flavor = "current_thread")] async fn task_id_output_destructor_handle_dropped_after_completion() { struct MyOutput { tx: Option>, } impl Drop for MyOutput { fn drop(&mut self) { let _ = self.tx.take().unwrap().send(task::id()); } } struct MyFuture { tx_output: Option>, tx_future: Option>, } impl Future for MyFuture { type Output = MyOutput; fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { let _ = self.tx_future.take().unwrap().send(()); Poll::Ready(MyOutput { tx: self.tx_output.take(), }) } } let (tx_output, mut rx_output) = oneshot::channel(); let (tx_future, rx_future) = oneshot::channel(); let handle = tokio::spawn(MyFuture { tx_output: Some(tx_output), tx_future: Some(tx_future), }); let id = handle.id(); rx_future.await.unwrap(); assert!(rx_output.try_recv().is_err()); drop(handle); assert_eq!(rx_output.await.unwrap(), id); } #[test] fn task_try_id_outside_task() { assert_eq!(None, task::try_id()); } #[cfg(not(target_os = "wasi"))] #[test] fn task_try_id_inside_block_on() { let rt = Runtime::new().unwrap(); rt.block_on(async { assert_eq!(None, task::try_id()); }); } #[tokio::test(flavor = "current_thread")] async fn task_id_spawn_local() { LocalSet::new() .run_until(async { task::spawn_local(async { println!("task id: {}", task::id()) }) .await .unwrap(); }) .await } #[tokio::test(flavor = "current_thread")] async fn task_id_nested_spawn_local() { LocalSet::new() .run_until(async { task::spawn_local(async { let parent_id = task::id(); LocalSet::new() .run_until(async { task::spawn_local(async move { assert_ne!(parent_id, task::id()); }) .await .unwrap(); }) .await; assert_eq!(parent_id, task::id()); }) .await .unwrap(); }) .await; } #[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_block_in_place_block_on_spawn() { use tokio::runtime::Builder; task::spawn(async { let parent_id = task::id(); task::block_in_place(move || { let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(rt.spawn(async move { assert_ne!(parent_id, task::id()); })) .unwrap(); }); assert_eq!(parent_id, task::id()); }) .await .unwrap(); } #[test] #[cfg_attr(not(panic = "unwind"), ignore)] fn task_id_outside_task_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = task::id(); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] #[cfg_attr(not(panic = "unwind"), ignore)] fn task_id_inside_block_on_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = Runtime::new().unwrap(); rt.block_on(async { task::id(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } tokio-1.43.1/tests/task_join_set.rs000064400000000000000000000200411046102023000154000ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(feature = "full")] use futures::future::FutureExt; use tokio::sync::oneshot; use tokio::task::JoinSet; use tokio::time::Duration; fn rt() -> tokio::runtime::Runtime { tokio::runtime::Builder::new_current_thread() .build() .unwrap() } #[tokio::test(start_paused = true)] async fn test_with_sleep() { let mut set = JoinSet::new(); for i in 0..10 { set.spawn(async move { i }); assert_eq!(set.len(), 1 + i); } set.detach_all(); assert_eq!(set.len(), 0); assert!(set.join_next().await.is_none()); for i in 0..10 { set.spawn(async move { tokio::time::sleep(Duration::from_secs(i as u64)).await; i }); assert_eq!(set.len(), 1 + i); } let mut seen = [false; 10]; while let Some(res) = set.join_next().await.transpose().unwrap() { seen[res] = true; } for was_seen in &seen { assert!(was_seen); } assert!(set.join_next().await.is_none()); // Do it again. for i in 0..10 { set.spawn(async move { tokio::time::sleep(Duration::from_secs(i as u64)).await; i }); } let mut seen = [false; 10]; while let Some(res) = set.join_next().await.transpose().unwrap() { seen[res] = true; } for was_seen in &seen { assert!(was_seen); } assert!(set.join_next().await.is_none()); } #[tokio::test] async fn test_abort_on_drop() { let mut set = JoinSet::new(); let mut recvs = Vec::new(); for _ in 0..16 { let (send, recv) = oneshot::channel::<()>(); recvs.push(recv); set.spawn(async { // This task will never complete on its own. futures::future::pending::<()>().await; drop(send); }); } drop(set); for recv in recvs { // The task is aborted soon and we will receive an error. assert!(recv.await.is_err()); } } #[tokio::test] async fn alternating() { let mut set = JoinSet::new(); assert_eq!(set.len(), 0); set.spawn(async {}); assert_eq!(set.len(), 1); set.spawn(async {}); assert_eq!(set.len(), 2); for _ in 0..16 { let () = set.join_next().await.unwrap().unwrap(); assert_eq!(set.len(), 1); set.spawn(async {}); assert_eq!(set.len(), 2); } } #[tokio::test(start_paused = true)] async fn abort_tasks() { let mut set = JoinSet::new(); let mut num_canceled = 0; let mut num_completed = 0; for i in 0..16 { let abort = set.spawn(async move { tokio::time::sleep(Duration::from_secs(i as u64)).await; i }); if i % 2 != 0 { // abort odd-numbered tasks. abort.abort(); } } loop { match set.join_next().await { Some(Ok(res)) => { num_completed += 1; assert_eq!(res % 2, 0); } Some(Err(e)) => { assert!(e.is_cancelled()); num_canceled += 1; } None => break, } } assert_eq!(num_canceled, 8); assert_eq!(num_completed, 8); } #[test] fn runtime_gone() { let mut set = JoinSet::new(); { let rt = rt(); set.spawn_on(async { 1 }, rt.handle()); drop(rt); } assert!(rt() .block_on(set.join_next()) .unwrap() .unwrap_err() .is_cancelled()); } #[tokio::test] async fn join_all() { let mut set: JoinSet = JoinSet::new(); for _ in 0..5 { set.spawn(async { 1 }); } let res: Vec = set.join_all().await; assert_eq!(res.len(), 5); for itm in res.into_iter() { assert_eq!(itm, 1) } } #[cfg(panic = "unwind")] #[tokio::test(start_paused = true)] async fn task_panics() { let mut set: JoinSet<()> = JoinSet::new(); let (tx, mut rx) = oneshot::channel(); assert_eq!(set.len(), 0); set.spawn(async move { tokio::time::sleep(Duration::from_secs(2)).await; tx.send(()).unwrap(); }); assert_eq!(set.len(), 1); set.spawn(async { tokio::time::sleep(Duration::from_secs(1)).await; panic!(); }); assert_eq!(set.len(), 2); let panic = tokio::spawn(set.join_all()).await.unwrap_err(); assert!(rx.try_recv().is_err()); assert!(panic.is_panic()); } #[tokio::test(start_paused = true)] async fn abort_all() { let mut set: JoinSet<()> = JoinSet::new(); for _ in 0..5 { set.spawn(futures::future::pending()); } for _ in 0..5 { set.spawn(async { tokio::time::sleep(Duration::from_secs(1)).await; }); } // The join set will now have 5 pending tasks and 5 ready tasks. tokio::time::sleep(Duration::from_secs(2)).await; set.abort_all(); assert_eq!(set.len(), 10); let mut count = 0; while let Some(res) = set.join_next().await { if let Err(err) = res { assert!(err.is_cancelled()); } count += 1; } assert_eq!(count, 10); assert_eq!(set.len(), 0); } // This ensures that `join_next` works correctly when the coop budget is // exhausted. #[tokio::test(flavor = "current_thread")] async fn join_set_coop() { // Large enough to trigger coop. const TASK_NUM: u32 = 1000; static SEM: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(0); let mut set = JoinSet::new(); for _ in 0..TASK_NUM { set.spawn(async { SEM.add_permits(1); }); } // Wait for all tasks to complete. // // Since this is a `current_thread` runtime, there's no race condition // between the last permit being added and the task completing. let _ = SEM.acquire_many(TASK_NUM).await.unwrap(); let mut count = 0; let mut coop_count = 0; loop { match set.join_next().now_or_never() { Some(Some(Ok(()))) => {} Some(Some(Err(err))) => panic!("failed: {err}"), None => { coop_count += 1; tokio::task::yield_now().await; continue; } Some(None) => break, } count += 1; } assert!(coop_count >= 1); assert_eq!(count, TASK_NUM); } #[tokio::test(flavor = "current_thread")] async fn try_join_next() { const TASK_NUM: u32 = 1000; let (send, recv) = tokio::sync::watch::channel(()); let mut set = JoinSet::new(); for _ in 0..TASK_NUM { let mut recv = recv.clone(); set.spawn(async move { recv.changed().await.unwrap() }); } drop(recv); assert!(set.try_join_next().is_none()); send.send_replace(()); send.closed().await; let mut count = 0; loop { match set.try_join_next() { Some(Ok(())) => { count += 1; } Some(Err(err)) => panic!("failed: {err}"), None => { break; } } } assert_eq!(count, TASK_NUM); } #[cfg(tokio_unstable)] #[tokio::test(flavor = "current_thread")] async fn try_join_next_with_id() { const TASK_NUM: u32 = 1000; let (send, recv) = tokio::sync::watch::channel(()); let mut set = JoinSet::new(); let mut spawned = std::collections::HashSet::with_capacity(TASK_NUM as usize); for _ in 0..TASK_NUM { let mut recv = recv.clone(); let handle = set.spawn(async move { recv.changed().await.unwrap() }); spawned.insert(handle.id()); } drop(recv); assert!(set.try_join_next_with_id().is_none()); send.send_replace(()); send.closed().await; let mut count = 0; let mut joined = std::collections::HashSet::with_capacity(TASK_NUM as usize); loop { match set.try_join_next_with_id() { Some(Ok((id, ()))) => { count += 1; joined.insert(id); } Some(Err(err)) => panic!("failed: {}", err), None => { break; } } } assert_eq!(count, TASK_NUM); assert_eq!(joined, spawned); } tokio-1.43.1/tests/task_local.rs000064400000000000000000000067371046102023000147000ustar 00000000000000#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::sync::oneshot; #[tokio::test(flavor = "multi_thread")] async fn local() { tokio::task_local! { static REQ_ID: u32; pub static FOO: bool; } let j1 = tokio::spawn(REQ_ID.scope(1, async move { assert_eq!(REQ_ID.get(), 1); assert_eq!(REQ_ID.get(), 1); })); let j2 = tokio::spawn(REQ_ID.scope(2, async move { REQ_ID.with(|v| { assert_eq!(REQ_ID.get(), 2); assert_eq!(*v, 2); }); tokio::time::sleep(std::time::Duration::from_millis(10)).await; assert_eq!(REQ_ID.get(), 2); })); let j3 = tokio::spawn(FOO.scope(true, async move { assert!(FOO.get()); })); j1.await.unwrap(); j2.await.unwrap(); j3.await.unwrap(); } #[tokio::test] async fn task_local_available_on_abort() { tokio::task_local! { static KEY: u32; } struct MyFuture { tx_poll: Option>, tx_drop: Option>, } impl Future for MyFuture { type Output = (); fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { if let Some(tx_poll) = self.tx_poll.take() { let _ = tx_poll.send(()); } Poll::Pending } } impl Drop for MyFuture { fn drop(&mut self) { let _ = self.tx_drop.take().unwrap().send(KEY.get()); } } let (tx_drop, rx_drop) = oneshot::channel(); let (tx_poll, rx_poll) = oneshot::channel(); let h = tokio::spawn(KEY.scope( 42, MyFuture { tx_poll: Some(tx_poll), tx_drop: Some(tx_drop), }, )); rx_poll.await.unwrap(); h.abort(); assert_eq!(rx_drop.await.unwrap(), 42); let err = h.await.unwrap_err(); if !err.is_cancelled() { if let Ok(panic) = err.try_into_panic() { std::panic::resume_unwind(panic); } else { panic!(); } } } #[tokio::test] async fn task_local_available_on_completion_drop() { tokio::task_local! { static KEY: u32; } struct MyFuture { tx: Option>, } impl Future for MyFuture { type Output = (); fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { Poll::Ready(()) } } impl Drop for MyFuture { fn drop(&mut self) { let _ = self.tx.take().unwrap().send(KEY.get()); } } let (tx, rx) = oneshot::channel(); let h = tokio::spawn(KEY.scope(42, MyFuture { tx: Some(tx) })); assert_eq!(rx.await.unwrap(), 42); h.await.unwrap(); } #[tokio::test] async fn take_value() { tokio::task_local! { static KEY: u32 } let fut = KEY.scope(1, async {}); let mut pinned = Box::pin(fut); assert_eq!(pinned.as_mut().take_value(), Some(1)); assert_eq!(pinned.as_mut().take_value(), None); } #[tokio::test] async fn poll_after_take_value_should_fail() { tokio::task_local! { static KEY: u32 } let fut = KEY.scope(1, async { let result = KEY.try_with(|_| {}); // The task local value no longer exists. assert!(result.is_err()); }); let mut fut = Box::pin(fut); fut.as_mut().take_value(); // Poll the future after `take_value` has been called fut.await; } tokio-1.43.1/tests/task_local_set.rs000064400000000000000000000510301046102023000155350ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(feature = "full")] use futures::{ future::{pending, ready}, FutureExt, }; use tokio::runtime; use tokio::sync::{mpsc, oneshot}; use tokio::task::{self, LocalSet}; use tokio::time; #[cfg(not(target_os = "wasi"))] use std::cell::Cell; use std::sync::atomic::AtomicBool; #[cfg(not(target_os = "wasi"))] use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; #[cfg(not(target_os = "wasi"))] use std::sync::atomic::Ordering::SeqCst; use std::time::Duration; #[tokio::test(flavor = "current_thread")] async fn local_current_thread_scheduler() { LocalSet::new() .run_until(async { task::spawn_local(async {}).await.unwrap(); }) .await; } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn local_threadpool() { thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); LocalSet::new() .run_until(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::spawn_local(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); }) .await .unwrap(); }) .await; } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn localset_future_threadpool() { thread_local! { static ON_LOCAL_THREAD: Cell = const { Cell::new(false) }; } ON_LOCAL_THREAD.with(|cell| cell.set(true)); let local = LocalSet::new(); local.spawn_local(async move { assert!(ON_LOCAL_THREAD.with(|cell| cell.get())); }); local.await; } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn localset_future_timers() { static RAN1: AtomicBool = AtomicBool::new(false); static RAN2: AtomicBool = AtomicBool::new(false); let local = LocalSet::new(); local.spawn_local(async move { time::sleep(Duration::from_millis(5)).await; RAN1.store(true, Ordering::SeqCst); }); local.spawn_local(async move { time::sleep(Duration::from_millis(10)).await; RAN2.store(true, Ordering::SeqCst); }); local.await; assert!(RAN1.load(Ordering::SeqCst)); assert!(RAN2.load(Ordering::SeqCst)); } #[tokio::test] async fn localset_future_drives_all_local_futs() { static RAN1: AtomicBool = AtomicBool::new(false); static RAN2: AtomicBool = AtomicBool::new(false); static RAN3: AtomicBool = AtomicBool::new(false); let local = LocalSet::new(); local.spawn_local(async move { task::spawn_local(async { task::yield_now().await; RAN3.store(true, Ordering::SeqCst); }); task::yield_now().await; RAN1.store(true, Ordering::SeqCst); }); local.spawn_local(async move { task::yield_now().await; RAN2.store(true, Ordering::SeqCst); }); local.await; assert!(RAN1.load(Ordering::SeqCst)); assert!(RAN2.load(Ordering::SeqCst)); assert!(RAN3.load(Ordering::SeqCst)); } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn local_threadpool_timer() { // This test ensures that runtime services like the timer are properly // set for the local task set. thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); LocalSet::new() .run_until(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); let join = task::spawn_local(async move { assert!(ON_RT_THREAD.with(|cell| cell.get())); time::sleep(Duration::from_millis(10)).await; assert!(ON_RT_THREAD.with(|cell| cell.get())); }); join.await.unwrap(); }) .await; } #[test] fn enter_guard_spawn() { let local = LocalSet::new(); let _guard = local.enter(); // Run the local task set. let join = task::spawn_local(async { true }); let rt = runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); local.block_on(&rt, async move { assert!(join.await.unwrap()); }); } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery #[test] // This will panic, since the thread that calls `block_on` cannot use // in-place blocking inside of `block_on`. #[should_panic] fn local_threadpool_blocking_in_place() { thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); let rt = runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); LocalSet::new().block_on(&rt, async { assert!(ON_RT_THREAD.with(|cell| cell.get())); let join = task::spawn_local(async move { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::block_in_place(|| {}); assert!(ON_RT_THREAD.with(|cell| cell.get())); }); join.await.unwrap(); }); } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn local_threadpool_blocking_run() { thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); LocalSet::new() .run_until(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); let join = task::spawn_local(async move { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::spawn_blocking(|| { assert!( !ON_RT_THREAD.with(|cell| cell.get()), "blocking must not run on the local task set's thread" ); }) .await .unwrap(); assert!(ON_RT_THREAD.with(|cell| cell.get())); }); join.await.unwrap(); }) .await; } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn all_spawns_are_local() { use futures::future; thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); LocalSet::new() .run_until(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); let handles = (0..128) .map(|_| { task::spawn_local(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); }) }) .collect::>(); for joined in future::join_all(handles).await { joined.unwrap(); } }) .await; } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn nested_spawn_is_local() { thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); LocalSet::new() .run_until(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::spawn_local(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::spawn_local(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::spawn_local(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); task::spawn_local(async { assert!(ON_RT_THREAD.with(|cell| cell.get())); }) .await .unwrap(); }) .await .unwrap(); }) .await .unwrap(); }) .await .unwrap(); }) .await; } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[test] fn join_local_future_elsewhere() { thread_local! { static ON_RT_THREAD: Cell = const { Cell::new(false) }; } ON_RT_THREAD.with(|cell| cell.set(true)); let rt = runtime::Runtime::new().unwrap(); let local = LocalSet::new(); local.block_on(&rt, async move { let (tx, rx) = oneshot::channel(); let join = task::spawn_local(async move { assert!( ON_RT_THREAD.with(|cell| cell.get()), "local task must run on local thread, no matter where it is awaited" ); rx.await.unwrap(); "hello world" }); let join2 = task::spawn(async move { assert!( !ON_RT_THREAD.with(|cell| cell.get()), "spawned task should be on a worker" ); tx.send(()).expect("task shouldn't have ended yet"); join.await.expect("task should complete successfully"); }); join2.await.unwrap() }); } // Tests for #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn localset_in_thread_local() { thread_local! { static LOCAL_SET: LocalSet = LocalSet::new(); } // holds runtime thread until end of main fn. let (_tx, rx) = oneshot::channel::<()>(); let handle = tokio::runtime::Handle::current(); std::thread::spawn(move || { LOCAL_SET.with(|local_set| { handle.block_on(local_set.run_until(async move { let _ = rx.await; })) }); }); } #[test] fn drop_cancels_tasks() { use std::rc::Rc; // This test reproduces issue #1842 let rt = rt(); let rc1 = Rc::new(()); let rc2 = rc1.clone(); let (started_tx, started_rx) = oneshot::channel(); let local = LocalSet::new(); local.spawn_local(async move { // Move this in let _rc2 = rc2; started_tx.send(()).unwrap(); futures::future::pending::<()>().await; }); local.block_on(&rt, async { started_rx.await.unwrap(); }); drop(local); drop(rt); assert_eq!(1, Rc::strong_count(&rc1)); } /// Runs a test function in a separate thread, and panics if the test does not /// complete within the specified timeout, or if the test function panics. /// /// This is intended for running tests whose failure mode is a hang or infinite /// loop that cannot be detected otherwise. fn with_timeout(timeout: Duration, f: impl FnOnce() + Send + 'static) { use std::sync::mpsc::RecvTimeoutError; let (done_tx, done_rx) = std::sync::mpsc::channel(); let thread = std::thread::spawn(move || { f(); // Send a message on the channel so that the test thread can // determine if we have entered an infinite loop: done_tx.send(()).unwrap(); }); // Since the failure mode of this test is an infinite loop, rather than // something we can easily make assertions about, we'll run it in a // thread. When the test thread finishes, it will send a message on a // channel to this thread. We'll wait for that message with a fairly // generous timeout, and if we don't receive it, we assume the test // thread has hung. // // Note that it should definitely complete in under a minute, but just // in case CI is slow, we'll give it a long timeout. match done_rx.recv_timeout(timeout) { Err(RecvTimeoutError::Timeout) => panic!( "test did not complete within {timeout:?} seconds, \ we have (probably) entered an infinite loop!", ), // Did the test thread panic? We'll find out for sure when we `join` // with it. Err(RecvTimeoutError::Disconnected) => {} // Test completed successfully! Ok(()) => {} } thread.join().expect("test thread should not panic!") } #[cfg_attr( target_os = "wasi", ignore = "`unwrap()` in `with_timeout()` panics on Wasi" )] #[test] fn drop_cancels_remote_tasks() { // This test reproduces issue #1885. with_timeout(Duration::from_secs(60), || { let (tx, mut rx) = mpsc::channel::<()>(1024); let rt = rt(); let local = LocalSet::new(); local.spawn_local(async move { while rx.recv().await.is_some() {} }); local.block_on(&rt, async { time::sleep(Duration::from_millis(1)).await; }); drop(tx); // This enters an infinite loop if the remote notified tasks are not // properly cancelled. drop(local); }); } #[cfg_attr( target_os = "wasi", ignore = "FIXME: `task::spawn_local().await.unwrap()` panics on Wasi" )] #[test] fn local_tasks_wake_join_all() { // This test reproduces issue #2460. with_timeout(Duration::from_secs(60), || { use futures::future::join_all; use tokio::task::LocalSet; let rt = rt(); let set = LocalSet::new(); let mut handles = Vec::new(); for _ in 1..=128 { handles.push(set.spawn_local(async move { tokio::task::spawn_local(async move {}).await.unwrap(); })); } rt.block_on(set.run_until(join_all(handles))); }); } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery #[test] fn local_tasks_are_polled_after_tick() { // This test depends on timing, so we run it up to five times. for _ in 0..4 { let res = std::panic::catch_unwind(local_tasks_are_polled_after_tick_inner); if res.is_ok() { // success return; } } // Test failed 4 times. Try one more time without catching panics. If it // fails again, the test fails. local_tasks_are_polled_after_tick_inner(); } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery #[tokio::main(flavor = "current_thread")] async fn local_tasks_are_polled_after_tick_inner() { // Reproduces issues #1899 and #1900 static RX1: AtomicUsize = AtomicUsize::new(0); static RX2: AtomicUsize = AtomicUsize::new(0); const EXPECTED: usize = 500; RX1.store(0, SeqCst); RX2.store(0, SeqCst); let (tx, mut rx) = mpsc::unbounded_channel(); let local = LocalSet::new(); local .run_until(async { let task2 = task::spawn(async move { // Wait a bit time::sleep(Duration::from_millis(10)).await; let mut oneshots = Vec::with_capacity(EXPECTED); // Send values for _ in 0..EXPECTED { let (oneshot_tx, oneshot_rx) = oneshot::channel(); oneshots.push(oneshot_tx); tx.send(oneshot_rx).unwrap(); } time::sleep(Duration::from_millis(10)).await; for tx in oneshots.drain(..) { tx.send(()).unwrap(); } loop { time::sleep(Duration::from_millis(20)).await; let rx1 = RX1.load(SeqCst); let rx2 = RX2.load(SeqCst); if rx1 == EXPECTED && rx2 == EXPECTED { break; } } }); while let Some(oneshot) = rx.recv().await { RX1.fetch_add(1, SeqCst); task::spawn_local(async move { oneshot.await.unwrap(); RX2.fetch_add(1, SeqCst); }); } task2.await.unwrap(); }) .await; } #[tokio::test] async fn acquire_mutex_in_drop() { use futures::future::pending; let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); let local = LocalSet::new(); local.spawn_local(async move { let _ = rx2.await; unreachable!(); }); local.spawn_local(async move { let _ = rx1.await; tx2.send(()).unwrap(); unreachable!(); }); // Spawn a task that will never notify local.spawn_local(async move { pending::<()>().await; tx1.send(()).unwrap(); }); // Tick the loop local .run_until(async { task::yield_now().await; }) .await; // Drop the LocalSet drop(local); } #[tokio::test] async fn spawn_wakes_localset() { let local = LocalSet::new(); futures::select! { _ = local.run_until(pending::<()>()).fuse() => unreachable!(), ret = async { local.spawn_local(ready(())).await.unwrap()}.fuse() => ret } } /// Checks that the task wakes up with `enter`. /// Reproduces . #[tokio::test] async fn sleep_with_local_enter_guard() { let local = LocalSet::new(); let _guard = local.enter(); let (tx, rx) = oneshot::channel(); local .run_until(async move { tokio::task::spawn_local(async move { time::sleep(Duration::ZERO).await; tx.send(()).expect("failed to send"); }); assert_eq!(rx.await, Ok(())); }) .await; } #[test] fn store_local_set_in_thread_local_with_runtime() { use tokio::runtime::Runtime; thread_local! { static CURRENT: RtAndLocalSet = RtAndLocalSet::new(); } struct RtAndLocalSet { rt: Runtime, local: LocalSet, } impl RtAndLocalSet { fn new() -> RtAndLocalSet { RtAndLocalSet { rt: tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(), local: LocalSet::new(), } } async fn inner_method(&self) { self.local .run_until(async move { tokio::task::spawn_local(async {}); }) .await } fn method(&self) { self.rt.block_on(self.inner_method()); } } CURRENT.with(|f| { f.method(); }); } #[cfg(tokio_unstable)] mod unstable { use tokio::runtime::UnhandledPanic; use tokio::task::LocalSet; #[tokio::test] #[should_panic( expected = "a spawned task panicked and the LocalSet is configured to shutdown on unhandled panic" )] async fn shutdown_on_panic() { LocalSet::new() .unhandled_panic(UnhandledPanic::ShutdownRuntime) .run_until(async { tokio::task::spawn_local(async { panic!("boom"); }); futures::future::pending::<()>().await; }) .await; } // This test compares that, when the task driving `run_until` has already // consumed budget, the `run_until` future has less budget than a "spawned" // task. // // "Budget" is a fuzzy metric as the Tokio runtime is able to change values // internally. This is why the test uses indirection to test this. #[tokio::test] async fn run_until_does_not_get_own_budget() { // Consume some budget tokio::task::consume_budget().await; LocalSet::new() .run_until(async { let spawned = tokio::spawn(async { let mut spawned_n = 0; { let mut spawned = tokio_test::task::spawn(async { loop { spawned_n += 1; tokio::task::consume_budget().await; } }); // Poll once assert!(!spawned.poll().is_ready()); } spawned_n }); let mut run_until_n = 0; { let mut run_until = tokio_test::task::spawn(async { loop { run_until_n += 1; tokio::task::consume_budget().await; } }); // Poll once assert!(!run_until.poll().is_ready()); } let spawned_n = spawned.await.unwrap(); assert_ne!(spawned_n, 0); assert_ne!(run_until_n, 0); assert!(spawned_n > run_until_n); }) .await } } fn rt() -> runtime::Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/task_panic.rs000064400000000000000000000055311046102023000146670ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] #![cfg(panic = "unwind")] use futures::future; use std::error::Error; use tokio::runtime::Builder; use tokio::task::{self, block_in_place}; mod support { pub mod panic; } use support::panic::test_panic; #[test] fn block_in_place_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = Builder::new_current_thread().enable_all().build().unwrap(); rt.block_on(async { block_in_place(|| {}); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn local_set_spawn_local_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _local = task::LocalSet::new(); task::spawn_local(async {}); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn local_set_block_on_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = Builder::new_current_thread().enable_all().build().unwrap(); let local = task::LocalSet::new(); rt.block_on(async { local.block_on(&rt, future::pending::<()>()); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn spawn_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { tokio::spawn(future::pending::<()>()); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn local_key_sync_scope_panic_caller() -> Result<(), Box> { tokio::task_local! { static NUMBER: u32; } let panic_location_file = test_panic(|| { NUMBER.sync_scope(1, || { NUMBER.with(|_| { NUMBER.sync_scope(1, || {}); }); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn local_key_with_panic_caller() -> Result<(), Box> { tokio::task_local! { static NUMBER: u32; } let panic_location_file = test_panic(|| { NUMBER.with(|_| {}); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn local_key_get_panic_caller() -> Result<(), Box> { tokio::task_local! { static NUMBER: u32; } let panic_location_file = test_panic(|| { NUMBER.get(); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } tokio-1.43.1/tests/task_yield_now.rs000064400000000000000000000012701046102023000155620ustar 00000000000000#![allow(unknown_lints, unexpected_cfgs)] #![cfg(all(feature = "full", not(target_os = "wasi"), tokio_unstable))] use tokio::task; use tokio_test::task::spawn; // `yield_now` is tested within the runtime in `rt_common`. #[test] fn yield_now_outside_of_runtime() { let mut task = spawn(async { task::yield_now().await; }); assert!(task.poll().is_pending()); assert!(task.is_woken()); assert!(task.poll().is_ready()); } #[tokio::test(flavor = "multi_thread")] async fn yield_now_external_executor_and_block_in_place() { let j = tokio::spawn(async { task::block_in_place(|| futures::executor::block_on(task::yield_now())); }); j.await.unwrap(); } tokio-1.43.1/tests/tcp_accept.rs000064400000000000000000000111231046102023000146520ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{mpsc, oneshot}; use tokio_test::assert_ok; use std::io; use std::net::{IpAddr, SocketAddr}; macro_rules! test_accept { ($(($ident:ident, $target:expr),)*) => { $( #[tokio::test] async fn $ident() { let listener = assert_ok!(TcpListener::bind($target).await); let addr = listener.local_addr().unwrap(); let (tx, rx) = oneshot::channel(); tokio::spawn(async move { let (socket, _) = assert_ok!(listener.accept().await); assert_ok!(tx.send(socket)); }); let cli = assert_ok!(TcpStream::connect(&addr).await); let srv = assert_ok!(rx.await); assert_eq!(cli.local_addr().unwrap(), srv.peer_addr().unwrap()); } )* } } test_accept! { (ip_str, "127.0.0.1:0"), (host_str, "localhost:0"), (socket_addr, "127.0.0.1:0".parse::().unwrap()), (str_port_tuple, ("127.0.0.1", 0)), (ip_port_tuple, ("127.0.0.1".parse::().unwrap(), 0)), } use std::pin::Pin; use std::sync::{ atomic::{AtomicUsize, Ordering::SeqCst}, Arc, }; use std::task::{Context, Poll}; use tokio_stream::{Stream, StreamExt}; struct TrackPolls<'a> { npolls: Arc, listener: &'a mut TcpListener, } impl<'a> Stream for TrackPolls<'a> { type Item = io::Result<(TcpStream, SocketAddr)>; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.npolls.fetch_add(1, SeqCst); self.listener.poll_accept(cx).map(Some) } } #[tokio::test] async fn no_extra_poll() { let mut listener = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = listener.local_addr().unwrap(); let (tx, rx) = oneshot::channel(); let (accepted_tx, mut accepted_rx) = mpsc::unbounded_channel(); tokio::spawn(async move { let mut incoming = TrackPolls { npolls: Arc::new(AtomicUsize::new(0)), listener: &mut listener, }; assert_ok!(tx.send(Arc::clone(&incoming.npolls))); while incoming.next().await.is_some() { accepted_tx.send(()).unwrap(); } }); let npolls = assert_ok!(rx.await); tokio::task::yield_now().await; // should have been polled exactly once: the initial poll assert_eq!(npolls.load(SeqCst), 1); let _ = assert_ok!(TcpStream::connect(&addr).await); accepted_rx.recv().await.unwrap(); // should have been polled twice more: once to yield Some(), then once to yield Pending assert_eq!(npolls.load(SeqCst), 1 + 2); } #[tokio::test] async fn accept_many() { use std::future::{poll_fn, Future}; use std::sync::atomic::AtomicBool; const N: usize = 50; let listener = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let listener = Arc::new(listener); let addr = listener.local_addr().unwrap(); let connected = Arc::new(AtomicBool::new(false)); let (pending_tx, mut pending_rx) = mpsc::unbounded_channel(); let (notified_tx, mut notified_rx) = mpsc::unbounded_channel(); for _ in 0..N { let listener = listener.clone(); let connected = connected.clone(); let pending_tx = pending_tx.clone(); let notified_tx = notified_tx.clone(); tokio::spawn(async move { let accept = listener.accept(); tokio::pin!(accept); let mut polled = false; poll_fn(|cx| { if !polled { polled = true; assert!(Pin::new(&mut accept).poll(cx).is_pending()); pending_tx.send(()).unwrap(); Poll::Pending } else if connected.load(SeqCst) { notified_tx.send(()).unwrap(); Poll::Ready(()) } else { Poll::Pending } }) .await; pending_tx.send(()).unwrap(); }); } // Wait for all tasks to have polled at least once for _ in 0..N { pending_rx.recv().await.unwrap(); } // Establish a TCP connection connected.store(true, SeqCst); let _sock = TcpStream::connect(addr).await.unwrap(); // Wait for all notifications for _ in 0..N { notified_rx.recv().await.unwrap(); } } tokio-1.43.1/tests/tcp_connect.rs000064400000000000000000000135221046102023000150510ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use tokio::net::{TcpListener, TcpStream}; use tokio::sync::oneshot; use tokio_test::assert_ok; use futures::join; #[tokio::test] async fn connect_v4() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); assert!(addr.is_ipv4()); let (tx, rx) = oneshot::channel(); tokio::spawn(async move { let (socket, addr) = assert_ok!(srv.accept().await); assert_eq!(addr, assert_ok!(socket.peer_addr())); assert_ok!(tx.send(socket)); }); let mine = assert_ok!(TcpStream::connect(&addr).await); let theirs = assert_ok!(rx.await); assert_eq!( assert_ok!(mine.local_addr()), assert_ok!(theirs.peer_addr()) ); assert_eq!( assert_ok!(theirs.local_addr()), assert_ok!(mine.peer_addr()) ); } #[tokio::test] async fn connect_v6() { let srv = assert_ok!(TcpListener::bind("[::1]:0").await); let addr = assert_ok!(srv.local_addr()); assert!(addr.is_ipv6()); let (tx, rx) = oneshot::channel(); tokio::spawn(async move { let (socket, addr) = assert_ok!(srv.accept().await); assert_eq!(addr, assert_ok!(socket.peer_addr())); assert_ok!(tx.send(socket)); }); let mine = assert_ok!(TcpStream::connect(&addr).await); let theirs = assert_ok!(rx.await); assert_eq!( assert_ok!(mine.local_addr()), assert_ok!(theirs.peer_addr()) ); assert_eq!( assert_ok!(theirs.local_addr()), assert_ok!(mine.peer_addr()) ); } #[tokio::test] async fn connect_addr_ip_string() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let addr = format!("127.0.0.1:{}", addr.port()); let server = async { assert_ok!(srv.accept().await); }; let client = async { assert_ok!(TcpStream::connect(addr).await); }; join!(server, client); } #[tokio::test] async fn connect_addr_ip_str_slice() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let addr = format!("127.0.0.1:{}", addr.port()); let server = async { assert_ok!(srv.accept().await); }; let client = async { assert_ok!(TcpStream::connect(&addr[..]).await); }; join!(server, client); } #[tokio::test] async fn connect_addr_host_string() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let addr = format!("localhost:{}", addr.port()); let server = async { assert_ok!(srv.accept().await); }; let client = async { assert_ok!(TcpStream::connect(addr).await); }; join!(server, client); } #[tokio::test] async fn connect_addr_ip_port_tuple() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let addr = (addr.ip(), addr.port()); let server = async { assert_ok!(srv.accept().await); }; let client = async { assert_ok!(TcpStream::connect(&addr).await); }; join!(server, client); } #[tokio::test] async fn connect_addr_ip_str_port_tuple() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let addr = ("127.0.0.1", addr.port()); let server = async { assert_ok!(srv.accept().await); }; let client = async { assert_ok!(TcpStream::connect(&addr).await); }; join!(server, client); } #[tokio::test] async fn connect_addr_host_str_port_tuple() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let addr = ("localhost", addr.port()); let server = async { assert_ok!(srv.accept().await); }; let client = async { assert_ok!(TcpStream::connect(&addr).await); }; join!(server, client); } /* * TODO: bring this back once TCP exposes HUP again * #[cfg(target_os = "linux")] mod linux { use tokio::net::{TcpListener, TcpStream}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio_test::assert_ok; use mio::unix::UnixReady; use futures_util::future::poll_fn; use std::io::Write; use std::time::Duration; use std::{net, thread}; #[tokio::test] fn poll_hup() { let addr = assert_ok!("127.0.0.1:0".parse()); let mut srv = assert_ok!(TcpListener::bind(&addr)); let addr = assert_ok!(srv.local_addr()); tokio::spawn(async move { let (mut client, _) = assert_ok!(srv.accept().await); assert_ok!(client.set_linger(Some(Duration::from_millis(0)))); assert_ok!(client.write_all(b"hello world").await); // TODO: Drop? }); /* let t = thread::spawn(move || { let mut client = assert_ok!(srv.accept()).0; client.set_linger(Some(Duration::from_millis(0))).unwrap(); client.write(b"hello world").unwrap(); thread::sleep(Duration::from_millis(200)); }); */ let mut stream = assert_ok!(TcpStream::connect(&addr).await); // Poll for HUP before reading. future::poll_fn(|| stream.poll_read_ready(UnixReady::hup().into())) .wait() .unwrap(); // Same for write half future::poll_fn(|| stream.poll_write_ready()) .wait() .unwrap(); let mut buf = vec![0; 11]; // Read the data future::poll_fn(|| stream.poll_read(&mut buf)) .wait() .unwrap(); assert_eq!(b"hello world", &buf[..]); t.join().unwrap(); } } */ tokio-1.43.1/tests/tcp_echo.rs000064400000000000000000000023771046102023000143440ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No socket on miri. use tokio::io::{self, AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::oneshot; use tokio_test::assert_ok; #[tokio::test] async fn echo_server() { const ITER: usize = 1024; let (tx, rx) = oneshot::channel(); let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); let msg = "foo bar baz"; tokio::spawn(async move { let mut stream = assert_ok!(TcpStream::connect(&addr).await); for _ in 0..ITER { // write assert_ok!(stream.write_all(msg.as_bytes()).await); // read let mut buf = [0; 11]; assert_ok!(stream.read_exact(&mut buf).await); assert_eq!(&buf[..], msg.as_bytes()); } assert_ok!(tx.send(())); }); let (mut stream, _) = assert_ok!(srv.accept().await); let (mut rd, mut wr) = stream.split(); let n = assert_ok!(io::copy(&mut rd, &mut wr).await); assert_eq!(n, (ITER * msg.len()) as u64); assert_ok!(rx.await); } tokio-1.43.1/tests/tcp_into_split.rs000064400000000000000000000074331046102023000156100ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use std::io::{Error, ErrorKind, Result}; use std::io::{Read, Write}; use std::{net, thread}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::try_join; #[tokio::test] async fn split() -> Result<()> { const MSG: &[u8] = b"split"; let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr()?; let (stream1, (mut stream2, _)) = try_join! { TcpStream::connect(&addr), listener.accept(), }?; let (mut read_half, mut write_half) = stream1.into_split(); let ((), (), ()) = try_join! { async { let len = stream2.write(MSG).await?; assert_eq!(len, MSG.len()); let mut read_buf = vec![0u8; 32]; let read_len = stream2.read(&mut read_buf).await?; assert_eq!(&read_buf[..read_len], MSG); Result::Ok(()) }, async { let len = write_half.write(MSG).await?; assert_eq!(len, MSG.len()); Ok(()) }, async { let mut read_buf = [0u8; 32]; let peek_len1 = read_half.peek(&mut read_buf[..]).await?; let peek_len2 = read_half.peek(&mut read_buf[..]).await?; assert_eq!(peek_len1, peek_len2); let read_len = read_half.read(&mut read_buf[..]).await?; assert_eq!(peek_len1, read_len); assert_eq!(&read_buf[..read_len], MSG); Ok(()) }, }?; Ok(()) } #[tokio::test] async fn reunite() -> Result<()> { let listener = net::TcpListener::bind("127.0.0.1:0")?; let addr = listener.local_addr()?; let handle = thread::spawn(move || { drop(listener.accept().unwrap()); drop(listener.accept().unwrap()); }); let stream1 = TcpStream::connect(&addr).await?; let (read1, write1) = stream1.into_split(); let stream2 = TcpStream::connect(&addr).await?; let (_, write2) = stream2.into_split(); let read1 = match read1.reunite(write2) { Ok(_) => panic!("Reunite should not succeed"), Err(err) => err.0, }; read1.reunite(write1).expect("Reunite should succeed"); handle.join().unwrap(); Ok(()) } /// Test that dropping the write half actually closes the stream. #[tokio::test] async fn drop_write() -> Result<()> { const MSG: &[u8] = b"split"; let listener = net::TcpListener::bind("127.0.0.1:0")?; let addr = listener.local_addr()?; let handle = thread::spawn(move || { let (mut stream, _) = listener.accept().unwrap(); stream.write_all(MSG).unwrap(); let mut read_buf = [0u8; 32]; let res = match stream.read(&mut read_buf) { Ok(0) => Ok(()), Ok(len) => Err(Error::new( ErrorKind::Other, format!("Unexpected read: {len} bytes."), )), Err(err) => Err(err), }; drop(stream); res }); let stream = TcpStream::connect(&addr).await?; let (mut read_half, write_half) = stream.into_split(); let mut read_buf = [0u8; 32]; let read_len = read_half.read(&mut read_buf[..]).await?; assert_eq!(&read_buf[..read_len], MSG); // drop it while the read is in progress std::thread::spawn(move || { thread::sleep(std::time::Duration::from_millis(10)); drop(write_half); }); match read_half.read(&mut read_buf[..]).await { Ok(0) => {} Ok(len) => panic!("Unexpected read: {len} bytes."), Err(err) => panic!("Unexpected error: {err}."), } handle.join().unwrap().unwrap(); Ok(()) } tokio-1.43.1/tests/tcp_into_std.rs000064400000000000000000000030071046102023000152400ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No socket on `miri`. use std::io::Read; use std::io::Result; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use tokio::net::TcpStream; #[tokio::test] async fn tcp_into_std() -> Result<()> { let mut data = [0u8; 12]; let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr().unwrap().to_string(); let handle = tokio::spawn(async { let stream: TcpStream = TcpStream::connect(addr).await.unwrap(); stream }); let (tokio_tcp_stream, _) = listener.accept().await?; let mut std_tcp_stream = tokio_tcp_stream.into_std()?; std_tcp_stream .set_nonblocking(false) .expect("set_nonblocking call failed"); let mut client = handle.await.expect("The task being joined has panicked"); client.write_all(b"Hello world!").await?; std_tcp_stream .read_exact(&mut data) .expect("std TcpStream read failed!"); assert_eq!(b"Hello world!", &data); // test back to tokio stream std_tcp_stream .set_nonblocking(true) .expect("set_nonblocking call failed"); let mut tokio_tcp_stream = TcpStream::from_std(std_tcp_stream)?; client.write_all(b"Hello tokio!").await?; let _size = tokio_tcp_stream.read_exact(&mut data).await?; assert_eq!(b"Hello tokio!", &data); Ok(()) } tokio-1.43.1/tests/tcp_peek.rs000064400000000000000000000017171046102023000143470ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use tokio::io::AsyncReadExt; use tokio::net::TcpStream; use tokio_test::assert_ok; use std::thread; use std::{io::Write, net}; #[tokio::test] async fn peek() { let listener = net::TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); let t = thread::spawn(move || assert_ok!(listener.accept()).0); let left = net::TcpStream::connect(addr).unwrap(); let mut right = t.join().unwrap(); let _ = right.write(&[1, 2, 3, 4]).unwrap(); let mut left: TcpStream = left.try_into().unwrap(); let mut buf = [0u8; 16]; let n = assert_ok!(left.peek(&mut buf).await); assert_eq!([1, 2, 3, 4], buf[..n]); let n = assert_ok!(left.read(&mut buf).await); assert_eq!([1, 2, 3, 4], buf[..n]); } tokio-1.43.1/tests/tcp_shutdown.rs000064400000000000000000000017001046102023000152660ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use tokio::io::{self, AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio_test::assert_ok; #[tokio::test] async fn shutdown() { let srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(srv.local_addr()); tokio::spawn(async move { let mut stream = assert_ok!(TcpStream::connect(&addr).await); assert_ok!(AsyncWriteExt::shutdown(&mut stream).await); let mut buf = [0u8; 1]; let n = assert_ok!(stream.read(&mut buf).await); assert_eq!(n, 0); }); let (mut stream, _) = assert_ok!(srv.accept().await); let (mut rd, mut wr) = stream.split(); let n = assert_ok!(io::copy(&mut rd, &mut wr).await); assert_eq!(n, 0); } tokio-1.43.1/tests/tcp_socket.rs000064400000000000000000000041071046102023000147070ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use std::time::Duration; use tokio::net::TcpSocket; use tokio_test::assert_ok; #[tokio::test] async fn basic_usage_v4() { // Create server let addr = assert_ok!("127.0.0.1:0".parse()); let srv = assert_ok!(TcpSocket::new_v4()); assert_ok!(srv.bind(addr)); let srv = assert_ok!(srv.listen(128)); // Create client & connect let addr = srv.local_addr().unwrap(); let cli = assert_ok!(TcpSocket::new_v4()); let _cli = assert_ok!(cli.connect(addr).await); // Accept let _ = assert_ok!(srv.accept().await); } #[tokio::test] async fn basic_usage_v6() { // Create server let addr = assert_ok!("[::1]:0".parse()); let srv = assert_ok!(TcpSocket::new_v6()); assert_ok!(srv.bind(addr)); let srv = assert_ok!(srv.listen(128)); // Create client & connect let addr = srv.local_addr().unwrap(); let cli = assert_ok!(TcpSocket::new_v6()); let _cli = assert_ok!(cli.connect(addr).await); // Accept let _ = assert_ok!(srv.accept().await); } #[tokio::test] async fn bind_before_connect() { // Create server let any_addr = assert_ok!("127.0.0.1:0".parse()); let srv = assert_ok!(TcpSocket::new_v4()); assert_ok!(srv.bind(any_addr)); let srv = assert_ok!(srv.listen(128)); // Create client & connect let addr = srv.local_addr().unwrap(); let cli = assert_ok!(TcpSocket::new_v4()); assert_ok!(cli.bind(any_addr)); let _cli = assert_ok!(cli.connect(addr).await); // Accept let _ = assert_ok!(srv.accept().await); } #[tokio::test] async fn basic_linger() { // Create server let addr = assert_ok!("127.0.0.1:0".parse()); let srv = assert_ok!(TcpSocket::new_v4()); assert_ok!(srv.bind(addr)); assert!(srv.linger().unwrap().is_none()); srv.set_linger(Some(Duration::new(0, 0))).unwrap(); assert_eq!(srv.linger().unwrap(), Some(Duration::new(0, 0))); } tokio-1.43.1/tests/tcp_split.rs000064400000000000000000000025571046102023000145610ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind // No `socket` on miri. use std::io::Result; use std::io::{Read, Write}; use std::{net, thread}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; #[tokio::test] async fn split() -> Result<()> { const MSG: &[u8] = b"split"; let listener = net::TcpListener::bind("127.0.0.1:0")?; let addr = listener.local_addr()?; let handle = thread::spawn(move || { let (mut stream, _) = listener.accept().unwrap(); stream.write_all(MSG).unwrap(); let mut read_buf = [0u8; 32]; let read_len = stream.read(&mut read_buf).unwrap(); assert_eq!(&read_buf[..read_len], MSG); }); let mut stream = TcpStream::connect(&addr).await?; let (mut read_half, mut write_half) = stream.split(); let mut read_buf = [0u8; 32]; let peek_len1 = read_half.peek(&mut read_buf[..]).await?; let peek_len2 = read_half.peek(&mut read_buf[..]).await?; assert_eq!(peek_len1, peek_len2); let read_len = read_half.read(&mut read_buf[..]).await?; assert_eq!(peek_len1, read_len); assert_eq!(&read_buf[..read_len], MSG); assert_eq!(write_half.write(MSG).await?, MSG.len()); handle.join().unwrap(); Ok(()) } tokio-1.43.1/tests/tcp_stream.rs000064400000000000000000000260251046102023000147150ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::io::{AsyncReadExt, AsyncWriteExt, Interest}; use tokio::net::{TcpListener, TcpStream}; use tokio::try_join; use tokio_test::task; use tokio_test::{assert_ok, assert_pending, assert_ready_ok}; use std::future::poll_fn; use std::io; use std::task::Poll; use std::time::Duration; #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn set_linger() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let stream = TcpStream::connect(listener.local_addr().unwrap()) .await .unwrap(); assert_ok!(stream.set_linger(Some(Duration::from_secs(1)))); assert_eq!(stream.linger().unwrap().unwrap().as_secs(), 1); assert_ok!(stream.set_linger(None)); assert!(stream.linger().unwrap().is_none()); } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn try_read_write() { const DATA: &[u8] = b"this is some data to write to the socket"; // Create listener let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); // Create socket pair let client = TcpStream::connect(listener.local_addr().unwrap()) .await .unwrap(); let (server, _) = listener.accept().await.unwrap(); let mut written = DATA.to_vec(); // Track the server receiving data let mut readable = task::spawn(server.readable()); assert_pending!(readable.poll()); // Write data. client.writable().await.unwrap(); assert_eq!(DATA.len(), client.try_write(DATA).unwrap()); // The task should be notified while !readable.is_woken() { tokio::task::yield_now().await; } // Fill the write buffer using non-vectored I/O loop { // Still ready let mut writable = task::spawn(client.writable()); assert_ready_ok!(writable.poll()); match client.try_write(DATA) { Ok(n) => written.extend(&DATA[..n]), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { break; } Err(e) => panic!("error = {e:?}"), } } { // Write buffer full let mut writable = task::spawn(client.writable()); assert_pending!(writable.poll()); // Drain the socket from the server end using non-vectored I/O let mut read = vec![0; written.len()]; let mut i = 0; while i < read.len() { server.readable().await.unwrap(); match server.try_read(&mut read[i..]) { Ok(n) => i += n, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("error = {e:?}"), } } assert_eq!(read, written); } written.clear(); client.writable().await.unwrap(); // Fill the write buffer using vectored I/O let data_bufs: Vec<_> = DATA.chunks(10).map(io::IoSlice::new).collect(); loop { // Still ready let mut writable = task::spawn(client.writable()); assert_ready_ok!(writable.poll()); match client.try_write_vectored(&data_bufs) { Ok(n) => written.extend(&DATA[..n]), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { break; } Err(e) => panic!("error = {e:?}"), } } { // Write buffer full let mut writable = task::spawn(client.writable()); assert_pending!(writable.poll()); // Drain the socket from the server end using vectored I/O let mut read = vec![0; written.len()]; let mut i = 0; while i < read.len() { server.readable().await.unwrap(); let mut bufs: Vec<_> = read[i..] .chunks_mut(0x10000) .map(io::IoSliceMut::new) .collect(); match server.try_read_vectored(&mut bufs) { Ok(n) => i += n, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("error = {e:?}"), } } assert_eq!(read, written); } // Now, we listen for shutdown drop(client); loop { let ready = server.ready(Interest::READABLE).await.unwrap(); if ready.is_read_closed() { return; } else { tokio::task::yield_now().await; } } } #[test] fn buffer_not_included_in_future() { use std::mem; const N: usize = 4096; let fut = async { let stream = TcpStream::connect("127.0.0.1:8080").await.unwrap(); loop { stream.readable().await.unwrap(); let mut buf = [0; N]; let n = stream.try_read(&mut buf[..]).unwrap(); if n == 0 { break; } } }; let n = mem::size_of_val(&fut); assert!(n < 1000); } macro_rules! assert_readable_by_polling { ($stream:expr) => { assert_ok!(poll_fn(|cx| $stream.poll_read_ready(cx)).await); }; } macro_rules! assert_not_readable_by_polling { ($stream:expr) => { poll_fn(|cx| { assert_pending!($stream.poll_read_ready(cx)); Poll::Ready(()) }) .await; }; } macro_rules! assert_writable_by_polling { ($stream:expr) => { assert_ok!(poll_fn(|cx| $stream.poll_write_ready(cx)).await); }; } macro_rules! assert_not_writable_by_polling { ($stream:expr) => { poll_fn(|cx| { assert_pending!($stream.poll_write_ready(cx)); Poll::Ready(()) }) .await; }; } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn poll_read_ready() { let (mut client, mut server) = create_pair().await; // Initial state - not readable. assert_not_readable_by_polling!(server); // There is data in the buffer - readable. assert_ok!(client.write_all(b"ping").await); assert_readable_by_polling!(server); // Readable until calls to `poll_read` return `Poll::Pending`. let mut buf = [0u8; 4]; assert_ok!(server.read_exact(&mut buf).await); assert_readable_by_polling!(server); read_until_pending(&mut server); assert_not_readable_by_polling!(server); // Detect the client disconnect. drop(client); assert_readable_by_polling!(server); } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn poll_write_ready() { let (mut client, server) = create_pair().await; // Initial state - writable. assert_writable_by_polling!(client); // No space to write - not writable. write_until_pending(&mut client); assert_not_writable_by_polling!(client); // Detect the server disconnect. drop(server); assert_writable_by_polling!(client); } async fn create_pair() -> (TcpStream, TcpStream) { let listener = assert_ok!(TcpListener::bind("127.0.0.1:0").await); let addr = assert_ok!(listener.local_addr()); let (client, (server, _)) = assert_ok!(try_join!(TcpStream::connect(&addr), listener.accept())); (client, server) } fn read_until_pending(stream: &mut TcpStream) -> usize { let mut buf = vec![0u8; 1024 * 1024]; let mut total = 0; loop { match stream.try_read(&mut buf) { Ok(n) => total += n, Err(err) => { assert_eq!(err.kind(), io::ErrorKind::WouldBlock); break; } } } total } fn write_until_pending(stream: &mut TcpStream) -> usize { let buf = vec![0u8; 1024 * 1024]; let mut total = 0; loop { match stream.try_write(&buf) { Ok(n) => total += n, Err(err) => { assert_eq!(err.kind(), io::ErrorKind::WouldBlock); break; } } } total } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn try_read_buf() { const DATA: &[u8] = b"this is some data to write to the socket"; // Create listener let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); // Create socket pair let client = TcpStream::connect(listener.local_addr().unwrap()) .await .unwrap(); let (server, _) = listener.accept().await.unwrap(); let mut written = DATA.to_vec(); // Track the server receiving data let mut readable = task::spawn(server.readable()); assert_pending!(readable.poll()); // Write data. client.writable().await.unwrap(); assert_eq!(DATA.len(), client.try_write(DATA).unwrap()); // The task should be notified while !readable.is_woken() { tokio::task::yield_now().await; } // Fill the write buffer loop { // Still ready let mut writable = task::spawn(client.writable()); assert_ready_ok!(writable.poll()); match client.try_write(DATA) { Ok(n) => written.extend(&DATA[..n]), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { break; } Err(e) => panic!("error = {e:?}"), } } { // Write buffer full let mut writable = task::spawn(client.writable()); assert_pending!(writable.poll()); // Drain the socket from the server end let mut read = Vec::with_capacity(written.len()); let mut i = 0; while i < read.capacity() { server.readable().await.unwrap(); match server.try_read_buf(&mut read) { Ok(n) => i += n, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("error = {e:?}"), } } assert_eq!(read, written); } // Now, we listen for shutdown drop(client); loop { let ready = server.ready(Interest::READABLE).await.unwrap(); if ready.is_read_closed() { return; } else { tokio::task::yield_now().await; } } } // read_closed is a best effort event, so test only for no false positives. #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn read_closed() { let (client, mut server) = create_pair().await; let mut ready_fut = task::spawn(client.ready(Interest::READABLE)); assert_pending!(ready_fut.poll()); assert_ok!(server.write_all(b"ping").await); let ready_event = assert_ok!(ready_fut.await); assert!(!ready_event.is_read_closed()); } // write_closed is a best effort event, so test only for no false positives. #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn write_closed() { let (mut client, mut server) = create_pair().await; // Fill the write buffer. let write_size = write_until_pending(&mut client); let mut ready_fut = task::spawn(client.ready(Interest::WRITABLE)); assert_pending!(ready_fut.poll()); // Drain the socket to make client writable. let mut read_size = 0; while read_size < write_size { server.readable().await.unwrap(); read_size += read_until_pending(&mut server); } let ready_event = assert_ok!(ready_fut.await); assert!(!ready_event.is_write_closed()); } tokio-1.43.1/tests/test_clock.rs000064400000000000000000000025601046102023000147040ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::time::{self, Duration, Instant}; #[tokio::test] async fn resume_lets_time_move_forward_instead_of_resetting_it() { let start = Instant::now(); time::pause(); time::advance(Duration::from_secs(10)).await; let advanced_by_ten_secs = Instant::now(); assert!(advanced_by_ten_secs - start > Duration::from_secs(10)); assert!(advanced_by_ten_secs - start < Duration::from_secs(11)); time::resume(); assert!(advanced_by_ten_secs < Instant::now()); assert!(Instant::now() - advanced_by_ten_secs < Duration::from_secs(1)); } #[tokio::test] async fn can_pause_after_resume() { let start = Instant::now(); time::pause(); time::advance(Duration::from_secs(10)).await; time::resume(); time::pause(); time::advance(Duration::from_secs(10)).await; assert!(Instant::now() - start > Duration::from_secs(20)); assert!(Instant::now() - start < Duration::from_secs(21)); } #[tokio::test] #[should_panic] async fn freezing_time_while_frozen_panics() { time::pause(); time::pause(); } #[tokio::test] #[should_panic] async fn advancing_time_when_time_is_not_frozen_panics() { time::advance(Duration::from_secs(1)).await; } #[tokio::test] #[should_panic] async fn resuming_time_when_not_frozen_panics() { time::pause(); time::resume(); time::resume(); } tokio-1.43.1/tests/time_interval.rs000064400000000000000000000404241046102023000154150ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use std::pin::Pin; use std::task::{Context, Poll}; use futures::{Stream, StreamExt}; use tokio::time::{self, Duration, Instant, Interval, MissedTickBehavior}; use tokio_test::{assert_pending, assert_ready, assert_ready_eq, task}; // Takes the `Interval` task, `start` variable, and optional time deltas // For each time delta, it polls the `Interval` and asserts that the result is // equal to `start` + the specific time delta. Then it asserts that the // `Interval` is pending. macro_rules! check_interval_poll { ($i:ident, $start:ident, $($delta:expr),*$(,)?) => { $( assert_ready_eq!(poll_next(&mut $i), $start + ms($delta)); )* assert_pending!(poll_next(&mut $i)); }; ($i:ident, $start:ident) => { check_interval_poll!($i, $start,); }; } #[tokio::test] #[should_panic] async fn interval_zero_duration() { let _ = time::interval_at(Instant::now(), ms(0)); } // Expected ticks: | 1 | 2 | 3 | 4 | 5 | 6 | // Actual ticks: | work -----| delay | work | work | work -| work -----| // Poll behavior: | | | | | | | | // | | | | | | | | // Ready(s) | | Ready(s + 2p) | | | | // Pending | Ready(s + 3p) | | | // Ready(s + p) Ready(s + 4p) | | // Ready(s + 5p) | // Ready(s + 6p) #[tokio::test(start_paused = true)] async fn burst() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(650)).await; check_interval_poll!(i, start, 600, 900); time::advance(ms(200)).await; check_interval_poll!(i, start); time::advance(ms(100)).await; check_interval_poll!(i, start, 1200); time::advance(ms(250)).await; check_interval_poll!(i, start, 1500); time::advance(ms(300)).await; check_interval_poll!(i, start, 1800); } // Expected ticks: | 1 | 2 | 3 | 4 | 5 | 6 | // Actual ticks: | work -----| delay | work -----| work -----| work -----| // Poll behavior: | | | | | | | | // | | | | | | | | // Ready(s) | | Ready(s + 2p) | | | | // Pending | Pending | | | // Ready(s + p) Ready(s + 2p + d) | | // Ready(s + 3p + d) | // Ready(s + 4p + d) #[tokio::test(start_paused = true)] async fn delay() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); i.set_missed_tick_behavior(MissedTickBehavior::Delay); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(650)).await; check_interval_poll!(i, start, 600); time::advance(ms(100)).await; check_interval_poll!(i, start); // We have to add one here for the same reason as is above. // Because `Interval` has reset its timer according to `Instant::now()`, // we have to go forward 1 more millisecond than is expected so that the // runtime realizes that it's time to resolve the timer. time::advance(ms(201)).await; // We add one because when using the `Delay` behavior, `Interval` // adds the `period` from `Instant::now()`, which will always be off by one // because we have to advance time by 1 (see above). check_interval_poll!(i, start, 1251); time::advance(ms(300)).await; // Again, we add one. check_interval_poll!(i, start, 1551); time::advance(ms(300)).await; check_interval_poll!(i, start, 1851); } // Expected ticks: | 1 | 2 | 3 | 4 | 5 | 6 | // Actual ticks: | work -----| delay | work ---| work -----| work -----| // Poll behavior: | | | | | | | // | | | | | | | // Ready(s) | | Ready(s + 2p) | | | // Pending | Ready(s + 4p) | | // Ready(s + p) Ready(s + 5p) | // Ready(s + 6p) #[tokio::test(start_paused = true)] async fn skip() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); i.set_missed_tick_behavior(MissedTickBehavior::Skip); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(650)).await; check_interval_poll!(i, start, 600); time::advance(ms(250)).await; check_interval_poll!(i, start, 1200); time::advance(ms(300)).await; check_interval_poll!(i, start, 1500); time::advance(ms(300)).await; check_interval_poll!(i, start, 1800); } #[tokio::test(start_paused = true)] async fn reset() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(100)).await; check_interval_poll!(i, start); i.reset(); time::advance(ms(250)).await; check_interval_poll!(i, start); time::advance(ms(50)).await; // We add one because when using `reset` method, `Interval` adds the // `period` from `Instant::now()`, which will always be off by one check_interval_poll!(i, start, 701); time::advance(ms(300)).await; check_interval_poll!(i, start, 1001); } #[tokio::test(start_paused = true)] async fn reset_immediately() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(100)).await; check_interval_poll!(i, start); i.reset_immediately(); // We add one because when using `reset` method, `Interval` adds the // `period` from `Instant::now()`, which will always be off by one check_interval_poll!(i, start, 401); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 701); } #[tokio::test(start_paused = true)] async fn reset_after() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(100)).await; check_interval_poll!(i, start); i.reset_after(Duration::from_millis(20)); // We add one because when using `reset` method, `Interval` adds the // `period` from `Instant::now()`, which will always be off by one time::advance(ms(20)).await; check_interval_poll!(i, start, 421); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 721); } #[tokio::test(start_paused = true)] async fn reset_at() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(100)).await; check_interval_poll!(i, start); i.reset_at(Instant::now() + Duration::from_millis(40)); // We add one because when using `reset` method, `Interval` adds the // `period` from `Instant::now()`, which will always be off by one time::advance(ms(40)).await; check_interval_poll!(i, start, 441); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 741); } #[tokio::test(start_paused = true)] async fn reset_at_bigger_than_interval() { let start = Instant::now(); // This is necessary because the timer is only so granular, and in order for // all our ticks to resolve, the time needs to be 1ms ahead of what we // expect, so that the runtime will see that it is time to resolve the timer time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); check_interval_poll!(i, start, 0); time::advance(ms(100)).await; check_interval_poll!(i, start); time::advance(ms(200)).await; check_interval_poll!(i, start, 300); time::advance(ms(100)).await; check_interval_poll!(i, start); i.reset_at(Instant::now() + Duration::from_millis(1000)); // Validate the interval does not tick until 1000ms have passed time::advance(ms(300)).await; check_interval_poll!(i, start); time::advance(ms(300)).await; check_interval_poll!(i, start); time::advance(ms(300)).await; check_interval_poll!(i, start); // We add one because when using `reset` method, `Interval` adds the // `period` from `Instant::now()`, which will always be off by one time::advance(ms(100)).await; check_interval_poll!(i, start, 1401); time::advance(ms(300)).await; check_interval_poll!(i, start, 1701); } fn poll_next(interval: &mut task::Spawn) -> Poll { interval.enter(|cx, mut interval| interval.poll_tick(cx)) } fn ms(n: u64) -> Duration { Duration::from_millis(n) } /// Helper struct to test the [tokio::time::Interval::poll_tick()] method. /// /// `poll_tick()` should register the waker in the context only if it returns /// `Poll::Pending`, not when returning `Poll::Ready`. This struct contains an /// interval timer and counts up on every tick when used as stream. When the /// counter is a multiple of four, it yields the current counter value. /// Depending on the value for `wake_on_pending`, it will reschedule itself when /// it returns `Poll::Pending` or not. When used with `wake_on_pending=false`, /// we expect that the stream stalls because the timer will **not** reschedule /// the next wake-up itself once it returned `Poll::Ready`. struct IntervalStreamer { counter: u32, timer: Interval, wake_on_pending: bool, } impl Stream for IntervalStreamer { type Item = u32; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = Pin::into_inner(self); if this.counter > 12 { return Poll::Ready(None); } match this.timer.poll_tick(cx) { Poll::Pending => Poll::Pending, Poll::Ready(_) => { this.counter += 1; if this.counter % 4 == 0 { Poll::Ready(Some(this.counter)) } else { if this.wake_on_pending { // Schedule this task for wake-up cx.waker().wake_by_ref(); } Poll::Pending } } } } } #[tokio::test(start_paused = true)] async fn stream_with_interval_poll_tick_self_waking() { let stream = IntervalStreamer { counter: 0, timer: tokio::time::interval(tokio::time::Duration::from_millis(10)), wake_on_pending: true, }; let (res_tx, mut res_rx) = tokio::sync::mpsc::channel(12); // Wrap task in timeout so that it will finish eventually even if the stream // stalls. tokio::spawn(tokio::time::timeout( tokio::time::Duration::from_millis(150), async move { tokio::pin!(stream); while let Some(item) = stream.next().await { res_tx.send(item).await.ok(); } }, )); let mut items = Vec::with_capacity(3); while let Some(result) = res_rx.recv().await { items.push(result); } // We expect the stream to yield normally and thus three items. assert_eq!(items, vec![4, 8, 12]); } #[tokio::test(start_paused = true)] async fn stream_with_interval_poll_tick_no_waking() { let stream = IntervalStreamer { counter: 0, timer: tokio::time::interval(tokio::time::Duration::from_millis(10)), wake_on_pending: false, }; let (res_tx, mut res_rx) = tokio::sync::mpsc::channel(12); // Wrap task in timeout so that it will finish eventually even if the stream // stalls. tokio::spawn(tokio::time::timeout( tokio::time::Duration::from_millis(150), async move { tokio::pin!(stream); while let Some(item) = stream.next().await { res_tx.send(item).await.ok(); } }, )); let mut items = Vec::with_capacity(0); while let Some(result) = res_rx.recv().await { items.push(result); } // We expect the stream to stall because it does not reschedule itself on // `Poll::Pending` and neither does [tokio::time::Interval] reschedule the // task when returning `Poll::Ready`. assert_eq!(items, vec![]); } #[tokio::test(start_paused = true)] async fn interval_doesnt_panic_max_duration_when_polling() { let mut timer = task::spawn(time::interval(Duration::MAX)); assert_ready!(timer.enter(|cx, mut timer| timer.poll_tick(cx))); } tokio-1.43.1/tests/time_panic.rs000064400000000000000000000044741046102023000146700ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery #![cfg(panic = "unwind")] use futures::future; use std::error::Error; use std::time::Duration; use tokio::runtime::{Builder, Runtime}; use tokio::time::{self, interval, interval_at, timeout, Instant}; mod support { pub mod panic; } use support::panic::test_panic; #[test] fn pause_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); rt.block_on(async { time::pause(); time::pause(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn resume_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let rt = current_thread(); rt.block_on(async { time::resume(); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn interval_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = interval(Duration::from_millis(0)); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn interval_at_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { let _ = interval_at(Instant::now(), Duration::from_millis(0)); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } #[test] fn timeout_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { // Runtime without `enable_time` so it has no current timer set. let rt = Builder::new_current_thread().build().unwrap(); rt.block_on(async { let _timeout = timeout(Duration::from_millis(5), future::pending::<()>()); }); }); // The panic location should be in this file assert_eq!(&panic_location_file.unwrap(), file!()); Ok(()) } fn current_thread() -> Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap() } tokio-1.43.1/tests/time_pause.rs000064400000000000000000000203531046102023000147050ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(not(miri))] // Too slow on miri. use rand::SeedableRng; use rand::{rngs::StdRng, Rng}; use tokio::time::{self, Duration, Instant, Sleep}; use tokio_test::{assert_elapsed, assert_pending, assert_ready, assert_ready_eq, task}; #[cfg(not(target_os = "wasi"))] use tokio_test::assert_err; use std::{ future::Future, pin::Pin, task::{Context, Poll}, }; #[tokio::test] async fn pause_time_in_main() { tokio::time::pause(); } #[tokio::test] async fn pause_time_in_task() { let t = tokio::spawn(async { tokio::time::pause(); }); t.await.unwrap(); } #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[should_panic] async fn pause_time_in_main_threads() { tokio::time::pause(); } #[cfg_attr(panic = "abort", ignore)] #[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn pause_time_in_spawn_threads() { let t = tokio::spawn(async { tokio::time::pause(); }); assert_err!(t.await); } #[test] fn paused_time_is_deterministic() { let run_1 = paused_time_stress_run(); let run_2 = paused_time_stress_run(); assert_eq!(run_1, run_2); } #[tokio::main(flavor = "current_thread", start_paused = true)] async fn paused_time_stress_run() -> Vec { let mut rng = StdRng::seed_from_u64(1); let mut times = vec![]; let start = Instant::now(); for _ in 0..10_000 { let sleep = rng.gen_range(Duration::from_secs(0)..Duration::from_secs(1)); time::sleep(sleep).await; times.push(start.elapsed()); } times } #[tokio::test(start_paused = true)] async fn advance_after_poll() { time::sleep(ms(1)).await; let start = Instant::now(); let mut sleep = task::spawn(time::sleep_until(start + ms(300))); assert_pending!(sleep.poll()); let before = Instant::now(); time::advance(ms(100)).await; assert_elapsed!(before, ms(100)); assert_pending!(sleep.poll()); } #[tokio::test(start_paused = true)] async fn sleep_no_poll() { let start = Instant::now(); // TODO: Skip this time::advance(ms(1)).await; let mut sleep = task::spawn(time::sleep_until(start + ms(300))); let before = Instant::now(); time::advance(ms(100)).await; assert_elapsed!(before, ms(100)); assert_pending!(sleep.poll()); } enum State { Begin, AwaitingAdvance(Pin>>), AfterAdvance, } struct Tester { sleep: Pin>, state: State, before: Option, poll: bool, } impl Future for Tester { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match &mut self.state { State::Begin => { if self.poll { assert_pending!(self.sleep.as_mut().poll(cx)); } self.before = Some(Instant::now()); let advance_fut = Box::pin(time::advance(ms(100))); self.state = State::AwaitingAdvance(advance_fut); self.poll(cx) } State::AwaitingAdvance(ref mut advance_fut) => match advance_fut.as_mut().poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(()) => { self.state = State::AfterAdvance; self.poll(cx) } }, State::AfterAdvance => { assert_elapsed!(self.before.unwrap(), ms(100)); assert_pending!(self.sleep.as_mut().poll(cx)); Poll::Ready(()) } } } } #[tokio::test(start_paused = true)] async fn sleep_same_task() { let start = Instant::now(); // TODO: Skip this time::advance(ms(1)).await; let sleep = Box::pin(time::sleep_until(start + ms(300))); Tester { sleep, state: State::Begin, before: None, poll: true, } .await; } #[tokio::test(start_paused = true)] async fn sleep_same_task_no_poll() { let start = Instant::now(); // TODO: Skip this time::advance(ms(1)).await; let sleep = Box::pin(time::sleep_until(start + ms(300))); Tester { sleep, state: State::Begin, before: None, poll: false, } .await; } #[tokio::test(start_paused = true)] async fn interval() { let start = Instant::now(); // TODO: Skip this time::advance(ms(1)).await; let mut i = task::spawn(time::interval_at(start, ms(300))); assert_ready_eq!(poll_next(&mut i), start); assert_pending!(poll_next(&mut i)); let before = Instant::now(); time::advance(ms(100)).await; assert_elapsed!(before, ms(100)); assert_pending!(poll_next(&mut i)); let before = Instant::now(); time::advance(ms(200)).await; assert_elapsed!(before, ms(200)); assert_ready_eq!(poll_next(&mut i), start + ms(300)); assert_pending!(poll_next(&mut i)); let before = Instant::now(); time::advance(ms(400)).await; assert_elapsed!(before, ms(400)); assert_ready_eq!(poll_next(&mut i), start + ms(600)); assert_pending!(poll_next(&mut i)); let before = Instant::now(); time::advance(ms(500)).await; assert_elapsed!(before, ms(500)); assert_ready_eq!(poll_next(&mut i), start + ms(900)); assert_ready_eq!(poll_next(&mut i), start + ms(1200)); assert_pending!(poll_next(&mut i)); } #[tokio::test(start_paused = true)] async fn test_time_advance_sub_ms() { let now = Instant::now(); let dur = Duration::from_micros(51_592); time::advance(dur).await; assert_eq!(now.elapsed(), dur); let now = Instant::now(); let dur = Duration::from_micros(1); time::advance(dur).await; assert_eq!(now.elapsed(), dur); } #[tokio::test(start_paused = true)] async fn test_time_advance_3ms_and_change() { let now = Instant::now(); let dur = Duration::from_micros(3_141_592); time::advance(dur).await; assert_eq!(now.elapsed(), dur); let now = Instant::now(); let dur = Duration::from_micros(3_123_456); time::advance(dur).await; assert_eq!(now.elapsed(), dur); } #[tokio::test(start_paused = true)] async fn regression_3710_with_submillis_advance() { let start = Instant::now(); time::advance(Duration::from_millis(1)).await; let mut sleep = task::spawn(time::sleep_until(start + Duration::from_secs(60))); assert_pending!(sleep.poll()); let before = Instant::now(); let dur = Duration::from_micros(51_592); time::advance(dur).await; assert_eq!(before.elapsed(), dur); assert_pending!(sleep.poll()); } #[tokio::test(start_paused = true)] async fn exact_1ms_advance() { let now = Instant::now(); let dur = Duration::from_millis(1); time::advance(dur).await; assert_eq!(now.elapsed(), dur); let now = Instant::now(); let dur = Duration::from_millis(1); time::advance(dur).await; assert_eq!(now.elapsed(), dur); } #[tokio::test(start_paused = true)] async fn advance_once_with_timer() { let mut sleep = task::spawn(time::sleep(Duration::from_millis(1))); assert_pending!(sleep.poll()); time::advance(Duration::from_micros(250)).await; assert_pending!(sleep.poll()); time::advance(Duration::from_micros(1500)).await; assert!(sleep.is_woken()); assert_ready!(sleep.poll()); } #[tokio::test(start_paused = true)] async fn advance_multi_with_timer() { // Round to the nearest ms // time::sleep(Duration::from_millis(1)).await; let mut sleep = task::spawn(time::sleep(Duration::from_millis(1))); assert_pending!(sleep.poll()); time::advance(Duration::from_micros(250)).await; assert_pending!(sleep.poll()); time::advance(Duration::from_micros(250)).await; assert_pending!(sleep.poll()); time::advance(Duration::from_micros(250)).await; assert_pending!(sleep.poll()); time::advance(Duration::from_micros(250)).await; assert!(sleep.is_woken()); assert_ready!(sleep.poll()); } fn poll_next(interval: &mut task::Spawn) -> Poll { interval.enter(|cx, mut interval| interval.poll_tick(cx)) } fn ms(n: u64) -> Duration { Duration::from_millis(n) } tokio-1.43.1/tests/time_rt.rs000064400000000000000000000040231046102023000142110ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::time::*; use std::sync::mpsc; #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] // Wasi doesn't support threads #[test] fn timer_with_threaded_runtime() { use tokio::runtime::Runtime; let rt = Runtime::new().unwrap(); let (tx, rx) = mpsc::channel(); rt.spawn(async move { let when = Instant::now() + Duration::from_millis(10); sleep_until(when).await; assert!(Instant::now() >= when); tx.send(()).unwrap(); }); rx.recv().unwrap(); } #[test] fn timer_with_current_thread_scheduler() { use tokio::runtime::Builder; let rt = Builder::new_current_thread().enable_all().build().unwrap(); let (tx, rx) = mpsc::channel(); rt.block_on(async move { let when = Instant::now() + Duration::from_millis(10); sleep_until(when).await; assert!(Instant::now() >= when); tx.send(()).unwrap(); }); rx.recv().unwrap(); } #[tokio::test] async fn starving() { use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; struct Starve + Unpin>(T, u64); impl + Unpin> Future for Starve { type Output = u64; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { if Pin::new(&mut self.0).poll(cx).is_ready() { return Poll::Ready(self.1); } self.1 += 1; cx.waker().wake_by_ref(); Poll::Pending } } let when = Instant::now() + Duration::from_millis(10); let starve = Starve(Box::pin(sleep_until(when)), 0); starve.await; assert!(Instant::now() >= when); } #[tokio::test] async fn timeout_value() { use tokio::sync::oneshot; let (_tx, rx) = oneshot::channel::<()>(); let now = Instant::now(); let dur = Duration::from_millis(10); let res = timeout(dur, rx).await; assert!(res.is_err()); assert!(Instant::now() >= now + dur); } tokio-1.43.1/tests/time_sleep.rs000064400000000000000000000203441046102023000147000ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(not(miri))] // Too slow on Miri. use std::future::Future; use std::task::Context; use futures::task::noop_waker_ref; use tokio::time::{self, Duration, Instant}; use tokio_test::{assert_elapsed, assert_pending, assert_ready, task}; #[tokio::test] async fn immediate_sleep() { time::pause(); let now = Instant::now(); // Ready! time::sleep_until(now).await; assert_elapsed!(now, ms(1)); } #[tokio::test] async fn is_elapsed() { time::pause(); let sleep = time::sleep(Duration::from_millis(10)); tokio::pin!(sleep); assert!(!sleep.is_elapsed()); assert!(futures::poll!(sleep.as_mut()).is_pending()); assert!(!sleep.is_elapsed()); sleep.as_mut().await; assert!(sleep.is_elapsed()); } #[tokio::test] async fn delayed_sleep_level_0() { time::pause(); for &i in &[1, 10, 60] { let now = Instant::now(); let dur = ms(i); time::sleep_until(now + dur).await; assert_elapsed!(now, dur); } } #[tokio::test] async fn sub_ms_delayed_sleep() { time::pause(); for _ in 0..5 { let now = Instant::now(); let deadline = now + ms(1) + Duration::new(0, 1); time::sleep_until(deadline).await; assert_elapsed!(now, ms(1)); } } #[tokio::test] async fn delayed_sleep_wrapping_level_0() { time::pause(); time::sleep(ms(5)).await; let now = Instant::now(); time::sleep_until(now + ms(60)).await; assert_elapsed!(now, ms(60)); } #[tokio::test] async fn reset_future_sleep_before_fire() { time::pause(); let now = Instant::now(); let mut sleep = task::spawn(Box::pin(time::sleep_until(now + ms(100)))); assert_pending!(sleep.poll()); let mut sleep = sleep.into_inner(); sleep.as_mut().reset(Instant::now() + ms(200)); sleep.await; assert_elapsed!(now, ms(200)); } #[tokio::test] async fn reset_past_sleep_before_turn() { time::pause(); let now = Instant::now(); let mut sleep = task::spawn(Box::pin(time::sleep_until(now + ms(100)))); assert_pending!(sleep.poll()); let mut sleep = sleep.into_inner(); sleep.as_mut().reset(now + ms(80)); sleep.await; assert_elapsed!(now, ms(80)); } #[tokio::test] async fn reset_past_sleep_before_fire() { time::pause(); let now = Instant::now(); let mut sleep = task::spawn(Box::pin(time::sleep_until(now + ms(100)))); assert_pending!(sleep.poll()); let mut sleep = sleep.into_inner(); time::sleep(ms(10)).await; sleep.as_mut().reset(now + ms(80)); sleep.await; assert_elapsed!(now, ms(80)); } #[tokio::test] async fn reset_future_sleep_after_fire() { time::pause(); let now = Instant::now(); let mut sleep = Box::pin(time::sleep_until(now + ms(100))); sleep.as_mut().await; assert_elapsed!(now, ms(100)); sleep.as_mut().reset(now + ms(110)); sleep.await; assert_elapsed!(now, ms(110)); } #[tokio::test] async fn reset_sleep_to_past() { time::pause(); let now = Instant::now(); let mut sleep = task::spawn(Box::pin(time::sleep_until(now + ms(100)))); assert_pending!(sleep.poll()); time::sleep(ms(50)).await; assert!(!sleep.is_woken()); sleep.as_mut().reset(now + ms(40)); // TODO: is this required? //assert!(sleep.is_woken()); assert_ready!(sleep.poll()); } #[cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery #[test] #[should_panic] fn creating_sleep_outside_of_context() { let now = Instant::now(); // This creates a delay outside of the context of a mock timer. This tests // that it will panic. let _fut = time::sleep_until(now + ms(500)); } #[tokio::test] async fn greater_than_max() { const YR_5: u64 = 5 * 365 * 24 * 60 * 60 * 1000; time::pause(); time::sleep_until(Instant::now() + ms(YR_5)).await; } #[tokio::test] async fn short_sleeps() { for _ in 0..1000 { tokio::time::sleep(std::time::Duration::from_millis(0)).await; } } #[tokio::test] async fn multi_long_sleeps() { tokio::time::pause(); for _ in 0..5u32 { tokio::time::sleep(Duration::from_secs( // about a year 365 * 24 * 3600, )) .await; } let deadline = tokio::time::Instant::now() + Duration::from_secs( // about 10 years 10 * 365 * 24 * 3600, ); tokio::time::sleep_until(deadline).await; assert!(tokio::time::Instant::now() >= deadline); } #[tokio::test] async fn long_sleeps() { tokio::time::pause(); let deadline = tokio::time::Instant::now() + Duration::from_secs( // about 10 years 10 * 365 * 24 * 3600, ); tokio::time::sleep_until(deadline).await; assert!(tokio::time::Instant::now() >= deadline); assert!(tokio::time::Instant::now() <= deadline + Duration::from_millis(1)); } #[tokio::test] async fn reset_after_firing() { let timer = tokio::time::sleep(std::time::Duration::from_millis(1)); tokio::pin!(timer); let deadline = timer.deadline(); timer.as_mut().await; assert_ready!(timer .as_mut() .poll(&mut Context::from_waker(noop_waker_ref()))); timer .as_mut() .reset(tokio::time::Instant::now() + std::time::Duration::from_secs(600)); assert_ne!(deadline, timer.deadline()); assert_pending!(timer .as_mut() .poll(&mut Context::from_waker(noop_waker_ref()))); assert_pending!(timer .as_mut() .poll(&mut Context::from_waker(noop_waker_ref()))); } const NUM_LEVELS: usize = 6; const MAX_DURATION: u64 = (1 << (6 * NUM_LEVELS)) - 1; #[tokio::test] async fn exactly_max() { time::pause(); time::sleep(ms(MAX_DURATION)).await; } #[tokio::test] async fn issue_5183() { time::pause(); let big = std::time::Duration::from_secs(u64::MAX / 10); // This is a workaround since awaiting sleep(big) will never finish. #[rustfmt::skip] tokio::select! { biased; _ = tokio::time::sleep(big) => {} _ = tokio::time::sleep(std::time::Duration::from_nanos(1)) => {} } } #[tokio::test] async fn no_out_of_bounds_close_to_max() { time::pause(); time::sleep(ms(MAX_DURATION - 1)).await; } fn ms(n: u64) -> Duration { Duration::from_millis(n) } #[tokio::test] async fn drop_after_reschedule_at_new_scheduled_time() { use futures::poll; tokio::time::pause(); let start = tokio::time::Instant::now(); let mut a = Box::pin(tokio::time::sleep(Duration::from_millis(5))); let mut b = Box::pin(tokio::time::sleep(Duration::from_millis(5))); let mut c = Box::pin(tokio::time::sleep(Duration::from_millis(10))); let _ = poll!(&mut a); let _ = poll!(&mut b); let _ = poll!(&mut c); b.as_mut().reset(start + Duration::from_millis(10)); a.await; drop(b); } #[tokio::test] async fn drop_from_wake() { use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::task::Context; let panicked = Arc::new(AtomicBool::new(false)); let list: Arc>>>> = Arc::new(Mutex::new(Vec::new())); let arc_wake = Arc::new(DropWaker(panicked.clone(), list.clone())); let arc_wake = futures::task::waker(arc_wake); tokio::time::pause(); { let mut lock = list.lock().unwrap(); for _ in 0..100 { let mut timer = Box::pin(tokio::time::sleep(Duration::from_millis(10))); let _ = timer.as_mut().poll(&mut Context::from_waker(&arc_wake)); lock.push(timer); } } tokio::time::sleep(Duration::from_millis(11)).await; assert!( !panicked.load(Ordering::SeqCst), "panicked when dropping timers" ); #[derive(Clone)] struct DropWaker( Arc, Arc>>>>, ); impl futures::task::ArcWake for DropWaker { fn wake_by_ref(arc_self: &Arc) { let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { *arc_self.1.lock().expect("panic in lock") = Vec::new() })); if result.is_err() { arc_self.0.store(true, Ordering::SeqCst); } } } } tokio-1.43.1/tests/time_timeout.rs000064400000000000000000000066061046102023000152630ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] use tokio::sync::oneshot; use tokio::time::{self, timeout, timeout_at, Instant}; use tokio_test::*; use futures::future::pending; use std::time::Duration; #[tokio::test] async fn simultaneous_deadline_future_completion() { // Create a future that is immediately ready let mut fut = task::spawn(timeout_at(Instant::now(), async {})); // Ready! assert_ready_ok!(fut.poll()); } #[cfg_attr(target_os = "wasi", ignore = "FIXME: `fut.poll()` panics on Wasi")] #[tokio::test] async fn completed_future_past_deadline() { // Wrap it with a deadline let mut fut = task::spawn(timeout_at(Instant::now() - ms(1000), async {})); // Ready! assert_ready_ok!(fut.poll()); } #[tokio::test] async fn future_and_deadline_in_future() { time::pause(); // Not yet complete let (tx, rx) = oneshot::channel(); // Wrap it with a deadline let mut fut = task::spawn(timeout_at(Instant::now() + ms(100), rx)); assert_pending!(fut.poll()); // Turn the timer, it runs for the elapsed time time::advance(ms(90)).await; assert_pending!(fut.poll()); // Complete the future tx.send(()).unwrap(); assert!(fut.is_woken()); assert_ready_ok!(fut.poll()).unwrap(); } #[tokio::test] async fn future_and_timeout_in_future() { time::pause(); // Not yet complete let (tx, rx) = oneshot::channel(); // Wrap it with a deadline let mut fut = task::spawn(timeout(ms(100), rx)); // Ready! assert_pending!(fut.poll()); // Turn the timer, it runs for the elapsed time time::advance(ms(90)).await; assert_pending!(fut.poll()); // Complete the future tx.send(()).unwrap(); assert_ready_ok!(fut.poll()).unwrap(); } #[tokio::test] async fn very_large_timeout() { time::pause(); // Not yet complete let (tx, rx) = oneshot::channel(); // copy-paste unstable `Duration::MAX` let duration_max = Duration::from_secs(u64::MAX) + Duration::from_nanos(999_999_999); // Wrap it with a deadline let mut fut = task::spawn(timeout(duration_max, rx)); // Ready! assert_pending!(fut.poll()); // Turn the timer, it runs for the elapsed time time::advance(Duration::from_secs(86400 * 365 * 10)).await; assert_pending!(fut.poll()); // Complete the future tx.send(()).unwrap(); assert_ready_ok!(fut.poll()).unwrap(); } #[tokio::test] async fn deadline_now_elapses() { use futures::future::pending; time::pause(); // Wrap it with a deadline let mut fut = task::spawn(timeout_at(Instant::now(), pending::<()>())); // Factor in jitter // TODO: don't require this time::advance(ms(1)).await; assert_ready_err!(fut.poll()); } #[tokio::test] async fn deadline_future_elapses() { time::pause(); // Wrap it with a deadline let mut fut = task::spawn(timeout_at(Instant::now() + ms(300), pending::<()>())); assert_pending!(fut.poll()); time::advance(ms(301)).await; assert!(fut.is_woken()); assert_ready_err!(fut.poll()); } fn ms(n: u64) -> Duration { Duration::from_millis(n) } #[tokio::test] async fn timeout_is_not_exhausted_by_future() { let fut = timeout(ms(1), async { let mut buffer = [0u8; 1]; loop { use tokio::io::AsyncReadExt; let _ = tokio::io::empty().read(&mut buffer).await; } }); assert!(fut.await.is_err()); } tokio-1.43.1/tests/tracing_sync.rs000064400000000000000000000227651046102023000152460ustar 00000000000000//! Tests for sync instrumentation. //! //! These tests ensure that the instrumentation for tokio //! synchronization primitives is correct. #![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(tokio_unstable, feature = "tracing", target_has_atomic = "64"))] use tokio::sync; use tracing_mock::{expect, subscriber}; #[tokio::test] async fn test_barrier_creates_span() { let barrier_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::barrier"); let size_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("size").with_value(&1_u64)); let arrived_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("arrived").with_value(&0_i64)); let (subscriber, handle) = subscriber::mock() .new_span( barrier_span .clone() .with_ancestry(expect::is_explicit_root()), ) .enter(&barrier_span) .event(size_event) .event(arrived_event) .exit(&barrier_span) .drop_span(&barrier_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); let _ = sync::Barrier::new(1); } handle.assert_finished(); } #[tokio::test] async fn test_mutex_creates_span() { let mutex_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::mutex"); let locked_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("locked").with_value(&false)); let batch_semaphore_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::batch_semaphore"); let batch_semaphore_permits_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("permits").with_value(&1u64)) .with_fields(expect::field("permits.op").with_value(&"override")); let (subscriber, handle) = subscriber::mock() .new_span(mutex_span.clone().with_ancestry(expect::is_explicit_root())) .enter(&mutex_span) .event(locked_event) .new_span( batch_semaphore_span .clone() .with_ancestry(expect::is_explicit_root()), ) .enter(&batch_semaphore_span) .event(batch_semaphore_permits_event) .exit(&batch_semaphore_span) .exit(&mutex_span) .drop_span(&mutex_span) .drop_span(&batch_semaphore_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); let _ = sync::Mutex::new(true); } handle.assert_finished(); } #[tokio::test] async fn test_oneshot_creates_span() { let oneshot_span_id = expect::id(); let oneshot_span = expect::span() .with_id(oneshot_span_id.clone()) .named("runtime.resource") .with_target("tokio::sync::oneshot"); let initial_tx_dropped_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("tx_dropped").with_value(&false)) .with_fields(expect::field("tx_dropped.op").with_value(&"override")); let final_tx_dropped_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("tx_dropped").with_value(&true)) .with_fields(expect::field("tx_dropped.op").with_value(&"override")); let initial_rx_dropped_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("rx_dropped").with_value(&false)) .with_fields(expect::field("rx_dropped.op").with_value(&"override")); let final_rx_dropped_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("rx_dropped").with_value(&true)) .with_fields(expect::field("rx_dropped.op").with_value(&"override")); let value_sent_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("value_sent").with_value(&false)) .with_fields(expect::field("value_sent.op").with_value(&"override")); let value_received_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("value_received").with_value(&false)) .with_fields(expect::field("value_received.op").with_value(&"override")); let async_op_span_id = expect::id(); let async_op_span = expect::span() .with_id(async_op_span_id.clone()) .named("runtime.resource.async_op") .with_target("tokio::sync::oneshot"); let async_op_poll_span = expect::span() .named("runtime.resource.async_op.poll") .with_target("tokio::sync::oneshot"); let (subscriber, handle) = subscriber::mock() .new_span( oneshot_span .clone() .with_ancestry(expect::is_explicit_root()), ) .enter(&oneshot_span) .event(initial_tx_dropped_event) .exit(&oneshot_span) .enter(&oneshot_span) .event(initial_rx_dropped_event) .exit(&oneshot_span) .enter(&oneshot_span) .event(value_sent_event) .exit(&oneshot_span) .enter(&oneshot_span) .event(value_received_event) .exit(&oneshot_span) .enter(&oneshot_span) .new_span( async_op_span .clone() .with_ancestry(expect::has_contextual_parent(&oneshot_span_id)), ) .exit(&oneshot_span) .enter(&async_op_span) .new_span( async_op_poll_span .clone() .with_ancestry(expect::has_contextual_parent(&async_op_span_id)), ) .exit(&async_op_span) .enter(&oneshot_span) .event(final_tx_dropped_event) .exit(&oneshot_span) .enter(&oneshot_span) .event(final_rx_dropped_event) .exit(&oneshot_span) .drop_span(oneshot_span) .drop_span(async_op_span) .drop_span(&async_op_poll_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); let _ = sync::oneshot::channel::(); } handle.assert_finished(); } #[tokio::test] async fn test_rwlock_creates_span() { let rwlock_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::rwlock"); let max_readers_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("max_readers").with_value(&0x1FFFFFFF_u64)); let write_locked_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("write_locked").with_value(&false)); let current_readers_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("current_readers").with_value(&0_i64)); let batch_semaphore_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::batch_semaphore"); let batch_semaphore_permits_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("permits").with_value(&1u64)) .with_fields(expect::field("permits.op").with_value(&"override")); let (subscriber, handle) = subscriber::mock() .new_span( rwlock_span .clone() .with_ancestry(expect::is_explicit_root()), ) .enter(rwlock_span.clone()) .event(max_readers_event) .event(write_locked_event) .event(current_readers_event) .exit(rwlock_span.clone()) .enter(rwlock_span.clone()) .new_span(batch_semaphore_span.clone()) .enter(batch_semaphore_span.clone()) .event(batch_semaphore_permits_event) .exit(batch_semaphore_span.clone()) .exit(rwlock_span.clone()) .drop_span(rwlock_span) .drop_span(batch_semaphore_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); let _ = sync::RwLock::new(true); } handle.assert_finished(); } #[tokio::test] async fn test_semaphore_creates_span() { let semaphore_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::semaphore"); let batch_semaphore_span = expect::span() .named("runtime.resource") .with_target("tokio::sync::batch_semaphore"); let batch_semaphore_permits_event = expect::event() .with_target("runtime::resource::state_update") .with_fields(expect::field("permits").with_value(&1u64)) .with_fields(expect::field("permits.op").with_value(&"override")); let (subscriber, handle) = subscriber::mock() .new_span( semaphore_span .clone() .with_ancestry(expect::is_explicit_root()), ) .enter(semaphore_span.clone()) .new_span(batch_semaphore_span.clone()) .enter(batch_semaphore_span.clone()) .event(batch_semaphore_permits_event) .exit(batch_semaphore_span.clone()) .exit(semaphore_span.clone()) .drop_span(semaphore_span) .drop_span(batch_semaphore_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); let _ = sync::Semaphore::new(1); } handle.assert_finished(); } tokio-1.43.1/tests/tracing_task.rs000064400000000000000000000124471046102023000152300ustar 00000000000000//! Tests for task instrumentation. //! //! These tests ensure that the instrumentation for task spawning and task //! lifecycles is correct. #![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(tokio_unstable, feature = "tracing", target_has_atomic = "64"))] use std::{mem, time::Duration}; use tokio::task; use tracing_mock::{expect, span::NewSpan, subscriber}; #[tokio::test] async fn task_spawn_creates_span() { let task_span = expect::span() .named("runtime.spawn") .with_target("tokio::task"); let (subscriber, handle) = subscriber::mock() .new_span(&task_span) .enter(&task_span) .exit(&task_span) // The task span is entered once more when it gets dropped .enter(&task_span) .exit(&task_span) .drop_span(task_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); tokio::spawn(futures::future::ready(())) .await .expect("failed to await join handle"); } handle.assert_finished(); } #[tokio::test] async fn task_spawn_loc_file_recorded() { let task_span = expect::span() .named("runtime.spawn") .with_target("tokio::task") .with_fields(expect::field("loc.file").with_value(&file!())); let (subscriber, handle) = subscriber::mock().new_span(task_span).run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); tokio::spawn(futures::future::ready(())) .await .expect("failed to await join handle"); } handle.assert_finished(); } #[tokio::test] async fn task_builder_name_recorded() { let task_span = expect_task_named("test-task"); let (subscriber, handle) = subscriber::mock().new_span(task_span).run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); task::Builder::new() .name("test-task") .spawn(futures::future::ready(())) .unwrap() .await .expect("failed to await join handle"); } handle.assert_finished(); } #[tokio::test] async fn task_builder_loc_file_recorded() { let task_span = expect::span() .named("runtime.spawn") .with_target("tokio::task") .with_fields(expect::field("loc.file").with_value(&file!())); let (subscriber, handle) = subscriber::mock().new_span(task_span).run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); task::Builder::new() .spawn(futures::future::ready(())) .unwrap() .await .expect("failed to await join handle"); } handle.assert_finished(); } #[tokio::test] async fn task_spawn_sizes_recorded() { let future = futures::future::ready(()); let size = mem::size_of_val(&future) as u64; let task_span = expect::span() .named("runtime.spawn") .with_target("tokio::task") // TODO(hds): check that original_size.bytes is NOT recorded when this can be done in // tracing-mock without listing every other field. .with_fields(expect::field("size.bytes").with_value(&size)); let (subscriber, handle) = subscriber::mock().new_span(task_span).run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); task::Builder::new() .spawn(future) .unwrap() .await .expect("failed to await join handle"); } handle.assert_finished(); } #[tokio::test] async fn task_big_spawn_sizes_recorded() { let future = { async fn big() { let mut a = [0_u8; N]; for (idx, item) in a.iter_mut().enumerate() { *item = (idx % 256) as u8; } tokio::time::sleep(Duration::from_millis(10)).await; for (idx, item) in a.iter_mut().enumerate() { assert_eq!(*item, (idx % 256) as u8); } } // This is larger than the release auto-boxing threshold big::<20_000>() }; fn boxed_size(_: &T) -> usize { mem::size_of::>() } let size = mem::size_of_val(&future) as u64; let boxed_size = boxed_size(&future); let task_span = expect::span() .named("runtime.spawn") .with_target("tokio::task") .with_fields( expect::field("size.bytes") .with_value(&boxed_size) .and(expect::field("original_size.bytes").with_value(&size)), ); let (subscriber, handle) = subscriber::mock().new_span(task_span).run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); task::Builder::new() .spawn(future) .unwrap() .await .expect("failed to await join handle"); } handle.assert_finished(); } /// Expect a task with name /// /// This is a convenience function to create the expectation for a new task /// with the `task.name` field set to the provided name. fn expect_task_named(name: &str) -> NewSpan { expect::span() .named("runtime.spawn") .with_target("tokio::task") .with_fields( expect::field("task.name").with_value(&tracing::field::debug(format_args!("{}", name))), ) } tokio-1.43.1/tests/tracing_time.rs000064400000000000000000000047711046102023000152250ustar 00000000000000//! Tests for time resource instrumentation. //! //! These tests ensure that the instrumentation for tokio //! synchronization primitives is correct. #![allow(unknown_lints, unexpected_cfgs)] #![warn(rust_2018_idioms)] #![cfg(all(tokio_unstable, feature = "tracing", target_has_atomic = "64"))] use std::time::Duration; use tracing_mock::{expect, subscriber}; #[tokio::test] async fn test_sleep_creates_span() { let sleep_span_id = expect::id(); let sleep_span = expect::span() .with_id(sleep_span_id.clone()) .named("runtime.resource") .with_target("tokio::time::sleep"); let state_update = expect::event() .with_target("runtime::resource::state_update") .with_fields( expect::field("duration") // FIXME(hds): This value isn't stable and doesn't seem to make sense. We're not // going to test on it until the resource instrumentation has been // refactored and we know that we're getting a stable value here. //.with_value(&(7_u64 + 1)) .and(expect::field("duration.op").with_value(&"override")), ); let async_op_span_id = expect::id(); let async_op_span = expect::span() .with_id(async_op_span_id.clone()) .named("runtime.resource.async_op") .with_target("tokio::time::sleep"); let async_op_poll_span = expect::span() .named("runtime.resource.async_op.poll") .with_target("tokio::time::sleep"); let (subscriber, handle) = subscriber::mock() .new_span(sleep_span.clone().with_ancestry(expect::is_explicit_root())) .enter(sleep_span.clone()) .event(state_update) .new_span( async_op_span .clone() .with_ancestry(expect::has_contextual_parent(&sleep_span_id)) .with_fields(expect::field("source").with_value(&"Sleep::new_timeout")), ) .exit(sleep_span.clone()) .enter(async_op_span.clone()) .new_span( async_op_poll_span .clone() .with_ancestry(expect::has_contextual_parent(&async_op_span_id)), ) .exit(async_op_span.clone()) .drop_span(async_op_span) .drop_span(async_op_poll_span) .drop_span(sleep_span) .run_with_handle(); { let _guard = tracing::subscriber::set_default(subscriber); _ = tokio::time::sleep(Duration::from_millis(7)); } handle.assert_finished(); } tokio-1.43.1/tests/udp.rs000064400000000000000000000466711046102023000133550ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi does not support bind or UDP // No `socket` on miri. use std::future::poll_fn; use std::io; use std::sync::Arc; use tokio::{io::ReadBuf, net::UdpSocket}; use tokio_test::assert_ok; const MSG: &[u8] = b"hello"; const MSG_LEN: usize = MSG.len(); #[tokio::test] async fn send_recv() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; sender.connect(receiver.local_addr()?).await?; receiver.connect(sender.local_addr()?).await?; sender.send(MSG).await?; let mut recv_buf = [0u8; 32]; let len = receiver.recv(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], MSG); Ok(()) } #[tokio::test] async fn send_recv_poll() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; sender.connect(receiver.local_addr()?).await?; receiver.connect(sender.local_addr()?).await?; poll_fn(|cx| sender.poll_send(cx, MSG)).await?; let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); poll_fn(|cx| receiver.poll_recv(cx, &mut read)).await?; assert_eq!(read.filled(), MSG); Ok(()) } #[tokio::test] async fn send_to_recv_from() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let receiver_addr = receiver.local_addr()?; sender.send_to(MSG, &receiver_addr).await?; let mut recv_buf = [0u8; 32]; let (len, addr) = receiver.recv_from(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], MSG); assert_eq!(addr, sender.local_addr()?); Ok(()) } #[tokio::test] async fn send_to_recv_from_poll() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let receiver_addr = receiver.local_addr()?; poll_fn(|cx| sender.poll_send_to(cx, MSG, receiver_addr)).await?; let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); let addr = poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?; assert_eq!(read.filled(), MSG); assert_eq!(addr, sender.local_addr()?); Ok(()) } #[tokio::test] async fn send_to_peek_from() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let receiver_addr = receiver.local_addr()?; poll_fn(|cx| sender.poll_send_to(cx, MSG, receiver_addr)).await?; // peek let mut recv_buf = [0u8; 32]; let (n, addr) = receiver.peek_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..n], MSG); assert_eq!(addr, sender.local_addr()?); // peek let mut recv_buf = [0u8; 32]; let (n, addr) = receiver.peek_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..n], MSG); assert_eq!(addr, sender.local_addr()?); let mut recv_buf = [0u8; 32]; let (n, addr) = receiver.recv_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..n], MSG); assert_eq!(addr, sender.local_addr()?); Ok(()) } #[tokio::test] async fn send_to_try_peek_from() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let receiver_addr = receiver.local_addr()?; poll_fn(|cx| sender.poll_send_to(cx, MSG, receiver_addr)).await?; // peek let mut recv_buf = [0u8; 32]; loop { match receiver.try_peek_from(&mut recv_buf) { Ok((n, addr)) => { assert_eq!(&recv_buf[..n], MSG); assert_eq!(addr, sender.local_addr()?); break; } Err(e) if e.kind() == io::ErrorKind::WouldBlock => { receiver.readable().await?; } Err(e) => return Err(e), } } // peek let mut recv_buf = [0u8; 32]; let (n, addr) = receiver.peek_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..n], MSG); assert_eq!(addr, sender.local_addr()?); let mut recv_buf = [0u8; 32]; let (n, addr) = receiver.recv_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..n], MSG); assert_eq!(addr, sender.local_addr()?); Ok(()) } #[tokio::test] async fn send_to_peek_from_poll() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let receiver_addr = receiver.local_addr()?; poll_fn(|cx| sender.poll_send_to(cx, MSG, receiver_addr)).await?; let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); let addr = poll_fn(|cx| receiver.poll_peek_from(cx, &mut read)).await?; assert_eq!(read.filled(), MSG); assert_eq!(addr, sender.local_addr()?); let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); poll_fn(|cx| receiver.poll_peek_from(cx, &mut read)).await?; assert_eq!(read.filled(), MSG); let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?; assert_eq!(read.filled(), MSG); Ok(()) } #[tokio::test] async fn peek_sender() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let sender_addr = sender.local_addr()?; let receiver_addr = receiver.local_addr()?; let msg = b"Hello, world!"; sender.send_to(msg, receiver_addr).await?; let peeked_sender = receiver.peek_sender().await?; assert_eq!(peeked_sender, sender_addr); // Assert that `peek_sender()` returns the right sender but // doesn't remove from the receive queue. let mut recv_buf = [0u8; 32]; let (read, received_sender) = receiver.recv_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..read], msg); assert_eq!(received_sender, peeked_sender); Ok(()) } #[tokio::test] async fn poll_peek_sender() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let sender_addr = sender.local_addr()?; let receiver_addr = receiver.local_addr()?; let msg = b"Hello, world!"; poll_fn(|cx| sender.poll_send_to(cx, msg, receiver_addr)).await?; let peeked_sender = poll_fn(|cx| receiver.poll_peek_sender(cx)).await?; assert_eq!(peeked_sender, sender_addr); // Assert that `poll_peek_sender()` returns the right sender but // doesn't remove from the receive queue. let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); let received_sender = poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?; assert_eq!(read.filled(), msg); assert_eq!(received_sender, peeked_sender); Ok(()) } #[tokio::test] async fn try_peek_sender() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; let sender_addr = sender.local_addr()?; let receiver_addr = receiver.local_addr()?; let msg = b"Hello, world!"; sender.send_to(msg, receiver_addr).await?; let peeked_sender = loop { match receiver.try_peek_sender() { Ok(peeked_sender) => break peeked_sender, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { receiver.readable().await?; } Err(e) => return Err(e), } }; assert_eq!(peeked_sender, sender_addr); // Assert that `try_peek_sender()` returns the right sender but // didn't remove from the receive queue. let mut recv_buf = [0u8; 32]; // We already peeked the sender so there must be data in the receive queue. let (read, received_sender) = receiver.try_recv_from(&mut recv_buf).unwrap(); assert_eq!(&recv_buf[..read], msg); assert_eq!(received_sender, peeked_sender); Ok(()) } #[tokio::test] async fn split() -> std::io::Result<()> { let socket = UdpSocket::bind("127.0.0.1:0").await?; let s = Arc::new(socket); let r = s.clone(); let addr = s.local_addr()?; tokio::spawn(async move { s.send_to(MSG, &addr).await.unwrap(); }); let mut recv_buf = [0u8; 32]; let (len, _) = r.recv_from(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], MSG); Ok(()) } #[tokio::test] async fn split_chan() -> std::io::Result<()> { // setup UdpSocket that will echo all sent items let socket = UdpSocket::bind("127.0.0.1:0").await?; let addr = socket.local_addr().unwrap(); let s = Arc::new(socket); let r = s.clone(); let (tx, mut rx) = tokio::sync::mpsc::channel::<(Vec, std::net::SocketAddr)>(1_000); tokio::spawn(async move { while let Some((bytes, addr)) = rx.recv().await { s.send_to(&bytes, &addr).await.unwrap(); } }); tokio::spawn(async move { let mut buf = [0u8; 32]; loop { let (len, addr) = r.recv_from(&mut buf).await.unwrap(); tx.send((buf[..len].to_vec(), addr)).await.unwrap(); } }); // test that we can send a value and get back some response let sender = UdpSocket::bind("127.0.0.1:0").await?; sender.send_to(MSG, addr).await?; let mut recv_buf = [0u8; 32]; let (len, _) = sender.recv_from(&mut recv_buf).await?; assert_eq!(&recv_buf[..len], MSG); Ok(()) } #[tokio::test] async fn split_chan_poll() -> std::io::Result<()> { // setup UdpSocket that will echo all sent items let socket = UdpSocket::bind("127.0.0.1:0").await?; let addr = socket.local_addr().unwrap(); let s = Arc::new(socket); let r = s.clone(); let (tx, mut rx) = tokio::sync::mpsc::channel::<(Vec, std::net::SocketAddr)>(1_000); tokio::spawn(async move { while let Some((bytes, addr)) = rx.recv().await { poll_fn(|cx| s.poll_send_to(cx, &bytes, addr)) .await .unwrap(); } }); tokio::spawn(async move { let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); loop { let addr = poll_fn(|cx| r.poll_recv_from(cx, &mut read)).await.unwrap(); tx.send((read.filled().to_vec(), addr)).await.unwrap(); } }); // test that we can send a value and get back some response let sender = UdpSocket::bind("127.0.0.1:0").await?; poll_fn(|cx| sender.poll_send_to(cx, MSG, addr)).await?; let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); let _ = poll_fn(|cx| sender.poll_recv_from(cx, &mut read)).await?; assert_eq!(read.filled(), MSG); Ok(()) } // # Note // // This test is purposely written such that each time `sender` sends data on // the socket, `receiver` awaits the data. On Unix, it would be okay waiting // until the end of the test to receive all the data. On Windows, this would // **not** be okay because it's resources are completion based (via IOCP). // If data is sent and not yet received, attempting to send more data will // result in `ErrorKind::WouldBlock` until the first operation completes. #[tokio::test] async fn try_send_spawn() { const MSG2: &[u8] = b"world!"; const MSG2_LEN: usize = MSG2.len(); let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); receiver .connect(sender.local_addr().unwrap()) .await .unwrap(); sender.writable().await.unwrap(); let sent = &sender .try_send_to(MSG, receiver.local_addr().unwrap()) .unwrap(); assert_eq!(sent, &MSG_LEN); let mut buf = [0u8; 32]; let mut received = receiver.recv(&mut buf[..]).await.unwrap(); sender .connect(receiver.local_addr().unwrap()) .await .unwrap(); let sent = &sender.try_send(MSG2).unwrap(); assert_eq!(sent, &MSG2_LEN); received += receiver.recv(&mut buf[..]).await.unwrap(); std::thread::spawn(move || { let sent = &sender.try_send(MSG).unwrap(); assert_eq!(sent, &MSG_LEN); }) .join() .unwrap(); received += receiver.recv(&mut buf[..]).await.unwrap(); assert_eq!(received, MSG_LEN * 2 + MSG2_LEN); } #[tokio::test] async fn try_send_recv() { // Create listener let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); // Create socket pair let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); // Connect the two client.connect(server.local_addr().unwrap()).await.unwrap(); server.connect(client.local_addr().unwrap()).await.unwrap(); for _ in 0..5 { loop { client.writable().await.unwrap(); match client.try_send(b"hello world") { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { server.readable().await.unwrap(); let mut buf = [0; 512]; match server.try_recv(&mut buf) { Ok(n) => { assert_eq!(n, 11); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } } #[tokio::test] async fn try_send_to_recv_from() { // Create listener let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let saddr = server.local_addr().unwrap(); // Create socket pair let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let caddr = client.local_addr().unwrap(); for _ in 0..5 { loop { client.writable().await.unwrap(); match client.try_send_to(b"hello world", saddr) { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { server.readable().await.unwrap(); let mut buf = [0; 512]; match server.try_recv_from(&mut buf) { Ok((n, addr)) => { assert_eq!(n, 11); assert_eq!(addr, caddr); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } } #[tokio::test] async fn try_recv_buf() { // Create listener let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); // Create socket pair let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); // Connect the two client.connect(server.local_addr().unwrap()).await.unwrap(); server.connect(client.local_addr().unwrap()).await.unwrap(); for _ in 0..5 { loop { client.writable().await.unwrap(); match client.try_send(b"hello world") { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { server.readable().await.unwrap(); let mut buf = Vec::with_capacity(512); match server.try_recv_buf(&mut buf) { Ok(n) => { assert_eq!(n, 11); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } } #[tokio::test] async fn recv_buf() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; sender.connect(receiver.local_addr()?).await?; receiver.connect(sender.local_addr()?).await?; sender.send(MSG).await?; let mut recv_buf = Vec::with_capacity(32); let len = receiver.recv_buf(&mut recv_buf).await?; assert_eq!(len, MSG_LEN); assert_eq!(&recv_buf[..len], MSG); Ok(()) } #[tokio::test] async fn try_recv_buf_from() { // Create listener let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let saddr = server.local_addr().unwrap(); // Create socket pair let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let caddr = client.local_addr().unwrap(); for _ in 0..5 { loop { client.writable().await.unwrap(); match client.try_send_to(b"hello world", saddr) { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { server.readable().await.unwrap(); let mut buf = Vec::with_capacity(512); match server.try_recv_buf_from(&mut buf) { Ok((n, addr)) => { assert_eq!(n, 11); assert_eq!(addr, caddr); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } } #[tokio::test] async fn recv_buf_from() -> std::io::Result<()> { let sender = UdpSocket::bind("127.0.0.1:0").await?; let receiver = UdpSocket::bind("127.0.0.1:0").await?; sender.connect(receiver.local_addr()?).await?; sender.send(MSG).await?; let mut recv_buf = Vec::with_capacity(32); let (len, caddr) = receiver.recv_buf_from(&mut recv_buf).await?; assert_eq!(len, MSG_LEN); assert_eq!(&recv_buf[..len], MSG); assert_eq!(caddr, sender.local_addr()?); Ok(()) } #[tokio::test] async fn poll_ready() { // Create listener let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let saddr = server.local_addr().unwrap(); // Create socket pair let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); let caddr = client.local_addr().unwrap(); for _ in 0..5 { loop { assert_ok!(poll_fn(|cx| client.poll_send_ready(cx)).await); match client.try_send_to(b"hello world", saddr) { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { assert_ok!(poll_fn(|cx| server.poll_recv_ready(cx)).await); let mut buf = Vec::with_capacity(512); match server.try_recv_buf_from(&mut buf) { Ok((n, addr)) => { assert_eq!(n, 11); assert_eq!(addr, caddr); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } } tokio-1.43.1/tests/uds_cred.rs000064400000000000000000000012621046102023000143400ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(all(unix, not(target_os = "dragonfly"), not(miri)))] // No `getsockopt` on miri. use tokio::net::UnixStream; use libc::getegid; use libc::geteuid; #[tokio::test] #[cfg_attr( target_os = "netbsd", ignore = "NetBSD does not support getpeereid() for sockets created by socketpair()" )] async fn test_socket_pair() { let (a, b) = UnixStream::pair().unwrap(); let cred_a = a.peer_cred().unwrap(); let cred_b = b.peer_cred().unwrap(); assert_eq!(cred_a, cred_b); let uid = unsafe { geteuid() }; let gid = unsafe { getegid() }; assert_eq!(cred_a.uid(), uid); assert_eq!(cred_a.gid(), gid); } tokio-1.43.1/tests/uds_datagram.rs000064400000000000000000000303271046102023000152070ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] use tokio::io::ReadBuf; use tokio::net::UnixDatagram; use tokio::try_join; use std::future::poll_fn; use std::io; use std::sync::Arc; async fn echo_server(socket: UnixDatagram) -> io::Result<()> { let mut recv_buf = vec![0u8; 1024]; loop { let (len, peer_addr) = socket.recv_from(&mut recv_buf[..]).await?; if let Some(path) = peer_addr.as_pathname() { socket.send_to(&recv_buf[..len], path).await?; } } } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn echo() -> io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); let client_path = dir.path().join("client.sock"); let server_socket = UnixDatagram::bind(server_path.clone())?; tokio::spawn(async move { let _ = echo_server(server_socket).await; }); { let socket = UnixDatagram::bind(&client_path).unwrap(); socket.connect(server_path)?; socket.send(b"ECHO").await?; let mut recv_buf = [0u8; 16]; let len = socket.recv(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], b"ECHO"); } Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn echo_from() -> io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); let client_path = dir.path().join("client.sock"); let server_socket = UnixDatagram::bind(server_path.clone())?; tokio::spawn(async move { let _ = echo_server(server_socket).await; }); { let socket = UnixDatagram::bind(&client_path).unwrap(); socket.connect(&server_path)?; socket.send(b"ECHO").await?; let mut recv_buf = [0u8; 16]; let (len, addr) = socket.recv_from(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], b"ECHO"); assert_eq!(addr.as_pathname(), Some(server_path.as_path())); } Ok(()) } // Even though we use sync non-blocking io we still need a reactor. #[tokio::test] #[cfg_attr(miri, ignore)] // No SOCK_DGRAM for `socketpair` in miri. async fn try_send_recv_never_block() -> io::Result<()> { let mut recv_buf = [0u8; 16]; let payload = b"PAYLOAD"; let mut count = 0; let (dgram1, dgram2) = UnixDatagram::pair()?; // Send until we hit the OS `net.unix.max_dgram_qlen`. loop { dgram1.writable().await.unwrap(); match dgram1.try_send(payload) { Err(err) => match (err.kind(), err.raw_os_error()) { (io::ErrorKind::WouldBlock, _) => break, (_, Some(libc::ENOBUFS)) => break, _ => { panic!("unexpected error {err:?}"); } }, Ok(len) => { assert_eq!(len, payload.len()); } } count += 1; } // Read every dgram we sent. while count > 0 { dgram2.readable().await.unwrap(); let len = dgram2.try_recv(&mut recv_buf[..])?; assert_eq!(len, payload.len()); assert_eq!(payload, &recv_buf[..len]); count -= 1; } let err = dgram2.try_recv(&mut recv_buf[..]).unwrap_err(); match err.kind() { io::ErrorKind::WouldBlock => (), _ => unreachable!("unexpected error {:?}", err), } Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn split() -> std::io::Result<()> { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("split.sock"); let s = Arc::new(UnixDatagram::bind(path.clone())?); let r = s.clone(); let msg = b"hello"; let ((), ()) = try_join! { async { s.send_to(msg, path).await?; io::Result::Ok(()) }, async { let mut recv_buf = [0u8; 32]; let (len, _) = r.recv_from(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], msg); Ok(()) }, }?; Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn send_to_recv_from_poll() -> std::io::Result<()> { let dir = tempfile::tempdir().unwrap(); let sender_path = dir.path().join("sender.sock"); let receiver_path = dir.path().join("receiver.sock"); let sender = UnixDatagram::bind(&sender_path)?; let receiver = UnixDatagram::bind(&receiver_path)?; let msg = b"hello"; poll_fn(|cx| sender.poll_send_to(cx, msg, &receiver_path)).await?; let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); let addr = poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?; assert_eq!(read.filled(), msg); assert_eq!(addr.as_pathname(), Some(sender_path.as_ref())); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn send_recv_poll() -> std::io::Result<()> { let dir = tempfile::tempdir().unwrap(); let sender_path = dir.path().join("sender.sock"); let receiver_path = dir.path().join("receiver.sock"); let sender = UnixDatagram::bind(&sender_path)?; let receiver = UnixDatagram::bind(&receiver_path)?; sender.connect(&receiver_path)?; receiver.connect(&sender_path)?; let msg = b"hello"; poll_fn(|cx| sender.poll_send(cx, msg)).await?; let mut recv_buf = [0u8; 32]; let mut read = ReadBuf::new(&mut recv_buf); poll_fn(|cx| receiver.poll_recv(cx, &mut read)).await?; assert_eq!(read.filled(), msg); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn try_send_to_recv_from() -> std::io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); let client_path = dir.path().join("client.sock"); // Create listener let server = UnixDatagram::bind(&server_path)?; // Create socket pair let client = UnixDatagram::bind(&client_path)?; for _ in 0..5 { loop { client.writable().await?; match client.try_send_to(b"hello world", &server_path) { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { server.readable().await?; let mut buf = [0; 512]; match server.try_recv_from(&mut buf) { Ok((n, addr)) => { assert_eq!(n, 11); assert_eq!(addr.as_pathname(), Some(client_path.as_ref())); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn try_recv_buf_from() -> std::io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); let client_path = dir.path().join("client.sock"); // Create listener let server = UnixDatagram::bind(&server_path)?; // Create socket pair let client = UnixDatagram::bind(&client_path)?; for _ in 0..5 { loop { client.writable().await?; match client.try_send_to(b"hello world", &server_path) { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { server.readable().await?; let mut buf = Vec::with_capacity(512); match server.try_recv_buf_from(&mut buf) { Ok((n, addr)) => { assert_eq!(n, 11); assert_eq!(addr.as_pathname(), Some(client_path.as_ref())); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn recv_buf_from() -> std::io::Result<()> { let tmp = tempfile::tempdir()?; // Bind each socket to a filesystem path let tx_path = tmp.path().join("tx"); let tx = UnixDatagram::bind(&tx_path)?; let rx_path = tmp.path().join("rx"); let rx = UnixDatagram::bind(&rx_path)?; let bytes = b"hello world"; tx.send_to(bytes, &rx_path).await?; let mut buf = Vec::with_capacity(24); let (size, addr) = rx.recv_buf_from(&mut buf).await?; let dgram = &buf[..size]; assert_eq!(dgram, bytes); assert_eq!(addr.as_pathname().unwrap(), &tx_path); Ok(()) } // Even though we use sync non-blocking io we still need a reactor. #[tokio::test] #[cfg_attr(miri, ignore)] // No SOCK_DGRAM for `socketpair` in miri. async fn try_recv_buf_never_block() -> io::Result<()> { let payload = b"PAYLOAD"; let mut count = 0; let (dgram1, dgram2) = UnixDatagram::pair()?; // Send until we hit the OS `net.unix.max_dgram_qlen`. loop { dgram1.writable().await.unwrap(); match dgram1.try_send(payload) { Err(err) => match (err.kind(), err.raw_os_error()) { (io::ErrorKind::WouldBlock, _) => break, (_, Some(libc::ENOBUFS)) => break, _ => { panic!("unexpected error {err:?}"); } }, Ok(len) => { assert_eq!(len, payload.len()); } } count += 1; } // Read every dgram we sent. while count > 0 { let mut recv_buf = Vec::with_capacity(16); dgram2.readable().await.unwrap(); let len = dgram2.try_recv_buf(&mut recv_buf)?; assert_eq!(len, payload.len()); assert_eq!(payload, &recv_buf[..len]); count -= 1; } let mut recv_buf = vec![0; 16]; let err = dgram2.try_recv_from(&mut recv_buf).unwrap_err(); match err.kind() { io::ErrorKind::WouldBlock => (), _ => unreachable!("unexpected error {:?}", err), } Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No SOCK_DGRAM for `socketpair` in miri. async fn recv_buf() -> std::io::Result<()> { // Create the pair of sockets let (sock1, sock2) = UnixDatagram::pair()?; // Since the sockets are paired, the paired send/recv // functions can be used let bytes = b"hello world"; sock1.send(bytes).await?; let mut buff = Vec::with_capacity(24); let size = sock2.recv_buf(&mut buff).await?; let dgram = &buff[..size]; assert_eq!(dgram, bytes); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` on miri. async fn poll_ready() -> io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); let client_path = dir.path().join("client.sock"); // Create listener let server = UnixDatagram::bind(&server_path)?; // Create socket pair let client = UnixDatagram::bind(&client_path)?; for _ in 0..5 { loop { poll_fn(|cx| client.poll_send_ready(cx)).await?; match client.try_send_to(b"hello world", &server_path) { Ok(n) => { assert_eq!(n, 11); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } loop { poll_fn(|cx| server.poll_recv_ready(cx)).await?; let mut buf = Vec::with_capacity(512); match server.try_recv_buf_from(&mut buf) { Ok((n, addr)) => { assert_eq!(n, 11); assert_eq!(addr.as_pathname(), Some(client_path.as_ref())); assert_eq!(&buf[0..11], &b"hello world"[..]); break; } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("{e:?}"), } } } Ok(()) } tokio-1.43.1/tests/uds_socket.rs000064400000000000000000000063551046102023000147230ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] use futures::future::try_join; use std::io; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::UnixSocket, }; #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn datagram_echo_server() -> io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); let client_path = dir.path().join("client.sock"); let server_socket = { let socket = UnixSocket::new_datagram()?; socket.bind(&server_path)?; socket.datagram()? }; tokio::spawn(async move { let mut recv_buf = vec![0u8; 1024]; loop { let (len, peer_addr) = server_socket.recv_from(&mut recv_buf[..]).await?; if let Some(path) = peer_addr.as_pathname() { server_socket.send_to(&recv_buf[..len], path).await?; } } #[allow(unreachable_code)] Ok::<(), io::Error>(()) }); { let socket = UnixSocket::new_datagram()?; socket.bind(&client_path).unwrap(); let socket = socket.datagram()?; socket.connect(server_path)?; socket.send(b"ECHO").await?; let mut recv_buf = [0u8; 16]; let len = socket.recv(&mut recv_buf[..]).await?; assert_eq!(&recv_buf[..len], b"ECHO"); } Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn listen_and_stream() -> std::io::Result<()> { let dir = tempfile::Builder::new().tempdir().unwrap(); let sock_path = dir.path().join("connect.sock"); let peer_path = dir.path().join("peer.sock"); let listener = { let sock = UnixSocket::new_stream()?; sock.bind(&sock_path)?; sock.listen(1024)? }; let accept = listener.accept(); let connect = { let sock = UnixSocket::new_stream()?; sock.bind(&peer_path)?; sock.connect(&sock_path) }; let ((mut server, _), mut client) = try_join(accept, connect).await?; assert_eq!( server.peer_addr().unwrap().as_pathname().unwrap(), &peer_path ); // Write to the client. client.write_all(b"hello").await?; drop(client); // Read from the server. let mut buf = vec![]; server.read_to_end(&mut buf).await?; assert_eq!(&buf, b"hello"); let len = server.read(&mut buf).await?; assert_eq!(len, 0); Ok(()) } #[tokio::test] #[cfg_attr(miri, ignore)] // No `socket` in miri. async fn assert_usage() -> std::io::Result<()> { let datagram_socket = UnixSocket::new_datagram()?; let result = datagram_socket .connect(std::path::PathBuf::new().join("invalid.sock")) .await; assert_eq!( result.unwrap_err().to_string(), "connect cannot be called on a datagram socket" ); let datagram_socket = UnixSocket::new_datagram()?; let result = datagram_socket.listen(1024); assert_eq!( result.unwrap_err().to_string(), "listen cannot be called on a datagram socket" ); let stream_socket = UnixSocket::new_stream()?; let result = stream_socket.datagram(); assert_eq!( result.unwrap_err().to_string(), "datagram cannot be called on a stream socket" ); Ok(()) } tokio-1.43.1/tests/uds_split.rs000064400000000000000000000024371046102023000145630ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(feature = "full")] #![cfg(unix)] #![cfg(not(miri))] // No `socket` in miri. use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::UnixStream; /// Checks that `UnixStream` can be split into a read half and a write half using /// `UnixStream::split` and `UnixStream::split_mut`. /// /// Verifies that the implementation of `AsyncWrite::poll_shutdown` shutdowns the stream for /// writing by reading to the end of stream on the other side of the connection. #[tokio::test] async fn split() -> std::io::Result<()> { let (mut a, mut b) = UnixStream::pair()?; let (mut a_read, mut a_write) = a.split(); let (mut b_read, mut b_write) = b.split(); let (a_response, b_response) = futures::future::try_join( send_recv_all(&mut a_read, &mut a_write, b"A"), send_recv_all(&mut b_read, &mut b_write, b"B"), ) .await?; assert_eq!(a_response, b"B"); assert_eq!(b_response, b"A"); Ok(()) } async fn send_recv_all( read: &mut (dyn AsyncRead + Unpin), write: &mut (dyn AsyncWrite + Unpin), input: &[u8], ) -> std::io::Result> { write.write_all(input).await?; write.shutdown().await?; let mut output = Vec::new(); read.read_to_end(&mut output).await?; Ok(output) } tokio-1.43.1/tests/uds_stream.rs000064400000000000000000000300161046102023000147150ustar 00000000000000#![cfg(feature = "full")] #![warn(rust_2018_idioms)] #![cfg(unix)] #![cfg(not(miri))] // No socket in miri. use std::io; #[cfg(target_os = "android")] use std::os::android::net::SocketAddrExt; #[cfg(target_os = "linux")] use std::os::linux::net::SocketAddrExt; use std::task::Poll; use tokio::io::{AsyncReadExt, AsyncWriteExt, Interest}; use tokio::net::{UnixListener, UnixStream}; use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task}; use futures::future::{poll_fn, try_join}; #[tokio::test] async fn accept_read_write() -> std::io::Result<()> { let dir = tempfile::Builder::new() .prefix("tokio-uds-tests") .tempdir() .unwrap(); let sock_path = dir.path().join("connect.sock"); let listener = UnixListener::bind(&sock_path)?; let accept = listener.accept(); let connect = UnixStream::connect(&sock_path); let ((mut server, _), mut client) = try_join(accept, connect).await?; // Write to the client. client.write_all(b"hello").await?; drop(client); // Read from the server. let mut buf = vec![]; server.read_to_end(&mut buf).await?; assert_eq!(&buf, b"hello"); let len = server.read(&mut buf).await?; assert_eq!(len, 0); Ok(()) } #[tokio::test] async fn shutdown() -> std::io::Result<()> { let dir = tempfile::Builder::new() .prefix("tokio-uds-tests") .tempdir() .unwrap(); let sock_path = dir.path().join("connect.sock"); let listener = UnixListener::bind(&sock_path)?; let accept = listener.accept(); let connect = UnixStream::connect(&sock_path); let ((mut server, _), mut client) = try_join(accept, connect).await?; // Shut down the client AsyncWriteExt::shutdown(&mut client).await?; // Read from the server should return 0 to indicate the channel has been closed. let mut buf = [0u8; 1]; let n = server.read(&mut buf).await?; assert_eq!(n, 0); Ok(()) } #[tokio::test] async fn try_read_write() -> std::io::Result<()> { let msg = b"hello world"; let dir = tempfile::tempdir()?; let bind_path = dir.path().join("bind.sock"); // Create listener let listener = UnixListener::bind(&bind_path)?; // Create socket pair let client = UnixStream::connect(&bind_path).await?; let (server, _) = listener.accept().await?; let mut written = msg.to_vec(); // Track the server receiving data let mut readable = task::spawn(server.readable()); assert_pending!(readable.poll()); // Write data. client.writable().await?; assert_eq!(msg.len(), client.try_write(msg)?); // The task should be notified while !readable.is_woken() { tokio::task::yield_now().await; } // Fill the write buffer using non-vectored I/O loop { // Still ready let mut writable = task::spawn(client.writable()); assert_ready_ok!(writable.poll()); match client.try_write(msg) { Ok(n) => written.extend(&msg[..n]), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { break; } Err(e) => panic!("error = {e:?}"), } } { // Write buffer full let mut writable = task::spawn(client.writable()); assert_pending!(writable.poll()); // Drain the socket from the server end using non-vectored I/O let mut read = vec![0; written.len()]; let mut i = 0; while i < read.len() { server.readable().await?; match server.try_read(&mut read[i..]) { Ok(n) => i += n, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("error = {e:?}"), } } assert_eq!(read, written); } written.clear(); client.writable().await.unwrap(); // Fill the write buffer using vectored I/O let msg_bufs: Vec<_> = msg.chunks(3).map(io::IoSlice::new).collect(); loop { // Still ready let mut writable = task::spawn(client.writable()); assert_ready_ok!(writable.poll()); match client.try_write_vectored(&msg_bufs) { Ok(n) => written.extend(&msg[..n]), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { break; } Err(e) => panic!("error = {e:?}"), } } { // Write buffer full let mut writable = task::spawn(client.writable()); assert_pending!(writable.poll()); // Drain the socket from the server end using vectored I/O let mut read = vec![0; written.len()]; let mut i = 0; while i < read.len() { server.readable().await?; let mut bufs: Vec<_> = read[i..] .chunks_mut(0x10000) .map(io::IoSliceMut::new) .collect(); match server.try_read_vectored(&mut bufs) { Ok(n) => i += n, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("error = {e:?}"), } } assert_eq!(read, written); } // Now, we listen for shutdown drop(client); loop { let ready = server.ready(Interest::READABLE).await?; if ready.is_read_closed() { break; } else { tokio::task::yield_now().await; } } Ok(()) } async fn create_pair() -> (UnixStream, UnixStream) { let dir = assert_ok!(tempfile::tempdir()); let bind_path = dir.path().join("bind.sock"); let listener = assert_ok!(UnixListener::bind(&bind_path)); let accept = listener.accept(); let connect = UnixStream::connect(&bind_path); let ((server, _), client) = assert_ok!(try_join(accept, connect).await); (client, server) } macro_rules! assert_readable_by_polling { ($stream:expr) => { assert_ok!(poll_fn(|cx| $stream.poll_read_ready(cx)).await); }; } macro_rules! assert_not_readable_by_polling { ($stream:expr) => { poll_fn(|cx| { assert_pending!($stream.poll_read_ready(cx)); Poll::Ready(()) }) .await; }; } macro_rules! assert_writable_by_polling { ($stream:expr) => { assert_ok!(poll_fn(|cx| $stream.poll_write_ready(cx)).await); }; } macro_rules! assert_not_writable_by_polling { ($stream:expr) => { poll_fn(|cx| { assert_pending!($stream.poll_write_ready(cx)); Poll::Ready(()) }) .await; }; } #[tokio::test] async fn poll_read_ready() { let (mut client, mut server) = create_pair().await; // Initial state - not readable. assert_not_readable_by_polling!(server); // There is data in the buffer - readable. assert_ok!(client.write_all(b"ping").await); assert_readable_by_polling!(server); // Readable until calls to `poll_read` return `Poll::Pending`. let mut buf = [0u8; 4]; assert_ok!(server.read_exact(&mut buf).await); assert_readable_by_polling!(server); read_until_pending(&mut server); assert_not_readable_by_polling!(server); // Detect the client disconnect. drop(client); assert_readable_by_polling!(server); } #[tokio::test] async fn poll_write_ready() { let (mut client, server) = create_pair().await; // Initial state - writable. assert_writable_by_polling!(client); // No space to write - not writable. write_until_pending(&mut client); assert_not_writable_by_polling!(client); // Detect the server disconnect. drop(server); assert_writable_by_polling!(client); } fn read_until_pending(stream: &mut UnixStream) { let mut buf = vec![0u8; 1024 * 1024]; loop { match stream.try_read(&mut buf) { Ok(_) => (), Err(err) => { assert_eq!(err.kind(), io::ErrorKind::WouldBlock); break; } } } } fn write_until_pending(stream: &mut UnixStream) { let buf = vec![0u8; 1024 * 1024]; loop { match stream.try_write(&buf) { Ok(_) => (), Err(err) => { assert_eq!(err.kind(), io::ErrorKind::WouldBlock); break; } } } } #[tokio::test] async fn try_read_buf() -> std::io::Result<()> { let msg = b"hello world"; let dir = tempfile::tempdir()?; let bind_path = dir.path().join("bind.sock"); // Create listener let listener = UnixListener::bind(&bind_path)?; // Create socket pair let client = UnixStream::connect(&bind_path).await?; let (server, _) = listener.accept().await?; let mut written = msg.to_vec(); // Track the server receiving data let mut readable = task::spawn(server.readable()); assert_pending!(readable.poll()); // Write data. client.writable().await?; assert_eq!(msg.len(), client.try_write(msg)?); // The task should be notified while !readable.is_woken() { tokio::task::yield_now().await; } // Fill the write buffer loop { // Still ready let mut writable = task::spawn(client.writable()); assert_ready_ok!(writable.poll()); match client.try_write(msg) { Ok(n) => written.extend(&msg[..n]), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { break; } Err(e) => panic!("error = {e:?}"), } } { // Write buffer full let mut writable = task::spawn(client.writable()); assert_pending!(writable.poll()); // Drain the socket from the server end let mut read = Vec::with_capacity(written.len()); let mut i = 0; while i < read.capacity() { server.readable().await?; match server.try_read_buf(&mut read) { Ok(n) => i += n, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue, Err(e) => panic!("error = {e:?}"), } } assert_eq!(read, written); } // Now, we listen for shutdown drop(client); loop { let ready = server.ready(Interest::READABLE).await?; if ready.is_read_closed() { break; } else { tokio::task::yield_now().await; } } Ok(()) } // https://github.com/tokio-rs/tokio/issues/3879 #[tokio::test] #[cfg(not(target_os = "macos"))] async fn epollhup() -> io::Result<()> { let dir = tempfile::Builder::new() .prefix("tokio-uds-tests") .tempdir() .unwrap(); let sock_path = dir.path().join("connect.sock"); let listener = UnixListener::bind(&sock_path)?; let connect = UnixStream::connect(&sock_path); tokio::pin!(connect); // Poll `connect` once. poll_fn(|cx| { use std::future::Future; assert_pending!(connect.as_mut().poll(cx)); Poll::Ready(()) }) .await; drop(listener); let err = connect.await.unwrap_err(); let errno = err.kind(); assert!( // As far as I can tell, whether we see ECONNREFUSED or ECONNRESET here // seems relatively inconsistent, at least on non-Linux operating // systems. The difference in meaning between these errnos is not // particularly well-defined, so let's just accept either. matches!( errno, io::ErrorKind::ConnectionRefused | io::ErrorKind::ConnectionReset ), "unexpected error kind: {errno:?} (expected ConnectionRefused or ConnectionReset)" ); Ok(()) } // test for https://github.com/tokio-rs/tokio/issues/6767 #[tokio::test] #[cfg(any(target_os = "linux", target_os = "android"))] async fn abstract_socket_name() { let socket_path = "\0aaa"; let listener = UnixListener::bind(socket_path).unwrap(); let accept = listener.accept(); let connect = UnixStream::connect(&socket_path); let ((stream, _), _) = try_join(accept, connect).await.unwrap(); let local_addr = stream.into_std().unwrap().local_addr().unwrap(); let abstract_path_name = local_addr.as_abstract_name().unwrap(); // `as_abstract_name` removes leading zero bytes assert_eq!(abstract_path_name, b"aaa"); } tokio-1.43.1/tests/unwindsafe.rs000064400000000000000000000021361046102023000147140ustar 00000000000000#![warn(rust_2018_idioms)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use std::panic::{RefUnwindSafe, UnwindSafe}; #[test] fn notify_is_unwind_safe() { is_unwind_safe::(); } #[test] fn join_handle_is_unwind_safe() { is_unwind_safe::>(); } #[test] fn net_types_are_unwind_safe() { is_unwind_safe::(); is_unwind_safe::(); is_unwind_safe::(); is_unwind_safe::(); } #[test] #[cfg(unix)] fn unix_net_types_are_unwind_safe() { is_unwind_safe::(); is_unwind_safe::(); is_unwind_safe::(); } #[test] #[cfg(windows)] fn windows_net_types_are_unwind_safe() { use tokio::net::windows::named_pipe::NamedPipeClient; use tokio::net::windows::named_pipe::NamedPipeServer; is_unwind_safe::(); is_unwind_safe::(); } fn is_unwind_safe() {}