fix(server): preserve caller tx on cancel

docs/server-runtime.md

@@ -463,6 +463,14 @@ thread. A failure while closing one owner must not skip cleanup of other owners
 server-info resources. Cleanup retains the first exception and attaches later cleanup
 failures as suppressed exceptions.
 
+Synchronous task cancellation preserves a transaction already open on its caller,
+including uncommitted graph writes. Cancellation still runs the normal completion and
+status callbacks, but their default cleanup does not close that existing transaction.
+The scheduler closes transactions opened by a cancellation callback when it had no
+transaction before starting, and a running task retains its worker owner until that
+worker's own cleanup finishes. The caller remains responsible for committing or rolling
+back its original transaction.
+
 Graph close attempts auth, scheduler and caller-transaction cleanup separately. If a
 scheduler remains registered or the graph reports active transactions, the graph remains
 open for a later close attempt. Pending local tasks retain their task database transaction

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/DistributedTaskScheduler.java

@@ -334,15 +334,18 @@ public <V> void cancel(HugeTask<V> task) {
         String context = TaskManager.getContext();
         AtomicBoolean cleaned = new AtomicBoolean(true);
         TaskManager.ContextCallable<Void> owned = new TaskManager.ContextCallable<>(() -> {
+            boolean callerTxOpen = this.graph.graph().tx().isOpen();
             try {
                 this.cancelTask(task, true);
                 return null;
             } finally {
-                try {
-                    this.graph.closeTx();
-                } catch (RuntimeException | Error error) {
-                    cleaned.set(false);
-                    throw error;
+                if (!callerTxOpen) {
+                    try {
+                        this.graph.closeTx();
+                    } catch (RuntimeException | Error error) {
+                        cleaned.set(false);
+                        throw error;
+                    }
                 }
             }
         });
@@ -924,10 +927,13 @@ private boolean cancelOwned(PendingTask pending, Callable<Boolean> action, boole
             pending.callbacks++;
         }
         TaskManager.ContextCallable<Boolean> owned = new TaskManager.ContextCallable<>(() -> {
+            boolean callerTxOpen = this.graph.graph().tx().isOpen();
             try {
                 return action.call();
             } finally {
-                this.closeOwner(pending);
+                if (!callerTxOpen) {
+                    this.closeOwner(pending);
+                }
             }
         });
         Callable<Boolean> cancellation = () -> {

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java

@@ -348,6 +348,7 @@ public <V> void initTaskCallable(HugeTask<V> task) {
     @Override
     public <V> void cancel(HugeTask<V> task) {
         E.checkArgumentNotNull(task, "Task can't be null");
+        boolean callerTxOpen = this.graph.graph().tx().isOpen();
         HugeTask<?> memTask;
         synchronized (this) {
             if (task.completed() || task.cancelling()) {
@@ -379,8 +380,10 @@ public <V> void cancel(HugeTask<V> task) {
             if (memTask != null && memTask.isCancelled()) {
                 this.cancelQueuedExecutions(task.id());
             }
-            // Keep the callback registered until its transaction really closes.
-            this.graph.closeTx();
+            // Keep the callback registered until its own transaction really closes.
+            if (!callerTxOpen) {
+                this.graph.closeTx();
+            }
             synchronized (this) {
                 this.cancellations--;
             }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/TaskCallable.java

@@ -56,7 +56,11 @@ public TaskCallable() {
     }
 
     protected void done() {
-        this.closeTx();
+        // Future cancellation runs done() on its caller, whose transaction may predate the callback.
+        // The scheduler closes callback-owned transactions and the actual worker's transaction separately.
+        if (this.task == null || !this.task.isCancelled()) {
+            this.closeTx();
+        }
     }
 
     protected void cancelled() {

hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/DistributedTaskAdmissionTest.java

@@ -31,6 +31,7 @@
 
 import org.apache.hugegraph.config.HugeConfig;
 import org.apache.hugegraph.dist.RegisterUtil;
+import org.apache.hugegraph.exception.HugeException;
 import org.apache.hugegraph.id.Id;
 import org.apache.hugegraph.id.IdGenerator;
 import org.apache.hugegraph.meta.MetaManager;
@@ -190,6 +191,168 @@ public void testBusyWorkerLeavesPersistentBacklogInDatabase() throws Exception {
         }
     }
 
+    @Test
+    public void testCancellationFailurePreservesCallerTransaction() throws Exception {
+        try (Fixture fixture = new Fixture("failed_caller_tx")) {
+            fixture.worker.armed.set(true);
+            AtomicBoolean reject = new AtomicBoolean(true);
+            IllegalStateException failure = new IllegalStateException("Cancellation rejected");
+            HugeTask<?> task = new HugeTask<>(IdGenerator.of(-9999951L), null, new EphemeralJob<Object>() {
+                @Override
+                public String type() {
+                    return "test";
+                }
+
+                @Override
+                public Object execute() {
+                    return null;
+                }
+            }) {
+                @Override
+                public boolean cancel(boolean interrupt) {
+                    if (reject.get()) {
+                        throw failure;
+                    }
+                    return super.cancel(interrupt);
+                }
+            };
+            task.type("test");
+            task.name("failed-cancellation-caller-tx");
+            fixture.scheduler.schedule(task);
+            Assert.assertTrue(fixture.worker.reached.await(10L, TimeUnit.SECONDS));
+            try {
+                this.openCallerTransaction(fixture.graph);
+                Assert.assertThrows(HugeException.class, () -> fixture.scheduler.cancel(task),
+                                    caught -> Assert.assertSame(failure, caught.getCause()));
+                Assert.assertFalse(task.isCancelled());
+                this.assertCallerTransactionCommits(fixture);
+                Assert.assertEquals(1, fixture.scheduler.pendingTasks());
+            } finally {
+                reject.set(false);
+                fixture.worker.release.countDown();
+            }
+            fixture.assertWorkerClean();
+        }
+    }
+
+    @Test
+    public void testNormalWorkerTransactionCleanupControl() throws Exception {
+        try (Fixture fixture = new Fixture("normal_worker_control")) {
+            AtomicBoolean opened = new AtomicBoolean();
+            HugeTask<?> task = new HugeTask<>(IdGenerator.of(-9999950L), null, new EphemeralJob<Object>() {
+                @Override
+                public String type() {
+                    return "test";
+                }
+
+                @Override
+                public Object execute() {
+                    this.graph().vertices().hasNext();
+                    opened.set(this.graph().tx().isOpen());
+                    return null;
+                }
+            });
+            task.type("test");
+            task.name("normal-worker-cleanup");
+            fixture.scheduler.schedule(task).get(10L, TimeUnit.SECONDS);
+            Assert.assertTrue(opened.get());
+            Assert.assertFalse(task.isCancelled());
+            Assert.assertEquals(TaskStatus.SUCCESS, task.status());
+            fixture.assertWorkerClean();
+            Assert.assertFalse(fixture.worker.submit(() -> fixture.graph.tx().isOpen()).get(10L, TimeUnit.SECONDS));
+        }
+    }
+
+    @Test
+    public void testQueuedDefaultCancellationPreservesCallerTransaction() throws Exception {
+        this.assertDefaultCancellationPreservesCallerTransaction(false);
+    }
+
+    @Test
+    public void testRunningDefaultCancellationPreservesCallerTransaction() throws Exception {
+        this.assertDefaultCancellationPreservesCallerTransaction(true);
+    }
+
+    private void assertDefaultCancellationPreservesCallerTransaction(boolean running) throws Exception {
+        try (Fixture fixture = new Fixture("caller_tx_" + running)) {
+            CountDownLatch started = new CountDownLatch(1);
+            CountDownLatch finish = new CountDownLatch(1);
+            if (!running) {
+                fixture.worker.armed.set(true);
+            }
+            HugeTask<?> task = new HugeTask<>(IdGenerator.of(running ? -9999921L : -9999922L), null,
+                    new EphemeralJob<Object>() {
+                        @Override
+                        public String type() {
+                            return "test";
+                        }
+
+                        @Override
+                        public Object execute() {
+                            this.graph().vertices().hasNext();
+                            started.countDown();
+                            awaitUninterruptibly(finish);
+                            return null;
+                        }
+                    });
+            task.type("test");
+            task.name("default-callback-caller-tx");
+            try {
+                fixture.scheduler.schedule(task);
+                Assert.assertTrue((running ? started : fixture.worker.reached)
+                                  .await(10L, TimeUnit.SECONDS));
+                this.openCallerTransaction(fixture.graph);
+                fixture.scheduler.cancel(task);
+                Assert.assertTrue(task.isCancelled());
+                this.assertCallerTransactionCommits(fixture);
+                if (running) {
+                    Assert.assertEquals(1, fixture.scheduler.pendingTasks());
+                }
+            } finally {
+                finish.countDown();
+                fixture.worker.release.countDown();
+            }
+            fixture.assertWorkerClean();
+            Assert.assertEquals(running ? 0L : 1L, started.getCount());
+        }
+    }
+
+    @Test
+    public void testCompletedCancellationPreservesCallerTransaction() throws Exception {
+        try (Fixture fixture = new Fixture("completed_caller_tx")) {
+            HugeTask<?> task = task(true, new AtomicBoolean(), 9999923L);
+            fixture.scheduler.schedule(task).get(10L, TimeUnit.SECONDS);
+            fixture.assertWorkerClean();
+            this.openCallerTransaction(fixture.graph);
+            fixture.scheduler.cancel(task);
+            Assert.assertFalse(task.isCancelled());
+            this.assertCallerTransactionCommits(fixture);
+            Assert.assertEquals(0, fixture.scheduler.pendingTasks());
+        }
+    }
+
+    private void openCallerTransaction(HugeGraph graph) {
+        graph.schema().vertexLabel("caller").useCustomizeStringId().create();
+        graph.tx().commit();
+        graph.addVertex(org.apache.tinkerpop.gremlin.structure.T.id, "pending",
+                        org.apache.tinkerpop.gremlin.structure.T.label, "caller");
+        Assert.assertTrue(graph.tx().isOpen());
+    }
+
+    private void assertCallerTransactionCommits(Fixture fixture) throws Exception {
+        Assert.assertTrue(fixture.graph.tx().isOpen());
+        Assert.assertTrue(fixture.graph.vertices("pending").hasNext());
+        fixture.graph.tx().commit();
+        Assert.assertFalse(fixture.graph.tx().isOpen());
+        Assert.assertTrue(fixture.database.submit(() -> {
+            try {
+                return fixture.graph.vertices("pending").hasNext();
+            } finally {
+                fixture.params.closeTx();
+            }
+        }).get(10L, TimeUnit.SECONDS));
+    }
+
     @Test
     public void testQueuedOuterCancellationPreservesCallerTransaction() throws Exception {
         try (Fixture fixture = new Fixture("queued_future")) {