Fix the bug that dataloader will only return one value when loading the same key.
@@ -187,9 +187,6 @@ impl<T> DataLoader<T> {
let prev_count = typed_requests.keys.len();
let keys = keys.collect::<HashSet<_>>();
typed_requests.keys.extend(keys.clone());
- if typed_requests.keys.len() == prev_count {
- return Ok(Default::default());
- }
let (tx, rx) = oneshot::channel();
typed_requests.pending.push((keys, tx));
if typed_requests.keys.len() >= self.max_batch_size {
@@ -226,8 +223,6 @@ mod tests {
use super::*;
use std::sync::Arc;
- #[async_std::test]
- async fn test_dataloader() {
struct MyLoader;
#[async_trait::async_trait]
@@ -252,6 +247,8 @@ mod tests {
}
}
+ #[async_std::test]
+ async fn test_dataloader() {
let loader = Arc::new(DataLoader::new(MyLoader).max_batch_size(10));
assert_eq!(
futures_util::future::try_join_all((0..100i32).map({
@@ -279,4 +276,25 @@ mod tests {
(0..100).map(Option::Some).collect::<Vec<_>>()
);
}
+
+ #[async_std::test]
+ async fn test_duplicate_keys() {
+ let loader = Arc::new(DataLoader::new(MyLoader).max_batch_size(10));
+ assert_eq!(
+ futures_util::future::try_join_all([1, 3, 5, 1, 7, 8, 3, 7].iter().copied().map({
+ let loader = loader.clone();
+ move |n| {
+ let loader = loader.clone();
+ async move { loader.load_one(n).await }
+ }
+ }))
+ .await
+ .unwrap(),
+ [1, 3, 5, 1, 7, 8, 3, 7]
+ .iter()
+ .copied()
+ .map(Option::Some)
+ .collect::<Vec<_>>()
+ );
+ }
}
