fix: add related bin log size into search stream reduce

docs/design_docs/segcore/segment_interface.md

@@ -4,7 +4,7 @@
 
 1. `get_row_count`: Get the number of entities in the segment
 2. `get_schema`: Get the corresponding collection schema in the segment
-3. `GetMemoryUsageInBytes`: Get memory usage of a segment
+3. `GetResourceUsage`: Get memory and disk usage of a segment.
 4. `Search(plan, placeholderGroup, timestamp) -> QueryResult`: Perform search operations according to the plan containing search parameters and predicate conditions, and return search results. Ensure that the time of all search results is before the specified timestamp(MVCC)
 5. `FillTargetEntry(plan, &queryResult)`: Fill the missing column data for search results based on target columns in the plan
 

internal/core/src/common/File.h

@@ -35,26 +35,61 @@
     static File
     Open(const std::string_view filepath, int flags) {
         int fd = open(filepath.data(), flags, S_IRUSR | S_IWUSR);
-        AssertInfo(fd != -1,
-                   "failed to create mmap file {}: {}",
-                   filepath,
-                   strerror(errno));
+        if (fd < 0) {
+            throw SegcoreError(ErrorCode::UnistdError,
+                               fmt::format("failed to open file at {}: {}",
+                                           filepath,
+                                           strerror(errno)));
+        }
         return File(fd);
     }
 
     int
     Descriptor() const {
+        if (fd_ < 0) {
+            throw SegcoreError(ErrorCode::UnistdError,
+                               "file descriptor is invalid");
+        }
         return fd_;
     }
 
     ssize_t
     Write(const void* buf, size_t size) {
-        return write(fd_, buf, size);
+        if (fd_ < 0) {
+            throw SegcoreError(ErrorCode::UnistdError,
+                               "file descriptor is invalid when writing file");
+        }
+
+        ssize_t cnt = write(fd_, buf, size);
+        if (cnt < 0) {
+            throw SegcoreError(ErrorCode::UnistdError,
+                               fmt::format("failed to write file: stderror: {}",
+                                           strerror(errno)));
+        }
+        if (cnt != size) {
+            throw SegcoreError(
+                ErrorCode::UnistdError,
+                fmt::format("short write to file: written: {}, expected: {}",
+                            cnt,
+                            size));
+        }
+        return cnt;
     }
 
     offset_t
     Seek(offset_t offset, int whence) {
-        return lseek(fd_, offset, whence);
+        if (fd_ < 0) {
+            throw SegcoreError(ErrorCode::UnistdError,
+                               "file descriptor is invalid when seek file");
+        }
+
+        offset_t position = lseek(fd_, offset, whence);
+        if (position < 0) {
+            throw SegcoreError(ErrorCode::UnistdError,
+                               fmt::format("failed to seek file: stderror: {}",
+                                           strerror(errno)));
+        }
+        return position;
     }
 
     void

internal/core/src/common/Types.h

@@ -507,6 +507,11 @@ struct TypeTraits<DataType::VECTOR_FLOAT> {
     static constexpr const char* Name = "VECTOR_FLOAT";
 };
 
+struct ResourceUsage {
+    size_t mem_size;
+    size_t disk_size;
+};
+
 }  // namespace milvus
 template <>
 struct fmt::formatter<milvus::DataType> : formatter<string_view> {

internal/core/src/common/type_c.h

@@ -16,6 +16,7 @@
 
 #pragma once
 #include <stdbool.h>
+#include <stddef.h>
 #include <stdint.h>
 
 #ifdef __cplusplus
@@ -63,6 +64,12 @@ typedef struct CStatus {
     const char* error_msg;
 } CStatus;
 
+// CSegmentResourceUsage is used to represent the resource usage of a segment, and used by CGO.
+typedef struct CSegmentResourceUsage {
+    size_t mem_size;
+    size_t disk_size;
+} CSegmentResourceUsage;
+
 typedef struct CProto {
     const void* proto_blob;
     int64_t proto_size;

internal/core/src/index/Index.h

@@ -89,6 +89,9 @@ class IndexBase {
         return index_type_;
     }
 
+    virtual ResourceUsage
+    GetResourceUsage() const = 0;
+
  protected:
     explicit IndexBase(IndexType index_type)
         : index_type_(std::move(index_type)) {

internal/core/src/index/InvertedIndexTantivy.cpp

@@ -447,6 +447,14 @@ InvertedIndexTantivy<T>::BuildWithRawData(size_t n,
     }
 }
 
+template <typename T>
+ResourceUsage
+InvertedIndexTantivy<T>::GetResourceUsage() const {
+    // TODO: we can't get the memory usage of inverted index now.
+    auto disk_size = disk_file_manager_->GetLocalFileSize();
+    return ResourceUsage{0, disk_size};
+}
+
 template class InvertedIndexTantivy<bool>;
 template class InvertedIndexTantivy<int8_t>;
 template class InvertedIndexTantivy<int16_t>;

internal/core/src/index/InvertedIndexTantivy.h

@@ -156,6 +156,9 @@ class InvertedIndexTantivy : public ScalarIndex<T> {
     const TargetBitmap
     RegexQuery(const std::string& pattern) override;
 
+    ResourceUsage
+    GetResourceUsage() const override;
+
  private:
     void
     finish();

internal/core/src/index/ScalarIndexSort.cpp

@@ -479,6 +479,15 @@ ScalarIndexSort<T>::ShouldSkip(const T lower_value,
     return true;
 }
 
+template <typename T>
+ResourceUsage
+ScalarIndexSort<T>::GetResourceUsage() const {
+    auto mem_size = data_.capacity() * sizeof(IndexStructure<T>) +
+                    idx_to_offsets_.capacity() * sizeof(int32_t);
+    // TODO: after enabling mmap, move the mem_size as disk_size.
+    return ResourceUsage{mem_size, 0};
+}
+
 template class ScalarIndexSort<bool>;
 template class ScalarIndexSort<int8_t>;
 template class ScalarIndexSort<int16_t>;

internal/core/src/index/ScalarIndexSort.h

@@ -100,6 +100,9 @@ class ScalarIndexSort : public ScalarIndex<T> {
         return true;
     }
 
+    ResourceUsage
+    GetResourceUsage() const override;
+
  private:
     bool
     ShouldSkip(const T lower_value, const T upper_value, const OpType op);

internal/core/src/index/StringIndexMarisa.cpp

@@ -37,6 +37,7 @@
 #include "index/Index.h"
 #include "storage/Util.h"
 #include "storage/space.h"
+#include "storage/LocalChunkManagerSingleton.h"
 
 namespace milvus::index {
 
@@ -256,18 +257,21 @@
     auto len = index->size;
 
     auto file = File::Open(file_name, O_RDWR | O_CREAT | O_EXCL);
-    auto written = file.Write(index->data.get(), len);
-    if (written != len) {
+    try {
+        auto written = file.Write(index->data.get(), len);
+    } catch (const SegcoreError& e) {
         file.Close();
         remove(file_name.c_str());
-        throw SegcoreError(
-            ErrorCode::UnistdError,
-            fmt::format("write index to fd error: {}", strerror(errno)));
+        throw;
     }
 
     file.Seek(0, SEEK_SET);
     if (config.contains(kEnableMmap)) {
         trie_.mmap(file_name.c_str());
+        auto local_chunk_manager =
+            milvus::storage::LocalChunkManagerSingleton::GetInstance()
+                .GetChunkManager();
+        resource_usage_.disk_size += local_chunk_manager->Size(file_name);
     } else {
         trie_.read(file.Descriptor());
     }
@@ -563,4 +567,10 @@
     return std::string(agent.key().ptr(), agent.key().length());
 }
 
+ResourceUsage
+StringIndexMarisa::GetResourceUsage() const {
+    // TODO: we cannot estimate the memory usage of marisa trie.
+    return resource_usage_;
+}
+
 }  // namespace milvus::index

internal/core/src/index/StringIndexMarisa.h

@@ -98,6 +98,9 @@ class StringIndexMarisa : public StringIndex {
         return true;
     }
 
+    ResourceUsage
+    GetResourceUsage() const override;
+
  private:
     void
     fill_str_ids(size_t n, const std::string* values);
@@ -123,6 +126,7 @@ class StringIndexMarisa : public StringIndex {
     bool built_ = false;
     std::shared_ptr<storage::MemFileManagerImpl> file_manager_;
     std::shared_ptr<milvus_storage::Space> space_;
+    ResourceUsage resource_usage_{};
 };
 
 using StringIndexMarisaPtr = std::unique_ptr<StringIndexMarisa>;

internal/core/src/index/VectorDiskIndex.cpp

@@ -510,6 +510,14 @@
     return load_config;
 }
 
+template <typename T>
+ResourceUsage
+VectorDiskAnnIndex<T>::GetResourceUsage() const {
+    // TODO: we can't get the memory usage of vector index now.
+    auto disk_size = file_manager_->GetLocalFileSize();
+    return ResourceUsage{0, disk_size};
+}
+
 template class VectorDiskAnnIndex<float>;
 template class VectorDiskAnnIndex<float16>;
 template class VectorDiskAnnIndex<bfloat16>;

internal/core/src/index/VectorDiskIndex.h

@@ -19,6 +19,7 @@
 #include <memory>
 #include <vector>
 
+#include "common/Types.h"
 #include "index/VectorIndex.h"
 #include "storage/DiskFileManagerImpl.h"
 #include "storage/space.h"
@@ -112,6 +113,9 @@ class VectorDiskAnnIndex : public VectorIndex {
                     const knowhere::Json& json,
                     const BitsetView& bitset) const override;
 
+    ResourceUsage
+    GetResourceUsage() const override;
+
  private:
     knowhere::Json
     update_load_json(const Config& config);

internal/core/src/index/VectorMemIndex.cpp

@@ -47,6 +47,7 @@
 #include "log/Log.h"
 #include "storage/DataCodec.h"
 #include "storage/MemFileManagerImpl.h"
+#include "storage/LocalChunkManagerSingleton.h"
 #include "storage/ThreadPools.h"
 #include "storage/space.h"
 #include "storage/Util.h"
@@ -760,14 +761,9 @@
                                "lost index slice data");
                     auto data = batch_data[file_name];
                     auto start_write_file = std::chrono::system_clock::now();
-                    auto written = file.Write(data->Data(), data->Size());
+                    file.Write(data->Data(), data->Size());
                     write_disk_duration_sum +=
                         (std::chrono::system_clock::now() - start_write_file);
-                    AssertInfo(
-                        written == data->Size(),
-                        fmt::format("failed to write index data to disk {}: {}",
-                                    filepath->data(),
-                                    strerror(errno)));
                 }
                 for (auto& file : batch) {
                     pending_index_files.erase(file);
@@ -831,6 +827,11 @@
     auto dim = index_.Dim();
     this->SetDim(index_.Dim());
 
+    auto local_chunk_manager =
+        milvus::storage::LocalChunkManagerSingleton::GetInstance()
+            .GetChunkManager();
+    resource_usage_.disk_size += local_chunk_manager->Size(filepath.value());
+
     auto ok = unlink(filepath->data());
     AssertInfo(ok == 0,
                "failed to unlink mmap index file {}: {}",
@@ -953,6 +954,14 @@
                strerror(errno));
     LOG_INFO("load vector index done");
 }
+
+template <typename T>
+ResourceUsage
+VectorMemIndex<T>::GetResourceUsage() const {
+    // TODO: we can't get the memory usage of vector index now.
+    return resource_usage_;
+}
+
 template class VectorMemIndex<float>;
 template class VectorMemIndex<uint8_t>;
 template class VectorMemIndex<float16>;

internal/core/src/index/VectorMemIndex.h

@@ -111,8 +111,13 @@ class VectorMemIndex : public VectorIndex {
     void
     LoadFromFileV2(const Config& config);
 
+    ResourceUsage
+    GetResourceUsage() const override;
+
  protected:
     Config config_;
+
+    ResourceUsage resource_usage_{};
     knowhere::Index<knowhere::IndexNode> index_;
     std::shared_ptr<storage::MemFileManagerImpl> file_manager_;
     std::shared_ptr<milvus_storage::Space> space_;