Allow MemtargetRemove to remove the last target with no issue

previously we got OUT_OF_HOST_MEMORY while removing it.
Extend the testing to cover this.

Author: lukaszstolarczuk

include/umf/memspace.h

@@ -46,9 +46,9 @@ umfMemoryProviderCreateFromMemspace(umf_const_memspace_handle_t hMemspace,
                                     umf_const_mempolicy_handle_t hPolicy,
                                     umf_memory_provider_handle_t *hProvider);
 ///
-/// \brief Creates new memspace from array of NUMA node ids.
+/// \brief Creates new memspace from an array of NUMA node ids.
 /// \param nodeIds array of NUMA node ids
-/// \param numIds size of the array
+/// \param numIds size of the array; it has to be greater than 0
 /// \param hMemspace [out] handle to the newly created memspace
 /// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
 ///

src/memspace.c

@@ -387,6 +387,7 @@ umfMemspaceMemtargetRemove(umf_memspace_handle_t hMemspace,
     if (!hMemspace || !hMemtarget) {
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
+
     unsigned i;
     for (i = 0; i < hMemspace->size; i++) {
         int cmp;
@@ -407,10 +408,16 @@ umfMemspaceMemtargetRemove(umf_memspace_handle_t hMemspace,
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    umf_memtarget_handle_t *newNodes =
-        umf_ba_global_alloc(sizeof(*newNodes) * (hMemspace->size - 1));
-    if (!newNodes) {
-        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    umf_memtarget_handle_t *newNodes = NULL;
+
+    if (hMemspace->size == 1) {
+        LOG_DEBUG("Removing the last memory target from the memspace.");
+    } else {
+        newNodes =
+            umf_ba_global_alloc(sizeof(*newNodes) * (hMemspace->size - 1));
+        if (!newNodes) {
+            return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        }
     }
 
     for (unsigned j = 0, z = 0; j < hMemspace->size; j++) {

test/memspaces/memtarget.cpp

@@ -145,3 +145,30 @@ TEST_F(test, memTargetInvalidRemove) {
     ret = umfMemspaceDestroy(memspace);
     EXPECT_EQ(ret, UMF_RESULT_SUCCESS);
 }
+
+TEST_F(test, memTargetRemoveAll) {
+    umf_const_memspace_handle_t const_memspace = umfMemspaceHostAllGet();
+    umf_memspace_handle_t memspace = nullptr;
+    umf_result_t ret = umfMemspaceClone(const_memspace, &memspace);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    ASSERT_NE(memspace, nullptr);
+    umf_const_memtarget_handle_t memtarget = nullptr;
+
+    // Remove all memtargets
+    size_t len = umfMemspaceMemtargetNum(memspace);
+    ASSERT_GT(len, 0);
+    size_t i = len - 1;
+    do {
+        memtarget = umfMemspaceMemtargetGet(memspace, i);
+        EXPECT_NE(memtarget, nullptr);
+        ret = umfMemspaceMemtargetRemove(memspace, memtarget);
+        ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    } while (i-- > 0);
+
+    // Try to remove the last one for the second time
+    ret = umfMemspaceMemtargetRemove(memspace, memtarget);
+    EXPECT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
+
+    ret = umfMemspaceDestroy(memspace);
+    EXPECT_EQ(ret, UMF_RESULT_SUCCESS);
+}