538: Added transform context used for handling the last succesfull page size and transform lock

Signed-off-by: Stevan Buzejic <buzejic.stevan@gmail.com>

Author: stevanbz

src/main/kotlin/org/opensearch/indexmanagement/opensearchapi/OpenSearchExtensions.kt

@@ -48,9 +48,11 @@ import org.opensearch.indexmanagement.indexstatemanagement.action.ShrinkAction
 import org.opensearch.indexmanagement.indexstatemanagement.model.ISMTemplate
 import org.opensearch.indexmanagement.indexstatemanagement.model.Policy
 import org.opensearch.indexmanagement.snapshotmanagement.model.SMMetadata
+import org.opensearch.indexmanagement.transform.util.TransformLockManager
 import org.opensearch.indexmanagement.util.NO_ID
 import org.opensearch.indexmanagement.util.SecurityUtils.Companion.DEFAULT_INJECT_ROLES
 import org.opensearch.indexmanagement.util.SecurityUtils.Companion.INTERNAL_REQUEST
+import org.opensearch.jobscheduler.spi.LockModel
 import org.opensearch.jobscheduler.spi.utils.LockService
 import org.opensearch.rest.RestStatus
 import org.opensearch.transport.RemoteTransportException
@@ -62,6 +64,7 @@ import kotlin.coroutines.resumeWithException
 import kotlin.coroutines.suspendCoroutine
 
 const val OPENDISTRO_SECURITY_PROTECTED_INDICES_CONF_REQUEST = "_opendistro_security_protected_indices_conf_request"
+private const val TIME_EXCEED_MESSAGE = "Time exceeded"
 
 fun contentParser(bytesReference: BytesReference): XContentParser {
     return XContentHelper.createParser(
@@ -184,6 +187,48 @@ suspend fun <T> BackoffPolicy.retry(
     } while (true)
 }
 
+/**
+ * Retries the given [block] of code as specified by the receiver [BackoffPolicy],
+ * if [block] throws an [OpenSearchException] that is retriable (502, 503, 504 or 500 with message Time exceeded).
+ *
+ * If all retries fail the final exception will be rethrown. Exceptions caught during intermediate retries are
+ * logged as warnings to [logger]. Similar to [org.opensearch.action.bulk.Retry], except these retries on
+ * 502, 503, 504, 429 error codes as well as 500 with Time exceeded. If the request is timeout, lock will be renewed
+ *
+ * @param logger - logger used to log intermediate failures
+ * @param transformLockManager - lock manager that stores current lock used in order to renew the lock if the request timed out
+ * @param retryOn - any additional [RestStatus] values that should be retried
+ * @param block - the block of code to retry. This should be a suspend function.
+ */
+suspend fun <T> BackoffPolicy.retry(
+    logger: Logger,
+    transformLockManager: TransformLockManager,
+    retryOn: List<RestStatus> = emptyList(),
+    block: suspend (backoff: TimeValue) -> T
+): T {
+    val iter = iterator()
+    var backoff: TimeValue = TimeValue.ZERO
+    do {
+        try {
+            return block(backoff)
+        } catch (e: OpenSearchException) {
+            if (iter.hasNext() && (e.isRetryable() || e.isTimedOut() || retryOn.contains(e.status()))) {
+                backoff = iter.next()
+                logger.warn("Operation failed. Retrying in $backoff.", e)
+                delay(backoff.millis)
+                // In the case of time out, renew the lock
+                if(e.isTimedOut()) {
+                    transformLockManager.renewLockForScheduledJob()
+                }
+            } else {
+                throw e
+            }
+        }
+    } while (true)
+}
+
+fun LockModel.lockExpirationInSeconds() = lockTime.epochSecond + lockDurationSeconds -  Instant.now().epochSecond
+
 /**
  * Retries on 502, 503 and 504 per elastic client's behavior: https://github.com/elastic/elasticsearch-net/issues/2061
  * 429 must be retried manually as it's not clear if it's ok to retry for requests other than Bulk requests.
@@ -192,6 +237,14 @@ fun OpenSearchException.isRetryable(): Boolean {
     return (status() in listOf(RestStatus.BAD_GATEWAY, RestStatus.SERVICE_UNAVAILABLE, RestStatus.GATEWAY_TIMEOUT))
 }
 
+/**
+ * Retries on 500 and Time exceeded message which means that the timeout occurred. In that case
+ * retry request with reduced size param and timeout param set based on the lock expiration
+ */
+fun OpenSearchException.isTimedOut(): Boolean {
+    return status()  == RestStatus.INTERNAL_SERVER_ERROR && TIME_EXCEED_MESSAGE == message
+}
+
 /**
  * Extension function for OpenSearch 6.3 and above that duplicates the OpenSearch 6.2 XContentBuilder.string() method.
  */

src/main/kotlin/org/opensearch/indexmanagement/transform/TransformRunner.kt

@@ -11,13 +11,11 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.launch
 import org.apache.logging.log4j.LogManager
-import org.opensearch.action.bulk.BackoffPolicy
 import org.opensearch.action.support.WriteRequest
 import org.opensearch.client.Client
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver
 import org.opensearch.cluster.service.ClusterService
 import org.opensearch.common.settings.Settings
-import org.opensearch.common.unit.TimeValue
 import org.opensearch.common.xcontent.NamedXContentRegistry
 import org.opensearch.index.shard.ShardId
 import org.opensearch.indexmanagement.opensearchapi.IndexManagementSecurityContext
@@ -31,10 +29,8 @@ import org.opensearch.indexmanagement.transform.model.ContinuousTransformStats
 import org.opensearch.indexmanagement.transform.model.Transform
 import org.opensearch.indexmanagement.transform.model.TransformMetadata
 import org.opensearch.indexmanagement.transform.model.initializeShardsToSearch
-import org.opensearch.indexmanagement.transform.settings.TransformSettings
-import org.opensearch.indexmanagement.util.acquireLockForScheduledJob
-import org.opensearch.indexmanagement.util.releaseLockForScheduledJob
-import org.opensearch.indexmanagement.util.renewLockForScheduledJob
+import org.opensearch.indexmanagement.transform.util.TransformContext
+import org.opensearch.indexmanagement.transform.util.TransformLockManager
 import org.opensearch.jobscheduler.spi.JobExecutionContext
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner
@@ -108,18 +104,18 @@ object TransformRunner :
         var newGlobalCheckpoints: Map<ShardId, Long>? = null
         var newGlobalCheckpointTime: Instant? = null
         var currentMetadata = metadata
-        val backoffPolicy = BackoffPolicy.exponentialBackoff(
-            TimeValue.timeValueMillis(TransformSettings.DEFAULT_RENEW_LOCK_RETRY_DELAY),
-            TransformSettings.DEFAULT_RENEW_LOCK_RETRY_COUNT
-        )
+
         val transformProcessedBucketLog = TransformProcessedBucketLog()
         var bucketsToTransform = BucketsToTransform(HashSet(), metadata)
-        var lock = acquireLockForScheduledJob(transform, context, backoffPolicy)
+
+        val transformContext = TransformContext.initTransformContext(transform, context)
+        val transformLockManager = transformContext.transformLockManager
         try {
             do {
                 when {
-                    lock == null -> {
-                        logger.warn("Cannot acquire lock for transform job ${transform.id}")
+                    transformLockManager.getLock() == null -> {
+                        val thread = Thread.currentThread().id
+                        logger.warn("Cannot acquire lock for transform job ${transform.id} and thread $thread")
                         // If we fail to get the lock we won't fail the job, instead we return early
                         return
                     }
@@ -128,6 +124,9 @@ object TransformRunner :
                         return
                     }
                     else -> {
+                        val thread = Thread.currentThread().id
+                        logger.warn("Thread $thread acquired a lock")
+
                         val validatedMetadata = validateTransform(transform, currentMetadata)
                         if (validatedMetadata.status == TransformMetadata.Status.FAILED) {
                             currentMetadata = validatedMetadata
@@ -147,15 +146,15 @@ object TransformRunner :
                             // If there are shards to search do it here
                             if (bucketsToTransform.currentShard != null) {
                                 // Computes aggregation on modified documents for current shard to get modified buckets
-                                bucketsToTransform = getBucketsToTransformIteration(transform, bucketsToTransform).also {
+                                bucketsToTransform = getBucketsToTransformIteration(transform, bucketsToTransform, transformContext).also {
                                     currentMetadata = it.metadata
                                 }
                                 // Filter out already processed buckets
                                 val modifiedBuckets = bucketsToTransform.modifiedBuckets.filter {
                                     transformProcessedBucketLog.isNotProcessed(it)
                                 }.toMutableSet()
                                 // Recompute modified buckets and update them in targetIndex
-                                currentMetadata = recomputeModifiedBuckets(transform, currentMetadata, modifiedBuckets)
+                                currentMetadata = recomputeModifiedBuckets(transform, currentMetadata, modifiedBuckets, transformContext)
                                 // Add processed buckets to 'processed set' so that we don't try to reprocess them again
                                 transformProcessedBucketLog.addBuckets(modifiedBuckets.toList())
                                 // Update TransformMetadata
@@ -164,16 +163,12 @@ object TransformRunner :
                             }
                         } else {
                             // Computes buckets from source index and stores them in targetIndex as transform docs
-                            currentMetadata = computeBucketsIteration(transform, currentMetadata)
+                            currentMetadata = computeBucketsIteration(transform, currentMetadata, transformContext)
                             // Update TransformMetadata
                             currentMetadata = transformMetadataService.writeMetadata(currentMetadata, true)
                         }
                         // we attempt to renew lock for every loop of transform
-                        val renewedLock = renewLockForScheduledJob(context, lock, backoffPolicy)
-                        if (renewedLock == null) {
-                            releaseLockForScheduledJob(context, lock)
-                        }
-                        lock = renewedLock
+                       transformLockManager.renewLockForScheduledJob()
                     }
                 }
             } while (bucketsToTransform.currentShard != null || currentMetadata.afterKey != null)
@@ -185,7 +180,7 @@ object TransformRunner :
                 failureReason = e.localizedMessage
             )
         } finally {
-            lock?.let {
+            if(transformLockManager.getLock() != null){
                 // Update the global checkpoints only after execution finishes successfully
                 if (transform.continuous && currentMetadata.status != TransformMetadata.Status.FAILED) {
                     currentMetadata = currentMetadata.copy(
@@ -198,20 +193,20 @@ object TransformRunner :
                     logger.info("Disabling the transform job ${transform.id}")
                     updateTransform(transform.copy(enabled = false, enabledAt = null))
                 }
-                releaseLockForScheduledJob(context, it)
+                transformLockManager.releaseLockForScheduledJob()
             }
         }
     }
 
-    private suspend fun getBucketsToTransformIteration(transform: Transform, bucketsToTransform: BucketsToTransform): BucketsToTransform {
+    private suspend fun getBucketsToTransformIteration(transform: Transform, bucketsToTransform: BucketsToTransform, transformContext: TransformContext): BucketsToTransform {
         var currentBucketsToTransform = bucketsToTransform
         val currentShard = bucketsToTransform.currentShard
         // Clear modified buckets from previous iteration
         currentBucketsToTransform.modifiedBuckets.clear()
 
         if (currentShard != null) {
             val shardLevelModifiedBuckets = withTransformSecurityContext(transform) {
-                transformSearchService.getShardLevelModifiedBuckets(transform, currentBucketsToTransform.metadata.afterKey, currentShard)
+                transformSearchService.getShardLevelModifiedBuckets(transform, currentBucketsToTransform.metadata.afterKey, currentShard, transformContext)
             }
             currentBucketsToTransform.modifiedBuckets.addAll(shardLevelModifiedBuckets.modifiedBuckets)
             val mergedSearchTime = currentBucketsToTransform.metadata.stats.searchTimeInMillis +
@@ -258,13 +253,15 @@ object TransformRunner :
     private suspend fun computeBucketsIteration(
         transform: Transform,
         metadata: TransformMetadata,
+        transformContext: TransformContext
     ): TransformMetadata {
 
         val transformSearchResult = withTransformSecurityContext(transform) {
             transformSearchService.executeCompositeSearch(
                 transform,
                 metadata.afterKey,
-                null
+                null,
+                transformContext
             )
         }
         val indexTimeInMillis = withTransformSecurityContext(transform) {
@@ -287,11 +284,15 @@ object TransformRunner :
     private suspend fun recomputeModifiedBuckets(
         transform: Transform,
         metadata: TransformMetadata,
-        modifiedBuckets: MutableSet<Map<String, Any>>
+        modifiedBuckets: MutableSet<Map<String, Any>>,
+        transformContext: TransformContext
     ): TransformMetadata {
         val updatedMetadata = if (modifiedBuckets.isNotEmpty()) {
+
+            Thread.sleep(61000)
+
             val transformSearchResult = withTransformSecurityContext(transform) {
-                transformSearchService.executeCompositeSearch(transform, null, modifiedBuckets)
+                transformSearchService.executeCompositeSearch(transform, null, modifiedBuckets, transformContext)
             }
             val indexTimeInMillis = withTransformSecurityContext(transform) {
                 transformIndexer.index(transformSearchResult.docsToIndex)