editoast: core_task: remove lock on valkey connection

editoast/cache/src/client.rs

@@ -36,6 +36,10 @@ pub enum Config {
 }
 
 impl Client {
+    pub fn app_version(&self) -> &str {
+        &self.app_version
+    }
+
     pub fn new(config: Config, app_version: &str) -> Self {
         Self {
             app_version: ArcStr::from(app_version),

editoast/cache/src/connection.rs

@@ -208,8 +208,8 @@ impl Connection {
 
         // Store the compressed values using mset
         if !compressed_items.is_empty() {
-            span!(Level::INFO, "Sending items to Redis")
-                .in_scope(|| self.mset::<_, _, ()>(&compressed_items))
+            span!(Level::INFO, "Sending items to Valkey")
+                .in_scope(async || self.mset::<_, _, ()>(&compressed_items).await)
                 .await?;
         }
         Ok(())
@@ -225,8 +225,8 @@ impl Connection {
 
         // Fetch the values from Redis
         let values = if !keys.is_empty() {
-            span!(Level::INFO, "Fetching values from Redis")
-                .in_scope(|| self.mget::<_, Vec<Option<Vec<u8>>>>(keys))
+            span!(Level::INFO, "Fetching values from Valkey")
+                .in_scope(async || self.mget::<_, Vec<Option<Vec<u8>>>>(keys).await)
                 .await?
         } else {
             // Avoid mget to fail if keys is empty

editoast/core_task/src/envs/pathfinding.rs

@@ -14,7 +14,6 @@ use itertools::Itertools;
 use ordered_float::OrderedFloat;
 use schemas::infra::TrackOffset;
 use schemas::rolling_stock::LoadingGaugeType;
-use tokio::sync::Mutex;
 
 use crate::CoreEnv;
 use crate::Correlated;
@@ -62,14 +61,14 @@ where
     /// Note that the receivers implement [trait futures::stream::Stream].
     pub fn run(
         self,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
         ready_trains_tx: Option<futures::channel::mpsc::UnboundedSender<TrainSet<Train>>>,
     ) -> PathfindingRun<Train> {
         use stream::StreamExt as _;
         let runner = Arc::new(Runner::new(self));
         let run = PathfindingRun::new(&runner);
         let paths = run.paths.clone();
-        tokio::spawn(runner.clone().stream(vkconn).fold(
+        tokio::spawn(runner.clone().stream(vk_client).fold(
             (paths, runner, ready_trains_tx),
             async |(paths, runner, ready_trains_tx),
                    Correlated {
@@ -97,7 +96,7 @@ where
     /// thus transferring ownership and allow easy path mutations.
     pub fn into_stream(
         self,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
     ) -> impl stream::Stream<
         Item = Correlated<
             TrainSet<Train>,
@@ -106,7 +105,7 @@ where
     > {
         use stream::StreamExt as _;
         let runner = Arc::new(Runner::new(self));
-        runner.clone().stream(vkconn).map(
+        runner.clone().stream(vk_client).map(
             move |Correlated {
                       correlation_key: input,
                       data: path,
@@ -296,7 +295,7 @@ where
 
     pub(in crate::envs) fn stream(
         self: Arc<Self>,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
     ) -> impl stream::Stream<
         Item = Correlated<
             PathfindingKey,
@@ -313,7 +312,7 @@ where
             })
             .collect_vec();
 
-        stream::iter(requests).run(vkconn, self.core_env.client.clone())
+        stream::iter(requests).run(vk_client, self.core_env.client.clone())
     }
 }
 
@@ -513,7 +512,7 @@ mod tests {
         let mut pfenv = PathfindingEnv::<usize>::new(CoreEnv::new_mock(mock));
         pfenv.extend([(1, train(1)), (2, train(2))]);
 
-        let vk = cache::Client::new_mock(
+        let vk = Arc::new(cache::Client::new_mock(
             vec![
                 mock_mget(vec![
                     (pfenv.key(1), Some(compress_json(&path(1)))),
@@ -527,13 +526,10 @@ mod tests {
                 ),
             ],
             "",
-        );
+        ));
 
         let (tx, rx) = futures::channel::mpsc::unbounded();
-        let pfrun = pfenv.run(
-            Arc::new(Mutex::new(vk.get_connection().await.unwrap())),
-            Some(tx),
-        );
+        let pfrun = pfenv.run(vk, Some(tx));
 
         // timeout to avoid blocking the CI if something goes wrong
         let trains = tokio::time::timeout(Duration::from_secs(1), async move {
@@ -580,14 +576,11 @@ mod tests {
         let mut pfenv = PathfindingEnv::<usize>::new(CoreEnv::new_mock(mock));
         pfenv.extend([(1, train(TRAIN)), (2, train(TRAIN))]);
 
-        let vk = cache::Client::new(cache::Config::NoCache, "");
+        let vk = Arc::new(cache::Client::new(cache::Config::NoCache, ""));
 
         use futures::StreamExt as _;
         let results = tokio::time::timeout(Duration::from_secs(1), async move {
-            pfenv
-                .into_stream(Arc::new(Mutex::new(vk.get_connection().await.unwrap())))
-                .collect::<Vec<_>>()
-                .await
+            pfenv.into_stream(vk).collect::<Vec<_>>().await
         })
         .await
         .unwrap();

editoast/core_task/src/envs/simulation.rs

@@ -10,7 +10,6 @@ use std::sync::Arc;
 
 use dashmap::DashMap;
 use futures::stream;
-use tokio::sync::Mutex;
 
 use crate::CoreEnv;
 use crate::Correlated;
@@ -95,13 +94,13 @@ where
     /// TODO: `SimulationEnv::run` to be able to retrieve paths as well.
     pub fn into_stream(
         self,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
     ) -> impl stream::Stream<
         Item = Correlated<TrainSet<Train>, Result<SimulationOutput, core_client::Error>>,
     > {
         use stream::StreamExt as _;
         let runner = Arc::new(Runner::new(self));
-        runner.clone().stream(vkconn).map(
+        runner.clone().stream(vk_client).map(
             move |Correlated {
                       correlation_key: input,
                       data: sim_output,
@@ -192,7 +191,7 @@ where
 
     pub(in crate::envs) fn stream(
         self: Arc<Self>,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
     ) -> impl stream::Stream<
         Item = Correlated<SimulationKey, Result<SimulationOutput, core_client::Error>>,
     > {
@@ -204,7 +203,7 @@ where
             Correlated<SimulationKey, Result<SimulationOutput, core_client::Error>>,
         >();
 
-        tokio::spawn(self.pathfinding_runner.clone().stream(vkconn.clone()).fold(
+        tokio::spawn(self.pathfinding_runner.clone().stream(vk_client.clone()).fold(
             // a trick to "move" context into the closure but keeping it AsyncFnMut
             (self.clone(), simulate_tx, return_tx.clone()),
             async |(runner, simulate_tx, return_tx), result| {
@@ -290,7 +289,7 @@ where
         use crate::TaskStreamExt as _;
         tokio::spawn(
             simulate_rx
-                .run(vkconn, self.core_env.client.clone())
+                .run(vk_client, self.core_env.client.clone())
                 .map(move |result|
                     match result {
                         Correlated {
@@ -526,7 +525,7 @@ mod tests {
         let mut simenv = SimulationEnv::<usize>::new(CoreEnv::new_mock(mock));
         simenv.extend([(1, train(1)), (2, train(2))]);
 
-        let vk = cache::Client::new_mock(
+        let vk = Arc::new(cache::Client::new_mock(
             vec![
                 mock_mget(vec![
                     (
@@ -553,11 +552,11 @@ mod tests {
                 ),
             ],
             "",
-        );
+        ));
 
         use futures::StreamExt as _;
         let simulations = simenv
-            .into_stream(Arc::new(Mutex::new(vk.get_connection().await.unwrap())))
+            .into_stream(vk)
             .flat_map(
                 |Correlated {
                      correlation_key,

editoast/core_task/src/lib.rs

@@ -24,7 +24,6 @@ use itertools::Itertools as _;
 use itertools::izip;
 use serde::de::DeserializeOwned;
 use serde::ser::Serialize;
-use tokio::sync::Mutex;
 use tracing::Instrument;
 
 /// Interface required by [SimulationEnv] and friends for [Correlated] keys
@@ -79,14 +78,16 @@ pub trait Task: Sized + Send {
     #[expect(async_fn_in_trait)] // not for public (ie. outside editoast) use, no auto traits bounds to specify on the resulting future
     async fn run(
         self,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
         ctx: Self::Context,
     ) -> Result<Self::Output, Self::Error> {
-        let (key, cache_entry) = {
-            let mut vkconn = vkconn.lock().await;
-            let key = self.key(vkconn.app_version());
-            let entry = vkconn.json_get::<Self::Output, _>(&key).await;
-            (key, entry)
+        let key = self.key(vk_client.app_version());
+        let cache_entry = match vk_client.get_connection().await {
+            Ok(mut vkconn) => vkconn
+                .json_get::<Self::Output, _>(&key)
+                .await
+                .map_err(|e| e.into()),
+            Err(e) => Err(e),
         };
         match cache_entry.unwrap_or_else(|e| {
             tracing::error!(?e, key, "cache read error — computing task output");
@@ -111,12 +112,8 @@ pub trait Task: Sized + Send {
                         tokio::spawn(
                             async move {
                                 use deadpool_redis::redis::AsyncCommands as _;
-                                if let Err(e) = vkconn
-                                    .lock()
-                                    .await
-                                    .set::<_, _, ()>(key.clone(), serialized)
-                                    .await
-                                {
+                                let mut vkconn = vk_client.get_connection().await.unwrap();
+                                if let Err(e) = vkconn.set::<_, _, ()>(&key, serialized).await {
                                     tracing::error!(?e, key, "cache write error");
                                 }
                             }
@@ -170,7 +167,7 @@ where
     /// The order of the results is not the same as the order of inputs.
     fn run(
         self,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
         ctx: T::Context,
     ) -> impl stream::Stream<
         Item = Correlated<CorrelationKey, Result<<T as Task>::Output, <T as Task>::Error>>,
@@ -187,7 +184,7 @@ where
     #[tracing::instrument(skip_all)]
     fn run(
         self,
-        vkconn: Arc<Mutex<cache::Connection>>,
+        vk_client: Arc<cache::Client>,
         ctx: <T as Task>::Context,
     ) -> impl stream::Stream<
         Item = Correlated<CorrelationKey, Result<<T as Task>::Output, <T as Task>::Error>>,
@@ -216,12 +213,13 @@ where
         let (cache_write_tx, mut cache_write_rx) =
             tokio::sync::mpsc::unbounded_channel::<(String, serde_json::Value)>();
         {
-            let vkconn = vkconn.clone();
+            let vk_client = vk_client.clone();
             // 'write_cache' task, writes input key-value pairs to cache, logging errors
             tokio::spawn(
                 async move {
                     while let Some((key, value)) = cache_write_rx.recv().await {
-                        if let Err(e) = vkconn.lock().await.json_set(key.clone(), &value).await {
+                        let mut vkconn = vk_client.get_connection().await.unwrap();
+                        if let Err(e) = vkconn.json_set(key.clone(), &value).await {
                             tracing::error!(?e, key, "task stream: cache write failure")
                         }
                     }
@@ -242,19 +240,17 @@ where
             // and sends a bunch of data to the 'aggregation' task
             tokio::spawn(
                 self.chunks(T::CACHE_READS_BATCH_SIZE)
-                    .zip(stream::repeat(vkconn.clone()))
+                    .zip(stream::repeat(vk_client.clone()))
                     .zip(stream::repeat(cache_read_tx))
-                    .for_each_concurrent(None, async move |((inputs, vkconn), cache_read_tx)| {
-                        let mut vk = vkconn.lock().await;
-
+                    .for_each_concurrent(None, async move |((inputs, vk_client), cache_read_tx)| {
                         // We sort the keys so that unit tests can predictably mock redis requests.
                         // That's because redis-test doesn't find a matching request in the list, but
                         // just pops the first one and asserts.
                         #[cfg(test)]
                         let inputs = inputs
                             .into_iter()
                             .map(|input| {
-                                let key = input.data.key(vk.app_version());
+                                let key = input.data.key(vk_client.app_version());
                                 (input, key)
                             })
                             .sorted_by_key(|(_, key)| key.clone())
@@ -267,14 +263,16 @@ where
                             .unzip::<_, _, Vec<_>, Vec<_>>();
                         let cache_keys = inputs
                             .iter()
-                            .map(|input| input.key(vk.app_version()))
+                            .map(|input| input.key(vk_client.app_version()))
                             .collect_vec();
 
                         // we have to clone because of json_get_bulk's API x Rust 2024 new rules
                         let keys = cache_keys.clone();
 
                         // Fetch from valkey or compute and write to valkey
-                        match vk.json_get_bulk::<T::Output, _>(keys.as_slice()).await {
+
+                        let mut vkconn = vk_client.get_connection().await.unwrap();
+                        match vkconn.json_get_bulk::<T::Output, _>(keys.as_slice()).await {
                             Ok(cached_values) => {
                                 for (value, correlation, key, input) in
                                     izip!(cached_values, correlation_keys, cache_keys, inputs)

editoast/src/views/path/pathfinding.rs

@@ -27,7 +27,6 @@ use schemas::train_schedule::PathItemLocation;
 use schemas::train_schedule::TrainScheduleLike;
 use serde::Deserialize;
 use serde::Serialize;
-use tokio::sync::Mutex;
 use tracing::debug;
 use tracing::info;
 use utoipa::ToSchema;
@@ -463,9 +462,8 @@ pub(in crate::views) async fn single_pathfinding_request(
     });
     pathfinding_env.extend([((), pathfinding_train)]);
 
-    let valkey_conn = Arc::new(Mutex::new(valkey_client.get_connection().await?));
     let result = match pathfinding_env
-        .into_stream(valkey_conn)
+        .into_stream(valkey_client)
         .collect::<Vec<_>>()
         .await
         .as_slice()

editoast/src/views/path/properties.rs

@@ -21,7 +21,6 @@ use geos::geojson::Geometry;
 use serde::Deserialize;
 use serde::Serialize;
 use std::hash::Hash;
-use std::sync::Arc;
 use utoipa::ToSchema;
 
 use crate::AppState;
@@ -107,13 +106,12 @@ pub(in crate::views) async fn post(
     .await?;
 
     use core_task::Task as _;
-    let vkconn = valkey_client.get_connection().await?;
     let path_properties = PathPropertiesRequest {
         track_section_ranges: &path_properties_input.track_section_ranges,
         infra: infra_id,
         expected_version: infra_version,
     }
-    .run(Arc::new(tokio::sync::Mutex::new(vkconn)), core_client)
+    .run(valkey_client, core_client)
     .await?;
 
     Ok(Json(PathProperties::from(path_properties)))

editoast/src/views/timetable/train_schedule.rs

@@ -43,7 +43,6 @@ use schemas::train_schedule::TrainScheduleLike as _;
 use serde::Deserialize;
 use serde::Serialize;
 use thiserror::Error;
-use tokio::sync::Mutex;
 use utoipa::IntoParams;
 use utoipa::ToSchema;
 
@@ -448,7 +447,7 @@ pub(in crate::views) async fn simulation_summary(
     // Populate the simulation environment and simulate
     simulation_env.extend(simulation_trains);
     let simulations = simulation_env
-        .into_stream(Arc::new(Mutex::new(valkey_client.get_connection().await?)))
+        .into_stream(valkey_client)
         .collect::<Vec<_>>()
         .await;