Introduce graphStore container in mem. tracker

Author: IoannisPanagiotas

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

@@ -0,0 +1,66 @@
+/*
+ * 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 org.neo4j.gds.api.graph.store.catalog.GraphStoreAddedEvent;
+import org.neo4j.gds.api.graph.store.catalog.GraphStoreRemovedEvent;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.Stream;
+
+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<>();
+
+    long addGraph(GraphStoreAddedEvent graphStoreAddedEvent){
+        var addedGraphMemory = graphStoreAddedEvent.memoryInBytes();
+        var graphsMemory = graphStoreReservedMemory.addAndGet(addedGraphMemory);
+        graphStoresMemory.putIfAbsent(graphStoreAddedEvent.user(), new ConcurrentHashMap<>());
+        graphStoresMemory.get(graphStoreAddedEvent.user()).put(graphStoreAddedEvent.graphName(),graphStoreAddedEvent.memoryInBytes());
+        return  graphsMemory;
+    }
+
+    long removeGraph(GraphStoreRemovedEvent graphStoreRemovedEvent){
+        var graphMemoryToRemove = graphStoreRemovedEvent.memoryInBytes();
+        var graphsMemoryAfterRemoval = graphStoreReservedMemory.addAndGet(-graphMemoryToRemove);
+        graphStoresMemory.get(graphStoreRemovedEvent.user()).remove(graphStoreRemovedEvent.graphName());
+        return  graphsMemoryAfterRemoval;
+    }
+
+    long graphStoreReservedMemory(){
+        return  graphStoreReservedMemory.get();
+    }
+
+    Stream<UserEntityMemory> listGraphs(String user){
+        return  graphStoresMemory
+            .getOrDefault(user, EMPTY_HASH_MAP)
+            .entrySet()
+            .stream()
+            .map( entry -> new UserEntityMemory(user, entry.getKey(), entry.getValue()));
+    }
+
+    Stream<UserEntityMemory> listGraphs(){
+       return   graphStoresMemory.keySet().stream().flatMap(this::listGraphs);
+    }
+
+}

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

@@ -31,18 +31,15 @@
 import org.neo4j.gds.core.utils.progress.UserTask;
 import org.neo4j.gds.logging.Log;
 
-import java.util.concurrent.atomic.AtomicLong;
 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 AtomicLong graphStoresMemory = new AtomicLong();
-
-    private final ObjectLongMap<JobId> memoryInUse = new ObjectLongHashMap<>();
-    private final Log log;
+    private final  GraphStoreMemoryContainer  graphStoreMemoryContainer = new GraphStoreMemoryContainer();
+    private final  ObjectLongMap<JobId> memoryInUse = new ObjectLongHashMap<>();
+    private final  Log log;
 
     public MemoryTracker(long initialMemory, Log log) {
         this.log = log;
@@ -71,7 +68,7 @@ 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() + graphStoresMemory.longValue());
+        return initialMemory - (reservedMemory.longValue() + graphStoreMemoryContainer.graphStoreReservedMemory());
     }
 
     @Override
@@ -100,24 +97,23 @@ private static void assertPositiveInitialMemory(long initialMemory) {
 
     @Override
     public void onGraphStoreAdded(GraphStoreAddedEvent graphStoreAddedEvent) {
-        var addedGraphMemory = graphStoreAddedEvent.memoryInBytes();
-        var graphsMemory = graphStoresMemory.addAndGet(addedGraphMemory);
+        var graphsMemory = graphStoreMemoryContainer.addGraph(graphStoreAddedEvent);
         log.debug(
             "Added graph %s, which added another %s bytes, now there are %s bytes occupied by projected graphs",
             graphStoreAddedEvent.graphName(),
-            addedGraphMemory,
+            graphStoreAddedEvent.memoryInBytes(),
             graphsMemory
         );
     }
 
     @Override
-    public void onGraphStoreRemoved(GraphStoreRemovedEvent graphStoreAddedEvent) {
-        var graphMemoryToRemove = graphStoreAddedEvent.memoryInBytes();
-        var graphsMemoryAfterRemoval = graphStoresMemory.addAndGet(-graphMemoryToRemove);
+    public void onGraphStoreRemoved(GraphStoreRemovedEvent graphStoreRemovedEvent) {
+
+       var graphsMemoryAfterRemoval= graphStoreMemoryContainer.removeGraph(graphStoreRemovedEvent);
         log.debug(
             "Removed graph %s, which freed %s bytes, there are still %s bytes occupied by projected graphs",
-            graphStoreAddedEvent.graphName(),
-            graphMemoryToRemove,
+            graphStoreRemovedEvent.graphName(),
+            graphStoreRemovedEvent.memoryInBytes(),
             graphsMemoryAfterRemoval
         );
     }

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

@@ -0,0 +1,23 @@
+/*
+ * 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;
+
+public record UserEntityMemory(String user, String graph, long memoryInBytes) {
+}

progress-tracking/src/test/java/org/neo4j/gds/mem/GraphStoreMemoryContainerTest.java

@@ -0,0 +1,75 @@
+/*
+ * 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 org.junit.jupiter.api.Test;
+import org.neo4j.gds.api.graph.store.catalog.GraphStoreAddedEvent;
+import org.neo4j.gds.api.graph.store.catalog.GraphStoreRemovedEvent;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class GraphStoreMemoryContainerTest {
+
+    @Test
+    void shouldAddGraphs(){
+         GraphStoreMemoryContainer graphStoreMemoryContainer=new GraphStoreMemoryContainer();
+         assertThat(graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("foo","DB","buzz",10))).isEqualTo(10);
+         assertThat(graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("foo","DB","buzz2",20))).isEqualTo(30);
+         assertThat(graphStoreMemoryContainer.graphStoreReservedMemory()).isEqualTo(30L);
+
+    }
+
+    @Test
+    void shouldRemoveGraphs(){
+        GraphStoreMemoryContainer graphStoreMemoryContainer=new GraphStoreMemoryContainer();
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("foo","DB","buzz",10));
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("foo","DB","buzz2",20));
+        assertThat(graphStoreMemoryContainer.removeGraph(new GraphStoreRemovedEvent("foo","DB","buzz2",20))).isEqualTo(10);
+    }
+
+    @Test
+    void shouldListForUser(){
+        GraphStoreMemoryContainer graphStoreMemoryContainer=new GraphStoreMemoryContainer();
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("Alice","DB","graph1",10));
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("Alice","DB","graph2",15));
+
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("Bob","DB","graph3",20));
+        var aliceList =graphStoreMemoryContainer.listGraphs("Alice").toList();
+        assertThat(aliceList).hasSize(2);
+        assertThat(aliceList.stream().map(UserEntityMemory::graph).toList()).containsExactlyInAnyOrder("graph1","graph2");
+        assertThat(aliceList.stream().map(UserEntityMemory::memoryInBytes).toList()).containsExactlyInAnyOrder(10L,15L);
+
+    }
+
+    @Test
+    void shouldListAll(){
+        GraphStoreMemoryContainer graphStoreMemoryContainer=new GraphStoreMemoryContainer();
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("Alice","DB","graph1",10));
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("Alice","DB","graph2",15));
+
+        graphStoreMemoryContainer.addGraph(new GraphStoreAddedEvent("Bob","DB","graph3",20));
+        var graphList =graphStoreMemoryContainer.listGraphs().toList();
+        assertThat(graphList).hasSize(3);
+        assertThat(graphList.stream().map(UserEntityMemory::graph).toList()).containsExactlyInAnyOrder("graph1","graph2","graph3");
+        assertThat(graphList.stream().map(UserEntityMemory::memoryInBytes).toList()).containsExactlyInAnyOrder(10L,15L,20L);
+
+    }
+
+}