Make sure augmented generation doesn't crash the server

Bukkit/src/main/java/com/github/intellectualsites/plotsquared/bukkit/generator/BukkitAugmentedGenerator.java

@@ -27,6 +27,6 @@ public class BukkitAugmentedGenerator extends BlockPopulator {
 
     @Override
     public void populate(@NotNull World world, @NotNull Random random, @NotNull Chunk source) {
-        AugmentedUtils.generate(world.getName(), source.getX(), source.getZ(), null);
+        AugmentedUtils.generate(source, world.getName(), source.getX(), source.getZ(), null);
     }
 }

Bukkit/src/main/java/com/github/intellectualsites/plotsquared/bukkit/util/block/BukkitLocalQueue.java

@@ -18,6 +18,7 @@ import com.sk89q.worldedit.world.block.BlockState;
 import io.papermc.lib.PaperLib;
 import lombok.NonNull;
 import org.bukkit.Bukkit;
+import org.bukkit.Chunk;
 import org.bukkit.Material;
 import org.bukkit.World;
 import org.bukkit.block.Biome;
@@ -25,6 +26,7 @@ import org.bukkit.block.Block;
 import org.bukkit.block.data.BlockData;
 
 import java.util.concurrent.ExecutionException;
+import java.util.function.Consumer;
 
 public class BukkitLocalQueue extends BasicLocalBlockQueue {
 
@@ -88,7 +90,7 @@ public class BukkitLocalQueue extends BasicLocalBlockQueue {
     @Override public final void setComponents(LocalChunk lc)
         throws ExecutionException, InterruptedException {
         setBaseBlocks(lc);
-        if (setBiome()) {
+        if (setBiome() && lc.biomes != null) {
             setBiomes(lc);
         }
     }
@@ -98,8 +100,7 @@ public class BukkitLocalQueue extends BasicLocalBlockQueue {
         if (worldObj == null) {
             throw new NullPointerException("World cannot be null.");
         }
-        PaperLib.getChunkAtAsync(worldObj, localChunk.getX(), localChunk.getZ(), true)
+        final Consumer<Chunk> chunkConsumer = chunk -> {
-            .thenAccept(chunk -> {
             for (int layer = 0; layer < localChunk.baseblocks.length; layer++) {
                 BaseBlock[] blocksLayer = localChunk.baseblocks[layer];
                 if (blocksLayer != null) {
@@ -131,7 +132,22 @@ public class BukkitLocalQueue extends BasicLocalBlockQueue {
                     }
                 }
             }
-            });
+        };
+        if (isForceSync()) {
+            chunkConsumer.accept(getChunk(worldObj, localChunk));
+        } else {
+            PaperLib.getChunkAtAsync(worldObj, localChunk.getX(), localChunk.getZ(), true).thenAccept(chunkConsumer);
+        }
+    }
+
+    private Chunk getChunk(final World world, final LocalChunk localChunk) {
+        if (this.getChunkObject() != null && this.getChunkObject() instanceof Chunk) {
+            final Chunk chunk = (Chunk) this.getChunkObject();
+            if (chunk.getWorld().equals(world) && chunk.getX() == localChunk.getX() && chunk.getZ() == localChunk.getZ()) {
+                return chunk;
+            }
+        }
+        return world.getChunkAt(localChunk.getX(), localChunk.getZ());
     }
 
     private void setMaterial(@NonNull final BlockState plotBlock, @NonNull final Block block) {
@@ -148,7 +164,10 @@ public class BukkitLocalQueue extends BasicLocalBlockQueue {
         if (worldObj == null) {
             throw new NullPointerException("World cannot be null.");
         }
-        PaperLib.getChunkAtAsync(worldObj, lc.getX(), lc.getZ(), true).thenAccept(chunk -> {
+        if (lc.biomes == null) {
+            throw new NullPointerException("Biomes cannot be null.");
+        }
+        final Consumer<Chunk> chunkConsumer = chunk -> {
             for (int x = 0; x < lc.biomes.length; x++) {
                 BiomeType[] biomeZ = lc.biomes[x];
                 if (biomeZ != null) {
@@ -157,13 +176,18 @@ public class BukkitLocalQueue extends BasicLocalBlockQueue {
                             BiomeType biomeType = biomeZ[z];
 
                             Biome biome = BukkitAdapter.adapt(biomeType);
-                            worldObj
+                            worldObj.setBiome((chunk.getX() << 4) + x, (chunk.getZ() << 4) + z,
-                                .setBiome((chunk.getX() << 4) + x, (chunk.getZ() << 4) + z, biome);
+                                biome);
                         }
                     }
                 }
             }
-        });
+        };
+        if (this.isForceSync()) {
+            chunkConsumer.accept(getChunk(worldObj, lc));
+        } else {
+            PaperLib.getChunkAtAsync(worldObj, lc.getX(), lc.getZ(), true).thenAccept(chunkConsumer);
+        }
     }
 
 }

Core/src/main/java/com/github/intellectualsites/plotsquared/plot/commands/Area.java

@@ -147,7 +147,7 @@ public class Area extends SubCommand {
                                             ChunkManager.largeRegionTask(world, region,
                                                 new RunnableVal<BlockVector2>() {
                                                     @Override public void run(BlockVector2 value) {
-                                                        AugmentedUtils.generate(world, value.getX(),
+                                                        AugmentedUtils.generate(null, world, value.getX(),
                                                             value.getZ(), null);
                                                     }
                                                 }, null);
@@ -445,7 +445,7 @@ public class Area extends SubCommand {
                     new RunnableVal<BlockVector2>() {
                         @Override public void run(BlockVector2 value) {
                             AugmentedUtils
-                                .generate(area.getWorldName(), value.getX(), value.getZ(), null);
+                                .generate(null, area.getWorldName(), value.getX(), value.getZ(), null);
                         }
                     }, () -> player.sendMessage("Regen complete"));
                 return true;

Core/src/main/java/com/github/intellectualsites/plotsquared/plot/generator/AugmentedUtils.java

@@ -16,6 +16,7 @@ import com.sk89q.worldedit.world.biome.BiomeType;
 import com.sk89q.worldedit.world.block.BlockState;
 import com.sk89q.worldedit.world.block.BlockTypes;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Set;
 
@@ -29,7 +30,7 @@ public class AugmentedUtils {
         enabled = true;
     }
 
-    public static boolean generate(@NotNull final String world, final int chunkX, final int chunkZ,
+    public static boolean generate(@Nullable Object chunkObject, @NotNull final String world, final int chunkX, final int chunkZ,
         LocalBlockQueue queue) {
         if (!enabled) {
             return false;
@@ -54,6 +55,7 @@ public class AugmentedUtils {
             // Mask
             if (queue == null) {
                 queue = GlobalBlockQueue.IMP.getNewQueue(world, false);
+                queue.setChunkObject(chunkObject);
             }
             LocalBlockQueue primaryMask;
             // coordinates
@@ -102,7 +104,7 @@ public class AugmentedUtils {
                             for (int y = 1; y < 128; y++) {
                                 queue.setBlock(rx, y, rz, air);
                             }
-                            canPlace[x][z] = can;
+                            canPlace[x][z] = true;
                             has = true;
                         }
                     }
@@ -134,6 +136,11 @@ public class AugmentedUtils {
                 }
                 toReturn = true;
             }
+            primaryMask.setChunkObject(chunkObject);
+            primaryMask.setForceSync(true);
+            secondaryMask.setChunkObject(chunkObject);
+            secondaryMask.setForceSync(true);
+
             ScopedLocalBlockQueue scoped = new ScopedLocalBlockQueue(secondaryMask,
                 new Location(area.getWorldName(), blockX, 0, blockZ),
                 new Location(area.getWorldName(), blockX + 15, 255, blockZ + 15));
@@ -141,6 +148,7 @@ public class AugmentedUtils {
             generator.populateChunk(scoped, area);
         }
         if (queue != null) {
+            queue.setForceSync(true);
             queue.flush();
         }
         return toReturn;

Core/src/main/java/com/github/intellectualsites/plotsquared/plot/util/block/BasicLocalBlockQueue.java

@@ -222,6 +222,7 @@ public abstract class BasicLocalBlockQueue extends LocalBlockQueue {
 
 
     public class BasicLocalChunk extends LocalChunk {
+
         public BasicLocalChunk(BasicLocalBlockQueue parent, int x, int z) {
             super(parent, x, z);
             baseblocks = new BaseBlock[16][];

Core/src/main/java/com/github/intellectualsites/plotsquared/plot/util/block/LocalBlockQueue.java

@@ -13,12 +13,18 @@ import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.world.biome.BiomeType;
 import com.sk89q.worldedit.world.block.BaseBlock;
 import com.sk89q.worldedit.world.block.BlockState;
+import lombok.Getter;
+import lombok.Setter;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Map;
 
 public abstract class LocalBlockQueue {
 
+    @Getter @Setter private boolean forceSync = false;
+    @Getter @Setter @Nullable private Object chunkObject;
+
     /**
      * Needed for compatibility with FAWE.
      *