66#include " memspace_helpers.hpp"
77#include " memspace_internal.h"
88#include " test_helpers.h"
9+ #include " utils_sanitizers.h"
910
1011#include < numa.h>
1112#include < numaif.h>
@@ -96,39 +97,53 @@ static void getAllocationPolicy(void *ptr, unsigned long maxNodeId, int &mode,
9697 allocNodeId = static_cast <size_t >(nodeId);
9798}
9899
99- TEST_F (memspaceHostAllProviderTest, memoryPolicyOOM) {
100+ TEST_F (memspaceHostAllProviderTest, allocsSpreadAcrossAllNumaNodes) {
101+ // This testcase is unsuitable for TSan.
102+ #ifdef __SANITIZE_THREAD__
103+ GTEST_SKIP ();
104+ #endif
105+
100106 // Arbitrary allocation size, should be big enough to avoid unnecessarily
101107 // prolonging the test execution.
102- size_t size = SIZE_4M * 128 ;
108+ size_t size = SIZE_4M;
103109 size_t alignment = 0 ;
104- std::vector<void *> allocs;
105110
106- enum umf_result_t umf_ret = UMF_RESULT_SUCCESS;
107- // Create allocations until OOM.
108- while (true ) {
111+ long long numaCombinedFreeSize = 0 ;
112+ // Gather free size of all numa nodes.
113+ for (auto &id : nodeIds) {
114+ long long numaFreeSize = 0 ;
115+ long long numaSize = numa_node_size64 (id, &numaFreeSize);
116+ UT_ASSERTne (numaSize, -1 );
117+ // We need the space for at least two allocations, so that we can
118+ // have some space left to avoid OOM killer.
119+ UT_ASSERT (numaFreeSize >= (long long )(2 * size));
120+
121+ numaCombinedFreeSize += numaFreeSize;
122+ }
123+
124+ umf_result_t umf_ret = UMF_RESULT_SUCCESS;
125+ // Create allocations until all the NUMA nodes until there's space only for
126+ // one allocation.
127+ std::vector<void *> allocs;
128+ std::unordered_set<size_t > allocNodeIds;
129+ while (numaCombinedFreeSize >= (long long )(2 * size)) {
109130 void *ptr = nullptr ;
110131 umf_ret = umfMemoryProviderAlloc (hProvider, size, alignment, &ptr);
111132 if (umf_ret != UMF_RESULT_SUCCESS) {
133+ UT_ASSERTeq (umf_ret, UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC);
134+ const char *msg = nullptr ;
135+ int32_t err = 0 ;
136+ umfMemoryProviderGetLastNativeError (hProvider, &msg, &err);
137+ // In this scenario, 'UMF_OS_RESULT_ERROR_ALLOC_FAILED' indicates OOM.
138+ UT_ASSERTeq (err, UMF_OS_RESULT_ERROR_ALLOC_FAILED);
112139 break ;
113140 }
114141
115142 UT_ASSERTne (ptr, nullptr );
116- allocs.push_back (ptr);
117- }
143+ // Access the allocation, so that all the pages associated with it are
144+ // allocated on available NUMA nodes.
145+ memset (ptr, 0xFF , size);
118146
119- UT_ASSERTeq (umf_ret, UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC);
120- const char *msg = nullptr ;
121- int32_t err = 0 ;
122- umfMemoryProviderGetLastNativeError (hProvider, &msg, &err);
123- // In this scenario, 'UMF_OS_RESULT_ERROR_ALLOC_FAILED' indicates OOM.
124- UT_ASSERTeq (err, UMF_OS_RESULT_ERROR_ALLOC_FAILED);
125-
126- // When allocating until OOM, the allocations should be distributed across
127- // all the NUMA nodes bound to 'HOST ALL' memspace, until each node runs
128- // out of memory.
129- UT_ASSERT (allocs.size () >= nodeIds.size ());
130- std::unordered_set<size_t > allocNodeIds;
131- for (auto &ptr : allocs) {
132147 int mode = -1 ;
133148 std::vector<size_t > boundNodeIds;
134149 size_t allocNodeId = SIZE_MAX;
@@ -151,8 +166,14 @@ TEST_F(memspaceHostAllProviderTest, memoryPolicyOOM) {
151166 auto it = std::find (nodeIds.begin (), nodeIds.end (), allocNodeId);
152167 UT_ASSERT (it != nodeIds.end ());
153168
169+ allocs.push_back (ptr);
154170 allocNodeIds.insert (allocNodeId);
155171
172+ numaCombinedFreeSize -= size;
173+ }
174+
175+ UT_ASSERT (allocs.size () >= nodeIds.size ());
176+ for (auto &ptr : allocs) {
156177 umf_ret = umfMemoryProviderFree (hProvider, ptr, size);
157178 UT_ASSERTeq (umf_ret, UMF_RESULT_SUCCESS);
158179 }
0 commit comments