1  // SPDX-License-Identifier: GPL-2.0
2  #include <stdlib.h>
3  #include <string.h>
4  #include <malloc.h>
5  #include <pthread.h>
6  #include <unistd.h>
7  #include <assert.h>
8  
9  #include <linux/gfp.h>
10  #include <linux/poison.h>
11  #include <linux/slab.h>
12  #include <linux/radix-tree.h>
13  #include <urcu/uatomic.h>
14  
15  int nr_allocated;
16  int preempt_count;
17  int test_verbose;
18  
19  struct kmem_cache {
20  	pthread_mutex_t lock;
21  	unsigned int size;
22  	unsigned int align;
23  	int nr_objs;
24  	void *objs;
25  	void (*ctor)(void *);
26  	unsigned int non_kernel;
27  	unsigned long nr_allocated;
28  	unsigned long nr_tallocated;
29  	bool exec_callback;
30  	void (*callback)(void *);
31  	void *private;
32  };
33  
kmem_cache_set_callback(struct kmem_cache * cachep,void (* callback)(void *))34  void kmem_cache_set_callback(struct kmem_cache *cachep, void (*callback)(void *))
35  {
36  	cachep->callback = callback;
37  }
38  
kmem_cache_set_private(struct kmem_cache * cachep,void * private)39  void kmem_cache_set_private(struct kmem_cache *cachep, void *private)
40  {
41  	cachep->private = private;
42  }
43  
kmem_cache_set_non_kernel(struct kmem_cache * cachep,unsigned int val)44  void kmem_cache_set_non_kernel(struct kmem_cache *cachep, unsigned int val)
45  {
46  	cachep->non_kernel = val;
47  }
48  
kmem_cache_get_alloc(struct kmem_cache * cachep)49  unsigned long kmem_cache_get_alloc(struct kmem_cache *cachep)
50  {
51  	return cachep->size * cachep->nr_allocated;
52  }
53  
kmem_cache_nr_allocated(struct kmem_cache * cachep)54  unsigned long kmem_cache_nr_allocated(struct kmem_cache *cachep)
55  {
56  	return cachep->nr_allocated;
57  }
58  
kmem_cache_nr_tallocated(struct kmem_cache * cachep)59  unsigned long kmem_cache_nr_tallocated(struct kmem_cache *cachep)
60  {
61  	return cachep->nr_tallocated;
62  }
63  
kmem_cache_zero_nr_tallocated(struct kmem_cache * cachep)64  void kmem_cache_zero_nr_tallocated(struct kmem_cache *cachep)
65  {
66  	cachep->nr_tallocated = 0;
67  }
68  
kmem_cache_alloc_lru(struct kmem_cache * cachep,struct list_lru * lru,int gfp)69  void *kmem_cache_alloc_lru(struct kmem_cache *cachep, struct list_lru *lru,
70  		int gfp)
71  {
72  	void *p;
73  
74  	if (cachep->exec_callback) {
75  		if (cachep->callback)
76  			cachep->callback(cachep->private);
77  		cachep->exec_callback = false;
78  	}
79  
80  	if (!(gfp & __GFP_DIRECT_RECLAIM)) {
81  		if (!cachep->non_kernel) {
82  			cachep->exec_callback = true;
83  			return NULL;
84  		}
85  
86  		cachep->non_kernel--;
87  	}
88  
89  	pthread_mutex_lock(&cachep->lock);
90  	if (cachep->nr_objs) {
91  		struct radix_tree_node *node = cachep->objs;
92  		cachep->nr_objs--;
93  		cachep->objs = node->parent;
94  		pthread_mutex_unlock(&cachep->lock);
95  		node->parent = NULL;
96  		p = node;
97  	} else {
98  		pthread_mutex_unlock(&cachep->lock);
99  		if (cachep->align)
100  			posix_memalign(&p, cachep->align, cachep->size);
101  		else
102  			p = malloc(cachep->size);
103  		if (cachep->ctor)
104  			cachep->ctor(p);
105  		else if (gfp & __GFP_ZERO)
106  			memset(p, 0, cachep->size);
107  	}
108  
109  	uatomic_inc(&cachep->nr_allocated);
110  	uatomic_inc(&nr_allocated);
111  	uatomic_inc(&cachep->nr_tallocated);
112  	if (kmalloc_verbose)
113  		printf("Allocating %p from slab\n", p);
114  	return p;
115  }
116  
__kmem_cache_free_locked(struct kmem_cache * cachep,void * objp)117  void __kmem_cache_free_locked(struct kmem_cache *cachep, void *objp)
118  {
119  	assert(objp);
120  	if (cachep->nr_objs > 10 || cachep->align) {
121  		memset(objp, POISON_FREE, cachep->size);
122  		free(objp);
123  	} else {
124  		struct radix_tree_node *node = objp;
125  		cachep->nr_objs++;
126  		node->parent = cachep->objs;
127  		cachep->objs = node;
128  	}
129  }
130  
kmem_cache_free_locked(struct kmem_cache * cachep,void * objp)131  void kmem_cache_free_locked(struct kmem_cache *cachep, void *objp)
132  {
133  	uatomic_dec(&nr_allocated);
134  	uatomic_dec(&cachep->nr_allocated);
135  	if (kmalloc_verbose)
136  		printf("Freeing %p to slab\n", objp);
137  	__kmem_cache_free_locked(cachep, objp);
138  }
139  
kmem_cache_free(struct kmem_cache * cachep,void * objp)140  void kmem_cache_free(struct kmem_cache *cachep, void *objp)
141  {
142  	pthread_mutex_lock(&cachep->lock);
143  	kmem_cache_free_locked(cachep, objp);
144  	pthread_mutex_unlock(&cachep->lock);
145  }
146  
kmem_cache_free_bulk(struct kmem_cache * cachep,size_t size,void ** list)147  void kmem_cache_free_bulk(struct kmem_cache *cachep, size_t size, void **list)
148  {
149  	if (kmalloc_verbose)
150  		pr_debug("Bulk free %p[0-%lu]\n", list, size - 1);
151  
152  	pthread_mutex_lock(&cachep->lock);
153  	for (int i = 0; i < size; i++)
154  		kmem_cache_free_locked(cachep, list[i]);
155  	pthread_mutex_unlock(&cachep->lock);
156  }
157  
kmem_cache_shrink(struct kmem_cache * cachep)158  void kmem_cache_shrink(struct kmem_cache *cachep)
159  {
160  }
161  
kmem_cache_alloc_bulk(struct kmem_cache * cachep,gfp_t gfp,size_t size,void ** p)162  int kmem_cache_alloc_bulk(struct kmem_cache *cachep, gfp_t gfp, size_t size,
163  			  void **p)
164  {
165  	size_t i;
166  
167  	if (kmalloc_verbose)
168  		pr_debug("Bulk alloc %lu\n", size);
169  
170  	pthread_mutex_lock(&cachep->lock);
171  	if (cachep->nr_objs >= size) {
172  		struct radix_tree_node *node;
173  
174  		for (i = 0; i < size; i++) {
175  			if (!(gfp & __GFP_DIRECT_RECLAIM)) {
176  				if (!cachep->non_kernel)
177  					break;
178  				cachep->non_kernel--;
179  			}
180  
181  			node = cachep->objs;
182  			cachep->nr_objs--;
183  			cachep->objs = node->parent;
184  			p[i] = node;
185  			node->parent = NULL;
186  		}
187  		pthread_mutex_unlock(&cachep->lock);
188  	} else {
189  		pthread_mutex_unlock(&cachep->lock);
190  		for (i = 0; i < size; i++) {
191  			if (!(gfp & __GFP_DIRECT_RECLAIM)) {
192  				if (!cachep->non_kernel)
193  					break;
194  				cachep->non_kernel--;
195  			}
196  
197  			if (cachep->align) {
198  				posix_memalign(&p[i], cachep->align,
199  					       cachep->size);
200  			} else {
201  				p[i] = malloc(cachep->size);
202  				if (!p[i])
203  					break;
204  			}
205  			if (cachep->ctor)
206  				cachep->ctor(p[i]);
207  			else if (gfp & __GFP_ZERO)
208  				memset(p[i], 0, cachep->size);
209  		}
210  	}
211  
212  	if (i < size) {
213  		size = i;
214  		pthread_mutex_lock(&cachep->lock);
215  		for (i = 0; i < size; i++)
216  			__kmem_cache_free_locked(cachep, p[i]);
217  		pthread_mutex_unlock(&cachep->lock);
218  		return 0;
219  	}
220  
221  	for (i = 0; i < size; i++) {
222  		uatomic_inc(&nr_allocated);
223  		uatomic_inc(&cachep->nr_allocated);
224  		uatomic_inc(&cachep->nr_tallocated);
225  		if (kmalloc_verbose)
226  			printf("Allocating %p from slab\n", p[i]);
227  	}
228  
229  	return size;
230  }
231  
232  struct kmem_cache *
kmem_cache_create(const char * name,unsigned int size,unsigned int align,unsigned int flags,void (* ctor)(void *))233  kmem_cache_create(const char *name, unsigned int size, unsigned int align,
234  		unsigned int flags, void (*ctor)(void *))
235  {
236  	struct kmem_cache *ret = malloc(sizeof(*ret));
237  
238  	pthread_mutex_init(&ret->lock, NULL);
239  	ret->size = size;
240  	ret->align = align;
241  	ret->nr_objs = 0;
242  	ret->nr_allocated = 0;
243  	ret->nr_tallocated = 0;
244  	ret->objs = NULL;
245  	ret->ctor = ctor;
246  	ret->non_kernel = 0;
247  	ret->exec_callback = false;
248  	ret->callback = NULL;
249  	ret->private = NULL;
250  	return ret;
251  }
252  
253  /*
254   * Test the test infrastructure for kem_cache_alloc/free and bulk counterparts.
255   */
test_kmem_cache_bulk(void)256  void test_kmem_cache_bulk(void)
257  {
258  	int i;
259  	void *list[12];
260  	static struct kmem_cache *test_cache, *test_cache2;
261  
262  	/*
263  	 * Testing the bulk allocators without aligned kmem_cache to force the
264  	 * bulk alloc/free to reuse
265  	 */
266  	test_cache = kmem_cache_create("test_cache", 256, 0, SLAB_PANIC, NULL);
267  
268  	for (i = 0; i < 5; i++)
269  		list[i] = kmem_cache_alloc(test_cache, __GFP_DIRECT_RECLAIM);
270  
271  	for (i = 0; i < 5; i++)
272  		kmem_cache_free(test_cache, list[i]);
273  	assert(test_cache->nr_objs == 5);
274  
275  	kmem_cache_alloc_bulk(test_cache, __GFP_DIRECT_RECLAIM, 5, list);
276  	kmem_cache_free_bulk(test_cache, 5, list);
277  
278  	for (i = 0; i < 12 ; i++)
279  		list[i] = kmem_cache_alloc(test_cache, __GFP_DIRECT_RECLAIM);
280  
281  	for (i = 0; i < 12; i++)
282  		kmem_cache_free(test_cache, list[i]);
283  
284  	/* The last free will not be kept around */
285  	assert(test_cache->nr_objs == 11);
286  
287  	/* Aligned caches will immediately free */
288  	test_cache2 = kmem_cache_create("test_cache2", 128, 128, SLAB_PANIC, NULL);
289  
290  	kmem_cache_alloc_bulk(test_cache2, __GFP_DIRECT_RECLAIM, 10, list);
291  	kmem_cache_free_bulk(test_cache2, 10, list);
292  	assert(!test_cache2->nr_objs);
293  
294  
295  }
296