Unbind alignment from MemoryAllocator

This makes it possible to specify the required memory alignment at allocation
time, instead of giving it as a constructor argument to the MemoryAllocator.

community/io/src/main/java/org/neo4j/io/mem/GrabAllocator.java

@@ -38,7 +38,6 @@ public final class GrabAllocator implements MemoryAllocator
      * The amount of memory that this memory manager can still allocate.
      */
     private long memoryReserve;
-    private final long alignment;
 
     private Grab grabs;
 
@@ -47,16 +46,10 @@ 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 alignment The byte multiple that the allocated pointers have to be aligned at.
      */
-    GrabAllocator( long expectedMaxMemory, long alignment )
+    GrabAllocator( long expectedMaxMemory )
     {
-        if ( alignment == 0 )
-        {
-            throw new IllegalArgumentException( "Alignment cannot be zero" );
-        }
         this.memoryReserve = expectedMaxMemory;
-        this.alignment = alignment;
     }
 
     @Override
@@ -66,7 +59,7 @@ public synchronized long usedMemory()
         Grab grab = grabs;
         while ( grab != null )
         {
-            sum += grab.nextAlignedPointer - grab.address;
+            sum += grab.nextPointer - grab.address;
             grab = grab.next;
         }
         return sum;
@@ -79,25 +72,29 @@ public synchronized long availableMemory()
         long availableInCurrentGrab = 0;
         if ( grab != null )
         {
-            availableInCurrentGrab = grab.limit - grab.nextAlignedPointer;
+            availableInCurrentGrab = grab.limit - grab.nextPointer;
         }
         return Math.max( memoryReserve, 0L ) + availableInCurrentGrab;
     }
 
     @Override
-    public synchronized long allocateAligned( long bytes )
+    public synchronized long allocateAligned( long bytes, long alignment )
     {
+        if ( alignment <= 0 )
+        {
+            throw new IllegalArgumentException( "Invalid alignment: " + alignment + "; alignment must be positive" );
+        }
         if ( bytes > GRAB_SIZE )
         {
             // This is a huge allocation. Put it in its own grab and keep any existing grab at the head.
             Grab nextGrab = grabs == null ? null : grabs.next;
-            Grab allocationGrab = new Grab( nextGrab, bytes, alignment );
+            Grab allocationGrab = new Grab( nextGrab, bytes );
             if ( !allocationGrab.canAllocate( bytes ) )
             {
                 allocationGrab.free();
-                allocationGrab = new Grab( nextGrab, bytes + alignment, alignment );
+                allocationGrab = new Grab( nextGrab, bytes + alignment );
             }
-            long allocation = allocationGrab.allocate( bytes );
+            long allocation = allocationGrab.allocate( bytes, alignment );
             grabs = grabs == null ? allocationGrab : grabs.setNext( allocationGrab );
             memoryReserve -= bytes;
             return allocation;
@@ -109,20 +106,20 @@ public synchronized long allocateAligned( long bytes )
             if ( desiredGrabSize < bytes )
             {
                 desiredGrabSize = bytes;
-                Grab grab = new Grab( grabs, desiredGrabSize, alignment );
+                Grab grab = new Grab( grabs, desiredGrabSize );
                 if ( grab.canAllocate( bytes ) )
                 {
                     memoryReserve -= desiredGrabSize;
                     grabs = grab;
-                    return grabs.allocate( bytes );
+                    return grabs.allocate( bytes, alignment );
                 }
                 grab.free();
                 desiredGrabSize = bytes + alignment;
             }
             memoryReserve -= desiredGrabSize;
-            grabs = new Grab( grabs, desiredGrabSize, alignment );
+            grabs = new Grab( grabs, desiredGrabSize );
         }
-        return grabs.allocate( bytes );
+        return grabs.allocate( bytes, alignment );
     }
 
     @Override
@@ -175,41 +172,39 @@ private static class Grab
         public final Grab next;
         private final long address;
         private final long limit;
-        private final long alignMask;
-        private long nextAlignedPointer;
+        private long nextPointer;
 
-        Grab( Grab next, long size, long alignment )
+        Grab( Grab next, long size )
         {
             this.next = next;
             this.address = allocateNativeMemory( size );
             this.limit = address + size;
-            this.alignMask = alignment - 1;
 
-            nextAlignedPointer = nextAligned( this.address );
+            nextPointer = address;
         }
 
-        Grab( Grab next, long address, long limit, long alignMask, long nextAlignedPointer )
+        Grab( Grab next, long address, long limit, long nextPointer )
         {
             this.next = next;
             this.address = address;
             this.limit = limit;
-            this.alignMask = alignMask;
-            this.nextAlignedPointer = nextAlignedPointer;
+            this.nextPointer = nextPointer;
         }
 
-        private long nextAligned( long pointer )
+        private long nextAligned( long pointer, long alignment )
         {
-            if ( (pointer & ~alignMask) == pointer )
+            long mask = alignment - 1;
+            if ( (pointer & ~mask) == pointer )
             {
                 return pointer;
             }
-            return (pointer + alignMask) & ~alignMask;
+            return (pointer + mask) & ~mask;
         }
 
-        long allocate( long bytes )
+        long allocate( long bytes, long alignment )
         {
-            long allocation = nextAlignedPointer;
-            nextAlignedPointer = nextAligned( nextAlignedPointer + bytes );
+            long allocation = nextAligned( nextPointer, alignment );
+            nextPointer = allocation + bytes;
             return allocation;
         }
 
@@ -220,19 +215,19 @@ void free()
 
         boolean canAllocate( long bytes )
         {
-            return nextAlignedPointer + bytes <= limit;
+            return nextPointer + bytes <= limit;
         }
 
         Grab setNext( Grab grab )
         {
-            return new Grab( grab, address, limit, alignMask, nextAlignedPointer );
+            return new Grab( grab, address, limit, nextPointer );
         }
 
         @Override
         public String toString()
         {
             long size = limit - address;
-            long reserve = nextAlignedPointer > limit ? 0 : limit - nextAlignedPointer;
+            long reserve = nextPointer > limit ? 0 : limit - nextPointer;
             double use = (1.0 - reserve / ((double) size)) * 100.0;
             return String.format( "Grab[size = %d bytes, reserve = %d bytes, use = %5.2f %%]", size, reserve, use );
         }

community/io/src/main/java/org/neo4j/io/mem/MemoryAllocator.java

@@ -24,9 +24,9 @@
  */
 public interface MemoryAllocator
 {
-    static MemoryAllocator createAllocator( long expectedMaxMemory, long alignment )
+    static MemoryAllocator createAllocator( long expectedMaxMemory )
     {
-        return new GrabAllocator( expectedMaxMemory, alignment );
+        return new GrabAllocator( expectedMaxMemory );
     }
 
     /**
@@ -42,8 +42,9 @@ static MemoryAllocator createAllocator( long expectedMaxMemory, long alignment )
     /**
      * Allocate a contiguous, aligned region of memory of the given size in bytes.
      * @param bytes the number of bytes to allocate.
+     * @param alignment The byte multiple that the allocated pointers have to be aligned at.
      * @return A pointer to the allocated memory.
      * @throws OutOfMemoryError if the requested memory could not be allocated.
      */
-    long allocateAligned( long bytes );
+    long allocateAligned( long bytes, long alignment );
 }

community/unsafe/src/main/java/org/neo4j/unsafe/impl/internal/dragons/NativeMemoryAllocationRefusedError.java → community/io/src/main/java/org/neo4j/io/mem/NativeMemoryAllocationRefusedError.java

@@ -17,7 +17,7 @@
  * 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.unsafe.impl.internal.dragons;
+package org.neo4j.io.mem;
 
 public class NativeMemoryAllocationRefusedError extends OutOfMemoryError
 {

community/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/MuninnPageCache.java

@@ -233,10 +233,10 @@ public MuninnPageCache(
 
         long alignment = swapperFactory.getRequiredBufferAlignment();
         long expectedMaxMemory = ((long) maxPages) * cachePageSize; // cast to long prevents overflow
-        MemoryAllocator memoryAllocator = MemoryAllocator.createAllocator( expectedMaxMemory, alignment );
+        MemoryAllocator memoryAllocator = MemoryAllocator.createAllocator( expectedMaxMemory );
         this.victimPage = VictimPageReference.getVictimPage( cachePageSize );
 
-        this.pages = new PageList( maxPages, cachePageSize, memoryAllocator, new SwapperSet(), victimPage );
+        this.pages = new PageList( maxPages, cachePageSize, memoryAllocator, new SwapperSet(), victimPage, alignment );
 
         setFreelistHead( new AtomicInteger() );
     }

community/io/src/main/java/org/neo4j/io/pagecache/impl/muninn/PageList.java

@@ -95,16 +95,19 @@ class PageList
     private final SwapperSet swappers;
     private final long victimPageAddress;
     private final long baseAddress;
+    private final long bufferAlignment;
 
-    PageList( int pageCount, int cachePageSize, MemoryAllocator memoryAllocator, SwapperSet swappers, long victimPageAddress )
+    PageList( int pageCount, int cachePageSize, MemoryAllocator memoryAllocator, SwapperSet swappers,
+              long victimPageAddress, long bufferAlignment )
     {
         this.pageCount = pageCount;
         this.cachePageSize = cachePageSize;
         this.memoryAllocator = memoryAllocator;
         this.swappers = swappers;
         this.victimPageAddress = victimPageAddress;
         long bytes = ((long) pageCount) * META_DATA_BYTES_PER_PAGE;
-        this.baseAddress = memoryAllocator.allocateAligned( bytes );
+        this.baseAddress = memoryAllocator.allocateAligned( bytes, Long.BYTES );
+        this.bufferAlignment = bufferAlignment;
         clearMemory( baseAddress, pageCount );
     }
 
@@ -123,6 +126,7 @@ class PageList
         this.swappers = pageList.swappers;
         this.victimPageAddress = pageList.victimPageAddress;
         this.baseAddress = pageList.baseAddress;
+        this.bufferAlignment = pageList.bufferAlignment;
     }
 
     private void clearMemory( long baseAddress, long pageCount )
@@ -142,13 +146,14 @@ private void clearMemory( long baseAddress, long pageCount )
 
     private void clearMemorySimple( long baseAddress, long pageCount )
     {
-        long address = baseAddress - 8;
+        long address = baseAddress - Long.BYTES;
+        long initialLockWord = OffHeapPageLock.initialLockWordWithExclusiveLock();
         for ( long i = 0; i < pageCount; i++ )
         {
-            UnsafeUtil.putLong( address += 8, OffHeapPageLock.initialLockWordWithExclusiveLock() ); // lock word
-            UnsafeUtil.putLong( address += 8, 0 ); // pointer
-            UnsafeUtil.putLong( address += 8, PageCursor.UNBOUND_PAGE_ID ); // file page id
-            UnsafeUtil.putLong( address += 8, 0 ); // rest
+            UnsafeUtil.putLong( address += Long.BYTES, initialLockWord );
+            UnsafeUtil.putLong( address += Long.BYTES, 0 ); // page buffer address pointer
+            UnsafeUtil.putLong( address += Long.BYTES, PageCursor.UNBOUND_PAGE_ID ); // file page id
+            UnsafeUtil.putLong( address += Long.BYTES, 0 ); // rest
         }
     }
 
@@ -306,7 +311,7 @@ public void initBuffer( long pageRef )
     {
         if ( getAddress( pageRef ) == 0L )
         {
-            long addr = memoryAllocator.allocateAligned( getCachePageSize() );
+            long addr = memoryAllocator.allocateAligned( getCachePageSize(), bufferAlignment );
             UnsafeUtil.putLong( offAddress( pageRef ), addr );
         }
     }

community/io/src/test/java/org/neo4j/io/mem/MemoryAllocatorTest.java

@@ -32,51 +32,51 @@
 
 public class MemoryAllocatorTest
 {
-    protected MemoryAllocator createAllocator( long expectedMaxMemory, long alignment )
+    protected MemoryAllocator createAllocator( long expectedMaxMemory )
     {
-        return MemoryAllocator.createAllocator( expectedMaxMemory, alignment );
+        return MemoryAllocator.createAllocator( expectedMaxMemory );
     }
 
     @Test
     public void allocatedPointerMustNotBeNull() throws Exception
     {
-        MemoryAllocator mman = createAllocator( 8 * PageCache.PAGE_SIZE, 8 );
-        long address = mman.allocateAligned( PageCache.PAGE_SIZE );
+        MemoryAllocator mman = createAllocator( 8 * PageCache.PAGE_SIZE );
+        long address = mman.allocateAligned( PageCache.PAGE_SIZE, 8 );
         assertThat( address, is( not( 0L ) ) );
     }
 
     @Test
     public void allocatedPointerMustBePageAligned() throws Exception
     {
-        MemoryAllocator mman = createAllocator( 8 * PageCache.PAGE_SIZE, UnsafeUtil.pageSize() );
-        long address = mman.allocateAligned( PageCache.PAGE_SIZE );
+        MemoryAllocator mman = createAllocator( 8 * PageCache.PAGE_SIZE );
+        long address = mman.allocateAligned( PageCache.PAGE_SIZE, UnsafeUtil.pageSize() );
         assertThat( address % UnsafeUtil.pageSize(), is( 0L ) );
     }
 
     @Test
     public void mustBeAbleToAllocatePastMemoryLimit() throws Exception
     {
-        MemoryAllocator mman = createAllocator( PageCache.PAGE_SIZE, 2 );
+        MemoryAllocator mman = createAllocator( PageCache.PAGE_SIZE );
         for ( int i = 0; i < 4100; i++ )
         {
-            assertThat( mman.allocateAligned( 1 ) % 2, is( 0L ) );
+            assertThat( mman.allocateAligned( 1, 2 ) % 2, is( 0L ) );
         }
         // Also asserts that no OutOfMemoryError is thrown.
     }
 
     @Test( expected = IllegalArgumentException.class )
     public void alignmentCannotBeZero() throws Exception
     {
-        createAllocator( PageCache.PAGE_SIZE, 0 );
+        createAllocator( PageCache.PAGE_SIZE ).allocateAligned( 8, 0 );
     }
 
     @Test
     public void mustBeAbleToAllocateSlabsLargerThanGrabSize() throws Exception
     {
-        MemoryAllocator mman = createAllocator( 32 * 1024 * 1024, 1 );
-        long page1 = mman.allocateAligned( UnsafeUtil.pageSize() );
-        long largeBlock = mman.allocateAligned( 1024 * 1024 ); // 1 MiB
-        long page2 = mman.allocateAligned( UnsafeUtil.pageSize() );
+        MemoryAllocator mman = createAllocator( 32 * 1024 * 1024 );
+        long page1 = mman.allocateAligned( UnsafeUtil.pageSize(), 1 );
+        long largeBlock = mman.allocateAligned( 1024 * 1024, 1 ); // 1 MiB
+        long page2 = mman.allocateAligned( UnsafeUtil.pageSize(), 1 );
         assertThat( page1, is( not( 0L ) ) );
         assertThat( largeBlock, is( not( 0L ) ) );
         assertThat( page2, is( not( 0L ) ) );
@@ -85,17 +85,17 @@ public void mustBeAbleToAllocateSlabsLargerThanGrabSize() throws Exception
     @Test
     public void allocatingMustIncreaseMemoryUsedAndDecreaseAvailableMemory() throws Exception
     {
-        MemoryAllocator mman = createAllocator( PageCache.PAGE_SIZE, 1 );
+        MemoryAllocator mman = createAllocator( PageCache.PAGE_SIZE );
         assertThat( mman.usedMemory(), is( 0L ) );
         assertThat( mman.availableMemory(), is( (long) PageCache.PAGE_SIZE ) );
         assertThat( mman.usedMemory() + mman.availableMemory(), is( (long) PageCache.PAGE_SIZE ) );
 
-        mman.allocateAligned( 32 );
+        mman.allocateAligned( 32, 1 );
         assertThat( mman.usedMemory(), is( greaterThanOrEqualTo( 32L ) ) );
         assertThat( mman.availableMemory(), is( lessThanOrEqualTo( PageCache.PAGE_SIZE - 32L ) ) );
         assertThat( mman.usedMemory() + mman.availableMemory(), is( (long) PageCache.PAGE_SIZE ) );
 
-        mman.allocateAligned( 32 );
+        mman.allocateAligned( 32, 1 );
         assertThat( mman.usedMemory(), is( greaterThanOrEqualTo( 64L ) ) );
         assertThat( mman.availableMemory(), is( lessThanOrEqualTo( PageCache.PAGE_SIZE - 32 - 32L ) ) );
         assertThat( mman.usedMemory() + mman.availableMemory(), is( (long) PageCache.PAGE_SIZE ) );

community/io/src/test/java/org/neo4j/io/pagecache/PageSwapperTest.java

@@ -71,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( ByteUnit.kibiBytes( 32 ), 1 );
+    private final MemoryAllocator mman = MemoryAllocator.createAllocator( ByteUnit.kibiBytes( 32 ) );
 
     protected abstract PageSwapperFactory swapperFactory() throws Exception;
 
@@ -95,7 +95,7 @@ protected int cachePageSize()
 
     protected long createPage( int cachePageSize )
     {
-        long address = mman.allocateAligned( cachePageSize + Integer.BYTES );
+        long address = mman.allocateAligned( cachePageSize + Integer.BYTES, 1 );
         UnsafeUtil.putInt( address, cachePageSize );
         return address + Integer.BYTES;
     }