Improved iterator for persistent memory allocator.

base/metrics/persistent_memory_allocator_unittest.cc

 // Copyright 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 <memory>
 
 #include "base/files/file.h"
 #include "base/files/file_util.h"
 #include "base/files/memory_mapped_file.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/memory/shared_memory.h"
 #include "base/metrics/histogram.h"
 #include "base/rand_util.h"
 #include "base/strings/safe_sprintf.h"
+#include "base/synchronization/condition_variable.h"
+#include "base/synchronization/lock.h"
 #include "base/threading/simple_thread.h"
 #include "testing/gmock/include/gmock/gmock.h"
 
 namespace {
 
 const uint32_t TEST_MEMORY_SIZE = 1 << 20;   // 1 MiB
 const uint32_t TEST_MEMORY_PAGE = 64 << 10;  // 64 KiB
 const uint32_t TEST_ID = 12345;
 const char TEST_NAME[] = "TestAllocator";
 
@@ skipping 35 matching lines @@
         mem_segment_.get(), TEST_MEMORY_SIZE, TEST_MEMORY_PAGE,
         TEST_ID, TEST_NAME, false));
     allocator_->CreateTrackingHistograms(allocator_->Name());
   }
 
   void TearDown() override {
     allocator_.reset();
   }
 
   unsigned CountIterables() {
-    PersistentMemoryAllocator::Iterator iter;
+    PersistentMemoryAllocator::Iterator iter(allocator_.get());
     uint32_t type;
     unsigned count = 0;
-    for (allocator_->CreateIterator(&iter);
-         allocator_->GetNextIterable(&iter, &type) != 0;) {
-      count++;
+    while (iter.GetNext(&type) != 0) {
+      ++count;
     }
     return count;
   }
 
  protected:
   std::unique_ptr<char[]> mem_segment_;
   std::unique_ptr<PersistentMemoryAllocator> allocator_;
 };
 
 TEST_F(PersistentMemoryAllocatorTest, AllocateAndIterate) {
@@ skipping 20 matching lines @@
   EXPECT_EQ(nullptr, allocator_->GetAsObject<TestObject2>(block1, 1));
   EXPECT_LE(sizeof(TestObject1), allocator_->GetAllocSize(block1));
   EXPECT_GT(sizeof(TestObject1) + kAllocAlignment,
             allocator_->GetAllocSize(block1));
   PersistentMemoryAllocator::MemoryInfo meminfo1;
   allocator_->GetMemoryInfo(&meminfo1);
   EXPECT_EQ(meminfo0.total, meminfo1.total);
   EXPECT_GT(meminfo0.free, meminfo1.free);
 
   // Ensure that the test-object can be made iterable.
-  PersistentMemoryAllocator::Iterator iter;
+  PersistentMemoryAllocator::Iterator iter1a(allocator_.get());
   uint32_t type;
-  allocator_->CreateIterator(&iter);
-  EXPECT_EQ(0U, allocator_->GetNextIterable(&iter, &type));
+  EXPECT_EQ(0U, iter1a.GetNext(&type));
   allocator_->MakeIterable(block1);
-  EXPECT_EQ(block1, allocator_->GetNextIterable(&iter, &type));
+  EXPECT_EQ(block1, iter1a.GetNext(&type));
   EXPECT_EQ(1U, type);
-  EXPECT_EQ(0U, allocator_->GetNextIterable(&iter, &type));
+  EXPECT_EQ(0U, iter1a.GetNext(&type));
 
   // Create second test-object and ensure everything is good and it cannot
   // be confused with test-object of another type.
   Reference block2 = allocator_->Allocate(sizeof(TestObject2), 2);
   EXPECT_NE(0U, block2);
   EXPECT_NE(nullptr, allocator_->GetAsObject<TestObject2>(block2, 2));
   EXPECT_EQ(nullptr, allocator_->GetAsObject<TestObject2>(block2, 1));
   EXPECT_LE(sizeof(TestObject2), allocator_->GetAllocSize(block2));
   EXPECT_GT(sizeof(TestObject2) + kAllocAlignment,
             allocator_->GetAllocSize(block2));
   PersistentMemoryAllocator::MemoryInfo meminfo2;
   allocator_->GetMemoryInfo(&meminfo2);
   EXPECT_EQ(meminfo1.total, meminfo2.total);
   EXPECT_GT(meminfo1.free, meminfo2.free);
 
   // Ensure that second test-object can also be made iterable.
   allocator_->MakeIterable(block2);
-  EXPECT_EQ(block2, allocator_->GetNextIterable(&iter, &type));
+  EXPECT_EQ(block2, iter1a.GetNext(&type));
   EXPECT_EQ(2U, type);
-  EXPECT_EQ(0U, allocator_->GetNextIterable(&iter, &type));
+  EXPECT_EQ(0U, iter1a.GetNext(&type));
 
   // Check that iteration can begin after an arbitrary location.
-  allocator_->CreateIterator(&iter, block1);
-  EXPECT_EQ(block2, allocator_->GetNextIterable(&iter, &type));
-  EXPECT_EQ(0U, allocator_->GetNextIterable(&iter, &type));
+  PersistentMemoryAllocator::Iterator iter1b(allocator_.get(), block1);
+  EXPECT_EQ(block2, iter1b.GetNext(&type));
+  EXPECT_EQ(0U, iter1b.GetNext(&type));
 
   // Ensure nothing has gone noticably wrong.
   EXPECT_FALSE(allocator_->IsFull());
   EXPECT_FALSE(allocator_->IsCorrupt());
 
   // Check the internal histogram record of used memory.
   allocator_->UpdateTrackingHistograms();
   std::unique_ptr<HistogramSamples> used_samples(
       allocator_->used_histogram_->SnapshotSamples());
   EXPECT_TRUE(used_samples);
@@ skipping 24 matching lines @@
   // Create second allocator (read/write) using the same memory segment.
   std::unique_ptr<PersistentMemoryAllocator> allocator2(
       new PersistentMemoryAllocator(mem_segment_.get(), TEST_MEMORY_SIZE,
                                     TEST_MEMORY_PAGE, 0, "", false));
   EXPECT_EQ(TEST_ID, allocator2->Id());
   EXPECT_FALSE(allocator2->used_histogram_);
   EXPECT_FALSE(allocator2->allocs_histogram_);
   EXPECT_NE(allocator2->allocs_histogram_, allocator_->allocs_histogram_);
 
   // Ensure that iteration and access through second allocator works.
-  allocator2->CreateIterator(&iter);
-  EXPECT_EQ(block1, allocator2->GetNextIterable(&iter, &type));
-  EXPECT_EQ(block2, allocator2->GetNextIterable(&iter, &type));
-  EXPECT_EQ(0U, allocator2->GetNextIterable(&iter, &type));
+  PersistentMemoryAllocator::Iterator iter2(allocator2.get());
+  EXPECT_EQ(block1, iter2.GetNext(&type));
+  EXPECT_EQ(block2, iter2.GetNext(&type));
+  EXPECT_EQ(0U, iter2.GetNext(&type));
   EXPECT_NE(nullptr, allocator2->GetAsObject<TestObject1>(block1, 1));
   EXPECT_NE(nullptr, allocator2->GetAsObject<TestObject2>(block2, 2));
 
   // Create a third allocator (read-only) using the same memory segment.
   std::unique_ptr<const PersistentMemoryAllocator> allocator3(
       new PersistentMemoryAllocator(mem_segment_.get(), TEST_MEMORY_SIZE,
                                     TEST_MEMORY_PAGE, 0, "", true));
   EXPECT_EQ(TEST_ID, allocator3->Id());
   EXPECT_FALSE(allocator3->used_histogram_);
   EXPECT_FALSE(allocator3->allocs_histogram_);
 
   // Ensure that iteration and access through third allocator works.
-  allocator3->CreateIterator(&iter);
-  EXPECT_EQ(block1, allocator3->GetNextIterable(&iter, &type));
-  EXPECT_EQ(block2, allocator3->GetNextIterable(&iter, &type));
-  EXPECT_EQ(0U, allocator3->GetNextIterable(&iter, &type));
+  PersistentMemoryAllocator::Iterator iter3(allocator3.get());
+  EXPECT_EQ(block1, iter3.GetNext(&type));
+  EXPECT_EQ(block2, iter3.GetNext(&type));
+  EXPECT_EQ(0U, iter3.GetNext(&type));
   EXPECT_NE(nullptr, allocator3->GetAsObject<TestObject1>(block1, 1));
   EXPECT_NE(nullptr, allocator3->GetAsObject<TestObject2>(block2, 2));
+
+  // Ensure that GetNextOfType works.
+  PersistentMemoryAllocator::Iterator iter1c(allocator_.get());
+  EXPECT_EQ(block2, iter1c.GetNextOfType(2));
+  EXPECT_EQ(0U, iter1c.GetNextOfType(2));
 }
 
 TEST_F(PersistentMemoryAllocatorTest, PageTest) {
   // This allocation will go into the first memory page.
   Reference block1 = allocator_->Allocate(TEST_MEMORY_PAGE / 2, 1);
   EXPECT_LT(0U, block1);
   EXPECT_GT(TEST_MEMORY_PAGE, block1);
 
   // This allocation won't fit in same page as previous block.
   Reference block2 =
@@ skipping 71 matching lines @@
   t4.Join();
   t5.Join();
 
   EXPECT_FALSE(allocator_->IsCorrupt());
   EXPECT_TRUE(allocator_->IsFull());
   EXPECT_EQ(CountIterables(),
             t1.iterable() + t2.iterable() + t3.iterable() + t4.iterable() +
             t5.iterable());
 }
 
+// A simple thread that counts objects by iterating through an allocator.
+class CounterThread : public SimpleThread {
+ public:
+  CounterThread(const std::string& name,
+                PersistentMemoryAllocator::Iterator* iterator,
+                Lock* lock,
+                ConditionVariable* condition)
+      : SimpleThread(name, Options()),
+        iterator_(iterator),
+        lock_(lock),
+        condition_(condition),
+        count_(0) {}
+
+  void Run() override {
+    // Wait so all threads can start at approximately the same time.
+    // Best performance comes from releasing a single worker which then
+    // releases the next, etc., etc.
+    {
+      AutoLock autolock(*lock_);
+      condition_->Wait();
+      condition_->Signal();
+    }
+
+    uint32_t type;
+    while (iterator_->GetNext(&type) != 0) {
+      ++count_;
+    }
+  }
+
+  unsigned count() { return count_; }
+
+ private:
+  PersistentMemoryAllocator::Iterator* iterator_;
+  Lock* lock_;
+  ConditionVariable* condition_;
+  unsigned count_;
+};
+
+// Ensure that parallel iteration returns the same number of objects as
+// single-threaded iteration.
+TEST_F(PersistentMemoryAllocatorTest, IteratorParallelismTest) {
+  // Fill the memory segment with random allocations.
+  unsigned iterable_count = 0;
+  for (;;) {
+    uint32_t size = RandInt(1, 99);
+    uint32_t type = RandInt(100, 999);
+    Reference block = allocator_->Allocate(size, type);
+    if (!block)
+      break;
+    allocator_->MakeIterable(block);
+    ++iterable_count;
+  }
+  EXPECT_FALSE(allocator_->IsCorrupt());
+  EXPECT_TRUE(allocator_->IsFull());
+  EXPECT_EQ(iterable_count, CountIterables());
+
+  PersistentMemoryAllocator::Iterator iter(allocator_.get());
+  Lock lock;
+  ConditionVariable condition(&lock);
+
+  CounterThread t1("t1", &iter, &lock, &condition);
+  CounterThread t2("t2", &iter, &lock, &condition);
+  CounterThread t3("t3", &iter, &lock, &condition);
+  CounterThread t4("t4", &iter, &lock, &condition);
+  CounterThread t5("t5", &iter, &lock, &condition);
+
+  t1.Start();
+  t2.Start();
+  t3.Start();
+  t4.Start();
+  t5.Start();
+
+  // This will release all the waiting threads.
+  condition.Signal();
+
+  t1.Join();
+  t2.Join();
+  t3.Join();
+  t4.Join();
+  t5.Join();
+
+  EXPECT_EQ(iterable_count,
+            t1.count() + t2.count() + t3.count() + t4.count() + t5.count());
+
+#if 0
+  // These ensure that the threads don't run sequentially. It shouldn't be
+  // enabled in general because it could lead to a flaky test if it happens
+  // simply by chance but it is useful during development to ensure that the
+  // test is working correctly.
+  EXPECT_NE(iterable_count, t1.count());
+  EXPECT_NE(iterable_count, t2.count());
+  EXPECT_NE(iterable_count, t3.count());
+  EXPECT_NE(iterable_count, t4.count());
+  EXPECT_NE(iterable_count, t5.count());
+#endif
+}
+
 // This test doesn't verify anything other than it doesn't crash. Its goal
 // is to find coding errors that aren't otherwise tested for, much like a
 // "fuzzer" would.
 // This test is suppsoed to fail on TSAN bot (crbug.com/579867).
 #if defined(THREAD_SANITIZER)
 #define MAYBE_CorruptionTest DISABLED_CorruptionTest
 #else
 #define MAYBE_CorruptionTest CorruptionTest
 #endif
 TEST_F(PersistentMemoryAllocatorTest, MAYBE_CorruptionTest) {
@@ skipping 102 matching lines @@
   std::unique_ptr<SharedMemory> shmem2(new SharedMemory(shared_handle,
                                                         /*readonly=*/true));
   ASSERT_TRUE(shmem2->Map(TEST_MEMORY_SIZE));
 
   SharedPersistentMemoryAllocator shalloc2(std::move(shmem2), 0, "", true);
   EXPECT_TRUE(shalloc2.IsReadonly());
   EXPECT_EQ(TEST_ID, shalloc2.Id());
   EXPECT_FALSE(shalloc2.IsFull());
   EXPECT_FALSE(shalloc2.IsCorrupt());
 
-  PersistentMemoryAllocator::Iterator iter2;
+  PersistentMemoryAllocator::Iterator iter2(&shalloc2);
   uint32_t type;
-  shalloc2.CreateIterator(&iter2);
-  EXPECT_EQ(r123, shalloc2.GetNextIterable(&iter2, &type));
-  EXPECT_EQ(r789, shalloc2.GetNextIterable(&iter2, &type));
-  EXPECT_EQ(0U, shalloc2.GetNextIterable(&iter2, &type));
+  EXPECT_EQ(r123, iter2.GetNext(&type));
+  EXPECT_EQ(r789, iter2.GetNext(&type));
+  EXPECT_EQ(0U, iter2.GetNext(&type));
 
   EXPECT_EQ(123U, shalloc2.GetType(r123));
   EXPECT_EQ(654U, shalloc2.GetType(r456));
   EXPECT_EQ(789U, shalloc2.GetType(r789));
 
   PersistentMemoryAllocator::MemoryInfo meminfo2;
   shalloc2.GetMemoryInfo(&meminfo2);
   EXPECT_EQ(meminfo1.total, meminfo2.total);
   EXPECT_EQ(meminfo1.free, meminfo2.free);
 
   // Read/write test.
   std::unique_ptr<SharedMemory> shmem3(new SharedMemory(shared_handle,
                                                         /*readonly=*/false));
   ASSERT_TRUE(shmem3->Map(TEST_MEMORY_SIZE));
 
   SharedPersistentMemoryAllocator shalloc3(std::move(shmem3), 0, "", false);
   EXPECT_FALSE(shalloc3.IsReadonly());
   EXPECT_EQ(TEST_ID, shalloc3.Id());
   EXPECT_FALSE(shalloc3.IsFull());
   EXPECT_FALSE(shalloc3.IsCorrupt());
 
-  PersistentMemoryAllocator::Iterator iter3;
-  shalloc3.CreateIterator(&iter3);
-  EXPECT_EQ(r123, shalloc3.GetNextIterable(&iter3, &type));
-  EXPECT_EQ(r789, shalloc3.GetNextIterable(&iter3, &type));
-  EXPECT_EQ(0U, shalloc3.GetNextIterable(&iter3, &type));
+  PersistentMemoryAllocator::Iterator iter3(&shalloc3);
+  EXPECT_EQ(r123, iter3.GetNext(&type));
+  EXPECT_EQ(r789, iter3.GetNext(&type));
+  EXPECT_EQ(0U, iter3.GetNext(&type));
 
   EXPECT_EQ(123U, shalloc3.GetType(r123));
   EXPECT_EQ(654U, shalloc3.GetType(r456));
   EXPECT_EQ(789U, shalloc3.GetType(r789));
 
   PersistentMemoryAllocator::MemoryInfo meminfo3;
   shalloc3.GetMemoryInfo(&meminfo3);
   EXPECT_EQ(meminfo1.total, meminfo3.total);
   EXPECT_EQ(meminfo1.free, meminfo3.free);
 
   // Interconnectivity test.
   Reference obj = shalloc3.Allocate(42, 42);
   ASSERT_TRUE(obj);
   shalloc3.MakeIterable(obj);
-  EXPECT_EQ(obj, shalloc2.GetNextIterable(&iter2, &type));
+  EXPECT_EQ(obj, iter2.GetNext(&type));
   EXPECT_EQ(42U, type);
 }
 
 
 //----- FilePersistentMemoryAllocator ------------------------------------------
 
 TEST(FilePersistentMemoryAllocatorTest, CreationTest) {
   ScopedTempDir temp_dir;
   ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
   FilePath file_path = temp_dir.path().AppendASCII("persistent_memory");
@@ skipping 23 matching lines @@
   EXPECT_TRUE(mmfile->IsValid());
   const size_t mmlength = mmfile->length();
   EXPECT_GE(meminfo1.total, mmlength);
 
   FilePersistentMemoryAllocator file(std::move(mmfile), 0, "");
   EXPECT_TRUE(file.IsReadonly());
   EXPECT_EQ(TEST_ID, file.Id());
   EXPECT_FALSE(file.IsFull());
   EXPECT_FALSE(file.IsCorrupt());
 
-  PersistentMemoryAllocator::Iterator iter;
+  PersistentMemoryAllocator::Iterator iter(&file);
   uint32_t type;
-  file.CreateIterator(&iter);
-  EXPECT_EQ(r123, file.GetNextIterable(&iter, &type));
-  EXPECT_EQ(r789, file.GetNextIterable(&iter, &type));
-  EXPECT_EQ(0U, file.GetNextIterable(&iter, &type));
+  EXPECT_EQ(r123, iter.GetNext(&type));
+  EXPECT_EQ(r789, iter.GetNext(&type));
+  EXPECT_EQ(0U, iter.GetNext(&type));
 
   EXPECT_EQ(123U, file.GetType(r123));
   EXPECT_EQ(654U, file.GetType(r456));
   EXPECT_EQ(789U, file.GetType(r789));
 
   PersistentMemoryAllocator::MemoryInfo meminfo2;
   file.GetMemoryInfo(&meminfo2);
   EXPECT_GE(meminfo1.total, meminfo2.total);
   EXPECT_GE(meminfo1.free, meminfo2.free);
   EXPECT_EQ(mmlength, meminfo2.total);
   EXPECT_EQ(0U, meminfo2.free);
 }
 
 TEST(FilePersistentMemoryAllocatorTest, AcceptableTest) {
   ScopedTempDir temp_dir;
   ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
   FilePath file_path_base = temp_dir.path().AppendASCII("persistent_memory_");
 
   LocalPersistentMemoryAllocator local(TEST_MEMORY_SIZE, TEST_ID, "");
-  local.Allocate(1, 1);
-  local.Allocate(11, 11);
+  local.MakeIterable(local.Allocate(1, 1));
+  local.MakeIterable(local.Allocate(11, 11));
   const size_t minsize = local.used();
   std::unique_ptr<char[]> garbage(new char[minsize]);
   RandBytes(garbage.get(), minsize);
 
   std::unique_ptr<MemoryMappedFile> mmfile;
   char filename[100];
   for (size_t filesize = minsize; filesize > 0; --filesize) {
     strings::SafeSPrintf(filename, "memory_%d_A", filesize);
     FilePath file_path = temp_dir.path().AppendASCII(filename);
     ASSERT_FALSE(PathExists(file_path));
     {
       File writer(file_path, File::FLAG_CREATE | File::FLAG_WRITE);
       ASSERT_TRUE(writer.IsValid());
       writer.Write(0, (const char*)local.data(), filesize);
     }
     ASSERT_TRUE(PathExists(file_path));
 
     mmfile.reset(new MemoryMappedFile());
     mmfile->Initialize(file_path);
     EXPECT_EQ(filesize, mmfile->length());
     if (FilePersistentMemoryAllocator::IsFileAcceptable(*mmfile)) {
       // Make sure construction doesn't crash. It will, however, cause
       // error messages warning about about a corrupted memory segment.
       FilePersistentMemoryAllocator allocator(std::move(mmfile), 0, "");
       // Also make sure that iteration doesn't crash.
-      PersistentMemoryAllocator::Iterator iter;
-      allocator.CreateIterator(&iter);
-      for (;;) {
-        Reference ref = allocator.GetNextIterable(&iter, 0);
-        if (!ref)
-          break;
+      PersistentMemoryAllocator::Iterator iter(&allocator);
+      uint32_t type_id;
+      Reference ref;
+      while ((ref = iter.GetNext(&type_id)) != 0) {
         const char* data = allocator.GetAsObject<char>(ref, 0);
         uint32_t type = allocator.GetType(ref);
         size_t size = allocator.GetAllocSize(ref);
         // Ensure compiler can't optimize-out above variables.
         (void)data;
         (void)type;
         (void)size;
-        // Ensure that corruption-detected flag gets properly set.
-        EXPECT_EQ(filesize != minsize, allocator.IsCorrupt());
       }
+      // Ensure that short files are detected as corrupt and full files are not.
+      EXPECT_EQ(filesize != minsize, allocator.IsCorrupt());
     } else {
       // For filesize >= minsize, the file must be acceptable. This
       // else clause (file-not-acceptable) should be reached only if
       // filesize < minsize.
       EXPECT_LT(filesize, minsize);
     }
 
     strings::SafeSPrintf(filename, "memory_%d_B", filesize);
     file_path = temp_dir.path().AppendASCII(filename);
     ASSERT_FALSE(PathExists(file_path));
@@ skipping 15 matching lines @@
     } else {
       // For filesize >= minsize, the file must be acceptable. This
       // else clause (file-not-acceptable) should be reached only if
       // filesize < minsize.
       EXPECT_GT(minsize, filesize);
     }
   }
 }
 
 }  // namespace base