minecraft-source/src/net/minecraft/server/level/ChunkTaskPriorityQueueSorter.java

package net.minecraft.server.level;

import org.apache.logging.log4j.LogManager;
import com.google.common.annotations.VisibleForTesting;
import com.mojang.datafixers.util.Either;
import java.util.stream.Stream;
import net.minecraft.Util;
import java.util.stream.Collector;
import java.util.concurrent.CompletableFuture;
import java.util.Optional;
import java.util.function.IntConsumer;
import net.minecraft.world.level.ChunkPos;
import java.util.function.IntSupplier;
import com.google.common.collect.Sets;
import java.util.stream.Collectors;
import java.util.concurrent.Executor;
import java.util.List;
import net.minecraft.util.thread.StrictQueue;
import net.minecraft.util.thread.ProcessorMailbox;
import java.util.Set;
import net.minecraft.util.Unit;
import java.util.function.Function;
import net.minecraft.util.thread.ProcessorHandle;
import java.util.Map;
import org.apache.logging.log4j.Logger;

public class ChunkTaskPriorityQueueSorter implements AutoCloseable, ChunkHolder.LevelChangeListener {
    private static final Logger LOGGER;
    private final Map<ProcessorHandle<?>, ChunkTaskPriorityQueue<? extends Function<ProcessorHandle<Unit>, ?>>> queues;
    private final Set<ProcessorHandle<?>> sleeping;
    private final ProcessorMailbox<StrictQueue.IntRunnable> mailbox;
    
    public ChunkTaskPriorityQueueSorter(final List<ProcessorHandle<?>> list, final Executor executor, final int integer) {
        final ChunkTaskPriorityQueue chunkTaskPriorityQueue;
        this.queues = list.stream().collect(Collectors.toMap(Function.identity(), air -> {
            new ChunkTaskPriorityQueue(air.name() + "_queue", integer);
            return chunkTaskPriorityQueue;
        }));
        this.sleeping = Sets.newHashSet(list);
        this.mailbox = new ProcessorMailbox<StrictQueue.IntRunnable>(new StrictQueue.FixedPriorityQueue(4), executor, "sorter");
    }
    
    public static Message<Runnable> message(final Runnable runnable, final long long2, final IntSupplier intSupplier) {
        return new Message<Runnable>(air -> () -> {
            runnable.run();
            air.tell(Unit.INSTANCE);
        }, long2, intSupplier);
    }
    
    public static Message<Runnable> message(final ChunkHolder wo, final Runnable runnable) {
        return message(runnable, wo.getPos().toLong(), wo::getQueueLevel);
    }
    
    public static Release release(final Runnable runnable, final long long2, final boolean boolean3) {
        return new Release(runnable, long2, boolean3);
    }
    
    public <T> ProcessorHandle<Message<T>> getProcessor(final ProcessorHandle<T> air, final boolean boolean2) {
        return this.mailbox.<ProcessorHandle<Message<T>>>ask(air3 -> new StrictQueue.IntRunnable(0, () -> {
            this.<T>getQueue(air);
            air3.tell(ProcessorHandle.of("chunk priority sorter around " + air.name(), a -> this.<T>submit(air, a.task, a.pos, a.level, boolean2)));
        })).join();
    }
    
    public ProcessorHandle<Release> getReleaseProcessor(final ProcessorHandle<Runnable> air) {
        return this.mailbox.<ProcessorHandle<Release>>ask(air2 -> new StrictQueue.IntRunnable(0, () -> air2.tell(ProcessorHandle.of("chunk priority sorter around " + air.name(), b -> this.<Runnable>release(air, b.pos, b.task, b.clearQueue))))).join();
    }
    
    @Override
    public void onLevelChange(final ChunkPos bje, final IntSupplier intSupplier, final int integer, final IntConsumer intConsumer) {
        final int integer2;
        this.mailbox.tell(new StrictQueue.IntRunnable(0, () -> {
            integer2 = intSupplier.getAsInt();
            this.queues.values().forEach(wq -> wq.resortChunkTasks(integer2, bje, integer));
            intConsumer.accept(integer);
        }));
    }
    
    private <T> void release(final ProcessorHandle<T> air, final long long2, final Runnable runnable, final boolean boolean4) {
        final ChunkTaskPriorityQueue<Function<ProcessorHandle<Unit>, Object>> wq7;
        this.mailbox.tell(new StrictQueue.IntRunnable(1, () -> {
            wq7 = this.getQueue(air);
            wq7.release(long2, boolean4);
            if (this.sleeping.remove(air)) {
                this.pollTask(wq7, (ProcessorHandle<Object>)air);
            }
            runnable.run();
        }));
    }
    
    private <T> void submit(final ProcessorHandle<T> air, final Function<ProcessorHandle<Unit>, T> function, final long long3, final IntSupplier intSupplier, final boolean boolean5) {
        final ChunkTaskPriorityQueue<Function<ProcessorHandle<Unit>, Object>> wq8;
        final int integer9;
        this.mailbox.tell(new StrictQueue.IntRunnable(2, () -> {
            wq8 = this.getQueue(air);
            integer9 = intSupplier.getAsInt();
            wq8.submit(Optional.<Function<ProcessorHandle<Unit>, Object>>of((Function<ProcessorHandle<Unit>, Object>)function), long3, integer9);
            if (boolean5) {
                wq8.submit(Optional.<Function<ProcessorHandle<Unit>, Object>>empty(), long3, integer9);
            }
            if (this.sleeping.remove(air)) {
                this.pollTask(wq8, (ProcessorHandle<Object>)air);
            }
        }));
    }
    
    private <T> void pollTask(final ChunkTaskPriorityQueue<Function<ProcessorHandle<Unit>, T>> wq, final ProcessorHandle<T> air) {
        final Stream<Either<Function<ProcessorHandle<Unit>, Object>, Runnable>> stream4;
        this.mailbox.tell(new StrictQueue.IntRunnable(3, () -> {
            stream4 = (Stream<Either<Function<ProcessorHandle<Unit>, Object>, Runnable>>)wq.pop();
            if (stream4 == null) {
                this.sleeping.add(air);
            }
            else {
                Util.sequence(stream4.map(either -> (CompletableFuture)either.map((Function)air::ask, runnable -> {
                    runnable.run();
                    return CompletableFuture.<Unit>completedFuture(Unit.INSTANCE);
                })).collect(Collectors.toList())).thenAccept(list -> this.pollTask((ChunkTaskPriorityQueue<Function<ProcessorHandle<Unit>, Object>>)wq, (ProcessorHandle<Object>)air));
            }
        }));
    }
    
    private <T> ChunkTaskPriorityQueue<Function<ProcessorHandle<Unit>, T>> getQueue(final ProcessorHandle<T> air) {
        final ChunkTaskPriorityQueue<? extends Function<ProcessorHandle<Unit>, ?>> wq3 = this.queues.get(air);
        if (wq3 == null) {
            throw Util.<IllegalArgumentException>pauseInIde(new IllegalArgumentException("No queue for: " + air));
        }
        return (ChunkTaskPriorityQueue<Function<ProcessorHandle<Unit>, T>>)wq3;
    }
    
    @VisibleForTesting
    public String getDebugStatus() {
        return this.queues.entrySet().stream().map(entry -> entry.getKey().name() + "=[" + ((ChunkTaskPriorityQueue)entry.getValue()).getAcquired().stream().map(long1 -> long1 + ":" + new ChunkPos(long1)).collect(Collectors.joining(",")) + "]").collect(Collectors.joining(",")) + ", s=" + this.sleeping.size();
    }
    
    @Override
    public void close() {
        this.queues.keySet().forEach(ProcessorHandle::close);
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
    
    public static final class Message<T> {
        private final Function<ProcessorHandle<Unit>, T> task;
        private final long pos;
        private final IntSupplier level;
        
        private Message(final Function<ProcessorHandle<Unit>, T> function, final long long2, final IntSupplier intSupplier) {
            this.task = function;
            this.pos = long2;
            this.level = intSupplier;
        }
    }
    
    public static final class Release {
        private final Runnable task;
        private final long pos;
        private final boolean clearQueue;
        
        private Release(final Runnable runnable, final long long2, final boolean boolean3) {
            this.task = runnable;
            this.pos = long2;
            this.clearQueue = boolean3;
        }
    }
}