/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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
 *
 *    http://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 org.elasticsearch.nio;

import java.io.IOException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.function.Consumer;
import java.util.function.Supplier;

public class EventHandler {

    protected final Consumer<Exception> exceptionHandler;
    private final Supplier<NioSelector> selectorSupplier;

    public EventHandler(Consumer<Exception> exceptionHandler, Supplier<NioSelector> selectorSupplier) {
        this.exceptionHandler = exceptionHandler;
        this.selectorSupplier = selectorSupplier;
    }

    /**
     * This method is called when a server channel signals it is ready to accept a connection. All of the
     * accept logic should occur in this call.
     *
     * @param context that can accept a connection
     */
    protected void acceptChannel(ServerChannelContext context) throws IOException {
        context.acceptChannels(selectorSupplier);
    }

    /**
     * This method is called when an attempt to accept a connection throws an exception.
     *
     * @param context that accepting a connection
     * @param exception that occurred
     */
    protected void acceptException(ServerChannelContext context, Exception exception) {
        context.handleException(exception);
    }

    /**
     * This method is called when a NioChannel is being registered with the selector. It should
     * only be called once per channel.
     *
     * @param context that was registered
     */
    protected void handleRegistration(ChannelContext<?> context) throws IOException {
        context.register();
        SelectionKey selectionKey = context.getSelectionKey();
        selectionKey.attach(context);
        if (context instanceof SocketChannelContext) {
            if (((SocketChannelContext) context).readyForFlush()) {
                SelectionKeyUtils.setConnectReadAndWriteInterested(context.getSelectionKey());
            } else {
                SelectionKeyUtils.setConnectAndReadInterested(context.getSelectionKey());
            }
        } else {
            assert context instanceof ServerChannelContext : "If not SocketChannelContext the context must be a ServerChannelContext";
            SelectionKeyUtils.setAcceptInterested(context.getSelectionKey());
        }
    }

    /**
     * This method is called when an attempt to register a channel throws an exception.
     *
     * @param context that was registered
     * @param exception that occurred
     */
    protected void registrationException(ChannelContext<?> context, Exception exception) {
        context.handleException(exception);
    }

    /**
     * This method is called when a NioSocketChannel has just been accepted or if it has receive an
     * OP_CONNECT event.
     *
     * @param context that was registered
     */
    protected void handleConnect(SocketChannelContext context) throws IOException {
        if (context.connect()) {
            SelectionKeyUtils.removeConnectInterested(context.getSelectionKey());
        }
    }

    /**
     * This method is called when an attempt to connect a channel throws an exception.
     *
     * @param context that was connecting
     * @param exception that occurred
     */
    protected void connectException(SocketChannelContext context, Exception exception) {
        context.handleException(exception);
    }

    /**
     * This method is called when a channel signals it is ready for be read. All of the read logic should
     * occur in this call.
     *
     * @param context that can be read
     */
    protected void handleRead(SocketChannelContext context) throws IOException {
        context.read();
    }

    /**
     * This method is called when an attempt to read from a channel throws an exception.
     *
     * @param context that was being read
     * @param exception that occurred
     */
    protected void readException(SocketChannelContext context, Exception exception) {
        context.handleException(exception);
    }

    /**
     * This method is called when a channel signals it is ready to receive writes. All of the write logic
     * should occur in this call.
     *
     * @param context that can be written to
     */
    protected void handleWrite(SocketChannelContext context) throws IOException {
        context.flushChannel();
    }

    /**
     * This method is called when an attempt to write to a channel throws an exception.
     *
     * @param context that was being written to
     * @param exception that occurred
     */
    protected void writeException(SocketChannelContext context, Exception exception) {
        context.handleException(exception);
    }

    /**
     * This method is called when a task or listener attached to a channel is available to run.
     *
     * @param task to handle
     */
    protected void handleTask(Runnable task) {
        task.run();
    }

    /**
     * This method is called when a task or listener attached to a channel operation throws an exception.
     *
     * @param exception that occurred
     */
    protected void taskException(Exception exception) {
        exceptionHandler.accept(exception);
    }

    /**
     * This method is called after events (READ, WRITE, CONNECT) have been handled for a channel.
     *
     * @param context that was handled
     */
    protected void postHandling(SocketChannelContext context) {
        if (context.selectorShouldClose()) {
            try {
                handleClose(context);
            } catch (IOException e) {
                closeException(context, e);
            }
        } else {
            SelectionKey selectionKey = context.getSelectionKey();
            boolean currentlyWriteInterested = SelectionKeyUtils.isWriteInterested(selectionKey);
            boolean pendingWrites = context.readyForFlush();
            if (currentlyWriteInterested == false && pendingWrites) {
                SelectionKeyUtils.setWriteInterested(selectionKey);
            } else if (currentlyWriteInterested && pendingWrites == false) {
                SelectionKeyUtils.removeWriteInterested(selectionKey);
            }
        }
    }

    /**
     * This method handles an IOException that was thrown during a call to {@link Selector#select(long)} or
     * {@link Selector#close()}.
     *
     * @param exception the exception
     */
    protected void selectorException(IOException exception) {
        exceptionHandler.accept(exception);
    }

    /**
     * This method handles an exception that was uncaught during a select loop.
     *
     * @param exception that was uncaught
     */
    protected void uncaughtException(Exception exception) {
        Thread thread = Thread.currentThread();
        thread.getUncaughtExceptionHandler().uncaughtException(thread, exception);
    }

    /**
     * This method handles the closing of an NioChannel
     *
     * @param context that should be closed
     */
    protected void handleClose(ChannelContext<?> context) throws IOException {
        context.closeFromSelector();
        assert context.isOpen() == false : "Should always be done as we are on the selector thread";
    }

    /**
     * This method is called when an attempt to close a channel throws an exception.
     *
     * @param context that was being closed
     * @param exception that occurred
     */
    protected void closeException(ChannelContext<?> context, Exception exception) {
        context.handleException(exception);
    }

    /**
     * This method is called when handling an event from a channel fails due to an unexpected exception.
     * An example would be if checking ready ops on a {@link java.nio.channels.SelectionKey} threw
     * {@link java.nio.channels.CancelledKeyException}.
     *
     * @param context that caused the exception
     * @param exception that was thrown
     */
    protected void genericChannelException(ChannelContext<?> context, Exception exception) {
        context.handleException(exception);
    }
}

@@ -6355,32 +6355,93 @@
        %7D else %7B%0A
+            boolean pendingWrites = context.readyForFlush();%0A
             Sele
@@ -6495,32 +6495,288 @@
 ();%0A            
+if (selectionKey == null) %7B%0A                if (pendingWrites) %7B%0A                    writeException(context, new IllegalStateException(%22Tried to write to an not yet registered channel%22));%0A                %7D%0A                return;%0A            %7D%0A            
 boolean currentl
@@ -6846,69 +6846,8 @@
 nKey
-);%0A            boolean pendingWrites = context.readyForFlush(
 );%0A 

Fix Exceptions in EventHandler#postHandling Breaking Select Loop (#44347)