From 7b672b97427c4776a3efc05882b4162b588de532 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 22 Nov 2024 12:41:15 +0000
Subject: [PATCH] [Enhancement] Opt memory tracker for FE (backport #53055)
 (#53143)

Co-authored-by: gengjun-git <gengjun@starrocks.com>
---
 .../com/starrocks/memory/MemoryTrackable.java   | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/fe/fe-core/src/main/java/com/starrocks/memory/MemoryTrackable.java b/fe/fe-core/src/main/java/com/starrocks/memory/MemoryTrackable.java
index e3470fec5f19d..af08bd4500f80 100644
--- a/fe/fe-core/src/main/java/com/starrocks/memory/MemoryTrackable.java
+++ b/fe/fe-core/src/main/java/com/starrocks/memory/MemoryTrackable.java
@@ -15,12 +15,19 @@
 package com.starrocks.memory;
 
 import com.starrocks.common.Pair;
-import org.apache.spark.util.SizeEstimator;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 public interface MemoryTrackable {
+    // The default implementation of estimateSize only calculate the shadow size of the object.
+    // The shadow size is the same for all instances of the specified class,
+    // so using CLASS_SIZE to cache the class's instance shadow size.
+    // the key is class name, the value is size
+    Map<String, Long> CLASS_SIZE = new ConcurrentHashMap<>();
+
     default long estimateSize() {
         List<Pair<List<Object>, Long>> samples = getSamples();
         long totalBytes = 0;
@@ -28,13 +35,19 @@ default long estimateSize() {
             List<Object> sampleObjects = pair.first;
             long size = pair.second;
             if (!sampleObjects.isEmpty()) {
-                totalBytes += (long) (((double) SizeEstimator.estimate(sampleObjects)) / sampleObjects.size() * size);
+                long sampleSize = sampleObjects.stream().mapToLong(this::getInstanceSize).sum();
+                totalBytes += (long) (((double) sampleSize) / sampleObjects.size() * size);
             }
         }
 
         return totalBytes;
     }
 
+    default long getInstanceSize(Object object) {
+        String className = object.getClass().getName();
+        return CLASS_SIZE.computeIfAbsent(className, s -> ClassLayout.parseInstance(object).instanceSize());
+    }
+
     Map<String, Long> estimateCount();
 
     // Samples for estimateSize() to calculate memory size;