Merge pull request #36393 from zeertzjq/rstream-close-cb

fix(channel): closing socket with pending writes leaks memory

src/nvim/channel.c

@@ -24,6 +24,7 @@
 #include "nvim/event/proc.h"
 #include "nvim/event/rstream.h"
 #include "nvim/event/socket.h"
+#include "nvim/event/stream.h"
 #include "nvim/event/wstream.h"
 #include "nvim/ex_cmds.h"
 #include "nvim/garray.h"
@@ -131,7 +132,7 @@ bool channel_close(uint64_t id, ChannelPart part, const char **error)
   case kChannelStreamProc:
     proc = &chan->stream.proc;
     if (part == kChannelPartStdin || close_main) {
-      wstream_may_close(&proc->in);
+      stream_may_close(&proc->in);
     }
     if (part == kChannelPartStdout || close_main) {
       rstream_may_close(&proc->out);
@@ -150,7 +151,7 @@ bool channel_close(uint64_t id, ChannelPart part, const char **error)
       rstream_may_close(&chan->stream.stdio.in);
     }
     if (part == kChannelPartStdout || close_main) {
-      wstream_may_close(&chan->stream.stdio.out);
+      stream_may_close(&chan->stream.stdio.out);
     }
     if (part == kChannelPartStderr) {
       *error = e_invstream;

src/nvim/event/proc.c

@@ -150,7 +150,7 @@ void proc_teardown(Loop *loop) FUNC_ATTR_NONNULL_ALL
 
 void proc_close_streams(Proc *proc) FUNC_ATTR_NONNULL_ALL
 {
-  wstream_may_close(&proc->in);
+  stream_may_close(&proc->in);
   rstream_may_close(&proc->out);
   rstream_may_close(&proc->err);
 }

src/nvim/event/rstream.c

@@ -36,6 +36,8 @@ void rstream_init(RStream *stream)
   stream->num_bytes = 0;
   stream->buffer = alloc_block();
   stream->read_pos = stream->write_pos = stream->buffer;
+  stream->s.close_cb = rstream_close_cb;
+  stream->s.close_cb_data = stream;
 }
 
 void rstream_start_inner(RStream *stream)
@@ -183,7 +185,7 @@ static void read_event(void **argv)
   stream->s.pending_reqs--;
   if (stream->s.closed && !stream->s.pending_reqs) {
     // Last pending read; free the stream.
-    stream_close_handle(&stream->s, true);
+    stream_close_handle(&stream->s);
   }
 }
 
@@ -231,7 +233,16 @@ static void invoke_read_cb(RStream *stream, bool eof)
   CREATE_EVENT(stream->s.events, read_event, stream);
 }
 
+static void rstream_close_cb(Stream *s, void *data)
+{
+  RStream *stream = data;
+  assert(stream && s == &stream->s);
+  if (stream->buffer) {
+    free_block(stream->buffer);
+  }
+}
+
 void rstream_may_close(RStream *stream)
 {
-  stream_may_close(&stream->s, true);
+  stream_may_close(&stream->s);
 }

src/nvim/event/stream.c

@@ -96,7 +96,7 @@ void stream_init(Loop *loop, Stream *stream, int fd, uv_stream_t *uvstream)
   stream->events = NULL;
 }
 
-void stream_may_close(Stream *stream, bool rstream)
+void stream_may_close(Stream *stream)
   FUNC_ATTR_NONNULL_ARG(1)
 {
   if (stream->closed) {
@@ -104,10 +104,6 @@ void stream_may_close(Stream *stream, bool rstream)
   }
   DLOG("closing Stream: %p", (void *)stream);
   stream->closed = true;
-  // TODO(justinmk): stream->close_cb is never actually invoked. Either remove it, or see if it can
-  // be used somewhere...
-  stream->close_cb = NULL;
-  stream->close_cb_data = NULL;
 
 #ifdef MSWIN
   if (UV_TTY == uv_guess_handle(stream->fd)) {
@@ -117,11 +113,11 @@ void stream_may_close(Stream *stream, bool rstream)
 #endif
 
   if (!stream->pending_reqs) {
-    stream_close_handle(stream, rstream);
+    stream_close_handle(stream);
   }  // Else: rstream.c:read_event() or wstream.c:write_cb() will call stream_close_handle().
 }
 
-void stream_close_handle(Stream *stream, bool rstream)
+void stream_close_handle(Stream *stream)
   FUNC_ATTR_NONNULL_ALL
 {
   uv_handle_t *handle = NULL;
@@ -139,19 +135,10 @@ void stream_close_handle(Stream *stream, bool rstream)
   assert(handle != NULL);
 
   if (!uv_is_closing(handle)) {
-    uv_close(handle, rstream ? rstream_close_cb : close_cb);
+    uv_close(handle, close_cb);
   }
 }
 
-static void rstream_close_cb(uv_handle_t *handle)
-{
-  RStream *stream = handle->data;
-  if (stream && stream->buffer) {
-    free_block(stream->buffer);
-  }
-  close_cb(handle);
-}
-
 static void close_cb(uv_handle_t *handle)
 {
   Stream *stream = handle->data;

src/nvim/event/wstream.c

@@ -155,7 +155,7 @@ static void write_cb(uv_write_t *req, int status)
 
   if (data->stream->closed && data->stream->pending_reqs == 0) {
     // Last pending write; free the stream.
-    stream_close_handle(data->stream, false);
+    stream_close_handle(data->stream);
   }
 
   xfree(data);
@@ -172,8 +172,3 @@ void wstream_release_wbuffer(WBuffer *buffer)
     xfree(buffer);
   }
 }
-
-void wstream_may_close(Stream *stream)
-{
-  stream_may_close(stream, false);
-}

src/nvim/os/shell.c

@@ -1271,7 +1271,7 @@ static void shell_write_cb(Stream *stream, void *data, int status)
     msg_schedule_semsg(_("E5677: Error writing input to shell-command: %s"),
                        uv_err_name(status));
   }
-  stream_may_close(stream, false);
+  stream_may_close(stream);
 }
 
 /// Applies 'shellxescape' (p_sxe) and 'shellxquote' (p_sxq) to a command.

test/functional/api/server_requests_spec.lua

@@ -305,6 +305,16 @@ describe('server -> client', function()
       eq(clientpid, fn.getpid())
       eq('howdy!', api.nvim_get_current_line())
 
+      -- sending notification and then closing channel immediately still works
+      n.exec_lua(function()
+        vim.rpcnotify(id, 'nvim_set_current_line', 'bye!')
+        vim.fn.chanclose(id)
+      end)
+
+      set_session(server)
+      eq(serverpid, fn.getpid())
+      eq('bye!', api.nvim_get_current_line())
+
       server:close()
       client:close()
     end

test/functional/ui/embed_spec.lua

@@ -282,6 +282,19 @@ describe('--embed UI', function()
       end,
     }
   end)
+
+  it('closing stdio with another remote UI does not leak memory #36392', function()
+    t.skip(t.is_os('win')) -- n.connect() hangs on Windows
+    clear({ args_rm = { '--headless' } })
+    Screen.new()
+    eq(1, #api.nvim_list_uis())
+    local server = api.nvim_get_vvar('servername')
+    local other_session = n.connect(server)
+    Screen.new(nil, nil, nil, other_session)
+    eq(2, #api.nvim_list_uis())
+    check_close()
+    other_session:close()
+  end)
 end)
 
 describe('--embed --listen UI', function()