Fix PQScratch memory leak (microsoft#522)

* fix memory leak * FIXED clang-format error * FIXED SSDQueryScratch Destroy OOM
yugabyte · Apr 16, 2024 · 3fca76c · 3fca76c
1 parent 87653ba
commit 3fca76c
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git 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
diff --git 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
@@ -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)