[Backport 2.4] Fix issue of skipping new partitions for source coordinator

data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/LeaseBasedSourceCoordinator.java

@@ -183,7 +183,7 @@ private void createPartitions(final List<PartitionIdentifier> partitionIdentifie
             final Optional<SourcePartitionStoreItem> optionalPartitionItem = sourceCoordinationStore.getSourcePartitionItem(sourceIdentifierWithPartitionType, partitionIdentifier.getPartitionKey());
 
             if (optionalPartitionItem.isPresent()) {
-                return;
+                continue;
             }
 
             final boolean partitionCreated = sourceCoordinationStore.tryCreatePartitionItem(

data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/LeaseBasedSourceCoordinatorTest.java

@@ -188,15 +188,17 @@ void getNextPartition_throws_UninitializedSourceCoordinatorException_when_called
     }
 
     @Test
-    void getNextPartition_calls_supplier_and_creates_partition_with_non_existing_item_when_partition_exists_and_is_created_successfully() {
+    void getNextPartition_calls_supplier_and_creates_partition_with_existing_then_non_existing_item_when_partition_exists_and_is_created_successfully() {
         final PartitionIdentifier partitionIdentifier = PartitionIdentifier.builder().withPartitionKey(UUID.randomUUID().toString()).build();
-        final Function<Map<String, Object>, List<PartitionIdentifier>> partitionCreationSupplier = (map) -> List.of(partitionIdentifier);
+        final PartitionIdentifier partitionIdentifierToSkip = PartitionIdentifier.builder().withPartitionKey(UUID.randomUUID().toString()).build();
+        final Function<Map<String, Object>, List<PartitionIdentifier>> partitionCreationSupplier = (map) -> List.of(partitionIdentifierToSkip, partitionIdentifier);
 
         given(sourceCoordinationStore.tryAcquireAvailablePartition(anyString(), anyString(), any())).willReturn(Optional.empty()).willReturn( Optional.empty());
         given(globalStateForPartitionCreationItem.getSourcePartitionStatus()).willReturn(SourcePartitionStatus.UNASSIGNED);
         given(globalStateForPartitionCreationItem.getPartitionOwner()).willReturn(null);
         given(sourceCoordinationStore.getSourcePartitionItem(fullSourceIdentifierForGlobalState, GLOBAL_STATE_SOURCE_PARTITION_KEY_FOR_CREATING_PARTITIONS)).willReturn(Optional.of(globalStateForPartitionCreationItem));
         doNothing().when(sourceCoordinationStore).tryUpdateSourcePartitionItem(globalStateForPartitionCreationItem);
+        given(sourceCoordinationStore.getSourcePartitionItem(fullSourceIdentifierForPartition, partitionIdentifierToSkip.getPartitionKey())).willReturn(Optional.of(sourcePartitionStoreItem));
         given(sourceCoordinationStore.getSourcePartitionItem(fullSourceIdentifierForPartition, partitionIdentifier.getPartitionKey())).willReturn(Optional.empty());
         given(sourceCoordinationStore.tryCreatePartitionItem(fullSourceIdentifierForPartition, partitionIdentifier.getPartitionKey(), SourcePartitionStatus.UNASSIGNED, 0L, null)).willReturn(true);
 

data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSource.java

@@ -17,6 +17,8 @@
 import org.opensearch.dataprepper.plugins.source.opensearch.worker.client.SearchAccessor;
 import org.opensearch.dataprepper.plugins.source.opensearch.worker.client.SearchAccessorStrategy;
 
+import java.util.Objects;
+
 @DataPrepperPlugin(name="opensearch", pluginType = Source.class, pluginConfigurationType = OpenSearchSourceConfiguration.class)
 public class OpenSearchSource implements Source<Record<Event>>, UsesSourceCoordination {
 
@@ -56,7 +58,9 @@ private void startProcess(final OpenSearchSourceConfiguration openSearchSourceCo
 
     @Override
     public void stop() {
-        openSearchService.stop();
+        if (Objects.nonNull(openSearchService)) {
+            openSearchService.stop();
+        }
     }
 
     @Override

data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/OpenSearchIndexPartitionCreationSupplier.java

@@ -76,6 +76,8 @@ private List<PartitionIdentifier> applyForOpenSearchClient(final Map<String, Obj
             return Collections.emptyList();
         }
 
+        LOG.debug("Found {} indices", indicesResponse.valueBody().size());
+
         return indicesResponse.valueBody().stream()
                 .filter(osIndicesRecord -> shouldIndexBeProcessed(osIndicesRecord.index()))
                 .map(indexRecord -> PartitionIdentifier.builder().withPartitionKey(indexRecord.index()).build())
@@ -91,6 +93,8 @@ private List<PartitionIdentifier> applyForElasticSearchClient(final Map<String,
             return Collections.emptyList();
         }
 
+        LOG.debug("Found {} indices", indicesResponse.valueBody().size());
+
         return indicesResponse.valueBody().stream()
                 .filter(esIndicesRecord -> shouldIndexBeProcessed(esIndicesRecord.index()))
                 .map(indexRecord -> PartitionIdentifier.builder().withPartitionKey(indexRecord.index()).build())

data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/S3ScanPartitionCreationSupplier.java

@@ -100,7 +100,6 @@ private List<PartitionIdentifier> listFilteredS3ObjectsForBucket(final List<Stri
                                                                      final LocalDateTime startDateTime,
                                                                      final LocalDateTime endDateTime,
                                                                      final Map<String, Object> globalStateMap) {
-
         Instant mostRecentLastModifiedTimestamp = globalStateMap.get(bucket) != null ? Instant.parse((String) globalStateMap.get(bucket)) : null;
         final List<PartitionIdentifier> allPartitionIdentifiers = new ArrayList<>();
         ListObjectsV2Response listObjectsV2Response = null;