Extend the kernel unit number allocator for allocating specific unit

numbers. This change adds a new function alloc_unr_specific() which
returns the requested unit number if it is free. If the number is
already allocated or out of the range, -1 is returned.

Update alloc_unr(9) manual page accordingly and add a MLINK for
alloc_unr_specific(9).

Discussed on:	freebsd-hackers

4 files changed, 181 insertions, 14 deletions

diff --git a/share/man/man9/Makefile b/share/man/man9/Makefile
index 4e8b0ce..6e93244 100644
--- a/share/man/man9/Makefile
+++ b/share/man/man9/Makefile
@@ -374,6 +374,7 @@ MAN=	accept_filter.9 \
 	zone.9
 
 MLINKS=	alloc_unr.9 alloc_unrl.9 \
+	alloc_unr.9 alloc_unr_specific.9 \
 	alloc_unr.9 delete_unrhdr.9 \
 	alloc_unr.9 free_unr.9 \
 	alloc_unr.9 new_unrhdr.9
diff --git a/share/man/man9/alloc_unr.9 b/share/man/man9/alloc_unr.9
index 2ba93e6..43c145f 100644
--- a/share/man/man9/alloc_unr.9
+++ b/share/man/man9/alloc_unr.9
@@ -24,13 +24,14 @@
 .\"
 .\" $FreeBSD$
 .\"
-.Dd February 7, 2010
+.Dd July 5, 2010
 .Dt ALLOC_UNR 9
 .Os
 .Sh NAME
 .Nm new_unrhdr ,
 .Nm delete_unrhdr ,
 .Nm alloc_unr ,
+.Nm alloc_unr_specific ,
 .Nm free_unr
 .Nd "kernel unit number allocator"
 .Sh SYNOPSIS
@@ -43,6 +44,8 @@
 .Fn alloc_unr "struct unrhdr *uh"
 .Ft int
 .Fn alloc_unrl "struct unrhdr *uh"
+.Ft int
+.Fn alloc_unr_specific "struct unrhdr *uh" "u_int item"
 .Ft void
 .Fn free_unr "struct unrhdr *uh" "u_int item"
 .Sh DESCRIPTION
@@ -81,6 +84,13 @@ is returned.
 Same as
 .Fn alloc_unr
 except that mutex is assumed to be already locked and thus is not used.
+.It Fn alloc_unr_specific uh item
+Allocate a specific unit number.
+This function allocates memory and thus may sleep.
+The allocated unit number is returned on success.
+If the specified number is already allocated or out of the range,
+.Li \-1
+is returned.
 .It Fn free_unr uh item
 Free a previously allocated unit number.
 This function may require allocating memory, and thus it can sleep.
diff --git a/sys/kern/subr_unit.c b/sys/kern/subr_unit.c
index c613383..84e7227 100644
--- a/sys/kern/subr_unit.c
+++ b/sys/kern/subr_unit.c
@@ -628,6 +628,132 @@ alloc_unr(struct unrhdr *uh)
 	return (i);
 }
 
+static int
+alloc_unr_specificl(struct unrhdr *uh, u_int item, void **p1, void **p2)
+{
+	struct unr *up, *upn;
+	struct unrb *ub;
+	u_int i, last, tl;
+
+	mtx_assert(uh->mtx, MA_OWNED);
+
+	if (item < uh->low + uh->first || item > uh->high)
+		return (-1);
+
+	up = TAILQ_FIRST(&uh->head);
+	/* Ideal split. */
+	if (up == NULL && item - uh->low == uh->first) {
+		uh->first++;
+		uh->last--;
+		uh->busy++;
+		check_unrhdr(uh, __LINE__);
+		return (item);
+	}
+
+	i = item - uh->low - uh->first;
+
+	if (up == NULL) {
+		up = new_unr(uh, p1, p2);
+		up->ptr = NULL;
+		up->len = i;
+		TAILQ_INSERT_TAIL(&uh->head, up, list);
+		up = new_unr(uh, p1, p2);
+		up->ptr = uh;
+		up->len = 1;
+		TAILQ_INSERT_TAIL(&uh->head, up, list);
+		uh->last = uh->high - uh->low - i;
+		uh->busy++;
+		check_unrhdr(uh, __LINE__);
+		return (item);
+	} else {
+		/* Find the item which contains the unit we want to allocate. */
+		TAILQ_FOREACH(up, &uh->head, list) {
+			if (up->len > i)
+				break;
+			i -= up->len;
+		}
+	}
+
+	if (up == NULL) {
+		if (i > 0) {
+			up = new_unr(uh, p1, p2);
+			up->ptr = NULL;
+			up->len = i;
+			TAILQ_INSERT_TAIL(&uh->head, up, list);
+		}
+		up = new_unr(uh, p1, p2);
+		up->ptr = uh;
+		up->len = 1;
+		TAILQ_INSERT_TAIL(&uh->head, up, list);
+		goto done;
+	}
+
+	if (is_bitmap(uh, up)) {
+		ub = up->ptr;
+		if (bit_test(ub->map, i) == 0) {
+			bit_set(ub->map, i);
+			ub->busy++;
+			goto done;
+		} else
+			return (-1);
+	} else if (up->ptr == uh)
+		return (-1);
+
+	KASSERT(up->ptr == NULL,
+	    ("alloc_unr_specificl: up->ptr != NULL (up=%p)", up));
+
+	/* Split off the tail end, if any. */
+	tl = up->len - (1 + i);
+	if (tl > 0) {
+		upn = new_unr(uh, p1, p2);
+		upn->ptr = NULL;
+		upn->len = tl;
+		TAILQ_INSERT_AFTER(&uh->head, up, upn, list);
+	}
+
+	/* Split off head end, if any */
+	if (i > 0) {
+		upn = new_unr(uh, p1, p2);
+		upn->len = i;
+		upn->ptr = NULL;
+		TAILQ_INSERT_BEFORE(up, upn, list);
+	}
+	up->len = 1;
+	up->ptr = uh;
+
+done:
+	last = uh->high - uh->low - (item - uh->low);
+	if (uh->last > last)
+		uh->last = last;
+	uh->busy++;
+	collapse_unr(uh, up);
+	check_unrhdr(uh, __LINE__);
+	return (item);
+}
+
+int
+alloc_unr_specific(struct unrhdr *uh, u_int item)
+{
+	void *p1, *p2;
+	int i;
+
+	WITNESS_WARN(WARN_GIANTOK | WARN_SLEEPOK, NULL, "alloc_unr_specific");
+
+	p1 = Malloc(sizeof(struct unr));
+	p2 = Malloc(sizeof(struct unr));
+
+	mtx_lock(uh->mtx);
+	i = alloc_unr_specificl(uh, item, &p1, &p2);
+	mtx_unlock(uh->mtx);
+
+	if (p1 != NULL)
+		Free(p1);
+	if (p2 != NULL)
+		Free(p2);
+
+	return (i);
+}
+
 /*
  * Free a unr.
  *
@@ -810,6 +936,42 @@ print_unrhdr(struct unrhdr *uh)
 	}
 }
 
+static void
+test_alloc_unr(struct unrhdr *uh, u_int i, char a[])
+{
+	int j;
+
+	if (a[i]) {
+		printf("F %u\n", i);
+		free_unr(uh, i);
+		a[i] = 0;
+	} else {
+		no_alloc = 1;
+		j = alloc_unr(uh);
+		if (j != -1) {
+			a[j] = 1;
+			printf("A %d\n", j);
+		}
+		no_alloc = 0;
+	}
+}
+
+static void
+test_alloc_unr_specific(struct unrhdr *uh, u_int i, char a[])
+{
+	int j;
+
+	j = alloc_unr_specific(uh, i);
+	if (j == -1) {
+		printf("F %u\n", i);
+		a[i] = 0;
+		free_unr(uh, i);
+	} else {
+		a[i] = 1;
+		printf("A %d\n", j);
+	}
+}
+
 /* Number of unrs to test */
 #define NN	10000
 
@@ -825,6 +987,7 @@ main(int argc __unused, const char **argv __unused)
 	print_unrhdr(uh);
 
 	memset(a, 0, sizeof a);
+	srandomdev();
 
 	fprintf(stderr, "sizeof(struct unr) %zu\n", sizeof(struct unr));
 	fprintf(stderr, "sizeof(struct unrb) %zu\n", sizeof(struct unrb));
@@ -838,19 +1001,11 @@ main(int argc __unused, const char **argv __unused)
 		if (a[i] && (j & 1))
 			continue;
 #endif
-		if (a[i]) {
-			printf("F %u\n", i);
-			free_unr(uh, i);
-			a[i] = 0;
-		} else {
-			no_alloc = 1;
-			i = alloc_unr(uh);
-			if (i != -1) {
-				a[i] = 1;
-				printf("A %u\n", i);
-			}
-			no_alloc = 0;
-		}
+		if ((random() & 1) != 0)
+			test_alloc_unr(uh, i, a);
+		else
+			test_alloc_unr_specific(uh, i, a);
+
 		if (1)	/* XXX: change this for detailed debug printout */
 			print_unrhdr(uh);
 		check_unrhdr(uh, __LINE__);
diff --git a/sys/sys/systm.h b/sys/sys/systm.h
index a1df1c8..0b6d052 100644
--- a/sys/sys/systm.h
+++ b/sys/sys/systm.h
@@ -363,6 +363,7 @@ void delete_unrhdr(struct unrhdr *uh);
 void clean_unrhdr(struct unrhdr *uh);
 void clean_unrhdrl(struct unrhdr *uh);
 int alloc_unr(struct unrhdr *uh);
+int alloc_unr_specific(struct unrhdr *uh, u_int item);
 int alloc_unrl(struct unrhdr *uh);
 void free_unr(struct unrhdr *uh, u_int item);