MigrationModule.java

package com.hypixel.hytale.server.core.modules.migrations;

import com.hypixel.hytale.common.plugin.PluginManifest;
import com.hypixel.hytale.common.util.CompletableFutureUtil;
import com.hypixel.hytale.component.SystemType;
import com.hypixel.hytale.logger.HytaleLogger;
import com.hypixel.hytale.math.util.ChunkUtil;
import com.hypixel.hytale.server.core.HytaleServer;
import com.hypixel.hytale.server.core.Options;
import com.hypixel.hytale.server.core.event.events.BootEvent;
import com.hypixel.hytale.server.core.plugin.JavaPlugin;
import com.hypixel.hytale.server.core.plugin.JavaPluginInit;
import com.hypixel.hytale.server.core.universe.Universe;
import com.hypixel.hytale.server.core.universe.world.World;
import com.hypixel.hytale.server.core.universe.world.chunk.WorldChunk;
import com.hypixel.hytale.server.core.universe.world.storage.ChunkStore;
import com.hypixel.hytale.server.core.universe.world.storage.IChunkLoader;
import com.hypixel.hytale.server.core.universe.world.storage.IChunkSaver;
import com.hypixel.hytale.server.core.universe.world.storage.component.ChunkSavingSystems;
import it.unimi.dsi.fastutil.longs.LongIterator;
import it.unimi.dsi.fastutil.longs.LongSet;
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
import java.nio.file.Path;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.logging.Level;
import javax.annotation.Nonnull;
import joptsimple.OptionSet;

public class MigrationModule extends JavaPlugin {
   public static final PluginManifest MANIFEST = PluginManifest.corePlugin(MigrationModule.class).build();
   protected static MigrationModule instance;
   @Nonnull
   private final Map<String, Function<Path, Migration>> migrationCtors = new Object2ObjectOpenHashMap<String, Function<Path, Migration>>();
   private SystemType<ChunkStore, ChunkColumnMigrationSystem> chunkColumnMigrationSystem;
   private SystemType<ChunkStore, ChunkSectionMigrationSystem> chunkSectionMigrationSystem;

   public MigrationModule(@Nonnull JavaPluginInit init) {
      super(init);
      instance = this;
   }

   public static MigrationModule get() {
      return instance;
   }

   protected void setup() {
      this.getEventRegistry().register(BootEvent.class, (event) -> {
         if (Options.getOptionSet().has(Options.MIGRATIONS)) {
            this.runMigrations();
            HytaleServer.get().shutdownServer();
         }
      });
      this.chunkColumnMigrationSystem = this.getChunkStoreRegistry().registerSystemType(ChunkColumnMigrationSystem.class);
      this.chunkSectionMigrationSystem = this.getChunkStoreRegistry().registerSystemType(ChunkSectionMigrationSystem.class);
   }

   public SystemType<ChunkStore, ChunkColumnMigrationSystem> getChunkColumnMigrationSystem() {
      return this.chunkColumnMigrationSystem;
   }

   public SystemType<ChunkStore, ChunkSectionMigrationSystem> getChunkSectionMigrationSystem() {
      return this.chunkSectionMigrationSystem;
   }

   public void register(String id, Function<Path, Migration> migration) {
      this.migrationCtors.put(id, migration);
   }

   public void runMigrations() {
      OptionSet optionSet = Options.getOptionSet();
      List<String> worldsToMigrate = optionSet.has(Options.MIGRATE_WORLDS) ? optionSet.valuesOf(Options.MIGRATE_WORLDS) : null;
      Map<String, Path> migrationMap = (Map)Options.getOptionSet().valueOf(Options.MIGRATIONS);
      List<Migration> migrations = new ObjectArrayList<Migration>();
      migrationMap.forEach((s, path) -> {
         Function<Path, Migration> migrationCtor = (Function)this.migrationCtors.get(s);
         if (migrationCtor != null) {
            migrations.add((Migration)migrationCtor.apply(path));
         }
      });
      if (!migrations.isEmpty()) {
         AtomicInteger worldsCount = new AtomicInteger();

         for(World world : Universe.get().getWorlds().values()) {
            String worldName = world.getName();
            if (worldsToMigrate == null || worldsToMigrate.contains(worldName)) {
               worldsCount.incrementAndGet();
               this.getLogger().at(Level.INFO).log("Starting to migrate world '%s'...", worldName);
               ChunkStore chunkComponentStore = world.getChunkStore();
               IChunkSaver saver = chunkComponentStore.getSaver();
               IChunkLoader loader = chunkComponentStore.getLoader();
               world.execute(() -> {
                  ChunkSavingSystems.Data data = (ChunkSavingSystems.Data)chunkComponentStore.getStore().getResource(ChunkStore.SAVE_RESOURCE);
                  data.isSaving = false;
                  data.waitForSavingChunks().whenComplete((aVoid, throwable) -> {
                     try {
                        LongSet chunks = loader.getIndexes();
                        this.getLogger().at(Level.INFO).log("Found %d chunks in world '%s'. Starting iteration...", chunks.size(), worldName);
                        List<CompletableFuture<?>> futures = new ObjectArrayList<CompletableFuture<?>>(chunks.size());
                        LongIterator iterator = chunks.iterator();

                        while(iterator.hasNext()) {
                           long index = iterator.nextLong();
                           int chunkX = ChunkUtil.xOfChunkIndex(index);
                           int chunkZ = ChunkUtil.zOfChunkIndex(index);
                           futures.add(loader.loadHolder(chunkX, chunkZ).thenCompose((holder) -> {
                              if (holder == null) {
                                 return CompletableFuture.completedFuture((Object)null);
                              } else {
                                 WorldChunk chunk = (WorldChunk)holder.getComponent(WorldChunk.getComponentType());
                                 migrations.forEach((migration) -> migration.run(chunk));
                                 return chunk.getNeedsSaving() ? saver.saveHolder(chunkX, chunkZ, holder) : CompletableFuture.completedFuture((Object)null);
                              }
                           }));
                        }

                        try {
                           CompletableFutureUtil.joinWithProgress(futures, (value, completed, max) -> this.getLogger().at(Level.INFO).log("Scanning + Migrating world '%s': %.2d% (%d of %d)", worldName, value, completed, max), 250, 750);
                        } catch (InterruptedException e) {
                           ((HytaleLogger.Api)this.getLogger().at(Level.SEVERE).withCause(e)).log("Interrupted while loading chunks:");
                           Thread.currentThread().interrupt();
                        }
                     } catch (Throwable t) {
                        ((HytaleLogger.Api)this.getLogger().at(Level.SEVERE).withCause(t)).log("Failed to migrate chunks!");
                     } finally {
                        data.isSaving = true;
                        this.getLogger().at(Level.INFO).log("%d world(s) left to migrate.", worldsCount.decrementAndGet());
                     }

                  });
               });
            }
         }

      }
   }
}