Remove blocking code in SearchDocument processing.

Original Pull Request #2094
Closes #2025

(cherry picked from commit c1a1ea9)

Author: sothawo

Diff for: src/main/java/org/springframework/data/elasticsearch/core/AbstractElasticsearchRestTransportTemplate.java

@@ -20,6 +20,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -109,7 +110,8 @@ public <T> List<SearchHits<T>> multiSearch(List<? extends Query> queries, Class<
 		List<SearchHits<T>> res = new ArrayList<>(queries.size());
 		int c = 0;
 		for (Query query : queries) {
-			res.add(callback.doWith(SearchDocumentResponse.from(items[c++].getResponse(), documentCallback::doWith)));
+			res.add(
+					callback.doWith(SearchDocumentResponse.from(items[c++].getResponse(), getEntityCreator(documentCallback))));
 		}
 		return res;
 	}
@@ -142,7 +144,7 @@ public List<SearchHits<?>> multiSearch(List<? extends Query> queries, List<Class
 					index);
 
 			SearchResponse response = items[c++].getResponse();
-			res.add(callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith)));
+			res.add(callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback))));
 		}
 		return res;
 	}
@@ -175,7 +177,7 @@ public List<SearchHits<?>> multiSearch(List<? extends Query> queries, List<Class
 					index);
 
 			SearchResponse response = items[c++].getResponse();
-			res.add(callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith)));
+			res.add(callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback))));
 		}
 		return res;
 	}
@@ -215,5 +217,9 @@ public SearchResponse suggest(SuggestBuilder suggestion, Class<?> clazz) {
 		return suggest(suggestion, getIndexCoordinatesFor(clazz));
 	}
 
+	protected <T> SearchDocumentResponse.EntityCreator<T> getEntityCreator(ReadDocumentCallback<T> documentCallback) {
+		return searchDocument -> CompletableFuture.completedFuture(documentCallback.doWith(searchDocument));
+	}
+
 	// endregion
 }

Diff for: src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 the original author or authors.
+ * Copyright 2013-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -319,7 +319,7 @@ public <T> SearchHits<T> search(Query query, Class<T> clazz, IndexCoordinates in
 		ReadDocumentCallback<T> documentCallback = new ReadDocumentCallback<T>(elasticsearchConverter, clazz, index);
 		SearchDocumentResponseCallback<SearchHits<T>> callback = new ReadSearchDocumentResponseCallback<>(clazz, index);
 
-		return callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith));
+		return callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback)));
 	}
 
 	@Override
@@ -336,7 +336,7 @@ public <T> SearchScrollHits<T> searchScrollStart(long scrollTimeInMillis, Query
 		ReadDocumentCallback<T> documentCallback = new ReadDocumentCallback<T>(elasticsearchConverter, clazz, index);
 		SearchDocumentResponseCallback<SearchScrollHits<T>> callback = new ReadSearchScrollDocumentResponseCallback<>(clazz,
 				index);
-		return callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith));
+		return callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback)));
 	}
 
 	@Override
@@ -351,7 +351,7 @@ public <T> SearchScrollHits<T> searchScrollContinue(@Nullable String scrollId, l
 		ReadDocumentCallback<T> documentCallback = new ReadDocumentCallback<T>(elasticsearchConverter, clazz, index);
 		SearchDocumentResponseCallback<SearchScrollHits<T>> callback = new ReadSearchScrollDocumentResponseCallback<>(clazz,
 				index);
-		return callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith));
+		return callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback)));
 	}
 
 	@Override
@@ -378,8 +378,8 @@ protected MultiSearchResponse.Item[] getMultiSearchResult(MultiSearchRequest req
 		Assert.isTrue(items.length == request.requests().size(), "Response should has same length with queries");
 		return items;
 	}
-	// endregion
 
+	// endregion
 	// region ClientCallback
 	/**
 	 * Callback interface to be used with {@link #execute(ClientCallback)} for operating directly on

Diff for: src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java

@@ -354,7 +354,7 @@ public <T> SearchHits<T> search(Query query, Class<T> clazz, IndexCoordinates in
 
 		ReadDocumentCallback<T> documentCallback = new ReadDocumentCallback<T>(elasticsearchConverter, clazz, index);
 		SearchDocumentResponseCallback<SearchHits<T>> callback = new ReadSearchDocumentResponseCallback<>(clazz, index);
-		return callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith));
+		return callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback)));
 	}
 
 	@Override
@@ -372,7 +372,7 @@ public <T> SearchScrollHits<T> searchScrollStart(long scrollTimeInMillis, Query
 		ReadDocumentCallback<T> documentCallback = new ReadDocumentCallback<T>(elasticsearchConverter, clazz, index);
 		SearchDocumentResponseCallback<SearchScrollHits<T>> callback = new ReadSearchScrollDocumentResponseCallback<>(clazz,
 				index);
-		return callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith));
+		return callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback)));
 	}
 
 	@Override
@@ -389,7 +389,7 @@ public <T> SearchScrollHits<T> searchScrollContinue(@Nullable String scrollId, l
 		ReadDocumentCallback<T> documentCallback = new ReadDocumentCallback<T>(elasticsearchConverter, clazz, index);
 		SearchDocumentResponseCallback<SearchScrollHits<T>> callback = new ReadSearchScrollDocumentResponseCallback<>(clazz,
 				index);
-		return callback.doWith(SearchDocumentResponse.from(response, documentCallback::doWith));
+		return callback.doWith(SearchDocumentResponse.from(response, getEntityCreator(documentCallback)));
 	}
 
 	@Override

Diff for: src/main/java/org/springframework/data/elasticsearch/core/ReactiveElasticsearchTemplate.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2021 the original author or authors.
+ * Copyright 2018-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,7 +23,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import org.elasticsearch.Version;
@@ -771,15 +770,18 @@ private Flux<SearchDocument> doFind(Query query, Class<?> clazz, IndexCoordinate
 		});
 	}
 
-	private Mono<SearchDocumentResponse> doFindForResponse(Query query, Class<?> clazz, IndexCoordinates index) {
+	private <T> Mono<SearchDocumentResponse> doFindForResponse(Query query, Class<?> clazz, IndexCoordinates index) {
 
 		return Mono.defer(() -> {
 			SearchRequest request = requestFactory.searchRequest(query, clazz, index);
 			request = prepareSearchRequest(request, false);
 
 			SearchDocumentCallback<?> documentCallback = new ReadSearchDocumentCallback<>(clazz, index);
+			// noinspection unchecked
+			SearchDocumentResponse.EntityCreator<T> entityCreator = searchDocument -> ((Mono<T>) documentCallback
+					.toEntity(searchDocument)).toFuture();
 
-			return doFindForResponse(request, searchDocument -> documentCallback.toEntity(searchDocument).block());
+			return doFindForResponse(request, entityCreator);
 		});
 	}
 
@@ -896,19 +898,18 @@ protected Flux<SearchDocument> doFind(SearchRequest request) {
 	 * Customization hook on the actual execution result {@link Mono}. <br />
 	 *
 	 * @param request the already prepared {@link SearchRequest} ready to be executed.
-	 * @param suggestEntityCreator
+	 * @param entityCreator
 	 * @return a {@link Mono} emitting the result of the operation converted to s {@link SearchDocumentResponse}.
 	 */
-	protected Mono<SearchDocumentResponse> doFindForResponse(SearchRequest request,
-			Function<SearchDocument, ? extends Object> suggestEntityCreator) {
+	protected <T> Mono<SearchDocumentResponse> doFindForResponse(SearchRequest request,
+			SearchDocumentResponse.EntityCreator<T> entityCreator) {
 
 		if (QUERY_LOGGER.isDebugEnabled()) {
 			QUERY_LOGGER.debug("Executing doFindForResponse: {}", request);
 		}
 
-		return Mono.from(execute(client1 -> client1.searchForResponse(request))).map(searchResponse -> {
-			return SearchDocumentResponse.from(searchResponse, suggestEntityCreator);
-		});
+		return Mono.from(execute(client -> client.searchForResponse(request)))
+				.map(searchResponse -> SearchDocumentResponse.from(searchResponse, entityCreator));
 	}
 
 	/**

Diff for: src/main/java/org/springframework/data/elasticsearch/core/document/SearchDocumentResponse.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,8 +17,11 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.text.Text;
@@ -38,13 +41,15 @@
 
 /**
  * This represents the complete search response from Elasticsearch, including the returned documents. Instances must be
- * created with the {@link #from(SearchResponse,Function)} method.
+ * created with the {@link #from(SearchResponse, EntityCreator)} method.
  *
  * @author Peter-Josef Meisch
  * @since 4.0
  */
 public class SearchDocumentResponse {
 
+	private static final Log LOGGER = LogFactory.getLog(SearchDocumentResponse.class);
+
 	private final long totalHits;
 	private final String totalHitsRelation;
 	private final float maxScore;
@@ -98,12 +103,11 @@ public Suggest getSuggest() {
 	 * creates a SearchDocumentResponse from the {@link SearchResponse}
 	 *
 	 * @param searchResponse must not be {@literal null}
-	 * @param suggestEntityCreator function to create an entity from a {@link SearchDocument}
+	 * @param entityCreator function to create an entity from a {@link SearchDocument}
 	 * @param <T> entity type
 	 * @return the SearchDocumentResponse
 	 */
-	public static <T> SearchDocumentResponse from(SearchResponse searchResponse,
-			Function<SearchDocument, T> suggestEntityCreator) {
+	public static <T> SearchDocumentResponse from(SearchResponse searchResponse, EntityCreator<T> entityCreator) {
 
 		Assert.notNull(searchResponse, "searchResponse must not be null");
 
@@ -112,7 +116,7 @@ public static <T> SearchDocumentResponse from(SearchResponse searchResponse,
 		Aggregations aggregations = searchResponse.getAggregations();
 		org.elasticsearch.search.suggest.Suggest suggest = searchResponse.getSuggest();
 
-		return from(searchHits, scrollId, aggregations, suggest, suggestEntityCreator);
+		return from(searchHits, scrollId, aggregations, suggest, entityCreator);
 	}
 
 	/**
@@ -122,14 +126,14 @@ public static <T> SearchDocumentResponse from(SearchResponse searchResponse,
 	 * @param scrollId scrollId
 	 * @param aggregations aggregations
 	 * @param suggestES the suggestion response from Elasticsearch
-	 * @param suggestEntityCreator function to create an entity from a {@link SearchDocument}
+	 * @param entityCreator function to create an entity from a {@link SearchDocument}
 	 * @param <T> entity type
 	 * @return the {@link SearchDocumentResponse}
 	 * @since 4.3
 	 */
 	public static <T> SearchDocumentResponse from(SearchHits searchHits, @Nullable String scrollId,
 			@Nullable Aggregations aggregations, @Nullable org.elasticsearch.search.suggest.Suggest suggestES,
-			Function<SearchDocument, T> suggestEntityCreator) {
+			EntityCreator<T> entityCreator) {
 
 		TotalHits responseTotalHits = searchHits.getTotalHits();
 
@@ -153,14 +157,14 @@ public static <T> SearchDocumentResponse from(SearchHits searchHits, @Nullable S
 			}
 		}
 
-		Suggest suggest = suggestFrom(suggestES, suggestEntityCreator);
+		Suggest suggest = suggestFrom(suggestES, entityCreator);
 		return new SearchDocumentResponse(totalHits, totalHitsRelation, maxScore, scrollId, searchDocuments, aggregations,
 				suggest);
 	}
 
 	@Nullable
 	private static <T> Suggest suggestFrom(@Nullable org.elasticsearch.search.suggest.Suggest suggestES,
-			Function<SearchDocument, T> entityCreator) {
+			EntityCreator<T> entityCreator) {
 
 		if (suggestES == null) {
 			return null;
@@ -219,7 +223,19 @@ private static <T> Suggest suggestFrom(@Nullable org.elasticsearch.search.sugges
 					List<CompletionSuggestion.Entry.Option<T>> options = new ArrayList<>();
 					for (org.elasticsearch.search.suggest.completion.CompletionSuggestion.Entry.Option optionES : entryES) {
 						SearchDocument searchDocument = optionES.getHit() != null ? DocumentAdapters.from(optionES.getHit()) : null;
-						T hitEntity = searchDocument != null ? entityCreator.apply(searchDocument) : null;
+
+						T hitEntity = null;
+
+						if (searchDocument != null) {
+							try {
+								hitEntity = entityCreator.apply(searchDocument).get();
+							} catch (Exception e) {
+								if (LOGGER.isWarnEnabled()) {
+									LOGGER.warn("Error creating entity from SearchDocument");
+								}
+							}
+						}
+
 						options.add(new CompletionSuggestion.Entry.Option<T>(textToString(optionES.getText()),
 								textToString(optionES.getHighlighted()), optionES.getScore(), optionES.collateMatch(),
 								optionES.getContexts(), scoreDocFrom(optionES.getDoc()), searchDocument, hitEntity));
@@ -254,4 +270,14 @@ private static ScoreDoc scoreDocFrom(@Nullable org.apache.lucene.search.ScoreDoc
 	private static String textToString(@Nullable Text text) {
 		return text != null ? text.string() : "";
 	}
+
+	/**
+	 * A function to convert a {@link SearchDocument} async into an entity. Asynchronous so that it can be used from the
+	 * imperative and the reactive code.
+	 *
+	 * @param <T> the entity type
+	 */
+	@FunctionalInterface
+	public interface EntityCreator<T> extends Function<SearchDocument, CompletableFuture<T>> {}
+
 }