Providing parameter '_total' when searching may lead to paging issues. (#5455)

* Initial failing test.

* Initial failing test.

* fix and changelog.

* applying spotless check.

* small code refactoring.

* test refactoring

---------

Co-authored-by: peartree <etienne.poirier@smilecdr.com>

hapi-fhir-docs/src/main/resources/ca/uhn/hapi/fhir/changelog/7_0_0/5454-search-with-total-interfers-with-bundle-pagination-.yaml

@@ -0,0 +1,5 @@
+---
+type: fix
+issue: 5454
+jira: SMILE-7295
+title: "Previously, searching with parameter '_total' could influence chunked query resultsets and subsequently, paged results. This is now fixed."

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/config/JpaConfig.java

@@ -572,12 +572,11 @@ public class JpaConfig {
 	@Scope("prototype")
 	public PersistedJpaSearchFirstPageBundleProvider newPersistedJpaSearchFirstPageBundleProvider(
 			RequestDetails theRequest,
-			Search theSearch,
 			SearchTask theSearchTask,
 			ISearchBuilder theSearchBuilder,
 			RequestPartitionId theRequestPartitionId) {
 		return new PersistedJpaSearchFirstPageBundleProvider(
-				theSearch, theSearchTask, theSearchBuilder, theRequest, theRequestPartitionId);
+				theSearchTask, theSearchBuilder, theRequest, theRequestPartitionId);
 	}
 
 	@Bean(name = RepositoryValidatingRuleBuilder.REPOSITORY_VALIDATING_RULE_BUILDER)

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/PersistedJpaBundleProvider.java

@@ -151,6 +151,11 @@ public class PersistedJpaBundleProvider implements IBundleProvider {
 		myRequestPartitionHelperSvc = theRequestPartitionHelperSvc;
 	}
 
+	@VisibleForTesting
+	public Search getSearchEntityForTesting() {
+		return getSearchEntity();
+	}
+
 	protected Search getSearchEntity() {
 		return mySearchEntity;
 	}

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/PersistedJpaBundleProviderFactory.java

@@ -55,14 +55,12 @@ public class PersistedJpaBundleProviderFactory {
 
 	public PersistedJpaSearchFirstPageBundleProvider newInstanceFirstPage(
 			RequestDetails theRequestDetails,
-			Search theSearch,
 			SearchTask theTask,
 			ISearchBuilder theSearchBuilder,
 			RequestPartitionId theRequestPartitionId) {
 		return (PersistedJpaSearchFirstPageBundleProvider) myApplicationContext.getBean(
 				JpaConfig.PERSISTED_JPA_SEARCH_FIRST_PAGE_BUNDLE_PROVIDER,
 				theRequestDetails,
-				theSearch,
 				theTask,
 				theSearchBuilder,
 				theRequestPartitionId);

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/PersistedJpaSearchFirstPageBundleProvider.java

@@ -21,7 +21,6 @@ package ca.uhn.fhir.jpa.search;
 
 import ca.uhn.fhir.interceptor.model.RequestPartitionId;
 import ca.uhn.fhir.jpa.dao.ISearchBuilder;
-import ca.uhn.fhir.jpa.entity.Search;
 import ca.uhn.fhir.jpa.entity.SearchTypeEnum;
 import ca.uhn.fhir.jpa.model.dao.JpaPid;
 import ca.uhn.fhir.jpa.model.search.SearchStatusEnum;
@@ -52,16 +51,14 @@ public class PersistedJpaSearchFirstPageBundleProvider extends PersistedJpaBundl
 	 */
 	@SuppressWarnings("rawtypes")
 	public PersistedJpaSearchFirstPageBundleProvider(
-			Search theSearch,
 			SearchTask theSearchTask,
 			ISearchBuilder theSearchBuilder,
 			RequestDetails theRequest,
 			RequestPartitionId theRequestPartitionId) {
-		super(theRequest, theSearch.getUuid());
+		super(theRequest, theSearchTask.getSearch());
 
-		assert theSearch.getSearchType() != SearchTypeEnum.HISTORY;
+		assert getSearchEntity().getSearchType() != SearchTypeEnum.HISTORY;
 
-		setSearchEntity(theSearch);
 		mySearchTask = theSearchTask;
 		mySearchBuilder = theSearchBuilder;
 		super.setRequestPartitionId(theRequestPartitionId);

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/SearchCoordinatorSvcImpl.java

@@ -110,7 +110,7 @@ public class SearchCoordinatorSvcImpl implements ISearchCoordinatorSvc<JpaPid> {
 	private final SearchStrategyFactory mySearchStrategyFactory;
 	private final ExceptionService myExceptionSvc;
 	private final BeanFactory myBeanFactory;
-	private final ConcurrentHashMap<String, SearchTask> myIdToSearchTask = new ConcurrentHashMap<>();
+	private ConcurrentHashMap<String, SearchTask> myIdToSearchTask = new ConcurrentHashMap<>();
 
 	private final Consumer<String> myOnRemoveSearchTask = myIdToSearchTask::remove;
 
@@ -162,6 +162,11 @@ public class SearchCoordinatorSvcImpl implements ISearchCoordinatorSvc<JpaPid> {
 		return myIdToSearchTask.keySet();
 	}
 
+	@VisibleForTesting
+	public void setIdToSearchTaskMapForUnitTests(ConcurrentHashMap<String, SearchTask> theIdToSearchTaskMap) {
+		myIdToSearchTask = theIdToSearchTaskMap;
+	}
+
 	@VisibleForTesting
 	public void setLoadingThrottleForUnitTests(Integer theLoadingThrottleForUnitTests) {
 		myLoadingThrottleForUnitTests = theLoadingThrottleForUnitTests;
@@ -571,7 +576,7 @@ public class SearchCoordinatorSvcImpl implements ISearchCoordinatorSvc<JpaPid> {
 		task.call();
 
 		PersistedJpaSearchFirstPageBundleProvider retVal = myPersistedJpaBundleProviderFactory.newInstanceFirstPage(
-				theRequestDetails, theSearch, task, theSb, theRequestPartitionId);
+				theRequestDetails, task, theSb, theRequestPartitionId);
 
 		ourLog.debug("Search initial phase completed in {}ms", w.getMillis());
 		return retVal;

hapi-fhir-jpaserver-test-dstu2/src/test/java/ca/uhn/fhir/jpa/search/SearchCoordinatorSvcImplTest.java

@@ -293,12 +293,11 @@ public class SearchCoordinatorSvcImplTest extends BaseSearchSvc {
 	}
 
 	private void initAsyncSearches() {
-		when(myPersistedJpaBundleProviderFactory.newInstanceFirstPage(nullable(RequestDetails.class), nullable(Search.class), nullable(SearchTask.class), nullable(ISearchBuilder.class), nullable(RequestPartitionId.class))).thenAnswer(t -> {
+		when(myPersistedJpaBundleProviderFactory.newInstanceFirstPage(nullable(RequestDetails.class), nullable(SearchTask.class), nullable(ISearchBuilder.class), nullable(RequestPartitionId.class))).thenAnswer(t -> {
 			RequestDetails requestDetails = t.getArgument(0, RequestDetails.class);
-			Search search = t.getArgument(1, Search.class);
+			SearchTask searchTask = t.getArgument(1, SearchTask.class);
-			SearchTask searchTask = t.getArgument(2, SearchTask.class);
+			ISearchBuilder<JpaPid> searchBuilder = t.getArgument(2, ISearchBuilder.class);
-			ISearchBuilder<JpaPid> searchBuilder = t.getArgument(3, ISearchBuilder.class);
+			PersistedJpaSearchFirstPageBundleProvider retVal = new PersistedJpaSearchFirstPageBundleProvider(searchTask, searchBuilder, requestDetails, null);
-			PersistedJpaSearchFirstPageBundleProvider retVal = new PersistedJpaSearchFirstPageBundleProvider(search, searchTask, searchBuilder, requestDetails, null);
 			retVal.setStorageSettingsForUnitTest(new JpaStorageSettings());
 			retVal.setTxServiceForUnitTest(myTransactionService);
 			retVal.setSearchCoordinatorSvcForUnitTest(mySvc);

hapi-fhir-jpaserver-test-r4/src/test/java/ca/uhn/fhir/jpa/dao/r4/FhirResourceDaoR4SearchOptimizedTest.java

@@ -8,7 +8,9 @@ import ca.uhn.fhir.jpa.entity.Search;
 import ca.uhn.fhir.jpa.model.entity.ResourceTable;
 import ca.uhn.fhir.jpa.model.search.SearchStatusEnum;
 import ca.uhn.fhir.jpa.search.PersistedJpaBundleProvider;
+import ca.uhn.fhir.jpa.search.PersistedJpaSearchFirstPageBundleProvider;
 import ca.uhn.fhir.jpa.search.SearchCoordinatorSvcImpl;
+import ca.uhn.fhir.jpa.search.builder.tasks.SearchTask;
 import ca.uhn.fhir.jpa.searchparam.MatchUrlService;
 import ca.uhn.fhir.jpa.searchparam.SearchParameterMap;
 import ca.uhn.fhir.jpa.test.BaseJpaR4Test;
@@ -29,7 +31,9 @@ import ca.uhn.fhir.rest.server.exceptions.InternalErrorException;
 import ca.uhn.fhir.test.utilities.ProxyUtil;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.checkerframework.checker.units.qual.A;
 import org.hl7.fhir.instance.model.api.IAnyResource;
+import org.hl7.fhir.instance.model.api.IBaseResource;
 import org.hl7.fhir.instance.model.api.IIdType;
 import org.hl7.fhir.r4.model.BodyStructure;
 import org.hl7.fhir.r4.model.CodeableConcept;
@@ -51,6 +55,8 @@ import org.hl7.fhir.r4.model.UriType;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.scheduling.concurrent.ThreadPoolExecutorFactoryBean;
 
@@ -58,6 +64,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.stream.Collectors;
@@ -102,6 +109,7 @@ public class FhirResourceDaoR4SearchOptimizedTest extends BaseJpaR4Test {
 	public final void after() {
 		mySearchCoordinatorSvcImpl.setLoadingThrottleForUnitTests(null);
 		mySearchCoordinatorSvcImpl.setSyncSizeForUnitTests(QueryParameterUtils.DEFAULT_SYNC_SIZE);
+		mySearchCoordinatorSvcImpl.setIdToSearchTaskMapForUnitTests(new ConcurrentHashMap<>());
 		myStorageSettings.setSearchPreFetchThresholds(new JpaStorageSettings().getSearchPreFetchThresholds());
 		myCaptureQueriesListener.setCaptureQueryStackTrace(false);
 		myStorageSettings.setIndexMissingFields(new JpaStorageSettings().getIndexMissingFields());
@@ -214,6 +222,39 @@ public class FhirResourceDaoR4SearchOptimizedTest extends BaseJpaR4Test {
 
 	}
 
+	@Test
+	public void testSearchCoordinatorSvc_whenExecutingSearchWithParamTotal_returnsBundleSynchronizedWithBackingSearchCapabilities(){
+		ArgumentCaptor<String> keyArgumentCaptor = ArgumentCaptor.forClass(String.class);
+		ArgumentCaptor<SearchTask> valueArgumentCaptor = ArgumentCaptor.forClass(SearchTask.class);
+		ConcurrentHashMap<String, SearchTask> spyingIdToSearchTaskMap = Mockito.spy(new ConcurrentHashMap<>());
+		mySearchCoordinatorSvcImpl.setIdToSearchTaskMapForUnitTests(spyingIdToSearchTaskMap);
+		create200Patients();
+
+		SearchParameterMap params = new SearchParameterMap();
+		params.add(Patient.SP_NAME, new StringParam("FAM"));
+		params.setSearchTotalMode(SearchTotalModeEnum.ACCURATE);
+
+		// calling dao.search will end up invoking the searchCoordinatorSvc.  based on the provided search parameters, the svc
+		// generates and triggers a searchTask which will create chunked resultsets. the searchTask make use of a searchEntity
+		// to keep track of search progress and key indicators like the search total count.
+		PersistedJpaSearchFirstPageBundleProvider results = (PersistedJpaSearchFirstPageBundleProvider) myPatientDao.search(params);
+
+		// to return the correct resources through method getResources(), the PersistedJpaSearchFirstPageBundleProvider generated by the
+		// searchCoordinatorSvc needs to access the same searchEntity that was used by the searchTask. this test ensures that the searchEntity
+		// operated upon by the searchTask is the same as the searchEntity that is found in the generated PersistedJpaSearchFirstPageBundleProvider.
+		Mockito.verify(spyingIdToSearchTaskMap, Mockito.times(1)).put(keyArgumentCaptor.capture(), valueArgumentCaptor.capture());
+
+		Search bundleProviderSearch = results.getSearchEntityForTesting();
+		Search backingSearch = valueArgumentCaptor.getValue().getSearch();
+
+		assertThat(bundleProviderSearch.getUuid(), equalTo(keyArgumentCaptor.getValue()));
+		assertThat(bundleProviderSearch.getUuid(), equalTo(backingSearch.getUuid()));
+
+		assertThat(bundleProviderSearch.getStatus(), equalTo(backingSearch.getStatus()));
+		assertThat(bundleProviderSearch.getId(), equalTo(backingSearch.getId()));
+
+	}
+
 	@Test
 	public void testFetchTotalAccurateForSlowLoading() {
 		create200Patients();