Merge pull request #120545 from stuartcarnie/metal_data_race

Metal: Fix data race when loading shader containers

drivers/metal/metal_objects_shared.cpp

@@ -641,8 +641,8 @@ static const char *SHADER_STAGE_NAMES[] = {
 	"comp", // RDC::SHADER_STAGE_COMPUTE
 };
 
-void ShaderCacheEntry::notify_free() const {
-	owner.shader_cache_free_entry(key);
+void ShaderCacheEntry::notify_free() {
+	owner.shader_cache_free_entry(this);
 }
 
 #pragma mark - MDLibrary

drivers/metal/metal_objects_shared.h

@@ -889,7 +889,7 @@ struct ShaderCacheEntry {
 	std::weak_ptr<MDLibrary> library;
 
 	/// Notify the cache that this entry is no longer needed.
-	void notify_free() const;
+	void notify_free();
 
 	ShaderCacheEntry(RenderingDeviceDriverMetal &p_owner, SHA256Digest p_key) :
 			owner(p_owner), key(p_key) {

drivers/metal/rendering_device_driver_metal.cpp

@@ -1042,13 +1042,34 @@ void RenderingDeviceDriverMetal::framebuffer_free(FramebufferID p_framebuffer) {
 
 #pragma mark - Shader
 
-void RenderingDeviceDriverMetal::shader_cache_free_entry(const SHA256Digest &key) {
-	if (ShaderCacheEntry **pentry = _shader_cache.getptr(key); pentry != nullptr) {
-		ShaderCacheEntry *entry = *pentry;
-		_shader_cache.erase(key);
-		entry->library.reset();
-		memdelete(entry);
+void RenderingDeviceDriverMetal::shader_cache_free_entry(ShaderCacheEntry *p_entry) {
+	MutexLock lock(_shader_cache_lock);
+	// Only remove the map slot if it still refers to this exact entry. A concurrent
+	// creation for the same hash may have replaced it, in which case that newer entry
+	// (and its library) must be left untouched.
+	if (ShaderCacheEntry **pentry = _shader_cache.getptr(p_entry->key); pentry != nullptr && *pentry == p_entry) {
+		_shader_cache.erase(p_entry->key);
 	}
+	memdelete(p_entry);
+}
+
+std::optional<std::shared_ptr<MDLibrary>> RenderingDeviceDriverMetal::shader_cache_get_library(const SHA256Digest &key) {
+	MutexLock lock(_shader_cache_lock);
+
+	if (ShaderCacheEntry **p = _shader_cache.getptr(key); p != nullptr) {
+		if (std::shared_ptr<MDLibrary> lib = (*p)->library.lock()) {
+			return lib;
+		}
+		// Library was released; remove stale cache entry and recreate.
+		_shader_cache.erase(key);
+	}
+
+	return std::nullopt;
+}
+
+void RenderingDeviceDriverMetal::shader_cache_set_entry(const SHA256Digest &key, ShaderCacheEntry *p_entry) {
+	MutexLock lock(_shader_cache_lock);
+	_shader_cache[key] = p_entry;
 }
 
 template <typename T, typename U>
@@ -1121,14 +1142,9 @@ RDD::ShaderID RenderingDeviceDriverMetal::shader_create_from_container(const Ref
 		if (shader.shader_stage == RDD::ShaderStage::SHADER_STAGE_COMPUTE) {
 			pipeline_type = PIPELINE_TYPE_COMPUTE;
 		}
-
-		if (ShaderCacheEntry **p = _shader_cache.getptr(shader_data.hash); p != nullptr) {
-			if (std::shared_ptr<MDLibrary> lib = (*p)->library.lock()) {
-				libraries[shader.shader_stage] = lib;
-				continue;
-			}
-			// Library was released; remove stale cache entry and recreate.
-			_shader_cache.erase(shader_data.hash);
+		if (std::optional<std::shared_ptr<MDLibrary>> lib = shader_cache_get_library(shader_data.hash); lib) {
+			libraries[shader.shader_stage] = std::move(*lib);
+			continue;
 		}
 
 		if (shader.code_decompressed_size > 0) {
@@ -1166,7 +1182,7 @@ RDD::ShaderID RenderingDeviceDriverMetal::shader_create_from_container(const Ref
 			library = MDLibrary::create(cd, device, source.get(), options.get(), _shader_load_strategy);
 		}
 
-		_shader_cache[shader_data.hash] = cd;
+		shader_cache_set_entry(shader_data.hash, cd);
 		libraries[shader.shader_stage] = library;
 	}
 

drivers/metal/rendering_device_driver_metal.h

@@ -36,6 +36,7 @@
 
 #include <Metal/Metal.hpp>
 
+#include <optional>
 #include <variant>
 
 class RenderingShaderContainerFormatMetal;
@@ -178,7 +179,10 @@ protected:
 	 * there are no more references to the MDLibrary associated with the cache entry.
 	 */
 	HashMap<SHA256Digest, ShaderCacheEntry *> _shader_cache;
-	void shader_cache_free_entry(const SHA256Digest &key);
+	Mutex _shader_cache_lock;
+	void shader_cache_free_entry(ShaderCacheEntry *p_entry);
+	std::optional<std::shared_ptr<MDLibrary>> shader_cache_get_library(const SHA256Digest &key);
+	void shader_cache_set_entry(const SHA256Digest &key, ShaderCacheEntry *p_entry);
 
 public:
 	virtual Error initialize(uint32_t p_device_index, uint32_t p_frame_count) override = 0;