Fixing hanging tasks for correlations (opensearch-project#874)

Signed-off-by: Megha Goyal <[email protected]>
goyamegh · Mar 5, 2024 · db025ce · db025ce
1 parent 172d58d
commit db025ce
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git 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<String, Object> 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/...main/java/org/opensearch/securityanalytics/transport/TransportCorrelateFindingAction.java b/...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;
@@ -517,6 +518,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<String, Object> hitSource = response.getHits().getHits()[0].getSourceAsMap();
                         long scoreTimestamp = (long) hitSource.get("scoreTimestamp");
@@ -655,6 +661,7 @@ public void onOperation() {
         }
 
         public void onFailures(Exception t) {
+            log.error("Exception occurred while processing correlations", t);
             if (counter.compareAndSet(false, true)) {
                 finishHim(t);
             }