From e9699c4d1194559df147fe069467fa9a56b4ad71 Mon Sep 17 00:00:00 2001
From: Tamas Nagy <tamas@tamasnagy.com>
Date: Tue, 20 Sep 2016 20:37:23 -0700
Subject: [PATCH] truncate excess output per execution request (#463)

* truncate excess output per execution request

fixes https://github.com/timholy/Images.jl/issues/558, fixes #455

* apply to both stderr and stdout; rename accessor function
---
 README.md              | 10 ++++++++++
 src/IJulia.jl          | 13 +++++++++++++
 src/execute_request.jl |  4 ++++
 src/stdio.jl           | 26 +++++++++++++++++++++-----
 4 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 3de0c56b..547f86d5 100644
--- a/README.md
+++ b/README.md
@@ -250,3 +250,13 @@ your `.julia` package directory) to change the line `verbose = false`
 at the top to `verbose = true` and `const capture_stderr = true` to
 `const capture_stderr = false`.  Then restart the kernel or open a new
 notebook and look for the error message when IJulia dies.
+
+### Preventing truncation of output
+
+The new default behavior of IJulia is to truncate stdout (via `show` or `println`)
+after 512kb. This to prevent browsers from getting bogged down when displaying the
+results. This limit can be increased to a custom value, like 1MB, as follows
+
+```
+IJulia.set_max_stdio(1 << 20)
+```
diff --git a/src/IJulia.jl b/src/IJulia.jl
index ddd946ca..371c5913 100644
--- a/src/IJulia.jl
+++ b/src/IJulia.jl
@@ -284,6 +284,19 @@ function clear_output(wait=false)
                                     Dict("wait" => wait)))
 end
 
+
+"""
+    set_max_stdio(max_output::Integer)
+
+Sets the maximum number of bytes, `max_output`, that can be written to stdout and
+stderr before getting truncated. A large value here allows a lot of output to be
+displayed in the notebook, potentially bogging down the browser.
+"""
+function set_max_stdio(max_output::Integer)
+    max_output_per_request[] = max_output
+end
+
+
 #######################################################################
 
 include("init.jl")
diff --git a/src/execute_request.jl b/src/execute_request.jl
index 734e3659..47abb43b 100644
--- a/src/execute_request.jl
+++ b/src/execute_request.jl
@@ -111,6 +111,9 @@ end
 
 # global variable so that display can be done in the correct Msg context
 execute_msg = Msg(["julia"], Dict("username"=>"julia", "session"=>"????"), Dict())
+# global variable tracking the number of bytes written in the current execution
+# request
+const stdio_bytes = Ref(0)
 
 function helpcode(code::AbstractString)
     code_ = strip(code)
@@ -132,6 +135,7 @@ function execute_request(socket, msg)
     @vprintln("EXECUTING ", code)
     global execute_msg = msg
     global n, In, Out, ans
+    stdio_bytes[] = 0
     silent = msg.content["silent"]
     store_history = get(msg.content, "store_history", !silent)
     empty!(execute_payloads)
diff --git a/src/stdio.jl b/src/stdio.jl
index 272d5e09..5dd43aa2 100644
--- a/src/stdio.jl
+++ b/src/stdio.jl
@@ -31,12 +31,18 @@ end
 #name=>iobuffer for each stream ("stdout","stderr") so they can be sent in flush
 const bufs = Dict{String,IOBuffer}()
 const stream_interval = 0.1
+# maximum number of bytes in libuv/os buffer before emptying
 const max_bytes = 10*1024
+# max output per code cell is 512 kb by default
+const max_output_per_request = Ref(1 << 19)
 
-"""Continually read from (size limited) Libuv/OS buffer into an (effectively unlimited) `IObuffer`
-to avoid problems when the Libuv/OS buffer gets full (https://github.com/JuliaLang/julia/issues/8789).
-Send data immediately when buffer contains more than `max_bytes` bytes. Otherwise, if data is available
-it will be sent every `stream_interval` seconds (see the Timers set up in watch_stdio)"""
+"""
+Continually read from (size limited) Libuv/OS buffer into an `IObuffer` to avoid problems when
+the Libuv/OS buffer gets full (https://github.com/JuliaLang/julia/issues/8789). Send data immediately
+when buffer contains more than `max_bytes` bytes. Otherwise, if data is available it will be sent every
+`stream_interval` seconds (see the Timers set up in watch_stdio). Truncate the output to `max_output_per_request`
+bytes per execution request since excessive output can bring browsers to a grinding halt.
+"""
 function watch_stream(rd::IO, name::AbstractString)
     task_local_storage(:IJulia_task, "read $name task")
     try
@@ -45,7 +51,17 @@ function watch_stream(rd::IO, name::AbstractString)
         while !eof(rd) # blocks until something is available
             nb = nb_available(rd)
             if nb > 0
-                write(buf, read(rd, nb))
+                stdio_bytes[] += nb
+                # if this stream has surpassed the maximum output limit then ignore future bytes
+                if stdio_bytes[] >= max_output_per_request[]
+                    read(rd, nb) # read from libuv/os buffer and discard
+                    if stdio_bytes[] - nb < max_output_per_request[]
+                        send_ipython(publish[], msg_pub(execute_msg, "stream",
+                            @compat Dict("name" => "stderr", "text" => "Excessive output truncated after $(stdio_bytes[]) bytes.")))
+                    end
+                else
+                    write(buf, read(rd, nb))
+                end
             end
             if buf.size > 0
                 if buf.size >= max_bytes