Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
1use serde::{Deserialize, Serialize};
2use thiserror::Error;
3
4#[derive(Debug, Serialize, Deserialize)]
5struct SnowflakeRequest {
6 count: u32,
7}
8
9#[derive(Error, Debug)]
10pub enum SnowflakeClientError {
11 #[error(transparent)]
12 IdGenerateError(#[from] reqwest::Error),
13
14 #[error("no snowflake id returned from snowflake service")]
15 NoIdGeneratedError,
16}
17
18pub struct SnowflakeClient {
19 snowflake_api_endpoint: String,
20 reqwest_client: reqwest::Client,
21 // buffer: SnowflakeBuffer,
22 // request_size: u32,
23}
24
25impl SnowflakeClient {
26 pub fn new(snowflake_api_endpoint: String) -> Self {
27 Self {
28 snowflake_api_endpoint,
29 reqwest_client: reqwest::Client::new(),
30 // note: to get a buffer working without hindering performance, we'd need to figure
31 // out a way to pass around the buffer without the dreaded Arc<Mutex<>>. because i dont wanna
32 // figure that out, and it'd be definitely overkill for now, i'm just gonna Arc<> and make the
33 // client only perform single id requests.
34 //
35 // buffer: SnowflakeBuffer::new(),
36 // // todo: it would be cool if this was calculated based on a rolling window
37 // // of the number of ids requested in the last 500ms or so . . but that's way
38 // // too overengineered for now lol
39 // request_size: 1,
40 }
41 }
42
43 pub async fn generate_id(&self) -> Result<u64, SnowflakeClientError> {
44 // osprey-snowflake api spec: https://github.com/ayubun/snowflake-id-worker?tab=readme-ov-file#api-spec
45 let mut snowflake_response: Vec<u64> = self
46 .reqwest_client
47 .post(format!("{}/generate", self.snowflake_api_endpoint))
48 .json(&SnowflakeRequest { count: 1 })
49 .send()
50 .await?
51 .json()
52 .await?;
53
54 let id = snowflake_response
55 .pop()
56 .ok_or(SnowflakeClientError::NoIdGeneratedError)?;
57
58 Ok(id)
59 }
60}
61
62// struct SnowflakeBuffer {
63// buffer: Vec<u64>,
64// last_buffer_fill: SystemTime,
65// buffer_ttl: Duration,
66// }
67
68// impl SnowflakeBuffer {
69// pub fn new() -> Self {
70// Self {
71// buffer: Vec::new(),
72// last_buffer_fill: UNIX_EPOCH,
73// // we don't want to use buffered ids older than 1 millisecond, since that
74// // is the precision of snowflake timestamps, and we want our snowflakes
75// // to be as close to the time of generation as possible.
76// buffer_ttl: Duration::from_millis(1),
77// }
78// }
79
80// /// returns the next id in the buffer, or `None` if the buffer
81// /// cannot provide an id (either empty or buffer is too old)
82// pub fn next_id(&mut self) -> Option<u64> {
83// if !self.buffer.is_empty() && self
84// .last_buffer_fill
85// .elapsed()
86// .unwrap_or_else(|_| {
87// UNIX_EPOCH
88// .elapsed()
89// .expect("invariant: literally impossible")
90// })
91// .as_nanos()
92// > self.buffer_ttl.as_nanos()
93// {
94// self.buffer.clear();
95// }
96// self.buffer.pop()
97// }
98
99// /// fills the buffer with the given ids.
100// pub fn fill(&mut self, ids: Vec<u64>) {
101// self.buffer = ids;
102// self.last_buffer_fill = SystemTime::now();
103// }
104// }