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 {
22 super::*,
23 crate::packet::{BytesPacket, BytesPacketBatch, Meta},
24 bytes::Bytes,
25 };
26
27 #[test]
28 fn test_batch_discard_random() {
29 agave_logger::setup();
30 let mut batch = BytesPacketBatch::new();
31 batch.resize(1, BytesPacket::new(Bytes::new(), Meta::default()));
32 let batch = PacketBatch::from(batch);
33 let num_batches = 100;
34 let mut batches = vec![batch; num_batches];
35 let max = 5;
36 discard_batches_randomly(&mut batches, max, num_batches);
37 assert_eq!(batches.len(), max);
38 }
39}