diff --git a/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/main/java/org/apache/geaflow/cluster/exception/ComponentExceptionSupervisor.java b/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/main/java/org/apache/geaflow/cluster/exception/ComponentExceptionSupervisor.java index 145174b2b..13a8cec25 100644 --- a/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/main/java/org/apache/geaflow/cluster/exception/ComponentExceptionSupervisor.java +++ b/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/main/java/org/apache/geaflow/cluster/exception/ComponentExceptionSupervisor.java @@ -21,6 +21,9 @@ import static org.apache.geaflow.cluster.constants.ClusterConstants.EXIT_CODE; +import com.google.common.annotations.VisibleForTesting; +import java.util.Objects; +import java.util.function.IntConsumer; import org.apache.geaflow.cluster.task.runner.AbstractTaskRunner; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -30,6 +33,16 @@ public class ComponentExceptionSupervisor extends AbstractTaskRunner 0) { + while (exitCode.get() == EXIT_NOT_CALLED && retry > 0) { SleepUtils.sleepMilliSecond(100); retry--; } @@ -165,4 +162,15 @@ public EventType getEventType() { return EventType.CREATE_TASK; } } + + private static class TestComponentExceptionSupervisor extends ComponentExceptionSupervisor { + + private TestComponentExceptionSupervisor(IntConsumer processExit) { + super(processExit); + } + + private static void install(IntConsumer processExit) { + setInstance(new TestComponentExceptionSupervisor(processExit)); + } + } } diff --git a/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/exception/ComponentUncaughtExceptionHandlerTest.java b/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/exception/ComponentUncaughtExceptionHandlerTest.java index 1bae92aca..38426508f 100644 --- a/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/exception/ComponentUncaughtExceptionHandlerTest.java +++ b/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/exception/ComponentUncaughtExceptionHandlerTest.java @@ -19,52 +19,44 @@ package org.apache.geaflow.cluster.exception; +import static org.apache.geaflow.cluster.constants.ClusterConstants.EXIT_CODE; + +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicBoolean; -import org.apache.geaflow.cluster.util.SystemExitSignalCatcher; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.geaflow.common.utils.ThreadUtil; import org.testng.Assert; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class ComponentUncaughtExceptionHandlerTest { - private static SecurityManager securityManager; - private static AtomicBoolean hasException = new AtomicBoolean(false); - - - @BeforeClass - public void before() { - securityManager = System.getSecurityManager(); - System.setSecurityManager(new SystemExitSignalCatcher(hasException)); - } - - @AfterClass - public void after() { - System.setSecurityManager(securityManager); - } - - @BeforeMethod - public void beforeMethod() { - hasException.set(false); - } - @Test public void testHandleExceptionInThreadPool() throws InterruptedException { - ComponentExceptionSupervisor.getInstance(); + AtomicInteger exitCode = new AtomicInteger(); + CountDownLatch exitCalled = new CountDownLatch(1); + ComponentExceptionSupervisor supervisor = new ComponentExceptionSupervisor(code -> { + exitCode.set(code); + exitCalled.countDown(); + }); + ComponentExceptionSupervisor.setInstance(supervisor); ExecutorService executorService = Executors.newFixedThreadPool(2, ThreadUtil.namedThreadFactory(true, "test-handler", new ComponentUncaughtExceptionHandler())); - executorService.execute(() -> { - throw new RuntimeException("test exception"); - }); - executorService.execute(ComponentExceptionSupervisor.getInstance()); - // wait async thread catch and handle exception - Thread.sleep(100); - Assert.assertTrue(hasException.get()); + try { + executorService.execute(() -> { + throw new RuntimeException("test exception"); + }); + executorService.execute(supervisor); + + Assert.assertTrue(exitCalled.await(1, TimeUnit.SECONDS)); + Assert.assertEquals(exitCode.get(), EXIT_CODE); + } finally { + supervisor.shutdown(); + executorService.shutdown(); + executorService.awaitTermination(1, TimeUnit.SECONDS); + } } } diff --git a/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/util/SystemExitSignalCatcher.java b/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/util/SystemExitSignalCatcher.java deleted file mode 100644 index bf9c9ee2b..000000000 --- a/geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/test/java/org/apache/geaflow/cluster/util/SystemExitSignalCatcher.java +++ /dev/null @@ -1,48 +0,0 @@ -/* - * 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. - */ - -package org.apache.geaflow.cluster.util; - -import java.security.Permission; -import java.util.concurrent.atomic.AtomicBoolean; -import org.apache.geaflow.common.exception.GeaflowRuntimeException; - -public class SystemExitSignalCatcher extends SecurityManager { - - private AtomicBoolean hasSignal; - - public SystemExitSignalCatcher(AtomicBoolean hasSignal) { - this.hasSignal = hasSignal; - } - - @Override - public void checkPermission(Permission perm) { - } - - @Override - public void checkPermission(Permission perm, Object context) { - } - - @Override - public void checkExit(int status) { - super.checkExit(status); - hasSignal.set(true); - throw new GeaflowRuntimeException("throw exception instead of exit process"); - } -}