diff --git a/src/support/arena.h b/src/support/arena.h index 79036326f632..6cf4a2818031 100644 --- a/src/support/arena.h +++ b/src/support/arena.h @@ -43,13 +43,6 @@ namespace tvm { namespace support { -namespace { -template // For lvalues (T is T&), -T&& forward(T&& param) { // take/return lvalue refs. - return static_cast(param); // For rvalues (T is T), -} // take/return rvalue refs. -} // namespace - /*! * \brief An arena page header. */ @@ -64,6 +57,21 @@ struct ArenaPageHeader { size_t offset; }; +/*! + * \brief A pending destructor call for one arena-allocated, non-trivially + * destructible object. Forms a singly-linked list of all such objects so + * their destructors can be invoked before the arena's pages are freed or + * recycled. + */ +struct ArenaDeleter { + /*! \brief The next pending destructor call. */ + ArenaDeleter* next; + /*! \brief The object to destroy. */ + void* obj; + /*! \brief Function that destroys obj (a type-erased call to ~T()). */ + void (*dtor)(void*); +}; + /*! * \brief Arena allocator that allocates memory from continuous * chunk and frees them all only during destruction. @@ -83,11 +91,13 @@ class GenericArena { /*! \brief Free all pages. */ void FreeAll() { + RunDeleters(); FreePageList(&head_); FreePageList(&free_list_); } /*! \brief Recycle all the pages in the arena */ void RecycleAll() { + RunDeleters(); // put all the current list to the free list. tail_->next = free_list_; // allocate the first in the free list to head @@ -115,18 +125,39 @@ class GenericArena { * \tparam Args Arguments to the constructor. * * \return The allocated object. - * \note The type T must be simple type, or only contain - * memory allocated from the same arena. - * Otherwise the destructor needs to be called explicitly. + * \note If T is not trivially destructible, its destructor is recorded and + * invoked automatically when the arena's pages are freed or recycled. */ template T* make(Args&&... args) { T* ptr = allocate_(); - new (ptr) T(forward(args)...); + new (ptr) T(std::forward(args)...); + if constexpr (!std::is_trivially_destructible::value) { + RegisterDeleter(ptr); + } return ptr; } private: + /*! + * \brief Record ptr's destructor to be called by RunDeleters, before the + * arena's pages are freed or recycled. + */ + template + void RegisterDeleter(T* ptr) { + ArenaDeleter* node = allocate_(); + node->obj = ptr; + node->dtor = [](void* p) { static_cast(p)->~T(); }; + node->next = deleters_; + deleters_ = node; + } + /*! \brief Invoke and clear all pending destructor calls registered so far. */ + void RunDeleters() { + for (ArenaDeleter* d = deleters_; d != nullptr; d = d->next) { + d->dtor(d->obj); + } + deleters_ = nullptr; + } /*! \brief internal page allocator. */ PageAllocator alloc_; /* \brief The head of the allocated list. */ @@ -135,6 +166,8 @@ class GenericArena { ArenaPageHeader* tail_{nullptr}; /* \brief List of free pages. */ ArenaPageHeader* free_list_{nullptr}; + /*! \brief Pending destructor calls for non-trivially-destructible objects. */ + ArenaDeleter* deleters_{nullptr}; /*! * \brief Align ptr by upper bound. * \param offset The offset value. diff --git a/tests/cpp/support_test.cc b/tests/cpp/support_test.cc index 87f14dce02ec..243b3f848cc4 100644 --- a/tests/cpp/support_test.cc +++ b/tests/cpp/support_test.cc @@ -20,6 +20,9 @@ #include #include +#include + +#include "../../src/support/arena.h" #include "../../src/support/utils.h" namespace tvm { @@ -43,5 +46,69 @@ TEST(StartsWithTests, Basic) { EXPECT_FALSE(::tvm::support::StartsWith("abc", "abcd")); } +namespace { +// A non-trivially-destructible type: destructing it has an observable +// side effect (incrementing a counter), unlike a plain-data struct. +struct DtorCounter { + explicit DtorCounter(int* counter) : counter(counter) {} + ~DtorCounter() { (*counter)++; } + int* counter; +}; +} // namespace + +TEST(ArenaTests, MakeRunsDestructorOnFreeAll) { + int destroyed = 0; + { + support::Arena arena; + for (int i = 0; i < 8; ++i) { + arena.make(&destroyed); + } + EXPECT_EQ(destroyed, 0); + // FreeAll() may be called directly (not only via ~Arena()), e.g. by + // MinRPCServer. It must run pending destructors itself rather than + // relying on the caller to do so, or on ~Arena() running afterwards. + arena.FreeAll(); + EXPECT_EQ(destroyed, 8); + } + // ~Arena() must not re-run the same destructors after an explicit + // FreeAll(), and must not touch the now-freed ArenaDeleter bookkeeping. + EXPECT_EQ(destroyed, 8); +} + +TEST(ArenaTests, MakeRunsDestructorOnRecycleAll) { + int destroyed = 0; + support::Arena arena; + arena.make(&destroyed); + arena.RecycleAll(); + EXPECT_EQ(destroyed, 1); + + arena.make(&destroyed); + EXPECT_EQ(destroyed, 1); +} + +TEST(ArenaTests, TrivialTypeUnaffected) { + support::Arena arena; + int* x = arena.make(42); + EXPECT_EQ(*x, 42); + // No crash/UB freeing an arena that only ever held trivially + // destructible objects. + arena.FreeAll(); +} + +namespace { +struct MoveOnlyHolder { + explicit MoveOnlyHolder(std::unique_ptr value) : value(std::move(value)) {} + std::unique_ptr value; +}; +} // namespace + +TEST(ArenaTests, MakeAcceptsMoveOnlyArgument) { + support::Arena arena; + // Regression test: make(std::make_unique<...>(...)) must not be ambiguous with std::forward + // via ADL + auto* holder = arena.make(std::make_unique(7)); + EXPECT_EQ(*holder->value, 7); +} + } // namespace test } // namespace tvm