Skip to content

Commit

Permalink
Fix more search response leaks (#103956)
Browse files Browse the repository at this point in the history
Some more mechanical fixing of leaked SearchResponse instances.
  • Loading branch information
original-brownbear authored Jan 5, 2024
1 parent 8efc72c commit 80a9508
Show file tree
Hide file tree
Showing 4 changed files with 407 additions and 275 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@

import org.apache.lucene.search.Query;
import org.apache.lucene.search.join.ScoreMode;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.bytes.BytesReference;
Expand Down Expand Up @@ -50,6 +49,7 @@
import static org.elasticsearch.index.query.QueryBuilders.scriptQuery;
import static org.elasticsearch.index.query.QueryBuilders.termQuery;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchHitsWithoutFailures;
import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
import static org.hamcrest.Matchers.equalTo;
Expand Down Expand Up @@ -407,25 +407,26 @@ public void testPercolateNamedQueries() {
""");

QueryBuilder query = new PercolateQueryBuilder("my_query", List.of(house1_doc, house2_doc), XContentType.JSON);
SearchResponse response = client().prepareSearch("houses").setQuery(query).get();
assertEquals(2, response.getHits().getTotalHits().value);
assertResponse(client().prepareSearch("houses").setQuery(query), response -> {
assertEquals(2, response.getHits().getTotalHits().value);

SearchHit[] hits = response.getHits().getHits();
assertThat(hits[0].getFields().get("_percolator_document_slot").getValues(), equalTo(Arrays.asList(0, 1)));
assertThat(
hits[0].getFields().get("_percolator_document_slot_0_matched_queries").getValues(),
equalTo(Arrays.asList("fireplace_query", "detached_query", "3_bedrooms_query"))
);
assertThat(
hits[0].getFields().get("_percolator_document_slot_1_matched_queries").getValues(),
equalTo(Arrays.asList("fireplace_query", "3_bedrooms_query"))
);
SearchHit[] hits = response.getHits().getHits();
assertThat(hits[0].getFields().get("_percolator_document_slot").getValues(), equalTo(Arrays.asList(0, 1)));
assertThat(
hits[0].getFields().get("_percolator_document_slot_0_matched_queries").getValues(),
equalTo(Arrays.asList("fireplace_query", "detached_query", "3_bedrooms_query"))
);
assertThat(
hits[0].getFields().get("_percolator_document_slot_1_matched_queries").getValues(),
equalTo(Arrays.asList("fireplace_query", "3_bedrooms_query"))
);

assertThat(hits[1].getFields().get("_percolator_document_slot").getValues(), equalTo(Arrays.asList(0)));
assertThat(
hits[1].getFields().get("_percolator_document_slot_0_matched_queries").getValues(),
equalTo(Arrays.asList("swimming_pool_query", "3_bedrooms_query"))
);
assertThat(hits[1].getFields().get("_percolator_document_slot").getValues(), equalTo(Arrays.asList(0)));
assertThat(
hits[1].getFields().get("_percolator_document_slot_0_matched_queries").getValues(),
equalTo(Arrays.asList("swimming_pool_query", "3_bedrooms_query"))
);
});
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -589,30 +589,33 @@ protected RequestWrapper<?> buildRequest(Hit doc) {
null,
SearchResponse.Clusters.EMPTY
);
try {
client.lastSearch.get().listener.onResponse(searchResponse);

client.lastSearch.get().listener.onResponse(searchResponse);
assertEquals(0, capturedDelay.get().seconds());
capturedCommand.get().run();

assertEquals(0, capturedDelay.get().seconds());
capturedCommand.get().run();
// So the next request is going to have to wait an extra 100 seconds or so (base was 10 seconds, so 110ish)
assertThat(client.lastScroll.get().request.scroll().keepAlive().seconds(), either(equalTo(110L)).or(equalTo(109L)));

// So the next request is going to have to wait an extra 100 seconds or so (base was 10 seconds, so 110ish)
assertThat(client.lastScroll.get().request.scroll().keepAlive().seconds(), either(equalTo(110L)).or(equalTo(109L)));
// Now we can simulate a response and check the delay that we used for the task
if (randomBoolean()) {
client.lastScroll.get().listener.onResponse(searchResponse);
assertEquals(99, capturedDelay.get().seconds());
} else {
// Let's rethrottle between the starting the scroll and getting the response
worker.rethrottle(10f);
client.lastScroll.get().listener.onResponse(searchResponse);
// The delay uses the new throttle
assertEquals(9, capturedDelay.get().seconds());
}

// Now we can simulate a response and check the delay that we used for the task
if (randomBoolean()) {
client.lastScroll.get().listener.onResponse(searchResponse);
assertEquals(99, capturedDelay.get().seconds());
} else {
// Let's rethrottle between the starting the scroll and getting the response
worker.rethrottle(10f);
client.lastScroll.get().listener.onResponse(searchResponse);
// The delay uses the new throttle
assertEquals(9, capturedDelay.get().seconds());
// Running the command ought to increment the delay counter on the task.
capturedCommand.get().run();
assertEquals(capturedDelay.get(), testTask.getStatus().getThrottled());
} finally {
searchResponse.decRef();
}

// Running the command ought to increment the delay counter on the task.
capturedCommand.get().run();
assertEquals(capturedDelay.get(), testTask.getStatus().getThrottled());
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.metrics.InternalHDRPercentiles;
Expand Down Expand Up @@ -211,11 +210,8 @@ private void setupTDigestHistogram(int compression) throws Exception {
}
client().admin().indices().refresh(new RefreshRequest("raw", "pre_agg")).get();

SearchResponse response = client().prepareSearch("raw").setTrackTotalHits(true).get();
assertEquals(numDocs, response.getHits().getTotalHits().value);

response = client().prepareSearch("pre_agg").get();
assertEquals(numDocs / frq, response.getHits().getTotalHits().value);
assertHitCount(client().prepareSearch("raw").setTrackTotalHits(true), numDocs);
assertHitCount(client().prepareSearch("pre_agg"), numDocs / frq);
}

public void testTDigestHistogram() throws Exception {
Expand All @@ -228,42 +224,50 @@ public void testTDigestHistogram() throws Exception {
.compression(compression)
.percentiles(10, 25, 50, 75);

SearchResponse responseRaw = client().prepareSearch("raw").addAggregation(builder).get();
SearchResponse responsePreAgg = client().prepareSearch("pre_agg").addAggregation(builder).get();
SearchResponse responseBoth = client().prepareSearch("raw", "pre_agg").addAggregation(builder).get();

InternalTDigestPercentiles percentilesRaw = responseRaw.getAggregations().get("agg");
InternalTDigestPercentiles percentilesPreAgg = responsePreAgg.getAggregations().get("agg");
InternalTDigestPercentiles percentilesBoth = responseBoth.getAggregations().get("agg");
for (int i = 1; i < 100; i++) {
assertEquals(percentilesRaw.percentile(i), percentilesPreAgg.percentile(i), 1.0);
assertEquals(percentilesRaw.percentile(i), percentilesBoth.percentile(i), 1.0);
}
assertResponse(
client().prepareSearch("raw").addAggregation(builder),
responseRaw -> assertResponse(
client().prepareSearch("pre_agg").addAggregation(builder),
responsePreAgg -> assertResponse(client().prepareSearch("raw", "pre_agg").addAggregation(builder), responseBoth -> {
InternalTDigestPercentiles percentilesRaw = responseRaw.getAggregations().get("agg");
InternalTDigestPercentiles percentilesPreAgg = responsePreAgg.getAggregations().get("agg");
InternalTDigestPercentiles percentilesBoth = responseBoth.getAggregations().get("agg");
for (int i = 1; i < 100; i++) {
assertEquals(percentilesRaw.percentile(i), percentilesPreAgg.percentile(i), 1.0);
assertEquals(percentilesRaw.percentile(i), percentilesBoth.percentile(i), 1.0);
}
})
)
);
}

public void testBoxplotHistogram() throws Exception {
int compression = TestUtil.nextInt(random(), 200, 300);
setupTDigestHistogram(compression);
BoxplotAggregationBuilder bpBuilder = new BoxplotAggregationBuilder("agg").field("inner.data").compression(compression);

SearchResponse bpResponseRaw = client().prepareSearch("raw").addAggregation(bpBuilder).get();
SearchResponse bpResponsePreAgg = client().prepareSearch("pre_agg").addAggregation(bpBuilder).get();
SearchResponse bpResponseBoth = client().prepareSearch("raw", "pre_agg").addAggregation(bpBuilder).get();

Boxplot bpRaw = bpResponseRaw.getAggregations().get("agg");
Boxplot bpPreAgg = bpResponsePreAgg.getAggregations().get("agg");
Boxplot bpBoth = bpResponseBoth.getAggregations().get("agg");
assertEquals(bpRaw.getMax(), bpPreAgg.getMax(), 0.0);
assertEquals(bpRaw.getMax(), bpBoth.getMax(), 0.0);
assertEquals(bpRaw.getMin(), bpPreAgg.getMin(), 0.0);
assertEquals(bpRaw.getMin(), bpBoth.getMin(), 0.0);
assertResponse(
client().prepareSearch("raw").addAggregation(bpBuilder),
bpResponseRaw -> assertResponse(
client().prepareSearch("pre_agg").addAggregation(bpBuilder),
bpResponsePreAgg -> assertResponse(client().prepareSearch("raw", "pre_agg").addAggregation(bpBuilder), bpResponseBoth -> {
Boxplot bpRaw = bpResponseRaw.getAggregations().get("agg");
Boxplot bpPreAgg = bpResponsePreAgg.getAggregations().get("agg");
Boxplot bpBoth = bpResponseBoth.getAggregations().get("agg");
assertEquals(bpRaw.getMax(), bpPreAgg.getMax(), 0.0);
assertEquals(bpRaw.getMax(), bpBoth.getMax(), 0.0);
assertEquals(bpRaw.getMin(), bpPreAgg.getMin(), 0.0);
assertEquals(bpRaw.getMin(), bpBoth.getMin(), 0.0);

assertEquals(bpRaw.getQ1(), bpPreAgg.getQ1(), 1.0);
assertEquals(bpRaw.getQ1(), bpBoth.getQ1(), 1.0);
assertEquals(bpRaw.getQ2(), bpPreAgg.getQ2(), 1.0);
assertEquals(bpRaw.getQ2(), bpBoth.getQ2(), 1.0);
assertEquals(bpRaw.getQ3(), bpPreAgg.getQ3(), 1.0);
assertEquals(bpRaw.getQ3(), bpBoth.getQ3(), 1.0);
assertEquals(bpRaw.getQ1(), bpPreAgg.getQ1(), 1.0);
assertEquals(bpRaw.getQ1(), bpBoth.getQ1(), 1.0);
assertEquals(bpRaw.getQ2(), bpPreAgg.getQ2(), 1.0);
assertEquals(bpRaw.getQ2(), bpBoth.getQ2(), 1.0);
assertEquals(bpRaw.getQ3(), bpPreAgg.getQ3(), 1.0);
assertEquals(bpRaw.getQ3(), bpBoth.getQ3(), 1.0);
})
)
);
}

@Override
Expand Down
Loading

0 comments on commit 80a9508

Please sign in to comment.