use async stream

Author: azkrishpy

include/aws/http/private/h1_encoder.h

@@ -33,6 +33,8 @@ struct aws_h1_encoder_message {
     struct aws_byte_buf outgoing_head_buf;
     /* Single stream used for unchunked body */
     struct aws_input_stream *body;
+    /* Single async stream used for unchunked body */
+    struct aws_async_input_stream *async_body;
 
     /* Pointer to list of `struct aws_h1_chunk`, used for chunked encoding.
      * List is owned by aws_h1_stream.
@@ -55,6 +57,7 @@ enum aws_h1_encoder_state {
     AWS_H1_ENCODER_STATE_HEAD,
     /* Write streaming body, without chunked encoding, because Content-Length is known */
     AWS_H1_ENCODER_STATE_UNCHUNKED_BODY_STREAM,
+    AWS_H1_ENCODER_STATE_ASYNC_WAITING,
     /* Write streaming body, with chunked encoding, because Content-Length is unknown */
     AWS_H1_ENCODER_STATE_CHUNKED_BODY_STREAM,
     AWS_H1_ENCODER_STATE_CHUNKED_BODY_STREAM_LAST_CHUNK,
@@ -81,6 +84,12 @@ struct aws_h1_encoder {
     uint64_t chunk_count;
     /* Encoder logs with this stream ptr as the ID, and passes this ptr to the chunk_complete callback */
     struct aws_http_stream *current_stream;
+    /* Future to record pending future of async body read */
+    struct aws_future_bool *pending_async_future;
+    /* Connection to schedule again while async reading. */
+    struct aws_h1_connection *connection;
+    /* async error recorded */
+    int async_error;
 };
 
 struct aws_h1_chunk *aws_h1_chunk_new(struct aws_allocator *allocator, const struct aws_http1_chunk_options *options);

include/aws/http/request_response.h

@@ -865,6 +865,13 @@ int aws_http_message_set_response_status(struct aws_http_message *response_messa
 AWS_HTTP_API
 struct aws_input_stream *aws_http_message_get_body_stream(const struct aws_http_message *message);
 
+/**
+ * Get the async body stream.
+ * Returns NULL if no async body stream is set.
+ */
+AWS_HTTP_API
+struct aws_async_input_stream *aws_http_message_get_async_body_stream(const struct aws_http_message *message);
+
 /**
  * Set the body stream.
  * NULL is an acceptable value for messages with no body.
@@ -874,6 +881,15 @@ struct aws_input_stream *aws_http_message_get_body_stream(const struct aws_http_
 AWS_HTTP_API
 void aws_http_message_set_body_stream(struct aws_http_message *message, struct aws_input_stream *body_stream);
 
+/**
+ * Set the async body stream.
+ * NULL is an acceptable value for messages with no body.
+ * Note: The message does NOT take ownership of the body stream.
+ * The stream must not be destroyed until the message is complete.
+ */
+AWS_HTTP_API
+void aws_http_message_set_async_body_stream(struct aws_http_message *message, struct aws_async_input_stream *async_body_stream);
+
 /**
  * aws_future<aws_http_message*>
  */

source/h1_connection.c

@@ -1113,6 +1113,15 @@ static void s_write_outgoing_stream(struct aws_h1_connection *connection, bool f
         goto error;
     }
 
+    if (connection->thread_data.encoder.async_error) {
+        /* Error receiving data asynchronously. Need to note when this is happening, but for now, abandon ship */
+        if (msg) {
+            aws_mem_release(msg->allocator, msg);
+        }
+        s_shutdown_due_to_error(connection, connection->thread_data.encoder.async_error);
+        return;
+    }
+
     if (msg->message_data.len > 0) {
         AWS_LOGF_TRACE(
             AWS_LS_HTTP_CONNECTION,
@@ -1130,7 +1139,13 @@ static void s_write_outgoing_stream(struct aws_h1_connection *connection, bool f
 
             goto error;
         }
-
+    } else if(connection->thread_data.encoder.message->async_body) {
+        AWS_LOGF_TRACE(
+            AWS_LS_HTTP_CONNECTION,
+            "id=%p: Outgoing async stream task is either complete or waiting on future. Never reschedule task.",
+            (void *)&connection->base);
+        aws_mem_release(msg->allocator, msg);
+        connection->thread_data.is_outgoing_stream_task_active = false;
     } else {
         /* If message is empty, warn that no work is being done
          * and reschedule the task to try again next tick.
@@ -1552,6 +1567,9 @@ static struct aws_h1_connection *s_connection_new(
 
     aws_h1_encoder_init(&connection->thread_data.encoder, alloc);
 
+    /* hacking around with adding connection to encoder. there should be a better way to do it */
+    connection->thread_data.encoder.connection = connection;
+
     aws_channel_task_init(
         &connection->outgoing_stream_task, s_outgoing_stream_task, connection, "http1_connection_outgoing_stream");
     aws_channel_task_init(

source/h1_encoder.c

@@ -2,11 +2,15 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0.
  */
+#include <aws/common/error.h>
+#include <aws/io/future.h>
 #include <aws/http/private/h1_encoder.h>
+#include <aws/http/private/h1_connection.h>
 #include <aws/http/private/strutil.h>
 #include <aws/http/status_code.h>
 #include <aws/io/logging.h>
 #include <aws/io/stream.h>
+#include <aws/io/async_stream.h>
 
 #include <inttypes.h>
 
@@ -249,6 +253,7 @@ int aws_h1_encoder_message_init_from_request(
     AWS_ZERO_STRUCT(*message);
 
     message->body = aws_input_stream_acquire(aws_http_message_get_body_stream(request));
+    message->async_body = aws_async_input_stream_acquire(aws_http_message_get_async_body_stream(request));
     message->pending_chunk_list = pending_chunk_list;
 
     struct aws_byte_cursor method;
@@ -693,6 +698,57 @@ static int s_switch_state(struct aws_h1_encoder *encoder, enum aws_h1_encoder_st
     return AWS_OP_SUCCESS;
 }
 
+static void s_on_async_body_read_complete(void *user_data) {
+    struct aws_h1_encoder *encoder = user_data;
+
+    struct aws_h1_connection *connection = encoder->connection;
+
+    int error = aws_future_bool_get_error(encoder->pending_async_future);
+    if (error) {
+        ENCODER_LOG(ERROR, encoder, "Encountered error after future was complete. Setting async_error, should be caught in the connection event loop.");
+        encoder->async_error = error;
+    }
+
+    bool eof = !error && aws_future_bool_get_result(encoder->pending_async_future);
+
+    aws_future_bool_release(encoder->pending_async_future);
+    encoder->pending_async_future = NULL;
+
+    if (eof) {
+        s_switch_state(encoder, AWS_H1_ENCODER_STATE_DONE);
+    } else {
+        ENCODER_LOG(DEBUG, encoder, "Error occurred or buffer was full but eof not reached. We have to initiate a new encode request with a new buffer.");
+        s_switch_state(encoder, AWS_H1_ENCODER_STATE_UNCHUNKED_BODY_STREAM);
+    }
+
+    aws_h1_connection_try_write_outgoing_stream(connection);
+}
+
+static int s_encode_stream_async(
+    struct aws_h1_encoder *encoder,
+    struct aws_byte_buf *dst,
+    struct aws_async_input_stream *stream) {
+
+    if (dst->capacity == dst->len) {
+        return AWS_OP_ERR;
+    }
+
+    ENCODER_LOG(TRACE, encoder, "Reading from async body stream.");
+
+    encoder->pending_async_future = aws_async_input_stream_read_to_fill(stream, dst);
+
+    if (aws_future_bool_is_done(encoder->pending_async_future)) {
+        s_on_async_body_read_complete(encoder);
+        return encoder->async_error;
+    }
+
+    aws_future_bool_register_callback(encoder->pending_async_future, s_on_async_body_read_complete, encoder);
+
+    s_switch_state(encoder, AWS_H1_ENCODER_STATE_ASYNC_WAITING);
+
+    return AWS_OP_SUCCESS;
+}
+
 /* Initial state. Waits until a new message is set */
 static int s_state_fn_init(struct aws_h1_encoder *encoder, struct aws_byte_buf *dst) {
     (void)dst;
@@ -720,7 +776,8 @@ static int s_state_fn_head(struct aws_h1_encoder *encoder, struct aws_byte_buf *
     aws_byte_buf_clean_up(&encoder->message->outgoing_head_buf);
 
     /* Pick next state */
-    if (encoder->message->body && encoder->message->content_length) {
+    /* Experimentally supporting async streams for unchunked requests.*/
+    if ((encoder->message->body || encoder->message->async_body) && encoder->message->content_length) {
         return s_switch_state(encoder, AWS_H1_ENCODER_STATE_UNCHUNKED_BODY_STREAM);
 
     } else if (encoder->message->body && encoder->message->has_chunked_encoding_header) {
@@ -736,6 +793,9 @@ static int s_state_fn_head(struct aws_h1_encoder *encoder, struct aws_byte_buf *
 
 /* Write out body with known Content-Length (not using chunked encoding). */
 static int s_state_fn_unchunked_body_stream(struct aws_h1_encoder *encoder, struct aws_byte_buf *dst) {
+    if (encoder->message->async_body) {
+        return s_encode_stream_async(encoder, dst, encoder->message->async_body);
+    }
     bool done;
     if (s_encode_stream(encoder, dst, encoder->message->body, encoder->message->content_length, &done)) {
         return AWS_OP_ERR;
@@ -750,6 +810,13 @@ static int s_state_fn_unchunked_body_stream(struct aws_h1_encoder *encoder, stru
     return s_switch_state(encoder, AWS_H1_ENCODER_STATE_DONE);
 }
 
+static int s_state_fn_async_waiting(struct aws_h1_encoder *encoder, struct aws_byte_buf *dst) {
+    (void) dst;
+    ENCODER_LOG(ERROR, encoder, "This point should never be reached. We should come back to the encoder only after the state has changed from ASYNC WAITING");
+
+    return AWS_OP_ERR;
+}
+
 /* Write out body (of unknown Content-Length) using chunked encoding.
  * Each pass through this state writes out 1 chunk of body data (or nothing at all). */
 static int s_state_fn_chunked_body_stream(struct aws_h1_encoder *encoder, struct aws_byte_buf *dst) {
@@ -994,6 +1061,7 @@ static struct encoder_state_def s_encoder_states[] = {
     [AWS_H1_ENCODER_STATE_INIT] = {.fn = s_state_fn_init, .name = "INIT"},
     [AWS_H1_ENCODER_STATE_HEAD] = {.fn = s_state_fn_head, .name = "HEAD"},
     [AWS_H1_ENCODER_STATE_UNCHUNKED_BODY_STREAM] = {.fn = s_state_fn_unchunked_body_stream, .name = "BODY"},
+    [AWS_H1_ENCODER_STATE_ASYNC_WAITING] = {.fn = s_state_fn_async_waiting, .name = "WAITING"},
     [AWS_H1_ENCODER_STATE_CHUNKED_BODY_STREAM] = {.fn = s_state_fn_chunked_body_stream, .name = "CHUNKED_BODY_STREAM"},
     [AWS_H1_ENCODER_STATE_CHUNKED_BODY_STREAM_LAST_CHUNK] =
         {.fn = s_state_fn_chunked_body_stream_last_chunk, .name = "LAST_CHUNK"},

source/request_response.c

@@ -13,6 +13,7 @@
 #include <aws/http/status_code.h>
 #include <aws/io/logging.h>
 #include <aws/io/stream.h>
+#include <aws/io/async_stream.h>
 
 #ifdef _MSC_VER
 #    pragma warning(disable : 4204) /* non-constant aggregate initializer */
@@ -453,6 +454,7 @@ struct aws_http_message {
     struct aws_allocator *allocator;
     struct aws_http_headers *headers;
     struct aws_input_stream *body_stream;
+    struct aws_async_input_stream *async_body_stream;
     struct aws_atomic_var refcount;
     enum aws_http_version http_version;
 
@@ -781,6 +783,17 @@ void aws_http_message_set_body_stream(struct aws_http_message *message, struct a
     }
 }
 
+void aws_http_message_set_async_body_stream(struct aws_http_message *message, struct aws_async_input_stream *async_body_stream) {
+    AWS_PRECONDITION(message);
+    /* release previous stream, if any */
+    aws_async_input_stream_release(message->async_body_stream);
+
+    message->async_body_stream = async_body_stream;
+    if (message->async_body_stream) {
+        aws_async_input_stream_acquire(message->async_body_stream);
+    }
+}
+
 int aws_http1_stream_write_chunk(struct aws_http_stream *http1_stream, const struct aws_http1_chunk_options *options) {
     AWS_PRECONDITION(http1_stream);
     AWS_PRECONDITION(http1_stream->vtable);
@@ -829,6 +842,11 @@ struct aws_input_stream *aws_http_message_get_body_stream(const struct aws_http_
     return message->body_stream;
 }
 
+struct aws_async_input_stream *aws_http_message_get_async_body_stream(const struct aws_http_message *message) {
+    AWS_PRECONDITION(message);
+    return message->async_body_stream;
+}
+
 struct aws_http_headers *aws_http_message_get_headers(const struct aws_http_message *message) {
     AWS_PRECONDITION(message);
     return message->headers;