Merge pull request #25200 from JuliaLang/jn/flush-uv

implement flush for libuv streams

Author: vtjnash

base/stream.jl

@@ -55,7 +55,17 @@ function eof(s::LibuvStream)
     return !isopen(s) && nb_available(s) <= 0
 end
 
-const DEFAULT_READ_BUFFER_SZ = 10485760           # 10 MB
+# Limit our default maximum read and buffer size,
+# to avoid DoS-ing ourself into an OOM situation
+const DEFAULT_READ_BUFFER_SZ = 10485760 # 10 MB
+
+# manually limit our write size, if the OS doesn't support full-size writes
+if Sys.iswindows()
+    const MAX_OS_WRITE = UInt(0x1FF0_0000) # 511 MB (determined semi-empirically, limited to 31 MB on XP)
+else
+    const MAX_OS_WRITE = UInt(typemax(Csize_t))
+end
+
 
 const StatusUninit      = 0 # handle is allocated, but not initialized
 const StatusInit        = 1 # handle is valid, but not connected/active
@@ -458,7 +468,10 @@ function notify_filled(buffer::IOBuffer, nread::Int)
     end
 end
 
-alloc_buf_hook(stream::LibuvStream, size::UInt) = alloc_request(stream.buffer, UInt(size))
+function alloc_buf_hook(stream::LibuvStream, size::UInt)
+    throttle = UInt(stream.throttle)
+    return alloc_request(stream.buffer, (size > throttle) ? throttle : size)
+end
 
 function uv_alloc_buf(handle::Ptr{Cvoid}, size::Csize_t, buf::Ptr{Cvoid})
     hd = uv_handle_data(handle)
@@ -477,6 +490,10 @@ function uv_alloc_buf(handle::Ptr{Cvoid}, size::Csize_t, buf::Ptr{Cvoid})
         if data == C_NULL
             newsize = 0
         end
+        # avoid aliasing of `nread` with `errno` in uv_readcb
+        # or exceeding the Win32 maximum uv_buf_t len
+        maxsize = @static Sys.iswindows() ? typemax(Cint) : typemax(Cssize_t)
+        newsize > maxsize && (newsize = maxsize)
     end
 
     ccall(:jl_uv_buf_set_base, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}), buf, data)
@@ -815,28 +832,21 @@ function readuntil(this::LibuvStream, c::UInt8)
 end
 
 uv_write(s::LibuvStream, p::Vector{UInt8}) = uv_write(s, pointer(p), UInt(sizeof(p)))
+
 function uv_write(s::LibuvStream, p::Ptr{UInt8}, n::UInt)
-    check_open(s)
-    uvw = Libc.malloc(_sizeof_uv_write)
-    uv_req_set_data(uvw, C_NULL) # in case we get interrupted before arriving at the wait call
-    err = ccall(:jl_uv_write,
-                Int32,
-                (Ptr{Cvoid}, Ptr{Cvoid}, UInt, Ptr{Cvoid}, Ptr{Cvoid}),
-                s, p, n, uvw,
-                uv_jl_writecb_task::Ptr{Cvoid})
-    if err < 0
-        Libc.free(uvw)
-        uv_error("write", err)
-    end
+    uvw = uv_write_async(s, p, n)
     ct = current_task()
     preserve_handle(ct)
     try
+        # wait for the last chunk to complete (or error)
+        # assume that any errors would be sticky,
+        # (so we don't need to monitor the error status of the intermediate writes)
         uv_req_set_data(uvw, ct)
         wait()
     finally
         if uv_req_data(uvw) != C_NULL
             # uvw is still alive,
-            # so make sure we don't get spurious notifications later
+            # so make sure we won't get spurious notifications later
             uv_req_set_data(uvw, C_NULL)
         else
             # done with uvw
@@ -847,6 +857,33 @@ function uv_write(s::LibuvStream, p::Ptr{UInt8}, n::UInt)
     return Int(n)
 end
 
+# helper function for uv_write that returns the uv_write_t struct for the write
+# rather than waiting on it
+function uv_write_async(s::LibuvStream, p::Ptr{UInt8}, n::UInt)
+    check_open(s)
+    while true
+        uvw = Libc.malloc(_sizeof_uv_write)
+        uv_req_set_data(uvw, C_NULL) # in case we get interrupted before arriving at the wait call
+        nwrite = min(n, MAX_OS_WRITE) # split up the write into chunks the OS can handle.
+        # TODO: use writev, when that is added to uv-win
+        err = ccall(:jl_uv_write,
+                    Int32,
+                    (Ptr{Cvoid}, Ptr{Cvoid}, UInt, Ptr{Cvoid}, Ptr{Cvoid}),
+                    s, p, nwrite, uvw,
+                    uv_jl_writecb_task::Ptr{Cvoid})
+        if err < 0
+            Libc.free(uvw)
+            uv_error("write", err)
+        end
+        n -= nwrite
+        p += nwrite
+        if n == 0
+            return uvw
+        end
+    end
+end
+
+
 # Optimized send
 # - smaller writes are buffered, final uv write on flush or when buffer full
 # - large isbits arrays are unbuffered and written directly
@@ -872,14 +909,15 @@ function unsafe_write(s::LibuvStream, p::Ptr{UInt8}, n::UInt)
 end
 
 function flush(s::LibuvStream)
-    if s.sendbuf === nothing
-        return
-    end
-    buf = s.sendbuf
-    if nb_available(buf) > 0
-        arr = take!(buf)        # Array of UInt8s
-        uv_write(s, arr)
+    if s.sendbuf !== nothing
+        buf = s.sendbuf
+        if nb_available(buf) > 0
+            arr = take!(buf)        # Array of UInt8s
+            uv_write(s, arr)
+            return
+        end
     end
+    uv_write(s, Ptr{UInt8}(Base.eventloop()), UInt(0)) # zero write from a random pointer to flush current queue
     return
 end
 
@@ -900,12 +938,13 @@ end
 function uv_writecb_task(req::Ptr{Cvoid}, status::Cint)
     d = uv_req_data(req)
     if d != C_NULL
-        uv_req_set_data(req, C_NULL)
+        uv_req_set_data(req, C_NULL) # let the Task know we got the writecb
+        t = unsafe_pointer_to_objref(d)::Task
         if status < 0
             err = UVError("write", status)
-            schedule(unsafe_pointer_to_objref(d)::Task, err, error=true)
+            schedule(t, err, error=true)
         else
-            schedule(unsafe_pointer_to_objref(d)::Task)
+            schedule(t)
         end
     else
         # no owner for this req, safe to just free it

src/gf.c

@@ -1301,6 +1301,7 @@ static void method_overwrite(jl_typemap_entry_t *newentry, jl_method_t *oldvalue
             jl_printf(s, " in module %s", jl_symbol_name(newmod->name));
         print_func_loc(s, method);
         jl_printf(s, ".\n");
+        jl_uv_flush(s);
     }
 }
 

src/jl_uv.c

@@ -98,19 +98,72 @@ static void jl_uv_closeHandle(uv_handle_t *handle)
     free(handle);
 }
 
-static void jl_uv_shutdownCallback(uv_shutdown_t *req, int status)
-{
-    /*
-     * This happens if the remote machine closes the connecition while we're
-     * in the shutdown request (in that case we call uv_close, thus cancelling
-     * this request).
-     */
-    if (status != UV__ECANCELED && !uv_is_closing((uv_handle_t*)req->handle)) {
-        uv_close((uv_handle_t*)req->handle, &jl_uv_closeHandle);
+static void jl_uv_flush_close_callback(uv_write_t *req, int status)
+{
+    uv_stream_t *stream = req->handle;
+    req->handle = NULL;
+    // ignore attempts to close the stream while attempting a graceful shutdown
+#ifdef _OS_WINDOWS_
+    if (stream->stream.conn.shutdown_req)
+#else
+    if (stream->shutdown_req)
+#endif
+    {
+        free(req);
+        return;
+    }
+    if (status == 0 && uv_is_writable(stream) && stream->write_queue_size != 0) {
+        // new data was written, wait for it to flush too
+        uv_buf_t buf;
+        buf.base = (char*)(req + 1);
+        buf.len = 0;
+        req->data = NULL;
+        if (uv_write(req, stream, &buf, 1, (uv_write_cb)jl_uv_flush_close_callback) == 0)
+            return;
+    }
+    if (!uv_is_closing((uv_handle_t*)stream)) { // avoid double-close on the stream
+        if (stream->type == UV_TTY)
+            uv_tty_set_mode((uv_tty_t*)stream, UV_TTY_MODE_NORMAL);
+        uv_close((uv_handle_t*)stream, &jl_uv_closeHandle);
     }
     free(req);
 }
 
+static void uv_flush_callback(uv_write_t *req, int status)
+{
+    *(int*)(req->data) = 1;
+    uv_stop(req->handle->loop);
+    free(req);
+}
+
+// Turn a normal write into a blocking write (primarly for use from C and gdb).
+// Warning: This calls uv_run, so it can have unbounded side-effects.
+// Be care where you call it from! - the libuv loop is also not reentrant.
+void jl_uv_flush(uv_stream_t *stream)
+{
+    if (stream == (void*)STDIN_FILENO ||
+        stream == (void*)STDOUT_FILENO ||
+        stream == (void*)STDERR_FILENO)
+        return;
+    if (stream->type != UV_TTY &&
+        stream->type != UV_TCP &&
+        stream->type != UV_NAMED_PIPE)
+        return;
+    while (uv_is_writable(stream) && stream->write_queue_size != 0) {
+        int fired = 0;
+	uv_buf_t buf;
+	buf.base = (char*)(&buf + 1);
+	buf.len = 0;
+        uv_write_t *write_req = (uv_write_t*)malloc(sizeof(uv_write_t));
+        write_req->data = (void*)&fired;
+        if (uv_write(write_req, stream, &buf, 1, uv_flush_callback) != 0)
+            return;
+        while (!fired) {
+            uv_run(uv_default_loop(), UV_RUN_DEFAULT);
+        }
+    }
+}
+
 // getters and setters
 JL_DLLEXPORT void *jl_uv_process_data(uv_process_t *p) { return p->data; }
 JL_DLLEXPORT void *jl_uv_buf_base(const uv_buf_t *buf) { return buf->base; }
@@ -196,39 +249,15 @@ JL_DLLEXPORT void jl_close_uv(uv_handle_t *handle)
         return;
     }
 
-    if (handle->type == UV_NAMED_PIPE || handle->type == UV_TCP) {
-        uv_stream_t *stream = (uv_stream_t*)handle;
-#ifdef _OS_WINDOWS_
-        if (stream->stream.conn.shutdown_req) {
-#else
-        if (stream->shutdown_req) {
-#endif
-            // don't close the stream while attempting a graceful shutdown
-            return;
-        }
-        if (uv_is_writable(stream) && stream->write_queue_size != 0) {
-            // attempt graceful shutdown of writable streams to give them a chance to flush first
-            // TODO: introduce a uv_drain cb API instead of abusing uv_shutdown in this way
-            uv_shutdown_t *req = (uv_shutdown_t*)malloc(sizeof(uv_shutdown_t));
-            req->data = 0;
-            /*
-             * We are explicitly ignoring the error here for the following reason:
-             * There is only two scenarios in which this returns an error:
-             * a) In case the stream is already shut down, in which case we're likely
-             *    in the process of closing this stream (since there's no other call to
-             *    uv_shutdown).
-             * b) In case the stream is already closed, in which case uv_close would
-             *    cause an assertion failure.
-             */
-            uv_shutdown(req, stream, &jl_uv_shutdownCallback);
-            return;
-        }
+    if (handle->type == UV_NAMED_PIPE || handle->type == UV_TCP || handle->type == UV_TTY) {
+        uv_write_t *req = (uv_write_t*)malloc(sizeof(uv_write_t));
+        req->handle = (uv_stream_t*)handle;
+        jl_uv_flush_close_callback(req, 0);
+        return;
     }
 
     if (!uv_is_closing(handle)) {
         // avoid double-closing the stream
-        if (handle->type == UV_TTY)
-            uv_tty_set_mode((uv_tty_t*)handle, UV_TTY_MODE_NORMAL);
         uv_close(handle, &jl_uv_closeHandle);
     }
 }

src/julia_internal.h

@@ -462,6 +462,7 @@ extern uv_loop_t *jl_io_loop;
 JL_DLLEXPORT void jl_uv_associate_julia_struct(uv_handle_t *handle,
                                                jl_value_t *data);
 JL_DLLEXPORT int jl_uv_fs_result(uv_fs_t *f);
+void jl_uv_flush(uv_stream_t *stream);
 
 typedef struct _typeenv {
     jl_tvar_t *var;

src/toplevel.c

@@ -816,7 +816,7 @@ JL_DLLEXPORT jl_value_t *jl_load(jl_module_t *module, const char *fname)
     if (module->istopmod) {
         jl_printf(JL_STDOUT, "%s\r\n", fname);
 #ifdef _OS_WINDOWS_
-        uv_run(uv_default_loop(), (uv_run_mode)1);
+        jl_uv_flush(JL_STDOUT);
 #endif
     }
     uv_stat_t stbuf;

test/read.jl

@@ -537,3 +537,18 @@ end # mktempdir() do dir
     @test countlines(file,'\n') == 4
     rm(file)
 end
+
+let p = Pipe()
+    Base.link_pipe(p, julia_only_read=true, julia_only_write=true)
+    t = @schedule read(p)
+    @sync begin
+        @async write(p, zeros(UInt16, 660_000))
+        for i = 1:typemax(UInt16)
+            @async write(p, UInt16(i))
+        end
+        @async close(p.in)
+    end
+    s = reinterpret(UInt16, wait(t))
+    @test length(s) == 660_000 + typemax(UInt16)
+    @test s[(end - typemax(UInt16)):end] == UInt16.(0:typemax(UInt16))
+end