DB-2993: export used direct memory

Author: Stefania Alborghetti

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

@@ -737,10 +737,10 @@ protected PoolChunk<ByteBuffer> newChunk(int pageSize, int maxOrder,
                 int pageShifts, int chunkSize) {
             if (directMemoryCacheAlignment == 0) {
                 return new PoolChunk<ByteBuffer>(this,
-                        allocateDirect(chunkSize), pageSize, maxOrder,
+                        PlatformDependent.allocateDirect(chunkSize), pageSize, maxOrder,
                         pageShifts, chunkSize, 0);
             }
-            final ByteBuffer memory = allocateDirect(chunkSize
+            final ByteBuffer memory = PlatformDependent.allocateDirect(chunkSize
                     + directMemoryCacheAlignment);
             return new PoolChunk<ByteBuffer>(this, memory, pageSize,
                     maxOrder, pageShifts, chunkSize,
@@ -751,26 +751,17 @@ protected PoolChunk<ByteBuffer> newChunk(int pageSize, int maxOrder,
         protected PoolChunk<ByteBuffer> newUnpooledChunk(int capacity) {
             if (directMemoryCacheAlignment == 0) {
                 return new PoolChunk<ByteBuffer>(this,
-                        allocateDirect(capacity), capacity, 0);
+                          PlatformDependent.allocateDirect(capacity), capacity, 0);
             }
-            final ByteBuffer memory = allocateDirect(capacity
+            final ByteBuffer memory = PlatformDependent.allocateDirect(capacity
                     + directMemoryCacheAlignment);
             return new PoolChunk<ByteBuffer>(this, memory, capacity,
                     offsetCacheLine(memory));
         }
 
-        private static ByteBuffer allocateDirect(int capacity) {
-            return PlatformDependent.useDirectBufferNoCleaner() ?
-                    PlatformDependent.allocateDirectNoCleaner(capacity) : ByteBuffer.allocateDirect(capacity);
-        }
-
         @Override
         protected void destroyChunk(PoolChunk<ByteBuffer> chunk) {
-            if (PlatformDependent.useDirectBufferNoCleaner()) {
-                PlatformDependent.freeDirectNoCleaner(chunk.memory);
-            } else {
-                PlatformDependent.freeDirectBuffer(chunk.memory);
-            }
+            PlatformDependent.freeDirectBuffer(chunk.memory);
         }
 
         @Override

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

@@ -64,7 +64,7 @@ public UnpooledDirectByteBuf(ByteBufAllocator alloc, int initialCapacity, int ma
         }
 
         this.alloc = alloc;
-        setByteBuffer(ByteBuffer.allocateDirect(initialCapacity));
+        setByteBuffer(allocateDirect(initialCapacity));
     }
 
     /**
@@ -103,14 +103,14 @@ protected UnpooledDirectByteBuf(ByteBufAllocator alloc, ByteBuffer initialBuffer
      * Allocate a new direct {@link ByteBuffer} with the given initialCapacity.
      */
     protected ByteBuffer allocateDirect(int initialCapacity) {
-        return ByteBuffer.allocateDirect(initialCapacity);
+        return PlatformDependent.allocateDirectJVM(initialCapacity);
     }
 
     /**
      * Free a direct {@link ByteBuffer}
      */
     protected void freeDirect(ByteBuffer buffer) {
-        PlatformDependent.freeDirectBuffer(buffer);
+        PlatformDependent.freeDirectBufferJVM(buffer);
     }
 
     private void setByteBuffer(ByteBuffer buffer) {

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

@@ -117,14 +117,14 @@ protected UnpooledUnsafeDirectByteBuf(ByteBufAllocator alloc, ByteBuffer initial
      * Allocate a new direct {@link ByteBuffer} with the given initialCapacity.
      */
     protected ByteBuffer allocateDirect(int initialCapacity) {
-        return ByteBuffer.allocateDirect(initialCapacity);
+        return PlatformDependent.allocateDirectJVM(initialCapacity);
     }
 
     /**
      * Free a direct {@link ByteBuffer}
      */
     protected void freeDirect(ByteBuffer buffer) {
-        PlatformDependent.freeDirectBuffer(buffer);
+        PlatformDependent.freeDirectBufferJVM(buffer);
     }
 
     final void setByteBuffer(ByteBuffer buffer, boolean tryFree) {

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

@@ -154,21 +154,16 @@ public Random current() {
 
         if (maxDirectMemory == 0 || !hasUnsafe() || !PlatformDependent0.hasDirectBufferNoCleanerConstructor()) {
             USE_DIRECT_BUFFER_NO_CLEANER = false;
-            DIRECT_MEMORY_COUNTER = null;
         } else {
             USE_DIRECT_BUFFER_NO_CLEANER = true;
             if (maxDirectMemory < 0) {
                 maxDirectMemory = maxDirectMemory0();
-                if (maxDirectMemory <= 0) {
-                    DIRECT_MEMORY_COUNTER = null;
-                } else {
-                    DIRECT_MEMORY_COUNTER = new AtomicLong();
-                }
-            } else {
-                DIRECT_MEMORY_COUNTER = new AtomicLong();
             }
         }
+
+        DIRECT_MEMORY_COUNTER = new AtomicLong();
         DIRECT_MEMORY_LIMIT = maxDirectMemory;
+
         logger.debug("-Dio.netty.maxDirectMemory: {} bytes", maxDirectMemory);
 
         int tryAllocateUninitializedArray =
@@ -380,14 +375,6 @@ public static <K, V> ConcurrentMap<K, V> newConcurrentHashMap(Map<? extends K, ?
         return new ConcurrentHashMap<K, V>(map);
     }
 
-    /**
-     * Try to deallocate the specified direct {@link ByteBuffer}. Please note this method does nothing if
-     * the current platform does not support this operation or the specified buffer is not a direct buffer.
-     */
-    public static void freeDirectBuffer(ByteBuffer buffer) {
-        CLEANER.freeDirectBuffer(buffer);
-    }
-
     public static long directBufferAddress(ByteBuffer buffer) {
         return PlatformDependent0.directBufferAddress(buffer);
     }
@@ -584,6 +571,45 @@ public static void setMemory(long address, long bytes, byte value) {
         PlatformDependent0.setMemory(address, bytes, value);
     }
 
+    public static ByteBuffer allocateDirect(int capacity) {
+        return USE_DIRECT_BUFFER_NO_CLEANER
+                ? PlatformDependent.allocateDirectNoCleaner(capacity)
+                : allocateDirectJVM(capacity);
+    }
+
+    public static void freeDirectBuffer(ByteBuffer buffer) {
+        if (USE_DIRECT_BUFFER_NO_CLEANER) {
+            PlatformDependent.freeDirectNoCleaner(buffer);
+        } else {
+            PlatformDependent.freeDirectBufferJVM(buffer);
+        }
+    }
+
+    /**
+     * Allocate a new {@link ByteBuffer} with the given {@code capacity}. {@link ByteBuffer}s allocated with
+     * this method <strong>MUST</strong> be deallocated via {@link #freeDirectNoCleaner(ByteBuffer)}.
+     */
+    public static ByteBuffer allocateDirectJVM(int capacity) {
+        incrementMemoryCounter(capacity);
+        try {
+            return ByteBuffer.allocateDirect(capacity);
+        } catch (Throwable e) {
+            decrementMemoryCounter(capacity);
+            throwException(e);
+            return null;
+        }
+    }
+
+    /**
+     * Try to deallocate the specified direct {@link ByteBuffer}. Please note this method does nothing if
+     * the current platform does not support this operation or the specified buffer is not a direct buffer.
+     */
+    public static void freeDirectBufferJVM(ByteBuffer buffer) {
+        int capacity = buffer.capacity();
+        CLEANER.freeDirectBuffer(buffer);
+        decrementMemoryCounter(capacity);
+    }
+
     /**
      * Allocate a new {@link ByteBuffer} with the given {@code capacity}. {@link ByteBuffer}s allocated with
      * this method <strong>MUST</strong> be deallocated via {@link #freeDirectNoCleaner(ByteBuffer)}.
@@ -632,32 +658,32 @@ public static void freeDirectNoCleaner(ByteBuffer buffer) {
     }
 
     private static void incrementMemoryCounter(int capacity) {
-        if (DIRECT_MEMORY_COUNTER != null) {
-            for (;;) {
-                long usedMemory = DIRECT_MEMORY_COUNTER.get();
-                long newUsedMemory = usedMemory + capacity;
-                if (newUsedMemory > DIRECT_MEMORY_LIMIT) {
-                    throw new OutOfDirectMemoryError("failed to allocate " + capacity
-                            + " byte(s) of direct memory (used: " + usedMemory + ", max: " + DIRECT_MEMORY_LIMIT + ')');
-                }
-                if (DIRECT_MEMORY_COUNTER.compareAndSet(usedMemory, newUsedMemory)) {
-                    break;
-                }
+        for (;;) {
+            long usedMemory = DIRECT_MEMORY_COUNTER.get();
+            long newUsedMemory = usedMemory + capacity;
+            if (DIRECT_MEMORY_LIMIT > 0 && newUsedMemory > DIRECT_MEMORY_LIMIT) {
+                throw new OutOfDirectMemoryError("failed to allocate " + capacity
+                        + " byte(s) of direct memory (used: " + usedMemory + ", max: " + DIRECT_MEMORY_LIMIT + ')');
+            }
+            if (DIRECT_MEMORY_COUNTER.compareAndSet(usedMemory, newUsedMemory)) {
+                break;
             }
         }
     }
 
     private static void decrementMemoryCounter(int capacity) {
-        if (DIRECT_MEMORY_COUNTER != null) {
-            long usedMemory = DIRECT_MEMORY_COUNTER.addAndGet(-capacity);
-            assert usedMemory >= 0;
-        }
+        long usedMemory = DIRECT_MEMORY_COUNTER.addAndGet(-capacity);
+        assert usedMemory >= 0 : "Used memory is negative: " + usedMemory;
     }
 
     public static boolean useDirectBufferNoCleaner() {
         return USE_DIRECT_BUFFER_NO_CLEANER;
     }
 
+    public static long usedDirectMemory() {
+        return DIRECT_MEMORY_COUNTER.get();
+    }
+
     /**
      * Compare two {@code byte} arrays for equality. For performance reasons no bounds checking on the
      * parameters is performed.

microbench/src/main/java/io/netty/microbench/handler/ssl/AbstractSslEngineBenchmark.java

@@ -109,12 +109,12 @@ ByteBuffer newBuffer(int size) {
         DIRECT {
             @Override
             ByteBuffer newBuffer(int size) {
-                return ByteBuffer.allocateDirect(size);
+                return PlatformDependent.allocateDirectJVM(size);
             }
 
             @Override
             void freeBuffer(ByteBuffer buffer) {
-                PlatformDependent.freeDirectBuffer(buffer);
+                PlatformDependent.freeDirectBufferJVM(buffer);
             }
         };