Add time threshold if suggestions are taking too much time

Author: ahornace

suggester/src/main/java/org/opengrok/suggest/Suggester.java

@@ -38,13 +38,15 @@
 import java.nio.file.Path;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -61,6 +63,8 @@ public final class Suggester implements Closeable {
 
     private static final String PROJECTS_DISABLED_KEY = "";
 
+    private static final int MAX_TIME_MS = 1000;
+
     private static final Logger logger = Logger.getLogger(Suggester.class.getName());
 
     private final Map<String, SuggesterProjectData> projectData = new ConcurrentHashMap<>();
@@ -79,6 +83,8 @@ public final class Suggester implements Closeable {
 
     private final Set<String> allowedFields;
 
+    private final ExecutorService executorService = Executors.newWorkStealingPool();
+
     /**
      * @param suggesterDir directory under which the suggester data should be created
      * @param resultSize maximum number of suggestions that should be returned
@@ -129,13 +135,13 @@ public void init(final Collection<NamedIndexDir> luceneIndexes) {
         synchronized (lock) {
             logger.log(Level.INFO, "Initializing suggester");
 
-            ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
+            ExecutorService executor = Executors.newWorkStealingPool();
 
             for (NamedIndexDir indexDir : luceneIndexes) {
-                submitInitIfIndexExists(executorService, indexDir);
+                submitInitIfIndexExists(executor, indexDir);
             }
 
-            shutdownAndAwaitTermination(executorService, "Suggester successfully initialized");
+            shutdownAndAwaitTermination(executor, "Suggester successfully initialized");
         }
     }
 
@@ -195,6 +201,7 @@ private void shutdownAndAwaitTermination(final ExecutorService executorService,
             logger.log(Level.INFO, logMessageOnSuccess);
         } catch (InterruptedException e) {
             logger.log(Level.SEVERE, "Interrupted while building suggesters", e);
+            Thread.currentThread().interrupt();
         }
     }
 
@@ -211,18 +218,18 @@ public void rebuild(final Collection<NamedIndexDir> indexDirs) {
         synchronized (lock) {
             logger.log(Level.INFO, "Rebuilding following suggesters: {0}", indexDirs);
 
-            ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
+            ExecutorService executor = Executors.newWorkStealingPool();
 
             for (NamedIndexDir indexDir : indexDirs) {
                 SuggesterProjectData data = this.projectData.get(indexDir.name);
                 if (data != null) {
-                    executorService.submit(getRebuildRunnable(data));
+                    executor.submit(getRebuildRunnable(data));
                 } else {
-                    submitInitIfIndexExists(executorService, indexDir);
+                    submitInitIfIndexExists(executor, indexDir);
                 }
             }
 
-            shutdownAndAwaitTermination(executorService, "Suggesters for " + indexDirs + " were successfully rebuilt");
+            shutdownAndAwaitTermination(executor, "Suggesters for " + indexDirs + " were successfully rebuilt");
         }
     }
 
@@ -287,8 +294,21 @@ public List<LookupResultItem> search(
                     indexReaders.get(0).getReader()));
         }
 
-        List<LookupResultItem> results = readers.parallelStream().flatMap(namedIndexReader -> {
+        List<LookupResultItem> results;
+        if (!SuggesterUtils.isComplexQuery(query, suggesterQuery)) { // use WFST for lone prefix
+            results = prefixLookup(readers, (SuggesterPrefixQuery) suggesterQuery);
+        } else {
+            results = complexLookup(readers, suggesterQuery, query);
+        }
+
+        return SuggesterUtils.combineResults(results, resultSize);
+    }
 
+    private List<LookupResultItem> prefixLookup(
+            final List<NamedIndexReader> readers,
+            final SuggesterPrefixQuery suggesterQuery
+    ) {
+        return readers.parallelStream().flatMap(namedIndexReader -> {
             SuggesterProjectData data = projectData.get(namedIndexReader.name);
             if (data == null) {
                 logger.log(Level.FINE, "{0} not yet initialized", namedIndexReader.name);
@@ -300,26 +320,45 @@ public List<LookupResultItem> search(
             }
 
             try {
-                if (!SuggesterUtils.isComplexQuery(query, suggesterQuery)) { // use WFST for lone prefix
-                    String prefix = ((SuggesterPrefixQuery) suggesterQuery).getPrefix().text();
-
-                    return data.lookup(suggesterQuery.getField(), prefix, resultSize)
-                            .stream()
-                            .map(item -> new LookupResultItem(item.key.toString(), namedIndexReader.name, item.value));
-                } else {
-                    SuggesterSearcher searcher = new SuggesterSearcher(namedIndexReader.reader, resultSize);
+                String prefix = suggesterQuery.getPrefix().text();
 
-                    List<LookupResultItem> resultItems = searcher.suggest(query, namedIndexReader.name, suggesterQuery,
-                            data.getSearchCounts(suggesterQuery.getField()));
-
-                    return resultItems.stream();
-                }
+                return data.lookup(suggesterQuery.getField(), prefix, resultSize)
+                        .stream()
+                        .map(item -> new LookupResultItem(item.key.toString(), namedIndexReader.name, item.value));
             } finally {
                 data.unlock();
             }
         }).collect(Collectors.toList());
+    }
 
-        return SuggesterUtils.combineResults(results, resultSize);
+    private List<LookupResultItem> complexLookup(
+            final List<NamedIndexReader> readers,
+            final SuggesterQuery suggesterQuery,
+            final Query query
+    ) {
+        List<LookupResultItem> results = new ArrayList<>(readers.size() * resultSize);
+        List<SuggesterSearchTask> searchTasks = new ArrayList<>(readers.size());
+        for (NamedIndexReader ir : readers) {
+            searchTasks.add(new SuggesterSearchTask(ir, query, suggesterQuery, results));
+        }
+
+        try {
+            executorService.invokeAll(searchTasks, MAX_TIME_MS, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            logger.log(Level.WARNING, "Interrupted while invoking suggester search", e);
+            Thread.currentThread().interrupt();
+        }
+
+        // wait for tasks to finish
+        for (SuggesterSearchTask searchTask : searchTasks) {
+            if (!searchTask.started) {
+                continue;
+            }
+            // "spin lock" – should be fast since all the tasks either finished or were interrupted
+            while (!searchTask.finished) {
+            }
+        }
+        return results;
     }
 
     /**
@@ -422,6 +461,7 @@ public List<Entry<BytesRef, Integer>> getSearchCounts(
      */
     @Override
     public void close() {
+        executorService.shutdownNow();
         projectData.values().forEach(f -> {
             try {
                 f.close();
@@ -431,6 +471,62 @@ public void close() {
         });
     }
 
+    private class SuggesterSearchTask implements Callable<Void> {
+
+        private final NamedIndexReader namedIndexReader;
+        private final Query query;
+        private final SuggesterQuery suggesterQuery;
+        private final List<LookupResultItem> results;
+
+        private volatile boolean finished = false;
+        private volatile boolean started = false;
+
+        SuggesterSearchTask(
+                final NamedIndexReader namedIndexReader,
+                final Query query,
+                final SuggesterQuery suggesterQuery,
+                final List<LookupResultItem> results
+        ) {
+            this.namedIndexReader = namedIndexReader;
+            this.query = query;
+            this.suggesterQuery = suggesterQuery;
+            this.results = results;
+        }
+
+        @Override
+        public Void call() {
+            try {
+                started = true;
+
+                SuggesterProjectData data = projectData.get(namedIndexReader.name);
+                if (data == null) {
+                    logger.log(Level.FINE, "{0} not yet initialized", namedIndexReader.name);
+                    return null;
+                }
+                boolean gotLock = data.tryLock();
+                if (!gotLock) { // do not wait for rebuild
+                    return null;
+                }
+
+                try {
+                    SuggesterSearcher searcher = new SuggesterSearcher(namedIndexReader.reader, resultSize);
+
+                    List<LookupResultItem> resultItems = searcher.suggest(query, namedIndexReader.name, suggesterQuery,
+                            data.getSearchCounts(suggesterQuery.getField()));
+
+                    synchronized (results) {
+                        results.addAll(resultItems);
+                    }
+                } finally {
+                    data.unlock();
+                }
+            } finally {
+                finished = true;
+            }
+            return null;
+        }
+    }
+
     /**
      * Model classes for holding project name and path to ist index directory.
      */

suggester/src/main/java/org/opengrok/suggest/SuggesterSearcher.java

@@ -47,7 +47,8 @@
 import org.opengrok.suggest.query.customized.CustomPhraseQuery;
 
 import java.io.IOException;
-import java.util.LinkedList;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.logging.Level;
@@ -65,6 +66,8 @@ class SuggesterSearcher extends IndexSearcher {
 
     private final int resultSize;
 
+    private boolean interrupted;
+
     /**
      * @param reader reader of the index for which to provide suggestions
      * @param resultSize size of the results
@@ -90,7 +93,7 @@ public List<LookupResultItem> suggest(
             final SuggesterQuery suggesterQuery,
             final PopularityCounter popularityCounter
     ) {
-        List<LookupResultItem> results = new LinkedList<>();
+        List<LookupResultItem> results = new ArrayList<>(resultSize * leafContexts.size());
 
         Query rewrittenQuery = null;
 
@@ -104,6 +107,9 @@ public List<LookupResultItem> suggest(
         }
 
         for (LeafReaderContext context : this.leafContexts) {
+            if (interrupted) {
+                break;
+            }
             try {
                 results.addAll(suggest(rewrittenQuery, context, project, suggesterQuery, popularityCounter));
             } catch (IOException e) {
@@ -125,6 +131,10 @@ private List<LookupResultItem> suggest(
             final SuggesterQuery suggesterQuery,
             final PopularityCounter searchCounts
     ) throws IOException {
+        if (Thread.interrupted()) {
+            interrupted = true;
+            return Collections.emptyList();
+        }
 
         boolean shouldLeaveOutSameTerms = shouldLeaveOutSameTerms(query, suggesterQuery);
         Set<BytesRef> tokensAlreadyIncluded = null;
@@ -148,13 +158,17 @@ private List<LookupResultItem> suggest(
 
         LookupPriorityQueue queue = new LookupPriorityQueue(resultSize);
 
-
         boolean needPositionsAndFrequencies = needPositionsAndFrequencies(query);
 
         PostingsEnum postingsEnum = null;
 
         BytesRef term = termsEnum.next();
         while (term != null) {
+            if (Thread.interrupted()) {
+                interrupted = true;
+                break;
+            }
+
             if (needPositionsAndFrequencies) {
                 postingsEnum = termsEnum.postings(postingsEnum, PostingsEnum.POSITIONS | PostingsEnum.FREQS);
             } else {