diff --git a/crates/modelardb_server/src/storage/compressed_data_buffer.rs b/crates/modelardb_server/src/storage/compressed_data_buffer.rs index 3086e757..ec9b8892 100644 --- a/crates/modelardb_server/src/storage/compressed_data_buffer.rs +++ b/crates/modelardb_server/src/storage/compressed_data_buffer.rs @@ -99,13 +99,13 @@ impl CompressedDataBuffer { } let mut compressed_segments_size = 0; - for compressed_segment_batch in compressed_segments.drain(0..) { - compressed_segments_size += - Self::size_of_compressed_segments(&compressed_segment_batch); - self.compressed_segments.push(compressed_segment_batch); - self.size_in_bytes += compressed_segments_size; + for compressed_segment in compressed_segments.drain(0..) { + compressed_segments_size += Self::size_of_compressed_segments(&compressed_segment); + self.compressed_segments.push(compressed_segment); } + self.size_in_bytes += compressed_segments_size; + self.batch_ids.extend(compressed_segment_batch.batch_ids); Ok(compressed_segments_size) @@ -165,11 +165,14 @@ mod tests { let mut compressed_data_buffer = CompressedDataBuffer::new(table::time_series_table_metadata_arc()); - compressed_data_buffer + let size_in_bytes = compressed_data_buffer .append_compressed_segment_batch(compressed_segment_batch()) .unwrap(); - assert!(compressed_data_buffer.size_in_bytes > 0); + // The batch contains two compressed segments, so the size of both is added to the buffer. + // The returned size must match since the caller reserves memory based on it. + assert_eq!(size_in_bytes, 2 * COMPRESSED_SEGMENTS_SIZE); + assert_eq!(compressed_data_buffer.size_in_bytes, size_in_bytes); } #[tokio::test] diff --git a/crates/modelardb_server/src/storage/compressed_data_manager.rs b/crates/modelardb_server/src/storage/compressed_data_manager.rs index a726149b..5a43a02d 100644 --- a/crates/modelardb_server/src/storage/compressed_data_manager.rs +++ b/crates/modelardb_server/src/storage/compressed_data_manager.rs @@ -378,13 +378,15 @@ mod tests { assert!(data_manager.compressed_data_buffers.contains_key(key)); assert_eq!(data_manager.compressed_queue.pop().unwrap(), key); - assert!( + + // The batch contains two compressed segments, so the size of both is added to the buffer. + assert_eq!( data_manager .compressed_data_buffers .get(key) .unwrap() - .size_in_bytes - > 0 + .size_in_bytes, + 2 * COMPRESSED_SEGMENTS_SIZE ); } @@ -408,13 +410,15 @@ mod tests { .await .unwrap(); - assert!( + // Each insert adds the size of the two compressed segments in the batch. + assert_eq!(previous_size, 2 * COMPRESSED_SEGMENTS_SIZE); + assert_eq!( data_manager .compressed_data_buffers .get(TIME_SERIES_TABLE_NAME) .unwrap() - .size_in_bytes - > previous_size + .size_in_bytes, + 4 * COMPRESSED_SEGMENTS_SIZE ); } @@ -499,10 +503,13 @@ mod tests { .await .unwrap(); - assert!( - -1 < data_manager + // The remaining memory was set to -1 above. Saving the buffer returns exactly the memory + // reserved for its two compressed segments. + assert_eq!( + data_manager .memory_pool - .remaining_compressed_memory_in_bytes() + .remaining_compressed_memory_in_bytes(), + 2 * COMPRESSED_SEGMENTS_SIZE as i64 - 1 ); } @@ -542,11 +549,12 @@ mod tests { .await .unwrap(); + // All the memory used for the buffer is returned when the buffer is saved. assert_eq!( data_manager .memory_pool .remaining_compressed_memory_in_bytes(), - 1405 + 0 ); // There should no longer be any compressed data in memory.