Mutate the stack traces of exceptions thrown in tasks to make them ea…

…sier to track

fabric-client-gametest-api-v1/src/client/java/net/fabricmc/fabric/api/client/gametest/v1/package-info.java

@@ -11,8 +11,11 @@
  *
  * <p>Client gametests run on the client gametest thread. Use the functions inside
  * {@link net.fabricmc.fabric.api.client.gametest.v1.ClientGameTestContext ClientGameTestContext} and other test helper
- * classes to run code on the correct thread. The game remains paused unless you explicitly unpause it using various
- * waiting functions such as
+ * classes to run code on the correct thread. Exceptions are transparently rethrown on the test thread, and their stack
+ * traces are mutated to include the async stack trace, to make them easy to track. You can disable this behavior by
+ * setting the {@code fabric.client.gametest.disableJoinAsyncStackTraces} system property.
+ *
+ * <p>The game remains paused unless you explicitly unpause it using various waiting functions such as
  * {@link net.fabricmc.fabric.api.client.gametest.v1.ClientGameTestContext#waitTick() ClientGameTestContext.waitTick()}.
  *
  * <p>A few changes have been made to how the vanilla game threads run, to make tests more reproducible. Notably, there

fabric-client-gametest-api-v1/src/client/java/net/fabricmc/fabric/impl/client/gametest/ThreadingImpl.java

@@ -70,6 +70,11 @@ private ThreadingImpl() {
 
 	private static final Logger LOGGER = LoggerFactory.getLogger("fabric-client-gametest-api-v1");
 
+	private static final boolean DISABLE_JOIN_ASYNC_STACK_TRACES = System.getProperty("fabric.client.gametest.disableJoinAsyncStackTraces") != null;
+	private static final String THREAD_IMPL_CLASS_NAME = ThreadingImpl.class.getName();
+	private static final String TASK_ON_THIS_THREAD_METHOD_NAME = "runTaskOnThisThread";
+	private static final String TASK_ON_OTHER_THREAD_METHOD_NAME = "runTaskOnOtherThread";
+
 	public static final int PHASE_TICK = 0;
 	public static final int PHASE_SERVER_TASKS = 1;
 	public static final int PHASE_CLIENT_TASKS = 2;
@@ -160,25 +165,25 @@ public static void checkOnGametestThread(String methodName) {
 		Preconditions.checkState(Thread.currentThread() == testThread, "%s can only be called from the client gametest thread", methodName);
 	}
 
-	@SuppressWarnings("unchecked")
 	public static <E extends Throwable> void runOnClient(FailableRunnable<E> action) throws E {
 		Preconditions.checkNotNull(action, "action");
 		checkOnGametestThread("runOnClient");
 		Preconditions.checkState(clientCanAcceptTasks, "runOnClient called when no client is running");
+		runTaskOnOtherThread(action, CLIENT_SEMAPHORE);
+	}
 
+	public static <E extends Throwable> void runOnServer(FailableRunnable<E> action) throws E {
+		Preconditions.checkNotNull(action, "action");
+		checkOnGametestThread("runOnServer");
+		Preconditions.checkState(serverCanAcceptTasks, "runOnServer called when no server is running");
+		runTaskOnOtherThread(action, SERVER_SEMAPHORE);
+	}
+
+	private static <E extends Throwable> void runTaskOnOtherThread(FailableRunnable<E> action, Semaphore clientOrServerSemaphore) throws E {
 		MutableObject<E> thrown = new MutableObject<>();
-		taskToRun = () -> {
-			try {
-				action.run();
-			} catch (Throwable e) {
-				thrown.setValue((E) e);
-			} finally {
-				taskToRun = null;
-				TEST_SEMAPHORE.release();
-			}
-		};
+		taskToRun = () -> runTaskOnThisThread(action, thrown);
 
-		CLIENT_SEMAPHORE.release();
+		clientOrServerSemaphore.release();
 
 		try {
 			TEST_SEMAPHORE.acquire();
@@ -187,39 +192,73 @@ public static <E extends Throwable> void runOnClient(FailableRunnable<E> action)
 		}
 
 		if (thrown.getValue() != null) {
+			joinAsyncStackTrace(thrown.getValue());
 			throw thrown.getValue();
 		}
 	}
 
 	@SuppressWarnings("unchecked")
-	public static <E extends Throwable> void runOnServer(FailableRunnable<E> action) throws E {
-		Preconditions.checkNotNull(action, "action");
-		checkOnGametestThread("runOnServer");
-		Preconditions.checkState(serverCanAcceptTasks, "runOnServer called when no server is running");
+	private static <E extends Throwable> void runTaskOnThisThread(FailableRunnable<E> action, MutableObject<E> thrown) {
+		try {
+			action.run();
+		} catch (Throwable e) {
+			thrown.setValue((E) e);
+		} finally {
+			taskToRun = null;
+			TEST_SEMAPHORE.release();
+		}
+	}
 
-		MutableObject<E> thrown = new MutableObject<>();
-		taskToRun = () -> {
-			try {
-				action.run();
-			} catch (Throwable e) {
-				thrown.setValue((E) e);
-			} finally {
-				taskToRun = null;
-				TEST_SEMAPHORE.release();
+	private static void joinAsyncStackTrace(Throwable e) {
+		if (DISABLE_JOIN_ASYNC_STACK_TRACES) {
+			return;
+		}
+
+		// find the end of the relevant part of the stack trace on the other thread
+		StackTraceElement[] taskStackTrace = e.getStackTrace();
+
+		if (taskStackTrace == null) {
+			return;
+		}
+
+		int taskRunOnThreadIndex = taskStackTrace.length - 1;
+
+		for (; taskRunOnThreadIndex >= 0; taskRunOnThreadIndex--) {
+			StackTraceElement element = taskStackTrace[taskRunOnThreadIndex];
+
+			if (THREAD_IMPL_CLASS_NAME.equals(element.getClassName()) && TASK_ON_THIS_THREAD_METHOD_NAME.equals(element.getMethodName())) {
+				break;
 			}
-		};
+		}
 
-		SERVER_SEMAPHORE.release();
+		if (taskRunOnThreadIndex == -1) {
+			// couldn't find stack trace element
+			return;
+		}
 
-		try {
-			TEST_SEMAPHORE.acquire();
-		} catch (InterruptedException e) {
-			throw new RuntimeException(e);
+		// find the start of the relevant part of the stack trace on the test thread
+		StackTraceElement[] testStackTrace = Thread.currentThread().getStackTrace();
+		int testRunOnThreadIndex = 0;
+
+		for (; testRunOnThreadIndex < testStackTrace.length; testRunOnThreadIndex++) {
+			StackTraceElement element = testStackTrace[testRunOnThreadIndex];
+
+			if (THREAD_IMPL_CLASS_NAME.equals(element.getClassName()) && TASK_ON_OTHER_THREAD_METHOD_NAME.equals(element.getMethodName())) {
+				break;
+			}
 		}
 
-		if (thrown.getValue() != null) {
-			throw thrown.getValue();
+		if (testRunOnThreadIndex == testStackTrace.length) {
+			// couldn't find stack trace element
+			return;
 		}
+
+		// join the stack traces
+		StackTraceElement[] joinedStackTrace = new StackTraceElement[(taskRunOnThreadIndex + 1) + 1 + (testStackTrace.length - taskRunOnThreadIndex)];
+		System.arraycopy(taskStackTrace, 0, joinedStackTrace, 0, taskRunOnThreadIndex + 1);
+		joinedStackTrace[taskRunOnThreadIndex + 1] = new StackTraceElement("Async Stack Trace", ".", null, 1);
+		System.arraycopy(testStackTrace, testRunOnThreadIndex, joinedStackTrace, taskRunOnThreadIndex + 2, testStackTrace.length - testRunOnThreadIndex);
+		e.setStackTrace(joinedStackTrace);
 	}
 
 	public static void runTick() {