Lab CS2: CollectorManager and Slice Reduction

Background

The intensive called the CollectorManager "the heart of the feature": newCollector() hands each slice its own mutable collector, the slices run in parallel with no shared state, and reduce(collectors) merges them into one shard-local result. This lab makes you own that contract. First you read OpenSearch's real wiring — ContextIndexSearcher and a concrete CollectorManager in core. Then you build a tiny standalone Lucene program: a CollectorManager that counts docs and tracks the max score, run it with an ExecutorService across slices, and verify reduce() merges correctly. Finally you'll prove the associativity invariant by deliberately breaking it.

Why This Matters for Contributors

When an aggregation produces wrong results under concurrent search, the bug is almost always in a CollectorManager.reduce that isn't associative/commutative, or in a collector that smuggles shared mutable state across slices. To fix one of those — and several real OpenSearch aggregations needed this fix to become slice-safe — you must be able to write a correct CollectorManager and see the failure mode of a broken one. This lab is that skill on a minimal, fast feedback loop, with no OpenSearch build in the way.

Prerequisites

Java 21+ (java -version).
A Lucene core jar on your classpath. Two easy ways: - From an apache/lucene checkout: ./gradlew :lucene:core:jar then find the jar under lucene/core/build/libs/. - From an OpenSearch checkout: the Lucene jars are already downloaded — find them with find ~/.gradle -name 'lucene-core-*.jar' | head.
You've read the intensive — the CollectorManager contract and the associativity argument especially.

Note: This lab uses Lucene directly, not OpenSearch. That's deliberate: CollectorManager is a Lucene type, and stripping away OpenSearch lets you see the contract with zero noise. Everything you learn transfers directly to ContextIndexSearcher, which is just IndexSearcher + a search-pool executor.

Step-by-Step Tasks

Step 1 — Read OpenSearch's real CollectorManager wiring

Before writing your own, read the real thing. Find ContextIndexSearcher and a concrete CollectorManager in core:

cd ~/src/OpenSearch
find server -name ContextIndexSearcher.java
# How the searcher decides to go concurrent and where it calls search(query, manager):
grep -n "CollectorManager\|search(\|slices\|getExecutor\|collectorManager" \
  server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java | head -25

# A concrete manager in the query phase (top-docs / count):
grep -rln "implements CollectorManager\|CollectorManager<" \
  server/src/main/java/org/opensearch/search/query/ \
  server/src/main/java/org/opensearch/search/aggregations/ | head
# Read its newCollector() and reduce():
F=$(grep -rl "implements CollectorManager" server/src/main/java/org/opensearch/search/query/ | head -1)
grep -n "newCollector\|reduce(" "$F"

Answer in your notes, from the code you just read:

Where does ContextIndexSearcher decide to run concurrently vs sequentially?
What does the concrete reduce() you found actually merge (a priority queue? a count? a bucket map?), and is that merge associative + commutative?

Step 2 — Build a tiny indexed dataset

A standalone program that writes a few segments. The commit() after each batch forces a new segment, so we get multiple leaves to slice:

// MakeIndex.java
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.store.*;
import java.nio.file.*;

public class MakeIndex {
  public static void main(String[] args) throws Exception {
    Path dir = Paths.get("/tmp/cs2-index");
    try (Directory d = FSDirectory.open(dir);
         IndexWriter w = new IndexWriter(d,
             new IndexWriterConfig(new StandardAnalyzer())
                 .setOpenMode(IndexWriterConfig.OpenMode.CREATE))) {
      int doc = 0;
      for (int seg = 0; seg < 6; seg++) {        // 6 commits -> ~6 segments
        int n = 1000 + seg * 500;
        for (int i = 0; i < n; i++, doc++) {
          Document dd = new Document();
          dd.add(new StringField("g", "k" + (doc % 50), Field.Store.NO));
          dd.add(new IntPoint("v", doc % 1000));
          dd.add(new TextField("body", "lucene segment slice reduce term " + (doc % 7), Field.Store.NO));
          w.addDocument(dd);
        }
        w.commit();   // <-- forces a segment boundary
      }
      System.out.println("indexed " + doc + " docs across ~6 segments");
    }
  }
}

LUCENE=$(find ~/.gradle -name 'lucene-core-*.jar' | head -1)
ANALYZ=$(find ~/.gradle -name 'lucene-analysis-common-*.jar' | head -1)
CP="$LUCENE:$ANALYZ"
javac -cp "$CP" MakeIndex.java
java  -cp ".:$CP" MakeIndex
# Confirm multiple segments:
ls /tmp/cs2-index | grep -E '\.si$' | wc -l   # number of segments

Step 3 — Write your own CollectorManager

Now the core of the lab. A CountMaxCollector accumulates a doc count and a max score; a CountMaxManager is the CollectorManager that creates one per slice and reduces them. Note the invariant: the collector holds mutable per-slice state; the manager holds none.

// CountMax.java
import org.apache.lucene.search.*;
import java.io.IOException;
import java.util.Collection;

/** Per-slice result: how many docs matched and the best score seen. */
record CountMax(long count, float maxScore) {}

/** A Collector with MUTABLE state -- one instance per slice, never shared. */
final class CountMaxCollector implements Collector {
  long count = 0;
  float maxScore = Float.NEGATIVE_INFINITY;

  @Override public ScoreMode scoreMode() { return ScoreMode.COMPLETE; } // need scores

  @Override public LeafCollector getLeafCollector(LeafReaderContext ctx) {
    return new LeafCollector() {
      private Scorable scorer;
      @Override public void setScorer(Scorable s) { this.scorer = s; }
      @Override public void collect(int doc) throws IOException {
        count++;
        float s = scorer.score();
        if (s > maxScore) maxScore = s;
      }
    };
  }

  CountMax result() { return new CountMax(count, maxScore); }
}

/** The CollectorManager: newCollector() per slice, reduce() merges them. */
final class CountMaxManager implements CollectorManager<CountMaxCollector, CountMax> {
  @Override public CountMaxCollector newCollector() {
    return new CountMaxCollector();           // fresh mutable state, no sharing
  }
  @Override public CountMax reduce(Collection<CountMaxCollector> collectors) {
    long total = 0;
    float best = Float.NEGATIVE_INFINITY;
    for (CountMaxCollector c : collectors) {  // order-independent: sum + max
      CountMax r = c.result();
      total += r.count();                     // SUM is associative + commutative
      best = Math.max(best, r.maxScore());     // MAX is associative + commutative
    }
    return new CountMax(total, best);
  }
}

Step 4 — Run it across slices with an executor

Wire it to an IndexSearcher constructed with an ExecutorService — that is exactly what ContextIndexSearcher does in concurrent mode. The searcher slices the leaves and submits one task per slice:

// RunCountMax.java
import org.apache.lucene.index.*;
import org.apache.lucene.search.*;
import org.apache.lucene.store.*;
import java.nio.file.*;
import java.util.concurrent.*;

public class RunCountMax {
  public static void main(String[] args) throws Exception {
    Path dir = Paths.get("/tmp/cs2-index");
    ExecutorService pool = Executors.newFixedThreadPool(
        Math.max(2, Runtime.getRuntime().availableProcessors()),
        r -> { Thread t = new Thread(r, "slice-worker"); t.setDaemon(true); return t; });

    try (Directory d = FSDirectory.open(dir);
         DirectoryReader reader = DirectoryReader.open(d)) {

      // Pass the executor -> concurrent slicing, just like ContextIndexSearcher.
      IndexSearcher searcher = new IndexSearcher(reader, pool);

      // How many slices did Lucene make? (its heuristic, run in YOUR program)
      IndexSearcher.LeafSlice[] slices = searcher.getSlices();
      System.out.println("leaves=" + reader.leaves().size()
                       + "  slices=" + (slices == null ? 1 : slices.length));

      // A query that matches a chunk of docs and produces scores.
      Query q = new org.apache.lucene.search.MatchAllDocsQuery();

      CountMax cm = searcher.search(q, new CountMaxManager());
      System.out.println("CONCURRENT  count=" + cm.count() + "  maxScore=" + cm.maxScore());

      // Sanity: a sequential count must agree.
      long seq = searcher.count(q);
      System.out.println("SEQUENTIAL  count=" + seq + "  agree=" + (seq == cm.count()));
    } finally {
      pool.shutdown();
    }
  }
}

javac -cp "$CP" CountMax.java RunCountMax.java
java  -cp ".:$CP" RunCountMax

Expected: slices > 1 (your 6 segments grouped into a few slices), and the concurrent count equals the sequential count. The reduce() summed the per-slice counts correctly.

Step 5 — Prove the slices really ran on different threads

Add a one-line print in getLeafCollector to see which thread scanned which segment — proof that slices are genuinely parallel and each has its own collector:

// inside CountMaxCollector.getLeafCollector, first line:
System.out.println("slice collector " + System.identityHashCode(CountMaxCollector.this)
    + " leaf " + ctx.ord + " on " + Thread.currentThread().getName());

Re-run. You'll see multiple distinct collector identity-hashes (one per slice, proving newCollector() ran per slice) and multiple slice-worker thread names (proving parallel execution). No two leaves in different slices share a collector instance.

Step 6 — Break associativity on purpose

Now feel the bug. Replace the manager's reduce with an order-dependent merge and watch it produce different answers across runs:

// BROKEN reduce: "the maxScore of whichever slice reduced first" -- order-dependent.
@Override public CountMax reduce(Collection<CountMaxCollector> collectors) {
  long total = 0;
  float firstSeen = Float.NaN;
  for (CountMaxCollector c : collectors) {
    total += c.result().count();
    if (Float.isNaN(firstSeen)) firstSeen = c.result().maxScore();  // WRONG: depends on order
  }
  return new CountMax(total, firstSeen);
}

Run it several times. The count stays correct (sum is still associative) but the reported "max" jumps around between runs because the collection's iteration order reflects nondeterministic slice completion. This is the entire class of concurrent-search aggregation bug, reproduced in 20 lines. Restore the correct Math.max version before moving on.

Warning: Single-threaded, this bug is invisible — pass a same-thread executor (or new IndexSearcher(reader) with no executor) and the order is fixed, so firstSeen is stable and the test passes. That's precisely why these bugs survive code review and only surface in production with concurrency on.

Step 7 — Tie it back to OpenSearch

Re-read the OpenSearch manager you found in Step 1 with new eyes:

cd ~/src/OpenSearch
# A top-docs style manager reduces priority queues -- find its merge:
grep -rn "TopDocs.merge\|TopScoreDocCollector\|CollectorManager\|reduce(" \
  server/src/main/java/org/opensearch/search/query/ | head

Confirm: its reduce merges priority queues with TopDocs.merge (or similar), which is associative/commutative under the sort comparator — the same discipline your CountMaxManager follows.

Deliverables

MakeIndex.java, CountMax.java, RunCountMax.java compiling and running.
Output showing slices > 1, multiple distinct collector identity-hashes, multiple slice-worker threads, and concurrent count == sequential count.
Output from the broken reduce showing the max wobbling across runs while count stays correct.
A short note: where in ContextIndexSearcher does OpenSearch make the same search(query, manager) call, and what does its real manager reduce?

Expected Output

leaves=6  slices=4
slice collector 1131040331 leaf 0 on slice-worker
slice collector 1830910733 leaf 3 on slice-worker
slice collector 1131040331 leaf 1 on slice-worker     # same collector reused within a slice
...
CONCURRENT  count=10500  maxScore=1.0
SEQUENTIAL  count=10500  agree=true

# with the BROKEN reduce, across 3 runs:
CONCURRENT maxScore=1.0
CONCURRENT maxScore=0.873...      # WOBBLES -- order-dependent reduce
CONCURRENT maxScore=1.0

Troubleshooting

Symptom	Cause	Fix
`slices=1`	Too few segments, or Lucene merged them	add more `commit()` batches in MakeIndex; check `*.si` count
`getSlices()` not found	Method name varies by Lucene version	grep the jar's `IndexSearcher` API or use `searcher.search(...)` directly and count via prints
Scores all `1.0`	`MatchAllDocsQuery` has constant score	use a `TermQuery` on `body` to get varied scores so the broken reduce visibly wobbles
`ClassNotFound` for analyzer	analysis-common jar not on classpath	`find ~/.gradle -name 'lucene-analysis-common-*.jar'` and add it
Concurrent count != sequential	A real bug in your collector (shared state)	ensure all mutable state lives in the collector, none in the manager

Stretch Goals

Write a top-K manager. Build a CollectorManager that returns the top 10 docs by score, with newCollector() returning a TopScoreDocCollector and reduce calling TopDocs.merge. Verify it matches searcher.search(q, TopScoreDocCollector...) sequential output.
A terms-count manager. Reduce a HashMap<BytesRef,Long> of per-term counts — the skeleton of a terms aggregation. Prove reduce (map merge by summing) is associative.
Force a single slice. Construct new IndexSearcher(reader) (no executor) and confirm reduce is still called (with one collector) — concurrency-off is just the one-slice case of the same contract.
Measure the reduce cost. Time reduce vs the scan for a cheap query and a heavy one; relate to the "tiny segments lose" trade-off from the intensive.

Coding Exercises

You wrote one CollectorManager already; these turn it into a graded suite that proves the associativity invariant the way a maintainer must before trusting a slice-safe aggregation. Build on your CountMax files — do not just re-run the demo.

(warm-up) Turn the demo into a JUnit test. Wrap Step 4 in a real test (CountMaxManagerTest) using JUnit 5 or OpenSearchTestCase if you have the test classpath. @BeforeAll builds the index; one @Test asserts concurrent.count() == searcher.count(q); a second asserts maxScore equals the sequential TopScoreDocCollector max. Run with ./gradlew test (or java -jar junit-platform-console-standalone.jar). This is Deliverable #2 as a green bar.
(warm-up) A property test for reduce. Write a test that builds N random per-slice CountMax partials, shuffles them, reduces each permutation, and asserts every permutation yields the identical result. Now drop in the broken reduce from Step 6 and watch the property test go red — you have automated the wobble you observed by hand. Keep both reduces behind a flag so the test documents the contrast.
(core) A top-K CollectorManager. Promote the first Stretch Goal to graded code: implement TopKManager whose newCollector() returns a TopScoreDocCollector and whose reduce calls TopDocs.merge. Write a test that runs a TermQuery on body (varied scores!) concurrently and asserts the merged top-10 doc ids and scores exactly match searcher.search(q, sequentialCollector). This is the real top-docs query phase in miniature.
(core) A terms-count CollectorManager. Implement TermsCountManager that reduces a HashMap<BytesRef,Long> (sum per key) — the skeleton of the terms aggregation. Write a test that asserts the concurrent per-term counts equal a single-threaded scan, and a second test that deliberately shares one mutable map across slices (move it into the manager) and asserts the result becomes wrong/non-deterministic under load. You have reproduced "shared mutable state across slices," the second class of concurrent-agg bug.
(core) Map the abstraction onto OpenSearch. In an OpenSearch checkout, find a real aggregation manager: rg -l "implements CollectorManager" server/src/main/java/org/opensearch/search/aggregations/ and read its reduce. Write a one-paragraph note (in a // comment block at the top of a small test) naming the class, what its reduce merges, and the associative operation it relies on — then write a JUnit assertion mirroring that merge on toy inputs (e.g. summing two InternalAggregation partials if the classpath allows, else your HashMap analogue). Cite ContextIndexSearcher.search(query, manager) as the call site (rg -n "search\\(.*[Mm]anager" server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java).
(advanced challenge) A generic, contract-checking harness. Build ManagerContractTester<C extends Collector, R> that, given any CollectorManager and a Lucene reader, (a) runs it concurrently and sequentially and asserts equal results, (b) shuffles the collector collection K times before reduce and asserts determinism, and (c) asserts each newCollector() returns a distinct instance (via an IdentityHashMap, catching shared-state managers). Run all four of your managers — CountMax, TopK, TermsCount, and the intentionally broken one — through it and show it passes the three good ones and fails the broken one with a clear message. This is a reusable slice-safety oracle, the tool you'd actually bring to a real concurrent-agg fix in Lab CS3's capstone.

Issues to Practice On

CollectorManager correctness is core-repo work, concentrated in Search and Search:Aggregations. The bugs are subtle, so reproduction discipline matters most.

gh issue list --repo opensearch-project/OpenSearch --label "good first issue" --state open
gh issue list --repo opensearch-project/OpenSearch --label "Search:Aggregations" --state open
gh issue list --repo opensearch-project/OpenSearch --search "concurrent aggregation reduce" --state open
gh issue list --repo opensearch-project/OpenSearch --label "flaky-test" --state open
# Labels drift — confirm before relying on one:
gh label list --repo opensearch-project/OpenSearch | rg -i "aggreg|search|flaky"

Two representative patterns:

"Aggregation X returns wrong results with concurrent search enabled." This is almost always a non-associative or non-commutative reduce, or a collector with shared state. Reproduce by toggling search.concurrent_segment_search.mode, locate the manager with rg "implements CollectorManager", fix the merge, and add a test that fails under concurrency and passes after — exactly your Exercise 2/4 pattern.
A flaky aggregation test that only fails occasionally. Run it under -Dtests.iters=100; the flake is the nondeterministic slice-completion order exposing an order-dependent reduce. The fix is a deterministic reduce plus an assertion that pins it.

Planted bug exercise. In an OpenSearch checkout, find a concrete aggregation's reduce (rg -n "public .* reduce\\(" server/src/main/java/org/opensearch/search/aggregations/). Introduce the exact bug from Step 6 — make the merge depend on iteration order (e.g. keep only the first partial's value for some field instead of combining). Run that aggregation's tests under concurrency and -Dtests.iters=50; watch which test flakes. Revert, then add an assertion that runs the same agg with mode=all and mode=none and requires identical results — the guard that would have caught it.

Etiquette: claim the issue first, reproduce before fixing, and ship a test + CHANGELOG.md entry + DCO git commit -s. See community interaction and the good-first-issue PR lab.

Validation / Self-check

You can write the CollectorManager contract from memory and explain why newCollector() per slice (not a shared collector) is the thread-safety mechanism.
You can name two associative+commutative reduces (sum, max, queue-merge) and explain why an order-dependent reduce is a latent bug.
You can demonstrate the bug: same code, correct count, wobbling max, and explain why it's invisible single-threaded.
You can point to where ContextIndexSearcher makes the same search(query, collectorManager) call OpenSearch uses in production.

Next: Lab CS3 — Benchmark and Tune Slicing.

Open-Source Engineer & Contributor