/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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.apache.flink.runtime.io.disk.iomanager;

import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.util.EnvironmentInformation;
import org.apache.flink.util.IOUtils;
import org.apache.flink.util.ShutdownHookUtil;

import java.io.IOException;
import java.lang.Thread.UncaughtExceptionHandler;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;

import static org.apache.flink.util.Preconditions.checkState;

/** A version of the {@link IOManager} that uses asynchronous I/O. */
public class IOManagerAsync extends IOManager implements UncaughtExceptionHandler {

    /** The writer threads used for asynchronous block oriented channel writing. */
    private final WriterThread[] writers;

    /** The reader threads used for asynchronous block oriented channel reading. */
    private final ReaderThread[] readers;

    /** Flag to signify that the IOManager has been shut down already */
    private final AtomicBoolean isShutdown = new AtomicBoolean();

    /** Shutdown hook to make sure that the directories are removed on exit */
    private final Thread shutdownHook;

    // -------------------------------------------------------------------------
    //               Constructors / Destructors
    // -------------------------------------------------------------------------

    /** Constructs a new asynchronous I/O manager, writing files to the system 's temp directory. */
    public IOManagerAsync() {
        this(EnvironmentInformation.getTemporaryFileDirectory());
    }

    /**
     * Constructs a new asynchronous I/O manager, writing file to the given directory.
     *
     * @param tempDir The directory to write temporary files to.
     */
    public IOManagerAsync(String tempDir) {
        this(new String[] {tempDir});
    }

    /**
     * Constructs a new asynchronous I/O manager, writing file round robin across the given
     * directories.
     *
     * @param tempDirs The directories to write temporary files to.
     */
    public IOManagerAsync(String[] tempDirs) {
        super(tempDirs);

        // start a write worker thread for each directory
        this.writers = new WriterThread[tempDirs.length];
        for (int i = 0; i < this.writers.length; i++) {
            final WriterThread t = new WriterThread();
            this.writers[i] = t;
            t.setName("IOManager writer thread #" + (i + 1));
            t.setDaemon(true);
            t.setUncaughtExceptionHandler(this);
            t.start();
        }

        // start a reader worker thread for each directory
        this.readers = new ReaderThread[tempDirs.length];
        for (int i = 0; i < this.readers.length; i++) {
            final ReaderThread t = new ReaderThread();
            this.readers[i] = t;
            t.setName("IOManager reader thread #" + (i + 1));
            t.setDaemon(true);
            t.setUncaughtExceptionHandler(this);
            t.start();
        }

        // install a shutdown hook that makes sure the temp directories get deleted
        this.shutdownHook =
                ShutdownHookUtil.addShutdownHook(this::close, getClass().getSimpleName(), LOG);
    }

    /**
     * Close method. Shuts down the reader and writer threads immediately, not waiting for their
     * pending requests to be served. This method waits until the threads have actually ceased their
     * operation.
     */
    @Override
    public void close() throws Exception {
        // mark shut down and exit if it already was shut down
        if (!isShutdown.compareAndSet(false, true)) {
            return;
        }

        // Remove shutdown hook to prevent resource leaks
        ShutdownHookUtil.removeShutdownHook(shutdownHook, getClass().getSimpleName(), LOG);

        if (LOG.isDebugEnabled()) {
            LOG.debug("Shutting down I/O manager.");
        }

        // close writing and reading threads with best effort and log problems
        // first notify all to close, then wait until all are closed

        List<AutoCloseable> closeables = new ArrayList<>(writers.length + readers.length + 2);

        for (WriterThread wt : writers) {
            closeables.add(getWriterThreadCloser(wt));
        }

        for (ReaderThread rt : readers) {
            closeables.add(getReaderThreadCloser(rt));
        }

        closeables.add(
                () -> {
                    try {
                        for (WriterThread wt : writers) {
                            wt.join();
                        }
                        for (ReaderThread rt : readers) {
                            rt.join();
                        }
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                    }
                });

        // make sure we call the super implementation in any case and at the last point,
        // because this will clean up the I/O directories
        closeables.add(super::close);

        IOUtils.closeAll(closeables);
    }

    private static AutoCloseable getWriterThreadCloser(WriterThread thread) {
        return () -> {
            try {
                thread.shutdown();
            } catch (Throwable t) {
                throw new IOException("Error while shutting down IO Manager writer thread.", t);
            }
        };
    }

    private static AutoCloseable getReaderThreadCloser(ReaderThread thread) {
        return () -> {
            try {
                thread.shutdown();
            } catch (Throwable t) {
                throw new IOException("Error while shutting down IO Manager reader thread.", t);
            }
        };
    }

    @Override
    public void uncaughtException(Thread t, Throwable e) {
        LOG.error(
                "IO Thread '"
                        + t.getName()
                        + "' terminated due to an exception. Shutting down I/O Manager.",
                e);
        try {
            close();
        } catch (Exception ex) {
            LOG.warn("IOManagerAsync did not shut down properly.", ex);
        }
    }

    // ------------------------------------------------------------------------
    //                        Reader / Writer instantiations
    // ------------------------------------------------------------------------

    @Override
    public BlockChannelWriter<MemorySegment> createBlockChannelWriter(
            FileIOChannel.ID channelID, LinkedBlockingQueue<MemorySegment> returnQueue)
            throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");
        return new AsynchronousBlockWriter(
                channelID, this.writers[channelID.getThreadNum()].requestQueue, returnQueue);
    }

    @Override
    public BlockChannelWriterWithCallback<MemorySegment> createBlockChannelWriter(
            FileIOChannel.ID channelID, RequestDoneCallback<MemorySegment> callback)
            throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");
        return new AsynchronousBlockWriterWithCallback(
                channelID, this.writers[channelID.getThreadNum()].requestQueue, callback);
    }

    /**
     * Creates a block channel reader that reads blocks from the given channel. The reader reads
     * asynchronously, such that a read request is accepted, carried out at some (close) point in
     * time, and the full segment is pushed to the given queue.
     *
     * @param channelID The descriptor for the channel to write to.
     * @param returnQueue The queue to put the full buffers into.
     * @return A block channel reader that reads from the given channel.
     * @throws IOException Thrown, if the channel for the reader could not be opened.
     */
    @Override
    public BlockChannelReader<MemorySegment> createBlockChannelReader(
            FileIOChannel.ID channelID, LinkedBlockingQueue<MemorySegment> returnQueue)
            throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");
        return new AsynchronousBlockReader(
                channelID, this.readers[channelID.getThreadNum()].requestQueue, returnQueue);
    }

    @Override
    public BufferFileWriter createBufferFileWriter(FileIOChannel.ID channelID) throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");

        return new AsynchronousBufferFileWriter(
                channelID, writers[channelID.getThreadNum()].requestQueue);
    }

    @Override
    public BufferFileReader createBufferFileReader(
            FileIOChannel.ID channelID, RequestDoneCallback<Buffer> callback) throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");

        return new AsynchronousBufferFileReader(
                channelID, readers[channelID.getThreadNum()].requestQueue, callback);
    }

    @Override
    public BufferFileSegmentReader createBufferFileSegmentReader(
            FileIOChannel.ID channelID, RequestDoneCallback<FileSegment> callback)
            throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");

        return new AsynchronousBufferFileSegmentReader(
                channelID, readers[channelID.getThreadNum()].requestQueue, callback);
    }

    /**
     * Creates a block channel reader that reads all blocks from the given channel directly in one
     * bulk. The reader draws segments to read the blocks into from a supplied list, which must
     * contain as many segments as the channel has blocks. After the reader is done, the list with
     * the full segments can be obtained from the reader.
     *
     * <p>If a channel is not to be read in one bulk, but in multiple smaller batches, a {@link
     * BlockChannelReader} should be used.
     *
     * @param channelID The descriptor for the channel to write to.
     * @param targetSegments The list to take the segments from into which to read the data.
     * @param numBlocks The number of blocks in the channel to read.
     * @return A block channel reader that reads from the given channel.
     * @throws IOException Thrown, if the channel for the reader could not be opened.
     */
    @Override
    public BulkBlockChannelReader createBulkBlockChannelReader(
            FileIOChannel.ID channelID, List<MemorySegment> targetSegments, int numBlocks)
            throws IOException {
        checkState(!isShutdown.get(), "I/O-Manager is shut down.");
        return new AsynchronousBulkBlockReader(
                channelID,
                this.readers[channelID.getThreadNum()].requestQueue,
                targetSegments,
                numBlocks);
    }

    // -------------------------------------------------------------------------
    //                             For Testing
    // -------------------------------------------------------------------------

    RequestQueue<ReadRequest> getReadRequestQueue(FileIOChannel.ID channelID) {
        return this.readers[channelID.getThreadNum()].requestQueue;
    }

    RequestQueue<WriteRequest> getWriteRequestQueue(FileIOChannel.ID channelID) {
        return this.writers[channelID.getThreadNum()].requestQueue;
    }

    // -------------------------------------------------------------------------
    //                           I/O Worker Threads
    // -------------------------------------------------------------------------

    /** A worker thread for asynchronous reads. */
    private static final class ReaderThread extends Thread {

        protected final RequestQueue<ReadRequest> requestQueue;

        private volatile boolean alive;

        // ---------------------------------------------------------------------
        // Constructors / Destructors
        // ---------------------------------------------------------------------

        protected ReaderThread() {
            this.requestQueue = new RequestQueue<ReadRequest>();
            this.alive = true;
        }

        /**
         * Shuts the thread down. This operation does not wait for all pending requests to be
         * served, halts the thread immediately. All buffers of pending requests are handed back to
         * their channel readers and an exception is reported to them, declaring their request queue
         * as closed.
         */
        protected void shutdown() {
            synchronized (this) {
                if (alive) {
                    alive = false;
                    requestQueue.close();
                    interrupt();
                }

                try {
                    join(1000);
                } catch (InterruptedException ignored) {
                }

                // notify all pending write requests that the thread has been shut down
                IOException ioex = new IOException("IO-Manager has been closed.");

                while (!this.requestQueue.isEmpty()) {
                    ReadRequest request = this.requestQueue.poll();
                    if (request != null) {
                        try {
                            request.requestDone(ioex);
                        } catch (Throwable t) {
                            IOManagerAsync.LOG.error(
                                    "The handler of the request complete callback threw an exception"
                                            + (t.getMessage() == null
                                                    ? "."
                                                    : ": " + t.getMessage()),
                                    t);
                        }
                    }
                }
            }
        }

        // ---------------------------------------------------------------------
        //                             Main loop
        // ---------------------------------------------------------------------

        @Override
        public void run() {

            while (alive) {

                // get the next buffer. ignore interrupts that are not due to a shutdown.
                ReadRequest request = null;
                while (alive && request == null) {
                    try {
                        request = this.requestQueue.take();
                    } catch (InterruptedException e) {
                        if (!this.alive) {
                            return;
                        } else {
                            IOManagerAsync.LOG.warn(
                                    Thread.currentThread() + " was interrupted without shutdown.");
                        }
                    }
                }

                // remember any IO exception that occurs, so it can be reported to the writer
                IOException ioex = null;

                try {
                    // read buffer from the specified channel
                    request.read();
                } catch (IOException e) {
                    ioex = e;
                } catch (Throwable t) {
                    ioex = new IOException("The buffer could not be read: " + t.getMessage(), t);
                    IOManagerAsync.LOG.error(
                            "I/O reading thread encountered an error"
                                    + (t.getMessage() == null ? "." : ": " + t.getMessage()),
                            t);
                }

                // invoke the processed buffer handler of the request issuing reader object
                try {
                    request.requestDone(ioex);
                } catch (Throwable t) {
                    IOManagerAsync.LOG.error(
                            "The handler of the request-complete-callback threw an exception"
                                    + (t.getMessage() == null ? "." : ": " + t.getMessage()),
                            t);
                }
            } // end while alive
        }
    } // end reading thread

    /** A worker thread that asynchronously writes the buffers to disk. */
    private static final class WriterThread extends Thread {

        protected final RequestQueue<WriteRequest> requestQueue;

        private volatile boolean alive;

        // ---------------------------------------------------------------------
        // Constructors / Destructors
        // ---------------------------------------------------------------------

        protected WriterThread() {
            this.requestQueue = new RequestQueue<WriteRequest>();
            this.alive = true;
        }

        /**
         * Shuts the thread down. This operation does not wait for all pending requests to be
         * served, halts the thread immediately. All buffers of pending requests are handed back to
         * their channel writers and an exception is reported to them, declaring their request queue
         * as closed.
         */
        protected void shutdown() {
            synchronized (this) {
                if (alive) {
                    alive = false;
                    requestQueue.close();
                    interrupt();
                }

                try {
                    join(1000);
                } catch (InterruptedException ignored) {
                }

                // notify all pending write requests that the thread has been shut down
                IOException ioex = new IOException("IO-Manager has been closed.");

                while (!this.requestQueue.isEmpty()) {
                    WriteRequest request = this.requestQueue.poll();
                    if (request != null) {
                        try {
                            request.requestDone(ioex);
                        } catch (Throwable t) {
                            IOManagerAsync.LOG.error(
                                    "The handler of the request complete callback threw an exception"
                                            + (t.getMessage() == null
                                                    ? "."
                                                    : ": " + t.getMessage()),
                                    t);
                        }
                    }
                }
            }
        }

        // ---------------------------------------------------------------------
        // Main loop
        // ---------------------------------------------------------------------

        @Override
        public void run() {

            while (this.alive) {

                WriteRequest request = null;

                // get the next buffer. ignore interrupts that are not due to a shutdown.
                while (alive && request == null) {
                    try {
                        request = requestQueue.take();
                    } catch (InterruptedException e) {
                        if (!this.alive) {
                            return;
                        } else {
                            IOManagerAsync.LOG.warn(
                                    Thread.currentThread() + " was interrupted without shutdown.");
                        }
                    }
                }

                // remember any IO exception that occurs, so it can be reported to the writer
                IOException ioex = null;

                try {
                    // write buffer to the specified channel
                    request.write();
                } catch (IOException e) {
                    ioex = e;
                } catch (Throwable t) {
                    ioex = new IOException("The buffer could not be written: " + t.getMessage(), t);
                    IOManagerAsync.LOG.error(
                            "I/O writing thread encountered an error"
                                    + (t.getMessage() == null ? "." : ": " + t.getMessage()),
                            t);
                }

                // invoke the processed buffer handler of the request issuing writer object
                try {
                    request.requestDone(ioex);
                } catch (Throwable t) {
                    IOManagerAsync.LOG.error(
                            "The handler of the request-complete-callback threw an exception"
                                    + (t.getMessage() == null ? "." : ": " + t.getMessage()),
                            t);
                }
            } // end while alive
        }
    }; // end writer thread
}
