Merge branch 'main' into connection-shutdown

include/aws/http/connection_manager.h

@@ -132,6 +132,13 @@ struct aws_http_connection_manager_options {
      */
     uint64_t connection_acquisition_timeout_ms;
 
+    /*
+     * If set to a non-zero value, aws_http_connection_manager_acquire_connection() calls will fail with
+     * AWS_ERROR_HTTP_CONNECTION_MANAGER_MAX_PENDING_ACQUISITIONS_EXCEEDED if there are already pending acquisitions
+     * equal to `max_pending_connection_acquisitions`.
+     */
+    uint64_t max_pending_connection_acquisitions;
+
     /**
      * THIS IS AN EXPERIMENTAL AND UNSTABLE API
      * (Optional)

include/aws/http/http.h

@@ -61,6 +61,7 @@ enum aws_http_errors {
     AWS_ERROR_HTTP_MANUAL_WRITE_HAS_COMPLETED,
     AWS_ERROR_HTTP_RESPONSE_FIRST_BYTE_TIMEOUT,
     AWS_ERROR_HTTP_CONNECTION_MANAGER_ACQUISITION_TIMEOUT,
+    AWS_ERROR_HTTP_CONNECTION_MANAGER_MAX_PENDING_ACQUISITIONS_EXCEEDED,
 
     AWS_ERROR_HTTP_END_RANGE = AWS_ERROR_ENUM_END_RANGE(AWS_C_HTTP_PACKAGE_ID)
 };

source/connection_manager.c

@@ -291,6 +291,8 @@ struct aws_http_connection_manager {
 
     uint64_t connection_acquisition_timeout_ms;
 
+    uint64_t max_pending_connection_acquisitions;
+
     /*
      * Task to cull idle connections.  This task is run periodically on the cull_event_loop if a non-zero
      * culling time interval is specified.
@@ -955,6 +957,7 @@ struct aws_http_connection_manager *aws_http_connection_manager_new(
     manager->enable_read_back_pressure = options->enable_read_back_pressure;
     manager->max_connection_idle_in_milliseconds = options->max_connection_idle_in_milliseconds;
     manager->connection_acquisition_timeout_ms = options->connection_acquisition_timeout_ms;
+    manager->max_pending_connection_acquisitions = options->max_pending_connection_acquisitions;
 
     if (options->proxy_ev_settings) {
         manager->proxy_ev_settings = *options->proxy_ev_settings;
@@ -1307,8 +1310,14 @@ void aws_http_connection_manager_acquire_connection(
     /* It's a use after free crime, we don't want to handle */
     AWS_FATAL_ASSERT(manager->state == AWS_HCMST_READY);
 
-    aws_linked_list_push_back(&manager->pending_acquisitions, &request->node);
-    ++manager->pending_acquisition_count;
+    if (manager->max_pending_connection_acquisitions == 0 ||
+        manager->pending_acquisition_count < manager->max_pending_connection_acquisitions) {
+        aws_linked_list_push_back(&manager->pending_acquisitions, &request->node);
+        ++manager->pending_acquisition_count;
+    } else {
+        request->error_code = AWS_ERROR_HTTP_CONNECTION_MANAGER_MAX_PENDING_ACQUISITIONS_EXCEEDED;
+        aws_linked_list_push_back(&work.completions, &request->node);
+    }
 
     s_aws_http_connection_manager_build_transaction(&work);
 

source/http.c

@@ -154,6 +154,9 @@ static struct aws_error_info s_errors[] = {
     AWS_DEFINE_ERROR_INFO_HTTP(
         AWS_ERROR_HTTP_CONNECTION_MANAGER_ACQUISITION_TIMEOUT,
         "Connection Manager failed to acquire a connection within the defined timeout."),
+    AWS_DEFINE_ERROR_INFO_HTTP(
+        AWS_ERROR_HTTP_CONNECTION_MANAGER_MAX_PENDING_ACQUISITIONS_EXCEEDED,
+        "Max pending acquisitions reached"),
 };
 /* clang-format on */
 

tests/CMakeLists.txt

@@ -541,6 +541,7 @@ add_net_test_case(test_connection_manager_many_http2_connections)
 add_net_test_case(test_connection_manager_acquire_release)
 add_net_test_case(test_connection_manager_close_and_release)
 add_net_test_case(test_connection_manager_acquire_release_mix)
+add_net_test_case(test_connection_manager_max_pending_acquisitions)
 
 # Integration test that requires proxy envrionment in us-east-1 region.
 # TODO: test the server name validation properly

tests/test_connection_manager.c

@@ -54,6 +54,7 @@ struct cm_tester_options {
     size_t max_connections;
     uint64_t max_connection_idle_in_ms;
     uint64_t connection_acquisition_timeout_ms;
+    uint64_t max_pending_connection_acquisitions;
 
     uint64_t starting_mock_time;
     bool http2;
@@ -228,6 +229,7 @@ static int s_cm_tester_init(struct cm_tester_options *options) {
         .shutdown_complete_callback = s_cm_tester_on_cm_shutdown_complete,
         .max_connection_idle_in_milliseconds = options->max_connection_idle_in_ms,
         .connection_acquisition_timeout_ms = options->connection_acquisition_timeout_ms,
+        .max_pending_connection_acquisitions = options->max_pending_connection_acquisitions,
         .http2_prior_knowledge = !options->use_tls && options->http2,
         .initial_settings_array = options->initial_settings_array,
         .num_initial_settings = options->num_initial_settings,
@@ -743,6 +745,36 @@ static int s_test_connection_manager_acquire_release_mix(struct aws_allocator *a
 }
 AWS_TEST_CASE(test_connection_manager_acquire_release_mix, s_test_connection_manager_acquire_release_mix);
 
+static int s_test_connection_manager_max_pending_acquisitions(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+
+    size_t num_connections = 2;
+    size_t num_pending_connections = 3;
+    struct cm_tester_options options = {
+        .allocator = allocator,
+        .max_connections = num_connections,
+        .max_pending_connection_acquisitions = num_connections,
+    };
+
+    ASSERT_SUCCESS(s_cm_tester_init(&options));
+
+    s_acquire_connections(num_connections + num_pending_connections);
+
+    ASSERT_SUCCESS(s_wait_on_connection_reply_count(num_connections + num_pending_connections));
+    ASSERT_UINT_EQUALS(num_pending_connections, s_tester.connection_errors);
+    for (size_t i = 0; i < num_pending_connections; i++) {
+        uint32_t error_code;
+        aws_array_list_get_at(&s_tester.connection_errors_list, &error_code, i);
+        ASSERT_UINT_EQUALS(AWS_ERROR_HTTP_CONNECTION_MANAGER_MAX_PENDING_ACQUISITIONS_EXCEEDED, error_code);
+    }
+    ASSERT_SUCCESS(s_release_connections(num_connections, false));
+
+    ASSERT_SUCCESS(s_cm_tester_clean_up());
+
+    return AWS_OP_SUCCESS;
+}
+AWS_TEST_CASE(test_connection_manager_max_pending_acquisitions, s_test_connection_manager_max_pending_acquisitions);
+
 static int s_aws_http_connection_manager_create_connection_sync_mock(
     const struct aws_http_client_connection_options *options) {
     struct cm_tester *tester = &s_tester;