Skip to content

Commit 771feac

Browse files
Davidding4718graytaylor0
Davidding4718
and
graytaylor0
authored
Add change to s3 scan source to create partitions after each page of … (#6006)
* Add change to s3 scan source to create partitions after each page of objects listed Signed-off-by: Taylor Gray <tylgry@amazon.com> (cherry picked from commit bf74cb7) * Renew the unit test to handle the two scans and seperate buckets Signed-off-by: Siqi Ding <dingdd@amazon.com> * Add global state ownership renewal during partition creation Signed-off-by: Siqi Ding <dingdd@amazon.com> * Add ownership renewal during S3 partition creation Prevents timeout during long-running operations by renewing global state ownership every 3 minutes Signed-off-by: Siqi Ding <dingdd@amazon.com> * Remove the each time update in ScanObjectWorker; Add the lastGlobalOwnershipRenewal related renewGlobalOwnershipIfNeeded to ensures ownership is maintained both before starting and throughout the entire partition creation process Signed-off-by: Siqi Ding <dingdd@amazon.com> * Make renewGlobalStateOwnershipForPartitionCreation private Signed-off-by: Siqi Ding <dingdd@amazon.com> --------- Signed-off-by: Taylor Gray <tylgry@amazon.com> Signed-off-by: Siqi Ding <dingdd@amazon.com> Co-authored-by: Taylor Gray <tylgry@amazon.com>
1 parent 0266303 commit 771feac

5 files changed

Lines changed: 213 additions & 82 deletions

File tree

‎data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/SourceCoordinator.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,12 @@ public interface SourceCoordinator<T> {
6363
*/
6464
Optional<SourcePartition<T>> getNextPartition(final Function<Map<String, Object>, List<PartitionIdentifier>> partitionCreationSupplier, final boolean forceSupplier);
6565

66+
/**
67+
* Can be used to directly create partitions for source coordination, as an alternative to relying on getNextPartition to create partitions.
68+
* @param partitionIdentifiers - The partitions to be created.
69+
*/
70+
void createPartitions(final List<PartitionIdentifier> partitionIdentifiers);
71+
6672
/**
6773
* Should be called by the source when it has fully processed a given partition
6874
* @throws org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionNotFoundException if the partition key could not be found in the distributed store

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ public class LeaseBasedSourceCoordinator<T> implements SourceCoordinator<T> {
6060
static final String PARTITION_NOT_OWNED_ERROR_COUNT = "partitionNotOwnedErrors";
6161
static final String PARTITION_UPDATE_ERROR_COUNT = "PartitionUpdateErrors";
6262
static final Duration DEFAULT_LEASE_TIMEOUT = Duration.ofMinutes(10);
63+
static final Duration DEFAULT_RENEW_TIMEOUT = Duration.ofMinutes(3);
64+
static final Duration OWNERSHIP_RENEWAL_INTERVAL = Duration.ofMinutes(3);
6365

6466
private static final String hostName;
6567
static final String PARTITION_TYPE = "PARTITION";
@@ -94,6 +96,7 @@ public class LeaseBasedSourceCoordinator<T> implements SourceCoordinator<T> {
9496
private final ReentrantLock lock;
9597

9698
private Instant lastSupplierRunTime;
99+
private Instant lastGlobalOwnershipRenewal;
97100

98101
static final Duration FORCE_SUPPLIER_AFTER_DURATION = Duration.ofMinutes(5);
99102

@@ -136,6 +139,7 @@ public LeaseBasedSourceCoordinator(final Class<T> partitionProgressStateClass,
136139
this.partitionsDeleted = pluginMetrics.counter(PARTITIONS_DELETED);
137140
this.lock = new ReentrantLock();
138141
this.lastSupplierRunTime = Instant.now();
142+
this.lastGlobalOwnershipRenewal = Instant.now();
139143
}
140144

141145
@Override
@@ -202,8 +206,13 @@ private Optional<SourcePartition<T>> getNextPartitionInternal(final Function<Map
202206
return Optional.of(sourcePartition);
203207
}
204208

205-
private void createPartitions(final List<PartitionIdentifier> partitionIdentifiers) {
209+
@Override
210+
public void createPartitions(final List<PartitionIdentifier> partitionIdentifiers) {
211+
renewGlobalOwnershipIfNeeded();
212+
206213
for (final PartitionIdentifier partitionIdentifier : partitionIdentifiers) {
214+
renewGlobalOwnershipIfNeeded();
215+
207216
final Optional<SourcePartitionStoreItem> optionalPartitionItem = sourceCoordinationStore.getSourcePartitionItem(sourceIdentifierWithPartitionType, partitionIdentifier.getPartitionKey());
208217

209218
if (optionalPartitionItem.isPresent()) {
@@ -514,6 +523,28 @@ private Optional<SourcePartitionStoreItem> acquireGlobalStateForPartitionCreatio
514523
return globalStateItemForPartitionCreation;
515524
}
516525

526+
private void renewGlobalStateOwnershipForPartitionCreation() {
527+
try {
528+
final Optional<SourcePartitionStoreItem> globalStateItem = sourceCoordinationStore.getSourcePartitionItem(
529+
sourceIdentifierWithGlobalStateType, GLOBAL_STATE_SOURCE_PARTITION_KEY_FOR_CREATING_PARTITIONS);
530+
531+
if (globalStateItem.isPresent() && ownerId.equals(globalStateItem.get().getPartitionOwner())) {
532+
globalStateItem.get().setPartitionOwnershipTimeout(Instant.now().plus(DEFAULT_LEASE_TIMEOUT));
533+
sourceCoordinationStore.tryUpdateSourcePartitionItem(globalStateItem.get());
534+
}
535+
} catch (final Exception e) {
536+
LOG.warn("Failed to renew global state ownership for partition creation", e);
537+
}
538+
}
539+
540+
private void renewGlobalOwnershipIfNeeded() {
541+
if (Instant.now().isAfter(lastGlobalOwnershipRenewal.plus(OWNERSHIP_RENEWAL_INTERVAL))) {
542+
LOG.info("Renewing global state ownership for partition creation");
543+
renewGlobalStateOwnershipForPartitionCreation();
544+
lastGlobalOwnershipRenewal = Instant.now();
545+
}
546+
}
547+
517548
private void giveUpAndSaveGlobalStateForPartitionCreation(final SourcePartitionStoreItem globalStateItemForPartitionCreation, final Map<String, Object> globalStateMap) {
518549

519550
globalStateItemForPartitionCreation.setSourcePartitionStatus(SourcePartitionStatus.UNASSIGNED);

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

Lines changed: 44 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.dataprepper.plugins.source.s3;
77

88
import org.opensearch.dataprepper.model.source.coordinator.PartitionIdentifier;
9+
import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator;
910
import org.opensearch.dataprepper.plugins.source.s3.configuration.FolderPartitioningOptions;
1011
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3ScanKeyPathOption;
1112
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3ScanSchedulingOptions;
@@ -53,19 +54,23 @@ public class S3ScanPartitionCreationSupplier implements Function<Map<String, Obj
5354

5455
private final boolean deleteS3ObjectsOnRead;
5556

57+
private final SourceCoordinator<S3SourceProgressState> sourceCoordinator;
58+
5659
public S3ScanPartitionCreationSupplier(final S3Client s3Client,
5760
final BucketOwnerProvider bucketOwnerProvider,
5861
final List<ScanOptions> scanOptionsList,
5962
final S3ScanSchedulingOptions schedulingOptions,
6063
final FolderPartitioningOptions folderPartitioningOptions,
61-
final boolean deleteS3ObjectsOnRead) {
64+
final boolean deleteS3ObjectsOnRead,
65+
final SourceCoordinator<S3SourceProgressState> sourceCoordinator) {
6266

6367
this.s3Client = s3Client;
6468
this.bucketOwnerProvider = bucketOwnerProvider;
6569
this.scanOptionsList = scanOptionsList;
6670
this.schedulingOptions = schedulingOptions;
6771
this.folderPartitioningOptions = folderPartitioningOptions;
6872
this.deleteS3ObjectsOnRead = deleteS3ObjectsOnRead;
73+
this.sourceCoordinator = sourceCoordinator;
6974
}
7075

7176
@Override
@@ -99,12 +104,13 @@ public List<PartitionIdentifier> apply(final Map<String, Object> globalStateMap)
99104
if (Objects.nonNull(s3ScanKeyPathOption) && Objects.nonNull(s3ScanKeyPathOption.getS3scanIncludePrefixOptions()))
100105
s3ScanKeyPathOption.getS3scanIncludePrefixOptions().forEach(includePath -> {
101106
listObjectsV2Request.prefix(includePath);
102-
objectsToProcess.addAll(listFilteredS3ObjectsForBucket(excludeItems, listObjectsV2Request,
103-
bucketName, scanOptions.getUseStartDateTime(), scanOptions.getUseEndDateTime(), globalStateMap));
107+
createFilteredS3ObjectPartitionsForBucket(excludeItems, listObjectsV2Request,
108+
bucketName, scanOptions.getUseStartDateTime(), scanOptions.getUseEndDateTime(), globalStateMap);
104109
});
105110
else
106-
objectsToProcess.addAll(listFilteredS3ObjectsForBucket(excludeItems, listObjectsV2Request,
107-
bucketName, scanOptions.getUseStartDateTime(), scanOptions.getUseEndDateTime(), globalStateMap));
111+
createFilteredS3ObjectPartitionsForBucket(excludeItems, listObjectsV2Request,
112+
bucketName, scanOptions.getUseStartDateTime(), scanOptions.getUseEndDateTime(), globalStateMap);
113+
108114
if (!bucketScanTime.containsKey(bucketName)) {
109115
bucketScanTime.put(bucketName, updatedScanTime.toString());
110116
}
@@ -117,22 +123,23 @@ public List<PartitionIdentifier> apply(final Map<String, Object> globalStateMap)
117123
globalStateMap.put(SCAN_COUNT, (Integer) globalStateMap.get(SCAN_COUNT) + 1);
118124
globalStateMap.put(LAST_SCAN_TIME, Instant.now().toEpochMilli());
119125

120-
return objectsToProcess;
126+
// Partitions are created by the supplier, so source coordinator does not need to attempt to create any partitions
127+
return Collections.emptyList();
121128
}
122129

123-
private List<PartitionIdentifier> listFilteredS3ObjectsForBucket(final List<String> excludeKeyPaths,
130+
private void createFilteredS3ObjectPartitionsForBucket(final List<String> excludeKeyPaths,
124131
final ListObjectsV2Request.Builder listObjectsV2Request,
125132
final String bucket,
126133
final LocalDateTime startDateTime,
127134
final LocalDateTime endDateTime,
128135
final Map<String, Object> globalStateMap) {
129136
final Instant previousScanTime = globalStateMap.get(bucket) != null ? Instant.parse((String) globalStateMap.get(bucket)) : null;
130137
final boolean isFirstScan = previousScanTime == null;
131-
final List<PartitionIdentifier> allPartitionIdentifiers = new ArrayList<>();
132138
ListObjectsV2Response listObjectsV2Response = null;
139+
133140
do {
134141
listObjectsV2Response = s3Client.listObjectsV2(listObjectsV2Request.fetchOwner(true).continuationToken(Objects.nonNull(listObjectsV2Response) ? listObjectsV2Response.nextContinuationToken() : null).build());
135-
allPartitionIdentifiers.addAll(listObjectsV2Response.contents().stream()
142+
final List<PartitionIdentifier> partitionsForPage = listObjectsV2Response.contents().stream()
136143
.filter(s3Object -> deleteS3ObjectsOnRead || isLastModifiedTimeAfterMostRecentScanForBucket(previousScanTime, s3Object))
137144
.map(s3Object -> Pair.of(s3Object.key(), instantToLocalDateTime(s3Object.lastModified())))
138145
.filter(keyTimestampPair -> !keyTimestampPair.left().endsWith("/"))
@@ -141,32 +148,17 @@ private List<PartitionIdentifier> listFilteredS3ObjectsForBucket(final List<Stri
141148
.filter(keyTimestampPair -> isKeyMatchedBetweenTimeRange(keyTimestampPair.right(), startDateTime, endDateTime, isFirstScan))
142149
.map(Pair::left)
143150
.map(objectKey -> PartitionIdentifier.builder().withPartitionKey(String.format(BUCKET_OBJECT_PARTITION_KEY_FORMAT, bucket, objectKey)).build())
144-
.collect(Collectors.toList()));
145-
151+
.collect(Collectors.toList());
146152
LOG.info("Found page of {} objects from bucket {}", listObjectsV2Response.keyCount(), bucket);
147-
} while (listObjectsV2Response.isTruncated());
148-
149-
if (folderPartitioningOptions != null) {
150-
final Set<PartitionIdentifier> folderPartitions = allPartitionIdentifiers.stream()
151-
.map(partitionIdentifier -> {
152-
final String fullObjectKey = partitionIdentifier.getPartitionKey();
153-
final String prefix = getPrefixWithDepth(fullObjectKey);
154-
if (prefix == null) {
155-
return null;
156-
}
157-
return PartitionIdentifier.builder().withPartitionKey(prefix).build();
158-
})
159-
.filter(Objects::nonNull)
160-
.collect(Collectors.toSet());
161-
162-
LOG.info("Running in folder_partitions mode at depth {}, found {} unique prefixes from {} objects", folderPartitioningOptions.getFolderDepth(), folderPartitions.size(), allPartitionIdentifiers.size());
163-
164-
return new ArrayList<>(folderPartitions);
165-
} else {
166-
LOG.info("Returning partitions for {} S3 objects from bucket {}", allPartitionIdentifiers.size(), bucket);
167-
}
168153

169-
return allPartitionIdentifiers;
154+
if (folderPartitioningOptions != null) {
155+
final List<PartitionIdentifier> folderPartitionsForPage = getFolderPartitionIdentifiers(partitionsForPage);
156+
sourceCoordinator.createPartitions(folderPartitionsForPage);
157+
} else {
158+
LOG.info("Creating partitions for {} S3 objects from bucket {}", partitionsForPage.size(), bucket);
159+
sourceCoordinator.createPartitions(partitionsForPage);
160+
}
161+
} while (listObjectsV2Response.isTruncated());
170162
}
171163

172164
private LocalDateTime instantToLocalDateTime(final Instant instant) {
@@ -264,4 +256,23 @@ private String getPrefixWithDepth(final String fullObjectKey) {
264256
int actualDepth = min(folderPartitioningOptions.getFolderDepth(), folders.length - 1);
265257
return String.join("/", Arrays.copyOfRange(folders, 0, actualDepth)) + "/";
266258
}
259+
260+
private List<PartitionIdentifier> getFolderPartitionIdentifiers(final List<PartitionIdentifier> objectPartitionIdentifiers) {
261+
262+
final Set<PartitionIdentifier> folderPartitions = objectPartitionIdentifiers.stream()
263+
.map(partitionIdentifier -> {
264+
final String fullObjectKey = partitionIdentifier.getPartitionKey();
265+
final String prefix = getPrefixWithDepth(fullObjectKey);
266+
if (prefix == null) {
267+
return null;
268+
}
269+
return PartitionIdentifier.builder().withPartitionKey(prefix).build();
270+
})
271+
.filter(Objects::nonNull)
272+
.collect(Collectors.toSet());
273+
274+
LOG.info("Running in folder_partitions mode at depth {}, found {} unique prefixes from {} objects", folderPartitioningOptions.getFolderDepth(), folderPartitions.size(), objectPartitionIdentifiers.size());
275+
276+
return new ArrayList<>(folderPartitions);
277+
}
267278
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ public ScanObjectWorker(final S3Client s3Client,
156156
this.folderPartitioningOptions = s3SourceConfig.getS3ScanScanOptions().getPartitioningOptions();
157157
this.acknowledgmentSetTimeout = s3SourceConfig.getS3ScanScanOptions().getAcknowledgmentTimeout();
158158

159-
this.partitionCreationSupplier = new S3ScanPartitionCreationSupplier(s3Client, bucketOwnerProvider, scanOptionsBuilderList, s3ScanSchedulingOptions, s3SourceConfig.getS3ScanScanOptions().getPartitioningOptions(), s3SourceConfig.isDeleteS3ObjectsOnRead());
159+
this.partitionCreationSupplier = new S3ScanPartitionCreationSupplier(s3Client, bucketOwnerProvider, scanOptionsBuilderList, s3ScanSchedulingOptions, s3SourceConfig.getS3ScanScanOptions().getPartitioningOptions(), s3SourceConfig.isDeleteS3ObjectsOnRead(), sourceCoordinator);
160160
this.acknowledgmentsRemainingForPartitions = new ConcurrentHashMap<>();
161161
this.objectsToDeleteForAcknowledgmentSets = new ConcurrentHashMap<>();
162162
}
@@ -373,7 +373,6 @@ private boolean shouldDeleteFolderPartition(final SourcePartition<S3SourceProgre
373373
return false;
374374
}
375375

376-
377376
private void processObjectsForFolderPartition(final List<S3ObjectReference> objectsToProcess,
378377
final SourcePartition<S3SourceProgressState> folderPartition) {
379378
int objectsProcessed = 0;

0 commit comments

Comments
 (0)