kmalloc_incref()
authorBarret Rhoden <brho@cs.berkeley.edu>
Fri, 20 Jun 2014 19:52:24 +0000 (12:52 -0700)
committerBarret Rhoden <brho@cs.berkeley.edu>
Fri, 20 Jun 2014 20:01:27 +0000 (13:01 -0700)
All kmalloc'd blobs are reference counted.  kmalloc() gives you one ref.
You can get more with kmalloc_incref().  Release them with kfree().

Note that if you krealloc, you may or may not get a new buffer.  If you
have a buffer with more than one ref, and if you realloc and get a new
buffer, the caller to krealloc gets the new buffer with refcnt == 1, and
the original buffer gets decreffed by 1.  At that point, both ref users
are pointing to different blobs with the same, but copied, contents.

kern/include/kmalloc.h
kern/src/kmalloc.c
kern/src/ktest/Kconfig.postboot
kern/src/ktest/pb_ktests.c

index 1f92f4d..794a9ad 100644 (file)
@@ -21,8 +21,9 @@ void* (DALLOC(size) kmalloc)(size_t size, int flags);
 void* (DALLOC(size) kzmalloc)(size_t size, int flags);
 void *kmalloc_align(size_t size, int flags, size_t align);
 void *kzmalloc_align(size_t size, int flags, size_t align);
-void* (DALLOC(size) krealloc)(void* buf, size_t size, int flags);
-void  (DFREE(addr) kfree)(void *addr);
+void *krealloc(void *buf, size_t size, int flags);
+void kmalloc_incref(void *buf);
+void kfree(void *buf);
 void kmalloc_canary_check(char *str);
 void *debug_canary;
 
@@ -52,6 +53,7 @@ struct kmalloc_tag {
                size_t num_pages WHEN(flags == KMALLOC_TAG_PAGES);
                uint64_t unused_force_align;
        };
+       struct kref kref;
        uint32_t canary;
        int flags;
 };
index 8923606..b6910c4 100644 (file)
@@ -30,6 +30,9 @@ static spinlock_t pages_list_lock = SPINLOCK_INITIALIZER;
 static page_list_t LCKD(&pages_list_lock)pages_list;
 
 struct kmem_cache *kmalloc_caches[NUM_KMALLOC_CACHES];
+
+static void __kfree_release(struct kref *kref);
+
 void kmalloc_init(void)
 {
        /* we want at least a 16 byte alignment of the tag so that the bufs kmalloc
@@ -69,6 +72,7 @@ void *kmalloc(size_t size, int flags)
                tag->flags = KMALLOC_TAG_PAGES;
                tag->num_pages = num_pgs;
                tag->canary = KMALLOC_CANARY;
+               kref_init(&tag->kref, __kfree_release, 1);
                return buf + sizeof(struct kmalloc_tag);
        }
        // else, alloc from the appropriate cache
@@ -80,6 +84,7 @@ void *kmalloc(size_t size, int flags)
        tag->flags = KMALLOC_TAG_CACHE;
        tag->my_cache = kmalloc_caches[cache_id];
        tag->canary = KMALLOC_CANARY;
+       kref_init(&tag->kref, __kfree_release, 1);
        return buf + sizeof(struct kmalloc_tag);
 }
 
@@ -192,23 +197,40 @@ void *krealloc(void* buf, size_t size, int flags)
        return nbuf;
 }
 
-void kfree(void *addr)
+/* Grabs a reference on a buffer.  Release with kfree().
+ *
+ * Note that a krealloc on a buffer with ref > 1 that needs a new, underlying
+ * buffer will result in two buffers existing.  In this case, the krealloc is a
+ * kmalloc and a kfree, but that kfree does not completely free since the
+ * original ref > 1. */
+void kmalloc_incref(void *buf)
 {
-       struct kmalloc_tag *tag;
-       void *orig_buf;
-       if (addr == NULL)
-               return;
-       if ((orig_buf = __get_unaligned_orig_buf(addr))) {
-               kfree(orig_buf);
-               return;
-       }
-       tag = __get_km_tag(addr);
+       void *orig_buf = __get_unaligned_orig_buf(buf);
+       buf = orig_buf ? orig_buf : buf;
+       /* if we want a smaller tag, we can extract the code from kref and manually
+        * set the release method in kfree. */
+       kref_get(&__get_km_tag(buf)->kref, 1);
+}
+
+static void __kfree_release(struct kref *kref)
+{
+       struct kmalloc_tag *tag = container_of(kref, struct kmalloc_tag, kref);
        if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_CACHE)
                kmem_cache_free(tag->my_cache, tag);
-       else if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_PAGES) {
+       else if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_PAGES)
                free_cont_pages(tag, LOG2_UP(tag->num_pages));
-       } else 
-               panic("[Italian Accent]: Che Cazzo! BO! Flag in kmalloc!!!");
+       else
+               panic("Bad flag 0x%x in %s", tag->flags, __FUNCTION__);
+}
+
+void kfree(void *buf)
+{
+       void *orig_buf;
+       if (buf == NULL)
+               return;
+       orig_buf = __get_unaligned_orig_buf(buf);
+       buf = orig_buf ? orig_buf : buf;
+       kref_put(&__get_km_tag(buf)->kref);
 }
 
 void kmalloc_canary_check(char *str)
index a5a2ed2..205647f 100644 (file)
@@ -235,3 +235,8 @@ config TEST_alarm
     default n
     help
         Run the alarm test
+
+config TEST_kmalloc_incref
+    depends on PB_KTESTS
+    bool "Kmalloc incref"
+    default n
index 4df0ba2..e7b165f 100644 (file)
@@ -1973,6 +1973,52 @@ bool test_alarm(void)
        return true;
 }
 
+bool test_kmalloc_incref(void)
+{
+       /* this test is a bit invasive of the kmalloc internals */
+       void *__get_unaligned_orig_buf(void *buf)
+       {
+               int *tag_flags = (int*)(buf - sizeof(int));
+               if ((*tag_flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_UNALIGN)
+                       return (buf - (*tag_flags >> KMALLOC_ALIGN_SHIFT));
+               else
+                       return 0;
+       }
+
+       bool test_buftag(void *b, struct kmalloc_tag *btag, char *str)
+       {
+               KT_ASSERT_M(str, kref_refcnt(&btag->kref) == 1);
+               kmalloc_incref(b);
+               KT_ASSERT_M(str, kref_refcnt(&btag->kref) == 2);
+               kfree(b);
+               KT_ASSERT_M(str, kref_refcnt(&btag->kref) == 1);
+               kfree(b);
+               /* dangerous read, it's been freed */
+               KT_ASSERT_M(str, kref_refcnt(&btag->kref) == 0);
+               return TRUE;
+       }
+
+       void *b1, *b2, *b2o;
+       struct kmalloc_tag *b1tag, *b2tag;
+
+       /* no realigned case */
+       b1 = kmalloc(55, 0);
+       KT_ASSERT(!__get_unaligned_orig_buf(b1));
+       b1tag = (struct kmalloc_tag*)(b1 - sizeof(struct kmalloc_tag));
+
+       /* realigned case.  alloc'd before b1's test, so we know we get different
+        * buffers. */
+       b2 = kmalloc_align(55, 0, 64);
+       b2o = __get_unaligned_orig_buf(b2);
+       KT_ASSERT(b2o);
+       b2tag = (struct kmalloc_tag*)(b2o - sizeof(struct kmalloc_tag));
+
+       test_buftag(b1, b1tag, "b1, no realign");
+       test_buftag(b2, b2tag, "b2, realigned");
+
+       return TRUE;
+}
+
 static struct ktest ktests[] = {
 #ifdef CONFIG_X86
        KTEST_REG(ipi_sending,        CONFIG_TEST_ipi_sending),
@@ -2010,7 +2056,8 @@ static struct ktest ktests[] = {
        KTEST_REG(apipe,              CONFIG_TEST_apipe),
        KTEST_REG(rwlock,             CONFIG_TEST_rwlock),
        KTEST_REG(rv,                 CONFIG_TEST_rv),
-       KTEST_REG(alarm,              CONFIG_TEST_alarm)
+       KTEST_REG(alarm,              CONFIG_TEST_alarm),
+       KTEST_REG(kmalloc_incref,     CONFIG_TEST_kmalloc_incref),
 };
 static int num_ktests = sizeof(ktests) / sizeof(struct ktest);
 linker_func_1(register_pb_ktests)