[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

@@ -47,34 +47,36 @@ public class MySqlClusterAggregator {
   }
 
   /**
-   * 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 the valid aggregated storage stats for all replicas of each partition.
-   * @param statsWrappers A map from instance name to {@link HostAccountStorageStatsWrapper}.
+   * 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(
-      Map<String, HostAccountStorageStatsWrapper> statsWrappers) throws IOException {
+      AccountStorageStatsIterator 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<>();
-    for (Map.Entry<String, HostAccountStorageStatsWrapper> statsWrapperEntry : statsWrappers.entrySet()) {
-      if (statsWrapperEntry.getValue() == null) {
+
+    while (accountStorageStatsIterator.hasNext()) {
+      Pair<String, HostAccountStorageStatsWrapper> statsWrapperEntry = accountStorageStatsIterator.next();
+      if (statsWrapperEntry.getSecond() == null) {
         continue;
       }
-      String instanceName = statsWrapperEntry.getKey();
-      HostAccountStorageStatsWrapper hostAccountStorageStatsWrapper = statsWrapperEntry.getValue();
+      String hostname = 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);
+          hostAccountStorageStatsWrapper.getHeader().getTimestamp(), hostname);
     }
     if (logger.isTraceEnabled()) {
       logger.trace("Combined raw HostAccountStorageStats {}",
@@ -95,28 +97,29 @@ Pair<AggregatedAccountStorageStats, AggregatedAccountStorageStats> aggregateHost
   }
 
   /**
-   * Aggregate all {@link HostPartitionClassStorageStatsWrapper} to generate two {@link AggregatedPartitionClassStorageStats}s. First
-   * {@link AggregatedPartitionClassStorageStats} is the sum of all {@link HostPartitionClassStorageStatsWrapper}s. The second
-   * {@link AggregatedPartitionClassStorageStats} is the valid aggregated storage stats for all replicas of each partition.
-   * @param statsWrappers A map from instance name to {@link HostPartitionClassStorageStatsWrapper}.
+   * 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(
-      Map<String, HostPartitionClassStorageStatsWrapper> statsWrappers) throws IOException {
+      PartitionClassStorageStatsIterator 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<>();
 
-    for (Map.Entry<String, HostPartitionClassStorageStatsWrapper> statsWrapperEntry : statsWrappers.entrySet()) {
-      if (statsWrapperEntry.getValue() == null) {
+    while (partitionClassStorageStatsIterator.hasNext()) {
+      Pair<String, HostPartitionClassStorageStatsWrapper> statsWrapperEntry = partitionClassStorageStatsIterator.next();
+      if (statsWrapperEntry.getSecond() == null) {
         continue;
       }
-      String instanceName = statsWrapperEntry.getKey();
-      HostPartitionClassStorageStatsWrapper hostPartitionClassStorageStatsWrapper = statsWrapperEntry.getValue();
+      String hostname = statsWrapperEntry.getFirst();
+      HostPartitionClassStorageStatsWrapper hostPartitionClassStorageStatsWrapper = statsWrapperEntry.getSecond();
       HostPartitionClassStorageStats hostPartitionClassStorageStats = hostPartitionClassStorageStatsWrapper.getStats();
       HostPartitionClassStorageStats hostPartitionClassStorageStatsCopy1 =
           new HostPartitionClassStorageStats(hostPartitionClassStorageStats);
@@ -126,7 +129,7 @@ Pair<AggregatedPartitionClassStorageStats, AggregatedPartitionClassStorageStats>
           hostPartitionClassStorageStatsCopy1.getStorageStats());
       selectRawHostPartitionClassStorageStatsMap(selectedHostPartitionClassStorageStatsMap,
           hostPartitionClassStorageStatsCopy2.getStorageStats(), partitionTimestampMap, partitionPhysicalStorageMap,
-          hostPartitionClassStorageStatsWrapper.getHeader().getTimestamp(), instanceName);
+          hostPartitionClassStorageStatsWrapper.getHeader().getTimestamp(), hostname);
     }
 
     if (logger.isTraceEnabled()) {

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

@@ -18,8 +18,6 @@
 import com.github.ambry.accountstats.AccountStatsStore;
 import com.github.ambry.commons.Callback;
 import com.github.ambry.config.ClusterMapConfig;
-import com.github.ambry.server.HostAccountStorageStatsWrapper;
-import com.github.ambry.server.HostPartitionClassStorageStatsWrapper;
 import com.github.ambry.server.StatsReportType;
 import com.github.ambry.server.storagestats.AggregatedAccountStorageStats;
 import com.github.ambry.server.storagestats.AggregatedPartitionClassStorageStats;
@@ -32,10 +30,8 @@
 import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import org.apache.helix.HelixManager;
 import org.apache.helix.task.Task;
 import org.apache.helix.task.TaskResult;
@@ -123,24 +119,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());
         }
@@ -227,61 +217,6 @@ private void removeInvalidAggregatedPartitionClassStats(AggregatedPartitionClass
     }
   }
 
-  /**
-   * Fetch account storage stats report for each instance in {@code instanceNames}. Each instance name is probably a fully qualified
-   * hostname with port number like this [hostname_portnumber]. It returns a map whose key is the instanceName and the value
-   * is the {@link HostAccountStorageStatsWrapper} for each instance.
-   * @param instanceNames The list of instance names to fetch account StatsWrapper.
-   * @return A map of {@link HostAccountStorageStatsWrapper} for each instance name.
-   * @throws Exception
-   */
-  private Map<String, HostAccountStorageStatsWrapper> fetchAccountStorageStatsWrapperForInstances(
-      List<String> instanceNames) throws Exception {
-    Map<String, HostAccountStorageStatsWrapper> statsWrappers = new HashMap<>();
-    for (String instanceName : instanceNames) {
-      Pair<String, Integer> pair = getHostNameAndPort(instanceName);
-      statsWrappers.put(instanceName,
-          accountStatsStore.queryHostAccountStorageStatsByHost(pair.getFirst(), pair.getSecond()));
-    }
-    return statsWrappers;
-  }
-
-  /**
-   * Fetch partition class storage stats report for each instance in {@code instanceNames}. Each instance name is probably a fully qualified
-   * hostname with port number like this [hostname_portnumber]. It returns a map whose key is the instanceName and the value
-   * is the {@link HostPartitionClassStorageStatsWrapper} for each instance.
-   * @param instanceNames The list of instance names to fetch partition class StatsWrapper.
-   * @return A map of {@link HostPartitionClassStorageStatsWrapper} for each instance name.
-   * @throws Exception
-   */
-  private Map<String, HostPartitionClassStorageStatsWrapper> fetchPartitionClassStorageStatsWrapperForInstances(
-      List<String> instanceNames) throws Exception {
-    Map<String, HostPartitionClassStorageStatsWrapper> statsWrappers = new HashMap<>();
-    Map<String, Set<Integer>> partitionNameAndIds = accountStatsStore.queryPartitionNameAndIds();
-    for (String instanceName : instanceNames) {
-      Pair<String, Integer> pair = getHostNameAndPort(instanceName);
-      statsWrappers.put(instanceName,
-          accountStatsStore.queryHostPartitionClassStorageStatsByHost(pair.getFirst(), pair.getSecond(),
-              partitionNameAndIds));
-    }
-    return statsWrappers;
-  }
-
-  private Pair<String, Integer> getHostNameAndPort(String instanceName) {
-    String hostname = instanceName;
-    int port = clusterMapConfig.clusterMapPort;
-    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);
-  }
-
   @Override
   public void cancel() {
 

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);
+  }
+}