Builtin thread-local object storage for AsyncTask

src/pocketmine/scheduler/AsyncPool.php

@@ -146,6 +146,7 @@ public function collectTasks(){
 
 				if(!$task->hasCancelledRun()){
 					$task->onCompletion($this->server);
+					$this->server->getScheduler()->removeLocalComplex($task);
 				}
 
 				$this->removeTask($task);

src/pocketmine/scheduler/AsyncTask.php

@@ -27,7 +27,7 @@
 /**
  * Class used to run async tasks in other threads.
  *
- * WARNING: Do not call PocketMine-MP API methods, or save objects from/on other Threads!!
+ * WARNING: Do not call PocketMine-MP API methods, or save objects (and arrays cotaining objects) from/on other Threads!!
  */
 abstract class AsyncTask extends Collectable{
 
@@ -42,6 +42,29 @@ abstract class AsyncTask extends Collectable{
 
 	private $crashed = false;
 
+	/**
+	 * Constructs a new instance of AsyncTask. Subclasses don't need to call this constructor unless an argument is to be passed. ONLY construct this class from the main thread.
+	 * <br>
+	 * If an argument is passed into this constructor, it will be stored in a thread-local storage (in ServerScheduler), which MUST be retrieved through {@link #fetchLocal} when {@link #onCompletion} is called.
+	 * Otherwise, a NOTICE level message will be raised and the reference will be removed after onCompletion exits.
+	 * <br>
+	 * If null or no argument is passed, do <em>not</em> call {@link #fetchLocal}, or an exception will be thrown.
+	 * <br>
+	 * WARNING: Use this method carefully. It might take a long time before an AsyncTask is completed. PocketMine will keep a strong reference to objects passed in this method.
+	 * This may result in a light memory leak. Usually this does not cause memory failure, but be aware that the object may be no longer usable when the AsyncTask completes.
+	 * (E.g. a {@link \pocketmine\Level} object is no longer usable because it is unloaded while the AsyncTask is executing, or even a plugin might be unloaded)
+	 * Since PocketMine keeps a strong reference, the objects are still valid, but the implementation is responsible for checking whether these objects are still usable.
+	 *
+	 * @param mixed $complexData the data to store, pass null to store nothing. Scalar types can be safely stored in class properties directly instead of using this thread-local storage.
+	 */
+	public function __construct($complexData = null){
+		if($complexData === null){
+			return;
+		}
+
+		Server::getInstance()->getScheduler()->storeLocalComplex($this, $complexData);
+	}
+
 	public function run(){
 		$this->result = null;
 
@@ -145,6 +168,24 @@ public function onCompletion(Server $server){
 
 	}
 
+	/**
+	 * Call this method from {@link #onCompletion} to fetch the data stored in the constructor, if any.
+	 *
+	 * @param Server $server default null
+	 *
+	 * @return mixed
+	 *
+	 * @throws \RuntimeException if no data were stored by this AsyncTask instance.
+	 */
+	protected function fetchLocal(Server $server = null){
+		if($server === null){
+			$server = Server::getInstance();
+			assert($server !== null, "Call this method only from the main thread!");
+		}
+
+		return $server->getScheduler()->fetchLocalComplex($this);
+	}
+
 	public function cleanObject(){
 		foreach($this as $p => $v){
 			if(!($v instanceof \Threaded)){
@@ -154,3 +195,4 @@ public function cleanObject(){
 	}
 
 }
+

src/pocketmine/scheduler/ServerScheduler.php

@@ -50,9 +50,13 @@ class ServerScheduler{
 	/** @var int */
 	protected $currentTick = 0;
 
+	/** @var \SplObjectStorage<AsyncTask> */
+	protected $objectStore;
+
 	public function __construct(){
 		$this->queue = new ReversePriorityQueue();
 		$this->asyncPool = new AsyncPool(Server::getInstance(), self::$WORKERS);
+		$this->objectStore = new \SplObjectStorage();
 	}
 
 	/**
@@ -91,6 +95,59 @@ public function scheduleAsyncTaskToWorker(AsyncTask $task, $worker){
 		$this->asyncPool->submitTaskToWorker($task, $worker);
 	}
 
+	/**
+	 * Stores any data that must not be passed to other threads or be serialized
+	 *
+	 * @internal Only call from AsyncTask.php
+	 *
+	 * @param AsyncTask     $for
+	 * @param object|array $cmplx
+	 *
+	 * @throws \RuntimeException if this method is called twice for the same instance of AsyncTask
+	 */
+	public function storeLocalComplex(AsyncTask $for, $cmplx){
+		if(isset($this->objectStore[$for])){
+			throw new \RuntimeException("Already storing a complex for this AsyncTask");
+		}
+		$this->objectStore[$for] = $cmplx;
+	}
+
+	/**
+	 * Fetches data that must not be passed to other threads or be serialized, previously stored with {@link #storeLocalComplex}
+	 *
+	 * @internal Only call from AsyncTask.php
+	 *
+	 * @param AsyncTask $for
+	 *
+	 * @throws \RuntimException if no data associated with this AsyncTask can be found
+	 */	
+	public function fetchLocalComplex(AsyncTask $for){
+		if(!isset($this->objectStore[$for])){
+			throw new \RuntimeException("Attempt to fetch undefined complex");
+		}
+		$cmplx = $this->objectStore[$for];
+		unset($this->objectStore[$for]);
+		return $cmplx;
+	}
+
+	/**
+	 * Makes sure no data stored from {@link #storeLocalComplex} is left for a specific AsyncTask
+	 *
+	 * @internal Only call from AsyncTask.php
+	 *
+	 * @param AsyncTask $for
+	 *
+	 * @return bool returns false if any data are removed from this call, true otherwise
+	 */	
+	public function removeLocalComplex(AsyncTask $for) : bool{
+		if(isset($this->objectStore[$for])){
+			Server::getInstance()->getLogger()->notice("AsyncTask " . get_class($for) . " stored local complex data but did not remove them after completion");
+			unset($this->objectStore[$for]);
+			return false;
+		}
+		return true;
+	}
+
 	public function getAsyncTaskPoolSize(){
 		return $this->asyncPool->getSize();
 	}