From 62aaf8783e4649dc04e648f71a6321fdb5cd7105 Mon Sep 17 00:00:00 2001
From: Yuan <yuanzhou@apache.org>
Date: Mon, 22 Sep 2025 12:37:58 +0100
Subject: [PATCH] [VL] reduce lock scope for objectstore

Signed-off-by: Yuan <yuanzhou@apache.org>
---
 cpp/core/utils/ObjectStore.cc | 8 ++++----
 cpp/core/utils/ObjectStore.h  | 7 ++++---
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/cpp/core/utils/ObjectStore.cc b/cpp/core/utils/ObjectStore.cc
index 05fd5843733..9d48ccdd8ac 100644
--- a/cpp/core/utils/ObjectStore.cc
+++ b/cpp/core/utils/ObjectStore.cc
@@ -21,8 +21,8 @@
 
 // static
 std::unique_ptr<gluten::ObjectStore> gluten::ObjectStore::create() {
-  static std::mutex mtx;
-  std::lock_guard<std::mutex> lock(mtx);
+  static std::shared_mutex mtx;
+  std::unique_lock lock(mtx);
   StoreHandle nextId = stores().nextId();
   auto store = std::unique_ptr<gluten::ObjectStore>(new gluten::ObjectStore(nextId));
   StoreHandle storeId = safeCast<StoreHandle>(stores().insert(store.get()));
@@ -51,7 +51,7 @@ gluten::ObjectStore::~ObjectStore() {
     }
     std::shared_ptr<void> tempObj;
     {
-      const std::lock_guard<std::mutex> lock(mtx_);
+      std::unique_lock lock(mtx_);
       // destructing in reversed order (the last added object destructed first)
       auto itr = aliveObjects_.rbegin();
       const ResourceHandle handle = (*itr).first;
@@ -75,7 +75,7 @@ gluten::ObjectStore::~ObjectStore() {
 }
 
 void gluten::ObjectStore::releaseInternal(gluten::ResourceHandle handle) {
-  const std::lock_guard<std::mutex> lock(mtx_);
+  std::unique_lock lock(mtx_);
   store_.erase(handle);
   aliveObjects_.erase(handle);
 }
diff --git a/cpp/core/utils/ObjectStore.h b/cpp/core/utils/ObjectStore.h
index 24fab4d423e..fbf3d08cf55 100644
--- a/cpp/core/utils/ObjectStore.h
+++ b/cpp/core/utils/ObjectStore.h
@@ -18,6 +18,7 @@
 #pragma once
 
 #include <map>
+#include <shared_mutex>
 #include "utils/ResourceMap.h"
 
 namespace gluten {
@@ -70,7 +71,7 @@ class ObjectStore {
 
   template <typename T>
   ObjectHandle save(std::shared_ptr<T> obj) {
-    const std::lock_guard<std::mutex> lock(mtx_);
+    std::unique_lock lock(mtx_);
     const std::string_view typeName = typeid(T).name();
     const size_t size = SafeSizeOf<T>::value;
     ResourceHandle handle = store_.insert(std::move(obj));
@@ -96,7 +97,7 @@ class ObjectStore {
 
   template <typename T>
   std::shared_ptr<T> retrieveInternal(ResourceHandle handle) {
-    const std::lock_guard<std::mutex> lock(mtx_);
+    std::shared_lock lock(mtx_);
     std::shared_ptr<void> object = store_.lookup(handle);
     // Programming carefully. This will lead to ub if wrong typename T was passed in.
     auto casted = std::static_pointer_cast<T>(object);
@@ -110,6 +111,6 @@ class ObjectStore {
   ResourceMap<std::shared_ptr<void>> store_;
   // Preserves handles of objects in the store in order, with additional attributes associated with them.
   std::map<ResourceHandle, ObjectDebugInfo> aliveObjects_{};
-  std::mutex mtx_;
+  std::shared_mutex mtx_;
 };
 } // namespace gluten