Avoid threading issues

AlexProgrammerDE · Dec 16, 2024 · b179d0d · b179d0d
1 parent ec8208a
commit b179d0d
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 34 deletions.
diff --git a/server/src/main/java/com/soulfiremc/server/plugins/AutoArmor.java b/server/src/main/java/com/soulfiremc/server/plugins/AutoArmor.java
@@ -101,7 +101,7 @@ private static void putOn(
       return;
     }
 
-    inventoryManager.connection().botControl().maybeRegister(() -> ControllingTask.staged(List.of(
+    inventoryManager.connection().botControl().maybeRegister(ControllingTask.staged(List.of(
       new ControllingTask.RunnableStage(inventoryManager::openPlayerInventory),
       new ControllingTask.RunnableStage(() -> inventoryManager.leftClickSlot(bestItemSlot)),
       new ControllingTask.WaitDelayStage(() -> 50L),

diff --git a/server/src/main/java/com/soulfiremc/server/plugins/AutoEat.java b/server/src/main/java/com/soulfiremc/server/plugins/AutoEat.java
@@ -81,13 +81,13 @@ public static void onJoined(BotJoinedEvent event) {
         }
 
         if (!playerInventory.isHeldItem(slot) && playerInventory.isHotbar(slot)) {
-          inventoryManager.connection().botControl().maybeRegister(() -> ControllingTask.staged(List.of(
+          inventoryManager.connection().botControl().maybeRegister(ControllingTask.staged(List.of(
             new ControllingTask.RunnableStage(() -> inventoryManager.changeHeldItem(playerInventory.toHotbarIndex(slot))),
             new ControllingTask.WaitDelayStage(() -> 50L),
             new ControllingTask.RunnableStage(() -> connection.botActionManager().useItemInHand(Hand.MAIN_HAND))
           )));
         } else if (playerInventory.isMainInventory(slot)) {
-          inventoryManager.connection().botControl().maybeRegister(() -> ControllingTask.staged(List.of(
+          inventoryManager.connection().botControl().maybeRegister(ControllingTask.staged(List.of(
             new ControllingTask.RunnableStage(inventoryManager::openPlayerInventory),
             new ControllingTask.RunnableStage(() -> inventoryManager.leftClickSlot(slot)),
             new ControllingTask.WaitDelayStage(() -> 50L),

diff --git a/server/src/main/java/com/soulfiremc/server/plugins/AutoTotem.java b/server/src/main/java/com/soulfiremc/server/plugins/AutoTotem.java
@@ -79,7 +79,7 @@ public static void onJoined(BotJoinedEvent event) {
           return;
         }
 
-        inventoryManager.connection().botControl().maybeRegister(() -> ControllingTask.staged(List.of(
+        inventoryManager.connection().botControl().maybeRegister(ControllingTask.staged(List.of(
           new ControllingTask.RunnableStage(inventoryManager::openPlayerInventory),
           new ControllingTask.RunnableStage(() -> inventoryManager.leftClickSlot(slot)),
           new ControllingTask.WaitDelayStage(() -> 50L),

diff --git a/server/src/main/java/com/soulfiremc/server/protocol/bot/BotControlAPI.java b/server/src/main/java/com/soulfiremc/server/protocol/bot/BotControlAPI.java
@@ -52,7 +52,7 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Comparator;
-import java.util.function.Supplier;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * This class is used to control the bot. The goal is to reduce friction for doing simple things.
@@ -62,63 +62,60 @@
 public class BotControlAPI {
   private final BotConnection connection;
   private final SecureRandom secureRandom = new SecureRandom();
+  private final AtomicReference<ControllingTask> controllingTask = new AtomicReference<>();
   @Getter
   @Setter
   private int attackCooldownTicks = 0;
-  private ControllingTask controllingTask;
 
   public void tick() {
     if (attackCooldownTicks > 0) {
       attackCooldownTicks--;
     }
 
-    if (controllingTask != null) {
-      controllingTask.tick();
+    var localTask = this.controllingTask.get();
+    if (localTask != null) {
+      localTask.tick();
 
-      if (controllingTask.isDone()) {
-        controllingTask.stop();
-        controllingTask = null;
+      if (localTask.isDone()) {
+        localTask.stop();
+        this.controllingTask.set(null);
       }
     }
   }
 
   public boolean stopControllingTask() {
-    if (controllingTask != null) {
-      controllingTask.stop();
-      controllingTask = null;
-
-      return true;
-    }
-
-    return false;
-  }
+    return this.controllingTask.updateAndGet(
+      current -> {
+        if (current != null) {
+          current.stop();
+          return null;
+        }
 
-  public boolean acceptsTask() {
-    return controllingTask == null;
+        return null;
+      }) != null;
   }
 
   public boolean activelyControlled() {
-    return controllingTask != null;
+    return this.controllingTask.get() != null;
   }
 
   public void registerControllingTask(ControllingTask task) {
-    if (controllingTask != null) {
-      controllingTask.stop();
-    }
+    this.controllingTask.updateAndGet(
+      current -> {
+        if (current != null) {
+          current.stop();
+        }
 
-    controllingTask = task;
+        return task;
+      });
   }
 
   public void unregisterControllingTask(ControllingTask task) {
-    if (controllingTask == task) {
-      controllingTask = null;
-    }
+    this.controllingTask.compareAndSet(task, null);
   }
 
-  public void maybeRegister(Supplier<ControllingTask> taskSupplier) {
-    if (acceptsTask()) {
-      registerControllingTask(taskSupplier.get());
-    }
+  public void maybeRegister(ControllingTask task) {
+    this.controllingTask.compareAndSet(null, task);
   }
 
   public boolean toggleFlight() {