feat(io): add Resize API for efficient storage allocation (#1571)

- Add Resize() method to basic_io.h with ResizeImpl support via SFINAE
- Implement ResizeImpl for async_io, buffer_io, and memory_block_io
- Replace Write-based expansion hack with Resize in datacell classes
- Add null pointer check for memory_block_io allocation

Signed-off-by: LHT129 <tianlan.lht@antgroup.com>

Author: LHT129

src/datacell/extra_info_datacell.h

@@ -54,9 +54,7 @@ class ExtraInfoDataCell : public ExtraInfoInterface {
         this->max_capacity_ = new_capacity;
         uint64_t io_size =
             static_cast<uint64_t>(new_capacity) * static_cast<uint64_t>(extra_info_size_);
-        uint8_t end_flag =
-            127;  // the value is meaningless, only to occupy the position for io allocate
-        this->io_->Write(&end_flag, 1, io_size);
+        this->io_->Resize(io_size);
     }
 
     void

src/datacell/flatten_datacell.h

@@ -92,9 +92,7 @@ class FlattenDataCell : public FlattenInterface {
         }
         this->max_capacity_ = new_capacity;
         uint64_t io_size = static_cast<uint64_t>(new_capacity) * static_cast<uint64_t>(code_size_);
-        uint8_t end_flag =
-            127;  // the value is meaningless, only to occupy the position for io allocate
-        this->io_->Write(&end_flag, 1, io_size);
+        this->io_->Resize(io_size);
     }
 
     void

src/datacell/graph_datacell.h

@@ -264,9 +264,7 @@ GraphDataCell<IOTmpl>::Resize(InnerIdType new_size) {
     }
     this->max_capacity_ = new_size;
     uint64_t io_size = static_cast<uint64_t>(new_size) * static_cast<uint64_t>(code_line_size_);
-    uint8_t end_flag =
-        127;  // the value is meaningless, only to occupy the position for io allocate
-    this->io_->Write(&end_flag, 1, io_size);
+    this->io_->Resize(io_size);
 }
 
 template <typename IOTmpl>

src/datacell/sparse_vector_datacell.h

@@ -80,10 +80,8 @@ class SparseVectorDataCell : public FlattenInterface {
         }
         uint64_t io_size = (new_capacity - total_count_) * max_code_size_ + current_offset_;
         this->max_capacity_ = new_capacity;
-        uint8_t end_flag =
-            127;  // the value is meaingless, only to occupy the position for io allocate
-        this->io_->Write(&end_flag, 1, io_size);
-        this->offset_io_->Write(&end_flag, 1, new_capacity * sizeof(uint32_t));
+        this->io_->Resize(io_size);
+        this->offset_io_->Resize(static_cast<uint64_t>(new_capacity) * sizeof(uint32_t));
     }
 
     void

src/io/async_io.cpp

@@ -81,6 +81,19 @@ AsyncIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) {
     fsync(wfd_);
 }
 
+void
+AsyncIO::ResizeImpl(uint64_t size) {
+#ifdef __APPLE__
+    auto ret = ftruncate(this->wfd_, static_cast<off_t>(size));
+#else
+    auto ret = ftruncate64(this->wfd_, static_cast<int64_t>(size));
+#endif
+    if (ret == -1) {
+        throw VsagException(ErrorType::INTERNAL_ERROR, "ftruncate failed");
+    }
+    this->size_ = size;
+}
+
 bool
 AsyncIO::ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const {
     bool need_release = true;

src/io/async_io.h

@@ -40,6 +40,9 @@ class AsyncIO : public BasicIO<AsyncIO> {
     void
     WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset);
 
+    void
+    ResizeImpl(uint64_t size);
+
     bool
     ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const;
 

src/io/basic_io.h

@@ -218,6 +218,24 @@ class BasicIO {
         }
     }
 
+    inline void
+    Resize(uint64_t size) {
+        if constexpr (has_ResizeImpl<IOTmpl>::value) {
+            return cast().ResizeImpl(size);
+        } else {
+            if (size <= this->size_) {
+                return;
+            }
+            ByteBuffer buffer(SERIALIZE_BUFFER_SIZE, this->allocator_);
+            uint64_t offset = this->size_;
+            while (offset < size) {
+                auto cur_size = std::min(SERIALIZE_BUFFER_SIZE, size - offset);
+                this->Write(buffer.data, cur_size, offset);
+                offset += cur_size;
+            }
+        }
+    }
+
     inline int64_t
     GetMemoryUsage() const {
         return this->size_;
@@ -314,5 +332,6 @@ class BasicIO {
                                  std::declval<uint64_t>())
     GENERATE_HAS_MEMBER_FUNCTION(ReleaseImpl, void, std::declval<const uint8_t*>())
     GENERATE_HAS_MEMBER_FUNCTION(InitIOImpl, void, std::declval<const IOParamPtr&>())
+    GENERATE_HAS_MEMBER_FUNCTION(ResizeImpl, void, std::declval<uint64_t>())
 };
 }  // namespace vsag

src/io/buffer_io.cpp

@@ -58,6 +58,19 @@ BufferIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) {
     }
 }
 
+void
+BufferIO::ResizeImpl(uint64_t size) {
+#ifdef __APPLE__
+    auto ret = ftruncate(this->fd_, static_cast<off_t>(size));
+#else
+    auto ret = ftruncate64(this->fd_, static_cast<int64_t>(size));
+#endif
+    if (ret == -1) {
+        throw VsagException(ErrorType::INTERNAL_ERROR, "ftruncate failed");
+    }
+    this->size_ = size;
+}
+
 bool
 BufferIO::ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const {
     if (size == 0) {

src/io/buffer_io.h

@@ -48,6 +48,9 @@ class BufferIO : public BasicIO<BufferIO> {
     void
     WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset);
 
+    void
+    ResizeImpl(uint64_t size);
+
     bool
     ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const;
 

src/io/memory_block_io.cpp

@@ -129,11 +129,25 @@ MemoryBlockIO::check_and_realloc(uint64_t size) {
     auto cur_block_size = this->blocks_.size();
     this->blocks_.reserve(new_block_count);
     while (cur_block_size < new_block_count) {
-        this->blocks_.emplace_back(static_cast<uint8_t*>(this->allocator_->Allocate(block_size_)));
+        auto* ptr = static_cast<uint8_t*>(this->allocator_->Allocate(block_size_));
+        if (ptr == nullptr) {
+            throw VsagException(ErrorType::INTERNAL_ERROR, "MemoryBlockIO allocation failed");
+        }
+        this->blocks_.emplace_back(ptr);
         ++cur_block_size;
     }
 }
 
+void
+MemoryBlockIO::ResizeImpl(uint64_t size) {
+    if (size <= this->size_) {
+        this->size_ = size;
+        return;
+    }
+    check_and_realloc(size);
+    this->size_ = size;
+}
+
 static int
 countr_zero(uint64_t x) {
     if (x == 0) {