feat(types,blahd): allow using mock clock for testing

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"

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"] }

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<T: Serialize> 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")

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"

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:

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.
+/// FIXME: Handle multi-threaded runtime.
-use std::time::SystemTime;
-
 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<i64> = 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!(
+        assert!(timestamp_ms < (1 << 48), "invalid timestamp");
-            0 < timestamp_ms && timestamp_ms < (1 << 48),
+        let timestamp_ms = timestamp_ms as i64;
-            "invalid timestamp",
+        let id = timestamp_ms << 16;
-        );
+        LAST_ID.with(|last_id| {
-        Id((timestamp_ms as i64) << 16)
+            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 {

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)]

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<T> {