From 11729974b17c879eb40db23640f590cb730986c0 Mon Sep 17 00:00:00 2001 From: LZD-PratyushBhatt Date: Sat, 8 Aug 2026 22:16:55 +0530 Subject: [PATCH] Add REST accessor for decoded WAGED assignments The controller persists the WAGED baseline and best possible assignments through ZkBucketDataAccessor, which GZIPs the serialized record and splits it across numbered bucket ZNodes. Reading those ZNodes over the existing /zookeeper API therefore returns opaque compressed chunks, and the propertyStore API cannot reach them at all because ASSIGNMENT_METADATA lives at the cluster root rather than under PROPERTYSTORE. The only other assignment API, /partitionAssignment, recomputes a what-if placement and never reads what was actually persisted. Add WagedAssignmentAccessor, which reassembles the buckets, decompresses, and deserializes each resource assignment server side: GET /clusters/{clusterId}/wagedAssignment/bestPossible GET /clusters/{clusterId}/wagedAssignment/baseline Supported query params: format=IdealStateFormat (default) | CurrentStateFormat resources, instances, partitions - comma separated allowlists, since these payloads are large on real clusters includeMetadata (default true) - the persisted write version, its ZK mtime, and the bucket metadata, so callers can tell how fresh the decoded assignment is A missing assignment returns 404 rather than an empty body, so callers can distinguish "WAGED never persisted here" from "assignment is empty". The path layout and the per-resource decode are exposed from AssignmentMetadataStore instead of being duplicated in helix-rest, so the reader cannot drift from the writer. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../waged/AssignmentMetadataStore.java | 35 +- .../helix/WagedAssignmentAccessor.java | 337 ++++++++++++++++++ .../server/TestWagedAssignmentAccessor.java | 248 +++++++++++++ 3 files changed, 615 insertions(+), 5 deletions(-) create mode 100644 helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/WagedAssignmentAccessor.java create mode 100644 helix-rest/src/test/java/org/apache/helix/rest/server/TestWagedAssignmentAccessor.java 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; + } +}