fix(server): fix some issues of the distributed scheduler

Tsukilc · Tsukilc · commit 1113520c394f · 2026-01-15T17:02:05.000+08:00
diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/DistributedTaskScheduler.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/DistributedTaskScheduler.java
@@ -19,7 +19,9 @@
 
 import java.util.Iterator;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -118,6 +120,11 @@ private static boolean sleep(long ms) {
     public void cronSchedule() {
         // Perform periodic scheduling tasks
 
+        // Check closed flag first to exit early
+        if (this.closed.get()) {
+            return;
+        }
+
         if (!this.graph.started() || this.graph.closed()) {
             return;
         }
@@ -253,6 +260,10 @@ public <V> Future<?> schedule(HugeTask<V> task) {
             return this.ephemeralTaskExecutor.submit(task);
         }
 
+        // Validate task state before saving to ensure correct exception type
+        E.checkState(task.type() != null, "Task type can't be null");
+        E.checkState(task.name() != null, "Task name can't be null");
+
         // Process schema task
         // Handle gremlin task
         // Handle OLAP calculation tasks
@@ -286,13 +297,25 @@ protected <V> void initTaskParams(HugeTask<V> task) {
 
     @Override
     public <V> void cancel(HugeTask<V> task) {
-        // Update status to CANCELLING
-        if (!task.completed()) {
-            // Task not completed, can only execute status not CANCELLING
-            this.updateStatus(task.id(), null, TaskStatus.CANCELLING);
-        } else {
-            LOG.info("cancel task({}) error, task has completed", task.id());
+        E.checkArgumentNotNull(task, "Task can't be null");
+
+        if (task.completed() || task.cancelling()) {
+            return;
         }
+
+        LOG.info("Cancel task '{}' in status {}", task.id(), task.status());
+
+        // Check if task is running locally, cancel it directly if so
+        HugeTask<?> runningTask = this.runningTasks.get(task.id());
+        if (runningTask != null) {
+            boolean cancelled = runningTask.cancel(true);
+            LOG.info("Cancel local running task '{}' result: {}", task.id(), cancelled);
+            return;
+        }
+
+        // Task not running locally, update status to CANCELLING
+        // for cronSchedule() or other nodes to handle
+        this.updateStatus(task.id(), null, TaskStatus.CANCELLING);
     }
 
     @Override
@@ -316,14 +339,25 @@ protected <V> HugeTask<V> deleteFromDB(Id id) {
 
     @Override
     public <V> HugeTask<V> delete(Id id, boolean force) {
-        if (!force) {
-            // Change status to DELETING, perform the deletion operation through automatic
-            // scheduling.
-            this.updateStatus(id, null, TaskStatus.DELETING);
+        HugeTask<?> task = this.taskWithoutResult(id);
+        if (task == null) {
             return null;
-        } else {
-            return this.deleteFromDB(id);
         }
+
+        if (!force) {
+            // Check task status: can't delete running tasks without force
+            if (!task.completed() && task.status() != TaskStatus.DELETING) {
+                throw new IllegalArgumentException(
+                        String.format("Can't delete incomplete task '%s' in status %s, " +
+                                      "Please try to cancel the task first",
+                                      id, task.status()));
+            }
+            // Already in DELETING status, delete directly from DB
+            // Completed tasks can also be deleted directly
+        }
+
+        // Delete from DB directly for completed/DELETING tasks or force=true
+        return this.deleteFromDB(id);
     }
 
     @Override
@@ -353,6 +387,18 @@ public boolean close() {
             cronFuture.cancel(false);
         }
 
+        // Wait for cron task to complete to ensure all transactions are closed
+        try {
+            cronFuture.get(schedulePeriod + 5, TimeUnit.SECONDS);
+        } catch (CancellationException e) {
+            // Task was cancelled, this is expected
+            LOG.debug("Cron task was cancelled");
+        } catch (TimeoutException e) {
+            LOG.warn("Cron task did not complete in time when closing scheduler");
+        } catch (ExecutionException | InterruptedException e) {
+            LOG.warn("Exception while waiting for cron task to complete", e);
+        }
+
         if (!this.taskDbExecutor.isShutdown()) {
             this.call(() -> {
                 try {
@@ -363,7 +409,10 @@ public boolean close() {
                 this.graph.closeTx();
             });
         }
-        return true;
+
+        //todo: serverInfoManager section should be removed in the future.
+        return this.serverManager().close();
+        //return true;
     }
 
     @Override
@@ -387,15 +436,17 @@ private <V> HugeTask<V> waitUntilTaskCompleted(Id id, long seconds,
         long passes = seconds * 1000 / intervalMs;
         HugeTask<V> task = null;
         for (long pass = 0; ; pass++) {
-            try {
-                task = this.taskWithoutResult(id);
-            } catch (NotFoundException e) {
-                if (task != null && task.completed()) {
-                    assert task.id().asLong() < 0L : task.id();
+            HugeTask<V> previousTask = task;
+            task = this.taskWithoutResult(id);
+            if (task == null) {
+                // Task not found in DB
+                if (previousTask != null && previousTask.completed()) {
+                    // Task was completed and then deleted (ephemeral task case)
+                    assert previousTask.id().asLong() < 0L : previousTask.id();
                     sleep(intervalMs);
-                    return task;
+                    return previousTask;
                 }
-                throw e;
+                throw new NotFoundException("Can't find task with id '%s'", id);
             }
             if (task.completed()) {
                 // Wait for task result being set after status is completed
@@ -466,6 +517,11 @@ private <V> V call(Callable<V> callable, ExecutorService executor) {
     protected boolean updateStatus(Id id, TaskStatus prestatus,
                                    TaskStatus status) {
         HugeTask<Object> task = this.taskWithoutResult(id);
+        if (task == null) {
+            // Task was already deleted by cronSchedule or another thread
+            LOG.info("Task '{}' not found, may have been deleted", id);
+            return false;
+        }
         initTaskParams(task);
         if (prestatus == null || task.status() == prestatus) {
             task.overwriteStatus(status);
diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/HugeTask.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/HugeTask.java
@@ -371,9 +371,7 @@ protected void done() {
     protected void set(V v) {
         String result = JsonUtil.toJson(v);
         checkPropertySize(result, P.RESULT);
-        if (!this.result(TaskStatus.SUCCESS, result)) {
-            assert this.completed();
-        }
+        assert this.result(TaskStatus.SUCCESS, result) || this.completed();
         // Will call done() and may cause to save to store
         super.set(v);
     }
@@ -742,19 +740,26 @@ private void checkPropertySize(int propertyLength, String propertyName) {
 
     public void syncWait() {
         // This method is just called by tests
+        /*
+         * For ephemeral tasks (negative ID), directly wait on the Future.
+         * Ephemeral tasks are not saved to DB, so we can't query them by ID.
+         * Since HugeTask extends FutureTask, we can directly wait for completion.
+         */
+        if (this.id().asLong() < 0) {
+            try {
+                this.get();
+            } catch (Exception e) {
+                throw new HugeException("Failed to wait for task '%s' completed",
+                                        e, this.id);
+            }
+            return;
+        }
+
+        // For normal tasks, wait through scheduler
         HugeTask<?> task = null;
         try {
             task = this.scheduler().waitUntilTaskCompleted(this.id());
         } catch (Throwable e) {
-            if (this.callable() instanceof EphemeralJob &&
-                e.getClass() == NotFoundException.class &&
-                e.getMessage().contains("Can't find task with id")) {
-                /*
-                 * The task with EphemeralJob won't saved in backends and
-                 * will be removed from memory when completed
-                 */
-                return;
-            }
             throw new HugeException("Failed to wait for task '%s' completed",
                                     e, this.id);
         }
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/MultiGraphsTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/MultiGraphsTest.java
@@ -248,7 +248,7 @@ public void testCreateGraphsWithInvalidNames() {
 
     @Test
     public void testCreateGraphsWithSameName() {
-        List<HugeGraph> graphs = openGraphs("g", "g", "G");
+        List<HugeGraph> graphs = openGraphs("gg", "gg", "GG");
         HugeGraph g1 = graphs.get(0);
         HugeGraph g2 = graphs.get(1);
         HugeGraph g3 = graphs.get(2);
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/TaskCoreTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/TaskCoreTest.java
@@ -192,13 +192,18 @@ public Object execute() throws Exception {
         Assert.assertEquals("test", task.type());
         Assert.assertFalse(task.completed());
 
-        HugeTask<?> task2 = scheduler.waitUntilTaskCompleted(task.id(), 10);
+        // Ephemeral tasks are node-local and not persisted to DB.
+        // Use Future.get() to wait for completion instead of ID-based lookup.
+        try {
+            task.get(10, java.util.concurrent.TimeUnit.SECONDS);
+        } catch (Exception e) {
+            throw new RuntimeException("Ephemeral task execution failed", e);
+        }
+
         Assert.assertEquals(TaskStatus.SUCCESS, task.status());
         Assert.assertEquals("{\"k1\":13579,\"k2\":\"24680\"}", task.result());
 
-        Assert.assertEquals(TaskStatus.SUCCESS, task2.status());
-        Assert.assertEquals("{\"k1\":13579,\"k2\":\"24680\"}", task2.result());
-
+        // Ephemeral tasks are not stored in DB, so these should throw NotFoundException
         Assert.assertThrows(NotFoundException.class, () -> {
             scheduler.waitUntilTaskCompleted(task.id(), 10);
         });
@@ -553,10 +558,12 @@ public void testGremlinJobAndCancel() throws TimeoutException {
         scheduler.cancel(task);
 
         task = scheduler.task(task.id());
-        System.out.println(scheduler.getClass());
-        if (scheduler.getClass().equals(DistributedTaskScheduler.class)) {
-            Assert.assertEquals(TaskStatus.CANCELLING, task.status());
-        }
+        // For DistributedTaskScheduler, local cancel may result in CANCELLED directly
+        // (task thread updates status after being interrupted)
+        // or CANCELLING (if task hasn't processed the interrupt yet)
+        Assert.assertTrue("Task status should be CANCELLING or CANCELLED, but was " + task.status(),
+                          task.status() == TaskStatus.CANCELLING ||
+                          task.status() == TaskStatus.CANCELLED);
 
         task = scheduler.waitUntilTaskCompleted(task.id(), 10);
         Assert.assertEquals(TaskStatus.CANCELLED, task.status());
@@ -628,9 +635,9 @@ public void testGremlinJobAndRestore() throws Exception {
         scheduler.cancel(task);
 
         task = scheduler.task(task.id());
-        if (scheduler.getClass().equals(DistributedTaskScheduler.class)) {
-            Assert.assertEquals(TaskStatus.CANCELLING, task.status());
-        }
+        Assert.assertTrue("Task status should be CANCELLING or CANCELLED, but was " + task.status(),
+                          task.status() == TaskStatus.CANCELLING ||
+                          task.status() == TaskStatus.CANCELLED);
 
         task = scheduler.waitUntilTaskCompleted(task.id(), 10);
         Assert.assertEquals(TaskStatus.CANCELLED, task.status());
@@ -640,36 +647,39 @@ public void testGremlinJobAndRestore() throws Exception {
         Assert.assertNull(task.result());
 
         HugeTask<Object> finalTask = task;
-        Assert.assertThrows(IllegalArgumentException.class, () -> {
-            Whitebox.invoke(scheduler.getClass(), "restore", scheduler,
-                    finalTask);
-        }, e -> {
-            Assert.assertContains("No need to restore completed task",
-                    e.getMessage());
-        });
 
-        HugeTask<Object> task2 = scheduler.task(task.id());
-        Assert.assertThrows(IllegalArgumentException.class, () -> {
+        // because Distributed do nothing in restore, so only test StandardTaskScheduler here
+        if (scheduler.getClass().equals(StandardTaskScheduler.class)) {
+            Assert.assertThrows(IllegalArgumentException.class, () -> {
+                Whitebox.invoke(scheduler.getClass(), "restore", scheduler,
+                                finalTask);
+            }, e -> {
+                Assert.assertContains("No need to restore completed task",
+                                      e.getMessage());
+            });
+
+            HugeTask<Object> task2 = scheduler.task(task.id());
+            Assert.assertThrows(IllegalArgumentException.class, () -> {
+                Whitebox.invoke(scheduler.getClass(), "restore", scheduler, task2);
+            }, e -> {
+                Assert.assertContains("No need to restore completed task",
+                                      e.getMessage());
+            });
+
+            Whitebox.setInternalState(task2, "status", TaskStatus.RUNNING);
             Whitebox.invoke(scheduler.getClass(), "restore", scheduler, task2);
-        }, e -> {
-            Assert.assertContains("No need to restore completed task",
-                    e.getMessage());
-        });
-
-        Whitebox.setInternalState(task2, "status", TaskStatus.RUNNING);
-        Whitebox.invoke(scheduler.getClass(), "restore", scheduler, task2);
 
-        Assert.assertThrows(IllegalArgumentException.class, () -> {
-            Whitebox.invoke(scheduler.getClass(), "restore", scheduler, task2);
-        }, e -> {
-            Assert.assertContains("is already in the queue", e.getMessage());
-        });
-
-        scheduler.waitUntilTaskCompleted(task2.id(), 10);
-        sleepAWhile(500);
-        Assert.assertEquals(10, task2.progress());
-        Assert.assertEquals(1, task2.retries());
-        Assert.assertEquals("100", task2.result());
+            Assert.assertThrows(IllegalArgumentException.class, () -> {
+                Whitebox.invoke(scheduler.getClass(), "restore", scheduler, task2);
+            }, e -> {
+                Assert.assertContains("is already in the queue", e.getMessage());
+            });
+            scheduler.waitUntilTaskCompleted(task2.id(), 10);
+            sleepAWhile(500);
+            Assert.assertEquals(10, task2.progress());
+            Assert.assertEquals(1, task2.retries());
+            Assert.assertEquals("100", task2.result());
+        }
     }
 
     private HugeTask<Object> runGremlinJob(String gremlin) {
