/*
 * Copyright (c) 2000, 2010, 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.
 */

/* @test
 * @summary Test selectors and socketchannels
 * @library ..
 * @key randomness
 */

import java.io.*;
import java.net.*;
import java.nio.*;
import java.nio.channels.*;
import java.nio.channels.spi.SelectorProvider;
import java.util.*;


public class SelectorTest {
    private static List clientList = new LinkedList();
    private static Random rnd = new Random();
    public static int NUM_CLIENTS = 30;
    public static int TEST_PORT = 31452;
    static PrintStream log = System.err;
    private static int FINISH_TIME = 30000;

    /*
     * Usage note
     *
     * java SelectorTest [server] [client <host>] [<port>]
     *
     * No arguments runs both client and server in separate threads
     * using the default port of 31452.
     *
     * client runs the client on this machine and connects to server
     * at the given IP address.
     *
     * server runs the server on localhost.
     */
    public static void main(String[] args) throws Exception {
        if (args.length == 0) {
            Server server = new Server(0);
            server.start();
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) { }
            InetSocketAddress isa
                = new InetSocketAddress(InetAddress.getLocalHost(), server.port());
            Client client = new Client(isa);
            client.start();
            if ((server.finish(FINISH_TIME) & client.finish(FINISH_TIME)) == 0)
                throw new Exception("Failure");
            log.println();

        } else if (args[0].equals("server")) {

            if (args.length > 1)
                TEST_PORT = Integer.parseInt(args[1]);
            Server server = new Server(TEST_PORT);
            server.start();
            if (server.finish(FINISH_TIME) == 0)
                throw new Exception("Failure");
            log.println();

        } else if (args[0].equals("client")) {

            if (args.length < 2) {
                log.println("No host specified: terminating.");
                return;
            }
            String ip = args[1];
            if (args.length > 2)
                TEST_PORT = Integer.parseInt(args[2]);
            InetAddress ia = InetAddress.getByName(ip);
            InetSocketAddress isa = new InetSocketAddress(ia, TEST_PORT);
            Client client = new Client(isa);
            client.start();
            if (client.finish(FINISH_TIME) == 0)
                throw new Exception("Failure");
            log.println();

        } else {
            System.out.println("Usage note:");
            System.out.println("java SelectorTest [server] [client <host>] [<port>]");
            System.out.println("No arguments runs both client and server in separate threads using the default port of 31452.");
            System.out.println("client runs the client on this machine and connects to the server specified.");
            System.out.println("server runs the server on localhost.");
        }
    }

    static class Client extends TestThread {
        InetSocketAddress isa;
        Client(InetSocketAddress isa) {
            super("Client", SelectorTest.log);
            this.isa = isa;
        }

        public void go() throws Exception {
            log.println("starting client...");
            for (int i=0; i<NUM_CLIENTS; i++)
                clientList.add(new RemoteEntity(i, isa, log));

            Collections.shuffle(clientList);

            log.println("created "+NUM_CLIENTS+" clients");
            do {
                for (Iterator i = clientList.iterator(); i.hasNext(); ) {
                    RemoteEntity re = (RemoteEntity) i.next();
                    if (re.cycle()) {
                        i.remove();
                    }
                }
                Collections.shuffle(clientList);
            } while (clientList.size() > 0);
        }
    }

    static class Server extends TestThread {
        private final ServerSocketChannel ssc;
        private List socketList = new ArrayList();
        private ServerSocket ss;
        private int connectionsAccepted = 0;
        private Selector pollSelector;
        private Selector acceptSelector;
        private Set pkeys;
        private Set pskeys;

        Server(int port) throws IOException {
            super("Server", SelectorTest.log);
            this.ssc = ServerSocketChannel.open().bind(new InetSocketAddress(port));
        }

        int port() {
            return ssc.socket().getLocalPort();
        }

        public void go() throws Exception {
            log.println("starting server...");
            acceptSelector = SelectorProvider.provider().openSelector();
            pollSelector = SelectorProvider.provider().openSelector();
            pkeys = pollSelector.keys();
            pskeys = pollSelector.selectedKeys();
            Set readyKeys = acceptSelector.selectedKeys();
            RequestHandler rh = new RequestHandler(pollSelector, log);
            Thread requestThread = new Thread(rh);

            requestThread.start();

            ssc.configureBlocking(false);
            SelectionKey acceptKey = ssc.register(acceptSelector,
                                                  SelectionKey.OP_ACCEPT);
            while(connectionsAccepted < SelectorTest.NUM_CLIENTS) {
                int keysAdded = acceptSelector.select(100);
                if (keysAdded > 0) {
                    Iterator i = readyKeys.iterator();
                    while(i.hasNext()) {
                        SelectionKey sk = (SelectionKey)i.next();
                        i.remove();
                        ServerSocketChannel nextReady =
                            (ServerSocketChannel)sk.channel();
                        SocketChannel sc = nextReady.accept();
                        connectionsAccepted++;
                        if (sc != null) {
                            sc.configureBlocking(false);
                            synchronized (pkeys) {
                               sc.register(pollSelector, SelectionKey.OP_READ);
                            }
                        } else {
                            throw new RuntimeException(
                                "Socket does not support Channels");
                        }
                    }
                }
            }
            acceptKey.cancel();
            requestThread.join();
            acceptSelector.close();
            pollSelector.close();
        }
    }
}

class RemoteEntity {
    private static Random rnd = new Random();
    int id;
    ByteBuffer data;
    int dataWrittenIndex;
    int totalDataLength;
    boolean initiated = false;
    boolean connected = false;
    boolean written = false;
    boolean acked = false;
    boolean closed = false;
    private SocketChannel sc;
    ByteBuffer ackBuffer;
    PrintStream log;
    InetSocketAddress server;

    RemoteEntity(int id, InetSocketAddress server, PrintStream log)
        throws Exception
    {
        int connectFailures = 0;
        this.id = id;
        this.log = log;
        this.server = server;

        sc = SocketChannel.open();
        sc.configureBlocking(false);

        // Prepare the data buffer to write out from this entity
        // Let's use both slow and fast buffers
        if (rnd.nextBoolean())
            data = ByteBuffer.allocateDirect(100);
        else
            data = ByteBuffer.allocate(100);
        String number = Integer.toString(id);
        if (number.length() == 1)
            number = "0"+number;
        String source = "Testing from " + number;
        data.put(source.getBytes("8859_1"));
        data.flip();
        totalDataLength = source.length();

        // Allocate an ack buffer
        ackBuffer = ByteBuffer.allocateDirect(10);
    }

    private void reset() throws Exception {
        sc.close();
        sc = SocketChannel.open();
        sc.configureBlocking(false);
    }

    private void connect() throws Exception {
        try {
            connected = sc.connect(server);
            initiated = true;
        }  catch (ConnectException e) {
            initiated = false;
            reset();
        }
    }

    private void finishConnect() throws Exception {
        try {
            connected = sc.finishConnect();
        }  catch (IOException e) {
            initiated = false;
            reset();
        }
    }

    int id() {
        return id;
    }

    boolean cycle() throws Exception {
        if (!initiated)
            connect();
        else if (!connected)
            finishConnect();
        else if (!written)
            writeCycle();
        else if (!acked)
            ackCycle();
        else if (!closed)
            close();
        return closed;
    }

    private void ackCycle() throws Exception {
        //log.println("acking from "+id);
        int bytesRead = sc.read(ackBuffer);
        if (bytesRead > 0) {
            acked = true;
        }
    }

    private void close() throws Exception {
        sc.close();
        closed = true;
    }

    private void writeCycle() throws Exception {
        log.println("writing from "+id);
        int numBytesToWrite = rnd.nextInt(10)+1;
        int newWriteTarget = dataWrittenIndex + numBytesToWrite;
        if (newWriteTarget > totalDataLength)
            newWriteTarget = totalDataLength;
        data.limit(newWriteTarget);
        int bytesWritten = sc.write(data);
        if (bytesWritten > 0)
            dataWrittenIndex += bytesWritten;
        if (dataWrittenIndex == totalDataLength) {
            written = true;
            sc.socket().shutdownOutput();
        }
    }

}


class RequestHandler implements Runnable {
    private static Random rnd = new Random();
    private Selector selector;
    private int connectionsHandled = 0;
    private HashMap dataBin = new HashMap();
    PrintStream log;

    public RequestHandler(Selector selector, PrintStream log) {
        this.selector = selector;
        this.log = log;
    }

    public void run() {
        log.println("starting request handler...");
        int connectionsAccepted = 0;

        Set nKeys = selector.keys();
        Set readyKeys = selector.selectedKeys();

        try {
            while(connectionsHandled < SelectorTest.NUM_CLIENTS) {
                int numKeys = selector.select(100);

                // Process channels with data
                synchronized (nKeys) {
                    if (readyKeys.size() > 0) {
                        Iterator i = readyKeys.iterator();
                        while(i.hasNext()) {
                            SelectionKey sk = (SelectionKey)i.next();
                            i.remove();
                            SocketChannel sc = (SocketChannel)sk.channel();
                            if (sc.isOpen())
                                read(sk, sc);
                        }
                    }
                }

                // Give other threads a chance to run
                if (numKeys == 0) {
                    try {
                        Thread.sleep(1);
                    } catch (Exception x) {}
                }
            }
        } catch (Exception e) {
            log.println("Unexpected error 1: "+e);
            e.printStackTrace();
        }
    }

    private void read(SelectionKey sk, SocketChannel sc) throws Exception {
        ByteBuffer bin = (ByteBuffer)dataBin.get(sc);
        if (bin == null) {
            if (rnd.nextBoolean())
                bin = ByteBuffer.allocateDirect(100);
            else
                bin = ByteBuffer.allocate(100);
            dataBin.put(sc, bin);
        }

        int bytesRead = 0;
        do {
            bytesRead = sc.read(bin);
        } while(bytesRead > 0);

        if (bytesRead == -1) {
            sk.interestOps(0);
            bin.flip();
            int size = bin.limit();
            byte[] data = new byte[size];
            for(int j=0; j<size; j++)
                data[j] = bin.get();
            String message = new String(data, "8859_1");
            connectionsHandled++;
            acknowledge(sc);
            log.println("Received >>>"+message + "<<<");
            log.println("Handled: "+connectionsHandled);
        }
    }

    private void acknowledge(SocketChannel sc) throws Exception {
            ByteBuffer ackBuffer = ByteBuffer.allocateDirect(10);
            String s = "ack";
            ackBuffer.put(s.getBytes("8859_1"));
            ackBuffer.flip();
            int bytesWritten = 0;
            while(bytesWritten == 0) {
                bytesWritten += sc.write(ackBuffer);
            }
            sc.close();
    }
}