ferroid/futures/
snowflake.rs

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

ferroid/futures/
snowflake.rs
