[ZEPPELIN-6129] Restore parallel paragraph execution for concurrent interpreters

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);
   }
 

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();
+  }
+}

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();

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;

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)}.
+ *
+ * <p>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<Object> {
+
+    private final CountDownLatch abortEntered = new CountDownLatch(1);
+    private final CountDownLatch releaseAbort = new CountDownLatch(1);
+    private volatile boolean runCalled = false;
+
+    BlockingAbortJob(String name) {
+      super(name, null);
+    }
+
+    @Override
+    protected Object jobRun() {
+      runCalled = true;
+      return null;
+    }
+
+    @Override
+    protected boolean jobAbort() {
+      abortEntered.countDown();
+      try {
+        releaseAbort.await(5, TimeUnit.SECONDS);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
+      return true;
+    }
+
+    @Override
+    public void setResult(Object result) {
+    }
+
+    @Override
+    public Object getReturn() {
+      return null;
+    }
+
+    @Override
+    public int progress() {
+      return 0;
+    }
+
+    @Override
+    public Map<String, Object> info() {
+      return Collections.emptyMap();
+    }
+  }
+}