response_first_byte_timeout_ms support (#456)

TingDaoK · graebm · web-flow · commit a082f8a2067e · 2023-10-26T11:28:05.000-07:00
Co-authored-by: Michael Graeb &lt;graebm@amazon.com&gt;
diff --git a/include/aws/http/connection.h b/include/aws/http/connection.h
@@ -307,6 +307,15 @@ struct aws_http_client_connection_options {
      */
     const struct aws_http_connection_monitoring_options *monitoring_options;
 
+    /**
+     * Optional (ignored if 0).
+     * After a request is fully sent, if the server does not begin responding within N milliseconds,
+     * then fail with AWS_ERROR_HTTP_RESPONSE_FIRST_BYTE_TIMEOUT.
+     * This can be overridden per-request by aws_http_make_request_options.response_first_byte_timeout_ms.
+     * TODO: Only supported in HTTP/1.1 now, support it in HTTP/2
+     */
+    uint64_t response_first_byte_timeout_ms;
+
     /**
      * Set to true to manually manage the flow-control window of each stream.
      *
diff --git a/include/aws/http/http.h b/include/aws/http/http.h
@@ -59,6 +59,7 @@ enum aws_http_errors {
     AWS_ERROR_HTTP_WEBSOCKET_PROTOCOL_ERROR,
     AWS_ERROR_HTTP_MANUAL_WRITE_NOT_ENABLED,
     AWS_ERROR_HTTP_MANUAL_WRITE_HAS_COMPLETED,
+    AWS_ERROR_HTTP_RESPONSE_FIRST_BYTE_TIMEOUT,
 
     AWS_ERROR_HTTP_END_RANGE = AWS_ERROR_ENUM_END_RANGE(AWS_C_HTTP_PACKAGE_ID)
 };
diff --git a/include/aws/http/private/connection_impl.h b/include/aws/http/private/connection_impl.h
@@ -103,7 +103,7 @@ struct aws_http_connection {
 
     union {
         struct aws_http_connection_client_data {
-            uint8_t delete_me; /* exists to prevent "empty struct" errors */
+            uint64_t response_first_byte_timeout_ms;
         } client;
 
         struct aws_http_connection_server_data {
@@ -133,6 +133,7 @@ struct aws_http_client_bootstrap {
     aws_http_on_client_connection_setup_fn *on_setup;
     aws_http_on_client_connection_shutdown_fn *on_shutdown;
     aws_http_proxy_request_transform_fn *proxy_request_transform;
+    uint64_t response_first_byte_timeout_ms;
 
     struct aws_http1_connection_options http1_options;
     struct aws_http2_connection_options http2_options; /* allocated with bootstrap */
diff --git a/include/aws/http/private/request_response_impl.h b/include/aws/http/private/request_response_impl.h
@@ -6,6 +6,7 @@
  * SPDX-License-Identifier: Apache-2.0.
  */
 
+#include <aws/common/task_scheduler.h>
 #include <aws/http/request_response.h>
 
 #include <aws/http/private/http_impl.h>
@@ -54,6 +55,10 @@ struct aws_http_stream {
     union {
         struct aws_http_stream_client_data {
             int response_status;
+            uint64_t response_first_byte_timeout_ms;
+            /* Using aws_task instead of aws_channel_task because, currently, channel-tasks can't be canceled.
+             * We only touch this from the connection's thread */
+            struct aws_task response_first_byte_timeout_task;
         } client;
         struct aws_http_stream_server_data {
             struct aws_byte_cursor request_method_str;
diff --git a/include/aws/http/request_response.h b/include/aws/http/request_response.h
@@ -300,6 +300,16 @@ struct aws_http_make_request_options {
      * when data has been supplied via `aws_http2_stream_write_data`
      */
     bool http2_use_manual_data_writes;
+
+    /**
+     * Optional (ignored if 0).
+     * After a request is fully sent, if the server does not begin responding within N milliseconds, then fail with
+     * AWS_ERROR_HTTP_RESPONSE_FIRST_BYTE_TIMEOUT.
+     * It override the connection level settings, when the request completes, the
+     * original monitoring options will be applied back to the connection.
+     * TODO: Only supported in HTTP/1.1 now, support it in HTTP/2
+     */
+    uint64_t response_first_byte_timeout_ms;
 };
 
 struct aws_http_request_handler_options {
diff --git a/source/connection.c b/source/connection.c
@@ -834,6 +834,8 @@ static void s_client_bootstrap_on_channel_setup(
     }
 
     http_bootstrap->connection->proxy_request_transform = http_bootstrap->proxy_request_transform;
+    http_bootstrap->connection->client_data->response_first_byte_timeout_ms =
+        http_bootstrap->response_first_byte_timeout_ms;
 
     AWS_LOGF_INFO(
         AWS_LS_HTTP_CONNECTION,
@@ -1073,6 +1075,7 @@ int aws_http_client_connect_internal(
     http_bootstrap->proxy_request_transform = proxy_request_transform;
     http_bootstrap->http1_options = *options.http1_options;
     http_bootstrap->http2_options = *options.http2_options;
+    http_bootstrap->response_first_byte_timeout_ms = options.response_first_byte_timeout_ms;
 
     /* keep a copy of the settings array if it's not NULL */
     if (options.http2_options->num_initial_settings > 0) {
diff --git a/source/h1_connection.c b/source/h1_connection.c
@@ -11,6 +11,7 @@
 #include <aws/http/private/h1_stream.h>
 #include <aws/http/private/request_response_impl.h>
 #include <aws/http/status_code.h>
+#include <aws/io/event_loop.h>
 #include <aws/io/logging.h>
 
 #include <inttypes.h>
@@ -535,6 +536,7 @@ static int s_aws_http1_switch_protocols(struct aws_h1_connection *connection) {
 static void s_stream_complete(struct aws_h1_stream *stream, int error_code) {
     struct aws_h1_connection *connection =
         AWS_CONTAINER_OF(stream->base.owning_connection, struct aws_h1_connection, base);
+    AWS_ASSERT(aws_channel_thread_is_callers_thread(connection->base.channel_slot->channel));
 
     /*
      * If this is the end of a successful CONNECT request, mark ourselves as pass-through since the proxy layer
@@ -547,6 +549,14 @@ static void s_stream_complete(struct aws_h1_stream *stream, int error_code) {
         }
     }
 
+    if (stream->base.client_data && stream->base.client_data->response_first_byte_timeout_task.fn != NULL) {
+        /* There is an outstanding response timeout task, but stream completed, we can cancel it now. We are
+         * safe to do it as we always on connection thread to schedule the task or cancel it */
+        struct aws_event_loop *connection_loop = aws_channel_get_event_loop(connection->base.channel_slot->channel);
+        /* The task will be zeroed out within the call */
+        aws_event_loop_cancel_task(connection_loop, &stream->base.client_data->response_first_byte_timeout_task);
+    }
+
     if (error_code != AWS_ERROR_SUCCESS) {
         if (stream->base.client_data && stream->is_incoming_message_done) {
             /* As a request that finished receiving the response, we ignore error and
@@ -721,6 +731,87 @@ static void s_client_update_incoming_stream_ptr(struct aws_h1_connection *connec
     s_set_incoming_stream_ptr(connection, desired);
 }
 
+static void s_http_stream_response_first_byte_timeout_task(
+    struct aws_task *task,
+    void *arg,
+    enum aws_task_status status) {
+    (void)task;
+    struct aws_h1_stream *stream = arg;
+    struct aws_http_connection *connection_base = stream->base.owning_connection;
+    /* zero-out task to indicate that it's no longer scheduled */
+    AWS_ZERO_STRUCT(stream->base.client_data->response_first_byte_timeout_task);
+
+    if (status == AWS_TASK_STATUS_CANCELED) {
+        return;
+    }
+
+    struct aws_h1_connection *connection = AWS_CONTAINER_OF(connection_base, struct aws_h1_connection, base);
+    /* Timeout happened, close the connection */
+    uint64_t response_first_byte_timeout_ms = stream->base.client_data->response_first_byte_timeout_ms == 0
+                                                  ? connection_base->client_data->response_first_byte_timeout_ms
+                                                  : stream->base.client_data->response_first_byte_timeout_ms;
+    AWS_LOGF_INFO(
+        AWS_LS_HTTP_CONNECTION,
+        "id=%p: Closing connection as timeout after request sent to the first byte received happened. "
+        "response_first_byte_timeout_ms is %" PRIu64 ".",
+        (void *)connection_base,
+        response_first_byte_timeout_ms);
+
+    /* Don't stop reading/writing immediately, let that happen naturally during the channel shutdown process. */
+    s_stop(
+        connection,
+        false /*stop_reading*/,
+        false /*stop_writing*/,
+        true /*schedule_shutdown*/,
+        AWS_ERROR_HTTP_RESPONSE_FIRST_BYTE_TIMEOUT);
+}
+
+static void s_set_outgoing_message_done(struct aws_h1_stream *stream) {
+    struct aws_http_connection *connection = stream->base.owning_connection;
+    struct aws_channel *channel = aws_http_connection_get_channel(connection);
+    AWS_ASSERT(aws_channel_thread_is_callers_thread(channel));
+
+    if (stream->is_outgoing_message_done) {
+        /* Already did the job */
+        return;
+    }
+
+    stream->is_outgoing_message_done = true;
+    AWS_ASSERT(stream->base.metrics.send_end_timestamp_ns == -1);
+    aws_high_res_clock_get_ticks((uint64_t *)&stream->base.metrics.send_end_timestamp_ns);
+    AWS_ASSERT(stream->base.metrics.send_start_timestamp_ns != -1);
+    AWS_ASSERT(stream->base.metrics.send_end_timestamp_ns >= stream->base.metrics.send_start_timestamp_ns);
+    stream->base.metrics.sending_duration_ns =
+        stream->base.metrics.send_end_timestamp_ns - stream->base.metrics.send_start_timestamp_ns;
+    if (stream->base.metrics.receive_start_timestamp_ns == -1) {
+        /* We haven't receive any message, schedule the response timeout task */
+
+        uint64_t response_first_byte_timeout_ms = 0;
+        if (stream->base.client_data != NULL && connection->client_data != NULL) {
+            response_first_byte_timeout_ms = stream->base.client_data->response_first_byte_timeout_ms == 0
+                                                 ? connection->client_data->response_first_byte_timeout_ms
+                                                 : stream->base.client_data->response_first_byte_timeout_ms;
+        }
+        if (response_first_byte_timeout_ms != 0) {
+            /* The task should not be initialized before. */
+            AWS_ASSERT(stream->base.client_data->response_first_byte_timeout_task.fn == NULL);
+            aws_task_init(
+                &stream->base.client_data->response_first_byte_timeout_task,
+                s_http_stream_response_first_byte_timeout_task,
+                stream,
+                "http_stream_response_first_byte_timeout_task");
+            uint64_t now_ns = 0;
+            aws_channel_current_clock_time(channel, &now_ns);
+            struct aws_event_loop *connection_loop = aws_channel_get_event_loop(channel);
+            aws_event_loop_schedule_task_future(
+                connection_loop,
+                &stream->base.client_data->response_first_byte_timeout_task,
+                now_ns + aws_timestamp_convert(
+                             response_first_byte_timeout_ms, AWS_TIMESTAMP_MILLIS, AWS_TIMESTAMP_NANOS, NULL));
+        }
+    }
+}
+
 /**
  * If necessary, update `outgoing_stream` so it is pointing at a stream
  * with data to send, or NULL if all streams are done sending data.
@@ -735,13 +826,7 @@ static struct aws_h1_stream *s_update_outgoing_stream_ptr(struct aws_h1_connecti
 
     /* If current stream is done sending data... */
     if (current && !aws_h1_encoder_is_message_in_progress(&connection->thread_data.encoder)) {
-        current->is_outgoing_message_done = true;
-        AWS_ASSERT(current->base.metrics.send_end_timestamp_ns == -1);
-        aws_high_res_clock_get_ticks((uint64_t *)&current->base.metrics.send_end_timestamp_ns);
-        AWS_ASSERT(current->base.metrics.send_start_timestamp_ns != -1);
-        AWS_ASSERT(current->base.metrics.send_end_timestamp_ns >= current->base.metrics.send_start_timestamp_ns);
-        current->base.metrics.sending_duration_ns =
-            current->base.metrics.send_end_timestamp_ns - current->base.metrics.send_start_timestamp_ns;
+        s_set_outgoing_message_done(current);
 
         /* RFC-7230 section 6.6: Tear-down.
          * If this was the final stream, don't allows any further streams to be sent */
@@ -1124,16 +1209,7 @@ static int s_decoder_on_header(const struct aws_h1_decoded_header *header, void
                         AWS_LS_HTTP_STREAM,
                         "id=%p: Received 'Connection: close' header, no more request data will be sent.",
                         (void *)&incoming_stream->base);
-                    incoming_stream->is_outgoing_message_done = true;
-                    AWS_ASSERT(incoming_stream->base.metrics.send_end_timestamp_ns == -1);
-                    aws_high_res_clock_get_ticks((uint64_t *)&incoming_stream->base.metrics.send_end_timestamp_ns);
-                    AWS_ASSERT(incoming_stream->base.metrics.send_start_timestamp_ns != -1);
-                    AWS_ASSERT(
-                        incoming_stream->base.metrics.send_end_timestamp_ns >=
-                        incoming_stream->base.metrics.send_start_timestamp_ns);
-                    incoming_stream->base.metrics.sending_duration_ns =
-                        incoming_stream->base.metrics.send_end_timestamp_ns -
-                        incoming_stream->base.metrics.send_start_timestamp_ns;
+                    s_set_outgoing_message_done(incoming_stream);
                 }
                 /* Stop writing right now.
                  * Shutdown will be scheduled after we finishing parsing the response */
@@ -1856,6 +1932,15 @@ static int s_try_process_next_stream_read_message(struct aws_h1_connection *conn
     if (incoming_stream->base.metrics.receive_start_timestamp_ns == -1) {
         /* That's the first time for the stream receives any message */
         aws_high_res_clock_get_ticks((uint64_t *)&incoming_stream->base.metrics.receive_start_timestamp_ns);
+        if (incoming_stream->base.client_data &&
+            incoming_stream->base.client_data->response_first_byte_timeout_task.fn != NULL) {
+            /* There is an outstanding response timeout task, as we already received the data, we can cancel it now. We
+             * are safe to do it as we always on connection thread to schedule the task or cancel it */
+            struct aws_event_loop *connection_loop = aws_channel_get_event_loop(connection->base.channel_slot->channel);
+            /* The task will be zeroed out within the call */
+            aws_event_loop_cancel_task(
+                connection_loop, &incoming_stream->base.client_data->response_first_byte_timeout_task);
+        }
     }
 
     /* As decoder runs, it invokes the internal s_decoder_X callbacks, which in turn invoke user callbacks.
diff --git a/source/h1_stream.c b/source/h1_stream.c
@@ -407,6 +407,7 @@ struct aws_h1_stream *aws_h1_stream_new_request(
 
     stream->base.client_data = &stream->base.client_or_server_data.client;
     stream->base.client_data->response_status = AWS_HTTP_STATUS_CODE_UNKNOWN;
+    stream->base.client_data->response_first_byte_timeout_ms = options->response_first_byte_timeout_ms;
     stream->base.on_metrics = options->on_metrics;
 
     /* Validate request and cache info that the encoder will eventually need */
diff --git a/source/http.c b/source/http.c
@@ -148,6 +148,9 @@ static struct aws_error_info s_errors[] = {
     AWS_DEFINE_ERROR_INFO_HTTP(
         AWS_ERROR_HTTP_MANUAL_WRITE_HAS_COMPLETED,
         "Manual write failed because manual writes are already completed."),
+    AWS_DEFINE_ERROR_INFO_HTTP(
+        AWS_ERROR_HTTP_RESPONSE_FIRST_BYTE_TIMEOUT,
+        "The server does not begin responding within the configuration after a request is fully sent."),
 };
 /* clang-format on */
 
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -144,6 +144,8 @@ add_test_case(h1_client_switching_protocols_fails_subsequent_requests)
 add_test_case(h1_client_switching_protocols_requires_downstream_handler)
 add_test_case(h1_client_connection_close_before_request_finishes)
 add_test_case(h1_client_response_close_connection_before_request_finishes)
+add_test_case(h1_client_response_first_byte_timeout_connection)
+add_test_case(h1_client_response_first_byte_timeout_request_override)
 
 add_test_case(strutil_trim_http_whitespace)
 add_test_case(strutil_is_http_token)
diff --git a/tests/test_h1_client.c b/tests/test_h1_client.c
