editoast: bulk get cached projection

OpenRailAssociation · Jul 5, 2024 · afb4e37 · afb4e37
1 parent c3d986c
commit afb4e37
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 7 deletions.
diff --git a/editoast/src/redis_utils.rs b/editoast/src/redis_utils.rs
@@ -77,6 +77,27 @@ impl RedisConnection {
         }
     }
 
+    /// Get a list of deserializable value from redis
+    #[tracing::instrument(name = "cache:get_bulk", skip(self), err)]
+    pub async fn json_get_bulk<T: DeserializeOwned, K: Debug + ToRedisArgs + Send + Sync>(
+        &mut self,
+        keys: K,
+    ) -> Result<Vec<Option<T>>> {
+        let values: Vec<Option<String>> = self.mget(keys).await?;
+        values
+            .into_iter()
+            .map(|value| match value {
+                Some(v) => match serde_json::from_str::<T>(&v) {
+                    Ok(value) => Ok(Some(value)),
+                    Err(_) => {
+                        Err(RedisError::from((ErrorKind::TypeError, "Expected valid json")).into())
+                    }
+                },
+                None => Ok(None),
+            })
+            .collect()
+    }
+
     /// Get a deserializable value from redis with expiry time
     #[tracing::instrument(name = "cache:get_with_expiration", skip(self), err)]
     pub async fn json_get_ex<T: DeserializeOwned, K: Debug + ToRedisArgs + Send + Sync>(

diff --git a/editoast/src/views/v2/train_schedule/projection.rs b/editoast/src/views/v2/train_schedule/projection.rs
@@ -194,8 +194,8 @@ async fn project_path(
     .await?;
 
     // 1. Retrieve cached projection
-    let mut hit_cache: HashMap<i64, CachedProjectPathTrainResult> = HashMap::new();
-    let mut miss_cache = HashMap::new();
+    let mut trains_hash_values = vec![];
+    let mut trains_details = HashMap::new();
 
     for (train, sim) in trains.iter().zip(simulations) {
         let pathfinding_result = pathfinding_from_train(
@@ -245,15 +245,24 @@ async fn project_path(
             &path_routes,
             &path_blocks,
         );
-        let projection: Option<CachedProjectPathTrainResult> = redis_conn
-            .json_get_ex(&hash, CACHE_PROJECTION_EXPIRATION)
-            .await?;
+        trains_hash_values.push(hash);
+        trains_details.insert(train.id, train_details);
+    }
+    let cached_projections: Vec<Option<CachedProjectPathTrainResult>> =
+        redis_conn.json_get_bulk(&trains_hash_values).await?;
+
+    let mut hit_cache: HashMap<i64, CachedProjectPathTrainResult> = HashMap::new();
+    let mut miss_cache = HashMap::new();
+    for ((train_id, train_details), projection) in
+        trains_details.into_iter().zip(cached_projections)
+    {
         if let Some(cached) = projection {
-            hit_cache.insert(train.id, cached);
+            hit_cache.insert(train_id, cached);
         } else {
-            miss_cache.insert(train.id, train_details);
+            miss_cache.insert(train_id, train_details.clone());
         }
     }
+
     info!(
         nb_hit = hit_cache.len(),
         nb_miss = miss_cache.len(),