helix-core/src/test/java/org/apache/helix/task/assigner/TestThreadCountBasedTaskAssigner.java - helix - Git at Google

 package org.apache.helix.task.assigner;

 /*
  * 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.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;

 import org.apache.helix.common.caches.TaskDataCache;
 import org.apache.helix.model.ClusterConfig;
 import org.apache.helix.model.InstanceConfig;
 import org.apache.helix.model.LiveInstance;
 import org.apache.helix.task.AssignableInstanceManager;
 import org.apache.helix.task.TaskConfig;
 import org.testng.Assert;
 import org.testng.annotations.Test;

 public class TestThreadCountBasedTaskAssigner extends AssignerTestBase {

   @Test
   public void testSuccessfulAssignment() {
     TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
     int taskCountPerType = 150;
     int instanceCount = 20;
     int threadCount = 50;
     AssignableInstanceManager assignableInstanceManager =
         createAssignableInstanceManager(instanceCount, threadCount);

     for (String quotaType : testQuotaTypes) {
       // Create tasks
       List<TaskConfig> tasks = createTaskConfigs(taskCountPerType);

       // Assign
       Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
           assignableInstanceManager.getAssignableInstanceNames(), tasks, quotaType);

       // Check success
       assertAssignmentResults(results.values(), true);

       // Check evenness
       for (AssignableInstance instance : assignableInstanceManager.getAssignableInstanceMap()
           .values()) {
         int assignedCount = instance.getUsedCapacity()
             .get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()).get(quotaType);
         Assert.assertTrue(assignedCount <= taskCountPerType / instanceCount + 1
             && assignedCount >= taskCountPerType / instanceCount);
       }
     }
   }

   @Test
   public void testAssignmentFailureNoInstance() {
     TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
     int taskCount = 10;
     List<TaskConfig> tasks = createTaskConfigs(taskCount);
     AssignableInstanceManager assignableInstanceManager = new AssignableInstanceManager();
     Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
         assignableInstanceManager.getAssignableInstanceNames(), tasks, "Dummy");
     Assert.assertEquals(results.size(), taskCount);
     for (TaskAssignResult result : results.values()) {
       Assert.assertFalse(result.isSuccessful());
       Assert.assertNull(result.getAssignableInstance());
       Assert.assertEquals(result.getFailureReason(),
           TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
     }
   }

   @Test
   public void testAssignmentFailureNoTask() {
     TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
     AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(1, 10);
     Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
         assignableInstanceManager.getAssignableInstanceNames(),
         Collections.<TaskConfig> emptyList(), AssignableInstance.DEFAULT_QUOTA_TYPE);
     Assert.assertTrue(results.isEmpty());
   }

   @Test
   public void testAssignmentFailureInsufficientQuota() {
     TaskAssigner assigner = new ThreadCountBasedTaskAssigner();

     // 10 * Type1 quota
     AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(2, 10);
     List<TaskConfig> tasks = createTaskConfigs(20);

     Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
         assignableInstanceManager.getAssignableInstanceNames(), tasks, testQuotaTypes[0]);
     int successCnt = 0;
     int failCnt = 0;
     for (TaskAssignResult rst : results.values()) {
       if (rst.isSuccessful()) {
         successCnt += 1;
       } else {
         failCnt += 1;
         Assert.assertEquals(rst.getFailureReason(),
             TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
       }
     }
     Assert.assertEquals(successCnt, 10);
     Assert.assertEquals(failCnt, 10);
   }

   @Test
   public void testAssignmentFailureDuplicatedTask() {
     TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
     AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(1, 20);
     List<TaskConfig> tasks = createTaskConfigs(10, false);

     // Duplicate all tasks
     tasks.addAll(createTaskConfigs(10, false));
     Collections.shuffle(tasks);

     Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
         assignableInstanceManager.getAssignableInstanceNames(), tasks, testQuotaTypes[0]);
     Assert.assertEquals(results.size(), 10);
     assertAssignmentResults(results.values(), true);
   }

   @Test(enabled = false, description = "Not enabling profiling tests")
   public void testAssignerProfiling() {
     int instanceCount = 1000;
     int taskCount = 50000;
     for (int batchSize : new int[] {10000, 5000, 2000, 1000, 500, 100}) {
       System.out.println("testing batch size: " + batchSize);
       profileAssigner(batchSize, instanceCount, taskCount);
     }
   }

   @Test
   public void testAssignmentToGivenInstances() {
     int totalNumberOfInstances = 10;
     int eligibleNumberOfInstances = 5;
     String instanceNameFormat = "instance-%s";

     TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
     AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(10, 20);
     List<TaskConfig> tasks = createTaskConfigs(100, false);
     Set<String> eligibleInstances = new HashSet<>();

     // Add only eligible number of instances
     for (int i = 0; i < eligibleNumberOfInstances; i++) {
       eligibleInstances.add(String.format(instanceNameFormat, i));
     }

     Map<String, TaskAssignResult> result = assigner.assignTasks(assignableInstanceManager,
         eligibleInstances, tasks, testQuotaTypes[0]);

     for (int i = 0; i < totalNumberOfInstances; i++) {
       String instance = String.format(instanceNameFormat, i);
       Set<String> test = assignableInstanceManager.getAssignableInstance(instance)
           .getCurrentAssignments();
       boolean isAssignmentEmpty = assignableInstanceManager.getAssignableInstance(instance)
           .getCurrentAssignments().isEmpty();
       // Check that assignment only took place to eligible number of instances and that assignment
       // did not happen to non-eligible AssignableInstances
       if (i < eligibleNumberOfInstances) {
         // Must have tasks assigned to these instances
         Assert.assertFalse(isAssignmentEmpty);
       } else {
         // These instances should have no tasks assigned to them
         Assert.assertTrue(isAssignmentEmpty);
       }
     }
   }

   private void profileAssigner(int assignBatchSize, int instanceCount, int taskCount) {
     int trail = 100;
     long totalTime = 0;
     for (int i = 0; i < trail; i++) {
       TaskAssigner assigner = new ThreadCountBasedTaskAssigner();

       // 50 * instanceCount number of tasks
       AssignableInstanceManager assignableInstanceManager =
           createAssignableInstanceManager(instanceCount, 100);
       List<TaskConfig> tasks = createTaskConfigs(taskCount);
       List<Map<String, TaskAssignResult>> allResults = new ArrayList<>();

       // Assign
       long start = System.currentTimeMillis();
       for (int j = 0; j < taskCount / assignBatchSize; j++) {
         allResults.add(assigner.assignTasks(assignableInstanceManager,
             assignableInstanceManager.getAssignableInstanceNames(),
             tasks.subList(j * assignBatchSize, (j + 1) * assignBatchSize), testQuotaTypes[0]));
       }
       long duration = System.currentTimeMillis() - start;
       totalTime += duration;

       // Validate
       for (Map<String, TaskAssignResult> results : allResults) {
         for (TaskAssignResult rst : results.values()) {
           Assert.assertTrue(rst.isSuccessful());
         }
       }
     }
     System.out.println("Average time: " + totalTime / trail + "ms");
   }

   private void assertAssignmentResults(Iterable<TaskAssignResult> results, boolean expected) {
     for (TaskAssignResult rst : results) {
       Assert.assertEquals(rst.isSuccessful(), expected);
     }
   }

   private List<TaskConfig> createTaskConfigs(int count) {
     return createTaskConfigs(count, true);
   }

   private List<TaskConfig> createTaskConfigs(int count, boolean randomID) {
     List<TaskConfig> tasks = new ArrayList<>();
     for (int i = 0; i < count; i++) {
       TaskConfig task =
           new TaskConfig(null, null, randomID ? UUID.randomUUID().toString() : "task-" + i, null);
       tasks.add(task);
     }
     return tasks;
   }

   private AssignableInstanceManager createAssignableInstanceManager(int count, int threadCount) {
     AssignableInstanceManager assignableInstanceManager = new AssignableInstanceManager();
     ClusterConfig clusterConfig = createClusterConfig(testQuotaTypes, testQuotaRatio, false);
     String instanceNameFormat = "instance-%s";
     Map<String, LiveInstance> liveInstanceMap = new HashMap<>();
     Map<String, InstanceConfig> instanceConfigMap = new HashMap<>();
     for (int i = 0; i < count; i++) {
       String instanceName = String.format(instanceNameFormat, i);
       liveInstanceMap.put(instanceName, createLiveInstance(
           new String[] { LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name() },
           new String[] { Integer.toString(threadCount) }, instanceName));
       instanceConfigMap.put(instanceName, new InstanceConfig(instanceName));
     }

     assignableInstanceManager
         .buildAssignableInstances(clusterConfig, new TaskDataCache(testClusterName),
             liveInstanceMap, instanceConfigMap);
     return assignableInstanceManager;
   }
 }
	package org.apache.helix.task.assigner;

	/*
	* 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.HashMap;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.UUID;

	import org.apache.helix.common.caches.TaskDataCache;
	import org.apache.helix.model.ClusterConfig;
	import org.apache.helix.model.InstanceConfig;
	import org.apache.helix.model.LiveInstance;
	import org.apache.helix.task.AssignableInstanceManager;
	import org.apache.helix.task.TaskConfig;
	import org.testng.Assert;
	import org.testng.annotations.Test;

	public class TestThreadCountBasedTaskAssigner extends AssignerTestBase {

	@Test
	public void testSuccessfulAssignment() {
	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
	int taskCountPerType = 150;
	int instanceCount = 20;
	int threadCount = 50;
	AssignableInstanceManager assignableInstanceManager =
	createAssignableInstanceManager(instanceCount, threadCount);

	for (String quotaType : testQuotaTypes) {
	// Create tasks
	List<TaskConfig> tasks = createTaskConfigs(taskCountPerType);

	// Assign
	Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
	assignableInstanceManager.getAssignableInstanceNames(), tasks, quotaType);

	// Check success
	assertAssignmentResults(results.values(), true);

	// Check evenness
	for (AssignableInstance instance : assignableInstanceManager.getAssignableInstanceMap()
	.values()) {
	int assignedCount = instance.getUsedCapacity()
	.get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()).get(quotaType);
	Assert.assertTrue(assignedCount <= taskCountPerType / instanceCount + 1
	&& assignedCount >= taskCountPerType / instanceCount);
	}
	}
	}

	@Test
	public void testAssignmentFailureNoInstance() {
	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
	int taskCount = 10;
	List<TaskConfig> tasks = createTaskConfigs(taskCount);
	AssignableInstanceManager assignableInstanceManager = new AssignableInstanceManager();
	Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
	assignableInstanceManager.getAssignableInstanceNames(), tasks, "Dummy");
	Assert.assertEquals(results.size(), taskCount);
	for (TaskAssignResult result : results.values()) {
	Assert.assertFalse(result.isSuccessful());
	Assert.assertNull(result.getAssignableInstance());
	Assert.assertEquals(result.getFailureReason(),
	TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
	}
	}

	@Test
	public void testAssignmentFailureNoTask() {
	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
	AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(1, 10);
	Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
	assignableInstanceManager.getAssignableInstanceNames(),
	Collections.<TaskConfig> emptyList(), AssignableInstance.DEFAULT_QUOTA_TYPE);
	Assert.assertTrue(results.isEmpty());
	}

	@Test
	public void testAssignmentFailureInsufficientQuota() {
	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();

	// 10 * Type1 quota
	AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(2, 10);
	List<TaskConfig> tasks = createTaskConfigs(20);

	Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
	assignableInstanceManager.getAssignableInstanceNames(), tasks, testQuotaTypes[0]);
	int successCnt = 0;
	int failCnt = 0;
	for (TaskAssignResult rst : results.values()) {
	if (rst.isSuccessful()) {
	successCnt += 1;
	} else {
	failCnt += 1;
	Assert.assertEquals(rst.getFailureReason(),
	TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
	}
	}
	Assert.assertEquals(successCnt, 10);
	Assert.assertEquals(failCnt, 10);
	}

	@Test
	public void testAssignmentFailureDuplicatedTask() {
	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
	AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(1, 20);
	List<TaskConfig> tasks = createTaskConfigs(10, false);

	// Duplicate all tasks
	tasks.addAll(createTaskConfigs(10, false));
	Collections.shuffle(tasks);

	Map<String, TaskAssignResult> results = assigner.assignTasks(assignableInstanceManager,
	assignableInstanceManager.getAssignableInstanceNames(), tasks, testQuotaTypes[0]);
	Assert.assertEquals(results.size(), 10);
	assertAssignmentResults(results.values(), true);
	}

	@Test(enabled = false, description = "Not enabling profiling tests")
	public void testAssignerProfiling() {
	int instanceCount = 1000;
	int taskCount = 50000;
	for (int batchSize : new int[] {10000, 5000, 2000, 1000, 500, 100}) {
	System.out.println("testing batch size: " + batchSize);
	profileAssigner(batchSize, instanceCount, taskCount);
	}
	}

	@Test
	public void testAssignmentToGivenInstances() {
	int totalNumberOfInstances = 10;
	int eligibleNumberOfInstances = 5;
	String instanceNameFormat = "instance-%s";

	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();
	AssignableInstanceManager assignableInstanceManager = createAssignableInstanceManager(10, 20);
	List<TaskConfig> tasks = createTaskConfigs(100, false);
	Set<String> eligibleInstances = new HashSet<>();

	// Add only eligible number of instances
	for (int i = 0; i < eligibleNumberOfInstances; i++) {
	eligibleInstances.add(String.format(instanceNameFormat, i));
	}

	Map<String, TaskAssignResult> result = assigner.assignTasks(assignableInstanceManager,
	eligibleInstances, tasks, testQuotaTypes[0]);

	for (int i = 0; i < totalNumberOfInstances; i++) {
	String instance = String.format(instanceNameFormat, i);
	Set<String> test = assignableInstanceManager.getAssignableInstance(instance)
	.getCurrentAssignments();
	boolean isAssignmentEmpty = assignableInstanceManager.getAssignableInstance(instance)
	.getCurrentAssignments().isEmpty();
	// Check that assignment only took place to eligible number of instances and that assignment
	// did not happen to non-eligible AssignableInstances
	if (i < eligibleNumberOfInstances) {
	// Must have tasks assigned to these instances
	Assert.assertFalse(isAssignmentEmpty);
	} else {
	// These instances should have no tasks assigned to them
	Assert.assertTrue(isAssignmentEmpty);
	}
	}
	}

	private void profileAssigner(int assignBatchSize, int instanceCount, int taskCount) {
	int trail = 100;
	long totalTime = 0;
	for (int i = 0; i < trail; i++) {
	TaskAssigner assigner = new ThreadCountBasedTaskAssigner();

	// 50 * instanceCount number of tasks
	AssignableInstanceManager assignableInstanceManager =
	createAssignableInstanceManager(instanceCount, 100);
	List<TaskConfig> tasks = createTaskConfigs(taskCount);
	List<Map<String, TaskAssignResult>> allResults = new ArrayList<>();

	// Assign
	long start = System.currentTimeMillis();
	for (int j = 0; j < taskCount / assignBatchSize; j++) {
	allResults.add(assigner.assignTasks(assignableInstanceManager,
	assignableInstanceManager.getAssignableInstanceNames(),
	tasks.subList(j * assignBatchSize, (j + 1) * assignBatchSize), testQuotaTypes[0]));
	}
	long duration = System.currentTimeMillis() - start;
	totalTime += duration;

	// Validate
	for (Map<String, TaskAssignResult> results : allResults) {
	for (TaskAssignResult rst : results.values()) {
	Assert.assertTrue(rst.isSuccessful());
	}
	}
	}
	System.out.println("Average time: " + totalTime / trail + "ms");
	}

	private void assertAssignmentResults(Iterable<TaskAssignResult> results, boolean expected) {
	for (TaskAssignResult rst : results) {
	Assert.assertEquals(rst.isSuccessful(), expected);
	}
	}

	private List<TaskConfig> createTaskConfigs(int count) {
	return createTaskConfigs(count, true);
	}

	private List<TaskConfig> createTaskConfigs(int count, boolean randomID) {
	List<TaskConfig> tasks = new ArrayList<>();
	for (int i = 0; i < count; i++) {
	TaskConfig task =
	new TaskConfig(null, null, randomID ? UUID.randomUUID().toString() : "task-" + i, null);
	tasks.add(task);
	}
	return tasks;
	}

	private AssignableInstanceManager createAssignableInstanceManager(int count, int threadCount) {
	AssignableInstanceManager assignableInstanceManager = new AssignableInstanceManager();
	ClusterConfig clusterConfig = createClusterConfig(testQuotaTypes, testQuotaRatio, false);
	String instanceNameFormat = "instance-%s";
	Map<String, LiveInstance> liveInstanceMap = new HashMap<>();
	Map<String, InstanceConfig> instanceConfigMap = new HashMap<>();
	for (int i = 0; i < count; i++) {
	String instanceName = String.format(instanceNameFormat, i);
	liveInstanceMap.put(instanceName, createLiveInstance(
	new String[] { LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name() },
	new String[] { Integer.toString(threadCount) }, instanceName));
	instanceConfigMap.put(instanceName, new InstanceConfig(instanceName));
	}

	assignableInstanceManager
	.buildAssignableInstances(clusterConfig, new TaskDataCache(testClusterName),
	liveInstanceMap, instanceConfigMap);
	return assignableInstanceManager;
	}
	}