Force all native allocations and deallocation to track

amount of requested and freed bytes
neo4j · Feb 1, 2018 · c1049c8 · c1049c8
1 parent 35938da
commit c1049c8
Show file tree

Hide file tree

Showing 23 changed files with 192 additions and 119 deletions.
diff --git a/community/io/src/main/java/org/neo4j/io/mem/GrabAllocator.java b/community/io/src/main/java/org/neo4j/io/mem/GrabAllocator.java
@@ -19,6 +19,7 @@
  */
 package org.neo4j.io.mem;
 
+import org.neo4j.memory.MemoryAllocationTracker;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
 import static org.neo4j.io.ByteUnit.kibiBytes;
@@ -39,6 +40,7 @@ public final class GrabAllocator implements MemoryAllocator
      * The amount of memory that this memory manager can still allocate.
      */
     private long memoryReserve;
+    private final MemoryAllocationTracker memoryTracker;
 
     private Grab grabs;
 
@@ -47,10 +49,12 @@ public final class GrabAllocator implements MemoryAllocator
      * given alignment size.
      * @param expectedMaxMemory The maximum amount of memory that this memory manager is expected to allocate. The
      * actual amount of memory used can end up greater than this value, if some of it gets wasted on alignment padding.
+     * @param memoryTracker memory usage tracker
      */
-    GrabAllocator( long expectedMaxMemory )
+    GrabAllocator( long expectedMaxMemory, MemoryAllocationTracker memoryTracker )
     {
         this.memoryReserve = expectedMaxMemory;
+        this.memoryTracker = memoryTracker;
     }
 
     @Override
@@ -93,12 +97,12 @@ public synchronized long allocateAligned( long bytes, long alignment )
                 // This is a huge allocation. Put it in its own grab and keep any existing grab at the head.
                 grabSize = bytes;
                 Grab nextGrab = grabs == null ? null : grabs.next;
-                Grab allocationGrab = new Grab( nextGrab, grabSize );
+                Grab allocationGrab = new Grab( nextGrab, grabSize, memoryTracker );
                 if ( !allocationGrab.canAllocate( bytes ) )
                 {
                     allocationGrab.free();
                     grabSize = bytes + alignment;
-                    allocationGrab = new Grab( nextGrab, grabSize );
+                    allocationGrab = new Grab( nextGrab, grabSize, memoryTracker );
                 }
                 long allocation = allocationGrab.allocate( bytes, alignment );
                 grabs = grabs == null ? allocationGrab : grabs.setNext( allocationGrab );
@@ -111,7 +115,7 @@ public synchronized long allocateAligned( long bytes, long alignment )
                 if ( grabSize < bytes )
                 {
                     grabSize = bytes;
-                    Grab grab = new Grab( grabs, grabSize );
+                    Grab grab = new Grab( grabs, grabSize, memoryTracker );
                     if ( grab.canAllocate( bytes ) )
                     {
                         memoryReserve -= grabSize;
@@ -121,7 +125,7 @@ public synchronized long allocateAligned( long bytes, long alignment )
                     grab.free();
                     grabSize = bytes + alignment;
                 }
-                grabs = new Grab( grabs, grabSize );
+                grabs = new Grab( grabs, grabSize, memoryTracker );
                 memoryReserve -= grabSize;
             }
             return grabs.allocate( bytes, alignment );
@@ -162,8 +166,13 @@ private void initCause( NativeMemoryAllocationRefusedError error, OutOfMemoryErr
     protected synchronized void finalize() throws Throwable
     {
         super.finalize();
-        Grab current = grabs;
+        close();
+    }
 
+    @Override
+    public void close()
+    {
+        Grab current = grabs;
         while ( current != null )
         {
             current.free();
@@ -176,22 +185,25 @@ private static class Grab
         public final Grab next;
         private final long address;
         private final long limit;
+        private final MemoryAllocationTracker memoryTracker;
         private long nextPointer;
 
-        Grab( Grab next, long size )
+        Grab( Grab next, long size, MemoryAllocationTracker memoryTracker )
         {
             this.next = next;
-            this.address = UnsafeUtil.allocateMemory( size );
+            this.address = UnsafeUtil.allocateMemory( size, memoryTracker );
             this.limit = address + size;
+            this.memoryTracker = memoryTracker;
             nextPointer = address;
         }
 
-        Grab( Grab next, long address, long limit, long nextPointer )
+        Grab( Grab next, long address, long limit, long nextPointer, MemoryAllocationTracker memoryTracker )
         {
             this.next = next;
             this.address = address;
             this.limit = limit;
             this.nextPointer = nextPointer;
+            this.memoryTracker = memoryTracker;
         }
 
         private long nextAligned( long pointer, long alignment )
@@ -213,7 +225,7 @@ long allocate( long bytes, long alignment )
 
         void free()
         {
-            UnsafeUtil.free( address );
+            UnsafeUtil.free( address, limit - address, memoryTracker );
         }
 
         boolean canAllocate( long bytes )
@@ -223,7 +235,7 @@ boolean canAllocate( long bytes )
 
         Grab setNext( Grab grab )
         {
-            return new Grab( grab, address, limit, nextPointer );
+            return new Grab( grab, address, limit, nextPointer, memoryTracker );
         }
 
         @Override

diff --git a/community/io/src/main/java/org/neo4j/io/mem/MemoryAllocator.java b/community/io/src/main/java/org/neo4j/io/mem/MemoryAllocator.java
@@ -20,15 +20,16 @@
 package org.neo4j.io.mem;
 
 import org.neo4j.io.ByteUnit;
+import org.neo4j.memory.MemoryAllocationTracker;
 
 /**
  * A MemoryAllocator is simple: it only allocates memory, until it itself is finalizable and frees it all in one go.
  */
 public interface MemoryAllocator
 {
-    static MemoryAllocator createAllocator( String expectedMemory )
+    static MemoryAllocator createAllocator( String expectedMemory, MemoryAllocationTracker memoryTracker )
     {
-        return new GrabAllocator( ByteUnit.parse( expectedMemory ) );
+        return new GrabAllocator( ByteUnit.parse( expectedMemory ), memoryTracker );
     }
 
     /**
@@ -49,4 +50,9 @@ static MemoryAllocator createAllocator( String expectedMemory )
      * @throws OutOfMemoryError if the requested memory could not be allocated.
      */
     long allocateAligned( long bytes, long alignment );
+
+    /**
+     * Close allocator and release all allocated resources
+     */
+    void close();
 }
diff --git a/community/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/MuninnPageCache.java b/community/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/MuninnPageCache.java
@@ -47,6 +47,8 @@
 import org.neo4j.io.pagecache.tracing.PageCacheTracer;
 import org.neo4j.io.pagecache.tracing.PageFaultEvent;
 import org.neo4j.io.pagecache.tracing.cursor.PageCursorTracerSupplier;
+import org.neo4j.memory.GlobalMemoryTracker;
+import org.neo4j.memory.MemoryAllocationTracker;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
 import static org.neo4j.util.FeatureToggles.flag;
@@ -219,7 +221,7 @@ public MuninnPageCache(
     {
         this( swapperFactory,
                 // Cast to long prevents overflow:
-                MemoryAllocator.createAllocator( "" + memoryRequiredForPages( maxPages ) ),
+                MemoryAllocator.createAllocator( "" + memoryRequiredForPages( maxPages ), GlobalMemoryTracker.INSTANCE ),
                 PAGE_SIZE,
                 pageCacheTracer,
                 pageCursorTracerSupplier );
@@ -261,6 +263,7 @@ public MuninnPageCache(
 
         // Expose the total number of pages
         pageCacheTracer.maxPages( maxPages );
+        MemoryAllocationTracker memoryTracker = GlobalMemoryTracker.INSTANCE;
 
         this.pageCacheId = pageCacheIdCounter.incrementAndGet();
         this.swapperFactory = swapperFactory;
@@ -270,7 +273,7 @@ public MuninnPageCache(
         this.pageCursorTracerSupplier = pageCursorTracerSupplier;
         this.printExceptionsOnClose = true;
         long alignment = swapperFactory.getRequiredBufferAlignment();
-        this.victimPage = VictimPageReference.getVictimPage( cachePageSize );
+        this.victimPage = VictimPageReference.getVictimPage( cachePageSize, memoryTracker );
         this.pages = new PageList( maxPages, cachePageSize, memoryAllocator, new SwapperSet(), victimPage, alignment );
 
         setFreelistHead( new AtomicInteger() );

diff --git a/...unity/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/StandalonePageCacheFactory.java b/...unity/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/StandalonePageCacheFactory.java
@@ -26,6 +26,7 @@
 import org.neo4j.io.pagecache.impl.SingleFilePageSwapperFactory;
 import org.neo4j.io.pagecache.tracing.PageCacheTracer;
 import org.neo4j.io.pagecache.tracing.cursor.DefaultPageCursorTracerSupplier;
+import org.neo4j.memory.GlobalMemoryTracker;
 
 /*
  * This class is an helper to allow to construct properly a page cache in the few places we need it without all
@@ -47,7 +48,7 @@ public static PageCache createPageCache( FileSystemAbstraction fileSystem )
 
         PageCacheTracer cacheTracer = PageCacheTracer.NULL;
         DefaultPageCursorTracerSupplier cursorTracerSupplier = DefaultPageCursorTracerSupplier.INSTANCE;
-        MemoryAllocator memoryAllocator = MemoryAllocator.createAllocator( "8 MiB" );
+        MemoryAllocator memoryAllocator = MemoryAllocator.createAllocator( "8 MiB", GlobalMemoryTracker.INSTANCE );
         return new MuninnPageCache( factory, memoryAllocator, cacheTracer, cursorTracerSupplier );
     }
 }
diff --git a/community/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/VictimPageReference.java b/community/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/VictimPageReference.java
@@ -19,6 +19,7 @@
  */
 package org.neo4j.io.pagecache.impl.muninn;
 
+import org.neo4j.memory.MemoryAllocationTracker;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
 class VictimPageReference
@@ -31,14 +32,14 @@ private VictimPageReference()
         // All state is static
     }
 
-    static synchronized long getVictimPage( int pageSize )
+    static synchronized long getVictimPage( int pageSize, MemoryAllocationTracker allocationTracker )
     {
         if ( victimPageSize < pageSize )
         {
             // Note that we NEVER free any old victim pages. This is important because we cannot tell
             // when we are done using them. Therefor, victim pages are allocated and stay allocated
             // until our process terminates.
-            victimPagePointer = UnsafeUtil.allocateMemory( pageSize );
+            victimPagePointer = UnsafeUtil.allocateMemory( pageSize, allocationTracker );
             victimPageSize = pageSize;
         }
         return victimPagePointer;

diff --git a/community/io/src/test/java/org/neo4j/io/mem/MemoryAllocatorTest.java b/community/io/src/test/java/org/neo4j/io/mem/MemoryAllocatorTest.java
@@ -21,13 +21,16 @@
 
 import org.junit.Test;
 
+import org.neo4j.io.ByteUnit;
 import org.neo4j.io.pagecache.PageCache;
+import org.neo4j.memory.LocalMemoryTracker;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.Matchers.not;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 
 public class MemoryAllocatorTest
@@ -37,7 +40,7 @@ public class MemoryAllocatorTest
 
     protected MemoryAllocator createAllocator( String expectedMaxMemory )
     {
-        return MemoryAllocator.createAllocator( expectedMaxMemory );
+        return MemoryAllocator.createAllocator( expectedMaxMemory, new LocalMemoryTracker() );
     }
 
     @Test
@@ -114,4 +117,20 @@ public void allocatingMustIncreaseMemoryUsedAndDecreaseAvailableMemory() throws
         assertThat( mman.usedMemory(), is( greaterThanOrEqualTo( 97L ) ) );
         assertThat( mman.availableMemory(), is( lessThanOrEqualTo( PageCache.PAGE_SIZE - 97L ) ) );
     }
+
+    @Test
+    public void trackMemoryAllocations()
+    {
+        LocalMemoryTracker memoryTracker = new LocalMemoryTracker();
+        MemoryAllocator allocator = MemoryAllocator.createAllocator( "2m", memoryTracker );
+
+        assertEquals( 0, memoryTracker.usedDirectMemory() );
+
+        long pointer = allocator.allocateAligned( ByteUnit.mebiBytes( 1 ), 1 );
+
+        assertEquals( ByteUnit.mebiBytes( 1 ), memoryTracker.usedDirectMemory() );
+
+        allocator.close();
+        assertEquals( 0, memoryTracker.usedDirectMemory() );
+    }
 }
diff --git a/community/io/src/test/java/org/neo4j/io/pagecache/PageSwapperTest.java b/community/io/src/test/java/org/neo4j/io/pagecache/PageSwapperTest.java
@@ -43,6 +43,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.neo4j.io.mem.MemoryAllocator;
+import org.neo4j.memory.LocalMemoryTracker;
 import org.neo4j.test.rule.TestDirectory;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
@@ -70,7 +71,7 @@ public abstract class PageSwapperTest
 
     private final ConcurrentLinkedQueue<PageSwapperFactory> openedFactories = new ConcurrentLinkedQueue<>();
     private final ConcurrentLinkedQueue<PageSwapper> openedSwappers = new ConcurrentLinkedQueue<>();
-    private final MemoryAllocator mman = MemoryAllocator.createAllocator( "32 KiB" );
+    private final MemoryAllocator mman = MemoryAllocator.createAllocator( "32 KiB", new LocalMemoryTracker() );
 
     protected abstract PageSwapperFactory swapperFactory() throws Exception;
 

diff --git a/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/LargePageListIT.java b/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/LargePageListIT.java
@@ -25,6 +25,8 @@
 
 import org.neo4j.io.ByteUnit;
 import org.neo4j.io.mem.MemoryAllocator;
+import org.neo4j.memory.GlobalMemoryTracker;
+import org.neo4j.memory.LocalMemoryTracker;
 
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertEquals;
@@ -43,9 +45,9 @@ public void veryLargePageListsMustBeFullyAccessible() throws Exception
         long pageCacheSize = ByteUnit.gibiBytes( 513 ) + pageSize;
         int pages = Math.toIntExact( pageCacheSize / pageSize );
 
-        MemoryAllocator mman = MemoryAllocator.createAllocator( "2 GiB" );
+        MemoryAllocator mman = MemoryAllocator.createAllocator( "2 GiB", new LocalMemoryTracker() );
         SwapperSet swappers = new SwapperSet();
-        long victimPage = VictimPageReference.getVictimPage( pageSize );
+        long victimPage = VictimPageReference.getVictimPage( pageSize, GlobalMemoryTracker.INSTANCE );
 
         PageList pageList = new PageList( pages, pageSize, mman, swappers, victimPage, Long.BYTES );
 

diff --git a/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/MuninnPageCacheFixture.java b/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/MuninnPageCacheFixture.java
@@ -27,6 +27,7 @@
 import org.neo4j.io.pagecache.PageSwapperFactory;
 import org.neo4j.io.pagecache.tracing.PageCacheTracer;
 import org.neo4j.io.pagecache.tracing.cursor.PageCursorTracerSupplier;
+import org.neo4j.memory.LocalMemoryTracker;
 
 public class MuninnPageCacheFixture extends PageCacheTestSupport.Fixture<MuninnPageCache>
 {
@@ -37,7 +38,8 @@ public MuninnPageCache createPageCache( PageSwapperFactory swapperFactory, int m
                                             PageCacheTracer tracer, PageCursorTracerSupplier cursorTracerSupplier )
     {
         long memory = MuninnPageCache.memoryRequiredForPages( maxPages );
-        MemoryAllocator allocator = MemoryAllocator.createAllocator( String.valueOf( memory ) );
+        MemoryAllocator allocator = MemoryAllocator.createAllocator( String.valueOf( memory ),
+                new LocalMemoryTracker() );
         return new MuninnPageCache( swapperFactory, allocator, tracer, cursorTracerSupplier );
     }
 

diff --git a/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/PageListTest.java b/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/PageListTest.java
@@ -50,6 +50,8 @@
 import org.neo4j.io.pagecache.tracing.FlushEvent;
 import org.neo4j.io.pagecache.tracing.FlushEventOpportunity;
 import org.neo4j.io.pagecache.tracing.PageFaultEvent;
+import org.neo4j.memory.GlobalMemoryTracker;
+import org.neo4j.memory.LocalMemoryTracker;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -87,7 +89,7 @@ public static Iterable<Object[]> parameters()
     public static void setUpStatics()
     {
         executor = Executors.newCachedThreadPool( new DaemonThreadFactory() );
-        mman = MemoryAllocator.createAllocator( "1 MiB" );
+        mman = MemoryAllocator.createAllocator( "1 MiB", new LocalMemoryTracker() );
     }
 
     @AfterClass
@@ -123,7 +125,7 @@ public PageListTest( int pageId )
     public void setUp()
     {
         swappers = new SwapperSet();
-        long victimPage = VictimPageReference.getVictimPage( pageSize );
+        long victimPage = VictimPageReference.getVictimPage( pageSize, GlobalMemoryTracker.INSTANCE );
         pageList = new PageList( pageIds.length, pageSize, mman, swappers, victimPage, ALIGNMENT );
         pageRef = pageList.deref( pageId );
         prevPageRef = pageList.deref( prevPageId );
@@ -134,7 +136,7 @@ public void setUp()
     public void mustExposePageCount() throws Exception
     {
         int pageCount;
-        long victimPage = VictimPageReference.getVictimPage( pageSize );
+        long victimPage = VictimPageReference.getVictimPage( pageSize, GlobalMemoryTracker.INSTANCE );
 
         pageCount = 3;
         assertThat( new PageList( pageCount, pageSize, mman, swappers, victimPage, ALIGNMENT ).getPageCount(),
@@ -1195,7 +1197,8 @@ public void unlockWriteAndTryTakeFlushLockMustInvalidateReadersOverlappingWithWr
     @Test
     public void mustExposeCachePageSize() throws Exception
     {
-        PageList list = new PageList( 0, 42, mman, swappers, VictimPageReference.getVictimPage( 42 ), ALIGNMENT );
+        PageList list = new PageList( 0, 42, mman, swappers,
+                VictimPageReference.getVictimPage( 42, GlobalMemoryTracker.INSTANCE ), ALIGNMENT );
         assertThat( list.getCachePageSize(), is( 42 ) );
     }
 

diff --git a/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/SequenceLockStressIT.java b/community/io/src/test/java/org/neo4j/io/pagecache/impl/muninn/SequenceLockStressIT.java
@@ -35,6 +35,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.neo4j.memory.GlobalMemoryTracker;
 import org.neo4j.test.rule.RepeatRule;
 import org.neo4j.unsafe.impl.internal.dragons.UnsafeUtil;
 
@@ -56,7 +57,7 @@ public static void shutDownExecutor()
     @Before
     public void allocateLock()
     {
-        lockAddr = UnsafeUtil.allocateMemory( Long.BYTES );
+        lockAddr = UnsafeUtil.allocateMemory( Long.BYTES, GlobalMemoryTracker.INSTANCE );
         UnsafeUtil.putLong( lockAddr, 0 );
     }