[Concurrent Segment Search] support forceTermination for terminate_after parameter

[jed326]

As a follow-up to #8306, we need to add support for forceTermination in the concurrent search case. Creating a new issue to track this as we need to track the number of docs across the threads and the implementation will require some discussion.

Describe the solution you'd like
forceTermination should be properly supported in concurrent segment search.

[jed326]

Some more background:

With #8306 we will properly supporting the timeout and cancel_after_time_interval search parameters. This goal of this issue is to properly support the terminate_after workflow.

In the non-concurrent case today, if the terminate_after parameter is set to not default

OpenSearch/server/src/main/java/org/opensearch/search/query/QueryPhase.java

Lines 225 to 231 in a11f98b

	if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) {
	// add terminate_after before the filter collectors
	// it will only be applied on documents accepted by these filter collectors
	collectors.add(createEarlyTerminationCollectorContext(searchContext.terminateAfter()));
	// this collector can filter documents during the collection
	hasFilterCollector = true;
	}

then an EarlyTerminatingCollector is used which will throw an EarlyTerminationException once the specified number of docs have been collected.

OpenSearch/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollector.java

Lines 76 to 83 in 1118dcf

	if (numCollected >= maxCountHits) {
	earlyTerminated = true;
	if (forceTermination) {
	throw new EarlyTerminationException("early termination [CountBased]");
	} else {
	throw new CollectionTerminatedException();
	}
	}

In the concurrent search case, we will not throw the EarlyTerminationException and instead throw CollectionTerminatedException

OpenSearch/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollectorManager.java

Lines 39 to 42 in 1118dcf

	@Override
	public EarlyTerminatingCollector newCollector() throws IOException {
	return new EarlyTerminatingCollector(manager.newCollector(), maxCountHits, false /* forced termination is not supported */);
	}

Which we can think of as a "soft" termination. collect will still be called for all the doc and the search request will not immediately exit out due to the EarlyTerminationException. In order to keep this behavior the same between the concurrent and non-concurrent cases, we will need a way to track the number of collected documents across all the threads and subsequently throw the EarlyTerminationException once the sum of all docs collected reaches the specified doc count.

[jed326]

Some potential solutions

1. Synchronize count across threads

We can use an AtomicInteger to synchronize the collected document count across all of the threads. However, numCollected is updated on each LeafCollector::collect call, so every single concurrent segment search thread will need to both read and update the document count at the same time.

OpenSearch/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollector.java

Lines 85 to 97 in 1118dcf

	@Override
	public void collect(int doc) throws IOException {
	if (++numCollected > maxCountHits) {
	earlyTerminated = true;
	if (forceTermination) {
	throw new EarlyTerminationException("early termination [CountBased]");
	} else {
	throw new CollectionTerminatedException();
	}
	}
	super.collect(doc);
	}
	};

We would need to do some benchmarking to see how this impacts performance, but intuitively this seems like not-insignificant overhead being introduced.

2. Keep "soft" termination behavior that exists today

Continue to throw CollectionTerminatedException, which will stop collecting docs after maxCountHits is reached for a given collector. All collectors will collect docs until maxCountHits in this scenario. This solution is what would happen if we do not make any changes.

3. Change (2) to do a "hard" termination.

Once one collector hits maxCountHits, stop all other collectors even if they have not hit maxCountHits then. This solution gets us some force termination behavior, however I don't think the behavior is very intuitive for the user.

4. Disallow terminate_after option when concurrent segment search is being used.

We could make concurrent segment search concretely not support forceTermination, at least for now, and revisit this at a later point as needed.

---

In terms of functionality I think solution 1 is best as it keeps the behavior the same between concurrent and non-concurrent search use cases. However, I am concerned about the potential performance hit doing such synchronization.

[jed326]

Tagging @reta @sohami @andrross for some feedback/thoughts. Thanks!

[andrross]

In terms of functionality I think solution 1 is best as it keeps the behavior the same between concurrent and non-concurrent search use cases. However, I am concerned about the potential performance hit doing such synchronization.

I agree with your analysis here. Any idea what the synchronization overhead might be?

I think option 4 is okay because it can be temporary and terminate_after can be added later if users need it. I'm also not totally opposed to option 2 where we essentially add the caveat that in the concurrent case terminate_after is per concurrent slice, not total count (I'm less sure this is a good idea). I think option 3 is my least preferred for the reason you cited.

There is potentially an option 1a where we don't keep an exact count (perhaps by updating the synchronized count every 10 or 100 or 1000 documents or whatever) in order to reduce the synchronization burden.