All Downloads are FREE. Search and download functionalities are using the official Maven repository.

dorkbox.executor.stream.PumpStreamHandler.kt Maven / Gradle / Ivy

/*
 * Copyright 2020 dorkbox, llc
 * Copyright (C) 2014 ZeroTurnaround 
 * Contains fragments of code from Apache Commons Exec, rights owned
 * by Apache Software Foundation (ASF).
 *
 * Licensed 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.
 *
 * NOTICE: This file originates from the Apache Commons Exec package.
 * It has been modified to fit our needs.
 *
 * The following is the original header of the file in Apache Commons Exec:
 *
 *   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.
 */

@file:Suppress("BlockingMethodInNonBlockingContext")
package dorkbox.executor.stream

import dorkbox.executor.Executor
import dorkbox.executor.stream.nopStreams.NopInputStream
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import org.slf4j.LoggerFactory
import org.slf4j.MDC
import java.io.IOException
import java.io.InputStream
import java.io.OutputStream

interface IOStream {
    /**
     * Closes this output stream and releases any system resources
     * associated with this stream. The general contract of `close`
     * is that it closes the output stream. A closed stream cannot perform
     * output operations and cannot be reopened.
     *
     *
     * The `close` method of `OutputStream` does nothing.
     *
     * @exception  IOException  if an I/O error occurs.
     */
    @Throws(IOException::class)
    fun close()

    /**
     * Transfer data between IO streams
     */
    suspend fun pump(length: Int, inputStream: InputStream)
}



@JvmInline
value class PumpedOutputStream(private val out: OutputStream) : IOStream {
    override suspend fun pump(length: Int, inputStream: InputStream) {
        (0 until length.coerceAtMost(PumpStreamHandler.DEFAULT_SIZE)).forEach { _ ->
            out.write(inputStream.read())
        }
        out.flush()
    }

    override fun close() {
        out.close()
    }
}

class PumpedOutputChannel(private val channel: Channel) : IOStream {
    override suspend fun pump(length: Int, inputStream: InputStream) {
        (0 until length.coerceAtMost(PumpStreamHandler.DEFAULT_SIZE)).forEach { _ ->
            channel.send(inputStream.read().toByte())
        }
    }

    override fun close() {
        channel.close()
    }
}


/**
 * Copies standard output and error of subprocesses to standard output and error
 * of the parent process. If output or error stream are set to null, any feedback
 * from that stream will be lost.
 *
 * @param out the output [OutputStream].
 * @param err the error  [OutputStream].
 * @param input the input [InputStream].
*/
class PumpStreamHandler(out: OutputStream = System.out,
                        err: OutputStream = System.err,
                        input: InputStream = NopInputStream.INPUT_STREAM,
                        asyncSupport: Boolean = false) : IOStreamHandler(out, err, input, asyncSupport) {

    companion object {
        private val log = LoggerFactory.getLogger(PumpStreamHandler::class.java)

        /**
         * the default size of the internal buffer for copying the streams
         */
        internal const val DEFAULT_SIZE = 1024

        /**
         * Poll the input stream so we don't block forever when trying to read from it
         */
        const val POLL_TIMEOUT = 250L
    }


    // control stream IO pumping
    @Volatile
    private var stop = false

    @Volatile
    private var finishedCleanly = false

    // coroutine job for pumping the IO streams
    private lateinit var pumpJob: Job
    private var pumpJobB: Job? = null
    private var pumpJobC: Job? = null


    // size of this channel is the same as the byte array sizes we pump data from
    internal val channel = Channel(DEFAULT_SIZE)

    private val output: IOStream
    private val error: IOStream

    init {
        if (asyncSupport) {
            // connect this to the receive part?
            output = PumpedOutputChannel(channel)
            error = PumpedOutputChannel(channel)
        } else {
            output = PumpedOutputStream(out)
            error = PumpedOutputStream(err)
        }
    }

    private fun runWithContext(action: suspend () -> Unit): Job {
        val contextMap: Map? = MDC.getCopyOfContextMap()
        return if (contextMap != null) {
            Executor.IO_DISPATCH.launch {
                MDC.setContextMap(contextMap)
                try {
                    action()
                } finally {
                    MDC.clear()
                }
            }
        } else {
            Executor.IO_DISPATCH.launch {
                action()
            }
        }
    }

    /**
     * Setup and start the IO stream processing for the subprocess
     *
     * @param process this is the process we are pumping IO for
     * @param separateErrorStream true to indicate we have separate error/output streams to pump.
     *                            false means error/output are both the "output" stream
     */
    override fun start(process: Process, separateErrorStream: Boolean, highPerformanceIO: Boolean) {
        if (highPerformanceIO) {
            pumpJob = runWithContext {
                runInThread(process)
            }
            pumpJobB = runWithContext {
                runOutThread(process)
            }

            // if we have a separate error stream, start up an I/O pumper for it
            if (separateErrorStream) {
                pumpJobC = runWithContext {
                    runErrThread(process)
                }
            }
        } else {
            pumpJob = runWithContext {
                runAllSingleThread(process, separateErrorStream)
            }
        }
    }

    @Suppress("BlockingMethodInNonBlockingContext", "DuplicatedCode")
    private suspend fun runAllSingleThread(process: Process, separateErrorStream: Boolean) = withContext(Dispatchers.IO) {
        var length: Int
        var readData: Boolean


        // these are the streams connected to the process I/O streams.
        // These are "flipped", since we write to the process input, and read from the process out/err
        val processIn: OutputStream
        val processOut: InputStream
        val processErr: InputStream

        try {
            // Set the [OutputStream] by means of which input can be sent to the process.
            processIn = process.outputStream

            // Set the [InputStream] from which to read the standard output of the process.
            processOut = process.inputStream

            // Set the [InputStream] from which to read the standard error of the process.
            processErr = if (separateErrorStream) {
                process.errorStream
            } else {
                // assign them to the same, since that's an easy check (without having to be a null check)
                processOut
            }
        } catch (e: IOException) {
            // hard abort if we can't do this
            process.destroy()
            throw e
        }


        try {
            // different than a while loop because we want to run this at least once
            // we pump each stream at a time in a single thread/coroutine
            while(true) {
                // keep track if we read any data. We suspend for a short timeout if we haven't ready anything
                readData = false

                // pump input -> process-input (NOTE: Can throw an exception if closed!)
                length = try { input.available() } catch (e: Exception) { 0 }  // NOTE: These checks prevent blocking forever on stream.read(...)
                if (length > 0) {
                    readData = true

                    (0 until length.coerceAtMost(DEFAULT_SIZE)).forEach { _ ->
                        processIn.write(input.read())
                    }
                    processIn.flush()
                }

                // pump process-output -> out (NOTE: Can throw an exception if closed!)
                length = try { processOut.available() } catch (e: Exception) { 0 }
                if (length > 0) {
                    readData = true
                    output.pump(length, processOut)
                }

                // pump process-error -> err  (NOTE: Can throw an exception if closed!)
                if (processErr !== processOut) {
                    length = try { processErr.available() } catch (e: Exception) { 0 }
                    if (length > 0) {
                        readData = true
                        error.pump(length, processErr)
                    }
                }

                // NOTE: We only ACTUALLY stop once all the data has been pumped.
                if (stop) {
                    //  - if the subprocess is aborted, then the process itself is killed
                    if (!finishedCleanly) {
                        break
                    }

                    //  - if the subprocess exits normally, make sure that all the data is pumped
                    if (readData) {
                        continue
                    }

                    // only ACTUALLY stop once all the data has be pumped
                    break
                }

                // DO NOT suspend if we have extra data to read!
                if (readData) {
                    continue
                }

                // suspend if we don't have any data to read.
                delay(POLL_TIMEOUT)
            }
        } catch (closingException: CancellationException) {
            // ignored, since we are explicitly closing the coroutine
        }
    }

    @Suppress("BlockingMethodInNonBlockingContext", "DuplicatedCode")
    private suspend fun runInThread(process: Process) = withContext(Dispatchers.IO) {
        val buf = ByteArray(DEFAULT_SIZE)
        var length: Int
        var readData: Boolean


        // these are the streams connected to the process I/O streams.
        // These are "flipped", since we write to the process input, and read from the process out/err
        val processIn: OutputStream

        try {
            // Set the [OutputStream] by means of which input can be sent to the process.
            processIn = process.outputStream
        } catch (e: IOException) {
            // hard abort if we can't do this
            process.destroy()
            throw e
        }



        try {
            // different than a while loop because we want to run this at least once
            // we pump each stream at a time in a single thread/coroutine
            while(true) {
                // keep track if we read any data. We suspend for a short timeout if we haven't ready anything
                readData = false

                // pump input -> process-input  (NOTE: Can throw an exception if closed!)
                length = try { input.available() } catch (e: Exception) { 0 }  // NOTE: These checks prevent blocking forever on stream.read(...)

                if (length > 0) {
                    readData = true
                    length = input.read(buf, 0, length.coerceAtMost(DEFAULT_SIZE))

                    processIn.write(buf, 0, length)
                    processIn.flush()
                }

                // NOTE: We only ACTUALLY stop once all the data has been pumped.
                if (stop) {
                    //  - if the subprocess is aborted, then the process itself is killed
                    if (!finishedCleanly) {
                        break
                    }

                    //  - if the subprocess exits normally, make sure that all the data is pumped
                    if (readData) {
                        continue
                    }

                    // only ACTUALLY stop once all the data has be pumped
                    break
                }

                // DO NOT suspend if we have extra data to read!
                if (readData) {
                    continue
                }

                // suspend if we don't have any data to read.
                delay(POLL_TIMEOUT)
            }
        } catch (closingException: CancellationException) {
            // ignored, since we are explicitly closing the coroutine
        }
    }


    @Suppress("BlockingMethodInNonBlockingContext", "DuplicatedCode")
    private suspend fun runOutThread(process: Process) = withContext(Dispatchers.IO) {
        var length: Int
        var readData: Boolean


        // these are the streams connected to the process I/O streams.
        // These are "flipped", since we write to the process input, and read from the process out/err
        val processOut: InputStream

        try {
            // Set the [InputStream] from which to read the standard output of the process.
            processOut = process.inputStream
        } catch (e: IOException) {
            // hard abort if we can't do this
            process.destroy()
            throw e
        }

        try {
            // different than a while loop because we want to run this at least once
            // we pump each stream at a time in a single thread/coroutine
            while(true) {
                // keep track if we read any data. We suspend for a short timeout if we haven't ready anything
                readData = false

                // pump process-output -> out  (NOTE: Can throw an exception if closed!)
                length = try { processOut.available() } catch (e: Exception) { 0 }
                if (length > 0) {
                    readData = true
                    output.pump(length, processOut)
                }

                // NOTE: We only ACTUALLY stop once all the data has been pumped.
                if (stop) {
                    //  - if the subprocess is aborted, then the process itself is killed
                    if (!finishedCleanly) {
                        break
                    }

                    //  - if the subprocess exits normally, make sure that all the data is pumped
                    if (readData) {
                        continue
                    }

                    // only ACTUALLY stop once all the data has be pumped
                    break
                }

                // DO NOT suspend if we have extra data to read!
                if (readData) {
                    continue
                }

                // suspend if we don't have any data to read.
                delay(POLL_TIMEOUT)
            }
        } catch (closingException: CancellationException) {
            // ignored, since we are explicitly closing the coroutine
            log.error("pow")
        }
    }

    @Suppress("BlockingMethodInNonBlockingContext", "DuplicatedCode")
    private suspend fun runErrThread(process: Process) = withContext(Dispatchers.IO) {
        var length: Int
        var readData: Boolean


        // these are the streams connected to the process I/O streams.
        // These are "flipped", since we write to the process input, and read from the process out/err
        val processErr: InputStream

        try {
            // Set the [InputStream] from which to read the standard error of the process.
            processErr = process.errorStream
        } catch (e: IOException) {
            // hard abort if we can't do this
            process.destroy()
            throw e
        }

        try {
            // different than a while loop because we want to run this at least once
            // we pump each stream at a time in a single thread/coroutine
            while(true) {
                // keep track if we read any data. We suspend for a short timeout if we haven't ready anything
                readData = false

                // pump process-error -> err  (NOTE: Can throw an exception if closed!)
                length = try { processErr.available() } catch (e: Exception) { 0 }
                if (length > 0) {
                    readData = true
                    error.pump(length, processErr)
                }

                // NOTE: We only ACTUALLY stop once all the data has been pumped.
                if (stop) {
                    //  - if the subprocess is aborted, then the process itself is killed
                    if (!finishedCleanly) {
                        break
                    }

                    //  - if the subprocess exits normally, make sure that all the data is pumped
                    if (readData) {
                        continue
                    }

                    // only ACTUALLY stop once all the data has be pumped
                    break
                }

                // DO NOT suspend if we have extra data to read!
                if (readData) {
                    continue
                }

                // suspend if we don't have any data to read.
                delay(POLL_TIMEOUT)
            }
        } catch (closingException: CancellationException) {
            // ignored, since we are explicitly closing the coroutine
        }
    }


    /**
     * Our subprocess has finished running. (when we "abort" the process, we interrupt the thread)
     */
    @Suppress("BlockingMethodInNonBlockingContext")
    override suspend fun stop(process: Process, finishedCleanly: Boolean) {
        this.stop = true
        this.finishedCleanly = finishedCleanly

        // wait for the pump job to end
        log.trace("waiting for IO pump job")

        try {
            pumpJob.join()
        } catch (ignored: InterruptedException) {
        }

        try {
            pumpJobB?.join()
        } catch (ignored: InterruptedException) {
        }

        try {
            pumpJobC?.join()
        } catch (ignored: InterruptedException) {
        }

        /**
         * Close the streams belonging to the given Process.
         *
         * @throws IOException
         */
        var caught: IOException? = null

        try {
            process.outputStream.close()
        } catch (e: IOException) {
            if (e.message == "Stream closed") {
                /**
                 * OutputStream's contract for the close() method: If the stream is already closed then invoking this method has no effect.
                 *
                 * When a UNIXProcess exits ProcessPipeOutputStream automatically closes its target FileOutputStream and replaces it with NullOutputStream.
                 * However the ProcessPipeOutputStream doesn't close itself at that moment.
                 * As ProcessPipeOutputStream extends BufferedOutputStream extends FilterOutputStream closing it flushes the buffer first.
                 * In Java 7 closing FilterOutputStream ignores any exception thrown by the target OutputStream. Since Java 8 these exceptions are now thrown.
                 *
                 * So since Java 8 after UNIXProcess detects the exit and there's something in the output buffer closing this stream throws IOException
                 * with message "Stream closed" from NullOutputStream.
                 */
                log.trace("Failed to close process output stream", e)
            }
            else {
                caught = add(caught, e)
            }
        }

        try {
            process.inputStream.close()
        } catch (e: IOException) {
            caught = add(caught, e)
        }

        try {
            process.errorStream.close()
        } catch (e: IOException) {
            caught = add(caught, e)
        }

        if (caught != null) {
            throw caught
        }
    }

    fun add(exception: IOException?, newException: IOException): IOException {
        if (exception == null) {
            return newException
        }

        exception.addSuppressed(newException)
        return exception
    }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy