Remove TaskRegistry interface

Author: soerenreichardt

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

@@ -24,6 +24,6 @@ public enum EmptyTaskRegistryFactory implements TaskRegistryFactory {
 
     @Override
     public TaskRegistry newInstance() {
-        return new LocalTaskRegistry("", EmptyTaskStore.INSTANCE);
+        return new TaskRegistry("", EmptyTaskStore.INSTANCE);
     }
 }

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

@@ -31,6 +31,6 @@ public class LocalTaskRegistryFactory implements TaskRegistryFactory {
 
     @Override
     public TaskRegistry newInstance() {
-        return new LocalTaskRegistry(username, taskStore);
+        return new TaskRegistry(username, taskStore);
     }
 }

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

@@ -19,11 +19,31 @@
  */
 package org.neo4j.gds.core.utils.progress;
 
+import org.jetbrains.annotations.TestOnly;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
 
-public interface TaskRegistry {
+public class TaskRegistry {
 
-    void registerTask(Task task);
+    private final String username;
+    private final TaskStore taskStore;
+    private final JobId jobId;
 
-    void unregisterTask();
+    @TestOnly
+    public TaskRegistry(TaskRegistry taskRegistry) {
+        this(taskRegistry.username, taskRegistry.taskStore);
+    }
+
+    public TaskRegistry(String username, TaskStore taskStore) {
+        this.username = username;
+        this.taskStore = taskStore;
+        this.jobId = new JobId();
+    }
+
+    public void registerTask(Task task) {
+        taskStore.store(username, jobId, task);
+    }
+
+    public void unregisterTask() {
+        taskStore.remove(username, jobId);
+    }
 }

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

@@ -29,7 +29,7 @@ class TaskRegistryTest {
     @Test
     void shouldStoreIncomingTasks() {
         var globalTaskStore = new GlobalTaskStore();
-        var taskRegistry1 = new LocalTaskRegistry("", globalTaskStore);
+        var taskRegistry1 = new TaskRegistry("", globalTaskStore);
 
         assertThat(globalTaskStore.isEmpty()).isTrue();
 
@@ -39,7 +39,7 @@ void shouldStoreIncomingTasks() {
         assertThat(globalTaskStore.query("")).containsValue(task1);
         assertThat(globalTaskStore.isEmpty()).isFalse();
 
-        var taskRegistry2 = new LocalTaskRegistry("", globalTaskStore);
+        var taskRegistry2 = new TaskRegistry("", globalTaskStore);
         var task2 = Tasks.leaf("task2");
         taskRegistry2.registerTask(task2);
 
@@ -50,7 +50,7 @@ void shouldStoreIncomingTasks() {
     @Test
     void shouldRemoveStoredTasks() {
         var globalTaskStore = new GlobalTaskStore();
-        var taskRegistry = new LocalTaskRegistry("", globalTaskStore);
+        var taskRegistry = new TaskRegistry("", globalTaskStore);
 
         var task = Tasks.leaf("task");
         taskRegistry.registerTask(task);

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

@@ -44,7 +44,7 @@
 import org.neo4j.gds.core.utils.mem.AllocationTracker;
 import org.neo4j.gds.core.utils.mem.MemoryRange;
 import org.neo4j.gds.core.utils.paged.HugeDoubleArray;
-import org.neo4j.gds.core.utils.progress.LocalTaskRegistry;
+import org.neo4j.gds.core.utils.progress.TaskRegistry;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.TaskProgressTracker;
 import org.neo4j.gds.extension.GdlExtension;
@@ -196,7 +196,7 @@ void cleanupProgressLogging() {
 
         var task = Pregel.progressTask(graph, config, computation.getClass().getSimpleName());
         var progressLogger =  new TestProgressLogger(task, config.concurrency());
-        var progressTracker = new TaskProgressTracker(task, progressLogger, () -> new LocalTaskRegistry("", taskStore));
+        var progressTracker = new TaskProgressTracker(task, progressLogger, () -> new TaskRegistry("", taskStore));
 
         var pregelAlgo = Pregel.create(
             graph,

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

@@ -24,7 +24,7 @@
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.catalog.GraphCreateProc;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
-import org.neo4j.gds.core.utils.progress.LocalTaskRegistry;
+import org.neo4j.gds.core.utils.progress.TaskRegistry;
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
 import org.neo4j.gds.test.TestProc;
 import org.neo4j.logging.NullLog;
@@ -52,7 +52,7 @@ void teardown() {
     @Test
     void cleanupTaskRegistryUnderRegularExecution() {
         var taskStore = new TestTaskStore();
-        var taskRegistryFactory = (TaskRegistryFactory) () -> new LocalTaskRegistry(getUsername(), taskStore);
+        var taskRegistryFactory = (TaskRegistryFactory) () -> new TaskRegistry(getUsername(), taskStore);
 
         var proc = new TestProc();
         proc.taskRegistryFactory = taskRegistryFactory;
@@ -67,7 +67,7 @@ void cleanupTaskRegistryUnderRegularExecution() {
     @Test
     void cleanupTaskRegistryWhenTheAlgorithmFails() {
         var taskStore = new TestTaskStore();
-        var taskRegistryFactory = (TaskRegistryFactory) () -> new LocalTaskRegistry(getUsername(), taskStore);
+        var taskRegistryFactory = (TaskRegistryFactory) () -> new TaskRegistry(getUsername(), taskStore);
 
         var proc = new TestProc();
         proc.taskRegistryFactory = taskRegistryFactory;

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

@@ -26,7 +26,6 @@
 import org.neo4j.gds.core.utils.progress.TaskRegistry;
 import org.neo4j.gds.core.utils.progress.TaskRegistryExtension;
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
-import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.core.utils.progress.tasks.TaskProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.Tasks;
 import org.neo4j.logging.Level;
@@ -86,17 +85,10 @@ public static class Bar {
         public Bar(String field) {this.field = field;}
     }
 
-    public static class NonReleasingTaskRegistry implements TaskRegistry {
-
-        private final TaskRegistry taskRegistry;
+    public static class NonReleasingTaskRegistry extends TaskRegistry {
 
         NonReleasingTaskRegistry(TaskRegistry taskRegistry) {
-            this.taskRegistry = taskRegistry;
-        }
-
-        @Override
-        public void registerTask(Task task) {
-            taskRegistry.registerTask(task);
+            super(taskRegistry);
         }
 
         @Override

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

@@ -150,7 +150,8 @@ public Stream<FastRPStreamProc.StreamResult> fakeStream(
             @Name(value = "graphName") Object graphNameOrConfig,
             @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration
         ) {
-            this.taskRegistryFactory = () -> new NonReleasingTaskRegistry(taskRegistryFactory.newInstance());
+            var taskRegistry = taskRegistryFactory.newInstance();
+            this.taskRegistryFactory = () -> new NonReleasingTaskRegistry(taskRegistry);
             return super.stream(graphNameOrConfig, configuration);
         }