From 2674ca933ef7caabe1fa216785efc4185c58bae4 Mon Sep 17 00:00:00 2001
From: Matthew <matik04@gmail.com>
Date: Sun, 15 Aug 2021 13:13:48 +0200
Subject: [PATCH] Move functionality from `TaskBuilder#schedule` to
 `Task#schedule` and `Task#scheduleForShutdown`

---
 .../java/net/minestom/server/timer/Task.java  | 14 +++++++++++++
 .../minestom/server/timer/TaskBuilder.java    | 21 ++-----------------
 2 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/src/main/java/net/minestom/server/timer/Task.java b/src/main/java/net/minestom/server/timer/Task.java
index a97d1ca1c..69d7d5ad9 100644
--- a/src/main/java/net/minestom/server/timer/Task.java
+++ b/src/main/java/net/minestom/server/timer/Task.java
@@ -100,6 +100,10 @@ public class Task implements Runnable {
      * Sets up the task for correct execution.
      */
     public void schedule() {
+        Int2ObjectMap<Task> tasks = this.schedulerManager.tasks;
+        synchronized (tasks) {
+            tasks.put(getId(), this);
+        }
         if(owningExtension != null) {
             this.schedulerManager.onScheduleFromExtension(owningExtension, this);
         }
@@ -108,6 +112,16 @@ public class Task implements Runnable {
                 this.schedulerManager.getTimerExecutionService().scheduleAtFixedRate(this, this.delay, this.repeat, TimeUnit.MILLISECONDS);
     }
 
+    public void scheduleForShutdown() {
+        Int2ObjectMap<Task> shutdownTasks = this.schedulerManager.shutdownTasks;
+        synchronized (shutdownTasks) {
+            shutdownTasks.put(getId(), this);
+        }
+        if (owningExtension != null) {
+            this.schedulerManager.onScheduleShutdownFromExtension(owningExtension, this);
+        }
+    }
+
     /**
      * Gets the current status of the task.
      *
diff --git a/src/main/java/net/minestom/server/timer/TaskBuilder.java b/src/main/java/net/minestom/server/timer/TaskBuilder.java
index d8cf7e6d5..cf527bf1e 100644
--- a/src/main/java/net/minestom/server/timer/TaskBuilder.java
+++ b/src/main/java/net/minestom/server/timer/TaskBuilder.java
@@ -144,27 +144,10 @@ public class TaskBuilder {
      */
     @NotNull
     public Task schedule() {
-        Task task = new Task(
-                this.schedulerManager,
-                this.runnable,
-                this.shutdown,
-                this.delay,
-                this.repeat,
-                this.isTransient,
-                this.owningExtension);
+        Task task = build();
         if (this.shutdown) {
-            Int2ObjectMap<Task> shutdownTasks = this.schedulerManager.shutdownTasks;
-            synchronized (shutdownTasks) {
-                shutdownTasks.put(task.getId(), task);
-            }
-            if (owningExtension != null) {
-                this.schedulerManager.onScheduleShutdownFromExtension(owningExtension, task);
-            }
+            task.scheduleForShutdown();
         } else {
-            Int2ObjectMap<Task> tasks = this.schedulerManager.tasks;
-            synchronized (tasks) {
-                tasks.put(task.getId(), task);
-            }
             task.schedule();
         }
         return task;