LUCENE-6772: MultiCollector now handles CollectionTerminatedException.

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1701231 13f79535-47bb-0310-9956-ffa450edef68

lucene/CHANGES.txt

@@ -155,6 +155,12 @@ Changes in Backwards Compatibility Policy
   If you have indexed text using those stemmers you may need to reindex.
   (Uwe Schindler, Robert Muir)
 
+Changes in Runtime Behavior
+
+* LUCENE-6772: MultiCollector now catches CollectionTerminatedException and
+  removes the collector that threw this exception from the list of sub
+  collectors to collect. (Adrien Grand)
+
 ======================= Lucene 5.3.0 =======================
 
 New Features

lucene/core/src/java/org/apache/lucene/search/MultiCollector.java

@@ -18,7 +18,9 @@ package org.apache.lucene.search;
  */
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
 import org.apache.lucene.index.LeafReaderContext;
 
@@ -112,21 +114,37 @@ public class MultiCollector implements Collector {
 
   @Override
   public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
-    final LeafCollector[] leafCollectors = new LeafCollector[collectors.length];
-    for (int i = 0; i < collectors.length; ++i) {
-      leafCollectors[i] = collectors[i].getLeafCollector(context);
+    final List<LeafCollector> leafCollectors = new ArrayList<>();
+    for (Collector collector : collectors) {
+      final LeafCollector leafCollector;
+      try {
+        leafCollector = collector.getLeafCollector(context);
+      } catch (CollectionTerminatedException e) {
+        // this leaf collector does not need this segment
+        continue;
+      }
+      leafCollectors.add(leafCollector);
+    }
+    switch (leafCollectors.size()) {
+      case 0:
+        throw new CollectionTerminatedException();
+      case 1:
+        return leafCollectors.get(0);
+      default:
+        return new MultiLeafCollector(leafCollectors, cacheScores);
     }
-    return new MultiLeafCollector(leafCollectors, cacheScores);
   }
 
   private static class MultiLeafCollector implements LeafCollector {
 
     private final boolean cacheScores;
     private final LeafCollector[] collectors;
+    private int numCollectors;
 
-    private MultiLeafCollector(LeafCollector[] collectors, boolean cacheScores) {
-      this.collectors = collectors;
+    private MultiLeafCollector(List<LeafCollector> collectors, boolean cacheScores) {
+      this.collectors = collectors.toArray(new LeafCollector[collectors.size()]);
       this.cacheScores = cacheScores;
+      this.numCollectors = this.collectors.length;
     }
 
     @Override
@@ -139,10 +157,28 @@ public class MultiCollector implements Collector {
       }
     }
 
+    private void removeCollector(int i) {
+      System.arraycopy(collectors, i + 1, collectors, i, numCollectors - i - 1);
+      --numCollectors;
+      collectors[numCollectors] = null;
+    }
+
     @Override
     public void collect(int doc) throws IOException {
-      for (LeafCollector c : collectors) {
-        c.collect(doc);
+      final LeafCollector[] collectors = this.collectors;
+      int numCollectors = this.numCollectors;
+      for (int i = 0; i < numCollectors; ) {
+        final LeafCollector collector = collectors[i];
+        try {
+          collector.collect(doc);
+          ++i;
+        } catch (CollectionTerminatedException e) {
+          removeCollector(i);
+          numCollectors = this.numCollectors;
+          if (numCollectors == 0) {
+            throw new CollectionTerminatedException();
+          }
+        }
       }
     }
 

lucene/core/src/test/org/apache/lucene/search/TestMultiCollector.java

@@ -0,0 +1,98 @@
+package org.apache.lucene.search;
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.RandomIndexWriter;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.LuceneTestCase;
+import org.apache.lucene.util.TestUtil;
+
+public class TestMultiCollector extends LuceneTestCase {
+
+  private static class TerminateAfterCollector extends FilterCollector {
+    
+    private int count = 0;
+    private final int terminateAfter;
+    
+    public TerminateAfterCollector(Collector in, int terminateAfter) {
+      super(in);
+      this.terminateAfter = terminateAfter;
+    }
+    
+    @Override
+    public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+      if (count >= terminateAfter) {
+        throw new CollectionTerminatedException();
+      }
+      final LeafCollector in = super.getLeafCollector(context);
+      return new FilterLeafCollector(in) {
+        @Override
+        public void collect(int doc) throws IOException {
+          if (count >= terminateAfter) {
+            throw new CollectionTerminatedException();
+          }
+          super.collect(doc);
+          count++;
+        }
+      };
+    }
+    
+  }
+
+  public void testCollectionTerminatedExceptionHandling() throws IOException {
+    final int iters = atLeast(3);
+    for (int iter = 0; iter < iters; ++iter) {
+      Directory dir = newDirectory();
+      RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+      final int numDocs = TestUtil.nextInt(random(), 100, 1000);
+      final Document doc = new Document();
+      for (int i = 0; i < numDocs; ++i) {
+        w.addDocument(doc);
+      }
+      final IndexReader reader = w.getReader();
+      w.close();
+      final IndexSearcher searcher = newSearcher(reader);
+      Map<TotalHitCountCollector, Integer> expectedCounts = new HashMap<>();
+      List<Collector> collectors = new ArrayList<>();
+      final int numCollectors = TestUtil.nextInt(random(), 1, 5);
+      for (int i = 0; i < numCollectors; ++i) {
+        final int terminateAfter = random().nextInt(numDocs + 10);
+        final int expectedCount = terminateAfter > numDocs ? numDocs : terminateAfter;
+        TotalHitCountCollector collector = new TotalHitCountCollector();
+        expectedCounts.put(collector, expectedCount);
+        collectors.add(new TerminateAfterCollector(collector, terminateAfter));
+      }
+      searcher.search(new MatchAllDocsQuery(), MultiCollector.wrap(collectors));
+      for (Map.Entry<TotalHitCountCollector, Integer> expectedCount : expectedCounts.entrySet()) {
+        assertEquals(expectedCount.getValue().intValue(), expectedCount.getKey().getTotalHits());
+      }
+      reader.close();
+      dir.close();
+    }
+  }
+
+}