diff --git a/include/aws/http/connection_manager.h b/include/aws/http/connection_manager.h index fd8d59d8..d3b28b76 100644 --- a/include/aws/http/connection_manager.h +++ b/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) diff --git a/include/aws/http/http.h b/include/aws/http/http.h index 1225ec1e..0c66b329 100644 --- a/include/aws/http/http.h +++ b/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) }; diff --git a/source/connection_manager.c b/source/connection_manager.c index efa768c0..5efe02a7 100644 --- a/source/connection_manager.c +++ b/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); diff --git a/source/http.c b/source/http.c index 1c68e40c..22467406 100644 --- a/source/http.c +++ b/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 */ diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5776180b..6cd95462 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -540,6 +540,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 diff --git a/tests/test_connection_manager.c b/tests/test_connection_manager.c index 998c1559..7715c7db 100644 --- a/tests/test_connection_manager.c +++ b/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;