From 0bff7ee025a31a917b251d71adfe11d6d3d6b64b Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 14 Oct 2026 08:45:36 +0800 Subject: [PATCH] fix(task): preserve caller tx on cancel Preserve caller transactions during cancellation. Keep worker cleanup in its execution path. Cover direct and failed cancellation with real tests. Document transaction ownership and fix the fixture. --- docs/server-runtime.md | 8 + .../task/DistributedTaskScheduler.java | 18 +- .../hugegraph/task/StandardTaskScheduler.java | 7 +- .../apache/hugegraph/task/TaskCallable.java | 6 +- .../DistributedTaskAdmissionTest.java | 163 ++++++++++++++++++ .../hugegraph/StandardTaskAdmissionTest.java | 162 +++++++++++++++++ .../core/TaskSchedulerServerInfoTest.java | 2 + 7 files changed, 357 insertions(+), 9 deletions(-) diff --git a/docs/server-runtime.md b/docs/server-runtime.md index f370951a7f..1d7f0c0802 100644 --- a/docs/server-runtime.md +++ b/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 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 index 0fe0e5b8f2..658596e240 100644 --- 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 @@ -334,15 +334,18 @@ public void cancel(HugeTask task) { String context = TaskManager.getContext(); AtomicBoolean cleaned = new AtomicBoolean(true); TaskManager.ContextCallable 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 action, boole pending.callbacks++; } TaskManager.ContextCallable 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 cancellation = () -> { diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java index 8d91069553..b53742a54e 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java @@ -348,6 +348,7 @@ public void initTaskCallable(HugeTask task) { @Override public void cancel(HugeTask 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 void cancel(HugeTask 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--; } diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/TaskCallable.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/TaskCallable.java index 5872acd9dd..9d7c9a6b79 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/TaskCallable.java +++ b/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() { diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/DistributedTaskAdmissionTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/DistributedTaskAdmissionTest.java index bddb0c3ec3..6ab384419b 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/DistributedTaskAdmissionTest.java +++ b/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() { + @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() { + @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() { + @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")) { diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/StandardTaskAdmissionTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/StandardTaskAdmissionTest.java index 1425693fa3..78a8e66b02 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/StandardTaskAdmissionTest.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/StandardTaskAdmissionTest.java @@ -45,6 +45,168 @@ public class StandardTaskAdmissionTest { + @Test + public void testCancellationFailurePreservesCallerTransaction() throws Exception { + try (Fixture fixture = new Fixture("standard_failed_caller_tx")) { + CountDownLatch occupied = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + fixture.worker.submit(() -> { + occupied.countDown(); + await(release); + }); + Assert.assertTrue(occupied.await(10L, TimeUnit.SECONDS)); + AtomicBoolean reject = new AtomicBoolean(true); + IllegalStateException failure = new IllegalStateException("Cancellation rejected"); + HugeTask task = new HugeTask<>(IdGenerator.of(9999952L), null, new EmptyJob()) { + @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); + try { + fixture.graph.schema().vertexLabel("caller").useCustomizeStringId().create(); + fixture.graph.tx().commit(); + fixture.graph.addVertex(T.id, "pending", T.label, "caller"); + Assert.assertThrows(IllegalStateException.class, () -> fixture.scheduler.cancel(task), + caught -> Assert.assertSame(failure, caught)); + Assert.assertFalse(task.isCancelled()); + Assert.assertTrue(fixture.graph.tx().isOpen()); + fixture.graph.tx().commit(); + Assert.assertTrue(fixture.database.submit(() -> { + try { + return fixture.graph.vertices("pending").hasNext(); + } finally { + fixture.params.closeTx(); + } + }).get(10L, TimeUnit.SECONDS)); + Assert.assertEquals(1, fixture.scheduler.pendingTasks()); + } finally { + reject.set(false); + release.countDown(); + } + fixture.worker.submit(() -> null).get(10L, TimeUnit.SECONDS); + Assert.assertEquals(0, fixture.scheduler.pendingTasks()); + } + } + + @Test + public void testNormalWorkerTransactionCleanupControl() throws Exception { + try (Fixture fixture = new Fixture("standard_normal_worker_control")) { + BlockingJob job = new BlockingJob(); + job.release.countDown(); + HugeTask task = fixture.task(job); + fixture.scheduler.schedule(task).get(10L, TimeUnit.SECONDS); + fixture.worker.submit(() -> null).get(10L, TimeUnit.SECONDS); + Assert.assertFalse(task.isCancelled()); + Assert.assertEquals(TaskStatus.SUCCESS, task.status()); + Assert.assertEquals(0L, job.started.getCount()); + Assert.assertEquals(0, fixture.scheduler.pendingTasks()); + ThreadLocal owners = Whitebox.getInternalState(fixture.graph.tx(), "transactions"); + Assert.assertNull(fixture.worker.submit(owners::get).get(10L, TimeUnit.SECONDS)); + Assert.assertFalse(fixture.worker.submit(() -> fixture.graph.tx().isOpen()).get(10L, TimeUnit.SECONDS)); + } + } + + @Test + public void testDirectDefaultDoneCancellationPreservesCallerTransaction() throws Exception { + try (Fixture fixture = new Fixture("standard_direct_done_control")) { + CountDownLatch occupied = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + fixture.worker.submit(() -> { + occupied.countDown(); + await(release); + }); + Assert.assertTrue(occupied.await(10L, TimeUnit.SECONDS)); + HugeTask task = fixture.task(new EmptyJob()); + try { + fixture.scheduler.schedule(task); + fixture.graph.schema().vertexLabel("caller").useCustomizeStringId().create(); + fixture.graph.tx().commit(); + fixture.graph.addVertex(T.id, "pending", T.label, "caller"); + // Direct FutureTask cancellation invokes the inherited TaskCallable.done before outer guards. + task.cancel(true); + Assert.assertTrue(fixture.graph.tx().isOpen()); + Assert.assertTrue(fixture.graph.vertices("pending").hasNext()); + fixture.graph.tx().commit(); + Assert.assertTrue(fixture.database.submit(() -> { + try { + return fixture.graph.vertices("pending").hasNext(); + } finally { + fixture.params.closeTx(); + } + }).get(10L, TimeUnit.SECONDS)); + } finally { + release.countDown(); + } + fixture.worker.submit(() -> null).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("standard_caller_tx_" + running)) { + CountDownLatch occupied = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + if (!running) { + fixture.worker.submit(() -> { + occupied.countDown(); + await(release); + }); + Assert.assertTrue(occupied.await(10L, TimeUnit.SECONDS)); + } + BlockingJob job = new BlockingJob(); + HugeTask task = fixture.task(job); + try { + fixture.scheduler.schedule(task); + if (running) { + Assert.assertTrue(job.started.await(10L, TimeUnit.SECONDS)); + } + fixture.graph.schema().vertexLabel("caller").useCustomizeStringId().create(); + fixture.graph.tx().commit(); + fixture.graph.addVertex(T.id, "pending", T.label, "caller"); + fixture.scheduler.cancel(task); + Assert.assertTrue(task.isCancelled()); + 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)); + if (running) { + Assert.assertEquals(1, fixture.scheduler.pendingTasks()); + } + } finally { + release.countDown(); + job.release.countDown(); + } + fixture.worker.submit(() -> null).get(10L, TimeUnit.SECONDS); + Assert.assertEquals(0, fixture.scheduler.pendingTasks()); + ThreadLocal owners = Whitebox.getInternalState(fixture.graph.tx(), "transactions"); + Assert.assertNull(fixture.worker.submit(owners::get).get(10L, TimeUnit.SECONDS)); + Assert.assertEquals(running ? 0L : 1L, job.started.getCount()); + } + } + @Test public void testCancelledWorkerFencesAdmissionUntilActualExit() throws Exception { try (Fixture fixture = new Fixture("standard_cancel_worker")) { diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/TaskSchedulerServerInfoTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/TaskSchedulerServerInfoTest.java index feef2f623c..8d6ea7db28 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/TaskSchedulerServerInfoTest.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/TaskSchedulerServerInfoTest.java @@ -45,6 +45,7 @@ import org.apache.hugegraph.type.define.NodeRole; import org.apache.hugegraph.unit.FakeObjects; import org.apache.hugegraph.util.ExecutorUtil; +import org.apache.tinkerpop.gremlin.structure.Transaction; import org.junit.Test; import org.mockito.Mockito; @@ -89,6 +90,7 @@ public void testDistributedCheckRequirementDoesNotNeedServerInfo() { public void testDistributedCancelTreatsMissingRetryTaskAsGone() { HugeGraph graph = Mockito.mock(HugeGraph.class); Mockito.when(graph.graphSpace()).thenReturn("DEFAULT"); + Mockito.when(graph.tx()).thenReturn(Mockito.mock(Transaction.class)); HugeGraphParams params = Mockito.mock(HugeGraphParams.class); Mockito.when(params.graph()).thenReturn(graph);