Merge pull request #9807 from FlorentinD/gds/progress-tracker-cleanup

Cleanup Graph Aggregation tasks in case of failure

Author: FlorentinD

cypher-aggregation/src/integrationTest/java/org/neo4j/gds/projection/ProductGraphAggregatorTest.java

@@ -23,12 +23,14 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.neo4j.gds.TestTaskStore;
 import org.neo4j.gds.api.DatabaseId;
 import org.neo4j.gds.core.loading.Capabilities.WriteMode;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
 import org.neo4j.gds.core.utils.progress.EmptyTaskStore;
 import org.neo4j.gds.logging.Log;
 import org.neo4j.gds.metrics.projections.ProjectionMetricsService;
+import org.neo4j.values.AnyValue;
 import org.neo4j.values.storable.NoValue;
 import org.neo4j.values.storable.Values;
 import org.neo4j.values.virtual.MapValue;
@@ -37,6 +39,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 class ProductGraphAggregatorTest {
 
@@ -85,14 +88,15 @@ void shouldImportHighNodeIds() {
     @MethodSource("emptyGraphNames")
     void shouldFailOnEmptyGraphName(String emptyGraphName, String description) {
 
+        TestTaskStore taskStore = new TestTaskStore();
         var aggregator = new ProductGraphAggregator(
             DatabaseId.random(),
             "neo4j",
             WriteMode.LOCAL,
             QueryEstimator.empty(),
             ExecutingQueryProvider.empty(),
             ProjectionMetricsService.DISABLED,
-            EmptyTaskStore.INSTANCE,
+            taskStore,
             Log.noOpLog()
         );
 
@@ -105,6 +109,8 @@ void shouldFailOnEmptyGraphName(String emptyGraphName, String description) {
                 MapValue.EMPTY,
                 NoValue.NO_VALUE
             )).withMessageContaining("`graphName` can not be null or blank");
+
+        assertThat(taskStore.tasks()).isEmpty();
     }
 
     private static Stream<Arguments> emptyGraphNames() {
@@ -115,4 +121,33 @@ private static Stream<Arguments> emptyGraphNames() {
             Arguments.of("   ", "spaces")
         );
     }
+
+    @Test
+    void shouldCleanupTaskOnFailure() {
+        TestTaskStore taskStore = new TestTaskStore();
+        var aggregator = new ProductGraphAggregator(
+            DatabaseId.random(),
+            "neo4j",
+            WriteMode.LOCAL,
+            QueryEstimator.empty(),
+            ExecutingQueryProvider.empty(),
+            ProjectionMetricsService.DISABLED,
+            taskStore,
+            Log.noOpLog()
+        );
+
+        assertThatThrownBy(() ->
+            aggregator.update(new AnyValue[] {
+                Values.stringValue("my-graph"),
+                Values.longValue(1L),
+                Values.stringValue("invalidID"),
+                MapValue.EMPTY,
+                MapValue.EMPTY,
+                NoValue.NO_VALUE }
+            ))
+            .hasCauseInstanceOf(IllegalArgumentException.class)
+            .hasMessageContaining("The node has to be either a NODE or an INTEGER, but got String");
+
+        assertThat(taskStore.tasks()).isEmpty();
+    }
 }

cypher-aggregation/src/main/java/org/neo4j/gds/projection/AlphaGraphAggregator.java

@@ -72,6 +72,7 @@ public void update(AnyValue[] input) throws ProcedureException {
                 NoValue.NO_VALUE
             );
         } catch (Exception e) {
+            super.onFailure();
             throw new ProcedureException(
                 Status.Procedure.ProcedureCallFailed,
                 e,

cypher-aggregation/src/main/java/org/neo4j/gds/projection/GraphAggregator.java

@@ -41,6 +41,7 @@
 import org.neo4j.gds.core.utils.progress.BatchingTaskProgressTracker;
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.TaskStore;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.TaskProgressTracker;
 import org.neo4j.gds.core.utils.warnings.EmptyUserLogRegistryFactory;
 import org.neo4j.gds.logging.Log;
@@ -97,6 +98,7 @@ abstract class GraphAggregator implements UserAggregationReducer, UserAggregatio
 
     // #result() may be called twice, we cache the result of the first call to return it again in the second invocation
     private @Nullable AggregationResult result;
+    private ProgressTracker progressTracker;
 
     GraphAggregator(
         DatabaseId databaseId,
@@ -213,7 +215,7 @@ private GraphImporter createGraphImporter(
             TaskRegistryFactory.local(username, taskStore),
             EmptyUserLogRegistryFactory.INSTANCE
         );
-        var progressTracker = BatchingTaskProgressTracker.create(internalProgressTracker, taskVolume, config.readConcurrency());
+        this.progressTracker = BatchingTaskProgressTracker.create(internalProgressTracker, taskVolume, config.readConcurrency());
 
         return new GraphImporter(
             config,
@@ -281,6 +283,7 @@ public AnyValue result() throws ProcedureException {
             projectionMetric.start();
             result = buildGraph();
         } catch (Exception e) {
+            this.onFailure();
             projectionMetric.failed(e);
             throw new ProcedureException(
                 Status.Procedure.ProcedureCallFailed,
@@ -305,6 +308,12 @@ public AnyValue result() throws ProcedureException {
         return builder.build();
     }
 
+    void onFailure() {
+        if (progressTracker != null) {
+            this.progressTracker.endSubTaskWithFailure();
+        }
+    }
+
     public @Nullable AggregationResult buildGraph() {
         var importer = this.importer;
         if (importer == null) {

cypher-aggregation/src/main/java/org/neo4j/gds/projection/ProductGraphAggregator.java

@@ -58,6 +58,7 @@ public void update(AnyValue[] input) throws ProcedureException {
                 input[5]
             );
         } catch (Exception e) {
+            super.onFailure();
             throw new ProcedureException(
                 Status.Procedure.ProcedureCallFailed,
                 e,

triplet-graph-builder/src/test/java/org/neo4j/gds/projection/GraphImporterTest.java

@@ -451,5 +451,7 @@ void shouldRegisterTaskAndLogProgress() {
         log.assertContainsMessage(TestLog.INFO, "Graph aggregation :: Build graph store :: Relationships :: Finished");
         log.assertContainsMessage(TestLog.INFO, "Graph aggregation :: Build graph store :: Finished");
         log.assertContainsMessage(TestLog.INFO, "Graph aggregation :: Finished");
-    }
+
+        assertThat(taskStore.tasks()).isEmpty();
+     }
 }