From 2858db430e80b03aa08cc66cdc28221e761f21c7 Mon Sep 17 00:00:00 2001
From: "Dylan K. Taylor" <odigiman@gmail.com>
Date: Wed, 10 Oct 2018 13:38:04 +0100
Subject: [PATCH] Fixed AsyncTask publishProgress() race condition on task exit

It's possible for a progress update to be lost due to the task finishing before the main thread found the progress update.
---
 src/pocketmine/scheduler/AsyncPool.php        | 14 +++-
 .../src/pmmp/TesterPlugin/Main.php            |  3 +-
 .../AsyncTaskPublishProgressRaceTest.php      | 68 +++++++++++++++++++
 3 files changed, 81 insertions(+), 4 deletions(-)
 create mode 100644 tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/tests/AsyncTaskPublishProgressRaceTest.php

diff --git a/src/pocketmine/scheduler/AsyncPool.php b/src/pocketmine/scheduler/AsyncPool.php
index 5588a8300..3cda006f2 100644
--- a/src/pocketmine/scheduler/AsyncPool.php
+++ b/src/pocketmine/scheduler/AsyncPool.php
@@ -289,12 +289,20 @@ class AsyncPool{
 	 */
 	public function collectTasks() : void{
 		foreach($this->tasks as $task){
-			if(!$task->isGarbage()){
-				$task->checkProgressUpdates($this->server);
-			}
+			$task->checkProgressUpdates($this->server);
 			if($task->isGarbage() and !$task->isRunning() and !$task->isCrashed()){
 				if(!$task->hasCancelledRun()){
 					try{
+						/*
+						 * It's possible for a task to submit a progress update and then finish before the progress
+						 * update is detected by the parent thread, so here we consume any missed updates.
+						 *
+						 * When this happens, it's possible for a progress update to arrive between the previous
+						 * checkProgressUpdates() call and the next isGarbage() call, causing progress updates to be
+						 * lost. Thus, it's necessary to do one last check here to make sure all progress updates have
+						 * been consumed before completing.
+						 */
+						$task->checkProgressUpdates($this->server);
 						$task->onCompletion($this->server);
 						if($task->removeDanglingStoredObjects()){
 							$this->logger->notice("AsyncTask " . get_class($task) . " stored local complex data but did not remove them after completion");
diff --git a/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/Main.php b/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/Main.php
index 12308b23e..c008b1498 100644
--- a/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/Main.php
+++ b/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/Main.php
@@ -44,7 +44,8 @@ class Main extends PluginBase implements Listener{
 
 		$this->waitingTests = [
 			new tests\AsyncTaskMemoryLeakTest($this),
-			new tests\AsyncTaskMainLoggerTest($this)
+			new tests\AsyncTaskMainLoggerTest($this),
+			new tests\AsyncTaskPublishProgressRaceTest($this)
 		];
 	}
 
diff --git a/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/tests/AsyncTaskPublishProgressRaceTest.php b/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/tests/AsyncTaskPublishProgressRaceTest.php
new file mode 100644
index 000000000..8a71f8eee
--- /dev/null
+++ b/tests/plugins/TesterPlugin/src/pmmp/TesterPlugin/tests/AsyncTaskPublishProgressRaceTest.php
@@ -0,0 +1,68 @@
+<?php
+
+/*
+ *
+ *  ____            _        _   __  __ _                  __  __ ____
+ * |  _ \ ___   ___| | _____| |_|  \/  (_)_ __   ___      |  \/  |  _ \
+ * | |_) / _ \ / __| |/ / _ \ __| |\/| | | '_ \ / _ \_____| |\/| | |_) |
+ * |  __/ (_) | (__|   <  __/ |_| |  | | | | | |  __/_____| |  | |  __/
+ * |_|   \___/ \___|_|\_\___|\__|_|  |_|_|_| |_|\___|     |_|  |_|_|
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * @author PocketMine Team
+ * @link http://www.pocketmine.net/
+ *
+ *
+*/
+
+declare(strict_types=1);
+
+namespace pmmp\TesterPlugin\tests;
+
+use pmmp\TesterPlugin\Test;
+use pocketmine\scheduler\AsyncTask;
+use pocketmine\Server;
+
+class AsyncTaskPublishProgressRaceTest extends Test{
+
+	public function getName() : string{
+		return "Verify progress updates work as expected when finishing task";
+	}
+
+	public function getDescription() : string{
+		return "Progress updates would be lost when finishing a task before its remaining progress updates were detected.";
+	}
+
+	public function run() : void{
+		//this test is racy, but it should fail often enough to be a pest if something is broken
+
+		$this->getPlugin()->getServer()->getAsyncPool()->submitTask(new class($this) extends AsyncTask{
+			private static $success = false;
+
+			public function __construct(AsyncTaskPublishProgressRaceTest $t){
+				$this->storeLocal($t);
+			}
+
+			public function onRun() : void{
+				$this->publishProgress("hello");
+			}
+
+			public function onProgressUpdate(Server $server, $progress) : void{
+				if($progress === "hello"){
+					// thread local on main thread
+					self::$success = true;
+				}
+			}
+
+			public function onCompletion(Server $server) : void{
+				/** @var AsyncTaskPublishProgressRaceTest $t */
+				$t = $this->fetchLocal();
+				$t->setResult(self::$success ? Test::RESULT_OK : Test::RESULT_FAILED);
+			}
+		});
+	}
+}