move page heap allocations with alignment into page heap

src/page_heap.cc

@@ -163,6 +163,40 @@ Span* PageHeap::New(Length n) {
   return SearchFreeAndLargeLists(n);
 }
 
+Span* PageHeap::NewAligned(Length n, Length align_pages) {
+  // Allocate extra pages and carve off an aligned portion
+  const Length alloc = n + align_pages;
+  if (alloc < n || alloc < align_pages) {
+    // overflow means we asked huge amounts, so lets trigger normal
+    // oom handling by asking enough to trigger oom.
+    Span* span = New(std::numeric_limits<Length>::max());
+    CHECK_CONDITION(span == nullptr);
+    return nullptr;
+  }
+  Span* span = New(alloc);
+  if (PREDICT_FALSE(span == nullptr)) return nullptr;
+
+  // Skip starting portion so that we end up aligned
+  Length skip = 0;
+  size_t align_bytes = align_pages << kPageShift;
+  while ((((span->start+skip) << kPageShift) & (align_bytes - 1)) != 0) {
+    skip++;
+  }
+  ASSERT(skip < alloc);
+  if (skip > 0) {
+    Span* rest = Split(span, skip);
+    Delete(span);
+    span = rest;
+  }
+
+  ASSERT(span->length >= n);
+  if (span->length > n) {
+    Span* trailer = Split(span, n);
+    Delete(trailer);
+  }
+  return span;
+}
+
 Span* PageHeap::AllocLarge(Length n) {
   Span *best = NULL;
   Span *best_normal = NULL;

src/page_heap.h

@@ -118,6 +118,10 @@ class PERFTOOLS_DLL_DECL PageHeap {
   // been rounded up already.
   Span* New(Length n);
 
+  // Same as above but with alignment. Requires page heap
+  // lock, like New above.
+  Span* NewAligned(Length n, Length align_pages);
+
   // Delete the span "[p, p+n-1]".
   // REQUIRES: span was returned by earlier call to New() and
   //           has not yet been deleted.
@@ -129,15 +133,9 @@ class PERFTOOLS_DLL_DECL PageHeap {
   //           and has not yet been deleted.
   void RegisterSizeClass(Span* span, uint32 sc);
 
-  // Split an allocated span into two spans: one of length "n" pages
+  Span* SplitForTest(Span* span, Length n) {
-  // followed by another span of length "span->length - n" pages.
+    return Split(span, n);
-  // Modifies "*span" to point to the first span of length "n" pages.
+  }
-  // Returns a pointer to the second span.
-  //
-  // REQUIRES: "0 < n < span->length"
-  // REQUIRES: span->location == IN_USE
-  // REQUIRES: span->sizeclass == 0
-  Span* Split(Span* span, Length n);
 
   // Return the descriptor for the specified page.  Returns NULL if
   // this PageID was not allocated previously.
@@ -280,6 +278,16 @@ class PERFTOOLS_DLL_DECL PageHeap {
   // Statistics on system, free, and unmapped bytes
   Stats stats_;
 
+  // Split an allocated span into two spans: one of length "n" pages
+  // followed by another span of length "span->length - n" pages.
+  // Modifies "*span" to point to the first span of length "n" pages.
+  // Returns a pointer to the second span.
+  //
+  // REQUIRES: "0 < n < span->length"
+  // REQUIRES: span->location == IN_USE
+  // REQUIRES: span->sizeclass == 0
+  Span* Split(Span* span, Length n);
+
   Span* SearchFreeAndLargeLists(Length n);
 
   bool GrowHeap(Length n);

src/tcmalloc.cc

@@ -1630,30 +1630,13 @@ void* do_memalign_pages(size_t align, size_t size) {
   // We will allocate directly from the page heap
   SpinLockHolder h(Static::pageheap_lock());
 
-  // Allocate extra pages and carve off an aligned portion
+  Span* span = Static::pageheap()->NewAligned(tcmalloc::pages(size),
-  const Length alloc = tcmalloc::pages(size + align);
+                                              tcmalloc::pages(align));
-  Span* span = Static::pageheap()->New(alloc);
+  if (span == nullptr) {
-  if (PREDICT_FALSE(span == NULL)) return NULL;
+    // errno was set inside page heap as necessary.
-
+    return nullptr;
-  // Skip starting portion so that we end up aligned
-  Length skip = 0;
-  while ((((span->start+skip) << kPageShift) & (align - 1)) != 0) {
-    skip++;
-  }
-  ASSERT(skip < alloc);
-  if (skip > 0) {
-    Span* rest = Static::pageheap()->Split(span, skip);
-    Static::pageheap()->Delete(span);
-    span = rest;
   }
 
-  // Skip trailing portion that we do not need to return
-  const Length needed = tcmalloc::pages(size);
-  ASSERT(span->length >= needed);
-  if (span->length > needed) {
-    Span* trailer = Static::pageheap()->Split(span, needed);
-    Static::pageheap()->Delete(trailer);
-  }
   return SpanToMallocResult(span);
 }
 

src/tests/page_heap_test.cc

@@ -9,6 +9,7 @@
 
 #include <stdio.h>
 
+#include <limits>
 #include <memory>
 
 #include "page_heap.h"
@@ -53,7 +54,7 @@ static void TestPageHeap_Stats() {
   CheckStats(ph.get(), 256, 0, 0);
 
   // Split span 's1' into 's1', 's2'.  Delete 's2'
-  tcmalloc::Span* s2 = ph->Split(s1, 128);
+  tcmalloc::Span* s2 = ph->SplitForTest(s1, 128);
   ph->Delete(s2);
   CheckStats(ph.get(), 256, 128, 0);
 
@@ -101,6 +102,13 @@ static void TestPageHeap_Limit() {
 
   std::unique_ptr<tcmalloc::PageHeap> ph(new tcmalloc::PageHeap());
 
+  // Lets also test if huge number of pages is ooming properly
+  {
+    auto res = ph->New(std::numeric_limits<Length>::max());
+    CHECK_EQ(res, nullptr);
+    CHECK_EQ(errno, ENOMEM);
+  }
+
   CHECK_EQ(kMaxPages, 1 << (20 - kPageShift));
 
   // We do not know much is taken from the system for other purposes,