Merge pull request #3738 from soerenreichardt/revisit-progtracker-release

Release ProgressTracker when baseTask has finished

Author: soerenreichardt

alpha/alpha-proc/src/test/java/org/neo4j/gds/walking/CollapsePathMutateProcTest.java

@@ -20,20 +20,21 @@
 package org.neo4j.gds.walking;
 
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.AlgoBaseProc;
 import org.neo4j.gds.AlgoBaseProcTest;
 import org.neo4j.gds.BaseProcTest;
 import org.neo4j.gds.GdsCypher;
 import org.neo4j.gds.MutateRelationshipsTest;
+import org.neo4j.gds.RelationshipType;
+import org.neo4j.gds.api.Relationships;
 import org.neo4j.gds.catalog.GraphCreateProc;
 import org.neo4j.gds.catalog.GraphWriteNodePropertiesProc;
 import org.neo4j.gds.core.CypherMapWrapper;
 import org.neo4j.gds.extension.Neo4jGraph;
 import org.neo4j.gds.impl.walking.CollapsePath;
 import org.neo4j.gds.impl.walking.CollapsePathConfig;
-import org.neo4j.gds.RelationshipType;
-import org.neo4j.gds.api.Relationships;
 import org.neo4j.kernel.internal.GraphDatabaseAPI;
 
 import java.util.List;
@@ -154,4 +155,8 @@ void testMutateYields() {
             )
         );
     }
+
+    @Disabled
+    @Override
+    public void shouldUnregisterTaskAfterComputation() {}
 }

core/src/main/java/org/neo4j/gds/core/utils/progress/GlobalTaskStore.java

@@ -34,11 +34,11 @@ public class GlobalTaskStore implements TaskStore, ThrowingFunction<Context, Tas
 
     private final Map<String, Map<JobId, Task>> registeredTasks;
 
-    GlobalTaskStore() {
+    public GlobalTaskStore() {
         this.registeredTasks = new ConcurrentHashMap<>();
     }
 
-    void store(String username, JobId jobId, Task task) {
+    public void store(String username, JobId jobId, Task task) {
         this.registeredTasks
             .computeIfAbsent(username, __ -> new ConcurrentHashMap<>())
             .put(jobId, task);

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

@@ -95,9 +95,12 @@ public void endSubTask() {
         var currentTask = requireCurrentTask();
         taskProgressLogger.logEndSubTask(currentTask, parentTask());
         currentTask.finish();
-        this.currentTask = nestedTasks.isEmpty()
-            ? Optional.empty()
-            : Optional.of(nestedTasks.pop());
+        if (nestedTasks.isEmpty()) {
+            this.currentTask = Optional.empty();
+            release();
+        } else {
+            this.currentTask = Optional.of(nestedTasks.pop());
+        }
     }
 
     @Override

core/src/test/java/org/neo4j/gds/core/utils/progress/GlobalTaskStoreTest.java

@@ -0,0 +1,38 @@
+/*
+ * 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.core.utils.progress;
+
+import org.junit.jupiter.api.Test;
+import org.neo4j.gds.core.utils.progress.tasks.Tasks;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+
+class GlobalTaskStoreTest {
+
+    @Test
+    void shouldBeIdempotentOnRemove() {
+        var taskStore = new GlobalTaskStore();
+        var jobId = new JobId();
+        taskStore.store("", jobId, Tasks.leaf("leaf"));
+        taskStore.remove("", jobId);
+        assertDoesNotThrow(() -> taskStore.remove("", jobId));
+    }
+
+}

pregel/src/test/java/org/neo4j/gds/beta/pregel/PregelTest.java

@@ -209,7 +209,7 @@ void cleanupProgressLogging() {
         pregelAlgo.run();
         pregelAlgo.release();
 
-        assertThat(taskRegistry.unregisterTaskCalls()).isEqualTo(1);
+        assertThat(taskRegistry.unregisterTaskCalls()).isGreaterThanOrEqualTo(1);
     }
 
     static Stream<Arguments> forkJoinAndPartitioning() {

proc/common/src/test/java/org/neo4j/gds/AlgorithmCleanupTest.java

@@ -58,7 +58,7 @@ void cleanupTaskRegistryUnderRegularExecution() {
         Map<String, Object> config = Map.of("writeProperty", "test");
 
         assertThatCode(() -> proc.stats("g", config)).doesNotThrowAnyException();
-        assertThat(taskRegistry.unregisterTaskCalls()).isEqualTo(1);
+        assertThat(taskRegistry.unregisterTaskCalls()).isGreaterThanOrEqualTo(1);
     }
 
     @Test
@@ -72,6 +72,6 @@ void cleanupTaskRegistryWhenTheAlgorithmFails() {
         Map<String, Object> config = Map.of("writeProperty", "test", "throwInCompute", true);
 
         assertThatThrownBy(() -> proc.stats("g", config)).isNotNull();
-        assertThat(taskRegistry.unregisterTaskCalls()).isEqualTo(1);
+        assertThat(taskRegistry.unregisterTaskCalls()).isGreaterThanOrEqualTo(1);
     }
 }

proc/test/src/main/java/org/neo4j/gds/AlgoBaseProcTest.java

@@ -42,6 +42,10 @@
 import org.neo4j.gds.core.TransactionContext;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
 import org.neo4j.gds.core.utils.progress.EmptyTaskRegistry;
+import org.neo4j.gds.core.utils.progress.GlobalTaskStore;
+import org.neo4j.gds.core.utils.progress.JobId;
+import org.neo4j.gds.core.utils.progress.LocalTaskRegistry;
+import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.core.write.NativeNodePropertyExporter;
 import org.neo4j.gds.core.write.NativeRelationshipExporter;
 import org.neo4j.gds.core.write.NativeRelationshipStreamExporter;
@@ -269,6 +273,55 @@ default void assertMissingProperty(String error, Runnable runnable) {
         assertThat(exception).hasMessageContaining(error);
     }
 
+    class InvocationCountingTaskStore extends GlobalTaskStore {
+        int registerTaskInvocations;
+
+        @Override
+        public void store(
+            String username, JobId jobId, Task task
+        ) {
+            super.store(username, jobId, task);
+            registerTaskInvocations++;
+        }
+    }
+
+    @Test
+    default void shouldUnregisterTaskAfterComputation() {
+        var taskStore = new InvocationCountingTaskStore();
+
+        String loadedGraphName = "loadedGraph";
+        GraphCreateConfig graphCreateConfig = withNameAndRelationshipProjections("", loadedGraphName, relationshipProjections());
+        applyOnProcedure(proc -> {
+            proc.taskRegistry = new LocalTaskRegistry("", taskStore);
+
+            GraphStore graphStore = graphLoader(graphCreateConfig).graphStore();
+            GraphStoreCatalog.set(
+                graphCreateConfig,
+                graphStore
+            );
+            Map<String, Object> configMap = createMinimalConfig(CypherMapWrapper.empty()).toMap();
+            AlgoBaseProc.ComputationResult<?, RESULT, CONFIG> computationResult1 = proc.compute(
+                loadedGraphName,
+                configMap,
+                releaseAlgorithm(),
+                true
+            );
+
+            AlgoBaseProc.ComputationResult<?, RESULT, CONFIG> computationResult2 = proc.compute(
+                loadedGraphName,
+                configMap,
+                releaseAlgorithm(),
+                true
+            );
+
+            // trigger consumption of stream return values
+            assertResultEquals(computationResult1.result(), computationResult2.result());
+
+            assertThat(taskStore.taskStream()).isEmpty();
+            assertThat(taskStore.registerTaskInvocations).isGreaterThan(1);
+        });
+    }
+
     @AllGraphStoreFactoryTypesTest
     default void testRunOnLoadedGraph(TestSupport.FactoryType factoryType) {
         // FIXME rethink this test for mutate

proc/test/src/main/java/org/neo4j/gds/test/TestAlgorithm.java

@@ -46,7 +46,6 @@ public TestAlgorithm(
         this.progressTracker = progressTracker;
     }
 
-
     @Override
     public TestAlgorithm compute() {
         progressTracker.beginSubTask();