Minor wrapper around memory

Co-authored-by: Veselin Nikolov <[email protected]>

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

@@ -66,20 +66,11 @@ public <CONFIGURATION extends AlgoBaseConfig> void assertAlgorithmCanRun(
         if (configuration.sudo()) return;
 
         try {
-            var memoryEstimation = estimationFactory.get();
 
-            var graphDimensions = graphDimensionFactory.create(graphStore, configuration);
-
-            var transformedGraphDimensions = dimensionTransformer.transform(graphDimensions);
-
-            var memoryTree = memoryEstimation.estimate(transformedGraphDimensions, configuration.concurrency());
-
-            var memoryRange = memoryTree.memoryUsage();
-
-            var bytesRequired = useMaxMemoryEstimation ? memoryRange.max : memoryRange.min;
+           var memoryRequirement = MemoryRequirement.create(estimationFactory,graphStore,graphDimensionFactory,dimensionTransformer,configuration,useMaxMemoryEstimation);
 
             try {
-                memoryGauge.tryToReserveMemory(bytesRequired);
+                memoryGauge.tryToReserveMemory(memoryRequirement.requiredMemory());
             } catch (MemoryReservationExceededException e) {
                 var message = StringFormatting.formatWithLocale(
                     "Memory required to run %s (%db) exceeds available memory (%db)",

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

@@ -0,0 +1,53 @@
+/*
+ * 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.applications.algorithms.machinery;
+
+import org.neo4j.gds.api.GraphStore;
+import org.neo4j.gds.applications.services.GraphDimensionFactory;
+import org.neo4j.gds.config.AlgoBaseConfig;
+import org.neo4j.gds.mem.MemoryEstimation;
+
+import java.util.function.Supplier;
+
+public record MemoryRequirement(long requiredMemory) {
+
+    static  <CONFIGURATION extends AlgoBaseConfig> MemoryRequirement create(
+        Supplier<MemoryEstimation> estimationFactory,
+        GraphStore graphStore,
+        GraphDimensionFactory graphDimensionFactory,
+        DimensionTransformer dimensionTransformer,
+        CONFIGURATION configuration,
+        boolean useMaxMemoryEstimation
+    ) {
+        var memoryEstimation = estimationFactory.get();
+
+        var graphDimensions = graphDimensionFactory.create(graphStore, configuration);
+
+        var transformedGraphDimensions = dimensionTransformer.transform(graphDimensions);
+
+        var memoryTree = memoryEstimation.estimate(transformedGraphDimensions, configuration.concurrency());
+
+        var memoryRange = memoryTree.memoryUsage();
+
+        var bytesRequired = useMaxMemoryEstimation ? memoryRange.max : memoryRange.min;
+
+        return new MemoryRequirement(bytesRequired);
+    }
+}