/*
 * Copyright (c) 2016, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

import java.io.IOException;

import java.lang.management.BufferPoolMXBean;
import java.lang.management.ManagementFactory;

import java.nio.ByteBuffer;

import java.nio.channels.FileChannel;

import java.nio.file.Path;
import java.nio.file.Paths;

import static java.nio.file.StandardOpenOption.CREATE;
import static java.nio.file.StandardOpenOption.TRUNCATE_EXISTING;
import static java.nio.file.StandardOpenOption.WRITE;

import java.util.List;
import java.util.Random;
import java.util.concurrent.CountDownLatch;

/*
 * @test
 * @requires sun.arch.data.model == "64"
 * @modules java.management
 * @build TestMaxCachedBufferSize
 * @run main/othervm TestMaxCachedBufferSize
 * @run main/othervm -Djdk.nio.maxCachedBufferSize=0 TestMaxCachedBufferSize
 * @run main/othervm -Djdk.nio.maxCachedBufferSize=2000 TestMaxCachedBufferSize
 * @run main/othervm -Djdk.nio.maxCachedBufferSize=100000 TestMaxCachedBufferSize
 * @run main/othervm -Djdk.nio.maxCachedBufferSize=10000000 TestMaxCachedBufferSize
 * @summary Test the implementation of the jdk.nio.maxCachedBufferSize property.
 */
public class TestMaxCachedBufferSize {
    private static final int DEFAULT_ITERS = 10 * 1000;
    private static final int DEFAULT_THREAD_NUM = 4;

    private static final int SMALL_BUFFER_MIN_SIZE =  4 * 1024;
    private static final int SMALL_BUFFER_MAX_SIZE = 64 * 1024;
    private static final int SMALL_BUFFER_DIFF_SIZE =
                                 SMALL_BUFFER_MAX_SIZE - SMALL_BUFFER_MIN_SIZE;

    private static final int LARGE_BUFFER_MIN_SIZE =      512 * 1024;
    private static final int LARGE_BUFFER_MAX_SIZE = 4 * 1024 * 1024;
    private static final int LARGE_BUFFER_DIFF_SIZE =
                                 LARGE_BUFFER_MAX_SIZE - LARGE_BUFFER_MIN_SIZE;

    private static final int LARGE_BUFFER_FREQUENCY = 100;

    private static final String FILE_NAME_PREFIX = "nio-out-file-";
    private static final int VERBOSE_PERIOD = 5 * 1000;

    private static int iters = DEFAULT_ITERS;
    private static int threadNum = DEFAULT_THREAD_NUM;

    private static BufferPoolMXBean getDirectPool() {
        final List<BufferPoolMXBean> pools =
                  ManagementFactory.getPlatformMXBeans(BufferPoolMXBean.class);
        for (BufferPoolMXBean pool : pools) {
            if (pool.getName().equals("direct")) {
                return pool;
            }
        }
        throw new Error("could not find direct pool");
    }
    private static final BufferPoolMXBean directPool = getDirectPool();

    // Each worker will do write operations on a file channel using
    // buffers of various sizes. The buffer size is randomly chosen to
    // be within a small or a large range. This way we can control
    // which buffers can be cached (all, only the small ones, or none)
    // by setting the jdk.nio.maxCachedBufferSize property.
    private static class Worker implements Runnable {
        private final int id;
        private final CountDownLatch finishLatch, exitLatch;
        private final Random random = new Random();
        private long smallBufferCount = 0;
        private long largeBufferCount = 0;

        private int getWriteSize() {
            int minSize = 0;
            int diff = 0;
            if (random.nextInt() % LARGE_BUFFER_FREQUENCY != 0) {
                // small buffer
                minSize = SMALL_BUFFER_MIN_SIZE;
                diff = SMALL_BUFFER_DIFF_SIZE;
                smallBufferCount += 1;
            } else {
                // large buffer
                minSize = LARGE_BUFFER_MIN_SIZE;
                diff = LARGE_BUFFER_DIFF_SIZE;
                largeBufferCount += 1;
            }
            return minSize + random.nextInt(diff);
        }

        private void loop() {
            final String fileName = String.format("%s%d", FILE_NAME_PREFIX, id);

            try {
                for (int i = 0; i < iters; i += 1) {
                    final int writeSize = getWriteSize();

                    // This will allocate a HeapByteBuffer. It should not
                    // be a direct buffer, otherwise the write() method on
                    // the channel below will not create a temporary
                    // direct buffer for the write.
                    final ByteBuffer buffer = ByteBuffer.allocate(writeSize);

                    // Put some random data on it.
                    while (buffer.hasRemaining()) {
                        buffer.put((byte) random.nextInt());
                    }
                    buffer.rewind();

                    final Path file = Paths.get(fileName);
                    try (FileChannel outChannel = FileChannel.open(file, CREATE, TRUNCATE_EXISTING, WRITE)) {
                        // The write() method will create a temporary
                        // direct buffer for the write and attempt to cache
                        // it. It's important that buffer is not a
                        // direct buffer, otherwise the temporary buffer
                        // will not be created.
                        long res = outChannel.write(buffer);
                    }

                    if ((i + 1) % VERBOSE_PERIOD == 0) {
                        System.out.printf(
                          " Worker %3d | %8d Iters | Small %8d Large %8d | Direct %4d / %7dK\n",
                          id, i + 1, smallBufferCount, largeBufferCount,
                          directPool.getCount(), directPool.getTotalCapacity() / 1024);
                    }
                }
            } catch (IOException e) {
                throw new Error("I/O error", e);
            } finally {
                finishLatch.countDown();
                try {
                    exitLatch.await();
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        }

        @Override
        public void run() {
            loop();
        }

        public Worker(int id, CountDownLatch finishLatch, CountDownLatch exitLatch) {
            this.id = id;
            this.finishLatch = finishLatch;
            this.exitLatch = exitLatch;
        }
    }

    public static void checkDirectBuffers(long expectedCount, long expectedMax) {
        final long directCount = directPool.getCount();
        final long directTotalCapacity = directPool.getTotalCapacity();
        System.out.printf("Direct %d / %dK\n",
                          directCount, directTotalCapacity / 1024);

        if (directCount > expectedCount) {
            throw new Error(String.format(
                "inconsistent direct buffer total count, expected = %d, found = %d",
                expectedCount, directCount));
        }

        if (directTotalCapacity > expectedMax) {
            throw new Error(String.format(
                "inconsistent direct buffer total capacity, expectex max = %d, found = %d",
                expectedMax, directTotalCapacity));
        }
    }

    public static void main(String[] args) {
        final String maxBufferSizeStr = System.getProperty("jdk.nio.maxCachedBufferSize");
        final long maxBufferSize =
            (maxBufferSizeStr != null) ? Long.valueOf(maxBufferSizeStr) : Long.MAX_VALUE;

        // We assume that the max cannot be equal to a size of a
        // buffer that can be allocated (makes sanity checking at the
        // end easier).
        if ((SMALL_BUFFER_MIN_SIZE <= maxBufferSize &&
                                     maxBufferSize <= SMALL_BUFFER_MAX_SIZE) ||
            (LARGE_BUFFER_MIN_SIZE <= maxBufferSize &&
                                     maxBufferSize <= LARGE_BUFFER_MAX_SIZE)) {
            throw new Error(String.format("max buffer size = %d not allowed",
                                          maxBufferSize));
        }

        System.out.printf("Threads %d | Iterations %d | MaxBufferSize %d\n",
                          threadNum, iters, maxBufferSize);
        System.out.println();

        final CountDownLatch finishLatch = new CountDownLatch(threadNum);
        final CountDownLatch exitLatch = new CountDownLatch(1);
        final Thread[] threads = new Thread[threadNum];
        for (int i = 0; i < threadNum; i += 1) {
            threads[i] = new Thread(new Worker(i, finishLatch, exitLatch));
            threads[i].start();
        }

        try {
            try {
                finishLatch.await();
            } catch (InterruptedException e) {
                throw new Error("finishLatch.await() interrupted!", e);
            }

            // There is an assumption here that, at this point, only the
            // cached DirectByteBuffers should be active. Given we
            // haven't used any other DirectByteBuffers in this test, this
            // should hold.
            //
            // Also note that we can only do the sanity checking at the
            // end and not during the run given that, at any time, there
            // could be buffers currently in use by some of the workers
            // that will not be cached.

            System.out.println();
            if (maxBufferSize < SMALL_BUFFER_MAX_SIZE) {
                // The max buffer size is smaller than all buffers that
                // were allocated. No buffers should have been cached.
                checkDirectBuffers(0, 0);
            } else if (maxBufferSize < LARGE_BUFFER_MIN_SIZE) {
                // The max buffer size is larger than all small buffers
                // but smaller than all large buffers that were
                // allocated. Only small buffers could have been cached.
                checkDirectBuffers(threadNum,
                                   (long) threadNum * (long) SMALL_BUFFER_MAX_SIZE);
            } else {
                // The max buffer size is larger than all buffers that
                // were allocated. All buffers could have been cached.
                checkDirectBuffers(threadNum,
                                   (long) threadNum * (long) LARGE_BUFFER_MAX_SIZE);
            }
        } finally {
            exitLatch.countDown();
            try {
                for (int i = 0; i < threadNum; i += 1) {
                    threads[i].join();
                }
            } catch (InterruptedException e) {
                // ignore
            }
        }
    }
}