Add failed state to Tasks and prevent PT release to mask failure

Author: breakanalysis

core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/ProgressTracker.java

@@ -52,6 +52,8 @@ default void logProgress() {
 
     void release();
 
+    void fail();
+
     class EmptyProgressTracker implements ProgressTracker {
 
         @Override
@@ -102,5 +104,10 @@ public ProgressLogger progressLogger() {
         @Override
         public void release() {
         }
+
+        @Override
+        public void fail() {
+
+        }
     }
 }

core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/Status.java

@@ -23,5 +23,6 @@ public enum Status {
     PENDING,
     RUNNING,
     FINISHED,
-    CANCELED
+    CANCELED,
+    FAILED
 }

core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/Task.java

@@ -178,6 +178,10 @@ public void setEstimatedMemoryRangeInBytes(MemoryRange memoryRangeInBytes) {
         this.estimatedMemoryRangeInBytes = memoryRangeInBytes;
     }
 
+    public void fail() {
+        this.status = Status.FAILED;
+    }
+
     protected Task nextSubTaskAfterValidation() {
         if (subTasks.stream().anyMatch(t -> t.status == Status.RUNNING)) {
             throw new IllegalStateException("Cannot move to next subtask, because some subtasks are still running");

core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTracker.java

@@ -127,7 +127,12 @@ public ProgressLogger progressLogger() {
     @Override
     public void release() {
         taskRegistry.unregisterTask();
-        validateTaskFinishedOrCanceled();
+        validateTaskNotRunning();
+    }
+
+    @Override
+    public void fail() {
+        baseTask.fail();
     }
 
     @TestOnly
@@ -150,7 +155,7 @@ private Task requireCurrentTask() {
         return currentTask.orElseThrow(() -> new IllegalStateException("No more running tasks"));
     }
 
-    private void validateTaskFinishedOrCanceled() {
+    private void validateTaskNotRunning() {
         if (baseTask.status() == Status.RUNNING) {
             var message = formatWithLocale(
                 "Attempted to release algorithm, but task %s is still running",

proc/centrality/src/test/java/org/neo4j/gds/testproc/FailingAlgorithm.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.testproc;
+
+import org.neo4j.gds.Algorithm;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+
+class FailingAlgorithm extends Algorithm<FailingAlgorithm, ProcedureThatFailsDuringTask.Output> {
+    FailingAlgorithm(ProgressTracker progressTracker) {
+        super();
+        this.progressTracker = progressTracker;
+    }
+
+    @Override
+    public ProcedureThatFailsDuringTask.Output compute() {
+        progressTracker.beginSubTask();
+        throw new IllegalStateException("Oops");
+    }
+
+    @Override
+    public FailingAlgorithm me() {
+        return this;
+    }
+
+    @Override
+    public void release() {
+
+    }
+}

proc/centrality/src/test/java/org/neo4j/gds/testproc/ProcedureFailTest.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.testproc;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.neo4j.gds.BaseProcTest;
+import org.neo4j.gds.extension.Neo4jGraph;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class ProcedureFailTest extends BaseProcTest {
+
+    @Neo4jGraph
+    public static final String DB_CYPHER = "CREATE (a:N)";
+
+    @BeforeEach
+    void setup() throws Exception {
+        registerProcedures(ProcedureThatFailsDuringTask.class);
+    }
+
+    @Test
+    void shouldFailWithIllegalStateException() {
+        assertThatThrownBy(() -> runQuery("CALL very.strange.procedure({nodeProjection: '*', relationshipProjection: '*'})"))
+            .getRootCause()
+            .isInstanceOf(IllegalStateException.class)
+            .hasMessageContaining("Oops");
+    }
+
+}

proc/centrality/src/test/java/org/neo4j/gds/testproc/ProcedureThatFailsDuringTask.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.testproc;
+
+import org.neo4j.gds.AlgoBaseProc;
+import org.neo4j.gds.AlgorithmFactory;
+import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.config.GraphCreateConfig;
+import org.neo4j.gds.core.CypherMapWrapper;
+import org.neo4j.gds.core.utils.mem.AllocationTracker;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.test.config.DummyConfig;
+import org.neo4j.gds.test.config.DummyConfigImpl;
+import org.neo4j.procedure.Mode;
+import org.neo4j.procedure.Name;
+import org.neo4j.procedure.Procedure;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+public class ProcedureThatFailsDuringTask extends AlgoBaseProc<FailingAlgorithm, ProcedureThatFailsDuringTask.Output, DummyConfig> {
+    @Procedure(name = "very.strange.procedure", mode = Mode.READ)
+    public Stream<Output> run(
+        @Name(value = "graphName") Object graphNameOrConfig,
+        @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration
+    ) {
+        var result = compute(graphNameOrConfig, configuration);
+        assert result.result() == null;
+        return Stream.empty();
+    }
+
+    @Override
+    protected DummyConfig newConfig(
+        String username,
+        Optional<String> graphName,
+        Optional<GraphCreateConfig> maybeImplicitCreate,
+        CypherMapWrapper config
+    ) {
+        return new DummyConfigImpl(graphName,  maybeImplicitCreate, username, config);
+    }
+
+    @Override
+    protected AlgorithmFactory<FailingAlgorithm, DummyConfig> algorithmFactory() {
+        return new AlgorithmFactory<FailingAlgorithm, DummyConfig>() {
+            @Override
+            protected String taskName() {
+                return "Failing Algorithm";
+            }
+
+            @Override
+            protected FailingAlgorithm build(
+                Graph graph,
+                DummyConfig configuration,
+                AllocationTracker allocationTracker,
+                ProgressTracker progressTracker
+            ) {
+                return new FailingAlgorithm(progressTracker);
+            }
+       };
+    }
+
+    public class Output {
+        public Object out;
+    }
+}

proc/common/src/main/java/org/neo4j/gds/AlgoBaseProc.java

@@ -232,6 +232,9 @@ protected ComputationResult<ALGO, ALGO_RESULT, CONFIG> compute(
             () -> {
                 try (ProgressTimer ignored = ProgressTimer.start(builder::computeMillis)) {
                     return algo.compute();
+                } catch (Exception e) {
+                    algo.progressTracker.fail();
+                    throw e;
                 } finally {
                     if (releaseAlgorithm) {
                         algo.progressTracker.release();

proc/test/src/main/java/org/neo4j/gds/test/config/DummyConfig.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.test.config;
+
+import org.neo4j.gds.annotation.Configuration;
+import org.neo4j.gds.annotation.ValueClass;
+import org.neo4j.gds.config.AlgoBaseConfig;
+
+@Configuration
+@ValueClass
+public interface DummyConfig extends AlgoBaseConfig {
+}