Small buffer (#422)

Co-authored-by: Michael Graeb <graebm@amazon.com>
awslabs · Apr 2, 2024 · 0ce756e · 0ce756e
1 parent 3085fd7
commit 0ce756e
Show file tree

Hide file tree

Showing 9 changed files with 83 additions and 13 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -32,7 +32,7 @@ jobs:
           - al2-x64
     steps:
     - name: Checkout Sources
-      uses: actions/checkout@v3
+      uses: actions/checkout@v4
         # We can't use the `uses: docker://image` version yet, GitHub lacks authentication for actions -> packages
     - name: Build ${{ env.PACKAGE_NAME }}
       run: |
@@ -137,7 +137,7 @@ jobs:
     runs-on: macos-13 # latest
     steps:
     - name: Checkout Sources
-      uses: actions/checkout@v3
+      uses: actions/checkout@v4
     - name: Build ${{ env.PACKAGE_NAME }} + consumers
       run: |
         python3 -c "from urllib.request import urlretrieve; urlretrieve('${{ env.BUILDER_HOST }}/${{ env.BUILDER_SOURCE }}/${{ env.BUILDER_VERSION }}/builder.pyz?run=${{ env.RUN }}', 'builder')"

diff --git a/.github/workflows/clang-format.yml b/.github/workflows/clang-format.yml
@@ -9,7 +9,7 @@ jobs:
 
     steps:
     - name: Checkout Sources
-      uses: actions/checkout@v1
+      uses: actions/checkout@v4
 
     - name: clang-format lint
       uses: DoozyX/clang-format-lint-action@v0.3.1

diff --git a/.github/workflows/codecov.yml b/.github/workflows/codecov.yml
@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - name: Checkout Sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Build ${{ env.PACKAGE_NAME }} + consumers
         run: |
           python3 -c "from urllib.request import urlretrieve; urlretrieve('${{ env.BUILDER_HOST }}/${{ env.BUILDER_SOURCE }}/${{ env.BUILDER_VERSION }}/builder.pyz?run=${{ env.RUN }}', 'builder')"

diff --git a/source/s3_auto_ranged_get.c b/source/s3_auto_ranged_get.c
@@ -11,6 +11,11 @@
 #include <aws/common/string.h>
 #include <inttypes.h>
 
+/* Dont use buffer pool when we know response size, and its below this number,
+ * i.e. when user provides explicit range that is small, ex. range = 1-100.
+ * Instead of going through the pool in that case, we just use a dynamic buffer
+ * for response (pre-mempool behavior). */
+const uint64_t s_min_size_response_for_pooling = 1 * 1024 * 1024;
 const uint32_t s_conservative_max_requests_in_flight = 8;
 const struct aws_byte_cursor g_application_xml_value = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("application/xml");
 
@@ -287,10 +292,20 @@ static bool s_s3_auto_ranged_get_update(
                             AWS_LS_S3_META_REQUEST,
                             "id=%p: Doing a ranged get to discover the size of the object and get the first part",
                             (void *)meta_request);
-                        ticket = aws_s3_buffer_pool_reserve(meta_request->client->buffer_pool, (size_t)first_part_size);
 
-                        if (ticket == NULL) {
-                            goto has_work_remaining;
+                        if (first_part_size >= s_min_size_response_for_pooling) {
+                            /* Note: explicitly reserving the whole part size
+                             * even if expect to receive less data. Pool will
+                             * reserve the whole part size for it anyways, so no
+                             * reason getting a smaller chunk. */
+                            ticket = aws_s3_buffer_pool_reserve(
+                                meta_request->client->buffer_pool, (size_t)meta_request->part_size);
+
+                            if (ticket == NULL) {
+                                goto has_work_remaining;
+                            }
+                        } else {
+                            ticket = NULL;
                         }
 
                         request = aws_s3_request_new(

diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c
@@ -1342,8 +1342,8 @@ static int s_s3_meta_request_headers_block_done(
  * Small helper to either do a static or dynamic append.
  * TODO: something like this would be useful in common.
  */
-static int s_response_body_append(bool is_dynamic, struct aws_byte_buf *buf, const struct aws_byte_cursor *data) {
-    return is_dynamic ? aws_byte_buf_append_dynamic(buf, data) : aws_byte_buf_append(buf, data);
+static int s_response_body_append(struct aws_byte_buf *buf, const struct aws_byte_cursor *data) {
+    return buf->allocator != NULL ? aws_byte_buf_append_dynamic(buf, data) : aws_byte_buf_append(buf, data);
 }
 
 static int s_s3_meta_request_incoming_body(
@@ -1381,8 +1381,7 @@ static int s_s3_meta_request_incoming_body(
     }
 
     if (request->send_data.response_body.capacity == 0) {
-        if (request->has_part_size_response_body && successful_response) {
-            AWS_FATAL_ASSERT(request->ticket);
+        if (request->has_part_size_response_body && request->ticket != NULL) {
             request->send_data.response_body =
                 aws_s3_buffer_pool_acquire_buffer(request->meta_request->client->buffer_pool, request->ticket);
         } else {
@@ -1393,7 +1392,7 @@ static int s_s3_meta_request_incoming_body(
 
     /* Note: not having part sized response body means the buffer is dynamic and
      * can grow. */
-    if (s_response_body_append(!request->has_part_size_response_body, &request->send_data.response_body, data)) {
+    if (s_response_body_append(&request->send_data.response_body, data)) {
         AWS_LOGF_ERROR(
             AWS_LS_S3_META_REQUEST,
             "id=%p: Request %p could not append to response body due to error %d (%s)",

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -282,6 +282,7 @@ if(ENABLE_MOCK_SERVER_TESTS)
     add_net_test_case(get_object_invalid_responses_mock_server)
     add_net_test_case(get_object_mismatch_checksum_responses_mock_server)
     add_net_test_case(get_object_throughput_failure_mock_server)
+    add_net_test_case(get_object_long_error_mock_server)
     add_net_test_case(upload_part_invalid_response_mock_server)
     add_net_test_case(upload_part_async_invalid_response_mock_server)
     add_net_test_case(resume_first_part_not_completed_mock_server)

diff --git a/tests/mock_s3_server/GetObject/get_object_long_error.json b/tests/mock_s3_server/GetObject/get_object_long_error.json
@@ -0,0 +1,9 @@
+{
+  "status": 403,
+  "headers": {
+    "Content-Type": "application/xml"
+  },
+  "body": [
+    "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<Error><Code>NoSuchKey</Code><Message>really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error really long error </Message><Resource>/mybucket/myfoto.jpg</Resource><RequestId>4442587FB7D0A2F9</RequestId></Error>"
+  ]
+}
diff --git a/tests/mock_s3_server/mock_s3_server.py b/tests/mock_s3_server/mock_s3_server.py
@@ -424,7 +424,9 @@ def handle_get_object(wrapper, request, parsed_path, head_request=False):
     else:
         RETRY_REQUEST_COUNT = 0
 
-    if parsed_path.path == "/get_object_invalid_response_missing_content_range" or parsed_path.path == "/get_object_invalid_response_missing_etags":
+    if (parsed_path.path == "/get_object_invalid_response_missing_content_range" or 
+        parsed_path.path == "/get_object_invalid_response_missing_etags" or 
+        parsed_path.path == "/get_object_long_error"):
         # Don't generate the body for those requests
         return response_config
 

diff --git a/tests/s3_mock_server_tests.c b/tests/s3_mock_server_tests.c
@@ -545,6 +545,50 @@ TEST_CASE(get_object_throughput_failure_mock_server) {
     return AWS_OP_SUCCESS;
 }
 
+TEST_CASE(get_object_long_error_mock_server) {
+    (void)ctx;
+
+    struct aws_s3_tester tester;
+    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester));
+    struct aws_s3_tester_client_options client_options = {
+        .part_size = 64 * 1024,
+        .tls_usage = AWS_S3_TLS_DISABLED,
+    };
+
+    struct aws_s3_client *client = NULL;
+    ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
+
+    struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/get_object_long_error");
+
+    struct aws_s3_tester_meta_request_options get_options = {
+        .allocator = allocator,
+        .meta_request_type = AWS_S3_META_REQUEST_TYPE_GET_OBJECT,
+        .client = client,
+        .get_options =
+            {
+                .object_path = object_path,
+            },
+        .default_type_options =
+            {
+                .mode = AWS_S3_TESTER_DEFAULT_TYPE_MODE_GET,
+            },
+        .mock_server = true,
+        .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_FAILURE,
+    };
+    struct aws_s3_meta_request_test_results out_results;
+    aws_s3_meta_request_test_results_init(&out_results, allocator);
+
+    ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &get_options, &out_results));
+
+    ASSERT_UINT_EQUALS(AWS_ERROR_S3_INVALID_RESPONSE_STATUS, out_results.finished_error_code);
+
+    aws_s3_meta_request_test_results_clean_up(&out_results);
+    aws_s3_client_release(client);
+    aws_s3_tester_clean_up(&tester);
+
+    return AWS_OP_SUCCESS;
+}
+
 static int s_test_upload_part_invalid_response_mock_server_ex(
     struct aws_allocator *allocator,
     bool async_input_stream) {