Fix max_pending_connection_acquisitions to respect connection pool si…

…ze (#485)
awslabs · Sep 10, 2024 · 4b1634b · 4b1634b
1 parent 4e74ab1
commit 4b1634b
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 10 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -6,7 +6,7 @@ on:
       - 'main'
 
 env:
-  BUILDER_VERSION: v0.9.62
+  BUILDER_VERSION: v0.9.64
   BUILDER_SOURCE: releases
   BUILDER_HOST: https://d19elf31gohf1l.cloudfront.net
   PACKAGE_NAME: aws-c-http

diff --git a/include/aws/http/connection_manager.h b/include/aws/http/connection_manager.h
@@ -134,8 +134,9 @@ struct aws_http_connection_manager_options {
 
     /*
      * 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`.
+     * AWS_ERROR_HTTP_CONNECTION_MANAGER_MAX_PENDING_ACQUISITIONS_EXCEEDED if the number of pending acquisitions
+     * reaches `max_pending_connection_acquisitions` after the connection pool has reached its capacity (i.e., all
+     * `num_connections` have been vended).
      */
     uint64_t max_pending_connection_acquisitions;
 

diff --git a/source/connection_manager.c b/source/connection_manager.c
@@ -1311,7 +1311,8 @@ void aws_http_connection_manager_acquire_connection(
     AWS_FATAL_ASSERT(manager->state == AWS_HCMST_READY);
 
     if (manager->max_pending_connection_acquisitions == 0 ||
-        manager->pending_acquisition_count < manager->max_pending_connection_acquisitions) {
+        manager->pending_acquisition_count + manager->internal_ref[AWS_HCMCT_VENDED_CONNECTION] <
+            manager->max_pending_connection_acquisitions + manager->max_connections) {
         aws_linked_list_push_back(&manager->pending_acquisitions, &request->node);
         ++manager->pending_acquisition_count;
     } else {

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -546,6 +546,7 @@ 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)
+add_net_test_case(test_connection_manager_max_pending_acquisitions_with_vended_connections)
 
 # 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
@@ -749,32 +749,81 @@ static int s_test_connection_manager_max_pending_acquisitions(struct aws_allocat
     (void)ctx;
 
     size_t num_connections = 2;
-    size_t num_pending_connections = 3;
+    size_t max_pending_connection_acquisitions = 2;
+    size_t num_connections_pending_error = 4;
     struct cm_tester_options options = {
         .allocator = allocator,
         .max_connections = num_connections,
-        .max_pending_connection_acquisitions = num_connections,
+        .max_pending_connection_acquisitions = max_pending_connection_acquisitions,
     };
 
     ASSERT_SUCCESS(s_cm_tester_init(&options));
 
-    s_acquire_connections(num_connections + num_pending_connections);
+    s_acquire_connections(num_connections + max_pending_connection_acquisitions + num_connections_pending_error);
 
-    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++) {
+    ASSERT_SUCCESS(s_wait_on_connection_reply_count(num_connections + num_connections_pending_error));
+    ASSERT_UINT_EQUALS(num_connections_pending_error, s_tester.connection_errors);
+    for (size_t i = 0; i < num_connections_pending_error; 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_wait_on_connection_reply_count(
+        num_connections + max_pending_connection_acquisitions + num_connections_pending_error));
+    ASSERT_SUCCESS(s_release_connections(max_pending_connection_acquisitions, 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_test_connection_manager_max_pending_acquisitions_with_vended_connections(
+    struct aws_allocator *allocator,
+    void *ctx) {
+    (void)ctx;
+
+    size_t num_connections = 2;
+    size_t max_pending_connection_acquisitions = 2;
+    size_t num_connections_pending_error = 4;
+    struct cm_tester_options options = {
+        .allocator = allocator,
+        .max_connections = num_connections,
+        .max_pending_connection_acquisitions = max_pending_connection_acquisitions,
+    };
+
+    ASSERT_SUCCESS(s_cm_tester_init(&options));
+
+    // fill the connection pool
+    s_acquire_connections(num_connections);
+    ASSERT_SUCCESS(s_wait_on_connection_reply_count(num_connections));
+
+    // try to acquire over max_pending_connection_acquisitions
+    s_acquire_connections(max_pending_connection_acquisitions + num_connections_pending_error);
+    ASSERT_SUCCESS(s_wait_on_connection_reply_count(num_connections + num_connections_pending_error));
+
+    ASSERT_UINT_EQUALS(num_connections_pending_error, s_tester.connection_errors);
+    for (size_t i = 0; i < num_connections_pending_error; 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_wait_on_connection_reply_count(
+        num_connections + max_pending_connection_acquisitions + num_connections_pending_error));
+    ASSERT_SUCCESS(s_release_connections(max_pending_connection_acquisitions, false));
+
+    ASSERT_SUCCESS(s_cm_tester_clean_up());
+
+    return AWS_OP_SUCCESS;
+}
+AWS_TEST_CASE(
+    test_connection_manager_max_pending_acquisitions_with_vended_connections,
+    s_test_connection_manager_max_pending_acquisitions_with_vended_connections);
+
 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;