Skip to content

Commit 5bfa986

Browse files
committed
feat: add interval_stream helper
1 parent 0f2d8ef commit 5bfa986

File tree

6 files changed

+270
-5
lines changed

6 files changed

+270
-5
lines changed

Cargo.lock

Lines changed: 165 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "toolbox"
3-
version = "0.1.0"
3+
version = "0.2.0"
44
authors = ["Oliver Chalk"]
55
edition = "2021"
66
readme = "README.md"
@@ -14,13 +14,18 @@ disallowed_methods = "warn"
1414

1515
[features]
1616
default = []
17-
tokio = ["dep:tokio"]
17+
named_task = ["tokio/rt"]
18+
interval_stream = ["tokio/time", "dep:futures"]
1819
tracing = ["dep:const_format", "dep:tracing", "dep:tracing-appender", "dep:tracing-subscriber"]
1920
version = ["dep:const_format"]
2021

2122
[dependencies]
2223
const_format = { version = "0.2.32", optional = true }
23-
tokio = { version = "1.0", features = ["rt"], optional = true }
24+
futures = { version = "0.3.31", optional = true }
25+
tokio = { version = "1.0", optional = true }
2426
tracing = { version = "0.1.40", optional = true }
2527
tracing-appender = { version = "0.2.3", optional = true }
2628
tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json"], optional = true }
29+
30+
[dev-dependencies]
31+
tokio-test = "0.4.4"

src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
pub mod fs;
2-
#[cfg(feature = "tokio")]
2+
#[cfg(any(feature = "named_task", feature = "interval_stream"))]
33
pub mod tokio;
44
#[cfg(feature = "tracing")]
55
pub mod tracing;

src/tokio/interval_stream.rs

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
use std::future::Future;
2+
use std::pin::pin;
3+
use std::task::Poll;
4+
5+
use futures::{ready, Stream};
6+
7+
/// Creates and polls a future on an interval, producing a stream.
8+
///
9+
/// # Example
10+
///
11+
/// ```rust
12+
/// # tokio_test::block_on(async {
13+
/// use futures::FutureExt;
14+
/// use futures::StreamExt;
15+
/// use toolbox::tokio::IntervalStream;
16+
///
17+
/// let slot_interval = tokio::time::interval(std::time::Duration::from_millis(1));
18+
/// let mut slot_stream = IntervalStream::new(
19+
/// slot_interval,
20+
/// Box::new(|| futures::future::ready("hello world").boxed()),
21+
/// );
22+
///
23+
/// assert_eq!(slot_stream.next().await.unwrap(), "hello world");
24+
/// assert_eq!(slot_stream.next().await.unwrap(), "hello world");
25+
/// # })
26+
/// ```
27+
pub struct IntervalStream<Fut>
28+
where
29+
Fut: Unpin,
30+
{
31+
interval: tokio::time::Interval,
32+
poll: Box<dyn Fn() -> Fut>,
33+
34+
in_progress: Option<Fut>,
35+
}
36+
37+
impl<Fut, Output> IntervalStream<Fut>
38+
where
39+
Fut: Future<Output = Output> + Unpin,
40+
{
41+
pub fn new(interval: tokio::time::Interval, poll: Box<dyn Fn() -> Fut>) -> Self {
42+
IntervalStream { interval, poll, in_progress: None }
43+
}
44+
}
45+
46+
impl<Fut, Output> Stream for IntervalStream<Fut>
47+
where
48+
Fut: Future<Output = Output> + Unpin,
49+
{
50+
type Item = Output;
51+
52+
fn poll_next(
53+
self: std::pin::Pin<&mut Self>,
54+
cx: &mut std::task::Context<'_>,
55+
) -> std::task::Poll<Option<Self::Item>> {
56+
let this = self.get_mut();
57+
58+
// Poll the current future if one already exists.
59+
if let Some(fut) = &mut this.in_progress {
60+
let output = ready!(pin!(fut).poll(cx));
61+
this.in_progress = None;
62+
63+
return Poll::Ready(Some(output));
64+
}
65+
66+
// Poll the interval to see if we should create a new future.
67+
ready!(this.interval.poll_tick(cx));
68+
69+
// Create a new future.
70+
let mut fut = (this.poll)();
71+
72+
// Poll the future.
73+
let pinned = pin!(&mut fut);
74+
let poll = pinned.poll(cx);
75+
match poll {
76+
Poll::Ready(output) => Poll::Ready(Some(output)),
77+
Poll::Pending => {
78+
this.in_progress = Some(fut);
79+
80+
Poll::Pending
81+
}
82+
}
83+
}
84+
}

src/tokio/mod.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
/// Creates and polls a future on an interval, producing a stream.
2+
#[cfg(feature = "interval_stream")]
3+
mod interval_stream;
4+
/// Allows attaching names to [`tokio::task`] to enable tracking which tasks are
5+
/// exiting.
6+
#[cfg(feature = "named_task")]
7+
mod named_task;
8+
9+
#[cfg(feature = "interval_stream")]
10+
pub use interval_stream::*;
11+
#[cfg(feature = "named_task")]
12+
pub use named_task::*;
File renamed without changes.

0 commit comments

Comments
 (0)