pax_global_header00006660000000000000000000000064150714505330014515gustar00rootroot0000000000000052 comment=701b04478f55b02e14abfbf8f8f15239238d7ffd dmitryvk-async-fn-stream-701b044/000077500000000000000000000000001507145053300166365ustar00rootroot00000000000000dmitryvk-async-fn-stream-701b044/.github/000077500000000000000000000000001507145053300201765ustar00rootroot00000000000000dmitryvk-async-fn-stream-701b044/.github/workflows/000077500000000000000000000000001507145053300222335ustar00rootroot00000000000000dmitryvk-async-fn-stream-701b044/.github/workflows/checks.yml000066400000000000000000000025011507145053300242140ustar00rootroot00000000000000on: push: branches: [main] pull_request: types: [opened, reopened, synchronize] concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} cancel-in-progress: true name: checks jobs: fmt: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Install rust (stable) uses: dtolnay/rust-toolchain@stable with: components: rustfmt - run: cargo fmt --check clippy: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Install rust (stable) uses: dtolnay/rust-toolchain@stable with: components: clippy - run: cargo clippy --all-targets --workspace -- --deny warnings test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Install rust (stable) uses: dtolnay/rust-toolchain@stable - run: cargo test --all-targets --workspace test-miri: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Install rust (nightly) uses: dtolnay/rust-toolchain@nightly with: components: miri - run: cargo miri test semver: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Check semver uses: obi1kenobi/cargo-semver-checks-action@v2 dmitryvk-async-fn-stream-701b044/.gitignore000066400000000000000000000000751507145053300206300ustar00rootroot00000000000000/target /Cargo.lock /perf.data /perf.data.old /flamegraph.svgdmitryvk-async-fn-stream-701b044/CHANGELOG.md000066400000000000000000000032301507145053300204450ustar00rootroot00000000000000# Changelog All notable changes to this project will be documented in this file. ## [0.3.2] - 2025-10-08 ### ๐Ÿš€ Features - Support `emit_result` for `try_fn_stream` (#16) ## [0.3.1] - 2025-09-30 ### ๐Ÿ› Bug Fixes - Restore Send and Sync for streams (#13) - Add missing auto-traits ### ๐Ÿ“š Documentation - Fix intra-doc links ### โšก Performance - Don't use mutex (#12) ## [0.3.0] - 2025-09-05 ### ๐Ÿš€ Features - [**breaking**] The yielding is done in CollectFuture instead of `.emit()` - Support internal concurrency within async stream closures - Improve support for concurrency (#9) ### ๐Ÿ› Bug Fixes - Call Waker from CollectFuture ### ๐Ÿšœ Refactor - Move synchronization logic to CollectFuture - Don't clear waker in CollectFuture - [**breaking**] Rename all instances of `Collect` to `Emit` (#10) ### ๐Ÿ“š Documentation - Document concurrency support ### โšก Performance - Avoid cloning Wakers when possible ### ๐Ÿงช Testing - Add tests for structured concurrency ### โš™๏ธ Miscellaneous Tasks - Enable clippy pedantic lints (#8) - Fix lints - Add a couple of benchmarks ## [0.2.2] - 2024-04-23 ### ๐Ÿš€ Features - Implement emit_err for try_fn_stream ### ๐Ÿšœ Refactor - Use internal_emit instead of duplicate code ### ๐Ÿ“š Documentation - Improve documentation and some panic messages ### ๐Ÿงช Testing - Add test for emit_err ### โš™๏ธ Miscellaneous Tasks - Bump version ## [0.2.1] - 2024-04-18 ### ๐Ÿšœ Refactor - Replace futures by futures-util and futures-executor ### ๐Ÿ“š Documentation - Replace futures by futures-util ### ๐Ÿงช Testing - Remove unnecessary import ## [0.1.0] - 2022-07-07 dmitryvk-async-fn-stream-701b044/Cargo.toml000066400000000000000000000013721507145053300205710ustar00rootroot00000000000000[package] name = "async-fn-stream" description = "Lightweight implementation of `async-stream` without macros" version = "0.3.2" edition = "2021" license = "MIT" homepage = "https://github.com/dmitryvk/async-fn-stream" repository = "https://github.com/dmitryvk/async-fn-stream" keywords = ["async", "stream"] [lints.clippy] pedantic = "deny" [[bench]] name = "bench" harness = false [profile.bench] debug = true [dependencies] futures-util = { version = "0.3", default-features = false, features = ["std"] } pin-project-lite = "0.2" smallvec = "1.15.1" [dev-dependencies] anyhow = "1" async-stream = "0.3.6" futures-executor = { version = "0.3", default-features = false, features = ["std", "thread-pool"] } tokio = { version = "1.0", features = ["full"] } dmitryvk-async-fn-stream-701b044/LICENSE000066400000000000000000000020361507145053300176440ustar00rootroot00000000000000Copyright 2022 Dmitry Kalyanov Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.dmitryvk-async-fn-stream-701b044/README.md000066400000000000000000000066241507145053300201250ustar00rootroot00000000000000A version of [async-stream](https://github.com/tokio-rs/async-stream) without macros. This crate provides generic implementations of `Stream` trait. `Stream` is an asynchronous version of `std::iter::Iterator`. Two functions are provided - `fn_stream` and `try_fn_stream`. # Usage ## Basic Usage If you need to create a stream that may result in error, use `try_fn_stream`, otherwise use `fn_stream`. To create a stream: 1. Invoke `fn_stream` or `try_fn_stream`, passing a closure (anonymous function). 2. Closure will accept an `emitter`. To return value from the stream, call `.emit(value)` on `emitter` and `.await` on its result. Once stream consumer has processed the value and called `.next()` on stream, `.await` will return. ## Returning errors `try_fn_stream` provides some conveniences for returning errors: 1. Errors can be return from closure via `return Err(...)` or the question mark (`?`) operator. This will end the stream. 2. An `emitter` also has an `emit_err()` method to return errors without ending the stream. ## Limitations `fn_stream` does not support cross-task streams, that is all stream values must be produced in the same task as the stream. Specifically, it is not supported to move `StreamEmitter` to another thread or tokio task. If your use case necessitates this, consider using async chanells for that. ## Advanced usage Internal concurrency is supported within `fn_stream` (see [examples/join.rs](examples/join.rs)). # Examples Finite stream of numbers ```rust use async_fn_stream::fn_stream; use futures_util::Stream; fn build_stream() -> impl Stream { fn_stream(|emitter| async move { for i in 0..3 { // yield elements from stream via `emitter` emitter.emit(i).await; } }) } ``` Read numbers from text file, with error handling ```rust use anyhow::Context; use async_fn_stream::try_fn_stream; use futures_util::{pin_mut, Stream, StreamExt}; use tokio::{ fs::File, io::{AsyncBufReadExt, BufReader}, }; fn read_numbers(file_name: String) -> impl Stream> { try_fn_stream(|emitter| async move { // Return errors via `?` operator. let file = BufReader::new(File::open(file_name).await.context("Failed to open file")?); pin_mut!(file); let mut line = String::new(); loop { line.clear(); let byte_count = file .read_line(&mut line) .await .context("Failed to read line")?; if byte_count == 0 { break; } for token in line.split_ascii_whitespace() { let Ok(number) = token.parse::() else { // Return errors via the `emit_err` method. emitter.emit_err( anyhow::anyhow!("Failed to convert string \"{token}\" to number") ).await; continue; }; emitter.emit(number).await; } } Ok(()) }) } ``` # Why not `async-stream`? [async-stream](https://github.com/tokio-rs/async-stream) is great! It has a nice syntax, but it is based on macros which brings some flaws: * proc-macros sometimes interact badly with IDEs such as `rust-analyzer` or `RustRover`. see e.g. * proc-macros may increase build times dmitryvk-async-fn-stream-701b044/benches/000077500000000000000000000000001507145053300202455ustar00rootroot00000000000000dmitryvk-async-fn-stream-701b044/benches/bench.rs000066400000000000000000000045761507145053300217060ustar00rootroot00000000000000use std::time::Instant; use async_stream::stream; use futures_util::{pin_mut, StreamExt}; #[tokio::main] async fn main() { let num_iters: u32 = std::env::var("ASYNC_FN_STREAM_BENCH_ITERS") .ok() .and_then(|s| s.parse().ok()) .unwrap_or(1); for _ in 0..num_iters { bench_async_fn_stream_sync_numbers().await; bench_async_fn_stream_sync_strings().await; bench_async_stream_sync_numbers().await; bench_async_stream_sync_strings().await; } } async fn bench_async_fn_stream_sync_numbers() { let start = Instant::now(); let stream = async_fn_stream::fn_stream(|emitter| async move { for i in 0..1_000_000 { emitter.emit(i).await; } }); pin_mut!(stream); let mut sum = 0u32; while let Some(i) = stream.next().await { sum = sum.wrapping_add(i); } assert!(sum > 0); let elapsed = start.elapsed(); println!("async-fn-stream sync numbers: {} us", elapsed.as_micros()); } async fn bench_async_fn_stream_sync_strings() { let start = Instant::now(); let stream = async_fn_stream::fn_stream(|emitter| async move { for i in 0..1_000_000 { emitter.emit(i.to_string()).await; } }); pin_mut!(stream); let mut sum = 0usize; while let Some(i) = stream.next().await { sum = sum.wrapping_add(i.len()); } assert!(sum > 0); let elapsed = start.elapsed(); println!("async-fn-stream sync strings: {} us", elapsed.as_micros()); } async fn bench_async_stream_sync_numbers() { let start = Instant::now(); let stream = stream! { for i in 0..1_000_000 { yield i; } }; pin_mut!(stream); let mut sum = 0u32; while let Some(i) = stream.next().await { sum = sum.wrapping_add(i); } assert!(sum > 0); let elapsed = start.elapsed(); println!("async-stream sync numbers: {} us", elapsed.as_micros()); } async fn bench_async_stream_sync_strings() { let start = Instant::now(); let stream = stream! { for i in 0..1_000_000 { yield i.to_string(); } }; pin_mut!(stream); let mut sum = 0usize; while let Some(i) = stream.next().await { sum = sum.wrapping_add(i.len()); } assert!(sum > 0); let elapsed = start.elapsed(); println!("async-stream sync strings: {} us", elapsed.as_micros()); } dmitryvk-async-fn-stream-701b044/cliff.toml000066400000000000000000000064771507145053300206340ustar00rootroot00000000000000# git-cliff ~ default configuration file # https://git-cliff.org/docs/configuration # # Lines starting with "#" are comments. # Configuration options are organized into tables and keys. # See documentation for more information on available options. [changelog] # changelog header header = """ # Changelog\n All notable changes to this project will be documented in this file.\n """ # template for the changelog body # https://keats.github.io/tera/docs/#introduction body = """ {% if version %}\ ## [{{ version | trim_start_matches(pat="v") }}] - {{ timestamp | date(format="%Y-%m-%d") }} {% else %}\ ## [unreleased] {% endif %}\ {% for group, commits in commits | group_by(attribute="group") %} ### {{ group | striptags | trim | upper_first }} {% for commit in commits %} - {% if commit.scope %}*({{ commit.scope }})* {% endif %}\ {% if commit.breaking %}[**breaking**] {% endif %}\ {{ commit.message | upper_first }}\ {% endfor %} {% endfor %}\n """ # template for the changelog footer footer = """ """ # remove the leading and trailing s trim = true # postprocessors postprocessors = [ # { pattern = '', replace = "https://github.com/orhun/git-cliff" }, # replace repository URL ] [git] # parse the commits based on https://www.conventionalcommits.org conventional_commits = true # filter out the commits that are not conventional filter_unconventional = true # process each line of a commit as an individual commit split_commits = false # regex for preprocessing the commit messages commit_preprocessors = [ # Replace issue numbers #{ pattern = '\((\w+\s)?#([0-9]+)\)', replace = "([#${2}](/issues/${2}))"}, # Check spelling of the commit with https://github.com/crate-ci/typos # If the spelling is incorrect, it will be automatically fixed. #{ pattern = '.*', replace_command = 'typos --write-changes -' }, ] # regex for parsing and grouping commits commit_parsers = [ { message = "^feat", group = "๐Ÿš€ Features" }, { message = "^fix", group = "๐Ÿ› Bug Fixes" }, { message = "^doc", group = "๐Ÿ“š Documentation" }, { message = "^perf", group = "โšก Performance" }, { message = "^refactor", group = "๐Ÿšœ Refactor" }, { message = "^style", group = "๐ŸŽจ Styling" }, { message = "^test", group = "๐Ÿงช Testing" }, { message = "^chore\\(release\\): prepare for", skip = true }, { message = "^chore\\(deps.*\\)", skip = true }, { message = "^chore\\(pr\\)", skip = true }, { message = "^chore\\(pull\\)", skip = true }, { message = "^chore|^ci", group = "โš™๏ธ Miscellaneous Tasks" }, { body = ".*security", group = "๐Ÿ›ก๏ธ Security" }, { message = "^revert", group = "โ—€๏ธ Revert" }, ] # protect breaking changes from being skipped due to matching a skipping commit_parser protect_breaking_commits = false # filter out the commits that are not matched by commit parsers filter_commits = true # regex for matching git tags # tag_pattern = "v[0-9].*" # regex for skipping tags # skip_tags = "" # regex for ignoring tags # ignore_tags = "" # sort the tags topologically topo_order = false # sort the commits inside sections by oldest/newest order sort_commits = "oldest" # limit the number of commits included in the changelog. # limit_commits = 42 dmitryvk-async-fn-stream-701b044/examples/000077500000000000000000000000001507145053300204545ustar00rootroot00000000000000dmitryvk-async-fn-stream-701b044/examples/file_numbers.rs000066400000000000000000000024711507145053300235000ustar00rootroot00000000000000use std::env::args; use anyhow::Context; use async_fn_stream::try_fn_stream; use futures_util::{pin_mut, Stream, StreamExt}; use tokio::{ fs::File, io::{AsyncBufReadExt, BufReader}, }; #[tokio::main] async fn main() -> Result<(), anyhow::Error> { let file_name = args().nth(1).unwrap(); let stream = read_numbers(file_name); pin_mut!(stream); while let Some(number) = stream.next().await { println!("number: {}", number?); } Ok(()) } fn read_numbers(file_name: String) -> impl Stream> { try_fn_stream(|emitter| async move { let file = BufReader::new(File::open(file_name).await.context("Failed to open file")?); pin_mut!(file); let mut line = String::new(); loop { line.clear(); let byte_count = file .read_line(&mut line) .await .context("Failed to read line")?; if byte_count == 0 { break; } for token in line.split_ascii_whitespace() { let number: i32 = token .parse() .with_context(|| format!("Failed to conver string \"{token}\" to number"))?; emitter.emit(number).await; } } Ok(()) }) } dmitryvk-async-fn-stream-701b044/examples/join.rs000066400000000000000000000017461507145053300217710ustar00rootroot00000000000000use async_fn_stream::fn_stream; use futures_util::{pin_mut, Stream, StreamExt}; fn build_stream() -> impl Stream { fn_stream(|emitter| async move { tokio::join!( async { for i in 0..3 { // yield elements from stream via `emitter` emitter.emit(i).await; } }, async { for i in 10..13 { // yield elements from stream via `emitter` emitter.emit(i).await; } } ); }) } async fn example() { let stream = build_stream(); pin_mut!(stream); let mut numbers = Vec::new(); while let Some(number) = stream.next().await { print!("{number} "); numbers.push(number); } println!(); numbers.sort_unstable(); assert_eq!(numbers, vec![0, 1, 2, 10, 11, 12]); } #[tokio::main] async fn main() { futures_executor::block_on(example()); } dmitryvk-async-fn-stream-701b044/examples/numbers.rs000066400000000000000000000012071507145053300224750ustar00rootroot00000000000000use async_fn_stream::fn_stream; use futures_util::{pin_mut, Stream, StreamExt}; fn build_stream() -> impl Stream { fn_stream(|emitter| async move { for i in 0..3 { // yield elements from stream via `emitter` emitter.emit(i).await; } }) } async fn example() { let stream = build_stream(); pin_mut!(stream); let mut numbers = Vec::new(); while let Some(number) = stream.next().await { print!("{number} "); numbers.push(number); } println!(); assert_eq!(numbers, vec![0, 1, 2]); } pub fn main() { futures_executor::block_on(example()); } dmitryvk-async-fn-stream-701b044/src/000077500000000000000000000000001507145053300174255ustar00rootroot00000000000000dmitryvk-async-fn-stream-701b044/src/lib.rs000066400000000000000000000670471507145053300205570ustar00rootroot00000000000000#![doc = include_str!("../README.md")] use std::{ cell::{Cell, UnsafeCell}, panic::{RefUnwindSafe, UnwindSafe}, pin::Pin, sync::Arc, task::{Poll, Waker}, }; use futures_util::{Future, Stream}; use pin_project_lite::pin_project; use smallvec::SmallVec; /// An intermediary that transfers values from stream to its consumer pub struct StreamEmitter { inner: Arc>>, } /// An intermediary that transfers values from stream to its consumer pub struct TryStreamEmitter { inner: Arc>>>, } thread_local! { /// A type-erased pointer to the `Inner<_>`, for which a call to `Stream::poll_next` is active. static ACTIVE_STREAM_INNER: Cell<*const ()> = const { Cell::new(std::ptr::null()) }; } /// A guard that ensures that `ACTIVE_STREAM_INNER` is returned to its previous value even in case of a panic. struct ActiveStreamPointerGuard { old_ptr: *const (), } impl ActiveStreamPointerGuard { fn set_active_ptr(ptr: *const ()) -> Self { let old_ptr = ACTIVE_STREAM_INNER.with(|thread_ptr| thread_ptr.replace(ptr)); Self { old_ptr } } } impl Drop for ActiveStreamPointerGuard { fn drop(&mut self) { ACTIVE_STREAM_INNER.with(|thread_ptr| thread_ptr.set(self.old_ptr)); } } /// SAFETY: /// `Inner` stores the state shared between `StreamEmitter`/`TryStreamEmitter` and `FnStream`/`TryFnStream`. /// It is stored within `Arc>`. Exclusive access to it is controlled using: /// 1) exclusive reference `Pin<&mut FnStream>` which is provided to `FnStream::poll_next` /// 2) `ACTIVE_STREAM_INNER`, which is itself managed by `FnStream::poll_next`. /// - `ACTIVE_STREAM_INNER` is only set when the corresponding `Pin<&mut FnStream>` is on the stack /// - `EmitFuture` can safely access its inner reference if it equals to `ACTIVE_STREAM_INNER` (all other accesses are invalid and lead to panics) struct Inner { // `stream_waker` is used to compare the waker of the stream future with the waker of the `Emit` future. // If the stream implementation does not have sub-executors, we don't have to push wakers to `pending_wakers`, which avoids cloning it. stream_waker: Option, // Due to internal concurrency, a single call to stream's future may yield multiple elements. // All elements are stored here and yielded from the stream before polling the future again. pending_values: SmallVec<[T; 1]>, pending_wakers: SmallVec<[Waker; 1]>, } // SAFETY: `FnStream` implements own synchronization unsafe impl + Send> Send for FnStream {} // SAFETY: `FnStream` implements own synchronization unsafe impl + Sync> Sync for FnStream {} impl + UnwindSafe> UnwindSafe for FnStream {} impl + RefUnwindSafe> RefUnwindSafe for FnStream { } // SAFETY: `TryFnStream` implements own synchronization unsafe impl> + Send> Send for TryFnStream { } // SAFETY: `TryFnStream` implements own synchronization unsafe impl> + Sync> Sync for TryFnStream { } impl> + UnwindSafe> UnwindSafe for TryFnStream { } impl> + RefUnwindSafe> RefUnwindSafe for TryFnStream { } // SAFETY: `StreamEmitter` implements own synchronization unsafe impl Send for StreamEmitter {} // SAFETY: `StreamEmitter` implements own synchronization unsafe impl Sync for StreamEmitter {} impl UnwindSafe for StreamEmitter {} impl RefUnwindSafe for StreamEmitter {} // SAFETY: `TryStreamEmitter` implements own synchronization unsafe impl Send for TryStreamEmitter {} // SAFETY: `TryStreamEmitter` implements own synchronization unsafe impl Sync for TryStreamEmitter {} impl UnwindSafe for TryStreamEmitter {} impl RefUnwindSafe for TryStreamEmitter {} // SAFETY: `EmitFuture` implements own synchronization. unsafe impl Send for EmitFuture<'_, T> {} // SAFETY: `EmitFuture` implements own synchronization unsafe impl Sync for EmitFuture<'_, T> {} impl UnwindSafe for EmitFuture<'_, T> {} impl RefUnwindSafe for EmitFuture<'_, T> {} pin_project! { /// Implementation of [`Stream`] trait created by [`fn_stream`]. pub struct FnStream> { #[pin] fut: Fut, inner: Arc>>, } } /// Create a new infallible stream which is implemented by `func`. /// /// Caller should pass an async function which will return successive stream elements via [`StreamEmitter::emit`]. /// /// # Example /// /// ```rust /// use async_fn_stream::fn_stream; /// use futures_util::Stream; /// /// fn build_stream() -> impl Stream { /// fn_stream(|emitter| async move { /// for i in 0..3 { /// // yield elements from stream via `emitter` /// emitter.emit(i).await; /// } /// }) /// } /// ``` pub fn fn_stream>( func: impl FnOnce(StreamEmitter) -> Fut, ) -> FnStream { FnStream::new(func) } impl> FnStream { fn new) -> Fut>(func: F) -> Self { let inner = Arc::new(UnsafeCell::new(Inner { stream_waker: None, pending_values: SmallVec::new(), pending_wakers: SmallVec::new(), })); let emitter = StreamEmitter { inner: inner.clone(), }; let fut = func(emitter); Self { fut, inner } } } impl> Stream for FnStream { type Item = T; fn poll_next( self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> Poll> { let this = self.project(); // SAFETY: // (see safety comment for `Inner`) // 1) we have no aliasing, since we're holding unique reference to `Self`, and // 2) `this.inner` is not deallocated for the duration of this method let inner = unsafe { &mut *this.inner.get() }; if let Some(value) = inner.pending_values.pop() { return Poll::Ready(Some(value)); } if !inner.pending_wakers.is_empty() { for waker in inner.pending_wakers.drain(..) { if !waker.will_wake(cx.waker()) { waker.wake(); } } } if let Some(stream_waker) = inner.stream_waker.as_mut() { stream_waker.clone_from(cx.waker()); } else { inner.stream_waker = Some(cx.waker().clone()); } // SAFETY: ensure that we're not holding a reference to this.inner _ = inner; // SAFETY for Inner: // - `ACTIVE_STREAM_INNER` now contains valid pointer to `this.inner` during the call to `fut.poll` // - `ACTIVE_STREAM_INNER` is restored after the call due to the use of guard let polling_ptr_guard = ActiveStreamPointerGuard::set_active_ptr(Arc::as_ptr(&*this.inner).cast()); let r = this.fut.poll(cx); drop(polling_ptr_guard); // SAFETY: // (see safety comment for `Inner`) // 1) we have no aliasing, since: // - we're holding unique reference to `Self`, and // - we removed the pointer from `ACTIVE_STREAM_INNER` // 2) `this.inner` is not deallocated for the duration of this method let inner = unsafe { &mut *this.inner.get() }; match r { std::task::Poll::Ready(()) => Poll::Ready(None), std::task::Poll::Pending => { if let Some(value) = inner.pending_values.pop() { Poll::Ready(Some(value)) } else { Poll::Pending } } } } } /// Create a new fallible stream which is implemented by `func`. /// /// Caller should pass an async function which can: /// /// - return successive stream elements via [`TryStreamEmitter::emit`] /// - return transient errors via [`TryStreamEmitter::emit_err`] /// - return fatal errors as [`Result::Err`] /// /// # Example /// ```rust /// use async_fn_stream::try_fn_stream; /// use futures_util::Stream; /// /// fn build_stream() -> impl Stream> { /// try_fn_stream(|emitter| async move { /// for i in 0..3 { /// // yield elements from stream via `emitter` /// emitter.emit(i).await; /// } /// /// // return errors view emitter without ending the stream /// emitter.emit_err(anyhow::anyhow!("An error happened")); /// /// // return errors from stream, ending the stream /// Err(anyhow::anyhow!("An error happened")) /// }) /// } /// ``` pub fn try_fn_stream>>( func: impl FnOnce(TryStreamEmitter) -> Fut, ) -> TryFnStream { TryFnStream::new(func) } pin_project! { /// Implementation of [`Stream`] trait created by [`try_fn_stream`]. pub struct TryFnStream>> { is_err: bool, #[pin] fut: Fut, inner: Arc>>>, } } impl>> TryFnStream { fn new) -> Fut>(func: F) -> Self { let inner = Arc::new(UnsafeCell::new(Inner { stream_waker: None, pending_values: SmallVec::new(), pending_wakers: SmallVec::new(), })); let emitter = TryStreamEmitter { inner: inner.clone(), }; let fut = func(emitter); Self { is_err: false, fut, inner, } } } impl>> Stream for TryFnStream { type Item = Result; fn poll_next( self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> Poll> { // TODO: merge the implementation with `FnStream` if self.is_err { return Poll::Ready(None); } let this = self.project(); // SAFETY: // (see safety comment for `Inner`) // 1) we have no aliasing, since we're holding unique reference to `Self`, and // 2) `this.inner` is not deallocated for the duration of this method let inner = unsafe { &mut *this.inner.get() }; if let Some(value) = inner.pending_values.pop() { return Poll::Ready(Some(value)); } if !inner.pending_wakers.is_empty() { for waker in inner.pending_wakers.drain(..) { if !waker.will_wake(cx.waker()) { waker.wake(); } } } if let Some(stream_waker) = inner.stream_waker.as_mut() { stream_waker.clone_from(cx.waker()); } else { inner.stream_waker = Some(cx.waker().clone()); } // SAFETY: ensure that we're not holding a reference to this.inner _ = inner; // SAFETY for Inner: // - `ACTIVE_STREAM_INNER` now contains valid pointer to `this.inner` during the call to `fut.poll` // - `ACTIVE_STREAM_INNER` is restored after the call due to the use of guard let polling_ptr_guard = ActiveStreamPointerGuard::set_active_ptr(Arc::as_ptr(&*this.inner).cast()); let r = this.fut.poll(cx); drop(polling_ptr_guard); // SAFETY: // (see safety comment for `Inner`) // 1) we have no aliasing, since: // - we're holding unique reference to `Self`, and // - we removed the pointer from `ACTIVE_STREAM_INNER` // 2) `this.inner` is not deallocated for the duration of this method let inner = unsafe { &mut *this.inner.get() }; match r { std::task::Poll::Ready(Ok(())) => Poll::Ready(None), std::task::Poll::Ready(Err(e)) => { *this.is_err = true; Poll::Ready(Some(Err(e))) } std::task::Poll::Pending => { if let Some(value) = inner.pending_values.pop() { Poll::Ready(Some(value)) } else { Poll::Pending } } } } } impl StreamEmitter { /// Emit value from a stream and wait until stream consumer calls [`futures_util::StreamExt::next`] again. /// /// # Panics /// Will panic if: /// * `emit` is called not in context of polling the stream #[must_use = "Ensure that emit() is awaited"] pub fn emit(&'_ self, value: T) -> EmitFuture<'_, T> { EmitFuture::new(&self.inner, value) } } impl TryStreamEmitter { /// Emit value from a stream and wait until stream consumer calls [`futures_util::StreamExt::next`] again. /// /// # Panics /// Will panic if: /// * `emit` is called not in context of polling the stream #[must_use = "Ensure that emit() is awaited"] pub fn emit(&'_ self, value: T) -> EmitFuture<'_, Result> { EmitFuture::new(&self.inner, Ok(value)) } /// Emit value from a stream and wait until stream consumer calls [`futures_util::StreamExt::next`] again. /// /// # Panics /// Will panic if: /// * `emit_result` is called not in context of polling the stream #[must_use = "Ensure that emit() is awaited"] pub fn emit_result(&'_ self, value: Result) -> EmitFuture<'_, Result> { EmitFuture::new(&self.inner, value) } /// Emit error from a stream and wait until stream consumer calls [`futures_util::StreamExt::next`] again. /// /// # Panics /// Will panic if: /// * `emit_err` is called not in context of polling the stream #[must_use = "Ensure that emit_err() is awaited"] pub fn emit_err(&'_ self, err: E) -> EmitFuture<'_, Result> { EmitFuture::new(&self.inner, Err(err)) } } pin_project! { /// Future returned from [`StreamEmitter::emit`]. pub struct EmitFuture<'a, T> { inner: &'a UnsafeCell>, value: Option, } } impl<'a, T> EmitFuture<'a, T> { fn new(inner: &'a UnsafeCell>, value: T) -> Self { Self { inner, value: Some(value), } } } impl Future for EmitFuture<'_, T> { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { let this = self.project(); assert!( ACTIVE_STREAM_INNER.get() == std::ptr::from_ref(*this.inner).cast::<()>(), "StreamEmitter::emit().await should only be called in the context of the corresponding `fn_stream()`/`try_fn_stream()`" ); // SAFETY: // 1) we hold a unique reference to `this.inner` since we verified that ACTIVE_STREAM_INNER == self.inner // - we're calling `{Try,}FnStream::poll_next` in this thread, which holds the only other reference to the same instance of `Inner` // - `{Try,}FnStream::poll_next` is not holding a reference to `self.inner` during the call to `fut.poll` // 2) `this.inner` is not deallocated for the duration of this method let inner = unsafe { &mut *this.inner.get() }; if let Some(value) = this.value.take() { inner.pending_values.push(value); let is_same_waker = if let Some(stream_waker) = inner.stream_waker.as_ref() { stream_waker.will_wake(cx.waker()) } else { false }; if !is_same_waker { inner.pending_wakers.push(cx.waker().clone()); } Poll::Pending } else if inner.pending_values.is_empty() { // stream only polls the future after draining `inner.pending_values`, so this check should not be necessary in theory; // this is just a safeguard against misuses; e.g. if a future calls `.emit().poll()` in a loop without yielding on `Poll::Pending`, // this would lead to overflow of `inner.pending_values` Poll::Ready(()) } else { Poll::Pending } } } #[cfg(test)] mod tests { use std::{io::ErrorKind, pin::pin}; use futures_util::{pin_mut, stream::FuturesUnordered, StreamExt}; use super::*; #[test] fn infallible_works() { futures_executor::block_on(async { let stream = fn_stream(|emitter| async move { eprintln!("stream 1"); emitter.emit(1).await; eprintln!("stream 2"); emitter.emit(2).await; eprintln!("stream 3"); }); pin_mut!(stream); assert_eq!(Some(1), stream.next().await); assert_eq!(Some(2), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn infallible_lifetime() { let a = 1; futures_executor::block_on(async { let b = 2; let a = &a; let b = &b; let stream = fn_stream(|emitter| async move { eprintln!("stream 1"); emitter.emit(a).await; eprintln!("stream 2"); emitter.emit(b).await; eprintln!("stream 3"); }); pin_mut!(stream); assert_eq!(Some(a), stream.next().await); assert_eq!(Some(b), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn infallible_unawaited_emit_is_ignored() { futures_executor::block_on(async { #[expect( unused_must_use, reason = "this code intentionally does not await emitter.emit()" )] let stream = fn_stream(|emitter| async move { emitter.emit(1)/* .await */; emitter.emit(2)/* .await */; emitter.emit(3).await; }); pin_mut!(stream); assert_eq!(Some(3), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn fallible_works() { futures_executor::block_on(async { let stream = try_fn_stream(|emitter| async move { eprintln!("try stream 1"); emitter.emit(1).await; eprintln!("try stream 2"); emitter.emit(2).await; eprintln!("try stream 3"); Err(std::io::Error::from(ErrorKind::Other)) }); pin_mut!(stream); assert_eq!(1, stream.next().await.unwrap().unwrap()); assert_eq!(2, stream.next().await.unwrap().unwrap()); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.is_none()); }); } #[test] fn fallible_emit_err_works() { futures_executor::block_on(async { let stream = try_fn_stream(|emitter| async move { eprintln!("try stream 1"); emitter.emit(1).await; eprintln!("try stream 2"); emitter.emit_result(Ok(2)).await; eprintln!("try stream 3"); emitter .emit_err(std::io::Error::from(ErrorKind::Other)) .await; eprintln!("try stream 4"); emitter .emit_result(Err(std::io::Error::from(ErrorKind::Other))) .await; eprintln!("try stream 5"); Err(std::io::Error::from(ErrorKind::Other)) }); pin_mut!(stream); assert_eq!(1, stream.next().await.unwrap().unwrap()); assert_eq!(2, stream.next().await.unwrap().unwrap()); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.is_none()); }); } #[test] fn method_async() { struct St { a: String, } impl St { async fn f1(&self) -> impl Stream { self.f2().await } #[allow(clippy::unused_async)] async fn f2(&self) -> impl Stream { fn_stream(|emitter| async move { emitter.emit(self.a.as_str()).await; emitter.emit(self.a.as_str()).await; emitter.emit(self.a.as_str()).await; }) } } futures_executor::block_on(async { let l = St { a: "qwe".to_owned(), }; let s = l.f1().await; let z: Vec<&str> = s.collect().await; assert_eq!(z, ["qwe", "qwe", "qwe"]); }); } #[test] fn tokio_join_one_works() { futures_executor::block_on(async { let stream = fn_stream(|emitter| async move { tokio::join!(async { emitter.emit(1).await },); emitter.emit(2).await; }); pin_mut!(stream); assert_eq!(Some(1), stream.next().await); assert_eq!(Some(2), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn tokio_join_many_works() { futures_executor::block_on(async { let stream = fn_stream(|emitter| async move { eprintln!("try stream 1"); tokio::join!( async { emitter.emit(1).await }, async { emitter.emit(2).await }, async { emitter.emit(3).await }, ); emitter.emit(4).await; }); pin_mut!(stream); for _ in 0..3 { let item = stream.next().await; assert!(matches!(item, Some(1..=3))); } assert_eq!(Some(4), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn tokio_futures_unordered_one_works() { futures_executor::block_on(async { let stream = fn_stream(|emitter| async move { let mut futs: FuturesUnordered<_> = (1..=1) .map(|i| { let emitter = &emitter; async move { emitter.emit(i).await } }) .collect(); while futs.next().await.is_some() {} emitter.emit(2).await; }); pin_mut!(stream); assert_eq!(Some(1), stream.next().await); assert_eq!(Some(2), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn tokio_futures_unordered_many_works() { futures_executor::block_on(async { let stream = fn_stream(|emitter| async move { let mut futs: FuturesUnordered<_> = (1..=3) .map(|i| { let emitter = &emitter; async move { emitter.emit(i).await } }) .collect(); while futs.next().await.is_some() {} emitter.emit(4).await; }); pin_mut!(stream); for _ in 1..=3 { let item = stream.next().await; assert!(matches!(item, Some(1..=3))); } assert_eq!(Some(4), stream.next().await); assert_eq!(None, stream.next().await); }); } #[test] fn infallible_nested_streams_work() { futures_executor::block_on(async { let mut stream = pin!(fn_stream(|emitter| async move { for i in 0..3 { let mut stream_2 = pin!(fn_stream(|emitter| async move { for j in 0..3 { emitter.emit(j).await; } })); while let Some(item) = stream_2.next().await { emitter.emit(3 * i + item).await; } } })); let mut sum = 0; while let Some(item) = stream.next().await { sum += item; } assert_eq!(sum, 36); }); } #[test] fn fallible_nested_streams_work() { futures_executor::block_on(async { let mut stream = pin!(try_fn_stream(|emitter| async move { for i in 0..3 { let mut stream_2 = pin!(try_fn_stream(|emitter| async move { for j in 0..3 { emitter.emit(j).await; } Ok::<_, ()>(()) })); while let Some(Ok(item)) = stream_2.next().await { emitter.emit(3 * i + item).await; } } Ok::<_, ()>(()) })); let mut sum = 0; while let Some(Ok(item)) = stream.next().await { sum += item; } assert_eq!(sum, 36); }); } #[test] #[should_panic( expected = "StreamEmitter::emit().await should only be called in the context of the corresponding `fn_stream()`/`try_fn_stream()`" )] fn infallible_bad_nested_emit_detected() { futures_executor::block_on(async { let mut stream = pin!(fn_stream(|emitter| async move { for i in 0..3 { let emitter_ref = &emitter; let mut stream_2 = pin!(fn_stream(|emitter_2| async move { emitter_2.emit(0).await; for j in 0..3 { emitter_ref.emit(j).await; } })); while let Some(item) = stream_2.next().await { emitter.emit(3 * i + item).await; } } })); let mut sum = 0; while let Some(item) = stream.next().await { sum += item; } assert_eq!(sum, 36); }); } #[test] #[should_panic( expected = "StreamEmitter::emit().await should only be called in the context of the corresponding `fn_stream()`/`try_fn_stream()`" )] fn fallible_bad_nested_emit_detected() { futures_executor::block_on(async { let mut stream = pin!(try_fn_stream(|emitter| async move { for i in 0..3 { let emitter_ref = &emitter; let mut stream_2 = pin!(try_fn_stream(|emitter_2| async move { emitter_2.emit(0).await; for j in 0..3 { emitter_ref.emit(j).await; } Ok::<_, ()>(()) })); while let Some(Ok(item)) = stream_2.next().await { emitter.emit(3 * i + item).await; } } Ok::<_, ()>(()) })); let mut sum = 0; while let Some(Ok(item)) = stream.next().await { sum += item; } assert_eq!(sum, 36); }); } }