WIP: stream: Implement cancellation support for uv_write_t

The ability to request cancellation of a pending or in-progress write
has been requested several times in libuv (at least joyent/libuv#1393 and
libuv#2051), but was not yet implemented. I am currently on a mission to
make Ctrl-C work nicely and reliably in Julia (JuliaLang/julia#60281).
Of course, this would require the ability to cancel in-progress writes
for sane semantics, so I have a renewed interest in this feature, which
this PR attempts to implement.

One primary problem with the API is that the existing callback does not
provide support for passing in the number of written bytes (since it
is expected to always complete). I see two options:

1. Create a new uv_write3 that takes a new cb type that does take this
   argument, or,
2. Create a new public `uv_write_t` field that this information can be
   read from.

This PR takes the first approach, but as it turns out the number of
bytes written needs to be stored anyway, so maybe there is not much
point to this choice.

As mentioned, we do need to store two extra bits of state (the total
number of bytes written and which kind of callback we have). To maintain
ABI compatibility we steal two pointers from the generic req_t reserve
pool. This preserves ABI compatibility, but is a bit awkward, because
it puts these new private fields into a different place than they would
otherwise be. I don't see that we have much choice though, other than
creating a completely new req_t subtype (which does not seem worth it).

The `uv_cancel` function returns `0` on success or `UV_EBUSY` if the
request was submitted using the `uv_write`/`uv_write2` API that does
not take the new callback signature (unless no bytes have been written,
in which case cancellation succeeds).

It also returns `0` if the request is already done. The thought here
was that it would be too racy to return an error code here, but I
think it would be fine to return EALREADY (although I don't know what
the caller would do with that information.

`uv_write3` also takes a flags parameter reserved for future use (just
in case).

The windows side is relatively straightforward in that the only thing
we really need to do is call `CancelIoEx` on the overlapped structure
that's already inside of our req. We do of course need the appropriate
accounting for the number of bytes written.

Discloure: Claude Code was used in the creation of this PR, although
dumb design decisions are probably by me.

I have done some minimal integrated testing of this on Linux.
On windows, I have run the included test, but have not tested
in-situ. I consider this WIP until I've had a chance to run our
full test suite on all platforms, but I wanted to make sure to
open this early for API feedback/concerns.

Author: Keno

CMakeLists.txt

@@ -662,6 +662,7 @@ if(LIBUV_BUILD_TESTS)
        test/test-tcp-unexpected-read.c
        test/test-tcp-write-after-connect.c
        test/test-tcp-write-fail.c
+       test/test-tcp-write-cancel.c
        test/test-tcp-write-queue-order.c
        test/test-tcp-write-to-half-open-connection.c
        test/test-tcp-writealot.c

Makefile.am

@@ -289,6 +289,7 @@ test_run_tests_SOURCES = test/blackhole-server.c \
                          test/test-tcp-try-write.c \
                          test/test-tcp-write-in-a-row.c \
                          test/test-tcp-try-write-error.c \
+                         test/test-tcp-write-cancel.c \
                          test/test-tcp-write-queue-order.c \
                          test/test-test-macros.c \
                          test/test-thread-equal.c \

include/uv.h

@@ -329,6 +329,7 @@ typedef void (*uv_read_cb)(uv_stream_t* stream,
                            ssize_t nread,
                            const uv_buf_t* buf);
 typedef void (*uv_write_cb)(uv_write_t* req, int status);
+typedef void (*uv_write3_cb)(uv_write_t* req, int status, size_t nwritten);
 typedef void (*uv_connect_cb)(uv_connect_t* req, int status);
 typedef void (*uv_shutdown_cb)(uv_shutdown_t* req, int status);
 typedef void (*uv_connection_cb)(uv_stream_t* server, int status);
@@ -437,7 +438,7 @@ UV_EXTERN char* uv_err_name_r(int err, char* buf, size_t buflen);
   /* read-only */                                                             \
   uv_req_type type;                                                           \
   /* private */                                                               \
-  void* reserved[6];                                                          \
+  void* reserved[4];                                                          \
   UV_REQ_PRIVATE_FIELDS                                                       \
 
 /* Abstract base class of all requests. */
@@ -567,6 +568,13 @@ UV_EXTERN int uv_try_write2(uv_stream_t* handle,
                             const uv_buf_t bufs[],
                             unsigned int nbufs,
                             uv_stream_t* send_handle);
+UV_EXTERN int uv_write3(uv_write_t* req,
+                        uv_stream_t* handle,
+                        const uv_buf_t bufs[],
+                        unsigned int nbufs,
+                        uv_stream_t* send_handle,
+                        unsigned int flags,
+                        uv_write3_cb cb);
 
 /* uv_write_t is a subclass of uv_req_t. */
 struct uv_write_s {

include/uv/unix.h

@@ -249,7 +249,14 @@ typedef struct {
 
 #define UV_REQ_TYPE_PRIVATE /* empty */
 
-#define UV_REQ_PRIVATE_FIELDS  /* empty */
+#define UV_REQ_PRIVATE_FIELDS                                                 \
+  union {                                                                     \
+    void* reserved2[2];                                                       \
+    struct {                                                                  \
+      size_t bytes_written;                                                   \
+      unsigned int write_flags;                                               \
+    } write_extra;                                                            \
+  };
 
 #define UV_PRIVATE_REQ_TYPES /* empty */
 
@@ -260,6 +267,10 @@ typedef struct {
   unsigned int nbufs;                                                         \
   int error;                                                                  \
   uv_buf_t bufsml[4];                                                         \
+  union {                                                                     \
+    uv_write_cb cb;                                                           \
+    uv_write3_cb cb3;                                                         \
+  } write_cb;                                                                 \
 
 #define UV_CONNECT_PRIVATE_FIELDS                                             \
   struct uv__queue queue;                                                     \

include/uv/win.h

@@ -379,13 +379,25 @@ typedef struct {
       WCHAR* name;                                                             \
     } connect;                                                                \
   } u;                                                                        \
-  struct uv_req_s* next_req;
+  struct uv_req_s* next_req;                                                  \
+  /* Extra fields for uv_write_t cancellation support */                      \
+  union {                                                                     \
+    void* reserved2[2];                                                       \
+    struct {                                                                  \
+      size_t bytes_written;                                                   \
+      unsigned int write_flags;                                               \
+    } write_extra;                                                            \
+  };
 
 #define UV_WRITE_PRIVATE_FIELDS \
   int coalesced;                \
   uv_buf_t write_buffer;        \
   HANDLE event_handle;          \
-  HANDLE wait_handle;
+  HANDLE wait_handle;           \
+  union {                       \
+    uv_write_cb cb;             \
+    uv_write3_cb cb3;           \
+  } write_cb;
 
 #define UV_CONNECT_PRIVATE_FIELDS                                             \
   /* empty */

src/queue.h

@@ -87,4 +87,16 @@ static inline void uv__queue_remove(struct uv__queue* q) {
   q->next->prev = q->prev;
 }
 
+static inline int uv__queue_contains(const struct uv__queue* h,
+                                     const struct uv__queue* q) {
+  const struct uv__queue* e;
+
+  uv__queue_foreach(e, h) {
+    if (e == q)
+      return 1;
+  }
+
+  return 0;
+}
+
 #endif /* QUEUE_H_ */

src/threadpool.c

@@ -21,7 +21,9 @@
 
 #include "uv-common.h"
 
-#if !defined(_WIN32)
+#if defined(_WIN32)
+# include "win/internal.h"
+#else
 # include "unix/internal.h"
 #endif
 
@@ -421,6 +423,8 @@ int uv_cancel(uv_req_t* req) {
     loop =  ((uv_work_t*) req)->loop;
     wreq = &((uv_work_t*) req)->work_req;
     break;
+  case UV_WRITE:
+    return uv__write_cancel((uv_write_t*) req);
   default:
     return UV_EINVAL;
   }

src/unix/internal.h

@@ -345,6 +345,7 @@ void uv__stream_init(uv_loop_t* loop, uv_stream_t* stream,
     uv_handle_type type);
 int uv__stream_open(uv_stream_t*, int fd, int flags);
 void uv__stream_destroy(uv_stream_t* stream);
+int uv__write_cancel(uv_write_t* req);
 #if defined(__APPLE__)
 int uv__stream_try_select(uv_stream_t* stream, int* fd);
 #endif /* defined(__APPLE__) */

src/unix/stream.c

@@ -692,6 +692,7 @@ static int uv__write_req_update(uv_stream_t* stream,
 
   assert(n <= stream->write_queue_size);
   stream->write_queue_size -= n;
+  req->write_extra.bytes_written += n;
 
   buf = req->bufs + req->write_index;
 
@@ -736,6 +737,46 @@ static void uv__write_req_finish(uv_write_t* req) {
 }
 
 
+int uv__write_cancel(uv_write_t* req) {
+  uv_stream_t* stream;
+  size_t remaining;
+
+  stream = req->handle;
+  if (stream == NULL)
+    return UV_EINVAL;
+
+  /* If the provided callback does not allow passing nwritten,
+   * don't allow cancellation of partial writes. Return EBUSY instead.
+   */
+  if (!(req->write_extra.write_flags & UV__WRITE_ALLOW_PARTIAL)) {
+    if (req->write_extra.bytes_written > 0)
+      return UV_EBUSY;
+  }
+
+  /* Already completed, treat as success */
+  if (!uv__queue_contains(&stream->write_queue, &req->queue))
+    return 0;
+
+  uv__queue_remove(&req->queue);
+
+  remaining = uv__count_bufs(req->bufs + req->write_index,
+                             req->nbufs - req->write_index);
+  assert(remaining <= stream->write_queue_size);
+  stream->write_queue_size -= remaining;
+
+  if (req->bufs != req->bufsml)
+    uv__free(req->bufs);
+  req->bufs = NULL;
+
+  req->error = UV_ECANCELED;
+  uv__queue_insert_tail(&stream->write_completed_queue, &req->queue);
+
+  uv__io_feed(stream->loop, &stream->io_watcher);
+
+  return 0;
+}
+
+
 static int uv__handle_fd(uv_handle_t* handle) {
   switch (handle->type) {
     case UV_NAMED_PIPE:
@@ -922,8 +963,13 @@ static void uv__write_callbacks(uv_stream_t* stream) {
     }
 
     /* NOTE: call callback AFTER freeing the request data. */
-    if (req->cb)
-      req->cb(req, req->error);
+    if (req->write_extra.write_flags & UV__WRITE_ALLOW_PARTIAL) {
+      if (req->write_cb.cb3 != NULL)
+        req->write_cb.cb3(req, req->error, req->write_extra.bytes_written);
+    } else {
+      if (req->write_cb.cb != NULL)
+        req->write_cb.cb(req, req->error);
+    }
   }
 }
 
@@ -1329,12 +1375,12 @@ static int uv__check_before_write(uv_stream_t* stream,
   return 0;
 }
 
-int uv_write2(uv_write_t* req,
-              uv_stream_t* stream,
-              const uv_buf_t bufs[],
-              unsigned int nbufs,
-              uv_stream_t* send_handle,
-              uv_write_cb cb) {
+static int uv_write_internal(uv_write_t* req,
+                            uv_stream_t* stream,
+                            const uv_buf_t bufs[],
+                            unsigned int nbufs,
+                            uv_stream_t* send_handle,
+                            unsigned int write_flags) {
   int empty_queue;
   int err;
 
@@ -1352,7 +1398,8 @@ int uv_write2(uv_write_t* req,
 
   /* Initialize the req */
   uv__req_init(stream->loop, req, UV_WRITE);
-  req->cb = cb;
+  req->write_extra.bytes_written = 0;
+  req->write_extra.write_flags = write_flags;
   req->handle = stream;
   req->error = 0;
   req->send_handle = send_handle;
@@ -1398,6 +1445,39 @@ int uv_write2(uv_write_t* req,
 }
 
 
+int uv_write2(uv_write_t* req,
+              uv_stream_t* stream,
+              const uv_buf_t bufs[],
+              unsigned int nbufs,
+              uv_stream_t* send_handle,
+              uv_write_cb cb) {
+  int err;
+
+  req->write_cb.cb = cb;
+  err = uv_write_internal(req, stream, bufs, nbufs, send_handle, 0);
+
+  return err;
+}
+
+int uv_write3(uv_write_t* req,
+              uv_stream_t* stream,
+              const uv_buf_t bufs[],
+              unsigned int nbufs,
+              uv_stream_t* send_handle,
+              unsigned int flags,
+              uv_write3_cb cb) {
+  int err;
+
+  if (flags != 0)
+    return UV_EINVAL;
+
+  req->write_cb.cb3 = cb;
+  err = uv_write_internal(req, stream, bufs, nbufs, send_handle,
+                          UV__WRITE_ALLOW_PARTIAL);
+
+  return err;
+}
+
 /* The buffers to be written must remain valid until the callback is called.
  * This is not required for the uv_buf_t array.
  */
@@ -1409,7 +1489,6 @@ int uv_write(uv_write_t* req,
   return uv_write2(req, handle, bufs, nbufs, NULL, cb);
 }
 
-
 int uv_try_write(uv_stream_t* stream,
                  const uv_buf_t bufs[],
                  unsigned int nbufs) {

src/uv-common.h

@@ -141,6 +141,11 @@ enum {
   UV_HANDLE_REAP                        = 0x10000000
 };
 
+/* Internal flags for uv_write_t requests */
+enum {
+  UV__WRITE_ALLOW_PARTIAL = 0x1  /* Request allows partial write on cancel */
+};
+
 static inline int uv__is_raw_tty_mode(uv_tty_mode_t m) {
   return m == UV_TTY_MODE_RAW || m == UV_TTY_MODE_RAW_VT;
 }