Only "data" frames affect websocket's read window now (#407)

Fix it so that a websocket with manual_window_management turned off will disable backpressure on its whole channel (previously, we kept backpressure on and just automatically maintained a constant window size).

Change it so that the websocket's read window is only affected by "data" frames (TEXT, BINARY, CONTINUATION) and not "control" frames (PING, PONG, CLOSE). The assumption here being that only "data" frames are generating the kind of traffic that the user is forwarding along and needs tight control over. This is similar to how HTTP/2 only counts "data" frames as affecting the window.

Add copious documentation, because this stuff is confusing

Author: graebm

include/aws/http/private/websocket_impl.h

@@ -78,6 +78,7 @@ struct aws_websocket_client_bootstrap_system_vtable {
     int (*aws_http_stream_activate)(struct aws_http_stream *stream);
     void (*aws_http_stream_release)(struct aws_http_stream *stream);
     struct aws_http_connection *(*aws_http_stream_get_connection)(const struct aws_http_stream *stream);
+    void (*aws_http_stream_update_window)(struct aws_http_stream *stream, size_t increment_size);
     int (*aws_http_stream_get_incoming_response_status)(const struct aws_http_stream *stream, int *out_status);
     struct aws_websocket *(*aws_websocket_handler_new)(const struct aws_websocket_handler_options *options);
 };

include/aws/http/websocket.h

@@ -97,6 +97,17 @@ typedef bool(aws_websocket_on_incoming_frame_begin_fn)(
  * Payload data will not be valid after this call, so copy if necessary.
  * The payload data is always unmasked at this point.
  *
+ * NOTE: If you created the websocket with `manual_window_management` set true, you must maintain the read window.
+ * Whenever the read window reaches 0, you will stop receiving anything.
+ * The websocket's `initial_window_size` determines the starting size of the read window.
+ * The read window shrinks as you receive the payload from "data" frames (TEXT, BINARY, and CONTINUATION).
+ * Use aws_websocket_increment_read_window() to increment the window again and keep frames flowing.
+ * Maintain a larger window to keep up high throughput.
+ * You only need to worry about the payload from "data" frames.
+ * The websocket automatically increments the window to account for any
+ * other incoming bytes, including other parts of a frame (opcode, payload-length, etc)
+ * and the payload of other frame types (PING, PONG, CLOSE).
+ *
  * Return true to proceed normally. If false is returned, the websocket will read no further data,
  * the frame will complete with an error-code, and the connection will close.
  */
@@ -186,8 +197,8 @@ struct aws_websocket_client_connection_options {
     struct aws_http_message *handshake_request;
 
     /**
-     * Initial window size for websocket.
-     * Required.
+     * Initial size of the websocket's read window.
+     * Ignored unless `manual_window_management` is true.
      * Set to 0 to prevent any incoming websocket frames until aws_websocket_increment_read_window() is called.
      */
     size_t initial_window_size;
@@ -241,11 +252,17 @@ struct aws_websocket_client_connection_options {
     /**
      * Set to true to manually manage the read window size.
      *
-     * If this is false, the connection will maintain a constant window size.
+     * If this is false, no backpressure is applied and frames will arrive as fast as possible.
      *
-     * If this is true, the caller must manually increment the window size using aws_websocket_increment_read_window().
-     * If the window is not incremented, it will shrink by the amount of payload data received. If the window size
-     * reaches 0, no further data will be received.
+     * If this is true, then whenever the read window reaches 0 you will stop receiving anything.
+     * The websocket's `initial_window_size` determines the starting size of the read window.
+     * The read window shrinks as you receive the payload from "data" frames (TEXT, BINARY, and CONTINUATION).
+     * Use aws_websocket_increment_read_window() to increment the window again and keep frames flowing.
+     * Maintain a larger window to keep up high throughput.
+     * You only need to worry about the payload from "data" frames.
+     * The websocket automatically increments the window to account for any
+     * other incoming bytes, including other parts of a frame (opcode, payload-length, etc)
+     * and the payload of other frame types (PING, PONG, CLOSE).
      */
     bool manual_window_management;
 
@@ -390,10 +407,21 @@ AWS_HTTP_API
 int aws_websocket_send_frame(struct aws_websocket *websocket, const struct aws_websocket_send_frame_options *options);
 
 /**
- * Manually increment the read window.
- * The read window shrinks as payload data is received, and reading stops when its size reaches 0.
- * Note that the read window can also be controlled from the aws_websocket_on_incoming_frame_payload_fn(),
- * callback, by manipulating the `out_increment_window` argument.
+ * Manually increment the read window to keep frames flowing.
+ *
+ * If the websocket was created with `manual_window_management` set true,
+ * then whenever the read window reaches 0 you will stop receiving data.
+ * The websocket's `initial_window_size` determines the starting size of the read window.
+ * The read window shrinks as you receive the payload from "data" frames (TEXT, BINARY, and CONTINUATION).
+ * Use aws_websocket_increment_read_window() to increment the window again and keep frames flowing.
+ * Maintain a larger window to keep up high throughput.
+ * You only need to worry about the payload from "data" frames.
+ * The websocket automatically increments the window to account for any
+ * other incoming bytes, including other parts of a frame (opcode, payload-length, etc)
+ * and the payload of other frame types (PING, PONG, CLOSE).
+ *
+ * If the websocket was created with `manual_window_management` set false, this function does nothing.
+ *
  * This function may be called from any thread.
  */
 AWS_HTTP_API

source/websocket.c

@@ -1238,6 +1238,9 @@ static int s_handler_process_read_message(
     struct aws_byte_cursor cursor = aws_byte_cursor_from_buf(&message->message_data);
     int err;
 
+    /* At the end of this function we'll bump the window back up by this amount.
+     * We start off assuming we'll re-open the window by the whole amount,
+     * but this number will go down if we process any payload data that ought to shrink the window */
     websocket->thread_data.incoming_message_window_update = message->message_data.len;
 
     AWS_LOGF_TRACE(
@@ -1382,15 +1385,16 @@ static int s_decoder_on_user_payload(struct aws_websocket *websocket, struct aws
         return aws_raise_error(AWS_ERROR_HTTP_CALLBACK_FAILURE);
     }
 
-    /* If user reduced window_update_size, reduce how much the websocket will update its window */
-    if (websocket->manual_window_update) {
-        size_t reduce = data.len;
+    /* If this is a "data" frame's payload, let the window shrink */
+    if (aws_websocket_is_data_frame(websocket->thread_data.current_incoming_frame->opcode) &&
+        websocket->manual_window_update) {
+
         websocket->thread_data.incoming_message_window_update -= data.len;
         AWS_LOGF_DEBUG(
             AWS_LS_HTTP_WEBSOCKET,
-            "id=%p: Incoming payload callback changed window update size, window will shrink by %zu.",
+            "id=%p: The read window is shrinking by %zu due to incoming payload from 'data' frame.",
             (void *)websocket,
-            reduce);
+            data.len);
     }
 
     return AWS_OP_SUCCESS;
@@ -1607,6 +1611,14 @@ void aws_websocket_increment_read_window(struct aws_websocket *websocket, size_t
         return;
     }
 
+    if (!websocket->manual_window_update) {
+        AWS_LOGF_DEBUG(
+            AWS_LS_HTTP_WEBSOCKET,
+            "id=%p: Ignoring window increment. Manual window management (aka read backpressure) is not enabled.",
+            (void *)websocket);
+        return;
+    }
+
     /* Schedule a task to do the increment.
      * If task is already scheduled, just increase size to be incremented */
     bool is_midchannel_handler = false;

source/websocket_bootstrap.c

@@ -29,6 +29,7 @@ static const struct aws_websocket_client_bootstrap_system_vtable s_default_syste
     .aws_http_stream_activate = aws_http_stream_activate,
     .aws_http_stream_release = aws_http_stream_release,
     .aws_http_stream_get_connection = aws_http_stream_get_connection,
+    .aws_http_stream_update_window = aws_http_stream_update_window,
     .aws_http_stream_get_incoming_response_status = aws_http_stream_get_incoming_response_status,
     .aws_websocket_handler_new = aws_websocket_handler_new,
 };
@@ -94,6 +95,10 @@ static int s_ws_bootstrap_on_handshake_response_header_block_done(
     struct aws_http_stream *stream,
     enum aws_http_header_block header_block,
     void *user_data);
+static int s_ws_bootstrap_on_handshake_response_body(
+    struct aws_http_stream *stream,
+    const struct aws_byte_cursor *data,
+    void *user_data);
 static void s_ws_bootstrap_on_stream_complete(struct aws_http_stream *stream, int error_code, void *user_data);
 
 int aws_websocket_client_connect(const struct aws_websocket_client_connection_options *options) {
@@ -188,10 +193,19 @@ int aws_websocket_client_connect(const struct aws_websocket_client_connection_op
     http_options.socket_options = options->socket_options;
     http_options.tls_options = options->tls_options;
     http_options.proxy_options = options->proxy_options;
-    http_options.initial_window_size = 1024; /* Adequate space for response data to trickle in */
 
-    /* TODO: websockets has issues if back-pressure is disabled on the whole channel. This should be fixed. */
-    http_options.manual_window_management = true;
+    if (options->manual_window_management) {
+        http_options.manual_window_management = true;
+
+        /* Give HTTP handler enough window to comfortably receive the handshake response.
+         *
+         * If the upgrade is unsuccessful, the HTTP window will shrink as the response body is received.
+         * In this case, we'll keep incrementing the window back to its original size so data keeps arriving.
+         *
+         * If the upgrade is successful, then the websocket handler is installed, and
+         * the HTTP handler will take over its own window management. */
+        http_options.initial_window_size = 1024;
+    }
 
     http_options.user_data = ws_bootstrap;
     http_options.on_setup = s_ws_bootstrap_on_http_setup;
@@ -309,6 +323,7 @@ static void s_ws_bootstrap_on_http_setup(struct aws_http_connection *http_connec
         .user_data = ws_bootstrap,
         .on_response_headers = s_ws_bootstrap_on_handshake_response_headers,
         .on_response_header_block_done = s_ws_bootstrap_on_handshake_response_header_block_done,
+        .on_response_body = s_ws_bootstrap_on_handshake_response_body,
         .on_complete = s_ws_bootstrap_on_stream_complete,
     };
 
@@ -544,6 +559,26 @@ static int s_ws_bootstrap_on_handshake_response_header_block_done(
     return AWS_OP_ERR;
 }
 
+/**
+ * Invoked as we receive the body of a failed response.
+ * This is never invoked if the handshake succeeds.
+ */
+static int s_ws_bootstrap_on_handshake_response_body(
+    struct aws_http_stream *stream,
+    const struct aws_byte_cursor *data,
+    void *user_data) {
+
+    struct aws_websocket_client_bootstrap *ws_bootstrap = user_data;
+
+    /* If we're managing the read window...
+     * bump the HTTP window back to its starting size, so that we keep receiving the whole response. */
+    if (ws_bootstrap->manual_window_update) {
+        s_system_vtable->aws_http_stream_update_window(stream, data->len);
+    }
+
+    return AWS_OP_SUCCESS;
+}
+
 static void s_ws_bootstrap_on_stream_complete(struct aws_http_stream *stream, int error_code, void *user_data) {
     /* Not checking error_code because a stream error ends the HTTP connection.
      * We'll deal with finishing setup and/or shutdown from the http-shutdown callback */

tests/CMakeLists.txt

@@ -206,7 +206,6 @@ add_test_case(websocket_handler_read_frames_complete_on_shutdown)
 add_test_case(websocket_handler_read_halts_if_begin_fn_returns_false)
 add_test_case(websocket_handler_read_halts_if_payload_fn_returns_false)
 add_test_case(websocket_handler_read_halts_if_complete_fn_returns_false)
-add_test_case(websocket_handler_window_reopens_by_default)
 add_test_case(websocket_handler_window_manual_increment)
 add_test_case(websocket_handler_window_manual_increment_off_thread)
 add_test_case(websocket_midchannel_sanity_check)

tests/test_websocket_bootstrap.c

@@ -30,6 +30,7 @@ static struct aws_http_stream *s_mock_http_connection_make_request(
 static int s_mock_http_stream_activate(struct aws_http_stream *stream);
 static void s_mock_http_stream_release(struct aws_http_stream *stream);
 static struct aws_http_connection *s_mock_http_stream_get_connection(const struct aws_http_stream *stream);
+static void s_mock_http_stream_update_window(struct aws_http_stream *stream, size_t increment_size);
 static int s_mock_http_stream_get_incoming_response_status(const struct aws_http_stream *stream, int *out_status);
 static struct aws_websocket *s_mock_websocket_handler_new(const struct aws_websocket_handler_options *options);
 
@@ -42,6 +43,7 @@ static const struct aws_websocket_client_bootstrap_system_vtable s_mock_system_v
     .aws_http_stream_activate = s_mock_http_stream_activate,
     .aws_http_stream_release = s_mock_http_stream_release,
     .aws_http_stream_get_connection = s_mock_http_stream_get_connection,
+    .aws_http_stream_update_window = s_mock_http_stream_update_window,
     .aws_http_stream_get_incoming_response_status = s_mock_http_stream_get_incoming_response_status,
     .aws_websocket_handler_new = s_mock_websocket_handler_new,
 };
@@ -112,6 +114,9 @@ static struct tester {
 
     bool websocket_shutdown_invoked;
     int websocket_shutdown_error_code;
+
+    /* Track the sum of all calls to aws_http_stream_update_window() */
+    size_t window_increment_total;
 } s_tester;
 
 static int s_tester_init(struct aws_allocator *alloc) {
@@ -308,6 +313,13 @@ static struct aws_http_connection *s_mock_http_stream_get_connection(const struc
     return s_mock_http_connection;
 }
 
+static void s_mock_http_stream_update_window(struct aws_http_stream *stream, size_t increment_size) {
+    AWS_FATAL_ASSERT(stream == s_mock_stream);
+    AWS_FATAL_ASSERT(!s_tester.http_connection_release_called);
+    AWS_FATAL_ASSERT(!s_tester.http_stream_release_called);
+    s_tester.window_increment_total += increment_size;
+}
+
 static int s_mock_http_stream_get_incoming_response_status(const struct aws_http_stream *stream, int *out_status) {
     AWS_FATAL_ASSERT(stream == s_mock_stream);
     AWS_FATAL_ASSERT(!s_tester.http_connection_release_called);

tests/test_websocket_handler.c

@@ -271,6 +271,7 @@ static bool s_on_incoming_frame_complete(
 static void s_set_readpush_frames(struct tester *tester, struct readpush_frame *frames, size_t num_frames) {
     tester->readpush_frames = frames;
     tester->num_readpush_frames = num_frames;
+    tester->readpush_frame_index = 0;
     for (size_t i = 0; i < num_frames; ++i) {
         struct readpush_frame *frame = &frames[i];
         frame->cursor = frame->payload;
@@ -1719,37 +1720,12 @@ TEST_CASE(websocket_handler_read_halts_if_complete_fn_returns_false) {
     return AWS_OP_SUCCESS;
 }
 
-TEST_CASE(websocket_handler_window_reopens_by_default) {
-    (void)ctx;
-    struct tester tester;
-    ASSERT_SUCCESS(s_tester_init(&tester, allocator));
-
-    struct readpush_frame pushing = {
-        .payload = aws_byte_cursor_from_c_str("Tore open the shutters and threw up the sash."),
-        .def =
-            {
-                .opcode = AWS_WEBSOCKET_OPCODE_TEXT,
-                .fin = true,
-            },
-    };
-
-    s_set_readpush_frames(&tester, &pushing, 1);
-    ASSERT_SUCCESS(s_do_readpush_all(&tester));
-
-    testing_channel_drain_queued_tasks(&tester.testing_channel);
-
-    uint64_t total_frame_size = aws_websocket_frame_encoded_size(&pushing.def);
-    ASSERT_UINT_EQUALS(total_frame_size, testing_channel_last_window_update(&tester.testing_channel));
-
-    ASSERT_SUCCESS(s_tester_clean_up(&tester));
-    return AWS_OP_SUCCESS;
-}
-
 static int s_window_manual_increment_common(struct aws_allocator *allocator, bool on_thread) {
     struct tester tester;
     s_tester_options.manual_window_update = true;
     ASSERT_SUCCESS(s_tester_init(&tester, allocator));
 
+    /* Push "data" frame to websocket */
     struct readpush_frame pushing = {
         .payload = aws_byte_cursor_from_c_str("Shrink, then open"),
         .def =
@@ -1777,6 +1753,21 @@ static int s_window_manual_increment_common(struct aws_allocator *allocator, boo
     testing_channel_drain_queued_tasks(&tester.testing_channel);
     ASSERT_UINT_EQUALS(pushing.def.payload_length, testing_channel_last_window_update(&tester.testing_channel));
 
+    /* Now push a control frame, and ensure the window automatically re-opens by the whole amount */
+    struct readpush_frame pushing_control_frame = {
+        .payload = aws_byte_cursor_from_c_str("free data"),
+        .def = {.opcode = AWS_WEBSOCKET_OPCODE_PONG, .fin = true},
+    };
+
+    s_set_readpush_frames(&tester, &pushing_control_frame, 1);
+    ASSERT_SUCCESS(s_do_readpush_all(&tester));
+    testing_channel_drain_queued_tasks(&tester.testing_channel);
+
+    ASSERT_UINT_EQUALS(
+        aws_websocket_frame_encoded_size(&pushing_control_frame.def),
+        testing_channel_last_window_update(&tester.testing_channel));
+
+    /* Done */
     ASSERT_SUCCESS(s_tester_clean_up(&tester));
     return AWS_OP_SUCCESS;
 }