Skip to content

Commit 2549bcd

Browse files
Fix/redis batch import (#42)
* fix: import model keys from redis cache in batches * fix: tests
1 parent f7ea329 commit 2549bcd

File tree

6 files changed

+36
-10
lines changed

6 files changed

+36
-10
lines changed

crates/routing-engine/src/indexer.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,10 @@ mod tests {
352352
todo!()
353353
}
354354

355-
async fn get_all_key_values(&self) -> Result<HashMap<String, String>, RedisClientError> {
355+
async fn get_all_key_values(
356+
&self,
357+
_: Option<usize>,
358+
) -> Result<HashMap<String, String>, RedisClientError> {
356359
todo!()
357360
}
358361
}

crates/routing-engine/src/routing_engine.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ use crate::{
1919
use crate::token_price::TokenPriceProvider;
2020
use crate::token_price::utils::{Errors, get_token_price};
2121

22+
const FETCH_REDIS_KEYS_BATCH_SIZE: usize = 50;
23+
2224
/// (from_chain, to_chain, from_token, to_token)
2325
#[derive(Debug)]
2426
struct PathQuery(u32, u32, String, String);
@@ -87,7 +89,7 @@ impl<PriceProvider: TokenPriceProvider> RoutingEngine<PriceProvider> {
8789

8890
/// Refresh the cache from Redis
8991
pub async fn refresh_cache(&self) {
90-
match self.redis_client.get_all_key_values().await {
92+
match self.redis_client.get_all_key_values(Some(FETCH_REDIS_KEYS_BATCH_SIZE)).await {
9193
Ok(kv_pairs) => {
9294
info!("Refreshing cache from Redis.");
9395
let mut cache = self.cache.write().await;
@@ -418,7 +420,10 @@ mod tests {
418420
unimplemented!()
419421
}
420422

421-
async fn get_all_key_values(&self) -> Result<HashMap<String, String>, RedisClientError> {
423+
async fn get_all_key_values(
424+
&self,
425+
_: Option<usize>,
426+
) -> Result<HashMap<String, String>, RedisClientError> {
422427
unimplemented!()
423428
}
424429
}

crates/routing-engine/src/settlement_engine.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -488,7 +488,10 @@ mod tests {
488488
unimplemented!()
489489
}
490490

491-
async fn get_all_key_values(&self) -> Result<HashMap<String, String>, RedisClientError> {
491+
async fn get_all_key_values(
492+
&self,
493+
_: Option<usize>,
494+
) -> Result<HashMap<String, String>, RedisClientError> {
492495
unimplemented!()
493496
}
494497
}

crates/routing-engine/src/token_price/coingecko.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,10 +144,10 @@ mod tests {
144144

145145
use async_trait::async_trait;
146146
use derive_more::Display;
147+
use serial_test::serial;
147148
use thiserror::Error;
148149

149150
use config::{Config, get_sample_config};
150-
use serial_test::serial;
151151
use storage::{KeyValueStore, RedisClientError};
152152

153153
use crate::CoingeckoClient;
@@ -192,7 +192,10 @@ mod tests {
192192
unimplemented!()
193193
}
194194

195-
async fn get_all_key_values(&self) -> Result<HashMap<String, String>, RedisClientError> {
195+
async fn get_all_key_values(
196+
&self,
197+
_: Option<usize>,
198+
) -> Result<HashMap<String, String>, RedisClientError> {
196199
unimplemented!()
197200
}
198201
}

crates/storage/src/lib.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,10 @@ pub trait KeyValueStore: Debug + Send + Sync {
2727

2828
async fn get_all_keys(&self) -> Result<Vec<String>, RedisClientError>;
2929

30-
async fn get_all_key_values(&self) -> Result<HashMap<String, String>, RedisClientError>;
30+
async fn get_all_key_values(
31+
&self,
32+
batch_size: Option<usize>,
33+
) -> Result<HashMap<String, String>, RedisClientError>;
3134
}
3235

3336
#[async_trait]

crates/storage/src/redis_client.rs

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,10 +58,19 @@ impl KeyValueStore for RedisClient {
5858
Ok(keys)
5959
}
6060

61-
async fn get_all_key_values(&self) -> Result<HashMap<String, String>, RedisClientError> {
61+
async fn get_all_key_values(
62+
&self,
63+
batch_size: Option<usize>,
64+
) -> Result<HashMap<String, String>, RedisClientError> {
6265
info!("Fetching all key-value pairs");
6366
let keys = self.get_all_keys().await?;
64-
let values: Vec<String> = self.connection.clone().mget(&keys).await?;
67+
68+
let batch_size = batch_size.unwrap_or(keys.len());
69+
let mut values = Vec::new();
70+
for batch in keys.chunks(batch_size) {
71+
values.extend(self.connection.clone().mget::<'_, _, Vec<_>>(batch).await?.into_iter());
72+
}
73+
6574
let kv_pairs = keys.into_iter().zip(values.into_iter()).collect();
6675
Ok(kv_pairs)
6776
}
@@ -167,7 +176,7 @@ mod tests {
167176
.unwrap();
168177

169178
// Fetch all key-values
170-
let key_values = client.get_all_key_values().await.unwrap();
179+
let key_values = client.get_all_key_values(None).await.unwrap();
171180

172181
assert_eq!(key_values.get("key1").unwrap(), "value1");
173182
assert_eq!(key_values.get("key2").unwrap(), "value2");

0 commit comments

Comments
 (0)