diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManager.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManager.java index f004b9fc03..69682b843b 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManager.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManager.java @@ -11,6 +11,7 @@ import org.opensearch.client.opensearch._types.query_dsl.TermsQueryField; import org.opensearch.client.opensearch.core.MsearchRequest; import org.opensearch.client.opensearch.core.MsearchResponse; +import org.opensearch.client.opensearch.core.msearch.MultiSearchResponseItem; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.plugins.sink.opensearch.BulkOperationWrapper; import org.opensearch.dataprepper.plugins.sink.opensearch.index.model.QueryManagerBulkOperation; @@ -19,6 +20,7 @@ import java.time.Duration; import java.time.Instant; +import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collection; import java.util.Iterator; @@ -127,12 +129,16 @@ public void run() { void runQueryLoop() { if (!bulkOperationsWaitingForQuery.isEmpty() && documentsCurrentlyBeingQueriedGauge.get() > 0) { - // Query for existing documents - final MsearchRequest msearchRequest = buildMultiSearchRequest(); + // Query for existing documents. Track the index key used for each search request, in request + // order, so responses can be correlated back to the correct pending-operations map. We cannot + // rely on hit.index() for this: a search against an alias (or datastream) returns the concrete + // backing index in the hit, which differs from the configured index/alias key. See #6902. + final List searchRequestIndexKeys = new ArrayList<>(); + final MsearchRequest msearchRequest = buildMultiSearchRequest(searchRequestIndexKeys); final MsearchResponse msearchResponse = queryForTermValues(msearchRequest); // Drop and Release Existing Documents - dropAndReleaseFoundEvents(msearchResponse); + dropAndReleaseFoundEvents(msearchResponse, searchRequestIndexKeys); // Move non-existing documents past query_duration to bulkOperationsReadyForIndex moveBulkRequestsThatHaveReachedQueryDuration(); @@ -183,30 +189,43 @@ public Set getAndClearBulkOperationsReadyToIndex() { } } - private MsearchRequest buildMultiSearchRequest() { + private MsearchRequest buildMultiSearchRequest(final List orderedIndexKeys) { + // Build the (index key, term-value chunk) pairs deterministically BEFORE constructing the request, + // so the order of orderedIndexKeys is guaranteed to align with the order in which searches are + // added to the msearch request (and therefore with the response order), independent of how or when + // the builder lambda below is invoked. This alignment is what lets dropAndReleaseFoundEvents map a + // response back to its configured index/alias key rather than the concrete backing index reported + // by hit.index(). See #6902. + final List>> searchRequestChunks = new ArrayList<>(); + for (final Map.Entry> entry : bulkOperationsWaitingForQuery.entrySet()) { + final String index = entry.getKey(); + final List values = getTermValues(entry.getValue().values()); + final int batchSize = 1000; + + LOG.info("Creating search requests for {} query term values in batches of {}", values.size(), batchSize); + for (int i = 0; i < values.size(); i += batchSize) { + final List chunk = values.subList(i, Math.min(i + batchSize, values.size())); + orderedIndexKeys.add(index); + searchRequestChunks.add(new AbstractMap.SimpleEntry<>(index, chunk)); + } + } + return MsearchRequest.of(m -> { - for (final Map.Entry> entry : bulkOperationsWaitingForQuery.entrySet()) { - final String index = entry.getKey(); - final List values = getTermValues(entry.getValue().values()); - final int batchSize = 1000; - - LOG.info("Creating search requests for {} query term values in batches of {}", values.size(), batchSize); - for (int i = 0; i < values.size(); i += batchSize) { - final List chunk = values.subList(i, Math.min(i + batchSize, values.size())); - - m.searches(s -> s - .header(h -> h.index(index)) - .body(b -> b - .size(chunk.size() * 2) - .source(source -> source.filter(f -> f.includes(queryTerm))) - .query(Query.of(q -> q - .terms(TermsQuery.of(t -> t - .field(queryTerm) - .terms(TermsQueryField.of(tf -> tf.value(chunk))) - )) - )) - )); - } + for (final Map.Entry> searchRequestChunk : searchRequestChunks) { + final String index = searchRequestChunk.getKey(); + final List chunk = searchRequestChunk.getValue(); + m.searches(s -> s + .header(h -> h.index(index)) + .body(b -> b + .size(chunk.size() * 2) + .source(source -> source.filter(f -> f.includes(queryTerm))) + .query(Query.of(q -> q + .terms(TermsQuery.of(t -> t + .field(queryTerm) + .terms(TermsQueryField.of(tf -> tf.value(chunk))) + )) + )) + )); } return m; }); @@ -257,36 +276,52 @@ private void moveBulkRequestsThatHaveReachedQueryDuration() { } } - private void dropAndReleaseFoundEvents(final MsearchResponse msearchResponse) { - msearchResponse.responses().forEach(response -> { + private void dropAndReleaseFoundEvents(final MsearchResponse msearchResponse, final List orderedIndexKeys) { + final List> responses = msearchResponse.responses(); + for (int responseIndex = 0; responseIndex < responses.size(); responseIndex++) { + final MultiSearchResponseItem response = responses.get(responseIndex); if (response.isFailure()) { LOG.error("Search response failed, potential for duplicate documents: {}", response.failure().error().toString()); - } else { - response.result().hits().hits().forEach(hit -> { - final String indexForHit = hit.index(); - final ObjectNode sourceForHit = hit.source(); - final String queryTermValue = sourceForHit.findValue(queryTerm).textValue(); + continue; + } - lockWaitingForQuery.lock(); - try { - final Map bulkOperationsForIndex = bulkOperationsWaitingForQuery.get(indexForHit); - final QueryManagerBulkOperation bulkOperationToRelease = bulkOperationsForIndex.get(queryTermValue); - if (bulkOperationToRelease == null) { - // Means two documents with the same query term value were found - LOG.warn("Bulk operation for term value {} with id {} is null, potentially a duplicate document", queryTermValue, hit.id()); - potentialDuplicatesDeleted.increment(); - } else { - LOG.debug("Found document with query term {}, dropping and releasing Event handle", queryTermValue); - bulkOperationToRelease.getBulkOperationWrapper().releaseEventHandle(true); - eventsDroppedAndReleasedCounter.increment(); - documentsCurrentlyBeingQueriedGauge.decrementAndGet(); - bulkOperationsForIndex.remove(queryTermValue); - } - } finally { - lockWaitingForQuery.unlock(); + // Correlate this response to the index key of the request that produced it (msearch preserves + // request order). hit.index() must not be used here: for aliases/datastreams the hit reports the + // concrete backing index, which differs from the configured index/alias key and would never match + // the pending-operations map, causing an NPE that wedges the query loop indefinitely. See #6902. + final String indexKey = responseIndex < orderedIndexKeys.size() ? orderedIndexKeys.get(responseIndex) : null; + final Map bulkOperationsForIndex = + indexKey == null ? null : bulkOperationsWaitingForQuery.get(indexKey); + + response.result().hits().hits().forEach(hit -> { + final ObjectNode sourceForHit = hit.source(); + final String queryTermValue = sourceForHit.findValue(queryTerm).textValue(); + + lockWaitingForQuery.lock(); + try { + if (bulkOperationsForIndex == null) { + // Defensive: no pending operations tracked for this index key. Skip rather than + // throwing so a single unexpected response can never NPE and block the query loop. + LOG.warn("No pending bulk operations found for index key {} (hit index {}); skipping release for term value {}", + indexKey, hit.index(), queryTermValue); + return; } - }); - } - }); + final QueryManagerBulkOperation bulkOperationToRelease = bulkOperationsForIndex.get(queryTermValue); + if (bulkOperationToRelease == null) { + // Means two documents with the same query term value were found + LOG.warn("Bulk operation for term value {} with id {} is null, potentially a duplicate document", queryTermValue, hit.id()); + potentialDuplicatesDeleted.increment(); + } else { + LOG.debug("Found document with query term {}, dropping and releasing Event handle", queryTermValue); + bulkOperationToRelease.getBulkOperationWrapper().releaseEventHandle(true); + eventsDroppedAndReleasedCounter.increment(); + documentsCurrentlyBeingQueriedGauge.decrementAndGet(); + bulkOperationsForIndex.remove(queryTermValue); + } + } finally { + lockWaitingForQuery.unlock(); + } + }); + } } } diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManagerTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManagerTest.java index bfcf58b71a..9319c91113 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManagerTest.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ExistingDocumentQueryManagerTest.java @@ -34,6 +34,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -107,7 +108,7 @@ void add_bulk_operation_and_found_in_query_drops_and_releases_event() throws IOE final MultiSearchItem multiSearchItem = mock(MultiSearchItem.class); final HitsMetadata hitsMetadata = mock(HitsMetadata.class); final Hit hit = mock(Hit.class); - when(hit.index()).thenReturn(index); + lenient().when(hit.index()).thenReturn(index); final ObjectNode objectNode = mock(ObjectNode.class); final JsonNode jsonNode = mock(JsonNode.class); @@ -141,6 +142,61 @@ void add_bulk_operation_and_found_in_query_drops_and_releases_event() throws IOE verifyNoMoreInteractions(indexConfiguration); } + @Test + void add_bulk_operation_and_found_in_query_when_hit_index_differs_from_configured_index_drops_and_releases_event() throws IOException { + // Regression test for #6902: when the sink writes to an alias (or datastream), the search hit reports + // the concrete backing index (e.g. "my-alias-000001"), which differs from the configured index/alias + // key ("my-alias"). Previously dropAndReleaseFoundEvents looked up the pending operations by + // hit.index(), which returned null and threw a NullPointerException, wedging the query loop and + // leaving the affected documents permanently stuck (never indexed, dropped, or sent to the DLQ). + final BulkOperationWrapper bulkOperationWrapper = mock(BulkOperationWrapper.class); + final String configuredIndex = UUID.randomUUID().toString(); + final String concreteBackingIndex = configuredIndex + "-000001"; + final String termValue = UUID.randomUUID().toString(); + when(bulkOperationWrapper.getTermValue()).thenReturn(termValue); + when(bulkOperationWrapper.getIndex()).thenReturn(configuredIndex); + + final MsearchResponse msearchResponse = mock(MsearchResponse.class); + final MultiSearchResponseItem responseItem = mock(MultiSearchResponseItem.class); + when(responseItem.isFailure()).thenReturn(false); + + final MultiSearchItem multiSearchItem = mock(MultiSearchItem.class); + final HitsMetadata hitsMetadata = mock(HitsMetadata.class); + final Hit hit = mock(Hit.class); + // The hit carries the concrete backing index, not the configured alias key. + lenient().when(hit.index()).thenReturn(concreteBackingIndex); + + final ObjectNode objectNode = mock(ObjectNode.class); + final JsonNode jsonNode = mock(JsonNode.class); + when(jsonNode.textValue()).thenReturn(termValue); + when(objectNode.findValue(queryTerm)).thenReturn(jsonNode); + when(hit.source()).thenReturn(objectNode); + + when(multiSearchItem.hits()).thenReturn(hitsMetadata); + when(hitsMetadata.hits()).thenReturn(List.of(hit)); + + when(responseItem.result()).thenReturn(multiSearchItem); + + when(msearchResponse.responses()).thenReturn(List.of(responseItem)); + + when(openSearchClient.msearch(any(MsearchRequest.class), eq(ObjectNode.class))) + .thenReturn(msearchResponse); + + final ExistingDocumentQueryManager objectUnderTest = createObjectUnderTest(); + + objectUnderTest.addBulkOperation(bulkOperationWrapper); + when(documentsCurrentlyQueried.get()).thenReturn(1); + + // Must not throw NullPointerException, and must drop + release the found duplicate. + objectUnderTest.runQueryLoop(); + + verify(eventsDroppedAndReleased).increment(); + verify(eventsAddedForQuerying).increment(); + verify(documentsCurrentlyQueried).incrementAndGet(); + verify(documentsCurrentlyQueried).decrementAndGet(); + verify(bulkOperationWrapper).releaseEventHandle(true); + } + @Test void add_bulk_operation_and_not_found_in_query_returns_as_ready_to_ingest() throws IOException, InterruptedException, NoSuchFieldException, IllegalAccessException { when(indexConfiguration.getQueryDuration()).thenReturn(Duration.ofMillis(1)); @@ -158,7 +214,7 @@ void add_bulk_operation_and_not_found_in_query_returns_as_ready_to_ingest() thro final MultiSearchItem multiSearchItem = mock(MultiSearchItem.class); final HitsMetadata hitsMetadata = mock(HitsMetadata.class); final Hit hit = mock(Hit.class); - when(hit.index()).thenReturn(index); + lenient().when(hit.index()).thenReturn(index); final ObjectNode objectNode = mock(ObjectNode.class); final JsonNode jsonNode = mock(JsonNode.class); @@ -261,7 +317,7 @@ void query_response_with_two_documents_with_same_term_value_tracks_duplicate_doc final MultiSearchItem multiSearchItem = mock(MultiSearchItem.class); final HitsMetadata hitsMetadata = mock(HitsMetadata.class); final Hit hit = mock(Hit.class); - when(hit.index()).thenReturn(index); + lenient().when(hit.index()).thenReturn(index); final ObjectNode objectNode = mock(ObjectNode.class); final JsonNode jsonNode = mock(JsonNode.class); @@ -270,7 +326,7 @@ void query_response_with_two_documents_with_same_term_value_tracks_duplicate_doc when(hit.source()).thenReturn(objectNode); final Hit duplicateHit = mock(Hit.class); - when(duplicateHit.index()).thenReturn(index); + lenient().when(duplicateHit.index()).thenReturn(index); when(duplicateHit.id()).thenReturn(UUID.randomUUID().toString()); when(jsonNode.textValue()).thenReturn(termValue);