diff --git a/Cargo.lock b/Cargo.lock index ab2c6c3..968abc8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -276,6 +276,7 @@ dependencies = [ "expect-test", "hex", "html-escape", + "mock_instant", "rand", "rusqlite", "serde", @@ -322,6 +323,7 @@ dependencies = [ "html-escape", "http-body-util", "humantime", + "mock_instant", "nix", "parking_lot", "paste", @@ -1320,6 +1322,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "mock_instant" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdcebb6db83796481097dedc7747809243cc81d9ed83e6a938b76d4ea0b249cf" + [[package]] name = "native-tls" version = "0.2.12" diff --git a/blah-types/Cargo.toml b/blah-types/Cargo.toml index c2c551d..cff810a 100644 --- a/blah-types/Cargo.toml +++ b/blah-types/Cargo.toml @@ -3,6 +3,10 @@ name = "blah-types" version = "0.0.0" edition = "2021" +[features] +default = [] +unsafe_use_mock_instant_for_testing = ["dep:mock_instant"] + [[bench]] name = "crypto_ops" harness = false @@ -13,6 +17,7 @@ bitflags_serde_shim = "0.2" ed25519-dalek = "2" hex = { version = "0.4", features = ["serde"] } html-escape = "0.2" +mock_instant = { version = "0.5", optional = true } rand = "0.8" rusqlite = { version = "0.32", optional = true } serde = { version = "1", features = ["derive"] } diff --git a/blah-types/src/crypto.rs b/blah-types/src/crypto.rs index 60934bb..2e97ce4 100644 --- a/blah-types/src/crypto.rs +++ b/blah-types/src/crypto.rs @@ -2,7 +2,6 @@ use std::fmt; use std::str::FromStr; -use std::time::SystemTime; use ed25519_dalek::{ Signature, SignatureError, Signer, SigningKey, VerifyingKey, PUBLIC_KEY_LENGTH, @@ -104,6 +103,12 @@ impl SignExt for T { } pub fn get_timestamp() -> u64 { + #[cfg(not(feature = "unsafe_use_mock_instant_for_testing"))] + use std::time::SystemTime; + + #[cfg(feature = "unsafe_use_mock_instant_for_testing")] + use mock_instant::thread_local::SystemTime; + SystemTime::now() .duration_since(SystemTime::UNIX_EPOCH) .expect("after UNIX epoch") diff --git a/blahd/Cargo.toml b/blahd/Cargo.toml index 85d0725..a1645a0 100644 --- a/blahd/Cargo.toml +++ b/blahd/Cargo.toml @@ -3,6 +3,10 @@ name = "blahd" version = "0.0.0" edition = "2021" +[features] +default = [] +unsafe_use_mock_instant_for_testing = ["dep:mock_instant", "blah-types/unsafe_use_mock_instant_for_testing"] + [dependencies] anyhow = "1" axum = { version = "0.7", features = ["ws"] } @@ -15,6 +19,7 @@ hex = { version = "0.4", features = ["serde"] } html-escape = "0.2" http-body-util = "0.1" humantime = "2" +mock_instant = { version = "0.5", optional = true } parking_lot = "0.12" # Maybe no better performance, just that we hate poisoning. ¯\_(ツ)_/¯ paste = "1.0.15" rand = "0.8" diff --git a/blahd/src/feed.rs b/blahd/src/feed.rs index 706d612..76678e0 100644 --- a/blahd/src/feed.rs +++ b/blahd/src/feed.rs @@ -1,7 +1,7 @@ //! Room feed generation. use std::fmt; use std::num::NonZero; -use std::time::{Duration, SystemTime}; +use std::time::Duration; use axum::http::header; use axum::response::{IntoResponse, Response}; @@ -43,7 +43,8 @@ pub trait FeedType { } fn timestamp_to_rfc3339(timestamp: u64) -> impl fmt::Display { - humantime::format_rfc3339(SystemTime::UNIX_EPOCH + Duration::from_secs(timestamp)) + // This only for formatting, thus always use the non-mock `SystemTime`. + humantime::format_rfc3339(std::time::SystemTime::UNIX_EPOCH + Duration::from_secs(timestamp)) } /// See: diff --git a/blahd/src/id.rs b/blahd/src/id.rs index 01ed17d..0c3bb20 100644 --- a/blahd/src/id.rs +++ b/blahd/src/id.rs @@ -1,10 +1,12 @@ +use std::cell::Cell; + /// Id generation. /// Ref: https://en.wikipedia.org/wiki/Snowflake_ID -/// FIXME: Currently we assume no more than one request in a single millisecond. -use std::time::SystemTime; - +/// FIXME: Handle multi-threaded runtime. use blah_types::Id; +use crate::utils::SystemTime; + pub fn timestamp_of_id(id: Id) -> u64 { (id.0 as u64 >> 16) / 1000 } @@ -16,17 +18,32 @@ pub trait IdExt { fn is_peer_chat(&self) -> bool; } +thread_local! { + static LAST_ID: Cell = const { Cell::new(0) }; +} + impl IdExt for Id { fn gen() -> Self { let timestamp = SystemTime::now() .duration_since(SystemTime::UNIX_EPOCH) .expect("after UNIX epoch"); let timestamp_ms = timestamp.as_millis(); - assert!( - 0 < timestamp_ms && timestamp_ms < (1 << 48), - "invalid timestamp", - ); - Id((timestamp_ms as i64) << 16) + assert!(timestamp_ms < (1 << 48), "invalid timestamp"); + let timestamp_ms = timestamp_ms as i64; + let id = timestamp_ms << 16; + LAST_ID.with(|last_id| { + let prev = last_id.get(); + if prev >> 16 != timestamp_ms { + // If not in the same millisecond, use the new timestamp as id. + last_id.set(id); + Id(id) + } else { + // Otherwise, try to increse the trailing counter. + assert!(prev < (1 << 16), "id counter overflow"); + last_id.set(prev + 1); + Id(prev + 1) + } + }) } fn gen_peer_chat_rid() -> Self { diff --git a/blahd/src/register.rs b/blahd/src/register.rs index 4cf04b9..5aba54d 100644 --- a/blahd/src/register.rs +++ b/blahd/src/register.rs @@ -1,5 +1,5 @@ use std::num::NonZero; -use std::time::{Duration, Instant}; +use std::time::Duration; use anyhow::{anyhow, ensure}; use axum::http::{HeaderMap, HeaderName, StatusCode}; @@ -15,6 +15,7 @@ use serde::Deserialize; use sha2::{Digest, Sha256}; use crate::database::TransactionOps; +use crate::utils::Instant; use crate::{ApiError, AppState, SERVER_AND_VERSION}; #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] diff --git a/blahd/src/utils.rs b/blahd/src/utils.rs index f3a1c0a..caf8ec2 100644 --- a/blahd/src/utils.rs +++ b/blahd/src/utils.rs @@ -1,6 +1,12 @@ use std::collections::{HashSet, VecDeque}; use std::hash::Hash; -use std::time::{Duration, Instant}; +use std::time::Duration; + +#[cfg(not(feature = "unsafe_use_mock_instant_for_testing"))] +pub use std::time::{Instant, SystemTime}; + +#[cfg(feature = "unsafe_use_mock_instant_for_testing")] +pub use mock_instant::thread_local::{Instant, SystemTime}; #[derive(Debug)] pub struct ExpiringSet {