/*
 * Copyright 2021 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:
 *
 *   https://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.netty.buffer.api.pool;

import io.netty.buffer.api.AllocatorControl.UntetheredMemory;
import io.netty.buffer.api.pool.PoolArena.SizeClass;
import io.netty.util.internal.MathUtil;
import io.netty.util.internal.ObjectPool;
import io.netty.util.internal.ObjectPool.Handle;
import io.netty.util.internal.PlatformDependent;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Queue;

import static io.netty.buffer.api.pool.PoolArena.SizeClass.Normal;
import static io.netty.buffer.api.pool.PoolArena.SizeClass.Small;
import static io.netty.util.internal.ObjectUtil.checkPositiveOrZero;

/**
 * Acts a Thread cache for allocations. This implementation is modelled after
 * <a href="https://people.freebsd.org/~jasone/jemalloc/bsdcan2006/jemalloc.pdf">jemalloc</a> and the described
 * techniques of
 * <a href="https://www.facebook.com/notes/facebook-engineering/scalable-memory-allocation-using-jemalloc/480222803919">
 * Scalable memory allocation using jemalloc</a>.
 */
final class PoolThreadCache {

    private static final InternalLogger logger = InternalLoggerFactory.getInstance(PoolThreadCache.class);
    private static final int INTEGER_SIZE_MINUS_ONE = Integer.SIZE - 1;

    final PoolArena arena;

    // Hold the caches for the different size classes, which are tiny, small and normal.
    private final MemoryRegionCache[] smallSubPageCaches;
    private final MemoryRegionCache[] normalCaches;

    private final int freeSweepAllocationThreshold;

    private int allocations;

    PoolThreadCache(PoolArena arena,
                    int smallCacheSize, int normalCacheSize, int maxCachedBufferCapacity,
                    int freeSweepAllocationThreshold) {
        checkPositiveOrZero(maxCachedBufferCapacity, "maxCachedBufferCapacity");
        this.freeSweepAllocationThreshold = freeSweepAllocationThreshold;
        this.arena = arena;
        if (arena != null) {
            // Create the caches for the heap allocations
            smallSubPageCaches = createSubPageCaches(
                    smallCacheSize, arena.numSmallSubpagePools);

            normalCaches = createNormalCaches(
                    normalCacheSize, maxCachedBufferCapacity, arena);

            arena.numThreadCaches.getAndIncrement();
        } else {
            // No heapArea is configured so just null out all caches
            smallSubPageCaches = null;
            normalCaches = null;
        }

        // Only check if there are caches in use.
        if ((smallSubPageCaches != null || normalCaches != null)
                && freeSweepAllocationThreshold < 1) {
            throw new IllegalArgumentException("freeSweepAllocationThreshold: "
                    + freeSweepAllocationThreshold + " (expected: > 0)");
        }
    }

    private static MemoryRegionCache[] createSubPageCaches(
            int cacheSize, int numCaches) {
        if (cacheSize > 0 && numCaches > 0) {
            MemoryRegionCache[] cache = new MemoryRegionCache[numCaches];
            for (int i = 0; i < cache.length; i++) {
                // TODO: maybe use cacheSize / cache.length
                cache[i] = new SubPageMemoryRegionCache(cacheSize);
            }
            return cache;
        } else {
            return null;
        }
    }

    private static MemoryRegionCache[] createNormalCaches(
            int cacheSize, int maxCachedBufferCapacity, PoolArena area) {
        if (cacheSize > 0 && maxCachedBufferCapacity > 0) {
            int max = Math.min(area.chunkSize, maxCachedBufferCapacity);

            // Create as many normal caches as we support based on how many sizeIdx we have and what the upper
            // bound is that we want to cache in general.
            List<MemoryRegionCache> cache = new ArrayList<>() ;
            for (int idx = area.numSmallSubpagePools; idx < area.nSizes && area.sizeIdx2size(idx) <= max ; idx++) {
                cache.add(new NormalMemoryRegionCache(cacheSize));
            }
            return cache.toArray(MemoryRegionCache[]::new);
        } else {
            return null;
        }
    }

    // val > 0
    static int log2(int val) {
        return INTEGER_SIZE_MINUS_ONE - Integer.numberOfLeadingZeros(val);
    }

    /**
     * Try to allocate a small buffer out of the cache. Returns {@code true} if successful {@code false} otherwise
     */
    UntetheredMemory allocateSmall(PooledAllocatorControl control, int size, int sizeIdx) {
        return allocate(cacheForSmall(sizeIdx), control, size);
    }

    /**
     * Try to allocate a normal buffer out of the cache. Returns {@code true} if successful {@code false} otherwise
     */
    UntetheredMemory allocateNormal(PoolArena area, PooledAllocatorControl control, int size, int sizeIdx) {
        return allocate(cacheForNormal(area, sizeIdx), control, size);
    }

    private UntetheredMemory allocate(MemoryRegionCache cache, PooledAllocatorControl control, int size) {
        if (cache == null) {
            // no cache found so just return false here
            return null;
        }
        UntetheredMemory allocated = cache.allocate(size, this, control);
        if (++allocations >= freeSweepAllocationThreshold) {
            allocations = 0;
            trim();
        }
        return allocated;
    }

    /**
     * Add {@link PoolChunk} and {@code handle} to the cache if there is enough room.
     * Returns {@code true} if it fit into the cache {@code false} otherwise.
     */
    boolean add(PoolArena area, PoolChunk chunk,
                long handle, int normCapacity, SizeClass sizeClass) {
        int sizeIdx = area.size2SizeIdx(normCapacity);
        MemoryRegionCache cache = cache(area, sizeIdx, sizeClass);
        if (cache == null) {
            return false;
        }
        return cache.add(chunk, handle, normCapacity);
    }

    private MemoryRegionCache cache(PoolArena area, int sizeIdx, SizeClass sizeClass) {
        if (sizeClass == Normal) {
            return cacheForNormal(area, sizeIdx);
        }
        if (sizeClass == Small) {
            return cacheForSmall(sizeIdx);
        }
        throw new AssertionError("Unexpected size class: " + sizeClass);
    }

    /**
     *  Should be called if the Thread that uses this cache is about to exist to release resources out of the cache
     */
    void free() {
        int numFreed = free(smallSubPageCaches) + free(normalCaches);

        if (numFreed > 0 && logger.isDebugEnabled()) {
            logger.debug("Freed {} thread-local buffer(s) from thread: {}", numFreed,
                    Thread.currentThread().getName());
        }

        if (arena != null) {
            arena.numThreadCaches.getAndDecrement();
        }
    }

    private static int free(MemoryRegionCache[] caches) {
        if (caches == null) {
            return 0;
        }

        int numFreed = 0;
        for (MemoryRegionCache c: caches) {
            numFreed += free(c);
        }
        return numFreed;
    }

    private static int free(MemoryRegionCache cache) {
        if (cache == null) {
            return 0;
        }
        return cache.free();
    }

    void trim() {
        trim(smallSubPageCaches);
        trim(normalCaches);
    }

    private static void trim(MemoryRegionCache[] caches) {
        if (caches == null) {
            return;
        }
        for (MemoryRegionCache c: caches) {
            trim(c);
        }
    }

    private static void trim(MemoryRegionCache cache) {
        if (cache == null) {
            return;
        }
        cache.trim();
    }

    private MemoryRegionCache cacheForSmall(int sizeIdx) {
        return cache(smallSubPageCaches, sizeIdx);
    }

    private MemoryRegionCache cacheForNormal(PoolArena area, int sizeIdx) {
        // We need to substract area.numSmallSubpagePools as sizeIdx is the overall index for all sizes.
        int idx = sizeIdx - area.numSmallSubpagePools;
        return cache(normalCaches, idx);
    }

    private static  MemoryRegionCache cache(MemoryRegionCache[] cache, int sizeIdx) {
        if (cache == null || sizeIdx > cache.length - 1) {
            return null;
        }
        return cache[sizeIdx];
    }

    /**
     * Cache used for buffers which are backed by SMALL size.
     */
    private static final class SubPageMemoryRegionCache extends MemoryRegionCache {
        SubPageMemoryRegionCache(int size) {
            super(size, Small);
        }

        @Override
        protected UntetheredMemory allocBuf(PoolChunk chunk, long handle, int size, PoolThreadCache threadCache,
                                  PooledAllocatorControl control) {
            return chunk.allocateBufferWithSubpage(handle, size, threadCache, control);
        }
    }

    /**
     * Cache used for buffers which are backed by NORMAL size.
     */
    private static final class NormalMemoryRegionCache extends MemoryRegionCache {
        NormalMemoryRegionCache(int size) {
            super(size, Normal);
        }

        @Override
        protected UntetheredMemory allocBuf(PoolChunk chunk, long handle, int size, PoolThreadCache threadCache,
                                  PooledAllocatorControl control) {
            return chunk.allocateBuffer(handle, size, threadCache, control);
        }
    }

    private abstract static class MemoryRegionCache {
        private final int size;
        private final Queue<Entry> queue;
        private final SizeClass sizeClass;
        private int allocations;

        MemoryRegionCache(int size, SizeClass sizeClass) {
            this.size = MathUtil.safeFindNextPositivePowerOfTwo(size);
            queue = PlatformDependent.newFixedMpscQueue(this.size);
            this.sizeClass = sizeClass;
        }

        /**
         * Allocate a new {@link UntetheredMemory} using the provided chunk and handle with the capacity restrictions.
         */
        protected abstract UntetheredMemory allocBuf(
                PoolChunk chunk, long handle, int size, PoolThreadCache threadCache, PooledAllocatorControl control);

        /**
         * Add to cache if not already full.
         */
        public final boolean add(PoolChunk chunk, long handle, int normCapacity) {
            Entry entry = newEntry(chunk, handle, normCapacity);
            boolean queued = queue.offer(entry);
            if (!queued) {
                // If it was not possible to cache the chunk, immediately recycle the entry
                entry.recycle();
            }

            return queued;
        }

        /**
         * Allocate something out of the cache if possible and remove the entry from the cache.
         */
        public final UntetheredMemory allocate(int size, PoolThreadCache threadCache, PooledAllocatorControl control) {
            Entry entry = queue.poll();
            if (entry == null) {
                return null;
            }
            UntetheredMemory buffer = allocBuf(entry.chunk, entry.handle, size, threadCache, control);
            entry.recycle();

            // allocations are not thread-safe which is fine as this is only called from the same thread all time.
            allocations++;
            return buffer;
        }

        /**
         * Clear out this cache and free up all previous cached {@link PoolChunk}s and {@code handle}s.
         */
        public final int free() {
            return free(Integer.MAX_VALUE);
        }

        private int free(int max) {
            int numFreed = 0;
            for (; numFreed < max; numFreed++) {
                Entry entry = queue.poll();
                if (entry != null) {
                    freeEntry(entry);
                } else {
                    // all cleared
                    return numFreed;
                }
            }
            return numFreed;
        }

        /**
         * Free up cached {@link PoolChunk}s if not allocated frequently enough.
         */
        public final void trim() {
            int free = size - allocations;
            allocations = 0;

            // We not even allocated all the number that are
            if (free > 0) {
                free(free);
            }
        }

        private  void freeEntry(Entry entry) {
            PoolChunk chunk = entry.chunk;
            long handle = entry.handle;

            entry.recycle();
            chunk.arena.freeChunk(chunk, handle, entry.normCapacity, sizeClass);
        }

        static final class Entry {
            final Handle<Entry> recyclerHandle;
            PoolChunk chunk;
            long handle = -1;
            int normCapacity;

            Entry(Handle<Entry> recyclerHandle) {
                this.recyclerHandle = recyclerHandle;
            }

            void recycle() {
                chunk = null;
                handle = -1;
                recyclerHandle.recycle(this);
            }
        }

        private static Entry newEntry(PoolChunk chunk, long handle, int normCapacity) {
            Entry entry = RECYCLER.get();
            entry.chunk = chunk;
            entry.handle = handle;
            entry.normCapacity = normCapacity;
            return entry;
        }

        private static final ObjectPool<Entry> RECYCLER = ObjectPool.newPool(handle -> new Entry(handle));
    }
}