/**
* 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
* <p/>
* http://www.apache.org/licenses/LICENSE-2.0
* <p/>
* 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.
*/
package org.apache.drill.exec.planner.fragment;
import com.google.common.collect.HashMultiset;
import com.google.common.collect.ImmutableList;
import mockit.Mocked;
import mockit.NonStrictExpectations;
import org.apache.drill.exec.physical.EndpointAffinity;
import org.apache.drill.exec.physical.base.PhysicalOperator;
import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
import org.junit.Test;
import java.util.Collections;
import java.util.List;
import static java.lang.Integer.MAX_VALUE;
import static org.apache.drill.exec.ExecConstants.SLICE_TARGET_DEFAULT;
import static org.apache.drill.exec.planner.fragment.HardAffinityFragmentParallelizer.INSTANCE;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
public class TestHardAffinityFragmentParallelizer {
// Create a set of test endpoints
private static final DrillbitEndpoint N1_EP1 = newDrillbitEndpoint("node1", 30010);
private static final DrillbitEndpoint N1_EP2 = newDrillbitEndpoint("node1", 30011);
private static final DrillbitEndpoint N2_EP1 = newDrillbitEndpoint("node2", 30010);
private static final DrillbitEndpoint N2_EP2 = newDrillbitEndpoint("node2", 30011);
private static final DrillbitEndpoint N3_EP1 = newDrillbitEndpoint("node3", 30010);
private static final DrillbitEndpoint N3_EP2 = newDrillbitEndpoint("node3", 30011);
private static final DrillbitEndpoint N4_EP2 = newDrillbitEndpoint("node4", 30011);
@Mocked private Fragment fragment;
@Mocked private PhysicalOperator root;
private static final DrillbitEndpoint newDrillbitEndpoint(String address, int port) {
return DrillbitEndpoint.newBuilder().setAddress(address).setControlPort(port).build();
}
private static final ParallelizationParameters newParameters(final long threshold, final int maxWidthPerNode,
final int maxGlobalWidth) {
return new ParallelizationParameters() {
@Override
public long getSliceTarget() {
return threshold;
}
@Override
public int getMaxWidthPerNode() {
return maxWidthPerNode;
}
@Override
public int getMaxGlobalWidth() {
return maxGlobalWidth;
}
/**
* {@link HardAffinityFragmentParallelizer} doesn't use affinity factor.
* @return
*/
@Override
public double getAffinityFactor() {
return 0.0f;
}
};
}
private final Wrapper newWrapper(double cost, int minWidth, int maxWidth, List<EndpointAffinity> epAffs) {
new NonStrictExpectations() {
{
fragment.getRoot(); result = root;
}
};
final Wrapper fragmentWrapper = new Wrapper(fragment, 1);
final Stats stats = fragmentWrapper.getStats();
stats.setDistributionAffinity(DistributionAffinity.HARD);
stats.addCost(cost);
stats.addMinWidth(minWidth);
stats.addMaxWidth(maxWidth);
stats.addEndpointAffinities(epAffs);
return fragmentWrapper;
}
@Test
public void simpleCase1() throws Exception {
final Wrapper wrapper = newWrapper(200, 1, 20, Collections.singletonList(new EndpointAffinity(N1_EP1, 1.0, true, MAX_VALUE)));
INSTANCE.parallelizeFragment(wrapper, newParameters(SLICE_TARGET_DEFAULT, 5, 20), null);
// Expect the fragment parallelization to be just one because:
// The cost (200) is below the threshold (SLICE_TARGET_DEFAULT) (which gives width of 200/10000 = ~1) and
assertEquals(1, wrapper.getWidth());
final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
assertEquals(1, assignedEps.size());
assertEquals(N1_EP1, assignedEps.get(0));
}
@Test
public void simpleCase2() throws Exception {
// Set the slice target to 1
final Wrapper wrapper = newWrapper(200, 1, 20, Collections.singletonList(new EndpointAffinity(N1_EP1, 1.0, true, MAX_VALUE)));
INSTANCE.parallelizeFragment(wrapper, newParameters(1, 5, 20), null);
// Expect the fragment parallelization to be 5:
// 1. the cost (200) is above the threshold (SLICE_TARGET_DEFAULT) (which gives 200/1=200 width) and
// 2. Max width per node is 5 (limits the width 200 to 5)
assertEquals(5, wrapper.getWidth());
final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
assertEquals(5, assignedEps.size());
for (DrillbitEndpoint ep : assignedEps) {
assertEquals(N1_EP1, ep);
}
}
@Test
public void multiNodeCluster1() throws Exception {
final Wrapper wrapper = newWrapper(200, 1, 20,
ImmutableList.of(
new EndpointAffinity(N1_EP1, 0.15, true, MAX_VALUE),
new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N2_EP1, 0.10, true, MAX_VALUE),
new EndpointAffinity(N3_EP2, 0.20, true, MAX_VALUE),
new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE)
));
INSTANCE.parallelizeFragment(wrapper, newParameters(SLICE_TARGET_DEFAULT, 5, 20), null);
// Expect the fragment parallelization to be 5 because:
// 1. The cost (200) is below the threshold (SLICE_TARGET_DEFAULT) (which gives width of 200/10000 = ~1) and
// 2. Number of mandoatory node assignments are 5 which overrides the cost based width of 1.
assertEquals(5, wrapper.getWidth());
// As there are 5 required eps and the width is 5, everyone gets assigned 1.
final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
assertEquals(5, assignedEps.size());
assertTrue(assignedEps.contains(N1_EP1));
assertTrue(assignedEps.contains(N1_EP2));
assertTrue(assignedEps.contains(N2_EP1));
assertTrue(assignedEps.contains(N3_EP2));
assertTrue(assignedEps.contains(N4_EP2));
}
@Test
public void multiNodeCluster2() throws Exception {
final Wrapper wrapper = newWrapper(200, 1, 20,
ImmutableList.of(
new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
));
INSTANCE.parallelizeFragment(wrapper, newParameters(1, 5, 20), null);
// Expect the fragment parallelization to be 20 because:
// 1. the cost (200) is above the threshold (SLICE_TARGET_DEFAULT) (which gives 200/1=200 width) and
// 2. Number of mandatory node assignments are 5 (current width 200 satisfies the requirement)
// 3. max fragment width is 20 which limits the width
assertEquals(20, wrapper.getWidth());
final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
assertEquals(20, assignedEps.size());
final HashMultiset<DrillbitEndpoint> counts = HashMultiset.create();
for(final DrillbitEndpoint ep : assignedEps) {
counts.add(ep);
}
// Each node gets at max 5.
assertTrue(counts.count(N1_EP2) <= 5);
assertTrue(counts.count(N2_EP2) <= 5);
assertTrue(counts.count(N3_EP1) <= 5);
assertTrue(counts.count(N4_EP2) <= 5);
assertTrue(counts.count(N1_EP1) <= 5);
}
@Test
public void multiNodeClusterNegative1() throws Exception {
final Wrapper wrapper = newWrapper(200, 1, 20,
ImmutableList.of(
new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
));
try {
INSTANCE.parallelizeFragment(wrapper, newParameters(1, 2, 2), null);
fail("Expected an exception, because max global query width (2) is less than the number of mandatory nodes (5)");
} catch (Exception e) {
// ok
}
}
@Test
public void multiNodeClusterNegative2() throws Exception {
final Wrapper wrapper = newWrapper(200, 1, 3,
ImmutableList.of(
new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
));
try {
INSTANCE.parallelizeFragment(wrapper, newParameters(1, 2, 2), null);
fail("Expected an exception, because max fragment width (3) is less than the number of mandatory nodes (5)");
} catch (Exception e) {
// ok
}
}
}