Improved iterator for persistent memory allocator.

base/metrics/persistent_memory_allocator.cc

 // Copyright (c) 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/metrics/persistent_memory_allocator.h"
 
 #include <assert.h>
 #include <algorithm>
 
 #include "base/files/memory_mapped_file.h"
@@ skipping 27 matching lines @@
 // TODO(bcwhite): When acceptable, consider moving flags to std::atomic<char>
 // types rather than combined bitfield.
 
 // Flags stored in the flags_ field of the SharedMetaData structure below.
 enum : int {
   kFlagCorrupt = 1 << 0,
   kFlagFull    = 1 << 1
 };
 
 bool CheckFlag(const volatile std::atomic<uint32_t>* flags, int flag) {
-  uint32_t loaded_flags = flags->load();
+  uint32_t loaded_flags = flags->load(std::memory_order_relaxed);
   return (loaded_flags & flag) != 0;
 }
 
 void SetFlag(volatile std::atomic<uint32_t>* flags, int flag) {
-  uint32_t loaded_flags = flags->load();
+  uint32_t loaded_flags = flags->load(std::memory_order_relaxed);
   for (;;) {
     uint32_t new_flags = (loaded_flags & ~flag) | flag;
     // In the failue case, actual "flags" value stored in loaded_flags.
     if (flags->compare_exchange_weak(loaded_flags, new_flags))
       break;
   }
 }
 
 }  // namespace
 
@@ skipping 48 matching lines @@
 
 // The "queue" block header is used to detect "last node" so that zero/null
 // can be used to indicate that it hasn't been added at all. It is part of
 // the SharedMetadata structure which itself is always located at offset zero.
 const PersistentMemoryAllocator::Reference
     PersistentMemoryAllocator::kReferenceQueue =
         offsetof(SharedMetadata, queue);
 const PersistentMemoryAllocator::Reference
     PersistentMemoryAllocator::kReferenceNull = 0;
 
+PersistentMemoryAllocator::Iterator::Iterator(
+    const PersistentMemoryAllocator* allocator)
+    : allocator_(allocator), last_record_(kReferenceQueue), record_count_(0) {}
+
+PersistentMemoryAllocator::Iterator::Iterator(
+    const PersistentMemoryAllocator* allocator,
+    Reference starting_after)
+    : allocator_(allocator), last_record_(starting_after), record_count_(0) {
+  // Ensure that the starting point is a valid, iterable block (meaning it can
+  // be read and has a non-zero "next" pointer).
+  const volatile BlockHeader* block =
+      allocator_->GetBlock(starting_after, 0, 0, false, false);
+  if (!block || block->next.load(std::memory_order_relaxed) == 0) {
+    NOTREACHED();
+    last_record_.store(kReferenceQueue, std::memory_order_release);
+  }
+}
+
+PersistentMemoryAllocator::Reference
+PersistentMemoryAllocator::Iterator::GetNext(uint32_t* type_return) {
+  // Make a copy of the existing count of found-records, acquiring all changes
+  // made to the allocator, notably "freeptr" (see comment in loop for why
+  // the load of that value cannot be moved above here) that occurred during
+  // any previous runs of this method, including those by parallel threads
+  // that interrupted it. It pairs with the Release at the end of this method.
+  //
+  // Otherwise, if the compiler were to arrange the two loads such that
+  // "count" was fetched _after_ "freeptr" then it would be possible for
+  // this thread to be interrupted between them and other threads perform
+  // multiple allocations, make-iterables, and iterations (with the included
+  // increment of |record_count_|) culminating in the check at the bottom
+  // mistakenly determining that a loop exists. Isn't this stuff fun?
+  uint32_t count = record_count_.load(std::memory_order_acquire);
+
+  Reference last = last_record_.load(std::memory_order_acquire);
+  Reference next;
+  while (true) {
+    const volatile BlockHeader* block =
+        allocator_->GetBlock(last, 0, 0, true, false);
+    if (!block)  // Invalid iterator state.
+      return kReferenceNull;
+
+    // The compiler and CPU can freely reorder all memory accesses on which
+    // there are no dependencies. It could, for example, move the load of
+    // "freeptr" to above this point because there are no explicit dependencies
+    // between it and "next". If it did, however, then another block could
+    // be queued after that but before the following load meaning there is
+    // one more queued block than the future "detect loop by having more
+    // blocks that could fit before freeptr" will allow.
+    //
+    // By "acquiring" the "next" value here, it's synchronized to the enqueue
+    // of the node which in turn is synchronized to the allocation (which sets
+    // freeptr). Thus, the scenario above cannot happen.
+    next = block->next.load(std::memory_order_acquire);
+    if (next == kReferenceQueue)  // No next allocation in queue.
+      return kReferenceNull;
+    block = allocator_->GetBlock(next, 0, 0, false, false);
+    if (!block) {  // Memory is corrupt.
+      allocator_->SetCorrupt();
+      return kReferenceNull;
+    }
+
+    // Update the "last_record" pointer to be the reference being returned.
+    // If it fails then another thread has already iterated past it so loop
+    // again. Failing will also load the existing value into "last" so there
+    // is no need to do another such load when the while-loop restarts. A
+    // "strong" compare-exchange is used because failing unnecessarily would
+    // mean repeating some fairly costly validations above.
+    if (last_record_.compare_exchange_strong(last, next)) {
+      *type_return = block->type_id;
+      break;
+    }
+  }
+
+  // Memory corruption could cause a loop in the list. Such must be detected
+  // so as to not cause an infinite loop in the caller. This is done by simply
+  // making sure it doesn't iterate more times than the absolute maximum
+  // number of allocations that could have been made. Callers are likely
+  // to loop multiple times before it is detected but at least it stops.
+  const uint32_t freeptr = std::min(
+      allocator_->shared_meta()->freeptr.load(std::memory_order_relaxed),
+      allocator_->mem_size_);
+  const uint32_t max_records =
+      freeptr / (sizeof(BlockHeader) + kAllocAlignment);
+  if (count > max_records) {
+    allocator_->SetCorrupt();
+    return kReferenceNull;
+  }
+
+  // Increment the count and release the changes made above. It pairs with
+  // the Acquire at the top of this method. Note that this operation is not
+  // strictly synchonized with fetching of the object to return, which would
+  // have to be done inside the loop and is somewhat complicated to achieve.
+  // It does not matter if it falls behind temporarily so long as it never
+  // gets ahead.
+  record_count_.fetch_add(1, std::memory_order_release);
+  return next;
+}
+
+PersistentMemoryAllocator::Reference
+PersistentMemoryAllocator::Iterator::GetNextOfType(uint32_t type_match) {
+  Reference ref;
+  uint32_t type_found;
+  while ((ref = GetNext(&type_found)) != 0) {
+    if (type_found == type_match)
+      return ref;
+  }
+  return kReferenceNull;
+}
 
 // static
 bool PersistentMemoryAllocator::IsMemoryAcceptable(const void* base,
                                                    size_t size,
                                                    size_t page_size,
                                                    bool readonly) {
   return ((base && reinterpret_cast<uintptr_t>(base) % kAllocAlignment == 0) &&
           (size >= sizeof(SharedMetadata) && size <= kSegmentMaxSize) &&
           (size >= kSegmentMinSize || readonly) &&
           (size % kAllocAlignment == 0 || readonly) &&
@@ skipping 38 matching lines @@
     }
 
     // This block is only executed when a completely new memory segment is
     // being initialized. It's unshared and single-threaded...
     volatile BlockHeader* const first_block =
         reinterpret_cast<volatile BlockHeader*>(mem_base_ +
                                                 sizeof(SharedMetadata));
     if (shared_meta()->cookie != 0 ||
         shared_meta()->size != 0 ||
         shared_meta()->version != 0 ||
-        shared_meta()->freeptr.load() != 0 ||
-        shared_meta()->flags.load() != 0 ||
+        shared_meta()->freeptr.load(std::memory_order_relaxed) != 0 ||
+        shared_meta()->flags.load(std::memory_order_relaxed) != 0 ||
         shared_meta()->id != 0 ||
         shared_meta()->name != 0 ||
         shared_meta()->tailptr != 0 ||
         shared_meta()->queue.cookie != 0 ||
-        shared_meta()->queue.next.load() != 0 ||
+        shared_meta()->queue.next.load(std::memory_order_relaxed) != 0 ||
         first_block->size != 0 ||
         first_block->cookie != 0 ||
         first_block->type_id != 0 ||
         first_block->next != 0) {
       // ...or something malicious has been playing with the metadata.
       NOTREACHED();
       SetCorrupt();
     }
 
     // This is still safe to do even if corruption has been detected.
     shared_meta()->cookie = kGlobalCookie;
     shared_meta()->size = mem_size_;
     shared_meta()->page_size = mem_page_;
     shared_meta()->version = kGlobalVersion;
     shared_meta()->id = id;
-    shared_meta()->freeptr.store(sizeof(SharedMetadata));
+    shared_meta()->freeptr.store(sizeof(SharedMetadata),
+                                 std::memory_order_release);
 
     // Set up the queue of iterable allocations.
     shared_meta()->queue.size = sizeof(BlockHeader);
     shared_meta()->queue.cookie = kBlockCookieQueue;
-    shared_meta()->queue.next.store(kReferenceQueue);
-    shared_meta()->tailptr.store(kReferenceQueue);
+    shared_meta()->queue.next.store(kReferenceQueue, std::memory_order_release);
+    shared_meta()->tailptr.store(kReferenceQueue, std::memory_order_release);
 
     // Allocate space for the name so other processes can learn it.
     if (!name.empty()) {
       const size_t name_length = name.length() + 1;
       shared_meta()->name = Allocate(name_length, 0);
       char* name_cstr = GetAsObject<char>(shared_meta()->name, 0);
       if (name_cstr)
         memcpy(name_cstr, name.data(), name.length());
     }
   } else {
     if (shared_meta()->size == 0 ||
         shared_meta()->version == 0 ||
-        shared_meta()->freeptr.load() == 0 ||
+        shared_meta()->freeptr.load(std::memory_order_relaxed) == 0 ||
         shared_meta()->tailptr == 0 ||
         shared_meta()->queue.cookie == 0 ||
-        shared_meta()->queue.next.load() == 0) {
+        shared_meta()->queue.next.load(std::memory_order_relaxed) == 0) {
       SetCorrupt();
     }
     if (!readonly) {
       // The allocator is attaching to a previously initialized segment of
       // memory. Make sure the embedded data matches what has been passed.
       if (shared_meta()->size != mem_size_ ||
           shared_meta()->page_size != mem_page_) {
         NOTREACHED();
         SetCorrupt();
       }
@@ skipping 39 matching lines @@
       "UMA.PersistentAllocator." + name_string + ".UsedPct", 1, 101, 21,
       HistogramBase::kUmaTargetedHistogramFlag);
 
   DCHECK(!allocs_histogram_);
   allocs_histogram_ = Histogram::FactoryGet(
       "UMA.PersistentAllocator." + name_string + ".Allocs", 1, 10000, 50,
       HistogramBase::kUmaTargetedHistogramFlag);
 }
 
 size_t PersistentMemoryAllocator::used() const {
-  return std::min(shared_meta()->freeptr.load(), mem_size_);
+  return std::min(shared_meta()->freeptr.load(std::memory_order_relaxed),
+                  mem_size_);
 }
 
 size_t PersistentMemoryAllocator::GetAllocSize(Reference ref) const {
   const volatile BlockHeader* const block = GetBlock(ref, 0, 0, false, false);
   if (!block)
     return 0;
   uint32_t size = block->size;
   // Header was verified by GetBlock() but a malicious actor could change
   // the value between there and here. Check it again.
   if (size <= sizeof(BlockHeader) || ref + size > mem_size_) {
@@ skipping 52 matching lines @@
     NOTREACHED();
     return kReferenceNull;
   }
 
   // Get the current start of unallocated memory. Other threads may
   // update this at any time and cause us to retry these operations.
   // This value should be treated as "const" to avoid confusion through
   // the code below but recognize that any failed compare-exchange operation
   // involving it will cause it to be loaded with a more recent value. The
   // code should either exit or restart the loop in that case.
-  /* const */ uint32_t freeptr = shared_meta()->freeptr.load();
+  /* const */ uint32_t freeptr =
+      shared_meta()->freeptr.load(std::memory_order_acquire);
 
   // Allocation is lockless so we do all our caculation and then, if saving
   // indicates a change has occurred since we started, scrap everything and
   // start over.
   for (;;) {
     if (IsCorrupt())
       return kReferenceNull;
 
     if (freeptr + size > mem_size_) {
       SetFlag(&shared_meta()->flags, kFlagFull);
@@ skipping 49 matching lines @@
 
     // Given that all memory was zeroed before ever being given to an instance
     // of this class and given that we only allocate in a monotomic fashion
     // going forward, it must be that the newly allocated block is completely
     // full of zeros. If we find anything in the block header that is NOT a
     // zero then something must have previously run amuck through memory,
     // writing beyond the allocated space and into unallocated space.
     if (block->size != 0 ||
         block->cookie != kBlockCookieFree ||
         block->type_id != 0 ||
-        block->next.load() != 0) {
+        block->next.load(std::memory_order_relaxed) != 0) {
       SetCorrupt();
       return kReferenceNull;
     }
 
     block->size = size;
     block->cookie = kBlockCookieAllocated;
     block->type_id = type_id;
     return freeptr;
   }
 }
 
 void PersistentMemoryAllocator::GetMemoryInfo(MemoryInfo* meminfo) const {
-  uint32_t remaining = std::max(mem_size_ - shared_meta()->freeptr.load(),
-                                (uint32_t)sizeof(BlockHeader));
+  uint32_t remaining = std::max(
+      mem_size_ - shared_meta()->freeptr.load(std::memory_order_relaxed),
+      (uint32_t)sizeof(BlockHeader));
   meminfo->total = mem_size_;
   meminfo->free = IsCorrupt() ? 0 : remaining - sizeof(BlockHeader);
 }
 
 void PersistentMemoryAllocator::MakeIterable(Reference ref) {
   DCHECK(!readonly_);
   if (IsCorrupt())
     return;
   volatile BlockHeader* block = GetBlock(ref, 0, 0, false, false);
   if (!block)  // invalid reference
@@ skipping 42 matching lines @@
       // check for crash/kill which means that this operation may also happen
       // even when the other thread is in perfect working order which is what
       // necessitates the CompareAndSwap above.
       shared_meta()->tailptr.compare_exchange_strong(tail, next,
                                                      std::memory_order_acq_rel,
                                                      std::memory_order_acquire);
     }
   }
 }
 
-void PersistentMemoryAllocator::CreateIterator(Iterator* state,
-                                               Reference starting_after) const {
-  if (starting_after) {
-    // Ensure that the starting point is a valid, iterable block.
-    const volatile BlockHeader* block =
-        GetBlock(starting_after, 0, 0, false, false);
-    if (!block || !block->next.load()) {
-      NOTREACHED();
-      starting_after = kReferenceQueue;
-    }
-  } else {
-    // A zero beginning is really the Queue reference.
-    starting_after = kReferenceQueue;
-  }
-
-  state->last = starting_after;
-  state->niter = 0;
-}
-
-PersistentMemoryAllocator::Reference PersistentMemoryAllocator::GetNextIterable(
-    Iterator* state,
-    uint32_t* type_id) const {
-  const volatile BlockHeader* block = GetBlock(state->last, 0, 0, true, false);
-  if (!block)  // invalid iterator state
-    return kReferenceNull;
-
-  // The compiler and CPU can freely reorder all memory accesses on which
-  // there are no dependencies. It could, for example, move the load of
-  // "freeptr" above this point because there are no explicit dependencies
-  // between it and "next". If it did, however, then another block could
-  // be queued after that but before the following load meaning there is
-  // one more queued block than the future "detect loop by having more
-  // blocks that could fit before freeptr" will allow.
-  //
-  // By "acquiring" the "next" value here, it's synchronized to the enqueue
-  // of the node which in turn is synchronized to the allocation (which sets
-  // freeptr). Thus, the scenario above cannot happen.
-  uint32_t next = block->next.load(std::memory_order_acquire);
-  block = GetBlock(next, 0, 0, false, false);
-  if (!block)  // no next allocation in queue
-    return kReferenceNull;
-
-  // Memory corruption could cause a loop in the list. We need to detect
-  // that so as to not cause an infinite loop in the caller. We do this
-  // simply by making sure we don't iterate more than the absolute maximum
-  // number of allocations that could have been made. Callers are likely
-  // to loop multiple times before it is detected but at least it stops.
-  uint32_t freeptr = std::min(
-      shared_meta()->freeptr.load(std::memory_order_acquire),
-      mem_size_);
-  if (state->niter > freeptr / (sizeof(BlockHeader) + kAllocAlignment)) {
-    SetCorrupt();
-    return kReferenceNull;
-  }
-
-  state->last = next;
-  state->niter++;
-  *type_id = block->type_id;
-
-  return next;
-}
-
 // The "corrupted" state is held both locally and globally (shared). The
 // shared flag can't be trusted since a malicious actor could overwrite it.
 // Because corruption can be detected during read-only operations such as
 // iteration, this method may be called by other "const" methods. In this
 // case, it's safe to discard the constness and modify the local flag and
 // maybe even the shared flag if the underlying data isn't actually read-only.
 void PersistentMemoryAllocator::SetCorrupt() const {
   LOG(ERROR) << "Corruption detected in shared-memory segment.";
-  const_cast<std::atomic<bool>*>(&corrupt_)->store(true);
+  const_cast<std::atomic<bool>*>(&corrupt_)->store(true,
+                                                   std::memory_order_relaxed);
   if (!readonly_) {
     SetFlag(const_cast<volatile std::atomic<uint32_t>*>(&shared_meta()->flags),
             kFlagCorrupt);
   }
 }
 
 bool PersistentMemoryAllocator::IsCorrupt() const {
-  if (corrupt_.load() || CheckFlag(&shared_meta()->flags, kFlagCorrupt)) {
+  if (corrupt_.load(std::memory_order_relaxed) ||
+      CheckFlag(&shared_meta()->flags, kFlagCorrupt)) {
     SetCorrupt();  // Make sure all indicators are set.
     return true;
   }
   return false;
 }
 
 bool PersistentMemoryAllocator::IsFull() const {
   return CheckFlag(&shared_meta()->flags, kFlagFull);
 }
 
@@ skipping 10 matching lines @@
   if (ref % kAllocAlignment != 0)
     return nullptr;
   if (ref < (queue_ok ? kReferenceQueue : sizeof(SharedMetadata)))
     return nullptr;
   size += sizeof(BlockHeader);
   if (ref + size > mem_size_)
     return nullptr;
 
   // Validation of referenced block-header.
   if (!free_ok) {
-    uint32_t freeptr = shared_meta()->freeptr.load();
+    uint32_t freeptr = std::min(
+        shared_meta()->freeptr.load(std::memory_order_relaxed), mem_size_);
     if (ref + size > freeptr)
       return nullptr;
     const volatile BlockHeader* const block =
         reinterpret_cast<volatile BlockHeader*>(mem_base_ + ref);
     if (block->size < size)
       return nullptr;
+    if (ref + block->size > freeptr)
+      return nullptr;
     if (ref != kReferenceQueue && block->cookie != kBlockCookieAllocated)
       return nullptr;
     if (type_id != 0 && block->type_id != type_id)
       return nullptr;
   }
 
   // Return pointer to block data.
   return reinterpret_cast<const volatile BlockHeader*>(mem_base_ + ref);
 }
 
@@ skipping 75 matching lines @@
 
 FilePersistentMemoryAllocator::~FilePersistentMemoryAllocator() {}
 
 // static
 bool FilePersistentMemoryAllocator::IsFileAcceptable(
     const MemoryMappedFile& file) {
   return IsMemoryAcceptable(file.data(), file.length(), 0, true);
 }
 
 }  // namespace base