From 3fca76c0e79ecc8a28b6b6f905c6bb47cb36f460 Mon Sep 17 00:00:00 2001
From: luyuncheng <luyuncheng@bytedance.com>
Date: Tue, 2 Apr 2024 21:14:33 +0800
Subject: [PATCH] Fix PQScratch memory leak (#522)

* fix memory leak

* FIXED clang-format error

* FIXED SSDQueryScratch Destroy OOM
---
 include/pq_scratch.h   |  1 +
 src/pq_flash_index.cpp |  4 ++++
 src/scratch.cpp        | 11 ++++++++++-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/include/pq_scratch.h b/include/pq_scratch.h
index 2aa90dbe1..95f1b1395 100644
--- a/include/pq_scratch.h
+++ b/include/pq_scratch.h
@@ -17,6 +17,7 @@ template <typename T> class PQScratch
 
     PQScratch(size_t graph_degree, size_t aligned_dim);
     void initialize(size_t dim, const T *query, const float norm = 1.0f);
+    virtual ~PQScratch();
 };
 
 } // namespace diskann
\ No newline at end of file
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
index 12b186fa3..68bcb8719 100644
--- a/src/pq_flash_index.cpp
+++ b/src/pq_flash_index.cpp
@@ -93,6 +93,10 @@ template <typename T, typename LabelT> PQFlashIndex<T, LabelT>::~PQFlashIndex()
     {
         delete[] _pts_to_labels;
     }
+    if (_medoids != nullptr)
+    {
+        delete[] _medoids;
+    }
 }
 
 template <typename T, typename LabelT> inline uint64_t PQFlashIndex<T, LabelT>::get_node_sector(uint64_t node_id)
diff --git a/src/scratch.cpp b/src/scratch.cpp
index c3836ccf1..1f8a34bb1 100644
--- a/src/scratch.cpp
+++ b/src/scratch.cpp
@@ -117,7 +117,7 @@ template <typename T> SSDQueryScratch<T>::~SSDQueryScratch()
     diskann::aligned_free((void *)sector_scratch);
     diskann::aligned_free((void *)this->_aligned_query_T);
 
-    delete[] this->_pq_scratch;
+    delete this->_pq_scratch;
 }
 
 template <typename T>
@@ -143,6 +143,15 @@ template <typename T> PQScratch<T>::PQScratch(size_t graph_degree, size_t aligne
     memset(rotated_query, 0, aligned_dim * sizeof(float));
 }
 
+template <typename T> PQScratch<T>::~PQScratch()
+{
+    diskann::aligned_free((void *)aligned_pq_coord_scratch);
+    diskann::aligned_free((void *)aligned_pqtable_dist_scratch);
+    diskann::aligned_free((void *)aligned_dist_scratch);
+    diskann::aligned_free((void *)aligned_query_float);
+    diskann::aligned_free((void *)rotated_query);
+}
+
 template <typename T> void PQScratch<T>::initialize(size_t dim, const T *query, const float norm)
 {
     for (size_t d = 0; d < dim; ++d)