diff --git a/jdbc/src/main/java/org/apache/zeppelin/jdbc/JDBCUserConfigurations.java b/jdbc/src/main/java/org/apache/zeppelin/jdbc/JDBCUserConfigurations.java index 311fb0bad0c..bcaea51e0ca 100644 --- a/jdbc/src/main/java/org/apache/zeppelin/jdbc/JDBCUserConfigurations.java +++ b/jdbc/src/main/java/org/apache/zeppelin/jdbc/JDBCUserConfigurations.java @@ -19,9 +19,9 @@ import java.sql.SQLException; import java.sql.Statement; -import java.util.HashMap; import java.util.Map; import java.util.Properties; +import java.util.concurrent.ConcurrentHashMap; /** * UserConfigurations for JDBC impersonation. @@ -33,7 +33,7 @@ public class JDBCUserConfigurations { private Boolean isSuccessful; public JDBCUserConfigurations() { - paragraphIdStatementMap = new HashMap<>(); + paragraphIdStatementMap = new ConcurrentHashMap<>(); } public void initStatementMap() throws SQLException { @@ -67,14 +67,26 @@ public void setUserProperty(UsernamePassword usernamePassword) { } public void saveStatement(String paragraphId, Statement statement) throws SQLException { + if (paragraphId == null) { + return; + } paragraphIdStatementMap.put(paragraphId, statement); } public void cancelStatement(String paragraphId) throws SQLException { - paragraphIdStatementMap.get(paragraphId).cancel(); + if (paragraphId == null) { + return; + } + Statement statement = paragraphIdStatementMap.get(paragraphId); + if (statement != null) { + statement.cancel(); + } } public void removeStatement(String paragraphId) { + if (paragraphId == null) { + return; + } paragraphIdStatementMap.remove(paragraphId); } diff --git a/jdbc/src/test/java/org/apache/zeppelin/jdbc/JDBCUserConfigurationsTest.java b/jdbc/src/test/java/org/apache/zeppelin/jdbc/JDBCUserConfigurationsTest.java new file mode 100644 index 00000000000..40fa032a6e4 --- /dev/null +++ b/jdbc/src/test/java/org/apache/zeppelin/jdbc/JDBCUserConfigurationsTest.java @@ -0,0 +1,68 @@ +/** + * 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.zeppelin.jdbc; + +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.sql.SQLException; +import java.sql.Statement; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +class JDBCUserConfigurationsTest { + + @Test + void cancelStatementBeforeSaveShouldNotThrowNPE() { + JDBCUserConfigurations jdbcUserConfigurations = new JDBCUserConfigurations(); + + assertDoesNotThrow(() -> jdbcUserConfigurations.cancelStatement("paragraph-not-registered")); + } + + @Test + void cancelStatementAfterSaveShouldCallCancelOnStatement() throws SQLException { + JDBCUserConfigurations jdbcUserConfigurations = new JDBCUserConfigurations(); + Statement statement = Mockito.mock(Statement.class); + jdbcUserConfigurations.saveStatement("paragraph-1", statement); + + jdbcUserConfigurations.cancelStatement("paragraph-1"); + + verify(statement).cancel(); + } + + @Test + void cancelStatementAfterRemoveShouldNotThrowNPE() throws SQLException { + JDBCUserConfigurations jdbcUserConfigurations = new JDBCUserConfigurations(); + Statement statement = Mockito.mock(Statement.class); + jdbcUserConfigurations.saveStatement("paragraph-1", statement); + jdbcUserConfigurations.removeStatement("paragraph-1"); + + assertDoesNotThrow(() -> jdbcUserConfigurations.cancelStatement("paragraph-1")); + verify(statement, never()).cancel(); + } + + @Test + void nullParagraphIdShouldBeNoOpAcrossAllMapOperations() throws SQLException { + JDBCUserConfigurations jdbcUserConfigurations = new JDBCUserConfigurations(); + Statement statement = Mockito.mock(Statement.class); + + assertDoesNotThrow(() -> jdbcUserConfigurations.saveStatement(null, statement)); + assertDoesNotThrow(() -> jdbcUserConfigurations.cancelStatement(null)); + assertDoesNotThrow(() -> jdbcUserConfigurations.removeStatement(null)); + verify(statement, never()).cancel(); + } +} diff --git a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/AbstractScheduler.java b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/AbstractScheduler.java index 7e99095b7ff..5bb3c82e020 100644 --- a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/AbstractScheduler.java +++ b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/AbstractScheduler.java @@ -77,7 +77,11 @@ public void submit(Job> job) { @Override public Job> cancel(String jobId) { Job> job = jobs.remove(jobId); - job.abort(); + // Synchronize on the same monitor as runJob()'s abort gate so that a cancellation + // happening right before the job is run is never missed (ZEPPELIN-6129). + synchronized (job) { + job.abort(); + } return job; } @@ -121,17 +125,21 @@ public void stop() { * @param runningJob */ protected void runJob(Job> runningJob) { - if (runningJob.isAborted()) { - LOGGER.info("Job {} is aborted", runningJob.getId()); - runningJob.setStatus(Job.Status.ABORT); - runningJob.aborted = false; - return; - } + // Synchronize the abort gate on the same monitor cancel() uses, so a cancellation + // submitted right before the job runs is never missed (ZEPPELIN-6129). + synchronized (runningJob) { + if (runningJob.isAborted()) { + LOGGER.info("Job {} is aborted", runningJob.getId()); + runningJob.setStatus(Job.Status.ABORT); + runningJob.aborted = false; + return; + } - LOGGER.info("Job {} started by scheduler {}", runningJob.getId(), name); - // Don't set RUNNING status when it is RemoteScheduler, update it via JobStatusPoller - if (!getClass().getSimpleName().equals("RemoteScheduler")) { - runningJob.setStatus(Job.Status.RUNNING); + LOGGER.info("Job {} started by scheduler {}", runningJob.getId(), name); + // Don't set RUNNING status when it is RemoteScheduler, update it via JobStatusPoller + if (!getClass().getSimpleName().equals("RemoteScheduler")) { + runningJob.setStatus(Job.Status.RUNNING); + } } runningJob.run(); Object jobResult = runningJob.getReturn(); diff --git a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/Job.java b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/Job.java index b0ed600f45a..d8b4a739b87 100644 --- a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/Job.java +++ b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/scheduler/Job.java @@ -89,7 +89,7 @@ public boolean isFailed() { private Date dateFinished; protected volatile Status status; - transient boolean aborted = false; + transient volatile boolean aborted = false; private volatile String errorMessage; private transient volatile Throwable exception; private transient JobListener listener; diff --git a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/scheduler/AbstractSchedulerAbortRaceTest.java b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/scheduler/AbstractSchedulerAbortRaceTest.java new file mode 100644 index 00000000000..e7f10d44514 --- /dev/null +++ b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/scheduler/AbstractSchedulerAbortRaceTest.java @@ -0,0 +1,168 @@ +/* + * 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.zeppelin.scheduler; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Covers the "abort right before run" race between a cancelling thread + * ({@link AbstractScheduler#cancel(String)}) and the scheduler thread that is about to invoke + * {@link AbstractScheduler#runJob(Job)}. + * + *
Honest limitation: a pure memory-visibility race on a non-volatile field cannot be + * reproduced deterministically without a specialized harness (e.g. jcstress) because it depends + * on JVM safepoints/JIT reordering. This class therefore verifies the observable behavior + * contract instead: (1) abort()/isAborted() agree, (2) a PENDING job aborted before runJob() is + * invoked never has its run() executed and ends in ABORT, and (3) cancel() and the runJob() gate + * are mutually exclusive on the same job monitor, which is a deterministic, latch-driven proof + * that the race window described in ZEPPELIN-6129 Task 2 is closed. + */ +class AbstractSchedulerAbortRaceTest { + + private FIFOScheduler scheduler; + + @AfterEach + void tearDown() { + if (scheduler != null) { + scheduler.stop(); + } + } + + @Test + void testAbortSetsIsAbortedTrue() { + SleepingJob job = new SleepingJob("abortJob", null, 5000); + + assertFalse(job.isAborted()); + job.abort(); + + assertTrue(job.isAborted()); + } + + @Test + void testCancelBeforeRunJobBlocksExecutionThroughSchedulerCancelPath() { + scheduler = new FIFOScheduler("cancel-gate-test"); + SleepingJob job = new SleepingJob("job1", null, 5000); + scheduler.submit(job); + + scheduler.cancel(job.getId()); + scheduler.runJob(job); + + assertEquals(Job.Status.ABORT, job.getStatus()); + assertNull(job.getReturn()); + } + + @Test + void testCancelAndRunJobGateAreMutuallyExclusiveOnJobMonitor() throws Exception { + scheduler = new FIFOScheduler("mutex-test"); + BlockingAbortJob job = new BlockingAbortJob("job1"); + scheduler.submit(job); + + Thread cancelThread = new Thread(() -> scheduler.cancel(job.getId()), "cancel-thread"); + cancelThread.start(); + + assertTrue(job.abortEntered.await(2, TimeUnit.SECONDS), + "cancel thread must reach jobAbort() and hold the job monitor"); + + Thread runJobThread = new Thread(() -> scheduler.runJob(job), "runjob-thread"); + runJobThread.start(); + + assertTrue(waitForState(runJobThread, Thread.State.BLOCKED, 2000), + "runJob() must block waiting for the same job monitor held by cancel()"); + assertFalse(job.runCalled, "job.run() must not start while cancel() still holds the monitor"); + + job.releaseAbort.countDown(); + cancelThread.join(2000); + runJobThread.join(2000); + + assertFalse(job.runCalled, "aborted job must never invoke run()"); + assertEquals(Job.Status.ABORT, job.getStatus()); + } + + private static boolean waitForState(Thread thread, Thread.State expected, long timeoutMs) + throws InterruptedException { + long deadline = System.currentTimeMillis() + timeoutMs; + while (System.currentTimeMillis() < deadline) { + if (thread.getState() == expected) { + return true; + } + Thread.sleep(10); + } + return thread.getState() == expected; + } + + /** + * Job whose {@code jobAbort()} blocks on a latch so the test can control exactly how long the + * cancelling thread holds the job monitor. + */ + private static class BlockingAbortJob extends Job