1use {
2 crate::packet::PacketBatch,
3 rand::{thread_rng, Rng},
4};
5
6pub fn discard_batches_randomly(
7 batches: &mut Vec<PacketBatch>,
8 max_packets: usize,
9 mut total_packets: usize,
10) -> usize {
11 while total_packets > max_packets {
12 let index = thread_rng().gen_range(0..batches.len());
13 let removed = batches.swap_remove(index);
14 total_packets = total_packets.saturating_sub(removed.len());
15 }
16 total_packets
17}
18
19#[cfg(test)]
20mod tests {
21 use {super::*, crate::packet::Packet};
22
23 #[test]
24 fn test_batch_discard_random() {
25 solana_logger::setup();
26 let mut batch = PacketBatch::default();
27 batch.resize(1, Packet::default());
28 let num_batches = 100;
29 let mut batches = vec![batch; num_batches];
30 let max = 5;
31 discard_batches_randomly(&mut batches, max, num_batches);
32 assert_eq!(batches.len(), max);
33 }
34}