[uma] Improve error reporing int DisjointPool

use get_last_allocation_error to return last error from allocation functions.
oneapi-src · Jun 26, 2023 · f4eb128 · f4eb128
1 parent b7fb527
commit f4eb128
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/source/common/uma_helpers.hpp b/source/common/uma_helpers.hpp
@@ -155,6 +155,11 @@ auto poolMakeUnique(uma_memory_provider_handle_t *providers,
  ret, pool_unique_handle_t(hPool, &umaPoolDestroy)};
 }
 
+template <typename Type> uma_result_t &getPoolLastStatusRef() {
+ static thread_local uma_result_t last_status = UMA_RESULT_SUCCESS;
+ return last_status;
+}
+
 } // namespace uma
 
 #endif /* UMA_HELPERS_H */
diff --git a/source/common/uma_pools/disjoint_pool.cpp b/source/common/uma_pools/disjoint_pool.cpp
@@ -74,6 +74,10 @@ static size_t AlignUp(size_t Val, size_t Alignment) {
  return (Val + Alignment - 1) & (~(Alignment - 1));
 }
 
+struct MemoryProviderError {
+ uma_result_t code;
+};
+
 DisjointPoolConfig::DisjointPoolConfig()
  : limits(std::make_shared<DisjointPoolConfig::SharedLimits>()) {}
 
@@ -371,7 +375,13 @@ Slab::~Slab() {
  } catch (std::exception &e) {
  std::cout << "DisjointPool: unexpected error: " << e.what() << "\n";
  }
- memoryProviderFree(bucket.getMemHandle(), MemPtr);
+
+ try {
+ memoryProviderFree(bucket.getMemHandle(), MemPtr);
+ } catch (MemoryProviderError &e) {
+ std::cout << "DisjointPool: error from memory provider: " << e.code
+ << "\n";
+ }
 }
 
 // Return the index of the first available chunk, SIZE_MAX otherwise
@@ -717,7 +727,7 @@ void Bucket::printStats(bool &TitlePrinted, const std::string &Label) {
  }
 }
 
-void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) {
+void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) try {
  void *Ptr;
 
  if (Size == 0) {
@@ -742,10 +752,13 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) {
  }
 
  return Ptr;
+} catch (MemoryProviderError &e) {
+ uma::getPoolLastStatusRef<DisjointPool>() = e.code;
+ return nullptr;
 }
 
 void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
- bool &FromPool) {
+ bool &FromPool) try {
  void *Ptr;
 
  if (Size == 0) {
@@ -789,6 +802,9 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
  }
 
  return AlignPtrUp(Ptr, Alignment);
+} catch (MemoryProviderError &e) {
+ uma::getPoolLastStatusRef<DisjointPool>() = e.code;
+ return nullptr;
 }
 
 Bucket &DisjointPool::AllocImpl::findBucket(size_t Size) {
@@ -803,7 +819,7 @@ Bucket &DisjointPool::AllocImpl::findBucket(size_t Size) {
  return *(*It);
 }
 
-void DisjointPool::AllocImpl::deallocate(void *Ptr, bool &ToPool) {
+void DisjointPool::AllocImpl::deallocate(void *Ptr, bool &ToPool) try {
  auto *SlabPtr = AlignPtrDown(Ptr, SlabMinSize());
 
  // Lock the map on read
@@ -846,6 +862,8 @@ void DisjointPool::AllocImpl::deallocate(void *Ptr, bool &ToPool) {
  // to some slab with an entry in the map. So we find a slab
  // but the range checks fail.
  memoryProviderFree(getMemHandle(), Ptr);
+} catch (MemoryProviderError &e) {
+ uma::getPoolLastStatusRef<DisjointPool>() = e.code;
 }
 
 void DisjointPool::AllocImpl::printStats(bool &TitlePrinted,
@@ -938,7 +956,7 @@ void DisjointPool::free(void *ptr) {
 }
 
 enum uma_result_t DisjointPool::get_last_allocation_error() {
- return UMA_RESULT_ERROR_UNKNOWN;
+ return uma::getPoolLastStatusRef<DisjointPool>();
 }
 
 DisjointPool::DisjointPool() {}