538: Created transfromContext class responsible for holding the context of the current transform job execution

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

Author: stevanbz

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

@@ -52,19 +52,22 @@ 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.tasks.TaskCancelledException
 import org.opensearch.transport.RemoteTransportException
 import java.io.IOException
 import java.time.Instant
+import java.util.regex.Pattern
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.resume
 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"
+
+// Timeout pattern used for checking the timeout message which is in unique format if the transform search timeout was set programmatically
+private val timeoutMessagePattern = Pattern.compile("cancelled task with reason: Cancellation timeout of (.*) is expired")
 
 fun contentParser(bytesReference: BytesReference): XContentParser {
     return XContentHelper.createParser(
@@ -193,14 +196,14 @@ suspend fun <T> BackoffPolicy.retry(
  *
  * 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
+ * 502, 503, 504 error codes as well as when TaskCancelledException is being raised as cause. 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(
+suspend fun <T> BackoffPolicy.retryTransformSearch(
     logger: Logger,
     transformLockManager: TransformLockManager,
     retryOn: List<RestStatus> = emptyList(),
@@ -212,22 +215,24 @@ suspend fun <T> BackoffPolicy.retry(
         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 {
+            if (!iter.hasNext() || !e.isRetryable(retryOn, e)) {
                 throw e
             }
+            backoff = iter.next()
+            logger.warn("Operation failed. Retrying in $backoff.", e)
+            delay(backoff.millis)
+            if (isTransformOperationTimedOut(e)) {
+                // In the case of time out, renew the lock
+                transformLockManager.renewLockForScheduledJob()
+            }
         }
     } while (true)
 }
 
-fun LockModel.lockExpirationInSeconds() = lockTime.epochSecond + lockDurationSeconds -  Instant.now().epochSecond
+private fun OpenSearchException.isRetryable(
+    retryOn: List<RestStatus>,
+    ex: OpenSearchException
+) = this.isRetryable() || isTransformOperationTimedOut(ex) || retryOn.contains(this.status())
 
 /**
  * Retries on 502, 503 and 504 per elastic client's behavior: https://github.com/elastic/elasticsearch-net/issues/2061
@@ -238,11 +243,17 @@ fun OpenSearchException.isRetryable(): Boolean {
 }
 
 /**
- * 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
+ * Retries on 408 or on TaskCancelledException once the message matches the given pattern.
+ * 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
+fun isTransformOperationTimedOut(ex: OpenSearchException): Boolean {
+    if (RestStatus.REQUEST_TIMEOUT == ex.status()) {
+        return true
+    }
+    if (ex.cause != null && ex.cause is TaskCancelledException) {
+        return timeoutMessagePattern.matcher((ex.cause as TaskCancelledException).message).matches()
+    }
+    return false
 }
 
 /**

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

@@ -108,25 +108,22 @@ object TransformRunner :
         val transformProcessedBucketLog = TransformProcessedBucketLog()
         var bucketsToTransform = BucketsToTransform(HashSet(), metadata)
 
-        val transformContext = TransformContext.initTransformContext(transform, context)
+        val transformContext = TransformContext(TransformLockManager(transform, context))
+        // Acquires the lock if there is no running job execution for the given transform; Lock is acquired per transform
         val transformLockManager = transformContext.transformLockManager
+        transformLockManager.acquireLockForScheduledJob()
         try {
             do {
                 when {
-                    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
+                    transformLockManager.lock == null -> {
+                        logger.warn("Cannot acquire lock for transform job ${transform.id}")
                         return
                     }
                     listOf(TransformMetadata.Status.STOPPED, TransformMetadata.Status.FINISHED).contains(metadata.status) -> {
                         logger.warn("Transform job ${transform.id} is in ${metadata.status} status. Skipping execution")
                         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
@@ -168,7 +165,7 @@ object TransformRunner :
                             currentMetadata = transformMetadataService.writeMetadata(currentMetadata, true)
                         }
                         // we attempt to renew lock for every loop of transform
-                       transformLockManager.renewLockForScheduledJob()
+                        transformLockManager.renewLockForScheduledJob()
                     }
                 }
             } while (bucketsToTransform.currentShard != null || currentMetadata.afterKey != null)
@@ -180,7 +177,7 @@ object TransformRunner :
                 failureReason = e.localizedMessage
             )
         } finally {
-            if(transformLockManager.getLock() != null){
+            transformLockManager.lock?.let {
                 // Update the global checkpoints only after execution finishes successfully
                 if (transform.continuous && currentMetadata.status != TransformMetadata.Status.FAILED) {
                     currentMetadata = currentMetadata.copy(
@@ -198,15 +195,24 @@ object TransformRunner :
         }
     }
 
-    private suspend fun getBucketsToTransformIteration(transform: Transform, bucketsToTransform: BucketsToTransform, transformContext: TransformContext): 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, transformContext)
+                transformSearchService.getShardLevelModifiedBuckets(
+                    transform,
+                    currentBucketsToTransform.metadata.afterKey,
+                    currentShard,
+                    transformContext
+                )
             }
             currentBucketsToTransform.modifiedBuckets.addAll(shardLevelModifiedBuckets.modifiedBuckets)
             val mergedSearchTime = currentBucketsToTransform.metadata.stats.searchTimeInMillis +
@@ -288,9 +294,6 @@ object TransformRunner :
         transformContext: TransformContext
     ): TransformMetadata {
         val updatedMetadata = if (modifiedBuckets.isNotEmpty()) {
-
-            Thread.sleep(61000)
-
             val transformSearchResult = withTransformSecurityContext(transform) {
                 transformSearchService.executeCompositeSearch(transform, null, modifiedBuckets, transformContext)
             }