[uma] Improve error reporing int DisjointPool

use get_last_allocation_error to return last error from allocation functions.
oneapi-src · Jun 27, 2023 · 438c0d3 · 438c0d3
1 parent fa8ec81
commit 438c0d3
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() {}