segment replication changes (#529)

Signed-off-by: Subhobrata Dey <[email protected]>
(cherry picked from commit af05189)

src/main/java/org/opensearch/securityanalytics/SecurityAnalyticsPlugin.java

@@ -12,6 +12,7 @@
 import java.util.function.Supplier;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.core.action.ActionResponse;
@@ -293,6 +294,7 @@ public void onFailure(Exception e) {
         // Trigger initialization of prepackaged rules by calling SearchRule API
         SearchRequest searchRequest = new SearchRequest(Rule.PRE_PACKAGED_RULES_INDEX);
         searchRequest.source(new SearchSourceBuilder().query(QueryBuilders.matchAllQuery()).size(0));
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
         client.execute(
                 SearchRuleAction.INSTANCE,
                 new SearchRuleRequest(true, searchRequest),

src/main/java/org/opensearch/securityanalytics/correlation/JoinEngine.java

@@ -9,6 +9,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.join.ScoreMode;
 import org.opensearch.OpenSearchStatusException;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.search.MultiSearchRequest;
 import org.opensearch.action.search.MultiSearchResponse;
@@ -85,6 +86,7 @@ public void onSearchDetectorResponse(Detector detector, Finding finding) {
         SearchRequest searchRequest = new SearchRequest();
         searchRequest.indices(CorrelationRule.CORRELATION_RULE_INDEX);
         searchRequest.source(searchSourceBuilder);
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
         client.search(searchRequest, new ActionListener<>() {
             @Override
@@ -144,6 +146,7 @@ private void getValidDocuments(String detectorType, List<String> indices, List<C
                 SearchRequest searchRequest = new SearchRequest();
                 searchRequest.indices(indices.toArray(new String[]{}));
                 searchRequest.source(searchSourceBuilder);
+                searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
                 validCorrelationRules.add(rule);
                 mSearchRequest.add(searchRequest);
@@ -221,6 +224,7 @@ private void searchFindingsByTimestamp(String detectorType, Map<String, List<Cor
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(DetectorMonitorConfig.getAllFindingsIndicesPattern(categoryToQueries.getKey()));
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
             mSearchRequest.add(searchRequest);
             categoryToQueriesPairs.add(new Pair<>(categoryToQueries.getKey(), categoryToQueries.getValue()));
         }
@@ -291,6 +295,7 @@ private void searchDocsWithFilterKeys(String detectorType, Map<String, DocSearch
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(docSearchCriteria.getValue().indices.toArray(new String[]{}));
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
             categories.add(docSearchCriteria.getKey());
             mSearchRequest.add(searchRequest);
@@ -359,6 +364,7 @@ private void getCorrelatedFindings(String detectorType, Map<String, List<String>
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(DetectorMonitorConfig.getAllFindingsIndicesPattern(relatedDocIds.getKey()));
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
             categories.add(relatedDocIds.getKey());
             mSearchRequest.add(searchRequest);

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.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.bulk.BulkRequest;
 import org.opensearch.action.bulk.BulkResponse;
@@ -74,6 +75,7 @@ public void insertCorrelatedFindings(String detectorType, Finding finding, Strin
         SearchRequest searchRequest = new SearchRequest();
         searchRequest.indices(CorrelationIndices.CORRELATION_INDEX);
         searchRequest.source(searchSourceBuilder);
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
         client.search(searchRequest, new ActionListener<>() {
             @Override
@@ -103,6 +105,7 @@ public void onResponse(SearchResponse response) {
                     SearchRequest searchRequest = new SearchRequest();
                     searchRequest.indices(CorrelationIndices.CORRELATION_INDEX);
                     searchRequest.source(searchSourceBuilder);
+                    searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
                     mSearchRequest.add(searchRequest);
                 }
@@ -240,6 +243,7 @@ public void insertOrphanFindings(String detectorType, Finding finding, float tim
         SearchRequest searchRequest = new SearchRequest();
         searchRequest.indices(CorrelationIndices.CORRELATION_INDEX);
         searchRequest.source(searchSourceBuilder);
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
         client.search(searchRequest, new ActionListener<>() {
             @Override
@@ -415,6 +419,7 @@ public void onFailure(Exception e) {
                             SearchRequest searchRequest = new SearchRequest();
                             searchRequest.indices(CorrelationIndices.CORRELATION_INDEX);
                             searchRequest.source(searchSourceBuilder);
+                            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
                             client.search(searchRequest, new ActionListener<>() {
                                 @Override

src/main/java/org/opensearch/securityanalytics/logtype/LogTypeService.java

@@ -25,6 +25,7 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.ResourceAlreadyExistsException;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.DocWriteRequest;
 import org.opensearch.action.admin.indices.create.CreateIndexRequest;
@@ -116,7 +117,7 @@ public void getAllLogTypes(ActionListener<List<String>> listener) {
                     .field(LOG_TYPES)
                     .size(MAX_LOG_TYPE_COUNT)
             ));
-            searchRequest.preference("_primary");
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
             client.search(
                 searchRequest,
                 ActionListener.delegateFailure(
@@ -374,6 +375,7 @@ public void getAllFieldMappings(ActionListener<List<FieldMappingDoc>> listener)
         SearchRequest searchRequest = new SearchRequest(LOG_TYPE_INDEX);
         searchRequest.source(new SearchSourceBuilder().query(QueryBuilders.boolQuery()
                 .mustNot(QueryBuilders.existsQuery("source"))).size(10000));
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
         client.search(
             searchRequest,
             ActionListener.delegateFailure(
@@ -407,6 +409,7 @@ public void getFieldMappingsByLogTypes(List<String> logTypes, ActionListener<Lis
                 QueryBuilders.termsQuery(LOG_TYPES, logTypes.toArray(new String[0])))
                 .size(10000)
         );
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
         client.search(
                 searchRequest,
                 ActionListener.delegateFailure(

src/main/java/org/opensearch/securityanalytics/resthandler/RestSearchCorrelationRuleAction.java

@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.rest.BaseRestHandler;
@@ -53,7 +54,8 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                 .version(true);
         SearchRequest searchRequest = new SearchRequest()
                 .source(searchSourceBuilder)
-                .indices(CorrelationRule.CORRELATION_RULE_INDEX);
+                .indices(CorrelationRule.CORRELATION_RULE_INDEX)
+                .preference(Preference.PRIMARY_FIRST.type());
 
         SearchCorrelationRuleRequest searchCorrelationRuleRequest = new SearchCorrelationRuleRequest(searchRequest);
         return channel -> client.execute(SearchCorrelationRuleAction.INSTANCE, searchCorrelationRuleRequest, new RestToXContentListener<>(channel));

src/main/java/org/opensearch/securityanalytics/resthandler/RestSearchDetectorAction.java

@@ -10,6 +10,7 @@
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.XContentFactory;
@@ -69,6 +70,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         SearchRequest searchRequest = new SearchRequest();
         searchRequest.source(searchSourceBuilder);
         searchRequest.indices(Detector.DETECTORS_INDEX);
+        searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
         SearchDetectorRequest searchDetectorRequest = new SearchDetectorRequest(searchRequest);
 

src/main/java/org/opensearch/securityanalytics/resthandler/RestSearchRuleAction.java

@@ -9,6 +9,7 @@
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.XContentFactory;
@@ -68,7 +69,8 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                 .version(true);
         SearchRequest searchRequest = new SearchRequest()
                 .source(searchSourceBuilder)
-                .indices(isPrepackaged ? Rule.PRE_PACKAGED_RULES_INDEX: Rule.CUSTOM_RULES_INDEX);
+                .indices(isPrepackaged ? Rule.PRE_PACKAGED_RULES_INDEX: Rule.CUSTOM_RULES_INDEX)
+                .preference(Preference.PRIMARY_FIRST.type());
 
         SearchRuleRequest searchRuleRequest = new SearchRuleRequest(isPrepackaged, searchRequest);
         return channel -> client.execute(SearchRuleAction.INSTANCE, searchRuleRequest, searchRuleResponse(channel));

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.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRunnable;
@@ -212,6 +213,7 @@ void start() {
                 SearchRequest searchRequest = new SearchRequest();
                 searchRequest.indices(Detector.DETECTORS_INDEX);
                 searchRequest.source(searchSourceBuilder);
+                searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
                 client.search(searchRequest, new ActionListener<>() {
                     @Override
@@ -291,6 +293,7 @@ public void getTimestampFeature(String detectorType, Map<String, List<String>> c
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(CorrelationIndices.CORRELATION_INDEX);
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
             client.search(searchRequest, new ActionListener<>() {
                 @Override

src/main/java/org/opensearch/securityanalytics/transport/TransportDeleteRuleAction.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.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.ActionRunnable;
 import org.opensearch.action.get.GetRequest;
@@ -144,7 +145,8 @@ private void onGetResponse(Rule rule) {
                                 .seqNoAndPrimaryTerm(true)
                                 .version(true)
                                 .query(queryBuilder)
-                                .size(10000));
+                                .size(10000))
+                        .preference(Preference.PRIMARY_FIRST.type());
 
                 client.search(searchRequest, new ActionListener<>() {
                     @Override

src/main/java/org/opensearch/securityanalytics/transport/TransportGetAlertsAction.java

@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.join.ScoreMode;
 import org.opensearch.OpenSearchStatusException;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
@@ -111,6 +112,7 @@ protected void doExecute(Task task, GetAlertsRequest request, ActionListener<Get
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(Detector.DETECTORS_INDEX);
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
             transportSearchDetectorAction.execute(new SearchDetectorRequest(searchRequest), new ActionListener<>() {
                 @Override

src/main/java/org/opensearch/securityanalytics/transport/TransportGetFindingsAction.java

@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.join.ScoreMode;
 import org.opensearch.OpenSearchStatusException;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
@@ -128,6 +129,7 @@ protected void doExecute(Task task, GetFindingsRequest request, ActionListener<G
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(Detector.DETECTORS_INDEX);
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
             transportSearchDetectorAction.execute(new SearchDetectorRequest(searchRequest), new ActionListener<>() {
                 @Override

src/main/java/org/opensearch/securityanalytics/transport/TransportIndexDetectorAction.java

@@ -10,6 +10,7 @@
 import org.apache.lucene.search.join.ScoreMode;
 import org.opensearch.common.SetOnce;
 import org.opensearch.OpenSearchStatusException;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.ActionRunnable;
 import org.opensearch.action.StepListener;
@@ -204,7 +205,9 @@ private void checkIndicesAndExecute(
         User user
     ) {
         String [] detectorIndices = request.getDetector().getInputs().stream().flatMap(detectorInput -> detectorInput.getIndices().stream()).toArray(String[]::new);
-        SearchRequest searchRequest =  new SearchRequest(detectorIndices).source(SearchSourceBuilder.searchSource().size(1).query(QueryBuilders.matchAllQuery()));;
+        SearchRequest searchRequest =  new SearchRequest(detectorIndices)
+                .source(SearchSourceBuilder.searchSource().size(1).query(QueryBuilders.matchAllQuery()))
+                .preference(Preference.PRIMARY_FIRST.type());
         client.search(searchRequest, new ActionListener<>() {
             @Override
             public void onResponse(SearchResponse searchResponse) {
@@ -1132,7 +1135,8 @@ public void importRules(IndexDetectorRequest request, ActionListener<List<IndexM
                             .seqNoAndPrimaryTerm(true)
                             .version(true)
                             .query(queryBuilder)
-                            .size(10000));
+                            .size(10000))
+                    .preference(Preference.PRIMARY_FIRST.type());
 
             client.search(searchRequest, new ActionListener<>() {
                 @Override
@@ -1191,7 +1195,8 @@ public void importCustomRules(Detector detector, DetectorInput detectorInput, Li
                             .seqNoAndPrimaryTerm(true)
                             .version(true)
                             .query(queryBuilder)
-                            .size(10000));
+                            .size(10000))
+                    .preference(Preference.PRIMARY_FIRST.type());
 
             client.search(searchRequest, new ActionListener<>() {
                 @Override

src/main/java/org/opensearch/securityanalytics/transport/TransportIndexRuleAction.java

@@ -19,6 +19,7 @@
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.client.Client;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.common.settings.Settings;
@@ -316,7 +317,8 @@ private void searchDetectors(String ruleId, ActionListener<SearchResponse> liste
                             .seqNoAndPrimaryTerm(true)
                             .version(true)
                             .query(queryBuilder)
-                            .size(10000));
+                            .size(10000))
+                    .preference(Preference.PRIMARY_FIRST.type());
 
             client.search(searchRequest, listener);
         }

src/main/java/org/opensearch/securityanalytics/transport/TransportListCorrelationAction.java

@@ -7,6 +7,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchStatusException;
+import org.opensearch.cluster.routing.Preference;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.ActionRunnable;
 import org.opensearch.action.search.SearchRequest;
@@ -112,6 +113,7 @@ void start() {
             SearchRequest searchRequest = new SearchRequest();
             searchRequest.indices(CorrelationIndices.CORRELATION_INDEX);
             searchRequest.source(searchSourceBuilder);
+            searchRequest.preference(Preference.PRIMARY_FIRST.type());
 
             client.search(searchRequest, new ActionListener<>() {
                 @Override