[umf] add out of memory test

Author: Damian Duy

test/unified_malloc_framework/common/pool.hpp

@@ -131,9 +131,7 @@ struct proxy_pool : public pool_base {
 
         memset(ptr, 0, num * size);
 
-        if (ptr) {
-            EXPECT_EQ_NOEXCEPT(ret, UMF_RESULT_SUCCESS);
-        }
+        umf::getPoolLastStatusRef<proxy_pool>() = ret;
         return ptr;
     }
     void *realloc(void *ptr, size_t size) noexcept {
@@ -145,9 +143,7 @@ struct proxy_pool : public pool_base {
     void *aligned_malloc(size_t size, size_t alignment) noexcept {
         void *ptr;
         auto ret = umfMemoryProviderAlloc(provider, size, alignment, &ptr);
-        if (ptr) {
-            EXPECT_EQ_NOEXCEPT(ret, UMF_RESULT_SUCCESS);
-        }
+        umf::getPoolLastStatusRef<proxy_pool>() = ret;
         return ptr;
     }
     size_t malloc_usable_size(void *ptr) noexcept {
@@ -156,7 +152,6 @@ struct proxy_pool : public pool_base {
     }
     enum umf_result_t free(void *ptr) noexcept {
         auto ret = umfMemoryProviderFree(provider, ptr, 0);
-        EXPECT_EQ_NOEXCEPT(ret, UMF_RESULT_SUCCESS);
         return ret;
     }
     enum umf_result_t get_last_allocation_error() {

test/unified_malloc_framework/common/provider.hpp

@@ -76,6 +76,28 @@ struct provider_malloc : public provider_base {
     const char *get_name() noexcept { return "malloc"; }
 };
 
+struct provider_mock_out_of_mem : public provider_base {
+    provider_malloc helper_prov;
+    int allocNum = 0;
+    umf_result_t initialize(int allocNum) noexcept {
+        this->allocNum = allocNum;
+        return UMF_RESULT_SUCCESS;
+    }
+    enum umf_result_t alloc(size_t size, size_t align, void **ptr) noexcept {
+        if (allocNum <= 0) {
+            *ptr = nullptr;
+            return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        }
+        allocNum--;
+
+        return helper_prov.alloc(size, align, ptr);
+    }
+    enum umf_result_t free(void *ptr, size_t size) noexcept {
+        return helper_prov.free(ptr, size);
+    }
+    const char *get_name() noexcept { return "mock_out_of_mem"; }
+};
+
 } // namespace umf_test
 
 #endif /* UMF_TEST_PROVIDER_HPP */

test/unified_malloc_framework/memoryPool.hpp

@@ -3,7 +3,9 @@
 // See LICENSE.TXT
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include "disjoint_pool.hpp"
 #include "pool.hpp"
+#include "provider.hpp"
 
 #include <array>
 #include <cstring>
@@ -53,6 +55,11 @@ struct umfMultiPoolTest : umfPoolTest {
     std::vector<umf::pool_unique_handle_t> pools;
 };
 
+struct umfMemTest : umfPoolTest {
+    void SetUp() override { umfPoolTest::SetUp(); }
+    void TearDown() override { umfPoolTest::TearDown(); }
+};
+
 TEST_P(umfPoolTest, allocFree) {
     static constexpr size_t allocSize = 64;
     auto *ptr = umfPoolMalloc(pool.get(), allocSize);
@@ -251,6 +258,31 @@ TEST_P(umfPoolTest, multiThreadedMallocFreeRandomSizes) {
     }
 }
 
+TEST_P(umfMemTest, outOfMem) {
+    static constexpr size_t allocSize = 16;
+    auto hPool = pool.get();
+
+    std::vector<void *> allocations;
+
+    while (true) {
+        allocations.emplace_back(umfPoolMalloc(hPool, allocSize));
+        if (allocations.back() == nullptr &&
+            umfPoolGetLastAllocationError(hPool) ==
+                UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY) {
+            break;
+        }
+        ASSERT_NE(allocations.back(), nullptr);
+    }
+
+    // remove last nullptr from the allocations vector
+    ASSERT_EQ(allocations.back(), nullptr);
+    allocations.pop_back();
+
+    for (auto allocation : allocations) {
+        umfPoolFree(hPool, allocation);
+    }
+}
+
 #ifdef UMF_ENABLE_POOL_TRACKING_TESTS
 // TODO: add similar tests for realloc/aligned_alloc, etc.
 // TODO: add multithreaded tests

test/unified_malloc_framework/memoryPoolAPI.cpp

@@ -82,7 +82,7 @@ TEST_F(test, memoryPoolTrace) {
     ASSERT_EQ(providerCalls.size(), provider_call_count);
 
     ret = umfPoolGetLastAllocationError(tracingPool.get());
-    ASSERT_EQ(ret, UMF_RESULT_ERROR_NOT_SUPPORTED);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
     ASSERT_EQ(poolCalls["get_last_native_error"], 1);
     ASSERT_EQ(poolCalls.size(), ++pool_call_count);
 
@@ -157,6 +157,15 @@ INSTANTIATE_TEST_SUITE_P(
             .second;
     }));
 
+INSTANTIATE_TEST_SUITE_P(
+    proxyPoolOOMTest, umfMemTest, ::testing::Values([] {
+        return umf::poolMakeUnique<umf_test::proxy_pool, 1>(
+                   {umf::memoryProviderMakeUnique<
+                        umf_test::provider_mock_out_of_mem>(10)
+                        .second})
+            .second;
+    }));
+
 ////////////////// Negative test cases /////////////////
 
 TEST_F(test, memoryPoolInvalidProvidersNullptr) {
@@ -254,10 +263,8 @@ TEST_F(test, getLastFailedMemoryProvider) {
     auto [ret, pool] = umf::poolMakeUnique<umf_test::proxy_pool>(&hProvider, 1);
     ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
 
-    ASSERT_EQ(umfGetLastFailedMemoryProvider(), nullptr);
     auto ptr = umfPoolMalloc(pool.get(), allocSize);
     ASSERT_NE(ptr, nullptr);
-    ASSERT_EQ(umfGetLastFailedMemoryProvider(), nullptr);
     umfPoolFree(pool.get(), ptr);
 
     // make provider return an error during allocation

test/unified_malloc_framework/umf_pools/disjoint_pool.cpp

@@ -33,6 +33,16 @@ static auto makePool() {
     return std::move(pool);
 }
 
+static auto makePoolOOMProvider() {
+    auto [ret, provider] =
+        umf::memoryProviderMakeUnique<umf_test::provider_mock_out_of_mem>(10);
+    EXPECT_EQ(ret, UMF_RESULT_SUCCESS);
+    auto [retp, pool] = umf::poolMakeUnique<usm::DisjointPool, 1>(
+        {std::move(provider)}, poolConfig());
+    EXPECT_EQ(retp, UMF_RESULT_SUCCESS);
+    return std::move(pool);
+}
+
 using umf_test::test;
 
 TEST_F(test, freeErrorPropagation) {
@@ -72,6 +82,9 @@ TEST_F(test, freeErrorPropagation) {
 INSTANTIATE_TEST_SUITE_P(disjointPoolTests, umfPoolTest,
                          ::testing::Values(makePool));
 
+INSTANTIATE_TEST_SUITE_P(disjointPoolTests, umfMemTest,
+                         ::testing::Values(makePoolOOMProvider));
+
 GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(umfMultiPoolTest);
 INSTANTIATE_TEST_SUITE_P(disjointMultiPoolTests, umfMultiPoolTest,
                          ::testing::Values(makePool));