diff --git a/editoast/cache/src/client.rs b/editoast/cache/src/client.rs index 4e608e9475e..583bddea27e 100644 --- a/editoast/cache/src/client.rs +++ b/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), diff --git a/editoast/cache/src/connection.rs b/editoast/cache/src/connection.rs index 57648bf8fc6..dfbc13ccf6c 100644 --- a/editoast/cache/src/connection.rs +++ b/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>>>(keys)) + span!(Level::INFO, "Fetching values from Valkey") + .in_scope(async || self.mget::<_, Vec>>>(keys).await) .await? } else { // Avoid mget to fail if keys is empty diff --git a/editoast/core_task/src/envs/pathfinding.rs b/editoast/core_task/src/envs/pathfinding.rs index 2b36cf4ef99..0582afb032b 100644 --- a/editoast/core_task/src/envs/pathfinding.rs +++ b/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>, + vk_client: Arc, ready_trains_tx: Option>>, ) -> PathfindingRun { 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>, + vk_client: Arc, ) -> impl stream::Stream< Item = Correlated< TrainSet, @@ -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, - vkconn: Arc>, + vk_client: Arc, ) -> 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::::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::::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::>() - .await + pfenv.into_stream(vk).collect::>().await }) .await .unwrap(); diff --git a/editoast/core_task/src/envs/simulation.rs b/editoast/core_task/src/envs/simulation.rs index 3fb406f6a56..4a80aa96039 100644 --- a/editoast/core_task/src/envs/simulation.rs +++ b/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>, + vk_client: Arc, ) -> impl stream::Stream< Item = Correlated, Result>, > { 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, - vkconn: Arc>, + vk_client: Arc, ) -> impl stream::Stream< Item = Correlated>, > { @@ -204,7 +203,7 @@ where Correlated>, >(); - 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::::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, diff --git a/editoast/core_task/src/lib.rs b/editoast/core_task/src/lib.rs index 18e92538152..eb19ebe2449 100644 --- a/editoast/core_task/src/lib.rs +++ b/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>, + vk_client: Arc, ctx: Self::Context, ) -> Result { - let (key, cache_entry) = { - let mut vkconn = vkconn.lock().await; - let key = self.key(vkconn.app_version()); - let entry = vkconn.json_get::(&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::(&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>, + vk_client: Arc, ctx: T::Context, ) -> impl stream::Stream< Item = Correlated::Output, ::Error>>, @@ -187,7 +184,7 @@ where #[tracing::instrument(skip_all)] fn run( self, - vkconn: Arc>, + vk_client: Arc, ctx: ::Context, ) -> impl stream::Stream< Item = Correlated::Output, ::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,11 +240,9 @@ 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. @@ -254,7 +250,7 @@ where 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::(keys.as_slice()).await { + + let mut vkconn = vk_client.get_connection().await.unwrap(); + match vkconn.json_get_bulk::(keys.as_slice()).await { Ok(cached_values) => { for (value, correlation, key, input) in izip!(cached_values, correlation_keys, cache_keys, inputs) diff --git a/editoast/src/views/path/pathfinding.rs b/editoast/src/views/path/pathfinding.rs index c15566bb076..87e90622380 100644 --- a/editoast/src/views/path/pathfinding.rs +++ b/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::>() .await .as_slice() diff --git a/editoast/src/views/path/properties.rs b/editoast/src/views/path/properties.rs index f4cd450d5c4..b5f178de9f6 100644 --- a/editoast/src/views/path/properties.rs +++ b/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))) diff --git a/editoast/src/views/timetable/train_schedule.rs b/editoast/src/views/timetable/train_schedule.rs index a9ebc7b8a8d..7cfc9b5b50b 100644 --- a/editoast/src/views/timetable/train_schedule.rs +++ b/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::>() .await;