From 7865851c02bc3b1cbe22f8e95011541893fd2618 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Danie=CC=88l=20de=20Kok?= Date: Tue, 6 Aug 2024 11:45:20 +0000 Subject: [PATCH] Test partially overlapping prefills. --- backends/v3/src/block_allocator.rs | 44 +++++++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/backends/v3/src/block_allocator.rs b/backends/v3/src/block_allocator.rs index 01b213d8..37c71653 100644 --- a/backends/v3/src/block_allocator.rs +++ b/backends/v3/src/block_allocator.rs @@ -393,7 +393,7 @@ mod tests { } #[test] - fn correctly_free_when_fully_overlapping_prefills_in_flight() { + fn correctly_free_fully_overlapping_prefills() { let mut cache = RadixAllocator::new(1, 10, None); let allocation1 = cache.allocate(4, Some(Arc::new(vec![0, 1, 2, 3]))).unwrap(); let allocation2 = cache.allocate(4, Some(Arc::new(vec![0, 1, 2, 3]))).unwrap(); @@ -407,4 +407,46 @@ mod tests { // 10 blocks, of which 1 reserved for health checks, 4 for the cached blocks. assert_eq!(cache.free_blocks.len(), 5); } + + #[test] + fn correctly_free_partially_overlapping_prefills() { + let mut cache = RadixAllocator::new(1, 20, None); + let allocation1 = cache.allocate(4, Some(Arc::new(vec![0, 1]))).unwrap(); + assert_eq!(allocation1.0, vec![16, 17, 18, 19]); + assert_eq!(allocation1.2, 0); + + cache.free(allocation1.0, allocation1.3); + + let allocation2 = cache + .allocate(8, Some(Arc::new(vec![0, 1, 2, 3, 4, 5]))) + .unwrap(); + assert_eq!(allocation2.0, vec![16, 17, 12, 13, 14, 15, 18, 19]); + assert_eq!(allocation2.2, 2); + + let allocation3 = cache + .allocate(8, Some(Arc::new(vec![0, 1, 2, 3, 6, 7]))) + .unwrap(); + assert_eq!(allocation3.0, vec![16, 17, 6, 7, 8, 9, 10, 11]); + assert_eq!(allocation3.2, 2); + + cache.free(allocation3.0, allocation3.3); + cache.free(allocation2.0, allocation2.3); + + // 20 blocks, of which 1 reserved for health checks, 6 for allocation3, 2 for allocation2. + assert_eq!(cache.free_blocks.len(), 11); + + let allocation4 = cache + .allocate(6, Some(Arc::new(vec![0, 1, 2, 3, 4, 5]))) + .unwrap(); + assert_eq!(allocation4.0, vec![16, 17, 6, 7, 14, 15]); + assert_eq!(allocation4.2, 6); + assert_eq!(cache.free_blocks.len(), 11); + + let allocation5 = cache + .allocate(6, Some(Arc::new(vec![0, 1, 2, 3, 6, 7]))) + .unwrap(); + assert_eq!(allocation5.0, vec![16, 17, 6, 7, 8, 9]); + assert_eq!(allocation5.2, 6); + assert_eq!(cache.free_blocks.len(), 11); + } }