Skip to content

Commit

Permalink
Fix concurrency bug in AbstractStringScriptFieldAutomatonQuery (#106678)
Browse files Browse the repository at this point in the history 
Back when we introduced queries against runtime fields, Elasticsearch did not support
inter-segment concurrency yet. At the time, it was fine to assume that segments will be
searched sequentially. AbstractStringScriptFieldAutomatonQuery used to have a BytesRefBuilder
instance shared across the segments, which gets re-initialized when each segment starts its work.
This is no longer possible with inter-segment concurrency.

Closes #105911
  • Loading branch information
@javanna
javanna committed Mar 26, 2024
1 parent ceb2701 commit 08d7542
Show file tree
Hide file tree
Showing 21 changed files with 260 additions and 199 deletions.
6 changes: 6 additions & 0 deletions docs/changelog/106678.yaml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
pr: 106678
summary: Fix concurrency bug in `AbstractStringScriptFieldAutomatonQuery`
area: Search
type: bug
issues:
- 105911
2 changes: 1 addition & 1 deletion server/src/main/java/org/elasticsearch/search/runtime/AbstractBooleanScriptFieldQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ abstract class AbstractBooleanScriptFieldQuery extends AbstractScriptFieldQuery<
}

@Override
protected boolean matches(BooleanFieldScript scriptContext, int docId) {
protected final boolean matches(BooleanFieldScript scriptContext, int docId) {
scriptContext.runForDoc(docId);
return matches(scriptContext.trues(), scriptContext.falses());
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/elasticsearch/search/runtime/AbstractDoubleScriptFieldQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ abstract class AbstractDoubleScriptFieldQuery extends AbstractScriptFieldQuery<D
}

@Override
protected boolean matches(DoubleFieldScript scriptContext, int docId) {
protected final boolean matches(DoubleFieldScript scriptContext, int docId) {
scriptContext.runForDoc(docId);
return matches(scriptContext.values(), scriptContext.count());
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/elasticsearch/search/runtime/AbstractGeoPointScriptFieldQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ abstract class AbstractGeoPointScriptFieldQuery extends AbstractScriptFieldQuery
}

@Override
protected boolean matches(GeoPointFieldScript scriptContext, int docId) {
protected final boolean matches(GeoPointFieldScript scriptContext, int docId) {
scriptContext.runForDoc(docId);
return matches(scriptContext.lats(), scriptContext.lons(), scriptContext.count());
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/elasticsearch/search/runtime/AbstractIpScriptFieldQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ abstract class AbstractIpScriptFieldQuery extends AbstractScriptFieldQuery<IpFie
}

@Override
protected boolean matches(IpFieldScript scriptContext, int docId) {
protected final boolean matches(IpFieldScript scriptContext, int docId) {
scriptContext.runForDoc(docId);
return matches(scriptContext.values(), scriptContext.count());
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/elasticsearch/search/runtime/AbstractLongScriptFieldQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ abstract class AbstractLongScriptFieldQuery extends AbstractScriptFieldQuery<Abs
}

@Override
protected boolean matches(AbstractLongFieldScript scriptContext, int docId) {
protected final boolean matches(AbstractLongFieldScript scriptContext, int docId) {
scriptContext.runForDoc(docId);
return AbstractLongScriptFieldQuery.this.matches(scriptContext.values(), scriptContext.count());
}
Expand Down
31 changes: 19 additions & 12 deletions server/src/main/java/org/elasticsearch/search/runtime/AbstractScriptFieldQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,18 +71,7 @@ public boolean isCacheable(LeafReaderContext ctx) {
public Scorer scorer(LeafReaderContext ctx) {
S scriptContext = scriptContextFunction.apply(ctx);
DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc());
TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
@Override
public boolean matches() {
return AbstractScriptFieldQuery.this.matches(scriptContext, approximation.docID());
}

@Override
public float matchCost() {
return MATCH_COST;
}
};
return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
return new ConstantScoreScorer(this, score(), scoreMode, createTwoPhaseIterator(scriptContext, approximation));
}

@Override
Expand All @@ -96,6 +85,24 @@ public Explanation explain(LeafReaderContext context, int doc) throws IOExceptio
};
}

/**
* Creates the {@link TwoPhaseIterator} for the current leaf reader. Override to inject custom behaviour or provide
* additional context to the matches method when needed.
*/
protected TwoPhaseIterator createTwoPhaseIterator(S scriptContext, DocIdSetIterator approximation) {
return new TwoPhaseIterator(approximation) {
@Override
public boolean matches() {
return AbstractScriptFieldQuery.this.matches(scriptContext, approximation.docID());
}

@Override
public float matchCost() {
return MATCH_COST;
}
};
}

protected abstract boolean matches(S scriptContext, int docId);

@Override
Expand Down
26 changes: 24 additions & 2 deletions ...c/main/java/org/elasticsearch/search/runtime/AbstractStringScriptFieldAutomatonQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@

package org.elasticsearch.search.runtime;

import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.search.TwoPhaseIterator;
import org.apache.lucene.util.BytesRefBuilder;
import org.apache.lucene.util.automaton.ByteRunAutomaton;
import org.elasticsearch.script.Script;
Expand All @@ -17,7 +19,6 @@
import java.util.List;

public abstract class AbstractStringScriptFieldAutomatonQuery extends AbstractStringScriptFieldQuery {
private final BytesRefBuilder scratch = new BytesRefBuilder();
private final ByteRunAutomaton automaton;

public AbstractStringScriptFieldAutomatonQuery(
Expand All @@ -31,7 +32,23 @@ public AbstractStringScriptFieldAutomatonQuery(
}

@Override
protected final boolean matches(List<String> values) {
protected TwoPhaseIterator createTwoPhaseIterator(StringFieldScript scriptContext, DocIdSetIterator approximation) {
BytesRefBuilder scratch = new BytesRefBuilder();
return new TwoPhaseIterator(approximation) {
@Override
public boolean matches() {
scriptContext.runForDoc(approximation.docID());
return AbstractStringScriptFieldAutomatonQuery.this.matches(scriptContext.getValues(), scratch);
}

@Override
public float matchCost() {
return MATCH_COST;
}
};
}

protected final boolean matches(List<String> values, BytesRefBuilder scratch) {
for (String value : values) {
scratch.copyChars(value);
if (automaton.run(scratch.bytes(), 0, scratch.length())) {
Expand All @@ -41,6 +58,11 @@ protected final boolean matches(List<String> values) {
return false;
}

@Override
protected final boolean matches(List<String> values) {
throw new UnsupportedOperationException();
}

@Override
public final void visit(QueryVisitor visitor) {
if (visitor.acceptField(fieldName())) {
Expand Down
41 changes: 21 additions & 20 deletions server/src/test/java/org/elasticsearch/index/mapper/BooleanScriptFieldTypeTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@

import static java.util.Collections.emptyMap;
import static java.util.Collections.singletonList;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;

Expand All @@ -73,8 +74,8 @@ protected ScriptFactory dummyScript() {
@Override
public void testDocValues() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true, false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true, false]}"))));
List<Long> results = new ArrayList<>();
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
Expand Down Expand Up @@ -104,16 +105,16 @@ public void collect(int doc) throws IOException {
};
}
});
assertThat(results, equalTo(List.of(1L, 0L, 1L)));
assertThat(results, containsInAnyOrder(1L, 0L, 1L));
}
}
}

@Override
public void testSort() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
BooleanScriptFieldData ifd = simpleMappedFieldType().fielddataBuilder(mockFielddataContext()).build(null, null);
Expand All @@ -128,8 +129,8 @@ public void testSort() throws IOException {
@Override
public void testUsedInScript() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
{
Expand Down Expand Up @@ -185,10 +186,10 @@ public double execute(ExplanationHolder explanation) {
@Override
public void testExistsQuery() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true, false]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": []}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true, false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": []}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
assertThat(searcher.count(simpleMappedFieldType().existsQuery(mockContext())), equalTo(3));
Expand All @@ -199,7 +200,7 @@ public void testExistsQuery() throws IOException {
@Override
public void testRangeQuery() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
MappedFieldType ft = simpleMappedFieldType();
Expand All @@ -210,7 +211,7 @@ public void testRangeQuery() throws IOException {
}
}
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
MappedFieldType ft = simpleMappedFieldType();
Expand All @@ -221,8 +222,8 @@ public void testRangeQuery() throws IOException {
}
}
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
MappedFieldType ft = simpleMappedFieldType();
Expand Down Expand Up @@ -269,7 +270,7 @@ protected Query randomRangeQuery(MappedFieldType ft, SearchExecutionContext ctx)
@Override
public void testTermQuery() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
assertThat(searcher.count(simpleMappedFieldType().termQuery(true, mockContext())), equalTo(1));
Expand All @@ -282,7 +283,7 @@ public void testTermQuery() throws IOException {
}
}
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
assertThat(searcher.count(simpleMappedFieldType().termQuery(false, mockContext())), equalTo(1));
Expand All @@ -305,7 +306,7 @@ protected Query randomTermQuery(MappedFieldType ft, SearchExecutionContext ctx)
@Override
public void testTermsQuery() throws IOException {
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [true]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
assertThat(searcher.count(simpleMappedFieldType().termsQuery(List.of(true, true), mockContext())), equalTo(1));
Expand All @@ -315,7 +316,7 @@ public void testTermsQuery() throws IOException {
}
}
try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
addDocument(iw, List.of(new StoredField("_source", new BytesRef("{\"foo\": [false]}"))));
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
assertThat(searcher.count(simpleMappedFieldType().termsQuery(List.of(false, false), mockContext())), equalTo(1));
Expand Down Expand Up @@ -364,7 +365,7 @@ public XContentParser parser() {
while (ctx.parser().nextToken() != Token.END_ARRAY) {
ootb.parse(ctx);
}
iw.addDocument(ctx.doc());
addDocument(iw, ctx.doc());
try (DirectoryReader reader = iw.getReader()) {
IndexSearcher searcher = newSearcher(reader);
assertSameCount(
Expand Down

0 comments on commit 08d7542

Please sign in to comment.