From cdac9a4d9d8dd544fcc84149bdc1a1721fade5fa Mon Sep 17 00:00:00 2001
From: Naganarasimha <naganarasimha_gr@apache.org>
Date: Fri, 9 Jun 2017 11:59:01 +0530
Subject: [PATCH] YARN-6682. Improve performance of AssignmentInformation
 datastructures. Contributed by Daryn Sharp.

---
 .../common/AssignmentInformation.java         | 61 ++++++++++---------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/AssignmentInformation.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/AssignmentInformation.java
index aad3bc71792..d3f2a73c810 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/AssignmentInformation.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/AssignmentInformation.java
@@ -19,22 +19,22 @@
 package org.apache.hadoop.yarn.server.resourcemanager.scheduler.common;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.Resource;
-import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer;
 
 @InterfaceAudience.Private
 @InterfaceStability.Unstable
 public class AssignmentInformation {
 
   public enum Operation {
-    ALLOCATION, RESERVATION
+    ALLOCATION, RESERVATION;
+    private static int SIZE = Operation.values().length;
+    static int size() {
+      return SIZE;
+    }
   }
 
   public static class AssignmentDetails {
@@ -47,23 +47,25 @@ public class AssignmentInformation {
     }
   }
 
-  private final Map<Operation, Integer> operationCounts;
-  private final Map<Operation, Resource> operationResources;
-  private final Map<Operation, List<AssignmentDetails>> operationDetails;
+  private final int[] operationCounts;
+  private final Resource[] operationResources;
+  private final List<AssignmentDetails>[] operationDetails;
 
+  @SuppressWarnings("unchecked")
   public AssignmentInformation() {
-    this.operationCounts = new HashMap<>();
-    this.operationResources = new HashMap<>();
-    this.operationDetails = new HashMap<>();
-    for (Operation op : Operation.values()) {
-      operationCounts.put(op, 0);
-      operationResources.put(op, Resource.newInstance(0, 0));
-      operationDetails.put(op, new ArrayList<AssignmentDetails>());
+    int numOps = Operation.size();
+    this.operationCounts = new int[numOps];
+    this.operationResources = new Resource[numOps];
+    this.operationDetails = new List[numOps];
+    for (int i=0; i < numOps; i++) {
+      operationCounts[i] = 0;
+      operationResources[i] = Resource.newInstance(0, 0);
+      operationDetails[i] = new ArrayList<AssignmentDetails>();
     }
   }
 
   public int getNumAllocations() {
-    return operationCounts.get(Operation.ALLOCATION);
+    return operationCounts[Operation.ALLOCATION.ordinal()];
   }
 
   public void incrAllocations() {
@@ -75,7 +77,7 @@ public class AssignmentInformation {
   }
 
   public int getNumReservations() {
-    return operationCounts.get(Operation.RESERVATION);
+    return operationCounts[Operation.RESERVATION.ordinal()];
   }
 
   public void incrReservations() {
@@ -87,20 +89,20 @@ public class AssignmentInformation {
   }
 
   private void increment(Operation op, int by) {
-    operationCounts.put(op, operationCounts.get(op) + by);
+    operationCounts[op.ordinal()] += by;
   }
 
   public Resource getAllocated() {
-    return operationResources.get(Operation.ALLOCATION);
+    return operationResources[Operation.ALLOCATION.ordinal()];
   }
 
   public Resource getReserved() {
-    return operationResources.get(Operation.RESERVATION);
+    return operationResources[Operation.RESERVATION.ordinal()];
   }
 
   private void addAssignmentDetails(Operation op, ContainerId containerId,
       String queue) {
-    operationDetails.get(op).add(new AssignmentDetails(containerId, queue));
+    getDetails(op).add(new AssignmentDetails(containerId, queue));
   }
 
   public void addAllocationDetails(ContainerId containerId, String queue) {
@@ -111,21 +113,22 @@ public class AssignmentInformation {
     addAssignmentDetails(Operation.RESERVATION, containerId, queue);
   }
 
+  private List<AssignmentDetails> getDetails(Operation op) {
+    return operationDetails[op.ordinal()];
+  }
+
   public List<AssignmentDetails> getAllocationDetails() {
-    return operationDetails.get(Operation.ALLOCATION);
+    return getDetails(Operation.ALLOCATION);
   }
 
   public List<AssignmentDetails> getReservationDetails() {
-    return operationDetails.get(Operation.RESERVATION);
+    return getDetails(Operation.RESERVATION);
   }
 
   private ContainerId getFirstContainerIdFromOperation(Operation op) {
-    if (null != operationDetails.get(op)) {
-      List<AssignmentDetails> assignDetails =
-          operationDetails.get(op);
-      if (!assignDetails.isEmpty()) {
-        return assignDetails.get(0).containerId;
-      }
+    List<AssignmentDetails> assignDetails = getDetails(op);
+    if (assignDetails != null && !assignDetails.isEmpty()) {
+      return assignDetails.get(0).containerId;
     }
     return null;
   }