Introduce setSteps and logSteps in ProgressTracker

Co-Authored-By: Adam Schill Collberg<adam.schill.collberg@protonmail.com>

Author: breakanalysis

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

@@ -61,6 +61,10 @@ default void logProgress() {
 
     void release();
 
+    void setSteps(long steps);
+
+    void logSteps(long steps);
+
     class EmptyProgressTracker implements ProgressTracker {
 
         @Override
@@ -108,6 +112,16 @@ public void logProgress(long value, String messageTemplate) {
         public void setVolume(long volume) {
         }
 
+        @Override
+        public void setSteps(long steps) {
+
+        }
+
+        @Override
+        public void logSteps(long steps) {
+
+        }
+
         @Override
         public void logDebug(String message) {
 

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

@@ -37,12 +37,16 @@
 
 public class TaskProgressTracker implements ProgressTracker {
 
+    private static final long UNKNOWN_STEPS = -1;
+
     private final Task baseTask;
     private final TaskRegistry taskRegistry;
     private final UserLogRegistry userLogRegistry;
     private final TaskProgressLogger taskProgressLogger;
     private final Stack<Task> nestedTasks;
     protected Optional<Task> currentTask;
+    private long currentTotalSteps;
+    private double progressLeftOvers;
 
     public TaskProgressTracker(Task baseTask, Log log, int concurrency, TaskRegistryFactory taskRegistryFactory) {
         this(baseTask, log, concurrency, new JobId(), taskRegistryFactory, EmptyUserLogRegistryFactory.INSTANCE);
@@ -56,6 +60,8 @@ public TaskProgressTracker(
         this.taskRegistry = taskRegistryFactory.newInstance(jobId);
         this.taskProgressLogger = new TaskProgressLogger(log, baseTask, concurrency);
         this.currentTask = Optional.empty();
+        this.currentTotalSteps = UNKNOWN_STEPS;
+        this.progressLeftOvers = 0;
         this.nestedTasks = new Stack<>();
         this.userLogRegistry = userLogRegistryFactory.newInstance();
     }
@@ -76,6 +82,8 @@ public void beginSubTask() {
         nextTask.start();
         taskProgressLogger.logBeginSubTask(nextTask, parentTask());
         currentTask = Optional.of(nextTask);
+        currentTotalSteps = UNKNOWN_STEPS;
+        progressLeftOvers = 0;
     }
 
     @Override
@@ -90,6 +98,26 @@ public void beginSubTask(long taskVolume) {
         setVolume(taskVolume);
     }
 
+    @Override
+    public void setSteps(long steps) {
+        if (steps <= 0) {
+            throw new IllegalStateException(formatWithLocale(
+                "Total steps for task must be at least 1 but was %d",
+                steps
+            ));
+        }
+        currentTotalSteps = steps;
+    }
+
+    @Override
+    public void logSteps(long steps) {
+        long volume = requireCurrentTask().getProgress().volume();
+        double progress = steps * volume / (double) currentTotalSteps + progressLeftOvers;
+        long longProgress = (long) progress;
+        progressLeftOvers = progress - longProgress;
+        logProgress(longProgress);
+    }
+
     @Override
     public void beginSubTask(String expectedTaskDescription, long taskVolume) {
         beginSubTask();

core/src/test/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTrackerTest.java

@@ -208,6 +208,24 @@ void shouldRegisterBaseTaskOnBaseTaskStart() {
         assertThat(taskStore.query("")).containsValue(task);
     }
 
+    @Test
+    void stepsShouldGiveProgress() {
+        var leafTask = Tasks.leaf("leaf", 100);
+        var progressTracker = progressTracker(leafTask);
+        
+        progressTracker.beginSubTask();
+        progressTracker.setSteps(13);
+        progressTracker.logProgress(3);
+        progressTracker.logSteps(1);
+        double expectedDoubleProgressFromFirstStep = 100.0 * 1.0 / 13.0;
+        long progressAfterFirstStep = leafTask.getProgress().progress();
+        assertThat(progressAfterFirstStep).isEqualTo((long) expectedDoubleProgressFromFirstStep + 3);
+
+        progressTracker.logProgress(1);
+        progressTracker.logSteps(4);
+        assertThat(leafTask.getProgress().progress()).isEqualTo(3 + 1 + (long) (100.0 * 5.0 / 13));
+    }
+
     private TaskProgressTracker progressTracker(Task task, Log log) {
         return new TaskProgressTracker(task, log, 1, EmptyTaskRegistryFactory.INSTANCE);
     }