Skip to content

Commit 120b63f

Browse files
author
lao
committed
Make pooled allocator support non-threadlocal use case, and integrated with mimalloc allocator
1 parent cb56eee commit 120b63f

9 files changed

Lines changed: 268 additions & 75 deletions

‎buffer/src/main/java/io/netty/buffer/MiByteBufAllocator.java‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package io.netty.buffer;
1717

18+
import io.netty.util.concurrent.FastThreadLocalThread;
1819
import io.netty.util.internal.PlatformDependent;
1920
import io.netty.util.internal.UnstableApi;
2021

@@ -34,6 +35,7 @@ public final class MiByteBufAllocator extends AbstractByteBufAllocator
3435

3536
private final MiMallocByteBufAllocator direct;
3637
private final MiMallocByteBufAllocator heap;
38+
private final PooledByteBufAllocator pooledAlloc;
3739

3840
public MiByteBufAllocator() {
3941
this(!PlatformDependent.isExplicitNoPreferDirect());
@@ -43,16 +45,25 @@ public MiByteBufAllocator(boolean preferDirect) {
4345
super(preferDirect);
4446
direct = new MiMallocByteBufAllocator(new DirectChunkAllocator(this));
4547
heap = new MiMallocByteBufAllocator(new HeapChunkAllocator(this));
48+
pooledAlloc = PooledByteBufAllocator.getNonThreadLocalAllocator();
4649
}
4750

4851
@Override
4952
protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) {
50-
return toLeakAwareBuffer(heap.allocate(initialCapacity, maxCapacity));
53+
if (FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
54+
return toLeakAwareBuffer(heap.allocate(initialCapacity, maxCapacity));
55+
} else {
56+
return pooledAlloc.heapBuffer(initialCapacity, maxCapacity);
57+
}
5158
}
5259

5360
@Override
5461
protected ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity) {
55-
return toLeakAwareBuffer(direct.allocate(initialCapacity, maxCapacity));
62+
if (FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
63+
return toLeakAwareBuffer(direct.allocate(initialCapacity, maxCapacity));
64+
} else {
65+
return pooledAlloc.directBuffer(initialCapacity, maxCapacity);
66+
}
5667
}
5768

5869
@Override
@@ -62,12 +73,12 @@ public boolean isDirectBufferPooled() {
6273

6374
@Override
6475
public long usedHeapMemory() {
65-
return heap.usedMemory();
76+
return heap.usedMemory() + pooledAlloc.usedHeapMemory();
6677
}
6778

6879
@Override
6980
public long usedDirectMemory() {
70-
return direct.usedMemory();
81+
return direct.usedMemory() + pooledAlloc.usedDirectMemory();
7182
}
7283

7384
long abandonedHeapSegmentCount() {

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

Lines changed: 81 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,20 @@
1717
package io.netty.buffer;
1818

1919
import io.netty.util.internal.CleanableDirectBuffer;
20+
import io.netty.util.internal.ObjectPool;
2021
import io.netty.util.internal.PlatformDependent;
2122
import io.netty.util.internal.StringUtil;
22-
2323
import java.nio.ByteBuffer;
2424
import java.util.ArrayList;
2525
import java.util.Collections;
2626
import java.util.List;
27+
import java.util.Queue;
2728
import java.util.concurrent.atomic.AtomicInteger;
2829
import java.util.concurrent.atomic.AtomicReference;
2930
import java.util.concurrent.atomic.LongAdder;
3031
import java.util.concurrent.locks.ReentrantLock;
31-
3232
import static io.netty.buffer.PoolChunk.isSubpage;
33+
import static io.netty.buffer.PooledByteBufAllocator.ARENA_BUFFER_QUEUE_CAPACITY_FOR_NON_THREAD_LOCAL;
3334
import static java.lang.Math.max;
3435

3536
abstract class PoolArena<T> implements PoolArenaMetric {
@@ -127,7 +128,7 @@ private PoolSubpage<T>[] newSubpagePoolArray(int size) {
127128
abstract boolean isDirect();
128129

129130
PooledByteBuf<T> allocate(PoolThreadCache cache, int reqCapacity, int maxCapacity) {
130-
PooledByteBuf<T> buf = newByteBuf(maxCapacity);
131+
PooledByteBuf<T> buf = newByteBuf(maxCapacity, cache);
131132
allocate(cache, buf, reqCapacity);
132133
return buf;
133134
}
@@ -150,7 +151,7 @@ private void allocate(PoolThreadCache cache, PooledByteBuf<T> buf, final int req
150151
private void tcacheAllocateSmall(PoolThreadCache cache, PooledByteBuf<T> buf, final int reqCapacity,
151152
final int sizeIdx) {
152153

153-
if (cache.allocateSmall(this, buf, reqCapacity, sizeIdx)) {
154+
if (cache.useThreadLocal() && cache.allocateSmall(this, buf, reqCapacity, sizeIdx)) {
154155
// was able to allocate out of the cache so move on
155156
return;
156157
}
@@ -190,7 +191,7 @@ private void tcacheAllocateSmall(PoolThreadCache cache, PooledByteBuf<T> buf, fi
190191

191192
private void tcacheAllocateNormal(PoolThreadCache cache, PooledByteBuf<T> buf, final int reqCapacity,
192193
final int sizeIdx) {
193-
if (cache.allocateNormal(this, buf, reqCapacity, sizeIdx)) {
194+
if (cache.useThreadLocal() && cache.allocateNormal(this, buf, reqCapacity, sizeIdx)) {
194195
// was able to allocate out of the cache so move on
195196
return;
196197
}
@@ -243,7 +244,8 @@ void free(PoolChunk<T> chunk, ByteBuffer nioBuffer, long handle, int normCapacit
243244
deallocationsHuge.increment();
244245
} else {
245246
SizeClass sizeClass = sizeClass(handle);
246-
if (cache != null && cache.add(this, chunk, nioBuffer, handle, normCapacity, sizeClass)) {
247+
if (cache != null && cache.useThreadLocal() &&
248+
cache.add(this, chunk, nioBuffer, handle, normCapacity, sizeClass)) {
247249
// cached so not free it.
248250
return;
249251
}
@@ -574,7 +576,7 @@ public long numPinnedBytes() {
574576

575577
protected abstract PoolChunk<T> newChunk(int pageSize, int maxPageIdx, int pageShifts, int chunkSize);
576578
protected abstract PoolChunk<T> newUnpooledChunk(int capacity);
577-
protected abstract PooledByteBuf<T> newByteBuf(int maxCapacity);
579+
protected abstract PooledByteBuf<T> newByteBuf(int maxCapacity, PoolThreadCache cache);
578580
protected abstract void memoryCopy(T src, int srcOffset, PooledByteBuf<T> dst, int length);
579581
protected abstract void destroyChunk(PoolChunk<T> chunk);
580582

@@ -661,10 +663,34 @@ private void destroyPoolChunkLists(PoolChunkList<T>... chunkLists) {
661663

662664
static final class HeapArena extends PoolArena<byte[]> {
663665
private final AtomicReference<PoolChunk<byte[]>> lastDestroyedChunk;
666+
private final Queue<PooledByteBuf<byte[]>> bufferQueue;
667+
private final ObjectPool.Handle<PooledHeapByteBuf> handle;
664668

665669
HeapArena(PooledByteBufAllocator parent, SizeClasses sizeClass) {
670+
this(parent, sizeClass, true);
671+
}
672+
673+
HeapArena(PooledByteBufAllocator parent, SizeClasses sizeClass, boolean useThreadLocal) {
666674
super(parent, sizeClass);
667675
lastDestroyedChunk = new AtomicReference<>();
676+
if (useThreadLocal) {
677+
bufferQueue = null;
678+
handle = new ObjectPool.Handle<PooledHeapByteBuf>() {
679+
@Override
680+
public void recycle(PooledHeapByteBuf self) {
681+
// noop
682+
}
683+
};
684+
} else {
685+
int qSize = Math.max(ARENA_BUFFER_QUEUE_CAPACITY_FOR_NON_THREAD_LOCAL, 2);
686+
bufferQueue = PlatformDependent.newFixedMpmcQueue(qSize);
687+
handle = new ObjectPool.Handle<PooledHeapByteBuf>() {
688+
@Override
689+
public void recycle(PooledHeapByteBuf self) {
690+
bufferQueue.offer(self);
691+
}
692+
};
693+
}
668694
}
669695

670696
private static byte[] newByteArray(int size) {
@@ -705,9 +731,18 @@ protected void destroyChunk(PoolChunk<byte[]> chunk) {
705731
}
706732

707733
@Override
708-
protected PooledByteBuf<byte[]> newByteBuf(int maxCapacity) {
709-
return HAS_UNSAFE ? PooledUnsafeHeapByteBuf.newUnsafeInstance(maxCapacity)
710-
: PooledHeapByteBuf.newInstance(maxCapacity);
734+
protected PooledByteBuf<byte[]> newByteBuf(int maxCapacity, PoolThreadCache threadCache) {
735+
if (threadCache.useThreadLocal()) {
736+
return HAS_UNSAFE ? PooledUnsafeHeapByteBuf.newUnsafeInstance(maxCapacity)
737+
: PooledHeapByteBuf.newInstance(maxCapacity);
738+
}
739+
PooledByteBuf<byte[]> buf = bufferQueue == null ? null : bufferQueue.poll();
740+
if (buf == null) {
741+
buf = HAS_UNSAFE ? PooledUnsafeHeapByteBuf.newInstanceNoThreadLocal(handle)
742+
: PooledHeapByteBuf.newInstanceNoThreadLocal(handle);
743+
}
744+
buf.reuse(maxCapacity);
745+
return buf;
711746
}
712747

713748
@Override
@@ -721,9 +756,33 @@ protected void memoryCopy(byte[] src, int srcOffset, PooledByteBuf<byte[]> dst,
721756
}
722757

723758
static final class DirectArena extends PoolArena<ByteBuffer> {
759+
private final Queue<PooledByteBuf<ByteBuffer>> bufferQueue;
760+
private final ObjectPool.Handle<PooledByteBuf<ByteBuffer>> handle;
724761

725762
DirectArena(PooledByteBufAllocator parent, SizeClasses sizeClass) {
763+
this(parent, sizeClass, true);
764+
}
765+
766+
DirectArena(PooledByteBufAllocator parent, SizeClasses sizeClass, boolean useThreadLocal) {
726767
super(parent, sizeClass);
768+
if (useThreadLocal) {
769+
bufferQueue = null;
770+
handle = new ObjectPool.Handle<PooledByteBuf<ByteBuffer>>() {
771+
@Override
772+
public void recycle(PooledByteBuf<ByteBuffer> self) {
773+
// noop
774+
}
775+
};
776+
} else {
777+
int qSize = Math.max(ARENA_BUFFER_QUEUE_CAPACITY_FOR_NON_THREAD_LOCAL, 2);
778+
bufferQueue = PlatformDependent.newFixedMpmcQueue(qSize);
779+
handle = new ObjectPool.Handle<PooledByteBuf<ByteBuffer>>() {
780+
@Override
781+
public void recycle(PooledByteBuf<ByteBuffer> self) {
782+
bufferQueue.offer(self);
783+
}
784+
};
785+
}
727786
}
728787

729788
@Override
@@ -774,12 +833,19 @@ protected void destroyChunk(PoolChunk<ByteBuffer> chunk) {
774833
}
775834

776835
@Override
777-
protected PooledByteBuf<ByteBuffer> newByteBuf(int maxCapacity) {
778-
if (HAS_UNSAFE) {
779-
return PooledUnsafeDirectByteBuf.newInstance(maxCapacity);
780-
} else {
781-
return PooledDirectByteBuf.newInstance(maxCapacity);
836+
protected PooledByteBuf<ByteBuffer> newByteBuf(int maxCapacity, PoolThreadCache threadCache) {
837+
if (threadCache.useThreadLocal()) {
838+
return HAS_UNSAFE ? PooledUnsafeDirectByteBuf.newInstance(maxCapacity)
839+
: PooledDirectByteBuf.newInstance(maxCapacity);
840+
}
841+
PooledByteBuf<ByteBuffer> buf = bufferQueue == null ? null : bufferQueue.poll();
842+
if (buf == null) {
843+
buf = HAS_UNSAFE ?
844+
PooledUnsafeDirectByteBuf.newInstanceNoThreadLocal(handle)
845+
: PooledDirectByteBuf.newInstanceNoThreadLocal(handle);
782846
}
847+
buf.reuse(maxCapacity);
848+
return buf;
783849
}
784850

785851
@Override

‎buffer/src/main/java/io/netty/buffer/PoolThreadCache.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,10 @@ final class PoolThreadCache {
6161
private final FreeOnFinalize freeOnFinalize;
6262

6363
private int allocations;
64+
// A special `PoolThreadCache` that is not stored in thread-local.
65+
static final PoolThreadCache THREAD_CACHE_WITHOUT_THREAD_LOCAL = new PoolThreadCache();
66+
final int smallCacheSize;
67+
final int normalCacheSize;
6468

6569
// TODO: Test if adding padding helps under contention
6670
//private long pad0, pad1, pad2, pad3, pad4, pad5, pad6, pad7;
@@ -72,6 +76,8 @@ final class PoolThreadCache {
7276
this.freeSweepAllocationThreshold = freeSweepAllocationThreshold;
7377
this.heapArena = heapArena;
7478
this.directArena = directArena;
79+
this.smallCacheSize = smallCacheSize;
80+
this.normalCacheSize = normalCacheSize;
7581
if (directArena != null) {
7682
smallSubPageDirectCaches = createSubPageCaches(smallCacheSize, directArena.sizeClass.nSubpages);
7783
normalDirectCaches = createNormalCaches(normalCacheSize, maxCachedBufferCapacity, directArena);
@@ -102,6 +108,20 @@ final class PoolThreadCache {
102108
freeOnFinalize = useFinalizer ? new FreeOnFinalize(this) : null;
103109
}
104110

111+
/** Creates a special `PoolThreadCache` that is not stored in thread-local. */
112+
PoolThreadCache() {
113+
this.freeSweepAllocationThreshold = 0;
114+
this.heapArena = null;
115+
this.directArena = null;
116+
this.smallSubPageDirectCaches = null;
117+
this.normalDirectCaches = null;
118+
this.smallSubPageHeapCaches = null;
119+
this.normalHeapCaches = null;
120+
this.freeOnFinalize = null;
121+
this.smallCacheSize = 0;
122+
this.normalCacheSize = 0;
123+
}
124+
105125
private static <T> MemoryRegionCache<T>[] createSubPageCaches(
106126
int cacheSize, int numCaches) {
107127
if (cacheSize > 0 && numCaches > 0) {
@@ -168,6 +188,10 @@ private boolean allocate(MemoryRegionCache<?> cache, PooledByteBuf buf, int reqC
168188
return allocated;
169189
}
170190

191+
boolean useThreadLocal() {
192+
return this != THREAD_CACHE_WITHOUT_THREAD_LOCAL;
193+
}
194+
171195
/**
172196
* Add {@link PoolChunk} and {@code handle} to the cache if there is enough room.
173197
* Returns {@code true} if it fit into the cache {@code false} otherwise.

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import io.netty.util.Recycler.EnhancedHandle;
2020
import io.netty.util.internal.ObjectPool.Handle;
21+
import io.netty.util.internal.ObjectUtil;
2122

2223
import java.io.IOException;
2324
import java.nio.ByteBuffer;
@@ -29,7 +30,7 @@
2930

3031
abstract class PooledByteBuf<T> extends AbstractReferenceCountedByteBuf {
3132

32-
private final EnhancedHandle<PooledByteBuf<T>> recyclerHandle;
33+
private final Handle<PooledByteBuf<T>> recyclerHandle;
3334

3435
protected PoolChunk<T> chunk;
3536
protected long handle;
@@ -44,7 +45,7 @@ abstract class PooledByteBuf<T> extends AbstractReferenceCountedByteBuf {
4445
@SuppressWarnings("unchecked")
4546
protected PooledByteBuf(Handle<? extends PooledByteBuf<T>> recyclerHandle, int maxCapacity) {
4647
super(maxCapacity);
47-
this.recyclerHandle = (EnhancedHandle<PooledByteBuf<T>>) recyclerHandle;
48+
this.recyclerHandle = ObjectUtil.checkNotNull((Handle<PooledByteBuf<T>>) recyclerHandle, "recyclerHandle");
4849
}
4950

5051
void init(PoolChunk<T> chunk, ByteBuffer nioBuffer,
@@ -181,7 +182,12 @@ protected final void deallocate() {
181182
tmpNioBuf = null;
182183
chunk = null;
183184
cache = null;
184-
this.recyclerHandle.unguardedRecycle(this);
185+
if (recyclerHandle instanceof EnhancedHandle) {
186+
EnhancedHandle<PooledByteBuf<T>> enhancedHandle = (EnhancedHandle<PooledByteBuf<T>>) recyclerHandle;
187+
enhancedHandle.unguardedRecycle(this);
188+
} else {
189+
recyclerHandle.recycle(this);
190+
}
185191
}
186192
}
187193

0 commit comments

Comments
 (0)