solana_perf/
discard.rs

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}