Skip to content

Commit 13bce60

Browse files
benchaplinbenchaplin
authored
Fix inner hits + aggregations concurrency bug (#128036)
Fork InnerHitSubContext instances before source is fetched in aggregations to prevent inter-segment race conditions. Relates to #122419
1 parent 4762111 commit 13bce60

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;
@@ -45,6 +48,7 @@
4548
import org.elasticsearch.search.lookup.FieldLookup;
4649
import org.elasticsearch.search.lookup.LeafSearchLookup;
4750
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
51+
import org.elasticsearch.search.sort.NestedSortBuilder;
4852
import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
4953
import org.elasticsearch.search.sort.SortBuilders;
5054
import org.elasticsearch.search.sort.SortOrder;
@@ -983,6 +987,67 @@ public void testTopHitsInNested() throws Exception {
983987
);
984988
}
985989

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

397397
private final NestedObjectMapper parentObjectMapper;
398398
private final NestedObjectMapper childObjectMapper;
399+
private final SearchExecutionContext searchExecutionContext;
399400

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

411430
@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
@@ -41,6 +41,8 @@
4141
import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
4242
import org.elasticsearch.search.aggregations.support.AggregationContext;
4343
import org.elasticsearch.search.fetch.FetchSearchResult;
44+
import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
45+
import org.elasticsearch.search.fetch.subphase.InnerHitsContext.InnerHitSubContext;
4446
import org.elasticsearch.search.internal.SubSearchContext;
4547
import org.elasticsearch.search.profile.ProfileResult;
4648
import org.elasticsearch.search.rescore.RescoreContext;
@@ -223,16 +225,40 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOE
223225
private static FetchSearchResult runFetchPhase(SubSearchContext subSearchContext, int[] docIdsToLoad) {
224226
// Fork the search execution context for each slice, because the fetch phase does not support concurrent execution yet.
225227
SearchExecutionContext searchExecutionContext = new SearchExecutionContext(subSearchContext.getSearchExecutionContext());
228+
// InnerHitSubContext is not thread-safe, so we fork it as well to support concurrent execution
229+
InnerHitsContext innerHitsContext = new InnerHitsContext(
230+
getForkedInnerHits(subSearchContext.innerHits().getInnerHits(), searchExecutionContext)
231+
);
232+
226233
SubSearchContext fetchSubSearchContext = new SubSearchContext(subSearchContext) {
227234
@Override
228235
public SearchExecutionContext getSearchExecutionContext() {
229236
return searchExecutionContext;
230237
}
238+
239+
@Override
240+
public InnerHitsContext innerHits() {
241+
return innerHitsContext;
242+
}
231243
};
244+
232245
fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null);
233246
return fetchSubSearchContext.fetchResult();
234247
}
235248

249+
private static Map<String, InnerHitSubContext> getForkedInnerHits(
250+
Map<String, InnerHitSubContext> originalInnerHits,
251+
SearchExecutionContext searchExecutionContext
252+
) {
253+
Map<String, InnerHitSubContext> forkedInnerHits = new HashMap<>();
254+
for (Map.Entry<String, InnerHitSubContext> entry : originalInnerHits.entrySet()) {
255+
var forkedContext = entry.getValue().copyWithSearchExecutionContext(searchExecutionContext);
256+
forkedInnerHits.put(entry.getKey(), forkedContext);
257+
}
258+
259+
return forkedInnerHits;
260+
}
261+
236262
@Override
237263
public InternalTopHits buildEmptyAggregation() {
238264
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
@@ -145,6 +145,7 @@
145145
import org.elasticsearch.search.fetch.FetchPhase;
146146
import org.elasticsearch.search.fetch.subphase.FetchDocValuesPhase;
147147
import org.elasticsearch.search.fetch.subphase.FetchSourcePhase;
148+
import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
148149
import org.elasticsearch.search.internal.AliasFilter;
149150
import org.elasticsearch.search.internal.ContextIndexSearcher;
150151
import org.elasticsearch.search.internal.SearchContext;
@@ -525,6 +526,7 @@ private SubSearchContext buildSubSearchContext(
525526
when(ctx.indexShard()).thenReturn(indexShard);
526527
when(ctx.newSourceLoader(null)).thenAnswer(inv -> searchExecutionContext.newSourceLoader(null, false));
527528
when(ctx.newIdLoader()).thenReturn(IdLoader.fromLeafStoredFieldLoader());
529+
when(ctx.innerHits()).thenReturn(new InnerHitsContext());
528530
var res = new SubSearchContext(ctx);
529531
releasables.add(res); // TODO: nasty workaround for not getting the standard resource handling behavior of a real search context
530532
return res;

0 commit comments

Comments
 (0)