Destroy callback for http stream (#390)

TingDaoK · web-flow · commit db4111947160 · 2022-09-23T14:13:21.000-07:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -6,7 +6,7 @@ on:
       - 'main'
 
 env:
-  BUILDER_VERSION: v0.9.15
+  BUILDER_VERSION: v0.9.18
   BUILDER_SOURCE: releases
   BUILDER_HOST: https://d19elf31gohf1l.cloudfront.net
   PACKAGE_NAME: aws-c-http
@@ -117,7 +117,7 @@ jobs:
       run: |
         python -c "from urllib.request import urlretrieve; urlretrieve('${{ env.BUILDER_HOST }}/${{ env.BUILDER_SOURCE }}/${{ env.BUILDER_VERSION }}/builder.pyz?run=${{ env.RUN }}', 'builder.pyz')"
         python builder.pyz build -p ${{ env.PACKAGE_NAME }} --cmake-extra=-DBUILD_SHARED_LIBS=ON
-  
+
   windows-app-verifier:
     runs-on: windows-2022 # latest
     steps:
diff --git a/include/aws/http/private/request_response_impl.h b/include/aws/http/private/request_response_impl.h
@@ -44,6 +44,7 @@ struct aws_http_stream {
     aws_http_on_incoming_header_block_done_fn *on_incoming_header_block_done;
     aws_http_on_incoming_body_fn *on_incoming_body;
     aws_http_on_stream_complete_fn *on_complete;
+    aws_http_on_stream_destroy_fn *on_destroy;
 
     struct aws_atomic_var refcount;
     enum aws_http_method request_method;
diff --git a/include/aws/http/request_response.h b/include/aws/http/request_response.h
@@ -182,11 +182,18 @@ typedef int(
 typedef int(aws_http_on_incoming_request_done_fn)(struct aws_http_stream *stream, void *user_data);
 
 /**
- * Invoked when request/response stream is complete, whether successful or unsuccessful
- * This is always invoked on the HTTP connection's event-loop thread.
+ * Invoked when request/response stream is completely destroyed.
+ * This may be invoked synchronously when aws_http_stream_release() is called.
+ * This is invoked even if the stream is never activated.
  */
 typedef void(aws_http_on_stream_complete_fn)(struct aws_http_stream *stream, int error_code, void *user_data);
 
+/**
+ * Invoked when request/response stream destroy completely.
+ * This can be invoked within the same thead who release the refcount on http stream.
+ */
+typedef void(aws_http_on_stream_destroy_fn)(void *user_data);
+
 /**
  * Options for creating a stream which sends a request from the client and receives a response from the server.
  */
@@ -234,6 +241,9 @@ struct aws_http_make_request_options {
      */
     aws_http_on_stream_complete_fn *on_complete;
 
+    /* Callback for when the request/response stream is completely destroyed. */
+    aws_http_on_stream_destroy_fn *on_destroy;
+
     /**
      * When using HTTP/2, request body data will be provided over time. The stream will only be polled for writing
      * when data has been supplied via `aws_http2_stream_write_data`
@@ -290,6 +300,9 @@ struct aws_http_request_handler_options {
      * See `aws_http_on_stream_complete_fn`.
      */
     aws_http_on_stream_complete_fn *on_complete;
+
+    /* Callback for when the request/response stream is completely destroyed. */
+    aws_http_on_stream_destroy_fn *on_destroy;
 };
 
 /**
diff --git a/source/h1_stream.c b/source/h1_stream.c
@@ -342,7 +342,8 @@ static struct aws_h1_stream *s_stream_new_common(
     aws_http_on_incoming_headers_fn *on_incoming_headers,
     aws_http_on_incoming_header_block_done_fn *on_incoming_header_block_done,
     aws_http_on_incoming_body_fn *on_incoming_body,
-    aws_http_on_stream_complete_fn on_complete) {
+    aws_http_on_stream_complete_fn *on_complete,
+    aws_http_on_stream_destroy_fn *on_destroy) {
 
     struct aws_h1_connection *connection = AWS_CONTAINER_OF(connection_base, struct aws_h1_connection, base);
 
@@ -359,6 +360,7 @@ static struct aws_h1_stream *s_stream_new_common(
     stream->base.on_incoming_header_block_done = on_incoming_header_block_done;
     stream->base.on_incoming_body = on_incoming_body;
     stream->base.on_complete = on_complete;
+    stream->base.on_destroy = on_destroy;
 
     aws_channel_task_init(
         &stream->cross_thread_work_task, s_stream_cross_thread_work_task, stream, "http1_stream_cross_thread_work");
@@ -384,7 +386,8 @@ struct aws_h1_stream *aws_h1_stream_new_request(
         options->on_response_headers,
         options->on_response_header_block_done,
         options->on_response_body,
-        options->on_complete);
+        options->on_complete,
+        options->on_destroy);
     if (!stream) {
         return NULL;
     }
@@ -430,7 +433,8 @@ struct aws_h1_stream *aws_h1_stream_new_request_handler(const struct aws_http_re
         options->on_request_headers,
         options->on_request_header_block_done,
         options->on_request_body,
-        options->on_complete);
+        options->on_complete,
+        options->on_destroy);
     if (!stream) {
         return NULL;
     }
diff --git a/source/h2_stream.c b/source/h2_stream.c
@@ -241,6 +241,7 @@ struct aws_h2_stream *aws_h2_stream_new_request(
     stream->base.on_incoming_header_block_done = options->on_response_header_block_done;
     stream->base.on_incoming_body = options->on_response_body;
     stream->base.on_complete = options->on_complete;
+    stream->base.on_destroy = options->on_destroy;
     stream->base.client_data = &stream->base.client_or_server_data.client;
     stream->base.client_data->response_status = AWS_HTTP_STATUS_CODE_UNKNOWN;
     struct aws_byte_cursor method;
diff --git a/source/http2_stream_manager.c b/source/http2_stream_manager.c
@@ -804,10 +804,17 @@ static void s_on_stream_complete(struct aws_http_stream *stream, int error_code,
         pending_stream_acquisition->options.on_complete(
             stream, error_code, pending_stream_acquisition->options.user_data);
     }
-    s_pending_stream_acquisition_destroy(pending_stream_acquisition);
     s_sm_connection_on_scheduled_stream_finishes(sm_connection, stream_manager);
 }
 
+static void s_on_stream_destroy(void *user_data) {
+    struct aws_h2_sm_pending_stream_acquisition *pending_stream_acquisition = user_data;
+    if (pending_stream_acquisition->options.on_destroy) {
+        pending_stream_acquisition->options.on_destroy(pending_stream_acquisition->options.user_data);
+    }
+    s_pending_stream_acquisition_destroy(pending_stream_acquisition);
+}
+
 /* Scheduled to happen from connection's thread */
 static void s_make_request_task(struct aws_channel_task *task, void *arg, enum aws_task_status status) {
     (void)task;
@@ -862,6 +869,7 @@ static void s_make_request_task(struct aws_channel_task *task, void *arg, enum a
         .on_response_header_block_done = s_on_incoming_header_block_done,
         .on_response_body = s_on_incoming_body,
         .on_complete = s_on_stream_complete,
+        .on_destroy = s_on_stream_destroy,
         .user_data = pending_stream_acquisition,
     };
     /* TODO: we could put the pending acquisition back to the list if the connection is not available for new request.
diff --git a/source/request_response.c b/source/request_response.c
@@ -1070,9 +1070,16 @@ void aws_http_stream_release(struct aws_http_stream *stream) {
     if (prev_refcount == 1) {
         AWS_LOGF_TRACE(AWS_LS_HTTP_STREAM, "id=%p: Final stream refcount released.", (void *)stream);
 
+        void *user_data = stream->user_data;
+        aws_http_on_stream_destroy_fn *on_destroy_callback = stream->on_destroy;
+
         struct aws_http_connection *owning_connection = stream->owning_connection;
         stream->vtable->destroy(stream);
 
+        if (on_destroy_callback) {
+            /* info user that destroy completed. */
+            on_destroy_callback(user_data);
+        }
         /* Connection needed to outlive stream, but it's free to go now */
         aws_http_connection_release(owning_connection);
     } else {
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -100,6 +100,8 @@ add_test_case(h1_client_request_send_multiple_chunked_encoding)
 add_test_case(h1_client_request_close_header_ends_connection)
 add_test_case(h1_client_request_close_header_with_pipelining)
 add_test_case(h1_client_request_close_header_with_chunked_encoding_and_pipelining)
+add_test_case(h1_client_stream_release_after_complete)
+add_test_case(h1_client_stream_release_before_complete)
 add_test_case(h1_client_response_get_1liner)
 add_test_case(h1_client_response_get_headers)
 add_test_case(h1_client_response_get_body)
@@ -375,6 +377,7 @@ add_h2_decoder_test_set(h2_decoder_err_bad_preface_from_client_3)
 
 add_test_case(h2_client_sanity_check)
 add_test_case(h2_client_stream_create)
+add_test_case(h2_client_stream_release_after_complete)
 add_test_case(h2_client_unactivated_stream_cleans_up)
 add_test_case(h2_client_connection_preface_sent)
 add_test_case(h2_client_auto_ping_ack)
diff --git a/tests/stream_test_helper.c b/tests/stream_test_helper.c
@@ -135,6 +135,14 @@ static void s_on_complete(struct aws_http_stream *stream, int error_code, void *
     aws_http_stream_get_incoming_response_status(stream, &tester->response_status);
 }
 
+static void s_on_destroy(void *user_data) {
+    struct client_stream_tester *tester = user_data;
+
+    /* Validate things are firing properly */
+    AWS_FATAL_ASSERT(!tester->destroyed);
+    tester->destroyed = true;
+}
+
 int client_stream_tester_init(
     struct client_stream_tester *tester,
     struct aws_allocator *alloc,
@@ -166,6 +174,7 @@ int client_stream_tester_init(
         .on_response_header_block_done = s_on_header_block_done,
         .on_response_body = s_on_body,
         .on_complete = s_on_complete,
+        .on_destroy = s_on_destroy,
     };
     tester->stream = aws_http_connection_make_request(options->connection, &request_options);
     ASSERT_NOT_NULL(tester->stream);
diff --git a/tests/stream_test_helper.h b/tests/stream_test_helper.h
@@ -43,6 +43,8 @@ struct client_stream_tester {
 
     /* Whether connection is open when on_complete fires */
     bool on_complete_connection_is_open;
+
+    bool destroyed;
 };
 
 struct client_stream_tester_options {
diff --git a/tests/test_h1_client.c b/tests/test_h1_client.c
@@ -1728,6 +1728,75 @@ static int s_stream_tester_init(
     return client_stream_tester_init(tester, main_tester->alloc, &options);
 }
 
+H1_CLIENT_TEST_CASE(h1_client_stream_release_after_complete) {
+    (void)ctx;
+    struct tester tester;
+    ASSERT_SUCCESS(s_tester_init(&tester, allocator));
+
+    /* send request */
+    struct aws_http_message *request = s_new_default_get_request(allocator);
+
+    struct client_stream_tester stream_tester;
+    ASSERT_SUCCESS(s_stream_tester_init(&stream_tester, &tester, request));
+
+    testing_channel_drain_queued_tasks(&tester.testing_channel);
+
+    /* Ensure the request can be destroyed after request is sent */
+    aws_http_message_destroy(request);
+
+    /* send response */
+    ASSERT_SUCCESS(testing_channel_push_read_str(&tester.testing_channel, "HTTP/1.1 204 No Content\r\n\r\n"));
+
+    testing_channel_drain_queued_tasks(&tester.testing_channel);
+
+    /* check result */
+    ASSERT_TRUE(stream_tester.complete);
+    ASSERT_FALSE(stream_tester.destroyed);
+
+    aws_http_stream_release(stream_tester.stream);
+    stream_tester.stream = NULL;
+    ASSERT_TRUE(stream_tester.destroyed);
+
+    /* clean up */
+    client_stream_tester_clean_up(&stream_tester);
+    ASSERT_SUCCESS(s_tester_clean_up(&tester));
+    return AWS_OP_SUCCESS;
+}
+
+H1_CLIENT_TEST_CASE(h1_client_stream_release_before_complete) {
+    (void)ctx;
+    struct tester tester;
+    ASSERT_SUCCESS(s_tester_init(&tester, allocator));
+
+    /* send request */
+    struct aws_http_message *request = s_new_default_get_request(allocator);
+
+    struct client_stream_tester stream_tester;
+    ASSERT_SUCCESS(s_stream_tester_init(&stream_tester, &tester, request));
+    aws_http_stream_release(stream_tester.stream);
+    stream_tester.stream = NULL;
+    ASSERT_FALSE(stream_tester.destroyed);
+
+    testing_channel_drain_queued_tasks(&tester.testing_channel);
+
+    /* Ensure the request can be destroyed after request is sent */
+    aws_http_message_destroy(request);
+
+    /* send response */
+    ASSERT_SUCCESS(testing_channel_push_read_str(&tester.testing_channel, "HTTP/1.1 204 No Content\r\n\r\n"));
+
+    testing_channel_drain_queued_tasks(&tester.testing_channel);
+
+    /* check result */
+    ASSERT_TRUE(stream_tester.complete);
+    ASSERT_TRUE(stream_tester.destroyed);
+
+    /* clean up */
+    client_stream_tester_clean_up(&stream_tester);
+    ASSERT_SUCCESS(s_tester_clean_up(&tester));
+    return AWS_OP_SUCCESS;
+}
+
 H1_CLIENT_TEST_CASE(h1_client_response_get_1liner) {
     (void)ctx;
     struct tester tester;
@@ -3561,11 +3630,17 @@ H1_CLIENT_TEST_CASE(h1_client_close_from_on_thread_makes_not_open) {
     return AWS_OP_SUCCESS;
 }
 
+static void s_unactivated_stream_cleans_up_on_destroy(void *data) {
+    bool *destroyed = data;
+    *destroyed = true;
+}
+
 H1_CLIENT_TEST_CASE(h1_client_unactivated_stream_cleans_up) {
     (void)ctx;
     struct tester tester;
     ASSERT_SUCCESS(s_tester_init(&tester, allocator));
     ASSERT_TRUE(aws_http_connection_is_open(tester.connection));
+    bool destroyed = false;
 
     struct aws_http_message *request = aws_http_message_new_request(allocator);
     ASSERT_SUCCESS(aws_http_message_set_request_method(request, aws_byte_cursor_from_c_str("GET")));
@@ -3574,13 +3649,17 @@ H1_CLIENT_TEST_CASE(h1_client_unactivated_stream_cleans_up) {
     struct aws_http_make_request_options options = {
         .self_size = sizeof(struct aws_http_make_request_options),
         .request = request,
+        .on_destroy = s_unactivated_stream_cleans_up_on_destroy,
+        .user_data = &destroyed,
     };
 
     struct aws_http_stream *stream = aws_http_connection_make_request(tester.connection, &options);
     aws_http_message_release(request);
     ASSERT_NOT_NULL(stream);
     /* we do not activate, that is the test. */
+    ASSERT_FALSE(destroyed);
     aws_http_stream_release(stream);
+    ASSERT_TRUE(destroyed);
     aws_http_connection_close(tester.connection);
     ASSERT_SUCCESS(s_tester_clean_up(&tester));
     return AWS_OP_SUCCESS;
diff --git a/tests/test_h2_client.c b/tests/test_h2_client.c
@@ -183,6 +183,52 @@ TEST_CASE(h2_client_stream_create) {
     return s_tester_clean_up();
 }
 
+static void s_stream_cleans_up_on_destroy(void *data) {
+    bool *destroyed = data;
+    *destroyed = true;
+}
+
+TEST_CASE(h2_client_stream_release_after_complete) {
+    ASSERT_SUCCESS(s_tester_init(allocator, ctx));
+
+    /* create request */
+    struct aws_http_message *request = aws_http2_message_new_request(allocator);
+    ASSERT_NOT_NULL(request);
+
+    struct aws_http_header headers[] = {
+        DEFINE_HEADER(":method", "GET"),
+        DEFINE_HEADER(":scheme", "https"),
+        DEFINE_HEADER(":path", "/"),
+    };
+    ASSERT_SUCCESS(aws_http_message_add_header_array(request, headers, AWS_ARRAY_SIZE(headers)));
+
+    bool destroyed = false;
+    struct aws_http_make_request_options options = {
+        .self_size = sizeof(options),
+        .request = request,
+        .on_destroy = s_stream_cleans_up_on_destroy,
+        .user_data = &destroyed,
+    };
+
+    struct aws_http_stream *stream = aws_http_connection_make_request(s_tester.connection, &options);
+    ASSERT_NOT_NULL(stream);
+    aws_http_stream_activate(stream);
+
+    /* shutdown channel so request can be released */
+    aws_channel_shutdown(s_tester.testing_channel.channel, AWS_ERROR_SUCCESS);
+    testing_channel_drain_queued_tasks(&s_tester.testing_channel);
+    ASSERT_TRUE(testing_channel_is_shutdown_completed(&s_tester.testing_channel));
+
+    /* release request */
+    ASSERT_FALSE(destroyed);
+    aws_http_stream_release(stream);
+    ASSERT_TRUE(destroyed);
+
+    aws_http_message_release(request);
+
+    return s_tester_clean_up();
+}
+
 TEST_CASE(h2_client_unactivated_stream_cleans_up) {
     ASSERT_SUCCESS(s_tester_init(allocator, ctx));
 
diff --git a/tests/test_stream_manager.c b/tests/test_stream_manager.c
