[HelixTask] Reduce memory footprint to minimize GC pauses

ambry-clustermap/src/main/java/com/github/ambry/clustermap/AccountStorageStatsIterator.java

@@ -0,0 +1,59 @@
+/**
+ * Copyright 2023 LinkedIn Corp. All rights reserved.
+ *
+ * Licensed 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.
+ */
+package com.github.ambry.clustermap;
+
+import com.github.ambry.accountstats.AccountStatsStore;
+import com.github.ambry.config.ClusterMapConfig;
+import com.github.ambry.server.HostAccountStorageStatsWrapper;
+import com.github.ambry.utils.Pair;
+import java.util.Iterator;
+import java.util.List;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This is an iterator for account storage stats fetched from MySQL
+ */
+public class AccountStorageStatsIterator implements Iterator<Pair<String, HostAccountStorageStatsWrapper>> {
+
+  private static final Logger logger = LoggerFactory.getLogger(AccountStorageStatsIterator.class);
+  private final Iterator<String> instances;
+  private final AccountStatsStore accountStatsStore;
+  private final ClusterMapConfig clusterMapConfig;
+
+  public AccountStorageStatsIterator(List<String> instances, AccountStatsStore accountStatsStore, ClusterMapConfig clusterMapConfig) {
+    this.instances = instances.iterator();
+    this.accountStatsStore = accountStatsStore;
+    this.clusterMapConfig = clusterMapConfig;
+  }
+
+  @Override
+  public boolean hasNext() {
+    return instances.hasNext();
+  }
+
+  @Override
+  public Pair<String, HostAccountStorageStatsWrapper> next() {
+    String hostname = instances.next();
+    try {
+      Pair<String, Integer> hostNameAndPort = TaskUtils.getHostNameAndPort(hostname, clusterMapConfig.clusterMapPort);
+      return new Pair<>(hostname,
+          accountStatsStore.queryHostAccountStorageStatsByHost(hostNameAndPort.getFirst(), hostNameAndPort.getSecond()));
+    } catch (Exception e) {
+      logger.error("Failed to get account storage stats for {}", hostname);
+      throw new RuntimeException(e);
+    }
+  }
+}

ambry-clustermap/src/main/java/com/github/ambry/clustermap/MySqlClusterAggregator.java

@@ -24,6 +24,7 @@
 import com.github.ambry.utils.Pair;
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.slf4j.Logger;
@@ -46,6 +47,57 @@ public class MySqlClusterAggregator {
     relevantTimePeriodInMs = TimeUnit.MINUTES.toMillis(relevantTimePeriodInMinutes);
   }
 
+  /**
+   * Aggregate all {@link HostAccountStorageStatsWrapper} to generate two {@link AggregatedAccountStorageStats}s.
+   * First {@link AggregatedAccountStorageStats} is the sum of all {@link HostAccountStorageStatsWrapper}s.
+   * The second {@link AggregatedAccountStorageStats} is an aggregated storage stats for all replicas of each partition.
+   * @param accountStorageStatsIterator An iterator for AccountStorageStatsIterator
+   * @return A {@link Pair} of {@link AggregatedAccountStorageStats}.
+   * @throws IOException
+   */
+  Pair<AggregatedAccountStorageStats, AggregatedAccountStorageStats> aggregateHostAccountStorageStatsWrappers(
+      Iterator accountStorageStatsIterator) throws IOException {
+
+    Map<Long, Map<Short, Map<Short, ContainerStorageStats>>> combinedHostAccountStorageStatsMap = new HashMap<>();
+    Map<Long, Map<Short, Map<Short, ContainerStorageStats>>> selectedHostAccountStorageStatsMap = new HashMap<>();
+    Map<Long, Long> partitionTimestampMap = new HashMap<>();
+    Map<Long, Long> partitionPhysicalStorageMap = new HashMap<>();
+
+    while (accountStorageStatsIterator.hasNext()) {
+      Pair<String, HostAccountStorageStatsWrapper> statsWrapperEntry =
+          (Pair<String, HostAccountStorageStatsWrapper>) accountStorageStatsIterator.next();
+      if (statsWrapperEntry.getSecond() == null) {
+        continue;
+      }
+      String instanceName = statsWrapperEntry.getFirst();
+      HostAccountStorageStatsWrapper hostAccountStorageStatsWrapper = statsWrapperEntry.getSecond();
+      HostAccountStorageStats hostAccountStorageStats = hostAccountStorageStatsWrapper.getStats();
+      HostAccountStorageStats hostAccountStorageStatsCopy1 = new HostAccountStorageStats(hostAccountStorageStats);
+      HostAccountStorageStats hostAccountStorageStatsCopy2 = new HostAccountStorageStats(hostAccountStorageStats);
+      combineRawHostAccountStorageStatsMap(combinedHostAccountStorageStatsMap,
+          hostAccountStorageStatsCopy1.getStorageStats());
+      selectRawHostAccountStorageStatsMap(selectedHostAccountStorageStatsMap,
+          hostAccountStorageStatsCopy2.getStorageStats(), partitionTimestampMap, partitionPhysicalStorageMap,
+          hostAccountStorageStatsWrapper.getHeader().getTimestamp(), instanceName);
+    }
+    if (logger.isTraceEnabled()) {
+      logger.trace("Combined raw HostAccountStorageStats {}",
+          mapper.writeValueAsString(combinedHostAccountStorageStatsMap));
+      logger.trace("Selected raw HostAccountStorageStats {}",
+          mapper.writeValueAsString(selectedHostAccountStorageStatsMap));
+    }
+
+    AggregatedAccountStorageStats combinedAggregated =
+        new AggregatedAccountStorageStats(aggregateHostAccountStorageStats(combinedHostAccountStorageStatsMap));
+    AggregatedAccountStorageStats selectedAggregated =
+        new AggregatedAccountStorageStats(aggregateHostAccountStorageStats(selectedHostAccountStorageStatsMap));
+    if (logger.isTraceEnabled()) {
+      logger.trace("Aggregated combined {}", mapper.writeValueAsString(combinedAggregated));
+      logger.trace("Aggregated selected {}", mapper.writeValueAsString(selectedAggregated));
+    }
+    return new Pair<>(combinedAggregated, selectedAggregated);
+  }
+
   /**
    * Aggregate all {@link HostAccountStorageStatsWrapper} to generate two {@link AggregatedAccountStorageStats}s. First
    * {@link AggregatedAccountStorageStats} is the sum of all {@link HostAccountStorageStatsWrapper}s. The second {@link AggregatedAccountStorageStats}
@@ -94,6 +146,61 @@ Pair<AggregatedAccountStorageStats, AggregatedAccountStorageStats> aggregateHost
     return new Pair<>(combinedAggregated, selectedAggregated);
   }
 
+  /**
+   * Aggregate all {@link HostPartitionClassStorageStatsWrapper} to generate two {@link AggregatedPartitionClassStorageStats}s.
+   * First {@link AggregatedPartitionClassStorageStats} is the sum of all {@link HostPartitionClassStorageStatsWrapper}s.
+   * Second {@link AggregatedPartitionClassStorageStats} is an aggregated storage stats for all replicas of each partition.
+   * @param partitionClassStorageStatsIterator An iterator for PartitionClassStorageStatsIterator
+   * @return A {@link Pair} of {@link AggregatedPartitionClassStorageStats}.
+   * @throws IOException
+   */
+  Pair<AggregatedPartitionClassStorageStats, AggregatedPartitionClassStorageStats> aggregateHostPartitionClassStorageStatsWrappers(
+      Iterator partitionClassStorageStatsIterator) throws IOException {
+    Map<String, Map<Long, Map<Short, Map<Short, ContainerStorageStats>>>> combinedHostPartitionClassStorageStatsMap =
+        new HashMap<>();
+    Map<String, Map<Long, Map<Short, Map<Short, ContainerStorageStats>>>> selectedHostPartitionClassStorageStatsMap =
+        new HashMap<>();
+    Map<Long, Long> partitionTimestampMap = new HashMap<>();
+    Map<Long, Long> partitionPhysicalStorageMap = new HashMap<>();
+
+    while (partitionClassStorageStatsIterator.hasNext()) {
+      Pair<String, HostPartitionClassStorageStatsWrapper> statsWrapperEntry =
+          (Pair<String, HostPartitionClassStorageStatsWrapper>) partitionClassStorageStatsIterator.next();
+      if (statsWrapperEntry.getSecond() == null) {
+        continue;
+      }
+      String hostname = statsWrapperEntry.getFirst();
+      HostPartitionClassStorageStatsWrapper hostPartitionClassStorageStatsWrapper = statsWrapperEntry.getSecond();
+      HostPartitionClassStorageStats hostPartitionClassStorageStats = hostPartitionClassStorageStatsWrapper.getStats();
+      HostPartitionClassStorageStats hostPartitionClassStorageStatsCopy1 =
+          new HostPartitionClassStorageStats(hostPartitionClassStorageStats);
+      HostPartitionClassStorageStats hostPartitionClassStorageStatsCopy2 =
+          new HostPartitionClassStorageStats(hostPartitionClassStorageStats);
+      combineRawHostPartitionClassStorageStatsMap(combinedHostPartitionClassStorageStatsMap,
+          hostPartitionClassStorageStatsCopy1.getStorageStats());
+      selectRawHostPartitionClassStorageStatsMap(selectedHostPartitionClassStorageStatsMap,
+          hostPartitionClassStorageStatsCopy2.getStorageStats(), partitionTimestampMap, partitionPhysicalStorageMap,
+          hostPartitionClassStorageStatsWrapper.getHeader().getTimestamp(), hostname);
+    }
+
+    if (logger.isTraceEnabled()) {
+      logger.trace("Combined raw HostPartitionClassStorageStats {}",
+          mapper.writeValueAsString(combinedHostPartitionClassStorageStatsMap));
+      logger.trace("Selected raw HostPartitionClassStorageStats {}",
+          mapper.writeValueAsString(selectedHostPartitionClassStorageStatsMap));
+    }
+
+    AggregatedPartitionClassStorageStats combinedAggregated = new AggregatedPartitionClassStorageStats(
+        aggregateHostPartitionClassStorageStats(combinedHostPartitionClassStorageStatsMap));
+    AggregatedPartitionClassStorageStats selectedAggregated = new AggregatedPartitionClassStorageStats(
+        aggregateHostPartitionClassStorageStats(selectedHostPartitionClassStorageStatsMap));
+    if (logger.isTraceEnabled()) {
+      logger.trace("Aggregated combined {}", mapper.writeValueAsString(combinedAggregated));
+      logger.trace("Aggregated selected {}", mapper.writeValueAsString(selectedAggregated));
+    }
+    return new Pair<>(combinedAggregated, selectedAggregated);
+  }
+
   /**
    * Aggregate all {@link HostPartitionClassStorageStatsWrapper} to generate two {@link AggregatedPartitionClassStorageStats}s. First
    * {@link AggregatedPartitionClassStorageStats} is the sum of all {@link HostPartitionClassStorageStatsWrapper}s. The second

ambry-clustermap/src/main/java/com/github/ambry/clustermap/MySqlReportAggregatorTask.java

@@ -33,6 +33,7 @@
 import java.time.format.DateTimeFormatter;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -123,24 +124,18 @@ public TaskResult run() {
     try {
       List<String> instanceNames = manager.getClusterManagmentTool().getInstancesInCluster(manager.getClusterName());
       if (statsReportType == StatsReportType.ACCOUNT_REPORT) {
-        Map<String, HostAccountStorageStatsWrapper> accountStatsWrappers =
-            fetchAccountStorageStatsWrapperForInstances(instanceNames);
-        fetchTimeMs.update(System.currentTimeMillis() - startTimeMs);
-        logger.info("Aggregating stats from " + accountStatsWrappers.size() + " hosts");
+        logger.info("Aggregating stats from " + instanceNames.size() + " hosts");
         Pair<AggregatedAccountStorageStats, AggregatedAccountStorageStats> results =
-            clusterAggregator.aggregateHostAccountStorageStatsWrappers(accountStatsWrappers);
+            clusterAggregator.aggregateHostAccountStorageStatsWrappers(new AccountStorageStatsIterator(instanceNames, accountStatsStore, clusterMapConfig));
         if (clusterMapConfig.clustermapEnableDeleteInvalidDataInMysqlAggregationTask) {
           removeInvalidAggregatedAccountAndContainerStats(results.getSecond());
         }
         accountStatsStore.storeAggregatedAccountStorageStats(results.getSecond());
         aggregatedAccountStorageStats = results.getFirst();
       } else if (statsReportType == StatsReportType.PARTITION_CLASS_REPORT) {
-        Map<String, HostPartitionClassStorageStatsWrapper> statsWrappers =
-            fetchPartitionClassStorageStatsWrapperForInstances(instanceNames);
-        fetchTimeMs.update(System.currentTimeMillis() - startTimeMs);
-        logger.info("Aggregating stats from " + statsWrappers.size() + " hosts");
+        logger.info("Aggregating stats from " + instanceNames.size() + " hosts");
         Pair<AggregatedPartitionClassStorageStats, AggregatedPartitionClassStorageStats> results =
-            clusterAggregator.aggregateHostPartitionClassStorageStatsWrappers(statsWrappers);
+            clusterAggregator.aggregateHostPartitionClassStorageStatsWrappers(new PartitionClassStorageStatsIterator(instanceNames, accountStatsStore, clusterMapConfig));
         if (clusterMapConfig.clustermapEnableDeleteInvalidDataInMysqlAggregationTask) {
           removeInvalidAggregatedPartitionClassStats(results.getSecond());
         }

ambry-clustermap/src/main/java/com/github/ambry/clustermap/PartitionClassStorageStatsIterator.java

@@ -0,0 +1,65 @@
+/**
+ * Copyright 2023 LinkedIn Corp. All rights reserved.
+ *
+ * Licensed 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.
+ */
+package com.github.ambry.clustermap;
+
+import com.github.ambry.accountstats.AccountStatsStore;
+import com.github.ambry.config.ClusterMapConfig;
+import com.github.ambry.server.HostPartitionClassStorageStatsWrapper;
+import com.github.ambry.utils.Pair;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This is an iterator for partition storage stats fetched from MySQL
+ */
+public class PartitionClassStorageStatsIterator implements
+                                                Iterator<Pair<String, HostPartitionClassStorageStatsWrapper>> {
+
+  private static final Logger logger = LoggerFactory.getLogger(PartitionClassStorageStatsIterator.class);
+  private final Iterator<String> instances;
+  private final Map<String, Set<Integer>> partitionNameAndIds;
+  private final AccountStatsStore accountStatsStore;
+  private final ClusterMapConfig clusterMapConfig;
+
+
+  public PartitionClassStorageStatsIterator(List<String> instances, AccountStatsStore accountStatsStore, ClusterMapConfig clusterMapConfig) throws Exception {
+    this.instances = instances.iterator();
+    this.partitionNameAndIds = accountStatsStore.queryPartitionNameAndIds();
+    this.accountStatsStore = accountStatsStore;
+    this.clusterMapConfig = clusterMapConfig;
+  }
+
+  @Override
+  public boolean hasNext() {
+    return instances.hasNext();
+  }
+
+  @Override
+  public Pair<String, HostPartitionClassStorageStatsWrapper> next() {
+    String hostname = instances.next();
+    try {
+      Pair<String, Integer> hostNameAndPort = TaskUtils.getHostNameAndPort(hostname, clusterMapConfig.clusterMapPort);
+      return new Pair<>(hostname,
+          accountStatsStore.queryHostPartitionClassStorageStatsByHost(
+              hostNameAndPort.getFirst(), hostNameAndPort.getSecond(), partitionNameAndIds));
+    } catch (Exception e) {
+      logger.error("Failed to get partition storage stats for {}", hostname);
+      throw new RuntimeException(e);
+    }
+  }
+}

ambry-clustermap/src/main/java/com/github/ambry/clustermap/TaskUtils.java

@@ -0,0 +1,41 @@
+/**
+ * Copyright 2023 LinkedIn Corp. All rights reserved.
+ *
+ * Licensed 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.
+ */
+package com.github.ambry.clustermap;
+
+import com.github.ambry.utils.Pair;
+
+public class TaskUtils {
+
+  /**
+   * Given an instance name of the form hostname_port, returns a pair of <hostname, port>.
+   * Given an instance name of the form hostname, returns a pair of <hostname, defaultPort>.
+   * @param instanceName Name of the instance machine
+   * @param defaultPort Default port used by machines
+   * @return Pair of <hostname, port>
+   */
+  protected static Pair<String, Integer> getHostNameAndPort(String instanceName, int defaultPort) {
+    String hostname = instanceName;
+    int port = defaultPort;
+    int ind = instanceName.lastIndexOf("_");
+    if (ind != -1) {
+      try {
+        port = Short.valueOf(instanceName.substring(ind + 1));
+        hostname = instanceName.substring(0, ind);
+      } catch (NumberFormatException e) {
+        // String after "_" is not a port number, then the hostname should be the instanceName
+      }
+    }
+    return new Pair<>(hostname, port);
+  }
+}