From 9c1b4fd1cdc445cd535d66ca67e58031d03daccf Mon Sep 17 00:00:00 2001
From: Covered123 <58715544+JavierLeon9966@users.noreply.github.com>
Date: Mon, 19 Apr 2021 09:41:51 -0300
Subject: [PATCH] Added CancelTaskException (#4186)

---
 src/scheduler/CancelTaskException.php         | 32 +++++++++++++
 src/scheduler/Task.php                        |  2 +
 src/scheduler/TaskHandler.php                 |  2 +
 tests/phpunit/scheduler/TaskSchedulerTest.php | 48 +++++++++++++++++++
 4 files changed, 84 insertions(+)
 create mode 100644 src/scheduler/CancelTaskException.php
 create mode 100644 tests/phpunit/scheduler/TaskSchedulerTest.php

diff --git a/src/scheduler/CancelTaskException.php b/src/scheduler/CancelTaskException.php
new file mode 100644
index 000000000..a500c4d78
--- /dev/null
+++ b/src/scheduler/CancelTaskException.php
@@ -0,0 +1,32 @@
+<?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 pocketmine\scheduler;
+
+/**
+ * This exception can be thrown from Task::onRun() to cancel the execution of the task.
+ * @see Task::onRun()
+ */
+final class CancelTaskException extends \Exception{
+
+}
diff --git a/src/scheduler/Task.php b/src/scheduler/Task.php
index af50cd5f4..0c4ea4469 100644
--- a/src/scheduler/Task.php
+++ b/src/scheduler/Task.php
@@ -49,6 +49,8 @@ abstract class Task{
 
 	/**
 	 * Actions to execute when run
+	 *
+	 * @throws CancelTaskException
 	 */
 	abstract public function onRun() : void;
 
diff --git a/src/scheduler/TaskHandler.php b/src/scheduler/TaskHandler.php
index 73fa15940..e95ee1fcb 100644
--- a/src/scheduler/TaskHandler.php
+++ b/src/scheduler/TaskHandler.php
@@ -115,6 +115,8 @@ class TaskHandler{
 		$this->timings->startTiming();
 		try{
 			$this->task->onRun();
+		}catch(CancelTaskException $e){
+			$this->cancel();
 		}finally{
 			$this->timings->stopTiming();
 		}
diff --git a/tests/phpunit/scheduler/TaskSchedulerTest.php b/tests/phpunit/scheduler/TaskSchedulerTest.php
new file mode 100644
index 000000000..bcb1ae97e
--- /dev/null
+++ b/tests/phpunit/scheduler/TaskSchedulerTest.php
@@ -0,0 +1,48 @@
+<?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 pocketmine\scheduler;
+
+use PHPUnit\Framework\TestCase;
+
+class TaskSchedulerTest extends TestCase{
+
+	/** @var TaskScheduler */
+	private $scheduler;
+
+	public function setUp() : void{
+		$this->scheduler = new TaskScheduler();
+	}
+
+	public function tearDown() : void{
+		$this->scheduler->shutdown();
+	}
+
+	public function testCancel() : void{
+		$task = $this->scheduler->scheduleTask(new ClosureTask(function() : void{
+			throw new CancelTaskException();
+		}));
+		$this->scheduler->mainThreadHeartbeat(0);
+		self::assertTrue($task->isCancelled(), "Task was not cancelled");
+	}
+}