Merge pull request #10890 from IoannisPanagiotas/mem-issues-my-fam

Remove memory from memoryTracker for incorrect project config

Author: vnickolov

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

@@ -42,14 +42,19 @@ public class MemoryUsageValidator {
     private final MemoryTracker memoryTracker;
     private final String username;
 
-    public MemoryUsageValidator(String username,MemoryTracker memoryTracker, boolean useMaxMemoryEstimation, Log log) {
+    public MemoryUsageValidator(
+        String username,
+        MemoryTracker memoryTracker,
+        boolean useMaxMemoryEstimation,
+        Log log
+    ) {
         this.log = log;
         this.useMaxMemoryEstimation = useMaxMemoryEstimation;
         this.memoryTracker = memoryTracker;
         this.username = username;
     }
 
-    public <C extends BaseConfig & JobIdConfig> MemoryRange tryValidateMemoryUsage(
+    public synchronized  <C extends BaseConfig & JobIdConfig> MemoryRange tryValidateMemoryUsage(
         String taskName,
         C config,
         Function<C, MemoryTreeWithDimensions> runEstimation

native-projection/src/main/java/org/neo4j/gds/projection/NativeFactory.java

@@ -41,8 +41,6 @@
 
 import java.util.Optional;
 
-import static org.neo4j.gds.projection.GraphDimensionsValidation.validate;
-
 final class NativeFactory extends CSRGraphStoreFactory<GraphProjectFromStoreConfig> {
 
     private final GraphProjectFromStoreConfig storeConfig;
@@ -64,18 +62,21 @@ static NativeFactory nativeFactory(
         return new NativeFactory(
             graphProjectFromStoreConfig,
             loadingContext,
-            dimensions
+            dimensions,
+            initProgressTracker(graphProjectFromStoreConfig, loadingContext, dimensions)
         );
     }
 
-    private NativeFactory(
+    // Package-private constructor used in tests
+    NativeFactory(
         GraphProjectFromStoreConfig graphProjectConfig,
         GraphLoaderContext loadingContext,
-        GraphDimensions graphDimensions
+        GraphDimensions graphDimensions,
+        ProgressTracker progressTracker
     ) {
         super(graphProjectConfig, ImmutableStaticCapabilities.of(WriteMode.LOCAL), loadingContext, graphDimensions);
         this.storeConfig = graphProjectConfig;
-        this.progressTracker = initProgressTracker();
+        this.progressTracker = progressTracker;
     }
 
     @Override
@@ -88,7 +89,10 @@ public MemoryEstimation estimateMemoryUsageAfterLoading() {
         return getMemoryEstimation(storeConfig.nodeProjections(), storeConfig.relationshipProjections(), false);
     }
 
-    private ProgressTracker initProgressTracker() {
+    private static ProgressTracker initProgressTracker(
+        GraphProjectFromStoreConfig graphProjectConfig,
+        GraphLoaderContext loadingContext, GraphDimensions dimensions
+    ) {
         long relationshipCount = graphProjectConfig
             .relationshipProjections()
             .projections()
@@ -133,11 +137,12 @@ private ProgressTracker initProgressTracker() {
 
     @Override
     public CSRGraphStore build() {
-        validate(dimensions, storeConfig);
-
-        var concurrency = graphProjectConfig.readConcurrency();
         try {
+            // Start the sub-task so it will be registered in the task store
             progressTracker.beginSubTask();
+            // `validate` can raise an exception which has to be handled in order to end the sub-task with failure
+            validate();
+            var concurrency = graphProjectConfig.readConcurrency();
             Nodes nodes = loadNodes(concurrency);
             RelationshipImportResult relationships = loadRelationships(nodes.idMap(), concurrency);
             CSRGraphStore graphStore = createGraphStore(nodes, relationships);
@@ -152,6 +157,11 @@ public CSRGraphStore build() {
         }
     }
 
+    // Allows for mocking the validation behaviour.
+    void validate() throws IllegalArgumentException {
+        GraphDimensionsValidation.validate(dimensions, storeConfig);
+    }
+
     private Nodes loadNodes(Concurrency concurrency) {
         var scanningNodesImporter = new ScanningNodesImporterBuilder()
             .concurrency(concurrency)

native-projection/src/test/java/org/neo4j/gds/projection/NativeFactoryTest.java

@@ -27,15 +27,24 @@
 import org.neo4j.gds.RelationshipProjections;
 import org.neo4j.gds.RelationshipType;
 import org.neo4j.gds.api.CSRGraphStoreFactory;
+import org.neo4j.gds.api.GraphLoaderContext;
 import org.neo4j.gds.core.GraphDimensions;
 import org.neo4j.gds.core.ImmutableGraphDimensions;
 import org.neo4j.gds.core.concurrency.Concurrency;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.mem.MemoryEstimation;
 import org.neo4j.gds.mem.MemoryTree;
 
 import java.util.concurrent.atomic.AtomicReference;
 
+import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 class NativeFactoryTest {
 
@@ -142,4 +151,26 @@ void memoryEstimationForMultipleProjections() {
         assertEquals(12_056_534_400L, estimate.memoryUsage().min);
         assertEquals(13_667_147_136L, estimate.memoryUsage().max);
     }
+
+    @Test
+    void shouldCleanTaskOnValidateFailure() {
+        var progressTrackerMock = mock(ProgressTracker.class);
+
+        var nativeFactorySpy = spy(new NativeFactory(
+            mock(GraphProjectFromStoreConfig.class),
+            mock(GraphLoaderContext.class),
+            mock(GraphDimensions.class),
+            progressTrackerMock
+        ));
+
+        doThrow(new IllegalArgumentException("Intentionally failing validation")).when(nativeFactorySpy).validate();
+
+        assertThatIllegalArgumentException()
+            .isThrownBy(nativeFactorySpy::build)
+            .withMessageContaining("Intentionally failing validation");
+
+        verify(progressTrackerMock, times(1)).beginSubTask();
+        verify(progressTrackerMock, times(1)).endSubTaskWithFailure();
+        verifyNoMoreInteractions(progressTrackerMock);
+    }
 }