Fix searching interrupts

9 files changed, 159 insertions, 101 deletions

diff --git a/forge/build.gradle b/forge/build.gradle
index 9ac210ac0..8fa846574 100644
--- a/forge/build.gradle
+++ b/forge/build.gradle
@@ -49,12 +49,6 @@ processServerComponentResources {
 }
 
 loom {
-    launches {
-        all {
-            arg "--mixin", "rei-jei-internals-workaround.mixins.json"
-        }
-    }
-
     forge {
         mixinConfig "rei.mixins.json"
     }
@@ -72,6 +66,7 @@ dependencies {
         shadowCommon(project(path: it, configuration: "transformProductionForge")) { transitive false }
     }
 
+    modRuntime("curse.maven:lazy-dfu-forge-460819:3544496")
     // modRuntime("curse.maven:chiselsbits-231095:3492889")
     // modRuntime("curse.maven:jumbofurnace-390880:3120970")
     // modRuntime("curse.maven:cyclic-239286:3221427")
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/BasicFilteringRuleImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/BasicFilteringRuleImpl.java
index a1a8b7558..1d51eb0f2 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/BasicFilteringRuleImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/BasicFilteringRuleImpl.java
@@ -34,19 +34,18 @@ import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.registry.ReloadStage;
 import me.shedaniel.rei.api.common.util.CollectionUtils;
 import me.shedaniel.rei.api.common.util.EntryStacks;
+import me.shedaniel.rei.impl.client.util.ThreadCreator;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.stream.Collectors;
 
 public enum BasicFilteringRuleImpl implements BasicFilteringRule<Pair<LongSet, LongSet>> {
     INSTANCE;
+    private static final ExecutorService EXECUTOR_SERVICE = new ThreadCreator("REI-BasicFiltering").asService();
     private final List<EntryStack<?>> hidden = new ArrayList<>(), shown = new ArrayList<>();
     
     @Override
@@ -73,7 +72,7 @@ public enum BasicFilteringRuleImpl implements BasicFilteringRule<Pair<LongSet, L
                         }
                     }
                     return output;
-                }));
+                }, EXECUTOR_SERVICE));
             }
             try {
                 CompletableFuture.allOf(completableFutures.toArray(new CompletableFuture[0])).get(5, TimeUnit.MINUTES);
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/ManualFilteringRule.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/ManualFilteringRule.java
index dcc1c9500..b7e06da86 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/ManualFilteringRule.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/ManualFilteringRule.java
@@ -32,16 +32,16 @@ import me.shedaniel.rei.api.client.entry.filtering.*;
 import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.util.CollectionUtils;
 import me.shedaniel.rei.api.common.util.EntryStacks;
+import me.shedaniel.rei.impl.client.util.ThreadCreator;
 
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.stream.Collectors;
 
 public class ManualFilteringRule implements FilteringRule<LongSet> {
+    private static final ExecutorService EXECUTOR_SERVICE = new ThreadCreator("REI-ManualFiltering").asService();
+    
     @Override
     public FilteringRuleType<? extends FilteringRule<LongSet>> getType() {
         return ManualFilteringRuleType.INSTANCE;
@@ -61,7 +61,7 @@ public class ManualFilteringRule implements FilteringRule<LongSet> {
                         }
                     }
                     return output;
-                }));
+                }, EXECUTOR_SERVICE));
             }
             try {
                 CompletableFuture.allOf(completableFutures.toArray(new CompletableFuture[0])).get(5, TimeUnit.MINUTES);
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java
index e53330953..7ab70acb3 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/entry/filtering/rules/SearchFilteringRule.java
@@ -30,6 +30,7 @@ import me.shedaniel.rei.api.client.entry.filtering.*;
 import me.shedaniel.rei.api.client.search.SearchFilter;
 import me.shedaniel.rei.api.client.search.SearchProvider;
 import me.shedaniel.rei.api.common.entry.EntryStack;
+import me.shedaniel.rei.impl.client.util.ThreadCreator;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
 import net.minecraft.util.StringUtil;
@@ -37,14 +38,12 @@ import net.minecraft.util.Unit;
 
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.function.Supplier;
 
 @Environment(EnvType.CLIENT)
 public class SearchFilteringRule implements FilteringRule<Unit> {
+    private static final ExecutorService EXECUTOR_SERVICE = new ThreadCreator("REI-SearchFiltering").asService();
     String filterStr;
     Supplier<SearchFilter> filter;
     boolean show;
@@ -102,9 +101,7 @@ public class SearchFilteringRule implements FilteringRule<Unit> {
                     }
                 }
                 return output;
-            }));
+            }, EXECUTOR_SERVICE));
         }
     }
-    
-    
 }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java
index 2eabb955c..d8c11ab1a 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/entrylist/EntryListSearchManager.java
@@ -46,6 +46,7 @@ import org.apache.logging.log4j.Level;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
+import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 
 public class EntryListSearchManager {
@@ -78,12 +79,15 @@ public class EntryListSearchManager {
         searchManager.updateFilter(searchTerm);
         if (searchManager.isDirty()) {
             searchManager.getAsync((list, filter) -> {
+                if (!filter.getFilter().equals(searchTerm)) return;
+                if (searchManager.filter() == null || searchManager.filter() != filter) return;
                 InternalLogger.getInstance().log(ConfigObject.getInstance().doDebugSearchTimeRequired() ? Level.INFO : Level.TRACE, "Search \"%s\" Used [%s]: %s", filter.getFilter(), Thread.currentThread().toString(), stopwatch.toString());
-                List</*EntryStack<?> | CollapsedStack*/ Object> finalList = collapse(copyAndOrder(list));
+                List</*EntryStack<?> | CollapsedStack*/ Object> finalList = collapse(copyAndOrder(list), () -> searchManager.filter() != null && searchManager.filter() == filter);
                 
                 InternalLogger.getInstance().log(ConfigObject.getInstance().doDebugSearchTimeRequired() ? Level.INFO : Level.TRACE, "Search \"%s\" Used and Applied [%s]: %s", filter.getFilter(), Thread.currentThread().toString(), stopwatch.stop().toString());
                 
-                Minecraft.getInstance().executeBlocking(() -> {
+                Minecraft.getInstance().submit(() -> {
+                    if (searchManager.filter() == null || searchManager.filter() != filter) return;
                     update.accept(finalList);
                 });
             });
@@ -104,7 +108,7 @@ public class EntryListSearchManager {
         return list;
     }
     
-    private List</*EntryStack<?> | CollapsedStack*/ Object> collapse(List<EntryStack<?>> stacks) {
+    private List</*EntryStack<?> | CollapsedStack*/ Object> collapse(List<EntryStack<?>> stacks, BooleanSupplier isValid) {
         CollapsibleEntryRegistryImpl collapsibleRegistry = (CollapsibleEntryRegistryImpl) CollapsibleEntryRegistry.getInstance();
         Map<CollapsibleEntryRegistryImpl.Entry, @Nullable CollapsedStack> entries = new HashMap<>();
         
@@ -112,8 +116,12 @@ public class EntryListSearchManager {
             entries.put(entry, null);
         }
         
+        if (!isValid.getAsBoolean()) return List.of();
+        
         List</*EntryStack<?> | CollapsedStack*/ Object> list = new ArrayList<>();
         
+        int i = 0;
+        
         for (EntryStack<?> stack : stacks) {
             long hashExact = EntryStacks.hashExact(stack);
             boolean matchedAny = false;
@@ -138,6 +146,8 @@ public class EntryListSearchManager {
                 }
             }
             
+            if (i++ % 50 == 0 && !isValid.getAsBoolean()) return List.of();
+            
             if (!matchedAny) {
                 list.add(stack);
             }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java
index 0a05c1351..da63bc0c9 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/AsyncSearchManager.java
@@ -24,27 +24,27 @@
 package me.shedaniel.rei.impl.client.search;
 
 import com.google.common.collect.Lists;
+import dev.architectury.platform.Platform;
 import me.shedaniel.rei.api.client.config.ConfigObject;
 import me.shedaniel.rei.api.client.search.SearchFilter;
 import me.shedaniel.rei.api.client.search.SearchProvider;
 import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.util.CollectionUtils;
-import me.shedaniel.rei.impl.common.InternalLogger;
+import me.shedaniel.rei.impl.client.util.ThreadCreator;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.AbstractMap;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.*;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 
 public class AsyncSearchManager {
-    private static final ThreadGroup GROUP = new ThreadGroup("REI-SearchManager");
-    private static final AtomicInteger THREAD_ID = new AtomicInteger(0);
+    private static final ExecutorService EXECUTOR_SERVICE = new ThreadCreator("REI-AsyncSearchManager").asService();
     private final Supplier<List<EntryStack<?>>> stacksProvider;
     private final Supplier<Predicate<EntryStack<?>>> additionalPredicateSupplier;
     private final UnaryOperator<EntryStack<?>> transformer;
@@ -58,34 +58,24 @@ public class AsyncSearchManager {
         this.transformer = transformer;
     }
     
-    private static Thread createThread(Runnable task) {
-        Thread thread = new Thread(GROUP, task, "REI-SearchManager-" + THREAD_ID.getAndIncrement());
-        thread.setDaemon(true);
-        thread.setUncaughtExceptionHandler(($, exception) -> {
-            if (!(exception instanceof InterruptedException) && !(exception instanceof CancellationException) && !(exception instanceof ThreadDeath)) {
-                InternalLogger.getInstance().throwException(exception);
-            }
-        });
-        return thread;
-    }
-    
     public void markDirty() {
         synchronized (AsyncSearchManager.this) {
             this.last = null;
         }
     }
     
-    private record ExecutorTuple(List<Thread> threads, SearchFilter filter, CompletableFuture<Map.Entry<List<EntryStack<?>>, SearchFilter>> future) {
+    @Nullable
+    public SearchFilter filter() {
+        return this.filter;
+    }
+    
+    private record ExecutorTuple(SearchFilter filter, CompletableFuture<Map.Entry<List<EntryStack<?>>, SearchFilter>> future) {
     }
     
     public void updateFilter(String filter) {
         if (this.filter == null || !this.filter.getFilter().equals(filter)) {
             if (this.executor != null) {
-                for (Thread thread : this.executor.threads()) {
-                    try {
-                        thread.stop();
-                    } catch (ThreadDeath ignored) {}
-                }
+                this.executor.future().cancel(Platform.isFabric());
             }
             this.executor = null;
             this.filter = SearchProvider.getInstance().createFilter(filter);
@@ -99,35 +89,20 @@ public class AsyncSearchManager {
     }
     
     public Future<?> getAsync(BiConsumer<List<EntryStack<?>>, SearchFilter> consumer) {
-        if (executor == null || executor.filter() != filter) {
-            if (executor != null) {
-                for (Thread thread : this.executor.threads()) {
-                    try {
-                        thread.stop();
-                    } catch (ThreadDeath ignored) {}
-                }
-                executor = null;
+        if (this.executor == null || this.executor.filter() != filter) {
+            if (this.executor != null) {
+                this.executor.future().cancel(Platform.isFabric());
             }
-            List<Thread> threads = new ArrayList<>();
-            executor = new ExecutorTuple(threads, filter, get(task -> {
-                Thread thread = createThread(task);
-                threads.add(thread);
-                thread.start();
-            }));
+            this.executor = new ExecutorTuple(filter, get(EXECUTOR_SERVICE));
         }
         SearchFilter savedFilter = filter;
-        ExecutorTuple tuple = executor;
-        return (executor = new ExecutorTuple(tuple.threads(), executor.filter(), executor.future().thenApplyAsync(result -> {
+        return (this.executor = new ExecutorTuple(this.executor.filter(), this.executor.future().thenApplyAsync(result -> {
             if (savedFilter == filter) {
                 consumer.accept(result.getKey(), result.getValue());
             }
             
             return result;
-        }, task -> {
-            Thread thread = createThread(task);
-            tuple.threads().add(thread);
-            thread.start();
-        }))).future();
+        }, EXECUTOR_SERVICE))).future();
     }
     
     public List<EntryStack<?>> getNow() {
@@ -163,23 +138,30 @@ public class AsyncSearchManager {
             UnaryOperator<EntryStack<?>> transformer, List<EntryStack<?>> stacks, Map.Entry<List<EntryStack<?>>, SearchFilter> last,
             AsyncSearchManager manager, Executor executor) {
         int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
+        boolean shouldAsync = ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4;
         
         if (!stacks.isEmpty()) {
             CompletableFuture<Void> preparationFuture = CompletableFuture.completedFuture(null);
             
             if (last == null || last.getValue() != filter) {
-                preparationFuture = CompletableFuture.runAsync(() -> {
+                Runnable prepare = () -> {
                     if (manager.filter == filter) {
                         filter.prepareFilter(stacks);
                     } else {
                         throw new CancellationException();
                     }
-                }, executor);
+                };
+                if (shouldAsync) {
+                    preparationFuture = CompletableFuture.runAsync(prepare, executor);
+                } else {
+                    prepare.run();
+                    preparationFuture = CompletableFuture.completedFuture(null);
+                }
             }
             
-            if (ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4) {
+            if (shouldAsync) {
                 List<CompletableFuture<List<EntryStack<?>>>> futures = Lists.newArrayList();
-                for (Iterable<EntryStack<?>> partitionStacks : CollectionUtils.partition(stacks, searchPartitionSize)) {
+                for (Iterable<EntryStack<?>> partitionStacks : CollectionUtils.partition(stacks, Math.max(searchPartitionSize, stacks.size() * 3 / Runtime.getRuntime().availableProcessors()))) {
                     futures.add(CompletableFuture.supplyAsync(() -> {
                         List<EntryStack<?>> filtered = Lists.newArrayList();
                         for (EntryStack<?> stack : partitionStacks) {
@@ -189,7 +171,7 @@ public class AsyncSearchManager {
                             if (manager.filter != filter) throw new CancellationException();
                         }
                         return filtered;
-                    }));
+                    }, executor));
                 }
                 return preparationFuture.thenCompose($ -> CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                                 .orTimeout(30, TimeUnit.SECONDS))
@@ -211,22 +193,16 @@ public class AsyncSearchManager {
                             return new AbstractMap.SimpleImmutableEntry<>(result, filter);
                         });
             } else {
-                return preparationFuture.thenApplyAsync($ -> {
-                            List<EntryStack<?>> list = new ArrayList<>();
-                            
-                            for (EntryStack<?> stack : stacks) {
-                                if (filter.test(stack) && additionalPredicate.test(stack)) {
-                                    list.add(transformer.apply(stack));
-                                }
-                                if (manager.filter != filter) throw new CancellationException();
-                            }
-                            
-                            return list;
-                        }, executor)
-                        .thenApply(result -> {
-                            return new AbstractMap.SimpleImmutableEntry<>(result, filter);
-                        });
+                List<EntryStack<?>> list = new ArrayList<>();
+                
+                for (EntryStack<?> stack : stacks) {
+                    if (filter.test(stack) && additionalPredicate.test(stack)) {
+                        list.add(transformer.apply(stack));
+                    }
+                    if (manager.filter != filter) throw new CancellationException();
+                }
                 
+                return CompletableFuture.completedFuture(new AbstractMap.SimpleImmutableEntry<>(list, filter));
             }
         }
         
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
index 7a54eae64..f38b9c9aa 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/search/argument/Argument.java
@@ -47,6 +47,8 @@ import me.shedaniel.rei.impl.client.search.IntRange;
 import me.shedaniel.rei.impl.client.search.argument.type.ArgumentType;
 import me.shedaniel.rei.impl.client.search.argument.type.ArgumentTypesRegistry;
 import me.shedaniel.rei.impl.client.search.result.ArgumentApplicableResult;
+import me.shedaniel.rei.impl.client.util.ThreadCreator;
+import me.shedaniel.rei.impl.common.InternalLogger;
 import me.shedaniel.rei.impl.common.util.HashedEntryStackWrapper;
 import net.fabricmc.api.EnvType;
 import net.fabricmc.api.Environment;
@@ -60,10 +62,7 @@ import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -71,6 +70,7 @@ import java.util.regex.Pattern;
 @ApiStatus.Internal
 @Environment(EnvType.CLIENT)
 public class Argument<T, R> {
+    private static final ExecutorService EXECUTOR_SERVICE = new ThreadCreator("REI-ArgumentCache").asService();
     public static final Short2ObjectMap<Long2ObjectMap<Object>> SEARCH_CACHE = Short2ObjectMaps.synchronize(new Short2ObjectOpenHashMap<>());
     private static final Object NO_CACHE = new Object();
     private static final AtomicReference<String> lastLanguage = new AtomicReference<>();
@@ -271,7 +271,7 @@ public class Argument<T, R> {
     }
     
     public static Long prepareStart = null;
-    public static Collection<EntryStack<?>> prepareStacks = null;
+    public static List<HashedEntryStackWrapper> prepareStacks = null;
     public static IntIntPair prepareStage = null;
     public static IntIntPair[] currentStages = null;
     
@@ -279,22 +279,34 @@ public class Argument<T, R> {
         if (prepareStage != null || currentStages != null) return;
         try {
             prepareStart = Util.getEpochMillis();
-            prepareStacks = stacks;
+            Long2ObjectMap<Object>[] caches = CollectionUtils.map(argumentTypes, Argument::getSearchCache).toArray(Long2ObjectMap[]::new);
+            prepareStacks = CollectionUtils.mapAndFilter(stacks, stack -> {
+                for (Long2ObjectMap<Object> cache : caches) {
+                    if (!cache.containsKey(stack.hashExact())) {
+                        return true;
+                    }
+                }
+                
+                return false;
+            }, HashedEntryStackWrapper::new);
+            if (prepareStacks.isEmpty()) {
+                return;
+            }
+            InternalLogger.getInstance().trace("Preparing " + prepareStacks.size() + " stacks for search arguments");
             prepareStage = new IntIntMutablePair(0, argumentTypes.size());
             currentStages = new IntIntPair[argumentTypes.size()];
-            List<HashedEntryStackWrapper> hashedStacks = CollectionUtils.map(stacks, HashedEntryStackWrapper::new);
             int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
-            boolean async = ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4;
+            boolean async = ConfigObject.getInstance().shouldAsyncSearch() && prepareStacks.size() > searchPartitionSize * 4;
             List<CompletableFuture<Long2ObjectMap<Object>>> futures = Lists.newArrayList();
             List<Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>>> pairs = Lists.newArrayList();
             
             for (ArgumentType<?, ?> argumentType : argumentTypes) {
                 prepareStage.first(prepareStage.firstInt() + 1);
                 Long2ObjectMap<Object> map = getSearchCache(argumentType);
-                IntIntPair currentStage = currentStages[prepareStage.firstInt() - 1] = new IntIntMutablePair(0, hashedStacks.size());
+                IntIntPair currentStage = currentStages[prepareStage.firstInt() - 1] = new IntIntMutablePair(0, prepareStacks.size());
                 
                 if (async) {
-                    for (Collection<HashedEntryStackWrapper> partitionStacks : CollectionUtils.partition(hashedStacks, searchPartitionSize)) {
+                    for (Collection<HashedEntryStackWrapper> partitionStacks : CollectionUtils.partition(prepareStacks, searchPartitionSize)) {
                         CompletableFuture<Long2ObjectMap<Object>> future = CompletableFuture.supplyAsync(() -> {
                             Long2ObjectMap<Object> out = new Long2ObjectArrayMap<>(searchPartitionSize + 1);
                             for (HashedEntryStackWrapper stack : partitionStacks) {
@@ -307,14 +319,14 @@ public class Argument<T, R> {
                                 }
                             }
                             return out;
-                        }).whenComplete((objectLong2ObjectMap, throwable) -> {
+                        }, EXECUTOR_SERVICE).whenComplete((objectLong2ObjectMap, throwable) -> {
                             currentStage.first(currentStage.firstInt() + partitionStacks.size());
                         });
                         futures.add(future);
                         pairs.add(Pair.of(argumentType, future));
                     }
                 } else {
-                    for (HashedEntryStackWrapper stack : hashedStacks) {
+                    for (HashedEntryStackWrapper stack : prepareStacks) {
                         currentStage.first(currentStage.firstInt() + 1);
                         
                         if (map.get(stack.hashExact()) == null) {
@@ -340,6 +352,8 @@ public class Argument<T, R> {
                     if (now != null) getSearchCache(pair.left()).putAll(now);
                 }
             }
+    
+            InternalLogger.getInstance().debug("Prepared " + prepareStacks.size() + " stacks for search arguments in " + (Util.getEpochMillis() - prepareStart) + "ms");
         } finally {
             prepareStart = null;
             prepareStacks = null;
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/util/ThreadCreator.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/util/ThreadCreator.java
new file mode 100644
index 000000000..6815d7c56
--- /dev/null
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/util/ThreadCreator.java
@@ -0,0 +1,64 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021, 2022 shedaniel
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package me.shedaniel.rei.impl.client.util;
+
+import me.shedaniel.rei.impl.common.InternalLogger;
+
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public final class ThreadCreator {
+    private final ThreadGroup group;
+    private final AtomicInteger threadId = new AtomicInteger(0);
+    
+    public ThreadCreator(String groupName) {
+        this.group = new ThreadGroup(groupName);
+    }
+    
+    public ThreadGroup group() {
+        return group;
+    }
+    
+    public AtomicInteger threadId() {
+        return threadId;
+    }
+    
+    public Thread create(Runnable task) {
+        Thread thread = new Thread(this.group(), task, this.group().getName() + "-" + this.threadId().getAndIncrement());
+        thread.setDaemon(true);
+        thread.setUncaughtExceptionHandler(($, exception) -> {
+            if (!(exception instanceof InterruptedException) && !(exception instanceof CancellationException) && !(exception instanceof ThreadDeath)) {
+                InternalLogger.getInstance().throwException(exception);
+            }
+        });
+        return thread;
+    }
+    
+    public ExecutorService asService() {
+        return new ThreadPoolExecutor(0, Runtime.getRuntime().availableProcessors() * 4,
+                0L, TimeUnit.SECONDS,
+                new SynchronousQueue<>(),
+                this::create);
+    }
+}
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java b/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java
index 74b0f2806..6ed68874e 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/test/REITestPlugin.java
@@ -64,7 +64,7 @@ public class REITestPlugin implements REIClientPlugin {
     
     @Override
     public void registerEntries(EntryRegistry registry) {
-        int times = 100;
+        int times = 10;
         for (Item item : Registry.ITEM) {
             EntryStack<ItemStack> base = EntryStacks.of(item);
             registry.addEntriesAfter(base, IntStream.range(0, times).mapToObj(value -> transformStack(EntryStacks.of(item))).collect(Collectors.toList()));
@@ -79,7 +79,10 @@ public class REITestPlugin implements REIClientPlugin {
     
     @Override
     public void registerCollapsibleEntries(CollapsibleEntryRegistry registry) {
+        int i = 0;
         for (Item item : Registry.ITEM) {
+            if (i++ % 10 != 0)
+                continue;
             registry.group(Registry.ITEM.getKey(item), new TextComponent(Registry.ITEM.getKey(item).toString()),
                     stack -> stack.getType() == VanillaEntryTypes.ITEM && stack.<ItemStack>castValue().is(item));
         }