diff --git a/router/client/src/sharded_client.rs b/router/client/src/sharded_client.rs index 7f0ec6f2..31f76312 100644 --- a/router/client/src/sharded_client.rs +++ b/router/client/src/sharded_client.rs @@ -2,7 +2,6 @@ use crate::Result; use crate::{Batch, Client, Generation}; use futures::future::join_all; -use futures::future::select_all; use tonic::transport::Uri; use tracing::instrument; @@ -60,9 +59,8 @@ impl ShardedClient { .iter_mut() .map(|client| Box::pin(client.prefill(batch.clone()))) .collect(); - // As soon as we receive one response, we can return as all shards will return the same - let (result, _, _) = select_all(futures).await; - result + // all shards return the same message + join_all(futures).await.pop().unwrap() } /// Generate one token for each request in the given cached batches @@ -79,8 +77,7 @@ impl ShardedClient { .iter_mut() .map(|client| Box::pin(client.decode(batches.clone()))) .collect(); - // As soon as we receive one response, we can return as all shards will return the same - let (result, _, _) = select_all(futures).await; - result + // all shards return the same message + join_all(futures).await.pop().unwrap() } }