From febc227dfb41e9da13e01026b86ab8ed4b3cec33 Mon Sep 17 00:00:00 2001
From: gengjun-git <gengjun@starrocks.com>
Date: Fri, 22 Nov 2024 19:06:42 +0800
Subject: [PATCH] [Enhancement] Opt memory tracker for FE (#53055)

## Why I'm doing:
`SizeEstimator.estimate` will calculate the retained size of the object, that is the size of the param object and all the referenced objects. This will lead to repeated counting when the same object is referenced by many objects, and result in very high resource usage when the param object refers to network object or GlobalMetaStore.

## What I'm doing:
Change to `ClassLayout` to only calculate the shadow size of the object.

Signed-off-by: gengjun-git <gengjun@starrocks.com>
(cherry picked from commit 4f60313baa6c3cf1365996dd82a0b17a86585ab0)
---
 .../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;