Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
package com.destroystokyo.paper.util.pooled;
import net.minecraft.server.MCUtil;
import org.apache.commons.lang3.mutable.MutableInt;
import java.util.ArrayDeque;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;
public final class PooledObjects<E> {
/**
* Wrapper for an object that will be have a cleaner registered for it, and may be automatically returned to pool.
*/
public class AutoReleased {
private final E object;
private final Runnable cleaner;
public AutoReleased(E object, Runnable cleaner) {
this.object = object;
this.cleaner = cleaner;
}
public final E getObject() {
return object;
}
public final Runnable getCleaner() {
return cleaner;
}
}
public static final PooledObjects<MutableInt> POOLED_MUTABLE_INTEGERS = new PooledObjects<>(MutableInt::new, 256);
private static final int BUCKETS = 8;
private final PooledObjectHandler<E> handler;
private final int bucketSize;
@SuppressWarnings("unchecked")
private final ArrayDeque<E>[] pools = new ArrayDeque[BUCKETS];
private final ReentrantLock[] locks = new ReentrantLock[BUCKETS];
private final AtomicLong bucketIdPool = new AtomicLong(0);
public PooledObjects(final PooledObjectHandler<E> handler, int maxPoolSize) {
if (handler == null) {
throw new NullPointerException("Handler must not be null");
}
if (maxPoolSize <= 0) {
throw new IllegalArgumentException("Max pool size must be greater-than 0");
}
int remainder = maxPoolSize % BUCKETS;
if (remainder > 0) {
maxPoolSize = maxPoolSize - remainder + BUCKETS;
}
this.handler = handler;
this.bucketSize = maxPoolSize / BUCKETS;
for (int i = 0; i < BUCKETS; i++) {
this.pools[i] = new ArrayDeque<>(bucketSize / 4);
this.locks[i] = new ReentrantLock();
}
}
public AutoReleased acquireCleaner(Object holder) {
return acquireCleaner(holder, this::release);
}
public AutoReleased acquireCleaner(Object holder, Consumer<E> releaser) {
E resource = acquire();
Runnable cleaner = MCUtil.registerCleaner(holder, resource, releaser);
return new AutoReleased(resource, cleaner);
}
public long size() {
long size = 0;
for (int i = 0; i < BUCKETS; i++) {
size += this.pools[i].size();
}
return size;
}
public E acquire() {
for (int base = (int) (this.bucketIdPool.getAndIncrement() % BUCKETS), i = 0; i < BUCKETS; i++ ) {
int bucketId = (base + i) % BUCKETS;
this.locks[bucketId].lock();
E value = this.pools[bucketId].poll();
this.locks[bucketId].unlock();
if (value != null) {
this.handler.onAcquire(value);
return value;
}
}
return this.handler.createNew();
}
public void release(final E value) {
int attempts = 3; // cap on contention
do {
// find least filled pool before locking
int smallest = -1;
for (int i = 0; i < BUCKETS; i++ ) {
ArrayDeque<E> pool = this.pools[i];
int size = pool.size();
if (size < this.bucketSize && (smallest == -1 || size < smallest || (size == smallest && ThreadLocalRandom.current().nextBoolean()))) {
smallest = i;
}
}
if (smallest == -1) return; // Can not find a pool to fill
ReentrantLock smallestLock = this.locks[smallest];
smallestLock.lock();
ArrayDeque<E> pool = this.pools[smallest];
if (pool.size() < this.bucketSize) {
this.handler.onRelease(value);
pool.push(value);
smallestLock.unlock();
return;
} else {
smallestLock.unlock();
}
} while (attempts-- > 0);
}
/** This object is restricted from interacting with any pool */
public interface PooledObjectHandler<E> {
/**
* Must return a non-null object
*/
E createNew();
default void onAcquire(final E value) {}
default void onRelease(final E value) {}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.