Cleanup on CypherAggregation::close in failure case

Author: FlorentinD

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

@@ -124,7 +124,7 @@ private static Stream<Arguments> emptyGraphNames() {
     }
 
     @Test
-    void shouldFailTaskOnFailure() {
+    void shouldFailTaskOnFailure() throws Exception {
         TestTaskStore taskStore = new TestTaskStore();
         var aggregator = new ProductGraphAggregator(
             DatabaseId.random(),
@@ -149,6 +149,9 @@ void shouldFailTaskOnFailure() {
             .hasCauseInstanceOf(IllegalArgumentException.class)
             .hasMessageContaining("The node has to be either a NODE or an INTEGER, but got String");
 
+        // assuming this gets called by Neo4j
+        aggregator.close();
+
         assertThat(taskStore.query())
             .map(i -> i.task().status())
             .containsExactly(Status.FAILED);

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

@@ -72,7 +72,6 @@ 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/CypherAggregation.java

@@ -98,7 +98,7 @@ public UserAggregationReducer createReducer(Context ctx) throws ProcedureExcepti
 
         var queryEstimator = QueryEstimator.fromTransaction(DatabaseTransactionContext.of(databaseService, transaction));
 
-        return new ProductGraphAggregator(
+        ProductGraphAggregator productGraphAggregator = new ProductGraphAggregator(
             DatabaseId.of(databaseService.databaseName()),
             username,
             writeMode,
@@ -108,5 +108,9 @@ public UserAggregationReducer createReducer(Context ctx) throws ProcedureExcepti
             taskStore,
             new LogAdapter(log)
         );
+
+        ctx.internalTransaction().registerCloseableResource(productGraphAggregator);
+
+        return productGraphAggregator;
     }
 }

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

@@ -70,7 +70,7 @@
 import static org.neo4j.gds.projection.GraphImporter.NO_TARGET_NODE;
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
 
-abstract class GraphAggregator implements UserAggregationReducer, UserAggregationUpdater {
+abstract class GraphAggregator implements UserAggregationReducer, UserAggregationUpdater, AutoCloseable {
 
     static final String SOURCE_NODE_PROPERTIES = "sourceNodeProperties";
     static final String SOURCE_NODE_LABELS = "sourceNodeLabels";
@@ -95,6 +95,7 @@ abstract class GraphAggregator implements UserAggregationReducer, UserAggregatio
     // Used for initializing the data and rel importers
     private final Lock lock;
     private final ExtractNodeId extractNodeId;
+    private final AtomicBoolean completedSuccessfully;
     private volatile @Nullable GraphImporter importer;
 
     // #result() may be called twice, we cache the result of the first call to return it again in the second invocation
@@ -123,6 +124,7 @@ abstract class GraphAggregator implements UserAggregationReducer, UserAggregatio
         this.lock = new ReentrantLock();
         this.configValidator = new ConfigValidator();
         this.extractNodeId = new ExtractNodeId();
+        this.completedSuccessfully = new AtomicBoolean(false);
     }
 
     void projectNextRelationship(
@@ -284,7 +286,6 @@ public AnyValue result() throws ProcedureException {
             projectionMetric.start();
             result = buildGraph();
         } catch (Exception e) {
-            this.onFailure();
             projectionMetric.failed(e);
             throw new ProcedureException(
                 Status.Procedure.ProcedureCallFailed,
@@ -306,13 +307,11 @@ public AnyValue result() throws ProcedureException {
         builder.add("projectMillis", Values.longValue(result.projectMillis()));
         builder.add("configuration", ValueUtils.asAnyValue(result.configuration()));
         builder.add("query", ValueUtils.asAnyValue(result.query()));
-        return builder.build();
-    }
+        MapValue projectResult = builder.build();
 
-    void onFailure() {
-        if (progressTracker != null) {
-            this.progressTracker.endSubTaskWithFailure();
-        }
+        this.completedSuccessfully.set(true);
+
+        return projectResult;
     }
 
     public @Nullable AggregationResult buildGraph() {
@@ -390,6 +389,13 @@ private static RelationshipType typeConfig(
         );
     }
 
+    @Override
+    public void close() throws Exception {
+        if (!completedSuccessfully.get() && progressTracker != null) {
+            this.progressTracker.endSubTaskWithFailure();
+        }
+    }
+
     public static final class ConfigValidator {
         private static final Set<String> DATA_CONFIG_KEYS = Set.of(
             SOURCE_NODE_PROPERTIES,

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

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