/*
 * Copyright 2012 The Netty Project
 *
 * The Netty Project licenses this file to you 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 org.jboss.netty.handler.ssl;

import javax.net.ssl.SSLEngine;
import java.nio.ByteBuffer;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * A {@link ByteBuffer} pool dedicated for {@link SslHandler} performance improvement.
 * <p>
 * In most cases, you won't need to create a new pool instance because {@link SslHandler} has a default pool
 * instance internally.
 * <p>
 * The reason why {@link SslHandler} requires a buffer pool is because the current {@link SSLEngine} implementation
 * always requires a 17KiB buffer for every 'wrap' and 'unwrap' operation.  In most cases, the actual size of the
 * required buffer is much smaller than that, and therefore allocating a 17KiB buffer for every 'wrap' and 'unwrap'
 * operation wastes a lot of memory bandwidth, resulting in the application performance degradation.
 */
public class SslBufferPool {

    // Add 1024 as a room for compressed data and another 1024 for Apache Harmony compatibility.
    private static final int MAX_PACKET_SIZE_ALIGNED = (OpenSslEngine.MAX_ENCRYPTED_PACKET_LENGTH / 128 + 1) * 128;

    private static final int DEFAULT_POOL_SIZE = MAX_PACKET_SIZE_ALIGNED * 1024;

    private final ByteBuffer preallocated;
    private final BlockingQueue<ByteBuffer> pool;
    private final int maxBufferCount;
    private final boolean allocateDirect;

    /**
     * The number of buffers allocated so far. Used only when {@link #preallocated} is null.
     */
    private final AtomicInteger numAllocations;

    /**
     * Creates a new buffer pool whose size is {@code 19267584}, which can hold {@code 1024} buffers.
     */
    public SslBufferPool() {
        this(DEFAULT_POOL_SIZE);
    }

    /**
     * Creates a new buffer pool whose size is {@code 19267584}, which can hold {@code 1024} buffers.
     *
     * @param preallocate {@code true} if and only if the buffers in this pool has to be pre-allocated
     *                    at construction time
     * @param allocateDirect {@code true} if and only if this pool has to allocate direct buffers.
     */
    public SslBufferPool(boolean preallocate, boolean allocateDirect) {
        this(DEFAULT_POOL_SIZE, preallocate, allocateDirect);
    }

    /**
     * Creates a new buffer pool.
     *
     * @param maxPoolSize the maximum number of bytes that this pool can hold
     */
    public SslBufferPool(int maxPoolSize) {
        this(maxPoolSize, false, false);
    }

    /**
     * Creates a new buffer pool.
     *
     * @param maxPoolSize the maximum number of bytes that this pool can hold
     * @param preallocate {@code true} if and only if the buffers in this pool has to be pre-allocated
     *                    at construction time
     * @param allocateDirect {@code true} if and only if this pool has to allocate direct buffers.
     */
    public SslBufferPool(int maxPoolSize, boolean preallocate, boolean allocateDirect) {
        if (maxPoolSize <= 0) {
            throw new IllegalArgumentException("maxPoolSize: " + maxPoolSize);
        }

        int maxBufferCount = maxPoolSize / MAX_PACKET_SIZE_ALIGNED;
        if (maxPoolSize % MAX_PACKET_SIZE_ALIGNED != 0) {
            maxBufferCount ++;
        }

        this.maxBufferCount = maxBufferCount;
        this.allocateDirect = allocateDirect;

        pool = new ArrayBlockingQueue<ByteBuffer>(maxBufferCount);

        if (preallocate) {
            preallocated = allocate(maxBufferCount * MAX_PACKET_SIZE_ALIGNED);
            numAllocations = null;
            for (int i = 0; i < maxBufferCount; i ++) {
                int pos = i * MAX_PACKET_SIZE_ALIGNED;
                preallocated.clear().position(pos).limit(pos + MAX_PACKET_SIZE_ALIGNED);
                pool.add(preallocated.slice());
            }
        } else {
            preallocated = null;
            numAllocations = new AtomicInteger();
        }
    }

    /**
     * Returns the maximum size of this pool in byte unit.  The returned value
     * can be somewhat different from what was specified in the constructor.
     */
    public int getMaxPoolSize() {
        return maxBufferCount * MAX_PACKET_SIZE_ALIGNED;
    }

    /**
     * Returns the number of bytes which were allocated but have not been
     * acquired yet.  You can estimate how optimal the specified maximum pool
     * size is from this value.  If it keeps returning {@code 0}, it means the
     * pool is getting exhausted.  If it keeps returns a unnecessarily big
     * value, it means the pool is wasting the heap space.
     */
    public int getUnacquiredPoolSize() {
        return pool.size() * MAX_PACKET_SIZE_ALIGNED;
    }

    /**
     * Acquire a new {@link ByteBuffer} out of the {@link SslBufferPool}
     *
     */
    public ByteBuffer acquireBuffer() {
        ByteBuffer buf;
        if (preallocated != null || numAllocations.get() >= maxBufferCount) {
            boolean interrupted = false;
            for (;;) {
                try {
                    buf = pool.take();
                    break;
                } catch (InterruptedException ignore) {
                    interrupted = true;
                }
            }
            if (interrupted) {
                Thread.currentThread().interrupt();
            }
        } else {
            buf = pool.poll();
            if (buf == null) {
                // Note that we can allocate more buffers than maxBufferCount.
                // We will discard the buffers allocated after numAllocations reached maxBufferCount in releaseBuffer().
                numAllocations.incrementAndGet();
                buf = allocate(OpenSslEngine.MAX_ENCRYPTED_PACKET_LENGTH);
            }
        }

        buf.clear();
        return buf;
    }

    /**
     * Release a previous acquired {@link ByteBuffer}
     */
    public void releaseBuffer(ByteBuffer buffer) {
        pool.offer(buffer);
    }

    private ByteBuffer allocate(int capacity) {
        if (allocateDirect) {
            return ByteBuffer.allocateDirect(capacity);
        } else {
            return ByteBuffer.allocate(capacity);
        }
    }
}
