Fix quic staked chunking (#27402)
This commit is contained in:
parent
d1522fc790
commit
b7b03cbb05
|
@ -320,7 +320,7 @@ impl ConnectionCache {
|
|||
}
|
||||
}
|
||||
|
||||
fn compute_max_parallel_chunks(&self) -> usize {
|
||||
fn compute_max_parallel_streams(&self) -> usize {
|
||||
let (client_type, stake, total_stake) =
|
||||
self.maybe_client_pubkey
|
||||
.map_or((ConnectionPeerType::Unstaked, 0, 0), |pubkey| {
|
||||
|
@ -370,7 +370,7 @@ impl ConnectionCache {
|
|||
BaseTpuConnection::Quic(Arc::new(QuicClient::new(
|
||||
endpoint.as_ref().unwrap().clone(),
|
||||
*addr,
|
||||
self.compute_max_parallel_chunks(),
|
||||
self.compute_max_parallel_streams(),
|
||||
)))
|
||||
};
|
||||
|
||||
|
@ -730,7 +730,7 @@ mod tests {
|
|||
solana_logger::setup();
|
||||
let mut connection_cache = ConnectionCache::default();
|
||||
assert_eq!(
|
||||
connection_cache.compute_max_parallel_chunks(),
|
||||
connection_cache.compute_max_parallel_streams(),
|
||||
QUIC_MAX_UNSTAKED_CONCURRENT_STREAMS
|
||||
);
|
||||
|
||||
|
@ -738,13 +738,13 @@ mod tests {
|
|||
let pubkey = Pubkey::new_unique();
|
||||
connection_cache.set_staked_nodes(&staked_nodes, &pubkey);
|
||||
assert_eq!(
|
||||
connection_cache.compute_max_parallel_chunks(),
|
||||
connection_cache.compute_max_parallel_streams(),
|
||||
QUIC_MAX_UNSTAKED_CONCURRENT_STREAMS
|
||||
);
|
||||
|
||||
staked_nodes.write().unwrap().total_stake = 10000;
|
||||
assert_eq!(
|
||||
connection_cache.compute_max_parallel_chunks(),
|
||||
connection_cache.compute_max_parallel_streams(),
|
||||
QUIC_MAX_UNSTAKED_CONCURRENT_STREAMS
|
||||
);
|
||||
|
||||
|
@ -754,7 +754,7 @@ mod tests {
|
|||
.pubkey_stake_map
|
||||
.insert(pubkey, 1);
|
||||
assert_eq!(
|
||||
connection_cache.compute_max_parallel_chunks(),
|
||||
connection_cache.compute_max_parallel_streams(),
|
||||
QUIC_MIN_STAKED_CONCURRENT_STREAMS
|
||||
);
|
||||
|
||||
|
@ -769,7 +769,7 @@ mod tests {
|
|||
.pubkey_stake_map
|
||||
.insert(pubkey, 1000);
|
||||
assert_ne!(
|
||||
connection_cache.compute_max_parallel_chunks(),
|
||||
connection_cache.compute_max_parallel_streams(),
|
||||
QUIC_MIN_STAKED_CONCURRENT_STREAMS
|
||||
);
|
||||
}
|
||||
|
|
|
@ -270,21 +270,21 @@ pub struct QuicClient {
|
|||
connection: Arc<Mutex<Option<QuicNewConnection>>>,
|
||||
addr: SocketAddr,
|
||||
stats: Arc<ClientStats>,
|
||||
num_chunks: usize,
|
||||
chunk_size: usize,
|
||||
}
|
||||
|
||||
impl QuicClient {
|
||||
pub fn new(
|
||||
endpoint: Arc<QuicLazyInitializedEndpoint>,
|
||||
addr: SocketAddr,
|
||||
num_chunks: usize,
|
||||
chunk_size: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
endpoint,
|
||||
connection: Arc::new(Mutex::new(None)),
|
||||
addr,
|
||||
stats: Arc::new(ClientStats::default()),
|
||||
num_chunks,
|
||||
chunk_size,
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -451,21 +451,6 @@ impl QuicClient {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn compute_chunk_length(num_buffers_to_chunk: usize, num_chunks: usize) -> usize {
|
||||
// The function is equivalent to checked div_ceil()
|
||||
// Also, if num_chunks == 0 || num_buffers_to_chunk == 0, return 1
|
||||
num_buffers_to_chunk
|
||||
.checked_div(num_chunks)
|
||||
.map_or(1, |value| {
|
||||
if num_buffers_to_chunk.checked_rem(num_chunks).unwrap_or(0) != 0 {
|
||||
value.saturating_add(1)
|
||||
} else {
|
||||
value
|
||||
}
|
||||
})
|
||||
.max(1)
|
||||
}
|
||||
|
||||
pub async fn send_batch<T>(
|
||||
&self,
|
||||
buffers: &[T],
|
||||
|
@ -498,8 +483,7 @@ impl QuicClient {
|
|||
// by just getting a reference to the NewConnection once
|
||||
let connection_ref: &NewConnection = &connection;
|
||||
|
||||
let chunk_len = Self::compute_chunk_length(buffers.len() - 1, self.num_chunks);
|
||||
let chunks = buffers[1..buffers.len()].iter().chunks(chunk_len);
|
||||
let chunks = buffers[1..buffers.len()].iter().chunks(self.chunk_size);
|
||||
|
||||
let futures: Vec<_> = chunks
|
||||
.into_iter()
|
||||
|
@ -608,29 +592,3 @@ impl TpuConnection for QuicTpuConnection {
|
|||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::nonblocking::quic_client::QuicClient;
|
||||
|
||||
#[test]
|
||||
fn test_transaction_batch_chunking() {
|
||||
assert_eq!(QuicClient::compute_chunk_length(0, 0), 1);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 0), 1);
|
||||
assert_eq!(QuicClient::compute_chunk_length(0, 10), 1);
|
||||
assert_eq!(QuicClient::compute_chunk_length(usize::MAX, usize::MAX), 1);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, usize::MAX), 1);
|
||||
assert!(QuicClient::compute_chunk_length(usize::MAX, 10) == (usize::MAX / 10) + 1);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 1), 10);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 2), 5);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 3), 4);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 4), 3);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 5), 2);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 6), 2);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 7), 2);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 8), 2);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 9), 2);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 10), 1);
|
||||
assert_eq!(QuicClient::compute_chunk_length(10, 11), 1);
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue