fix: remove thread locals to fix unsoundness outlined in #1

Author: decahedron1

.github/workflows/test.yml

@@ -28,15 +28,3 @@ jobs:
       - uses: Swatinem/rust-cache@v1
       - name: Run tests
         run: cargo test
-  test-no-std:
-    name: Run tests (unstable-thread-local)
-    runs-on: ubuntu-24.04
-    steps:
-      - uses: actions/checkout@v4
-      - name: Install nightly Rust toolchain
-        uses: dtolnay/rust-toolchain@nightly
-        with:
-          toolchain: nightly-2024-11-11
-      - uses: Swatinem/rust-cache@v1
-      - name: Run tests
-        run: cargo test --features unstable-thread-local

Cargo.toml

@@ -17,4 +17,3 @@ tokio = { version = "1", features = [ "full", "macros" ] }
 futures = "0.3"
 
 [features]
-unstable-thread-local = []

README.md

@@ -8,9 +8,9 @@ use futures_util::{pin_mut, stream::StreamExt};
 
 #[tokio::main]
 async fn main() {
-	let stream = async_stream(|r#yield| async move {
+	let stream = async_stream(|yielder| async move {
 		for i in 0..3 {
-			r#yield(i).await;
+			yielder.r#yield(i).await;
 		}
 	});
 	pin_mut!(stream);
@@ -21,4 +21,4 @@ async fn main() {
 ```
 
 ## `#![no_std]` support
-`async-stream-lite` supports `#![no_std]` on nightly Rust (due to the usage of [the unstable `#[thread_local]` attribute](https://doc.rust-lang.org/beta/unstable-book/language-features/thread-local.html)). To enable `#![no_std]` support, enable the `unstable-thread-local` feature.
+`async-stream-lite` supports `#![no_std]`, but requires `alloc`.

src/lib.rs

@@ -1,15 +1,16 @@
 #![allow(clippy::tabs_in_doc_comments)]
-#![cfg_attr(feature = "unstable-thread-local", feature(thread_local))]
-#![cfg_attr(all(not(test), feature = "unstable-thread-local"), no_std)]
+#![cfg_attr(not(test), no_std)]
 
+extern crate alloc;
 extern crate core;
 
+use alloc::sync::{Arc, Weak};
 use core::{
 	cell::Cell,
 	future::Future,
 	marker::PhantomData,
 	pin::Pin,
-	ptr,
+	sync::atomic::{AtomicBool, Ordering},
 	task::{Context, Poll}
 };
 
@@ -19,91 +20,96 @@ use futures_core::stream::{FusedStream, Stream};
 mod tests;
 mod r#try;
 
-#[cfg(not(feature = "unstable-thread-local"))]
-thread_local! {
-	static STORE: Cell<*mut ()> = const { Cell::new(ptr::null_mut()) };
+pub(crate) struct SharedStore<T> {
+	entered: AtomicBool,
+	cell: Cell<Option<T>>
 }
-#[cfg(feature = "unstable-thread-local")]
-#[thread_local]
-static STORE: Cell<*mut ()> = Cell::new(ptr::null_mut());
 
-pub(crate) fn r#yield<T>(value: T) -> YieldFut<T> {
-	YieldFut { value: Some(value) }
+impl<T> Default for SharedStore<T> {
+	fn default() -> Self {
+		Self {
+			entered: AtomicBool::new(false),
+			cell: Cell::new(None)
+		}
+	}
+}
+
+impl<T> SharedStore<T> {
+	pub fn has_value(&self) -> bool {
+		unsafe { &*self.cell.as_ptr() }.is_some()
+	}
+}
+
+unsafe impl<T> Sync for SharedStore<T> {}
+
+pub struct Yielder<T> {
+	pub(crate) store: Weak<SharedStore<T>>
+}
+
+impl<T> Yielder<T> {
+	pub fn r#yield(&self, value: T) -> YieldFut<'_, T> {
+		#[cold]
+		fn invalid_usage() -> ! {
+			panic!("attempted to use async_stream_lite yielder outside of stream context or across threads")
+		}
+
+		let Some(store) = self.store.upgrade() else {
+			invalid_usage();
+		};
+		if !store.entered.load(Ordering::Relaxed) {
+			invalid_usage();
+		}
+
+		store.cell.replace(Some(value));
+
+		YieldFut { store, _p: PhantomData }
+	}
 }
 
 /// Future returned by an [`AsyncStream`]'s yield function.
 ///
 /// This future must be `.await`ed inside the generator in order for the item to be yielded by the stream.
 #[must_use = "stream will not yield this item unless the future returned by yield is awaited"]
-pub struct YieldFut<T> {
-	value: Option<T>
+pub struct YieldFut<'y, T> {
+	store: Arc<SharedStore<T>>,
+	_p: PhantomData<&'y ()>
 }
 
-impl<T> Unpin for YieldFut<T> {}
+impl<T> Unpin for YieldFut<'_, T> {}
 
-impl<T> Future for YieldFut<T> {
+impl<T> Future for YieldFut<'_, T> {
 	type Output = ();
 
-	fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
-		if self.value.is_none() {
+	fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+		if !self.store.has_value() {
 			return Poll::Ready(());
 		}
 
-		fn op<T>(cell: &Cell<*mut ()>, value: &mut Option<T>) {
-			let ptr = cell.get().cast::<Option<T>>();
-			let option_ref = unsafe { ptr.as_mut() }.expect("attempted to use async_stream yielder outside of stream context or across threads");
-			if option_ref.is_none() {
-				*option_ref = value.take();
-			}
-		}
-
-		#[cfg(not(feature = "unstable-thread-local"))]
-		return STORE.with(|cell| {
-			op(cell, &mut self.value);
-			Poll::Pending
-		});
-		#[cfg(feature = "unstable-thread-local")]
-		{
-			op(&STORE, &mut self.value);
-			Poll::Pending
-		}
+		Poll::Pending
 	}
 }
 
-struct Enter<'a, T> {
-	_p: PhantomData<&'a T>,
-	prev: *mut ()
+struct Enter<'s, T> {
+	store: &'s SharedStore<T>
 }
 
-fn enter<T>(dst: &mut Option<T>) -> Enter<'_, T> {
-	fn op<T>(cell: &Cell<*mut ()>, dst: &mut Option<T>) -> *mut () {
-		let prev = cell.get();
-		cell.set((dst as *mut Option<T>).cast::<()>());
-		prev
-	}
-	#[cfg(not(feature = "unstable-thread-local"))]
-	let prev = STORE.with(|cell| op(cell, dst));
-	#[cfg(feature = "unstable-thread-local")]
-	let prev = op(&STORE, dst);
-	Enter { _p: PhantomData, prev }
+fn enter<T>(store: &SharedStore<T>) -> Enter<'_, T> {
+	store.entered.store(true, Ordering::Relaxed);
+	Enter { store }
 }
 
 impl<T> Drop for Enter<'_, T> {
 	fn drop(&mut self) {
-		#[cfg(not(feature = "unstable-thread-local"))]
-		STORE.with(|cell| cell.set(self.prev));
-		#[cfg(feature = "unstable-thread-local")]
-		STORE.set(self.prev);
+		self.store.entered.store(false, Ordering::Relaxed);
 	}
 }
 
 pin_project_lite::pin_project! {
 	/// A [`Stream`] created from an asynchronous generator-like function.
 	///
 	/// To create an [`AsyncStream`], use the [`async_stream`] function.
-	#[derive(Debug)]
 	pub struct AsyncStream<T, U> {
-		_p: PhantomData<T>,
+		store: Arc<SharedStore<T>>,
 		done: bool,
 		#[pin]
 		generator: U
@@ -131,16 +137,15 @@ where
 			return Poll::Ready(None);
 		}
 
-		let mut dst = None;
 		let res = {
-			let _enter = enter(&mut dst);
+			let _enter = enter(&me.store);
 			me.generator.poll(cx)
 		};
 
 		*me.done = res.is_ready();
 
-		if dst.is_some() {
-			return Poll::Ready(dst.take());
+		if me.store.has_value() {
+			return Poll::Ready(me.store.cell.take());
 		}
 
 		if *me.done { Poll::Ready(None) } else { Poll::Pending }
@@ -153,16 +158,16 @@ where
 
 /// Create an asynchronous [`Stream`] from an asynchronous generator function.
 ///
-/// The provided function will be given a "yielder" function, which, when called, causes the stream to yield an item:
+/// The provided function will be given a [`Yielder`], which, when called, causes the stream to yield an item:
 /// ```
 /// use async_stream_lite::async_stream;
 /// use futures::{pin_mut, stream::StreamExt};
 ///
 /// #[tokio::main]
 /// async fn main() {
-/// 	let stream = async_stream(|r#yield| async move {
+/// 	let stream = async_stream(|yielder| async move {
 /// 		for i in 0..3 {
-/// 			r#yield(i).await;
+/// 			yielder.r#yield(i).await;
 /// 		}
 /// 	});
 /// 	pin_mut!(stream);
@@ -181,9 +186,9 @@ where
 /// };
 ///
 /// fn zero_to_three() -> impl Stream<Item = u32> {
-/// 	async_stream(|r#yield| async move {
+/// 	async_stream(|yielder| async move {
 /// 		for i in 0..3 {
-/// 			r#yield(i).await;
+/// 			yielder.r#yield(i).await;
 /// 		}
 /// 	})
 /// }
@@ -207,9 +212,9 @@ where
 /// };
 ///
 /// fn zero_to_three() -> BoxStream<'static, u32> {
-/// 	Box::pin(async_stream(|r#yield| async move {
+/// 	Box::pin(async_stream(|yielder| async move {
 /// 		for i in 0..3 {
-/// 			r#yield(i).await;
+/// 			yielder.r#yield(i).await;
 /// 		}
 /// 	}))
 /// }
@@ -232,18 +237,18 @@ where
 /// };
 ///
 /// fn zero_to_three() -> impl Stream<Item = u32> {
-/// 	async_stream(|r#yield| async move {
+/// 	async_stream(|yielder| async move {
 /// 		for i in 0..3 {
-/// 			r#yield(i).await;
+/// 			yielder.r#yield(i).await;
 /// 		}
 /// 	})
 /// }
 ///
 /// fn double<S: Stream<Item = u32>>(input: S) -> impl Stream<Item = u32> {
-/// 	async_stream(|r#yield| async move {
+/// 	async_stream(|yielder| async move {
 /// 		pin_mut!(input);
 /// 		while let Some(value) = input.next().await {
-/// 			r#yield(value * 2).await;
+/// 			yielder.r#yield(value * 2).await;
 /// 		}
 /// 	})
 /// }
@@ -261,15 +266,12 @@ where
 /// See also [`try_async_stream`], a variant of [`async_stream`] which supports try notation (`?`).
 pub fn async_stream<T, F, U>(generator: F) -> AsyncStream<T, U>
 where
-	F: FnOnce(fn(value: T) -> YieldFut<T>) -> U,
+	F: FnOnce(Yielder<T>) -> U,
 	U: Future<Output = ()>
 {
-	let generator = generator(r#yield::<T>);
-	AsyncStream {
-		_p: PhantomData,
-		done: false,
-		generator
-	}
+	let store = Arc::new(SharedStore::default());
+	let generator = generator(Yielder { store: Arc::downgrade(&store) });
+	AsyncStream { store, done: false, generator }
 }
 
 pub use self::r#try::{TryAsyncStream, try_async_stream};