Try to prevent async search cache from getting stuck

2 files changed, 54 insertions, 50 deletions

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 f3a4cb2ca..94b00cd05 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
@@ -241,72 +241,74 @@ public class Argument<T, R> {
     
     public static void prepareFilter(Collection<EntryStack<?>> stacks, Collection<ArgumentType<?, ?>> argumentTypes) {
         if (prepareStage != null || currentStages != null) return;
-        prepareStart = Util.getEpochMillis();
-        prepareStacks = stacks;
+        try {
+            prepareStart = Util.getEpochMillis();
+            prepareStacks = stacks;
         prepareStage = new MutablePair<>(0, argumentTypes.size());
         currentStages = new MutablePair[argumentTypes.size()];
-        List<HashedEntryStackWrapper> hashedStacks = CollectionUtils.map(stacks, HashedEntryStackWrapper::new);
-        int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
-        boolean async = ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4;
-        List<CompletableFuture<Long2ObjectMap<Object>>> futures = Lists.newArrayList();
-        List<Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>>> pairs = Lists.newArrayList();
-        
-        for (ArgumentType<?, ?> argumentType : argumentTypes) {
+            List<HashedEntryStackWrapper> hashedStacks = CollectionUtils.map(stacks, HashedEntryStackWrapper::new);
+            int searchPartitionSize = ConfigObject.getInstance().getAsyncSearchPartitionSize();
+            boolean async = ConfigObject.getInstance().shouldAsyncSearch() && stacks.size() > searchPartitionSize * 4;
+            List<CompletableFuture<Long2ObjectMap<Object>>> futures = Lists.newArrayList();
+            List<Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>>> pairs = Lists.newArrayList();
+            
+            for (ArgumentType<?, ?> argumentType : argumentTypes) {
             prepareStage.setLeft(prepareStage.getLeft() + 1);
-            Long2ObjectMap<Object> map = getSearchCache(argumentType);
+                Long2ObjectMap<Object> map = getSearchCache(argumentType);
             MutablePair<Integer, Integer> currentStage = currentStages[prepareStage.getLeft() - 1] = new MutablePair<>(0, hashedStacks.size());
-            
-            if (async) {
-                for (Collection<HashedEntryStackWrapper> partitionStacks : CollectionUtils.partition(hashedStacks, searchPartitionSize)) {
-                    CompletableFuture<Long2ObjectMap<Object>> future = CompletableFuture.supplyAsync(() -> {
-                        Long2ObjectMap<Object> out = new Long2ObjectArrayMap<>(searchPartitionSize + 1);
-                        for (HashedEntryStackWrapper stack : partitionStacks) {
-                            if (map.get(stack.hashExact()) == null) {
-                                Object data = argumentType.cacheData(stack.unwrap());
-                                
-                                if (data != null) {
-                                    out.put(stack.hashExact(), data);
+                
+                if (async) {
+                    for (Collection<HashedEntryStackWrapper> partitionStacks : CollectionUtils.partition(hashedStacks, searchPartitionSize)) {
+                        CompletableFuture<Long2ObjectMap<Object>> future = CompletableFuture.supplyAsync(() -> {
+                            Long2ObjectMap<Object> out = new Long2ObjectArrayMap<>(searchPartitionSize + 1);
+                            for (HashedEntryStackWrapper stack : partitionStacks) {
+                                if (map.get(stack.hashExact()) == null) {
+                                    Object data = argumentType.cacheData(stack.unwrap());
+                                    
+                                    if (data != null) {
+                                        out.put(stack.hashExact(), data);
+                                    }
                                 }
                             }
-                        }
-                        return out;
-                    }).whenComplete((objectLong2ObjectMap, throwable) -> {
+                            return out;
+                        }).whenComplete((objectLong2ObjectMap, throwable) -> {
                         currentStage.setLeft(currentStage.getLeft() + partitionStacks.size());
-                    });
-                    futures.add(future);
-                    pairs.add(Pair.of(argumentType, future));
-                }
-            } else {
-                for (HashedEntryStackWrapper stack : hashedStacks) {
+                        });
+                        futures.add(future);
+                        pairs.add(Pair.of(argumentType, future));
+                    }
+                } else {
+                    for (HashedEntryStackWrapper stack : hashedStacks) {
                     currentStage.setLeft(currentStage.getLeft() + 1);
-                    
-                    if (map.get(stack.hashExact()) == null) {
-                        Object data = argumentType.cacheData(stack.unwrap());
                         
-                        if (data != null) {
-                            map.put(stack.hashExact(), data);
+                        if (map.get(stack.hashExact()) == null) {
+                            Object data = argumentType.cacheData(stack.unwrap());
+                            
+                            if (data != null) {
+                                map.put(stack.hashExact(), data);
+                            }
                         }
                     }
                 }
             }
-        }
-        
-        if (async) {
-            try {
-                CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(10, TimeUnit.SECONDS);
-            } catch (InterruptedException | ExecutionException | TimeoutException e) {
-                e.printStackTrace();
-            }
-            for (Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>> pair : pairs) {
+            
+            if (async) {
+                try {
+                    CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(10, TimeUnit.SECONDS);
+                } catch (InterruptedException | ExecutionException | TimeoutException e) {
+                    e.printStackTrace();
+                }
+                for (Pair<ArgumentType<?, ?>, CompletableFuture<Long2ObjectMap<Object>>> pair : pairs) {
                 Long2ObjectMap<Object> now = pair.getRight().getNow(null);
                 if (now != null) getSearchCache(pair.getLeft()).putAll(now);
+                }
             }
+        } finally {
+            prepareStart = null;
+            prepareStacks = null;
+            prepareStage = null;
+            currentStages = null;
         }
-    
-        prepareStart = null;
-        prepareStacks = null;
-        prepareStage = null;
-        currentStages = null;
     }
     
     public ArgumentType<?, ?> getArgument() {
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/client/SearchFilterPrepareWatcher.java b/runtime/src/main/java/me/shedaniel/rei/plugin/client/SearchFilterPrepareWatcher.java
index 73d182fe2..03d445459 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/client/SearchFilterPrepareWatcher.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/client/SearchFilterPrepareWatcher.java
@@ -57,7 +57,9 @@ public class SearchFilterPrepareWatcher implements HintProvider {
             int currentStageTotal = currentStage == null ? 0 : currentStage.getRight();
                 double prepareStageProgress = prepareStageTotal == 0 ? 0 : prepareStageCurrent / (double) prepareStageTotal;
                 double currentStageProgress = currentStageTotal == 0 ? 0 : currentStageCurrent / (double) currentStageTotal;
-                lastProcess = prepareStageTotal == 0 ? 0 : Math.max(0, prepareStageProgress - (1 - currentStageProgress) / prepareStageTotal);
+                double lastProcess = prepareStageTotal == 0 ? 0 : Math.max(0, prepareStageProgress - (1 - currentStageProgress) / prepareStageTotal);
+                if (lastProcess < 0.05 || lastProcess > 0.95) return Collections.emptyList();
+                this.lastProcess = lastProcess;
                 return ImmutableList.of(new TranslatableComponent("text.rei.caching.search"),
                         new TranslatableComponent("text.rei.caching.search.step", prepareStageCurrent, prepareStageTotal, Math.round(lastProcess * 100)));
             }