Merge pull request #4 from matrix-org/quenting/ulid

Author: sandhose

Cargo.toml

@@ -21,7 +21,7 @@ tokio = { version = "1.21.2", features = ["sync", "time"] }
 tower = { version = "0.4.13", features = ["util"] }
 tower-http = { version = "0.3.4", features = ["cors", "limit", "set-header"] }
 tracing = "0.1.37"
-uuid = { version = "1.1.2", features = ["v4", "fast-rng", "serde"] }
+ulid = { version = "1.0.0", features = ["serde"] }
 
 [dev-dependencies]
 hyper = "0.14.20"

src/lib.rs

@@ -42,13 +42,13 @@ use headers::{
 };
 use mime::Mime;
 use sha2::Digest;
-use tokio::sync::RwLock;
+use tokio::sync::{Mutex, RwLock};
 use tower_http::{
     cors::{Any, CorsLayer},
     limit::RequestBodyLimitLayer,
     set_header::SetResponseHeaderLayer,
 };
-use uuid::Uuid;
+use ulid::Ulid;
 
 struct Session {
     hash: [u8; 32],
@@ -114,23 +114,34 @@ impl Session {
 #[derive(Clone, Default)]
 struct Sessions {
     // TODO: is that global lock alright?
-    inner: Arc<RwLock<HashMap<Uuid, Session>>>,
+    inner: Arc<RwLock<HashMap<Ulid, Session>>>,
     ttl: Duration,
+    generator: Arc<Mutex<ulid::Generator>>,
 }
 
 impl Sessions {
-    async fn insert(self, id: Uuid, session: Session, ttl: Duration) {
+    async fn insert(self, id: Ulid, session: Session, ttl: Duration) {
         self.inner.write().await.insert(id, session);
         // TODO: cancel this task when an item gets deleted
         tokio::task::spawn(async move {
             tokio::time::sleep(ttl).await;
             self.inner.write().await.remove(&id);
         });
     }
+
+    async fn generate_id(&self) -> Ulid {
+        self.generator
+            .lock()
+            .await
+            .generate()
+            // This would panic the thread if too many IDs (more than 2^40) are generated on the same
+            // millisecond, which is very unlikely
+            .expect("Failed to generate random ID")
+    }
 }
 
 impl Deref for Sessions {
-    type Target = RwLock<HashMap<Uuid, Session>>;
+    type Target = RwLock<HashMap<Ulid, Session>>;
 
     fn deref(&self) -> &Self::Target {
         &self.inner
@@ -159,8 +170,9 @@ async fn new_session(
     payload: Bytes,
 ) -> impl IntoResponse {
     let ttl = sessions.ttl;
-    // TODO: should we use something else? Check for colisions?
-    let id = Uuid::new_v4();
+
+    let id = sessions.generate_id().await;
+
     let content_type =
         content_type.map_or(mime::APPLICATION_OCTET_STREAM, |TypedHeader(c)| c.into());
     let session = Session::new(payload, content_type, ttl);
@@ -172,7 +184,7 @@ async fn new_session(
     (StatusCode::CREATED, headers, additional_headers)
 }
 
-async fn delete_session(State(sessions): State<Sessions>, Path(id): Path<Uuid>) -> StatusCode {
+async fn delete_session(State(sessions): State<Sessions>, Path(id): Path<Ulid>) -> StatusCode {
     if sessions.write().await.remove(&id).is_some() {
         StatusCode::NO_CONTENT
     } else {
@@ -182,7 +194,7 @@ async fn delete_session(State(sessions): State<Sessions>, Path(id): Path<Uuid>)
 
 async fn update_session(
     State(sessions): State<Sessions>,
-    Path(id): Path<Uuid>,
+    Path(id): Path<Ulid>,
     content_type: Option<TypedHeader<ContentType>>,
     if_match: Option<TypedHeader<IfMatch>>,
     payload: Bytes,
@@ -206,7 +218,7 @@ async fn update_session(
 
 async fn get_session(
     State(sessions): State<Sessions>,
-    Path(id): Path<Uuid>,
+    Path(id): Path<Ulid>,
     if_none_match: Option<TypedHeader<IfNoneMatch>>,
 ) -> Response {
     let sessions = sessions.read().await;
@@ -241,6 +253,7 @@ where
     let sessions = Sessions {
         inner: Arc::default(),
         ttl,
+        generator: Arc::default(),
     };
 
     let state = AppState::new(sessions);
@@ -383,6 +396,40 @@ mod tests {
         assert_eq!(response.status(), StatusCode::NOT_FOUND);
     }
 
+    #[tokio::test]
+    async fn test_monotonically_increasing() {
+        let ttl = Duration::from_secs(60);
+        let app = router("/", ttl, 4096);
+
+        // Prepare a thousand requests
+        let mut requests = Vec::with_capacity(1000);
+        for _ in 0..requests.capacity() {
+            requests.push(
+                app.clone()
+                    .oneshot(Request::post("/").body(String::new()).unwrap()),
+            );
+        }
+
+        // Run them all in order
+        let mut responses = Vec::with_capacity(requests.len());
+        for fut in requests {
+            responses.push(fut.await);
+        }
+
+        // Get the location out of them
+        let ids: Vec<_> = responses
+            .iter()
+            .map(|res| {
+                let res = res.as_ref().unwrap();
+                assert_eq!(res.status(), StatusCode::CREATED);
+                res.headers().get(LOCATION).unwrap().to_str().unwrap()
+            })
+            .collect();
+
+        // Check that all the IDs are monotonically increasing
+        assert!(ids.windows(2).all(|loc| loc[0] < loc[1]));
+    }
+
     #[tokio::test]
     async fn test_post_max_bytes() {
         let ttl = Duration::from_secs(60);