[umf] add out of memory test

oneapi-src · Aug 17, 2023 · 90738c4 · 90738c4
1 parent 333cfc2
commit 90738c4
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 10 deletions.
diff --git a/test/unified_malloc_framework/common/pool.hpp b/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() {

diff --git a/test/unified_malloc_framework/common/provider.hpp b/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 */
diff --git a/test/unified_malloc_framework/memoryPool.hpp b/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,27 @@ 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);
+ }
+
+ 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

diff --git a/test/unified_malloc_framework/memoryPoolAPI.cpp b/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

diff --git a/test/unified_malloc_framework/umf_pools/disjoint_pool.cpp b/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));