5 files changed, 53 insertions, 35 deletions

diff --git a/src/nvim/os/channel.c b/src/nvim/os/channel.c
index 2923ab0912..2584340950 100644
--- a/src/nvim/os/channel.c
+++ b/src/nvim/os/channel.c
@@ -107,12 +107,11 @@ bool channel_send_event(uint64_t id, char *type, typval_T *data)
   msgpack_packer packer;
   msgpack_packer_init(&packer, &msgpack_event_buffer, msgpack_sbuffer_write);
   msgpack_rpc_notification(event_type, event_data, &packer);
-  char *bytes = xmemdup(msgpack_event_buffer.data, msgpack_event_buffer.size);
 
   wstream_write(channel->data.streams.write,
-                bytes,
-                msgpack_event_buffer.size,
-                true);
+                wstream_new_buffer(msgpack_event_buffer.data,
+                                   msgpack_event_buffer.size,
+                                   true));
 
   msgpack_rpc_free_object(event_data);
   msgpack_sbuffer_clear(&msgpack_event_buffer);
@@ -158,9 +157,9 @@ static void parse_msgpack(RStream *rstream, void *data, bool eof)
     // Perform the call
     msgpack_rpc_call(channel->id, &unpacked.data, &response);
     wstream_write(channel->data.streams.write,
-                  xmemdup(channel->sbuffer->data, channel->sbuffer->size),
-                  channel->sbuffer->size,
-                  true);
+                  wstream_new_buffer(channel->sbuffer->data,
+                                     channel->sbuffer->size,
+                                     true));
 
     // Clear the buffer for future calls
     msgpack_sbuffer_clear(channel->sbuffer);
diff --git a/src/nvim/os/job.c b/src/nvim/os/job.c
index ac79a5b882..c4a9c85d1d 100644
--- a/src/nvim/os/job.c
+++ b/src/nvim/os/job.c
@@ -234,7 +234,7 @@ bool job_write(int id, char *data, uint32_t len)
     return false;
   }
 
-  if (!wstream_write(job->in, data, len, true)) {
+  if (!wstream_write(job->in, wstream_new_buffer(data, len, false))) {
     job_stop(job->id);
     return false;
   }
diff --git a/src/nvim/os/wstream.c b/src/nvim/os/wstream.c
index 22524e2430..57afdd0e8f 100644
--- a/src/nvim/os/wstream.c
+++ b/src/nvim/os/wstream.c
@@ -20,14 +20,14 @@ struct wstream {
   bool freed;
 };
 
+struct wbuffer {
+  size_t refcount, size;
+  char *data;
+};
+
 typedef struct {
   WStream *wstream;
-  // Buffer containing data to be written
-  char *buffer;
-  // Size of the buffer
-  size_t length;
-  // If it's our responsibility to free the buffer
-  bool free;
+  WBuffer *buffer;
 } WriteData;
 
 static void write_cb(uv_write_t *req, int status);
@@ -59,51 +59,60 @@ void wstream_set_stream(WStream *wstream, uv_stream_t *stream)
   wstream->stream = stream;
 }
 
-bool wstream_write(WStream *wstream, char *buffer, size_t length, bool free)
+bool wstream_write(WStream *wstream, WBuffer *buffer)
 {
   WriteData *data;
   uv_buf_t uvbuf;
   uv_write_t *req;
 
-  if (wstream->freed) {
-    // Don't accept write requests after the WStream instance was freed
-    return false;
-  }
+  // This should not be called after a wstream was freed
+  assert(!wstream->freed);
 
-  if (wstream->curmem + length > wstream->maxmem) {
+  if (wstream->curmem + buffer->size > wstream->maxmem) {
     return false;
   }
 
-  if (free) {
-    // We should only account for buffers that are ours to free
-    wstream->curmem += length;
-  }
-
+  buffer->refcount++;
+  wstream->curmem += buffer->size;
   data = xmalloc(sizeof(WriteData));
   data->wstream = wstream;
   data->buffer = buffer;
-  data->length = length;
-  data->free = free;
   req = xmalloc(sizeof(uv_write_t));
   req->data = data;
-  uvbuf.base = buffer;
-  uvbuf.len = length;
+  uvbuf.base = buffer->data;
+  uvbuf.len = buffer->size;
   wstream->pending_reqs++;
   uv_write(req, wstream->stream, &uvbuf, 1, write_cb);
 
   return true;
 }
 
+WBuffer *wstream_new_buffer(char *data, size_t size, bool copy)
+{
+  WBuffer *rv = xmalloc(sizeof(WBuffer));
+  rv->size = size;
+  rv->refcount = 0;
+
+  if (copy) {
+    rv->data = xmemdup(data, size);
+  } else {
+    rv->data = data;
+  }
+
+  return rv;
+}
+
 static void write_cb(uv_write_t *req, int status)
 {
   WriteData *data = req->data;
 
   free(req);
+  data->wstream->curmem -= data->buffer->size;
 
-  if (data->free) {
+  if (!--data->buffer->refcount) {
     // Free the data written to the stream
+    free(data->buffer->data);
     free(data->buffer);
-    data->wstream->curmem -= data->length;
   }
 
   data->wstream->pending_reqs--;
diff --git a/src/nvim/os/wstream.h b/src/nvim/os/wstream.h
index 2cef16e6cd..1f61f6afd0 100644
--- a/src/nvim/os/wstream.h
+++ b/src/nvim/os/wstream.h
@@ -31,10 +31,19 @@ void wstream_set_stream(WStream *wstream, uv_stream_t *stream);
 ///
 /// @param wstream The `WStream` instance
 /// @param buffer The buffer which contains data to be written
-/// @param length Number of bytes that should be written from `buffer`
-/// @param free If true, `buffer` will be freed after the write is complete
-/// @return true if the data was successfully queued, false otherwise.
-bool wstream_write(WStream *wstream, char *buffer, size_t length, bool free);
+/// @return false if the write failed
+bool wstream_write(WStream *wstream, WBuffer *buffer);
+
+/// Creates a WBuffer object for holding output data. Instances of this
+/// object can be reused across WStream instances, and the memory is freed
+/// automatically when no longer needed(it tracks the number of references
+/// internally)
+///
+/// @param data Data stored by the WBuffer
+/// @param size The size of the data array
+/// @param copy If true, the data will be copied into the WBuffer
+/// @return The allocated WBuffer instance
+WBuffer *wstream_new_buffer(char *data, size_t size, bool copy);
 
 #endif  // NVIM_OS_WSTREAM_H
 
diff --git a/src/nvim/os/wstream_defs.h b/src/nvim/os/wstream_defs.h
index e044ecfd8e..a7565c9bc7 100644
--- a/src/nvim/os/wstream_defs.h
+++ b/src/nvim/os/wstream_defs.h
@@ -1,6 +1,7 @@
 #ifndef NVIM_OS_WSTREAM_DEFS_H
 #define NVIM_OS_WSTREAM_DEFS_H
 
+typedef struct wbuffer WBuffer;
 typedef struct wstream WStream;
 
 #endif  // NVIM_OS_WSTREAM_DEFS_H