Tasks - keep tasks after completion (#10044)

* Introduce task retention

Keep completed tasks around for X time (setting based).
Maintain previous behaviour for gds.listProgress and Aura Shutdown/Backup

* Fix tests

* Rename task ttl to retentionPeriod

Co-authored-by: Sören Reichardt <soren.reichardt@neo4j.com>

* Fix test on rebase

* Reuse logic to wait for taskStore to get cleaned up

We can be a bit relaxed on this side-effect check. As long as it happens its fine. Avoiding race conditions for instance on cancel()

* Fix resultStore and exportJobId order

* Remove accidentially commited test

* Format code

* Document setting also for prod deployment

* Define export and resultStore jobId for v2 endpoints

even graphstore exports go through the resultStore

---------

Co-authored-by: Sören Reichardt <soren.reichardt@neo4j.com>

Author: FlorentinD

algo-common/src/main/java/org/neo4j/gds/result/CommunityStatistics.java

@@ -188,7 +188,7 @@ public static CommunityStats communityStats(
                     concurrency
                 );
             }
-        } catch (Exception e){
+        } catch (Exception e) {
             if (e.getMessage().contains("is out of bounds for histogram, current covered range")) {
                 return  new CommunityStats(0,Optional.empty(), computeMilliseconds.get(), false);
             } else {

algo/src/test/java/org/neo4j/gds/traversal/RandomWalkTest.java

@@ -25,6 +25,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.neo4j.gds.TaskStoreHelper;
 import org.neo4j.gds.TestProgressTracker;
 import org.neo4j.gds.TestSupport;
 import org.neo4j.gds.applications.algorithms.machinery.RequestScopedDependencies;
@@ -48,15 +49,12 @@
 import org.neo4j.gds.logging.GdsTestLog;
 import org.neo4j.gds.termination.TerminationFlag;
 
-import java.time.Instant;
-import java.time.temporal.ChronoUnit;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.TreeSet;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.locks.LockSupport;
 import java.util.stream.Collectors;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -516,7 +514,7 @@ void progressLogging() throws InterruptedException {
                 .build();
 
             var log = new GdsTestLog();
-            var taskStore = new PerDatabaseTaskStore();
+            var taskStore = new PerDatabaseTaskStore(Duration.ZERO);
 
             var requestScopedDependencies = RequestScopedDependencies.builder()
                 .taskRegistryFactory(TaskRegistryFactory.local("rw", taskStore))
@@ -538,7 +536,7 @@ void progressLogging() throws InterruptedException {
                 assertThat(randomWalksStream).hasSize(5000);
             });
 
-            awaitEmptyTaskStore(taskStore);
+            TaskStoreHelper.awaitEmptyTaskStore(taskStore);
 
             assertThat(log.getMessages(TestLog.INFO))
                 .extracting(removingThreadId())
@@ -571,7 +569,7 @@ void shouldLogProgressOnWeightedGraph() {
                 .build();
 
             var log = new GdsTestLog();
-            var taskStore = new PerDatabaseTaskStore();
+            var taskStore = new PerDatabaseTaskStore(Duration.ZERO);
 
             var requestScopedDependencies = RequestScopedDependencies.builder()
                 .taskRegistryFactory(TaskRegistryFactory.local("rw", taskStore))
@@ -592,7 +590,7 @@ void shouldLogProgressOnWeightedGraph() {
                 assertThat(randomWalksStream).hasSize(5000);
             });
 
-            awaitEmptyTaskStore(taskStore);
+            TaskStoreHelper.awaitEmptyTaskStore(taskStore);
 
             assertThat(log.getMessages(TestLog.INFO))
                 .extracting(removingThreadId())
@@ -615,9 +613,9 @@ void shouldLogProgressOnWeightedGraph() {
         }
 
         @Test
-        void shouldLeaveNoTasksBehind() {
+        void shouldLeaveNoOngoingTasksBehind() {
             var config = RandomWalkStreamConfigImpl.builder().build();
-            var taskStore = new PerDatabaseTaskStore();
+            var taskStore = new PerDatabaseTaskStore(Duration.ZERO);
 
             var requestScopedDependencies = RequestScopedDependencies.builder()
                 .taskRegistryFactory(TaskRegistryFactory.local("rw", taskStore))
@@ -638,26 +636,11 @@ void shouldLeaveNoTasksBehind() {
             //noinspection ResultOfMethodCallIgnored
             result.count();
 
-            awaitEmptyTaskStore(taskStore);
+            TaskStoreHelper.awaitEmptyTaskStore(taskStore);
 
             // the task store should now be empty
-            assertThat(taskStore.isEmpty()).isTrue();
+            assertThat(taskStore.queryRunning()).isEmpty();
         }
 
-        private void awaitEmptyTaskStore(PerDatabaseTaskStore taskStore) {
-            // there is a race condition between the thread consuming the result,
-            // and the thread scheduled to end the last subtask
-            long timeoutInSeconds = 5 * (TestSupport.CI ? 5 : 1);
-            var deadline = Instant.now().plus(timeoutInSeconds, ChronoUnit.SECONDS);
-
-            // On my machine (TM) with 1000 iterations this never fails and each run is 10ish or 100ish ms
-            while (Instant.now().isBefore(deadline)) {
-                if (taskStore.isEmpty()) {
-                    break;
-                }
-
-                LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(100));
-            }
-        }
     }
 }

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/Neo4jDatabaseNodePropertyWriter.java

@@ -168,6 +168,9 @@ static NodePropertiesWritten writeNodeProperties(
         try {
             exporter.write(nodeProperties);
             propertiesWritten.setValue(exporter.propertiesWritten());
+        } catch (Exception e) {
+            progressTracker.endSubTaskWithFailure();
+            throw e;
         } finally {
             progressTracker.release();
         }

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/NodePropertyWriter.java

@@ -96,7 +96,10 @@ public NodePropertiesWritten writeNodeProperties(
 
         try {
             return writeNodeProperties(nodePropertyExporter, nodeProperties);
-        } finally {
+        } catch (Exception e) {
+            progressTracker.endSubTaskWithFailure();
+            throw e;
+        }finally {
             progressTracker.release();
         }
     }

applications/algorithms/node-embeddings/src/test/java/org/neo4j/gds/applications/algorithms/embeddings/NodeEmbeddingAlgorithmsGraphSageTest.java

@@ -34,6 +34,7 @@
 import org.neo4j.gds.logging.GdsTestLog;
 import org.neo4j.gds.termination.TerminationFlag;
 
+import java.time.Duration;
 import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -76,7 +77,7 @@ void testLogProgressForGraphSageTrain() {
             GraphSageTrainTask.create(graph, config),
             config.username(),
             config.jobId(),
-            new PerDatabaseTaskStore(),
+            new PerDatabaseTaskStore(Duration.ofMinutes(1)),
             new LoggerForProgressTrackingAdapter(log)
         );
 

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/WriteNodePropertiesApplication.java

@@ -146,11 +146,12 @@ private static long writeNodeProperties(
                 .collect(Collectors.toList());
 
             exporter.write(writeNodeProperties);
+            progressTracker.endSubTask();
 
             return exporter.propertiesWritten();
-
-        } finally {
-            progressTracker.endSubTask();
+        } catch (Exception e) {
+            progressTracker.endSubTaskWithFailure();
+            throw e;
         }
     }
 }

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/WriteRelationshipsApplication.java

@@ -94,6 +94,7 @@ WriteRelationshipResult compute(
                 // result
                 return builder.build();
             } catch (RuntimeException e) {
+                progressTracker.endSubTaskWithFailure();
                 loggers.log().warn("Writing relationships failed", e);
                 throw e;
             }

applications/operations/src/main/java/org/neo4j/gds/applications/operations/OperationsApplications.java

@@ -64,6 +64,8 @@ public Stream<UserTask> listProgress() {
 
         if (user.isAdmin()) return taskStore.query();
 
+        // TODO only fetch ongoing tasks here and name listProgressOngoing
+        //      (add listProgress which fetches all tasks)
         return taskStore.query(user.getUsername());
     }
 

core-write/src/main/java/org/neo4j/gds/core/write/NativeNodeLabelExporter.java

@@ -86,8 +86,10 @@ public void write(String nodeLabel) {
             } else {
                 writeSequential(nodeLabelToken);
             }
-        } finally {
             progressTracker.endSubTask();
+        } catch (Exception e) {
+           progressTracker.endSubTaskWithFailure();
+           throw e;
         }
     }
 

core-write/src/main/java/org/neo4j/gds/core/write/NativeNodePropertyExporter.java

@@ -122,8 +122,10 @@ public void write(Collection<NodeProperty> nodeProperties) {
             } else {
                 writeSequential(resolvedNodeProperties);
             }
-        } finally {
             progressTracker.endSubTask();
+        } catch (Exception e) {
+            progressTracker.endSubTaskWithFailure();
+            throw e;
         }
     }