scripts

commit 2523880dc2e339eef0e00e76848721f4b2d27812
parent 859527cd00e6cb2e44b434b8c8db0a64f2ba8b4f
Author: sin <sin@2f30.org>
Date:   Fri,  2 Aug 2013 10:20:09 +0100

Make rballoc thread-safe

Diffstat:
M
random/rballoc.c
 | 
76
++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------------

1 file changed, 64 insertions(+), 12 deletions(-)
diff --git a/random/rballoc.c b/random/rballoc.c
@@ -6,6 +6,7 @@
 #include <string.h>
 #include <errno.h>
 #include "tree.h"
+#include "spinlock.h"
 
 enum { PAGESIZE = 4096 };
 
@@ -17,6 +18,7 @@ struct node {
 
 RB_HEAD(free_tree, node) ft = RB_INITIALIZER(&ft);
 RB_HEAD(alloc_tree, node) at = RB_INITIALIZER(&at);
+static spinlock_t rblock;
 
 static int ft_cmp(struct node *a, struct node *b);
 RB_PROTOTYPE(free_tree, node, entry, ft_cmp)
@@ -109,6 +111,7 @@ malloc(size_t siz)
 
 	if (!siz)
 		return NULL;
+	lock(&rblock);
 	/* Lookup in the free tree for a block greater
 	 * than or equal to `siz' bytes */
 	n.siz = siz;
@@ -117,20 +120,25 @@ malloc(size_t siz)
 		/* No available free block, create a new block
 		 * and add it to the alloc tree */
 		an = alloc_object(sizeof(*an));
-		if (!an)
+		if (!an) {
+			unlock(&rblock);
 			return NULL;
+		}
 		p = alloc_block(siz);
 		if (!p) {
 			free_object(an, sizeof(*an));
+			unlock(&rblock);
 			return NULL;
 		}
 		an->buf = p;
 		an->siz = siz;
 		RB_INSERT(alloc_tree, &at, an);
+		unlock(&rblock);
 		return an->buf;
 	}
 	an = RB_REMOVE(free_tree, &ft, res);
 	RB_INSERT(alloc_tree, &at, an);
+	unlock(&rblock);
 	return an->buf;
 }
 
@@ -138,28 +146,66 @@ void *
 realloc(void *oldp, size_t siz)
 {
 	struct node n, *res;
-	void *p;
+	struct node *oldan, *newan;
+	struct node *fn;
 
 	if (!oldp)
 		return malloc(siz);
-	if (!siz && oldp) {
-		free(oldp);
+	if (!siz) {
+		if (oldp)
+			free(oldp);
 		return NULL;
 	}
+	lock(&rblock);
 	n.buf = oldp;
 	res = RB_FIND(alloc_tree, &at, &n);
 	if (res) {
 		/* If we were asked to shrink the allocated space
 		 * just re-use it */
-		if (res->siz >= siz)
+		if (res->siz >= siz) {
+			unlock(&rblock);
 			return res->buf;
-		p = malloc(siz);
-		if (!p)
-			return NULL;
-		memcpy(p, res->buf, siz < res->siz ? siz : res->siz);
-		free(res->buf);
-		return p;
+		}
+		oldan = res;
+		/* Lookup in the free tree for a block greater
+		 * than or equal to `siz' bytes */
+		n.siz = siz;
+		res = RB_NFIND(free_tree, &ft, &n);
+		if (!res) {
+			/* No available free block, create a new block
+			 * and add it to the alloc tree */
+			newan = alloc_object(sizeof(*newan));
+			if (!newan) {
+				unlock(&rblock);
+				return NULL;
+			}
+			newan->buf = alloc_block(siz);
+			if (!newan->buf) {
+				free_object(newan, sizeof(*newan));
+				unlock(&rblock);
+				return NULL;
+			}
+			newan->siz = siz;
+			RB_INSERT(alloc_tree, &at, newan);
+		} else {
+			/* Grab the block from the free tree instead */
+			newan = RB_REMOVE(free_tree, &ft, res);
+			RB_INSERT(alloc_tree, &at, newan);
+		}
+		/* Copy over the contents at `oldp' */
+		memcpy(newan->buf, oldan->buf,
+		       siz < oldan->siz ? siz : oldan->siz);
+		/* Free `oldp' */
+		n.buf = oldan;
+		res = RB_FIND(alloc_tree, &at, &n);
+		if (res) {
+			fn = RB_REMOVE(alloc_tree, &at, res);
+			RB_INSERT(free_tree, &ft, fn);
+		}
+		unlock(&rblock);
+		return newan->buf;
 	}
+	unlock(&rblock);
 	return NULL;
 }
 
@@ -182,12 +228,14 @@ free(void *p)
 
 	if (!p)
 		return;
+	lock(&rblock);
 	n.buf = p;
 	res = RB_FIND(alloc_tree, &at, &n);
 	if (res) {
 		fn = RB_REMOVE(alloc_tree, &at, res);
 		RB_INSERT(free_tree, &ft, fn);
 	}
+	unlock(&rblock);
 }
 
 int
@@ -218,10 +266,14 @@ malloc_usable_size(void *p)
 
 	if (!p)
 		return 0;
+	lock(&rblock);
 	n.buf = p;
 	res = RB_FIND(alloc_tree, &at, &n);
-	if (res)
+	if (res) {
+		unlock(&rblock);
 		return res->siz;
+	}
+	unlock(&rblock);
 	return 0;
 }