feat(buffer_object): Add resize method to BufferIO (#57)

Fixes #40

Author: Leahlijuan

src/ml_flashpoint/checkpoint_object_manager/buffer_io.py

@@ -404,6 +404,46 @@ def close(self, truncate: bool = True) -> None:
             # Set position to -1, a common convention for a closed stream.
             self._pos = -1
 
+    def resize(self, new_size: int) -> None:
+        """Resizes the buffer to the new size (including metadata size).
+
+        Args:
+            new_size: The new size of the buffer in bytes. Must be >= METADATA_SIZE.
+        """
+        self._check_validity("write")
+        _LOGGER.info("Resizing BufferIO from %d to %d bytes.", len(self._mv), new_size)
+
+        if new_size < METADATA_SIZE:
+            raise ValueError(f"New size must be at least {METADATA_SIZE} bytes, got {new_size}.")
+
+        # 1. Release the memoryview
+        if self._mv:
+            self._mv.release()
+            self._mv = None
+
+        # 2. Call C++ resize
+        try:
+            self.buffer_obj.resize(new_size)
+        except Exception:
+            _LOGGER.exception("Failed to resize underlying BufferObject.")
+            raise
+
+        # 3. Recreate the memoryview
+        try:
+            self._mv = memoryview(self.buffer_obj)
+        except Exception:
+            _LOGGER.exception("Failed to recreate memoryview after resize.")
+            raise ValueError("Failed to recreate memoryview after resize.")
+
+        # 4. Re-map metadata
+        # Since the buffer might have moved in memory, we need to refresh the metadata view.
+        try:
+            # We are always in writable mode if we are resizing (resize checks !readonly)
+            self._metadata = BufferMetadataType.from_buffer(self._mv[:METADATA_SIZE])
+        except Exception as e:
+            _LOGGER.exception("Failed to recreate metadata object from buffer slice.")
+            raise IOError(f"Could not initialize metadata from buffer: {e}") from e
+
     # --- Properties and Context Manager ---
 
     @property

src/ml_flashpoint/checkpoint_object_manager/buffer_object/bindings.cpp

@@ -53,6 +53,9 @@ PYBIND11_MODULE(buffer_object_ext, m) {
            "Gets the buffer object capacity.",
            py::call_guard<py::gil_scoped_release>())
 
+      .def("resize", &BufferObject::resize, "Resizes the buffer object.",
+           py::arg("new_capacity"), py::call_guard<py::gil_scoped_release>())
+
       .def(
           "get_data_ptr",
           [](const BufferObject& self) -> std::uintptr_t {

src/ml_flashpoint/checkpoint_object_manager/buffer_object/buffer_helper.cpp

@@ -270,5 +270,48 @@ absl::Status unmap_and_close(int fd, void* data_ptr, size_t data_size,
   return absl::OkStatus();
 }
 
+absl::Status resize_mmap(int fd, size_t new_size, void*& data_ptr,
+                         size_t& curr_size) {
+  // @brief Resizes the memory map.
+  //
+  // This function assumes `data_ptr` is currently mapped with
+  // `curr_size`. It unmaps it, ftruncates the file, and remaps it with
+  // `new_size`.
+
+  if (new_size == curr_size) {
+    return absl::OkStatus();
+  }
+  if (fd == -1) {
+    return absl::InvalidArgumentError("Invalid file descriptor for resize.");
+  }
+  if (data_ptr == MAP_FAILED || curr_size == 0) {
+    return absl::InvalidArgumentError(
+        "Invalid data pointer or size for resize.");
+  }
+
+  // 1. Unmap existing
+  if (munmap(data_ptr, curr_size) == -1) {
+    return ErrnoToStatus("munmap() failed during resize");
+  }
+  data_ptr = nullptr;
+
+  // 2. Truncate
+  if (ftruncate(fd, new_size) == -1) {
+    return ErrnoToStatus("ftruncate() failed during resize");
+  }
+
+  // 3. Mmap new size
+  void* ptr = mmap(NULL, new_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+  if (ptr == MAP_FAILED) {
+    return ErrnoToStatus("mmap() failed during resize");
+  }
+
+  data_ptr = ptr;
+  curr_size = new_size;
+
+  LOG(INFO) << "Successfully resized mmap to " << new_size;
+  return absl::OkStatus();
+}
+
 }  // namespace
    // ml_flashpoint::checkpoint_object_manager::buffer_object::internal

src/ml_flashpoint/checkpoint_object_manager/buffer_object/buffer_helper.h

@@ -45,6 +45,11 @@ absl::Status open_file_and_mmap_ro(const std::string& object_id, int& out_fd,
 // Unmaps memory, optionally truncates the file, and closes the file descriptor.
 absl::Status unmap_and_close(int fd, void* data_ptr, size_t data_size,
                              std::optional<size_t> truncate_size);
+// Resizes the file to `new_size` and remaps it into memory,
+// updating `data_ptr` to point to the resized memory buffer.
+// If new_size == curr_size, this is a no-op and `data_ptr` is returned as is.
+absl::Status resize_mmap(int fd, size_t new_size, void*& data_ptr,
+                         size_t& curr_size);
 
 };  // namespace
     // ml_flashpoint::checkpoint_object_manager::buffer_object::internal

src/ml_flashpoint/checkpoint_object_manager/buffer_object/buffer_object.cpp

@@ -172,4 +172,27 @@ void BufferObject::close(std::optional<size_t> truncate_size) noexcept {
   fd_ = -1;
   data_ptr_ = nullptr;
   capacity_ = 0;
+}
+
+void BufferObject::resize(size_t new_capacity) {
+  if (is_closed()) {
+    throw std::runtime_error("Cannot resize a closed buffer.");
+  }
+  if (is_readonly()) {
+    throw std::runtime_error("Cannot resize a read-only buffer.");
+  }
+  if (new_capacity == 0) {
+    throw std::runtime_error("Cannot resize buffer to 0.");
+  }
+
+  LOG(INFO) << "BufferObject::resize: Resizing object_id=" << object_id_
+            << " from " << capacity_ << " to " << new_capacity;
+
+  absl::Status status =
+      ml_flashpoint::checkpoint_object_manager::buffer_object::internal::
+          resize_mmap(fd_, new_capacity, data_ptr_, capacity_);
+
+  if (!status.ok()) {
+    throw std::runtime_error(status.ToString());
+  }
 }

src/ml_flashpoint/checkpoint_object_manager/buffer_object/buffer_object.h

@@ -71,6 +71,9 @@ class BufferObject {
   // close the buffer object, optionally truncates the file before closing
   void close(std::optional<size_t> truncate_size = std::nullopt) noexcept;
 
+  // Resizes the buffer to the new capacity.
+  void resize(size_t new_capacity);
+
  private:
   std::string object_id_;  // The unique identifier of the buffer, typically the
                            // filename.

tests/checkpoint_object_manager/buffer_object/buffer_helper_test.cpp

@@ -611,5 +611,120 @@ TEST_F(BufferHelperTest, UnmapAndCloseFailsWithValidPtrAndZeroSize) {
   ASSERT_EQ(errno, EBADF);
 }
 
+// --- Tests for resize_mmap ---
+
+// Verifies that resizing to a larger size succeeds and preserves data.
+struct ResizeParams {
+  size_t initial_size;
+  size_t new_size;
+};
+
+class ResizeMmapCombinedTest
+    : public BufferHelperTest,
+      public ::testing::WithParamInterface<ResizeParams> {};
+
+TEST_P(ResizeMmapCombinedTest, ResizeMmapSucceeds) {
+  // Given
+  const ResizeParams& params = GetParam();
+  const size_t initial_size = params.initial_size;
+  const size_t new_size = params.new_size;
+  const std::string content = "Some data to preserve";
+
+  int fd = -1;
+  size_t out_size = 0;
+  void* ptr = MAP_FAILED;
+  absl::Status status = create_file_and_mmap(test_path_, initial_size, fd,
+                                             out_size, ptr, /*overwrite=*/true);
+  ASSERT_TRUE(status.ok());
+
+  // Ensure content fits in initial size
+  ASSERT_LE(content.size(), initial_size);
+  std::memcpy(ptr, content.c_str(), content.size());
+
+  // When
+  status = resize_mmap(fd, new_size, ptr, out_size);
+
+  // Then
+  ASSERT_TRUE(status.ok());
+  EXPECT_EQ(out_size, new_size);
+  EXPECT_NE(ptr, nullptr);
+  EXPECT_NE(ptr, MAP_FAILED);
+
+  // Verify content preserved (up to the new size)
+  size_t expected_len = std::min(content.size(), new_size);
+  EXPECT_EQ(std::memcmp(ptr, content.c_str(), expected_len), 0);
+
+  // Cleanup
+  SafeUnmapAndClose(fd, ptr, out_size);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ResizeMmapCombinedTests, ResizeMmapCombinedTest,
+    ::testing::Values(
+        // Larger
+        ResizeParams{1024, 2048},  // Aligned -> Aligned
+        ResizeParams{1024, 2011},  // Aligned -> Unaligned
+        ResizeParams{1025, 2048},  // Unaligned -> Aligned
+        ResizeParams{1025, 2011},  // Unaligned -> Unaligned
+        ResizeParams{4096, 8192},  // Page -> Page
+        ResizeParams{4096, 4097},  // Page -> Unaligned
+        // Smaller
+        ResizeParams{2048, 1024},  // Aligned -> Aligned
+        ResizeParams{2048, 1011},  // Aligned -> Unaligned
+        ResizeParams{2049, 1024},  // Unaligned -> Aligned
+        ResizeParams{2049, 1011},  // Unaligned -> Unaligned
+        ResizeParams{8192, 4096},  // Page -> Page
+        ResizeParams{8192, 4097},  // Page -> Unaligned
+        // Same
+        ResizeParams{1024, 1024},  // Aligned -> Aligned (Same)
+        ResizeParams{4096, 4096},  // Page -> Page (Same)
+        ResizeParams{1025, 1025}   // Unaligned -> Unaligned (Same)
+        ));
+
+// Verifies that resize fails with invalid fd
+TEST_F(BufferHelperTest, ResizeMmapFailsOnInvalidFd) {
+  // Given
+  void* ptr = nullptr;
+  size_t size = 1024;
+
+  // When
+  absl::Status status = resize_mmap(-1, 2048, ptr, size);
+
+  // Then
+  EXPECT_FALSE(status.ok());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+}
+
+// Verifies failure when ftruncate fails (e.g., read-only fd)
+TEST_F(BufferHelperTest, ResizeMmapFailsOnFtruncateFailure) {
+  // Given
+  CreateEmptyFile(test_path_);
+  CreateFileWithContent(test_path_, "data");
+
+  // Open as Read-Only
+  int fd = open(test_path_.c_str(), O_RDONLY);
+  ASSERT_NE(fd, -1);
+
+  // Map it (read-only map)
+  struct stat sb;
+  fstat(fd, &sb);
+  size_t size = sb.st_size;
+  void* ptr = mmap(NULL, size, PROT_READ, MAP_SHARED, fd, 0);
+  ASSERT_NE(ptr, MAP_FAILED);
+
+  // When
+  absl::Status status = resize_mmap(fd, size * 2, ptr, size);
+
+  // Then
+  EXPECT_FALSE(status.ok());
+  EXPECT_THAT(status.message(), testing::HasSubstr("ftruncate() failed"));
+
+  // Cleanup
+  if (ptr != nullptr && ptr != MAP_FAILED) {
+    munmap(ptr, size);
+  }
+  close(fd);
+}
+
 }  // namespace
    // ml_flashpoint::checkpoint_object_manager::buffer_object::internal

tests/checkpoint_object_manager/buffer_object/buffer_object_test.cpp

@@ -407,3 +407,60 @@ TEST_F(BufferObjectTest, MoveAssignmentToSelfIsSafe) {
   EXPECT_EQ(obj.get_id(), original_id);
   EXPECT_EQ(obj.is_readonly(), original_readonly);
 }
+// --- Tests for resize ---
+
+TEST_F(BufferObjectTest, ResizeSucceeds) {
+  const size_t initial_capacity = 1024;
+  const size_t new_capacity = 2048;
+
+  BufferObject buffer(test_path_, initial_capacity);
+  ASSERT_EQ(buffer.get_capacity(), initial_capacity);
+
+  ASSERT_NO_THROW(buffer.resize(new_capacity));
+  EXPECT_EQ(buffer.get_capacity(), new_capacity);
+  EXPECT_FALSE(buffer.is_closed());
+
+  // Verify size on disk
+  EXPECT_EQ(std::filesystem::file_size(test_path_), new_capacity);
+}
+
+TEST_F(BufferObjectTest, ResizePreservesData) {
+  const size_t initial_capacity = 1024;
+  const size_t new_capacity = 2048;
+  const std::string content = "Important Data";
+
+  BufferObject buffer(test_path_, initial_capacity);
+  std::memcpy(buffer.get_data_ptr(), content.c_str(), content.size());
+
+  buffer.resize(new_capacity);
+
+  // Verify data is still there
+  EXPECT_EQ(std::memcmp(buffer.get_data_ptr(), content.c_str(), content.size()),
+            0);
+}
+
+TEST_F(BufferObjectTest, ResizeFailsOnClosedBuffer) {
+  BufferObject buffer(test_path_, 1024);
+  buffer.close();
+  ASSERT_TRUE(buffer.is_closed());
+
+  ASSERT_THROW(buffer.resize(2048), std::runtime_error);
+}
+
+TEST_F(BufferObjectTest, ResizeFailsOnReadOnlyBuffer) {
+  // Create file
+  {
+    BufferObject writer(test_path_, 1024);
+  }
+
+  // Open RO
+  BufferObject reader(test_path_);
+  ASSERT_TRUE(reader.is_readonly());
+
+  ASSERT_THROW(reader.resize(2048), std::runtime_error);
+}
+
+TEST_F(BufferObjectTest, ResizeFailsOnZeroCapacity) {
+  BufferObject buffer(test_path_, 1024);
+  ASSERT_THROW(buffer.resize(0), std::runtime_error);
+}