diff --git a/src/generic/kmalloc.c b/src/generic/kmalloc.c
index 7529420..e1b3008 100644
--- a/src/generic/kmalloc.c
+++ b/src/generic/kmalloc.c
@@ -11,17 +11,28 @@
 /* Supports this many different kmalloc sizes */
 #define KMALLOC_POOLS_MAX	5
 
+struct kmalloc_pool_head {
+	struct list_head cache_list;
+	int occupied;
+	int total_caches;
+	int cache_size;
+};
+
 struct kmalloc_mempool {
 	int total;
-	struct list_head pool_head[KMALLOC_POOLS_MAX];
+	struct kmalloc_pool_head pool_head[KMALLOC_POOLS_MAX];
 	struct mutex kmalloc_mutex;
 };
 struct kmalloc_mempool km_pool;
 
 void init_kmalloc()
 {
-	for (int i = 0; i < KMALLOC_POOLS_MAX; i++)
-		INIT_LIST_HEAD(&km_pool.pool_head[i]);
+	for (int i = 0; i < KMALLOC_POOLS_MAX; i++) {
+		INIT_LIST_HEAD(&km_pool.pool_head[i].cache_list);
+		km_pool.pool_head[i].occupied = 0;
+		km_pool.pool_head[i].total_caches = 0;
+		km_pool.pool_head[i].cache_size = 0;
+	}
 	mutex_init(&km_pool.kmalloc_mutex);
 }
 
@@ -33,27 +44,32 @@ void init_kmalloc()
  */
 void *kmalloc(int size)
 {
-	struct mem_cache *cache, *n;
+	struct mem_cache *cache;
 	int right_sized_pool_idx = -1;
        	int index;
 
-	/* Search all existing pools for this size, and if found, free bufs */
+	BUG_ON(!size); /* It is a kernel bug if size is 0 */
+
 	for (int i = 0; i < km_pool.total; i++) {
-		list_for_each_entry_safe(cache, n, &km_pool.pool_head[i], list) {
-			if (cache->struct_size == size) {
-				right_sized_pool_idx = i;
+		/* Check if this pool has right size */
+		if (km_pool.pool_head[i].cache_size == size) {
+			right_sized_pool_idx = i;
+			/*
+			 * Found the pool, now see if any
+			 * cache has available slots
+			 */
+			list_for_each_entry(cache, &km_pool.pool_head[i].cache_list,
+					    list) {
 				if (cache->free)
 					return mem_cache_alloc(cache);
 				else
-					continue;
-			} else
-				break;
+					break;
+			}
 		}
 	}
 
 	/*
-	 * No such pool list already available at hand, and we don't have room
-	 * for new pool lists.
+	 * All pools are allocated and none has requested size
 	 */
 	if ((right_sized_pool_idx < 0) &&
 	    (km_pool.total == KMALLOC_POOLS_MAX - 1)) {
@@ -62,17 +78,21 @@ void *kmalloc(int size)
 		BUG();
 	}
 
+	/* A pool exists with given size? (But no cache in it is free) */
 	if (right_sized_pool_idx >= 0)
 		index = right_sized_pool_idx;
-	else
-		index = km_pool.total++;
+	else /* No pool of this size, allocate new by incrementing total */
+		index = km_pool.total++; 
 
 	/* Only allow up to page size */
 	BUG_ON(size >= PAGE_SIZE);
 	BUG_ON(!(cache = mem_cache_init(alloc_page(), PAGE_SIZE,
 					size, 0)));
 	printk("%s: Created new cache for size %d\n", __FUNCTION__, size);
-	list_add(&cache->list, &km_pool.pool_head[index]);
+	list_add(&cache->list, &km_pool.pool_head[index].cache_list);
+	km_pool.pool_head[index].occupied = 1;
+	km_pool.pool_head[index].total_caches++;
+	km_pool.pool_head[index].cache_size = size;
 	return mem_cache_alloc(cache);
 }
 
@@ -85,20 +105,22 @@ int kfree(void *p)
 	struct mem_cache *cache, *tmp;
 
 	for (int i = 0; i < km_pool.total; i++)
-		list_for_each_entry_safe(cache, tmp, &km_pool.pool_head[i], list)
+		list_for_each_entry_safe(cache, tmp,
+					 &km_pool.pool_head[i].cache_list,
+					 list) {
 			if (!mem_cache_free(cache, p)) {
 				if (mem_cache_is_empty(cache)) {
+					km_pool.pool_head[i].total_caches--;
 					list_del(&cache->list);
 					free_page(cache);
 					/*
-					 * Total remains the same unless all
-					 * caches are freed on that pool
+					 * Total remains the same but slot
+					 * may have no caches left.
 					 */
-					if (list_empty(&km_pool.pool_head[i]))
-						km_pool.total--;
 				}
 				return 0;
 			}
+		}
 	return -1;
 }