Skip to content

Commit

Permalink
proposal
Browse files Browse the repository at this point in the history
  • Loading branch information
GrantPSpencer committed Feb 6, 2025
1 parent 33a28e7 commit 87466be
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
Expand All @@ -31,7 +30,6 @@
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Maps;
import org.apache.helix.HelixRebalanceException;
import org.apache.helix.controller.rebalancer.waged.RebalanceAlgorithm;
Expand Down Expand Up @@ -59,11 +57,13 @@ class ConstraintBasedAlgorithm implements RebalanceAlgorithm {
private static final Logger LOG = LoggerFactory.getLogger(ConstraintBasedAlgorithm.class);
private final List<HardConstraint> _hardConstraints;
private final Map<SoftConstraint, Float> _softConstraints;
private boolean _enabledConstraintLogging;

ConstraintBasedAlgorithm(List<HardConstraint> hardConstraints,
Map<SoftConstraint, Float> softConstraints) {
_hardConstraints = hardConstraints;
_softConstraints = softConstraints;
_enabledConstraintLogging = false;
}

@Override
Expand Down Expand Up @@ -128,7 +128,7 @@ private Optional<AssignableNode> getNodeWithHighestPoints(AssignableReplica repl
// need to record all the failure reasons and it gives us the ability to debug/fix the runtime
// cluster environment
for (HardConstraint hardConstraint : _hardConstraints) {
if (!hardConstraint.isAssignmentValid(candidateNode, replica, clusterContext)) {
if (!hardConstraint.isAssignmentValid(candidateNode, replica, clusterContext, _enabledConstraintLogging)) {
hardConstraintFailures.computeIfAbsent(candidateNode, node -> new ArrayList<>())
.add(hardConstraint);
isValid = false;
Expand All @@ -138,15 +138,12 @@ private Optional<AssignableNode> getNodeWithHighestPoints(AssignableReplica repl
}).collect(Collectors.toList());

if (candidateNodes.isEmpty()) {
LOG.info("Found no eligible candidate nodes. Enabling hard constraint level logging for cluster: {}", clusterContext.getClusterName());
enableFullLoggingForCluster();
setConstraintLevelLoggingEnabled(true, clusterContext.getClusterName());
optimalAssignment.recordAssignmentFailure(replica,
Maps.transformValues(hardConstraintFailures, this::convertFailureReasons));
return Optional.empty();
}

LOG.info("Disabling hard constraint level logging for cluster: {}", clusterContext.getClusterName());
removeFullLoggingForCluster();
setConstraintLevelLoggingEnabled(false, clusterContext.getClusterName());

return candidateNodes.parallelStream().map(node -> new HashMap.SimpleEntry<>(node,
getAssignmentNormalizedScore(node, replica, clusterContext)))
Expand Down Expand Up @@ -187,20 +184,13 @@ private List<String> convertFailureReasons(List<HardConstraint> hardConstraints)
}

/**
* Enables logging of failures in all hard constraints
*/
private void enableFullLoggingForCluster() {
for (HardConstraint hardConstraint : _hardConstraints) {
hardConstraint.setEnableLogging(true);
}
}

/**
* Removes the cluster from full logging in all hard constraints (if added previously)
* Enables logging of failures for all hard constraints
*/
private void removeFullLoggingForCluster() {
for (HardConstraint hardConstraint : _hardConstraints) {
hardConstraint.setEnableLogging(false);
private void setConstraintLevelLoggingEnabled(boolean enabled, String clusterName) {
if (_enabledConstraintLogging != enabled) {
LOG.info("Logging of hard constraint is now: {}, triggered by cluster {}",
enabled ? "enabled" : "disabled", clusterName);
_enabledConstraintLogging = enabled;
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,12 @@ abstract class HardConstraint {
abstract boolean isAssignmentValid(AssignableNode node, AssignableReplica replica,
ClusterContext clusterContext);

boolean isAssignmentValid(AssignableNode node, AssignableReplica replica,
ClusterContext clusterContext, boolean loggingEnabled) {
enableLogging = loggingEnabled;
return isAssignmentValid(node, replica, clusterContext);
}

/**
* Return class name by default as description if it's explanatory enough, child class could override
* the method and add more detailed descriptions
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,18 +27,19 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.apache.helix.HelixRebalanceException;
import org.apache.helix.controller.rebalancer.waged.model.AssignableNode;
import org.apache.helix.controller.rebalancer.waged.model.AssignableReplica;
import org.apache.helix.controller.rebalancer.waged.model.ClusterContext;
import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
import org.apache.helix.controller.rebalancer.waged.model.ClusterModelTestHelper;
import org.apache.helix.controller.rebalancer.waged.model.OptimalAssignment;
import org.testng.Assert;
import org.testng.annotations.Test;

import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
Expand All @@ -50,7 +51,8 @@ public class TestConstraintBasedAlgorithm {
public void testCalculateNoValidAssignment() throws IOException {
HardConstraint mockHardConstraint = mock(HardConstraint.class);
SoftConstraint mockSoftConstraint = mock(SoftConstraint.class);
when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(false);
when(mockHardConstraint.isAssignmentValid(any(AssignableNode.class), any(AssignableReplica.class), any(
ClusterContext.class), anyBoolean())).thenReturn(false);
when(mockSoftConstraint.getAssignmentNormalizedScore(any(), any(), any())).thenReturn(1.0);
ConstraintBasedAlgorithm algorithm =
new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint),
Expand All @@ -62,8 +64,8 @@ public void testCalculateNoValidAssignment() throws IOException {
Assert.assertEquals(ex.getFailureType(), HelixRebalanceException.Type.FAILED_TO_CALCULATE);
}

verify(mockHardConstraint, times(1)).setEnableLogging(eq(true));
verify(mockHardConstraint, times(1)).isAssignmentValid(any(), any(), any());
verify(mockHardConstraint, times(1)).isAssignmentValid(any(AssignableNode.class), any(AssignableReplica.class), any(
ClusterContext.class), eq(false));
}

@Test
Expand All @@ -84,19 +86,25 @@ public void testCalculateNoValidAssignmentFirstAndThenRecovery() throws IOExcept
Assert.assertEquals(ex.getFailureType(), HelixRebalanceException.Type.FAILED_TO_CALCULATE);
}

verify(mockHardConstraint, times(1)).setEnableLogging(eq(true));
verify(mockHardConstraint, times(1)).isAssignmentValid(any(), any(), any());
verify(mockHardConstraint, times(1)).isAssignmentValid(any(AssignableNode.class), any(AssignableReplica.class), any(
ClusterContext.class), eq(false));

// calling again for recovery (no exception)
algorithm.calculate(clusterModel);
verify(mockHardConstraint, atLeastOnce()).setEnableLogging(eq(false));
try {
algorithm.calculate(clusterModel);
} catch (HelixRebalanceException ex) {
Assert.assertEquals(ex.getFailureType(), HelixRebalanceException.Type.FAILED_TO_CALCULATE);
}
verify(mockHardConstraint, times(1)).isAssignmentValid(any(AssignableNode.class), any(AssignableReplica.class), any(
ClusterContext.class), eq(true));
}

@Test
public void testCalculateWithValidAssignment() throws IOException, HelixRebalanceException {
HardConstraint mockHardConstraint = mock(HardConstraint.class);
SoftConstraint mockSoftConstraint = mock(SoftConstraint.class);
when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(true);
when(mockHardConstraint.isAssignmentValid(any(AssignableNode.class), any(AssignableReplica.class), any(
ClusterContext.class), anyBoolean())).thenReturn(true);
when(mockSoftConstraint.getAssignmentNormalizedScore(any(), any(), any())).thenReturn(1.0);
ConstraintBasedAlgorithm algorithm =
new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint),
Expand All @@ -111,7 +119,8 @@ public void testCalculateWithValidAssignment() throws IOException, HelixRebalanc
public void testCalculateScoreDeterminism() throws IOException, HelixRebalanceException {
HardConstraint mockHardConstraint = mock(HardConstraint.class);
SoftConstraint mockSoftConstraint = mock(SoftConstraint.class);
when(mockHardConstraint.isAssignmentValid(any(), any(), any())).thenReturn(true);
when(mockHardConstraint.isAssignmentValid(any(AssignableNode.class), any(AssignableReplica.class), any(
ClusterContext.class), anyBoolean())).thenReturn(true);
when(mockSoftConstraint.getAssignmentNormalizedScore(any(), any(), any())).thenReturn(1.0);
ConstraintBasedAlgorithm algorithm =
new ConstraintBasedAlgorithm(ImmutableList.of(mockHardConstraint),
Expand Down

0 comments on commit 87466be

Please sign in to comment.