/*

 * Copyright 2009 Google Inc.  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.net.InetSocketAddress;

import java.net.SocketAddress;

import java.nio.channels.SelectionKey;

import java.nio.channels.Selector;

import java.nio.channels.ServerSocketChannel;

import java.nio.channels.SocketChannel;

import java.util.ArrayList;

import java.util.Iterator;

import java.util.List;

/**

 * Reproduces O(N^2) behavior of JDK6/7 select() call. This happens when

 * a selector has many unprocessed updates to its interest set (e.g. adding

 * OP_READ on a bunch of newly accepted sockets). The O(N^2) is triggered

 * by cancelling a number of selection keys (or just closing a few sockets).

 * In this case, select() will first go through the list of cancelled keys

 * and try to deregister them. That deregistration is O(N^2) over the list

 * of unprocessed updates to the interest set.

 *

 * <p> This O(N^2) behavior is a BUG in JVM and should be fixed.

 *

 * <p> The test first creates initCount connections, and adds them

 * to the server epoll set. It then creates massCount connections,

 * registers interest (causing updateList to be populated with massCount*2

 * elements), but does not add them to epoll set (that would've cleared

 * updateList). The test then closes initCount connections, thus populating

 * deregistration queue. The subsequent call to selectNow() will first process

 * deregistration queue, performing O(N^2) over updateList size,

 * equal to massCount * 2.

 *

 * <p> Note that connect rate is artificially slowed down to compensate

 * for what I believe is a Linux bug, where too high of a connection rate

 * ends up in SYN's being dropped and then slow retransmits.

 *

 * @author Igor Chernyshev

 */

public class LotsOfCancels {

    static long testStartTime;

    public static void main(String[] args) throws Exception {

        // the final select should run in less than 1000ms.

        runTest(500, 2700, 1000);

    }

    static void log(String msg) {

        System.out.println(getLogPrefix() + msg);

    }

    static String getLogPrefix() {

        return durationMillis(testStartTime) + ": ";

    }

    /**

     * Returns the elapsed time since startNanos, in milliseconds.

     * @param startNanos the start time; this must be a value returned

     * by {@link System.nanoTime}

     */

    static long durationMillis(long startNanos) {

        return (System.nanoTime() - startNanos) / (1000L * 1000L);

    }

    static void runTest(int initCount, int massCount, int maxSelectTime)

            throws Exception {

        testStartTime = System.nanoTime();

        InetSocketAddress address = new InetSocketAddress("127.0.0.1", 7359);

        // Create server channel, add it to selector and run epoll_ctl.

        log("Setting up server");

        Selector serverSelector = Selector.open();

        ServerSocketChannel server = ServerSocketChannel.open();

        server.configureBlocking(false);

        server.socket().bind(address, 5000);

        server.register(serverSelector, SelectionKey.OP_ACCEPT);

        serverSelector.selectNow();

        log("Setting up client");

        ClientThread client = new ClientThread(address);

        client.start();

        Thread.sleep(100);

        // Set up initial set of client sockets.

        log("Starting initial client connections");

        client.connectClients(initCount);

        Thread.sleep(500);  // Wait for client connections to arrive

        // Accept all initial client sockets, add to selector and run

        // epoll_ctl.

        log("Accepting initial connections");

        List<SocketChannel> serverChannels1 =

            acceptAndAddAll(serverSelector, server, initCount);

        if (serverChannels1.size() != initCount) {

            throw new Exception("Accepted " + serverChannels1.size() +

                                " instead of " + initCount);

        }

        serverSelector.selectNow();

        // Set up mass set of client sockets.

        log("Requesting mass client connections");

        client.connectClients(massCount);

        Thread.sleep(500);  // Wait for client connections to arrive

        // Accept all mass client sockets, add to selector and do NOT

        // run epoll_ctl.

        log("Accepting mass connections");

        List<SocketChannel> serverChannels2 =

            acceptAndAddAll(serverSelector, server, massCount);

        if (serverChannels2.size() != massCount) {

            throw new Exception("Accepted " + serverChannels2.size() +

                                " instead of " + massCount);

        }

        // Close initial set of sockets.

        log("Closing initial connections");

        closeAll(serverChannels1);

        // Now get the timing of select() call.

        log("Running the final select call");

        long startTime = System.nanoTime();

        serverSelector.selectNow();

        long duration = durationMillis(startTime);

        log("Init count = " + initCount +

            ", mass count = " + massCount +

            ", duration = " + duration + "ms");

        if (duration > maxSelectTime) {

            System.out.println

                ("\n\n\n\n\nFAILURE: The final selectNow() took " +

                 duration + "ms " +

                 "- seems like O(N^2) bug is still here\n\n");

            System.exit(1);

        }

    }

    static List<SocketChannel> acceptAndAddAll(Selector selector,

                                               ServerSocketChannel server,

                                               int expected)

            throws Exception {

        int retryCount = 0;

        int acceptCount = 0;

        List<SocketChannel> channels = new ArrayList<SocketChannel>();

        while (channels.size() < expected) {

            SocketChannel channel = server.accept();

            if (channel == null) {

                log("accept() returned null " +

                    "after accepting " + acceptCount + " more connections");

                acceptCount = 0;

                if (retryCount < 10) {

                    // See if more new sockets got stacked behind.

                    retryCount++;

                    Thread.sleep(500);

                    continue;

                }

                break;

            }

            retryCount = 0;

            acceptCount++;

            channel.configureBlocking(false);

            channel.register(selector, SelectionKey.OP_READ);

            channels.add(channel);

        }

        // Cause an additional updateList entry per channel.

        for (SocketChannel channel : channels) {

            channel.register(selector, SelectionKey.OP_WRITE);

        }

        return channels;

    }

    static void closeAll(List<SocketChannel> channels)

            throws Exception {

        for (SocketChannel channel : channels) {

            channel.close();

        }

    }

    static class ClientThread extends Thread {

        private final SocketAddress address;

        private final Selector selector;

        private int connectionsNeeded;

        private int totalCreated;

        ClientThread(SocketAddress address) throws Exception {

            this.address = address;

            selector = Selector.open();

            setDaemon(true);

        }

        void connectClients(int count) throws Exception {

            synchronized (this) {

                connectionsNeeded += count;

            }

            selector.wakeup();

        }

        @Override

        public void run() {

            try {

                handleClients();

            } catch (Throwable e) {

                e.printStackTrace();

                System.exit(1);

            }

        }

        private void handleClients() throws Exception {

            int selectCount = 0;

            while (true) {

                int createdCount = 0;

                synchronized (this) {

                    if (connectionsNeeded > 0) {

                        while (connectionsNeeded > 0 && createdCount < 20) {

                            connectionsNeeded--;

                            createdCount++;

                            totalCreated++;

                            SocketChannel channel = SocketChannel.open();

                            channel.configureBlocking(false);

                            channel.connect(address);

                            if (!channel.finishConnect()) {

                                channel.register(selector,

                                                 SelectionKey.OP_CONNECT);

                            }

                        }

                        log("Started total of " +

                            totalCreated + " client connections");

                        Thread.sleep(200);

                    }

                }

                if (createdCount > 0) {

                    selector.selectNow();

                } else {

                    selectCount++;

                    long startTime = System.nanoTime();

                    selector.select();

                    long duration = durationMillis(startTime);

                    log("Exited clientSelector.select(), loop #"

                        + selectCount + ", duration = " + duration + "ms");

                }

                int keyCount = -1;

                Iterator<SelectionKey> keys =

                    selector.selectedKeys().iterator();

                while (keys.hasNext()) {

                    SelectionKey key = keys.next();

                    synchronized (key) {

                        keyCount++;

                        keys.remove();

                        if (!key.isValid()) {

                            log("Ignoring client key #" + keyCount);

                            continue;

                        }

                        int readyOps = key.readyOps();

                        if (readyOps == SelectionKey.OP_CONNECT) {

                            key.interestOps(0);

                            ((SocketChannel) key.channel()).finishConnect();

                        } else {

                            log("readyOps() on client key #" + keyCount +

                                " returned " + readyOps);

                        }

                    }

                }

            }

        }

    }

}