chore: prepare Tokio v1.20.5 (#5731)

Some of these changes will be progressively reverted as we merge this
into newer branches.

.cargo/audit.toml

@@ -2,6 +2,7 @@
 
 [advisories]
 ignore = [
-    # https://github.com/tokio-rs/tokio/issues/4177
-    "RUSTSEC-2020-0159",
+    # We depend on nix 0.22 only via mio-aio, a dev-dependency.
+    # https://github.com/tokio-rs/tokio/pull/4255#issuecomment-974786349
+    "RUSTSEC-2021-0119",
 ]

.cargo/config

@@ -0,0 +1,2 @@
+# [build]
+# rustflags = ["--cfg", "tokio_unstable"]

.circleci/config.yml

@@ -5,8 +5,8 @@ jobs:
       image: ubuntu-2004:202101-01
     resource_class: arm.medium
     environment:
-      # Change to pin rust versino
-      RUST_STABLE: stable
+      # Change to pin rust version
+      RUST_STABLE: 1.62.1
     steps:
       - checkout
       - run:

.cirrus.yml

@@ -1,5 +1,11 @@
+only_if: $CIRRUS_TAG == '' && ($CIRRUS_PR != '' || $CIRRUS_BRANCH == 'master' || $CIRRUS_BRANCH =~ 'tokio-.*')
+auto_cancellation: $CIRRUS_BRANCH != 'master' && $CIRRUS_BRANCH !=~ 'tokio-.*'
 freebsd_instance:
-  image: freebsd-12-2-release-amd64
+  image_family: freebsd-13-1
+env:
+  RUST_STABLE: 1.62.1
+  RUST_NIGHTLY: nightly-2022-03-21
+  RUSTFLAGS: -D warnings
 
 # Test FreeBSD in a full VM on cirrus-ci.com.  Test the i686 target too, in the
 # same VM.  The binary will be built in 32-bit mode, but will execute on a
@@ -8,9 +14,9 @@ freebsd_instance:
 task:
   name: FreeBSD 64-bit
   setup_script:
-    - pkg install -y bash curl
+    - pkg install -y bash
     - curl https://sh.rustup.rs -sSf --output rustup.sh
-    - sh rustup.sh -y --profile minimal --default-toolchain stable
+    - sh rustup.sh -y --profile minimal --default-toolchain $RUST_STABLE
     - . $HOME/.cargo/env
     - |
       echo "~~~~ rustc --version ~~~~"
@@ -22,12 +28,12 @@ task:
 task:
   name: FreeBSD docs
   env:
-    RUSTFLAGS: --cfg docsrs
-    RUSTDOCFLAGS: --cfg docsrs
+    RUSTFLAGS: --cfg docsrs --cfg tokio_unstable
+    RUSTDOCFLAGS: --cfg docsrs --cfg tokio_unstable -Dwarnings
   setup_script:
-    - pkg install -y bash curl
+    - pkg install -y bash
     - curl https://sh.rustup.rs -sSf --output rustup.sh
-    - sh rustup.sh -y --profile minimal --default-toolchain nightly-2021-10-25
+    - sh rustup.sh -y --profile minimal --default-toolchain $RUST_NIGHTLY
     - . $HOME/.cargo/env
     - |
       echo "~~~~ rustc --version ~~~~"
@@ -39,9 +45,9 @@ task:
 task:
   name: FreeBSD 32-bit
   setup_script:
-    - pkg install -y bash curl
+    - pkg install -y bash
     - curl https://sh.rustup.rs -sSf --output rustup.sh
-    - sh rustup.sh -y --profile minimal --default-toolchain stable
+    - sh rustup.sh -y --profile minimal --default-toolchain $RUST_STABLE
     - . $HOME/.cargo/env
     - rustup target add i686-unknown-freebsd
     - |

.clippy.toml

@@ -1 +1 @@
-msrv = "1.45"
+msrv = "1.49"

.github/labeler.yml

@@ -1,9 +1,8 @@
 
 R-loom:
-- ./tokio/src/sync/*
-- ./tokio/src/sync/**/*
-- ./tokio-util/src/sync/*
-- ./tokio-util/src/sync/**/*
-- ./tokio/src/runtime/*
-- ./tokio/src/runtime/**/*
-
+- tokio/src/sync/*
+- tokio/src/sync/**/*
+- tokio-util/src/sync/*
+- tokio-util/src/sync/**/*
+- tokio/src/runtime/*
+- tokio/src/runtime/**/*

.github/workflows/audit.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     if: "!contains(github.event.head_commit.message, 'ci skip')"
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
 
       - name: Audit Check
         uses: actions-rs/audit-check@v1

.github/workflows/ci.yml

@@ -7,9 +7,25 @@ on:
 name: CI
 
 env:
+  RUSTFLAGS: -Dwarnings
   RUST_BACKTRACE: 1
-  nightly: nightly-2021-10-25
-  minrust: 1.45.2
+  # Change to specific Rust release to pin
+  rust_stable: 1.62.1
+  rust_nightly: nightly-2022-04-18
+  rust_clippy: 1.56.0
+  # When updating this, also update:
+  # - README.md
+  # - tokio/README.md
+  # - CONTRIBUTING.md
+  # - tokio/Cargo.toml
+  # - tokio-util/Cargo.toml
+  # - tokio-test/Cargo.toml
+  # - tokio-stream/Cargo.toml
+  rust_min: 1.49.0
+
+defaults:
+  run:
+    shell: bash
 
 jobs:
   # Depends on all action sthat are required for a "successful" CI run.
@@ -19,16 +35,19 @@ jobs:
     needs:
       - test
       - test-unstable
+      - test-parking_lot
       - miri
       - cross
       - features
       - minrust
       - fmt
-      - clippy
       - docs
       - valgrind
       - loom-compile
       - check-readme
+      - test-hyper
+      - wasm32-unknown-unknown
+      - wasm32-wasi
     steps:
       - run: exit 0
 
@@ -42,12 +61,17 @@ jobs:
           - ubuntu-latest
           - macos-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+            toolchain: ${{ env.rust_stable }}
+            override: true
       - name: Install Rust
         run: rustup update stable
       - uses: Swatinem/rust-cache@v1
       - name: Install cargo-hack
-        run: cargo install cargo-hack
+        run: cargo install cargo-hack --version 0.5.26
 
       # Run `tokio` with `full` features. This excludes testing utilities which
       # can alter the runtime behavior of Tokio.
@@ -64,6 +88,11 @@ jobs:
         run: cargo hack test --each-feature
         working-directory: tests-integration
 
+      # Run macro build tests
+      - name: test tests-build --each-feature
+        run: cargo hack test --each-feature
+        working-directory: tests-build
+
       # Build benchmarks. Run of benchmarks is done by bench.yml workflow.
       - name: build benches
         run: cargo build --benches
@@ -71,13 +100,40 @@ jobs:
         # bench.yml workflow runs benchmarks only on linux.
         if: startsWith(matrix.os, 'ubuntu')
 
+  test-parking_lot:
+    # The parking_lot crate has a feature called send_guard which changes when
+    # some of its types are Send. Tokio has some measures in place to prevent
+    # this from affecting when Tokio types are Send, and this test exists to
+    # ensure that those measures are working.
+    #
+    # This relies on the potentially affected Tokio type being listed in
+    # `tokio/tokio/tests/async_send_sync.rs`.
+    name: compile tests with parking lot send guards
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+            toolchain: ${{ env.rust_stable }}
+            override: true
+      - uses: Swatinem/rust-cache@v1
+      - name: Enable parking_lot send_guard feature
+        # Inserts the line "plsend = ["parking_lot/send_guard"]" right after [features]
+        run: sed -i '/\[features\]/a plsend = ["parking_lot/send_guard"]' tokio/Cargo.toml
+      - name: Compile tests with all features enabled
+        run: cargo build --workspace --all-features --tests
+
   valgrind:
     name: valgrind
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update stable
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+            toolchain: ${{ env.rust_stable }}
+            override: true
       - uses: Swatinem/rust-cache@v1
 
       - name: Install Valgrind
@@ -113,54 +169,64 @@ jobs:
           - ubuntu-latest
           - macos-latest
     steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update stable
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+            toolchain: ${{ env.rust_stable }}
+            override: true
       - uses: Swatinem/rust-cache@v1
-
       # Run `tokio` with "unstable" cfg flag.
       - name: test tokio full --cfg unstable
         run: cargo test --all-features
         working-directory: tokio
         env:
-          RUSTFLAGS: --cfg tokio_unstable
+          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
+          # in order to run doctests for unstable features, we must also pass
+          # the unstable cfg to RustDoc
+          RUSTDOCFLAGS: --cfg tokio_unstable
 
   miri:
     name: miri
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_nightly }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: ${{ env.nightly }}
+          toolchain: nightly-2022-07-10
+          components: miri
           override: true
       - uses: Swatinem/rust-cache@v1
-      - name: Install Miri
-        run: |
-          set -e
-          rustup component add miri
-          cargo miri setup
-          rm -rf tokio/tests
-
       - name: miri
-        run: cargo miri test --features rt,rt-multi-thread,sync task
+        # Many of tests in tokio/tests and doctests use #[tokio::test] or
+        # #[tokio::main] that calls epoll_create1 that Miri does not support.
+        run: cargo miri test --features full --lib --no-fail-fast
         working-directory: tokio
-  san:
-    name: san
+        env:
+          MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-tag-raw-pointers
+          PROPTEST_CASES: 10
+
+  asan:
+    name: asan
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install llvm
+        # Required to resolve symbols in sanitizer output
+        run: sudo apt-get install -y llvm
+      - name: Install Rust ${{ env.rust_nightly }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: ${{ env.nightly }}
+          toolchain: ${{ env.rust_nightly }}
           override: true
       - uses: Swatinem/rust-cache@v1
       - name: asan
-        run: cargo test --all-features --target x86_64-unknown-linux-gnu --lib -- --test-threads 1
-        working-directory: tokio
+        run: cargo test --workspace --all-features --target x86_64-unknown-linux-gnu --tests -- --test-threads 1
         env:
           RUSTFLAGS: -Z sanitizer=address
-          ASAN_OPTIONS: detect_leaks=0
+          # Ignore `trybuild` errors as they are irrelevant and flaky on nightly
+          TRYBUILD: overwrite
 
   cross:
     name: cross
@@ -173,67 +239,78 @@ jobs:
           - powerpc64-unknown-linux-gnu
           - mips-unknown-linux-gnu
           - arm-linux-androideabi
+          - mipsel-unknown-linux-musl
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: stable
+          toolchain: ${{ env.rust_stable }}
           target: ${{ matrix.target }}
           override: true
-      - uses: Swatinem/rust-cache@v1
       - uses: actions-rs/cargo@v1
         with:
           use-cross: true
           command: check
-          args: --workspace --target ${{ matrix.target }}
+          args: --workspace --all-features --target ${{ matrix.target }}
+      - uses: actions-rs/cargo@v1
+        with:
+          use-cross: true
+          command: check
+          args: --workspace --all-features --target ${{ matrix.target }}
+        env:
+          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
 
   features:
     name: features
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_nightly }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: ${{ env.nightly }}
+          toolchain: ${{ env.rust_nightly }}
+          target: ${{ matrix.target }}
           override: true
       - uses: Swatinem/rust-cache@v1
       - name: Install cargo-hack
-        run: cargo install cargo-hack
-
+        run: cargo install cargo-hack --version 0.5.26
       - name: check --each-feature
         run: cargo hack check --all --each-feature -Z avoid-dev-deps
-
       # Try with unstable feature flags
       - name: check --each-feature --unstable
         run: cargo hack check --all --each-feature -Z avoid-dev-deps
         env:
-          RUSTFLAGS: --cfg tokio_unstable
+          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
 
   minrust:
     name: minrust
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_min }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: ${{ env.minrust }}
+          toolchain: ${{ env.rust_min }}
           override: true
       - uses: Swatinem/rust-cache@v1
-      - name: "test --workspace --all-features"
-        run: cargo check --workspace --all-features
+      - name: "test --all-features"
+        run: cargo check --all-features
+        working-directory: tokio
 
   minimal-versions:
     name: minimal-versions
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_nightly }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: ${{ env.nightly }}
+          toolchain: ${{ env.rust_nightly }}
           override: true
       - uses: Swatinem/rust-cache@v1
       - name: Install cargo-hack
-        run: cargo install cargo-hack
+        run: cargo install cargo-hack --version 0.5.26
       - name: "check --all-features -Z minimal-versions"
         run: |
           # Remove dev-dependencies from Cargo.toml to prevent the next `cargo update`
@@ -242,18 +319,29 @@ jobs:
           # Update Cargo.lock to minimal version dependencies.
           cargo update -Z minimal-versions
           cargo hack check --all-features --ignore-private
+      - name: "check --all-features --unstable -Z minimal-versions"
+        env:
+          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
+        run: |
+          # Remove dev-dependencies from Cargo.toml to prevent the next `cargo update`
+          # from determining minimal versions based on dev-dependencies.
+          cargo hack --remove-dev-deps --workspace
+          # Update Cargo.lock to minimal version dependencies.
+          cargo update -Z minimal-versions
+          cargo hack check --all-features --ignore-private
 
   fmt:
     name: fmt
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update stable
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{ env.rust_stable }}
+          override: true
+          components: rustfmt
       - uses: Swatinem/rust-cache@v1
-      - name: Install rustfmt
-        run: rustup component add rustfmt
-
       # Check fmt
       - name: "rustfmt --check"
         # Workaround for rust-lang/cargo#7732
@@ -263,59 +351,135 @@ jobs:
             exit 1
           fi
 
-  clippy:
-    name: clippy
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update 1.52.1 && rustup default 1.52.1
-      - uses: Swatinem/rust-cache@v1
-      - name: Install clippy
-        run: rustup component add clippy
-
-      # Run clippy
-      - name: "clippy --all"
-        run: cargo clippy --all --tests --all-features
-
   docs:
     name: docs
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_nightly }}
+        uses: actions-rs/toolchain@v1
         with:
-          toolchain: ${{ env.nightly }}
+          toolchain: ${{ env.rust_nightly }}
           override: true
       - uses: Swatinem/rust-cache@v1
       - name: "doc --lib --all-features"
         run: cargo doc --lib --no-deps --all-features --document-private-items
         env:
-          RUSTFLAGS: --cfg docsrs
-          RUSTDOCFLAGS: --cfg docsrs
+          RUSTFLAGS: --cfg docsrs --cfg tokio_unstable
+          RUSTDOCFLAGS: --cfg docsrs --cfg tokio_unstable -Dwarnings
 
   loom-compile:
     name: build loom tests
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update stable
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{ env.rust_stable }}
+          override: true
       - uses: Swatinem/rust-cache@v1
       - name: build --cfg loom
         run: cargo test --no-run --lib --features full
         working-directory: tokio
         env:
-          RUSTFLAGS: --cfg loom --cfg tokio_unstable
+          RUSTFLAGS: --cfg loom --cfg tokio_unstable -Dwarnings
 
   check-readme:
     name: Check README
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Verify that both READMEs are identical
         run: diff README.md tokio/README.md
 
       - name: Verify that Tokio version is up to date in README
         working-directory: tokio
         run: grep -q "$(sed '/^version = /!d' Cargo.toml | head -n1)" README.md
+
+  test-hyper:
+    name: Test hyper
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os:
+          - windows-latest
+          - ubuntu-latest
+          - macos-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{ env.rust_stable }}
+          override: true
+      - uses: Swatinem/rust-cache@v1
+      - name: Test hyper
+        run: |
+          set -x
+          git clone https://github.com/hyperium/hyper.git
+          cd hyper
+          # checkout the latest release because HEAD maybe contains breakage.
+          tag=$(git describe --abbrev=0 --tags)
+          git checkout "${tag}"
+          echo '[workspace]' >>Cargo.toml
+          echo '[patch.crates-io]' >>Cargo.toml
+          echo 'tokio = { path = "../tokio" }' >>Cargo.toml
+          echo 'tokio-util = { path = "../tokio-util" }' >>Cargo.toml
+          echo 'tokio-stream = { path = "../tokio-stream" }' >>Cargo.toml
+          echo 'tokio-test = { path = "../tokio-test" }' >>Cargo.toml
+          git diff
+          cargo test --features full
+
+  wasm32-unknown-unknown:
+    name: test tokio for wasm32-unknown-unknown
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: ${{ env.rust_stable }}
+          override: true
+      - uses: Swatinem/rust-cache@v1
+      - name: Install wasm-pack
+        run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
+      - name: test tokio
+        run: wasm-pack test --node -- --features "macros sync"
+        working-directory: tokio
+
+  wasm32-wasi:
+    name: wasm32-wasi
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust stable
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+          override: true
+      - uses: Swatinem/rust-cache@v1
+
+      # Install dependencies
+      - name: Install cargo-hack
+        run: cargo install cargo-hack
+
+      - name: Install wasm32-wasi target
+        run: rustup target add wasm32-wasi
+
+      - name: Install wasmtime
+        run: cargo install wasmtime-cli
+
+      - name: Install cargo-wasi
+        run: cargo install cargo-wasi
+
+      # TODO: Expand this when full WASI support lands.
+      #   Currently, this is a bare bones regression test
+      #   for features that work today with wasi.
+
+      - name: test tests-integration --features wasi-rt
+        # TODO: this should become: `cargo hack wasi test --each-feature`
+        run: cargo wasi test --test rt_yield --features wasi-rt
+        working-directory: tests-integration
+        env:
+          RUSTFLAGS: ""

.github/workflows/labeler.yml

@@ -11,3 +11,4 @@ jobs:
     - uses: actions/labeler@v3
       with:
         repo-token: "${{ secrets.GITHUB_TOKEN }}"
+        sync-labels: true

.github/workflows/loom.yml

@@ -8,7 +8,10 @@ on:
 name: Loom
 
 env:
+  RUSTFLAGS: -Dwarnings
   RUST_BACKTRACE: 1
+  # Change to specific Rust release to pin
+  rust_stable: 1.62.1
 
 jobs:
   loom:
@@ -26,14 +29,17 @@ jobs:
           - loom_pool::group_d
           - time::driver
     steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update stable
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+            toolchain: ${{ env.rust_stable }}
+            override: true
       - uses: Swatinem/rust-cache@v1
       - name: loom ${{ matrix.scope }}
         run: cargo test --lib --release --features full -- --nocapture $SCOPE
         working-directory: tokio
         env:
-          RUSTFLAGS: --cfg loom --cfg tokio_unstable
+          RUSTFLAGS: --cfg loom --cfg tokio_unstable -Dwarnings
           LOOM_MAX_PREEMPTIONS: 2
           SCOPE: ${{ matrix.scope }}

.github/workflows/pr-audit.yml

@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     if: "!contains(github.event.head_commit.message, 'ci skip')"
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
 
       - name: Install cargo-audit
         uses: actions-rs/cargo@v1

.github/workflows/stress-test.yml

@@ -5,6 +5,12 @@ on:
     branches:
       - master
 
+env:
+  RUSTFLAGS: -Dwarnings
+  RUST_BACKTRACE: 1
+  # Change to specific Rust release to pin
+  rust_stable: 1.62.1
+
 jobs:
   stess-test:
     name: Stress Test
@@ -14,9 +20,12 @@ jobs:
         stress-test:
           - simple_echo_tcp
     steps:
-      - uses: actions/checkout@v2
-      - name: Install Rust
-        run: rustup update stable
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: actions-rs/toolchain@v1
+        with:
+            toolchain: ${{ env.rust_stable }}
+            override: true
       - uses: Swatinem/rust-cache@v1
       - name: Install Valgrind
         run: |

.gitignore

@@ -1,2 +1,4 @@
 target
 Cargo.lock
+
+.cargo/config.toml

CONTRIBUTING.md

@@ -131,6 +131,24 @@ cargo check --all-features
 cargo test --all-features
 ```
 
+Clippy must be run using the MSRV, so Tokio can avoid having to `#[allow]` new
+lints whose fixes would be incompatible with the current MSRV:
+
+<!--
+When updating this, also update:
+- .github/workflows/ci.yml
+- README.md
+- tokio/README.md
+- tokio/Cargo.toml
+- tokio-util/Cargo.toml
+- tokio-test/Cargo.toml
+- tokio-stream/Cargo.toml
+-->
+
+```
+cargo +1.49.0 clippy --all-features
+```
+
 When building documentation normally, the markers that list the features
 required for various parts of Tokio are missing. To build the documentation
 correctly, use this command:
@@ -139,6 +157,14 @@ correctly, use this command:
 RUSTDOCFLAGS="--cfg docsrs" cargo +nightly doc --all-features
 ```
 
+To build documentation including Tokio's unstable features, it is necessary to
+pass `--cfg tokio_unstable` to both RustDoc *and* rustc. To build the
+documentation for unstable features, use this command:
+
+```
+RUSTDOCFLAGS="--cfg docsrs --cfg tokio_unstable" RUSTFLAGS="--cfg tokio_unstable" cargo +nightly doc --all-features
+```
+
 There is currently a [bug in cargo] that means documentation cannot be built
 from the root of the workspace. If you `cd` into the `tokio` subdirectory the
 command shown above will work.
@@ -165,6 +191,12 @@ LOOM_MAX_PREEMPTIONS=1 RUSTFLAGS="--cfg loom" \
     cargo test --lib --release --features full -- --test-threads=1 --nocapture
 ```
 
+You can run miri tests with
+```
+MIRIFLAGS="-Zmiri-disable-isolation -Zmiri-tag-raw-pointers" PROPTEST_CASES=10 \
+    cargo +nightly miri test --features full --lib
+```
+
 ### Tests
 
 If the change being proposed alters code (as opposed to only documentation for

Cargo.lock

@@ -0,0 +1,32 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "predicates"
+version = "1.0.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f49cfaf7fdaa3bfacc6fa3e7054e65148878354a5cfddcf661df4c851f8021df"
+dependencies = [
+ "difference",
+ "float-cmp",
+ "normalize-line-endings",
+ "predicates-core",
+ "regex",
+]
+
+[[package]]
+name = "predicates-core"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72f883590242d3c6fc5bf50299011695fa6590c2c70eac95ee1bdb9a733ad1a2"
+
+[[package]]
+name = "predicates-tree"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54ff541861505aabf6ea722d2131ee980b8276e10a1297b94e896dd8b621850d"
+dependencies = [
+ "predicates-core",
+ "termtree",
+]

Cargo.toml

@@ -11,5 +11,6 @@ members = [
   "benches",
   "examples",
   "stress-test",
+  "tests-build",
   "tests-integration",
 ]

Cross.toml

@@ -0,0 +1,4 @@
+[build.env]
+passthrough = [
+    "RUSTFLAGS",
+]

LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2021 Tokio Contributors
+Copyright (c) 2022 Tokio Contributors
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

README.md

@@ -56,7 +56,7 @@ Make sure you activated the full features of the tokio crate on Cargo.toml:
 
 ```toml
 [dependencies]
-tokio = { version = "1.13.0", features = ["full"] }
+tokio = { version = "1.20.5", features = ["full"] }
 ```
 Then, on your main.rs:
 
@@ -163,9 +163,21 @@ several other libraries, including:
 
 ## Supported Rust Versions
 
-Tokio is built against the latest stable release. The minimum supported version
-is 1.45.  The current Tokio version is not guaranteed to build on Rust versions
-earlier than the minimum supported version.
+<!--
+When updating this, also update:
+- .github/workflows/ci.yml
+- CONTRIBUTING.md
+- README.md
+- tokio/README.md
+- tokio/Cargo.toml
+- tokio-util/Cargo.toml
+- tokio-test/Cargo.toml
+- tokio-stream/Cargo.toml
+-->
+
+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.
 
 ## Release schedule
 
@@ -180,17 +192,18 @@ 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.8.x` - LTS release until February 2022.
+ * `1.14.x` - LTS release until June 2022.
+ * `1.18.x` - LTS release until January 2023
 
 Each LTS release will continue to receive backported fixes for at least half a
 year. If you wish to use a fixed minor release in your project, we recommend
 that you use an LTS release.
 
 To use a fixed minor version, you can specify the version with a tilde. For
-example, to specify that you wish to use the newest `1.8.x` patch release, you
+example, to specify that you wish to use the newest `1.14.x` patch release, you
 can use the following dependency specification:
 ```text
-tokio = { version = "~1.8", features = [...] }
+tokio = { version = "~1.14", features = [...] }
 ```
 
 ## License

benches/Cargo.toml

@@ -9,7 +9,7 @@ tokio = { version = "1.5.0", path = "../tokio", features = ["full"] }
 bencher = "0.1.5"
 
 [dev-dependencies]
-tokio-util = { version = "0.6.6", path = "../tokio-util", features = ["full"] }
+tokio-util = { version = "0.7.0", path = "../tokio-util", features = ["full"] }
 tokio-stream = { path = "../tokio-stream" }
 
 [target.'cfg(unix)'.dependencies]

benches/fs.rs

@@ -21,7 +21,7 @@ fn rt() -> tokio::runtime::Runtime {
 const BLOCK_COUNT: usize = 1_000;
 
 const BUFFER_SIZE: usize = 4096;
-const DEV_ZERO: &'static str = "/dev/zero";
+const DEV_ZERO: &str = "/dev/zero";
 
 fn async_read_codec(b: &mut Bencher) {
     let rt = rt();

examples/Cargo.toml

@@ -5,14 +5,14 @@ publish = false
 edition = "2018"
 
 # If you copy one of the examples into a new project, you should be using
-# [dependencies] instead.
+# [dependencies] instead, and delete the **path**.
 [dev-dependencies]
-tokio = { version = "1.0.0", path = "../tokio",features = ["full", "tracing"] }
-tokio-util = { version = "0.6.3", path = "../tokio-util",features = ["full"] }
+tokio = { version = "1.0.0", path = "../tokio", features = ["full", "tracing"] }
+tokio-util = { version = "0.7.0", path = "../tokio-util", features = ["full"] }
 tokio-stream = { version = "0.1", path = "../tokio-stream" }
 
 tracing = "0.1"
-tracing-subscriber = { version = "0.2.7", default-features = false, features = ["fmt", "ansi", "env-filter", "chrono", "tracing-log"] }
+tracing-subscriber = { version = "0.3.1", default-features = false, features = ["fmt", "ansi", "env-filter", "tracing-log"] }
 bytes = "1.0.0"
 futures = { version = "0.3.0", features = ["thread-pool"]}
 http = "0.2"
@@ -21,7 +21,6 @@ serde_derive = "1.0"
 serde_json = "1.0"
 httparse = "1.0"
 httpdate = "1.0"
-once_cell = "1.5.2"
 rand = "0.8.3"
 
 [target.'cfg(windows)'.dev-dependencies.winapi]
@@ -71,11 +70,6 @@ path = "udp-codec.rs"
 name = "tinyhttp"
 path = "tinyhttp.rs"
 
-[[example]]
-name = "custom-executor"
-path = "custom-executor.rs"
-
-
 [[example]]
 name = "custom-executor-tokio-context"
 path = "custom-executor-tokio-context.rs"

examples/tinydb.rs

@@ -149,7 +149,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
 }
 
 fn handle_request(line: &str, db: &Arc<Database>) -> Response {
-    let request = match Request::parse(&line) {
+    let request = match Request::parse(line) {
         Ok(req) => req,
         Err(e) => return Response::Error { msg: e },
     };

netlify.toml

@@ -0,0 +1,16 @@
+[build]
+  command = """
+    rustup install nightly --profile minimal && cargo doc --no-deps --all-features
+    """
+  publish = "target/doc"
+
+[build.environment]
+  RUSTDOCFLAGS="""
+    --cfg docsrs \
+    --cfg tokio_unstable \
+    """
+  RUSTFLAGS="--cfg tokio_unstable --cfg docsrs"
+
+[[redirects]]
+  from = "/"
+  to = "/tokio"

tests-build/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "tests-build"
+version = "0.1.0"
+authors = ["Tokio Contributors <[email protected]>"]
+edition = "2018"
+publish = false
+
+[features]
+full = ["tokio/full"]
+rt = ["tokio/rt", "tokio/macros"]
+
+[dependencies]
+tokio = { path = "../tokio", optional = true }
+
+[dev-dependencies]
+trybuild = "1.0"

tests-build/README.md

@@ -0,0 +1,2 @@
+Tests the various combination of feature flags. This is broken out to a separate
+crate to work around limitations with cargo features.

tests-build/src/lib.rs

@@ -0,0 +1,2 @@
+#[cfg(feature = "tokio")]
+pub use tokio;

tests-build/tests/fail/macros_core_no_default.rs

@@ -0,0 +1,6 @@
+use tests_build::tokio;
+
+#[tokio::main]
+async fn my_fn() {}
+
+fn main() {}

tests-build/tests/fail/macros_core_no_default.stderr

@@ -0,0 +1,7 @@
+error: The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled.
+ --> tests/fail/macros_core_no_default.rs:3:1
+  |
+3 | #[tokio::main]
+  | ^^^^^^^^^^^^^^
+  |
+  = note: this error originates in the attribute macro `tokio::main` (in Nightly builds, run with -Z macro-backtrace for more info)

tests-build/tests/fail/macros_dead_code.rs

@@ -0,0 +1,8 @@
+#![deny(dead_code)]
+
+use tests_build::tokio;
+
+#[tokio::main]
+async fn f() {}
+
+fn main() {}

tests-build/tests/fail/macros_dead_code.stderr

@@ -0,0 +1,11 @@
+error: function is never used: `f`
+ --> $DIR/macros_dead_code.rs:6:10
+  |
+6 | async fn f() {}
+  |          ^
+  |
+note: the lint level is defined here
+ --> $DIR/macros_dead_code.rs:1:9
+  |
+1 | #![deny(dead_code)]
+  |         ^^^^^^^^^

tests-build/tests/fail/macros_invalid_input.rs

@@ -0,0 +1,51 @@
+#![deny(duplicate_macro_attributes)]
+
+use tests_build::tokio;
+
+#[tokio::main]
+fn main_is_not_async() {}
+
+#[tokio::main(foo)]
+async fn main_attr_has_unknown_args() {}
+
+#[tokio::main(threadpool::bar)]
+async fn main_attr_has_path_args() {}
+
+#[tokio::test]
+fn test_is_not_async() {}
+
+#[tokio::test(foo)]
+async fn test_attr_has_args() {}
+
+#[tokio::test(foo = 123)]
+async fn test_unexpected_attr() {}
+
+#[tokio::test(flavor = 123)]
+async fn test_flavor_not_string() {}
+
+#[tokio::test(flavor = "foo")]
+async fn test_unknown_flavor() {}
+
+#[tokio::test(flavor = "multi_thread", start_paused = false)]
+async fn test_multi_thread_with_start_paused() {}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = "foo")]
+async fn test_worker_threads_not_int() {}
+
+#[tokio::test(flavor = "current_thread", worker_threads = 4)]
+async fn test_worker_threads_and_current_thread() {}
+
+#[tokio::test(crate = 456)]
+async fn test_crate_not_ident_int() {}
+
+#[tokio::test(crate = "456")]
+async fn test_crate_not_ident_invalid() {}
+
+#[tokio::test(crate = "abc::edf")]
+async fn test_crate_not_ident_path() {}
+
+#[tokio::test]
+#[test]
+async fn test_has_second_test_attr() {}
+
+fn main() {}

tests-build/tests/fail/macros_invalid_input.stderr

@@ -0,0 +1,101 @@
+error: the `async` keyword is missing from the function declaration
+ --> $DIR/macros_invalid_input.rs:6:1
+  |
+6 | fn main_is_not_async() {}
+  | ^^
+
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
+ --> $DIR/macros_invalid_input.rs:8:15
+  |
+8 | #[tokio::main(foo)]
+  |               ^^^
+
+error: Must have specified ident
+  --> $DIR/macros_invalid_input.rs:11:15
+   |
+11 | #[tokio::main(threadpool::bar)]
+   |               ^^^^^^^^^^^^^^^
+
+error: the `async` keyword is missing from the function declaration
+  --> $DIR/macros_invalid_input.rs:15:1
+   |
+15 | fn test_is_not_async() {}
+   | ^^
+
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
+  --> $DIR/macros_invalid_input.rs:17:15
+   |
+17 | #[tokio::test(foo)]
+   |               ^^^
+
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
+  --> $DIR/macros_invalid_input.rs:20:15
+   |
+20 | #[tokio::test(foo = 123)]
+   |               ^^^^^^^^^
+
+error: Failed to parse value of `flavor` as string.
+  --> $DIR/macros_invalid_input.rs:23:24
+   |
+23 | #[tokio::test(flavor = 123)]
+   |                        ^^^
+
+error: No such runtime flavor `foo`. The runtime flavors are `current_thread` and `multi_thread`.
+  --> $DIR/macros_invalid_input.rs:26:24
+   |
+26 | #[tokio::test(flavor = "foo")]
+   |                        ^^^^^
+
+error: The `start_paused` option requires the `current_thread` runtime flavor. Use `#[tokio::test(flavor = "current_thread")]`
+  --> $DIR/macros_invalid_input.rs:29:55
+   |
+29 | #[tokio::test(flavor = "multi_thread", start_paused = false)]
+   |                                                       ^^^^^
+
+error: Failed to parse value of `worker_threads` as integer.
+  --> $DIR/macros_invalid_input.rs:32:57
+   |
+32 | #[tokio::test(flavor = "multi_thread", worker_threads = "foo")]
+   |                                                         ^^^^^
+
+error: The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[tokio::test(flavor = "multi_thread")]`
+  --> $DIR/macros_invalid_input.rs:35:59
+   |
+35 | #[tokio::test(flavor = "current_thread", worker_threads = 4)]
+   |                                                           ^
+
+error: Failed to parse value of `crate` as ident.
+  --> $DIR/macros_invalid_input.rs:38:23
+   |
+38 | #[tokio::test(crate = 456)]
+   |                       ^^^
+
+error: Failed to parse value of `crate` as ident: "456"
+  --> $DIR/macros_invalid_input.rs:41:23
+   |
+41 | #[tokio::test(crate = "456")]
+   |                       ^^^^^
+
+error: Failed to parse value of `crate` as ident: "abc::edf"
+  --> $DIR/macros_invalid_input.rs:44:23
+   |
+44 | #[tokio::test(crate = "abc::edf")]
+   |                       ^^^^^^^^^^
+
+error: second test attribute is supplied
+  --> $DIR/macros_invalid_input.rs:48:1
+   |
+48 | #[test]
+   | ^^^^^^^
+
+error: duplicated attribute
+  --> $DIR/macros_invalid_input.rs:48:1
+   |
+48 | #[test]
+   | ^^^^^^^
+   |
+note: the lint level is defined here
+  --> $DIR/macros_invalid_input.rs:1:9
+   |
+1  | #![deny(duplicate_macro_attributes)]
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^

tests-build/tests/fail/macros_type_mismatch.rs

@@ -0,0 +1,35 @@
+use tests_build::tokio;
+
+#[tokio::main]
+async fn missing_semicolon_or_return_type() {
+    Ok(())
+}
+
+#[tokio::main]
+async fn missing_return_type() {
+    return Ok(());
+}
+
+#[tokio::main]
+async fn extra_semicolon() -> Result<(), ()> {
+    /* TODO(taiki-e): help message still wrong
+    help: try using a variant of the expected enum
+       |
+    23 |     Ok(Ok(());)
+       |
+    23 |     Err(Ok(());)
+       |
+    */
+    Ok(());
+}
+
+// https://github.com/tokio-rs/tokio/issues/4635
+#[allow(redundant_semicolons)]
+#[rustfmt::skip]
+#[tokio::main]
+async fn issue_4635() {
+    return 1;
+    ;
+}
+
+fn main() {}

tests-build/tests/fail/macros_type_mismatch.stderr

@@ -0,0 +1,47 @@
+error[E0308]: mismatched types
+ --> $DIR/macros_type_mismatch.rs:5:5
+  |
+4 | async fn missing_semicolon_or_return_type() {
+  |                                             - help: a return type might be missing here: `-> _`
+5 |     Ok(())
+  |     ^^^^^^ expected `()`, found enum `Result`
+  |
+  = note: expected unit type `()`
+                  found enum `Result<(), _>`
+
+error[E0308]: mismatched types
+  --> $DIR/macros_type_mismatch.rs:10:5
+   |
+9  | async fn missing_return_type() {
+   |                                - help: a return type might be missing here: `-> _`
+10 |     return Ok(());
+   |     ^^^^^^^^^^^^^^ expected `()`, found enum `Result`
+   |
+   = note: expected unit type `()`
+                   found enum `Result<(), _>`
+
+error[E0308]: mismatched types
+  --> $DIR/macros_type_mismatch.rs:23:5
+   |
+14 | async fn extra_semicolon() -> Result<(), ()> {
+   |                               -------------- expected `Result<(), ()>` because of return type
+...
+23 |     Ok(());
+   |     ^^^^^^^ expected enum `Result`, found `()`
+   |
+   = note:   expected enum `Result<(), ()>`
+           found unit type `()`
+help: try adding an expression at the end of the block
+   |
+23 ~     Ok(());;
+24 +     Ok(())
+   |
+
+error[E0308]: mismatched types
+  --> $DIR/macros_type_mismatch.rs:32:5
+   |
+30 | async fn issue_4635() {
+   |                       - help: try adding a return type: `-> i32`
+31 |     return 1;
+32 |     ;
+   |     ^ expected `()`, found integer

tests-build/tests/macros.rs

@@ -0,0 +1,27 @@
+#[test]
+fn compile_fail_full() {
+    let t = trybuild::TestCases::new();
+
+    #[cfg(feature = "full")]
+    t.pass("tests/pass/forward_args_and_output.rs");
+
+    #[cfg(feature = "full")]
+    t.pass("tests/pass/macros_main_return.rs");
+
+    #[cfg(feature = "full")]
+    t.pass("tests/pass/macros_main_loop.rs");
+
+    #[cfg(feature = "full")]
+    t.compile_fail("tests/fail/macros_invalid_input.rs");
+
+    #[cfg(feature = "full")]
+    t.compile_fail("tests/fail/macros_dead_code.rs");
+
+    #[cfg(feature = "full")]
+    t.compile_fail("tests/fail/macros_type_mismatch.rs");
+
+    #[cfg(all(feature = "rt", not(feature = "full")))]
+    t.compile_fail("tests/fail/macros_core_no_default.rs");
+
+    drop(t);
+}

tests-build/tests/macros_clippy.rs

@@ -0,0 +1,7 @@
+#[cfg(feature = "full")]
+#[tokio::test]
+async fn test_with_semicolon_without_return_type() {
+    #![deny(clippy::semicolon_if_nothing_returned)]
+
+    dbg!(0);
+}

tests-build/tests/pass/forward_args_and_output.rs

@@ -0,0 +1,13 @@
+use tests_build::tokio;
+
+fn main() {}
+
+// arguments and output type is forwarded so other macros can access them
+
+#[tokio::test]
+async fn test_fn_has_args(_x: u8) {}
+
+#[tokio::test]
+async fn test_has_output() -> Result<(), Box<dyn std::error::Error>> {
+    Ok(())
+}

tests-build/tests/pass/macros_main_loop.rs

@@ -0,0 +1,14 @@
+use tests_build::tokio;
+
+#[tokio::main]
+async fn main() -> Result<(), ()> {
+    loop {
+        if !never() {
+            return Ok(());
+        }
+    }
+}
+
+fn never() -> bool {
+    std::time::Instant::now() > std::time::Instant::now()
+}

tests-build/tests/pass/macros_main_return.rs

@@ -0,0 +1,6 @@
+use tests_build::tokio;
+
+#[tokio::main]
+async fn main() -> Result<(), ()> {
+    return Ok(());
+}

tests-integration/Cargo.toml

@@ -7,6 +7,7 @@ publish = false
 
 [[bin]]
 name = "test-cat"
+required-features = ["rt-process-io-util"]
 
 [[bin]]
 name = "test-mem"
@@ -16,11 +17,30 @@ required-features = ["rt-net"]
 name = "test-process-signal"
 required-features = ["rt-process-signal"]
 
+[[test]]
+name = "macros_main"
+
+[[test]]
+name = "macros_pin"
+
+[[test]]
+name = "macros_select"
+
+[[test]]
+name = "rt_yield"
+required-features = ["rt", "macros", "sync"]
+
 [features]
+rt-process-io-util = ["tokio/rt", "tokio/macros", "tokio/process", "tokio/io-util", "tokio/io-std"]
 # For mem check
 rt-net = ["tokio/rt", "tokio/rt-multi-thread", "tokio/net"]
 # For test-process-signal
-rt-process-signal = ["rt", "tokio/process", "tokio/signal"]
+rt-process-signal = ["rt-net", "tokio/process", "tokio/signal"]
+# For testing wasi + rt/macros/sync features
+#
+# This is an explicit feature so we can use `cargo hack` testing single features
+# instead of all possible permutations.
+wasi-rt = ["rt", "macros", "sync"]
 
 full = [
     "macros",

tests-integration/src/bin/test-cat.rs

@@ -1,20 +1,14 @@
 //! A cat-like utility that can be used as a subprocess to test I/O
 //! stream communication.
 
-use std::io;
-use std::io::Write;
+use tokio::io::AsyncWriteExt;
 
-fn main() {
-    let stdin = io::stdin();
-    let mut stdout = io::stdout();
-    let mut line = String::new();
-    loop {
-        line.clear();
-        stdin.read_line(&mut line).unwrap();
-        if line.is_empty() {
-            break;
-        }
-        stdout.write_all(line.as_bytes()).unwrap();
-    }
-    stdout.flush().unwrap();
+#[tokio::main(flavor = "current_thread")]
+async fn main() {
+    let mut stdin = tokio::io::stdin();
+    let mut stdout = tokio::io::stdout();
+
+    tokio::io::copy(&mut stdin, &mut stdout).await.unwrap();
+
+    stdout.flush().await.unwrap();
 }

tests-integration/tests/macros_main.rs

@@ -1,4 +1,4 @@
-#![cfg(all(feature = "macros", feature = "rt"))]
+#![cfg(all(feature = "macros", feature = "rt-multi-thread"))]
 
 #[tokio::main]
 async fn basic_main() -> usize {

tests-integration/tests/process_stdio.rs

@@ -8,12 +8,22 @@ use tokio_test::assert_ok;
 
 use futures::future::{self, FutureExt};
 use std::convert::TryInto;
-use std::env;
 use std::io;
 use std::process::{ExitStatus, Stdio};
 
+// so, we need to change this back as a test, but for now this doesn't work because of:
+// https://github.com/rust-lang/rust/pull/95469
+//
+// undo when this is closed: https://github.com/tokio-rs/tokio/issues/4802
+
+// fn cat() -> Command {
+//     let mut cmd = Command::new(std::env!("CARGO_BIN_EXE_test-cat"));
+//     cmd.stdin(Stdio::piped()).stdout(Stdio::piped());
+//     cmd
+// }
+
 fn cat() -> Command {
-    let mut cmd = Command::new(env!("CARGO_BIN_EXE_test-cat"));
+    let mut cmd = Command::new("cat");
     cmd.stdin(Stdio::piped()).stdout(Stdio::piped());
     cmd
 }

tests-integration/tests/rt_yield.rs

@@ -0,0 +1,41 @@
+use tokio::sync::oneshot;
+use tokio::task;
+
+async fn spawn_send() {
+    let (tx, rx) = oneshot::channel();
+
+    let task = tokio::spawn(async {
+        for _ in 0..10 {
+            task::yield_now().await;
+        }
+
+        tx.send("done").unwrap();
+    });
+
+    assert_eq!("done", rx.await.unwrap());
+    task.await.unwrap();
+}
+
+#[tokio::main(flavor = "current_thread")]
+async fn entry_point() {
+    spawn_send().await;
+}
+
+#[tokio::test]
+async fn test_macro() {
+    spawn_send().await;
+}
+
+#[test]
+fn main_macro() {
+    entry_point();
+}
+
+#[test]
+fn manual_rt() {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap();
+
+    rt.block_on(async { spawn_send().await });
+}

tokio-macros/CHANGELOG.md

@@ -1,3 +1,23 @@
+# 1.8.0 (June 4th, 2022)
+
+- macros: always emit return statement ([#4636])
+- macros: support setting a custom crate name for `#[tokio::main]` and `#[tokio::test]` ([#4613])
+
+[#4613]: https://github.com/tokio-rs/tokio/pull/4613
+[#4636]: https://github.com/tokio-rs/tokio/pull/4636
+
+# 1.7.0 (December 15th, 2021)
+
+- macros: address remaining `clippy::semicolon_if_nothing_returned` warning ([#4252])
+
+[#4252]: https://github.com/tokio-rs/tokio/pull/4252
+
+# 1.6.0 (November 16th, 2021)
+
+- macros: fix mut patterns in `select!` macro ([#4211])
+
+[#4211]: https://github.com/tokio-rs/tokio/pull/4211
+
 # 1.5.1 (October 29th, 2021)
 
 - macros: fix type resolution error in `#[tokio::main]` ([#4176])

tokio-macros/Cargo.toml

@@ -2,17 +2,15 @@
 name = "tokio-macros"
 # When releasing to crates.io:
 # - Remove path dependencies
-# - Update doc url
-#   - Cargo.toml
 # - Update CHANGELOG.md.
 # - Create "tokio-macros-1.0.x" git tag.
-version = "1.5.1"
+version = "1.8.0"
 edition = "2018"
+rust-version = "1.49"
 authors = ["Tokio Contributors <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/tokio-rs/tokio"
 homepage = "https://tokio.rs"
-documentation = "https://docs.rs/tokio-macros/1.5.1/tokio_macros"
 description = """
 Tokio's proc macros.
 """

tokio-macros/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2021 Tokio Contributors
+Copyright (c) 2022 Tokio Contributors
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

tokio-macros/src/entry.rs

@@ -1,5 +1,5 @@
 use proc_macro::TokenStream;
-use proc_macro2::Span;
+use proc_macro2::{Ident, Span};
 use quote::{quote, quote_spanned, ToTokens};
 use syn::parse::Parser;
 
@@ -29,6 +29,7 @@ struct FinalConfig {
     flavor: RuntimeFlavor,
     worker_threads: Option<usize>,
     start_paused: Option<bool>,
+    crate_name: Option<String>,
 }
 
 /// Config used in case of the attribute not being able to build a valid config
@@ -36,6 +37,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
     flavor: RuntimeFlavor::CurrentThread,
     worker_threads: None,
     start_paused: None,
+    crate_name: None,
 };
 
 struct Configuration {
@@ -45,6 +47,7 @@ struct Configuration {
     worker_threads: Option<(usize, Span)>,
     start_paused: Option<(bool, Span)>,
     is_test: bool,
+    crate_name: Option<String>,
 }
 
 impl Configuration {
@@ -59,6 +62,7 @@ impl Configuration {
             worker_threads: None,
             start_paused: None,
             is_test,
+            crate_name: None,
         }
     }
 
@@ -104,6 +108,15 @@ impl Configuration {
         Ok(())
     }
 
+    fn set_crate_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> {
+        if self.crate_name.is_some() {
+            return Err(syn::Error::new(span, "`crate` set multiple times."));
+        }
+        let name_ident = parse_ident(name, span, "crate")?;
+        self.crate_name = Some(name_ident.to_string());
+        Ok(())
+    }
+
     fn macro_name(&self) -> &'static str {
         if self.is_test {
             "tokio::test"
@@ -151,6 +164,7 @@ impl Configuration {
         };
 
         Ok(FinalConfig {
+            crate_name: self.crate_name.clone(),
             flavor,
             worker_threads,
             start_paused,
@@ -185,6 +199,27 @@ fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::E
     }
 }
 
+fn parse_ident(lit: syn::Lit, span: Span, field: &str) -> Result<Ident, syn::Error> {
+    match lit {
+        syn::Lit::Str(s) => {
+            let err = syn::Error::new(
+                span,
+                format!(
+                    "Failed to parse value of `{}` as ident: \"{}\"",
+                    field,
+                    s.value()
+                ),
+            );
+            let path = s.parse::<syn::Path>().map_err(|_| err.clone())?;
+            path.get_ident().cloned().ok_or(err)
+        }
+        _ => Err(syn::Error::new(
+            span,
+            format!("Failed to parse value of `{}` as ident.", field),
+        )),
+    }
+}
+
 fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
     match bool {
         syn::Lit::Bool(b) => Ok(b.value),
@@ -243,9 +278,15 @@ fn build_config(
                         let msg = "Attribute `core_threads` is renamed to `worker_threads`";
                         return Err(syn::Error::new_spanned(namevalue, msg));
                     }
+                    "crate" => {
+                        config.set_crate_name(
+                            namevalue.lit.clone(),
+                            syn::spanned::Spanned::span(&namevalue.lit),
+                        )?;
+                    }
                     name => {
                         let msg = format!(
-                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`",
+                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`",
                             name,
                         );
                         return Err(syn::Error::new_spanned(namevalue, msg));
@@ -275,7 +316,7 @@ fn build_config(
                         format!("The `{}` attribute requires an argument.", name)
                     }
                     name => {
-                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", name)
+                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`", name)
                     }
                 };
                 return Err(syn::Error::new_spanned(path, msg));
@@ -313,12 +354,16 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To
         (start, end)
     };
 
+    let crate_name = config.crate_name.as_deref().unwrap_or("tokio");
+
+    let crate_ident = Ident::new(crate_name, last_stmt_start_span);
+
     let mut rt = match config.flavor {
         RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
-            tokio::runtime::Builder::new_current_thread()
+            #crate_ident::runtime::Builder::new_current_thread()
         },
         RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
-            tokio::runtime::Builder::new_multi_thread()
+            #crate_ident::runtime::Builder::new_multi_thread()
         },
     };
     if let Some(v) = config.worker_threads {
@@ -338,29 +383,17 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To
 
     let body = &input.block;
     let brace_token = input.block.brace_token;
-    let (tail_return, tail_semicolon) = match body.stmts.last() {
-        Some(syn::Stmt::Semi(expr, _)) => match expr {
-            syn::Expr::Return(_) => (quote! { return }, quote! { ; }),
-            _ => match &input.sig.output {
-                syn::ReturnType::Type(_, ty) if matches!(&**ty, syn::Type::Tuple(ty) if ty.elems.is_empty()) =>
-                {
-                    (quote! {}, quote! { ; }) // unit
-                }
-                syn::ReturnType::Default => (quote! {}, quote! { ; }), // unit
-                syn::ReturnType::Type(..) => (quote! {}, quote! {}),   // ! or another
-            },
-        },
-        _ => (quote! {}, quote! {}),
-    };
     input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
         {
             let body = async #body;
-            #[allow(clippy::expect_used)]
-            #tail_return #rt
-                .enable_all()
-                .build()
-                .expect("Failed building the Runtime")
-                .block_on(body)#tail_semicolon
+            #[allow(clippy::expect_used, clippy::diverging_sub_expression)]
+            {
+                return #rt
+                    .enable_all()
+                    .build()
+                    .expect("Failed building the Runtime")
+                    .block_on(body);
+            }
         }
     })
     .expect("Parsing failure");

tokio-macros/src/lib.rs

@@ -5,7 +5,6 @@
     rust_2018_idioms,
     unreachable_pub
 )]
-#![cfg_attr(docsrs, deny(rustdoc::broken_intra_doc_links))]
 #![doc(test(
     no_crate_inject,
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables))
@@ -169,12 +168,32 @@ use proc_macro::TokenStream;
 ///
 /// Note that `start_paused` requires the `test-util` feature to be enabled.
 ///
-/// ### NOTE:
+/// ### Rename package
 ///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::main(crate = "tokio1")]
+/// async fn main() {
+///     println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// fn main() {
+///     tokio1::runtime::Builder::new_multi_thread()
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             println!("Hello world");
+///         })
+/// }
+/// ```
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
@@ -214,23 +233,52 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
-/// ### NOTE:
+/// ### Rename package
 ///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::main(crate = "tokio1")]
+/// async fn main() {
+///     println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// fn main() {
+///     tokio1::runtime::Builder::new_multi_thread()
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             println!("Hello world");
+///         })
+/// }
+/// ```
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::main(args, item, false)
 }
 
-/// Marks async function to be executed by runtime, suitable to test environment
+/// Marks async function to be executed by runtime, suitable to test environment.
+/// This macro helps set up a `Runtime` without requiring the user to use
+/// [Runtime](../tokio/runtime/struct.Runtime.html) or
+/// [Builder](../tokio/runtime/struct.Builder.html) directly.
 ///
-/// ## Usage
+/// Note: This macro is designed to be simplistic and targets applications that
+/// do not require a complex setup. If the provided functionality is not
+/// sufficient, you may be interested in using
+/// [Builder](../tokio/runtime/struct.Builder.html), which provides a more
+/// powerful interface.
 ///
-/// ### Multi-thread runtime
+/// # Multi-threaded runtime
+///
+/// To use the multi-threaded runtime, the macro can be configured using
 ///
 /// ```no_run
 /// #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
@@ -239,9 +287,17 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
-/// ### Using default
+/// The `worker_threads` option configures the number of worker threads, and
+/// defaults to the number of cpus on the system. This is the default
+/// flavor.
 ///
-/// The default test runtime is single-threaded.
+/// Note: The multi-threaded runtime requires the `rt-multi-thread` feature
+/// flag.
+///
+/// # Current thread runtime
+///
+/// The default test runtime is single-threaded. Each test gets a
+/// separate current-thread runtime.
 ///
 /// ```no_run
 /// #[tokio::test]
@@ -250,6 +306,81 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
+/// ## Usage
+///
+/// ### Using the multi-thread runtime
+///
+/// ```no_run
+/// #[tokio::test(flavor = "multi_thread")]
+/// async fn my_test() {
+///     assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+///     tokio::runtime::Builder::new_multi_thread()
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             assert!(true);
+///         })
+/// }
+/// ```
+///
+/// ### Using current thread runtime
+///
+/// ```no_run
+/// #[tokio::test]
+/// async fn my_test() {
+///     assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+///     tokio::runtime::Builder::new_current_thread()
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             assert!(true);
+///         })
+/// }
+/// ```
+///
+/// ### Set number of worker threads
+///
+/// ```no_run
+/// #[tokio::test(flavor ="multi_thread", worker_threads = 2)]
+/// async fn my_test() {
+///     assert!(true);
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+///     tokio::runtime::Builder::new_multi_thread()
+///         .worker_threads(2)
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             assert!(true);
+///         })
+/// }
+/// ```
+///
 /// ### Configure the runtime to start with time paused
 ///
 /// ```no_run
@@ -259,14 +390,34 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[test]
+/// fn my_test() {
+///     tokio::runtime::Builder::new_current_thread()
+///         .enable_all()
+///         .start_paused(true)
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             assert!(true);
+///         })
+/// }
+/// ```
+///
 /// Note that `start_paused` requires the `test-util` feature to be enabled.
 ///
-/// ### NOTE:
+/// ### Rename package
 ///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::test(crate = "tokio1")]
+/// async fn my_test() {
+///     println!("Hello world");
+/// }
+/// ```
 #[proc_macro_attribute]
 pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::test(args, item, true)
@@ -282,13 +433,6 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
 ///     assert!(true);
 /// }
 /// ```
-///
-/// ### NOTE:
-///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
 #[proc_macro_attribute]
 pub fn test_rt(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::test(args, item, false)
@@ -329,3 +473,11 @@ pub fn test_fail(_args: TokenStream, _item: TokenStream) -> TokenStream {
 pub fn select_priv_declare_output_enum(input: TokenStream) -> TokenStream {
     select::declare_output_enum(input)
 }
+
+/// Implementation detail of the `select!` macro. This macro is **not** intended
+/// to be used as part of the public API and is permitted to change.
+#[proc_macro]
+#[doc(hidden)]
+pub fn select_priv_clean_pattern(input: TokenStream) -> TokenStream {
+    select::clean_pattern_macro(input)
+}

tokio-macros/src/select.rs

@@ -41,3 +41,70 @@ pub(crate) fn declare_output_enum(input: TokenStream) -> TokenStream {
         pub(super) type Mask = #mask;
     })
 }
+
+pub(crate) fn clean_pattern_macro(input: TokenStream) -> TokenStream {
+    // If this isn't a pattern, we return the token stream as-is. The select!
+    // macro is using it in a location requiring a pattern, so an error will be
+    // emitted there.
+    let mut input: syn::Pat = match syn::parse(input.clone()) {
+        Ok(it) => it,
+        Err(_) => return input,
+    };
+
+    clean_pattern(&mut input);
+    quote::ToTokens::into_token_stream(input).into()
+}
+
+// Removes any occurrences of ref or mut in the provided pattern.
+fn clean_pattern(pat: &mut syn::Pat) {
+    match pat {
+        syn::Pat::Box(_box) => {}
+        syn::Pat::Lit(_literal) => {}
+        syn::Pat::Macro(_macro) => {}
+        syn::Pat::Path(_path) => {}
+        syn::Pat::Range(_range) => {}
+        syn::Pat::Rest(_rest) => {}
+        syn::Pat::Verbatim(_tokens) => {}
+        syn::Pat::Wild(_underscore) => {}
+        syn::Pat::Ident(ident) => {
+            ident.by_ref = None;
+            ident.mutability = None;
+            if let Some((_at, pat)) = &mut ident.subpat {
+                clean_pattern(&mut *pat);
+            }
+        }
+        syn::Pat::Or(or) => {
+            for case in or.cases.iter_mut() {
+                clean_pattern(case);
+            }
+        }
+        syn::Pat::Slice(slice) => {
+            for elem in slice.elems.iter_mut() {
+                clean_pattern(elem);
+            }
+        }
+        syn::Pat::Struct(struct_pat) => {
+            for field in struct_pat.fields.iter_mut() {
+                clean_pattern(&mut field.pat);
+            }
+        }
+        syn::Pat::Tuple(tuple) => {
+            for elem in tuple.elems.iter_mut() {
+                clean_pattern(elem);
+            }
+        }
+        syn::Pat::TupleStruct(tuple) => {
+            for elem in tuple.pat.elems.iter_mut() {
+                clean_pattern(elem);
+            }
+        }
+        syn::Pat::Reference(reference) => {
+            reference.mutability = None;
+            clean_pattern(&mut *reference.pat);
+        }
+        syn::Pat::Type(type_pat) => {
+            clean_pattern(&mut *type_pat.pat);
+        }
+        _ => {}
+    }
+}

tokio-stream/CHANGELOG.md

@@ -1,3 +1,33 @@
+# 0.1.9 (June 4, 2022)
+
+- deps: upgrade `tokio-util` dependency to `0.7.x` ([#3762])
+- stream: add `StreamExt::map_while` ([#4351])
+- stream: add `StreamExt::then` ([#4355])
+- stream: add cancel-safety docs to `StreamExt::next` and `try_next` ([#4715])
+- stream: expose `Elapsed` error ([#4502])
+- stream: expose `Timeout` ([#4601])
+- stream: implement `Extend` for `StreamMap` ([#4272])
+- sync: add `Clone` to `RecvError` types ([#4560])
+
+[#3762]: https://github.com/tokio-rs/tokio/pull/3762
+[#4272]: https://github.com/tokio-rs/tokio/pull/4272
+[#4351]: https://github.com/tokio-rs/tokio/pull/4351
+[#4355]: https://github.com/tokio-rs/tokio/pull/4355
+[#4502]: https://github.com/tokio-rs/tokio/pull/4502
+[#4560]: https://github.com/tokio-rs/tokio/pull/4560
+[#4601]: https://github.com/tokio-rs/tokio/pull/4601
+[#4715]: https://github.com/tokio-rs/tokio/pull/4715
+
+# 0.1.8 (October 29, 2021)
+
+- stream: add `From<Receiver<T>>` impl for receiver streams ([#4080])
+- stream: impl `FromIterator` for `StreamMap` ([#4052])
+- signal: make windows docs for signal module show up on unix builds ([#3770])
+
+[#3770]: https://github.com/tokio-rs/tokio/pull/3770
+[#4052]: https://github.com/tokio-rs/tokio/pull/4052
+[#4080]: https://github.com/tokio-rs/tokio/pull/4080
+
 # 0.1.7 (July 7, 2021)
 
 ### Fixed

tokio-stream/Cargo.toml

@@ -2,17 +2,15 @@
 name = "tokio-stream"
 # When releasing to crates.io:
 # - Remove path dependencies
-# - Update doc url
-#   - Cargo.toml
 # - Update CHANGELOG.md.
 # - Create "tokio-stream-0.1.x" git tag.
-version = "0.1.7"
+version = "0.1.9"
 edition = "2018"
+rust-version = "1.49"
 authors = ["Tokio Contributors <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/tokio-rs/tokio"
 homepage = "https://tokio.rs"
-documentation = "https://docs.rs/tokio-stream/0.1.7/tokio_stream"
 description = """
 Utilities to work with `Stream` and `tokio`.
 """
@@ -31,11 +29,12 @@ signal = ["tokio/signal"]
 futures-core = { version = "0.3.0" }
 pin-project-lite = "0.2.0"
 tokio = { version = "1.8.0", path = "../tokio", features = ["sync"] }
-tokio-util = { version = "0.6.3", path = "../tokio-util", optional = true }
+tokio-util = { version = "0.7.0", path = "../tokio-util", optional = true }
 
 [dev-dependencies]
 tokio = { version = "1.2.0", path = "../tokio", features = ["full", "test-util"] }
 async-stream = "0.3"
+parking_lot = "0.12.0"
 tokio-test = { path = "../tokio-test" }
 futures = { version = "0.3", default-features = false }
 

tokio-stream/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2021 Tokio Contributors
+Copyright (c) 2022 Tokio Contributors
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

tokio-stream/src/lib.rs

@@ -10,7 +10,6 @@
     unreachable_pub
 )]
 #![cfg_attr(docsrs, feature(doc_cfg))]
-#![cfg_attr(docsrs, deny(rustdoc::broken_intra_doc_links))]
 #![doc(test(
     no_crate_inject,
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables))
@@ -78,6 +77,9 @@ pub mod wrappers;
 
 mod stream_ext;
 pub use stream_ext::{collect::FromStream, StreamExt};
+cfg_time! {
+    pub use stream_ext::timeout::{Elapsed, Timeout};
+}
 
 mod empty;
 pub use empty::{empty, Empty};

tokio-stream/src/stream_ext.rs

@@ -1,3 +1,4 @@
+use core::future::Future;
 use futures_core::Stream;
 
 mod all;
@@ -27,6 +28,9 @@ use fuse::Fuse;
 mod map;
 use map::Map;
 
+mod map_while;
+use map_while::MapWhile;
+
 mod merge;
 use merge::Merge;
 
@@ -39,21 +43,26 @@ use skip::Skip;
 mod skip_while;
 use skip_while::SkipWhile;
 
-mod try_next;
-use try_next::TryNext;
-
 mod take;
 use take::Take;
 
 mod take_while;
 use take_while::TakeWhile;
 
+mod then;
+use then::Then;
+
+mod try_next;
+use try_next::TryNext;
+
 cfg_time! {
-    mod timeout;
+    pub(crate) mod timeout;
     use timeout::Timeout;
     use tokio::time::Duration;
     mod throttle;
     use throttle::{throttle, Throttle};
+    mod chunks_timeout;
+    use chunks_timeout::ChunksTimeout;
 }
 
 /// An extension trait for the [`Stream`] trait that provides a variety of
@@ -106,6 +115,12 @@ pub trait StreamExt: Stream {
     /// pinning it to the stack using the `pin_mut!` macro from the `pin_utils`
     /// crate.
     ///
+    /// # Cancel safety
+    ///
+    /// This method is cancel safe. The returned future only
+    /// holds onto a reference to the underlying stream,
+    /// so dropping it will never lose a value.
+    ///
     /// # Examples
     ///
     /// ```
@@ -142,6 +157,12 @@ pub trait StreamExt: Stream {
     /// an [`Option<Result<T, E>>`](Option), making for easy use
     /// with the [`?`](std::ops::Try) operator.
     ///
+    /// # Cancel safety
+    ///
+    /// This method is cancel safe. The returned future only
+    /// holds onto a reference to the underlying stream,
+    /// so dropping it will never lose a value.
+    ///
     /// # Examples
     ///
     /// ```
@@ -197,6 +218,93 @@ pub trait StreamExt: Stream {
         Map::new(self, f)
     }
 
+    /// Map this stream's items to a different type for as long as determined by
+    /// the provided closure. A stream of the target type will be returned,
+    /// which will yield elements until the closure returns `None`.
+    ///
+    /// The provided closure is executed over all elements of this stream as
+    /// they are made available, until it returns `None`. It is executed inline
+    /// with calls to [`poll_next`](Stream::poll_next). Once `None` is returned,
+    /// the underlying stream will not be polled again.
+    ///
+    /// Note that this function consumes the stream passed into it and returns a
+    /// wrapped version of it, similar to the [`Iterator::map_while`] method in the
+    /// standard library.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// use tokio_stream::{self as stream, StreamExt};
+    ///
+    /// let stream = stream::iter(1..=10);
+    /// let mut stream = stream.map_while(|x| {
+    ///     if x < 4 {
+    ///         Some(x + 3)
+    ///     } else {
+    ///         None
+    ///     }
+    /// });
+    /// assert_eq!(stream.next().await, Some(4));
+    /// assert_eq!(stream.next().await, Some(5));
+    /// assert_eq!(stream.next().await, Some(6));
+    /// assert_eq!(stream.next().await, None);
+    /// # }
+    /// ```
+    fn map_while<T, F>(self, f: F) -> MapWhile<Self, F>
+    where
+        F: FnMut(Self::Item) -> Option<T>,
+        Self: Sized,
+    {
+        MapWhile::new(self, f)
+    }
+
+    /// Maps this stream's items asynchronously to a different type, returning a
+    /// new stream of the resulting type.
+    ///
+    /// The provided closure is executed over all elements of this stream as
+    /// they are made available, and the returned future is executed. Only one
+    /// future is executed at the time.
+    ///
+    /// Note that this function consumes the stream passed into it and returns a
+    /// wrapped version of it, similar to the existing `then` methods in the
+    /// standard library.
+    ///
+    /// Be aware that if the future is not `Unpin`, then neither is the `Stream`
+    /// returned by this method. To handle this, you can use `tokio::pin!` as in
+    /// the example below or put the stream in a `Box` with `Box::pin(stream)`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// use tokio_stream::{self as stream, StreamExt};
+    ///
+    /// async fn do_async_work(value: i32) -> i32 {
+    ///     value + 3
+    /// }
+    ///
+    /// let stream = stream::iter(1..=3);
+    /// let stream = stream.then(do_async_work);
+    ///
+    /// tokio::pin!(stream);
+    ///
+    /// assert_eq!(stream.next().await, Some(4));
+    /// assert_eq!(stream.next().await, Some(5));
+    /// assert_eq!(stream.next().await, Some(6));
+    /// # }
+    /// ```
+    fn then<F, Fut>(self, f: F) -> Then<Self, Fut, F>
+    where
+        F: FnMut(Self::Item) -> Fut,
+        Fut: Future,
+        Self: Sized,
+    {
+        Then::new(self, f)
+    }
+
     /// Combine two streams into one by interleaving the output of both as it
     /// is produced.
     ///
@@ -899,6 +1007,63 @@ pub trait StreamExt: Stream {
     {
         throttle(duration, self)
     }
+
+    /// Batches the items in the given stream using a maximum duration and size for each batch.
+    ///
+    /// This stream returns the next batch of items in the following situations:
+    ///  1. The inner stream has returned at least `max_size` many items since the last batch.
+    ///  2. The time since the first item of a batch is greater than the given duration.
+    ///  3. The end of the stream is reached.
+    ///
+    /// The length of the returned vector is never empty or greater than the maximum size. Empty batches
+    /// will not be emitted if no items are received upstream.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `max_size` is zero
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use std::time::Duration;
+    /// use tokio::time;
+    /// use tokio_stream::{self as stream, StreamExt};
+    /// use futures::FutureExt;
+    ///
+    /// #[tokio::main]
+    /// # async fn _unused() {}
+    /// # #[tokio::main(flavor = "current_thread", start_paused = true)]
+    /// async fn main() {
+    ///     let iter = vec![1, 2, 3, 4].into_iter();
+    ///     let stream0 = stream::iter(iter);
+    ///
+    ///     let iter = vec![5].into_iter();
+    ///     let stream1 = stream::iter(iter)
+    ///          .then(move |n| time::sleep(Duration::from_secs(5)).map(move |_| n));
+    ///
+    ///     let chunk_stream = stream0
+    ///         .chain(stream1)
+    ///         .chunks_timeout(3, Duration::from_secs(2));
+    ///     tokio::pin!(chunk_stream);
+    ///
+    ///     // a full batch was received
+    ///     assert_eq!(chunk_stream.next().await, Some(vec![1,2,3]));
+    ///     // deadline was reached before max_size was reached
+    ///     assert_eq!(chunk_stream.next().await, Some(vec![4]));
+    ///     // last element in the stream
+    ///     assert_eq!(chunk_stream.next().await, Some(vec![5]));
+    /// }
+    /// ```
+    #[cfg(feature = "time")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "time")))]
+    #[track_caller]
+    fn chunks_timeout(self, max_size: usize, duration: Duration) -> ChunksTimeout<Self>
+    where
+        Self: Sized,
+    {
+        assert!(max_size > 0, "`max_size` must be non-zero.");
+        ChunksTimeout::new(self, max_size, duration)
+    }
 }
 
 impl<St: ?Sized> StreamExt for St where St: Stream {}
@@ -906,10 +1071,10 @@ impl<St: ?Sized> StreamExt for St where St: Stream {}
 /// Merge the size hints from two streams.
 fn merge_size_hints(
     (left_low, left_high): (usize, Option<usize>),
-    (right_low, right_hign): (usize, Option<usize>),
+    (right_low, right_high): (usize, Option<usize>),
 ) -> (usize, Option<usize>) {
     let low = left_low.saturating_add(right_low);
-    let high = match (left_high, right_hign) {
+    let high = match (left_high, right_high) {
         (Some(h1), Some(h2)) => h1.checked_add(h2),
         _ => None,
     };

tokio-stream/src/stream_ext/chunks_timeout.rs

@@ -0,0 +1,84 @@
+use crate::stream_ext::Fuse;
+use crate::Stream;
+use tokio::time::{sleep, Instant, Sleep};
+
+use core::future::Future;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+use pin_project_lite::pin_project;
+use std::time::Duration;
+
+pin_project! {
+    /// Stream returned by the [`chunks_timeout`](super::StreamExt::chunks_timeout) method.
+    #[must_use = "streams do nothing unless polled"]
+    #[derive(Debug)]
+    pub struct ChunksTimeout<S: Stream> {
+        #[pin]
+        stream: Fuse<S>,
+        #[pin]
+        deadline: Sleep,
+        duration: Duration,
+        items: Vec<S::Item>,
+        cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475
+    }
+}
+
+impl<S: Stream> ChunksTimeout<S> {
+    pub(super) fn new(stream: S, max_size: usize, duration: Duration) -> Self {
+        ChunksTimeout {
+            stream: Fuse::new(stream),
+            deadline: sleep(duration),
+            duration,
+            items: Vec::with_capacity(max_size),
+            cap: max_size,
+        }
+    }
+}
+
+impl<S: Stream> Stream for ChunksTimeout<S> {
+    type Item = Vec<S::Item>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut me = self.as_mut().project();
+        loop {
+            match me.stream.as_mut().poll_next(cx) {
+                Poll::Pending => break,
+                Poll::Ready(Some(item)) => {
+                    if me.items.is_empty() {
+                        me.deadline.as_mut().reset(Instant::now() + *me.duration);
+                        me.items.reserve_exact(*me.cap);
+                    }
+                    me.items.push(item);
+                    if me.items.len() >= *me.cap {
+                        return Poll::Ready(Some(std::mem::take(me.items)));
+                    }
+                }
+                Poll::Ready(None) => {
+                    // Returning Some here is only correct because we fuse the inner stream.
+                    let last = if me.items.is_empty() {
+                        None
+                    } else {
+                        Some(std::mem::take(me.items))
+                    };
+
+                    return Poll::Ready(last);
+                }
+            }
+        }
+
+        if !me.items.is_empty() {
+            ready!(me.deadline.poll(cx));
+            return Poll::Ready(Some(std::mem::take(me.items)));
+        }
+
+        Poll::Pending
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let chunk_len = if self.items.is_empty() { 0 } else { 1 };
+        let (lower, upper) = self.stream.size_hint();
+        let lower = (lower / self.cap).saturating_add(chunk_len);
+        let upper = upper.and_then(|x| x.checked_add(chunk_len));
+        (lower, upper)
+    }
+}

tokio-stream/src/stream_ext/collect.rs

@@ -66,17 +66,17 @@ where
         use Poll::Ready;
 
         loop {
-            let mut me = self.as_mut().project();
+            let me = self.as_mut().project();
 
             let item = match ready!(me.stream.poll_next(cx)) {
                 Some(item) => item,
                 None => {
-                    return Ready(U::finalize(sealed::Internal, &mut me.collection));
+                    return Ready(U::finalize(sealed::Internal, me.collection));
                 }
             };
 
-            if !U::extend(sealed::Internal, &mut me.collection, item) {
-                return Ready(U::finalize(sealed::Internal, &mut me.collection));
+            if !U::extend(sealed::Internal, me.collection, item) {
+                return Ready(U::finalize(sealed::Internal, me.collection));
             }
         }
     }

tokio-stream/src/stream_ext/map_while.rs

@@ -0,0 +1,52 @@
+use crate::Stream;
+
+use core::fmt;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// Stream for the [`map_while`](super::StreamExt::map_while) method.
+    #[must_use = "streams do nothing unless polled"]
+    pub struct MapWhile<St, F> {
+        #[pin]
+        stream: St,
+        f: F,
+    }
+}
+
+impl<St, F> fmt::Debug for MapWhile<St, F>
+where
+    St: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MapWhile")
+            .field("stream", &self.stream)
+            .finish()
+    }
+}
+
+impl<St, F> MapWhile<St, F> {
+    pub(super) fn new(stream: St, f: F) -> Self {
+        MapWhile { stream, f }
+    }
+}
+
+impl<St, F, T> Stream for MapWhile<St, F>
+where
+    St: Stream,
+    F: FnMut(St::Item) -> Option<T>,
+{
+    type Item = T;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
+        let me = self.project();
+        let f = me.f;
+        me.stream.poll_next(cx).map(|opt| opt.and_then(f))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (_, upper) = self.stream.size_hint();
+        (0, upper)
+    }
+}

tokio-stream/src/stream_ext/next.rs

@@ -8,6 +8,13 @@ use pin_project_lite::pin_project;
 
 pin_project! {
     /// Future for the [`next`](super::StreamExt::next) method.
+    ///
+    /// # Cancel safety
+    ///
+    /// This method is cancel safe. It only
+    /// holds onto a reference to the underlying stream,
+    /// so dropping it will never lose a value.
+    ///
     #[derive(Debug)]
     #[must_use = "futures do nothing unless you `.await` or poll them"]
     pub struct Next<'a, St: ?Sized> {

tokio-stream/src/stream_ext/then.rs

@@ -0,0 +1,83 @@
+use crate::Stream;
+
+use core::fmt;
+use core::future::Future;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// Stream for the [`then`](super::StreamExt::then) method.
+    #[must_use = "streams do nothing unless polled"]
+    pub struct Then<St, Fut, F> {
+        #[pin]
+        stream: St,
+        #[pin]
+        future: Option<Fut>,
+        f: F,
+    }
+}
+
+impl<St, Fut, F> fmt::Debug for Then<St, Fut, F>
+where
+    St: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Then")
+            .field("stream", &self.stream)
+            .finish()
+    }
+}
+
+impl<St, Fut, F> Then<St, Fut, F> {
+    pub(super) fn new(stream: St, f: F) -> Self {
+        Then {
+            stream,
+            future: None,
+            f,
+        }
+    }
+}
+
+impl<St, F, Fut> Stream for Then<St, Fut, F>
+where
+    St: Stream,
+    Fut: Future,
+    F: FnMut(St::Item) -> Fut,
+{
+    type Item = Fut::Output;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Fut::Output>> {
+        let mut me = self.project();
+
+        loop {
+            if let Some(future) = me.future.as_mut().as_pin_mut() {
+                match future.poll(cx) {
+                    Poll::Ready(item) => {
+                        me.future.set(None);
+                        return Poll::Ready(Some(item));
+                    }
+                    Poll::Pending => return Poll::Pending,
+                }
+            }
+
+            match me.stream.as_mut().poll_next(cx) {
+                Poll::Ready(Some(item)) => {
+                    me.future.set(Some((me.f)(item)));
+                }
+                Poll::Ready(None) => return Poll::Ready(None),
+                Poll::Pending => return Poll::Pending,
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let future_len = if self.future.is_some() { 1 } else { 0 };
+        let (lower, upper) = self.stream.size_hint();
+
+        let lower = lower.saturating_add(future_len);
+        let upper = upper.and_then(|upper| upper.checked_add(future_len));
+
+        (lower, upper)
+    }
+}

tokio-stream/src/stream_ext/try_next.rs

@@ -9,6 +9,12 @@ use pin_project_lite::pin_project;
 
 pin_project! {
     /// Future for the [`try_next`](super::StreamExt::try_next) method.
+    ///
+    /// # Cancel safety
+    ///
+    /// This method is cancel safe. It only
+    /// holds onto a reference to the underlying stream,
+    /// so dropping it will never lose a value.
     #[derive(Debug)]
     #[must_use = "futures do nothing unless you `.await` or poll them"]
     pub struct TryNext<'a, St: ?Sized> {

tokio-stream/src/stream_map.rs

@@ -585,6 +585,15 @@ where
     }
 }
 
+impl<K, V> Extend<(K, V)> for StreamMap<K, V> {
+    fn extend<T>(&mut self, iter: T)
+    where
+        T: IntoIterator<Item = (K, V)>,
+    {
+        self.entries.extend(iter);
+    }
+}
+
 mod rand {
     use std::cell::Cell;
 

tokio-stream/src/wrappers/broadcast.rs

@@ -14,11 +14,11 @@ use std::task::{Context, Poll};
 /// [`Stream`]: trait@crate::Stream
 #[cfg_attr(docsrs, doc(cfg(feature = "sync")))]
 pub struct BroadcastStream<T> {
-    inner: ReusableBoxFuture<(Result<T, RecvError>, Receiver<T>)>,
+    inner: ReusableBoxFuture<'static, (Result<T, RecvError>, Receiver<T>)>,
 }
 
 /// An error returned from the inner stream of a [`BroadcastStream`].
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone)]
 pub enum BroadcastStreamRecvError {
     /// The receiver lagged too far behind. Attempting to receive again will
     /// return the oldest message still retained by the channel.

tokio-stream/src/wrappers/watch.rs

@@ -49,7 +49,7 @@ use tokio::sync::watch::error::RecvError;
 /// [`Stream`]: trait@crate::Stream
 #[cfg_attr(docsrs, doc(cfg(feature = "sync")))]
 pub struct WatchStream<T> {
-    inner: ReusableBoxFuture<(Result<(), RecvError>, Receiver<T>)>,
+    inner: ReusableBoxFuture<'static, (Result<(), RecvError>, Receiver<T>)>,
 }
 
 async fn make_future<T: Clone + Send + Sync>(

tokio-stream/tests/chunks_timeout.rs

@@ -0,0 +1,84 @@
+#![warn(rust_2018_idioms)]
+#![cfg(all(feature = "time", feature = "sync", feature = "io-util"))]
+
+use tokio::time;
+use tokio_stream::{self as stream, StreamExt};
+use tokio_test::assert_pending;
+use tokio_test::task;
+
+use futures::FutureExt;
+use std::time::Duration;
+
+#[tokio::test(start_paused = true)]
+async fn usage() {
+    let iter = vec![1, 2, 3].into_iter();
+    let stream0 = stream::iter(iter);
+
+    let iter = vec![4].into_iter();
+    let stream1 =
+        stream::iter(iter).then(move |n| time::sleep(Duration::from_secs(3)).map(move |_| n));
+
+    let chunk_stream = stream0
+        .chain(stream1)
+        .chunks_timeout(4, Duration::from_secs(2));
+
+    let mut chunk_stream = task::spawn(chunk_stream);
+
+    assert_pending!(chunk_stream.poll_next());
+    time::advance(Duration::from_secs(2)).await;
+    assert_eq!(chunk_stream.next().await, Some(vec![1, 2, 3]));
+
+    assert_pending!(chunk_stream.poll_next());
+    time::advance(Duration::from_secs(2)).await;
+    assert_eq!(chunk_stream.next().await, Some(vec![4]));
+}
+
+#[tokio::test(start_paused = true)]
+async fn full_chunk_with_timeout() {
+    let iter = vec![1, 2].into_iter();
+    let stream0 = stream::iter(iter);
+
+    let iter = vec![3].into_iter();
+    let stream1 =
+        stream::iter(iter).then(move |n| time::sleep(Duration::from_secs(1)).map(move |_| n));
+
+    let iter = vec![4].into_iter();
+    let stream2 =
+        stream::iter(iter).then(move |n| time::sleep(Duration::from_secs(3)).map(move |_| n));
+
+    let chunk_stream = stream0
+        .chain(stream1)
+        .chain(stream2)
+        .chunks_timeout(3, Duration::from_secs(2));
+
+    let mut chunk_stream = task::spawn(chunk_stream);
+
+    assert_pending!(chunk_stream.poll_next());
+    time::advance(Duration::from_secs(2)).await;
+    assert_eq!(chunk_stream.next().await, Some(vec![1, 2, 3]));
+
+    assert_pending!(chunk_stream.poll_next());
+    time::advance(Duration::from_secs(2)).await;
+    assert_eq!(chunk_stream.next().await, Some(vec![4]));
+}
+
+#[tokio::test]
+#[ignore]
+async fn real_time() {
+    let iter = vec![1, 2, 3, 4].into_iter();
+    let stream0 = stream::iter(iter);
+
+    let iter = vec![5].into_iter();
+    let stream1 =
+        stream::iter(iter).then(move |n| time::sleep(Duration::from_secs(5)).map(move |_| n));
+
+    let chunk_stream = stream0
+        .chain(stream1)
+        .chunks_timeout(3, Duration::from_secs(2));
+
+    let mut chunk_stream = task::spawn(chunk_stream);
+
+    assert_eq!(chunk_stream.next().await, Some(vec![1, 2, 3]));
+    assert_eq!(chunk_stream.next().await, Some(vec![4]));
+    assert_eq!(chunk_stream.next().await, Some(vec![5]));
+}

tokio-stream/tests/stream_panic.rs

@@ -0,0 +1,55 @@
+#![warn(rust_2018_idioms)]
+#![cfg(feature = "time")]
+
+use parking_lot::{const_mutex, Mutex};
+use std::error::Error;
+use std::panic;
+use std::sync::Arc;
+use tokio::time::Duration;
+use tokio_stream::{self as stream, StreamExt};
+
+fn test_panic<Func: FnOnce() + panic::UnwindSafe>(func: Func) -> Option<String> {
+    static PANIC_MUTEX: Mutex<()> = const_mutex(());
+
+    {
+        let _guard = PANIC_MUTEX.lock();
+        let panic_file: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
+
+        let prev_hook = panic::take_hook();
+        {
+            let panic_file = panic_file.clone();
+            panic::set_hook(Box::new(move |panic_info| {
+                let panic_location = panic_info.location().unwrap();
+                panic_file
+                    .lock()
+                    .clone_from(&Some(panic_location.file().to_string()));
+            }));
+        }
+
+        let result = panic::catch_unwind(func);
+        // Return to the previously set panic hook (maybe default) so that we get nice error
+        // messages in the tests.
+        panic::set_hook(prev_hook);
+
+        if result.is_err() {
+            panic_file.lock().clone()
+        } else {
+            None
+        }
+    }
+}
+
+#[test]
+fn stream_chunks_timeout_panic_caller() -> Result<(), Box<dyn Error>> {
+    let panic_location_file = test_panic(|| {
+        let iter = vec![1, 2, 3].into_iter();
+        let stream0 = stream::iter(iter);
+
+        let _chunk_stream = stream0.chunks_timeout(0, Duration::from_secs(2));
+    });
+
+    // The panic location should be in this file
+    assert_eq!(&panic_location_file.unwrap(), file!());
+
+    Ok(())
+}

tokio-stream/tests/stream_timeout.rs

@@ -1,10 +1,10 @@
-#![cfg(feature = "full")]
+#![cfg(all(feature = "time", feature = "sync", feature = "io-util"))]
 
 use tokio::time::{self, sleep, Duration};
 use tokio_stream::{self, StreamExt};
 use tokio_test::*;
 
-use futures::StreamExt as _;
+use futures::stream;
 
 async fn maybe_sleep(idx: i32) -> i32 {
     if idx % 2 == 0 {

tokio-stream/tests/time_throttle.rs

@@ -1,5 +1,5 @@
 #![warn(rust_2018_idioms)]
-#![cfg(feature = "full")]
+#![cfg(all(feature = "time", feature = "sync", feature = "io-util"))]
 
 use tokio::time;
 use tokio_stream::StreamExt;

tokio-test/Cargo.toml

@@ -2,17 +2,15 @@
 name = "tokio-test"
 # When releasing to crates.io:
 # - Remove path dependencies
-# - Update doc url
-#   - Cargo.toml
 # - Update CHANGELOG.md.
 # - Create "tokio-test-0.4.x" git tag.
 version = "0.4.2"
 edition = "2018"
+rust-version = "1.49"
 authors = ["Tokio Contributors <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/tokio-rs/tokio"
 homepage = "https://tokio.rs"
-documentation = "https://docs.rs/tokio-test/0.4.2/tokio_test"
 description = """
 Testing utilities for Tokio- and futures-based code
 """
@@ -20,7 +18,7 @@ categories = ["asynchronous", "testing"]
 
 [dependencies]
 tokio = { version = "1.2.0", path = "../tokio", features = ["rt", "sync", "time", "test-util"] }
-tokio-stream = { version = "0.1", path = "../tokio-stream" }
+tokio-stream = { version = "0.1.1", path = "../tokio-stream" }
 async-stream = "0.3"
 
 bytes = "1.0.0"

tokio-test/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2021 Tokio Contributors
+Copyright (c) 2022 Tokio Contributors
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

tokio-test/src/lib.rs

@@ -4,7 +4,6 @@
     rust_2018_idioms,
     unreachable_pub
 )]
-#![cfg_attr(docsrs, deny(rustdoc::broken_intra_doc_links))]
 #![doc(test(
     no_crate_inject,
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables))

tokio-util/CHANGELOG.md

@@ -1,6 +1,79 @@
+# 0.7.3 (June 4, 2022)
+
+### Changed
+
+- tracing: don't require default tracing features ([#4592])
+- util: simplify implementation of `ReusableBoxFuture` ([#4675])
+
+### Added (unstable)
+
+- task: add `JoinMap` ([#4640], [#4697])
+
+[#4592]: https://github.com/tokio-rs/tokio/pull/4592
+[#4640]: https://github.com/tokio-rs/tokio/pull/4640
+[#4675]: https://github.com/tokio-rs/tokio/pull/4675
+[#4697]: https://github.com/tokio-rs/tokio/pull/4697
+
+# 0.7.2 (May 14, 2022)
+
+This release contains a rewrite of `CancellationToken` that fixes a memory leak. ([#4652])
+
+[#4652]: https://github.com/tokio-rs/tokio/pull/4652
+
+# 0.7.1 (February 21, 2022)
+
+### Added
+
+- codec: add `length_field_type` to `LengthDelimitedCodec` builder ([#4508])
+- io: add `StreamReader::into_inner_with_chunk()` ([#4559])
+
+### Changed
+
+- switch from log to tracing ([#4539])
+
+### Fixed
+
+- sync: fix waker update condition in `CancellationToken` ([#4497])
+- bumped tokio dependency to 1.6 to satisfy minimum requirements ([#4490])
+
+[#4490]: https://github.com/tokio-rs/tokio/pull/4490
+[#4497]: https://github.com/tokio-rs/tokio/pull/4497
+[#4508]: https://github.com/tokio-rs/tokio/pull/4508
+[#4539]: https://github.com/tokio-rs/tokio/pull/4539
+[#4559]: https://github.com/tokio-rs/tokio/pull/4559
+
+# 0.7.0 (February 9, 2022)
+
+### Added
+
+- task: add `spawn_pinned` ([#3370])
+- time: add `shrink_to_fit` and `compact` methods to `DelayQueue` ([#4170])
+- codec: improve `Builder::max_frame_length` docs ([#4352])
+- codec: add mutable reference getters for codecs to pinned `Framed` ([#4372])
+- net: add generic trait to combine `UnixListener` and `TcpListener` ([#4385])
+- codec: implement `Framed::map_codec` ([#4427])
+- codec: implement `Encoder<BytesMut>` for `BytesCodec` ([#4465])
+
+### Changed
+
+- sync: add lifetime parameter to `ReusableBoxFuture` ([#3762])
+- sync: refactored `PollSender<T>` to fix a subtly broken `Sink<T>` implementation ([#4214])
+- time: remove error case from the infallible `DelayQueue::poll_elapsed` ([#4241])
+
+[#3370]: https://github.com/tokio-rs/tokio/pull/3370
+[#4170]: https://github.com/tokio-rs/tokio/pull/4170
+[#4352]: https://github.com/tokio-rs/tokio/pull/4352
+[#4372]: https://github.com/tokio-rs/tokio/pull/4372
+[#4385]: https://github.com/tokio-rs/tokio/pull/4385
+[#4427]: https://github.com/tokio-rs/tokio/pull/4427
+[#4465]: https://github.com/tokio-rs/tokio/pull/4465
+[#3762]: https://github.com/tokio-rs/tokio/pull/3762
+[#4214]: https://github.com/tokio-rs/tokio/pull/4214
+[#4241]: https://github.com/tokio-rs/tokio/pull/4241
+
 # 0.6.10 (May 14, 2021)
 
-This is a backport for a memory leak in `CancellationToken`. ([#4652])
+This is a backport for the memory leak in `CancellationToken` that was originally fixed in 0.7.2. ([#4652])
 
 [#4652]: https://github.com/tokio-rs/tokio/pull/4652
 

tokio-util/Cargo.toml

@@ -2,17 +2,15 @@
 name = "tokio-util"
 # When releasing to crates.io:
 # - Remove path dependencies
-# - Update doc url
-#   - Cargo.toml
 # - Update CHANGELOG.md.
-# - Create "tokio-util-0.6.x" git tag.
-version = "0.6.10"
+# - Create "tokio-util-0.7.x" git tag.
+version = "0.7.3"
 edition = "2018"
+rust-version = "1.49"
 authors = ["Tokio Contributors <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/tokio-rs/tokio"
 homepage = "https://tokio.rs"
-documentation = "https://docs.rs/tokio-util/0.6.9/tokio_util"
 description = """
 Additional utilities for working with Tokio.
 """
@@ -27,25 +25,27 @@ full = ["codec", "compat", "io-util", "time", "net", "rt"]
 
 net = ["tokio/net"]
 compat = ["futures-io",]
-codec = []
+codec = ["tracing"]
 time = ["tokio/time","slab"]
 io = []
 io-util = ["io", "tokio/rt", "tokio/io-util"]
-rt = ["tokio/rt"]
+rt = ["tokio/rt", "tokio/sync", "futures-util", "hashbrown"]
 
 __docs_rs = ["futures-util"]
 
 [dependencies]
-tokio = { version = "1.7.0", path = "../tokio", features = ["sync"] }
-
+tokio = { version = "1.19.0", path = "../tokio", features = ["sync"] }
 bytes = "1.0.0"
 futures-core = "0.3.0"
 futures-sink = "0.3.0"
 futures-io = { version = "0.3.0", optional = true }
 futures-util = { version = "0.3.0", optional = true }
-log = "0.4"
 pin-project-lite = "0.2.0"
-slab = { version = "0.4.1", optional = true } # Backs `DelayQueue`
+slab = { version = "0.4.4", optional = true } # Backs `DelayQueue`
+tracing = { version = "0.1.25", default-features = false, features = ["std"], optional = true }
+
+[target.'cfg(tokio_unstable)'.dependencies]
+hashbrown = { version = "0.12.0", optional = true }
 
 [dev-dependencies]
 tokio = { version = "1.0.0", path = "../tokio", features = ["full"] }
@@ -55,7 +55,12 @@ tokio-stream = { version = "0.1", path = "../tokio-stream" }
 async-stream = "0.3.0"
 futures = "0.3.0"
 futures-test = "0.3.5"
+parking_lot = "0.12.0"
 
 [package.metadata.docs.rs]
 all-features = true
-rustdoc-args = ["--cfg", "docsrs"]
+# enable unstable features in the documentation
+rustdoc-args = ["--cfg", "docsrs", "--cfg", "tokio_unstable"]
+# it's necessary to _also_ pass `--cfg tokio_unstable` to rustc, or else
+# dependencies will not be enabled, and the docs build will fail.
+rustc-args = ["--cfg", "docsrs", "--cfg", "tokio_unstable"]

tokio-util/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2021 Tokio Contributors
+Copyright (c) 2022 Tokio Contributors
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

tokio-util/src/codec/bytes_codec.rs

@@ -74,3 +74,13 @@ impl Encoder<Bytes> for BytesCodec {
         Ok(())
     }
 }
+
+impl Encoder<BytesMut> for BytesCodec {
+    type Error = io::Error;
+
+    fn encode(&mut self, data: BytesMut, buf: &mut BytesMut) -> Result<(), io::Error> {
+        buf.reserve(data.len());
+        buf.put(data);
+        Ok(())
+    }
+}

tokio-util/src/codec/decoder.rs

@@ -20,7 +20,7 @@ use std::io;
 /// it's possible to temporarily read 0 bytes by reaching EOF.
 ///
 /// In these cases `decode_eof` will be called until it signals
-/// fullfillment of all closing frames by returning `Ok(None)`.
+/// fulfillment of all closing frames by returning `Ok(None)`.
 /// After that, repeated attempts to read from the [`Framed`] or [`FramedRead`]
 /// will not invoke `decode` or `decode_eof` again, until data can be read
 /// during a retry.

tokio-util/src/codec/framed.rs

@@ -204,6 +204,35 @@ impl<T, U> Framed<T, U> {
         &mut self.inner.codec
     }
 
+    /// Maps the codec `U` to `C`, preserving the read and write buffers
+    /// wrapped by `Framed`.
+    ///
+    /// Note that care should be taken to not tamper with the underlying codec
+    /// as it may corrupt the stream of frames otherwise being worked with.
+    pub fn map_codec<C, F>(self, map: F) -> Framed<T, C>
+    where
+        F: FnOnce(U) -> C,
+    {
+        // This could be potentially simplified once rust-lang/rust#86555 hits stable
+        let parts = self.into_parts();
+        Framed::from_parts(FramedParts {
+            io: parts.io,
+            codec: map(parts.codec),
+            read_buf: parts.read_buf,
+            write_buf: parts.write_buf,
+            _priv: (),
+        })
+    }
+
+    /// Returns a mutable reference to the underlying codec wrapped by
+    /// `Framed`.
+    ///
+    /// Note that care should be taken to not tamper with the underlying codec
+    /// as it may corrupt the stream of frames otherwise being worked with.
+    pub fn codec_pin_mut(self: Pin<&mut Self>) -> &mut U {
+        self.project().inner.project().codec
+    }
+
     /// Returns a reference to the read buffer.
     pub fn read_buffer(&self) -> &BytesMut {
         &self.inner.state.read.buffer

tokio-util/src/codec/framed_impl.rs

@@ -7,12 +7,12 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use bytes::BytesMut;
 use futures_core::ready;
 use futures_sink::Sink;
-use log::trace;
 use pin_project_lite::pin_project;
 use std::borrow::{Borrow, BorrowMut};
 use std::io;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use tracing::trace;
 
 pin_project! {
     #[derive(Debug)]
@@ -278,7 +278,7 @@ where
 
         while !pinned.state.borrow_mut().buffer.is_empty() {
             let WriteFrame { buffer } = pinned.state.borrow_mut();
-            trace!("writing; remaining={}", buffer.len());
+            trace!(remaining = buffer.len(), "writing;");
 
             let n = ready!(poll_write_buf(pinned.inner.as_mut(), cx, buffer))?;
 

tokio-util/src/codec/framed_read.rs

@@ -108,6 +108,32 @@ impl<T, D> FramedRead<T, D> {
         &mut self.inner.codec
     }
 
+    /// Maps the decoder `D` to `C`, preserving the read buffer
+    /// wrapped by `Framed`.
+    pub fn map_decoder<C, F>(self, map: F) -> FramedRead<T, C>
+    where
+        F: FnOnce(D) -> C,
+    {
+        // This could be potentially simplified once rust-lang/rust#86555 hits stable
+        let FramedImpl {
+            inner,
+            state,
+            codec,
+        } = self.inner;
+        FramedRead {
+            inner: FramedImpl {
+                inner,
+                state,
+                codec: map(codec),
+            },
+        }
+    }
+
+    /// Returns a mutable reference to the underlying decoder.
+    pub fn decoder_pin_mut(self: Pin<&mut Self>) -> &mut D {
+        self.project().inner.project().codec
+    }
+
     /// Returns a reference to the read buffer.
     pub fn read_buffer(&self) -> &BytesMut {
         &self.inner.state.buffer

tokio-util/src/codec/framed_write.rs

@@ -88,6 +88,32 @@ impl<T, E> FramedWrite<T, E> {
         &mut self.inner.codec
     }
 
+    /// Maps the encoder `E` to `C`, preserving the write buffer
+    /// wrapped by `Framed`.
+    pub fn map_encoder<C, F>(self, map: F) -> FramedWrite<T, C>
+    where
+        F: FnOnce(E) -> C,
+    {
+        // This could be potentially simplified once rust-lang/rust#86555 hits stable
+        let FramedImpl {
+            inner,
+            state,
+            codec,
+        } = self.inner;
+        FramedWrite {
+            inner: FramedImpl {
+                inner,
+                state,
+                codec: map(codec),
+            },
+        }
+    }
+
+    /// Returns a mutable reference to the underlying encoder.
+    pub fn encoder_pin_mut(self: Pin<&mut Self>) -> &mut E {
+        self.project().inner.project().codec
+    }
+
     /// Returns a reference to the write buffer.
     pub fn write_buffer(&self) -> &BytesMut {
         &self.inner.state.buffer

tokio-util/src/codec/length_delimited.rs

@@ -84,7 +84,7 @@
 //! # fn bind_read<T: AsyncRead>(io: T) {
 //! LengthDelimitedCodec::builder()
 //!     .length_field_offset(0) // default value
-//!     .length_field_length(2)
+//!     .length_field_type::<u16>()
 //!     .length_adjustment(0)   // default value
 //!     .num_skip(0) // Do not strip frame header
 //!     .new_read(io);
@@ -118,7 +118,7 @@
 //! # fn bind_read<T: AsyncRead>(io: T) {
 //! LengthDelimitedCodec::builder()
 //!     .length_field_offset(0) // default value
-//!     .length_field_length(2)
+//!     .length_field_type::<u16>()
 //!     .length_adjustment(0)   // default value
 //!     // `num_skip` is not needed, the default is to skip
 //!     .new_read(io);
@@ -150,7 +150,7 @@
 //! # fn bind_read<T: AsyncRead>(io: T) {
 //! LengthDelimitedCodec::builder()
 //!     .length_field_offset(0) // default value
-//!     .length_field_length(2)
+//!     .length_field_type::<u16>()
 //!     .length_adjustment(-2)  // size of head
 //!     .num_skip(0)
 //!     .new_read(io);
@@ -228,7 +228,7 @@
 //! # fn bind_read<T: AsyncRead>(io: T) {
 //! LengthDelimitedCodec::builder()
 //!     .length_field_offset(1) // length of hdr1
-//!     .length_field_length(2)
+//!     .length_field_type::<u16>()
 //!     .length_adjustment(1)  // length of hdr2
 //!     .num_skip(3) // length of hdr1 + LEN
 //!     .new_read(io);
@@ -274,7 +274,7 @@
 //! # fn bind_read<T: AsyncRead>(io: T) {
 //! LengthDelimitedCodec::builder()
 //!     .length_field_offset(1) // length of hdr1
-//!     .length_field_length(2)
+//!     .length_field_type::<u16>()
 //!     .length_adjustment(-3)  // length of hdr1 + LEN, negative
 //!     .num_skip(3)
 //!     .new_read(io);
@@ -350,7 +350,7 @@
 //! # fn write_frame<T: AsyncWrite>(io: T) {
 //! # let _ =
 //! LengthDelimitedCodec::builder()
-//!     .length_field_length(2)
+//!     .length_field_type::<u16>()
 //!     .new_write(io);
 //! # }
 //! # pub fn main() {}
@@ -379,7 +379,7 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use std::error::Error as StdError;
 use std::io::{self, Cursor};
-use std::{cmp, fmt};
+use std::{cmp, fmt, mem};
 
 /// Configure length delimited `LengthDelimitedCodec`s.
 ///
@@ -629,6 +629,24 @@ impl Default for LengthDelimitedCodec {
 
 // ===== impl Builder =====
 
+mod builder {
+    /// Types that can be used with `Builder::length_field_type`.
+    pub trait LengthFieldType {}
+
+    impl LengthFieldType for u8 {}
+    impl LengthFieldType for u16 {}
+    impl LengthFieldType for u32 {}
+    impl LengthFieldType for u64 {}
+
+    #[cfg(any(
+        target_pointer_width = "8",
+        target_pointer_width = "16",
+        target_pointer_width = "32",
+        target_pointer_width = "64",
+    ))]
+    impl LengthFieldType for usize {}
+}
+
 impl Builder {
     /// Creates a new length delimited codec builder with default configuration
     /// values.
@@ -642,7 +660,7 @@ impl Builder {
     /// # fn bind_read<T: AsyncRead>(io: T) {
     /// LengthDelimitedCodec::builder()
     ///     .length_field_offset(0)
-    ///     .length_field_length(2)
+    ///     .length_field_type::<u16>()
     ///     .length_adjustment(0)
     ///     .num_skip(0)
     ///     .new_read(io);
@@ -746,7 +764,7 @@ impl Builder {
         }
     }
 
-    /// Sets the max frame length
+    /// Sets the max frame length in bytes
     ///
     /// This configuration option applies to both encoding and decoding. The
     /// default value is 8MB.
@@ -767,7 +785,7 @@ impl Builder {
     ///
     /// # fn bind_read<T: AsyncRead>(io: T) {
     /// LengthDelimitedCodec::builder()
-    ///     .max_frame_length(8 * 1024)
+    ///     .max_frame_length(8 * 1024 * 1024)
     ///     .new_read(io);
     /// # }
     /// # pub fn main() {}
@@ -777,6 +795,42 @@ impl Builder {
         self
     }
 
+    /// Sets the unsigned integer type used to represent the length field.
+    ///
+    /// The default type is [`u32`]. The max type is [`u64`] (or [`usize`] on
+    /// 64-bit targets).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tokio::io::AsyncRead;
+    /// use tokio_util::codec::LengthDelimitedCodec;
+    ///
+    /// # fn bind_read<T: AsyncRead>(io: T) {
+    /// LengthDelimitedCodec::builder()
+    ///     .length_field_type::<u32>()
+    ///     .new_read(io);
+    /// # }
+    /// # pub fn main() {}
+    /// ```
+    ///
+    /// Unlike [`Builder::length_field_length`], this does not fail at runtime
+    /// and instead produces a compile error:
+    ///
+    /// ```compile_fail
+    /// # use tokio::io::AsyncRead;
+    /// # use tokio_util::codec::LengthDelimitedCodec;
+    /// # fn bind_read<T: AsyncRead>(io: T) {
+    /// LengthDelimitedCodec::builder()
+    ///     .length_field_type::<u128>()
+    ///     .new_read(io);
+    /// # }
+    /// # pub fn main() {}
+    /// ```
+    pub fn length_field_type<T: builder::LengthFieldType>(&mut self) -> &mut Self {
+        self.length_field_length(mem::size_of::<T>())
+    }
+
     /// Sets the number of bytes used to represent the length field
     ///
     /// The default value is `4`. The max value is `8`.
@@ -878,7 +932,7 @@ impl Builder {
     /// # pub fn main() {
     /// LengthDelimitedCodec::builder()
     ///     .length_field_offset(0)
-    ///     .length_field_length(2)
+    ///     .length_field_type::<u16>()
     ///     .length_adjustment(0)
     ///     .num_skip(0)
     ///     .new_codec();
@@ -902,7 +956,7 @@ impl Builder {
     /// # fn bind_read<T: AsyncRead>(io: T) {
     /// LengthDelimitedCodec::builder()
     ///     .length_field_offset(0)
-    ///     .length_field_length(2)
+    ///     .length_field_type::<u16>()
     ///     .length_adjustment(0)
     ///     .num_skip(0)
     ///     .new_read(io);
@@ -925,7 +979,7 @@ impl Builder {
     /// # use tokio_util::codec::LengthDelimitedCodec;
     /// # fn write_frame<T: AsyncWrite>(io: T) {
     /// LengthDelimitedCodec::builder()
-    ///     .length_field_length(2)
+    ///     .length_field_type::<u16>()
     ///     .new_write(io);
     /// # }
     /// # pub fn main() {}
@@ -947,7 +1001,7 @@ impl Builder {
     /// # fn write_frame<T: AsyncRead + AsyncWrite>(io: T) {
     /// # let _ =
     /// LengthDelimitedCodec::builder()
-    ///     .length_field_length(2)
+    ///     .length_field_type::<u16>()
     ///     .new_framed(io);
     /// # }
     /// # pub fn main() {}

tokio-util/src/io/stream_reader.rs

@@ -84,13 +84,24 @@ where
     }
 
     /// Do we have a chunk and is it non-empty?
-    fn has_chunk(self: Pin<&mut Self>) -> bool {
-        if let Some(chunk) = self.project().chunk {
+    fn has_chunk(&self) -> bool {
+        if let Some(ref chunk) = self.chunk {
             chunk.remaining() > 0
         } else {
             false
         }
     }
+
+    /// Consumes this `StreamReader`, returning a Tuple consisting
+    /// of the underlying stream and an Option of the internal buffer,
+    /// which is Some in case the buffer contains elements.
+    pub fn into_inner_with_chunk(self) -> (S, Option<B>) {
+        if self.has_chunk() {
+            (self.inner, self.chunk)
+        } else {
+            (self.inner, None)
+        }
+    }
 }
 
 impl<S, B> StreamReader<S, B> {
@@ -118,6 +129,10 @@ impl<S, B> StreamReader<S, B> {
     /// Consumes this `BufWriter`, returning the underlying stream.
     ///
     /// Note that any leftover data in the internal buffer is lost.
+    /// If you additionally want access to the internal buffer use
+    /// [`into_inner_with_chunk`].
+    ///
+    /// [`into_inner_with_chunk`]: crate::io::StreamReader::into_inner_with_chunk
     pub fn into_inner(self) -> S {
         self.inner
     }

tokio-util/src/io/sync_bridge.rs

@@ -85,9 +85,10 @@ impl<T: Unpin> SyncIoBridge<T> {
     ///
     /// Use e.g. `SyncIoBridge::new(Box::pin(src))`.
     ///
-    /// # Panic
+    /// # Panics
     ///
     /// This will panic if called outside the context of a Tokio runtime.
+    #[track_caller]
     pub fn new(src: T) -> Self {
         Self::new_with_handle(src, tokio::runtime::Handle::current())
     }

tokio-util/src/lib.rs

@@ -5,7 +5,6 @@
     rust_2018_idioms,
     unreachable_pub
 )]
-#![cfg_attr(docsrs, deny(rustdoc::broken_intra_doc_links))]
 #![doc(test(
     no_crate_inject,
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables))
@@ -31,6 +30,7 @@ cfg_codec! {
 
 cfg_net! {
     pub mod udp;
+    pub mod net;
 }
 
 cfg_compat! {
@@ -43,6 +43,7 @@ cfg_io! {
 
 cfg_rt! {
     pub mod context;
+    pub mod task;
 }
 
 cfg_time! {
@@ -114,6 +115,9 @@ mod util {
 
         let n = {
             let dst = buf.chunk_mut();
+
+            // Safety: `chunk_mut()` returns a `&mut UninitSlice`, and `UninitSlice` is a
+            // transparent wrapper around `[MaybeUninit<u8>]`.
             let dst = unsafe { &mut *(dst as *mut _ as *mut [MaybeUninit<u8>]) };
             let mut buf = ReadBuf::uninit(dst);
             let ptr = buf.filled().as_ptr();

tokio-util/src/net/mod.rs

@@ -0,0 +1,97 @@
+//! TCP/UDP/Unix helpers for tokio.
+
+use crate::either::Either;
+use std::future::Future;
+use std::io::Result;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+#[cfg(unix)]
+pub mod unix;
+
+/// A trait for a listener: `TcpListener` and `UnixListener`.
+pub trait Listener {
+    /// The stream's type of this listener.
+    type Io: tokio::io::AsyncRead + tokio::io::AsyncWrite;
+    /// The socket address type of this listener.
+    type Addr;
+
+    /// Polls to accept a new incoming connection to this listener.
+    fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>>;
+
+    /// Accepts a new incoming connection from this listener.
+    fn accept(&mut self) -> ListenerAcceptFut<'_, Self>
+    where
+        Self: Sized,
+    {
+        ListenerAcceptFut { listener: self }
+    }
+
+    /// Returns the local address that this listener is bound to.
+    fn local_addr(&self) -> Result<Self::Addr>;
+}
+
+impl Listener for tokio::net::TcpListener {
+    type Io = tokio::net::TcpStream;
+    type Addr = std::net::SocketAddr;
+
+    fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>> {
+        Self::poll_accept(self, cx)
+    }
+
+    fn local_addr(&self) -> Result<Self::Addr> {
+        self.local_addr().map(Into::into)
+    }
+}
+
+/// Future for accepting a new connection from a listener.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct ListenerAcceptFut<'a, L> {
+    listener: &'a mut L,
+}
+
+impl<'a, L> Future for ListenerAcceptFut<'a, L>
+where
+    L: Listener,
+{
+    type Output = Result<(L::Io, L::Addr)>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.listener.poll_accept(cx)
+    }
+}
+
+impl<L, R> Either<L, R>
+where
+    L: Listener,
+    R: Listener,
+{
+    /// Accepts a new incoming connection from this listener.
+    pub async fn accept(&mut self) -> Result<Either<(L::Io, L::Addr), (R::Io, R::Addr)>> {
+        match self {
+            Either::Left(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok(Either::Left((stream, addr)))
+            }
+            Either::Right(listener) => {
+                let (stream, addr) = listener.accept().await?;
+                Ok(Either::Right((stream, addr)))
+            }
+        }
+    }
+
+    /// Returns the local address that this listener is bound to.
+    pub fn local_addr(&self) -> Result<Either<L::Addr, R::Addr>> {
+        match self {
+            Either::Left(listener) => {
+                let addr = listener.local_addr()?;
+                Ok(Either::Left(addr))
+            }
+            Either::Right(listener) => {
+                let addr = listener.local_addr()?;
+                Ok(Either::Right(addr))
+            }
+        }
+    }
+}

tokio-util/src/net/unix/mod.rs

@@ -0,0 +1,18 @@
+//! Unix domain socket helpers.
+
+use super::Listener;
+use std::io::Result;
+use std::task::{Context, Poll};
+
+impl Listener for tokio::net::UnixListener {
+    type Io = tokio::net::UnixStream;
+    type Addr = tokio::net::unix::SocketAddr;
+
+    fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<(Self::Io, Self::Addr)>> {
+        Self::poll_accept(self, cx)
+    }
+
+    fn local_addr(&self) -> Result<Self::Addr> {
+        self.local_addr().map(Into::into)
+    }
+}

tokio-util/src/sync/mod.rs

@@ -4,7 +4,7 @@ mod cancellation_token;
 pub use cancellation_token::{guard::DropGuard, CancellationToken, WaitForCancellationFuture};
 
 mod mpsc;
-pub use mpsc::PollSender;
+pub use mpsc::{PollSendError, PollSender};
 
 mod poll_semaphore;
 pub use poll_semaphore::PollSemaphore;

tokio-util/src/sync/mpsc.rs

@@ -1,221 +1,284 @@
-use futures_core::ready;
 use futures_sink::Sink;
 use std::pin::Pin;
-use std::sync::Arc;
 use std::task::{Context, Poll};
-use tokio::sync::mpsc::{error::SendError, Sender};
+use std::{fmt, mem};
+use tokio::sync::mpsc::OwnedPermit;
+use tokio::sync::mpsc::Sender;
 
 use super::ReusableBoxFuture;
 
-// This implementation was chosen over something based on permits because to get a
-// `tokio::sync::mpsc::Permit` out of the `inner` future, you must transmute the
-// lifetime on the permit to `'static`.
+/// Error returned by the `PollSender` when the channel is closed.
+#[derive(Debug)]
+pub struct PollSendError<T>(Option<T>);
+
+impl<T> PollSendError<T> {
+    /// Consumes the stored value, if any.
+    ///
+    /// If this error was encountered when calling `start_send`/`send_item`, this will be the item
+    /// that the caller attempted to send.  Otherwise, it will be `None`.
+    pub fn into_inner(self) -> Option<T> {
+        self.0
+    }
+}
+
+impl<T> fmt::Display for PollSendError<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(fmt, "channel closed")
+    }
+}
+
+impl<T: fmt::Debug> std::error::Error for PollSendError<T> {}
+
+#[derive(Debug)]
+enum State<T> {
+    Idle(Sender<T>),
+    Acquiring,
+    ReadyToSend(OwnedPermit<T>),
+    Closed,
+}
 
 /// A wrapper around [`mpsc::Sender`] that can be polled.
 ///
 /// [`mpsc::Sender`]: tokio::sync::mpsc::Sender
 #[derive(Debug)]
 pub struct PollSender<T> {
-    /// is none if closed
-    sender: Option<Arc<Sender<T>>>,
-    is_sending: bool,
-    inner: ReusableBoxFuture<Result<(), SendError<T>>>,
+    sender: Option<Sender<T>>,
+    state: State<T>,
+    acquire: ReusableBoxFuture<'static, Result<OwnedPermit<T>, PollSendError<T>>>,
 }
 
-// By reusing the same async fn for both Some and None, we make sure every
-// future passed to ReusableBoxFuture has the same underlying type, and hence
-// the same size and alignment.
-async fn make_future<T>(data: Option<(Arc<Sender<T>>, T)>) -> Result<(), SendError<T>> {
+// Creates a future for acquiring a permit from the underlying channel.  This is used to ensure
+// there's capacity for a send to complete.
+//
+// By reusing the same async fn for both `Some` and `None`, we make sure every future passed to
+// ReusableBoxFuture has the same underlying type, and hence the same size and alignment.
+async fn make_acquire_future<T>(
+    data: Option<Sender<T>>,
+) -> Result<OwnedPermit<T>, PollSendError<T>> {
     match data {
-        Some((sender, value)) => sender.send(value).await,
-        None => unreachable!(
-            "This future should not be pollable, as is_sending should be set to false."
-        ),
+        Some(sender) => sender
+            .reserve_owned()
+            .await
+            .map_err(|_| PollSendError(None)),
+        None => unreachable!("this future should not be pollable in this state"),
     }
 }
 
 impl<T: Send + 'static> PollSender<T> {
-    /// Create a new `PollSender`.
+    /// Creates a new `PollSender`.
     pub fn new(sender: Sender<T>) -> Self {
         Self {
-            sender: Some(Arc::new(sender)),
-            is_sending: false,
-            inner: ReusableBoxFuture::new(make_future(None)),
+            sender: Some(sender.clone()),
+            state: State::Idle(sender),
+            acquire: ReusableBoxFuture::new(make_acquire_future(None)),
         }
     }
 
-    /// Start sending a new item.
+    fn take_state(&mut self) -> State<T> {
+        mem::replace(&mut self.state, State::Closed)
+    }
+
+    /// Attempts to prepare the sender to receive a value.
     ///
-    /// This method panics if a send is currently in progress. To ensure that no
-    /// send is in progress, call `poll_send_done` first until it returns
-    /// `Poll::Ready`.
+    /// This method must be called and return `Poll::Ready(Ok(()))` prior to each call to
+    /// `send_item`.
     ///
-    /// If this method returns an error, that indicates that the channel is
-    /// closed. Note that this method is not guaranteed to return an error if
-    /// the channel is closed, but in that case the error would be reported by
-    /// the first call to `poll_send_done`.
-    pub fn start_send(&mut self, value: T) -> Result<(), SendError<T>> {
-        if self.is_sending {
-            panic!("start_send called while not ready.");
-        }
-        match self.sender.clone() {
-            Some(sender) => {
-                self.inner.set(make_future(Some((sender, value))));
-                self.is_sending = true;
-                Ok(())
+    /// This method returns `Poll::Ready` once the underlying channel is ready to receive a value,
+    /// by reserving a slot in the channel for the item to be sent. If this method returns
+    /// `Poll::Pending`, the current task is registered to be notified (via
+    /// `cx.waker().wake_by_ref()`) when `poll_reserve` should be called again.
+    ///
+    /// # Errors
+    ///
+    /// If the channel is closed, an error will be returned.  This is a permanent state.
+    pub fn poll_reserve(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), PollSendError<T>>> {
+        loop {
+            let (result, next_state) = match self.take_state() {
+                State::Idle(sender) => {
+                    // Start trying to acquire a permit to reserve a slot for our send, and
+                    // immediately loop back around to poll it the first time.
+                    self.acquire.set(make_acquire_future(Some(sender)));
+                    (None, State::Acquiring)
+                }
+                State::Acquiring => match self.acquire.poll(cx) {
+                    // Channel has capacity.
+                    Poll::Ready(Ok(permit)) => {
+                        (Some(Poll::Ready(Ok(()))), State::ReadyToSend(permit))
+                    }
+                    // Channel is closed.
+                    Poll::Ready(Err(e)) => (Some(Poll::Ready(Err(e))), State::Closed),
+                    // Channel doesn't have capacity yet, so we need to wait.
+                    Poll::Pending => (Some(Poll::Pending), State::Acquiring),
+                },
+                // We're closed, either by choice or because the underlying sender was closed.
+                s @ State::Closed => (Some(Poll::Ready(Err(PollSendError(None)))), s),
+                // We're already ready to send an item.
+                s @ State::ReadyToSend(_) => (Some(Poll::Ready(Ok(()))), s),
+            };
+
+            self.state = next_state;
+            if let Some(result) = result {
+                return result;
             }
-            None => Err(SendError(value)),
         }
     }
 
-    /// If a send is in progress, poll for its completion. If no send is in progress,
-    /// this method returns `Poll::Ready(Ok(()))`.
+    /// Sends an item to the channel.
     ///
-    /// This method can return the following values:
+    /// Before calling `send_item`, `poll_reserve` must be called with a successful return
+    /// value of `Poll::Ready(Ok(()))`.
     ///
-    ///  - `Poll::Ready(Ok(()))` if the in-progress send has been completed, or there is
-    ///    no send in progress (even if the channel is closed).
-    ///  - `Poll::Ready(Err(err))` if the in-progress send failed because the channel has
-    ///    been closed.
-    ///  - `Poll::Pending` if a send is in progress, but it could not complete now.
+    /// # Errors
     ///
-    /// When this method returns `Poll::Pending`, the current task is scheduled
-    /// to receive a wakeup when the message is sent, or when the entire channel
-    /// is closed (but not if just this sender is closed by
-    /// `close_this_sender`). Note that on multiple calls to `poll_send_done`,
-    /// only the `Waker` from the `Context` passed to the most recent call is
-    /// scheduled to receive a wakeup.
+    /// If the channel is closed, an error will be returned.  This is a permanent state.
     ///
-    /// If this method returns `Poll::Ready`, then `start_send` is guaranteed to
-    /// not panic.
-    pub fn poll_send_done(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), SendError<T>>> {
-        if !self.is_sending {
-            return Poll::Ready(Ok(()));
-        }
+    /// # Panics
+    ///
+    /// If `poll_reserve` was not successfully called prior to calling `send_item`, then this method
+    /// will panic.
+    #[track_caller]
+    pub fn send_item(&mut self, value: T) -> Result<(), PollSendError<T>> {
+        let (result, next_state) = match self.take_state() {
+            State::Idle(_) | State::Acquiring => {
+                panic!("`send_item` called without first calling `poll_reserve`")
+            }
+            // We have a permit to send our item, so go ahead, which gets us our sender back.
+            State::ReadyToSend(permit) => (Ok(()), State::Idle(permit.send(value))),
+            // We're closed, either by choice or because the underlying sender was closed.
+            State::Closed => (Err(PollSendError(Some(value))), State::Closed),
+        };
 
-        let result = self.inner.poll(cx);
-        if result.is_ready() {
-            self.is_sending = false;
-        }
-        if let Poll::Ready(Err(_)) = &result {
-            self.sender = None;
-        }
+        // Handle deferred closing if `close` was called between `poll_reserve` and `send_item`.
+        self.state = if self.sender.is_some() {
+            next_state
+        } else {
+            State::Closed
+        };
         result
     }
 
-    /// Check whether the channel is ready to send more messages now.
+    /// Checks whether this sender is been closed.
     ///
-    /// If this method returns `true`, then `start_send` is guaranteed to not
-    /// panic.
-    ///
-    /// If the channel is closed, this method returns `true`.
-    pub fn is_ready(&self) -> bool {
-        !self.is_sending
-    }
-
-    /// Check whether the channel has been closed.
+    /// The underlying channel that this sender was wrapping may still be open.
     pub fn is_closed(&self) -> bool {
-        match &self.sender {
-            Some(sender) => sender.is_closed(),
-            None => true,
-        }
+        matches!(self.state, State::Closed) || self.sender.is_none()
     }
 
-    /// Clone the underlying `Sender`.
+    /// Gets a reference to the `Sender` of the underlying channel.
     ///
-    /// If this method returns `None`, then the channel is closed. (But it is
-    /// not guaranteed to return `None` if the channel is closed.)
-    pub fn clone_inner(&self) -> Option<Sender<T>> {
-        self.sender.as_ref().map(|sender| (&**sender).clone())
+    /// If `PollSender` has been closed, `None` is returned. The underlying channel that this sender
+    /// was wrapping may still be open.
+    pub fn get_ref(&self) -> Option<&Sender<T>> {
+        self.sender.as_ref()
     }
 
-    /// Access the underlying `Sender`.
+    /// Closes this sender.
     ///
-    /// If this method returns `None`, then the channel is closed. (But it is
-    /// not guaranteed to return `None` if the channel is closed.)
-    pub fn inner_ref(&self) -> Option<&Sender<T>> {
-        self.sender.as_deref()
-    }
-
-    // This operation is supported because it is required by the Sink trait.
-    /// Close this sender. No more messages can be sent from this sender.
+    /// No more messages will be able to be sent from this sender, but the underlying channel will
+    /// remain open until all senders have dropped, or until the [`Receiver`] closes the channel.
     ///
-    /// Note that this only closes the channel from the view-point of this
-    /// sender. The channel remains open until all senders have gone away, or
-    /// until the [`Receiver`] closes the channel.
-    ///
-    /// If there is a send in progress when this method is called, that send is
-    /// unaffected by this operation, and `poll_send_done` can still be called
-    /// to complete that send.
+    /// If a slot was previously reserved by calling `poll_reserve`, then a final call can be made
+    /// to `send_item` in order to consume the reserved slot.  After that, no further sends will be
+    /// possible.  If you do not intend to send another item, you can release the reserved slot back
+    /// to the underlying sender by calling [`abort_send`].
     ///
+    /// [`abort_send`]: crate::sync::PollSender::abort_send
     /// [`Receiver`]: tokio::sync::mpsc::Receiver
-    pub fn close_this_sender(&mut self) {
+    pub fn close(&mut self) {
+        // Mark ourselves officially closed by dropping our main sender.
         self.sender = None;
+
+        // If we're already idle, closed, or we haven't yet reserved a slot, we can quickly
+        // transition to the closed state.  Otherwise, leave the existing permit in place for the
+        // caller if they want to complete the send.
+        match self.state {
+            State::Idle(_) => self.state = State::Closed,
+            State::Acquiring => {
+                self.acquire.set(make_acquire_future(None));
+                self.state = State::Closed;
+            }
+            _ => {}
+        }
     }
 
-    /// Abort the current in-progress send, if any.
+    /// Aborts the current in-progress send, if any.
     ///
-    /// Returns `true` if a send was aborted.
+    /// Returns `true` if a send was aborted.  If the sender was closed prior to calling
+    /// `abort_send`, then the sender will remain in the closed state, otherwise the sender will be
+    /// ready to attempt another send.
     pub fn abort_send(&mut self) -> bool {
-        if self.is_sending {
-            self.inner.set(make_future(None));
-            self.is_sending = false;
-            true
-        } else {
-            false
-        }
+        // We may have been closed in the meantime, after a call to `poll_reserve` already
+        // succeeded.  We'll check if `self.sender` is `None` to see if we should transition to the
+        // closed state when we actually abort a send, rather than resetting ourselves back to idle.
+
+        let (result, next_state) = match self.take_state() {
+            // We're currently trying to reserve a slot to send into.
+            State::Acquiring => {
+                // Replacing the future drops the in-flight one.
+                self.acquire.set(make_acquire_future(None));
+
+                // If we haven't closed yet, we have to clone our stored sender since we have no way
+                // to get it back from the acquire future we just dropped.
+                let state = match self.sender.clone() {
+                    Some(sender) => State::Idle(sender),
+                    None => State::Closed,
+                };
+                (true, state)
+            }
+            // We got the permit.  If we haven't closed yet, get the sender back.
+            State::ReadyToSend(permit) => {
+                let state = if self.sender.is_some() {
+                    State::Idle(permit.release())
+                } else {
+                    State::Closed
+                };
+                (true, state)
+            }
+            s => (false, s),
+        };
+
+        self.state = next_state;
+        result
     }
 }
 
 impl<T> Clone for PollSender<T> {
-    /// Clones this `PollSender`. The resulting clone will not have any
-    /// in-progress send operations, even if the current `PollSender` does.
+    /// Clones this `PollSender`.
+    ///
+    /// The resulting `PollSender` will have an initial state identical to calling `PollSender::new`.
     fn clone(&self) -> PollSender<T> {
+        let (sender, state) = match self.sender.clone() {
+            Some(sender) => (Some(sender.clone()), State::Idle(sender)),
+            None => (None, State::Closed),
+        };
+
         Self {
-            sender: self.sender.clone(),
-            is_sending: false,
-            inner: ReusableBoxFuture::new(async { unreachable!() }),
+            sender,
+            state,
+            // We don't use `make_acquire_future` here because our relaxed bounds on `T` are not
+            // compatible with the transitive bounds required by `Sender<T>`.
+            acquire: ReusableBoxFuture::new(async { unreachable!() }),
         }
     }
 }
 
 impl<T: Send + 'static> Sink<T> for PollSender<T> {
-    type Error = SendError<T>;
+    type Error = PollSendError<T>;
 
-    /// This is equivalent to calling [`poll_send_done`].
-    ///
-    /// [`poll_send_done`]: PollSender::poll_send_done
     fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        Pin::into_inner(self).poll_send_done(cx)
+        Pin::into_inner(self).poll_reserve(cx)
     }
 
-    /// This is equivalent to calling [`poll_send_done`].
-    ///
-    /// [`poll_send_done`]: PollSender::poll_send_done
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        Pin::into_inner(self).poll_send_done(cx)
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
     }
 
-    /// This is equivalent to calling [`start_send`].
-    ///
-    /// [`start_send`]: PollSender::start_send
     fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
-        Pin::into_inner(self).start_send(item)
+        Pin::into_inner(self).send_item(item)
     }
 
-    /// This method will first flush the `PollSender`, and then close it by
-    /// calling [`close_this_sender`].
-    ///
-    /// If a send fails while flushing because the [`Receiver`] has gone away,
-    /// then this function returns an error. The channel is still successfully
-    /// closed in this situation.
-    ///
-    /// [`close_this_sender`]: PollSender::close_this_sender
-    /// [`Receiver`]: tokio::sync::mpsc::Receiver
-    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        ready!(self.as_mut().poll_flush(cx))?;
-
-        Pin::into_inner(self).close_this_sender();
+    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Pin::into_inner(self).close();
         Poll::Ready(Ok(()))
     }
 }

tokio-util/src/sync/poll_semaphore.rs

@@ -12,7 +12,7 @@ use super::ReusableBoxFuture;
 /// [`Semaphore`]: tokio::sync::Semaphore
 pub struct PollSemaphore {
     semaphore: Arc<Semaphore>,
-    permit_fut: Option<ReusableBoxFuture<Result<OwnedSemaphorePermit, AcquireError>>>,
+    permit_fut: Option<ReusableBoxFuture<'static, Result<OwnedSemaphorePermit, AcquireError>>>,
 }
 
 impl PollSemaphore {

tokio-util/src/sync/reusable_box.rs

@@ -1,33 +1,29 @@
 use std::alloc::Layout;
+use std::fmt;
 use std::future::Future;
-use std::panic::AssertUnwindSafe;
+use std::marker::PhantomData;
+use std::mem::{self, ManuallyDrop};
 use std::pin::Pin;
-use std::ptr::{self, NonNull};
+use std::ptr;
 use std::task::{Context, Poll};
-use std::{fmt, panic};
 
-/// A reusable `Pin<Box<dyn Future<Output = T> + Send>>`.
+/// A reusable `Pin<Box<dyn Future<Output = T> + Send + 'a>>`.
 ///
 /// This type lets you replace the future stored in the box without
 /// reallocating when the size and alignment permits this.
-pub struct ReusableBoxFuture<T> {
-    boxed: NonNull<dyn Future<Output = T> + Send>,
+pub struct ReusableBoxFuture<'a, T> {
+    boxed: Pin<Box<dyn Future<Output = T> + Send + 'a>>,
 }
 
-impl<T> ReusableBoxFuture<T> {
+impl<'a, T> ReusableBoxFuture<'a, T> {
     /// Create a new `ReusableBoxFuture<T>` containing the provided future.
     pub fn new<F>(future: F) -> Self
     where
-        F: Future<Output = T> + Send + 'static,
+        F: Future<Output = T> + Send + 'a,
     {
-        let boxed: Box<dyn Future<Output = T> + Send> = Box::new(future);
-
-        let boxed = Box::into_raw(boxed);
-
-        // SAFETY: Box::into_raw does not return null pointers.
-        let boxed = unsafe { NonNull::new_unchecked(boxed) };
-
-        Self { boxed }
+        Self {
+            boxed: Box::pin(future),
+        }
     }
 
     /// Replace the future currently stored in this box.
@@ -36,7 +32,7 @@ impl<T> ReusableBoxFuture<T> {
     /// different from the layout of the currently stored future.
     pub fn set<F>(&mut self, future: F)
     where
-        F: Future<Output = T> + Send + 'static,
+        F: Future<Output = T> + Send + 'a,
     {
         if let Err(future) = self.try_set(future) {
             *self = Self::new(future);
@@ -50,64 +46,31 @@ impl<T> ReusableBoxFuture<T> {
     /// future.
     pub fn try_set<F>(&mut self, future: F) -> Result<(), F>
     where
-        F: Future<Output = T> + Send + 'static,
+        F: Future<Output = T> + Send + 'a,
     {
-        // SAFETY: The pointer is not dangling.
-        let self_layout = {
-            let dyn_future: &(dyn Future<Output = T> + Send) = unsafe { self.boxed.as_ref() };
-            Layout::for_value(dyn_future)
-        };
-
-        if Layout::new::<F>() == self_layout {
-            // SAFETY: We just checked that the layout of F is correct.
-            unsafe {
-                self.set_same_layout(future);
-            }
-
-            Ok(())
-        } else {
-            Err(future)
+        // If we try to inline the contents of this function, the type checker complains because
+        // the bound `T: 'a` is not satisfied in the call to `pending()`. But by putting it in an
+        // inner function that doesn't have `T` as a generic parameter, we implicitly get the bound
+        // `F::Output: 'a` transitively through `F: 'a`, allowing us to call `pending()`.
+        #[inline(always)]
+        fn real_try_set<'a, F>(
+            this: &mut ReusableBoxFuture<'a, F::Output>,
+            future: F,
+        ) -> Result<(), F>
+        where
+            F: Future + Send + 'a,
+        {
+            // future::Pending<T> is a ZST so this never allocates.
+            let boxed = mem::replace(&mut this.boxed, Box::pin(Pending(PhantomData)));
+            reuse_pin_box(boxed, future, |boxed| this.boxed = Pin::from(boxed))
         }
-    }
 
-    /// Set the current future.
-    ///
-    /// # Safety
-    ///
-    /// This function requires that the layout of the provided future is the
-    /// same as `self.layout`.
-    unsafe fn set_same_layout<F>(&mut self, future: F)
-    where
-        F: Future<Output = T> + Send + 'static,
-    {
-        // Drop the existing future, catching any panics.
-        let result = panic::catch_unwind(AssertUnwindSafe(|| {
-            ptr::drop_in_place(self.boxed.as_ptr());
-        }));
-
-        // Overwrite the future behind the pointer. This is safe because the
-        // allocation was allocated with the same size and alignment as the type F.
-        let self_ptr: *mut F = self.boxed.as_ptr() as *mut F;
-        ptr::write(self_ptr, future);
-
-        // Update the vtable of self.boxed. The pointer is not null because we
-        // just got it from self.boxed, which is not null.
-        self.boxed = NonNull::new_unchecked(self_ptr);
-
-        // If the old future's destructor panicked, resume unwinding.
-        match result {
-            Ok(()) => {}
-            Err(payload) => {
-                panic::resume_unwind(payload);
-            }
-        }
+        real_try_set(self, future)
     }
 
     /// Get a pinned reference to the underlying future.
     pub fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
-        // SAFETY: The user of this box cannot move the box, and we do not move it
-        // either.
-        unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
+        self.boxed.as_mut()
     }
 
     /// Poll the future stored inside this box.
@@ -116,7 +79,7 @@ impl<T> ReusableBoxFuture<T> {
     }
 }
 
-impl<T> Future for ReusableBoxFuture<T> {
+impl<T> Future for ReusableBoxFuture<'_, T> {
     type Output = T;
 
     /// Poll the future stored inside this box.
@@ -125,27 +88,84 @@ impl<T> Future for ReusableBoxFuture<T> {
     }
 }
 
-// The future stored inside ReusableBoxFuture<T> must be Send.
-unsafe impl<T> Send for ReusableBoxFuture<T> {}
-
 // The only method called on self.boxed is poll, which takes &mut self, so this
 // struct being Sync does not permit any invalid access to the Future, even if
 // the future is not Sync.
-unsafe impl<T> Sync for ReusableBoxFuture<T> {}
+unsafe impl<T> Sync for ReusableBoxFuture<'_, T> {}
 
-// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
-impl<T> Unpin for ReusableBoxFuture<T> {}
-
-impl<T> Drop for ReusableBoxFuture<T> {
-    fn drop(&mut self) {
-        unsafe {
-            drop(Box::from_raw(self.boxed.as_ptr()));
-        }
-    }
-}
-
-impl<T> fmt::Debug for ReusableBoxFuture<T> {
+impl<T> fmt::Debug for ReusableBoxFuture<'_, T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("ReusableBoxFuture").finish()
     }
 }
+
+fn reuse_pin_box<T: ?Sized, U, O, F>(boxed: Pin<Box<T>>, new_value: U, callback: F) -> Result<O, U>
+where
+    F: FnOnce(Box<U>) -> O,
+{
+    let layout = Layout::for_value::<T>(&*boxed);
+    if layout != Layout::new::<U>() {
+        return Err(new_value);
+    }
+
+    // SAFETY: We don't ever construct a non-pinned reference to the old `T` from now on, and we
+    // always drop the `T`.
+    let raw: *mut T = Box::into_raw(unsafe { Pin::into_inner_unchecked(boxed) });
+
+    // When dropping the old value panics, we still want to call `callback` — so move the rest of
+    // the code into a guard type.
+    let guard = CallOnDrop::new(|| {
+        let raw: *mut U = raw.cast::<U>();
+        unsafe { raw.write(new_value) };
+
+        // SAFETY:
+        // - `T` and `U` have the same layout.
+        // - `raw` comes from a `Box` that uses the same allocator as this one.
+        // - `raw` points to a valid instance of `U` (we just wrote it in).
+        let boxed = unsafe { Box::from_raw(raw) };
+
+        callback(boxed)
+    });
+
+    // Drop the old value.
+    unsafe { ptr::drop_in_place(raw) };
+
+    // Run the rest of the code.
+    Ok(guard.call())
+}
+
+struct CallOnDrop<O, F: FnOnce() -> O> {
+    f: ManuallyDrop<F>,
+}
+
+impl<O, F: FnOnce() -> O> CallOnDrop<O, F> {
+    fn new(f: F) -> Self {
+        let f = ManuallyDrop::new(f);
+        Self { f }
+    }
+    fn call(self) -> O {
+        let mut this = ManuallyDrop::new(self);
+        let f = unsafe { ManuallyDrop::take(&mut this.f) };
+        f()
+    }
+}
+
+impl<O, F: FnOnce() -> O> Drop for CallOnDrop<O, F> {
+    fn drop(&mut self) {
+        let f = unsafe { ManuallyDrop::take(&mut self.f) };
+        f();
+    }
+}
+
+/// The same as `std::future::Pending<T>`; we can't use that type directly because on rustc
+/// versions <1.60 it didn't unconditionally implement `Send`.
+// FIXME: use `std::future::Pending<T>` once the MSRV is >=1.60
+struct Pending<T>(PhantomData<fn() -> T>);
+
+impl<T> Future for Pending<T> {
+    type Output = T;
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Pending
+    }
+}

tokio-util/src/task/join_map.rs

@@ -0,0 +1,817 @@
+use hashbrown::hash_map::RawEntryMut;
+use hashbrown::HashMap;
+use std::borrow::Borrow;
+use std::collections::hash_map::RandomState;
+use std::fmt;
+use std::future::Future;
+use std::hash::{BuildHasher, Hash, Hasher};
+use tokio::runtime::Handle;
+use tokio::task::{AbortHandle, Id, JoinError, JoinSet, LocalSet};
+
+/// A collection of tasks spawned on a Tokio runtime, associated with hash map
+/// keys.
+///
+/// This type is very similar to the [`JoinSet`] type in `tokio::task`, with the
+/// addition of a  set of keys associated with each task. These keys allow
+/// [cancelling a task][abort] or [multiple tasks][abort_matching] in the
+/// `JoinMap` based on   their keys, or [test whether a task corresponding to a
+/// given key exists][contains] in the `JoinMap`.
+///
+/// In addition, when tasks in the `JoinMap` complete, they will return the
+/// associated key along with the value returned by the task, if any.
+///
+/// A `JoinMap` can be used to await the completion of some or all of the tasks
+/// in the map. The map is not ordered, and the tasks will be returned in the
+/// order they complete.
+///
+/// All of the tasks must have the same return type `V`.
+///
+/// When the `JoinMap` is dropped, all tasks in the `JoinMap` are immediately aborted.
+///
+/// **Note**: This type depends on Tokio's [unstable API][unstable]. See [the
+/// documentation on unstable features][unstable] for details on how to enable
+/// Tokio's unstable features.
+///
+/// # Examples
+///
+/// Spawn multiple tasks and wait for them:
+///
+/// ```
+/// use tokio_util::task::JoinMap;
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let mut map = JoinMap::new();
+///
+///     for i in 0..10 {
+///         // Spawn a task on the `JoinMap` with `i` as its key.
+///         map.spawn(i, async move { /* ... */ });
+///     }
+///
+///     let mut seen = [false; 10];
+///
+///     // When a task completes, `join_next` returns the task's key along
+///     // with its output.
+///     while let Some((key, res)) = map.join_next().await {
+///         seen[key] = true;
+///         assert!(res.is_ok(), "task {} completed successfully!", key);
+///     }
+///
+///     for i in 0..10 {
+///         assert!(seen[i]);
+///     }
+/// }
+/// ```
+///
+/// Cancel tasks based on their keys:
+///
+/// ```
+/// use tokio_util::task::JoinMap;
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let mut map = JoinMap::new();
+///
+///     map.spawn("hello world", async move { /* ... */ });
+///     map.spawn("goodbye world", async move { /* ... */});
+///
+///     // Look up the "goodbye world" task in the map and abort it.
+///     let aborted = map.abort("goodbye world");
+///
+///     // `JoinMap::abort` returns `true` if a task existed for the
+///     // provided key.
+///     assert!(aborted);
+///
+///     while let Some((key, res)) = map.join_next().await {
+///         if key == "goodbye world" {
+///             // The aborted task should complete with a cancelled `JoinError`.
+///             assert!(res.unwrap_err().is_cancelled());
+///         } else {
+///             // Other tasks should complete normally.
+///             assert!(res.is_ok());
+///         }
+///     }
+/// }
+/// ```
+///
+/// [`JoinSet`]: tokio::task::JoinSet
+/// [unstable]: tokio#unstable-features
+/// [abort]: fn@Self::abort
+/// [abort_matching]: fn@Self::abort_matching
+/// [contains]: fn@Self::contains_key
+#[cfg_attr(docsrs, doc(cfg(all(feature = "rt", tokio_unstable))))]
+pub struct JoinMap<K, V, S = RandomState> {
+    /// A map of the [`AbortHandle`]s of the tasks spawned on this `JoinMap`,
+    /// indexed by their keys and task IDs.
+    ///
+    /// The [`Key`] type contains both the task's `K`-typed key provided when
+    /// spawning tasks, and the task's IDs. The IDs are stored here to resolve
+    /// hash collisions when looking up tasks based on their pre-computed hash
+    /// (as stored in the `hashes_by_task` map).
+    tasks_by_key: HashMap<Key<K>, AbortHandle, S>,
+
+    /// A map from task IDs to the hash of the key associated with that task.
+    ///
+    /// This map is used to perform reverse lookups of tasks in the
+    /// `tasks_by_key` map based on their task IDs. When a task terminates, the
+    /// ID is provided to us by the `JoinSet`, so we can look up the hash value
+    /// of that task's key, and then remove it from the `tasks_by_key` map using
+    /// the raw hash code, resolving collisions by comparing task IDs.
+    hashes_by_task: HashMap<Id, u64, S>,
+
+    /// The [`JoinSet`] that awaits the completion of tasks spawned on this
+    /// `JoinMap`.
+    tasks: JoinSet<V>,
+}
+
+/// A [`JoinMap`] key.
+///
+/// This holds both a `K`-typed key (the actual key as seen by the user), _and_
+/// a task ID, so that hash collisions between `K`-typed keys can be resolved
+/// using either `K`'s `Eq` impl *or* by checking the task IDs.
+///
+/// This allows looking up a task using either an actual key (such as when the
+/// user queries the map with a key), *or* using a task ID and a hash (such as
+/// when removing completed tasks from the map).
+#[derive(Debug)]
+struct Key<K> {
+    key: K,
+    id: Id,
+}
+
+impl<K, V> JoinMap<K, V> {
+    /// Creates a new empty `JoinMap`.
+    ///
+    /// The `JoinMap` is initially created with a capacity of 0, so it will not
+    /// allocate until a task is first spawned on it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    /// let map: JoinMap<&str, i32> = JoinMap::new();
+    /// ```
+    #[inline]
+    #[must_use]
+    pub fn new() -> Self {
+        Self::with_hasher(RandomState::new())
+    }
+
+    /// Creates an empty `JoinMap` with the specified capacity.
+    ///
+    /// The `JoinMap` will be able to hold at least `capacity` tasks without
+    /// reallocating.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    /// let map: JoinMap<&str, i32> = JoinMap::with_capacity(10);
+    /// ```
+    #[inline]
+    #[must_use]
+    pub fn with_capacity(capacity: usize) -> Self {
+        JoinMap::with_capacity_and_hasher(capacity, Default::default())
+    }
+}
+
+impl<K, V, S: Clone> JoinMap<K, V, S> {
+    /// Creates an empty `JoinMap` which will use the given hash builder to hash
+    /// keys.
+    ///
+    /// The created map has the default initial capacity.
+    ///
+    /// Warning: `hash_builder` is normally randomly generated, and
+    /// is designed to allow `JoinMap` to be resistant to attacks that
+    /// cause many collisions and very poor performance. Setting it
+    /// manually using this function can expose a DoS attack vector.
+    ///
+    /// The `hash_builder` passed should implement the [`BuildHasher`] trait for
+    /// the `JoinMap` to be useful, see its documentation for details.
+    #[inline]
+    #[must_use]
+    pub fn with_hasher(hash_builder: S) -> Self {
+        Self::with_capacity_and_hasher(0, hash_builder)
+    }
+
+    /// Creates an empty `JoinMap` with the specified capacity, using `hash_builder`
+    /// to hash the keys.
+    ///
+    /// The `JoinMap` will be able to hold at least `capacity` elements without
+    /// reallocating. If `capacity` is 0, the `JoinMap` will not allocate.
+    ///
+    /// Warning: `hash_builder` is normally randomly generated, and
+    /// is designed to allow HashMaps to be resistant to attacks that
+    /// cause many collisions and very poor performance. Setting it
+    /// manually using this function can expose a DoS attack vector.
+    ///
+    /// The `hash_builder` passed should implement the [`BuildHasher`] trait for
+    /// the `JoinMap`to be useful, see its documentation for details.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// use tokio_util::task::JoinMap;
+    /// use std::collections::hash_map::RandomState;
+    ///
+    /// let s = RandomState::new();
+    /// let mut map = JoinMap::with_capacity_and_hasher(10, s);
+    /// map.spawn(1, async move { "hello world!" });
+    /// # }
+    /// ```
+    #[inline]
+    #[must_use]
+    pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
+        Self {
+            tasks_by_key: HashMap::with_capacity_and_hasher(capacity, hash_builder.clone()),
+            hashes_by_task: HashMap::with_capacity_and_hasher(capacity, hash_builder),
+            tasks: JoinSet::new(),
+        }
+    }
+
+    /// Returns the number of tasks currently in the `JoinMap`.
+    pub fn len(&self) -> usize {
+        let len = self.tasks_by_key.len();
+        debug_assert_eq!(len, self.hashes_by_task.len());
+        len
+    }
+
+    /// Returns whether the `JoinMap` is empty.
+    pub fn is_empty(&self) -> bool {
+        let empty = self.tasks_by_key.is_empty();
+        debug_assert_eq!(empty, self.hashes_by_task.is_empty());
+        empty
+    }
+
+    /// Returns the number of tasks the map can hold without reallocating.
+    ///
+    /// This number is a lower bound; the `JoinMap` might be able to hold
+    /// more, but is guaranteed to be able to hold at least this many.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// let map: JoinMap<i32, i32> = JoinMap::with_capacity(100);
+    /// assert!(map.capacity() >= 100);
+    /// ```
+    #[inline]
+    pub fn capacity(&self) -> usize {
+        let capacity = self.tasks_by_key.capacity();
+        debug_assert_eq!(capacity, self.hashes_by_task.capacity());
+        capacity
+    }
+}
+
+impl<K, V, S> JoinMap<K, V, S>
+where
+    K: Hash + Eq,
+    V: 'static,
+    S: BuildHasher,
+{
+    /// Spawn the provided task and store it in this `JoinMap` with the provided
+    /// key.
+    ///
+    /// If a task previously existed in the `JoinMap` for this key, that task
+    /// will be cancelled and replaced with the new one. The previous task will
+    /// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
+    /// *not* return a cancelled [`JoinError`] for that task.
+    ///
+    /// # Panics
+    ///
+    /// This method panics if called outside of a Tokio runtime.
+    ///
+    /// [`join_next`]: Self::join_next
+    #[track_caller]
+    pub fn spawn<F>(&mut self, key: K, task: F)
+    where
+        F: Future<Output = V>,
+        F: Send + 'static,
+        V: Send,
+    {
+        let task = self.tasks.spawn(task);
+        self.insert(key, task)
+    }
+
+    /// Spawn the provided task on the provided runtime and store it in this
+    /// `JoinMap` with the provided key.
+    ///
+    /// If a task previously existed in the `JoinMap` for this key, that task
+    /// will be cancelled and replaced with the new one. The previous task will
+    /// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
+    /// *not* return a cancelled [`JoinError`] for that task.
+    ///
+    /// [`join_next`]: Self::join_next
+    #[track_caller]
+    pub fn spawn_on<F>(&mut self, key: K, task: F, handle: &Handle)
+    where
+        F: Future<Output = V>,
+        F: Send + 'static,
+        V: Send,
+    {
+        let task = self.tasks.spawn_on(task, handle);
+        self.insert(key, task);
+    }
+
+    /// Spawn the provided task on the current [`LocalSet`] and store it in this
+    /// `JoinMap` with the provided key.
+    ///
+    /// If a task previously existed in the `JoinMap` for this key, that task
+    /// will be cancelled and replaced with the new one. The previous task will
+    /// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
+    /// *not* return a cancelled [`JoinError`] for that task.
+    ///
+    /// # Panics
+    ///
+    /// This method panics if it is called outside of a `LocalSet`.
+    ///
+    /// [`LocalSet`]: tokio::task::LocalSet
+    /// [`join_next`]: Self::join_next
+    #[track_caller]
+    pub fn spawn_local<F>(&mut self, key: K, task: F)
+    where
+        F: Future<Output = V>,
+        F: 'static,
+    {
+        let task = self.tasks.spawn_local(task);
+        self.insert(key, task);
+    }
+
+    /// Spawn the provided task on the provided [`LocalSet`] and store it in
+    /// this `JoinMap` with the provided key.
+    ///
+    /// If a task previously existed in the `JoinMap` for this key, that task
+    /// will be cancelled and replaced with the new one. The previous task will
+    /// be removed from the `JoinMap`; a subsequent call to [`join_next`] will
+    /// *not* return a cancelled [`JoinError`] for that task.
+    ///
+    /// [`LocalSet`]: tokio::task::LocalSet
+    /// [`join_next`]: Self::join_next
+    #[track_caller]
+    pub fn spawn_local_on<F>(&mut self, key: K, task: F, local_set: &LocalSet)
+    where
+        F: Future<Output = V>,
+        F: 'static,
+    {
+        let task = self.tasks.spawn_local_on(task, local_set);
+        self.insert(key, task)
+    }
+
+    fn insert(&mut self, key: K, abort: AbortHandle) {
+        let hash = self.hash(&key);
+        let id = abort.id();
+        let map_key = Key {
+            id: id.clone(),
+            key,
+        };
+
+        // Insert the new key into the map of tasks by keys.
+        let entry = self
+            .tasks_by_key
+            .raw_entry_mut()
+            .from_hash(hash, |k| k.key == map_key.key);
+        match entry {
+            RawEntryMut::Occupied(mut occ) => {
+                // There was a previous task spawned with the same key! Cancel
+                // that task, and remove its ID from the map of hashes by task IDs.
+                let Key { id: prev_id, .. } = occ.insert_key(map_key);
+                occ.insert(abort).abort();
+                let _prev_hash = self.hashes_by_task.remove(&prev_id);
+                debug_assert_eq!(Some(hash), _prev_hash);
+            }
+            RawEntryMut::Vacant(vac) => {
+                vac.insert(map_key, abort);
+            }
+        };
+
+        // Associate the key's hash with this task's ID, for looking up tasks by ID.
+        let _prev = self.hashes_by_task.insert(id, hash);
+        debug_assert!(_prev.is_none(), "no prior task should have had the same ID");
+    }
+
+    /// Waits until one of the tasks in the map completes and returns its
+    /// output, along with the key corresponding to that task.
+    ///
+    /// Returns `None` if the map is empty.
+    ///
+    /// # Cancel Safety
+    ///
+    /// This method is cancel safe. If `join_next` is used as the event in a [`tokio::select!`]
+    /// statement and some other branch completes first, it is guaranteed that no tasks were
+    /// removed from this `JoinMap`.
+    ///
+    /// # Returns
+    ///
+    /// This function returns:
+    ///
+    ///  * `Some((key, Ok(value)))` if one of the tasks in this `JoinMap` has
+    ///    completed. The `value` is the return value of that ask, and `key` is
+    ///    the key associated with the task.
+    ///  * `Some((key, Err(err))` if one of the tasks in this JoinMap` has
+    ///    panicked or been aborted. `key` is the key associated  with the task
+    ///    that panicked or was aborted.
+    ///  * `None` if the `JoinMap` is empty.
+    ///
+    /// [`tokio::select!`]: tokio::select
+    #[doc(alias = "join_one")]
+    pub async fn join_next(&mut self) -> Option<(K, Result<V, JoinError>)> {
+        let (res, id) = match self.tasks.join_next_with_id().await {
+            Some(Ok((id, output))) => (Ok(output), id),
+            Some(Err(e)) => {
+                let id = e.id();
+                (Err(e), id)
+            }
+            None => return None,
+        };
+        let key = self.remove_by_id(id)?;
+        Some((key, res))
+    }
+
+    #[doc(hidden)]
+    #[deprecated(since = "0.7.4", note = "renamed to `JoinMap::join_next`.")]
+    pub async fn join_one(&mut self) -> Option<(K, Result<V, JoinError>)> {
+        self.join_next().await
+    }
+
+    /// Aborts all tasks and waits for them to finish shutting down.
+    ///
+    /// Calling this method is equivalent to calling [`abort_all`] and then calling [`join_next`] in
+    /// a loop until it returns `None`.
+    ///
+    /// This method ignores any panics in the tasks shutting down. When this call returns, the
+    /// `JoinMap` will be empty.
+    ///
+    /// [`abort_all`]: fn@Self::abort_all
+    /// [`join_next`]: fn@Self::join_next
+    pub async fn shutdown(&mut self) {
+        self.abort_all();
+        while self.join_next().await.is_some() {}
+    }
+
+    /// Abort the task corresponding to the provided `key`.
+    ///
+    /// If this `JoinMap` contains a task corresponding to `key`, this method
+    /// will abort that task and return `true`. Otherwise, if no task exists for
+    /// `key`, this method returns `false`.
+    ///
+    /// # Examples
+    ///
+    /// Aborting a task by key:
+    ///
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let mut map = JoinMap::new();
+    ///
+    /// map.spawn("hello world", async move { /* ... */ });
+    /// map.spawn("goodbye world", async move { /* ... */});
+    ///
+    /// // Look up the "goodbye world" task in the map and abort it.
+    /// map.abort("goodbye world");
+    ///
+    /// while let Some((key, res)) = map.join_next().await {
+    ///     if key == "goodbye world" {
+    ///         // The aborted task should complete with a cancelled `JoinError`.
+    ///         assert!(res.unwrap_err().is_cancelled());
+    ///     } else {
+    ///         // Other tasks should complete normally.
+    ///         assert!(res.is_ok());
+    ///     }
+    /// }
+    /// # }
+    /// ```
+    ///
+    /// `abort` returns `true` if a task was aborted:
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let mut map = JoinMap::new();
+    ///
+    /// map.spawn("hello world", async move { /* ... */ });
+    /// map.spawn("goodbye world", async move { /* ... */});
+    ///
+    /// // A task for the key "goodbye world" should exist in the map:
+    /// assert!(map.abort("goodbye world"));
+    ///
+    /// // Aborting a key that does not exist will return `false`:
+    /// assert!(!map.abort("goodbye universe"));
+    /// # }
+    /// ```
+    pub fn abort<Q: ?Sized>(&mut self, key: &Q) -> bool
+    where
+        Q: Hash + Eq,
+        K: Borrow<Q>,
+    {
+        match self.get_by_key(key) {
+            Some((_, handle)) => {
+                handle.abort();
+                true
+            }
+            None => false,
+        }
+    }
+
+    /// Aborts all tasks with keys matching `predicate`.
+    ///
+    /// `predicate` is a function called with a reference to each key in the
+    /// map. If it returns `true` for a given key, the corresponding task will
+    /// be cancelled.
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// # // use the current thread rt so that spawned tasks don't
+    /// # // complete in the background before they can be aborted.
+    /// # #[tokio::main(flavor = "current_thread")]
+    /// # async fn main() {
+    /// let mut map = JoinMap::new();
+    ///
+    /// map.spawn("hello world", async move {
+    ///     // ...
+    ///     # tokio::task::yield_now().await; // don't complete immediately, get aborted!
+    /// });
+    /// map.spawn("goodbye world", async move {
+    ///     // ...
+    ///     # tokio::task::yield_now().await; // don't complete immediately, get aborted!
+    /// });
+    /// map.spawn("hello san francisco", async move {
+    ///     // ...
+    ///     # tokio::task::yield_now().await; // don't complete immediately, get aborted!
+    /// });
+    /// map.spawn("goodbye universe", async move {
+    ///     // ...
+    ///     # tokio::task::yield_now().await; // don't complete immediately, get aborted!
+    /// });
+    ///
+    /// // Abort all tasks whose keys begin with "goodbye"
+    /// map.abort_matching(|key| key.starts_with("goodbye"));
+    ///
+    /// let mut seen = 0;
+    /// while let Some((key, res)) = map.join_next().await {
+    ///     seen += 1;
+    ///     if key.starts_with("goodbye") {
+    ///         // The aborted task should complete with a cancelled `JoinError`.
+    ///         assert!(res.unwrap_err().is_cancelled());
+    ///     } else {
+    ///         // Other tasks should complete normally.
+    ///         assert!(key.starts_with("hello"));
+    ///         assert!(res.is_ok());
+    ///     }
+    /// }
+    ///
+    /// // All spawned tasks should have completed.
+    /// assert_eq!(seen, 4);
+    /// # }
+    /// ```
+    pub fn abort_matching(&mut self, mut predicate: impl FnMut(&K) -> bool) {
+        // Note: this method iterates over the tasks and keys *without* removing
+        // any entries, so that the keys from aborted tasks can still be
+        // returned when calling `join_next` in the future.
+        for (Key { ref key, .. }, task) in &self.tasks_by_key {
+            if predicate(key) {
+                task.abort();
+            }
+        }
+    }
+
+    /// Returns `true` if this `JoinMap` contains a task for the provided key.
+    ///
+    /// If the task has completed, but its output hasn't yet been consumed by a
+    /// call to [`join_next`], this method will still return `true`.
+    ///
+    /// [`join_next`]: fn@Self::join_next
+    pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool
+    where
+        Q: Hash + Eq,
+        K: Borrow<Q>,
+    {
+        self.get_by_key(key).is_some()
+    }
+
+    /// Returns `true` if this `JoinMap` contains a task with the provided
+    /// [task ID].
+    ///
+    /// If the task has completed, but its output hasn't yet been consumed by a
+    /// call to [`join_next`], this method will still return `true`.
+    ///
+    /// [`join_next`]: fn@Self::join_next
+    /// [task ID]: tokio::task::Id
+    pub fn contains_task(&self, task: &Id) -> bool {
+        self.get_by_id(task).is_some()
+    }
+
+    /// Reserves capacity for at least `additional` more tasks to be spawned
+    /// on this `JoinMap` without reallocating for the map of task keys. The
+    /// collection may reserve more space to avoid frequent reallocations.
+    ///
+    /// Note that spawning a task will still cause an allocation for the task
+    /// itself.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the new allocation size overflows [`usize`].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// let mut map: JoinMap<&str, i32> = JoinMap::new();
+    /// map.reserve(10);
+    /// ```
+    #[inline]
+    pub fn reserve(&mut self, additional: usize) {
+        self.tasks_by_key.reserve(additional);
+        self.hashes_by_task.reserve(additional);
+    }
+
+    /// Shrinks the capacity of the `JoinMap` as much as possible. It will drop
+    /// down as much as possible while maintaining the internal rules
+    /// and possibly leaving some space in accordance with the resize policy.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// let mut map: JoinMap<i32, i32> = JoinMap::with_capacity(100);
+    /// map.spawn(1, async move { 2 });
+    /// map.spawn(3, async move { 4 });
+    /// assert!(map.capacity() >= 100);
+    /// map.shrink_to_fit();
+    /// assert!(map.capacity() >= 2);
+    /// # }
+    /// ```
+    #[inline]
+    pub fn shrink_to_fit(&mut self) {
+        self.hashes_by_task.shrink_to_fit();
+        self.tasks_by_key.shrink_to_fit();
+    }
+
+    /// Shrinks the capacity of the map with a lower limit. It will drop
+    /// down no lower than the supplied limit while maintaining the internal rules
+    /// and possibly leaving some space in accordance with the resize policy.
+    ///
+    /// If the current capacity is less than the lower limit, this is a no-op.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// use tokio_util::task::JoinMap;
+    ///
+    /// let mut map: JoinMap<i32, i32> = JoinMap::with_capacity(100);
+    /// map.spawn(1, async move { 2 });
+    /// map.spawn(3, async move { 4 });
+    /// assert!(map.capacity() >= 100);
+    /// map.shrink_to(10);
+    /// assert!(map.capacity() >= 10);
+    /// map.shrink_to(0);
+    /// assert!(map.capacity() >= 2);
+    /// # }
+    /// ```
+    #[inline]
+    pub fn shrink_to(&mut self, min_capacity: usize) {
+        self.hashes_by_task.shrink_to(min_capacity);
+        self.tasks_by_key.shrink_to(min_capacity)
+    }
+
+    /// Look up a task in the map by its key, returning the key and abort handle.
+    fn get_by_key<'map, Q: ?Sized>(&'map self, key: &Q) -> Option<(&'map Key<K>, &'map AbortHandle)>
+    where
+        Q: Hash + Eq,
+        K: Borrow<Q>,
+    {
+        let hash = self.hash(key);
+        self.tasks_by_key
+            .raw_entry()
+            .from_hash(hash, |k| k.key.borrow() == key)
+    }
+
+    /// Look up a task in the map by its task ID, returning the key and abort handle.
+    fn get_by_id<'map>(&'map self, id: &Id) -> Option<(&'map Key<K>, &'map AbortHandle)> {
+        let hash = self.hashes_by_task.get(id)?;
+        self.tasks_by_key
+            .raw_entry()
+            .from_hash(*hash, |k| &k.id == id)
+    }
+
+    /// Remove a task from the map by ID, returning the key for that task.
+    fn remove_by_id(&mut self, id: Id) -> Option<K> {
+        // Get the hash for the given ID.
+        let hash = self.hashes_by_task.remove(&id)?;
+
+        // Remove the entry for that hash.
+        let entry = self
+            .tasks_by_key
+            .raw_entry_mut()
+            .from_hash(hash, |k| k.id == id);
+        let (Key { id: _key_id, key }, handle) = match entry {
+            RawEntryMut::Occupied(entry) => entry.remove_entry(),
+            _ => return None,
+        };
+        debug_assert_eq!(_key_id, id);
+        debug_assert_eq!(id, handle.id());
+        self.hashes_by_task.remove(&id);
+        Some(key)
+    }
+
+    /// Returns the hash for a given key.
+    #[inline]
+    fn hash<Q: ?Sized>(&self, key: &Q) -> u64
+    where
+        Q: Hash,
+    {
+        let mut hasher = self.tasks_by_key.hasher().build_hasher();
+        key.hash(&mut hasher);
+        hasher.finish()
+    }
+}
+
+impl<K, V, S> JoinMap<K, V, S>
+where
+    V: 'static,
+{
+    /// Aborts all tasks on this `JoinMap`.
+    ///
+    /// This does not remove the tasks from the `JoinMap`. To wait for the tasks to complete
+    /// cancellation, you should call `join_next` in a loop until the `JoinMap` is empty.
+    pub fn abort_all(&mut self) {
+        self.tasks.abort_all()
+    }
+
+    /// Removes all tasks from this `JoinMap` without aborting them.
+    ///
+    /// The tasks removed by this call will continue to run in the background even if the `JoinMap`
+    /// is dropped. They may still be aborted by key.
+    pub fn detach_all(&mut self) {
+        self.tasks.detach_all();
+        self.tasks_by_key.clear();
+        self.hashes_by_task.clear();
+    }
+}
+
+// Hand-written `fmt::Debug` implementation in order to avoid requiring `V:
+// Debug`, since no value is ever actually stored in the map.
+impl<K: fmt::Debug, V, S> fmt::Debug for JoinMap<K, V, S> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // format the task keys and abort handles a little nicer by just
+        // printing the key and task ID pairs, without format the `Key` struct
+        // itself or the `AbortHandle`, which would just format the task's ID
+        // again.
+        struct KeySet<'a, K: fmt::Debug, S>(&'a HashMap<Key<K>, AbortHandle, S>);
+        impl<K: fmt::Debug, S> fmt::Debug for KeySet<'_, K, S> {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                f.debug_map()
+                    .entries(self.0.keys().map(|Key { key, id }| (key, id)))
+                    .finish()
+            }
+        }
+
+        f.debug_struct("JoinMap")
+            // The `tasks_by_key` map is the only one that contains information
+            // that's really worth formatting for the user, since it contains
+            // the tasks' keys and IDs. The other fields are basically
+            // implementation details.
+            .field("tasks", &KeySet(&self.tasks_by_key))
+            .finish()
+    }
+}
+
+impl<K, V> Default for JoinMap<K, V> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+// === impl Key ===
+
+impl<K: Hash> Hash for Key<K> {
+    // Don't include the task ID in the hash.
+    #[inline]
+    fn hash<H: Hasher>(&self, hasher: &mut H) {
+        self.key.hash(hasher);
+    }
+}
+
+// Because we override `Hash` for this type, we must also override the
+// `PartialEq` impl, so that all instances with the same hash are equal.
+impl<K: PartialEq> PartialEq for Key<K> {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.key == other.key
+    }
+}
+
+impl<K: Eq> Eq for Key<K> {}

tokio-util/src/task/mod.rs

@@ -0,0 +1,10 @@
+//! Extra utilities for spawning tasks
+
+#[cfg(tokio_unstable)]
+mod join_map;
+mod spawn_pinned;
+pub use spawn_pinned::LocalPoolHandle;
+
+#[cfg(tokio_unstable)]
+#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
+pub use join_map::JoinMap;

tokio-util/src/task/spawn_pinned.rs

@@ -0,0 +1,436 @@
+use futures_util::future::{AbortHandle, Abortable};
+use std::fmt;
+use std::fmt::{Debug, Formatter};
+use std::future::Future;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use tokio::runtime::Builder;
+use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
+use tokio::sync::oneshot;
+use tokio::task::{spawn_local, JoinHandle, LocalSet};
+
+/// A cloneable handle to a local pool, used for spawning `!Send` tasks.
+///
+/// Internally the local pool uses a [`tokio::task::LocalSet`] for each worker thread
+/// in the pool. Consequently you can also use [`tokio::task::spawn_local`] (which will
+/// execute on the same thread) inside the Future you supply to the various spawn methods
+/// of `LocalPoolHandle`,
+///
+/// [`tokio::task::LocalSet`]: tokio::task::LocalSet
+/// [`tokio::task::spawn_local`]: tokio::task::spawn_local
+///
+/// # Examples
+///
+/// ```
+/// use std::rc::Rc;
+/// use tokio::{self, task };
+/// use tokio_util::task::LocalPoolHandle;
+///
+/// #[tokio::main(flavor = "current_thread")]
+/// async fn main() {
+///     let pool = LocalPoolHandle::new(5);
+///
+///     let output = pool.spawn_pinned(|| {
+///         // `data` is !Send + !Sync
+///         let data = Rc::new("local data");
+///         let data_clone = data.clone();
+///
+///         async move {
+///             task::spawn_local(async move {
+///                 println!("{}", data_clone);
+///             });
+///     
+///             data.to_string()
+///         }   
+///     }).await.unwrap();
+///     println!("output: {}", output);
+/// }
+/// ```
+///
+#[derive(Clone)]
+pub struct LocalPoolHandle {
+    pool: Arc<LocalPool>,
+}
+
+impl LocalPoolHandle {
+    /// Create a new pool of threads to handle `!Send` tasks. Spawn tasks onto this
+    /// pool via [`LocalPoolHandle::spawn_pinned`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if the pool size is less than one.
+    #[track_caller]
+    pub fn new(pool_size: usize) -> LocalPoolHandle {
+        assert!(pool_size > 0);
+
+        let workers = (0..pool_size)
+            .map(|_| LocalWorkerHandle::new_worker())
+            .collect();
+
+        let pool = Arc::new(LocalPool { workers });
+
+        LocalPoolHandle { pool }
+    }
+
+    /// Returns the number of threads of the Pool.
+    #[inline]
+    pub fn num_threads(&self) -> usize {
+        self.pool.workers.len()
+    }
+
+    /// Returns the number of tasks scheduled on each worker. The indices of the
+    /// worker threads correspond to the indices of the returned `Vec`.
+    pub fn get_task_loads_for_each_worker(&self) -> Vec<usize> {
+        self.pool
+            .workers
+            .iter()
+            .map(|worker| worker.task_count.load(Ordering::SeqCst))
+            .collect::<Vec<_>>()
+    }
+
+    /// Spawn a task onto a worker thread and pin it there so it can't be moved
+    /// off of the thread. Note that the future is not [`Send`], but the
+    /// [`FnOnce`] which creates it is.
+    ///
+    /// # Examples
+    /// ```
+    /// use std::rc::Rc;
+    /// use tokio_util::task::LocalPoolHandle;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     // Create the local pool
+    ///     let pool = LocalPoolHandle::new(1);
+    ///
+    ///     // Spawn a !Send future onto the pool and await it
+    ///     let output = pool
+    ///         .spawn_pinned(|| {
+    ///             // Rc is !Send + !Sync
+    ///             let local_data = Rc::new("test");
+    ///
+    ///             // This future holds an Rc, so it is !Send
+    ///             async move { local_data.to_string() }
+    ///         })
+    ///         .await
+    ///         .unwrap();
+    ///
+    ///     assert_eq!(output, "test");
+    /// }
+    /// ```
+    pub fn spawn_pinned<F, Fut>(&self, create_task: F) -> JoinHandle<Fut::Output>
+    where
+        F: FnOnce() -> Fut,
+        F: Send + 'static,
+        Fut: Future + 'static,
+        Fut::Output: Send + 'static,
+    {
+        self.pool
+            .spawn_pinned(create_task, WorkerChoice::LeastBurdened)
+    }
+
+    /// Differs from `spawn_pinned` only in that you can choose a specific worker thread
+    /// of the pool, whereas `spawn_pinned` chooses the worker with the smallest
+    /// number of tasks scheduled.
+    ///
+    /// A worker thread is chosen by index. Indices are 0 based and the largest index
+    /// is given by `num_threads() - 1`
+    ///
+    /// # Panics
+    ///
+    /// This method panics if the index is out of bounds.
+    ///
+    /// # Examples
+    ///
+    /// This method can be used to spawn a task on all worker threads of the pool:
+    ///
+    /// ```
+    /// use tokio_util::task::LocalPoolHandle;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     const NUM_WORKERS: usize = 3;
+    ///     let pool = LocalPoolHandle::new(NUM_WORKERS);
+    ///     let handles = (0..pool.num_threads())
+    ///         .map(|worker_idx| {
+    ///             pool.spawn_pinned_by_idx(
+    ///                 || {
+    ///                     async {
+    ///                         "test"
+    ///                     }
+    ///                 },
+    ///                 worker_idx,
+    ///             )
+    ///         })
+    ///         .collect::<Vec<_>>();
+    ///
+    ///     for handle in handles {
+    ///         handle.await.unwrap();
+    ///     }
+    /// }
+    /// ```
+    ///
+    #[track_caller]
+    pub fn spawn_pinned_by_idx<F, Fut>(&self, create_task: F, idx: usize) -> JoinHandle<Fut::Output>
+    where
+        F: FnOnce() -> Fut,
+        F: Send + 'static,
+        Fut: Future + 'static,
+        Fut::Output: Send + 'static,
+    {
+        self.pool
+            .spawn_pinned(create_task, WorkerChoice::ByIdx(idx))
+    }
+}
+
+impl Debug for LocalPoolHandle {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.write_str("LocalPoolHandle")
+    }
+}
+
+enum WorkerChoice {
+    LeastBurdened,
+    ByIdx(usize),
+}
+
+struct LocalPool {
+    workers: Vec<LocalWorkerHandle>,
+}
+
+impl LocalPool {
+    /// Spawn a `?Send` future onto a worker
+    #[track_caller]
+    fn spawn_pinned<F, Fut>(
+        &self,
+        create_task: F,
+        worker_choice: WorkerChoice,
+    ) -> JoinHandle<Fut::Output>
+    where
+        F: FnOnce() -> Fut,
+        F: Send + 'static,
+        Fut: Future + 'static,
+        Fut::Output: Send + 'static,
+    {
+        let (sender, receiver) = oneshot::channel();
+        let (worker, job_guard) = match worker_choice {
+            WorkerChoice::LeastBurdened => self.find_and_incr_least_burdened_worker(),
+            WorkerChoice::ByIdx(idx) => self.find_worker_by_idx(idx),
+        };
+        let worker_spawner = worker.spawner.clone();
+
+        // Spawn a future onto the worker's runtime so we can immediately return
+        // a join handle.
+        worker.runtime_handle.spawn(async move {
+            // Move the job guard into the task
+            let _job_guard = job_guard;
+
+            // Propagate aborts via Abortable/AbortHandle
+            let (abort_handle, abort_registration) = AbortHandle::new_pair();
+            let _abort_guard = AbortGuard(abort_handle);
+
+            // Inside the future we can't run spawn_local yet because we're not
+            // in the context of a LocalSet. We need to send create_task to the
+            // LocalSet task for spawning.
+            let spawn_task = Box::new(move || {
+                // Once we're in the LocalSet context we can call spawn_local
+                let join_handle =
+                    spawn_local(
+                        async move { Abortable::new(create_task(), abort_registration).await },
+                    );
+
+                // Send the join handle back to the spawner. If sending fails,
+                // we assume the parent task was canceled, so cancel this task
+                // as well.
+                if let Err(join_handle) = sender.send(join_handle) {
+                    join_handle.abort()
+                }
+            });
+
+            // Send the callback to the LocalSet task
+            if let Err(e) = worker_spawner.send(spawn_task) {
+                // Propagate the error as a panic in the join handle.
+                panic!("Failed to send job to worker: {}", e);
+            }
+
+            // Wait for the task's join handle
+            let join_handle = match receiver.await {
+                Ok(handle) => handle,
+                Err(e) => {
+                    // We sent the task successfully, but failed to get its
+                    // join handle... We assume something happened to the worker
+                    // and the task was not spawned. Propagate the error as a
+                    // panic in the join handle.
+                    panic!("Worker failed to send join handle: {}", e);
+                }
+            };
+
+            // Wait for the task to complete
+            let join_result = join_handle.await;
+
+            match join_result {
+                Ok(Ok(output)) => output,
+                Ok(Err(_)) => {
+                    // Pinned task was aborted. But that only happens if this
+                    // task is aborted. So this is an impossible branch.
+                    unreachable!(
+                        "Reaching this branch means this task was previously \
+                         aborted but it continued running anyways"
+                    )
+                }
+                Err(e) => {
+                    if e.is_panic() {
+                        std::panic::resume_unwind(e.into_panic());
+                    } else if e.is_cancelled() {
+                        // No one else should have the join handle, so this is
+                        // unexpected. Forward this error as a panic in the join
+                        // handle.
+                        panic!("spawn_pinned task was canceled: {}", e);
+                    } else {
+                        // Something unknown happened (not a panic or
+                        // cancellation). Forward this error as a panic in the
+                        // join handle.
+                        panic!("spawn_pinned task failed: {}", e);
+                    }
+                }
+            }
+        })
+    }
+
+    /// Find the worker with the least number of tasks, increment its task
+    /// count, and return its handle. Make sure to actually spawn a task on
+    /// the worker so the task count is kept consistent with load.
+    ///
+    /// A job count guard is also returned to ensure the task count gets
+    /// decremented when the job is done.
+    fn find_and_incr_least_burdened_worker(&self) -> (&LocalWorkerHandle, JobCountGuard) {
+        loop {
+            let (worker, task_count) = self
+                .workers
+                .iter()
+                .map(|worker| (worker, worker.task_count.load(Ordering::SeqCst)))
+                .min_by_key(|&(_, count)| count)
+                .expect("There must be more than one worker");
+
+            // Make sure the task count hasn't changed since when we choose this
+            // worker. Otherwise, restart the search.
+            if worker
+                .task_count
+                .compare_exchange(
+                    task_count,
+                    task_count + 1,
+                    Ordering::SeqCst,
+                    Ordering::Relaxed,
+                )
+                .is_ok()
+            {
+                return (worker, JobCountGuard(Arc::clone(&worker.task_count)));
+            }
+        }
+    }
+
+    #[track_caller]
+    fn find_worker_by_idx(&self, idx: usize) -> (&LocalWorkerHandle, JobCountGuard) {
+        let worker = &self.workers[idx];
+        worker.task_count.fetch_add(1, Ordering::SeqCst);
+
+        (worker, JobCountGuard(Arc::clone(&worker.task_count)))
+    }
+}
+
+/// Automatically decrements a worker's job count when a job finishes (when
+/// this gets dropped).
+struct JobCountGuard(Arc<AtomicUsize>);
+
+impl Drop for JobCountGuard {
+    fn drop(&mut self) {
+        // Decrement the job count
+        let previous_value = self.0.fetch_sub(1, Ordering::SeqCst);
+        debug_assert!(previous_value >= 1);
+    }
+}
+
+/// Calls abort on the handle when dropped.
+struct AbortGuard(AbortHandle);
+
+impl Drop for AbortGuard {
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
+type PinnedFutureSpawner = Box<dyn FnOnce() + Send + 'static>;
+
+struct LocalWorkerHandle {
+    runtime_handle: tokio::runtime::Handle,
+    spawner: UnboundedSender<PinnedFutureSpawner>,
+    task_count: Arc<AtomicUsize>,
+}
+
+impl LocalWorkerHandle {
+    /// Create a new worker for executing pinned tasks
+    fn new_worker() -> LocalWorkerHandle {
+        let (sender, receiver) = unbounded_channel();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("Failed to start a pinned worker thread runtime");
+        let runtime_handle = runtime.handle().clone();
+        let task_count = Arc::new(AtomicUsize::new(0));
+        let task_count_clone = Arc::clone(&task_count);
+
+        std::thread::spawn(|| Self::run(runtime, receiver, task_count_clone));
+
+        LocalWorkerHandle {
+            runtime_handle,
+            spawner: sender,
+            task_count,
+        }
+    }
+
+    fn run(
+        runtime: tokio::runtime::Runtime,
+        mut task_receiver: UnboundedReceiver<PinnedFutureSpawner>,
+        task_count: Arc<AtomicUsize>,
+    ) {
+        let local_set = LocalSet::new();
+        local_set.block_on(&runtime, async {
+            while let Some(spawn_task) = task_receiver.recv().await {
+                // Calls spawn_local(future)
+                (spawn_task)();
+            }
+        });
+
+        // If there are any tasks on the runtime associated with a LocalSet task
+        // that has already completed, but whose output has not yet been
+        // reported, let that task complete.
+        //
+        // Since the task_count is decremented when the runtime task exits,
+        // reading that counter lets us know if any such tasks completed during
+        // the call to `block_on`.
+        //
+        // Tasks on the LocalSet can't complete during this loop since they're
+        // stored on the LocalSet and we aren't accessing it.
+        let mut previous_task_count = task_count.load(Ordering::SeqCst);
+        loop {
+            // This call will also run tasks spawned on the runtime.
+            runtime.block_on(tokio::task::yield_now());
+            let new_task_count = task_count.load(Ordering::SeqCst);
+            if new_task_count == previous_task_count {
+                break;
+            } else {
+                previous_task_count = new_task_count;
+            }
+        }
+
+        // It's now no longer possible for a task on the runtime to be
+        // associated with a LocalSet task that has completed. Drop both the
+        // LocalSet and runtime to let tasks on the runtime be cancelled if and
+        // only if they are still on the LocalSet.
+        //
+        // Drop the LocalSet task first so that anyone awaiting the runtime
+        // JoinHandle will see the cancelled error after the LocalSet task
+        // destructor has completed.
+        drop(local_set);
+        drop(runtime);
+    }
+}

tokio-util/src/time/delay_queue.rs

@@ -7,10 +7,15 @@
 use crate::time::wheel::{self, Wheel};
 
 use futures_core::ready;
-use tokio::time::{error::Error, sleep_until, Duration, Instant, Sleep};
+use tokio::time::{sleep_until, Duration, Instant, Sleep};
 
+use core::ops::{Index, IndexMut};
 use slab::Slab;
 use std::cmp;
+use std::collections::HashMap;
+use std::convert::From;
+use std::fmt;
+use std::fmt::Debug;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::pin::Pin;
@@ -67,7 +72,6 @@ use std::task::{self, Poll, Waker};
 /// Using `DelayQueue` to manage cache entries.
 ///
 /// ```rust,no_run
-/// use tokio::time::error::Error;
 /// use tokio_util::time::{DelayQueue, delay_queue};
 ///
 /// use futures::ready;
@@ -103,13 +107,12 @@ use std::task::{self, Poll, Waker};
 ///         }
 ///     }
 ///
-///     fn poll_purge(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
-///         while let Some(res) = ready!(self.expirations.poll_expired(cx)) {
-///             let entry = res?;
+///     fn poll_purge(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+///         while let Some(entry) = ready!(self.expirations.poll_expired(cx)) {
 ///             self.entries.remove(entry.get_ref());
 ///         }
 ///
-///         Poll::Ready(Ok(()))
+///         Poll::Ready(())
 ///     }
 /// }
 /// ```
@@ -128,7 +131,7 @@ use std::task::{self, Poll, Waker};
 #[derive(Debug)]
 pub struct DelayQueue<T> {
     /// Stores data associated with entries
-    slab: Slab<Data<T>>,
+    slab: SlabStorage<T>,
 
     /// Lookup structure tracking all delays in the queue
     wheel: Wheel<Stack<T>>,
@@ -152,6 +155,216 @@ pub struct DelayQueue<T> {
     waker: Option<Waker>,
 }
 
+#[derive(Default)]
+struct SlabStorage<T> {
+    inner: Slab<Data<T>>,
+
+    // A `compact` call requires a re-mapping of the `Key`s that were changed
+    // during the `compact` call of the `slab`. Since the keys that were given out
+    // cannot be changed retroactively we need to keep track of these re-mappings.
+    // The keys of `key_map` correspond to the old keys that were given out and
+    // the values to the `Key`s that were re-mapped by the `compact` call.
+    key_map: HashMap<Key, KeyInternal>,
+
+    // Index used to create new keys to hand out.
+    next_key_index: usize,
+
+    // Whether `compact` has been called, necessary in order to decide whether
+    // to include keys in `key_map`.
+    compact_called: bool,
+}
+
+impl<T> SlabStorage<T> {
+    pub(crate) fn with_capacity(capacity: usize) -> SlabStorage<T> {
+        SlabStorage {
+            inner: Slab::with_capacity(capacity),
+            key_map: HashMap::new(),
+            next_key_index: 0,
+            compact_called: false,
+        }
+    }
+
+    // Inserts data into the inner slab and re-maps keys if necessary
+    pub(crate) fn insert(&mut self, val: Data<T>) -> Key {
+        let mut key = KeyInternal::new(self.inner.insert(val));
+        let key_contained = self.key_map.contains_key(&key.into());
+
+        if key_contained {
+            // It's possible that a `compact` call creates capacitiy in `self.inner` in
+            // such a way that a `self.inner.insert` call creates a `key` which was
+            // previously given out during an `insert` call prior to the `compact` call.
+            // If `key` is contained in `self.key_map`, we have encountered this exact situation,
+            // We need to create a new key `key_to_give_out` and include the relation
+            // `key_to_give_out` -> `key` in `self.key_map`.
+            let key_to_give_out = self.create_new_key();
+            assert!(!self.key_map.contains_key(&key_to_give_out.into()));
+            self.key_map.insert(key_to_give_out.into(), key);
+            key = key_to_give_out;
+        } else if self.compact_called {
+            // Include an identity mapping in `self.key_map` in order to allow us to
+            // panic if a key that was handed out is removed more than once.
+            self.key_map.insert(key.into(), key);
+        }
+
+        key.into()
+    }
+
+    // Re-map the key in case compact was previously called.
+    // Note: Since we include identity mappings in key_map after compact was called,
+    // we have information about all keys that were handed out. In the case in which
+    // compact was called and we try to remove a Key that was previously removed
+    // we can detect invalid keys if no key is found in `key_map`. This is necessary
+    // in order to prevent situations in which a previously removed key
+    // corresponds to a re-mapped key internally and which would then be incorrectly
+    // removed from the slab.
+    //
+    // Example to illuminate this problem:
+    //
+    // Let's assume our `key_map` is {1 -> 2, 2 -> 1} and we call remove(1). If we
+    // were to remove 1 again, we would not find it inside `key_map` anymore.
+    // If we were to imply from this that no re-mapping was necessary, we would
+    // incorrectly remove 1 from `self.slab.inner`, which corresponds to the
+    // handed-out key 2.
+    pub(crate) fn remove(&mut self, key: &Key) -> Data<T> {
+        let remapped_key = if self.compact_called {
+            match self.key_map.remove(key) {
+                Some(key_internal) => key_internal,
+                None => panic!("invalid key"),
+            }
+        } else {
+            (*key).into()
+        };
+
+        self.inner.remove(remapped_key.index)
+    }
+
+    pub(crate) fn shrink_to_fit(&mut self) {
+        self.inner.shrink_to_fit();
+        self.key_map.shrink_to_fit();
+    }
+
+    pub(crate) fn compact(&mut self) {
+        if !self.compact_called {
+            for (key, _) in self.inner.iter() {
+                self.key_map.insert(Key::new(key), KeyInternal::new(key));
+            }
+        }
+
+        let mut remapping = HashMap::new();
+        self.inner.compact(|_, from, to| {
+            remapping.insert(from, to);
+            true
+        });
+
+        // At this point `key_map` contains a mapping for every element.
+        for internal_key in self.key_map.values_mut() {
+            if let Some(new_internal_key) = remapping.get(&internal_key.index) {
+                *internal_key = KeyInternal::new(*new_internal_key);
+            }
+        }
+
+        if self.key_map.capacity() > 2 * self.key_map.len() {
+            self.key_map.shrink_to_fit();
+        }
+
+        self.compact_called = true;
+    }
+
+    // Tries to re-map a `Key` that was given out to the user to its
+    // corresponding internal key.
+    fn remap_key(&self, key: &Key) -> Option<KeyInternal> {
+        let key_map = &self.key_map;
+        if self.compact_called {
+            key_map.get(&*key).copied()
+        } else {
+            Some((*key).into())
+        }
+    }
+
+    fn create_new_key(&mut self) -> KeyInternal {
+        while self.key_map.contains_key(&Key::new(self.next_key_index)) {
+            self.next_key_index = self.next_key_index.wrapping_add(1);
+        }
+
+        KeyInternal::new(self.next_key_index)
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    pub(crate) fn capacity(&self) -> usize {
+        self.inner.capacity()
+    }
+
+    pub(crate) fn clear(&mut self) {
+        self.inner.clear();
+        self.key_map.clear();
+        self.compact_called = false;
+    }
+
+    pub(crate) fn reserve(&mut self, additional: usize) {
+        self.inner.reserve(additional);
+
+        if self.compact_called {
+            self.key_map.reserve(additional);
+        }
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    pub(crate) fn contains(&self, key: &Key) -> bool {
+        let remapped_key = self.remap_key(key);
+
+        match remapped_key {
+            Some(internal_key) => self.inner.contains(internal_key.index),
+            None => false,
+        }
+    }
+}
+
+impl<T> fmt::Debug for SlabStorage<T>
+where
+    T: fmt::Debug,
+{
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if fmt.alternate() {
+            fmt.debug_map().entries(self.inner.iter()).finish()
+        } else {
+            fmt.debug_struct("Slab")
+                .field("len", &self.len())
+                .field("cap", &self.capacity())
+                .finish()
+        }
+    }
+}
+
+impl<T> Index<Key> for SlabStorage<T> {
+    type Output = Data<T>;
+
+    fn index(&self, key: Key) -> &Self::Output {
+        let remapped_key = self.remap_key(&key);
+
+        match remapped_key {
+            Some(internal_key) => &self.inner[internal_key.index],
+            None => panic!("Invalid index {}", key.index),
+        }
+    }
+}
+
+impl<T> IndexMut<Key> for SlabStorage<T> {
+    fn index_mut(&mut self, key: Key) -> &mut Data<T> {
+        let remapped_key = self.remap_key(&key);
+
+        match remapped_key {
+            Some(internal_key) => &mut self.inner[internal_key.index],
+            None => panic!("Invalid index {}", key.index),
+        }
+    }
+}
+
 /// An entry in `DelayQueue` that has expired and been removed.
 ///
 /// Values are returned by [`DelayQueue::poll_expired`].
@@ -176,15 +389,23 @@ pub struct Expired<T> {
 ///
 /// [`DelayQueue`]: struct@DelayQueue
 /// [`DelayQueue::insert`]: method@DelayQueue::insert
-#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Key {
     index: usize,
 }
 
+// Whereas `Key` is given out to users that use `DelayQueue`, internally we use
+// `KeyInternal` as the key type in order to make the logic of mapping between keys
+// as a result of `compact` calls clearer.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+struct KeyInternal {
+    index: usize,
+}
+
 #[derive(Debug)]
 struct Stack<T> {
     /// Head of the stack
-    head: Option<usize>,
+    head: Option<Key>,
     _p: PhantomData<fn() -> T>,
 }
 
@@ -201,10 +422,10 @@ struct Data<T> {
     expired: bool,
 
     /// Next entry in the stack
-    next: Option<usize>,
+    next: Option<Key>,
 
     /// Previous entry in the stack
-    prev: Option<usize>,
+    prev: Option<Key>,
 }
 
 /// Maximum number of entries the queue can handle
@@ -253,7 +474,7 @@ impl<T> DelayQueue<T> {
     pub fn with_capacity(capacity: usize) -> DelayQueue<T> {
         DelayQueue {
             wheel: Wheel::new(),
-            slab: Slab::with_capacity(capacity),
+            slab: SlabStorage::with_capacity(capacity),
             expired: Stack::default(),
             delay: None,
             wheel_now: 0,
@@ -310,6 +531,7 @@ impl<T> DelayQueue<T> {
     /// [`reset`]: method@Self::reset
     /// [`Key`]: struct@Key
     /// [type]: #
+    #[track_caller]
     pub fn insert_at(&mut self, value: T, when: Instant) -> Key {
         assert!(self.slab.len() < MAX_ENTRIES, "max entries exceeded");
 
@@ -348,16 +570,13 @@ impl<T> DelayQueue<T> {
             }
         }
 
-        Key::new(key)
+        key
     }
 
     /// Attempts to pull out the next value of the delay queue, registering the
     /// current task for wakeup if the value is not yet available, and returning
     /// `None` if the queue is exhausted.
-    pub fn poll_expired(
-        &mut self,
-        cx: &mut task::Context<'_>,
-    ) -> Poll<Option<Result<Expired<T>, Error>>> {
+    pub fn poll_expired(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Expired<T>>> {
         if !self
             .waker
             .as_ref()
@@ -368,18 +587,16 @@ impl<T> DelayQueue<T> {
         }
 
         let item = ready!(self.poll_idx(cx));
-        Poll::Ready(item.map(|result| {
-            result.map(|idx| {
-                let data = self.slab.remove(idx);
-                debug_assert!(data.next.is_none());
-                debug_assert!(data.prev.is_none());
+        Poll::Ready(item.map(|key| {
+            let data = self.slab.remove(&key);
+            debug_assert!(data.next.is_none());
+            debug_assert!(data.prev.is_none());
 
-                Expired {
-                    key: Key::new(idx),
-                    data: data.inner,
-                    deadline: self.start + Duration::from_millis(data.when),
-                }
-            })
+            Expired {
+                key,
+                data: data.inner,
+                deadline: self.start + Duration::from_millis(data.when),
+            }
         }))
     }
 
@@ -433,11 +650,13 @@ impl<T> DelayQueue<T> {
     /// [`reset`]: method@Self::reset
     /// [`Key`]: struct@Key
     /// [type]: #
+    #[track_caller]
     pub fn insert(&mut self, value: T, timeout: Duration) -> Key {
         self.insert_at(value, Instant::now() + timeout)
     }
 
-    fn insert_idx(&mut self, when: u64, key: usize) {
+    #[track_caller]
+    fn insert_idx(&mut self, when: u64, key: Key) {
         use self::wheel::{InsertError, Stack};
 
         // Register the deadline with the timer wheel
@@ -458,14 +677,15 @@ impl<T> DelayQueue<T> {
     /// # Panics
     ///
     /// Panics if the key is not contained in the expired queue or the wheel.
+    #[track_caller]
     fn remove_key(&mut self, key: &Key) {
         use crate::time::wheel::Stack;
 
         // Special case the `expired` queue
-        if self.slab[key.index].expired {
-            self.expired.remove(&key.index, &mut self.slab);
+        if self.slab[*key].expired {
+            self.expired.remove(key, &mut self.slab);
         } else {
-            self.wheel.remove(&key.index, &mut self.slab);
+            self.wheel.remove(key, &mut self.slab);
         }
     }
 
@@ -497,11 +717,12 @@ impl<T> DelayQueue<T> {
     /// assert_eq!(*item.get_ref(), "foo");
     /// # }
     /// ```
+    #[track_caller]
     pub fn remove(&mut self, key: &Key) -> Expired<T> {
         let prev_deadline = self.next_deadline();
 
         self.remove_key(key);
-        let data = self.slab.remove(key.index);
+        let data = self.slab.remove(key);
 
         let next_deadline = self.next_deadline();
         if prev_deadline != next_deadline {
@@ -553,16 +774,17 @@ impl<T> DelayQueue<T> {
     /// // "foo" is now scheduled to be returned in 10 seconds
     /// # }
     /// ```
+    #[track_caller]
     pub fn reset_at(&mut self, key: &Key, when: Instant) {
         self.remove_key(key);
 
         // Normalize the deadline. Values cannot be set to expire in the past.
         let when = self.normalize_deadline(when);
 
-        self.slab[key.index].when = when;
-        self.slab[key.index].expired = false;
+        self.slab[*key].when = when;
+        self.slab[*key].expired = false;
 
-        self.insert_idx(when, key.index);
+        self.insert_idx(when, *key);
 
         let next_deadline = self.next_deadline();
         if let (Some(ref mut delay), Some(deadline)) = (&mut self.delay, next_deadline) {
@@ -571,6 +793,50 @@ impl<T> DelayQueue<T> {
         }
     }
 
+    /// Shrink the capacity of the slab, which `DelayQueue` uses internally for storage allocation.
+    /// This function is not guaranteed to, and in most cases, won't decrease the capacity of the slab
+    /// to the number of elements still contained in it, because elements cannot be moved to a different
+    /// index. To decrease the capacity to the size of the slab use [`compact`].
+    ///
+    /// This function can take O(n) time even when the capacity cannot be reduced or the allocation is
+    /// shrunk in place. Repeated calls run in O(1) though.
+    ///
+    /// [`compact`]: method@Self::compact
+    pub fn shrink_to_fit(&mut self) {
+        self.slab.shrink_to_fit();
+    }
+
+    /// Shrink the capacity of the slab, which `DelayQueue` uses internally for storage allocation,
+    /// to the number of elements that are contained in it.
+    ///
+    /// This methods runs in O(n).
+    ///
+    /// # Examples
+    ///
+    /// Basic usage
+    ///
+    /// ```rust
+    /// use tokio_util::time::DelayQueue;
+    /// use std::time::Duration;
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let mut delay_queue = DelayQueue::with_capacity(10);
+    ///
+    /// let key1 = delay_queue.insert(5, Duration::from_secs(5));
+    /// let key2 = delay_queue.insert(10, Duration::from_secs(10));
+    /// let key3 = delay_queue.insert(15, Duration::from_secs(15));
+    ///
+    /// delay_queue.remove(&key2);
+    ///
+    /// delay_queue.compact();
+    /// assert_eq!(delay_queue.capacity(), 2);
+    /// # }
+    /// ```
+    pub fn compact(&mut self) {
+        self.slab.compact();
+    }
+
     /// Returns the next time to poll as determined by the wheel
     fn next_deadline(&mut self) -> Option<Instant> {
         self.wheel
@@ -613,6 +879,7 @@ impl<T> DelayQueue<T> {
     /// // "foo"is now scheduled to be returned in 10 seconds
     /// # }
     /// ```
+    #[track_caller]
     pub fn reset(&mut self, key: &Key, timeout: Duration) {
         self.reset_at(key, Instant::now() + timeout);
     }
@@ -718,7 +985,12 @@ impl<T> DelayQueue<T> {
     /// assert!(delay_queue.capacity() >= 11);
     /// # }
     /// ```
+    #[track_caller]
     pub fn reserve(&mut self, additional: usize) {
+        assert!(
+            self.slab.capacity() + additional <= MAX_ENTRIES,
+            "max queue capacity exceeded"
+        );
         self.slab.reserve(additional);
     }
 
@@ -750,13 +1022,13 @@ impl<T> DelayQueue<T> {
     /// should be returned.
     ///
     /// A slot should be returned when the associated deadline has been reached.
-    fn poll_idx(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Result<usize, Error>>> {
+    fn poll_idx(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Key>> {
         use self::wheel::Stack;
 
         let expired = self.expired.pop(&mut self.slab);
 
         if expired.is_some() {
-            return Poll::Ready(expired.map(Ok));
+            return Poll::Ready(expired);
         }
 
         loop {
@@ -776,7 +1048,7 @@ impl<T> DelayQueue<T> {
             self.delay = self.next_deadline().map(|when| Box::pin(sleep_until(when)));
 
             if let Some(idx) = wheel_idx {
-                return Poll::Ready(Some(Ok(idx)));
+                return Poll::Ready(Some(idx));
             }
 
             if self.delay.is_none() {
@@ -808,7 +1080,7 @@ impl<T> Default for DelayQueue<T> {
 impl<T> futures_core::Stream for DelayQueue<T> {
     // DelayQueue seems much more specific, where a user may care that it
     // has reached capacity, so return those errors instead of panicking.
-    type Item = Result<Expired<T>, Error>;
+    type Item = Expired<T>;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
         DelayQueue::poll_expired(self.get_mut(), cx)
@@ -816,9 +1088,9 @@ impl<T> futures_core::Stream for DelayQueue<T> {
 }
 
 impl<T> wheel::Stack for Stack<T> {
-    type Owned = usize;
-    type Borrowed = usize;
-    type Store = Slab<Data<T>>;
+    type Owned = Key;
+    type Borrowed = Key;
+    type Store = SlabStorage<T>;
 
     fn is_empty(&self) -> bool {
         self.head.is_none()
@@ -837,28 +1109,30 @@ impl<T> wheel::Stack for Stack<T> {
         }
 
         store[item].next = old;
-        self.head = Some(item)
+        self.head = Some(item);
     }
 
     fn pop(&mut self, store: &mut Self::Store) -> Option<Self::Owned> {
-        if let Some(idx) = self.head {
-            self.head = store[idx].next;
+        if let Some(key) = self.head {
+            self.head = store[key].next;
 
             if let Some(idx) = self.head {
                 store[idx].prev = None;
             }
 
-            store[idx].next = None;
-            debug_assert!(store[idx].prev.is_none());
+            store[key].next = None;
+            debug_assert!(store[key].prev.is_none());
 
-            Some(idx)
+            Some(key)
         } else {
             None
         }
     }
 
+    #[track_caller]
     fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store) {
-        assert!(store.contains(*item));
+        let key = *item;
+        assert!(store.contains(item));
 
         // Ensure that the entry is in fact contained by the stack
         debug_assert!({
@@ -867,29 +1141,31 @@ impl<T> wheel::Stack for Stack<T> {
             let mut contains = false;
 
             while let Some(idx) = next {
+                let data = &store[idx];
+
                 if idx == *item {
                     debug_assert!(!contains);
                     contains = true;
                 }
 
-                next = store[idx].next;
+                next = data.next;
             }
 
             contains
         });
 
-        if let Some(next) = store[*item].next {
-            store[next].prev = store[*item].prev;
+        if let Some(next) = store[key].next {
+            store[next].prev = store[key].prev;
         }
 
-        if let Some(prev) = store[*item].prev {
-            store[prev].next = store[*item].next;
+        if let Some(prev) = store[key].prev {
+            store[prev].next = store[key].next;
         } else {
-            self.head = store[*item].next;
+            self.head = store[key].next;
         }
 
-        store[*item].next = None;
-        store[*item].prev = None;
+        store[key].next = None;
+        store[key].prev = None;
     }
 
     fn when(item: &Self::Borrowed, store: &Self::Store) -> u64 {
@@ -912,6 +1188,24 @@ impl Key {
     }
 }
 
+impl KeyInternal {
+    pub(crate) fn new(index: usize) -> KeyInternal {
+        KeyInternal { index }
+    }
+}
+
+impl From<Key> for KeyInternal {
+    fn from(item: Key) -> Self {
+        KeyInternal::new(item.index)
+    }
+}
+
+impl From<KeyInternal> for Key {
+    fn from(item: KeyInternal) -> Self {
+        Key::new(item.index)
+    }
+}
+
 impl<T> Expired<T> {
     /// Returns a reference to the inner value.
     pub fn get_ref(&self) -> &T {

tokio-util/src/time/wheel/mod.rs

@@ -6,6 +6,7 @@ mod stack;
 pub(crate) use self::stack::Stack;
 
 use std::borrow::Borrow;
+use std::fmt::Debug;
 use std::usize;
 
 /// Timing wheel implementation.
@@ -117,6 +118,7 @@ where
     }
 
     /// Remove `item` from the timing wheel.
+    #[track_caller]
     pub(crate) fn remove(&mut self, item: &T::Borrowed, store: &mut T::Store) {
         let when = T::when(item, store);
 

tokio-util/src/time/wheel/stack.rs

@@ -1,4 +1,6 @@
 use std::borrow::Borrow;
+use std::cmp::Eq;
+use std::hash::Hash;
 
 /// Abstracts the stack operations needed to track timeouts.
 pub(crate) trait Stack: Default {
@@ -6,7 +8,7 @@ pub(crate) trait Stack: Default {
     type Owned: Borrow<Self::Borrowed>;
 
     /// Borrowed item
-    type Borrowed;
+    type Borrowed: Eq + Hash;
 
     /// Item storage, this allows a slab to be used instead of just the heap
     type Store;

tokio-util/src/udp/frame.rs

@@ -81,17 +81,21 @@ where
             }
 
             // We're out of data. Try and fetch more data to decode
-            let addr = unsafe {
-                // Convert `&mut [MaybeUnit<u8>]` to `&mut [u8]` because we will be
-                // writing to it via `poll_recv_from` and therefore initializing the memory.
-                let buf = &mut *(pin.rd.chunk_mut() as *mut _ as *mut [MaybeUninit<u8>]);
+            let addr = {
+                // Safety: `chunk_mut()` returns a `&mut UninitSlice`, and `UninitSlice` is a
+                // transparent wrapper around `[MaybeUninit<u8>]`.
+                let buf = unsafe { &mut *(pin.rd.chunk_mut() as *mut _ as *mut [MaybeUninit<u8>]) };
                 let mut read = ReadBuf::uninit(buf);
                 let ptr = read.filled().as_ptr();
                 let res = ready!(pin.socket.borrow().poll_recv_from(cx, &mut read));
 
                 assert_eq!(ptr, read.filled().as_ptr());
                 let addr = res?;
-                pin.rd.advance_mut(read.filled().len());
+
+                // Safety: This is guaranteed to be the number of initialized (and read) bytes due
+                // to the invariants provided by `ReadBuf::filled`.
+                unsafe { pin.rd.advance_mut(read.filled().len()) };
+
                 addr
             };
 

tokio-util/tests/codecs.rs

@@ -38,6 +38,9 @@ fn bytes_encoder() {
     codec
         .encode(Bytes::from_static(&[0; INITIAL_CAPACITY + 1]), &mut buf)
         .unwrap();
+    codec
+        .encode(BytesMut::from(&b"hello"[..]), &mut buf)
+        .unwrap();
 }
 
 #[test]

tokio-util/tests/framed.rs

@@ -12,7 +12,10 @@ use std::task::{Context, Poll};
 const INITIAL_CAPACITY: usize = 8 * 1024;
 
 /// Encode and decode u32 values.
-struct U32Codec;
+#[derive(Default)]
+struct U32Codec {
+    read_bytes: usize,
+}
 
 impl Decoder for U32Codec {
     type Item = u32;
@@ -24,6 +27,7 @@ impl Decoder for U32Codec {
         }
 
         let n = buf.split_to(4).get_u32();
+        self.read_bytes += 4;
         Ok(Some(n))
     }
 }
@@ -39,6 +43,38 @@ impl Encoder<u32> for U32Codec {
     }
 }
 
+/// Encode and decode u64 values.
+#[derive(Default)]
+struct U64Codec {
+    read_bytes: usize,
+}
+
+impl Decoder for U64Codec {
+    type Item = u64;
+    type Error = io::Error;
+
+    fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<u64>> {
+        if buf.len() < 8 {
+            return Ok(None);
+        }
+
+        let n = buf.split_to(8).get_u64();
+        self.read_bytes += 8;
+        Ok(Some(n))
+    }
+}
+
+impl Encoder<u64> for U64Codec {
+    type Error = io::Error;
+
+    fn encode(&mut self, item: u64, dst: &mut BytesMut) -> io::Result<()> {
+        // Reserve space
+        dst.reserve(8);
+        dst.put_u64(item);
+        Ok(())
+    }
+}
+
 /// This value should never be used
 struct DontReadIntoThis;
 
@@ -63,18 +99,39 @@ impl tokio::io::AsyncRead for DontReadIntoThis {
 
 #[tokio::test]
 async fn can_read_from_existing_buf() {
-    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec);
+    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec::default());
     parts.read_buf = BytesMut::from(&[0, 0, 0, 42][..]);
 
     let mut framed = Framed::from_parts(parts);
     let num = assert_ok!(framed.next().await.unwrap());
 
     assert_eq!(num, 42);
+    assert_eq!(framed.codec().read_bytes, 4);
+}
+
+#[tokio::test]
+async fn can_read_from_existing_buf_after_codec_changed() {
+    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec::default());
+    parts.read_buf = BytesMut::from(&[0, 0, 0, 42, 0, 0, 0, 0, 0, 0, 0, 84][..]);
+
+    let mut framed = Framed::from_parts(parts);
+    let num = assert_ok!(framed.next().await.unwrap());
+
+    assert_eq!(num, 42);
+    assert_eq!(framed.codec().read_bytes, 4);
+
+    let mut framed = framed.map_codec(|codec| U64Codec {
+        read_bytes: codec.read_bytes,
+    });
+    let num = assert_ok!(framed.next().await.unwrap());
+
+    assert_eq!(num, 84);
+    assert_eq!(framed.codec().read_bytes, 12);
 }
 
 #[test]
 fn external_buf_grows_to_init() {
-    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec);
+    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec::default());
     parts.read_buf = BytesMut::from(&[0, 0, 0, 42][..]);
 
     let framed = Framed::from_parts(parts);
@@ -85,7 +142,7 @@ fn external_buf_grows_to_init() {
 
 #[test]
 fn external_buf_does_not_shrink() {
-    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec);
+    let mut parts = FramedParts::new(DontReadIntoThis, U32Codec::default());
     parts.read_buf = BytesMut::from(&vec![0; INITIAL_CAPACITY * 2][..]);
 
     let framed = Framed::from_parts(parts);