Allow OOM strategy for RBTree

Author: jdksjolen

src/hotspot/share/utilities/rbTree.hpp

@@ -468,13 +468,17 @@ class RBTree : public AbstractRBTree<K, RBNode<K, V>, COMPARATOR> {
 
   RBNode<K, V>* allocate_node(const K& key) {
     void* node_place = _allocator.allocate(sizeof(RBNode<K, V>));
-    assert(node_place != nullptr, "rb-tree allocator must exit on failure");
+    if (node_place == nullptr) {
+      return nullptr;
+    }
     return new (node_place) RBNode<K, V>(key);
   }
 
   RBNode<K, V>* allocate_node(const K& key, const V& val) {
     void* node_place = _allocator.allocate(sizeof(RBNode<K, V>));
-    assert(node_place != nullptr, "rb-tree allocator must exit on failure");
+    if (node_place == nullptr) {
+      return nullptr;
+    }
     return new (node_place) RBNode<K, V>(key, val);
   }
 
@@ -485,16 +489,21 @@ class RBTree : public AbstractRBTree<K, RBNode<K, V>, COMPARATOR> {
 
   // Inserts a node with the given key/value into the tree,
   // if the key already exist, the value is updated instead.
-  void upsert(const K& key, const V& val, const RBNode<K, V>* hint_node = nullptr) {
+  // Returns false if and only if allocation of a new node failed.
+  bool upsert(const K& key, const V& val, const RBNode<K, V>* hint_node = nullptr) {
     Cursor node_cursor = cursor(key, hint_node);
     RBNode<K, V>* node = node_cursor.node();
     if (node != nullptr) {
       node->set_val(val);
-      return;
+      return true;
     }
 
     node = allocate_node(key, val);
+    if (node == nullptr) {
+      return false;
+    }
     insert_at_cursor(node, node_cursor);
+    return true;
   }
 
   // Finds the value of the node associated with the given key.
@@ -545,12 +554,12 @@ class RBTree : public AbstractRBTree<K, RBNode<K, V>, COMPARATOR> {
   }
 };
 
-template <MemTag mem_tag>
+template <MemTag mem_tag, AllocFailType strategy>
 class RBTreeCHeapAllocator {
 public:
   void* allocate(size_t sz) {
     void* allocation = os::malloc(sz, mem_tag);
-    if (allocation == nullptr) {
+    if (allocation == nullptr && strategy == AllocFailStrategy::EXIT_OOM) {
       vm_exit_out_of_memory(sz, OOM_MALLOC_ERROR,
                             "red-black tree failed allocation");
     }
@@ -560,8 +569,8 @@ class RBTreeCHeapAllocator {
   void free(void* ptr) { os::free(ptr); }
 };
 
-template <typename K, typename V, typename COMPARATOR, MemTag mem_tag>
-using RBTreeCHeap = RBTree<K, V, COMPARATOR, RBTreeCHeapAllocator<mem_tag>>;
+template <typename K, typename V, typename COMPARATOR, MemTag mem_tag, AllocFailStrategy strategy = AllocFailStrategy::EXIT_OOM>
+using RBTreeCHeap = RBTree<K, V, COMPARATOR, RBTreeCHeapAllocator<mem_tag, strategy>>;
 
 template <typename K, typename COMPARATOR>
 using IntrusiveRBTree = AbstractRBTree<K, IntrusiveRBNode, COMPARATOR>;

test/hotspot/gtest/utilities/test_rbtree.cpp

@@ -1079,3 +1079,15 @@ TEST_VM_F(RBTreeTest, VerifyItThroughStressTest) {
   }
 }
 
+struct OomAllocator {
+  void* allocate(size_t sz) {
+    return nullptr;
+  }
+  void free(void* ptr) {}
+};
+TEST_VM_F(RBTreeTest, AllocatorMayReturnNull) {
+  RBTree<int, int, Cmp, OomAllocator> rbtree;
+  bool success = rbtree.upsert(5, 5);
+  EXPECT_EQ(false, success);
+  // The test didn't exit the VM, so it was succesful.
+}