/*
 * Copyright (c) 2015 Twitter, Inc. All rights reserved.
 * Licensed under the Apache License v2.0
 * http://www.apache.org/licenses/LICENSE-2.0
 */

package com.twitter.whiskey.nio;

import com.twitter.whiskey.util.Clock;
import com.twitter.whiskey.util.DefaultClock;

import java.io.IOException;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.Comparator;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;


// TODO: consider simply implementing ScheduledExecutorService

/**
 * Manages a single thread for handling event-based execution generated by both the
 * application interface and socket connections.
 *
 * @author Michael Schore
 */
public class RunLoop implements Executor {

    private static final RunLoop INSTANCE = new RunLoop();

    private final ConcurrentLinkedQueue<Runnable> tasks = new ConcurrentLinkedQueue<>();
    private final PriorityBlockingQueue<ScheduledRunnable> scheduledTasks = new PriorityBlockingQueue<>(1, ScheduledRunnable.Comparator());
    private final Selector selector;
    private final Clock clock;
    private final RunLoopThread thread = new RunLoopThread();

    private int executions = 0;
    private int loops = 0;

    private AtomicBoolean running = new AtomicBoolean(false);
    private AtomicBoolean signal = new AtomicBoolean(false);
    private volatile boolean selecting = false;

    public RunLoop() {
        this(new DefaultClock());
    }

    public RunLoop(Clock clock) {

        try {
            selector = Selector.open();
        } catch (IOException e) {
            // Apparently this is platform specific: it can in theory happen on Windows, but
            // never on Linux. On Android... who knows? Not much we can do if it happens.
            throw new RuntimeException(e);
        }
        this.clock = clock;
    }

    public static RunLoop instance() {
        return INSTANCE;
    }

    public void startThread() {
        if (running.compareAndSet(false, true)) thread.start();
    }

    public void stopThread() throws InterruptedException {
        if (running.compareAndSet(true, false)) {
            selector.wakeup();
            thread.join();
        }
    }

    /**
     * Executes the runnable on the internal RunLoopThread, waking the selector if necessary.
     */
    @Override
    public void execute(Runnable command) {
        tasks.add(command);
        // We don't need to perform the CAS unless we're already selecting, due to the check
        // prior to entering a blocking select for tasks in the task queue.
        if (selecting && signal.compareAndSet(false, true)) {
            selector.wakeup();
        }
    }

    /**
     * Schedules the runnable for future execution on the internal RunLoopThread.
     *
     * @param command the runnable to execute
     * @param delay the delay to wait before execution
     * @param unit time unit of the specified delay
     */
    public void schedule(Runnable command, long delay, TimeUnit unit) {
        schedule(command, delay, 0, unit);
    }

    /**
     * Schedules the runnable for future execution on the internal RunLoopThread.
     *
     * @param command the runnable to execute
     * @param delay the delay to wait before execution
     * @param tolerance window of time past the specified delay that execution is still allowed. If
     *                  this is set to 0, then the command will still be executed regardless of how
     *                  much additional time may have passed.
     * @param unit time unit of both the specified delay and tolerance
     */
    public void schedule(Runnable command, long delay, long tolerance, TimeUnit unit) {

        if (delay < 0 || tolerance < 0) {
            throw new IllegalArgumentException();
        }

        long triggerPoint = clock.now() + unit.toMillis(delay);
        scheduledTasks.add(new ScheduledRunnable(command, triggerPoint, unit.toMillis(tolerance)));
        // Ensure we don't hit the race where a new scheduled task is added right before the
        // RunLoop enters a blocking select.
        if (signal.compareAndSet(false, true)) {
            selector.wakeup();
        }
    }

     /**
     * Registers socket events to be handled on the internal RunLoopThread. This method is
     * unsynchronized, and may block if not called from the internal RunLoopThread.
     */
    @SuppressWarnings("ResourceType")
    public void register(int interestSet, Selectable selectable) {
        if (interestSet == 0) {
            return;
        }

        try {
            SelectionKey key = selectable.getChannel().register(selector, interestSet);
            key.attach(selectable);
            // TODO: implement unregister and map the key internally instead of passing it back
            selectable.setSelectionKey(key);
        } catch(ClosedChannelException e) {
            selectable.onClose(e);
        }
    }

    void wake() {
        if (signal.compareAndSet(false, true)) {
            selector.wakeup();
        }
    }

    void run(final boolean blocking) {

        loops++;
        signal.set(false);

        // Check standard tasks
        Runnable currentTask;
        while ((currentTask = tasks.poll()) != null) {
            executions++;
            currentTask.run();
        }

        long selectTimeout = 0;

        // Check scheduled tasks and setup maximum delay
        ScheduledRunnable nextScheduledTask;
        while (!scheduledTasks.isEmpty()) {
            nextScheduledTask = scheduledTasks.peek();
            long now = clock.now();
            if (nextScheduledTask.tolerance > 0 &&
                nextScheduledTask.triggerPoint <= now - nextScheduledTask.tolerance) {
                // Discard the task - we missed the tolerance window
                scheduledTasks.poll();
            } else if (nextScheduledTask.triggerPoint <= now) {
                // It's time to run the task
                executions++;
                nextScheduledTask.run();
                scheduledTasks.poll();
            } else {
                // Determine the select timeout and break
                selectTimeout = nextScheduledTask.triggerPoint - now;
                break;
            }
        }

        int readyChannels = 0;

        // Select
        try {
            selecting = true;
            if (blocking && tasks.isEmpty() && !signal.get()) {
                readyChannels = selector.select(selectTimeout);
            } else {
                readyChannels = selector.selectNow();
            }
            selecting = false;
        } catch (IOException e) {
            // Recovery would have to involve re-registering all sockets
            // on a new Selector. Consider this fatal for now.
            throw new RuntimeException(e);
        }

        if (readyChannels > 0) {
            Set<SelectionKey> selected = selector.selectedKeys();
            for (Iterator<SelectionKey> iterator = selected.iterator(); iterator.hasNext(); ) {
                SelectionKey key = iterator.next();
                iterator.remove();

                Object attachment = key.attachment();
                if (attachment instanceof Selectable) {
                    Selectable selectable = (Selectable) attachment;
                    if (key.isConnectable()) {
                        executions++;
                        selectable.onConnect();
                    } else if (key.isReadable()) {
                        executions++;
                        selectable.onReadable();
                    } else if (key.isWritable()) {
                        executions++;
                        selectable.onWriteable();
                    }
                }
            }
        }
    }

    // TODO: allow cancellation of scheduled tasks
    private class RunLoopThread extends Thread {

        @Override
        public void run() {
            while (running.get()) {
                RunLoop.this.run(true);
            }
        }
    }

    private static class ScheduledRunnable implements Runnable {
        public long triggerPoint;
        public long tolerance;
        private Runnable runnable;

        ScheduledRunnable(Runnable runnable, long triggerPoint, long tolerance) {
            this.runnable = runnable;
            this.triggerPoint = triggerPoint;
            this.tolerance = tolerance;
        }

        @Override
        public void run() {
            runnable.run();
        }

        public static Comparator<ScheduledRunnable> Comparator() {
            return new Comparator<ScheduledRunnable>() {
                @Override
                public int compare(ScheduledRunnable lhs, ScheduledRunnable rhs) {
                    return (int)(lhs.triggerPoint - rhs.triggerPoint);
                }
            };
        }
    }
}

@@ -8116,16 +8116,88 @@
 Thread %7B
+%0A        public RunLoopThread() %7B%0A            setDaemon(true);%0A        %7D
 %0A%0A      

use daemon thread to avoid blocking clean shutdown