1/*
2 * simple memory allocator, backed by mmap() so that it hands out memory
3 * that can be shared across processes and threads
4 */
5#include <sys/mman.h>
6#include <stdio.h>
7#include <stdlib.h>
8#include <assert.h>
9#include <string.h>
10#include <unistd.h>
11#include <inttypes.h>
12#include <sys/types.h>
13#include <limits.h>
14#include <fcntl.h>
15
16#include "fio.h"
17#include "mutex.h"
18#include "arch/arch.h"
19#include "os/os.h"
20#include "smalloc.h"
21#include "log.h"
22
23#define SMALLOC_REDZONE		/* define to detect memory corruption */
24
25#define SMALLOC_BPB	32	/* block size, bytes-per-bit in bitmap */
26#define SMALLOC_BPI	(sizeof(unsigned int) * 8)
27#define SMALLOC_BPL	(SMALLOC_BPB * SMALLOC_BPI)
28
29#define INITIAL_SIZE	16*1024*1024	/* new pool size */
30#define INITIAL_POOLS	8		/* maximum number of pools to setup */
31
32#define MAX_POOLS	16
33
34#define SMALLOC_PRE_RED		0xdeadbeefU
35#define SMALLOC_POST_RED	0x5aa55aa5U
36
37unsigned int smalloc_pool_size = INITIAL_SIZE;
38#ifdef SMALLOC_REDZONE
39static const int int_mask = sizeof(int) - 1;
40#endif
41
42struct pool {
43	struct fio_mutex *lock;			/* protects this pool */
44	void *map;				/* map of blocks */
45	unsigned int *bitmap;			/* blocks free/busy map */
46	size_t free_blocks;		/* free blocks */
47	size_t nr_blocks;			/* total blocks */
48	size_t next_non_full;
49	size_t mmap_size;
50};
51
52struct block_hdr {
53	size_t size;
54#ifdef SMALLOC_REDZONE
55	unsigned int prered;
56#endif
57};
58
59static struct pool mp[MAX_POOLS];
60static unsigned int nr_pools;
61static unsigned int last_pool;
62
63static inline int ptr_valid(struct pool *pool, void *ptr)
64{
65	unsigned int pool_size = pool->nr_blocks * SMALLOC_BPL;
66
67	return (ptr >= pool->map) && (ptr < pool->map + pool_size);
68}
69
70static inline size_t size_to_blocks(size_t size)
71{
72	return (size + SMALLOC_BPB - 1) / SMALLOC_BPB;
73}
74
75static int blocks_iter(struct pool *pool, unsigned int pool_idx,
76		       unsigned int idx, size_t nr_blocks,
77		       int (*func)(unsigned int *map, unsigned int mask))
78{
79
80	while (nr_blocks) {
81		unsigned int this_blocks, mask;
82		unsigned int *map;
83
84		if (pool_idx >= pool->nr_blocks)
85			return 0;
86
87		map = &pool->bitmap[pool_idx];
88
89		this_blocks = nr_blocks;
90		if (this_blocks + idx > SMALLOC_BPI) {
91			this_blocks = SMALLOC_BPI - idx;
92			idx = SMALLOC_BPI - this_blocks;
93		}
94
95		if (this_blocks == SMALLOC_BPI)
96			mask = -1U;
97		else
98			mask = ((1U << this_blocks) - 1) << idx;
99
100		if (!func(map, mask))
101			return 0;
102
103		nr_blocks -= this_blocks;
104		idx = 0;
105		pool_idx++;
106	}
107
108	return 1;
109}
110
111static int mask_cmp(unsigned int *map, unsigned int mask)
112{
113	return !(*map & mask);
114}
115
116static int mask_clear(unsigned int *map, unsigned int mask)
117{
118	assert((*map & mask) == mask);
119	*map &= ~mask;
120	return 1;
121}
122
123static int mask_set(unsigned int *map, unsigned int mask)
124{
125	assert(!(*map & mask));
126	*map |= mask;
127	return 1;
128}
129
130static int blocks_free(struct pool *pool, unsigned int pool_idx,
131		       unsigned int idx, size_t nr_blocks)
132{
133	return blocks_iter(pool, pool_idx, idx, nr_blocks, mask_cmp);
134}
135
136static void set_blocks(struct pool *pool, unsigned int pool_idx,
137		       unsigned int idx, size_t nr_blocks)
138{
139	blocks_iter(pool, pool_idx, idx, nr_blocks, mask_set);
140}
141
142static void clear_blocks(struct pool *pool, unsigned int pool_idx,
143			 unsigned int idx, size_t nr_blocks)
144{
145	blocks_iter(pool, pool_idx, idx, nr_blocks, mask_clear);
146}
147
148static int find_next_zero(int word, int start)
149{
150	assert(word != -1U);
151	word >>= start;
152	return ffz(word) + start;
153}
154
155static bool add_pool(struct pool *pool, unsigned int alloc_size)
156{
157	int bitmap_blocks;
158	int mmap_flags;
159	void *ptr;
160
161	if (nr_pools == MAX_POOLS)
162		return false;
163
164#ifdef SMALLOC_REDZONE
165	alloc_size += sizeof(unsigned int);
166#endif
167	alloc_size += sizeof(struct block_hdr);
168	if (alloc_size < INITIAL_SIZE)
169		alloc_size = INITIAL_SIZE;
170
171	/* round up to nearest full number of blocks */
172	alloc_size = (alloc_size + SMALLOC_BPL - 1) & ~(SMALLOC_BPL - 1);
173	bitmap_blocks = alloc_size / SMALLOC_BPL;
174	alloc_size += bitmap_blocks * sizeof(unsigned int);
175	pool->mmap_size = alloc_size;
176
177	pool->nr_blocks = bitmap_blocks;
178	pool->free_blocks = bitmap_blocks * SMALLOC_BPB;
179
180	mmap_flags = OS_MAP_ANON;
181#ifdef CONFIG_ESX
182	mmap_flags |= MAP_PRIVATE;
183#else
184	mmap_flags |= MAP_SHARED;
185#endif
186	ptr = mmap(NULL, alloc_size, PROT_READ|PROT_WRITE, mmap_flags, -1, 0);
187
188	if (ptr == MAP_FAILED)
189		goto out_fail;
190
191	pool->map = ptr;
192	pool->bitmap = (void *) ptr + (pool->nr_blocks * SMALLOC_BPL);
193	memset(pool->bitmap, 0, bitmap_blocks * sizeof(unsigned int));
194
195	pool->lock = fio_mutex_init(FIO_MUTEX_UNLOCKED);
196	if (!pool->lock)
197		goto out_fail;
198
199	nr_pools++;
200	return true;
201out_fail:
202	log_err("smalloc: failed adding pool\n");
203	if (pool->map)
204		munmap(pool->map, pool->mmap_size);
205	return false;
206}
207
208void sinit(void)
209{
210	bool ret;
211	int i;
212
213	for (i = 0; i < INITIAL_POOLS; i++) {
214		ret = add_pool(&mp[nr_pools], smalloc_pool_size);
215		if (!ret)
216			break;
217	}
218
219	/*
220	 * If we added at least one pool, we should be OK for most
221	 * cases.
222	 */
223	assert(i);
224}
225
226static void cleanup_pool(struct pool *pool)
227{
228	/*
229	 * This will also remove the temporary file we used as a backing
230	 * store, it was already unlinked
231	 */
232	munmap(pool->map, pool->mmap_size);
233
234	if (pool->lock)
235		fio_mutex_remove(pool->lock);
236}
237
238void scleanup(void)
239{
240	unsigned int i;
241
242	for (i = 0; i < nr_pools; i++)
243		cleanup_pool(&mp[i]);
244}
245
246#ifdef SMALLOC_REDZONE
247static void *postred_ptr(struct block_hdr *hdr)
248{
249	uintptr_t ptr;
250
251	ptr = (uintptr_t) hdr + hdr->size - sizeof(unsigned int);
252	ptr = (uintptr_t) PTR_ALIGN(ptr, int_mask);
253
254	return (void *) ptr;
255}
256
257static void fill_redzone(struct block_hdr *hdr)
258{
259	unsigned int *postred = postred_ptr(hdr);
260
261	hdr->prered = SMALLOC_PRE_RED;
262	*postred = SMALLOC_POST_RED;
263}
264
265static void sfree_check_redzone(struct block_hdr *hdr)
266{
267	unsigned int *postred = postred_ptr(hdr);
268
269	if (hdr->prered != SMALLOC_PRE_RED) {
270		log_err("smalloc pre redzone destroyed!\n"
271			" ptr=%p, prered=%x, expected %x\n",
272				hdr, hdr->prered, SMALLOC_PRE_RED);
273		assert(0);
274	}
275	if (*postred != SMALLOC_POST_RED) {
276		log_err("smalloc post redzone destroyed!\n"
277			"  ptr=%p, postred=%x, expected %x\n",
278				hdr, *postred, SMALLOC_POST_RED);
279		assert(0);
280	}
281}
282#else
283static void fill_redzone(struct block_hdr *hdr)
284{
285}
286
287static void sfree_check_redzone(struct block_hdr *hdr)
288{
289}
290#endif
291
292static void sfree_pool(struct pool *pool, void *ptr)
293{
294	struct block_hdr *hdr;
295	unsigned int i, idx;
296	unsigned long offset;
297
298	if (!ptr)
299		return;
300
301	ptr -= sizeof(*hdr);
302	hdr = ptr;
303
304	assert(ptr_valid(pool, ptr));
305
306	sfree_check_redzone(hdr);
307
308	offset = ptr - pool->map;
309	i = offset / SMALLOC_BPL;
310	idx = (offset % SMALLOC_BPL) / SMALLOC_BPB;
311
312	fio_mutex_down(pool->lock);
313	clear_blocks(pool, i, idx, size_to_blocks(hdr->size));
314	if (i < pool->next_non_full)
315		pool->next_non_full = i;
316	pool->free_blocks += size_to_blocks(hdr->size);
317	fio_mutex_up(pool->lock);
318}
319
320void sfree(void *ptr)
321{
322	struct pool *pool = NULL;
323	unsigned int i;
324
325	if (!ptr)
326		return;
327
328	for (i = 0; i < nr_pools; i++) {
329		if (ptr_valid(&mp[i], ptr)) {
330			pool = &mp[i];
331			break;
332		}
333	}
334
335	if (pool) {
336		sfree_pool(pool, ptr);
337		return;
338	}
339
340	log_err("smalloc: ptr %p not from smalloc pool\n", ptr);
341}
342
343static void *__smalloc_pool(struct pool *pool, size_t size)
344{
345	size_t nr_blocks;
346	unsigned int i;
347	unsigned int offset;
348	unsigned int last_idx;
349	void *ret = NULL;
350
351	fio_mutex_down(pool->lock);
352
353	nr_blocks = size_to_blocks(size);
354	if (nr_blocks > pool->free_blocks)
355		goto fail;
356
357	i = pool->next_non_full;
358	last_idx = 0;
359	offset = -1U;
360	while (i < pool->nr_blocks) {
361		unsigned int idx;
362
363		if (pool->bitmap[i] == -1U) {
364			i++;
365			pool->next_non_full = i;
366			last_idx = 0;
367			continue;
368		}
369
370		idx = find_next_zero(pool->bitmap[i], last_idx);
371		if (!blocks_free(pool, i, idx, nr_blocks)) {
372			idx += nr_blocks;
373			if (idx < SMALLOC_BPI)
374				last_idx = idx;
375			else {
376				last_idx = 0;
377				while (idx >= SMALLOC_BPI) {
378					i++;
379					idx -= SMALLOC_BPI;
380				}
381			}
382			continue;
383		}
384		set_blocks(pool, i, idx, nr_blocks);
385		offset = i * SMALLOC_BPL + idx * SMALLOC_BPB;
386		break;
387	}
388
389	if (i < pool->nr_blocks) {
390		pool->free_blocks -= nr_blocks;
391		ret = pool->map + offset;
392	}
393fail:
394	fio_mutex_up(pool->lock);
395	return ret;
396}
397
398static void *smalloc_pool(struct pool *pool, size_t size)
399{
400	size_t alloc_size = size + sizeof(struct block_hdr);
401	void *ptr;
402
403	/*
404	 * Round to int alignment, so that the postred pointer will
405	 * be naturally aligned as well.
406	 */
407#ifdef SMALLOC_REDZONE
408	alloc_size += sizeof(unsigned int);
409	alloc_size = (alloc_size + int_mask) & ~int_mask;
410#endif
411
412	ptr = __smalloc_pool(pool, alloc_size);
413	if (ptr) {
414		struct block_hdr *hdr = ptr;
415
416		hdr->size = alloc_size;
417		fill_redzone(hdr);
418
419		ptr += sizeof(*hdr);
420		memset(ptr, 0, size);
421	}
422
423	return ptr;
424}
425
426void *smalloc(size_t size)
427{
428	unsigned int i, end_pool;
429
430	if (size != (unsigned int) size)
431		return NULL;
432
433	i = last_pool;
434	end_pool = nr_pools;
435
436	do {
437		for (; i < end_pool; i++) {
438			void *ptr = smalloc_pool(&mp[i], size);
439
440			if (ptr) {
441				last_pool = i;
442				return ptr;
443			}
444		}
445		if (last_pool) {
446			end_pool = last_pool;
447			last_pool = i = 0;
448			continue;
449		}
450
451		break;
452	} while (1);
453
454	log_err("smalloc: OOM. Consider using --alloc-size to increase the "
455		"shared memory available.\n");
456	return NULL;
457}
458
459void *scalloc(size_t nmemb, size_t size)
460{
461	return smalloc(nmemb * size);
462}
463
464char *smalloc_strdup(const char *str)
465{
466	char *ptr = NULL;
467
468	ptr = smalloc(strlen(str) + 1);
469	if (ptr)
470		strcpy(ptr, str);
471	return ptr;
472}
473