ferroid/futures/
snowflake.rs

1use super::SleepProvider;
2use crate::{IdGenStatus, Result, SnowflakeGenerator, SnowflakeId, TimeSource, ToU64};
3use core::{
4    fmt,
5    future::Future,
6    marker::PhantomData,
7    pin::Pin,
8    task::{Context, Poll},
9    time::Duration,
10};
11use pin_project_lite::pin_project;
12
13/// Extension trait for asynchronously generating Snowflake IDs.
14///
15/// This trait enables `SnowflakeGenerator` types to yield IDs in a
16/// `Future`-based context by awaiting until the generator is ready to produce a
17/// new ID.
18///
19/// The default implementation uses [`SnowflakeGeneratorFuture`] and a specified
20/// [`SleepProvider`] to yield when the generator is not yet ready.
21pub trait SnowflakeGeneratorAsyncExt<ID, T>
22where
23    ID: SnowflakeId,
24    T: TimeSource<ID::Ty>,
25{
26    type Err: fmt::Debug;
27
28    /// Returns a future that resolves to the next available Snowflake ID.
29    ///
30    /// If the generator is not ready to issue a new ID immediately, the future
31    /// will sleep for the amount of time indicated by the generator and retry.
32    ///
33    /// # Errors
34    ///
35    /// This future may return an error if the generator encounters one.
36    fn try_next_id_async<S>(&self) -> impl Future<Output = Result<ID, Self::Err>>
37    where
38        S: SleepProvider;
39}
40
41impl<G, ID, T> SnowflakeGeneratorAsyncExt<ID, T> for G
42where
43    G: SnowflakeGenerator<ID, T>,
44    ID: SnowflakeId,
45    T: TimeSource<ID::Ty>,
46{
47    type Err = G::Err;
48
49    #[allow(clippy::future_not_send)]
50    fn try_next_id_async<'a, S>(&'a self) -> impl Future<Output = Result<ID, Self::Err>>
51    where
52        S: SleepProvider,
53    {
54        SnowflakeGeneratorFuture::<'a, G, ID, T, S>::new(self)
55    }
56}
57
58pin_project! {
59    /// A future that polls a [`SnowflakeGenerator`] until it is ready to
60    /// produce an ID.
61    ///
62    /// This future handles `Pending` responses by sleeping for a recommended
63    /// amount of time before polling the generator again.
64    #[must_use = "futures do nothing unless you `.await` or poll them"]
65    pub struct SnowflakeGeneratorFuture<'a, G, ID, T, S>
66    where
67        G: SnowflakeGenerator<ID, T>,
68        ID: SnowflakeId,
69        T: TimeSource<ID::Ty>,
70        S: SleepProvider,
71    {
72        generator: &'a G,
73        #[pin]
74        sleep: Option<S::Sleep>,
75        _idt: PhantomData<(ID, T)>
76    }
77}
78
79impl<'a, G, ID, T, S> SnowflakeGeneratorFuture<'a, G, ID, T, S>
80where
81    G: SnowflakeGenerator<ID, T>,
82    ID: SnowflakeId,
83    T: TimeSource<ID::Ty>,
84    S: SleepProvider,
85{
86    /// Constructs a new [`SnowflakeGeneratorFuture`] from a given generator.
87    ///
88    /// This does not immediately begin polling the generator; instead, it will
89    /// attempt to produce an ID when `.poll()` is called.
90    pub const fn new(generator: &'a G) -> Self {
91        Self {
92            generator,
93            sleep: None,
94            _idt: PhantomData,
95        }
96    }
97}
98impl<G, ID, T, S> Future for SnowflakeGeneratorFuture<'_, G, ID, T, S>
99where
100    G: SnowflakeGenerator<ID, T>,
101    ID: SnowflakeId,
102    T: TimeSource<ID::Ty>,
103    S: SleepProvider,
104{
105    type Output = Result<ID, G::Err>;
106
107    /// Polls the generator for a new ID.
108    ///
109    /// If the generator is not ready, this will register the task waker and
110    /// sleep for the time recommended by the generator before polling again.
111    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
112        let mut this = self.project();
113
114        if let Some(sleep) = this.sleep.as_mut().as_pin_mut() {
115            match sleep.poll(cx) {
116                Poll::Pending => {
117                    return Poll::Pending;
118                }
119                Poll::Ready(()) => {
120                    this.sleep.set(None);
121                }
122            }
123        }
124        match this.generator.try_next_id()? {
125            IdGenStatus::Ready { id } => Poll::Ready(Ok(id)),
126            IdGenStatus::Pending { yield_for } => {
127                let sleep_fut = S::sleep_for(Duration::from_millis(yield_for.to_u64()));
128                this.sleep.as_mut().set(Some(sleep_fut));
129                cx.waker().wake_by_ref();
130                Poll::Pending
131            }
132        }
133    }
134}
ferroid/futures/snowflake.rs

ferroid/futures/
snowflake.rs
