From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Spottedleaf <Spottedleaf@users.noreply.github.com>
Date: Mon, 15 May 2023 11:34:28 -0700
Subject: [PATCH] Mark POI/Entity load tasks as completed before releasing
 scheduling lock

It must be marked as completed during that lock hold since the
waiters field is set to null. Thus, any other thread attempting
a cancellation will fail to remove from waiters. Also, any
other thread attempting to cancel may set the completed field
to true which would cause accept() to fail as well.

Completion was always designed to happen while holding the
scheduling lock to prevent these race conditions. The code
was originally set up to complete while not holding the
scheduling lock to avoid invoking callbacks while holding the
lock, however the access to the completion field was not
considered.

Resolve this by marking the callback as completed during the
lock, but invoking the accept() function after releasing
the lock. This will prevent any cancellation attempts to be
blocked, and allow the current thread to complete the callback
without any issues.

diff --git a/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java b/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java
index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 100644
--- a/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java
+++ b/src/main/java/io/papermc/paper/chunk/system/scheduling/NewChunkHolder.java
@@ -0,0 +0,0 @@ public final class NewChunkHolder {
                     LOGGER.error("Unhandled entity data load exception, data data will be lost: ", result.right());
                 }
 
+                // Folia start - mark these tasks as completed before releasing the scheduling lock
+                for (final GenericDataLoadTaskCallback callback : waiters) {
+                    callback.markCompleted();
+                }
+                // Folia end - mark these tasks as completed before releasing the scheduling lock
+
                 completeWaiters = waiters;
             } else {
                 // cancelled
@@ -0,0 +0,0 @@ public final class NewChunkHolder {
         // avoid holding the scheduling lock while completing
         if (completeWaiters != null) {
             for (final GenericDataLoadTaskCallback callback : completeWaiters) {
-                callback.accept(result);
+                callback.acceptCompleted(result); // Folia - mark these tasks as completed before releasing the scheduling lock
             }
         }
 
@@ -0,0 +0,0 @@ public final class NewChunkHolder {
                     LOGGER.error("Unhandled poi load exception, poi data will be lost: ", result.right());
                 }
 
+                // Folia start - mark these tasks as completed before releasing the scheduling lock
+                for (final GenericDataLoadTaskCallback callback : waiters) {
+                    callback.markCompleted();
+                }
+                // Folia end - mark these tasks as completed before releasing the scheduling lock
+
                 completeWaiters = waiters;
             } else {
                 // cancelled
@@ -0,0 +0,0 @@ public final class NewChunkHolder {
         // avoid holding the scheduling lock while completing
         if (completeWaiters != null) {
             for (final GenericDataLoadTaskCallback callback : completeWaiters) {
-                callback.accept(result);
+                callback.acceptCompleted(result); // Folia - mark these tasks as completed before releasing the scheduling lock
             }
         }
         this.scheduler.schedulingLock.lock();
@@ -0,0 +0,0 @@ public final class NewChunkHolder {
         }
     }
 
-    public static abstract class GenericDataLoadTaskCallback implements Cancellable, Consumer<GenericDataLoadTask.TaskResult<?, Throwable>> {
+    public static abstract class GenericDataLoadTaskCallback implements Cancellable { // Folia - mark callbacks as completed before unlocking scheduling lock
 
         protected final Consumer<GenericDataLoadTask.TaskResult<?, Throwable>> consumer;
         protected final NewChunkHolder chunkHolder;
@@ -0,0 +0,0 @@ public final class NewChunkHolder {
             return this.completed = true;
         }
 
-        @Override
-        public void accept(final GenericDataLoadTask.TaskResult<?, Throwable> result) {
+        // Folia start - mark callbacks as completed before unlocking scheduling lock
+        // must hold scheduling lock
+        void markCompleted() {
+            if (this.completed) {
+                throw new IllegalStateException("May not be completed here");
+            }
+            this.completed = true;
+        }
+        // Folia end - mark callbacks as completed before unlocking scheduling lock
+
+        // Folia - mark callbacks as completed before unlocking scheduling lock
+        void acceptCompleted(final GenericDataLoadTask.TaskResult<?, Throwable> result) {
             if (result != null) {
-                if (this.setCompleted()) {
+                if (this.completed) { // Folia - mark callbacks as completed before unlocking scheduling lock
                     this.consumer.accept(result);
                 } else {
-                    throw new IllegalStateException("Cannot be cancelled at this point");
+                    throw new IllegalStateException("Cannot be uncompleted at this point"); // Folia - mark callbacks as completed before unlocking scheduling lock
                 }
             } else {
                 throw new NullPointerException("Result cannot be null (cancelled)");