Use instance of LockService instantiated in JobScheduler through Guice

build.gradle

@@ -65,6 +65,7 @@ plugins {
     id "com.netflix.nebula.ospackage" version "11.10.0"
     id 'java-library'
     id "de.undercouch.download" version "5.6.0"
+    id "org.gradle.test-retry" version "1.6.2"
 }
 
 apply plugin: 'opensearch.opensearchplugin'
@@ -249,6 +250,10 @@ task integTest(type: RestIntegTestTask) {
     description = "Run tests against a cluster"
     testClassesDirs = sourceSets.test.output.classesDirs
     classpath = sourceSets.test.runtimeClasspath
+    retry {
+        failOnPassedAfterRetry = false
+        maxRetries = 3
+    }
 }
 tasks.named("check").configure { dependsOn(integTest) }
 

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

@@ -14,7 +14,10 @@
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.common.lifecycle.Lifecycle;
 import org.opensearch.common.lifecycle.LifecycleComponent;
+import org.opensearch.common.lifecycle.LifecycleListener;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.IndexScopedSettings;
 import org.opensearch.common.settings.Setting;
@@ -37,6 +40,7 @@
 import org.opensearch.jobscheduler.spi.JobSchedulerExtension;
 import org.opensearch.jobscheduler.spi.ScheduledJobParser;
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner;
+import org.opensearch.jobscheduler.spi.utils.LockService;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.ClusterPlugin;
 import org.opensearch.plugins.EnginePlugin;
@@ -283,6 +287,8 @@ public class SecurityAnalyticsPlugin extends Plugin implements ActionPlugin, Map
 
     private SATIFSourceConfigService saTifSourceConfigService;
 
+    private TIFLockService threatIntelLockService;
+
     @Override
     public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings settings) {
         List<SystemIndexDescriptor> descriptors = List.of(
@@ -321,7 +327,7 @@ public Collection<Object> createComponents(Client client,
         DetectorThreatIntelService detectorThreatIntelService = new DetectorThreatIntelService(threatIntelFeedDataService, client, xContentRegistry);
         TIFJobParameterService tifJobParameterService = new TIFJobParameterService(client, clusterService);
         TIFJobUpdateService tifJobUpdateService = new TIFJobUpdateService(clusterService, tifJobParameterService, threatIntelFeedDataService, builtInTIFMetadataLoader);
-        TIFLockService threatIntelLockService = new TIFLockService(clusterService, client);
+        threatIntelLockService = new TIFLockService(clusterService, client);
         saTifSourceConfigService = new SATIFSourceConfigService(client, clusterService, threadPool, xContentRegistry, threatIntelLockService);
         STIX2IOCFetchService stix2IOCFetchService = new STIX2IOCFetchService(client, clusterService);
         SATIFSourceConfigManagementService saTifSourceConfigManagementService = new SATIFSourceConfigManagementService(saTifSourceConfigService, threatIntelLockService, stix2IOCFetchService, xContentRegistry, clusterService);
@@ -344,7 +350,7 @@ public Collection<Object> createComponents(Client client,
 
     @Override
     public Collection<Class<? extends LifecycleComponent>> getGuiceServiceClasses() {
-        return List.of(DetectorIndexManagementService.class, BuiltinLogTypeLoader.class);
+        return List.of(DetectorIndexManagementService.class, BuiltinLogTypeLoader.class, GuiceHolder.class);
     }
 
     @Override
@@ -575,6 +581,9 @@ public void onFailure(Exception e) {
                 log.warn("Failed to initialize LogType config index and builtin log types");
             }
         });
+
+        LockService lockService = GuiceHolder.getLockService();
+        threatIntelLockService.initialize(lockService);
     }
 
     @NonNull
@@ -584,4 +593,39 @@ public Map<String, RemoteMonitorRunner> getMonitorTypesToMonitorRunners() {
                 THREAT_INTEL_MONITOR_TYPE, ThreatIntelMonitorRunner.getMonitorRunner()
         );
     }
+
+    public static class GuiceHolder implements LifecycleComponent {
+
+        private static LockService lockService;
+
+        @Inject
+        public GuiceHolder(final LockService lockService) {
+            GuiceHolder.lockService = lockService;
+        }
+
+        static LockService getLockService() {
+            return lockService;
+        }
+
+        @Override
+        public void close() {}
+
+        @Override
+        public Lifecycle.State lifecycleState() {
+            return null;
+        }
+
+        @Override
+        public void addLifecycleListener(LifecycleListener listener) {}
+
+        @Override
+        public void removeLifecycleListener(LifecycleListener listener) {}
+
+        @Override
+        public void start() {}
+
+        @Override
+        public void stop() {}
+
+    }
 }

src/main/java/org/opensearch/securityanalytics/threatIntel/common/TIFLockService.java

@@ -31,7 +31,7 @@ public class TIFLockService {
     public static final long LOCK_DURATION_IN_SECONDS = 300l;
     public static final long RENEW_AFTER_IN_SECONDS = 120l;
     private final ClusterService clusterService;
-    private final LockService lockService;
+    private LockService lockService;
 
 
     /**
@@ -42,7 +42,10 @@ public class TIFLockService {
      */
     public TIFLockService(final ClusterService clusterService, final Client client) {
         this.clusterService = clusterService;
-        this.lockService = new LockService(client, clusterService);
+    }
+
+    public void initialize(final LockService lockService) {
+        this.lockService = lockService;
     }
 
     /**

src/test/java/org/opensearch/securityanalytics/correlation/CorrelationEngineRestApiIT.java

@@ -679,6 +679,8 @@ public void testBasicCorrelationEngineWorkflowWithFieldBasedRulesOnMultipleLogTy
         );
     }
 
+    // broken by https://github.com/opensearch-project/common-utils/pull/829
+    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/common-utils/pull/829")
     public  void testBasicCorrelationEngineWorkflowWithIndexPatterns() throws IOException, InterruptedException {
         updateClusterSetting(SecurityAnalyticsSettings.ENABLE_AUTO_CORRELATIONS.getKey(), "false");
 

src/test/java/org/opensearch/securityanalytics/mapper/MapperRestApiIT.java

@@ -715,6 +715,8 @@ public void testCreateMappings_withIndexPattern_existing_indexTemplate_update_su
         assertTrue(props.containsKey("destination.port"));
     }
 
+    // broken by https://github.com/opensearch-project/common-utils/pull/829
+    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/common-utils/pull/829")
     public void testCreateMappings_withIndexPattern_differentMappings_indexTemplateCleanup_success() throws IOException, InterruptedException {
         String indexName1 = "test_index_1";
         String indexName2 = "test_index_2";

src/test/java/org/opensearch/securityanalytics/threatIntel/ThreatIntelTestCase.java

@@ -52,6 +52,7 @@
 import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 import org.opensearch.securityanalytics.TestHelpers;
@@ -102,7 +103,7 @@ public void prepareThreatIntelTestCase() {
         client = new NoOpNodeClient(this.getTestName());
         verifyingClient = spy(new VerifyingClient(this.getTestName()));
         clusterSettings = new ClusterSettings(settings, new HashSet<>(SecurityAnalyticsSettings.settings()));
-        lockService = new LockService(client, clusterService);
+        lockService = mock(LockService.class);
         ingestMetadata = new IngestMetadata(Collections.emptyMap());
         when(metadata.custom(IngestMetadata.TYPE)).thenReturn(ingestMetadata);
         when(clusterService.getSettings()).thenReturn(Settings.EMPTY);

src/test/java/org/opensearch/securityanalytics/threatIntel/common/ThreatIntelLockServiceTests.java

@@ -14,59 +14,63 @@
 
 import org.junit.Assert;
 import org.junit.Before;
+import org.mockito.Mockito;
 import org.opensearch.action.DocWriteResponse;
 import org.opensearch.action.update.UpdateRequest;
 import org.opensearch.action.update.UpdateResponse;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.jobscheduler.spi.LockModel;
+import org.opensearch.jobscheduler.spi.utils.LockService;
 import org.opensearch.securityanalytics.threatIntel.ThreatIntelTestCase;
 import org.opensearch.securityanalytics.TestHelpers;
 
 public class ThreatIntelLockServiceTests extends ThreatIntelTestCase {
     private TIFLockService threatIntelLockService;
-    private TIFLockService noOpsLockService;
 
     @Before
     public void init() {
         threatIntelLockService = new TIFLockService(clusterService, verifyingClient);
-        noOpsLockService = new TIFLockService(clusterService, client);
-    }
-
-    public void testAcquireLock_whenValidInput_thenSucceed() {
-        // Cannot test because LockService is final class
-        // Simply calling method to increase coverage
-        noOpsLockService.acquireLock(TestHelpers.randomLowerCaseString(), randomPositiveLong(), mock(ActionListener.class));
+        threatIntelLockService.initialize(lockService);
     }
 
     public void testAcquireLock_whenCalled_thenNotBlocked() {
         long expectedDurationInMillis = 1000;
+
+        Mockito.doAnswer(inv -> {
+                    ActionListener<LockModel> listener = inv.getArgument(3);
+                    listener.onResponse(null);          // or listener.onFailure(ex);
+                    return null;                        // because the real method is void
+                })
+                .when(lockService)
+                .acquireLockWithId(
+                        Mockito.any(), // jobIndexName you expect
+                        Mockito.any(), // lockDurationSeconds you expect
+                        Mockito.any(), // lockId you expect
+                        Mockito.any()  // listener – generics erase to ActionListener
+                );
         Instant before = Instant.now();
         threatIntelLockService.acquireLock(null, null, ActionListener.wrap(
-                r -> fail("Should not have been blocked"), e -> {
+                r -> {
                     Instant after = Instant.now();
                     assertTrue(after.toEpochMilli() - before.toEpochMilli() < expectedDurationInMillis);
-                }
-        ));
-    }
-
-    public void testReleaseLock_whenValidInput_thenSucceed() {
-        // Cannot test because LockService is final class
-        // Simply calling method to increase coverage
-        LockModel lockModel = new LockModel(
-                TestHelpers.randomLowerCaseString(),
-                TestHelpers.randomLowerCaseString(),
-                Instant.now(),
-                LOCK_DURATION_IN_SECONDS,
-                false
-        );
-        noOpsLockService.releaseLock(lockModel, ActionListener.wrap(
-                Assert::assertFalse, e -> fail()
+                }, e -> fail("Should not have failed")
         ));
     }
 
     public void testRenewLock_whenCalled_thenNotBlocked() {
         long expectedDurationInMillis = 1000;
+
+        Mockito.doAnswer(inv -> {
+                    ActionListener<LockModel> listener = inv.getArgument(1);
+                    listener.onResponse(null);          // or listener.onFailure(ex);
+                    return null;                        // because the real method is void
+                })
+                .when(lockService)
+                .renewLock(
+                        Mockito.any(), // lockModel
+                        Mockito.any()  // listener – generics erase to ActionListener
+                );
         Instant before = Instant.now();
         assertNull(threatIntelLockService.renewLock(null));
         Instant after = Instant.now();