[GLUTEN-2974] [VL] Minor refactor for AllocationListener and VeloxMemoryMamanger to non copyable and non assignable

cpp/core/CMakeLists.txt

@@ -182,15 +182,17 @@ file(COPY utils/metrics.h DESTINATION ${root_directory}/releases/include/common/
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-deprecated-declarations -Wno-attributes")
 set(SPARK_COLUMNAR_PLUGIN_SRCS
-        jni/JniWrapper.cc
+        ${SUBSTRAIT_PROTO_SRCS}
+        ${GLUTEN_PROTO_SRCS}
+        compute/ProtobufUtils.cc
         compute/Backend.cc
         compute/ResultIterator.cc
         config/GlutenConfig.cc
+        jni/JniWrapper.cc
+        memory/AllocationListener.cc
         memory/MemoryAllocator.cc
         memory/ArrowMemoryPool.cc
-        ${SUBSTRAIT_PROTO_SRCS}
-        ${GLUTEN_PROTO_SRCS}
-        compute/ProtobufUtils.cc
+        memory/ColumnarBatch.cc
         operators/writer/ArrowWriter.cc
         shuffle/reader.cc
         shuffle/ShuffleWriter.cc
@@ -203,7 +205,6 @@ set(SPARK_COLUMNAR_PLUGIN_SRCS
         shuffle/LocalPartitionWriter.cc
         shuffle/rss/RemotePartitionWriter.cc
         shuffle/rss/CelebornPartitionWriter.cc
-        memory/ColumnarBatch.cc
         utils/Compression.cc
         utils/StringUtil.cc)
 

cpp/core/compute/Backend.h

@@ -75,7 +75,7 @@ class Backend : public std::enable_shared_from_this<Backend> {
   }
 
   virtual MemoryManager*
-  getMemoryManager(const std::string& name, std::shared_ptr<MemoryAllocator>, std::shared_ptr<AllocationListener>) {
+  getMemoryManager(const std::string& name, std::shared_ptr<MemoryAllocator>, std::unique_ptr<AllocationListener>) {
     throw GlutenException("getMemoryManager not implemented.");
   }
 

cpp/core/jni/JniCommon.h

@@ -24,7 +24,7 @@
 
 #include "compute/ProtobufUtils.h"
 #include "config/GlutenConfig.h"
-#include "memory/ArrowMemoryPool.h"
+#include "memory/AllocationListener.h"
 #include "utils/compression.h"
 #include "utils/exception.h"
 
@@ -166,27 +166,29 @@ class SparkAllocationListener final : public gluten::AllocationListener {
  public:
   SparkAllocationListener(
       JavaVM* vm,
-      jobject javaListener,
-      jmethodID javaReserveMethod,
-      jmethodID javaUnreserveMethod,
+      jobject jListenerLocalRef,
+      jmethodID jReserveMethod,
+      jmethodID jUnreserveMethod,
       int64_t blockSize)
-      : vm_(vm),
-        javaReserveMethod_(javaReserveMethod),
-        javaUnreserveMethod_(javaUnreserveMethod),
-        blockSize_(blockSize) {
+      : vm_(vm), jReserveMethod_(jReserveMethod), jUnreserveMethod_(jUnreserveMethod), blockSize_(blockSize) {
     JNIEnv* env;
     attachCurrentThreadAsDaemonOrThrow(vm_, &env);
-    javaListener_ = env->NewGlobalRef(javaListener);
+    jListenerGlobalRef_ = env->NewGlobalRef(jListenerLocalRef);
   }
 
+  SparkAllocationListener(const SparkAllocationListener&) = delete;
+  SparkAllocationListener(SparkAllocationListener&&) = delete;
+  SparkAllocationListener& operator=(const SparkAllocationListener&) = delete;
+  SparkAllocationListener& operator=(SparkAllocationListener&&) = delete;
+
   ~SparkAllocationListener() override {
     JNIEnv* env;
     if (vm_->GetEnv(reinterpret_cast<void**>(&env), jniVersion) != JNI_OK) {
       std::cerr << "SparkAllocationListener#~SparkAllocationListener(): "
                 << "JNIEnv was not attached to current thread" << std::endl;
       return;
     }
-    env->DeleteGlobalRef(javaListener_);
+    env->DeleteGlobalRef(jListenerGlobalRef_);
   }
 
   void allocationChanged(int64_t size) override {
@@ -213,25 +215,25 @@ class SparkAllocationListener final : public gluten::AllocationListener {
   }
 
   void updateReservation(int64_t diff) {
-    JNIEnv* env;
-    attachCurrentThreadAsDaemonOrThrow(vm_, &env);
     int64_t granted = reserve(diff);
     if (granted == 0) {
       return;
     }
+    JNIEnv* env;
+    attachCurrentThreadAsDaemonOrThrow(vm_, &env);
     if (granted < 0) {
-      env->CallLongMethod(javaListener_, javaUnreserveMethod_, -granted);
+      env->CallLongMethod(jListenerGlobalRef_, jUnreserveMethod_, -granted);
       checkException(env);
       return;
     }
-    env->CallLongMethod(javaListener_, javaReserveMethod_, granted);
+    env->CallLongMethod(jListenerGlobalRef_, jReserveMethod_, granted);
     checkException(env);
   }
 
   JavaVM* vm_;
-  jobject javaListener_;
-  jmethodID javaReserveMethod_;
-  jmethodID javaUnreserveMethod_;
+  jobject jListenerGlobalRef_;
+  jmethodID jReserveMethod_;
+  jmethodID jUnreserveMethod_;
   int64_t blockSize_;
   int64_t blocksReserved_ = 0L;
   int64_t bytesReserved_ = 0L;

cpp/core/jni/JniWrapper.cc

@@ -29,6 +29,7 @@
 #include "operators/writer/Datasource.h"
 
 #include <arrow/c/bridge.h>
+#include "memory/AllocationListener.h"
 #include "operators/serializer/ColumnarBatchSerializer.h"
 #include "shuffle/LocalPartitionWriter.h"
 #include "shuffle/PartitionWriterCreator.h"
@@ -1135,11 +1136,11 @@ JNIEXPORT jlong JNICALL Java_io_glutenproject_memory_nmm_NativeMemoryManager_cre
     throw gluten::GlutenException("Allocator does not exist or has been closed");
   }
 
+  auto listener = std::make_unique<SparkAllocationListener>(
+      vm, jlistener, reserveMemoryMethod, unreserveMemoryMethod, reservationBlockSize);
   auto name = jStringToCString(env, jname);
   auto backend = createBackend();
-  auto listener = std::make_shared<SparkAllocationListener>(
-      vm, jlistener, reserveMemoryMethod, unreserveMemoryMethod, reservationBlockSize);
-  auto manager = backend->getMemoryManager(name, *allocator, listener);
+  auto manager = backend->getMemoryManager(name, *allocator, std::move(listener));
   return reinterpret_cast<jlong>(manager);
   JNI_METHOD_END(-1L)
 }

cpp/core/memory/AllocationListener.cc

@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "AllocationListener.h"
+
+namespace gluten {
+
+class NoopAllocationListener : public gluten::AllocationListener {
+ public:
+  void allocationChanged(int64_t diff) override {
+    // no-op
+  }
+};
+
+std::unique_ptr<AllocationListener> AllocationListener::noop() {
+  return std::make_unique<NoopAllocationListener>();
+};
+
+} // namespace gluten

cpp/core/memory/AllocationListener.h

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <memory>
+
+namespace gluten {
+
+class AllocationListener {
+ public:
+  static std::unique_ptr<AllocationListener> noop();
+
+  virtual ~AllocationListener() = default;
+
+  // Value of diff can be either positive or negative
+  virtual void allocationChanged(int64_t diff) = 0;
+
+ protected:
+  AllocationListener() = default;
+};
+
+} // namespace gluten

cpp/core/memory/ArrowMemoryPool.h

@@ -17,6 +17,8 @@
 
 #pragma once
 
+#include "arrow/memory_pool.h"
+
 #include "MemoryAllocator.h"
 
 namespace gluten {

cpp/core/memory/MemoryAllocator.cc

@@ -104,14 +104,6 @@ int64_t ListenableMemoryAllocator::getBytes() const {
   return bytes_;
 }
 
-MemoryAllocator* ListenableMemoryAllocator::delegatedAllocator() {
-  return delegated_;
-}
-
-std::shared_ptr<AllocationListener> ListenableMemoryAllocator::listener() {
-  return listener_;
-}
-
 bool StdMemoryAllocator::allocate(int64_t size, void** out) {
   *out = std::malloc(size);
   bytes_ += size;
@@ -168,18 +160,6 @@ int64_t StdMemoryAllocator::getBytes() const {
   return bytes_;
 }
 
-class NoopAllocationListener : public gluten::AllocationListener {
- public:
-  void allocationChanged(int64_t diff) override {
-    // no-op
-  }
-};
-
-AllocationListener* AllocationListener::noop() {
-  static AllocationListener* instance = new NoopAllocationListener();
-  return instance;
-};
-
 std::shared_ptr<MemoryAllocator> defaultMemoryAllocator() {
 #if defined(GLUTEN_ENABLE_HBM)
   static std::shared_ptr<MemoryAllocator> alloc = HbwMemoryAllocator::newInstance();

cpp/core/memory/MemoryAllocator.h

@@ -23,7 +23,7 @@
 #include <memory>
 #include <utility>
 
-#include "arrow/memory_pool.h"
+#include "memory/AllocationListener.h"
 
 namespace gluten {
 
@@ -43,23 +43,10 @@ class MemoryAllocator {
   virtual int64_t getBytes() const = 0;
 };
 
-class AllocationListener {
- public:
-  static AllocationListener* noop();
-
-  virtual ~AllocationListener() = default;
-
-  // Value of diff can be either positive or negative
-  virtual void allocationChanged(int64_t diff) = 0;
-
- protected:
-  AllocationListener() = default;
-};
-
 class ListenableMemoryAllocator final : public MemoryAllocator {
  public:
-  explicit ListenableMemoryAllocator(MemoryAllocator* delegated, std::shared_ptr<AllocationListener> listener)
-      : delegated_(delegated), listener_(std::move(listener)) {}
+  explicit ListenableMemoryAllocator(MemoryAllocator* delegated, AllocationListener* listener)
+      : delegated_(delegated), listener_(listener) {}
 
  public:
   bool allocate(int64_t size, void** out) override;
@@ -74,15 +61,11 @@ class ListenableMemoryAllocator final : public MemoryAllocator {
 
   bool free(void* p, int64_t size) override;
 
-  MemoryAllocator* delegatedAllocator();
-
-  std::shared_ptr<AllocationListener> listener();
-
   int64_t getBytes() const override;
 
  private:
   MemoryAllocator* delegated_;
-  std::shared_ptr<AllocationListener> listener_;
+  AllocationListener* listener_;
   std::atomic_int64_t bytes_{0};
 };
 

cpp/velox/compute/VeloxBackend.h

@@ -59,8 +59,8 @@ class VeloxBackend final : public Backend {
   MemoryManager* getMemoryManager(
       const std::string& name,
       std::shared_ptr<MemoryAllocator> allocator,
-      std::shared_ptr<AllocationListener> listener) override {
-    return new VeloxMemoryManager(name, allocator, listener);
+      std::unique_ptr<AllocationListener> listener) override {
+    return new VeloxMemoryManager(name, allocator, std::move(listener));
   }
 
   // FIXME This is not thread-safe?

cpp/velox/memory/VeloxMemoryManager.cc

@@ -195,9 +195,9 @@ velox::memory::IMemoryManager::Options VeloxMemoryManager::getOptions(
 VeloxMemoryManager::VeloxMemoryManager(
     std::string name,
     std::shared_ptr<MemoryAllocator> allocator,
-    std::shared_ptr<AllocationListener> listener)
-    : MemoryManager(), name_(name), listener_(listener) {
-  glutenAlloc_ = std::make_unique<ListenableMemoryAllocator>(allocator.get(), listener_);
+    std::unique_ptr<AllocationListener> listener)
+    : MemoryManager(), name_(name), listener_(std::move(listener)) {
+  glutenAlloc_ = std::make_unique<ListenableMemoryAllocator>(allocator.get(), listener_.get());
   arrowPool_ = std::make_shared<ArrowMemoryPool>(glutenAlloc_.get());
 
   auto options = getOptions(allocator);

cpp/velox/memory/VeloxMemoryManager.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "memory/AllocationListener.h"
 #include "memory/MemoryAllocator.h"
 #include "memory/MemoryManager.h"
 #include "velox/common/memory/Memory.h"
@@ -29,7 +30,12 @@ class VeloxMemoryManager final : public MemoryManager {
   explicit VeloxMemoryManager(
       std::string name,
       std::shared_ptr<MemoryAllocator> allocator,
-      std::shared_ptr<AllocationListener> listener);
+      std::unique_ptr<AllocationListener> listener);
+
+  VeloxMemoryManager(const VeloxMemoryManager&) = delete;
+  VeloxMemoryManager(VeloxMemoryManager&&) = delete;
+  VeloxMemoryManager& operator=(const VeloxMemoryManager&) = delete;
+  VeloxMemoryManager& operator=(VeloxMemoryManager&&) = delete;
 
   std::shared_ptr<facebook::velox::memory::MemoryPool> getAggregateMemoryPool() const {
     return veloxAggregatePool_;
@@ -56,7 +62,7 @@ class VeloxMemoryManager final : public MemoryManager {
 
   // This is a listenable allocator used for arrow.
   std::unique_ptr<MemoryAllocator> glutenAlloc_;
-  std::shared_ptr<AllocationListener> listener_;
+  std::unique_ptr<AllocationListener> listener_;
   std::shared_ptr<arrow::MemoryPool> arrowPool_;
 
   std::unique_ptr<facebook::velox::memory::MemoryManager> veloxMemoryManager_;
@@ -67,9 +73,7 @@ class VeloxMemoryManager final : public MemoryManager {
 /// Not tracked by Spark and should only used in test or validation.
 inline std::shared_ptr<gluten::VeloxMemoryManager> getDefaultMemoryManager() {
   static auto memoryManager = std::make_shared<gluten::VeloxMemoryManager>(
-      "test",
-      gluten::defaultMemoryAllocator(),
-      std::shared_ptr<gluten::AllocationListener>(gluten::AllocationListener::noop()));
+      "test", gluten::defaultMemoryAllocator(), gluten::AllocationListener::noop());
   return memoryManager;
 }