Skip to content

Commit

Permalink
add enumMap instead of list for tracking taskResourceUsageTrackets
Browse files Browse the repository at this point in the history
Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
  • Loading branch information
kaushalmahi12 committed Apr 30, 2024
1 parent aa4fd2b commit bf11c85
Showing 1 changed file with 7 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import org.opensearch.tasks.TaskCancellation;

import java.util.ArrayList;
import java.util.EnumMap;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicLong;
Expand All @@ -29,10 +30,10 @@ public class TaskResourceUsageTrackers {
private TaskResourceUsageTracker cpuUsageTracker;
private TaskResourceUsageTracker heapUsageTracker;
private TaskResourceUsageTracker elapsedTimeTracker;
private final List<TaskResourceUsageTracker> all;
private final EnumMap<TaskResourceUsageTrackerType, TaskResourceUsageTracker> all;

public TaskResourceUsageTrackers() {
all = new ArrayList<>(3);
all = new EnumMap<>(TaskResourceUsageTrackerType.class);
}

/**
Expand All @@ -41,7 +42,7 @@ public TaskResourceUsageTrackers() {
*/
public void addCpuUsageTracker(final TaskResourceUsageTracker cpuUsageTracker) {
this.cpuUsageTracker = cpuUsageTracker;
all.add(cpuUsageTracker);
all.put(TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, cpuUsageTracker);
}

/**
Expand All @@ -50,7 +51,7 @@ public void addCpuUsageTracker(final TaskResourceUsageTracker cpuUsageTracker) {
*/
public void addHeapUsageTracker(final TaskResourceUsageTracker heapUsageTracker) {
this.heapUsageTracker = heapUsageTracker;
all.add(heapUsageTracker);
all.put(TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER, heapUsageTracker);
}

/**
Expand All @@ -59,7 +60,7 @@ public void addHeapUsageTracker(final TaskResourceUsageTracker heapUsageTracker)
*/
public void addElapsedTimeTracker(final TaskResourceUsageTracker elapsedTimeTracker) {
this.elapsedTimeTracker = elapsedTimeTracker;
all.add(elapsedTimeTracker);
all.put(TaskResourceUsageTrackerType.ELAPSED_TIME_TRACKER, elapsedTimeTracker);
}

/**
Expand Down Expand Up @@ -91,7 +92,7 @@ public Optional<TaskResourceUsageTracker> getElapsedTimeTracker() {
* @return
*/
public List<TaskResourceUsageTracker> all() {
return all;
return new ArrayList<>(all.values());
}

/**
Expand Down

0 comments on commit bf11c85

Please sign in to comment.