Fix calculation of compressed data buffer size

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]

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.