From 1accfeb110a1b2b7e4edbd901dcce5dbe8403ef3 Mon Sep 17 00:00:00 2001
From: 884244693 <884244693@qq.com>
Date: Sat, 8 Feb 2025 04:18:49 +0000
Subject: [PATCH] [fix] ivf calc center id core cause by ivf center id cache
 heap use after free

---
 src/share/vector_index/ob_vector_index_util.cpp |  7 ++++---
 src/share/vector_index/ob_vector_index_util.h   | 14 +++++++++-----
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/share/vector_index/ob_vector_index_util.cpp b/src/share/vector_index/ob_vector_index_util.cpp
index 3db14ee1bc..c94fa0b7d2 100644
--- a/src/share/vector_index/ob_vector_index_util.cpp
+++ b/src/share/vector_index/ob_vector_index_util.cpp
@@ -2697,7 +2697,7 @@ ObExprVecIvfCenterIdCache* ObVectorIndexUtil::get_ivf_center_id_cache_ctx(const
       void *cache_ctx_buf = NULL;
       ret = exec_ctx->create_expr_op_ctx(id, sizeof(ObExprVecIvfCenterIdCtx), cache_ctx_buf);
       if (OB_SUCC(ret) && OB_NOT_NULL(cache_ctx_buf)) {
-        cache_ctx = new (cache_ctx_buf) ObExprVecIvfCenterIdCtx(exec_ctx->get_allocator());
+        cache_ctx = new (cache_ctx_buf) ObExprVecIvfCenterIdCtx();
       }
     }
   }
@@ -2715,7 +2715,7 @@ void ObVectorIndexUtil::get_ivf_pq_center_id_cache_ctx(const uint64_t& id, sql::
       void *cache_ctx_buf = NULL;
       ret = exec_ctx->create_expr_op_ctx(id, sizeof(ObExprVecIvfCenterIdCtx), cache_ctx_buf);
       if (OB_SUCC(ret) && OB_NOT_NULL(cache_ctx_buf)) {
-        cache_ctx = new (cache_ctx_buf) ObExprVecIvfCenterIdCtx(exec_ctx->get_allocator());
+        cache_ctx = new (cache_ctx_buf) ObExprVecIvfCenterIdCtx();
       }
     }
   }
@@ -2747,7 +2747,8 @@ int ObVectorIndexUtil::get_ivf_aux_info(share::ObPluginVectorIndexService *servi
           LOG_WARN("failed to get centers from cache", K(ret));
         }
       } else {
-        if (OB_FAIL(service->get_ivf_aux_info(table_id, tablet_id, allocator, centers))) {
+        cache->reuse();
+        if (OB_FAIL(service->get_ivf_aux_info(table_id, tablet_id, cache->get_allocator(), centers))) {
           LOG_WARN("failed to get centers", K(ret));
         } else if (OB_FAIL(cache->update_cache(table_id, tablet_id, centers))) {
           LOG_WARN("failed to update ivf center id cache", K(ret));
diff --git a/src/share/vector_index/ob_vector_index_util.h b/src/share/vector_index/ob_vector_index_util.h
index fac0e5e2c1..320cef03a6 100644
--- a/src/share/vector_index/ob_vector_index_util.h
+++ b/src/share/vector_index/ob_vector_index_util.h
@@ -122,10 +122,11 @@ public:
 class ObExprVecIvfCenterIdCache
 {
 public:
-  ObExprVecIvfCenterIdCache(common::ObIAllocator &allocator)
+  ObExprVecIvfCenterIdCache()
     : table_id_(ObCommonID::INVALID_ID),
       tablet_id_(),
-      centers_()
+      centers_(),
+      allocator_("IvfCIdCache", OB_MALLOC_NORMAL_BLOCK_SIZE, MTL_ID())
   {}
   virtual ~ObExprVecIvfCenterIdCache() {}
   bool hit(ObTableID table_id, ObTabletID tablet_id) { return table_id == table_id_ && tablet_id == tablet_id_; }
@@ -136,10 +137,13 @@ public:
     tablet_id_ = tablet_id;
     return centers_.assign(centers);
   }
+  ObArenaAllocator &get_allocator() { return allocator_; }
+  void reuse() { table_id_ = ObCommonID::INVALID_ID; tablet_id_.reset(); centers_.reuse(); allocator_.reuse(); }
 private:
   ObTableID table_id_;
   ObTabletID tablet_id_;
   ObSEArray<float*, 8> centers_;
+  ObArenaAllocator allocator_;
 };
 
 class ObVectorIndexUtil final
@@ -147,10 +151,10 @@ class ObVectorIndexUtil final
   class ObExprVecIvfCenterIdCtx : public sql::ObExprOperatorCtx
   {
   public:
-    ObExprVecIvfCenterIdCtx(common::ObIAllocator &allocator)
+    ObExprVecIvfCenterIdCtx()
       : ObExprOperatorCtx(),
-        cache_(allocator),
-        pq_cache_(allocator)
+        cache_(),
+        pq_cache_()
     {}
     virtual ~ObExprVecIvfCenterIdCtx() {}
     ObExprVecIvfCenterIdCache *get_cache() { return &cache_; }