Repository: aurora Updated Branches: refs/heads/master 92f6d9f64 -> 2ac16f015
Normalize state endpoint to reduce API payload size. Reviewed at https://reviews.apache.org/r/59565/ Project: http://git-wip-us.apache.org/repos/asf/aurora/repo Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/2ac16f01 Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/2ac16f01 Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/2ac16f01 Branch: refs/heads/master Commit: 2ac16f015cfb7601d78e9610745f3cf28e03af72 Parents: 92f6d9f Author: David McLaughlin <[email protected]> Authored: Thu May 25 10:37:59 2017 -0700 Committer: David McLaughlin <[email protected]> Committed: Thu May 25 10:37:59 2017 -0700 ---------------------------------------------------------------------- .../org/apache/aurora/scheduler/http/State.java | 54 +++++++++++++++++++- .../apache/aurora/scheduler/http/StateTest.java | 51 +++++++++++++----- 2 files changed, 91 insertions(+), 14 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/aurora/blob/2ac16f01/src/main/java/org/apache/aurora/scheduler/http/State.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/aurora/scheduler/http/State.java b/src/main/java/org/apache/aurora/scheduler/http/State.java index dc356b7..6d1b400 100644 --- a/src/main/java/org/apache/aurora/scheduler/http/State.java +++ b/src/main/java/org/apache/aurora/scheduler/http/State.java @@ -13,6 +13,11 @@ */ package org.apache.aurora.scheduler.http; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; import javax.inject.Inject; import javax.ws.rs.GET; import javax.ws.rs.Path; @@ -20,13 +25,17 @@ import javax.ws.rs.Produces; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Multimap; import org.apache.aurora.scheduler.preemptor.ClusterState; import org.apache.aurora.scheduler.preemptor.ClusterStateImpl; import org.apache.aurora.scheduler.preemptor.PreemptionVictim; +import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.toList; import static org.apache.aurora.scheduler.http.api.GsonMessageBodyHandler.GSON; @@ -37,6 +46,47 @@ import static org.apache.aurora.scheduler.http.api.GsonMessageBodyHandler.GSON; public class State { private final ClusterState clusterState; + @VisibleForTesting + static String taskKey(ITaskConfig config) { + return String.format("%s/%s/%s-%d", + config.getJob().getRole(), + config.getJob().getEnvironment(), + config.getJob().getName(), + config.hashCode()); + } + + /** + * Class that normalizes the cluster state by removing instance-specific information and moving + * tasks into a lookup table. This reduces the total size of the payload from O(RUNNING_TASKS) to + * O(DISTINCT_RUNNING_TASK_CONFIGS). + */ + private static class NormalizedClusterState { + private final Map<String, ITaskConfig> taskConfigs; + private final Map<String, List<String>> agents; + + NormalizedClusterState( + Map<String, ITaskConfig> taskConfigMap, + Map<String, List<String>> agentTasksMap) { + + this.taskConfigs = requireNonNull(taskConfigMap); + this.agents = requireNonNull(agentTasksMap); + } + + static NormalizedClusterState fromClusterState(Multimap<String, PreemptionVictim> state) { + Map<String, ITaskConfig> tasks = new HashMap<>(); + ImmutableMap.Builder<String, List<String>> agents = new ImmutableMap.Builder<>(); + for (Entry<String, Collection<PreemptionVictim>> entry: state.asMap().entrySet()) { + for (PreemptionVictim victim: entry.getValue()) { + tasks.putIfAbsent(taskKey(victim.getConfig()), victim.getConfig()); + } + agents.put( + entry.getKey(), + entry.getValue().stream().map(e -> taskKey(e.getConfig())).collect(toList())); + } + return new NormalizedClusterState(tasks, agents.build()); + } + } + @Inject State(ClusterStateImpl clusterState) { this.clusterState = requireNonNull(clusterState); @@ -45,7 +95,7 @@ public class State { @GET @Produces(MediaType.APPLICATION_JSON) public Response getState() { - Multimap<String, PreemptionVictim> state = clusterState.getSlavesToActiveTasks(); - return Response.ok(GSON.toJson(state.asMap())).build(); + return Response.ok(GSON.toJson(NormalizedClusterState.fromClusterState( + clusterState.getSlavesToActiveTasks()))).build(); } } http://git-wip-us.apache.org/repos/asf/aurora/blob/2ac16f01/src/test/java/org/apache/aurora/scheduler/http/StateTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/http/StateTest.java b/src/test/java/org/apache/aurora/scheduler/http/StateTest.java index 0af8b0d..0685d6e 100644 --- a/src/test/java/org/apache/aurora/scheduler/http/StateTest.java +++ b/src/test/java/org/apache/aurora/scheduler/http/StateTest.java @@ -13,12 +13,16 @@ */ package org.apache.aurora.scheduler.http; +import java.util.List; + import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMultimap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSetMultimap; +import com.google.common.collect.Iterators; import com.google.common.collect.Multimap; import org.apache.aurora.common.testing.easymock.EasyMockTest; @@ -28,9 +32,11 @@ import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.scheduler.preemptor.ClusterStateImpl; import org.apache.aurora.scheduler.preemptor.PreemptionVictim; import org.apache.aurora.scheduler.storage.entities.IAssignedTask; +import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import org.junit.Before; import org.junit.Test; +import static org.apache.aurora.scheduler.http.State.taskKey; import static org.easymock.EasyMock.expect; import static org.junit.Assert.assertEquals; @@ -53,29 +59,50 @@ public class StateTest extends EasyMockTest { return victims.build(); } - private IAssignedTask makeTask(String taskId, String agentId) { + private IAssignedTask makeTask(String taskId, String agentId, TaskConfig taskConfig) { return IAssignedTask.build(new AssignedTask() .setSlaveId(agentId) .setSlaveHost(agentId + "-host") .setTaskId(taskId) - .setTask(new TaskConfig().setJob(new JobKey("role", "env", "job")))); + .setTask(taskConfig)); + } + + private static List<String> getAgentTasks(JsonNode tasks) { + return ImmutableList.copyOf(Iterators.transform(tasks.elements(), JsonNode::asText)); } @Test public void testJson() throws Exception { + TaskConfig task1 = new TaskConfig().setJob(new JobKey("role", "env", "job")); + TaskConfig task2 = new TaskConfig().setJob(new JobKey("role", "env", "job")); + TaskConfig task3 = new TaskConfig().setJob(new JobKey("role", "env", "another-job")); + String task1Key = taskKey(ITaskConfig.build(task1)); + String task2Key = taskKey(ITaskConfig.build(task2)); + String task3Key = taskKey(ITaskConfig.build(task3)); + + // Tests: + // Same task config on multiple hosts. + // Different tasks configs for the same job. + // Multiple job keys. Multimap<String, PreemptionVictim> expected = createState( - makeTask("task1", "agent1"), - makeTask("task2", "agent1"), - makeTask("task3", "agent2"), - makeTask("task4", "agent3")); + makeTask("task1", "agent1", task1), + makeTask("task2", "agent1", task1), + makeTask("task3", "agent2", task2), + makeTask("task4", "agent3", task3), + makeTask("task5", "agent2", task1)); + expect(clusterState.getSlavesToActiveTasks()).andReturn(expected); control.replay(); JsonNode result = new ObjectMapper().readTree((String) state.getState().getEntity()); - assertEquals(ImmutableList.of("agent1", "agent2", "agent3"), - ImmutableList.copyOf(result.fieldNames())); - assertEquals(2, ((ArrayNode) result.get("agent1")).size()); - assertEquals(1, ((ArrayNode) result.get("agent2")).size()); - assertEquals(1, ((ArrayNode) result.get("agent3")).size()); + ObjectNode tasks = (ObjectNode) result.get("taskConfigs"); + + assertEquals(ImmutableSet.of(task1Key, task2Key, task3Key), + ImmutableSet.copyOf(tasks.fieldNames())); + + JsonNode agents = result.get("agents"); + assertEquals(ImmutableList.of(task1Key, task1Key), getAgentTasks(agents.get("agent1"))); + assertEquals(ImmutableList.of(task2Key, task1Key), getAgentTasks(agents.get("agent2"))); + assertEquals(ImmutableList.of(task3Key), getAgentTasks(agents.get("agent3"))); } }
