Fix alignment handling for pooled direct buffers (#11106)

Motivation:
Alignment handling was broken, and basically turned into a fixed offset into each allocation address regardless of its initial value, instead of ensuring that the allocated address is either aligned or bumped to the nearest alignment offset.
The brokenness of the alignment handling extended so far, that overlapping ByteBuf instances could even be created, as was seen in #11101.

Modification:
Instead of fixing the per-allocation pointer bump, we now ensure that 1) the minimum page size is a whole multiple of the alignment, and 2) the reference memory for each chunk is bumped to the nearest aligned address, and finally 3) ensured that the reservations are whole multiples of the alignment, thus ensuring that the next allocation automatically occurs from an aligned address.

Incidentally, (3) above comes for free because the reservations are in whole pages, and in (1) we ensured that pages are sized in whole multiples of the alignment.

In order to ensure that the memory for a chunk is aligned, we introduce some new PlatformDependent infrastructure.
The PlatformDependent.alignDirectBuffer will produce a slice of the given buffer, and the slice will have an address that is aligned.
This method is plainly available on ByteBuffer in Java 9 onwards, but for pre-9 we have to use Unsafe, which means it can fail and might not be available on all platforms.
Attempts to create a PooledByteBufAllocator that uses alignment, when this is not supported, will throw an exception.
Luckily, I think use of aligned allocations are rare.

Result:
Aligned pooled byte bufs now work correctly, and never have any overlap.

Fixes #11101

buffer/src/main/java/io/netty/buffer/PoolArena.java

@@ -41,7 +41,6 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
 
     final int numSmallSubpagePools;
     final int directMemoryCacheAlignment;
-    final int directMemoryCacheAlignmentMask;
     private final PoolSubpage<T>[] smallSubpagePools;
 
     private final PoolChunkList<T> q050;
@@ -77,7 +76,6 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
         super(pageSize, pageShifts, chunkSize, cacheAlignment);
         this.parent = parent;
         directMemoryCacheAlignment = cacheAlignment;
-        directMemoryCacheAlignmentMask = cacheAlignment - 1;
 
         numSmallSubpagePools = nSubpages;
         smallSubpagePools = newSubpagePoolArray(numSmallSubpagePools);
@@ -152,7 +150,7 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
             return;
         }
 
-        /**
+        /*
          * Synchronize on the head. This is needed as {@link PoolChunk#allocateSubpage(int)} and
          * {@link PoolChunk#free(long)} may modify the doubly linked list as well.
          */
@@ -235,7 +233,7 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
         }
     }
 
-    private SizeClass sizeClass(long handle) {
+    private static SizeClass sizeClass(long handle) {
         return isSubpage(handle) ? SizeClass.Small : SizeClass.Normal;
     }
 
@@ -559,12 +557,13 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
 
         @Override
         protected PoolChunk<byte[]> newChunk(int pageSize, int maxPageIdx, int pageShifts, int chunkSize) {
-            return new PoolChunk<byte[]>(this, newByteArray(chunkSize), pageSize, pageShifts, chunkSize, maxPageIdx, 0);
+            return new PoolChunk<byte[]>(
+                    this, null, newByteArray(chunkSize), pageSize, pageShifts, chunkSize, maxPageIdx);
         }
 
         @Override
         protected PoolChunk<byte[]> newUnpooledChunk(int capacity) {
-            return new PoolChunk<byte[]>(this, newByteArray(capacity), capacity, 0);
+            return new PoolChunk<byte[]>(this, null, newByteArray(capacity), capacity);
         }
 
         @Override
@@ -601,43 +600,31 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
             return true;
         }
 
-        // mark as package-private, only for unit test
-        int offsetCacheLine(ByteBuffer memory) {
-            // We can only calculate the offset if Unsafe is present as otherwise directBufferAddress(...) will
-            // throw an NPE.
-            int remainder = HAS_UNSAFE
-                    ? (int) (PlatformDependent.directBufferAddress(memory) & directMemoryCacheAlignmentMask)
-                    : 0;
-
-            // offset = alignment - address & (alignment - 1)
-            return directMemoryCacheAlignment - remainder;
-        }
-
         @Override
         protected PoolChunk<ByteBuffer> newChunk(int pageSize, int maxPageIdx,
             int pageShifts, int chunkSize) {
             if (directMemoryCacheAlignment == 0) {
-                return new PoolChunk<ByteBuffer>(this,
-                        allocateDirect(chunkSize), pageSize, pageShifts,
-                        chunkSize, maxPageIdx, 0);
+                ByteBuffer memory = allocateDirect(chunkSize);
+                return new PoolChunk<ByteBuffer>(this, memory, memory, pageSize, pageShifts,
+                        chunkSize, maxPageIdx);
             }
-            final ByteBuffer memory = allocateDirect(chunkSize
-                    + directMemoryCacheAlignment);
-            return new PoolChunk<ByteBuffer>(this, memory, pageSize,
-                    pageShifts, chunkSize, maxPageIdx,
-                    offsetCacheLine(memory));
+
+            final ByteBuffer base = allocateDirect(chunkSize + directMemoryCacheAlignment);
+            final ByteBuffer memory = PlatformDependent.alignDirectBuffer(base, directMemoryCacheAlignment);
+            return new PoolChunk<ByteBuffer>(this, base, memory, pageSize,
+                    pageShifts, chunkSize, maxPageIdx);
         }
 
         @Override
         protected PoolChunk<ByteBuffer> newUnpooledChunk(int capacity) {
             if (directMemoryCacheAlignment == 0) {
-                return new PoolChunk<ByteBuffer>(this,
-                        allocateDirect(capacity), capacity, 0);
+                ByteBuffer memory = allocateDirect(capacity);
+                return new PoolChunk<ByteBuffer>(this, memory, memory, capacity);
             }
-            final ByteBuffer memory = allocateDirect(capacity
-                    + directMemoryCacheAlignment);
-            return new PoolChunk<ByteBuffer>(this, memory, capacity,
-                    offsetCacheLine(memory));
+
+            final ByteBuffer base = allocateDirect(capacity + directMemoryCacheAlignment);
+            final ByteBuffer memory = PlatformDependent.alignDirectBuffer(base, directMemoryCacheAlignment);
+            return new PoolChunk<ByteBuffer>(this, base, memory, capacity);
         }
 
         private static ByteBuffer allocateDirect(int capacity) {
@@ -648,9 +635,9 @@ abstract class PoolArena<T> extends SizeClasses implements PoolArenaMetric {
         @Override
         protected void destroyChunk(PoolChunk<ByteBuffer> chunk) {
             if (PlatformDependent.useDirectBufferNoCleaner()) {
-                PlatformDependent.freeDirectNoCleaner(chunk.memory);
+                PlatformDependent.freeDirectNoCleaner((ByteBuffer) chunk.base);
             } else {
-                PlatformDependent.freeDirectBuffer(chunk.memory);
+                PlatformDependent.freeDirectBuffer((ByteBuffer) chunk.base);
             }
         }
 

buffer/src/main/java/io/netty/buffer/PoolChunk.java

@@ -141,9 +141,9 @@ final class PoolChunk<T> implements PoolChunkMetric {
     static final int RUN_OFFSET_SHIFT = SIZE_BIT_LENGTH + SIZE_SHIFT;
 
     final PoolArena<T> arena;
+    final Object base;
     final T memory;
     final boolean unpooled;
-    final int offset;
 
     /**
      * store the first page and last page of each avail run
@@ -181,14 +181,14 @@ final class PoolChunk<T> implements PoolChunkMetric {
     //private long pad0, pad1, pad2, pad3, pad4, pad5, pad6, pad7;
 
     @SuppressWarnings("unchecked")
-    PoolChunk(PoolArena<T> arena, T memory, int pageSize, int pageShifts, int chunkSize, int maxPageIdx, int offset) {
+    PoolChunk(PoolArena<T> arena, Object base, T memory, int pageSize, int pageShifts, int chunkSize, int maxPageIdx) {
         unpooled = false;
         this.arena = arena;
+        this.base = base;
         this.memory = memory;
         this.pageSize = pageSize;
         this.pageShifts = pageShifts;
         this.chunkSize = chunkSize;
-        this.offset = offset;
         freeBytes = chunkSize;
 
         runsAvail = newRunsAvailqueueArray(maxPageIdx);
@@ -204,11 +204,11 @@ final class PoolChunk<T> implements PoolChunkMetric {
     }
 
     /** Creates a special chunk that is not pooled. */
-    PoolChunk(PoolArena<T> arena, T memory, int size, int offset) {
+    PoolChunk(PoolArena<T> arena, Object base, T memory, int size) {
         unpooled = true;
         this.arena = arena;
+        this.base = base;
         this.memory = memory;
-        this.offset = offset;
         pageSize = 0;
         pageShifts = 0;
         runsAvailMap = null;
@@ -569,9 +569,8 @@ final class PoolChunk<T> implements PoolChunkMetric {
         assert s.doNotDestroy;
         assert reqCapacity <= s.elemSize;
 
-        buf.init(this, nioBuffer, handle,
-                 (runOffset << pageShifts) + bitmapIdx * s.elemSize + offset,
-                 reqCapacity, s.elemSize, threadCache);
+        int offset = (runOffset << pageShifts) + bitmapIdx * s.elemSize;
+        buf.init(this, nioBuffer, handle, offset, reqCapacity, s.elemSize, threadCache);
     }
 
     @Override

buffer/src/main/java/io/netty/buffer/PooledByteBuf.java

@@ -52,7 +52,7 @@ abstract class PooledByteBuf<T> extends AbstractReferenceCountedByteBuf {
     }
 
     void initUnpooled(PoolChunk<T> chunk, int length) {
-        init0(chunk, null, 0, chunk.offset, length, length, null);
+        init0(chunk, null, 0, 0, length, length, null);
     }
 
     private void init0(PoolChunk<T> chunk, ByteBuffer nioBuffer,

buffer/src/main/java/io/netty/buffer/PooledByteBufAllocator.java

@@ -63,15 +63,19 @@ public class PooledByteBufAllocator extends AbstractByteBufAllocator implements
     };
 
     static {
+        int defaultAlignment = SystemPropertyUtil.getInt(
+                "io.netty.allocator.directMemoryCacheAlignment", 0);
         int defaultPageSize = SystemPropertyUtil.getInt("io.netty.allocator.pageSize", 8192);
         Throwable pageSizeFallbackCause = null;
         try {
-            validateAndCalculatePageShifts(defaultPageSize);
+            validateAndCalculatePageShifts(defaultPageSize, defaultAlignment);
         } catch (Throwable t) {
             pageSizeFallbackCause = t;
             defaultPageSize = 8192;
+            defaultAlignment = 0;
         }
         DEFAULT_PAGE_SIZE = defaultPageSize;
+        DEFAULT_DIRECT_MEMORY_CACHE_ALIGNMENT = defaultAlignment;
 
         int defaultMaxOrder = SystemPropertyUtil.getInt("io.netty.allocator.maxOrder", 11);
         Throwable maxOrderFallbackCause = null;
@@ -142,9 +146,6 @@ public class PooledByteBufAllocator extends AbstractByteBufAllocator implements
         DEFAULT_USE_CACHE_FOR_ALL_THREADS = SystemPropertyUtil.getBoolean(
                 "io.netty.allocator.useCacheForAllThreads", true);
 
-        DEFAULT_DIRECT_MEMORY_CACHE_ALIGNMENT = SystemPropertyUtil.getInt(
-                "io.netty.allocator.directMemoryCacheAlignment", 0);
-
         // Use 1023 by default as we use an ArrayDeque as backing storage which will then allocate an internal array
         // of 1024 elements. Otherwise we would allocate 2048 and only use 1024 which is wasteful.
         DEFAULT_MAX_CACHED_BYTEBUFFERS_PER_CHUNK = SystemPropertyUtil.getInt(
@@ -266,6 +267,17 @@ public class PooledByteBufAllocator extends AbstractByteBufAllocator implements
         threadCache = new PoolThreadLocalCache(useCacheForAllThreads);
         this.smallCacheSize = smallCacheSize;
         this.normalCacheSize = normalCacheSize;
+
+        if (directMemoryCacheAlignment != 0) {
+            if (!PlatformDependent.hasAlignDirectByteBuffer()) {
+                throw new UnsupportedOperationException("Buffer alignment is not supported. " +
+                        "Either Unsafe or ByteBuffer.alignSlice() must be available.");
+            }
+
+            // Ensure page size is a whole multiple of the alignment, or bump it to the next whole multiple.
+            pageSize = (int) PlatformDependent.align(pageSize, directMemoryCacheAlignment);
+        }
+
         chunkSize = validateAndCalculateChunkSize(pageSize, maxOrder);
 
         checkPositiveOrZero(nHeapArena, "nHeapArena");
@@ -281,7 +293,7 @@ public class PooledByteBufAllocator extends AbstractByteBufAllocator implements
                     + directMemoryCacheAlignment + " (expected: power of two)");
         }
 
-        int pageShifts = validateAndCalculatePageShifts(pageSize);
+        int pageShifts = validateAndCalculatePageShifts(pageSize, directMemoryCacheAlignment);
 
         if (nHeapArena > 0) {
             heapArenas = newArenaArray(nHeapArena);
@@ -321,15 +333,20 @@ public class PooledByteBufAllocator extends AbstractByteBufAllocator implements
         return new PoolArena[size];
     }
 
-    private static int validateAndCalculatePageShifts(int pageSize) {
+    private static int validateAndCalculatePageShifts(int pageSize, int alignment) {
         if (pageSize < MIN_PAGE_SIZE) {
-            throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: " + MIN_PAGE_SIZE + ")");
+            throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: " + MIN_PAGE_SIZE + ')');
         }
 
         if ((pageSize & pageSize - 1) != 0) {
             throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: power of 2)");
         }
 
+        if (pageSize < alignment) {
+            throw new IllegalArgumentException("Alignment cannot be greater than page size. " +
+                    "Alignment: " + alignment + ", page size: " + pageSize + '.');
+        }
+
         // Logarithm base 2. At this point we know that pageSize is a power of two.
         return Integer.SIZE - 1 - Integer.numberOfLeadingZeros(pageSize);
     }

buffer/src/test/java/io/netty/buffer/AbstractPooledByteBufTest.java

@@ -125,4 +125,18 @@ public abstract class AbstractPooledByteBufTest extends AbstractByteBufTest {
         assertTrue(buf.isContiguous());
         buf.release();
     }
+
+    @Test
+    public void distinctBuffersMustNotOverlap() {
+        ByteBuf a = newBuffer(16384);
+        ByteBuf b = newBuffer(65536);
+        a.setByte(a.capacity() - 1, 1);
+        b.setByte(0, 2);
+        try {
+            assertEquals(1, a.getByte(a.capacity() - 1));
+        } finally {
+            a.release();
+            b.release();
+        }
+    }
 }

buffer/src/test/java/io/netty/buffer/AlignedPooledByteBufAllocatorTest.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2021 The Netty Project
+ *
+ * The Netty Project licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package io.netty.buffer;
+
+public class AlignedPooledByteBufAllocatorTest extends PooledByteBufAllocatorTest {
+    @Override
+    protected PooledByteBufAllocator newAllocator(boolean preferDirect) {
+        int directMemoryCacheAlignment = 1;
+        return new PooledByteBufAllocator(
+                preferDirect,
+                PooledByteBufAllocator.defaultNumHeapArena(),
+                PooledByteBufAllocator.defaultNumDirectArena(),
+                PooledByteBufAllocator.defaultPageSize(),
+                11,
+                PooledByteBufAllocator.defaultSmallCacheSize(),
+                64,
+                PooledByteBufAllocator.defaultUseCacheForAllThreads(),
+                directMemoryCacheAlignment);
+    }
+}

buffer/src/test/java/io/netty/buffer/PoolArenaTest.java

@@ -103,26 +103,6 @@ public class PoolArenaTest {
         }
     }
 
-    @Test
-    public void testDirectArenaOffsetCacheLine() throws Exception {
-        assumeTrue(PlatformDependent.hasUnsafe());
-        int capacity = 5;
-        int alignment = 128;
-
-        for (int i = 0; i < 1000; i++) {
-            ByteBuffer bb = PlatformDependent.useDirectBufferNoCleaner()
-                    ? PlatformDependent.allocateDirectNoCleaner(capacity + alignment)
-                    : ByteBuffer.allocateDirect(capacity + alignment);
-
-            PoolArena.DirectArena arena = new PoolArena.DirectArena(null, 512, 9, 512, alignment);
-            int offset = arena.offsetCacheLine(bb);
-            long address = PlatformDependent.directBufferAddress(bb);
-
-            Assert.assertEquals(0, (offset + address) & (alignment - 1));
-            PlatformDependent.freeDirectBuffer(bb);
-        }
-    }
-
     @Test
     public void testAllocationCounter() {
         final PooledByteBufAllocator allocator = new PooledByteBufAllocator(

buffer/src/test/java/io/netty/buffer/PooledAlignedBigEndianDirectByteBufTest.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2021 The Netty Project
+ *
+ * The Netty Project licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package io.netty.buffer;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+
+import java.nio.ByteOrder;
+
+import static org.junit.Assert.assertSame;
+
+public class PooledAlignedBigEndianDirectByteBufTest extends PooledBigEndianDirectByteBufTest {
+    private static final int directMemoryCacheAlignment = 1;
+    private static PooledByteBufAllocator allocator;
+
+    @BeforeClass
+    public static void setUpAllocator() {
+        allocator = new PooledByteBufAllocator(
+                true,
+                PooledByteBufAllocator.defaultNumHeapArena(),
+                PooledByteBufAllocator.defaultNumDirectArena(),
+                PooledByteBufAllocator.defaultPageSize(),
+                11,
+                PooledByteBufAllocator.defaultSmallCacheSize(),
+                64,
+                PooledByteBufAllocator.defaultUseCacheForAllThreads(),
+                directMemoryCacheAlignment);
+    }
+
+    @AfterClass
+    public static void releaseAllocator() {
+        allocator = null;
+    }
+
+    @Override
+    protected ByteBuf alloc(int length, int maxCapacity) {
+        ByteBuf buffer = allocator.directBuffer(length, maxCapacity);
+        assertSame(ByteOrder.BIG_ENDIAN, buffer.order());
+        return buffer;
+    }
+}

common/src/main/java/io/netty/util/internal/PlatformDependent.java

@@ -768,6 +768,32 @@ public final class PlatformDependent {
         decrementMemoryCounter(capacity);
     }
 
+    public static boolean hasAlignDirectByteBuffer() {
+        return hasUnsafe() || PlatformDependent0.hasAlignSliceMethod();
+    }
+
+    public static ByteBuffer alignDirectBuffer(ByteBuffer buffer, int alignment) {
+        if (!buffer.isDirect()) {
+            throw new IllegalArgumentException("Cannot get aligned slice of non-direct byte buffer.");
+        }
+        if (PlatformDependent0.hasAlignSliceMethod()) {
+            return PlatformDependent0.alignSlice(buffer, alignment);
+        }
+        if (hasUnsafe()) {
+            long address = directBufferAddress(buffer);
+            long aligned = align(address, alignment);
+            buffer.position((int) (aligned - address));
+            return buffer.slice();
+        }
+        // We don't have enough information to be able to align any buffers.
+        throw new UnsupportedOperationException("Cannot align direct buffer. " +
+                "Needs either Unsafe or ByteBuffer.alignSlice method available.");
+    }
+
+    public static long align(long value, int alignment) {
+        return Pow2.align(value, alignment);
+    }
+
     private static void incrementMemoryCounter(int capacity) {
         if (DIRECT_MEMORY_COUNTER != null) {
             long newUsedMemory = DIRECT_MEMORY_COUNTER.addAndGet(capacity);

common/src/main/java/io/netty/util/internal/PlatformDependent0.java

@@ -46,6 +46,7 @@ final class PlatformDependent0 {
     private static final Constructor<?> DIRECT_BUFFER_CONSTRUCTOR;
     private static final Throwable EXPLICIT_NO_UNSAFE_CAUSE = explicitNoUnsafeCause0();
     private static final Method ALLOCATE_ARRAY_METHOD;
+    private static final Method ALIGN_SLICE;
     private static final int JAVA_VERSION = javaVersion0();
     private static final boolean IS_ANDROID = isAndroid0();
 
@@ -398,6 +399,21 @@ final class PlatformDependent0 {
             ALLOCATE_ARRAY_METHOD = allocateArrayMethod;
         }
 
+        if (javaVersion() > 9) {
+            ALIGN_SLICE = (Method) AccessController.doPrivileged(new PrivilegedAction<Object>() {
+                @Override
+                public Object run() {
+                    try {
+                        return ByteBuffer.class.getDeclaredMethod("alignedSlice", int.class);
+                    } catch (Exception e) {
+                        return null;
+                    }
+                }
+            });
+        } else {
+            ALIGN_SLICE = null;
+        }
+
         INTERNAL_UNSAFE = internalUnsafe;
 
         logger.debug("java.nio.DirectByteBuffer.<init>(long, int): {}",
@@ -474,6 +490,20 @@ final class PlatformDependent0 {
         return newDirectBuffer(UNSAFE.allocateMemory(Math.max(1, capacity)), capacity);
     }
 
+    static boolean hasAlignSliceMethod() {
+        return ALIGN_SLICE != null;
+    }
+
+    static ByteBuffer alignSlice(ByteBuffer buffer, int alignment) {
+        try {
+            return (ByteBuffer) ALIGN_SLICE.invoke(buffer, alignment);
+        } catch (IllegalAccessException e) {
+            throw new Error(e);
+        } catch (InvocationTargetException e) {
+            throw new Error(e);
+        }
+    }
+
     static boolean hasAllocateArrayMethod() {
         return ALLOCATE_ARRAY_METHOD != null;
     }