XFEServerManager

package com.xfestudio.xfeservermanager.core.rollback;

import com.xfestudio.xfeservermanager.api.world.WorldChangeRecord;
import com.xfestudio.xfeservermanager.api.world.WorldStatePayload;
import java.time.Clock;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;

/** Incremental, checkpointed executor intended to be called once per server tick. */
public final class RollbackCoordinator {
    private final WorldAccessPort world;
    private final RollbackCheckpointPort checkpoints;
    private final Clock clock;
    private final NanoClock nanoClock;
    private final Map<UUID, RollbackOperation> operations = new HashMap<>();
    private final Map<UUID, Set<RollbackPlanner.ChunkCoordinate>> acquiredChunks = new HashMap<>();

    public RollbackCoordinator(WorldAccessPort world, RollbackCheckpointPort checkpoints,
                               Clock clock, NanoClock nanoClock) {
        this.world = Objects.requireNonNull(world, "world");
        this.checkpoints = Objects.requireNonNull(checkpoints, "checkpoints");
        this.clock = Objects.requireNonNull(clock, "clock");
        this.nanoClock = Objects.requireNonNull(nanoClock, "nanoClock");
        for (RollbackOperation loaded : checkpoints.loadIncomplete()) {
            RollbackOperation recovered = loaded.status() == RollbackStatus.RUNNING
                    ? copy(loaded, RollbackStatus.PAUSED, loaded.cursor(), loaded.appliedEventIds(),
                            loaded.failures(), loaded.inFlightEventId())
                    : loaded;
            checkpoints.save(recovered);
            operations.put(recovered.operationId(), recovered);
        }
    }

    public synchronized RollbackOperation begin(RollbackPlan plan) {
        var operation = new RollbackOperation(UUID.randomUUID(), plan, RollbackStatus.PREVIEWED, 0,
                List.of(), List.of(), clock.instant());
        return store(operation);
    }

    public synchronized Optional<RollbackOperation> find(UUID operationId) {
        return Optional.ofNullable(operations.get(operationId));
    }

    public synchronized RollbackOperation run(UUID operationId) {
        RollbackOperation current = require(operationId);
        if (current.status() != RollbackStatus.PREVIEWED && current.status() != RollbackStatus.PAUSED) {
            throw new IllegalStateException("operation cannot run from " + current.status());
        }
        return store(copy(current, RollbackStatus.RUNNING, current.cursor(), current.appliedEventIds(),
                current.failures(), current.inFlightEventId()));
    }

    public synchronized RollbackOperation pause(UUID operationId) {
        RollbackOperation current = require(operationId);
        if (current.status() != RollbackStatus.RUNNING) throw new IllegalStateException("only running operations can pause");
        RollbackOperation paused = store(copy(current, RollbackStatus.PAUSED, current.cursor(),
                current.appliedEventIds(), current.failures(), current.inFlightEventId()));
        releaseAll(operationId);
        return paused;
    }

    public synchronized RollbackOperation cancel(UUID operationId) {
        RollbackOperation current = require(operationId);
        if (isTerminal(current.status())) throw new IllegalStateException("operation is already terminal");
        if (current.inFlightEventId() != null) current = reconcileInFlight(current);
        RollbackOperation cancelled = store(copy(current, RollbackStatus.CANCELLED, current.cursor(),
                current.appliedEventIds(), current.failures(), null));
        releaseAll(cancelled.operationId());
        return cancelled;
    }

    public synchronized RollbackOperation tick(UUID operationId, TickBudget budget) {
        RollbackOperation current = require(operationId);
        if (current.status() != RollbackStatus.RUNNING) return current;
        try {
            if (current.inFlightEventId() != null) current = reconcileInFlight(current);
            long started = nanoClock.nanoTime();
            int cursor = current.cursor();
            var applied = new ArrayList<>(current.appliedEventIds());
            var failures = new ArrayList<>(current.failures());
            var visitedChunks = new HashSet<RollbackPlanner.ChunkCoordinate>();

            while (cursor < current.plan().items().size()) {
                if (nanoClock.nanoTime() - started >= budget.maxNanos()) break;
                RollbackPlanItem item = current.plan().items().get(cursor);
                if (!item.executable(current.plan().force())) {
                    cursor++;
                    continue;
                }
                RollbackPlanner.ChunkCoordinate chunk = RollbackPlanner.ChunkCoordinate.of(item.entry().change());
                if (!visitedChunks.contains(chunk) && visitedChunks.size() >= budget.maxDistinctChunks()) break;
                visitedChunks.add(chunk);
                if (!ensureLoaded(current, chunk)) {
                    failures.add(failure(item, "CHUNK_LOAD_FAILED", "temporary chunk load was refused"));
                    cursor++;
                    continue;
                }

                // Persist intent before mutating the world. If the process dies after
                // apply but before the next checkpoint, reconcileInFlight recognizes
                // the target state and retains redo provenance.
                current = store(copy(current, RollbackStatus.RUNNING, cursor, applied, failures,
                        item.entry().change().eventId()));
                applyOne(current, item, applied, failures);
                cursor++;
                current = store(copy(current, RollbackStatus.RUNNING, cursor, applied, failures, null));
            }

            RollbackStatus status = cursor == current.plan().items().size()
                    ? RollbackStatus.COMPLETED : RollbackStatus.RUNNING;
            return store(copy(current, status, cursor, applied, failures, null));
        } finally {
            // Tickets are scoped to one tick. This bounds leaks on pause, failure,
            // checkpoint exceptions, or construction of a replacement coordinator.
            releaseAll(operationId);
        }
    }

    private RollbackOperation reconcileInFlight(RollbackOperation operation) {
        int cursor = operation.cursor();
        RollbackPlanItem item = operation.plan().items().get(cursor);
        WorldStatePayload actual = world.currentState(item.entry().change());
        var applied = new ArrayList<>(operation.appliedEventIds());
        var failures = new ArrayList<>(operation.failures());
        if (Objects.equals(actual, item.targetState())) {
            if (!applied.contains(operation.inFlightEventId())) applied.add(operation.inFlightEventId());
            cursor++;
        } else if (!operation.plan().force() && !Objects.equals(actual, item.expectedCurrent())) {
            failures.add(failure(item, "CONFLICT_AFTER_RECOVERY",
                    "state differs from both the journal expectation and rollback target"));
            cursor++;
        }
        return store(copy(operation, operation.status(), cursor, applied, failures, null));
    }

    private void applyOne(RollbackOperation operation, RollbackPlanItem item,
                          List<UUID> applied, List<RollbackFailure> failures) {
        WorldChangeRecord change = item.entry().change();
        WorldStatePayload actual = world.currentState(change);
        if (!operation.plan().force() && !Objects.equals(actual, item.expectedCurrent())) {
            failures.add(failure(item, "CONFLICT", "current state changed after preview"));
            return;
        }
        var compatibility = world.compatibility(change, item.targetState());
        if (!compatibility.compatible()) {
            failures.add(failure(item, "INCOMPATIBLE", compatibility.detail()));
            return;
        }
        try {
            world.apply(change, item.targetState());
            applied.add(change.eventId());
        } catch (WorldMutationException exception) {
            failures.add(failure(item, "APPLY_FAILED", exception.getMessage()));
        }
    }

    private boolean ensureLoaded(RollbackOperation operation, RollbackPlanner.ChunkCoordinate chunk) {
        if (world.isChunkLoaded(chunk.dimensionId(), chunk.x(), chunk.z())) return true;
        if (!operation.plan().loadUnloadedChunks()) return false;
        Set<RollbackPlanner.ChunkCoordinate> held = acquiredChunks.computeIfAbsent(operation.operationId(), ignored -> new HashSet<>());
        if (held.contains(chunk)) return true;
        if (held.size() >= RollbackPlanner.MAX_TEMPORARY_CHUNKS) return false;
        if (!world.acquireChunk(chunk.dimensionId(), chunk.x(), chunk.z())) return false;
        held.add(chunk);
        return true;
    }

    private void releaseAll(UUID operationId) {
        for (RollbackPlanner.ChunkCoordinate chunk : acquiredChunks.getOrDefault(operationId, Set.of())) {
            world.releaseChunk(chunk.dimensionId(), chunk.x(), chunk.z());
        }
        acquiredChunks.remove(operationId);
    }

    private RollbackOperation store(RollbackOperation operation) {
        checkpoints.save(operation);
        operations.put(operation.operationId(), operation);
        return operation;
    }

    private RollbackOperation require(UUID operationId) {
        RollbackOperation operation = operations.get(operationId);
        if (operation == null) throw new IllegalArgumentException("unknown rollback operation " + operationId);
        return operation;
    }

    private RollbackOperation copy(RollbackOperation operation, RollbackStatus status, int cursor,
                                   List<UUID> applied, List<RollbackFailure> failures) {
        return copy(operation, status, cursor, applied, failures, null);
    }

    private RollbackOperation copy(RollbackOperation operation, RollbackStatus status, int cursor,
                                   List<UUID> applied, List<RollbackFailure> failures, UUID inFlightEventId) {
        return new RollbackOperation(operation.operationId(), operation.plan(), status, cursor,
                applied, failures, inFlightEventId, clock.instant());
    }

    private static RollbackFailure failure(RollbackPlanItem item, String code, String detail) {
        return new RollbackFailure(item.entry().change().eventId(), code, detail);
    }

    private static boolean isTerminal(RollbackStatus status) {
        return status == RollbackStatus.CANCELLED || status == RollbackStatus.COMPLETED || status == RollbackStatus.FAILED;
    }
}