intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 205 insertions, 0 deletions

diff --git a/contrib/restricted/aws/aws-c-io/source/message_pool.c b/contrib/restricted/aws/aws-c-io/source/message_pool.c
new file mode 100644
index 00000000000..3090f57e6bf
--- /dev/null
+++ b/contrib/restricted/aws/aws-c-io/source/message_pool.c
@@ -0,0 +1,205 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+#include <aws/io/message_pool.h>
+
+#include <aws/common/thread.h>
+
+int aws_memory_pool_init(
+    struct aws_memory_pool *mempool,
+    struct aws_allocator *alloc,
+    uint16_t ideal_segment_count,
+    size_t segment_size) {
+
+    mempool->alloc = alloc;
+    mempool->ideal_segment_count = ideal_segment_count;
+    mempool->segment_size = segment_size;
+    mempool->data_ptr = aws_mem_calloc(alloc, ideal_segment_count, sizeof(void *));
+    if (!mempool->data_ptr) {
+        return AWS_OP_ERR;
+    }
+
+    aws_array_list_init_static(&mempool->stack, mempool->data_ptr, ideal_segment_count, sizeof(void *));
+
+    for (uint16_t i = 0; i < ideal_segment_count; ++i) {
+        void *memory = aws_mem_acquire(alloc, segment_size);
+        if (memory) {
+            aws_array_list_push_back(&mempool->stack, &memory);
+        } else {
+            goto clean_up;
+        }
+    }
+
+    return AWS_OP_SUCCESS;
+
+clean_up:
+    aws_memory_pool_clean_up(mempool);
+    return AWS_OP_ERR;
+}
+
+void aws_memory_pool_clean_up(struct aws_memory_pool *mempool) {
+    void *cur = NULL;
+
+    while (aws_array_list_length(&mempool->stack) > 0) {
+        /* the only way this fails is not possible since I already checked the length. */
+        aws_array_list_back(&mempool->stack, &cur);
+        aws_array_list_pop_back(&mempool->stack);
+        aws_mem_release(mempool->alloc, cur);
+    }
+
+    aws_array_list_clean_up(&mempool->stack);
+    aws_mem_release(mempool->alloc, mempool->data_ptr);
+}
+
+void *aws_memory_pool_acquire(struct aws_memory_pool *mempool) {
+    void *back = NULL;
+    if (aws_array_list_length(&mempool->stack) > 0) {
+        aws_array_list_back(&mempool->stack, &back);
+        aws_array_list_pop_back(&mempool->stack);
+
+        return back;
+    }
+
+    void *mem = aws_mem_acquire(mempool->alloc, mempool->segment_size);
+    return mem;
+}
+
+void aws_memory_pool_release(struct aws_memory_pool *mempool, void *to_release) {
+    size_t pool_size = aws_array_list_length(&mempool->stack);
+
+    if (pool_size >= mempool->ideal_segment_count) {
+        aws_mem_release(mempool->alloc, to_release);
+        return;
+    }
+
+    aws_array_list_push_back(&mempool->stack, &to_release);
+}
+
+struct message_pool_allocator {
+    struct aws_allocator base_allocator;
+    struct aws_message_pool *msg_pool;
+};
+
+void *s_message_pool_mem_acquire(struct aws_allocator *allocator, size_t size) {
+    (void)allocator;
+    (void)size;
+
+    /* no one should ever call this ever. */
+    AWS_ASSERT(0);
+    return NULL;
+}
+
+void s_message_pool_mem_release(struct aws_allocator *allocator, void *ptr) {
+    struct message_pool_allocator *msg_pool_alloc = allocator->impl;
+
+    aws_message_pool_release(msg_pool_alloc->msg_pool, (struct aws_io_message *)ptr);
+}
+
+static size_t MSG_OVERHEAD = sizeof(struct aws_io_message) + sizeof(struct message_pool_allocator);
+
+int aws_message_pool_init(
+    struct aws_message_pool *msg_pool,
+    struct aws_allocator *alloc,
+    struct aws_message_pool_creation_args *args) {
+
+    msg_pool->alloc = alloc;
+
+    size_t msg_data_size = args->application_data_msg_data_size + MSG_OVERHEAD;
+
+    if (aws_memory_pool_init(
+            &msg_pool->application_data_pool, alloc, args->application_data_msg_count, msg_data_size)) {
+        return AWS_OP_ERR;
+    }
+
+    size_t small_blk_data_size = args->small_block_msg_data_size + MSG_OVERHEAD;
+
+    if (aws_memory_pool_init(&msg_pool->small_block_pool, alloc, args->small_block_msg_count, small_blk_data_size)) {
+        aws_memory_pool_clean_up(&msg_pool->application_data_pool);
+        return AWS_OP_ERR;
+    }
+
+    return AWS_OP_SUCCESS;
+}
+
+void aws_message_pool_clean_up(struct aws_message_pool *msg_pool) {
+    aws_memory_pool_clean_up(&msg_pool->application_data_pool);
+    aws_memory_pool_clean_up(&msg_pool->small_block_pool);
+    AWS_ZERO_STRUCT(*msg_pool);
+}
+
+struct message_wrapper {
+    struct aws_io_message message;
+    struct message_pool_allocator msg_allocator;
+    uint8_t buffer_start[1];
+};
+
+struct aws_io_message *aws_message_pool_acquire(
+    struct aws_message_pool *msg_pool,
+    enum aws_io_message_type message_type,
+    size_t size_hint) {
+
+    struct message_wrapper *message_wrapper = NULL;
+    size_t max_size = 0;
+    switch (message_type) {
+        case AWS_IO_MESSAGE_APPLICATION_DATA:
+            if (size_hint > msg_pool->small_block_pool.segment_size - MSG_OVERHEAD) {
+                message_wrapper = aws_memory_pool_acquire(&msg_pool->application_data_pool);
+                max_size = msg_pool->application_data_pool.segment_size - MSG_OVERHEAD;
+            } else {
+                message_wrapper = aws_memory_pool_acquire(&msg_pool->small_block_pool);
+                max_size = msg_pool->small_block_pool.segment_size - MSG_OVERHEAD;
+            }
+            break;
+        default:
+            AWS_ASSERT(0);
+            aws_raise_error(AWS_IO_CHANNEL_UNKNOWN_MESSAGE_TYPE);
+            return NULL;
+    }
+
+    if (!message_wrapper) {
+        return NULL;
+    }
+
+    message_wrapper->message.message_type = message_type;
+    message_wrapper->message.message_tag = 0;
+    message_wrapper->message.user_data = NULL;
+    message_wrapper->message.copy_mark = 0;
+    message_wrapper->message.on_completion = NULL;
+    /* the buffer shares the allocation with the message. It's the bit at the end. */
+    message_wrapper->message.message_data.buffer = message_wrapper->buffer_start;
+    message_wrapper->message.message_data.len = 0;
+    message_wrapper->message.message_data.capacity = size_hint <= max_size ? size_hint : max_size;
+
+    /* set the allocator ptr */
+    message_wrapper->msg_allocator.base_allocator.impl = &message_wrapper->msg_allocator;
+    message_wrapper->msg_allocator.base_allocator.mem_acquire = s_message_pool_mem_acquire;
+    message_wrapper->msg_allocator.base_allocator.mem_realloc = NULL;
+    message_wrapper->msg_allocator.base_allocator.mem_release = s_message_pool_mem_release;
+    message_wrapper->msg_allocator.msg_pool = msg_pool;
+
+    message_wrapper->message.allocator = &message_wrapper->msg_allocator.base_allocator;
+    return &message_wrapper->message;
+}
+
+void aws_message_pool_release(struct aws_message_pool *msg_pool, struct aws_io_message *message) {
+
+    memset(message->message_data.buffer, 0, message->message_data.len);
+    message->allocator = NULL;
+
+    struct message_wrapper *wrapper = AWS_CONTAINER_OF(message, struct message_wrapper, message);
+
+    switch (message->message_type) {
+        case AWS_IO_MESSAGE_APPLICATION_DATA:
+            if (message->message_data.capacity > msg_pool->small_block_pool.segment_size - MSG_OVERHEAD) {
+                aws_memory_pool_release(&msg_pool->application_data_pool, wrapper);
+            } else {
+                aws_memory_pool_release(&msg_pool->small_block_pool, wrapper);
+            }
+            break;
+        default:
+            AWS_ASSERT(0);
+            aws_raise_error(AWS_IO_CHANNEL_UNKNOWN_MESSAGE_TYPE);
+    }
+}