[#32] Improved output capturing and added tests

ileasile · ileasile · commit 0f89932c4fe1 · 2019-12-18T19:37:49.000+03:00
diff --git a/src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt b/src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt
@@ -36,7 +36,12 @@ data class KernelConfig(
         val signatureKey: String,
         val pollingIntervalMillis: Long = 100,
         val scriptClasspath: List<File> = emptyList(),
-        val resolverConfig: ResolverConfig?
+        val resolverConfig: ResolverConfig?,
+
+        val captureOutput: Boolean = true,
+        val captureBufferTimeLimitMs: Int = 100,
+        val captureBufferMaxSize: Int = 1000,
+        val cellOutputMaxSize: Int = 100000
 )
 
 val protocolVersion = "5.3"
diff --git a/src/main/kotlin/org/jetbrains/kotlin/jupyter/protocol.kt b/src/main/kotlin/org/jetbrains/kotlin/jupyter/protocol.kt
@@ -15,14 +15,25 @@ enum class ResponseState {
     Ok, Error
 }
 
+enum class JupyterOutType {
+    STDOUT, STDERR;
+    fun optionName() = name.toLowerCase()
+}
+
 data class ResponseWithMessage(val state: ResponseState, val result: MimeTypedResult?, val displays: List<MimeTypedResult> = emptyList(), val stdOut: String? = null, val stdErr: String? = null) {
     val hasStdOut: Boolean = stdOut != null && stdOut.isNotEmpty()
     val hasStdErr: Boolean = stdErr != null && stdErr.isNotEmpty()
 }
 
+fun JupyterConnection.Socket.sendOut(msg:Message, stream: JupyterOutType, text: String) {
+    connection.iopub.send(makeReplyMessage(msg, header = makeHeader("stream", msg),
+            content = jsonObject(
+                    "name" to stream.optionName(),
+                    "text" to text)))
+}
+
 fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJupyter?, executionCount: AtomicLong) {
-    val msgType = msg.header!!["msg_type"]
-    when (msgType) {
+    when (msg.header!!["msg_type"]) {
         "kernel_info_request" ->
             sendWrapped(msg, makeReplyMessage(msg, "kernel_info_reply",
                     content = jsonObject(
@@ -72,18 +83,11 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
                 }
             }
 
-            fun sendOut(stream: String, text: String) {
-                connection.iopub.send(makeReplyMessage(msg, header = makeHeader("stream", msg),
-                        content = jsonObject(
-                                "name" to stream,
-                                "text" to text)))
-            }
-
             if (res.hasStdOut) {
-                sendOut("stdout", res.stdOut!!)
+                sendOut(msg, JupyterOutType.STDOUT, res.stdOut!!)
             }
             if (res.hasStdErr) {
-                sendOut("stderr", res.stdErr!!)
+                sendOut(msg, JupyterOutType.STDERR, res.stdErr!!)
             }
 
             when (res.state) {
@@ -166,13 +170,54 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
     }
 }
 
-class CapturingOutputStream(val stdout: PrintStream, val captureOutput: Boolean) : OutputStream() {
+class CapturingOutputStream(private val stdout: PrintStream,
+                            private val maxOutputSize: Int,
+                            private val captureOutput: Boolean,
+                            private val maxBufferSize: Int,
+                            private val maxBufferLifeTimeMs: Int,
+                            val onCaptured: (String) -> Unit) : OutputStream() {
     val capturedOutput = ByteArrayOutputStream()
+    private var time = System.currentTimeMillis()
+    private var overallOutputSize = 0
+
+    private fun shouldSend(b: Int): Boolean {
+        val c = b.toChar()
+        if (c == '\n' || c == '\r')
+            return true
+        if (capturedOutput.size() >= maxBufferSize)
+            return true
+
+        val currentTime = System.currentTimeMillis()
+        if (currentTime - time >= maxBufferLifeTimeMs) {
+            time = currentTime
+            return true
+        }
+        return false
+    }
 
     override fun write(b: Int) {
+        if (++overallOutputSize > maxOutputSize) {
+            throw OutputLimitExceededException()
+        }
+
         stdout.write(b)
-        if (captureOutput) capturedOutput.write(b)
+
+        if (captureOutput) {
+            capturedOutput.write(b)
+            if (shouldSend(b)) {
+                flush()
+            }
+        }
+    }
+
+    override fun flush() {
+        if (capturedOutput.size() > 0) {
+            onCaptured(capturedOutput.toString("UTF-8"))
+            capturedOutput.reset()
+        }
     }
+
+    class OutputLimitExceededException(message: String = "Cell output limit exceeded"): Exception(message)
 }
 
 fun Any.toMimeTypedResult(): MimeTypedResult? = when (this) {
@@ -186,10 +231,19 @@ fun JupyterConnection.evalWithIO(body: () -> EvalResult?): ResponseWithMessage {
     val out = System.out
     val err = System.err
 
-    // TODO: make configuration option of whether to pipe back stdout and stderr
-    // TODO: make a configuration option to limit the total stdout / stderr possibly returned (in case it goes wild...)
-    val forkedOut = CapturingOutputStream(out, true)
-    val forkedError = CapturingOutputStream(err, false)
+    fun getCapturingStream(stream: PrintStream, outType: JupyterOutType): CapturingOutputStream {
+        return CapturingOutputStream(
+                stream,
+                config.cellOutputMaxSize,
+                config.captureOutput,
+                config.captureBufferMaxSize,
+                config.captureBufferTimeLimitMs) { text ->
+            this.iopub.sendOut(contextMessage!!, outType, text)
+        }
+    }
+
+    val forkedOut = getCapturingStream(out, JupyterOutType.STDOUT)
+    val forkedError = getCapturingStream(err, JupyterOutType.STDERR)
 
     System.setOut(PrintStream(forkedOut, true, "UTF-8"))
     System.setErr(PrintStream(forkedError, true, "UTF-8"))
@@ -202,26 +256,26 @@ fun JupyterConnection.evalWithIO(body: () -> EvalResult?): ResponseWithMessage {
             if (exec == null) {
                 ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), null, "NO REPL!")
             } else {
-                val stdOut = forkedOut.capturedOutput.toString("UTF-8").emptyWhenNull()
-                val stdErr = forkedError.capturedOutput.toString("UTF-8").emptyWhenNull()
+                forkedOut.flush()
+                forkedError.flush()
 
                 try {
                     var result: MimeTypedResult? = null
-                    var displays = exec.displayValues.mapNotNull { it.toMimeTypedResult() }
+                    val displays = exec.displayValues.mapNotNull { it.toMimeTypedResult() }.toMutableList()
                     if (exec.resultValue is DisplayResult) {
                         val resultDisplay = exec.resultValue.value.toMimeTypedResult()
                         if (resultDisplay != null)
                             displays += resultDisplay
                     } else result = exec.resultValue?.toMimeTypedResult()
-                    ResponseWithMessage(ResponseState.Ok, result, displays, stdOut, stdErr)
+                    ResponseWithMessage(ResponseState.Ok, result, displays, null, null)
                 } catch (e: Exception) {
-                    ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), stdOut,
-                            joinLines(stdErr, "error:  Unable to convert result to a string: ${e}"))
+                    ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), null,
+                            "error:  Unable to convert result to a string: $e")
                 }
             }
         } catch (ex: ReplCompilerException) {
-            val stdOut = forkedOut.capturedOutput.toString("UTF-8").emptyWhenNull()
-            val stdErr = forkedError.capturedOutput.toString("UTF-8").emptyWhenNull()
+            forkedOut.flush()
+            forkedError.flush()
 
             // handle runtime vs. compile time and send back correct format of response, now we just send text
             /*
@@ -232,10 +286,10 @@ fun JupyterConnection.evalWithIO(body: () -> EvalResult?): ResponseWithMessage {
                    'traceback' : list(str), # traceback frames as strings
                 }
              */
-            ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), stdOut,
-                    joinLines(stdErr, ex.errorResult.message))
+            ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), null,
+                    ex.errorResult.message)
         } catch (ex: ReplEvalRuntimeException) {
-            val stdOut = forkedOut.capturedOutput.toString("UTF-8").emptyWhenNull()
+            forkedOut.flush()
 
             // handle runtime vs. compile time and send back correct format of response, now we just send text
             /*
@@ -262,7 +316,7 @@ fun JupyterConnection.evalWithIO(body: () -> EvalResult?): ResponseWithMessage {
                     }
                 }
             }
-            ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), stdOut, stdErr.toString())
+            ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), null, stdErr.toString())
         }
     } finally {
         System.setIn(`in`)
@@ -271,7 +325,4 @@ fun JupyterConnection.evalWithIO(body: () -> EvalResult?): ResponseWithMessage {
     }
 }
 
-fun joinLines(vararg parts: String): String = parts.filter(String::isNotBlank).joinToString("\n")
 fun String.nullWhenEmpty(): String? = if (this.isBlank()) null else this
-fun String?.emptyWhenNull(): String = if (this == null || this.isBlank()) "" else this
-
diff --git a/src/test/kotlin/org/jetbrains/kotlin/jupyter/test/capturingStreamTests.kt b/src/test/kotlin/org/jetbrains/kotlin/jupyter/test/capturingStreamTests.kt
@@ -0,0 +1,90 @@
+package org.jetbrains.kotlin.jupyter.test
+
+import org.jetbrains.kotlin.jupyter.CapturingOutputStream
+import org.junit.Assert.assertArrayEquals
+import org.junit.Assert.assertEquals
+import org.junit.Test
+import java.io.OutputStream
+import java.io.PrintStream
+
+class CapturingStreamTests {
+    private val nullOStream = object: OutputStream() {
+        override fun write(b: Int) {
+        }
+    }
+
+    private fun getStream(stdout: OutputStream = nullOStream,
+                          maxOutputSize: Int = 1000,
+                          captureOutput: Boolean = true,
+                          maxBufferSize: Int = 1000,
+                          maxBufferLifeTimeMs: Int = 1000,
+                          onCaptured: (String) -> Unit = {}): CapturingOutputStream {
+
+        val printStream = PrintStream(stdout, false, "UTF-8")
+        return CapturingOutputStream(printStream, maxOutputSize, captureOutput,
+                maxBufferSize, maxBufferLifeTimeMs, onCaptured)
+    }
+
+    @Test
+    fun testMaxOutputSizeOk() {
+        val s = getStream(maxOutputSize = 6)
+        s.write("kotlin".toByteArray())
+    }
+
+    @Test(expected = CapturingOutputStream.OutputLimitExceededException::class)
+    fun testMaxOutputSizeError() {
+        val s = getStream(maxOutputSize = 3)
+        s.write("java".toByteArray())
+    }
+
+    @Test
+    fun testOutputCapturingFlag() {
+        val contents = "abc".toByteArray()
+
+        val s1 = getStream(captureOutput = false)
+        s1.write(contents)
+        assertEquals(0, s1.capturedOutput.size())
+
+        val s2 = getStream(captureOutput = true)
+        s2.write(contents)
+        assertArrayEquals(contents, s2.capturedOutput.toByteArray())
+    }
+
+    @Test
+    fun testMaxBufferSize() {
+        val contents = "0123456789\nfortran".toByteArray()
+        val expected = arrayOf("012", "345", "678", "9\n", "for", "tra", "n")
+
+        var i = 0
+        val s = getStream(maxBufferSize = 3) {
+            assertEquals(expected[i], it)
+            ++i
+        }
+
+        s.write(contents)
+        s.flush()
+
+        assertEquals(expected.size, i)
+    }
+    
+    @Test
+    fun testMaxBufferLifeTime() {
+        val strings = arrayOf("c ", "a", "ada ", "b", "scala ", "c")
+        val expected = arrayOf("c a", "ada b", "scala c")
+
+        var i = 0
+        val s = getStream(maxBufferLifeTimeMs = 1000) {
+            assertEquals(expected[i], it)
+            ++i
+        }
+
+        strings.forEach {
+            Thread.sleep(600)
+            s.write(it.toByteArray())
+        }
+
+        s.flush()
+
+        assertEquals(expected.size, i)
+    }
+}
