Fix Node UI Storage Threading Issues

Since the same node tree can be used in modifiers on different objects,
there can be multiple threads writing to the maps in the node tree UI
storage at the same time. The additions for attribute name hints and
error messages made it so this would often cause a crash or at least
an ASAN report. This patch adds locks to prevent multiple threads
from using the maps concurrently.

In a brief test I actually didn't observe a crash without the global
`bNodeTree` UI storage mutex, but I think it's necessary for the change
to be correct, and I did notice some unfreed memory without it anyway.
Ideally it would be in a node tree runtime struct though.

Differential Revision: https://developer.blender.org/D10577

2 files changed, 18 insertions, 1 deletions

diff --git a/source/blender/blenkernel/BKE_node_ui_storage.hh b/source/blender/blenkernel/BKE_node_ui_storage.hh
index 231eb11d473..a49ff988272 100644
--- a/source/blender/blenkernel/BKE_node_ui_storage.hh
+++ b/source/blender/blenkernel/BKE_node_ui_storage.hh
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <mutex>
+
 #include "BLI_hash.hh"
 #include "BLI_map.hh"
 #include "BLI_session_uuid.h"
@@ -82,6 +84,7 @@ struct NodeUIStorage {
 
 struct NodeTreeUIStorage {
   blender::Map<NodeTreeEvaluationContext, blender::Map<std::string, NodeUIStorage>> context_map;
+  std::mutex context_map_mutex;
 };
 
 const NodeUIStorage *BKE_node_tree_ui_storage_get_from_context(const bContext *C,
diff --git a/source/blender/blenkernel/intern/node_ui_storage.cc b/source/blender/blenkernel/intern/node_ui_storage.cc
index 397f54ea7e1..6e0253eca31 100644
--- a/source/blender/blenkernel/intern/node_ui_storage.cc
+++ b/source/blender/blenkernel/intern/node_ui_storage.cc
@@ -16,6 +16,8 @@
 
 #include "CLG_log.h"
 
+#include <mutex>
+
 #include "BLI_map.hh"
 #include "BLI_string_ref.hh"
 #include "BLI_vector.hh"
@@ -33,10 +35,20 @@ using blender::Map;
 using blender::StringRef;
 using blender::Vector;
 
+/* Use a global mutex because otherwise it would have to be stored directly in the
+ * bNodeTree struct in DNA. This could change if the node tree had a runtime struct. */
+static std::mutex global_ui_storage_mutex;
+
 static void ui_storage_ensure(bNodeTree &ntree)
 {
+  /* As an optimization, only acquire a lock if the UI storage doesn't exist,
+   * because it only needs to be allocated once for every node tree. */
   if (ntree.ui_storage == nullptr) {
-    ntree.ui_storage = new NodeTreeUIStorage();
+    std::lock_guard<std::mutex> lock(global_ui_storage_mutex);
+    /* Check again-- another thread may have allocated the storage while this one waited. */
+    if (ntree.ui_storage == nullptr) {
+      ntree.ui_storage = new NodeTreeUIStorage();
+    }
   }
 }
 
@@ -74,6 +86,7 @@ void BKE_nodetree_ui_storage_free_for_context(bNodeTree &ntree,
 {
   NodeTreeUIStorage *ui_storage = ntree.ui_storage;
   if (ui_storage != nullptr) {
+    std::lock_guard<std::mutex> lock(ui_storage->context_map_mutex);
     ui_storage->context_map.remove(context);
   }
 }
@@ -116,6 +129,7 @@ static NodeUIStorage &node_ui_storage_ensure(bNodeTree &ntree,
   ui_storage_ensure(ntree);
   NodeTreeUIStorage &ui_storage = *ntree.ui_storage;
 
+  std::lock_guard<std::mutex> lock(ui_storage.context_map_mutex);
   Map<std::string, NodeUIStorage> &node_tree_ui_storage =
       ui_storage.context_map.lookup_or_add_default(context);