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

 package org.apache.helix.task;

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

 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 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.assigner.AssignableInstance;
 import org.apache.helix.task.assigner.TaskAssignResult;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;

 public class TestAssignableInstanceManager {
   private static final int NUM_PARTICIPANTS = 3;
   private static final int NUM_JOBS = 3;
   private static final int NUM_TASKS = 3;
   private static final String CLUSTER_NAME = "TestCluster_0";
   private static final String INSTANCE_PREFIX = "Instance_";
   private static final String JOB_PREFIX = "Job_";
   private static final String TASK_PREFIX = "Task_";

   private ClusterConfig _clusterConfig;
   private MockTaskDataCache _taskDataCache;
   private AssignableInstanceManager _assignableInstanceManager;
   private Map<String, LiveInstance> _liveInstances;
   private Map<String, InstanceConfig> _instanceConfigs;
   private Set<String> _taskIDs; // To keep track of what tasks were created

   @BeforeClass
   public void beforeClass() {
     System.out.println(
         "START " + this.getClass().getSimpleName() + " at " + new Date(System.currentTimeMillis()));
     _clusterConfig = new ClusterConfig(CLUSTER_NAME);
     _taskDataCache = new MockTaskDataCache(CLUSTER_NAME);
     _liveInstances = new HashMap<>();
     _instanceConfigs = new HashMap<>();
     _taskIDs = new HashSet<>();

     // Populate live instances and their corresponding instance configs
     for (int i = 0; i < NUM_PARTICIPANTS; i++) {
       String instanceName = INSTANCE_PREFIX + i;
       LiveInstance liveInstance = new LiveInstance(instanceName);
       InstanceConfig instanceConfig = new InstanceConfig(instanceName);
       _liveInstances.put(instanceName, liveInstance);
       _instanceConfigs.put(instanceName, instanceConfig);
     }

     // Populate taskDataCache with JobConfigs and JobContexts
     for (int i = 0; i < NUM_JOBS; i++) {
       String jobName = JOB_PREFIX + i;

       // Create a JobConfig
       JobConfig.Builder jobConfigBuilder = new JobConfig.Builder();
       List<TaskConfig> taskConfigs = new ArrayList<>();
       for (int j = 0; j < NUM_TASKS; j++) {
         String taskID = jobName + "_" + TASK_PREFIX + j;
         TaskConfig.Builder taskConfigBuilder = new TaskConfig.Builder();
         taskConfigBuilder.setTaskId(taskID);
         _taskIDs.add(taskID);
         taskConfigs.add(taskConfigBuilder.build());
       }

       jobConfigBuilder.setJobId(jobName);
       jobConfigBuilder.addTaskConfigs(taskConfigs);
       jobConfigBuilder.setCommand("MOCK");
       jobConfigBuilder.setWorkflow("WORKFLOW");
       _taskDataCache.addJobConfig(jobName, jobConfigBuilder.build());

       // Create a JobContext
       ZNRecord znRecord = new ZNRecord(JOB_PREFIX + "context_" + i);
       JobContext jobContext = new MockJobContext(znRecord, _liveInstances, _taskIDs);
       _taskDataCache.addJobContext(jobName, jobContext);
       _taskIDs.clear();
     }

     // Create an AssignableInstanceManager and build
     _assignableInstanceManager = new AssignableInstanceManager();
     _assignableInstanceManager.buildAssignableInstances(_clusterConfig, _taskDataCache,
         _liveInstances, _instanceConfigs);
   }

   @Test
   public void testGetAssignableInstanceMap() {
     Map<String, AssignableInstance> assignableInstanceMap =
         _assignableInstanceManager.getAssignableInstanceMap();
     for (String liveInstance : _liveInstances.keySet()) {
       Assert.assertTrue(assignableInstanceMap.containsKey(liveInstance));
     }
   }

   @Test(dependsOnMethods = "testGetAssignableInstanceMap")
   public void testGetTaskAssignResultMap() {
     Map<String, TaskAssignResult> taskAssignResultMap =
         _assignableInstanceManager.getTaskAssignResultMap();
     for (String taskID : _taskIDs) {
       Assert.assertTrue(taskAssignResultMap.containsKey(taskID));
     }
   }

   @Test(dependsOnMethods = "testGetTaskAssignResultMap")
   public void testUpdateAssignableInstances() {
     Map<String, LiveInstance> newLiveInstances = new HashMap<>();
     Map<String, InstanceConfig> newInstanceConfigs = new HashMap<>();

     // A brand new set of LiveInstances
     for (int i = NUM_PARTICIPANTS; i < NUM_PARTICIPANTS + 3; i++) {
       String instanceName = INSTANCE_PREFIX + i;
       newLiveInstances.put(instanceName, new LiveInstance(instanceName));
       newInstanceConfigs.put(instanceName, new InstanceConfig(instanceName));
     }

     _assignableInstanceManager.updateAssignableInstances(_clusterConfig, newLiveInstances,
         newInstanceConfigs);

     // Check that the assignable instance map contains new instances and there are no
     // TaskAssignResults due to previous live instances being removed
     Assert.assertEquals(_assignableInstanceManager.getTaskAssignResultMap().size(), 0);
     Assert.assertEquals(_assignableInstanceManager.getAssignableInstanceMap().size(),
         newLiveInstances.size());
     for (String instance : newLiveInstances.keySet()) {
       Assert
           .assertTrue(_assignableInstanceManager.getAssignableInstanceMap().containsKey(instance));
     }
   }

   public class MockTaskDataCache extends TaskDataCache {
     private Map<String, JobConfig> _jobConfigMap;
     private Map<String, WorkflowConfig> _workflowConfigMap;
     private Map<String, JobContext> _jobContextMap;
     private Map<String, WorkflowContext> _workflowContextMap;

     public MockTaskDataCache(String clusterName) {
       super(clusterName);
       _jobConfigMap = new HashMap<>();
       _workflowConfigMap = new HashMap<>();
       _jobContextMap = new HashMap<>();
       _workflowContextMap = new HashMap<>();
     }

     public void addJobConfig(String jobName, JobConfig jobConfig) {
       _jobConfigMap.put(jobName, jobConfig);
     }

     public void addJobContext(String jobName, JobContext jobContext) {
       _jobContextMap.put(jobName, jobContext);
     }

     public void addWorkflowConfig(String workflowName, WorkflowConfig workflowConfig) {
       _workflowConfigMap.put(workflowName, workflowConfig);
     }

     public void addWorkflowContext(String workflowName, WorkflowContext workflowContext) {
       _workflowContextMap.put(workflowName, workflowContext);
     }

     @Override
     public JobContext getJobContext(String jobName) {
       return _jobContextMap.get(jobName);
     }

     @Override
     public Map<String, JobConfig> getJobConfigMap() {
       return _jobConfigMap;
     }

     @Override
     public Map<String, WorkflowConfig> getWorkflowConfigMap() {
       return _workflowConfigMap;
     }

     public Map<String, JobContext> getJobContextMap() {
       return _jobContextMap;
     }

     public Map<String, WorkflowContext> getWorkflowContextMap() {
       return _workflowContextMap;
     }
   }

   public class MockJobContext extends JobContext {
     private Set<Integer> _taskPartitionSet;
     private Map<Integer, TaskPartitionState> _taskPartitionStateMap;
     private Map<Integer, String> _partitionToTaskIDMap;
     private Map<Integer, String> _taskToInstanceMap;

     public MockJobContext(ZNRecord record, Map<String, LiveInstance> liveInstanceMap,
         Set<String> taskIDs) {
       super(record);
       _taskPartitionSet = new HashSet<>();
       _taskPartitionStateMap = new HashMap<>();
       _partitionToTaskIDMap = new HashMap<>();
       _taskToInstanceMap = new HashMap<>();

       List<String> taskIDList = new ArrayList<>(taskIDs);
       for (int i = 0; i < taskIDList.size(); i++) {
         _taskPartitionSet.add(i);
         _taskPartitionStateMap.put(i, TaskPartitionState.RUNNING);
         _partitionToTaskIDMap.put(i, taskIDList.get(i));
         String someInstance = liveInstanceMap.keySet().iterator().next();
         _taskToInstanceMap.put(i, someInstance);
       }
     }

     @Override
     public Set<Integer> getPartitionSet() {
       return _taskPartitionSet;
     }

     @Override
     public TaskPartitionState getPartitionState(int p) {
       return _taskPartitionStateMap.get(p);
     }

     @Override
     public String getAssignedParticipant(int p) {
       return _taskToInstanceMap.get(p);
     }

     @Override
     public String getTaskIdForPartition(int p) {
       return _partitionToTaskIDMap.get(p);
     }
   }
 }
	package org.apache.helix.task;

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

	import org.apache.helix.zookeeper.datamodel.ZNRecord;
	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.assigner.AssignableInstance;
	import org.apache.helix.task.assigner.TaskAssignResult;
	import org.testng.Assert;
	import org.testng.annotations.BeforeClass;
	import org.testng.annotations.Test;

	public class TestAssignableInstanceManager {
	private static final int NUM_PARTICIPANTS = 3;
	private static final int NUM_JOBS = 3;
	private static final int NUM_TASKS = 3;
	private static final String CLUSTER_NAME = "TestCluster_0";
	private static final String INSTANCE_PREFIX = "Instance_";
	private static final String JOB_PREFIX = "Job_";
	private static final String TASK_PREFIX = "Task_";

	private ClusterConfig _clusterConfig;
	private MockTaskDataCache _taskDataCache;
	private AssignableInstanceManager _assignableInstanceManager;
	private Map<String, LiveInstance> _liveInstances;
	private Map<String, InstanceConfig> _instanceConfigs;
	private Set<String> _taskIDs; // To keep track of what tasks were created

	@BeforeClass
	public void beforeClass() {
	System.out.println(
	"START " + this.getClass().getSimpleName() + " at " + new Date(System.currentTimeMillis()));
	_clusterConfig = new ClusterConfig(CLUSTER_NAME);
	_taskDataCache = new MockTaskDataCache(CLUSTER_NAME);
	_liveInstances = new HashMap<>();
	_instanceConfigs = new HashMap<>();
	_taskIDs = new HashSet<>();

	// Populate live instances and their corresponding instance configs
	for (int i = 0; i < NUM_PARTICIPANTS; i++) {
	String instanceName = INSTANCE_PREFIX + i;
	LiveInstance liveInstance = new LiveInstance(instanceName);
	InstanceConfig instanceConfig = new InstanceConfig(instanceName);
	_liveInstances.put(instanceName, liveInstance);
	_instanceConfigs.put(instanceName, instanceConfig);
	}

	// Populate taskDataCache with JobConfigs and JobContexts
	for (int i = 0; i < NUM_JOBS; i++) {
	String jobName = JOB_PREFIX + i;

	// Create a JobConfig
	JobConfig.Builder jobConfigBuilder = new JobConfig.Builder();
	List<TaskConfig> taskConfigs = new ArrayList<>();
	for (int j = 0; j < NUM_TASKS; j++) {
	String taskID = jobName + "_" + TASK_PREFIX + j;
	TaskConfig.Builder taskConfigBuilder = new TaskConfig.Builder();
	taskConfigBuilder.setTaskId(taskID);
	_taskIDs.add(taskID);
	taskConfigs.add(taskConfigBuilder.build());
	}

	jobConfigBuilder.setJobId(jobName);
	jobConfigBuilder.addTaskConfigs(taskConfigs);
	jobConfigBuilder.setCommand("MOCK");
	jobConfigBuilder.setWorkflow("WORKFLOW");
	_taskDataCache.addJobConfig(jobName, jobConfigBuilder.build());

	// Create a JobContext
	ZNRecord znRecord = new ZNRecord(JOB_PREFIX + "context_" + i);
	JobContext jobContext = new MockJobContext(znRecord, _liveInstances, _taskIDs);
	_taskDataCache.addJobContext(jobName, jobContext);
	_taskIDs.clear();
	}

	// Create an AssignableInstanceManager and build
	_assignableInstanceManager = new AssignableInstanceManager();
	_assignableInstanceManager.buildAssignableInstances(_clusterConfig, _taskDataCache,
	_liveInstances, _instanceConfigs);
	}

	@Test
	public void testGetAssignableInstanceMap() {
	Map<String, AssignableInstance> assignableInstanceMap =
	_assignableInstanceManager.getAssignableInstanceMap();
	for (String liveInstance : _liveInstances.keySet()) {
	Assert.assertTrue(assignableInstanceMap.containsKey(liveInstance));
	}
	}

	@Test(dependsOnMethods = "testGetAssignableInstanceMap")
	public void testGetTaskAssignResultMap() {
	Map<String, TaskAssignResult> taskAssignResultMap =
	_assignableInstanceManager.getTaskAssignResultMap();
	for (String taskID : _taskIDs) {
	Assert.assertTrue(taskAssignResultMap.containsKey(taskID));
	}
	}

	@Test(dependsOnMethods = "testGetTaskAssignResultMap")
	public void testUpdateAssignableInstances() {
	Map<String, LiveInstance> newLiveInstances = new HashMap<>();
	Map<String, InstanceConfig> newInstanceConfigs = new HashMap<>();

	// A brand new set of LiveInstances
	for (int i = NUM_PARTICIPANTS; i < NUM_PARTICIPANTS + 3; i++) {
	String instanceName = INSTANCE_PREFIX + i;
	newLiveInstances.put(instanceName, new LiveInstance(instanceName));
	newInstanceConfigs.put(instanceName, new InstanceConfig(instanceName));
	}

	_assignableInstanceManager.updateAssignableInstances(_clusterConfig, newLiveInstances,
	newInstanceConfigs);

	// Check that the assignable instance map contains new instances and there are no
	// TaskAssignResults due to previous live instances being removed
	Assert.assertEquals(_assignableInstanceManager.getTaskAssignResultMap().size(), 0);
	Assert.assertEquals(_assignableInstanceManager.getAssignableInstanceMap().size(),
	newLiveInstances.size());
	for (String instance : newLiveInstances.keySet()) {
	Assert
	.assertTrue(_assignableInstanceManager.getAssignableInstanceMap().containsKey(instance));
	}
	}

	public class MockTaskDataCache extends TaskDataCache {
	private Map<String, JobConfig> _jobConfigMap;
	private Map<String, WorkflowConfig> _workflowConfigMap;
	private Map<String, JobContext> _jobContextMap;
	private Map<String, WorkflowContext> _workflowContextMap;

	public MockTaskDataCache(String clusterName) {
	super(clusterName);
	_jobConfigMap = new HashMap<>();
	_workflowConfigMap = new HashMap<>();
	_jobContextMap = new HashMap<>();
	_workflowContextMap = new HashMap<>();
	}

	public void addJobConfig(String jobName, JobConfig jobConfig) {
	_jobConfigMap.put(jobName, jobConfig);
	}

	public void addJobContext(String jobName, JobContext jobContext) {
	_jobContextMap.put(jobName, jobContext);
	}

	public void addWorkflowConfig(String workflowName, WorkflowConfig workflowConfig) {
	_workflowConfigMap.put(workflowName, workflowConfig);
	}

	public void addWorkflowContext(String workflowName, WorkflowContext workflowContext) {
	_workflowContextMap.put(workflowName, workflowContext);
	}

	@Override
	public JobContext getJobContext(String jobName) {
	return _jobContextMap.get(jobName);
	}

	@Override
	public Map<String, JobConfig> getJobConfigMap() {
	return _jobConfigMap;
	}

	@Override
	public Map<String, WorkflowConfig> getWorkflowConfigMap() {
	return _workflowConfigMap;
	}

	public Map<String, JobContext> getJobContextMap() {
	return _jobContextMap;
	}

	public Map<String, WorkflowContext> getWorkflowContextMap() {
	return _workflowContextMap;
	}
	}

	public class MockJobContext extends JobContext {
	private Set<Integer> _taskPartitionSet;
	private Map<Integer, TaskPartitionState> _taskPartitionStateMap;
	private Map<Integer, String> _partitionToTaskIDMap;
	private Map<Integer, String> _taskToInstanceMap;

	public MockJobContext(ZNRecord record, Map<String, LiveInstance> liveInstanceMap,
	Set<String> taskIDs) {
	super(record);
	_taskPartitionSet = new HashSet<>();
	_taskPartitionStateMap = new HashMap<>();
	_partitionToTaskIDMap = new HashMap<>();
	_taskToInstanceMap = new HashMap<>();

	List<String> taskIDList = new ArrayList<>(taskIDs);
	for (int i = 0; i < taskIDList.size(); i++) {
	_taskPartitionSet.add(i);
	_taskPartitionStateMap.put(i, TaskPartitionState.RUNNING);
	_partitionToTaskIDMap.put(i, taskIDList.get(i));
	String someInstance = liveInstanceMap.keySet().iterator().next();
	_taskToInstanceMap.put(i, someInstance);
	}
	}

	@Override
	public Set<Integer> getPartitionSet() {
	return _taskPartitionSet;
	}

	@Override
	public TaskPartitionState getPartitionState(int p) {
	return _taskPartitionStateMap.get(p);
	}

	@Override
	public String getAssignedParticipant(int p) {
	return _taskToInstanceMap.get(p);
	}

	@Override
	public String getTaskIdForPartition(int p) {
	return _partitionToTaskIDMap.get(p);
	}
	}
	}