Post-review fixes

Author: ileasile

README.md

@@ -92,14 +92,9 @@ The following line magics are supported:
  - `%use <lib1>, <lib2> ...` - injects code for supported libraries: artifact resolution, default imports, initialization code, type renderers
  - `%trackClasspath` - logs any changes of current classpath. Useful for debugging artifact resolution failures
  - `%trackExecution` - logs pieces of code that are going to be executed. Useful for debugging of libraries support
- - `%output [--max-cell-size=N] [--max-buffer=N] [--max-buffer-newline=N] [--max-time=N] [--no-stdout] [--reset-to-defaults]` - 
- output capturing settings.
-     - `max-cell-size` specifies the characters count which may be printed to stdout. Default is 100000.
-     - `max-buffer` - max characters count stored in internal buffer before being sent to client. Default is 10000.
-     - `max-buffer-newline` - same as above, but trigger happens only if newline character was encountered. Default is 100.
-     - `max-time` - max time in milliseconds before the buffer is sent to client. Default is 100.
-     - `no-stdout` - don't capture output. Default is false.
-     - `reset-to-defaults` - reset all output settings that were set with magics to defaults
+ - `%output [options]` - output capturing settings.
+ 
+ See detailed info about line magics [here](doc/magics.md).
 
 ### Supported Libraries
 

doc/magics.md

@@ -0,0 +1,14 @@
+# Line magics
+The following line magics are supported:
+ - `%use <lib1>, <lib2> ...` - injects code for supported libraries: artifact resolution, default imports, initialization code, type renderers
+ - `%trackClasspath` - logs any changes of current classpath. Useful for debugging artifact resolution failures
+ - `%trackExecution` - logs pieces of code that are going to be executed. Useful for debugging of libraries support
+ - `%output [--max-cell-size=N] [--max-buffer=N] [--max-buffer-newline=N] [--max-time=N] [--no-stdout] [--reset-to-defaults]` - 
+ output capturing settings.
+     - `max-cell-size` specifies the characters count which may be printed to stdout. Default is 100000.
+     - `max-buffer` - max characters count stored in internal buffer before being sent to client. Default is 10000.
+     - `max-buffer-newline` - same as above, but trigger happens only if newline character was encountered. Default is 100.
+     - `max-time` - max time in milliseconds before the buffer is sent to client. Default is 100.
+     - `no-stdout` - don't capture output. Default is false.
+     - `reset-to-defaults` - reset all output settings that were set with magics to defaults
+ 

src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt

@@ -42,7 +42,7 @@ enum class JupyterSockets {
 
 data class OutputConfig(
         var captureOutput: Boolean = true,
-        var captureBufferTimeLimitMs: Int = 100,
+        var captureBufferTimeLimitMs: Long = 100,
         var captureBufferMaxSize: Int = 1000,
         var cellOutputMaxSize: Int = 100000,
         var captureNewlineBufferSize: Int = 100

src/main/kotlin/org/jetbrains/kotlin/jupyter/magics.kt

@@ -5,6 +5,7 @@ import com.github.ajalt.clikt.parameters.options.default
 import com.github.ajalt.clikt.parameters.options.flag
 import com.github.ajalt.clikt.parameters.options.option
 import com.github.ajalt.clikt.parameters.types.int
+import com.github.ajalt.clikt.parameters.types.long
 import org.jetbrains.kotlin.jupyter.repl.spark.ClassWriter
 
 enum class ReplLineMagics(val desc: String, val argumentsUsage: String? = null, val visibleInHelp: Boolean = true) {
@@ -28,7 +29,7 @@ fun processMagics(repl: ReplForJupyter, code: String): String {
             val max: Int by option("--max-cell-size", help = "Maximum cell output").int().default(conf.cellOutputMaxSize)
             val maxBuffer: Int by option("--max-buffer", help = "Maximum buffer size").int().default(conf.captureBufferMaxSize)
             val maxBufferNewline: Int by option("--max-buffer-newline", help = "Maximum buffer size when newline got").int().default(conf.captureNewlineBufferSize)
-            val maxTimeInterval: Int by option("--max-time", help = "Maximum time wait for output to accumulate").int().default(conf.captureBufferTimeLimitMs)
+            val maxTimeInterval: Long by option("--max-time", help = "Maximum time wait for output to accumulate").long().default(conf.captureBufferTimeLimitMs)
             val dontCaptureStdout: Boolean by option("--no-stdout", help = "Don't capture output").flag(default = !conf.captureOutput)
             val reset: Boolean by option("--reset-to-defaults", help = "Reset to defaults").flag()
 

src/main/kotlin/org/jetbrains/kotlin/jupyter/protocol.kt

@@ -4,12 +4,14 @@ import com.beust.klaxon.JsonObject
 import jupyter.kotlin.DisplayResult
 import jupyter.kotlin.MimeTypedResult
 import jupyter.kotlin.textResult
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.config.KotlinCompilerVersion
 import java.io.ByteArrayOutputStream
 import java.io.OutputStream
 import java.io.PrintStream
 import java.lang.reflect.InvocationTargetException
 import java.util.concurrent.atomic.AtomicLong
+import kotlin.concurrent.timer
 
 enum class ResponseState {
     Ok, Error
@@ -177,47 +179,70 @@ class CapturingOutputStream(private val stdout: PrintStream,
                             private val conf: OutputConfig,
                             private val captureOutput: Boolean,
                             val onCaptured: (String) -> Unit) : OutputStream() {
-    val capturedOutput = ByteArrayOutputStream()
-    private var time = System.currentTimeMillis()
+    private val capturedLines = ByteArrayOutputStream()
+    private val capturedNewLine = ByteArrayOutputStream()
     private var overallOutputSize = 0
     private var newlineFound = false
 
-    private fun shouldSend(b: Int): Boolean {
+    private val timer = timer(
+            initialDelay = conf.captureBufferTimeLimitMs,
+            period = conf.captureBufferTimeLimitMs,
+            action = {
+                flush()
+            })
+
+    val contents: ByteArray
+        @TestOnly
+        get() = capturedLines.toByteArray() + capturedNewLine.toByteArray()
+
+    private fun sendIfNeeded(b: Int) {
         val c = b.toChar()
-        newlineFound = newlineFound || c == '\n' || c == '\r'
-        if (newlineFound && capturedOutput.size() >= conf.captureNewlineBufferSize)
-            return true
-        if (capturedOutput.size() >= conf.captureBufferMaxSize)
-            return true
-
-        val currentTime = System.currentTimeMillis()
-        if (currentTime - time >= conf.captureBufferTimeLimitMs) {
-            time = currentTime
-            return true
+        if (c == '\n' || c == '\r') {
+            newlineFound = true
+            capturedNewLine.writeTo(capturedLines)
+            capturedNewLine.reset()
         }
-        return false
+
+        val size = capturedLines.size() + capturedNewLine.size()
+
+        if (newlineFound && size >= conf.captureNewlineBufferSize)
+            return flushBuffers(capturedLines)
+        if (size >= conf.captureBufferMaxSize)
+            return flush()
     }
 
     override fun write(b: Int) {
         ++overallOutputSize
         stdout.write(b)
 
         if (captureOutput && overallOutputSize <= conf.cellOutputMaxSize) {
-            capturedOutput.write(b)
-            if (shouldSend(b)) {
-                flush()
-            }
+            capturedNewLine.write(b)
+            sendIfNeeded(b)
         }
     }
 
-    override fun flush() {
+    private fun resetBuffer(buffer: ByteArrayOutputStream): String {
+        val str = buffer.toString("UTF-8")
+        buffer.reset()
+        return str
+    }
+
+    private fun flushBuffers(vararg buffers: ByteArrayOutputStream) {
         newlineFound = false
-        if (capturedOutput.size() > 0) {
-            val str = capturedOutput.toString("UTF-8")
-            capturedOutput.reset()
+        val str = buffers.map(this::resetBuffer).reduce { acc, s -> acc + s }
+        if (str.isNotEmpty()) {
             onCaptured(str)
         }
     }
+
+    override fun flush() {
+        flushBuffers(capturedLines, capturedNewLine)
+    }
+
+    override fun close() {
+        super.close()
+        timer.cancel()
+    }
 }
 
 fun Any.toMimeTypedResult(): MimeTypedResult? = when (this) {
@@ -301,7 +326,6 @@ fun JupyterConnection.evalWithIO(maybeConfig: OutputConfig?, body: () -> EvalRes
              */
             val stdErr = StringBuilder()
             with(stdErr) {
-                forkedError.capturedOutput.toString("UTF-8")?.nullWhenEmpty()?.also { appendln(it) }
                 val cause = ex.errorResult.cause
                 if (cause == null) appendln(ex.errorResult.message)
                 else {
@@ -318,10 +342,10 @@ fun JupyterConnection.evalWithIO(maybeConfig: OutputConfig?, body: () -> EvalRes
             ResponseWithMessage(ResponseState.Error, textResult("Error!"), emptyList(), null, stdErr.toString())
         }
     } finally {
+        forkedOut.close()
+        forkedError.close()
         System.setIn(`in`)
         System.setErr(err)
         System.setOut(out)
     }
 }
-
-fun String.nullWhenEmpty(): String? = if (this.isBlank()) null else this

src/test/kotlin/org/jetbrains/kotlin/jupyter/test/capturingStreamTests.kt

@@ -16,7 +16,7 @@ class CapturingStreamTests {
 
     private fun getStream(stdout: OutputStream = nullOStream,
                           captureOutput: Boolean = true,
-                          maxBufferLifeTimeMs: Int = 1000,
+                          maxBufferLifeTimeMs: Long = 1000,
                           maxBufferSize: Int = 1000,
                           maxOutputSize: Int = 1000,
                           maxBufferNewlineSize: Int = 1,
@@ -37,7 +37,7 @@ class CapturingStreamTests {
     fun testMaxOutputSizeError() {
         val s = getStream(maxOutputSize = 3)
         s.write("java".toByteArray())
-        assertArrayEquals("jav".toByteArray(), s.capturedOutput.toByteArray())
+        assertArrayEquals("jav".toByteArray(), s.contents)
     }
 
     @Test
@@ -46,11 +46,11 @@ class CapturingStreamTests {
 
         val s1 = getStream(captureOutput = false)
         s1.write(contents)
-        assertEquals(0, s1.capturedOutput.size())
+        assertEquals(0, s1.contents.size)
 
         val s2 = getStream(captureOutput = true)
         s2.write(contents)
-        assertArrayEquals(contents, s2.capturedOutput.toByteArray())
+        assertArrayEquals(contents, s2.contents)
     }
 
     @Test
@@ -73,7 +73,7 @@ class CapturingStreamTests {
     @Test
     fun testNewlineBufferSize() {
         val contents = "12345\n12\n3451234567890".toByteArray()
-        val expected = arrayOf("12345\n", "12\n345", "123456789", "0")
+        val expected = arrayOf("12345\n", "12\n", "345123456", "7890")
 
         var i = 0
         val s = getStream(maxBufferSize = 9, maxBufferNewlineSize = 6) {
@@ -89,8 +89,8 @@ class CapturingStreamTests {
 
     @Test
     fun testMaxBufferLifeTime() {
-        val strings = arrayOf("c ", "a", "ada ", "b", "scala ", "c")
-        val expected = arrayOf("c a", "ada b", "scala c")
+        val strings = arrayOf("11", "22", "33", "44", "55", "66")
+        val expected = arrayOf("1122", "3344", "5566")
 
         var i = 0
         val s = getStream(maxBufferLifeTimeMs = 1000) {
@@ -99,7 +99,7 @@ class CapturingStreamTests {
         }
 
         strings.forEach {
-            Thread.sleep(600)
+            Thread.sleep(450)
             s.write(it.toByteArray())
         }
 

src/test/kotlin/org/jetbrains/kotlin/jupyter/test/executeTests.kt

@@ -12,7 +12,6 @@ fun Message.type(): String {
 
 class ExecuteTests : KernelServerTestsBase() {
 
-    @Synchronized
     private fun doExecute(code : String, hasResult: Boolean = true, ioPubChecker : (ZMQ.Socket) -> Unit = {}) : Any? {
         val context = ZMQ.context(1)
         val shell = context.socket(ZMQ.REQ)