Extract writing boilerplate into wstream.c module

5 files changed, 162 insertions, 7 deletions

diff --git a/src/os/rstream.h b/src/os/rstream.h
index 2ca85bdf23..4678889238 100644
--- a/src/os/rstream.h
+++ b/src/os/rstream.h
@@ -36,12 +36,6 @@ void rstream_free(RStream *rstream);
 /// @param stream The new `uv_stream_t` instance
 void rstream_set_stream(RStream *rstream, uv_stream_t *stream);
 
-/// Sets the underlying `uv_file_t` instance
-///
-/// @param rstream The `RStream` instance
-/// @param stream The new `uv_stream_t` instance
-void rstream_set_stream(RStream *rstream, uv_stream_t *stream);
-
 /// Sets the underlying file descriptor that will be read from. Only pipes
 /// and regular files are supported for now.
 ///
diff --git a/src/os/rstream_defs.h b/src/os/rstream_defs.h
index 62c910d041..ca3035adbf 100644
--- a/src/os/rstream_defs.h
+++ b/src/os/rstream_defs.h
@@ -3,7 +3,7 @@
 
 typedef struct rstream RStream;
 
-/// Function called when the RStream receives data
+/// Type of function called when the RStream receives data
 ///
 /// @param rstream The RStream instance
 /// @param data State associated with the RStream instance
diff --git a/src/os/wstream.c b/src/os/wstream.c
new file mode 100644
index 0000000000..c7984d8266
--- /dev/null
+++ b/src/os/wstream.c
@@ -0,0 +1,114 @@
+#include <stdint.h>
+#include <stdbool.h>
+
+#include <uv.h>
+
+#include "os/wstream.h"
+#include "os/wstream_defs.h"
+#include "vim.h"
+#include "memory.h"
+
+struct wstream {
+  uv_stream_t *stream;
+  // Memory currently used by pending buffers
+  uint32_t curmem;
+  // Maximum memory used by this instance
+  uint32_t maxmem;
+  // Number of pending requests
+  uint32_t pending_reqs;
+  bool freed;
+};
+
+typedef struct {
+  WStream *wstream;
+  // Buffer containing data to be written
+  char *buffer;
+  // Size of the buffer
+  uint32_t length;
+  // If it's our responsibility to free the buffer
+  bool free;
+} WriteData;
+
+static void write_cb(uv_write_t *req, int status);
+
+WStream * wstream_new(uint32_t maxmem)
+{
+  WStream *rv = xmalloc(sizeof(WStream));
+  rv->maxmem = maxmem;
+  rv->stream = NULL;
+  rv->curmem = 0;
+  rv->pending_reqs = 0;
+  rv->freed = false;
+
+  return rv;
+}
+
+void wstream_free(WStream *wstream)
+{
+  if (!wstream->pending_reqs) {
+    free(wstream);
+  } else {
+    wstream->freed = true;
+  }
+}
+
+void wstream_set_stream(WStream *wstream, uv_stream_t *stream)
+{
+  stream->data = wstream;
+  wstream->stream = stream;
+}
+
+bool wstream_write(WStream *wstream, char *buffer, uint32_t length, bool free)
+{
+  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;
+  }
+
+  if (wstream->curmem + length > wstream->maxmem) {
+    return false;
+  }
+
+  if (free) {
+    // We should only account for buffers that are ours to free
+    wstream->curmem += length;
+  }
+
+  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;
+  wstream->pending_reqs++;
+  uv_write(req, wstream->stream, &uvbuf, 1, write_cb);
+
+  return true;
+}
+
+static void write_cb(uv_write_t *req, int status)
+{
+  WriteData *data = req->data;
+
+  free(req);
+
+  if (data->free) {
+    // Free the data written to the stream
+    free(data->buffer);
+    data->wstream->curmem -= data->length;
+  }
+
+  if (data->wstream->freed && --data->wstream->pending_reqs == 0) {
+    // Last pending write, free the wstream;
+    free(data->wstream);
+  }
+
+  free(data);
+}
diff --git a/src/os/wstream.h b/src/os/wstream.h
new file mode 100644
index 0000000000..4a557ffd9f
--- /dev/null
+++ b/src/os/wstream.h
@@ -0,0 +1,40 @@
+#ifndef NEOVIM_OS_WSTREAM_H
+#define NEOVIM_OS_WSTREAM_H
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <uv.h>
+
+#include "os/wstream_defs.h"
+
+/// Creates a new WStream instance. A WStream encapsulates all the boilerplate
+/// necessary for writing to a libuv stream.
+///
+/// @param maxmem Maximum amount memory used by this `WStream` instance.
+/// @return The newly-allocated `WStream` instance
+WStream * wstream_new(uint32_t maxmem);
+
+/// Frees all memory allocated for a WStream instance
+///
+/// @param wstream The `WStream` instance
+void wstream_free(WStream *wstream);
+
+/// Sets the underlying `uv_stream_t` instance
+///
+/// @param wstream The `WStream` instance
+/// @param stream The new `uv_stream_t` instance
+void wstream_set_stream(WStream *wstream, uv_stream_t *stream);
+
+/// Queues data for writing to the backing file descriptor of a `WStream`
+/// instance. This will fail if the write would cause the WStream use more
+/// memory than specified by `maxmem`.
+///
+/// @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, uint32_t length, bool free);
+
+#endif  // NEOVIM_OS_WSTREAM_H
+
diff --git a/src/os/wstream_defs.h b/src/os/wstream_defs.h
new file mode 100644
index 0000000000..59d57365fa
--- /dev/null
+++ b/src/os/wstream_defs.h
@@ -0,0 +1,7 @@
+#ifndef NEOVIM_OS_WSTREAM_DEFS_H
+#define NEOVIM_OS_WSTREAM_DEFS_H
+
+typedef struct wstream WStream;
+
+#endif  // NEOVIM_OS_WSTREAM_DEFS_H
+