From 6d8d2cc9dc2ab7c7341cb6c07f6005636dab9d58 Mon Sep 17 00:00:00 2001
From: Player <player@player.to>
Date: Sun, 3 Nov 2019 18:16:29 +0100
Subject: [PATCH] Fix Fabric-side thread safety for Tags obtained from
 TagRegistry

---
 .../fabricmc/fabric/api/tag/TagRegistry.java  | 14 +----
 .../impl/tag/extension/TagDelegate.java       | 57 +++++++++++++++----
 2 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/api/tag/TagRegistry.java b/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/api/tag/TagRegistry.java
index 1ec17ee90..a4754ec51 100644
--- a/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/api/tag/TagRegistry.java
+++ b/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/api/tag/TagRegistry.java
@@ -41,19 +41,7 @@ public final class TagRegistry {
 	private TagRegistry() { }
 
 	public static <T> Tag<T> create(Identifier id, Supplier<TagContainer<T>> containerSupplier) {
-		return new TagDelegate<T>(id, null) {
-			private TagContainer<T> container;
-
-			@Override
-			protected void onAccess() {
-				TagContainer<T> currContainer = containerSupplier.get();
-
-				if (container != currContainer) {
-					container = currContainer;
-					delegate = container.getOrCreate(this.getId());
-				}
-			}
-		};
+		return new TagDelegate<>(id, containerSupplier);
 	}
 
 	public static Tag<Block> block(Identifier id) {
diff --git a/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/impl/tag/extension/TagDelegate.java b/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/impl/tag/extension/TagDelegate.java
index e0c0d338e..2f3d077ab 100644
--- a/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/impl/tag/extension/TagDelegate.java
+++ b/fabric-tag-extensions-v0/src/main/java/net/fabricmc/fabric/impl/tag/extension/TagDelegate.java
@@ -17,35 +17,68 @@
 package net.fabricmc.fabric.impl.tag.extension;
 
 import java.util.Collection;
+import java.util.function.Supplier;
 
 import net.minecraft.tag.Tag;
+import net.minecraft.tag.TagContainer;
 import net.minecraft.util.Identifier;
 
-public class TagDelegate<T> extends Tag<T> {
-	protected Tag<T> delegate;
+public final class TagDelegate<T> extends Tag<T> {
+	private final Supplier<TagContainer<T>> containerSupplier;
+	private volatile Target<T> target;
 
-	public TagDelegate(Identifier id, Tag<T> delegate) {
+	public TagDelegate(Identifier id, Supplier<TagContainer<T>> containerSupplier) {
 		super(id);
-		this.delegate = delegate;
-	}
 
-	protected void onAccess() { }
+		this.containerSupplier = containerSupplier;
+	}
 
 	@Override
 	public boolean contains(T var1) {
-		onAccess();
-		return delegate.contains(var1);
+		return getTag().contains(var1);
 	}
 
 	@Override
 	public Collection<T> values() {
-		onAccess();
-		return delegate.values();
+		return getTag().values();
 	}
 
 	@Override
 	public Collection<Tag.Entry<T>> entries() {
-		onAccess();
-		return delegate.entries();
+		return getTag().entries();
+	}
+
+	/**
+	 * Retrieve the tag this delegate is pointing to, computing it if missing or outdated.
+	 *
+	 * <p>Thread safety is being ensured by using an immutable holder object for consistently retrieving both result
+	 * and condition, volatile for safe publishing and assuming TagContainer.getOrCreate is safe to call concurrently.
+	 *
+	 * <p>It should be possible to exploit a benign data race on this.target by removing volatile, but this option
+	 * hasn't been chosen yet since a performance problem in the area is yet to be proven.
+	 */
+	private Tag<T> getTag() {
+		Target<T> target = this.target;
+		TagContainer<T> reqContainer = containerSupplier.get();
+		Tag<T> ret;
+
+		if (target == null || target.container != reqContainer) {
+			ret = reqContainer.getOrCreate(getId());
+			this.target = new Target<>(reqContainer, ret);
+		} else {
+			ret = target.tag;
+		}
+
+		return ret;
+	}
+
+	private static final class Target<T> {
+		Target(TagContainer<T> container, Tag<T> tag) {
+			this.container = container;
+			this.tag = tag;
+		}
+
+		final TagContainer<T> container;
+		final Tag<T> tag;
 	}
 }