From e09167a0580d6c5ec979fc4280a56990658bfae4 Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Mon, 11 Mar 2024 17:02:24 -0700 Subject: [PATCH] Fixing hanging tasks for correlations (#874) (#891) Signed-off-by: Megha Goyal (cherry picked from commit db025ce69c0201798b9e862a9156a658d2f2d241) Co-authored-by: Megha Goyal <56077967+goyamegh@users.noreply.github.com> (cherry picked from commit 73de979513b3c04ec7a6164131a2d5bdf6895450) --- .../correlation/VectorEmbeddingsEngine.java | 8 +++++++- .../transport/TransportCorrelateFindingAction.java | 7 +++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/opensearch/securityanalytics/correlation/VectorEmbeddingsEngine.java b/src/main/java/org/opensearch/securityanalytics/correlation/VectorEmbeddingsEngine.java index 9a423f6fb..931f3321c 100644 --- a/src/main/java/org/opensearch/securityanalytics/correlation/VectorEmbeddingsEngine.java +++ b/src/main/java/org/opensearch/securityanalytics/correlation/VectorEmbeddingsEngine.java @@ -7,6 +7,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.OpenSearchStatusException; +import org.opensearch.ResourceNotFoundException; import org.opensearch.cluster.routing.Preference; import org.opensearch.core.action.ActionListener; import org.opensearch.action.bulk.BulkRequest; @@ -84,6 +85,11 @@ public void onResponse(SearchResponse response) { correlateFindingAction.onFailures(new OpenSearchStatusException("Search request timed out", RestStatus.REQUEST_TIMEOUT)); } + if (response.getHits().getHits().length == 0) { + correlateFindingAction.onFailures( + new ResourceNotFoundException("Failed to find hits in metadata index for finding id {}", finding.getId())); + } + Map hitSource = response.getHits().getHits()[0].getSourceAsMap(); long counter = Long.parseLong(hitSource.get("counter").toString()); @@ -125,7 +131,7 @@ public void onResponse(MultiSearchResponse items) { continue; } - long totalHits = response.getResponse().getHits().getTotalHits().value; + long totalHits = response.getResponse().getHits().getHits().length; totalNeighbors += totalHits; for (int idx = 0; idx < totalHits; ++idx) { diff --git a/src/main/java/org/opensearch/securityanalytics/transport/TransportCorrelateFindingAction.java b/src/main/java/org/opensearch/securityanalytics/transport/TransportCorrelateFindingAction.java index b7a906159..8e417c649 100644 --- a/src/main/java/org/opensearch/securityanalytics/transport/TransportCorrelateFindingAction.java +++ b/src/main/java/org/opensearch/securityanalytics/transport/TransportCorrelateFindingAction.java @@ -8,6 +8,7 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.search.join.ScoreMode; import org.opensearch.OpenSearchStatusException; +import org.opensearch.ResourceNotFoundException; import org.opensearch.cluster.routing.Preference; import org.opensearch.core.action.ActionListener; import org.opensearch.action.ActionRequest; @@ -515,6 +516,11 @@ public void onFailure(Exception e) { client.search(searchRequest, new ActionListener<>() { @Override public void onResponse(SearchResponse response) { + if (response.getHits().getHits().length == 0) { + onFailures(new ResourceNotFoundException( + "Failed to find hits in metadata index for finding id {}", request.getFinding().getId())); + } + String id = response.getHits().getHits()[0].getId(); Map hitSource = response.getHits().getHits()[0].getSourceAsMap(); long scoreTimestamp = (long) hitSource.get("scoreTimestamp"); @@ -653,6 +659,7 @@ public void onOperation() { } public void onFailures(Exception t) { + log.error("Exception occurred while processing correlations", t); if (counter.compareAndSet(false, true)) { finishHim(t); }