Proposal to fix reuse endpoint with refcount can reach zero (#183)

*Issue*

- There was a race condition when the endpoint ref drops to zero, the other thread can bump it up to one again.
- There was a work around to the race condition is to allow this from happening and handle it by the caller to check the refcount when clean up with lock and decide to not really clean up the resource if the refcount bumps up to one.

*Description of changes:*

- I believe the workaround described above is not a good solution
- The caller should be responsible to not bump up the ref once it drops to zero
- Add the helper for caller to hold lock to clean the resource before drop the refcount if needed.

Author: TingDaoK

include/aws/s3/private/s3_client_impl.h

@@ -31,23 +31,13 @@ enum aws_s3_connection_finish_code {
     AWS_S3_CONNECTION_FINISH_CODE_RETRY,
 };
 
-/* Callback for the owner of the endpoint when the ref count of this endpoint hits zero. This gives the owner a
- * chance to clean up any references to the endpoint, but also to short circuit clean-up. This can be needed in the case
- * of a synced table, where before the lock of the table can be grabbed, a reference for the endpoint is added by a
- * different critical section.*/
-typedef bool(aws_s3_endpoint_ref_zero_fn)(struct aws_s3_endpoint *endpoint);
-
 /* Callback for the owner of the endpoint when the endpoint has completely cleaned up. */
 typedef void(aws_s3_endpoint_shutdown_fn)(void *user_data);
 
 struct aws_s3_endpoint_options {
     /* URL of the host that this endpoint refers to. */
     struct aws_string *host_name;
 
-    /* Callback for the owner of the endpoint when the endpoint's refcount hits zero. (More details in the typedef of
-     * this callback.)*/
-    aws_s3_endpoint_ref_zero_fn *ref_count_zero_callback;
-
     /* Callback for when this endpoint completely shuts down. */
     aws_s3_endpoint_shutdown_fn *shutdown_callback;
 
@@ -83,13 +73,13 @@ struct aws_s3_endpoint {
     /* Connection manager that manages all connections to this endpoint. */
     struct aws_http_connection_manager *http_connection_manager;
 
-    /* Callback for the owner of the endpoint when the endpoint's refcount hits zero. (More details in the typedef of
-     * this callback.)*/
-    aws_s3_endpoint_ref_zero_fn *ref_count_zero_callback;
-
     /* Callback for when this endpoint completely shuts down. */
     aws_s3_endpoint_shutdown_fn *shutdown_callback;
 
+    /* True, if the endpoint is created by client. So, it need to be thread safe to manage the refcount via
+     * `aws_s3_client_endpoint_release` */
+    bool handled_by_client;
+
     void *user_data;
 };
 
@@ -129,8 +119,6 @@ struct aws_s3_client_vtable {
 
     void (*process_work)(struct aws_s3_client *client);
 
-    bool (*endpoint_ref_count_zero)(struct aws_s3_endpoint *endpoint);
-
     void (*endpoint_shutdown_callback)(void *user_data);
 
     void (*finish_destroy)(struct aws_s3_client *client);
@@ -326,6 +314,11 @@ struct aws_s3_endpoint *aws_s3_endpoint_acquire(struct aws_s3_endpoint *endpoint
 AWS_S3_API
 void aws_s3_endpoint_release(struct aws_s3_endpoint *endpoint);
 
+/* If the endpoint is created by s3 client, it will be managed by the client via a hash table that need to be protected
+ * by lock. A lock will be acquired within the call, never invoke with lock held */
+AWS_S3_API
+void aws_s3_client_endpoint_release(struct aws_s3_client *client, struct aws_s3_endpoint *endpoint);
+
 AWS_S3_API
 extern const uint32_t g_max_num_connections_per_vip;
 

source/s3_client.c

@@ -108,8 +108,6 @@ static void s_s3_client_process_work_task(struct aws_task *task, void *arg, enum
 
 static void s_s3_client_process_work_default(struct aws_s3_client *client);
 
-static bool s_s3_client_endpoint_ref_count_zero(struct aws_s3_endpoint *endpoint);
-
 static void s_s3_client_endpoint_shutdown_callback(void *user_data);
 
 /* Default factory function for creating a meta request. */
@@ -123,7 +121,6 @@ static struct aws_s3_client_vtable s_s3_client_default_vtable = {
     .get_host_address_count = aws_host_resolver_get_host_address_count,
     .schedule_process_work_synced = s_s3_client_schedule_process_work_synced_default,
     .process_work = s_s3_client_process_work_default,
-    .endpoint_ref_count_zero = s_s3_client_endpoint_ref_count_zero,
     .endpoint_shutdown_callback = s_s3_client_endpoint_shutdown_callback,
     .finish_destroy = s_s3_client_finish_destroy_default,
 };
@@ -682,7 +679,6 @@ struct aws_s3_meta_request *aws_s3_client_make_meta_request(
         if (was_created) {
             struct aws_s3_endpoint_options endpoint_options = {
                 .host_name = endpoint_host_name,
-                .ref_count_zero_callback = client->vtable->endpoint_ref_count_zero,
                 .shutdown_callback = client->vtable->endpoint_shutdown_callback,
                 .client_bootstrap = client->client_bootstrap,
                 .tls_connection_options = is_https ? client->tls_connection_options : NULL,
@@ -699,7 +695,7 @@ struct aws_s3_meta_request *aws_s3_client_make_meta_request(
                 error_occurred = true;
                 goto unlock;
             }
-
+            endpoint->handled_by_client = true;
             endpoint_hash_element->value = endpoint;
             ++client->synced_data.num_endpoints_allocated;
         } else {
@@ -736,33 +732,6 @@ struct aws_s3_meta_request *aws_s3_client_make_meta_request(
     return meta_request;
 }
 
-static bool s_s3_client_endpoint_ref_count_zero(struct aws_s3_endpoint *endpoint) {
-    AWS_PRECONDITION(endpoint);
-
-    struct aws_s3_client *client = endpoint->user_data;
-    AWS_PRECONDITION(client);
-
-    bool clean_up_endpoint = false;
-
-    /* BEGIN CRITICAL SECTION */
-    {
-        aws_s3_client_lock_synced_data(client);
-
-        /* It is possible that before we were able to acquire the lock here, the critical section used for looking up
-         * the endpoint in the table and assigning it to a new meta request was called in a different thread. To handle
-         * this case, we check the ref count before removing it.*/
-        if (aws_atomic_load_int(&endpoint->ref_count.ref_count) == 0) {
-            aws_hash_table_remove(&client->synced_data.endpoints, endpoint->host_name, NULL, NULL);
-            clean_up_endpoint = true;
-        }
-
-        aws_s3_client_unlock_synced_data(client);
-    }
-    /* END CRITICAL SECTION */
-
-    return clean_up_endpoint;
-}
-
 static void s_s3_client_endpoint_shutdown_callback(void *user_data) {
     struct aws_s3_client *client = user_data;
     AWS_PRECONDITION(client);
@@ -1677,7 +1646,8 @@ void aws_s3_client_notify_connection_finished(
     aws_retry_token_release(connection->retry_token);
     connection->retry_token = NULL;
 
-    aws_s3_endpoint_release(connection->endpoint);
+    /* The endpoint must be created by client here */
+    aws_s3_client_endpoint_release(client, connection->endpoint);
     connection->endpoint = NULL;
 
     aws_mem_release(client->allocator, connection);

source/s3_endpoint.c

@@ -104,7 +104,6 @@ struct aws_s3_endpoint *aws_s3_endpoint_new(
         goto error_cleanup;
     }
 
-    endpoint->ref_count_zero_callback = options->ref_count_zero_callback;
     endpoint->shutdown_callback = options->shutdown_callback;
     endpoint->user_data = options->user_data;
 
@@ -213,14 +212,29 @@ void aws_s3_endpoint_release(struct aws_s3_endpoint *endpoint) {
     aws_ref_count_release(&endpoint->ref_count);
 }
 
+void aws_s3_client_endpoint_release(struct aws_s3_client *client, struct aws_s3_endpoint *endpoint) {
+    AWS_PRECONDITION(endpoint);
+    AWS_PRECONDITION(client);
+    AWS_PRECONDITION(endpoint->handled_by_client);
+
+    /* BEGIN CRITICAL SECTION */
+    {
+        aws_s3_client_lock_synced_data(client);
+        /* The last refcount to release */
+        if (aws_atomic_load_int(&endpoint->ref_count.ref_count) == 1) {
+            aws_hash_table_remove(&client->synced_data.endpoints, endpoint->host_name, NULL, NULL);
+        }
+        aws_s3_client_unlock_synced_data(client);
+    }
+    /* END CRITICAL SECTION */
+
+    aws_ref_count_release(&endpoint->ref_count);
+}
+
 static void s_s3_endpoint_ref_count_zero(void *user_data) {
     struct aws_s3_endpoint *endpoint = user_data;
     AWS_PRECONDITION(endpoint);
 
-    if (endpoint->ref_count_zero_callback != NULL && !endpoint->ref_count_zero_callback(endpoint)) {
-        return;
-    }
-
     if (endpoint->http_connection_manager != NULL) {
         struct aws_http_connection_manager *http_connection_manager = endpoint->http_connection_manager;
         endpoint->http_connection_manager = NULL;

source/s3_meta_request.c

@@ -370,6 +370,8 @@ static void s_s3_meta_request_destroy(void *user_data) {
 
     aws_cached_signing_config_destroy(meta_request->cached_signing_config);
     aws_mutex_clean_up(&meta_request->synced_data.lock);
+    /* For a meta request created by client, the endpoint should always be cleaned up by meta request finish call. Just
+     * use regular release here for those meta request not really went through the client */
     aws_s3_endpoint_release(meta_request->endpoint);
     aws_s3_client_release(meta_request->client);
 
@@ -1402,7 +1404,8 @@ void aws_s3_meta_request_finish_default(struct aws_s3_meta_request *meta_request
 
     aws_s3_meta_request_result_clean_up(meta_request, &finish_result);
 
-    aws_s3_endpoint_release(meta_request->endpoint);
+    /* The endpoint must be created by client here */
+    aws_s3_client_endpoint_release(meta_request->client, meta_request->endpoint);
     meta_request->endpoint = NULL;
 
     aws_s3_client_release(meta_request->client);

tests/CMakeLists.txt

@@ -127,10 +127,8 @@ add_net_test_case(test_s3_default_sending_meta_request_user_agent)
 add_net_test_case(test_s3_range_requests)
 add_net_test_case(test_s3_not_satisfiable_range)
 
-add_net_test_case(test_s3_endpoint_ref)
 add_net_test_case(test_s3_bad_endpoint)
 add_net_test_case(test_s3_different_endpoints)
-add_net_test_case(test_s3_endpoint_resurrect)
 
 add_test_case(test_s3_replace_quote_entities)
 add_test_case(test_s3_parse_content_range_response_header)

tests/s3_endpoint_tests.c

@@ -8,128 +8,6 @@
 #include <aws/io/channel_bootstrap.h>
 #include <aws/testing/aws_test_harness.h>
 
-static bool s_test_s3_endpoint_ref_zero_ref_callback(struct aws_s3_endpoint *endpoint) {
-    struct aws_s3_tester *tester = endpoint->user_data;
-
-    if (aws_s3_tester_inc_counter1(tester) == 1) {
-        return false;
-    }
-
-    return true;
-}
-
-static void s_test_s3_endpoint_ref_shutdown(void *user_data) {
-    struct aws_s3_tester *tester = user_data;
-
-    aws_s3_tester_inc_counter1(tester);
-}
-
-AWS_TEST_CASE(test_s3_endpoint_ref, s_test_s3_endpoint_ref)
-static int s_test_s3_endpoint_ref(struct aws_allocator *allocator, void *ctx) {
-    (void)ctx;
-
-    struct aws_s3_tester tester;
-    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester, false));
-    aws_s3_tester_set_counter1_desired(&tester, 3);
-
-    struct aws_string *host_name =
-        aws_s3_tester_build_endpoint_string(allocator, &g_test_public_bucket_name, &g_test_s3_region);
-
-    struct aws_s3_endpoint_options endpoint_options = {
-        .host_name = host_name,
-        .ref_count_zero_callback = s_test_s3_endpoint_ref_zero_ref_callback,
-        .shutdown_callback = s_test_s3_endpoint_ref_shutdown,
-        .client_bootstrap = tester.client_bootstrap,
-        .tls_connection_options = NULL,
-        .dns_host_address_ttl_seconds = 1,
-        .user_data = &tester,
-        .max_connections = 4,
-    };
-
-    struct aws_s3_endpoint *endpoint = aws_s3_endpoint_new(allocator, &endpoint_options);
-
-    ASSERT_TRUE(endpoint->http_connection_manager != NULL);
-
-    /* During the first release, s_test_s3_endpoint_ref_zero_ref_callback will return false, which will mean it does not
-     * get cleaned up. */
-    aws_s3_endpoint_release(endpoint);
-    ASSERT_TRUE(aws_atomic_load_int(&endpoint->ref_count.ref_count) == 0);
-
-    /* Perform an acquire, bumping the ref count to 1. */
-    aws_s3_endpoint_acquire(endpoint);
-    ASSERT_TRUE(aws_atomic_load_int(&endpoint->ref_count.ref_count) == 1);
-
-    /* During the secnod release, s_test_s3_endpoint_ref_zero_ref_callback will return true, causing clean up to
-     * occur.*/
-    aws_s3_endpoint_release(endpoint);
-
-    /* Wait for shutdown callback to increment the counter, hitting 3. */
-    aws_s3_tester_wait_for_counters(&tester);
-
-    aws_string_destroy(host_name);
-    aws_s3_tester_clean_up(&tester);
-
-    return 0;
-}
-
-static bool s_test_s3_endpoint_resurrect_endpoint_ref_count_zero(struct aws_s3_endpoint *endpoint) {
-    struct aws_s3_client *client = endpoint->user_data;
-    struct aws_s3_tester *tester = client->shutdown_callback_user_data;
-
-    bool acquire_and_release_endpoint = false;
-
-    if (aws_s3_tester_inc_counter1(tester) == 1) {
-        acquire_and_release_endpoint = true;
-        aws_s3_endpoint_acquire(endpoint);
-    }
-
-    struct aws_s3_client_vtable *original_client_vtable =
-        aws_s3_tester_get_client_vtable_patch(tester, 0)->original_vtable;
-
-    bool result = original_client_vtable->endpoint_ref_count_zero(endpoint);
-
-    if (acquire_and_release_endpoint) {
-        aws_s3_endpoint_release(endpoint);
-    }
-
-    return result;
-}
-
-AWS_TEST_CASE(test_s3_endpoint_resurrect, s_test_s3_endpoint_resurrect)
-static int s_test_s3_endpoint_resurrect(struct aws_allocator *allocator, void *ctx) {
-    (void)ctx;
-
-    struct aws_s3_tester tester;
-    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester, false));
-
-    struct aws_s3_tester_client_options client_options;
-    AWS_ZERO_STRUCT(client_options);
-
-    struct aws_s3_client *client = NULL;
-    ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
-
-    struct aws_s3_client_vtable *patched_client_vtable = aws_s3_tester_patch_client_vtable(&tester, client, NULL);
-    patched_client_vtable->endpoint_ref_count_zero = s_test_s3_endpoint_resurrect_endpoint_ref_count_zero;
-
-    struct aws_s3_tester_meta_request_options options = {
-        .allocator = allocator,
-        .client = client,
-        .meta_request_type = AWS_S3_META_REQUEST_TYPE_GET_OBJECT,
-        .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_SUCCESS,
-        .get_options =
-            {
-                .object_path = g_pre_existing_object_1MB,
-            },
-    };
-
-    ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &options, NULL));
-
-    aws_s3_client_release(client);
-    aws_s3_tester_clean_up(&tester);
-
-    return 0;
-}
-
 AWS_TEST_CASE(test_s3_different_endpoints, s_test_s3_different_endpoints)
 static int s_test_s3_different_endpoints(struct aws_allocator *allocator, void *ctx) {
     (void)ctx;

tests/s3_tester.c

@@ -595,12 +595,6 @@ static void s_s3_client_process_work_empty(struct aws_s3_client *client) {
     (void)client;
 }
 
-static bool s_s3_client_endpoint_ref_count_zero_empty(struct aws_s3_endpoint *endpoint) {
-    AWS_PRECONDITION(endpoint);
-    (void)endpoint;
-    return true;
-}
-
 static void s_s3_client_endpoint_shutdown_callback_empty(void *user_data) {
     AWS_PRECONDITION(user_data);
     (void)user_data;
@@ -618,7 +612,6 @@ struct aws_s3_client_vtable g_aws_s3_client_mock_vtable = {
     .get_host_address_count = s_s3_client_get_host_address_count_empty,
     .schedule_process_work_synced = s_s3_client_schedule_process_work_synced_empty,
     .process_work = s_s3_client_process_work_empty,
-    .endpoint_ref_count_zero = s_s3_client_endpoint_ref_count_zero_empty,
     .endpoint_shutdown_callback = s_s3_client_endpoint_shutdown_callback_empty,
     .finish_destroy = s_s3_client_finish_destroy_empty,
 };