Makes TaskStateModel gets properly shutdown when the JVM exists Signed-off-by: Yinan Li <[email protected]>
Project: http://git-wip-us.apache.org/repos/asf/helix/repo Commit: http://git-wip-us.apache.org/repos/asf/helix/commit/6ae7891e Tree: http://git-wip-us.apache.org/repos/asf/helix/tree/6ae7891e Diff: http://git-wip-us.apache.org/repos/asf/helix/diff/6ae7891e Branch: refs/heads/master Commit: 6ae7891e1754190219cc47940013aff215abacf8 Parents: eb82c95 Author: Yinan Li <[email protected]> Authored: Tue Sep 8 17:14:55 2015 -0700 Committer: Yinan Li <[email protected]> Committed: Wed Sep 9 12:40:58 2015 -0700 ---------------------------------------------------------------------- .../helix/task/TaskStateModelFactory.java | 28 ++++++- .../helix/task/TaskStateModelFactoryTest.java | 82 ++++++++++++++++++++ 2 files changed, 108 insertions(+), 2 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/helix/blob/6ae7891e/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java ---------------------------------------------------------------------- diff --git a/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java b/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java index 23d6160..43da591 100644 --- a/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java +++ b/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java @@ -19,6 +19,7 @@ package org.apache.helix.task; * under the License. */ +import java.util.Deque; import java.util.Map; import org.apache.helix.HelixManager; @@ -26,20 +27,43 @@ import org.apache.helix.api.StateTransitionHandlerFactory; import org.apache.helix.api.id.PartitionId; import org.apache.helix.api.id.ResourceId; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Lists; + /** * Factory class for {@link TaskStateModel}. */ public class TaskStateModelFactory extends StateTransitionHandlerFactory<TaskStateModel> { + private final HelixManager _manager; private final Map<String, TaskFactory> _taskFactoryRegistry; + private final Deque<TaskStateModel> _taskStateModels; public TaskStateModelFactory(HelixManager manager, Map<String, TaskFactory> taskFactoryRegistry) { _manager = manager; - _taskFactoryRegistry = taskFactoryRegistry; + _taskFactoryRegistry = ImmutableMap.copyOf(taskFactoryRegistry); + _taskStateModels = Lists.newLinkedList(); } @Override public TaskStateModel createStateTransitionHandler(ResourceId resourceId, PartitionId partitionId) { - return new TaskStateModel(_manager, _taskFactoryRegistry); + final TaskStateModel taskStateModel = new TaskStateModel(_manager, _taskFactoryRegistry); + _taskStateModels.push(taskStateModel); + return taskStateModel; + } + + /** + * Shutdown this {@link TaskStateModelFactory} instance. + * + * <p> + * This method shuts down all of the {@link TaskStateModel} instances created by this + * {@link TaskStateModelFactory} instance. + * </p> + */ + public void shutdown() { + while (!_taskStateModels.isEmpty()) { + TaskStateModel taskStateModel = _taskStateModels.pop(); + taskStateModel.shutdown(); + } } } http://git-wip-us.apache.org/repos/asf/helix/blob/6ae7891e/helix-core/src/test/java/org/apache/helix/task/TaskStateModelFactoryTest.java ---------------------------------------------------------------------- diff --git a/helix-core/src/test/java/org/apache/helix/task/TaskStateModelFactoryTest.java b/helix-core/src/test/java/org/apache/helix/task/TaskStateModelFactoryTest.java new file mode 100644 index 0000000..d0f6a41 --- /dev/null +++ b/helix-core/src/test/java/org/apache/helix/task/TaskStateModelFactoryTest.java @@ -0,0 +1,82 @@ +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.List; +import java.util.Map; + +import org.apache.helix.Mocks; +import org.apache.helix.api.id.PartitionId; +import org.apache.helix.api.id.ResourceId; +import org.testng.Assert; +import org.testng.annotations.Test; + +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +/** + * Unit tests for {@link TaskStateModelFactory}. + * + * @author liyinan926 + */ +@Test(groups = { "org.apache.helix.task" }) +public class TaskStateModelFactoryTest { + + @Test + public void testCreateAndShutdownTaskStateModels() { + Map<String, TaskFactory> taskFactoryMap = Maps.newHashMap(); + taskFactoryMap.put("Test", new TestTaskFactory()); + TaskStateModelFactory taskStateModelFactory = new TaskStateModelFactory(new Mocks.MockManager(), taskFactoryMap); + + List<TaskStateModel> taskStateModelList = Lists.newArrayList(); + for (int i = 0; i < 3; i++) { + ResourceId resourceId = ResourceId.from("Task_" + i); + taskStateModelList.add( + taskStateModelFactory.createStateTransitionHandler(resourceId, PartitionId.from(resourceId, ""))); + } + + taskStateModelFactory.shutdown(); + + for (TaskStateModel taskStateModel : taskStateModelList) { + Assert.assertTrue(taskStateModel.isShutdown()); + } + } + + private static class TestTaskFactory implements TaskFactory { + + @Override + public Task createNewTask(TaskCallbackContext context) { + return new TestTask(); + } + + private static class TestTask implements Task { + + @Override + public TaskResult run() { + return new TaskResult(TaskResult.Status.COMPLETED, ""); + } + + @Override + public void cancel() { + + } + } + } +}
