/*
* Copyright 2013 LinkedIn, Inc
*
* Licensed 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.
*/
package voldemort.client.rebalance;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;
import org.junit.Test;
import org.mockito.Mockito;
import voldemort.ClusterTestUtils;
import voldemort.ServerTestUtils;
import voldemort.client.protocol.admin.AdminClient;
import voldemort.client.rebalance.task.StealerBasedRebalanceTask;
import voldemort.cluster.Cluster;
public class RebalanceSchedulerTest {
private static List<StealerBasedRebalanceTask> sbTaskList = new ArrayList<StealerBasedRebalanceTask>();
private static ExecutorService service = Executors.newFixedThreadPool(3);
private static RebalanceScheduler scheduler = new RebalanceScheduler(service, 10);
private static AdminClient adminClient;
private final Semaphore donorPermit = new Semaphore(1);
private final RebalanceBatchPlanProgressBar progressBar = new RebalanceBatchPlanProgressBar(0,
1,
1);
@Test
public void test() {
RebalanceScheduler mockedScheduler = Mockito.spy(scheduler);
Cluster zzCurrent = ClusterTestUtils.getZZCluster();
adminClient = ServerTestUtils.getAdminClient(zzCurrent);
Map<String, List<Integer>> outerMap = new HashMap<String, List<Integer>>();
List<Integer> someList = Arrays.asList(0, 1, 2);
outerMap.put("storeA", someList);
int stealerId = 0;
int donorId = 1;
RebalanceTaskInfo partitionsInfo = new RebalanceTaskInfo(stealerId,
donorId,
(HashMap<String, List<Integer>>) outerMap,
zzCurrent);
StealerBasedRebalanceTask sbTask = new StealerBasedRebalanceTask(0,
0,
partitionsInfo,
donorPermit,
adminClient,
progressBar,
mockedScheduler);
sbTaskList.add(sbTask);
sbTaskList.add(sbTask);
sbTaskList.add(sbTask);
sbTaskList.add(sbTask);
sbTaskList.add(sbTask);
mockedScheduler.initializeLatch(sbTaskList.size());
mockedScheduler.populateTasksByStealer(sbTaskList);
// In the beginning both stealer and donor are idle so scheduler should
// return the scheduled
// task
StealerBasedRebalanceTask scheduledTask = mockedScheduler.scheduleNextTask(false);
org.junit.Assert.assertNotNull(sbTask);
org.junit.Assert.assertEquals(sbTask, scheduledTask);
mockedScheduler.removeNodesFromWorkerList(Arrays.asList(stealerId, donorId));
// Now lets remove the donor from the worker list so that the donor is
// idle and add the
// stealer to the worker list. The scheduler should return null as it
// won't be able to
// schedule the task.
mockedScheduler.addNodesToWorkerList(Arrays.asList(stealerId));
mockedScheduler.removeNodesFromWorkerList(Arrays.asList(donorId));
org.junit.Assert.assertEquals(mockedScheduler.scheduleNextTask(false), null);
// This time stealer doesn't have any work but donor does
mockedScheduler.addNodesToWorkerList(Arrays.asList(donorId));
mockedScheduler.removeNodesFromWorkerList(Arrays.asList(stealerId));
org.junit.Assert.assertEquals(mockedScheduler.scheduleNextTask(false), null);
// And now both stealer and donor have work to do
mockedScheduler.addNodesToWorkerList(Arrays.asList(donorId));
mockedScheduler.addNodesToWorkerList(Arrays.asList(stealerId));
org.junit.Assert.assertEquals(mockedScheduler.scheduleNextTask(false), null);
// And stealer and donor are idle again
mockedScheduler.removeNodesFromWorkerList(Arrays.asList(stealerId, donorId));
StealerBasedRebalanceTask nextscheduledTask = mockedScheduler.scheduleNextTask(false);
org.junit.Assert.assertNotNull(sbTask);
org.junit.Assert.assertEquals(sbTask, nextscheduledTask);
}
}