From a370caca0fb026890e6a8a893591658fec6a6fb6 Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Wed, 10 Jul 2024 19:15:44 +0100 Subject: [PATCH] address review comments --- .../search/retriever/RetrieverRewriteIT.java | 76 +++---------------- .../action/search/TransportSearchAction.java | 16 ++-- 2 files changed, 22 insertions(+), 70 deletions(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java index 4eb4b2248bce5..8f28e5f279cf6 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java @@ -12,32 +12,28 @@ import org.apache.lucene.util.SetOnce; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.search.SearchRequest; +import org.elasticsearch.action.search.SearchRequestBuilder; import org.elasticsearch.action.search.SearchResponse; import org.elasticsearch.index.query.QueryBuilders; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.plugins.Plugin; -import org.elasticsearch.plugins.SearchPlugin; import org.elasticsearch.search.MockSearchService; import org.elasticsearch.search.builder.SearchSourceBuilder; import org.elasticsearch.test.ESIntegTestCase; -import org.elasticsearch.xcontent.ConstructingObjectParser; -import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.test.hamcrest.ElasticsearchAssertions; import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xcontent.XContentParser; import org.junit.Before; import java.io.IOException; import java.util.Collection; import java.util.List; -import java.util.concurrent.ExecutionException; -import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; import static org.hamcrest.Matchers.equalTo; public class RetrieverRewriteIT extends ESIntegTestCase { @Override protected Collection> nodePlugins() { - return List.of(AssertingPlugin.class, MockSearchService.TestPlugin.class); + return List.of(MockSearchService.TestPlugin.class); } private static String INDEX_DOCS = "docs"; @@ -60,78 +56,35 @@ public void setup() throws Exception { refresh(INDEX_QUERIES); } - public void testRewrite() throws ExecutionException, InterruptedException { + public void testRewrite() { SearchSourceBuilder source = new SearchSourceBuilder(); StandardRetrieverBuilder standard = new StandardRetrieverBuilder(); standard.queryBuilder = QueryBuilders.termQuery(ID_FIELD, "doc_0"); source.retriever(new AssertingRetrieverBuilder(standard)); - SearchRequest req = new SearchRequest(INDEX_DOCS, INDEX_QUERIES).source(source); - SearchResponse resp = client().search(req).get(); - try { + SearchRequestBuilder req = client().prepareSearch(INDEX_DOCS, INDEX_QUERIES).setSource(source); + ElasticsearchAssertions.assertResponse(req, resp -> { assertNull(resp.pointInTimeId()); assertThat(resp.getHits().getTotalHits().value, equalTo(1L)); assertThat(resp.getHits().getTotalHits().relation, equalTo(TotalHits.Relation.EQUAL_TO)); assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_0")); - } finally { - resp.decRef(); - } + }); } - public void testRewriteCompound() throws ExecutionException, InterruptedException { + public void testRewriteCompound() { SearchSourceBuilder source = new SearchSourceBuilder(); source.retriever(new AssertingCompoundRetrieverBuilder("query_0")); - SearchRequest req = new SearchRequest(INDEX_DOCS, INDEX_QUERIES).source(source); - SearchResponse resp = client().search(req).get(); - try { + SearchRequestBuilder req = client().prepareSearch(INDEX_DOCS, INDEX_QUERIES).setSource(source); + ElasticsearchAssertions.assertResponse(req, resp -> { assertNull(resp.pointInTimeId()); assertThat(resp.getHits().getTotalHits().value, equalTo(1L)); assertThat(resp.getHits().getTotalHits().relation, equalTo(TotalHits.Relation.EQUAL_TO)); assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_2")); - } finally { - resp.decRef(); - } - } - - public static class AssertingPlugin extends Plugin implements SearchPlugin { - public AssertingPlugin() {} - - @Override - public List> getRetrievers() { - return List.of( - new RetrieverSpec(AssertingRetrieverBuilder.NAME, AssertingRetrieverBuilder::fromXContent), - new RetrieverSpec(AssertingCompoundRetrieverBuilder.NAME, AssertingCompoundRetrieverBuilder::fromXContent) - ); - } - } - - public static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( - AssertingCompoundRetrieverBuilder.NAME, - args -> new AssertingRetrieverBuilder((RetrieverBuilder) args[0]) - ); - - public static final ConstructingObjectParser PARSER_COMPOUND = - new ConstructingObjectParser<>( - AssertingCompoundRetrieverBuilder.NAME, - args -> new AssertingCompoundRetrieverBuilder((String) args[0]) - ); - - static { - RetrieverBuilder.declareBaseParserFields(AssertingRetrieverBuilder.NAME, PARSER); - PARSER.declareObject(constructorArg(), RetrieverBuilder::parseInnerRetrieverBuilder, new ParseField("retriever")); - - RetrieverBuilder.declareBaseParserFields(AssertingCompoundRetrieverBuilder.NAME, PARSER_COMPOUND); - PARSER_COMPOUND.declareString(constructorArg(), new ParseField("id")); + }); } private static class AssertingRetrieverBuilder extends RetrieverBuilder { - static final String NAME = "asserting"; - private final RetrieverBuilder innerRetriever; - public static AssertingRetrieverBuilder fromXContent(XContentParser parser, RetrieverParserContext context) throws IOException { - return PARSER.apply(parser, context); - } - private AssertingRetrieverBuilder(RetrieverBuilder innerRetriever) { this.innerRetriever = innerRetriever; } @@ -177,16 +130,9 @@ protected int doHashCode() { } private static class AssertingCompoundRetrieverBuilder extends RetrieverBuilder { - static final String NAME = "asserting_compound"; - private final String id; private final SetOnce innerRetriever; - public static AssertingCompoundRetrieverBuilder fromXContent(XContentParser parser, RetrieverParserContext context) - throws IOException { - return PARSER_COMPOUND.apply(parser, context); - } - private AssertingCompoundRetrieverBuilder(String id) { this.id = id; this.innerRetriever = new SetOnce<>(null); diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java index eaddfc6d6592e..370bd49c0523b 100644 --- a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java @@ -490,19 +490,25 @@ void executeRequest( }); final SearchSourceBuilder source = original.source(); if (shouldOpenPIT(source)) { - openPIT(original, searchService.getDefaultKeepAliveInMillis(), listener.delegateFailureAndWrap((delegate, resp) -> { + openPIT(client, original, searchService.getDefaultKeepAliveInMillis(), listener.delegateFailureAndWrap((delegate, resp) -> { + // We set the keep alive to -1 to indicate that we don't need the pit id in the response. + // This is needed since we delete the pit prior to sending the response so the id doesn't exist anymore. source.pointInTimeBuilder(new PointInTimeBuilder(resp.getPointInTimeId()).setKeepAlive(TimeValue.MINUS_ONE)); executeRequest(task, original, new ActionListener<>() { @Override public void onResponse(SearchResponse response) { // we need to close the PIT first so we delay the release of the response to after the closing response.incRef(); - closePIT(original.source().pointInTimeBuilder(), () -> ActionListener.respondAndRelease(listener, response)); + closePIT( + client, + original.source().pointInTimeBuilder(), + () -> ActionListener.respondAndRelease(listener, response) + ); } @Override public void onFailure(Exception e) { - closePIT(original.source().pointInTimeBuilder(), () -> listener.onFailure(e)); + closePIT(client, original.source().pointInTimeBuilder(), () -> listener.onFailure(e)); } }, searchPhaseProvider); })); @@ -529,7 +535,7 @@ private boolean shouldOpenPIT(SearchSourceBuilder source) { return retriever != null && retriever.isCompound(); } - private void openPIT(SearchRequest request, long keepAliveMillis, ActionListener listener) { + static void openPIT(Client client, SearchRequest request, long keepAliveMillis, ActionListener listener) { OpenPointInTimeRequest pitReq = new OpenPointInTimeRequest(request.indices()).indicesOptions(request.indicesOptions()) .preference(request.preference()) .routing(request.routing()) @@ -537,7 +543,7 @@ private void openPIT(SearchRequest request, long keepAliveMillis, ActionListener client.execute(TransportOpenPointInTimeAction.TYPE, pitReq, listener); } - private void closePIT(PointInTimeBuilder pit, Runnable next) { + static void closePIT(Client client, PointInTimeBuilder pit, Runnable next) { client.execute( TransportClosePointInTimeAction.TYPE, new ClosePointInTimeRequest(pit.getEncodedId()),