track use of correlation tokens from a RR operation submission through decref/destruction

Author: sbSteveK

include/aws/mqtt/mqtt.h

@@ -93,7 +93,6 @@ enum aws_mqtt_error {
     AWS_ERROR_MQTT_REQUEST_RESPONSE_MODELED_SERVICE_ERROR,
     AWS_ERROR_MQTT_REQUEST_RESPONSE_PAYLOAD_PARSE_ERROR,
     AWS_ERROR_MQTT_REQUEST_RESPONSE_INVALID_RESPONSE_PATH,
-    AWS_ERROR_MQTT_REQUEST_RESPONSE_DUPLICATE_CORRELATION_TOKEN,
 
     AWS_ERROR_END_MQTT_RANGE = AWS_ERROR_ENUM_END_RANGE(AWS_C_MQTT_PACKAGE_ID),
 };

source/mqtt.c

@@ -269,9 +269,6 @@ bool aws_mqtt_is_valid_topic_filter(const struct aws_byte_cursor *topic_filter)
             AWS_DEFINE_ERROR_INFO_MQTT(
                  AWS_ERROR_MQTT_REQUEST_RESPONSE_INVALID_RESPONSE_PATH,
                  "Request-response operation failed due to arrival on an unknown response path"),
-            AWS_DEFINE_ERROR_INFO_MQTT(
-                 AWS_ERROR_MQTT_REQUEST_RESPONSE_DUPLICATE_CORRELATION_TOKEN,
-                 "Request-response operation failed due to a duplicate correlation token."),
         };
 /* clang-format on */
 #undef AWS_DEFINE_ERROR_INFO_MQTT

source/request-response/request_response_client.c

@@ -273,6 +273,9 @@ struct aws_mqtt_request_response_client {
      * Map from cursor (correlation token) -> request operation
      */
     struct aws_hash_table operations_by_correlation_tokens;
+
+    /* Track correlation tokens currently in use by operations in queued or active states */
+    struct aws_hash_table correlation_tokens_in_use;
 };
 
 struct aws_mqtt_request_response_client *aws_mqtt_request_response_client_acquire_internal(
@@ -318,6 +321,7 @@ static void s_mqtt_request_response_client_final_destroy(struct aws_mqtt_request
 
     aws_mqtt_request_response_client_subscriptions_clean_up(&client->subscriptions);
     aws_hash_table_clean_up(&client->operations_by_correlation_tokens);
+    aws_hash_table_clean_up(&client->correlation_tokens_in_use);
 
     aws_event_loop_group_release_from_event_loop(client->loop);
     aws_mem_release(client->allocator, client);
@@ -1102,6 +1106,15 @@ static struct aws_mqtt_request_response_client *s_aws_mqtt_request_response_clie
         NULL,
         NULL);
 
+    aws_hash_table_init(
+        &rr_client->correlation_tokens_in_use,
+        allocator,
+        MQTT_RR_CLIENT_OPERATION_TABLE_DEFAULT_SIZE,
+        aws_hash_byte_cursor_ptr,
+        aws_mqtt_byte_cursor_hash_equality,
+        NULL,
+        NULL);
+
     aws_linked_list_init(&rr_client->operation_queue);
 
     aws_task_init(
@@ -1234,25 +1247,17 @@ static int s_add_request_operation_to_response_path_table(
 }
 
 static bool s_is_correlation_token_in_use(
-    struct aws_mqtt_request_response_client *client,
-    struct aws_mqtt_rr_client_operation *operation) {
+    const struct aws_mqtt_request_response_client *client,
+    const struct aws_byte_cursor *correlation_token) {
     struct aws_hash_element *elem = NULL;
-    aws_hash_table_find(
-        &client->operations_by_correlation_tokens,
-        &operation->storage.request_storage.options.correlation_token,
-        &elem);
+    aws_hash_table_find(&client->correlation_tokens_in_use, correlation_token, &elem);
     return elem != NULL;
 }
 
 static int s_add_request_operation_to_correlation_token_table(
     struct aws_mqtt_request_response_client *client,
     struct aws_mqtt_rr_client_operation *operation) {
 
-    if (s_is_correlation_token_in_use(client, operation)) {
-        aws_raise_error(AWS_ERROR_MQTT_REQUEST_RESPONSE_DUPLICATE_CORRELATION_TOKEN);
-        return AWS_OP_ERR;
-    }
-
     return aws_hash_table_put(
         &client->operations_by_correlation_tokens,
         &operation->storage.request_storage.options.correlation_token,
@@ -1632,6 +1637,17 @@ static bool s_are_request_operation_options_valid(
         return false;
     }
 
+    if (request_options->correlation_token.len > 0) {
+        // do a correlation token check
+        if (s_is_correlation_token_in_use(client, &request_options->correlation_token)) {
+            AWS_LOGF_ERROR(
+                AWS_LS_MQTT_REQUEST_RESPONSE,
+                "(%p) rr client request options - correlation token is already in use.",
+                (void *)client);
+            return false;
+        }
+    }
+
     return true;
 }
 
@@ -1712,6 +1728,15 @@ static void s_mqtt_rr_client_submit_operation(struct aws_task *task, void *arg,
     aws_mqtt_rr_client_operation_release(operation);
 }
 
+static void s_remove_correlation_token_from_in_use(struct aws_mqtt_rr_client_operation *operation) {
+    if (operation == NULL || operation->type != AWS_MRROT_REQUEST) {
+        return;
+    }
+    struct aws_mqtt_request_response_client *client = operation->client_internal_ref;
+    aws_hash_table_remove(
+        &client->correlation_tokens_in_use, &operation->storage.request_storage.options.correlation_token, NULL, NULL);
+}
+
 static void s_aws_mqtt_streaming_operation_storage_clean_up(struct aws_mqtt_streaming_operation_storage *storage) {
     aws_byte_buf_clean_up(&storage->operation_data);
 }
@@ -1993,6 +2018,11 @@ int aws_mqtt_request_response_client_submit_request(
         return aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
     }
 
+    // Store the correlation token to prevent allowing a duplicate correlation token to be used.
+    if (request_options->correlation_token.len > 0) {
+        aws_hash_table_put(&client->correlation_tokens_in_use, &request_options->correlation_token, NULL, NULL);
+    }
+
     uint64_t now = 0;
     if (aws_high_res_clock_get_ticks(&now)) {
         return aws_raise_error(AWS_ERROR_CLOCK_FAILURE);
@@ -2144,6 +2174,7 @@ struct aws_mqtt_rr_client_operation *aws_mqtt_rr_client_operation_acquire(
 struct aws_mqtt_rr_client_operation *aws_mqtt_rr_client_operation_release(
     struct aws_mqtt_rr_client_operation *operation) {
     if (operation != NULL) {
+        s_remove_correlation_token_from_in_use(operation);
         aws_ref_count_release(&operation->ref_count);
     }
 

tests/request-response/request_response_client_tests.c

@@ -3337,20 +3337,17 @@ static int s_rrc_request_response_failure_duplicate_correlation_token_fn(struct
         false));
 
     struct aws_byte_cursor record_key_2 = aws_byte_cursor_from_c_str("testkey2");
-    ASSERT_SUCCESS(s_rrc_test_submit_test_request(
+    s_rrc_test_submit_test_request(
         &fixture,
         RRC_PHDT_FAILURE_DUPLICATE_CORRELATION_TOKEN,
-        "test",
+        "test2",
         record_key_2,
-        "test/accepted",
+        "test2/accepted",
         "token1",
         NULL,
-        false));
-
-    s_rrc_wait_on_request_completion(&fixture, record_key_2);
+        false);
 
-    ASSERT_SUCCESS(s_rrc_verify_request_completion(
-        &fixture, record_key_2, AWS_ERROR_MQTT_REQUEST_RESPONSE_DUPLICATE_CORRELATION_TOKEN, NULL, NULL));
+    ASSERT_INT_EQUALS(aws_last_error(), AWS_ERROR_INVALID_ARGUMENT);
 
     s_aws_rr_client_test_fixture_clean_up(&fixture, false);