From 743c056f9d871647d5524e9a8fdea735a5939b14 Mon Sep 17 00:00:00 2001
From: Will Jhun <36206676+wjhun@users.noreply.github.com>
Date: Fri, 30 Jun 2023 10:57:24 -0400
Subject: [PATCH] mcache: allow use of arbitrary sizes for malloc-style
 interface

This modifies the mcache to allow using a malloc/free-style interface (where
the allocation size is not specified on a free()) when allocation sizes exceed
the largest objcache size. Such use would previously result in an error when
attempting to free such large allocations. To accomplish this, the mcache now
maintains a mapping of large allocation addresses to their corresponding
sizes. When a deallocate with an unspecified (-1ull) size occurs, a table
lookup/removal is made to check if the address corresponds to such a
"fallback" allocation and, if so, recover the allocation size.

The cost of using this fallback table is an insertion when making a
fallback (large) allocation and a table lookup/removal when deallocating a
fallback allocation of a known size, or on any deallocation (free) of an
unknown size. This should not impact the vast majority of mcache uses, for
which 1) the allocation size falls within the domain of the contained
objcaches and 2) the size parameter is valid on deallocation.
---
 src/runtime/heap/mcache.c | 80 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 71 insertions(+), 9 deletions(-)

diff --git a/src/runtime/heap/mcache.c b/src/runtime/heap/mcache.c
index 6bb66c9b6..b3da6ddb5 100644
--- a/src/runtime/heap/mcache.c
+++ b/src/runtime/heap/mcache.c
@@ -27,6 +27,7 @@ typedef struct mcache {
     u64 allocated;
     u64 parent_threshold;
     tuple mgmt;
+    table fallbacks;
 } *mcache;
 
 u64 mcache_alloc(heap h, bytes b)
@@ -41,19 +42,29 @@ u64 mcache_alloc(heap h, bytes b)
     rputs(": ");
 #endif
     if (b > m->parent_threshold) {
+        if (!m->fallbacks) {
+            m->fallbacks = allocate_table(m->meta, identity_key, pointer_equal);
+            if (m->fallbacks == INVALID_ADDRESS) {
+                rputs("mcache_alloc: failed to allocate fallbacks table\n");
+                return INVALID_PHYSICAL;
+            }
+        }
         u64 size = pad(b, m->parent->pagesize);
         u64 a = allocate_u64(m->parent, size);
         if (a != INVALID_PHYSICAL) {
             m->allocated += size;
+            table_set(m->fallbacks, pointer_from_u64(a), pointer_from_u64(b));
+        }
 #ifdef MCACHE_DEBUG
-            rputs("fallback to parent, size ");
-            print_u64(size);
-            rputs(", addr ");
-            print_u64(a);
-            rputs("\n");
+        rputs("fallback to parent, size ");
+        print_u64(b);
+        rputs(", padded to ");
+        print_u64(size);
+        rputs(", addr ");
+        print_u64(a);
+        rputs("\n");
 #endif
-            return a;
-        }
+        return a;
     }
 
     /* Could become a binary search if search set is large... */
@@ -102,15 +113,52 @@ void mcache_dealloc(heap h, u64 a, bytes b)
 #endif
 
     mcache m = (mcache)h;
+    u64 size = 0;
 
+    /* The fallback table tracks allocations that fall back to the parent
+       heap. This allows use of a "malloc-style" interface to the mcache in
+       which the allocation size is not specified on a deallocate. The cost of
+       this is a table insertion when making a fallback allocation and a table
+       lookup/removal when deallocating a fallback allocation of a known size,
+       or on any deallocation (free) of an unknown size. */
     if (b != -1ull && b > m->parent_threshold) {
-        u64 size = pad(b, m->parent->pagesize);
+        if (!m->fallbacks) {
+            rputs("mcache_dealloc: fallbacks table not allocated\n");
+            return;
+        }
+        size = u64_from_pointer(table_remove(m->fallbacks, pointer_from_u64(a)));
+        if (!size) {
+            rputs("mcache_dealloc: address ");
+            print_u64(a);
+            rputs(" (size ");
+            print_u64(b);
+            rputs(") not found in fallback table\n");
+            return;
+        }
+        if (size != b) {
+            rputs("mcache_dealloc: address ");
+            print_u64(a);
+            rputs(" (given size ");
+            print_u64(b);
+            rputs(") does not match alloc size (");
+            print_u64(size);
+            rputs("\n");
+        }
+        size = pad(size, m->parent->pagesize);
+    }
+    if (b == -1ull && m->fallbacks) {
+        size = u64_from_pointer(table_remove(m->fallbacks, pointer_from_u64(a)));
+        if (size > 0)
+            size = pad(size, m->parent->pagesize);
+    }
+
+    if (size) {
 #ifdef MCACHE_DEBUG
         rputs("dealloc size ");
         print_u64(b);
         rputs(", pagesize ");
         print_u64(m->parent->pagesize);
-        rputs(", parent alloc, padded size ");
+        rputs(", parent alloc size ");
         print_u64(size);
         rputs("\n");
 #endif
@@ -178,6 +226,19 @@ void destroy_mcache(heap h)
 	if (o)
 	    o->destroy(o);
     }
+    if (m->fallbacks) {
+        table_foreach(m->fallbacks, p, size) {
+#ifdef MCACHE_DEBUG
+            rputs("   dealloc fallback @ ");
+            print_u64(u64_from_pointer(p));
+            rputs(", size ");
+            print_u64(u64_from_pointer(size));
+            rputs("\n");
+#endif
+            deallocate(m->parent, p, u64_from_pointer(size));
+        }
+        deallocate_table(m->fallbacks);
+    }
     deallocate(m->meta, m, sizeof(struct mcache));
 }
 
@@ -291,6 +352,7 @@ heap allocate_mcache(heap meta, heap parent, int min_order, int max_order, bytes
     m->allocated = 0;
     m->parent_threshold = U64_FROM_BIT(max_order);
     m->mgmt = 0;
+    m->fallbacks = 0;
 
     for(int i = 0, order = min_order; order <= max_order; i++, order++) {
 	u64 obj_size = U64_FROM_BIT(order);