├── .gitignore
├── CODE_OF_CONDUCT.md
├── src
    ├── lib.rs
    ├── io.rs
    ├── clock.rs
    └── limiter.rs
├── benches
    └── standard_clock.rs
├── LICENSE-MIT
├── .github
    └── workflows
    │   └── rust.yml
├── Cargo.toml
├── README.md
└── LICENSE-Apache


/.gitignore:
--------------------------------------------------------------------------------
1 | /target
2 | **/*.rs.bk
3 | Cargo.lock
4 | 


--------------------------------------------------------------------------------
/CODE_OF_CONDUCT.md:
--------------------------------------------------------------------------------
1 | # Code of Conduct
2 | 
3 | We follow the [CNCF Code of Conduct](https://github.com/cncf/foundation/blob/main/code-of-conduct.md).
4 | 
5 | Please contact the [CNCF Code of Conduct Committee](mailto:conduct@cncf.io) in order to report violations of the Code of Conduct.
6 | 


--------------------------------------------------------------------------------
/src/lib.rs:
--------------------------------------------------------------------------------
 1 | // Copyright 2019 TiKV Project Authors. Licensed under MIT or Apache-2.0.
 2 | 
 3 | #![warn(
 4 |     elided_lifetimes_in_paths,
 5 |     future_incompatible,
 6 |     missing_debug_implementations,
 7 |     missing_docs,
 8 |     single_use_lifetimes,
 9 |     trivial_casts,
10 |     trivial_numeric_casts,
11 |     unsafe_code,
12 |     unused_qualifications,
13 |     variant_size_differences,
14 |     clippy::all,
15 |     clippy::pedantic
16 | )]
17 | #![allow(clippy::module_name_repetitions, clippy::must_use_candidate)]
18 | #![doc = include_str!("../README.md")]
19 | 
20 | pub mod clock;
21 | #[cfg(feature = "futures-io")]
22 | mod io;
23 | pub mod limiter;
24 | 
25 | pub use limiter::{Limiter, Resource};
26 | 


--------------------------------------------------------------------------------
/benches/standard_clock.rs:
--------------------------------------------------------------------------------
 1 | // Copyright 2019 TiKV Project Authors. Licensed under MIT or Apache-2.0.
 2 | 
 3 | use async_speed_limit::Limiter;
 4 | use criterion::{criterion_group, criterion_main, Criterion};
 5 | 
 6 | // Intel(R) Xeon(R) CPU E5-2630 v4 @ 2.20GHz
 7 | // infinity speed          time:   [32.527 ns 32.667 ns 32.818 ns]
 8 | //                         change: [-0.2189% +0.3899% +1.0013%] (p = 0.20 > 0.05)
 9 | fn criterion_benchmark_infinity_speed(c: &mut Criterion) {
10 |     let limiter = <Limiter>::new(f64::INFINITY);
11 | 
12 |     c.bench_function("infinity speed", |b| {
13 |         b.iter(|| futures_executor::block_on(limiter.consume(1)))
14 |     });
15 | }
16 | 
17 | criterion_group!(benches, criterion_benchmark_infinity_speed);
18 | criterion_main!(benches);
19 | 


--------------------------------------------------------------------------------
/LICENSE-MIT:
--------------------------------------------------------------------------------
 1 | Copyright 2019 TiKV Project Authors
 2 | 
 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of
 4 | this software and associated documentation files (the "Software"), to deal in
 5 | the Software without restriction, including without limitation the rights to
 6 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 7 | the Software, and to permit persons to whom the Software is furnished to do so,
 8 | subject to the following conditions:
 9 | 
10 | The above copyright notice and this permission notice shall be included in all
11 | copies or substantial portions of the Software.
12 | 
13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
15 | FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
16 | COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
17 | IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
18 | CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
19 | 


--------------------------------------------------------------------------------
/.github/workflows/rust.yml:
--------------------------------------------------------------------------------
 1 | name: Rust
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - master
 7 |   pull_request:
 8 |     branches:
 9 |       - master
10 | 
11 | jobs:
12 |   test:
13 |     runs-on: ubuntu-latest
14 |     timeout-minutes: 5
15 |     strategy:
16 |       fail-fast: true
17 |       matrix:
18 |         rustup:
19 |           - toolchain: stable
20 |             components: "clippy, rustfmt"
21 |           - toolchain: nightly
22 |             components: ""
23 |     steps:
24 |     - uses: actions/checkout@v2
25 |     - uses: actions-rs/toolchain@v1
26 |       name: Install Rust
27 |       with:
28 |         toolchain: ${{ matrix.rustup.toolchain }}
29 |         profile: minimal
30 |         components: ${{ matrix.rustup.components }}
31 |         default: true
32 |     - name: Clippy
33 |       run: cargo clippy
34 |       if: contains(matrix.rustup.components, 'clippy')
35 |     - name: Format
36 |       run: cargo fmt -- --check
37 |       if: contains(matrix.rustup.components, 'rustfmt')
38 |     - name: Test (no features)
39 |       run: cargo test --no-default-features
40 |     - name: Test (all features)
41 |       run: cargo test --all-features
42 |       if: matrix.rustup.toolchain == 'nightly'
43 | 


--------------------------------------------------------------------------------
/Cargo.toml:
--------------------------------------------------------------------------------
 1 | [package]
 2 | name = "async-speed-limit"
 3 | version = "0.4.2"
 4 | authors = ["The TiKV Project Developers"]
 5 | license = "MIT OR Apache-2.0"
 6 | edition = "2018"
 7 | repository = "https://github.com/tikv/async-speed-limit"
 8 | description = "Asynchronously speed-limiting multiple byte streams"
 9 | readme = "README.md"
10 | 
11 | [package.metadata.docs.rs]
12 | all-features = true
13 | 
14 | [dependencies]
15 | futures-timer = { version = "3.0.2", optional = true }
16 | futures-io = { version = "0.3.19", optional = true }
17 | futures-core = { version = "0.3.1", optional = true }
18 | pin-project-lite = "0.2.0"
19 | tokio = { version = "1", optional = true }
20 | 
21 | [features]
22 | default = ["futures-io", "fused-future", "standard-clock"]
23 | standard-clock = ["futures-timer"]
24 | fused-future = ["futures-core"]
25 | read-initializer = [] # This feature has been deprecated. It is only kept for downstream compatibility.
26 | 
27 | [dev-dependencies]
28 | futures-executor = "0.3.1"
29 | futures-util = { version = "0.3.1", default-features = false, features = ["std", "channel", "io", "io-compat"] }
30 | rand = "0.8.0"
31 | tokio = { version = "1", features = ["io-util", "rt", "macros"] }
32 | tokio-util = { version = "0.6.9", features = ["codec"] }
33 | criterion = "0.3"
34 | 
35 | [[bench]]
36 | name = "standard_clock"
37 | harness = false
38 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | async-speed-limit
 2 | =================
 3 | 
 4 | [![Build status](https://github.com/tikv/async-speed-limit/workflows/Rust/badge.svg)](https://github.com/tikv/async-speed-limit/actions?query=workflow%3ARust)
 5 | [![Latest Version](https://img.shields.io/crates/v/async-speed-limit.svg)](https://crates.io/crates/async-speed-limit)
 6 | [![Documentation](https://img.shields.io/badge/api-rustdoc-blue.svg)](https://docs.rs/async-speed-limit)
 7 | 
 8 | Asynchronously speed-limiting multiple byte streams (`AsyncRead` and `AsyncWrite`).
 9 | 
10 | ## Usage
11 | 
12 | ```rust
13 | use async_speed_limit::Limiter;
14 | use futures_util::{
15 |     future::try_join,
16 |     io::{self, AsyncRead, AsyncWrite},
17 | };
18 | use std::{marker::Unpin, pin::Pin};
19 | 
20 | #[cfg(feature = "standard-clock")]
21 | async fn copy_both_slowly(
22 |     r1: impl AsyncRead,
23 |     w1: &mut (impl AsyncWrite + Unpin),
24 |     r2: impl AsyncRead,
25 |     w2: &mut (impl AsyncWrite + Unpin),
26 | ) -> io::Result<()> {
27 |     // create a speed limiter of 1.0 KiB/s.
28 |     let limiter = <Limiter>::new(1024.0);
29 | 
30 |     // limit both readers by the same queue
31 |     // (so 1.0 KiB/s is the combined maximum speed)
32 |     let r1 = limiter.clone().limit(r1);
33 |     let r2 = limiter.limit(r2);
34 | 
35 |     // do the copy.
36 |     try_join(io::copy(r1, w1), io::copy(r2, w2)).await?;
37 |     Ok(())
38 | }
39 | ```
40 | 
41 | ## Algorithm
42 | 
43 | async-speed-limit imposes the maximum speed limit using the [token bucket]
44 | algorithm with a single queue. Transmission is throttled by adding a delay
45 | proportional to the consumed tokens *after* it is sent. Because of this, the
46 | traffic flow will have high burstiness around when the token bucket is full.
47 | 
48 | The time needed to refill the bucket from empty to full is called the
49 | *refill period*. Reducing the refill period also reduces burstiness, but there
50 | would be more sleeps. The default value (0.1 s) should be good enough for most
51 | situations.
52 | 
53 | [token bucket]: https://en.wikipedia.org/wiki/Token_bucket
54 | 
55 | ## Tokio 0.1
56 | 
57 | async-speed-limit is designed for "Futures 0.3". This package does not directly
58 | support Tokio 0.1, but you can use futures-util's [`compat` module] to bridge
59 | the types.
60 | 
61 | [futures-timer]: https://crates.io/crates/futures-timer
62 | [`compat` module]: https://docs.rs/futures-util/0.3/futures_util/compat/index.html
63 | 
64 | ## Cargo features
65 | 
66 | | Name                         | Dependency      | Effect                                                                                    |
67 | |------------------------------|-----------------|-------------------------------------------------------------------------------------------|
68 | | **standard-clock** (default) | [futures-timer] | Enables the `clock::StandardClock` struct.                                                |
69 | | **fused-future** (default)   | [futures-core]  | Implements `FusedFuture` on `limiter::Consume`.                                           |
70 | | **futures-io** (default)     | [futures-io]    | Implements `AsyncRead` and `AsyncWrite` on `limiter::Resource`.                           |
71 | | ~~**read-initializer**~~     | -               | Deprecated and does nothing. This feature has been removed since `futures-io 0.3.19`.     |
72 | 
73 | [futures-core]: https://crates.io/crates/futures-core
74 | [futures-io]: https://crates.io/crates/futures-io
75 | 
76 | ## License
77 | 
78 | async-speed-limit is dual-licensed under
79 | 
80 | * Apache 2.0 license ([LICENSE-Apache](./LICENSE-Apache) or <http://www.apache.org/licenses/LICENSE-2.0>)
81 | * MIT license ([LICENSE-MIT](./LICENSE-MIT) or <https://opensource.org/licenses/MIT>)
82 | 


--------------------------------------------------------------------------------
/src/io.rs:
--------------------------------------------------------------------------------
  1 | // Copyright 2019 TiKV Project Authors. Licensed under MIT or Apache-2.0.
  2 | 
  3 | use crate::{clock::Clock, limiter::Resource};
  4 | use std::{
  5 |     io::{self, IoSlice, IoSliceMut},
  6 |     pin::Pin,
  7 |     task::{Context, Poll},
  8 | };
  9 | 
 10 | fn length_of_result_usize<B>(a: &io::Result<usize>, _: &B) -> usize {
 11 |     *a.as_ref().unwrap_or(&0)
 12 | }
 13 | 
 14 | impl<R: futures_io::AsyncRead, C: Clock> futures_io::AsyncRead for Resource<R, C> {
 15 |     fn poll_read(
 16 |         self: Pin<&mut Self>,
 17 |         cx: &mut Context<'_>,
 18 |         buf: &mut [u8],
 19 |     ) -> Poll<io::Result<usize>> {
 20 |         self.poll_limited(cx, (), length_of_result_usize, |r, cx, ()| {
 21 |             r.poll_read(cx, buf)
 22 |         })
 23 |     }
 24 | 
 25 |     fn poll_read_vectored(
 26 |         self: Pin<&mut Self>,
 27 |         cx: &mut Context<'_>,
 28 |         bufs: &mut [IoSliceMut<'_>],
 29 |     ) -> Poll<io::Result<usize>> {
 30 |         self.poll_limited(cx, (), length_of_result_usize, |r, cx, ()| {
 31 |             r.poll_read_vectored(cx, bufs)
 32 |         })
 33 |     }
 34 | }
 35 | 
 36 | impl<R: futures_io::AsyncWrite, C: Clock> futures_io::AsyncWrite for Resource<R, C> {
 37 |     fn poll_write(
 38 |         self: Pin<&mut Self>,
 39 |         cx: &mut Context<'_>,
 40 |         buf: &[u8],
 41 |     ) -> Poll<io::Result<usize>> {
 42 |         self.poll_limited(cx, (), length_of_result_usize, |r, cx, ()| {
 43 |             r.poll_write(cx, buf)
 44 |         })
 45 |     }
 46 | 
 47 |     fn poll_write_vectored(
 48 |         self: Pin<&mut Self>,
 49 |         cx: &mut Context<'_>,
 50 |         bufs: &[IoSlice<'_>],
 51 |     ) -> Poll<io::Result<usize>> {
 52 |         self.poll_limited(cx, (), length_of_result_usize, |r, cx, ()| {
 53 |             r.poll_write_vectored(cx, bufs)
 54 |         })
 55 |     }
 56 | 
 57 |     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
 58 |         self.get_pin_mut().poll_flush(cx)
 59 |     }
 60 | 
 61 |     fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
 62 |         self.get_pin_mut().poll_close(cx)
 63 |     }
 64 | }
 65 | 
 66 | #[cfg(feature = "tokio")]
 67 | impl<R: tokio::io::AsyncRead, C: Clock> tokio::io::AsyncRead for Resource<R, C> {
 68 |     fn poll_read(
 69 |         self: Pin<&mut Self>,
 70 |         cx: &mut Context<'_>,
 71 |         buf: &mut tokio::io::ReadBuf<'_>,
 72 |     ) -> Poll<io::Result<()>> {
 73 |         let filled = buf.filled().len();
 74 | 
 75 |         self.poll_limited(
 76 |             cx,
 77 |             buf,
 78 |             |res: &io::Result<()>, buf| {
 79 |                 res.is_ok()
 80 |                     .then(|| buf.filled().len() - filled)
 81 |                     .unwrap_or_default()
 82 |             },
 83 |             |r, cx, buf| r.poll_read(cx, buf),
 84 |         )
 85 |     }
 86 | }
 87 | 
 88 | #[cfg(feature = "tokio")]
 89 | impl<R: tokio::io::AsyncWrite, C: Clock> tokio::io::AsyncWrite for Resource<R, C> {
 90 |     fn poll_write(
 91 |         self: Pin<&mut Self>,
 92 |         cx: &mut Context<'_>,
 93 |         buf: &[u8],
 94 |     ) -> Poll<io::Result<usize>> {
 95 |         self.poll_limited(cx, (), length_of_result_usize, |r, cx, _| {
 96 |             r.poll_write(cx, buf)
 97 |         })
 98 |     }
 99 | 
100 |     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
101 |         self.get_pin_mut().poll_flush(cx)
102 |     }
103 | 
104 |     fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
105 |         self.get_pin_mut().poll_shutdown(cx)
106 |     }
107 | }
108 | 
109 | #[cfg(test)]
110 | mod tests {
111 |     use super::*;
112 |     use crate::{
113 |         clock::{ManualClock, Nanoseconds},
114 |         Limiter,
115 |     };
116 |     use futures_executor::LocalPool;
117 |     use futures_util::{
118 |         io::{copy_buf, BufReader},
119 |         task::SpawnExt,
120 |     };
121 |     use rand::{thread_rng, RngCore};
122 | 
123 |     #[test]
124 |     fn limited_read() {
125 |         let mut pool = LocalPool::new();
126 |         let sp = pool.spawner();
127 | 
128 |         let limiter = Limiter::<ManualClock>::new(512.0);
129 |         let clock = limiter.clock();
130 | 
131 |         sp.spawn({
132 |             let limiter = limiter.clone();
133 |             let clock = clock.clone();
134 |             async move {
135 |                 let mut src = vec![0_u8; 1024];
136 |                 thread_rng().fill_bytes(&mut src);
137 |                 let mut dst = Vec::new();
138 | 
139 |                 let read = BufReader::with_capacity(256, limiter.limit(&*src));
140 |                 let count = copy_buf(read, &mut dst).await.unwrap();
141 | 
142 |                 assert_eq!(clock.now(), Nanoseconds(2_000_000_000));
143 |                 assert_eq!(count, src.len() as u64);
144 |                 assert!(src == dst);
145 |             }
146 |         })
147 |         .unwrap();
148 | 
149 |         clock.set_time(Nanoseconds(0));
150 |         pool.run_until_stalled();
151 |         assert_eq!(limiter.total_bytes_consumed(), 256);
152 | 
153 |         clock.set_time(Nanoseconds(500_000_000));
154 |         pool.run_until_stalled();
155 |         assert_eq!(limiter.total_bytes_consumed(), 512);
156 | 
157 |         clock.set_time(Nanoseconds(1_000_000_000));
158 |         pool.run_until_stalled();
159 |         assert_eq!(limiter.total_bytes_consumed(), 768);
160 | 
161 |         clock.set_time(Nanoseconds(1_500_000_000));
162 |         pool.run_until_stalled();
163 |         assert_eq!(limiter.total_bytes_consumed(), 1024);
164 | 
165 |         clock.set_time(Nanoseconds(2_000_000_000));
166 |         pool.run_until_stalled();
167 | 
168 |         assert!(!pool.try_run_one());
169 |     }
170 | 
171 |     #[test]
172 |     fn unlimited_read() {
173 |         let mut pool = LocalPool::new();
174 |         let sp = pool.spawner();
175 | 
176 |         let limiter = Limiter::<ManualClock>::new(std::f64::INFINITY);
177 | 
178 |         sp.spawn({
179 |             async move {
180 |                 let mut src = vec![0_u8; 1024];
181 |                 thread_rng().fill_bytes(&mut src);
182 |                 let mut dst = Vec::new();
183 | 
184 |                 let read = BufReader::with_capacity(256, limiter.limit(&*src));
185 |                 let count = copy_buf(read, &mut dst).await.unwrap();
186 | 
187 |                 assert_eq!(count, src.len() as u64);
188 |                 assert!(src == dst);
189 |             }
190 |         })
191 |         .unwrap();
192 | 
193 |         pool.run_until_stalled();
194 |         assert!(!pool.try_run_one());
195 |     }
196 | 
197 |     #[test]
198 |     fn limited_write() {
199 |         let mut pool = LocalPool::new();
200 |         let sp = pool.spawner();
201 | 
202 |         let limiter = Limiter::<ManualClock>::new(512.0);
203 |         let clock = limiter.clock();
204 | 
205 |         sp.spawn({
206 |             let limiter = limiter.clone();
207 |             let clock = clock.clone();
208 |             async move {
209 |                 let mut src = vec![0_u8; 1024];
210 |                 thread_rng().fill_bytes(&mut src);
211 | 
212 |                 let read = BufReader::with_capacity(256, &*src);
213 |                 let mut write = limiter.limit(Vec::new());
214 |                 let count = copy_buf(read, &mut write).await.unwrap();
215 | 
216 |                 assert_eq!(clock.now(), Nanoseconds(1_500_000_000));
217 |                 assert_eq!(count, src.len() as u64);
218 |                 assert!(src == write.into_inner());
219 |             }
220 |         })
221 |         .unwrap();
222 | 
223 |         clock.set_time(Nanoseconds(0));
224 |         pool.run_until_stalled();
225 |         assert_eq!(limiter.total_bytes_consumed(), 256);
226 | 
227 |         clock.set_time(Nanoseconds(500_000_000));
228 |         pool.run_until_stalled();
229 |         assert_eq!(limiter.total_bytes_consumed(), 512);
230 | 
231 |         clock.set_time(Nanoseconds(1_000_000_000));
232 |         pool.run_until_stalled();
233 |         assert_eq!(limiter.total_bytes_consumed(), 768);
234 | 
235 |         clock.set_time(Nanoseconds(1_500_000_000));
236 |         pool.run_until_stalled();
237 |         assert_eq!(limiter.total_bytes_consumed(), 1024);
238 | 
239 |         clock.set_time(Nanoseconds(2_000_000_000));
240 |         pool.run_until_stalled();
241 | 
242 |         assert!(!pool.try_run_one());
243 |     }
244 | }
245 | 
246 | #[cfg(test)]
247 | #[cfg(all(feature = "tokio", feature = "standard-clock"))]
248 | mod tokio_tests {
249 |     use std::{
250 |         io,
251 |         time::{Duration, Instant},
252 |     };
253 | 
254 |     use crate::Limiter;
255 | 
256 |     use futures_util::{future::ok, TryStreamExt as _};
257 |     use tokio::io::{copy, repeat, sink, AsyncReadExt as _, AsyncWriteExt as _};
258 |     use tokio_util::codec::{BytesCodec, FramedRead};
259 | 
260 |     #[tokio::test]
261 |     async fn limited_read() -> io::Result<()> {
262 |         let limiter = <Limiter>::new(32768.0);
263 | 
264 |         let start_time = Instant::now();
265 | 
266 |         let reader = repeat(50).take(65536);
267 |         let mut reader = limiter.limit(reader);
268 | 
269 |         let mut sink = sink();
270 |         let total = copy(&mut reader, &mut sink).await?;
271 |         sink.shutdown().await?;
272 | 
273 |         let elapsed = start_time.elapsed();
274 | 
275 |         assert!(
276 |             Duration::from_millis(1900) <= elapsed && elapsed <= Duration::from_millis(2100),
277 |             "elapsed = {:?}",
278 |             elapsed
279 |         );
280 |         assert_eq!(total, 65536);
281 | 
282 |         Ok(())
283 |     }
284 | 
285 |     #[tokio::test]
286 |     async fn unlimited_read() -> io::Result<()> {
287 |         let limiter = <Limiter>::new(std::f64::INFINITY);
288 | 
289 |         let start_time = Instant::now();
290 | 
291 |         let reader = repeat(50).take(65536);
292 |         let mut reader = limiter.limit(reader);
293 |         let mut sink = sink();
294 | 
295 |         let total = copy(&mut reader, &mut sink).await?;
296 |         sink.shutdown().await?;
297 | 
298 |         let elapsed = start_time.elapsed();
299 | 
300 |         assert!(
301 |             elapsed <= Duration::from_millis(100),
302 |             "elapsed = {:?}",
303 |             elapsed
304 |         );
305 |         assert_eq!(total, 65536);
306 | 
307 |         Ok(())
308 |     }
309 | 
310 |     #[tokio::test]
311 |     async fn limited_read_byte_stream() -> io::Result<()> {
312 |         let limiter = <Limiter>::new(30000.0);
313 | 
314 |         let start_time = Instant::now();
315 | 
316 |         let reader = repeat(50).take(60000);
317 |         let reader = limiter.limit(reader);
318 | 
319 |         let total = FramedRead::new(reader, BytesCodec::new())
320 |             .try_fold(0, |i, j| {
321 |                 assert!(j.iter().all(|b| *b == 50_u8), "{} / {:?}", i, j);
322 |                 ok(i + j.len())
323 |             })
324 |             .await?;
325 | 
326 |         let elapsed = start_time.elapsed();
327 | 
328 |         assert!(
329 |             Duration::from_millis(1900) <= elapsed && elapsed <= Duration::from_millis(2100),
330 |             "elapsed = {:?}",
331 |             elapsed
332 |         );
333 |         assert_eq!(total, 60000);
334 | 
335 |         Ok(())
336 |     }
337 | }
338 | 


--------------------------------------------------------------------------------
/src/clock.rs:
--------------------------------------------------------------------------------
  1 | // Copyright 2019 TiKV Project Authors. Licensed under MIT or Apache-2.0.
  2 | 
  3 | //! Clocks
  4 | 
  5 | #[cfg(feature = "standard-clock")]
  6 | use futures_timer::Delay;
  7 | #[cfg(feature = "standard-clock")]
  8 | use std::time::Instant;
  9 | use std::{
 10 |     convert::TryInto,
 11 |     fmt::Debug,
 12 |     future::Future,
 13 |     marker::Unpin,
 14 |     mem,
 15 |     ops::{Add, Sub},
 16 |     pin::Pin,
 17 |     sync::{
 18 |         atomic::{AtomicU64, Ordering},
 19 |         Arc, Mutex,
 20 |     },
 21 |     task::{Context, Poll, Waker},
 22 |     time::Duration,
 23 | };
 24 | 
 25 | /// A `Clock` controls the passing of time.
 26 | ///
 27 | /// [`Limiter`](crate::Limiter) uses [`sleep()`](Clock::sleep()) to impose speed
 28 | /// limit, and it relies on the current and past timestamps to determine how
 29 | /// long to sleep. Both of these time-related features are encapsulated into
 30 | /// this `Clock` trait.
 31 | ///
 32 | /// # Implementing
 33 | ///
 34 | /// The [`StandardClock`] should be enough in most situation. However, these are
 35 | /// cases for a custom clock, e.g. use a coarse clock instead of the standard
 36 | /// high-precision clock, or use a specialized future associated with an
 37 | /// executor instead of the generic `futures-timer`.
 38 | ///
 39 | /// Types implementing `Clock` must be cheap to clone (e.g. using `Arc`), and
 40 | /// the default value must be ready to use.
 41 | pub trait Clock: Clone + Default {
 42 |     /// Type to represent a point of time.
 43 |     ///
 44 |     /// Subtracting two instances should return the duration elapsed between
 45 |     /// them. The subtraction must never block or panic when they are properly
 46 |     /// ordered.
 47 |     type Instant: Copy + Sub<Output = Duration>;
 48 | 
 49 |     /// Future type returned by [`sleep()`](Clock::sleep()).
 50 |     type Delay: Future<Output = ()> + Unpin;
 51 | 
 52 |     /// Returns the current time instant. It should be monotonically increasing,
 53 |     /// but not necessarily high-precision or steady.
 54 |     ///
 55 |     /// This function must never block or panic.
 56 |     fn now(&self) -> Self::Instant;
 57 | 
 58 |     /// Asynchronously sleeps the current task for the given duration.
 59 |     ///
 60 |     /// This method should return a future which fulfilled as `()` after a
 61 |     /// duration of `dur`. It should *not* block the current thread.
 62 |     ///
 63 |     /// `sleep()` is often called with a duration of 0 s. This should result in
 64 |     /// a future being resolved immediately.
 65 |     fn sleep(&self, dur: Duration) -> Self::Delay;
 66 | }
 67 | 
 68 | /// A `BlockingClock` is a [`Clock`] which supports synchronous sleeping.
 69 | pub trait BlockingClock: Clock {
 70 |     /// Sleeps and blocks the current thread for the given duration.
 71 |     fn blocking_sleep(&self, dur: Duration);
 72 | }
 73 | 
 74 | /// The physical clock using [`std::time::Instant`].
 75 | ///
 76 | /// The sleeping future is based on [`futures-timer`]. Blocking sleep uses
 77 | /// [`std::thread::sleep()`].
 78 | ///
 79 | /// [`futures-timer`]: https://docs.rs/futures-timer/
 80 | #[cfg(feature = "standard-clock")]
 81 | #[derive(Copy, Clone, Debug, Default)]
 82 | pub struct StandardClock;
 83 | 
 84 | #[cfg(feature = "standard-clock")]
 85 | impl Clock for StandardClock {
 86 |     type Instant = Instant;
 87 |     type Delay = Delay;
 88 | 
 89 |     fn now(&self) -> Self::Instant {
 90 |         Instant::now()
 91 |     }
 92 | 
 93 |     fn sleep(&self, dur: Duration) -> Self::Delay {
 94 |         Delay::new(dur)
 95 |     }
 96 | }
 97 | 
 98 | #[cfg(feature = "standard-clock")]
 99 | impl BlockingClock for StandardClock {
100 |     fn blocking_sleep(&self, dur: Duration) {
101 |         std::thread::sleep(dur);
102 |     }
103 | }
104 | 
105 | /// Number of nanoseconds since an arbitrary epoch.
106 | ///
107 | /// This is the instant type of [`ManualClock`].
108 | #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
109 | pub struct Nanoseconds(pub u64);
110 | 
111 | impl Sub for Nanoseconds {
112 |     type Output = Duration;
113 |     fn sub(self, other: Self) -> Duration {
114 |         Duration::from_nanos(self.0 - other.0)
115 |     }
116 | }
117 | 
118 | impl Add<Duration> for Nanoseconds {
119 |     type Output = Self;
120 |     fn add(self, other: Duration) -> Self {
121 |         let dur: u64 = other
122 |             .as_nanos()
123 |             .try_into()
124 |             .expect("cannot increase more than 2^64 ns");
125 |         Self(self.0 + dur)
126 |     }
127 | }
128 | 
129 | /// The future returned by [`ManualClock`]`::sleep()`.
130 | #[derive(Debug)]
131 | pub struct ManualDelay {
132 |     clock: Arc<ManualClockContent>,
133 |     timeout: Nanoseconds,
134 | }
135 | 
136 | impl Future for ManualDelay {
137 |     type Output = ();
138 | 
139 |     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
140 |         let now = self.clock.now();
141 |         if now >= self.timeout {
142 |             Poll::Ready(())
143 |         } else {
144 |             self.clock.register(cx);
145 |             Poll::Pending
146 |         }
147 |     }
148 | }
149 | 
150 | /// Internal, shared part of [`ManualClock`]. `ManualClock` itself is an `Arc`
151 | /// of `ManualClockContent`.
152 | #[derive(Default, Debug)]
153 | struct ManualClockContent {
154 |     now: AtomicU64,
155 |     wakers: Mutex<Vec<Waker>>,
156 | }
157 | 
158 | impl ManualClockContent {
159 |     fn now(&self) -> Nanoseconds {
160 |         Nanoseconds(self.now.load(Ordering::SeqCst))
161 |     }
162 | 
163 |     fn set_time(&self, time: u64) {
164 |         let old_time = self.now.swap(time, Ordering::SeqCst);
165 |         assert!(old_time <= time, "cannot move the time backwards");
166 | 
167 |         let wakers = { mem::take(&mut *self.wakers.lock().unwrap()) };
168 |         wakers.into_iter().for_each(Waker::wake);
169 |     }
170 | 
171 |     fn register(&self, cx: &mut Context<'_>) {
172 |         self.wakers.lock().unwrap().push(cx.waker().clone());
173 |     }
174 | }
175 | 
176 | /// A [`Clock`] where the passage of time can be manually controlled.
177 | ///
178 | /// This type is mainly used for testing behavior of speed limiter only.
179 | ///
180 | /// This clock only supports up to 2<sup>64</sup> ns (about 584.5 years).
181 | ///
182 | /// # Examples
183 | ///
184 | /// ```rust
185 | /// use async_speed_limit::clock::{Clock, ManualClock, Nanoseconds};
186 | ///
187 | /// let clock = ManualClock::new();
188 | /// assert_eq!(clock.now(), Nanoseconds(0));
189 | /// clock.set_time(Nanoseconds(1_000_000_000));
190 | /// assert_eq!(clock.now(), Nanoseconds(1_000_000_000));
191 | /// ```
192 | #[derive(Default, Debug, Clone)]
193 | pub struct ManualClock(Arc<ManualClockContent>);
194 | 
195 | impl ManualClock {
196 |     /// Creates a new clock with time set to 0.
197 |     pub fn new() -> Self {
198 |         Self::default()
199 |     }
200 | 
201 |     /// Set the current time of this clock to the given value.
202 |     ///
203 |     /// # Panics
204 |     ///
205 |     /// Since [`now()`](Clock::now()) must be monotonically increasing, if the
206 |     /// new time is less than the previous time, this function will panic.
207 |     pub fn set_time(&self, time: Nanoseconds) {
208 |         self.0.set_time(time.0);
209 |     }
210 | }
211 | 
212 | impl Clock for ManualClock {
213 |     type Instant = Nanoseconds;
214 |     type Delay = ManualDelay;
215 | 
216 |     fn now(&self) -> Self::Instant {
217 |         self.0.now()
218 |     }
219 | 
220 |     fn sleep(&self, dur: Duration) -> Self::Delay {
221 |         ManualDelay {
222 |             timeout: self.0.now() + dur,
223 |             clock: self.0.clone(),
224 |         }
225 |     }
226 | }
227 | 
228 | #[cfg(test)]
229 | mod tests {
230 |     use super::*;
231 |     use futures_executor::LocalPool;
232 |     use futures_util::task::SpawnExt;
233 |     use std::sync::{
234 |         atomic::{AtomicUsize, Ordering},
235 |         Arc,
236 |     };
237 | 
238 |     #[test]
239 |     fn manual_clock_basics() {
240 |         let clock = ManualClock::new();
241 |         let t1 = clock.now();
242 |         assert_eq!(t1, Nanoseconds(0));
243 | 
244 |         clock.set_time(Nanoseconds(1_000_000_000));
245 | 
246 |         let t2 = clock.now();
247 |         assert_eq!(t2, Nanoseconds(1_000_000_000));
248 |         assert_eq!(t2 - t1, Duration::from_secs(1));
249 | 
250 |         clock.set_time(Nanoseconds(1_000_000_007));
251 | 
252 |         let t3 = clock.now();
253 |         assert_eq!(t3, Nanoseconds(1_000_000_007));
254 |         assert_eq!(t3 - t2, Duration::from_nanos(7));
255 |     }
256 | 
257 |     #[test]
258 |     fn manual_clock_sleep() {
259 |         let counter = Arc::new(AtomicUsize::new(0));
260 |         let clock = ManualClock::new();
261 |         let mut pool = LocalPool::new();
262 |         let sp = pool.spawner();
263 | 
264 |         // expected sequence:
265 |         //
266 |         //   t=0    .     .
267 |         //   t=1    .     .
268 |         //   t=2    +1    .
269 |         //   t=3    .     +16
270 |         //   t=4    .     +64
271 |         //   t=5    +4    .
272 | 
273 |         sp.spawn({
274 |             let counter = counter.clone();
275 |             let clock = clock.clone();
276 |             async move {
277 |                 clock.sleep(Duration::from_secs(2)).await;
278 |                 counter.fetch_add(1, Ordering::Relaxed);
279 |                 clock.sleep(Duration::from_secs(3)).await;
280 |                 counter.fetch_add(4, Ordering::Relaxed);
281 |             }
282 |         })
283 |         .unwrap();
284 | 
285 |         sp.spawn({
286 |             let counter = counter.clone();
287 |             let clock = clock.clone();
288 |             async move {
289 |                 clock.sleep(Duration::from_secs(3)).await;
290 |                 counter.fetch_add(16, Ordering::Relaxed);
291 |                 clock.sleep(Duration::from_secs(1)).await;
292 |                 counter.fetch_add(64, Ordering::Relaxed);
293 |             }
294 |         })
295 |         .unwrap();
296 | 
297 |         clock.set_time(Nanoseconds(0));
298 |         pool.run_until_stalled();
299 |         assert_eq!(counter.load(Ordering::Relaxed), 0);
300 | 
301 |         clock.set_time(Nanoseconds(1_000_000_000));
302 |         pool.run_until_stalled();
303 |         assert_eq!(counter.load(Ordering::Relaxed), 0);
304 | 
305 |         clock.set_time(Nanoseconds(2_000_000_000));
306 |         pool.run_until_stalled();
307 |         assert_eq!(counter.load(Ordering::Relaxed), 1);
308 | 
309 |         clock.set_time(Nanoseconds(3_000_000_000));
310 |         pool.run_until_stalled();
311 |         assert_eq!(counter.load(Ordering::Relaxed), 17);
312 | 
313 |         clock.set_time(Nanoseconds(4_000_000_000));
314 |         pool.run_until_stalled();
315 |         assert_eq!(counter.load(Ordering::Relaxed), 81);
316 | 
317 |         clock.set_time(Nanoseconds(5_000_000_000));
318 |         pool.run_until_stalled();
319 |         assert_eq!(counter.load(Ordering::Relaxed), 85);
320 | 
321 |         // all futures should be exhausted.
322 |         assert!(!pool.try_run_one());
323 |     }
324 | 
325 |     #[test]
326 |     #[cfg(feature = "standard-clock")]
327 |     fn standard_clock() {
328 |         let res = futures_executor::block_on(async {
329 |             let init = StandardClock.now();
330 |             StandardClock.sleep(Duration::from_secs(1)).await;
331 |             StandardClock.now() - init
332 |         });
333 |         assert!(
334 |             Duration::from_millis(900) <= res && res <= Duration::from_millis(1100),
335 |             "standard clock slept too long at {:?}",
336 |             res
337 |         );
338 |     }
339 | }
340 | 


--------------------------------------------------------------------------------
/LICENSE-Apache:
--------------------------------------------------------------------------------
  1 | Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "{}"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright {}
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/src/limiter.rs:
--------------------------------------------------------------------------------
   1 | // Copyright 2019 TiKV Project Authors. Licensed under MIT or Apache-2.0.
   2 | 
   3 | //! Speed limiter
   4 | 
   5 | #[cfg(feature = "standard-clock")]
   6 | use crate::clock::StandardClock;
   7 | use crate::clock::{BlockingClock, Clock};
   8 | use pin_project_lite::pin_project;
   9 | use std::{
  10 |     future::Future,
  11 |     mem,
  12 |     ops::Sub,
  13 |     pin::Pin,
  14 |     sync::Arc,
  15 |     sync::{
  16 |         atomic::{AtomicBool, AtomicUsize, Ordering},
  17 |         Mutex,
  18 |     },
  19 |     task::{Context, Poll},
  20 |     time::Duration,
  21 | };
  22 | 
  23 | /// Stores the current state of the limiter.
  24 | #[derive(Debug, Clone, Copy)]
  25 | struct Bucket<I> {
  26 |     /// Last updated instant of the bucket. This is used to compare with the
  27 |     /// current instant to deduce the current bucket value.
  28 |     last_updated: I,
  29 |     /// The speed limit (unit: B/s).
  30 |     speed_limit: f64,
  31 |     /// Time needed to refill the entire bucket (unit: s).
  32 |     refill: f64,
  33 |     /// The number of bytes the bucket is carrying at the time `last_updated`.
  34 |     /// This value can be negative.
  35 |     value: f64,
  36 |     /// The minimum duration to wait if self.value is smaller than 0 after
  37 |     /// call `self.consume` (unit: s).
  38 |     /// By default, is the same as `refill`.
  39 |     min_wait: f64,
  40 | }
  41 | 
  42 | impl<I> Bucket<I> {
  43 |     /// Returns the maximum number of bytes this bucket can carry.
  44 |     fn capacity(&self) -> f64 {
  45 |         self.speed_limit * self.refill
  46 |     }
  47 | 
  48 |     /// Consumes the given number of bytes from the bucket.
  49 |     ///
  50 |     /// Returns the duration we need for the consumed bytes to recover.
  51 |     ///
  52 |     /// This method should only be called when the speed is finite.
  53 |     fn consume(&mut self, size: f64) -> Duration {
  54 |         self.value -= size;
  55 |         if self.value > 0.0 {
  56 |             Duration::from_secs(0)
  57 |         } else {
  58 |             let sleep_secs = self.min_wait - self.value / self.speed_limit;
  59 |             Duration::from_secs_f64(sleep_secs)
  60 |         }
  61 |     }
  62 | 
  63 |     /// Reverts the previous consumption of the given number of bytes.
  64 |     ///
  65 |     /// This method should only be called when the speed is finite
  66 |     fn unconsume(&mut self, size: f64) {
  67 |         self.value += size;
  68 |     }
  69 | 
  70 |     /// Changes the speed limit.
  71 |     ///
  72 |     /// The current value will be raised or lowered so that the number of
  73 |     /// consumed bytes remains constant.
  74 |     fn set_speed_limit(&mut self, new_speed_limit: f64) {
  75 |         let old_capacity = self.capacity();
  76 |         self.speed_limit = new_speed_limit;
  77 |         if new_speed_limit.is_finite() {
  78 |             let new_capacity = self.capacity();
  79 |             if old_capacity.is_finite() {
  80 |                 self.value += new_capacity - old_capacity;
  81 |             } else {
  82 |                 self.value = new_capacity;
  83 |             }
  84 |         }
  85 |     }
  86 | }
  87 | 
  88 | impl<I: Copy + Sub<Output = Duration>> Bucket<I> {
  89 |     /// Refills the bucket to match the value at current time.
  90 |     ///
  91 |     /// This method should only be called when the speed is finite.
  92 |     fn refill(&mut self, now: I) {
  93 |         let elapsed = (now - self.last_updated).as_secs_f64();
  94 |         let refilled = self.speed_limit * elapsed;
  95 |         self.value = self.capacity().min(self.value + refilled);
  96 |         self.last_updated = now;
  97 |     }
  98 | }
  99 | 
 100 | /// Builder for [`Limiter`].
 101 | ///
 102 | /// # Examples
 103 | ///
 104 | #[cfg_attr(feature = "standard-clock", doc = "```rust")]
 105 | #[cfg_attr(not(feature = "standard-clock"), doc = "```ignore")]
 106 | /// use async_speed_limit::Limiter;
 107 | /// use std::time::Duration;
 108 | ///
 109 | /// let limiter = <Limiter>::builder(1_048_576.0)
 110 | ///     .refill(Duration::from_millis(100))
 111 | ///     .build();
 112 | /// # drop(limiter);
 113 | /// ```
 114 | #[derive(Debug)]
 115 | pub struct Builder<C: Clock> {
 116 |     clock: C,
 117 |     bucket: Bucket<C::Instant>,
 118 |     min_wait: Option<f64>,
 119 | }
 120 | 
 121 | impl<C: Clock> Builder<C> {
 122 |     /// Creates a new limiter builder.
 123 |     ///
 124 |     /// Use [infinity](`std::f64::INFINITY`) to make the speed unlimited.
 125 |     pub fn new(speed_limit: f64) -> Self {
 126 |         let clock = C::default();
 127 |         let mut result = Self {
 128 |             bucket: Bucket {
 129 |                 last_updated: clock.now(),
 130 |                 speed_limit: 0.0,
 131 |                 refill: 0.1,
 132 |                 value: 0.0,
 133 |                 min_wait: 0.1,
 134 |             },
 135 |             clock,
 136 |             min_wait: None,
 137 |         };
 138 |         result.speed_limit(speed_limit);
 139 |         result
 140 |     }
 141 | 
 142 |     /// Sets the speed limit of the limiter.
 143 |     ///
 144 |     /// Use [infinity](`std::f64::INFINITY`) to make the speed unlimited.
 145 |     ///
 146 |     /// # Panics
 147 |     ///
 148 |     /// The speed limit must be positive. Panics if the speed limit is negative,
 149 |     /// zero, or NaN.
 150 |     pub fn speed_limit(&mut self, speed_limit: f64) -> &mut Self {
 151 |         assert!(speed_limit > 0.0, "speed limit must be positive");
 152 |         self.bucket.speed_limit = speed_limit;
 153 |         self
 154 |     }
 155 | 
 156 |     /// Sets the refill period of the limiter.
 157 |     ///
 158 |     /// The default value is 0.1 s, which should be good for most use cases. The
 159 |     /// refill period is ignored if the speed is [infinity](`std::f64::INFINITY`).
 160 |     ///
 161 |     /// # Panics
 162 |     ///
 163 |     /// The duration must not be zero, otherwise this method panics.
 164 |     pub fn refill(&mut self, dur: Duration) -> &mut Self {
 165 |         assert!(
 166 |             dur > Duration::from_secs(0),
 167 |             "refill duration must not be zero"
 168 |         );
 169 |         self.bucket.refill = dur.as_secs_f64();
 170 |         self
 171 |     }
 172 | 
 173 |     /// Sets the minimum wait duration when the speed limit was exceeded.
 174 |     ///
 175 |     /// The default value is same as the refill period.
 176 |     pub fn min_wait(&mut self, dur: Duration) -> &mut Self {
 177 |         self.min_wait = Some(dur.as_secs_f64());
 178 |         self
 179 |     }
 180 | 
 181 |     /// Sets the clock instance used by the limiter.
 182 |     pub fn clock(&mut self, clock: C) -> &mut Self {
 183 |         self.clock = clock;
 184 |         self
 185 |     }
 186 | 
 187 |     /// Builds the limiter.
 188 |     pub fn build(&mut self) -> Limiter<C> {
 189 |         self.bucket.value = self.bucket.capacity();
 190 |         self.bucket.last_updated = self.clock.now();
 191 |         let is_unlimited = self.bucket.speed_limit.is_infinite();
 192 |         let min_wait = self.min_wait.unwrap_or(self.bucket.refill);
 193 |         self.bucket.min_wait = min_wait;
 194 |         Limiter {
 195 |             bucket: Arc::new(Mutex::new(self.bucket)),
 196 |             clock: mem::take(&mut self.clock),
 197 |             total_bytes_consumed: Arc::new(AtomicUsize::new(0)),
 198 |             is_unlimited: Arc::new(AtomicBool::new(is_unlimited)),
 199 |         }
 200 |     }
 201 | }
 202 | 
 203 | macro_rules! declare_limiter {
 204 |     ($($default_clock:tt)*) => {
 205 |         /// A type to control the maximum speed limit of multiple streams.
 206 |         ///
 207 |         /// When a `Limiter` is cloned, the instances would share the same
 208 |         /// queue. Multiple tasks can cooperatively respect a global speed limit
 209 |         /// via clones. Cloning a `Limiter` is cheap (equals to cloning two
 210 |         /// `Arc`s).
 211 |         ///
 212 |         /// The speed limit is imposed by awaiting
 213 |         /// [`consume()`](Limiter::consume()). The method returns a future which
 214 |         /// sleeps until rate falls below the limit.
 215 |         ///
 216 |         /// # Examples
 217 |         ///
 218 |         /// Upload some small files atomically in parallel, while maintaining a
 219 |         /// global speed limit of 1 MiB/s.
 220 |         ///
 221 |         #[cfg_attr(feature = "standard-clock", doc = "```rust")]
 222 |         #[cfg_attr(not(feature = "standard-clock"), doc = "```ignore")]
 223 |         /// use async_speed_limit::Limiter;
 224 |         /// use futures_util::future::try_join_all;
 225 |         ///
 226 |         /// # async {
 227 |         /// # let files = &[""];
 228 |         /// # async fn upload(file: &str) -> Result<(), ()> { Ok(()) }
 229 |         /// let limiter = <Limiter>::new(1_048_576.0);
 230 |         /// let processes = files
 231 |         ///     .iter()
 232 |         ///     .map(|file| {
 233 |         ///         let limiter = limiter.clone();
 234 |         ///         async move {
 235 |         ///             limiter.consume(file.len()).await;
 236 |         ///             upload(file).await?;
 237 |         ///             Ok(())
 238 |         ///         }
 239 |         ///     });
 240 |         /// try_join_all(processes).await?;
 241 |         /// # Ok::<_, ()>(()) };
 242 |         /// ```
 243 |         #[derive(Debug, Clone)]
 244 |         pub struct Limiter<C: Clock $($default_clock)*> {
 245 |             /// State of the limiter.
 246 |             // TODO avoid using Arc<Mutex>?
 247 |             bucket: Arc<Mutex<Bucket<C::Instant>>>,
 248 |             /// Clock used for time calculation.
 249 |             clock: C,
 250 |             /// Statistics of the number of bytes consumed for record. When this
 251 |             /// number reaches `usize::MAX` it will wrap around.
 252 |             total_bytes_consumed: Arc<AtomicUsize>,
 253 |             /// A flag indicates unlimited speed.
 254 |             is_unlimited: Arc<AtomicBool>,
 255 |         }
 256 |     }
 257 | }
 258 | 
 259 | #[cfg(feature = "standard-clock")]
 260 | declare_limiter! { = StandardClock }
 261 | 
 262 | #[cfg(not(feature = "standard-clock"))]
 263 | declare_limiter! {}
 264 | 
 265 | impl<C: Clock> Limiter<C> {
 266 |     /// Creates a new speed limiter.
 267 |     ///
 268 |     /// Use [infinity](`std::f64::INFINITY`) to make the speed unlimited.
 269 |     pub fn new(speed_limit: f64) -> Self {
 270 |         Builder::new(speed_limit).build()
 271 |     }
 272 | 
 273 |     /// Makes a [`Builder`] for further configurating this limiter.
 274 |     ///
 275 |     /// Use [infinity](`std::f64::INFINITY`) to make the speed unlimited.
 276 |     pub fn builder(speed_limit: f64) -> Builder<C> {
 277 |         Builder::new(speed_limit)
 278 |     }
 279 | 
 280 |     /// Returns the clock associated with this limiter.
 281 |     pub fn clock(&self) -> &C {
 282 |         &self.clock
 283 |     }
 284 | 
 285 |     /// Dynamically changes the speed limit. The new limit applies to all clones
 286 |     /// of this instance.
 287 |     ///
 288 |     /// Use [infinity](`std::f64::INFINITY`) to make the speed unlimited.
 289 |     ///
 290 |     /// This change will not affect any tasks scheduled _before_ this call.
 291 |     pub fn set_speed_limit(&self, speed_limit: f64) {
 292 |         debug_assert!(speed_limit > 0.0, "speed limit must be positive");
 293 |         self.bucket.lock().unwrap().set_speed_limit(speed_limit);
 294 |         self.is_unlimited
 295 |             .store(speed_limit.is_infinite(), Ordering::Relaxed);
 296 |     }
 297 | 
 298 |     /// Returns the current speed limit.
 299 |     ///
 300 |     /// This method returns [infinity](`std::f64::INFINITY`) if the speed is
 301 |     /// unlimited.
 302 |     pub fn speed_limit(&self) -> f64 {
 303 |         self.bucket.lock().unwrap().speed_limit
 304 |     }
 305 | 
 306 |     /// Obtains the total number of bytes consumed by this limiter so far.
 307 |     ///
 308 |     /// If more than `usize::MAX` bytes have been consumed, the count will wrap
 309 |     /// around.
 310 |     pub fn total_bytes_consumed(&self) -> usize {
 311 |         self.total_bytes_consumed.load(Ordering::Relaxed)
 312 |     }
 313 | 
 314 |     /// Resets the total number of bytes consumed to 0.
 315 |     pub fn reset_statistics(&self) {
 316 |         self.total_bytes_consumed.store(0, Ordering::Relaxed);
 317 |     }
 318 | 
 319 |     /// Consumes several bytes from the speed limiter, returns the duration
 320 |     /// needed to sleep to maintain the speed limit.
 321 |     pub fn consume_duration(&self, byte_size: usize) -> Duration {
 322 |         self.total_bytes_consumed
 323 |             .fetch_add(byte_size, Ordering::Relaxed);
 324 | 
 325 |         if self.is_unlimited.load(Ordering::Relaxed) {
 326 |             return Duration::from_secs(0);
 327 |         }
 328 | 
 329 |         #[allow(clippy::cast_precision_loss)]
 330 |         let size = byte_size as f64;
 331 | 
 332 |         // Using a lock should be fine,
 333 |         // as we're not blocking for a long time.
 334 |         let mut bucket = self.bucket.lock().unwrap();
 335 |         bucket.refill(self.clock.now());
 336 |         bucket.consume(size)
 337 |     }
 338 | 
 339 |     /// Reverts the consumption of the given bytes size.
 340 |     pub fn unconsume(&self, byte_size: usize) {
 341 |         self.total_bytes_consumed
 342 |             .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
 343 |                 Some(x.saturating_sub(byte_size))
 344 |             })
 345 |             .unwrap();
 346 | 
 347 |         if !self.is_unlimited.load(Ordering::Relaxed) {
 348 |             #[allow(clippy::cast_precision_loss)]
 349 |             let size = byte_size as f64;
 350 | 
 351 |             let mut bucket = self.bucket.lock().unwrap();
 352 |             bucket.unconsume(size);
 353 |         }
 354 |     }
 355 | 
 356 |     /// Consumes several bytes from the speed limiter.
 357 |     ///
 358 |     /// The consumption happens at the beginning, *before* the speed limit is
 359 |     /// applied. The returned future is fulfilled after the speed limit is
 360 |     /// satified.
 361 |     pub fn consume(&self, byte_size: usize) -> Consume<C, ()> {
 362 |         let sleep_dur = self.consume_duration(byte_size);
 363 |         // TODO use Duration::is_zero after `duration_zero` is stable.
 364 |         let future = if sleep_dur == Duration::from_secs(0) {
 365 |             None
 366 |         } else {
 367 |             Some(self.clock.sleep(sleep_dur))
 368 |         };
 369 |         Consume {
 370 |             future,
 371 |             result: Some(()),
 372 |         }
 373 |     }
 374 | 
 375 |     /// Wraps a streaming resource with speed limiting. See documentation of
 376 |     /// [`Resource`] for details.
 377 |     ///
 378 |     /// If you want to reuse the limiter after calling this function, `clone()`
 379 |     /// the limiter first.
 380 |     pub fn limit<R>(self, resource: R) -> Resource<R, C> {
 381 |         Resource::new(self, resource)
 382 |     }
 383 | 
 384 |     /// Returns the number of active clones of this limiter.
 385 |     ///
 386 |     /// Currently only used for testing, and thus not exported.
 387 |     #[cfg(test)]
 388 |     fn shared_count(&self) -> usize {
 389 |         Arc::strong_count(&self.bucket)
 390 |     }
 391 | }
 392 | 
 393 | impl<C: BlockingClock> Limiter<C> {
 394 |     /// Consumes several bytes, and sleeps the current thread to maintain the
 395 |     /// speed limit.
 396 |     ///
 397 |     /// The consumption happens at the beginning, *before* the speed limit is
 398 |     /// applied. This method blocks the current thread (e.g. using
 399 |     /// [`std::thread::sleep()`] given a [`StandardClock`]), and *must not* be
 400 |     /// used in `async` context.
 401 |     ///
 402 |     /// Prefer using this method instead of
 403 |     /// [`futures_executor::block_on`]`(limiter.`[`consume`](Limiter::consume())`(size))`.
 404 |     ///
 405 |     /// [`futures_executor::block_on`]: https://docs.rs/futures-executor/0.3/futures_executor/fn.block_on.html
 406 |     pub fn blocking_consume(&self, byte_size: usize) {
 407 |         let sleep_dur = self.consume_duration(byte_size);
 408 |         self.clock.blocking_sleep(sleep_dur);
 409 |     }
 410 | }
 411 | 
 412 | /// The future returned by [`Limiter::consume()`].
 413 | #[derive(Debug)]
 414 | pub struct Consume<C: Clock, R> {
 415 |     future: Option<C::Delay>,
 416 |     result: Option<R>,
 417 | }
 418 | 
 419 | #[allow(clippy::use_self)] // https://github.com/rust-lang/rust-clippy/issues/3410
 420 | impl<C: Clock, R> Consume<C, R> {
 421 |     /// Replaces the return value of the future.
 422 |     pub fn map<T, F: FnOnce(R) -> T>(self, f: F) -> Consume<C, T> {
 423 |         Consume {
 424 |             future: self.future,
 425 |             result: self.result.map(f),
 426 |         }
 427 |     }
 428 | }
 429 | 
 430 | impl<C: Clock, R: Unpin> Future for Consume<C, R> {
 431 |     type Output = R;
 432 | 
 433 |     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
 434 |         let this = self.get_mut();
 435 |         let is_ready = match &mut this.future {
 436 |             Some(future) => Pin::new(future).poll(cx).is_ready(),
 437 |             None => true,
 438 |         };
 439 |         if is_ready {
 440 |             if let Some(value) = this.result.take() {
 441 |                 return Poll::Ready(value);
 442 |             }
 443 |         }
 444 |         Poll::Pending
 445 |     }
 446 | }
 447 | 
 448 | #[cfg(feature = "fused-future")]
 449 | impl<C: Clock, R: Unpin> futures_core::future::FusedFuture for Consume<C, R> {
 450 |     fn is_terminated(&self) -> bool {
 451 |         self.result.is_none()
 452 |     }
 453 | }
 454 | 
 455 | pin_project! {
 456 |     /// A speed-limited wrapper of a byte stream.
 457 |     ///
 458 |     /// The `Resource` can be used to limit speed of
 459 |     ///
 460 |     /// * [`AsyncRead`](futures_io::AsyncRead)
 461 |     /// * [`AsyncWrite`](futures_io::AsyncWrite)
 462 |     ///
 463 |     /// Just like [`Limiter`], the delay is inserted *after* the data are sent
 464 |     /// or received, in which we know the exact amount of bytes transferred to
 465 |     /// give an accurate delay. The instantaneous speed can exceed the limit if
 466 |     /// many read/write tasks are started simultaneously. Therefore, restricting
 467 |     /// the concurrency is also important to avoid breaching the constraint.
 468 |     pub struct Resource<R, C: Clock> {
 469 |         limiter: Limiter<C>,
 470 |         #[pin]
 471 |         resource: R,
 472 |         waiter: Option<Consume<C, ()>>,
 473 |     }
 474 | }
 475 | 
 476 | impl<R, C: Clock> Resource<R, C> {
 477 |     /// Creates a new speed-limited resource.
 478 |     ///
 479 |     /// To make the resouce have unlimited speed, set the speed of [`Limiter`]
 480 |     /// to [infinity](`std::f64::INFINITY`).
 481 |     pub fn new(limiter: Limiter<C>, resource: R) -> Self {
 482 |         Self {
 483 |             limiter,
 484 |             resource,
 485 |             waiter: None,
 486 |         }
 487 |     }
 488 | 
 489 |     /// Unwraps this value, returns the underlying resource.
 490 |     pub fn into_inner(self) -> R {
 491 |         self.resource
 492 |     }
 493 | 
 494 |     /// Gets a reference to the underlying resource.
 495 |     ///
 496 |     /// It is inadvisable to directly operate the underlying resource.
 497 |     pub fn get_ref(&self) -> &R {
 498 |         &self.resource
 499 |     }
 500 | 
 501 |     /// Gets a mutable reference to the underlying resource.
 502 |     ///
 503 |     /// It is inadvisable to directly operate the underlying resource.
 504 |     pub fn get_mut(&mut self) -> &mut R {
 505 |         &mut self.resource
 506 |     }
 507 | 
 508 |     /// Gets a pinned reference to the underlying resource.
 509 |     ///
 510 |     /// It is inadvisable to directly operate the underlying resource.
 511 |     pub fn get_pin_ref(self: Pin<&Self>) -> Pin<&R> {
 512 |         self.project_ref().resource
 513 |     }
 514 | 
 515 |     /// Gets a pinned mutable reference to the underlying resource.
 516 |     ///
 517 |     /// It is inadvisable to directly operate the underlying resource.
 518 |     pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> {
 519 |         self.project().resource
 520 |     }
 521 | }
 522 | 
 523 | impl<R, C: Clock> Resource<R, C> {
 524 |     /// Wraps a poll function with a delay after it.
 525 |     ///
 526 |     /// This method calls the given `poll` function until it is fulfilled. After
 527 |     /// that, the result is saved into this `Resource` instance (therefore
 528 |     /// different `poll_***` calls should not be interleaving), while returning
 529 |     /// `Pending` until the limiter has completely consumed the result.
 530 |     #[allow(dead_code)]
 531 |     pub(crate) fn poll_limited<T, B>(
 532 |         self: Pin<&mut Self>,
 533 |         cx: &mut Context<'_>,
 534 |         mut buf: B,
 535 |         length: impl FnOnce(&T, &B) -> usize,
 536 |         poll: impl FnOnce(Pin<&mut R>, &mut Context<'_>, &mut B) -> Poll<T>,
 537 |     ) -> Poll<T> {
 538 |         let this = self.project();
 539 | 
 540 |         if let Some(waiter) = this.waiter {
 541 |             let res = Pin::new(waiter).poll(cx);
 542 |             if res.is_pending() {
 543 |                 return Poll::Pending;
 544 |             }
 545 |             *this.waiter = None;
 546 |         }
 547 | 
 548 |         let res = poll(this.resource, cx, &mut buf);
 549 |         if let Poll::Ready(obj) = &res {
 550 |             let len = length(obj, &buf);
 551 |             if len > 0 {
 552 |                 *this.waiter = Some(this.limiter.consume(len));
 553 |             }
 554 |         }
 555 |         res
 556 |     }
 557 | }
 558 | 
 559 | //------------------------------------------------------------------------------
 560 | 
 561 | #[cfg(test)]
 562 | mod tests_with_manual_clock {
 563 |     use super::*;
 564 |     use crate::clock::{Clock, ManualClock, Nanoseconds};
 565 |     use futures_executor::LocalPool;
 566 |     use futures_util::task::SpawnExt;
 567 |     use std::{future::Future, thread::panicking};
 568 | 
 569 |     /// Part of the `Fixture` which is to be shared with the spawned tasks.
 570 |     #[derive(Clone)]
 571 |     struct SharedFixture {
 572 |         limiter: Limiter<ManualClock>,
 573 |     }
 574 | 
 575 |     impl SharedFixture {
 576 |         fn now(&self) -> u64 {
 577 |             self.limiter.clock().now().0
 578 |         }
 579 | 
 580 |         fn sleep(&self, nanos: u64) -> impl Future<Output = ()> + '_ {
 581 |             self.limiter.clock().sleep(Duration::from_nanos(nanos))
 582 |         }
 583 | 
 584 |         fn consume(&self, bytes: usize) -> impl Future<Output = ()> + '_ {
 585 |             self.limiter.consume(bytes)
 586 |         }
 587 | 
 588 |         fn unconsume(&self, bytes: usize) {
 589 |             self.limiter.unconsume(bytes);
 590 |         }
 591 |     }
 592 | 
 593 |     /// The test fixture used by all test cases.
 594 |     struct Fixture {
 595 |         shared: SharedFixture,
 596 |         pool: LocalPool,
 597 |     }
 598 | 
 599 |     impl Fixture {
 600 |         fn new() -> Self {
 601 |             Self::with_min_wait(Duration::from_secs(1))
 602 |         }
 603 | 
 604 |         fn with_min_wait(min_wait: Duration) -> Self {
 605 |             Self {
 606 |                 shared: SharedFixture {
 607 |                     limiter: Limiter::builder(512.0)
 608 |                         .refill(Duration::from_secs(1))
 609 |                         .min_wait(min_wait)
 610 |                         .build(),
 611 |                 },
 612 |                 pool: LocalPool::new(),
 613 |             }
 614 |         }
 615 | 
 616 |         fn spawn<F, G>(&self, f: F)
 617 |         where
 618 |             F: FnOnce(SharedFixture) -> G,
 619 |             G: Future<Output = ()> + Send + 'static,
 620 |         {
 621 |             self.pool.spawner().spawn(f(self.shared.clone())).unwrap();
 622 |         }
 623 | 
 624 |         fn set_time(&mut self, time: u64) {
 625 |             self.shared.limiter.clock().set_time(Nanoseconds(time));
 626 |             self.pool.run_until_stalled();
 627 |         }
 628 | 
 629 |         fn set_speed_limit(&self, limit: f64) {
 630 |             self.shared.limiter.set_speed_limit(limit);
 631 |         }
 632 | 
 633 |         fn total_bytes_consumed(&self) -> usize {
 634 |             self.shared.limiter.total_bytes_consumed()
 635 |         }
 636 |     }
 637 | 
 638 |     impl Drop for Fixture {
 639 |         fn drop(&mut self) {
 640 |             if !panicking() {
 641 |                 // the count is 1 only if all spawned futures are finished.
 642 |                 assert_eq!(self.shared.limiter.shared_count(), 1);
 643 |             }
 644 |         }
 645 |     }
 646 | 
 647 |     #[test]
 648 |     fn under_limit_single_thread() {
 649 |         let mut fx = Fixture::new();
 650 | 
 651 |         fx.spawn(|sfx| async move {
 652 |             sfx.consume(50).await;
 653 |             assert_eq!(sfx.now(), 0);
 654 |             sfx.consume(51).await;
 655 |             assert_eq!(sfx.now(), 0);
 656 |             sfx.consume(52).await;
 657 |             assert_eq!(sfx.now(), 0);
 658 |             sfx.consume(53).await;
 659 |             assert_eq!(sfx.now(), 0);
 660 |             sfx.consume(54).await;
 661 |             assert_eq!(sfx.now(), 0);
 662 |             sfx.consume(55).await;
 663 |             assert_eq!(sfx.now(), 0);
 664 |         });
 665 | 
 666 |         fx.set_time(0);
 667 |         assert_eq!(fx.total_bytes_consumed(), 315);
 668 |     }
 669 | 
 670 |     #[test]
 671 |     fn over_limit_single_thread() {
 672 |         let mut fx = Fixture::new();
 673 | 
 674 |         fx.spawn(|sfx| {
 675 |             async move {
 676 |                 sfx.consume(200).await;
 677 |                 assert_eq!(sfx.now(), 0);
 678 |                 sfx.consume(201).await;
 679 |                 assert_eq!(sfx.now(), 0);
 680 |                 sfx.consume(202).await;
 681 |                 assert_eq!(sfx.now(), 1_177_734_375);
 682 |                 // 1_177_734_375 ns = (200+201+202)/512 seconds
 683 | 
 684 |                 sfx.consume(203).await;
 685 |                 assert_eq!(sfx.now(), 1_177_734_375);
 686 |                 sfx.consume(204).await;
 687 |                 assert_eq!(sfx.now(), 1_177_734_375);
 688 |                 sfx.consume(205).await;
 689 |                 assert_eq!(sfx.now(), 2_373_046_875);
 690 |             }
 691 |         });
 692 | 
 693 |         fx.set_time(0);
 694 |         assert_eq!(fx.total_bytes_consumed(), 603);
 695 |         fx.set_time(1_177_734_374);
 696 |         assert_eq!(fx.total_bytes_consumed(), 603);
 697 |         fx.set_time(1_177_734_375);
 698 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 699 |         fx.set_time(2_373_046_874);
 700 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 701 |         fx.set_time(2_373_046_875);
 702 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 703 |     }
 704 | 
 705 |     #[test]
 706 |     fn over_limit_single_thread_with_min_wait() {
 707 |         let mut fx = Fixture::with_min_wait(Duration::from_millis(100));
 708 | 
 709 |         fx.spawn(|sfx| {
 710 |             async move {
 711 |                 sfx.consume(200).await;
 712 |                 assert_eq!(sfx.now(), 0);
 713 |                 sfx.consume(201).await;
 714 |                 assert_eq!(sfx.now(), 0);
 715 |                 sfx.consume(202).await;
 716 |                 assert_eq!(sfx.now(), 277_734_375);
 717 |                 // 277_734_375 ns = 100_000_000 + (200+201+202-512)/512 seconds
 718 | 
 719 |                 sfx.consume(203).await;
 720 |                 assert_eq!(sfx.now(), 674_218_750);
 721 |                 sfx.consume(204).await;
 722 |                 assert_eq!(sfx.now(), 1_072_656_250);
 723 |                 sfx.consume(205).await;
 724 |                 assert_eq!(sfx.now(), 1_473_046_875);
 725 |             }
 726 |         });
 727 | 
 728 |         fx.set_time(0);
 729 |         assert_eq!(fx.total_bytes_consumed(), 603);
 730 |         fx.set_time(277_734_374);
 731 |         assert_eq!(fx.total_bytes_consumed(), 603);
 732 |         fx.set_time(277_734_375);
 733 |         assert_eq!(fx.total_bytes_consumed(), 806);
 734 |         fx.set_time(674_218_750);
 735 |         assert_eq!(fx.total_bytes_consumed(), 1010);
 736 |         fx.set_time(1_072_656_250);
 737 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 738 |         fx.set_time(1_473_046_875);
 739 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 740 |     }
 741 | 
 742 |     #[test]
 743 |     fn over_limit_multi_thread() {
 744 |         let mut fx = Fixture::new();
 745 | 
 746 |         // Due to how LocalPool does scheduling, the first task is always polled
 747 |         // before the second task. Nevertheless, the second task can still send
 748 |         // stuff using the timing difference.
 749 | 
 750 |         fx.spawn(|sfx| async move {
 751 |             sfx.consume(200).await;
 752 |             assert_eq!(sfx.now(), 0);
 753 |             sfx.consume(202).await;
 754 |             assert_eq!(sfx.now(), 0);
 755 |             sfx.consume(204).await;
 756 |             assert_eq!(sfx.now(), 1_183_593_750);
 757 |             sfx.consume(206).await;
 758 |             assert_eq!(sfx.now(), 1_183_593_750);
 759 |             sfx.consume(208).await;
 760 |             assert_eq!(sfx.now(), 2_384_765_625);
 761 |         });
 762 |         fx.spawn(|sfx| async move {
 763 |             sfx.consume(201).await;
 764 |             assert_eq!(sfx.now(), 1_576_171_875);
 765 |             sfx.consume(203).await;
 766 |             assert_eq!(sfx.now(), 2_781_250_000);
 767 |             sfx.consume(205).await;
 768 |             assert_eq!(sfx.now(), 2_781_250_000);
 769 |             sfx.consume(207).await;
 770 |             assert_eq!(sfx.now(), 2_781_250_000);
 771 |             sfx.consume(209).await;
 772 |             assert_eq!(sfx.now(), 3_994_140_625);
 773 |         });
 774 | 
 775 |         fx.set_time(0);
 776 |         assert_eq!(fx.total_bytes_consumed(), 807);
 777 |         fx.set_time(1_183_593_749);
 778 |         assert_eq!(fx.total_bytes_consumed(), 807);
 779 |         fx.set_time(1_183_593_750);
 780 |         assert_eq!(fx.total_bytes_consumed(), 1221);
 781 |         fx.set_time(1_576_171_874);
 782 |         assert_eq!(fx.total_bytes_consumed(), 1221);
 783 |         fx.set_time(1_576_171_875);
 784 |         assert_eq!(fx.total_bytes_consumed(), 1424);
 785 |         fx.set_time(2_384_765_624);
 786 |         assert_eq!(fx.total_bytes_consumed(), 1424);
 787 |         fx.set_time(2_384_765_625);
 788 |         assert_eq!(fx.total_bytes_consumed(), 1424);
 789 |         fx.set_time(2_781_249_999);
 790 |         assert_eq!(fx.total_bytes_consumed(), 1424);
 791 |         fx.set_time(2_781_250_000);
 792 |         assert_eq!(fx.total_bytes_consumed(), 2045);
 793 |         fx.set_time(3_994_140_624);
 794 |         assert_eq!(fx.total_bytes_consumed(), 2045);
 795 |         fx.set_time(3_994_140_625);
 796 |         assert_eq!(fx.total_bytes_consumed(), 2045);
 797 |     }
 798 | 
 799 |     #[test]
 800 |     fn over_limit_multi_thread_2() {
 801 |         let mut fx = Fixture::new();
 802 | 
 803 |         fx.spawn(|sfx| async move {
 804 |             sfx.consume(300).await;
 805 |             assert_eq!(sfx.now(), 0);
 806 |             sfx.consume(301).await;
 807 |             assert_eq!(sfx.now(), 1_173_828_125);
 808 |             sfx.consume(302).await;
 809 |             assert_eq!(sfx.now(), 1_173_828_125);
 810 |             sfx.consume(303).await;
 811 |             assert_eq!(sfx.now(), 2_550_781_250);
 812 |             sfx.consume(304).await;
 813 |             assert_eq!(sfx.now(), 2_550_781_250);
 814 |         });
 815 |         fx.spawn(|sfx| async move {
 816 |             sfx.consume(100).await;
 817 |             assert_eq!(sfx.now(), 1_369_140_625);
 818 |             sfx.consume(101).await;
 819 |             assert_eq!(sfx.now(), 2_748_046_875);
 820 |             sfx.consume(102).await;
 821 |             assert_eq!(sfx.now(), 2_748_046_875);
 822 |             sfx.consume(103).await;
 823 |             assert_eq!(sfx.now(), 2_748_046_875);
 824 |             sfx.consume(104).await;
 825 |             assert_eq!(sfx.now(), 3_945_312_500);
 826 |         });
 827 | 
 828 |         fx.set_time(0);
 829 |         assert_eq!(fx.total_bytes_consumed(), 701);
 830 |         fx.set_time(1_173_828_125);
 831 |         assert_eq!(fx.total_bytes_consumed(), 1306);
 832 |         fx.set_time(1_369_140_625);
 833 |         assert_eq!(fx.total_bytes_consumed(), 1407);
 834 |         fx.set_time(2_550_781_250);
 835 |         assert_eq!(fx.total_bytes_consumed(), 1711);
 836 |         fx.set_time(2_748_046_875);
 837 |         assert_eq!(fx.total_bytes_consumed(), 2020);
 838 |         fx.set_time(3_945_312_500);
 839 |         assert_eq!(fx.total_bytes_consumed(), 2020);
 840 |     }
 841 | 
 842 |     #[test]
 843 |     fn over_limit_multi_thread_yielded() {
 844 |         let mut fx = Fixture::new();
 845 | 
 846 |         // we're adding 1ns sleeps between each consume() to act as yield points,
 847 |         // so the consume() are evenly distributed, and can take advantage of
 848 |         // single bursting.
 849 | 
 850 |         fx.spawn(|sfx| async move {
 851 |             sfx.consume(300).await;
 852 |             assert_eq!(sfx.now(), 0);
 853 |             sfx.sleep(1).await;
 854 |             sfx.consume(301).await;
 855 |             assert_eq!(sfx.now(), 1_369_140_625);
 856 |             sfx.sleep(1).await;
 857 |             sfx.consume(302).await;
 858 |             assert_eq!(sfx.now(), 1_369_140_626);
 859 |             sfx.sleep(1).await;
 860 |             sfx.consume(303).await;
 861 |             assert_eq!(sfx.now(), 2_748_046_875);
 862 |             sfx.sleep(1).await;
 863 |             sfx.consume(304).await;
 864 |             assert_eq!(sfx.now(), 2_748_046_876);
 865 |         });
 866 |         fx.spawn(|sfx| async move {
 867 |             sfx.consume(100).await;
 868 |             assert_eq!(sfx.now(), 0);
 869 |             sfx.sleep(1).await;
 870 |             sfx.consume(101).await;
 871 |             assert_eq!(sfx.now(), 1_566_406_250);
 872 |             sfx.sleep(1).await;
 873 |             sfx.consume(102).await;
 874 |             assert_eq!(sfx.now(), 2_947_265_625);
 875 |             sfx.sleep(1).await;
 876 |             sfx.consume(103).await;
 877 |             assert_eq!(sfx.now(), 2_947_265_626);
 878 |             sfx.sleep(1).await;
 879 |             sfx.consume(104).await;
 880 |             assert_eq!(sfx.now(), 2_947_265_627);
 881 |         });
 882 | 
 883 |         fx.set_time(0);
 884 |         assert_eq!(fx.total_bytes_consumed(), 400);
 885 |         fx.set_time(1);
 886 |         assert_eq!(fx.total_bytes_consumed(), 802);
 887 |         fx.set_time(1_369_140_625);
 888 |         assert_eq!(fx.total_bytes_consumed(), 802);
 889 |         fx.set_time(1_369_140_626);
 890 |         assert_eq!(fx.total_bytes_consumed(), 1104);
 891 |         fx.set_time(1_566_406_250);
 892 |         assert_eq!(fx.total_bytes_consumed(), 1407);
 893 |         fx.set_time(1_566_406_251);
 894 |         assert_eq!(fx.total_bytes_consumed(), 1509);
 895 |         fx.set_time(2_748_046_875);
 896 |         assert_eq!(fx.total_bytes_consumed(), 1509);
 897 |         fx.set_time(2_748_046_876);
 898 |         assert_eq!(fx.total_bytes_consumed(), 1813);
 899 |         fx.set_time(2_947_265_625);
 900 |         assert_eq!(fx.total_bytes_consumed(), 1813);
 901 |         fx.set_time(2_947_265_626);
 902 |         assert_eq!(fx.total_bytes_consumed(), 1916);
 903 |         fx.set_time(2_947_265_627);
 904 |         assert_eq!(fx.total_bytes_consumed(), 2020);
 905 |     }
 906 | 
 907 |     #[test]
 908 |     fn unconsume() {
 909 |         let mut fx = Fixture::new();
 910 | 
 911 |         fx.spawn(|sfx| async move {
 912 |             sfx.consume(200).await;
 913 |             assert_eq!(sfx.now(), 0);
 914 |             sfx.consume(201).await;
 915 |             assert_eq!(sfx.now(), 0);
 916 |             sfx.unconsume(200);
 917 |             sfx.consume(202).await;
 918 |             assert_eq!(sfx.now(), 0);
 919 |             sfx.consume(200).await;
 920 |             assert_eq!(sfx.now(), 1_177_734_375);
 921 | 
 922 |             sfx.consume(203).await;
 923 |             assert_eq!(sfx.now(), 1_177_734_375);
 924 |             sfx.consume(204).await;
 925 |             assert_eq!(sfx.now(), 1_177_734_375);
 926 |             sfx.consume(205).await;
 927 |             assert_eq!(sfx.now(), 2_373_046_875);
 928 |             sfx.unconsume(2000);
 929 |         });
 930 | 
 931 |         fx.set_time(0);
 932 |         assert_eq!(fx.total_bytes_consumed(), 603);
 933 |         fx.set_time(1_177_734_374);
 934 |         assert_eq!(fx.total_bytes_consumed(), 603);
 935 |         fx.set_time(1_177_734_375);
 936 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 937 |         fx.set_time(2_373_046_874);
 938 |         assert_eq!(fx.total_bytes_consumed(), 1215);
 939 |         fx.set_time(2_373_046_875);
 940 |         assert_eq!(fx.total_bytes_consumed(), 0);
 941 |     }
 942 | 
 943 |     /// Ensures the speed limiter won't forget to enforce until a long pause
 944 |     /// i.e. we're observing the _maximum_ speed, not the _average_ speed.
 945 |     #[test]
 946 |     fn hiatus() {
 947 |         let mut fx = Fixture::new();
 948 | 
 949 |         fx.spawn(|sfx| async move {
 950 |             sfx.consume(400).await;
 951 |             assert_eq!(sfx.now(), 0);
 952 |             sfx.consume(401).await;
 953 |             assert_eq!(sfx.now(), 1_564_453_125);
 954 | 
 955 |             sfx.sleep(10_000_000_000).await;
 956 |             assert_eq!(sfx.now(), 11_564_453_125);
 957 | 
 958 |             sfx.consume(402).await;
 959 |             assert_eq!(sfx.now(), 11_564_453_125);
 960 |             sfx.consume(403).await;
 961 |             assert_eq!(sfx.now(), 13_136_718_750);
 962 |         });
 963 | 
 964 |         fx.set_time(0);
 965 |         assert_eq!(fx.total_bytes_consumed(), 801);
 966 |         fx.set_time(1_564_453_125);
 967 |         assert_eq!(fx.total_bytes_consumed(), 801);
 968 |         fx.set_time(11_564_453_125);
 969 |         assert_eq!(fx.total_bytes_consumed(), 1606);
 970 |         fx.set_time(13_136_718_750);
 971 |         assert_eq!(fx.total_bytes_consumed(), 1606);
 972 |     }
 973 | 
 974 |     // Ensures we could still send something much higher than the speed limit
 975 |     #[test]
 976 |     fn burst() {
 977 |         let mut fx = Fixture::new();
 978 | 
 979 |         fx.spawn(|sfx| async move {
 980 |             sfx.consume(5000).await;
 981 |             assert_eq!(sfx.now(), 9_765_625_000);
 982 |             sfx.consume(5001).await;
 983 |             assert_eq!(sfx.now(), 19_533_203_125);
 984 |             sfx.consume(5002).await;
 985 |             assert_eq!(sfx.now(), 29_302_734_375);
 986 |         });
 987 | 
 988 |         fx.set_time(0);
 989 |         assert_eq!(fx.total_bytes_consumed(), 5000);
 990 |         fx.set_time(9_765_625_000);
 991 |         assert_eq!(fx.total_bytes_consumed(), 10001);
 992 |         fx.set_time(19_533_203_125);
 993 |         assert_eq!(fx.total_bytes_consumed(), 15003);
 994 |         fx.set_time(29_302_734_375);
 995 |         assert_eq!(fx.total_bytes_consumed(), 15003);
 996 |     }
 997 | 
 998 |     #[test]
 999 |     fn change_speed_limit() {
1000 |         let mut fx = Fixture::new();
1001 | 
1002 |         // we try to send 5120 bytes at granularity of 256 bytes each.
1003 |         fx.spawn(|sfx| async move {
1004 |             for _ in 0..20 {
1005 |                 sfx.consume(256).await;
1006 |             }
1007 |         });
1008 | 
1009 |         // at first, we will send 512 B/s.
1010 |         fx.set_time(0);
1011 |         assert_eq!(fx.total_bytes_consumed(), 512);
1012 |         fx.set_time(500_000_000);
1013 |         assert_eq!(fx.total_bytes_consumed(), 512);
1014 |         fx.set_time(1_000_000_000);
1015 |         assert_eq!(fx.total_bytes_consumed(), 1024);
1016 |         fx.set_time(1_500_000_000);
1017 |         assert_eq!(fx.total_bytes_consumed(), 1024);
1018 | 
1019 |         // decrease the speed to 256 B/s
1020 |         fx.set_speed_limit(256.0);
1021 |         fx.set_time(1_500_000_001);
1022 |         assert_eq!(fx.total_bytes_consumed(), 1024);
1023 | 
1024 |         fx.set_time(2_000_000_000);
1025 |         assert_eq!(fx.total_bytes_consumed(), 1280);
1026 |         fx.set_time(2_500_000_000);
1027 |         assert_eq!(fx.total_bytes_consumed(), 1280);
1028 |         fx.set_time(3_000_000_000);
1029 |         assert_eq!(fx.total_bytes_consumed(), 1280);
1030 |         fx.set_time(3_500_000_000);
1031 |         assert_eq!(fx.total_bytes_consumed(), 1280);
1032 |         fx.set_time(4_000_000_000);
1033 |         assert_eq!(fx.total_bytes_consumed(), 1536);
1034 |         fx.set_time(4_500_000_000);
1035 |         assert_eq!(fx.total_bytes_consumed(), 1536);
1036 | 
1037 |         // increase the speed to 1024 B/s
1038 |         fx.set_speed_limit(1024.0);
1039 |         fx.set_time(4_500_000_001);
1040 |         assert_eq!(fx.total_bytes_consumed(), 1536);
1041 | 
1042 |         fx.set_time(5_000_000_000);
1043 |         assert_eq!(fx.total_bytes_consumed(), 2560);
1044 |         fx.set_time(5_500_000_000);
1045 |         assert_eq!(fx.total_bytes_consumed(), 2560);
1046 |         fx.set_time(6_000_000_000);
1047 |         assert_eq!(fx.total_bytes_consumed(), 3584);
1048 |         fx.set_time(6_500_000_000);
1049 |         assert_eq!(fx.total_bytes_consumed(), 3584);
1050 |         fx.set_time(7_000_000_000);
1051 |         assert_eq!(fx.total_bytes_consumed(), 4608);
1052 |         fx.set_time(7_500_000_000);
1053 |         assert_eq!(fx.total_bytes_consumed(), 4608);
1054 |         fx.set_time(8_000_000_000);
1055 |         assert_eq!(fx.total_bytes_consumed(), 5120);
1056 |     }
1057 | 
1058 |     /// Ensures lots of small takes won't prevent scheduling of a large take.
1059 |     #[test]
1060 |     fn thousand_cuts() {
1061 |         let mut fx = Fixture::new();
1062 | 
1063 |         fx.spawn(|sfx| async move {
1064 |             for _ in 0..64 {
1065 |                 sfx.consume(16).await;
1066 |             }
1067 |         });
1068 | 
1069 |         fx.spawn(|sfx| async move {
1070 |             sfx.consume(555).await;
1071 |             assert_eq!(sfx.now(), 2_083_984_375);
1072 |             sfx.consume(556).await;
1073 |             assert_eq!(sfx.now(), 3_201_171_875);
1074 |         });
1075 | 
1076 |         fx.set_time(0);
1077 |         assert_eq!(fx.total_bytes_consumed(), 1067);
1078 |         fx.set_time(1_000_000_000);
1079 |         assert_eq!(fx.total_bytes_consumed(), 1083);
1080 |         fx.set_time(2_000_000_000);
1081 |         assert_eq!(fx.total_bytes_consumed(), 1083);
1082 |         fx.set_time(2_083_984_375);
1083 |         assert_eq!(fx.total_bytes_consumed(), 1639);
1084 |         fx.set_time(3_000_000_000);
1085 |         assert_eq!(fx.total_bytes_consumed(), 2055);
1086 |         fx.set_time(3_201_171_875);
1087 |         assert_eq!(fx.total_bytes_consumed(), 2055);
1088 |         fx.set_time(4_000_000_000);
1089 |         assert_eq!(fx.total_bytes_consumed(), 2055);
1090 |         fx.set_time(4_169_921_875);
1091 |         assert_eq!(fx.total_bytes_consumed(), 2135);
1092 |     }
1093 | 
1094 |     #[test]
1095 |     fn set_infinite_speed_limit() {
1096 |         let mut fx = Fixture::new();
1097 | 
1098 |         fx.spawn(|sfx| async move {
1099 |             for _ in 0..1000 {
1100 |                 sfx.consume(512).await;
1101 |             }
1102 |             sfx.sleep(1).await;
1103 |             for _ in 0..1000 {
1104 |                 sfx.consume(512).await;
1105 |             }
1106 |             sfx.sleep(1).await;
1107 |             sfx.consume(512).await;
1108 |             sfx.consume(512).await;
1109 |         });
1110 | 
1111 |         fx.set_time(0);
1112 |         assert_eq!(fx.total_bytes_consumed(), 512);
1113 |         fx.set_time(1_000_000_000);
1114 |         assert_eq!(fx.total_bytes_consumed(), 1024);
1115 | 
1116 |         // change speed limit to infinity...
1117 |         fx.set_speed_limit(std::f64::INFINITY);
1118 | 
1119 |         // should not affect tasks still waiting
1120 |         fx.set_time(1_500_000_000);
1121 |         assert_eq!(fx.total_bytes_consumed(), 1024);
1122 | 
1123 |         // but all future consumptions will be unlimited.
1124 |         fx.set_time(2_000_000_000);
1125 |         assert_eq!(fx.total_bytes_consumed(), 512_000);
1126 | 
1127 |         // should act normal for keeping speed limit at infinity.
1128 |         fx.set_speed_limit(std::f64::INFINITY);
1129 |         fx.set_time(2_000_000_001);
1130 |         assert_eq!(fx.total_bytes_consumed(), 1_024_000);
1131 | 
1132 |         // reducing speed limit to normal.
1133 |         fx.set_speed_limit(512.0);
1134 |         fx.set_time(2_000_000_002);
1135 |         assert_eq!(fx.total_bytes_consumed(), 1_024_512);
1136 |         fx.set_time(3_000_000_002);
1137 |         assert_eq!(fx.total_bytes_consumed(), 1_025_024);
1138 |         fx.set_time(4_000_000_002);
1139 |         assert_eq!(fx.total_bytes_consumed(), 1_025_024);
1140 |     }
1141 | }
1142 | 
1143 | #[cfg(test)]
1144 | #[cfg(feature = "standard-clock")]
1145 | mod tests_with_standard_clock {
1146 |     use super::*;
1147 |     use futures_executor::LocalPool;
1148 |     use futures_util::{future::join_all, task::SpawnExt};
1149 |     use rand::{thread_rng, Rng};
1150 |     use std::time::Instant;
1151 | 
1152 |     // This test case is ported from RocksDB.
1153 |     #[test]
1154 |     fn rate() {
1155 |         eprintln!("tests_with_standard_clock::rate() will run for 20 seconds, please be patient");
1156 | 
1157 |         let mut pool = LocalPool::new();
1158 |         let sp = pool.spawner();
1159 | 
1160 |         for &i in &[1, 2, 4, 8, 16] {
1161 |             let target = i * 10_240;
1162 | 
1163 |             let limiter = <Limiter>::new(target as f64);
1164 |             for &speed_limit in &[target, target * 2] {
1165 |                 limiter.reset_statistics();
1166 |                 limiter.set_speed_limit(speed_limit as f64);
1167 |                 let start = Instant::now();
1168 | 
1169 |                 let handles = (0..i).map(|_| {
1170 |                     let limiter = limiter.clone();
1171 |                     sp.spawn_with_handle(async move {
1172 |                         // tests for 2 seconds.
1173 |                         let until = Instant::now() + Duration::from_secs(2);
1174 |                         while Instant::now() < until {
1175 |                             let size = thread_rng().gen_range(1..=target / 10);
1176 |                             limiter.consume(size).await;
1177 |                         }
1178 |                     })
1179 |                     .unwrap()
1180 |                 });
1181 | 
1182 |                 pool.run_until(join_all(handles));
1183 |                 assert_eq!(limiter.shared_count(), 1);
1184 | 
1185 |                 let elapsed = start.elapsed();
1186 |                 let speed = limiter.total_bytes_consumed() as f64 / elapsed.as_secs_f64();
1187 |                 let diff_ratio = speed / speed_limit as f64;
1188 |                 eprintln!(
1189 |                     "rate: {} threads, expected speed {} B/s, actual speed {:.0} B/s, elapsed {:?}",
1190 |                     i, speed_limit, speed, elapsed
1191 |                 );
1192 |                 assert!((0.80..=1.25).contains(&diff_ratio));
1193 |                 assert!(elapsed <= Duration::from_secs(4));
1194 |             }
1195 |         }
1196 |     }
1197 | 
1198 |     #[test]
1199 |     fn block() {
1200 |         eprintln!("tests_with_standard_clock::block() will run for 20 seconds, please be patient");
1201 | 
1202 |         for &i in &[1, 2, 4, 8, 16] {
1203 |             let target = i * 10_240;
1204 | 
1205 |             let limiter = <Limiter>::new(target as f64);
1206 |             for &speed_limit in &[target, target * 2] {
1207 |                 limiter.reset_statistics();
1208 |                 limiter.set_speed_limit(speed_limit as f64);
1209 |                 let start = Instant::now();
1210 | 
1211 |                 let handles = (0..i)
1212 |                     .map(|_| {
1213 |                         let limiter = limiter.clone();
1214 |                         std::thread::spawn(move || {
1215 |                             // tests for 2 seconds.
1216 |                             let until = Instant::now() + Duration::from_secs(2);
1217 |                             while Instant::now() < until {
1218 |                                 let size = thread_rng().gen_range(1..=target / 10);
1219 |                                 limiter.blocking_consume(size);
1220 |                             }
1221 |                         })
1222 |                     })
1223 |                     .collect::<Vec<_>>();
1224 | 
1225 |                 for jh in handles {
1226 |                     jh.join().unwrap();
1227 |                 }
1228 | 
1229 |                 assert_eq!(limiter.shared_count(), 1);
1230 | 
1231 |                 let elapsed = start.elapsed();
1232 |                 let speed = limiter.total_bytes_consumed() as f64 / elapsed.as_secs_f64();
1233 |                 let diff_ratio = speed / speed_limit as f64;
1234 |                 eprintln!(
1235 |                     "block: {} threads, expected speed {} B/s, actual speed {:.0} B/s, elapsed {:?}",
1236 |                     i, speed_limit, speed, elapsed
1237 |                 );
1238 |                 assert!((0.80..=1.25).contains(&diff_ratio));
1239 |                 assert!(elapsed <= Duration::from_secs(4));
1240 |             }
1241 |         }
1242 |     }
1243 | }
1244 | 


--------------------------------------------------------------------------------