Skip to content

Commit 77ecdb9

Browse files
jiajunwangJiajun Wang
authored andcommitted
Adjust the expected replica count according to fault zone count. (#476)
The rebalancer should determine the expected replica count according to the fault zone instead of the node count only.
1 parent ca8a114 commit 77ecdb9

6 files changed

Lines changed: 76 additions & 71 deletions

File tree

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/AssignableNode.java

Lines changed: 25 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -64,12 +64,10 @@ public class AssignableNode implements Comparable<AssignableNode> {
6464
* @param clusterConfig
6565
* @param instanceConfig
6666
* @param instanceName
67-
* @param existingAssignment A collection of replicas that have been pre-allocated to the node.
6867
*/
69-
AssignableNode(ClusterConfig clusterConfig, InstanceConfig instanceConfig, String instanceName,
70-
Collection<AssignableReplica> existingAssignment) {
68+
AssignableNode(ClusterConfig clusterConfig, InstanceConfig instanceConfig, String instanceName) {
7169
_instanceName = instanceName;
72-
refresh(clusterConfig, instanceConfig, existingAssignment);
70+
refresh(clusterConfig, instanceConfig);
7371
}
7472

7573
private void reset() {
@@ -88,10 +86,8 @@ private void reset() {
8886
* subject to change. If the assumption is no longer true, this function should become private.
8987
* @param clusterConfig - the Cluster Config of the cluster where the node is located
9088
* @param instanceConfig - the Instance Config of the node
91-
* @param existingAssignment - all the existing replicas that are current assigned to the node
9289
*/
93-
private void refresh(ClusterConfig clusterConfig, InstanceConfig instanceConfig,
94-
Collection<AssignableReplica> existingAssignment) {
90+
private void refresh(ClusterConfig clusterConfig, InstanceConfig instanceConfig) {
9591
reset();
9692

9793
Map<String, Integer> instanceCapacity = fetchInstanceCapacity(clusterConfig, instanceConfig);
@@ -101,8 +97,29 @@ private void refresh(ClusterConfig clusterConfig, InstanceConfig instanceConfig,
10197
_disabledPartitionsMap = instanceConfig.getDisabledPartitionsMap();
10298
_maxCapacity = instanceCapacity;
10399
_maxPartition = clusterConfig.getMaxPartitionsPerInstance();
100+
}
101+
102+
/**
103+
* This function should only be used to assign a set of new partitions that are not allocated on
104+
* this node.
105+
* Using this function avoids the overhead of updating capacity repeatedly.
106+
*/
107+
void assignNewBatch(Collection<AssignableReplica> replicas) {
108+
Map<String, Integer> totalPartitionCapacity = new HashMap<>();
109+
for (AssignableReplica replica : replicas) {
110+
addToAssignmentRecord(replica);
111+
// increment the capacity requirement according to partition's capacity configuration.
112+
for (Map.Entry<String, Integer> capacity : replica.getCapacity().entrySet()) {
113+
totalPartitionCapacity.compute(capacity.getKey(),
114+
(key, totalValue) -> (totalValue == null) ? capacity.getValue()
115+
: totalValue + capacity.getValue());
116+
}
117+
}
104118

105-
assignNewBatch(existingAssignment);
119+
// Update the global state after all single replications' calculation is done.
120+
for (String key : totalPartitionCapacity.keySet()) {
121+
updateCapacityAndUtilization(key, totalPartitionCapacity.get(key));
122+
}
106123
}
107124

108125
/**
@@ -314,29 +331,6 @@ private String computeFaultZone(ClusterConfig clusterConfig, InstanceConfig inst
314331
}
315332
}
316333

317-
/**
318-
* This function should only be used to assign a set of new partitions that are not allocated on
319-
* this node.
320-
* Using this function avoids the overhead of updating capacity repeatedly.
321-
*/
322-
private void assignNewBatch(Collection<AssignableReplica> replicas) {
323-
Map<String, Integer> totalPartitionCapacity = new HashMap<>();
324-
for (AssignableReplica replica : replicas) {
325-
addToAssignmentRecord(replica);
326-
// increment the capacity requirement according to partition's capacity configuration.
327-
for (Map.Entry<String, Integer> capacity : replica.getCapacity().entrySet()) {
328-
totalPartitionCapacity.compute(capacity.getKey(),
329-
(key, totalValue) -> (totalValue == null) ? capacity.getValue()
330-
: totalValue + capacity.getValue());
331-
}
332-
}
333-
334-
// Update the global state after all single replications' calculation is done.
335-
for (String key : totalPartitionCapacity.keySet()) {
336-
updateCapacityAndUtilization(key, totalPartitionCapacity.get(key));
337-
}
338-
}
339-
340334
/**
341335
* @throws HelixException if the replica has already been assigned to the node.
342336
*/

helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModelProvider.java

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,15 @@ public static ClusterModel generateClusterModel(ResourceControllerDataProvider d
6262
Map<HelixConstants.ChangeType, Set<String>> clusterChanges,
6363
Map<String, ResourceAssignment> baselineAssignment,
6464
Map<String, ResourceAssignment> bestPossibleAssignment) {
65+
// Construct all the assignable nodes and initialize with the allocated replicas.
66+
Set<AssignableNode> assignableNodes =
67+
parseAllNodes(dataProvider.getClusterConfig(), dataProvider.getInstanceConfigMap(),
68+
activeInstances);
69+
6570
// Generate replica objects for all the resource partitions.
6671
// <resource, replica set>
6772
Map<String, Set<AssignableReplica>> replicaMap =
68-
parseAllReplicas(dataProvider, resourceMap, activeInstances.size());
73+
parseAllReplicas(dataProvider, resourceMap, assignableNodes);
6974

7075
// Check if the replicas need to be reassigned.
7176
Map<String, Set<AssignableReplica>> allocatedReplicas =
@@ -74,10 +79,9 @@ public static ClusterModel generateClusterModel(ResourceControllerDataProvider d
7479
findToBeAssignedReplicas(replicaMap, clusterChanges, activeInstances,
7580
bestPossibleAssignment, allocatedReplicas);
7681

77-
// Construct all the assignable nodes and initialize with the allocated replicas.
78-
Set<AssignableNode> assignableNodes =
79-
parseAllNodes(dataProvider.getClusterConfig(), dataProvider.getInstanceConfigMap(),
80-
activeInstances, allocatedReplicas);
82+
// Update the allocated replicas to the assignable nodes.
83+
assignableNodes.stream().forEach(node -> node.assignNewBatch(
84+
allocatedReplicas.getOrDefault(node.getInstanceName(), Collections.emptySet())));
8185

8286
// Construct and initialize cluster context.
8387
ClusterContext context = new ClusterContext(
@@ -171,15 +175,13 @@ private static Set<AssignableReplica> findToBeAssignedReplicas(
171175
* @param clusterConfig The cluster configuration.
172176
* @param instanceConfigMap A map of all the instance configuration.
173177
* @param activeInstances All the instances that are online and enabled.
174-
* @param allocatedReplicas A map of all the assigned replicas, which will not be reassigned during the rebalance.
175178
* @return A map of assignable node set, <InstanceName, node set>.
176179
*/
177180
private static Set<AssignableNode> parseAllNodes(ClusterConfig clusterConfig,
178-
Map<String, InstanceConfig> instanceConfigMap, Set<String> activeInstances,
179-
Map<String, Set<AssignableReplica>> allocatedReplicas) {
181+
Map<String, InstanceConfig> instanceConfigMap, Set<String> activeInstances) {
180182
return activeInstances.stream().map(
181183
instanceName -> new AssignableNode(clusterConfig, instanceConfigMap.get(instanceName),
182-
instanceName, allocatedReplicas.getOrDefault(instanceName, Collections.emptySet())))
184+
instanceName))
183185
.collect(Collectors.toSet());
184186
}
185187

@@ -188,11 +190,12 @@ private static Set<AssignableNode> parseAllNodes(ClusterConfig clusterConfig,
188190
*
189191
* @param dataProvider The cluster status cache that contains the current cluster status.
190192
* @param resourceMap All the valid resources that are managed by the rebalancer.
193+
* @param assignableNodes All the active assignable nodes.
191194
* @return A map of assignable replica set, <ResourceName, replica set>.
192195
*/
193196
private static Map<String, Set<AssignableReplica>> parseAllReplicas(
194197
ResourceControllerDataProvider dataProvider, Map<String, Resource> resourceMap,
195-
int instanceCount) {
198+
Set<AssignableNode> assignableNodes) {
196199
Map<String, Set<AssignableReplica>> totalReplicaMap = new HashMap<>();
197200
ClusterConfig clusterConfig = dataProvider.getClusterConfig();
198201

@@ -211,8 +214,11 @@ private static Map<String, Set<AssignableReplica>> parseAllReplicas(
211214
is.getStateModelDefRef(), resourceName));
212215
}
213216

217+
int activeFaultZoneCount =
218+
assignableNodes.stream().map(node -> node.getFaultZone()).collect(Collectors.toSet())
219+
.size();
214220
Map<String, Integer> stateCountMap =
215-
def.getStateCountMap(instanceCount, is.getReplicaCount(instanceCount));
221+
def.getStateCountMap(activeFaultZoneCount, is.getReplicaCount(assignableNodes.size()));
216222

217223
for (String partition : is.getPartitionSet()) {
218224
for (Map.Entry<String, Integer> entry : stateCountMap.entrySet()) {

helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/ClusterModelTestHelper.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,7 @@ private Set<AssignableNode> generateNodes(ResourceControllerDataProvider testCac
4343
Set<AssignableNode> nodeSet = new HashSet<>();
4444
testCache.getInstanceConfigMap().values().stream()
4545
.forEach(config -> nodeSet.add(new AssignableNode(testCache.getClusterConfig(),
46-
testCache.getInstanceConfigMap().get(_testInstanceId), config.getInstanceName(),
47-
Collections.emptyList())));
46+
testCache.getInstanceConfigMap().get(_testInstanceId), config.getInstanceName())));
4847
return nodeSet;
4948
}
5049
}

helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestAssignableNode.java

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ public void testNormalUsage() throws IOException {
6464
expectedCapacityMap.put("item3", 30);
6565

6666
AssignableNode assignableNode = new AssignableNode(testCache.getClusterConfig(),
67-
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId, assignmentSet);
67+
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId);
68+
assignableNode.assignNewBatch(assignmentSet);
6869
Assert.assertEquals(assignableNode.getAssignedPartitionsMap(), expectedAssignment);
6970
Assert.assertEquals(assignableNode.getAssignedReplicaCount(), 4);
7071
Assert.assertEquals(assignableNode.getHighestCapacityUtilization(), 16.0 / 20.0, 0.005);
@@ -167,8 +168,7 @@ public void testReleaseNoPartition() throws IOException {
167168
ResourceControllerDataProvider testCache = setupClusterDataCache();
168169

169170
AssignableNode assignableNode = new AssignableNode(testCache.getClusterConfig(),
170-
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId,
171-
Collections.emptyList());
171+
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId);
172172
AssignableReplica removingReplica = new AssignableReplica(testCache.getClusterConfig(),
173173
testCache.getResourceConfig(_resourceNames.get(1)), _partitionNames.get(2) + "non-exist",
174174
"MASTER", 1);
@@ -183,7 +183,8 @@ public void testAssignDuplicateReplica() throws IOException {
183183
Set<AssignableReplica> assignmentSet = generateReplicas(testCache);
184184

185185
AssignableNode assignableNode = new AssignableNode(testCache.getClusterConfig(),
186-
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId, assignmentSet);
186+
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId);
187+
assignableNode.assignNewBatch(assignmentSet);
187188
AssignableReplica duplicateReplica = new AssignableReplica(testCache.getClusterConfig(),
188189
testCache.getResourceConfig(_resourceNames.get(0)), _partitionNames.get(0), "SLAVE", 2);
189190
assignableNode.assign(duplicateReplica);
@@ -206,8 +207,7 @@ public void testParseFaultZoneNotFound() throws IOException {
206207
when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
207208

208209
new AssignableNode(testCache.getClusterConfig(),
209-
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId,
210-
Collections.emptyList());
210+
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId);
211211
}
212212

213213
@Test
@@ -227,8 +227,7 @@ public void testParseFaultZone() throws IOException {
227227
when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
228228

229229
AssignableNode assignableNode = new AssignableNode(testCache.getClusterConfig(),
230-
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId,
231-
Collections.emptyList());
230+
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId);
232231

233232
Assert.assertEquals(assignableNode.getFaultZone(), "2/");
234233

@@ -245,8 +244,7 @@ public void testParseFaultZone() throws IOException {
245244
when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
246245

247246
assignableNode = new AssignableNode(testCache.getClusterConfig(),
248-
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId,
249-
Collections.emptyList());
247+
testCache.getInstanceConfigMap().get(_testInstanceId), _testInstanceId);
250248

251249
Assert.assertEquals(assignableNode.getFaultZone(), "2/testInstance/");
252250
}
@@ -259,8 +257,7 @@ public void testDefaultInstanceCapacity() {
259257
InstanceConfig testInstanceConfig = new InstanceConfig("testInstanceConfigId");
260258

261259
AssignableNode assignableNode =
262-
new AssignableNode(testClusterConfig, testInstanceConfig, _testInstanceId,
263-
Collections.emptyList());
260+
new AssignableNode(testClusterConfig, testInstanceConfig, _testInstanceId);
264261
Assert.assertEquals(assignableNode.getMaxCapacity(), _capacityDataMap);
265262
}
266263

@@ -274,7 +271,6 @@ public void testIncompleteInstanceCapacity() {
274271
InstanceConfig testInstanceConfig = new InstanceConfig("testInstanceConfigId");
275272
testInstanceConfig.setInstanceCapacityMap(_capacityDataMap);
276273

277-
new AssignableNode(testClusterConfig, testInstanceConfig, _testInstanceId,
278-
Collections.emptyList());
274+
new AssignableNode(testClusterConfig, testInstanceConfig, _testInstanceId);
279275
}
280276
}

helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestClusterModel.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,7 @@ Set<AssignableNode> generateNodes(ResourceControllerDataProvider testCache) {
4343
Set<AssignableNode> nodeSet = new HashSet<>();
4444
testCache.getInstanceConfigMap().values().stream().forEach(config -> nodeSet.add(
4545
new AssignableNode(testCache.getClusterConfig(),
46-
testCache.getInstanceConfigMap().get(_testInstanceId), config.getInstanceName(),
47-
Collections.emptyList())));
46+
testCache.getInstanceConfigMap().get(_testInstanceId), config.getInstanceName())));
4847
return nodeSet;
4948
}
5049

helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestClusterModelProvider.java

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,14 @@
1919
* under the License.
2020
*/
2121

22+
import java.io.IOException;
23+
import java.util.Collections;
24+
import java.util.HashMap;
25+
import java.util.HashSet;
26+
import java.util.Map;
27+
import java.util.Set;
28+
import java.util.stream.Collectors;
29+
2230
import org.apache.helix.HelixConstants;
2331
import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
2432
import org.apache.helix.controller.rebalancer.waged.WagedRebalancer;
@@ -34,14 +42,6 @@
3442
import org.testng.annotations.BeforeClass;
3543
import org.testng.annotations.Test;
3644

37-
import java.io.IOException;
38-
import java.util.Collections;
39-
import java.util.HashMap;
40-
import java.util.HashSet;
41-
import java.util.Map;
42-
import java.util.Set;
43-
import java.util.stream.Collectors;
44-
4545
import static org.mockito.Matchers.anyString;
4646
import static org.mockito.Mockito.when;
4747

@@ -111,7 +111,18 @@ public void testGenerateClusterModel() throws IOException {
111111
Assert.assertEquals(
112112
clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
113113
.collect(Collectors.toSet()), _instances);
114-
// Shall have 2 resources and 12 replicas
114+
// Shall have 2 resources and 4 replicas, since all nodes are in the same fault zone.
115+
Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
116+
Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
117+
.allMatch(replicaSet -> replicaSet.size() == 4));
118+
119+
// Adjust instance fault zone, so they have different fault zones.
120+
testCache.getInstanceConfigMap().values().stream()
121+
.forEach(config -> config.setZoneId(config.getInstanceName()));
122+
clusterModel = ClusterModelProvider.generateClusterModel(testCache, _resourceNames.stream()
123+
.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
124+
_instances, Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
125+
// Shall have 2 resources and 12 replicas after fault zone adjusted.
115126
Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
116127
Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
117128
.allMatch(replicaSet -> replicaSet.size() == 12));
@@ -197,10 +208,10 @@ public void testGenerateClusterModel() throws IOException {
197208
_instances, Collections.singletonMap(HelixConstants.ChangeType.RESOURCE_CONFIG,
198209
Collections.singleton(changedResourceName)), Collections.emptyMap(),
199210
bestPossibleAssignment);
200-
// There should be no existing assignment for all the resource except for resource2.
211+
// There should be no existing assignment for all the resource except for resource2
201212
Assert.assertEquals(clusterModel.getContext().getAssignmentForFaultZoneMap().size(), 1);
202213
Map<String, Set<String>> resourceAssignmentMap =
203-
clusterModel.getContext().getAssignmentForFaultZoneMap().get(_testFaultZoneId);
214+
clusterModel.getContext().getAssignmentForFaultZoneMap().get(_testInstanceId);
204215
// Should be only resource2 in the map
205216
Assert.assertEquals(resourceAssignmentMap.size(), 1);
206217
for (String resource : _resourceNames) {

0 commit comments

Comments
 (0)