1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
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);
}
}
|
