1 /*
2 * Copyright (c) 2005-2012 The DragonFly Project.
3 * Copyright (c) 2013 François Tigeot
4 * Copyright (c) 2013 Matthew Dillon
5 * All rights reserved.
6 *
7 * This code is derived from software contributed to The DragonFly Project
8 * by Jeffrey Hsu.
9 *
10 * Redistribution and use in source and binary forms, with or without
11 * modification, are permitted provided that the following conditions
12 * are met:
13 *
14 * 1. Redistributions of source code must retain the above copyright
15 * notice, this list of conditions and the following disclaimer.
16 * 2. Redistributions in binary form must reproduce the above copyright
17 * notice, this list of conditions and the following disclaimer in
18 * the documentation and/or other materials provided with the
19 * distribution.
20 * 3. Neither the name of The DragonFly Project nor the names of its
21 * contributors may be used to endorse or promote products derived
22 * from this software without specific, prior written permission.
23 *
24 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
25 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
26 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
27 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
28 * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
29 * INCIDENTAL, SPECIAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES (INCLUDING,
30 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
31 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
32 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
33 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
34 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
35 * SUCH DAMAGE.
36 */
37
38 #ifdef USERLAND_TEST
39 /*
40 * Testing:
41 *
42 * cc -I. -DUSERLAND_TEST libkern/linux_idr.c -o /tmp/idr -g
43 */
44
45 #define _KERNEL
46 #define _KERNEL_STRUCTURES
47 #define KLD_MODULE
48 #include <stdio.h>
49 #include <stdlib.h>
50 #include <unistd.h>
51 #include <string.h>
52 #include <limits.h>
53 #include <assert.h>
54 #include <sys/idr.h>
55 #include <sys/errno.h>
56
57 #undef MALLOC_DEFINE
58 #define MALLOC_DEFINE(a, b, c)
59 #define lwkt_gettoken(x)
60 #define lwkt_reltoken(x)
61 #undef kmalloc
62 #define kmalloc(bytes, zone, flags) calloc(bytes, 1)
63 #define lwkt_token_init(a, b)
64 #define lwkt_token_uninit(a)
65 #define kfree(ptr, flags) free(ptr)
66 #define KKASSERT(a)
67 #define panic(str, ...) assert(0)
68 #define min(a, b) (((a) < (b)) ? (a) : (b))
69 #define max(a, b) (((a) > (b)) ? (a) : (b))
70
71 int
main(int ac,char ** av)72 main(int ac, char **av)
73 {
74 char buf[256];
75 struct idr idr;
76 intptr_t generation = 0x0;
77 int error;
78 int id;
79
80 idr_init(&idr);
81
82 printf("cmd> ");
83 fflush(stdout);
84 while (fgets(buf, sizeof(buf), stdin) != NULL) {
85 if (sscanf(buf, "a %d", &id) == 1) {
86 for (;;) {
87 if (idr_pre_get(&idr, 0) == 0) {
88 fprintf(stderr, "pre_get failed\n");
89 exit(1);
90 }
91 error = idr_get_new_above(&idr,
92 (void *)generation,
93 id, &id);
94 if (error == -EAGAIN)
95 continue;
96 if (error) {
97 fprintf(stderr, "get_new err %d\n",
98 error);
99 exit(1);
100 }
101 printf("allocated %d value %08x\n",
102 id, (int)generation);
103 ++generation;
104 break;
105 }
106 } else if (sscanf(buf, "f %d", &id) == 1) {
107 idr_remove(&idr, id);
108 } else if (sscanf(buf, "g %d", &id) == 1) {
109 void *res = idr_find(&idr, id);
110 printf("find %d res %p\n", id, res);
111 }
112 printf("cmd> ");
113 fflush(stdout);
114 }
115 return 0;
116 }
117
118 #else
119
120 #include <sys/idr.h>
121 #include <sys/kernel.h>
122 #include <sys/libkern.h>
123 #include <sys/malloc.h>
124 #include <sys/param.h>
125 #include <sys/systm.h>
126 #include <sys/spinlock2.h>
127 #include <sys/limits.h>
128
129 #endif
130
131 /* Must be 2^n - 1 */
132 #define IDR_DEFAULT_SIZE 255
133
134 MALLOC_DEFINE(M_IDR, "idr", "Integer ID management");
135
136 static void idr_grow(struct idr *idp, int want);
137 static void idr_reserve(struct idr *idp, int id, int incr);
138 static int idr_find_free(struct idr *idp, int want, int lim);
139
140 /*
141 * Number of nodes in right subtree, including the root.
142 */
143 static __inline int
right_subtree_size(int n)144 right_subtree_size(int n)
145 {
146 return (n ^ (n | (n + 1)));
147 }
148
149 /*
150 * Bigger ancestor.
151 */
152 static __inline int
right_ancestor(int n)153 right_ancestor(int n)
154 {
155 return (n | (n + 1));
156 }
157
158 /*
159 * Smaller ancestor.
160 */
161 static __inline int
left_ancestor(int n)162 left_ancestor(int n)
163 {
164 return ((n & (n + 1)) - 1);
165 }
166
167 static __inline void
idrfixup(struct idr * idp,int id)168 idrfixup(struct idr *idp, int id)
169 {
170 if (id < idp->idr_freeindex) {
171 idp->idr_freeindex = id;
172 }
173 while (idp->idr_lastindex >= 0 &&
174 idp->idr_nodes[idp->idr_lastindex].data == NULL
175 ) {
176 --idp->idr_lastindex;
177 }
178 }
179
180 static __inline struct idr_node *
idr_get_node(struct idr * idp,int id)181 idr_get_node(struct idr *idp, int id)
182 {
183 struct idr_node *idrnp;
184 if (id < 0 || id >= idp->idr_count)
185 return (NULL);
186 idrnp = &idp->idr_nodes[id];
187 if (idrnp->allocated == 0)
188 return (NULL);
189 return (idrnp);
190 }
191
192 static void
idr_reserve(struct idr * idp,int id,int incr)193 idr_reserve(struct idr *idp, int id, int incr)
194 {
195 while (id >= 0) {
196 idp->idr_nodes[id].allocated += incr;
197 KKASSERT(idp->idr_nodes[id].allocated >= 0);
198 id = left_ancestor(id);
199 }
200 }
201
202 static int
idr_find_free(struct idr * idp,int want,int lim)203 idr_find_free(struct idr *idp, int want, int lim)
204 {
205 int id, rsum, rsize, node;
206
207 /*
208 * Search for a free descriptor starting at the higher
209 * of want or fd_freefile. If that fails, consider
210 * expanding the ofile array.
211 *
212 * NOTE! the 'allocated' field is a cumulative recursive allocation
213 * count. If we happen to see a value of 0 then we can shortcut
214 * our search. Otherwise we run through through the tree going
215 * down branches we know have free descriptor(s) until we hit a
216 * leaf node. The leaf node will be free but will not necessarily
217 * have an allocated field of 0.
218 */
219 /* move up the tree looking for a subtree with a free node */
220 for (id = max(want, idp->idr_freeindex); id < min(idp->idr_count, lim);
221 id = right_ancestor(id)) {
222 if (idp->idr_nodes[id].allocated == 0)
223 return (id);
224
225 rsize = right_subtree_size(id);
226 if (idp->idr_nodes[id].allocated == rsize)
227 continue; /* right subtree full */
228
229 /*
230 * Free fd is in the right subtree of the tree rooted at fd.
231 * Call that subtree R. Look for the smallest (leftmost)
232 * subtree of R with an unallocated fd: continue moving
233 * down the left branch until encountering a full left
234 * subtree, then move to the right.
235 */
236 for (rsum = 0, rsize /= 2; rsize > 0; rsize /= 2) {
237 node = id + rsize;
238 rsum += idp->idr_nodes[node].allocated;
239 if (idp->idr_nodes[id].allocated == rsum + rsize) {
240 id = node; /* move to the right */
241 if (idp->idr_nodes[node].allocated == 0)
242 return (id);
243 rsum = 0;
244 }
245 }
246 return (id);
247 }
248 return (-1);
249 }
250
251 /*
252 * Blocking pre-get support, allows callers to use idr_pre_get() in
253 * combination with idr_get_new_above() such that idr_get_new_above()
254 * can be called safely with a spinlock held.
255 *
256 * Returns 0 on failure, 1 on success.
257 *
258 * Caller must hold a blockable lock.
259 */
260 int
idr_pre_get(struct idr * idp,__unused unsigned gfp_mask)261 idr_pre_get(struct idr *idp, __unused unsigned gfp_mask)
262 {
263 int want = idp->idr_maxwant;
264 int lim = INT_MAX;
265 int result = 1; /* success */
266 int id;
267
268 KKASSERT(mycpu->gd_spinlocks == 0);
269 lwkt_gettoken(&idp->idr_token);
270 for (;;) {
271 /*
272 * Grow if necessary (or if forced by the loop)
273 */
274 if (want >= idp->idr_count)
275 idr_grow(idp, want);
276
277 /*
278 * Check if a spot is available, break and return 0 if true,
279 * unless the available spot is beyond our limit. It is
280 * possible to exceed the limit due to the way array growth
281 * works.
282 *
283 * XXX we assume that the caller uses a consistent <sid> such
284 * that the idr_maxwant field is correct, otherwise we
285 * may believe that a slot is available but the caller then
286 * fails in idr_get_new_above() and loops.
287 */
288 id = idr_find_free(idp, idp->idr_maxwant, lim);
289 if (id != -1) {
290 if (id >= lim)
291 result = 0; /* failure */
292 break;
293 }
294
295 /*
296 * Return ENOSPC if our limit has been reached, otherwise
297 * loop and force growth.
298 */
299 if (idp->idr_count >= lim) {
300 result = 0; /* failure */
301 break;
302 }
303 want = idp->idr_count;
304 }
305 lwkt_reltoken(&idp->idr_token);
306 return result;
307 }
308
309 /*
310 * Allocate an integer. If -EAGAIN is returned the caller should loop
311 * and call idr_pre_get() with no locks held, and then retry the call
312 * to idr_get_new_above().
313 *
314 * Can be safely called with spinlocks held.
315 */
316 int
idr_get_new_above(struct idr * idp,void * ptr,int sid,int * id)317 idr_get_new_above(struct idr *idp, void *ptr, int sid, int *id)
318 {
319 int resid;
320
321 /*
322 * NOTE! Because the idp is initialized with a non-zero count,
323 * sid might be < idp->idr_count but idr_maxwant might not
324 * yet be initialized. So check both cases.
325 */
326 lwkt_gettoken(&idp->idr_token);
327 if (sid >= idp->idr_count || idp->idr_maxwant < sid) {
328 idp->idr_maxwant = max(idp->idr_maxwant, sid);
329 lwkt_reltoken(&idp->idr_token);
330 return -EAGAIN;
331 }
332
333 resid = idr_find_free(idp, sid, INT_MAX);
334 if (resid == -1) {
335 lwkt_reltoken(&idp->idr_token);
336 return -EAGAIN;
337 }
338
339 if (resid >= idp->idr_count)
340 panic("idr_get_new_above(): illegal resid %d", resid);
341 if (resid > idp->idr_lastindex)
342 idp->idr_lastindex = resid;
343 if (sid <= idp->idr_freeindex)
344 idp->idr_freeindex = resid;
345 *id = resid;
346 idr_reserve(idp, resid, 1);
347 idp->idr_nodes[resid].data = ptr;
348
349 lwkt_reltoken(&idp->idr_token);
350 return (0);
351 }
352
353 /*
354 * start: minimum id, inclusive
355 * end: maximum id, exclusive or INT_MAX if end is negative
356 */
357 int
idr_alloc(struct idr * idp,void * ptr,int start,int end,unsigned gfp_mask)358 idr_alloc(struct idr *idp, void *ptr, int start, int end, unsigned gfp_mask)
359 {
360 int lim = end > 0 ? end - 1 : INT_MAX;
361 int want = start;
362 int result, id;
363
364 if (start < 0)
365 return -EINVAL;
366
367 if (lim < start)
368 return -ENOSPC;
369
370 lwkt_gettoken(&idp->idr_token);
371
372 grow_again:
373 if (want >= idp->idr_count)
374 idr_grow(idp, want);
375
376 /*
377 * Check if a spot is available, break and return 0 if true,
378 * unless the available spot is beyond our limit. It is
379 * possible to exceed the limit due to the way array growth
380 * works.
381 */
382 id = idr_find_free(idp, start, INT_MAX);
383 if (id == -1) {
384 want = idp->idr_count;
385 goto grow_again;
386 }
387
388 if (id > lim) {
389 result = -ENOSPC;
390 goto done;
391 }
392
393 if (id >= idp->idr_count)
394 panic("idr_alloc(): illegal resid %d", id);
395 if (id > idp->idr_lastindex)
396 idp->idr_lastindex = id;
397 if (start <= idp->idr_freeindex)
398 idp->idr_freeindex = id;
399 result = id;
400 idr_reserve(idp, id, 1);
401 idp->idr_nodes[id].data = ptr;
402
403 done:
404 lwkt_reltoken(&idp->idr_token);
405 return result;
406 }
407
408 int
idr_get_new(struct idr * idp,void * ptr,int * id)409 idr_get_new(struct idr *idp, void *ptr, int *id)
410 {
411 return idr_get_new_above(idp, ptr, 0, id);
412 }
413
414 /*
415 * Grow the file table so it can hold through descriptor (want).
416 *
417 * Caller must hold a blockable lock.
418 */
419 static void
idr_grow(struct idr * idp,int want)420 idr_grow(struct idr *idp, int want)
421 {
422 struct idr_node *oldnodes, *newnodes;
423 int nf;
424
425 /* We want 2^n - 1 descriptors */
426 nf = idp->idr_count;
427 do {
428 nf = 2 * nf + 1;
429 } while (nf <= want);
430
431 #ifdef USERLAND_TEST
432 printf("idr_grow: %d -> %d\n", idp->idr_count, nf);
433 #endif
434
435 /* Allocate a new zero'ed node array */
436 newnodes = kmalloc(nf * sizeof(struct idr_node),
437 M_IDR, M_ZERO | M_WAITOK);
438
439 /* We might race another grow */
440 if (nf <= idp->idr_count) {
441 kfree(newnodes, M_IDR);
442 return;
443 }
444
445 /*
446 * Copy existing nodes to the beginning of the new array
447 */
448 oldnodes = idp->idr_nodes;
449 if (oldnodes) {
450 bcopy(oldnodes, newnodes,
451 idp->idr_count * sizeof(struct idr_node));
452 }
453 idp->idr_nodes = newnodes;
454 idp->idr_count = nf;
455
456 if (oldnodes) {
457 kfree(oldnodes, M_IDR);
458 }
459 idp->idr_nexpands++;
460 }
461
462 void *
idr_remove(struct idr * idp,int id)463 idr_remove(struct idr *idp, int id)
464 {
465 void *ptr;
466
467 lwkt_gettoken(&idp->idr_token);
468 if (id < 0 || id >= idp->idr_count) {
469 lwkt_reltoken(&idp->idr_token);
470 return NULL;
471 }
472 if (idp->idr_nodes[id].allocated == 0) {
473 lwkt_reltoken(&idp->idr_token);
474 return NULL;
475 }
476 ptr = idp->idr_nodes[id].data;
477 idp->idr_nodes[id].data = NULL;
478 idr_reserve(idp, id, -1);
479 idrfixup(idp, id);
480 lwkt_reltoken(&idp->idr_token);
481
482 return ptr;
483 }
484
485 /*
486 * Remove all int allocations, leave array intact.
487 *
488 * Caller must hold a blockable lock (or be in a context where holding
489 * the spinlock is not relevant).
490 */
491 void
idr_remove_all(struct idr * idp)492 idr_remove_all(struct idr *idp)
493 {
494 lwkt_gettoken(&idp->idr_token);
495 bzero(idp->idr_nodes, idp->idr_count * sizeof(struct idr_node));
496 idp->idr_lastindex = -1;
497 idp->idr_freeindex = 0;
498 idp->idr_nexpands = 0;
499 idp->idr_maxwant = 0;
500 lwkt_reltoken(&idp->idr_token);
501 }
502
503 void
idr_destroy(struct idr * idp)504 idr_destroy(struct idr *idp)
505 {
506 lwkt_token_uninit(&idp->idr_token);
507 if (idp->idr_nodes) {
508 kfree(idp->idr_nodes, M_IDR);
509 idp->idr_nodes = NULL;
510 }
511 bzero(idp, sizeof(*idp));
512 }
513
514 void *
idr_find(struct idr * idp,int id)515 idr_find(struct idr *idp, int id)
516 {
517 void *ret;
518
519 if (id < 0 || id >= idp->idr_count) {
520 ret = NULL;
521 } else if (idp->idr_nodes[id].allocated == 0) {
522 ret = NULL;
523 } else {
524 ret = idp->idr_nodes[id].data;
525 }
526 return ret;
527 }
528
529 int
idr_for_each(struct idr * idp,int (* fn)(int id,void * p,void * data),void * data)530 idr_for_each(struct idr *idp, int (*fn)(int id, void *p, void *data),
531 void *data)
532 {
533 int i, error = 0;
534 struct idr_node *nodes;
535
536 nodes = idp->idr_nodes;
537 for (i = 0; i < idp->idr_count; i++) {
538 if (nodes[i].data != NULL && nodes[i].allocated > 0) {
539 error = fn(i, nodes[i].data, data);
540 if (error != 0)
541 break;
542 }
543 }
544 return error;
545 }
546
547 void *
idr_replace(struct idr * idp,void * ptr,int id)548 idr_replace(struct idr *idp, void *ptr, int id)
549 {
550 struct idr_node *idrnp;
551 void *ret;
552
553 lwkt_gettoken(&idp->idr_token);
554 idrnp = idr_get_node(idp, id);
555 if (idrnp == NULL) {
556 ret = NULL;
557 } else {
558 ret = idrnp->data;
559 idrnp->data = ptr;
560 }
561 lwkt_reltoken(&idp->idr_token);
562 return (ret);
563 }
564
565 void
idr_init(struct idr * idp)566 idr_init(struct idr *idp)
567 {
568 bzero(idp, sizeof(struct idr));
569 idp->idr_nodes = kmalloc(IDR_DEFAULT_SIZE * sizeof(struct idr_node),
570 M_IDR, M_WAITOK | M_ZERO);
571 idp->idr_count = IDR_DEFAULT_SIZE;
572 idp->idr_lastindex = -1;
573 idp->idr_maxwant = 0;
574 lwkt_token_init(&idp->idr_token, "idr token");
575 }
576