Introduce task container in mem. tracker

Author: IoannisPanagiotas

progress-tracking/src/main/java/org/neo4j/gds/mem/GraphStoreMemoryContainer.java

@@ -30,7 +30,7 @@ class GraphStoreMemoryContainer {
 
     private final ConcurrentHashMap<String, ConcurrentHashMap<String,Long>> graphStoresMemory = new ConcurrentHashMap<>();
     private  final AtomicLong graphStoreReservedMemory = new AtomicLong();
-    private static final ConcurrentHashMap<String,Long> EMPTY_HASH_MAP   =new ConcurrentHashMap<>();
+    private static final ConcurrentHashMap<String,Long> EMPTY_HASH_MAP   = new ConcurrentHashMap<>();
 
     long addGraph(GraphStoreAddedEvent graphStoreAddedEvent){
         var addedGraphMemory = graphStoreAddedEvent.memoryInBytes();

progress-tracking/src/main/java/org/neo4j/gds/mem/MemoryTracker.java

@@ -19,9 +19,6 @@
  */
 package org.neo4j.gds.mem;
 
-import com.carrotsearch.hppc.ObjectLongHashMap;
-import com.carrotsearch.hppc.ObjectLongMap;
-import com.carrotsearch.hppc.procedures.LongProcedure;
 import org.neo4j.gds.api.graph.store.catalog.GraphStoreAddedEvent;
 import org.neo4j.gds.api.graph.store.catalog.GraphStoreAddedEventListener;
 import org.neo4j.gds.api.graph.store.catalog.GraphStoreRemovedEvent;
@@ -31,14 +28,12 @@
 import org.neo4j.gds.core.utils.progress.UserTask;
 import org.neo4j.gds.logging.Log;
 
-import java.util.concurrent.atomic.LongAdder;
-
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
 
 public class MemoryTracker implements TaskStoreListener, GraphStoreAddedEventListener, GraphStoreRemovedEventListener {
     private final long initialMemory;
     private final  GraphStoreMemoryContainer  graphStoreMemoryContainer = new GraphStoreMemoryContainer();
-    private final  ObjectLongMap<JobId> memoryInUse = new ObjectLongHashMap<>();
+    private final  TaskMemoryContainer taskMemoryContainer = new TaskMemoryContainer();
     private final  Log log;
 
     public MemoryTracker(long initialMemory, Log log) {
@@ -53,7 +48,7 @@ public long initialMemory() {
 
     public synchronized void track(JobId jobId, long memoryEstimate) {
         log.debug("Tracking %s:  %s bytes", jobId.asString(), memoryEstimate);
-        memoryInUse.put(jobId, memoryEstimate);
+        taskMemoryContainer.reserve(jobId, memoryEstimate);
         log.debug("Available memory after tracking task: %s bytes", availableMemory());
     }
 
@@ -66,22 +61,21 @@ public synchronized void tryToTrack(JobId jobId, long memoryEstimate) throws Mem
     }
 
     public synchronized long availableMemory() {
-        var reservedMemory = new LongAdder();
-        memoryInUse.values().forEach((LongProcedure) reservedMemory::add);
-        return initialMemory - (reservedMemory.longValue() + graphStoreMemoryContainer.graphStoreReservedMemory());
+        return initialMemory - graphStoreMemoryContainer.graphStoreReservedMemory() - taskMemoryContainer.taskReservedMemory();
     }
 
     @Override
     public void onTaskAdded(UserTask userTask) {
         // do nothing, we add the memory explicitly prior to execution
+        taskMemoryContainer.addTask(userTask);
     }
 
     @Override
     public synchronized void onTaskRemoved(UserTask userTask) {
         var taskDescription = userTask.task().description();
         log.debug("Removing task: %s", taskDescription);
         var jobId = userTask.jobId();
-        var removed = memoryInUse.remove(jobId);
+        var removed= taskMemoryContainer.removeTask(userTask);
         log.debug("Removed task %s (%s):  %s bytes", taskDescription, jobId.asString(), removed);
         log.debug("Available memory after removing task: %s bytes", availableMemory());
         log.debug("Done removing task: %s", taskDescription);

progress-tracking/src/main/java/org/neo4j/gds/mem/TaskMemoryContainer.java

@@ -0,0 +1,71 @@
+/*
+ * 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.mem;
+
+import com.carrotsearch.hppc.ObjectLongHashMap;
+import com.carrotsearch.hppc.ObjectLongMap;
+import org.neo4j.gds.core.utils.progress.JobId;
+import org.neo4j.gds.core.utils.progress.UserTask;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.Stream;
+
+class TaskMemoryContainer {
+
+    private final ConcurrentHashMap<String,ConcurrentHashMap<UserTask,Long>> memoryInUse = new ConcurrentHashMap<>();
+    private final AtomicLong allocatedMemory = new AtomicLong();
+    private final ObjectLongMap<JobId>  temporaryMap =new ObjectLongHashMap<>();
+    private static final ConcurrentHashMap<UserTask,Long> EMPTY_HASH_MAP = new ConcurrentHashMap<>();
+
+    void reserve(JobId jobId,long memoryAmount){
+        temporaryMap.put(jobId,memoryAmount);
+        allocatedMemory.addAndGet(memoryAmount);
+    }
+
+    void addTask(UserTask task){
+        var memoryAmount = temporaryMap.remove(task.jobId());
+        memoryInUse.putIfAbsent(task.username(), new ConcurrentHashMap<>());
+        memoryInUse.get(task.username()).put(task,memoryAmount);
+    }
+
+    long removeTask(UserTask task){
+         var mem=  memoryInUse.getOrDefault(task.username(), EMPTY_HASH_MAP).remove(task);
+         allocatedMemory.addAndGet(-mem);
+         return  mem;
+    }
+
+    long taskReservedMemory(){
+        return  allocatedMemory.get();
+    }
+
+    Stream<UserEntityMemory> listTasks(String user){
+        return  memoryInUse
+            .getOrDefault(user, EMPTY_HASH_MAP)
+            .entrySet()
+            .stream()
+            .map( entry -> new UserEntityMemory(user, entry.getKey().task().description(), entry.getValue()));
+    }
+
+    Stream<UserEntityMemory> listTasks(){
+        return  memoryInUse.keySet().stream().flatMap(this::listTasks);
+    }
+
+}