PooledObjects.java

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) {}
    }
}