fix: remove thread local

Author: decahedron1

Cargo.toml

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

README.md

@@ -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,21 +20,49 @@ 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());
+
+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> {
-	_p: PhantomData<T>
+	pub(crate) store: Weak<SharedStore<T>>
 }
 
 impl<T> Yielder<T> {
 	pub fn r#yield(&self, value: T) -> YieldFut<'_, T> {
-		YieldFut { value: Some(value), _p: PhantomData }
+		#[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 }
 	}
 }
 
@@ -42,7 +71,7 @@ impl<T> Yielder<T> {
 /// 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<'y, T> {
-	value: Option<T>,
+	store: Arc<SharedStore<T>>,
 	_p: PhantomData<&'y ()>
 }
 
@@ -51,66 +80,36 @@ impl<T> Unpin 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
@@ -138,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 }
@@ -271,12 +269,9 @@ where
 	F: FnOnce(Yielder<T>) -> U,
 	U: Future<Output = ()>
 {
-	let generator = generator(Yielder { _p: PhantomData });
-	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};

src/tests.rs

@@ -310,3 +310,18 @@ async fn multithreaded() {
 	}
 	join_all(futures).await;
 }
+
+#[tokio::test]
+#[should_panic = "attempted to use async_stream_lite yielder outside of stream context or across threads"]
+async fn test_move_yielder() {
+	let mut slot = None;
+	let s = async_stream(|yielder: Yielder<()>| async {
+		slot.replace(yielder);
+	});
+	pin_mut!(s);
+
+	let _ = s.next().await;
+	drop(s);
+
+	slot.take().unwrap().r#yield(()).await;
+}

src/try.rs

@@ -1,3 +1,4 @@
+use alloc::sync::Arc;
 use core::{
 	future::Future,
 	marker::PhantomData,
@@ -7,18 +8,18 @@ use core::{
 
 use futures_core::stream::{FusedStream, Stream};
 
-use crate::{Yielder, enter};
+use crate::{SharedStore, Yielder, enter};
 
 pin_project_lite::pin_project! {
 	/// A [`Stream`] created from a fallible, asynchronous generator-like function.
 	///
 	/// To create a [`TryAsyncStream`], use the [`try_async_stream`] function. See also [`crate::AsyncStream`].
-	#[derive(Debug)]
 	pub struct TryAsyncStream<T, E, U> {
-		_p: PhantomData<(T, E)>,
+		store: Arc<SharedStore<T>>,
 		done: bool,
 		#[pin]
-		generator: U
+		generator: U,
+		_p: PhantomData<E>
 	}
 }
 
@@ -43,18 +44,17 @@ 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 let Poll::Ready(Err(e)) = res {
 			return Poll::Ready(Some(Err(e)));
-		} else if dst.is_some() {
-			return Poll::Ready(dst.take().map(Ok));
+		} else if me.store.has_value() {
+			return Poll::Ready(me.store.cell.take().map(Ok));
 		}
 
 		if *me.done { Poll::Ready(None) } else { Poll::Pending }
@@ -96,11 +96,13 @@ where
 	F: FnOnce(Yielder<T>) -> U,
 	U: Future<Output = Result<(), E>>
 {
-	let generator = generator(Yielder { _p: PhantomData });
+	let store = Arc::new(SharedStore::default());
+	let generator = generator(Yielder { store: Arc::downgrade(&store) });
 	TryAsyncStream {
-		_p: PhantomData,
+		store,
 		done: false,
-		generator
+		generator,
+		_p: PhantomData
 	}
 }