io.undertow.server.handlers.cache.LimitedBufferSlicePool Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of undertow-core Show documentation
There is a newer version: 2.3.18.Final
/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * Licensed 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
 *
 *     http://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.undertow.server.handlers.cache;

import static org.wildfly.common.Assert.checkNotNullParam;

import org.xnio.BufferAllocator;

import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;

/**
 * A limited buffer pooled allocator.  This pool uses a series of buffer regions to back the
 * returned pooled buffers.  When the buffer is no longer needed, it should be freed back into the pool; failure
 * to do so will cause the corresponding buffer area to be unavailable until the buffer is garbage-collected.
 *
 * @author David M. Lloyd
 * @author Jason T. Greene
 */
public final class LimitedBufferSlicePool {

    private static final AtomicIntegerFieldUpdater regionUpdater = AtomicIntegerFieldUpdater.newUpdater(LimitedBufferSlicePool.class, "regionsUsed");
    private final Queue sliceQueue = new ConcurrentLinkedQueue<>();
    private final BufferAllocator allocator;
    private final int bufferSize;
    private final int buffersPerRegion;
    private final int maxRegions;
    private volatile int regionsUsed;


    /**
     * Construct a new instance.
     *
     * @param allocator the buffer allocator to use
     * @param bufferSize the size of each buffer
     * @param maxRegionSize the maximum region size for each backing buffer
     * @param maxRegions the maximum regions to create, zero for unlimited
     */
    public LimitedBufferSlicePool(final BufferAllocator allocator, final int bufferSize, final int maxRegionSize, final int maxRegions) {
        if (bufferSize <= 0) {
            throw new IllegalArgumentException("Buffer size must be greater than zero");
        }
        if (maxRegionSize < bufferSize) {
            throw new IllegalArgumentException("Maximum region size must be greater than or equal to the buffer size");
        }
        buffersPerRegion = maxRegionSize / bufferSize;
        this.bufferSize = bufferSize;
        this.allocator = checkNotNullParam("allocator", allocator);
        this.maxRegions = maxRegions;
    }

     /**
     * Construct a new instance.
     *
     * @param allocator the buffer allocator to use
     * @param bufferSize the size of each buffer
     * @param maxRegionSize the maximum region size for each backing buffer
     */
    public LimitedBufferSlicePool(BufferAllocator allocator, int bufferSize, int maxRegionSize) {
        this(allocator, bufferSize, maxRegionSize, 0);
    }


    /**
     * Construct a new instance, using a direct buffer allocator.
     *
     * @param bufferSize the size of each buffer
     * @param maxRegionSize the maximum region size for each backing buffer
     */
    public LimitedBufferSlicePool(final int bufferSize, final int maxRegionSize) {
        this(BufferAllocator.DIRECT_BYTE_BUFFER_ALLOCATOR, bufferSize, maxRegionSize);
    }

    /**
     * Allocates a new byte buffer if possible
     *
     * @return new buffer or null if none available
     **/
    public PooledByteBuffer allocate() {
        final Queue sliceQueue = this.sliceQueue;
        final Slice slice = sliceQueue.poll();
        if (slice == null && (maxRegions <= 0 || regionUpdater.getAndIncrement(this) < maxRegions)) {
            final int bufferSize = this.bufferSize;
            final int buffersPerRegion = this.buffersPerRegion;
            final ByteBuffer region = allocator.allocate(buffersPerRegion * bufferSize);
            int idx = bufferSize;
            for (int i = 1; i < buffersPerRegion; i ++) {
                sliceQueue.add(new Slice(region, idx, bufferSize));
                idx += bufferSize;
            }
            final Slice newSlice = new Slice(region, 0, bufferSize);
            return new PooledByteBuffer(newSlice, newSlice.slice(), sliceQueue);
        }
        if (slice == null) {
            return null;
        }
        return new PooledByteBuffer(slice, slice.slice(), sliceQueue);
    }

    public boolean canAllocate(int slices) {
        if (regionsUsed < maxRegions)
            return true;

        if (sliceQueue.isEmpty())
            return false;

        Iterator iterator = sliceQueue.iterator();
        for (int i = 0; i < slices; i++) {
            if (! iterator.hasNext()) {
                return false;
            }
            try {
                iterator.next();
            } catch (NoSuchElementException e) {
                return false;
            }
        }

        return true;
    }

    public static final class PooledByteBuffer {
        private final Slice region;
        private final Queue slices;
        volatile ByteBuffer buffer;

        private static final AtomicReferenceFieldUpdater bufferUpdater = AtomicReferenceFieldUpdater.newUpdater(PooledByteBuffer.class, ByteBuffer.class, "buffer");

        private PooledByteBuffer(final Slice region, final ByteBuffer buffer, final Queue slices) {
            this.region = region;
            this.buffer = buffer;
            this.slices = slices;
        }

        public void free() {
            if (bufferUpdater.getAndSet(this, null) != null) {
                // trust the user, repool the buffer
                slices.add(region);
            }
        }

        public ByteBuffer getBuffer() {
            final ByteBuffer buffer = this.buffer;
            if (buffer == null) {
                throw new IllegalStateException();
            }
            return buffer;
        }

        public String toString() {
            return "Pooled buffer " + buffer;
        }
    }

    private static final class Slice {
        private final ByteBuffer parent;
        private final int start;
        private final int size;

        private Slice(final ByteBuffer parent, final int start, final int size) {
            this.parent = parent;
            this.start = start;
            this.size = size;
        }

        ByteBuffer slice() {
            return ((ByteBuffer)parent.duplicate().position(start).limit(start+size)).slice();
        }
    }
}