From 98af18921aa3c274ef7ece03005337b58df3da96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Johan=20Sj=C3=B6len?= <jsjolen@openjdk.org>
Date: Mon, 1 Sep 2025 09:24:52 +0000
Subject: [PATCH] 8366456: Allow AllocFailStrategy for RBTree

Reviewed-by: cnorrbin, aboldtch
---
 src/hotspot/share/utilities/rbTree.hpp       | 26 +++++++----
 test/hotspot/gtest/utilities/test_rbtree.cpp | 46 ++++++++++++--------
 2 files changed, 46 insertions(+), 26 deletions(-)

diff --git a/src/hotspot/share/utilities/rbTree.hpp b/src/hotspot/share/utilities/rbTree.hpp
index 405fa3e9ae9..cc52cec3fe0 100644
--- a/src/hotspot/share/utilities/rbTree.hpp
+++ b/src/hotspot/share/utilities/rbTree.hpp
@@ -61,7 +61,6 @@
 // used for extra validation can optionally be provided. This should return:
 //     - true if a < b
 //     - false otherwise
-// ALLOCATOR must check for oom and exit, as RBTree does not handle the allocation failing.
 // K needs to be of a type that is trivially destructible.
 // The tree will call a value's destructor when its node is removed.
 // Nodes are address stable and will not change during its lifetime.
@@ -468,13 +467,17 @@ public:
 
   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 +488,21 @@ public:
 
   // 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 +553,12 @@ public:
   }
 };
 
-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 +568,8 @@ public:
   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, AllocFailType 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>;
diff --git a/test/hotspot/gtest/utilities/test_rbtree.cpp b/test/hotspot/gtest/utilities/test_rbtree.cpp
index e1be83bfd58..d8394de017e 100644
--- a/test/hotspot/gtest/utilities/test_rbtree.cpp
+++ b/test/hotspot/gtest/utilities/test_rbtree.cpp
@@ -986,23 +986,23 @@ struct IntCmp {
 };
 
 TEST_VM(RBTreeTestNonFixture, TestPrintIntegerTree) {
-  typedef RBTree<int, unsigned, IntCmp, RBTreeCHeapAllocator<mtTest> > TreeType;
-    TreeType tree;
-    const int i1 = 82924;
-    const char* const s1 = "[82924] = 1";
-    const int i2 = -13591;
-    const char* const s2 = "[-13591] = 2";
-    const int i3 = 0;
-    const char* const s3 = "[0] = 3";
-    tree.upsert(i1, 1U);
-    tree.upsert(i2, 2U);
-    tree.upsert(i3, 3U);
-    stringStream ss;
-    tree.print_on(&ss);
-    const char* const N = nullptr;
-    ASSERT_NE(strstr(ss.base(), s1), N);
-    ASSERT_NE(strstr(ss.base(), s2), N);
-    ASSERT_NE(strstr(ss.base(), s3), N);
+  using TreeType = RBTreeCHeap<int, unsigned, IntCmp, mtTest>;
+  TreeType tree;
+  const int i1 = 82924;
+  const char* const s1 = "[82924] = 1";
+  const int i2 = -13591;
+  const char* const s2 = "[-13591] = 2";
+  const int i3 = 0;
+  const char* const s3 = "[0] = 3";
+  tree.upsert(i1, 1U);
+  tree.upsert(i2, 2U);
+  tree.upsert(i3, 3U);
+  stringStream ss;
+  tree.print_on(&ss);
+  const char* const N = nullptr;
+  ASSERT_NE(strstr(ss.base(), s1), N);
+  ASSERT_NE(strstr(ss.base(), s2), N);
+  ASSERT_NE(strstr(ss.base(), s3), N);
 }
 
 TEST_VM_F(RBTreeTest, IntrusiveTest) {
@@ -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.
+}