Adding subshard work items on lease expiry

DataGenerator/src/main/java/org/opensearch/migrations/data/WorkloadGenerator.java

@@ -53,7 +53,8 @@ private List<CompletableFuture<?>> generateDocs(String indexName, Workload workl
                 log.atTrace().setMessage("Created doc for index {}: {}")
                     .addArgument(indexName)
                     .addArgument(doc::toString).log();
-                return new BulkDocSection(indexName + "_" + docIdCounter.incrementAndGet(), indexName, null, doc.toString());
+                var docId = docIdCounter.incrementAndGet();
+                return new BulkDocSection(indexName + "_" + docId, indexName, null, doc.toString(), null);
             })
             .collect(Collectors.toList());
 

DocumentsFromSnapshotMigration/src/main/java/org/opensearch/migrations/RfsMigrateDocuments.java

@@ -5,8 +5,13 @@
 import java.nio.file.Paths;
 import java.time.Clock;
 import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import org.opensearch.migrations.bulkload.common.DefaultSourceRepoAccessor;
 import org.opensearch.migrations.bulkload.common.DocumentReindexer;
@@ -20,13 +25,16 @@
 import org.opensearch.migrations.bulkload.common.http.ConnectionContext;
 import org.opensearch.migrations.bulkload.models.IndexMetadata;
 import org.opensearch.migrations.bulkload.models.ShardMetadata;
+import org.opensearch.migrations.bulkload.tracing.IWorkCoordinationContexts;
 import org.opensearch.migrations.bulkload.workcoordination.CoordinateWorkHttpClient;
 import org.opensearch.migrations.bulkload.workcoordination.IWorkCoordinator;
 import org.opensearch.migrations.bulkload.workcoordination.LeaseExpireTrigger;
 import org.opensearch.migrations.bulkload.workcoordination.OpenSearchWorkCoordinator;
 import org.opensearch.migrations.bulkload.workcoordination.ScopedWorkCoordinator;
+import org.opensearch.migrations.bulkload.workcoordination.WorkItemTimeProvider;
 import org.opensearch.migrations.bulkload.worker.DocumentsRunner;
 import org.opensearch.migrations.bulkload.worker.ShardWorkPreparer;
+import org.opensearch.migrations.bulkload.worker.WorkItemCursor;
 import org.opensearch.migrations.cluster.ClusterProviderRegistry;
 import org.opensearch.migrations.reindexer.tracing.RootDocumentMigrationContext;
 import org.opensearch.migrations.tracing.ActiveContextTracker;
@@ -45,6 +53,7 @@
 import com.beust.jcommander.ParameterException;
 import com.beust.jcommander.ParametersDelegate;
 import lombok.Getter;
+import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.slf4j.MDC;
 
@@ -55,6 +64,9 @@ public class RfsMigrateDocuments {
     public static final int TOLERABLE_CLIENT_SERVER_CLOCK_DIFFERENCE_SECONDS = 5;
     public static final String LOGGING_MDC_WORKER_ID = "workerId";
 
+    // Increase successor nextAcquisitionLeaseExponent if shard setup takes more than 10% of lease total time
+    private static final double SHARD_SETUP_LEASE_DURATION_THRESHOLD = 0.1;
+
     public static final String DEFAULT_DOCUMENT_TRANSFORMATION_CONFIG = "[" +
             "  {" +
             "    \"JsonTransformerForDocumentTypeRemovalProvider\":\"\"" +
@@ -270,11 +282,28 @@ public static void main(String[] args) throws Exception {
         }
         IJsonTransformer docTransformer = new TransformationLoader().getTransformerFactoryLoader(docTransformerConfig);
 
-        try (var processManager = new LeaseExpireTrigger(RfsMigrateDocuments::exitOnLeaseTimeout, Clock.systemUTC());
-             var workCoordinator = new OpenSearchWorkCoordinator(
+        var workItemRef = new AtomicReference<IWorkCoordinator.WorkItemAndDuration>();
+        var progressCursor = new AtomicReference<WorkItemCursor>();
+        var cancellationRunnableRef = new AtomicReference<Runnable>();
+        var workItemTimeProvider = new WorkItemTimeProvider();
+        try (var workCoordinator = new OpenSearchWorkCoordinator(
                  new CoordinateWorkHttpClient(connectionContext),
                  TOLERABLE_CLIENT_SERVER_CLOCK_DIFFERENCE_SECONDS,
-                 workerId)
+                 workerId,
+                Clock.systemUTC(),
+                workItemRef::set);
+             var processManager = new LeaseExpireTrigger(
+                w -> exitOnLeaseTimeout(
+                        workItemRef,
+                        workCoordinator,
+                        w,
+                        progressCursor,
+                        workItemTimeProvider,
+                        arguments.initialLeaseDuration,
+                        () -> Optional.ofNullable(cancellationRunnableRef.get()).ifPresent(Runnable::run),
+                        context.getWorkCoordinationContext()::createSuccessorWorkItemsContext),
+                Clock.systemUTC()
+            );
         ) {
             MDC.put(LOGGING_MDC_WORKER_ID, workerId); // I don't see a need to clean this up since we're in main
             OpenSearchClient targetClient = new OpenSearchClient(connectionContext);
@@ -307,6 +336,7 @@ public static void main(String[] args) throws Exception {
             run(
                 LuceneDocumentsReader.getFactory(sourceResourceProvider),
                 reindexer,
+                progressCursor,
                 workCoordinator,
                 arguments.initialLeaseDuration,
                 processManager,
@@ -316,7 +346,9 @@ public static void main(String[] args) throws Exception {
                 sourceResourceProvider.getShardMetadata(),
                 unpackerFactory,
                 arguments.maxShardSizeBytes,
-                context);
+                context,
+                cancellationRunnableRef,
+                workItemTimeProvider);
         } catch (NoWorkLeftException e) {
             log.atWarn().setMessage("No work left to acquire.  Exiting with error code to signal that.").log();
             System.exit(NO_WORK_LEFT_EXIT_CODE);
@@ -326,11 +358,105 @@ public static void main(String[] args) throws Exception {
         }
     }
 
-    private static void exitOnLeaseTimeout(String workItemId) {
-        log.error("Terminating RfsMigrateDocuments because the lease has expired for " + workItemId);
+    @SneakyThrows
+    private static void exitOnLeaseTimeout(
+            AtomicReference<IWorkCoordinator.WorkItemAndDuration> workItemRef,
+            IWorkCoordinator coordinator,
+            String workItemId, 
+            AtomicReference<WorkItemCursor> progressCursorRef,
+            WorkItemTimeProvider workItemTimeProvider,
+            Duration initialLeaseDuration,
+            Runnable cancellationRunnable,
+            Supplier<IWorkCoordinationContexts.ICreateSuccessorWorkItemsContext> contextSupplier
+    ) {
+        log.atWarn().setMessage("Terminating RfsMigrateDocuments because the lease has expired for {}")
+                .addArgument(workItemId)
+                .log();
+        if (progressCursorRef.get() != null) {
+            log.atWarn().setMessage("Progress cursor set, cancelling active doc migration if still running").log();
+            cancellationRunnable.run();
+            // Get a new progressCursor after cancellation for most up-to-date checkpoint
+            var progressCursor = progressCursorRef.get();
+            log.atWarn().setMessage("Progress cursor: {}")
+                    .addArgument(progressCursor).log();
+            var workItemAndDuration = workItemRef.get();
+            if (workItemAndDuration == null) {
+                throw new IllegalStateException("Unexpected state with progressCursor set without a" +
+                        "work item");
+            }
+            log.atWarn().setMessage("Work Item and Duration: {}").addArgument(workItemAndDuration)
+                    .log();
+            log.atWarn().setMessage("Work Item: {}").addArgument(workItemAndDuration.getWorkItem())
+                    .log();
+            var successorWorkItemIds = getSuccessorWorkItemIds(workItemAndDuration, progressCursor);
+            log.atWarn().setMessage("Successor Work Ids: {}").addArgument(String.join(", ", successorWorkItemIds))
+                    .log();
+            var successorNextAcquisitionLeaseExponent = getSuccessorNextAcquisitionLeaseExponent(workItemTimeProvider, initialLeaseDuration, workItemAndDuration.getLeaseExpirationTime());
+            coordinator.createSuccessorWorkItemsAndMarkComplete(
+                workItemId,
+                successorWorkItemIds,
+                successorNextAcquisitionLeaseExponent,
+                contextSupplier
+            );
+        } else {
+            log.atWarn().setMessage("No progress cursor to create successor work items from. This can happen when" +
+                    "downloading and unpacking shard takes longer than the lease").log();
+            log.atWarn().setMessage("Skipping creation of successor work item to retry the existing one with more time")
+                    .log();
+        }
+
         System.exit(PROCESS_TIMED_OUT_EXIT_CODE);
     }
 
+    protected static int getSuccessorNextAcquisitionLeaseExponent(WorkItemTimeProvider workItemTimeProvider, Duration initialLeaseDuration,
+                                       Instant leaseExpirationTime) {
+        if (workItemTimeProvider.getLeaseAcquisitionTimeRef().get() == null ||
+            workItemTimeProvider.getDocumentMigraionStartTimeRef().get() == null) {
+            throw new IllegalStateException("Unexpected state with either leaseAquisitionTime or" +
+                    "documentMigrationStartTime as null while creating successor work item");
+        }
+        var leaseAcquisitionTime = workItemTimeProvider.getLeaseAcquisitionTimeRef().get();
+        var documentMigrationStartTime = workItemTimeProvider.getDocumentMigraionStartTimeRef().get();
+        var leaseDuration = Duration.between(leaseAcquisitionTime, leaseExpirationTime);
+        var leaseDurationFactor = (double) leaseDuration.toMillis() / initialLeaseDuration.toMillis();
+        // 2 ^ n = leaseDurationFactor <==> log2(leaseDurationFactor) = n, n >= 0
+        var existingNextAcquisitionLeaseExponent = Math.max(Math.round(Math.log(leaseDurationFactor) / Math.log(2)), 0);
+        var shardSetupDuration = Duration.between(leaseAcquisitionTime, documentMigrationStartTime);
+        var successorShardNextAcquisitionLeaseExponent = (int) (existingNextAcquisitionLeaseExponent + (((double) shardSetupDuration.toMillis() / leaseDuration.toMillis() > SHARD_SETUP_LEASE_DURATION_THRESHOLD) ? 1 : 0));
+
+        log.atDebug().setMessage("SuccessorNextAcquisitionLeaseExponent calculated values:" +
+                "\nleaseAcquisitionTime:{}" +
+                "\ndocumentMigrationStartTime:{}" +
+                "\nleaseDuration:{}" +
+                "\nleaseDurationFactor:{}" +
+                "\nexistingNextAcquisitionLeaseExponent:{}" +
+                "\nshardSetupDuration:{}" +
+                "\nsuccessorShardNextAcquisitionLeaseExponent:{}")
+                .addArgument(leaseAcquisitionTime)
+                .addArgument(documentMigrationStartTime)
+                .addArgument(leaseDuration)
+                .addArgument(leaseDurationFactor)
+                .addArgument(existingNextAcquisitionLeaseExponent)
+                .addArgument(shardSetupDuration)
+                .addArgument(successorShardNextAcquisitionLeaseExponent)
+                .log();
+
+        return successorShardNextAcquisitionLeaseExponent;
+    }
+
+    private static ArrayList<String> getSuccessorWorkItemIds(IWorkCoordinator.WorkItemAndDuration workItemAndDuration, WorkItemCursor progressCursor) {
+        if (workItemAndDuration == null) {
+            throw new IllegalStateException("Unexpected worker coordination state. Expected workItem set when progressCursor not null.");
+        }
+        var workItem = workItemAndDuration.getWorkItem();
+        var successorWorkItem = new IWorkCoordinator.WorkItemAndDuration
+                .WorkItem(workItem.getIndexName(), workItem.getShardNumber(),
+                progressCursor.getDocId() + 1);
+        ArrayList<String> successorWorkItemIds = new ArrayList<>();
+        successorWorkItemIds.add(successorWorkItem.toString());
+        return successorWorkItemIds;
+    }
+
     private static RootDocumentMigrationContext makeRootContext(Args arguments, String workerId) {
         var compositeContextTracker = new CompositeContextTracker(
             new ActiveContextTracker(),
@@ -346,6 +472,7 @@ private static RootDocumentMigrationContext makeRootContext(Args arguments, Stri
 
     public static DocumentsRunner.CompletionStatus run(Function<Path, LuceneDocumentsReader> readerFactory,
                                                        DocumentReindexer reindexer,
+                                                       AtomicReference<WorkItemCursor> progressCursor,
                                                        IWorkCoordinator workCoordinator,
                                                        Duration maxInitialLeaseDuration,
                                                        LeaseExpireTrigger leaseExpireTrigger,
@@ -355,7 +482,9 @@ public static DocumentsRunner.CompletionStatus run(Function<Path, LuceneDocument
                                                        ShardMetadata.Factory shardMetadataFactory,
                                                        SnapshotShardUnpacker.Factory unpackerFactory,
                                                        long maxShardSizeBytes,
-                                                       RootDocumentMigrationContext rootDocumentContext)
+                                                       RootDocumentMigrationContext rootDocumentContext,
+                                                       AtomicReference<Runnable> cancellationRunnable,
+                                                       WorkItemTimeProvider timeProvider)
         throws IOException, InterruptedException, NoWorkLeftException
     {
         var scopedWorkCoordinator = new ScopedWorkCoordinator(workCoordinator, leaseExpireTrigger);
@@ -370,14 +499,22 @@ public static DocumentsRunner.CompletionStatus run(Function<Path, LuceneDocument
         )) {
             throw new NoWorkLeftException("No work items are pending/all work items have been processed.  Returning.");
         }
-        var runner = new DocumentsRunner(scopedWorkCoordinator, maxInitialLeaseDuration, (name, shard) -> {
-            var shardMetadata = shardMetadataFactory.fromRepo(snapshotName, name, shard);
-            log.info("Shard size: " + shardMetadata.getTotalSizeBytes());
-            if (shardMetadata.getTotalSizeBytes() > maxShardSizeBytes) {
-                throw new DocumentsRunner.ShardTooLargeException(shardMetadata.getTotalSizeBytes(), maxShardSizeBytes);
-            }
-            return shardMetadata;
-        }, unpackerFactory, readerFactory, reindexer);
+        var runner = new DocumentsRunner(scopedWorkCoordinator,
+            maxInitialLeaseDuration,
+            reindexer,
+            unpackerFactory,
+            (name, shard) -> {
+                var shardMetadata = shardMetadataFactory.fromRepo(snapshotName, name, shard);
+                log.info("Shard size: " + shardMetadata.getTotalSizeBytes());
+                if (shardMetadata.getTotalSizeBytes() > maxShardSizeBytes) {
+                    throw new DocumentsRunner.ShardTooLargeException(shardMetadata.getTotalSizeBytes(), maxShardSizeBytes);
+                }
+                return shardMetadata;
+            },
+            readerFactory,
+            progressCursor::set,
+            cancellationRunnable::set,
+            timeProvider);
         return runner.migrateNextShard(rootDocumentContext::createReindexContext);
     }
 

DocumentsFromSnapshotMigration/src/test/java/org/opensearch/migrations/RfsMigrateDocumentsTest.java

@@ -0,0 +1,75 @@
+package org.opensearch.migrations;
+
+import java.time.Duration;
+import java.time.Instant;
+
+import org.opensearch.migrations.bulkload.workcoordination.WorkItemTimeProvider;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+
+class RfsMigrateDocumentsTest {
+
+
+    private static class TestClass extends RfsMigrateDocuments {
+        public static int getSuccessorNextAcquisitionLeaseExponent(WorkItemTimeProvider workItemTimeProvider, Duration initialLeaseDuration,
+                                                   Instant leaseExpirationTime) {
+            return RfsMigrateDocuments.getSuccessorNextAcquisitionLeaseExponent(workItemTimeProvider, initialLeaseDuration, leaseExpirationTime);
+        }
+    }
+
+    @Test
+    public void testGetSuccessorNextAcquisitionLeaseExponent_LessThanThreshold() {
+        WorkItemTimeProvider workItemTimeProvider = new WorkItemTimeProvider();
+
+        var shardPrepTime = Duration.ofSeconds(59);
+        var initialShardAttempts = 0;
+        var initialLeaseMultiple =  (int) Math.pow(2, initialShardAttempts);
+
+        workItemTimeProvider.getLeaseAcquisitionTimeRef().set(Instant.EPOCH);
+        workItemTimeProvider.getDocumentMigraionStartTimeRef().set(Instant.EPOCH.plus(shardPrepTime));
+        Duration initialLeaseDuration = Duration.ofMinutes(10);
+        Instant leaseExpirationTime = Instant.EPOCH.plus(initialLeaseDuration.multipliedBy(initialLeaseMultiple));
+
+        int successorAttempts = TestClass.getSuccessorNextAcquisitionLeaseExponent(workItemTimeProvider, initialLeaseDuration, leaseExpirationTime);
+
+        Assertions.assertEquals(initialShardAttempts, successorAttempts, "Should return initialShardAttempts + 1 when shard prep time is less than 10% of lease duration");
+    }
+
+    @Test
+    public void testGetSuccessorNextAcquisitionLeaseExponent_EqualToThreshold() {
+        WorkItemTimeProvider workItemTimeProvider = new WorkItemTimeProvider();
+
+        var shardPrepTime = Duration.ofSeconds(60);
+        var initialShardAttempts = 0;
+        var initialLeaseMultiple =  (int) Math.pow(2, initialShardAttempts);
+
+        workItemTimeProvider.getLeaseAcquisitionTimeRef().set(Instant.EPOCH);
+        workItemTimeProvider.getDocumentMigraionStartTimeRef().set(Instant.EPOCH.plus(shardPrepTime));
+        Duration initialLeaseDuration = Duration.ofMinutes(10);
+        Instant leaseExpirationTime = Instant.EPOCH.plus(initialLeaseDuration.multipliedBy(initialLeaseMultiple));
+
+        int successorAttempts = TestClass.getSuccessorNextAcquisitionLeaseExponent(workItemTimeProvider, initialLeaseDuration, leaseExpirationTime);
+
+        Assertions.assertEquals(initialShardAttempts, successorAttempts, "Should return initialShardAttempts when shard prep time is equal to 10% of lease duration");
+    }
+
+    @Test
+    public void testGetSuccessorNextAcquisitionLeaseExponent_ExceedsThreshold() {
+        WorkItemTimeProvider workItemTimeProvider = new WorkItemTimeProvider();
+
+        var shardPrepTime = Duration.ofSeconds(61);
+        var initialShardAttempts = 0;
+        var initialLeaseMultiple =  (int) Math.pow(2, initialShardAttempts);
+
+        workItemTimeProvider.getLeaseAcquisitionTimeRef().set(Instant.EPOCH);
+        workItemTimeProvider.getDocumentMigraionStartTimeRef().set(Instant.EPOCH.plus(shardPrepTime));
+        Duration initialLeaseDuration = Duration.ofMinutes(10);
+        Instant leaseExpirationTime = Instant.EPOCH.plus(initialLeaseDuration.multipliedBy(initialLeaseMultiple));
+
+        int successorAttempts = TestClass.getSuccessorNextAcquisitionLeaseExponent(workItemTimeProvider, initialLeaseDuration, leaseExpirationTime);
+
+        Assertions.assertEquals(initialShardAttempts + 1, successorAttempts, "Should return initialShardAttempts + 1 when shard prep time is greater than to 10% of lease duration");
+    }
+}