Merge pull request #5672 from BryanFauble/plfm-9726-search-convergence-timeout

PLFM-9726: Fix SearchIndex stuck in CREATING when AOSS convergence probe times out

Author: BryanFauble

integration-test/src/test/java/org/sagebionetworks/ITSearchQueryTest.java

@@ -61,6 +61,7 @@
 import org.sagebionetworks.repo.model.search.dsl.TermsAggregation;
 import org.sagebionetworks.repo.model.search.table.SearchAutocompleteRequest;
 import org.sagebionetworks.repo.model.search.table.SearchIndexQuery;
+import org.sagebionetworks.util.TimeUtils;
 
 /**
  * Integration tests for the SearchIndex query path.
@@ -420,8 +421,13 @@ public void testAutocompleteWithEdgeNgram() throws Exception {
 						.setQuery(new Query().setMatch_bool_prefix(
 								Map.of("geneName", new MatchBoolPrefixFieldOptions().setQuery("BRC")))));
 
-		// call under test
-		SearchQueryResults autocompleteResults = synapse.searchAutocomplete(autocompleteRequest);
+		// call under test — AOSS is eventually consistent per query type and per replica, so
+		// the match_all wait above does not guarantee this match_bool_prefix autocomplete sees
+		// every document yet. Poll the autocomplete itself until both BRCA hits surface.
+		SearchQueryResults autocompleteResults = waitForQuery(
+			() -> synapse.searchAutocomplete(autocompleteRequest),
+			r -> r.getHits() != null && r.getHits().size() >= 2,
+			"at least 2 autocomplete hits for 'BRC' (BRCA1, BRCA2)");
 		assertNotNull(autocompleteResults);
 		assertNotNull(autocompleteResults.getHits());
 		assertTrue(autocompleteResults.getHits().size() >= 2,
@@ -453,6 +459,44 @@ public void testStartSearchQueryWithUnsupportedKey() {
 		assertEquals("JSON Element in Entity is Unsupported: notPartOfSpecification", message);
 	}
 
+	/**
+	 * A search/autocomplete call against the live cluster that may throw {@link SynapseException}.
+	 */
+	@FunctionalInterface
+	private interface QueryCall {
+		SearchQueryResults call() throws SynapseException;
+	}
+
+	/**
+	 * Poll a synchronous search/autocomplete call until {@code condition} holds, then return the
+	 * matching result. AOSS is eventually consistent per query type and per replica, so a freshly
+	 * built index that already answers a {@code match_all} probe may transiently under-return for a
+	 * different query shape (autocomplete, match, post_filter). Tests asserting on a specific
+	 * synchronous query must poll that exact query rather than trusting a prior wait. Mirrors the
+	 * {@code waitForSearch}/{@code waitForSearchHits} helpers in {@code OpenSearchManagerImplAutoWiredTest}.
+	 *
+	 * @param queryCall   the call under test
+	 * @param condition   the readiness predicate on the result (e.g. expected hit count reached)
+	 * @param description what is being waited for, used in the timeout message
+	 * @return the result that satisfied {@code condition}
+	 */
+	private SearchQueryResults waitForQuery(QueryCall queryCall,
+			java.util.function.Predicate<SearchQueryResults> condition, String description) {
+		SearchQueryResults[] holder = { null };
+		boolean ready = TimeUtils.waitForExponential(MAX_QUERY_TIMEOUT_MS, 1000L, null, (v) -> {
+			try {
+				SearchQueryResults r = queryCall.call();
+				holder[0] = r;
+				return r != null && condition.test(r);
+			} catch (SynapseException e) {
+				// index_not_found / still-propagating — keep polling
+				return false;
+			}
+		});
+		assertTrue(ready, "Timed out waiting for " + description);
+		return holder[0];
+	}
+
 	private void grantPublicRead(String entityId) throws SynapseException {
 		AccessControlList acl = synapse.getACL(entityId);
 		ResourceAccess publicAccess = new ResourceAccess();

services/repository-managers/src/main/java/org/sagebionetworks/repo/manager/search/OpenSearchManager.java

@@ -93,28 +93,6 @@ Optional<String> createIndex(String indexName, List<ColumnModel> columns,
 	 */
 	void waitForIndexWritable(String indexName) throws RecoverableMessageException;
 
-	/**
-	 * Block until {@code indexName} reports a document count of at least
-	 * {@code expectedCount}. AOSS is eventually consistent: a {@code bulkIndex} call can
-	 * return success while the documents are not yet visible to {@code _search} or
-	 * {@code _count}. Polling {@code _count} (allowed on AOSS under
-	 * {@code aoss:ReadDocument}) is the only convergence signal available.
-	 *
-	 * <p>The comparison is {@code actual >= expectedCount} rather than strict equality so
-	 * a leftover readiness-probe sentinel (whose cleanup is best-effort, see
-	 * {@link #waitForIndexWritable}) does not strand convergence one short.</p>
-	 *
-	 * <p>No-op when {@code expectedCount == 0} — an empty index reports zero immediately
-	 * and skipping the call avoids an unnecessary round-trip.</p>
-	 *
-	 * @param indexName     the OpenSearch index name
-	 * @param expectedCount the document count the index must reach before returning
-	 * @throws RecoverableMessageException when the index does not converge within the
-	 *         retry budget, so the SearchIndex lifecycle message goes back on SQS for a
-	 *         later attempt instead of being marked permanently FAILED
-	 */
-	void waitForDocumentCount(String indexName, long expectedCount) throws RecoverableMessageException;
-
 	/**
 	 * Execute a search query against the OpenSearch index. The {@code options} set controls
 	 * which sections of the OpenSearch request are populated: omitting HITS switches the

services/repository-managers/src/main/java/org/sagebionetworks/repo/manager/search/OpenSearchManagerImpl.java

@@ -30,8 +30,6 @@
 import org.opensearch.client.opensearch._types.mapping.Property;
 import org.opensearch.client.opensearch.core.BulkRequest;
 import org.opensearch.client.opensearch.core.BulkResponse;
-import org.opensearch.client.opensearch.core.CountRequest;
-import org.opensearch.client.opensearch.core.CountResponse;
 import org.opensearch.client.opensearch.core.DeleteRequest;
 import org.opensearch.client.opensearch.core.IndexRequest;
 import org.opensearch.client.opensearch.core.SearchResponse;
@@ -105,13 +103,6 @@ public class OpenSearchManagerImpl implements OpenSearchManager {
 	static int VALIDATE_MAX_RETRIES = 10;
 	static long VALIDATE_INITIAL_BACKOFF_MS = 1000L;
 
-	// Convergence probe for a freshly-finished bulk index. AOSS acknowledges bulk writes
-	// before the documents are visible to _search or _count, so we poll _count until the
-	// index reports the expected number of documents. Same budget as the writability probe
-	// for consistency. Non-final so unit tests can lower the values and avoid real sleeps.
-	static int COUNT_PROBE_MAX_RETRIES = 10;
-	static long COUNT_PROBE_INITIAL_BACKOFF_MS = 10000L;
-
 	// Cleanup retry for the readiness-probe sentinel. AOSS doesn't honor refresh=wait_for,
 	// so a single delete that fails on a transient network blip would orphan the sentinel
 	// (visible only to MATCH_ALL queries since _row_id = -1 cannot collide with real ids,
@@ -543,55 +534,6 @@ public void waitForIndexWritable(String indexName) throws RecoverableMessageExce
 		}
 	}
 
-	@Override
-	public void waitForDocumentCount(String indexName, long expectedCount) throws RecoverableMessageException {
-		// Empty indexes report zero immediately; skip the round-trip.
-		if (expectedCount <= 0L) {
-			return;
-		}
-		final int[] attempt = {0};
-		final long[] lastObserved = {-1L};
-		try {
-			TimeUtils.waitForExponentialMaxRetry(COUNT_PROBE_MAX_RETRIES, COUNT_PROBE_INITIAL_BACKOFF_MS,
-					() -> {
-						attempt[0]++;
-						try {
-							CountResponse response = openSearchClient.count(CountRequest.of(r -> r
-									.index(indexName)));
-							long actual = response.count();
-							lastObserved[0] = actual;
-							// >= rather than == so a leftover readiness-probe sentinel cannot
-							// permanently strand convergence one short.
-							if (actual >= expectedCount) {
-								return Boolean.TRUE;
-							}
-							LOG.warn("Index {} not yet converged (attempt {}/{}): {} of {} documents visible",
-									indexName, attempt[0], COUNT_PROBE_MAX_RETRIES, actual, expectedCount);
-							throw new RetryException("count " + actual + " of " + expectedCount);
-						} catch (OpenSearchException e) {
-							LOG.warn("Index {} count probe failed (attempt {}/{}): {}",
-									indexName, attempt[0], COUNT_PROBE_MAX_RETRIES, describeError(e.error()));
-							throw new RetryException(e);
-						} catch (IOException e) {
-							LOG.warn("Index {} count probe failed (attempt {}/{}): {}",
-									indexName, attempt[0], COUNT_PROBE_MAX_RETRIES, e.getMessage());
-							throw new RetryException(e);
-						}
-					});
-		} catch (RetryException e) {
-			LOG.error("Index {} did not converge to expected count after {} attempts ({} of {})",
-					indexName, COUNT_PROBE_MAX_RETRIES, lastObserved[0], expectedCount);
-			throw new RecoverableMessageException(
-					"AOSS index " + indexName + " did not converge to expected count ("
-							+ lastObserved[0] + " of " + expectedCount + ") within the retry budget",
-					e.getCause());
-		} catch (RuntimeException e) {
-			throw e;
-		} catch (Exception e) {
-			throw new RuntimeException("Failed convergence probe for search index: " + indexName, e);
-		}
-	}
-
 	@Override
 	public long bulkIndex(String indexName, List<BulkOperation> operations) {
 		// Callers must hand us idempotent ops (index/delete with explicit _id). When a

services/repository-managers/src/main/java/org/sagebionetworks/repo/manager/search/SearchIndexLifecycleManagerImpl.java

@@ -313,13 +313,6 @@ private void buildIndex(ProgressCallback progressCallback, String entityId, Long
 			// failure is not in the search index's configuration.
 			throw e;
 		} catch (Throwable e) {
-			// SearchIndexRowHandler.close() tunnels RecoverableMessageException through
-			// IOException because RowHandler.close() can't declare anything else. Unwrap
-			// here so a convergence-probe timeout reaches the do-not-mark-FAILED branch
-			// above on retry instead of permanently failing the index.
-			if (e instanceof IOException && e.getCause() instanceof RecoverableMessageException) {
-				throw (RecoverableMessageException) e.getCause();
-			}
 			// Another worker is currently deleting this same AOSS index. Translate
 			// to a recoverable SQS retry: by the time the retry runs, the winning
 			// delete has finished and our deleteIndex on retry no-ops via
@@ -616,17 +609,6 @@ private void flush() {
 		@Override
 		public void close() throws IOException {
 			flush();
-			// AOSS bulk writes acknowledge before the documents are visible to _search or
-			// _count. Block until the index reports the row count we streamed, so the
-			// SearchIndex isn't flipped to ACTIVE while a query would still under-return.
-			try {
-				client.waitForDocumentCount(indexName, totalRows);
-			} catch (RecoverableMessageException e) {
-				// RowHandler.close() can only declare IOException; tunnel the recoverable
-				// signal through it. The lifecycle build path unwraps before its FAILED
-				// branch so the message goes back on SQS instead of marking the index FAILED.
-				throw new IOException(e);
-			}
 		}
 	}
 

services/repository-managers/src/test/java/org/sagebionetworks/repo/manager/search/OpenSearchManagerImplAutoWiredTest.java

@@ -1236,9 +1236,10 @@ public void testSearchWithPostFilter() {
 				.setSize(10L)
 				.setFrom(0L);
 
-		// call under test — post_filter narrows hits; aggregations stay at full population
-		SearchQueryResults results = openSearchManager.search(indexName, body, columns,
-				EnumSet.of(SearchQueryPart.HITS, SearchQueryPart.TOTAL_HITS));
+		// call under test — post_filter narrows hits; aggregations stay at full population.
+		// Poll until post_filter returns the expected 2 ACTIVE hits.
+		SearchQueryResults results = waitForSearchHits(body, columns,
+				EnumSet.of(SearchQueryPart.HITS, SearchQueryPart.TOTAL_HITS), 2);
 
 		assertEquals(2L, results.getTotalHits(),
 				"totalHits must reflect post_filter narrowing — only ACTIVE rows");
@@ -1285,9 +1286,10 @@ public void testSearchWithHighlight() {
 				.setSize(10L)
 				.setFrom(0L);
 
-		// call under test — highlight payload round-trips and SearchHit.highlights is populated
-		SearchQueryResults results = openSearchManager.search(indexName, body, columns,
-				EnumSet.of(SearchQueryPart.HITS, SearchQueryPart.TOTAL_HITS));
+		// call under test — highlight payload round-trips and SearchHit.highlights is populated.
+		// Poll until the match query returns all 3 hits.
+		SearchQueryResults results = waitForSearchHits(body, columns,
+				EnumSet.of(SearchQueryPart.HITS, SearchQueryPart.TOTAL_HITS), 3);
 
 		assertEquals(3L, results.getTotalHits());
 		assertNotNull(results.getHits());

services/repository-managers/src/test/java/org/sagebionetworks/repo/manager/search/OpenSearchManagerImplTest.java

@@ -136,7 +136,6 @@ private static IndexSettingsAnalysis toAnalysis(String settingsJson) {
 
 	private long originalBulkInitialBackoffMs;
 	private long originalProbeInitialBackoffMs;
-	private long originalCountProbeInitialBackoffMs;
 	private long originalSentinelCleanupInitialBackoffMs;
 
 	@BeforeEach
@@ -147,8 +146,6 @@ public void setUp() {
 		OpenSearchManagerImpl.BULK_INDEX_INITIAL_BACKOFF_MS = 1L;
 		originalProbeInitialBackoffMs = OpenSearchManagerImpl.INDEX_WRITABLE_INITIAL_BACKOFF_MS;
 		OpenSearchManagerImpl.INDEX_WRITABLE_INITIAL_BACKOFF_MS = 1L;
-		originalCountProbeInitialBackoffMs = OpenSearchManagerImpl.COUNT_PROBE_INITIAL_BACKOFF_MS;
-		OpenSearchManagerImpl.COUNT_PROBE_INITIAL_BACKOFF_MS = 1L;
 		originalSentinelCleanupInitialBackoffMs = OpenSearchManagerImpl.SENTINEL_CLEANUP_INITIAL_BACKOFF_MS;
 		OpenSearchManagerImpl.SENTINEL_CLEANUP_INITIAL_BACKOFF_MS = 1L;
 	}
@@ -157,7 +154,6 @@ public void setUp() {
 	public void tearDown() {
 		OpenSearchManagerImpl.BULK_INDEX_INITIAL_BACKOFF_MS = originalBulkInitialBackoffMs;
 		OpenSearchManagerImpl.INDEX_WRITABLE_INITIAL_BACKOFF_MS = originalProbeInitialBackoffMs;
-		OpenSearchManagerImpl.COUNT_PROBE_INITIAL_BACKOFF_MS = originalCountProbeInitialBackoffMs;
 		OpenSearchManagerImpl.SENTINEL_CLEANUP_INITIAL_BACKOFF_MS = originalSentinelCleanupInitialBackoffMs;
 	}
 
@@ -1718,117 +1714,6 @@ public void testWaitForIndexWritableSentinelCleanupRetriesAndSucceeds() throws E
 		verify(openSearchClient, times(2)).delete(argThat((DeleteRequest req) -> req != null));
 	}
 
-	// --- waitForDocumentCount ---
-
-	@Test
-	public void testWaitForDocumentCountWithImmediateMatchReturns() throws Exception {
-		// Happy path: AOSS reports the expected count on the first attempt; no retry, no log noise.
-		when(openSearchClient.count(any(org.opensearch.client.opensearch.core.CountRequest.class)))
-				.thenReturn(org.opensearch.client.opensearch.core.CountResponse.of(b -> b
-						.count(5L)
-						.shards(s -> s.total(1).successful(1).failed(0))));
-
-		// call under test
-		manager.waitForDocumentCount("search-index-syn1", 5L);
-
-		verify(openSearchClient, times(1))
-				.count(any(org.opensearch.client.opensearch.core.CountRequest.class));
-	}
-
-	@Test
-	public void testWaitForDocumentCountConvergesAfterUnderCount() throws Exception {
-		// AOSS is eventually consistent: the first probe sees fewer documents than were
-		// bulk-indexed; the second sees the full count and returns.
-		when(openSearchClient.count(any(org.opensearch.client.opensearch.core.CountRequest.class)))
-				.thenReturn(countResponse(2L))
-				.thenReturn(countResponse(5L));
-
-		// call under test
-		manager.waitForDocumentCount("search-index-syn1", 5L);
-
-		verify(openSearchClient, times(2))
-				.count(any(org.opensearch.client.opensearch.core.CountRequest.class));
-	}
-
-	@Test
-	public void testWaitForDocumentCountTreatsExcessAsConverged() throws Exception {
-		// >= rather than == so a leftover readiness-probe sentinel cannot strand convergence
-		// one short. An excess count is therefore a successful return, not a retry.
-		when(openSearchClient.count(any(org.opensearch.client.opensearch.core.CountRequest.class)))
-				.thenReturn(countResponse(6L));
-
-		// call under test
-		manager.waitForDocumentCount("search-index-syn1", 5L);
-
-		verify(openSearchClient, times(1))
-				.count(any(org.opensearch.client.opensearch.core.CountRequest.class));
-	}
-
-	@Test
-	public void testWaitForDocumentCountExhaustsRetriesAndThrowsRecoverableMessageException() throws Exception {
-		// Persistent under-count: every attempt sees fewer documents than expected. The probe
-		// must throw RecoverableMessageException so the lifecycle worker re-queues the message
-		// without writing FAILED — convergence is transient by definition.
-		when(openSearchClient.count(any(org.opensearch.client.opensearch.core.CountRequest.class)))
-				.thenReturn(countResponse(2L));
-
-		// call under test
-		RecoverableMessageException ex = assertThrows(RecoverableMessageException.class,
-				() -> manager.waitForDocumentCount("search-index-syn1", 5L));
-
-		assertTrue(ex.getMessage().contains("did not converge"), ex.getMessage());
-		assertTrue(ex.getMessage().contains("2 of 5"), ex.getMessage());
-		verify(openSearchClient, times(OpenSearchManagerImpl.COUNT_PROBE_MAX_RETRIES))
-				.count(any(org.opensearch.client.opensearch.core.CountRequest.class));
-	}
-
-	@Test
-	public void testWaitForDocumentCountWithIOExceptionExhaustsRetries() throws Exception {
-		when(openSearchClient.count(any(org.opensearch.client.opensearch.core.CountRequest.class)))
-				.thenThrow(new IOException("connection reset"));
-
-		// call under test
-		assertThrows(RecoverableMessageException.class,
-				() -> manager.waitForDocumentCount("search-index-syn1", 5L));
-
-		verify(openSearchClient, times(OpenSearchManagerImpl.COUNT_PROBE_MAX_RETRIES))
-				.count(any(org.opensearch.client.opensearch.core.CountRequest.class));
-	}
-
-	@Test
-	public void testWaitForDocumentCountWithOpenSearchExceptionExhaustsRetries() throws Exception {
-		// AOSS may briefly return index_not_found_exception (or any other transient error) from
-		// _count while shards are still propagating the freshly-written documents.
-		ErrorResponse notFound = ErrorResponse.of(er -> er
-				.error(ErrorCause.of(e -> e.type("index_not_found_exception").reason("no such index")))
-				.status(404));
-		when(openSearchClient.count(any(org.opensearch.client.opensearch.core.CountRequest.class)))
-				.thenThrow(new OpenSearchException(notFound));
-
-		// call under test
-		assertThrows(RecoverableMessageException.class,
-				() -> manager.waitForDocumentCount("search-index-syn1", 5L));
-
-		verify(openSearchClient, times(OpenSearchManagerImpl.COUNT_PROBE_MAX_RETRIES))
-				.count(any(org.opensearch.client.opensearch.core.CountRequest.class));
-	}
-
-	@Test
-	public void testWaitForDocumentCountWithZeroExpectedShortCircuits() throws Exception {
-		// An empty SearchIndex needs no probe — _count would just return 0 and we'd skip the
-		// retry path anyway. Avoid the round-trip entirely.
-		// call under test
-		manager.waitForDocumentCount("search-index-syn1", 0L);
-
-		verifyZeroInteractions(openSearchClient);
-	}
-
-	private static org.opensearch.client.opensearch.core.CountResponse countResponse(long count) {
-		return org.opensearch.client.opensearch.core.CountResponse.of(b -> b
-				.count(count)
-				.shards(s -> s.total(1).successful(1).failed(0)));
-	}
-
 	// --- per-document fallback on partial batch failure ---
 
 	@Test