aws_explicit_aligned_allocator_new (#1147)

Author: TingDaoK

include/aws/common/allocator.h

@@ -228,6 +228,19 @@ size_t aws_small_block_allocator_page_size(struct aws_allocator *sba_allocator);
 AWS_COMMON_API
 size_t aws_small_block_allocator_page_size_available(struct aws_allocator *sba_allocator);
 
+/*
+ * Create an aligned allocator with an explicit alignment. Always align the allocated buffer with the passed-in
+ * alignment value.
+ */
+AWS_COMMON_API
+struct aws_allocator *aws_explicit_aligned_allocator_new(size_t alignment);
+
+/*
+ * Destroys a customized aligned allocator instance and frees its memory.
+ */
+AWS_COMMON_API
+void aws_explicit_aligned_allocator_destroy(struct aws_allocator *aligned_alloc);
+
 AWS_EXTERN_C_END
 AWS_POP_SANE_WARNING_LEVEL
 

source/allocator.c

@@ -35,20 +35,28 @@ bool aws_allocator_is_valid(const struct aws_allocator *alloc) {
 }
 
 static void *s_aligned_malloc(struct aws_allocator *allocator, size_t size) {
-    (void)allocator;
-    /* larger allocations should be aligned so that AVX and friends can avoid
-     * the extra preamble during unaligned versions of memcpy/memset on big buffers
-     * This will also accelerate hardware CRC and SHA on ARM chips
-     *
-     * 64 byte alignment for > page allocations on 64 bit systems
-     * 32 byte alignment for > page allocations on 32 bit systems
-     * 16 byte alignment for <= page allocations on 64 bit systems
-     * 8 byte alignment for <= page allocations on 32 bit systems
-     *
-     * We use PAGE_SIZE as the boundary because we are not aware of any allocations of
-     * this size or greater that are not data buffers
-     */
-    const size_t alignment = sizeof(void *) * (size > (size_t)PAGE_SIZE ? 8 : 2);
+    size_t alignment = 0;
+    if (allocator->impl != NULL) {
+        alignment = (size_t)allocator->impl;
+    } else {
+        /**
+         * For implicit alignment.
+         * larger allocations should be aligned so that AVX and friends can avoid
+         * the extra preamble during unaligned versions of memcpy/memset on big buffers
+         * This will also accelerate hardware CRC and SHA on ARM chips
+         *
+         * 64 byte alignment for > page allocations on 64 bit systems
+         * 32 byte alignment for > page allocations on 32 bit systems
+         * 16 byte alignment for <= page allocations on 64 bit systems
+         * 8 byte alignment for <= page allocations on 32 bit systems
+         *
+         * We use PAGE_SIZE as the boundary because we are not aware of any allocations of
+         * this size or greater that are not data buffers.
+         *
+         * Unless there is a customized alignment size.
+         */
+        alignment = sizeof(void *) * (size > (size_t)PAGE_SIZE ? 8 : 2);
+    }
 #if !defined(_WIN32)
     void *result = NULL;
     int err = posix_memalign(&result, alignment, size);
@@ -146,26 +154,48 @@ static void *s_non_aligned_calloc(struct aws_allocator *allocator, size_t num, s
     return mem;
 }
 
-static struct aws_allocator default_allocator = {
+static struct aws_allocator s_default_allocator = {
     .mem_acquire = s_non_aligned_malloc,
     .mem_release = s_non_aligned_free,
     .mem_realloc = s_non_aligned_realloc,
     .mem_calloc = s_non_aligned_calloc,
 };
 
 struct aws_allocator *aws_default_allocator(void) {
-    return &default_allocator;
+    return &s_default_allocator;
 }
 
-static struct aws_allocator aligned_allocator = {
+static struct aws_allocator s_implicit_aligned_allocator = {
     .mem_acquire = s_aligned_malloc,
     .mem_release = s_aligned_free,
     .mem_realloc = s_aligned_realloc,
     .mem_calloc = s_aligned_calloc,
 };
 
 struct aws_allocator *aws_aligned_allocator(void) {
-    return &aligned_allocator;
+    return &s_implicit_aligned_allocator;
+}
+
+struct aws_allocator *aws_explicit_aligned_allocator_new(size_t customized_alignment) {
+    if (customized_alignment == 0 || (customized_alignment & (customized_alignment - 1)) != 0 ||
+        customized_alignment % sizeof(void *) != 0) {
+        /**
+         * the alignment must be a power of two and a multiple of sizeof(void *) and non-zero.
+         */
+        aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+        return NULL;
+    }
+    struct aws_allocator *aligned_alloc = aws_mem_calloc(aws_default_allocator(), 1, sizeof(struct aws_allocator));
+    *aligned_alloc = s_implicit_aligned_allocator;
+    aligned_alloc->impl = (void *)customized_alignment;
+    return aligned_alloc;
+}
+
+void aws_explicit_aligned_allocator_destroy(struct aws_allocator *aligned_alloc) {
+    if (!aligned_alloc) {
+        return;
+    }
+    aws_mem_release(aws_default_allocator(), aligned_alloc);
 }
 
 void *aws_mem_acquire(struct aws_allocator *allocator, size_t size) {

tests/CMakeLists.txt

@@ -327,6 +327,9 @@ add_test_case(default_threaded_reallocs)
 add_test_case(default_threaded_allocs_and_frees)
 add_test_case(aligned_threaded_reallocs)
 add_test_case(aligned_threaded_allocs_and_frees)
+add_test_case(explicit_aligned_sanitize)
+add_test_case(explicit_aligned_threaded_reallocs)
+add_test_case(explicit_aligned_threaded_allocs_and_frees)
 
 add_test_case(test_memtrace_none)
 add_test_case(test_memtrace_count)

tests/alloc_test.c

@@ -350,7 +350,7 @@ static int s_default_threaded_allocs_and_frees(struct aws_allocator *allocator,
 AWS_TEST_CASE(default_threaded_allocs_and_frees, s_default_threaded_allocs_and_frees)
 
 /*
- * No align allocator tests.
+ * aligned allocator tests.
  */
 static int s_aligned_threaded_reallocs(struct aws_allocator *allocator, void *ctx) {
     (void)allocator;
@@ -381,3 +381,66 @@ static int s_aligned_threaded_allocs_and_frees(struct aws_allocator *allocator,
     return 0;
 }
 AWS_TEST_CASE(aligned_threaded_allocs_and_frees, s_aligned_threaded_allocs_and_frees)
+
+static int s_explicit_aligned_sanitize(struct aws_allocator *allocator, void *ctx) {
+    (void)allocator;
+    (void)ctx;
+
+    struct aws_allocator *aligned_alloc = aws_explicit_aligned_allocator_new(1);
+    ASSERT_NULL(aligned_alloc);
+    ASSERT_UINT_EQUALS(aws_last_error(), AWS_ERROR_INVALID_ARGUMENT);
+
+    aligned_alloc = aws_explicit_aligned_allocator_new(3 * sizeof(void *));
+    ASSERT_NULL(aligned_alloc);
+    ASSERT_UINT_EQUALS(aws_last_error(), AWS_ERROR_INVALID_ARGUMENT);
+
+    aligned_alloc = aws_explicit_aligned_allocator_new(0);
+    ASSERT_NULL(aligned_alloc);
+    ASSERT_UINT_EQUALS(aws_last_error(), AWS_ERROR_INVALID_ARGUMENT);
+
+    size_t aligned_size = 1024;
+    aligned_alloc = aws_explicit_aligned_allocator_new(aligned_size);
+    ASSERT_NOT_NULL(aligned_alloc);
+    void *test = aws_mem_acquire(aligned_alloc, sizeof(void *));
+    ASSERT_TRUE((uintptr_t)test % aligned_size == 0);
+    aws_mem_release(aligned_alloc, test);
+
+    aws_explicit_aligned_allocator_destroy(aligned_alloc);
+
+    return 0;
+}
+AWS_TEST_CASE(explicit_aligned_sanitize, s_explicit_aligned_sanitize)
+
+static int s_explicit_aligned_threaded_reallocs(struct aws_allocator *allocator, void *ctx) {
+    (void)allocator;
+    (void)ctx;
+    srand(15);
+    struct aws_allocator *aligned_alloc = aws_explicit_aligned_allocator_new(512);
+
+    struct aws_allocator *alloc = aws_mem_tracer_new(aligned_alloc, NULL, AWS_MEMTRACE_STACKS, 8);
+
+    s_thread_test(alloc, s_threaded_realloc_worker, alloc);
+
+    aws_mem_tracer_destroy(alloc);
+    aws_explicit_aligned_allocator_destroy(aligned_alloc);
+
+    return 0;
+}
+AWS_TEST_CASE(explicit_aligned_threaded_reallocs, s_explicit_aligned_threaded_reallocs)
+
+static int s_explicit_aligned_threaded_allocs_and_frees(struct aws_allocator *allocator, void *ctx) {
+    (void)allocator;
+    (void)ctx;
+    srand(99);
+    struct aws_allocator *aligned_alloc = aws_explicit_aligned_allocator_new(512);
+
+    struct aws_allocator *alloc = aws_mem_tracer_new(aligned_alloc, NULL, AWS_MEMTRACE_STACKS, 8);
+
+    s_thread_test(alloc, s_threaded_alloc_worker, alloc);
+
+    aws_mem_tracer_destroy(alloc);
+    aws_explicit_aligned_allocator_destroy(aligned_alloc);
+
+    return 0;
+}
+AWS_TEST_CASE(explicit_aligned_threaded_allocs_and_frees, s_explicit_aligned_threaded_allocs_and_frees)