From 5e6f22b925add09dc20b7d53bf8f177a582a4ed1 Mon Sep 17 00:00:00 2001
From: Adrien Grand <jpountz@gmail.com>
Date: Mon, 8 Feb 2016 16:41:42 +0100
Subject: [PATCH] LUCENE-7002: Fixed MultiCollector to not throw a NPE if
 setScorer is called after one of the sub collectors is done collecting.

---
 lucene/CHANGES.txt                            |  3 +
 .../apache/lucene/search/MultiCollector.java  |  3 +-
 .../lucene/search/TestMultiCollector.java     | 71 +++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index c2aa7722036..b729f770f72 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -235,6 +235,9 @@ Bug Fixes
 * LUCENE-6998: Fix a couple places to better detect truncated index files
   as corruption.  (Robert Muir, Mike McCandless)
 
+* LUCENE-7002: Fixed MultiCollector to not throw a NPE if setScorer is called
+  after one of the sub collectors is done collecting. (John Wang, Adrien Grand)
+
 Other
 
 * LUCENE-6924: Upgrade randomizedtesting to 2.3.2. (Dawid Weiss)
diff --git a/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java b/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java
index 71a2be27244..81cd5942f56 100644
--- a/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java
+++ b/lucene/core/src/java/org/apache/lucene/search/MultiCollector.java
@@ -152,7 +152,8 @@ public class MultiCollector implements Collector {
       if (cacheScores) {
         scorer = new ScoreCachingWrappingScorer(scorer);
       }
-      for (LeafCollector c : collectors) {
+      for (int i = 0; i < numCollectors; ++i) {
+        final LeafCollector c = collectors[i];
         c.setScorer(scorer);
       }
     }
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestMultiCollector.java b/lucene/core/src/test/org/apache/lucene/search/TestMultiCollector.java
index 9a142ffd826..4ef54dac3d1 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestMultiCollector.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestMultiCollector.java
@@ -19,9 +19,12 @@ package org.apache.lucene.search;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.lucene.document.Document;
 import org.apache.lucene.index.IndexReader;
@@ -63,6 +66,27 @@ public class TestMultiCollector extends LuceneTestCase {
     
   }
 
+  private static class SetScorerCollector extends FilterCollector {
+
+    private final AtomicBoolean setScorerCalled;
+
+    public SetScorerCollector(Collector in, AtomicBoolean setScorerCalled) {
+      super(in);
+      this.setScorerCalled = setScorerCalled;
+    }
+
+    @Override
+    public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+      return new FilterLeafCollector(super.getLeafCollector(context)) {
+        @Override
+        public void setScorer(Scorer scorer) throws IOException {
+          super.setScorer(scorer);
+          setScorerCalled.set(true);
+        }
+      };
+    }
+  }
+
   public void testCollectionTerminatedExceptionHandling() throws IOException {
     final int iters = atLeast(3);
     for (int iter = 0; iter < iters; ++iter) {
@@ -95,4 +119,51 @@ public class TestMultiCollector extends LuceneTestCase {
     }
   }
 
+  public void testSetScorerAfterCollectionTerminated() throws IOException {
+    Collector collector1 = new TotalHitCountCollector();
+    Collector collector2 = new TotalHitCountCollector();
+
+    AtomicBoolean setScorerCalled1 = new AtomicBoolean();
+    collector1 = new SetScorerCollector(collector1, setScorerCalled1);
+    
+    AtomicBoolean setScorerCalled2 = new AtomicBoolean();
+    collector2 = new SetScorerCollector(collector2, setScorerCalled2);
+
+    collector1 = new TerminateAfterCollector(collector1, 1);
+    collector2 = new TerminateAfterCollector(collector2, 2);
+
+    Scorer scorer = new FakeScorer();
+
+    List<Collector> collectors = Arrays.asList(collector1, collector2);
+    Collections.shuffle(collectors, random());
+    Collector collector = MultiCollector.wrap(collectors);
+
+    LeafCollector leafCollector = collector.getLeafCollector(null);
+    leafCollector.setScorer(scorer);
+    assertTrue(setScorerCalled1.get());
+    assertTrue(setScorerCalled2.get());
+
+    leafCollector.collect(0);
+    leafCollector.collect(1);
+
+    setScorerCalled1.set(false);
+    setScorerCalled2.set(false);
+    leafCollector.setScorer(scorer);
+    assertFalse(setScorerCalled1.get());
+    assertTrue(setScorerCalled2.get());
+
+    try {
+      leafCollector.collect(1);
+      fail();
+    } catch (CollectionTerminatedException e) {
+      // expected
+    }
+
+    setScorerCalled1.set(false);
+    setScorerCalled2.set(false);
+    leafCollector.setScorer(scorer);
+    assertFalse(setScorerCalled1.get());
+    assertFalse(setScorerCalled2.get());
+  }
+
 }