forked from apache/helix
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support Simple Greedy Rebalance Strategy
Support Simple Greedy Rebalance Strategy for RoundRobin assignment. It also supported with global limitation on how many partition assigned for each of the instance.
- Loading branch information
Showing
5 changed files
with
320 additions
and
0 deletions.
There are no files selected for viewing
91 changes: 91 additions & 0 deletions
91
helix-core/src/main/java/org/apache/helix/controller/common/CapacityNode.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| package org.apache.helix.controller.common; | ||
|
|
||
| /* | ||
| * 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.HashMap; | ||
| import java.util.HashSet; | ||
| import java.util.Map; | ||
| import java.util.Set; | ||
|
|
||
| /** | ||
| * A Node is an entity that can serve capacity recording purpose. It has a capacity and knowledge | ||
| * of partitions assigned to it, so it can decide if it can receive additional partitions. | ||
| */ | ||
| public class CapacityNode { | ||
| private int _currentlyAssigned; | ||
| private int _capacity; | ||
| private final String _id; | ||
| private final Map<String, Set<String>> _partitionMap; | ||
|
|
||
| public CapacityNode(String id) { | ||
| _partitionMap = new HashMap<>(); | ||
| _currentlyAssigned = 0; | ||
| this._id = id; | ||
| } | ||
|
|
||
| /** | ||
| * Check if this replica can be legally added to this node | ||
| * | ||
| * @param resource The resource to assign | ||
| * @param partition The partition to assign | ||
| * @return true if the assignment can be made, false otherwise | ||
| */ | ||
| public boolean canAdd(String resource, String partition) { | ||
| if (_currentlyAssigned >= _capacity || (_partitionMap.containsKey(resource) | ||
| && _partitionMap.get(resource).contains(partition))) { | ||
| return false; | ||
| } | ||
| _partitionMap.computeIfAbsent(resource, k -> new HashSet<>()).add(partition); | ||
| _currentlyAssigned++; | ||
| return true; | ||
| } | ||
|
|
||
| /** | ||
| * Set the capacity of this node | ||
| * @param capacity The capacity to set | ||
| */ | ||
| public void setCapacity(int capacity) { | ||
| _capacity = capacity; | ||
| } | ||
|
|
||
| /** | ||
| * Get the ID of this node | ||
| * @return The ID of this node | ||
| */ | ||
| public String getId() { | ||
| return _id; | ||
| } | ||
|
|
||
| /** | ||
| * Get number of partitions currently assigned to this node | ||
| * @return The number of partitions currently assigned to this node | ||
| */ | ||
| public int getCurrentlyAssigned() { | ||
| return _currentlyAssigned; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| StringBuilder sb = new StringBuilder(); | ||
| sb.append("##########\nname=").append(_id).append("\nassigned:").append(_currentlyAssigned) | ||
| .append("\ncapacity:").append(_capacity); | ||
| return sb.toString(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
100 changes: 100 additions & 0 deletions
100
...rc/main/java/org/apache/helix/controller/rebalancer/strategy/GreedyRebalanceStrategy.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| package org.apache.helix.controller.rebalancer.strategy; | ||
|
|
||
| /* | ||
| * 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.ArrayList; | ||
| import java.util.Collections; | ||
| import java.util.Comparator; | ||
| import java.util.HashSet; | ||
| import java.util.LinkedHashMap; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.Set; | ||
|
|
||
| import org.apache.helix.controller.common.CapacityNode; | ||
| import org.apache.helix.zookeeper.datamodel.ZNRecord; | ||
| import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| public class GreedyRebalanceStrategy implements RebalanceStrategy<ResourceControllerDataProvider> { | ||
| private static Logger logger = LoggerFactory.getLogger(GreedyRebalanceStrategy.class); | ||
| private String _resourceName; | ||
| private List<String> _partitions; | ||
| private LinkedHashMap<String, Integer> _states; | ||
|
|
||
| public GreedyRebalanceStrategy() { | ||
| } | ||
|
|
||
| @Override | ||
| public void init(String resourceName, final List<String> partitions, | ||
| final LinkedHashMap<String, Integer> states, int maximumPerNode) { | ||
| _resourceName = resourceName; | ||
| _partitions = partitions; | ||
| _states = states; | ||
| } | ||
|
|
||
| @Override | ||
| public ZNRecord computePartitionAssignment(final List<String> allNodes, final List<String> liveNodes, | ||
| final Map<String, Map<String, String>> currentMapping, ResourceControllerDataProvider clusterData) { | ||
| int numReplicas = countStateReplicas(); | ||
| ZNRecord znRecord = new ZNRecord(_resourceName); | ||
| if (liveNodes.size() == 0) { | ||
| return znRecord; | ||
| } | ||
|
|
||
| if (clusterData.getSimpleCapacitySet() == null) { | ||
| logger.warn("No capacity set for resource: " + _resourceName); | ||
| return znRecord; | ||
| } | ||
|
|
||
| List<CapacityNode> assignableNodes = new ArrayList<>(clusterData.getSimpleCapacitySet()); | ||
| Collections.sort(assignableNodes, Comparator.comparing(CapacityNode::getId)); | ||
|
|
||
| for (int i = 0, index = 0; i < _partitions.size(); i++) { | ||
| int startIndex = index; | ||
| List<String> preferenceList = new ArrayList<>(); | ||
| for (int j = 0; j < numReplicas; j++) { | ||
| if (index - startIndex >= assignableNodes.size()) { | ||
| logger.warn("No enough assignable nodes for resource: " + _resourceName); | ||
| break; | ||
| } | ||
| while (index - startIndex < assignableNodes.size()) { | ||
| CapacityNode node = assignableNodes.get(index++ % assignableNodes.size()); | ||
| if (node.canAdd(_resourceName, _partitions.get(i))) { | ||
| preferenceList.add(node.getId()); | ||
| break; | ||
| } | ||
| } | ||
| } | ||
| znRecord.setListField(_partitions.get(i), preferenceList); | ||
| } | ||
|
|
||
| return znRecord; | ||
| } | ||
|
|
||
| private int countStateReplicas() { | ||
| int total = 0; | ||
| for (Integer count : _states.values()) { | ||
| total += count; | ||
| } | ||
| return total; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
85 changes: 85 additions & 0 deletions
85
...ore/src/test/java/org/apache/helix/controller/rebalancer/TestGreedyRebalanceStrategy.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| package org.apache.helix.controller.rebalancer; | ||
|
|
||
| /* | ||
| * 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.ArrayList; | ||
| import java.util.Arrays; | ||
| import java.util.HashSet; | ||
| import java.util.LinkedHashMap; | ||
| import java.util.List; | ||
| import java.util.Set; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| import org.apache.helix.controller.common.CapacityNode; | ||
| import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider; | ||
| import org.apache.helix.controller.rebalancer.strategy.GreedyRebalanceStrategy; | ||
| import org.apache.helix.model.ClusterConfig; | ||
| import org.apache.helix.zookeeper.datamodel.ZNRecord; | ||
| import org.mockito.Mockito; | ||
| import org.testng.Assert; | ||
| import org.testng.annotations.Test; | ||
|
|
||
| import static org.mockito.Mockito.when; | ||
|
|
||
| public class TestGreedyRebalanceStrategy { | ||
| private static final String TEST_CLUSTER_NAME = "TestCluster"; | ||
| private static final String TEST_RESOURCE_PREFIX = "TestResource_"; | ||
|
|
||
| @Test | ||
| public void testAssignmentWithGlobalPartitionLimit() { | ||
|
|
||
| ResourceControllerDataProvider clusterDataCache = | ||
| Mockito.mock(ResourceControllerDataProvider.class); | ||
| LinkedHashMap<String, Integer> states = new LinkedHashMap<String, Integer>(2); | ||
| states.put("OFFLINE", 0); | ||
| states.put("ONLINE", 1); | ||
|
|
||
| Set<CapacityNode> capacityNodeSet = new HashSet<>(); | ||
| for (int i = 0; i < 5; i++) { | ||
| CapacityNode capacityNode = new CapacityNode("Node-" + i); | ||
| capacityNode.setCapacity(1); | ||
| capacityNodeSet.add(capacityNode); | ||
| } | ||
|
|
||
| List<String> liveNodes = | ||
| capacityNodeSet.stream().map(CapacityNode::getId).collect(Collectors.toList()); | ||
|
|
||
| List<String> partitions = new ArrayList<>(); | ||
| for (int i = 0; i < 3; i++) { | ||
| partitions.add(TEST_RESOURCE_PREFIX + "0_" + i); | ||
| } | ||
| when(clusterDataCache.getSimpleCapacitySet()).thenReturn(capacityNodeSet); | ||
|
|
||
| GreedyRebalanceStrategy greedyRebalanceStrategy = new GreedyRebalanceStrategy(); | ||
| greedyRebalanceStrategy.init(TEST_RESOURCE_PREFIX + 0, partitions, states, 1); | ||
| greedyRebalanceStrategy.computePartitionAssignment(null, liveNodes, null, clusterDataCache); | ||
|
|
||
| partitions = new ArrayList<>(); | ||
| for (int i = 0; i < 2; i++) { | ||
| partitions.add(TEST_RESOURCE_PREFIX + "1_" + i); | ||
| } | ||
| greedyRebalanceStrategy = new GreedyRebalanceStrategy(); | ||
| greedyRebalanceStrategy.init(TEST_RESOURCE_PREFIX + 1, partitions, states, 1); | ||
| greedyRebalanceStrategy.computePartitionAssignment(null, liveNodes, null, clusterDataCache); | ||
|
|
||
| Assert.assertEquals( | ||
| capacityNodeSet.stream().filter(node -> node.getCurrentlyAssigned() != 1).count(), 0); | ||
| } | ||
| } |