slab: simple test for custom alignment and unaligned size

Makefile

@@ -105,6 +105,12 @@ build/tests/test_tlsf: tests/test_tlsf.c $(TEST_TLSF_DEPS) build/tests
 	@"$(CC)" $(CFLAGS) -c -o $@.o $< $(CPPFLAGS)
 	@"$(CC)" -o $@ $@.o $(TEST_TLSF_DEPS)
 
+TEST_SLAB_DEPS=build/src/test.o build/src/x86_64/address.o build/src/slab.o
+build/tests/test_slab: tests/test_slab.c $(TEST_SLAB_DEPS) build/tests
+	@printf "CC %s\n" $@
+	@"$(CC)" $(CFLAGS) -c -o $@.o $< $(CPPFLAGS)
+	@"$(CC)" -o $@ $@.o $(TEST_SLAB_DEPS)
+
 build/tests:
 	mkdir -p $@
 

include/slab.h

@@ -10,7 +10,11 @@
 /// Per-cache descriptor: one cache per object kind
 /// Usually this is allocated statically somewhere.
 struct slab_cache {
-    size_t obj_size;            // bytes per object
+    size_t obj_size;
+    size_t obj_size_align8;
+    size_t obj_align;
+    size_t obj_align_shift;
+    size_t num_slots_per_slab;
     struct slab *empty_slabs;   // slabs with all slots free
     struct slab *partial_slabs; // slabs with some free slots
     struct slab *full_slabs;    // slabs with no free slots
@@ -39,12 +43,21 @@ struct free_stack {
 /// Initializes the cache, without allocating any memory.
 /// `obj_size` must be > 0.
 /// constructor and destructor may be `NULL`, in which case allocated objects are not initialized.
+// if not specified, default alignment is 8 (suitable for most data structures).
+
 void slab_cache_init(
     struct slab_cache *cache,
     size_t obj_size,
     void (*constructor)(void *),
     void (*destructor)(void *));
 
+void slab_cache_init_with_align(
+    struct slab_cache *cache,
+    size_t obj_size,
+    size_t obj_align,
+    void (*constructor)(void*),
+    void (*destructor)(void*));
+
 /// Allocates one element from the cache.
 /// Returns `NULL` on memory exhaustion.
 void *slab_cache_alloc(struct slab_cache *cache);

include/tlsf.h

@@ -20,12 +20,6 @@
 #define LV2_BITMAP_SET(bitmap, second) BIT_SET(bitmap, second)
 #define LV2_BITMAP_CLR(bitmap, second) BIT_CLR(bitmap, second)
 
-#define IS_ALIGNED_16(sz) (((uint64_t)(sz) & (uint64_t)15) == 0)
-#define ROUND_UP(n, shift) ((((uint64_t)n) + ((1ull << shift) - 1)) & ~(uint64_t)((1ull << shift) - 1))
-#define ROUND_UP_16(n) ROUND_UP(n, 4)
-
-#define IS_POWER_OF_2(n) (__builtin_popcountl(n) == 1)
-
 #define INFO_BIT_FREE 0
 #define INFO_BIT_IS_LAST_PHYSICAL 1
 #define INFO_SIZE_MASK (~3ull)
@@ -37,13 +31,6 @@
 #define INFO_UNPACK_IS_LAST_PHYSICAL(info) BIT_GET(info, INFO_BIT_IS_LAST_PHYSICAL)
 #define INFO_UNPACK_IS_FREE(info) BIT_GET(info, INFO_BIT_FREE)
 
-// Highest bit set (0-indexed).
-// - 4 (100) has highest bit 2.
-// - 10 (1010) has highest bit 3.
-// The following holds: (1 << HIGHEST_BIT_SET(sz)) <= sz
-// Don't use for `sz == 0`.
-#define HIGHEST_BIT_SET(sz) (63 - __builtin_clzl(sz))
-
 struct tlsf_pool {
     size_t sli_shift;
     // this stuff is mostly for debug

include/util.h

@@ -33,4 +33,26 @@
 #define BIT_SET(bitmap, bit) ((bitmap) |= ((__typeof__(bitmap))1 << (bit)))
 #define BIT_CLR(bitmap, bit) ((bitmap) &= ~((__typeof__(bitmap))1 << (bit)))
 
+#define IS_ALIGNED_16(sz) (((uint64_t)(sz) & (uint64_t)15) == 0)
+#define ROUND_UP(n, shift) ((((uint64_t)n) + ((1ull << shift) - 1)) & ~(uint64_t)((1ull << shift) - 1))
+#define ROUND_UP_8(n) ROUND_UP(n, 3)
+#define ROUND_UP_16(n) ROUND_UP(n, 4)
+
+// TODO gcc says popcount cannot be found. Use manual version for now.
+#ifdef POPCOUNT_USE_BUILTIN
+#define IS_POWER_OF_2(n) (__builtin_popcountl(n) == 1)
+#else
+#define IS_POWER_OF_2(n) ((n ^ (n - 1)) > (n - 1))
+#endif
+
+
+// Highest bit set (0-indexed).
+// - 4 (100) has highest bit 2.
+// - 10 (1010) has highest bit 3.
+// The following holds: (1 << HIGHEST_BIT_SET(sz)) <= sz
+// Don't use for `sz == 0`.
+#define HIGHEST_BIT_SET(sz) (63 - __builtin_clzl((uint64_t)sz))
+
+#define LOG2(n) HIGHEST_BIT_SET(n)
+
 #endif

src/kernel.c

@@ -101,7 +101,7 @@ void main(void) {
         virtio_blk_queue_read(&ctx, 0, on_read_done, (void*)(intptr_t)0);
         uint8_t *write = kern_alloc(512);
         for (int i = 0; i < 512; i++) {
-            write[i] = i*2;
+            write[i] = i*3;
         }
         virtio_blk_queue_write(&ctx, 2, write, on_read_done, (void*)(intptr_t)2);
         kern_free(write);

src/slab.c

@@ -3,32 +3,52 @@
 #include "ram.h"
 #include "std.h"
 #include "bootparam.h"
-
+#include "util.h"
-#define EFF_SPACE_IN_SLAB (PAGE_SIZE - sizeof(struct slab))
-#define SLOT_SIZE(obj_size) ((obj_size) + sizeof(struct free_stack))
-#define SLOTS_PER_SLAB(obj_size) (EFF_SPACE_IN_SLAB / SLOT_SIZE(obj_size))
-
-#define ROUND_UP_8(sz) (((sz) + 7) & ~(size_t)0x7)
 
 #define SLAB_INIT_MAGIC 0x736c6162696e6974
 
+#define EFF_SPACE_IN_SLAB(obj_align_shift) (PG_SIZE - ROUND_UP(sizeof(struct slab), obj_align_shift))
+#define SLOT_SIZE(obj_size, obj_align_shift) ROUND_UP(ROUND_UP_8(obj_size) + sizeof(struct free_stack), obj_align_shift)
+#define NUM_SLOTS_PER_SLAB(obj_size, obj_align_shift) EFF_SPACE_IN_SLAB(obj_align_shift) / SLOT_SIZE(obj_size, obj_align_shift)
+
+void slab_cache_init_with_align(
+    struct slab_cache *cache,
+    size_t obj_size,
+    size_t obj_align,
+    void (*constructor)(void*),
+    void (*destructor)(void*))
+{
+    ASSERT(cache != NULL);
+    ASSERT(obj_size > 0);
+    ASSERT(obj_align > 0);
+    ASSERT(IS_POWER_OF_2(obj_align));
+
+    cache->obj_size = obj_size;
+    cache->obj_size_align8 = ROUND_UP_8(obj_size);
+    cache->obj_align = obj_align;
+    size_t obj_align_shift = LOG2(obj_align);
+    cache->obj_align_shift = obj_align_shift;
+    cache->num_slots_per_slab = NUM_SLOTS_PER_SLAB(obj_size, obj_align_shift);
+    // NOTE: we allow one slot per slab, even though it makes some things more complicated
+    if (cache->num_slots_per_slab == 0) {
+        PANIC("slab: object too big.");
+    }
+
+    cache->empty_slabs = NULL;
+    cache->partial_slabs = NULL;
+    cache->full_slabs = NULL;
+    cache->constructor = constructor;
+    cache->destructor = destructor;
+    cache->magic = SLAB_INIT_MAGIC;
+}
+
 void slab_cache_init(
     struct slab_cache *cache,
     size_t obj_size,
     void (*constructor)(void *),
     void (*destructor)(void *))
 {
-    ASSERT(cache != NULL);
+    slab_cache_init_with_align(cache, obj_size, 8, constructor, destructor);
-    ASSERT(obj_size > 0);
-    ASSERT(SLOT_SIZE(obj_size) <= EFF_SPACE_IN_SLAB);
-
-    cache->obj_size = ROUND_UP_8(obj_size);
-    cache->empty_slabs = NULL;
-    cache->partial_slabs = NULL;
-    cache->full_slabs = NULL;
-    cache->constructor = constructor;
-    cache->destructor = destructor;
-    cache->magic = SLAB_INIT_MAGIC;
 }
 
 #define SLAB_MAGIC 0x736c61626d616765
@@ -42,29 +62,34 @@ bool slab_cache_new_slab(struct slab_cache *cache) {
     if (!ram_alloc_frame(&ppn, RAM_PAGE_NORMAL)) {
         return false;
     }
-    void *raw = ppn_to_pointer(ppn);
+    uint8_t *raw = ppn_to_pointer(ppn);
+    uint8_t * const page_start = raw;
+    uint8_t * const page_end = page_start + PG_SIZE;
 
-    struct slab *as_slab = raw;
+    struct slab *as_slab = (struct slab *)raw;
-    as_slab->free_count = SLOTS_PER_SLAB(cache->obj_size);
 
     // build the free-slot stack in the bytes after the header
-    char *slots_base = (char *)raw + sizeof(struct slab);
+    raw += sizeof(struct slab);
-    ASSERT((intptr_t)slots_base % 8 == 0);
-
     struct free_stack *head = NULL;
 
-    for (size_t i = 0; i < as_slab->free_count; i++) {
+    for (size_t i = 0; i < cache->num_slots_per_slab; i++) {
-        // the stack structure lives after the corresponding object
+        raw = (uint8_t*)ROUND_UP(raw, cache->obj_align_shift);
-        void *obj = slots_base + i * SLOT_SIZE(cache->obj_size);
+        uint8_t *end_of_obj = raw + cache->obj_size_align8 + sizeof(struct free_stack);
+        ASSERT(end_of_obj < page_end);
+
         if (cache->constructor) {
-            cache->constructor(obj);
+            cache->constructor(raw);
         }
-        struct free_stack *node = (struct free_stack*)((char*)obj + cache->obj_size);
+        raw += cache->obj_size_align8;
-        node->magic = SLOT_MAGIC_FREE;
+        struct free_stack *stack = (struct free_stack *)raw;
-        node->next = head;
+        stack->magic = SLOT_MAGIC_FREE;
-        head = node;
+        stack->next = head;
+        head = stack;
+        raw += sizeof(struct free_stack);
     }
+
     as_slab->free_stack = head;
+    as_slab->free_count = cache->num_slots_per_slab;
 
     as_slab->cache = cache;
     as_slab->magic = SLAB_MAGIC;
@@ -86,7 +111,7 @@ static void *slab_obj_alloc(struct slab *slab) {
     // move to beginning of object
     ASSERT(free->magic == SLOT_MAGIC_FREE);
     free->magic = SLOT_MAGIC_OCCUPIED;
-    return (char*)free - slab->cache->obj_size;
+    return (char*)free - slab->cache->obj_size_align8;
 }
 
 bool slab_allocation_asserted;
@@ -127,7 +152,7 @@ void *slab_cache_alloc(struct slab_cache *cache) {
 
     // move previously empty slab to partial or full list
     cache->empty_slabs = slab->next;
-    if (SLOTS_PER_SLAB(cache->obj_size) > 1) {
+    if (cache->num_slots_per_slab > 1) {
         slab->next = cache->partial_slabs;
         cache->partial_slabs = slab;
     } else {
@@ -167,10 +192,10 @@ static void slab_cache_remove_slab_from_list(
 static void slab_obj_free(struct slab *slab, void *obj) {
     ASSERT(slab != NULL);
     ASSERT(obj != NULL);
-    ASSERT(slab->free_count < SLOTS_PER_SLAB(slab->cache->obj_size));
+    ASSERT(slab->free_count < slab->cache->num_slots_per_slab);
     ASSERT(PAGE_BASE(obj) == (void*)slab);
 
-    struct free_stack *node = (struct free_stack *)((char*)obj + slab->cache->obj_size);
+    struct free_stack *node = (struct free_stack *)((char*)obj + slab->cache->obj_size_align8);
 
     // check if `obj` is actually a valid object pointer
     ASSERT((char*)node + sizeof(struct free_stack) <= (char*)slab + PAGE_SIZE);
@@ -189,15 +214,20 @@ static void slab_cache_return_slab_to_kernel(
 {
     ASSERT(cache != NULL);
     ASSERT(slab != NULL);
-    ASSERT(slab->free_count == SLOTS_PER_SLAB(cache->obj_size));
+    ASSERT(slab->free_count == cache->num_slots_per_slab);
-    char *slots_base = (char *)slab + sizeof(struct slab);
+    uint8_t *raw = (uint8_t*)slab;
-    ASSERT((intptr_t)slots_base % 8 == 0);
+    uint8_t * const page_start = raw;
+    uint8_t * const page_end = page_start + PG_SIZE;
+    raw += sizeof(struct slab);
 
-    for (size_t i = 0; i < slab->free_count; i++) {
+    for (size_t i = 0; i < cache->num_slots_per_slab; i++) {
-        void *obj = slots_base + i * SLOT_SIZE(cache->obj_size);
+        raw = (uint8_t*)ROUND_UP(raw, cache->obj_align_shift);
+        uint8_t *end_of_obj = raw + cache->obj_size_align8 + sizeof(struct free_stack);
+        ASSERT(end_of_obj < page_end);
         if (cache->destructor) {
-            cache->destructor(obj);
+            cache->destructor(raw);
         }
+        raw = end_of_obj;
     }
 
     ram_free(ppn_from_aligned_pa(pa_from_pointer(slab)));
@@ -219,7 +249,7 @@ void slab_cache_free(void *obj) {
     // slab is either full or partial
     // afterwards it's either partial or empty
 
-    if (slab->free_count == SLOTS_PER_SLAB(slab->cache->obj_size)) {
+    if (slab->free_count == slab->cache->num_slots_per_slab) {
         // empty now
         if (was_partial_before) {
             slab_cache_remove_slab_from_list(slab->cache, slab, &slab->cache->partial_slabs);

src/virtio-block.c

@@ -214,14 +214,15 @@ void virtio_blk_on_used_notif(struct virtio_blk_context *ctx) {
     }
 }
 
-void virtio_blk_queue_read(struct virtio_blk_context *ctx, uint64_t sector, void (*cb)(uint8_t*,void*), void *userdata) {
+static void virtio_blk_io_request_setup(struct virtio_blk_context *ctx, uint64_t sector, uint8_t *write_data, void (*cb)(uint8_t*,void*), void *userdata) {
+    unsigned op = (write_data == NULL) ? VIRTIO_BLK_T_IN : VIRTIO_BLK_T_OUT;
     if (!slab_cache_initialized) {
         slab_cache_init(&request_cache, sizeof(struct queued_op), NULL, NULL);
         slab_cache_initialized = true;
     }
     struct queued_op *r = slab_cache_alloc(&request_cache);
     LIST_INSERT_TAIL(request_head, r);
-    r->op = VIRTIO_BLK_T_IN;
+    r->op = op;
     r->cb = cb;
     r->userdata = userdata;
 
@@ -238,11 +239,15 @@ void virtio_blk_queue_read(struct virtio_blk_context *ctx, uint64_t sector, void
     LE16_WRITE(desc0->flags, VIRTQ_DESC_F_NEXT);
     LE16_WRITE(desc0->next, r->desc_idx[1]);
 
-    // descriptor 1: data in (device-writable)
+    // descriptor 1: data (readable or writeable depending on request)
     struct virtq_desc *desc1 = &ctx->queue_desc[r->desc_idx[1]];
     LE64_WRITE(desc1->addr, buf_addr + 16);
     LE32_WRITE(desc1->len, 512);
-    LE16_WRITE(desc1->flags, VIRTQ_DESC_F_WRITE | VIRTQ_DESC_F_NEXT);
+    uint16_t flags = VIRTQ_DESC_F_NEXT;
+    if (op == VIRTIO_BLK_T_IN) {
+        flags |= VIRTQ_DESC_F_WRITE;
+    }
+    LE16_WRITE(desc1->flags, flags);
     LE16_WRITE(desc1->next, r->desc_idx[2]);
 
     // descriptor 2: status (device-writable)
@@ -256,7 +261,13 @@ void virtio_blk_queue_read(struct virtio_blk_context *ctx, uint64_t sector, void
 
     // fill request
     struct virtio_blk_req *req = buf_ptr;
-    virtio_blk_req_init_in(req, sector);
+    LE32_WRITE(req->type, op);
+    LE32_WRITE(req->reserved, 0);
+    LE64_WRITE(req->sector, sector);
+    if (op == VIRTIO_BLK_T_OUT) {
+        ASSERT(write_data != NULL);
+        memcpy(buf_ptr + 16, write_data, 512);
+    }
 
     // add chain head to avail queue
     struct virtq_avail *avail = ctx->avail;
@@ -270,58 +281,11 @@ void virtio_blk_queue_read(struct virtio_blk_context *ctx, uint64_t sector, void
     VIRTIO_BLK_NOTIFY(*ctx);
 }
 
+void virtio_blk_queue_read(struct virtio_blk_context *ctx, uint64_t sector, void (*cb)(uint8_t*,void*), void *userdata) {
+    virtio_blk_io_request_setup(ctx, sector, NULL, cb, userdata);
+}
+
 void virtio_blk_queue_write(struct virtio_blk_context *ctx, uint64_t sector, uint8_t *write_data, void (*cb)(uint8_t*,void*), void *userdata) {
-    if (!slab_cache_initialized) {
+    ASSERT(write_data != NULL);
-        slab_cache_init(&request_cache, sizeof(struct queued_op), NULL, NULL);
+    virtio_blk_io_request_setup(ctx, sector, write_data, cb, userdata);
-        slab_cache_initialized = true;
-    }
-    struct queued_op *r = slab_cache_alloc(&request_cache);
-    LIST_INSERT_TAIL(request_head, r);
-    r->op = VIRTIO_BLK_T_OUT;
-    r->cb = cb;
-    r->userdata = userdata;
-
-    ASSERT(virtio_blk_find_free_descriptors(ctx, 3, r->desc_idx));
-    struct ppn buf_ppn;
-    ASSERT(ram_alloc_frame_zeroed(&buf_ppn, RAM_PAGE_NORMAL));
-    uint64_t buf_addr = ppn_to_value(buf_ppn);
-    uint8_t *buf_ptr = ppn_to_pointer(buf_ppn);
-
-    // descriptor 0: header (device-readable)
-    struct virtq_desc *desc0 = &ctx->queue_desc[r->desc_idx[0]];
-    LE64_WRITE(desc0->addr, buf_addr);
-    LE32_WRITE(desc0->len, 16);
-    LE16_WRITE(desc0->flags, VIRTQ_DESC_F_NEXT);
-    LE16_WRITE(desc0->next, r->desc_idx[1]);
-
-    // descriptor 1: data out (device-readable)
-    struct virtq_desc *desc1 = &ctx->queue_desc[r->desc_idx[1]];
-    LE64_WRITE(desc1->addr, buf_addr + 16);
-    LE32_WRITE(desc1->len, 512);
-    LE16_WRITE(desc1->flags, VIRTQ_DESC_F_NEXT);
-    LE16_WRITE(desc1->next, r->desc_idx[2]);
-
-    // descriptor 2: status (device-writable)
-    struct virtq_desc *desc2 = &ctx->queue_desc[r->desc_idx[2]];
-    LE64_WRITE(desc2->addr, buf_addr + 16 + 512);
-    LE32_WRITE(desc2->len, 1);
-    LE16_WRITE(desc2->flags, VIRTQ_DESC_F_WRITE);
-    LE16_WRITE(desc2->next, 0);
-
-    r->buf_ppn = buf_ppn;
-
-    // fill request
-    virtio_blk_req_init_out((struct virtio_blk_req*)buf_ptr, sector);
-    memcpy(buf_ptr + 16, write_data, 512);
-
-    // add chain head to avail queue
-    struct virtq_avail *avail = ctx->avail;
-    LE16_WRITE(avail->ring[LE16_READ(avail->idx)], r->desc_idx[0]);
-    __asm__ __volatile__("" ::: "memory");
-    LE16_WRITE(avail->idx, LE16_READ(avail->idx) + 1); // TODO modulo?
-    __asm__ __volatile__("" ::: "memory");
-    // TODO here, check if notifications are disabled
-
-    // send avail buffer notification for read
-    VIRTIO_BLK_NOTIFY(*ctx);
 }

tests/test_slab.c

@@ -0,0 +1,74 @@
+#define _POSIX_C_SOURCE 200112L
+#include <stdlib.h>
+#include "slab.h"
+#include "test.h"
+#include "address.h"
+#include "ram.h"
+#include "string.h"
+
+extern uint64_t identity_mapping_start;
+
+bool ram_alloc_frame(struct ppn *ppn_out, enum frame_size size) {
+    TEST_ASSERT(size == RAM_PAGE_NORMAL);
+    void *addr;
+    TEST_ASSERT(posix_memalign(&addr, 4096, 4096) == 0);
+    *ppn_out = ppn_from_aligned_pa(pa_from_value((uint64_t)addr));
+    return true;
+}
+
+void ram_free(struct ppn ppn) {
+    free(ppn_to_pointer(ppn));
+}
+
+size_t num_ctor_calls = 0;
+size_t num_dtor_calls = 0;
+
+void constructor(void *obj) {
+    TEST_ASSERT(((uint64_t)obj & 0x3f) == 0);
+    uint8_t *p = obj;
+    for (int i = 0; i < 100; i++) {
+        p[i] = i;
+    }
+    num_ctor_calls++;
+}
+
+void destructor(void *obj) {
+    TEST_ASSERT(((uint64_t)obj & 0x3f) == 0);
+    num_dtor_calls++;
+}
+
+int main() {
+    identity_mapping_start = 0;
+    struct slab_cache cache;
+    TEST_SHOULD_PANIC(slab_cache_alloc(&cache));
+    slab_cache_init_with_align(&cache, 100, 64, constructor, destructor);
+
+    uint8_t *ptr0 = slab_cache_alloc(&cache);
+    TEST_ASSERT(ptr0 != NULL);
+    TEST_ASSERT(((uint64_t)ptr0 & 0x3f) == 0);
+
+    uint8_t *ptr1 = slab_cache_alloc(&cache);
+    TEST_ASSERT(ptr1 != NULL);
+    TEST_ASSERT(((uint64_t)ptr1 & 0x3f) == 0);
+
+    for (int i = 0; i < 100; i++) {
+        TEST_ASSERT(ptr0[i] == i);
+        TEST_ASSERT(ptr1[i] == i);
+    }
+
+    memset(ptr0, 42, 100);
+    memset(ptr1, 43, 100);
+
+    for (int i = 0; i < 100; i++) {
+        TEST_ASSERT(ptr0[i] == 42);
+        TEST_ASSERT(ptr1[i] == 43);
+    }
+
+    slab_cache_free(ptr0);
+    slab_cache_free(ptr1);
+    slab_cache_destroy(&cache);
+
+    TEST_ASSERT(num_ctor_calls > 0 && num_ctor_calls == num_dtor_calls);
+}
+
+// TODO test where we write past object and then expect panic on free