diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/AssignmentMetadataStore.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/AssignmentMetadataStore.java index 157bb0ae4c..0078b519a8 100644 --- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/AssignmentMetadataStore.java +++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/AssignmentMetadataStore.java @@ -61,8 +61,30 @@ public class AssignmentMetadataStore { protected AssignmentMetadataStore(BucketDataAccessor bucketDataAccessor, String clusterName) { _dataAccessor = bucketDataAccessor; - _baselinePath = String.format(BASELINE_TEMPLATE, clusterName, ASSIGNMENT_METADATA_KEY); - _bestPossiblePath = String.format(BEST_POSSIBLE_TEMPLATE, clusterName, ASSIGNMENT_METADATA_KEY); + _baselinePath = getBaselinePath(clusterName); + _bestPossiblePath = getBestPossiblePath(clusterName); + } + + /** + * Returns the metadata store path holding the persisted WAGED baseline assignment. + * Exposed so that read-only consumers, such as the REST layer, can locate the assignment without + * duplicating the path layout. + * @param clusterName the cluster whose baseline assignment is being located + * @return the bucketized root path of the baseline assignment + */ + public static String getBaselinePath(String clusterName) { + return String.format(BASELINE_TEMPLATE, clusterName, ASSIGNMENT_METADATA_KEY); + } + + /** + * Returns the metadata store path holding the persisted WAGED best possible assignment. + * Exposed so that read-only consumers, such as the REST layer, can locate the assignment without + * duplicating the path layout. + * @param clusterName the cluster whose best possible assignment is being located + * @return the bucketized root path of the best possible assignment + */ + public static String getBestPossiblePath(String clusterName) { + return String.format(BEST_POSSIBLE_TEMPLATE, clusterName, ASSIGNMENT_METADATA_KEY); } public Map getBaseline() { @@ -234,10 +256,13 @@ private HelixProperty combineAssignments(String name, /** * Returns a Map of (ResourceName, ResourceAssignment) pairs. - * @param property - * @return + * This is the inverse of {@link #combineAssignments(String, Map)} and is exposed so that + * read-only consumers, such as the REST layer, decode a persisted assignment with the exact same + * contract the controller used to encode it. + * @param property the combined assignment read from the bucketized metadata store + * @return the per-resource assignments */ - private Map splitAssignments(HelixProperty property) { + public static Map splitAssignments(HelixProperty property) { Map assignmentMap = new HashMap<>(); // Convert each resource's assignment String into a ResourceAssignment object and put it in a // map diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/WagedAssignmentAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/WagedAssignmentAccessor.java new file mode 100644 index 0000000000..b9d074b9a9 --- /dev/null +++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/WagedAssignmentAccessor.java @@ -0,0 +1,337 @@ +package org.apache.helix.rest.server.resources.helix; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.stream.Collectors; +import javax.ws.rs.DefaultValue; +import javax.ws.rs.GET; +import javax.ws.rs.Path; +import javax.ws.rs.PathParam; +import javax.ws.rs.QueryParam; +import javax.ws.rs.core.Response; + +import com.codahale.metrics.annotation.ResponseMetered; +import com.codahale.metrics.annotation.Timed; +import org.apache.helix.AccessOption; +import org.apache.helix.BaseDataAccessor; +import org.apache.helix.HelixProperty; +import org.apache.helix.controller.rebalancer.waged.AssignmentMetadataStore; +import org.apache.helix.model.Partition; +import org.apache.helix.model.ResourceAssignment; +import org.apache.helix.rest.common.HttpConstants; +import org.apache.helix.rest.server.filters.ClusterAuth; +import org.apache.helix.zookeeper.zkclient.exception.ZkNoNodeException; +import org.apache.zookeeper.data.Stat; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Read-only access to the WAGED assignments the controller persisted in the assignment metadata + * store, decoded into plain JSON. + *

+ * The controller writes these assignments through + * {@link org.apache.helix.manager.zk.ZkBucketDataAccessor}, which GZIPs the serialized record and + * splits it across numbered bucket ZNodes. Reading the raw ZNodes, for example through + * {@code /zookeeper/{path}?command=getBinaryData}, therefore returns opaque compressed chunks that + * a caller cannot interpret. This accessor performs the bucket reassembly, decompression, and + * per-resource deserialization server side and returns the resulting partition placements. + *

+ * Note that this reports what was last persisted, which is the controller's own view of its + * latest computation. It is not a live recomputation, and it may lag the in-memory state of the + * active controller. Use {@code /clusters/{clusterId}/partitionAssignment} for what-if computation + * instead. + */ +@ClusterAuth +@Path("/clusters/{clusterId}/wagedAssignment") +public class WagedAssignmentAccessor extends AbstractHelixResource { + private static final Logger LOG = LoggerFactory.getLogger(WagedAssignmentAccessor.class); + + // Bookkeeping ZNodes written by ZkBucketDataAccessor alongside the versioned payload. + private static final String LAST_SUCCESSFUL_WRITE_KEY = "LAST_SUCCESSFUL_WRITE"; + private static final String LAST_WRITE_KEY = "LAST_WRITE"; + private static final String METADATA_KEY = "METADATA"; + + private static final String CLUSTER_ID_FIELD = "cluster"; + private static final String ASSIGNMENT_TYPE_FIELD = "assignmentType"; + private static final String FORMAT_FIELD = "format"; + private static final String METADATA_FIELD = "metadata"; + private static final String ASSIGNMENT_FIELD = "assignment"; + + public enum AssignmentType { + BASELINE, + BEST_POSSIBLE + } + + /** + * Result shape. Mirrors the options offered by + * {@link ResourceAssignmentOptimizerAccessor} so both assignment APIs can be consumed the + * same way. + */ + public enum AssignmentFormat { + /** resource -> partition -> instance -> state. */ + IdealStateFormat, + /** instance -> resource -> partition -> state. */ + CurrentStateFormat + } + + /** + * Sample HTTP URL: + * {@code GET /clusters/{clusterId}/wagedAssignment/bestPossible?format=IdealStateFormat&resources=db0,db1} + *

+ * Returns the decoded best possible assignment, which is the placement the WAGED rebalancer + * converged on and handed to the rest of the controller pipeline. + * + * @param clusterId the cluster to read + * @param formatStr {@link AssignmentFormat}, defaults to {@code IdealStateFormat} + * @param resources optional comma separated resource allowlist + * @param instances optional comma separated instance allowlist + * @param partitions optional comma separated partition allowlist + * @param includeMetadata whether to include the persisted write bookkeeping, defaults to true + * @return the decoded assignment + */ + @ResponseMetered(name = HttpConstants.READ_REQUEST) + @Timed(name = HttpConstants.READ_REQUEST) + @GET + @Path("bestPossible") + public Response getBestPossibleAssignment(@PathParam("clusterId") String clusterId, + @QueryParam("format") @DefaultValue("IdealStateFormat") String formatStr, + @QueryParam("resources") String resources, @QueryParam("instances") String instances, + @QueryParam("partitions") String partitions, + @QueryParam("includeMetadata") @DefaultValue("true") boolean includeMetadata) { + return getAssignment(clusterId, AssignmentType.BEST_POSSIBLE, formatStr, resources, instances, + partitions, includeMetadata); + } + + /** + * Sample HTTP URL: + * {@code GET /clusters/{clusterId}/wagedAssignment/baseline?format=CurrentStateFormat} + *

+ * Returns the decoded baseline assignment, which is the steady state placement WAGED computes + * ignoring transient conditions such as instances being temporarily down. + * + * @param clusterId the cluster to read + * @param formatStr {@link AssignmentFormat}, defaults to {@code IdealStateFormat} + * @param resources optional comma separated resource allowlist + * @param instances optional comma separated instance allowlist + * @param partitions optional comma separated partition allowlist + * @param includeMetadata whether to include the persisted write bookkeeping, defaults to true + * @return the decoded assignment + */ + @ResponseMetered(name = HttpConstants.READ_REQUEST) + @Timed(name = HttpConstants.READ_REQUEST) + @GET + @Path("baseline") + public Response getBaselineAssignment(@PathParam("clusterId") String clusterId, + @QueryParam("format") @DefaultValue("IdealStateFormat") String formatStr, + @QueryParam("resources") String resources, @QueryParam("instances") String instances, + @QueryParam("partitions") String partitions, + @QueryParam("includeMetadata") @DefaultValue("true") boolean includeMetadata) { + return getAssignment(clusterId, AssignmentType.BASELINE, formatStr, resources, instances, + partitions, includeMetadata); + } + + private Response getAssignment(String clusterId, AssignmentType assignmentType, String formatStr, + String resources, String instances, String partitions, boolean includeMetadata) { + AssignmentFormat format; + try { + format = AssignmentFormat.valueOf(formatStr); + } catch (IllegalArgumentException e) { + return badRequest(String.format("Invalid format: %s. Supported formats are %s", formatStr, + Arrays.toString(AssignmentFormat.values()))); + } + + Set resourceFilter = parseFilter(resources); + Set instanceFilter = parseFilter(instances); + Set partitionFilter = parseFilter(partitions); + String rootPath = getAssignmentPath(clusterId, assignmentType); + + Map assignments; + try { + HelixProperty combined = + getZkBucketDataAccessor().compressedBucketRead(rootPath, HelixProperty.class); + assignments = AssignmentMetadataStore.splitAssignments(combined); + } catch (ZkNoNodeException e) { + // WAGED has never persisted this assignment for the cluster, so there is nothing to decode. + return notFound(String.format( + "No %s assignment found at %s. The cluster may have no WAGED resources, or the " + + "controller may not have persisted an assignment yet.", assignmentType, rootPath)); + } catch (Exception e) { + LOG.error("Failed to read {} assignment for cluster {} at path {}", assignmentType, clusterId, + rootPath, e); + return serverError(e); + } + + Map response = new LinkedHashMap<>(); + response.put(CLUSTER_ID_FIELD, clusterId); + response.put(ASSIGNMENT_TYPE_FIELD, assignmentType.name()); + response.put(FORMAT_FIELD, format.name()); + if (includeMetadata) { + response.put(METADATA_FIELD, readWriteMetadata(rootPath)); + } + response.put(ASSIGNMENT_FIELD, + format == AssignmentFormat.CurrentStateFormat ? toCurrentStateFormat(assignments, + resourceFilter, instanceFilter, partitionFilter) + : toIdealStateFormat(assignments, resourceFilter, instanceFilter, partitionFilter)); + return JSONRepresentation(response); + } + + private static String getAssignmentPath(String clusterId, AssignmentType assignmentType) { + return assignmentType == AssignmentType.BASELINE ? AssignmentMetadataStore + .getBaselinePath(clusterId) : AssignmentMetadataStore.getBestPossiblePath(clusterId); + } + + private static Set parseFilter(String csv) { + if (csv == null || csv.trim().isEmpty()) { + return Collections.emptySet(); + } + return Arrays.stream(csv.split(",")).map(String::trim).filter(s -> !s.isEmpty()) + .collect(Collectors.toSet()); + } + + private static boolean included(Set filter, String value) { + return filter.isEmpty() || filter.contains(value); + } + + /** + * Reads the bookkeeping ZNodes the bucketized write left behind so callers can tell how fresh the + * decoded assignment is and how large the persisted payload was. Best effort: a partially written + * or concurrently garbage collected path yields whatever could be read rather than failing the + * whole request. + */ + private Map readWriteMetadata(String rootPath) { + Map metadata = new LinkedHashMap<>(); + BaseDataAccessor accessor = getByteArrayDataAccessor(); + try { + String lastSuccessfulWrite = readString(accessor, rootPath + "/" + LAST_SUCCESSFUL_WRITE_KEY); + metadata.put("lastSuccessfulWriteVersion", lastSuccessfulWrite); + metadata.put("lastWriteVersion", readString(accessor, rootPath + "/" + LAST_WRITE_KEY)); + + Stat[] stats = + accessor.getStats(Collections.singletonList(rootPath + "/" + LAST_SUCCESSFUL_WRITE_KEY), + AccessOption.PERSISTENT); + if (stats != null && stats.length > 0 && stats[0] != null) { + metadata.put("lastSuccessfulWriteTimeMs", stats[0].getMtime()); + } + + if (lastSuccessfulWrite != null) { + byte[] rawBucketMetadata = accessor + .get(rootPath + "/" + lastSuccessfulWrite + "/" + METADATA_KEY, null, + AccessOption.PERSISTENT); + if (rawBucketMetadata != null) { + metadata.put("bucketMetadata", OBJECT_MAPPER.readValue(rawBucketMetadata, Map.class)); + } + } + } catch (Exception e) { + // Metadata is diagnostic only, so never fail the assignment read because of it. + LOG.warn("Failed to read assignment write metadata at path {}", rootPath, e); + } + return metadata; + } + + private static String readString(BaseDataAccessor accessor, String path) { + byte[] bytes = accessor.get(path, null, AccessOption.PERSISTENT); + return bytes == null ? null : new String(bytes); + } + + /** + * Builds resource -> partition -> instance -> state, matching the layout of an IdealState + * map field. + */ + private static Map>> toIdealStateFormat( + Map assignments, Set resourceFilter, + Set instanceFilter, Set partitionFilter) { + Map>> result = new TreeMap<>(); + for (Map.Entry entry : assignments.entrySet()) { + String resource = entry.getKey(); + if (!included(resourceFilter, resource)) { + continue; + } + Map> partitionMap = new TreeMap<>(); + for (Partition partition : entry.getValue().getMappedPartitions()) { + String partitionName = partition.getPartitionName(); + if (!included(partitionFilter, partitionName)) { + continue; + } + Map replicaMap = + filterReplicas(entry.getValue().getReplicaMap(partition), instanceFilter); + if (!replicaMap.isEmpty()) { + partitionMap.put(partitionName, replicaMap); + } + } + if (!partitionMap.isEmpty()) { + result.put(resource, partitionMap); + } + } + return result; + } + + /** + * Builds instance -> resource -> partition -> state, which is the convenient shape when + * asking what a single host is expected to carry. + */ + private static Map>> toCurrentStateFormat( + Map assignments, Set resourceFilter, + Set instanceFilter, Set partitionFilter) { + Map>> result = new TreeMap<>(); + for (Map.Entry entry : assignments.entrySet()) { + String resource = entry.getKey(); + if (!included(resourceFilter, resource)) { + continue; + } + for (Partition partition : entry.getValue().getMappedPartitions()) { + String partitionName = partition.getPartitionName(); + if (!included(partitionFilter, partitionName)) { + continue; + } + for (Map.Entry replica : entry.getValue().getReplicaMap(partition) + .entrySet()) { + if (!included(instanceFilter, replica.getKey())) { + continue; + } + result.computeIfAbsent(replica.getKey(), k -> new TreeMap<>()) + .computeIfAbsent(resource, k -> new TreeMap<>()) + .put(partitionName, replica.getValue()); + } + } + } + return result; + } + + private static Map filterReplicas(Map replicaMap, + Set instanceFilter) { + if (instanceFilter.isEmpty()) { + return new TreeMap<>(replicaMap); + } + Map filtered = new TreeMap<>(); + for (Map.Entry replica : replicaMap.entrySet()) { + if (instanceFilter.contains(replica.getKey())) { + filtered.put(replica.getKey(), replica.getValue()); + } + } + return filtered; + } +} diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestWagedAssignmentAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestWagedAssignmentAccessor.java new file mode 100644 index 0000000000..1beb2f0e7a --- /dev/null +++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestWagedAssignmentAccessor.java @@ -0,0 +1,248 @@ +package org.apache.helix.rest.server; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import javax.ws.rs.core.Response; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.common.collect.ImmutableMap; +import org.apache.helix.AccessOption; +import org.apache.helix.HelixProperty; +import org.apache.helix.TestHelper; +import org.apache.helix.controller.rebalancer.waged.AssignmentMetadataStore; +import org.apache.helix.manager.zk.ZkBucketDataAccessor; +import org.apache.helix.model.Partition; +import org.apache.helix.model.ResourceAssignment; +import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordJacksonSerializer; +import org.apache.helix.zookeeper.zkclient.serialize.ZkSerializer; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +/** + * Verifies that the WAGED assignment REST accessor decodes what the controller persisted through + * the bucketized, GZIP compressed assignment metadata store. + */ +public class TestWagedAssignmentAccessor extends AbstractTestClass { + private static final String TEST_CLUSTER = "TestCluster_0"; + private static final String UNKNOWN_CLUSTER = "NonExistentWagedAssignmentCluster"; + private static final String RESOURCE_0 = "wagedDb0"; + private static final String RESOURCE_1 = "wagedDb1"; + private static final String INSTANCE_0 = "wagedInstance0"; + private static final String INSTANCE_1 = "wagedInstance1"; + private static final String MASTER = "MASTER"; + private static final String SLAVE = "SLAVE"; + private static final ZkSerializer SERIALIZER = new ZNRecordJacksonSerializer(); + + private ZkBucketDataAccessor _bucketDataAccessor; + + @BeforeClass + public void beforeClass() throws IOException { + _bucketDataAccessor = new ZkBucketDataAccessor(ZK_ADDR); + _bucketDataAccessor.compressedBucketWrite( + AssignmentMetadataStore.getBestPossiblePath(TEST_CLUSTER), + combineAssignments("BEST_POSSIBLE", buildBestPossibleAssignment())); + _bucketDataAccessor.compressedBucketWrite(AssignmentMetadataStore.getBaselinePath(TEST_CLUSTER), + combineAssignments("BASELINE", buildBaselineAssignment())); + } + + @AfterClass + public void afterClass() { + if (_bucketDataAccessor != null) { + _bucketDataAccessor.close(); + } + _baseAccessor.remove("/" + TEST_CLUSTER + "/ASSIGNMENT_METADATA", AccessOption.PERSISTENT); + } + + @Test + public void testGetBestPossibleAssignmentIsDecoded() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode node = getAssignmentNode("bestPossible", null); + + Assert.assertEquals(node.get("cluster").asText(), TEST_CLUSTER); + Assert.assertEquals(node.get("assignmentType").asText(), "BEST_POSSIBLE"); + Assert.assertEquals(node.get("format").asText(), "IdealStateFormat"); + + JsonNode assignment = node.get("assignment"); + Assert.assertEquals(assignment.size(), 2); + Assert.assertEquals(assignment.get(RESOURCE_0).get(RESOURCE_0 + "_0").get(INSTANCE_0).asText(), + MASTER); + Assert.assertEquals(assignment.get(RESOURCE_0).get(RESOURCE_0 + "_0").get(INSTANCE_1).asText(), + SLAVE); + Assert.assertEquals(assignment.get(RESOURCE_0).get(RESOURCE_0 + "_1").get(INSTANCE_1).asText(), + MASTER); + Assert.assertEquals(assignment.get(RESOURCE_1).get(RESOURCE_1 + "_0").get(INSTANCE_0).asText(), + MASTER); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testGetBaselineAssignmentIsDecoded() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode node = getAssignmentNode("baseline", null); + + Assert.assertEquals(node.get("assignmentType").asText(), "BASELINE"); + JsonNode assignment = node.get("assignment"); + Assert.assertEquals(assignment.size(), 1); + Assert.assertEquals(assignment.get(RESOURCE_0).get(RESOURCE_0 + "_0").get(INSTANCE_1).asText(), + MASTER); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testCurrentStateFormat() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode node = + getAssignmentNode("bestPossible", ImmutableMap.of("format", "CurrentStateFormat")); + + Assert.assertEquals(node.get("format").asText(), "CurrentStateFormat"); + JsonNode assignment = node.get("assignment"); + // Inverted: instance -> resource -> partition -> state + Assert.assertEquals(assignment.size(), 2); + Assert.assertEquals( + assignment.get(INSTANCE_0).get(RESOURCE_0).get(RESOURCE_0 + "_0").asText(), MASTER); + Assert.assertEquals( + assignment.get(INSTANCE_0).get(RESOURCE_1).get(RESOURCE_1 + "_0").asText(), MASTER); + Assert.assertEquals( + assignment.get(INSTANCE_1).get(RESOURCE_0).get(RESOURCE_0 + "_1").asText(), MASTER); + Assert.assertFalse(assignment.get(INSTANCE_1).has(RESOURCE_1)); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testResourceFilter() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode assignment = + getAssignmentNode("bestPossible", ImmutableMap.of("resources", RESOURCE_1)) + .get("assignment"); + + Assert.assertEquals(assignment.size(), 1); + Assert.assertTrue(assignment.has(RESOURCE_1)); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testInstanceFilterDropsEmptyResources() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode assignment = + getAssignmentNode("bestPossible", ImmutableMap.of("instances", INSTANCE_1)) + .get("assignment"); + + // wagedDb1 only lives on instance0, so it drops out entirely. + Assert.assertEquals(assignment.size(), 1); + Assert.assertEquals(assignment.get(RESOURCE_0).get(RESOURCE_0 + "_0").size(), 1); + Assert.assertEquals(assignment.get(RESOURCE_0).get(RESOURCE_0 + "_0").get(INSTANCE_1).asText(), + SLAVE); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testPartitionFilter() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode assignment = + getAssignmentNode("bestPossible", ImmutableMap.of("partitions", RESOURCE_0 + "_1")) + .get("assignment"); + + Assert.assertEquals(assignment.size(), 1); + Assert.assertEquals(assignment.get(RESOURCE_0).size(), 1); + Assert.assertTrue(assignment.get(RESOURCE_0).has(RESOURCE_0 + "_1")); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testWriteMetadata() throws IOException { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + JsonNode metadata = getAssignmentNode("bestPossible", null).get("metadata"); + + Assert.assertNotNull(metadata); + Assert.assertNotNull(metadata.get("lastSuccessfulWriteVersion")); + Assert.assertTrue(metadata.get("lastSuccessfulWriteTimeMs").asLong() > 0); + Assert.assertTrue(metadata.get("bucketMetadata").get("DATA_SIZE").asInt() > 0); + + // Metadata can be turned off for callers that only want the placement. + JsonNode noMetadata = + getAssignmentNode("bestPossible", ImmutableMap.of("includeMetadata", "false")); + Assert.assertFalse(noMetadata.has("metadata")); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testInvalidFormatIsRejected() { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + get("clusters/" + TEST_CLUSTER + "/wagedAssignment/bestPossible", + ImmutableMap.of("format", "NotAFormat"), Response.Status.BAD_REQUEST.getStatusCode(), true); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + @Test + public void testMissingAssignmentReturnsNotFound() { + System.out.println("Start test :" + TestHelper.getTestMethodName()); + get("clusters/" + UNKNOWN_CLUSTER + "/wagedAssignment/bestPossible", null, + Response.Status.NOT_FOUND.getStatusCode(), true); + System.out.println("End test :" + TestHelper.getTestMethodName()); + } + + private JsonNode getAssignmentNode(String assignmentType, Map queryParams) + throws IOException { + String body = get("clusters/" + TEST_CLUSTER + "/wagedAssignment/" + assignmentType, + queryParams, Response.Status.OK.getStatusCode(), true); + return OBJECT_MAPPER.readTree(body); + } + + private static Map buildBestPossibleAssignment() { + Map assignments = new HashMap<>(); + ResourceAssignment resource0 = new ResourceAssignment(RESOURCE_0); + resource0.addReplicaMap(new Partition(RESOURCE_0 + "_0"), + ImmutableMap.of(INSTANCE_0, MASTER, INSTANCE_1, SLAVE)); + resource0.addReplicaMap(new Partition(RESOURCE_0 + "_1"), ImmutableMap.of(INSTANCE_1, MASTER)); + assignments.put(RESOURCE_0, resource0); + + ResourceAssignment resource1 = new ResourceAssignment(RESOURCE_1); + resource1.addReplicaMap(new Partition(RESOURCE_1 + "_0"), ImmutableMap.of(INSTANCE_0, MASTER)); + assignments.put(RESOURCE_1, resource1); + return assignments; + } + + private static Map buildBaselineAssignment() { + Map assignments = new HashMap<>(); + ResourceAssignment resource0 = new ResourceAssignment(RESOURCE_0); + resource0.addReplicaMap(new Partition(RESOURCE_0 + "_0"), + ImmutableMap.of(INSTANCE_1, MASTER, INSTANCE_0, SLAVE)); + assignments.put(RESOURCE_0, resource0); + return assignments; + } + + /** + * Mirrors {@code AssignmentMetadataStore#combineAssignments} so the test exercises the exact wire + * format the controller writes. + */ + private static HelixProperty combineAssignments(String name, + Map assignmentMap) { + HelixProperty property = new HelixProperty(name); + assignmentMap.forEach((resource, assignment) -> property.getRecord() + .setSimpleField(resource, new String(SERIALIZER.serialize(assignment.getRecord())))); + return property; + } +}