Address PR feedback

Signed-off-by: Ryan Bogan <[email protected]>
opensearch-project · Dec 1, 2023 · ac5df23 · ac5df23
1 parent 613a28e
commit ac5df23
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/knn/training/TrainingJobClusterStateListener.java b/src/main/java/org/opensearch/knn/training/TrainingJobClusterStateListener.java
@@ -185,10 +185,13 @@ private List<String> searchModelIds() throws IOException, InterruptedException {
         modelDao.search(new SearchRequest(), new ActionListener<SearchResponse>() {
             @Override
             public void onResponse(SearchResponse searchResponse) {
-                for (SearchHit searchHit : searchResponse.getHits().getHits()) {
-                    modelIds.add(searchHit.getId());
+                try {
+                    for (SearchHit searchHit : searchResponse.getHits().getHits()) {
+                        modelIds.add(searchHit.getId());
+                    }
+                } finally {
+                    latch.countDown();
                 }
-                latch.countDown();
             }
 
             @Override