diff options
| author | Alice Ryhl <aliceryhl@google.com> | 2023-11-02 12:11:04 +0000 |
|---|---|---|
| committer | Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com> | 2023-11-02 12:11:04 +0000 |
| commit | 5b1c750c7506e5054e3373bd1b0d4a65560a4783 (patch) | |
| tree | 5f95a8dac05d0cce254136de01ec466bf0f54f36 | |
| parent | c8fc590cc95e6cfd8a866f9f0208299e919a96e4 (diff) | |
| parent | 26142af839f9f62dfaafa33818e4cfdf3582fb86 (diff) | |
| download | tokio-5b1c750c7506e5054e3373bd1b0d4a65560a4783.tar.gz | |
Update Tokio to the latest LTS release v1.32.0 am: 26142af839
Original change: https://android-review.googlesource.com/c/platform/external/rust/crates/tokio/+/2815234
Change-Id: I63172610b8f3055b9ef75fd63391892a121d01d9
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
341 files changed, 22721 insertions, 5711 deletions
@@ -23,16 +23,15 @@ rust_library { host_supported: true, crate_name: "tokio", cargo_env_compat: true, - cargo_pkg_version: "1.25.0", + cargo_pkg_version: "1.32.0", srcs: ["src/lib.rs"], - edition: "2018", + edition: "2021", features: [ "bytes", "fs", "io-util", "libc", "macros", - "memchr", "mio", "net", "num_cpus", @@ -42,12 +41,10 @@ rust_library { "sync", "time", "tokio-macros", - "windows-sys", ], rustlibs: [ "libbytes", "liblibc", - "libmemchr", "libmio", "libnum_cpus", "libpin_project_lite", @@ -68,16 +65,15 @@ rust_library { host_supported: true, crate_name: "tokio", cargo_env_compat: true, - cargo_pkg_version: "1.25.0", + cargo_pkg_version: "1.32.0", srcs: ["src/lib.rs"], - edition: "2018", + edition: "2021", features: [ "bytes", "fs", "io-util", "libc", "macros", - "memchr", "mio", "net", "num_cpus", @@ -88,12 +84,10 @@ rust_library { "test-util", "time", "tokio-macros", - "windows-sys", ], rustlibs: [ "libbytes", "liblibc", - "libmemchr", "libmio", "libnum_cpus", "libpin_project_lite", @@ -115,7 +109,7 @@ rust_defaults { cargo_env_compat: true, test_suites: ["general-tests"], auto_gen_config: true, - edition: "2018", + edition: "2021", features: [ "bytes", "fs", diff --git a/CHANGELOG.md b/CHANGELOG.md index 39a57fd..eb40bc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,375 @@ +# 1.32.0 (August 16, 2023) + +### Fixed + +- sync: fix potential quadradic 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.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 @@ -334,6 +706,16 @@ wasm32-wasi target is given unstable support for the `net` feature. [#4956]: https://github.com/tokio-rs/tokio/pull/4956 [#4959]: https://github.com/tokio-rs/tokio/pull/4959 +# 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. @@ -480,6 +862,14 @@ This release fixes a bug in `Notified::enable`. ([#4747]) [#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 @@ -10,10 +10,10 @@ # See Cargo.toml.orig for the original contents. [package] -edition = "2018" -rust-version = "1.49" +edition = "2021" +rust-version = "1.63" name = "tokio" -version = "1.25.0" +version = "1.32.0" authors = ["Tokio Contributors <team@tokio.rs>"] description = """ An event-driven, non-blocking I/O platform for writing asynchronous I/O @@ -36,13 +36,13 @@ repository = "https://github.com/tokio-rs/tokio" [package.metadata.docs.rs] all-features = true -rustdoc-args = [ - "--cfg", - "docsrs", +rustc-args = [ "--cfg", "tokio_unstable", ] -rustc-args = [ +rustdoc-args = [ + "--cfg", + "docsrs", "--cfg", "tokio_unstable", ] @@ -57,13 +57,10 @@ features = [ version = "1.0.0" optional = true -[dependencies.memchr] -version = "2.2" -optional = true - [dependencies.mio] -version = "0.8.4" +version = "0.8.6" optional = true +default-features = false [dependencies.num_cpus] version = "1.8.0" @@ -74,10 +71,10 @@ version = "0.12.0" optional = true [dependencies.pin-project-lite] -version = "0.2.0" +version = "0.2.11" [dependencies.tokio-macros] -version = "1.7.0" +version = "~2.1.0" optional = true [dev-dependencies.async-stream] @@ -90,18 +87,12 @@ features = ["async-await"] [dev-dependencies.mockall] version = "0.11.1" -[dev-dependencies.tempfile] -version = "3.1.0" - [dev-dependencies.tokio-stream] version = "0.1" [dev-dependencies.tokio-test] version = "0.4.0" -[build-dependencies.autocfg] -version = "1.1" - [features] default = [] fs = [] @@ -120,10 +111,7 @@ full = [ "time", ] io-std = [] -io-util = [ - "memchr", - "bytes", -] +io-util = ["bytes"] macros = ["tokio-macros"] net = [ "libc", @@ -171,41 +159,37 @@ test-util = [ ] time = [] -[target."cfg(all(any(target_arch = \"wasm32\", target_arch = \"wasm64\"), not(target_os = \"wasi\")))".dev-dependencies.wasm-bindgen-test] +[target."cfg(all(target_family = \"wasm\", not(target_os = \"wasi\")))".dev-dependencies.wasm-bindgen-test] version = "0.3.0" -[target."cfg(docsrs)".dependencies.windows-sys] -version = "0.42.0" -features = [ - "Win32_Foundation", - "Win32_Security_Authorization", -] - [target."cfg(loom)".dev-dependencies.loom] -version = "0.5.2" +version = "0.7" features = [ "futures", "checkpoint", ] -[target."cfg(not(all(any(target_arch = \"wasm32\", target_arch = \"wasm64\"), target_os = \"unknown\")))".dev-dependencies.rand] +[target."cfg(not(all(target_family = \"wasm\", target_os = \"unknown\")))".dev-dependencies.rand] version = "0.8.0" -[target."cfg(not(any(target_arch = \"wasm32\", target_arch = \"wasm64\")))".dependencies.socket2] -version = "0.4.4" +[target."cfg(not(target_family = \"wasm\"))".dependencies.socket2] +version = "0.5.3" features = ["all"] optional = true -[target."cfg(not(any(target_arch = \"wasm32\", target_arch = \"wasm64\")))".dev-dependencies.proptest] -version = "1" +[target."cfg(not(target_family = \"wasm\"))".dev-dependencies.socket2] +version = "0.5.3" -[target."cfg(not(any(target_arch = \"wasm32\", target_arch = \"wasm64\")))".dev-dependencies.socket2] -version = "0.4" +[target."cfg(not(target_family = \"wasm\"))".dev-dependencies.tempfile] +version = "3.1.0" [target."cfg(target_os = \"freebsd\")".dev-dependencies.mio-aio] version = "0.7.0" features = ["tokio"] +[target."cfg(tokio_taskdump)".dependencies.backtrace] +version = "0.3.58" + [target."cfg(tokio_unstable)".dependencies.tracing] version = "0.1.25" features = ["std"] @@ -213,7 +197,7 @@ optional = true default-features = false [target."cfg(unix)".dependencies.libc] -version = "0.2.42" +version = "0.2.145" optional = true [target."cfg(unix)".dependencies.signal-hook-registry] @@ -221,7 +205,7 @@ version = "1.1.1" optional = true [target."cfg(unix)".dev-dependencies.libc] -version = "0.2.42" +version = "0.2.145" [target."cfg(unix)".dev-dependencies.nix] version = "0.26" @@ -232,8 +216,12 @@ features = [ default-features = false [target."cfg(windows)".dependencies.windows-sys] -version = "0.42.0" +version = "0.48" optional = true -[target."cfg(windows)".dev-dependencies.ntapi] -version = "0.3.6" +[target."cfg(windows)".dev-dependencies.windows-sys] +version = "0.48" +features = [ + "Win32_Foundation", + "Win32_Security_Authorization", +] diff --git a/Cargo.toml.orig b/Cargo.toml.orig index 0f6d30a..19384b7 100644 --- a/Cargo.toml.orig +++ b/Cargo.toml.orig @@ -6,9 +6,9 @@ name = "tokio" # - README.md # - Update CHANGELOG.md. # - Create "v1.x.y" git tag. -version = "1.25.0" -edition = "2018" -rust-version = "1.49" +version = "1.32.0" +edition = "2021" +rust-version = "1.63" authors = ["Tokio Contributors <team@tokio.rs>"] license = "MIT" readme = "README.md" @@ -42,7 +42,7 @@ full = [ ] fs = [] -io-util = ["memchr", "bytes"] +io-util = ["bytes"] # stdin, stdout, stderr io-std = [] macros = ["tokio-macros"] @@ -93,74 +93,71 @@ time = [] # a few releases. stats = [] -[build-dependencies] -autocfg = "1.1" - [dependencies] -tokio-macros = { version = "1.7.0", path = "../tokio-macros", optional = true } +tokio-macros = { version = "~2.1.0", path = "../tokio-macros", optional = true } -pin-project-lite = "0.2.0" +pin-project-lite = "0.2.11" # Everything else is optional... bytes = { version = "1.0.0", optional = true } -memchr = { version = "2.2", optional = true } -mio = { version = "0.8.4", optional = true } +mio = { version = "0.8.6", optional = true, default-features = false } num_cpus = { version = "1.8.0", optional = true } parking_lot = { version = "0.12.0", optional = true } -[target.'cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))'.dependencies] -socket2 = { version = "0.4.4", optional = true, features = [ "all" ] } +[target.'cfg(not(target_family = "wasm"))'.dependencies] +socket2 = { version = "0.5.3", 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.25", 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.42", optional = true } +libc = { version = "0.2.145", optional = true } signal-hook-registry = { version = "1.1.1", optional = true } [target.'cfg(unix)'.dev-dependencies] -libc = { version = "0.2.42" } +libc = { version = "0.2.145" } nix = { version = "0.26", default-features = false, features = ["fs", "socket"] } [target.'cfg(windows)'.dependencies.windows-sys] -version = "0.42.0" +version = "0.48" optional = true -[target.'cfg(docsrs)'.dependencies.windows-sys] -version = "0.42.0" +[target.'cfg(windows)'.dev-dependencies.windows-sys] +version = "0.48" features = [ "Win32_Foundation", "Win32_Security_Authorization", ] -[target.'cfg(windows)'.dev-dependencies.ntapi] -version = "0.3.6" - [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" -tempfile = "3.1.0" async-stream = "0.3" -[target.'cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))'.dev-dependencies] -proptest = "1" -socket2 = "0.4" +[target.'cfg(not(target_family = "wasm"))'.dev-dependencies] +socket2 = "0.5.3" +tempfile = "3.1.0" -[target.'cfg(not(all(any(target_arch = "wasm32", target_arch = "wasm64"), target_os = "unknown")))'.dev-dependencies] +[target.'cfg(not(all(target_family = "wasm", target_os = "unknown")))'.dev-dependencies] rand = "0.8.0" -[target.'cfg(all(any(target_arch = "wasm32", target_arch = "wasm64"), not(target_os = "wasi")))'.dev-dependencies] +[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.7.0", features = ["tokio"] } [target.'cfg(loom)'.dev-dependencies] -loom = { version = "0.5.2", features = ["futures", "checkpoint"] } +loom = { version = "0.7", features = ["futures", "checkpoint"] } [package.metadata.docs.rs] all-features = true @@ -56,7 +56,7 @@ Make sure you activated the full features of the tokio crate on Cargo.toml: ```toml [dependencies] -tokio = { version = "1.25.0", features = ["full"] } +tokio = { version = "1.32.0", features = ["full"] } ``` Then, on your main.rs: @@ -187,7 +187,21 @@ When updating this, also update: 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.49.0. +released at least six months ago. The current MSRV is 1.63. + +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.30 to now - 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 @@ -202,8 +216,8 @@ 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.18.x` - LTS release until June 2023 - * `1.20.x` - LTS release until September 2023. + * `1.20.x` - LTS release until September 2023. (MSRV 1.49) + * `1.25.x` - LTS release until March 2024. (MSRV 1.49) 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 @@ -216,6 +230,12 @@ can use the following dependency specification: tokio = { version = "~1.18", 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. + ## License This project is licensed under the [MIT license]. diff --git a/TEST_MAPPING b/TEST_MAPPING index 63733b6..22a0fd7 100644 --- a/TEST_MAPPING +++ b/TEST_MAPPING @@ -15,6 +15,12 @@ }, { "path": "packages/modules/DnsResolver" + }, + { + "path": "system/security/keystore2" + }, + { + "path": "system/security/keystore2/legacykeystore" } ], "presubmit": [ diff --git a/build.rs b/build.rs deleted file mode 100644 index ddade28..0000000 --- a/build.rs +++ /dev/null @@ -1,185 +0,0 @@ -use autocfg::AutoCfg; - -const CONST_THREAD_LOCAL_PROBE: &str = r#" -{ - thread_local! { - static MY_PROBE: usize = const { 10 }; - } - - MY_PROBE.with(|val| *val) -} -"#; - -const ADDR_OF_PROBE: &str = r#" -{ - let my_var = 10; - ::std::ptr::addr_of!(my_var) -} -"#; - -const CONST_MUTEX_NEW_PROBE: &str = r#" -{ - static MY_MUTEX: ::std::sync::Mutex<i32> = ::std::sync::Mutex::new(1); - *MY_MUTEX.lock().unwrap() -} -"#; - -const TARGET_HAS_ATOMIC_PROBE: &str = r#" -{ - #[cfg(target_has_atomic = "ptr")] - let _ = (); -} -"#; - -const TARGET_ATOMIC_U64_PROBE: &str = r#" -{ - #[allow(unused_imports)] - use std::sync::atomic::AtomicU64 as _; -} -"#; - -fn main() { - let mut enable_const_thread_local = false; - let mut enable_addr_of = false; - let mut enable_target_has_atomic = false; - let mut enable_const_mutex_new = false; - let mut target_needs_atomic_u64_fallback = false; - - match AutoCfg::new() { - Ok(ac) => { - // These checks prefer to call only `probe_rustc_version` if that is - // enough to determine whether the feature is supported. This is - // because the `probe_expression` call involves a call to rustc, - // which the `probe_rustc_version` call avoids. - - // Const-initialized thread locals were stabilized in 1.59. - if ac.probe_rustc_version(1, 60) { - enable_const_thread_local = true; - } else if ac.probe_rustc_version(1, 59) { - // This compiler claims to be 1.59, but there are some nightly - // compilers that claim to be 1.59 without supporting the - // feature. Explicitly probe to check if code using them - // compiles. - // - // The oldest nightly that supports the feature is 2021-12-06. - if ac.probe_expression(CONST_THREAD_LOCAL_PROBE) { - enable_const_thread_local = true; - } - } - - // The `addr_of` and `addr_of_mut` macros were stabilized in 1.51. - if ac.probe_rustc_version(1, 52) { - enable_addr_of = true; - } else if ac.probe_rustc_version(1, 51) { - // This compiler claims to be 1.51, but there are some nightly - // compilers that claim to be 1.51 without supporting the - // feature. Explicitly probe to check if code using them - // compiles. - // - // The oldest nightly that supports the feature is 2021-01-31. - if ac.probe_expression(ADDR_OF_PROBE) { - enable_addr_of = true; - } - } - - // The `target_has_atomic` cfg was stabilized in 1.60. - if ac.probe_rustc_version(1, 61) { - enable_target_has_atomic = true; - } else if ac.probe_rustc_version(1, 60) { - // This compiler claims to be 1.60, but there are some nightly - // compilers that claim to be 1.60 without supporting the - // feature. Explicitly probe to check if code using them - // compiles. - // - // The oldest nightly that supports the feature is 2022-02-11. - if ac.probe_expression(TARGET_HAS_ATOMIC_PROBE) { - enable_target_has_atomic = true; - } - } - - // If we can't tell using `target_has_atomic`, tell if the target - // has `AtomicU64` by trying to use it. - if !enable_target_has_atomic && !ac.probe_expression(TARGET_ATOMIC_U64_PROBE) { - target_needs_atomic_u64_fallback = true; - } - - // The `Mutex::new` method was made const in 1.63. - if ac.probe_rustc_version(1, 64) { - enable_const_mutex_new = true; - } else if ac.probe_rustc_version(1, 63) { - // This compiler claims to be 1.63, but there are some nightly - // compilers that claim to be 1.63 without supporting the - // feature. Explicitly probe to check if code using them - // compiles. - // - // The oldest nightly that supports the feature is 2022-06-20. - if ac.probe_expression(CONST_MUTEX_NEW_PROBE) { - enable_const_mutex_new = true; - } - } - } - - Err(e) => { - // If we couldn't detect the compiler version and features, just - // print a warning. This isn't a fatal error: we can still build - // Tokio, we just can't enable cfgs automatically. - println!( - "cargo:warning=tokio: failed to detect compiler features: {}", - e - ); - } - } - - if !enable_const_thread_local { - // To disable this feature on compilers that support it, you can - // explicitly pass this flag with the following environment variable: - // - // RUSTFLAGS="--cfg tokio_no_const_thread_local" - autocfg::emit("tokio_no_const_thread_local") - } - - if !enable_addr_of { - // To disable this feature on compilers that support it, you can - // explicitly pass this flag with the following environment variable: - // - // RUSTFLAGS="--cfg tokio_no_addr_of" - autocfg::emit("tokio_no_addr_of") - } - - if !enable_target_has_atomic { - // To disable this feature on compilers that support it, you can - // explicitly pass this flag with the following environment variable: - // - // RUSTFLAGS="--cfg tokio_no_target_has_atomic" - autocfg::emit("tokio_no_target_has_atomic") - } - - if !enable_const_mutex_new { - // To disable this feature on compilers that support it, you can - // explicitly pass this flag with the following environment variable: - // - // RUSTFLAGS="--cfg tokio_no_const_mutex_new" - autocfg::emit("tokio_no_const_mutex_new") - } - - if target_needs_atomic_u64_fallback { - // To disable this feature on compilers that support it, you can - // explicitly pass this flag with the following environment variable: - // - // RUSTFLAGS="--cfg tokio_no_atomic_u64" - autocfg::emit("tokio_no_atomic_u64") - } - - let target = ::std::env::var("TARGET").unwrap_or_default(); - - // We emit cfgs instead of using `target_family = "wasm"` that requires Rust 1.54. - // Note that these cfgs are unavailable in `Cargo.toml`. - if target.starts_with("wasm") { - autocfg::emit("tokio_wasm"); - if target.contains("wasi") { - autocfg::emit("tokio_wasi"); - } else { - autocfg::emit("tokio_wasm_not_wasi"); - } - } -} diff --git a/cargo2android_tests.bp b/cargo2android_tests.bp index 44bc4c4..a9c19af 100644 --- a/cargo2android_tests.bp +++ b/cargo2android_tests.bp @@ -4,7 +4,7 @@ rust_defaults { cargo_env_compat: true, test_suites: ["general-tests"], auto_gen_config: true, - edition: "2018", + edition: "2021", features: [ "bytes", "fs", diff --git a/patches/io_mem_stream.patch b/patches/io_mem_stream.patch deleted file mode 100644 index c21ce18..0000000 --- a/patches/io_mem_stream.patch +++ /dev/null @@ -1,12 +0,0 @@ -diff --git a/tests/io_mem_stream.rs b/tests/io_mem_stream.rs -index 01baa53..520391a 100644 ---- a/tests/io_mem_stream.rs -+++ b/tests/io_mem_stream.rs -@@ -63,6 +63,7 @@ async fn disconnect() { - } - - #[tokio::test] -+#[cfg(not(target_os = "android"))] - async fn disconnect_reader() { - let (a, mut b) = duplex(2); - diff --git a/patches/macros_join.rs.patch b/patches/macros_join.rs.patch deleted file mode 100644 index f041b14..0000000 --- a/patches/macros_join.rs.patch +++ /dev/null @@ -1,12 +0,0 @@ -diff --git a/tests/macros_join.rs b/tests/macros_join.rs -index 16e7c43..4441582 100644 ---- a/tests/macros_join.rs -+++ b/tests/macros_join.rs -@@ -66,6 +66,7 @@ async fn two_await() { - - #[test] - #[cfg(target_pointer_width = "64")] -+#[ignore = "Android: ignore until the compiler is updated. aliceryhl@ says these tests assume latest stable compiler."] - fn join_size() { - use futures::future; - use std::mem; diff --git a/patches/panic_unwind_tests.patch b/patches/panic_unwind_tests.patch index c11065b..845de51 100644 --- a/patches/panic_unwind_tests.patch +++ b/patches/panic_unwind_tests.patch @@ -1,53 +1,44 @@ -diff --git a/patches/panic_unwind_tests.patch b/patches/panic_unwind_tests.patch -index adf431e..e69de29 100644 ---- a/patches/panic_unwind_tests.patch -+++ b/patches/panic_unwind_tests.patch -@@ -1,24 +0,0 @@ --diff --git a/tests/sync_broadcast.rs b/tests/sync_broadcast.rs --index 9aa3484..53ee7d8 100644 ----- a/tests/sync_broadcast.rs --+++ b/tests/sync_broadcast.rs --@@ -292,6 +292,7 @@ fn capacity_too_big() { -- -- #[test] --+#[cfg(panic = "unwind")] -- #[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding -- fn panic_in_clone() { -- use std::panic::{self, AssertUnwindSafe}; -- --diff --git a/tests/sync_watch.rs b/tests/sync_watch.rs --index 34f9b78..e8eacce 100644 ----- a/tests/sync_watch.rs --+++ b/tests/sync_watch.rs --@@ -214,6 +214,7 @@ fn reopened_after_subscribe() { -- -- #[test] --+#[cfg(panic = "unwind")] -- #[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding -- fn send_modify_panic() { -- let (tx, mut rx) = watch::channel("one"); -- -diff --git a/tests/sync_broadcast.rs b/tests/sync_broadcast.rs -index 67c378b..cd66924 100644 ---- a/tests/sync_broadcast.rs -+++ b/tests/sync_broadcast.rs -@@ -291,6 +291,7 @@ fn capacity_too_big() { - } +diff --git a/tests/rt_common.rs b/tests/rt_common.rs +index f5dce54..abca8dd 100644 +--- a/tests/rt_common.rs ++++ b/tests/rt_common.rs +@@ -845,6 +845,7 @@ rt_test! { + } + #[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(); +diff --git a/tests/task_abort.rs b/tests/task_abort.rs +index 4798246..481cc96 100644 +--- a/tests/task_abort.rs ++++ b/tests/task_abort.rs +@@ -176,6 +176,7 @@ fn test_abort_wakes_task_3964() { + /// Checks that aborting a task whose destructor panics does not allow the + /// panic to escape the task. #[test] +#[cfg(panic = "unwind")] - #[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding - fn panic_in_clone() { - use std::panic::{self, AssertUnwindSafe}; -diff --git a/tests/sync_watch.rs b/tests/sync_watch.rs -index 34f9b78..d4f8ce8 100644 ---- a/tests/sync_watch.rs -+++ b/tests/sync_watch.rs -@@ -213,6 +213,7 @@ fn reopened_after_subscribe() { + fn test_abort_task_that_panics_on_drop_contained() { + let rt = Builder::new_current_thread().enable_time().build().unwrap(); + +@@ -199,6 +200,7 @@ fn test_abort_task_that_panics_on_drop_contained() { + + /// 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(); + +diff --git a/tests/task_blocking.rs b/tests/task_blocking.rs +index 62e20f6..4594498 100644 +--- a/tests/task_blocking.rs ++++ b/tests/task_blocking.rs +@@ -114,6 +114,7 @@ fn can_enter_current_thread_rt_from_within_block_in_place() { } #[test] +#[cfg(panic = "unwind")] - #[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding - fn send_modify_panic() { - let (tx, mut rx) = watch::channel("one"); + fn useful_panic_message_when_dropping_rt_in_rt() { + use std::panic::{catch_unwind, AssertUnwindSafe}; + diff --git a/patches/rt_common.patch b/patches/rt_common.patch deleted file mode 100644 index 1444cfe..0000000 --- a/patches/rt_common.patch +++ /dev/null @@ -1,12 +0,0 @@ -diff --git a/tests/rt_common.rs b/tests/rt_common.rs -index cb1d0f6..e5fc7a9 100644 ---- a/tests/rt_common.rs -+++ b/tests/rt_common.rs -@@ -647,6 +647,7 @@ rt_test! { - } - - #[test] -+ #[cfg(not(target_os = "android"))] - fn panic_in_task() { - let rt = rt(); - let (tx, rx) = oneshot::channel(); diff --git a/patches/task_abort.patch b/patches/task_abort.patch deleted file mode 100644 index df05ccb..0000000 --- a/patches/task_abort.patch +++ /dev/null @@ -1,20 +0,0 @@ -diff --git a/tests/task_abort.rs b/tests/task_abort.rs -index cdaa405..ec0eed7 100644 ---- a/tests/task_abort.rs -+++ b/tests/task_abort.rs -@@ -180,6 +180,7 @@ fn test_abort_wakes_task_3964() { - /// Checks that aborting a task whose destructor panics does not allow the - /// panic to escape the task. - #[test] -+#[cfg(not(target_os = "android"))] - fn test_abort_task_that_panics_on_drop_contained() { - let rt = Builder::new_current_thread().enable_time().build().unwrap(); - -@@ -204,6 +205,7 @@ fn test_abort_task_that_panics_on_drop_contained() { - - /// Checks that aborting a task whose destructor panics has the expected result. - #[test] -+#[cfg(not(target_os = "android"))] - fn test_abort_task_that_panics_on_drop_returned() { - let rt = Builder::new_current_thread().enable_time().build().unwrap(); - diff --git a/patches/task_blocking.patch b/patches/task_blocking.patch deleted file mode 100644 index 7f4f7d4..0000000 --- a/patches/task_blocking.patch +++ /dev/null @@ -1,12 +0,0 @@ -diff --git a/tests/task_blocking.rs b/tests/task_blocking.rs -index 82bef8a..d9514d2 100644 ---- a/tests/task_blocking.rs -+++ b/tests/task_blocking.rs -@@ -114,6 +114,7 @@ fn can_enter_basic_rt_from_within_block_in_place() { - } - - #[test] -+#[cfg(not(target_os = "android"))] - fn useful_panic_message_when_dropping_rt_in_rt() { - use std::panic::{catch_unwind, AssertUnwindSafe}; - diff --git a/src/doc/os.rs b/src/doc/os.rs index 0ddf869..cded8b9 100644 --- a/src/doc/os.rs +++ b/src/doc/os.rs @@ -11,9 +11,12 @@ pub mod windows { /// 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::FromRawHandle::from_raw_handle](https://doc.rust-lang.org/std/os/windows/io/trait.AsRawHandle.html#tymethod.as_raw_handle) + /// 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; } @@ -22,5 +25,44 @@ pub mod windows { /// 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<'_>; + } } } diff --git a/src/fs/file.rs b/src/fs/file.rs index d513beb..1e88f30 100644 --- a/src/fs/file.rs +++ b/src/fs/file.rs @@ -3,7 +3,7 @@ //! [`File`]: File use self::State::*; -use crate::fs::asyncify; +use crate::fs::{asyncify, OpenOptions}; use crate::io::blocking::Buf; use crate::io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}; use crate::sync::Mutex; @@ -124,8 +124,6 @@ impl File { /// /// See [`OpenOptions`] for more details. /// - /// [`OpenOptions`]: super::OpenOptions - /// /// # Errors /// /// This function will return an error if called from outside of the Tokio @@ -167,8 +165,6 @@ impl File { /// /// See [`OpenOptions`] for more details. /// - /// [`OpenOptions`]: super::OpenOptions - /// /// # Errors /// /// Results in an error if called from outside of the Tokio runtime or if @@ -199,6 +195,37 @@ impl File { Ok(File::from_std(std_file)) } + /// 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`][std] to a [`tokio::fs::File`][file]. /// /// [std]: std::fs::File @@ -399,6 +426,7 @@ impl File { /// # } /// ``` pub async fn try_clone(&self) -> io::Result<File> { + 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)) @@ -498,6 +526,7 @@ impl AsyncRead for File { cx: &mut Context<'_>, dst: &mut ReadBuf<'_>, ) -> Poll<io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); @@ -594,6 +623,7 @@ impl AsyncSeek for File { } fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> { + ready!(crate::trace::trace_leaf(cx)); let inner = self.inner.get_mut(); loop { @@ -629,6 +659,7 @@ impl AsyncWrite for File { cx: &mut Context<'_>, src: &[u8], ) -> Poll<io::Result<usize>> { + ready!(crate::trace::trace_leaf(cx)); let me = self.get_mut(); let inner = me.inner.get_mut(); @@ -695,11 +726,13 @@ impl AsyncWrite for File { } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { + 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<Result<(), io::Error>> { + ready!(crate::trace::trace_leaf(cx)); self.poll_flush(cx) } } @@ -726,23 +759,44 @@ impl std::os::unix::io::AsRawFd for File { } #[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)] -impl std::os::windows::io::AsRawHandle for File { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - self.std.as_raw_handle() +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() + } } -} -#[cfg(windows)] -impl std::os::windows::io::FromRawHandle for File { - unsafe fn from_raw_handle(handle: std::os::windows::io::RawHandle) -> Self { - StdFile::from_raw_handle(handle).into() + 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() + } } } @@ -750,8 +804,18 @@ impl Inner { async fn complete_inflight(&mut self) { use crate::future::poll_fn; - if let Err(e) = poll_fn(|cx| Pin::new(&mut *self).poll_flush(cx)).await { - self.last_write_err = Some(e.kind()); + 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, } } diff --git a/src/fs/mod.rs b/src/fs/mod.rs index 3afefc6..f6d9605 100644 --- a/src/fs/mod.rs +++ b/src/fs/mod.rs @@ -102,6 +102,9 @@ 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; @@ -112,9 +115,7 @@ feature! { pub use self::symlink::symlink; } -feature! { - #![windows] - +cfg_windows! { mod symlink_dir; pub use self::symlink_dir::symlink_dir; diff --git a/src/fs/open_options.rs b/src/fs/open_options.rs index c600ef5..1035102 100644 --- a/src/fs/open_options.rs +++ b/src/fs/open_options.rs @@ -10,6 +10,11 @@ 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 @@ -399,8 +404,6 @@ impl OpenOptions { feature! { #![unix] - use std::os::unix::fs::OpenOptionsExt; - impl OpenOptions { /// Sets the mode bits that a new file will be created with. /// @@ -464,11 +467,7 @@ feature! { } } -feature! { - #![windows] - - use std::os::windows::fs::OpenOptionsExt; - +cfg_windows! { impl OpenOptions { /// Overrides the `dwDesiredAccess` argument to the call to [`CreateFile`] /// with the specified value. diff --git a/src/fs/read_dir.rs b/src/fs/read_dir.rs index 9471e8c..def735b 100644 --- a/src/fs/read_dir.rs +++ b/src/fs/read_dir.rs @@ -35,9 +35,9 @@ pub async fn read_dir(path: impl AsRef<Path>) -> io::Result<ReadDir> { asyncify(|| -> io::Result<ReadDir> { let mut std = std::fs::read_dir(path)?; let mut buf = VecDeque::with_capacity(CHUNK_SIZE); - ReadDir::next_chunk(&mut buf, &mut std); + let remain = ReadDir::next_chunk(&mut buf, &mut std); - Ok(ReadDir(State::Idle(Some((buf, std))))) + Ok(ReadDir(State::Idle(Some((buf, std, remain))))) }) .await } @@ -66,8 +66,8 @@ pub struct ReadDir(State); #[derive(Debug)] enum State { - Idle(Option<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir)>), - Pending(JoinHandle<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir)>), + Idle(Option<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir, bool)>), + Pending(JoinHandle<(VecDeque<io::Result<DirEntry>>, std::fs::ReadDir, bool)>), } impl ReadDir { @@ -103,38 +103,35 @@ impl ReadDir { loop { match self.0 { State::Idle(ref mut data) => { - let (buf, _) = data.as_mut().unwrap(); + 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(); + let (mut buf, mut std, _) = data.take().unwrap(); self.0 = State::Pending(spawn_blocking(move || { - ReadDir::next_chunk(&mut buf, &mut std); - (buf, std) + let remain = ReadDir::next_chunk(&mut buf, &mut std); + (buf, std, remain) })); } State::Pending(ref mut rx) => { - let (mut buf, std) = ready!(Pin::new(rx).poll(cx))?; - - let ret = match buf.pop_front() { - Some(Ok(x)) => Ok(Some(x)), - Some(Err(e)) => Err(e), - None => Ok(None), - }; - - self.0 = State::Idle(Some((buf, std))); - - return Poll::Ready(ret); + self.0 = State::Idle(Some(ready!(Pin::new(rx).poll(cx))?)); } } } } - fn next_chunk(buf: &mut VecDeque<io::Result<DirEntry>>, std: &mut std::fs::ReadDir) { - for ret in std.by_ref().take(CHUNK_SIZE) { + fn next_chunk(buf: &mut VecDeque<io::Result<DirEntry>>, 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 { @@ -142,7 +139,9 @@ impl ReadDir { target_os = "solaris", target_os = "illumos", target_os = "haiku", - target_os = "vxworks" + target_os = "vxworks", + target_os = "nto", + target_os = "vita", )))] file_type: std.file_type().ok(), std: Arc::new(std), @@ -152,6 +151,8 @@ impl ReadDir { break; } } + + true } } @@ -201,7 +202,9 @@ pub struct DirEntry { target_os = "solaris", target_os = "illumos", target_os = "haiku", - target_os = "vxworks" + target_os = "vxworks", + target_os = "nto", + target_os = "vita", )))] file_type: Option<FileType>, std: Arc<std::fs::DirEntry>, @@ -332,7 +335,9 @@ impl DirEntry { target_os = "solaris", target_os = "illumos", target_os = "haiku", - target_os = "vxworks" + target_os = "vxworks", + target_os = "nto", + target_os = "vita", )))] if let Some(file_type) = self.file_type { return Ok(file_type); diff --git a/src/fs/symlink_dir.rs b/src/fs/symlink_dir.rs index 736e762..6753c25 100644 --- a/src/fs/symlink_dir.rs +++ b/src/fs/symlink_dir.rs @@ -10,7 +10,7 @@ use std::path::Path; /// /// This is an async version of [`std::os::windows::fs::symlink_dir`][std] /// -/// [std]: std::os::windows::fs::symlink_dir +/// [std]: https://doc.rust-lang.org/std/os/windows/fs/fn.symlink_dir.html pub async fn symlink_dir(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); diff --git a/src/fs/symlink_file.rs b/src/fs/symlink_file.rs index 07d8e60..623352a 100644 --- a/src/fs/symlink_file.rs +++ b/src/fs/symlink_file.rs @@ -10,7 +10,7 @@ use std::path::Path; /// /// This is an async version of [`std::os::windows::fs::symlink_file`][std] /// -/// [std]: std::os::windows::fs::symlink_file +/// [std]: https://doc.rust-lang.org/std/os/windows/fs/fn.symlink_file.html pub async fn symlink_file(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> io::Result<()> { let src = src.as_ref().to_owned(); let dst = dst.as_ref().to_owned(); diff --git a/src/fs/try_exists.rs b/src/fs/try_exists.rs new file mode 100644 index 0000000..069518b --- /dev/null +++ b/src/fs/try_exists.rs @@ -0,0 +1,34 @@ +use 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<Path>) -> io::Result<bool> { + let path = path.as_ref().to_owned(); + // std's Path::try_exists is not available for current Rust min supported version. + // Current implementation is based on its internal implementation instead. + match asyncify(move || std::fs::metadata(path)).await { + Ok(_) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(error), + } +} diff --git a/src/fuzz.rs b/src/fuzz.rs new file mode 100644 index 0000000..d89718f --- /dev/null +++ b/src/fuzz.rs @@ -0,0 +1 @@ +pub use crate::util::linked_list::tests::fuzz_linked_list; diff --git a/src/io/async_fd.rs b/src/io/async_fd.rs index 92fc6b3..ae8f5c6 100644 --- a/src/io/async_fd.rs +++ b/src/io/async_fd.rs @@ -1,4 +1,4 @@ -use crate::io::Interest; +use crate::io::{Interest, Ready}; use crate::runtime::io::{ReadyEvent, Registration}; use crate::runtime::scheduler; @@ -65,8 +65,8 @@ use std::{task::Context, task::Poll}; /// # Examples /// /// This example shows how to turn [`std::net::TcpStream`] asynchronous using -/// `AsyncFd`. It implements `read` as an async fn, and `AsyncWrite` as a trait -/// to show how to implement both approaches. +/// `AsyncFd`. It implements the read/write operations both as an `async fn` +/// and using the IO traits [`AsyncRead`] and [`AsyncWrite`]. /// /// ```no_run /// use futures::ready; @@ -74,7 +74,7 @@ use std::{task::Context, task::Poll}; /// use std::net::TcpStream; /// use std::pin::Pin; /// use std::task::{Context, Poll}; -/// use tokio::io::AsyncWrite; +/// use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; /// use tokio::io::unix::AsyncFd; /// /// pub struct AsyncTcpStream { @@ -99,6 +99,39 @@ use std::{task::Context, task::Poll}; /// } /// } /// } +/// +/// pub async fn write(&self, buf: &[u8]) -> io::Result<usize> { +/// 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<io::Result<()>> { +/// 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 { @@ -139,8 +172,12 @@ use std::{task::Context, task::Poll}; /// [`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 pub struct AsyncFd<T: AsRawFd> { 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<T>, } @@ -166,13 +203,14 @@ pub struct AsyncFdReadyMutGuard<'a, T: AsRawFd> { event: Option<ReadyEvent>, } -const ALL_INTEREST: Interest = Interest::READABLE.add(Interest::WRITABLE); - impl<T: AsRawFd> AsyncFd<T> { /// 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 @@ -185,11 +223,12 @@ impl<T: AsRawFd> AsyncFd<T> { where T: AsRawFd, { - Self::with_interest(inner, ALL_INTEREST) + Self::with_interest(inner, Interest::READABLE | Interest::WRITABLE) } - /// Creates new instance as `new` with additional ability to customize interest, - /// allowing to specify whether file descriptor will be polled for read, write or both. + /// 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 /// @@ -234,13 +273,12 @@ impl<T: AsRawFd> AsyncFd<T> { } fn take_inner(&mut self) -> Option<T> { - let fd = self.inner.as_ref().map(AsRawFd::as_raw_fd); + let inner = self.inner.take()?; + let fd = inner.as_raw_fd(); - if let Some(fd) = fd { - let _ = self.registration.deregister(&mut SourceFd(&fd)); - } + let _ = self.registration.deregister(&mut SourceFd(&fd)); - self.inner.take() + Some(inner) } /// Deregisters this file descriptor and returns ownership of the backing @@ -282,11 +320,10 @@ impl<T: AsRawFd> AsyncFd<T> { ) -> Poll<io::Result<AsyncFdReadyGuard<'a, T>>> { let event = ready!(self.registration.poll_read_ready(cx))?; - Ok(AsyncFdReadyGuard { + Poll::Ready(Ok(AsyncFdReadyGuard { async_fd: self, event: Some(event), - }) - .into() + })) } /// Polls for read readiness. @@ -320,11 +357,10 @@ impl<T: AsRawFd> AsyncFd<T> { ) -> Poll<io::Result<AsyncFdReadyMutGuard<'a, T>>> { let event = ready!(self.registration.poll_read_ready(cx))?; - Ok(AsyncFdReadyMutGuard { + Poll::Ready(Ok(AsyncFdReadyMutGuard { async_fd: self, event: Some(event), - }) - .into() + })) } /// Polls for write readiness. @@ -360,11 +396,10 @@ impl<T: AsRawFd> AsyncFd<T> { ) -> Poll<io::Result<AsyncFdReadyGuard<'a, T>>> { let event = ready!(self.registration.poll_write_ready(cx))?; - Ok(AsyncFdReadyGuard { + Poll::Ready(Ok(AsyncFdReadyGuard { async_fd: self, event: Some(event), - }) - .into() + })) } /// Polls for write readiness. @@ -398,14 +433,102 @@ impl<T: AsRawFd> AsyncFd<T> { ) -> Poll<io::Result<AsyncFdReadyMutGuard<'a, T>>> { let event = ready!(self.registration.poll_write_ready(cx))?; - Ok(AsyncFdReadyMutGuard { + Poll::Ready(Ok(AsyncFdReadyMutGuard { async_fd: self, event: Some(event), - }) - .into() + })) } - async fn readiness(&self, interest: Interest) -> io::Result<AsyncFdReadyGuard<'_, T>> { + /// 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<dyn Error>> { + /// 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<AsyncFdReadyGuard<'_, T>> { let event = self.registration.readiness(interest).await?; Ok(AsyncFdReadyGuard { @@ -414,7 +537,94 @@ impl<T: AsRawFd> AsyncFd<T> { }) } - async fn readiness_mut( + /// 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<dyn Error>> { + /// 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<AsyncFdReadyMutGuard<'_, T>> { @@ -436,7 +646,7 @@ impl<T: AsRawFd> AsyncFd<T> { /// [`AsyncFdReadyGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn readable<'a>(&'a self) -> io::Result<AsyncFdReadyGuard<'a, T>> { - self.readiness(Interest::READABLE).await + self.ready(Interest::READABLE).await } /// Waits for the file descriptor to become readable, returning a @@ -447,7 +657,7 @@ impl<T: AsRawFd> AsyncFd<T> { /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn readable_mut<'a>(&'a mut self) -> io::Result<AsyncFdReadyMutGuard<'a, T>> { - self.readiness_mut(Interest::READABLE).await + self.ready_mut(Interest::READABLE).await } /// Waits for the file descriptor to become writable, returning a @@ -460,7 +670,7 @@ impl<T: AsRawFd> AsyncFd<T> { /// [`AsyncFdReadyGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn writable<'a>(&'a self) -> io::Result<AsyncFdReadyGuard<'a, T>> { - self.readiness(Interest::WRITABLE).await + self.ready(Interest::WRITABLE).await } /// Waits for the file descriptor to become writable, returning a @@ -471,7 +681,110 @@ impl<T: AsRawFd> AsyncFd<T> { /// resource mutably when handling the [`AsyncFdReadyMutGuard`]. #[allow(clippy::needless_lifetimes)] // The lifetime improves rustdoc rendering. pub async fn writable_mut<'a>(&'a mut self) -> io::Result<AsyncFdReadyMutGuard<'a, T>> { - self.readiness_mut(Interest::WRITABLE).await + 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<T> { inner: T } + /// # impl<T> Dox<T> { + /// # async fn writable(&self) -> Result<&Self> { + /// # Ok(self) + /// # } + /// # fn try_io<R>(&self, _: impl FnMut(&T) -> Result<R>) -> Result<Result<R>> { + /// # panic!() + /// # } + /// async fn async_io<R>(&self, mut f: impl FnMut(&T) -> io::Result<R>) -> io::Result<R> { + /// 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<R>( + &self, + interest: Interest, + mut f: impl FnMut(&T) -> io::Result<R>, + ) -> io::Result<R> { + 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<R>( + &mut self, + interest: Interest, + mut f: impl FnMut(&mut T) -> io::Result<R>, + ) -> io::Result<R> { + self.registration + .async_io(interest, || f(self.inner.as_mut().unwrap())) + .await } } @@ -481,6 +794,12 @@ impl<T: AsRawFd> AsRawFd for AsyncFd<T> { } } +impl<T: AsRawFd> std::os::unix::io::AsFd for AsyncFd<T> { + fn as_fd(&self) -> std::os::unix::io::BorrowedFd<'_> { + unsafe { std::os::unix::io::BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + impl<T: std::fmt::Debug + AsRawFd> std::fmt::Debug for AsyncFd<T> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AsyncFd") @@ -496,22 +815,117 @@ impl<T: AsRawFd> Drop for AsyncFd<T> { } impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { - /// Indicates to tokio that the file descriptor is no longer ready. The - /// internal readiness flag will be cleared, and tokio will wait for the + /// 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. - /// - /// [`drop`]: method@std::mem::drop 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`. + /// + /// # 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<dyn Error>> { + /// 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. /// @@ -521,6 +935,20 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { // 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 @@ -536,6 +964,43 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { /// `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"))] @@ -545,14 +1010,11 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { ) -> Result<io::Result<R>, TryIoError> { let result = f(self.async_fd); - if let Err(e) = result.as_ref() { - if e.kind() == io::ErrorKind::WouldBlock { + match result { + Err(err) if err.kind() == io::ErrorKind::WouldBlock => { self.clear_ready(); + Err(TryIoError(())) } - } - - match result { - Err(err) if err.kind() == io::ErrorKind::WouldBlock => Err(TryIoError(())), result => Ok(result), } } @@ -569,22 +1031,117 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyGuard<'a, Inner> { } impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { - /// Indicates to tokio that the file descriptor is no longer ready. The - /// internal readiness flag will be cleared, and tokio will wait for the + /// 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. - /// - /// [`drop`]: method@std::mem::drop 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`. + /// + /// # 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<dyn Error>> { + /// 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. /// @@ -594,6 +1151,20 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { // 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 @@ -616,14 +1187,11 @@ impl<'a, Inner: AsRawFd> AsyncFdReadyMutGuard<'a, Inner> { ) -> Result<io::Result<R>, TryIoError> { let result = f(self.async_fd); - if let Err(e) = result.as_ref() { - if e.kind() == io::ErrorKind::WouldBlock { + match result { + Err(err) if err.kind() == io::ErrorKind::WouldBlock => { self.clear_ready(); + Err(TryIoError(())) } - } - - match result { - Err(err) if err.kind() == io::ErrorKind::WouldBlock => Err(TryIoError(())), result => Ok(result), } } diff --git a/src/io/interest.rs b/src/io/interest.rs index 013c114..9256bd2 100644 --- a/src/io/interest.rs +++ b/src/io/interest.rs @@ -5,13 +5,28 @@ 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(mio::Interest); +pub struct Interest(usize); impl Interest { // The non-FreeBSD definitions in this block are active only when @@ -19,30 +34,41 @@ impl Interest { cfg_aio! { /// Interest for POSIX AIO. #[cfg(target_os = "freebsd")] - pub const AIO: Interest = Interest(mio::Interest::AIO); + pub const AIO: Interest = Interest(AIO); /// Interest for POSIX AIO. #[cfg(not(target_os = "freebsd"))] - pub const AIO: Interest = Interest(mio::Interest::READABLE); + pub const AIO: Interest = Interest(READABLE); /// Interest for POSIX AIO lio_listio events. #[cfg(target_os = "freebsd")] - pub const LIO: Interest = Interest(mio::Interest::LIO); + pub const LIO: Interest = Interest(LIO); /// Interest for POSIX AIO lio_listio events. #[cfg(not(target_os = "freebsd"))] - pub const LIO: Interest = Interest(mio::Interest::READABLE); + pub const LIO: Interest = Interest(READABLE); } /// Interest in all readable events. /// /// Readable interest includes read-closed events. - pub const READABLE: Interest = Interest(mio::Interest::READABLE); + pub const READABLE: Interest = Interest(READABLE); /// Interest in all writable events. /// /// Writable interest includes write-closed events. - pub const WRITABLE: Interest = Interest(mio::Interest::WRITABLE); + 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. /// @@ -58,7 +84,7 @@ impl Interest { /// assert!(both.is_readable()); /// ``` pub const fn is_readable(self) -> bool { - self.0.is_readable() + self.0 & READABLE != 0 } /// Returns true if the value includes writable interest. @@ -75,7 +101,53 @@ impl Interest { /// assert!(both.is_writable()); /// ``` pub const fn is_writable(self) -> bool { - self.0.is_writable() + 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. @@ -92,18 +164,69 @@ impl Interest { /// assert!(BOTH.is_readable()); /// assert!(BOTH.is_writable()); pub const fn add(self, other: Interest) -> Interest { - Interest(self.0.add(other.0)) + Self(self.0 | other.0) } // This function must be crate-private to avoid exposing a `mio` dependency. - pub(crate) const fn to_mio(self) -> mio::Interest { - self.0 + pub(crate) fn to_mio(self) -> mio::Interest { + fn mio_add(wrapped: &mut Option<mio::Interest>, 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 registeration. + // + // 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, } } @@ -121,12 +244,67 @@ impl ops::BitOr for Interest { impl ops::BitOrAssign for Interest { #[inline] fn bitor_assign(&mut self, other: Self) { - self.0 = (*self | other).0; + *self = *self | other } } impl fmt::Debug for Interest { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(fmt) + 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(()) } } diff --git a/src/io/mod.rs b/src/io/mod.rs index f48035a..0fd6cc2 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -130,19 +130,23 @@ //! 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 Sink/Stream +//! ## Conversion to and from Stream/Sink //! -//! It is often convenient to encapsulate the reading and writing of -//! bytes and instead work with a [`Sink`] or [`Stream`] of some data -//! type that is encoded as bytes and/or decoded from bytes. Tokio -//! provides some utility traits in the [tokio-util] crate that -//! abstract the asynchronous buffering that is required and allows -//! you to write [`Encoder`] and [`Decoder`] functions working with a -//! buffer of bytes, and then use that ["codec"] to transform anything -//! that implements [`AsyncRead`] and [`AsyncWrite`] into a `Sink`/`Stream` of -//! your structured data. +//! It is often convenient to encapsulate the reading and writing of bytes in a +//! [`Stream`] or [`Sink`] of data. //! -//! [tokio-util]: https://docs.rs/tokio-util/0.6/tokio_util/codec/index.html +//! 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 //! @@ -167,9 +171,11 @@ //! [`AsyncWrite`]: trait@AsyncWrite //! [`AsyncReadExt`]: trait@AsyncReadExt //! [`AsyncWriteExt`]: trait@AsyncWriteExt -//! ["codec"]: https://docs.rs/tokio-util/0.6/tokio_util/codec/index.html -//! [`Encoder`]: https://docs.rs/tokio-util/0.6/tokio_util/codec/trait.Encoder.html -//! [`Decoder`]: https://docs.rs/tokio-util/0.6/tokio_util/codec/trait.Decoder.html +//! ["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 @@ -217,11 +223,11 @@ cfg_io_driver_impl! { pub use ready::Ready; } - #[cfg_attr(tokio_wasi, allow(unused_imports))] + #[cfg_attr(target_os = "wasi", allow(unused_imports))] mod poll_evented; #[cfg(not(loom))] - #[cfg_attr(tokio_wasi, allow(unused_imports))] + #[cfg_attr(target_os = "wasi", allow(unused_imports))] pub(crate) use poll_evented::PollEvented; } diff --git a/src/io/poll_evented.rs b/src/io/poll_evented.rs index dfe9ae3..cb5bffd 100644 --- a/src/io/poll_evented.rs +++ b/src/io/poll_evented.rs @@ -73,6 +73,9 @@ cfg_io_driver! { impl<E: Source> PollEvented<E> { /// 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. @@ -121,7 +124,7 @@ impl<E: Source> PollEvented<E> { } /// Returns a reference to the registration. - #[cfg(any(feature = "net"))] + #[cfg(feature = "net")] pub(crate) fn registration(&self) -> &Registration { &self.registration } @@ -162,8 +165,10 @@ feature! { match self.io.as_ref().unwrap().read(b) { Ok(n) => { // if we read a partially full buffer, this is sufficient on unix to show - // that the socket buffer has been drained - if n > 0 && (!cfg!(windows) && n < len) { + // that the socket buffer has been drained. 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 < len) { self.registration.clear_readiness(evt); } @@ -193,8 +198,10 @@ feature! { 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 - if n > 0 && (!cfg!(windows) && n < buf.len()) { + // 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); } diff --git a/src/io/read_buf.rs b/src/io/read_buf.rs index 0dc595a..283d96e 100644 --- a/src/io/read_buf.rs +++ b/src/io/read_buf.rs @@ -270,6 +270,33 @@ impl<'a> ReadBuf<'a> { } } +#[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") diff --git a/src/io/ready.rs b/src/io/ready.rs index ef135c4..919f199 100644 --- a/src/io/ready.rs +++ b/src/io/ready.rs @@ -1,5 +1,7 @@ #![cfg_attr(not(feature = "net"), allow(unreachable_pub))] +use crate::io::interest::Interest; + use std::fmt; use std::ops; @@ -7,6 +9,9 @@ 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. /// @@ -31,8 +36,22 @@ impl Ready { /// 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. - pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED); + #[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 { @@ -65,6 +84,17 @@ impl Ready { 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 } @@ -144,6 +174,38 @@ impl Ready { 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 @@ -172,35 +234,39 @@ impl Ready { pub(crate) fn as_usize(self) -> usize { self.0 } -} -cfg_io_readiness! { - use crate::io::Interest; - - impl Ready { - pub(crate) fn from_interest(interest: Interest) -> Ready { - let mut ready = Ready::EMPTY; - - if interest.is_readable() { - ready |= Ready::READABLE; - ready |= Ready::READ_CLOSED; - } + pub(crate) fn from_interest(interest: Interest) -> Ready { + let mut ready = Ready::EMPTY; - if interest.is_writable() { - ready |= Ready::WRITABLE; - ready |= Ready::WRITE_CLOSED; - } + if interest.is_readable() { + ready |= Ready::READABLE; + ready |= Ready::READ_CLOSED; + } - ready + if interest.is_writable() { + ready |= Ready::WRITABLE; + ready |= Ready::WRITE_CLOSED; } - pub(crate) fn intersection(self, interest: Interest) -> Ready { - Ready(self.0 & Ready::from_interest(interest).0) + #[cfg(any(target_os = "linux", target_os = "android"))] + if interest.is_priority() { + ready |= Ready::PRIORITY; + ready |= Ready::READ_CLOSED; } - pub(crate) fn satisfies(self, interest: Interest) -> bool { - self.0 & Ready::from_interest(interest).0 != 0 + 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 } } @@ -240,11 +306,17 @@ impl ops::Sub<Ready> for Ready { impl fmt::Debug for Ready { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Ready") - .field("is_readable", &self.is_readable()) + 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()) - .finish() + .field("is_error", &self.is_error()); + + #[cfg(any(target_os = "linux", target_os = "android"))] + fmt.field("is_priority", &self.is_priority()); + + fmt.finish() } } diff --git a/src/io/split.rs b/src/io/split.rs index f067b65..63f0960 100644 --- a/src/io/split.rs +++ b/src/io/split.rs @@ -35,9 +35,12 @@ cfg_io_util! { where T: AsyncRead + AsyncWrite, { + let is_write_vectored = stream.is_write_vectored(); + let inner = Arc::new(Inner { locked: AtomicBool::new(false), stream: UnsafeCell::new(stream), + is_write_vectored, }); let rd = ReadHalf { @@ -53,6 +56,7 @@ cfg_io_util! { struct Inner<T> { locked: AtomicBool, stream: UnsafeCell<T>, + is_write_vectored: bool, } struct Guard<'a, T> { @@ -131,6 +135,19 @@ impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_shutdown(cx) } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<Result<usize, io::Error>> { + let mut inner = ready!(self.inner.poll_lock(cx)); + inner.stream_pin().poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored + } } impl<T> Inner<T> { diff --git a/src/io/stderr.rs b/src/io/stderr.rs index 2f624fb..e55cb16 100644 --- a/src/io/stderr.rs +++ b/src/io/stderr.rs @@ -74,16 +74,37 @@ cfg_io_std! { } #[cfg(unix)] -impl std::os::unix::io::AsRawFd for Stderr { - fn as_raw_fd(&self) -> std::os::unix::io::RawFd { - std::io::stderr().as_raw_fd() +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)] -impl std::os::windows::io::AsRawHandle for Stderr { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - std::io::stderr().as_raw_handle() +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()) } + } } } diff --git a/src/io/stdin.rs b/src/io/stdin.rs index c9578f1..640cb0b 100644 --- a/src/io/stdin.rs +++ b/src/io/stdin.rs @@ -49,16 +49,37 @@ cfg_io_std! { } #[cfg(unix)] -impl std::os::unix::io::AsRawFd for Stdin { - fn as_raw_fd(&self) -> std::os::unix::io::RawFd { - std::io::stdin().as_raw_fd() +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)] -impl std::os::windows::io::AsRawHandle for Stdin { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - std::io::stdin().as_raw_handle() +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()) } + } } } diff --git a/src/io/stdio_common.rs b/src/io/stdio_common.rs index b1cc61d..06da761 100644 --- a/src/io/stdio_common.rs +++ b/src/io/stdio_common.rs @@ -176,6 +176,7 @@ mod tests { } #[test] + #[cfg_attr(miri, ignore)] fn test_splitter() { let data = str::repeat("â–ˆ", MAX_BUF); let mut wr = super::SplitByUtf8BoundaryIfWindows::new(TextMockWriter); @@ -189,6 +190,7 @@ mod tests { } #[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 diff --git a/src/io/stdout.rs b/src/io/stdout.rs index a08ed01..7636fe9 100644 --- a/src/io/stdout.rs +++ b/src/io/stdout.rs @@ -73,16 +73,37 @@ cfg_io_std! { } #[cfg(unix)] -impl std::os::unix::io::AsRawFd for Stdout { - fn as_raw_fd(&self) -> std::os::unix::io::RawFd { - std::io::stdout().as_raw_fd() +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)] -impl std::os::windows::io::AsRawHandle for Stdout { - fn as_raw_handle(&self) -> std::os::windows::io::RawHandle { - std::io::stdout().as_raw_handle() +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()) } + } } } diff --git a/src/io/util/async_buf_read_ext.rs b/src/io/util/async_buf_read_ext.rs index b241e35..50ea2d9 100644 --- a/src/io/util/async_buf_read_ext.rs +++ b/src/io/util/async_buf_read_ext.rs @@ -146,7 +146,7 @@ cfg_io_util! { /// [`next_line`] method. /// * Use [`tokio_util::codec::LinesCodec`][LinesCodec]. /// - /// [LinesCodec]: https://docs.rs/tokio-util/0.6/tokio_util/codec/struct.LinesCodec.html + /// [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 diff --git a/src/io/util/async_read_ext.rs b/src/io/util/async_read_ext.rs index df5445c..7588f82 100644 --- a/src/io/util/async_read_ext.rs +++ b/src/io/util/async_read_ext.rs @@ -27,7 +27,7 @@ cfg_io_util! { ) => { $( $(#[$outer])* - fn $name<'a>(&'a mut self) -> $($fut)*<&'a mut Self> where Self: Unpin { + fn $name(&mut self) -> $($fut)*<&mut Self> where Self: Unpin { $($fut)*::new(self) } )* @@ -230,10 +230,13 @@ cfg_io_util! { /// let mut buffer = BytesMut::with_capacity(10); /// /// assert!(buffer.is_empty()); + /// assert!(buffer.capacity() >= 10); /// - /// // read up to 10 bytes, note that the return value is not needed - /// // to access the data that was read as `buffer`'s internal - /// // cursor is updated. + /// // 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[..]); @@ -292,7 +295,8 @@ cfg_io_util! { /// #[tokio::main] /// async fn main() -> io::Result<()> { /// let mut f = File::open("foo.txt").await?; - /// let mut buffer = [0; 10]; + /// let len = 10; + /// let mut buffer = vec![0; len]; /// /// // read exactly 10 bytes /// f.read_exact(&mut buffer).await?; diff --git a/src/io/util/async_write_ext.rs b/src/io/util/async_write_ext.rs index dfdde82..a9812d5 100644 --- a/src/io/util/async_write_ext.rs +++ b/src/io/util/async_write_ext.rs @@ -30,7 +30,7 @@ cfg_io_util! { ) => { $( $(#[$outer])* - fn $name<'a>(&'a mut self, n: $ty) -> $($fut)*<&'a mut Self> where Self: Unpin { + fn $name(&mut self, n: $ty) -> $($fut)*<&mut Self> where Self: Unpin { $($fut)*::new(self, n) } )* diff --git a/src/io/util/copy.rs b/src/io/util/copy.rs index 47dad89..5586124 100644 --- a/src/io/util/copy.rs +++ b/src/io/util/copy.rs @@ -153,14 +153,22 @@ cfg_io_util! { /// /// 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. + /// `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 /// diff --git a/src/io/util/copy_bidirectional.rs b/src/io/util/copy_bidirectional.rs index c93060b..e1a7db1 100644 --- a/src/io/util/copy_bidirectional.rs +++ b/src/io/util/copy_bidirectional.rs @@ -1,8 +1,8 @@ use super::copy::CopyBuffer; +use crate::future::poll_fn; use crate::io::{AsyncRead, AsyncWrite}; -use std::future::Future; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; @@ -13,13 +13,6 @@ enum TransferState { Done(u64), } -struct CopyBidirectional<'a, A: ?Sized, B: ?Sized> { - a: &'a mut A, - b: &'a mut B, - a_to_b: TransferState, - b_to_a: TransferState, -} - fn transfer_one_direction<A, B>( cx: &mut Context<'_>, state: &mut TransferState, @@ -48,35 +41,6 @@ where } } } - -impl<'a, A, B> Future for CopyBidirectional<'a, A, B> -where - A: AsyncRead + AsyncWrite + Unpin + ?Sized, - B: AsyncRead + AsyncWrite + Unpin + ?Sized, -{ - type Output = io::Result<(u64, u64)>; - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - // Unpack self into mut refs to each field to avoid borrow check issues. - let CopyBidirectional { - a, - b, - a_to_b, - b_to_a, - } = &mut *self; - - let a_to_b = transfer_one_direction(cx, a_to_b, &mut *a, &mut *b)?; - let b_to_a = transfer_one_direction(cx, b_to_a, &mut *b, &mut *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))) - } -} - /// Copies data in both directions between `a` and `b`. /// /// This function returns a future that will read from both streams, @@ -110,11 +74,18 @@ where A: AsyncRead + AsyncWrite + Unpin + ?Sized, B: AsyncRead + AsyncWrite + Unpin + ?Sized, { - CopyBidirectional { - a, - b, - a_to_b: TransferState::Running(CopyBuffer::new()), - b_to_a: TransferState::Running(CopyBuffer::new()), - } + let mut a_to_b = TransferState::Running(CopyBuffer::new()); + let mut b_to_a = TransferState::Running(CopyBuffer::new()); + 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 } diff --git a/src/io/util/copy_buf.rs b/src/io/util/copy_buf.rs index 6831580..c23fc9a 100644 --- a/src/io/util/copy_buf.rs +++ b/src/io/util/copy_buf.rs @@ -24,11 +24,17 @@ cfg_io_util! { /// /// 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. + /// `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 /// diff --git a/src/io/util/empty.rs b/src/io/util/empty.rs index 9e648f8..b96fabb 100644 --- a/src/io/util/empty.rs +++ b/src/io/util/empty.rs @@ -53,6 +53,7 @@ impl AsyncRead for Empty { cx: &mut Context<'_>, _: &mut ReadBuf<'_>, ) -> Poll<io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(())) } @@ -61,6 +62,7 @@ impl AsyncRead for Empty { impl AsyncBufRead for Empty { #[inline] fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { + ready!(crate::trace::trace_leaf(cx)); ready!(poll_proceed_and_make_progress(cx)); Poll::Ready(Ok(&[])) } diff --git a/src/io/util/mem.rs b/src/io/util/mem.rs index 31884b3..5b404c2 100644 --- a/src/io/util/mem.rs +++ b/src/io/util/mem.rs @@ -233,6 +233,7 @@ impl AsyncRead for Pipe { cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll<std::io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = self.poll_read_internal(cx, buf); @@ -249,6 +250,7 @@ impl AsyncRead for Pipe { cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll<std::io::Result<()>> { + ready!(crate::trace::trace_leaf(cx)); self.poll_read_internal(cx, buf) } } @@ -261,6 +263,7 @@ impl AsyncWrite for Pipe { cx: &mut task::Context<'_>, buf: &[u8], ) -> Poll<std::io::Result<usize>> { + ready!(crate::trace::trace_leaf(cx)); let coop = ready!(crate::runtime::coop::poll_proceed(cx)); let ret = self.poll_write_internal(cx, buf); @@ -277,6 +280,7 @@ impl AsyncWrite for Pipe { cx: &mut task::Context<'_>, buf: &[u8], ) -> Poll<std::io::Result<usize>> { + ready!(crate::trace::trace_leaf(cx)); self.poll_write_internal(cx, buf) } } diff --git a/src/io/util/read_to_end.rs b/src/io/util/read_to_end.rs index f4a564d..8edba2a 100644 --- a/src/io/util/read_to_end.rs +++ b/src/io/util/read_to_end.rs @@ -1,11 +1,11 @@ use crate::io::util::vec_with_initialized::{into_read_buf_parts, VecU8, VecWithInitialized}; -use crate::io::AsyncRead; +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; +use std::mem::{self, MaybeUninit}; use std::pin::Pin; use std::task::{Context, Poll}; @@ -67,16 +67,47 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>( // 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. - buf.reserve(32); + // 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 = buf.get_read_buf(); + 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<u8>; 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); - let filled_before = read_buf.filled().len(); - let poll_result = read.poll_read(cx, &mut read_buf); - let filled_after = read_buf.filled().len(); - let n = filled_after - filled_before; + // 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); @@ -87,11 +118,11 @@ fn poll_read_to_end<V: VecU8, R: AsyncRead + ?Sized>( // 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!(filled_before, filled_after); + debug_assert_eq!(n, 0); Poll::Pending } Poll::Ready(Err(err)) => { - debug_assert_eq!(filled_before, filled_after); + debug_assert_eq!(n, 0); Poll::Ready(Err(err)) } Poll::Ready(Ok(())) => Poll::Ready(Ok(n)), diff --git a/src/io/util/read_until.rs b/src/io/util/read_until.rs index 90a0e8a..fb6fb22 100644 --- a/src/io/util/read_until.rs +++ b/src/io/util/read_until.rs @@ -1,4 +1,5 @@ use crate::io::AsyncBufRead; +use crate::util::memchr; use pin_project_lite::pin_project; use std::future::Future; diff --git a/src/io/util/vec_with_initialized.rs b/src/io/util/vec_with_initialized.rs index a9b94e3..f527492 100644 --- a/src/io/util/vec_with_initialized.rs +++ b/src/io/util/vec_with_initialized.rs @@ -28,6 +28,7 @@ pub(crate) struct VecWithInitialized<V> { // The number of initialized bytes in the vector. // Always between `vec.len()` and `vec.capacity()`. num_initialized: usize, + starting_capacity: usize, } impl VecWithInitialized<Vec<u8>> { @@ -47,6 +48,7 @@ where // to its length are initialized. Self { num_initialized: vec.as_mut().len(), + starting_capacity: vec.as_ref().capacity(), vec, } } @@ -111,6 +113,15 @@ where 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 { @@ -384,7 +384,33 @@ //! [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 //! -//! ## WASM support +//! ## 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: @@ -430,26 +456,9 @@ compile_error! { "Tokio requires the platform pointer width to be 32, 64, or 128 bits" } -// Ensure that our build script has correctly set cfg flags for wasm. -// -// Each condition is written all(a, not(b)). This should be read as -// "if a, then we must also have b". -#[cfg(any( - all(target_arch = "wasm32", not(tokio_wasm)), - all(target_arch = "wasm64", not(tokio_wasm)), - all(target_family = "wasm", not(tokio_wasm)), - all(target_os = "wasi", not(tokio_wasm)), - all(target_os = "wasi", not(tokio_wasi)), - all(target_os = "wasi", tokio_wasm_not_wasi), - all(tokio_wasm, not(any(target_arch = "wasm32", target_arch = "wasm64"))), - all(tokio_wasm_not_wasi, not(tokio_wasm)), - all(tokio_wasi, not(tokio_wasm)) -))] -compile_error!("Tokio's build script has incorrectly detected wasm."); - #[cfg(all( not(tokio_unstable), - tokio_wasm, + target_family = "wasm", any( feature = "fs", feature = "io-std", @@ -461,6 +470,21 @@ compile_error!("Tokio's build script has incorrectly detected wasm."); ))] 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(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 @@ -526,6 +550,40 @@ 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<Output = ()> { + 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 @@ -631,3 +689,7 @@ cfg_macros! { #[cfg(feature = "io-util")] #[cfg(test)] fn is_unpin<T: Unpin>() {} + +/// fuzz test (fuzz_linked_list) +#[cfg(fuzzing)] +pub mod fuzz; diff --git a/src/loom/mocked.rs b/src/loom/mocked.rs index 56dc1a0..d40e2c1 100644 --- a/src/loom/mocked.rs +++ b/src/loom/mocked.rs @@ -15,6 +15,7 @@ pub(crate) mod sync { } #[inline] + #[track_caller] pub(crate) fn lock(&self) -> MutexGuard<'_, T> { self.0.lock().unwrap() } diff --git a/src/loom/std/barrier.rs b/src/loom/std/barrier.rs new file mode 100644 index 0000000..a3f0ca0 --- /dev/null +++ b/src/loom/std/barrier.rs @@ -0,0 +1,217 @@ +//! 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<BarrierState>, + 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<BarrierWaitResult> { + // 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 + } +} diff --git a/src/loom/std/mod.rs b/src/loom/std/mod.rs index 6bd1ad9..0c611af 100644 --- a/src/loom/std/mod.rs +++ b/src/loom/std/mod.rs @@ -4,8 +4,9 @@ mod atomic_u16; mod atomic_u32; mod atomic_u64; mod atomic_usize; +mod barrier; mod mutex; -#[cfg(feature = "parking_lot")] +#[cfg(all(feature = "parking_lot", not(miri)))] mod parking_lot; mod unsafe_cell; @@ -55,17 +56,17 @@ pub(crate) mod sync { // internal use. Note however that some are not _currently_ named by // consuming code. - #[cfg(feature = "parking_lot")] + #[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(feature = "parking_lot"))] + #[cfg(not(all(feature = "parking_lot", not(miri))))] #[allow(unused_imports)] pub(crate) use std::sync::{Condvar, MutexGuard, RwLock, RwLockReadGuard, WaitTimeoutResult}; - #[cfg(not(feature = "parking_lot"))] + #[cfg(not(all(feature = "parking_lot", not(miri))))] pub(crate) use crate::loom::std::mutex::Mutex; pub(crate) mod atomic { @@ -76,6 +77,8 @@ pub(crate) mod sync { pub(crate) use std::sync::atomic::{fence, AtomicBool, AtomicPtr, AtomicU8, Ordering}; } + + pub(crate) use super::barrier::Barrier; } pub(crate) mod sys { diff --git a/src/loom/std/mutex.rs b/src/loom/std/mutex.rs index 076f786..7b8f9ba 100644 --- a/src/loom/std/mutex.rs +++ b/src/loom/std/mutex.rs @@ -13,7 +13,6 @@ impl<T> Mutex<T> { } #[inline] - #[cfg(not(tokio_no_const_mutex_new))] pub(crate) const fn const_new(t: T) -> Mutex<T> { Mutex(sync::Mutex::new(t)) } diff --git a/src/loom/std/parking_lot.rs b/src/loom/std/parking_lot.rs index e3af258..9b9a81d 100644 --- a/src/loom/std/parking_lot.rs +++ b/src/loom/std/parking_lot.rs @@ -52,8 +52,7 @@ impl<T> Mutex<T> { } #[inline] - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - #[cfg_attr(docsrs, doc(cfg(all(feature = "parking_lot",))))] + #[cfg(not(all(loom, test)))] pub(crate) const fn const_new(t: T) -> Mutex<T> { Mutex(PhantomData, parking_lot::const_mutex(t)) } diff --git a/src/loom/std/unsafe_cell.rs b/src/loom/std/unsafe_cell.rs index 66c1d79..3d6513b 100644 --- a/src/loom/std/unsafe_cell.rs +++ b/src/loom/std/unsafe_cell.rs @@ -6,10 +6,12 @@ impl<T> UnsafeCell<T> { UnsafeCell(std::cell::UnsafeCell::new(data)) } + #[inline(always)] pub(crate) fn with<R>(&self, f: impl FnOnce(*const T) -> R) -> R { f(self.0.get()) } + #[inline(always)] pub(crate) fn with_mut<R>(&self, f: impl FnOnce(*mut T) -> R) -> R { f(self.0.get()) } diff --git a/src/macros/addr_of.rs b/src/macros/addr_of.rs index 51d4889..9d28158 100644 --- a/src/macros/addr_of.rs +++ b/src/macros/addr_of.rs @@ -1,7 +1,6 @@ //! 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. -#[cfg(not(tokio_no_addr_of))] macro_rules! generate_addr_of_methods { ( impl<$($gen:ident)*> $struct_name:ty {$( @@ -21,33 +20,3 @@ macro_rules! generate_addr_of_methods { )*} }; } - -// The `addr_of_mut!` macro is only available for MSRV at least 1.51.0. This -// version of the macro uses a workaround for older versions of rustc. -#[cfg(tokio_no_addr_of)] -macro_rules! generate_addr_of_methods { - ( - impl<$($gen:ident)*> $struct_name:ty {$( - $(#[$attrs:meta])* - $vis:vis unsafe fn $fn_name:ident(self: NonNull<Self>) -> NonNull<$field_type:ty> { - &self$(.$field_name:tt)+ - } - )*} - ) => { - impl<$($gen)*> $struct_name {$( - $(#[$attrs])* - $vis unsafe fn $fn_name(me: ::core::ptr::NonNull<Self>) -> ::core::ptr::NonNull<$field_type> { - let me = me.as_ptr(); - let me_u8 = me as *mut u8; - - let field_offset = { - let me_ref = &*me; - let field_ref_u8 = (&me_ref $(.$field_name)+ ) as *const $field_type as *const u8; - field_ref_u8.offset_from(me_u8) - }; - - ::core::ptr::NonNull::new_unchecked(me_u8.offset(field_offset).cast()) - } - )*} - }; -} diff --git a/src/macros/cfg.rs b/src/macros/cfg.rs index 1c66d24..443cc14 100644 --- a/src/macros/cfg.rs +++ b/src/macros/cfg.rs @@ -13,6 +13,30 @@ macro_rules! feature { } } +/// 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 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)*) => { @@ -61,7 +85,7 @@ macro_rules! cfg_fs { ($($item:item)*) => { $( #[cfg(feature = "fs")] - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] #[cfg_attr(docsrs, doc(cfg(feature = "fs")))] $item )* @@ -252,7 +276,7 @@ macro_rules! cfg_process { #[cfg(feature = "process")] #[cfg_attr(docsrs, doc(cfg(feature = "process")))] #[cfg(not(loom))] - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] $item )* } @@ -281,7 +305,7 @@ macro_rules! cfg_signal { #[cfg(feature = "signal")] #[cfg_attr(docsrs, doc(cfg(feature = "signal")))] #[cfg(not(loom))] - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] $item )* } @@ -348,7 +372,7 @@ macro_rules! cfg_not_rt { macro_rules! cfg_rt_multi_thread { ($($item:item)*) => { $( - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] #[cfg_attr(docsrs, doc(cfg(feature = "rt-multi-thread")))] $item )* @@ -361,6 +385,44 @@ macro_rules! cfg_not_rt_multi_thread { } } +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)*) => { $( @@ -461,14 +523,7 @@ macro_rules! cfg_not_coop { macro_rules! cfg_has_atomic_u64 { ($($item:item)*) => { $( - #[cfg_attr( - not(tokio_no_target_has_atomic), - cfg(all(target_has_atomic = "64", not(tokio_no_atomic_u64)) - ))] - #[cfg_attr( - tokio_no_target_has_atomic, - cfg(not(tokio_no_atomic_u64)) - )] + #[cfg(target_has_atomic = "64")] $item )* } @@ -477,14 +532,7 @@ macro_rules! cfg_has_atomic_u64 { macro_rules! cfg_not_has_atomic_u64 { ($($item:item)*) => { $( - #[cfg_attr( - not(tokio_no_target_has_atomic), - cfg(any(not(target_has_atomic = "64"), tokio_no_atomic_u64) - ))] - #[cfg_attr( - tokio_no_target_has_atomic, - cfg(tokio_no_atomic_u64) - )] + #[cfg(not(target_has_atomic = "64"))] $item )* } @@ -493,13 +541,7 @@ macro_rules! cfg_not_has_atomic_u64 { macro_rules! cfg_has_const_mutex_new { ($($item:item)*) => { $( - #[cfg(all( - not(all(loom, test)), - any( - feature = "parking_lot", - not(tokio_no_const_mutex_new) - ) - ))] + #[cfg(not(all(loom, test)))] $item )* } @@ -508,13 +550,7 @@ macro_rules! cfg_has_const_mutex_new { macro_rules! cfg_not_has_const_mutex_new { ($($item:item)*) => { $( - #[cfg(not(all( - not(all(loom, test)), - any( - feature = "parking_lot", - not(tokio_no_const_mutex_new) - ) - )))] + #[cfg(all(loom, test))] $item )* } @@ -523,7 +559,7 @@ macro_rules! cfg_not_has_const_mutex_new { macro_rules! cfg_not_wasi { ($($item:item)*) => { $( - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] $item )* } @@ -532,7 +568,7 @@ macro_rules! cfg_not_wasi { macro_rules! cfg_is_wasm_not_wasi { ($($item:item)*) => { $( - #[cfg(tokio_wasm_not_wasi)] + #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] $item )* } diff --git a/src/macros/join.rs b/src/macros/join.rs index 7e85203..8a01986 100644 --- a/src/macros/join.rs +++ b/src/macros/join.rs @@ -158,7 +158,9 @@ macro_rules! join { // ===== Entry point ===== - ( $($e:expr),* $(,)?) => { + ( $($e:expr),+ $(,)?) => { $crate::join!(@{ () (0) } $($e,)*) }; + + () => { async {}.await } } diff --git a/src/macros/loom.rs b/src/macros/loom.rs index d57d9fb..fa2653c 100644 --- a/src/macros/loom.rs +++ b/src/macros/loom.rs @@ -1,11 +1,7 @@ macro_rules! if_loom { ($($t:tt)*) => {{ #[cfg(loom)] - const LOOM: bool = true; - #[cfg(not(loom))] - const LOOM: bool = false; - - if LOOM { + { $($t)* } }} diff --git a/src/macros/mod.rs b/src/macros/mod.rs index 57678c6..82f42db 100644 --- a/src/macros/mod.rs +++ b/src/macros/mod.rs @@ -23,10 +23,6 @@ cfg_trace! { mod trace; } -#[macro_use] -#[cfg(feature = "rt")] -pub(crate) mod scoped_tls; - cfg_macros! { #[macro_use] mod select; diff --git a/src/macros/scoped_tls.rs b/src/macros/scoped_tls.rs deleted file mode 100644 index ed74b32..0000000 --- a/src/macros/scoped_tls.rs +++ /dev/null @@ -1,77 +0,0 @@ -use crate::loom::thread::LocalKey; - -use std::cell::Cell; -use std::marker; - -/// Sets a reference as a thread-local. -macro_rules! scoped_thread_local { - ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => ( - $(#[$attrs])* - $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty> - = $crate::macros::scoped_tls::ScopedKey { - inner: { - tokio_thread_local!(static FOO: ::std::cell::Cell<*const ()> = const { - std::cell::Cell::new(::std::ptr::null()) - }); - &FOO - }, - _marker: ::std::marker::PhantomData, - }; - ) -} - -/// Type representing a thread local storage key corresponding to a reference -/// to the type parameter `T`. -pub(crate) struct ScopedKey<T> { - pub(crate) inner: &'static LocalKey<Cell<*const ()>>, - pub(crate) _marker: marker::PhantomData<T>, -} - -unsafe impl<T> Sync for ScopedKey<T> {} - -impl<T> ScopedKey<T> { - /// Inserts a value into this scoped thread local storage slot for a - /// duration of a closure. - pub(crate) fn set<F, R>(&'static self, t: &T, f: F) -> R - where - F: FnOnce() -> R, - { - struct Reset { - key: &'static LocalKey<Cell<*const ()>>, - val: *const (), - } - - impl Drop for Reset { - fn drop(&mut self) { - self.key.with(|c| c.set(self.val)); - } - } - - let prev = self.inner.with(|c| { - let prev = c.get(); - c.set(t as *const _ as *const ()); - prev - }); - - let _reset = Reset { - key: self.inner, - val: prev, - }; - - f() - } - - /// Gets a value out of this scoped variable. - pub(crate) fn with<F, R>(&'static self, f: F) -> R - where - F: FnOnce(Option<&T>) -> R, - { - let val = self.inner.with(|c| c.get()); - - if val.is_null() { - f(None) - } else { - unsafe { f(Some(&*(val as *const T))) } - } - } -} diff --git a/src/macros/select.rs b/src/macros/select.rs index 7ba04a0..31c9b3a 100644 --- a/src/macros/select.rs +++ b/src/macros/select.rs @@ -131,6 +131,13 @@ /// 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 diff --git a/src/macros/thread_local.rs b/src/macros/thread_local.rs index 74be99a..6d4a6aa 100644 --- a/src/macros/thread_local.rs +++ b/src/macros/thread_local.rs @@ -10,23 +10,9 @@ macro_rules! tokio_thread_local { ($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } } } -#[cfg(not(tokio_no_const_thread_local))] #[cfg(not(all(loom, test)))] macro_rules! tokio_thread_local { ($($tts:tt)+) => { ::std::thread_local!{ $($tts)+ } } } - -#[cfg(tokio_no_const_thread_local)] -#[cfg(not(all(loom, test)))] -macro_rules! tokio_thread_local { - ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty = const { $expr:expr } $(;)?) => { - ::std::thread_local! { - $(#[$attrs])* - $vis static $name: $ty = $expr; - } - }; - - ($($tts:tt)+) => { ::std::thread_local!{ $($tts)+ } } -} diff --git a/src/macros/try_join.rs b/src/macros/try_join.rs index 597cd5d..7b12370 100644 --- a/src/macros/try_join.rs +++ b/src/macros/try_join.rs @@ -210,7 +210,9 @@ macro_rules! try_join { // ===== Entry point ===== - ( $($e:expr),* $(,)?) => { + ( $($e:expr),+ $(,)?) => { $crate::try_join!(@{ () (0) } $($e,)*) }; + + () => { async { Ok(()) }.await } } diff --git a/src/net/tcp/listener.rs b/src/net/tcp/listener.rs index 4441313..28da34a 100644 --- a/src/net/tcp/listener.rs +++ b/src/net/tcp/listener.rs @@ -5,7 +5,6 @@ cfg_not_wasi! { use crate::net::{to_socket_addrs, ToSocketAddrs}; } -use std::convert::TryFrom; use std::fmt; use std::io; use std::net::{self, SocketAddr}; @@ -282,7 +281,7 @@ impl TcpListener { .map(|raw_socket| unsafe { std::net::TcpListener::from_raw_socket(raw_socket) }) } - #[cfg(tokio_wasi)] + #[cfg(target_os = "wasi")] { use std::os::wasi::io::{FromRawFd, IntoRawFd}; self.io @@ -407,10 +406,16 @@ mod sys { 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(tokio_wasi)] + #[cfg(target_os = "wasi")] mod sys { use super::TcpListener; use std::os::wasi::prelude::*; @@ -420,17 +425,27 @@ cfg_unstable! { self.io.as_raw_fd() } } + + impl AsFd for TcpListener { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } } -#[cfg(windows)] -mod sys { - use super::TcpListener; - use std::os::windows::prelude::*; +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()) } + } + } } diff --git a/src/net/tcp/socket.rs b/src/net/tcp/socket.rs index 09349fe..1497bd6 100644 --- a/src/net/tcp/socket.rs +++ b/src/net/tcp/socket.rs @@ -5,11 +5,13 @@ use std::io; use std::net::SocketAddr; #[cfg(unix)] -use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; -#[cfg(windows)] -use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}; +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`. @@ -398,6 +400,51 @@ impl TcpSocket { 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<dyn std::error::Error>> { + /// let socket = TcpSocket::new_v4()?; + /// + /// println!("{:?}", socket.nodelay()?); + /// # 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<dyn std::error::Error>> { + /// let stream = TcpSocket::new_v4()?; + /// + /// stream.set_nodelay(true)?; + /// # Ok(()) + /// # } + /// ``` + pub fn nodelay(&self) -> io::Result<bool> { + self.inner.nodelay() + } + /// Gets the value of the `IP_TOS` option for this socket. /// /// For more information about this option, see [`set_tos`]. @@ -406,7 +453,7 @@ impl TcpSocket { /// 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.4.2/src/socket2/socket.rs.html#1178 + // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1464 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", @@ -433,7 +480,7 @@ impl TcpSocket { /// /// **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.4.2/src/socket2/socket.rs.html#1178 + // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1446 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", @@ -472,7 +519,7 @@ impl TcpSocket { /// works for some socket types, particularly `AF_INET` sockets. /// /// If `interface` is `None` or an empty string it removes the binding. - #[cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] + #[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")))) @@ -738,6 +785,13 @@ impl AsRawFd for TcpSocket { } #[cfg(unix)] +impl AsFd for TcpSocket { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} + +#[cfg(unix)] impl FromRawFd for TcpSocket { /// Converts a `RawFd` to a `TcpSocket`. /// @@ -758,30 +812,35 @@ impl IntoRawFd for TcpSocket { } } -#[cfg(windows)] -impl IntoRawSocket for TcpSocket { - fn into_raw_socket(self) -> RawSocket { - self.inner.into_raw_socket() +cfg_windows! { + impl IntoRawSocket for TcpSocket { + fn into_raw_socket(self) -> RawSocket { + self.inner.into_raw_socket() + } } -} -#[cfg(windows)] -impl AsRawSocket for TcpSocket { - fn as_raw_socket(&self) -> RawSocket { - self.inner.as_raw_socket() + impl AsRawSocket for TcpSocket { + fn as_raw_socket(&self) -> RawSocket { + self.inner.as_raw_socket() + } } -} -#[cfg(windows)] -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 } + 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 } + } } } diff --git a/src/net/tcp/split.rs b/src/net/tcp/split.rs index 343d4fd..1a1e222 100644 --- a/src/net/tcp/split.rs +++ b/src/net/tcp/split.rs @@ -141,9 +141,9 @@ impl ReadHalf<'_> { /// 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. + /// 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 @@ -153,6 +153,9 @@ impl ReadHalf<'_> { /// /// 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 @@ -275,9 +278,9 @@ impl ReadHalf<'_> { impl WriteHalf<'_> { /// 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. + /// 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 @@ -287,6 +290,9 @@ impl WriteHalf<'_> { /// /// 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 diff --git a/src/net/tcp/split_owned.rs b/src/net/tcp/split_owned.rs index 53fc5f0..6771d64 100644 --- a/src/net/tcp/split_owned.rs +++ b/src/net/tcp/split_owned.rs @@ -196,9 +196,9 @@ impl OwnedReadHalf { /// 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. + /// 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 @@ -208,6 +208,9 @@ impl OwnedReadHalf { /// /// 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 @@ -357,9 +360,9 @@ impl OwnedWriteHalf { /// 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. + /// 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 @@ -369,6 +372,9 @@ impl OwnedWriteHalf { /// /// 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 diff --git a/src/net/tcp/stream.rs b/src/net/tcp/stream.rs index b17d33f..707f73d 100644 --- a/src/net/tcp/stream.rs +++ b/src/net/tcp/stream.rs @@ -8,7 +8,6 @@ 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::convert::TryFrom; use std::fmt; use std::io; use std::net::{Shutdown, SocketAddr}; @@ -259,7 +258,7 @@ impl TcpStream { .map(|raw_socket| unsafe { std::net::TcpStream::from_raw_socket(raw_socket) }) } - #[cfg(tokio_wasi)] + #[cfg(target_os = "wasi")] { use std::os::wasi::io::{FromRawFd, IntoRawFd}; self.io @@ -1016,6 +1015,42 @@ impl TcpStream { .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<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + 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. @@ -1342,21 +1377,31 @@ mod sys { self.io.as_raw_fd() } } + + impl AsFd for TcpStream { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } -#[cfg(windows)] -mod sys { - use super::TcpStream; - use std::os::windows::prelude::*; +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, tokio_wasi))] +#[cfg(all(tokio_unstable, target_os = "wasi"))] mod sys { use super::TcpStream; use std::os::wasi::prelude::*; @@ -1366,4 +1411,10 @@ mod sys { self.io.as_raw_fd() } } + + impl AsFd for TcpStream { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } diff --git a/src/net/udp.rs b/src/net/udp.rs index 213d914..547aaae 100644 --- a/src/net/udp.rs +++ b/src/net/udp.rs @@ -1,7 +1,6 @@ use crate::io::{Interest, PollEvented, ReadBuf, Ready}; use crate::net::{to_socket_addrs, ToSocketAddrs}; -use std::convert::TryFrom; use std::fmt; use std::io; use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr}; @@ -826,7 +825,7 @@ impl UdpSocket { /// 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 + /// 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. /// @@ -882,10 +881,12 @@ impl UdpSocket { /// Tries to receive data from the stream into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. /// - /// The function must be called with valid byte array buf of sufficient size + /// 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()`. /// @@ -932,28 +933,87 @@ impl UdpSocket { let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] 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<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().async_io(Interest::READABLE, || { + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] 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. /// - /// The function must be called with valid byte array buf of sufficient size + /// 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 @@ -996,10 +1056,10 @@ impl UdpSocket { let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; - // Safety: We trust `UdpSocket::recv_from` to have filled up `n` bytes in the - // buffer. 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); } @@ -1007,6 +1067,62 @@ impl UdpSocket { 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<B: BufMut>(&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<u8>] 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 @@ -1177,6 +1293,15 @@ impl UdpSocket { /// 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() @@ -1201,6 +1326,15 @@ impl UdpSocket { /// # 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<'_>, @@ -1226,13 +1360,23 @@ impl UdpSocket { /// Tries to receive 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 + /// 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 @@ -1319,6 +1463,42 @@ impl UdpSocket { .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<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .async_io(interest, || self.io.try_io(&mut f)) + .await + } + /// 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. @@ -1331,6 +1511,17 @@ impl UdpSocket { /// 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 @@ -1349,6 +1540,9 @@ impl UdpSocket { /// 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() @@ -1357,7 +1551,7 @@ impl UdpSocket { } /// Receives data from the socket, without removing it from the input queue. - /// On success, returns the number of bytes read. + /// On success, returns the sending address of the datagram. /// /// # Notes /// @@ -1371,6 +1565,17 @@ impl UdpSocket { /// 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: @@ -1382,6 +1587,9 @@ impl UdpSocket { /// # 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<'_>, @@ -1404,6 +1612,117 @@ impl UdpSocket { 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<SocketAddr> { + 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<io::Result<SocketAddr>> { + 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<SocketAddr> { + self.io + .registration() + .try_io(Interest::READABLE, || self.peek_sender_inner()) + } + + #[inline] + fn peek_sender_inner(&self) -> io::Result<SocketAddr> { + 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`]. @@ -1536,7 +1855,7 @@ impl UdpSocket { /// 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.4.2/src/socket2/socket.rs.html#1178 + // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1464 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", @@ -1563,7 +1882,7 @@ impl UdpSocket { /// /// **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.4.2/src/socket2/socket.rs.html#1178 + // https://docs.rs/socket2/0.5.3/src/socket2/socket.rs.html#1446 #[cfg(not(any( target_os = "fuchsia", target_os = "redox", @@ -1602,7 +1921,7 @@ impl UdpSocket { /// works for some socket types, particularly `AF_INET` sockets. /// /// If `interface` is `None` or an empty string it removes the binding. - #[cfg(all(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] + #[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")))) @@ -1691,7 +2010,7 @@ impl fmt::Debug for UdpSocket { } } -#[cfg(all(unix))] +#[cfg(unix)] mod sys { use super::UdpSocket; use std::os::unix::prelude::*; @@ -1701,16 +2020,27 @@ mod sys { self.io.as_raw_fd() } } + + impl AsFd for UdpSocket { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } + } } -#[cfg(windows)] -mod sys { - use super::UdpSocket; - use std::os::windows::prelude::*; +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()) } + } + } } diff --git a/src/net/unix/datagram/socket.rs b/src/net/unix/datagram/socket.rs index 76c6b19..caebf8e 100644 --- a/src/net/unix/datagram/socket.rs +++ b/src/net/unix/datagram/socket.rs @@ -1,11 +1,10 @@ use crate::io::{Interest, PollEvented, ReadBuf, Ready}; use crate::net::unix::SocketAddr; -use std::convert::TryFrom; use std::fmt; use std::io; use std::net::Shutdown; -use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; use std::task::{Context, Poll}; @@ -90,6 +89,7 @@ cfg_net_unix! { /// # Ok(()) /// # } /// ``` + #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixDatagram { io: PollEvented<mio::net::UnixDatagram>, } @@ -806,6 +806,8 @@ impl UnixDatagram { 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 @@ -865,9 +867,64 @@ impl UnixDatagram { 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 + /// ``` + /// # use std::error::Error; + /// # #[tokio::main] + /// # async fn main() -> Result<(), Box<dyn Error>> { + /// 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<B: BufMut>(&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<u8>] 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 @@ -925,6 +982,52 @@ impl UnixDatagram { 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 + /// ``` + /// # use std::error::Error; + /// # #[tokio::main] + /// # async fn main() -> Result<(), Box<dyn Error>> { + /// 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<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().async_io(Interest::READABLE, || { + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] 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. @@ -1260,6 +1363,42 @@ impl UnixDatagram { .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<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io + .registration() + .async_io(interest, || self.io.try_io(&mut f)) + .await + } + /// Returns the local address that this socket is bound to. /// /// # Examples @@ -1436,3 +1575,9 @@ impl AsRawFd for UnixDatagram { self.io.as_raw_fd() } } + +impl AsFd for UnixDatagram { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} diff --git a/src/net/unix/listener.rs b/src/net/unix/listener.rs index 9887f73..036acd0 100644 --- a/src/net/unix/listener.rs +++ b/src/net/unix/listener.rs @@ -1,10 +1,9 @@ use crate::io::{Interest, PollEvented}; use crate::net::unix::{SocketAddr, UnixStream}; -use std::convert::TryFrom; use std::fmt; use std::io; -use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; use std::task::{Context, Poll}; @@ -44,6 +43,7 @@ cfg_net_unix! { /// } /// } /// ``` + #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixListener { io: PollEvented<mio::net::UnixListener>, } @@ -208,3 +208,9 @@ impl AsRawFd for UnixListener { self.io.as_raw_fd() } } + +impl AsFd for UnixListener { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} diff --git a/src/net/unix/mod.rs b/src/net/unix/mod.rs index 97b6327..a49b70a 100644 --- a/src/net/unix/mod.rs +++ b/src/net/unix/mod.rs @@ -1,4 +1,4 @@ -//! Unix domain socket utility types. +//! 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. @@ -22,6 +22,8 @@ pub(crate) use stream::UnixStream; mod ucred; pub use ucred::UCred; +pub mod pipe; + /// A type representing process and process group IDs. #[allow(non_camel_case_types)] pub type uid_t = u32; diff --git a/src/net/unix/pipe.rs b/src/net/unix/pipe.rs new file mode 100644 index 0000000..efe9b5e --- /dev/null +++ b/src/net/unix/pipe.rs @@ -0,0 +1,1218 @@ +//! 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::{FileTypeExt, OpenOptionsExt}; +use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; +use std::path::Path; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + use bytes::BufMut; +} + +/// 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`][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 +/// [result]: std::io::Result +/// +/// # 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<dyn Error>> { +/// 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<dyn Error>> { +/// 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<dyn Error>> { + /// 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<P: AsRef<Path>>(&self, path: P) -> io::Result<Receiver> { + 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<P: AsRef<Path>>(&self, path: P) -> io::Result<Sender> { + let file = self.open(path.as_ref(), PipeEnd::Sender)?; + Sender::from_file_unchecked(file) + } + + fn open(&self, path: &Path, pipe_end: PipeEnd) -> io::Result<File> { + 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_fifo(&file)? { + 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<dyn std::error::Error>> { +/// // 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<dyn std::error::Error>> { +/// 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<mio_pipe::Sender>, +} + +impl Sender { + fn from_mio(mio_tx: mio_pipe::Sender) -> io::Result<Sender> { + 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(mut file: File) -> io::Result<Sender> { + if !is_fifo(&file)? { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); + } + + let flags = get_file_flags(&file)?; + if has_write_access(flags) { + set_nonblocking(&mut file, flags)?; + Sender::from_file_unchecked(file) + } 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<dyn Error>> { + /// 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> { + let raw_fd = file.into_raw_fd(); + let mio_tx = unsafe { mio_pipe::Sender::from_raw_fd(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<Ready> { + 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<io::Result<()>> { + 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<usize> { + 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<usize> { + self.io + .registration() + .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) + } +} + +impl AsyncWrite for Sender { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.io.poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.io.poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + 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<dyn std::error::Error>> { +/// 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<dyn Error>> { +/// 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<mio_pipe::Receiver>, +} + +impl Receiver { + fn from_mio(mio_rx: mio_pipe::Receiver) -> io::Result<Receiver> { + 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(mut file: File) -> io::Result<Receiver> { + if !is_fifo(&file)? { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "not a pipe")); + } + + let flags = get_file_flags(&file)?; + if has_read_access(flags) { + set_nonblocking(&mut file, flags)?; + Receiver::from_file_unchecked(file) + } 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<dyn Error>> { + /// 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> { + let raw_fd = file.into_raw_fd(); + let mio_rx = unsafe { mio_pipe::Receiver::from_raw_fd(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<Ready> { + 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<io::Result<()>> { + 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<usize> { + 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<usize> { + 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<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + 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<u8>] 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) + }) + } + } +} + +impl AsyncRead for Receiver { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + // 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 file is a FIFO +fn is_fifo(file: &File) -> io::Result<bool> { + Ok(file.metadata()?.file_type().is_fifo()) +} + +/// Gets file descriptor's flags by fcntl. +fn get_file_flags(file: &File) -> io::Result<libc::c_int> { + let fd = file.as_raw_fd(); + let flags = unsafe { libc::fcntl(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's flags with O_NONBLOCK by fcntl. +fn set_nonblocking(file: &mut File, current_flags: libc::c_int) -> io::Result<()> { + let fd = file.as_raw_fd(); + + let flags = current_flags | libc::O_NONBLOCK; + + if flags != current_flags { + let ret = unsafe { libc::fcntl(fd, libc::F_SETFL, flags) }; + if ret < 0 { + return Err(io::Error::last_os_error()); + } + } + + Ok(()) +} diff --git a/src/net/unix/split.rs b/src/net/unix/split.rs index 816a257..6fc7067 100644 --- a/src/net/unix/split.rs +++ b/src/net/unix/split.rs @@ -55,9 +55,20 @@ pub(crate) fn split(stream: &mut UnixStream) -> (ReadHalf<'_>, WriteHalf<'_>) { impl ReadHalf<'_> { /// Wait 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. + /// 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 /// @@ -178,9 +189,9 @@ impl ReadHalf<'_> { impl WriteHalf<'_> { /// 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. + /// 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 @@ -188,6 +199,11 @@ impl WriteHalf<'_> { /// [`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 diff --git a/src/net/unix/split_owned.rs b/src/net/unix/split_owned.rs index 2cb561d..d1c6f93 100644 --- a/src/net/unix/split_owned.rs +++ b/src/net/unix/split_owned.rs @@ -110,9 +110,9 @@ impl OwnedReadHalf { /// 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. + /// 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 @@ -120,6 +120,11 @@ impl OwnedReadHalf { /// [`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 @@ -267,9 +272,9 @@ impl OwnedWriteHalf { /// 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. + /// 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 @@ -277,6 +282,11 @@ impl OwnedWriteHalf { /// [`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 diff --git a/src/net/unix/stream.rs b/src/net/unix/stream.rs index c249bf4..241bcb2 100644 --- a/src/net/unix/stream.rs +++ b/src/net/unix/stream.rs @@ -5,11 +5,10 @@ use crate::net::unix::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf}; use crate::net::unix::ucred::{self, UCred}; use crate::net::unix::SocketAddr; -use std::convert::TryFrom; use std::fmt; use std::io::{self, Read, Write}; use std::net::Shutdown; -use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd}; use std::os::unix::net; use std::path::Path; use std::pin::Pin; @@ -33,6 +32,7 @@ cfg_net_unix! { /// /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown /// [`UnixListener::accept`]: crate::net::UnixListener::accept + #[cfg_attr(docsrs, doc(alias = "uds"))] pub struct UnixStream { io: PollEvented<mio::net::UnixStream>, } @@ -706,6 +706,42 @@ impl UnixStream { .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<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + 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 @@ -1000,3 +1036,9 @@ impl AsRawFd for UnixStream { self.io.as_raw_fd() } } + +impl AsFd for UnixStream { + fn as_fd(&self) -> BorrowedFd<'_> { + unsafe { BorrowedFd::borrow_raw(self.as_raw_fd()) } + } +} diff --git a/src/net/unix/ucred.rs b/src/net/unix/ucred.rs index 3cb61d7..43f008b 100644 --- a/src/net/unix/ucred.rs +++ b/src/net/unix/ucred.rs @@ -31,10 +31,15 @@ impl UCred { } } -#[cfg(any(target_os = "linux", target_os = "android", target_os = "openbsd"))] +#[cfg(any( + target_os = "linux", + target_os = "redox", + target_os = "android", + target_os = "openbsd" +))] pub(crate) use self::impl_linux::get_peer_cred; -#[cfg(any(target_os = "netbsd"))] +#[cfg(target_os = "netbsd")] pub(crate) use self::impl_netbsd::get_peer_cred; #[cfg(any(target_os = "dragonfly", target_os = "freebsd"))] @@ -46,7 +51,18 @@ 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(any(target_os = "linux", target_os = "android", target_os = "openbsd"))] +#[cfg(target_os = "aix")] +pub(crate) use self::impl_aix::get_peer_cred; + +#[cfg(target_os = "espidf")] +pub(crate) use self::impl_noproc::get_peer_cred; + +#[cfg(any( + target_os = "linux", + target_os = "redox", + target_os = "android", + target_os = "openbsd" +))] pub(crate) mod impl_linux { use crate::net::unix::{self, UnixStream}; @@ -55,7 +71,7 @@ pub(crate) mod impl_linux { #[cfg(target_os = "openbsd")] use libc::sockpeercred as ucred; - #[cfg(any(target_os = "linux", target_os = "android"))] + #[cfg(any(target_os = "linux", target_os = "redox", target_os = "android"))] use libc::ucred; pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> { @@ -98,7 +114,7 @@ pub(crate) mod impl_linux { } } -#[cfg(any(target_os = "netbsd"))] +#[cfg(target_os = "netbsd")] pub(crate) mod impl_netbsd { use crate::net::unix::{self, UnixStream}; @@ -250,3 +266,45 @@ pub(crate) mod impl_solaris { } } } + +#[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<super::UCred> { + 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(target_os = "espidf")] +pub(crate) mod impl_noproc { + use crate::net::unix::UnixStream; + use std::io; + + pub(crate) fn get_peer_cred(_sock: &UnixStream) -> io::Result<super::UCred> { + Ok(super::UCred { + uid: 0, + gid: 0, + pid: None, + }) + } +} diff --git a/src/net/windows/named_pipe.rs b/src/net/windows/named_pipe.rs index 9ede94e..a03e1d0 100644 --- a/src/net/windows/named_pipe.rs +++ b/src/net/windows/named_pipe.rs @@ -10,7 +10,7 @@ use std::ptr; use std::task::{Context, Poll}; use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; -use crate::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle}; +use crate::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, RawHandle}; cfg_io_util! { use bytes::BufMut; @@ -851,6 +851,39 @@ impl NamedPipeServer { ) -> io::Result<R> { 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<R>( + &self, + interest: Interest, + f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().async_io(interest, f).await + } } impl AsyncRead for NamedPipeServer { @@ -895,6 +928,12 @@ impl AsRawHandle for NamedPipeServer { } } +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`]. @@ -1601,6 +1640,39 @@ impl NamedPipeClient { ) -> io::Result<R> { 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<R>( + &self, + interest: Interest, + f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().async_io(interest, f).await + } } impl AsyncRead for NamedPipeClient { @@ -1645,17 +1717,10 @@ impl AsRawHandle for NamedPipeClient { } } -// Helper to set a boolean flag as a bitfield. -macro_rules! bool_flag { - ($f:expr, $t:expr, $flag:expr) => {{ - let current = $f; - - if $t { - $f = current | $flag; - } else { - $f = current & !$flag; - }; - }}; +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 @@ -1665,8 +1730,17 @@ macro_rules! bool_flag { /// See [`ServerOptions::create`]. #[derive(Debug, Clone)] pub struct ServerOptions { - open_mode: u32, - pipe_mode: u32, + // 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, @@ -1687,8 +1761,14 @@ impl ServerOptions { /// ``` pub fn new() -> ServerOptions { ServerOptions { - open_mode: windows_sys::PIPE_ACCESS_DUPLEX | windows_sys::FILE_FLAG_OVERLAPPED, - pipe_mode: windows_sys::PIPE_TYPE_BYTE | windows_sys::PIPE_REJECT_REMOTE_CLIENTS, + 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, @@ -1701,14 +1781,11 @@ impl ServerOptions { /// The default pipe mode is [`PipeMode::Byte`]. See [`PipeMode`] for /// documentation of what each mode means. /// - /// This corresponding to specifying [`dwPipeMode`]. + /// 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 { - let is_msg = matches!(pipe_mode, PipeMode::Message); - // Pipe mode is implemented as a bit flag 0x4. Set is message and unset - // is byte. - bool_flag!(self.pipe_mode, is_msg, windows_sys::PIPE_TYPE_MESSAGE); + self.pipe_mode = pipe_mode; self } @@ -1804,7 +1881,7 @@ impl ServerOptions { /// # Ok(()) } /// ``` pub fn access_inbound(&mut self, allowed: bool) -> &mut Self { - bool_flag!(self.open_mode, allowed, windows_sys::PIPE_ACCESS_INBOUND); + self.access_inbound = allowed; self } @@ -1902,7 +1979,7 @@ impl ServerOptions { /// # Ok(()) } /// ``` pub fn access_outbound(&mut self, allowed: bool) -> &mut Self { - bool_flag!(self.open_mode, allowed, windows_sys::PIPE_ACCESS_OUTBOUND); + self.access_outbound = allowed; self } @@ -1970,11 +2047,7 @@ impl ServerOptions { /// [`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 { - bool_flag!( - self.open_mode, - first, - windows_sys::FILE_FLAG_FIRST_PIPE_INSTANCE - ); + self.first_pipe_instance = first; self } @@ -2056,7 +2129,7 @@ impl ServerOptions { /// /// [`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 { - bool_flag!(self.open_mode, requested, windows_sys::WRITE_DAC); + self.write_dac = requested; self } @@ -2066,7 +2139,7 @@ impl ServerOptions { /// /// [`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 { - bool_flag!(self.open_mode, requested, windows_sys::WRITE_OWNER); + self.write_owner = requested; self } @@ -2076,11 +2149,7 @@ impl ServerOptions { /// /// [`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 { - bool_flag!( - self.open_mode, - requested, - windows_sys::ACCESS_SYSTEM_SECURITY - ); + self.access_system_security = requested; self } @@ -2091,11 +2160,7 @@ impl ServerOptions { /// /// [`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 { - bool_flag!( - self.pipe_mode, - reject, - windows_sys::PIPE_REJECT_REMOTE_CLIENTS - ); + self.reject_remote_clients = reject; self } @@ -2241,10 +2306,46 @@ impl ServerOptions { ) -> io::Result<NamedPipeServer> { 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(), - self.open_mode, - self.pipe_mode, + open_mode, + pipe_mode, self.max_instances, self.out_buffer_size, self.in_buffer_size, @@ -2266,8 +2367,10 @@ impl ServerOptions { /// See [`ClientOptions::open`]. #[derive(Debug, Clone)] pub struct ClientOptions { - desired_access: u32, + generic_read: bool, + generic_write: bool, security_qos_flags: u32, + pipe_mode: PipeMode, } impl ClientOptions { @@ -2286,9 +2389,11 @@ impl ClientOptions { /// ``` pub fn new() -> Self { Self { - desired_access: windows_sys::GENERIC_READ | windows_sys::GENERIC_WRITE, + generic_read: true, + generic_write: true, security_qos_flags: windows_sys::SECURITY_IDENTIFICATION | windows_sys::SECURITY_SQOS_PRESENT, + pipe_mode: PipeMode::Byte, } } @@ -2299,7 +2404,7 @@ impl ClientOptions { /// [`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 { - bool_flag!(self.desired_access, allowed, windows_sys::GENERIC_READ); + self.generic_read = allowed; self } @@ -2310,7 +2415,7 @@ impl ClientOptions { /// [`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 { - bool_flag!(self.desired_access, allowed, windows_sys::GENERIC_WRITE); + self.generic_write = allowed; self } @@ -2341,6 +2446,15 @@ impl ClientOptions { 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 @@ -2419,13 +2533,24 @@ impl ClientOptions { ) -> io::Result<NamedPipeClient> { 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(), - self.desired_access, + desired_access, 0, attrs as *mut _, windows_sys::OPEN_EXISTING, @@ -2437,6 +2562,16 @@ impl ClientOptions { 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 _) } @@ -2553,48 +2688,3 @@ unsafe fn named_pipe_info(handle: RawHandle) -> io::Result<PipeInfo> { max_instances, }) } - -#[cfg(test)] -mod test { - use self::windows_sys::{PIPE_REJECT_REMOTE_CLIENTS, PIPE_TYPE_BYTE, PIPE_TYPE_MESSAGE}; - use super::*; - - #[test] - fn opts_default_pipe_mode() { - let opts = ServerOptions::new(); - assert_eq!(opts.pipe_mode, PIPE_TYPE_BYTE | PIPE_REJECT_REMOTE_CLIENTS); - } - - #[test] - fn opts_unset_reject_remote() { - let mut opts = ServerOptions::new(); - opts.reject_remote_clients(false); - assert_eq!(opts.pipe_mode & PIPE_REJECT_REMOTE_CLIENTS, 0); - } - - #[test] - fn opts_set_pipe_mode_maintains_reject_remote_clients() { - let mut opts = ServerOptions::new(); - opts.pipe_mode(PipeMode::Byte); - assert_eq!(opts.pipe_mode, PIPE_TYPE_BYTE | PIPE_REJECT_REMOTE_CLIENTS); - - opts.reject_remote_clients(false); - opts.pipe_mode(PipeMode::Byte); - assert_eq!(opts.pipe_mode, PIPE_TYPE_BYTE); - - opts.reject_remote_clients(true); - opts.pipe_mode(PipeMode::Byte); - assert_eq!(opts.pipe_mode, PIPE_TYPE_BYTE | PIPE_REJECT_REMOTE_CLIENTS); - - opts.reject_remote_clients(false); - opts.pipe_mode(PipeMode::Message); - assert_eq!(opts.pipe_mode, PIPE_TYPE_MESSAGE); - - opts.reject_remote_clients(true); - opts.pipe_mode(PipeMode::Message); - assert_eq!( - opts.pipe_mode, - PIPE_TYPE_MESSAGE | PIPE_REJECT_REMOTE_CLIENTS - ); - } -} diff --git a/src/process/mod.rs b/src/process/mod.rs index 66e4212..5cf06fe 100644 --- a/src/process/mod.rs +++ b/src/process/mod.rs @@ -156,7 +156,6 @@ //! ```no_run //! use tokio::join; //! use tokio::process::Command; -//! use std::convert::TryInto; //! use std::process::Stdio; //! //! #[tokio::main] @@ -217,7 +216,7 @@ //! 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 regards +//! 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 @@ -245,22 +244,24 @@ mod kill; use crate::io::{AsyncRead, AsyncWrite, ReadBuf}; use crate::process::kill::Kill; -use std::convert::TryInto; use std::ffi::OsStr; use std::future::Future; use std::io; -#[cfg(unix)] -use std::os::unix::process::CommandExt; -#[cfg(windows)] -use std::os::windows::io::{AsRawHandle, RawHandle}; -#[cfg(windows)] -use std::os::windows::process::CommandExt; use std::path::Path; use std::pin::Pin; use std::process::{Command as StdCommand, ExitStatus, Output, Stdio}; use std::task::Context; use std::task::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 @@ -309,7 +310,8 @@ impl Command { /// /// ```no_run /// use tokio::process::Command; - /// let command = Command::new("sh"); + /// let mut command = Command::new("sh"); + /// # let _ = command.output(); // assert borrow checker /// ``` /// /// [rust-lang/rust#37519]: https://github.com/rust-lang/rust/issues/37519 @@ -328,16 +330,20 @@ impl Command { /// Only one argument can be passed per use. So instead of: /// /// ```no_run - /// tokio::process::Command::new("sh") - /// .arg("-C /path/to/repo"); + /// 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 - /// tokio::process::Command::new("sh") - /// .arg("-C") - /// .arg("/path/to/repo"); + /// 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`]. @@ -349,11 +355,15 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") + /// let output = Command::new("ls") /// .arg("-l") - /// .arg("-a"); + /// .arg("-a") + /// .output().await.unwrap(); + /// # } + /// /// ``` pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Command { self.std.arg(arg); @@ -371,10 +381,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .args(&["-l", "-a"]); + /// let output = Command::new("ls") + /// .args(&["-l", "-a"]) + /// .output().await.unwrap(); + /// # } /// ``` pub fn args<I, S>(&mut self, args: I) -> &mut Command where @@ -385,6 +398,17 @@ impl Command { 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<S: AsRef<OsStr>>(&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, @@ -395,10 +419,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .env("PATH", "/bin"); + /// let output = Command::new("ls") + /// .env("PATH", "/bin") + /// .output().await.unwrap(); + /// # } /// ``` pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Command where @@ -416,6 +443,7 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::{Stdio}; /// use std::env; @@ -426,11 +454,13 @@ impl Command { /// k == "TERM" || k == "TZ" || k == "LANG" || k == "PATH" /// ).collect(); /// - /// let command = Command::new("printenv") + /// let output = Command::new("printenv") /// .stdin(Stdio::null()) /// .stdout(Stdio::inherit()) /// .env_clear() - /// .envs(&filtered_env); + /// .envs(&filtered_env) + /// .output().await.unwrap(); + /// # } /// ``` pub fn envs<I, K, V>(&mut self, vars: I) -> &mut Command where @@ -449,10 +479,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .env_remove("PATH"); + /// let output = Command::new("ls") + /// .env_remove("PATH") + /// .output().await.unwrap(); + /// # } /// ``` pub fn env_remove<K: AsRef<OsStr>>(&mut self, key: K) -> &mut Command { self.std.env_remove(key); @@ -466,10 +499,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .env_clear(); + /// let output = Command::new("ls") + /// .env_clear() + /// .output().await.unwrap(); + /// # } /// ``` pub fn env_clear(&mut self) -> &mut Command { self.std.env_clear(); @@ -493,10 +529,13 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .current_dir("/bin"); + /// let output = Command::new("ls") + /// .current_dir("/bin") + /// .output().await.unwrap(); + /// # } /// ``` pub fn current_dir<P: AsRef<Path>>(&mut self, dir: P) -> &mut Command { self.std.current_dir(dir); @@ -516,11 +555,14 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use std::process::{Stdio}; /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .stdin(Stdio::null()); + /// let output = Command::new("ls") + /// .stdin(Stdio::null()) + /// .output().await.unwrap(); + /// # } /// ``` pub fn stdin<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Command { self.std.stdin(cfg); @@ -540,11 +582,14 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::Stdio; /// - /// let command = Command::new("ls") - /// .stdout(Stdio::null()); + /// let output = Command::new("ls") + /// .stdout(Stdio::null()) + /// .output().await.unwrap(); + /// # } /// ``` pub fn stdout<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Command { self.std.stdout(cfg); @@ -564,11 +609,14 @@ impl Command { /// Basic usage: /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// use std::process::{Stdio}; /// - /// let command = Command::new("ls") - /// .stderr(Stdio::null()); + /// let output = Command::new("ls") + /// .stderr(Stdio::null()) + /// .output().await.unwrap(); + /// # } /// ``` pub fn stderr<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Command { self.std.stderr(cfg); @@ -595,7 +643,7 @@ impl Command { /// 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 makes no additional guarantees are made with regards + /// 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 @@ -606,16 +654,16 @@ impl Command { self } - /// 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 - #[cfg(windows)] - #[cfg_attr(docsrs, doc(cfg(windows)))] - pub fn creation_flags(&mut self, flags: u32) -> &mut Command { - self.std.creation_flags(flags); - 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 @@ -707,10 +755,13 @@ impl Command { /// [`tokio::process::Command`]: crate::process::Command /// /// ```no_run + /// # async fn test() { // allow using await /// use tokio::process::Command; /// - /// let command = Command::new("ls") - /// .process_group(0); + /// let output = Command::new("ls") + /// .process_group(0) + /// .output().await.unwrap(); + /// # } /// ``` #[cfg(unix)] #[cfg(tokio_unstable)] @@ -772,7 +823,7 @@ impl Command { /// 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 regards + /// 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 @@ -953,6 +1004,7 @@ where type Output = Result<T, E>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); @@ -1043,13 +1095,14 @@ impl Child { } } - /// Extracts the raw handle of the process associated with this child while - /// it is still running. Returns `None` if the child has exited. - #[cfg(windows)] - pub fn raw_handle(&self) -> Option<RawHandle> { - match &self.child { - FusedChild::Child(c) => Some(c.inner.as_raw_handle()), - 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<RawHandle> { + match &self.child { + FusedChild::Child(c) => Some(c.inner.as_raw_handle()), + FusedChild::Done(_) => None, + } } } @@ -1284,7 +1337,7 @@ impl ChildStdin { } impl ChildStdout { - /// Creates an asynchronous `ChildStderr` from a synchronous one. + /// Creates an asynchronous `ChildStdout` from a synchronous one. /// /// # Errors /// @@ -1388,53 +1441,100 @@ impl TryInto<Stdio> for ChildStderr { } #[cfg(unix)] +#[cfg_attr(docsrs, doc(cfg(unix)))] mod sys { - use std::os::unix::io::{AsRawFd, RawFd}; + use std::{ + io, + os::unix::io::{AsFd, AsRawFd, BorrowedFd, OwnedFd, RawFd}, + }; use super::{ChildStderr, ChildStdin, ChildStdout}; - impl AsRawFd for ChildStdin { - fn as_raw_fd(&self) -> RawFd { - self.inner.as_raw_fd() - } - } + macro_rules! impl_traits { + ($type:ty) => { + impl $type { + /// Convert into [`OwnedFd`]. + pub fn into_owned_fd(self) -> io::Result<OwnedFd> { + self.inner.into_owned_fd() + } + } - impl AsRawFd for ChildStdout { - fn as_raw_fd(&self) -> RawFd { - self.inner.as_raw_fd() - } - } + impl AsRawFd for $type { + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } + } - impl AsRawFd for ChildStderr { - 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(windows)] -mod sys { - use std::os::windows::io::{AsRawHandle, RawHandle}; +#[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<OwnedHandle> { + self.inner.into_owned_handle() + } + } - use super::{ChildStderr, ChildStdin, ChildStdout}; + impl AsRawHandle for $type { + fn as_raw_handle(&self) -> RawHandle { + self.inner.as_raw_handle() + } + } - impl AsRawHandle for ChildStdin { - 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()) } + } + } + }; } - impl AsRawHandle for ChildStdout { - fn as_raw_handle(&self) -> RawHandle { - self.inner.as_raw_handle() - } - } + #[cfg(docsrs)] + macro_rules! impl_traits { + ($type:ty) => { + impl $type { + /// Convert into [`OwnedHandle`]. + pub fn into_owned_handle(self) -> io::Result<OwnedHandle> { + todo!("For doc generation only") + } + } - impl AsRawHandle for ChildStderr { - fn as_raw_handle(&self) -> RawHandle { - self.inner.as_raw_handle() - } + 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)))] diff --git a/src/process/unix/mod.rs b/src/process/unix/mod.rs index 78c792c..b9c2d78 100644 --- a/src/process/unix/mod.rs +++ b/src/process/unix/mod.rs @@ -39,7 +39,7 @@ use std::fmt; use std::fs::File; use std::future::Future; use std::io; -use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +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; @@ -194,7 +194,13 @@ impl AsRawFd for Pipe { } } -pub(crate) fn convert_to_stdio(io: ChildStdio) -> io::Result<Stdio> { +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<File> { let mut fd = io.inner.into_inner()?.fd; // Ensure that the fd to be inherited is set to *blocking* mode, as this @@ -203,7 +209,11 @@ pub(crate) fn convert_to_stdio(io: ChildStdio) -> io::Result<Stdio> { // change it to nonblocking mode. set_nonblocking(&mut fd, false)?; - Ok(Stdio::from(fd)) + Ok(fd) +} + +pub(crate) fn convert_to_stdio(io: ChildStdio) -> io::Result<Stdio> { + convert_to_blocking_file(io).map(Stdio::from) } impl Source for Pipe { @@ -234,6 +244,12 @@ pub(crate) struct ChildStdio { inner: PollEvented<Pipe>, } +impl ChildStdio { + pub(super) fn into_owned_fd(self) -> io::Result<OwnedFd> { + 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) @@ -246,6 +262,12 @@ impl AsRawFd for ChildStdio { } } +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>, diff --git a/src/process/windows.rs b/src/process/windows.rs index 0e4c7dd..792a9c9 100644 --- a/src/process/windows.rs +++ b/src/process/windows.rs @@ -24,7 +24,7 @@ use std::fmt; use std::fs::File as StdFile; use std::future::Future; use std::io; -use std::os::windows::prelude::{AsRawHandle, IntoRawHandle, RawHandle}; +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}; @@ -36,10 +36,9 @@ use windows_sys::{ DuplicateHandle, BOOLEAN, DUPLICATE_SAME_ACCESS, HANDLE, INVALID_HANDLE_VALUE, }, Win32::System::Threading::{ - GetCurrentProcess, RegisterWaitForSingleObject, UnregisterWaitEx, WT_EXECUTEINWAITTHREAD, - WT_EXECUTEONLYONCE, + GetCurrentProcess, RegisterWaitForSingleObject, UnregisterWaitEx, INFINITE, + WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE, }, - Win32::System::WindowsProgramming::INFINITE, }; #[must_use = "futures do nothing unless polled"] @@ -196,6 +195,12 @@ pub(crate) struct ChildStdio { io: Blocking<ArcFile>, } +impl ChildStdio { + pub(super) fn into_owned_handle(self) -> io::Result<OwnedHandle> { + convert_to_file(self).map(OwnedHandle::from) + } +} + impl AsRawHandle for ChildStdio { fn as_raw_handle(&self) -> RawHandle { self.raw.as_raw_handle() @@ -241,13 +246,15 @@ where Ok(ChildStdio { raw, io }) } -pub(crate) fn convert_to_stdio(child_stdio: ChildStdio) -> io::Result<Stdio> { +fn convert_to_file(child_stdio: ChildStdio) -> io::Result<StdFile> { 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)) - .map(Stdio::from) + Arc::try_unwrap(raw).or_else(|raw| duplicate_handle(&*raw)) +} + +pub(crate) fn convert_to_stdio(child_stdio: ChildStdio) -> io::Result<Stdio> { + convert_to_file(child_stdio).map(Stdio::from) } fn duplicate_handle<T: AsRawHandle>(io: &T) -> io::Result<StdFile> { diff --git a/src/runtime/blocking/pool.rs b/src/runtime/blocking/pool.rs index e9f6b66..3377867 100644 --- a/src/runtime/blocking/pool.rs +++ b/src/runtime/blocking/pool.rs @@ -173,7 +173,7 @@ const KEEP_ALIVE: Duration = Duration::from_secs(10); /// Tasks will be scheduled as non-mandatory, meaning they may not get executed /// in case of runtime shutdown. #[track_caller] -#[cfg_attr(tokio_wasi, allow(dead_code))] +#[cfg_attr(target_os = "wasi", allow(dead_code))] pub(crate) fn spawn_blocking<F, R>(func: F) -> JoinHandle<R> where F: FnOnce() -> R + Send + 'static, @@ -371,7 +371,9 @@ impl Spawner { task.name = %name.unwrap_or_default(), task.id = id.as_u64(), "fn" = %std::any::type_name::<F>(), - spawn.location = %format_args!("{}:{}:{}", location.file(), location.line(), location.column()), + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), ); fut.instrument(span) }; diff --git a/src/runtime/blocking/schedule.rs b/src/runtime/blocking/schedule.rs index edf775b..8dfe5fd 100644 --- a/src/runtime/blocking/schedule.rs +++ b/src/runtime/blocking/schedule.rs @@ -23,8 +23,10 @@ impl BlockingSchedule { scheduler::Handle::CurrentThread(handle) => { handle.driver.clock.inhibit_auto_advance(); } - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] scheduler::Handle::MultiThread(_) => {} + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + scheduler::Handle::MultiThreadAlt(_) => {} } } BlockingSchedule { @@ -43,8 +45,10 @@ impl task::Schedule for BlockingSchedule { handle.driver.clock.allow_auto_advance(); handle.driver.unpark(); } - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] scheduler::Handle::MultiThread(_) => {} + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + scheduler::Handle::MultiThreadAlt(_) => {} } } None diff --git a/src/runtime/builder.rs b/src/runtime/builder.rs index ea0df2e..03f1678 100644 --- a/src/runtime/builder.rs +++ b/src/runtime/builder.rs @@ -1,5 +1,5 @@ use crate::runtime::handle::Handle; -use crate::runtime::{blocking, driver, Callback, Runtime}; +use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime}; use crate::util::rand::{RngSeed, RngSeedGenerator}; use std::fmt; @@ -82,11 +82,19 @@ pub struct Builder { pub(super) keep_alive: Option<Duration>, /// How many ticks before pulling a task from the global/remote queue? - pub(super) global_queue_interval: u32, + /// + /// 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<u32>, /// 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. @@ -95,6 +103,12 @@ pub struct Builder { /// 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, } @@ -185,8 +199,10 @@ pub(crate) type ThreadNameFn = std::sync::Arc<dyn Fn() -> String + Send + Sync + #[derive(Clone, Copy)] pub(crate) enum Kind { CurrentThread, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] MultiThread, + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + MultiThreadAlt, } impl Builder { @@ -205,7 +221,7 @@ impl Builder { #[cfg(not(loom))] const EVENT_INTERVAL: u32 = 61; - Builder::new(Kind::CurrentThread, 31, EVENT_INTERVAL) + Builder::new(Kind::CurrentThread, EVENT_INTERVAL) } cfg_not_wasi! { @@ -216,7 +232,27 @@ impl Builder { #[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, 61) + 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) + } } } @@ -224,7 +260,7 @@ impl Builder { /// values. /// /// Configuration methods can be chained on the return value. - pub(crate) fn new(kind: Kind, global_queue_interval: u32, event_interval: u32) -> Builder { + pub(crate) fn new(kind: Kind, event_interval: u32) -> Builder { Builder { kind, @@ -260,14 +296,24 @@ impl Builder { // Defaults for these values depend on the scheduler kind, so we get them // as parameters. - global_queue_interval, + 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: Default::default(), + disable_lifo_slot: false, } } @@ -358,9 +404,17 @@ impl Builder { /// Specifies the limit for additional threads spawned by the Runtime. /// /// These threads are used for blocking operations like tasks spawned - /// through [`spawn_blocking`]. 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`]. + /// 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. /// @@ -374,6 +428,11 @@ impl Builder { /// 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] @@ -638,8 +697,10 @@ impl Builder { pub fn build(&mut self) -> io::Result<Runtime> { match &self.kind { Kind::CurrentThread => self.build_current_thread_runtime(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Kind::MultiThread => self.build_threaded_runtime(), + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Kind::MultiThreadAlt => self.build_alt_threaded_runtime(), } } @@ -647,8 +708,10 @@ impl Builder { driver::Cfg { enable_pause_time: match self.kind { Kind::CurrentThread => true, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Kind::MultiThread => false, + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Kind::MultiThreadAlt => false, }, enable_io: self.enable_io, enable_time: self.enable_time, @@ -706,7 +769,7 @@ impl Builder { /// # } /// ``` pub fn global_queue_interval(&mut self, val: u32) -> &mut Self { - self.global_queue_interval = val; + self.global_queue_interval = Some(val); self } @@ -877,6 +940,141 @@ impl Builder { } } + cfg_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_count_histogram_` builder methods to configure the + /// histogram details. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_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_count_histogram_num_buckets(), 10); + /// # assert_eq!(m.poll_count_histogram_bucket_range(0), us(0)..us(100)); + /// # assert_eq!(m.poll_count_histogram_bucket_range(1), us(100)..us(200)); + /// ``` + /// + /// [`Handle::metrics()`]: crate::runtime::Handle::metrics + /// [`Instant::now()`]: std::time::Instant::now + pub fn enable_metrics_poll_count_histogram(&mut self) -> &mut Self { + self.metrics_poll_count_histogram_enable = true; + self + } + + /// 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}; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .metrics_poll_count_histogram_scale(HistogramScale::Log) + /// .build() + /// .unwrap(); + /// ``` + pub fn metrics_poll_count_histogram_scale(&mut self, histogram_scale: crate::runtime::HistogramScale) -> &mut Self { + self.metrics_poll_count_histogram.scale = histogram_scale; + 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; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .metrics_poll_count_histogram_resolution(Duration::from_micros(100)) + /// .build() + /// .unwrap(); + /// ``` + 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.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; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_metrics_poll_count_histogram() + /// .metrics_poll_count_histogram_buckets(15) + /// .build() + /// .unwrap(); + /// ``` + pub fn metrics_poll_count_histogram_buckets(&mut self, buckets: usize) -> &mut Self { + self.metrics_poll_count_histogram.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<Runtime> { use crate::runtime::scheduler::{self, CurrentThread}; use crate::runtime::{runtime::Scheduler, Config}; @@ -905,10 +1103,12 @@ impl Builder { after_unpark: self.after_unpark.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(), }, ); @@ -922,6 +1122,14 @@ impl Builder { blocking_pool, )) } + + fn metrics_poll_count_histogram_builder(&self) -> Option<HistogramBuilder> { + if self.metrics_poll_count_histogram_enable { + Some(self.metrics_poll_count_histogram.clone()) + } else { + None + } + } } cfg_io_driver! { @@ -1046,10 +1254,12 @@ cfg_rt_multi_thread! { after_unpark: self.after_unpark.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(), }, ); @@ -1061,6 +1271,49 @@ cfg_rt_multi_thread! { Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool)) } + + cfg_unstable! { + fn build_alt_threaded_runtime(&mut self) -> io::Result<Runtime> { + 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())?; + + // 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(), + 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)) + } + } } } diff --git a/src/runtime/config.rs b/src/runtime/config.rs index 39eb1cf..d846a0d 100644 --- a/src/runtime/config.rs +++ b/src/runtime/config.rs @@ -1,14 +1,20 @@ -#![cfg_attr(any(not(feature = "full"), tokio_wasm), allow(dead_code))] +#![cfg_attr( + any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"), + allow(dead_code) +)] use crate::runtime::Callback; 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: u32, + pub(crate) global_queue_interval: Option<u32>, /// 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>, @@ -28,6 +34,9 @@ pub(crate) struct Config { /// deterministic way. pub(crate) seed_generator: RngSeedGenerator, + /// How to build poll time histograms + pub(crate) metrics_poll_count_histogram: Option<crate::runtime::HistogramBuilder>, + #[cfg(tokio_unstable)] /// How to respond to unhandled task panics. pub(crate) unhandled_panic: crate::runtime::UnhandledPanic, diff --git a/src/runtime/context.rs b/src/runtime/context.rs index fef53ca..5943e9a 100644 --- a/src/runtime/context.rs +++ b/src/runtime/context.rs @@ -4,14 +4,33 @@ use crate::runtime::coop; use std::cell::Cell; #[cfg(any(feature = "rt", feature = "macros"))] -use crate::util::rand::{FastRand, RngSeed}; +use crate::util::rand::FastRand; cfg_rt! { - use crate::runtime::{scheduler, task::Id, Defer}; + mod blocking; + pub(crate) use blocking::{disallow_block_in_place, try_enter_blocking_region, BlockingRegionGuard}; - use std::cell::RefCell; - use std::marker::PhantomData; - use std::time::Duration; + 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 { @@ -21,7 +40,11 @@ struct Context { /// Handle to the runtime scheduler running on the current thread. #[cfg(feature = "rt")] - handle: RefCell<Option<scheduler::Handle>>, + current: current::HandleCell, + + /// Handle to the scheduler's internal "context" + #[cfg(feature = "rt")] + scheduler: Scoped<scheduler::Context>, #[cfg(feature = "rt")] current_task_id: Cell<Option<Id>>, @@ -34,21 +57,25 @@ struct Context { #[cfg(feature = "rt")] runtime: Cell<EnterRuntime>, - /// Yielded task wakers are stored here and notified after resource drivers - /// are polled. - #[cfg(feature = "rt")] - defer: RefCell<Option<Defer>>, - #[cfg(any(feature = "rt", feature = "macros"))] - rng: FastRand, + rng: Cell<Option<FastRand>>, /// Tracks the amount of "work" a task may still do before yielding back to /// the sheduler budget: Cell<coop::Budget>, + + #[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 = { + static CONTEXT: Context = const { Context { #[cfg(feature = "rt")] thread_id: Cell::new(None), @@ -56,7 +83,12 @@ tokio_thread_local! { /// Tracks the current runtime handle to use when spawning, /// accessing drivers, etc... #[cfg(feature = "rt")] - handle: RefCell::new(None), + 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), @@ -68,20 +100,35 @@ tokio_thread_local! { #[cfg(feature = "rt")] runtime: Cell::new(EnterRuntime::NotEntered), - #[cfg(feature = "rt")] - defer: RefCell::new(None), - #[cfg(any(feature = "rt", feature = "macros"))] - rng: FastRand::new(RngSeed::new()), + 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(feature = "macros")] +#[cfg(any(feature = "macros", all(feature = "sync", feature = "rt")))] pub(crate) fn thread_rng_n(n: u32) -> u32 { - CONTEXT.with(|ctx| ctx.rng.fastrand_n(n)) + 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<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> { @@ -89,9 +136,7 @@ pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, A } cfg_rt! { - use crate::runtime::{ThreadId, TryCurrentError}; - - use std::fmt; + use crate::runtime::ThreadId; pub(crate) fn thread_id() -> Result<ThreadId, AccessError> { CONTEXT.try_with(|ctx| { @@ -106,48 +151,6 @@ cfg_rt! { }) } - #[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, - } - - #[derive(Debug)] - #[must_use] - pub(crate) struct SetCurrentGuard { - old_handle: Option<scheduler::Handle>, - old_seed: RngSeed, - } - - /// 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, - - /// If true, then this is the root runtime guard. It is possible to nest - /// runtime guards by using `block_in_place` between the calls. We need - /// to track the root guard as this is the guard responsible for freeing - /// the deferred task queue. - is_root: bool, - } - - /// Guard tracking that a caller has entered a blocking region. - #[must_use] - pub(crate) struct BlockingRegionGuard { - _p: PhantomData<RefCell<()>>, - } - - pub(crate) struct DisallowBlockInPlaceGuard(bool); - pub(crate) fn set_current_task_id(id: Option<Id>) -> Option<Id> { CONTEXT.try_with(|ctx| ctx.current_task_id.replace(id)).unwrap_or(None) } @@ -156,265 +159,33 @@ cfg_rt! { CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None) } - pub(crate) fn try_current() -> Result<scheduler::Handle, TryCurrentError> { - match CONTEXT.try_with(|ctx| ctx.handle.borrow().clone()) { - Ok(Some(handle)) => Ok(handle), - Ok(None) => Err(TryCurrentError::new_no_context()), - Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()), - } - } - - /// Sets this [`Handle`] as the current active [`Handle`]. - /// - /// [`Handle`]: crate::runtime::scheduler::Handle - pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> { - CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok() - } - - - /// 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) -> EnterRuntimeGuard { - if let Some(enter) = try_enter_runtime(handle, allow_block_in_place) { - return enter; - } - - 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." - ); - } - - /// Tries to enter a runtime context, returns `None` if already in a runtime - /// context. - fn try_enter_runtime(handle: &scheduler::Handle, allow_block_in_place: bool) -> Option<EnterRuntimeGuard> { - CONTEXT.with(|c| { - if c.runtime.get().is_entered() { - None - } else { - // Set the entered flag - c.runtime.set(EnterRuntime::Entered { allow_block_in_place }); - - // Initialize queue to track yielded tasks - let mut defer = c.defer.borrow_mut(); - - let is_root = if defer.is_none() { - *defer = Some(Defer::new()); - true - } else { - false - }; - - Some(EnterRuntimeGuard { - blocking: BlockingRegionGuard::new(), - handle: c.set_current(handle), - is_root, - }) - } - }) - } - - pub(crate) fn try_enter_blocking_region() -> Option<BlockingRegionGuard> { - CONTEXT.try_with(|c| { - if c.runtime.get().is_entered() { - None + pub(crate) fn defer(waker: &Waker) { + with_scheduler(|maybe_scheduler| { + if let Some(scheduler) = maybe_scheduler { + scheduler.defer(waker); } 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 + // Called from outside of the runtime, immediately wake the + // task. + waker.wake_by_ref(); } }); - - DisallowBlockInPlaceGuard(reset) } - pub(crate) fn with_defer<R>(f: impl FnOnce(&mut Defer) -> R) -> Option<R> { - CONTEXT.with(|c| { - let mut defer = c.defer.borrow_mut(); - defer.as_mut().map(f) - }) - } - - impl Context { - fn set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard { - let rng_seed = handle.seed_generator().next_seed(); - - let old_handle = self.handle.borrow_mut().replace(handle.clone()); - let old_seed = self.rng.replace_seed(rng_seed); - - SetCurrentGuard { - old_handle, - old_seed, - } - } - } - - impl Drop for SetCurrentGuard { - fn drop(&mut self) { - CONTEXT.with(|ctx| { - *ctx.handle.borrow_mut() = self.old_handle.take(); - ctx.rng.replace_seed(self.old_seed.clone()); - }); - } + pub(super) fn set_scheduler<R>(v: &scheduler::Context, f: impl FnOnce() -> R) -> R { + CONTEXT.with(|c| c.scheduler.set(v, f)) } - 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); - - if self.is_root { - *c.defer.borrow_mut() = None; - } - }); - } - } - - impl BlockingRegionGuard { - 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<F>(&mut self, f: F) -> Result<F::Output, AccessError> - 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<F>(&mut self, f: F, timeout: Duration) -> Result<F::Output, ()> - 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(()); - } - - // Wake any yielded tasks before parking in order to avoid - // blocking. - with_defer(|defer| defer.wake()); - - 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, - }); - } - }) - } - } - } - - impl EnterRuntime { - pub(crate) fn is_entered(self) -> bool { - matches!(self, EnterRuntime::Entered { .. }) - } - } -} - -// Forces the current "entered" state to be cleared while the closure -// is executed. -// -// # Warning -// -// This is hidden for a reason. Do not use without fully understanding -// executors. Misusing can easily cause your program to deadlock. -cfg_rt_multi_thread! { - /// Returns true if in a runtime context. - pub(crate) fn current_enter_context() -> EnterRuntime { - CONTEXT.with(|c| c.runtime.get()) + #[track_caller] + pub(super) fn with_scheduler<R>(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R { + CONTEXT.with(|c| c.scheduler.with(f)) } - pub(crate) fn exit_runtime<F: FnOnce() -> 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); - }); - } + cfg_taskdump! { + /// SAFETY: Callers of this function must ensure that trace frames always + /// form a valid linked list. + pub(crate) unsafe fn with_trace<R>(f: impl FnOnce(&trace::Context) -> R) -> Option<R> { + CONTEXT.try_with(|c| f(&c.trace)).ok() } - - 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() } } diff --git a/src/runtime/context/blocking.rs b/src/runtime/context/blocking.rs new file mode 100644 index 0000000..8ae4f57 --- /dev/null +++ b/src/runtime/context/blocking.rs @@ -0,0 +1,121 @@ +use 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<NotSendOrSync>, +} + +pub(crate) struct DisallowBlockInPlaceGuard(bool); + +pub(crate) fn try_enter_blocking_region() -> Option<BlockingRegionGuard> { + 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<F>(&mut self, f: F) -> Result<F::Output, AccessError> + 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<F>(&mut self, f: F, timeout: Duration) -> Result<F::Output, ()> + 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, + }); + } + }) + } + } +} diff --git a/src/runtime/context/current.rs b/src/runtime/context/current.rs new file mode 100644 index 0000000..c3dc5c8 --- /dev/null +++ b/src/runtime/context/current.rs @@ -0,0 +1,99 @@ +use 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<scheduler::Handle>, + + // The depth for this guard + depth: usize, + + // Don't let the type move across threads. + _p: PhantomData<SyncNotSend>, +} + +pub(super) struct HandleCell { + /// Current handle + handle: RefCell<Option<scheduler::Handle>>, + + /// Tracks the number of nested calls to `try_set_current`. + depth: Cell<usize>, +} + +/// Sets this [`Handle`] as the current active [`Handle`]. +/// +/// [`Handle`]: crate::runtime::scheduler::Handle +pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> { + CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok() +} + +pub(crate) fn with_current<F, R>(f: F) -> Result<R, TryCurrentError> +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(); + + if depth == usize::MAX { + panic!("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); + }); + } +} diff --git a/src/runtime/context/runtime.rs b/src/runtime/context/runtime.rs new file mode 100644 index 0000000..f2e2989 --- /dev/null +++ b/src/runtime/context/runtime.rs @@ -0,0 +1,99 @@ +use 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<F, R>(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 { .. }) + } +} diff --git a/src/runtime/context/runtime_mt.rs b/src/runtime/context/runtime_mt.rs new file mode 100644 index 0000000..728caea --- /dev/null +++ b/src/runtime/context/runtime_mt.rs @@ -0,0 +1,36 @@ +use 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<F: FnOnce() -> 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() +} diff --git a/src/runtime/context/scoped.rs b/src/runtime/context/scoped.rs new file mode 100644 index 0000000..7b202a1 --- /dev/null +++ b/src/runtime/context/scoped.rs @@ -0,0 +1,56 @@ +use std::cell::Cell; +use std::ptr; + +/// Scoped thread-local storage +pub(super) struct Scoped<T> { + pub(super) inner: Cell<*const T>, +} + +impl<T> Scoped<T> { + pub(super) const fn new() -> Scoped<T> { + Scoped { + inner: Cell::new(ptr::null()), + } + } + + /// Inserts a value into the scoped cell for the duration of the closure + pub(super) fn set<F, R>(&self, t: &T, f: F) -> R + where + F: FnOnce() -> R, + { + struct Reset<'a, T> { + cell: &'a Cell<*const T>, + prev: *const T, + } + + impl<T> 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<F, R>(&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)) } + } + } +} diff --git a/src/runtime/coop.rs b/src/runtime/coop.rs index 0ba137a..15a4d98 100644 --- a/src/runtime/coop.rs +++ b/src/runtime/coop.rs @@ -36,6 +36,11 @@ use crate::runtime::context; #[derive(Debug, Copy, Clone)] pub(crate) struct Budget(Option<u8>); +pub(crate) struct BudgetDecrement { + success: bool, + hit_zero: bool, +} + impl Budget { /// Budget assigned to a task on each poll. /// @@ -164,7 +169,7 @@ cfg_coop! { /// that the budget empties appropriately. /// /// Note that `RestoreOnPending` restores the budget **as it was before `poll_proceed`**. - /// Therefore, if the budget is _fCURRENT.withurther_ adjusted between when `poll_proceed` returns and + /// 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] @@ -172,9 +177,17 @@ cfg_coop! { context::budget(|cell| { let mut budget = cell.get(); - if budget.decrement() { + 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(); @@ -183,19 +196,43 @@ cfg_coop! { }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained())))) } + cfg_rt! { + cfg_metrics! { + #[inline(always)] + fn inc_budget_forced_yield_count() { + let _ = context::with_current(|handle| { + handle.scheduler_metrics().inc_budget_forced_yield_count(); + }); + } + } + + cfg_not_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) -> bool { + fn decrement(&mut self) -> BudgetDecrement { if let Some(num) = &mut self.0 { if *num > 0 { *num -= 1; - true + + let hit_zero = *num == 0; + + BudgetDecrement { success: true, hit_zero } } else { - false + BudgetDecrement { success: false, hit_zero: false } } } else { - true + BudgetDecrement { success: true, hit_zero: false } } } @@ -209,7 +246,7 @@ cfg_coop! { mod test { use super::*; - #[cfg(tokio_wasm_not_wasi)] + #[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; fn get() -> Budget { diff --git a/src/runtime/defer.rs b/src/runtime/defer.rs deleted file mode 100644 index 4078512..0000000 --- a/src/runtime/defer.rs +++ /dev/null @@ -1,27 +0,0 @@ -use std::task::Waker; - -pub(crate) struct Defer { - deferred: Vec<Waker>, -} - -impl Defer { - pub(crate) fn new() -> Defer { - Defer { - deferred: Default::default(), - } - } - - pub(crate) fn defer(&mut self, waker: Waker) { - self.deferred.push(waker); - } - - pub(crate) fn is_empty(&self) -> bool { - self.deferred.is_empty() - } - - pub(crate) fn wake(&mut self) { - for waker in self.deferred.drain(..) { - waker.wake(); - } - } -} diff --git a/src/runtime/driver.rs b/src/runtime/driver.rs index 4fb6b87..0474c2b 100644 --- a/src/runtime/driver.rs +++ b/src/runtime/driver.rs @@ -2,7 +2,10 @@ // 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(not(feature = "full"), allow(dead_code))] +#![cfg_attr( + any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"), + allow(dead_code) +)] use crate::runtime::park::{ParkThread, UnparkThread}; @@ -45,8 +48,7 @@ impl Driver { 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.clone()); + let (time_driver, time_handle) = create_time_driver(cfg.enable_time, io_stack, &clock); Ok(( Self { inner: time_driver }, @@ -59,6 +61,10 @@ impl Driver { )) } + pub(crate) fn is_enabled(&self) -> bool { + self.inner.is_enabled() + } + pub(crate) fn park(&mut self, handle: &Handle) { self.inner.park(handle) } @@ -111,10 +117,8 @@ impl Handle { .expect("A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.") } - cfg_test_util! { - pub(crate) fn clock(&self) -> &Clock { - &self.clock - } + pub(crate) fn clock(&self) -> &Clock { + &self.clock } } } @@ -157,6 +161,13 @@ cfg_io_driver! { } 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), @@ -220,6 +231,11 @@ cfg_not_io_driver! { 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 + } } } @@ -289,7 +305,7 @@ cfg_time! { fn create_time_driver( enable: bool, io_stack: IoStack, - clock: Clock, + clock: &Clock, ) -> (TimeDriver, TimeHandle) { if enable { let (driver, handle) = crate::runtime::time::Driver::new(io_stack, clock); @@ -301,6 +317,13 @@ cfg_time! { } 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), @@ -337,7 +360,7 @@ cfg_not_time! { fn create_time_driver( _enable: bool, io_stack: IoStack, - _clock: Clock, + _clock: &Clock, ) -> (TimeDriver, TimeHandle) { (io_stack, ()) } diff --git a/src/runtime/dump.rs b/src/runtime/dump.rs new file mode 100644 index 0000000..994b7f9 --- /dev/null +++ b/src/runtime/dump.rs @@ -0,0 +1,76 @@ +//! Snapshots of runtime state. +//! +//! See [Handle::dump][crate::runtime::Handle::dump]. + +use std::fmt; + +/// 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<Task>, +} + +/// A snapshot of a task. +/// +/// See [Handle::dump][crate::runtime::Handle::dump]. +#[derive(Debug)] +pub struct Task { + trace: Trace, +} + +/// An execution trace of a task's last poll. +/// +/// See [Handle::dump][crate::runtime::Handle::dump]. +#[derive(Debug)] +pub struct Trace { + inner: super::task::trace::Trace, +} + +impl Dump { + pub(crate) fn new(tasks: Vec<Task>) -> 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<Item = &Task> { + self.tasks.iter() + } +} + +impl Task { + pub(crate) fn new(trace: super::task::trace::Trace) -> Self { + Self { + trace: Trace { inner: trace }, + } + } + + /// 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) + } +} diff --git a/src/runtime/handle.rs b/src/runtime/handle.rs index c5dc65f..ebc7bba 100644 --- a/src/runtime/handle.rs +++ b/src/runtime/handle.rs @@ -1,3 +1,5 @@ +#[cfg(tokio_unstable)] +use crate::runtime; use crate::runtime::{context, scheduler, RuntimeFlavor}; /// Handle to the runtime. @@ -35,9 +37,44 @@ pub struct EnterGuard<'a> { 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. + /// 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 @@ -109,7 +146,9 @@ impl Handle { /// /// Contrary to `current`, this never panics pub fn try_current() -> Result<Self, TryCurrentError> { - context::try_current().map(|inner| Handle { inner }) + context::with_current(|inner| Handle { + inner: inner.clone(), + }) } /// Spawns a future onto the Tokio runtime. @@ -152,7 +191,7 @@ impl Handle { self.spawn_named(future, None) } - /// Runs the provided function on an executor dedicated to blocking. + /// Runs the provided function on an executor dedicated to blocking /// operations. /// /// # Examples @@ -252,19 +291,24 @@ impl Handle { /// [`tokio::time`]: crate::time #[track_caller] pub fn block_on<F: Future>(&self, future: F) -> 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", None, super::task::Id::next().as_u64()); // Enter the runtime context. This sets the current driver handles and // prevents blocking an existing runtime. - let mut enter = context::enter_runtime(&self.inner, true); - - // Block on the future - enter - .blocking - .block_on(future) - .expect("failed to park thread") + context::enter_runtime(&self.inner, true, |blocking| { + blocking.block_on(future).expect("failed to park thread") + }) } #[track_caller] @@ -274,6 +318,14 @@ impl Handle { 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", _name, id.as_u64()); self.inner.spawn(future, id) @@ -303,8 +355,42 @@ impl Handle { pub fn runtime_flavor(&self) -> RuntimeFlavor { match self.inner { scheduler::Handle::CurrentThread(_) => RuntimeFlavor::CurrentThread, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] scheduler::Handle::MultiThread(_) => RuntimeFlavor::MultiThread, + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + 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(all(feature = "rt-multi-thread", not(target_os = "wasi")))] + scheduler::Handle::MultiThread(handle) => handle.owned_id(), + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + scheduler::Handle::MultiThreadAlt(handle) => handle.owned_id(), + }; + owned_id.into() } } } @@ -321,6 +407,163 @@ cfg_metrics! { } } +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 + /// └╼ <tokio::util::trace::InstrumentedAsyncOp<F> 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<T>::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 + /// └╼ <tokio::sync::notify::Notified as core::future::future::Future>::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"), + } + } + } + + cfg_rt_multi_thread! { + /// Spawn a new thread and asynchronously await on its result. + async fn spawn_thread<F>(f: F) -> <F as Future>::Output + where + F: Future + Send + 'static, + <F as Future>::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 { diff --git a/src/runtime/id.rs b/src/runtime/id.rs new file mode 100644 index 0000000..58551d4 --- /dev/null +++ b/src/runtime/id.rs @@ -0,0 +1,46 @@ +use std::fmt; +use std::num::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<NonZeroU64> for Id { + fn from(value: NonZeroU64) -> Self { + Id(value) + } +} + +impl fmt::Display for Id { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} diff --git a/src/runtime/io/driver.rs b/src/runtime/io/driver.rs new file mode 100644 index 0000000..755cb9d --- /dev/null +++ b/src/runtime/io/driver.rs @@ -0,0 +1,280 @@ +// 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 { + /// Tracks the number of times `turn` is called. It is safe for this to wrap + /// as it is mostly used to determine when to call `compact()`. + tick: u8, + + /// 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<registration_set::Synced>, + + /// 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(u8), + Clear(u8), +} + +const TOKEN_WAKEUP: mio::Token = mio::Token(0); +const TOKEN_SIGNAL: mio::Token = mio::Token(1); + +fn _assert_kinds() { + fn _assert<T: Send + Sync>() {} + + _assert::<Handle>(); +} + +// ===== 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 { + tick: 0, + 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<Duration>) { + debug_assert!(!handle.registrations.is_shutdown(&handle.synced.lock())); + + self.tick = self.tick.wrapping_add(1); + + 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); + // Use std::ptr::from_exposed_addr when stable + let ptr: *const ScheduledIo = token.0 as *const _; + + // 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<ScheduledIo>` so we can safely cast this to a ref. + let io: &ScheduledIo = unsafe { &*ptr }; + + io.set_readiness(Tick::Set(self.tick), |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<Arc<ScheduledIo>> { + let scheduled_io = self.registrations.allocate(&mut self.synced.lock())?; + let token = scheduled_io.token(); + + // TODO: if this returns an err, the `ScheduledIo` leaks... + self.registry.register(source, token, interest.to_mio())?; + + // 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<ScheduledIo>, + 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, + } + } +} diff --git a/src/runtime/io/driver/signal.rs b/src/runtime/io/driver/signal.rs new file mode 100644 index 0000000..ea3ef07 --- /dev/null +++ b/src/runtime/io/driver/signal.rs @@ -0,0 +1,22 @@ +use 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 + } +} diff --git a/src/runtime/io/mod.rs b/src/runtime/io/mod.rs index 2e578b6..e7c721d 100644 --- a/src/runtime/io/mod.rs +++ b/src/runtime/io/mod.rs @@ -1,344 +1,16 @@ #![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 crate::io::interest::Interest; -use crate::io::ready::Ready; -use crate::runtime::driver; -use crate::util::slab::{self, Slab}; -use crate::{loom::sync::RwLock, util::bit}; - use metrics::IoDriverMetrics; - -use std::fmt; -use std::io; -use std::time::Duration; - -/// I/O driver, backed by Mio. -pub(crate) struct Driver { - /// Tracks the number of times `turn` is called. It is safe for this to wrap - /// as it is mostly used to determine when to call `compact()`. - tick: u8, - - /// 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, - - /// Primary slab handle containing the state for each resource registered - /// with this driver. - resources: Slab<ScheduledIo>, - - /// 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, - - /// Allocates `ScheduledIo` handles when creating new resources. - io_dispatch: RwLock<IoDispatcher>, - - /// Used to wake up the reactor from a call to `turn`. - /// Not supported on Wasi due to lack of threading support. - #[cfg(not(tokio_wasi))] - waker: mio::Waker, - - pub(crate) metrics: IoDriverMetrics, -} - -#[derive(Debug)] -pub(crate) struct ReadyEvent { - tick: u8, - pub(crate) ready: Ready, - is_shutdown: bool, -} - -struct IoDispatcher { - allocator: slab::Allocator<ScheduledIo>, - is_shutdown: bool, -} - -#[derive(Debug, Eq, PartialEq, Clone, Copy)] -enum Direction { - Read, - Write, -} - -enum Tick { - Set(u8), - Clear(u8), -} - -// TODO: Don't use a fake token. Instead, reserve a slot entry for the wakeup -// token. -const TOKEN_WAKEUP: mio::Token = mio::Token(1 << 31); -const TOKEN_SIGNAL: mio::Token = mio::Token(1 + (1 << 31)); - -const ADDRESS: bit::Pack = bit::Pack::least_significant(24); - -// Packs the generation value in the `readiness` field. -// -// The generation prevents a race condition where a slab slot is reused for a -// new socket while the I/O driver is about to apply a readiness event. The -// generation value is checked when setting new readiness. If the generation do -// not match, then the readiness event is discarded. -const GENERATION: bit::Pack = ADDRESS.then(7); - -fn _assert_kinds() { - fn _assert<T: Send + Sync>() {} - - _assert::<Handle>(); -} - -// ===== 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(tokio_wasi))] - let waker = mio::Waker::new(poll.registry(), TOKEN_WAKEUP)?; - let registry = poll.registry().try_clone()?; - - let slab = Slab::new(); - let allocator = slab.allocator(); - - let driver = Driver { - tick: 0, - signal_ready: false, - events: mio::Events::with_capacity(nevents), - poll, - resources: slab, - }; - - let handle = Handle { - registry, - io_dispatch: RwLock::new(IoDispatcher::new(allocator)), - #[cfg(not(tokio_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(); - - if handle.shutdown() { - self.resources.for_each(|io| { - // If a task is waiting on the I/O resource, notify it that the - // runtime is being shutdown. And shutdown will clear all wakers. - io.shutdown(); - }); - } - } - - fn turn(&mut self, handle: &Handle, max_wait: Option<Duration>) { - // How often to call `compact()` on the resource slab - const COMPACT_INTERVAL: u8 = 255; - - self.tick = self.tick.wrapping_add(1); - - if self.tick == COMPACT_INTERVAL { - self.resources.compact() - } - - 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(tokio_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 { - Self::dispatch( - &mut self.resources, - self.tick, - token, - Ready::from_mio(event), - ); - ready_count += 1; - } - } - - handle.metrics.incr_ready_count_by(ready_count); - } - - fn dispatch(resources: &mut Slab<ScheduledIo>, tick: u8, token: mio::Token, ready: Ready) { - let addr = slab::Address::from_usize(ADDRESS.unpack(token.0)); - - let io = match resources.get(addr) { - Some(io) => io, - None => return, - }; - - let res = io.set_readiness(Some(token.0), Tick::Set(tick), |curr| curr | ready); - - if res.is_err() { - // token no longer valid! - return; - } - - io.wake(ready); - } -} - -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(tokio_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<slab::Ref<ScheduledIo>> { - let (address, shared) = self.allocate()?; - - let token = GENERATION.pack(shared.generation(), ADDRESS.pack(address.as_usize(), 0)); - - self.registry - .register(source, mio::Token(token), interest.to_mio())?; - - self.metrics.incr_fd_count(); - - Ok(shared) - } - - /// Deregisters an I/O resource from the reactor. - pub(super) fn deregister_source(&self, source: &mut impl mio::event::Source) -> io::Result<()> { - self.registry.deregister(source)?; - - self.metrics.dec_fd_count(); - - Ok(()) - } - - /// shutdown the dispatcher. - fn shutdown(&self) -> bool { - let mut io = self.io_dispatch.write().unwrap(); - if io.is_shutdown { - return false; - } - io.is_shutdown = true; - true - } - - fn allocate(&self) -> io::Result<(slab::Address, slab::Ref<ScheduledIo>)> { - let io = self.io_dispatch.read().unwrap(); - if io.is_shutdown { - return Err(io::Error::new( - io::ErrorKind::Other, - crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR, - )); - } - io.allocator.allocate().ok_or_else(|| { - io::Error::new( - io::ErrorKind::Other, - "reactor at max registered I/O resources", - ) - }) - } -} - -impl fmt::Debug for Handle { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Handle") - } -} - -// ===== impl IoDispatcher ===== - -impl IoDispatcher { - fn new(allocator: slab::Allocator<ScheduledIo>) -> Self { - Self { - allocator, - is_shutdown: false, - } - } -} - -impl Direction { - pub(super) fn mask(self) -> Ready { - match self { - Direction::Read => Ready::READABLE | Ready::READ_CLOSED, - Direction::Write => Ready::WRITABLE | Ready::WRITE_CLOSED, - } - } -} - -// Signal handling -cfg_signal_internal_and_unix! { - 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 - } - } -} diff --git a/src/runtime/io/platform.rs b/src/runtime/io/platform.rs deleted file mode 100644 index 6b27988..0000000 --- a/src/runtime/io/platform.rs +++ /dev/null @@ -1,44 +0,0 @@ -pub(crate) use self::sys::*; - -#[cfg(unix)] -mod sys { - use mio::unix::UnixReady; - use mio::Ready; - - pub(crate) fn hup() -> Ready { - UnixReady::hup().into() - } - - pub(crate) fn is_hup(ready: Ready) -> bool { - UnixReady::from(ready).is_hup() - } - - pub(crate) fn error() -> Ready { - UnixReady::error().into() - } - - pub(crate) fn is_error(ready: Ready) -> bool { - UnixReady::from(ready).is_error() - } -} - -#[cfg(windows)] -mod sys { - use mio::Ready; - - pub(crate) fn hup() -> Ready { - Ready::empty() - } - - pub(crate) fn is_hup(_: Ready) -> bool { - false - } - - pub(crate) fn error() -> Ready { - Ready::empty() - } - - pub(crate) fn is_error(_: Ready) -> bool { - false - } -} diff --git a/src/runtime/io/registration.rs b/src/runtime/io/registration.rs index 140b924..7595898 100644 --- a/src/runtime/io/registration.rs +++ b/src/runtime/io/registration.rs @@ -3,10 +3,10 @@ use crate::io::interest::Interest; use crate::runtime::io::{Direction, Handle, ReadyEvent, ScheduledIo}; use crate::runtime::scheduler; -use crate::util::slab; use mio::event::Source; use std::io; +use std::sync::Arc; use std::task::{Context, Poll}; cfg_io_driver! { @@ -45,10 +45,12 @@ cfg_io_driver! { #[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: slab::Ref<ScheduledIo>, + shared: Arc<ScheduledIo>, } } @@ -95,7 +97,7 @@ impl 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(io) + self.handle().deregister_source(&self.shared, io) } pub(crate) fn clear_readiness(&self, event: ReadyEvent) { @@ -116,7 +118,7 @@ impl Registration { // 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(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] pub(crate) fn poll_read_io<R>( &self, cx: &mut Context<'_>, @@ -144,6 +146,7 @@ impl Registration { cx: &mut Context<'_>, direction: Direction, ) -> Poll<io::Result<ReadyEvent>> { + 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)); @@ -198,6 +201,33 @@ impl Registration { } } + pub(crate) async fn readiness(&self, interest: Interest) -> io::Result<ReadyEvent> { + let ev = self.shared.readiness(interest).await; + + if ev.is_shutdown { + return Err(gone()); + } + + Ok(ev) + } + + pub(crate) async fn async_io<R>( + &self, + interest: Interest, + mut f: impl FnMut() -> io::Result<R>, + ) -> io::Result<R> { + loop { + let event = self.readiness(interest).await?; + + match f() { + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + self.clear_readiness(event); + } + x => return x, + } + } + } + fn handle(&self) -> &Handle { self.handle.driver().io() } @@ -222,30 +252,3 @@ fn gone() -> io::Error { crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR, ) } - -cfg_io_readiness! { - impl Registration { - pub(crate) async fn readiness(&self, interest: Interest) -> io::Result<ReadyEvent> { - let ev = self.shared.readiness(interest).await; - - if ev.is_shutdown { - return Err(gone()) - } - - Ok(ev) - } - - pub(crate) async fn async_io<R>(&self, interest: Interest, mut f: impl FnMut() -> io::Result<R>) -> io::Result<R> { - loop { - let event = self.readiness(interest).await?; - - match f() { - Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { - self.clear_readiness(event); - } - x => return x, - } - } - } - } -} diff --git a/src/runtime/io/registration_set.rs b/src/runtime/io/registration_set.rs new file mode 100644 index 0000000..028eb2e --- /dev/null +++ b/src/runtime/io/registration_set.rs @@ -0,0 +1,134 @@ +use 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<Arc<ScheduledIo>, 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<Arc<ScheduledIo>>, +} + +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<Arc<ScheduledIo>> { + 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<ScheduledIo>) -> 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<Arc<ScheduledIo>> { + 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) { + for io in synced.pending_release.drain(..) { + // safety: the registration is part of our list + let _ = unsafe { synced.registrations.remove(io.as_ref().into()) }; + } + + self.num_pending_release.store(0, Release); + } +} + +// Safety: `Arc` pins the inner data +unsafe impl linked_list::Link for Arc<ScheduledIo> { + type Handle = Arc<ScheduledIo>; + type Target = ScheduledIo; + + fn as_raw(handle: &Self::Handle) -> NonNull<ScheduledIo> { + // safety: Arc::as_ptr never returns null + unsafe { NonNull::new_unchecked(Arc::as_ptr(handle) as *mut _) } + } + + unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Arc<ScheduledIo> { + // safety: the linked list currently owns a ref count + unsafe { Arc::from_raw(ptr.as_ptr() as *const _) } + } + + unsafe fn pointers( + target: NonNull<Self::Target>, + ) -> NonNull<linked_list::Pointers<ScheduledIo>> { + NonNull::new_unchecked(target.as_ref().linked_list_pointers.get()) + } +} diff --git a/src/runtime/io/scheduled_io.rs b/src/runtime/io/scheduled_io.rs index 197a4e0..ddce4b3 100644 --- a/src/runtime/io/scheduled_io.rs +++ b/src/runtime/io/scheduled_io.rs @@ -1,43 +1,122 @@ -use super::{ReadyEvent, Tick}; use 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::slab::Entry; +use crate::util::linked_list::{self, LinkedList}; use crate::util::WakeList; -use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; +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}; -use super::Direction; - -cfg_io_readiness! { - use crate::util::linked_list::{self, LinkedList}; - - use std::cell::UnsafeCell; - use std::future::Future; - use std::marker::PhantomPinned; - use std::pin::Pin; - use std::ptr::NonNull; -} - /// 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, riscv64, 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/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_riscv64.go#L7 +// - 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 +// +// riscv32 is assumed not to exceed the cache line size of riscv64. +#[cfg_attr( + any( + target_arch = "arm", + target_arch = "mips", + target_arch = "mips64", + target_arch = "riscv32", + target_arch = "riscv64", + 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, 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 +// +// 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 = "riscv32", + target_arch = "riscv64", + target_arch = "sparc", + target_arch = "hexagon", + target_arch = "m68k", + target_arch = "s390x", + )), + repr(align(64)) +)] pub(crate) struct ScheduledIo { - /// Packs the resource's readiness with the resource's generation. + pub(super) linked_list_pointers: UnsafeCell<linked_list::Pointers<Self>>, + + /// Packs the resource's readiness and I/O driver latest tick. readiness: AtomicUsize, waiters: Mutex<Waiters>, } -cfg_io_readiness! { - type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; -} +type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; #[derive(Debug, Default)] struct Waiters { - #[cfg(feature = "net")] /// List of all current waiters. list: WaitList, @@ -48,83 +127,64 @@ struct Waiters { writer: Option<Waker>, } -cfg_io_readiness! { - #[derive(Debug)] - struct Waiter { - pointers: linked_list::Pointers<Waiter>, +#[derive(Debug)] +struct Waiter { + pointers: linked_list::Pointers<Waiter>, - /// The waker for this task. - waker: Option<Waker>, + /// The waker for this task. + waker: Option<Waker>, - /// The interest this waiter is waiting on. - interest: Interest, + /// The interest this waiter is waiting on. + interest: Interest, - is_ready: bool, + is_ready: bool, - /// Should never be `!Unpin`. - _p: PhantomPinned, - } + /// Should never be `!Unpin`. + _p: PhantomPinned, +} - generate_addr_of_methods! { - impl<> Waiter { - unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { - &self.pointers - } +generate_addr_of_methods! { + impl<> Waiter { + unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { + &self.pointers } } +} - /// Future returned by `readiness()`. - struct Readiness<'a> { - scheduled_io: &'a ScheduledIo, +/// Future returned by `readiness()`. +struct Readiness<'a> { + scheduled_io: &'a ScheduledIo, - state: State, + state: State, - /// Entry in the waiter `LinkedList`. - waiter: UnsafeCell<Waiter>, - } + /// Entry in the waiter `LinkedList`. + waiter: UnsafeCell<Waiter>, +} - enum State { - Init, - Waiting, - Done, - } +enum State { + Init, + Waiting, + Done, } // The `ScheduledIo::readiness` (`AtomicUsize`) is packed full of goodness. // -// | shutdown | generation | driver tick | readiness | -// |----------+------------+--------------+-----------| -// | 1 bit | 7 bits + 8 bits + 16 bits | +// | shutdown | driver tick | readiness | +// |----------+-------------+-----------| +// | 1 bit | 8 bits + 16 bits | const READINESS: bit::Pack = bit::Pack::least_significant(16); const TICK: bit::Pack = READINESS.then(8); -const GENERATION: bit::Pack = TICK.then(7); - -const SHUTDOWN: bit::Pack = GENERATION.then(1); - -#[test] -fn test_generations_assert_same() { - assert_eq!(super::GENERATION, GENERATION); -} +const SHUTDOWN: bit::Pack = TICK.then(1); // ===== impl ScheduledIo ===== -impl Entry for ScheduledIo { - fn reset(&self) { - let state = self.readiness.load(Acquire); - - let generation = GENERATION.unpack(state); - let next = GENERATION.pack_lossy(generation + 1, 0); - - self.readiness.store(next, Release); - } -} - impl Default for ScheduledIo { fn default() -> ScheduledIo { ScheduledIo { + linked_list_pointers: UnsafeCell::new(linked_list::Pointers::new()), readiness: AtomicUsize::new(0), waiters: Mutex::new(Default::default()), } @@ -132,8 +192,9 @@ impl Default for ScheduledIo { } impl ScheduledIo { - pub(crate) fn generation(&self) -> usize { - GENERATION.unpack(self.readiness.load(Acquire)) + pub(crate) fn token(&self) -> mio::Token { + // use `expose_addr` when stable + mio::Token(self as *const _ as usize) } /// Invoked when the IO driver is shut down; forces this ScheduledIo into a @@ -148,61 +209,39 @@ impl ScheduledIo { /// the current value, returning the previous readiness value. /// /// # Arguments - /// - `token`: the token for this `ScheduledIo`. /// - `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. - /// - /// # Returns - /// - /// If the given token's generation no longer matches the `ScheduledIo`'s - /// generation, then the corresponding IO resource has been removed and - /// replaced with a new resource. In that case, this method returns `Err`. - /// Otherwise, this returns the previous readiness. - pub(super) fn set_readiness( - &self, - token: Option<usize>, - tick: Tick, - f: impl Fn(Ready) -> Ready, - ) -> Result<(), ()> { + pub(super) fn set_readiness(&self, tick: Tick, f: impl Fn(Ready) -> Ready) { let mut current = self.readiness.load(Acquire); - loop { - let current_generation = GENERATION.unpack(current); - - if let Some(token) = token { - // Check that the generation for this access is still the - // current one. - if GENERATION.unpack(token) != current_generation { - return Err(()); - } - } + // The shutdown bit should not be set + debug_assert_eq!(0, SHUTDOWN.unpack(current)); - // Mask out the tick/generation bits so that the modifying - // function doesn't see them. + loop { + // Mask out the tick bits so that the modifying function doesn't see + // them. let current_readiness = Ready::from_usize(current); let new = f(current_readiness); - let packed = match tick { + let next = match tick { Tick::Set(t) => TICK.pack(t as usize, new.as_usize()), Tick::Clear(t) => { if TICK.unpack(current) as u8 != t { // Trying to clear readiness with an old event! - return Err(()); + return; } TICK.pack(t as usize, new.as_usize()) } }; - let next = GENERATION.pack(current_generation, packed); - match self .readiness .compare_exchange(current, next, AcqRel, Acquire) { - Ok(_) => return Ok(()), + Ok(_) => return, // we lost the race, retry! Err(actual) => current = actual, } @@ -238,7 +277,6 @@ impl ScheduledIo { } } - #[cfg(feature = "net")] 'outer: loop { let mut iter = waiters.list.drain_filter(|w| ready.satisfies(w.interest)); @@ -351,9 +389,7 @@ impl ScheduledIo { // 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; - - // result isn't important - let _ = self.set_readiness(None, Tick::Clear(event.tick), |curr| curr - mask_no_closed); + self.set_readiness(Tick::Clear(event.tick), |curr| curr - mask_no_closed); } pub(crate) fn clear_wakers(&self) { @@ -372,187 +408,193 @@ impl Drop for ScheduledIo { unsafe impl Send for ScheduledIo {} unsafe impl Sync for ScheduledIo {} -cfg_io_readiness! { - 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, - }), - } - } +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 } - unsafe impl linked_list::Link for Waiter { - type Handle = NonNull<Waiter>; - type Target = Waiter; - - fn as_raw(handle: &NonNull<Waiter>) -> NonNull<Waiter> { - *handle - } - - unsafe fn from_raw(ptr: NonNull<Waiter>) -> NonNull<Waiter> { - ptr - } - - unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { - Waiter::addr_of_pointers(target) + // 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, + }), } } +} - // ===== impl Readiness ===== - - impl Future for Readiness<'_> { - type Output = ReadyEvent; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - use std::sync::atomic::Ordering::SeqCst; +unsafe impl linked_list::Link for Waiter { + type Handle = NonNull<Waiter>; + type Target = Waiter; - let (scheduled_io, state, waiter) = unsafe { - let me = self.get_unchecked_mut(); - (&me.scheduled_io, &mut me.state, &me.waiter) - }; + fn as_raw(handle: &NonNull<Waiter>) -> NonNull<Waiter> { + *handle + } - loop { - match *state { - State::Init => { - // Optimistically check existing readiness - let curr = scheduled_io.readiness.load(SeqCst); - let ready = Ready::from_usize(READINESS.unpack(curr)); - let is_shutdown = SHUTDOWN.unpack(curr) != 0; - - // Safety: `waiter.interest` never changes - let interest = unsafe { (*waiter.get()).interest }; - 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 }); - } + unsafe fn from_raw(ptr: NonNull<Waiter>) -> NonNull<Waiter> { + ptr + } - // Wasn't ready, take the lock (and check again while locked). - let mut waiters = scheduled_io.waiters.lock(); + unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> { + Waiter::addr_of_pointers(target) + } +} - let curr = scheduled_io.readiness.load(SeqCst); - let mut ready = Ready::from_usize(READINESS.unpack(curr)); - let is_shutdown = SHUTDOWN.unpack(curr) != 0; +// ===== impl Readiness ===== - if is_shutdown { - ready = Ready::ALL; - } +impl Future for Readiness<'_> { + type Output = ReadyEvent; - let ready = ready.intersection(interest); + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + use std::sync::atomic::Ordering::SeqCst; - 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 }); - } + let (scheduled_io, state, waiter) = unsafe { + let me = self.get_unchecked_mut(); + (&me.scheduled_io, &mut me.state, &me.waiter) + }; - // Not ready even after locked, insert into list... + loop { + match *state { + State::Init => { + // Optimistically check existing readiness + let curr = scheduled_io.readiness.load(SeqCst); + let ready = Ready::from_usize(READINESS.unpack(curr)); + let is_shutdown = SHUTDOWN.unpack(curr) != 0; + + // Safety: `waiter.interest` never changes + let interest = unsafe { (*waiter.get()).interest }; + 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, + }); + } - // Safety: called while locked - unsafe { - (*waiter.get()).waker = Some(cx.waker().clone()); - } + // Wasn't ready, take the lock (and check again while locked). + let mut waiters = scheduled_io.waiters.lock(); - // 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. - if !w.waker.as_ref().unwrap().will_wake(cx.waker()) { - w.waker = Some(cx.waker().clone()); - } - - return Poll::Pending; - } + let curr = scheduled_io.readiness.load(SeqCst); + let mut ready = Ready::from_usize(READINESS.unpack(curr)); + let is_shutdown = SHUTDOWN.unpack(curr) != 0; - // 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); + if is_shutdown { + ready = Ready::ALL; } - 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; + let ready = ready.intersection(interest); - // 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`. + if !ready.is_empty() || is_shutdown { + // Currently 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 curr_ready = Ready::from_usize(READINESS.unpack(curr)); - let ready = curr_ready.intersection(w.interest); - + *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. + if !w.waker.as_ref().unwrap().will_wake(cx.waker()) { + w.waker = Some(cx.waker().clone()); + } - impl Drop for Readiness<'_> { - fn drop(&mut self) { - let mut waiters = self.scheduled_io.waiters.lock(); + return Poll::Pending; + } - // Safety: `waiter` is only ever stored in `waiters` - unsafe { - waiters - .list - .remove(NonNull::new_unchecked(self.waiter.get())) - }; + // 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 curr_ready = Ready::from_usize(READINESS.unpack(curr)); + let ready = curr_ready.intersection(w.interest); + + return Poll::Ready(ReadyEvent { + tick, + ready, + is_shutdown, + }); + } + } } } +} - unsafe impl Send for Readiness<'_> {} - unsafe impl Sync for Readiness<'_> {} +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<'_> {} diff --git a/src/runtime/metrics/batch.rs b/src/runtime/metrics/batch.rs index 4e6b28d..1d0f3de 100644 --- a/src/runtime/metrics/batch.rs +++ b/src/runtime/metrics/batch.rs @@ -1,8 +1,7 @@ -use crate::runtime::WorkerMetrics; +use crate::runtime::metrics::{HistogramBatch, WorkerMetrics}; -use std::convert::TryFrom; use std::sync::atomic::Ordering::Relaxed; -use std::time::Instant; +use std::time::{Duration, Instant}; pub(crate) struct MetricsBatch { /// Number of times the worker parked. @@ -33,11 +32,26 @@ pub(crate) struct MetricsBatch { /// The total busy duration in nanoseconds. busy_duration_total: u64, - last_resume_time: Instant, + + /// 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<PollTimer>, +} + +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() -> MetricsBatch { + pub(crate) fn new(worker_metrics: &WorkerMetrics) -> MetricsBatch { + let now = Instant::now(); + MetricsBatch { park_count: 0, noop_count: 0, @@ -48,11 +62,19 @@ impl MetricsBatch { local_schedule_count: 0, overflow_count: 0, busy_duration_total: 0, - last_resume_time: Instant::now(), + 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) { + 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.noop_count.store(self.noop_count, Relaxed); worker.steal_count.store(self.steal_count, Relaxed); @@ -69,6 +91,11 @@ impl MetricsBatch { .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. @@ -80,22 +107,38 @@ impl MetricsBatch { } else { self.poll_count_on_last_park = self.poll_count; } - - let busy_duration = self.last_resume_time.elapsed(); - let busy_duration = u64::try_from(busy_duration.as_nanos()).unwrap_or(u64::MAX); - self.busy_duration_total += busy_duration; } - pub(crate) fn returned_from_park(&mut self) { - self.last_resume_time = Instant::now(); + /// Start processing a batch of tasks + pub(crate) fn start_processing_scheduled_tasks(&mut self) { + self.processing_scheduled_tasks_started_at = Instant::now(); } - pub(crate) fn inc_local_schedule_count(&mut self) { - self.local_schedule_count += 1; + /// 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); } - pub(crate) fn incr_poll_count(&mut self) { + /// 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; } } @@ -114,3 +157,7 @@ cfg_rt_multi_thread! { } } } + +fn duration_as_u64(dur: Duration) -> u64 { + u64::try_from(dur.as_nanos()).unwrap_or(u64::MAX) +} diff --git a/src/runtime/metrics/histogram.rs b/src/runtime/metrics/histogram.rs new file mode 100644 index 0000000..976f54f --- /dev/null +++ b/src/runtime/metrics/histogram.rs @@ -0,0 +1,502 @@ +use crate::loom::sync::atomic::{AtomicU64, Ordering::Relaxed}; + +use std::cmp; +use std::ops::Range; + +#[derive(Debug)] +pub(crate) struct Histogram { + /// The histogram buckets + buckets: Box<[AtomicU64]>, + + /// Bucket scale, linear or log + scale: HistogramScale, + + /// Minimum resolution + resolution: u64, +} + +#[derive(Debug, Clone)] +pub(crate) struct HistogramBuilder { + /// Histogram scale + pub(crate) scale: HistogramScale, + + /// Must be a power of 2 + pub(crate) resolution: u64, + + /// Number of buckets + pub(crate) num_buckets: usize, +} + +#[derive(Debug)] +pub(crate) struct HistogramBatch { + buckets: Box<[u64]>, + scale: HistogramScale, + resolution: u64, +} + +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, + } +} + +impl Histogram { + pub(crate) fn num_buckets(&self) -> usize { + self.buckets.len() + } + + pub(crate) fn get(&self, bucket: usize) -> u64 { + self.buckets[bucket].load(Relaxed) + } + + pub(crate) fn bucket_range(&self, bucket: usize) -> Range<u64> { + match self.scale { + HistogramScale::Log => Range { + start: if bucket == 0 { + 0 + } else { + self.resolution << (bucket - 1) + }, + end: if bucket == self.buckets.len() - 1 { + u64::MAX + } else { + self.resolution << bucket + }, + }, + HistogramScale::Linear => Range { + start: self.resolution * bucket as u64, + end: if bucket == self.buckets.len() - 1 { + u64::MAX + } else { + self.resolution * (bucket as u64 + 1) + }, + }, + } + } +} + +impl HistogramBatch { + pub(crate) fn from_histogram(histogram: &Histogram) -> HistogramBatch { + let buckets = vec![0; histogram.buckets.len()].into_boxed_slice(); + + HistogramBatch { + buckets, + scale: histogram.scale, + resolution: histogram.resolution, + } + } + + 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.scale, histogram.scale); + debug_assert_eq!(self.resolution, histogram.resolution); + 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 { + match self.scale { + HistogramScale::Linear => { + let max = self.buckets.len() - 1; + cmp::min(value / self.resolution, max as u64) as usize + } + HistogramScale::Log => { + let max = self.buckets.len() - 1; + + if value < self.resolution { + 0 + } else { + let significant_digits = 64 - value.leading_zeros(); + let bucket_digits = 64 - (self.resolution - 1).leading_zeros(); + cmp::min(significant_digits as usize - bucket_digits as usize, max) + } + } + } + } +} + +impl HistogramBuilder { + pub(crate) fn new() -> HistogramBuilder { + HistogramBuilder { + scale: HistogramScale::Linear, + // Resolution is in nanoseconds. + resolution: 100_000, + num_buckets: 10, + } + } + + pub(crate) fn build(&self) -> Histogram { + let mut resolution = self.resolution; + + assert!(resolution > 0); + + if matches!(self.scale, HistogramScale::Log) { + resolution = resolution.next_power_of_two(); + } + + Histogram { + buckets: (0..self.num_buckets) + .map(|_| AtomicU64::new(0)) + .collect::<Vec<_>>() + .into_boxed_slice(), + resolution, + scale: self.scale, + } + } +} + +impl Default for HistogramBuilder { + fn default() -> HistogramBuilder { + HistogramBuilder::new() + } +} + +#[cfg(test)] +mod test { + use super::*; + + macro_rules! assert_bucket_eq { + ($h:expr, $bucket:expr, $val:expr) => {{ + assert_eq!($h.buckets[$bucket], $val); + }}; + } + + #[test] + fn log_scale_resolution_1() { + let h = HistogramBuilder { + scale: HistogramScale::Log, + resolution: 1, + num_buckets: 10, + } + .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); + } + + #[test] + fn log_scale_resolution_2() { + let h = HistogramBuilder { + scale: HistogramScale::Log, + resolution: 2, + num_buckets: 10, + } + .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 = HistogramBuilder { + scale: HistogramScale::Linear, + resolution: 1, + num_buckets: 10, + } + .build(); + + 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 = HistogramBuilder { + scale: HistogramScale::Linear, + resolution: 100, + num_buckets: 10, + } + .build(); + + 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 = HistogramBuilder { + scale: HistogramScale::Linear, + resolution: 100, + num_buckets: 10, + } + .build(); + + 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]); + } + } +} diff --git a/src/runtime/metrics/mock.rs b/src/runtime/metrics/mock.rs index c388dc0..e4bb3a9 100644 --- a/src/runtime/metrics/mock.rs +++ b/src/runtime/metrics/mock.rs @@ -6,6 +6,9 @@ pub(crate) struct WorkerMetrics {} pub(crate) struct MetricsBatch {} +#[derive(Clone, Default)] +pub(crate) struct HistogramBuilder {} + impl SchedulerMetrics { pub(crate) fn new() -> Self { Self {} @@ -20,19 +23,27 @@ impl WorkerMetrics { 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) {} } impl MetricsBatch { - pub(crate) fn new() -> Self { + pub(crate) fn new(_: &WorkerMetrics) -> Self { Self {} } - pub(crate) fn submit(&mut self, _to: &WorkerMetrics) {} + pub(crate) fn submit(&mut self, _to: &WorkerMetrics, _mean_poll_time: u64) {} pub(crate) fn about_to_park(&mut self) {} - pub(crate) fn returned_from_park(&mut self) {} - pub(crate) fn incr_poll_count(&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! { diff --git a/src/runtime/metrics/mod.rs b/src/runtime/metrics/mod.rs index 4b96f1b..88be4a5 100644 --- a/src/runtime/metrics/mod.rs +++ b/src/runtime/metrics/mod.rs @@ -12,6 +12,11 @@ cfg_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; + mod runtime; #[allow(unreachable_pub)] // rust-lang/rust#57411 pub use runtime::RuntimeMetrics; @@ -31,5 +36,5 @@ cfg_metrics! { cfg_not_metrics! { mod mock; - pub(crate) use mock::{SchedulerMetrics, WorkerMetrics, MetricsBatch}; + pub(crate) use mock::{SchedulerMetrics, WorkerMetrics, MetricsBatch, HistogramBuilder}; } diff --git a/src/runtime/metrics/runtime.rs b/src/runtime/metrics/runtime.rs index d29cb3d..95b517c 100644 --- a/src/runtime/metrics/runtime.rs +++ b/src/runtime/metrics/runtime.rs @@ -1,5 +1,6 @@ use crate::runtime::Handle; +use std::ops::Range; use std::sync::atomic::Ordering::Relaxed; use std::time::Duration; @@ -68,7 +69,26 @@ impl RuntimeMetrics { self.handle.inner.num_blocking_threads() } - /// Returns the number of idle threads, which hve spawned by the runtime + /// Returns the number of active tasks in the runtime. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main] + /// async fn main() { + /// let metrics = Handle::current().metrics(); + /// + /// let n = metrics.active_tasks_count(); + /// println!("Runtime has {} active tasks", n); + /// } + /// ``` + pub fn active_tasks_count(&self) -> usize { + self.handle.inner.active_tasks_count() + } + + /// Returns the number of idle threads, which have spawned by the runtime /// for `spawn_blocking` calls. /// /// # Examples @@ -124,6 +144,21 @@ impl RuntimeMetrics { .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 @@ -544,6 +579,237 @@ impl RuntimeMetrics { 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_count_histogram()`] when building the + /// runtime. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let enabled = metrics.poll_count_histogram_enabled(); + /// + /// println!("Tracking task poll time distribution: {:?}", enabled); + /// }); + /// } + /// ``` + /// + /// [`enable_metrics_poll_count_histogram()`]: crate::runtime::Builder::enable_metrics_poll_count_histogram + /// [`Instant::now()`]: std::time::Instant::now + pub fn poll_count_histogram_enabled(&self) -> bool { + self.handle + .inner + .worker_metrics(0) + .poll_count_histogram + .is_some() + } + + /// Returns the number of histogram buckets tracking the distribution of + /// task poll times. + /// + /// This value is configured by calling + /// [`metrics_poll_count_histogram_buckets()`] when building the runtime. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let buckets = metrics.poll_count_histogram_num_buckets(); + /// + /// println!("Histogram buckets: {:?}", buckets); + /// }); + /// } + /// ``` + /// + /// [`metrics_poll_count_histogram_buckets()`]: + /// crate::runtime::Builder::metrics_poll_count_histogram_buckets + pub fn poll_count_histogram_num_buckets(&self) -> usize { + self.handle + .inner + .worker_metrics(0) + .poll_count_histogram + .as_ref() + .map(|histogram| histogram.num_buckets()) + .unwrap_or_default() + } + + /// Returns the range of task poll times tracked by the given bucket. + /// + /// This value is configured by calling + /// [`metrics_poll_count_histogram_resolution()`] when building the runtime. + /// + /// # Panics + /// + /// The method panics if `bucket` represents an invalid bucket index, i.e. + /// is greater than or equal to `poll_count_histogram_num_buckets()`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::{self, Handle}; + /// + /// fn main() { + /// runtime::Builder::new_current_thread() + /// .enable_metrics_poll_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let buckets = metrics.poll_count_histogram_num_buckets(); + /// + /// for i in 0..buckets { + /// let range = metrics.poll_count_histogram_bucket_range(i); + /// println!("Histogram bucket {} range: {:?}", i, range); + /// } + /// }); + /// } + /// ``` + /// + /// [`metrics_poll_count_histogram_resolution()`]: + /// crate::runtime::Builder::metrics_poll_count_histogram_resolution + #[track_caller] + pub fn poll_count_histogram_bucket_range(&self, bucket: usize) -> Range<Duration> { + 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() + } + + /// 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_count_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_count_histogram() + /// .build() + /// .unwrap() + /// .block_on(async { + /// let metrics = Handle::current().metrics(); + /// let buckets = metrics.poll_count_histogram_num_buckets(); + /// + /// for worker in 0..metrics.num_workers() { + /// for i in 0..buckets { + /// let count = metrics.poll_count_histogram_bucket_count(worker, i); + /// println!("Poll count {}", count); + /// } + /// } + /// }); + /// } + /// ``` + /// + /// [`poll_count_histogram_bucket_range()`]: crate::runtime::RuntimeMetrics::poll_count_histogram_bucket_range + #[track_caller] + pub fn poll_count_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() + } + + /// 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`. /// diff --git a/src/runtime/metrics/scheduler.rs b/src/runtime/metrics/scheduler.rs index d1ba3b6..d9f8edf 100644 --- a/src/runtime/metrics/scheduler.rs +++ b/src/runtime/metrics/scheduler.rs @@ -11,12 +11,14 @@ use crate::loom::sync::atomic::{AtomicU64, Ordering::Relaxed}; pub(crate) struct SchedulerMetrics { /// Number of tasks that are scheduled from outside the runtime. pub(super) remote_schedule_count: AtomicU64, + pub(super) budget_forced_yield_count: AtomicU64, } impl SchedulerMetrics { pub(crate) fn new() -> SchedulerMetrics { SchedulerMetrics { remote_schedule_count: AtomicU64::new(0), + budget_forced_yield_count: AtomicU64::new(0), } } @@ -24,4 +26,9 @@ impl SchedulerMetrics { pub(crate) fn inc_remote_schedule_count(&self) { self.remote_schedule_count.fetch_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.fetch_add(1, Relaxed); + } } diff --git a/src/runtime/metrics/worker.rs b/src/runtime/metrics/worker.rs index a40c76e..b53a86b 100644 --- a/src/runtime/metrics/worker.rs +++ b/src/runtime/metrics/worker.rs @@ -1,5 +1,7 @@ use crate::loom::sync::atomic::Ordering::Relaxed; use crate::loom::sync::atomic::{AtomicU64, AtomicUsize}; +use crate::runtime::metrics::Histogram; +use crate::runtime::Config; /// Retrieve runtime worker metrics. /// @@ -26,6 +28,9 @@ pub(crate) struct WorkerMetrics { /// Number of tasks the worker polled. pub(crate) poll_count: AtomicU64, + /// EWMA task poll time, in nanoseconds. + pub(crate) mean_poll_time: AtomicU64, + /// Amount of time the worker spent doing work vs. parking. pub(crate) busy_duration_total: AtomicU64, @@ -38,9 +43,21 @@ pub(crate) struct WorkerMetrics { /// Number of tasks currently in the local queue. Used only by the /// current-thread scheduler. pub(crate) queue_depth: AtomicUsize, + + /// If `Some`, tracks the the number of polls by duration range. + pub(super) poll_count_histogram: Option<Histogram>, } 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 { park_count: AtomicU64::new(0), @@ -48,10 +65,12 @@ impl WorkerMetrics { steal_count: AtomicU64::new(0), steal_operations: AtomicU64::new(0), poll_count: AtomicU64::new(0), + mean_poll_time: AtomicU64::new(0), overflow_count: AtomicU64::new(0), busy_duration_total: AtomicU64::new(0), local_schedule_count: AtomicU64::new(0), queue_depth: AtomicUsize::new(0), + poll_count_histogram: None, } } diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index b6f43ea..e3369cb 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -156,9 +156,11 @@ //! multi-thread scheduler spawns threads to schedule tasks and for `spawn_blocking` //! calls. //! -//! While the `Runtime` is active, threads may shutdown after periods of being -//! idle. Once `Runtime` is dropped, all runtime threads are forcibly shutdown. -//! Any tasks that have not yet completed will be dropped. +//! 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 @@ -173,7 +175,7 @@ // At the top due to macros #[cfg(test)] -#[cfg(not(tokio_wasm))] +#[cfg(not(target_family = "wasm"))] #[macro_use] mod tests; @@ -210,7 +212,7 @@ cfg_rt! { use config::Config; mod blocking; - #[cfg_attr(tokio_wasi, allow(unused_imports))] + #[cfg_attr(target_os = "wasi", allow(unused_imports))] pub(crate) use blocking::spawn_blocking; cfg_trace! { @@ -224,12 +226,18 @@ cfg_rt! { 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 defer; - pub(crate) use defer::Defer; + cfg_taskdump! { + pub mod dump; + pub use dump::Dump; + } mod handle; pub use handle::{EnterGuard, Handle, TryCurrentError}; @@ -242,9 +250,9 @@ cfg_rt! { cfg_metrics! { mod metrics; - pub use metrics::RuntimeMetrics; + pub use metrics::{RuntimeMetrics, HistogramScale}; - pub(crate) use metrics::{MetricsBatch, SchedulerMetrics, WorkerMetrics}; + pub(crate) use metrics::{MetricsBatch, SchedulerMetrics, WorkerMetrics, HistogramBuilder}; cfg_net! { pub(crate) use metrics::IoDriverMetrics; @@ -253,7 +261,7 @@ cfg_rt! { cfg_not_metrics! { pub(crate) mod metrics; - pub(crate) use metrics::{SchedulerMetrics, WorkerMetrics, MetricsBatch}; + pub(crate) use metrics::{SchedulerMetrics, WorkerMetrics, MetricsBatch, HistogramBuilder}; } /// After thread starts / before thread stops diff --git a/src/runtime/park.rs b/src/runtime/park.rs index dc86f42..3e43745 100644 --- a/src/runtime/park.rs +++ b/src/runtime/park.rs @@ -67,9 +67,9 @@ impl ParkThread { CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst)); // Wasm doesn't have threads, so just sleep. - #[cfg(not(tokio_wasm))] + #[cfg(not(target_family = "wasm"))] self.inner.park_timeout(duration); - #[cfg(tokio_wasm)] + #[cfg(target_family = "wasm")] std::thread::sleep(duration); } @@ -222,7 +222,6 @@ impl UnparkThread { use crate::loom::thread::AccessError; use std::future::Future; use std::marker::PhantomData; -use std::mem; use std::rc::Rc; use std::task::{RawWaker, RawWakerVTable, Waker}; @@ -284,11 +283,6 @@ impl CachedParkThread { return Ok(v); } - // Wake any yielded tasks before parking in order to avoid - // blocking. - #[cfg(feature = "rt")] - crate::runtime::context::with_defer(|defer| defer.wake()); - self.park(); } } @@ -322,16 +316,12 @@ unsafe fn unparker_to_raw_waker(unparker: Arc<Inner>) -> RawWaker { } unsafe fn clone(raw: *const ()) -> RawWaker { - let unparker = Inner::from_raw(raw); - - // Increment the ref count - mem::forget(unparker.clone()); - - unparker_to_raw_waker(unparker) + Arc::increment_strong_count(raw as *const Inner); + unparker_to_raw_waker(Inner::from_raw(raw)) } unsafe fn drop_waker(raw: *const ()) { - let _ = Inner::from_raw(raw); + drop(Inner::from_raw(raw)); } unsafe fn wake(raw: *const ()) { @@ -340,11 +330,8 @@ unsafe fn wake(raw: *const ()) { } unsafe fn wake_by_ref(raw: *const ()) { - let unparker = Inner::from_raw(raw); - unparker.unpark(); - - // We don't actually own a reference to the unparker - mem::forget(unparker); + let raw = raw as *const Inner; + (*raw).unpark(); } #[cfg(loom)] diff --git a/src/runtime/runtime.rs b/src/runtime/runtime.rs index 1985673..5d79685 100644 --- a/src/runtime/runtime.rs +++ b/src/runtime/runtime.rs @@ -9,6 +9,10 @@ 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. @@ -24,31 +28,69 @@ cfg_rt_multi_thread! { /// /// # Shutdown /// -/// Shutting down the runtime is done by dropping the value. The current -/// thread will block until the shut down operation has completed. +/// 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. /// -/// * Drain any scheduled work queues. -/// * Drop any futures that have not yet completed. -/// * Drop the reactor. +/// Blocking functions spawned through [`Runtime::spawn_blocking`] keep running +/// until they return. /// -/// Once the reactor has dropped, any outstanding I/O resources bound to -/// that reactor will no longer function. Calling any method on them will -/// result in an error. +/// 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 /// -/// The Tokio runtime implements `Sync` and `Send` to allow you to wrap it -/// in a `Arc`. Most fn take `&self` to allow you to call them concurrently -/// across multiple threads. +/// There are several ways to establish shared access to a Tokio runtime: +/// +/// * Using an <code>[Arc]\<Runtime></code>. +/// * Using a [`Handle`]. +/// * Entering the runtime context. /// -/// Calls to `shutdown` and `shutdown_timeout` require exclusive ownership of -/// the runtime type and this can be achieved via `Arc::try_unwrap` when only -/// one strong count reference is left over. +/// Using an <code>[Arc]\<Runtime></code> 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 <code>[Arc]\<Runtime></code> and [`Handle`] is that +/// an <code>[Arc]\<Runtime></code> 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 <code>[Arc]\<Runtime></code>, +/// 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 +/// [`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 @@ -71,6 +113,9 @@ pub enum RuntimeFlavor { CurrentThread, /// The flavor that executes tasks across multiple threads. MultiThread, + /// The flavor that executes tasks across multiple threads. + #[cfg(tokio_unstable)] + MultiThreadAlt, } /// The runtime scheduler is either a multi-thread or a current-thread executor. @@ -80,8 +125,12 @@ pub(super) enum Scheduler { CurrentThread(CurrentThread), /// Execute tasks across multiple threads. - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] MultiThread(MultiThread), + + /// Execute tasks across multiple threads. + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + MultiThreadAlt(MultiThreadAlt), } impl Runtime { @@ -227,6 +276,13 @@ impl Runtime { /// 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 @@ -268,6 +324,15 @@ impl Runtime { /// [handle]: fn@Handle::block_on #[track_caller] pub fn block_on<F: Future>(&self, future: F) -> 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, @@ -280,8 +345,10 @@ impl Runtime { match &self.scheduler { Scheduler::CurrentThread(exec) => exec.block_on(&self.handle.inner, future), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Scheduler::MultiThread(exec) => exec.block_on(&self.handle.inner, future), + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Scheduler::MultiThreadAlt(exec) => exec.block_on(&self.handle.inner, future), } } @@ -322,18 +389,9 @@ impl Runtime { } /// Shuts down the runtime, waiting for at most `duration` for all spawned - /// task to shutdown. - /// - /// Usually, dropping a `Runtime` handle is sufficient as tasks are able to - /// shutdown in a timely fashion. However, dropping a `Runtime` will wait - /// indefinitely for all tasks to terminate, and there are cases where a long - /// blocking task has been spawned, which can block dropping `Runtime`. + /// work to stop. /// - /// In this case, calling `shutdown_timeout` with an explicit wait timeout - /// can work. The `shutdown_timeout` will signal all tasks to shutdown and - /// will wait for at most `duration` for all spawned tasks to terminate. If - /// `timeout` elapses before all tasks are dropped, the function returns and - /// outstanding tasks are potentially leaked. + /// See the [struct level documentation](Runtime#shutdown) for more details. /// /// # Examples /// @@ -362,7 +420,7 @@ impl Runtime { self.blocking_pool.shutdown(Some(duration)); } - /// Shuts down the runtime, without waiting for any spawned tasks to shutdown. + /// 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 @@ -373,6 +431,8 @@ impl Runtime { /// 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))`. /// /// ``` @@ -403,12 +463,18 @@ impl Drop for Runtime { let _guard = context::try_set_current(&self.handle.inner); current_thread.shutdown(&self.handle.inner); } - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] 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", not(target_os = "wasi")))] + 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); + } } } } diff --git a/src/runtime/scheduler/block_in_place.rs b/src/runtime/scheduler/block_in_place.rs new file mode 100644 index 0000000..803ff45 --- /dev/null +++ b/src/runtime/scheduler/block_in_place.rs @@ -0,0 +1,21 @@ +use crate::runtime::scheduler; + +#[track_caller] +pub(crate) fn block_in_place<F, R>(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) +} diff --git a/src/runtime/scheduler/current_thread.rs b/src/runtime/scheduler/current_thread/mod.rs index 375e47c..30b17c0 100644 --- a/src/runtime/scheduler/current_thread.rs +++ b/src/runtime/scheduler/current_thread/mod.rs @@ -1,10 +1,10 @@ use crate::future::poll_fn; use crate::loom::sync::atomic::AtomicBool; -use crate::loom::sync::{Arc, Mutex}; +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}; -use crate::runtime::{blocking, context, scheduler, Config}; -use crate::runtime::{MetricsBatch, SchedulerMetrics, WorkerMetrics}; +use crate::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics}; use crate::sync::notify::Notify; use crate::util::atomic_cell::AtomicCell; use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef}; @@ -15,6 +15,7 @@ use std::fmt; use std::future::Future; use std::sync::atomic::Ordering::{AcqRel, Release}; use std::task::Poll::{Pending, Ready}; +use std::task::Waker; use std::time::Duration; /// Executes tasks on the current thread @@ -46,7 +47,7 @@ pub(crate) struct Handle { /// a function that will perform the scheduling work and acts as a capability token. struct Core { /// Scheduler run queue - tasks: VecDeque<task::Notified<Arc<Handle>>>, + tasks: VecDeque<Notified>, /// Current tick tick: u32, @@ -59,6 +60,9 @@ struct Core { /// 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, @@ -66,8 +70,8 @@ struct Core { /// Scheduler state shared between threads. struct Shared { - /// Remote run queue. None if the `Runtime` has been dropped. - queue: Mutex<Option<VecDeque<task::Notified<Arc<Handle>>>>>, + /// Remote run queue + inject: Inject<Arc<Handle>>, /// Collection of all active tasks spawned onto this executor. owned: OwnedTasks<Arc<Handle>>, @@ -86,20 +90,29 @@ struct Shared { } /// Thread-local context. -struct Context { +/// +/// pub(crate) to store in `runtime::context`. +pub(crate) struct Context { /// Scheduler handle handle: Arc<Handle>, /// Scheduler core, enabling the holder of `Context` to execute the /// scheduler. core: RefCell<Option<Box<Core>>>, + + /// Deferred tasks, usually ones that called `task::yield_now()`. + pub(crate) defer: Defer, } +type Notified = task::Notified<Arc<Handle>>; + /// Initial queue capacity. const INITIAL_CAPACITY: usize = 64; -// Tracks the current CurrentThread. -scoped_thread_local!(static CURRENT: Context); +/// 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( @@ -109,14 +122,21 @@ impl CurrentThread { seed_generator: RngSeedGenerator, config: Config, ) -> (CurrentThread, Arc<Handle>) { + let worker_metrics = WorkerMetrics::from_config(&config); + + // 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 { shared: Shared { - queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))), + inject: Inject::new(), owned: OwnedTasks::new(), woken: AtomicBool::new(false), config, scheduler_metrics: SchedulerMetrics::new(), - worker_metrics: WorkerMetrics::new(), + worker_metrics, }, driver: driver_handle, blocking_spawner, @@ -127,7 +147,8 @@ impl CurrentThread { tasks: VecDeque::with_capacity(INITIAL_CAPACITY), tick: 0, driver: Some(driver), - metrics: MetricsBatch::new(), + metrics: MetricsBatch::new(&handle.shared.worker_metrics), + global_queue_interval, unhandled_panic: false, }))); @@ -143,48 +164,49 @@ impl CurrentThread { pub(crate) fn block_on<F: Future>(&self, handle: &scheduler::Handle, future: F) -> F::Output { pin!(future); - let mut enter = crate::runtime::context::enter_runtime(handle, false); - 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) { - return core.block_on(future); - } else { - let notified = self.notify.notified(); - pin!(notified); - - if let Some(out) = enter - .blocking - .block_on(poll_fn(|cx| { - if notified.as_mut().poll(cx).is_ready() { - return Ready(None); - } + 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) { + 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)); - } + if let Ready(out) = future.as_mut().poll(cx) { + return Ready(Some(out)); + } - Pending - })) - .expect("Failed to `Enter::block_on`") - { - return out; + Pending + })) + .expect("Failed to `Enter::block_on`") + { + return out; + } } } - } + }) } fn take_core(&self, handle: &Arc<Handle>) -> Option<CoreGuard<'_>> { let core = self.core.take()?; Some(CoreGuard { - context: Context { + context: scheduler::Context::CurrentThread(Context { handle: handle.clone(), core: RefCell::new(Some(core)), - }, + defer: Defer::new(), + }), scheduler: self, }) } @@ -201,42 +223,58 @@ impl CurrentThread { None => panic!("Oh no! We never placed the Core back, this is a bug!"), }; - core.enter(|mut core, _context| { - // 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(); + // Check that the thread-local is not being destroyed + let tls_available = context::with_current(|_| ()).is_ok(); - // Drain local queue - // We already shut down every task, so we just need to drop the task. - while let Some(task) = core.pop_task(handle) { - drop(task); - } + 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); + } + } +} - // Drain remote queue and set it to None - let remote_queue = handle.shared.queue.lock().take(); +fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> { + // 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(); - // Using `Option::take` to replace the shared queue with `None`. - // We already shut down every task, so we just need to drop the task. - if let Some(remote_queue) = remote_queue { - for task in remote_queue { - drop(task); - } - } + // 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); + } - assert!(handle.shared.owned.is_empty()); + // Close the injection queue + handle.shared.inject.close(); - // Submit metrics - core.metrics.submit(&handle.shared.worker_metrics); + // Drain remote queue + while let Some(task) = handle.shared.inject.pop() { + drop(task); + } - // Shutdown the resource drivers - if let Some(driver) = core.driver.as_mut() { - driver.shutdown(&handle.driver); - } + assert!(handle.shared.owned.is_empty()); - (core, ()) - }); + // 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 { @@ -248,7 +286,23 @@ impl fmt::Debug for CurrentThread { // ===== impl Core ===== impl Core { - fn pop_task(&mut self, handle: &Handle) -> Option<task::Notified<Arc<Handle>>> { + /// 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<Notified> { + 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<Notified> { let ret = self.tasks.pop_front(); handle .shared @@ -257,7 +311,7 @@ impl Core { ret } - fn push_task(&mut self, handle: &Handle, task: task::Notified<Arc<Handle>>) { + fn push_task(&mut self, handle: &Handle, task: Notified) { self.tasks.push_back(task); self.metrics.inc_local_schedule_count(); handle @@ -265,14 +319,18 @@ impl Core { .worker_metrics .set_queue_depth(self.tasks.len()); } -} -fn did_defer_tasks() -> bool { - context::with_defer(|deferred| !deferred.is_empty()).unwrap() + fn submit_metrics(&mut self, handle: &Handle) { + self.metrics.submit(&handle.shared.worker_metrics, 0); + } } -fn wake_deferred_tasks() { - context::with_defer(|deferred| deferred.wake()); +#[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 ===== @@ -281,8 +339,10 @@ impl Context { /// Execute the closure with the given scheduler core stored in the /// thread-local context. fn run_task<R>(&self, mut core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) { - core.metrics.incr_poll_count(); - self.enter(core, || crate::runtime::coop::budget(f)) + 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, @@ -303,15 +363,14 @@ impl Context { if core.tasks.is_empty() { // Park until the thread is signaled core.metrics.about_to_park(); - core.metrics.submit(&handle.shared.worker_metrics); + core.submit_metrics(handle); let (c, _) = self.enter(core, || { driver.park(&handle.driver); - wake_deferred_tasks(); + self.defer.wake(); }); core = c; - core.metrics.returned_from_park(); } if let Some(f) = &handle.shared.config.after_unpark { @@ -330,10 +389,11 @@ impl Context { fn park_yield(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> { let mut driver = core.driver.take().expect("driver missing"); - core.metrics.submit(&handle.shared.worker_metrics); + core.submit_metrics(handle); + let (mut core, _) = self.enter(core, || { driver.park_timeout(&handle.driver, Duration::from_millis(0)); - wake_deferred_tasks(); + self.defer.wake(); }); core.driver = Some(driver); @@ -353,6 +413,10 @@ impl Context { 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 ===== @@ -377,11 +441,58 @@ impl Handle { handle } - fn pop(&self) -> Option<task::Notified<Arc<Handle>>> { - match self.shared.queue.lock().as_mut() { - Some(queue) => queue.pop_front(), - None => None, - } + /// 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(dump::Task::new) + .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<Notified> { + self.shared.inject.pop() } fn waker_ref(me: &Arc<Self>) -> WakerRef<'_> { @@ -404,14 +515,7 @@ cfg_metrics! { } pub(crate) fn injection_queue_depth(&self) -> usize { - // TODO: avoid having to lock. The multi-threaded injection queue - // could probably be used here. - self.shared - .queue - .lock() - .as_ref() - .map(|queue| queue.len()) - .unwrap_or(0) + self.shared.inject.len() } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { @@ -419,6 +523,10 @@ cfg_metrics! { &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() } @@ -430,6 +538,20 @@ cfg_metrics! { pub(crate) fn blocking_queue_depth(&self) -> usize { self.blocking_spawner.queue_depth() } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.shared.owned.active_tasks_count() + } + } +} + +cfg_unstable! { + use std::num::NonZeroU64; + + impl Handle { + pub(crate) fn owned_id(&self) -> NonZeroU64 { + self.shared.owned.id + } } } @@ -447,8 +569,10 @@ impl Schedule for Arc<Handle> { } fn schedule(&self, task: task::Notified<Self>) { - CURRENT.with(|maybe_cx| match maybe_cx { - Some(cx) if Arc::ptr_eq(self, &cx.handle) => { + 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 @@ -461,14 +585,9 @@ impl Schedule for Arc<Handle> { // Track that a task was scheduled from **outside** of the runtime. self.shared.scheduler_metrics.inc_remote_schedule_count(); - // If the queue is None, then the runtime has shut down. We - // don't need to do anything with the notification in that case. - let mut guard = self.shared.queue.lock(); - if let Some(queue) = guard.as_mut() { - queue.push_back(task); - drop(guard); - self.driver.unpark(); - } + // Schedule the task + self.shared.inject.push(task); + self.driver.unpark(); } }); } @@ -482,11 +601,14 @@ impl Schedule for Arc<Handle> { // Do nothing } UnhandledPanic::ShutdownRuntime => { + use scheduler::Context::CurrentThread; + // 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(|maybe_cx| match maybe_cx { - Some(cx) if Arc::ptr_eq(self, &cx.handle) => { + // `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 @@ -520,7 +642,7 @@ impl Wake for Handle { /// Used to ensure we always place the `Core` value back into its slot in /// `CurrentThread`, even if the future panics. struct CoreGuard<'a> { - context: Context, + context: scheduler::Context, scheduler: &'a CurrentThread, } @@ -533,6 +655,8 @@ impl CoreGuard<'_> { pin!(future); + core.metrics.start_processing_scheduled_tasks(); + 'outer: loop { let handle = &context.handle; @@ -554,25 +678,23 @@ impl CoreGuard<'_> { return (core, None); } - // Get and increment the current tick - let tick = core.tick; - core.tick = core.tick.wrapping_add(1); + core.tick(); - let entry = if tick % handle.shared.config.global_queue_interval == 0 { - handle.pop().or_else(|| core.tasks.pop_front()) - } else { - core.tasks.pop_front().or_else(|| handle.pop()) - }; + let entry = core.next_task(handle); let task = match entry { Some(entry) => entry, None => { - core = if did_defer_tasks() { + 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; } @@ -587,9 +709,13 @@ impl CoreGuard<'_> { 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(); } }); @@ -608,13 +734,15 @@ impl CoreGuard<'_> { where F: FnOnce(Box<Core>, &Context) -> (Box<Core>, R), { + let context = self.context.expect_current_thread(); + // Remove `core` from `context` to pass into the closure. - let core = self.context.core.borrow_mut().take().expect("core missing"); + let core = context.core.borrow_mut().take().expect("core missing"); // Call the closure and place `core` back - let (core, ret) = CURRENT.set(&self.context, || f(core, &self.context)); + let (core, ret) = context::set_scheduler(&self.context, || f(core, context)); - *self.context.core.borrow_mut() = Some(core); + *context.core.borrow_mut() = Some(core); ret } @@ -622,7 +750,9 @@ impl CoreGuard<'_> { impl Drop for CoreGuard<'_> { fn drop(&mut self) { - if let Some(core) = self.context.core.borrow_mut().take() { + 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); diff --git a/src/runtime/scheduler/defer.rs b/src/runtime/scheduler/defer.rs new file mode 100644 index 0000000..a4be8ef --- /dev/null +++ b/src/runtime/scheduler/defer.rs @@ -0,0 +1,43 @@ +use std::cell::RefCell; +use std::task::Waker; + +pub(crate) struct Defer { + deferred: RefCell<Vec<Waker>>, +} + +impl Defer { + pub(crate) fn new() -> Defer { + Defer { + deferred: Default::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<Waker> { + let mut deferred = self.deferred.borrow_mut(); + std::mem::take(&mut *deferred) + } +} diff --git a/src/runtime/scheduler/inject.rs b/src/runtime/scheduler/inject.rs new file mode 100644 index 0000000..39976fc --- /dev/null +++ b/src/runtime/scheduler/inject.rs @@ -0,0 +1,72 @@ +//! 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; +} + +cfg_metrics! { + 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<T: 'static> { + shared: Shared<T>, + synced: Mutex<Synced>, +} + +impl<T: 'static> Inject<T> { + pub(crate) fn new() -> Inject<T> { + 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<T>) { + let mut synced = self.synced.lock(); + // safety: passing correct `Synced` + unsafe { self.shared.push(&mut synced, task) } + } + + pub(crate) fn pop(&self) -> Option<task::Notified<T>> { + if self.shared.is_empty() { + return None; + } + + let mut synced = self.synced.lock(); + // safety: passing correct `Synced` + unsafe { self.shared.pop(&mut synced) } + } +} diff --git a/src/runtime/scheduler/inject/metrics.rs b/src/runtime/scheduler/inject/metrics.rs new file mode 100644 index 0000000..76f045f --- /dev/null +++ b/src/runtime/scheduler/inject/metrics.rs @@ -0,0 +1,7 @@ +use super::Inject; + +impl<T: 'static> Inject<T> { + pub(crate) fn len(&self) -> usize { + self.shared.len() + } +} diff --git a/src/runtime/scheduler/inject/pop.rs b/src/runtime/scheduler/inject/pop.rs new file mode 100644 index 0000000..4e6d5d3 --- /dev/null +++ b/src/runtime/scheduler/inject/pop.rs @@ -0,0 +1,55 @@ +use 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<T>, +} + +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<T>; + + fn next(&mut self) -> Option<Self::Item> { + 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<usize>) { + (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() {} + } +} diff --git a/src/runtime/scheduler/inject/rt_multi_thread.rs b/src/runtime/scheduler/inject/rt_multi_thread.rs new file mode 100644 index 0000000..1d5f040 --- /dev/null +++ b/src/runtime/scheduler/inject/rt_multi_thread.rs @@ -0,0 +1,113 @@ +use super::{Shared, Synced}; + +use crate::runtime::scheduler::Lock; +use crate::runtime::task; + +use std::sync::atomic::Ordering::Release; + +impl<'a> Lock<Synced> for &'a mut Synced { + type Handle = &'a mut Synced; + + fn lock(self) -> Self::Handle { + self + } +} + +impl AsMut<Synced> for Synced { + fn as_mut(&mut self) -> &mut Synced { + self + } +} + +impl<T: 'static> Shared<T> { + /// 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<L, I>(&self, shared: L, mut iter: I) + where + L: Lock<Synced>, + I: Iterator<Item = task::Notified<T>>, + { + 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<L>( + &self, + shared: L, + batch_head: task::RawTask, + batch_tail: task::RawTask, + num: usize, + ) where + L: Lock<Synced>, + { + 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::<T>::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); + } +} diff --git a/src/runtime/scheduler/inject/shared.rs b/src/runtime/scheduler/inject/shared.rs new file mode 100644 index 0000000..2d29486 --- /dev/null +++ b/src/runtime/scheduler/inject/shared.rs @@ -0,0 +1,124 @@ +use 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<T: 'static> { + /// Number of pending tasks in the queue. This helps prevent unnecessary + /// locking in the hot path. + pub(super) len: AtomicUsize, + + _p: PhantomData<T>, +} + +unsafe impl<T> Send for Shared<T> {} +unsafe impl<T> Sync for Shared<T> {} + +impl<T: 'static> Shared<T> { + pub(crate) fn new() -> (Shared<T>, 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, + all(feature = "rt-multi-thread", not(target_os = "wasi")) + ))] + 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<T>) { + 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<task::Notified<T>> { + 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) + } +} diff --git a/src/runtime/scheduler/inject/synced.rs b/src/runtime/scheduler/inject/synced.rs new file mode 100644 index 0000000..45f6038 --- /dev/null +++ b/src/runtime/scheduler/inject/synced.rs @@ -0,0 +1,41 @@ +#![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<task::RawTask>, + + /// Linked-list tail. + pub(super) tail: Option<task::RawTask>, +} + +unsafe impl Send for Synced {} +unsafe impl Sync for Synced {} + +impl Synced { + pub(super) fn pop<T: 'static>(&mut self) -> Option<task::Notified<T>> { + 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() + } +} diff --git a/src/runtime/scheduler/lock.rs b/src/runtime/scheduler/lock.rs new file mode 100644 index 0000000..0901c2b --- /dev/null +++ b/src/runtime/scheduler/lock.rs @@ -0,0 +1,6 @@ +/// A lock (mutex) yielding generic data. +pub(crate) trait Lock<T> { + type Handle: AsMut<T>; + + fn lock(self) -> Self::Handle; +} diff --git a/src/runtime/scheduler/mod.rs b/src/runtime/scheduler/mod.rs index f45d8a8..d02c027 100644 --- a/src/runtime/scheduler/mod.rs +++ b/src/runtime/scheduler/mod.rs @@ -1,11 +1,28 @@ cfg_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; } 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; @@ -15,9 +32,12 @@ pub(crate) enum Handle { #[cfg(feature = "rt")] CurrentThread(Arc<current_thread::Handle>), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] MultiThread(Arc<multi_thread::Handle>), + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + MultiThreadAlt(Arc<multi_thread_alt::Handle>), + // 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"))] @@ -25,6 +45,17 @@ pub(crate) enum Handle { Disabled, } +#[cfg(feature = "rt")] +pub(super) enum Context { + CurrentThread(current_thread::Context), + + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] + MultiThread(multi_thread::Context), + + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + MultiThreadAlt(multi_thread_alt::Context), +} + impl Handle { #[cfg_attr(not(feature = "full"), allow(dead_code))] pub(crate) fn driver(&self) -> &driver::Handle { @@ -32,9 +63,12 @@ impl Handle { #[cfg(feature = "rt")] Handle::CurrentThread(ref h) => &h.driver, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Handle::MultiThread(ref h) => &h.driver, + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Handle::MultiThreadAlt(ref h) => &h.driver, + #[cfg(not(feature = "rt"))] Handle::Disabled => unreachable!(), } @@ -48,23 +82,33 @@ cfg_rt! { 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(all(feature = "rt-multi-thread", not(target_os = "wasi")))] + $ty::MultiThread($h) => $e, + + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + $ty::MultiThreadAlt($h) => $e, + } + } + } impl Handle { #[track_caller] pub(crate) fn current() -> Handle { - match context::try_current() { + match context::with_current(Clone::clone) { Ok(handle) => handle, Err(e) => panic!("{}", e), } } pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner { - match self { - Handle::CurrentThread(h) => &h.blocking_spawner, - - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(h) => &h.blocking_spawner, - } + match_flavor!(self, Handle(h) => &h.blocking_spawner) } pub(crate) fn spawn<F>(&self, future: F, id: Id) -> JoinHandle<F::Output> @@ -75,8 +119,11 @@ cfg_rt! { match self { Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id), + + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Handle::MultiThreadAlt(h) => multi_thread_alt::Handle::spawn(h, future, id), } } @@ -84,27 +131,36 @@ cfg_rt! { match *self { Handle::CurrentThread(_) => {}, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Handle::MultiThread(ref h) => h.shutdown(), + + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Handle::MultiThreadAlt(ref h) => h.shutdown(), } } pub(crate) fn seed_generator(&self) -> &RngSeedGenerator { - match self { - Handle::CurrentThread(h) => &h.seed_generator, - - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(h) => &h.seed_generator, - } + match_flavor!(self, Handle(h) => &h.seed_generator) } pub(crate) fn as_current_thread(&self) -> &Arc<current_thread::Handle> { match self { Handle::CurrentThread(handle) => handle, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] _ => panic!("not a CurrentThread handle"), } } + + cfg_rt_multi_thread! { + cfg_unstable! { + pub(crate) fn expect_multi_thread_alt(&self) -> &Arc<multi_thread_alt::Handle> { + match self { + Handle::MultiThreadAlt(handle) => handle, + _ => panic!("not a `MultiThreadAlt` handle"), + } + } + } + } } cfg_metrics! { @@ -114,64 +170,77 @@ cfg_rt! { pub(crate) fn num_workers(&self) -> usize { match self { Handle::CurrentThread(_) => 1, - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))] Handle::MultiThread(handle) => handle.num_workers(), + #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))] + Handle::MultiThreadAlt(handle) => handle.num_workers(), } } pub(crate) fn num_blocking_threads(&self) -> usize { - match self { - Handle::CurrentThread(handle) => handle.num_blocking_threads(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.num_blocking_threads(), - } + match_flavor!(self, Handle(handle) => handle.num_blocking_threads()) } pub(crate) fn num_idle_blocking_threads(&self) -> usize { - match self { - Handle::CurrentThread(handle) => handle.num_idle_blocking_threads(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.num_idle_blocking_threads(), - } + match_flavor!(self, Handle(handle) => handle.num_idle_blocking_threads()) + } + + pub(crate) fn active_tasks_count(&self) -> usize { + match_flavor!(self, Handle(handle) => handle.active_tasks_count()) } pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { - match self { - Handle::CurrentThread(handle) => handle.scheduler_metrics(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.scheduler_metrics(), - } + match_flavor!(self, Handle(handle) => handle.scheduler_metrics()) } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { - match self { - Handle::CurrentThread(handle) => handle.worker_metrics(worker), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.worker_metrics(worker), - } + match_flavor!(self, Handle(handle) => handle.worker_metrics(worker)) } pub(crate) fn injection_queue_depth(&self) -> usize { - match self { - Handle::CurrentThread(handle) => handle.injection_queue_depth(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.injection_queue_depth(), - } + match_flavor!(self, Handle(handle) => handle.injection_queue_depth()) } pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { - match self { - Handle::CurrentThread(handle) => handle.worker_metrics(worker).queue_depth(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.worker_local_queue_depth(worker), - } + 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(all(feature = "rt-multi-thread", not(target_os = "wasi")))] + _ => 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 { - Handle::CurrentThread(handle) => handle.blocking_queue_depth(), - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] - Handle::MultiThread(handle) => handle.blocking_queue_depth(), + 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`") + } } } } diff --git a/src/runtime/scheduler/multi_thread/counters.rs b/src/runtime/scheduler/multi_thread/counters.rs new file mode 100644 index 0000000..50bcc11 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/counters.rs @@ -0,0 +1,62 @@ +#[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::*; diff --git a/src/runtime/scheduler/multi_thread/handle.rs b/src/runtime/scheduler/multi_thread/handle.rs index 69a4620..568eb80 100644 --- a/src/runtime/scheduler/multi_thread/handle.rs +++ b/src/runtime/scheduler/multi_thread/handle.rs @@ -9,6 +9,14 @@ use crate::util::RngSeedGenerator; use std::fmt; +cfg_metrics! { + mod metrics; +} + +cfg_taskdump! { + mod taskdump; +} + /// Handle to the multi thread scheduler pub(crate) struct Handle { /// Task spawner @@ -45,48 +53,18 @@ impl Handle { { let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); - if let Some(notified) = notified { - me.schedule_task(notified, false); - } + me.schedule_option_task_without_yield(notified); handle } } -cfg_metrics! { - use crate::runtime::{SchedulerMetrics, WorkerMetrics}; +cfg_unstable! { + use std::num::NonZeroU64; impl Handle { - pub(crate) fn num_workers(&self) -> usize { - self.shared.worker_metrics.len() - } - - 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 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() + pub(crate) fn owned_id(&self) -> NonZeroU64 { + self.shared.owned.id } } } diff --git a/src/runtime/scheduler/multi_thread/handle/metrics.rs b/src/runtime/scheduler/multi_thread/handle/metrics.rs new file mode 100644 index 0000000..838694f --- /dev/null +++ b/src/runtime/scheduler/multi_thread/handle/metrics.rs @@ -0,0 +1,41 @@ +use 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 { + self.blocking_spawner.num_threads() + } + + pub(crate) fn num_idle_blocking_threads(&self) -> usize { + self.blocking_spawner.num_idle_threads() + } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.shared.owned.active_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() + } +} diff --git a/src/runtime/scheduler/multi_thread/handle/taskdump.rs b/src/runtime/scheduler/multi_thread/handle/taskdump.rs new file mode 100644 index 0000000..477d857 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/handle/taskdump.rs @@ -0,0 +1,26 @@ +use 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 + } +} diff --git a/src/runtime/scheduler/multi_thread/idle.rs b/src/runtime/scheduler/multi_thread/idle.rs index a57bf6a..834bc2b 100644 --- a/src/runtime/scheduler/multi_thread/idle.rs +++ b/src/runtime/scheduler/multi_thread/idle.rs @@ -1,7 +1,7 @@ //! Coordinates idling workers use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::Mutex; +use crate::runtime::scheduler::multi_thread::Shared; use std::fmt; use std::sync::atomic::Ordering::{self, SeqCst}; @@ -13,13 +13,16 @@ pub(super) struct Idle { /// Used as a fast-path to avoid acquiring the lock when needed. state: AtomicUsize, - /// Sleeping workers - sleepers: Mutex<Vec<usize>>, - /// Total number of workers. num_workers: usize, } +/// Data synchronized by the scheduler mutex +pub(super) struct Synced { + /// Sleeping workers + sleepers: Vec<usize>, +} + const UNPARK_SHIFT: usize = 16; const UNPARK_MASK: usize = !SEARCH_MASK; const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1; @@ -28,19 +31,24 @@ const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1; struct State(usize); impl Idle { - pub(super) fn new(num_workers: usize) -> Idle { + pub(super) fn new(num_workers: usize) -> (Idle, Synced) { let init = State::new(num_workers); - Idle { + let idle = Idle { state: AtomicUsize::new(init.into()), - sleepers: Mutex::new(Vec::with_capacity(num_workers)), 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) -> Option<usize> { + pub(super) fn worker_to_notify(&self, shared: &Shared) -> Option<usize> { // 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. @@ -55,7 +63,7 @@ impl Idle { } // Acquire the lock - let mut sleepers = self.sleepers.lock(); + let mut lock = shared.synced.lock(); // Check again, now that the lock is acquired if !self.notify_should_wakeup() { @@ -67,7 +75,7 @@ impl Idle { State::unpark_one(&self.state, 1); // Get the worker to unpark - let ret = sleepers.pop(); + let ret = lock.idle.sleepers.pop(); debug_assert!(ret.is_some()); ret @@ -75,15 +83,20 @@ impl Idle { /// Returns `true` if the worker needs to do a final check for submitted /// work. - pub(super) fn transition_worker_to_parked(&self, worker: usize, is_searching: bool) -> bool { + pub(super) fn transition_worker_to_parked( + &self, + shared: &Shared, + worker: usize, + is_searching: bool, + ) -> bool { // Acquire the lock - let mut sleepers = self.sleepers.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 - sleepers.push(worker); + lock.idle.sleepers.push(worker); ret } @@ -113,8 +126,9 @@ impl Idle { /// 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, worker_id: usize) -> bool { - let mut sleepers = self.sleepers.lock(); + 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 { @@ -131,9 +145,9 @@ impl Idle { } /// Returns `true` if `worker_id` is contained in the sleep set. - pub(super) fn is_parked(&self, worker_id: usize) -> bool { - let sleepers = self.sleepers.lock(); - sleepers.contains(&worker_id) + 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 { diff --git a/src/runtime/scheduler/multi_thread/mod.rs b/src/runtime/scheduler/multi_thread/mod.rs index 47cd1f3..d85a0ae 100644 --- a/src/runtime/scheduler/multi_thread/mod.rs +++ b/src/runtime/scheduler/multi_thread/mod.rs @@ -1,18 +1,39 @@ //! 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::Launch; +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; @@ -62,11 +83,9 @@ impl MultiThread { where F: Future, { - let mut enter = crate::runtime::context::enter_runtime(handle, true); - enter - .blocking - .block_on(future) - .expect("failed to park thread") + 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) { diff --git a/src/runtime/scheduler/multi_thread/overflow.rs b/src/runtime/scheduler/multi_thread/overflow.rs new file mode 100644 index 0000000..ab66481 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/overflow.rs @@ -0,0 +1,26 @@ +use crate::runtime::task; + +#[cfg(test)] +use std::cell::RefCell; + +pub(crate) trait Overflow<T: 'static> { + fn push(&self, task: task::Notified<T>); + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<T>>; +} + +#[cfg(test)] +impl<T: 'static> Overflow<T> for RefCell<Vec<task::Notified<T>>> { + fn push(&self, task: task::Notified<T>) { + self.borrow_mut().push(task); + } + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<T>>, + { + self.borrow_mut().extend(iter); + } +} diff --git a/src/runtime/scheduler/multi_thread/park.rs b/src/runtime/scheduler/multi_thread/park.rs index 6bdbff9..0a00ea0 100644 --- a/src/runtime/scheduler/multi_thread/park.rs +++ b/src/runtime/scheduler/multi_thread/park.rs @@ -4,7 +4,6 @@ use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::{Arc, Condvar, Mutex}; -use crate::loom::thread; use crate::runtime::driver::{self, Driver}; use crate::util::TryLock; @@ -104,18 +103,14 @@ impl Unparker { impl Inner { /// Parks the current thread for at most `dur`. fn park(&self, handle: &driver::Handle) { - for _ in 0..3 { - // 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; - } - - thread::yield_now(); + // 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() { diff --git a/src/runtime/scheduler/multi_thread/queue.rs b/src/runtime/scheduler/multi_thread/queue.rs index faf56db..dd66fa2 100644 --- a/src/runtime/scheduler/multi_thread/queue.rs +++ b/src/runtime/scheduler/multi_thread/queue.rs @@ -2,8 +2,8 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::Arc; -use crate::runtime::task::{self, Inject}; -use crate::runtime::MetricsBatch; +use crate::runtime::scheduler::multi_thread::{Overflow, Stats}; +use crate::runtime::task; use std::mem::{self, MaybeUninit}; use std::ptr; @@ -105,9 +105,18 @@ pub(crate) fn local<T: 'static>() -> (Steal<T>, Local<T>) { } impl<T> Local<T> { - /// Returns true if the queue has entries that can be stolen. - pub(crate) fn is_stealable(&self) -> bool { - !self.inner.is_empty() + /// 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 @@ -118,12 +127,66 @@ impl<T> Local<T> { !self.inner.is_empty() } - /// Pushes a task to the back of the local queue, skipping the LIFO slot. - pub(crate) fn push_back( + /// 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<Item = task::Notified<T>>) { + 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 curent 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<O: Overflow<T>>( &mut self, mut task: task::Notified<T>, - inject: &Inject<T>, - metrics: &mut MetricsBatch, + overflow: &O, + stats: &mut Stats, ) { let tail = loop { let head = self.inner.head.load(Acquire); @@ -138,12 +201,12 @@ impl<T> Local<T> { } else if steal != real { // Concurrently stealing, this will free up capacity, so only // push the task onto the inject queue - inject.push(task); + 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, inject, metrics) { + match self.push_overflow(task, real, tail, overflow, stats) { Ok(_) => return, // Lost the race, try again Err(v) => { @@ -153,6 +216,11 @@ impl<T> Local<T> { } }; + self.push_back_finish(task, tail); + } + + // Second half of `push_back` + fn push_back_finish(&self, task: task::Notified<T>, tail: UnsignedShort) { // Map the position to a slot index. let idx = tail as usize & MASK; @@ -179,13 +247,13 @@ impl<T> Local<T> { /// workers "missed" some of the tasks during a steal, they will get /// another opportunity. #[inline(never)] - fn push_overflow( + fn push_overflow<O: Overflow<T>>( &mut self, task: task::Notified<T>, head: UnsignedShort, tail: UnsignedShort, - inject: &Inject<T>, - metrics: &mut MetricsBatch, + overflow: &O, + stats: &mut Stats, ) -> Result<(), task::Notified<T>> { /// How many elements are we taking from the local queue. /// @@ -267,10 +335,10 @@ impl<T> Local<T> { head: head as UnsignedLong, i: 0, }; - inject.push_batch(batch_iter.chain(std::iter::once(task))); + overflow.push_batch(batch_iter.chain(std::iter::once(task))); // Add 1 to factor in the task currently being scheduled. - metrics.incr_overflow_count(); + stats.incr_overflow_count(); Ok(()) } @@ -326,7 +394,7 @@ impl<T> Steal<T> { pub(crate) fn steal_into( &self, dst: &mut Local<T>, - dst_metrics: &mut MetricsBatch, + dst_stats: &mut Stats, ) -> Option<task::Notified<T>> { // Safety: the caller is the only thread that mutates `dst.tail` and // holds a mutable reference. @@ -352,8 +420,8 @@ impl<T> Steal<T> { return None; } - dst_metrics.incr_steal_count(n as u16); - dst_metrics.incr_steal_operations(); + dst_stats.incr_steal_count(n as u16); + dst_stats.incr_steal_operations(); // We are returning a task here n -= 1; @@ -501,6 +569,13 @@ impl<T> Drop for Local<T> { } impl<T> Inner<T> { + 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); diff --git a/src/runtime/scheduler/multi_thread/stats.rs b/src/runtime/scheduler/multi_thread/stats.rs new file mode 100644 index 0000000..3b8c502 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/stats.rs @@ -0,0 +1,140 @@ +use 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, + + /// 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; + + cmp::max( + // 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 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(); + } +} diff --git a/src/runtime/scheduler/multi_thread/trace.rs b/src/runtime/scheduler/multi_thread/trace.rs new file mode 100644 index 0000000..7b4aeb5 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/trace.rs @@ -0,0 +1,61 @@ +use 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<Option<Dump>>, +} + +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<Dump> { + 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; + } + } +} diff --git a/src/runtime/scheduler/multi_thread/trace_mock.rs b/src/runtime/scheduler/multi_thread/trace_mock.rs new file mode 100644 index 0000000..2c17a4e --- /dev/null +++ b/src/runtime/scheduler/multi_thread/trace_mock.rs @@ -0,0 +1,11 @@ +pub(super) struct TraceStatus {} + +impl TraceStatus { + pub(super) fn new(_: usize) -> Self { + Self {} + } + + pub(super) fn trace_requested(&self) -> bool { + false + } +} diff --git a/src/runtime/scheduler/multi_thread/worker.rs b/src/runtime/scheduler/multi_thread/worker.rs index 148255a..8f3181f 100644 --- a/src/runtime/scheduler/multi_thread/worker.rs +++ b/src/runtime/scheduler/multi_thread/worker.rs @@ -59,17 +59,33 @@ use crate::loom::sync::{Arc, Mutex}; use crate::runtime; use crate::runtime::context; -use crate::runtime::scheduler::multi_thread::{queue, Handle, Idle, Parker, Unparker}; -use crate::runtime::task::{Inject, OwnedTasks}; +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; use crate::runtime::{ - blocking, coop, driver, scheduler, task, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics, + blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics, }; use crate::util::atomic_cell::AtomicCell; use crate::util::rand::{FastRand, RngSeedGenerator}; use std::cell::RefCell; +use std::task::Waker; use std::time::Duration; +cfg_metrics! { + mod metrics; +} + +cfg_taskdump! { + mod taskdump; +} + +cfg_not_taskdump! { + mod taskdump_mock; +} + /// A scheduler worker pub(super) struct Worker { /// Reference to scheduler's handle @@ -90,10 +106,14 @@ struct Core { /// 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 message passing patterns and - /// helps to reduce latency. + /// next (LIFO). This is an optimization for improving locality which + /// benefits message passing patterns and helps to reduce latency. lifo_slot: Option<Notified>, + /// 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<Arc<Handle>>, @@ -104,21 +124,27 @@ struct Core { /// 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<Parker>, - /// Batching metrics so they can be submitted to RuntimeMetrics. - metrics: MetricsBatch, + /// 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(super) struct Shared { +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]>, @@ -126,13 +152,16 @@ pub(super) struct Shared { /// 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 - inject: Inject<Arc<Handle>>, + pub(super) inject: inject::Shared<Arc<Handle>>, /// Coordinates idle workers idle: Idle, /// Collection of all active tasks spawned onto this executor. - pub(super) owned: OwnedTasks<Arc<Handle>>, + pub(crate) owned: OwnedTasks<Arc<Handle>>, + + /// Data synchronized by the scheduler mutex + pub(super) synced: Mutex<Synced>, /// Cores that have observed the shutdown signal /// @@ -141,6 +170,9 @@ pub(super) struct Shared { #[allow(clippy::vec_box)] // we're moving an already-boxed value shutdown_cores: Mutex<Vec<Box<Core>>>, + /// The number of cores that have observed the trace signal. + pub(super) trace_status: TraceStatus, + /// Scheduler configuration options config: Config, @@ -148,24 +180,43 @@ pub(super) struct Shared { 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. - steal: queue::Steal<Arc<Handle>>, + pub(super) steal: queue::Steal<Arc<Handle>>, /// Unparks the associated worker thread unpark: Unparker, } /// Thread-local context -struct Context { +pub(crate) struct Context { /// Worker worker: Arc<Worker>, /// Core data core: RefCell<Option<Box<Core>>>, + + /// Tasks to wake after resource drivers are polled. This is mostly to + /// handle yielded tasks. + pub(crate) defer: Defer, } /// Starts the workers @@ -182,8 +233,11 @@ type Task = task::Task<Arc<Handle>>; /// A notified task handle type Notified = task::Notified<Arc<Handle>>; -// Tracks thread-local state -scoped_thread_local!(static CURRENT: Context); +/// Value picked out of thin-air. Running the LIFO slot a handful of times +/// seemms 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, @@ -203,32 +257,47 @@ pub(super) fn create( 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), - metrics: MetricsBatch::new(), - rand: FastRand::new(config.seed_generator.next_seed()), + 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(WorkerMetrics::new()); + 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 { shared: Shared { remotes: remotes.into_boxed_slice(), - inject: Inject::new(), - idle: Idle::new(size), + inject, + idle, owned: OwnedTasks::new(), + 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, @@ -254,28 +323,34 @@ where F: FnOnce() -> R, { // Try to steal the worker core back - struct Reset(coop::Budget); + struct Reset { + take_core: bool, + budget: coop::Budget, + } impl Drop for Reset { fn drop(&mut self) { - CURRENT.with(|maybe_cx| { + with_current(|maybe_cx| { if let Some(cx) = maybe_cx { - let core = cx.worker.core.take(); - let mut cx_core = cx.core.borrow_mut(); - assert!(cx_core.is_none()); - *cx_core = core; + if self.take_core { + let core = cx.worker.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); + coop::set(self.budget); } }); } } let mut had_entered = false; + let mut take_core = false; - let setup_result = CURRENT.with(|maybe_cx| { + let setup_result = with_current(|maybe_cx| { match ( crate::runtime::context::current_enter_context(), maybe_cx.is_some(), @@ -325,6 +400,10 @@ where None => return Ok(()), }; + // 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()); @@ -351,7 +430,10 @@ where if had_entered { // Unset the current task's budget. Blocking sections are not // constrained by task budgets. - let _reset = Reset(coop::stop()); + let _reset = Reset { + take_core, + budget: coop::stop(), + }; crate::runtime::context::exit_runtime(f) } else { @@ -392,29 +474,47 @@ fn run(worker: Arc<Worker>) { }; let handle = scheduler::Handle::MultiThread(worker.handle.clone()); - let _enter = crate::runtime::context::enter_runtime(&handle, true); - // Set the worker context. - let cx = Context { - worker, - core: RefCell::new(None), - }; - - CURRENT.set(&cx, || { - // 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. - wake_deferred_tasks(); + 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<Core>) -> 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(); @@ -427,13 +527,18 @@ impl Context { 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 did_defer_tasks() { + core = if !self.defer.is_empty() { self.park_timeout(core, Some(Duration::from_millis(0))) } else { self.park(core) @@ -455,13 +560,21 @@ impl Context { // 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 - core.metrics.incr_poll_count(); *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. @@ -470,40 +583,89 @@ impl Context { // by another worker. let mut core = match self.core.borrow_mut().take() { Some(core) => core, - None => return Err(()), + 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 => return Ok(core), + None => { + self.reset_lifo_enabled(&mut core); + core.stats.end_poll(); + return Ok(core); + } }; - if coop::has_budget_remaining() { - // Run the LIFO task, then loop - core.metrics.incr_poll_count(); - *self.core.borrow_mut() = Some(core); - let task = self.worker.handle.shared.owned.assert_owner(task); - task.run(); - } else { + 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(task, self.worker.inject(), &mut core.metrics); + 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<Core>) -> Box<Core> { 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 @@ -526,10 +688,9 @@ impl Context { } if core.transition_to_parked(&self.worker) { - while !core.is_shutdown { - core.metrics.about_to_park(); + while !core.is_shutdown && !core.is_traced { + core.stats.about_to_park(); core = self.park_timeout(core, None); - core.metrics.returned_from_park(); // Run regularly scheduled maintenance core.maintenance(&self.worker); @@ -547,6 +708,8 @@ impl Context { } fn park_timeout(&self, mut core: Box<Core>, duration: Option<Duration>) -> Box<Core> { + self.assert_lifo_enabled_is_correct(&core); + // Take the parker out of core let mut park = core.park.take().expect("park missing"); @@ -560,7 +723,7 @@ impl Context { park.park(&self.worker.handle.driver); } - wake_deferred_tasks(); + self.defer.wake(); // Remove `core` from context core = self.core.borrow_mut().take().expect("core missing"); @@ -568,14 +731,16 @@ impl Context { // Place `park` back in `core` core.park = Some(park); - // If there are tasks available to steal, but this worker is not - // looking for tasks to steal, notify another worker. - if !core.is_searching && core.run_queue.is_stealable() { - self.worker.handle.notify_parked(); + if core.should_notify_others() { + self.worker.handle.notify_parked_local(); } core } + + pub(crate) fn defer(&self, waker: &Waker) { + self.defer.defer(waker); + } } impl Core { @@ -586,10 +751,57 @@ impl Core { /// Return the next notified task available to this worker. fn next_task(&mut self, worker: &Worker) -> Option<Notified> { - if self.tick % worker.handle.shared.config.global_queue_interval == 0 { - worker.inject().pop().or_else(|| self.next_local_task()) + 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 { - self.next_local_task().or_else(|| worker.inject().pop()) + 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 nto 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 immedietly + let ret = tasks.next(); + + // Push the rest of the on the run queue + self.run_queue.push_back(tasks); + + ret } } @@ -622,14 +834,14 @@ impl Core { let target = &worker.handle.shared.remotes[i]; if let Some(task) = target .steal - .steal_into(&mut self.run_queue, &mut self.metrics) + .steal_into(&mut self.run_queue, &mut self.stats) { return Some(task); } } // Fallback on checking the global queue - worker.handle.shared.inject.pop() + worker.handle.next_remote_task() } fn transition_to_searching(&mut self, worker: &Worker) -> bool { @@ -649,23 +861,36 @@ impl Core { 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.lifo_slot.is_some() || self.run_queue.has_tasks() { + 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.index, self.is_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. @@ -680,18 +905,27 @@ impl Core { /// Returns `true` if the transition happened. fn transition_from_parked(&mut self, worker: &Worker) -> bool { - // If a task is in the lifo slot, then we must unpark regardless of + // If a task is in the lifo slot/run queue, then we must unpark regardless of // being notified - if self.lifo_slot.is_some() { + 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.index); + 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.index) { + if worker + .handle + .shared + .idle + .is_parked(&worker.handle.shared, worker.index) + { return false; } @@ -702,12 +936,18 @@ impl Core { /// Runs maintenance work such as checking the pool's state. fn maintenance(&mut self, worker: &Worker) { - self.metrics + self.stats .submit(&worker.handle.shared.worker_metrics[worker.index]); if !self.is_shutdown { // Check if the scheduler has been shutdown - self.is_shutdown = worker.inject().is_closed(); + 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(); } } @@ -717,7 +957,7 @@ impl Core { // Signal to all tasks to shut down. worker.handle.shared.owned.close_and_shutdown_all(); - self.metrics + self.stats .submit(&worker.handle.shared.worker_metrics[worker.index]); } @@ -731,11 +971,24 @@ impl Core { 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); + + debug_assert!(next > 1); + + // Smooth out jitter + if 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<Arc<Handle>> { + fn inject(&self) -> &inject::Shared<Arc<Handle>> { &self.handle.shared.inject } } @@ -757,7 +1010,7 @@ impl task::Schedule for Arc<Handle> { impl Handle { pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { - CURRENT.with(|maybe_cx| { + 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) { @@ -770,22 +1023,27 @@ impl Handle { } // Otherwise, use the inject queue. - self.shared.inject.push(task); - self.shared.scheduler_metrics.inc_remote_schedule_count(); - self.notify_parked(); + self.push_remote_task(task); + self.notify_parked_remote(); }) } + pub(super) fn schedule_option_task_without_yield(&self, task: Option<Notified>) { + if let Some(task) = task { + self.schedule_task(task, false); + } + } + fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { - core.metrics.inc_local_schedule_count(); + 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 || self.shared.config.disable_lifo_slot { + let should_notify = if is_yield || !core.lifo_enabled { core.run_queue - .push_back(task, &self.shared.inject, &mut core.metrics); + .push_back_or_overflow(task, self, &mut core.stats); true } else { // Push to the LIFO slot @@ -794,7 +1052,7 @@ impl Handle { if let Some(prev) = prev { core.run_queue - .push_back(prev, &self.shared.inject, &mut core.metrics); + .push_back_or_overflow(prev, self, &mut core.stats); } core.lifo_slot = Some(task); @@ -806,23 +1064,56 @@ impl Handle { // 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(); + self.notify_parked_local(); + } + } + + fn next_remote_task(&self) -> Option<Notified> { + 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() { + if self + .shared + .inject + .close(&mut self.shared.synced.lock().inject) + { self.notify_all(); } } - fn notify_parked(&self) { - if let Some(index) = self.shared.idle.worker_to_notify() { + 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_all(&self) { + 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); } @@ -831,13 +1122,13 @@ impl Handle { fn notify_if_work_pending(&self) { for remote in &self.shared.remotes[..] { if !remote.steal.is_empty() { - self.notify_parked(); + self.notify_parked_local(); return; } } if !self.shared.inject.is_empty() { - self.notify_parked(); + self.notify_parked_local(); } } @@ -845,7 +1136,7 @@ impl Handle { 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(); + self.notify_parked_local(); } } @@ -870,7 +1161,7 @@ impl Handle { // Drain the injection queue // // We already shut down every task, so we can simply drop the tasks. - while let Some(task) = self.shared.inject.pop() { + while let Some(task) = self.next_remote_task() { drop(task); } } @@ -880,22 +1171,56 @@ impl Handle { } } -fn did_defer_tasks() -> bool { - context::with_defer(|deferred| !deferred.is_empty()).unwrap() +impl Overflow<Arc<Handle>> for Handle { + fn push(&self, task: task::Notified<Arc<Handle>>) { + self.push_remote_task(task); + } + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<Arc<Handle>>>, + { + unsafe { + self.shared.inject.push_batch(self, iter); + } + } } -fn wake_deferred_tasks() { - context::with_defer(|deferred| deferred.wake()); +pub(crate) struct InjectGuard<'a> { + lock: crate::loom::sync::MutexGuard<'a, Synced>, } -cfg_metrics! { - impl Shared { - pub(super) fn injection_queue_depth(&self) -> usize { - self.inject.len() - } +impl<'a> AsMut<inject::Synced> for InjectGuard<'a> { + fn as_mut(&mut self) -> &mut inject::Synced { + &mut self.lock.inject + } +} + +impl<'a> Lock<inject::Synced> for &'a Handle { + type Handle = InjectGuard<'a>; - pub(super) fn worker_local_queue_depth(&self, worker: usize) -> usize { - self.remotes[worker].steal.len() + fn lock(self) -> Self::Handle { + InjectGuard { + lock: self.shared.synced.lock(), } } } + +#[track_caller] +fn with_current<R>(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), + }) +} + +// `u32::abs_diff` is not available on Tokio's MSRV. +fn abs_diff(a: u32, b: u32) -> u32 { + if a > b { + a - b + } else { + b - a + } +} diff --git a/src/runtime/scheduler/multi_thread/worker/metrics.rs b/src/runtime/scheduler/multi_thread/worker/metrics.rs new file mode 100644 index 0000000..a9a5ab3 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/worker/metrics.rs @@ -0,0 +1,11 @@ +use 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() + } +} diff --git a/src/runtime/scheduler/multi_thread/worker/taskdump.rs b/src/runtime/scheduler/multi_thread/worker/taskdump.rs new file mode 100644 index 0000000..d310d9f --- /dev/null +++ b/src/runtime/scheduler/multi_thread/worker/taskdump.rs @@ -0,0 +1,79 @@ +use 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<Core>) -> Box<Core> { + 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<Arc<Handle>> { + 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 + } +} diff --git a/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs b/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs new file mode 100644 index 0000000..24c5600 --- /dev/null +++ b/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs @@ -0,0 +1,7 @@ +use super::{Core, Handle}; + +impl Handle { + pub(super) fn trace_core(&self, core: Box<Core>) -> Box<Core> { + core + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/counters.rs b/src/runtime/scheduler/multi_thread_alt/counters.rs new file mode 100644 index 0000000..edda0d4 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/counters.rs @@ -0,0 +1,166 @@ +#[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::*; diff --git a/src/runtime/scheduler/multi_thread_alt/handle.rs b/src/runtime/scheduler/multi_thread_alt/handle.rs new file mode 100644 index 0000000..d746bca --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/handle.rs @@ -0,0 +1,75 @@ +use crate::future::Future; +use crate::loom::sync::Arc; +use crate::runtime::scheduler::multi_thread_alt::worker; +use crate::runtime::{ + blocking, driver, + task::{self, JoinHandle}, +}; +use crate::util::RngSeedGenerator; + +use std::fmt; + +cfg_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, +} + +impl Handle { + /// Spawns a future onto the thread pool + pub(crate) fn spawn<F>(me: &Arc<Self>, future: F, id: task::Id) -> JoinHandle<F::Output> + 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<T>(me: &Arc<Self>, future: T, id: task::Id) -> JoinHandle<T::Output> + where + T: Future + Send + 'static, + T::Output: Send + 'static, + { + let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); + + 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() + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/handle/metrics.rs b/src/runtime/scheduler/multi_thread_alt/handle/metrics.rs new file mode 100644 index 0000000..838694f --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/handle/metrics.rs @@ -0,0 +1,41 @@ +use 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 { + self.blocking_spawner.num_threads() + } + + pub(crate) fn num_idle_blocking_threads(&self) -> usize { + self.blocking_spawner.num_idle_threads() + } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.shared.owned.active_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() + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/handle/taskdump.rs b/src/runtime/scheduler/multi_thread_alt/handle/taskdump.rs new file mode 100644 index 0000000..477d857 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/handle/taskdump.rs @@ -0,0 +1,26 @@ +use 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 + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/idle.rs b/src/runtime/scheduler/multi_thread_alt/idle.rs new file mode 100644 index 0000000..ae2fb8b --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/idle.rs @@ -0,0 +1,423 @@ +//! 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<AtomicUsize>, +} + +pub(super) struct Snapshot { + chunks: Vec<usize>, +} + +/// Data synchronized by the scheduler mutex +pub(super) struct Synced { + /// Worker IDs that are currently sleeping + sleepers: Vec<usize>, + + /// Cores available for workers + available_cores: Vec<Box<Core>>, +} + +impl Idle { + pub(super) fn new(cores: Vec<Box<Core>>, 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<Box<Core>> { + 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<usize>, + 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<Core>) { + // 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<Core>]) -> 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<Core>]) { + 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<Core>]) -> 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() +} diff --git a/src/runtime/scheduler/multi_thread_alt/mod.rs b/src/runtime/scheduler/multi_thread_alt/mod.rs new file mode 100644 index 0000000..e30c9b4 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/mod.rs @@ -0,0 +1,91 @@ +//! 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<F>(&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() + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/overflow.rs b/src/runtime/scheduler/multi_thread_alt/overflow.rs new file mode 100644 index 0000000..ab66481 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/overflow.rs @@ -0,0 +1,26 @@ +use crate::runtime::task; + +#[cfg(test)] +use std::cell::RefCell; + +pub(crate) trait Overflow<T: 'static> { + fn push(&self, task: task::Notified<T>); + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<T>>; +} + +#[cfg(test)] +impl<T: 'static> Overflow<T> for RefCell<Vec<task::Notified<T>>> { + fn push(&self, task: task::Notified<T>) { + self.borrow_mut().push(task); + } + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<T>>, + { + self.borrow_mut().extend(iter); + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/park.rs b/src/runtime/scheduler/multi_thread_alt/park.rs new file mode 100644 index 0000000..0a00ea0 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/park.rs @@ -0,0 +1,232 @@ +//! 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<Inner>, +} + +pub(crate) struct Unparker { + inner: Arc<Inner>, +} + +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<Shared>, +} + +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<Driver>, +} + +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(); + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/queue.rs b/src/runtime/scheduler/multi_thread_alt/queue.rs new file mode 100644 index 0000000..e41e9fd --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/queue.rs @@ -0,0 +1,595 @@ +//! 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: <https://github.com/tokio-rs/tokio/issues/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<T: 'static> { + inner: Arc<Inner<T>>, +} + +/// Consumer handle. May be used from many threads. +pub(crate) struct Steal<T: 'static>(Arc<Inner<T>>); + +#[repr(align(128))] +pub(crate) struct Inner<T: 'static> { + /// 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<MaybeUninit<task::Notified<T>>>]>, + + mask: usize, +} + +unsafe impl<T> Send for Inner<T> {} +unsafe impl<T> Sync for Inner<T> {} + +/// Create a new local run-queue +pub(crate) fn local<T: 'static>(capacity: usize) -> (Steal<T>, Local<T>) { + 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<T> Local<T> { + /// 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<Item = task::Notified<T>>) { + 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 curent 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<O: Overflow<T>>( + &mut self, + mut task: task::Notified<T>, + 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<T>, 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<O: Overflow<T>>( + &mut self, + task: task::Notified<T>, + head: UnsignedShort, + tail: UnsignedShort, + overflow: &O, + stats: &mut Stats, + ) -> Result<(), task::Notified<T>> { + // 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<MaybeUninit<task::Notified<T>>>], + mask: usize, + head: UnsignedLong, + i: UnsignedLong, + num: UnsignedShort, + } + impl<'a, T: 'static> Iterator for BatchTaskIter<'a, T> { + type Item = task::Notified<T>; + + #[inline] + fn next(&mut self) -> Option<task::Notified<T>> { + 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<task::Notified<T>> { + 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<T> Steal<T> { + 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<T>, + dst_stats: &mut Stats, + ) -> Option<task::Notified<T>> { + // 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<T>, 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_metrics! { + impl<T> Steal<T> { + pub(crate) fn len(&self) -> usize { + self.0.len() as _ + } + } +} + +impl<T> Clone for Steal<T> { + fn clone(&self) -> Steal<T> { + Steal(self.0.clone()) + } +} + +impl<T> Drop for Local<T> { + fn drop(&mut self) { + if !std::thread::panicking() { + assert!(self.pop().is_none(), "queue not empty"); + } + } +} + +impl<T> Inner<T> { + 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::<UnsignedShort>() * 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::<UnsignedShort>() * 8)) +} diff --git a/src/runtime/scheduler/multi_thread_alt/stats.rs b/src/runtime/scheduler/multi_thread_alt/stats.rs new file mode 100644 index 0000000..228e797 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/stats.rs @@ -0,0 +1,171 @@ +use 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( + // 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 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(); + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/trace.rs b/src/runtime/scheduler/multi_thread_alt/trace.rs new file mode 100644 index 0000000..cc65a48 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/trace.rs @@ -0,0 +1,61 @@ +use 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<Option<Dump>>, +} + +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<Dump> { + 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; + } + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/trace_mock.rs b/src/runtime/scheduler/multi_thread_alt/trace_mock.rs new file mode 100644 index 0000000..2c17a4e --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/trace_mock.rs @@ -0,0 +1,11 @@ +pub(super) struct TraceStatus {} + +impl TraceStatus { + pub(super) fn new(_: usize) -> Self { + Self {} + } + + pub(super) fn trace_requested(&self) -> bool { + false + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/worker.rs b/src/runtime/scheduler/multi_thread_alt/worker.rs new file mode 100644 index 0000000..c9a0555 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/worker.rs @@ -0,0 +1,1595 @@ +//! 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::context; +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; +use crate::runtime::{blocking, coop, driver, task, Config, SchedulerMetrics, WorkerMetrics}; +use crate::util::atomic_cell::AtomicCell; +use crate::util::rand::{FastRand, RngSeedGenerator}; + +use std::cell::{Cell, RefCell}; +use std::cmp; +use std::task::Waker; +use std::time::Duration; + +cfg_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<usize>, + + /// 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<Notified>, + + /// The worker-local run queue. + run_queue: queue::Local<Arc<Handle>>, + + /// 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<Arc<Handle>>, + + /// Coordinates idle workers + idle: Idle, + + /// Collection of all active tasks spawned onto this executor. + pub(super) owned: OwnedTasks<Arc<Handle>>, + + /// Data synchronized by the scheduler mutex + pub(super) synced: Mutex<Synced>, + + /// Power's Tokio's I/O, timers, etc... the responsibility of polling the + /// driver is shared across workers. + driver: AtomicCell<Driver>, + + /// Condition variables used to unblock worker threads. Each worker thread + /// has its own condvar it waits on. + pub(super) condvars: Vec<Condvar>, + + /// 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<Option<Box<Core>>>, + + /// 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<Box<Core>>, + + /// The driver goes here when shutting down + shutdown_driver: Option<Box<Driver>>, + + /// 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<Arc<Handle>>, +} + +/// Thread-local context +pub(crate) struct Context { + // Current scheduler's handle + handle: Arc<Handle>, + + /// Worker index + index: usize, + + /// True when the LIFO slot is enabled + lifo_enabled: Cell<bool>, + + /// Core data + core: RefCell<Option<Box<Core>>>, + + /// Used to pass cores to other threads when `block_in_place` is called + handoff_core: Arc<AtomicCell<Core>>, + + /// Tasks to wake after resource drivers are polled. This is mostly to + /// handle yielded tasks. + pub(crate) defer: RefCell<Vec<Notified>>, +} + +/// 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<Box<Core>, ()>; +type NextTaskResult = Result<(Option<Notified>, Box<Core>), ()>; + +/// A task handle +type Task = task::Task<Arc<Handle>>; + +/// A notified task handle +type Notified = task::Notified<Arc<Handle>>; + +/// Value picked out of thin-air. Running the LIFO slot a handful of times +/// seemms 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 { + shared: Shared { + remotes: remotes.into_boxed_slice(), + inject, + idle, + owned: OwnedTasks::new(), + 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, R>(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<Handle>, + handoff_core: Arc<AtomicCell<Core>>, + 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 assinged to us + self.wait_for_core(cx, synced)? + } + } + }; + + 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<Box<Core>> { + 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); + + 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); + debug_assert!(self.global_queue_interval > 1); + + // 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<Core>) -> 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<Core>) -> 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<Notified> { + 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<Core>) -> 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<Notified> { + 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 immedietly + 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<Core>) -> 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<Notified> { + 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<Core>, 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<Core>, + 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<Core>) -> 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<Core>) -> 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<Core>) -> 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<Core>) -> 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<Core>) -> 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); + + debug_assert!(next > 1); + + // Smooth out jitter + if 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 + } +} + +impl Core { + fn next_local_task(&mut self) -> Option<Notified> { + self.next_lifo_task().or_else(|| self.run_queue.pop()) + } + + fn next_lifo_task(&mut self) -> Option<Notified> { + self.lifo_slot.take() + } +} + +impl Shared { + fn next_remote_task_synced(&self, synced: &mut Synced) -> Option<Notified> { + // 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<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<Arc<Handle>>>, + { + unsafe { + self.inject.push_batch(self, iter); + } + } + + fn push_remote_task_batch_synced<I>(&self, synced: &mut Synced, iter: I) + where + I: Iterator<Item = task::Notified<Arc<Handle>>>, + { + unsafe { + self.inject.push_batch(&mut synced.inject, iter); + } + } + + fn take_driver(&self) -> Option<Box<Driver>> { + 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<Core>) { + self.owned.close_and_shutdown_all(); + + 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<Arc<Handle>> for Shared { + fn push(&self, task: task::Notified<Arc<Handle>>) { + self.push_remote_task(&mut self.synced.lock(), task); + } + + fn push_batch<I>(&self, iter: I) + where + I: Iterator<Item = task::Notified<Arc<Handle>>>, + { + self.push_remote_task_batch(iter) + } +} + +impl<'a> Lock<inject::Synced> for &'a Shared { + type Handle = SyncedGuard<'a>; + + fn lock(self) -> Self::Handle { + SyncedGuard { + lock: self.synced.lock(), + } + } +} + +impl<'a> Lock<Synced> for &'a Shared { + type Handle = SyncedGuard<'a>; + + fn lock(self) -> Self::Handle { + SyncedGuard { + lock: self.synced.lock(), + } + } +} + +impl task::Schedule for Arc<Handle> { + fn release(&self, task: &Task) -> Option<Task> { + self.shared.owned.remove(task) + } + + fn schedule(&self, task: Notified) { + self.shared.schedule_task(task, false); + } + + fn yield_now(&self, task: Notified) { + self.shared.schedule_task(task, true); + } +} + +impl AsMut<Synced> 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<inject::Synced> for SyncedGuard<'a> { + fn as_mut(&mut self) -> &mut inject::Synced { + &mut self.lock.inject + } +} + +impl<'a> AsMut<Synced> for SyncedGuard<'a> { + fn as_mut(&mut self) -> &mut Synced { + &mut self.lock + } +} + +#[track_caller] +fn with_current<R>(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), + }) +} + +// `u32::abs_diff` is not available on Tokio's MSRV. +fn abs_diff(a: u32, b: u32) -> u32 { + if a > b { + a - b + } else { + b - a + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/worker/metrics.rs b/src/runtime/scheduler/multi_thread_alt/worker/metrics.rs new file mode 100644 index 0000000..a9a5ab3 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/worker/metrics.rs @@ -0,0 +1,11 @@ +use 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() + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/worker/taskdump.rs b/src/runtime/scheduler/multi_thread_alt/worker/taskdump.rs new file mode 100644 index 0000000..7cf69c4 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/worker/taskdump.rs @@ -0,0 +1,79 @@ +use 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<Core>) -> Box<Core> { + 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<Arc<Handle>> { + 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 + } +} diff --git a/src/runtime/scheduler/multi_thread_alt/worker/taskdump_mock.rs b/src/runtime/scheduler/multi_thread_alt/worker/taskdump_mock.rs new file mode 100644 index 0000000..24c5600 --- /dev/null +++ b/src/runtime/scheduler/multi_thread_alt/worker/taskdump_mock.rs @@ -0,0 +1,7 @@ +use super::{Core, Handle}; + +impl Handle { + pub(super) fn trace_core(&self, core: Box<Core>) -> Box<Core> { + core + } +} diff --git a/src/runtime/task/core.rs b/src/runtime/task/core.rs index bcccc69..d62ea96 100644 --- a/src/runtime/task/core.rs +++ b/src/runtime/task/core.rs @@ -17,6 +17,7 @@ use crate::runtime::task::state::State; use crate::runtime::task::{Id, Schedule}; use crate::util::linked_list; +use std::num::NonZeroU64; use std::pin::Pin; use std::ptr::NonNull; use std::task::{Context, Poll, Waker}; @@ -28,6 +29,94 @@ use std::task::{Context, Poll, Waker}; /// /// 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, riscv64, 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/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_riscv64.go#L7 +// - 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 +// +// riscv32 is assumed not to exceed the cache line size of riscv64. +#[cfg_attr( + any( + target_arch = "arm", + target_arch = "mips", + target_arch = "mips64", + target_arch = "riscv32", + target_arch = "riscv64", + 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, 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 +// +// 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 = "riscv32", + target_arch = "riscv64", + target_arch = "sparc", + target_arch = "hexagon", + target_arch = "m68k", + target_arch = "s390x", + )), + repr(align(64)) +)] #[repr(C)] pub(super) struct Cell<T: Future, S> { /// Hot task state data @@ -76,7 +165,7 @@ pub(crate) struct Header { /// 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 zero if it has never been + /// 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 @@ -85,7 +174,7 @@ pub(crate) struct Header { /// 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<u64>, + pub(super) owner_id: UnsafeCell<Option<NonZeroU64>>, /// The tracing ID for this instrumented task. #[cfg(all(tokio_unstable, feature = "tracing"))] @@ -123,17 +212,32 @@ impl<T: Future, S: Schedule> Cell<T, S> { /// 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<Cell<T, S>> { + // 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<tracing::Id>, + ) -> 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::<T, S>(); let result = Box::new(Cell { - header: Header { + header: new_header( state, - queue_next: UnsafeCell::new(None), - vtable: raw::vtable::<T, S>(), - owner_id: UnsafeCell::new(0), + vtable, #[cfg(all(tokio_unstable, feature = "tracing"))] tracing_id, - }, + ), core: Core { scheduler, stage: CoreStage { @@ -141,26 +245,33 @@ impl<T: Future, S: Schedule> Cell<T, S> { }, task_id, }, - trailer: Trailer { - waker: UnsafeCell::new(None), - owned: linked_list::Pointers::new(), - }, + trailer: Trailer::new(), }); #[cfg(debug_assertions)] { - let trailer_addr = (&result.trailer) as *const Trailer as usize; - let trailer_ptr = unsafe { Header::get_trailer(NonNull::from(&result.header)) }; - assert_eq!(trailer_addr, trailer_ptr.as_ptr() as usize); - - let scheduler_addr = (&result.core.scheduler) as *const S as usize; - let scheduler_ptr = - unsafe { Header::get_scheduler::<S>(NonNull::from(&result.header)) }; - assert_eq!(scheduler_addr, scheduler_ptr.as_ptr() as usize); - - let id_addr = (&result.core.task_id) as *const Id as usize; - let id_ptr = unsafe { Header::get_id_ptr(NonNull::from(&result.header)) }; - assert_eq!(id_addr, id_ptr.as_ptr() as usize); + // Using a separate function for this code avoids instantiating it separately for every `T`. + unsafe fn check<S>(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::<S>(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 @@ -278,23 +389,19 @@ impl<T: Future, S: Schedule> Core<T, S> { } } -cfg_rt_multi_thread! { - impl Header { - pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) { - self.queue_next.with_mut(|ptr| *ptr = next); - } +impl Header { + pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) { + self.queue_next.with_mut(|ptr| *ptr = next); } -} -impl Header { // safety: The caller must guarantee exclusive access to this field, and - // must ensure that the id is either 0 or the id of the OwnedTasks + // 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: u64) { - self.owner_id.with_mut(|ptr| *ptr = owner); + 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) -> u64 { + pub(super) fn get_owner_id(&self) -> Option<NonZeroU64> { // 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) } @@ -358,6 +465,13 @@ impl Header { } impl Trailer { + fn new() -> Self { + Trailer { + waker: UnsafeCell::new(None), + owned: linked_list::Pointers::new(), + } + } + pub(super) unsafe fn set_waker(&self, waker: Option<Waker>) { self.waker.with_mut(|ptr| { *ptr = waker; diff --git a/src/runtime/task/harness.rs b/src/runtime/task/harness.rs index 8e3c3d1..8bfd57e 100644 --- a/src/runtime/task/harness.rs +++ b/src/runtime/task/harness.rs @@ -2,8 +2,9 @@ use 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::{JoinError, Notified, RawTask, Schedule, Task}; +use crate::runtime::task::{Id, JoinError, Notified, RawTask, Schedule, Task}; +use std::any::Any; use std::mem; use std::mem::ManuallyDrop; use std::panic; @@ -192,8 +193,17 @@ where 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::<T, S>(&header_ptr); + let waker_ref = waker_ref::<S>(&header_ptr); let cx = Context::from_waker(&waker_ref); let res = poll_future(self.core(), cx); @@ -202,17 +212,13 @@ where return PollFuture::Complete; } - match self.state().transition_to_idle() { - TransitionToIdle::Ok => PollFuture::Done, - TransitionToIdle::OkNotified => PollFuture::Notified, - TransitionToIdle::OkDealloc => PollFuture::Dealloc, - TransitionToIdle::Cancelled => { - // The transition to idle failed because the task was - // cancelled during the poll. - cancel_task(self.core()); - 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()); @@ -447,13 +453,16 @@ fn cancel_task<T: Future, S: Schedule>(core: &Core<T, S>) { 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<dyn Any + Send + 'static>>, +) -> JoinError { match res { - Ok(()) => { - core.store_output(Err(JoinError::cancelled(core.task_id))); - } - Err(panic) => { - core.store_output(Err(JoinError::panic(core.task_id, panic))); - } + Ok(()) => JoinError::cancelled(task_id), + Err(panic) => JoinError::panic(task_id, panic), } } @@ -482,10 +491,7 @@ fn poll_future<T: Future, S: Schedule>(core: &Core<T, S>, cx: Context<'_>) -> Po let output = match output { Ok(Poll::Pending) => return Poll::Pending, Ok(Poll::Ready(output)) => Ok(output), - Err(panic) => { - core.scheduler.unhandled_panic(); - Err(JoinError::panic(core.task_id, panic)) - } + Err(panic) => Err(panic_to_error(&core.scheduler, core.task_id, panic)), }; // Catch and ignore panics if the future panics on drop. @@ -499,3 +505,13 @@ fn poll_future<T: Future, S: Schedule>(core: &Core<T, S>, cx: Context<'_>) -> Po Poll::Ready(()) } + +#[cold] +fn panic_to_error<S: Schedule>( + scheduler: &S, + task_id: Id, + panic: Box<dyn Any + Send + 'static>, +) -> JoinError { + scheduler.unhandled_panic(); + JoinError::panic(task_id, panic) +} diff --git a/src/runtime/task/inject.rs b/src/runtime/task/inject.rs deleted file mode 100644 index 1585e13..0000000 --- a/src/runtime/task/inject.rs +++ /dev/null @@ -1,220 +0,0 @@ -//! Inject queue used to send wakeups to a work-stealing scheduler - -use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::Mutex; -use crate::runtime::task; - -use std::marker::PhantomData; -use std::ptr::NonNull; -use std::sync::atomic::Ordering::{Acquire, Release}; - -/// 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<T: 'static> { - /// Pointers to the head and tail of the queue. - pointers: Mutex<Pointers>, - - /// Number of pending tasks in the queue. This helps prevent unnecessary - /// locking in the hot path. - len: AtomicUsize, - - _p: PhantomData<T>, -} - -struct Pointers { - /// True if the queue is closed. - is_closed: bool, - - /// Linked-list head. - head: Option<NonNull<task::Header>>, - - /// Linked-list tail. - tail: Option<NonNull<task::Header>>, -} - -unsafe impl<T> Send for Inject<T> {} -unsafe impl<T> Sync for Inject<T> {} - -impl<T: 'static> Inject<T> { - pub(crate) fn new() -> Inject<T> { - Inject { - pointers: Mutex::new(Pointers { - is_closed: false, - head: None, - tail: None, - }), - len: AtomicUsize::new(0), - _p: PhantomData, - } - } - - pub(crate) fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Closes the injection queue, returns `true` if the queue is open when the - /// transition is made. - pub(crate) fn close(&self) -> bool { - let mut p = self.pointers.lock(); - - if p.is_closed { - return false; - } - - p.is_closed = true; - true - } - - pub(crate) fn is_closed(&self) -> bool { - self.pointers.lock().is_closed - } - - pub(crate) fn len(&self) -> usize { - self.len.load(Acquire) - } - - /// Pushes a value into the queue. - /// - /// This does nothing if the queue is closed. - pub(crate) fn push(&self, task: task::Notified<T>) { - // Acquire queue lock - let mut p = self.pointers.lock(); - - if p.is_closed { - return; - } - - // safety: only mutated with the lock held - let len = unsafe { self.len.unsync_load() }; - let task = task.into_raw(); - - // The next pointer should already be null - debug_assert!(get_next(task).is_none()); - - if let Some(tail) = p.tail { - // safety: Holding the Notified for a task guarantees exclusive - // access to the `queue_next` field. - set_next(tail, Some(task)); - } else { - p.head = Some(task); - } - - p.tail = Some(task); - - self.len.store(len + 1, Release); - } - - /// Pushes several values into the queue. - #[inline] - pub(crate) fn push_batch<I>(&self, mut iter: I) - where - I: Iterator<Item = task::Notified<T>>, - { - 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. - set_next(prev, Some(next)); - prev = next; - counter += 1; - }); - - // Now that the tasks are linked together, insert them into the - // linked list. - self.push_batch_inner(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] - fn push_batch_inner( - &self, - batch_head: NonNull<task::Header>, - batch_tail: NonNull<task::Header>, - num: usize, - ) { - debug_assert!(get_next(batch_tail).is_none()); - - let mut p = self.pointers.lock(); - - if let Some(tail) = p.tail { - set_next(tail, Some(batch_head)); - } else { - p.head = Some(batch_head); - } - - p.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 = unsafe { self.len.unsync_load() }; - - self.len.store(len + num, Release); - } - - pub(crate) fn pop(&self) -> Option<task::Notified<T>> { - // Fast path, if len == 0, then there are no values - if self.is_empty() { - return None; - } - - let mut p = self.pointers.lock(); - - // It is possible to hit null here if another thread popped the last - // task between us checking `len` and acquiring the lock. - let task = p.head?; - - p.head = get_next(task); - - if p.head.is_none() { - p.tail = None; - } - - set_next(task, None); - - // Decrement 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. - self.len - .store(unsafe { self.len.unsync_load() } - 1, Release); - - // safety: a `Notified` is pushed into the queue and now it is popped! - Some(unsafe { task::Notified::from_raw(task) }) - } -} - -impl<T: 'static> Drop for Inject<T> { - fn drop(&mut self) { - if !std::thread::panicking() { - assert!(self.pop().is_none(), "queue not empty"); - } - } -} - -fn get_next(header: NonNull<task::Header>) -> Option<NonNull<task::Header>> { - unsafe { header.as_ref().queue_next.with(|ptr| *ptr) } -} - -fn set_next(header: NonNull<task::Header>, val: Option<NonNull<task::Header>>) { - unsafe { - header.as_ref().set_next(val); - } -} diff --git a/src/runtime/task/join.rs b/src/runtime/task/join.rs index 11c4b9b..ee39258 100644 --- a/src/runtime/task/join.rs +++ b/src/runtime/task/join.rs @@ -182,28 +182,28 @@ impl<T> JoinHandle<T> { /// ```rust /// use tokio::time; /// - /// #[tokio::main] - /// 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()); - /// } + /// # #[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 pub fn abort(&self) { @@ -220,9 +220,8 @@ impl<T> JoinHandle<T> { /// ```rust /// use tokio::time; /// - /// # #[tokio::main(flavor = "current_thread")] + /// # #[tokio::main(flavor = "current_thread", start_paused = true)] /// # async fn main() { - /// # time::pause(); /// let handle1 = tokio::spawn(async { /// // do some stuff here /// }); @@ -252,7 +251,41 @@ impl<T> JoinHandle<T> { } /// Returns a new `AbortHandle` that can be used to remotely abort this task. - pub(crate) fn abort_handle(&self) -> super::AbortHandle { + /// + /// 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<task::AbortHandle> = 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 + pub fn abort_handle(&self) -> super::AbortHandle { self.raw.ref_inc(); super::AbortHandle::new(self.raw) } @@ -280,6 +313,7 @@ impl<T> Future for JoinHandle<T> { type Output = super::Result<T>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + ready!(crate::trace::trace_leaf(cx)); let mut ret = Poll::Pending; // Keep track of task budget diff --git a/src/runtime/task/list.rs b/src/runtime/task/list.rs index 159c13e..3a1fcce 100644 --- a/src/runtime/task/list.rs +++ b/src/runtime/task/list.rs @@ -10,13 +10,14 @@ use crate::future::Future; use crate::loom::cell::UnsafeCell; use crate::loom::sync::Mutex; use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task}; -use crate::util::linked_list::{Link, LinkedList}; +use crate::util::linked_list::{CountedLinkedList, Link, LinkedList}; 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 zero for tasks not owned by any list. +// 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 @@ -28,10 +29,10 @@ cfg_has_atomic_u64! { static NEXT_OWNED_TASKS_ID: AtomicU64 = AtomicU64::new(1); - fn get_next_id() -> u64 { + fn get_next_id() -> NonZeroU64 { loop { let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed); - if id != 0 { + if let Some(id) = NonZeroU64::new(id) { return id; } } @@ -43,23 +44,27 @@ cfg_not_has_atomic_u64! { static NEXT_OWNED_TASKS_ID: AtomicU32 = AtomicU32::new(1); - fn get_next_id() -> u64 { + fn get_next_id() -> NonZeroU64 { loop { let id = NEXT_OWNED_TASKS_ID.fetch_add(1, Ordering::Relaxed); - if id != 0 { - return u64::from(id); + if let Some(id) = NonZeroU64::new(u64::from(id)) { + return id; } } } } pub(crate) struct OwnedTasks<S: 'static> { - inner: Mutex<OwnedTasksInner<S>>, - id: u64, + inner: Mutex<CountedOwnedTasksInner<S>>, + pub(crate) id: NonZeroU64, +} +struct CountedOwnedTasksInner<S: 'static> { + list: CountedLinkedList<Task<S>, <Task<S> as Link>::Target>, + closed: bool, } pub(crate) struct LocalOwnedTasks<S: 'static> { inner: UnsafeCell<OwnedTasksInner<S>>, - id: u64, + pub(crate) id: NonZeroU64, _not_send_or_sync: PhantomData<*const ()>, } struct OwnedTasksInner<S: 'static> { @@ -70,8 +75,8 @@ struct OwnedTasksInner<S: 'static> { impl<S: 'static> OwnedTasks<S> { pub(crate) fn new() -> Self { Self { - inner: Mutex::new(OwnedTasksInner { - list: LinkedList::new(), + inner: Mutex::new(CountedOwnedTasksInner { + list: CountedLinkedList::new(), closed: false, }), id: get_next_id(), @@ -92,7 +97,15 @@ impl<S: 'static> OwnedTasks<S> { T::Output: Send + '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<S>, notified: Notified<S>) -> Option<Notified<S>> + where + S: Schedule, + { unsafe { // safety: We just created the task, so we have exclusive access // to the field. @@ -104,10 +117,10 @@ impl<S: 'static> OwnedTasks<S> { drop(lock); drop(notified); task.shutdown(); - (join, None) + None } else { lock.list.push_front(task); - (join, Some(notified)) + Some(notified) } } @@ -115,7 +128,7 @@ impl<S: 'static> OwnedTasks<S> { /// a LocalNotified, giving the thread permission to poll this task. #[inline] pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> { - assert_eq!(task.header().get_owner_id(), self.id); + 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. @@ -153,12 +166,14 @@ impl<S: 'static> OwnedTasks<S> { } } + pub(crate) fn active_tasks_count(&self) -> usize { + self.inner.lock().list.count() + } + pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> { - let task_id = task.header().get_owner_id(); - if task_id == 0 { - // The task is unowned. - return None; - } + // 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); @@ -172,6 +187,18 @@ impl<S: 'static> OwnedTasks<S> { } } +cfg_taskdump! { + impl<S: 'static> OwnedTasks<S> { + /// Locks the tasks, and calls `f` on an iterator over them. + pub(crate) fn for_each<F>(&self, f: F) + where + F: FnMut(&Task<S>) + { + self.inner.lock().list.for_each(f) + } + } +} + impl<S: 'static> LocalOwnedTasks<S> { pub(crate) fn new() -> Self { Self { @@ -229,11 +256,9 @@ impl<S: 'static> LocalOwnedTasks<S> { } pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> { - let task_id = task.header().get_owner_id(); - if task_id == 0 { - // The task is unowned. - return None; - } + // 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); @@ -247,7 +272,7 @@ impl<S: 'static> LocalOwnedTasks<S> { /// it to a LocalNotified, giving the thread permission to poll this task. #[inline] pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> { - assert_eq!(task.header().get_owner_id(), self.id); + 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 @@ -278,7 +303,7 @@ impl<S: 'static> LocalOwnedTasks<S> { } } -#[cfg(all(test))] +#[cfg(test)] mod tests { use super::*; @@ -287,11 +312,9 @@ mod tests { #[test] fn test_id_not_broken() { let mut last_id = get_next_id(); - assert_ne!(last_id, 0); for _ in 0..1000 { let next_id = get_next_id(); - assert_ne!(next_id, 0); assert!(last_id < next_id); last_id = next_id; } diff --git a/src/runtime/task/mod.rs b/src/runtime/task/mod.rs index 55131ac..abf7cc2 100644 --- a/src/runtime/task/mod.rs +++ b/src/runtime/task/mod.rs @@ -182,11 +182,6 @@ mod id; #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] pub use id::{id, try_id, Id}; -cfg_rt_multi_thread! { - mod inject; - pub(super) use self::inject::Inject; -} - #[cfg(feature = "rt")] mod abort; mod join; @@ -200,13 +195,17 @@ mod list; pub(crate) use self::list::{LocalOwnedTasks, OwnedTasks}; mod raw; -use self::raw::RawTask; +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; @@ -333,13 +332,28 @@ cfg_rt! { } impl<S: 'static> Task<S> { - unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { + unsafe fn new(raw: RawTask) -> Task<S> { Task { - raw: RawTask::from_raw(ptr), + raw, _p: PhantomData, } } + unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { + 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() } @@ -355,25 +369,17 @@ impl<S: 'static> Notified<S> { } } -cfg_rt_multi_thread! { - impl<S: 'static> Notified<S> { - unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> { - Notified(Task::from_raw(ptr)) - } - } - - impl<S: 'static> Task<S> { - fn into_raw(self) -> NonNull<Header> { - let ret = self.raw.header_ptr(); - mem::forget(self); - ret - } +impl<S: 'static> Notified<S> { + pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> { + Notified(Task::new(ptr)) } +} - impl<S: 'static> Notified<S> { - fn into_raw(self) -> NonNull<Header> { - self.0.into_raw() - } +impl<S: 'static> Notified<S> { + pub(crate) fn into_raw(self) -> RawTask { + let raw = self.0.raw; + mem::forget(self); + raw } } @@ -398,7 +404,7 @@ impl<S: Schedule> LocalNotified<S> { impl<S: Schedule> UnownedTask<S> { // Used in test of the inject queue. #[cfg(test)] - #[cfg_attr(tokio_wasm, allow(dead_code))] + #[cfg_attr(target_family = "wasm", allow(dead_code))] pub(super) fn into_notified(self) -> Notified<S> { Notified(self.into_task()) } diff --git a/src/runtime/task/raw.rs b/src/runtime/task/raw.rs index b9700ae..8078859 100644 --- a/src/runtime/task/raw.rs +++ b/src/runtime/task/raw.rs @@ -6,7 +6,7 @@ use std::ptr::NonNull; use std::task::{Poll, Waker}; /// Raw task handle -pub(super) struct RawTask { +pub(crate) struct RawTask { ptr: NonNull<Header>, } @@ -195,7 +195,7 @@ impl RawTask { } /// Safety: mutual exclusion is required to call this function. - pub(super) fn poll(self) { + pub(crate) fn poll(self) { let vtable = self.header().vtable; unsafe { (vtable.poll)(self.ptr) } } @@ -240,6 +240,27 @@ impl RawTask { 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<RawTask> { + 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<RawTask>) { + self.header().set_next(val.map(|task| task.ptr)); + } } impl Clone for RawTask { diff --git a/src/runtime/task/state.rs b/src/runtime/task/state.rs index 7728312..12f5449 100644 --- a/src/runtime/task/state.rs +++ b/src/runtime/task/state.rs @@ -88,7 +88,7 @@ pub(super) enum TransitionToNotifiedByVal { } #[must_use] -pub(super) enum TransitionToNotifiedByRef { +pub(crate) enum TransitionToNotifiedByRef { DoNothing, Submit, } @@ -270,6 +270,22 @@ impl State { }) } + /// Transitions the state to `NOTIFIED`, unconditionally increasing the ref count. + #[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) { + self.fetch_update_action(|mut snapshot| { + snapshot.set_notified(); + snapshot.ref_inc(); + ((), 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 diff --git a/src/runtime/task/trace/mod.rs b/src/runtime/task/trace/mod.rs new file mode 100644 index 0000000..af2c564 --- /dev/null +++ b/src/runtime/task/trace/mod.rs @@ -0,0 +1,336 @@ +use crate::loom::sync::Arc; +use crate::runtime::context; +use crate::runtime::scheduler::{self, current_thread, Inject}; + +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}; + +type Backtrace = Vec<BacktraceFrame>; +type SymbolTrace = Vec<Symbol>; + +/// The ambiant backtracing context. +pub(crate) struct Context { + /// The address of [`Trace::root`] establishes an upper unwinding bound on + /// the backtraces in `Trace`. + active_frame: Cell<Option<NonNull<Frame>>>, + /// The place to stash backtraces. + collector: Cell<Option<Trace>>, +} + +/// 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<NonNull<Frame>>, +} + +/// 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<Backtrace>, +} + +pin_project_lite::pin_project! { + #[derive(Debug, Clone)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub(crate) struct Root<T> { + #[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, R>(f: F) -> Option<R> + where + F: FnOnce(&Self) -> R, + { + crate::runtime::context::with_trace(f) + } + + unsafe fn with_current_frame<F, R>(f: F) -> R + where + F: FnOnce(&Cell<Option<NonNull<Frame>>>) -> R, + { + Self::try_with_current(|context| f(&context.active_frame)).expect(FAIL_NO_THREAD_LOCAL) + } + + fn with_current_collector<F, R>(f: F) -> R + where + F: FnOnce(&Cell<Option<Trace>>) -> 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) + } + } +} + +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, R>(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<F>(future: F) -> Root<F> { + Root { future } + } +} + +/// 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(all(feature = "rt-multi-thread", not(target_os = "wasi")))] + scheduler::Context::MultiThread(s) => s.defer.defer(cx.waker()), + #[cfg(all( + tokio_unstable, + feature = "rt-multi-thread", + not(target_os = "wasi") + ))] + 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<F: FnOnce() -> R, R>(f: F) -> impl Drop { + use std::mem::ManuallyDrop; + + struct Defer<F: FnOnce() -> R, R>(ManuallyDrop<F>); + + impl<F: FnOnce() -> R, R> Drop for Defer<F, R> { + #[inline(always)] + fn drop(&mut self) { + unsafe { + ManuallyDrop::take(&mut self.0)(); + } + } + } + + Defer(ManuallyDrop::new(f)) +} + +impl<T: Future> Future for Root<T> { + type Output = T::Output; + + #[inline(never)] + fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { + // 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<Arc<current_thread::Handle>>, + local: &mut VecDeque<Notified<Arc<current_thread::Handle>>>, + injection: &Inject<Arc<current_thread::Handle>>, +) -> Vec<Trace> { + // clear the local and injection queues + local.clear(); + + while let Some(task) = injection.pop() { + drop(task); + } + + // notify each task + let mut tasks = vec![]; + owned.for_each(|task| { + // set the notified bit + task.as_raw().state().transition_to_notified_for_tracing(); + // store the raw tasks into a vec + tasks.push(task.as_raw()); + }); + + tasks + .into_iter() + .map(|task| { + let ((), trace) = Trace::capture(|| task.poll()); + trace + }) + .collect() +} + +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<Arc<multi_thread::Handle>>, + local: &mut multi_thread::queue::Local<Arc<multi_thread::Handle>>, + synced: &Mutex<Synced>, + injection: &Shared<Arc<multi_thread::Handle>>, + ) -> Vec<Trace> { + // clear the local queue + while let Some(notified) = local.pop() { + drop(notified); + } + + // clear the injection queue + let mut synced = synced.lock(); + while let Some(notified) = injection.pop(&mut synced.inject) { + drop(notified); + } + + drop(synced); + + // notify each task + let mut traces = vec![]; + owned.for_each(|task| { + // set the notified bit + task.as_raw().state().transition_to_notified_for_tracing(); + + // trace the task + let ((), trace) = Trace::capture(|| task.as_raw().poll()); + traces.push(trace); + + // reschedule the task + let _ = task.as_raw().state().transition_to_notified_by_ref(); + task.as_raw().schedule(); + }); + + traces + } +} diff --git a/src/runtime/task/trace/symbol.rs b/src/runtime/task/trace/symbol.rs new file mode 100644 index 0000000..49d7ba3 --- /dev/null +++ b/src/runtime/task/trace/symbol.rs @@ -0,0 +1,92 @@ +use 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<H>(&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(()) + } +} diff --git a/src/runtime/task/trace/tree.rs b/src/runtime/task/trace/tree.rs new file mode 100644 index 0000000..7e6f8ef --- /dev/null +++ b/src/runtime/task/trace/tree.rs @@ -0,0 +1,126 @@ +use 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<Symbol>, + + /// The adjacency list of symbols in the execution tree(s). + edges: HashMap<Symbol, HashSet<Symbol>>, +} + +impl Tree { + /// Constructs a [`Tree`] from [`Trace`] + pub(super) fn from_trace(trace: Trace) -> Self { + let mut roots: HashSet<Symbol> = HashSet::default(); + let mut edges: HashMap<Symbol, HashSet<Symbol>> = 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<impl ExactSizeIterator<Item = &Symbol>> { + Some(self.edges.get(frame)?.iter()) + } + + /// Format this [`Tree`] as a textual tree. + fn display<W: fmt::Write>( + &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 +} diff --git a/src/runtime/task/waker.rs b/src/runtime/task/waker.rs index b5f5ace..2a1568f 100644 --- a/src/runtime/task/waker.rs +++ b/src/runtime/task/waker.rs @@ -1,4 +1,3 @@ -use crate::future::Future; use crate::runtime::task::{Header, RawTask, Schedule}; use std::marker::PhantomData; @@ -14,9 +13,8 @@ pub(super) struct WakerRef<'a, S: 'static> { /// Returns a `WakerRef` which avoids having to preemptively increase the /// refcount if there is no need to do so. -pub(super) fn waker_ref<T, S>(header: &NonNull<Header>) -> WakerRef<'_, S> +pub(super) fn waker_ref<S>(header: &NonNull<Header>) -> WakerRef<'_, S> where - T: Future, S: Schedule, { // `Waker::will_wake` uses the VTABLE pointer as part of the check. This diff --git a/src/runtime/tests/inject.rs b/src/runtime/tests/inject.rs new file mode 100644 index 0000000..ccead5e --- /dev/null +++ b/src/runtime/tests/inject.rs @@ -0,0 +1,54 @@ +use 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); + } +} diff --git a/src/runtime/tests/loom_current_thread_scheduler.rs b/src/runtime/tests/loom_current_thread.rs index a772603..edda6e4 100644 --- a/src/runtime/tests/loom_current_thread_scheduler.rs +++ b/src/runtime/tests/loom_current_thread.rs @@ -1,3 +1,5 @@ +mod yield_now; + use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Arc; use crate::loom::thread; diff --git a/src/runtime/tests/loom_yield.rs b/src/runtime/tests/loom_current_thread/yield_now.rs index ba506e5..ba506e5 100644 --- a/src/runtime/tests/loom_yield.rs +++ b/src/runtime/tests/loom_current_thread/yield_now.rs diff --git a/src/runtime/tests/loom_pool.rs b/src/runtime/tests/loom_multi_thread.rs index b3ecd43..c5980c2 100644 --- a/src/runtime/tests/loom_pool.rs +++ b/src/runtime/tests/loom_multi_thread.rs @@ -1,3 +1,7 @@ +mod queue; +mod shutdown; +mod yield_now; + /// Full runtime loom tests. These are heavy tests and take significant time to /// run on CI. /// @@ -323,8 +327,7 @@ mod group_d { // Spawn a task let c2 = c1.clone(); pool.spawn(track(async move { - gated().await; - gated().await; + multi_gated().await; if 1 == c1.fetch_add(1, Relaxed) { done_tx1.assert_send(()); @@ -333,8 +336,7 @@ mod group_d { // Spawn a second task pool.spawn(track(async move { - gated().await; - gated().await; + multi_gated().await; if 1 == c2.fetch_add(1, Relaxed) { done_tx2.assert_send(()); @@ -349,14 +351,12 @@ mod group_d { 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 gated() -> impl Future<Output = &'static str> { - gated2(false) -} - fn gated2(thread: bool) -> impl Future<Output = &'static str> { use loom::thread; use std::sync::Arc; @@ -394,6 +394,38 @@ fn gated2(thread: bool) -> impl Future<Output = &'static str> { }) } +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<T: Future>(f: T) -> Track<T> { Track { inner: f, diff --git a/src/runtime/tests/loom_queue.rs b/src/runtime/tests/loom_multi_thread/queue.rs index fc93bf3..0d81828 100644 --- a/src/runtime/tests/loom_queue.rs +++ b/src/runtime/tests/loom_multi_thread/queue.rs @@ -1,24 +1,27 @@ -use crate::runtime::scheduler::multi_thread::queue; -use crate::runtime::task::Inject; -use crate::runtime::tests::NoopSchedule; -use crate::runtime::MetricsBatch; +use 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 = Inject::new(); - let mut metrics = MetricsBatch::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); let th = thread::spawn(move || { - let mut metrics = MetricsBatch::new(); + 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 metrics).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -34,8 +37,8 @@ fn basic() { for _ in 0..2 { for _ in 0..2 { - let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + let (task, _) = unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } if local.pop().is_some() { @@ -43,17 +46,15 @@ fn basic() { } // Push another task - let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + let (task, _) = unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); while local.pop().is_some() { n += 1; } } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); @@ -65,15 +66,15 @@ fn basic() { fn steal_overflow() { loom::model(|| { let (steal, mut local) = queue::local(); - let inject = Inject::new(); - let mut metrics = MetricsBatch::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); let th = thread::spawn(move || { - let mut metrics = MetricsBatch::new(); + let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; - if steal.steal_into(&mut local, &mut metrics).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -87,16 +88,16 @@ fn steal_overflow() { let mut n = 0; // push a task, pop a task - let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + 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, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + let (task, _) = unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } n += th.join().unwrap(); @@ -105,9 +106,7 @@ fn steal_overflow() { n += 1; } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); assert_eq!(7, n); }); @@ -118,10 +117,10 @@ fn multi_stealer() { const NUM_TASKS: usize = 5; fn steal_tasks(steal: queue::Steal<NoopSchedule>) -> usize { - let mut metrics = MetricsBatch::new(); + let mut stats = new_stats(); let (_, mut local) = queue::local(); - if steal.steal_into(&mut local, &mut metrics).is_none() { + if steal.steal_into(&mut local, &mut stats).is_none() { return 0; } @@ -136,13 +135,13 @@ fn multi_stealer() { loom::model(|| { let (steal, mut local) = queue::local(); - let inject = Inject::new(); - let mut metrics = MetricsBatch::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); // Push work for _ in 0..NUM_TASKS { - let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + let (task, _) = unowned(async {}); + local.push_back_or_overflow(task, &inject, &mut stats); } let th1 = { @@ -158,9 +157,7 @@ fn multi_stealer() { n += 1; } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); n += th1.join().unwrap(); n += th2.join().unwrap(); @@ -172,25 +169,25 @@ fn multi_stealer() { #[test] fn chained_steal() { loom::model(|| { - let mut metrics = MetricsBatch::new(); + let mut stats = new_stats(); let (s1, mut l1) = queue::local(); let (s2, mut l2) = queue::local(); - let inject = Inject::new(); + let inject = RefCell::new(vec![]); // Load up some tasks for _ in 0..4 { - let (task, _) = super::unowned(async {}); - l1.push_back(task, &inject, &mut metrics); + let (task, _) = unowned(async {}); + l1.push_back_or_overflow(task, &inject, &mut stats); - let (task, _) = super::unowned(async {}); - l2.push_back(task, &inject, &mut metrics); + 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 metrics = MetricsBatch::new(); + let mut stats = new_stats(); let (_, mut local) = queue::local(); - s1.steal_into(&mut local, &mut metrics); + s1.steal_into(&mut local, &mut stats); while local.pop().is_some() {} }); @@ -198,12 +195,11 @@ fn chained_steal() { // Drain our tasks, then attempt to steal while l1.pop().is_some() {} - s2.steal_into(&mut l1, &mut metrics); + s2.steal_into(&mut l1, &mut stats); th.join().unwrap(); while l1.pop().is_some() {} while l2.pop().is_some() {} - while inject.pop().is_some() {} }); } diff --git a/src/runtime/tests/loom_shutdown_join.rs b/src/runtime/tests/loom_multi_thread/shutdown.rs index 6fbc4bf..6fbc4bf 100644 --- a/src/runtime/tests/loom_shutdown_join.rs +++ b/src/runtime/tests/loom_multi_thread/shutdown.rs diff --git a/src/runtime/tests/loom_multi_thread/yield_now.rs b/src/runtime/tests/loom_multi_thread/yield_now.rs new file mode 100644 index 0000000..ba506e5 --- /dev/null +++ b/src/runtime/tests/loom_multi_thread/yield_now.rs @@ -0,0 +1,37 @@ +use 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() +} diff --git a/src/runtime/tests/loom_multi_thread_alt.rs b/src/runtime/tests/loom_multi_thread_alt.rs new file mode 100644 index 0000000..1b9c3b4 --- /dev/null +++ b/src/runtime/tests/loom_multi_thread_alt.rs @@ -0,0 +1,570 @@ +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::future::poll_fn; +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::Future; +use std::pin::Pin; +use std::sync::atomic::Ordering::{Relaxed, SeqCst}; +use std::task::{Context, Poll}; + +mod atomic_take { + use loom::sync::atomic::AtomicBool; + use std::mem::MaybeUninit; + use std::sync::atomic::Ordering::SeqCst; + + pub(super) struct AtomicTake<T> { + inner: MaybeUninit<T>, + taken: AtomicBool, + } + + impl<T> AtomicTake<T> { + pub(super) fn new(value: T) -> Self { + Self { + inner: MaybeUninit::new(value), + taken: AtomicBool::new(false), + } + } + + pub(super) fn take(&self) -> Option<T> { + // 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<T> Drop for AtomicTake<T> { + fn drop(&mut self) { + drop(self.take()); + } + } +} + +#[derive(Clone)] +struct AtomicOneshot<T> { + value: std::sync::Arc<atomic_take::AtomicTake<oneshot::Sender<T>>>, +} +impl<T> AtomicOneshot<T> { + fn new(sender: oneshot::Sender<T>) -> 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<std::sync::atomic::AtomicIsize>; + 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<Output = &'static str> { + 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<T: Future>(f: T) -> Track<T> { + Track { + inner: f, + arc: Arc::new(()), + } +} + +pin_project! { + struct Track<T> { + #[pin] + inner: T, + // Arc is used to hook into loom's leak tracking. + arc: Arc<()>, + } +} + +impl<T> Track<T> { + fn into_inner(self) -> T { + self.inner + } +} + +impl<T: Future> Future for Track<T> { + type Output = Track<T::Output>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let me = self.project(); + + Poll::Ready(Track { + inner: ready!(me.inner.poll(cx)), + arc: me.arc.clone(), + }) + } +} diff --git a/src/runtime/tests/loom_multi_thread_alt/queue.rs b/src/runtime/tests/loom_multi_thread_alt/queue.rs new file mode 100644 index 0000000..0d81828 --- /dev/null +++ b/src/runtime/tests/loom_multi_thread_alt/queue.rs @@ -0,0 +1,205 @@ +use 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<NoopSchedule>) -> 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() {} + }); +} diff --git a/src/runtime/tests/loom_multi_thread_alt/shutdown.rs b/src/runtime/tests/loom_multi_thread_alt/shutdown.rs new file mode 100644 index 0000000..6fbc4bf --- /dev/null +++ b/src/runtime/tests/loom_multi_thread_alt/shutdown.rs @@ -0,0 +1,28 @@ +use 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()); + }); +} diff --git a/src/runtime/tests/loom_multi_thread_alt/yield_now.rs b/src/runtime/tests/loom_multi_thread_alt/yield_now.rs new file mode 100644 index 0000000..ba506e5 --- /dev/null +++ b/src/runtime/tests/loom_multi_thread_alt/yield_now.rs @@ -0,0 +1,37 @@ +use 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() +} diff --git a/src/runtime/tests/mod.rs b/src/runtime/tests/mod.rs index 4e7c245..0ba7480 100644 --- a/src/runtime/tests/mod.rs +++ b/src/runtime/tests/mod.rs @@ -52,17 +52,20 @@ mod unowned_wrapper { cfg_loom! { mod loom_blocking; - mod loom_current_thread_scheduler; + mod loom_current_thread; + mod loom_join_set; mod loom_local; + mod loom_multi_thread; + mod loom_multi_thread_alt; mod loom_oneshot; - mod loom_pool; - mod loom_queue; - mod loom_shutdown_join; - mod loom_join_set; - mod loom_yield; + + // 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))] diff --git a/src/runtime/tests/queue.rs b/src/runtime/tests/queue.rs index 68d2e89..5df92b7 100644 --- a/src/runtime/tests/queue.rs +++ b/src/runtime/tests/queue.rs @@ -1,18 +1,18 @@ -use crate::runtime::scheduler::multi_thread::queue; -use crate::runtime::task::{self, Inject, Schedule, Task}; -use crate::runtime::MetricsBatch; +use crate::runtime::scheduler::multi_thread::{queue, Stats}; +use crate::runtime::task::{self, Schedule, Task}; +use std::cell::RefCell; use std::thread; use std::time::Duration; #[allow(unused)] macro_rules! assert_metrics { - ($metrics:ident, $field:ident == $v:expr) => {{ + ($stats:ident, $field:ident == $v:expr) => {{ use crate::runtime::WorkerMetrics; use std::sync::atomic::Ordering::Relaxed; let worker = WorkerMetrics::new(); - $metrics.submit(&worker); + $stats.submit(&worker); let expect = $v; let actual = worker.$field.load(Relaxed); @@ -21,46 +21,87 @@ macro_rules! assert_metrics { }}; } +fn new_stats() -> Stats { + use crate::runtime::WorkerMetrics; + Stats::new(&WorkerMetrics::new()) +} + #[test] -fn fits_256() { +fn fits_256_one_at_a_time() { let (_, mut local) = queue::local(); - let inject = Inject::new(); - let mut metrics = MetricsBatch::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); for _ in 0..256 { let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + local.push_back_or_overflow(task, &inject, &mut stats); } cfg_metrics! { - assert_metrics!(metrics, overflow_count == 0); + assert_metrics!(stats, overflow_count == 0); } - assert!(inject.pop().is_none()); + 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::<Vec<_>>(); + 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::<Vec<_>>(); + + 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 = Inject::new(); - let mut metrics = MetricsBatch::new(); + let inject = RefCell::new(vec![]); + let mut stats = new_stats(); for _ in 0..257 { let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + local.push_back_or_overflow(task, &inject, &mut stats); } cfg_metrics! { - assert_metrics!(metrics, overflow_count == 1); + assert_metrics!(stats, overflow_count == 1); } let mut n = 0; - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); while local.pop().is_some() { n += 1; @@ -71,21 +112,21 @@ fn overflow() { #[test] fn steal_batch() { - let mut metrics = MetricsBatch::new(); + let mut stats = new_stats(); let (steal1, mut local1) = queue::local(); let (_, mut local2) = queue::local(); - let inject = Inject::new(); + let inject = RefCell::new(vec![]); for _ in 0..4 { let (task, _) = super::unowned(async {}); - local1.push_back(task, &inject, &mut metrics); + local1.push_back_or_overflow(task, &inject, &mut stats); } - assert!(steal1.steal_into(&mut local2, &mut metrics).is_some()); + assert!(steal1.steal_into(&mut local2, &mut stats).is_some()); cfg_metrics! { - assert_metrics!(metrics, steal_count == 2); + assert_metrics!(stats, steal_count == 2); } for _ in 0..1 { @@ -117,19 +158,19 @@ fn stress1() { const NUM_PUSH: usize = normal_or_miri(500, 10); const NUM_POP: usize = normal_or_miri(250, 10); - let mut metrics = MetricsBatch::new(); + let mut stats = new_stats(); for _ in 0..NUM_ITER { let (steal, mut local) = queue::local(); - let inject = Inject::new(); + let inject = RefCell::new(vec![]); let th = thread::spawn(move || { - let mut metrics = MetricsBatch::new(); + 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 metrics).is_some() { + if steal.steal_into(&mut local, &mut stats).is_some() { n += 1; } @@ -141,7 +182,7 @@ fn stress1() { } cfg_metrics! { - assert_metrics!(metrics, steal_count == n as _); + assert_metrics!(stats, steal_count == n as _); } n @@ -152,7 +193,7 @@ fn stress1() { for _ in 0..NUM_LOCAL { for _ in 0..NUM_PUSH { let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + local.push_back_or_overflow(task, &inject, &mut stats); } for _ in 0..NUM_POP { @@ -164,9 +205,7 @@ fn stress1() { } } - while inject.pop().is_some() { - n += 1; - } + n += inject.borrow_mut().drain(..).count(); n += th.join().unwrap(); @@ -180,14 +219,14 @@ fn stress2() { const NUM_TASKS: usize = normal_or_miri(1_000_000, 50); const NUM_STEAL: usize = normal_or_miri(1_000, 10); - let mut metrics = MetricsBatch::new(); + let mut stats = new_stats(); for _ in 0..NUM_ITER { let (steal, mut local) = queue::local(); - let inject = Inject::new(); + let inject = RefCell::new(vec![]); let th = thread::spawn(move || { - let mut stats = MetricsBatch::new(); + let mut stats = new_stats(); let (_, mut local) = queue::local(); let mut n = 0; @@ -210,15 +249,13 @@ fn stress2() { for i in 0..NUM_TASKS { let (task, _) = super::unowned(async {}); - local.push_back(task, &inject, &mut metrics); + local.push_back_or_overflow(task, &inject, &mut stats); if i % 128 == 0 && local.pop().is_some() { num_pop += 1; } - while inject.pop().is_some() { - num_pop += 1; - } + num_pop += inject.borrow_mut().drain(..).count(); } num_pop += th.join().unwrap(); @@ -227,9 +264,7 @@ fn stress2() { num_pop += 1; } - while inject.pop().is_some() { - num_pop += 1; - } + num_pop += inject.borrow_mut().drain(..).count(); assert_eq!(num_pop, NUM_TASKS); } diff --git a/src/runtime/tests/task.rs b/src/runtime/tests/task.rs index a79c0f5..0485bba 100644 --- a/src/runtime/tests/task.rs +++ b/src/runtime/tests/task.rs @@ -1,11 +1,10 @@ use crate::runtime::task::{self, unowned, Id, JoinHandle, OwnedTasks, Schedule, Task}; use crate::runtime::tests::NoopSchedule; -use crate::util::TryLock; use std::collections::VecDeque; use std::future::Future; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; struct AssertDropHandle { is_dropped: Arc<AtomicBool>, @@ -243,7 +242,7 @@ fn with(f: impl FnOnce(Runtime)) { let rt = Runtime(Arc::new(Inner { owned: OwnedTasks::new(), - core: TryLock::new(Core { + core: Mutex::new(Core { queue: VecDeque::new(), }), })); @@ -256,7 +255,7 @@ fn with(f: impl FnOnce(Runtime)) { struct Runtime(Arc<Inner>); struct Inner { - core: TryLock<Core>, + core: Mutex<Core>, owned: OwnedTasks<Runtime>, } @@ -264,7 +263,7 @@ struct Core { queue: VecDeque<task::Notified<Runtime>>, } -static CURRENT: TryLock<Option<Runtime>> = TryLock::new(None); +static CURRENT: Mutex<Option<Runtime>> = Mutex::new(None); impl Runtime { fn spawn<T>(&self, future: T) -> JoinHandle<T::Output> diff --git a/src/runtime/time/entry.rs b/src/runtime/time/entry.rs index f0d613a..798d3c1 100644 --- a/src/runtime/time/entry.rs +++ b/src/runtime/time/entry.rs @@ -72,6 +72,10 @@ 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 = u64::MAX - 2; /// This structure holds the current shared state of the timer - its scheduled /// time (if registered), or otherwise the result of the timer completing, as @@ -94,7 +98,7 @@ pub(super) struct StateCell { /// without holding the driver lock is undefined behavior. result: UnsafeCell<TimerResult>, /// The currently-registered waker - waker: CachePadded<AtomicWaker>, + waker: AtomicWaker, } impl Default for StateCell { @@ -114,7 +118,7 @@ impl StateCell { Self { state: AtomicU64::new(STATE_DEREGISTERED), result: UnsafeCell::new(Ok(())), - waker: CachePadded(AtomicWaker::new()), + waker: AtomicWaker::new(), } } @@ -126,7 +130,7 @@ impl StateCell { fn when(&self) -> Option<u64> { let cur_state = self.state.load(Ordering::Relaxed); - if cur_state == u64::MAX { + if cur_state == STATE_DEREGISTERED { None } else { Some(cur_state) @@ -139,7 +143,7 @@ impl StateCell { // 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.0.register_by_ref(waker); + self.waker.register_by_ref(waker); self.read_state() } @@ -227,7 +231,7 @@ impl StateCell { self.state.store(STATE_DEREGISTERED, Ordering::Release); - self.waker.0.take_waker() + self.waker.take_waker() } /// Marks the timer as registered (poll will return None) and sets the @@ -298,9 +302,11 @@ pub(crate) struct TimerEntry { /// This is manipulated only under the inner mutex. TODO: Can we use loom /// cells for this? inner: StdUnsafeCell<TimerShared>, - /// Initial deadline for the timer. This is used to register on the first + /// Deadline for the timer. This is used to register on the first /// poll, as we can't register prior to being pinned. - initial_deadline: Option<Instant>, + deadline: Instant, + /// Whether the deadline has been registered. + registered: bool, /// Ensure the type is !Unpin _m: std::marker::PhantomPinned, } @@ -329,11 +335,20 @@ pub(super) type EntryList = crate::util::linked_list::LinkedList<TimerShared, Ti /// frontend (`Entry`) and driver backend. /// /// Note that this structure is located inside the `TimerEntry` structure. -#[derive(Debug)] -#[repr(C)] pub(crate) struct TimerShared { - /// Data manipulated by the driver thread itself, only. - driver_state: CachePadded<TimerSharedPadded>, + /// 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<TimerShared>, + + /// 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, + + /// The true expiration time. Set by the timer future, read by the driver. + true_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 @@ -343,10 +358,23 @@ pub(crate) struct TimerShared { _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("when", &self.true_when.load(Ordering::Relaxed)) + .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<Self>) -> NonNull<linked_list::Pointers<TimerShared>> { - &self.driver_state.0.pointers + &self.pointers } } } @@ -354,8 +382,10 @@ generate_addr_of_methods! { impl TimerShared { pub(super) fn new() -> Self { Self { + cached_when: AtomicU64::new(0), + true_when: AtomicU64::new(0), + pointers: linked_list::Pointers::new(), state: StateCell::default(), - driver_state: CachePadded(TimerSharedPadded::new()), _p: PhantomPinned, } } @@ -363,7 +393,7 @@ impl TimerShared { /// 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.driver_state.0.cached_when.load(Ordering::Relaxed) + self.cached_when.load(Ordering::Relaxed) } /// Gets the true time-of-expiration value, and copies it into the cached @@ -374,10 +404,7 @@ impl TimerShared { pub(super) unsafe fn sync_when(&self) -> u64 { let true_when = self.true_when(); - self.driver_state - .0 - .cached_when - .store(true_when, Ordering::Relaxed); + self.cached_when.store(true_when, Ordering::Relaxed); true_when } @@ -387,10 +414,7 @@ impl TimerShared { /// 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.driver_state - .0 - .cached_when - .store(when, Ordering::Relaxed); + self.cached_when.store(when, Ordering::Relaxed); } /// Returns the true time-of-expiration value, with relaxed memory ordering. @@ -405,7 +429,7 @@ impl TimerShared { /// in the timer wheel. pub(super) unsafe fn set_expiration(&self, t: u64) { self.state.set_expiration(t); - self.driver_state.0.cached_when.store(t, Ordering::Relaxed); + self.cached_when.store(t, Ordering::Relaxed); } /// Sets the true time-of-expiration only if it is after the current. @@ -429,48 +453,6 @@ impl TimerShared { } } -/// Additional shared state between the driver and the timer which is cache -/// padded. This contains the information that the driver thread accesses most -/// frequently to minimize contention. In particular, we move it away from the -/// waker, as the waker is updated on every poll. -struct TimerSharedPadded { - /// 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<TimerShared>, - - /// 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, - - /// The true expiration time. Set by the timer future, read by the driver. - true_when: AtomicU64, -} - -impl std::fmt::Debug for TimerSharedPadded { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TimerSharedPadded") - .field("when", &self.true_when.load(Ordering::Relaxed)) - .field("cached_when", &self.cached_when.load(Ordering::Relaxed)) - .finish() - } -} - -impl TimerSharedPadded { - fn new() -> Self { - Self { - cached_when: AtomicU64::new(0), - true_when: AtomicU64::new(0), - pointers: linked_list::Pointers::new(), - } - } -} - -unsafe impl Send for TimerShared {} -unsafe impl Sync for TimerShared {} - unsafe impl linked_list::Link for TimerShared { type Handle = TimerHandle; @@ -504,7 +486,8 @@ impl TimerEntry { Self { driver, inner: StdUnsafeCell::new(TimerShared::new()), - initial_deadline: Some(deadline), + deadline, + registered: false, _m: std::marker::PhantomPinned, } } @@ -513,8 +496,12 @@ impl TimerEntry { unsafe { &*self.inner.get() } } + pub(crate) fn deadline(&self) -> Instant { + self.deadline + } + pub(crate) fn is_elapsed(&self) -> bool { - !self.inner().state.might_be_registered() && self.initial_deadline.is_none() + !self.inner().state.might_be_registered() && self.registered } /// Cancels and deregisters the timer. This operation is irreversible. @@ -544,8 +531,9 @@ impl TimerEntry { unsafe { self.driver().clear_entry(NonNull::from(self.inner())) }; } - pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant) { - unsafe { self.as_mut().get_unchecked_mut() }.initial_deadline = None; + pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant, reregister: bool) { + unsafe { self.as_mut().get_unchecked_mut() }.deadline = new_time; + unsafe { self.as_mut().get_unchecked_mut() }.registered = reregister; let tick = self.driver().time_source().deadline_to_tick(new_time); @@ -553,9 +541,11 @@ impl TimerEntry { return; } - unsafe { - self.driver() - .reregister(&self.driver.driver().io, tick, self.inner().into()); + if reregister { + unsafe { + self.driver() + .reregister(&self.driver.driver().io, tick, self.inner().into()); + } } } @@ -567,8 +557,9 @@ impl TimerEntry { panic!("{}", crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR); } - if let Some(deadline) = self.initial_deadline { - self.as_mut().reset(deadline); + if !self.registered { + let deadline = self.deadline; + self.as_mut().reset(deadline, true); } let this = unsafe { self.get_unchecked_mut() }; @@ -579,6 +570,11 @@ impl TimerEntry { 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 { @@ -646,74 +642,3 @@ impl Drop for TimerEntry { unsafe { Pin::new_unchecked(self) }.as_mut().cancel() } } - -// Copied from [crossbeam/cache_padded](https://github.com/crossbeam-rs/crossbeam/blob/fa35346b7c789bba045ad789e894c68c466d1779/crossbeam-utils/src/cache_padded.rs#L62-L127) -// -// 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, and riscv64 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/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_riscv64.go#L7 -#[cfg_attr( - any( - target_arch = "arm", - target_arch = "mips", - target_arch = "mips64", - target_arch = "riscv64", - ), - 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 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 -// -// 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 = "riscv64", - target_arch = "s390x", - )), - repr(align(64)) -)] -#[derive(Debug, Default)] -struct CachePadded<T>(T); diff --git a/src/runtime/time/mod.rs b/src/runtime/time/mod.rs index f81cab8..423ad79 100644 --- a/src/runtime/time/mod.rs +++ b/src/runtime/time/mod.rs @@ -8,7 +8,7 @@ mod entry; pub(crate) use entry::TimerEntry; -use entry::{EntryList, TimerHandle, TimerShared}; +use entry::{EntryList, TimerHandle, TimerShared, MAX_SAFE_MILLIS_DURATION}; mod handle; pub(crate) use self::handle::Handle; @@ -125,7 +125,7 @@ impl Driver { /// 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) -> (Driver, Handle) { + pub(crate) fn new(park: IoStack, clock: &Clock) -> (Driver, Handle) { let time_source = TimeSource::new(clock); let handle = Handle { @@ -186,7 +186,7 @@ impl Driver { match next_wake { Some(when) => { - let now = handle.time_source.now(); + 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. @@ -214,13 +214,13 @@ impl Driver { } // Process pending timers after waking up - handle.process(); + 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 = &handle.time_source.clock; + let clock = rt_handle.clock(); if clock.can_auto_advance() { self.park.park_timeout(rt_handle, Duration::from_secs(0)); @@ -231,7 +231,9 @@ impl Driver { // advance the clock. if !handle.did_wake() { // Simulate advancing time - clock.advance(duration); + if let Err(msg) = clock.advance(duration) { + panic!("{}", msg); + } } } else { self.park.park_timeout(rt_handle, duration); @@ -248,8 +250,8 @@ impl Driver { impl Handle { /// Runs timer related logic, and returns the next wakeup time - pub(self) fn process(&self) { - let now = self.time_source().now(); + pub(self) fn process(&self, clock: &Clock) { + let now = self.time_source().now(clock); self.process_at_time(now) } diff --git a/src/runtime/time/source.rs b/src/runtime/time/source.rs index e6788ed..4647bc4 100644 --- a/src/runtime/time/source.rs +++ b/src/runtime/time/source.rs @@ -1,19 +1,16 @@ +use super::MAX_SAFE_MILLIS_DURATION; use crate::time::{Clock, Duration, Instant}; -use std::convert::TryInto; - /// A structure which handles conversion from Instants to u64 timestamps. #[derive(Debug)] pub(crate) struct TimeSource { - pub(crate) clock: Clock, start_time: Instant, } impl TimeSource { - pub(crate) fn new(clock: Clock) -> Self { + pub(crate) fn new(clock: &Clock) -> Self { Self { start_time: clock.now(), - clock, } } @@ -29,14 +26,14 @@ impl TimeSource { .unwrap_or_else(|| Duration::from_secs(0)); let ms = dur.as_millis(); - ms.try_into().unwrap_or(u64::MAX) + ms.try_into().unwrap_or(MAX_SAFE_MILLIS_DURATION) } pub(crate) fn tick_to_duration(&self, t: u64) -> Duration { Duration::from_millis(t) } - pub(crate) fn now(&self) -> u64 { - self.instant_to_tick(self.clock.now()) + pub(crate) fn now(&self, clock: &Clock) -> u64 { + self.instant_to_tick(clock.now()) } } diff --git a/src/runtime/time/tests/mod.rs b/src/runtime/time/tests/mod.rs index 88c7d76..e7ab222 100644 --- a/src/runtime/time/tests/mod.rs +++ b/src/runtime/time/tests/mod.rs @@ -1,4 +1,4 @@ -#![cfg(not(tokio_wasi))] +#![cfg(not(target_os = "wasi"))] use std::{task::Context, time::Duration}; @@ -62,12 +62,13 @@ fn single_timer() { thread::yield_now(); - let handle = handle.inner.driver().time(); + 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. - handle.process_at_time(handle.time_source().now() + 2_000_000_000); + time.process_at_time(time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) @@ -97,10 +98,11 @@ fn drop_timer() { thread::yield_now(); - let handle = handle.inner.driver().time(); + let time = handle.inner.driver().time(); + let clock = handle.inner.driver().clock(); // advance 2s in the future. - handle.process_at_time(handle.time_source().now() + 2_000_000_000); + time.process_at_time(time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) @@ -132,10 +134,11 @@ fn change_waker() { thread::yield_now(); - let handle = handle.inner.driver().time(); + let time = handle.inner.driver().time(); + let clock = handle.inner.driver().clock(); // advance 2s - handle.process_at_time(handle.time_source().now() + 2_000_000_000); + time.process_at_time(time.time_source().now(clock) + 2_000_000_000); jh.join().unwrap(); }) @@ -161,7 +164,7 @@ fn reset_future() { .as_mut() .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref())); - entry.as_mut().reset(start + Duration::from_secs(2)); + entry.as_mut().reset(start + Duration::from_secs(2), true); // shouldn't complete before 2s block_on(futures::future::poll_fn(|cx| { diff --git a/src/runtime/time/wheel/mod.rs b/src/runtime/time/wheel/mod.rs index c3ba364..bf13b7b 100644 --- a/src/runtime/time/wheel/mod.rs +++ b/src/runtime/time/wheel/mod.rs @@ -148,23 +148,13 @@ impl Wheel { return Some(handle); } - // under what circumstances is poll.expiration Some vs. None? - let expiration = self.next_expiration().and_then(|expiration| { - if expiration.deadline > now { - None - } else { - Some(expiration) - } - }); - - match expiration { - Some(ref expiration) if expiration.deadline > now => return None, - Some(ref expiration) => { + match self.next_expiration() { + Some(ref expiration) if expiration.deadline <= now => { self.process_expiration(expiration); self.set_elapsed(expiration.deadline); } - None => { + _ => { // 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 diff --git a/src/signal/unix.rs b/src/signal/unix.rs index e5345fd..ae5c130 100644 --- a/src/signal/unix.rs +++ b/src/signal/unix.rs @@ -292,7 +292,11 @@ fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { } } -/// A stream of events for receiving a particular type of OS signal. +/// 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 @@ -307,7 +311,7 @@ fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { /// 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 stream corresponds to +/// 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 @@ -345,11 +349,11 @@ fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result<()> { /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { /// // An infinite stream of hangup signals. -/// let mut stream = signal(SignalKind::hangup())?; +/// let mut sig = signal(SignalKind::hangup())?; /// /// // Print whenever a HUP signal is received /// loop { -/// stream.recv().await; +/// sig.recv().await; /// println!("got signal HUP"); /// } /// } @@ -360,7 +364,7 @@ pub struct Signal { inner: RxFuture, } -/// Creates a new stream which will receive notifications when the current +/// 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. diff --git a/src/signal/windows.rs b/src/signal/windows.rs index 730f95d..2f70f98 100644 --- a/src/signal/windows.rs +++ b/src/signal/windows.rs @@ -22,7 +22,7 @@ pub(crate) use self::imp::{OsExtraData, OsStorage}; #[path = "windows/stub.rs"] mod imp; -/// Creates a new stream which receives "ctrl-c" notifications sent to the +/// Creates a new listener which receives "ctrl-c" notifications sent to the /// process. /// /// # Examples @@ -32,12 +32,12 @@ mod imp; /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-C events. -/// let mut stream = ctrl_c()?; +/// // 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() { -/// stream.recv().await; +/// signal.recv().await; /// println!("got CTRL-C. {} more to exit", countdown); /// } /// @@ -50,14 +50,18 @@ pub fn ctrl_c() -> io::Result<CtrlC> { }) } -/// Represents a stream which receives "ctrl-c" notifications sent to the process +/// Represents a listener which receives "ctrl-c" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// -/// A notification to this process notifies *all* streams listening for +/// 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 stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// 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, @@ -66,7 +70,7 @@ pub struct CtrlC { impl CtrlC { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by the listener. /// /// # Examples /// @@ -75,12 +79,11 @@ impl CtrlC { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-C events. - /// let mut stream = ctrl_c()?; + /// let mut signal = ctrl_c()?; /// /// // Print whenever a CTRL-C event is received. /// for countdown in (0..3).rev() { - /// stream.recv().await; + /// signal.recv().await; /// println!("got CTRL-C. {} more to exit", countdown); /// } /// @@ -94,7 +97,7 @@ impl CtrlC { /// 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 stream. + /// `None` is returned if no more events can be received. /// /// # Examples /// @@ -124,14 +127,18 @@ impl CtrlC { } } -/// Represents a stream which receives "ctrl-break" notifications sent to the process +/// Represents a listener which receives "ctrl-break" notifications sent to the process /// via `SetConsoleCtrlHandler`. /// -/// A notification to this process notifies *all* streams listening for +/// 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 stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// 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, @@ -140,7 +147,7 @@ pub struct CtrlBreak { impl CtrlBreak { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -149,12 +156,12 @@ impl CtrlBreak { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-BREAK events. - /// let mut stream = ctrl_break()?; + /// // A listener of CTRL-BREAK events. + /// let mut signal = ctrl_break()?; /// /// // Print whenever a CTRL-BREAK event is received. /// loop { - /// stream.recv().await; + /// signal.recv().await; /// println!("got signal CTRL-BREAK"); /// } /// } @@ -166,7 +173,7 @@ impl CtrlBreak { /// 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 stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -196,7 +203,7 @@ impl CtrlBreak { } } -/// Creates a new stream which receives "ctrl-break" notifications sent to the +/// Creates a new listener which receives "ctrl-break" notifications sent to the /// process. /// /// # Examples @@ -206,12 +213,12 @@ impl CtrlBreak { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-BREAK events. -/// let mut stream = ctrl_break()?; +/// // A listener of CTRL-BREAK events. +/// let mut signal = ctrl_break()?; /// /// // Print whenever a CTRL-BREAK event is received. /// loop { -/// stream.recv().await; +/// signal.recv().await; /// println!("got signal CTRL-BREAK"); /// } /// } @@ -222,7 +229,7 @@ pub fn ctrl_break() -> io::Result<CtrlBreak> { }) } -/// Creates a new stream which receives "ctrl-close" notifications sent to the +/// Creates a new listener which receives "ctrl-close" notifications sent to the /// process. /// /// # Examples @@ -232,12 +239,12 @@ pub fn ctrl_break() -> io::Result<CtrlBreak> { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-CLOSE events. -/// let mut stream = ctrl_close()?; +/// // 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() { -/// stream.recv().await; +/// signal.recv().await; /// println!("got CTRL-CLOSE. {} more to exit", countdown); /// } /// @@ -250,14 +257,14 @@ pub fn ctrl_close() -> io::Result<CtrlClose> { }) } -/// Represents a stream which receives "ctrl-close" notitifications sent to the process +/// Represents a listener which receives "ctrl-close" notitifications sent to the process /// via 'SetConsoleCtrlHandler'. /// -/// A notification to this process notifies *all* streams listening for +/// 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 stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// 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, @@ -266,7 +273,7 @@ pub struct CtrlClose { impl CtrlClose { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -275,11 +282,11 @@ impl CtrlClose { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-CLOSE events. - /// let mut stream = ctrl_close()?; + /// // A listener of CTRL-CLOSE events. + /// let mut signal = ctrl_close()?; /// /// // Print whenever a CTRL-CLOSE event is received. - /// stream.recv().await; + /// signal.recv().await; /// println!("got CTRL-CLOSE. Cleaning up before exiting"); /// /// Ok(()) @@ -292,7 +299,7 @@ impl CtrlClose { /// 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 stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -322,7 +329,7 @@ impl CtrlClose { } } -/// Creates a new stream which receives "ctrl-shutdown" notifications sent to the +/// Creates a new listener which receives "ctrl-shutdown" notifications sent to the /// process. /// /// # Examples @@ -332,10 +339,10 @@ impl CtrlClose { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-SHUTDOWN events. -/// let mut stream = ctrl_shutdown()?; +/// // A listener of CTRL-SHUTDOWN events. +/// let mut signal = ctrl_shutdown()?; /// -/// stream.recv().await; +/// signal.recv().await; /// println!("got CTRL-SHUTDOWN. Cleaning up before exiting"); /// /// Ok(()) @@ -347,14 +354,14 @@ pub fn ctrl_shutdown() -> io::Result<CtrlShutdown> { }) } -/// Represents a stream which receives "ctrl-shutdown" notitifications sent to the process +/// Represents a listener which receives "ctrl-shutdown" notitifications sent to the process /// via 'SetConsoleCtrlHandler'. /// -/// A notification to this process notifies *all* streams listening for +/// 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 stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// 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, @@ -363,7 +370,7 @@ pub struct CtrlShutdown { impl CtrlShutdown { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -372,11 +379,11 @@ impl CtrlShutdown { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-SHUTDOWN events. - /// let mut stream = ctrl_shutdown()?; + /// // A listener of CTRL-SHUTDOWN events. + /// let mut signal = ctrl_shutdown()?; /// /// // Print whenever a CTRL-SHUTDOWN event is received. - /// stream.recv().await; + /// signal.recv().await; /// println!("got CTRL-SHUTDOWN. Cleaning up before exiting"); /// /// Ok(()) @@ -389,7 +396,7 @@ impl CtrlShutdown { /// 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 stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -419,7 +426,7 @@ impl CtrlShutdown { } } -/// Creates a new stream which receives "ctrl-logoff" notifications sent to the +/// Creates a new listener which receives "ctrl-logoff" notifications sent to the /// process. /// /// # Examples @@ -429,10 +436,10 @@ impl CtrlShutdown { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { -/// // An infinite stream of CTRL-LOGOFF events. -/// let mut stream = ctrl_logoff()?; +/// // A listener of CTRL-LOGOFF events. +/// let mut signal = ctrl_logoff()?; /// -/// stream.recv().await; +/// signal.recv().await; /// println!("got CTRL-LOGOFF. Cleaning up before exiting"); /// /// Ok(()) @@ -444,14 +451,14 @@ pub fn ctrl_logoff() -> io::Result<CtrlLogoff> { }) } -/// Represents a stream which receives "ctrl-logoff" notitifications sent to the process +/// Represents a listener which receives "ctrl-logoff" notitifications sent to the process /// via 'SetConsoleCtrlHandler'. /// -/// A notification to this process notifies *all* streams listening for +/// 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 stream may only receive one item about the two notifications. -#[must_use = "streams do nothing unless polled"] +/// 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, @@ -460,7 +467,7 @@ pub struct CtrlLogoff { impl CtrlLogoff { /// Receives the next signal notification event. /// - /// `None` is returned if no more events can be received by this stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// @@ -469,11 +476,11 @@ impl CtrlLogoff { /// /// #[tokio::main] /// async fn main() -> Result<(), Box<dyn std::error::Error>> { - /// // An infinite stream of CTRL-LOGOFF events. - /// let mut stream = ctrl_logoff()?; + /// // An listener of CTRL-LOGOFF events. + /// let mut signal = ctrl_logoff()?; /// /// // Print whenever a CTRL-LOGOFF event is received. - /// stream.recv().await; + /// signal.recv().await; /// println!("got CTRL-LOGOFF. Cleaning up before exiting"); /// /// Ok(()) @@ -486,7 +493,7 @@ impl CtrlLogoff { /// 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 stream. + /// `None` is returned if no more events can be received by this listener. /// /// # Examples /// diff --git a/src/signal/windows/sys.rs b/src/signal/windows/sys.rs index f8133e0..26e6bdf 100644 --- a/src/signal/windows/sys.rs +++ b/src/signal/windows/sys.rs @@ -1,4 +1,3 @@ -use std::convert::TryFrom; use std::io; use std::sync::Once; diff --git a/src/sync/barrier.rs b/src/sync/barrier.rs index 2ce1d73..29b6d4e 100644 --- a/src/sync/barrier.rs +++ b/src/sync/barrier.rs @@ -132,6 +132,8 @@ impl Barrier { 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. diff --git a/src/sync/batch_semaphore.rs b/src/sync/batch_semaphore.rs index 57493f4..e0dafaa 100644 --- a/src/sync/batch_semaphore.rs +++ b/src/sync/batch_semaphore.rs @@ -178,14 +178,9 @@ impl Semaphore { /// Creates a new semaphore with the initial number of permits. /// /// Maximum number of permits on 32-bit platforms is `1<<29`. - /// - /// If the specified number of permits exceeds the maximum permit amount - /// Then the value will get clamped to the maximum number of permits. - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - pub(crate) const fn const_new(mut permits: usize) -> Self { - // NOTE: assertions and by extension panics are still being worked on: https://github.com/rust-lang/rust/issues/74925 - // currently we just clamp the permit count when it exceeds the max - permits &= Self::MAX_PERMITS; + #[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), @@ -198,6 +193,19 @@ impl Semaphore { } } + /// 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(), + } + } + /// Returns the current number of available permits. pub(crate) fn available_permits(&self) -> usize { self.permits.load(Acquire) >> Self::PERMIT_SHIFT @@ -264,7 +272,7 @@ impl Semaphore { match self.permits.compare_exchange(curr, next, AcqRel, Acquire) { Ok(_) => { - // TODO: Instrument once issue has been solved} + // TODO: Instrument once issue has been solved return Ok(()); } Err(actual) => curr = actual, @@ -444,6 +452,7 @@ impl Semaphore { } assert_eq!(acquired, 0); + let mut old_waker = None; // Otherwise, register the waker & enqueue the node. node.waker.with_mut(|waker| { @@ -455,7 +464,7 @@ impl Semaphore { .map(|waker| !waker.will_wake(cx.waker())) .unwrap_or(true) { - *waker = Some(cx.waker().clone()); + old_waker = std::mem::replace(waker, Some(cx.waker().clone())); } }); @@ -468,6 +477,8 @@ impl Semaphore { waiters.queue.push_front(node); } + drop(waiters); + drop(old_waker); Pending } diff --git a/src/sync/broadcast.rs b/src/sync/broadcast.rs index 1c6b2ca..42cde81 100644 --- a/src/sync/broadcast.rs +++ b/src/sync/broadcast.rs @@ -4,7 +4,7 @@ //! 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 also `Send` or `Sync` respectively. +//! 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 @@ -54,6 +54,10 @@ //! 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 @@ -114,8 +118,9 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::{Arc, Mutex, RwLock, RwLockReadGuard}; -use crate::util::linked_list::{self, LinkedList}; +use crate::loom::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard}; +use crate::util::linked_list::{self, GuardedLinkedList, LinkedList}; +use crate::util::WakeList; use std::fmt; use std::future::Future; @@ -361,6 +366,17 @@ struct Waiter { _p: PhantomPinned, } +impl Waiter { + fn new() -> Self { + Self { + queued: false, + waker: None, + pointers: linked_list::Pointers::new(), + _p: PhantomPinned, + } + } +} + generate_addr_of_methods! { impl<> Waiter { unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> { @@ -439,42 +455,13 @@ const MAX_RECEIVERS: usize = usize::MAX >> 2; /// This will panic if `capacity` is equal to `0` or larger /// than `usize::MAX / 2`. #[track_caller] -pub fn channel<T: Clone>(mut capacity: usize) -> (Sender<T>, Receiver<T>) { - assert!(capacity > 0, "capacity is empty"); - assert!(capacity <= usize::MAX >> 1, "requested capacity too large"); - - // 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(RwLock::new(Slot { - rem: AtomicUsize::new(0), - pos: (i as u64).wrapping_sub(capacity as u64), - val: UnsafeCell::new(None), - })); - } - - let shared = Arc::new(Shared { - buffer: buffer.into_boxed_slice(), - mask: capacity - 1, - tail: Mutex::new(Tail { - pos: 0, - rx_cnt: 1, - closed: false, - waiters: LinkedList::new(), - }), - num_tx: AtomicUsize::new(1), - }); - +pub fn channel<T: Clone>(capacity: usize) -> (Sender<T>, Receiver<T>) { + // 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: shared.clone(), + shared: tx.shared.clone(), next: 0, }; - - let tx = Sender { shared }; - (tx, rx) } @@ -485,6 +472,65 @@ unsafe impl<T: Send> Send for Receiver<T> {} unsafe impl<T: Send> Sync for Receiver<T> {} impl<T> Sender<T> { + /// Creates the sending-half of the [`broadcast`] channel. + /// + /// See documentation of [`broadcast::channel`] for errors when calling this function. + /// + /// [`broadcast`]: crate::sync::broadcast + /// [`broadcast::channel`]: crate::sync::broadcast + #[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(RwLock::new(Slot { + rem: AtomicUsize::new(0), + pos: (i as u64).wrapping_sub(capacity as u64), + val: UnsafeCell::new(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. /// @@ -496,7 +542,8 @@ impl<T> Sender<T> { /// /// 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 before receiving the message. + /// a receiver may drop or lag ([see lagging](self#lagging)) before receiving + /// the message. /// /// # Note /// @@ -565,12 +612,10 @@ impl<T> Sender<T> { // Release the slot lock before notifying the receivers. drop(slot); - tail.notify_rx(); - - // 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. - drop(tail); + // 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) } @@ -735,11 +780,34 @@ impl<T> Sender<T> { 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; - tail.notify_rx(); + self.shared.notify_rx(tail); } } @@ -760,18 +828,110 @@ fn new_receiver<T>(shared: Arc<Shared<T>>) -> Receiver<T> { Receiver { shared, next } } -impl Tail { - fn notify_rx(&mut self) { - while let Some(mut waiter) = self.waiters.pop_back() { - // Safety: `waiters` lock is still held. - let waiter = unsafe { waiter.as_mut() }; +/// 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<Waiter, <Waiter as linked_list::Link>::Target>, + is_empty: bool, + shared: &'a Shared<T>, +} + +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() {} + } + } +} - assert!(waiter.queued); - waiter.queued = false; +impl<'a, T> WaitersList<'a, T> { + fn new( + unguarded_list: LinkedList<Waiter, <Waiter as linked_list::Link>::Target>, + guard: Pin<&'a Waiter>, + shared: &'a Shared<T>, + ) -> Self { + let guard_ptr = NonNull::from(guard.get_ref()); + let list = unguarded_list.into_guarded(guard_ptr); + WaitersList { + list, + is_empty: false, + shared, + } + } - let waker = waiter.waker.take().unwrap(); - waker.wake(); + /// 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<NonNull<Waiter>> { + 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<T> Shared<T> { + 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(mut waiter) => { + // Safety: `tail` lock is still held. + let waiter = unsafe { waiter.as_mut() }; + + assert!(waiter.queued); + waiter.queued = false; + + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + 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(); } } @@ -860,6 +1020,29 @@ impl<T> Receiver<T> { 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, @@ -880,6 +1063,8 @@ impl<T> Receiver<T> { // the slot lock. drop(slot); + let mut old_waker = None; + let mut tail = self.shared.tail.lock(); // Acquire slot lock again @@ -912,7 +1097,10 @@ impl<T> Receiver<T> { match (*ptr).waker { Some(ref w) if w.will_wake(waker) => {} _ => { - (*ptr).waker = Some(waker.clone()); + old_waker = std::mem::replace( + &mut (*ptr).waker, + Some(waker.clone()), + ); } } @@ -924,6 +1112,11 @@ impl<T> Receiver<T> { } } + // Drop the old waker after releasing the locks. + drop(slot); + drop(tail); + drop(old_waker); + return Err(TryRecvError::Empty); } @@ -1106,6 +1299,33 @@ impl<T: Clone> Receiver<T> { 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<T, RecvError> { + crate::future::block_on(self.recv()) + } } impl<T> Drop for Receiver<T> { @@ -1164,6 +1384,8 @@ where type Output = Result<T, RecvError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> { + ready!(crate::trace::trace_leaf(cx)); + let (receiver, waiter) = self.project(); let guard = match receiver.recv_ref(Some((waiter, cx.waker()))) { @@ -1252,3 +1474,41 @@ impl<'a, T> Drop for RecvGuard<'a, T> { } fn is_unpin<T: Unpin>() {} + +#[cfg(not(loom))] +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn receiver_count_on_sender_constructor() { + let sender = Sender::<i32>::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::<i32>(16); + assert_eq!(sender.receiver_count(), 1); + + let _rx_2 = rx.resubscribe(); + assert_eq!(sender.receiver_count(), 2); + } +} diff --git a/src/sync/mod.rs b/src/sync/mod.rs index 8fba196..70fd9b9 100644 --- a/src/sync/mod.rs +++ b/src/sync/mod.rs @@ -449,7 +449,7 @@ cfg_sync! { pub mod mpsc; mod mutex; - pub use mutex::{Mutex, MutexGuard, TryLockError, OwnedMutexGuard, MappedMutexGuard}; + pub use mutex::{Mutex, MutexGuard, TryLockError, OwnedMutexGuard, MappedMutexGuard, OwnedMappedMutexGuard}; pub(crate) mod notify; pub use notify::Notify; diff --git a/src/sync/mpsc/bounded.rs b/src/sync/mpsc/bounded.rs index 8babdc7..e870ae5 100644 --- a/src/sync/mpsc/bounded.rs +++ b/src/sync/mpsc/bounded.rs @@ -18,7 +18,7 @@ use std::task::{Context, Poll}; /// 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/0.6/tokio_util/sync/struct.PollSender.html +/// [`PollSender`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSender.html pub struct Sender<T> { chan: chan::Tx<T, Semaphore>, } @@ -326,6 +326,7 @@ impl<T> Receiver<T> { /// ``` #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(&mut self) -> Option<T> { crate::future::block_on(self.recv()) } @@ -696,6 +697,7 @@ impl<T> Sender<T> { /// ``` #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "send_blocking"))] pub fn blocking_send(&self, value: T) -> Result<(), SendError<T>> { crate::future::block_on(self.send(value)) } @@ -859,6 +861,8 @@ impl<T> Sender<T> { } async fn reserve_inner(&self) -> Result<(), SendError<()>> { + crate::trace::async_trace_leaf().await; + match self.chan.semaphore().semaphore.acquire(1).await { Ok(_) => Ok(()), Err(_) => Err(SendError(())), diff --git a/src/sync/mpsc/chan.rs b/src/sync/mpsc/chan.rs index edd3e95..f24bb03 100644 --- a/src/sync/mpsc/chan.rs +++ b/src/sync/mpsc/chan.rs @@ -6,6 +6,7 @@ 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::process; @@ -46,18 +47,18 @@ pub(crate) trait Semaphore { } pub(super) struct Chan<T, S> { + /// Handle to the push half of the lock-free list. + tx: CachePadded<list::Tx<T>>, + + /// Receiver waker. Notified when a value is pushed into the channel. + rx_waker: CachePadded<AtomicWaker>, + /// Notifies all tasks listening for the receiver being dropped. notify_rx_closed: Notify, - /// Handle to the push half of the lock-free list. - tx: list::Tx<T>, - /// Coordinates access to channel's capacity. semaphore: S, - /// Receiver waker. Notified when a value is pushed into the channel. - rx_waker: AtomicWaker, - /// Tracks the number of outstanding sender handles. /// /// When this drops to zero, the send half of the channel is closed. @@ -73,9 +74,9 @@ where { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Chan") - .field("tx", &self.tx) + .field("tx", &*self.tx) .field("semaphore", &self.semaphore) - .field("rx_waker", &self.rx_waker) + .field("rx_waker", &*self.rx_waker) .field("tx_count", &self.tx_count) .field("rx_fields", &"...") .finish() @@ -108,9 +109,9 @@ pub(crate) fn channel<T, S: Semaphore>(semaphore: S) -> (Tx<T, S>, Rx<T, S>) { let chan = Arc::new(Chan { notify_rx_closed: Notify::new(), - tx, + tx: CachePadded::new(tx), semaphore, - rx_waker: AtomicWaker::new(), + rx_waker: CachePadded::new(AtomicWaker::new()), tx_count: AtomicUsize::new(1), rx_fields: UnsafeCell::new(RxFields { list: rx, @@ -242,6 +243,8 @@ impl<T, S: Semaphore> Rx<T, S> { pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> { use super::block::Read::*; + ready!(crate::trace::trace_leaf(cx)); + // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); diff --git a/src/sync/mpsc/error.rs b/src/sync/mpsc/error.rs index 1c789da..e418856 100644 --- a/src/sync/mpsc/error.rs +++ b/src/sync/mpsc/error.rs @@ -4,22 +4,28 @@ use std::error::Error; use std::fmt; /// Error returned by the `Sender`. -#[derive(Debug)] +#[derive(PartialEq, Eq, Clone, Copy)] pub struct SendError<T>(pub T); +impl<T> fmt::Debug for SendError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SendError").finish_non_exhaustive() + } +} + impl<T> fmt::Display for SendError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } -impl<T: fmt::Debug> std::error::Error for SendError<T> {} +impl<T> Error for SendError<T> {} // ===== TrySendError ===== /// This enumeration is the list of the possible error outcomes for the /// [try_send](super::Sender::try_send) method. -#[derive(Debug, Eq, PartialEq)] +#[derive(PartialEq, Eq, Clone, Copy)] pub enum TrySendError<T> { /// The data could not be sent on the channel because the channel is /// currently full and sending would require blocking. @@ -30,7 +36,14 @@ pub enum TrySendError<T> { Closed(T), } -impl<T: fmt::Debug> Error for TrySendError<T> {} +impl<T> fmt::Debug for TrySendError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + TrySendError::Full(..) => "Full(..)".fmt(f), + TrySendError::Closed(..) => "Closed(..)".fmt(f), + } + } +} impl<T> fmt::Display for TrySendError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -45,6 +58,8 @@ impl<T> fmt::Display for TrySendError<T> { } } +impl<T> Error for TrySendError<T> {} + impl<T> From<SendError<T>> for TrySendError<T> { fn from(src: SendError<T>) -> TrySendError<T> { TrySendError::Closed(src.0) @@ -96,7 +111,7 @@ impl Error for RecvError {} cfg_time! { // ===== SendTimeoutError ===== - #[derive(Debug, Eq, PartialEq)] + #[derive(PartialEq, Eq, Clone, Copy)] /// Error returned by [`Sender::send_timeout`](super::Sender::send_timeout)]. pub enum SendTimeoutError<T> { /// The data could not be sent on the channel because the channel is @@ -108,7 +123,14 @@ cfg_time! { Closed(T), } - impl<T: fmt::Debug> Error for SendTimeoutError<T> {} + impl<T> fmt::Debug for SendTimeoutError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + SendTimeoutError::Timeout(..) => "Timeout(..)".fmt(f), + SendTimeoutError::Closed(..) => "Closed(..)".fmt(f), + } + } + } impl<T> fmt::Display for SendTimeoutError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -122,4 +144,6 @@ cfg_time! { ) } } + + impl<T> Error for SendTimeoutError<T> {} } diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs index 33889fa..b2af084 100644 --- a/src/sync/mpsc/mod.rs +++ b/src/sync/mpsc/mod.rs @@ -33,7 +33,8 @@ //! //! 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. +//! result in an error. Additionally, all unread messages will be drained from the +//! channel and dropped. //! //! # Clean Shutdown //! diff --git a/src/sync/mpsc/unbounded.rs b/src/sync/mpsc/unbounded.rs index 5010204..cd83fc1 100644 --- a/src/sync/mpsc/unbounded.rs +++ b/src/sync/mpsc/unbounded.rs @@ -243,6 +243,7 @@ impl<T> UnboundedReceiver<T> { /// ``` #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(&mut self) -> Option<T> { crate::future::block_on(self.recv()) } diff --git a/src/sync/mutex.rs b/src/sync/mutex.rs index 024755c..0700c26 100644 --- a/src/sync/mutex.rs +++ b/src/sync/mutex.rs @@ -6,9 +6,10 @@ 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, marker, mem}; +use std::{fmt, mem, ptr}; /// An asynchronous `Mutex`-like type. /// @@ -141,8 +142,11 @@ pub struct Mutex<T: ?Sized> { /// /// 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<T>, @@ -163,7 +167,10 @@ pub struct MutexGuard<'a, T: ?Sized> { /// point `lock` will succeed yet again. /// /// [`Arc`]: std::sync::Arc +#[clippy::has_significant_drop] pub struct OwnedMutexGuard<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: Arc<Mutex<T>>, @@ -174,12 +181,71 @@ pub struct OwnedMutexGuard<T: ?Sized> { /// 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: marker::PhantomData<&'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<T: ?Sized, U: ?Sized = T> { + // 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<Mutex<T>>, +} + +/// 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<T>, +} + +/// A helper type used when taking apart a `OwnedMutexGuard` without running +/// its Drop implementation. +struct OwnedMutexGuardInner<T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + lock: Arc<Mutex<T>>, +} + +/// 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<T: ?Sized, U: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + data: *mut U, + lock: Arc<Mutex<T>>, } // As long as T: Send, it's fine to send and share Mutex<T> between threads. @@ -192,6 +258,19 @@ unsafe impl<T> Sync for OwnedMutexGuard<T> 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<T, U> Sync for OwnedMappedMutexGuard<T, U> +where + T: ?Sized + Send + Sync, + U: ?Sized + Send + Sync, +{ +} +unsafe impl<T, U> Send for OwnedMappedMutexGuard<T, U> +where + T: ?Sized + Send, + U: ?Sized + Send, +{ +} + /// Error returned from the [`Mutex::try_lock`], [`RwLock::try_read`] and /// [`RwLock::try_write`] functions. /// @@ -299,8 +378,7 @@ impl<T: ?Sized> Mutex<T> { /// /// static LOCK: Mutex<i32> = Mutex::const_new(5); /// ``` - #[cfg(all(feature = "parking_lot", not(all(loom, test)),))] - #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] + #[cfg(not(all(loom, test)))] pub const fn const_new(t: T) -> Self where T: Sized, @@ -337,15 +415,27 @@ impl<T: ?Sized> Mutex<T> { /// } /// ``` 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"))] - trace::async_op( - || self.acquire(), + let acquire_fut = trace::async_op( + move || acquire_fut, self.resource_span.clone(), "Mutex::lock", "poll", false, - ) - .await; + ); + + #[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(|| { @@ -355,14 +445,7 @@ impl<T: ?Sized> Mutex<T> { ); }); - #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] - self.acquire().await; - - MutexGuard { - lock: self, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span: self.resource_span.clone(), - } + guard } /// Blockingly locks this `Mutex`. When the lock has been acquired, function returns a @@ -414,6 +497,7 @@ impl<T: ?Sized> Mutex<T> { /// ``` #[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()) } @@ -508,37 +592,44 @@ impl<T: ?Sized> Mutex<T> { /// [`Arc`]: std::sync::Arc pub async fn lock_owned(self: Arc<Self>) -> OwnedMutexGuard<T> { #[cfg(all(tokio_unstable, feature = "tracing"))] - trace::async_op( - || self.acquire(), - self.resource_span.clone(), + 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, - ) - .await; + ); + + #[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(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = true, ); }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - - #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] - self.acquire().await; - - OwnedMutexGuard { - lock: self, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - } + 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. @@ -566,6 +657,12 @@ impl<T: ?Sized> Mutex<T> { pub fn try_lock(&self) -> Result<MutexGuard<'_, T>, 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!( @@ -574,11 +671,7 @@ impl<T: ?Sized> Mutex<T> { ); }); - Ok(MutexGuard { - lock: self, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span: self.resource_span.clone(), - }) + Ok(guard) } Err(_) => Err(TryLockError(())), } @@ -635,22 +728,21 @@ impl<T: ?Sized> Mutex<T> { pub fn try_lock_owned(self: Arc<Self>) -> Result<OwnedMutexGuard<T>, 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"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", locked = true, ); }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - - Ok(OwnedMutexGuard { - lock: self, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - }) + Ok(guard) } Err(_) => Err(TryLockError(())), } @@ -710,6 +802,17 @@ where // === 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. @@ -746,12 +849,13 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = &this.lock.s; - mem::forget(this); + let inner = this.skip_drop(); MappedMutexGuard { - s, + s: &inner.lock.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, } } @@ -796,12 +900,13 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = &this.lock.s; - mem::forget(this); + let inner = this.skip_drop(); Ok(MappedMutexGuard { - s, + s: &inner.lock.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, }) } @@ -834,6 +939,8 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> { impl<T: ?Sized> 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!( @@ -841,7 +948,6 @@ impl<T: ?Sized> Drop for MutexGuard<'_, T> { locked = false, ); }); - self.lock.s.release(1); } } @@ -873,6 +979,116 @@ impl<T: ?Sized + fmt::Display> fmt::Display for MutexGuard<'_, T> { // === impl OwnedMutexGuard === impl<T: ?Sized> OwnedMutexGuard<T> { + fn skip_drop(self) -> OwnedMutexGuardInner<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. + 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<U, F>(mut this: Self, f: F) -> OwnedMappedMutexGuard<T, U> + where + 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<U, F>(mut this: Self, f: F) -> Result<OwnedMappedMutexGuard<T, U>, 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(OwnedMappedMutexGuard { + data, + lock: inner.lock, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, + }) + } + /// Returns a reference to the original `Arc<Mutex>`. /// /// ``` @@ -903,6 +1119,8 @@ impl<T: ?Sized> OwnedMutexGuard<T> { impl<T: ?Sized> Drop for OwnedMutexGuard<T> { fn drop(&mut self) { + self.lock.s.release(1); + #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { tracing::trace!( @@ -910,7 +1128,6 @@ impl<T: ?Sized> Drop for OwnedMutexGuard<T> { locked = false, ); }); - self.lock.s.release(1) } } @@ -942,6 +1159,16 @@ impl<T: ?Sized + fmt::Display> fmt::Display for OwnedMutexGuard<T> { // === 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. @@ -956,12 +1183,13 @@ impl<'a, T: ?Sized> MappedMutexGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = this.s; - mem::forget(this); + let inner = this.skip_drop(); MappedMutexGuard { - s, + s: inner.s, data, - marker: marker::PhantomData, + marker: PhantomData, + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: inner.resource_span, } } @@ -983,19 +1211,28 @@ impl<'a, T: ?Sized> MappedMutexGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = this.s; - mem::forget(this); + let inner = this.skip_drop(); Ok(MappedMutexGuard { - s, + s: inner.s, data, - marker: marker::PhantomData, + 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) + self.s.release(1); + + #[cfg(all(tokio_unstable, feature = "tracing"))] + self.resource_span.in_scope(|| { + tracing::trace!( + target: "runtime::resource::state_update", + locked = false, + ); + }); } } @@ -1023,3 +1260,111 @@ impl<'a, T: ?Sized + fmt::Display> fmt::Display for MappedMutexGuard<'a, T> { fmt::Display::fmt(&**self, f) } } + +// === impl OwnedMappedMutexGuard === + +impl<T: ?Sized, U: ?Sized> OwnedMappedMutexGuard<T, U> { + fn skip_drop(self) -> OwnedMappedMutexGuardInner<T, U> { + 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<S, F>(mut this: Self, f: F) -> OwnedMappedMutexGuard<T, S> + 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<S, F>(mut this: Self, f: F) -> Result<OwnedMappedMutexGuard<T, S>, 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<T: ?Sized, U: ?Sized> Drop for OwnedMappedMutexGuard<T, U> { + 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<T: ?Sized, U: ?Sized> Deref for OwnedMappedMutexGuard<T, U> { + type Target = U; + fn deref(&self) -> &Self::Target { + unsafe { &*self.data } + } +} + +impl<T: ?Sized, U: ?Sized> DerefMut for OwnedMappedMutexGuard<T, U> { + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.data } + } +} + +impl<T: ?Sized, U: ?Sized + fmt::Debug> fmt::Debug for OwnedMappedMutexGuard<T, U> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&**self, f) + } +} + +impl<T: ?Sized, U: ?Sized + fmt::Display> fmt::Display for OwnedMappedMutexGuard<T, U> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&**self, f) + } +} diff --git a/src/sync/notify.rs b/src/sync/notify.rs index efe16f9..bf00ca3 100644 --- a/src/sync/notify.rs +++ b/src/sync/notify.rs @@ -5,21 +5,22 @@ // 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, LinkedList}; +use crate::util::linked_list::{self, GuardedLinkedList, LinkedList}; use crate::util::WakeList; -use std::cell::UnsafeCell; 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::SeqCst; +use std::sync::atomic::Ordering::{self, Acquire, Relaxed, Release, SeqCst}; use std::task::{Context, Poll, Waker}; type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; +type GuardedWaitList = GuardedLinkedList<Waiter, <Waiter as linked_list::Link>::Target>; /// Notifies a single task to wake up. /// @@ -198,37 +199,51 @@ type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>; /// [`Semaphore`]: crate::sync::Semaphore #[derive(Debug)] pub struct Notify { - // This uses 2 bits to store one of `EMPTY`, + // `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<WaitList>, } -#[derive(Debug, Clone, Copy)] -enum NotificationType { - // Notification triggered by calling `notify_waiters` - AllWaiters, - // Notification triggered by calling `notify_one` - OneWaiter, -} - #[derive(Debug)] struct Waiter { /// Intrusive linked-list pointers. pointers: linked_list::Pointers<Waiter>, - /// Waiting task's waker. - waker: Option<Waker>, + /// 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<Option<Waker>>, - /// `true` if the notification has been assigned to this waiter. - notified: Option<NotificationType>, + /// Notification for this waiter. + /// * 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<Self>) -> NonNull<linked_list::Pointers<Waiter>> { @@ -237,6 +252,109 @@ generate_addr_of_methods! { } } +// No notification. +const NOTIFICATION_NONE: usize = 0; + +// Notification type used by `notify_one`. +const NOTIFICATION_ONE: usize = 1; + +// Notification type used by `notify_waiters`. +const NOTIFICATION_ALL: usize = 2; + +/// Notification for a `Waiter`. +/// This struct is equivalent to `Option<Notification>`, 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) { + self.0.store(notification as usize, Release); + } + + fn load(&self, ordering: Ordering) -> Option<Notification> { + match self.0.load(ordering) { + NOTIFICATION_NONE => None, + NOTIFICATION_ONE => Some(Notification::One), + 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 Notification { + One = NOTIFICATION_ONE, + All = NOTIFICATION_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<NonNull<Waiter>> { + 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 @@ -249,8 +367,11 @@ pub struct Notified<'a> { /// 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: UnsafeCell<Waiter>, + waiter: Waiter, } unsafe impl<'a> Send for Notified<'a> {} @@ -258,7 +379,7 @@ unsafe impl<'a> Sync for Notified<'a> {} #[derive(Debug)] enum State { - Init(usize), + Init, Waiting, Done, } @@ -322,8 +443,7 @@ impl Notify { /// /// static NOTIFY: Notify = Notify::const_new(); /// ``` - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] + #[cfg(not(all(loom, test)))] pub const fn const_new() -> Notify { Notify { state: AtomicUsize::new(0), @@ -383,17 +503,13 @@ impl Notify { /// ``` pub fn notified(&self) -> Notified<'_> { // we load the number of times notify_waiters - // was called and store that in our initial state + // was called and store that in the future. let state = self.state.load(SeqCst); Notified { notify: self, - state: State::Init(state >> NOTIFY_WAITERS_SHIFT), - waiter: UnsafeCell::new(Waiter { - pointers: linked_list::Pointers::new(), - waker: None, - notified: None, - _p: PhantomPinned, - }), + state: State::Init, + notify_waiters_calls: get_num_notify_waiters_calls(state), + waiter: Waiter::new(), } } @@ -500,12 +616,9 @@ impl Notify { /// } /// ``` pub fn notify_waiters(&self) { - let mut wakers = WakeList::new(); - - // 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 + // 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); @@ -516,23 +629,43 @@ impl Notify { return; } - // At this point, it is guaranteed that the state will not - // concurrently change, as holding the lock is required to - // transition **out** of `WAITING`. + // 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 waiters.pop_back() { - Some(mut waiter) => { - // Safety: `waiters` lock is still held. - let waiter = unsafe { waiter.as_mut() }; - - assert!(waiter.notified.is_none()); - - waiter.notified = Some(NotificationType::AllWaiters); - - if let Some(waker) = waiter.waker.take() { + 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; @@ -540,20 +673,17 @@ impl Notify { } } + // 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(); } - // All waiters will be notified, the state must be transitioned to - // `EMPTY`. As transitioning **from** `WAITING` requires the lock to be - // held, a `store` is sufficient. - let new = set_state(inc_num_notify_waiters_calls(curr), EMPTY); - self.state.store(new, SeqCst); - // Release the lock before notifying drop(waiters); @@ -592,15 +722,16 @@ fn notify_locked(waiters: &mut WaitList, state: &AtomicUsize, curr: usize) -> Op // transition **out** of `WAITING`. // // Get a pending waiter - let mut waiter = waiters.pop_back().unwrap(); + let waiter = waiters.pop_back().unwrap(); - // Safety: `waiters` lock is still held. - let waiter = unsafe { waiter.as_mut() }; + // Safety: we never make mutable references to waiters. + let waiter = unsafe { waiter.as_ref() }; - assert!(waiter.notified.is_none()); + // Safety: we hold the lock, so we can access the waker. + let waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }; - waiter.notified = Some(NotificationType::OneWaiter); - let waker = waiter.waker.take(); + // This waiter is unlinked and will not be shared ever again, release it. + waiter.notification.store_release(Notification::One); if waiters.is_empty() { // As this the **final** waiter in the list, the state @@ -730,26 +861,32 @@ impl Notified<'_> { /// 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, &UnsafeCell<Waiter>) { + fn project(self: Pin<&mut Self>) -> (&Notify, &mut State, &usize, &Waiter) { unsafe { - // Safety: both `notify` and `state` are `Unpin`. + // Safety: `notify`, `state` and `notify_waiters_calls` are `Unpin`. is_unpin::<&Notify>(); - is_unpin::<AtomicUsize>(); + is_unpin::<State>(); + is_unpin::<usize>(); let me = self.get_unchecked_mut(); - (me.notify, &mut me.state, &me.waiter) + ( + me.notify, + &mut me.state, + &me.notify_waiters_calls, + &me.waiter, + ) } } fn poll_notified(self: Pin<&mut Self>, waker: Option<&Waker>) -> Poll<()> { use State::*; - let (notify, state, waiter) = self.project(); + let (notify, state, notify_waiters_calls, waiter) = self.project(); - loop { + 'outer_loop: loop { match *state { - Init(initial_notify_waiters_calls) => { + Init => { let curr = notify.state.load(SeqCst); // Optimistically try acquiring a pending notification @@ -763,7 +900,7 @@ impl Notified<'_> { if res.is_ok() { // Acquired the notification *state = Done; - return Poll::Ready(()); + continue 'outer_loop; } // Clone the waker before locking, a waker clone can be @@ -779,9 +916,9 @@ impl Notified<'_> { // if notify_waiters has been called after the future // was created, then we are done - if get_num_notify_waiters_calls(curr) != initial_notify_waiters_calls { + if get_num_notify_waiters_calls(curr) != *notify_waiters_calls { *state = Done; - return Poll::Ready(()); + continue 'outer_loop; } // Transition the state to WAITING. @@ -817,7 +954,7 @@ impl Notified<'_> { Ok(_) => { // Acquired the notification *state = Done; - return Poll::Ready(()); + continue 'outer_loop; } Err(actual) => { assert_eq!(get_state(actual), EMPTY); @@ -829,52 +966,109 @@ impl Notified<'_> { } } + 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 { - (*waiter.get()).waker = waker; + old_waker = + waiter.waker.with_mut(|v| std::mem::replace(&mut *v, waker)); } } // Insert the waiter into the linked list - // - // safety: pointers from `UnsafeCell` are never null. - waiters.push_front(unsafe { NonNull::new_unchecked(waiter.get()) }); + waiters.push_front(NonNull::from(waiter)); *state = Waiting; + drop(waiters); + drop(old_waker); + return Poll::Pending; } 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. + #[cfg(tokio_taskdump)] + if let Some(waker) = waker { + let mut ctx = Context::from_waker(waker); + 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()) }); - let waiters = notify.waiters.lock(); + waiter.notification.clear(); + *state = Done; + return Poll::Ready(()); + } - // Safety: called while locked - let w = unsafe { &mut *waiter.get() }; + // 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(); - if w.notified.is_some() { - // Our waker has been notified. Reset the fields and - // remove it from the list. - w.waker = None; - w.notified = None; + // 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 = 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 = Done; } else { - // Update the waker, if necessary. - if let Some(waker) = waker { - let should_update = match w.waker.as_ref() { - Some(current_waker) => !current_waker.will_wake(waker), - None => true, - }; - if should_update { - w.waker = Some(waker.clone()); - } + // 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; } @@ -884,8 +1078,16 @@ impl Notified<'_> { // is helpful to visualize the scope of the critical // section. drop(waiters); + + // Drop the old waker after releasing the lock. + drop(old_waker); } Done => { + #[cfg(tokio_taskdump)] + if let Some(waker) = waker { + let mut ctx = Context::from_waker(waker); + ready!(crate::trace::trace_leaf(&mut ctx)); + } return Poll::Ready(()); } } @@ -906,7 +1108,7 @@ impl Drop for Notified<'_> { use State::*; // Safety: The type only transitions to a "Waiting" state when pinned. - let (notify, state, waiter) = unsafe { Pin::new_unchecked(self).project() }; + 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 @@ -915,11 +1117,16 @@ impl Drop for Notified<'_> { 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: the waiter is only added to `waiters` by virtue of it - // being the only `LinkedList` available to the type. - unsafe { waiters.remove(NonNull::new_unchecked(waiter.get())) }; + // 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); @@ -929,13 +1136,7 @@ impl Drop for Notified<'_> { // 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. - // - // Safety: with the entry removed from the linked list, there can be - // no concurrent access to the entry - if matches!( - unsafe { (*waiter.get()).notified }, - Some(NotificationType::OneWaiter) - ) { + if notification == Some(Notification::One) { if let Some(waker) = notify_locked(&mut waiters, ¬ify.state, notify_state) { drop(waiters); waker.wake(); diff --git a/src/sync/once_cell.rs b/src/sync/once_cell.rs index 081efae..8a5463a 100644 --- a/src/sync/once_cell.rs +++ b/src/sync/once_cell.rs @@ -114,12 +114,10 @@ impl<T> Drop for OnceCell<T> { impl<T> From<T> for OnceCell<T> { fn from(value: T) -> Self { - let semaphore = Semaphore::new(0); - semaphore.close(); OnceCell { value_set: AtomicBool::new(true), value: UnsafeCell::new(MaybeUninit::new(value)), - semaphore, + semaphore: Semaphore::new_closed(), } } } @@ -139,6 +137,10 @@ impl<T> OnceCell<T> { /// 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 stablised, + // we can made this function available in const context, + // by creating `Semaphore::const_new_closed`. pub fn new_with(value: Option<T>) -> Self { if let Some(v) = value { OnceCell::from(v) @@ -171,8 +173,7 @@ impl<T> OnceCell<T> { /// assert_eq!(*result, 2); /// } /// ``` - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] + #[cfg(not(all(loom, test)))] pub const fn const_new() -> Self { OnceCell { value_set: AtomicBool::new(false), @@ -301,6 +302,8 @@ impl<T> OnceCell<T> { F: FnOnce() -> Fut, Fut: Future<Output = T>, { + crate::trace::async_trace_leaf().await; + if self.initialized() { // SAFETY: The OnceCell has been fully initialized. unsafe { self.get_unchecked() } @@ -349,6 +352,8 @@ impl<T> OnceCell<T> { F: FnOnce() -> Fut, Fut: Future<Output = Result<T, E>>, { + crate::trace::async_trace_leaf().await; + if self.initialized() { // SAFETY: The OnceCell has been fully initialized. unsafe { Ok(self.get_unchecked()) } diff --git a/src/sync/oneshot.rs b/src/sync/oneshot.rs index fcd7a32..af3cc85 100644 --- a/src/sync/oneshot.rs +++ b/src/sync/oneshot.rs @@ -12,6 +12,10 @@ //! 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 //! //! ``` @@ -786,6 +790,8 @@ impl<T> Sender<T> { /// } /// ``` 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)); @@ -795,7 +801,7 @@ impl<T> Sender<T> { if state.is_closed() { coop.made_progress(); - return Poll::Ready(()); + return Ready(()); } if state.is_tx_task_set() { @@ -1056,6 +1062,7 @@ impl<T> Receiver<T> { /// ``` #[track_caller] #[cfg(feature = "sync")] + #[cfg_attr(docsrs, doc(alias = "recv_blocking"))] pub fn blocking_recv(self) -> Result<T, RecvError> { crate::future::block_on(self) } @@ -1125,6 +1132,7 @@ impl<T> Inner<T> { } fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> { + ready!(crate::trace::trace_leaf(cx)); // Keep track of task budget let coop = ready!(crate::runtime::coop::poll_proceed(cx)); diff --git a/src/sync/rwlock.rs b/src/sync/rwlock.rs index 0492e4e..ba464ea 100644 --- a/src/sync/rwlock.rs +++ b/src/sync/rwlock.rs @@ -5,7 +5,6 @@ use crate::util::trace; use std::cell::UnsafeCell; use std::marker; use std::marker::PhantomData; -use std::mem::ManuallyDrop; use std::sync::Arc; pub(crate) mod owned_read_guard; @@ -86,7 +85,6 @@ const MAX_READS: u32 = 10; /// [`RwLockWriteGuard`]: struct@RwLockWriteGuard /// [`Send`]: trait@std::marker::Send /// [_write-preferring_]: https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock#Priority_policies -#[derive(Debug)] pub struct RwLock<T: ?Sized> { #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span, @@ -336,8 +334,7 @@ impl<T: ?Sized> RwLock<T> { /// /// static LOCK: RwLock<i32> = RwLock::const_new(5); /// ``` - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] + #[cfg(not(all(loom, test)))] pub const fn const_new(value: T) -> RwLock<T> where T: Sized, @@ -361,13 +358,13 @@ impl<T: ?Sized> RwLock<T> { /// /// static LOCK: RwLock<i32> = RwLock::const_with_max_readers(5, 1024); /// ``` - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] - pub const fn const_with_max_readers(value: T, mut max_reads: u32) -> RwLock<T> + #[cfg(not(all(loom, test)))] + pub const fn const_with_max_readers(value: T, max_reads: u32) -> RwLock<T> where T: Sized, { - max_reads &= MAX_READS; + assert!(max_reads <= MAX_READS); + RwLock { mr: max_reads, c: UnsafeCell::new(value), @@ -423,23 +420,33 @@ impl<T: ?Sized> RwLock<T> { /// } /// ``` 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 inner = trace::async_op( - || self.s.acquire(1), + let acquire_fut = trace::async_op( + move || acquire_fut, self.resource_span.clone(), "RwLock::read", "poll", false, ); - #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let inner = self.s.acquire(1); - - inner.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!() - }); + #[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(|| { @@ -450,13 +457,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - RwLockReadGuard { - s: &self.s, - data: self.c.get(), - marker: marker::PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span: self.resource_span.clone(), - } + guard } /// Blockingly locks this `RwLock` with shared read access. @@ -565,25 +566,38 @@ impl<T: ?Sized> RwLock<T> { /// ``` pub async fn read_owned(self: Arc<Self>) -> OwnedRwLockReadGuard<T> { #[cfg(all(tokio_unstable, feature = "tracing"))] - let inner = trace::async_op( - || self.s.acquire(1), - self.resource_span.clone(), + 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, ); - #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let inner = self.s.acquire(1); - - inner.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!() - }); + #[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(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, @@ -591,16 +605,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - - OwnedRwLockReadGuard { - data: self.c.get(), - lock: ManuallyDrop::new(self), - _p: PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - } + guard } /// Attempts to acquire this `RwLock` with shared read access. @@ -642,6 +647,14 @@ impl<T: ?Sized> RwLock<T> { 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!( @@ -651,13 +664,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - Ok(RwLockReadGuard { - s: &self.s, - data: self.c.get(), - marker: marker::PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span: self.resource_span.clone(), - }) + Ok(guard) } /// Attempts to acquire this `RwLock` with shared read access. @@ -705,8 +712,16 @@ impl<T: ?Sized> RwLock<T> { 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"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, @@ -714,16 +729,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - - Ok(OwnedRwLockReadGuard { - data: self.c.get(), - lock: ManuallyDrop::new(self), - _p: PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - }) + Ok(guard) } /// Locks this `RwLock` with exclusive write access, causing the current @@ -755,23 +761,34 @@ impl<T: ?Sized> RwLock<T> { ///} /// ``` pub async fn write(&self) -> RwLockWriteGuard<'_, T> { + let acquire_fut = async { + self.s.acquire(self.mr).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 inner = trace::async_op( - || self.s.acquire(self.mr), + let acquire_fut = trace::async_op( + move || acquire_fut, self.resource_span.clone(), "RwLock::write", "poll", false, ); - #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let inner = self.s.acquire(self.mr); - - inner.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!() - }); + #[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(|| { @@ -782,14 +799,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - 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(), - } + guard } /// Blockingly locks this `RwLock` with exclusive write access. @@ -884,25 +894,39 @@ impl<T: ?Sized> RwLock<T> { /// ``` pub async fn write_owned(self: Arc<Self>) -> OwnedRwLockWriteGuard<T> { #[cfg(all(tokio_unstable, feature = "tracing"))] - let inner = trace::async_op( - || self.s.acquire(self.mr), - self.resource_span.clone(), + let resource_span = self.resource_span.clone(); + + let acquire_fut = async { + self.s.acquire(self.mr).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, ); - #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let inner = self.s.acquire(self.mr); - - inner.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!() - }); + #[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(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = true, @@ -910,17 +934,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - - OwnedRwLockWriteGuard { - permits_acquired: self.mr, - data: self.c.get(), - lock: ManuallyDrop::new(self), - _p: PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - } + guard } /// Attempts to acquire this `RwLock` with exclusive write access. @@ -953,6 +967,15 @@ impl<T: ?Sized> RwLock<T> { 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!( @@ -962,14 +985,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - Ok(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(), - }) + Ok(guard) } /// Attempts to acquire this `RwLock` with exclusive write access. @@ -1009,8 +1025,17 @@ impl<T: ?Sized> RwLock<T> { 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"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = true, @@ -1018,17 +1043,7 @@ impl<T: ?Sized> RwLock<T> { ) }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - - Ok(OwnedRwLockWriteGuard { - permits_acquired: self.mr, - data: self.c.get(), - lock: ManuallyDrop::new(self), - _p: PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - }) + Ok(guard) } /// Returns a mutable reference to the underlying data. @@ -1078,3 +1093,17 @@ where Self::new(T::default()) } } + +impl<T: ?Sized> std::fmt::Debug for RwLock<T> +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!("<locked>")), + }; + d.finish() + } +} diff --git a/src/sync/rwlock/owned_read_guard.rs b/src/sync/rwlock/owned_read_guard.rs index 27b71bd..273e7b8 100644 --- a/src/sync/rwlock/owned_read_guard.rs +++ b/src/sync/rwlock/owned_read_guard.rs @@ -1,10 +1,7 @@ use crate::sync::rwlock::RwLock; -use std::fmt; use std::marker::PhantomData; -use std::mem; -use std::mem::ManuallyDrop; -use std::ops; 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. @@ -14,16 +11,40 @@ use std::sync::Arc; /// /// [`read_owned`]: method@crate::sync::RwLock::read_owned /// [`RwLock`]: struct@crate::sync::RwLock +#[clippy::has_significant_drop] pub struct OwnedRwLockReadGuard<T: ?Sized, U: ?Sized = T> { + // 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, - // ManuallyDrop allows us to destructure into this field without running the destructor. - pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>, + pub(super) lock: Arc<RwLock<T>>, pub(super) data: *const U, pub(super) _p: PhantomData<T>, } +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<T: ?Sized, U: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + lock: Arc<RwLock<T>>, + data: *const U, +} + impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { + fn skip_drop(self) -> Inner<T, U> { + 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. @@ -52,23 +73,19 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { /// # } /// ``` #[inline] - pub fn map<F, V: ?Sized>(mut this: Self, f: F) -> OwnedRwLockReadGuard<T, V> + pub fn map<F, V: ?Sized>(this: Self, f: F) -> OwnedRwLockReadGuard<T, V> where F: FnOnce(&U) -> &V, { let data = f(&*this) as *const V; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); OwnedRwLockReadGuard { - lock: ManuallyDrop::new(lock), + lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, } } @@ -103,7 +120,7 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { /// # } /// ``` #[inline] - pub fn try_map<F, V: ?Sized>(mut this: Self, f: F) -> Result<OwnedRwLockReadGuard<T, V>, Self> + pub fn try_map<F, V: ?Sized>(this: Self, f: F) -> Result<OwnedRwLockReadGuard<T, V>, Self> where F: FnOnce(&U) -> Option<&V>, { @@ -111,18 +128,14 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> { Some(data) => data as *const V, None => return Err(this), }; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); Ok(OwnedRwLockReadGuard { - lock: ManuallyDrop::new(lock), + lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, }) } } @@ -156,7 +169,6 @@ where impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockReadGuard<T, U> { fn drop(&mut self) { self.lock.s.release(1); - unsafe { ManuallyDrop::drop(&mut self.lock) }; #[cfg(all(tokio_unstable, feature = "tracing"))] self.resource_span.in_scope(|| { diff --git a/src/sync/rwlock/owned_write_guard.rs b/src/sync/rwlock/owned_write_guard.rs index dbedab4..a8ce4a1 100644 --- a/src/sync/rwlock/owned_write_guard.rs +++ b/src/sync/rwlock/owned_write_guard.rs @@ -1,11 +1,9 @@ use crate::sync::rwlock::owned_read_guard::OwnedRwLockReadGuard; use crate::sync::rwlock::owned_write_guard_mapped::OwnedRwLockMappedWriteGuard; use crate::sync::rwlock::RwLock; -use std::fmt; use std::marker::PhantomData; -use std::mem::{self, ManuallyDrop}; -use std::ops; 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. @@ -15,17 +13,43 @@ use std::sync::Arc; /// /// [`write_owned`]: method@crate::sync::RwLock::write_owned /// [`RwLock`]: struct@crate::sync::RwLock +#[clippy::has_significant_drop] pub struct OwnedRwLockWriteGuard<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, - // ManuallyDrop allows us to destructure into this field without running the destructor. - pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>, + pub(super) lock: Arc<RwLock<T>>, pub(super) data: *mut T, pub(super) _p: PhantomData<T>, } +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<T: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + permits_acquired: u32, + lock: Arc<RwLock<T>>, + data: *const T, +} + impl<T: ?Sized> OwnedRwLockWriteGuard<T> { + fn skip_drop(self) -> Inner<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. + 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. /// @@ -64,24 +88,90 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, } } - /// Attempts to make a new [`OwnedRwLockMappedWriteGuard`] for a component + /// 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<F, U: ?Sized>(this: Self, f: F) -> OwnedRwLockReadGuard<T, U> + 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`. /// @@ -128,24 +218,99 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { Some(data) => data as *mut U, None => return Err(this), }; - let permits_acquired = this.permits_acquired; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); Ok(OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + 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<F, U: ?Sized>( + this: Self, + f: F, + ) -> Result<OwnedRwLockReadGuard<T, U>, 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. @@ -191,15 +356,22 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { /// assert_eq!(*lock.read().await, 2, "second writer obtained write lock"); /// # } /// ``` - pub fn downgrade(mut self) -> OwnedRwLockReadGuard<T> { - let lock = unsafe { ManuallyDrop::take(&mut self.lock) }; - let data = self.data; - let to_release = (self.permits_acquired - 1) as usize; + pub fn downgrade(self) -> OwnedRwLockReadGuard<T> { + 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 - lock.s.release(to_release); + let to_release = (this.permits_acquired - 1) as usize; + guard.lock.s.release(to_release); + #[cfg(all(tokio_unstable, feature = "tracing"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, @@ -208,7 +380,7 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { }); #[cfg(all(tokio_unstable, feature = "tracing"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, @@ -216,18 +388,7 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> { ) }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(self); - - OwnedRwLockReadGuard { - lock: ManuallyDrop::new(lock), - data, - _p: PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - } + guard } } @@ -266,6 +427,7 @@ where impl<T: ?Sized> Drop for OwnedRwLockWriteGuard<T> { 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!( @@ -274,6 +436,5 @@ impl<T: ?Sized> Drop for OwnedRwLockWriteGuard<T> { write_locked.op = "override", ) }); - unsafe { ManuallyDrop::drop(&mut self.lock) }; } } diff --git a/src/sync/rwlock/owned_write_guard_mapped.rs b/src/sync/rwlock/owned_write_guard_mapped.rs index 55a24d9..9f49521 100644 --- a/src/sync/rwlock/owned_write_guard_mapped.rs +++ b/src/sync/rwlock/owned_write_guard_mapped.rs @@ -1,9 +1,7 @@ use crate::sync::rwlock::RwLock; -use std::fmt; use std::marker::PhantomData; -use std::mem::{self, ManuallyDrop}; -use std::ops; 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. @@ -14,17 +12,43 @@ use std::sync::Arc; /// /// [mapping]: method@crate::sync::OwnedRwLockWriteGuard::map /// [`OwnedRwLockWriteGuard`]: struct@crate::sync::OwnedRwLockWriteGuard +#[clippy::has_significant_drop] pub struct OwnedRwLockMappedWriteGuard<T: ?Sized, U: ?Sized = T> { + // 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, - // ManuallyDrop allows us to destructure into this field without running the destructor. - pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>, + pub(super) lock: Arc<RwLock<T>>, pub(super) data: *mut U, pub(super) _p: PhantomData<T>, } +#[allow(dead_code)] // Unused fields are still used in Drop. +struct Inner<T: ?Sized, U: ?Sized> { + #[cfg(all(tokio_unstable, feature = "tracing"))] + resource_span: tracing::Span, + permits_acquired: u32, + lock: Arc<RwLock<T>>, + data: *const U, +} + impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> { + fn skip_drop(self) -> Inner<T, U> { + 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. /// @@ -63,20 +87,15 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> { F: FnOnce(&mut U) -> &mut V, { let data = f(&mut *this) as *mut V; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, } } @@ -125,20 +144,15 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> { Some(data) => data as *mut V, None => return Err(this), }; - let lock = unsafe { ManuallyDrop::take(&mut this.lock) }; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); Ok(OwnedRwLockMappedWriteGuard { - permits_acquired, - lock: ManuallyDrop::new(lock), + permits_acquired: this.permits_acquired, + lock: this.lock, data, _p: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, }) } } @@ -178,6 +192,7 @@ where impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockMappedWriteGuard<T, U> { 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!( @@ -186,6 +201,5 @@ impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockMappedWriteGuard<T, U> { write_locked.op = "override", ) }); - unsafe { ManuallyDrop::drop(&mut self.lock) }; } } diff --git a/src/sync/rwlock/read_guard.rs b/src/sync/rwlock/read_guard.rs index f5fc1d6..a04b595 100644 --- a/src/sync/rwlock/read_guard.rs +++ b/src/sync/rwlock/read_guard.rs @@ -1,8 +1,6 @@ use crate::sync::batch_semaphore::Semaphore; -use std::fmt; -use std::marker; -use std::mem; -use std::ops; +use std::marker::PhantomData; +use std::{fmt, mem, ops}; /// RAII structure used to release the shared read access of a lock when /// dropped. @@ -12,16 +10,39 @@ use std::ops; /// /// [`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: marker::PhantomData<&'a 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 @@ -61,18 +82,14 @@ impl<'a, T: ?Sized> RwLockReadGuard<'a, T> { F: FnOnce(&T) -> &U, { let data = f(&*this) as *const U; - let s = this.s; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); RwLockReadGuard { - s, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, } } @@ -120,18 +137,14 @@ impl<'a, T: ?Sized> RwLockReadGuard<'a, T> { Some(data) => data as *const U, None => return Err(this), }; - let s = this.s; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); Ok(RwLockReadGuard { - s, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, }) } } diff --git a/src/sync/rwlock/write_guard.rs b/src/sync/rwlock/write_guard.rs index cefa183..d405fc2 100644 --- a/src/sync/rwlock/write_guard.rs +++ b/src/sync/rwlock/write_guard.rs @@ -1,10 +1,8 @@ use crate::sync::batch_semaphore::Semaphore; use crate::sync::rwlock::read_guard::RwLockReadGuard; use crate::sync::rwlock::write_guard_mapped::RwLockMappedWriteGuard; -use std::fmt; -use std::marker; -use std::mem; -use std::ops; +use std::marker::PhantomData; +use std::{fmt, mem, ops}; /// RAII structure used to release the exclusive write access of a lock when /// dropped. @@ -14,17 +12,42 @@ use std::ops; /// /// [`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: marker::PhantomData<&'a 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 @@ -67,23 +90,96 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = this.s; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, } } - /// Attempts to make a new [`RwLockMappedWriteGuard`] for a component of + /// 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<F, U: ?Sized>(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`. /// @@ -134,22 +230,102 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = this.s; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); + Ok(RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + 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<F, U: ?Sized>(this: Self, f: F) -> Result<RwLockReadGuard<'a, U>, 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. @@ -198,12 +374,21 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { /// /// [`RwLock`]: struct@crate::sync::RwLock pub fn downgrade(self) -> RwLockReadGuard<'a, T> { - let RwLockWriteGuard { s, data, .. } = self; - let to_release = (self.permits_acquired - 1) as usize; + 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 - s.release(to_release); + let to_release = (this.permits_acquired - 1) as usize; + this.s.release(to_release); + #[cfg(all(tokio_unstable, feature = "tracing"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", write_locked = false, @@ -212,7 +397,7 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { }); #[cfg(all(tokio_unstable, feature = "tracing"))] - self.resource_span.in_scope(|| { + guard.resource_span.in_scope(|| { tracing::trace!( target: "runtime::resource::state_update", current_readers = 1, @@ -220,18 +405,7 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> { ) }); - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = self.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(self); - - RwLockReadGuard { - s, - data, - marker: marker::PhantomData, - #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, - } + guard } } diff --git a/src/sync/rwlock/write_guard_mapped.rs b/src/sync/rwlock/write_guard_mapped.rs index b5c644a..7705189 100644 --- a/src/sync/rwlock/write_guard_mapped.rs +++ b/src/sync/rwlock/write_guard_mapped.rs @@ -1,8 +1,6 @@ use crate::sync::batch_semaphore::Semaphore; -use std::fmt; -use std::marker; -use std::mem; -use std::ops; +use std::marker::PhantomData; +use std::{fmt, mem, ops}; /// RAII structure used to release the exclusive write access of a lock when /// dropped. @@ -13,16 +11,41 @@ use std::ops; /// /// [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: marker::PhantomData<&'a 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 @@ -64,20 +87,15 @@ impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> { F: FnOnce(&mut T) -> &mut U, { let data = f(&mut *this) as *mut U; - let s = this.s; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + resource_span: this.resource_span, } } @@ -131,22 +149,20 @@ impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> { Some(data) => data as *mut U, None => return Err(this), }; - let s = this.s; - let permits_acquired = this.permits_acquired; - #[cfg(all(tokio_unstable, feature = "tracing"))] - let resource_span = this.resource_span.clone(); - // NB: Forget to avoid drop impl from being called. - mem::forget(this); + let this = this.skip_drop(); Ok(RwLockMappedWriteGuard { - permits_acquired, - s, + permits_acquired: this.permits_acquired, + s: this.s, data, - marker: marker::PhantomData, + marker: PhantomData, #[cfg(all(tokio_unstable, feature = "tracing"))] - resource_span, + 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<T: ?Sized> ops::Deref for RwLockMappedWriteGuard<'_, T> { diff --git a/src/sync/semaphore.rs b/src/sync/semaphore.rs index 6e5a1a8..cb77021 100644 --- a/src/sync/semaphore.rs +++ b/src/sync/semaphore.rs @@ -73,7 +73,7 @@ use std::sync::Arc; /// } /// ``` /// -/// [`PollSemaphore`]: https://docs.rs/tokio-util/0.6/tokio_util/sync/struct.PollSemaphore.html +/// [`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 { @@ -89,6 +89,7 @@ pub struct Semaphore { /// /// [`acquire`]: crate::sync::Semaphore::acquire() #[must_use] +#[clippy::has_significant_drop] #[derive(Debug)] pub struct SemaphorePermit<'a> { sem: &'a Semaphore, @@ -101,6 +102,7 @@ pub struct SemaphorePermit<'a> { /// /// [`acquire_owned`]: crate::sync::Semaphore::acquire_owned() #[must_use] +#[clippy::has_significant_drop] #[derive(Debug)] pub struct OwnedSemaphorePermit { sem: Arc<Semaphore>, @@ -123,7 +125,7 @@ fn bounds() { } impl Semaphore { - /// The maximum number of permits which a semaphore can hold. It is `usize::MAX >>> 3`. + /// 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; @@ -170,20 +172,22 @@ impl Semaphore { /// /// static SEM: Semaphore = Semaphore::const_new(10); /// ``` - /// - #[cfg(all(feature = "parking_lot", not(all(loom, test))))] - #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] + #[cfg(not(all(loom, test)))] pub const fn const_new(permits: usize) -> Self { - #[cfg(all(tokio_unstable, feature = "tracing"))] - return Self { + Self { ll_sem: ll::Semaphore::const_new(permits), + #[cfg(all(tokio_unstable, feature = "tracing"))] resource_span: tracing::Span::none(), - }; + } + } - #[cfg(any(not(tokio_unstable), not(feature = "tracing")))] - return Self { - ll_sem: ll::Semaphore::const_new(permits), - }; + /// 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(), + } } /// Returns the current number of available permits. @@ -674,6 +678,11 @@ impl OwnedSemaphorePermit { self.permits += other.permits; other.permits = 0; } + + /// Returns the [`Semaphore`] from which this permit was acquired. + pub fn semaphore(&self) -> &Arc<Semaphore> { + &self.sem + } } impl Drop for SemaphorePermit<'_> { diff --git a/src/sync/tests/atomic_waker.rs b/src/sync/tests/atomic_waker.rs index 8ebfb91..989175a 100644 --- a/src/sync/tests/atomic_waker.rs +++ b/src/sync/tests/atomic_waker.rs @@ -12,7 +12,7 @@ impl AssertSync for AtomicWaker {} impl AssertSend for Waker {} impl AssertSync for Waker {} -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[test] @@ -37,7 +37,7 @@ fn wake_without_register() { } #[test] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn atomic_waker_panic_safe() { use std::panic; use std::ptr; diff --git a/src/sync/tests/loom_notify.rs b/src/sync/tests/loom_notify.rs index d484a75..a4ded1d 100644 --- a/src/sync/tests/loom_notify.rs +++ b/src/sync/tests/loom_notify.rs @@ -4,6 +4,11 @@ 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(|| { @@ -138,3 +143,189 @@ fn notify_drop() { 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::<Vec<_>>(); + + 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::<Vec<_>>(); + + 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::<Vec<_>>(); + + // 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(); + }); +} diff --git a/src/sync/tests/loom_watch.rs b/src/sync/tests/loom_watch.rs index c575b5b..51589cd 100644 --- a/src/sync/tests/loom_watch.rs +++ b/src/sync/tests/loom_watch.rs @@ -2,6 +2,7 @@ use crate::sync::watch; use loom::future::block_on; use loom::thread; +use std::sync::Arc; #[test] fn smoke() { @@ -34,3 +35,56 @@ fn smoke() { 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(); + }); +} diff --git a/src/sync/tests/notify.rs b/src/sync/tests/notify.rs index eb0da8f..323b5c6 100644 --- a/src/sync/tests/notify.rs +++ b/src/sync/tests/notify.rs @@ -1,10 +1,9 @@ use crate::sync::Notify; use std::future::Future; -use std::mem::ManuallyDrop; use std::sync::Arc; use std::task::{Context, RawWaker, RawWakerVTable, Waker}; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; #[test] @@ -12,16 +11,16 @@ 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 arc = ManuallyDrop::new(Arc::<Notify>::from_raw(data as *const Notify)); + let ptr = data as *const Notify; + Arc::<Notify>::increment_strong_count(ptr); // Or some other arbitrary code that shouldn't be executed while the // Notify wait list is locked. - arc.notify_one(); - let _arc_clone: ManuallyDrop<_> = arc.clone(); + (*ptr).notify_one(); RawWaker::new(data, VTABLE) } unsafe fn drop_w(data: *const ()) { - let _ = Arc::<Notify>::from_raw(data as *const Notify); + drop(Arc::<Notify>::from_raw(data as *const Notify)); } unsafe fn wake(_data: *const ()) { @@ -46,6 +45,45 @@ fn notify_clones_waker_before_lock() { 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(Arc<Notify>); + + impl ArcWake for PanickingWaker { + fn wake_by_ref(_arc_self: &Arc<Self>) { + 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(); @@ -63,7 +101,7 @@ fn notify_simple() { } #[test] -#[cfg(not(tokio_wasm))] +#[cfg(not(target_family = "wasm"))] fn watch_test() { let rt = crate::runtime::Builder::new_current_thread() .build() diff --git a/src/sync/tests/semaphore_batch.rs b/src/sync/tests/semaphore_batch.rs index d4e35aa..391797b 100644 --- a/src/sync/tests/semaphore_batch.rs +++ b/src/sync/tests/semaphore_batch.rs @@ -1,7 +1,9 @@ use crate::sync::batch_semaphore::Semaphore; use tokio_test::*; -#[cfg(tokio_wasm_not_wasi)] +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] @@ -169,11 +171,15 @@ fn poll_acquire_one_zero_permits() { } #[test] +fn max_permits_doesnt_panic() { + Semaphore::new(MAX_PERMITS); +} + +#[test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn validates_max_permits() { - use std::usize; - Semaphore::new((usize::MAX >> 2) + 1); + Semaphore::new(MAX_PERMITS + 1); } #[test] @@ -252,3 +258,32 @@ fn cancel_acquire_releases_permits() { 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(Option<OwnedSemaphorePermit>); + + impl ArcWake for ReleaseOnDrop { + fn wake_by_ref(_arc_self: &Arc<Self>) {} + } + + 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()); + } +} diff --git a/src/sync/watch.rs b/src/sync/watch.rs index 6db46b7..ad64539 100644 --- a/src/sync/watch.rs +++ b/src/sync/watch.rs @@ -39,6 +39,9 @@ //! 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 @@ -58,6 +61,7 @@ use crate::sync::notify::Notify; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::atomic::Ordering::Relaxed; use crate::loom::sync::{Arc, RwLock, RwLockReadGuard}; +use std::fmt; use std::mem; use std::ops; use std::panic; @@ -166,7 +170,6 @@ impl<'a, T> Ref<'a, T> { } } -#[derive(Debug)] struct Shared<T> { /// The most recent value. value: RwLock<T>, @@ -181,30 +184,49 @@ struct Shared<T> { ref_count_rx: AtomicUsize, /// Notifies waiting receivers that the value changed. - notify_rx: Notify, + notify_rx: big_notify::BigNotify, /// Notifies any task listening for `Receiver` dropped events. notify_tx: Notify, } +impl<T: fmt::Debug> fmt::Debug for Shared<T> { + 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(Debug)] + #[derive(PartialEq, Eq, Clone, Copy)] pub struct SendError<T>(pub T); // ===== impl SendError ===== - impl<T: fmt::Debug> fmt::Display for SendError<T> { + impl<T> fmt::Debug for SendError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SendError").finish_non_exhaustive() + } + } + + impl<T> fmt::Display for SendError<T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { write!(fmt, "channel closed") } } - impl<T: fmt::Debug> std::error::Error for SendError<T> {} + impl<T> Error for SendError<T> {} /// Error produced when receiving a change notification. #[derive(Debug, Clone)] @@ -218,7 +240,63 @@ pub mod error { } } - impl std::error::Error for RecvError {} + impl Error for RecvError {} +} + +mod big_notify { + use super::*; + 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: 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: 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, 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}; @@ -320,7 +398,7 @@ pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) { value: RwLock::new(init), state: AtomicState::new(), ref_count_rx: AtomicUsize::new(1), - notify_rx: Notify::new(), + notify_rx: big_notify::BigNotify::new(), notify_tx: Notify::new(), }); @@ -524,18 +602,93 @@ impl<T> Receiver<T> { /// } /// ``` pub async fn changed(&mut self) -> Result<(), error::RecvError> { + changed_impl(&self.shared, &mut self.version).await + } + + /// Waits for a value that satisifes 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. + /// + /// [`Receiver::changed()`]: crate::sync::watch::Receiver::changed + /// + /// # Examples + /// + /// ``` + /// use tokio::sync::watch; + /// + /// #[tokio::main] + /// + /// async fn main() { + /// let (tx, _rx) = watch::channel("hello"); + /// + /// tx.send("goodbye").unwrap(); + /// + /// // here we subscribe to a second receiver + /// // now in case of using `changed` we would have + /// // to first check the current value and then wait + /// // for changes or else `changed` would hang. + /// let mut rx2 = tx.subscribe(); + /// + /// // in place of changed we have use `wait_for` + /// // which would automatically check the current value + /// // and wait for changes until the closure returns true. + /// assert!(rx2.wait_for(|val| *val == "goodbye").await.is_ok()); + /// assert_eq!(*rx2.borrow(), "goodbye"); + /// } + /// ``` + pub async fn wait_for( + &mut self, + mut f: impl FnMut(&T) -> bool, + ) -> Result<Ref<'_, T>, error::RecvError> { + let mut closed = false; 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 = self.shared.notify_rx.notified(); + { + let inner = self.shared.value.read().unwrap(); - if let Some(ret) = maybe_changed(&self.shared, &mut self.version) { - return ret; + let new_version = self.shared.state.load().version(); + let has_changed = self.version != new_version; + self.version = new_version; + + if (!closed || has_changed) && f(&inner) { + return Ok(Ref { inner, has_changed }); + } } - notified.await; - // loop around again in case the wake-up was spurious + if closed { + return Err(error::RecvError(())); + } + + // Wait for the value to change. + closed = changed_impl(&self.shared, &mut self.version).await.is_err(); } } @@ -584,6 +737,27 @@ fn maybe_changed<T>( None } +async fn changed_impl<T>( + shared: &Shared<T>, + 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<T> Clone for Receiver<T> { fn clone(&self) -> Self { let version = self.version; @@ -636,7 +810,7 @@ impl<T> Sender<T> { /// Modifies the watched value **unconditionally** in-place, /// notifying all receivers. /// - /// This can useful for modifying the watched value, without + /// 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. /// @@ -676,7 +850,7 @@ impl<T> Sender<T> { /// Modifies the watched value **conditionally** in-place, /// notifying all receivers only if modified. /// - /// This can useful for modifying the watched value, without + /// 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. /// @@ -867,6 +1041,8 @@ impl<T> Sender<T> { /// } /// ``` pub async fn closed(&self) { + crate::trace::async_trace_leaf().await; + while self.receiver_count() > 0 { let notified = self.shared.notify_tx.notified(); diff --git a/src/task/blocking.rs b/src/task/blocking.rs index 46756a9..1cce466 100644 --- a/src/task/blocking.rs +++ b/src/task/blocking.rs @@ -75,7 +75,7 @@ cfg_rt_multi_thread! { where F: FnOnce() -> R, { - crate::runtime::scheduler::multi_thread::block_in_place(f) + crate::runtime::scheduler::block_in_place(f) } } @@ -146,7 +146,7 @@ cfg_rt! { /// [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/0.6/tokio_util/io/struct.SyncIoBridge.html + /// [`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 diff --git a/src/task/builder.rs b/src/task/builder.rs index 91c400c..306cc39 100644 --- a/src/task/builder.rs +++ b/src/task/builder.rs @@ -80,7 +80,7 @@ impl<'a> Builder<'a> { /// /// This method panics if called outside of a Tokio runtime. /// - /// See [`task::spawn`](crate::task::spawn) for + /// See [`task::spawn`](crate::task::spawn()) for /// more details. #[track_caller] pub fn spawn<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>> diff --git a/src/task/consume_budget.rs b/src/task/consume_budget.rs index 1212cfc..e7432ff 100644 --- a/src/task/consume_budget.rs +++ b/src/task/consume_budget.rs @@ -33,6 +33,7 @@ pub async fn consume_budget() { let mut status = Poll::Pending; crate::future::poll_fn(move |cx| { + ready!(crate::trace::trace_leaf(cx)); if status.is_ready() { return status; } diff --git a/src/task/join_set.rs b/src/task/join_set.rs index e6d8d62..4eb15a2 100644 --- a/src/task/join_set.rs +++ b/src/task/join_set.rs @@ -201,6 +201,67 @@ impl<T: 'static> JoinSet<T> { 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<F>(&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<F>(&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<T>) -> AbortHandle { let abort = jh.abort_handle(); let mut entry = self.inner.insert_idle(jh); @@ -301,7 +362,7 @@ impl<T: 'static> JoinSet<T> { /// This can happen if the [coop budget] is reached. /// /// [coop budget]: crate::task#cooperative-scheduling - fn poll_join_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<T, JoinError>>> { + pub fn poll_join_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<T, JoinError>>> { // 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()) { @@ -358,7 +419,8 @@ impl<T: 'static> JoinSet<T> { /// [coop budget]: crate::task#cooperative-scheduling /// [task ID]: crate::task::Id #[cfg(tokio_unstable)] - fn poll_join_next_with_id( + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub fn poll_join_next_with_id( &mut self, cx: &mut Context<'_>, ) -> Poll<Option<Result<(Id, T), JoinError>>> { diff --git a/src/task/local.rs b/src/task/local.rs index 0675faa..59b68b8 100644 --- a/src/task/local.rs +++ b/src/task/local.rs @@ -1,6 +1,8 @@ //! 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}; use crate::runtime::{context, ThreadId}; use crate::sync::AtomicWaker; @@ -34,14 +36,14 @@ cfg_rt! { /// async fn main() { /// // `Rc` does not implement `Send`, and thus may not be sent between /// // threads safely. - /// let unsend_data = Rc::new("my unsend data..."); + /// let nonsend_data = Rc::new("my nonsend data..."); /// - /// let unsend_data = unsend_data.clone(); - /// // Because the `async` block here moves `unsend_data`, the future is `!Send`. + /// 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!("{}", unsend_data); + /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// } @@ -60,18 +62,18 @@ cfg_rt! { /// /// #[tokio::main] /// async fn main() { - /// let unsend_data = Rc::new("my unsend data..."); + /// 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 unsend_data = unsend_data.clone(); + /// 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!("{}", unsend_data); + /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// }).await; @@ -94,18 +96,18 @@ cfg_rt! { /// /// #[tokio::main] /// async fn main() { - /// let unsend_data = Rc::new("world"); + /// let nonsend_data = Rc::new("world"); /// let local = task::LocalSet::new(); /// - /// let unsend_data2 = unsend_data.clone(); + /// let nonsend_data2 = nonsend_data.clone(); /// local.spawn_local(async move { /// // ... - /// println!("hello {}", unsend_data2) + /// println!("hello {}", nonsend_data2) /// }); /// /// local.spawn_local(async move { /// time::sleep(time::Duration::from_millis(100)).await; - /// println!("goodbye {}", unsend_data) + /// println!("goodbye {}", nonsend_data) /// }); /// /// // ... @@ -299,7 +301,7 @@ cfg_rt! { /// /// 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 - /// use `spawn_local` if you want to stay within the `LocalSet`. + /// `spawn_local` if you want to stay within the `LocalSet`. /// /// # Examples /// @@ -309,15 +311,15 @@ cfg_rt! { /// /// #[tokio::main] /// async fn main() { - /// let unsend_data = Rc::new("my unsend data..."); + /// let nonsend_data = Rc::new("my nonsend data..."); /// /// let local = task::LocalSet::new(); /// /// // Run the local task set. /// local.run_until(async move { - /// let unsend_data = unsend_data.clone(); + /// let nonsend_data = nonsend_data.clone(); /// task::spawn_local(async move { - /// println!("{}", unsend_data); + /// println!("{}", nonsend_data); /// // ... /// }).await.unwrap(); /// }).await; @@ -785,6 +787,30 @@ cfg_unstable! { .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() + } } } diff --git a/src/task/spawn.rs b/src/task/spawn.rs index 66b0d67..23a46d7 100644 --- a/src/task/spawn.rs +++ b/src/task/spawn.rs @@ -1,4 +1,3 @@ -use crate::runtime::Handle; use crate::task::JoinHandle; use std::future::Future; @@ -16,6 +15,10 @@ cfg_rt! { /// 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. @@ -178,10 +181,26 @@ cfg_rt! { T: Future + Send + 'static, T::Output: Send + 'static, { - use crate::runtime::task; + 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", name, id.as_u64()); - let handle = Handle::current(); - handle.inner.spawn(task, id) + + match context::with_current(|handle| handle.spawn(task, id)) { + Ok(join_handle) => join_handle, + Err(e) => panic!("{}", e), + } } } diff --git a/src/task/task_local.rs b/src/task/task_local.rs index d3b108f..237c4d8 100644 --- a/src/task/task_local.rs +++ b/src/task/task_local.rs @@ -1,3 +1,4 @@ +use pin_project_lite::pin_project; use std::cell::RefCell; use std::error::Error; use std::future::Future; @@ -47,7 +48,6 @@ macro_rules! task_local { } #[doc(hidden)] -#[cfg(not(tokio_no_const_thread_local))] #[macro_export] macro_rules! __task_local_inner { ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => { @@ -62,22 +62,6 @@ macro_rules! __task_local_inner { }; } -#[doc(hidden)] -#[cfg(tokio_no_const_thread_local)] -#[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<Option<$t>> = 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. @@ -299,36 +283,53 @@ impl<T: 'static> fmt::Debug for LocalKey<T> { } } -/// 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; -/// # } -/// ``` -// Doesn't use pin_project due to custom Drop. -pub struct TaskLocalFuture<T, F> -where - T: 'static, -{ - local: &'static LocalKey<T>, - slot: Option<T>, - future: Option<F>, - _pinned: PhantomPinned, +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<T, F> + where + T: 'static, + { + local: &'static LocalKey<T>, + slot: Option<T>, + #[pin] + future: Option<F>, + #[pin] + _pinned: PhantomPinned, + } + + impl<T: 'static, F> PinnedDrop for TaskLocalFuture<T, F> { + fn drop(this: Pin<&mut Self>) { + let this = this.project(); + if mem::needs_drop::<F>() && this.future.is_some() { + // Drop the future while the task-local is set, if possible. Otherwise + // the future is dropped normally when the `Option<F>` field drops. + let mut future = this.future; + let _ = this.local.scope_inner(this.slot, || { + future.set(None); + }); + } + } + } } impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> { @@ -336,23 +337,21 @@ impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> { #[track_caller] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { - // safety: The TaskLocalFuture struct is `!Unpin` so there is no way to - // move `self.future` from now on. - let this = unsafe { Pin::into_inner_unchecked(self) }; - let mut future_opt = unsafe { Pin::new_unchecked(&mut this.future) }; - - let res = - this.local - .scope_inner(&mut 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) + 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); } - None => None, - }); + Some(res) + } + None => None, + }); match res { Ok(Some(res)) => res, @@ -362,19 +361,6 @@ impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> { } } -impl<T: 'static, F> Drop for TaskLocalFuture<T, F> { - fn drop(&mut self) { - if mem::needs_drop::<F>() && self.future.is_some() { - // Drop the future while the task-local is set, if possible. Otherwise - // the future is dropped normally when the `Option<F>` field drops. - let future = &mut self.future; - let _ = self.local.scope_inner(&mut self.slot, || { - *future = None; - }); - } - } -} - impl<T: 'static, F> fmt::Debug for TaskLocalFuture<T, F> where T: fmt::Debug, diff --git a/src/task/yield_now.rs b/src/task/yield_now.rs index 7b61dd8..428d124 100644 --- a/src/task/yield_now.rs +++ b/src/task/yield_now.rs @@ -46,21 +46,15 @@ pub async fn yield_now() { 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; - let defer = context::with_defer(|rt| { - rt.defer(cx.waker().clone()); - }); - - if defer.is_none() { - // Not currently in a runtime, just notify ourselves - // immediately. - cx.waker().wake_by_ref(); - } + context::defer(cx.waker()); Poll::Pending } diff --git a/src/time/clock.rs b/src/time/clock.rs index cd11a67..091cf4b 100644 --- a/src/time/clock.rs +++ b/src/time/clock.rs @@ -29,32 +29,53 @@ cfg_not_test_util! { cfg_test_util! { use crate::time::{Duration, Instant}; - use crate::loom::sync::{Arc, Mutex}; + use crate::loom::sync::Mutex; + use crate::loom::sync::atomic::Ordering; + use std::sync::atomic::AtomicBool as StdAtomicBool; cfg_rt! { - fn clock() -> Option<Clock> { + #[track_caller] + fn with_clock<R>(f: impl FnOnce(Option<&Clock>) -> Result<R, &'static str>) -> R { use crate::runtime::Handle; - match Handle::try_current() { - Ok(handle) => Some(handle.inner.driver().clock().clone()), - Err(ref e) if e.is_missing_context() => None, + 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! { - fn clock() -> Option<Clock> { - None + #[track_caller] + fn with_clock<R>(f: impl FnOnce(Option<&Clock>) -> Result<R, &'static str>) -> R { + match f(None) { + Ok(ret) => ret, + Err(msg) => panic!("{}", msg), + } } } /// A handle to a source of time. - #[derive(Debug, Clone)] + #[derive(Debug)] pub(crate) struct Clock { - inner: Arc<Mutex<Inner>>, + inner: Mutex<Inner>, } + // 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. @@ -107,8 +128,12 @@ cfg_test_util! { /// [`advance`]: crate::time::advance #[track_caller] pub fn pause() { - let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); - clock.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. @@ -122,14 +147,21 @@ cfg_test_util! { /// runtime. #[track_caller] pub fn resume() { - let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); - let mut inner = clock.inner.lock(); + with_clock(|maybe_clock| { + let clock = match maybe_clock { + Some(clock) => clock, + None => return Err("time cannot be frozen from outside the Tokio runtime"), + }; - if inner.unfrozen.is_some() { - panic!("time is not frozen"); - } + let mut inner = clock.inner.lock(); - inner.unfrozen = Some(std::time::Instant::now()); + if inner.unfrozen.is_some() { + return Err("time is not frozen"); + } + + inner.unfrozen = Some(std::time::Instant::now()); + Ok(()) + }) } /// Advances time. @@ -164,19 +196,31 @@ cfg_test_util! { /// /// [`sleep`]: fn@crate::time::sleep pub async fn advance(duration: Duration) { - let clock = clock().expect("time cannot be frozen from outside the Tokio runtime"); - clock.advance(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 let Some(clock) = clock() { - clock.now() - } else { - Instant::from_std(std::time::Instant::now()) + 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 { @@ -186,34 +230,43 @@ cfg_test_util! { let now = std::time::Instant::now(); let clock = Clock { - inner: Arc::new(Mutex::new(Inner { + inner: Mutex::new(Inner { enable_pausing, base: now, unfrozen: Some(now), auto_advance_inhibit_count: 0, - })), + }), }; if start_paused { - clock.pause(); + if let Err(msg) = clock.pause() { + panic!("{}", msg); + } } clock } - #[track_caller] - pub(crate) fn pause(&self) { + pub(crate) fn pause(&self) -> Result<(), &'static str> { let mut inner = self.inner.lock(); if !inner.enable_pausing { drop(inner); // avoid poisoning the lock - panic!("`time::pause()` requires the `current_thread` Tokio runtime. \ + return Err("`time::pause()` requires the `current_thread` Tokio runtime. \ This is the default Runtime used by `#[tokio::test]."); } - let elapsed = inner.unfrozen.as_ref().expect("time is already frozen").elapsed(); + // 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`). @@ -232,15 +285,15 @@ cfg_test_util! { inner.unfrozen.is_none() && inner.auto_advance_inhibit_count == 0 } - #[track_caller] - pub(crate) fn advance(&self, duration: Duration) { + pub(crate) fn advance(&self, duration: Duration) -> Result<(), &'static str> { let mut inner = self.inner.lock(); if inner.unfrozen.is_some() { - panic!("time is not frozen"); + return Err("time is not frozen"); } inner.base += duration; + Ok(()) } pub(crate) fn now(&self) -> Instant { diff --git a/src/time/instant.rs b/src/time/instant.rs index f184929..14cf6e5 100644 --- a/src/time/instant.rs +++ b/src/time/instant.rs @@ -188,7 +188,7 @@ impl ops::Sub<Duration> for Instant { type Output = Instant; fn sub(self, rhs: Duration) -> Instant { - Instant::from_std(self.std - rhs) + Instant::from_std(std::time::Instant::sub(self.std, rhs)) } } diff --git a/src/time/interval.rs b/src/time/interval.rs index ea8b393..48f81af 100644 --- a/src/time/interval.rs +++ b/src/time/interval.rs @@ -2,10 +2,10 @@ use crate::future::poll_fn; use crate::time::{sleep_until, Duration, Instant, Sleep}; use crate::util::trace; +use std::future::Future; use std::panic::Location; use std::pin::Pin; use std::task::{Context, Poll}; -use std::{convert::TryInto, future::Future}; /// Creates new [`Interval`] that yields with interval of `period`. The first /// tick completes immediately. The default [`MissedTickBehavior`] is @@ -387,7 +387,7 @@ impl Default for MissedTickBehavior { /// 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 +/// [`sleep`]: crate::time::sleep() #[derive(Debug)] pub struct Interval { /// Future that completes the next time the `Interval` yields a value. @@ -482,7 +482,10 @@ impl Interval { timeout + self.period }; - self.delay.as_mut().reset(next); + // 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) @@ -492,6 +495,8 @@ impl Interval { /// /// This method ignores [`MissedTickBehavior`] strategy. /// + /// This is equivalent to calling `reset_at(Instant::now() + period)`. + /// /// # Examples /// /// ``` @@ -518,6 +523,107 @@ impl Interval { 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 diff --git a/src/time/sleep.rs b/src/time/sleep.rs index 0a012e2..9809571 100644 --- a/src/time/sleep.rs +++ b/src/time/sleep.rs @@ -220,6 +220,7 @@ pin_project! { /// [`select!`]: ../macro.select.html /// [`tokio::pin!`]: ../macro.pin.html // Alias for old name in 0.2 + #[project(!Unpin)] #[cfg_attr(docsrs, doc(alias = "Delay"))] #[derive(Debug)] #[must_use = "futures do nothing unless you `.await` or poll them"] @@ -235,7 +236,6 @@ pin_project! { cfg_trace! { #[derive(Debug)] struct Inner { - deadline: Instant, ctx: trace::AsyncOpTracingCtx, } } @@ -243,7 +243,6 @@ cfg_trace! { cfg_not_trace! { #[derive(Debug)] struct Inner { - deadline: Instant, } } @@ -261,10 +260,11 @@ impl Sleep { #[cfg(all(tokio_unstable, feature = "tracing"))] let inner = { + 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()); + 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!( @@ -296,11 +296,11 @@ impl Sleep { resource_span, }; - Inner { deadline, ctx } + Inner { ctx } }; #[cfg(not(all(tokio_unstable, feature = "tracing")))] - let inner = Inner { deadline }; + let inner = Inner {}; Sleep { inner, entry } } @@ -311,7 +311,7 @@ impl Sleep { /// Returns the instant at which the future will complete. pub fn deadline(&self) -> Instant { - self.inner.deadline + self.entry.deadline() } /// Returns `true` if `Sleep` has elapsed. @@ -354,10 +354,22 @@ impl Sleep { 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); - (me.inner).deadline = deadline; + me.entry.as_mut().reset(deadline, true); #[cfg(all(tokio_unstable, feature = "tracing"))] { @@ -370,8 +382,9 @@ impl Sleep { 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(); + let now = time_source.now(clock); let deadline_tick = time_source.deadline_to_tick(deadline); deadline_tick.saturating_sub(now) }; @@ -388,6 +401,8 @@ impl Sleep { fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> { 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!( diff --git a/src/time/timeout.rs b/src/time/timeout.rs index 3bb98ea..52ab989 100644 --- a/src/time/timeout.rs +++ b/src/time/timeout.rs @@ -28,6 +28,11 @@ use std::task::{self, Poll}; /// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, 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 /// @@ -100,6 +105,11 @@ where /// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, 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 /// diff --git a/src/util/bit.rs b/src/util/bit.rs index a43c2c2..72e25e4 100644 --- a/src/util/bit.rs +++ b/src/util/bit.rs @@ -37,14 +37,6 @@ impl Pack { (base & !self.mask) | (value << self.shift) } - /// Packs the value with `base`, losing any bits of `value` that fit. - /// - /// If `value` is larger than the max value that can be represented by the - /// allotted width, the most significant bits are truncated. - pub(crate) fn pack_lossy(&self, value: usize, base: usize) -> usize { - self.pack(value & self.max_value(), base) - } - pub(crate) fn unpack(&self, src: usize) -> usize { unpack(src, self.mask, self.shift) } diff --git a/src/util/cacheline.rs b/src/util/cacheline.rs new file mode 100644 index 0000000..64fd5cc --- /dev/null +++ b/src/util/cacheline.rs @@ -0,0 +1,95 @@ +#![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, mips64, and riscv64 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/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_riscv64.go#L7 +#[cfg_attr( + any( + target_arch = "arm", + target_arch = "mips", + target_arch = "mips64", + target_arch = "riscv64", + ), + 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 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 +// +// 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 = "riscv64", + target_arch = "s390x", + )), + repr(align(64)) +)] +pub(crate) struct CachePadded<T> { + value: T, +} + +impl<T> CachePadded<T> { + /// Pads and aligns a value to the length of a cache line. + pub(crate) fn new(value: T) -> CachePadded<T> { + CachePadded::<T> { value } + } +} + +impl<T> Deref for CachePadded<T> { + type Target = T; + + fn deref(&self) -> &T { + &self.value + } +} + +impl<T> DerefMut for CachePadded<T> { + fn deref_mut(&mut self) -> &mut T { + &mut self.value + } +} diff --git a/src/util/idle_notified_set.rs b/src/util/idle_notified_set.rs index ce8ff9e..19b81e2 100644 --- a/src/util/idle_notified_set.rs +++ b/src/util/idle_notified_set.rs @@ -421,7 +421,7 @@ impl<T: 'static> Wake for ListEntry<T> { // We move ourself to the notified list. let me = unsafe { // Safety: We just checked that we are in this particular list. - lock.idle.remove(NonNull::from(&**me)).unwrap() + lock.idle.remove(ListEntry::as_raw(me)).unwrap() }; lock.notified.push_front(me); @@ -460,3 +460,22 @@ unsafe impl<T> linked_list::Link for ListEntry<T> { 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(); + } +} diff --git a/src/util/linked_list.rs b/src/util/linked_list.rs index b46bd6d..1f9bdf4 100644 --- a/src/util/linked_list.rs +++ b/src/util/linked_list.rs @@ -178,8 +178,12 @@ impl<L: Link> LinkedList<L, L::Target> { /// /// # Safety /// - /// The caller **must** ensure that `node` is currently contained by - /// `self` or not contained by any other list. + /// 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<L::Target>) -> Option<L::Handle> { if let Some(prev) = L::pointers(node).as_ref().get_prev() { debug_assert_eq!(L::pointers(prev).as_ref().get_next(), Some(node)); @@ -224,6 +228,53 @@ impl<L: Link> fmt::Debug for LinkedList<L, L::Target> { } } +// ===== impl CountedLinkedList ==== + +// Delegates operations to the base LinkedList implementation, and adds a counter to the elements +// in the list. +pub(crate) struct CountedLinkedList<L: Link, T> { + list: LinkedList<L, T>, + count: usize, +} + +impl<L: Link> CountedLinkedList<L, L::Target> { + pub(crate) fn new() -> CountedLinkedList<L, L::Target> { + CountedLinkedList { + list: LinkedList::new(), + count: 0, + } + } + + pub(crate) fn push_front(&mut self, val: L::Handle) { + self.list.push_front(val); + self.count += 1; + } + + pub(crate) fn pop_back(&mut self) -> Option<L::Handle> { + let val = self.list.pop_back(); + if val.is_some() { + self.count -= 1; + } + val + } + + pub(crate) fn is_empty(&self) -> bool { + self.list.is_empty() + } + + pub(crate) unsafe fn remove(&mut self, node: NonNull<L::Target>) -> Option<L::Handle> { + let val = self.list.remove(node); + if val.is_some() { + self.count -= 1; + } + val + } + + pub(crate) fn count(&self) -> usize { + self.count + } +} + #[cfg(any( feature = "fs", feature = "rt", @@ -246,7 +297,7 @@ impl<L: Link> Default for LinkedList<L, L::Target> { // ===== impl DrainFilter ===== -cfg_io_readiness! { +cfg_io_driver_impl! { pub(crate) struct DrainFilter<'a, T: Link, F> { list: &'a mut LinkedList<T, T::Target>, filter: F, @@ -256,7 +307,7 @@ cfg_io_readiness! { impl<T: Link> LinkedList<T, T::Target> { pub(crate) fn drain_filter<F>(&mut self, filter: F) -> DrainFilter<'_, T, F> where - F: FnMut(&mut T::Target) -> bool, + F: FnMut(&T::Target) -> bool, { let curr = self.head; DrainFilter { @@ -270,7 +321,7 @@ cfg_io_readiness! { impl<'a, T, F> Iterator for DrainFilter<'a, T, F> where T: Link, - F: FnMut(&mut T::Target) -> bool, + F: FnMut(&T::Target) -> bool, { type Item = T::Handle; @@ -290,6 +341,126 @@ cfg_io_readiness! { } } +cfg_taskdump! { + impl<T: Link> CountedLinkedList<T, T::Target> { + pub(crate) fn for_each<F>(&mut self, f: F) + where + F: FnMut(&T::Handle), + { + self.list.for_each(f) + } + } + + impl<T: Link> LinkedList<T, T::Target> { + pub(crate) fn for_each<F>(&mut self, mut f: F) + where + F: FnMut(&T::Handle), + { + use std::mem::ManuallyDrop; + + 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<L, T> { + /// Pointer to the guard node. + guard: NonNull<T>, + + /// Node type marker. + _marker: PhantomData<*const L>, + } + + impl<U, L: Link<Handle = NonNull<U>>> LinkedList<L, L::Target> { + /// 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<L, L::Target> { + // `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<L: Link> GuardedLinkedList<L, L::Target> { + fn tail(&self) -> Option<NonNull<L::Target>> { + 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<L::Handle> { + 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<T> Pointers<T> { @@ -352,9 +523,9 @@ impl<T> fmt::Debug for Pointers<T> { } } -#[cfg(test)] +#[cfg(any(test, fuzzing))] #[cfg(not(loom))] -mod tests { +pub(crate) mod tests { use super::*; use std::pin::Pin; @@ -413,6 +584,7 @@ mod tests { } } + #[cfg(test)] macro_rules! assert_clean { ($e:ident) => {{ assert!($e.pointers.get_next().is_none()); @@ -420,6 +592,7 @@ mod tests { }}; } + #[cfg(test)] macro_rules! assert_ptr_eq { ($a:expr, $b:expr) => {{ // Deal with mapping a Pin<&mut T> -> Option<NonNull<T>> @@ -623,31 +796,42 @@ mod tests { } } - #[cfg(not(tokio_wasm))] - proptest::proptest! { - #[test] - fn fuzz_linked_list(ops: Vec<usize>) { - run_fuzz(ops); + #[test] + fn count() { + let mut list = CountedLinkedList::<&Entry, <&Entry as Link>::Target>::new(); + assert_eq!(0, list.count()); + + let a = entry(5); + let b = entry(7); + list.push_front(a.as_ref()); + list.push_front(b.as_ref()); + assert_eq!(2, list.count()); + + list.pop_back(); + assert_eq!(1, list.count()); + + unsafe { + list.remove(ptr(&b)); } + assert_eq!(0, list.count()); } - #[cfg(not(tokio_wasm))] - fn run_fuzz(ops: Vec<usize>) { - use std::collections::VecDeque; - - #[derive(Debug)] + /// 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 % 3 { + .map(|i| match i % 3u8 { 0 => Op::Push, 1 => Op::Pop, - 2 => Op::Remove(i / 3), + 2 => Op::Remove((i / 3u8) as usize), _ => unreachable!(), }) .collect::<Vec<_>>(); diff --git a/src/util/markers.rs b/src/util/markers.rs new file mode 100644 index 0000000..7031fb6 --- /dev/null +++ b/src/util/markers.rs @@ -0,0 +1,8 @@ +/// Marker for types that are `Sync` but not `Send` +pub(crate) struct SyncNotSend(*mut ()); + +unsafe impl Sync for SyncNotSend {} + +cfg_rt! { + pub(crate) struct NotSendOrSync(*mut ()); +} diff --git a/src/util/memchr.rs b/src/util/memchr.rs new file mode 100644 index 0000000..44fd2da --- /dev/null +++ b/src/util/memchr.rs @@ -0,0 +1,74 @@ +//! 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<usize> { + haystack.iter().position(|val| needle == *val) +} + +#[cfg(all(unix, feature = "libc"))] +pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> { + 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); + } + } +} diff --git a/src/util/mod.rs b/src/util/mod.rs index 9f6119a..7d4cd5f 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,17 +1,11 @@ cfg_io_driver! { pub(crate) mod bit; - pub(crate) mod slab; } #[cfg(feature = "rt")] pub(crate) mod atomic_cell; -#[cfg(any( - feature = "rt", - feature = "signal", - feature = "process", - tokio_no_const_mutex_new, -))] +#[cfg(any(feature = "rt", feature = "signal", feature = "process"))] pub(crate) mod once_cell; #[cfg(any( @@ -76,3 +70,10 @@ cfg_rt_multi_thread! { pub(crate) mod trace; pub(crate) mod error; + +#[cfg(feature = "io-util")] +pub(crate) mod memchr; + +pub(crate) mod markers; + +pub(crate) mod cacheline; diff --git a/src/util/once_cell.rs b/src/util/once_cell.rs index 1925f0a..71fc007 100644 --- a/src/util/once_cell.rs +++ b/src/util/once_cell.rs @@ -19,7 +19,7 @@ impl<T> OnceCell<T> { } } - /// Get the value inside this cell, intiailizing it using the provided + /// 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 diff --git a/src/util/pad.rs b/src/util/pad.rs deleted file mode 100644 index bf0913c..0000000 --- a/src/util/pad.rs +++ /dev/null @@ -1,52 +0,0 @@ -use core::fmt; -use core::ops::{Deref, DerefMut}; - -#[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 -#[cfg_attr(target_arch = "x86_64", repr(align(128)))] -#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))] -pub(crate) struct CachePadded<T> { - value: T, -} - -unsafe impl<T: Send> Send for CachePadded<T> {} -unsafe impl<T: Sync> Sync for CachePadded<T> {} - -impl<T> CachePadded<T> { - pub(crate) fn new(t: T) -> CachePadded<T> { - CachePadded::<T> { value: t } - } -} - -impl<T> Deref for CachePadded<T> { - type Target = T; - - fn deref(&self) -> &T { - &self.value - } -} - -impl<T> DerefMut for CachePadded<T> { - fn deref_mut(&mut self) -> &mut T { - &mut self.value - } -} - -impl<T: fmt::Debug> fmt::Debug for CachePadded<T> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CachePadded") - .field("value", &self.value) - .finish() - } -} - -impl<T> From<T> for CachePadded<T> { - fn from(t: T) -> Self { - CachePadded::new(t) - } -} diff --git a/src/util/rand.rs b/src/util/rand.rs index 749da6b..d96c8d3 100644 --- a/src/util/rand.rs +++ b/src/util/rand.rs @@ -1,48 +1,9 @@ -use std::cell::Cell; - cfg_rt! { - 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<FastRand>, - } - - impl RngSeedGenerator { - /// Returns a new generator from the provided seed. - pub(crate) fn new(seed: RngSeed) -> Self { - Self { - state: Mutex::new(FastRand::new(seed)), - } - } - - /// Returns the next seed in the sequence. - pub(crate) fn next_seed(&self) -> RngSeed { - let rng = self - .state - .lock() - .expect("RNG seed generator is internally corrupt"); + mod rt; + pub(crate) use rt::RngSeedGenerator; - 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()) - } + cfg_unstable! { + mod rt_unstable; } } @@ -57,31 +18,25 @@ pub struct RngSeed { 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: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf> +/// This generator passes the SmallCrush suite, part of TestU01 framework: +/// <http://simul.iro.umontreal.ca/testu01/tu01.html> +#[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()) } - cfg_unstable! { - /// Generates a seed from the provided byte slice. - /// - /// # Example - /// - /// ``` - /// # use tokio::runtime::RngSeed; - /// let seed = RngSeed::from_bytes(b"make me a seed"); - /// ``` - #[cfg(feature = "rt")] - pub fn from_bytes(bytes: &[u8]) -> Self { - use std::{collections::hash_map::DefaultHasher, hash::Hasher}; - - let mut hasher = DefaultHasher::default(); - hasher.write(bytes); - Self::from_u64(hasher.finish()) - } - } - fn from_u64(seed: u64) -> Self { let one = (seed >> 32) as u32; let mut two = seed as u32; @@ -98,60 +53,42 @@ impl RngSeed { Self { s, r } } } -/// 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: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf> -/// This generator passes the SmallCrush suite, part of TestU01 framework: -/// <http://simul.iro.umontreal.ca/testu01/tu01.html> -#[derive(Debug)] -pub(crate) struct FastRand { - one: Cell<u32>, - two: Cell<u32>, -} 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 new(seed: RngSeed) -> FastRand { + pub(crate) fn from_seed(seed: RngSeed) -> FastRand { FastRand { - one: Cell::new(seed.s), - two: Cell::new(seed.r), + one: seed.s, + two: seed.r, } } - /// 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. - #[cfg(feature = "rt")] - pub(crate) fn replace_seed(&self, seed: RngSeed) -> RngSeed { - let old_seed = RngSeed::from_pair(self.one.get(), self.two.get()); - - self.one.replace(seed.s); - self.two.replace(seed.r); - - old_seed - } - - #[cfg(any(feature = "macros", feature = "rt-multi-thread"))] - pub(crate) fn fastrand_n(&self, n: u32) -> u32 { + #[cfg(any( + feature = "macros", + feature = "rt-multi-thread", + 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(&self) -> u32 { - let mut s1 = self.one.get(); - let s0 = self.two.get(); + 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.set(s0); - self.two.set(s1); + self.one = s0; + self.two = s1; s0.wrapping_add(s1) } diff --git a/src/util/rand/rt.rs b/src/util/rand/rt.rs new file mode 100644 index 0000000..6584ba4 --- /dev/null +++ b/src/util/rand/rt.rs @@ -0,0 +1,61 @@ +use 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<FastRand>, +} + +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 + } +} diff --git a/src/util/rand/rt_unstable.rs b/src/util/rand/rt_unstable.rs new file mode 100644 index 0000000..3a8613e --- /dev/null +++ b/src/util/rand/rt_unstable.rs @@ -0,0 +1,20 @@ +use 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()) + } +} diff --git a/src/util/slab.rs b/src/util/slab.rs deleted file mode 100644 index 0e16e40..0000000 --- a/src/util/slab.rs +++ /dev/null @@ -1,854 +0,0 @@ -#![cfg_attr(not(feature = "rt"), allow(dead_code))] - -use crate::loom::cell::UnsafeCell; -use crate::loom::sync::atomic::{AtomicBool, AtomicUsize}; -use crate::loom::sync::{Arc, Mutex}; -use crate::util::bit; -use std::fmt; -use std::mem; -use std::ops; -use std::ptr; -use std::sync::atomic::Ordering::Relaxed; - -/// Amortized allocation for homogeneous data types. -/// -/// The slab pre-allocates chunks of memory to store values. It uses a similar -/// growing strategy as `Vec`. When new capacity is needed, the slab grows by -/// 2x. -/// -/// # Pages -/// -/// Unlike `Vec`, growing does not require moving existing elements. Instead of -/// being a continuous chunk of memory for all elements, `Slab` is an array of -/// arrays. The top-level array is an array of pages. Each page is 2x bigger -/// than the previous one. When the slab grows, a new page is allocated. -/// -/// Pages are lazily initialized. -/// -/// # Allocating -/// -/// When allocating an object, first previously used slots are reused. If no -/// previously used slot is available, a new slot is initialized in an existing -/// page. If all pages are full, then a new page is allocated. -/// -/// When an allocated object is released, it is pushed into it's page's free -/// list. Allocating scans all pages for a free slot. -/// -/// # Indexing -/// -/// The slab is able to index values using an address. Even when the indexed -/// object has been released, it is still safe to index. This is a key ability -/// for using the slab with the I/O driver. Addresses are registered with the -/// OS's selector and I/O resources can be released without synchronizing with -/// the OS. -/// -/// # Compaction -/// -/// `Slab::compact` will release pages that have been allocated but are no -/// longer used. This is done by scanning the pages and finding pages with no -/// allocated objects. These pages are then freed. -/// -/// # Synchronization -/// -/// The `Slab` structure is able to provide (mostly) unsynchronized reads to -/// values stored in the slab. Insertions and removals are synchronized. Reading -/// objects via `Ref` is fully unsynchronized. Indexing objects uses amortized -/// synchronization. -/// -pub(crate) struct Slab<T> { - /// Array of pages. Each page is synchronized. - pages: [Arc<Page<T>>; NUM_PAGES], - - /// Caches the array pointer & number of initialized slots. - cached: [CachedPage<T>; NUM_PAGES], -} - -/// Allocate values in the associated slab. -pub(crate) struct Allocator<T> { - /// Pages in the slab. The first page has a capacity of 16 elements. Each - /// following page has double the capacity of the previous page. - /// - /// Each returned `Ref` holds a reference count to this `Arc`. - pages: [Arc<Page<T>>; NUM_PAGES], -} - -/// References a slot in the slab. Indexing a slot using an `Address` is memory -/// safe even if the slot has been released or the page has been deallocated. -/// However, it is not guaranteed that the slot has not been reused and is now -/// represents a different value. -/// -/// The I/O driver uses a counter to track the slot's generation. Once accessing -/// the slot, the generations are compared. If they match, the value matches the -/// address. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub(crate) struct Address(usize); - -/// An entry in the slab. -pub(crate) trait Entry: Default { - /// Resets the entry's value and track the generation. - fn reset(&self); -} - -/// A reference to a value stored in the slab. -pub(crate) struct Ref<T> { - value: *const Value<T>, -} - -/// Maximum number of pages a slab can contain. -const NUM_PAGES: usize = 19; - -/// Minimum number of slots a page can contain. -const PAGE_INITIAL_SIZE: usize = 32; -const PAGE_INDEX_SHIFT: u32 = PAGE_INITIAL_SIZE.trailing_zeros() + 1; - -/// A page in the slab. -struct Page<T> { - /// Slots. - slots: Mutex<Slots<T>>, - - // Number of slots currently being used. This is not guaranteed to be up to - // date and should only be used as a hint. - used: AtomicUsize, - - // Set to `true` when the page has been allocated. - allocated: AtomicBool, - - // The number of slots the page can hold. - len: usize, - - // Length of all previous pages combined. - prev_len: usize, -} - -struct CachedPage<T> { - /// Pointer to the page's slots. - slots: *const Slot<T>, - - /// Number of initialized slots. - init: usize, -} - -/// Page state. -struct Slots<T> { - /// Slots. - slots: Vec<Slot<T>>, - - head: usize, - - /// Number of slots currently in use. - used: usize, -} - -unsafe impl<T: Sync> Sync for Page<T> {} -unsafe impl<T: Sync> Send for Page<T> {} -unsafe impl<T: Sync> Sync for CachedPage<T> {} -unsafe impl<T: Sync> Send for CachedPage<T> {} -unsafe impl<T: Sync> Sync for Ref<T> {} -unsafe impl<T: Sync> Send for Ref<T> {} - -/// A slot in the slab. Contains slot-specific metadata. -/// -/// `#[repr(C)]` guarantees that the struct starts w/ `value`. We use pointer -/// math to map a value pointer to an index in the page. -#[repr(C)] -struct Slot<T> { - /// Pointed to by `Ref`. - value: UnsafeCell<Value<T>>, - - /// Next entry in the free list. - next: u32, - - /// Makes miri happy by making mutable references not take exclusive access. - /// - /// Could probably also be fixed by replacing `slots` with a raw-pointer - /// based equivalent. - _pin: std::marker::PhantomPinned, -} - -/// Value paired with a reference to the page. -struct Value<T> { - /// Value stored in the value. - value: T, - - /// Pointer to the page containing the slot. - /// - /// A raw pointer is used as this creates a ref cycle. - page: *const Page<T>, -} - -impl<T> Slab<T> { - /// Create a new, empty, slab. - pub(crate) fn new() -> Slab<T> { - // Initializing arrays is a bit annoying. Instead of manually writing - // out an array and every single entry, `Default::default()` is used to - // initialize the array, then the array is iterated and each value is - // initialized. - let mut slab = Slab { - pages: Default::default(), - cached: Default::default(), - }; - - let mut len = PAGE_INITIAL_SIZE; - let mut prev_len: usize = 0; - - for page in &mut slab.pages { - let page = Arc::get_mut(page).unwrap(); - page.len = len; - page.prev_len = prev_len; - len *= 2; - prev_len += page.len; - - // Ensure we don't exceed the max address space. - debug_assert!( - page.len - 1 + page.prev_len < (1 << 24), - "max = {:b}", - page.len - 1 + page.prev_len - ); - } - - slab - } - - /// Returns a new `Allocator`. - /// - /// The `Allocator` supports concurrent allocation of objects. - pub(crate) fn allocator(&self) -> Allocator<T> { - Allocator { - pages: self.pages.clone(), - } - } - - /// Returns a reference to the value stored at the given address. - /// - /// `&mut self` is used as the call may update internal cached state. - pub(crate) fn get(&mut self, addr: Address) -> Option<&T> { - let page_idx = addr.page(); - let slot_idx = self.pages[page_idx].slot(addr); - - // If the address references a slot that was last seen as uninitialized, - // the `CachedPage` is updated. This requires acquiring the page lock - // and updating the slot pointer and initialized offset. - if self.cached[page_idx].init <= slot_idx { - self.cached[page_idx].refresh(&self.pages[page_idx]); - } - - // If the address **still** references an uninitialized slot, then the - // address is invalid and `None` is returned. - if self.cached[page_idx].init <= slot_idx { - return None; - } - - // Get a reference to the value. The lifetime of the returned reference - // is bound to `&self`. The only way to invalidate the underlying memory - // is to call `compact()`. The lifetimes prevent calling `compact()` - // while references to values are outstanding. - // - // The referenced data is never mutated. Only `&self` references are - // used and the data is `Sync`. - Some(self.cached[page_idx].get(slot_idx)) - } - - /// Calls the given function with a reference to each slot in the slab. The - /// slot may not be in-use. - /// - /// This is used by the I/O driver during the shutdown process to notify - /// each pending task. - pub(crate) fn for_each(&mut self, mut f: impl FnMut(&T)) { - for page_idx in 0..self.pages.len() { - // It is required to avoid holding the lock when calling the - // provided function. The function may attempt to acquire the lock - // itself. If we hold the lock here while calling `f`, a deadlock - // situation is possible. - // - // Instead of iterating the slots directly in `page`, which would - // require holding the lock, the cache is updated and the slots are - // iterated from the cache. - self.cached[page_idx].refresh(&self.pages[page_idx]); - - for slot_idx in 0..self.cached[page_idx].init { - f(self.cached[page_idx].get(slot_idx)); - } - } - } - - // Release memory back to the allocator. - // - // If pages are empty, the underlying memory is released back to the - // allocator. - pub(crate) fn compact(&mut self) { - // Iterate each page except the very first one. The very first page is - // never freed. - for (idx, page) in self.pages.iter().enumerate().skip(1) { - if page.used.load(Relaxed) != 0 || !page.allocated.load(Relaxed) { - // If the page has slots in use or the memory has not been - // allocated then it cannot be compacted. - continue; - } - - let mut slots = match page.slots.try_lock() { - Some(slots) => slots, - // If the lock cannot be acquired due to being held by another - // thread, don't try to compact the page. - _ => continue, - }; - - if slots.used > 0 || slots.slots.capacity() == 0 { - // The page is in use or it has not yet been allocated. Either - // way, there is no more work to do. - continue; - } - - page.allocated.store(false, Relaxed); - - // Remove the slots vector from the page. This is done so that the - // freeing process is done outside of the lock's critical section. - let vec = mem::take(&mut slots.slots); - slots.head = 0; - - // Drop the lock so we can drop the vector outside the lock below. - drop(slots); - - debug_assert!( - self.cached[idx].slots.is_null() || self.cached[idx].slots == vec.as_ptr(), - "cached = {:?}; actual = {:?}", - self.cached[idx].slots, - vec.as_ptr(), - ); - - // Clear cache - self.cached[idx].slots = ptr::null(); - self.cached[idx].init = 0; - - drop(vec); - } - } -} - -impl<T> fmt::Debug for Slab<T> { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - debug(fmt, "Slab", &self.pages[..]) - } -} - -impl<T: Entry> Allocator<T> { - /// Allocate a new entry and return a handle to the entry. - /// - /// Scans pages from smallest to biggest, stopping when a slot is found. - /// Pages are allocated if necessary. - /// - /// Returns `None` if the slab is full. - pub(crate) fn allocate(&self) -> Option<(Address, Ref<T>)> { - // Find the first available slot. - for page in &self.pages[..] { - if let Some((addr, val)) = Page::allocate(page) { - return Some((addr, val)); - } - } - - None - } -} - -impl<T> fmt::Debug for Allocator<T> { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - debug(fmt, "slab::Allocator", &self.pages[..]) - } -} - -impl<T> ops::Deref for Ref<T> { - type Target = T; - - fn deref(&self) -> &T { - // Safety: `&mut` is never handed out to the underlying value. The page - // is not freed until all `Ref` values are dropped. - unsafe { &(*self.value).value } - } -} - -impl<T> Drop for Ref<T> { - fn drop(&mut self) { - // Safety: `&mut` is never handed out to the underlying value. The page - // is not freed until all `Ref` values are dropped. - let _ = unsafe { (*self.value).release() }; - } -} - -impl<T: fmt::Debug> fmt::Debug for Ref<T> { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - (**self).fmt(fmt) - } -} - -impl<T: Entry> Page<T> { - // Allocates an object, returns the ref and address. - // - // `self: &Arc<Page<T>>` is avoided here as this would not work with the - // loom `Arc`. - fn allocate(me: &Arc<Page<T>>) -> Option<(Address, Ref<T>)> { - // Before acquiring the lock, use the `used` hint. - if me.used.load(Relaxed) == me.len { - return None; - } - - // Allocating objects requires synchronization - let mut locked = me.slots.lock(); - - if locked.head < locked.slots.len() { - // Re-use an already initialized slot. - // - // Help out the borrow checker - let locked = &mut *locked; - - // Get the index of the slot at the head of the free stack. This is - // the slot that will be reused. - let idx = locked.head; - let slot = &locked.slots[idx]; - - // Update the free stack head to point to the next slot. - locked.head = slot.next as usize; - - // Increment the number of used slots - locked.used += 1; - me.used.store(locked.used, Relaxed); - - // Reset the slot - slot.value.with(|ptr| unsafe { (*ptr).value.reset() }); - - // Return a reference to the slot - Some((me.addr(idx), locked.gen_ref(idx, me))) - } else if me.len == locked.slots.len() { - // The page is full - None - } else { - // No initialized slots are available, but the page has more - // capacity. Initialize a new slot. - let idx = locked.slots.len(); - - if idx == 0 { - // The page has not yet been allocated. Allocate the storage for - // all page slots. - locked.slots.reserve_exact(me.len); - } - - // Initialize a new slot - locked.slots.push(Slot { - value: UnsafeCell::new(Value { - value: Default::default(), - page: Arc::as_ptr(me), - }), - next: 0, - _pin: std::marker::PhantomPinned, - }); - - // Increment the head to indicate the free stack is empty - locked.head += 1; - - // Increment the number of used slots - locked.used += 1; - me.used.store(locked.used, Relaxed); - me.allocated.store(true, Relaxed); - - debug_assert_eq!(locked.slots.len(), locked.head); - - Some((me.addr(idx), locked.gen_ref(idx, me))) - } - } -} - -impl<T> Page<T> { - /// Returns the slot index within the current page referenced by the given - /// address. - fn slot(&self, addr: Address) -> usize { - addr.0 - self.prev_len - } - - /// Returns the address for the given slot. - fn addr(&self, slot: usize) -> Address { - Address(slot + self.prev_len) - } -} - -impl<T> Default for Page<T> { - fn default() -> Page<T> { - Page { - used: AtomicUsize::new(0), - allocated: AtomicBool::new(false), - slots: Mutex::new(Slots { - slots: Vec::new(), - head: 0, - used: 0, - }), - len: 0, - prev_len: 0, - } - } -} - -impl<T> Page<T> { - /// Release a slot into the page's free list. - fn release(&self, value: *const Value<T>) { - let mut locked = self.slots.lock(); - - let idx = locked.index_for(value); - locked.slots[idx].next = locked.head as u32; - locked.head = idx; - locked.used -= 1; - - self.used.store(locked.used, Relaxed); - } -} - -impl<T> CachedPage<T> { - /// Refreshes the cache. - fn refresh(&mut self, page: &Page<T>) { - let slots = page.slots.lock(); - - if !slots.slots.is_empty() { - self.slots = slots.slots.as_ptr(); - self.init = slots.slots.len(); - } - } - - /// Gets a value by index. - fn get(&self, idx: usize) -> &T { - assert!(idx < self.init); - - // Safety: Pages are allocated concurrently, but are only ever - // **deallocated** by `Slab`. `Slab` will always have a more - // conservative view on the state of the slot array. Once `CachedPage` - // sees a slot pointer and initialized offset, it will remain valid - // until `compact()` is called. The `compact()` function also updates - // `CachedPage`. - unsafe { - let slot = self.slots.add(idx); - let value = slot as *const Value<T>; - - &(*value).value - } - } -} - -impl<T> Default for CachedPage<T> { - fn default() -> CachedPage<T> { - CachedPage { - slots: ptr::null(), - init: 0, - } - } -} - -impl<T> Slots<T> { - /// Maps a slot pointer to an offset within the current page. - /// - /// The pointer math removes the `usize` index from the `Ref` struct, - /// shrinking the struct to a single pointer size. The contents of the - /// function is safe, the resulting `usize` is bounds checked before being - /// used. - /// - /// # Panics - /// - /// panics if the provided slot pointer is not contained by the page. - fn index_for(&self, slot: *const Value<T>) -> usize { - use std::mem; - - assert_ne!(self.slots.capacity(), 0, "page is unallocated"); - - let base = self.slots.as_ptr() as usize; - let slot = slot as usize; - let width = mem::size_of::<Slot<T>>(); - - assert!(slot >= base, "unexpected pointer"); - - let idx = (slot - base) / width; - assert!(idx < self.slots.len() as usize); - - idx - } - - /// Generates a `Ref` for the slot at the given index. This involves bumping the page's ref count. - fn gen_ref(&self, idx: usize, page: &Arc<Page<T>>) -> Ref<T> { - assert!(idx < self.slots.len()); - mem::forget(page.clone()); - - let vec_ptr = self.slots.as_ptr(); - let slot: *const Slot<T> = unsafe { vec_ptr.add(idx) }; - let value: *const Value<T> = slot as *const Value<T>; - - Ref { value } - } -} - -impl<T> Value<T> { - /// Releases the slot, returning the `Arc<Page<T>>` logically owned by the ref. - fn release(&self) -> Arc<Page<T>> { - // Safety: called by `Ref`, which owns an `Arc<Page<T>>` instance. - let page = unsafe { Arc::from_raw(self.page) }; - page.release(self as *const _); - page - } -} - -impl Address { - fn page(self) -> usize { - // Since every page is twice as large as the previous page, and all page - // sizes are powers of two, we can determine the page index that - // contains a given address by shifting the address down by the smallest - // page size and looking at how many twos places necessary to represent - // that number, telling us what power of two page size it fits inside - // of. We can determine the number of twos places by counting the number - // of leading zeros (unused twos places) in the number's binary - // representation, and subtracting that count from the total number of - // bits in a word. - let slot_shifted = (self.0 + PAGE_INITIAL_SIZE) >> PAGE_INDEX_SHIFT; - (bit::pointer_width() - slot_shifted.leading_zeros()) as usize - } - - pub(crate) const fn as_usize(self) -> usize { - self.0 - } - - pub(crate) fn from_usize(src: usize) -> Address { - Address(src) - } -} - -fn debug<T>(fmt: &mut fmt::Formatter<'_>, name: &str, pages: &[Arc<Page<T>>]) -> fmt::Result { - let mut capacity = 0; - let mut len = 0; - - for page in pages { - if page.allocated.load(Relaxed) { - capacity += page.len; - len += page.used.load(Relaxed); - } - } - - fmt.debug_struct(name) - .field("len", &len) - .field("capacity", &capacity) - .finish() -} - -#[cfg(all(test, not(loom)))] -mod test { - use super::*; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering::SeqCst; - - struct Foo { - cnt: AtomicUsize, - id: AtomicUsize, - } - - impl Default for Foo { - fn default() -> Foo { - Foo { - cnt: AtomicUsize::new(0), - id: AtomicUsize::new(0), - } - } - } - - impl Entry for Foo { - fn reset(&self) { - self.cnt.fetch_add(1, SeqCst); - } - } - - #[test] - fn insert_remove() { - let mut slab = Slab::<Foo>::new(); - let alloc = slab.allocator(); - - let (addr1, foo1) = alloc.allocate().unwrap(); - foo1.id.store(1, SeqCst); - assert_eq!(0, foo1.cnt.load(SeqCst)); - - let (addr2, foo2) = alloc.allocate().unwrap(); - foo2.id.store(2, SeqCst); - assert_eq!(0, foo2.cnt.load(SeqCst)); - - assert_eq!(1, slab.get(addr1).unwrap().id.load(SeqCst)); - assert_eq!(2, slab.get(addr2).unwrap().id.load(SeqCst)); - - drop(foo1); - - assert_eq!(1, slab.get(addr1).unwrap().id.load(SeqCst)); - - let (addr3, foo3) = alloc.allocate().unwrap(); - assert_eq!(addr3, addr1); - assert_eq!(1, foo3.cnt.load(SeqCst)); - foo3.id.store(3, SeqCst); - assert_eq!(3, slab.get(addr3).unwrap().id.load(SeqCst)); - - drop(foo2); - drop(foo3); - - slab.compact(); - - // The first page is never released - assert!(slab.get(addr1).is_some()); - assert!(slab.get(addr2).is_some()); - assert!(slab.get(addr3).is_some()); - } - - #[test] - fn insert_many() { - const MANY: usize = normal_or_miri(10_000, 50); - - let mut slab = Slab::<Foo>::new(); - let alloc = slab.allocator(); - let mut entries = vec![]; - - for i in 0..MANY { - let (addr, val) = alloc.allocate().unwrap(); - val.id.store(i, SeqCst); - entries.push((addr, val)); - } - - for (i, (addr, v)) in entries.iter().enumerate() { - assert_eq!(i, v.id.load(SeqCst)); - assert_eq!(i, slab.get(*addr).unwrap().id.load(SeqCst)); - } - - entries.clear(); - - for i in 0..MANY { - let (addr, val) = alloc.allocate().unwrap(); - val.id.store(MANY - i, SeqCst); - entries.push((addr, val)); - } - - for (i, (addr, v)) in entries.iter().enumerate() { - assert_eq!(MANY - i, v.id.load(SeqCst)); - assert_eq!(MANY - i, slab.get(*addr).unwrap().id.load(SeqCst)); - } - } - - #[test] - fn insert_drop_reverse() { - let mut slab = Slab::<Foo>::new(); - let alloc = slab.allocator(); - let mut entries = vec![]; - - for i in 0..normal_or_miri(10_000, 100) { - let (addr, val) = alloc.allocate().unwrap(); - val.id.store(i, SeqCst); - entries.push((addr, val)); - } - - for _ in 0..10 { - // Drop 1000 in reverse - for _ in 0..normal_or_miri(1_000, 10) { - entries.pop(); - } - - // Check remaining - for (i, (addr, v)) in entries.iter().enumerate() { - assert_eq!(i, v.id.load(SeqCst)); - assert_eq!(i, slab.get(*addr).unwrap().id.load(SeqCst)); - } - } - } - - #[test] - fn no_compaction_if_page_still_in_use() { - let mut slab = Slab::<Foo>::new(); - let alloc = slab.allocator(); - let mut entries1 = vec![]; - let mut entries2 = vec![]; - - for i in 0..normal_or_miri(10_000, 100) { - let (addr, val) = alloc.allocate().unwrap(); - val.id.store(i, SeqCst); - - if i % 2 == 0 { - entries1.push((addr, val, i)); - } else { - entries2.push(val); - } - } - - drop(entries2); - - for (addr, _, i) in &entries1 { - assert_eq!(*i, slab.get(*addr).unwrap().id.load(SeqCst)); - } - } - - const fn normal_or_miri(normal: usize, miri: usize) -> usize { - if cfg!(miri) { - miri - } else { - normal - } - } - - #[test] - fn compact_all() { - let mut slab = Slab::<Foo>::new(); - let alloc = slab.allocator(); - let mut entries = vec![]; - - for _ in 0..2 { - entries.clear(); - - for i in 0..normal_or_miri(10_000, 100) { - let (addr, val) = alloc.allocate().unwrap(); - val.id.store(i, SeqCst); - - entries.push((addr, val)); - } - - let mut addrs = vec![]; - - for (addr, _) in entries.drain(..) { - addrs.push(addr); - } - - slab.compact(); - - // The first page is never freed - for addr in &addrs[PAGE_INITIAL_SIZE..] { - assert!(slab.get(*addr).is_none()); - } - } - } - - #[test] - fn issue_3014() { - let mut slab = Slab::<Foo>::new(); - let alloc = slab.allocator(); - let mut entries = vec![]; - - for _ in 0..normal_or_miri(5, 2) { - entries.clear(); - - // Allocate a few pages + 1 - for i in 0..(32 + 64 + 128 + 1) { - let (addr, val) = alloc.allocate().unwrap(); - val.id.store(i, SeqCst); - - entries.push((addr, val, i)); - } - - for (addr, val, i) in &entries { - assert_eq!(*i, val.id.load(SeqCst)); - assert_eq!(*i, slab.get(*addr).unwrap().id.load(SeqCst)); - } - - // Release the last entry - entries.pop(); - - // Compact - slab.compact(); - - // Check all the addresses - - for (addr, val, i) in &entries { - assert_eq!(*i, val.id.load(SeqCst)); - assert_eq!(*i, slab.get(*addr).unwrap().id.load(SeqCst)); - } - } - } -} diff --git a/src/util/trace.rs b/src/util/trace.rs index 76e8a6c..be5602d 100644 --- a/src/util/trace.rs +++ b/src/util/trace.rs @@ -11,18 +11,22 @@ cfg_trace! { #[inline] #[track_caller] pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>, id: u64) -> Instrumented<F> { + #[track_caller] + fn get_span(kind: &'static str, name: Option<&str>, id: u64) -> tracing::Span { + let location = std::panic::Location::caller(); + tracing::trace_span!( + target: "tokio::task", + "runtime.spawn", + %kind, + task.name = %name.unwrap_or_default(), + task.id = id, + loc.file = location.file(), + loc.line = location.line(), + loc.col = location.column(), + ) + } use tracing::instrument::Instrument; - let location = std::panic::Location::caller(); - let span = tracing::trace_span!( - target: "tokio::task", - "runtime.spawn", - %kind, - task.name = %name.unwrap_or_default(), - task.id = id, - loc.file = location.file(), - loc.line = location.line(), - loc.col = location.column(), - ); + let span = get_span(kind, name, id); task.instrument(span) } diff --git a/src/util/wake.rs b/src/util/wake.rs index 5526cbc..c872ce5 100644 --- a/src/util/wake.rs +++ b/src/util/wake.rs @@ -50,16 +50,8 @@ fn waker_vtable<W: Wake>() -> &'static RawWakerVTable { ) } -unsafe fn inc_ref_count<T: Wake>(data: *const ()) { - // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop - let arc = ManuallyDrop::new(Arc::<T>::from_raw(data as *const T)); - - // Now increase refcount, but don't drop new refcount either - let _arc_clone: ManuallyDrop<_> = arc.clone(); -} - unsafe fn clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker { - inc_ref_count::<T>(data); + Arc::<T>::increment_strong_count(data as *const T); RawWaker::new(data, waker_vtable::<T>()) } diff --git a/tests/_require_full.rs b/tests/_require_full.rs index 4b9698a..d33943a 100644 --- a/tests/_require_full.rs +++ b/tests/_require_full.rs @@ -1,4 +1,4 @@ -#[cfg(not(any(feature = "full", tokio_wasm)))] +#[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 diff --git a/tests/async_send_sync.rs b/tests/async_send_sync.rs index e46d5c8..dfd26f9 100644 --- a/tests/async_send_sync.rs +++ b/tests/async_send_sync.rs @@ -130,7 +130,7 @@ macro_rules! assert_value { macro_rules! cfg_not_wasi { ($($item:item)*) => { $( - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] $item )* } @@ -263,6 +263,19 @@ mod unix_datagram { 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::*; @@ -334,6 +347,15 @@ assert_value!(tokio::sync::OnceCell<YY>: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard<NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard<YN>: Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard<YY>: Send & Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<NN,NN>: !Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<NN,YN>: !Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<NN,YY>: !Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<YN,NN>: !Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<YN,YN>: Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<YN,YY>: Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<YY,NN>: !Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<YY,YN>: Send & !Sync & Unpin); +assert_value!(tokio::sync::OwnedMappedMutexGuard<YY,YY>: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockMappedWriteGuard<NN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockMappedWriteGuard<YN>: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedRwLockMappedWriteGuard<YY>: Send & Sync & Unpin); @@ -510,7 +532,7 @@ async_assert_fn!(tokio::task::unconstrained(BoxFutureSend<()>): Send & !Sync & U 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::EnterGuard<'_>: !Send & Sync & Unpin); assert_value!(tokio::runtime::Handle: Send & Sync & Unpin); assert_value!(tokio::runtime::Runtime: Send & Sync & Unpin); diff --git a/tests/buffered.rs b/tests/buffered.rs index 4251c3f..27e75d1 100644 --- a/tests/buffered.rs +++ b/tests/buffered.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support bind() +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support bind() use tokio::net::TcpListener; use tokio_test::assert_ok; diff --git a/tests/dump.rs b/tests/dump.rs new file mode 100644 index 0000000..658ee4b --- /dev/null +++ b/tests/dump.rs @@ -0,0 +1,99 @@ +#![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 trace = task.trace().to_string(); + eprintln!("\n\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 trace = task.trace().to_string(); + eprintln!("\n\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() => {}, + ); + }); +} diff --git a/tests/fs.rs b/tests/fs.rs index ba38b71..9f739f6 100644 --- a/tests/fs.rs +++ b/tests/fs.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support file operations +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support file operations use tokio::fs; use tokio_test::assert_ok; diff --git a/tests/fs_canonicalize_dir.rs b/tests/fs_canonicalize_dir.rs new file mode 100644 index 0000000..e7f6c68 --- /dev/null +++ b/tests/fs_canonicalize_dir.rs @@ -0,0 +1,22 @@ +#![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:\\")); +} diff --git a/tests/fs_copy.rs b/tests/fs_copy.rs index 04678cf..3311710 100644 --- a/tests/fs_copy.rs +++ b/tests/fs_copy.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support file operations +#![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tempfile::tempdir; use tokio::fs; diff --git a/tests/fs_dir.rs b/tests/fs_dir.rs index f197a40..3f28c1a 100644 --- a/tests/fs_dir.rs +++ b/tests/fs_dir.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support directory operations +#![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}; @@ -46,6 +46,27 @@ async fn build_dir() { } #[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"); @@ -85,3 +106,21 @@ async fn read_inherent() { 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()); +} diff --git a/tests/fs_file.rs b/tests/fs_file.rs index 603ccad..40bd4fc 100644 --- a/tests/fs_file.rs +++ b/tests/fs_file.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support file operations +#![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use std::io::prelude::*; use tempfile::NamedTempFile; @@ -87,13 +87,93 @@ async fn coop() { 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] #[cfg(unix)] -async fn unix_fd() { +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(); @@ -102,6 +182,26 @@ async fn unix_fd() { } #[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; diff --git a/tests/fs_link.rs b/tests/fs_link.rs index d198abc..48da5f6 100644 --- a/tests/fs_link.rs +++ b/tests/fs_link.rs @@ -1,10 +1,9 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support file operations +#![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations use tokio::fs; -use std::io::prelude::*; -use std::io::BufReader; +use std::io::Write; use tempfile::tempdir; #[tokio::test] @@ -13,24 +12,23 @@ async fn test_hard_link() { let src = dir.path().join("src.txt"); let dst = dir.path().join("dst.txt"); - { - let mut file = std::fs::File::create(&src).unwrap(); - file.write_all(b"hello").unwrap(); - } + std::fs::File::create(&src) + .unwrap() + .write_all(b"hello") + .unwrap(); - let dst_2 = dst.clone(); + fs::hard_link(&src, &dst).await.unwrap(); - assert!(fs::hard_link(src, dst_2.clone()).await.is_ok()); + std::fs::File::create(&src) + .unwrap() + .write_all(b"new-data") + .unwrap(); - let mut content = String::new(); + let content = fs::read(&dst).await.unwrap(); + assert_eq!(content, b"new-data"); - { - let file = std::fs::File::open(dst).unwrap(); - let mut reader = BufReader::new(file); - reader.read_to_string(&mut content).unwrap(); - } - - assert!(content == "hello"); + // test that this is not a symlink: + assert!(fs::read_link(&dst).await.is_err()); } #[cfg(unix)] @@ -40,25 +38,20 @@ async fn test_symlink() { let src = dir.path().join("src.txt"); let dst = dir.path().join("dst.txt"); - { - let mut file = std::fs::File::create(&src).unwrap(); - file.write_all(b"hello").unwrap(); - } - - let src_2 = src.clone(); - let dst_2 = dst.clone(); - - assert!(fs::symlink(src_2.clone(), dst_2.clone()).await.is_ok()); + std::fs::File::create(&src) + .unwrap() + .write_all(b"hello") + .unwrap(); - let mut content = String::new(); + fs::symlink(&src, &dst).await.unwrap(); - { - let file = std::fs::File::open(dst.clone()).unwrap(); - let mut reader = BufReader::new(file); - reader.read_to_string(&mut content).unwrap(); - } + std::fs::File::create(&src) + .unwrap() + .write_all(b"new-data") + .unwrap(); - assert!(content == "hello"); + 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); diff --git a/tests/fs_open_options.rs b/tests/fs_open_options.rs new file mode 100644 index 0000000..41cfb45 --- /dev/null +++ b/tests/fs_open_options.rs @@ -0,0 +1,80 @@ +#![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() { + // TESTING HACK: use Debug output to check the stored data + assert!(format!("{:?}", OpenOptions::new().mode(0o644)).contains("mode: 420 ")); +} + +#[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,") + ); +} diff --git a/tests/fs_open_options_windows.rs b/tests/fs_open_options_windows.rs new file mode 100644 index 0000000..f09e8b4 --- /dev/null +++ b/tests/fs_open_options_windows.rs @@ -0,0 +1,51 @@ +#![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,")); +} diff --git a/tests/fs_remove_dir_all.rs b/tests/fs_remove_dir_all.rs new file mode 100644 index 0000000..5c71bdf --- /dev/null +++ b/tests/fs_remove_dir_all.rs @@ -0,0 +1,31 @@ +#![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"), + }; +} diff --git a/tests/fs_remove_file.rs b/tests/fs_remove_file.rs new file mode 100644 index 0000000..ea47721 --- /dev/null +++ b/tests/fs_remove_file.rs @@ -0,0 +1,24 @@ +#![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"), + }; +} diff --git a/tests/fs_rename.rs b/tests/fs_rename.rs new file mode 100644 index 0000000..91bb39e --- /dev/null +++ b/tests/fs_rename.rs @@ -0,0 +1,28 @@ +#![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"), + }; +} diff --git a/tests/fs_symlink_dir_windows.rs b/tests/fs_symlink_dir_windows.rs new file mode 100644 index 0000000..bcfba21 --- /dev/null +++ b/tests/fs_symlink_dir_windows.rs @@ -0,0 +1,31 @@ +#![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); +} diff --git a/tests/fs_symlink_file_windows.rs b/tests/fs_symlink_file_windows.rs new file mode 100644 index 0000000..4bc032e --- /dev/null +++ b/tests/fs_symlink_file_windows.rs @@ -0,0 +1,24 @@ +#![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); +} diff --git a/tests/fs_try_exists.rs b/tests/fs_try_exists.rs new file mode 100644 index 0000000..dca554d --- /dev/null +++ b/tests/fs_try_exists.rs @@ -0,0 +1,44 @@ +#![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 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 + ); + } +} diff --git a/tests/io_async_fd.rs b/tests/io_async_fd.rs index 5a6875e..7abd592 100644 --- a/tests/io_async_fd.rs +++ b/tests/io_async_fd.rs @@ -599,3 +599,214 @@ fn driver_shutdown_wakes_poll_race() { assert_err!(futures::executor::block_on(poll_writable(&afd_a))); } } + +#[tokio::test] +#[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<S: AsRawFd>(stream: &S, data: &[u8]) -> io::Result<usize> { + 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] +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] +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(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<libc::c_int> { + // 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(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::<libc::sockaddr_in>::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 { + Err(std::io::Error::last_os_error()).unwrap() + } + }); + + let fd = AsyncFd::new(socket).unwrap(); + + let guard = fd.ready(Interest::ERROR).await.unwrap(); + assert_eq!(guard.ready(), Ready::ERROR); +} diff --git a/tests/io_copy_bidirectional.rs b/tests/io_copy_bidirectional.rs index c549675..10eba31 100644 --- a/tests/io_copy_bidirectional.rs +++ b/tests/io_copy_bidirectional.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support bind() +#![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}; diff --git a/tests/io_driver.rs b/tests/io_driver.rs index 97018e0..2e2c847 100644 --- a/tests/io_driver.rs +++ b/tests/io_driver.rs @@ -1,6 +1,6 @@ #![warn(rust_2018_idioms)] // Wasi does not support panic recovery or threading -#![cfg(all(feature = "full", not(tokio_wasi)))] +#![cfg(all(feature = "full", not(target_os = "wasi")))] use tokio::net::TcpListener; use tokio::runtime; diff --git a/tests/io_driver_drop.rs b/tests/io_driver_drop.rs index 0a384d4..c318263 100644 --- a/tests/io_driver_drop.rs +++ b/tests/io_driver_drop.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support bind use tokio::net::TcpListener; use tokio::runtime; diff --git a/tests/io_fill_buf.rs b/tests/io_fill_buf.rs index 62c3f1b..534417c 100644 --- a/tests/io_fill_buf.rs +++ b/tests/io_fill_buf.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support file operations +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support file operations use tempfile::NamedTempFile; use tokio::fs::File; diff --git a/tests/io_mem_stream.rs b/tests/io_mem_stream.rs index 3d9941d..e79ec4e 100644 --- a/tests/io_mem_stream.rs +++ b/tests/io_mem_stream.rs @@ -63,7 +63,6 @@ async fn disconnect() { } #[tokio::test] -#[cfg(not(target_os = "android"))] async fn disconnect_reader() { let (a, mut b) = duplex(2); diff --git a/tests/io_panic.rs b/tests/io_panic.rs index 922d7c0..9b2bf6a 100644 --- a/tests/io_panic.rs +++ b/tests/io_panic.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use std::task::{Context, Poll}; use std::{error::Error, pin::Pin}; diff --git a/tests/io_read.rs b/tests/io_read.rs index 7f74c5e..6bea0ac 100644 --- a/tests/io_read.rs +++ b/tests/io_read.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![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; diff --git a/tests/io_read_buf.rs b/tests/io_read_buf.rs index 0328168..49a4f86 100644 --- a/tests/io_read_buf.rs +++ b/tests/io_read_buf.rs @@ -34,3 +34,61 @@ async fn read_buf() { 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); +} diff --git a/tests/io_read_to_end.rs b/tests/io_read_to_end.rs index 171e6d6..29b60be 100644 --- a/tests/io_read_to_end.rs +++ b/tests/io_read_to_end.rs @@ -5,6 +5,7 @@ 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() { @@ -76,3 +77,48 @@ async fn read_to_end_uninit() { test.read_to_end(&mut buf).await.unwrap(); assert_eq!(buf.len(), 33); } + +#[tokio::test] +async fn read_to_end_doesnt_grow_with_capacity() { + let arr: Vec<u8> = (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); +} diff --git a/tests/io_split.rs b/tests/io_split.rs index 4cbd2a2..9f17c9e 100644 --- a/tests/io_split.rs +++ b/tests/io_split.rs @@ -1,7 +1,9 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery -use tokio::io::{split, AsyncRead, AsyncWrite, ReadBuf, ReadHalf, WriteHalf}; +use tokio::io::{ + split, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf, ReadHalf, WriteHalf, +}; use std::io; use std::pin::Pin; @@ -36,6 +38,18 @@ impl AsyncWrite for RW { fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { Poll::Ready(Ok(())) } + + fn poll_write_vectored( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _bufs: &[io::IoSlice<'_>], + ) -> Poll<Result<usize, io::Error>> { + Poll::Ready(Ok(2)) + } + + fn is_write_vectored(&self) -> bool { + true + } } #[test] @@ -77,3 +91,26 @@ fn unsplit_err2() { 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()); + }); +} diff --git a/tests/io_take.rs b/tests/io_take.rs index fba01d2..539f17f 100644 --- a/tests/io_take.rs +++ b/tests/io_take.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use std::pin::Pin; use std::task::{Context, Poll}; diff --git a/tests/join_handle_panic.rs b/tests/join_handle_panic.rs index bc34fd0..c25a5f8 100644 --- a/tests/join_handle_panic.rs +++ b/tests/join_handle_panic.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery struct PanicsOnDrop; diff --git a/tests/macros_join.rs b/tests/macros_join.rs index 9e4d234..37dd05f 100644 --- a/tests/macros_join.rs +++ b/tests/macros_join.rs @@ -2,13 +2,13 @@ #![allow(clippy::disallowed_names)] use std::sync::Arc; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] #[cfg(target_pointer_width = "64")] use wasm_bindgen_test::wasm_bindgen_test as test; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::{oneshot, Semaphore}; @@ -66,7 +66,6 @@ async fn two_await() { #[test] #[cfg(target_pointer_width = "64")] -#[ignore = "Android: ignore until the compiler is updated. aliceryhl@ says these tests assume latest stable compiler."] fn join_size() { use futures::future; use std::mem; @@ -154,3 +153,9 @@ async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { *poll_order.lock().unwrap() ); } + +#[tokio::test] +#[allow(clippy::unit_cmp)] +async fn empty_join() { + assert_eq!(tokio::join!(), ()); +} diff --git a/tests/macros_pin.rs b/tests/macros_pin.rs index e99a3ee..2de68ad 100644 --- a/tests/macros_pin.rs +++ b/tests/macros_pin.rs @@ -1,9 +1,9 @@ #![cfg(feature = "macros")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; async fn one() {} diff --git a/tests/macros_rename_test.rs b/tests/macros_rename_test.rs index 90a86f9..0c5eba2 100644 --- a/tests/macros_rename_test.rs +++ b/tests/macros_rename_test.rs @@ -1,10 +1,14 @@ -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threading +#![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() @@ -24,3 +28,8 @@ fn crate_rename_main() { 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); +} diff --git a/tests/macros_select.rs b/tests/macros_select.rs index 26d6fec..68a607b 100644 --- a/tests/macros_select.rs +++ b/tests/macros_select.rs @@ -1,10 +1,10 @@ #![cfg(feature = "macros")] #![allow(clippy::disallowed_names)] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::oneshot; @@ -206,47 +206,56 @@ async fn nested() { assert_eq!(res, 3); } -#[maybe_tokio_test] #[cfg(target_pointer_width = "64")] -async fn struct_size() { +mod pointer_64_tests { + use super::maybe_tokio_test; use futures::future; use std::mem; - let fut = async { - let ready = future::ready(0i32); + #[maybe_tokio_test] + async fn struct_size_1() { + let fut = async { + let ready = future::ready(0i32); - tokio::select! { - _ = ready => {}, - } - }; + tokio::select! { + _ = ready => {}, + } + }; - assert_eq!(mem::size_of_val(&fut), 40); + assert_eq!(mem::size_of_val(&fut), 32); + } - let fut = async { - let ready1 = future::ready(0i32); - let ready2 = future::ready(0i32); + #[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 => {}, - } - }; + tokio::select! { + _ = ready1 => {}, + _ = ready2 => {}, + } + }; - assert_eq!(mem::size_of_val(&fut), 48); + assert_eq!(mem::size_of_val(&fut), 40); + } - let fut = async { - let ready1 = future::ready(0i32); - let ready2 = future::ready(0i32); - let ready3 = future::ready(0i32); + #[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 => {}, - } - }; + tokio::select! { + _ = ready1 => {}, + _ = ready2 => {}, + _ = ready3 => {}, + } + }; - assert_eq!(mem::size_of_val(&fut), 56); + assert_eq!(mem::size_of_val(&fut), 48); + } } #[maybe_tokio_test] @@ -624,7 +633,7 @@ mod unstable { } #[test] - #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + #[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() diff --git a/tests/macros_test.rs b/tests/macros_test.rs index 85279b7..b5095e3 100644 --- a/tests/macros_test.rs +++ b/tests/macros_test.rs @@ -1,4 +1,4 @@ -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threading +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threading use tokio::test; diff --git a/tests/macros_try_join.rs b/tests/macros_try_join.rs index 6c43222..c8ed00b 100644 --- a/tests/macros_try_join.rs +++ b/tests/macros_try_join.rs @@ -6,10 +6,10 @@ use std::sync::Arc; use tokio::sync::{oneshot, Semaphore}; use tokio_test::{assert_pending, assert_ready, task}; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; #[maybe_tokio_test] @@ -183,3 +183,8 @@ async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() { *poll_order.lock().unwrap() ); } + +#[tokio::test] +async fn empty_try_join() { + assert_eq!(tokio::try_join!() as Result<_, ()>, Ok(())); +} diff --git a/tests/net_bind_resource.rs b/tests/net_bind_resource.rs index 1c604aa..05fd160 100644 --- a/tests/net_bind_resource.rs +++ b/tests/net_bind_resource.rs @@ -1,9 +1,8 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support panic recovery or bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery or bind use tokio::net::TcpListener; -use std::convert::TryFrom; use std::net; #[test] diff --git a/tests/net_lookup_host.rs b/tests/net_lookup_host.rs index 2b346e6..667030c 100644 --- a/tests/net_lookup_host.rs +++ b/tests/net_lookup_host.rs @@ -1,4 +1,4 @@ -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support direct socket operations +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support direct socket operations use tokio::net; use tokio_test::assert_ok; diff --git a/tests/net_named_pipe.rs b/tests/net_named_pipe.rs index c421224..02eda48 100644 --- a/tests/net_named_pipe.rs +++ b/tests/net_named_pipe.rs @@ -2,13 +2,11 @@ #![cfg(all(windows))] use std::io; -use std::mem; -use std::os::windows::io::AsRawHandle; use std::time::Duration; -use tokio::io::AsyncWriteExt; +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, NO_ERROR, UNICODE_STRING}; +use windows_sys::Win32::Foundation::{ERROR_NO_DATA, ERROR_PIPE_BUSY}; #[tokio::test] async fn test_named_pipe_client_drop() -> io::Result<()> { @@ -16,8 +14,6 @@ async fn test_named_pipe_client_drop() -> io::Result<()> { let mut server = ServerOptions::new().create(PIPE_NAME)?; - assert_eq!(num_instances("test-named-pipe-client-drop")?, 1); - let client = ClientOptions::new().open(PIPE_NAME)?; server.connect().await?; @@ -327,17 +323,51 @@ async fn test_named_pipe_multi_client_ready() -> io::Result<()> { Ok(()) } -// This tests what happens when a client tries to disconnect. +// This tests that message mode works as expected. #[tokio::test] async fn test_named_pipe_mode_message() -> io::Result<()> { - const PIPE_NAME: &str = r"\\.\pipe\test-named-pipe-mode-message"; + // 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 server = ServerOptions::new() - .pipe_mode(PipeMode::Message) - .create(PIPE_NAME)?; + let mut client = ClientOptions::new().pipe_mode(mode).open(&pipe_name)?; - let _ = ClientOptions::new().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(()) } @@ -367,53 +397,3 @@ async fn test_named_pipe_access() -> io::Result<()> { } Ok(()) } - -fn num_instances(pipe_name: impl AsRef<str>) -> io::Result<u32> { - use ntapi::ntioapi; - - let mut name = pipe_name.as_ref().encode_utf16().collect::<Vec<_>>(); - let mut name = UNICODE_STRING { - Length: (name.len() * mem::size_of::<u16>()) as u16, - MaximumLength: (name.len() * mem::size_of::<u16>()) as u16, - Buffer: name.as_mut_ptr(), - }; - let root = std::fs::File::open(r"\\.\Pipe\")?; - let mut io_status_block = unsafe { mem::zeroed() }; - let mut file_directory_information = [0_u8; 1024]; - - let status = unsafe { - ntioapi::NtQueryDirectoryFile( - root.as_raw_handle(), - std::ptr::null_mut(), - None, - std::ptr::null_mut(), - &mut io_status_block, - &mut file_directory_information as *mut _ as *mut _, - 1024, - ntioapi::FileDirectoryInformation, - 0, - &mut name as *mut _ as _, - 0, - ) - }; - - if status as u32 != NO_ERROR { - return Err(io::Error::last_os_error()); - } - - let info = unsafe { - mem::transmute::<_, &ntioapi::FILE_DIRECTORY_INFORMATION>(&file_directory_information) - }; - let raw_name = unsafe { - std::slice::from_raw_parts( - info.FileName.as_ptr(), - info.FileNameLength as usize / mem::size_of::<u16>(), - ) - }; - let name = String::from_utf16(raw_name).unwrap(); - let num_instances = unsafe { *info.EndOfFile.QuadPart() }; - - assert_eq!(name, pipe_name.as_ref()); - - Ok(num_instances as u32) -} diff --git a/tests/net_panic.rs b/tests/net_panic.rs index fc2209a..7980761 100644 --- a/tests/net_panic.rs +++ b/tests/net_panic.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] +#![cfg(all(feature = "full", not(target_os = "wasi")))] use std::error::Error; use tokio::net::{TcpListener, TcpStream}; diff --git a/tests/net_unix_pipe.rs b/tests/net_unix_pipe.rs new file mode 100644 index 0000000..c96d6e7 --- /dev/null +++ b/tests/net_unix_pipe.rs @@ -0,0 +1,429 @@ +#![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<TempFifo> { + 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<Path> for TempFifo { + fn as_ref(&self) -> &Path { + self.path.as_ref() + } +} + +#[tokio::test] +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")] +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<Path>, 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] +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")] +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] +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] +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] +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] +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<T: AsRawFd>(fd: &T) -> io::Result<bool> { + let flags = nix::fcntl::fcntl(fd.as_raw_fd(), nix::fcntl::F_GETFL)?; + Ok((flags & libc::O_NONBLOCK) != 0) +} + +#[tokio::test] +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] +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] +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] +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(()) +} diff --git a/tests/no_rt.rs b/tests/no_rt.rs index 747fab6..89c7ce0 100644 --- a/tests/no_rt.rs +++ b/tests/no_rt.rs @@ -1,4 +1,4 @@ -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use tokio::net::TcpStream; use tokio::sync::oneshot; diff --git a/tests/process_kill_on_drop.rs b/tests/process_kill_on_drop.rs index 658e4ad..d919b1a 100644 --- a/tests/process_kill_on_drop.rs +++ b/tests/process_kill_on_drop.rs @@ -12,7 +12,7 @@ use tokio_test::assert_ok; #[tokio::test] async fn kill_on_drop() { let mut cmd = Command::new("bash"); - cmd.args(&[ + cmd.args([ "-c", " # Fork another child that won't get killed diff --git a/tests/process_smoke.rs b/tests/process_smoke.rs index 81d2020..635b951 100644 --- a/tests/process_smoke.rs +++ b/tests/process_smoke.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi cannot run system commands +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi cannot run system commands use tokio::process::Command; use tokio_test::assert_ok; diff --git a/tests/rt_basic.rs b/tests/rt_basic.rs index 6caf0a4..789c047 100644 --- a/tests/rt_basic.rs +++ b/tests/rt_basic.rs @@ -178,7 +178,7 @@ fn drop_tasks_in_context() { } #[test] -#[cfg_attr(tokio_wasi, ignore = "Wasi does not support panic recovery")] +#[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] #[should_panic(expected = "boom")] fn wake_in_drop_after_panic() { let (tx, rx) = oneshot::channel::<()>(); @@ -239,7 +239,7 @@ fn spawn_two() { } } -#[cfg_attr(tokio_wasi, ignore = "WASI: std::thread::spawn not supported")] +#[cfg_attr(target_os = "wasi", ignore = "WASI: std::thread::spawn not supported")] #[test] fn spawn_remote() { let rt = rt(); @@ -276,7 +276,7 @@ fn spawn_remote() { } #[test] -#[cfg_attr(tokio_wasi, ignore = "Wasi does not support panic recovery")] +#[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." )] @@ -315,7 +315,7 @@ mod unstable { } #[test] - #[cfg_attr(tokio_wasi, ignore = "Wasi does not support panic recovery")] + #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] fn spawns_do_nothing() { use std::sync::Arc; @@ -344,7 +344,7 @@ mod unstable { } #[test] - #[cfg_attr(tokio_wasi, ignore = "Wasi does not support panic recovery")] + #[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}; diff --git a/tests/rt_common.rs b/tests/rt_common.rs index 3892998..abca8dd 100644 --- a/tests/rt_common.rs +++ b/tests/rt_common.rs @@ -21,7 +21,7 @@ macro_rules! rt_test { } } - #[cfg(not(tokio_wasi))] // Wasi doesn't support threads + #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads mod threaded_scheduler_4_threads { $($t)* @@ -37,7 +37,7 @@ macro_rules! rt_test { } } - #[cfg(not(tokio_wasi))] // Wasi doesn't support threads + #[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads mod threaded_scheduler_1_thread { $($t)* @@ -52,6 +52,40 @@ macro_rules! rt_test { .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<Runtime> { + 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<Runtime> { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap() + .into() + } + } } } @@ -695,12 +729,34 @@ rt_test! { /// Tests that yielded tasks are not scheduled until **after** resource /// drivers are polled. /// - /// Note: we may have to delete this test as it is not necessarily reliable. /// The OS does not guarantee when I/O events are delivered, so there may be - /// more yields than anticipated. + /// 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"))] 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; @@ -727,14 +783,16 @@ rt_test! { barrier.wait(); - tokio::spawn(async move { + let (fail_test, fail_test_recv) = oneshot::channel::<()>(); + + let jh = tokio::spawn(async move { // Create a TCP litener let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::join!( async { - // Done blocking intentionally + // Done in a blocking manner intentionally. let _socket = std::net::TcpStream::connect(addr).unwrap(); // Yield until connected @@ -744,7 +802,12 @@ rt_test! { cnt += 1; if cnt >= 10 { - panic!("yielded too many times; TODO: delete this test?"); + // 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; } } }, @@ -753,8 +816,20 @@ rt_test! { flag.store(true, SeqCst); } ); - }).await.unwrap(); - }); + }); + + // 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 { + // Check for panics in spawned task. + jh.abort(); + jh.await.unwrap(); + } + + success + }) } #[cfg(not(target_os="wasi"))] // Wasi does not support threads @@ -769,9 +844,9 @@ rt_test! { assert_err!(rx.try_recv()); } - #[cfg_attr(tokio_wasi, ignore = "Wasi does not support threads or panic recovery")] + #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support threads or panic recovery")] + #[cfg(panic = "unwind")] #[test] - #[cfg(not(target_os = "android"))] fn panic_in_task() { let rt = rt(); let (tx, rx) = oneshot::channel(); @@ -799,7 +874,7 @@ rt_test! { #[test] #[should_panic] - #[cfg_attr(tokio_wasi, ignore = "Wasi does not support panic recovery")] + #[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!() }); @@ -876,10 +951,6 @@ rt_test! { #[test] fn wake_while_rt_is_dropping() { use tokio::sync::Barrier; - use core::sync::atomic::{AtomicBool, Ordering}; - - let drop_triggered = Arc::new(AtomicBool::new(false)); - let set_drop_triggered = drop_triggered.clone(); struct OnDrop<F: FnMut()>(F); @@ -891,56 +962,39 @@ rt_test! { let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); - let (tx3, rx3) = oneshot::channel(); - let barrier = Arc::new(Barrier::new(4)); + let barrier = Arc::new(Barrier::new(3)); let barrier1 = barrier.clone(); let barrier2 = barrier.clone(); - let barrier3 = 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()); - tx3.send(()).unwrap(); }); rt.spawn(async move { - let h1 = tokio::runtime::Handle::current(); - // When this task is dropped, we'll be "closing remotes". - // We spawn a new task that owns the `tx1`, to move its Drop - // out of here. - // - // Importantly, the oneshot 1 has a waker already stored, so - // the eventual drop here will try to re-schedule again. - let mut opt_tx1 = Some(tx1); + let mut tx1 = Some(tx1); let _d = OnDrop(move || { - let tx1 = opt_tx1.take().unwrap(); - h1.spawn(async move { - tx1.send(()).unwrap(); - }); - // Just a sanity check that this entire thing actually happened - set_drop_triggered.store(true, Ordering::Relaxed); + let _ = tx1.take().unwrap().send(()); }); - let _ = tokio::join!(rx2, barrier2.wait()); - }); - rt.spawn(async move { - let _ = tokio::join!(rx3, barrier3.wait()); - // We'll never get here, but once task 3 drops, this will - // force task 2 to re-schedule since it's waiting on oneshot 2. - tx2.send(()).unwrap(); + // 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 + // Drop the rt. Regardless of which task is dropped first, its destructor will wake the + // other task. drop(rt); - - // Make sure that the spawn actually happened - assert!(drop_triggered.load(Ordering::Relaxed)); } #[cfg(not(target_os="wasi"))] // Wasi doesn't support UDP or bind() @@ -1029,7 +1083,7 @@ rt_test! { // See https://github.com/rust-lang/rust/issues/74875 #[test] #[cfg(not(windows))] - #[cfg_attr(tokio_wasi, ignore = "Wasi does not support threads")] + #[cfg_attr(target_os = "wasi", ignore = "Wasi does not support threads")] fn runtime_in_thread_local() { use std::cell::RefCell; use std::thread; @@ -1074,7 +1128,7 @@ rt_test! { tx.send(()).unwrap(); } - #[cfg(not(tokio_wasi))] // Wasi does not support bind + #[cfg(not(target_os = "wasi"))] // Wasi does not support bind #[test] fn local_set_block_on_socket() { let rt = rt(); @@ -1096,7 +1150,7 @@ rt_test! { }); } - #[cfg(not(tokio_wasi))] // Wasi does not support bind + #[cfg(not(target_os = "wasi"))] // Wasi does not support bind #[test] fn local_set_client_server_block_on() { let rt = rt(); @@ -1110,7 +1164,7 @@ rt_test! { assert_err!(rx.try_recv()); } - #[cfg(not(tokio_wasi))] // Wasi does not support bind + #[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); @@ -1277,4 +1331,39 @@ rt_test! { } }); } + + #[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(); + } + } } diff --git a/tests/rt_handle.rs b/tests/rt_handle.rs new file mode 100644 index 0000000..5d4b53f --- /dev/null +++ b/tests/rt_handle.rs @@ -0,0 +1,90 @@ +#![warn(rust_2018_idioms)] +#![cfg(feature = "full")] + +use tokio::runtime::Runtime; + +#[test] +fn basic_enter() { + let rt1 = rt(); + let rt2 = rt(); + + let enter1 = rt1.enter(); + let enter2 = rt2.enter(); + + drop(enter2); + drop(enter1); +} + +#[test] +#[should_panic] +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] +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"))] +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(); +} + +#[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() +} diff --git a/tests/rt_handle_block_on.rs b/tests/rt_handle_block_on.rs index 5ec783e..54b9002 100644 --- a/tests/rt_handle_block_on.rs +++ b/tests/rt_handle_block_on.rs @@ -10,10 +10,10 @@ use std::time::Duration; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use tokio::{net, time}; -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads macro_rules! multi_threaded_rt_test { ($($t:tt)*) => { mod threaded_scheduler_4_threads_only { @@ -46,7 +46,7 @@ macro_rules! multi_threaded_rt_test { } } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] macro_rules! rt_test { ($($t:tt)*) => { mod current_thread_scheduler { @@ -126,7 +126,7 @@ fn unbounded_mpsc_channel() { }) } -#[cfg(not(tokio_wasi))] // Wasi doesn't support file operations or bind +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support file operations or bind rt_test! { use tokio::fs; // ==== spawn blocking futures ====== @@ -419,7 +419,7 @@ rt_test! { } } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] multi_threaded_rt_test! { #[cfg(unix)] #[test] @@ -482,7 +482,7 @@ multi_threaded_rt_test! { // ==== utils ====== /// Create a new multi threaded runtime -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] fn new_multi_thread(n: usize) -> Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(n) @@ -513,7 +513,7 @@ where f(); } - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] { let rt = new_multi_thread(1); let _enter = rt.enter(); @@ -523,7 +523,7 @@ where f(); } - #[cfg(not(tokio_wasi))] + #[cfg(not(target_os = "wasi"))] { let rt = new_multi_thread(4); let _enter = rt.enter(); diff --git a/tests/rt_metrics.rs b/tests/rt_metrics.rs index fdb2fb5..42faab3 100644 --- a/tests/rt_metrics.rs +++ b/tests/rt_metrics.rs @@ -1,9 +1,13 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", tokio_unstable, not(tokio_wasi)))] +#![cfg(all(feature = "full", tokio_unstable, not(target_os = "wasi")))] +use std::future::Future; use std::sync::{Arc, Mutex}; +use std::task::Poll; +use tokio::macros::support::poll_fn; use tokio::runtime::Runtime; +use tokio::task::consume_budget; use tokio::time::{self, Duration}; #[test] @@ -78,6 +82,23 @@ fn blocking_queue_depth() { } #[test] +fn active_tasks_count() { + let rt = current_thread(); + let metrics = rt.metrics(); + assert_eq!(0, metrics.active_tasks_count()); + rt.spawn(async move { + assert_eq!(1, metrics.active_tasks_count()); + }); + + let rt = threaded(); + let metrics = rt.metrics(); + assert_eq!(0, metrics.active_tasks_count()); + rt.spawn(async move { + assert_eq!(1, metrics.active_tasks_count()); + }); +} + +#[test] fn remote_schedule_count() { use std::thread; @@ -194,24 +215,37 @@ fn worker_steal_count() { } #[test] -fn worker_poll_count() { +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(async {}).await.unwrap(); + 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_count_histogram_enabled()); + for i in 0..10 { + assert_eq!(0, metrics.poll_count_histogram_bucket_count(0, i)); + } let rt = threaded(); let metrics = rt.metrics(); rt.block_on(async { for _ in 0..N { - tokio::spawn(async {}).await.unwrap(); + tokio::spawn(task()).await.unwrap(); } }); drop(rt); @@ -221,6 +255,132 @@ fn worker_poll_count() { .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_count_histogram_enabled()); + for n in 0..metrics.num_workers() { + for i in 0..10 { + assert_eq!(0, metrics.poll_count_histogram_bucket_count(n, i)); + } + } +} + +#[test] +fn worker_poll_count_histogram() { + const N: u64 = 5; + + let rts = [ + tokio::runtime::Builder::new_current_thread() + .enable_all() + .enable_metrics_poll_count_histogram() + .metrics_poll_count_histogram_scale(tokio::runtime::HistogramScale::Linear) + .metrics_poll_count_histogram_buckets(3) + .metrics_poll_count_histogram_resolution(Duration::from_millis(50)) + .build() + .unwrap(), + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .enable_metrics_poll_count_histogram() + .metrics_poll_count_histogram_scale(tokio::runtime::HistogramScale::Linear) + .metrics_poll_count_histogram_buckets(3) + .metrics_poll_count_histogram_resolution(Duration::from_millis(50)) + .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_count_histogram_num_buckets(); + + assert!(metrics.poll_count_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_count_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_count_histogram() + .metrics_poll_count_histogram_scale(tokio::runtime::HistogramScale::Linear) + .metrics_poll_count_histogram_buckets(3) + .metrics_poll_count_histogram_resolution(us(50)) + .build() + .unwrap(); + let metrics = rt.metrics(); + + assert_eq!(metrics.poll_count_histogram_bucket_range(0), us(0)..us(50)); + assert_eq!( + metrics.poll_count_histogram_bucket_range(1), + us(50)..us(100) + ); + assert_eq!(metrics.poll_count_histogram_bucket_range(2), us(100)..max); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .enable_metrics_poll_count_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_count_histogram_bucket_range(0), us(0)..a); + assert_eq!(metrics.poll_count_histogram_bucket_range(1), a..b); + assert_eq!(metrics.poll_count_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_count_histogram_scale(tokio::runtime::HistogramScale::Linear) + .metrics_poll_count_histogram_buckets(3) + .metrics_poll_count_histogram_resolution(Duration::from_millis(50)) + .build() + .unwrap(), + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .metrics_poll_count_histogram_scale(tokio::runtime::HistogramScale::Linear) + .metrics_poll_count_histogram_buckets(3) + .metrics_poll_count_histogram_resolution(Duration::from_millis(50)) + .build() + .unwrap(), + ]; + + for rt in rts { + let metrics = rt.metrics(); + assert!(!metrics.poll_count_histogram_enabled()); + } } #[test] @@ -363,10 +523,20 @@ fn injection_queue_depth() { // First we need to block the runtime workers let (tx1, rx1) = std::sync::mpsc::channel(); let (tx2, rx2) = std::sync::mpsc::channel(); + let (tx3, rx3) = std::sync::mpsc::channel(); + let rx3 = Arc::new(Mutex::new(rx3)); rt.spawn(async move { rx1.recv().unwrap() }); rt.spawn(async move { rx2.recv().unwrap() }); + // Spawn some more to make sure there are items + for _ in 0..10 { + let rx = rx3.clone(); + rt.spawn(async move { + rx.lock().unwrap().recv().unwrap(); + }); + } + thread::spawn(move || { handle.spawn(async {}); }) @@ -375,7 +545,11 @@ fn injection_queue_depth() { let n = metrics.injection_queue_depth(); assert!(1 <= n, "{}", n); - assert!(3 >= n, "{}", n); + assert!(15 >= n, "{}", n); + + for _ in 0..10 { + tx3.send(()).unwrap(); + } tx1.send(()).unwrap(); tx2.send(()).unwrap(); @@ -433,6 +607,78 @@ fn worker_local_queue_depth() { }); } +#[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() { @@ -479,3 +725,7 @@ fn threaded() -> Runtime { .build() .unwrap() } + +fn us(n: u64) -> Duration { + Duration::from_micros(n) +} diff --git a/tests/rt_panic.rs b/tests/rt_panic.rs index f9a684f..57f3274 100644 --- a/tests/rt_panic.rs +++ b/tests/rt_panic.rs @@ -1,6 +1,6 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "full")] -#![cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery +#![cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery use futures::future; use std::error::Error; diff --git a/tests/rt_threaded.rs b/tests/rt_threaded.rs index c598418..83df6ed 100644 --- a/tests/rt_threaded.rs +++ b/tests/rt_threaded.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] +#![cfg(all(feature = "full", not(target_os = "wasi")))] use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; @@ -30,7 +30,8 @@ fn single_thread() { let _ = runtime::Builder::new_multi_thread() .enable_all() .worker_threads(1) - .build(); + .build() + .unwrap(); } #[test] @@ -161,6 +162,32 @@ fn many_multishot_futures() { } #[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] fn spawn_shutdown() { let rt = rt(); let (tx, rx) = mpsc::channel(); @@ -561,6 +588,154 @@ 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<AtomicBool>, counter: Arc<AtomicUsize>, 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() } @@ -587,4 +762,22 @@ mod unstable { .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()); + } } diff --git a/tests/rt_threaded_alt.rs b/tests/rt_threaded_alt.rs new file mode 100644 index 0000000..3e3ac07 --- /dev/null +++ b/tests/rt_threaded_alt.rs @@ -0,0 +1,717 @@ +#![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 futures::future::poll_fn; +use std::future::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(tokio_unstable)] + { + $( $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] +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<AtomicUsize>); + + 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(); +} + +#[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<Waker>, + } + + struct MyFuture { + shared: Arc<Mutex<Shared>>, + 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<AtomicBool>, counter: Arc<AtomicUsize>, 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() +} diff --git a/tests/rt_time_start_paused.rs b/tests/rt_time_start_paused.rs new file mode 100644 index 0000000..283f474 --- /dev/null +++ b/tests/rt_time_start_paused.rs @@ -0,0 +1,14 @@ +#![cfg(all(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()); + } +} diff --git a/tests/signal_no_rt.rs b/tests/signal_no_rt.rs index ffcc665..db284e1 100644 --- a/tests/signal_no_rt.rs +++ b/tests/signal_no_rt.rs @@ -4,7 +4,7 @@ use tokio::signal::unix::{signal, SignalKind}; -#[cfg_attr(tokio_wasi, ignore = "Wasi does not support panic recovery")] +#[cfg_attr(target_os = "wasi", ignore = "Wasi does not support panic recovery")] #[test] #[should_panic] fn no_runtime_panics_creating_signals() { diff --git a/tests/sync_barrier.rs b/tests/sync_barrier.rs index 2001c35..ac5977f 100644 --- a/tests/sync_barrier.rs +++ b/tests/sync_barrier.rs @@ -2,7 +2,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::Barrier; diff --git a/tests/sync_broadcast.rs b/tests/sync_broadcast.rs index cd66924..16b9a0a 100644 --- a/tests/sync_broadcast.rs +++ b/tests/sync_broadcast.rs @@ -2,7 +2,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::broadcast; @@ -276,14 +276,14 @@ fn send_no_rx() { #[test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn zero_capacity() { broadcast::channel::<()>(0); } #[test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn capacity_too_big() { use std::usize; @@ -292,7 +292,7 @@ fn capacity_too_big() { #[test] #[cfg(panic = "unwind")] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn panic_in_clone() { use std::panic::{self, AssertUnwindSafe}; @@ -563,7 +563,7 @@ fn sender_len() { } #[test] -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] fn sender_len_random() { use rand::Rng; @@ -587,3 +587,57 @@ fn sender_len_random() { 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<Self>) {} + } + + // 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(()); +} diff --git a/tests/sync_errors.rs b/tests/sync_errors.rs index bf0a6cf..4e43c8f 100644 --- a/tests/sync_errors.rs +++ b/tests/sync_errors.rs @@ -1,7 +1,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; fn is_error<T: std::error::Error + Send + Sync>() {} diff --git a/tests/sync_mpsc.rs b/tests/sync_mpsc.rs index 6e87096..d2b7078 100644 --- a/tests/sync_mpsc.rs +++ b/tests/sync_mpsc.rs @@ -2,20 +2,21 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; -#[cfg(tokio_wasm_not_wasi)] +#[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::sync::Arc; use tokio::sync::mpsc; use tokio::sync::mpsc::error::{TryRecvError, TrySendError}; -#[cfg(not(tokio_wasm_not_wasi))] -use tokio::test as maybe_tokio_test; use tokio_test::*; -#[cfg(not(tokio_wasm))] +#[cfg(not(target_family = "wasm"))] mod support { pub(crate) mod mpsc_stream; } @@ -87,7 +88,7 @@ async fn reserve_disarm() { } #[tokio::test] -#[cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads async fn send_recv_stream_with_buffer() { use tokio_stream::StreamExt; @@ -154,7 +155,7 @@ async fn start_send_past_cap() { #[test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn buffer_gteq_one() { mpsc::channel::<i32>(0); } @@ -191,7 +192,7 @@ async fn async_send_recv_unbounded() { } #[tokio::test] -#[cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads async fn send_recv_stream_unbounded() { use tokio_stream::StreamExt; @@ -452,7 +453,7 @@ fn unconsumed_messages_are_dropped() { } #[test] -#[cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads fn blocking_recv() { let (tx, mut rx) = mpsc::channel::<u8>(1); @@ -470,14 +471,14 @@ fn blocking_recv() { #[tokio::test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[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(tokio_wasi)))] // Wasi doesn't support threads +#[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads fn blocking_send() { let (tx, mut rx) = mpsc::channel::<u8>(1); @@ -495,7 +496,7 @@ fn blocking_send() { #[tokio::test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[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(()); @@ -648,7 +649,7 @@ async fn recv_timeout() { #[test] #[should_panic = "there is no reactor running, must be called from the context of a Tokio 1.x runtime"] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn recv_timeout_panic() { use futures::future::FutureExt; use tokio::time::Duration; diff --git a/tests/sync_mpsc_weak.rs b/tests/sync_mpsc_weak.rs index 0fdfc00..fad4c72 100644 --- a/tests/sync_mpsc_weak.rs +++ b/tests/sync_mpsc_weak.rs @@ -2,7 +2,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use std::sync::atomic::AtomicUsize; diff --git a/tests/sync_mutex.rs b/tests/sync_mutex.rs index 1e35a55..f423c82 100644 --- a/tests/sync_mutex.rs +++ b/tests/sync_mutex.rs @@ -1,12 +1,12 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::Mutex; diff --git a/tests/sync_mutex_owned.rs b/tests/sync_mutex_owned.rs index ba472fe..badcef3 100644 --- a/tests/sync_mutex_owned.rs +++ b/tests/sync_mutex_owned.rs @@ -1,12 +1,12 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::Mutex; diff --git a/tests/sync_notify.rs b/tests/sync_notify.rs index aa58039..e31b9d4 100644 --- a/tests/sync_notify.rs +++ b/tests/sync_notify.rs @@ -1,7 +1,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::Notify; diff --git a/tests/sync_once_cell.rs b/tests/sync_once_cell.rs index 38dfa7c..d5a6947 100644 --- a/tests/sync_once_cell.rs +++ b/tests/sync_once_cell.rs @@ -4,7 +4,11 @@ use std::mem; use std::ops::Drop; 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() { @@ -102,184 +106,170 @@ fn from() { assert_eq!(*cell.get().unwrap(), 2); } -#[cfg(feature = "parking_lot")] -mod parking_lot { - use super::*; - - use tokio::runtime; - use tokio::sync::SetError; - use tokio::time; - - use std::time::Duration; - - async fn func1() -> u32 { - 5 - } - - async fn func2() -> u32 { - time::sleep(Duration::from_millis(1)).await; - 10 - } - - async fn func_err() -> Result<u32, ()> { - Err(()) - } +async fn func1() -> u32 { + 5 +} - async fn func_ok() -> Result<u32, ()> { - Ok(10) - } +async fn func2() -> u32 { + time::sleep(Duration::from_millis(1)).await; + 10 +} - async fn func_panic() -> u32 { - time::sleep(Duration::from_millis(1)).await; - panic!(); - } +async fn func_err() -> Result<u32, ()> { + Err(()) +} - 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 func_ok() -> Result<u32, ()> { + Ok(10) +} - async fn advance_time_and_set( - cell: &'static OnceCell<u32>, - v: u32, - ) -> Result<(), SetError<u32>> { - time::advance(Duration::from_millis(1)).await; - cell.set(v) - } +async fn func_panic() -> u32 { + time::sleep(Duration::from_millis(1)).await; + panic!(); +} - #[test] - fn get_or_init() { - let rt = runtime::Builder::new_current_thread() - .enable_time() - .start_paused(true) - .build() - .unwrap(); +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 +} - static ONCE: OnceCell<u32> = OnceCell::const_new(); +async fn advance_time_and_set(cell: &'static OnceCell<u32>, v: u32) -> Result<(), SetError<u32>> { + time::advance(Duration::from_millis(1)).await; + cell.set(v) +} - 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 }); +#[test] +fn get_or_init() { + let rt = runtime::Builder::new_current_thread() + .enable_time() + .start_paused(true) + .build() + .unwrap(); - time::advance(Duration::from_millis(1)).await; - time::resume(); + static ONCE: OnceCell<u32> = OnceCell::const_new(); - let result1 = handle1.await.unwrap(); - let result2 = handle2.await.unwrap(); + 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 }); - assert_eq!(*result1, 5); - assert_eq!(*result2, 5); - }); - } + time::advance(Duration::from_millis(1)).await; + time::resume(); - #[test] - fn get_or_init_panic() { - let rt = runtime::Builder::new_current_thread() - .enable_time() - .build() - .unwrap(); + let result1 = handle1.await.unwrap(); + let result2 = handle2.await.unwrap(); - static ONCE: OnceCell<u32> = OnceCell::const_new(); + assert_eq!(*result1, 5); + assert_eq!(*result2, 5); + }); +} - rt.block_on(async { - time::pause(); +#[test] +fn get_or_init_panic() { + let rt = runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); - let handle1 = rt.spawn(async { ONCE.get_or_init(func1).await }); - let handle2 = rt.spawn(async { ONCE.get_or_init(func_panic).await }); + static ONCE: OnceCell<u32> = OnceCell::const_new(); - time::advance(Duration::from_millis(1)).await; + rt.block_on(async { + time::pause(); - let result1 = handle1.await.unwrap(); - let result2 = handle2.await.unwrap(); + let handle1 = rt.spawn(async { ONCE.get_or_init(func1).await }); + let handle2 = rt.spawn(async { ONCE.get_or_init(func_panic).await }); - assert_eq!(*result1, 5); - assert_eq!(*result2, 5); - }); - } + time::advance(Duration::from_millis(1)).await; - #[test] - fn set_and_get() { - let rt = runtime::Builder::new_current_thread() - .enable_time() - .build() - .unwrap(); + let result1 = handle1.await.unwrap(); + let result2 = handle2.await.unwrap(); - static ONCE: OnceCell<u32> = OnceCell::const_new(); + assert_eq!(*result1, 5); + assert_eq!(*result2, 5); + }); +} - rt.block_on(async { - let _ = rt.spawn(async { ONCE.set(5) }).await; - let value = ONCE.get().unwrap(); - assert_eq!(*value, 5); - }); - } +#[test] +fn set_and_get() { + let rt = runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + + static ONCE: OnceCell<u32> = 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<u32> = OnceCell::const_new(); - let uninit = ONCE.get(); - assert!(uninit.is_none()); - } +#[test] +fn get_uninit() { + static ONCE: OnceCell<u32> = OnceCell::const_new(); + let uninit = ONCE.get(); + assert!(uninit.is_none()); +} - #[test] - fn set_twice() { - static ONCE: OnceCell<u32> = OnceCell::const_new(); +#[test] +fn set_twice() { + static ONCE: OnceCell<u32> = 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()); - } + 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(); +#[test] +fn set_while_initializing() { + let rt = runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); - static ONCE: OnceCell<u32> = OnceCell::const_new(); + static ONCE: OnceCell<u32> = OnceCell::const_new(); - rt.block_on(async { - time::pause(); + 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 }); + 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; + time::advance(Duration::from_millis(2)).await; - let result1 = handle1.await.unwrap(); - let result2 = handle2.await.unwrap(); + let result1 = handle1.await.unwrap(); + let result2 = handle2.await.unwrap(); - assert_eq!(*result1, 5); - assert!(result2.err().unwrap().is_initializing_err()); - }); - } + 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(); +#[test] +fn get_or_try_init() { + let rt = runtime::Builder::new_current_thread() + .enable_time() + .start_paused(true) + .build() + .unwrap(); - static ONCE: OnceCell<u32> = OnceCell::const_new(); + static ONCE: OnceCell<u32> = 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 }); + 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(); + time::advance(Duration::from_millis(1)).await; + time::resume(); - let result1 = handle1.await.unwrap(); - assert!(result1.is_err()); + let result1 = handle1.await.unwrap(); + assert!(result1.is_err()); - let result2 = handle2.await.unwrap(); - assert_eq!(*result2.unwrap(), 10); - }); - } + let result2 = handle2.await.unwrap(); + assert_eq!(*result2.unwrap(), 10); + }); } diff --git a/tests/sync_oneshot.rs b/tests/sync_oneshot.rs index 15b6923..163d50d 100644 --- a/tests/sync_oneshot.rs +++ b/tests/sync_oneshot.rs @@ -1,12 +1,12 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))] use tokio::test as maybe_tokio_test; use tokio::sync::oneshot; @@ -179,7 +179,7 @@ fn explicit_close_try_recv() { #[test] #[should_panic] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn close_try_recv_poll() { let (_tx, rx) = oneshot::channel::<i32>(); let mut rx = task::spawn(rx); diff --git a/tests/sync_panic.rs b/tests/sync_panic.rs index 6c23664..bcdaaae 100644 --- a/tests/sync_panic.rs +++ b/tests/sync_panic.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] +#![cfg(all(feature = "full", not(target_os = "wasi")))] use std::{error::Error, sync::Arc}; use tokio::{ diff --git a/tests/sync_rwlock.rs b/tests/sync_rwlock.rs index 948ec13..f711804 100644 --- a/tests/sync_rwlock.rs +++ b/tests/sync_rwlock.rs @@ -1,19 +1,19 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; -#[cfg(not(tokio_wasm_not_wasi))] +#[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; +use tokio::sync::{RwLock, RwLockWriteGuard}; use tokio_test::task::spawn; use tokio_test::{assert_pending, assert_ready}; @@ -172,7 +172,7 @@ async fn write_order() { } // A single RwLock is contested by tasks in multiple threads -#[cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#[cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread", worker_threads = 8)] async fn multithreaded() { use futures::stream::{self, StreamExt}; @@ -279,3 +279,53 @@ fn try_read_try_write() { 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 _); +} diff --git a/tests/sync_semaphore.rs b/tests/sync_semaphore.rs index 3d47ed0..40a5a08 100644 --- a/tests/sync_semaphore.rs +++ b/tests/sync_semaphore.rs @@ -1,6 +1,6 @@ #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use std::sync::Arc; @@ -78,7 +78,7 @@ fn merge() { } #[test] -#[cfg(not(tokio_wasm))] // No stack unwinding on wasm targets +#[cfg(not(target_family = "wasm"))] // No stack unwinding on wasm targets #[should_panic] fn merge_unrelated_permits() { let sem1 = Arc::new(Semaphore::new(3)); @@ -118,7 +118,7 @@ fn add_max_amount_permits() { assert_eq!(s.available_permits(), tokio::sync::Semaphore::MAX_PERMITS); } -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding #[test] #[should_panic] fn add_more_than_max_amount_permits1() { @@ -126,7 +126,7 @@ fn add_more_than_max_amount_permits1() { s.add_permits(tokio::sync::Semaphore::MAX_PERMITS); } -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding #[test] #[should_panic] fn add_more_than_max_amount_permits2() { @@ -135,7 +135,7 @@ fn add_more_than_max_amount_permits2() { s.add_permits(1); } -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding #[test] #[should_panic] fn panic_when_exceeds_maxpermits() { diff --git a/tests/sync_semaphore_owned.rs b/tests/sync_semaphore_owned.rs index f694576..d4b12d4 100644 --- a/tests/sync_semaphore_owned.rs +++ b/tests/sync_semaphore_owned.rs @@ -1,6 +1,6 @@ #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use std::sync::Arc; @@ -104,7 +104,7 @@ fn merge() { } #[test] -#[cfg(not(tokio_wasm))] // No stack unwinding on wasm targets +#[cfg(not(target_family = "wasm"))] // No stack unwinding on wasm targets #[should_panic] fn merge_unrelated_permits() { let sem1 = Arc::new(Semaphore::new(3)); diff --git a/tests/sync_watch.rs b/tests/sync_watch.rs index d4f8ce8..dab57aa 100644 --- a/tests/sync_watch.rs +++ b/tests/sync_watch.rs @@ -2,7 +2,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "sync")] -#[cfg(tokio_wasm_not_wasi)] +#[cfg(all(target_family = "wasm", not(target_os = "wasi")))] use wasm_bindgen_test::wasm_bindgen_test as test; use tokio::sync::watch; @@ -214,7 +214,7 @@ fn reopened_after_subscribe() { #[test] #[cfg(panic = "unwind")] -#[cfg(not(tokio_wasm))] // wasm currently doesn't support unwinding +#[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding fn send_modify_panic() { let (tx, mut rx) = watch::channel("one"); diff --git a/tests/task_abort.rs b/tests/task_abort.rs index 8a1b007..481cc96 100644 --- a/tests/task_abort.rs +++ b/tests/task_abort.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery use std::sync::Arc; use std::thread::sleep; @@ -176,7 +176,7 @@ fn test_abort_wakes_task_3964() { /// Checks that aborting a task whose destructor panics does not allow the /// panic to escape the task. #[test] -#[cfg(not(target_os = "android"))] +#[cfg(panic = "unwind")] fn test_abort_task_that_panics_on_drop_contained() { let rt = Builder::new_current_thread().enable_time().build().unwrap(); @@ -200,7 +200,7 @@ fn test_abort_task_that_panics_on_drop_contained() { /// Checks that aborting a task whose destructor panics has the expected result. #[test] -#[cfg(not(target_os = "android"))] +#[cfg(panic = "unwind")] fn test_abort_task_that_panics_on_drop_returned() { let rt = Builder::new_current_thread().enable_time().build().unwrap(); diff --git a/tests/task_blocking.rs b/tests/task_blocking.rs index d82a0e0..4594498 100644 --- a/tests/task_blocking.rs +++ b/tests/task_blocking.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads use tokio::{runtime, task, time}; use tokio_test::assert_ok; @@ -114,7 +114,7 @@ fn can_enter_current_thread_rt_from_within_block_in_place() { } #[test] -#[cfg(not(target_os = "android"))] +#[cfg(panic = "unwind")] fn useful_panic_message_when_dropping_rt_in_rt() { use std::panic::{catch_unwind, AssertUnwindSafe}; diff --git a/tests/task_id.rs b/tests/task_id.rs index d7b7c0c..e10f24b 100644 --- a/tests/task_id.rs +++ b/tests/task_id.rs @@ -2,21 +2,21 @@ #![allow(clippy::declare_interior_mutable_const)] #![cfg(all(feature = "full", tokio_unstable))] -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use std::error::Error; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use tokio::runtime::{Builder, Runtime}; use tokio::sync::oneshot; use tokio::task::{self, Id, LocalSet}; -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] mod support { pub mod panic; } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use support::panic::test_panic; #[tokio::test(flavor = "current_thread")] @@ -26,7 +26,7 @@ async fn task_id_spawn() { .unwrap(); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "current_thread")] async fn task_id_spawn_blocking() { task::spawn_blocking(|| println!("task id: {}", task::id())) @@ -43,7 +43,7 @@ async fn task_id_collision_current_thread() { assert_ne!(id1.unwrap(), id2.unwrap()); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_collision_multi_thread() { let handle1 = tokio::spawn(async { task::id() }); @@ -64,7 +64,7 @@ async fn task_ids_match_current_thread() { handle.await.unwrap(); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_ids_match_multi_thread() { let (tx, rx) = oneshot::channel(); @@ -76,7 +76,7 @@ async fn task_ids_match_multi_thread() { handle.await.unwrap(); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_future_destructor_completion() { struct MyFuture { @@ -104,7 +104,7 @@ async fn task_id_future_destructor_completion() { assert_eq!(rx.await.unwrap(), id); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_future_destructor_abort() { struct MyFuture { @@ -210,7 +210,7 @@ fn task_try_id_outside_task() { assert_eq!(None, task::try_id()); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[test] fn task_try_id_inside_block_on() { let rt = Runtime::new().unwrap(); @@ -253,7 +253,7 @@ async fn task_id_nested_spawn_local() { .await; } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[tokio::test(flavor = "multi_thread")] async fn task_id_block_in_place_block_on_spawn() { task::spawn(async { @@ -273,7 +273,7 @@ async fn task_id_block_in_place_block_on_spawn() { .unwrap(); } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[test] fn task_id_outside_task_panic_caller() -> Result<(), Box<dyn Error>> { let panic_location_file = test_panic(|| { @@ -286,7 +286,7 @@ fn task_id_outside_task_panic_caller() -> Result<(), Box<dyn Error>> { Ok(()) } -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] #[test] fn task_id_inside_block_on_panic_caller() -> Result<(), Box<dyn Error>> { let panic_location_file = test_panic(|| { diff --git a/tests/task_join_set.rs b/tests/task_join_set.rs index b1b6cf9..d236aa1 100644 --- a/tests/task_join_set.rs +++ b/tests/task_join_set.rs @@ -1,6 +1,7 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full"))] +#![cfg(feature = "full")] +use futures::future::FutureExt; use tokio::sync::oneshot; use tokio::task::JoinSet; use tokio::time::Duration; @@ -184,52 +185,45 @@ async fn abort_all() { assert_eq!(set.len(), 0); } -#[cfg(feature = "parking_lot")] -mod parking_lot { - use super::*; +// 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; - use futures::future::FutureExt; + static SEM: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(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(); - let mut set = JoinSet::new(); + for _ in 0..TASK_NUM { + set.spawn(async { + SEM.add_permits(1); + }); + } - 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(); - // 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, + 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; } - - count += 1; + Some(None) => break, } - assert!(coop_count >= 1); - assert_eq!(count, TASK_NUM); + + count += 1; } + assert!(coop_count >= 1); + assert_eq!(count, TASK_NUM); } diff --git a/tests/task_local.rs b/tests/task_local.rs index 949a40c..a9ffaa1 100644 --- a/tests/task_local.rs +++ b/tests/task_local.rs @@ -1,4 +1,4 @@ -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threads #![allow(clippy::declare_interior_mutable_const)] use std::future::Future; use std::pin::Pin; diff --git a/tests/task_local_set.rs b/tests/task_local_set.rs index 2da87f5..baca58b 100644 --- a/tests/task_local_set.rs +++ b/tests/task_local_set.rs @@ -11,13 +11,13 @@ use tokio::sync::{mpsc, oneshot}; use tokio::task::{self, LocalSet}; use tokio::time; -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use std::cell::Cell; use std::sync::atomic::AtomicBool; -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -#[cfg(not(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use std::sync::atomic::Ordering::SeqCst; use std::time::Duration; @@ -30,7 +30,7 @@ async fn local_current_thread_scheduler() { .await; } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn local_threadpool() { thread_local! { @@ -51,7 +51,7 @@ async fn local_threadpool() { .await; } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn localset_future_threadpool() { thread_local! { @@ -67,7 +67,7 @@ async fn localset_future_threadpool() { local.await; } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[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); @@ -112,7 +112,7 @@ async fn localset_future_drives_all_local_futs() { assert!(RAN3.load(Ordering::SeqCst)); } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[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 @@ -151,7 +151,7 @@ fn enter_guard_spawn() { }); } -#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery +#[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`. @@ -178,7 +178,7 @@ fn local_threadpool_blocking_in_place() { }); } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn local_threadpool_blocking_run() { thread_local! { @@ -207,7 +207,7 @@ async fn local_threadpool_blocking_run() { .await; } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn all_spawns_are_local() { use futures::future; @@ -234,7 +234,7 @@ async fn all_spawns_are_local() { .await; } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn nested_spawn_is_local() { thread_local! { @@ -270,7 +270,7 @@ async fn nested_spawn_is_local() { .await; } -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[test] fn join_local_future_elsewhere() { thread_local! { @@ -307,7 +307,7 @@ fn join_local_future_elsewhere() { } // Tests for <https://github.com/tokio-rs/tokio/issues/4973> -#[cfg(not(tokio_wasi))] // Wasi doesn't support threads +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support threads #[tokio::test(flavor = "multi_thread")] async fn localset_in_thread_local() { thread_local! { @@ -398,7 +398,10 @@ fn with_timeout(timeout: Duration, f: impl FnOnce() + Send + 'static) { thread.join().expect("test thread should not panic!") } -#[cfg_attr(tokio_wasi, ignore = "`unwrap()` in `with_timeout()` panics on Wasi")] +#[cfg_attr( + target_os = "wasi", + ignore = "`unwrap()` in `with_timeout()` panics on Wasi" +)] #[test] fn drop_cancels_remote_tasks() { // This test reproduces issue #1885. @@ -422,7 +425,7 @@ fn drop_cancels_remote_tasks() { } #[cfg_attr( - tokio_wasi, + target_os = "wasi", ignore = "FIXME: `task::spawn_local().await.unwrap()` panics on Wasi" )] #[test] @@ -446,7 +449,7 @@ fn local_tasks_wake_join_all() { }); } -#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery +#[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. @@ -463,7 +466,7 @@ fn local_tasks_are_polled_after_tick() { local_tasks_are_polled_after_tick_inner(); } -#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery +#[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 diff --git a/tests/task_panic.rs b/tests/task_panic.rs index e4cedce..2ca21b1 100644 --- a/tests/task_panic.rs +++ b/tests/task_panic.rs @@ -1,6 +1,6 @@ #![warn(rust_2018_idioms)] #![allow(clippy::declare_interior_mutable_const)] -#![cfg(all(feature = "full", not(tokio_wasi)))] +#![cfg(all(feature = "full", not(target_os = "wasi")))] use futures::future; use std::error::Error; diff --git a/tests/task_yield_now.rs b/tests/task_yield_now.rs new file mode 100644 index 0000000..b16bca5 --- /dev/null +++ b/tests/task_yield_now.rs @@ -0,0 +1,16 @@ +#![cfg(all(feature = "full", 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()); +} diff --git a/tests/tcp_accept.rs b/tests/tcp_accept.rs index ba4a498..d547c48 100644 --- a/tests/tcp_accept.rs +++ b/tests/tcp_accept.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{mpsc, oneshot}; diff --git a/tests/tcp_connect.rs b/tests/tcp_connect.rs index f238442..269ba8e 100644 --- a/tests/tcp_connect.rs +++ b/tests/tcp_connect.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::net::{TcpListener, TcpStream}; use tokio::sync::oneshot; @@ -35,7 +35,6 @@ async fn connect_v4() { } #[tokio::test] -#[cfg(not(tokio_no_ipv6))] async fn connect_v6() { let srv = assert_ok!(TcpListener::bind("[::1]:0").await); let addr = assert_ok!(srv.local_addr()); diff --git a/tests/tcp_echo.rs b/tests/tcp_echo.rs index f47d4ac..dfd4dd7 100644 --- a/tests/tcp_echo.rs +++ b/tests/tcp_echo.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::io::{self, AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; diff --git a/tests/tcp_into_split.rs b/tests/tcp_into_split.rs index 010984a..2a03069 100644 --- a/tests/tcp_into_split.rs +++ b/tests/tcp_into_split.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use std::io::{Error, ErrorKind, Result}; use std::io::{Read, Write}; diff --git a/tests/tcp_into_std.rs b/tests/tcp_into_std.rs index 320d994..58996f7 100644 --- a/tests/tcp_into_std.rs +++ b/tests/tcp_into_std.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use std::io::Read; use std::io::Result; diff --git a/tests/tcp_peek.rs b/tests/tcp_peek.rs index b712023..b2d3eda 100644 --- a/tests/tcp_peek.rs +++ b/tests/tcp_peek.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::io::AsyncReadExt; use tokio::net::TcpStream; @@ -7,7 +7,7 @@ use tokio::net::TcpStream; use tokio_test::assert_ok; use std::thread; -use std::{convert::TryInto, io::Write, net}; +use std::{io::Write, net}; #[tokio::test] async fn peek() { diff --git a/tests/tcp_shutdown.rs b/tests/tcp_shutdown.rs index 1d477f3..86c1485 100644 --- a/tests/tcp_shutdown.rs +++ b/tests/tcp_shutdown.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::io::{self, AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; diff --git a/tests/tcp_socket.rs b/tests/tcp_socket.rs index 66309d3..faead95 100644 --- a/tests/tcp_socket.rs +++ b/tests/tcp_socket.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use std::time::Duration; use tokio::net::TcpSocket; @@ -24,7 +24,6 @@ async fn basic_usage_v4() { } #[tokio::test] -#[cfg(not(tokio_no_ipv6))] async fn basic_usage_v6() { // Create server let addr = assert_ok!("[::1]:0".parse()); diff --git a/tests/tcp_split.rs b/tests/tcp_split.rs index 335b21f..32aaa1f 100644 --- a/tests/tcp_split.rs +++ b/tests/tcp_split.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use std::io::Result; use std::io::{Read, Write}; diff --git a/tests/tcp_stream.rs b/tests/tcp_stream.rs index 31fe3ba..725a601 100644 --- a/tests/tcp_stream.rs +++ b/tests/tcp_stream.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support bind +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support bind use tokio::io::{AsyncReadExt, AsyncWriteExt, Interest}; use tokio::net::{TcpListener, TcpStream}; diff --git a/tests/time_interval.rs b/tests/time_interval.rs index 186582e..9b7a6fd 100644 --- a/tests/time_interval.rs +++ b/tests/time_interval.rs @@ -1,10 +1,12 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "full")] -use tokio::time::{self, Duration, Instant, MissedTickBehavior}; -use tokio_test::{assert_pending, assert_ready_eq, task}; +use std::pin::Pin; +use std::task::{Context, Poll}; -use std::task::Poll; +use futures::{Stream, StreamExt}; +use tokio::time::{self, Duration, Instant, Interval, MissedTickBehavior}; +use tokio_test::{assert_pending, 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 @@ -202,6 +204,154 @@ async fn reset() { check_interval_poll!(i, start, 1001); } +#[tokio::test(start_paused = true)] +async fn reset_immediatelly() { + 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<time::Interval>) -> Poll<Instant> { interval.enter(|cx, mut interval| interval.poll_tick(cx)) } @@ -209,3 +359,113 @@ fn poll_next(interval: &mut task::Spawn<time::Interval>) -> Poll<Instant> { 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<Option<Self::Item>> { + 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![]); +} diff --git a/tests/time_panic.rs b/tests/time_panic.rs index aaff11b..02b928b 100644 --- a/tests/time_panic.rs +++ b/tests/time_panic.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support panic recovery use futures::future; use std::error::Error; diff --git a/tests/time_pause.rs b/tests/time_pause.rs index c772e38..6251b4b 100644 --- a/tests/time_pause.rs +++ b/tests/time_pause.rs @@ -6,7 +6,7 @@ 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(tokio_wasi))] +#[cfg(not(target_os = "wasi"))] use tokio_test::assert_err; use std::{ @@ -29,14 +29,14 @@ async fn pause_time_in_task() { t.await.unwrap(); } -#[cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#[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(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads +#[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 { diff --git a/tests/time_rt.rs b/tests/time_rt.rs index 20f9e18..13f888c 100644 --- a/tests/time_rt.rs +++ b/tests/time_rt.rs @@ -5,7 +5,7 @@ use tokio::time::*; use std::sync::mpsc; -#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] // Wasi doesn't support threads +#[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; diff --git a/tests/time_sleep.rs b/tests/time_sleep.rs index 4174a73..4afcc12 100644 --- a/tests/time_sleep.rs +++ b/tests/time_sleep.rs @@ -168,7 +168,7 @@ async fn reset_sleep_to_past() { assert_ready!(sleep.poll()); } -#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery +#[cfg(not(target_os = "wasi"))] // Wasi doesn't support panic recovery #[test] #[should_panic] fn creating_sleep_outside_of_context() { @@ -268,6 +268,20 @@ async fn exactly_max() { } #[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; diff --git a/tests/time_timeout.rs b/tests/time_timeout.rs index be6b8fb..ec871cf 100644 --- a/tests/time_timeout.rs +++ b/tests/time_timeout.rs @@ -17,7 +17,7 @@ async fn simultaneous_deadline_future_completion() { assert_ready_ok!(fut.poll()); } -#[cfg_attr(tokio_wasi, ignore = "FIXME: `fut.poll()` panics on Wasi")] +#[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 diff --git a/tests/udp.rs b/tests/udp.rs index 2b6ab4d..eea281c 100644 --- a/tests/udp.rs +++ b/tests/udp.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support bind or UDP +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support bind or UDP use futures::future::poll_fn; use std::io; @@ -107,6 +107,45 @@ async fn send_to_peek_from() -> std::io::Result<()> { } #[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?; @@ -135,6 +174,92 @@ async fn send_to_peek_from_poll() -> std::io::Result<()> { } #[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); @@ -400,6 +525,23 @@ async fn try_recv_buf() { } #[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(); @@ -443,6 +585,23 @@ async fn try_recv_buf_from() { } #[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(); diff --git a/tests/uds_datagram.rs b/tests/uds_datagram.rs index c08bd45..ad22a0b 100644 --- a/tests/uds_datagram.rs +++ b/tests/uds_datagram.rs @@ -277,6 +277,28 @@ async fn try_recv_buf_from() -> std::io::Result<()> { Ok(()) } +#[tokio::test] +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] async fn try_recv_buf_never_block() -> io::Result<()> { @@ -326,6 +348,24 @@ async fn try_recv_buf_never_block() -> io::Result<()> { } #[tokio::test] +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] async fn poll_ready() -> io::Result<()> { let dir = tempfile::tempdir().unwrap(); let server_path = dir.path().join("server.sock"); diff --git a/tests/unwindsafe.rs b/tests/unwindsafe.rs index 3e63820..8ab6654 100644 --- a/tests/unwindsafe.rs +++ b/tests/unwindsafe.rs @@ -1,5 +1,5 @@ #![warn(rust_2018_idioms)] -#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi does not support panic recovery +#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery use std::panic::{RefUnwindSafe, UnwindSafe}; |