[Enhancement] Optimize memory tracker (#49841)

## Why I'm doing: Currently, memory statistics consume a lot of performance because they are all calculated each time. ## What I'm doing: Get some samples to estimate the manager's memory, even though it is inaccurate, but it's helpful to solve memory leak bugs. Signed-off-by: gengjun-git <[email protected]>
StarRocks · Sep 4, 2024 · f0cb5e9 · f0cb5e9
1 parent a2db4b6
commit f0cb5e9
Show file tree

Hide file tree

Showing 43 changed files with 674 additions and 159 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/backup/BackupHandler.java b/fe/fe-core/src/main/java/com/starrocks/backup/BackupHandler.java
@@ -91,6 +91,7 @@
 import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -744,6 +745,10 @@ public void removeOldJobs() throws DdlException {
     public Map<String, Long> estimateCount() {
         return ImmutableMap.of("BackupOrRestoreJob", (long) dbIdToBackupOrRestoreJob.size());
     }
-}
-
 
+    @Override
+    public List<Pair<List<Object>, Long>> getSamples() {
+        List<Object> jobSamples = new ArrayList<>(dbIdToBackupOrRestoreJob.values());
+        return Lists.newArrayList(Pair.create(jobSamples, (long) dbIdToBackupOrRestoreJob.size()));
+    }
+}
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/Database.java b/fe/fe-core/src/main/java/com/starrocks/catalog/Database.java
@@ -72,6 +72,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -730,4 +731,21 @@ public void setExist(boolean exist) {
     public boolean getExist() {
         return exist;
     }
+
+    public List<PhysicalPartition> getPartitionSamples() {
+        return this.idToTable.values()
+                .stream()
+                .filter(table -> table instanceof OlapTable)
+                .map(table -> ((OlapTable) table).getPartitionSample())
+                .filter(Objects::nonNull)
+                .collect(Collectors.toList());
+    }
+
+    public int getOlapPartitionsCount() {
+        return this.idToTable.values()
+                .stream()
+                .filter(table -> table instanceof OlapTable)
+                .mapToInt(table -> ((OlapTable) table).getPartitionsCount())
+                .sum();
+    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/OlapTable.java b/fe/fe-core/src/main/java/com/starrocks/catalog/OlapTable.java
@@ -3529,4 +3529,15 @@ public void unlockCreatePartition(String partitionName) {
         }
     }
 
+    public int getPartitionsCount() {
+        return physicalPartitionIdToPartitionId.size();
+    }
+
+    public PhysicalPartition getPartitionSample() {
+        if (!idToPartition.isEmpty()) {
+            return idToPartition.values().iterator().next().getSubPartitions().iterator().next();
+        } else {
+            return null;
+        }
+    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/TabletInvertedIndex.java b/fe/fe-core/src/main/java/com/starrocks/catalog/TabletInvertedIndex.java
@@ -888,5 +888,24 @@ public Map<String, Long> estimateCount() {
                                "TabletCount", getTabletCount(),
                                "ReplicateCount", getReplicaCount());
     }
-}
 
+    @Override
+    public List<Pair<List<Object>, Long>> getSamples() {
+        readLock();
+        try {
+            List<Object> tabletMetaSamples = tabletMetaMap.values()
+                    .stream()
+                    .limit(1)
+                    .collect(Collectors.toList());
+
+            List<Object> longSamples = Lists.newArrayList(0L);
+            long longSize = tabletMetaMap.size() + replicaToTabletMap.size() * 2L + forceDeleteTablets.size() * 4L
+                    + replicaMetaTable.size() * 2L + backingReplicaMetaTable.size() * 2L;
+
+            return Lists.newArrayList(Pair.create(tabletMetaSamples, (long) tabletMetaMap.size()),
+                    Pair.create(longSamples, longSize));
+        } finally {
+            readUnlock();
+        }
+    }
+}
diff --git a/fe/fe-core/src/main/java/com/starrocks/common/Config.java b/fe/fe-core/src/main/java/com/starrocks/common/Config.java
@@ -1340,7 +1340,7 @@ public class Config extends ConfigBase {
      * If set to true, memory tracker feature will open
      */
     @ConfField(mutable = true)
-    public static boolean memory_tracker_enable = false;
+    public static boolean memory_tracker_enable = true;
 
     /**
      * Decide how often to track the memory usage of the FE process

diff --git a/fe/fe-core/src/main/java/com/starrocks/common/util/ProfileManager.java b/fe/fe-core/src/main/java/com/starrocks/common/util/ProfileManager.java
@@ -39,10 +39,10 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.starrocks.common.Config;
+import com.starrocks.common.Pair;
 import com.starrocks.memory.MemoryTrackable;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.apache.spark.util.SizeEstimator;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -53,6 +53,7 @@
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock.ReadLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock.WriteLock;
+import java.util.stream.Collectors;
 
 /*
  * if you want to visit the atrribute(such as queryID,defaultDb)
@@ -79,20 +80,12 @@ public class ProfileManager implements MemoryTrackable {
     public static final String VARIABLES = "Variables";
     public static final String PROFILE_COLLECT_TIME = "Collect Profile Time";
 
+    private static final int MEMORY_PROFILE_SAMPLES = 10;
+
     public static final ArrayList<String> PROFILE_HEADERS = new ArrayList<>(
             Arrays.asList(QUERY_ID, USER, DEFAULT_DB, SQL_STATEMENT, QUERY_TYPE,
                     START_TIME, END_TIME, TOTAL_TIME, QUERY_STATE));
 
-    @Override
-    public long estimateSize() {
-        return SizeEstimator.estimate(profileMap) + SizeEstimator.estimate(loadProfileMap);
-    }
-
-    @Override
-    public Map<String, Long> estimateCount() {
-        return ImmutableMap.of("QueryProfile", (long) profileMap.size(),
-                               "LoadProfile", (long) loadProfileMap.size());
-    }
 
     public static class ProfileElement {
         public Map<String, String> infoStrings = Maps.newHashMap();
@@ -298,19 +291,27 @@ public List<ProfileElement> getAllProfileElements() {
         return result;
     }
 
-    public long getQueryProfileCount() {
-        readLock.lock();
-        try {
-            return profileMap.size();
-        } finally {
-            readLock.unlock();
-        }
+    @Override
+    public Map<String, Long> estimateCount() {
+        return ImmutableMap.of("QueryProfile", (long) profileMap.size(),
+                "LoadProfile", (long) loadProfileMap.size());
     }
 
-    public long getLoadProfileCount() {
+    @Override
+    public List<Pair<List<Object>, Long>> getSamples() {
         readLock.lock();
         try {
-            return loadProfileMap.size();
+            List<Object> profileSamples = profileMap.values()
+                    .stream()
+                    .limit(MEMORY_PROFILE_SAMPLES)
+                    .collect(Collectors.toList());
+            List<Object> loadProfileSamples = loadProfileMap.values()
+                    .stream()
+                    .limit(MEMORY_PROFILE_SAMPLES)
+                    .collect(Collectors.toList());
+
+            return Lists.newArrayList(Pair.create(profileSamples, (long) profileMap.size()),
+                    Pair.create(loadProfileSamples, (long) loadProfileMap.size()));
         } finally {
             readLock.unlock();
         }

diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/CatalogConnector.java b/fe/fe-core/src/main/java/com/starrocks/connector/CatalogConnector.java
@@ -14,9 +14,11 @@
 
 package com.starrocks.connector;
 
+import com.starrocks.common.Pair;
 import com.starrocks.connector.informationschema.InformationSchemaConnector;
 import com.starrocks.connector.metadata.TableMetaConnector;
 
+import java.util.List;
 import java.util.Map;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -56,16 +58,20 @@ public boolean supportMemoryTrack() {
     }
 
     @Override
-    public long estimateSize() {
-        return normalConnector.estimateSize();
+    public Map<String, Long> estimateCount() {
+        return normalConnector.estimateCount();
     }
 
     @Override
-    public Map<String, Long> estimateCount() {
-        return normalConnector.estimateCount();
+    public List<Pair<List<Object>, Long>> getSamples() {
+        return normalConnector.getSamples();
     }
 
     public String normalConnectorClassName() {
-        return normalConnector.getClass().getSimpleName();
+        if (normalConnector instanceof LazyConnector) {
+            return ((LazyConnector) normalConnector).getRealConnectorClassName();
+        } else {
+            return normalConnector.getClass().getSimpleName();
+        }
     }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/Connector.java b/fe/fe-core/src/main/java/com/starrocks/connector/Connector.java
@@ -14,9 +14,15 @@
 
 package com.starrocks.connector;
 
+import com.starrocks.common.Pair;
 import com.starrocks.connector.config.ConnectorConfig;
 import com.starrocks.memory.MemoryTrackable;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 public interface Connector extends MemoryTrackable {
     /**
      * Get the connector meta of connector
@@ -44,4 +50,12 @@ default void bindConfig(ConnectorConfig config) {
     default boolean supportMemoryTrack() {
         return false;
     }
+
+    default Map<String, Long> estimateCount() {
+        return new HashMap<>();
+    }
+
+    default List<Pair<List<Object>, Long>> getSamples() {
+        return new ArrayList<>();
+    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/LazyConnector.java b/fe/fe-core/src/main/java/com/starrocks/connector/LazyConnector.java
@@ -14,10 +14,14 @@
 
 package com.starrocks.connector;
 
+import com.starrocks.common.Pair;
 import com.starrocks.connector.exception.StarRocksConnectorException;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.List;
+import java.util.Map;
+
 public class LazyConnector implements Connector {
     private static final Logger LOG = LogManager.getLogger(LazyConnector.class);
     private Connector delegate;
@@ -56,4 +60,30 @@ public void shutdown() {
             }
         }
     }
+
+    @Override
+    public boolean supportMemoryTrack() {
+        initIfNeeded();
+        return delegate.supportMemoryTrack();
+    }
+
+    @Override
+    public Map<String, Long> estimateCount() {
+        initIfNeeded();
+        return delegate.estimateCount();
+    }
+
+    @Override
+    public List<Pair<List<Object>, Long>> getSamples() {
+        initIfNeeded();
+        return delegate.getSamples();
+    }
+
+    public String getRealConnectorClassName() {
+        if (delegate != null) {
+            return delegate.getClass().getSimpleName();
+        } else {
+            return LazyConnector.class.getSimpleName();
+        }
+    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/delta/CachingDeltaLakeMetastore.java b/fe/fe-core/src/main/java/com/starrocks/connector/delta/CachingDeltaLakeMetastore.java
@@ -19,6 +19,7 @@
 import com.starrocks.catalog.Database;
 import com.starrocks.catalog.Table;
 import com.starrocks.common.Config;
+import com.starrocks.common.Pair;
 import com.starrocks.connector.DatabaseTableName;
 import com.starrocks.connector.exception.StarRocksConnectorException;
 import com.starrocks.connector.metastore.CachingMetastore;
@@ -37,11 +38,13 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.stream.Collectors;
 
 import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
 
 public class CachingDeltaLakeMetastore extends CachingMetastore implements IDeltaLakeMetastore {
     private static final Logger LOG = LogManager.getLogger(CachingDeltaLakeMetastore.class);
+    private static final int MEMORY_META_SAMPLES = 10;
 
     public final IDeltaLakeMetastore delegate;
     private final Map<DatabaseTableName, Long> lastAccessTimeMap;
@@ -180,4 +183,29 @@ public void invalidateAll() {
             ((DeltaLakeMetastore) delegate).invalidateAll();
         }
     }
+
+    @Override
+    public List<Pair<List<Object>, Long>> getSamples() {
+        List<Object> dbSamples = databaseCache.asMap().values()
+                .stream()
+                .limit(MEMORY_META_SAMPLES)
+                .collect(Collectors.toList());
+        List<Object> tableSamples = tableCache.asMap().values()
+                .stream()
+                .limit(MEMORY_META_SAMPLES)
+                .collect(Collectors.toList());
+
+        List<Pair<List<Object>, Long>> samples = delegate.getSamples();
+        samples.add(Pair.create(dbSamples, databaseCache.size()));
+        samples.add(Pair.create(tableSamples, tableCache.size()));
+        return samples;
+    }
+
+    @Override
+    public Map<String, Long> estimateCount() {
+        Map<String, Long> delegateCount = Maps.newHashMap(delegate.estimateCount());
+        delegateCount.put("databaseCache", databaseCache.size());
+        delegateCount.put("tableCache", tableCache.size());
+        return delegateCount;
+    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/delta/DeltaLakeConnector.java b/fe/fe-core/src/main/java/com/starrocks/connector/delta/DeltaLakeConnector.java
@@ -14,6 +14,7 @@
 
 package com.starrocks.connector.delta;
 
+import com.starrocks.common.Pair;
 import com.starrocks.connector.Connector;
 import com.starrocks.connector.ConnectorContext;
 import com.starrocks.connector.ConnectorMetadata;
@@ -24,6 +25,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
@@ -81,4 +83,19 @@ public void onCreate() {
         updateProcessor.ifPresent(processor -> GlobalStateMgr.getCurrentState().getConnectorTableMetadataProcessor()
                         .registerCacheUpdateProcessor(catalogName, updateProcessor.get()));
     }
+
+    @Override
+    public boolean supportMemoryTrack() {
+        return metastore != null;
+    }
+
+    @Override
+    public List<Pair<List<Object>, Long>> getSamples() {
+        return metastore.getSamples();
+    }
+
+    @Override
+    public Map<String, Long> estimateCount() {
+        return metastore.estimateCount();
+    }
 }