Browse Source

Introduce TMemoryPool::TBookmark

Данный PR добавляет класс TMemoryPool::TBookmark, запомниющий текущее состояние пула при создании и восстанавливающий это состояние в деструкторе.
В случае, если в пуле успел смениться текущий чанк, то очищается только последний чанк.
commit_hash:cfeb77912864d25d5d8c34739915c504e289770a
ppavel96 1 month ago
parent
commit
1fea87d262
2 changed files with 61 additions and 0 deletions
  1. 23 0
      util/memory/pool.h
  2. 38 0
      util/memory/pool_ut.cpp

+ 23 - 0
util/memory/pool.h

@@ -131,6 +131,29 @@ public:
         }
     };
 
+    // When a bookmark is destroyed, the memory pool returns to the state when the bookmark was created.
+    class TBookmark {
+    public:
+        inline TBookmark(TMemoryPool& memoryPool)
+            : OwnerPoolRef_(memoryPool)
+            , BookmarkChunk_(memoryPool.Current_)
+            , BookmarkChunkDataSize_(memoryPool.Current_->DataSize())
+        {
+        }
+
+        inline ~TBookmark() {
+            OwnerPoolRef_.Current_->ResetChunk();
+            if (OwnerPoolRef_.Current_ == BookmarkChunk_) {
+                Y_UNUSED(BookmarkChunk_->Allocate(BookmarkChunkDataSize_));
+            }
+        }
+
+    private:
+        TMemoryPool& OwnerPoolRef_;
+        TMemoryPool::TChunk* BookmarkChunk_;
+        size_t BookmarkChunkDataSize_;
+    };
+
     inline TMemoryPool(size_t initial, IGrowPolicy* grow = TExpGrow::Instance(), IAllocator* alloc = TDefaultAllocator::Instance(), const TOptions& options = TOptions())
         : Current_(&Empty_)
         , BlockSize_(initial)

+ 38 - 0
util/memory/pool_ut.cpp

@@ -84,6 +84,7 @@ class TMemPoolTest: public TTestBase {
     UNIT_TEST(TestLargeStartingAlign)
     UNIT_TEST(TestMoveAlloc)
     UNIT_TEST(TestRoundUpToNextPowerOfTwoOption)
+    UNIT_TEST(TestMemoryPoolBookmark)
     UNIT_TEST_SUITE_END();
 
 private:
@@ -280,6 +281,43 @@ private:
         pool.Allocate(1);
         UNIT_ASSERT_VALUES_EQUAL(2 * EXPECTED_ALLOCATION_SIZE, allocator.GetSize());
     }
+
+    void TestMemoryPoolBookmark() {
+        TCheckedAllocator alloc;
+
+        {
+            TMemoryPool pool(200U, TMemoryPool::TExpGrow::Instance(), &alloc);
+            ui64* someData = pool.Allocate<ui64>();
+            static const ui64 TESTING{0x123456789ABCDEF};
+            *someData = TESTING;
+
+            const auto ma = pool.MemoryAllocated();
+            const auto chunkOverhead = ma - sizeof(ui64);
+            const auto firstChunkTotal = pool.MemoryWaste() + ma;
+
+            // Allocate some memory in pool but not enough to need new chunks:
+            {
+                TMemoryPool::TBookmark bookmarkA(pool);
+                for (size_t i = 0U; i != 10; ++i) {
+                    UNIT_ASSERT(pool.Allocate<ui64>());
+                }
+            }
+            UNIT_ASSERT_VALUES_EQUAL(pool.MemoryAllocated(), ma);
+            UNIT_ASSERT_VALUES_EQUAL(*someData, TESTING);
+
+            // Allocate some memory in pool enough to need a new single chunk:
+            {
+                TMemoryPool::TBookmark bookmarkB(pool);
+                for (size_t i = 0U; i != 50; ++i) {
+                    UNIT_ASSERT(pool.Allocate<ui64>());
+                }
+            }
+            UNIT_ASSERT_VALUES_EQUAL(pool.MemoryAllocated(), firstChunkTotal + chunkOverhead); // The last (second) chunk is completely free
+            UNIT_ASSERT_VALUES_EQUAL(*someData, TESTING);
+        }
+
+        alloc.CheckAtEnd();
+    }
 };
 
 UNIT_TEST_SUITE_REGISTRATION(TMemPoolTest);