diff --git a/Cargo.lock b/Cargo.lock index bcd8ed6..c541c00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -24,7 +24,9 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" name = "fastid" version = "0.0.1" dependencies = [ + "once_cell", "pyo3", + "rs-snowflake", "ulid", ] @@ -229,6 +231,12 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rs-snowflake" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e60ef3b82994702bbe4e134d98aadca4b49ed04440148985678d415c68127666" + [[package]] name = "scopeguard" version = "1.2.0" diff --git a/Cargo.toml b/Cargo.toml index 99661a9..f21555b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "fastid" -version = "0.0.1" +version = "0.0.2" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -12,4 +12,6 @@ crate-type = ["cdylib"] [dependencies] ulid="1.1.0" -pyo3="0.20.0" \ No newline at end of file +pyo3="0.20.0" +rs-snowflake = "0.6.0" +once_cell = "1.19.0" diff --git a/fastid.pyi b/fastid.pyi index 1e65986..d7d36ba 100644 --- a/fastid.pyi +++ b/fastid.pyi @@ -1,2 +1,8 @@ def ulid() -> str: - ... \ No newline at end of file + ... + +def snowflake_str(machine_id: int = 1, node_id: int = 1) -> str: + ... + +def snowflake_int(machine_id: int = 1, node_id: int = 1) -> int: + ... diff --git a/python/test/test_fastid.py b/python/test/test_fastid.py index 1728e08..ac5d3aa 100644 --- a/python/test/test_fastid.py +++ b/python/test/test_fastid.py @@ -6,5 +6,11 @@ def test_ulid(): print() - print(timeit.timeit(lambda: microsecond.new().str, number=10000)) - print(timeit.timeit(fastid.ulid, number=10000)) + print("ulid-py ", timeit.timeit(lambda: microsecond.new().str, number=10000)) + print("fastid.ulid", timeit.timeit(lambda: fastid.ulid(), number=10000)) + + +def test_snowflake(): + print() + print("fastid.snowflake_int", timeit.timeit(lambda: fastid.snowflake_int(), number=10000)) + print("fastid.snowflake_str", timeit.timeit(lambda: fastid.snowflake_str(), number=10000)) diff --git a/src/lib.rs b/src/lib.rs index d50d217..aa67f99 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,14 +1,44 @@ +use std::sync::Mutex; +use once_cell::sync::Lazy; use pyo3::prelude::*; -use ulid::{Ulid}; +use snowflake::SnowflakeIdGenerator; +use ulid::Ulid; +static SNOWFLAKE_GENERATOR: Lazy> = Lazy::new(|| { + Mutex::new(SnowflakeIdGenerator::new(1, 1)) +}); #[pyfunction] -#[pyo3(name="ulid")] +#[pyo3(name = "ulid")] fn get_ulid() -> PyResult { Ok(Ulid::new().to_string()) } + +fn get_snowflake(machine_id: Option, node_id: Option) -> i64 { + let mut sg = SNOWFLAKE_GENERATOR.lock().unwrap(); + sg.node_id = node_id.unwrap_or(1); + sg.machine_id = machine_id.unwrap_or(1); + sg.real_time_generate() +} +#[pyfunction] +#[pyo3(name = "snowflake_str")] +fn get_snowflake_str(machine_id: Option, node_id: Option) -> PyResult { + Ok(format!("{:X}", get_snowflake(machine_id, node_id))) +} + +#[pyfunction] +#[pyo3(name = "snowflake_int")] +fn get_snowflake_int(machine_id: Option, node_id: Option) -> PyResult { + Ok(get_snowflake(machine_id, node_id)) +} + + #[pymodule] fn fastid(_py: Python, m: &PyModule) -> PyResult<()> { m.add_function(wrap_pyfunction!(get_ulid, m)?)?; + + m.add_function(wrap_pyfunction!(get_snowflake_int, m)?)?; + m.add_function(wrap_pyfunction!(get_snowflake_str, m)?)?; + Ok(()) }