From 2420dee8be087da00dc9f3522de6d6838cf9d5cf Mon Sep 17 00:00:00 2001
From: TheNewHEROBRINE <TheNewHEROBRINEX@users.noreply.github.com>
Date: Wed, 6 Dec 2023 14:40:09 +0100
Subject: [PATCH] AsyncTask: Fix retrieval of null data from the thread-local
 storage (#6176)

---
 src/scheduler/AsyncTask.php               |  3 ++-
 tests/phpunit/scheduler/AsyncPoolTest.php | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/scheduler/AsyncTask.php b/src/scheduler/AsyncTask.php
index ba5cc424c..b4c3ce20d 100644
--- a/src/scheduler/AsyncTask.php
+++ b/src/scheduler/AsyncTask.php
@@ -28,6 +28,7 @@ use pmmp\thread\Thread as NativeThread;
 use pmmp\thread\ThreadSafe;
 use pmmp\thread\ThreadSafeArray;
 use pocketmine\thread\NonThreadSafeValue;
+use function array_key_exists;
 use function assert;
 use function igbinary_serialize;
 use function igbinary_unserialize;
@@ -230,7 +231,7 @@ abstract class AsyncTask extends Runnable{
 	 */
 	protected function fetchLocal(string $key){
 		$id = spl_object_id($this);
-		if(!isset(self::$threadLocalStorage[$id][$key])){
+		if(!isset(self::$threadLocalStorage[$id]) || !array_key_exists($key, self::$threadLocalStorage[$id])){
 			throw new \InvalidArgumentException("No matching thread-local data found on this thread");
 		}
 
diff --git a/tests/phpunit/scheduler/AsyncPoolTest.php b/tests/phpunit/scheduler/AsyncPoolTest.php
index a8a15146e..54c8ccafd 100644
--- a/tests/phpunit/scheduler/AsyncPoolTest.php
+++ b/tests/phpunit/scheduler/AsyncPoolTest.php
@@ -121,4 +121,23 @@ class AsyncPoolTest extends TestCase{
 			usleep(50 * 1000);
 		}
 	}
+
+	public function testNullComplexDataFetch() : void{
+		$this->pool->submitTask(new class extends AsyncTask{
+			public function __construct(){
+				$this->storeLocal("null", null);
+			}
+
+			public function onRun() : void{
+				//dummy
+			}
+
+			public function onCompletion() : void{
+				AsyncPoolTest::assertNull($this->fetchLocal("null"));
+			}
+		});
+		while($this->pool->collectTasks()){
+			usleep(50 * 1000);
+		}
+	}
 }