Implement handling strategy for retryable vs non-retryable exceptons in workerPartition

Signed-off-by: Vecheka Chhourn <[email protected]>

Author: vecheka

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365SourceConfigTest.java

@@ -101,4 +101,10 @@ void testNegativeDurationRange() throws Exception {
 
         assertEquals(0, config.getLookBackHours());
     }
+
+    @Test
+    void testDefaultDurationValues() {
+        assertEquals(Duration.ofDays(30), config.getDurationToGiveUpRetry());
+        assertEquals(Duration.ofDays(1), config.getDurationToDelayRetry());
+    }
 }

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/CrawlerSourceConfig.java

@@ -1,12 +1,21 @@
 package org.opensearch.dataprepper.plugins.source.source_crawler.base;
 
+import java.time.Duration;
+
 /**
  * Marker interface to all the SAAS connectors configuration
  */
 public interface CrawlerSourceConfig {
 
     int DEFAULT_NUMBER_OF_WORKERS = 1;
 
+    /*
+     * Retry settings for non-retrayble exceptions in workerPartition
+     * default to 30 days to giveup retry; and 1 day to delay retry
+     */
+    Duration DEFAULT_MAX_DURATION_TO_GIVEUP_RETRY = Duration.ofDays(30);
+    Duration DEFAULT_MAX_DURATION_TO_DELAY_RETRY = Duration.ofDays(1);
+
     /**
      * Number of worker threads enabled for this source
      *
@@ -20,4 +29,20 @@ public interface CrawlerSourceConfig {
      * @return boolean indicating acknowledgement state
      */
     boolean isAcknowledgments();
+
+    /**
+     * Duration to give up retrying workerPartition's work on non-retrayble exceptions
+     * @return Duration indicating max duration to give up retrying 
+     */
+    default Duration getDurationToGiveUpRetry() {
+        return DEFAULT_MAX_DURATION_TO_GIVEUP_RETRY;
+    }
+
+    /**
+     * Duration to retry workerPartition's work on non-retrayble exceptions
+     * @return Duration indicating max duration to delay retrying
+     */
+    default Duration getDurationToDelayRetry() {
+        return DEFAULT_MAX_DURATION_TO_DELAY_RETRY;
+    }
 }

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/coordination/scheduler/WorkerScheduler.java

@@ -12,14 +12,17 @@
 import org.opensearch.dataprepper.plugins.source.source_crawler.base.Crawler;
 import org.opensearch.dataprepper.plugins.source.source_crawler.base.CrawlerSourceConfig;
 import org.opensearch.dataprepper.plugins.source.source_crawler.base.SaasWorkerProgressState;
+import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.state.DimensionalTimeSliceWorkerProgressState;
 import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.partition.SaasSourcePartition;
+import org.opensearch.dataprepper.plugins.source.source_crawler.exception.SaaSCrawlerException;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.time.Duration;
 import java.util.Optional;
+import java.time.Instant;
 
 /**
  * Worker class for executing the partitioned work created while crawling a source.
@@ -75,10 +78,10 @@ public void run() {
         log.info("Worker thread started");
         log.info("Processing Partitions");
         while (!Thread.currentThread().isInterrupted()) {
+            Optional<EnhancedSourcePartition> partition = Optional.empty();
             try {
                 // Get the next available partition from the coordinator
-                Optional<EnhancedSourcePartition> partition =
-                        sourceCoordinator.acquireAvailablePartition(SaasSourcePartition.PARTITION_TYPE);
+                partition = sourceCoordinator.acquireAvailablePartition(SaasSourcePartition.PARTITION_TYPE);
                 if (partition.isPresent()) {
                     // Process the partition (source extraction logic)
                     processPartition(partition.get(), buffer);
@@ -94,8 +97,18 @@ public void run() {
                     }
                 }
             } catch (Exception e) {
-                // TODO: will be in a followup to handle retry strategy differently for non-retryable exceptions
-                backoffRetry(e);
+                this.parititionsFailedCounter.increment();
+                // always default to backoffRetry strategy
+                boolean shouldUseBackoffRetry = true;
+                if (e instanceof SaaSCrawlerException) {
+                    SaaSCrawlerException saasException = (SaaSCrawlerException) e;
+                    if (!saasException.isRetryable()) {
+                        shouldUseBackoffRetry = delayRetry(partition, e);
+                    }
+                }
+                if (shouldUseBackoffRetry) {
+                    backoffRetry(e);
+                }
             }
         }
         log.warn("SourceItemWorker Scheduler is interrupted, looks like shutdown has triggered");
@@ -106,15 +119,67 @@ public void run() {
      * @param e - exception thrown by workerScheduler
      */
     private void backoffRetry(Exception e) {
-        this.parititionsFailedCounter.increment();
-        log.error("Error processing partition", e);
+        log.error("[Retryable Exception] Error processing partition", e);
         try {
             Thread.sleep(RETRY_BACKOFF_ON_EXCEPTION_MILLIS);
         } catch (InterruptedException ex) {
             log.warn("Thread interrupted while waiting to retry due to {}", ex.getMessage());
         }
     }
 
+    /**
+     * Delay retry by X Duration (current default = 1 day) for all non-retryble exceptions up to X days (current default = 30 days)
+     * @param sourcePartition - information on WorkerPartition state
+     * @param ex - exception thrown by workerScheduler
+     * @return boolean: true if we should fallback to backoffRetry
+     */
+    private boolean delayRetry(Optional<EnhancedSourcePartition> sourcePartition, Exception ex) {
+        log.error("[Non-Retryable Exception] Error processing worker partition. Will delay retry with the configured duration", ex);
+        try {
+            SaasSourcePartition workerPartition = (SaasSourcePartition) sourcePartition.get();
+            boolean isWorkerPartitionLeaseExtended = false;
+            if (workerPartition != null) {
+                SaasWorkerProgressState progressState = (SaasWorkerProgressState) workerPartition.getProgressState().get();
+                // TODO: ideally we should add partitionCreationTime for all type of SaasWorkerProgressState
+                if (progressState instanceof DimensionalTimeSliceWorkerProgressState) {
+                    DimensionalTimeSliceWorkerProgressState workerProgressState = (DimensionalTimeSliceWorkerProgressState) progressState;
+                    updateWorkerPartition(workerProgressState.getPartitionCreationTime(), workerPartition);
+                    isWorkerPartitionLeaseExtended = true;
+                }
+            }
+
+            // other SaasWorkerProgressState types (not DimensionalTimeSliceWorkerProgressState) should never use delayRetry()
+            // to be safe, fallback to default retry strategy
+            if (!isWorkerPartitionLeaseExtended) {
+                return true;
+            }
+            return false;
+        } catch (Exception e) {
+            log.error("Error updating workerPartition ", e);
+            // on exception, do not interrupt thread and retry again
+            return false;
+        }
+    }
+
+
+    /**
+     * Update the workerPartition if the partitionCreationTime <= max days to keep retrying (current default = 30 days) on nonretryable exceptions.
+     * Otherwise, give up the workerPartition.
+     * @param partitionCreationTime - timestamp in epoch when the worker partition was first created
+     * @param workerPartition - information on WorkerPartition state
+     */
+    private void updateWorkerPartition(Instant partitionCreationTime, SaasSourcePartition workerPartition) {
+        log.info("Updating workerPartition {}", workerPartition.getPartitionKey());
+        Duration age = Duration.between(partitionCreationTime, Instant.now());
+        if (age.compareTo(this.sourceConfig.getDurationToGiveUpRetry()) <= 0) {
+            log.info("Partition {} is within or equal to the configured max duration, scheduling retry", workerPartition.getPartitionKey());
+            sourceCoordinator.saveProgressStateForPartition(workerPartition, this.sourceConfig.getDurationToDelayRetry());
+        } else {
+            log.info("Partition {} is older than the configured max duration, giving up", workerPartition.getPartitionKey());
+            sourceCoordinator.giveUpPartition(workerPartition);
+        }
+    }
+
     private void processPartition(EnhancedSourcePartition partition, Buffer<Record<Event>> buffer) {
         // Implement your source extraction logic here
         // Update the partition state or commit the partition as needed