Skip to content

Commit ba77068

Browse files
benchaplinbenchaplin
committed
Fix inner hits + aggregations concurrency bug (elastic#128036)
Fork InnerHitSubContext instances before source is fetched in aggregations to prevent inter-segment race conditions. Relates to elastic#122419
1 parent 86b2daa commit ba77068

File tree

8 files changed

+219
-1
lines changed

8 files changed

+219
-1
lines changed

‎docs/changelog/128036.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
pr: 128036
2+
summary: Fix inner hits + aggregations concurrency bug
3+
area: Search
4+
type: bug
5+
issues:
6+
- 122419

‎modules/parent-join/src/internalClusterTest/java/org/elasticsearch/join/query/InnerHitsIT.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.elasticsearch.script.ScriptType;
2727
import org.elasticsearch.search.SearchHit;
2828
import org.elasticsearch.search.SearchHits;
29+
import org.elasticsearch.search.aggregations.AggregationBuilder;
30+
import org.elasticsearch.search.aggregations.metrics.TopHits;
2931
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
3032
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
3133
import org.elasticsearch.search.sort.FieldSortBuilder;
@@ -51,6 +53,7 @@
5153
import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO;
5254
import static org.elasticsearch.join.query.JoinQueryBuilders.hasChildQuery;
5355
import static org.elasticsearch.join.query.JoinQueryBuilders.hasParentQuery;
56+
import static org.elasticsearch.search.aggregations.AggregationBuilders.topHits;
5457
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
5558
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
5659
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCountAndNoFailures;
@@ -64,6 +67,7 @@
6467
import static org.hamcrest.Matchers.containsString;
6568
import static org.hamcrest.Matchers.equalTo;
6669
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
70+
import static org.hamcrest.Matchers.lessThanOrEqualTo;
6771
import static org.hamcrest.Matchers.notNullValue;
6872
import static org.hamcrest.Matchers.nullValue;
6973

@@ -698,4 +702,68 @@ public void testTooHighResultWindow() {
698702
)
699703
);
700704
}
705+
706+
public void testTopHitsOnParentChild() throws Exception {
707+
assertAcked(
708+
prepareCreate("idx").setMapping(
709+
jsonBuilder().startObject()
710+
.startObject("_doc")
711+
.startObject("properties")
712+
.startObject("id")
713+
.field("type", "keyword")
714+
.endObject()
715+
.startObject("join_field")
716+
.field("type", "join")
717+
.startObject("relations")
718+
.field("parent", new String[] { "child1", "child2" })
719+
.endObject()
720+
.endObject()
721+
.endObject()
722+
.endObject()
723+
.endObject()
724+
)
725+
);
726+
ensureGreen("idx");
727+
728+
List<IndexRequestBuilder> requestBuilders = new ArrayList<>();
729+
int numDocs = scaledRandomIntBetween(10, 100);
730+
int child1 = 0;
731+
int child2 = 0;
732+
int[] child1InnerObjects = new int[numDocs];
733+
int[] child2InnerObjects = new int[numDocs];
734+
for (int parent = 0; parent < numDocs; parent++) {
735+
String parentId = String.format(Locale.ENGLISH, "p_%03d", parent);
736+
requestBuilders.add(createIndexRequest("idx", "parent", parentId, null));
737+
738+
int numChildDocs = child1InnerObjects[parent] = scaledRandomIntBetween(1, numDocs);
739+
int limit = child1 + numChildDocs;
740+
for (; child1 < limit; child1++) {
741+
requestBuilders.add(createIndexRequest("idx", "child1", String.format(Locale.ENGLISH, "c1_%04d", child1), parentId));
742+
}
743+
numChildDocs = child2InnerObjects[parent] = scaledRandomIntBetween(1, numDocs);
744+
limit = child2 + numChildDocs;
745+
for (; child2 < limit; child2++) {
746+
requestBuilders.add(createIndexRequest("idx", "child2", String.format(Locale.ENGLISH, "c2_%04d", child2), parentId));
747+
}
748+
}
749+
750+
indexRandom(true, requestBuilders);
751+
ensureSearchable();
752+
753+
QueryBuilder hasChildQuery = hasChildQuery("child2", matchAllQuery(), ScoreMode.None).innerHit(new InnerHitBuilder().setSize(2));
754+
AggregationBuilder topHitsAgg = topHits("top-children").size(3);
755+
756+
assertNoFailuresAndResponse(prepareSearch("idx").setQuery(hasChildQuery).addAggregation(topHitsAgg), response -> {
757+
assertHitCount(response, numDocs);
758+
759+
TopHits topHits = response.getAggregations().get("top-children");
760+
SearchHits hits = topHits.getHits();
761+
assertThat(hits.getHits().length, equalTo(3));
762+
763+
for (SearchHit hit : hits) {
764+
SearchHits innerHits = hit.getInnerHits().get("child2");
765+
assertThat(innerHits.getHits().length, lessThanOrEqualTo(2));
766+
}
767+
});
768+
}
701769
}

‎modules/parent-join/src/main/java/org/elasticsearch/join/query/ParentChildInnerHitContextBuilder.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,35 @@ static final class JoinFieldInnerHitSubContext extends InnerHitsContext.InnerHit
8888
private final String typeName;
8989
private final boolean fetchChildInnerHits;
9090
private final Joiner joiner;
91+
private final SearchExecutionContext searchExecutionContext;
9192

9293
JoinFieldInnerHitSubContext(String name, SearchContext context, String typeName, boolean fetchChildInnerHits, Joiner joiner) {
9394
super(name, context);
9495
this.typeName = typeName;
9596
this.fetchChildInnerHits = fetchChildInnerHits;
9697
this.joiner = joiner;
98+
this.searchExecutionContext = null;
99+
}
100+
101+
JoinFieldInnerHitSubContext(
102+
JoinFieldInnerHitSubContext joinFieldInnerHitSubContext,
103+
SearchExecutionContext searchExecutionContext
104+
) {
105+
super(joinFieldInnerHitSubContext);
106+
this.typeName = joinFieldInnerHitSubContext.typeName;
107+
this.fetchChildInnerHits = joinFieldInnerHitSubContext.fetchChildInnerHits;
108+
this.joiner = joinFieldInnerHitSubContext.joiner;
109+
this.searchExecutionContext = searchExecutionContext;
110+
}
111+
112+
@Override
113+
public JoinFieldInnerHitSubContext copyWithSearchExecutionContext(SearchExecutionContext searchExecutionContext) {
114+
return new JoinFieldInnerHitSubContext(this, searchExecutionContext);
115+
}
116+
117+
@Override
118+
public SearchExecutionContext getSearchExecutionContext() {
119+
return searchExecutionContext != null ? searchExecutionContext : super.getSearchExecutionContext();
97120
}
98121

99122
@Override

‎server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/metrics/TopHitsIT.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@
1919
import org.elasticsearch.common.settings.Settings;
2020
import org.elasticsearch.index.IndexSettings;
2121
import org.elasticsearch.index.fielddata.ScriptDocValues;
22+
import org.elasticsearch.index.query.InnerHitBuilder;
2223
import org.elasticsearch.index.query.MatchAllQueryBuilder;
24+
import org.elasticsearch.index.query.QueryBuilder;
2325
import org.elasticsearch.index.query.QueryBuilders;
2426
import org.elasticsearch.index.seqno.SequenceNumbers;
2527
import org.elasticsearch.plugins.Plugin;
@@ -30,6 +32,7 @@
3032
import org.elasticsearch.script.ScriptType;
3133
import org.elasticsearch.search.SearchHit;
3234
import org.elasticsearch.search.SearchHits;
35+
import org.elasticsearch.search.aggregations.AggregationBuilder;
3336
import org.elasticsearch.search.aggregations.Aggregator.SubAggCollectionMode;
3437
import org.elasticsearch.search.aggregations.BucketOrder;
3538
import org.elasticsearch.search.aggregations.InternalAggregation;
@@ -46,6 +49,7 @@
4649
import org.elasticsearch.search.lookup.FieldLookup;
4750
import org.elasticsearch.search.lookup.LeafSearchLookup;
4851
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
52+
import org.elasticsearch.search.sort.NestedSortBuilder;
4953
import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
5054
import org.elasticsearch.search.sort.SortBuilders;
5155
import org.elasticsearch.search.sort.SortOrder;
@@ -984,6 +988,67 @@ public void testTopHitsInNested() throws Exception {
984988
);
985989
}
986990

991+
public void testTopHitsOnInnerHits() {
992+
QueryBuilder nestedQuery = nestedQuery("comments", matchQuery("comments.message", "text"), ScoreMode.Avg).innerHit(
993+
new InnerHitBuilder().setSize(2)
994+
);
995+
AggregationBuilder topHitsAgg = topHits("top-comments").size(3)
996+
.sort(SortBuilders.fieldSort("comments.date").order(SortOrder.ASC).setNestedSort(new NestedSortBuilder("comments")));
997+
998+
assertNoFailuresAndResponse(prepareSearch("articles").setQuery(nestedQuery).addAggregation(topHitsAgg), response -> {
999+
TopHits topHits = response.getAggregations().get("top-comments");
1000+
SearchHits hits = topHits.getHits();
1001+
assertThat(hits.getHits().length, equalTo(3));
1002+
1003+
for (SearchHit hit : hits) {
1004+
SearchHits innerHits = hit.getInnerHits().get("comments");
1005+
assertThat(innerHits.getHits().length, lessThanOrEqualTo(2));
1006+
for (SearchHit innerHit : innerHits) {
1007+
assertThat(innerHit.getNestedIdentity().getField().string(), equalTo("comments"));
1008+
Map<String, Object> source = innerHit.getSourceAsMap();
1009+
assertTrue(source.containsKey("message"));
1010+
assertFalse(source.containsKey("reviewers"));
1011+
}
1012+
}
1013+
});
1014+
}
1015+
1016+
public void testTopHitsOnMultipleNestedInnerHits() {
1017+
QueryBuilder doubleNestedQuery = nestedQuery(
1018+
"comments",
1019+
nestedQuery("comments.reviewers", matchQuery("comments.reviewers.name", "user c"), ScoreMode.Avg).innerHit(
1020+
new InnerHitBuilder()
1021+
),
1022+
ScoreMode.Avg
1023+
).innerHit(new InnerHitBuilder("review"));
1024+
AggregationBuilder topHitsAgg = topHits("top-reviewers").size(2)
1025+
.sort(SortBuilders.fieldSort("comments.date").order(SortOrder.ASC).setNestedSort(new NestedSortBuilder("comments")));
1026+
1027+
assertNoFailuresAndResponse(prepareSearch("articles").setQuery(doubleNestedQuery).addAggregation(topHitsAgg), response -> {
1028+
TopHits topHits = response.getAggregations().get("top-reviewers");
1029+
SearchHits hits = topHits.getHits();
1030+
assertThat(hits.getHits().length, equalTo(1));
1031+
1032+
SearchHit hit = hits.getAt(0);
1033+
SearchHits innerHits = hit.getInnerHits().get("review");
1034+
assertThat(innerHits.getHits().length, equalTo(2));
1035+
1036+
assertThat(innerHits.getAt(0).getId(), equalTo("1"));
1037+
assertThat(innerHits.getAt(0).getNestedIdentity().getField().string(), equalTo("comments"));
1038+
assertThat(innerHits.getAt(0).getNestedIdentity().getOffset(), equalTo(0));
1039+
Map<String, Object> source0 = innerHits.getAt(0).getSourceAsMap();
1040+
assertTrue(source0.containsKey("message"));
1041+
assertTrue(source0.containsKey("reviewers"));
1042+
1043+
assertThat(innerHits.getAt(1).getId(), equalTo("1"));
1044+
assertThat(innerHits.getAt(1).getNestedIdentity().getField().string(), equalTo("comments"));
1045+
assertThat(innerHits.getAt(1).getNestedIdentity().getOffset(), equalTo(1));
1046+
Map<String, Object> source1 = innerHits.getAt(1).getSourceAsMap();
1047+
assertTrue(source1.containsKey("message"));
1048+
assertTrue(source1.containsKey("reviewers"));
1049+
});
1050+
}
1051+
9871052
public void testUseMaxDocInsteadOfSize() throws Exception {
9881053
updateIndexSettings(
9891054
Settings.builder().put(IndexSettings.MAX_INNER_RESULT_WINDOW_SETTING.getKey(), ArrayUtil.MAX_ARRAY_LENGTH),

‎server/src/main/java/org/elasticsearch/index/query/NestedQueryBuilder.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -397,6 +397,7 @@ static final class NestedInnerHitSubContext extends InnerHitsContext.InnerHitSub
397397

398398
private final NestedObjectMapper parentObjectMapper;
399399
private final NestedObjectMapper childObjectMapper;
400+
private final SearchExecutionContext searchExecutionContext;
400401

401402
NestedInnerHitSubContext(
402403
String name,
@@ -407,6 +408,24 @@ static final class NestedInnerHitSubContext extends InnerHitsContext.InnerHitSub
407408
super(name, context);
408409
this.parentObjectMapper = parentObjectMapper;
409410
this.childObjectMapper = childObjectMapper;
411+
this.searchExecutionContext = null;
412+
}
413+
414+
NestedInnerHitSubContext(NestedInnerHitSubContext nestedInnerHitSubContext, SearchExecutionContext searchExecutionContext) {
415+
super(nestedInnerHitSubContext);
416+
this.parentObjectMapper = nestedInnerHitSubContext.parentObjectMapper;
417+
this.childObjectMapper = nestedInnerHitSubContext.childObjectMapper;
418+
this.searchExecutionContext = searchExecutionContext;
419+
}
420+
421+
@Override
422+
public NestedInnerHitSubContext copyWithSearchExecutionContext(SearchExecutionContext searchExecutionContext) {
423+
return new NestedInnerHitSubContext(this, searchExecutionContext);
424+
}
425+
426+
@Override
427+
public SearchExecutionContext getSearchExecutionContext() {
428+
return searchExecutionContext != null ? searchExecutionContext : super.getSearchExecutionContext();
410429
}
411430

412431
@Override

‎server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregator.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939
import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
4040
import org.elasticsearch.search.aggregations.support.AggregationContext;
4141
import org.elasticsearch.search.fetch.FetchSearchResult;
42+
import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
43+
import org.elasticsearch.search.fetch.subphase.InnerHitsContext.InnerHitSubContext;
4244
import org.elasticsearch.search.internal.SubSearchContext;
4345
import org.elasticsearch.search.profile.ProfileResult;
4446
import org.elasticsearch.search.rescore.RescoreContext;
@@ -219,16 +221,40 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOE
219221
private static FetchSearchResult runFetchPhase(SubSearchContext subSearchContext, int[] docIdsToLoad) {
220222
// Fork the search execution context for each slice, because the fetch phase does not support concurrent execution yet.
221223
SearchExecutionContext searchExecutionContext = new SearchExecutionContext(subSearchContext.getSearchExecutionContext());
224+
// InnerHitSubContext is not thread-safe, so we fork it as well to support concurrent execution
225+
InnerHitsContext innerHitsContext = new InnerHitsContext(
226+
getForkedInnerHits(subSearchContext.innerHits().getInnerHits(), searchExecutionContext)
227+
);
228+
222229
SubSearchContext fetchSubSearchContext = new SubSearchContext(subSearchContext) {
223230
@Override
224231
public SearchExecutionContext getSearchExecutionContext() {
225232
return searchExecutionContext;
226233
}
234+
235+
@Override
236+
public InnerHitsContext innerHits() {
237+
return innerHitsContext;
238+
}
227239
};
240+
228241
fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null);
229242
return fetchSubSearchContext.fetchResult();
230243
}
231244

245+
private static Map<String, InnerHitSubContext> getForkedInnerHits(
246+
Map<String, InnerHitSubContext> originalInnerHits,
247+
SearchExecutionContext searchExecutionContext
248+
) {
249+
Map<String, InnerHitSubContext> forkedInnerHits = new HashMap<>();
250+
for (Map.Entry<String, InnerHitSubContext> entry : originalInnerHits.entrySet()) {
251+
var forkedContext = entry.getValue().copyWithSearchExecutionContext(searchExecutionContext);
252+
forkedInnerHits.put(entry.getKey(), forkedContext);
253+
}
254+
255+
return forkedInnerHits;
256+
}
257+
232258
@Override
233259
public InternalTopHits buildEmptyAggregation() {
234260
TopDocs topDocs;

‎server/src/main/java/org/elasticsearch/search/fetch/subphase/InnerHitsContext.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import org.apache.lucene.search.Weight;
2323
import org.apache.lucene.util.Bits;
2424
import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
25+
import org.elasticsearch.index.query.SearchExecutionContext;
2526
import org.elasticsearch.search.SearchHit;
2627
import org.elasticsearch.search.internal.SearchContext;
2728
import org.elasticsearch.search.internal.SubSearchContext;
@@ -44,7 +45,7 @@ public InnerHitsContext() {
4445
this.innerHits = new HashMap<>();
4546
}
4647

47-
InnerHitsContext(Map<String, InnerHitSubContext> innerHits) {
48+
public InnerHitsContext(Map<String, InnerHitSubContext> innerHits) {
4849
this.innerHits = Objects.requireNonNull(innerHits);
4950
}
5051

@@ -84,6 +85,14 @@ protected InnerHitSubContext(String name, SearchContext context) {
8485
this.context = context;
8586
}
8687

88+
public InnerHitSubContext(InnerHitSubContext innerHitSubContext) {
89+
super(innerHitSubContext);
90+
this.name = innerHitSubContext.name;
91+
this.context = innerHitSubContext.context;
92+
}
93+
94+
public abstract InnerHitSubContext copyWithSearchExecutionContext(SearchExecutionContext searchExecutionContext);
95+
8796
public abstract TopDocsAndMaxScore topDocs(SearchHit hit) throws IOException;
8897

8998
public String getName() {

‎test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@
143143
import org.elasticsearch.search.fetch.FetchPhase;
144144
import org.elasticsearch.search.fetch.subphase.FetchDocValuesPhase;
145145
import org.elasticsearch.search.fetch.subphase.FetchSourcePhase;
146+
import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
146147
import org.elasticsearch.search.internal.AliasFilter;
147148
import org.elasticsearch.search.internal.ContextIndexSearcher;
148149
import org.elasticsearch.search.internal.SearchContext;
@@ -515,6 +516,7 @@ private SubSearchContext buildSubSearchContext(
515516
when(ctx.indexShard()).thenReturn(indexShard);
516517
when(ctx.newSourceLoader(null)).thenAnswer(inv -> searchExecutionContext.newSourceLoader(null, false));
517518
when(ctx.newIdLoader()).thenReturn(IdLoader.fromLeafStoredFieldLoader());
519+
when(ctx.innerHits()).thenReturn(new InnerHitsContext());
518520
var res = new SubSearchContext(ctx);
519521
releasables.add(res); // TODO: nasty workaround for not getting the standard resource handling behavior of a real search context
520522
return res;

0 commit comments

Comments
 (0)