From 3a094ec94586c95eff6a1c3c19b68d7171e2282e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B0=A7?= <57036417+YReddice@users.noreply.github.com>
Date: Sat, 12 Oct 2024 07:50:10 +0800
Subject: [PATCH] fix diskann memory leak in diskann (#2019)

### What problem does this PR solve?

Fixes #1995 mentioned memory leak in diskann.

Issue link:#1988

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 .../knn_diskann/inner/diskann_utils.cppm      | 20 ++++++++++++-------
 .../knn_index/knn_diskann/vamana_alg.cppm     |  7 ++++++-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/storage/knn_index/knn_diskann/inner/diskann_utils.cppm b/src/storage/knn_index/knn_diskann/inner/diskann_utils.cppm
index d03dbcb1b8..44f4fcf1df 100644
--- a/src/storage/knn_index/knn_diskann/inner/diskann_utils.cppm
+++ b/src/storage/knn_index/knn_diskann/inner/diskann_utils.cppm
@@ -72,7 +72,7 @@ public:
     void Insert(const Neighbor &nbr) {
         if (size_ == capacity_ && data_[size_ - 1] < nbr)
             return;
-        
+
         auto it = std::lower_bound(data_.begin(), data_.begin() + size_, nbr);
         // if (it != data_.start()+size_ && (*it).id == nbr.id) {
         //     return;
@@ -315,7 +315,7 @@ public:
         return ret;
     }
 
-    bool Empty() { return (this->size() == 0); }
+    bool Empty() { return (this->Size() == 0); }
 
     // PUSH BACK
     void Push(T &new_val) {
@@ -380,17 +380,23 @@ public:
     T *ScratchSpace() { return scratch_; }
 
     ~ScratchStoreManager() {
-        scratch_->Clear();
-        scratch_pool_.Push(scratch_);
-        scratch_pool_.PushNotifyAll();
+        if (scratch_ != nullptr) {
+            scratch_->Clear();
+            scratch_pool_.Push(scratch_);
+            scratch_pool_.PushNotifyAll();
+        }
     }
 
     void Destroy() {
+        if (scratch_!= nullptr) {
+            delete scratch_;
+            scratch_ = nullptr;
+        }
         while (!scratch_pool_.Empty()) {
-            auto scratch = scratch_pool_.pop();
+            auto scratch = scratch_pool_.Pop();
             while (scratch == nullptr) {
                 scratch_pool_.WaitForPushNotify();
-                scratch = scratch_pool_.pop();
+                scratch = scratch_pool_.Pop();
             }
             delete scratch;
         }
diff --git a/src/storage/knn_index/knn_diskann/vamana_alg.cppm b/src/storage/knn_index/knn_diskann/vamana_alg.cppm
index d011f7a635..ffab8c4c0b 100644
--- a/src/storage/knn_index/knn_diskann/vamana_alg.cppm
+++ b/src/storage/knn_index/knn_diskann/vamana_alg.cppm
@@ -79,7 +79,12 @@ public:
         }
         return *this;
     }
-    ~MemVamana() = default;
+    ~MemVamana() {
+        if (!query_scratch_.Empty()) {
+            ScratchStoreManager<InMemQueryScratch> manager(query_scratch_);
+            manager.Destroy();
+        }
+    }
 
     static This Make(SizeT L, SizeT R, SizeT dim, SizeT point_num) {
         This ret(L, R, dim, point_num);