Enhancements to PartitionAlloc to support threading and arbitrary allocation sizes.

Source/wtf/PartitionAlloc.h

 /*
  * Copyright (C) 2013 Google Inc. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
  * met:
  *
  *     * Redistributions of source code must retain the above copyright
  * notice, this list of conditions and the following disclaimer.
  *     * Redistributions in binary form must reproduce the above
@@ skipping 23 matching lines @@
 // DESCRIPTION
 // partitionAlloc() and partitionFree() are approximately analagous
 // to malloc() and free().
 // The main difference is that a PartitionRoot object must be supplied to
 // partitionAlloc(), representing a specific "heap partition" that will
 // be used to satisfy the allocation. Different partitions are guaranteed to
 // exist in separate address spaces, including being separate from the main
 // system heap. If the contained objects are all freed, physical memory is
 // returned to the system but the address space remains reserved.
 //
-// Allocations using this API are only supported from the Blink main thread.
-// This is ASSERT()ed in Debug builds.
+// Allocations and frees against a single partition must be single threaded.
+// Allocations must not exceed a max size, typically 4088 at this time.
+// Allocation sizes must be aligned to the system pointer size.
+// The separate APIs partitionAllocGeneric and partitionFreeGeneric are
+// provided, and they do not have the above three restrictions. In return, you
+// take a small performance hit and are also obliged to keep track of
+// allocation sizes, and pass them to partitionFreeGeneric.
 //
 // This allocator is designed to be extremely fast, thanks to the following
 // properties and design:
 // - Just a single (reasonably predicatable) branch in the hot / fast path for
 // both allocating and (significantly) freeing.
 // - A minimal number of operations in the hot / fast path, with the slow paths
 // in separate functions, leading to the possibility of inlining.
 // - Each partition page (which is usually multiple physical pages) has a header
 // structure which allows fast mapping of free() address to an underlying
 // bucket.
-// - No support for threading yet, leading to simpler design.
+// - Supports a lock-free API for fast performance in single-threaded cases.
 // - The freelist for a given bucket is split across a number of partition
 // pages, enabling various simple tricks to try and minimize fragmentation.
 // - Fine-grained bucket sizes leading to less waste and better packing.
 //
 // The following security properties are provided at this time:
 // - Linear overflows cannot corrupt into the partition.
 // - Linear overflows cannot corrupt out of the partition.
 // - Freed pages will only be re-used within the partition.
 // - Freed pages will only hold same-sized objects when re-used.
 // - Dereference of freelist pointer will fault.
 //
 // The following security properties could be investigated in the future:
 // - No double-free detection (tcmalloc has some but it may be only a detection
 // and not a defense).
 // - No randomness in freelist pointers.
 // - Per-object bucketing (instead of per-size) is mostly available at the API,
 // but not used yet.
 // - No randomness of freelist entries or bucket position.
+// - No specific protection against corruption of page header metadata.
 
 #include "wtf/Assertions.h"
-#include "wtf/MainThread.h"
+#include "wtf/Atomics.h"
+#include "wtf/FastMalloc.h"
 
 #include <stdlib.h>
 
 namespace WTF {
 
 // Allocation granularity of sizeof(void*) bytes.
-static const size_t kBucketShift = (sizeof(void*) == 8) ? 3 : 2;
-// Support allocations up to kMaxAllocation bytes.
-static const size_t kMaxAllocation = 4096;
-static const size_t kNumBuckets = kMaxAllocation / (1 << kBucketShift);
+static const size_t kAllocationGranularity = sizeof(void*);
+static const size_t kAllocationGranularityMask = kAllocationGranularity - 1;
+static const size_t kBucketShift = (kAllocationGranularity == 8) ? 3 : 2;
+// Supports allocations up to 4088 (one bucket is used for metadata).
+static const size_t kMaxAllocationOrder = 12;
+static const size_t kMaxAllocation = (1 << kMaxAllocationOrder) - kAllocationGranularity;
+static const size_t kNumBuckets = (1 << kMaxAllocationOrder) / (1 << kBucketShift);
 // Underlying partition storage pages are a power-of-two size. It is typical
 // for a partition page to be based on multiple system pages. We rarely deal
 // with system pages. Most references to "page" refer to partition pages. We
 // do also have the concept of "super pages" -- these are the underlying
 // system allocations we make. Super pages can typically fit multiple
 // partition pages inside them. See PageAllocator.h for more details on
 // super pages.
 static const size_t kPartitionPageSize = 1 << 14; // 16KB
 static const size_t kPartitionPageOffsetMask = kPartitionPageSize - 1;
 static const size_t kPartitionPageBaseMask = ~kPartitionPageOffsetMask;
@@ skipping 21 matching lines @@
 };
 
 struct PartitionBucket {
     PartitionRoot* root;
     PartitionPageHeader* currPage;
     PartitionFreepagelistEntry* freePages;
     size_t numFullPages;
 };
 
 struct PartitionRoot {
+    int lock;
     PartitionPageHeader seedPage;
     PartitionBucket seedBucket;
     PartitionBucket buckets[kNumBuckets];
     char* nextSuperPage;
     char* nextPartitionPage;
     char* nextPartitionPageEnd;
     bool initialized;
 };
 
 WTF_EXPORT void partitionAllocInit(PartitionRoot*);
 WTF_EXPORT bool partitionAllocShutdown(PartitionRoot*);
 
 WTF_EXPORT NEVER_INLINE void* partitionAllocSlowPath(PartitionBucket*);
 WTF_EXPORT NEVER_INLINE void partitionFreeSlowPath(PartitionPageHeader*);
+WTF_EXPORT NEVER_INLINE void* partitionReallocGeneric(PartitionRoot*, void*, size_t, size_t);
 
 ALWAYS_INLINE PartitionFreelistEntry* partitionFreelistMask(PartitionFreelistEntry* ptr)
 {
     // For now, use a simple / fast mask that guarantees an invalid pointer in
     // case it gets used as a vtable pointer.
     // The one attack we're fully mitigating is where an object is freed and its
     // vtable used where the attacker doesn't get the chance to run allocations
     // between the free and use.
     // We're deliberately not trying to defend against OOB reads or writes.
     uintptr_t masked = ~reinterpret_cast<uintptr_t>(ptr);
@@ skipping 19 matching lines @@
     if (LIKELY(ret != 0)) {
         page->freelistHead = partitionFreelistMask(ret->next);
         page->numAllocatedSlots++;
         return ret;
     }
     return partitionAllocSlowPath(bucket);
 }
 
 ALWAYS_INLINE void* partitionAlloc(PartitionRoot* root, size_t size)
 {
-    ASSERT(isMainThread());
-    ASSERT(size);
     size_t index = size >> kBucketShift;
     ASSERT(index < kNumBuckets);
     ASSERT(size == index << kBucketShift);
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
     return malloc(size);
 #else
     PartitionBucket* bucket = &root->buckets[index];
     return partitionBucketAlloc(bucket);
 #endif
 }
 
-ALWAYS_INLINE void partitionFree(void* ptr)
+ALWAYS_INLINE PartitionPageHeader* partitionPointerToPage(void* ptr)
 {
-    ASSERT(isMainThread());
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
-    free(ptr);
-#else
     uintptr_t pointerAsUint = reinterpret_cast<uintptr_t>(ptr);
     // Checks that the pointer is after the page header. You can't free the
     // page header!
     ASSERT((pointerAsUint & kPartitionPageOffsetMask) >= sizeof(PartitionPageHeader));
     PartitionPageHeader* page = reinterpret_cast<PartitionPageHeader*>(pointerAsUint & kPartitionPageBaseMask);
     // Checks that the pointer is a multiple of bucket size.
     ASSERT(!(((pointerAsUint & kPartitionPageOffsetMask) - sizeof(PartitionPageHeader)) % partitionBucketSize(page->bucket)));
+    return page;
+}
+
+ALWAYS_INLINE void partitionFreeWithPage(void* ptr, PartitionPageHeader* page)
+{
     PartitionFreelistEntry* entry = static_cast<PartitionFreelistEntry*>(ptr);
     entry->next = partitionFreelistMask(page->freelistHead);
     page->freelistHead = entry;
     --page->numAllocatedSlots;
     if (UNLIKELY(page->numAllocatedSlots <= 0))
         partitionFreeSlowPath(page);
+}
+
+ALWAYS_INLINE void partitionFree(void* ptr)
+{
+#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+    free(ptr);
+#else
+    PartitionPageHeader* page = partitionPointerToPage(ptr);
+    partitionFreeWithPage(ptr, page);
 #endif
 }
 
+ALWAYS_INLINE void* partitionAllocGeneric(PartitionRoot* root, size_t size)
+{
+#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+    return malloc(size);
+#else
+    if (LIKELY(size <= kMaxAllocation)) {
+        size += kAllocationGranularityMask;
+        size &= ~kAllocationGranularityMask;
+        while (atomicIncrement(&root->lock) != 1)
+            atomicDecrement(&root->lock);

[abarth-chromium, 2013/08/02 23:58:04]

Is this better than TCMalloc_SpinLock?

In any case, it's a more general primitive that we should factor out into its
own header.

You should test this on Mac, where atomicIncrement is stupidly expensive.

+        void* ret = partitionAlloc(root, size);
+        atomicDecrement(&root->lock);
+        return ret;
+    }
+    return WTF::fastMalloc(size);
+#endif
+}
+
+ALWAYS_INLINE void partitionFreeGeneric(void* ptr, size_t size)
+{
+#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+    free(ptr);
+#else
+    if (LIKELY(size <= kMaxAllocation)) {
+        PartitionPageHeader* page = partitionPointerToPage(ptr);
+        PartitionRoot* root = page->bucket->root;
+        while (atomicIncrement(&root->lock) != 1)
+            atomicDecrement(&root->lock);
+        partitionFreeWithPage(ptr, page);
+        atomicDecrement(&root->lock);
+        return;
+    }
+    return WTF::fastFree(ptr);
+#endif
+}
+
 } // namespace WTF
 
 using WTF::PartitionRoot;
 using WTF::partitionAllocInit;
 using WTF::partitionAllocShutdown;
 using WTF::partitionAlloc;
 using WTF::partitionFree;
+using WTF::partitionAllocGeneric;
+using WTF::partitionFreeGeneric;
+using WTF::partitionReallocGeneric;
 
 #endif // WTF_PartitionAlloc_h